Saving and loading objects

jaxrts allows saving and restoring jaxrts.plasma_state.PlasmaState, jaxrts.setup.Setup and others to / from the disk.

The underlying functionality is implemented in jaxrts.saving, and relies on converting objects into serializable dictionaries. The results are stored as .json files.

Note that the preferred way to construct a PlasmaState or other objects is by just doing so programmatically. Many examples for doing so are given in the example gallery.

Warning

Saved files are not guaranteed to be forward-compatible, as you just store the instructions how to build an object, and not the object itself.

Saving an object

The following code write a jaxrts.plasma_state.PlasmaState to the file state.json.

from jax import numpy as jnp
import jaxrts
ureg = jaxrts.ureg

state = jaxrts.PlasmaState(
    ions=[jaxrts.Element("C")],
    Z_free=jnp.array([3.5]),
    mass_density=jnp.array([3.5]) * ureg.gram / ureg.centimeter**3,
    T_e=10 * ureg.electron_volt / ureg.k_B,
)

with open("state.json", "w") as f:
    jaxrts.saving.dump(state, f, indent=2)

Content of state.json

{
  "_type": "PlasmaState",
  "value": [
    {
      "Z_free": {
        "_type": "Array",
        "value": [
          3.5
        ]
      },
      "mass_density": {
        "_type": "Quantity",
        "value": [
          {
            "_type": "Array",
            "value": [
              3.5
            ]
          },
          [
            [
              "gram",
              1
            ],
            [
              "centimeter",
              -3
            ]
          ]
        ]
      },
      "T_e": {
        "_type": "Quantity",
        "value": [
          10.0,
          [
            [
              "electron_volt",
              1
            ],
            [
              "boltzmann_constant",
              -1
            ]
          ]
        ]
      },
      "T_i": {
        "_type": "Quantity",
        "value": [
          {
            "_type": "Array",
            "value": [
              10.0
            ]
          },
          [
            [
              "electron_volt",
              1
            ],
            [
              "boltzmann_constant",
              -1
            ]
          ]
        ]
      },
      "ion_core_radius": {
        "_type": "Quantity",
        "value": [
          {
            "_type": "Array",
            "value": [
              -1.0
            ]
          },
          [
            [
              "angstrom",
              1
            ]
          ]
        ]
      },
      "models": {
        "screening length": {
          "_type": "Model",
          "value": [
            "DebyeHueckelScreeningLength",
            [
              [],
              [
                "screening length"
              ]
            ]
          ]
        },
        "ee-lfc": {
          "_type": "Model",
          "value": [
            "LFCConstant",
            [
              [
                0.0
              ],
              [
                "ee-lfc"
              ]
            ]
          ]
        },
        "free-bound scattering": {
          "_type": "Model",
          "value": [
            "Neglect",
            [
              [],
              [
                "free-bound scattering"
              ]
            ]
          ]
        }
      }
    },
    {
      "ions": [
        {
          "_type": "Element",
          "value": "C"
        }
      ]
    }
  ]
}

This file contains all the information required to create the plasma state. Apart from Callables, which are stored using the dill packge [McKerns and Aivazis, 2010, McKerns et al., 2011], the file is human-readable and could be edited.

Loading an object

Warning

Similar to pickle, dill is not secure. Load only data you trust. In regard to jaxrts, this especially means that you should only load a jaxrts.setup.Setup from trustworthy sources, because the jaxrts.setup.Setup.instrument is stored using dill.

To load a plasma state, we also have to specify the :py:class:jpu.UnitRegistry instance, that should be used. This should be jaxrts.ureg.

with open("state.json") as f:
    state = jaxrts.saving.load(f, unit_reg = jaxrts.ureg)

 print(state)

ions                     :       C
mass densities (g/cc)    :    3.50
ionization               :    3.50
ion temperatures (eV)    :    10.0
e- temperature (eV)      :    10.0

Models attached
===============
screening length         : DebyeHueckelScreeningLength
ee-lfc                   : LFCConstant
free-bound scattering    : Neglect

Warning

Loading requires that all referenced jaxrts.models.Model classes are known to the importer. If custom models were used, they have to be registered by using the additional_mappings argument of jaxrts.saving.load().