"""Functions for computing kinetic energy."""
import numpy as np
import xarray as xr
from movement.kinematics.kinematics import compute_velocity
from movement.utils.logging import logger
from movement.utils.vector import compute_norm
from movement.validators.arrays import validate_dims_coords
[docs]
def compute_kinetic_energy(
position: xr.DataArray,
keypoints: list | None = None,
masses: dict | None = None,
decompose: bool = False,
) -> xr.DataArray:
r"""Compute kinetic energy per individual.
We consider each individual's set of keypoints (pose) as a classical
system of particles in physics (see Notes).
Parameters
----------
position
The input data containing position information, with ``time``,
``space`` and ``keypoints`` as required dimensions.
keypoints
A list of keypoint names to include in the computation.
By default, all are used.
masses
A dictionary mapping keypoint names to masses, e.g.
{"snout": 1.2, "tail": 0.8}.
By default, unit mass is assumed for all keypoints.
decompose
If True, the kinetic energy is decomposed into "translational" and
"internal" components (see Notes). This requires at least two keypoints
per individual, but more would be desirable for a meaningful
decomposition. The default is False, meaning the total kinetic energy
is returned.
Returns
-------
xarray.DataArray
A data array containing the kinetic energy per individual, for every
time point. Note that the output array lacks ``space`` and
``keypoints`` dimensions.
If ``decompose=True`` an extra ``energy`` dimension is added,
with coordinates ``translational`` and ``internal``.
Notes
-----
Considering a given individual at time point :math:`t` as a system of
keypoint particles, its total kinetic energy :math:`T_{total}` is given by:
.. math:: T_{total} = \sum_{i} \frac{1}{2} m_i \| \mathbf{v}_i(t) \|^2
where :math:`m_i` is the mass of the :math:`i`-th keypoint and
:math:`\mathbf{v}_i(t)` is its velocity at time :math:`t`.
From Samuel König's second theorem, we can decompose :math:`T_{total}`
into:
- Translational kinetic energy: the kinetic energy of the individual's
total mass :math:`M` moving with the centre of mass velocity;
- Internal kinetic energy: the kinetic energy of the keypoints moving
relative to the individual's centre of mass.
We compute translational kinetic energy :math:`T_{trans}` as follows:
.. math:: T_{trans} = \frac{1}{2} M \| \mathbf{v}_{cm}(t) \|^2
where :math:`M = \sum_{i} m_i` is the total mass of the individual
and :math:`\mathbf{v}_{cm}(t) = \frac{1}{M} \sum_{i} m_i \mathbf{v}_i(t)`
is the velocity of the centre of mass at time :math:`t`
(computed as the weighted mean of keypoint velocities).
Internal kinetic energy :math:`T_{int}` is derived as the difference
between the total and translational components:
.. math:: T_{int} = T_{total} - T_{trans}
Examples
--------
>>> from movement.kinematics import compute_kinetic_energy
>>> import numpy as np
>>> import xarray as xr
Compute total kinetic energy:
>>> position = xr.DataArray(
... np.random.rand(3, 2, 4, 2),
... coords={
... "time": np.arange(3),
... "individuals": ["id0", "id1"],
... "keypoints": ["snout", "spine", "tail_base", "tail_tip"],
... "space": ["x", "y"],
... },
... dims=["time", "individuals", "keypoints", "space"],
... )
>>> kinetic_energy_total = compute_kinetic_energy(position)
>>> kinetic_energy_total
<xarray.DataArray (time: 3, individuals: 2)> Size: 48B
0.6579 0.7394 0.1304 0.05152 0.2436 0.5719
Coordinates:
* time (time) int64 24B 0 1 2
* individuals (individuals) <U3 24B 'id0' 'id1'
Compute kinetic energy decomposed into translational
and internal components:
>>> kinetic_energy = compute_kinetic_energy(position, decompose=True)
>>> kinetic_energy
<xarray.DataArray (time: 3, individuals: 2, energy: 2)> Size: 96B
0.0172 1.318 0.02069 0.6498 0.02933 ... 0.1716 0.07829 0.7942 0.06901 0.857
Coordinates:
* time (time) int64 24B 0 1 2
* individuals (individuals) <U3 24B 'id0' 'id1'
* energy (energy) <U13 104B 'translational' 'internal'
Select the 'internal' component:
>>> kinetic_energy_internal = kinetic_energy.sel(energy="internal")
Use unequal keypoint masses and exclude an unreliable keypoint
(e.g. "tail_tip"):
>>> masses = {"snout": 1.2, "spine": 0.8, "tail_base": 1.0}
>>> kinetic_energy = compute_kinetic_energy(
... position,
... keypoints=["snout", "spine", "tail_base"],
... masses=masses,
... decompose=True,
... )
"""
# Validate required dimensions and coordinate labels
validate_dims_coords(
position, {"time": [], "space": ["x", "y"], "keypoints": []}
)
# Subset keypoints if requested
if keypoints is not None:
position = position.sel(keypoints=keypoints)
# Validate that at least 2 keypoints exist for decomposition
if decompose and position.sizes["keypoints"] < 2:
raise logger.error(
ValueError(
"At least 2 keypoints are required to decompose "
"kinetic energy into translational and internal components."
)
)
# Compute velocity from position
velocity = compute_velocity(position)
# Initialise unit weights
weights = xr.DataArray(
np.ones(position.sizes["keypoints"]),
dims=["keypoints"],
coords={"keypoints": position.coords["keypoints"]},
)
# Update weights with keypoint masses, if provided
if masses:
for keypoint, mass in masses.items():
weights.loc[keypoint] = mass
# Compute total KE
weighted_ke = 0.5 * weights * (compute_norm(velocity) ** 2)
ke_total = weighted_ke.sum(dim="keypoints")
if not decompose:
return ke_total
else:
# Compute translational KE based on centre of mass velocity
v_cm = (velocity * weights.expand_dims(space=["x", "y"])).sum(
dim="keypoints"
) / weights.sum()
ke_trans = 0.5 * weights.sum() * compute_norm(v_cm) ** 2
# Internal KE
ke_int = ke_total - ke_trans
# Format output
ke = xr.concat([ke_trans, ke_int], dim="energy")
ke = ke.assign_coords(energy=["translational", "internal"])
ke = ke.transpose("time", ..., "energy")
return ke
