atomkit

Atom-level analysis toolkit for molecular dynamics trajectories.

Install

pip install atomkit
# or
uv pip install atomkit

Features

SpatialGrid

4D CSR-indexed spatial grid (space + time) for fast region queries on LAMMPS trajectories. Stores as HDF5 with zstd compression and lazy loading (mmap).

CLI:

# Convert LAMMPS trajectory to HDF5 (all timesteps)
atomkit convert simulation.lammpstrj output.h5

# Options
atomkit convert simulation.lammpstrj -c 4.0              # cell size 4Å
atomkit convert simulation.lammpstrj -t 0:100            # timesteps 0-99
atomkit convert simulation.lammpstrj -t 0:1000:10        # every 10th timestep
atomkit convert simulation.lammpstrj -t 0,50,100         # specific timesteps
atomkit convert simulation.lammpstrj --coords unwrapped  # use xu,yu,zu columns
atomkit convert simulation.lammpstrj --coords wrapped    # use x,y,z columns

# Field filtering (smaller files, faster conversion)
atomkit convert traj.lammpstrj out.h5 -f stress_xx,stress_yy,stress_zz  # only these
atomkit convert traj.lammpstrj out.h5 --exclude-fields vx,vy,vz         # exclude these

# Inspect files
atomkit info output.h5                         # show HDF5 grid info
atomkit info simulation.lammpstrj --list-fields  # list fields in LAMMPS file

Python:

from atomkit import SpatialGrid, Region

# Create from LAMMPS file (loads all timesteps by default)
grid = SpatialGrid.from_lammps('simulation.lammpstrj', cell_size=4.0)
grid.save('data.h5')

# Coordinate type: "auto" (default), "unwrapped", "wrapped", or "scaled"
# - unwrapped (xu,yu,zu): actual positions, tracks displacement outside box
# - wrapped (x,y,z): positions wrapped into simulation box
# - scaled (xs,ys,zs): fractional coordinates (0-1)
grid = SpatialGrid.from_lammps('traj.lammpstrj', coord_type='unwrapped')

# Query with 4D regions (returns read-only numpy arrays)
with SpatialGrid.load('data.h5') as grid:
    # Region bounds: (min, max) tuple, single value, or omit for unbounded

    # Single timestep (t=100 means timestep VALUE 100)
    data = grid.query(Region(t=100))

    # Spatial box, all timesteps
    data = grid.query(Region(x=(0, 50), y=(0, 50), z=(0, 50)))

    # Full 4D query
    data = grid.query(Region(x=(0, 50), y=(0, 50), z=(0, 50), t=(0, 1000)))

    # Slice at a point (all cells containing x=25)
    data = grid.query(Region(x=25.0, t=100))

    # Everything
    data = grid.query()  # or Region()

    # Access fields
    data['coords']       # (N, 3) atom positions
    data['stress']       # (N,) stress values
    data['_timestep']    # (N,) which timestep each atom belongs to
    data['_source_idx']  # (N,) original file indices

    # Per-timestep analysis
    for t in np.unique(data['_timestep']):
        mask = data['_timestep'] == t
        print(f"t={t}: mean stress = {data['stress'][mask].mean()}")

    # Fast approximate count (no field reads)
    n_approx = grid.count(Region(x=(0, 50), y=(0, 50), z=(0, 50)))

    # Exact vs cell-level query
    data = grid.query(region)                  # default, exact bounds
    data = grid.query(region, cell_level=True) # faster, includes full boundary cells

    # Add fields later
    grid.add_field('velocity', vel_array)

Region

4D axis-aligned bounding box for space-time queries:

from atomkit import Region

# Flexible bounds specification:
Region(x=(0, 10), y=(0, 10), z=(0, 10))  # Spatial box, all timesteps
Region(t=100)                             # Single timestep, all space
Region(x=5.0)                             # YZ plane at x=5 (slice query)
Region()                                  # Everything (unbounded)

# Region operations
region = Region(x=(0, 100), y=(0, 100), z=(0, 100), t=(0, 1000))
region.volume()                           # Spatial volume
region.subdivide(nx=10, ny=10, nz=10)     # Split into sub-regions
region.with_time(500)                     # Same space, different time
region.expand(padding=5.0)                # Grow bounds

Grid dimensions:

• grid.n_timesteps - number of timesteps
• grid.n_atoms - atoms per timestep
• grid.grid_shape - (nx, ny, nz) spatial cells
• grid.timestep_values - actual timestep values from trajectory

SourceBox

Original simulation box metadata (bounds, tilt, boundary conditions):

# Consolidated box info from LAMMPS
grid.source_box.bounds      # (xlo, xhi, ylo, yhi, zlo, zhi)
grid.source_box.tilt        # (xy, xz, yz) for triclinic boxes
grid.source_box.boundary    # "pp pp pp" (periodic)
grid.source_box.is_triclinic
grid.source_box.is_valid
grid.source_box.contains(x, y, z)  # handles sheared boxes

Cell Aggregates (CellsView)

O(1) cell-level queries using precomputed cumsum arrays:

# Field selection and slicing (both orderings work)
grid.cells["stress"]                    # Select field
grid.cells["stress"][0, :, :, 5]        # Select + slice
grid.cells[0, :, :, 5]["stress"]        # Slice + select

# O(1) reductions (uses cumsum internally)
grid.cells["stress"].sum()              # Total sum (scalar)
grid.cells["stress"].sum("t", "z")      # Sum over t,z → 2D array (x, y)
grid.cells["stress"].mean()             # Mean = sum/count
grid.cells["stress"].count()            # Atom count in region
grid.cells.count("z")                   # Count projection along z

# Sliced queries (O(1) box sums)
view = grid.cells["stress"][0:5, 10:20, :, 5:15]
view.sum()                              # Sum in this 4D box
view.mean()                             # Mean in this box
view.count()                            # Atoms in this box

# Raw arrays (for custom analysis)
grid.cells["stress"].np                 # Raw sum array (t, nx, ny, nz)
grid.cells.counts                       # Atom counts array
grid.cells.fields                       # List of available fields

# Multiple fields at once
grid.cells.sum()                        # {"stress": ..., "velocity": ...}

See docs/cells-api.md for detailed documentation.

GridView

Create views into subregions (shares underlying data, no copy):

# View by coordinate bounds
view = grid.view(x=(0, 50), z=(10, 100))
view.counts       # sliced counts array
view.box_bounds   # adjusted bounds
view.grid_shape   # subset shape

# View constrained to source box
view = grid.view_source_box()

Trimming

Filter atoms outside source box during load:

# Trim atoms outside LAMMPS box bounds (useful for unwrapped coords)
grid = SpatialGrid.from_lammps('traj.dump', trim_to_source_box=True)

Future Extensions

Marching Cubes / Void Visualization

For visualizing empty space (cracks, voids, delamination) and computing volumes:

# Potential future API
verts, faces = grid.void_mesh(Region(t=100), threshold=0.5)
volume = grid.void_volume(Region(t=100), threshold=0.5)

Grid Alignment

For snapping user-selected regions to grid cell boundaries:

# Potential future API
aligned = AlignedRegion.snap(region, grid, mode='enclosing')