kups.application.mcmc ¶

@no_jax_tracing
def analyze_mcmc(
    fixed: IsMCMCFixedData,
    per_step: IsMCMCStepData,
    n_blocks: int | None = None,
) -> dict[SystemId, MCMCAnalysisResult]:
    """Analyze MCMC simulation results from pre-loaded data.

    Computes energy, loading (average particle counts), and heat of adsorption
    per species using block averaging for error estimation. Analysis is
    performed independently for each system.

    Args:
        fixed: Fixed (one-shot) logged data containing system metadata.
        per_step: Per-step logged data containing energies and particle counts.
        n_blocks: Number of blocks for error estimation. ``None`` uses
            :func:`~kups.core.utils.block_average.optimal_block_average`.

    Returns:
        Per-system analysis results keyed by ``SystemId``.
    """
    system_keys = fixed.systems.keys
    count_keys = per_step.particle_count.keys
    count_data = per_step.particle_count.data
    temperatures = fixed.systems.data.temperature
    all_energy = per_step.systems.data.potential_energy
    guest_stress = per_step.systems.data.guest_stress
    stress_potential = guest_stress.potential
    stress_tail = guest_stress.tail_correction
    stress_ideal = guest_stress.ideal_gas
    all_stress = stress_potential + stress_tail + stress_ideal

    results: dict[SystemId, MCMCAnalysisResult] = {}
    for i, sys_id in enumerate(system_keys):
        motif_cols = [j for j, label in enumerate(count_keys) if label[0] == sys_id]
        counts = count_data[:, motif_cols].reshape(-1, len(motif_cols))
        energy = all_energy[:, i].reshape(-1)
        s = all_stress[:, i]
        has_stress = bool(jnp.any(s != 0))
        temperature = float(temperatures[i])
        results[sys_id] = _analyze_single_system(
            energy,
            counts,
            temperature,
            n_blocks,
            stress=s if has_stress else None,
            stress_potential=stress_potential[:, i]
            if has_stress and stress_potential is not None
            else None,
            stress_tail_correction=stress_tail[:, i]
            if has_stress and stress_tail is not None
            else None,
            stress_ideal_gas=stress_ideal[:, i]
            if has_stress and stress_ideal is not None
            else None,
        )

    return results

@no_jax_tracing
def analyze_mcmc_file(
    hdf5_path: str | Path,
    n_blocks: int | None = None,
) -> dict[SystemId, MCMCAnalysisResult]:
    """Analyze MCMC simulation results from an HDF5 file.

    Convenience wrapper that reads the HDF5 file and delegates to
    :func:`analyze_mcmc`.

    Args:
        hdf5_path: Path to HDF5 output file from
            :func:`~kups.application.mcmc.simulation.run_mcmc`.
        n_blocks: Number of blocks for error estimation. ``None`` uses
            :func:`~kups.core.utils.block_average.optimal_block_average`.

    Returns:
        Per-system analysis results keyed by ``SystemId``.
    """
    with HDF5StorageReader[MCMCLoggedData](hdf5_path) as reader:
        fixed = reader.focus_group(lambda s: s.fixed)[...]
        per_step = reader.focus_group(lambda s: s.per_step)[...]

    return analyze_mcmc(fixed, per_step, n_blocks)

def make_mcmc_logged_data[S: IsMCMCState](
    state: S,
    stress_fn: StateProperty[S, Table[SystemId, StressResult]] | None = None,
) -> MCMCLoggedData:
    """Create MCMC logging config that logs host particles once and adsorbate buffer per step.

    Host particles are identified by having out-of-bounds group indices
    (they belong to no molecular group). This works for any number of
    host systems regardless of how particles are interleaved in the buffer.

    Args:
        state: The initial MCMC state (after padding).
        stress_fn: Optional stress tensor function. When provided, the
            guest stress tensor is logged each cycle.

    Returns:
        Configured logging data with host/adsorbate split.
    """

    particles = state.particles
    is_occupied = np.asarray(particles.occupation)
    is_host = ~particles.data.group.valid_mask & is_occupied

    host_positions = np.where(is_host)[0]
    ads_positions = np.where(~is_host)[0]

    host_idx = Index(particles.keys, jnp.asarray(host_positions))
    ads_idx = Index(particles.keys, jnp.asarray(ads_positions))
    host_keys = tuple(particles.keys[int(i)] for i in host_positions)
    ads_keys = tuple(particles.keys[int(i)] for i in ads_positions)

    def fixed_view(state: S) -> MCMCFixedData:
        return MCMCFixedData(
            systems=state.systems,
            motifs=state.motifs,
            host_particles=Table(host_keys, state.particles[host_idx]),
        )

    def step_view(state: S) -> MCMCStepData:
        counts = Index.combine(state.groups.data.system, state.groups.data.motif).counts
        ewald = state.ewald_parameters.cache
        if stress_fn is not None:
            guest_stress = stress_fn(jax.random.key(0), state)
        else:
            z = jnp.zeros((len(state.systems), 3, 3))
            guest_stress = Table(state.systems.keys, StressResult(z, z, z))
        sys_keys = state.systems.keys
        system_step = MCMCSystemStepData(
            potential_energy=state.systems.data.potential_energy,
            lj_energy=state.lj_parameters.cache.total_energies.data,
            ewald_short_range_energy=ewald.short_range.total_energies.data,
            ewald_long_range_energy=ewald.long_range.total_energies.data,
            ewald_self_energy=ewald.self_interaction.total_energies.data,
            ewald_exclusion_energy=ewald.exclusion.total_energies.data,
            translation_step_width=state.translation_params.map_data(
                lambda p: p.value
            ).data,
            rotation_step_width=state.rotation_params.map_data(lambda p: p.value).data,
            translation_acceptance=state.translation_params.map_data(
                lambda p: p.history.values.mean(axis=-1)
            ).data,
            rotation_acceptance=state.rotation_params.map_data(
                lambda p: p.history.values.mean(axis=-1)
            ).data,
            reinsertion_acceptance=state.reinsertion_params.map_data(
                lambda p: p.history.values.mean(axis=-1)
            ).data,
            exchange_acceptance=state.exchange_params.map_data(
                lambda p: p.history.values.mean(axis=-1)
            ).data,
            guest_stress=guest_stress.data,
        )
        return MCMCStepData(
            particles=Table(ads_keys, state.particles[ads_idx]),
            groups=state.groups,
            particle_count=counts,
            systems=Table(sys_keys, system_step),
        )

    return MCMCLoggedData(
        fixed=WriterGroupConfig(fixed_view, Once()),
        per_step=WriterGroupConfig(step_view, EveryNStep(1)),
    )

def run_mcmc[State](
    key: Array,
    propagator: Propagator[State],
    state: State,
    config: RunConfig,
    logged_data: MCMCLoggedData,
) -> State:
    """Run a µVT MCMC simulation with warmup and production phases.

    Args:
        key: JAX PRNG key.
        propagator: Propagator, e.g. from :func:`~kups.application.simulations.mcmc_rigid.make_propagator`.
        state: Initial simulation state.
        config: Run configuration.
        logged_data: Logging configuration with host/adsorbate split.

    Returns:
        Final simulation state after production run.
    """
    chain = key_chain(key)
    logging.info("Warming up (%d cycles)...", config.num_warmup_cycles)
    state = run_warmup_cycles(next(chain), propagator, state, config.num_warmup_cycles)
    logging.info("Production run (%d cycles)...", config.num_cycles)
    logger = CompositeLogger(
        HDF5StorageWriter(config.out_file, logged_data, state, config.num_cycles),
        TqdmLogger(config.num_cycles),
    )
    state = run_simulation_cycles(
        next(chain), propagator, state, config.num_cycles, logger
    )
    logging.info("Done.")
    return state