.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/compute_polar_coordinates.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_polar_coordinates.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_polar_coordinates.py:

Express 2D vectors in polar coordinates
============================================

Compute a vector representing head direction and express it in polar
coordinates.

.. GENERATED FROM PYTHON SOURCE LINES 8-10

Imports
-------

.. GENERATED FROM PYTHON SOURCE LINES 10-21

.. code-block:: Python


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

    from movement import sample_data
    from movement.io import load_poses
    from movement.utils.vector import cart2pol, pol2cart








.. GENERATED FROM PYTHON SOURCE LINES 22-26

Load sample dataset
------------------------
In this tutorial, we will use a sample dataset with a single individual
(a mouse) and six keypoints.

.. GENERATED FROM PYTHON SOURCE LINES 26-38

.. code-block:: Python


    ds_path = sample_data.fetch_dataset_paths(
        "SLEAP_single-mouse_EPM.analysis.h5"
    )["poses"]
    ds = load_poses.from_sleap_file(ds_path, fps=None)  # force time_unit = frames

    print(ds)
    print("-----------------------------")
    print(f"Individuals: {ds.individuals.values}")
    print(f"Keypoints: {ds.keypoints.values}")






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

 .. code-block:: none

    <xarray.Dataset> Size: 1MB
    Dimensions:      (time: 18485, individuals: 1, keypoints: 6, space: 2)
    Coordinates:
      * time         (time) int64 148kB 0 1 2 3 4 ... 18480 18481 18482 18483 18484
      * individuals  (individuals) <U12 48B 'individual_0'
      * keypoints    (keypoints) <U9 216B 'snout' 'left_ear' ... 'tail_end'
      * space        (space) <U1 8B 'x' 'y'
    Data variables:
        position     (time, individuals, keypoints, space) float32 887kB nan ... ...
        confidence   (time, individuals, keypoints) float32 444kB nan nan ... 0.7607
    Attributes:
        fps:              None
        time_unit:        frames
        source_software:  SLEAP
        source_file:      /home/runner/.movement/data/poses/SLEAP_single-mouse_EP...
        ds_type:          poses
    -----------------------------
    Individuals: ['individual_0']
    Keypoints: ['snout' 'left_ear' 'right_ear' 'centre' 'tail_base' 'tail_end']




.. GENERATED FROM PYTHON SOURCE LINES 39-42

The loaded dataset ``ds`` contains two data variables:``position`` and
``confidence``. Both are stored as data arrays. In this tutorial, we will
use only ``position``:

.. GENERATED FROM PYTHON SOURCE LINES 42-45

.. code-block:: Python

    position = ds.position









.. GENERATED FROM PYTHON SOURCE LINES 46-52

Compute head vector
---------------------
To demonstrate how polar coordinates can be useful in behavioural analyses,
we will compute the head vector of the mouse.

We define it as the vector from the midpoint between the ears to the snout.

.. GENERATED FROM PYTHON SOURCE LINES 52-66

.. code-block:: Python


    # compute the midpoint between the ears
    midpoint_ears = position.sel(keypoints=["left_ear", "right_ear"]).mean(
        dim="keypoints"
    )

    # compute the head vector
    head_vector = position.sel(keypoints="snout") - midpoint_ears

    # drop the keypoints dimension
    # (otherwise the `head_vector` data array retains a `snout` keypoint from the
    # operation above)
    head_vector = head_vector.drop_vars("keypoints")








.. GENERATED FROM PYTHON SOURCE LINES 67-74

Visualise the head trajectory
--------------------------------------
We can plot the data to check that our computation of the head vector is
correct.

We can start by plotting the trajectory of the midpoint between the ears. We
will refer to this as the head trajectory.

.. GENERATED FROM PYTHON SOURCE LINES 74-95

.. code-block:: Python


    fig, ax = plt.subplots(1, 1)
    mouse_name = ds.individuals.values[0]

    sc = ax.scatter(
        midpoint_ears.sel(individuals=mouse_name, space="x"),
        midpoint_ears.sel(individuals=mouse_name, space="y"),
        s=15,
        c=midpoint_ears.time,
        cmap="viridis",
        marker="o",
    )

    ax.axis("equal")
    ax.set_xlabel("x (pixels)")
    ax.set_ylabel("y (pixels)")
    ax.invert_yaxis()
    ax.set_title(f"Head trajectory ({mouse_name})")
    fig.colorbar(sc, ax=ax, label=f"time ({ds.attrs['time_unit']})")
    fig.show()




.. image-sg:: /examples/images/sphx_glr_compute_polar_coordinates_001.png
   :alt: Head trajectory (individual_0)
   :srcset: /examples/images/sphx_glr_compute_polar_coordinates_001.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 96-101

We can see that the majority of the head trajectory data is within a
cruciform shape. This is because the dataset is of a mouse moving on an
`Elevated Plus Maze <https://en.wikipedia.org/wiki/Elevated_plus_maze>`_.
We can actually verify this is the case by overlaying the head
trajectory on the sample frame of the dataset.

.. GENERATED FROM PYTHON SOURCE LINES 101-133

.. code-block:: Python


    # read sample frame
    frame_path = sample_data.fetch_dataset_paths(
        "SLEAP_single-mouse_EPM.analysis.h5"
    )["frame"]
    im = plt.imread(frame_path)


    # plot sample frame
    fig, ax = plt.subplots(1, 1)
    ax.imshow(im)

    # plot head trajectory with semi-transparent markers
    sc = ax.scatter(
        midpoint_ears.sel(individuals=mouse_name, space="x"),
        midpoint_ears.sel(individuals=mouse_name, space="y"),
        s=15,
        c=midpoint_ears.time,
        cmap="viridis",
        marker="o",
        alpha=0.05,  # transparency
    )

    ax.axis("equal")
    ax.set_xlabel("x (pixels)")
    ax.set_ylabel("y (pixels)")
    # No need to invert the y-axis now, since the image is plotted
    # using a pixel coordinate system with origin on the top left of the image
    ax.set_title(f"Head trajectory ({mouse_name})")

    fig.show()




.. image-sg:: /examples/images/sphx_glr_compute_polar_coordinates_002.png
   :alt: Head trajectory (individual_0)
   :srcset: /examples/images/sphx_glr_compute_polar_coordinates_002.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 134-136

The overlaid plot suggests the mouse spends most of its time in the
covered arms of the maze.

.. GENERATED FROM PYTHON SOURCE LINES 138-143

Visualise the head vector
---------------------------
To visually check our computation of the head vector, it is easier to select
a subset of the data. We can focus on the trajectory of the head when the
mouse is within a small rectangular area and time window.

.. GENERATED FROM PYTHON SOURCE LINES 143-152

.. code-block:: Python


    # area of interest
    xmin, ymin = 600, 665  # pixels
    x_delta, y_delta = 125, 100  # pixels

    # time window
    time_window = range(1650, 1671)  # frames









.. GENERATED FROM PYTHON SOURCE LINES 153-154

For that subset of the data, we now plot the head vector.

.. GENERATED FROM PYTHON SOURCE LINES 154-222

.. code-block:: Python


    fig, ax = plt.subplots(1, 1)
    mouse_name = ds.individuals.values[0]

    # plot midpoint between the ears, and color based on time
    sc = ax.scatter(
        midpoint_ears.sel(individuals=mouse_name, space="x", time=time_window),
        midpoint_ears.sel(individuals=mouse_name, space="y", time=time_window),
        s=50,
        c=midpoint_ears.time[time_window],
        cmap="viridis",
        marker="*",
    )

    # plot snout, and color based on time
    sc = ax.scatter(
        position.sel(
            individuals=mouse_name, space="x", time=time_window, keypoints="snout"
        ),
        position.sel(
            individuals=mouse_name, space="y", time=time_window, keypoints="snout"
        ),
        s=50,
        c=position.time[time_window],
        cmap="viridis",
        marker="o",
    )

    # plot the computed head vector
    ax.quiver(
        midpoint_ears.sel(individuals=mouse_name, space="x", time=time_window),
        midpoint_ears.sel(individuals=mouse_name, space="y", time=time_window),
        head_vector.sel(individuals=mouse_name, space="x", time=time_window),
        head_vector.sel(individuals=mouse_name, space="y", time=time_window),
        angles="xy",
        scale=1,
        scale_units="xy",
        headwidth=7,
        headlength=9,
        headaxislength=9,
        color="gray",
    )

    ax.axis("equal")
    ax.set_xlim(xmin, xmin + x_delta)
    ax.set_ylim(ymin, ymin + y_delta)
    ax.set_xlabel("x (pixels)")
    ax.set_ylabel("y (pixels)")
    ax.set_title(f"Zoomed in head vector ({mouse_name})")
    ax.invert_yaxis()
    fig.colorbar(
        sc,
        ax=ax,
        label=f"time ({ds.attrs['time_unit']})",
        ticks=list(time_window)[0::2],
    )

    ax.legend(
        [
            "midpoint_ears",
            "snout",
            "head_vector",
        ],
        loc="best",
    )

    fig.show()




.. image-sg:: /examples/images/sphx_glr_compute_polar_coordinates_003.png
   :alt: Zoomed in head vector (individual_0)
   :srcset: /examples/images/sphx_glr_compute_polar_coordinates_003.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 223-225

From the plot we can confirm the head vector goes from the midpoint between
the ears to the snout, as we defined it.

.. GENERATED FROM PYTHON SOURCE LINES 228-233

Express the head vector in polar coordinates
-------------------------------------------------------------
A convenient way to inspect the orientation of a vector in 2D is by
expressing it in polar coordinates. We can do this with the vector function
``cart2pol``:

.. GENERATED FROM PYTHON SOURCE LINES 233-237

.. code-block:: Python

    head_vector_polar = cart2pol(head_vector)

    print(head_vector_polar)





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

 .. code-block:: none

    <xarray.DataArray 'position' (time: 18485, individuals: 1, space_pol: 2)> Size: 148kB
    nan nan nan nan nan nan nan nan ... 1.247 18.14 1.126 20.92 1.19 15.59 1.28
    Coordinates:
      * time         (time) int64 148kB 0 1 2 3 4 ... 18480 18481 18482 18483 18484
      * individuals  (individuals) <U12 48B 'individual_0'
      * space_pol    (space_pol) <U3 24B 'rho' 'phi'




.. GENERATED FROM PYTHON SOURCE LINES 238-252

Notice how the resulting array has a ``space_pol`` dimension with two
coordinates: ``rho`` and ``phi``. These are the polar coordinates of the
head vector.

The coordinate ``rho`` is the norm (i.e., magnitude, length) of the vector.
In our case, the distance from the midpoint between the ears to the snout.
The coordinate ``phi`` is the orientation of the head vector relative to the
positive x-axis, and ranges from -``pi`` to ``pi``
(following the `atan2 <https://en.wikipedia.org/wiki/Atan2>`_ convention).

In our coordinate system, ``phi`` will be
positive if the shortest path from the positive x-axis to the vector is
clockwise. Conversely, ``phi`` will be negative if the shortest path from
the positive x-axis to the vector is anti-clockwise.

.. GENERATED FROM PYTHON SOURCE LINES 254-259

Histogram of ``rho`` values
----------------------------
Since ``rho`` is the distance between the ears' midpoint and the snout,
we would expect ``rho`` to be approximately constant in this data. We can
check this by plotting a histogram of its values across the whole clip.

.. GENERATED FROM PYTHON SOURCE LINES 259-293

.. code-block:: Python


    fig, ax = plt.subplots(1, 1)

    # plot histogram using xarray's built-in histogram function
    rho_data = head_vector_polar.sel(individuals=mouse_name, space_pol="rho")
    rho_data.plot.hist(bins=50, ax=ax, edgecolor="lightgray", linewidth=0.5)

    # add mean
    ax.axvline(
        x=rho_data.mean().values,
        c="b",
        linestyle="--",
    )


    # add median
    ax.axvline(
        x=rho_data.median().values,
        c="r",
        linestyle="-",
    )

    # add legend
    ax.legend(
        [
            f"mean = {np.nanmean(rho_data):.2f} pixels",
            f"median = {np.nanmedian(rho_data):.2f} pixels",
        ],
        loc="best",
    )
    ax.set_ylabel("count")
    ax.set_xlabel("rho (pixels)")
    fig.show()




.. image-sg:: /examples/images/sphx_glr_compute_polar_coordinates_004.png
   :alt: individuals = individual_0, space_pol = rho
   :srcset: /examples/images/sphx_glr_compute_polar_coordinates_004.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 294-297

We can see that there is some spread in the value of ``rho`` in this
dataset. This may be due to noise in the detection of the head keypoints,
or due to the mouse tipping its snout upwards during the recording.

.. GENERATED FROM PYTHON SOURCE LINES 299-303

Histogram of ``phi`` values
-------------------------------------
We can also explore which ``phi`` values are most common in the dataset with
a circular histogram.

.. GENERATED FROM PYTHON SOURCE LINES 303-334

.. code-block:: Python



    # compute number of bins
    bin_width_deg = 5  # width of the bins in degrees
    n_bins = int(360 / bin_width_deg)

    # initialise figure with polar projection
    fig = plt.figure()
    ax = fig.add_subplot(projection="polar")

    # plot histogram using xarray's built-in histogram function
    head_vector_polar.sel(individuals=mouse_name, space_pol="phi").plot.hist(
        bins=np.linspace(-np.pi, np.pi, n_bins + 1), ax=ax
    )

    # axes settings
    ax.set_title("phi histogram")
    ax.set_theta_direction(-1)  # theta increases in clockwise direction
    ax.set_theta_offset(0)  # set zero at the right
    ax.set_xlabel("")  # remove default x-label from xarray's plot.hist()

    # set xticks to match the phi values in degrees
    n_xtick_edges = 9
    ax.set_xticks(np.linspace(0, 2 * np.pi, n_xtick_edges)[:-1])
    xticks_in_deg = (
        list(range(0, 180 + 45, 45)) + list(range(0, -180, -45))[-1:0:-1]
    )
    ax.set_xticklabels([str(t) + "\N{DEGREE SIGN}" for t in xticks_in_deg])

    fig.show()




.. image-sg:: /examples/images/sphx_glr_compute_polar_coordinates_005.png
   :alt: phi histogram
   :srcset: /examples/images/sphx_glr_compute_polar_coordinates_005.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 336-338

The ``phi`` circular histogram shows that the head vector appears at a
variety of orientations in this dataset.

.. GENERATED FROM PYTHON SOURCE LINES 340-347

Polar plot of the head vector within a time window
---------------------------------------------------
We can also use a polar plot to represent the head vector in time. This way
we can visualise the head vector in a coordinate system that translates with
the mouse but is always  parallel to the pixel coordinate system.
Again, this will be easier to visualise if we focus on a
small time window.

.. GENERATED FROM PYTHON SOURCE LINES 347-386

.. code-block:: Python


    # select phi values within a time window
    phi = head_vector_polar.sel(
        individuals=mouse_name,
        space_pol="phi",
        time=time_window,
    ).values

    # plot tip of the head vector within that window, and color based on time
    fig = plt.figure()
    ax = fig.add_subplot(projection="polar")
    sc = ax.scatter(
        phi,
        np.ones_like(phi),  # assign a constant value rho=1 for visualization
        c=time_window,
        cmap="viridis",
        s=50,
    )

    # axes settings
    ax.set_theta_direction(-1)  # theta increases in clockwise direction
    ax.set_theta_offset(0)  # set zero at the right
    cax = fig.colorbar(
        sc,
        ax=ax,
        label=f"time ({ds.attrs['time_unit']})",
        ticks=list(time_window)[0::2],
    )

    # set xticks to match the phi values in degrees
    n_xtick_edges = 9
    ax.set_xticks(np.linspace(0, 2 * np.pi, n_xtick_edges)[:-1])
    xticks_in_deg = (
        list(range(0, 180 + 45, 45)) + list(range(0, -180, -45))[-1:0:-1]
    )
    ax.set_xticklabels([str(t) + "\N{DEGREE SIGN}" for t in xticks_in_deg])

    fig.show()




.. image-sg:: /examples/images/sphx_glr_compute_polar_coordinates_006.png
   :alt: compute polar coordinates
   :srcset: /examples/images/sphx_glr_compute_polar_coordinates_006.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 387-392

In the polar plot above, the midpoint between the ears is at the centre of
the plot. The tip of the head vector (the ``snout``) is represented with
color markers at a constant ``rho`` value of 1. Markers are colored by frame.
The polar plot shows how in this small time window of 20 frames,
the head of the mouse turned anti-clockwise.

.. GENERATED FROM PYTHON SOURCE LINES 394-398

Convert polar coordinates to cartesian
------------------------------------------
``movement`` also provides a ``pol2cart`` convenience function to transform
a vector in polar coordinates back to cartesian.

.. GENERATED FROM PYTHON SOURCE LINES 398-402

.. code-block:: Python

    head_vector_cart = pol2cart(head_vector_polar)

    print(head_vector_cart)





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

 .. code-block:: none

    <xarray.DataArray 'position' (time: 18485, individuals: 1, space: 2)> Size: 148kB
    nan nan nan nan nan nan nan nan ... 17.93 7.802 16.38 7.777 19.42 4.465 14.94
    Coordinates:
      * time         (time) int64 148kB 0 1 2 3 4 ... 18480 18481 18482 18483 18484
      * individuals  (individuals) <U12 48B 'individual_0'
      * space        (space) <U1 8B 'x' 'y'




.. GENERATED FROM PYTHON SOURCE LINES 403-405

Note that the resulting `head_vector_cart` array has a ``space`` dimension
with two coordinates: ``x`` and ``y``.


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

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


.. _sphx_glr_download_examples_compute_polar_coordinates.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_polar_coordinates.ipynb
        :alt: Launch binder
        :width: 150 px

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

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

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

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

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

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


.. only:: html

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

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