.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/compute_kinematics.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_compute_kinematics.py>`
        to download the full example code. or to run this example in your browser via Binder

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_compute_kinematics.py:

Compute and visualise kinematics.
====================================

Compute displacement, velocity and acceleration, and
visualise the results.

.. GENERATED FROM PYTHON SOURCE LINES 10-12

Imports
-------

.. GENERATED FROM PYTHON SOURCE LINES 12-21

.. code-block:: Python


    # For interactive plots: install ipympl with `pip install ipympl` and uncomment
    # the following line in your notebook
    # %matplotlib widget
    from matplotlib import pyplot as plt

    from movement import sample_data
    from movement.utils.vector import compute_norm








.. GENERATED FROM PYTHON SOURCE LINES 22-26

Load sample dataset
------------------------
First, we load an example dataset. In this case, we select the
``SLEAP_three-mice_Aeon_proofread`` sample data.

.. GENERATED FROM PYTHON SOURCE LINES 26-32

.. code-block:: Python

    ds = sample_data.fetch_dataset(
        "SLEAP_three-mice_Aeon_proofread.analysis.h5",
    )

    print(ds)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    <xarray.Dataset> Size: 27kB
    Dimensions:      (time: 601, individuals: 3, keypoints: 1, space: 2)
    Coordinates:
      * time         (time) float64 5kB 0.0 0.02 0.04 0.06 ... 11.96 11.98 12.0
      * individuals  (individuals) <U10 120B 'AEON3B_NTP' 'AEON3B_TP1' 'AEON3B_TP2'
      * keypoints    (keypoints) <U8 32B 'centroid'
      * space        (space) <U1 8B 'x' 'y'
    Data variables:
        position     (time, individuals, keypoints, space) float32 14kB 770.3 ......
        confidence   (time, individuals, keypoints) float32 7kB nan nan ... nan nan
    Attributes:
        fps:              50.0
        time_unit:        seconds
        source_software:  SLEAP
        source_file:      /home/runner/.movement/data/poses/SLEAP_three-mice_Aeon...
        ds_type:          poses
        frame_path:       /home/runner/.movement/data/frames/three-mice_Aeon_fram...
        video_path:       None




.. GENERATED FROM PYTHON SOURCE LINES 33-38

We can see in the printed description of the dataset ``ds`` that
the data was acquired at 50 fps, and the time axis is expressed in seconds.
It includes data for three individuals(``AEON3B_NTP``, ``AEON3B_TP1``,
and ``AEON3B_TP2``), and only one keypoint called ``centroid`` was tracked
in ``x`` and ``y`` dimensions.

.. GENERATED FROM PYTHON SOURCE LINES 40-44

The loaded dataset ``ds`` contains two data arrays:
``position`` and ``confidence``.
To compute displacement, velocity and acceleration, we will need the
``position`` one:

.. GENERATED FROM PYTHON SOURCE LINES 44-47

.. code-block:: Python

    position = ds.position









.. GENERATED FROM PYTHON SOURCE LINES 48-52

Visualise the data
---------------------------
First, let's visualise the trajectories of the mice in the XY plane,
colouring them by individual.

.. GENERATED FROM PYTHON SOURCE LINES 52-73

.. code-block:: Python


    fig, ax = plt.subplots(1, 1)
    for mouse_name, col in zip(
        position.individuals.values, ["r", "g", "b"], strict=False
    ):
        ax.plot(
            position.sel(individuals=mouse_name, space="x"),
            position.sel(individuals=mouse_name, space="y"),
            linestyle="-",
            marker=".",
            markersize=2,
            linewidth=0.5,
            c=col,
            label=mouse_name,
        )
        ax.invert_yaxis()
        ax.set_xlabel("x (pixels)")
        ax.set_ylabel("y (pixels)")
        ax.axis("equal")
        ax.legend()




.. image-sg:: /examples/images/sphx_glr_compute_kinematics_001.png
   :alt: compute kinematics
   :srcset: /examples/images/sphx_glr_compute_kinematics_001.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 74-78

We can see that the trajectories of the three mice are close to a circular
arc. Notice that the x and y axes are set to equal scales, and that the
origin of the coordinate system is at the top left of the image. This
follows the convention for SLEAP and most image processing tools.

.. GENERATED FROM PYTHON SOURCE LINES 80-81

We can also color the data points based on their timestamps:

.. GENERATED FROM PYTHON SOURCE LINES 81-98

.. code-block:: Python

    fig, axes = plt.subplots(3, 1, sharey=True)
    for mouse_name, ax in zip(position.individuals.values, axes, strict=False):
        sc = ax.scatter(
            position.sel(individuals=mouse_name, space="x"),
            position.sel(individuals=mouse_name, space="y"),
            s=2,
            c=position.time,
            cmap="viridis",
        )
        ax.invert_yaxis()
        ax.set_title(mouse_name)
        ax.set_xlabel("x (pixels)")
        ax.set_ylabel("y (pixels)")
        ax.axis("equal")
        fig.colorbar(sc, ax=ax, label="time (s)")
    fig.tight_layout()




.. image-sg:: /examples/images/sphx_glr_compute_kinematics_002.png
   :alt: AEON3B_NTP, AEON3B_TP1, AEON3B_TP2
   :srcset: /examples/images/sphx_glr_compute_kinematics_002.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 99-103

These plots show that for this snippet of the data,
two of the mice (``AEON3B_NTP`` and ``AEON3B_TP1``)
moved around the circle in clockwise direction, and
the third mouse (``AEON3B_TP2``) followed an anti-clockwise direction.

.. GENERATED FROM PYTHON SOURCE LINES 105-109

We can also easily plot the components of the position vector against time
using ``xarray``'s built-in plotting methods. We use
:meth:`xarray.DataArray.squeeze` to
remove the dimension of length 1 from the data (the ``keypoints`` dimension).

.. GENERATED FROM PYTHON SOURCE LINES 109-112

.. code-block:: Python

    position.squeeze().plot.line(x="time", row="individuals", aspect=2, size=2.5)
    plt.gcf().show()




.. image-sg:: /examples/images/sphx_glr_compute_kinematics_003.png
   :alt: individuals = AEON3B_NTP, individuals = AEON3B_TP1, individuals = AEON3B_TP2
   :srcset: /examples/images/sphx_glr_compute_kinematics_003.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 113-116

If we use ``xarray``'s plotting function, the axes units are automatically
taken from the data array. In our case, ``time`` is expressed in seconds,
and the ``x`` and ``y`` coordinates of the ``position`` are in pixels.

.. GENERATED FROM PYTHON SOURCE LINES 118-124

Compute displacement
---------------------
We can start off by computing the distance travelled by the mice along
their trajectories.
For this, we can use the ``compute_displacement`` method of the
``move`` accessor.

.. GENERATED FROM PYTHON SOURCE LINES 124-126

.. code-block:: Python

    displacement = ds.move.compute_displacement()








.. GENERATED FROM PYTHON SOURCE LINES 127-130

This method will return a data array equivalent to the ``position`` one,
but holding displacement data along the ``space`` axis, rather than
position data.

.. GENERATED FROM PYTHON SOURCE LINES 132-134

Notice that we could also compute the displacement (and all the other
kinematic variables) using the :mod:`movement.analysis.kinematics` module:

.. GENERATED FROM PYTHON SOURCE LINES 136-140

.. code-block:: Python

    import movement.analysis.kinematics as kin

    displacement_kin = kin.compute_displacement(position)








.. GENERATED FROM PYTHON SOURCE LINES 141-144

The ``displacement`` data array holds, for a given individual and keypoint
at timestep ``t``, the vector that goes from its previous position at time
``t-1`` to its current position at time ``t``.

.. GENERATED FROM PYTHON SOURCE LINES 146-150

And what happens at ``t=0``, since there is no previous timestep?
We define the displacement vector at time ``t=0`` to be the zero vector.
This way the shape of the ``displacement`` data array is the
same as the  ``position`` array:

.. GENERATED FROM PYTHON SOURCE LINES 150-153

.. code-block:: Python

    print(f"Shape of position: {position.shape}")
    print(f"Shape of displacement: {displacement.shape}")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Shape of position: (601, 3, 1, 2)
    Shape of displacement: (601, 3, 1, 2)




.. GENERATED FROM PYTHON SOURCE LINES 154-156

We can visualise these displacement vectors with a quiver plot. In this case
we focus on the mouse ``AEON3B_TP2``:

.. GENERATED FROM PYTHON SOURCE LINES 156-193

.. code-block:: Python

    mouse_name = "AEON3B_TP2"

    fig = plt.figure()
    ax = fig.add_subplot()

    # plot position data
    sc = ax.scatter(
        position.sel(individuals=mouse_name, space="x"),
        position.sel(individuals=mouse_name, space="y"),
        s=15,
        c=position.time,
        cmap="viridis",
    )

    # plot displacement vectors: at t, vector from t-1 to t
    ax.quiver(
        position.sel(individuals=mouse_name, space="x"),
        position.sel(individuals=mouse_name, space="y"),
        displacement.sel(individuals=mouse_name, space="x"),
        displacement.sel(individuals=mouse_name, space="y"),
        angles="xy",
        scale=1,
        scale_units="xy",
        headwidth=7,
        headlength=9,
        headaxislength=9,
    )

    ax.axis("equal")
    ax.set_xlim(450, 575)
    ax.set_ylim(950, 1075)
    ax.set_xlabel("x (pixels)")
    ax.set_ylabel("y (pixels)")
    ax.set_title(f"Zoomed in trajectory of {mouse_name}")
    ax.invert_yaxis()
    fig.colorbar(sc, ax=ax, label="time (s)")




.. image-sg:: /examples/images/sphx_glr_compute_kinematics_004.png
   :alt: Zoomed in trajectory of AEON3B_TP2
   :srcset: /examples/images/sphx_glr_compute_kinematics_004.png
   :class: sphx-glr-single-img


.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    <matplotlib.colorbar.Colorbar object at 0x7fbbf8eaa250>



.. GENERATED FROM PYTHON SOURCE LINES 194-198

Notice that this figure is not very useful as a visual check:
we can see that there are vectors defined for each point in
the trajectory, but we have no easy way to verify they are indeed
the displacement vectors from ``t-1`` to ``t``.

.. GENERATED FROM PYTHON SOURCE LINES 200-207

If instead we plot
the opposite of the displacement vector, we will see that at every time
``t``, the vectors point to the position at ``t-1``.
Remember that the displacement vector is defined as the vector at
time ``t``, that goes from the previous position ``t-1`` to the
current position at ``t``. Therefore, the opposite vector will point
from the position point at ``t``, to the position point at ``t-1``.

.. GENERATED FROM PYTHON SOURCE LINES 209-211

We can easily do this by flipping the sign of the displacement vector in
the plot above:

.. GENERATED FROM PYTHON SOURCE LINES 211-248

.. code-block:: Python

    mouse_name = "AEON3B_TP2"

    fig = plt.figure()
    ax = fig.add_subplot()

    # plot position data
    sc = ax.scatter(
        position.sel(individuals=mouse_name, space="x"),
        position.sel(individuals=mouse_name, space="y"),
        s=15,
        c=position.time,
        cmap="viridis",
    )

    # plot displacement vectors: at t, vector from t-1 to t
    ax.quiver(
        position.sel(individuals=mouse_name, space="x"),
        position.sel(individuals=mouse_name, space="y"),
        -displacement.sel(individuals=mouse_name, space="x"),  # flipped sign
        -displacement.sel(individuals=mouse_name, space="y"),  # flipped sign
        angles="xy",
        scale=1,
        scale_units="xy",
        headwidth=7,
        headlength=9,
        headaxislength=9,
    )
    ax.axis("equal")
    ax.set_xlim(450, 575)
    ax.set_ylim(950, 1075)
    ax.set_xlabel("x (pixels)")
    ax.set_ylabel("y (pixels)")
    ax.set_title(f"Zoomed in trajectory of {mouse_name}")
    ax.invert_yaxis()
    fig.colorbar(sc, ax=ax, label="time (s)")





.. image-sg:: /examples/images/sphx_glr_compute_kinematics_005.png
   :alt: Zoomed in trajectory of AEON3B_TP2
   :srcset: /examples/images/sphx_glr_compute_kinematics_005.png
   :class: sphx-glr-single-img


.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    <matplotlib.colorbar.Colorbar object at 0x7fbbf8d34090>



.. GENERATED FROM PYTHON SOURCE LINES 249-251

Now we can visually verify that indeed the displacement vector
connects the previous and current positions as expected.

.. GENERATED FROM PYTHON SOURCE LINES 253-255

With the displacement data we can compute the distance travelled by the
mouse along its trajectory.

.. GENERATED FROM PYTHON SOURCE LINES 255-269

.. code-block:: Python


    # length of each displacement vector
    displacement_vectors_lengths = compute_norm(
        displacement.sel(individuals=mouse_name)
    )

    # sum the lengths of all displacement vectors (in pixels)
    total_displacement = displacement_vectors_lengths.sum(dim="time").values[0]

    print(
        f"The mouse {mouse_name}'s trajectory is {total_displacement:.2f} "
        "pixels long"
    )





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    The mouse AEON3B_TP2's trajectory is 1640.09 pixels long




.. GENERATED FROM PYTHON SOURCE LINES 270-274

Compute velocity
----------------
We can easily compute the velocity vectors for all individuals in our data
array:

.. GENERATED FROM PYTHON SOURCE LINES 274-276

.. code-block:: Python

    velocity = ds.move.compute_velocity()








.. GENERATED FROM PYTHON SOURCE LINES 277-281

The ``velocity`` method will return a data array equivalent to the
``position`` one, but holding velocity data along the ``space`` axis, rather
than position data. Notice how ``xarray`` nicely deals with the different
individuals and spatial dimensions for us! ✨

.. GENERATED FROM PYTHON SOURCE LINES 283-287

We can plot the components of the velocity vector against time
using ``xarray``'s built-in plotting methods. We use
:meth:`xarray.DataArray.squeeze` to
remove the dimension of length 1 from the data (the ``keypoints`` dimension).

.. GENERATED FROM PYTHON SOURCE LINES 287-291

.. code-block:: Python


    velocity.squeeze().plot.line(x="time", row="individuals", aspect=2, size=2.5)
    plt.gcf().show()




.. image-sg:: /examples/images/sphx_glr_compute_kinematics_006.png
   :alt: individuals = AEON3B_NTP, individuals = AEON3B_TP1, individuals = AEON3B_TP2
   :srcset: /examples/images/sphx_glr_compute_kinematics_006.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 292-298

The components of the velocity vector seem noisier than the components of
the position vector.
This is expected, since we are deriving the velocity using differences in
position (which is somewhat noisy), over small stepsizes.
More specifically, we use numpy's gradient implementation, which
uses second order central differences.

.. GENERATED FROM PYTHON SOURCE LINES 300-302

We can also visualise the speed, as the magnitude (norm)
of the velocity vector:

.. GENERATED FROM PYTHON SOURCE LINES 302-313

.. code-block:: Python

    fig, axes = plt.subplots(3, 1, sharex=True, sharey=True)
    for mouse_name, ax in zip(velocity.individuals.values, axes, strict=False):
        # compute the magnitude of the velocity vector for one mouse
        speed_one_mouse = compute_norm(velocity.sel(individuals=mouse_name))
        # plot speed against time
        ax.plot(speed_one_mouse)
        ax.set_title(mouse_name)
        ax.set_xlabel("time (s)")
        ax.set_ylabel("speed (px/s)")
    fig.tight_layout()




.. image-sg:: /examples/images/sphx_glr_compute_kinematics_007.png
   :alt: AEON3B_NTP, AEON3B_TP1, AEON3B_TP2
   :srcset: /examples/images/sphx_glr_compute_kinematics_007.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 314-316

To visualise the direction of the velocity vector at each timestep, we can
use a quiver plot:

.. GENERATED FROM PYTHON SOURCE LINES 316-346

.. code-block:: Python

    mouse_name = "AEON3B_TP2"
    fig = plt.figure()
    ax = fig.add_subplot()
    # plot trajectory (position data)
    sc = ax.scatter(
        position.sel(individuals=mouse_name, space="x"),
        position.sel(individuals=mouse_name, space="y"),
        s=15,
        c=position.time,
        cmap="viridis",
    )
    # plot velocity vectors
    ax.quiver(
        position.sel(individuals=mouse_name, space="x"),
        position.sel(individuals=mouse_name, space="y"),
        velocity.sel(individuals=mouse_name, space="x"),
        velocity.sel(individuals=mouse_name, space="y"),
        angles="xy",
        scale=2,
        scale_units="xy",
        color="r",
    )
    ax.axis("equal")
    ax.set_xlabel("x (pixels)")
    ax.set_ylabel("y (pixels)")
    ax.set_title(f"Velocity quiver plot for {mouse_name}")
    ax.invert_yaxis()
    fig.colorbar(sc, ax=ax, label="time (s)")
    fig.show()




.. image-sg:: /examples/images/sphx_glr_compute_kinematics_008.png
   :alt: Velocity quiver plot for AEON3B_TP2
   :srcset: /examples/images/sphx_glr_compute_kinematics_008.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 347-349

Here we scaled the length of vectors to half of their actual value
(``scale=2``) for easier visualisation.

.. GENERATED FROM PYTHON SOURCE LINES 351-354

Compute acceleration
---------------------
We can compute the acceleration of the data with an equivalent method:

.. GENERATED FROM PYTHON SOURCE LINES 354-356

.. code-block:: Python

    accel = ds.move.compute_acceleration()








.. GENERATED FROM PYTHON SOURCE LINES 357-359

and plot of the components of the acceleration vector ``ax``, ``ay`` per
individual:

.. GENERATED FROM PYTHON SOURCE LINES 359-377

.. code-block:: Python

    fig, axes = plt.subplots(3, 1, sharex=True, sharey=True)
    for mouse_name, ax in zip(accel.individuals.values, axes, strict=False):
        # plot x-component of acceleration vector
        ax.plot(
            accel.sel(individuals=mouse_name, space=["x"]).squeeze(),
            label="ax",
        )
        # plot y-component of acceleration vector
        ax.plot(
            accel.sel(individuals=mouse_name, space=["y"]).squeeze(),
            label="ay",
        )
        ax.set_title(mouse_name)
        ax.set_xlabel("time (s)")
        ax.set_ylabel("speed (px/s**2)")
        ax.legend(loc="center right", bbox_to_anchor=(1.07, 1.07))
    fig.tight_layout()




.. image-sg:: /examples/images/sphx_glr_compute_kinematics_009.png
   :alt: AEON3B_NTP, AEON3B_TP1, AEON3B_TP2
   :srcset: /examples/images/sphx_glr_compute_kinematics_009.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 378-380

The can also represent the magnitude (norm) of the acceleration vector
for each individual:

.. GENERATED FROM PYTHON SOURCE LINES 380-391

.. code-block:: Python

    fig, axes = plt.subplots(3, 1, sharex=True, sharey=True)
    for mouse_name, ax in zip(accel.individuals.values, axes, strict=False):
        # compute magnitude of the acceleration vector for one mouse
        accel_one_mouse = compute_norm(accel.sel(individuals=mouse_name))

        # plot acceleration against time
        ax.plot(accel_one_mouse)
        ax.set_title(mouse_name)
        ax.set_xlabel("time (s)")
        ax.set_ylabel("accel (px/s**2)")
    fig.tight_layout()



.. image-sg:: /examples/images/sphx_glr_compute_kinematics_010.png
   :alt: AEON3B_NTP, AEON3B_TP1, AEON3B_TP2
   :srcset: /examples/images/sphx_glr_compute_kinematics_010.png
   :class: sphx-glr-single-img






.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 2.379 seconds)


.. _sphx_glr_download_examples_compute_kinematics.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: binder-badge

      .. image:: images/binder_badge_logo.svg
        :target: https://mybinder.org/v2/gh/neuroinformatics-unit/movement/gh-pages?filepath=notebooks/examples/compute_kinematics.ipynb
        :alt: Launch binder
        :width: 150 px

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: compute_kinematics.ipynb <compute_kinematics.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: compute_kinematics.py <compute_kinematics.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: compute_kinematics.zip <compute_kinematics.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_