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 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 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 self.get_structure() self.get_composition() self.get_validity() self.get_fingerprints() def get_structure(self): if min(self.lengths.tolist()) < 0: self.constructed = False self.invalid_reason = 'non_positive_lattice' 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: # counts crystal as invalid if fingerprint cannot be constructed. 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 = [c.valid for c in self.preds] 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): return self.get_match_rate_and_rms() 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_comp_diversity(self): comp_fps = [c.comp_fp for c in self.valid_samples] comp_fps = CompScaler.transform(comp_fps) comp_div = get_fp_pdist(comp_fps) return {'comp_div': comp_div} def get_struct_diversity(self): return {'struct_div': get_fp_pdist([c.struct_fp for c in self.valid_samples])} 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_comp_diversity()) metrics.update(self.get_struct_diversity()) metrics.update(self.get_density_wdist()) metrics.update(self.get_num_elem_wdist()) metrics.update(self.get_prop_wdist()) print(metrics) metrics.update(self.get_coverage()) return metrics class OptEval(object): def __init__(self, crys, num_opt=100, eval_model_name=None): """ crys is a list of length ( * ), where is the number of different initialization for optimizing crystals, and is the number of saved crystals for each intialzation. default to minimize the property. """ step_opt = int(len(crys) / num_opt) self.crys = crys self.step_opt = step_opt self.num_opt = num_opt self.eval_model_name = eval_model_name def get_success_rate(self): valid_indices = np.array([c.valid for c in self.crys]) valid_indices = valid_indices.reshape(self.step_opt, self.num_opt) valid_x, valid_y = valid_indices.nonzero() props = np.ones([self.step_opt, self.num_opt]) * np.inf valid_crys = [c for c in self.crys if c.valid] if len(valid_crys) == 0: sr_5, sr_10, sr_15 = 0, 0, 0 else: pred_props = prop_model_eval(self.eval_model_name, [ c.dict for c in valid_crys]) percentiles = Percentiles[self.eval_model_name] props[valid_x, valid_y] = pred_props best_props = props.min(axis=0) sr_5 = (best_props <= percentiles[0]).mean() sr_10 = (best_props <= percentiles[1]).mean() sr_15 = (best_props <= percentiles[2]).mean() return {'SR5': sr_5, 'SR10': sr_10, 'SR15': sr_15} def get_metrics(self): return self.get_success_rate() 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) 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 main(args): all_metrics = {} cfg = load_config(args.root_path) eval_model_name = cfg.data.eval_model_name if 'recon' in args.tasks: recon_file_path = get_file_paths(args.root_path, 'recon', args.label) crys_array_list, true_crystal_array_list = get_crystal_array_list( recon_file_path) pred_crys = p_map(lambda x: Crystal(x), crys_array_list) gt_crys = p_map(lambda x: Crystal(x), true_crystal_array_list) rec_evaluator = RecEval(pred_crys, gt_crys) recon_metrics = rec_evaluator.get_metrics() all_metrics.update(recon_metrics) 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) gen_crys = p_map(lambda x: Crystal(x), crys_array_list) if 'recon' not in args.tasks: _, 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) if 'opt' in args.tasks: opt_file_path = get_file_paths(args.root_path, 'opt', args.label) crys_array_list, _ = get_crystal_array_list(opt_file_path) opt_crys = p_map(lambda x: Crystal(x), crys_array_list) opt_evaluator = OptEval(opt_crys, eval_model_name=eval_model_name) opt_metrics = opt_evaluator.get_metrics() all_metrics.update(opt_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) # only overwrite metrics computed in the new run. 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=['recon', 'gen', 'opt']) args = parser.parse_args() main(args)