From ebae6afed63eadeac696cd2fa33bf0d7ba15f1ec Mon Sep 17 00:00:00 2001
From: niksirbi <niko.sirbiladze@gmail.com>
Date: Fri, 23 Aug 2024 17:12:35 +0100
Subject: [PATCH] wip: implement compute_speed and compute_distance_traveled

---
 movement/analysis/kinematics.py | 67 +++++++++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/movement/analysis/kinematics.py b/movement/analysis/kinematics.py
index 15375a534..7da57c544 100644
--- a/movement/analysis/kinematics.py
+++ b/movement/analysis/kinematics.py
@@ -3,6 +3,7 @@
 import xarray as xr
 
 from movement.utils.logging import log_error
+from movement.utils.vector import compute_norm
 
 
 def compute_displacement(data: xr.DataArray) -> xr.DataArray:
@@ -84,6 +85,72 @@ def compute_acceleration(data: xr.DataArray) -> xr.DataArray:
     return _compute_approximate_time_derivative(data, order=2)
 
 
+def compute_speed(data: xr.DataArray) -> xr.DataArray:
+    """Compute instantaneous speed at each time point.
+
+    Speed is a scalar quantity computed as the Euclidean norm (magnitude)
+    of the velocity vector at each time point.
+
+
+    Parameters
+    ----------
+    data : xarray.DataArray
+        The input data containing position information in Cartesian
+        coordinates, with ``time`` and ``space`` as dimensions.
+
+    Returns
+    -------
+    xarray.DataArray
+        An xarray DataArray containing the computed speed. Will have
+        the same dimensions as the input data, except for ``space``
+        which will be removed.
+
+    """
+    _validate_time_dimension(data)
+    return compute_norm(compute_velocity(data))
+
+
+def compute_distance_travelled(
+    data: xr.DataArray,
+    start: float | None = None,
+    stop: float | None = None,
+) -> xr.DataArray:
+    """Compute the distance travelled between two time points.
+
+    The distance travelled is defined as the cumulative sum of the norms
+    (magnitudes) of the displacement vectors between two time points
+    ``start`` and ``stop``, which should be given in the time units of the
+    data array. If not provided, the minimum and maximum time points
+    in the data array are used.
+
+    Parameters
+    ----------
+    data : xarray.DataArray
+        The input data containing position information in Cartesian
+        coordinates, with ``time`` and ``space`` as dimensions.
+    start : float, optional
+        The start time point to compute the distance travelled from.
+        If None (default), the minimum time point in the data is used.
+    stop : float, optional
+        The stop time point to compute the distance travelled to.
+        If None (default), the maximum time point in the data is used.
+
+    Returns
+    -------
+    xarray.DataArray
+        An xarray DataArray containing the computed distance travelled.
+        Will have the same dimensions as the input data, except for ``time``
+        and ``space`` which will be removed.
+
+    """
+    _validate_time_dimension(data)
+    start = data.time.min() if start is None else start
+    stop = data.time.max() if stop is None else stop
+    data = data.sel(time=slice(start, stop))
+    displacement_norm = compute_norm(compute_displacement(data))
+    return displacement_norm.sum(dim="time")
+
+
 def _compute_approximate_time_derivative(
     data: xr.DataArray, order: int
 ) -> xr.DataArray: