Rename species to atomic_numbers (#19)

we had species for historical reasons, now
cleanly renamed to atomic_numbers everywhere

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: sirmarcel

src/petjax/model.py

@@ -47,7 +47,7 @@ def __call__(
         R_ij,
         centers,
         neighbors,
-        species,
+        atomic_numbers,
         reverse,
         pair_mask,
         atom_mask,
@@ -65,7 +65,16 @@ def __call__(
             max_atomic_number=self.max_atomic_number,
             attention_temperature=self.attention_temperature,
             name="backbone",
-        )(R_ij, centers, neighbors, species, reverse, pair_mask, atom_mask, pair_cutoffs)
+        )(
+            R_ij,
+            centers,
+            neighbors,
+            atomic_numbers,
+            reverse,
+            pair_mask,
+            atom_mask,
+            pair_cutoffs,
+        )
 
         predictions = Energy(d_head=self.d_head, name="energy_head")(
             node, edge, cutoffs, pair_mask, atom_mask
@@ -87,7 +96,9 @@ def __call__(
             force_scale = self.param(
                 "force_scale", nn.initializers.ones, (self.max_atomic_number + 1,)
             )
-            out["forces"] = forces * force_scale[species][:, None] * atom_mask[:, None]
+            out["forces"] = (
+                forces * force_scale[atomic_numbers][:, None] * atom_mask[:, None]
+            )
         if self.direct_stress:
             stress = DirectStress(d_head=self.d_head, name="stress_head")(
                 node, edge, cutoffs, pair_mask, atom_mask
@@ -121,7 +132,7 @@ def __call__(
         R_ij,
         centers,
         neighbors,
-        species,
+        atomic_numbers,
         reverse,
         pair_mask,
         atom_mask,
@@ -130,7 +141,7 @@ def __call__(
         d_pet = self.d_pet
         d_node = self.d_node
         P = R_ij.shape[0]
-        N = species.shape[0]
+        N = atomic_numbers.shape[0]
         n = P // N
 
         r_ij = safe_norm(R_ij, axis=-1)
@@ -149,11 +160,11 @@ def __call__(
         # Initial edge features; species embeddings index by atomic number Z
         # (table size max_atomic_number + 1, row 0 unused).
         edge_embed = nn.Embed(self.max_atomic_number + 1, d_pet, name="edge_embedder")
-        messages = edge_embed(species)[neighbors] * pair_mask[..., None]
+        messages = edge_embed(atomic_numbers)[neighbors] * pair_mask[..., None]
 
         # Node embedding (feedforward: persists across layers)
         node_embed = nn.Embed(self.max_atomic_number + 1, d_node, name="node_embedders_0")
-        node = node_embed(species)[:, None, :] * atom_mask[:, None, None]
+        node = node_embed(atomic_numbers)[:, None, :] * atom_mask[:, None, None]
 
         for layer_idx in range(self.num_gnn_layers):
             # Geometric features
@@ -179,7 +190,9 @@ def __call__(
                     d_pet,
                     name=f"gnn_layers_{layer_idx}_neighbor_embed",
                 )
-                neighbor_feats = neighbor_embed(species)[neighbors] * pair_mask[..., None]
+                neighbor_feats = (
+                    neighbor_embed(atomic_numbers)[neighbors] * pair_mask[..., None]
+                )
                 tokens_flat = masked(
                     MLP(
                         (d_pet, d_pet),

src/petjax/select.py

@@ -56,7 +56,7 @@ def truncate(
         structure["others"],
         structure["reverse"],
         structure["pair_mask"],
-        structure["species"],
+        structure["atomic_numbers"],
         structure["atom_mask"],
         structure["k_sel_sizer"].shape[-1],
         num_neighbors_adaptive,
@@ -73,7 +73,7 @@ def truncate_edges(
     others,
     reverse,
     pair_mask,
-    species,
+    atomic_numbers,
     atom_mask,
     k_sel,
     num_neighbors_adaptive,
@@ -97,7 +97,7 @@ def truncate_edges(
     couples pairs only through their center/other atoms. Requires ``centers``
     non-decreasing; ``k_sel`` must be a static int under jit.
     """
-    N = species.shape[0]
+    N = atomic_numbers.shape[0]
 
     pair_cutoffs, selected = _select_edges(
         R_ij,
@@ -121,7 +121,7 @@ def truncate_edges(
         "R_ij": R_ij[sel_to_pair],
         "centers": centers[sel_to_pair],
         "neighbors": others[sel_to_pair],
-        "species": species,
+        "atomic_numbers": atomic_numbers,
         "reverse": slot[reverse[sel_to_pair]],
         "pair_mask": pair_mask_sel,
         "atom_mask": atom_mask,
@@ -130,7 +130,7 @@ def truncate_edges(
     return truncated, overflow
 
 
-def pack_edges(R_ij, centers, others, reverse, pair_mask, species, atom_mask, k):
+def pack_edges(R_ij, centers, others, reverse, pair_mask, atomic_numbers, atom_mask, k):
     """Fixed-width pack on a flat NL with precomputed displacements — no
     selection: every unmasked pair goes into the rectangular ``[N * k]``
     layout. Returns the truncated dict keyed like ``truncate_edges``'s, with
@@ -141,15 +141,15 @@ def pack_edges(R_ij, centers, others, reverse, pair_mask, species, atom_mask, k)
     Same layout requirements as ``truncate_edges``: ``centers``
     non-decreasing, ``k`` a static int under jit.
     """
-    N = species.shape[0]
+    N = atomic_numbers.shape[0]
     slot, sel_to_pair, pair_mask_sel, overflow = _pack_selected_to_flat(
         pair_mask, centers, N, k
     )
     truncated = {
         "R_ij": R_ij[sel_to_pair],
         "centers": centers[sel_to_pair],
         "neighbors": others[sel_to_pair],
-        "species": species,
+        "atomic_numbers": atomic_numbers,
         "reverse": slot[reverse[sel_to_pair]],
         "pair_mask": pair_mask_sel,
         "atom_mask": atom_mask,

src/petjax/structure.py

@@ -103,8 +103,8 @@ def to_structure(
     pair_mask[:n_pair_raw] = True
     reverse[:n_pair_raw] = reverse_sparse.astype(int_dtype)
 
-    species = np.zeros(N_padded, dtype=int_dtype)
-    species[:n_atoms] = atoms.get_atomic_numbers()  # embeddings index by atomic number
+    atomic_numbers = np.zeros(N_padded, dtype=int_dtype)
+    atomic_numbers[:n_atoms] = atoms.get_atomic_numbers()
 
     atom_mask = np.zeros(N_padded, dtype=bool)
     atom_mask[:n_atoms] = True
@@ -119,7 +119,7 @@ def to_structure(
     return {
         "positions": positions,
         "cell": cell,
-        "species": species,
+        "atomic_numbers": atomic_numbers,
         "atom_mask": atom_mask,
         "centers": centers,
         "others": others,

tests/test_calculator.py

@@ -252,7 +252,7 @@ def test_pair_cutoffs_none_uses_static_cutoff(model_data, mini_xyz):
         structure["others"],
         structure["reverse"],
         structure["pair_mask"],
-        structure["species"],
+        structure["atomic_numbers"],
         structure["atom_mask"],
         k_sel,
         config["num_neighbors_adaptive"],

tests/test_select.py

@@ -32,7 +32,7 @@ def _pack(structure, k):
         structure["others"],
         structure["reverse"],
         structure["pair_mask"],
-        structure["species"],
+        structure["atomic_numbers"],
         structure["atom_mask"],
         k,
     )