DiscreteCallable

class pmrf.models.adapters.callable.DiscreteCallable(frequency: Any, fn: Any, theta: Any, *, z0: complex = 50 + 0j, name: str | None = None, metadata: Any = None, kind: str = 's')

Bases: AbstractSingleDiscreteProperty

A model that predicts its output at a discrete set of frequencies already known to the model using an arbitrary callable.

This class can be used to wrap external machine learning architectures (Equinox/Parax/other).

Parameters:
  • fn (Callable[[], jnp.ndarray] | Callable[[jnp.ndarray], jnp.ndarray]) – The underlying callable model which predicts the response. May either be a function or a callable PyTree (e.g. equinox.Module) which will be frozen. Must either accept no parameters or an array of shape (nparams,) depending on if theta is None, and return an array of shape (nfreq, nports, nports).

  • theta (Param) – Parameters to pass to fn of length nparams. Can be None for models that contain their own parax.Parameter objects. All parameters, including fixed parameters, are passed.

discrete_matrix() Array

The primary discrete data.

fn: Callable[[], Array] | Callable[[Array], Array]

The underlying callable model

theta: AbstractVariable | Inexact[jaxlib._jax.Array, '...']

Parameters to pass to fn