fennol.models.embeddings.mace

  1import functools
  2import math
  3import jax
  4import jax.numpy as jnp
  5import flax.linen as nn
  6from typing import Sequence, Dict, Union, ClassVar, Optional, Set
  7import dataclasses
  8
  9from ..misc.encodings import RadialBasis
 10from ...utils.activations import activation_from_str
 11
 12
 13try:
 14    import e3nn_jax as e3nn
 15
 16    E3NN_AVAILABLE = True
 17    E3NN_EXCEPTION = None
 18    Irreps = e3nn.Irreps
 19    Irrep = e3nn.Irrep
 20except Exception as e:
 21    E3NN_AVAILABLE = False
 22    E3NN_EXCEPTION = e
 23    e3nn = None
 24
 25    class Irreps(tuple):
 26        pass
 27
 28    class Irrep(tuple):
 29        pass
 30
 31
 32class MACE(nn.Module):
 33    """MACE equivariant message passing neural network.
 34
 35    adapted from MACE-jax github repo by M. Geiger and I. Batatia
 36    
 37    T. Plé reordered some operations and changed defaults to match the recent mace-torch version 
 38    -> compatibility with pretrained torch models requires some work on the parameters:
 39        - normalization of activation functions in e3nn differ between jax and pytorch => need rescaling
 40        - multiplicity ordering and signs of U matrices in SymmetricContraction differ => need to reorder and flip signs in the weight tensors
 41        - we use a maximum Z instead of a list of species => need to adapt species-dependent parameters
 42    
 43    References:
 44        - I. Batatia et al. "MACE: Higher order equivariant message passing neural networks for fast and accurate force fields." Advances in Neural Information Processing Systems 35 (2022): 11423-11436.
 45        https://doi.org/10.48550/arXiv.2206.07697
 46        - I. Batatia et al. "The design space of e(3)-equivariant atom-centered interatomic potentials." arXiv preprint arXiv:2205.06643 (2022).
 47        https://doi.org/10.48550/arXiv.2205.06643
 48
 49    """
 50    _graphs_properties: Dict
 51    output_irreps: Union[Irreps, str] = "1x0e"
 52    """The output irreps of the model."""
 53    hidden_irreps: Union[Irreps, str] = "128x0e + 128x1o"
 54    """The hidden irreps of the model."""
 55    readout_mlp_irreps: Union[Irreps, str] = "16x0e"
 56    """The hidden irreps of the readout MLP."""
 57    graph_key: str = "graph"
 58    """The key in the input dictionary that corresponds to the molecular graph to use."""
 59    output_key: Optional[str] = None
 60    """The key of the embedding in the output dictionary."""
 61    avg_num_neighbors: float = 1.0
 62    """The expected average number of neighbors."""
 63    ninteractions: int = 2
 64    """The number of interaction layers."""
 65    num_features: Optional[int] = None
 66    """The number of features per node. default gcd of hidden_irreps multiplicities"""
 67    radial_basis: dict = dataclasses.field(
 68        default_factory=lambda: {"basis": "bessel", "dim": 8, "trainable": False}
 69    )
 70    """The dictionary of parameters for radial basis functions. See `fennol.models.misc.encodings.RadialBasis`."""
 71    lmax: int = 1
 72    """The maximum angular momentum to consider."""
 73    correlation: int = 3
 74    """The correlation order at each layer."""
 75    activation: str = "silu"
 76    """The activation function to use."""
 77    symmetric_tensor_product_basis: bool = False
 78    """Whether to use the symmetric tensor product basis."""
 79    interaction_irreps: Union[Irreps, str] = "o3_restricted"
 80    skip_connection_first_layer: bool = True
 81    radial_network_hidden: Sequence[int] = dataclasses.field(
 82        default_factory=lambda: [64, 64, 64]
 83    )
 84    scalar_output: bool = False
 85    zmax: int = 86
 86    """The maximum atomic number to consider."""
 87    convolution_mode: int = 1
 88
 89    FID: ClassVar[str] = "MACE"
 90
 91    @nn.compact
 92    def __call__(self, inputs):
 93        if not E3NN_AVAILABLE:
 94            raise E3NN_EXCEPTION
 95
 96        species_indices = inputs["species"]
 97        graph = inputs[self.graph_key]
 98        distances = graph["distances"]
 99        vec = e3nn.IrrepsArray("1o", graph["vec"])
100        switch = graph["switch"]
101        edge_src = graph["edge_src"]
102        edge_dst = graph["edge_dst"]
103
104        output_irreps = e3nn.Irreps(self.output_irreps)
105        hidden_irreps = e3nn.Irreps(self.hidden_irreps)
106        readout_mlp_irreps = e3nn.Irreps(self.readout_mlp_irreps)
107
108        # extract or set num_features
109        if self.num_features is None:
110            num_features = functools.reduce(math.gcd, (mul for mul, _ in hidden_irreps))
111            hidden_irreps = e3nn.Irreps(
112                [(mul // num_features, ir) for mul, ir in hidden_irreps]
113            )
114        else:
115            num_features = self.num_features
116
117        # get interaction irreps
118        if self.interaction_irreps == "o3_restricted":
119            interaction_irreps = e3nn.Irreps.spherical_harmonics(self.lmax)
120        elif self.interaction_irreps == "o3_full":
121            interaction_irreps = e3nn.Irreps(e3nn.Irrep.iterator(self.lmax))
122        else:
123            interaction_irreps = e3nn.Irreps(self.interaction_irreps)
124        convol_irreps = num_features * interaction_irreps
125
126        # convert species to internal indices
127        # maxidx = max(PERIODIC_TABLE_REV_IDX.values())
128        # conv_tensor = [0] * (maxidx + 2)
129        # if isinstance(self.species_order, str):
130        #     species_order = [el.strip() for el in self.species_order.split(",")]
131        # else:
132        #     species_order = [el for el in self.species_order]
133        # for i, s in enumerate(species_order):
134        #     conv_tensor[PERIODIC_TABLE_REV_IDX[s]] = i
135        # species_indices = jnp.asarray(conv_tensor, dtype=jnp.int32)[species]
136        num_species = self.zmax + 2
137
138        # species encoding
139        encoding_irreps: e3nn.Irreps = (
140            (num_features * hidden_irreps).filter("0e").regroup()
141        )
142        species_encoding = self.param(
143            "species_encoding",
144            lambda key, shape: jax.nn.standardize(
145                jax.random.normal(key, shape, dtype=jnp.float32)
146            ),
147            (num_species, encoding_irreps.dim),
148        )[species_indices]
149        # convert to IrrepsArray
150        node_feats = e3nn.IrrepsArray(encoding_irreps, species_encoding)
151
152        # radial embedding
153        cutoff = self._graphs_properties[self.graph_key]["cutoff"]
154        radial_embedding = (
155            RadialBasis(
156                **{
157                    **self.radial_basis,
158                    "end": cutoff,
159                    "name": f"RadialBasis",
160                }
161            )(distances)
162            * switch[:, None]
163        )
164
165        # spherical harmonics
166        assert self.convolution_mode in [0,1,2], "convolution_mode must be 0, 1 or 2"
167        if self.convolution_mode == 0:
168            Yij = e3nn.spherical_harmonics(range(0, self.lmax + 1), vec, True)
169        elif self.convolution_mode == 1:
170            Yij = e3nn.spherical_harmonics(range(1, self.lmax + 1), vec, True)
171
172        outputs = []
173        node_feats_all = []
174        for layer in range(self.ninteractions):
175            first = layer == 0
176            last = layer == self.ninteractions - 1
177
178            layer_irreps = num_features * (
179                hidden_irreps if not last else hidden_irreps.filter(output_irreps)
180            )
181
182            # Linear skip connection
183            sc = None
184            if not first or self.skip_connection_first_layer:
185                sc = e3nn.flax.Linear(
186                    layer_irreps,
187                    num_indexed_weights=num_species,
188                    name=f"skip_tp_{layer}",
189                    force_irreps_out=True,
190                )(species_indices, node_feats)
191
192            ################################################
193            # Interaction block (Message passing convolution)
194            node_feats = e3nn.flax.Linear(node_feats.irreps, name=f"linear_up_{layer}")(
195                node_feats
196            )
197
198
199            messages = node_feats[edge_src]
200            if self.convolution_mode == 0:
201                messages = e3nn.tensor_product(
202                    messages,
203                    Yij,
204                    filter_ir_out=convol_irreps,
205                    regroup_output=True,
206                )
207            elif self.convolution_mode == 1:
208                messages = e3nn.concatenate(
209                    [
210                        messages.filter(convol_irreps),
211                        e3nn.tensor_product(
212                            messages,
213                            Yij,
214                            filter_ir_out=convol_irreps,
215                        ),
216                        # e3nn.tensor_product_with_spherical_harmonics(
217                        #     messages, vectors, self.max_ell
218                        # ).filter(convol_irreps),
219                    ]
220                ).regroup()
221            else:
222                messages = e3nn.tensor_product_with_spherical_harmonics(
223                    messages, vec, self.lmax
224                ).filter(convol_irreps).regroup()
225
226            # mix = FullyConnectedNet(
227            #     [*self.radial_network_hidden, messages.irreps.num_irreps],
228            #     activation=activation_from_str(self.activation),
229            #     name=f"radial_network_{layer}",
230            #     use_bias=False,
231            # )(radial_embedding)
232            mix = e3nn.flax.MultiLayerPerceptron(
233                [*self.radial_network_hidden, messages.irreps.num_irreps],
234                act=activation_from_str(self.activation),
235                output_activation=False,
236                name=f"radial_network_{layer}",
237                gradient_normalization="element",
238            )(
239                radial_embedding
240            )
241
242            messages = messages * mix
243            node_feats = (
244                e3nn.IrrepsArray.zeros(
245                    messages.irreps, node_feats.shape[:1], messages.dtype
246                )
247                .at[edge_dst]
248                .add(messages)
249            )
250            # print("irreps_mid jax",node_feats.irreps)
251            # jax.debug.print("node_feats={n}", n=jnp.sum(node_feats.array,axis=0)[550:570])
252
253            node_feats = (
254                e3nn.flax.Linear(convol_irreps, name=f"linear_dn_{layer}")(node_feats)
255                / self.avg_num_neighbors
256            )
257
258            if first and not self.skip_connection_first_layer:
259                node_feats = e3nn.flax.Linear(
260                    node_feats.irreps,
261                    num_indexed_weights=num_species,
262                    name=f"skip_tp_{layer}",
263                )(species_indices, node_feats)
264
265            ################################################
266            # Equivariant product basis block
267
268            # symmetric contractions
269            node_feats = SymmetricContraction(
270                keep_irrep_out={ir for _, ir in layer_irreps},
271                correlation=self.correlation,
272                num_species=num_species,
273                gradient_normalization="element",  # NOTE: This is to copy mace-torch
274                symmetric_tensor_product_basis=self.symmetric_tensor_product_basis,
275            )(
276                node_feats, species_indices
277            )
278
279
280            node_feats = e3nn.flax.Linear(
281                layer_irreps, name=f"linear_contraction_{layer}"
282            )(node_feats)
283
284
285            if sc is not None:
286                # add skip connection
287                node_feats = node_feats + sc
288
289
290            ################################################
291            
292            # Readout block
293            if last:
294                num_vectors = readout_mlp_irreps.filter(drop=["0e", "0o"]).num_irreps
295                layer_out = e3nn.flax.Linear(
296                    (readout_mlp_irreps + e3nn.Irreps(f"{num_vectors}x0e")).simplify(),
297                    name=f"hidden_linear_readout_last",
298                )(node_feats)
299                layer_out = e3nn.gate(
300                    layer_out,
301                    even_act=activation_from_str(self.activation),
302                    even_gate_act=None,
303                )
304                layer_out = e3nn.flax.Linear(
305                    output_irreps, name=f"linear_readout_last"
306                )(layer_out)
307            else:
308                layer_out = e3nn.flax.Linear(
309                    output_irreps,
310                    name=f"linear_readout_{layer}",
311                )(node_feats)
312
313            if self.scalar_output:
314                layer_out = layer_out.filter("0e").array
315
316            outputs.append(layer_out)
317            node_feats_all.append(node_feats.filter("0e").array)
318
319        if self.scalar_output:
320            output = jnp.stack(outputs, axis=1)
321        else:
322            output = e3nn.stack(outputs, axis=1)
323
324        node_feats_all = jnp.concatenate(node_feats_all, axis=-1)
325
326        output_key = self.output_key if self.output_key is not None else self.name
327        return {
328            **inputs,
329            output_key: output,
330            output_key + "_node_feats": node_feats_all,
331        }
332
333
334class SymmetricContraction(nn.Module):
335
336    correlation: int
337    keep_irrep_out: Set[Irrep]
338    num_species: int
339    gradient_normalization: Union[str, float]
340    symmetric_tensor_product_basis: bool
341
342    @nn.compact
343    def __call__(self, input, index):
344        if not E3NN_AVAILABLE:
345            raise E3NN_EXCEPTION
346
347        if self.gradient_normalization is None:
348            gradient_normalization = e3nn.config("gradient_normalization")
349        else:
350            gradient_normalization = self.gradient_normalization
351        if isinstance(gradient_normalization, str):
352            gradient_normalization = {"element": 0.0, "path": 1.0}[
353                gradient_normalization
354            ]
355
356        keep_irrep_out = self.keep_irrep_out
357        if isinstance(keep_irrep_out, str):
358            keep_irrep_out = e3nn.Irreps(keep_irrep_out)
359            assert all(mul == 1 for mul, _ in keep_irrep_out)
360
361        keep_irrep_out = {e3nn.Irrep(ir) for ir in keep_irrep_out}
362
363        input = input.mul_to_axis().remove_nones()
364
365        ### PREPARE WEIGHTS
366        ws = []
367        Us = []
368        for order in range(1, self.correlation + 1):  # correlation, ..., 1
369            if self.symmetric_tensor_product_basis:
370                U = e3nn.reduced_symmetric_tensor_product_basis(
371                    input.irreps, order, keep_ir=keep_irrep_out
372                )
373            else:
374                U = e3nn.reduced_tensor_product_basis(
375                    [input.irreps] * order, keep_ir=keep_irrep_out
376                )
377            # U = U / order  # normalization TODO(mario): put back after testing
378            # NOTE(mario): The normalization constants (/order and /mul**0.5)
379            # has been numerically checked to be correct.
380
381            # TODO(mario) implement norm_p
382            Us.append(U)
383
384            wsorder = []
385            for (mul, ir_out), u in zip(U.irreps, U.list):
386                u = u.astype(input.array.dtype)
387                # u: ndarray [(irreps_x.dim)^order, multiplicity, ir_out.dim]
388
389                w = self.param(
390                    f"w{order}_{ir_out}",
391                    nn.initializers.normal(
392                        stddev=(mul**-0.5) ** (1.0 - gradient_normalization)
393                    ),
394                    (self.num_species, mul, input.shape[-2]),
395                )
396                w = w * (mul**-0.5) ** gradient_normalization
397                wsorder.append(w)
398            ws.append(wsorder)
399
400        def fn(input: e3nn.IrrepsArray, index: jnp.ndarray):
401            # - This operation is parallel on the feature dimension (but each feature has its own parameters)
402            # This operation is an efficient implementation of
403            # vmap(lambda w, x: FunctionalLinear(irreps_out)(w, concatenate([x, tensor_product(x, x), tensor_product(x, x, x), ...])))(w, x)
404            # up to x power self.correlation
405            assert input.ndim == 2  # [num_features, irreps_x.dim]
406            assert index.ndim == 0  # int
407
408            out = dict()
409            x_ = input.array
410
411            for order in range(self.correlation, 0, -1):  # correlation, ..., 1
412
413                U = Us[order - 1]
414
415                # ((w3 x + w2) x + w1) x
416                #  \-----------/
417                #       out
418
419                for ii, ((mul, ir_out), u) in enumerate(zip(U.irreps, U.list)):
420                    u = u.astype(x_.dtype)
421                    # u: ndarray [(irreps_x.dim)^order, multiplicity, ir_out.dim]
422
423                    w = ws[order - 1][ii][index]
424                    if ir_out not in out:
425                        out[ir_out] = (
426                            "special",
427                            jnp.einsum("...jki,kc,cj->c...i", u, w, x_),
428                        )  # [num_features, (irreps_x.dim)^(oder-1), ir_out.dim]
429                    else:
430                        out[ir_out] += jnp.einsum(
431                            "...ki,kc->c...i", u, w
432                        )  # [num_features, (irreps_x.dim)^order, ir_out.dim]
433
434                # ((w3 x + w2) x + w1) x
435                #  \----------------/
436                #         out (in the normal case)
437
438                for ir_out in out:
439                    if isinstance(out[ir_out], tuple):
440                        out[ir_out] = out[ir_out][1]
441                        continue  # already done (special case optimization above)
442
443                    out[ir_out] = jnp.einsum(
444                        "c...ji,cj->c...i", out[ir_out], x_
445                    )  # [num_features, (irreps_x.dim)^(oder-1), ir_out.dim]
446
447                # ((w3 x + w2) x + w1) x
448                #  \-------------------/
449                #           out
450
451            # out[irrep_out] : [num_features, ir_out.dim]
452            irreps_out = e3nn.Irreps(sorted(out.keys()))
453            return e3nn.IrrepsArray.from_list(
454                irreps_out,
455                [out[ir][:, None, :] for (_, ir) in irreps_out],
456                (input.shape[0],),
457            )
458
459        # Treat batch indices using vmap
460        shape = jnp.broadcast_shapes(input.shape[:-2], index.shape)
461        input = input.broadcast_to(shape + input.shape[-2:])
462        index = jnp.broadcast_to(index, shape)
463
464        fn_mapped = fn
465        for _ in range(input.ndim - 2):
466            fn_mapped = jax.vmap(fn_mapped)
467
468        return fn_mapped(input, index).axis_to_mul()
469
470
471# class SymmetricContraction(nn.Module):
472
473#     correlation: int
474#     keep_irrep_out: Set[Irrep]
475#     num_species: int
476#     gradient_normalization: Union[str, float]
477#     symmetric_tensor_product_basis: bool
478
479#     @nn.compact
480#     def __call__(self, input: IrrepsArray, index: jnp.ndarray):
481#         if not E3NN_AVAILABLE:
482#             raise E3NN_EXCEPTION
483
484#         if self.gradient_normalization is None:
485#             gradient_normalization = e3nn.config("gradient_normalization")
486#         else:
487#             gradient_normalization = self.gradient_normalization
488#         if isinstance(gradient_normalization, str):
489#             gradient_normalization = {"element": 0.0, "path": 1.0}[
490#                 gradient_normalization
491#             ]
492
493#         keep_irrep_out = self.keep_irrep_out
494#         if isinstance(keep_irrep_out, str):
495#             keep_irrep_out = e3nn.Irreps(keep_irrep_out)
496#             assert all(mul == 1 for mul, _ in keep_irrep_out)
497
498#         keep_irrep_out = {e3nn.Irrep(ir) for ir in keep_irrep_out}
499
500#         onehot = jnp.eye(self.num_species)[index]
501
502#         ### PREPARE WEIGHTS
503#         ws = []
504#         us = []
505#         for ir_out in keep_irrep_out:
506#             usorder = []
507#             wsorder = []
508#             for order in range(1, self.correlation + 1):  # correlation, ..., 1
509#                 if self.symmetric_tensor_product_basis:
510#                     U = e3nn.reduced_symmetric_tensor_product_basis(
511#                         input.irreps, order, keep_ir=[ir_out]
512#                     )
513#                 else:
514#                     U = e3nn.reduced_tensor_product_basis(
515#                         [input.irreps] * order, keep_ir=[ir_out]
516#                     )
517#                 u = jnp.moveaxis(U.list[0].astype(input.array.dtype), -1, 0)
518#                 usorder.append(u)
519
520#                 mul, _ = U.irreps[0]
521#                 w = self.param(
522#                     f"w{order}_{ir_out}",
523#                     nn.initializers.normal(
524#                         stddev=(mul**-0.5) ** (1.0 - gradient_normalization)
525#                     ),
526#                     (self.num_species, mul, input.shape[-2]),
527#                 )
528#                 w = w * (mul**-0.5) ** gradient_normalization
529#                 wsorder.append(w)
530#             ws.append(wsorder)
531#             us.append(usorder)
532
533#         x = input.array
534
535#         outs = []
536#         for i, ir in enumerate(keep_irrep_out):
537#             w = ws[i][-1]  # [index]
538#             u = us[i][-1]
539#             out = jnp.einsum("...jk,ekc,bcj,be->bc...", u, w, x, onehot)
540
541#             for order in range(self.correlation - 1, 0, -1):
542#                 w = ws[i][order - 1]  # [index]
543#                 u = us[i][order - 1]
544
545#                 c_tensor = jnp.einsum("...k,ekc,be->bc...", u, w, onehot) + out
546#                 out = jnp.einsum("bc...j,bcj->bc...", c_tensor, x)
547
548#             outs.append(out.reshape(x.shape[0], -1))
549
550#         out = jnp.concatenate(outs, axis=-1)
551
552#         return e3nn.IrrepsArray(input.shape[1] * e3nn.Irreps(keep_irrep_out), out)