k3-node 1.0.0__py3-none-any.whl

k3_node/models/dimenet_utils.py

@@ -0,0 +1,153 @@
+import math
+import numpy as np
+import sympy as sym
+from scipy import special as sp
+from scipy.optimize import brentq
+
+
+def Jn(r, n):
+    return np.sqrt(np.pi / (2 * r)) * sp.jv(n + 0.5, r)
+
+
+def Jn_zeros(n, k):
+    zerosj = np.zeros((n, k), dtype="float32")
+    zerosj[0] = np.arange(1, k + 1) * np.pi
+    points = np.arange(1, k + n) * np.pi
+    racines = np.zeros(k + n - 1, dtype="float32")
+    for i in range(1, n):
+        for j in range(k + n - 1 - i):
+            foo = brentq(Jn, points[j], points[j + 1], (i,))
+            racines[j] = foo
+        points = racines
+        zerosj[i][:k] = racines[:k]
+
+    return zerosj
+
+
+def spherical_bessel_formulas(n):
+    x = sym.symbols("x")
+
+    f = [sym.sin(x) / x]
+    a = sym.sin(x) / x
+    for i in range(1, n):
+        b = sym.diff(a, x) / x
+        f += [sym.simplify(b * (-x) ** i)]
+        a = sym.simplify(b)
+    return f
+
+
+def bessel_basis(n, k):
+    zeros = Jn_zeros(n, k)
+    normalizer = []
+    for order in range(n):
+        normalizer_tmp = []
+        for i in range(k):
+            normalizer_tmp += [0.5 * Jn(zeros[order, i], order + 1) ** 2]
+        normalizer_tmp = 1 / np.array(normalizer_tmp) ** 0.5
+        normalizer += [normalizer_tmp]
+
+    f = spherical_bessel_formulas(n)
+    x = sym.symbols("x")
+    bess_basis = []
+    for order in range(n):
+        bess_basis_tmp = []
+        for i in range(k):
+            bess_basis_tmp += [
+                sym.simplify(
+                    normalizer[order][i] * f[order].subs(x, zeros[order, i] * x)
+                )
+            ]
+        bess_basis += [bess_basis_tmp]
+    return bess_basis
+
+
+def sph_harm_prefactor(k, m):
+    return (
+        (2 * k + 1)
+        * math.factorial(k - abs(m))
+        / (4 * np.pi * math.factorial(k + abs(m)))
+    ) ** 0.5
+
+
+def associated_legendre_polynomials(k, zero_m_only=True):
+    r"""Helper function to calculate Y_l^m."""
+    z = sym.symbols("z")
+    P_l_m = [[0] * (j + 1) for j in range(k)]
+
+    P_l_m[0][0] = 1
+    if k > 0:
+        P_l_m[1][0] = z
+
+        for j in range(2, k):
+            P_l_m[j][0] = sym.simplify(
+                ((2 * j - 1) * z * P_l_m[j - 1][0] - (j - 1) * P_l_m[j - 2][0])
+                / j
+            )
+        if not zero_m_only:
+            for i in range(1, k):
+                P_l_m[i][i] = sym.simplify(
+                    (1 - 2 * i) * P_l_m[i - 1][i - 1] * (1 - z**2) ** 0.5
+                )
+                if i + 1 < k:
+                    P_l_m[i + 1][i] = sym.simplify(
+                        (2 * i + 1) * z * P_l_m[i][i]
+                    )
+                for j in range(i + 2, k):
+                    P_l_m[j][i] = sym.simplify(
+                        (
+                            (2 * j - 1) * z * P_l_m[j - 1][i]
+                            - (i + j - 1) * P_l_m[j - 2][i]
+                        )
+                        / (j - i)
+                    )
+
+    return P_l_m
+
+
+def real_sph_harm(k, zero_m_only=True, spherical_coordinates=True):
+    if not zero_m_only:
+        S_m = [0]
+        C_m = [1]
+        for i in range(1, k):
+            x = sym.symbols("x")
+            y = sym.symbols("y")
+            S_m += [x * S_m[i - 1] + y * C_m[i - 1]]
+            C_m += [x * C_m[i - 1] - y * S_m[i - 1]]
+
+    P_l_m = associated_legendre_polynomials(k, zero_m_only)
+    if spherical_coordinates:
+        theta = sym.symbols("theta")
+        z = sym.symbols("z")
+        for i in range(len(P_l_m)):
+            for j in range(len(P_l_m[i])):
+                if not isinstance(P_l_m[i][j], int):
+                    P_l_m[i][j] = P_l_m[i][j].subs(z, sym.cos(theta))
+        if not zero_m_only:
+            phi = sym.symbols("phi")
+            for i in range(len(S_m)):
+                S_m[i] = S_m[i].subs(x, sym.sin(theta) * sym.cos(phi)).subs(
+                    y, sym.sin(theta) * sym.sin(phi)
+                )
+            for i in range(len(C_m)):
+                C_m[i] = C_m[i].subs(x, sym.sin(theta) * sym.cos(phi)).subs(
+                    y, sym.sin(theta) * sym.sin(phi)
+                )
+
+    Y_func_l_m = [["0"] * (2 * j + 1) for j in range(k)]
+    for i in range(k):
+        Y_func_l_m[i][0] = sym.simplify(sph_harm_prefactor(i, 0) * P_l_m[i][0])
+
+    if not zero_m_only:
+        for i in range(1, k):
+            for j in range(1, i + 1):
+                Y_func_l_m[i][j] = sym.simplify(
+                    2**0.5 * sph_harm_prefactor(i, j) * C_m[j] * P_l_m[i][j]
+                )
+        for i in range(1, k):
+            for j in range(1, i + 1):
+                Y_func_l_m[i][-j] = sym.simplify(
+                    2**0.5 * sph_harm_prefactor(i, -j) * S_m[j] * P_l_m[i][j]
+                )
+
+    return Y_func_l_m
+

k3_node/models/gnnff.py

@@ -0,0 +1,263 @@
+from typing import Optional
+import numpy as np
+import keras
+from keras import ops
+
+from k3_node.layers.aggr import SumAggregation
+from k3_node.layers.pool import radius_graph
+from k3_node.models.schnet import ShiftedSoftplus
+from k3_node.models.dimenet import triplets
+
+
+class GaussianFilter(keras.layers.Layer):
+    r"""Gaussian filter for edge distances.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import GaussianFilter
+
+        dist = np.array([0.9, 1.5, 3.2], dtype="float32")
+        print(tuple(GaussianFilter(start=0.0, stop=5.0, num_gaussians=10)(dist).shape))  # (3, 10)
+        ```
+    """
+    def __init__(self, start: float = 0.0, stop: float = 5.0, num_gaussians: int = 50, **kwargs):
+        super().__init__(**kwargs)
+        offset = np.linspace(start, stop, num_gaussians, dtype=np.float32)
+        diff = float(offset[1] - offset[0])
+        self.coeff = -0.5 / (diff ** 2)
+        self.offset = self.add_weight(
+            name="offset",
+            shape=(num_gaussians,),
+            initializer=keras.initializers.Constant(offset),
+            trainable=False,
+            dtype="float32",
+        )
+
+    def call(self, dist):
+        dist = ops.expand_dims(dist, -1) - ops.expand_dims(self.offset, 0)
+        return ops.exp(self.coeff * ops.power(dist, 2))
+
+
+class NodeBlock(keras.layers.Layer):
+    """Layer ``NodeBlock``.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import NodeBlock, triplets
+
+        edge_index = np.array([[0, 1, 0, 2, 1, 2], [1, 0, 2, 0, 2, 1]])  # a triangle, both directions
+        col, row, idx_i, idx_j, idx_k, idx_kj, idx_ji = triplets(edge_index, num_nodes=3)
+        node_emb = np.random.rand(3, 16).astype("float32")
+        edge_emb = np.random.rand(6, 16).astype("float32")
+
+        block = NodeBlock(hidden_node_channels=16, hidden_edge_channels=16)
+        print(tuple(block(node_emb, edge_emb, row).shape))  # (3, 16): updated node embeddings
+        ```
+    """
+    def __init__(self, hidden_node_channels: int, hidden_edge_channels: int, **kwargs):
+        super().__init__(**kwargs)
+        self.hidden_node_channels = hidden_node_channels
+        self.lin_c1 = keras.layers.Dense(2 * hidden_node_channels)
+        self.bn_c1 = keras.layers.BatchNormalization(momentum=0.9, epsilon=1e-5)
+        self.bn = keras.layers.BatchNormalization(momentum=0.9, epsilon=1e-5)
+        self.sum_aggr = SumAggregation()
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def call(self, node_emb, edge_emb, i, training=None):
+        node_i = ops.take(node_emb, i, axis=0)
+        c1 = ops.concatenate([node_i, edge_emb], axis=1)
+        c1 = self.bn_c1(self.lin_c1(c1), training=training)
+        c1_filter = ops.sigmoid(c1[:, :self.hidden_node_channels])
+        c1_core = ops.tanh(c1[:, self.hidden_node_channels:])
+        num_nodes = ops.shape(node_emb)[0]
+        c1_emb = self.sum_aggr(c1_filter * c1_core, index=i, dim_size=num_nodes)
+        c1_emb = self.bn(c1_emb, training=training)
+        return ops.tanh(node_emb + c1_emb)
+
+
+class EdgeBlock(keras.layers.Layer):
+    """Layer ``EdgeBlock``.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import EdgeBlock, triplets
+
+        edge_index = np.array([[0, 1, 0, 2, 1, 2], [1, 0, 2, 0, 2, 1]])  # a triangle, both directions
+        col, row, idx_i, idx_j, idx_k, idx_kj, idx_ji = triplets(edge_index, num_nodes=3)
+        node_emb = np.random.rand(3, 16).astype("float32")
+        edge_emb = np.random.rand(6, 16).astype("float32")
+
+        block = EdgeBlock(hidden_node_channels=16, hidden_edge_channels=16)
+        out = block(node_emb, edge_emb, row, col, idx_i, idx_j, idx_k, idx_ji, idx_kj)
+        print(tuple(out.shape))  # (6, 16): updated edge embeddings
+        ```
+    """
+    def __init__(self, hidden_node_channels: int, hidden_edge_channels: int, **kwargs):
+        super().__init__(**kwargs)
+        self.hidden_edge_channels = hidden_edge_channels
+        self.lin_c2 = keras.layers.Dense(2 * hidden_edge_channels)
+        self.lin_c3 = keras.layers.Dense(2 * hidden_edge_channels)
+        self.bn_c2 = keras.layers.BatchNormalization(momentum=0.9, epsilon=1e-5)
+        self.bn_c3 = keras.layers.BatchNormalization(momentum=0.9, epsilon=1e-5)
+        self.bn_c2_2 = keras.layers.BatchNormalization(momentum=0.9, epsilon=1e-5)
+        self.bn_c3_2 = keras.layers.BatchNormalization(momentum=0.9, epsilon=1e-5)
+        self.sum_aggr = SumAggregation()
+
+    def build(self, input_shape=None):
+        self.built = True
+
+    def call(
+        self,
+        node_emb,
+        edge_emb,
+        i,
+        j,
+        idx_i,
+        idx_j,
+        idx_k,
+        idx_ji,
+        idx_kj,
+        training=None,
+    ):
+        node_i = ops.take(node_emb, i, axis=0)
+        node_j = ops.take(node_emb, j, axis=0)
+        c2 = node_i * node_j
+        c2 = self.bn_c2(self.lin_c2(c2), training=training)
+        c2_filter = ops.sigmoid(c2[:, :self.hidden_edge_channels])
+        c2_core = ops.tanh(c2[:, self.hidden_edge_channels:])
+        c2_emb = self.bn_c2_2(c2_filter * c2_core, training=training)
+
+        node_idx_i = ops.take(node_emb, idx_i, axis=0)
+        node_idx_j = ops.take(node_emb, idx_j, axis=0)
+        node_idx_k = ops.take(node_emb, idx_k, axis=0)
+        edge_idx_ji = ops.take(edge_emb, idx_ji, axis=0)
+        edge_idx_kj = ops.take(edge_emb, idx_kj, axis=0)
+        c3 = ops.concatenate([node_idx_i, node_idx_j, node_idx_k, edge_idx_ji, edge_idx_kj], axis=1)
+        c3 = self.bn_c3(self.lin_c3(c3), training=training)
+        c3_filter = ops.sigmoid(c3[:, :self.hidden_edge_channels])
+        c3_core = ops.tanh(c3[:, self.hidden_edge_channels:])
+        c3_emb = self.sum_aggr(c3_filter * c3_core, index=idx_ji, dim_size=ops.shape(edge_emb)[0])
+        c3_emb = self.bn_c3_2(c3_emb, training=training)
+
+        return ops.tanh(edge_emb + c2_emb + c3_emb)
+
+
+class GNNFF(keras.layers.Layer):
+    r"""The Graph Neural Network Force Field (GNNFF) from the
+    `"Accurate and scalable graph neural network force field and molecular
+    dynamics with direct force architecture"
+    <https://www.nature.com/articles/s41524-021-00543-3>`_ paper.
+    :class:`GNNFF` directly predicts atomic forces from automatically
+    extracted features of the local atomic environment.
+
+    Args:
+        hidden_node_channels (int): Hidden node embedding size.
+        hidden_edge_channels (int): Hidden edge embedding size.
+        num_layers (int): Number of message passing blocks.
+        cutoff (float, optional): Cutoff distance. (default: 5.0)
+        max_num_neighbors (int, optional): Maximum neighbors per node. (default: 32)
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import GNNFF
+
+        z = np.array([6, 8, 1, 1, 1])  # atomic numbers of a small molecule
+        pos = np.random.rand(5, 3).astype("float32") * 2.0  # 3D coordinates (Angstrom)
+
+        model = GNNFF(hidden_node_channels=16, hidden_edge_channels=16, num_layers=2)
+        forces = model(z, pos)  # predicts a 3D force vector for every atom
+        print(tuple(forces.shape))  # (5, 3)
+        ```
+    """
+    def __init__(
+        self,
+        hidden_node_channels: int,
+        hidden_edge_channels: int,
+        num_layers: int,
+        cutoff: float = 5.0,
+        max_num_neighbors: int = 32,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cutoff = cutoff
+        self.max_num_neighbors = max_num_neighbors
+        self.hidden_node_channels = hidden_node_channels
+        self.hidden_edge_channels = hidden_edge_channels
+        self.num_layers = num_layers
+
+        self.node_emb = keras.Sequential([
+            keras.layers.Embedding(95, hidden_node_channels),
+            ShiftedSoftplus(),
+            keras.layers.Dense(hidden_node_channels),
+            ShiftedSoftplus(),
+            keras.layers.Dense(hidden_node_channels),
+        ])
+        self.edge_emb = GaussianFilter(0.0, 5.0, hidden_edge_channels)
+
+        self.node_blocks = [
+            NodeBlock(hidden_node_channels, hidden_edge_channels)
+            for _ in range(num_layers)
+        ]
+        self.edge_blocks = [
+            EdgeBlock(hidden_node_channels, hidden_edge_channels)
+            for _ in range(num_layers)
+        ]
+
+        self.force_predictor = keras.Sequential([
+            keras.layers.Dense(hidden_edge_channels),
+            ShiftedSoftplus(),
+            keras.layers.Dense(hidden_edge_channels),
+            ShiftedSoftplus(),
+            keras.layers.Dense(1),
+        ])
+        self.sum_aggr = SumAggregation()
+
+    def build(self, input_shape=None):
+        if hasattr(self.node_emb, "built") and not self.node_emb.built:
+            self.node_emb.build((None,))
+        if hasattr(self.edge_emb, "built") and not self.edge_emb.built:
+            self.edge_emb.build((None,))
+        for block in self.node_blocks:
+            if hasattr(block, "built") and not block.built:
+                block.build()
+        for block in self.edge_blocks:
+            if hasattr(block, "built") and not block.built:
+                block.build()
+        if hasattr(self.force_predictor, "built") and not self.force_predictor.built:
+            self.force_predictor.build((None, self.hidden_edge_channels))
+        self.built = True
+
+    def call(self, z, pos, batch=None, training=None):
+        edge_index = radius_graph(
+            pos,
+            r=self.cutoff,
+            batch=batch,
+            max_num_neighbors=self.max_num_neighbors,
+        )
+
+        num_nodes = ops.shape(z)[0]
+        i, j, idx_i, idx_j, idx_k, idx_kj, idx_ji = triplets(edge_index, num_nodes=num_nodes)
+
+        pos_i = ops.take(pos, i, axis=0)
+        pos_j = ops.take(pos, j, axis=0)
+        diff = pos_i - pos_j
+        dist = ops.sqrt(ops.sum(ops.power(diff, 2), axis=-1))
+        unit_vec = diff / ops.expand_dims(dist + 1e-8, -1)
+
+        z = ops.cast(z, "int32")
+        node_emb = self.node_emb(z)
+        edge_emb = self.edge_emb(dist)
+
+        for node_block, edge_block in zip(self.node_blocks, self.edge_blocks):
+            node_emb = node_block(node_emb, edge_emb, i, training=training)
+            edge_emb = edge_block(node_emb, edge_emb, i, j, idx_i, idx_j, idx_k, idx_ji, idx_kj, training=training)
+
+        force = self.force_predictor(edge_emb) * unit_vec
+        return self.sum_aggr(force, index=i, dim_size=num_nodes)
+