model/nn/atomwise.py

# This file is a part of the `nequip` package. Please see LICENSE and README at the root for information on using it.
import torch
import torch.nn.functional

from e3nn.o3._linear import Linear

from onescience.datapipes.materials.nequip import AtomicDataDict
from onescience.datapipes.materials.nequip._key_registry import get_field_type
from ._graph_mixin import GraphModuleMixin
from .utils import scatter
from .model_modifier_utils import model_modifier, replace_submodules
from onescience.utils.nequip.internal.global_dtype import _GLOBAL_DTYPE

from typing import Optional, List, Dict, Union


class AtomwiseOperation(GraphModuleMixin, torch.nn.Module):
    def __init__(self, operation, field: str, irreps_in=None):
        super().__init__()
        self.operation = operation
        self.field = field
        self._init_irreps(
            irreps_in=irreps_in,
            my_irreps_in={field: operation.irreps_in},
            irreps_out={field: operation.irreps_out},
        )

    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
        data[self.field] = self.operation(data[self.field])
        return data


class AtomwiseLinear(GraphModuleMixin, torch.nn.Module):
    def __init__(
        self,
        field: str = AtomicDataDict.NODE_FEATURES_KEY,
        out_field: Optional[str] = None,
        irreps_in=None,
        irreps_out=None,
    ):
        super().__init__()
        self.field = field
        out_field = out_field if out_field is not None else field
        self.out_field = out_field
        if irreps_out is None:
            irreps_out = irreps_in[field]

        self._init_irreps(
            irreps_in=irreps_in,
            required_irreps_in=[field],
            irreps_out={out_field: irreps_out},
        )
        self.linear = Linear(
            irreps_in=self.irreps_in[field], irreps_out=self.irreps_out[out_field]
        )

    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
        data[self.out_field] = self.linear(data[self.field])
        return data


class AtomwiseReduce(GraphModuleMixin, torch.nn.Module):
    constant: float

    def __init__(
        self,
        field: str,
        out_field: Optional[str] = None,
        reduce="sum",
        avg_num_atoms=None,
        irreps_in={},
    ):
        super().__init__()
        assert reduce in ("sum", "mean", "normalized_sum")
        self.constant = 1.0
        if reduce == "normalized_sum":
            assert avg_num_atoms is not None
            self.constant = float(avg_num_atoms) ** -0.5
            reduce = "sum"
        self.reduce = reduce
        self.field = field
        self.out_field = f"{reduce}_{field}" if out_field is None else out_field
        self._init_irreps(
            irreps_in=irreps_in,
            irreps_out=(
                {self.out_field: irreps_in[self.field]}
                if self.field in irreps_in
                else {}
            ),
        )

    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
        field = data[self.field]
        if AtomicDataDict.BATCH_KEY in data:
            result = scatter(
                field,
                data[AtomicDataDict.BATCH_KEY],
                dim=0,
                dim_size=AtomicDataDict.num_frames(data),
                reduce=self.reduce,
            )
        else:
            # We can significantly simplify and avoid scatters
            if self.reduce == "sum":
                result = field.sum(dim=0, keepdim=True)
            elif self.reduce == "mean":
                result = field.mean(dim=0, keepdim=True)
            else:
                assert False
        if self.constant != 1.0:
            result = result * self.constant
        data[self.out_field] = result
        return data


class PerTypeScaleShift(GraphModuleMixin, torch.nn.Module):
    """Scale and/or shift a predicted per-atom property based on (learnable) per-species/type parameters.

    Note that scaling/shifting is always done casting into the global dtype (``float64``), even if ``model_dtype`` is a lower precision.

    If a single scalar is provided for scales/shifts, a shortcut implementation is used. Otherwise, a more expensive implementation that assigns separate scales/shifts to each atom type is used.

    If scales/shifts are trainable, the more expensive implementation that assigns separate scales/shifts to each atom type is used, even if a single scalar was provided for the initialization.
    """

    field: str
    out_field: str
    has_scales: bool
    has_shifts: bool
    scales_trainble: bool
    shifts_trainable: bool

    def __init__(
        self,
        type_names: List[str],
        field: str,
        out_field: Optional[str] = None,
        scales: Optional[Union[float, Dict[str, float]]] = None,
        shifts: Optional[Union[float, Dict[str, float]]] = None,
        scales_trainable: bool = False,
        shifts_trainable: bool = False,
        irreps_in={},
    ):
        super().__init__()
        self.type_names = type_names
        self.num_types = len(type_names)

        # === fields and irreps ===
        self.field = field
        self.out_field = field if out_field is None else out_field
        assert get_field_type(self.field) == "node"
        assert get_field_type(self.out_field) == "node"

        self._init_irreps(
            irreps_in=irreps_in,
            my_irreps_in={self.field: "0e"},  # input to shift must be a single scalar
            irreps_out={self.out_field: irreps_in[self.field]},
        )

        # === dtype ===
        self.out_dtype = _GLOBAL_DTYPE

        # === preprocess scales and shifts ===
        # we only accept single values or dicts
        # lists are no longer supported
        if isinstance(scales, list) or isinstance(shifts, list):
            raise ValueError(
                "\n\nLists are no longer supported for per-type energy scales and shifts. Please use dicts that map from the model's `type_names` as keys to the relevant scale or shift values. For example, the following\n\n  per_type_energy_shifts: [1, 2, 3]\n\nshould be changed to\n\n  per_type_energy_shifts:\n    C: 1\n    H: 2\n    O: 3\n\n"
            )

        # single valued case
        if isinstance(scales, float) or isinstance(scales, int):
            scales = [scales]
        if isinstance(shifts, float) or isinstance(shifts, int):
            shifts = [shifts]

        # dict case
        if isinstance(scales, dict):
            assert set(self.type_names) == set(scales.keys())
            scales = [scales[name] for name in self.type_names]
        if isinstance(shifts, dict):
            assert set(self.type_names) == set(shifts.keys())
            shifts = [shifts[name] for name in self.type_names]

        # we convert everything to lists at this point for conversion into `torch.Tensor`s
        for sc_vars in (scales, shifts):
            if sc_vars is not None:
                assert isinstance(sc_vars, list)

        # === scales ===
        self.has_scales = scales is not None
        self.scales_trainable = scales_trainable
        if self.has_scales:
            scales = torch.as_tensor(scales, dtype=self.out_dtype)
            if self.scales_trainable and scales.numel() == 1:
                # effective no-op if self.num_types == 1
                scales = (
                    torch.ones(self.num_types, dtype=scales.dtype, device=scales.device)
                    * scales
                )
            assert scales.shape == (self.num_types,) or scales.numel() == 1, (
                f"Scales expected to have shape ({self.num_types},), but found {scales.shape}"
            )
            scales = scales.reshape(-1, 1)
            if self.scales_trainable:
                self.scales = torch.nn.Parameter(scales)
            else:
                self.register_buffer("scales", scales)
        else:
            self.register_buffer("scales", torch.Tensor())
        self.scales_shortcut = self.scales.numel() == 1

        # === shifts ===
        self.has_shifts = shifts is not None
        self.shifts_trainable = shifts_trainable
        if self.has_shifts:
            shifts = torch.as_tensor(shifts, dtype=self.out_dtype)
            if self.shifts_trainable and shifts.numel() == 1:
                # effective no-op if self.num_types == 1
                shifts = (
                    torch.ones(self.num_types, dtype=shifts.dtype, device=shifts.device)
                    * shifts
                )
            assert shifts.shape == (self.num_types,) or shifts.numel() == 1, (
                f"Shifts expected to have shape ({self.num_types},), but found {shifts.shape}"
            )
            shifts = shifts.reshape(-1, 1)
            if self.shifts_trainable:
                self.shifts = torch.nn.Parameter(shifts)
            else:
                self.register_buffer("shifts", shifts)
        else:
            self.register_buffer("shifts", torch.Tensor())
        self.shifts_shortcut = self.shifts.numel() == 1

    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
        """"""
        # shortcut if no scales or shifts found (only dtype promotion performed)
        if not (self.has_scales or self.has_shifts):
            data[self.out_field] = data[self.field].to(self.out_dtype)
            return data

        # === set up ===
        in_field = data[self.field]
        types = data[AtomicDataDict.ATOM_TYPE_KEY].view(-1)
        # to account for local-ghost truncation in ML-IAP
        types = types[: in_field.size(0)]

        if self.has_scales:
            if self.scales_shortcut:
                scales = self.scales
            else:
                scales = torch.nn.functional.embedding(types, self.scales)
        else:
            scales = self.scales  # dummy for torchscript

        if self.has_shifts:
            if self.shifts_shortcut:
                shifts = self.shifts
            else:
                shifts = torch.nn.functional.embedding(types, self.shifts)
        else:
            shifts = self.shifts  # dummy for torchscript

        # === explicit cast ===
        in_field = in_field.to(self.out_dtype)

        # === scale/shift ===
        if self.has_scales and self.has_shifts:
            # we can used an FMA for performance
            # addcmul computes
            # input + tensor1 * tensor2 elementwise
            # it will promote to widest dtype, which comes from shifts/scales
            in_field = torch.addcmul(shifts, scales, in_field)
        else:
            # fallback path for mix of enabled shifts and scales
            # multiplication / addition promotes dtypes already, so no cast is needed
            if self.has_scales:
                in_field = scales * in_field
            if self.has_shifts:
                in_field = shifts + in_field

        data[self.out_field] = in_field
        return data

    @model_modifier(persistent=True, private=False)
    @classmethod
    def modify_PerTypeScaleShift(
        cls,
        model,
        scales: Optional[Union[float, Dict[str, float]]] = None,
        shifts: Optional[Union[float, Dict[str, float]]] = None,
        scales_trainable: bool = False,
        shifts_trainable: bool = False,
    ):
        """Modify per-type scales and shifts of a model.

        The new ``scales`` and ``shifts`` should be provided as dicts.
        The keys must correspond to the ``type_names`` registered in the model being modified, and may not include all the possible ``type_names`` of the original model.
        For example, if one uses a pretrained model with 50 atom types, and seeks to only modify 3 per-atom shifts to be consistent with a fine-tuning dataset's DFT settings, one could use

        .. code-block:: yaml

            shifts:
              C: 1.23
              H: 0.12
              O: 2.13

        In this case, the per-type atomic energy shifts of the original model will be used for every other atom type, except for atom types with the new shifts specified.

        For more details on fine-tuning, see https://nequip.readthedocs.io/en/latest/guide/training-techniques/fine_tuning.html

        Args:
            scales: the new per-type atomic energy scales
            shifts: the new per-type atomic energy shifts (e.g. isolated atom energies of a dataset used for fine-tuning)
            scales_trainable (bool): whether the new scales are trainable
            shifts_trainable (bool): whether the new shifts are trainable
        """

        def _helper(sc_var, vname, old):
            # get original dict values
            orig_sc_var = getattr(old, vname).detach().cpu().reshape(-1).tolist()
            # handle special case of single-valued shortcut
            if len(orig_sc_var) != len(old.type_names):
                assert len(orig_sc_var) == 1
                orig_sc_var = orig_sc_var * len(old.type_names)
            new_sc_var = {name: val for name, val in zip(old.type_names, orig_sc_var)}
            if sc_var is not None:
                # preprocess to list if single number
                if isinstance(sc_var, float) or isinstance(sc_var, int):
                    sc_var = {name: sc_var for name in old.type_names}
                assert isinstance(sc_var, dict)
                assert all(k in old.type_names for k in sc_var.keys()), (
                    f"Provided `{vname}` dict keys ({sc_var.keys()}) do not match the expected type names of the model ({old.type_names})."
                )
                # update original model's dict with new dict entries
                new_sc_var.update(sc_var)
            # if no new values provided, we default to the original model's dict entries
            return new_sc_var

        def factory(old):
            return cls(
                type_names=old.type_names,
                field=old.field,
                out_field=old.out_field,
                scales=_helper(scales, "scales", old),
                shifts=_helper(shifts, "shifts", old),
                scales_trainable=scales_trainable,
                shifts_trainable=shifts_trainable,
                irreps_in=old.irreps_in,
            )

        return replace_submodules(model, cls, factory)

    def __repr__(self) -> str:
        return f"{self.__class__.__name__} \n  scales: {_format_type_vals(self.scales.reshape(-1).tolist(), self.type_names)}\n  shifts: {_format_type_vals(self.shifts.reshape(-1).tolist(), self.type_names)}"


def _format_type_vals(
    vals: List[float], type_names: List[str], element_formatter: str = ".6f"
) -> str:
    if vals is None or not vals:
        return f"[{', '.join(type_names)}: None]"

    if len(vals) == 1:
        return (f"[{', '.join(type_names)}: {{:{element_formatter}}}]").format(vals[0])
    elif len(vals) == len(type_names):
        return (
            "["
            + ", ".join(
                f"{{{i}[0]}}: {{{i}[1]:{element_formatter}}}" for i in range(len(vals))
            )
            + "]"
        ).format(*zip(type_names, vals))
    else:
        raise ValueError(
            f"Don't know how to format vals=`{vals}` for types {type_names} with element_formatter=`{element_formatter}`"
        )