Module

class pmrf.modules.Module(*, name: str | None = None, metadata: Any = None)

Bases: Module

Base class for parameter-aware objects in ParamRF.

A module is an immutable JAX PyTree that may contain ParamRF parameters, RF models, and other modules. Unlike pmrf.Model, it does not imply an RF response or a number of ports.

at(target: Callable[[Self], T] | str | tuple[str, ...] | list[str]) → Lens[Self, T]

A functional interface for module manipulation.

This wraps equinox.tree_at() using an optic. Pass a callable, a string parameter name, or a tuple of names selecting the values to inspect or replace, then call methods such as .get(), .set(), or .apply().

Updates are surgical: replacement values bypass dataclass converters and validation. Callers must therefore preserve field invariants and pass fully constructed parameters where required.

Examples

>>> import pmrf as prf
>>> from pmrf.models import Resistor
>>> module = Resistor(R=50.0, name="res")
>>> module.at("res.R").get()
50.0
>>> updated = module.at("res.R").set(prf.Unconstrained(100.0))
Returns:

An optic focused on the root of this module.

Return type:

Lens

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

A functional interface for mapping over a module.

This wraps jax.tree.map(). To map parameters, pass is_target=pmrf.is_param.

Examples

>>> import pmrf as prf
>>> from pmrf.models import Resistor, Capacitor
>>> module = Resistor(R=50.0) ** Capacitor(C=1e-12)
>>> scaled = module.map(lambda p: p * 2.0, is_target=prf.is_param)
Returns:

The mapped module.

Return type:

Module

named_params(full_params: bool = False, free_only: bool = False, namespace_separator: str = '_') → dict[str, float | Array | TypeAliasForwardRef('Param')]

Return a named dictionary of parameters in the module.

Parameters and modules can be given names upon construction. If no custom names are present, standard Python attribute paths are used. Named modules collapse the path to their left into a namespace, whilst a named parameter collapses its path to the nearest named module or the root. This supports flat parameter names, flat module names, or fully nested namespaces.

Parameters:

  • full_params (bool, default=False) – Return full parameter objects instead of their resulting values.

  • free_only (bool, default=False) – Return only free parameters.

  • namespace_separator (str, default="_") – Separator used to join named module namespaces.

Returns:

Parameter names mapped to values or parameter objects.

Return type:

dict[str, Any]

tied(target: ~typing.Callable[[~typing.Any], ~typing.Any] | str | tuple[str, ...] | list[str], source: ~typing.Callable[[~typing.Any], ~typing.Any] | str | tuple[str, ...] | list[str], tie_fn: ~typing.Callable[[~typing.Any], ~typing.Any] = <function Module.<lambda>>) → Module

Tie parameters or sub-modules within this module together.

The target is hidden from optimizers and reconstructed from the source when the module is unwrapped. Targets and sources may be structural callables or resolved parameter names.

Parameters:

  • target – Callable or parameter name selecting the value to replace.

  • source – Callable or parameter name selecting the value it is derived from.

  • tie_fn – Transformation applied to the source value. Defaults to identity.

Returns:

A wrapped module with the relationship applied during unwrapping.

Return type:

Module

metadata: Any = None

Arbitrary metadata stored alongside the module.

name: str | None = None

A name for the module.