kups.application.md.data ¶

@dataclass
class BAOABLangevinParams:
    """Control parameters for the BAOAB Langevin (NVT) integrator.

    Attributes:
        time_step: Integration timestep ``Δt``, shape ``(n_systems,)``.
        temperature: Target temperature ``T`` (K), shape ``(n_systems,)``.
        friction_coefficient: Langevin friction ``γ`` (1/time), shape ``(n_systems,)``.
    """

    time_step: Array
    temperature: Array
    friction_coefficient: Array

@dataclass
class BAOABNPTLangevinParams:
    r"""Control parameters for the BAOAB NPT Langevin integrator.

    Implements the fully-flexible-cell extended-variable NPT Langevin
    formulation of Gao, Fang & Wang, *Sampling the isothermal-isobaric
    ensemble by Langevin dynamics*, JCP 2016 (arxiv 1601.01044). The atom
    side reuses the existing Langevin friction ``γ``; the cell side adds a
    fictitious mass tensor and a separate Langevin friction.

    Attributes:
        time_step: Integration timestep ``Δt``, shape ``(n_systems,)``.
        temperature: Target temperature ``T`` (K), shape ``(n_systems,)``.
        friction_coefficient: Atom Langevin friction ``γ`` (1/time),
            shape ``(n_systems,)``. Reused from BAOAB-Langevin.
        target_pressure: Target pressure ``P₀`` (energy/length³),
            shape ``(n_systems,)``.
        pressure_coupling_time: Barostat coupling time ``τ_P`` (time),
            shape ``(n_systems,)``. Drives the auto-derived
            :attr:`barostat_mass` per Gao Eq. (14)/(15).
        compressibility: Isothermal compressibility ``β`` (length³/energy),
            shape ``(n_systems,)``. Drives :attr:`barostat_mass`.
        barostat_mass: Fictitious cell mass tensor $M_{\alpha\beta}$
            (mass·length²), lower-triangular, shape ``(n_systems, 3, 3)``.
            Auto-derived from ``compressibility``, ``pressure_coupling_time``
            and the initial cell at construction time.
        barostat_friction: Langevin friction on cell DOFs $\gamma_{\alpha\beta}$
            (1/time), lower-triangular, shape ``(n_systems, 3, 3)``. Defaults
            to ``friction_coefficient`` broadcast across the 6 cell DOFs per
            Gao Remark 3.
    """

    time_step: Array
    temperature: Array
    friction_coefficient: Array
    target_pressure: Array
    pressure_coupling_time: Array
    compressibility: Array
    barostat_mass: Array
    barostat_friction: Array

@dataclass
class CSVRNPTParams:
    """Control parameters for the CSVR-NPT integrator.

    Attributes:
        time_step: Integration timestep ``Δt``, shape ``(n_systems,)``.
        temperature: Target temperature ``T`` (K), shape ``(n_systems,)``.
        thermostat_time_constant: CSVR coupling time ``τ`` (time), shape ``(n_systems,)``.
        target_pressure: Target pressure ``P₀`` (energy/length³), shape ``(n_systems,)``.
        pressure_coupling_time: Barostat coupling time ``τ_P`` (time), shape ``(n_systems,)``.
        compressibility: Isothermal compressibility ``β`` (length³/energy), shape ``(n_systems,)``.
        minimum_scale_factor: Minimum barostat scale factor, shape ``(n_systems,)``.
    """

    time_step: Array
    temperature: Array
    thermostat_time_constant: Array
    target_pressure: Array
    pressure_coupling_time: Array
    compressibility: Array
    minimum_scale_factor: Array

@dataclass
class CSVRParams:
    """Control parameters for the CSVR (NVT) integrator.

    Attributes:
        time_step: Integration timestep ``Δt``, shape ``(n_systems,)``.
        temperature: Target temperature ``T`` (K), shape ``(n_systems,)``.
        thermostat_time_constant: CSVR coupling time ``τ`` (time), shape ``(n_systems,)``.
    """

    time_step: Array
    temperature: Array
    thermostat_time_constant: Array

@dataclass
class MDParticles(Particles):
    """Particle state for molecular dynamics simulations.

    Extends :class:`Particles` with gradient, momenta, and derived
    kinematic quantities needed by MD integrators.

    Attributes:
        position_gradients: Energy gradient w.r.t. positions, shape ``(n_atoms, 3)``.
        momenta: Particle momenta, shape ``(n_atoms, 3)``.
        exclusion: Per-particle exclusion index (defaults to one group per
            atom via :func:`default_exclusion` if not supplied).
    """

    position_gradients: Array
    momenta: Array
    exclusion: Index[ExclusionId] = field(default=None, kw_only=True)  # type: ignore

    def __post_init__(self) -> None:
        if self.exclusion is None:
            object.__setattr__(self, "exclusion", default_exclusion(len(self.charges)))

    @property
    def forces(self) -> Array:
        """Negative position gradient, shape ``(n_atoms, 3)``."""
        return -self.position_gradients

    @property
    def velocities(self) -> Array:
        """Velocities derived from momenta and masses, shape ``(n_atoms, 3)``."""
        return self.momenta / self.masses[..., None]

    @property
    def kinetic_energy(self) -> Array:
        """Per-particle kinetic energy, shape ``(n_atoms,)``."""
        return particle_kinetic_energy(self.momenta, self.masses)

@dataclass
class MDSystems:
    r"""Per-system state for molecular dynamics simulations.

    Attributes:
        cell: Cell geometry for each system.
        integrator_params: Bundled integrator control parameters; concrete shape
            (e.g. :class:`VerletParams`, :class:`BAOABLangevinParams`,
            :class:`CSVRParams`, :class:`CSVRNPTParams`,
            :class:`BAOABNPTLangevinParams`) is chosen to match the selected
            integrator.
        cell_gradients: Energy gradient w.r.t. the cell, stored as a
            :class:`Cell` (the ``vectors`` leaf holds the
            shape-``(n_systems, 3, 3)`` gradient used by
            :attr:`stress_tensor`).
        cell_momentum: Extended-variable cell-momentum tensor $p^h$,
            lower-triangular ``(n_systems, 3, 3)``. Only meaningful for
            integrators that drive cell dynamics via an explicit conjugate
            momentum (e.g. BAOAB NPT Langevin); zero otherwise.
        potential_energy: Total potential energy per system (eV), shape ``(n_systems,)``.
    """

    cell: Cell[Any]
    integrator_params: IntegratorParams
    cell_gradients: Cell[Any]
    cell_momentum: Array
    potential_energy: Array

class MdParameters(BaseModel):
    """Physical and numerical parameters for an MD simulation."""

    temperature: float
    """Target temperature (K)."""
    time_step: float
    """Integration timestep (fs)."""
    friction_coefficient: float
    """Langevin friction coefficient (1/fs)."""
    thermostat_time_constant: float
    """CSVR thermostat coupling time (fs)."""
    target_pressure: float
    """Target pressure for NPT barostat (Pa)."""
    pressure_coupling_time: float
    """Barostat coupling time (fs)."""
    compressibility: float
    """Isothermal compressibility (1/Pa)."""
    minimum_scale_factor: float
    """Minimum allowed box scaling factor per barostat step (dimensionless)."""
    integrator: Integrator
    """Integration algorithm to use."""
    initialize_momenta: bool = False
    """If True, initialize momenta from Maxwell-Boltzmann distribution."""

class MdRunConfig(BaseModel):
    """Run configuration for an MD simulation."""

    out_file: str | Path
    """Path to the output HDF5 file."""
    num_steps: int
    """Number of production steps."""
    num_warmup_steps: int
    """Number of warmup steps before production."""
    block_size: int = 1
    """Production steps fused per on-device dispatch (1 = per-step). Only the last frame of
    each block is saved, so the trajectory has num_steps // block_size frames."""
    seed: int | None
    """Random seed for reproducibility. None for time-based."""

@dataclass
class MdState:
    """Force-field-agnostic molecular-dynamics state.

    The potential is built with its parameters at construction time (via the
    adapters' ``parameters=``), so no force-field field lives on the state.
    """

    particles: Table[ParticleId, MDParticles]
    systems: Table[SystemId, MDSystems]
    neighborlist_params: UniversalNeighborlistParameters
    step: Array

@dataclass
class VerletParams:
    """Control parameters for the NVE Velocity Verlet integrator.

    Attributes:
        time_step: Integration timestep ``Δt`` (internal time units), shape ``(n_systems,)``.
    """

    time_step: Array