API Reference

kdescent

class kdescent.KCalc

__init__(training_x, training_weights=None, num_kernels=20, bandwidth_factor=0.4, num_fourier_kernels=20, fourier_range_factor=4.0, covariant_kernels=True, comm=None)

This KDE object is the fundamental building block of kdescent. It can be used to compare randomized evaluations of the PDF and ECF by training data to model predictions.

Parameters:

• training_x (array-like) – Training data of shape (n_data, n_features)

• training_weights (array-like, optional) – Training weights of shape (n_data,), by default None

• num_kernels (int, optional) – Number of KDE kernels to appriximate the PDF, by default 20

• bandwidth_factor (float, optional) – Increase or decrease the kernel bandwidth, by default 0.4

• num_fourier_kernels (int, optional) – Number of points in k-space to evaluate the ECF, by default 20

• fourier_range_factor (float, optional) – Increase or decrease the Fourier search space, by default 4.0

• covariant_kernels (bool, optional) – By default (True), kernels will align with the principle components of the training data, which can blow up kernel count values in nearly degenerate subspaces. Set False to prevent this

• comm (MPI Communicator, optional) – For parallel computing, this guarantees consistent kernel placements by all MPI ranks within the comm, by default None. WARNING: Do not pass in an MPI communicator if you plan on wrapping kernel drawing with a JIT-compiled function. In this case, be very careful to pass identical randkeys for each MPI rank

compare_fourier_counts(randkey, x, weights=None)

Return randomly-placed evaluations of the ECF (Empirical Characteristic Function = Fourier-transformed PDF)

Parameters:

• x (array-like) – Model data of shape (n_model_data, n_features)

• weights (array-like, optional) – Effective counts with shape (n_model_data,). If supplied, the ECF will be weighted as sum(weights * exp^(…)) at each evaluation in k-space instead of simply sum(exp^(…))

Returns:

• prediction (jnp.ndarray (complex-valued)) – CF evaluations measured on x. Has shape (num_kernels,)

• truth (jnp.ndarray (complex-valued)) – CF evaluations measured on training_x. This is always different due to the random evaluation kernels. Has shape (num_kernels,)

compare_kde_counts(randkey, x, weights=None)

Realize kernel centers and return all kernel-weighted counts

Parameters:

• x (array-like) – Model data of shape (n_model_data, n_features)

• weights (array-like, optional) – Effective counts with shape (n_model_data,). If supplied, function will return sum(weights * kernel_weights) within each kernel instead of simply sum(kernel_weights)

Returns:

• prediction (jnp.ndarray) – KDE counts measured on x. Has shape (num_kernels,)

• truth (jnp.ndarray) – KDE counts measured on training_x. This is always different due to the random kernel placements. Has shape (num_kernels,)

kdescent.adam(lossfunc, guess, nsteps=100, param_bounds=None, learning_rate=0.01, randkey=1, const_randkey=False, **other_kwargs)

Perform gradient descent

Parameters:

• lossfunc (callable) – Function to be minimized via gradient descent. Must be compatible with jax.jit and jax.grad. Must have signature f(params, **other_kwargs)

• guess (array-like) – The starting parameters.

• nsteps (int, optional) – Number of gradient descent iterations to perform, by default 100

• param_bounds (Sequence, optional) – Lower and upper bounds of each parameter of “shape” (ndim, 2). Pass None as the bound for each unbounded parameter, by default None

• learning_rate (float, optional) – Initial Adam learning rate, by default 0.05

• randkey (int, optional) – Random seed or key, by default 1. If not None, lossfunc must accept the “randkey” keyword argument, e.g. lossfunc(params, randkey=key)

• const_randkey (bool, optional) – By default (False), randkey is regenerated at each gradient descent iteration. Remove this behavior by setting const_randkey=True

Returns:

params – List of params throughout the entire gradient descent, of shape (nsteps, n_param)

Return type:

jnp.array

kdescent.bfgs(lossfunc, guess, maxsteps=100, param_bounds=None, randkey=None)

Run BFGS to descend the gradient and optimize the model parameters, given an initial guess. Stochasticity must be held fixed via a random key

Parameters:

• lossfunc (callable) – Function to be minimized via gradient descent. Must be compatible with jax.jit and jax.grad. Must have signature f(params, **other_kwargs)

• guess (array-like) – The starting parameters.

• maxsteps (int, optional) – The maximum number of steps to take, by default 100.

• param_bounds (Sequence, optional) – Lower and upper bounds of each parameter of “shape” (ndim, 2). Pass None as the bound for each unbounded parameter, by default None

• randkey (int | PRNG Key, optional) – Since BFGS requires a deterministic function, this key will be passed to calc_loss_and_grad_from_params() as the “randkey” kwarg as a constant at every iteration, by default None

Returns:

messagestr

describes reason of termination

successboolean

True if converged

funfloat

minimum loss found

xarray

parameters at minimum loss found

jacarray

gradient of loss at minimum loss found

nfevint

number of function evaluations

nitint

number of gradient descent iterations

Return type:

OptimizeResult (contains the following attributes)