fennol.models.misc.misc

  1import flax.linen as nn
  2from typing import Any, Sequence, Callable, Union, ClassVar, Optional, Dict, List
  3import jax.numpy as jnp
  4import jax
  5import numpy as np
  6from functools import partial
  7from ...utils.activations import activation_from_str
  8from ...utils.periodic_table import (
  9    CHEMICAL_PROPERTIES,
 10    PERIODIC_TABLE,
 11    PERIODIC_TABLE_REV_IDX,
 12)
 13
 14
 15def apply_switch(x: jax.Array, switch: jax.Array):
 16    """@private Multiply a switch array to an array of values."""
 17    shape = x.shape
 18    return (
 19        jnp.expand_dims(x, axis=-1).reshape(*switch.shape, -1) * switch[..., None]
 20    ).reshape(shape)
 21
 22
 23class ApplySwitch(nn.Module):
 24    """Multiply an edge array by a switch array.
 25
 26    FID: APPLY_SWITCH
 27    """
 28
 29    key: str
 30    """The key of the input array."""
 31    switch_key: Optional[str] = None
 32    """The key of the switch array."""
 33    graph_key: Optional[str] = None
 34    """The key of the graph containing the switch."""
 35    output_key: Optional[str] = None
 36    """The key of the output array. If None, the input key is used."""
 37
 38    FID: ClassVar[str] = "APPLY_SWITCH"
 39
 40    @nn.compact
 41    def __call__(self, inputs) -> Any:
 42        if self.graph_key is not None:
 43            graph = inputs[self.graph_key]
 44            switch = graph["switch"]
 45        elif self.switch_key is not None:
 46            switch = inputs[self.switch_key]
 47        else:
 48            raise ValueError("Either graph_key or switch_key must be specified")
 49
 50        x = inputs[self.key]
 51        output = apply_switch(x, switch)
 52        output_key = self.key if self.output_key is None else self.output_key
 53        return {**inputs, output_key: output}
 54
 55
 56class AtomToEdge(nn.Module):
 57    """Map atom-wise values to edge-wise values.
 58
 59    FID: ATOM_TO_EDGE
 60
 61    By default, we map the destination atom value to the edge. This can be changed by setting `use_source` to True.
 62    """
 63
 64    _graphs_properties: Dict
 65    key: str
 66    """The key of the input atom-wise array."""
 67    output_key: Optional[str] = None
 68    """The key of the output edge-wise array. If None, the input key is used."""
 69    graph_key: str = "graph"
 70    """The key of the graph containing the edges."""
 71    switch: bool = False
 72    """Whether to apply a switch to the edge values."""
 73    switch_key: Optional[str] = None
 74    """The key of the switch array. If None, the switch is taken from the graph."""
 75    use_source: bool = False
 76    """Whether to use the source atom value instead of the destination atom value."""
 77
 78    FID: ClassVar[str] = "ATOM_TO_EDGE"
 79
 80    @nn.compact
 81    def __call__(self, inputs) -> Any:
 82        graph = inputs[self.graph_key]
 83        nat = inputs["species"].shape[0]
 84        edge_src, edge_dst = graph["edge_src"], graph["edge_dst"]
 85
 86        x = inputs[self.key]
 87        if self.use_source:
 88            x_edge = x[edge_src]
 89        else:
 90            x_edge = x[edge_dst]
 91
 92        if self.switch:
 93            switch = (
 94                graph["switch"] if self.switch_key is None else inputs[self.switch_key]
 95            )
 96            x_edge = apply_switch(x_edge, switch)
 97
 98        output_key = self.key if self.output_key is None else self.output_key
 99        return {**inputs, output_key: x_edge}
100
101
102class ScatterEdges(nn.Module):
103    """Reduce an edge array to atoms by summing over neighbors.
104
105    FID: SCATTER_EDGES
106    """
107
108    _graphs_properties: Dict
109    key: str
110    """The key of the input edge-wise array."""
111    output_key: Optional[str] = None
112    """The key of the output atom-wise array. If None, the input key is used."""
113    graph_key: str = "graph"
114    """The key of the graph containing the edges."""
115    switch: bool = False
116    """Whether to apply a switch to the edge values before summing."""
117    switch_key: Optional[str] = None
118    """The key of the switch array. If None, the switch is taken from the graph."""
119    antisymmetric: bool = False
120
121    FID: ClassVar[str] = "SCATTER_EDGES"
122
123    @nn.compact
124    def __call__(self, inputs) -> Any:
125        graph = inputs[self.graph_key]
126        nat = inputs["species"].shape[0]
127        x = inputs[self.key]
128
129        if self.switch:
130            switch = (
131                graph["switch"] if self.switch_key is None else inputs[self.switch_key]
132            )
133            x = apply_switch(x, switch)
134
135        edge_src, edge_dst = graph["edge_src"], graph["edge_dst"]
136        output = jax.ops.segment_sum(
137            x, edge_src, nat
138        )  # jnp.zeros((nat, *x.shape[1:])).at[edge_src].add(x,mode="drop")
139        if self.antisymmetric:
140            output = output - jax.ops.segment_sum(x, edge_dst, nat)
141
142        output_key = self.key if self.output_key is None else self.output_key
143        return {**inputs, output_key: output}
144
145
146class EdgeConcatenate(nn.Module):
147    """Concatenate the source and destination atom values of an edge.
148
149    FID: EDGE_CONCATENATE
150    """
151
152    _graphs_properties: Dict
153    key: str
154    """The key of the input atom-wise array."""
155    output_key: Optional[str] = None
156    """The key of the output edge-wise array. If None, the input key is used."""
157    graph_key: str = "graph"
158    """The key of the graph containing the edges."""
159    switch: bool = False
160    """Whether to apply a switch to the edge values."""
161    switch_key: Optional[str] = None
162    """The key of the switch array. If None, the switch is taken from the graph."""
163    axis: int = -1
164    """The axis along which to concatenate the atom values."""
165
166    FID: ClassVar[str] = "EDGE_CONCATENATE"
167
168    @nn.compact
169    def __call__(self, inputs) -> Any:
170        graph = inputs[self.graph_key]
171        edge_src, edge_dst = graph["edge_src"], graph["edge_dst"]
172        nat = inputs["species"].shape[0]
173        xi = inputs[self.key]
174
175        assert self._graphs_properties[self.graph_key][
176            "directed"
177        ], "EdgeConcatenate only works for directed graphs"
178        assert xi.shape[0] == nat, "Shape mismatch, xi.shape[0] != nat"
179
180        xij = jnp.concatenate([xi[edge_src], xi[edge_dst]], axis=self.axis)
181
182        if self.switch:
183            switch = (
184                graph["switch"] if self.switch_key is None else inputs[self.switch_key]
185            )
186            xij = apply_switch(xij, switch)
187
188        output_key = self.name if self.output_key is None else self.output_key
189        return {**inputs, output_key: xij}
190
191
192class ScatterSystem(nn.Module):
193    """Reduce an atom-wise array to a system-wise array by summing over atoms (in the batch).
194
195    FID: SCATTER_SYSTEM
196    """
197
198    key: str
199    """The key of the input atom-wise array."""
200    output_key: Optional[str] = None
201    """The key of the output system-wise array. If None, the input key is used."""
202    average: bool = False
203    """Wether to divide by the number of atoms in the system."""
204
205    FID: ClassVar[str] = "SCATTER_SYSTEM"
206
207    @nn.compact
208    def __call__(self, inputs) -> Any:
209        batch_index = inputs["batch_index"]
210        x = inputs[self.key]
211        assert (
212            x.shape[0] == batch_index.shape[0]
213        ), f"Shape mismatch {x.shape[0]} != {batch_index.shape[0]}"
214        nsys = inputs["natoms"].shape[0]
215        if self.average:
216            shape = [batch_index.shape[0]] + (x.ndim - 1) * [1]
217            x = x / inputs["natoms"][batch_index].reshape(shape)
218
219        output = jax.ops.segment_sum(x, batch_index, nsys)
220
221        output_key = self.key if self.output_key is None else self.output_key
222        return {**inputs, output_key: output}
223
224
225class SystemToAtoms(nn.Module):
226    """Broadcast a system-wise array to an atom-wise array.
227
228    FID: SYSTEM_TO_ATOMS
229    """
230
231    key: str
232    """The key of the input system-wise array."""
233    output_key: Optional[str] = None
234    """The key of the output atom-wise array. If None, the input key is used."""
235
236    FID: ClassVar[str] = "SYSTEM_TO_ATOMS"
237
238    @nn.compact
239    def __call__(self, inputs) -> Any:
240        batch_index = inputs["batch_index"]
241        x = inputs[self.key]
242        output = x[batch_index]
243
244        output_key = self.key if self.output_key is None else self.output_key
245        return {**inputs, output_key: output}
246
247
248class SumAxis(nn.Module):
249    """Sum an array along an axis.
250
251    FID: SUM_AXIS
252    """
253
254    key: str
255    """The key of the input array."""
256    axis: Union[None, int, Sequence[int]] = None
257    """The axis along which to sum the array."""
258    output_key: Optional[str] = None
259    """The key of the output array. If None, the input key is used."""
260    norm: Optional[str] = None
261    """Normalization of the sum. Can be 'dim', 'sqrt', or 'none'."""
262
263    FID: ClassVar[str] = "SUM_AXIS"
264
265    @nn.compact
266    def __call__(self, inputs) -> Any:
267        x = inputs[self.key]
268        output = jnp.sum(x, axis=self.axis)
269        if self.norm is not None:
270            norm = self.norm.lower()
271            if norm == "dim":
272                dim = np.prod(x.shape[self.axis])
273                output = output / dim
274            elif norm == "sqrt":
275                dim = np.prod(x.shape[self.axis])
276                output = output / dim**0.5
277            elif norm == "none":
278                pass
279            else:
280                raise ValueError(f"Unknown norm {norm}")
281        output_key = self.key if self.output_key is None else self.output_key
282        return {**inputs, output_key: output}
283
284
285class Split(nn.Module):
286    """Split an array along an axis.
287
288    FID: SPLIT
289    """
290
291    key: str
292    """The key of the input array."""
293    output_keys: Sequence[str]
294    """The keys of the output arrays."""
295    axis: int = -1
296    """The axis along which to split the array."""
297    sizes: Union[int, Sequence[int]] = 1
298    """The sizes of the splits."""
299    squeeze: bool = True
300    """Whether to remove the axis in the output if the size is 1."""
301
302    FID: ClassVar[str] = "SPLIT"
303
304    @nn.compact
305    def __call__(self, inputs) -> Any:
306        x = inputs[self.key]
307
308        if isinstance(self.sizes, int):
309            split_size = [self.sizes] * len(self.output_keys)
310        else:
311            split_size = self.sizes
312        if len(split_size) == len(self.output_keys):
313            assert (
314                sum(split_size) == x.shape[self.axis]
315            ), f"Split sizes {split_size} do not match input shape"
316            split_size = split_size[:-1]
317        assert (
318            len(split_size) == len(self.output_keys) - 1
319        ), f"Wrong number of split sizes {split_size} for {len(self.output_keys)} outputs"
320        split_indices = np.cumsum(split_size)
321        outs = {}
322
323        for k, v in zip(self.output_keys, jnp.split(x, split_indices, axis=self.axis)):
324            outs[k] = (
325                jnp.squeeze(v, axis=self.axis)
326                if self.squeeze and v.shape[self.axis] == 1
327                else v
328            )
329
330        return {**inputs, **outs}
331
332
333class Concatenate(nn.Module):
334    """Concatenate a list of arrays along an axis.
335
336    FID: CONCATENATE
337    """
338
339    keys: Sequence[str]
340    """The keys of the input arrays."""
341    axis: int = -1
342    """The axis along which to concatenate the arrays."""
343    output_key: Optional[str] = None
344    """The key of the output array. If None, the name of the module is used."""
345
346    FID: ClassVar[str] = "CONCATENATE"
347
348    @nn.compact
349    def __call__(self, inputs) -> Any:
350        output = jnp.concatenate([inputs[k] for k in self.keys], axis=self.axis)
351        output_key = self.output_key if self.output_key is not None else self.name
352        return {**inputs, output_key: output}
353
354
355class Activation(nn.Module):
356    """Apply an element-wise activation function to an array.
357
358    FID: ACTIVATION
359    """
360
361    key: str
362    """The key of the input array."""
363    activation: Union[Callable, str]
364    """The activation function or its name."""
365    scale_out: float = 1.0
366    """Output scaling factor."""
367    shift_out: float = 0.0
368    """Output shift."""
369    output_key: Optional[str] = None
370    """The key of the output array. If None, the input key is used."""
371
372    FID: ClassVar[str] = "ACTIVATION"
373
374    @nn.compact
375    def __call__(self, inputs) -> Any:
376        x = inputs[self.key]
377        activation = (
378            activation_from_str(self.activation)
379            if isinstance(self.activation, str)
380            else self.activation
381        )
382        output = self.scale_out * activation(x) + self.shift_out
383        output_key = self.output_key if self.output_key is not None else self.key
384        return {**inputs, output_key: output}
385
386class LayerNorm(nn.Module):
387    """Layer normalization module.
388
389    FID: LAYER_NORM
390    """
391
392    key: Optional[str] = None
393    """The key of the input."""
394    output_key: Optional[str] = None
395    """The key of the output. If None, it is the same as the key."""
396    axis: Union[int, Sequence[int]] = -1
397    """The axis to normalize."""
398    epsilon: float = 1e-6
399    """The epsilon for numerical stability."""
400    scale: float = 1.0
401    """The scale for the normalization."""
402    shift: float = 0.0
403    """The shift for the normalization."""
404
405    FID: ClassVar[str] = "LAYER_NORM"
406
407    @nn.compact
408    def __call__(self, inputs):
409        if isinstance(inputs, dict):
410            if self.key is None:
411                raise ValueError("Key must be specified for LayerNorm")
412            x = inputs[self.key]
413        else:
414            x = inputs
415        mu = jnp.mean(x, axis=self.axis, keepdims=True)
416        dx = x-mu
417        var = jnp.mean(dx ** 2, axis=self.axis, keepdims=True)
418        sig = (self.epsilon + var) ** (-0.5)
419        out = self.scale * (sig * dx) + self.shift
420
421        output_key = self.output_key if self.output_key is not None else self.key
422        
423        if isinstance(inputs, dict):
424            return {**inputs, output_key: out}
425        
426        return out
427        
428
429
430class Scale(nn.Module):
431    """Scale an array by a constant factor.
432
433    FID: SCALE
434    """
435
436    key: str
437    """The key of the input array."""
438    scale: float
439    """The (initial) scaling factor."""
440    output_key: Optional[str] = None
441    """The key of the output array. If None, the input key is used."""
442    trainable: bool = False
443    """Whether the scaling factor is trainable."""
444
445    FID: ClassVar[str] = "SCALE"
446
447    @nn.compact
448    def __call__(self, inputs) -> Any:
449        x = inputs[self.key]
450
451        if self.trainable:
452            scale = self.param("scale", lambda rng: jnp.asarray(self.scale))
453        else:
454            scale = self.scale
455
456        output = scale * x
457        output_key = self.output_key if self.output_key is not None else self.key
458        return {**inputs, output_key: output}
459
460
461class Add(nn.Module):
462    """Add together a list of arrays.
463
464    FID: ADD
465    """
466
467    keys: Sequence[str]
468    """The keys of the input arrays."""
469    output_key: Optional[str] = None
470    """The key of the output array. If None, the name of the module is used."""
471
472    FID: ClassVar[str] = "ADD"
473
474    @nn.compact
475    def __call__(self, inputs) -> Any:
476        output = 0
477        for k in self.keys:
478            output = output + inputs[k]
479
480        output_key = self.output_key if self.output_key is not None else self.name
481        return {**inputs, output_key: output}
482
483
484class Multiply(nn.Module):
485    """Element-wise-multiply together a list of arrays.
486
487    FID: MULTIPLY
488    """
489
490    keys: Sequence[str]
491    """The keys of the input arrays."""
492    output_key: Optional[str] = None
493    """The key of the output array. If None, the name of the module is used."""
494
495    FID: ClassVar[str] = "MULTIPLY"
496
497    @nn.compact
498    def __call__(self, inputs) -> Any:
499        output = 1
500        for k in self.keys:
501            output = output * inputs[k]
502
503        output_key = self.output_key if self.output_key is not None else self.name
504        return {**inputs, output_key: output}
505
506
507class Transpose(nn.Module):
508    """Transpose an array.
509
510    FID: TRANSPOSE
511    """
512
513    key: str
514    """The key of the input array."""
515    axes: Sequence[int]
516    """The permutation of the axes. See `jax.numpy.transpose` for more details."""
517    output_key: Optional[str] = None
518    """The key of the output array. If None, the input key is used."""
519
520    FID: ClassVar[str] = "TRANSPOSE"
521
522    @nn.compact
523    def __call__(self, inputs) -> Any:
524        output = jnp.transpose(inputs[self.key], axes=self.axes)
525        output_key = self.output_key if self.output_key is not None else self.key
526        return {**inputs, output_key: output}
527
528
529class Reshape(nn.Module):
530    """Reshape an array.
531
532    FID: RESHAPE
533    """
534
535    key: str
536    """The key of the input array."""
537    shape: Sequence[Union[int,str]]
538    """The shape of the output array."""
539    output_key: Optional[str] = None
540    """The key of the output array. If None, the input key is used."""
541
542    FID: ClassVar[str] = "RESHAPE"
543
544    @nn.compact
545    def __call__(self, inputs) -> Any:
546        shape = []
547        for s in self.shape:
548            if isinstance(s,int):
549                shape.append(s)
550                continue
551
552            if isinstance(s,str):
553                s_=s.lower().strip()
554                if s_ in ["natoms" ,"nat","natom","n_atoms","atoms"]:
555                    shape.append(inputs["species"].shape[0])
556                    continue
557                
558                if s_ in ["nsys","nbatch","nsystems","n_sys","n_systems","n_batch"]:
559                    shape.append(inputs["natoms"].shape[0])
560                    continue
561                
562                s_ = s.strip().split("[")
563                key = s_[0]
564                if key in inputs:
565                    axis = int(s_[1].split("]")[0])
566                    shape.append(inputs[key].shape[axis])
567                    continue
568
569            raise ValueError(f"Error parsing shape component {s}")
570
571        output = jnp.reshape(inputs[self.key], shape)
572        output_key = self.output_key if self.output_key is not None else self.key
573        return {**inputs, output_key: output}
574
575
576class ChemicalConstant(nn.Module):
577    """Map atomic species to a constant value.
578
579    FID: CHEMICAL_CONSTANT
580    """
581
582    value: Union[str, List[float], float, Dict]
583    """The constant value or a dictionary of values for each element."""
584    output_key: Optional[str] = None
585    """The key of the output array. If None, the name of the module is used."""
586    trainable: bool = False
587    """Whether the constant is trainable."""
588
589    FID: ClassVar[str] = "CHEMICAL_CONSTANT"
590
591    @nn.compact
592    def __call__(self, inputs) -> Any:
593        if isinstance(self.value, str):
594            constant = CHEMICAL_PROPERTIES[self.value.upper()]
595        elif isinstance(self.value, list) or isinstance(self.value, tuple):
596            constant = list(self.value)
597        elif isinstance(self.value, float):
598            constant = [self.value] * len(PERIODIC_TABLE)
599        elif hasattr(self.value, "items"):
600            constant = [0.0] * len(PERIODIC_TABLE)
601            for k, v in self.value.items():
602                constant[PERIODIC_TABLE_REV_IDX[k]] = v
603        else:
604            raise ValueError(f"Unknown constant type {type(self.value)}")
605
606        if self.trainable:
607            constant = self.param(
608                "constant", lambda rng: jnp.asarray(constant, dtype=jnp.float32)
609            )
610        else:
611            constant = jnp.asarray(constant, dtype=jnp.float32)
612        output = constant[inputs["species"]]
613        output_key = self.output_key if self.output_key is not None else self.name
614        return {**inputs, output_key: output}
615
616
617class SwitchFunction(nn.Module):
618    """Compute a switch array from an array of distances and a cutoff.
619
620    FID: SWITCH_FUNCTION
621    """
622
623    cutoff: Optional[float] = None
624    """The cutoff distance. If None, the cutoff is taken from the graph."""
625    switch_start: float = 0.0
626    """The proportion of the cutoff distance at which the switch function starts."""
627    graph_key: Optional[str] = "graph"
628    """The key of the graph containing the distances and edge mask."""
629    output_key: Optional[str] = None
630    """The key of the output switch array. If None, it is added to the graph."""
631    switch_type: str = "cosine"
632    """The type of switch function. Can be 'cosine', 'polynomial', or 'exponential'."""
633    p: Optional[float] = None
634    """ The parameter of the switch function. If None, it is fixed to the default for each `switch_type`."""
635    trainable: bool = False
636    """Whether the switch parameter is trainable."""
637
638    FID: ClassVar[str] = "SWITCH_FUNCTION"
639
640    @nn.compact
641    def __call__(self, inputs) -> Any:
642        if self.graph_key is not None:
643            graph = inputs[self.graph_key]
644            distances, edge_mask = graph["distances"], graph["edge_mask"]
645            if self.cutoff is not None:
646                edge_mask = jnp.logical_and(edge_mask, (distances < self.cutoff))
647                cutoff = self.cutoff
648            else:
649                cutoff = graph["cutoff"]
650        else:
651            # distances = inputs
652            if len(inputs) == 3:
653                distances, edge_mask, cutoff = inputs
654            else:
655                distances, edge_mask = inputs
656                assert (
657                    self.cutoff is not None
658                ), "cutoff must be specified if no graph is given"
659                # edge_mask = distances < self.cutoff
660                cutoff = self.cutoff
661
662        if self.switch_start > 1.0e-5:
663            assert (
664                self.switch_start < 1.0
665            ), "switch_start is a proportion of cutoff and must be smaller than 1."
666            cutoff_in = self.switch_start * cutoff
667            x = distances - cutoff_in
668            end = cutoff - cutoff_in
669        else:
670            x = distances
671            end = cutoff
672
673        switch_type = self.switch_type.lower()
674        if switch_type == "cosine":
675            p = self.p if self.p is not None else 1.0
676            if self.trainable:
677                p = self.param("p", lambda rng: jnp.asarray(p, dtype=jnp.float32))
678            switch = (0.5 * jnp.cos(x * (jnp.pi / end)) + 0.5) ** p
679
680        elif switch_type == "polynomial":
681            p = self.p if self.p is not None else 3.0
682            if self.trainable:
683                p = self.param("p", lambda rng: jnp.asarray(p, dtype=jnp.float32))
684            d = x / end
685            switch = (
686                1.0
687                - 0.5 * (p + 1) * (p + 2) * d**p
688                + p * (p + 2) * d ** (p + 1)
689                - 0.5 * p * (p + 1) * d ** (p + 2)
690            )
691
692        elif switch_type == "exponential":
693            p = self.p if self.p is not None else 1.0
694            if self.trainable:
695                p = self.param("p", lambda rng: jnp.asarray(p, dtype=jnp.float32))
696            r2 = x**2
697            c2 = end**2
698            switch = jnp.exp(-p * r2 / (c2 - r2))
699        
700        elif switch_type == "hard":
701            switch = jnp.where(distances < cutoff, 1.0, 0.0)
702        else:
703            raise ValueError(f"Unknown switch function {switch_type}")
704
705        if self.switch_start > 1.0e-5:
706            switch = jnp.where(distances < cutoff_in, 1.0, switch)
707
708        switch = jnp.where(edge_mask, switch, 0.0)
709
710        if self.graph_key is not None:
711            if self.output_key is not None:
712                return {**inputs, self.output_key: switch}
713            else:
714                return {**inputs, self.graph_key: {**graph, "switch": switch}}
715        else:
716            return switch  # , edge_mask