Model (pmrf.Model)

class Model(*, name=None, z0=50 + 0j)

Bases: Module

Overview

This base class is used to represent any computable RF network, referred to in ParamRF as a “Model”. This class can be overridden for defining complex models, or can be utilized indirectly by combining models already provided in pmrf.models.

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 new models by sub-classing the model and adding custom parameters and/or sub-models

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

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

• Use with_xxx functions to modify fields, models and parameters within the model e.g. with_params(), with_fields().

• Use “past tense” functions to modify the model in conjunction with another model or data e.g. terminated(), flipped().

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

Methods & Properties Summary

Defaults / Primary

Method / Property	Description
DEFAULT_NAMED_PARAMS	Mapping from parameter name to Parameter.
DEFAULT_PARAM_NAMES	Default parameter names for the model.
DEFAULT_PARAMS	Default parameters for the model.
primary_function	The primary function (s or a) as a callable.
primary_property	The primary property (e.g. "s", "a") as a string.

Introspection Properties

Method / Property	Description
number_of_ports	Number of ports.
nports	Alias of number_of_ports.
port_tuples	All (m, n) port index pairs.
num_params	Number of free parameters.
num_flat_params	Number of free, flattened parameters.

Core API

Method	Description
__call__()	Build the model. Should be overridden by sub-classes.
primary()	Dispatch to the primary function for the given frequency.
s()	Scattering (S) parameter matrix.
a()	ABCD parameter matrix.
z()	Impedance (Z) parameter matrix.
y()	Admittance (Y) parameter matrix.
s_jacobian()	Jacobian of the S-parameters w.r.t free parameters.
a_jacobian()	Jacobian of the ABCD-parameters w.r.t free parameters.
z_jacobian()	Jacobian of the Z-parameters w.r.t free parameters.
y_jacobian()	Jacobian of the Y-parameters w.r.t free parameters.

Function Tools

Method	Description
func_jacobian()	Calculate the Jacobian of an arbitrary function w.r.t parameters.
func_sensitivity()	Calculate the sensitivity of an arbitrary function w.r.t parameters.
func_samples()	Evaluate an arbitrary function over parameter samples.

Model Inspection & Manipulation

Method	Description
children()	Returns the immediate submodels.
submodels()	Returns all nested submodels (depth-first).
partition()	Partition model into parameters and static trees.
flipped()	Return a version of the model with ports flipped.
renumbered()	Return a version of the model with ports renumbered.
terminated()	Return a new model terminated by another (e.g. load).
sampled()	Return a new model with parameters drawn from this model’s distribution.

Parameter Inspection

Method	Description
named_params()	Named model parameter objects as a dict.
named_param_values()	Named model parameter values as a dict of jax arrays.
param_names()	Model parameter names as a list.
param()	A single model parameter object by name.
params()	Model parameters as a list.
param_value()	A single model parameter value by name.
param_values()	Model parameter values as a list of jax arrays.
named_flat_params()	Named flattened model parameter objects as a dict.
named_flat_param_values()	Named flattened model parameter values as a dict.
flat_param_names()	Flattened parameter names as a list.
flat_params()	Flattened parameters as a list.
flat_param_values()	Flattened model parameter values as a flat array.
flat_param_bounds()	Flattened model parameter bounds as jax arrays.
param_groups()	Return all parameter groups relevant to this model.
distribution()	Joint distribution over (flattened) parameters.

Parameter Manipulation

Method	Description
with_params()	Return a model with parameters updated.
with_mapped_params()	Apply a map function to parameters.
with_fixed_params()	Return a model with specified parameters fixed.
with_free_params()	Return a model with specified parameters free.
with_free_params_only()	Return a model with ONLY the specified parameters free.
with_all_params_fixed()	Return a model with all parameters fixed.
with_all_params_free()	Return a model with all parameters free.

Parameter Group Manipulation

Method	Description
with_param_groups()	Return a model with parameter groups appended.
with_demoted_param_groups()	Recursively demote parameter groups to deepest submodel.
with_no_param_groups()	Return a model with all parameter groups removed.

Distribution Manipulation

Field & Model Manipulation

Method	Description
with_defaults()	Return this model type with default initialization args.
with_models()	Combines this model with free parameters in other models.
with_fields()	Return a copy with dataclass-style field replacements.
with_name()	Return a copy of this model with a different name.
with_submodel_fields()	Dataclass-style field replacements on a nested sub-model.
with_free_submodels()	Free all parameters in the given submodels.
with_free_submodels_only()	Returns a model with ONLY the specified submodels freed.
with_fixed_submodels()	Fix all parameters in the given submodels.
with_tied_submodels()	Tie submodels structurally to a shared model.
tied()	Return the model with self tied to a shared model.
with_injected_params()	Inject parameters from a shared model into target submodels.

Plotting, File, & Conversion Utilities

Method	Description
plot_func()	Evaluate and plot an arbitrary function of the model.
plot_func_samples()	Evaluate and plot a function over parameter samples.
to_skrf()	Convert the model at frequencies to an skrf.Network.
export_touchstone()	Export the model response to a Touchstone file.

ivar name:

An optional name for the model instance.

vartype name:

str or None

ivar separator:

The separator character used when flattening nested parameter names (default is ‘_’).

vartype separator:

str

ivar metadata:

A dictionary for storing arbitrary metadata associated with the model.

vartype metadata:

dict

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()

DEFAULT_NAMED_PARAMS = {}

DEFAULT_PARAMS = []

DEFAULT_PARAM_NAMES = []

a(freq)

ABCD parameter matrix.

If a different parameter type is primary, this converts it to A.

Parameters:

freq (Frequency) – Frequency grid.

Returns:

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

Return type:

jnp.ndarray

a_jacobian(freq)

Calculate the Jacobian of the ABCD-parameters with respect to free parameters.

See func_jacobian().

Parameters:

• freq (Frequency) – The frequency grid to evaluate the ABCD-parameters over.

• self (Self)

Returns:

A dictionary mapping flat parameter names to their gradient arrays. Each array has shape (n_freqs, n_ports, n_ports).

Return type:

dict[str, jnp.ndarray]

children()

Returns the immediate submodels.

Return type:

list[Model]

copy()

Returns a deepcopy of self.

Return type:

Model

Parameters:

self (Self)

distribution(param_groups=True)

Joint distribution over (flattened) parameters.

Parameters:

param_groups (bool, optional) – Whether or not to use the internal parameter groups to create the joint distribution. Defaults to True.

Return type:

JointParameterDistribution

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

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

Parameters:

• filename (str)

• frequency (Frequency | skrf.Frequency)

• 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

flat_param_bounds(**kwargs)

Return flattened model parameter bounds as jax arrays.

Note that a minimum and maximum percentile is used to get the bounds for any non-uniform distribution.

Equivalent to getting the bounds from distribution(), which key-word arguments are forwarded to.

Return type:

tuple[Array, Array]

flat_param_names(*args, **kwargs)

Return flattened parameter names as a list.

See named_flat_params().

Return type:

list[str]

flat_param_values(*args, **kwargs)

Return flattened model parameter values as a jax arrays.

See named_flat_param_values().

Return type:

Array

flat_params(*args, **kwargs)

Return flattened parameters as a list.

See named_flat_params().

Return type:

list[Parameter]

flipped(**kwargs)

Return a version of the model with ports flipped.

See pmrf.models.composite.transformed.Flipped.

Return type:

Model

func_jacobian(func, freq)

Calculate the Jacobian of an arbitrary function with respect to free parameters.

This uses forward-mode automatic differentiation to compute the gradients of the provided function with respect to each free parameter in the model.

Parameters:

• func (Callable[[Model, Frequency], jnp.ndarray]) – Function to differentiate. Must take a Model and a Frequency object and return a jnp.ndarray of any shape.

• freq (Frequency) – The frequency grid to evaluate the function over.

• self (Self)

Returns:

A dictionary mapping flat parameter names to their gradient arrays. Each array has the same shape as the output of func.

Return type:

dict[str, jnp.ndarray]

func_samples(func, freq, *, key, num_samples=1000)

Evaluates an arbitrary function over samples drawn from the model’s distribution.

Parameters:

• func (Callable[[Model], jnp.ndarray]) – A function that takes a Model instance and returns a JAX array.

• prng_key (Array) – JAX random key for sampling.

• num_samples (int, default=1000) – Number of models to sample from the joint distribution.

• freq (Frequency)

• key (Array)

Returns:

The function evaluated over all samples. Shape will be (num_samples, *func_output_shape).

Return type:

jnp.ndarray

func_sensitivity(func, freq)

Calculate the relative (normalized) sensitivity of an arbitrary function.

This computes the fractional change in the function’s output given a fractional change in each free parameter. Mathematically, it evaluates:

\[S_{rel} = \frac{\partial y}{\partial \theta} \frac{\theta}{y}\]

Parameters:

• func (Callable[[Model, Frequency], jnp.ndarray]) – Function to evaluate. Must take a Model and a Frequency object and return a jnp.ndarray of any shape.

• freq (Frequency) – The frequency grid to evaluate the function over.

• self (Self)

Returns:

A dictionary mapping flat parameter names to their normalized sensitivity arrays. Each array has the same shape as the output of func.

Return type:

dict[str, jnp.ndarray]

name: str | None = None

named_flat_param_values(scaled=False, return_floats=False, **kwargs)

Named flattened model parameter values as a dict of jax arrays.

See named_flat_params().

Parameters:

• scaled (bool, default=False) – Whether or not to scale the returned values by the parameter scales.

• **kwargs – Additional key-word arguments as in named_params().

Return type:

dict[str, jnp.ndarray]

named_flat_params(include_fixed=False, submodels=None)

Named flattened model parameters as a dict.

Flat parameters are a de-vectorized version of the internal parameters of the model. The returned parameter objects therefore are not necessarily equal to the internal model objects.

Keys are fully-qualified parameter names with de-vectorized suffixes added. The order matches the internal flattened array order.

Parameters:

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

• 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]

named_param_values(scaled=False, **kwargs)

Named model parameter values as a dict of jax arrays.

See named_params().

Parameters:

• scaled (bool, default=False) – Whether or not to scale the returned values by the parameter scales.

• **kwargs – Additional key-word arguments as in named_params().

Return type:

dict[str, jnp.ndarray]

named_params(param_filter=None, *, include_fixed=False, submodels=None)

Named model parameters as a dict.

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

Parameters:

• param_filter (str | Sequence[str] | Parameter | Sequence[Parameter] | Callable[[str], bool], default=None) – A filter indicating which parameters to return. For the default case, all parameters are returned.

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

• 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]

property nports: int

Alias of number_of_ports.

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 number_of_ports: int

Number of ports.

Return type:

int

param(name, *args, **kwargs)

Return a single model parameter by name.

See named_params().

Parameters:

name (str)

Return type:

Parameter

param_groups(include_fixed=False, explicit_only=False)

Return all parameter groups relevant to this model, including submodels.

This function recursively traverses submodels to collect their parameter groups, adjusting parameter names to match the current model’s scope.

Priority is given to groups defined in the parent model. If a parameter is grouped explicitly in self._param_groups, it will be removed from any groups returned by submodels.

Parameters:

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

Return type:

list[ParameterGroup]

param_names(*args, **kwargs)

Return model parameter names as a list.

See named_params().

Return type:

list[str]

param_value(name, *args, **kwargs)

Return a single model parameter value by name as a single jax array.

See named_param_values().

Parameters:

name (str)

Return type:

Array

param_values(*args, **kwargs)

Return model parameter values as a list of jax arrays.

See named_param_values().

Return type:

list[Array]

params(*args, **kwargs)

Return model parameters as a list.

See named_params().

Return type:

list[Parameter]

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 (Self)

Returns:

(params_tree, static_tree)

Return type:

(Self, Self)

plot_func(func, freq, *, ax=None, label=None, color=None, **kwargs)

Evaluate and plot an arbitrary function of the current model.

This method evaluates the provided function using the model’s current parameter values and plots the resulting response over frequency.

Parameters:

• func (Callable[[Model, Frequency], jnp.ndarray]) – Function to evaluate. Must take a Model and a Frequency object and return a jnp.ndarray of shape (n_freqs,).

• freq (Frequency) – Frequency grid to evaluate over.

• ax (matplotlib.axes.Axes, optional) – Axes to plot on. If None, the current axes (plt.gca()) are used.

• label (str, optional) – Label for the plotted line (used in legends).

• color (str, optional) – Color for the line. If None, uses the matplotlib color cycle.

• **kwargs (dict) – Additional keyword arguments forwarded to matplotlib.pyplot.plot (e.g., linestyle, linewidth, alpha).

Returns:

The axes containing the plot.

Return type:

matplotlib.axes.Axes

plot_func_samples(func, freq, *, key, num_samples=1000, contours=True, ax=None, label=None, color='C0', alpha=0.1)

Evaluate and plot a function over samples from the parameter distribution.

This method draws samples from the model’s joint parameter distribution, evaluates the provided function for each sample, and plots the resulting responses over frequency.

Parameters:

• func (Callable[[Model, Frequency], jnp.ndarray]) – Function to evaluate. Must take a Model and a Frequency object and return a jnp.ndarray of shape (n_freqs,).

• freq (Frequency) – Frequency grid to evaluate over.

• key (Array) – PRNG key for sampling the distribution.

• num_samples (int, default=1000) – Number of samples to draw.

• contours (bool, default=True) – If True, plots the mean response and filled contours corresponding to 1, 2, and 3 standard deviations. If False, plots all individual sample responses as transparent lines.

• ax (matplotlib.axes.Axes, optional) – Axes to plot on. If None, the current axes (plt.gca()) are used.

• label (str, optional) – Label for the mean line (used in legends).

• color (str, default='C0') – Base color for the lines and shaded regions.

• alpha (float, default=0.1) – Transparency of the individual lines (when contours=False).

Returns:

The axes containing the plot.

Return type:

matplotlib.axes.Axes

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

All (m, n) port index pairs.

Return type:

list[tuple[int, int]]

primary(freq)

Dispatch to the primary function for the given frequency.

Parameters:

freq (Frequency)

Return type:

Array

property primary_function: Callable[[Frequency], Array]

The primary function (s or a) as a callable.

The primary function is the first overridden among PRIMARY_PROPERTIES, unless __call__ is overridden, in which case the primary function of the built model is returned.

Return type:

Callable[[Frequency], jnp.ndarray]

Raises:

NotImplementedError – If no primary property is overridden.

property primary_property: str

The primary property (e.g. "s", "a") as a string.

The primary property is the first overridden among PRIMARY_PROPERTIES, unless __call__ is overridden, in which case the primary property of the built model is returned.

Return type:

str

Raises:

NotImplementedError – If no primary property is overridden.

renumbered(from_ports, to_ports=None, **kwargs)

Return a version of the model with ports renumbered.

See pmrf.models.composite.transformed.Renumbered.

from_portstuple[int]

The original port indices that map to to_ports.

to_portstuple[int]

The new port indices.

Return type:

Model

Parameters:

• from_ports (tuple[int])

• to_ports (tuple[int])

s(freq)

Scattering parameter matrix.

If a different parameter type (a, z, y) is primary, this converts it to S.

Note that, in ParamRF, the power wave definition of S-parameters should be used. If you have a formulation in terms of another definition (such as traveling waves), simply use pmrf.rf_functions.s2s() (or pmrf.rf_functions.renormalize_s() if you need to change impedance too).

Parameters:

freq (Frequency) – Frequency grid.

Returns:

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

Return type:

jnp.ndarray

s_jacobian(freq)

Calculate the Jacobian of the S-parameters with respect to free parameters.

See func_jacobian().

Parameters:

• freq (Frequency) – The frequency grid to evaluate the S-parameters over.

• self (Self)

Returns:

A dictionary mapping flat parameter names to their gradient arrays. Each array has shape (n_freqs, n_ports, n_ports).

Return type:

dict[str, jnp.ndarray]

sampled(key=None, **kwargs)

Returns a new model with parameters sampled from this parameter’s distribution.

Return type:

Model

submodels()

Returns all nested submodels (depth-first), excluding self.

Return type:

list[Model]

terminated(load=None, **kwargs)

Returns a new model that contains this model terminated in another.

See pmrf.models.composite.transformed.Terminated.

Parameters:

load (Model | str, optional) – Load network. Can be ‘short’ or ‘open’ as aliases for SHORT and OPEN. Defaults to a SHORT.

Return type:

Model

tied(shared_model)

Return the model with self structurally tied to a shared model.

Same as with_tied_submodels but operates directly on self.

Parameters:

• shared_model (Model) – The external model to tie the submodels to.

• self (Self)

Returns:

A new model instance with self tied to the specified model.

Return type:

Self

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 (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

with_all_params_fixed(**kwargs)

Returns a model with all parameters fixed.

This is an alias for calling with_free_params() with fix_others=True and no parameters passed.

See `with_free_params`().

Parameters:

self (Self)

Return type:

Self

with_all_params_free(**kwargs)

Returns a model with all parameters free.

This is an alias for calling with_free_params() with all parameters passed.

See `with_free_params`().

Parameters:

self (Self)

Return type:

Self

classmethod with_defaults(*args, **kwargs)

Return this model type with default initialization arguments.

This method is very useful in utilizing an existing model with default values, without having to create a new model type via inheritance.

Arguments are forwarded as if they were passed to __init__.

Return type:

type[Model]

with_demoted_param_groups()

Recursively demote parameter groups to the deepest possible submodel.

This method identifies parameter groups where every parameter belongs to the same immediate submodel. It moves those groups to the submodel, stripping the prefix. It then recursively calls this method on the submodels to ensure groups continue moving down the hierarchy as far as possible.

Returns:

A new model instance with parameter groups distributed to their lowest relevant submodels.

Return type:

Self

Parameters:

self (Self)

with_fields(*args, **kwargs)

Return a copy of this model with dataclass-style field replacements.

Parameters are forwarded to dataclasses.replace().

Parameters:

self (Self)

Return type:

Self

with_fixed_params(param_filter, free_others=False, **kwargs)

Return a model with specified parameters fixed.

This maps each parameter in the filter, calling Parameter.as_fixed() on each.

See with_mapped_params().

Parameters:

• free_others (bool, default=False) – Also free all parameters not in the filter.

• self (Self)

• param_filter (str | Sequence[str] | Parameter | Sequence[Parameter] | Callable[[str], bool])

Return type:

Self

with_fixed_submodels(submodels)

Fix all parameters in the given submodels.

Submodels parameters are obtained using param_names()., and subsequently fixed using `with_fixed_params`().

Parameters:

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

• self (Self)

Return type:

Self

with_free_params(param_filter, *, fix_others=False, **kwargs)

Free the specified parameters.

This maps each parameter in the filter, calling Parameter.as_free() on each.

See with_mapped_params().

Parameters:

• fix_others (bool, default=False) – Also fix all parameters not in the filter.

• self (Self)

• param_filter (str | Sequence[str] | Parameter | Sequence[Parameter] | Callable[[str], bool])

Return type:

Self

with_free_params_only(param_filter, **kwargs)

Returns a model with only the specified parameters freed.

This is an alias for calling with_free_params() with fix_others=True.

See `with_free_params`().

Parameters:

• self (Self)

• param_filter (str | list[str] | Callable[[str], bool])

Return type:

Self

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

Free all parameters in the given submodels.

Submodels parameters are obtained using param_names()., and subsequently freed using `with_free_params`().

Parameters:

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

• include_fixed (bool, default=False) – Also free parameters that are currently fixed in the submodels.

• fix_others (bool, default=False) – Fix all other submodels.

• self (Self)

Return type:

Self

with_free_submodels_only(*args, **kwargs)

Returns a model with only the specified submodels freed.

This is an alias for calling with_free_submodels() with fix_others=True.

See `with_free_params`().

Parameters:

self (Self)

Return type:

Self

with_injected_params(submodel_attrs, shared_model)

Return a copy of the model with free parameters from a shared model injected into target submodels.

This method dynamically overrides the parameter values of internal submodels using the values from an external shared model. It is designed to be called during the forward pass (e.g., inside __call__) to enforce hard equality constraints on parameters that were structurally tied using with_tied_submodels().

Parameters:

• submodel_attrs (str | Sequence[str]) – The attribute name(s) of the internal submodel(s) to inject parameters into.

• shared_model (Model) – The external model providing the free parameter values.

• self (Self)

Returns:

A new model instance with the updated parameter values injected.

Return type:

Self

with_mapped_distributions(mapper, dist_filter=None, *, map_others=None, param_groups=False)

Return a model with a function applied to its parameter distributions.

This method allows for bulk-updates of distributions, such as widening variances or changing distribution types.

If param_groups is False, the mapping is applied to the distributions of individual parameters (flattened).

If param_groups is True, the mapping is applied to the distributions of ParameterGroup objects. This mode is recursive: it will traverse the model tree and apply the mapping to all explicit parameter groups in all submodels.

Parameters:

• mapper (Callable[[Distribution], Distribution]) – Function that takes a distribution and returns a new one.

• dist_filter (Callable[[Distribution], bool] | None, default=None) – A predicate function. If provided, the mapping is only applied to distributions where dist_filter(dist) is True. If None, applies to all.

• map_others (Callable[[Distribution], Distribution] | None, default=None) – An optional map to apply to all distributions NOT in the filter.

• param_groups (bool, default=False) – If True, map distributions on parameter groups (recursively). If False, map distributions on individual parameters (flat).

• self (Self)

Returns:

A new model with updated distributions.

Return type:

Self

with_mapped_params(mapper, param_filter=None, *, map_others=None, prefixes=False, include_fixed=False)

Return a model with specified parameters mapped.

Parameters:

• mapper (Callable[[Parameter], Parameter]) – The map to apply to each parameter in the filter (or all if no filter).

• param_filter (str | Sequence[str] | Callable[[str], bool] | None, default=None) – Parameter names to map. If None, applies mapper to all parameters.

• map_others (Callable[[Parameter], Parameter] | None, default=None) – An optional map to apply to all parameters NOT in the filter.

• prefixes (bool, default=False) – Specifies that, when a string or list of strings is passed in param_filter, these must be interpreted as parameter prefixes to map and not full path names. Defaults to False.

• self (Self)

• include_fixed (bool)

Return type:

Self

with_models(models)

Combines this model with free parameters in other models.

This is useful to combine separate models obtained from fitting the same initial model with different free parameters.

Parameters:

• models (Model or Sequence[Model]) – The other models to combine this model with.

• self (Self)

Return type:

Model

with_name(name)

Return a copy of this model with a different name.

Parameters:

• self (Self)

• name (str | None)

Return type:

Self

with_no_param_groups()

Return a new model with all parameter groups removed recursively.

This clears the _param_groups of the current model and traverses all nested submodels (and sequences of submodels) to remove their parameter groups as well.

Returns:

A new model instance with no parameter groups.

Return type:

Self

Parameters:

self (Self)

with_param_groups(param_groups)

Return a model with parameter groups appended, replacing existing relationships.

This method implements an “atomic replacement” policy. If any parameter in an existing group is claimed by a new group, the entire existing group is removed.

This ensures that groups defining joint distributions are not left in an invalid broken state (e.g. having a dimension removed). Parameters that were in the removed group but not in the new group will revert to being ungrouped (handled by param_groups as singleton groups).

Parameters:

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

• self (Self)

Return type:

Self

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

Return a new model with parameters updated.

This is a multi-purpose function that updates parameters differently based on the types pass.

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=True) – 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 (Self)

Return type:

Self

Raises:

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

with_submodel_fields(submodel, *args, **kwargs)

Return a copy of this model with dataclass-style field replacements on a nested sub-model.

Parameters are forwarded to dataclasses.replace().

Parameters:

• submodel (str | Sequence[str]) – The name of the submodel (or sequence of names) to traverse. Can be a single string with a path e.g. ‘submodel1.submodel2’, or a list of submodels e.g. [‘submodel1’, ‘submodel2’].

• self (Self)

Return type:

Self

with_tied_submodels(submodel_attrs, shared_model)

Return a copy of the model with specified submodels structurally tied to a shared model.

This method prepares submodels to act as structural proxies during optimization. It prevents the optimizer from seeing duplicate free parameters by fixing the target submodel’s parameters. If the target submodel’s type matches the shared model, only the overlapping free parameters are fixed. If the types differ, the target submodel is entirely replaced by a fully fixed copy of the shared model.

This is typically used in __post_init__ to set up the model structure, and paired with with_injected_params() during the forward pass.

Parameters:

• submodel_attrs (str | Sequence[str]) – The attribute name(s) of the internal submodel(s) to tie.

• shared_model (Model) – The external model to tie the submodels to.

• self (Self)

Returns:

A new model instance with the specified submodels tied.

Return type:

Self

with_uniform_distributions(percentage, param_filter=None, *, respect_bounds=False, remove_param_groups=True, zero_values='keep', **kwargs)

Return a model with uniform distributions set centered on current parameter values.

The distributions are defined with bounds calculated as value * (1.0 +/- percentage).

Parameters:

• percentage (float) – The fractional width of the uniform distribution (e.g. 0.1 = 10%).

• filter (str | Sequence[str] | Callable[[str], bool], default=None) – The parameters to be updated with new uniform distributions. For the default case, all are updated.

• respect_bounds (bool, default=False) – Whether or not the min and max bounds of the current distributions should be respected. If True, new bounds will not go larger than past these bounds.

• remove_param_groups (bool, default=True) – Whether to remove parameter groups recursively when setting the uniform distributions. Otherwise, the joint distribution of the model may not be the desired uniform distribution.

• zero_values (str, default='keep') – How to treat zero values. Currently the only option is to keep them and their bounds as is.

• param_filter (str | Sequence[str] | Parameter | Sequence[Parameter] | Callable[[str], bool])

Returns:

A new model with updated parameter distributions.

Return type:

Self

y(freq)

Admittance (Y) parameter matrix.

If a different parameter type is primary, this converts it to Y.

Parameters:

freq (Frequency) – Frequency grid.

Returns:

Y matrix with shape (nf, n, n).

Return type:

jnp.ndarray

y_jacobian(freq)

Calculate the Jacobian of the Y-parameters with respect to free parameters.

See func_jacobian().

Parameters:

• freq (Frequency) – The frequency grid to evaluate the Y-parameters over.

• self (Self)

Returns:

A dictionary mapping flat parameter names to their gradient arrays. Each array has shape (n_freqs, n_ports, n_ports).

Return type:

dict[str, jnp.ndarray]

z(freq)

Impedance (Z) parameter matrix.

If a different parameter type is primary, this converts it to Z.

Parameters:

freq (Frequency) – Frequency grid.

Returns:

Z matrix with shape (nf, n, n).

Return type:

jnp.ndarray

z0: complex = (50+0j)

z_jacobian(freq)

Calculate the Jacobian of the Z-parameters with respect to free parameters.

See func_jacobian().

Parameters:

• freq (Frequency) – The frequency grid to evaluate the Z-parameters over.

• self (Self)

Returns:

A dictionary mapping flat parameter names to their gradient arrays. Each array has shape (n_freqs, n_ports, n_ports).

Return type:

dict[str, jnp.ndarray]

Parameters:

• name (str | None)

• z0 (complex)