# 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 e3nn.o3._irreps import Irreps from e3nn.util.jit import compile_mode from onescience.datapipes.materials.nequip import AtomicDataDict from ._graph_mixin import GraphModuleMixin from .model_modifier_utils import model_modifier, replace_submodules @compile_mode("unsupported") class PartialForceOutput(GraphModuleMixin, torch.nn.Module): r"""Generate partial and total forces from an energy model. Args: func: the energy model vectorize: the vectorize option to ``torch.autograd.functional.jacobian``, false by default since it doesn't work well. """ vectorize: bool def __init__( self, func: GraphModuleMixin, vectorize: bool = False, vectorize_warnings: bool = False, ): super().__init__() self.func = func self.vectorize = vectorize if vectorize_warnings: # See https://pytorch.org/docs/stable/generated/torch.autograd.functional.jacobian.html torch._C._debug_only_display_vmap_fallback_warnings(True) # check and init irreps self._init_irreps( irreps_in=func.irreps_in, my_irreps_in={AtomicDataDict.PER_ATOM_ENERGY_KEY: Irreps("0e")}, irreps_out=func.irreps_out, ) self.irreps_out[AtomicDataDict.PARTIAL_FORCE_KEY] = Irreps("1o") self.irreps_out[AtomicDataDict.FORCE_KEY] = Irreps("1o") def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type: data = data.copy() out_data = {} def wrapper(pos: torch.Tensor) -> torch.Tensor: """Wrapper from pos to atomic energy""" nonlocal data, out_data data[AtomicDataDict.POSITIONS_KEY] = pos out_data = self.func(data) return out_data[AtomicDataDict.PER_ATOM_ENERGY_KEY].squeeze(-1) pos = data[AtomicDataDict.POSITIONS_KEY] partial_forces = torch.autograd.functional.jacobian( func=wrapper, inputs=pos, create_graph=self.training, # needed to allow gradients of this output during training vectorize=self.vectorize, ) partial_forces = partial_forces.negative() # output is [n_at, n_at, 3] out_data[AtomicDataDict.PARTIAL_FORCE_KEY] = partial_forces out_data[AtomicDataDict.FORCE_KEY] = partial_forces.sum(dim=0) return out_data @compile_mode("script") class ForceStressOutput(GraphModuleMixin, torch.nn.Module): r"""Compute forces (and stress if cell is provided) using autograd of an energy model. See: Knuth et. al. Comput. Phys. Commun 190, 33-50, 2015 https://pure.mpg.de/rest/items/item_2085135_9/component/file_2156800/content Args: func: the energy model to wrap """ do_derivatives: bool def __init__(self, func: GraphModuleMixin, do_derivatives: bool = True): super().__init__() self.func = func self.do_derivatives = do_derivatives # check and init irreps self._init_irreps( irreps_in=self.func.irreps_in.copy(), irreps_out=self.func.irreps_out.copy(), ) self.irreps_out[AtomicDataDict.FORCE_KEY] = "1o" self.irreps_out[AtomicDataDict.STRESS_KEY] = "1o" self.irreps_out[AtomicDataDict.VIRIAL_KEY] = "1o" self.irreps_out[AtomicDataDict.EDGE_FORCE_KEY] = "1o" # for torchscript compat self.register_buffer("_empty", torch.Tensor()) def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type: # short-circuit if not self.do_derivatives: return self.func(data) # === LOGIC BRANCHING NOTES === # if edge vectors not present, we assume that positions are present # and proceed with the usual procedure to compute forces, virials, stress # else, we compute edge forces # NOTE: if edge vectors are not present, we assume that it is for non-batched inference with no cell # at the point of making this change, it is specifically for LAMMPS-MLIAP compatibility if AtomicDataDict.EDGE_VECTORS_KEY not in data: if AtomicDataDict.BATCH_KEY in data: batch = data[AtomicDataDict.BATCH_KEY] num_batch: int = AtomicDataDict.num_frames(data) else: # Special case for efficiency batch = self._empty num_batch: int = 1 pos = data[AtomicDataDict.POSITIONS_KEY] has_cell: bool = AtomicDataDict.CELL_KEY in data if has_cell: orig_cell = data[AtomicDataDict.CELL_KEY] # Make the cell per-batch cell = orig_cell.view(-1, 3, 3).expand(num_batch, 3, 3) data[AtomicDataDict.CELL_KEY] = cell else: # torchscript orig_cell = self._empty cell = self._empty # Add the displacements # the GradientOutput will make them require grad # See SchNetPack code: # https://github.com/atomistic-machine-learning/schnetpack/blob/master/src/schnetpack/atomistic/model.py#L45 # SchNetPack issue: # https://github.com/atomistic-machine-learning/schnetpack/issues/165 # Paper they worked from: # Knuth et. al. Comput. Phys. Commun 190, 33-50, 2015 # https://pure.mpg.de/rest/items/item_2085135_9/component/file_2156800/content if num_batch > 1: displacement = torch.zeros( (num_batch, 3, 3), dtype=pos.dtype, device=pos.device, ) else: displacement = torch.zeros( (3, 3), dtype=pos.dtype, device=pos.device, ) displacement.requires_grad_(True) data["_displacement"] = displacement # in the above paper, the infinitesimal distortion is *symmetric* # so we symmetrize the displacement before applying it to # the positions/cell # This is not strictly necessary (reasoning thanks to Mario): # the displacement's asymmetric 1o term corresponds to an # infinitesimal rotation, which should not affect the final # output (invariance). # That said, due to numerical error, this will never be # exactly true. So, we symmetrize the deformation to # take advantage of this understanding and not rely on # the invariance here: symmetric_displacement = 0.5 * ( displacement + displacement.transpose(-1, -2) ) did_pos_req_grad: bool = pos.requires_grad pos.requires_grad_(True) if num_batch > 1: # bmm is natom in batch # batched [natom, 1, 3] @ [natom, 3, 3] -> [natom, 1, 3] -> [natom, 3] data[AtomicDataDict.POSITIONS_KEY] = pos + torch.bmm( pos.unsqueeze(-2), torch.index_select(symmetric_displacement, 0, batch), ).squeeze(-2) else: # (num_atoms, 3), (3, 3) -> (num_atoms, 3) data[AtomicDataDict.POSITIONS_KEY] = pos + torch.sum( pos.view(-1, 3, 1) * symmetric_displacement, 1 ) # assert torch.equal(pos, data[AtomicDataDict.POSITIONS_KEY]) # we only displace the cell if we have one: if has_cell: # bmm is num_batch in batch # here we apply the distortion to the cell as well # this is critical also for the correctness # if we didn't symmetrize the distortion, since without this # there would then be an infinitesimal rotation of the positions # but not cell, and it thus wouldn't be global and have # no effect due to equivariance/invariance. if num_batch > 1: # [n_batch, 3, 3] @ [n_batch, 3, 3] data[AtomicDataDict.CELL_KEY] = cell + torch.bmm( cell, symmetric_displacement ) else: # [3, 3] @ [3, 3] --- enforced to these shapes data[AtomicDataDict.CELL_KEY] = ( cell.view(3, 3) + torch.sum(cell.view(3, 3, 1) * symmetric_displacement, 1) ).view(1, 3, 3) # Call model and get gradients data = self.func(data) grads = torch.autograd.grad( [data[AtomicDataDict.TOTAL_ENERGY_KEY].sum()], [pos, data["_displacement"]], create_graph=self.training, # needed to allow gradients of this output during training ) # Put negative sign on forces forces = grads[0] if forces is None: # condition needed to unwrap optional for torchscript assert False, "failed to compute forces autograd" forces = torch.neg(forces) data[AtomicDataDict.FORCE_KEY] = forces # Store virial virial = grads[1] if virial is None: # condition needed to unwrap optional for torchscript assert False, "failed to compute virial autograd" virial = virial.view(num_batch, 3, 3) # we only compute the stress (1/V * virial) if we have a cell whose volume we can compute if has_cell: # ^ can only scale by cell volume if we have one...: # Rescale stress tensor # See https://github.com/atomistic-machine-learning/schnetpack/blob/master/src/schnetpack/atomistic/output_modules.py#L180 # See also https://en.wikipedia.org/wiki/Triple_product # See also https://gitlab.com/ase/ase/-/blob/master/ase/cell.py, # which uses np.abs(np.linalg.det(cell)) # First dim is batch, second is vec, third is xyz # Note the .abs(), since volume should always be positive # det is equal to a dot (b cross c) volume = torch.linalg.det(cell).abs().unsqueeze(-1) # NOTE: to support batching periodic and non-periodic structures together, # the data processing stage is responsible for ensuring that: # 1. non-periodic systems have a finite dummy cell to prevent infs in the division below # 2. stress labels for non-periodic systems are NaN and handled with `ignore_nan` in loss and metrics stress = virial / volume.view(num_batch, 1, 1) data[AtomicDataDict.CELL_KEY] = orig_cell else: stress = self._empty # torchscript data[AtomicDataDict.STRESS_KEY] = stress # see discussion in https://github.com/libAtoms/QUIP/issues/227 about sign convention # (and conventions docs page) # they say the standard convention is virial = -stress x volume # looking above this means that we need to pick up another negative sign for the virial # to fit this equation with the stress computed above virial = torch.neg(virial) data[AtomicDataDict.VIRIAL_KEY] = virial # Remove helper del data["_displacement"] if not did_pos_req_grad: # don't give later modules one that does pos.requires_grad_(False) else: # we differentiate wrt EDGE_VECTORS_KEY directly in this branch # NOTE: we only consider the case of non-batched inference, without a cell # so no batching, no training considerations, no cell # make `edge_vectors` requires grad edge_vectors = data[AtomicDataDict.EDGE_VECTORS_KEY] edge_vectors.requires_grad_(True) data[AtomicDataDict.EDGE_VECTORS_KEY] = edge_vectors # do energy model forward and backward data = self.func(data) edge_forces = torch.autograd.grad( [data[AtomicDataDict.TOTAL_ENERGY_KEY].sum()], [edge_vectors], # no training arg because we only consider inference )[0] # assert needed for TorchScript assert edge_forces is not None # NOTE: there shouldn't be a sign flip to match LAMMPS convention data[AtomicDataDict.EDGE_FORCE_KEY] = edge_forces return data @model_modifier(persistent=True, private=False) @classmethod def enable_ForceStressOutput(cls, model): """Enable force and stress computation.""" def factory(old): new = cls(func=old.func, do_derivatives=True) return new return replace_submodules(model, cls, factory) @model_modifier(persistent=True, private=False) @classmethod def disable_ForceStressOutput(cls, model): """Disable force and stress computation.""" def factory(old): new = cls(func=old.func, do_derivatives=False) return new return replace_submodules(model, cls, factory)