Improving Performance

Overview

Questions

  • How can I write custom actions to be as efficient as possible?

Objectives

  • Mention common means for improving performance.

  • Demonstrate using the local snapshot API for increased performance.

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(cpu, seed=1)

# Create a simple cubic configuration of particles
N = 12  # particles per box direction
box_L = 50  # 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)

sim.state.thermalize_particle_momenta(hoomd.filter.All(), 1.)

lj = md.pair.LJ(nlist=md.nlist.Cell())
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)

sim.operations += integrator


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)


energy = GaussianVariant(0.1, 0.001)
sim.run(0)
rng = np.random.default_rng(1245)

General Guidelines

When trying to improve the performance of custom actions, the first step is always to profile the class. Python comes with profiling tools that can be used to determine bottlenecks in a custom action’s performance. In addition, there are many external visualization and profiling tools available. However, after profiling here are some tips that should help improve performance.

  • State.get_snapshot aggregates data across MPI ranks and is \(O(n)\) to construct and setting the state to a new snapshot \(O(n)\) as well. However, hoomd.State.cpu_local_snaphshot or hoomd.State.gpu_local_snapshot are on order \(O(1)\) to construct and modifying data in a local snapshot is \(O(1)\) as well.

  • HOOMD-blue makes use of properties heavily. Since users can change the system state in Python at any point, we must recompute many of these quantities every time they are queried. If you are using something like hoomd.md.pair.LJ.energies multiple times, it will be more performant to first store the values and then use that copy.

  • Avoid for loops for numerical calculation. Try to utilize NumPy broadcasting or existing functions in NumPy or Scipy on the CPU or CuPy on the GPU.

Improve InsertEnergyUpdater

As an example, we will improve the performance of the InsertEnergyUpdater. Specifically we will change to use the cpu_local_snapshot to update particle velocity. We will use the %%timeit magic function for timing the simulation’s run time before and after our optimization. To highlight the differnce, we will run the updater every timestep.

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

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

    @property
    def energy(self):
        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
        ])
[3]:
energy_action = InsertEnergyUpdater(energy)
energy_operation = hoomd.update.CustomUpdater(action=energy_action, trigger=1)
sim.operations.updaters.append(energy_operation)
[4]:
%%timeit
sim.run(100)
111 ms ± 1.98 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

We now show the profile for the optimized code which uses the cpu_local_snapshot for updating velocities.

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

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

    @property
    def energy(self):
        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 attach(self, simulation):
        self._state = simulation.state
        self._comm = simulation.device.communicator

    def detach(self):
        del self._state
        del self._comm

    def act(self, timestep):
        part_tag = rng.integers(self._state.N_particles)
        direction = self._get_direction()
        energy = self.energy(timestep)
        with self._state.cpu_local_snapshot as snap:
            # We restrict the computation to the MPI
            # rank containing the particle if applicable.
            # By checking if multiple MPI ranks exist first
            # we can avoid for checking inclusion of a tag id
            # in an array.
            if (self._comm.num_ranks <= 1 or part_tag in snap.particles.tag):
                i = snap.particles.rtag[part_tag]
                mass = snap.particles.mass[i]
                magnitude = np.sqrt(2 * energy / mass)
                velocity = direction * magnitude
                old_velocity = snap.particles.velocity[i]
                new_velocity = old_velocity + velocity
                snap.particles.velocity[i] = new_velocity

    @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
        ])
[6]:
# Create and add our modified custom updater
sim.operations -= energy_operation
energy_action = InsertEnergyUpdater(energy)
energy_operation = hoomd.update.CustomUpdater(action=energy_action, trigger=1)
sim.operations.updaters.append(energy_operation)
[7]:
%%timeit
sim.run(100)
35.6 ms ± 855 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

As can be seen the simulation with the new EnergyInsertUpdater about an order of magnitude faster with a system size of \(12^3 = 1728\), by virtue of the local snapshot modification having \(O(1)\) time complexity. At larger system sizes this change will grow to be even more substantial.

This concludes the tutorial on custom actions in Python. For more information see the API documentation.