Half supercell error when using OrbModel with batched inputs #247
Description
OrbModel wraps orb-models. An optimization used is half_supercell. While I'm not sure what is exactly performed, the orb repo only allows this optimization with volume of a cell > 1000 (https://github.com/orbital-materials/orb-models/blob/4c3bbdd11fbb08c95455be70aaf7057a20fd0a2d/orb_models/forcefield/featurization_utilities.py#L331). This is checked in OrbModel but it uses the maximum value, while it should ensure all crystals have a volume greater than 1000.
Minimum code to reproduce the error.
from ase.build import bulk
from torch_sim.models.orb import OrbModel
from torch_sim.state import initialize_state
import torch_sim as ts
from torch_sim.integrators import nvt
import torch
# ORB model imports
from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator
device = "cpu"
input = bulk("Ti", "hcp", a=2.95, c=4.68) * 4
input2 = bulk("Ti", "hcp", a=2.95, c=4.68) * 3
input = initialize_state([input, input2], dtype=torch.float32, device=device)
cells = input.cell
# compute volume
print(torch.det(cells)) # tensor([2257.3582, 952.3230])
orbff = pretrained.orb_v3_conservative_inf_omat(
device=device,
precision="float32-highest",
compile=False
)
orb_model = OrbModel(model=orbff, dtype=torch.get_default_dtype())
temperature = 300.0 # Kelvin
timestep = 0.001
pressure = -1
# Run batched NPT MD
results = ts.integrate(
system=input,
model=orb_model,
integrator=nvt.nvt_langevin,
n_steps=10,
temperature=temperature,
timestep=timestep,
)