Skip to content

kups.application.simulations.relax_lj

Lennard-Jones structure relaxation entry point.

Config

Bases: BaseModel

Top-level configuration for LJ relaxation.

Source code in src/kups/application/simulations/relax_lj.py 
62
63
64
65
66
67
68
class Config(BaseModel):
    """Top-level configuration for LJ relaxation."""

    run: RelaxRunConfig
    relax: RelaxParameters
    lj: LjConfig
    inp_file: str | Path

LjConfig

Bases: BaseModel

Lennard-Jones potential parameters.

Source code in src/kups/application/simulations/relax_lj.py 
53
54
55
56
57
58
59
class LjConfig(BaseModel):
    """Lennard-Jones potential parameters."""

    tail_correction: bool
    cutoff: float
    parameters: dict[str, tuple[float | None, float | None]]
    mixing_rule: MixingRule

RelaxLjState

Simulation state for Lennard-Jones relaxation.

Source code in src/kups/application/simulations/relax_lj.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
@dataclass
class RelaxLjState:
    """Simulation state for Lennard-Jones relaxation."""

    particles: Table[ParticleId, RelaxParticles]
    systems: Table[SystemId, RelaxSystems]
    neighborlist_params: UniversalNeighborlistParameters
    opt_state: optax.OptState
    step: Array
    lj_parameters: LennardJonesParameters

    @property
    def neighborlist(self) -> NearestNeighborList:
        return DenseNearestNeighborList.from_state(self)

init_state(config, opt_init)

Initialise an LJ relaxation state from configuration.

Parameters:

Name Type Description Default
config Config

Simulation configuration.

 required
opt_init OptInit

Optimizer state initialiser from the relaxation propagator.

 required

Returns:

Type Description
RelaxLjState

Fully constructed LJ relaxation state.
Source code in src/kups/application/simulations/relax_lj.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def init_state(config: Config, opt_init: OptInit) -> RelaxLjState:
    """Initialise an LJ relaxation state from configuration.

    Args:
        config: Simulation configuration.
        opt_init: Optimizer state initialiser from the relaxation propagator.

    Returns:
        Fully constructed LJ relaxation state.
    """
    lj_params = LennardJonesParameters.from_dict(
        cutoff=config.lj.cutoff,
        parameters=config.lj.parameters,
        mixing_rule=config.lj.mixing_rule,
    )
    particles, systems = relax_state_from_ase(config.inp_file)

    neighborlist_params = UniversalNeighborlistParameters.estimate(
        particles.data.system.counts, systems, lj_params.cutoff
    )
    opt_state = opt_init(
        (particles.data.positions, systems.data.unitcell.lattice_vectors)
    )
    return RelaxLjState(
        particles=particles,
        systems=systems,
        neighborlist_params=neighborlist_params,
        opt_state=opt_state,
        step=jnp.array([0]),
        lj_parameters=lj_params,
    )

main()

CLI entry point for LJ relaxation.

Source code in src/kups/application/simulations/relax_lj.py 
140
141
142
143
144
145
def main() -> None:
    """CLI entry point for LJ relaxation."""
    cli = NanoArgs(Config)
    config = cli.parse()
    run(config)
    rich.print(analyze_relax_file(config.run.out_file))

run(config)

Run an LJ structure relaxation from the given configuration.

Parameters:

Name Type Description Default
config Config

Simulation configuration.

 required
Source code in src/kups/application/simulations/relax_lj.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def run(config: Config) -> None:
    """Run an LJ structure relaxation from the given configuration.

    Args:
        config: Simulation configuration.
    """
    key = jax.random.key(config.run.seed or time.time_ns())
    state_lens = identity_lens(RelaxLjState)
    optimizer = make_optimizer(config.relax.optimizer)
    potential = make_lennard_jones_from_state(
        state_lens, compute_position_and_unitcell_gradients=True
    )
    propagator, opt_init = make_relax_propagator(
        state_lens, potential, optimizer, config.relax.optimize_cell
    )
    state = init_state(config, opt_init)
    logging.info("Starting relaxation")
    run_relax(key, propagator, state, config.run)