Visualization with plato

In this notebook, we run a Lennard-Jones simulation, color particles according to their local density computed with freud, and display the results with plato. Note that plato has multiple backends – see the plato documentation for information about each backend and the features it supports.

import itertools

import hoomd
import numpy as np
from hoomd import md


def make_simple_cubic_snapshot(a, n):
    """Make a snapshot with a simple cubic lattice.

    Args:
        a (float): Lattice spacing
        n (int): Number of particles

    Returns:
        hoomd.Snapshot: The initial system snapshot.
    """
    k = int(np.ceil(n ** (1 / 3)))
    L = k * a
    x = np.linspace(-L / 2, L / 2, k, endpoint=False)
    position = list(itertools.product(x, repeat=3))
    position = position[:n]

    snap = hoomd.Snapshot()
    snap.particles.N = n
    snap.particles.types = ["A"]
    snap.particles.typeid[:] = [0] * n
    snap.particles.position[:] = position
    snap.configuration.box = [L, L, L, 0, 0, 0]

    return snap


# Create an 8x8x8 simple cubic lattice
sim = hoomd.Simulation(hoomd.device.CPU())
sim.create_state_from_snapshot(make_simple_cubic_snapshot(a=1.5, n=1000))
sim.seed = 42

# Specify Lennard-Jones interactions between particle pairs
nl = hoomd.md.nlist.Cell(buffer=0.2)
lj = hoomd.md.pair.LJ(default_r_cut=3.0, nlist=nl)
lj.params[("A", "A")] = dict(epsilon=1.0, sigma=1.0)

# Integrate at constant temperature
nvt = md.methods.ConstantVolume(
    filter=hoomd.filter.All(), thermostat=md.methods.thermostats.Bussi(kT=0.01, tau=0.5)
)
integrator = md.Integrator(dt=0.005, methods=[nvt], forces=[lj])
sim.operations.integrator = integrator
sim.state.thermalize_particle_momenta(filter=hoomd.filter.All(), kT=0.01)

# Run for 10,000 time steps
sim.run(10e3)
snap = sim.state.get_snapshot()

Now we import the modules needed for visualization.

import freud
import matplotlib.cm
import numpy as np
import plato

# For interactive scenes, use:
import plato.draw.pythreejs as draw
from matplotlib.colors import Normalize

# For static scenes, use:
# import plato.draw.fresnel as draw

This code sets up the plato Scene object with the particles and colors computed above.

positions = snap.particles.position
box = freud.Box.from_box(snap.configuration.box)
ld = freud.density.LocalDensity(3.0, 1.0)
ld.compute(system=snap)
colors = matplotlib.cm.viridis(Normalize()(ld.density))
radii = np.ones(len(positions)) * 0.5
box_prim = draw.Box.from_box(box, width=0.2)
sphere_prim = draw.Spheres(
    positions=snap.particles.position, radii=radii, colors=colors, vertex_count=32
)
scene = draw.Scene((sphere_prim, box_prim), zoom=1.5)

Click and drag the 3D scene below - it’s interactive!

scene.show()