tutorials/tensorclass_particles

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

Simulating particles with a tensorclass

A TensorClass can hold the state of a computation, not only a dataset. This tutorial writes a small gravity simulation: a Particles class with a position, a velocity, and a mass for each particle, and a step method that moves the particles forward in time. The same method runs on one system of particles or on a batch of independent systems. The tutorial is for readers who write simulations or other numerical code in PyTorch, and who keep the state in several tensors that must stay aligned.

What you will learn:

• How to define a tensorclass whose batch dimensions are the particles, with methods that run on any number of leading dimensions.

• How to run the same step on one system or on a batch of systems.

• How to compile step with torch.compile().

• How to stack the states of each step into a trajectory, index it, and save it with memmap.

import tempfile
from pathlib import Path

import torch
from tensordict import TensorClass

torch.manual_seed(0)

<torch._C.Generator object at 0x7fafda560c50>

The state of a system

Each particle has a position and a velocity in three dimensions, and a mass. The Particles class below declares these three fields. Its batch size is the shape of the particles: [n] for one system of n particles, and [b, n] for b systems. Each field has the batch dimensions first, then its own dimensions: position and velocity end with a dimension of size 3, and mass has none.

The methods treat the last batch dimension as the particles of one system, and broadcast over the dimensions before it:

• acceleration sums the gravitational pull of all the other particles on each particle. The constant EPS softens the pull at short distances, so that two close particles do not get a very large acceleration.

• step moves the state forward by dt with the velocity Verlet method. It builds each new Particles with replace(), which returns a copy with some fields replaced, so the state that you call step on does not change.

• energy returns the kinetic plus the potential energy of each system. For a small dt, it stays almost constant from step to step.

The gravitational constant is 1, and the units are arbitrary.

EPS = 0.2


class Particles(TensorClass):
    position: torch.Tensor
    velocity: torch.Tensor
    mass: torch.Tensor

    @classmethod
    def random(cls, *batch_size):
        return cls(
            position=torch.randn(*batch_size, 3),
            velocity=0.1 * torch.randn(*batch_size, 3),
            mass=torch.rand(batch_size) + 0.5,
            batch_size=batch_size,
        )

    def displacement(self):
        # The vector from particle i to particle j, for each pair (i, j).
        return self.position.unsqueeze(-3) - self.position.unsqueeze(-2)

    def acceleration(self):
        d = self.displacement()
        inv_r3 = (d.pow(2).sum(-1) + EPS**2) ** -1.5
        # A particle does not pull itself: its displacement d[..., i, i] is 0.
        return (d * (self.mass.unsqueeze(-2) * inv_r3).unsqueeze(-1)).sum(-2)

    def step(self, dt):
        velocity = self.velocity + 0.5 * dt * self.acceleration()
        moved = self.replace(position=self.position + dt * velocity)
        velocity = velocity + 0.5 * dt * moved.acceleration()
        return moved.replace(velocity=velocity)

    def energy(self):
        kinetic = 0.5 * (self.mass * self.velocity.pow(2).sum(-1)).sum(-1)
        r = (self.displacement().pow(2).sum(-1) + EPS**2).sqrt()
        pair = self.mass.unsqueeze(-1) * self.mass.unsqueeze(-2) / r
        # Count each pair once, and leave out i == j.
        potential = -pair.triu(diagonal=1).sum((-2, -1))
        return kinetic + potential

One system

A system of 8 particles has batch_size=[8]. The representation shows the shape of each field:

system = Particles.random(8)
print(system)

Particles(batch_size=[8], fields={
    mass: f32[],
    position: f32[3],
    velocity: f32[3]})

step returns the next state, with the same batch size. The masses do not change, and replace reuses the same mass tensor.

dt = 0.01
next_system = system.step(dt)
assert next_system.batch_size == torch.Size([8])
assert next_system.mass is system.mass
assert not torch.equal(next_system.position, system.position)

A batch of systems

Four independent systems of 8 particles each make a Particles with batch_size=[4, 8]. The step method runs on it without a change: it computes the pairwise forces inside each system, and not between systems.

systems = Particles.random(4, 8)
next_systems = systems.step(dt)
assert next_systems.batch_size == torch.Size([4, 8])
assert systems.energy().shape == torch.Size([4])

An index on the first dimension selects one system, with all its fields. The step of one system on its own gives the same result as its slice of the batched step:

torch.testing.assert_close(systems[2].step(dt).position, next_systems[2].position)

Compiling the step

torch.compile() traces through a tensorclass: the compiled graph holds only the tensor operations of step. To compile the method, compile the function Particles.step and pass the state as its first argument.

The first call compiles the function, which takes a few seconds. Later calls with the same batch size reuse the compiled code. A call with a new batch size can compile again: torch.compile, torch.func, torch.export, and CUDA graphs lists what makes a function recompile. The compiled code fuses the element-wise operations of step into fewer kernels.

compiled_step = torch.compile(Particles.step)

torch.testing.assert_close(compiled_step(systems, dt).position, next_systems.position)

A trajectory

The loop below runs 200 steps, and keeps each state in a list. torch.stack() stacks the 201 states, the first one included, into one Particles with batch_size=[201, 4, 8]: time, system, particle.

n_steps = 200
states = [systems]
for _ in range(n_steps):
    states.append(compiled_step(states[-1], dt))
trajectory = torch.stack(states)
assert isinstance(trajectory, Particles)
assert trajectory.batch_size == torch.Size([n_steps + 1, 4, 8])

The trajectory is a Particles like the others, so the methods and the indexing work on it too. trajectory[-1] is the last state of the four systems, and trajectory[:, 0] is the history of the first system. trajectory.energy() has one value for each time step and each system. The relative change of the energy over the run stays small:

assert trajectory[-1].batch_size == torch.Size([4, 8])
assert trajectory[:, 0].batch_size == torch.Size([n_steps + 1, 8])

energy = trajectory.energy()
assert energy.shape == torch.Size([n_steps + 1, 4])
drift = ((energy - energy[0]) / energy[0]).abs().max()
print(f"largest relative change of the energy: {drift:.1e}")

largest relative change of the energy: 6.5e-03

Saving and loading the trajectory

memmap() writes each field to a file in a directory, as a memory-mapped tensor. Particles.load_memmap maps the files back into a Particles. Call load_memmap on the class itself when the class is defined in a script or a notebook, as here: another program could not import it from __main__.

The loaded fields are memory-mapped, so an index reads only the part of the files that it needs, such as the last state:

with tempfile.TemporaryDirectory() as tmpdir:
    path = Path(tmpdir) / "trajectory"
    trajectory.memmap(path)
    loaded = Particles.load_memmap(path)
    assert isinstance(loaded, Particles)
    assert (loaded == trajectory).all()
    torch.testing.assert_close(loaded[-1].position, trajectory[-1].position)

A long run might not fit in memory. To write each state to disk as the simulation runs, allocate the files first with memmap_like(), then assign each state to its time step. expand() gives the shape of the whole run without allocating it.

with tempfile.TemporaryDirectory() as tmpdir:
    path = Path(tmpdir) / "long_run"
    on_disk = systems.expand(n_steps + 1, *systems.batch_size).memmap_like(path)
    state = systems
    on_disk[0] = state
    for t in range(1, n_steps + 1):
        state = compiled_step(state, dt)
        on_disk[t] = state
    torch.testing.assert_close(
        Particles.load_memmap(path).position, trajectory.position
    )

Conclusion

You can now hold the state of a simulation in a TensorClass with a typed field for each quantity. Its methods run on one system, a batch of systems, or a whole trajectory, because they use only the last batch dimension. You can compile a method with torch.compile(), stack the states into a trajectory with torch.stack(), and save it to disk with memmap.

Further reading

• TensorClass: a dataclass of tensors explains how to define and use tensorclasses.

• torch.compile, torch.func, torch.export, and CUDA graphs lists what works with torch.compile(), and what makes a function compile again.

• Save and load tensordicts compares the ways to save a tensordict or a tensorclass.

• Filling a TensorDict row by row fills a tensordict step by step.

• Using tensorclasses for datasets uses a tensorclass to load a dataset.