File size: 6,671 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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | # This file is a part of the `nequip` package. Please see LICENSE and README at the root for information on using it.
import torch
from onescience.datapipes.materials.nequip import AtomicDataDict
from onescience.utils.nequip.internal.aoti_metadata import NEQUIP_CUSTOM_OPS_LIBS_KEY
from ._graph_mixin import GraphModuleMixin
from typing import List, Dict, Any, Optional, Final
R_MAX_KEY: Final[str] = "r_max"
PER_EDGE_TYPE_CUTOFF_KEY: Final[str] = "per_edge_type_cutoff"
TYPE_NAMES_KEY: Final[str] = "type_names"
NUM_TYPES_KEY: Final[str] = "num_types"
MODEL_DTYPE_KEY: Final[str] = "model_dtype"
def _model_metadata_from_config(model_config: Dict[str, str]) -> Dict[str, str]:
model_metadata_dict = {}
# manually process everything
model_metadata_dict[MODEL_DTYPE_KEY] = model_config[MODEL_DTYPE_KEY]
model_metadata_dict[TYPE_NAMES_KEY] = " ".join(model_config[TYPE_NAMES_KEY])
model_metadata_dict[NUM_TYPES_KEY] = str(len(model_config[TYPE_NAMES_KEY]))
model_metadata_dict[R_MAX_KEY] = str(model_config[R_MAX_KEY])
if model_config.get(PER_EDGE_TYPE_CUTOFF_KEY, None) is not None:
from .embedding.utils import cutoff_partialdict_to_str
model_metadata_dict[PER_EDGE_TYPE_CUTOFF_KEY] = cutoff_partialdict_to_str(
model_config[PER_EDGE_TYPE_CUTOFF_KEY],
model_config[TYPE_NAMES_KEY],
model_config[R_MAX_KEY],
)
return model_metadata_dict
class GraphModel(GraphModuleMixin, torch.nn.Module):
"""Top-level module for any complete `nequip` model.
Manages top-level rescaling, dtypes, and more.
Args:
model (GraphModuleMixin): model to wrap
model_input_fields (Dict[str, Any]): input fields and their irreps
"""
model_input_fields: List[str]
is_graph_model: Final[bool] = True
is_compile_graph_model: Final[bool] = False
# ^ to identify `GraphModel` types from `nequip-package`d models (see https://pytorch.org/docs/stable/package.html#torch-package-sharp-edges)
_metadata: Dict[str, str]
def __init__(
self,
model: GraphModuleMixin,
model_config: Optional[Dict[str, str]] = None,
model_input_fields: Dict[str, Any] = {},
) -> None:
super().__init__()
irreps_in = {
# Things that always make sense as inputs:
AtomicDataDict.POSITIONS_KEY: "1o",
AtomicDataDict.EDGE_INDEX_KEY: None,
AtomicDataDict.EDGE_TRANSPOSE_PERM_KEY: None,
AtomicDataDict.EDGE_CELL_SHIFT_KEY: None,
AtomicDataDict.EDGE_VECTORS_KEY: "1o",
AtomicDataDict.CELL_KEY: "1o", # 3 of them, but still
AtomicDataDict.BATCH_KEY: None,
AtomicDataDict.NUM_NODES_KEY: None,
AtomicDataDict.ATOM_TYPE_KEY: None,
# for LAMMPS ML-IAP
AtomicDataDict.LMP_MLIAP_DATA_KEY: None,
AtomicDataDict.NUM_LOCAL_GHOST_NODES_KEY: None,
}
model_input_fields = AtomicDataDict._fix_irreps_dict(model_input_fields)
irreps_in.update(model_input_fields)
self._init_irreps(irreps_in=irreps_in, irreps_out=model.irreps_out)
for k, irreps in model.irreps_in.items():
if self.irreps_in.get(k, None) != irreps:
raise RuntimeError(
f"Model has `{k}` in its irreps_in with irreps `{irreps}`, but `{k}` is missing from/has inconsistent irreps in model_input_fields of `{self.irreps_in.get(k, 'missing')}`"
)
self.model = model
self.model_input_fields = list(self.irreps_in.keys())
# the following logic is for backward compatibility and to simplify unittests
self.model_dtype = torch.get_default_dtype()
self._metadata = {}
self.type_names = []
if model_config is not None:
self._metadata = _model_metadata_from_config(model_config)
self.type_names = self._metadata[TYPE_NAMES_KEY].split(" ")
model_dtype = {"float32": torch.float32, "float64": torch.float64}[
self._metadata[MODEL_DTYPE_KEY]
]
assert self.model_dtype == model_dtype
@property
@torch.jit.unused
def metadata(self) -> Dict[str, str]:
"""Get model metadata, including dynamic contributions from modules.
Collects metadata from all modules that override ``_get_metadata_contributions()``.
Dynamic contributions can override static config values.
Expected to be queried for inference workflows (but not for ``nequip-package``).
"""
out = self._metadata.copy()
# collect dynamic metadata from module tree
contributed_keys = {} # track which module provided each key
for name, module in self.model.named_modules():
if (
hasattr(module, "_is_graph_module_mixin")
and module._is_graph_module_mixin
):
contributions = module._get_metadata_contributions()
if not contributions:
continue
# detect conflicts between multiple modules
for key in contributions:
if key in contributed_keys:
raise ValueError(
f"Metadata conflict: modules '{contributed_keys[key]}' "
f"and '{name}' both contribute key '{key}'"
)
contributed_keys[key] = name
# update metadata (overrides static values if keys overlap)
out.update(contributions)
# update r_max if dynamic per-edge-type cutoffs were contributed
if PER_EDGE_TYPE_CUTOFF_KEY in contributed_keys:
cutoff_values = [float(x) for x in out[PER_EDGE_TYPE_CUTOFF_KEY].split()]
out[R_MAX_KEY] = str(max(cutoff_values))
# collect custom ops libs that need to be imported at AOTI load time
custom_ops_libs: set = set()
for m in self.model.modules():
custom_ops_libs.update(getattr(m, "_nequip_custom_ops_libs", ()))
if custom_ops_libs:
out[NEQUIP_CUSTOM_OPS_LIBS_KEY] = " ".join(sorted(custom_ops_libs))
return out
def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
# restrict the input data to allowed keys to prevent the model from directly using the dict from the outside,
# preventing weird pass-by-reference bugs
new_data: AtomicDataDict.Type = {}
for k in self.model_input_fields:
if k in data:
new_data[k] = data[k]
return self.model(new_data)
|