Skip to content

kups.application.simulations.relax

Unified structure-relaxation 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 RelaxState and driver serve every backend.

Config

Bases: BaseModel

Top-level configuration for a structure relaxation run.

Source code in src/kups/application/simulations/relax.py 
59
60
61
62
63
64
class Config(BaseModel):
    """Top-level configuration for a structure relaxation run."""

    run: RelaxRunConfig
    potential: PotentialConfig
    inp_files: tuple[str | Path, ...]

main()

CLI entry point for structure relaxation.

Source code in src/kups/application/simulations/relax.py 
104
105
106
107
108
109
110
def main() -> None:
    """CLI entry point for structure relaxation."""
    cli = NanoArgs(Config)
    config = cli.parse()
    rich.print(config)
    run(config)
    rich.print(analyze_relax_file(config.run.out_file))

run(config)

Run a structure relaxation for the configured force field.

Source code in src/kups/application/simulations/relax.py 
 67
 68
 69
 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
100
101
def run(config: Config) -> None:
    """Run a structure relaxation for the configured force field."""
    key = jax.random.key(config.run.seed or time.time_ns())
    state_lens = identity_lens(RelaxState)
    optimizer = make_optimizer(config.run.optimizer)
    gradient = FRECHET_FILTER if config.run.optimize_cell else POSITIONS_ONLY
    potential, cutoff = config.potential.build(state_lens, gradient)
    propagator, opt_init = make_relax_propagator(
        state_lens, potential, optimizer, gradient
    )

    all_particles: list[Table[ParticleId, RelaxParticles]] = []
    all_systems: list[Table[SystemId, RelaxSystems]] = []
    for inp_file in config.inp_files:
        logging.info(f"Loading structure from {inp_file}")
        particles_i, systems_i = relax_state_from_ase(inp_file)
        all_particles.append(particles_i)
        all_systems.append(systems_i)
    particles, systems = Table.union(all_particles, all_systems)

    # Torch MLFF models need extra neighbor-list capacity headroom.
    multiplier = 2.0 if isinstance(config.potential, MaceConfig | UmaConfig) else 1.0
    neighborlist_params = UniversalNeighborlistParameters.estimate(
        particles.data.system.counts, systems, cutoff, multiplier=multiplier
    )
    opt_state = opt_init(particles, systems)
    state = RelaxState(
        particles=particles,
        systems=systems,
        neighborlist_params=neighborlist_params,
        opt_state=opt_state,
        step=jnp.array([0]),
    )
    logging.info("Starting relaxation")
    run_relax(key, propagator, state, config.run)