k3-node 1.0.0__py3-none-any.whl

k3_node/applications/materials/readout.py

@@ -0,0 +1,235 @@
+"""Readout layers for materials models."""
+
+from __future__ import annotations
+
+from typing import Sequence, Optional, Union
+import keras
+from keras import layers, ops
+from k3_node.layers.pool import global_add_pool, global_mean_pool, global_max_pool
+from .core import MLP, GatedMLP, get_activation, infer_num_graphs
+from k3_node.ops.segment import segment_max, segment_sum
+
+
+class ReduceReadOut(layers.Layer):
+    """Pool node features into graph features via sum, mean, or max reduction.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import MatGLReduceReadOut
+
+        node_feat = np.random.rand(6, 16).astype("float32")
+        batch = np.array([0, 0, 0, 1, 1, 1])  # two structures with 3 atoms each
+
+        print(tuple(MatGLReduceReadOut(op="mean")(node_feat, batch=batch, num_graphs=2).shape))  # (2, 16)
+        ```
+    """
+
+    def __init__(self, op: str = "mean", field: str = "node_feat", **kwargs):
+        super().__init__(**kwargs)
+        self.op = op.lower()
+        self.field = field
+
+    def call(self, node_feat, batch=None, num_graphs=None):
+        n_graphs = infer_num_graphs(batch=batch, num_graphs=num_graphs)
+        if self.op == "mean":
+            return global_mean_pool(node_feat, batch=batch, size=n_graphs)
+        elif self.op == "sum":
+            return global_add_pool(node_feat, batch=batch, size=n_graphs)
+        elif self.op == "max":
+            return global_max_pool(node_feat, batch=batch, size=n_graphs)
+        else:
+            raise ValueError(f"Unsupported reduction op: {self.op}")
+
+
+class WeightedReadOut(layers.Layer):
+    """Feed node features through a GatedMLP to predict atomic properties.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import MatGLWeightedReadOut
+
+        node_feat = np.random.rand(6, 16).astype("float32")
+        batch = np.array([0, 0, 0, 1, 1, 1])  # two structures with 3 atoms each
+
+        print(tuple(MatGLWeightedReadOut(in_feats=16, dims=[16], num_targets=1)(node_feat).shape))  # (6, 1): per-atom outputs
+        ```
+    """
+
+    def __init__(self, in_feats: int, dims: Sequence[int], num_targets: int = 1, **kwargs):
+        super().__init__(**kwargs)
+        self.in_feats = in_feats
+        self.dims = list(dims)
+        self.num_targets = num_targets
+        self.gated = GatedMLP(in_feats=in_feats, dims=[*self.dims, num_targets], activate_last=False)
+
+    def call(self, node_feat):
+        return self.gated(node_feat)
+
+
+class WeightedAtomReadOut(layers.Layer):
+    """Weighted atom readout for whole-graph properties with normalized learned weights.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import MatGLWeightedAtomReadOut
+
+        node_feat = np.random.rand(6, 16).astype("float32")
+        batch = np.array([0, 0, 0, 1, 1, 1])  # two structures with 3 atoms each
+
+        readout = MatGLWeightedAtomReadOut(in_feats=16, dims=[16, 8])
+        print(tuple(readout(node_feat, batch=batch, num_graphs=2).shape))  # (2, 8)
+        ```
+    """
+
+    def __init__(
+        self,
+        in_feats: int,
+        dims: Sequence[int],
+        activation: str = "swish",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_feats = in_feats
+        self.dims = list(dims)
+        self.mlp = MLP([in_feats, *self.dims], activation=activation, activate_last=True)
+        self.weight_mlp = MLP([in_feats, *self.dims[:-1], 1], activation=activation, activate_last=False)
+
+    def call(self, node_feat, batch=None, num_graphs=None):
+        if batch is None:
+            batch = ops.zeros((ops.shape(node_feat)[0],), dtype="int32")
+        batch = ops.cast(batch, "int32")
+        n_graphs = infer_num_graphs(batch=batch, num_graphs=num_graphs)
+
+        updated_field = self.mlp(node_feat)
+        weights = ops.sigmoid(self.weight_mlp(node_feat))  # [num_nodes, 1]
+
+        weight_sum = segment_sum(weights, batch, num_segments=n_graphs)  # [num_graphs, 1]
+        ws_len = ops.shape(weight_sum)[0]
+        b_safe = ops.clip(batch, 0, ops.maximum(ws_len - 1, 0))
+        weight_sum_per_node = ops.take(weight_sum, b_safe, axis=0)
+        factor = weights / ops.maximum(weight_sum_per_node, 1e-8)  # [num_nodes, 1]
+
+        return global_add_pool(factor * updated_field, batch=batch, size=n_graphs)
+
+
+class Set2SetReadOut(layers.Layer):
+    """Iterative content-based attention pooling (Set2Set) for nodes.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import MatGLSet2SetReadOut
+
+        node_feat = np.random.rand(6, 16).astype("float32")
+        batch = np.array([0, 0, 0, 1, 1, 1])  # two structures with 3 atoms each
+
+        print(tuple(MatGLSet2SetReadOut(in_channels=16)(node_feat, batch=batch, num_graphs=2).shape))  # (2, 32)
+        ```
+    """
+
+    def __init__(self, in_channels: int, processing_steps: int = 3, num_layers: int = 1, **kwargs):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = 2 * in_channels
+        self.processing_steps = processing_steps
+        self.num_layers = num_layers
+        self.lstm_cell = layers.LSTMCell(in_channels)
+
+    def build(self, input_shape=None):
+        self.lstm_cell.build((None, self.out_channels))
+        super().build(input_shape)
+
+    def call(self, x, batch=None, num_graphs=None):
+        if batch is None:
+            batch = ops.zeros((ops.shape(x)[0],), dtype="int32")
+        else:
+            batch = ops.cast(batch, "int32")
+        n_graphs = infer_num_graphs(batch=batch, num_graphs=num_graphs)
+
+        h = [
+            ops.zeros((n_graphs, self.in_channels), dtype=x.dtype),
+            ops.zeros((n_graphs, self.in_channels), dtype=x.dtype),
+        ]
+        q_star = ops.zeros((n_graphs, self.out_channels), dtype=x.dtype)
+        b_safe = ops.clip(batch, 0, ops.maximum(n_graphs - 1, 0))
+
+        for _ in range(self.processing_steps):
+            q, h = self.lstm_cell(q_star, h)
+            q_taken = ops.take(q, b_safe, axis=0)
+            e = ops.sum(x * q_taken, axis=-1, keepdims=True)
+
+            max_e = segment_max(e, batch, num_segments=n_graphs)
+            max_e_taken = ops.take(max_e, b_safe, axis=0)
+            exp_e = ops.exp(e - max_e_taken)
+
+            sum_exp_e = segment_sum(exp_e, batch, num_segments=n_graphs)
+            sum_exp_e_taken = ops.take(sum_exp_e, b_safe, axis=0)
+            alpha = exp_e / ops.maximum(sum_exp_e_taken, 1e-12)
+
+            r = segment_sum(alpha * x, batch, num_segments=n_graphs)
+            q_star = ops.concatenate([q, r], axis=-1)
+
+        return q_star
+
+
+class EdgeSet2Set(layers.Layer):
+    """Iterative content-based attention pooling (Set2Set) for edge features.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import MatGLEdgeSet2Set
+
+        edge_feat = np.random.rand(8, 16).astype("float32")
+        edge_batch = np.array([0, 0, 0, 0, 1, 1, 1, 1])  # bonds of two structures
+        print(tuple(MatGLEdgeSet2Set(input_dim=16)(edge_feat, edge_batch=edge_batch, num_graphs=2).shape))  # (2, 32)
+        ```
+    """
+
+    def __init__(self, input_dim: int, n_iters: int = 3, n_layers: int = 1, **kwargs):
+        super().__init__(**kwargs)
+        self.input_dim = input_dim
+        self.output_dim = 2 * input_dim
+        self.n_iters = n_iters
+        self.n_layers = n_layers
+        self.lstm_cell = layers.LSTMCell(input_dim)
+
+    def build(self, input_shape=None):
+        self.lstm_cell.build((None, self.output_dim))
+        super().build(input_shape)
+
+    def call(self, edge_feat, edge_batch=None, num_graphs=None):
+        if edge_batch is None:
+            edge_batch = ops.zeros((ops.shape(edge_feat)[0],), dtype="int32")
+        else:
+            edge_batch = ops.cast(edge_batch, "int32")
+        n_graphs = infer_num_graphs(batch=edge_batch, num_graphs=num_graphs)
+
+        h = [
+            ops.zeros((n_graphs, self.input_dim), dtype=edge_feat.dtype),
+            ops.zeros((n_graphs, self.input_dim), dtype=edge_feat.dtype),
+        ]
+        q_star = ops.zeros((n_graphs, self.output_dim), dtype=edge_feat.dtype)
+        eb_safe = ops.clip(edge_batch, 0, ops.maximum(n_graphs - 1, 0))
+
+        for _ in range(self.n_iters):
+            q, h = self.lstm_cell(q_star, h)
+            q_taken = ops.take(q, eb_safe, axis=0)
+            e = ops.sum(edge_feat * q_taken, axis=-1, keepdims=True)
+
+            max_e = segment_max(e, edge_batch, num_segments=n_graphs)
+            max_e_taken = ops.take(max_e, eb_safe, axis=0)
+            exp_e = ops.exp(e - max_e_taken)
+
+            sum_exp_e = segment_sum(exp_e, edge_batch, num_segments=n_graphs)
+            sum_exp_e_taken = ops.take(sum_exp_e, eb_safe, axis=0)
+            alpha = exp_e / ops.maximum(sum_exp_e_taken, 1e-12)
+
+            r = segment_sum(alpha * edge_feat, edge_batch, num_segments=n_graphs)
+            q_star = ops.concatenate([q, r], axis=-1)
+
+        return q_star
+

k3_node/applications/materials/so3net.py

@@ -0,0 +1,234 @@
+"""Multi-backend Keras 3 implementation of SO3Net."""
+
+from __future__ import annotations
+
+import math
+from typing import Sequence, Optional, Union, Tuple, Dict, Any, Literal
+import keras
+from keras import layers, ops
+import numpy as np
+
+from .core import MLP, scatter_add, infer_num_graphs
+from .basis import (
+    BondExpansion,
+    RadialBesselFunction,
+    compute_pair_vector_and_distance,
+    cosine_cutoff,
+)
+from .readout import WeightedAtomReadOut, ReduceReadOut
+
+
+class RealSphericalHarmonics(layers.Layer):
+    """Computes real spherical harmonics up to order lmax for 3D unit vectors.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import RealSphericalHarmonics
+
+        vec = np.random.rand(8, 3).astype("float32")  # bond vectors
+        print(tuple(RealSphericalHarmonics(lmax=2)(vec).shape))  # (8, 9): (lmax + 1)^2 harmonics per bond
+        ```
+    """
+
+    def __init__(self, lmax: int = 2, **kwargs):
+        super().__init__(**kwargs)
+        self.lmax = lmax
+        self.num_components = (lmax + 1) ** 2
+
+    def call(self, vec):
+        norm = ops.sqrt(ops.maximum(ops.sum(vec ** 2, axis=-1, keepdims=True), 1e-12))
+        u = vec / norm
+        x, y, z = u[..., 0:1], u[..., 1:2], u[..., 2:3]
+
+        c00 = 0.5 * math.sqrt(1.0 / math.pi)
+        y00 = ops.broadcast_to(ops.convert_to_tensor(c00, dtype=vec.dtype), ops.shape(x))
+        comps = [y00]
+
+        if self.lmax >= 1:
+            c1 = math.sqrt(3.0 / (4.0 * math.pi))
+            comps.extend([c1 * y, c1 * z, c1 * x])
+
+        if self.lmax >= 2:
+            c2_0 = 0.5 * math.sqrt(5.0 / (4.0 * math.pi))
+            c2_1 = math.sqrt(15.0 / (4.0 * math.pi))
+            c2_2 = 0.5 * math.sqrt(15.0 / (4.0 * math.pi))
+            comps.extend([
+                c2_1 * (x * y),
+                c2_1 * (y * z),
+                c2_0 * (3.0 * z ** 2 - 1.0),
+                c2_1 * (x * z),
+                c2_2 * (x ** 2 - y ** 2),
+            ])
+
+        return ops.concatenate(comps[:self.num_components], axis=-1)
+
+
+class SO3Convolution(layers.Layer):
+    """Equivariant interaction convolution layer for SO3Net.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import SO3Convolution
+
+        edge_index = np.array([[0, 1, 1, 2, 2, 3, 3, 0], [1, 0, 2, 1, 3, 2, 0, 3]])  # 4 atoms, 8 bonds
+        node_feat = np.random.rand(4, 16).astype("float32")
+        edge_feat = np.random.rand(8, 16).astype("float32")
+        sh_feats = np.random.rand(8, 9).astype("float32")  # spherical harmonics of each bond (lmax=2)
+        edge_weight = np.random.rand(8).astype("float32") * 3.0
+
+        conv = SO3Convolution(units=16, lmax=2, num_rbf=16)
+        print(tuple(conv(edge_index, node_feat, edge_feat, sh_feats, edge_weight).shape))  # (4, 16)
+        ```
+    """
+
+    def __init__(self, units: int, lmax: int = 2, num_rbf: int = 32, cutoff: float = 5.0, activation: str = "swish", **kwargs):
+        super().__init__(**kwargs)
+        self.units = units
+        self.lmax = lmax
+        self.cutoff = float(cutoff)
+        self.num_sh = (lmax + 1) ** 2
+        self.rbf_proj = layers.Dense(units, use_bias=True)
+        self.node_proj = layers.Dense(units, use_bias=True)
+        self.sh_proj = layers.Dense(units, use_bias=False)
+        self.out_proj = layers.Dense(units, use_bias=True)
+
+    def call(self, edge_index, node_feats, edge_attr, sh_feats, edge_weight):
+        src = ops.cast(edge_index[0], "int32")
+        dst = ops.cast(edge_index[1], "int32")
+        num_nodes = ops.shape(node_feats)[0]
+
+        C = ops.expand_dims(cosine_cutoff(edge_weight, self.cutoff), axis=-1)
+        rbf_h = ops.silu(self.rbf_proj(edge_attr)) * C
+        node_h = ops.take(self.node_proj(node_feats), dst, axis=0)
+
+        scalar_msg = rbf_h * node_h
+        msg = scatter_add(scalar_msg, src, num_segments=num_nodes)
+        return node_feats + ops.silu(self.out_proj(msg))
+
+
+class SO3Net(keras.Model):
+    """SO(3)-equivariant representation model using spherical harmonics.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import SO3Net
+
+        # A 4-atom structure: positions, bonds (listed in both directions) and atomic numbers
+        structure = {
+            "pos": np.array([[0.0, 0.0, 0.0], [1.0, 0.5, 0.0], [0.5, 1.2, 0.8], [1.5, 1.5, 1.0]], dtype="float32"),
+            "edge_index": np.array([[0, 1, 1, 2, 2, 3, 3, 0], [1, 0, 2, 1, 3, 2, 0, 3]]),
+            "line_edge_index": np.array([[0, 1, 2, 3], [1, 2, 3, 0]]),  # bond pairs forming angles
+            "node_type": np.array([6, 8, 1, 6]),  # atomic numbers
+            "batch": np.zeros(4, dtype="int32"),  # all atoms belong to structure 0
+            "state_attr": np.zeros((1, 2), dtype="float32"),  # global state features
+        }
+
+        model = SO3Net(units=16, nblocks=2, lmax=2, num_rbf=16)
+        energy = model(structure)  # predicted property (e.g. energy) of the structure
+        print(tuple(energy.shape))  # (1,)
+        ```
+    """
+
+    def __init__(
+        self,
+        units: int = 64,
+        nblocks: int = 3,
+        lmax: int = 2,
+        num_rbf: int = 32,
+        cutoff: float = 5.0,
+        ntypes_node: int = 95,
+        ntargets: int = 1,
+        readout_type: Literal["weighted_atom", "reduce_atom"] = "weighted_atom",
+        activation_type: str = "swish",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cutoff = cutoff
+        self.units = units
+        self.lmax = lmax
+
+        self.rbf = RadialBesselFunction(max_n=num_rbf, cutoff=cutoff)
+        self.sh = RealSphericalHarmonics(lmax=lmax)
+        self.embedding = layers.Embedding(ntypes_node, units)
+
+        self.convs = [
+            SO3Convolution(units=units, lmax=lmax, num_rbf=num_rbf, cutoff=cutoff, activation=activation_type)
+            for _ in range(nblocks)
+        ]
+
+        if readout_type == "weighted_atom":
+            self.readout = WeightedAtomReadOut(units, dims=[units, ntargets], activation=activation_type)
+        else:
+            self.readout = ReduceReadOut(op="mean")
+            self.final_mlp = MLP([units, ntargets], activation=activation_type, activate_last=False)
+
+    def _unpack_inputs(self, inputs):
+        if isinstance(inputs, dict):
+            pos = inputs.get("pos")
+            edge_index = inputs.get("edge_index")
+            node_type = inputs.get("node_type", inputs.get("z"))
+            pbc_offshift = inputs.get("pbc_offshift", None)
+            batch = inputs.get("batch", None)
+            num_graphs = inputs.get("num_graphs", None)
+            state_attr = inputs.get("state_attr", None)
+            return pos, edge_index, node_type, pbc_offshift, batch, num_graphs, state_attr
+        elif isinstance(inputs, (tuple, list)):
+            pos = inputs[0]
+            edge_index = inputs[1]
+            node_type = inputs[2]
+            pbc_offshift = inputs[3] if len(inputs) > 3 else None
+            batch = inputs[4] if len(inputs) > 4 else None
+            num_graphs = inputs[5] if len(inputs) > 5 else None
+            state_attr = inputs[6] if len(inputs) > 6 else None
+            return pos, edge_index, node_type, pbc_offshift, batch, num_graphs, state_attr
+        return inputs, None, None, None, None, None, None
+
+    def call(self, inputs, edge_index=None, node_type=None, pbc_offshift=None, batch=None, num_graphs=None, state_attr=None):
+        if edge_index is None:
+            (
+                pos,
+                edge_index,
+                node_type_in,
+                pbc_offshift_in,
+                batch_in,
+                num_graphs_in,
+                state_attr_in,
+            ) = self._unpack_inputs(inputs)
+            if node_type is None:
+                node_type = node_type_in
+            if pbc_offshift is None:
+                pbc_offshift = pbc_offshift_in
+            if batch is None:
+                batch = batch_in
+            if num_graphs is None:
+                num_graphs = num_graphs_in
+            if state_attr is None:
+                state_attr = state_attr_in
+        else:
+            pos = inputs
+
+        num_nodes = ops.shape(pos)[0]
+        if batch is None:
+            batch = ops.zeros((num_nodes,), dtype="int32")
+        else:
+            batch = ops.cast(batch, "int32")
+        n_graphs = infer_num_graphs(batch=batch, num_graphs=num_graphs, state_attr=state_attr)
+
+        vec, bond_dists = compute_pair_vector_and_distance(pos, edge_index, pbc_offshift)
+        edge_attr = self.rbf(bond_dists)
+        sh_attr = self.sh(vec)
+        node_feats = self.embedding(ops.cast(node_type, "int32"))
+
+        for conv in self.convs:
+            node_feats = conv(edge_index, node_feats, edge_attr, sh_attr, bond_dists)
+
+        if isinstance(self.readout, WeightedAtomReadOut):
+            out = self.readout(node_feats, batch=batch, num_graphs=n_graphs)
+        else:
+            pooled = self.readout(node_feats, batch=batch, num_graphs=n_graphs)
+            out = self.final_mlp(pooled)
+
+        return ops.squeeze(out, axis=-1)