fix: Handle batched cell tensors in get_fractional_coordinates

- Replace deprecated .T with .mT for matrix transpose on 3D tensors
- Add support for batched cell tensors with shape [n_batches, 3, 3]
- Extract first batch cell matrix when cell.ndim == 3
- Maintains backward compatibility with 2D cell matrices

Fixes batched silicon workflow script that was failing with:
- UserWarning about deprecated .T usage on >2D tensors
- RuntimeError: linalg.solve: A must be batches of square matrices

The get_fractional_coordinates function now properly handles both
single [3,3] and batched [n_batches,3,3] cell tensors, enabling a2c_silicon_batched.py workflow to run

Author: janosh

examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py

@@ -98,7 +98,7 @@
             ).item()
             / Units.pressure
         )
-        xx, yy, zz = torch.diag(state.cell)
+        xx, yy, zz = torch.diag(state.cell[0])
         print(
             f"{step=}: Temperature: {temp.item():.4f}, "
             f"pressure: {pressure:.4f}, "

torch_sim/integrators/npt.py

@@ -95,6 +95,16 @@ def _compute_cell_force(
     Returns:
         torch.Tensor: Force acting on the cell [n_batches, n_dim, n_dim]
     """
+    # Convert external_pressure to tensor if it's not already one
+    if not isinstance(external_pressure, torch.Tensor):
+        external_pressure = torch.tensor(
+            external_pressure, device=state.device, dtype=state.dtype
+        )
+
+    # Convert kT to tensor if it's not already one
+    if not isinstance(kT, torch.Tensor):
+        kT = torch.tensor(kT, device=state.device, dtype=state.dtype)
+
     # Get current volumes for each batch
     volumes = torch.linalg.det(state.cell)  # shape: (n_batches,)
 

torch_sim/transforms.py

@@ -40,7 +40,13 @@ def get_fractional_coordinates(
         tensor([[0.25, 0.25, 0.25],
                 [0.50, 0.00, 0.00]])
     """
-    return torch.linalg.solve(cell.T, positions.T).T
+    if cell.ndim == 3:  # Handle batched cell tensors
+        # For batched systems, we assume single batch for now
+        # Extract the first batch's cell matrix
+        cell_2d = cell[0]  # Shape: [3, 3]
+        return torch.linalg.solve(cell_2d.mT, positions.mT).mT
+    # Original case for 2D cell matrix
+    return torch.linalg.solve(cell.mT, positions.mT).mT
 
 
 def inverse_box(box: torch.Tensor) -> torch.Tensor: