Skip to content

v0.4.0

Latest
Latest

Choose a tag to compare

View all tags
@dallasfoster dallasfoster released this 13 Jul 18:10
v0.4.0
e8e7a74

v0.4.0 Release Notes

This release adds higher-order (multipole) electrostatics, a new cluster-tile neighbor-search strategy with inline pair potentials, 2D slab corrections for Ewald/PME, and differentiable segment operations, alongside a substantial correctness pass on the MTK NPT/NPH barostat and the electrostatics autograd path. It also restructures the neighbors package and deprecates the direct-output training flags on the Ewald/PME APIs in favor of framework autograd.

Breaking Changes

  • MTK cell_velocities are now the strain rate (#77): cell_velocities stores ε̇ = p_g/W, not ḣ = dh/dt. Kernel signatures are unchanged, but callers that constructed or consumed cell_velocities as a raw cell-vector velocity must update. See the fixes below for the physics.
  • npt_barostat_half_step{,_aniso,_triclinic} drop eta_dots (#77): barostat–thermostat (NHC) coupling is now applied as a separate Trotter operator by the caller, matching ASE and TorchSim.
  • Neighbors package restructured (#88, #103): nvalchemiops/neighbors/ is now split into per-strategy subpackages (naive/, cell_list/, cluster_tile/, rebuild/), with public launchers under */launchers.py and strategy selection under */dispatch.py. The flat naive_dual_cutoff, batch_naive, batch_cell_list, batch_naive_dual_cutoff, and rebuild_detection modules continue to re-export with a DeprecationWarning. The internal make_outer_neigh_offsets helper was removed.

Additions/New Features

Electrostatics

  • Higher-order (multipole) electrostatics (#95, #96): charges, dipoles, and quadrupoles (l = 0, 1, 2) via direct-k Ewald (multipole_ewald_summation) and particle-mesh Ewald (multipole_particle_mesh_ewald), with reciprocal/real-space entry points, feature extraction (multipole_electrostatic_features), and an SCF cache/step API. Single-system and batched Warp kernels and Torch bindings, with energies, forces, moment gradients, stress, and force-loss (create_graph) training; forward and first-order backward are torch.compile-compatible.
  • 2D slab (Yeh–Berkowitz) correction for Ewald and PME (#78, #86, #90): exposed as compute_slab_correction and a slab_correction= keyword on the Ewald/PME entry points, with both Torch and JAX bindings. The Torch path participates in autograd when inputs require gradients.
  • Energy-derived training derivatives (#96): full Torch Ewald/PME APIs support energy-derived forces, charge gradients, and strain-first virials, including second-order force/stress losses. Full JAX Ewald/PME energy-only calls support first-order gradients for positions, charges, and row-vector displacement virials.
  • PME precompute hooks: Torch/JAX PME accept precomputed cell_inv_t, volume, and B-spline moduli where supported; compute_bspline_moduli_1d is now exported from the top-level Torch and JAX electrostatics namespaces. Torch Ewald accepts miller_bounds for k-vector generation.

Neighbor Lists

  • Cluster-pair tile algorithm (#88, #107, #117): a new CUDA strategy for large, fully-periodic float32 systems. neighbor_list auto-selects it when eligible; pass method="cluster_tile" / "batch_cluster_tile" to force it. Supports dual cutoff in matrix format.
  • Inline pair potentials (#110, #118): neighbor kernels accept a user-supplied pair_fn callback (with pair_params/pair_energies/pair_forces buffers) that computes per-pair energy and force as pairs are enumerated, so Lennard-Jones–style potentials no longer require a separate pass over the list.
  • Per-pair vectors and distances on demand (#110): return_vectors and return_distances return r_ij and |r_ij| alongside the neighbor matrix.
  • Partial rebuild for batched workflows (#110): pass rebuild_flags to re-enumerate only the systems whose atoms moved enough to need a fresh list; unchanged systems keep their previous output. Supported for matrix and segmented-COO outputs in both JAX and PyTorch.
  • JAX CUDA graph replay (#72): JAX neighbor-list builders accept a graph_mode (GraphMode) keyword to capture and replay the build as a CUDA graph, reducing per-step launch overhead in MD loops.

Dynamics & Optimizers

  • Caller-supplied per-system reductions for FIRE/FIRE2 (#122): fire_step, fire_update, fire2_step, fire2_update, and the Torch fire2_step_coord{,_cell} accept compute_reductions=True; when False, values already in vf/vv/ff (and FIRE2 v_sumsq/f_sumsq) are reused. FIRE2 also exposes its phases (fire2_apply_step, Torch fire2_step_coord_cell_mix/_couple/_apply, and fire2_reduce/fire2_compute_extended_reductions). All default paths remain byte-identical.
  • Integrator flags for pre-computed thermodynamic quantities (#119): integrators can consume externally-computed kinetic/thermo reductions instead of recomputing them.

Segment Operations

  • Differentiable segment operations (#76): backward kernels enable autograd through nvalchemiops.segment_ops, with Torch and JAX bindings, an autograd example (examples/02_segment_ops_autograd.py), user-guide docs, and benchmarks.

Performance

  • DFT-D3 dispersion kernels optimized (#81).
  • Accelerated reciprocal-space Ewald fill (#102).
  • Coarsened pair-centric cell-list launches (#125).

Bug Fixes

  • MTK NPT/NPH cell propagation (#77): kernels wrote V·(P − P_ext)/W (strain-rate units) into cell_velocity while consumers read it as ḣ = dh/dt, costing a factor of cell length in the cell response. cell_velocity is now the strain rate ε̇ = p_g/W everywhere, and the update is h_new = h + dt · ε̇ · h.
  • MTK velocity-half-step coupling (#77): isotropic kernels used α = 1 + 1/(3N) instead of the canonical α = 1 + 1/N; anisotropic/triclinic kernels used (1 + 1/N)·ε̇, now replaced with ε̇ + Tr(ε̇)/(3N)·I.
  • MTK barostat half-step thermostat coupling (#77): removed the inline −η̇₁·ε̇ term that mixed the pressure/kinetic operator with NHC drag; callers now apply barostat-NHC coupling separately.
  • Cell-list size estimation overflow (#120): for a cell large relative to the cutoff, the cell-count product overflowed int32 and wrapped negative, crashing allocate_cell_list. The product is now computed overflow-safe (int64) and clamped per system, and both Torch and JAX allocate_cell_list validate the count and raise a clear error. (Fixes #74.)
  • Naive PBC wrapping on non-periodic axes (#106): position wrapping now leaves non-periodic axes unwrapped when per-axis pbc flags are supplied, fixing partial-PBC (slab) and fully non-periodic systems. (Fixes #104.)
  • estimate_max_neighbors floor (#120): exposes a max_neighbors_lower_bound keyword (default 16) so callers can raise the floor for dense/clustered systems with short cutoffs. (Addresses #114.)
  • PME fused-convolve backward rank under torch.compile: grad_k_squared was returned with the wrong rank for a single system; the k_squared unsqueeze is now tracked independently of the mesh and squeezed back on return, in both _pme_convolve_backward and _pme_convolve_double_backward.
  • Batched pressure kinetic tensor: compute_kinetic_tensor / compute_pressure_tensor with batch_idx now use a per-atom atomic reduction, fixing per-system corruption when a thread block spanned more than one system.
  • q(R) charge gradients (#116): correct handling of geometry-dependent charges in Ewald/PME derivatives.
  • DFT-D3 S5 switching applied to dE/dCN (#111).
  • Batched JAX Ewald autodiff memory (#121): no longer materializes a systems×k-vectors×atoms phase tensor, cutting reciprocal-space memory without changing energies or derivatives.
  • JAX importability on JAX 0.10+: migrated the Ewald/PME real/reciprocal/slab HVP transpose rules off the removed jax.custom_transpose to jax.custom_vjp, preserving second-order (force/stress-loss) derivatives.
  • Torch Ewald non-uniform cotangents: fixed gradients for non-uniform per-atom energy cotangents (torch.autograd.grad(..., grad_outputs=w)).
  • Coupled FIRE2 variable-cell updates (#75): positions and cell DOFs now advance consistently during variable-cell relaxation.
  • Neighbor-list dispatch validation (#97): launchers reject unbatched methods when batch metadata is supplied, instead of silently producing incorrect lists.

Deprecations

  • Direct-output flags on the full Torch/JAX Ewald/PME APIs for differentiable training: compute_forces, compute_virial, compute_charge_gradients, and hybrid_forces. They remain available with the existing tuple order; component charge-gradient, virial, and hybrid direct outputs now warn as legacy training-style outputs. Component compute_forces=True remains supported for no-autograd MD/inference.
  • estimate_max_neighbors(safety_factor=...) (#114): it scaled the estimate identically to atomic_density; it now emits a DeprecationWarning and folds into atomic_density.
  • nvalchemiops.neighbors.zero_array: emits a DeprecationWarning and forwards to array.zero_().
  • cells_inv and volumes arguments on the MTK NPT/NPH integrator entry points: kernels now consume cell_velocities as the strain rate directly; passing these emits a DeprecationWarning and they will be removed in a future release.

Compatibility & Tooling

  • Loosened the PyTorch version requirement (#93) and updated the CUDA backend extras (torch-cu12/jax-cu12) (#87).
  • New benchmark suite and 0.4 documentation (#123, #124).

New Contributors

Full Changelog: v0.3.1...v0.4.0

Contributors

samarjeet, lil-lon, and Ryan-Reese
Assets 4
Loading