kups.potential.classical.lennard_jones ¶

@jit
def global_lennard_jones_tail_correction_energy(
    inp: GlobalTailCorrectedLennardJonesInput,
) -> WithPatch[Table[SystemId, Energy], IdPatch[Any]]:
    """Compute analytical long-range tail correction energy."""
    q1, q2, _volume, n_graphs = _global_tail_correction_common(inp)
    result = _tail_energy(q1, q2)
    total_energies = Table.arange(result, label=SystemId)
    assert len(total_energies) == n_graphs
    return WithPatch(total_energies, IdPatch[Any]())

@jit
def global_lennard_jones_tail_correction_pressure(
    inp: GlobalTailCorrectedLennardJonesInput,
) -> WithPatch[Table[SystemId, Energy], IdPatch[Any]]:
    """Compute analytical long-range tail correction for pressure."""
    q1, q2, volume, n_graphs = _global_tail_correction_common(inp)
    result = (16 / 3) * jnp.pi / volume * (q2 / 3 * 2 - q1)
    total_pressure = Table.arange(result, label=SystemId)
    assert len(total_pressure) == n_graphs
    return WithPatch(total_pressure, IdPatch[Any]())

@jit
def lennard_jones_pair_kernel(
    parameters: LennardJonesParameters,
    labels_i: Index[Label],
    labels_j: Index[Label],
    rij: Array,
    r2: Array,
    system: Index[SystemId],
    /,
) -> Array:
    """Lennard-Jones pair kernel shared by graph and fused evaluation.

    Callers apply the cutoff; pair inputs broadcast to a common shape.

    Args:
        parameters: LJ mixing tables.
        labels_i: Species labels of the left atoms.
        labels_j: Species labels of the right atoms.
        rij: Difference vectors (unused).
        r2: Squared distances.
        system: Left-side system ids (unused).

    Returns:
        Pair energies.
    """
    del rij, system
    species_i = labels_i.indices_in(parameters.labels)
    species_j = labels_j.indices_in(parameters.labels)
    n_species = len(parameters.labels)
    if (
        species_i.size * n_species <= species_j.size
        and n_species <= MAX_SELECTED_SPECIES
    ):
        # Few query rows against many keys: select mixing rows per species.
        epsilon_rows = parameters.epsilon[species_i]
        sigma_rows = parameters.sigma[species_i]
        # A where chain: jnp.select's argmax over the matches is slower here.
        epsilon = jnp.zeros((), epsilon_rows.dtype)
        sigma = jnp.zeros((), sigma_rows.dtype)
        for t in range(n_species):
            match = species_j == t
            epsilon = jnp.where(match, epsilon_rows[..., t], epsilon)
            sigma = jnp.where(match, sigma_rows[..., t], sigma)
    else:
        epsilon = parameters.epsilon[species_i, species_j]
        sigma = parameters.sigma[species_i, species_j]
    c6 = (sigma**2 / r2) ** 3
    return 4 * epsilon * (c6**2 - c6)

def make_global_lennard_jones_tail_correction_potential[State, Gradients, Hessians](
    particles_view: View[State, Table[ParticleId, IsLJGraphParticles]],
    systems_view: View[State, Table[SystemId, HasCell[AnyPeriodicity]]],
    parameter_view: View[State, GlobalTailCorrectedLennardJonesParameters],
    gradient_lens: Lens[GCLJInp, Gradients],
    hessian_lens: Lens[Gradients, Hessians],
    hessian_idx_view: View[State, Hessians],
    patch_idx_view: View[State, PotentialOut[Gradients, Hessians]] | None = None,
    out_cache_lens: Lens[State, KahanSummand[PotentialOut[Gradients, Hessians]]]
    | None = None,
    *,
    composition: RigidBodyComposition[State, HasLabels] | None = None,
) -> Potential[State, Gradients, Hessians, Patch[State]]:
    """Create the LJ tail correction, optionally prepared for rigid bodies.

    ``composition`` projects the same tail formula onto rigid-body counts once
    at construction. It requires fixed parameters and cells and no derivatives.
    """
    constructor = GraphInputConstructor(
        GraphConstructor(
            particles=particles_view,
            systems=systems_view,
            neighborlist=lambda _: EmptyNeighborList[Literal[0]](),
            probe=None,
        ),
        parameter_view=parameter_view,
    )
    if composition is not None:
        initial = composition.initial_state
        coefficients = _molecular_tail_coefficients(
            composition,
            particles_view(initial).data,
            parameter_view(initial),
            systems_view(initial).data.cell.volume,
        )
        return composition.potential(
            coefficients,
            gradient_lens(constructor(initial, None)),
            hessian_lens,
            out_cache_lens,
            patch_idx_view,
        )
    return PotentialFromEnergy(
        energy_fn=global_lennard_jones_tail_correction_energy,
        composer=FullSumComposer(constructor),
        gradient_lens=gradient_lens,
        hessian_lens=hessian_lens,
        hessian_idx_view=hessian_idx_view,
        cache_lens=out_cache_lens,
        patch_idx_view=patch_idx_view,
    )

def make_global_lennard_jones_tail_correction_pressure[State](
    particles_view: View[State, Table[ParticleId, IsLJGraphParticles]],
    systems_view: View[State, Table[SystemId, HasCell[AnyPeriodicity]]],
    parameter_view: View[State, GlobalTailCorrectedLennardJonesParameters],
) -> StateProperty[State, Table[SystemId, Array]]:
    """Create long-range pressure correction for Lennard-Jones systems."""
    graph_constructor = GraphConstructor(
        particles=particles_view,
        systems=systems_view,
        neighborlist=lambda _: EmptyNeighborList[Literal[0]](),
        probe=None,
    )

    def pressure(key: Array, state: State) -> Table[SystemId, Array]:
        del key
        params = parameter_view(state)
        graph = graph_constructor(state, None)
        return global_lennard_jones_tail_correction_pressure(
            GraphPotentialInput(params, graph)
        ).data

    return pressure

def make_lennard_jones_potential[
    State,
    Ptch: Patch[Any],
    Gradients,
    Hessians,
](
    particles_view: View[State, Table[ParticleId, IsLJGraphParticles]],
    systems_view: View[State, Table[SystemId, HasCell[AnyPeriodicity]]],
    neighborlist_view: View[State, NeighborList[Literal[2]]],
    parameter_view: View[State, LennardJonesParameters],
    probe: Probe[State, Ptch, IsParticleProbe[IsLJGraphParticles]] | None,
    gradient_lens: Lens[LJRadiusInp, Gradients],
    hessian_lens: Lens[Gradients, Hessians],
    hessian_idx_view: View[State, Hessians],
    patch_idx_view: View[State, PotentialOut[Gradients, Hessians]] | None = None,
    out_cache_lens: Lens[State, KahanSummand[PotentialOut[Gradients, Hessians]]]
    | None = None,
) -> Potential[State, Gradients, Hessians, Ptch]:
    """Create a standard Lennard-Jones potential with sharp cutoff."""
    graph_fn = GraphConstructor(
        particles=particles_view,
        systems=systems_view,
        neighborlist=neighborlist_view,
        probe=probe,
    )
    composer = LocalSumComposer(
        GraphInputConstructor(
            graph_constructor=graph_fn,
            parameter_view=parameter_view,
        )
    )
    return PotentialFromEnergy(
        composer=composer,
        energy_fn=lennard_jones_energy,
        gradient_lens=gradient_lens,
        hessian_lens=hessian_lens,
        hessian_idx_view=hessian_idx_view,
        cache_lens=out_cache_lens,
        patch_idx_view=patch_idx_view,
    )

def make_pair_tail_corrected_lennard_jones_potential[
    State,
    Ptch: Patch[Any],
    Gradients,
    Hessians,
](
    particles_view: View[State, Table[ParticleId, IsLJGraphParticles]],
    systems_view: View[State, Table[SystemId, HasCell[AnyPeriodicity]]],
    neighborlist_view: View[State, NeighborList[Literal[2]]],
    parameter_view: View[State, PairTailCorrectedLennardJonesParameters],
    probe: Probe[State, Ptch, IsParticleProbe[IsLJGraphParticles]] | None,
    gradient_lens: Lens[PCLJInp, Gradients],
    hessian_lens: Lens[Gradients, Hessians],
    hessian_idx_view: View[State, Hessians],
    patch_idx_view: View[State, PotentialOut[Gradients, Hessians]] | None = None,
    out_cache_lens: Lens[State, KahanSummand[PotentialOut[Gradients, Hessians]]]
    | None = None,
) -> Potential[State, Gradients, Hessians, Ptch]:
    """Create a Lennard-Jones potential with smooth pairwise tail correction."""
    radius_graph_fn = GraphConstructor(
        particles=particles_view,
        systems=systems_view,
        neighborlist=neighborlist_view,
        probe=probe,
    )
    composer = LocalSumComposer(
        GraphInputConstructor(
            graph_constructor=radius_graph_fn,
            parameter_view=parameter_view,
        )
    )
    return PotentialFromEnergy(
        composer=composer,
        energy_fn=pair_tail_corrected_lennard_jones_energy,
        gradient_lens=gradient_lens,
        hessian_lens=hessian_lens,
        hessian_idx_view=hessian_idx_view,
        cache_lens=out_cache_lens,
        patch_idx_view=patch_idx_view,
    )

@jit
def pair_tail_corrected_lennard_jones_energy(
    inp: PairTailCorrectedLennardJonesInput,
) -> WithPatch[Table[SystemId, Energy], IdPatch[Any]]:
    """Compute Lennard-Jones energy with smooth pairwise tail correction."""
    graph = inp.graph
    r: Array = jnp.linalg.norm(graph.edge_shifts, axis=(-2, -1))
    edge_energy = lennard_jones_energy.edge_energies(inp)

    batch = graph.edge_batch_mask
    r_tr = inp.parameters.truncation_radius[batch]
    r_cut = inp.parameters.cutoff[batch]
    mask = r > r_tr
    remove = r >= r_cut
    factor1 = ((r_cut**2) - r**2) ** 2
    factor2 = 2 * r**2 + (r_cut**2 - 3 * r_tr**2)
    div = (r_cut**2 - r_tr**2) ** 3
    corrected_edge_energy = jnp.where(
        mask, edge_energy * factor1 * factor2 / div, edge_energy
    )
    corrected_edge_energy = jnp.where(remove, 0.0, corrected_edge_energy)
    total_energies = batch.sum_over(corrected_edge_energy) / 2
    return WithPatch(total_energies, IdPatch[Any]())