Skip to content

kups.application.simulations.md

Unified molecular-dynamics entry point.

The force field is selected by the potential config (LJ / tojax-MLFF / torch-MLFF). The potential is built directly with its parameters, so a single force-field-agnostic MdState and driver serve every backend.

Config

Bases: BaseModel

Top-level configuration for a molecular-dynamics run.

Source code in src/kups/application/simulations/md.py 
61
62
63
64
65
66
67
class Config(BaseModel):
    """Top-level configuration for a molecular-dynamics run."""

    run: MdRunConfig
    md: MdParameters
    potential: PotentialConfig
    inp_files: tuple[str | Path, ...]

main()

CLI entry point for molecular-dynamics simulations.

Source code in src/kups/application/simulations/md.py 
102
103
104
105
106
107
108
def main() -> None:
    """CLI entry point for molecular-dynamics simulations."""
    cli = NanoArgs(Config)
    config = cli.parse()
    rich.print(config)
    run(config)
    rich.print(analyze_md_file(config.run.out_file))

run(config)

Run a molecular-dynamics simulation for the configured force field.

Source code in src/kups/application/simulations/md.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def run(config: Config) -> None:
    """Run a molecular-dynamics simulation for the configured force field."""
    seed = config.run.seed or time.time_ns()
    chain = key_chain(jax.random.key(seed))
    state_lens = identity_lens(MdState)
    potential, cutoff = config.potential.build(state_lens, POSITIONS_AND_CELL)
    propagator = make_md_propagator(state_lens, config.md.integrator, potential)

    mb_key = next(chain) if config.md.initialize_momenta else None
    all_particles: list[Table[ParticleId, MDParticles]] = []
    all_systems: list[Table[SystemId, MDSystems]] = []
    for inp_file in config.inp_files:
        logging.info(f"Loading structure from {inp_file}")
        particles_i, systems_i = md_state_from_ase(inp_file, config.md, key=mb_key)
        all_particles.append(particles_i)
        all_systems.append(systems_i)
    particles, systems = Table.union(all_particles, all_systems)

    base = 1 if isinstance(config.potential, TojaxPotentialConfig) else 2
    multiplier = 2.0 if isinstance(config.potential, MaceConfig | UmaConfig) else 1.0
    neighborlist_params = UniversalNeighborlistParameters.estimate(
        particles.data.system.counts, systems, cutoff, base=base, multiplier=multiplier
    )
    state = MdState(
        particles=particles,
        systems=systems,
        neighborlist_params=neighborlist_params,
        step=jnp.array([0]),
    )
    run_md(next(chain), propagator, state, config.run)