Spaces:
Paused
Paused
| import os | |
| import numpy as np | |
| import itertools | |
| import shutil | |
| from collections import Counter | |
| from pymatgen.core import Structure | |
| from smact import element_dictionary, metals, neutral_ratios | |
| from smact.screening import pauling_test | |
| from pymatgen.analysis.structure_matcher import StructureMatcher | |
| from tqdm import tqdm | |
| # --- FIX: Use relative paths for Hugging Face --- | |
| input_folder = "generation_cifs" | |
| output_folder = "valid_generation_cifs" | |
| unique_folder = "unique_structures" | |
| novel_folder = "novel_structures" | |
| # ------------------------------------------------ | |
| def get_safe_space_group_number(structure): | |
| try: | |
| spg_info = structure.get_space_group_info() | |
| return spg_info[1] if spg_info else None | |
| except Exception: | |
| return None | |
| def structure_validity(structure, cutoff=0.5): | |
| dist_mat = structure.distance_matrix | |
| dist_mat = dist_mat + np.diag(np.ones(dist_mat.shape[0]) * (cutoff + 10.)) | |
| if dist_mat.min() < cutoff or structure.volume < 0.1: | |
| return False | |
| spg_number = get_safe_space_group_number(structure) | |
| if spg_number is None: | |
| return False | |
| return True | |
| def smact_validity(atom_types, use_pauling_test=True, include_alloys=True): | |
| elem_counter = Counter(atom_types) | |
| elems = [(elem, elem_counter[elem]) for elem in sorted(elem_counter.keys())] | |
| comp, elem_counts = list(zip(*elems)) | |
| elem_counts = np.array(elem_counts) | |
| elem_counts = elem_counts / np.gcd.reduce(elem_counts) | |
| count = tuple(elem_counts.astype("int").tolist()) | |
| elem_symbols = tuple(comp) | |
| try: | |
| space = element_dictionary(elem_symbols) | |
| except: | |
| return False # Handle elements not in SMACT | |
| smact_elems = [e[1] for e in space.items()] | |
| electronegs = [e.pauling_eneg for e in smact_elems] | |
| ox_combos = [e.oxidation_states for e in smact_elems] | |
| if len(set(elem_symbols)) == 1: | |
| return True | |
| if include_alloys: | |
| is_metal_list = [elem_s in metals for elem_s in elem_symbols] | |
| if all(is_metal_list): | |
| return True | |
| threshold = np.max(count) | |
| oxn = 1 | |
| for oxc in ox_combos: | |
| oxn *= len(oxc) | |
| if oxn > 1e7: | |
| return False | |
| for ox_states in itertools.product(*ox_combos): | |
| stoichs = [(c,) for c in count] | |
| cn_e, cn_r = neutral_ratios(ox_states, stoichs=stoichs, threshold=threshold) | |
| if cn_e: | |
| if use_pauling_test: | |
| electroneg_OK = pauling_test(ox_states, electronegs) | |
| else: | |
| electroneg_OK = True | |
| if electroneg_OK: | |
| return True | |
| return False | |
| def get_atom_types_from_structure(structure): | |
| return [str(element) for element in structure.species] | |
| # Create directories | |
| os.makedirs(output_folder, exist_ok=True) | |
| os.makedirs(unique_folder, exist_ok=True) | |
| os.makedirs(novel_folder, exist_ok=True) | |
| total = 0 | |
| structure_valid_count = 0 | |
| composition_valid_count = 0 | |
| both_valid_count = 0 | |
| if os.path.exists(input_folder): | |
| for filename in os.listdir(input_folder): | |
| if filename.endswith(".cif"): | |
| total += 1 | |
| full_path = os.path.join(input_folder, filename) | |
| try: | |
| structure = Structure.from_file(full_path) | |
| atom_types = get_atom_types_from_structure(structure) | |
| struct_valid = structure_validity(structure) | |
| comp_valid = smact_validity(atom_types) | |
| if struct_valid: | |
| structure_valid_count += 1 | |
| if comp_valid: | |
| composition_valid_count += 1 | |
| if struct_valid and comp_valid: | |
| both_valid_count += 1 | |
| shutil.copy(full_path, os.path.join(output_folder, filename)) | |
| except Exception as e: | |
| print(f"Error processing {filename}: {e}") | |
| print(f"\n=== Validity Summary ===") | |
| print(f"Total CIFs checked: {total}") | |
| print(f"Structural validity passed: {structure_valid_count}") | |
| print(f"Compositional validity passed: {composition_valid_count}") | |
| print(f"Both valid (saved): {both_valid_count}") | |