kups.potential.mliap.local ¶

class EdgeFunction(Protocol):
    """Protocol for edge/message function.

    Computes edge messages `(n_edges, msg_dim)` from source/target node
    embeddings `(n_edges, embed_dim)` and displacement vectors `(n_edges, 3)`.
    """

    def __call__(
        self, node1: Array, node2: Array, difference_vectors: Array
    ) -> Array: ...

class IsLocalMLIAPGraphParticles(
    IsLocalMLIAPParticleData, IsRadiusGraphPoints, Protocol
):
    """Combined protocol for local MLIAP particles in radius graph context."""

    ...

class IsLocalMLIAPParticleData(HasPositionsAndAtomicNumbers, HasSystemIndex, Protocol):
    """Protocol for particle data required by LocalMLIAP.

    Must provide positions, atomic numbers, and system index.
    """

    ...

@dataclass
class LocalMLIAPCache:
    """Cache for incremental energy updates.

    Stores intermediate values from the last full computation to enable
    efficient incremental updates when only a subset of atoms change.

    Attributes:
        node_init: Cached node embeddings from init_function, shape `(n_atoms, embed_dim)`
        msg_sum: Cached aggregated messages per node, shape `(n_atoms, msg_dim)`
    """

    node_init: Array
    msg_sum: Array

@dataclass
class LocalMLIAPComposer[
    State,
    P: IsLocalMLIAPGraphParticles,
    S: HasCell[AnyPeriodicity],
    Ptch: Patch[Any],
]:
    """Composes simulation state into LocalMLIAP input.

    Extracts particles, edges, and model config from state, handling both
    full computation (patch=None) and incremental updates (patch provided).

    Type Parameters:
        State: Simulation state type
        P: Particle data type (positions + system + inclusion/exclusion + atomic numbers)
        S: System data type (cell + cutoff)
        Ptch: Patch type for incremental updates

    Attributes:
        particles: View to extract indexed particle data from state
        systems: View to extract indexed system data from state
        neighborlist: View to extract a cutoff-bound neighbor list from state
        model: Lens to access model config in state
        probe: Probe to detect particle changes from patch
    """

    particles: View[State, Table[ParticleId, P]] = field(static=True)
    systems: View[State, Table[SystemId, S]] = field(static=True)
    neighborlist: View[State, NeighborList[Literal[2]]] = field(static=True)
    model: Lens[State, LocalMLIAPData] = field(static=True)
    probe: Probe[State, Ptch, IsGraphProbe[P, Literal[2]]] | None = field(static=True)

    def __call__(
        self, state: State, patch: Ptch | None
    ) -> Sum[LocalMLIAPInput[State, P, S]]:
        """Compose state and patch into LocalMLIAP input.

        Args:
            state: Current simulation state
            patch: Proposed changes (None for full computation)

        Returns:
            Sum containing single LocalMLIAPInput summand
        """

        radius_graph_constr = GraphConstructor(
            self.particles,
            self.systems,
            self.neighborlist,
            self.probe,
        )
        pc = PointCloud(self.particles(state), self.systems(state))
        conf = self.model.get(state)
        cache_lens = self.model.focus(lambda x: x.cache)
        if patch is None:
            graph = radius_graph_constr(state, None)
            result = LocalMLIAPInput(pc, None, graph.edges, None, conf, cache_lens)
        else:
            assert self.probe is not None
            new_graph = radius_graph_constr(state, patch)
            old_graph = radius_graph_constr(state, patch, old_graph=True)
            result = LocalMLIAPInput(
                pc,
                self.probe(state, patch).particles,
                new_graph.edges,
                old_graph.edges,
                conf,
                cache_lens,
            )
        return Sum(Summand(result))

@dataclass
class LocalMLIAPData:
    """Configuration for a local MLIAP model.

    Bundles the model functions, cutoff, and cache together with a lens
    for updating the cache in the state.

    Attributes:
        cutoff: Interaction cutoff radius [Å], shape `(n_systems,)`
        init_function: Maps atomic numbers to node embeddings
        edge_function: Computes messages from node pairs and displacements
        readout_function: Computes per-atom energies from embeddings and messages
        cache: Cached values for incremental updates
    """

    cutoff: Table[SystemId, Array]
    init_function: NodeInitFunction = field(static=True)
    edge_function: EdgeFunction = field(static=True)
    readout_function: ReadoutFunction = field(static=True)
    cache: LocalMLIAPCache

    @staticmethod
    def from_zip_file(zip_file: str | Path, n_atoms: int) -> "LocalMLIAPData":
        """Load a local MLIAP model from a zip archive.

        Expects ``node_init.jax``, ``edge.jax``, ``readout.jax``,
        and ``metadata.json`` (with ``cutoff`` and ``precision`` keys).

        Args:
            zip_file: Path to the ``.zip`` archive.
            n_atoms: Number of atoms to allocate cache for.

        Returns:
            Loaded local MLIAP model with initialized cache.
        """
        with zipfile.ZipFile(zip_file, "r") as zf:
            with zf.open("metadata.json") as f:
                metadata = json.loads(f.read().decode())
            with zf.open("node_init.jax") as f:
                exported_init = jax.export.deserialize(f.read())  # type: ignore
            with zf.open("edge.jax") as f:
                exported_edge = jax.export.deserialize(f.read())  # type: ignore
            with zf.open("readout.jax") as f:
                exported_readout = jax.export.deserialize(f.read())  # type: ignore

        convert = _make_dtype_converter(metadata["precision"])

        def init_fn(atomic_numbers: Array) -> Array:
            return sequential_vmap_with_vjp(exported_init.call)(convert(atomic_numbers))

        def edge_fn(node1: Array, node2: Array, difference_vectors: Array) -> Array:
            fn = sequential_vmap_with_vjp(exported_edge.call)
            return fn(convert(node1), convert(node2), convert(difference_vectors))

        def readout_fn(node_emb: Array, msg_sum: Array) -> Array:
            return sequential_vmap_with_vjp(exported_readout.call)(
                convert(node_emb), convert(msg_sum)
            )

        # Infer dimensions
        sample_atomic_numbers = jnp.zeros((1,), dtype=int)
        node_emb_dim = jax.eval_shape(init_fn, sample_atomic_numbers).shape[-1]
        sample_node = jnp.zeros((1, node_emb_dim))
        msg_dim = jax.eval_shape(
            edge_fn, sample_node, sample_node, jnp.zeros((1, 3))
        ).shape[-1]

        return LocalMLIAPData(
            cutoff=Table((SystemId(0),), jnp.array([metadata["cutoff"]], float)),
            init_function=init_fn,
            edge_function=edge_fn,
            readout_function=readout_fn,
            cache=LocalMLIAPCache(
                node_init=jnp.zeros((n_atoms, node_emb_dim)),
                msg_sum=jnp.zeros((n_atoms, msg_dim)),
            ),
        )

@dataclass
class LocalMLIAPInput[
    State,
    P: IsLocalMLIAPGraphParticles,
    S: HasCell[AnyPeriodicity],
]:
    """Input bundle for local MLIAP energy computation.

    Contains all data needed to compute energies, supporting both full
    computation and incremental updates.

    Type Parameters:
        State: Simulation state type
        P: Particle data type (positions, atomic numbers, system, inclusion/exclusion)
        S: System data type (must have cell)

    Attributes:
        point_cloud: Current particle positions and systems
        point_cloud_changes: Changed particles for incremental update (None for full)
        edges: Current edges within cutoff
        edges_deleted: Old edges to remove for incremental update (None for full)
        config: Model configuration with functions and cache
        cache_lens: Lens to access/update cache in state
    """

    point_cloud: PointCloud[P, S]
    point_cloud_changes: WithIndices[ParticleId, P] | None
    edges: Edges[Literal[2]]
    edges_deleted: Edges[Literal[2]] | None
    config: LocalMLIAPData
    cache_lens: Lens[State, LocalMLIAPCache] = field(static=True)

@dataclass
class LocalMLIAPPatch[State](Patch[State]):
    """Patch to update the MLIAP cache in state.

    Applied after energy computation to update cached node embeddings
    and aggregated messages for systems where moves were accepted.

    Type Parameters:
        State: Simulation state type

    Attributes:
        cache: New cache values to apply
        system_idx: System index for masking
        lens: Lens to access cache in state
    """

    cache: LocalMLIAPCache
    system_idx: Index[SystemId]
    lens: Lens[State, LocalMLIAPCache] = field(static=True)

    def __call__(self, state: State, accept: Accept) -> State:
        """Apply cache update to state.

        Args:
            state: Current simulation state
            accept: Boolean mask per system indicating accepted moves

        Returns:
            Updated state with new cache values where mask is True
        """
        new_cache = self.cache
        mask = accept[self.system_idx]
        return self.lens.apply(
            state,
            lambda cache: tree_map(
                lambda a, b: where_broadcast_last(mask, a, b), new_cache, cache
            ),
        )

class NodeInitFunction(Protocol):
    """Protocol for node initialization function.

    Maps atomic numbers `(n_atoms,)` to node embeddings `(n_atoms, embed_dim)`.
    """

    def __call__(self, atomic_numbers: Array) -> Array: ...

class ReadoutFunction(Protocol):
    """Protocol for readout function.

    Computes per-atom energies `(n_atoms,)` from node embeddings and
    aggregated messages, both of shape `(n_atoms, embed_dim)`.
    """

    def __call__(self, node_emb: Array, msg_sum: Array) -> Array: ...