MoSeq Pipeline Tutorial¶

This notebook provides a tutorial on how to use the MoSeq pipeline to analyze behavioral data. The pipeline is a tool for taking keypoint pose estimations and extracting behavioral syllables.

*Note: Moseq is an optional dependency within the spyglass package. For installation instructions, see the setup tutorial

Here is a schematic showing the tables used in this pipeline. The Basic steps are:

Model Training

Define training data in PoseGroup

Define Moseq model and training parameters in MoseqModelParams

Combine a set of training parameters and training data in MoseqModelSelection

Populate MoseqModel to train

Convert pose data to behavioral syllables

Combine a trained model from MoseqModel and a pose data from PositionOutput in MoseqSyllableSelection

Populate MoseqSyllable to apply the trained model to the selected data

moseq_outline.png|2000x900

Accessing the keypoint (pose) data¶

In the spyglass architecture, keypoint tracking is performed in the Position module, and can be accessed through PositionOutput.fetch_pose_dataframe(). In this tutorial, we are using a set of unpublished data from the Frank lab. For a tutorial on running keypoint extraction in spyglass, see the DLC tutorial notebook.

We can access an example set of keypoint pose data here:

In [ ]:

Copied!





%load_ext autoreload
%autoreload 2
from spyglass.position.position_merge import PositionOutput

# Key defining the DLC data we are using
pose_key = {
    "nwb_file_name": "SC100020230912_.nwb",
    "epoch": 9,
    "video_file_num": 14,
    "project_name": "sideHomeOfficial",
    "dlc_model_name": "sideHomeOfficial_tutorial_00",
    "dlc_model_params_name": "default",
    "task_mode": "trigger",
}

# Fetch the pose data for demo purposes
merge_key = (PositionOutput.DLCPosV1 & pose_key).fetch1("KEY")
pose_df = (PositionOutput & merge_key).fetch_pose_dataframe()
pose_df
%load_ext autoreload
%autoreload 2
from spyglass.position.position_merge import PositionOutput

# Key defining the DLC data we are using
pose_key = {
    "nwb_file_name": "SC100020230912_.nwb",
    "epoch": 9,
    "video_file_num": 14,
    "project_name": "sideHomeOfficial",
    "dlc_model_name": "sideHomeOfficial_tutorial_00",
    "dlc_model_params_name": "default",
    "task_mode": "trigger",
}

# Fetch the pose data for demo purposes
merge_key = (PositionOutput.DLCPosV1 & pose_key).fetch1("KEY")
pose_df = (PositionOutput & merge_key).fetch_pose_dataframe()
pose_df

Out[ ]:

	earL				earR				forelimbL		...	tailBase		tailMid				tailTip
	video_frame_ind	x	y	likelihood	video_frame_ind	x	y	likelihood	video_frame_ind	x	...	y	likelihood	video_frame_ind	x	y	likelihood	video_frame_ind	x	y	likelihood
time
1.694621e+09	0.0	8072.609849	3759.525550	0.825574	0.0	7883.074692	3895.790402	0.028869	0.0	7656.537329	...	4283.328800	0.999663	0.0	4959.599074	4041.634348	0.845290	0.0	4956.257257	4052.028750	0.008635
1.694621e+09	1.0	8261.998988	3770.095027	0.863721	1.0	8225.103732	3798.036042	0.216548	1.0	7601.948953	...	4330.051234	0.999686	1.0	5007.866176	4112.705698	0.919689	1.0	4959.488217	4086.425580	0.011417
1.694621e+09	2.0	8439.114777	3755.394062	0.904633	2.0	8147.386840	3806.549055	0.710284	2.0	7613.715879	...	4325.570084	0.999876	2.0	5076.785832	4131.337436	0.978790	2.0	5091.458472	4159.736284	0.019404
1.694621e+09	3.0	8076.690548	3533.975641	0.496637	3.0	8238.187752	3744.631928	0.976969	3.0	7688.181299	...	4370.041688	0.999769	3.0	5155.758285	4173.914220	0.956584	3.0	5145.784112	4157.734515	0.022773
1.694621e+09	4.0	8093.694597	3484.686411	0.742271	4.0	8284.158606	3675.837536	0.889797	4.0	7745.837035	...	4458.699207	0.999916	4.0	5424.312777	4266.901987	0.949616	4.0	5419.347178	4256.722217	0.042030
...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...
1.694622e+09	23859.0	8323.301275	4157.022004	0.002994	23859.0	8323.682192	4137.912654	0.002425	23859.0	7615.597512	...	4085.460346	0.002940	23859.0	7892.239171	4094.556699	0.019522	23859.0	7542.600577	3850.607998	0.037190
1.694622e+09	23860.0	8322.580462	4164.396172	0.000736	23860.0	8319.023281	4144.923485	0.001018	23860.0	7612.563848	...	4092.408667	0.000326	23860.0	2442.091901	833.311661	0.002724	23860.0	8329.425252	4117.408073	0.025406
1.694622e+09	23861.0	8313.878944	4136.119901	0.002430	23861.0	7611.079736	4317.360829	0.003122	23861.0	7631.078869	...	4087.278981	0.001913	23861.0	7888.670269	4099.420475	0.014236	23861.0	7535.556049	3857.031337	0.048863
1.694622e+09	23862.0	8326.235314	4170.750165	0.000500	23862.0	8321.127116	4151.999268	0.001903	23862.0	7620.843915	...	4209.775628	0.000241	23862.0	7885.540888	4096.438577	0.000256	23862.0	7531.894848	3858.573320	0.039602
1.694622e+09	23863.0	7906.806815	4096.058880	0.002833	23863.0	7560.711728	3926.797069	0.002349	23863.0	7911.580490	...	4079.302183	0.003016	23863.0	7877.772617	4084.137146	0.026200	23863.0	7535.153156	3859.044338	0.049477

23864 rows × 60 columns

To train a moseq model, we first need to define the epochs of pose data we will train on as well as the bodyparts to use within the model. We define this in the PoseGroup table below.

Note that training can be run using data from multiple epochs by passing a list of merge ids to create_group

In [ ]:

Copied!





from spyglass.behavior.v1.core import PoseGroup

# Define the group name and bodyparts to include in the Moseq model
group_name = "tutorial_group"
merge_ids = [(PositionOutput & merge_key).fetch("merge_id")[0]]
bodyparts = [
    "forelimbL",
    "forelimbR",
    "nose",
    "spine1",
    "spine3",
    "spine5",
    "tailBase",
]

# Create the group
PoseGroup().create_group(
    group_name,
    merge_ids,
    bodyparts,
)

# Look at the group in the database
group_key = {"pose_group_name": group_name}
PoseGroup() & group_key
from spyglass.behavior.v1.core import PoseGroup

# Define the group name and bodyparts to include in the Moseq model
group_name = "tutorial_group"
merge_ids = [(PositionOutput & merge_key).fetch("merge_id")[0]]
bodyparts = [
    "forelimbL",
    "forelimbR",
    "nose",
    "spine1",
    "spine3",
    "spine5",
    "tailBase",
]

# Create the group
PoseGroup().create_group(
    group_name,
    merge_ids,
    bodyparts,
)

# Look at the group in the database
group_key = {"pose_group_name": group_name}
PoseGroup() & group_key

Out[ ]:

pose_group_name	bodyparts list of body parts to include in the pose
tutorial_group	=BLOB=

Total: 1

Defining the Moseq Model¶

Next, we make an entry into the MoseqModelParams table. The information in this is used to initialize the moseq model and includes hyperparameters for model training as well as allows you to begin training from an existing model in the database (discussed more below). Relevant parameters can be found in the Moseq documentation

** Note: All bodyparts in the PoseGroup entry will be used in the model

** Note: When training on multiple animals of different sizes (e.g. male and female rats), The Moseq documentation suggests scaling keypoint distances by each animal's size. This can be done here using the optional parameter normalize=True.

In [ ]:

Copied!





from spyglass.behavior.v1.moseq import (
    MoseqModel,
    MoseqModelParams,
    MoseqModelSelection,
)

model_params_name = "tutorial_kappa4_mini"
params = {}
# the skeleton list defines pairs of bodyparts that are linked by and edge
params["skeleton"] = [
    ["nose", "spine1"],
    ["spine1", "forelimbL"],
    ["spine1", "forelimbR"],
    ["spine1", "spine3"],
    ["spine3", "spine5"],
    ["spine5", "tailBase"],
]
# kappa affects the distribution of syllable durations, likely needs tuning for each dataset
params["kappa"] = 1e4
# num_ar_iters is the number of iterations of the autoregressive model for warm-up
params["num_ar_iters"] = 50
# num_epochs is the number of epochs to train the model
params["num_epochs"] = 50
# anterior and posterior bodyparts are used to define the orientation of the animal
params["anterior_bodyparts"] = ["nose"]
params["posterior_bodyparts"] = ["tailBase"]
# Optional: set number of PCs to use; based on number needed to explain target_variance in data
params["target_variance"] = 0.9
# Optional: scale the keypoint distances by animal size
params["normalize"] = False

MoseqModelParams().insert1(
    {"model_params_name": model_params_name, "model_params": params},
    skip_duplicates=True,
)

MoseqModelParams() & {"model_params_name": model_params_name}
from spyglass.behavior.v1.moseq import (
    MoseqModel,
    MoseqModelParams,
    MoseqModelSelection,
)

model_params_name = "tutorial_kappa4_mini"
params = {}
# the skeleton list defines pairs of bodyparts that are linked by and edge
params["skeleton"] = [
    ["nose", "spine1"],
    ["spine1", "forelimbL"],
    ["spine1", "forelimbR"],
    ["spine1", "spine3"],
    ["spine3", "spine5"],
    ["spine5", "tailBase"],
]
# kappa affects the distribution of syllable durations, likely needs tuning for each dataset
params["kappa"] = 1e4
# num_ar_iters is the number of iterations of the autoregressive model for warm-up
params["num_ar_iters"] = 50
# num_epochs is the number of epochs to train the model
params["num_epochs"] = 50
# anterior and posterior bodyparts are used to define the orientation of the animal
params["anterior_bodyparts"] = ["nose"]
params["posterior_bodyparts"] = ["tailBase"]
# Optional: set number of PCs to use; based on number needed to explain target_variance in data
params["target_variance"] = 0.9
# Optional: scale the keypoint distances by animal size
params["normalize"] = False

MoseqModelParams().insert1(
    {"model_params_name": model_params_name, "model_params": params},
    skip_duplicates=True,
)

MoseqModelParams() & {"model_params_name": model_params_name}

Out[ ]:

model_params_name	model_params
tutorial_kappa4_mini	=BLOB=

Total: 1

To train the model, we link a set of model params with training data in PoseGroup using the MoseqModelSelection table.

In [7]:

Copied!





MoseqModelSelection().insert1(
    {
        "model_params_name": model_params_name,
        "pose_group_name": group_name,
    },
    skip_duplicates=True,
)

MoseqModelSelection() & group_key
MoseqModelSelection().insert1(
    {
        "model_params_name": model_params_name,
        "pose_group_name": group_name,
    },
    skip_duplicates=True,
)

MoseqModelSelection() & group_key

Out[7]:

pose_group_name	model_params_name
tutorial_group	tutorial_kappa4_mini

Total: 1

We can then train the model by populating the corresponding MoseqModel entry. This will load the keypoint data, format it for moseq, and then train according to the setting in the MoseqModelParams entry

In [8]:

Copied!





model_key = {
    "model_params_name": model_params_name,
    "pose_group_name": group_name,
}
MoseqModel().populate(model_key)
model_key = {
    "model_params_name": model_params_name,
    "pose_group_name": group_name,
}
MoseqModel().populate(model_key)

The directory
`/stelmo/nwb/moseq/projects/d9bbbaed42eca56d999882d9f65a62b8` already
exists. Use `overwrite=True` or pick a different name
Outputs will be saved to /stelmo/nwb/moseq/projects/d9bbbaed42eca56d99
9882d9f65a62b8/d9bbbaed42eca56d999882d9f65a62b8_ar

 49%|█████████████████▏                 | 25/51 [00:39<00:31,  1.20s/it]

No description has been provided for this image

 98%|██████████████████████████████████▎| 50/51 [01:11<00:01,  1.20s/it]

100%|███████████████████████████████████| 51/51 [01:15<00:00,  1.47s/it]

Outputs will be saved to /stelmo/nwb/moseq/projects/d9bbbaed42eca56d99
9882d9f65a62b8/d9bbbaed42eca56d999882d9f65a62b8_ar

 49%|█████████████████▏                 | 25/51 [01:17<00:33,  1.29s/it]

 98%|██████████████████████████████████▎| 50/51 [01:50<00:01,  1.28s/it]

100%|███████████████████████████████████| 51/51 [01:52<00:00,  2.21s/it]
Reindexing: 100%|█████████████| 5/5 [00:04<00:00,  1.14model snapshot/s]

Out[8]:

{'success_count': 1, 'error_list': []}

The model is now trained and accessible through the the MoseqModel table.

In [9]:

Copied!

trained_model = MoseqModel().fetch_model(model_key)
trained_model
trained_model = MoseqModel().fetch_model(model_key)
trained_model

Out[9]:

{'hypparams': {'ar_hypparams': {'K_0': array([[10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
           10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,  0.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., 10.,
            0.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
           10.,  0.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0., 10.,  0.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0., 10.,  0.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0., 10.,  0.],
          [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,
            0.,  0.,  0.,  0., 10.]]),
   'K_0_scale': 10.0,
   'M_0': array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],
          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0.]]),
   'S_0': array([[0.01, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
          [0.  , 0.01, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
          [0.  , 0.  , 0.01, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
          [0.  , 0.  , 0.  , 0.01, 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
          [0.  , 0.  , 0.  , 0.  , 0.01, 0.  , 0.  , 0.  , 0.  , 0.  ],
          [0.  , 0.  , 0.  , 0.  , 0.  , 0.01, 0.  , 0.  , 0.  , 0.  ],
          [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.01, 0.  , 0.  , 0.  ],
          [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.01, 0.  , 0.  ],
          [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.01, 0.  ],
          [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.01]]),
   'S_0_scale': 0.01,
   'latent_dim': 10,
   'nlags': 3,
   'nu_0': 12,
   'num_states': 100},
  'cen_hypparams': {'sigmasq_loc': 0.5},
  'obs_hypparams': {'nu_s': 5,
   'nu_sigma': 100000.0,
   'sigmasq_0': 0.1,
   'sigmasq_C': 0.1},
  'trans_hypparams': {'alpha': 5.7,
   'gamma': 1000.0,
   'kappa': 10000.0,
   'num_states': 100}},
 'noise_prior': array([[[  85.09396632,  545.57184238,    3.26912788, ...,
             3.74473452,    3.16806043,    3.16017819],
         [  51.81106111,  230.71030275,    3.16274819, ...,
             4.15726905,    3.16711123,    3.16010703],
         [   9.51601252,   20.12595475,    3.16250782, ...,
             3.90254457,    3.16632719,    3.15950745],
         ...,
         [   3.16452054,    3.17060538,    3.16077972, ...,
             3.1607654 ,  785.3479052 ,    3.15928394],
         [   3.16425008,    3.17144215,    3.16506481, ...,
             3.16049012,  655.20722205,    3.15939231],
         [   3.16317888,    3.17151249,    3.16149601, ...,
             3.16017743,  807.39627946,    3.15929523]],
 
        [[   5.20974151,    3.31407256,    3.1595703 , ...,
             3.15975227,  708.58301097,    3.16154236],
         [   6.40493147,    3.40467013,    3.16134418, ...,
             3.15987271,  701.15784208,    3.16455056],
         [   4.25600686,    3.76226764,    3.16022562, ...,
             3.15967511,  890.53950356,    3.16566386],
         ...,
         [2767.0914515 ,  402.89713042,   83.81455615, ...,
          2666.74427478, 1826.81322958, 1494.65445934],
         [2521.47500505,  381.62513487,   17.9536452 , ...,
          2578.18433345, 1932.01796793, 1351.12016796],
         [2754.9473198 ,  190.44217276,   81.71405037, ...,
          2385.74517734, 2056.9868742 , 1693.76296152]],
 
        [[2190.08948711, 2940.58474568, 1753.14427333, ...,
          2643.83088435,  814.70096164, 3073.43228434],
         [2388.36032373, 3033.67586441, 1896.94175212, ...,
          2646.92204068,  309.39586923, 3103.78751556],
         [2306.57000661, 3006.92541804, 2555.19489222, ...,
          2604.64101735,  451.72114915, 3092.30964595],
         ...,
         [2598.07497597, 2618.88649289,   44.46441351, ...,
          2953.46512017, 2279.78258046,  787.20009603],
         [2598.07497597, 2618.88649289,   44.46441351, ...,
          2953.46512017, 2279.78258046,  787.20009603],
         [2598.07497597, 2618.88649289,   44.46441351, ...,
          2953.46512017, 2279.78258046,  787.20009603]]]),
 'params': {'Ab': array([[[-4.46215747e-02, -4.53323925e-03,  1.02696845e-02, ...,
           -1.50481767e-02,  1.05553969e-02, -7.47031645e-04],
          [ 1.39363177e-02, -4.96932401e-02, -1.40526030e-02, ...,
            3.31950219e-03, -1.52263178e-02, -4.68263864e-03],
          [ 1.87693434e-02,  4.31404947e-03,  1.31345930e-02, ...,
           -3.09320059e-02, -6.81530174e-03,  2.33054162e-03],
          ...,
          [-1.73346178e-03,  3.25186133e-02, -2.24636027e-02, ...,
           -9.57252490e-03,  2.72123789e-02, -2.05331225e-03],
          [-1.03397165e-01, -6.41257996e-02, -1.13890819e-01, ...,
            1.50646948e+00,  6.24669921e-02,  3.52949623e-03],
          [-5.29651197e-02,  1.47762826e-02, -1.33999709e-01, ...,
            1.06972384e-01,  1.51488002e+00,  9.46307640e-03]],
  
         [[-2.13619140e-02, -1.57649549e-02, -1.33870458e-02, ...,
            6.12128953e-03, -6.55556637e-03,  1.86060962e-03],
          [-1.57592526e-02, -5.83593048e-02, -2.17692430e-02, ...,
            2.26258549e-03, -1.30057417e-02,  5.40178285e-04],
          [-9.09339298e-05,  3.72411797e-02,  4.08292988e-03, ...,
           -9.07597665e-03,  7.08824090e-03, -3.49221750e-05],
          ...,
          [ 4.00352792e-02, -5.89283466e-02, -2.91629823e-02, ...,
           -4.29674156e-03,  8.31131848e-03, -5.17947600e-03],
          [-6.45401859e-03, -1.80467195e-02,  1.01441745e-02, ...,
            1.02722678e+00,  7.33979129e-02, -3.94744363e-03],
          [ 9.64488981e-02,  4.27237027e-02,  7.60997101e-03, ...,
            6.51252221e-02,  1.09660911e+00, -1.06188676e-02]],
  
         [[-3.09182887e-02, -6.08457768e-02,  5.52088384e-02, ...,
            5.57328905e-02, -6.26678935e-02,  6.75059759e-02],
          [-1.12463320e-01,  2.68419322e-01, -8.88289149e-03, ...,
            3.13660232e-02, -3.44783438e-02,  2.72020360e-03],
          [ 2.21014432e-02, -2.98582149e-01, -1.49517018e-01, ...,
           -5.21899806e-02,  1.04014149e-02, -9.96887294e-02],
          ...,
          [-1.23527489e-01,  1.04868097e+00, -2.61604486e-02, ...,
            4.03090640e-02,  4.41300129e-02, -3.34485232e-02],
          [ 1.14870705e+00,  6.70275282e-01,  2.48247901e-03, ...,
            1.64212400e+00, -3.87876421e-02,  4.95449813e-02],
          [ 1.46460856e+00,  1.03663535e+00, -8.13304897e-03, ...,
            5.62013062e-01,  1.14839741e+00,  2.30373929e-01]],
  
         ...,
  
         [[ 8.32889078e-02, -1.03230867e-02,  4.12325948e-02, ...,
            1.09747963e-03, -1.29044488e-02, -9.86686091e-02],
          [ 3.61856313e-02, -1.94593889e-01,  2.50268862e-01, ...,
           -2.33580546e-01, -2.37767264e-01, -9.02075193e-02],
          [-9.45697457e-02,  4.58398097e-01, -4.37062992e-01, ...,
            4.83871180e-01,  3.22974140e-01,  1.48208696e-01],
          ...,
          [ 9.40950154e-01,  2.66593400e-01,  3.33808601e-01, ...,
            2.29041678e-02,  3.94342385e-02, -1.61493055e-01],
          [ 6.75259431e-02, -1.07783588e+00,  9.08221510e-01, ...,
           -3.81786233e-01, -4.06974075e-01,  1.66693021e-01],
          [ 4.88735399e-01, -1.04874205e+00,  1.67199567e+00, ...,
           -1.06963189e+00, -6.60511564e-01, -9.30878714e-01]],
  
         [[-1.47667529e-01,  2.87801507e-01, -1.74778146e-02, ...,
            5.92254917e-02, -2.03781867e-01,  2.00252955e-01],
          [ 2.38535318e-01, -4.62023171e-02, -1.55241824e-01, ...,
            4.68484853e-01, -3.20663488e-01, -1.53428673e-01],
          [ 5.03322597e-02, -3.12576405e-01, -1.87434243e-01, ...,
            2.20344350e-01, -5.74237885e-02, -1.76909509e-01],
          ...,
          [ 2.30818781e-01, -3.73750666e-01, -1.21179209e-01, ...,
            2.14414008e-01, -8.75369055e-02, -3.52200419e-01],
          [-1.17729662e-01,  3.24681821e-01, -2.45180972e-02, ...,
            1.05605275e+00,  1.87835544e-01,  1.56487795e-01],
          [ 1.05059429e-01, -1.79732361e-01,  4.09082034e-02, ...,
            3.03130532e-01,  7.52518836e-01, -2.03629356e-01]],
  
         [[-1.11081781e+00,  1.73561118e-01, -3.57181391e-01, ...,
            1.40392096e-01, -3.20443452e-01, -4.03708396e-01],
          [ 1.04530655e-01, -6.95985847e-02, -1.00066412e-02, ...,
           -2.08042951e-01,  3.71784704e-02,  1.65457618e-01],
          [ 6.20923767e-01,  1.35799020e-01,  2.69155484e-01, ...,
           -2.14791917e-01,  8.17008665e-02,  5.73279779e-01],
          ...,
          [ 1.75361885e-01,  1.38586333e-01, -1.02925809e-01, ...,
           -9.92891071e-02,  7.87585099e-02,  2.66775931e-02],
          [ 4.81746856e-01, -6.69681827e-02,  2.45275615e-01, ...,
            8.76546707e-01,  3.04595879e-01,  1.16160234e-01],
          [-1.28292390e-01, -6.40350697e-04, -3.82403435e-01, ...,
            2.18771233e-01,  4.98059374e-01,  2.52783135e-01]]]),
  'Cd': array([[ 110.61934405,  -48.44133239, -110.43641753,   42.28271775,
            27.59427701,   53.40522535,  -26.55020875,   78.95180276,
           135.66160738,  108.19984483,  -25.889416  ],
         [-120.63154601,  108.77988962,  -65.14076825,  -44.95867244,
           -26.80421665,   38.23124737,  -29.08911006,  -25.05752639,
            80.32691892,   61.5798872 ,   67.35861504],
         [  70.49551226,  267.02624135, -274.52749732,  427.79740469,
          -192.33593785, -102.03366681,  109.80016318,    9.45430359,
            -7.97263196,    3.96264469,  801.94978128],
         [-116.45855324,  428.37821141, -178.91574472, -230.22132077,
            26.56990908,  -50.17780542,   26.74812666,  135.28691243,
             5.68673262,  -30.85510541,  -66.20716828],
         [ 103.68266935,  131.32268763,  -57.81201121,  255.30405463,
           161.48856257,  291.79993006,  -72.97507689,   74.18033605,
            -2.43259075,  -59.01374239,   40.31959058],
         [ 263.85558565,  651.88742896,  -39.63950311,  -93.59360415,
           -51.09814681,  -14.05255614, -125.04778065, -127.37150514,
            33.16243537,  -18.14431375,  -38.32128396],
         [ 219.39440085,  323.89449403,  488.30915666,  101.81600047,
           217.21174085,  -73.65996798,  121.22427751,   22.45539522,
            30.2700296 ,   16.30614508, -387.28569786],
         [ 817.07019125, -153.93380544, -213.75464402, -181.95987697,
            18.62167913,   13.33176217,   78.17887754,   38.29397356,
           -23.44035221,   -3.53859823, -570.94414708],
         [-118.87137255, -161.21908945, -295.11585752,   31.34903559,
           300.40349113, -104.61625737,   50.18812602,  -97.80960766,
            95.87221463,  -61.40884581, -862.80631933],
         [-261.78579526,  236.50258635, -153.27255087, -127.85103665,
            82.85474131,  190.90241718,  166.10562267,  -80.8930017 ,
           -67.20858796,   57.99975884,   41.24871111],
         [  24.37001477,   61.2851621 , -156.06975795,   94.85876796,
           251.83136966, -135.48880474, -123.17669106,   20.34639992,
          -122.07705374,   77.24670198, -938.67702471],
         [-283.45248673,  141.49033784,  -54.49428241,  -89.73697468,
            15.88145629,  -69.20348696,    9.37856834,  142.054567  ,
            11.64867087,  -27.54618465,   71.41824193]]),
  'Q': array([[[ 5.22323266e-04,  2.38219648e-04, -7.64533216e-05, ...,
           -3.54780839e-04,  3.64061000e-04, -1.20570922e-04],
          [ 2.38219648e-04,  1.57099147e-03, -9.00414571e-04, ...,
            1.26064741e-03, -7.72914889e-04, -1.96548196e-03],
          [-7.64533216e-05, -9.00414571e-04,  1.77357178e-03, ...,
           -1.77251325e-03, -1.42176449e-03, -2.16642619e-03],
          ...,
          [-3.54780839e-04,  1.26064741e-03, -1.77251325e-03, ...,
            8.01750515e-03,  5.91458483e-03,  7.74333492e-03],
          [ 3.64061000e-04, -7.72914889e-04, -1.42176449e-03, ...,
            5.91458483e-03,  3.45359360e-02,  4.03727681e-02],
          [-1.20570922e-04, -1.96548196e-03, -2.16642619e-03, ...,
            7.74333492e-03,  4.03727681e-02,  5.80333017e-02]],
  
         [[ 3.31474497e-05,  2.02067040e-05,  5.56939464e-06, ...,
           -4.08588281e-05,  6.05403399e-06, -6.22883249e-05],
          [ 2.02067040e-05,  8.68541967e-05, -4.51726332e-05, ...,
            4.61097997e-05, -5.89398694e-05, -1.50475320e-04],
          [ 5.56939464e-06, -4.51726332e-05,  8.91198527e-05, ...,
           -9.44945150e-05,  8.01738405e-06, -5.00656069e-06],
          ...,
          [-4.08588281e-05,  4.61097997e-05, -9.44945150e-05, ...,
            3.97333914e-04,  6.44063291e-05,  2.06727301e-04],
          [ 6.05403399e-06, -5.89398694e-05,  8.01738405e-06, ...,
            6.44063291e-05,  1.02190908e-03,  1.09474998e-03],
          [-6.22883249e-05, -1.50475320e-04, -5.00656069e-06, ...,
            2.06727301e-04,  1.09474998e-03,  2.04870109e-03]],
  
         [[ 3.25205969e-02,  1.38548689e-02, -2.58941736e-02, ...,
           -4.93145319e-03,  3.23276987e-02,  3.37514162e-02],
          [ 1.38548689e-02,  9.64636620e-02, -3.44370324e-02, ...,
            8.76950173e-02, -2.37622195e-02, -7.23729778e-02],
          [-2.58941736e-02, -3.44370324e-02,  2.22375973e-01, ...,
           -9.18979881e-02, -3.69116581e-03, -6.18768163e-02],
          ...,
          [-4.93145319e-03,  8.76950173e-02, -9.18979881e-02, ...,
            4.03202003e-01,  2.98432337e-01,  2.82168477e-01],
          [ 3.23276987e-02, -2.37622195e-02, -3.69116581e-03, ...,
            2.98432337e-01,  1.36318320e+00,  1.56400162e+00],
          [ 3.37514162e-02, -7.23729778e-02, -6.18768163e-02, ...,
            2.82168477e-01,  1.56400162e+00,  2.13169697e+00]],
  
         ...,
  
         [[ 1.79604242e-04,  1.27351143e-03, -2.25790153e-03, ...,
            7.75913592e-04,  3.84717100e-03,  6.79729988e-03],
          [ 1.27351143e-03,  2.03362913e-02, -3.40675881e-02, ...,
           -2.45835743e-03,  5.64636841e-02,  1.01328317e-01],
          [-2.25790153e-03, -3.40675881e-02,  5.77847132e-02, ...,
            2.07663016e-03, -9.61974320e-02, -1.70930692e-01],
          ...,
          [ 7.75913592e-04, -2.45835743e-03,  2.07663016e-03, ...,
            1.77763685e-02, -1.58924316e-03, -7.59561515e-03],
          [ 3.84717100e-03,  5.64636841e-02, -9.61974320e-02, ...,
           -1.58924316e-03,  1.63112233e-01,  2.84523886e-01],
          [ 6.79729988e-03,  1.01328317e-01, -1.70930692e-01, ...,
           -7.59561515e-03,  2.84523886e-01,  5.11180722e-01]],
  
         [[ 4.19953155e-03,  1.82549766e-03, -5.20307328e-04, ...,
           -1.69053399e-03,  2.09307285e-03,  1.87663048e-04],
          [ 1.82549766e-03,  9.56672131e-03,  4.39992941e-03, ...,
            1.18415286e-03, -1.33309392e-03,  6.54167133e-03],
          [-5.20307328e-04,  4.39992941e-03,  3.57057569e-03, ...,
            1.78052921e-03, -2.14305039e-03,  3.59208661e-03],
          ...,
          [-1.69053399e-03,  1.18415286e-03,  1.78052921e-03, ...,
            5.15793815e-03, -3.70914941e-03,  1.27456149e-03],
          [ 2.09307285e-03, -1.33309392e-03, -2.14305039e-03, ...,
           -3.70914941e-03,  7.52703347e-03, -2.49393267e-03],
          [ 1.87663048e-04,  6.54167133e-03,  3.59208661e-03, ...,
            1.27456149e-03, -2.49393267e-03,  6.03493136e-03]],
  
         [[ 4.47974636e-02, -2.41775888e-03, -2.08146852e-02, ...,
           -4.04892445e-03, -2.20512672e-02,  1.61999196e-02],
          [-2.41775888e-03,  2.44941120e-03,  2.13244753e-03, ...,
            6.57751193e-04,  1.16383368e-03,  1.96372134e-04],
          [-2.08146852e-02,  2.13244753e-03,  1.17140105e-02, ...,
            2.53656055e-03,  1.01085709e-02, -5.95090238e-03],
          ...,
          [-4.04892445e-03,  6.57751193e-04,  2.53656055e-03, ...,
            2.52487315e-03,  1.65271923e-03, -1.24978183e-03],
          [-2.20512672e-02,  1.16383368e-03,  1.01085709e-02, ...,
            1.65271923e-03,  1.21833260e-02, -9.18403857e-03],
          [ 1.61999196e-02,  1.96372134e-04, -5.95090238e-03, ...,
           -1.24978183e-03, -9.18403857e-03,  1.10270944e-02]]]),
  'betas': array([0.04317816, 0.028366  , 0.02888231, 0.02118798, 0.01116526,
         0.01187804, 0.01514762, 0.00821975, 0.01763858, 0.01284291,
         0.0080156 , 0.0134484 , 0.00961658, 0.01352408, 0.00941192,
         0.00751887, 0.01560189, 0.00797555, 0.01159084, 0.00949425,
         0.01619717, 0.01152705, 0.00868463, 0.01128475, 0.00986286,
         0.0134457 , 0.00962391, 0.0056718 , 0.00991697, 0.00879776,
         0.01228751, 0.01095335, 0.00723672, 0.01201365, 0.00883775,
         0.01401049, 0.0086412 , 0.01130961, 0.01098993, 0.01192615,
         0.01338405, 0.00657512, 0.00892606, 0.00440874, 0.00870401,
         0.00904339, 0.00934786, 0.01111604, 0.00985228, 0.00794174,
         0.01270171, 0.00932687, 0.00978221, 0.01061355, 0.00656652,
         0.00526869, 0.00841172, 0.00778256, 0.01041194, 0.0077332 ,
         0.00909155, 0.0092678 , 0.00700253, 0.00694302, 0.00714616,
         0.00523147, 0.00819137, 0.00715518, 0.00884735, 0.00834463,
         0.00839971, 0.00963812, 0.00762802, 0.00582899, 0.00376967,
         0.01144135, 0.00984846, 0.00874191, 0.01102249, 0.00707488,
         0.00532071, 0.00725955, 0.00715207, 0.01087727, 0.00630275,
         0.01037624, 0.0076442 , 0.00741542, 0.00895867, 0.00614022,
         0.00667958, 0.01034766, 0.00669939, 0.00720433, 0.00468195,
         0.00526503, 0.00768182, 0.00500747, 0.00822967, 0.00431947]),
  'pi': array([[9.56206902e-01, 2.50635321e-02, 5.55510642e-03, ...,
          3.87638759e-05, 6.88387045e-07, 6.42214088e-10],
         [1.89003719e-02, 9.74310609e-01, 3.14674523e-04, ...,
          1.67155886e-63, 4.67272774e-12, 5.30926140e-21],
         [6.57331452e-03, 4.82610991e-04, 9.89367031e-01, ...,
          1.25046373e-22, 1.57697767e-05, 8.74097093e-23],
         ...,
         [2.15085570e-07, 3.70888018e-08, 4.59160781e-05, ...,
          9.99553989e-01, 2.95582665e-13, 1.06948007e-28],
         [2.06812514e-06, 9.80696049e-10, 5.56881463e-08, ...,
          5.02544894e-34, 9.99691367e-01, 2.07875748e-49],
         [4.76330874e-05, 5.71148688e-06, 9.09098113e-07, ...,
          5.10396789e-14, 1.13576205e-19, 9.99524096e-01]]),
  'sigmasq': array([1., 1., 1., 1., 1., 1., 1.])},
 'seed': array([3525552227, 2853420638], dtype=uint32),
 'states': {'h': array([[-3.01630931,  0.25392622, -0.12851277, ..., -0.60136257,
          -0.59470782, -0.58911573],
         [-0.38468165, -0.32258805, -0.26653508, ...,  2.70670462,
           2.70505852,  2.80828143],
         [ 0.39536775,  0.22733078, -0.43050853, ..., -0.60921516,
          -0.67321651, -0.61746923]]),
  's': array([[[3.83954311e+05, 3.89445864e+04, 2.29676262e+05, ...,
           6.45159473e+05, 9.45585765e+05, 7.93946703e+05],
          [4.25003237e+05, 3.69566826e+05, 2.07632262e+05, ...,
           1.07826626e+06, 2.43499497e+05, 4.60355937e+03],
          [1.26161280e+06, 3.51740566e+01, 7.69397041e+00, ...,
           2.31953912e+05, 5.31034475e+00, 6.31718988e+00],
          ...,
          [2.48841510e+00, 5.26640919e+02, 2.32301145e+00, ...,
           9.39493020e+00, 1.75081687e+05, 4.23306268e+00],
          [2.41758370e+00, 4.79760871e+02, 2.59643738e+00, ...,
           3.17620016e+00, 1.50887385e+06, 3.35861446e+00],
          [3.47027566e+00, 3.44640861e+02, 4.64903249e+00, ...,
           4.41433358e+01, 3.71696508e+05, 1.44514508e+00]],
  
         [[7.90716928e+02, 1.14940238e+04, 3.68670906e+04, ...,
           5.57325149e+04, 6.24870362e+04, 3.42441126e+04],
          [3.67647135e+04, 8.16057644e+03, 1.17213123e+04, ...,
           2.98832872e+03, 6.39509012e+03, 1.08326933e+04],
          [7.74976604e+00, 7.88645749e+01, 3.68411900e+00, ...,
           2.70770373e+00, 1.54077163e+03, 3.36101016e+01],
          ...,
          [3.45956028e+05, 4.22973295e+02, 6.80088559e+04, ...,
           3.40571513e+03, 9.66549292e+02, 4.07161670e+05],
          [4.35022174e+05, 3.19230536e+02, 3.14864068e+04, ...,
           1.44371616e+03, 1.02457604e+03, 1.95459529e+04],
          [6.69751824e+03, 1.93612628e+02, 5.90358065e+04, ...,
           4.28835134e+03, 8.47550254e+02, 5.97505435e+04]],
  
         [[1.49195416e+05, 6.55973400e+04, 2.39010489e+05, ...,
           7.23467075e+06, 6.38718750e+05, 1.01157347e+06],
          [2.29847649e+05, 9.29814629e+04, 1.96449769e+05, ...,
           1.35708010e+06, 1.07345814e+05, 1.15807167e+05],
          [2.36285066e+05, 1.34801244e+03, 1.49732880e+05, ...,
           9.38064567e+06, 3.80540924e+02, 3.49921669e+03],
          ...,
          [1.29484416e+03, 6.28183217e+03, 3.65135411e+01, ...,
           1.13410643e+03, 1.85214878e+03, 7.45955144e+02],
          [1.28706321e+03, 2.02995159e+03, 2.88443979e+01, ...,
           1.31461749e+03, 1.21317477e+03, 7.89773862e+02],
          [1.67540251e+03, 1.42800674e+03, 2.69271607e+01, ...,
           9.11623966e+02, 1.26890885e+03, 2.52465571e+02]]]),
  'v': array([[[8054.86320214, 3700.95023036],
          [8054.61147444, 3701.1117163 ],
          [8054.42136108, 3700.19900057],
          ...,
          [2844.40987411, 4379.82435366],
          [2844.99288697, 4379.19280095],
          [2845.83718607, 4378.87909341]],
  
         [[3027.9212918 , 4373.8164837 ],
          [3027.57058411, 4373.55394959],
          [3028.2037225 , 4373.78167417],
          ...,
          [7441.60248752, 3403.55756446],
          [7442.21654844, 3402.12275889],
          [7441.81785375, 3402.21740407]],
  
         [[6827.47557041, 3504.51746425],
          [6826.54078115, 3504.78429073],
          [6827.17267898, 3506.07451299],
          ...,
          [7846.73716954, 4109.15264239],
          [7847.01352953, 4107.02166753],
          [7848.80501905, 4105.41235436]]]),
  'x': array([[[ 1.75551722,  2.59696241, -2.85210746, ...,  1.23149155,
            1.41515904, -1.76716145],
          [-1.97406926,  1.59455599,  2.10559232, ...,  2.00031353,
            2.32647064,  0.05627159],
          [-1.17714702, -0.60871185,  2.0550534 , ..., -1.02056125,
           -4.49706204, -4.59715032],
          ...,
          [-0.2012797 , -0.55066649,  0.54944769, ...,  0.38586197,
           -1.52423826,  0.61919123],
          [-0.20761287, -0.56321661,  0.5529417 , ...,  0.37407516,
           -1.48829803,  0.65311427],
          [-0.21230706, -0.57060573,  0.53895116, ...,  0.38130593,
           -1.50271046,  0.56449664]],
  
         [[ 0.04232713, -1.21919344, -0.26415194, ...,  0.35034371,
           -0.41619882,  0.09500112],
          [-0.23312999, -0.54076603, -0.53938168, ..., -0.05790709,
            1.20346539, -2.98602427],
          [-0.45529968, -0.87805472, -0.16911672, ..., -0.02316295,
            1.93925452, -3.02893195],
          ...,
          [ 0.95221294, -0.44032102, -0.02132928, ..., -1.07228898,
           -0.11845274, -0.30913148],
          [ 0.9897017 , -0.40597589,  0.01207887, ..., -1.16395874,
           -0.22400794, -0.36448998],
          [ 1.07186228, -0.29178944,  0.09389472, ..., -1.93081567,
           -0.94669651,  0.20405419]],
  
         [[ 5.50779846, -0.60799916,  0.41967545, ...,  2.91357466,
            2.45901126,  3.73700994],
          [ 1.4327214 , -1.43784035, -1.44625745, ..., -2.21158007,
            0.47911447,  1.74034002],
          [-0.44565534, -2.26108735, -7.05460819, ..., -2.77927338,
           -2.30177168,  0.22916206],
          ...,
          [ 0.47010155, -0.30374975, -0.03907097, ..., -0.6225977 ,
           -0.33285099,  0.28286693],
          [ 0.47268796, -0.30421846, -0.0320453 , ..., -0.64556003,
           -0.32510917,  0.29499986],
          [ 0.47042952, -0.30616302, -0.03615157, ..., -0.6089098 ,
           -0.36669708,  0.25123201]]]),
  'z': array([[ 2,  2, 30, ...,  1,  1,  1],
         [59, 59, 59, ...,  0,  0,  2],
         [ 6,  6,  6, ..., 29, 29, 29]])}}

We can also analyze components of the trained model; ie. the pca breakdown of the pose skeleton

In [6]:

Copied!

table = MoseqModel() & model_key
table.analyze_pca()
table = MoseqModel() & model_key
table.analyze_pca()

>=90.0% of variance exlained by 6 components.

as well as average trajectories for each syllable

In [11]:

Copied!

table.generate_trajectory_plots()
table.generate_trajectory_plots()

Saving trajectory plots to /stelmo/nwb/moseq/projects/d9bbbaed42eca56d999882d9f65a62b8/d9bbbaed42eca56d999882d9f65a62b8_ar/trajectory_plots

Generating trajectory plots: 100%|████████| 3/3 [00:01<00:00,  2.02it/s]

And example videos of each syllable. These are saved as mp4 files in the passed output_dir

In [ ]:

Copied!

output_dir = "/path/to/save/videos/"
table.generate_grid_movies(output_dir=output_dir)
output_dir = "/path/to/save/videos/"
table.generate_grid_movies(output_dir=output_dir)

Run data through the trained model¶

Now that we have a trained model, we can use it to convert pose data into behavioral syllables. We do so by combining a trained model with an epoch of pose data, and then applying the populate command

In [17]:

Copied!





%load_ext autoreload
%autoreload 2
from spyglass.behavior.v1.moseq import MoseqSyllableSelection, MoseqSyllable

# Make a selection table entry defining the pose data and moseq model to use
pose_key = (
    PoseGroup().Pose().fetch("pose_merge_id", as_dict=True)[0]
)  # can also use data outside of the training epochs
key = {**model_key, **pose_key, "num_iters": 3}
MoseqSyllableSelection().insert1(key, skip_duplicates=True)

# Run populate to apply the model to the pose data
MoseqSyllable().populate(key)
MoseqSyllable()
%load_ext autoreload
%autoreload 2
from spyglass.behavior.v1.moseq import MoseqSyllableSelection, MoseqSyllable

# Make a selection table entry defining the pose data and moseq model to use
pose_key = (
    PoseGroup().Pose().fetch("pose_merge_id", as_dict=True)[0]
)  # can also use data outside of the training epochs
key = {**model_key, **pose_key, "num_iters": 3}
MoseqSyllableSelection().insert1(key, skip_duplicates=True)

# Run populate to apply the model to the pose data
MoseqSyllable().populate(key)
MoseqSyllable()

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

Applying model: 100%|█████████████████████| 3/3 [00:04<00:00,  1.37s/it]

Saved results to /stelmo/nwb/moseq/projects/d9bbbaed42eca56d999882d9f6
5a62b8/d9bbbaed42eca56d999882d9f65a62b8_ar/results.h5

/home/sambray/mambaforge-pypy3/envs/spyglass-moseq-1/lib/python3.9/site-packages/hdmf/utils.py:668: UserWarning:

Ignoring cached namespace 'hdmf-common' version 1.5.1 because version 1.8.0 is already loaded.

/home/sambray/mambaforge-pypy3/envs/spyglass-moseq-1/lib/python3.9/site-packages/hdmf/utils.py:668: UserWarning:

Ignoring cached namespace 'hdmf-experimental' version 0.2.0 because version 0.5.0 is already loaded.

[15:05:48][INFO] Spyglass: Writing new NWB file SC100020230912_LKNI5NCTIH.nwb
INFO:spyglass:Writing new NWB file SC100020230912_LKNI5NCTIH.nwb

Out[17]:

pose_merge_id	pose_group_name	model_params_name	analysis_file_name name of the file	moseq_object_id
1f3fdc6d-d8d5-c283-846a-4830e85dcb32	tutorial_group	tutorial_kappa4_mini	SC100020230912_LKNI5NCTIH.nwb	e550ee2f-4744-46ac-a6f9-f9c87315c0ed

Total: 1

In [18]:

Copied!





import matplotlib.pyplot as plt

moseq_df = (MoseqSyllable() & key).fetch1_dataframe()
moseq_df
ind = slice(1000, 3000)

fig, ax = plt.subplots(2, 1, figsize=(8, 5), sharex=True)
t = moseq_df.index.values[ind]
ax[0].plot(t, moseq_df["centroid x"].values[ind], label="x")
ax[0].plot(t, moseq_df["centroid y"].values[ind], label="y")
ax[1].scatter(
    t,
    moseq_df.syllable.values[ind],
    s=10,
    c=moseq_df.syllable.values[ind],
    cmap="tab20",
)
ax[0].set_ylabel("centroid")
ax[1].set_ylabel("syllable")
ax[1].set_xlabel("time (s)")
import matplotlib.pyplot as plt

moseq_df = (MoseqSyllable() & key).fetch1_dataframe()
moseq_df
ind = slice(1000, 3000)

fig, ax = plt.subplots(2, 1, figsize=(8, 5), sharex=True)
t = moseq_df.index.values[ind]
ax[0].plot(t, moseq_df["centroid x"].values[ind], label="x")
ax[0].plot(t, moseq_df["centroid y"].values[ind], label="y")
ax[1].scatter(
    t,
    moseq_df.syllable.values[ind],
    s=10,
    c=moseq_df.syllable.values[ind],
    cmap="tab20",
)
ax[0].set_ylabel("centroid")
ax[1].set_ylabel("syllable")
ax[1].set_xlabel("time (s)")

Out[18]:

Text(0.5, 0, 'time (s)')

This concludes the tutorial for basic usage of the Moseq pipeline. Next, we will look at usage for extending training from a base model and leveraging spyglass's relational database to easily sweep through model hyperparameters

Extend model training¶

There are many cases where you may want to begin trining from an existing model rather than begin from scratch. Examples include continuing training an incompletely converged entry, or using a pre-trained base model and refining it to a particular set of animals/ imaging conditions (ie. transfer learning).

The spyglass moseq pipeline allows for this style of iterative training. To do so, we will define a new entry in MoseqModelParams using the make_training_extension_params method. This entry will have the same params as those used in model_key, except it will point to the model_key entry for the initial_model

In [19]:

Copied!





# Insert a training extension model entry
extension_params = MoseqModelParams().make_training_extension_params(
    model_key, num_epochs=100, skip_duplicates=True
)
print("initial model: ", extension_params["model_params"]["initial_model"])
new_params_key = {
    "model_params_name": extension_params["model_params_name"],
}
MoseqModelParams() & "model_params_name LIKE '%tutorial_kappa4%'"
# Insert a training extension model entry
extension_params = MoseqModelParams().make_training_extension_params(
    model_key, num_epochs=100, skip_duplicates=True
)
print("initial model: ", extension_params["model_params"]["initial_model"])
new_params_key = {
    "model_params_name": extension_params["model_params_name"],
}
MoseqModelParams() & "model_params_name LIKE '%tutorial_kappa4%'"

initial model:  {'pose_group_name': 'tutorial_group', 'model_params_name': 'tutorial_kappa4_mini'}

Out[19]:

model_params_name	model_params
tutorial_grouptutorial_kappa4_mini_extension001	=BLOB=
tutorial_kappa4_mini	=BLOB=

Total: 2

This extension model can then be trained following the same steps as before

In [20]:

Copied!





new_model_key = {
    **new_params_key,
    "pose_group_name": model_key["pose_group_name"],
}
MoseqModelSelection().insert1(new_model_key, skip_duplicates=True)
MoseqModel().populate(new_model_key)
new_model_key = {
    **new_params_key,
    "pose_group_name": model_key["pose_group_name"],
}
MoseqModelSelection().insert1(new_model_key, skip_duplicates=True)
MoseqModel().populate(new_model_key)

Outputs will be saved to /stelmo/nwb/moseq/projects/45522739aed0cce276
67e45026e3e32c/45522739aed0cce27667e45026e3e32c

 25%|████████▍                         | 25/101 [00:34<01:42,  1.35s/it]

 50%|████████████████▊                 | 50/101 [01:09<01:09,  1.35s/it]

 74%|█████████████████████████▏        | 75/101 [01:44<00:35,  1.35s/it]

 99%|████████████████████████████████▋| 100/101 [02:18<00:01,  1.36s/it]

100%|█████████████████████████████████| 101/101 [02:20<00:00,  1.40s/it]
Reindexing: 100%|█████████████| 5/5 [00:01<00:00,  4.67model snapshot/s]

Hyperparameter search (kappa)¶

The relational database structure makes it relatively easy to train and organize multiple models on the same data. Here we demonstrate leveraging this architecture to test values of kappa in the moseq model. The kappa value determines the rate of syllable transitions, with larger values corresponding to longer syllables (moseq docs). This value will likely need tuned for your specific data to achieve a syllable distribution at appropriate timescales.

To do so we will make a set of parameter entries with varying values of kappa and then training an initial model for each. Looking at the results above, we see a shorter median distribution (~3 frames = 100ms) than we would like (recommended ~400ms) . We will therefore try several parameter sets with larger kappa values.

In [26]:

Copied!





original_params = (MoseqModelParams() & model_key).fetch1("model_params")

new_params_key_list = []
for i in [5, 6, 7, 8]:
    new_params = original_params.copy()
    new_params["kappa"] = 10**i
    new_params["num_epochs"] = 100
    new_model_params_name = f"tutorial_kappa{i}_mini"
    new_params_key = {
        "model_params_name": new_model_params_name,
        "model_params": new_params,
    }
    new_params_key_list.append(new_params_key)
    MoseqModelParams().insert1(new_params_key, skip_duplicates=True)

MoseqModelParams() & "model_params_name LIKE '%tutorial_kappa%'"
original_params = (MoseqModelParams() & model_key).fetch1("model_params")

new_params_key_list = []
for i in [5, 6, 7, 8]:
    new_params = original_params.copy()
    new_params["kappa"] = 10**i
    new_params["num_epochs"] = 100
    new_model_params_name = f"tutorial_kappa{i}_mini"
    new_params_key = {
        "model_params_name": new_model_params_name,
        "model_params": new_params,
    }
    new_params_key_list.append(new_params_key)
    MoseqModelParams().insert1(new_params_key, skip_duplicates=True)

MoseqModelParams() & "model_params_name LIKE '%tutorial_kappa%'"

Out[26]:

model_params_name	model_params
tutorial_grouptutorial_kappa4_mini_extension001	=BLOB=
tutorial_kappa4_mini	=BLOB=
tutorial_kappa5_mini	=BLOB=
tutorial_kappa6_mini	=BLOB=
tutorial_kappa7_mini	=BLOB=
tutorial_kappa8_mini	=BLOB=

Total: 6

We can now train a model for each of these entries. We are training several different models here, so depending on your hardware, now may be a good time for a coffee break.

In [27]:

Copied!





for new_params_key in new_params_key_list:
    new_model_key = {
        "model_params_name": new_params_key["model_params_name"],
        "pose_group_name": model_key["pose_group_name"],
    }
    MoseqModelSelection().insert1(new_model_key, skip_duplicates=True)
MoseqModel().populate()
for new_params_key in new_params_key_list:
    new_model_key = {
        "model_params_name": new_params_key["model_params_name"],
        "pose_group_name": model_key["pose_group_name"],
    }
    MoseqModelSelection().insert1(new_model_key, skip_duplicates=True)
MoseqModel().populate()

Outputs will be saved to /stelmo/nwb/moseq/projects/411c8460ca6c2e53b6
e62826b7e49bd7/411c8460ca6c2e53b6e62826b7e49bd7_ar

 49%|█████████████████▏                 | 25/51 [00:32<00:33,  1.27s/it]

 98%|██████████████████████████████████▎| 50/51 [01:05<00:01,  1.27s/it]

100%|███████████████████████████████████| 51/51 [01:07<00:00,  1.32s/it]
/home/sambray/mambaforge-pypy3/envs/spyglass-moseq-1/lib/python3.9/site-packages/keypoint_moseq/fitting.py:589: UserWarning:

'kappa' with <class 'int'> will be cast to <class 'float'>

Outputs will be saved to /stelmo/nwb/moseq/projects/411c8460ca6c2e53b6
e62826b7e49bd7/411c8460ca6c2e53b6e62826b7e49bd7_ar

 25%|████████▍                         | 25/101 [00:33<01:40,  1.32s/it]

 50%|████████████████▊                 | 50/101 [01:06<01:07,  1.32s/it]

 74%|█████████████████████████▏        | 75/101 [01:40<00:34,  1.32s/it]

 99%|████████████████████████████████▋| 100/101 [02:14<00:01,  1.32s/it]

100%|█████████████████████████████████| 101/101 [02:16<00:00,  1.35s/it]
Reindexing: 100%|█████████████| 7/7 [00:01<00:00,  4.74model snapshot/s]

Outputs will be saved to /stelmo/nwb/moseq/projects/8b46f4d1412b3b5343
c7568134c8ecac/8b46f4d1412b3b5343c7568134c8ecac_ar

 49%|█████████████████▏                 | 25/51 [00:31<00:33,  1.27s/it]

 98%|██████████████████████████████████▎| 50/51 [01:04<00:01,  1.27s/it]

100%|███████████████████████████████████| 51/51 [01:06<00:00,  1.30s/it]
/home/sambray/mambaforge-pypy3/envs/spyglass-moseq-1/lib/python3.9/site-packages/keypoint_moseq/fitting.py:589: UserWarning:

'kappa' with <class 'int'> will be cast to <class 'float'>

Outputs will be saved to /stelmo/nwb/moseq/projects/8b46f4d1412b3b5343
c7568134c8ecac/8b46f4d1412b3b5343c7568134c8ecac_ar

 25%|████████▍                         | 25/101 [00:33<01:40,  1.32s/it]

 50%|████████████████▊                 | 50/101 [01:07<01:07,  1.32s/it]

 74%|█████████████████████████▏        | 75/101 [01:40<00:34,  1.32s/it]

 99%|████████████████████████████████▋| 100/101 [02:14<00:01,  1.32s/it]

100%|█████████████████████████████████| 101/101 [02:16<00:00,  1.35s/it]
Reindexing: 100%|█████████████| 7/7 [00:01<00:00,  4.86model snapshot/s]

Outputs will be saved to /stelmo/nwb/moseq/projects/aef642d24bf39d00d3
4f4f187f20c545/aef642d24bf39d00d34f4f187f20c545_ar

 49%|█████████████████▏                 | 25/51 [00:32<00:33,  1.27s/it]

 98%|██████████████████████████████████▎| 50/51 [01:04<00:01,  1.27s/it]

100%|███████████████████████████████████| 51/51 [01:06<00:00,  1.31s/it]
/home/sambray/mambaforge-pypy3/envs/spyglass-moseq-1/lib/python3.9/site-packages/keypoint_moseq/fitting.py:589: UserWarning:

'kappa' with <class 'int'> will be cast to <class 'float'>

Outputs will be saved to /stelmo/nwb/moseq/projects/aef642d24bf39d00d3
4f4f187f20c545/aef642d24bf39d00d34f4f187f20c545_ar

 25%|████████▍                         | 25/101 [00:33<01:40,  1.32s/it]

 50%|████████████████▊                 | 50/101 [01:06<01:07,  1.32s/it]

 74%|█████████████████████████▏        | 75/101 [01:40<00:34,  1.32s/it]

 99%|████████████████████████████████▋| 100/101 [02:14<00:01,  1.32s/it]

100%|█████████████████████████████████| 101/101 [02:16<00:00,  1.35s/it]
Reindexing: 100%|█████████████| 7/7 [00:01<00:00,  4.72model snapshot/s]

Outputs will be saved to /stelmo/nwb/moseq/projects/3230c67ed1b7a904a1
73d32a721f8e4c/3230c67ed1b7a904a173d32a721f8e4c_ar

 49%|█████████████████▏                 | 25/51 [00:31<00:33,  1.27s/it]

 98%|██████████████████████████████████▎| 50/51 [01:04<00:01,  1.27s/it]

100%|███████████████████████████████████| 51/51 [01:06<00:00,  1.30s/it]
/home/sambray/mambaforge-pypy3/envs/spyglass-moseq-1/lib/python3.9/site-packages/keypoint_moseq/fitting.py:589: UserWarning:

'kappa' with <class 'int'> will be cast to <class 'float'>

Outputs will be saved to /stelmo/nwb/moseq/projects/3230c67ed1b7a904a1
73d32a721f8e4c/3230c67ed1b7a904a173d32a721f8e4c_ar

 25%|████████▍                         | 25/101 [00:33<01:40,  1.32s/it]

 50%|████████████████▊                 | 50/101 [01:07<01:07,  1.32s/it]

 74%|█████████████████████████▏        | 75/101 [01:40<00:34,  1.32s/it]

 99%|████████████████████████████████▋| 100/101 [02:14<00:01,  1.32s/it]

100%|█████████████████████████████████| 101/101 [02:16<00:00,  1.35s/it]
Reindexing: 100%|█████████████| 7/7 [00:01<00:00,  4.82model snapshot/s]

You can then choose the model that best matches your syllable duration of interest and continue training it using the training extension described above

In [28]:

Copied!





# print out link to the pdf of training results (includes syllable durations)
for new_params_key in new_params_key_list:
    new_model_key = {
        "model_params_name": new_params_key["model_params_name"],
        "pose_group_name": model_key["pose_group_name"],
    }
    training_results_path = MoseqModel().get_training_progress_path(
        new_model_key
    )
    print(
        f"{new_model_key['model_params_name']} training results: {training_results_path}"
    )
# print out link to the pdf of training results (includes syllable durations)
for new_params_key in new_params_key_list:
    new_model_key = {
        "model_params_name": new_params_key["model_params_name"],
        "pose_group_name": model_key["pose_group_name"],
    }
    training_results_path = MoseqModel().get_training_progress_path(
        new_model_key
    )
    print(
        f"{new_model_key['model_params_name']} training results: {training_results_path}"
    )

tutorial_kappa5_mini training results: /stelmo/nwb/moseq/projects/411c8460ca6c2e53b6e62826b7e49bd7/411c8460ca6c2e53b6e62826b7e49bd7_ar/fitting_progress.pdf
tutorial_kappa6_mini training results: /stelmo/nwb/moseq/projects/8b46f4d1412b3b5343c7568134c8ecac/8b46f4d1412b3b5343c7568134c8ecac_ar/fitting_progress.pdf
tutorial_kappa7_mini training results: /stelmo/nwb/moseq/projects/aef642d24bf39d00d34f4f187f20c545/aef642d24bf39d00d34f4f187f20c545_ar/fitting_progress.pdf
tutorial_kappa8_mini training results: /stelmo/nwb/moseq/projects/3230c67ed1b7a904a173d32a721f8e4c/3230c67ed1b7a904a173d32a721f8e4c_ar/fitting_progress.pdf

In [ ]: