$ show deepmodeling/DMFF
DMFF (Deep Modeling Force Field) is a JAX-based Python library for differentiable molecular force field parameterization and energy/force evaluation. It supports classical force fields (bonded, Lennard-Jones, Coulomb), advanced polarizable multipole models (ADMP with PME, Thole damping, Slater damping), and machine learning potentials (sGNN, EANN). The library provides a generator pattern for constructing differentiable energy functions from XML force field definitions, an operator pattern for atom typing and virtual site placement, and a parameter management system (ParamSet) that integrates with JAX's automatic differentiation for gradient-based optimization. It also includes an OpenMM plugin (C++/CUDA) that can load trained JAX models via TensorFlow for production molecular dynamics simulations. The architecture is modular: the core `dmff` package contains subpackages for ADMP (polarizable multipole), classical forces, sGNN/EANN, common utilities (neighbor lists, constants), and an API layer (Hamiltonian, topology, parameter sets). The backend includes a C++ OpenMM plugin with Reference and CUDA platforms, serialization support, and a TensorFlow export path via `jax2tf`. The library is designed for research in force field development, enabling parameter fitting, free energy calculations (MBAR), and integration with machine learning models. While primarily a software library, it has strong research characteristics: it is accompanied by a published paper, includes example workflows for water models and PEG systems, and provides reference data for validation. The codebase demonstrates high craftsmanship in areas like PME implementation, subgraph neural network construction, and differentiable programming patterns, though some parts (e.g., the OpenMM plugin's TensorFlow dependency) show potential drift risks.
The LLM identified 4 files worth highlighting for reuse:
Implements the ADMPPmeForce class for multipolar polarizable electrostatics with particle mesh Ewald (PME) summation, Thole damping, and polarization. The class provides a `get_energy` method that returns a JAX-differentiable energy function. It handles multipole moments up to quadrupole, reciprocal space B-spline interpolation, and self-energy corrections. This is a complex, self-contained implementation of a state-of-the-art electrostatic model, reusable for any polarizable force field requiring PME.
Provides three neighbor list implementations: `NeighborListFreud` (using the Freud library), `NoCutoffNeighborList` (full pair list), and `NoPeriodicNeighborList` (for non-periodic systems). Each implements a consistent interface with `allocate` and `update` methods, returning pair indices and shifts. The abstraction allows easy swapping of neighbor list strategies and is reusable across different force field modules.
Defines the `ParamSet` class that stores force field parameters as a dictionary of JAX arrays, registered as a JAX pytree for automatic differentiation. Supports parameter masking (trainable vs. fixed) via `update_mask`, enabling selective optimization. The class provides methods for serialization, updating from external sources, and integration with the generator pattern. This is a clean, reusable pattern for managing differentiable parameters in any JAX-based optimization workflow.
Implements `TopGraph` and `TopSubGraph` classes for building molecular topological graphs for subgraph neural network (sGNN) energy prediction. Handles atom typification, chirality labeling, canonical order generation via permutation groups, and feature extraction (internal coordinates, atomic numbers, bond types). The `get_canonical_orders_wt_permutation_grps` method efficiently enumerates symmetric atom orders to reduce feature redundancy. This is a sophisticated, self-contained component for graph-based molecular representation, reusable for any GNN-based force field.