Model (pmrf.Model)

class pmrf.Model(*, z0: complex = 50 + 0j, name: str | None = None, metadata: Any = None)

Bases: Module

Base class for RF models.

Derived from this class to define your own, custom model.

This class should not be instantiated directly. It is created internally in ParamRF when models are built compositionally, or can be inherited from, in which case at least one of pmrf.Model.s(), pmrf.Model.a(), pmrf.Model.y(), pmrf.Model.z(), pmrf.Model.build(), or pmrf.Model.primary_matrix() should be overridden.

The model is a Equinox Module (an immutable dataclass) and a JAX PyTree. Parameters are declared using standard dataclass field syntax and should be annotated with type pmrf.Param and field specifier pmrf.param(). For more details in parameter definitions, see pmrf.parameters.

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

  • Use pmrf.Model.at and methods such as terminated() and flipped() to create modified versions of your model.

Methods & Properties Summary

Core Methods

Method

Description

s()

Scattering (S) parameter matrix.

a()

ABCD parameter matrix.

z()

Impedance (Z) parameter matrix.

y()

Admittance (Y) parameter matrix.

build()

Build the model. Can be overridden for advanced models.

primary_matrix()

Return the primary matrix. Can be overridden for dynamic dispatch.

primary_property

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

Helper Methods

Method

Description

number_of_ports

Number of ports.

nports

Alias of number_of_ports.

port_tuples

All (m, n) port index pairs.

named_params()

Extracts all named parameters in the model.

Model Transformation

Method

Description

at()

Modify, filter or inspect a value at some path in the model.

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

tied()

Tie certain parameters/sub-models together.

File & Conversion Utilities

Method

Description

to_skrf()

Convert the model at frequencies to an skrf.Network.

export_touchstone()

Export the model response to a Touchstone file.

Examples

A general PiCLC network model:

import jax.numpy as jnp
import pmrf as prf

class PiCLC(prf.Model):
    C1: prf.Param = prf.param()
    L:  prf.Param = prf.param()
    C2: prf.Param = prf.param()

    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 built in build using cascading:

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

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

    def build(self) -> prf.Model:
        return self.res ** self.ind ** self.cap.terminated()
at(where: Callable[[Self], S]) → Lens[Self, S]

(experimental) A functional interface for model manipulation.

This is a wrapper around equinox.tree_at via the jax-optix library.

Pass in a callable that accepts your model and returns the attributes you would like to retrieve/modify. Then, use methods like .get() and .set() to retrieve values or an updated model.

Warning: Note that updates are “surgical”, so the value is replaced as-is i.e. without any automatic converters. When replacing parameters, ensure to pass in a fully constructed parameter and not a float.

Examples

>>> import pmrf as prf
>>> from pmrf.models import Resistor
>>> model = Resistor(R=50.0)
>>> # Retrieve a value using the lens
>>> model.at(lambda m: m.R).get()
50.0
>>> # Return a new model instance with the updated value
>>> updated_model = model.at(lambda m: m.R).set(100.0)
Returns:

A lens object focused on the root of the current instance.

Return type:

Lens

cascaded(other, **kwargs) → Model

Cascade this model with another, returning a new model.

See pmrf.models.composite.interconnected.Cascade.

Return type:

Model

export_touchstone(filename: str, frequency: Frequency | Any, sigma: float = 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

flipped(**kwargs) → Model

Return a version of the model with ports flipped.

See pmrf.models.composite.transformed.Flipped.

Return type:

Model

map(fn: Callable[[Any], Any], is_leaf: Callable | None = None) → Self

(experimental) A functional interface for model mapping.

This is a wrapper around jax.tree.map.

To map parameters, pass is_leaf=prf.is_param.

Examples

>>> import pmrf as prf
>>> from pmrf.models import Resistor, Capacitor
>>> model = Resistor(R=50.0) ** Capacitor(C=1e-12)
>>> # Scale all parameters in the model by a factor of 2
>>> scaled_model = model.map(lambda p: p * 2.0, is_leaf=prf.is_param)

Returns the mapped model.

named_params(include_fixed: bool = False, unwrap: bool = True, namespace_separator: str = '_') → dict[str, AbstractVariable | Inexact[jaxlib._jax.Array, '...']]

Returns a named dictionary of parameters in the model.

Parameters and models can be given names upon construction. The naming convention is as follows:

  1. If no names are present in the path of a parameter, its Python path is used.

  2. If there are any named models in the path of a parameter, the path up until the left of that model is collapsed, forming a namespace prefix. If multiple models in the path have names, they are joined using the supplied namespace separator.

  3. If the parameter itself is named, its path to the left is collapsed, either to the root or to the first named model.

This enables you to choose your own naming convention:

  1. For flat parameter names across your entire root model, name all your parameters but none of your models.

  2. For flat model names across your entire root model, name all your leaf models but none of your parameters or composite models.

  3. For fully nested naming, name all of your models and optionally your parameters.

Parameters:

  • include_fixed (bool) – Whether to include fixed parameters in the returned dictionary. Defaults to False.

  • unwrap (bool) – Unwraps the parameters into raw floats/arrays. Defaults to True. To inspect internal parameter states (e.g. distributions, fixed etc.) pass unwrap=False.

  • namespace_separator (str) – The separator to use to create a parameter namespace using model names.

Returns:

A dictionary mapping string paths (e.g., ‘.ind.value’) to their corresponding JAX arrays or parameter objects.

Return type:

dict[str, Any]

pathed_params(include_fixed: bool = False, unwrap: bool = True, keystr: bool = False, separator: str | None = None) → list[tuple[Any, AbstractVariable | Inexact[jaxlib._jax.Array, '...']]]

Returns the parameters as a list of tuples alongside their paths.

The paths represent JAX tree paths.

Parameters:

  • include_fixed (bool, default=False) – Whether to include fixed parameters in the returned dictionary. Defaults to False.

  • unwrap (bool, default=True) – Unwraps the parameters into raw floats/arrays. Defaults to True. To inspect internal parameter states (e.g. distributions, fixed etc.) pass unwrap=False.

  • keystr (bool, default=False) – Whether equivalent strings should be returned as opposed to full JAX paths. Defaults to False.

  • separator (str, optional) – The separator to use if keystr is True.

primary_matrix(freq: Frequency) → Array

The primary matrix (e.g. s, a etc.) as a function of frequency.

The primary matrix represents the matrix returned by pmrf.Model.primary_property, which is either overridden by sub-classes, or is the first proprerty directly overriden out of pmrf.Model.s(), pmrf.Model.a(), pmrf.Model.y(), pmrf.Model.z() (in that order), unless :meth:pmrf.Model.build is overridden, in which case the primary matrix of the built model is returned.

This method can also be overriden itself in order to to dynamically implement one of the matrices as opposed to overriding it explicitly.

Return type:

jnp.ndarray

Raises:

NotImplementedError – If no primary property is overridden.

renumbered(from_ports: tuple[int], to_ports: tuple[int] = None, **kwargs) → Model

Return a version of the model with ports renumbered.

See pmrf.models.composite.transformed.Renumbered.

Parameters:

  • from_ports (tuple[int]) – The original port indices that map to to_ports.

  • to_ports (tuple[int]) – The new port indices.

Return type:

Model

terminated(other: Model = 'short', **kwargs) → Model

Terminate this model in another, returning a new model.

See pmrf.models.composite.transformed.Terminated.

Parameters:

other (Model | str, optional) – The model to terminate this one in. Can be literals ‘short’, ‘open’ or any model with half the ports of this one. Defaults to a ‘short’.

Return type:

Model

tied(target: ~typing.Callable[[~typing.Any], ~typing.Any], source: ~typing.Callable[[~typing.Any], ~typing.Any], tie_fn: ~typing.Callable[[~typing.Any], ~typing.Any] = <function Model.<lambda>>, **kwargs) → Model

Tie parameters or sub-models within this model together.

See pmrf.models.composite.wrapped.Tied for more details.

Examples

>>> import pmrf as prf
>>> from pmrf.models import Resistor, Capacitor
>>>
>>> rc = Resistor(R=50.0) ** Capacitor(C=1.0e-12)
>>>
>>> # Tie the resistor's R to always be 50e12 times the capacitor's C
>>> tied_rc = rc.tied(
...     target=lambda m: m.models[0].R,
...     source=lambda m: m.models[1].C,
...     tie_fn=lambda c: c * 50e12
... )
>>>
>>> # The optimizer will now only see the Capacitor's C parameter.
>>> # When evaluated, R will automatically track C.
Parameters:

  • target (callable) – A callable extracting the parameter to be overwritten (e.g., lambda m: m.resistor.R).

  • source (callable) – A callable extracting the parameter to draw the value from (e.g., lambda m: m.capacitor.C).

  • tie_fn (callable, optional) – An optional transformation function applied to the source before injecting it into the target. Defaults to the identity function (lambda x: x).

Return type:

Model

to_skrf(frequency: Frequency | Any, sigma=0.0, **kwargs) → Network

Convert the model at frequencies to an skrf.Network.

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

Parameters:

  • frequency (pmrf.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

metadata: Any = None

Arbitrary metadata to store alongside the model.

name: str | None = None

A name for the model.

property nports: int

Alias of number_of_ports.

property number_of_ports: int

Number of ports.

Return type:

int

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

All (m, n) port index pairs.

Return type:

list[tuple[int, int]]

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 build 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.

z0: complex = (50+0j)

The characteristic impedance of the model. NB: Mixing impedances across models is not fully supported.