Handle node/graph features/targets in from_ase_atoms_list

Author: vsimkus

orb_models/forcefield/forcefield_adapter.py

@@ -209,6 +209,7 @@ def from_ase_atoms_list(
         if len(atoms) == 1:
             return self.from_ase_atoms(
                 atoms[0],
+                max_num_neighbors=max_num_neighbors,
                 edge_method=edge_method,
                 wrap=wrap,
                 device=device,
@@ -311,6 +312,7 @@ def from_ase_atoms_list(
             "positions": positions,
             "atomic_numbers": atomic_numbers,
             "atomic_numbers_embedding": atomic_numbers_embedding,
+            "atom_identity": torch.cat([torch.arange(n, dtype=torch.long) for n in n_atoms], dim=0),
         }
         edge_feats = {
             "vectors": edge_vectors,
@@ -320,6 +322,17 @@ def from_ase_atoms_list(
             "cell": cells,
             "pbc": pbcs,
         }
+
+        # Merge extra features from atoms.info
+        node_feats.update(_batch_info_tensors(atoms, "node_features"))
+        edge_feats.update(_batch_info_tensors(atoms, "edge_features"))
+        graph_feats.update(_batch_info_tensors(atoms, "graph_features", system_level=True))
+
+        # Collect targets from atoms.info
+        node_targets = _batch_info_tensors(atoms, "node_targets")
+        edge_targets = _batch_info_tensors(atoms, "edge_targets")
+        system_targets = _batch_info_tensors(atoms, "graph_targets", system_level=True)
+
         # Collect charge and spin: all-or-nothing semantics
         charge_spin_list = [_get_charge_and_spin(a) for a in atoms]
         has_charge_spin = [bool(cs) for cs in charge_spin_list]
@@ -341,9 +354,9 @@ def from_ase_atoms_list(
             node_features=node_feats,
             edge_features=edge_feats,
             system_features=graph_feats,
-            node_targets={},
-            edge_targets={},
-            system_targets={},
+            node_targets=node_targets,
+            edge_targets=edge_targets,
+            system_targets=system_targets,
             system_id=None,
             fix_atoms=fix_atoms,
             tags=tags,
@@ -482,6 +495,39 @@ def is_compatible_with(self, other: AbstractAtomsAdapter):
         return True
 
 
+def _batch_info_tensors(
+    atoms_list: list[ase.Atoms],
+    info_key: str,
+    system_level: bool = False,
+) -> dict:
+    """Collect tensor dicts from atoms.info[info_key] across a list and batch them.
+
+    For system-level tensors, matches from_ase_atoms behavior: unsqueeze non-scalars
+    (numel > 1) before concatenating, so cat produces [N, ...] for non-scalars
+    and [N] for scalars.
+    """
+    dicts = [a.info.get(info_key, {}) for a in atoms_list]
+    all_keys = [set(d.keys()) for d in dicts]
+    keys = set().union(*all_keys)
+    if not keys:
+        return {}
+    if any(ks != keys for ks in all_keys):
+        raise ValueError(
+            f"All atoms must have the same set of keys in info['{info_key}']. "
+            f"Got: {[sorted(ks) for ks in all_keys]}"
+        )
+    out: dict = {}
+    for k in keys:
+        values = [d[k] for d in dicts]
+        if any(v is None for v in values):
+            out[k] = None
+        elif system_level:
+            out[k] = torch.cat([v.unsqueeze(0) if v.numel() > 1 else v for v in values], dim=0)
+        else:
+            out[k] = torch.cat(values, dim=0)
+    return out
+
+
 def _get_charge_and_spin(atoms: ase.Atoms | ts.SimState) -> dict[str, torch.Tensor]:
     out = {}
     if isinstance(atoms, ase.Atoms) and ("charge" in atoms.info or "spin" in atoms.info):

tests/common/atoms/test_graph_batch.py

@@ -131,6 +131,26 @@ def test_random_batching():
     assert batched.equals(AtomGraphs.batch(graphs))
 
 
+def test_batch_none_values_collapse():
+    """Batching graphs where a target value is None in one graph produces None for that key."""
+    g1 = graph()
+    g2 = graph()
+    g2.node_targets["node_target"] = None
+
+    batched = AtomGraphs.batch([g1, g2])
+    assert batched.node_targets["node_target"] is None
+
+
+def test_batch_mismatched_keys_raises():
+    """Batching graphs with different feature/target keys raises."""
+    g1 = graph()
+    g2 = graph()
+    g2.system_targets["extra"] = torch.tensor([1.0])
+
+    with pytest.raises(ValueError, match="same nested structure"):
+        AtomGraphs.batch([g1, g2])
+
+
 def test_refeaturization_pos_substitution(dataset_and_loader):
     dataset = dataset_and_loader[0]
     datapoint = dataset[0]

tests/forcefield/test_forcefield_adapter.py

@@ -196,15 +196,23 @@ def test_forcefield_adapter_requires_both_spin_and_charge():
 
 def test_from_ase_atoms_list_parallel_equivalence():
     """Test that from_ase_atoms_list produces equivalent results to sequential processing."""
-    atoms_list = [
-        Atoms(
+    n_atoms = 3
+    atoms_list = []
+    for i in range(4):
+        a = Atoms(
             "H2O",
             positions=np.array([[0, 0, 0], [0, 1, 0], [1, 0, 0]]) + i * 0.1,
             pbc=True,
             cell=np.diag([5, 5, 5]),
         )
-        for i in range(4)
-    ]
+        a.info["graph_features"] = {"bandgap": torch.tensor([float(i)])}
+        a.info["graph_targets"] = {
+            "energy": torch.tensor([float(i) * 0.5]),
+            "stress": torch.randn(6),
+        }
+        a.info["node_features"] = {"mulliken_charges": torch.randn(n_atoms, 1)}
+        a.info["node_targets"] = {"forces": torch.randn(n_atoms, 3)}
+        atoms_list.append(a)
     adapter = ForcefieldAtomsAdapter(radius=6.0, max_num_neighbors=20)
 
     batch_result = adapter.from_ase_atoms_list(atoms_list)
@@ -226,6 +234,88 @@ def test_from_ase_atoms_list_parallel_equivalence():
         sd = sequential_batch.edge_features["vectors"][start:end].cpu().norm(dim=1).sort()[0]
         assert torch.allclose(bd, sd, atol=1e-4)
 
+    # Verify features/targets keys match between batched and sequential
+    assert batch_result.node_features.keys() == sequential_batch.node_features.keys()
+    assert batch_result.system_features.keys() == sequential_batch.system_features.keys()
+    assert batch_result.system_targets.keys() == sequential_batch.system_targets.keys()
+    assert batch_result.node_targets.keys() == sequential_batch.node_targets.keys()
+
+    # Verify feature/target values
+    torch.testing.assert_close(
+        batch_result.system_features["bandgap"].cpu(),
+        sequential_batch.system_features["bandgap"].cpu(),
+    )
+    torch.testing.assert_close(
+        batch_result.system_targets["energy"].cpu(),
+        sequential_batch.system_targets["energy"].cpu(),
+    )
+    torch.testing.assert_close(
+        batch_result.system_targets["stress"].cpu(),
+        sequential_batch.system_targets["stress"].cpu(),
+    )
+    torch.testing.assert_close(
+        batch_result.node_features["mulliken_charges"].cpu(),
+        sequential_batch.node_features["mulliken_charges"].cpu(),
+    )
+    torch.testing.assert_close(
+        batch_result.node_targets["forces"].cpu(),
+        sequential_batch.node_targets["forces"].cpu(),
+    )
+
+
+def test_from_ase_atoms_list_inconsistent_info_raises():
+    """Test that from_ase_atoms_list raises when atoms have inconsistent info keys."""
+    adapter = ForcefieldAtomsAdapter(radius=6.0, max_num_neighbors=20)
+
+    a0 = Atoms(
+        "H2O",
+        positions=np.array([[0, 0, 0], [0, 1, 0], [1, 0, 0]]),
+        pbc=True,
+        cell=np.diag([5, 5, 5]),
+    )
+    a0.info["node_targets"] = {"forces": torch.randn(3, 3)}
+    a0.info["graph_targets"] = {"energy": torch.tensor([1.0])}
+
+    a1 = Atoms(
+        "H2O",
+        positions=np.array([[0, 0, 0], [0, 1, 0], [1, 0, 0]]) + 0.1,
+        pbc=True,
+        cell=np.diag([5, 5, 5]),
+    )
+
+    with pytest.raises(ValueError, match="same set of keys"):
+        adapter.from_ase_atoms_list([a0, a1])
+
+
+def test_from_ase_atoms_list_none_values_collapse_to_none():
+    """Test that if any atom has None for a key, the batched result is None for that key."""
+    adapter = ForcefieldAtomsAdapter(radius=6.0, max_num_neighbors=20)
+    atoms_list = []
+    for i in range(3):
+        a = Atoms(
+            "H2O",
+            positions=np.array([[0, 0, 0], [0, 1, 0], [1, 0, 0]]) + i * 0.1,
+            pbc=True,
+            cell=np.diag([5, 5, 5]),
+        )
+        a.info["node_targets"] = {
+            "forces": None if i == 0 else torch.randn(3, 3),
+        }
+        a.info["graph_targets"] = {
+            "energy": torch.tensor([float(i)]),
+            "stress": None if i == 1 else torch.randn(6),
+        }
+        atoms_list.append(a)
+
+    batch_result = adapter.from_ase_atoms_list(atoms_list)
+
+    assert batch_result.node_targets["forces"] is None
+    assert batch_result.system_targets["stress"] is None
+    torch.testing.assert_close(
+        batch_result.system_targets["energy"].cpu(),
+        torch.tensor([0.0, 1.0, 2.0]),
+    )
+
 
 def test_from_ase_atoms_list_nonperiodic():
     """Test from_ase_atoms_list with non-periodic systems."""