Index

A

align_egocentric() (in module jax_moseq.models.keypoint_slds.alignment) angle_to_rotation_matrix() (in module jax_moseq.models.keypoint_slds.alignment)

apply_rotation() (in module jax_moseq.models.keypoint_slds.alignment) ar_to_lds_dynamics() (in module jax_moseq.utils.kalman) ar_to_lds_emissions() (in module jax_moseq.utils.kalman)

B

batch() (in module jax_moseq.utils.utils)

C

center_embedding() (in module jax_moseq.models.keypoint_slds.alignment) compute_squared_error() (in module jax_moseq.models.keypoint_slds.gibbs) (in module jax_moseq.models.slds.gibbs)

concatenate_stateseqs() (in module jax_moseq.utils.utils) continuous_stateseq_log_prob() (in module jax_moseq.models.arhmm.log_prob) count_transitions() (in module jax_moseq.utils.transitions)

D

discrete_stateseq_log_prob() (in module jax_moseq.models.arhmm.log_prob)

E

estimate_aligned() (in module jax_moseq.models.keypoint_slds.alignment)

estimate_coordinates() (in module jax_moseq.models.keypoint_slds.alignment) estimate_error() (in module jax_moseq.models.keypoint_slds.initialize)

F

fit_pca() (in module jax_moseq.models.keypoint_slds.alignment) (in module jax_moseq.utils.utils)

G

get_durations() (in module jax_moseq.utils.utils) get_frequencies() (in module jax_moseq.utils.utils)

get_lags() (in module jax_moseq.utils.autoregression) get_mixed_map_gpus() (in module jax_moseq.utils.utils) get_mixed_map_iters() (in module jax_moseq.utils.utils)

I

init_ar_params() (in module jax_moseq.models.arhmm.initialize) init_continuous_stateseqs() (in module jax_moseq.models.slds.initialize) init_hdp_transitions() (in module jax_moseq.utils.transitions) init_hyperparams() (in module jax_moseq.models.arhmm.initialize) (in module jax_moseq.models.keypoint_slds.initialize) (in module jax_moseq.models.slds.initialize) init_model() (in module jax_moseq.models.arhmm.initialize) (in module jax_moseq.models.keypoint_slds.initialize) (in module jax_moseq.models.slds.initialize)

init_obs_params() (in module jax_moseq.models.slds.initialize) init_params() (in module jax_moseq.models.arhmm.initialize) (in module jax_moseq.models.keypoint_slds.initialize) (in module jax_moseq.models.slds.initialize) init_states() (in module jax_moseq.models.arhmm.initialize) (in module jax_moseq.models.keypoint_slds.initialize) (in module jax_moseq.models.slds.initialize) inverse_rigid_transform() (in module jax_moseq.models.keypoint_slds.alignment)

J

jax_io() (in module jax_moseq.utils.utils) jax_moseq.models.arhmm.gibbs module jax_moseq.models.arhmm.initialize module jax_moseq.models.arhmm.log_prob module jax_moseq.models.keypoint_slds.alignment module jax_moseq.models.keypoint_slds.gibbs module jax_moseq.models.keypoint_slds.initialize module jax_moseq.models.keypoint_slds.log_prob module

jax_moseq.models.slds.gibbs module jax_moseq.models.slds.initialize module jax_moseq.models.slds.log_prob module jax_moseq.utils.autoregression module jax_moseq.utils.distributions module jax_moseq.utils.kalman module jax_moseq.utils.transitions module jax_moseq.utils.utils module

K

kalman_sample() (in module jax_moseq.utils.kalman)

L

location_log_prob() (in module jax_moseq.models.keypoint_slds.log_prob) log_joint_likelihood() (in module jax_moseq.models.arhmm.log_prob) (in module jax_moseq.models.keypoint_slds.log_prob) (in module jax_moseq.models.slds.log_prob)

M

marginal_log_likelihood() (in module jax_moseq.models.arhmm.log_prob) mixed_map() (in module jax_moseq.utils.utils) model_likelihood() (in module jax_moseq.models.arhmm.log_prob) (in module jax_moseq.models.keypoint_slds.log_prob) (in module jax_moseq.models.slds.log_prob) module jax_moseq.models.arhmm.gibbs jax_moseq.models.arhmm.initialize jax_moseq.models.arhmm.log_prob jax_moseq.models.keypoint_slds.alignment jax_moseq.models.keypoint_slds.gibbs jax_moseq.models.keypoint_slds.initialize jax_moseq.models.keypoint_slds.log_prob jax_moseq.models.slds.gibbs jax_moseq.models.slds.initialize jax_moseq.models.slds.log_prob jax_moseq.utils.autoregression jax_moseq.utils.distributions jax_moseq.utils.kalman jax_moseq.utils.transitions jax_moseq.utils.utils

O

obs_log_prob() (in module jax_moseq.models.keypoint_slds.log_prob) (in module jax_moseq.models.slds.log_prob)

P

pad_affine() (in module jax_moseq.utils.utils) pad_along_axis() (in module jax_moseq.utils.utils)

preprocess_for_pca() (in module jax_moseq.models.keypoint_slds.alignment) psd_inv() (in module jax_moseq.utils.utils) psd_solve() (in module jax_moseq.utils.utils)

R

resample_ar_params() (in module jax_moseq.models.arhmm.gibbs) resample_continuous_stateseqs() (in module jax_moseq.models.keypoint_slds.gibbs) (in module jax_moseq.models.slds.gibbs) resample_dir_transitions() (in module jax_moseq.utils.transitions) resample_discrete_stateseqs() (in module jax_moseq.models.arhmm.gibbs) resample_hdp_transitions() (in module jax_moseq.utils.transitions) resample_heading() (in module jax_moseq.models.keypoint_slds.gibbs) resample_location() (in module jax_moseq.models.keypoint_slds.gibbs) resample_model() (in module jax_moseq.models.arhmm.gibbs) (in module jax_moseq.models.keypoint_slds.gibbs) (in module jax_moseq.models.slds.gibbs)

resample_obs_variance() (in module jax_moseq.models.keypoint_slds.gibbs) (in module jax_moseq.models.slds.gibbs) resample_obs_variance_from_sqerr() (in module jax_moseq.models.slds.gibbs) resample_scales() (in module jax_moseq.models.keypoint_slds.gibbs) (in module jax_moseq.models.slds.gibbs) resample_scales_from_sqerr() (in module jax_moseq.models.slds.gibbs) rigid_transform() (in module jax_moseq.models.keypoint_slds.alignment)

S

sample_betas() (in module jax_moseq.utils.transitions) sample_dir_transitions() (in module jax_moseq.utils.transitions) sample_hdp_transitions() (in module jax_moseq.utils.transitions) sample_hmm_stateseq() (in module jax_moseq.utils.distributions) sample_pi() (in module jax_moseq.utils.transitions)

scale_log_prob() (in module jax_moseq.models.slds.log_prob) set_mixed_map_gpus() (in module jax_moseq.utils.utils) set_mixed_map_iters() (in module jax_moseq.utils.utils) state_cross_likelihoods() (in module jax_moseq.models.arhmm.log_prob) stateseq_marginals() (in module jax_moseq.models.arhmm.gibbs) symmetrize() (in module jax_moseq.utils.utils)

T

to_vanilla_slds() (in module jax_moseq.models.keypoint_slds.alignment)

U

unbatch() (in module jax_moseq.utils.utils)

V

vector_to_angle() (in module jax_moseq.models.keypoint_slds.alignment)

W

wrap_angle() (in module jax_moseq.utils.utils)