.. 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 9-11

Imports
-------

.. GENERATED FROM PYTHON SOURCE LINES 11-22

.. 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

    import movement.kinematics as kin
    from movement import sample_data
    from movement.plots import plot_centroid_trajectory
    from movement.utils.vector import compute_norm








.. GENERATED FROM PYTHON SOURCE LINES 23-27

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 27-33

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




.. GENERATED FROM PYTHON SOURCE LINES 34-39

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 41-45

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 45-48

.. code-block:: Python

    position = ds.position









.. GENERATED FROM PYTHON SOURCE LINES 49-57

Visualise the data
---------------------------
First, let's visualise the trajectories of the mice in the XY plane,
colouring them by individual.
We use :func:`movement.plots.plot_centroid_trajectory` which is a wrapper
around :func:`matplotlib.pyplot.scatter` that simplifies plotting the
trajectories of individuals in the dataset.
The fig and ax objects returned can be used to further customise the plot.

.. GENERATED FROM PYTHON SOURCE LINES 57-82

.. code-block:: Python


    # Create a single figure and axes
    fig, ax = plt.subplots(1, 1)
    # Invert y-axis so (0,0) is in the top-left,
    # matching typical image coordinate systems
    ax.invert_yaxis()
    # Plot trajectories for each mouse on the same axes
    for mouse_name, col in zip(
        position.individuals.values,
        ["r", "g", "b"],  # colours
        strict=False,
    ):
        plot_centroid_trajectory(
            position,
            individual=mouse_name,
            ax=ax,  # Use the same axes for all plots
            c=col,
            marker="o",
            s=10,
            alpha=0.2,
            label=mouse_name,
        )
        ax.legend().set_alpha(1)
    fig.show()




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





.. GENERATED FROM PYTHON SOURCE LINES 83-87

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 89-91

By default :func:`plot_centroid_trajectory()<movement.plots.\
plot_centroid_trajectory>` colours data points based on their timestamps:

.. GENERATED FROM PYTHON SOURCE LINES 91-107

.. code-block:: Python

    fig, axes = plt.subplots(3, 1, sharey=True)
    for mouse_name, ax in zip(position.individuals.values, axes, strict=False):
        ax.invert_yaxis()
        fig, ax = plot_centroid_trajectory(
            position,
            individual=mouse_name,
            ax=ax,
            s=2,
        )
        ax.set_title(f"Trajectory {mouse_name}")
        ax.set_xlabel("x (pixels)")
        ax.set_ylabel("y (pixels)")
        ax.collections[0].colorbar.set_label("Time (frames)")
    fig.tight_layout()
    fig.show()




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





.. GENERATED FROM PYTHON SOURCE LINES 108-112

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 114-118

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 118-121

.. 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 122-125

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 127-134

Compute displacement
---------------------
The :mod:`movement.kinematics` module
provides functions to compute various kinematic quantities,
such as displacement, velocity, and acceleration.
We can start off by computing the distance travelled by the mice along
their trajectories:

.. GENERATED FROM PYTHON SOURCE LINES 134-137

.. code-block:: Python


    displacement = kin.compute_displacement(position)








.. GENERATED FROM PYTHON SOURCE LINES 138-146

The :func:`movement.kinematics.compute_displacement`
function will return a data array equivalent to the ``position`` one,
but holding displacement data along the ``space`` axis, rather than
position data.

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 148-152

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 152-155

.. 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, 2, 1, 3)
    Shape of displacement: (601, 2, 1, 3)




.. GENERATED FROM PYTHON SOURCE LINES 156-158

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 158-195

.. 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 0x7fa4b949fe00>



.. GENERATED FROM PYTHON SOURCE LINES 196-200

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 202-209

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 211-213

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

.. GENERATED FROM PYTHON SOURCE LINES 213-250

.. 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 0x7fa4b8a8e030>



.. GENERATED FROM PYTHON SOURCE LINES 251-253

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

.. GENERATED FROM PYTHON SOURCE LINES 255-257

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

.. GENERATED FROM PYTHON SOURCE LINES 257-271

.. 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 272-276

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

.. GENERATED FROM PYTHON SOURCE LINES 276-278

.. code-block:: Python

    velocity = kin.compute_velocity(position)








.. GENERATED FROM PYTHON SOURCE LINES 279-284

The :func:`movement.kinematics.compute_velocity`
function 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 286-290

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 290-294

.. 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 295-301

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 303-305

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

.. GENERATED FROM PYTHON SOURCE LINES 305-316

.. 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 317-319

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

.. GENERATED FROM PYTHON SOURCE LINES 319-349

.. 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 350-352

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

.. GENERATED FROM PYTHON SOURCE LINES 354-358

Compute acceleration
---------------------
Let's now compute the acceleration for all individuals in our data
array:

.. GENERATED FROM PYTHON SOURCE LINES 358-360

.. code-block:: Python

    accel = kin.compute_acceleration(position)








.. GENERATED FROM PYTHON SOURCE LINES 361-363

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

.. GENERATED FROM PYTHON SOURCE LINES 363-381

.. 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 382-384

We can also compute and visualise the magnitude (norm) of the
acceleration vector for each individual:

.. GENERATED FROM PYTHON SOURCE LINES 384-395

.. 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.070 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>`_