Ghost Loading + Validation/Viz v3

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]
 

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

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

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

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, )

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()

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()