NumPyroNSFitter

class NumPyroNSFitter(model, *, likelihood_kind=None, likelihood_params=None, feature_sigmas=None, **kwargs)[source]

Bases: NumPyroFitter

Bayesian fitter using NumPyro Nested Sampling.

This backend utilizes numpyro.contrib.nested_sampling.NestedSampler to compute both the posterior samples and the log-evidence of the model.

Parameters:
execute(targets, *, constructor_kwargs=None, terminated_kwargs=None, seed=42, fitted_params='mean', **kwargs)[source]

Execute the Nested Sampling run.

NB: This method should not be called directly. Call run() instead.

Parameters:

  • targets (jax.numpy.ndarray) – The extracted target features to fit against.

  • constructor_kwargs (dict, optional) – Keyword arguments passed directly to the NestedSampler constructor (e.g., num_live_points).

  • terminated_kwargs (dict, optional) – Keyword arguments passed to control the termination criteria of the sampling loop.

  • seed (int, default=42) – The PRNG key seed for JAX’s random number generator.

  • fitted_params ({'mean'}, default='mean') – How to select the final point estimates for the returned model’s parameters from the posterior samples.

  • **kwargs – Additional arguments. num_samples is extracted to determine how many samples to draw from the final nested sampling weights.

Returns:

The fitted model and a dictionary containing the generated samples.

Return type:

tuple[Model, dict[str, jax.numpy.ndarray]]

cdf(theta)

Evaluate the combined cumulative distribution function (CDF).

Parameters:

theta (jax.numpy.ndarray) – The parameter values. Note that this 1D array must contain the model parameters followed sequentially by the likelihood noise parameters.

Returns:

The combined CDF probabilities mapped between \(0\) and \(1\).

Return type:

jax.numpy.ndarray

icdf(u)

Evaluate the combined inverse cumulative distribution function (ICDF).

Parameters:

u (jax.numpy.ndarray) – The probability values. Note that this 1D array corresponds to the probabilities for the model parameters followed by the likelihood noise parameters.

Returns:

The physical parameter values evaluated from the prior distributions.

Return type:

jax.numpy.ndarray

log_likelihood(theta, target)

Evaluate the log-likelihood of the target data.

This handles expanding 1D parameters into the 2D format expected by the vmapped feature extractor, and computes the probability density of the target data against the selected Gaussian or Multivariate Gaussian distribution.

Parameters:

  • theta (jax.numpy.ndarray) – The concatenated 1D array containing the model parameters followed by the likelihood noise standard deviations (\(\sigma\)).

  • target (jax.numpy.ndarray) – The extracted target features (measurement data) to evaluate against.

Returns:

The scalar log-likelihood probability.

Return type:

jax.numpy.ndarray

log_prior(theta)

Evaluate the total log-prior probability.

This lazily compiles the JAX graph to sum the log-prior probabilities of both the underlying model parameters and the added likelihood noise parameters.

Parameters:

theta (jax.numpy.ndarray) – The concatenated 1D array containing the model parameters followed by the likelihood parameters.

Returns:

The scalar log-prior probability.

Return type:

jax.numpy.ndarray

make_numpyro_model(targets)

Construct a dynamic NumPyro model function that mirrors the BayesianFitter logic.

This method traces the prior distributions for both the physical model parameters and the likelihood noise parameters using numpyro.sample. It then concatenates them and injects the lazily compiled ParamRF log-likelihood graph directly into the trace as a numpyro.factor node.

Parameters:

targets (jax.numpy.ndarray) – The extracted measurement data to condition the likelihood on.

Returns:

A dynamic parameter-less function representing the complete NumPyro probabilistic model, ready to be passed to a sampling kernel.

Return type:

callable

model_features(theta)

Extract the RF features from the model for a given set of parameters.

This function maps the parameters into the model, simulates it over the defined frequency band, and extracts the target specifications. The entire extraction pipeline is vectorized over the batch dimension and lazily compiled via jax.jit(jax.vmap(...)).

Parameters:

theta (jax.numpy.ndarray) – A 1D array of a single parameter set, or a 2D array representing a batch of parameters.

Returns:

The extracted model features. Matches the batch dimension of theta.

Return type:

jax.numpy.ndarray

Raises:

RuntimeError – If frequency or features were not provided during initialization.

property num_params: int

Total number of active parameters (model free parameters + likelihood noise parameters).

Type:

int

static read_results(stream)

Reconstruct the NumPyro sample dictionary from the binary stream.

Parameters:

stream (BytesIO)

Return type:

Any

run(measured, **kwargs)

Execute the Bayesian fitting routine.

This method intercepts the standard run sequence to automatically resolve the target features, likelihood kind, and noise priors based on the shape and type of the provided measurement data before passing execution to the backend.

Parameters:

  • measured (str or skrf.Network or NetworkCollection) – The measurement data to condition the likelihood on.

  • **kwargs – Additional arguments forwarded to the specific backend solver.

Returns:

The fitted model and the raw results object.

Return type:

tuple[Model, FitResults]

static write_results(stream, results)

Save the NumPyro sample dictionary as compressed numpy arrays.

This overrides the default JSON serialization, providing significantly faster write times and smaller file sizes for large MCMC traces.

Parameters: