BlackJAXNSFitter

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

Bases: BayesianFitter

Bayesian fitter using the BlackJAX nested slice sampling algorithm.

This backend leverages blackjax.nss for fully JAX-compiled nested sampling. It provides both the global evidence (logZ) and posterior samples, natively utilizing hardware acceleration (GPU/TPU) when available.

Parameters:
execute(target, *, fitted_params='maximum-likelihood', nlive_factor=None, num_delete=None, num_inner_steps=None, logZ_convergence=-3.0, seed=0)[source]

Execute the nested sampling run using BlackJAX.

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

Parameters:

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

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

  • nlive_factor (int, optional) – A multiplier to determine the total number of live points. The total number of live points (n_live) evaluates to nlive_factor * num_params. Defaults to 25.

  • num_delete (int, optional) – The number of dead points to replace per sampling iteration. If not provided, it defaults to 10% of n_live on CPU and 50% on GPU/TPU.

  • num_inner_steps (int, optional) – The number of MCMC steps used to generate a new live point. Defaults to 3 * num_params.

  • logZ_convergence (float, default=-3.0) – The convergence threshold for the log-evidence. The sampling loop stops when the estimated remaining evidence is less than this threshold.

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

  • **kwargs – Additional keyword arguments.

Returns:

The fitted model (with parameter groups updated to the full posterior) and the raw anesthetic.NestedSamples object.

Return type:

tuple[Model, Any]

static write_results(stream, results)[source]

Encode anesthetic NestedSamples into a CSV stream for HDF5 serialization.

Parameters:
static read_results(stream)[source]

Reconstruct anesthetic NestedSamples from a CSV stream.

Parameters:

stream (BytesIO)

Return type:

Any

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

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

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]