Model (pmrf.Model)

class pmrf.models.Model(*, name=None, separator='_', _z0=(50+0j), _param_groups=<factory>)

Bases: Module

Overview

Base class representing a computable RF network, referred to in ParamRF as a “Model”. This class is abstract and should not be instantiated directly. Derive from Model and override one of the primary property functions (e.g. __call__(), s(), a()).

The model is an Equinox Module (immutable, dataclass-like) and is treated as a JAX PyTree. Parameters are declared using standard dataclass field syntax with types like pmrf.Parameter.

Usage

• Define sub-classes with custom parameters and sub-models

• Construct models by passing parameters and/or submodels to the initializer (like a dataclass).

• Use with_xxx functions to update the model e.g. with_params(), with_replaced().

• Retrieve parameter information via methods such as named_params(), param_names(), flat_params(), etc..

See also the pmrf.fitting and pmrf.sampling modules for details on model fitting and sampling.

Examples

A PiCLC network (“foundational” model with fixed parameters and equations):

import jax.numpy as jnp
import pmrf as prf

class PiCLC(prf.Model):
    C1: prf.Parameter = 1.0e-12
    L:  prf.Parameter = 1.0e-9
    C2: prf.Parameter = 1.0e-12

    def a(self, freq: prf.Frequency) -> jnp.ndarray:
        w = freq.w
        Y1, Y2, Y3 = (1j * w * self.C1), (1j * w * self.C2), 1 / (1j * w * self.L)
        return jnp.array([
            [1 + Y2 / Y3,        1 / Y3],
            [Y1 + Y2 + Y1*Y2/Y3, 1 + Y1 / Y3],
        ]).transpose(2, 0, 1)

An RLC network (“circuit” model with free parameters built using cascading)

import pmrf as prf
from pmrf.models import Resistor, Capacitor, Inductor, Cascade
from pmrf.parameters import Uniform

class RLC(prf.Model):
    res: Resistor = Resistor(Uniform(9.0, 11.0))
    ind: Inductor = Inductor(Uniform(0.0, 10.0, scale=1e-9))
    cap: Capacitor = Capacitor(Uniform(0.0, 10.0, scale=1e-12))

    def __call__(self) -> prf.Model:
        return self.res ** self.ind ** self.cap.terminated()

name: str | None = None

separator: str = '_'

DEFAULT_PARAMS = {}

DEFAULT_PARAM_NAMES = []

property primary_function: Callable[[Frequency], Array]

Primary function (s or a) as a callable.

Return type:

Callable[[Frequency], jnp.ndarray]

property primary_property: str

Name of the primary property (e.g. "s", "a").

The primary property is the first overridden among PRIMARY_PROPERTIES. If __call__ is overridden, delegation occurs to the returned model.

Return type:

str

Raises:

NotImplementedError – If no primary property is overridden.

property number_of_ports: int

Number of ports.

Return type:

int

property num_flat_params: int

Number of free, flattened parameters.

Return type:

int

property num_params: int

Number of free parameters.

Return type:

int

property nports: int

Alias of number_of_ports.

property port_tuples: list[tuple[int, int]]

All (m, n) port index pairs.

Return type:

list[tuple[int, int]]

property z0: Array

Internal characteristic impedance.

Return type:

jnp.ndarray

__call__()

Build a compositional circuit representing this model.

Return type:

Model

Raises:

NotImplementedError – In the base class; override in derived classes to build a compositional representation.

primary(freq)

Dispatch to the primary function for the given frequency.

Parameters:

freq (Frequency)

Return type:

Array

a(freq)

ABCD parameter matrix.

If only s() is implemented, converts via pmrf.functions.conversions.s2a().

Parameters:

freq (Frequency) – Frequency grid.

Returns:

ABCD matrix with shape (nf, 2, 2).

Return type:

jnp.ndarray

Raises:

NotImplementedError – If neither a nor s is implemented in a subclass.

s(freq)

Scattering parameter matrix.

If only a() is implemented, converts via pmrf.functions.conversions.a2s().

Parameters:

freq (Frequency) – Frequency grid.

Returns:

S-parameter matrix with shape (nf, n, n).

Return type:

jnp.ndarray

Raises:

NotImplementedError – If neither a nor s is implemented in a subclass.

children()

Immediate submodels.

Return type:

list[Model]

named_children()

Immediate submodels keyed by their name attribute.

Return type:

dict[str, Model]

submodels()

All nested submodels (depth-first), excluding self.

Return type:

list[Model]

connected(others, ports)

Parameters:

• others (Model | Sequence[Model])

• ports (Sequence[int | Sequence[int]])

Return type:

Model

flipped()

Return a version of the model with ports flipped.

Return type:

Model

terminated(load=None)

Terminate a 2-port model in a 1-port load (default: SHORT).

Parameters:

load (Model, optional) – Load network. Defaults to a SHORT.

Return type:

Model

copy()

Parameters:

self (ModelT)

Return type:

ModelT

partition(include_fixed=False, param_objects=False)

Partition model into (parameters, static) trees.

This is useful for internal use, or for inspecting the model and its parameters.

Parameters:

• include_fixed (bool, default=False) – Include fixed parameters in the parameter tree.

• param_objects (bool, default=False) – If True, keep full Parameter objects; otherwise filter to .value.

• self (ModelT)

Returns:

(params_tree, static_tree)

Return type:

(ModelT, ModelT)

params(*args, **kwargs)

Return type:

list[Parameter]

named_params(include_fixed=False, flat=False, flat_params=False, values=False, scaled_values=False, submodels=None)

Return model parameters as a dict (or flattened array).

Keys are fully-qualified parameter names. The order matches the internal flattened array order.

Parameters:

• include_fixed (bool, default=False) – Include fixed parameters.

• flat (bool, default=False) – If True, return a 1D jnp.ndarray of parameter values.

• flat_params (bool, default=False) – If True, return a dict of manufactured parameter objects (flattened) instead of an array.

• submodels (Model | Sequence[Model] | str | Sequence[str] | None, optional) – Restrict to parameters used by the given submodel(s). If strings are provided, getattr(self, name) is used.

Return type:

dict[str, Parameter] or jnp.ndarray

flat_params(*args, **kwargs)

Shorthand for named_params(flat=True, ...).

Return type:

Array

flat_param_names()

Flat parameter names (matching flat_params() order).

Return type:

list[str]

submodel_params(submodels, **kwargs)

Parameters:

submodels (Model | Sequence[Model] | str | Sequence[str])

param_names(*args, **kwargs)

Parameter names for current (possibly filtered) selection.

param_groups(include_fixed=False)

Return all parameter groups relevant to this model.

This function is useful for fitters that need to take into account e.g. any constraints or correlated priors. Only groups that were create with the same flat flag are returned.

Currently, only groups for the current model are considered (i.e. groups in children models are ignored).

Parameters:

include_fixed (bool, default=False) – Include groups involving fixed parameters.

Return type:

list[ParameterGroup]

distribution()

Joint distribution over (flattened) parameters.

Return type:

Distribution

classmethod with_defaults(*args, **kwargs)

with_params(params=None, check_missing=False, check_unknown=False, fix_others=False, include_fixed=False, **param_kwargs)

Return a new model with core parameters updated.

Parameters:

• params (dict[str, Parameter] | dict[str, float] | jnp.ndarray | None, optional) – Parameter updates. If an array, all values must be provided (matching flat_params order). You may also pass keyword args.

• check_missing (bool, default=False) – Require that all model parameters are specified.

• check_unknown (bool, default=False) – Error if unknown parameter keys are provided.

• fix_others (bool, default=False) – Fix any parameters not explicitly passed.

• include_fixed (bool, default=False) – Include fixed parameters when interpreting params mapping.

• **param_kwargs (dict) – Additional parameter updates by name.

• self (ModelT)

Return type:

ModelT

Raises:

Exception – If shape/order mismatches, unknown/missing names (when checked), or if arrays are found outside of Parameters.

with_flat_params(*args, **kwargs)

Alias for with_params() when passing a flat array.

with_fields(*args, **kwargs)

Return a copy with dataclass-style field replacements.

Parameters:

self (ModelT)

Return type:

ModelT

with_param_groups(param_groups)

Return a a model with parameter groups appended.

Parmaeter groups can be used to specify relationships between parameters in the model, such as joint priors.

Parameters:

• param_groups (ParameterGroup or list[ParameterGroup]) – Group(s) to add.

• self (ModelT)

Return type:

ModelT

with_fixed_params(params, check_unknown=True)

Return a model with specified parameters fixed.

Parameters:

• params (list[str]) – Parameter names to fix.

• check_unknown (bool, default=True) – Error if any provided name does not exist.

• self (ModelT)

Return type:

ModelT

with_free_params(params, fix_others=False, check_unknown=True)

Free the specified parameters.

Parameters:

• params (Sequence[str] | Sequence[Parameter]) – Parameters to set free.

• fix_others (bool, default=True) – Fix parameters not specified.

• check_unknown (bool, default=True) – Error if any provided name does not exist.

• self (ModelT)

Return type:

ModelT

with_free_submodels(submodels, include_fixed=False, fix_others=False)

Free all parameters in the given submodels.

Submodels parameters are obtained using Model.submodel_params, and subsequently freed using Model.with_free_params.

Parameters:

• free_submodels (Model | Sequence[Model] | str | Sequence[str]) – Submodels whose parameters should be free.

• self (ModelT)

• submodels (Model | Sequence[Model] | str | Sequence[str])

Return type:

ModelT

with_fixed_submodels(submodels)

Fixed all parameters in the given submodels.

Submodels parameters are obtained using Model.submodel_params, and subsequently fixed using Model.with_fixed_params.

Parameters:

• free_submodels (Model | Sequence[Model] | str | Sequence[str]) – Submodels whose parameters should be free.

• self (ModelT)

• submodels (Model | Sequence[Model] | str | Sequence[str])

Return type:

ModelT

to_skrf(frequency, sigma=0.0, **kwargs)

Convert the model at frequencies to an skrf.Network.

The active primary property (self.primary_property) is used.

Parameters:

• frequency (pmrf.Frequency | skrf.Frequency) – Frequency grid.

• sigma (float, default=0.0) – If nonzero, add complex Gaussian noise with stdev sigma to s.

• **kwargs – Forwarded to skrf.Network constructor.

Return type:

skrf.Network

write_hdf(group)

Write model into an HDF5 group.

Parameters:

group (h5py.Group) – Target group. Two subgroups are created: raw and params.

classmethod read_hdf(group)

Read model from an HDF5 group.

Parameters:

group (h5py.Group) – Group previously created via write_hdf().

Returns:

The decoded model, or None if decoding fails.

Return type:

ModelT | None

save(target)

Serialize the model to disk in pickle format.

Parameters:

• filepath (str) – Destination file path.

• target (str | BinaryIO)

classmethod load(source)

Load a model from a pickled file.

Parameters:

• filepath (str) – Source JSON path.

• source (str | BinaryIO)

Return type:

ModelT

export_touchstone(filename, frequency, sigma=0.0, **skrf_kwargs)

Export the model response to a Touchstone file via scikit-rf.

Parameters:

• frequency (Frequency | skrf.Frequency)

• filename (str)

• sigma (float, default=0.0) – Additive complex noise std for S-parameters.

• **skrf_kwargs – Forwarded to skrf.Network.write_touchstone().

Returns:

Return value of Network.write_touchstone.

Return type:

Any

Parameters:

• name (str | None)

• separator (str)

• _z0 (complex)

• _param_groups (list[ParameterGroup])

Module-local helpers

These helpers live in pmrf.model but are re-exported at pmrf. Prefer the canonical docs at the root package page.

pmrf.models.wrap(func, *args, as_numpy=False)

Wrap a function/method taking (model, frequency, *args, **kwargs) into one that accepts flat arrays.

The wrapper accepts (theta, [f], *args, **kwargs) where theta is a flat parameter array, and optionally f is a frequency array with a provided unit.

Supported calls

• wrap(userfunc, model, "MHz")

• wrap(userfunc, model, freq)

• wrap(model.userfunc, "MHz")

• wrap(model.userfunc, freq)

If a Frequency object is passed, it is closed over and the wrapped function only accepts theta.

param func:

Function or bound method taking (model, frequency, ...).

type func:

Callable

param *args:

Either (model, Frequency|unit) or just (Frequency|unit) if func is bound.

type *args:

tuple

param as_numpy:

If True, JIT the wrapper and convert inputs/outputs to NumPy arrays.

type as_numpy:

bool, default=False

returns:

Wrapped function accepting (theta_array, [f_array], *args, **kwargs).

rtype:

Callable

raises TypeError:

If the frequency argument is neither a unit string nor Frequency.

Parameters:

• func (Callable)

• as_numpy (bool)

Return type:

Callable