fix: perf damage in data processing by jnp

Author: bhcao

src/mlip/data/helpers/edge_vectors.py

@@ -25,6 +25,7 @@ def get_edge_vectors(
     shifts: np.ndarray,  # [n_edges, 3]
     cell: Optional[np.ndarray],  # [n_graph, 3, 3]
     n_edge: np.ndarray,  # [n_graph]
+    use_np: bool = False,
 ) -> Tuple[np.ndarray, np.ndarray]:
     """Compute positions of sender and receiver nodes of each edge.
 
@@ -55,6 +56,8 @@ def get_edge_vectors(
                 Output `S` of `ase.neighborlist.primitive_neighbor_list`.
         cell: The cell of each graph. Array of shape ``[n_graph, 3, 3]``.
         n_edge: The number of edges of each graph. Array of shape ``[n_graph]``.
+        use_np: Whether to use `numpy` or `jax.numpy` for the computation. Default
+                is `False`, which means `jax.numpy` is used.
 
     Returns:
         The positions of the sender and receiver nodes of each edge.
@@ -63,15 +66,21 @@ def get_edge_vectors(
     vectors_receivers = positions[receivers]  # [n_edges, 3]
 
     if cell is not None:
-        num_edges = receivers.shape[0]
-        shifts = jnp.einsum(
+        if use_np:
+            np_ = np
+            kwargs_num_edges = {}
+        else:
+            np_ = jnp
+            kwargs_num_edges = {"total_repeat_length": receivers.shape[0]}
+
+        shifts = np_.einsum(
             "ei,eij->ej",
             shifts,  # [n_edges, 3]
-            jnp.repeat(
+            np_.repeat(
                 cell,  # [n_graph, 3, 3]
                 n_edge,  # [n_graph]
                 axis=0,
-                total_repeat_length=num_edges,
+                **kwargs_num_edges,
             ),  # [n_edges, 3, 3]
         )  # [n_edges, 3]
         vectors_senders -= shifts  # minus sign to match results with ASE
@@ -86,6 +95,7 @@ def get_edge_relative_vectors(
     shifts: np.ndarray,  # [n_edges, 3]
     cell: Optional[np.ndarray],  # [n_graph, 3, 3]
     n_edge: np.ndarray,  # [n_graph]
+    use_np: bool = False,
 ) -> np.ndarray:
     """Compute the relative edge vectors from senders to receivers.
 
@@ -108,6 +118,8 @@ def get_edge_relative_vectors(
                 functionality, and labelled `S` by ASE.
         cell: The unit cells of each graph, an array of shape `[n_graph, 3, 3]`.
         n_edge: The number of edges for each graph, an array of shape `[n_graph]`.
+        use_np: Whether to use `numpy` or `jax.numpy` for the computation. Default
+                is `False`, which means `jax.numpy` is used.
 
     Returns:
         The relative edge vectors, labelled `D` by ASE.
@@ -119,5 +131,6 @@ def get_edge_relative_vectors(
         shifts=shifts,
         cell=cell,
         n_edge=n_edge,
+        use_np=use_np,
     )
     return vectors_receivers - vectors_senders

src/mlip/data/helpers/neighbor_analysis.py

@@ -14,7 +14,6 @@
 
 import logging
 
-import jax
 import jraph
 import numpy as np
 from tqdm_loggable.auto import tqdm
@@ -59,14 +58,14 @@ def compute_avg_min_neighbor_distance(graphs: list[jraph.GraphsTuple]) -> float:
     log_interval = max(1, len(graphs) // LOG_PROPORTION)
 
     for i, graph in enumerate(graphs):
-        jit_get_edge_relative_vectors = jax.jit(get_edge_relative_vectors)
-        vectors = jit_get_edge_relative_vectors(
+        vectors = get_edge_relative_vectors(
             graph.nodes.positions,
             graph.senders,
             graph.receivers,
             graph.edges.shifts,
             graph.globals.cell,
             graph.n_edge,
+            use_np=True,
         )
         length = np.linalg.norm(vectors, axis=-1)
         min_neighbor_distances.append(length.min())