Custom Writer

Overview

Questions

  • How could I write a custom trajectory writer?

Objectives

  • Show an example custom writer

Boilerplate Code

[1]:
import h5py
import hoomd
import hoomd.hpmc as hpmc
import numpy as np

cpu = hoomd.device.CPU()
sim = hoomd.Simulation(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)

Problem

For this section, we will demonstrate writing a custom trajectory writer using h5py. We will start by implementing the ability to store positions, timesteps, and box dimensions in an HDF5 file.

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

    def __init__(self, filename, mode):
        self.filename = filename
        if mode not in {'w', 'w-', 'x', 'a', 'r+'}:
            raise ValueError("mode must be writtable")
        self.file = h5py.File(filename, mode)
        self.write_metadata()
        frames = list(self.file.keys())
        if frames:
            self._cur_frame = max(map(int, frames)) + 1
        else:
            self._cur_frame = 1

    def write_metadata(self):
        """Write the file metadata that defines the type of hdf5 file"""
        if 'app' in self.file.attrs:
            if self.file.attrs.app != 'hoomd-v3':
                raise RuntimeError(
                    'HDF5 file metadata "app" is not "hoomd-v3".')
        else:
            self.file.attrs.app = 'hoomd-v3'

        if 'version' not in self.file.attrs:
            self.file.attrs.version = '1.0'

    def act(self, timestep):
        """Write out a new frame to the trajectory."""
        new_frame = self.file.create_group(str(self._cur_frame))
        self._cur_frame += 1
        positions = new_frame.create_dataset('positions',
                                             (self._state.N_particles, 3),
                                             dtype='f8')
        snapshot = self._state.get_snapshot()
        positions[:] = snapshot.particles.position
        new_frame.attrs['timestep'] = timestep
        box_array = np.concatenate((self._state.box.L, self._state.box.tilts))
        new_frame.attrs['box'] = box_array

    def __del__(self):
        self.file.close()

Define a function that creates a HDF5Writer wrapped in a custom writer.

This function will make creating our custom writer easier. We will now add an HPMC sphere integrator and our custom writer to our simulation and run for 1000 steps.

(Note that the ‘w’ mode will truncate any existing file.)

[3]:
h5_writer = hoomd.write.CustomWriter(action=HDF5Writer('traj.h5', 'w'),
                                     trigger=100)
integrator = hpmc.integrate.Sphere()
integrator.shape['A'] = {'diameter': 1.}

sim.operations += integrator
sim.operations += h5_writer

sim.run(1000)

We have run the simulation, and our HDF5 file has been written. Lets check the groups our file contains now.

[4]:
list(h5_writer.file.keys())
[4]:
['1', '10', '2', '3', '4', '5', '6', '7', '8', '9']

Ten frames have been written as expected. Let’s check the properties from the last frame and compare them to the simulation currently. We will open the file again in read only mode to check these properties. First we flush the open HDF5 file to ensure the data has been written to the OS buffer at least.

[5]:
h5_writer.file.flush()

with h5py.File('traj.h5', 'r') as traj:
    assert traj['10'].attrs['timestep'] == sim.timestep
    box_array = np.concatenate((sim.state.box.L, sim.state.box.tilts))
    assert np.allclose(traj['10'].attrs['box'], box_array)
    snapshot = sim.state.get_snapshot()
    assert np.allclose(snapshot.particles.position, traj['10']['positions'][:])

Expanding on HDF5Writer

Our HDF5Writer class is already sufficient for storing the trajectory. However, there are plenty of other features we could add. Examples include utilizing the HOOMD-blue logging subsystem to allow logging data to the HDF5 file, and adding support for MPI. Also, we could also add support for other system properties such as images, velocities, and others. We will focus on adding this feature.

We need to decide on a method of specifying properties to write. We will use a tuple system where we signify the property we want to store using a tuple that nests into a snapshot object. For example to write images we will use the tuple ('particles', 'image') to signify we want to store images. We will let an user pass in a list of tuples of any length to specify what they want to store. (Positions will always be stored, and we will move them to the particles group).

[6]:
class HDF5Writer(hoomd.custom.Action):

    def __init__(self, filename, mode, properties):
        self.filename = filename
        self.properties = set(properties) | {('particles', 'position')}
        if mode not in {'w', 'w-', 'x', 'a', 'r+'}:
            raise ValueError("mode must be writtable")
        self.file = h5py.File(filename, mode)
        self.write_metadata()
        frames = list(self.file.keys())
        if frames:
            self._cur_frame = max(map(int, frames)) + 1
        else:
            self._cur_frame = 1

    def write_metadata(self):
        """Write the file metadata that defines the type of hdf5 file"""
        if 'app' in self.file.attrs:
            if self.file.attrs.app != 'hoomd-v3':
                raise RuntimeError(
                    'HDF5 file metadata "app" is not "hoomd-v3".')
        else:
            self.file.attrs.app = 'hoomd-v3'

        if 'version' not in self.file.attrs:
            self.file.attrs.version = '1.0'

    def _set_property(self, base_group, prop):
        # Get data array
        data = self._state.get_snapshot()
        for name in prop:
            data = getattr(data, name)
        # Get dataset
        use_group = base_group
        for name in prop[:-1]:
            if name not in use_group:
                use_group = base_group.create_group(name)
            else:
                use_group = base_group[name]
        dataset = use_group.create_dataset(prop[-1],
                                           data.shape,
                                           dtype=str(data.dtype))
        dataset[:] = data

    def act(self, timestep):
        """Write out a new frame to the trajectory."""
        new_frame = self.file.create_group(str(self._cur_frame))
        self._cur_frame += 1
        for prop in self.properties:
            self._set_property(new_frame, prop)
        new_frame.attrs['timestep'] = timestep
        box_array = np.concatenate((self._state.box.L, self._state.box.tilts))
        new_frame.attrs['box'] = box_array

    def __del__(self):
        self.file.close()

We will now use our extended trajectory writer to write out particle images as well.

[7]:
h5_writer.file.close()
sim.operations -= h5_writer
h5_writer = hoomd.write.CustomWriter(action=HDF5Writer(
    'traj.h5', 'w', [('particles', 'image')]),
                                     trigger=100)
sim.operations.writers.append(h5_writer)
sim.run(1000)

To see that this worked we will check the first frame for particle images.

[8]:
h5_writer.file.flush()

with h5py.File('traj.h5', 'r') as traj:
    display(traj['1']['particles']['image'][:10])
    display(traj['1']['particles']['position'][:10])
array([[-1,  0,  0],
       [-1,  0,  0],
       [-1, -1,  0],
       [-1,  0,  0],
       [-1, -1,  0],
       [ 0,  0,  0],
       [ 0,  0,  0],
       [ 0,  0,  0],
       [ 0,  0,  0],
       [ 0,  0,  0]], dtype=int32)
array([[ 7.08916892, -9.1162494 , -9.03166465],
       [ 7.17894751, -7.55565503, -4.1420896 ],
       [ 6.84048969,  9.85875838, -0.14319672],
       [ 9.42302572, -7.66224406,  1.71042043],
       [ 4.04383384,  8.15467659,  9.35673311],
       [-5.21819354, -9.57761671, -7.17922194],
       [-6.56869188, -9.00928178, -7.91171588],
       [-1.41025576, -9.14286987, -3.21326451],
       [-3.29261443, -8.20593309,  2.56455928],
       [-2.02993862, -3.93072604,  4.98365   ]])

We could continue add more features such as argument validation in the constructor, support for the logging subsystem of HOOMD-blue, a classmethod, or a number of other things. However, these are left as exercises. This section has shown a non-trivial application of the custom action feature in HOOMD-blue for custom writers.