Core Primitives

The core primitives in ParamRF are the pmrf.Model, pmrf.Frequency and pmrf.Param classes. Other classes, such as JAX arrays and pmrf.evaluators, help bring the rest of the library together.

The Model

pmrf.Model represents the base class for any RF model. All built-in models and components, such as Resistor, PhaseLine, CoaxialLine etc. inherit from this class.

Under the hood, ParamRF uses the Equinox library to interoperate with JAX. This means that pmrf.Model is an Equinox Module, a JAX PyTree and a Python dataclass. If these concepts are completely foreign to you, do not worry. The practical consequences of this are:

Models are immutable. Models cannot be mutated (model.x = 5.0 will throw an error). Rather, they represent “pure functions” which store their parameters alongside. This means that models should never contain any fields that can be derived from other fields (state should be “pure”), and should also not reference other objects in the model (using the same parameter object twice will result in two separate parameters!). Instead, pmrf.Model.at() can be used to manipulate a model. Although this approach may feel new to some, the end result is improved optimization performance, differentiation capabilities, and added flexibility.

Models are lazy. Since models only store their parameters (and not their S-matrix or frequency), they do not perform any computation at initialization. Rather, their response is only evaluated when one of their methods (e.g. pmrf.Model.s()) is passed a Frequency object as input. This is in contrast to other, purely object-oriented libraries (such as scikit-rf).

Models are JAX-native. Models are simply JAX PyTrees, meaning they can be passed around in any JAX context. This means that they can be compiled just-in-time (JIT) for enhanced performance and simulation on other platforms (GPUs, TPUs etc) and also used for many advanced JAX features (such as for vectorization via jax.vmap() and differentiation via jax.jacfwd()). See the JAX Overview section for more details.

For more information on how to combine and build custom models, see the Examples page.

Frequency, Parameters, and jnp.ndarray

The pmrf.Frequency class defines the axis over which models are evaluated. Ultimately, this is a lightweight wrapper around a JAX array (commonly imported as jnp.ndarray). Those unfamiliar with JAX can see either the JAX Overview section or have a look at JAX’s own quickstart guide. However, for those seeking a TLDR: the API is very similar to NumPy’s np.ndarray, with a few “rough edges”, for example control flow is handled quite differently.

All JAX arrays are treated as potential free parameters in ParamRF. However, in RF modeling, it is very common to want to specify parameter scaling, bounds, constraints or a prior probability distribution. To accomplish this, ParamRF exposes a parameter wrapper class pmrf.Param. This class eagerly casts to a jnp.ndarray, meaning that you can conveniently treat parameters as if they were regular arrays in your equations.

Power users may want to have a look at the Parax and distreqx documentation, which ParamRF uses for parameters and constraints under the hood.

The Evaluator and other classes

AbstractEvaluator is an interface for a callable class that can be used to “evaluate” a model over frequency. Its output is a jnp.ndarray (tensor or scalar). Evaluators can be used for extraction of model features, encapsulating loss/error/likelihood functions, and more. For example, to create a goal-oriented objective function for an optimization, the Goal evaluator can be used.

Evaluators are created automatically when fitting routines are called, such as minimize() or sample(). For example, specifying 's21_db' as a fitting feature for fit() creates a Feature evaluator, whilst specifying a loss or likelihood function creates a TargetLoss or MarginalLogLikelihood evaluator respectively.

Other core class interfaces include those in losses, likelihoods, discrepancy_models and noise_models. These help glue the rest of the library together, as well as enable more advanced features such as hyperparameter-based optimization and Gaussian process discrepancy modeling. See the Tutorials section for more information.