| from collections import Counter |
| import argparse |
| import os |
| import json |
|
|
| import numpy as np |
| from pathlib import Path |
| from tqdm import tqdm |
| from p_tqdm import p_map |
| from scipy.stats import wasserstein_distance |
| import pandas as pd |
|
|
| from pymatgen.core.structure import Structure |
| from pymatgen.core.composition import Composition |
| from pymatgen.core.lattice import Lattice |
| from pymatgen.analysis.structure_matcher import StructureMatcher |
| from matminer.featurizers.site.fingerprint import CrystalNNFingerprint |
| from matminer.featurizers.composition.composite import ElementProperty |
|
|
| from pyxtal import pyxtal |
|
|
| import pickle |
|
|
| import sys |
| sys.path.append('.') |
|
|
| from diffcsp.eval_utils import ( |
| smact_validity, structure_validity, CompScaler, get_fp_pdist, |
| load_config, load_data, get_crystals_list, prop_model_eval, compute_cov) |
|
|
| CrystalNNFP = CrystalNNFingerprint.from_preset("ops") |
| CompFP = ElementProperty.from_preset('magpie') |
|
|
| Percentiles = { |
| 'mp20': np.array([-3.17562208, -2.82196882, -2.52814761]), |
| 'carbon': np.array([-154.527093, -154.45865733, -154.44206825]), |
| 'perovskite': np.array([0.43924842, 0.61202443, 0.7364607]), |
| } |
|
|
| COV_Cutoffs = { |
| 'mp20': {'struc': 0.4, 'comp': 10.}, |
| 'carbon': {'struc': 0.2, 'comp': 4.}, |
| 'perovskite': {'struc': 0.2, 'comp': 4}, |
| } |
|
|
|
|
| class Crystal(object): |
|
|
| def __init__(self, crys_array_dict): |
| self.frac_coords = crys_array_dict['frac_coords'] |
| self.atom_types = crys_array_dict['atom_types'] |
| self.lengths = crys_array_dict['lengths'] |
| self.angles = crys_array_dict['angles'] |
| self.dict = crys_array_dict |
| if len(self.atom_types.shape) > 1: |
| self.dict['atom_types'] = (np.argmax(self.atom_types, axis=-1) + 1) |
| self.atom_types = (np.argmax(self.atom_types, axis=-1) + 1) |
|
|
| self.get_structure() |
| self.get_composition() |
| self.get_validity() |
| |
|
|
|
|
| def get_structure(self): |
| if min(self.lengths.tolist()) < 0: |
| self.constructed = False |
| self.invalid_reason = 'non_positive_lattice' |
| if np.isnan(self.lengths).any() or np.isnan(self.angles).any() or np.isnan(self.frac_coords).any(): |
| self.constructed = False |
| self.invalid_reason = 'nan_value' |
| else: |
| try: |
| self.structure = Structure( |
| lattice=Lattice.from_parameters( |
| *(self.lengths.tolist() + self.angles.tolist())), |
| species=self.atom_types, coords=self.frac_coords, coords_are_cartesian=False) |
| self.constructed = True |
| except Exception: |
| self.constructed = False |
| self.invalid_reason = 'construction_raises_exception' |
| if self.structure.volume < 0.1: |
| self.constructed = False |
| self.invalid_reason = 'unrealistically_small_lattice' |
|
|
| def get_composition(self): |
| elem_counter = Counter(self.atom_types) |
| composition = [(elem, elem_counter[elem]) |
| for elem in sorted(elem_counter.keys())] |
| elems, counts = list(zip(*composition)) |
| counts = np.array(counts) |
| counts = counts / np.gcd.reduce(counts) |
| self.elems = elems |
| self.comps = tuple(counts.astype('int').tolist()) |
|
|
| def get_validity(self): |
| self.comp_valid = smact_validity(self.elems, self.comps) |
| if self.constructed: |
| self.struct_valid = structure_validity(self.structure) |
| else: |
| self.struct_valid = False |
| self.valid = self.comp_valid and self.struct_valid |
|
|
| def get_fingerprints(self): |
| elem_counter = Counter(self.atom_types) |
| comp = Composition(elem_counter) |
| self.comp_fp = CompFP.featurize(comp) |
| try: |
| site_fps = [CrystalNNFP.featurize( |
| self.structure, i) for i in range(len(self.structure))] |
| except Exception: |
| |
| self.valid = False |
| self.comp_fp = None |
| self.struct_fp = None |
| return |
| self.struct_fp = np.array(site_fps).mean(axis=0) |
|
|
|
|
| class RecEval(object): |
|
|
| def __init__(self, pred_crys, gt_crys, stol=0.5, angle_tol=10, ltol=0.3): |
| assert len(pred_crys) == len(gt_crys) |
| self.matcher = StructureMatcher( |
| stol=stol, angle_tol=angle_tol, ltol=ltol) |
| self.preds = pred_crys |
| self.gts = gt_crys |
|
|
| def get_match_rate_and_rms(self): |
| def process_one(pred, gt, is_valid): |
| if not is_valid: |
| return None |
| try: |
| rms_dist = self.matcher.get_rms_dist( |
| pred.structure, gt.structure) |
| rms_dist = None if rms_dist is None else rms_dist[0] |
| return rms_dist |
| except Exception: |
| return None |
| validity = [c1.valid and c2.valid for c1,c2 in zip(self.preds, self.gts)] |
|
|
| rms_dists = [] |
| for i in tqdm(range(len(self.preds))): |
| rms_dists.append(process_one( |
| self.preds[i], self.gts[i], validity[i])) |
| rms_dists = np.array(rms_dists) |
| match_rate = sum(rms_dists != None) / len(self.preds) |
| mean_rms_dist = rms_dists[rms_dists != None].mean() |
| return {'match_rate': match_rate, |
| 'rms_dist': mean_rms_dist} |
|
|
| def get_metrics(self): |
| metrics = {} |
| metrics.update(self.get_match_rate_and_rms()) |
| return metrics |
|
|
|
|
| class RecEvalBatch(object): |
|
|
| def __init__(self, pred_crys, gt_crys, stol=0.5, angle_tol=10, ltol=0.3): |
| self.matcher = StructureMatcher( |
| stol=stol, angle_tol=angle_tol, ltol=ltol) |
| self.preds = pred_crys |
| self.gts = gt_crys |
| self.batch_size = len(self.preds) |
|
|
| def get_match_rate_and_rms(self): |
| def process_one(pred, gt, is_valid): |
| if not is_valid: |
| return None |
| try: |
| rms_dist = self.matcher.get_rms_dist( |
| pred.structure, gt.structure) |
| rms_dist = None if rms_dist is None else rms_dist[0] |
| return rms_dist |
| except Exception: |
| return None |
|
|
| rms_dists = [] |
| self.all_rms_dis = np.zeros((self.batch_size, len(self.gts))) |
| for i in tqdm(range(len(self.preds[0]))): |
| tmp_rms_dists = [] |
| for j in range(self.batch_size): |
| rmsd = process_one(self.preds[j][i], self.gts[i], self.preds[j][i].valid) |
| self.all_rms_dis[j][i] = rmsd |
| if rmsd is not None: |
| tmp_rms_dists.append(rmsd) |
| if len(tmp_rms_dists) == 0: |
| rms_dists.append(None) |
| else: |
| rms_dists.append(np.min(tmp_rms_dists)) |
| |
| rms_dists = np.array(rms_dists) |
| match_rate = sum(rms_dists != None) / len(self.preds[0]) |
| mean_rms_dist = rms_dists[rms_dists != None].mean() |
| return {'match_rate': match_rate, |
| 'rms_dist': mean_rms_dist} |
|
|
| def get_metrics(self): |
| metrics = {} |
| metrics.update(self.get_match_rate_and_rms()) |
| return metrics |
|
|
|
|
|
|
| class GenEval(object): |
|
|
| def __init__(self, pred_crys, gt_crys, n_samples=1000, eval_model_name=None): |
| self.crys = pred_crys |
| self.gt_crys = gt_crys |
| self.n_samples = n_samples |
| self.eval_model_name = eval_model_name |
|
|
| valid_crys = [c for c in pred_crys if c.valid] |
| if len(valid_crys) >= n_samples: |
| sampled_indices = np.random.choice( |
| len(valid_crys), n_samples, replace=False) |
| self.valid_samples = [valid_crys[i] for i in sampled_indices] |
| else: |
| raise Exception( |
| f'not enough valid crystals in the predicted set: {len(valid_crys)}/{n_samples}') |
|
|
| def get_validity(self): |
| comp_valid = np.array([c.comp_valid for c in self.crys]).mean() |
| struct_valid = np.array([c.struct_valid for c in self.crys]).mean() |
| valid = np.array([c.valid for c in self.crys]).mean() |
| return {'comp_valid': comp_valid, |
| 'struct_valid': struct_valid, |
| 'valid': valid} |
|
|
|
|
| def get_density_wdist(self): |
| pred_densities = [c.structure.density for c in self.valid_samples] |
| gt_densities = [c.structure.density for c in self.gt_crys] |
| wdist_density = wasserstein_distance(pred_densities, gt_densities) |
| return {'wdist_density': wdist_density} |
|
|
|
|
| def get_num_elem_wdist(self): |
| pred_nelems = [len(set(c.structure.species)) |
| for c in self.valid_samples] |
| gt_nelems = [len(set(c.structure.species)) for c in self.gt_crys] |
| wdist_num_elems = wasserstein_distance(pred_nelems, gt_nelems) |
| return {'wdist_num_elems': wdist_num_elems} |
|
|
| def get_prop_wdist(self): |
| if self.eval_model_name is not None: |
| pred_props = prop_model_eval(self.eval_model_name, [ |
| c.dict for c in self.valid_samples]) |
| gt_props = prop_model_eval(self.eval_model_name, [ |
| c.dict for c in self.gt_crys]) |
| wdist_prop = wasserstein_distance(pred_props, gt_props) |
| return {'wdist_prop': wdist_prop} |
| else: |
| return {'wdist_prop': None} |
|
|
| def get_coverage(self): |
| cutoff_dict = COV_Cutoffs[self.eval_model_name] |
| (cov_metrics_dict, combined_dist_dict) = compute_cov( |
| self.crys, self.gt_crys, |
| struc_cutoff=cutoff_dict['struc'], |
| comp_cutoff=cutoff_dict['comp']) |
| return cov_metrics_dict |
|
|
| def get_metrics(self): |
| metrics = {} |
| metrics.update(self.get_validity()) |
| metrics.update(self.get_density_wdist()) |
| |
| metrics.update(self.get_num_elem_wdist()) |
| metrics.update(self.get_coverage()) |
| return metrics |
|
|
|
|
| def get_file_paths(root_path, task, label='', suffix='pt'): |
| if args.label == '': |
| out_name = f'eval_{task}.{suffix}' |
| else: |
| out_name = f'eval_{task}_{label}.{suffix}' |
| out_name = os.path.join(root_path, out_name) |
| return out_name |
|
|
|
|
| def get_crystal_array_list(file_path, batch_idx=0): |
| data = load_data(file_path) |
| if batch_idx == -1: |
| batch_size = data['frac_coords'].shape[0] |
| crys_array_list = [] |
| for i in range(batch_size): |
| tmp_crys_array_list = get_crystals_list( |
| data['frac_coords'][i], |
| data['atom_types'][i], |
| data['lengths'][i], |
| data['angles'][i], |
| data['num_atoms'][i]) |
| crys_array_list.append(tmp_crys_array_list) |
| elif batch_idx == -2: |
| crys_array_list = get_crystals_list( |
| data['frac_coords'], |
| data['atom_types'], |
| data['lengths'], |
| data['angles'], |
| data['num_atoms']) |
| else: |
| crys_array_list = get_crystals_list( |
| data['frac_coords'][batch_idx], |
| data['atom_types'][batch_idx], |
| data['lengths'][batch_idx], |
| data['angles'][batch_idx], |
| data['num_atoms'][batch_idx]) |
|
|
| if 'input_data_batch' in data: |
| batch = data['input_data_batch'] |
| if isinstance(batch, dict): |
| true_crystal_array_list = get_crystals_list( |
| batch['frac_coords'], batch['atom_types'], batch['lengths'], |
| batch['angles'], batch['num_atoms']) |
| else: |
| true_crystal_array_list = get_crystals_list( |
| batch.frac_coords, batch.atom_types, batch.lengths, |
| batch.angles, batch.num_atoms) |
| else: |
| true_crystal_array_list = None |
|
|
| return crys_array_list, true_crystal_array_list |
|
|
|
|
| def get_gt_crys_ori(cif): |
| structure = Structure.from_str(cif,fmt='cif') |
| lattice = structure.lattice |
| crys_array_dict = { |
| 'frac_coords':structure.frac_coords, |
| 'atom_types':np.array([_.Z for _ in structure.species]), |
| 'lengths': np.array(lattice.abc), |
| 'angles': np.array(lattice.angles) |
| } |
| return Crystal(crys_array_dict) |
|
|
| def main(args): |
| all_metrics = {} |
|
|
| cfg = load_config(args.root_path) |
| eval_model_name = cfg.data.eval_model_name |
|
|
| if 'gen' in args.tasks: |
|
|
| gen_file_path = get_file_paths(args.root_path, 'gen', args.label) |
| recon_file_path = get_file_paths(args.root_path, 'recon', args.label) |
| crys_array_list, _ = get_crystal_array_list(gen_file_path, batch_idx = -2) |
| gen_crys = p_map(lambda x: Crystal(x), crys_array_list) |
| if args.gt_file != '': |
| csv = pd.read_csv(args.gt_file) |
| gt_crys = p_map(get_gt_crys_ori, csv['cif']) |
| else: |
| _, true_crystal_array_list = get_crystal_array_list( |
| recon_file_path) |
| gt_crys = p_map(lambda x: Crystal(x), true_crystal_array_list) |
| gen_evaluator = GenEval( |
| gen_crys, gt_crys, eval_model_name=eval_model_name) |
| gen_metrics = gen_evaluator.get_metrics() |
| all_metrics.update(gen_metrics) |
|
|
|
|
| else: |
|
|
| recon_file_path = get_file_paths(args.root_path, 'diff', args.label) |
| batch_idx = -1 if args.multi_eval else 0 |
| crys_array_list, true_crystal_array_list = get_crystal_array_list( |
| recon_file_path, batch_idx = batch_idx) |
| if args.gt_file != '': |
| csv = pd.read_csv(args.gt_file) |
| gt_crys = p_map(get_gt_crys_ori, csv['cif']) |
| else: |
| gt_crys = p_map(lambda x: Crystal(x), true_crystal_array_list) |
|
|
| if not args.multi_eval: |
| pred_crys = p_map(lambda x: Crystal(x), crys_array_list) |
| else: |
| pred_crys = [] |
| for i in range(len(crys_array_list)): |
| print(f"Processing batch {i}") |
| pred_crys.append(p_map(lambda x: Crystal(x), crys_array_list[i])) |
|
|
|
|
| if 'csp' in args.tasks: |
|
|
| if args.multi_eval: |
| rec_evaluator = RecEvalBatch(pred_crys, gt_crys) |
| else: |
| rec_evaluator = RecEval(pred_crys, gt_crys) |
|
|
| recon_metrics = rec_evaluator.get_metrics() |
|
|
| all_metrics.update(recon_metrics) |
|
|
| |
|
|
| print(all_metrics) |
|
|
| if args.label == '': |
| metrics_out_file = 'eval_metrics.json' |
| else: |
| metrics_out_file = f'eval_metrics_{args.label}.json' |
| metrics_out_file = os.path.join(args.root_path, metrics_out_file) |
|
|
| |
| if Path(metrics_out_file).exists(): |
| with open(metrics_out_file, 'r') as f: |
| written_metrics = json.load(f) |
| if isinstance(written_metrics, dict): |
| written_metrics.update(all_metrics) |
| else: |
| with open(metrics_out_file, 'w') as f: |
| json.dump(all_metrics, f) |
| if isinstance(written_metrics, dict): |
| with open(metrics_out_file, 'w') as f: |
| json.dump(written_metrics, f) |
| else: |
| with open(metrics_out_file, 'w') as f: |
| json.dump(all_metrics, f) |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--root_path', required=True) |
| parser.add_argument('--label', default='') |
| parser.add_argument('--tasks', nargs='+', default=['csp', 'gen']) |
| parser.add_argument('--gt_file',default='') |
| parser.add_argument('--multi_eval',action='store_true') |
| args = parser.parse_args() |
| main(args) |
|
|
