add batch dim to a2c functions

Author: janosh

examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py

@@ -50,7 +50,7 @@
     compute_forces=True,
     compute_stress=False,
     dtype=dtype,
-    enable_cueq=True,
+    enable_cueq=torch.cuda.is_available(),
 )
 state = ts.io.atoms_to_state(si_dc, device=device, dtype=dtype)
 

torch_sim/models/mace.py

@@ -304,7 +304,24 @@ def forward(  # noqa: C901
         # TODO (AG): Currently doesn't work for batched neighbor lists
         for b in range(self.n_systems):
             batch_mask = state.batch == b
-            # Calculate neighbor list for this system
+            # Extract cell for this batch and transpose to row vector format
+            cell_for_nl = state.row_vector_cell[b]
+
+            # Ensure cell has correct shape [3, 3] for neighbor list
+            if cell_for_nl.ndim == 3:
+                cell_for_nl = cell_for_nl.squeeze(0)
+
+            # Safety checks to prevent tensor confusion
+            if state.positions[batch_mask].shape[1] != 3:
+                raise ValueError(
+                    f"positions should have shape [n_atoms, 3], got "
+                    f"{state.positions[batch_mask].shape}"
+                )
+            if cell_for_nl.shape != (3, 3):
+                raise ValueError(
+                    f"cell should have shape [3, 3], got {cell_for_nl.shape}"
+                )
+
             edge_idx, shifts_idx = self.neighbor_list_fn(
                 positions=state.positions[batch_mask],
                 cell=state.row_vector_cell[b],

torch_sim/neighbors.py

@@ -537,6 +537,12 @@ def vesin_nl_ts(
     References:
           https://github.com/Luthaf/vesin
     """
+    # Defensive checks to catch parameter confusion
+    if positions.ndim != 2 or positions.shape[1] != 3:
+        raise ValueError(f"positions must have shape [n_atoms, 3], got {positions.shape}")
+    if cell.ndim != 2 or cell.shape[0] != 3 or cell.shape[1] != 3:
+        raise ValueError(f"cell must have shape [3, 3], got {cell.shape}")
+
     device = positions.device
     dtype = positions.dtype
 

torch_sim/workflows/a2c.py

@@ -280,11 +280,17 @@ def random_packed_structure(
         diameter = get_diameter(composition)
         print(f"Using random pack diameter of {diameter}")
 
+    # Ensure cell has batch dimension [1, 3, 3] if it doesn't already
+    if cell.ndim == 2:
+        cell = cell.unsqueeze(0)  # Add batch dimension
+
     # Perform overlap minimization if diameter is specified
     if diameter is not None:
         print("Reduce atom overlap using the soft_sphere potential")
         # Convert fractional to cartesian coordinates
-        positions_cart = torch.matmul(positions, cell)
+        positions_cart = torch.matmul(
+            positions, cell.squeeze(0)
+        )  # Use first (and only) batch
 
         # Initialize soft sphere potential calculator
         model = SoftSphereModel(
@@ -296,23 +302,30 @@ def random_packed_structure(
         )
 
         # Dummy atomic numbers
-        atomic_numbers = torch.ones_like(positions_cart, device=device, dtype=torch.int)
+        atomic_numbers = torch.ones(N_atoms, device=device, dtype=torch.int)
+
+        # Create batch tensor for single system
+        batch = torch.zeros(N_atoms, device=device, dtype=torch.long)
 
         # Set up FIRE optimizer with unit masses
         state = ts.SimState(
             positions=positions_cart,
             masses=torch.ones(N_atoms, device=device, dtype=dtype),
             atomic_numbers=atomic_numbers,
-            cell=cell,
+            cell=cell,  # Keep batch dimension
             pbc=True,
+            batch=batch,
         )
         fire_init, fire_update = fire(model=model)
         state = fire_init(state)
         print(f"Initial energy: {state.energy.item():.4f}")
         # Run FIRE optimization until convergence or max iterations
         for _step in range(max_iter):
             # Check if minimum distance criterion is met (95% of target diameter)
-            if min_distance(state.positions, cell, distance_tolerance) > diameter * 0.95:
+            if (
+                min_distance(state.positions, cell.squeeze(0), distance_tolerance)
+                > diameter * 0.95
+            ):
                 break
 
             if log is not None:
@@ -321,6 +334,20 @@ def random_packed_structure(
             state = fire_update(state)
 
         print(f"Final energy: {state.energy.item():.4f}")
+    else:
+        # If no optimization, still create a proper state with batch dimensions
+        positions_cart = torch.matmul(positions, cell.squeeze(0))
+        atomic_numbers = torch.ones(N_atoms, device=device, dtype=torch.int)
+        batch = torch.zeros(N_atoms, device=device, dtype=torch.long)
+
+        state = ts.SimState(
+            positions=positions_cart,
+            masses=torch.ones(N_atoms, device=device, dtype=dtype),
+            atomic_numbers=atomic_numbers,
+            cell=cell,  # Keep batch dimension
+            pbc=True,
+            batch=batch,
+        )
 
     if log is not None:
         return state, log
@@ -408,11 +435,17 @@ def random_packed_structure_multi(
         diameter_matrix = get_diameter_matrix(composition, device=device, dtype=dtype)
         print(f"Using random pack diameter matrix:\n{diameter_matrix.cpu().numpy()}")
 
+    # Ensure cell has batch dimension [1, 3, 3] if it doesn't already
+    if cell.ndim == 2:
+        cell = cell.unsqueeze(0)  # Add batch dimension
+
     # Perform overlap minimization if diameter matrix is specified
     if diameter_matrix is not None:
         print("Reduce atom overlap using the soft_sphere potential")
         # Convert fractional to cartesian coordinates
-        positions_cart = torch.matmul(positions, cell)
+        positions_cart = torch.matmul(
+            positions, cell.squeeze(0)
+        )  # Use first (and only) batch
 
         # Initialize multi-species soft sphere potential calculator
         model = SoftSphereMultiModel(
@@ -425,14 +458,18 @@ def random_packed_structure_multi(
         )
 
         # Dummy atomic numbers
-        atomic_numbers = torch.ones_like(positions_cart, device=device, dtype=torch.int)
+        atomic_numbers = torch.ones(N_atoms, device=device, dtype=torch.int)
+
+        # Create batch tensor for single system
+        batch = torch.zeros(N_atoms, device=device, dtype=torch.long)
 
         state_dict = ts.SimState(
             positions=positions_cart,
             masses=torch.ones(N_atoms, device=device, dtype=dtype),
             atomic_numbers=atomic_numbers,
-            cell=cell,
+            cell=cell,  # Keep batch dimension
             pbc=True,
+            batch=batch,
         )
         # Set up FIRE optimizer with unit masses for all atoms
         fire_init, fire_update = fire(model=model)
@@ -441,11 +478,25 @@ def random_packed_structure_multi(
         # Run FIRE optimization until convergence or max iterations
         for _step in range(max_iter):
             # Check if minimum distance criterion is met (95% of smallest target diameter)
-            min_dist = min_distance(state.positions, cell, distance_tolerance)
+            min_dist = min_distance(state.positions, cell.squeeze(0), distance_tolerance)
             if min_dist > diameter_matrix.min() * 0.95:
                 break
             state = fire_update(state)
         print(f"Final energy: {state.energy.item():.4f}")
+    else:
+        # If no optimization, still create a proper state with batch dimensions
+        positions_cart = torch.matmul(positions, cell.squeeze(0))
+        atomic_numbers = torch.ones(N_atoms, device=device, dtype=torch.int)
+        batch = torch.zeros(N_atoms, device=device, dtype=torch.long)
+
+        state = ts.SimState(
+            positions=positions_cart,
+            masses=torch.ones(N_atoms, device=device, dtype=dtype),
+            atomic_numbers=atomic_numbers,
+            cell=cell,  # Keep batch dimension
+            pbc=True,
+            batch=batch,
+        )
 
     return state
 
@@ -472,8 +523,8 @@ def valid_subcell(
     Args:
         positions: Atomic positions tensor of shape [n_atoms, 3], where each row contains
             the (x,y,z) coordinates of an atom.
-        cell: Unit cell tensor of shape [3, 3] containing the three lattice vectors that
-            define the periodic boundary conditions.
+        cell: Unit cell tensor of shape [3, 3] or [1, 3, 3] containing the three lattice
+            vectors that define the periodic boundary conditions.
         initial_energy: Total energy of the structure before relaxation, in eV.
         final_energy: Total energy of the structure after relaxation, in eV.
         e_tol: Energy tolerance for comparing initial and final energies, in eV.
@@ -510,8 +561,9 @@ def valid_subcell(
         return False
 
     # Check minimum interatomic distances to detect atomic fusion
-    # Uses periodic boundary conditions via min_distance function
-    min_dist = min_distance(positions, cell, distance_tolerance)
+    # Handle both batched and unbatched cell tensors
+    cell_for_min_dist = cell.squeeze(0) if cell.ndim == 3 else cell
+    min_dist = min_distance(positions, cell_for_min_dist, distance_tolerance)
     if min_dist < fusion_distance:
         print("Bad structure! Fusion found.")
         return False
@@ -645,15 +697,18 @@ def subcells_to_structures(
         candidates: List of (ids, lower_bound, upper_bound)
         tuples from get_subcells_to_crystallize
         fractional_positions: Fractional coordinates of atoms
-        cell: Unit cell tensor
+        cell: Unit cell tensor of shape [3, 3] or [1, 3, 3]
         species: List of atomic species symbols
 
     Returns:
         list[tuple[torch.Tensor, torch.Tensor, list[str]]]: Each tuple contains:
             - fractional_positions: Fractional coordinates of atoms
-            - cell: Unit cell tensor
+            - cell: Unit cell tensor with proper batch dimensions
             - species: atomic species symbols
     """
+    # Handle both batched and unbatched cell tensors
+    cell_2d = cell.squeeze(0) if cell.ndim == 3 else cell
+
     list_subcells = []
     for ids, lower_bound, upper_bound in candidates:
         # Get positions of atoms in this subcell
@@ -666,7 +721,11 @@ def subcells_to_structures(
         new_frac_pos = new_frac_pos / (upper_bound - lower_bound)
 
         # Calculate new cell parameters
-        new_cell = cell * (upper_bound - lower_bound).unsqueeze(0)
+        new_cell = cell_2d * (upper_bound - lower_bound).unsqueeze(0)
+
+        # Add batch dimension to maintain consistency
+        if cell.ndim == 3:  # Original cell had batch dimension
+            new_cell = new_cell.unsqueeze(0)
 
         # Get species for these atoms and convert tensor indices to list/numpy array
         # before indexing species list
@@ -721,12 +780,18 @@ def get_unit_cell_relaxed_structure(
         tuple containing:
             - UnitCellFIREState: Final state containing relaxed positions, cell and more
             - dict: Logger with energy and stress trajectories
-            - float: Final energy in eV
-            - float: Final pressure in eV/ų
     """
     # Get device and dtype from model
     device, dtype = model.device, model.dtype
 
+    # Ensure state has proper batch dimensions
+    if state.cell.ndim == 2:
+        state.cell = state.cell.unsqueeze(0)  # Add batch dimension
+
+    # Ensure batch tensor exists and has correct shape
+    if state.batch is None:
+        state.batch = torch.zeros(len(state.positions), device=device, dtype=torch.long)
+
     logger = {
         "energy": torch.zeros((max_iter, state.n_batches), device=device, dtype=dtype),
         "stress": torch.zeros(
@@ -769,7 +834,7 @@ def step_fn(
         f"Final energy: {[f'{e:.4f}' for e in final_energy]} eV, "
         f"Final pressure: {[f'{p:.4f}' for p in final_pressure]} eV/A^3"
     )
-    return state, logger, final_energy, final_pressure
+    return state, logger
 
 
 def get_relaxed_structure(
@@ -791,12 +856,18 @@ def get_relaxed_structure(
         tuple containing:
             - FIREState: Final state containing relaxed positions and other quantities
             - dict: Logger with energy trajectory
-            - float: Final energy in eV
-            - float: Final pressure in eV/ų
     """
     # Get device and dtype from model
     device, dtype = model.device, model.dtype
 
+    # Ensure state has proper batch dimensions
+    if state.cell.ndim == 2:
+        state.cell = state.cell.unsqueeze(0)  # Add batch dimension
+
+    # Ensure batch tensor exists and has correct shape
+    if state.batch is None:
+        state.batch = torch.zeros(len(state.positions), device=device, dtype=torch.long)
+
     logger = {"energy": torch.zeros((max_iter, 1), device=device, dtype=dtype)}
 
     results = model(state)
@@ -816,9 +887,7 @@ def step_fn(idx: int, state: FireState, logger: dict) -> tuple[FireState, dict]:
 
     # Get final results
     model.compute_stress = True
-    final_results = model(
-        positions=state.positions, cell=state.cell, atomic_numbers=state.atomic_numbers
-    )
+    final_results = model(state)
 
     final_energy = final_results["energy"].item()
     final_stress = final_results["stress"]
@@ -827,4 +896,4 @@ def step_fn(idx: int, state: FireState, logger: dict) -> tuple[FireState, dict]:
         f"Final energy: {final_energy:.4f} eV, "
         f"Final pressure: {final_pressure:.4f} eV/A^3"
     )
-    return state, logger, final_energy, final_pressure
+    return state, logger