kups.application.simulations.mcmc_rigid ¶

def init_cell_tables(
    particles: Table[ParticleId, MCMCParticles],
    motifs: Table[MotifParticleId, MotifParticles],
    systems: Table[SystemId, MCMCSystems],
    lj_parameters: GlobalTailCorrectedLennardJonesParameters,
    ewald_parameters: EwaldParameters,
    dtype: DTypeLike = jnp.float64,
) -> FusedPotentialCache[tuple[Index[Label] | Array, ...]]:
    """Initialize the shared pair cache from the simulation's input data."""
    lj: PairEnergy[
        tuple[GlobalTailCorrectedLennardJonesParameters, EwaldParameters],
        MCMCParticles,
        Index[Label],
    ] = lennard_jones_pair.with_particles(lambda p: p).with_parameters(lambda p: p[0])
    parameters = (lj_parameters, ewald_parameters)
    cloud = PointCloud(particles, systems)
    if _is_charged(particles, motifs):
        return FusedPotentialCache.create(
            lj
            + ewald_short_range_pair.with_particles(lambda p: p).with_parameters(
                lambda p: p[1]
            ),
            parameters,
            cloud,
            dtype=dtype,
        )
    return FusedPotentialCache.create(
        PairEnergySum.from_term(lj), parameters, cloud, dtype=dtype
    )

def init_state(key: Array, config: Config) -> MCMCState:
    """Initialize the full MCMC state from configuration."""
    chain = key_chain(key)
    host_states = [
        mcmc_state_from_config(next(chain), host, config.adsorbates)
        for host in config.hosts
    ]
    if not host_states:
        raise ValueError("At least one host must be provided.")
    motifs = host_states[-1][3]
    particles, groups, system = Table.union(
        [s[0] for s in host_states],
        [s[1] for s in host_states],
        [s[2] for s in host_states],
    )
    logging.info(
        f"Initialized state with {len(particles)} particles, "
        f"{len(groups)} molecules, across {len(system)} systems."
    )
    max_adsorbates = estimate_max_adsorbates(particles, motifs, system)
    if config.run.exchange_prob <= 0:
        # Without insertions no free slots are needed, except one to keep an
        # empty group buffer non-empty.
        free = jnp.zeros_like(max_adsorbates.data)
        if len(groups) == 0:
            free = free.at[0].set(1)
        max_adsorbates = max_adsorbates.set_data(free)
    n_sys = len(system)
    lj_params = config.lj.make_parameters(
        particles.data.labels.keys + motifs.data.labels.keys
    )
    blocking_spheres = BlockingSpheresParameters.from_data(
        [host.blocking_spheres for host in config.hosts]
    )
    max_motif_size = motifs.data.motif.max_count
    assert max_motif_size is not None
    particles, groups, motifs, system = unify_keys_by_cls(
        (particles, groups, motifs, system)
    )
    num_buffer_particles = max_adsorbates * max_motif_size
    particles, groups = add_buffers(
        (particles, int(num_buffer_particles.data.sum())),
        (groups, int(max_adsorbates.data.sum())),
    )

    ewald_params = EwaldParameters.make(
        particles,
        system,
        epsilon_total=config.ewald.precision,
        real_cutoff=config.ewald.real_cutoff,
        reciprocal_dtype=config.kernel_dtype,
    )
    n_kvecs = ewald_params.reciprocal_lattice_shifts.data.shape[1]

    neighborlist_params = UniversalNeighborlistParameters.estimate(
        particles.data.system.counts + num_buffer_particles,
        system,
        tree_map(jnp.maximum, lj_params.cutoff, ewald_params.cutoff),
    )
    if blocking_spheres.radii.shape[0] > 0:
        # Systems without spheres have no radius to size from, so use the batch-wide max.
        max_radius = Table((SystemId(0),), blocking_spheres.radii.max(keepdims=True))
        blocking_nlist = UniversalNeighborlistParameters.estimate(
            particles.data.system.counts + num_buffer_particles,
            system,
            Table.broadcast_to(max_radius, system),
        )
    else:
        blocking_nlist = UniversalNeighborlistParameters(0, 0, 0, 0)
    logging.info(f"Estimated neighbor list parameters: {neighborlist_params}")
    min_half_box = float(system.data.cell.perpendicular_lengths.min() / 2)
    cell_tables = init_cell_tables(
        particles, motifs, system, lj_params, ewald_params, config.kernel_dtype
    )
    return MCMCState(
        particles=particles,
        groups=groups,
        motifs=motifs,
        systems=system,
        neighborlist_params=neighborlist_params,
        blocking_spheres_neighborlist_params=blocking_nlist,
        lj_parameters=WithCache(
            lj_params,
            KahanSummand.init(
                PotentialOut(
                    Table.arange(jnp.zeros(n_sys), label=SystemId), EMPTY, EMPTY
                )
            ),
        ),
        ewald_parameters=WithCache(ewald_params, EwaldCache.make(n_sys, n_kvecs)),
        blocking_spheres_parameters=blocking_spheres,
        translation_params=Table.arange(
            ParameterSchedulerState.create(n_sys, upper_bound=min_half_box),
            label=SystemId,
        ),
        rotation_params=Table.arange(
            ParameterSchedulerState.create(n_sys), label=SystemId
        ),
        reinsertion_params=Table.arange(
            ParameterSchedulerState.create(n_sys), label=SystemId
        ),
        exchange_params=Table.arange(
            ParameterSchedulerState.create(n_sys), label=SystemId
        ),
        cell_tables=cell_tables,
    )

def make_guest_stress(
    state: MCMCState,
) -> StateProperty[MCMCState, Table[SystemId, StressResult]]:
    """Build guest-only stress tensor function from MCMCState."""
    state_lens = identity_lens(MCMCState)
    # The tail correction is deliberately not differentiated here: its energy
    # depends on the cell through the volume, so ``dU_tail/dh`` would add
    # ``(U_tail / V) * I`` to the configurational stress on top of the
    # closed-form ``p_tail`` added below, double-counting it.
    potentials: list[
        Potential[MCMCState, PositionsAndCell, EmptyType, Patch[MCMCState]]
    ] = [
        make_lennard_jones_from_state(state_lens, gradient=POSITIONS_AND_CELL),
    ]
    if state.is_charged:
        potentials.append(
            make_ewald_from_state(
                state_lens,
                gradient=POSITIONS_AND_CELL,
                include_exclusion_mask=True,
            )
        )
    potential = sum_potentials(*potentials)
    potential = CachedPotential(
        potential,
        lens(
            lambda x: PotentialOut(
                x.systems.map_data(lambda s: s.potential_energy.value),
                PositionsAndCell(
                    x.particles.set_data(x.particles.data.position_gradients),
                    x.systems.set_data(x.systems.data.cell_gradients),
                ),
                EMPTY,
            )
        ),
        lambda x: PotentialOut(
            x.systems.index,  # type: ignore
            PositionsAndCell(x.particles.data.system, x.systems.index),  # type: ignore[arg-type]
            EMPTY,
        ),
    )
    propagator = PotentialAsPropagator(potential)

    def stress_fn(key: Array, state: MCMCState) -> Table[SystemId, StressResult]:
        chain = key_chain(key)
        # For stress calculations, we only consider the guest particles, so we take a "guest-only view" of the state.
        guest_only_state = state.guest_only
        state_with_derivatives = propagator(next(chain), guest_only_state)
        config_stress = molecular_virial_stress_from_state(
            next(chain), state_with_derivatives
        )
        p_tail = global_lennard_jones_tail_correction_pressure_from_state(
            next(chain), state
        )
        group_counts = state.groups.data.system.counts
        p_ideal = ideal_gas_pressure(group_counts, state.systems)
        return Table(
            state.systems.keys,
            StressResult(
                potential=config_stress.data,
                tail_correction=p_tail.data[:, None, None] * jnp.eye(3),
                ideal_gas=p_ideal.data[:, None, None] * jnp.eye(3),
            ),
        )

    return stress_fn

def make_propagator(
    state: MCMCState,
    config: RunConfig,
) -> tuple[Propagator[MCMCState], Propagator[MCMCState]]:
    """Build MCMC initialization and cycles."""
    state_lens = identity_lens(MCMCState)

    potential = fuse_pair_potentials(
        _make_potential(state),
        state,
        lambda s: s.particles,
        lambda s: s.systems,
        state_lens.focus(lambda s: s.cell_tables),
        probe=lambda s, p: p.particles,
        max_queries_per_system=state.motifs.data.motif.max_count,
    )
    potential, probability_ratio = make_muvt_probability_ratio(state_lens, potential)
    propagator = make_gcmc_mcmc_propagator(
        state_lens,
        MCMCStateUpdate.from_changes,
        probability_ratio,
        translation_weight=config.translation_prob,
        rotation_weight=config.rotation_prob,
        reinsertion_weight=config.reinsertion_prob,
        exchange_weight=config.exchange_prob,
    )
    propagator = LoopPropagator(
        propagator,
        lambda x: jnp.maximum(
            x.groups.data.system.counts.data.max(), config.min_cycle_length
        ),
    )
    init_prop = ResetOnErrorPropagator(PotentialAsPropagator(potential))
    return init_prop, ResetOnErrorPropagator(propagator)