TorusImaging1DSpline#

class torusimaging.TorusImaging1DSpline(label_knots, e_knots, e_signs, regularization_func=None, units=<UnitSystem (kpc, Myr, solMass, rad)>, elliptical_angle_convention='pv')#

Bases: TorusImaging1D

A version of the TorusImaging1D model that uses splines to model the label function and the Fourier coefficient \(e_m\) functions.

Parameters:

  • label_knots (Union[Array, ndarray, bool, number, bool, int, float, complex]) – The spline knot locations for the label function.

  • e_knots (dict[int, Union[Array, ndarray, bool, number, bool, int, float, complex]]) – A dictionary keyed by m integers with values as the spline knot locations for the e functions.

  • e_signs (dict[int, float | int]) – A dictionary keyed by m integers with values as the signs of the e functions.

  • regularization_func (Callable[[Any], Array] | None) – A function that takes in a TorusImaging1DSpline instance and a parameter dictionary and returns an additional regularization term to add to the log-likelihood.

  • units (UnitSystem) – A Gala gala.units.UnitSystem instance.

  • elliptical_angle_convention (Literal['pv', 'vp']) – A string that indicates the order of position and velocity in the elliptical angle definition. “pv” indicates tan{theta_e} = pos / vel, and “vp” indicates tan{theta_e} = vel / pos.

Methods Summary

auto_init(binned_data, label_knots, e_knots)

check_e_funcs(e_params, r_e_max)

Check that the parameter values and functions used for the e functions are valid given the condition that d(r)/d(r_e) > 0.

compute_action_angle(pos, vel, params[, ...])

Compute the vertical period, action, and angle given input phase-space coordinates.

compute_elliptical(pos, vel, params)

Compute the elliptical radius \(r_e\) (r_e) and angle :math:` heta_e'` (theta_e)

estimate_init_params(binned_data, bounds)

get_acceleration(pos, params)

Compute the acceleration as a function of position in the limit as velocity goes to zero

get_acceleration_deriv(pos, params)

Compute the derivative of the acceleration with respect to position as a function of position in the limit as velocity goes to zero

get_crlb(params, data[, objective, inv])

Returns the Cramer-Rao lower bound matrix for the parameters evaluated at the input parameter values.

get_crlb_error_samples(params, data[, ...])

Generate Gaussian samples of parameter values centered on the input parameter values with covariance matrix set by the Cramer-Rao lower bound matrix.

get_crlb_uncertainties(params, data[, objective])

Compute the uncertainties on the parameters using the diagonal of the Cramer-Rao lower bound matrix (see get_crlb()).

get_label(pos, vel, params)

Compute the model predicted label value given the input phase-space coordinates

ln_gaussian_likelihood(params, pos, vel, ...)

Compute the log-likelihood of the Gaussian likelihood function.

ln_poisson_likelihood(params, pos, vel, counts)

Compute the log-likelihood of the Poisson likelihood function.

mcmc_run_label(binned_data, p0[, bounds, ...])

Currently only supports uniform priors on all parameters, specified by the input bounds.

objective_gaussian(params, pos, vel, label, ...)

objective_poisson(params, pos, vel, counts)

optimize(params0, objective[, bounds, ...])

Optimize the model parameters given the input data using jaxopt.ScipyboundedMinimize.

unpack_bounds(bounds)

Split a bounds dictionary that is specified like: {"key": (lower, upper)} into two bounds dictionaries for the lower and upper bounds separately, e.g., for the example above: {"key": lower} and {"key": upper}.

Methods Documentation

classmethod auto_init(binned_data, label_knots, e_knots, e_signs=None, regularization_func=None, units=<UnitSystem (kpc, Myr, solMass, rad)>, label_knots_spacing_power=1.0, e_knots_spacing_power=1.0, re_max_factor=1.0, bounds=None, **kwargs)#
Parameters:

  • binned_data (dict[str, Union[Array, ndarray, bool, number, bool, int, float, complex]]) – A dictionary with keys “pos”, “vel”, “label”, “label_err”.

  • label_knots (Union[int, _Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]) – Either an integer number of knots to use, or an array of knot positions.

  • e_knots (dict[int, Union[int, _Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]]) – A dictionary keyed by the m order of the e function, with values either the number of knots to use, or an array of knot positions.

  • e_signs (dict[int, float | int] | None) – A dictionary keyed by the m order of the e function, with values 1 or -1 to represent the sign of the gradient of the e function.

  • regularization_func (Callable[[Any], Array] | bool | None) – A function that takes in two arguments: a TorusImaging1DSpline instance and a parameter dictionary and returns an additional regularization term to add to the log-likelihood. If not specified, this defaults to the torusimaging.model_spline.regularization_function_default() and additional arguments to that function must be specified here. The default regularization function tries to enforce smoothness on the splines, and that the density is positive. It requires the following keyword arguments: label_l2_sigma, label_smooth_sigma, e_l2_sigmas, e_smooth_sigmas. If False, no regularization is applied.

  • units (UnitSystem) – A Gala gala.units.UnitSystem instance.

  • label_knots_spacing_power (float) – If label_knots is an integer number of knots, this controls the spacing of the knots. The knots are placed between 0 and re_max_factor times the estimated maximum elliptical radius in the data, with spacing such that the spacing in r^``label_knots_spacing_power`` is uniform.

  • e_knots_spacing_power (float) – If the values in e_knots are integers number of knots, this controls the spacing of the knots. The knots are placed between 0 and re_max_factor times the estimated maximum elliptical radius in the data, with spacing such that the spacing in r^``e_knots_spacing_power`` is uniform.

  • re_max_factor (float) – A factor to multiply the estimated maximum elliptical radius in the data by to set the maximum knot position.

  • bounds (dict[str, Any] | None) – If provided, this dictionary of parameter bounds is used instead of estimating bounds automatically.

  • **kwargs (Any) – All other keyword arguments are passed to the constructor.

Return type:

tuple[TorusImaging1DSpline, dict[str, Any], TorusImaging1DParams]

check_e_funcs(e_params, r_e_max)#

Check that the parameter values and functions used for the e functions are valid given the condition that d(r)/d(r_e) > 0.

Parameters:
Return type:

tuple[bool, ndarray[tuple[Any, ...], dtype[TypeVar(_ScalarT, bound= generic)]]]

compute_action_angle(pos, vel, params, N_grid=32, Bisection_kwargs=None)#

Compute the vertical period, action, and angle given input phase-space coordinates.

Parameters:

  • pos (Annotated[Quantity]) – The position values.

  • vel (Annotated[Quantity]) – The velocity values.

  • params (TorusImaging1DParams) – A dictionary of model parameters.

  • N_grid (int) – The number of grid points to use in estimating the action integral.

  • Bisection_kwargs (dict[str, Any] | None)

Return type:

QTable

compute_elliptical(pos, vel, params)#

Compute the elliptical radius \(r_e\) (r_e) and angle :math:` heta_e’` (theta_e)

Parameters:

  • pos (Annotated[Quantity]) – The position values.

  • vel (Annotated[Quantity]) – The velocity values.

  • params (TorusImaging1DParams) – A dictionary of model parameters.

Return type:

tuple[Quantity, Quantity]

estimate_init_params(binned_data, bounds)#
Parameters:
Return type:

TorusImaging1DParams

get_acceleration(pos, params)#

Compute the acceleration as a function of position in the limit as velocity goes to zero

Parameters:

  • pos (Annotated[Quantity]) – The position values.

  • params (TorusImaging1DParams) – A dictionary of model parameters.

Return type:

Quantity

get_acceleration_deriv(pos, params)#

Compute the derivative of the acceleration with respect to position as a function of position in the limit as velocity goes to zero

Parameters:

  • pos (Annotated[Quantity]) – The position values.

  • params (TorusImaging1DParams) – A dictionary of model parameters.

Return type:

Quantity

get_crlb(params, data, objective='gaussian', inv=False)#

Returns the Cramer-Rao lower bound matrix for the parameters evaluated at the input parameter values.

To instead return the Fisher information matrix, specify inv=True.

Parameters:
Return type:

ndarray[tuple[Any, ...], dtype[TypeVar(_ScalarT, bound= generic)]]

get_crlb_error_samples(params, data, objective='gaussian', size=1, seed=None, list_of_samples=True)#

Generate Gaussian samples of parameter values centered on the input parameter values with covariance matrix set by the Cramer-Rao lower bound matrix.

Parameters:
Return type:

list[dict] | dict[str, Union[dict, _Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]]

get_crlb_uncertainties(params, data, objective='gaussian')#

Compute the uncertainties on the parameters using the diagonal of the Cramer-Rao lower bound matrix (see get_crlb()).

Parameters:
Return type:

dict[str, Union[dict, _Buffer, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], complex, bytes, str, _NestedSequence[complex | bytes | str]]]

get_label(pos, vel, params)#

Compute the model predicted label value given the input phase-space coordinates

Parameters:

  • pos (Annotated[Quantity])

  • vel (Annotated[Quantity])

  • params (TorusImaging1DParams)

Return type:

Array

ln_gaussian_likelihood(params, pos, vel, label, label_err)#

Compute the log-likelihood of the Gaussian likelihood function.

Note: the input position and velocity arrays must already be converted to the unit system of the model.

Parameters:

  • params (TorusImaging1DParams)

  • pos (Array)

  • vel (Array)

  • label (Array)

  • label_err (Array)

Return type:

Array

ln_poisson_likelihood(params, pos, vel, counts)#

Compute the log-likelihood of the Poisson likelihood function. This should be used when the label you are modeling is the log-number of stars per pixel, i.e. the phase-space density itself.

Note: the input position and velocity arrays must already be converted to the unit system of the model.

Parameters:

  • params (TorusImaging1DParams)

  • pos (Array)

  • vel (Array)

  • counts (Array)

Return type:

Array

mcmc_run_label(binned_data, p0, bounds=None, rng_seed=0, num_steps=1000, num_warmup=1000)#

Currently only supports uniform priors on all parameters, specified by the input bounds.

Parameters:

  • binned_data (dict) – A dictionary containing the binned label moment data.

  • p0 (dict) – The initial values of the parameters.

  • bounds (tuple[dict] | None) – The bounds on the parameters, used to define uniform priors on the parameters. This can either be a tuple of dictionaries, or a dictionary of tuples (keyed by parameter names) to specify the lower and upper bounds for each parameter.

  • rng_seed (int) – The random number generator seed.

  • num_steps (int) – The number of MCMC steps to take.

  • num_warmup (int) – The number of warmup or burn-in steps to take to tune the NUTS sampler.

Return type:

tuple[Any, list[dict]]

Returns:

  • state – The HMCState object returned by BlackJAX.

  • mcmc_samples – A list of dictionaries containing the parameter values for each MCMC sample.

objective_gaussian(params, pos, vel, label, label_err)#
Parameters:

  • params (TorusImaging1DParams)

  • pos (Array)

  • vel (Array)

  • label (Array)

  • label_err (Array)

objective_poisson(params, pos, vel, counts)#
Parameters:
optimize(params0, objective, bounds=None, jaxopt_kwargs=None, **data)#

Optimize the model parameters given the input data using jaxopt.ScipyboundedMinimize.

Parameters:

  • params0 (dict) – The initial values of the parameters.

  • objective (Literal['poisson', 'gaussian']) – The string name of the objective function to use (either “poisson” or “gaussian”).

  • bounds (tuple[dict] | None) – The bounds on the parameters. This can either be a tuple of dictionaries, or a dictionary of tuples (keyed by parameter names) to specify the lower and upper bounds for each parameter.

  • jaxopt_kwargs (dict | None) – Any keyword arguments passed to jaxopt.ScipyBoundedMinimize.

  • **data (Union[Quantity, Array, ndarray, bool, number, bool, int, float, complex]) – Passed through to the objective function.

Return type:

OptStep

classmethod unpack_bounds(bounds)#

Split a bounds dictionary that is specified like: {“key”: (lower, upper)} into two bounds dictionaries for the lower and upper bounds separately, e.g., for the example above: {“key”: lower} and {“key”: upper}.

Parameters:

bounds (dict)

Return type:

tuple[dict]