ParamRF: A JAX-native Framework for Declarative Circuit Modelling
For RF engineers and researchers, ParamRF provides a flexible, high-performance framework for circuit modeling and optimization, but it is an incremental tool built on existing technologies (JAX, Equinox).
ParamRF is a JAX-native Python library for parametric RF circuit modeling that compiles circuits into optimized algebraic functions, enabling GPU/TPU acceleration and automatic differentiation. It demonstrates fitting models to measured data using built-in optimizers and Bayesian samplers.
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.