Feat/allow different temperatures ts integrate

[thomasloux]

Summary

Change ts.integrate to accept temperature of size (n_system,), so that its system can have its own temperature.

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
• Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.

[thomasloux]

I didn't add in tests, but it's working:

import numpy as np
import torch
from ase.build import bulk
from pymatgen.core import Structure

import torch_sim as ts
from torch_sim.models.lennard_jones import LennardJonesModel
from torch_sim.trajectory import TorchSimTrajectory, TrajectoryReporter
from torch_sim.units import MetalUnits

n=4
n_steps = 1500
trajectory_file = [f"lj_trajectory{i}.h5md" for i in range(n)]
# report potential energy every 10 steps and kinetic energy every 20 steps
prop_calculators = {
    10: {"potential_energy": lambda state: state.energy},
    20: {
        "kinetic_energy": lambda state: ts.calc_kinetic_energy(
            momenta=state.momenta, masses=state.masses
        ),
        "temperature": lambda state: state.calc_temperature(),
    },
}

reporter = TrajectoryReporter(
    trajectory_file,
    # report state every 10 steps
    state_frequency=10,
    prop_calculators=prop_calculators,
)


lj_model = LennardJonesModel(
    sigma=2.0,  # Å, typical for Si-Si interaction
    epsilon=0.1,  # eV, typical for Si-Si interaction
    device=torch.device("cpu"),
    dtype=torch.float64,
)

si_atoms = bulk("Si", "fcc", a=5.43, cubic=True)
temperatures = (torch.arange(n) + 1) * 500
# temperatures = temperatures.unsqueeze(0).repeat(n_steps, 1)

final_state = ts.integrate(
    system=[si_atoms.copy() for _ in range(n)],
    model=lj_model,
    integrator=ts.Integrator.nvt_langevin,
    n_steps=n_steps,
    temperature=temperatures,
    timestep=0.002,
    trajectory_reporter=reporter,
)

temperatures = []
for filename in trajectory_file:
    with TorchSimTrajectory(filename) as traj:
        temperatures.append(traj.get_array("temperature").flatten())

temperatures = np.stack(temperatures, axis=1)  # shape (n_steps, n)
import matplotlib.pyplot as plt
# smoothen temperatures
temperatures = np.stack([np.convolve(temperatures[:, i], np.ones(20)/20, mode='valid') for i in range(temperatures.shape[1])], axis=1)
plt.plot(temperatures)

[orionarcher]

Awesome! This is a nice feature. I'd suggest adding a bit more testing fo the different temperature shapes though.

[orionarcher]

I honestly might just throw a well-documented error here and demand the user provide a 2D tensor if n_systems == n_steps. It's a rare edge case and the consequences of incorrectly guessing the default is potentially many hours of wasted debugging.

[thomasloux]

I added a warning, let me know if you want to be more strict and raise an error

[orionarcher]

Sounds good to me

[thomasloux]

Note for me: I need to add a test to check that it works with the autobatcher

[thomasloux]

@orionarcher I modified the autobatcher.bin_index which was a list[dict[int, float]] which I think is just the results of the original use of the binpacking from the first version of autobatcher #39. I changed the bin_index to actually be list[list[indices]]. Otherwise
for state, indices in batcher: indices would be a super weird like {1: float, 2: float} and so you can't use it directly like: myTensor[indices]

[orionarcher]

Would be nice to have a test with a 2D array but otherwise LGTM

[orionarcher]

Ready to merge @thomasloux?