diff --git a/omg/analysis/analysis.py b/omg/analysis/analysis.py index 1a628dc..ae330b5 100644 --- a/omg/analysis/analysis.py +++ b/omg/analysis/analysis.py @@ -40,6 +40,18 @@ def get_bonds(atoms: Atoms, covalent_increase_factor: float = 1.25) -> List[List List of bonds for the atoms in the structure. :rtype: List[List[int]] """ + # Filter out ghost atoms before computing bonds + # Ghost atoms are non-physical and should not be included in coordination calculations + if "is_ghost" in atoms.arrays: + is_ghost = atoms.arrays["is_ghost"].astype(bool) + atoms = atoms[~is_ghost] + else: + # Fallback: filter by atomic number (ghost atoms have number <= 0 or > MAX_ATOM_NUM) + valid_mask = (atoms.numbers > 0) & (atoms.numbers <= MAX_ATOM_NUM) + atoms = atoms[valid_mask] + # If all atoms were ghost atoms, return empty results + if len(atoms) == 0: + return [] distances = atoms.get_all_distances(mic=True) cr = [covalent_radii[number] for number in atoms.numbers] @@ -75,6 +87,18 @@ def get_coordination_numbers(atoms: Atoms, covalent_increase_factor: float = 1.2 List of coordination numbers for the atoms in the structure. :rtype: List[int] """ + # Filter out ghost atoms before computing coordination numbers + # Ghost atoms are non-physical and should not be included in coordination calculations + if "is_ghost" in atoms.arrays: + is_ghost = atoms.arrays["is_ghost"].astype(bool) + atoms = atoms[~is_ghost] + else: + # Fallback: filter by atomic number (ghost atoms have number <= 0 or > MAX_ATOM_NUM) + valid_mask = (atoms.numbers > 0) & (atoms.numbers <= MAX_ATOM_NUM) + atoms = atoms[valid_mask] + # If all atoms were ghost atoms, return empty results + if len(atoms) == 0: + return [] bonds = get_bonds(atoms, covalent_increase_factor) return [len(b) for b in bonds] diff --git a/omg/analysis/valid_atoms.py b/omg/analysis/valid_atoms.py index b8c70e2..52dffc6 100644 --- a/omg/analysis/valid_atoms.py +++ b/omg/analysis/valid_atoms.py @@ -76,6 +76,26 @@ def __init__(self, atoms: Atoms, volume_check_cutoff: float = 0.1, structure_che if upper_narity_limit is not None and upper_narity_limit < 1: raise ValueError("The upper n-arity limit must be at least 1.") self._upper_narity_limit = upper_narity_limit + # Filter out ghost atoms before validation + # Ghost atoms are non-physical and should not be included in metrics + if "is_ghost" in atoms.arrays: + is_ghost = atoms.arrays["is_ghost"].astype(bool) + atoms = atoms[~is_ghost] + else: + # Fallback: filter by atomic number (ghost atoms have number <= 0 or > MAX_ATOM_NUM) + valid_mask = (atoms.numbers > 0) & (atoms.numbers <= MAX_ATOM_NUM) + atoms = atoms[valid_mask] + # If all atoms were ghost atoms, mark as invalid + if len(atoms) == 0: + self._volume_valid = False + self._structure_valid = False + self._composition_valid = False + self._fingerprint_valid = False + self._composition_fingerprint, self._structure_fingerprint = None, None + self._atoms = atoms + self._structure = None + self._composition = None + return self._atoms = atoms self._structure = AseAtomsAdaptor.get_structure(atoms) self._composition = self._structure.composition diff --git a/omg/conf_examples/csp_ghosted_linede_mp_20_j05.yaml b/omg/conf_examples/csp_ghosted_linede_mp_20_j05.yaml new file mode 100644 index 0000000..dc00157 --- /dev/null +++ b/omg/conf_examples/csp_ghosted_linede_mp_20_j05.yaml @@ -0,0 +1,109 @@ +model: + si: + class_path: omg.si.stochastic_interpolants.StochasticInterpolants + init_args: + stochastic_interpolants: + # chemical species + - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity + # fractional coordinates + - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant + init_args: + interpolant: omg.si.interpolants.PeriodicLinearInterpolant + gamma: null + epsilon: null + differential_equation_type: "ODE" + integrator_kwargs: + method: "euler" + velocity_annealing_factor: 10.182659004291072 + correct_center_of_mass_motion: true + # lattice vectors + - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant + init_args: + interpolant: omg.si.interpolants.LinearInterpolant + gamma: null + epsilon: null + differential_equation_type: "ODE" + integrator_kwargs: + method: "euler" + velocity_annealing_factor: 1.824475401606087 + correct_center_of_mass_motion: false + data_fields: + # if the order of the data_fields changes, + # the order of the above StochasticInterpolant inputs must also change + - "species" + - "pos" + - "cell" + integration_time_steps: 210 + relative_si_costs: + species_loss: 0.0 + pos_loss_b: 0.9994149341846618 + cell_loss_b: 0.0005850658153382233 + sampler: + class_path: omg.sampler.IndependentSampler + init_args: + pos_distribution: + class_path: omg.sampler.position_distributions.UniformPositionDistribution + cell_distribution: + class_path: omg.sampler.cell_distributions.InformedLatticeDistribution + init_args: + dataset_name: mp_20 + species_distribution: + class_path: omg.sampler.species_distributions.MirrorSpecies + model: + class_path: omg.model.model.Model + init_args: + encoder: + class_path: omg.model.encoders.cspnet_full.CSPNetFull + head: + class_path: omg.model.heads.pass_through.PassThrough + time_embedder: + class_path: omg.model.model_utils.SinusoidalTimeEmbeddings + init_args: + dim: 256 + use_min_perm_dist: False + float_32_matmul_precision: "high" + validation_mode: "loss" + dataset_name: "mp_20_ghosted_j05" +data: + train_dataset: + class_path: omg.datamodule.StructureDataset + init_args: + file_path: "data/mp_20_ghosted_j05/train_ghosted.lmdb" + lazy_storage: True + niggli_reduce: False + val_dataset: + class_path: omg.datamodule.StructureDataset + init_args: + file_path: "data/mp_20_ghosted_j05/val_ghosted.lmdb" + lazy_storage: True + niggli_reduce: False + pred_dataset: + class_path: omg.datamodule.StructureDataset + init_args: + file_path: "data/mp_20_ghosted_j05/test_ghosted.lmdb" + lazy_storage: True + niggli_reduce: False + batch_size: 256 +trainer: + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + filename: "best_val_loss_total" + save_top_k: 1 + monitor: "val_loss_total" + save_weights_only: false + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + save_top_k: -1 # Store every checkpoint after 100 epochs. + monitor: "val_loss_total" + every_n_epochs: 100 + save_weights_only: false + gradient_clip_val: 0.5 + num_sanity_val_steps: 0 + precision: "32-true" + max_epochs: 2000 + enable_progress_bar: true +optimizer: + class_path: torch.optim.Adam + init_args: + lr: 0.0006689636445843722 diff --git a/omg/conf_examples/csp_ghosted_linode_mp_20_j05.yaml b/omg/conf_examples/csp_ghosted_linode_mp_20_j05.yaml new file mode 100644 index 0000000..dc00157 --- /dev/null +++ b/omg/conf_examples/csp_ghosted_linode_mp_20_j05.yaml @@ -0,0 +1,109 @@ +model: + si: + class_path: omg.si.stochastic_interpolants.StochasticInterpolants + init_args: + stochastic_interpolants: + # chemical species + - class_path: omg.si.single_stochastic_interpolant_identity.SingleStochasticInterpolantIdentity + # fractional coordinates + - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant + init_args: + interpolant: omg.si.interpolants.PeriodicLinearInterpolant + gamma: null + epsilon: null + differential_equation_type: "ODE" + integrator_kwargs: + method: "euler" + velocity_annealing_factor: 10.182659004291072 + correct_center_of_mass_motion: true + # lattice vectors + - class_path: omg.si.single_stochastic_interpolant.SingleStochasticInterpolant + init_args: + interpolant: omg.si.interpolants.LinearInterpolant + gamma: null + epsilon: null + differential_equation_type: "ODE" + integrator_kwargs: + method: "euler" + velocity_annealing_factor: 1.824475401606087 + correct_center_of_mass_motion: false + data_fields: + # if the order of the data_fields changes, + # the order of the above StochasticInterpolant inputs must also change + - "species" + - "pos" + - "cell" + integration_time_steps: 210 + relative_si_costs: + species_loss: 0.0 + pos_loss_b: 0.9994149341846618 + cell_loss_b: 0.0005850658153382233 + sampler: + class_path: omg.sampler.IndependentSampler + init_args: + pos_distribution: + class_path: omg.sampler.position_distributions.UniformPositionDistribution + cell_distribution: + class_path: omg.sampler.cell_distributions.InformedLatticeDistribution + init_args: + dataset_name: mp_20 + species_distribution: + class_path: omg.sampler.species_distributions.MirrorSpecies + model: + class_path: omg.model.model.Model + init_args: + encoder: + class_path: omg.model.encoders.cspnet_full.CSPNetFull + head: + class_path: omg.model.heads.pass_through.PassThrough + time_embedder: + class_path: omg.model.model_utils.SinusoidalTimeEmbeddings + init_args: + dim: 256 + use_min_perm_dist: False + float_32_matmul_precision: "high" + validation_mode: "loss" + dataset_name: "mp_20_ghosted_j05" +data: + train_dataset: + class_path: omg.datamodule.StructureDataset + init_args: + file_path: "data/mp_20_ghosted_j05/train_ghosted.lmdb" + lazy_storage: True + niggli_reduce: False + val_dataset: + class_path: omg.datamodule.StructureDataset + init_args: + file_path: "data/mp_20_ghosted_j05/val_ghosted.lmdb" + lazy_storage: True + niggli_reduce: False + pred_dataset: + class_path: omg.datamodule.StructureDataset + init_args: + file_path: "data/mp_20_ghosted_j05/test_ghosted.lmdb" + lazy_storage: True + niggli_reduce: False + batch_size: 256 +trainer: + callbacks: + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + filename: "best_val_loss_total" + save_top_k: 1 + monitor: "val_loss_total" + save_weights_only: false + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + save_top_k: -1 # Store every checkpoint after 100 epochs. + monitor: "val_loss_total" + every_n_epochs: 100 + save_weights_only: false + gradient_clip_val: 0.5 + num_sanity_val_steps: 0 + precision: "32-true" + max_epochs: 2000 + enable_progress_bar: true +optimizer: + class_path: torch.optim.Adam + init_args: + lr: 0.0006689636445843722 diff --git a/omg/data/mp_20_ghosted_j05/test_ghosted.lmdb b/omg/data/mp_20_ghosted_j05/test_ghosted.lmdb new file mode 100644 index 0000000..82af5fa Binary files /dev/null and b/omg/data/mp_20_ghosted_j05/test_ghosted.lmdb differ diff --git a/omg/data/mp_20_ghosted_j05/train_ghosted.lmdb b/omg/data/mp_20_ghosted_j05/train_ghosted.lmdb new file mode 100644 index 0000000..096283d Binary files /dev/null and b/omg/data/mp_20_ghosted_j05/train_ghosted.lmdb differ diff --git a/omg/data/mp_20_ghosted_j05/val_ghosted.lmdb b/omg/data/mp_20_ghosted_j05/val_ghosted.lmdb new file mode 100644 index 0000000..1efb021 Binary files /dev/null and b/omg/data/mp_20_ghosted_j05/val_ghosted.lmdb differ diff --git a/omg/datamodule/omg_data.py b/omg/datamodule/omg_data.py index 833d1ce..c0893cb 100644 --- a/omg/datamodule/omg_data.py +++ b/omg/datamodule/omg_data.py @@ -48,7 +48,13 @@ def __init__(self, structure: Optional[Structure] = None) -> None: self.property = None else: self.n_atoms = torch.tensor(len(structure.atomic_numbers)) # Shape: (1, ) - self.species = structure.atomic_numbers # Shape: (n_atoms, ) + # Handle ghost atoms: convert -1 to fixed global label (119) to keep distinct from mask (0) + species = structure.atomic_numbers.clone() + ghost_mask = species < 0 + if ghost_mask.any(): + GLOBAL_GHOST_LABEL = 119 # Fixed label for all ghost atoms + species[ghost_mask] = GLOBAL_GHOST_LABEL + self.species = species # Shape: (n_atoms, ) self.cell = structure.cell.unsqueeze(0) # Shape: (1, 3, 3). self.pos = structure.pos # Shape: (n_atoms, 3). self.pos_is_fractional = torch.tensor(structure.pos_is_fractional, dtype=torch.bool) # Shape: (1, ) diff --git a/omg/model/encoders/cspnet_full.py b/omg/model/encoders/cspnet_full.py index ee02b87..272133d 100644 --- a/omg/model/encoders/cspnet_full.py +++ b/omg/model/encoders/cspnet_full.py @@ -39,7 +39,10 @@ def __init__( self.ip = ip self.hidden_dim = hidden_dim self.species_shift = 1 - self.node_embedding = nn.Embedding(max_atoms, hidden_dim) + # Embedding size increased to accommodate ghost atoms (label 119) + # With species_shift=1, species 119 maps to index 118, so we need at least 119 entries + # Using max_atoms + 20 = 120 to safely handle species up to 119 + self.node_embedding = nn.Embedding(max_atoms + 20, hidden_dim) self.atom_latent_emb = nn.Linear(hidden_dim + latent_dim, hidden_dim) if act_fn == 'silu': self.act_fn = nn.SiLU() diff --git a/omg/model/encoders/diffcsp_copies.py b/omg/model/encoders/diffcsp_copies.py index 15ec806..fc20479 100644 --- a/omg/model/encoders/diffcsp_copies.py +++ b/omg/model/encoders/diffcsp_copies.py @@ -645,10 +645,13 @@ def __init__( self.ip = ip self.smooth = smooth + # Embedding size increased to accommodate ghost atoms (label 119) + # With species_shift=1, species 119 maps to index 118, so we need at least 119 entries + # Using max_atoms + 20 = 120 to safely handle species up to 119 if self.smooth: - self.node_embedding = nn.Linear(max_atoms, hidden_dim) + self.node_embedding = nn.Linear(max_atoms + 20, hidden_dim) else: - self.node_embedding = nn.Embedding(max_atoms, hidden_dim) + self.node_embedding = nn.Embedding(max_atoms + 20, hidden_dim) self.atom_latent_emb = nn.Linear(hidden_dim + latent_dim, hidden_dim) if act_fn == 'silu': self.act_fn = nn.SiLU() diff --git a/omg/omg_trainer.py b/omg/omg_trainer.py index b308581..3784799 100644 --- a/omg/omg_trainer.py +++ b/omg/omg_trainer.py @@ -135,12 +135,23 @@ def _load_dataset_atoms(dataset: OMGDataset) -> list[Atoms]: assert len(struc.species) == struc.pos.shape[0] assert struc.pos.shape[1] == 3 assert struc.cell[0].shape == (3, 3) + species_np = struc.species.cpu().numpy() + numbers_for_ase = species_np.copy() + # Handle ghost atoms: species=-1 should be converted to max_real_Z+1 for ASE compatibility + is_ghost = species_np == -1 + if is_ghost.any(): + real_numbers = species_np[~is_ghost] + max_real = int(real_numbers.max()) if real_numbers.size else 0 + ghost_label = max_real + 1 + numbers_for_ase[is_ghost] = ghost_label if bool(struc.pos_is_fractional): - atoms = Atoms(numbers=struc.species, scaled_positions=struc.pos, cell=struc.cell[0], + atoms = Atoms(numbers=numbers_for_ase, scaled_positions=struc.pos, cell=struc.cell[0], pbc=(True, True, True)) else: - atoms = Atoms(numbers=struc.species, positions=struc.pos, cell=struc.cell[0], + atoms = Atoms(numbers=numbers_for_ase, positions=struc.pos, cell=struc.cell[0], pbc=(True, True, True)) + if is_ghost.any(): + atoms.set_array("is_ghost", is_ghost) all_ref_atoms.append(atoms) return all_ref_atoms @@ -171,17 +182,40 @@ def _plot_to_pdf(reference: Sequence[Atoms], initial: Optional[Sequence[Atoms]], Filename for the storage of the symmetric structures. :type symmetry_filename: Path """ + def filter_ghost_atoms(atoms_list: Sequence[Atoms]) -> list[Atoms]: + """Filter ghost atoms from a list of Atoms objects.""" + filtered = [] + for atoms in atoms_list: + if "is_ghost" in atoms.arrays: + is_ghost = atoms.arrays["is_ghost"].astype(bool) + atoms = atoms[~is_ghost] + else: + valid_mask = (atoms.numbers > 0) & (atoms.numbers <= MAX_ATOM_NUM) + atoms = atoms[valid_mask] + filtered.append(atoms) + return filtered + fractional_coordinates_corrector = PeriodicBoundaryConditionsCorrector(min_value=0.0, max_value=1.0) + # Filter ghost atoms before visualization + reference_atoms = filter_ghost_atoms(reference) + generated_atoms = filter_ghost_atoms(generated) + if initial is not None: + initial_atoms = filter_ghost_atoms(initial) + else: + initial_atoms = None + # Keep ASE Atoms versions of certain inputs - reference_atoms = reference - generated_atoms = generated + reference_atoms_for_analysis = reference_atoms + generated_atoms_for_analysis = generated_atoms # Convert to Data - reference = convert_ase_atoms_to_data(reference) - if initial is not None: - initial = convert_ase_atoms_to_data(initial) - generated = convert_ase_atoms_to_data(generated) + reference = convert_ase_atoms_to_data(reference_atoms) + if initial_atoms is not None: + initial = convert_ase_atoms_to_data(initial_atoms) + else: + initial = None + generated = convert_ase_atoms_to_data(generated_atoms) # List of volumes of all test structures. ref_vol = [] @@ -253,7 +287,7 @@ def _plot_to_pdf(reference: Sequence[Atoms], initial: Optional[Sequence[Atoms]], ref_root_mean_square_distances.append(float(torch.sqrt(ds.mean()))) ref_sg_fail = 0 - for struc in reference_atoms: + for struc in reference_atoms_for_analysis: ref_avg_cn.append(np.mean(get_coordination_numbers(struc))) cn_dict = get_coordination_numbers_species(struc) @@ -275,7 +309,7 @@ def _plot_to_pdf(reference: Sequence[Atoms], initial: Optional[Sequence[Atoms]], ref_crystal_sys[cs] = 0 ref_crystal_sys[cs] += 1 print("Number of times space group identification failed for prediction dataset: " - "{}/{}".format(ref_sg_fail, len(reference_atoms))) + "{}/{}".format(ref_sg_fail, len(reference_atoms_for_analysis))) # List of volumes of all generated structures. vol = [] @@ -344,7 +378,7 @@ def _plot_to_pdf(reference: Sequence[Atoms], initial: Optional[Sequence[Atoms]], sg_fail = 0 sg_fail_F = 0 - for struc in generated_atoms: + for struc in generated_atoms_for_analysis: avg_cn.append(np.mean(get_coordination_numbers(struc))) cn_dict = get_coordination_numbers_species(struc) @@ -393,9 +427,9 @@ def _plot_to_pdf(reference: Sequence[Atoms], initial: Optional[Sequence[Atoms]], write(symmetry_filename_F, sym_struc_F, format='extxyz', append=True) print("Number of times space group identification failed for generated dataset (var_prec = True): " - "{}/{} total".format(sg_fail, len(generated_atoms))) + "{}/{} total".format(sg_fail, len(generated_atoms_for_analysis))) print("Number of times space group identification failed for generated dataset (var_prec = False): " - "{}/{} total".format(sg_fail_F, len(generated_atoms))) + "{}/{} total".format(sg_fail_F, len(generated_atoms_for_analysis))) # Plot with PdfPages(plot_name) as pdf: diff --git a/omg/utils.py b/omg/utils.py index 347828e..f75c5a2 100644 --- a/omg/utils.py +++ b/omg/utils.py @@ -80,12 +80,24 @@ def xyz_saver(data: Union[OMGData, list[OMGData]], filename: Path) -> None: batch_size = len(d.n_atoms) for i in range(batch_size): lower, upper = d.ptr[i * 1], d.ptr[(i * 1) + 1] + species_slice = d.species[lower:upper].cpu().numpy() + MAX_Z = 118 # highest element ASE knows about + safe_species = species_slice.copy() + ghost_mask = species_slice <= 0 + ghost_mask |= species_slice > MAX_Z + safe_species[ghost_mask] = 0 # map ghosts to dummy element 'X' if d.pos_is_fractional[i]: - atoms.append(Atoms(numbers=d.species[lower:upper], scaled_positions=d.pos[lower:upper, :], - cell=d.cell[i, :, :], pbc=True)) + atom = Atoms(numbers=safe_species, scaled_positions=d.pos[lower:upper, :], + cell=d.cell[i, :, :], pbc=True) else: - atoms.append(Atoms(numbers=d.species[lower:upper], positions=d.pos[lower:upper, :], - cell=d.cell[i, :, :], pbc=True)) + atom = Atoms(numbers=safe_species, positions=d.pos[lower:upper, :], + cell=d.cell[i, :, :], pbc=True) + if ghost_mask.any(): + atom.new_array("is_ghost", ghost_mask.astype(np.int8)) + ghost_numbers = np.full_like(safe_species, -1) + ghost_numbers[ghost_mask] = species_slice[ghost_mask] + atom.new_array("ghost_atomic_number", ghost_numbers) + atoms.append(atom) write(filename, atoms, append=True)