Egocentric coordinate transformation

[talmo]

Is your feature request related to a problem? Please describe. For many analyses, it's super nice to have a way to align the animal such that the centroid is at the origin and the head is aligned to the x-axis. This is sometimes referred to as "egocentrization" since it places all the points in the animal-centric coordinate frame.

This is useful for:

• Gait analysis, since stride segmentation and gait-relative kinematic features are done in the egocentric space (e.g., Pereira et al. (2019), DeAngelis et al. (2020), Hession et al. (2023))
• Coarsely estimating visual inputs from 2D information (e.g., Walter & Couzin (2021))
• Preprocessing for downstream clustering/behavior segmentation (e.g., Weinreb et al. (2024))

Describe the solution you'd like

Here's the nicely vectorized implementation we've used for a while:

def signed_angle(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """Finds the signed angle between two 2D vectors a and b.

    Args:
        a: Array of shape (n, 2).
        b: Array of shape (n, 2).

    Returns:
        The signed angles in degrees in vector of shape (n, 2).

        This angle is positive if a is rotated clockwise to align to b and negative if
        this rotation is counter-clockwise.
    """
    a = a / np.linalg.norm(a, axis=1, keepdims=True)
    b = b / np.linalg.norm(b, axis=1, keepdims=True)
    theta = np.arccos(np.around(np.sum(a * b, axis=1), decimals=4))
    cross = np.cross(a, b, axis=1)
    sign = np.zeros(cross.shape)
    sign[cross >= 0] = -1
    sign[cross < 0] = 1
    return np.rad2deg(theta) * sign

def normalize_to_egocentric(
    x: np.ndarray,
    rel_to: np.ndarray | None = None,
    ctr_ind: int = 0,
    fwd_ind: int = 1,
    return_angles: bool = False,
):
    """Normalize pose estimates to egocentric coordinates.

    Args:
        x: Pose of shape (joints, 2) or (time, joints, 2)
        rel_to: Pose to align x with of shape (joints, 2) or (time, joints, 2). Defaults
            to x if not specified. Useful for aligning other individuals relative to the
            centered animal.
        ctr_ind: Index of centroid joint. Defaults to 0.
        fwd_ind: Index of "forward" joint (e.g., head). Defaults to 1.
        return_angles: If True, return angles with the aligned coordinates.

    Returns:
        Egocentrically aligned poses of the same shape as the input.

        If return_angles is True, also returns a vector of angles.
    """

    if rel_to is None:
        rel_to = x

    is_singleton = (x.ndim == 2) and (rel_to.ndim == 2)

    if x.ndim == 2:
        x = np.expand_dims(x, axis=0)
    if rel_to.ndim == 2:
        rel_to = np.expand_dims(rel_to, axis=0)

    # Find egocentric forward coordinates.
    ctr = rel_to[..., ctr_ind, :]  # (t, 2)
    fwd = rel_to[..., fwd_ind, :]  # (t, 2)
    ego_fwd = fwd - ctr

    # Compute angle.
    ang = np.arctan2(
        ego_fwd[..., 1], ego_fwd[..., 0]
    )  # arctan2(y, x) -> radians in [-pi, pi]
    ca = np.cos(ang)  # (t,)
    sa = np.sin(ang)  # (t,)

    # Build rotation matrix.
    rot = np.zeros([len(ca), 3, 3], dtype=ca.dtype)
    rot[..., 0, 0] = ca
    rot[..., 0, 1] = -sa
    rot[..., 1, 0] = sa
    rot[..., 1, 1] = ca
    rot[..., 2, 2] = 1

    # Center.
    x = x - np.expand_dims(ctr, axis=1)

    # Pad, rotate and crop.
    x = np.pad(x, ((0, 0), (0, 0), (0, 1)), "constant", constant_values=1) @ rot
    x = x[..., :2]

    if is_singleton:
        x = x[0]

    if return_angles:
        return x, ang
    else:
        return x

Describe alternatives you've considered We've done it ourselves :)

Additional context I know there's been some discussion on APIs and etc., so I don't want to step on your style guide. Given some guidance, happy to shoot a PR or have a student work on it :)

[niksirbi]

Thanks @talmo for such a thorough description. This is indeed a very useful feature and we've been planning to add it from the project's conception.

Your sketched out implementation makes sense to me conceptually, and it could fit in well into the codebase, with a few modifications. The functions would have to accept and return xarray.DataArray objects (i.e. data variables) instead of numpy arrays, similar to the existing functions in our kinematics and filtering modules.

The best place for the main function would probably be a new transforms.py module, as we are likely to add more coordinate transforms in future. I envision something along the lines of from movement.transforms import to_egocentric. Utilities like signed_angle can go into utils/vector.py.

I'm happy for you or someone from your team to have a shot at opening a draft pull request to get this started. If the need arises, I'd be also willing to hop on a call to explain our thoughts on the architecture. If our ongoing API changes get in the way of merging, we can always step in and help with that.