fennol.models.embeddings.raster

  1import jax
  2import jax.numpy as jnp
  3import flax.linen as nn
  4from typing import Sequence, Dict, Union, ClassVar, Optional
  5import numpy as np
  6from ...utils.periodic_table import D3_COV_RADII
  7import dataclasses
  8from ..misc.encodings import RadialBasis, SpeciesEncoding
  9from ...utils.spherical_harmonics import generate_spherical_harmonics
 10from ..misc.e3 import  ChannelMixing
 11from ..misc.nets import FullyConnectedNet,BlockIndexNet
 12from ...utils.activations import activation_from_str
 13from ...utils.atomic_units import au
 14from ...utils.initializers import initializer_from_str
 15
 16class RaSTER(nn.Module):
 17    """ Range-Separated Transformer with Equivariant Representations
 18
 19    FID : RASTER
 20
 21    """
 22
 23    _graphs_properties: Dict
 24    dim: int = 176
 25    """The dimension of the output embedding."""
 26    nlayers: int = 2
 27    """The number of message-passing layers."""
 28    att_dim: int = 16
 29    """The dimension of the attention heads."""
 30    scal_heads: int = 16
 31    """The number of scalar attention heads."""
 32    tens_heads: int = 4
 33    """The number of tensor attention heads."""
 34    lmax: int = 3
 35    """The maximum angular momentum to consider."""
 36    normalize_vec: bool = True
 37    """Whether to normalize the vector features before computing spherical harmonics."""
 38    att_activation: str = "identity"
 39    """The activation function to use for the attention coefficients."""
 40    activation: str = "swish"
 41    """The activation function to use for the update network."""
 42    update_hidden: Sequence[int] = ()
 43    """The hidden layers for the update network."""
 44    update_bias: bool = True
 45    """Whether to use bias in the update network."""
 46    positional_activation: str = "swish"
 47    """The activation function to use for the positional embedding network."""
 48    positional_bias: bool = True
 49    """Whether to use bias in the positional embedding network."""
 50    switch_before_net: bool = False
 51    """Whether to apply the switch function to the radial basis before the edge neural network."""
 52    ignore_parity: bool = False
 53    """Whether to ignore the parity of the spherical harmonics when constructing the relative positional encoding."""
 54    additive_positional: bool = False
 55    """Whether to use additive relative positional encoding. If False, multiplicative relative positional encoding is used."""
 56    edge_value: bool = False
 57    """Whether to use edge values in the attention mechanism."""
 58    layer_normalization: bool = True
 59    """Whether to use layer normalization of atomic embeddings."""
 60    graph_key: str = "graph"
 61    """ The key in the input dictionary that corresponds to the radial graph."""
 62    embedding_key: str = "embedding"
 63    """ The key in the output dictionary that corresponds to the embedding."""
 64    radial_basis: dict = dataclasses.field(
 65        default_factory=lambda: {"start": 0.8, "basis": "gaussian", "dim": 16}
 66    )
 67    """The dictionary of parameters for radial basis functions. See `fennol.models.misc.encodings.RadialBasis`."""
 68    species_encoding: str | dict = dataclasses.field(
 69        default_factory=lambda: {"dim": 16, "trainable": True, "encoding": "random"}
 70    )
 71    """The dictionary of parameters for species encoding. See `fennol.models.misc.encodings.SpeciesEncoding`."""
 72    graph_lode: Optional[str] = None
 73    """The key in the input dictionary that corresponds to the long-range graph."""
 74    lmax_lode: int = 0
 75    """The maximum angular momentum for the long-range features."""
 76    lode_rshort: Optional[float] = None
 77    """The short-range cutoff for the long-range features."""
 78    lode_dshort: float = 2.0
 79    """The width of the short-range cutoff for the long-range features."""
 80    lode_extra_powers: Sequence[int] = ()
 81    """The extra powers to include in the long-range features."""
 82    a_lode: float = -1.0
 83    """The damping parameter for the long-range features. If negative, the damping is trainable with initial value abs(a_lode)."""
 84    block_index_key: Optional[str] = None
 85    """The key in the input dictionary that corresponds to the block index for the MoE network. If None, a normal neural network is used."""
 86    lode_channels: int = 1
 87    """The number of channels for the long-range features."""
 88    switch_cov_start: float = 0.5
 89    """The start of close-range covalent switch (in units of covalent radii)."""
 90    switch_cov_end: float = 0.6
 91    """The end of close-range covalent switch (in units of covalent radii)."""
 92    normalize_keys: bool = False
 93    """Whether to normalize queries and keys in the attention mechanism."""
 94    normalize_components: bool = False
 95    """Whether to normalize the components before the update network."""
 96    keep_all_layers: bool = False
 97    """Whether to return the stacked scalar embeddings from all message-passing layers."""
 98    kernel_init: Optional[str] = None
 99    
100    FID: ClassVar[str] = "RASTER"
101
102    @nn.compact
103    def __call__(self, inputs):
104        species = inputs["species"]
105
106        ## SETUP LAYER NORMALIZATION
107        def _layer_norm(x):
108            mu = jnp.mean(x, axis=-1, keepdims=True)
109            dx = x - mu
110            var = jnp.mean(dx**2, axis=-1, keepdims=True)
111            sig = (1.0e-6 + var) ** (-0.5)
112            return dx * sig
113
114        if self.layer_normalization:
115            layer_norm = _layer_norm
116        else:
117            layer_norm = lambda x: x
118        
119        if self.normalize_keys:
120            ln_qk = _layer_norm
121        else:
122            ln_qk = lambda x: x
123
124        kernel_init = initializer_from_str(self.kernel_init)
125
126        ## SPECIES ENCODING
127        if isinstance(self.species_encoding, str):
128            Zi = inputs[self.species_encoding]
129        else:
130            Zi = SpeciesEncoding(**self.species_encoding)(species)
131
132        ## INITIALIZE SCALAR FEATURES
133        xi = layer_norm(nn.Dense(self.dim, use_bias=False,name="species_linear",kernel_init=kernel_init)(Zi))
134
135        # RADIAL GRAPH
136        graph = inputs[self.graph_key]
137        distances = graph["distances"]
138        switch = graph["switch"]
139        edge_src = graph["edge_src"]
140        edge_dst = graph["edge_dst"]
141        vec = (
142            graph["vec"] / graph["distances"][:, None]
143            if self.normalize_vec
144            else graph["vec"]
145        )
146        ## CLOSE-RANGE SWITCH
147        use_switch_cov = False
148        if self.switch_cov_end > 0 and self.switch_cov_start > 0:
149            use_switch_cov = True
150            assert self.switch_cov_start < self.switch_cov_end, f"switch_cov_start {self.switch_cov_start} must be smaller than switch_cov_end {self.switch_cov_end}"
151            assert self.switch_cov_start > 0 and self.switch_cov_end < 1, f"switch_cov_start {self.switch_cov_start} and switch_cov_end {self.switch_cov_end} must be between 0 and 1"
152            rc = jnp.array(D3_COV_RADII*au.ANG)[species]
153            rcij = rc[edge_src] + rc[edge_dst]
154            rstart = rcij * self.switch_cov_start
155            rend = rcij * self.switch_cov_end
156            switch_short = (distances >= rend) + 0.5*(1-jnp.cos(jnp.pi*(distances - rstart)/(rend-rstart)))*(distances > rstart)*(distances < rend)
157            switch = switch * switch_short
158
159        ## COMPUTE SPHERICAL HARMONICS ON EDGES
160        Yij = generate_spherical_harmonics(lmax=self.lmax, normalize=False)(vec)[:,None,:]
161        nrep = np.array([2 * l + 1 for l in range(self.lmax + 1)])
162        ls = np.arange(self.lmax + 1).repeat(nrep)
163            
164        parity = jnp.array((-1) ** ls[None,None,:])
165        if self.ignore_parity:
166            parity = -jnp.ones_like(parity)
167
168        ## INITIALIZE TENSOR FEATURES
169        Vi = 0. #jnp.zeros((Zi.shape[0],self.tens_heads, Yij.shape[1]))
170
171        # RADIAL BASIS
172        cutoff = self._graphs_properties[self.graph_key]["cutoff"]
173        radial_terms = RadialBasis(
174            **{
175                **self.radial_basis,
176                "end": cutoff,
177                "name": f"RadialBasis",
178            }
179        )(distances)
180        if self.switch_before_net:
181            radial_terms = radial_terms * switch[:, None]
182        elif use_switch_cov:
183            radial_terms = radial_terms * switch_short[:, None]
184
185        ## INITIALIZE LODE
186        do_lode = self.graph_lode is not None
187        if do_lode:
188            ## LONG-RANGE GRAPH
189            graph_lode = inputs[self.graph_lode]
190            switch_lode = graph_lode["switch"][:, None]
191            edge_src_lr, edge_dst_lr = graph_lode["edge_src"], graph_lode["edge_dst"]
192            r = graph_lode["distances"][:, None]
193            rc = self._graphs_properties[self.graph_lode]["cutoff"]
194
195            lmax_lr = self.lmax_lode
196            equivariant_lode = lmax_lr > 0
197            assert lmax_lr >= 0, f"lmax_lode must be >= 0, got {lmax_lr}"
198            assert (
199                lmax_lr <= self.lmax
200            ), f"lmax_lode must be <= lmax for multipole interaction, got {lmax_lr} > {self.lmax}"
201            nrep_lr = np.array([2 * l + 1 for l in range(lmax_lr + 1)], dtype=np.int32)
202            if equivariant_lode:
203                ls_lr = np.arange(lmax_lr + 1)
204            else:
205                ls_lr = np.array([0])
206
207            ## PARAMETERS FOR THE LR RADIAL BASIS
208            nextra_powers = len(self.lode_extra_powers)
209            if nextra_powers > 0:
210                ls_lr = np.concatenate([self.lode_extra_powers, ls_lr])
211
212            if self.a_lode > 0:
213                a = self.a_lode**2
214            else:
215                a = (
216                    self.param(
217                        "a_lr",
218                        lambda key: jnp.asarray([-self.a_lode] * ls_lr.shape[0])[
219                            None, :
220                        ],
221                    )
222                    ** 2
223                )
224            rc2a = rc**2 + a
225            ls_lr = 0.5 * (ls_lr[None, :] + 1)
226            ### minimal radial basis for long range (damped coulomb)
227            eij_lr = (
228                1.0 / (r**2 + a) ** ls_lr
229                - 1.0 / rc2a**ls_lr
230                + (r - rc) * rc * (2 * ls_lr) / rc2a ** (ls_lr + 1)
231            ) * switch_lode
232
233            if self.lode_rshort is not None:
234                rs = self.lode_rshort
235                d = self.lode_dshort
236                switch_short = 0.5 * (1 - jnp.cos(jnp.pi * (r - rs) / d)) * (r > rs) * (
237                    r < rs + d
238                ) + (r >= rs + d)
239                eij_lr = eij_lr * switch_short
240
241            dim_lr = 1
242            if nextra_powers > 0:
243                eij_lr_extra = eij_lr[:, :nextra_powers]
244                eij_lr = eij_lr[:, nextra_powers:]
245                dim_lr += nextra_powers
246
247            if equivariant_lode:
248                ## SPHERICAL HARMONICS ON LONG-RANGE GRAPH
249                eij_lr = eij_lr.repeat(nrep_lr, axis=-1)
250                Yij_lr = generate_spherical_harmonics(lmax=lmax_lr, normalize=False)(
251                    graph_lode["vec"] / r
252                )
253                dim_lr += lmax_lr
254                eij_lr = eij_lr * Yij_lr
255                del Yij_lr
256        
257
258        if self.keep_all_layers:
259            xis = []
260
261        ### START MESSAGE PASSING ITERATIONS
262        for layer in range(self.nlayers):
263            ## GATHER SCALAR EDGE FEATURES
264            u = [radial_terms]
265            if layer > 0:
266                ## edge-tensor contraction
267                xij2 = (Vi[edge_dst] + (parity* Vi)[edge_src]) * Yij
268                for l in range(self.lmax + 1):
269                    u.append((xij2[:,:, l**2 : (l + 1) ** 2]).sum(axis=-1))
270            ur = jnp.concatenate(u, axis=-1)
271
272            ## BUILD RELATIVE POSITIONAL ENCODING
273            if self.edge_value:
274                nout = 2
275            else:
276                nout = 1
277            w = FullyConnectedNet(
278                [2 * self.att_dim, nout*self.att_dim],
279                activation=self.positional_activation,
280                use_bias=self.positional_bias,
281                name=f"positional_encoding_{layer}",
282            )(ur).reshape(radial_terms.shape[0],nout, self.att_dim)
283            if self.edge_value:
284                w,vij = jnp.split(w, 2, axis=1)
285
286            nls = self.lmax + 1 if layer == 0 else 2 * (self.lmax + 1)
287
288
289            ## QUERY, KEY, VALUE
290            q = ln_qk(nn.Dense((self.scal_heads + nls*self.tens_heads) * self.att_dim, use_bias=False,name=f"queries_{layer}",kernel_init=kernel_init)(
291                xi
292            ).reshape(xi.shape[0], self.scal_heads + nls*self.tens_heads, self.att_dim))
293            k = nn.Dense((self.scal_heads + nls*self.tens_heads) * self.att_dim, use_bias=False, name=f"keys_{layer}",kernel_init=kernel_init)(
294                xi
295            ).reshape(xi.shape[0], self.scal_heads + nls*self.tens_heads, self.att_dim)
296
297            v = nn.Dense(self.scal_heads * self.att_dim, use_bias=False, name=f"values_{layer}",kernel_init=kernel_init)(xi).reshape(
298                xi.shape[0], self.scal_heads, self.att_dim
299            )
300
301            ## ATTENTION COEFFICIENTS
302            if self.additive_positional:
303                wk = ln_qk(w + k[edge_dst])
304            else:
305                wk = ln_qk(w * k[edge_dst])
306
307            act = activation_from_str(self.att_activation)
308            aij = (
309                act((q[edge_src] * wk).sum(axis=-1) / (self.att_dim**0.5))
310                * switch[:, None]
311            )
312
313            aijl = aij[:, : self.tens_heads*(self.lmax + 1)].reshape(-1,self.tens_heads,self.lmax+1).repeat(nrep, axis=-1)
314            if layer > 0:
315                aijl1 = aij[:, self.tens_heads*(self.lmax + 1) : self.tens_heads*nls].reshape(-1,self.tens_heads,self.lmax+1).repeat(nrep, axis=-1)
316            aij = aij[:, self.tens_heads*nls:, None]
317
318            if self.edge_value:
319                ## EDGE VALUES
320                if self.additive_positional:
321                    vij = vij + v[edge_dst]
322                else:
323                    vij = vij * v[edge_dst]
324            else:
325                ## MOVE DEST VALUES TO EDGE
326                vij = v[edge_dst]
327
328            ## SCALAR ATTENDED FEATURES
329            vai = jax.ops.segment_sum(
330                aij * vij,
331                edge_src,
332                num_segments=xi.shape[0],
333            )
334            vai = vai.reshape(xi.shape[0], -1)
335
336            ### TENSOR ATTENDED FEATURES
337            uij = aijl * Yij
338            if layer > 0:
339                uij = uij + aijl1 * Vi[edge_dst]
340            Vi = Vi + jax.ops.segment_sum(uij, edge_src, num_segments=Zi.shape[0])
341
342            ## SELF SCALAR FEATURES
343            si = nn.Dense(self.att_dim, use_bias=False, name=f"self_values_{layer}",kernel_init=kernel_init)(xi)
344
345            components = [si, vai]
346
347            ### CONTRACT TENSOR FEATURES TO BUILD INVARIANTS
348            if self.tens_heads == 1:
349                Vi2 = Vi**2
350            else:
351                Vi2 = Vi * ChannelMixing(self.lmax, self.tens_heads, name=f"extract_mixing_{layer}")(Vi)
352            for l in range(self.lmax + 1):
353                norm = 1.0 / (2 * l + 1)
354                components.append(
355                    (Vi2[:,:, l**2 : (l + 1) ** 2]).sum(axis=-1) * norm
356                )
357
358            ### LODE (~ LONG-RANGE ATTENTION)
359            if do_lode and layer == self.nlayers - 1:
360                assert self.lode_channels <= self.tens_heads
361                zj = nn.Dense(self.lode_channels*dim_lr, use_bias=False, name=f"lode_values_{layer}",kernel_init=kernel_init)(xi).reshape(
362                    xi.shape[0], self.lode_channels, dim_lr
363                )
364                if nextra_powers > 0:
365                    zj_extra = zj[:,:, :nextra_powers]
366                    zj = zj[:, :, nextra_powers:]
367                    xi_lr_extra = jax.ops.segment_sum(
368                        eij_lr_extra[:,None,:] * zj_extra[edge_dst_lr],
369                        edge_src_lr,
370                        species.shape[0],
371                    ).reshape(species.shape[0],-1)
372                    components.append(xi_lr_extra)
373                if equivariant_lode:
374                    zj = zj.repeat(nrep_lr, axis=-1)
375                Vi_lr = jax.ops.segment_sum(
376                    eij_lr[:,None,:] * zj[edge_dst_lr], edge_src_lr, species.shape[0]
377                )
378                components.append(Vi_lr[:,: , 0])
379                if equivariant_lode:
380                    Mi_lr = Vi[:,:self.lode_channels, : (lmax_lr + 1) ** 2] * Vi_lr
381                    for l in range(1, lmax_lr + 1):
382                        norm = 1.0 / (2 * l + 1)
383                        components.append(
384                            Mi_lr[:, :,l**2 : (l + 1) ** 2].sum(axis=-1)
385                            * norm
386                        )
387
388            ### CONCATENATE UPDATE COMPONENTS
389            components = jnp.concatenate(components, axis=-1)
390            if self.normalize_components:
391                components = _layer_norm(components)
392            ### COMPUTE UPDATE
393            if self.block_index_key is not None:
394                ## MoE neural network from block index
395                block_index = inputs[self.block_index_key]
396                updi = BlockIndexNet(
397                        output_dim=self.dim + self.tens_heads*(self.lmax + 1),
398                        hidden_neurons=self.update_hidden,
399                        activation=self.activation,
400                        use_bias=self.update_bias,
401                        name=f"update_net_{layer}",
402                        kernel_init=kernel_init,
403                    )((species,components, block_index))
404            else:
405                updi = FullyConnectedNet(
406                        [*self.update_hidden, self.dim + self.tens_heads*(self.lmax + 1)],
407                        activation=self.activation,
408                        use_bias=self.update_bias,
409                        name=f"update_net_{layer}",
410                        kernel_init=kernel_init,
411                    )(components)
412                
413            ## UPDATE ATOM FEATURES
414            xi = layer_norm(xi + updi[:,:self.dim])
415            Vi = Vi * (1 + updi[:,self.dim:]).reshape(-1,self.tens_heads,self.lmax+1).repeat(nrep, axis=-1)
416            if self.tens_heads > 1:
417                Vi = ChannelMixing(self.lmax, self.tens_heads,name=f"update_mixing_{layer}")(Vi)
418
419            if self.keep_all_layers:
420                ## STORE ALL LAYERS
421                xis.append(xi)
422
423
424        output = {**inputs, self.embedding_key: xi, self.embedding_key + "_tensor": Vi}
425        if self.keep_all_layers:
426            output[self.embedding_key+'_layers'] = jnp.stack(xis,axis=1)
427        return output