File size: 1,646 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 | import torch
from onescience.datapipes.materials.nequip import AtomicDataDict
from ._graph_mixin import GraphModuleMixin
from .model_modifier_utils import replace_submodules, model_modifier
class GhostExchangeModule(GraphModuleMixin, torch.nn.Module):
"""Base class for ghost atom exchange modules."""
def __init__(
self,
field: str = AtomicDataDict.NODE_FEATURES_KEY,
irreps_in={},
):
super().__init__()
self.field = field
self._init_irreps(
irreps_in=irreps_in,
my_irreps_in={field: irreps_in[field]},
irreps_out={field: irreps_in[field]},
)
def forward(
self,
data: AtomicDataDict.Type,
ghost_included: bool,
) -> AtomicDataDict.Type:
raise NotImplementedError("Subclasses must implement forward method")
class NoOpGhostExchangeModule(GhostExchangeModule):
"""Base ghost exchange module that performs a no-op."""
def forward(
self,
data: AtomicDataDict.Type,
ghost_included: bool,
) -> AtomicDataDict.Type:
return data
@model_modifier(persistent=True, private=True)
@classmethod
def enable_LAMMPSMLIAPGhostExchange(cls, model):
"""Enable LAMMPS ML-IAP ghost exchange for inference in LAMMPS ML-IAP."""
from ._ghost_exchange_lmp_mliap import LAMMPSMLIAPGhostExchangeModule
def factory(old):
new = LAMMPSMLIAPGhostExchangeModule(
field=old.field,
irreps_in=old.irreps_in,
)
return new
return replace_submodules(model, cls, factory)
|