# SPDX-License-Identifier: LGPL-3.0-or-later from abc import ( abstractmethod, ) from collections.abc import ( Callable, ) from typing import ( Any, ) import array_api_compat import numpy as np from deepmd.dpmodel import ( DEFAULT_PRECISION, PRECISION_DICT, NativeOP, ) from deepmd.dpmodel.array_api import ( Array, ) from deepmd.dpmodel.common import ( get_xp_precision, to_numpy_array, ) from deepmd.dpmodel.output_def import ( OutputVariableDef, ) from deepmd.dpmodel.utils import ( AtomExcludeMask, FittingNet, NetworkCollection, ) from deepmd.dpmodel.utils.seed import ( child_seed, ) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) from deepmd.utils.env_mat_stat import ( StatItem, ) from deepmd.utils.finetune import ( get_index_between_two_maps, map_atom_exclude_types, ) from deepmd.utils.path import ( DPPath, ) from deepmd.utils.stat_file import ( load_required_items, ) from .base_fitting import ( BaseFitting, ) class GeneralFitting(NativeOP, BaseFitting): r"""General fitting class. Parameters ---------- var_name The name of the output variable. ntypes The number of atom types. dim_descrpt The dimension of the input descriptor. neuron Number of neurons :math:`N` in each hidden layer of the fitting net bias_atom_e Average energy per atom for each element. resnet_dt Time-step `dt` in the resnet construction: :math:`y = x + dt * \phi (Wx + b)` numb_fparam Number of frame parameter numb_aparam Number of atomic parameter rcond The condition number for the regression of atomic energy. tot_ener_zero Force the total energy to zero. Useful for the charge fitting. trainable If the weights of fitting net are trainable. Suppose that we have :math:`N_l` hidden layers in the fitting net, this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable. activation_function The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN| precision The precision of the embedding net parameters. Supported options are |PRECISION| layer_name : list[Optional[str]], optional The name of the each layer. If two layers, either in the same fitting or different fittings, have the same name, they will share the same neural network parameters. use_aparam_as_mask: bool, optional If True, the atomic parameters will be used as a mask that determines the atom is real/virtual. And the aparam will not be used as the atomic parameters for embedding. mixed_types If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types. exclude_types: list[int] Atomic contributions of the excluded atom types are set zero. remove_vaccum_contribution: list[bool], optional Remove vacuum contribution before the bias is added. The list assigned each type. For `mixed_types` provide `[True]`, otherwise it should be a list of the same length as `ntypes` signaling if or not removing the vacuum contribution for the atom types in the list. vacuum_ref: bool Reference the network output of every atom to the output of the same network for an isolated atom of the same type under the same frame parameters, atomic parameters and case embedding. The output of an atom without neighbors is then exactly ``bias_atom_e`` of its type. The call takes the vacuum descriptor of every type until :meth:`fold_vacuum_reference` has folded the reference into the bias or stored the table. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. seed: Optional[Union[int, list[int]]] Random seed for initializing the network parameters. default_fparam: list[float], optional The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net. """ # Deployment constants kept out of checkpoints; see # :meth:`fold_vacuum_reference`. A subclass extends this tuple rather than # replacing it. CONFIG_DERIVED_ARRAYS = ("vacuum_table", "folded_reference") def __init__( self, var_name: str, ntypes: int, dim_descrpt: int, neuron: list[int] = [120, 120, 120], resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, bias_atom_e: Array | None = None, rcond: float | None = None, tot_ener_zero: bool = False, trainable: list[bool] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, spin: Any = None, mixed_types: bool = True, exclude_types: list[int] = [], remove_vaccum_contribution: list[bool] | None = None, vacuum_ref: bool = False, type_map: list[str] | None = None, seed: int | list[int] | None = None, default_fparam: list[float] | None = None, ) -> None: self.var_name = var_name self.ntypes = ntypes self.dim_descrpt = dim_descrpt self.neuron = neuron self.resnet_dt = resnet_dt self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam self.dim_case_embd = dim_case_embd self.default_fparam = default_fparam self.rcond = rcond self.tot_ener_zero = tot_ener_zero self.trainable = trainable self.type_map = type_map if self.trainable is None: self.trainable = [True for ii in range(len(self.neuron) + 1)] if isinstance(self.trainable, bool): self.trainable = [self.trainable] * (len(self.neuron) + 1) self.activation_function = activation_function self.precision = precision if self.precision.lower() not in PRECISION_DICT: raise ValueError( f"Unsupported precision '{self.precision}'. Supported options are: {list(PRECISION_DICT.keys())}" ) self.prec = PRECISION_DICT[self.precision.lower()] self.layer_name = layer_name self.use_aparam_as_mask = use_aparam_as_mask self.spin = spin self.mixed_types = mixed_types # order matters, should be place after the assignment of ntypes self.reinit_exclude(exclude_types) if self.spin is not None: raise NotImplementedError("spin is not supported") self.remove_vaccum_contribution = remove_vaccum_contribution self.vacuum_ref = vacuum_ref self.eval_return_middle_output = False net_dim_out = self._net_out_dim() # init constants if bias_atom_e is None: self.bias_atom_e = np.zeros( [self.ntypes, net_dim_out], dtype=GLOBAL_NP_FLOAT_PRECISION ) else: assert bias_atom_e.shape == (self.ntypes, net_dim_out) self.bias_atom_e = bias_atom_e.astype(GLOBAL_NP_FLOAT_PRECISION) if self.numb_fparam > 0: self.fparam_avg = np.zeros(self.numb_fparam, dtype=self.prec) self.fparam_inv_std = np.ones(self.numb_fparam, dtype=self.prec) else: self.fparam_avg, self.fparam_inv_std = None, None if self.numb_aparam > 0: self.aparam_avg = np.zeros(self.numb_aparam, dtype=self.prec) self.aparam_inv_std = np.ones(self.numb_aparam, dtype=self.prec) else: self.aparam_avg, self.aparam_inv_std = None, None if self.dim_case_embd > 0: self.case_embd = np.zeros(self.dim_case_embd, dtype=self.prec) else: self.case_embd = None self.vacuum_table = None self.folded_reference = None if self.default_fparam is not None: if self.numb_fparam > 0: assert len(self.default_fparam) == self.numb_fparam, ( "default_fparam length mismatch!" ) self.default_fparam_tensor = np.array(self.default_fparam, dtype=self.prec) else: self.default_fparam_tensor = None # init networks in_dim = ( self.dim_descrpt + self.numb_fparam + (0 if self.use_aparam_as_mask else self.numb_aparam) + self.dim_case_embd ) self.nets = NetworkCollection( 1 if not self.mixed_types else 0, self.ntypes, network_type="fitting_network", networks=[ FittingNet( in_dim, net_dim_out, self.neuron, self.activation_function, self.resnet_dt, self.precision, bias_out=True, seed=child_seed(seed, ii), trainable=trainable, ) for ii in range(self.ntypes if not self.mixed_types else 1) ], ) def compute_input_stats( self, merged: Callable[[], list[dict]] | list[dict], protection: float = 1e-2, stat_file_path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. Parameters ---------- merged : Union[Callable[[], list[dict]], list[dict]] - list[dict]: A list of data samples from various data systems. Each element, `merged[i]`, is a data dictionary containing `keys`: `numpy.ndarray` originating from the `i`-th data system. - Callable[[], list[dict]]: A lazy function that returns data samples in the above format only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once. protection : float Divided-by-zero protection stat_file_path : Optional[DPPath] The path to the stat file. """ self._param_stats: dict[str, list[StatItem]] = {} if self.numb_fparam == 0 and self.numb_aparam == 0: # skip data statistics return # stat fparam if self.numb_fparam > 0: cached = load_required_items(stat_file_path, ["fparam"]) if cached is not None: fparam_stats = self._load_param_stats( cached["fparam"], "fparam", self.numb_fparam ) else: sampled = merged() if callable(merged) else merged for ii, frame in enumerate(sampled): if "find_fparam" not in frame: raise ValueError( f"numb_fparam > 0 but fparam is not acquired " f"for system {ii}." ) if not frame["find_fparam"]: raise ValueError( f"numb_fparam > 0 but no fparam data is provided " f"for system {ii}." ) xp_fp = array_api_compat.array_namespace(sampled[0]["fparam"]) cat_data = xp_fp.concat([frame["fparam"] for frame in sampled], axis=0) cat_data = xp_fp.reshape(cat_data, (-1, self.numb_fparam)) fparam_stats = [ StatItem( number=cat_data.shape[0], sum=float(xp_fp.sum(cat_data[:, ii])), squared_sum=float(xp_fp.sum(cat_data[:, ii] ** 2)), ) for ii in range(self.numb_fparam) ] if stat_file_path is not None: self._save_param_stats_to_file( stat_file_path, "fparam", fparam_stats ) self._param_stats["fparam"] = fparam_stats fparam_avg = np.array( [s.compute_avg() for s in fparam_stats], dtype=np.float64 ) fparam_std = np.array( [s.compute_std(protection=protection) for s in fparam_stats], dtype=np.float64, ) fparam_inv_std = 1.0 / fparam_std xp = array_api_compat.array_namespace(self.fparam_avg) self.fparam_avg = xp.asarray( fparam_avg, dtype=self.fparam_avg.dtype, device=array_api_compat.device(self.fparam_avg), ) self.fparam_inv_std = xp.asarray( fparam_inv_std, dtype=self.fparam_inv_std.dtype, device=array_api_compat.device(self.fparam_inv_std), ) # stat aparam if self.numb_aparam > 0: cached = load_required_items(stat_file_path, ["aparam"]) if cached is not None: aparam_stats = self._load_param_stats( cached["aparam"], "aparam", self.numb_aparam ) else: sampled = merged() if callable(merged) else merged for ii, frame in enumerate(sampled): if "find_aparam" not in frame: raise ValueError( f"numb_aparam > 0 but aparam is not acquired " f"for system {ii}." ) if not frame["find_aparam"]: raise ValueError( f"numb_aparam > 0 but no aparam data is provided " f"for system {ii}." ) xp_ap = array_api_compat.array_namespace(sampled[0]["aparam"]) sys_sumv = [] sys_sumv2 = [] sys_sumn = [] for ss_ in [frame["aparam"] for frame in sampled]: ss = xp_ap.reshape(ss_, (-1, self.numb_aparam)) sys_sumv.append(xp_ap.sum(ss, axis=0)) sys_sumv2.append(xp_ap.sum(ss * ss, axis=0)) sys_sumn.append(ss.shape[0]) sumv = xp_ap.sum(xp_ap.stack(sys_sumv), axis=0) sumv2 = xp_ap.sum(xp_ap.stack(sys_sumv2), axis=0) sumn = sum(sys_sumn) aparam_stats = [ StatItem( number=sumn, sum=float(sumv[ii]), squared_sum=float(sumv2[ii]), ) for ii in range(self.numb_aparam) ] if stat_file_path is not None: self._save_param_stats_to_file( stat_file_path, "aparam", aparam_stats ) self._param_stats["aparam"] = aparam_stats aparam_avg = np.array( [s.compute_avg() for s in aparam_stats], dtype=np.float64 ) aparam_std = np.array( [s.compute_std(protection=protection) for s in aparam_stats], dtype=np.float64, ) aparam_inv_std = 1.0 / aparam_std xp = array_api_compat.array_namespace(self.aparam_avg) self.aparam_avg = xp.asarray( aparam_avg, dtype=self.aparam_avg.dtype, device=array_api_compat.device(self.aparam_avg), ) self.aparam_inv_std = xp.asarray( aparam_inv_std, dtype=self.aparam_inv_std.dtype, device=array_api_compat.device(self.aparam_inv_std), ) @staticmethod def _save_param_stats_to_file( stat_file_path: DPPath, name: str, stats: list[StatItem], ) -> None: stat_file_path.mkdir(exist_ok=True, parents=True) fp = stat_file_path / name arr = np.array([[s.number, s.sum, s.squared_sum] for s in stats]) fp.save_numpy(arr) @staticmethod def _load_param_stats( arr: np.ndarray, name: str, numb: int, ) -> list[StatItem]: if arr.shape != (numb, 3): raise ValueError( f"Invalid {name} statistics shape {arr.shape}; expected ({numb}, 3)." ) return [ StatItem(number=arr[ii][0], sum=arr[ii][1], squared_sum=arr[ii][2]) for ii in range(numb) ] def get_param_stats(self) -> dict[str, list[StatItem]]: """Get the stored fparam/aparam statistics (populated by compute_input_stats).""" return getattr(self, "_param_stats", {}) @abstractmethod def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" pass def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.numb_fparam def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.numb_aparam def has_default_fparam(self) -> bool: """Check if the fitting has default frame parameters.""" return self.default_fparam is not None def get_default_fparam(self) -> list[float] | None: """Get the default frame parameters.""" return self.default_fparam def set_return_middle_output(self, enable: bool) -> None: """Enable or disable returning the middle (pre-last-layer) output. When enabled, the fitting network's ``call`` method will include a ``"middle_output"`` key in the returned dict, containing the hidden-layer activations before the final linear layer. Shape: ``[nframes, nloc, neuron[-1]]``. Raises ------ ValueError If ``enable`` is True but ``neuron`` is empty (no hidden layers). """ if enable and len(self.neuron) == 0: raise ValueError( "middle_output requires at least one hidden layer (neuron=[])" ) self.eval_return_middle_output = enable def _middle_output_def(self) -> list[OutputVariableDef]: """Return extra OutputVariableDefs for middle_output when enabled.""" if self.eval_return_middle_output and len(self.neuron) > 0: return [ OutputVariableDef( "middle_output", [self.neuron[-1]], reducible=False, r_differentiable=False, c_differentiable=False, ), ] return [] def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. Only atoms with selected atom types have atomic contribution to the result of the model. If returning an empty list, all atom types are selected. """ return [ii for ii in range(self.ntypes) if ii not in self.exclude_types] def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map def set_case_embd(self, case_idx: int) -> None: """ Set the case embedding of this fitting net by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ self.case_embd = np.eye(self.dim_case_embd, dtype=self.prec)[case_idx] def change_type_map( self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. """ assert self.type_map is not None, ( "'type_map' must be defined when performing type changing!" ) assert self.mixed_types, "Only models in mixed types can perform type changing!" remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) self.type_map = type_map self.ntypes = len(type_map) self.reinit_exclude(map_atom_exclude_types(self.exclude_types, remap_index)) if has_new_type: xp = array_api_compat.array_namespace(self.bias_atom_e) extend_shape = [len(type_map), *list(self.bias_atom_e.shape[1:])] extend_bias_atom_e = xp.zeros( extend_shape, dtype=self.bias_atom_e.dtype, device=array_api_compat.device(self.bias_atom_e), ) self.bias_atom_e = xp.concat([self.bias_atom_e, extend_bias_atom_e], axis=0) self.bias_atom_e = self.bias_atom_e[remap_index] # the stored references belong to the old type map self.vacuum_table = None self.folded_reference = None def __setitem__(self, key: str, value: Any) -> None: if key in ["bias_atom_e"]: self.bias_atom_e = value elif key in ["fparam_avg"]: self.fparam_avg = value elif key in ["fparam_inv_std"]: self.fparam_inv_std = value elif key in ["aparam_avg"]: self.aparam_avg = value elif key in ["aparam_inv_std"]: self.aparam_inv_std = value elif key in ["case_embd"]: self.case_embd = value elif key in ["scale"]: self.scale = value elif key in ["default_fparam_tensor"]: self.default_fparam_tensor = value else: raise KeyError(key) def __getitem__(self, key: str) -> Any: if key in ["bias_atom_e"]: return self.bias_atom_e elif key in ["fparam_avg"]: return self.fparam_avg elif key in ["fparam_inv_std"]: return self.fparam_inv_std elif key in ["aparam_avg"]: return self.aparam_avg elif key in ["aparam_inv_std"]: return self.aparam_inv_std elif key in ["case_embd"]: return self.case_embd elif key in ["scale"]: return self.scale elif key in ["default_fparam_tensor"]: return self.default_fparam_tensor else: raise KeyError(key) def reinit_exclude( self, exclude_types: list[int] = [], ) -> None: self.exclude_types = exclude_types self.emask = AtomExcludeMask(self.ntypes, self.exclude_types) def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", "@version": 5, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, "neuron": self.neuron, "resnet_dt": self.resnet_dt, "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, "default_fparam": self.default_fparam, "rcond": self.rcond, "activation_function": self.activation_function, "precision": self.precision, "mixed_types": self.mixed_types, "exclude_types": self.exclude_types, "vacuum_ref": self.vacuum_ref, "nets": self.nets.serialize(), "@variables": { "bias_atom_e": to_numpy_array(self.bias_atom_e), "case_embd": to_numpy_array(self.case_embd), "fparam_avg": to_numpy_array(self.fparam_avg), "fparam_inv_std": to_numpy_array(self.fparam_inv_std), "aparam_avg": to_numpy_array(self.aparam_avg), "aparam_inv_std": to_numpy_array(self.aparam_inv_std), }, "type_map": self.type_map, # not supported "tot_ener_zero": self.tot_ener_zero, "trainable": self.trainable, "layer_name": self.layer_name, "use_aparam_as_mask": self.use_aparam_as_mask, "spin": self.spin, } @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() data.pop("@class") data.pop("type") variables = data.pop("@variables") nets = data.pop("nets") obj = cls(**data) for kk in variables.keys(): obj[kk] = variables[kk] obj.nets = NetworkCollection.deserialize(nets) return obj def needs_vacuum_descriptor(self) -> bool: """Whether the call takes the vacuum descriptor of every type from the descriptor. A referencing fitting takes it until :meth:`fold_vacuum_reference` has folded the reference into the bias or stored the table. """ return self.vacuum_ref and self.vacuum_table is None def uniform_conditioning(self) -> bool: """Whether every atom receives the same conditioning columns. Frame parameters vary between frames and atomic parameters between atoms, while the case embedding is shared by all atoms of a call. """ return self.numb_fparam == 0 and ( self.numb_aparam == 0 or self.use_aparam_as_mask ) def conditioning_columns( self, descriptor: Array, fparam: Array | None, aparam: Array | None, ) -> Array | None: """Normalized conditioning columns appended to the descriptor of every atom. Parameters ---------- descriptor : Array Descriptor with shape (nf, nloc, nd); it supplies the frame and atom counts, the dtype and the device of the columns. fparam : Array, optional Frame parameters with shape (nf, numb_fparam). The default frame parameter is used when omitted. aparam : Array, optional Atomic parameters with shape (nf, nloc, numb_aparam). Returns ------- Array or None The frame parameters, atomic parameters and case embedding of every atom, normalized and concatenated, with shape (nf, nloc, ncond); None when the descriptor alone is the fitting input. """ xp = array_api_compat.array_namespace(descriptor) nf, nloc, _ = descriptor.shape device = array_api_compat.device(descriptor) # Fitting statistics remain NumPy for portable serialization; the # constants are materialized next to the runtime descriptor instead of # asking a backend namespace to reshape a foreign array directly. columns = [] if self.numb_fparam > 0: if fparam is None: assert self.default_fparam_tensor is not None default_fparam = xp.asarray( self.default_fparam_tensor, dtype=descriptor.dtype, device=device ) fparam = xp.tile( xp.reshape(default_fparam, (1, self.numb_fparam)), (nf, 1) ) try: fparam = xp.reshape(fparam, (nf, 1, self.numb_fparam)) except (ValueError, RuntimeError) as e: raise ValueError( f"input fparam: cannot reshape {fparam.shape} " f"into ({nf}, {self.numb_fparam})." ) from e fparam_device = array_api_compat.device(fparam) fparam_avg = xp.asarray( self.fparam_avg, dtype=fparam.dtype, device=fparam_device ) fparam_inv_std = xp.asarray( self.fparam_inv_std, dtype=fparam.dtype, device=fparam_device ) fparam = (fparam - fparam_avg) * fparam_inv_std columns.append(xp.broadcast_to(fparam, (nf, nloc, self.numb_fparam))) if self.numb_aparam > 0 and not self.use_aparam_as_mask: assert aparam is not None, "aparam should not be None" try: aparam = xp.reshape(aparam, (nf, nloc, self.numb_aparam)) except (ValueError, RuntimeError) as e: raise ValueError( f"input aparam: cannot reshape {aparam.shape} " f"into ({nf}, {nloc}, {self.numb_aparam})." ) from e aparam_device = array_api_compat.device(aparam) aparam_avg = xp.asarray( self.aparam_avg, dtype=aparam.dtype, device=aparam_device ) aparam_inv_std = xp.asarray( self.aparam_inv_std, dtype=aparam.dtype, device=aparam_device ) columns.append((aparam - aparam_avg) * aparam_inv_std) if self.dim_case_embd > 0: assert self.case_embd is not None case_embd = xp.asarray( self.case_embd, dtype=descriptor.dtype, device=device ) columns.append( xp.broadcast_to( xp.reshape(case_embd, (1, 1, -1)), (nf, nloc, self.dim_case_embd) ) ) if len(columns) == 0: return None return xp.concat(columns, axis=-1) def vacuum_input( self, vacuum_descriptor: Array | None, atype: Array, cond: Array | None, ) -> Array | None: """Fitting input rows of the vacuum references. The vacuum reference of an atom is an isolated atom of its type under the conditioning of the atom itself. With frame or atomic parameters the conditioning differs between atoms and the rows follow the atoms, with shape (nf, nloc, in_dim). Otherwise one row per type suffices, with shape (ntypes, in_dim), and :meth:`vacuum_output` gathers the network output by type. Parameters ---------- vacuum_descriptor : Array, optional Descriptor of an isolated atom of every type with shape (ntypes, dim_descrpt); None takes the table stored by :meth:`fold_vacuum_reference`. atype : Array Atom types with shape (nf, nloc). cond : Array, optional Conditioning columns of the atoms with shape (nf, nloc, ncond). Returns ------- Array or None The reference rows, or None when ``vacuum_ref`` is off. Raises ------ ValueError If ``vacuum_ref`` is on and the vacuum descriptor is missing or has the wrong shape. """ if not self.vacuum_ref: return None xp = array_api_compat.array_namespace(atype) if vacuum_descriptor is None: if self.vacuum_table is None: raise ValueError( "vacuum_ref requires the vacuum descriptor of every atom type" ) vacuum_descriptor = xp.asarray( self.vacuum_table, dtype=get_xp_precision(xp, self.precision), device=array_api_compat.device(atype), ) if tuple(vacuum_descriptor.shape) != (self.ntypes, self.dim_descrpt): raise ValueError( f"vacuum descriptor of shape {tuple(vacuum_descriptor.shape)} " f"does not match ({self.ntypes}, {self.dim_descrpt})" ) if not self.uniform_conditioning(): assert cond is not None nf, nloc = atype.shape x_vac = xp.reshape( xp.take(vacuum_descriptor, xp.reshape(atype, (-1,)), axis=0), (nf, nloc, self.dim_descrpt), ) return xp.concat([x_vac, cond], axis=-1) if self.dim_case_embd == 0: return vacuum_descriptor assert self.case_embd is not None case_embd = xp.asarray( self.case_embd, dtype=vacuum_descriptor.dtype, device=array_api_compat.device(vacuum_descriptor), ) return xp.concat( [ vacuum_descriptor, xp.broadcast_to( xp.reshape(case_embd, (1, -1)), (self.ntypes, self.dim_case_embd) ), ], axis=-1, ) def vacuum_output(self, vacuum_property: Array, atype: Array) -> Array: """Network output of the vacuum reference of every atom. Parameters ---------- vacuum_property : Array Network output on the rows of :meth:`vacuum_input`. atype : Array Atom types with shape (nf, nloc). Returns ------- Array The reference output of every atom with shape (nf, nloc, dim_out). """ if not self.uniform_conditioning(): return vacuum_property xp = array_api_compat.array_namespace(vacuum_property, atype) nf, nloc = atype.shape return xp.reshape( xp.take(vacuum_property, xp.reshape(atype, (-1,)), axis=0), (nf, nloc, vacuum_property.shape[-1]), ) def vacuum_property(self, vacuum_descriptor: Array | None) -> Array: """Network output on the vacuum descriptor of every type. Defined under uniform conditioning, where the output of a reference depends on its type alone. Parameters ---------- vacuum_descriptor : Array Descriptor of an isolated atom of every type with shape (ntypes, dim_descrpt). Returns ------- Array The reference output of every type with shape (ntypes, dim_out), in the precision of ``vacuum_descriptor``. """ if vacuum_descriptor is None: vacuum_descriptor = self.vacuum_table if vacuum_descriptor is None: raise ValueError( "vacuum_ref requires the vacuum descriptor of every atom type" ) xp = array_api_compat.array_namespace(vacuum_descriptor) atype = xp.arange( self.ntypes, dtype=xp.int64, device=array_api_compat.device(vacuum_descriptor), ) xx_vac = self.vacuum_input( xp.astype(vacuum_descriptor, get_xp_precision(xp, self.precision)), atype, None, ) if self.mixed_types: out = self.nets[()](xx_vac) else: out = xp.concat( [ self.nets[(type_i,)](xx_vac[type_i : type_i + 1, :]) for type_i in range(self.ntypes) ], axis=0, ) return xp.astype(out, vacuum_descriptor.dtype) def fold_vacuum_reference(self, vacuum_descriptor: Array) -> None: """Fold the vacuum reference into the fitting so the call needs no reference atoms. Under uniform conditioning the reference output of an atom is a constant of its type, so subtracting it from ``bias_atom_e`` yields the same outputs as the referenced call and the option is switched off; the subtracted outputs are kept as ``folded_reference`` so that :meth:`readout_reference` still knows the bias the fold moved. With frame or atomic parameters the reference output varies between atoms, so the vacuum descriptor is stored instead and the call evaluates the references from the stored table. Both are deployment constants: an exported model bakes them, checkpoints leave them out, and a fitting loaded from a checkpoint takes the reference from the descriptor again. Parameters ---------- vacuum_descriptor : Array Descriptor of an isolated atom of every type with shape (ntypes, dim_descrpt). """ if not self.vacuum_ref: return if not self.uniform_conditioning(): self.vacuum_table = to_numpy_array(vacuum_descriptor) return reference = to_numpy_array(self.vacuum_property(vacuum_descriptor)).astype( GLOBAL_NP_FLOAT_PRECISION ) self.folded_reference = reference self["bias_atom_e"] = to_numpy_array(self.bias_atom_e) - reference self.vacuum_ref = False def readout_reference(self) -> Array: """Per-type constant the readout gate of a bridging window fades to. It is the fitting bias ``bias_atom_e`` of every type, which under the vacuum reference is the output of an isolated atom of that type. A fold under uniform conditioning moves the reference outputs out of the bias; the stored copy restores the bias the gate reads, so the gate is the same function of the atom before and after the fold. Returns ------- Array The reference of every type with shape (ntypes, dim_out). """ if self.folded_reference is None: return self.bias_atom_e xp = array_api_compat.array_namespace(self.bias_atom_e) return self.bias_atom_e + xp.asarray( self.folded_reference, dtype=self.bias_atom_e.dtype, device=array_api_compat.device(self.bias_atom_e), ) def fold_case_embedding(self) -> None: """Fold the case embedding into the first layer so the call needs no case columns. The case embedding is the same trailing block of input columns for every atom, so its product with the rows of the first-layer weight that read it is a constant of the pre-activation. Adding that product to the bias and removing the rows yields a fitting without the case embedding that gives the same outputs. Export folds the fixed case of a deployed model, so that the frozen branch of a multi-task model with a shared fitting network is a plain fitting network. The fold rebuilds the first layers in place, so it applies to a fitting that no other branch shares. A first layer whose residual adds its input to the output also carries the case embedding past the activation; such a fitting keeps it. """ if self.dim_case_embd == 0: return networks = [ self.nets[ii] for ii in range(self.ntypes if not self.mixed_types else 1) ] for net in networks: width_in, width_out = net[0].w.shape if net[0].resnet and width_out in (width_in, 2 * width_in): return case_embd = to_numpy_array(self.case_embd).astype(np.float64) for net in networks: data = net[0].serialize() variables = data["@variables"] weight = variables["w"] kept = weight.shape[0] - self.dim_case_embd variables["b"] = ( variables["b"].astype(np.float64) + case_embd @ weight[kept:].astype(np.float64) ).astype(weight.dtype) variables["w"] = weight[:kept] # The residual is inactive at the original input width, checked # above; clearing it keeps the narrower input from activating it. data["resnet"] = False net[0] = type(net[0]).deserialize(data) net.in_dim = kept self.dim_case_embd = 0 self.case_embd = None def _call_common( self, descriptor: Array, atype: Array, gr: Array | None = None, g2: Array | None = None, h2: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, vacuum_descriptor: Array | None = None, node_gate: Array | None = None, return_middle_output: bool = False, ) -> dict[str, Array]: """Calculate the fitting. Parameters ---------- descriptor input descriptor. shape: nf x nloc x nd atype the atom type. shape: nf x nloc gr The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3 g2 The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng h2 The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3 fparam The frame parameter. shape: nf x nfp. nfp being `numb_fparam` aparam The atomic parameter. shape: nf x nloc x nap. nap being `numb_aparam` vacuum_descriptor The descriptor of an isolated atom of every type, required by ``vacuum_ref``. shape: ntypes x nd node_gate Per-atom source gate of a bridged descriptor; the learned part of the output fades with it. shape: nf x nloc x 1 return_middle_output Whether this call also returns the input of the output layer under the ``middle_output`` key: the activation after the last hidden layer, or the fitting input for a fitting without hidden layers. The key is declared by the output definition only when :meth:`set_return_middle_output` requests it for every call. """ xp = array_api_compat.array_namespace(descriptor, atype) nf, nloc, nd = descriptor.shape net_dim_out = self._net_out_dim() # check input dim if nd != self.dim_descrpt: raise ValueError( "get an input descriptor of dim {nd}," "which is not consistent with {self.dim_descrpt}." ) return_middle_output = return_middle_output or self.eval_return_middle_output # === Step 1. Assemble the fitting input === # The conditioning columns are shared by the atoms and by their vacuum # references, so that the reference of an atom differs from the atom # in its descriptor only. xx = descriptor cond = self.conditioning_columns(descriptor, fparam, aparam) # ``remove_vaccum_contribution`` subtracts the network output for a zero # descriptor under the same conditioning columns. xx_zeros = ( None if self.remove_vaccum_contribution is None else xp.zeros_like(xx) ) if cond is not None: xx = xp.concat([xx, cond], axis=-1) if xx_zeros is not None: xx_zeros = xp.concat([xx_zeros, cond], axis=-1) xx_vac = self.vacuum_input(vacuum_descriptor, atype, cond) # === Step 2. Evaluate the fitting networks === # Each network evaluates its hidden layers once, and the output layer # reads their activation, which is the middle output. results: dict[str, Array] = {} if not self.mixed_types: outs = xp.zeros( [nf, nloc, net_dim_out], dtype=get_xp_precision(xp, self.precision), device=array_api_compat.device(descriptor), ) if return_middle_output: middle_dim = self.neuron[-1] if len(self.neuron) > 0 else xx.shape[-1] middle_outs = xp.zeros( [nf, nloc, middle_dim], dtype=get_xp_precision(xp, self.precision), device=array_api_compat.device(descriptor), ) for type_i in range(self.ntypes): mask = xp.tile( xp.reshape((atype == type_i), (nf, nloc, 1)), (1, 1, net_dim_out), ) hidden = self.nets[(type_i,)].call_until_last(xx) atom_property = self.nets[(type_i,)].call_last(hidden) if self.remove_vaccum_contribution is not None and not ( len(self.remove_vaccum_contribution) > type_i and not self.remove_vaccum_contribution[type_i] ): assert xx_zeros is not None atom_property -= self.nets[(type_i,)](xx_zeros) if xx_vac is not None: # The mask below keeps the atoms of type ``type_i`` alone, so # the network of the type runs on its own reference row when # the references are per type, and on the per-atom rows # otherwise. reference = ( xx_vac[type_i : type_i + 1] if self.uniform_conditioning() else xx_vac ) atom_property = atom_property - self.nets[(type_i,)](reference) atom_property = xp.where( mask, atom_property, xp.zeros_like(atom_property) ) outs = outs + atom_property # Shape is [nframes, natoms[0], 1] if return_middle_output: mid_mask = xp.tile( xp.reshape((atype == type_i), (nf, nloc, 1)), (1, 1, middle_dim), ) middle_outs = middle_outs + xp.where( mid_mask, hidden, xp.zeros_like(hidden) ) else: middle_outs = self.nets[()].call_until_last(xx) outs = self.nets[()].call_last(middle_outs) if xx_zeros is not None: outs -= self.nets[()](xx_zeros) if xx_vac is not None: outs = outs - self.vacuum_output(self.nets[()](xx_vac), atype) bias_atom_e = xp.asarray( self.bias_atom_e, dtype=outs.dtype, device=array_api_compat.device(outs), ) outs += xp.reshape( xp.take( bias_atom_e, xp.reshape(atype, (-1,)), axis=0, ), (nf, nloc, net_dim_out), ) # === Step 3. Readout gate of a bridging window === # The learned part of the output is its deviation from the bias of # the type (see :meth:`readout_reference`), which is the output of an # isolated atom whenever the fitting references its atoms. A bridged # descriptor hands over the per-atom source gate, and the learned part # fades with it: an atom of a frozen pair keeps only the bias of its # type. if node_gate is not None: reference = xp.reshape( xp.take( xp.asarray( self.readout_reference(), dtype=outs.dtype, device=array_api_compat.device(outs), ), xp.reshape(atype, (-1,)), axis=0, ), (nf, nloc, net_dim_out), ) outs = reference + xp.astype(node_gate, outs.dtype) * (outs - reference) # nf x nloc exclude_mask = self.emask.build_type_exclude_mask(atype) exclude_mask = xp.astype(exclude_mask, xp.bool) # nf x nloc x nod outs = xp.where(exclude_mask[:, :, None], outs, xp.zeros_like(outs)) results[self.var_name] = outs if return_middle_output: results["middle_output"] = middle_outs return results def call_graph( self, descriptor: Array, atype: Array, gr: Array | None = None, g2: Array | None = None, h2: Array | None = None, fparam: Array | None = None, aparam: Array | None = None, vacuum_descriptor: Array | None = None, node_gate: Array | None = None, return_middle_output: bool = False, ) -> dict[str, Array]: """Graph-native (flat node axis) fitting forward. The node axis is flat ``(N,)``. This reuses the dense forward by treating the node axis as ``nf'=N`` single-atom frames (``nloc'=1``) -- an internal, encapsulated workaround, verified bit-identical to the dense call. Parameters ---------- descriptor input descriptor. N x nd atype the atom type. N gr equivariant single-particle representation. N x ng x 3 g2 the rotationally invariant pair-partical representation. unused by this fitting; passed through to the dense call. h2 the rotationally equivariant pair-partical representation. unused by this fitting; passed through to the dense call. fparam NODE-level frame parameter (already gathered by frame_id). N x nfp aparam atomic parameter. N x nap vacuum_descriptor the descriptor of an isolated atom of every type, required by ``vacuum_ref``. ntypes x nd node_gate per-atom source gate of a bridged descriptor; the learned part of the output fades with it. N return_middle_output whether this call also returns the input of the output layer under the ``middle_output`` key: the activation after the last hidden layer, or the fitting input for a fitting without hidden layers. Returns ------- result_dict the fitting result on the flat node axis. each value N x *shape """ import array_api_compat xp = array_api_compat.array_namespace(descriptor, atype) n, nd = descriptor.shape d1 = xp.reshape(descriptor, (n, 1, nd)) a1 = xp.reshape(atype, (n, 1)) g1 = None if gr is None else xp.reshape(gr, (n, 1, gr.shape[-2], 3)) if aparam is not None and len(aparam.shape) != 2: # enforce the flat contract loudly: a rectangular (nf, nloc, nda) # aparam with nf*nloc == N would silently reshape into the right # element order here but hand torch.export an unprovable # N == nf*nloc relation (and misalign rows for any other layout). raise ValueError( "graph-route aparam must be flat (N, nda) on the node axis; " f"got a rank-{len(aparam.shape)} array of shape {aparam.shape}" ) ap1 = None if aparam is None else xp.reshape(aparam, (n, 1, aparam.shape[-1])) # fparam: dense API expects (nf, nfp); here nf'=N single-atom frames, so the # node-level (N, nfp) IS the per-(pseudo)frame param -- tiled over nloc'=1. # Only referencing fittings take the keyword; it travels with a table. vacuum_kwargs = ( {} if vacuum_descriptor is None else {"vacuum_descriptor": vacuum_descriptor} ) gate_kwargs = ( {} if node_gate is None else {"node_gate": xp.reshape(node_gate, (n, 1, 1))} ) middle_kwargs = {"return_middle_output": True} if return_middle_output else {} ret = self.__call__( d1, a1, gr=g1, g2=g2, h2=h2, fparam=fparam, aparam=ap1, **vacuum_kwargs, **gate_kwargs, **middle_kwargs, ) return {kk: xp.reshape(vv, (n, *vv.shape[2:])) for kk, vv in ret.items()}