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.