Remove unbatched code

.pre-commit-config.yaml

@@ -7,7 +7,7 @@ default_install_hook_types: [pre-commit, commit-msg]
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.8
+    rev: v0.11.12
     hooks:
       - id: ruff
         args: [--fix]
@@ -32,7 +32,7 @@ repos:
         stages: [pre-commit, commit-msg]
 
   - repo: https://github.com/igorshubovych/markdownlint-cli
-    rev: v0.44.0
+    rev: v0.45.0
     hooks:
       - id: markdownlint
         # MD013: line length

README.md

@@ -58,7 +58,7 @@ final_state = ts.integrate(
     n_steps=50,
     timestep=0.002,
     temperature=1000,
-    integrator=ts.nvt_langevin,
+    integrator=ts.integrators.nvt_langevin,
     trajectory_reporter=dict(filenames=trajectory_files, state_frequency=10),
 )
 final_atoms_list = final_state.to_atoms()

examples/readme.md

@@ -3,7 +3,7 @@
 Examples are provided in two foms:
 
 * Tutorials are intended to provide pedagogical walkthroughs of TorchSim's core functionality
-* Scripts are a holdover from early torch-sim development, they are not currently recommended as a learning resource. See issue [here](https://github.com/Radical-AI/torch-sim/issues/109).
+* Scripts are a holdover from early torch-sim development, they are not currently recommended as a learning resource. See issue [issue 109](https://github.com/Radical-AI/torch-sim/issues/109).
 
 ## Tutorial Formatting
 

examples/scripts/1_Introduction/1.1_Lennard_Jones.py

@@ -11,7 +11,6 @@
 import torch
 
 from torch_sim.models.lennard_jones import LennardJonesModel
-from torch_sim.unbatched.models.lennard_jones import UnbatchedLennardJonesModel
 
 
 # Set up the device and data type
@@ -61,41 +60,7 @@
 #  - sigma: distance at which potential is zero (3.405 Å for Ar)
 #  - epsilon: depth of potential well (0.0104 eV for Ar)
 #  - cutoff: distance beyond which interactions are ignored (typically 2.5*sigma)
-model = UnbatchedLennardJonesModel(
-    use_neighbor_list=True,
-    cutoff=2.5 * 3.405,
-    sigma=3.405,
-    epsilon=0.0104,
-    device=device,
-    dtype=dtype,
-    compute_forces=True,
-    compute_stress=True,
-    per_atom_energies=True,
-    per_atom_stresses=True,
-)
-
-# Print system information
-print(f"Positions: {positions.shape}")
-print(f"Cell: {cell.shape}")
-
-state = dict(
-    positions=positions,
-    cell=cell,
-    atomic_numbers=atomic_numbers,
-    pbc=True,
-)
-# Run the simulation and get results
-results = model(state)
-
-# Print the results
-print(f"Energy: {float(results['energy']):.4f}")
-print(f"Forces: {results['forces']}")
-print(f"Stress: {results['stress']}")
-print(f"Energies: {results['energies']}")
-print(f"Stresses: {results['stresses']}")
-
-# Batched model
-batched_model = LennardJonesModel(
+model = LennardJonesModel(
     use_neighbor_list=True,
     cutoff=2.5 * 3.405,
     sigma=3.405,
@@ -117,7 +82,7 @@
 )
 
 # Run the simulation and get results
-results = batched_model(state)
+results = model(state)
 
 # Print the results
 print(f"Energy: {results['energy']}")

examples/scripts/1_Introduction/1.2_MACE.py

@@ -1,18 +1,17 @@
-"""MACE simple single system example."""
+"""Minimal MACE batched example."""
 
 # /// script
 # dependencies = [
 #     "mace-torch>=0.3.12",
 # ]
 # ///
 
+import numpy as np
 import torch
 from ase.build import bulk
 from mace.calculators.foundations_models import mace_mp
 
-import torch_sim as ts
-from torch_sim.models.mace import MaceUrls
-from torch_sim.unbatched.models.mace import UnbatchedMaceModel
+from torch_sim.models.mace import MaceModel, MaceUrls
 
 
 # Set device and data type
@@ -27,30 +26,78 @@
     device=device,
 )
 
-# Option 2: Load from local file (comment out Option 1 to use this)
+# Option 2: Load the compiled model from the local file
 # MODEL_PATH = "../../../checkpoints/MACE/mace-mpa-0-medium.model"
 # loaded_model = torch.load(MODEL_PATH, map_location=device)
 
 # Create diamond cubic Silicon
 si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2))
+atoms_list = [si_dc, si_dc]
 
-# Initialize the unbatched MACE model
-model = UnbatchedMaceModel(
+batched_model = MaceModel(
+    # Pass the raw model
     model=loaded_model,
+    # Or load from compiled model
+    # model=compiled_model,
     device=device,
     compute_forces=True,
     compute_stress=True,
     dtype=dtype,
     enable_cueq=False,
 )
 
-# Convert ASE atoms to state
-state = ts.io.atoms_to_state(si_dc, device=device, dtype=dtype)
+# First we will create a concatenated positions array
+# This will have shape (16, 3) which is concatenated from two 8 atom systems
+positions_numpy = np.concatenate([atoms.positions for atoms in atoms_list])
 
-# Run inference
-results = model(state)
+# stack cell vectors into a (2, 3, 3) array where the first index is batch dimension
+cell_numpy = np.stack([atoms.cell.array for atoms in atoms_list])
 
-# Print results
-print(f"Energy: {float(results['energy']):.4f}")
-print(f"Forces: {results['forces']}")
-print(f"Stress: {results['stress']}")
+# concatenate atomic numbers into a (16,) array
+atomic_numbers_numpy = np.concatenate(
+    [atoms.get_atomic_numbers() for atoms in atoms_list]
+)
+
+# convert to tensors
+positions = torch.tensor(positions_numpy, device=device, dtype=dtype)
+cell = torch.tensor(cell_numpy, device=device, dtype=dtype)
+atomic_numbers = torch.tensor(atomic_numbers_numpy, device=device, dtype=torch.int)
+
+# create batch index array of shape (16,) which is 0 for first 8 atoms, 1 for last 8 atoms
+atoms_per_batch = torch.tensor(
+    [len(atoms) for atoms in atoms_list], device=device, dtype=torch.int
+)
+batch = torch.repeat_interleave(
+    torch.arange(len(atoms_per_batch), device=device), atoms_per_batch
+)
+
+# You can see their shapes are as expected
+print(f"Positions: {positions.shape}")
+print(f"Cell: {cell.shape}")
+print(f"Atomic numbers: {atomic_numbers.shape}")
+print(f"Batch: {batch.shape}")
+
+# Now we can pass them to the model
+results = batched_model(
+    dict(
+        positions=positions,
+        cell=cell,
+        atomic_numbers=atomic_numbers,
+        batch=batch,
+        pbc=True,
+    )
+)
+
+# The energy has shape (n_batches,) as the structures in a batch
+print(f"Energy: {results['energy'].shape}")
+
+# The forces have shape (n_atoms, 3) same as positions
+print(f"Forces: {results['forces'].shape}")
+
+# The stress has shape (n_batches, 3, 3) same as cell
+print(f"Stress: {results['stress'].shape}")
+
+# Check if the energy, forces, and stress are the same for the Si system across the batch
+print(torch.max(torch.abs(results["energy"][0] - results["energy"][1])))
+print(torch.max(torch.abs(results["forces"][0] - results["forces"][1])))
+print(torch.max(torch.abs(results["stress"][0] - results["stress"][1])))

examples/scripts/1_Introduction/1.4_Fairchem.py → examples/scripts/1_Introduction/1.3_Fairchem.py

@@ -3,7 +3,7 @@
 
 # /// script
 # dependencies = [
-#     "fairchem-core",
+#     "fairchem-core==1.10.0",
 # ]
 # ///