fennol.models.fennix

  1from typing import Any, Sequence, Callable, Union, Optional, Tuple, Dict
  2from copy import deepcopy
  3import dataclasses
  4from collections import OrderedDict
  5
  6import jax
  7import jax.numpy as jnp
  8import flax.linen as nn
  9import numpy as np
 10from flax import serialization
 11from flax.core.frozen_dict import freeze, unfreeze, FrozenDict
 12from ..utils import AtomicUnits as au
 13
 14from .preprocessing import (
 15    GraphGenerator,
 16    PreprocessingChain,
 17    JaxConverter,
 18    atom_unpadding,
 19    check_input,
 20)
 21from .modules import MODULES, PREPROCESSING, FENNIXModules
 22
 23
 24@dataclasses.dataclass
 25class FENNIX:
 26    """
 27    Static wrapper for FENNIX models
 28
 29    The underlying model is a `fennol.models.modules.FENNIXModules` built from the `modules` dictionary
 30    which references registered modules in `fennol.models.modules.MODULES` and provides the parameters for initialization.
 31
 32    Since the model is static and contains variables, it must be initialized right away with either
 33    `example_data`, `variables` or `rng_key`. If `variables` is provided, it is used directly. If `example_data`
 34    is provided, the model is initialized with `example_data` and the resulting variables are stored
 35    in the wrapper. If only `rng_key` is provided, the model is initialized with a dummy system and the resulting.
 36    """
 37
 38    cutoff: Union[float, None]
 39    modules: FENNIXModules
 40    variables: Dict
 41    preprocessing: PreprocessingChain
 42    _apply: Callable[[Dict, Dict], Dict]
 43    _total_energy: Callable[[Dict, Dict], Tuple[jnp.ndarray, Dict]]
 44    _energy_and_forces: Callable[[Dict, Dict], Tuple[jnp.ndarray, jnp.ndarray, Dict]]
 45    _input_args: Dict
 46    _graphs_properties: Dict
 47    preproc_state: Dict
 48    energy_terms: Optional[Sequence[str]] = None
 49    _initializing: bool = True
 50    use_atom_padding: bool = False
 51
 52    def __init__(
 53        self,
 54        cutoff: float,
 55        modules: OrderedDict,
 56        preprocessing: OrderedDict = OrderedDict(),
 57        example_data=None,
 58        rng_key: Optional[jax.random.PRNGKey] = None,
 59        variables: Optional[dict] = None,
 60        energy_terms: Optional[Sequence[str]] = None,
 61        use_atom_padding: bool = False,
 62        graph_config: Dict = {},
 63        energy_unit: str = "Ha",
 64        **kwargs,
 65    ) -> None:
 66        """Initialize the FENNIX model
 67
 68        Arguments:
 69        ----------
 70        cutoff: float
 71            The cutoff radius for the model
 72        modules: OrderedDict
 73            The dictionary defining the sequence of FeNNol modules and their parameters.
 74        preprocessing: OrderedDict
 75            The dictionary defining the sequence of preprocessing modules and their parameters.
 76        example_data: dict
 77            Example data to initialize the model. If not provided, a dummy system is generated.
 78        rng_key: jax.random.PRNGKey
 79            The random key to initialize the model. If not provided, jax.random.PRNGKey(0) is used (should be avoided).
 80        variables: dict
 81            The variables of the model (i.e. weights, biases and all other tunable parameters).
 82            If not provided, the variables are initialized (usually at random)
 83        energy_terms: Sequence[str]
 84            The energy terms in the model output that will be summed to compute the total energy.
 85            If None, the total energy is always zero (useful for non-PES models).
 86        use_atom_padding: bool
 87            If True, the model will use atom padding for the input data.
 88            This is useful when one plans to frequently change the number of atoms in the system (for example during training).
 89        graph_config: dict
 90            Edit the graph configuration. Mostly used to change a long-range cutoff as a function of a simulation box size.
 91
 92        """
 93        self._input_args = {
 94            "cutoff": cutoff,
 95            "modules": OrderedDict(modules),
 96            "preprocessing": OrderedDict(preprocessing),
 97            "energy_unit": energy_unit,
 98        }
 99        self.energy_unit = energy_unit
100        self.Ha_to_model_energy = au.get_multiplier(energy_unit)
101        self.cutoff = cutoff
102        self.energy_terms = energy_terms
103        self.use_atom_padding = use_atom_padding
104
105        # add non-differentiable/non-jittable modules
106        preprocessing = deepcopy(preprocessing)
107        if cutoff is None:
108            preprocessing_modules = []
109        else:
110            prep_keys = list(preprocessing.keys())
111            graph_params = {"cutoff": cutoff, "graph_key": "graph"}
112            if len(prep_keys) > 0 and prep_keys[0] == "graph":
113                graph_params = {
114                    **graph_params,
115                    **preprocessing.pop("graph"),
116                }
117            graph_params = {**graph_params, **graph_config}
118
119            preprocessing_modules = [
120                GraphGenerator(**graph_params),
121            ]
122
123        for name, params in preprocessing.items():
124            key = str(params.pop("module_name")) if "module_name" in params else name
125            key = str(params.pop("FID")) if "FID" in params else key
126            mod = PREPROCESSING[key.upper()](**freeze(params))
127            preprocessing_modules.append(mod)
128
129        self.preprocessing = PreprocessingChain(
130            tuple(preprocessing_modules), use_atom_padding
131        )
132        graphs_properties = self.preprocessing.get_graphs_properties()
133        self._graphs_properties = freeze(graphs_properties)
134        # add preprocessing modules that should be differentiated/jitted
135        mods = [(JaxConverter, {})] + self.preprocessing.get_processors()
136        # mods = self.preprocessing.get_processors(return_list=True)
137
138        # build the model
139        modules = deepcopy(modules)
140        modules_names = []
141        for name, params in modules.items():
142            key = str(params.pop("module_name")) if "module_name" in params else name
143            key = str(params.pop("FID")) if "FID" in params else key
144            if name in modules_names:
145                raise ValueError(f"Module {name} already exists")
146            modules_names.append(name)
147            params["name"] = name
148            mod = MODULES[key.upper()]
149            fields = [f.name for f in dataclasses.fields(mod)]
150            if "_graphs_properties" in fields:
151                params["_graphs_properties"] = graphs_properties
152            if "_energy_unit" in fields:
153                params["_energy_unit"] = energy_unit
154            mods.append((mod, params))
155
156        self.modules = FENNIXModules(mods)
157
158        self.__apply = self.modules.apply
159        self._apply = jax.jit(self.modules.apply)
160
161        self.set_energy_terms(energy_terms)
162
163        # initialize the model
164
165        inputs, rng_key = self.reinitialize_preprocessing(rng_key, example_data)
166
167        if variables is not None:
168            self.variables = variables
169        elif rng_key is not None:
170            self.variables = self.modules.init(rng_key, inputs)
171        else:
172            raise ValueError(
173                "Either variables or a jax.random.PRNGKey must be provided for initialization"
174            )
175
176        self._initializing = False
177
178    def set_energy_terms(
179        self, energy_terms: Union[Sequence[str], None], jit: bool = True
180    ) -> None:
181        """Set the energy terms to be computed by the model and prepare the energy and force functions."""
182        object.__setattr__(self, "energy_terms", energy_terms)
183        if isinstance(energy_terms, str):
184            energy_terms = [energy_terms]
185        
186        if energy_terms is None or len(energy_terms) == 0:
187
188            def total_energy(variables, data):
189                out = self.__apply(variables, data)
190                coords = out["coordinates"]
191                nsys = out["natoms"].shape[0]
192                nat = coords.shape[0]
193                dtype = coords.dtype
194                e = jnp.zeros(nsys, dtype=dtype)
195                eat = jnp.zeros(nat, dtype=dtype)
196                out["total_energy"] = e
197                out["atomic_energies"] = eat
198                return e, out
199
200            def energy_and_forces(variables, data):
201                e, out = total_energy(variables, data)
202                f = jnp.zeros_like(out["coordinates"])
203                out["forces"] = f
204                return e, f, out
205
206            def energy_and_forces_and_virial(variables, data):
207                e, f, out = energy_and_forces(variables, data)
208                v = jnp.zeros(
209                    (out["natoms"].shape[0], 3, 3), dtype=out["coordinates"].dtype
210                )
211                out["virial_tensor"] = v
212                return e, f, v, out
213
214        else:
215            # build the energy and force functions
216            def total_energy(variables, data):
217                out = self.__apply(variables, data)
218                atomic_energies = 0.0
219                system_energies = 0.0
220                species = out["species"]
221                nsys = out["natoms"].shape[0]
222                for term in self.energy_terms:
223                    e = out[term]
224                    if e.ndim > 1 and e.shape[-1] == 1:
225                        e = jnp.squeeze(e, axis=-1)
226                    if e.shape[0] == nsys and nsys != species.shape[0]:
227                        system_energies += e
228                        continue
229                    assert e.shape == species.shape
230                    atomic_energies += e
231                # atomic_energies = jnp.squeeze(atomic_energies, axis=-1)
232                if isinstance(atomic_energies, jnp.ndarray):
233                    if "true_atoms" in out:
234                        atomic_energies = jnp.where(
235                            out["true_atoms"], atomic_energies, 0.0
236                        )
237                    out["atomic_energies"] = atomic_energies
238                    energies = jax.ops.segment_sum(
239                        atomic_energies,
240                        data["batch_index"],
241                        num_segments=len(data["natoms"]),
242                    )
243                else:
244                    energies = 0.0
245
246                if isinstance(system_energies, jnp.ndarray):
247                    if "true_sys" in out:
248                        system_energies = jnp.where(
249                            out["true_sys"], system_energies, 0.0
250                        )
251                    out["system_energies"] = system_energies
252
253                out["total_energy"] = energies + system_energies
254                return out["total_energy"], out
255
256            def energy_and_forces(variables, data):
257                def _etot(variables, coordinates):
258                    energy, out = total_energy(
259                        variables, {**data, "coordinates": coordinates}
260                    )
261                    return energy.sum(), out
262
263                de, out = jax.grad(_etot, argnums=1, has_aux=True)(
264                    variables, data["coordinates"]
265                )
266                out["forces"] = -de
267
268                return out["total_energy"], out["forces"], out
269
270            # def energy_and_forces_and_virial(variables, data):
271            #     x = data["coordinates"]
272            #     batch_index = data["batch_index"]
273            #     if "cells" in data:
274            #         cells = data["cells"]
275            #         ## cells is a nbatchx3x3 matrix which lines are cell vectors (i.e. cells[0,0,:] is the first cell vector of the first system)
276
277            #         def _etot(variables, coordinates, cells):
278            #             reciprocal_cells = jnp.linalg.inv(cells)
279            #             energy, out = total_energy(
280            #                 variables,
281            #                 {
282            #                     **data,
283            #                     "coordinates": coordinates,
284            #                     "cells": cells,
285            #                     "reciprocal_cells": reciprocal_cells,
286            #                 },
287            #             )
288            #             return energy.sum(), out
289
290            #         (dedx, dedcells), out = jax.grad(_etot, argnums=(1, 2), has_aux=True)(
291            #             variables, x, cells
292            #         )
293            #         f= -dedx
294            #         out["forces"] = f
295            #     else:
296            #         _,f,out = energy_and_forces(variables, data)
297
298            #     vir = -jax.ops.segment_sum(
299            #         f[:, :, None] * x[:, None, :],
300            #         batch_index,
301            #         num_segments=len(data["natoms"]),
302            #     )
303
304            #     if "cells" in data:
305            #         # dvir = jax.vmap(jnp.matmul)(dedcells, cells.transpose(0, 2, 1))
306            #         dvir = jnp.einsum("...ki,...kj->...ij", dedcells, cells)
307            #         nsys = data["natoms"].shape[0]
308            #         if cells.shape[0]==1 and nsys>1:
309            #             dvir = dvir / nsys
310            #         vir = vir + dvir
311
312            #     out["virial_tensor"] = vir
313
314            #     return out["total_energy"], f, vir, out
315
316            def energy_and_forces_and_virial(variables, data):
317                x = data["coordinates"]
318                scaling = jnp.asarray(
319                    np.eye(3)[None, :, :].repeat(data["natoms"].shape[0], axis=0)
320                )
321                def _etot(variables, coordinates, scaling):
322                    batch_index = data["batch_index"]
323                    coordinates = jax.vmap(jnp.matmul)(
324                        coordinates, scaling[batch_index]
325                    )
326                    inputs = {**data, "coordinates": coordinates}
327                    if "cells" in data:
328                        ## cells is a nbatchx3x3 matrix which lines are cell vectors (i.e. cells[0,0,:] is the first cell vector of the first system)
329                        cells = jax.vmap(jnp.matmul)(data["cells"], scaling)
330                        reciprocal_cells = jnp.linalg.inv(cells)
331                        inputs["cells"] = cells
332                        inputs["reciprocal_cells"] = reciprocal_cells
333                    energy, out = total_energy(variables, inputs)
334                    return energy.sum(), out
335
336                (dedx, vir), out = jax.grad(_etot, argnums=(1, 2), has_aux=True)(
337                    variables, x, scaling
338                )
339                f = -dedx
340                out["forces"] = f
341                out["virial_tensor"] = vir
342
343                return out["total_energy"], f, vir, out
344
345        object.__setattr__(self, "_total_energy_raw", total_energy)
346        if jit:
347            object.__setattr__(self, "_total_energy", jax.jit(total_energy))
348            object.__setattr__(self, "_energy_and_forces", jax.jit(energy_and_forces))
349            object.__setattr__(
350                self,
351                "_energy_and_forces_and_virial",
352                jax.jit(energy_and_forces_and_virial),
353            )
354        else:
355            object.__setattr__(self, "_total_energy", total_energy)
356            object.__setattr__(self, "_energy_and_forces", energy_and_forces)
357            object.__setattr__(
358                self, "_energy_and_forces_and_virial", energy_and_forces_and_virial
359            )
360
361    def get_gradient_function(
362        self,
363        *gradient_keys: Sequence[str],
364        jit: bool = True,
365        variables_as_input: bool = False,
366    ):
367        """Return a function that computes the energy and the gradient of the energy with respect to the keys in gradient_keys"""
368
369        def _energy_gradient(variables, data):
370            def _etot(variables, inputs):
371                if "strain" in inputs:
372                    scaling = inputs["strain"]
373                    batch_index = data["batch_index"]
374                    coordinates = inputs["coordinates"] if "coordinates" in inputs else data["coordinates"]
375                    coordinates = jax.vmap(jnp.matmul)(
376                        coordinates, scaling[batch_index]
377                    )
378                    inputs = {**inputs, "coordinates": coordinates}
379                    if "cells" in inputs or "cells" in data:
380                        cells = inputs["cells"] if "cells" in inputs else data["cells"]
381                        cells = jax.vmap(jnp.matmul)(cells, scaling)
382                        inputs["cells"] = cells
383                if "cells" in inputs:
384                    reciprocal_cells = jnp.linalg.inv(inputs["cells"])
385                    inputs = {**inputs, "reciprocal_cells": reciprocal_cells}
386                energy, out = self._total_energy_raw(variables, {**data, **inputs})
387                return energy.sum(), out
388
389            if "strain" in gradient_keys and "strain" not in data:
390                data = {**data, "strain": jnp.array(np.eye(3)[None, :, :].repeat(data["natoms"].shape[0], axis=0))}
391            inputs = {k: data[k] for k in gradient_keys}
392            de, out = jax.grad(_etot, argnums=1, has_aux=True)(variables, inputs)
393
394            return (
395                out["total_energy"],
396                de,
397                {**out, **{"dEd_" + k: de[k] for k in gradient_keys}},
398            )
399
400        if variables_as_input:
401            energy_gradient = _energy_gradient
402        else:
403
404            def energy_gradient(data):
405                return _energy_gradient(self.variables, data)
406
407        if jit:
408            return jax.jit(energy_gradient)
409        else:
410            return energy_gradient
411
412    def preprocess(self,use_gpu=False, verbose=False,**inputs) -> Dict[str, Any]:
413        """apply preprocessing to the input data
414
415        !!! This is not a pure function => do not apply jax transforms !!!"""
416        if self.preproc_state is None:
417            out, _ = self.reinitialize_preprocessing(example_data=inputs)
418        elif use_gpu:
419            do_check_input = self.preproc_state.get("check_input", True)
420            if do_check_input:
421                inputs = check_input(inputs)
422            preproc_state, inputs = self.preprocessing.atom_padding(
423                self.preproc_state, inputs
424            )
425            inputs = self.preprocessing.process(preproc_state, inputs)
426            preproc_state, state_up, out, overflow = (
427                self.preprocessing.check_reallocate(
428                    preproc_state, inputs
429                )
430            )
431            if verbose and overflow:
432                print("GPU preprocessing: nblist overflow => reallocating nblist")
433                print("size updates:", state_up)
434        else:
435            preproc_state, out = self.preprocessing(self.preproc_state, inputs)
436
437        object.__setattr__(self, "preproc_state", preproc_state)
438        return out
439
440    def reinitialize_preprocessing(
441        self, rng_key: Optional[jax.random.PRNGKey] = None, example_data=None
442    ) -> None:
443        ### TODO ###
444        if rng_key is None:
445            rng_key_pre = jax.random.PRNGKey(0)
446        else:
447            rng_key, rng_key_pre = jax.random.split(rng_key)
448
449        if example_data is None:
450            rng_key_sys, rng_key_pre = jax.random.split(rng_key_pre)
451            example_data = self.generate_dummy_system(rng_key_sys, n_atoms=10)
452
453        preproc_state, inputs = self.preprocessing.init_with_output(example_data)
454        object.__setattr__(self, "preproc_state", preproc_state)
455        return inputs, rng_key
456
457    def __call__(self, variables: Optional[dict] = None, gpu_preprocessing=False,**inputs) -> Dict[str, Any]:
458        """Apply the FENNIX model (preprocess + modules)
459
460        !!! This is not a pure function => do not apply jax transforms !!!
461        if you want to apply jax transforms, use  self._apply(variables, inputs) which is pure and preprocess the input using self.preprocess
462        """
463        if variables is None:
464            variables = self.variables
465        inputs = self.preprocess(use_gpu=gpu_preprocessing,**inputs)
466        output = self._apply(variables, inputs)
467        if self.use_atom_padding:
468            output = atom_unpadding(output)
469        return output
470
471    def total_energy(
472        self, variables: Optional[dict] = None, unit: Union[float, str] = None, gpu_preprocessing=False,**inputs
473    ) -> Tuple[jnp.ndarray, Dict]:
474        """compute the total energy of the system
475
476        !!! This is not a pure function => do not apply jax transforms !!!
477        if you want to apply jax transforms, use self._total_energy(variables, inputs) which is pure and preprocess the input using self.preprocess
478        """
479        if variables is None:
480            variables = self.variables
481        inputs = self.preprocess(use_gpu=gpu_preprocessing,**inputs)
482        # def print_shape(path,value):
483        #     if isinstance(value,jnp.ndarray):
484        #         print(path,value.shape)
485        #     else:
486        #         print(path,value)
487        # jax.tree_util.tree_map_with_path(print_shape,inputs)
488        _, output = self._total_energy(variables, inputs)
489        if self.use_atom_padding:
490            output = atom_unpadding(output)
491        e = output["total_energy"]
492        if unit is not None:
493            model_energy_unit = self.Ha_to_model_energy
494            if isinstance(unit, str):
495                unit = au.get_multiplier(unit)
496            e = e * (unit / model_energy_unit)
497        return e, output
498
499    def energy_and_forces(
500        self, variables: Optional[dict] = None, unit: Union[float, str] = None,gpu_preprocessing=False,**inputs
501    ) -> Tuple[jnp.ndarray, jnp.ndarray, Dict]:
502        """compute the total energy and forces of the system
503
504        !!! This is not a pure function => do not apply jax transforms !!!
505        if you want to apply jax transforms, use  self._energy_and_forces(variables, inputs) which is pure and preprocess the input using self.preprocess
506        """
507        if variables is None:
508            variables = self.variables
509        inputs = self.preprocess(use_gpu=gpu_preprocessing,**inputs)
510        _, _, output = self._energy_and_forces(variables, inputs)
511        if self.use_atom_padding:
512            output = atom_unpadding(output)
513        e = output["total_energy"]
514        f = output["forces"]
515        if unit is not None:
516            model_energy_unit = self.Ha_to_model_energy
517            if isinstance(unit, str):
518                unit = au.get_multiplier(unit)
519            e = e * (unit / model_energy_unit)
520            f = f * (unit / model_energy_unit)
521        return e, f, output
522
523    def energy_and_forces_and_virial(
524        self, variables: Optional[dict] = None, unit: Union[float, str] = None,gpu_preprocessing=False, **inputs
525    ) -> Tuple[jnp.ndarray, jnp.ndarray, Dict]:
526        """compute the total energy and forces of the system
527
528        !!! This is not a pure function => do not apply jax transforms !!!
529        if you want to apply jax transforms, use  self._energy_and_forces_and_virial(variables, inputs) which is pure and preprocess the input using self.preprocess
530        """
531        if variables is None:
532            variables = self.variables
533        inputs = self.preprocess(use_gpu=gpu_preprocessing,**inputs)
534        _, _, _, output = self._energy_and_forces_and_virial(variables, inputs)
535        if self.use_atom_padding:
536            output = atom_unpadding(output)
537        e = output["total_energy"]
538        f = output["forces"]
539        v = output["virial_tensor"]
540        if unit is not None:
541            model_energy_unit = self.Ha_to_model_energy
542            if isinstance(unit, str):
543                unit = au.get_multiplier(unit)
544            e = e * (unit / model_energy_unit)
545            f = f * (unit / model_energy_unit)
546            v = v * (unit / model_energy_unit)
547        return e, f, v, output
548
549    def remove_atom_padding(self, output):
550        """remove atom padding from the output"""
551        return atom_unpadding(output)
552
553    def get_model(self) -> Tuple[FENNIXModules, Dict]:
554        """return the model and its variables"""
555        return self.modules, self.variables
556
557    def get_preprocessing(self) -> Tuple[PreprocessingChain, Dict]:
558        """return the preprocessing chain and its state"""
559        return self.preprocessing, self.preproc_state
560
561    def __setattr__(self, __name: str, __value: Any) -> None:
562        if __name == "variables":
563            if __value is not None:
564                if not (
565                    isinstance(__value, dict)
566                    or isinstance(__value, OrderedDict)
567                    or isinstance(__value, FrozenDict)
568                ):
569                    raise ValueError(f"{__name} must be a dict")
570                object.__setattr__(self, __name, JaxConverter()(__value))
571            else:
572                raise ValueError(f"{__name} cannot be None")
573        elif __name == "preproc_state":
574            if __value is not None:
575                if not (
576                    isinstance(__value, dict)
577                    or isinstance(__value, OrderedDict)
578                    or isinstance(__value, FrozenDict)
579                ):
580                    raise ValueError(f"{__name} must be a FrozenDict")
581                object.__setattr__(self, __name, freeze(JaxConverter()(__value)))
582            else:
583                raise ValueError(f"{__name} cannot be None")
584
585        elif self._initializing:
586            object.__setattr__(self, __name, __value)
587        else:
588            raise ValueError(f"{__name} attribute of FENNIX model is immutable.")
589
590    def generate_dummy_system(
591        self, rng_key: jax.random.PRNGKey, box_size=None, n_atoms: int = 10
592    ) -> Dict[str, Any]:
593        """
594        Generate dummy system for initialization
595        """
596        if box_size is None:
597            box_size = 2 * self.cutoff
598            for g in self._graphs_properties.values():
599                cutoff = g["cutoff"]
600                if cutoff is not None:
601                    box_size = min(box_size, 2 * g["cutoff"])
602        coordinates = np.array(
603            jax.random.uniform(rng_key, (n_atoms, 3), maxval=box_size), dtype=np.float64
604        )
605        species = np.ones((n_atoms,), dtype=np.int32)
606        batch_index = np.zeros((n_atoms,), dtype=np.int32)
607        natoms = np.array([n_atoms], dtype=np.int32)
608        return {
609            "species": species,
610            "coordinates": coordinates,
611            # "graph": graph,
612            "batch_index": batch_index,
613            "natoms": natoms,
614        }
615
616    def summarize(
617        self, rng_key: jax.random.PRNGKey = None, example_data=None, **kwargs
618    ) -> str:
619        """Summarize the model architecture and parameters"""
620        if rng_key is None:
621            head = "Summarizing with example data:\n"
622            rng_key = jax.random.PRNGKey(0)
623        if example_data is None:
624            head = "Summarizing with dummy 10 atoms system:\n"
625            rng_key, rng_key_sys = jax.random.split(rng_key)
626            example_data = self.generate_dummy_system(rng_key_sys, n_atoms=10)
627        rng_key, rng_key_pre = jax.random.split(rng_key)
628        _, inputs = self.preprocessing.init_with_output(example_data)
629        return head + nn.tabulate(self.modules, rng_key, **kwargs)(inputs)
630
631    def to_dict(self):
632        """return a dictionary representation of the model"""
633        return {
634            **self._input_args,
635            "energy_terms": self.energy_terms,
636            "variables": deepcopy(self.variables),
637        }
638
639    def save(self, filename):
640        """save the model to a file"""
641        state_dict = self.to_dict()
642        state_dict["preprocessing"] = [
643            [k, v] for k, v in state_dict["preprocessing"].items()
644        ]
645        state_dict["modules"] = [[k, v] for k, v in state_dict["modules"].items()]
646        with open(filename, "wb") as f:
647            f.write(serialization.msgpack_serialize(state_dict))
648
649    @classmethod
650    def load(
651        cls,
652        filename,
653        use_atom_padding=False,
654        graph_config={},
655    ):
656        """load a model from a file"""
657        with open(filename, "rb") as f:
658            state_dict = serialization.msgpack_restore(f.read())
659        state_dict["preprocessing"] = {k: v for k, v in state_dict["preprocessing"]}
660        state_dict["modules"] = {k: v for k, v in state_dict["modules"]}
661        return cls(
662            **state_dict,
663            graph_config=graph_config,
664            use_atom_padding=use_atom_padding,
665        )