atomkit 0.2.0

Atom-level analysis toolkit for molecular dynamics trajectories

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

# Inspect file
atomkit info output.h5

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

Precomputed per-cell statistics (sum/min/max/mean) for fast analysis:

# Precomputed for all numeric fields during grid construction
grid.cells.counts              # (t, nx, ny, nz) atom counts
grid.cells["stress"].sum       # (t, nx, ny, nz) sum per cell
grid.cells["stress"].min       # min per cell
grid.cells["stress"].max       # max per cell
grid.cells["stress"].mean      # mean per cell (sum/counts)
grid.cells.fields              # list of fields with aggregates

# Slicing: get 2D slice at z_idx=5
mean_slice = grid.cells["stress"].mean[0, :, :, 5]  # (nx, ny)
count_slice = grid.counts[0, :, :, 5]

# Projection: sum/mean along an axis
count_proj_z = grid.counts[0].sum(axis=2)           # (nx, ny) - project along z
mean_proj_z = grid.cells["stress"].mean[0].mean(axis=2)

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')