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]:
import numbers
import hoomd
import numpy
cpu = hoomd.device.CPU()
simulation = 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 = numpy.meshgrid(
*(numpy.linspace(-box_L / 2, box_L / 2, N, endpoint=False),) * 3
)
positions = numpy.array((x.ravel(), y.ravel(), z.ravel())).T
snap.particles.position[:] = positions
snap.particles.types = ['A']
snap.particles.typeid[:] = 0
simulation.create_state_from_snapshot(snap)
rng = numpy.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 = numpy.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] = new_velocity
self._state.set_snapshot(snap)
@staticmethod
def _get_direction():
theta, z = rng.random(2)
theta *= 2 * numpy.pi
z = 2 * (z - 0.5)
return numpy.array(
[
numpy.sqrt(1 - (z * z)) * numpy.cos(theta),
numpy.sqrt(1 - (z * z)) * numpy.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]:
simulation.state.thermalize_particle_momenta(filter=hoomd.filter.All(), kT=1.0)
lj = hoomd.md.pair.LJ(nlist=hoomd.md.nlist.Cell(buffer=0.4))
lj.params[('A', 'A')] = {'epsilon': 1.0, 'sigma': 1.0}
lj.r_cut[('A', 'A')] = 2.5
integrator = hoomd.md.Integrator(
methods=[hoomd.md.methods.ConstantVolume(hoomd.filter.All())], forces=[lj], dt=0.005
)
thermo = hoomd.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)
simulation.operations += integrator
simulation.operations += thermo
simulation.operations += table
# Create and add our custom updater
energy_operation = hoomd.update.CustomUpdater(
action=InsertEnergyUpdater(10.0), trigger=100
)
simulation.operations += energy_operation
[4]:
simulation.run(1000)
kinetic_energy potential_energy total_energy
187.44685 -0.13298 187.31387
195.16975 -5.51636 189.65338
209.79330 -7.55467 202.23864
219.90712 -8.52161 211.38551
229.98365 -10.48876 219.49489
239.84999 -8.49731 231.35268
246.07432 -10.46920 235.60512
246.31498 -10.08009 236.23489
253.43790 -6.76563 246.67227
269.87772 -9.59299 260.28473
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, numbers.Number):
self._energy = hoomd.variant.Constant(new_energy)
elif isinstance(new_energy, hoomd.variant.Variant):
self._energy = new_energy
else:
message = 'energy must be a variant or real number.'
raise ValueError(message)
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 = numpy.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] = new_velocity
self._state.set_snapshot(snap)
@staticmethod
def _get_direction():
theta, z = rng.random(2)
theta *= 2 * numpy.pi
z = 2 * (z - 0.5)
return numpy.array(
[
numpy.sqrt(1 - (z * z)) * numpy.cos(theta),
numpy.sqrt(1 - (z * z)) * numpy.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.0, std=2.0)
sample_energies = numpy.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]:
simulation.operations.updaters.remove(energy_operation)
# Create and add our custom updater
energy_operation = hoomd.update.CustomUpdater(
action=InsertEnergyUpdater(energy), trigger=100
)
simulation.operations.updaters.append(energy_operation)
simulation.run(1000)
285.96876 -10.45450 275.51427
303.84204 -8.40413 295.43791
320.12541 -6.86564 313.25977
333.13502 -12.74600 320.38903
346.94076 -13.49587 333.44489
359.78511 -13.30828 346.47683
363.05476 -5.18245 357.87231
385.26368 -9.41100 375.85268
384.85984 -2.13054 382.72930
400.94032 -9.80655 391.13377
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.