ParamRF: A JAX-native Framework for Declarative Circuit Modelling

TL;DR

ParamRF offers a JAX-native framework for declarative RF circuit modelling that leverages Equinox and PyTrees to enable high-performance, differentiable frequency-domain modelling. The paper details an object-oriented, immutable modelling paradigm with compositional capabilities, and demonstrates integration with both classical optimization and Bayesian inference backends for fitting and analysis. It presents automatic differentiation of model outputs with respect to frequency and parameters, and provides a unified fitting interface across backends such as SciPy, Optax, NumPyro, PolyChord, and BlackJAX. An end-to-end example fitting a PhysicalCoaxial model to measured data illustrates practical applicability and ecosystem compatibility, positioning ParamRF as a complementary tool within the scientific Python landscape and outlining avenues for future expansion.

Abstract

This work introduces ParamRF: a Python library for efficient, parametric modelling of radio frequency (RF) circuits. Built on top of the next-generation computational library JAX, as well as the object-oriented wrapper Equinox, the framework provides an easy-to-use, declarative modelling interface, without sacrificing performance. By representing circuits as JAX PyTrees and leveraging just-in-time compilation, models are compiled as pure functions into an optimized, algebraic graph. Since the resultant functions are JAX-native, this allows computation on CPUs, GPUs, or TPUs, providing integration with a wide range of solvers. Further, thanks to JAX's automatic differentiation, gradients with respect to both frequency and circuit parameters can be calculated for any circuit model outputs. This allows for more efficient optimization, as well as exciting new analysis opportunities. We showcase ParamRF's typical use-case of fitting a model to measured data via its built-in fitting engines, which include classical optimizers like L-BFGS and SLSQP, as well as modern Bayesian samplers such as PolyChord and BlackJAX. The result is a flexible framework for frequency-domain circuit modelling, fitting and analysis.