File size: 2,468 Bytes
3e02ab8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | # This file is a part of the `nequip` package. Please see LICENSE and README at the root for information on using it.
from typing import List, Optional
import torch
from e3nn.o3._irreps import Irreps
from onescience.datapipes.materials.nequip import AtomicDataDict
from ._graph_mixin import GraphModuleMixin
class Concat(GraphModuleMixin, torch.nn.Module):
"""Concatenate multiple fields into one."""
def __init__(self, in_fields: List[str], out_field: str, irreps_in={}):
super().__init__()
self.in_fields = list(in_fields)
self.out_field = out_field
self._init_irreps(irreps_in=irreps_in, required_irreps_in=self.in_fields)
self.irreps_out[self.out_field] = sum(
(self.irreps_in[k] for k in self.in_fields), Irreps()
)
def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
data[self.out_field] = torch.cat([data[k] for k in self.in_fields], dim=-1)
return data
class ApplyFactor(GraphModuleMixin, torch.nn.Module):
"""Applies factor to field."""
def __init__(
self,
in_field: str,
factor: float,
out_field: Optional[str] = None,
irreps_in={},
):
super().__init__()
self.in_field = in_field
self.out_field = in_field if out_field is None else out_field
self.factor = factor
self._init_irreps(irreps_in=irreps_in)
self.irreps_out[self.out_field] = self.irreps_in[self.in_field]
def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
data[self.out_field] = self.factor * data[self.in_field]
return data
class SaveForOutput(torch.nn.Module, GraphModuleMixin):
"""Copy a field and disconnect it from the autograd graph.
Copy a field and disconnect it from the autograd graph, storing it under another key for inspection as part of the models output.
Args:
field: the field to save
out_field: the key to put the saved copy in
"""
field: str
out_field: str
def __init__(self, field: str, out_field: str, irreps_in=None):
super().__init__()
self._init_irreps(irreps_in=irreps_in)
self.irreps_out[out_field] = self.irreps_in[field]
self.field = field
self.out_field = out_field
def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
data[self.out_field] = data[self.field].detach().clone()
return data
|