k3-node 1.0.0__py3-none-any.whl

k3_node/applications/materials/core.py

@@ -0,0 +1,351 @@
+"""Core layers and mathematical primitives for materials models."""
+
+from __future__ import annotations
+
+import math
+from typing import Sequence, Callable, Optional, Union
+import keras
+from keras import layers, ops
+import numpy as np
+from k3_node.ops.segment import segment_sum
+
+
+class SoftPlus2(layers.Layer):
+    """SoftPlus2 activation: log(exp(x) + 1) - log(2). Zero at the origin."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.shift = float(math.log(2.0))
+
+    def call(self, x):
+        return ops.softplus(x) - self.shift
+
+
+class SoftExponential(layers.Layer):
+    """Soft exponential activation with learnable alpha."""
+
+    def __init__(self, alpha: float = 0.0, **kwargs):
+        super().__init__(**kwargs)
+        self.init_alpha = float(alpha)
+        self._eps = 1e-6
+
+    def build(self, input_shape=None):
+        self.alpha = self.add_weight(
+            name="alpha",
+            shape=(),
+            initializer=keras.initializers.Constant(self.init_alpha),
+            trainable=True,
+            dtype="float32",
+        )
+        super().build(input_shape)
+
+    def call(self, x):
+        alpha = self.alpha
+        near_zero = ops.abs(alpha) < self._eps
+        safe_alpha = ops.where(near_zero, 1.0, alpha)
+
+        neg_log_arg = ops.where(alpha < 0.0, 1.0 - alpha * (x + alpha), ops.ones_like(x))
+        neg_log_arg = ops.maximum(neg_log_arg, self._eps)
+        neg = -ops.log(neg_log_arg) / safe_alpha
+        pos = ops.expm1(safe_alpha * x) / safe_alpha + safe_alpha
+
+        out = ops.where(alpha < 0.0, neg, pos)
+        return ops.where(near_zero, x, out)
+
+
+def get_activation(act: Union[str, Callable, layers.Layer, None]):
+    """Resolve activation specification into a callable or Keras layer."""
+    if act is None:
+        return lambda x: x
+    if isinstance(act, str):
+        act_lower = act.lower()
+        if act_lower in ("swish", "silu"):
+            return ops.silu
+        elif act_lower == "softplus2":
+            return SoftPlus2()
+        elif act_lower == "softexp":
+            return SoftExponential()
+        elif act_lower == "softplus":
+            return ops.softplus
+        elif act_lower == "tanh":
+            return ops.tanh
+        elif act_lower == "sigmoid":
+            return ops.sigmoid
+        elif act_lower == "relu":
+            return ops.relu
+        elif act_lower in ("identity", "linear", "none"):
+            return lambda x: x
+        return keras.activations.get(act)
+    if isinstance(act, layers.Layer):
+        return act
+    return act
+
+
+class MLP(layers.Layer):
+    """Multi-layer perceptron compatible with multi-backend Keras 3."""
+
+    def __init__(
+        self,
+        dims: Sequence[int],
+        activation: Union[str, Callable, layers.Layer, None] = "swish",
+        activate_last: bool = False,
+        bias_last: bool = True,
+        use_bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.dims = list(dims)
+        self.activate_last = activate_last
+        self.bias_last = bias_last
+        self.use_bias = use_bias
+        self.act_fn = get_activation(activation)
+
+        self.dense_layers = []
+        for i in range(len(self.dims) - 1):
+            is_last = (i == len(self.dims) - 2)
+            bias = self.bias_last if is_last else self.use_bias
+            dense = layers.Dense(self.dims[i + 1], use_bias=bias)
+            if self.dims[i] is not None:
+                dense.build((None, self.dims[i]))
+            self.dense_layers.append(dense)
+
+    def call(self, x):
+        for i, layer in enumerate(self.dense_layers):
+            x = layer(x)
+            is_last = (i == len(self.dense_layers) - 1)
+            if not is_last or self.activate_last:
+                x = self.act_fn(x)
+        return x
+
+
+class GatedMLP(layers.Layer):
+    """Gated multi-layer perceptron: layer(x) * sigmoid(gate(x))."""
+
+    def __init__(
+        self,
+        in_feats: int,
+        dims: Sequence[int],
+        activate_last: bool = True,
+        use_bias: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_feats = in_feats
+        self.dims = [in_feats, *dims]
+        self.activate_last = activate_last
+        self.use_bias = use_bias
+
+        self.val_layers = []
+        self.gate_layers = []
+        for i in range(len(self.dims) - 1):
+            out_dim = self.dims[i + 1]
+            val = layers.Dense(out_dim, use_bias=use_bias)
+            gate = layers.Dense(out_dim, use_bias=use_bias)
+            if self.dims[i] is not None:
+                val.build((None, self.dims[i]))
+                gate.build((None, self.dims[i]))
+            self.val_layers.append(val)
+            self.gate_layers.append(gate)
+
+    def call(self, x):
+        h_val = x
+        h_gate = x
+        for i in range(len(self.val_layers)):
+            is_last = (i == len(self.val_layers) - 1)
+            h_val = self.val_layers[i](h_val)
+            h_gate = self.gate_layers[i](h_gate)
+            if not is_last:
+                h_val = ops.silu(h_val)
+                h_gate = ops.silu(h_gate)
+            else:
+                if self.activate_last:
+                    h_val = ops.silu(h_val)
+                h_gate = ops.sigmoid(h_gate)
+        return h_val * h_gate
+
+
+class EmbeddingBlock(layers.Layer):
+    """Embeddings for nodes (atoms), edges (bonds), and global states."""
+
+    def __init__(
+        self,
+        degree_rbf: int,
+        dim_node_embedding: int,
+        dim_edge_embedding: Optional[int] = None,
+        dim_state_embedding: Optional[int] = None,
+        dim_state_feats: Optional[int] = None,
+        ntypes_node: Optional[int] = None,
+        ntypes_state: Optional[int] = None,
+        include_state: bool = False,
+        activation: Union[str, Callable, None] = "swish",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.degree_rbf = degree_rbf
+        self.dim_node_embedding = dim_node_embedding
+        self.dim_edge_embedding = dim_edge_embedding
+        self.dim_state_embedding = dim_state_embedding
+        self.dim_state_feats = dim_state_feats
+        self.ntypes_node = ntypes_node
+        self.ntypes_state = ntypes_state
+        self.include_state = include_state
+
+        if ntypes_node is not None:
+            self.layer_node_embedding = layers.Embedding(ntypes_node, dim_node_embedding)
+        else:
+            self.layer_node_embedding = layers.Dense(dim_node_embedding, use_bias=False)
+
+        if dim_edge_embedding is not None:
+            self.layer_edge_embedding = MLP([degree_rbf, dim_edge_embedding], activation=activation, activate_last=True)
+        else:
+            self.layer_edge_embedding = None
+
+        if include_state:
+            if ntypes_state is not None and dim_state_embedding is not None:
+                self.layer_state_embedding = layers.Embedding(ntypes_state, dim_state_embedding)
+            elif dim_state_feats is not None:
+                self.layer_state_embedding = layers.Dense(dim_state_feats, use_bias=False)
+            elif dim_state_embedding is not None:
+                self.layer_state_embedding = layers.Dense(dim_state_embedding, use_bias=False)
+            else:
+                self.layer_state_embedding = None
+        else:
+            self.layer_state_embedding = None
+
+    def call(self, node_attr, edge_attr, state_attr=None):
+        if self.ntypes_node is not None:
+            node_idx = ops.clip(ops.cast(node_attr, "int32"), 0, self.ntypes_node - 1)
+            node_feat = self.layer_node_embedding(node_idx)
+        else:
+            node_feat = self.layer_node_embedding(node_attr)
+
+        edge_feat = self.layer_edge_embedding(edge_attr) if self.layer_edge_embedding is not None else edge_attr
+
+        state_feat = None
+        if self.include_state and state_attr is not None and self.layer_state_embedding is not None:
+            if self.ntypes_state is not None:
+                state_idx = ops.clip(ops.cast(state_attr, "int32"), 0, self.ntypes_state - 1)
+                state_feat = self.layer_state_embedding(state_idx)
+            else:
+                state_feat = self.layer_state_embedding(state_attr)
+        elif self.include_state and state_attr is not None:
+            state_feat = state_attr
+
+        return node_feat, edge_feat, state_feat
+
+
+# --- Tensor math utilities for TensorNet ---
+
+def vector_to_skewtensor(vector):
+    """Create skew-symmetric 3x3 tensor from a 3D vector.
+
+    ```
+    [0, -v_z, v_y]
+    [v_z, 0, -v_x]
+    [-v_y, v_x, 0]
+    ```
+    """
+    vector = ops.convert_to_tensor(vector)
+    vx = vector[..., 0]
+    vy = vector[..., 1]
+    vz = vector[..., 2]
+    zero = ops.zeros_like(vx)
+
+    row0 = ops.stack([zero, -vz, vy], axis=-1)
+    row1 = ops.stack([vz, zero, -vx], axis=-1)
+    row2 = ops.stack([-vy, vx, zero], axis=-1)
+    return ops.stack([row0, row1, row2], axis=-2)
+
+
+def vector_to_symtensor(vector):
+    """Create symmetric traceless tensor from outer product of 3D vector with itself."""
+    vector = ops.convert_to_tensor(vector)
+    v_col = ops.expand_dims(vector, axis=-1)
+    v_row = ops.expand_dims(vector, axis=-2)
+    outer = ops.matmul(v_col, v_row)
+    # trace = outer[..., 0, 0] + outer[..., 1, 1] + outer[..., 2, 2]
+    trace = outer[..., 0, 0] + outer[..., 1, 1] + outer[..., 2, 2]
+    eye3 = ops.eye(3, dtype=outer.dtype)
+    scalars = ops.expand_dims(ops.expand_dims(trace / 3.0, axis=-1), axis=-1) * eye3
+    ndim = len(outer.shape)
+    axes = list(range(ndim - 2)) + [ndim - 1, ndim - 2]
+    sym = 0.5 * (outer + ops.transpose(outer, axes=axes))
+    return sym - scalars
+
+
+def decompose_tensor(tensor):
+    """Decompose 3x3 Cartesian tensor into scalar (I), skew-symmetric (A), and symmetric traceless (S)."""
+    tensor = ops.convert_to_tensor(tensor)
+    trace = tensor[..., 0, 0] + tensor[..., 1, 1] + tensor[..., 2, 2]
+    eye3 = ops.eye(3, dtype=tensor.dtype)
+    scalars = ops.expand_dims(ops.expand_dims(trace / 3.0, axis=-1), axis=-1) * eye3
+    ndim = len(tensor.shape)
+    axes = list(range(ndim - 2)) + [ndim - 1, ndim - 2]
+    transposed = ops.transpose(tensor, axes=axes)
+    skew = 0.5 * (tensor - transposed)
+    sym = 0.5 * (tensor + transposed)
+    traceless = sym - scalars
+    return scalars, skew, traceless
+
+
+def new_radial_tensor(scalars, skew, traceless, f_I, f_A, f_S):
+    """Multiply irreducible tensor components by radial invariant features."""
+    scalars_out = ops.expand_dims(ops.expand_dims(f_I, axis=-1), axis=-1) * scalars
+    skew_out = ops.expand_dims(ops.expand_dims(f_A, axis=-1), axis=-1) * skew
+    traceless_out = ops.expand_dims(ops.expand_dims(f_S, axis=-1), axis=-1) * traceless
+    return scalars_out, skew_out, traceless_out
+
+
+def tensor_norm(tensor):
+    """Computes Frobenius norm squared across last two dimensions (3, 3)."""
+    return ops.sum(tensor ** 2, axis=(-2, -1))
+
+
+def scatter_add(x, index, num_segments: int):
+    """Scatter sum x elements into segments indicated by index."""
+    index = ops.cast(index, "int32")
+    return segment_sum(x, index, num_segments=num_segments)
+
+
+def scatter_mean(x, index, num_segments: int):
+    """Scatter mean x elements into segments indicated by index."""
+    index = ops.cast(index, "int32")
+    sums = segment_sum(x, index, num_segments=num_segments)
+    ones = ops.ones_like(x[..., :1])
+    counts = segment_sum(ones, index, num_segments=num_segments)
+    counts = ops.maximum(counts, 1.0)
+    return sums / counts
+
+
+def infer_num_graphs(batch=None, num_graphs=None, state_attr=None):
+    """Infer the number of graphs in a batch across backends (JAX, PyTorch, TF)."""
+    if num_graphs is not None:
+        return num_graphs
+    if state_attr is not None:
+        shape = ops.shape(state_attr)
+        if len(shape) == 1:
+            return 1
+        return shape[0]
+    if batch is None:
+        return 1
+    shape = ops.shape(batch)
+    if len(shape) > 0 and shape[0] == 0:
+        return 0
+    try:
+        val = ops.convert_to_numpy(ops.max(batch))
+        return int(val) + 1
+    except Exception:
+        pass
+    try:
+        if hasattr(batch, "__getitem__"):
+            val = batch[-1]
+            if hasattr(val, "item"):
+                return int(val.item()) + 1
+            return int(val) + 1
+    except Exception:
+        pass
+    try:
+        return ops.cast(ops.max(batch) + 1, "int32")
+    except Exception:
+        return 1
+

k3_node/applications/materials/grace.py

@@ -0,0 +1,246 @@
+"""Multi-backend Keras 3 implementation of GRACE (Graph Atomic Cluster Expansion)."""
+
+from __future__ import annotations
+
+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 (
+    ChebyshevRadialBasis,
+    compute_pair_vector_and_distance,
+    polynomial_cutoff,
+)
+from .so3net import RealSphericalHarmonics
+from .readout import ReduceReadOut
+
+
+class GraceSPBasis(layers.Layer):
+    """Single-particle ACE basis aggregation for atomic clusters.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import GraceSPBasis
+
+        edge_index = np.array([[0, 1, 1, 2, 2, 3, 3, 0], [1, 0, 2, 1, 3, 2, 0, 3]])
+        node_type = np.array([6, 8, 1, 6])
+        rad_basis = np.random.rand(8, 6).astype("float32")  # radial basis per bond
+        sh_basis = np.random.rand(8, 9).astype("float32")  # spherical harmonics per bond (lmax=2)
+        layer = GraceSPBasis(n_rad_base=6, lmax=2, embedding_size=8)
+        print(tuple(layer(edge_index, node_type, rad_basis, sh_basis, num_nodes=4).shape))  # (4, 9, 8): atomic basis A_i
+        ```
+    """
+
+    def __init__(
+        self,
+        n_rad_base: int = 6,
+        lmax: int = 2,
+        embedding_size: int = 8,
+        ntypes_node: int = 95,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.n_rad_base = n_rad_base
+        self.lmax = lmax
+        self.num_sh = (lmax + 1) ** 2
+        self.elem_emb = layers.Embedding(ntypes_node, embedding_size)
+        self.radial_proj = layers.Dense(embedding_size, use_bias=False)
+
+    def call(self, edge_index, node_type, rad_basis, sh_basis, num_nodes: Optional[int] = None):
+        src = ops.cast(edge_index[0], "int32")
+        dst = ops.cast(edge_index[1], "int32")
+
+        # Neighbor chemical indicator
+        neigh_z = ops.take(node_type, dst, axis=0)
+        c_j = self.elem_emb(ops.cast(neigh_z, "int32"))
+
+        # Radial modulation
+        R_nl = self.radial_proj(rad_basis)
+        radial_atom = c_j * R_nl
+
+        # Outer product with spherical harmonics: [E, num_sh, emb_size]
+        sp_msg = ops.expand_dims(sh_basis, axis=-1) * ops.expand_dims(radial_atom, axis=1)
+
+        # Aggregate into central atom
+        A_i = scatter_add(sp_msg, src, num_segments=num_nodes)
+        return A_i
+
+
+class GraceACEStack(layers.Layer):
+    """Multi-order Atomic Cluster Expansion stack accumulating rotational invariants.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import GraceACEStack
+
+        A_i = np.random.rand(4, 9, 8).astype("float32")  # atomic basis from GraceSPBasis
+        print(tuple(GraceACEStack(embedding_size=8, max_order=2)(A_i).shape))  # (4, 16): body-order invariants
+        ```
+    """
+
+    def __init__(
+        self,
+        embedding_size: int = 8,
+        max_order: int = 3,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.embedding_size = embedding_size
+        self.max_order = max_order
+
+    def call(self, A_i):
+        A_00 = A_i[:, 0, :]
+        invariants = [A_00]
+
+        if self.max_order >= 2:
+            A2 = ops.sum(A_i ** 2, axis=1)
+            invariants.append(A2)
+
+        if self.max_order >= 3:
+            A3 = A_00 * ops.sum(A_i ** 2, axis=1)
+            invariants.append(A3)
+
+        return ops.concatenate(invariants, axis=-1)
+
+
+class GRACE(keras.Model):
+    """Graph Atomic Cluster Expansion (GRACE) foundational interatomic potential.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import GRACE
+
+        # 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 = GRACE(cutoff=5.0, n_rad_base=6, lmax=2, embedding_size=8, max_order=2, nblocks=2,
+                      readout_hidden=(16,))
+        energy = model(structure)  # predicted property (e.g. energy) of the structure
+        print(tuple(energy.shape))  # (1,)
+        ```
+    """
+
+    def __init__(
+        self,
+        cutoff: float = 5.0,
+        n_rad_base: int = 6,
+        lmax: int = 2,
+        embedding_size: int = 8,
+        max_order: int = 3,
+        nblocks: int = 2,
+        readout_hidden: Sequence[int] = (32,),
+        ntypes_node: int = 95,
+        activation_type: str = "swish",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cutoff = cutoff
+        self.nblocks = nblocks
+
+        self.chebyshev = ChebyshevRadialBasis(nfunc=n_rad_base, cutoff=cutoff, cutoff_exponent=5)
+        self.sh = RealSphericalHarmonics(lmax=lmax)
+
+        self.sp_bases = []
+        self.ace_stacks = []
+        self.readouts = []
+        inv_dim = embedding_size * min(max_order, 3)
+
+        for _ in range(nblocks):
+            self.sp_bases.append(
+                GraceSPBasis(
+                    n_rad_base=n_rad_base,
+                    lmax=lmax,
+                    embedding_size=embedding_size,
+                    ntypes_node=ntypes_node,
+                )
+            )
+            self.ace_stacks.append(
+                GraceACEStack(
+                    embedding_size=embedding_size,
+                    max_order=max_order,
+                )
+            )
+            self.readouts.append(
+                MLP([inv_dim, *readout_hidden, 1], activation=activation_type, activate_last=False)
+            )
+
+        self.graph_pool = ReduceReadOut(op="sum")
+
+    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)
+        rad_basis = self.chebyshev(bond_dists)
+        sh_basis = self.sh(vec)
+
+        total_atomic_energy = ops.zeros((num_nodes, 1), dtype=pos.dtype)
+        for i in range(self.nblocks):
+            A_i = self.sp_bases[i](edge_index, node_type, rad_basis, sh_basis, num_nodes=num_nodes)
+            invariants = self.ace_stacks[i](A_i)
+            e_block = self.readouts[i](invariants)
+            total_atomic_energy = total_atomic_energy + e_block
+
+        total_energy = self.graph_pool(total_atomic_energy, batch=batch, num_graphs=n_graphs)
+        return ops.squeeze(total_energy, axis=-1)
+