Custom Updater#

Overview#

Questions#

  • How can I modify the state of a system in a custom updater?

Objectives#

  • Show an example of a non-trival custom updater.

Boilerplate Code#

[1]:
from numbers import Number

import hoomd
import hoomd.md as md
import numpy as np

cpu = hoomd.device.CPU()
sim = hoomd.Simulation(device=cpu, seed=1)

# Create a simple cubic configuration of particles
N = 5  # particles per box direction
box_L = 20  # box dimension

snap = hoomd.Snapshot(cpu.communicator)
snap.configuration.box = [box_L] * 3 + [0, 0, 0]
snap.particles.N = N**3
x, y, z = np.meshgrid(*(np.linspace(-box_L / 2, box_L / 2, N, endpoint=False),)
                      * 3)
positions = np.array((x.ravel(), y.ravel(), z.ravel())).T
snap.particles.position[:] = positions
snap.particles.types = ['A']
snap.particles.typeid[:] = 0

sim.create_state_from_snapshot(snap)
rng = np.random.default_rng(1245)

Problem#

In this section, we will show how to create a custom updater that modifies the system state. To show this, we will create a custom updater that adds a prescribed amount of energy to a single particle simulating the bombardment of radioactive material into our system. For this problem, we pick a random particle and modify its velocity according to the radiation energy in a random direction.

[2]:
class InsertEnergyUpdater(hoomd.custom.Action):

    def __init__(self, energy):
        self.energy = energy

    def act(self, timestep):
        snap = self._state.get_snapshot()
        if snap.communicator.rank == 0:
            particle_i = rng.integers(snap.particles.N)
            mass = snap.particles.mass[particle_i]
            direction = self._get_direction()
            magnitude = np.sqrt(2 * self.energy / mass)
            velocity = direction * magnitude
            old_velocity = snap.particles.velocity[particle_i]
            new_velocity = old_velocity + velocity
            snap.particles.velocity[particle_i] = velocity
        self._state.set_snapshot(snap)

    @staticmethod
    def _get_direction():
        theta, z = rng.random(2)
        theta *= 2 * np.pi
        z = 2 * (z - 0.5)
        return np.array([
            np.sqrt(1 - (z * z)) * np.cos(theta),
            np.sqrt(1 - (z * z)) * np.sin(theta), z
        ])

We will now use our custom updater with an NVE integrator. Particles will interact via a Lennard-Jones potential. Using the Table writer and a hoomd.logging.Logger, we will monitor the energy, which should be increasing as we are adding energy to the system. We will also thermalize our system to a kT == 1.

[3]:
sim.state.thermalize_particle_momenta(filter=hoomd.filter.All(), kT=1.)

lj = md.pair.LJ(nlist=md.nlist.Cell(buffer=0.4))
lj.params[('A', 'A')] = {'epsilon': 1., 'sigma': 1.}
lj.r_cut[('A', 'A')] = 2.5
integrator = md.Integrator(methods=[md.methods.NVE(hoomd.filter.All())],
                           forces=[lj],
                           dt=0.005)

thermo = md.compute.ThermodynamicQuantities(hoomd.filter.All())
logger = hoomd.logging.Logger(categories=['scalar'])
logger.add(thermo, ['kinetic_energy', 'potential_energy'])
logger['total_energy'] = (
    lambda: thermo.kinetic_energy + thermo.potential_energy, 'scalar')

table = hoomd.write.Table(100, logger, max_header_len=1)

sim.operations += integrator
sim.operations += thermo
sim.operations += table

# Create and add our custom updater
energy_operation = hoomd.update.CustomUpdater(action=InsertEnergyUpdater(10.),
                                              trigger=100)

sim.operations += energy_operation
[4]:
sim.run(1000)
 kinetic_energy  potential_energy   total_energy
   189.55469         -0.16021        189.39447
   203.13934         -5.51636        197.62298
   214.05941         -7.90628        206.15314
   219.49181         -8.47534        211.01647
   230.38656         -9.71804        220.66852
   237.66638         -9.07038        228.59601
   245.73067        -10.18110        235.54957
   254.95301        -12.21494        242.73808
   258.55741         -6.01446        252.54296
   269.38334         -8.12332        261.26002

As we can see the total energy of the system is indeed increasing. The energy isn’t increasing by 10 every time since we are adding the velocity in a random direction which may be against the current velocity.

Improving upon our Custom Action#

Maybe we want to allow for the energy to be from a distribution. HOOMD-blue has a concept called a variant which allows for quantities that vary over time. Let’s change the InsertEnergyupdater to use variants and create a custom variant that grabs a random number from a Gaussian distribution. (If you don’t understand the variant code, that is fine. We are just using this to showcase how you can iteratively improve custom actions).

Note:

Variant objects model a parameter as a function of the timestep, so to get the value for a particular timestep we have to call the variant. For more information see the documentation for hoomd.variant.

[5]:
class InsertEnergyUpdater(hoomd.custom.Action):

    def __init__(self, energy):
        self._energy = energy

    @property
    def energy(self):
        """A `hoomd.variant.Variant` object."""
        return self._energy

    @energy.setter
    def energy(self, new_energy):
        if isinstance(new_energy, Number):
            self._energy = hoomd.variant.Constant(new_energy)
        elif isinstance(new_energy, hoomd.variant.Variant):
            self._energy = new_energy
        else:
            raise ValueError("energy must be a variant or real number.")

    def act(self, timestep):
        snap = self._state.get_snapshot()
        if snap.communicator.rank == 0:
            particle_i = rng.integers(snap.particles.N)
            mass = snap.particles.mass[particle_i]
            direction = self._get_direction()
            magnitude = np.sqrt(2 * self.energy(timestep) / mass)
            velocity = direction * magnitude
            old_velocity = snap.particles.velocity[particle_i]
            new_velocity = old_velocity + velocity
            snap.particles.velocity[particle_i] = velocity
        self._state.set_snapshot(snap)

    @staticmethod
    def _get_direction():
        theta, z = rng.random(2)
        theta *= 2 * np.pi
        z = 2 * (z - 0.5)
        return np.array([
            np.sqrt(1 - (z * z)) * np.cos(theta),
            np.sqrt(1 - (z * z)) * np.sin(theta), z
        ])


class GaussianVariant(hoomd.variant.Variant):

    def __init__(self, mean, std):
        hoomd.variant.Variant.__init__(self)
        self.mean = mean
        self.std = std

    def __call__(self, timestep):
        return rng.normal(self.mean, self.std)

We briefly show that the Gaussian Variant works.

[6]:
energy = GaussianVariant(mean=10., std=2.)
sample_energies = np.array([energy(0) for _ in range(1000)])
f"Mean: {sample_energies.mean()}, std. dev. {sample_energies.std()}"
[6]:
'Mean: 10.069550459202723, std. dev. 1.9965744919420398'

We now use the updated InsertEnergyUpdater in the simulation.

[7]:
sim.operations.updaters.remove(energy_operation)
# Create and add our custom updater
energy_operation = hoomd.update.CustomUpdater(
    action=InsertEnergyUpdater(energy), trigger=100)
sim.operations.updaters.append(energy_operation)
sim.run(1000)
   279.83085        -11.21077        268.62009
   289.14791         -9.62083        279.52708
   302.04414         -9.22880        292.81534
   312.06778        -11.05103        301.01675
   324.69061        -11.95850        312.73211
   332.29403         -9.80310        322.49093
   345.55571        -14.83428        330.72143
   357.07548        -16.49285        340.58263
   357.62391        -11.70456        345.91935
   372.21959        -15.91459        356.30500

We could continue to improve upon this updater and the execution of this operation. However, this suffices to showcase the ability of non-trivial updaters to affect the simulation state.