kups.potential.classical.internal ¶

@dataclass
class InternalEnergies[State, StatePatch: Patch](
    Potential[State, EmptyType, EmptyType, StatePatch]
):
    """Potential providing fixed internal energies for molecular motifs.

    Computes total energy by summing precomputed motif energies for all molecules
    in each system. Used to add/subtract intramolecular contributions in rigid
    body simulations where internal geometries are fixed.

    Type Parameters:
        State: Simulation state type
        StatePatch: Patch type for state updates

    Attributes:
        motifs: Lens to indexed motif data
        motif_potential_out: Lens to precomputed motif energies

    Note:
        Currently does not support gradients or Hessians (rigid molecules).
    """

    motifs: View[State, Table[ParticleId, MotifData]] = field(static=True)
    motif_potential_out: View[State, Energy] = field(static=True)

    def __call__(
        self, state: State, patch: StatePatch | None = None
    ) -> WithPatch[PotentialOut[EmptyType, EmptyType], Patch[State]]:
        sys_idx = self.motifs(state).data.system
        if patch is not None:
            accept = Table(sys_idx.keys, jnp.ones(sys_idx.num_labels, dtype=jnp.bool))
            state = patch(state, accept)
        motifs = self.motifs(state)
        motif_energies = self.motif_potential_out(state)

        out_energies = motifs.data.system.sum_over(motif_energies[motifs.data.values])
        return WithPatch(PotentialOut(out_energies, EMPTY, EMPTY), IdPatch())

@dataclass
class MotifData:
    """Motif data with system assignment.

    Attributes:
        values: Motif index values (integers)
        system: System assignment for each motif entry
    """

    values: Array
    system: Index[SystemId]