`kups.core` ¶

Core library for JAX-based particle simulations.

This package provides the foundational components for building molecular dynamics and Monte Carlo simulations with JAX. The architecture emphasizes composability, type safety, and compatibility with JAX transformations.

Module Organization¶

Data Structures and Manipulation¶

data: PyTree wrappers (Batched, Table) for structured data
lens: Functional lenses for accessing and modifying nested structures
unitcell: Periodic boundary conditions and crystallographic unit cells

Computation and Simulation¶

potential: Energy calculations with gradients and Hessians
propagator: State evolution (MD integrators, MC moves, composition)
neighborlist: Efficient neighbor search algorithms (cell lists, dense)

System Management¶

assertion: Runtime validation with automatic fixing capabilities
capacity: Dynamic array resizing with automatic capacity management
patch: State modification protocols and implementations
result: Result types with runtime assertion tracking

Utilities¶

parameter_scheduler: Adaptive parameter tuning (step sizes, temperatures)
constants: Physical constants and unit conversions
utils: Mathematical utilities, JAX helpers, functional programming tools
storage: HDF5-based trajectory logging

Design Philosophy¶

Composability: Small, focused components that combine easily
Type Safety: Extensive use of generics and protocols
JAX Integration: All components work with jit, grad, vmap
Immutability: Functional updates via lenses and patches
Performance: Optimized algorithms with automatic capacity management

Common Patterns¶

# Build a potential from components
potential = sum_potentials(
    bonded_potential,
    lennard_jones_potential,
    coulomb_potential
)

# Compose propagators
propagator = compose_propagators(
    neighbor_list_update,
    monte_carlo_move,
    parameter_adjustment
)

# Run simulation with logging
with h5py.File("traj.h5", "w") as f:
    writer = HDF5StorageWriter.init(f, config, state, total_steps)
    with writer.background_writer() as bg:
        for step in range(total_steps):
            state = propagator(key, state)
            bg.write(state, step)

See individual module documentation for detailed API references:

kups.core ¶