TorusImaging1D#

class torusimaging.TorusImaging1D(label_func, e_funcs, regularization_func=None, units=<UnitSystem (kpc, Myr, solMass, rad)>)#

Bases: object

A flexible and customizable interface for fitting and MCMC sampling an Orbital Torus Imaging model. This implementation assumes that you are working in a 1 degree of freedom phase space with position coordinate q and velocity coordinate p.

Notation:

\(\Omega_0\) or Omega0: A scale frequency used to compute the elliptical radius r_e. This is the asymptotic orbital frequency at zero action.
\(r_e\) or r_e: The elliptical radius \(r_e = \sqrt{q^2\, \Omega_0 + p^2 \, \Omega_0^{-1}}\).
\(\theta_e\) or theta_e: The elliptical angle defined as \(\tan{\theta_e} = \frac{q}{p}\,\Omega_0\).
\(r\) or r: The distorted elliptical radius \(r = r_e \, f(r_e, \theta_e)\), which is close to \(\sqrt{J}\) (the action) and so we sometimes call it the “proxy action” below. \(f(\cdot)\) is the distortion function defined below.
\(f(r_e, \theta_e)\): The distortion function is a Fourier expansion, defined as: \(f(r_e, \theta_e) = 1 + \sum_m e_m(r_e)\,\cos(m\,\theta_e)\)
\(J\) or J: The action.
\(\theta\) or theta: The conjugate angle.

Parameters:

label_func (Callable[[float], float]) – A function that specifies the dependence of the label function on the distorted radius \(r\).
e_funcs (dict[int, Callable[[float], float]]) – A dictionary that provides functions that specify the dependence of the Fourier distortion coefficients \(e_m(r_e)\). Keys should be the (integer) “m” order of the distortion term (for the distortion function), and values should be Python callable objects that can be passed to jax.jit(). The first argument of each of these functions should be the elliptical radius \(r_e\) or re.
regularization_func (Optional[Callable[[Any], float]]) – An optional function that computes a regularization term to add to the log-likelihood function when optimizing.
units (UnitSystem) – The unit system to work in. Default is to use the “galactic” unit system from Gala: (kpc, Myr, Msun, radian).

Methods Summary

`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`)
`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, ...)	param params:
`objective_poisson`(params, pos, vel, counts)	param params:
`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

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:

e_params (NewType(EParams, dict[int, dict[str, Array | float]])) –
r_e_max (float) –

Return type:

tuple[bool, ndarray[Any, dtype[TypeVar(_ScalarType_co, bound= generic, covariant=True)]]]

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 (Quantity) – The position values.
vel (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 (Optional[dict[str, Any]]) –

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 (Quantity) – The position values.
vel (Quantity) – The velocity values.
params (TorusImaging1DParams) – A dictionary of model parameters.

Return type:

tuple[Quantity, Quantity]

get_acceleration(pos, params)#

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

Parameters:

pos (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 (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:

params (TorusImaging1DParams) –
data (dict[str, Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]]) –
objective (str) –
inv (bool) –

Return type:

ndarray[Any, dtype[TypeVar(_ScalarType_co, bound= generic, covariant=True)]]

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:

params (TorusImaging1DParams) –
data (dict[str, Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]]) –
objective (str) –
size (int) –
seed (Optional[int]) –
list_of_samples (bool) –

Return type:

list[dict] | dict[str, Union[dict, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]]

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:

params (TorusImaging1DParams) –
data (dict[str, Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]]) –
objective (str) –

Return type:

dict[str, Union[dict, _SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]]

get_label(pos, vel, params)#

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

Parameters:

pos (Quantity) –
vel (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 (Optional[tuple[dict]]) – 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:

params (TorusImaging1DParams) –
pos (Array) –
vel (Array) –
counts (Union[_SupportsArray[dtype[Any]], _NestedSequence[_SupportsArray[dtype[Any]]], bool, int, float, complex, str, bytes, _NestedSequence[Union[bool, int, float, complex, str, bytes]]]) –

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 (Optional[tuple[dict]]) – 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 (Optional[dict]) – 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]