fennol.models.embeddings.aimnet

  1import jax
  2import jax.numpy as jnp
  3import flax.linen as nn
  4from typing import Sequence, Dict, Callable, ClassVar,Union
  5import numpy as np
  6from ...utils.periodic_table import PERIODIC_TABLE
  7from ..misc.encodings import SpeciesEncoding
  8from ..misc.nets import FullyConnectedNet
  9import dataclasses
 10
 11
 12class AIMNet(nn.Module):
 13    """Atom-In-Molecule Network message-passing embedding
 14
 15    FID : AIMNET
 16
 17    ### Reference
 18    Roman Zubatyuk et al. ,Accurate and transferable multitask prediction of chemical properties with an atoms-in-molecules neural network.Sci. Adv.5,eaav6490(2019).DOI:10.1126/sciadv.aav6490
 19    """
 20
 21    _graphs_properties: Dict
 22    graph_angle_key: str
 23    """ The key in the input dictionary that corresponds to the angular graph. """
 24    nlayers: int = 3
 25    """ The number of message-passing layers."""
 26    zmax: int = 86
 27    """ The maximum atomic number to allocate AFV."""
 28    radial_eta: float = 16.0
 29    """ Controls the width of the gaussian sensity functions in radial AEV."""
 30    angular_eta: float = 8.0
 31    """ Controls the width of the gaussian sensity functions in angular AEV."""
 32    radial_dist_divisions: int = 16
 33    """ Number of basis function to encode ditance in radial AEV."""
 34    angular_dist_divisions: int = 4
 35    """ Number of basis function to encode ditance in angular AEV."""
 36    zeta: float = 32.0
 37    """ The power parameter in angle embedding."""
 38    angle_sections: int = 4
 39    """ The number of angle sections."""
 40    radial_start: float = 0.8
 41    """ The starting distance in radial AEV."""
 42    angular_start: float = 0.8
 43    """ The starting distance in angular AEV."""
 44    embedding_key: str = "embedding"
 45    """ The key to use for the output embedding in the returned dictionary."""
 46    graph_key: str = "graph"
 47    """ The key in the input dictionary that corresponds to the radial graph."""
 48    keep_all_layers: bool = False
 49    """ If True, the output will contain the embeddings from all layers."""
 50
 51    activation: Union[Callable, str] = "swish"
 52    """ The activation function to use."""
 53    combination_neurons: Sequence[int] = dataclasses.field(
 54        default_factory=lambda: [256, 128, 16]
 55    )
 56    """ The number of neurons in the AFV combination network."""
 57
 58    embedding_neurons: Sequence[int] = dataclasses.field(
 59        default_factory=lambda: [512, 256, 256]
 60    )
 61    """ The number of neurons in the embedding network."""
 62    interaction_neurons: Sequence[int] = dataclasses.field(
 63        default_factory=lambda: [256, 256, 128]
 64    )
 65    """ The number of neurons in the interaction network."""
 66    afv_neurons: Sequence[int] = dataclasses.field(
 67        default_factory=lambda: [256, 256, 16]
 68    )
 69    """ The number of neurons in the AFV update network. The last number of neurons defines the size of AFV."""
 70
 71    FID: ClassVar[str] = "AIMNET"
 72
 73    @nn.compact
 74    def __call__(self, inputs):
 75        """Forward pass of the AIMNet model."""
 76        species = inputs["species"]
 77
 78        # species encoding (AFV)
 79        afv_dim = self.afv_neurons[-1]
 80        afv = SpeciesEncoding(dim=afv_dim, zmax=self.zmax, encoding="random")(species)
 81
 82        # Radial graph
 83        graph = inputs[self.graph_key]
 84        distances = graph["distances"]
 85        switch = graph["switch"]
 86        edge_src = graph["edge_src"]
 87        edge_dst = graph["edge_dst"]
 88
 89        # Radial AEV
 90        cutoff = self._graphs_properties[self.graph_key]["cutoff"]
 91        shiftR = jnp.asarray(
 92            np.linspace(self.radial_start, cutoff, self.radial_dist_divisions + 1)[
 93                None, :-1
 94            ],
 95            dtype=distances.dtype,
 96        )
 97        x2 = self.radial_eta * (distances[:, None] - shiftR) ** 2
 98        radial_terms = 0.25 * jnp.exp(-x2) * switch[:, None]
 99
100        # Angular graph
101        graph = inputs[self.graph_angle_key]
102        edge_dst_a = graph["edge_dst"]
103        angles = graph["angles"]
104        distances = graph["distances"]
105        central_atom = graph["central_atom"]
106        angle_src, angle_dst = graph["angle_src"], graph["angle_dst"]
107        angle_atom_1 = edge_dst_a[angle_src]
108        angle_atom_2 = edge_dst_a[angle_dst]
109        switch = graph["switch"]
110        d12 = 0.5 * (distances[angle_src] + distances[angle_dst])[:, None]
111
112        # Angular AEV parameters
113        angular_cutoff = self._graphs_properties[self.graph_angle_key]["cutoff"]
114        angle_start = np.pi / (2 * self.angle_sections)
115        shiftZ = jnp.asarray(
116            (np.linspace(0, np.pi, self.angle_sections + 1) + angle_start)[None, :-1],
117            dtype=distances.dtype,
118        )
119        shiftA = jnp.asarray(
120            np.linspace(
121                self.angular_start, angular_cutoff, self.angular_dist_divisions + 1
122            )[None, :-1],
123            dtype=distances.dtype,
124        )
125
126        # Angular AEV
127        factor1 = (0.5 + 0.5 * jnp.cos(angles[:, None] - shiftZ)) ** self.zeta
128        factor2 = jnp.exp(-self.angular_eta * (d12 - shiftA) ** 2)
129        angular_terms = (
130            (factor1[:, None, :] * factor2[:, :, None]).reshape(
131                -1, self.angle_sections * self.angular_dist_divisions
132            )
133            * 2
134            * (switch[angle_src] * switch[angle_dst])[:, None]
135        )
136
137        if self.keep_all_layers:
138            mis = []
139        for layer in range(self.nlayers):
140            # combine pair info
141            Gij = (radial_terms[:, None, :] * afv[edge_dst, :, None]).reshape(
142                radial_terms.shape[0], -1
143            )
144            Gri = jax.ops.segment_sum(Gij, edge_src, species.shape[0])
145
146            # combine triplet info
147            afv1 = afv[angle_atom_1]
148            afv2 = afv[angle_atom_2]
149            afv12 = jnp.concatenate((afv1 * afv2, afv1 + afv2), axis=-1)
150            afv_ang = FullyConnectedNet(
151                self.combination_neurons,
152                activation=self.activation,
153                name=f"combination_net_{layer}",
154            )(afv12)
155            Gijk = (angular_terms[:, None, :] * afv_ang[:, :, None]).reshape(
156                angular_terms.shape[0], -1
157            )
158            Gai = jax.ops.segment_sum(Gijk, central_atom, species.shape[0])
159
160            # environment field
161            fi = FullyConnectedNet(
162                self.embedding_neurons,
163                activation=self.activation,
164                name=f"embedding_net_{layer}",
165            )(jnp.concatenate((Gri, Gai), axis=-1))
166
167            # update AFV
168            dafv = FullyConnectedNet(
169                self.afv_neurons,
170                activation=self.activation,
171                name=f"afv_update_net_{layer}",
172            )(fi)
173            afv = afv + dafv
174
175            if self.keep_all_layers or layer == self.nlayers - 1:
176                # embedding
177                mi = FullyConnectedNet(
178                    self.interaction_neurons,
179                    activation=self.activation,
180                    name=f"interaction_net_{layer}",
181                )(fi)
182                if self.keep_all_layers:
183                    mis.append(mi)
184
185        embedding_key = (
186            self.embedding_key if self.embedding_key is not None else self.name
187        )
188
189        output = {**inputs, embedding_key: mi, embedding_key + "_afv": afv}
190        if self.keep_all_layers:
191            output[embedding_key + "_layers"] = jnp.stack(mis, axis=1)
192
193        return output