Skip to content

kups.potential.common.direct

Direct potential for models providing precomputed gradients.

Wraps model functions that directly produce energy and gradients (e.g., PyTorch force fields) into the kUPS Potential protocol. Unlike PotentialFromEnergy which uses autodiff, this passes through whatever gradients/Hessians the model provides.

DirectPotential

Bases: Potential[State, Gradients, Hessians, StatePatch]

Potential wrapping models that directly produce gradients.

For models providing precomputed gradients (e.g., PyTorch force fields). Passes through whatever gradients and Hessians the model provides (typically Hessians=EmptyType).

Source code in src/kups/potential/common/direct.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
@dataclass
class DirectPotential[
    State,
    Input,
    Gradients,
    Hessians,
    StatePatch: Patch,
](Potential[State, Gradients, Hessians, StatePatch]):
    """Potential wrapping models that directly produce gradients.

    For models providing precomputed gradients (e.g., PyTorch force fields).
    Passes through whatever gradients and Hessians the model provides
    (typically Hessians=EmptyType).
    """

    direct_potential_fn: DirectPotentialFunction[State, Input, Gradients, Hessians] = (
        field(static=True)
    )
    composer: SumComposer[State, Input, StatePatch] = field(static=True)
    cache_lens: Lens[State, PotentialOut[Gradients, Hessians]] | None = field(
        static=True
    )
    patch_idx_view: View[State, PotentialOut[Gradients, Hessians]] | None = field(
        static=True
    )

    @jit
    def __call__(
        self, state: State, patch: StatePatch | None = None
    ) -> WithPatch[PotentialOut[Gradients, Hessians], Patch[State]]:
        dp_plan = self.composer(state, patch)
        assert len(dp_plan) > 0, "At least one configuration must be added."

        outs: list[PotentialOut[Gradients, Hessians]] = []
        patches: list[Patch[State]] = []

        for inp, weight in dp_plan:
            result = self.direct_potential_fn(inp)
            outs.append(weight * result.data)
            patches.append(result.patch)

        if dp_plan.add_previous_total:
            assert self.cache_lens is not None
            outs.append(self.cache_lens.get(state))

        total = kahan_summation(*outs)[0]

        if self.cache_lens is not None:
            assert self.patch_idx_view is not None
            cache_patch = IndexLensPatch(
                total, self.patch_idx_view(state), self.cache_lens
            )
        else:
            cache_patch = IdPatch()

        out_patch = ComposedPatch((cache_patch, *patches))
        return WithPatch(total, out_patch)

DirectPotentialFunction

Bases: Protocol

Protocol for functions returning PotentialOut directly.

Source code in src/kups/potential/common/direct.py 
20
21
22
23
24
25
class DirectPotentialFunction[State, Input, Gradients, Hessians](Protocol):
    """Protocol for functions returning PotentialOut directly."""

    def __call__(
        self, inp: Input
    ) -> WithPatch[PotentialOut[Gradients, Hessians], Patch[State]]: ...