Creating your own model

In jaxrts, we strongly encourage community-driven collaboration. Users are welcome to create, implement, or contribute new models to the codebase, targeting specific components of the simulation. This section provides a brief overview of how to implement a new model within the framework.

To showcase this, we’ll demonstrate how to create a simple, hypothetical model for the Local Field correction (LFC) that always returns \(\pi\), for any plasma conditions.

First, create a low-level function that calculates the respective quantity from a set of parameters

import jax.numpy as jnp
from .units import Quantity
...


   def lfc_pi(
      Zi: float,
      ne: Quantity,
      ni: Quantity,
      Te: Quantity,
      Ti: Quantity,
      Zbar: float | None = None
   ) -> Quantity:

   return jnp.pi

Next, create the corresponding jaxrts.models.Model class. Use the existing implementations for the same model type as a reference for structure and style.

import jax
import jaxrts
from jaxrts.setup import Setup
...

class PiLFC(jaxrts.models.Model):
   """
   Hypothetical LFC Model, in which the LFC is always pi.
   """

   # The allowed model keys for a plasma state.
   allowed_keys = ["ee-lfc"]
   __name__ = "PiLFC"

   # Citations keys for reference
   cite_keys = ["JohnDoe.2025"]

   def __init__(self):
      super().__init__()

   # This function is required by every model.
   @jax.jit
   def evaluate(self, plasma_state: "PlasmaState", setup: Setup) -> Quantity:
      return lfc_pi(
            plasma_state.Z_free,
            plasma_state.n_e,
            plasma_state.n_i,
            plasma_state.T_e,
            plasma_state.T_i,
            Zbar=plasma_state.Z_free
      )

   # The following is required to jit a Model
   # Here, 'children' are attributes of the class that can be traced
   # using jax, e.g. plasma_state's, floats etc., while aux_data are static arguments.
   def _tree_flatten(self):
      children = ()
      aux_data = (self.model_key)  # static values
      return (children, aux_data)

   @classmethod
   def _tree_unflatten(cls, aux_data, children):
      obj = object.__new__(cls)
      (obj.model_key) = aux_data
      return obj

Finally, register your new model by calling

jax.tree_util.register_pytree_node(
    PiLFC,
    PiLFC._tree_flatten,
    PiLFC._tree_unflatten,
)

Congratulations – you’ve successfully created your own model! You can add it to a jaxrts.plasmastate.PlasmaState by assigning it to the "ee-lfc" key:

state["ee-lfc"] = PiLFC()

Note

A jaxrts.models.Model provides the jaxrts.models.Model.check() and jaxrts.models.Model.prepare() methods. The former should be used to raise errors e.g., if the model is only applicable to one component systems. The other can be used in order to modify the passed jaxrts.plasmastate.PlasmaState, e.g., to set sane defaults for other, subsequent Models using jaxrts.plasmastate.PlasmaState.update_default_model().

Subclasses of Model

As some jaxrts.models.Model types require further methods, jaxrts.models provides subclasses for specific keys / models. Find a list bellow: