kups.application.mcmc.data ¶

class AdsorbateConfig(BaseModel):
    """Configuration for a single adsorbate species.

    Thermodynamic parameters feed Peng-Robinson EOS for fugacity calculation.
    Molecular geometry is defined by ``positions`` and ``symbols``; charges,
    masses, and atomic numbers are auto-derived from ASE if omitted.
    """

    critical_temperature: float
    """Critical temperature (K) for Peng-Robinson EOS."""
    critical_pressure: float
    """Critical pressure (Pa) for Peng-Robinson EOS."""
    acentric_factor: float
    """Acentric factor (dimensionless)."""
    positions: tuple[tuple[float, float, float], ...]
    """Atom positions within the molecule, relative to its centre of mass (Ang)."""
    symbols: tuple[str, ...]
    """Element symbols for each atom in the molecule."""
    charges: tuple[float, ...] = ()
    """Partial charges on each atom (e). Defaults to zeros."""
    masses: tuple[float, ...] = ()
    """Atomic masses (amu). Derived from symbols via ASE if omitted."""
    atomic_numbers: tuple[int, ...] = ()
    """Atomic numbers. Derived from symbols via ASE if omitted."""

    @model_validator(mode="before")
    @classmethod
    def _fill_defaults_from_ase(cls, data: dict) -> dict:  # type: ignore[override]
        """Derive charges, masses, and atomic numbers from symbols when absent."""
        symbols = data.get("symbols", ())
        if not data.get("charges"):
            data["charges"] = tuple(0.0 for _ in symbols)
        if not data.get("masses") or not data.get("atomic_numbers"):
            zs = [ase.data.atomic_numbers.get(s.split("_")[0], 0) for s in symbols]
            if not data.get("masses"):
                data["masses"] = tuple(
                    float(ase.data.atomic_masses[z]) if z > 0 else 0.0 for z in zs
                )
            if not data.get("atomic_numbers"):
                data["atomic_numbers"] = tuple(zs)
        return data

    @property
    def as_particles(self) -> Table[MotifParticleId, MotifParticles]:
        """Convert this adsorbate config into indexed particles with motif data."""
        n = len(self.symbols)
        assert self.masses is not None
        assert self.atomic_numbers is not None
        assert self.charges is not None
        return Table.arange(
            MotifParticles(
                positions=jnp.asarray(self.positions),
                masses=jnp.asarray(self.masses),
                atomic_numbers=jnp.asarray(self.atomic_numbers),
                charges=jnp.asarray(self.charges),
                labels=Index.new(list(map(Label, self.symbols))),
                motif=Index.integer(
                    jnp.zeros(n, dtype=int),
                    label=MotifId,
                    max_count=len(self.positions),
                ),
            ),
            label=MotifParticleId,
        )

class EwaldCfg(BaseModel):
    """Ewald summation settings for electrostatic interactions."""

    enabled: bool = True
    cutoff: float = 12.0
    """Real-space cutoff (Ang)."""
    precision: float = 1e-6
    """Target error tolerance."""

class HostConfig(BaseModel):
    """Configuration for a single host framework (e.g. zeolite, MOF).

    Each host gets its own system with unit cell, thermodynamic conditions,
    and optional initial adsorbate placement.
    """

    cif_file: str
    """Path to the host structure CIF file."""
    pressure: float
    """System pressure (Pa)."""
    temperature: float
    """System temperature (K)."""
    init_adsorbates: tuple[int, ...] = (0,)
    """Number of initial adsorbates per species to randomly insert."""
    adsorbate_composition: tuple[float, ...] = (1.0,)
    """Mole fractions for each adsorbate component (must sum to 1)."""
    adsorbate_interaction: tuple[tuple[float, ...], ...] = ((0.0,),)
    """Binary interaction parameters k_ij (n_ads x n_ads)."""
    unitcell_replication: int | tuple[int, int, int] | None = None
    """(nx, ny, nz) replication of the unit cell. Auto-computed if None."""

class LJConfig(BaseModel):
    """Lennard-Jones force-field parameters."""

    parameters: dict[str, tuple[float, float]]
    """Mapping symbol -> (epsilon/kB [K], sigma [Ang]) for each element type."""
    cutoff: float = 12.0
    """Real-space cutoff (Ang)."""
    tail_correction: bool = True
    """Apply long-range tail correction."""

@dataclass
class MCMCGroup:
    """Molecular group state for rigid-body MCMC.

    Attributes:
        system: Index mapping each group to its parent system.
        motif: Index mapping each group to its adsorbate species.
    """

    system: Index[SystemId]
    motif: Index[MotifId]

    @property
    def labels(self) -> Index[Label]:
        """Group labels derived from motif index."""
        return Index(tuple(Label(str(m)) for m in self.motif.keys), self.motif.indices)

@dataclass
class MCMCParticles(Particles):
    """Particle state for rigid-body MCMC with group membership.

    Attributes:
        group: Index mapping each particle to its molecular group.
    """

    group: Index[GroupId]
    motif: Index[MotifParticleId]
    position_gradients: Array = field(default=None)  # type: ignore[assignment]

    def __post_init__(self):
        if self.position_gradients is None:
            object.__setattr__(
                self, "position_gradients", jnp.zeros_like(self.positions)
            )

    @property
    def inclusion(self) -> Index[InclusionId]:
        """Inclusion index derived from the system index."""
        return Index(
            tuple(InclusionId(lab) for lab in self.system.keys),
            self.system.indices,
            self.system.max_count,
        )

    @property
    def exclusion(self) -> Index[ExclusionId]:
        """Exclusion index derived from the group index."""
        return Index(
            tuple(ExclusionId(lab) for lab in self.group.keys),
            self.group.indices,
            self.group.max_count,
        )

    def guest_only(self) -> MCMCParticles:
        # Buffered sanitizes the the data to only include particles with valid group membership.
        return Buffered.arange(self, view=lambda x: x.group).data

@dataclass
class MCMCSystems:
    """Per-system thermodynamic state for MCMC simulations.

    Attributes:
        unitcell: Unit cell geometry, batched shape ``(n_systems,)``.
        temperature: Temperature (K), shape ``(n_systems,)``.
        potential_energy: Total potential energy per system (eV), shape ``(n_systems,)``.
        log_fugacity: Log fugacity per species (dimensionless), shape ``(n_systems, n_species)``.
    """

    unitcell: UnitCell
    temperature: Array
    potential_energy: Array
    log_fugacity: Array
    unitcell_gradients: UnitCell = field(default=None)  # type: ignore[assignment]

    def __post_init__(self):
        if self.unitcell_gradients is None:
            object.__setattr__(
                self, "unitcell_gradients", tree_zeros_like(self.unitcell)
            )

    @property
    def log_activity(self) -> Array:
        """Log activity: log(fugacity / kT), shape ``(n_systems, n_species)``."""
        return (
            self.log_fugacity - jnp.log(self.temperature * BOLTZMANN_CONSTANT)[:, None]
        )

@dataclass
class MotifParticles:
    """Template particle data for an adsorbate motif (molecule template).

    Stores the reference geometry, atomic properties, and species assignment
    for each atom in an adsorbate molecule. Used as a lookup table during
    exchange moves to populate new particle data.

    Attributes:
        positions: Reference atom positions relative to COM (Ang), shape ``(n_atoms, 3)``.
        masses: Atomic masses (amu), shape ``(n_atoms,)``.
        atomic_numbers: Atomic numbers, shape ``(n_atoms,)``.
        charges: Partial charges (e), shape ``(n_atoms,)``.
        labels: Per-atom string labels.
        motif: Index mapping atoms to their species (MotifId).
    """

    positions: Array
    masses: Array
    atomic_numbers: Array
    charges: Array
    labels: Index[Label]
    motif: Index[MotifId]

class RunConfig(BaseModel):
    """Run configuration for MCMC simulation."""

    out_file: str | Path
    """Path to the output HDF5 file."""
    num_cycles: int
    """Number of production cycles."""
    num_warmup_cycles: int
    """Number of warmup cycles (not written to output)."""
    min_cycle_length: int
    translation_prob: float = 1 / 6
    rotation_prob: float = 1 / 6
    reinsertion_prob: float = 1 / 6
    exchange_prob: float = 1 / 2
    """Ordered move sequence for the palindrome propagator."""
    seed: int | None = None
    """Random seed. ``None`` uses a time-based seed."""

class StressResult(NamedTuple):
    """Decomposed stress tensor with potential, tail correction, and ideal gas components."""

    potential: Array
    tail_correction: Array
    ideal_gas: Array