normSq

ssapy.utils.normSq(arr)[source][source]

Compute the squared norm of an array over the last axis while preserving leading axes.

This function calculates the squared norm of vectors along the last axis of the input array using Einstein summation notation (np.einsum). It is designed to be portable and efficient, offering better performance than np.linalg.norm() for this specific use case.

Parameters:

arrarray-like: Input array where the squared norm is computed along the last axis. The array can have any shape, with the last axis representing the vector components.

Returns:

norm_squaredarray-like: The squared norm of the input array computed along the last axis. The output shape matches the input shape, excluding the last axis.

Notes:

This implementation is more portable and faster than using np.linalg.norm() for computing norms over the last axis, especially for large arrays.
It is slightly less specialized compared to methods optimized for fixed shapes (e.g., shape (3,)), but it generalizes well to arrays of arbitrary shape.
The computation uses np.einsum(”…i,…i”, arr, arr), which efficiently performs the summation of squared components along the last axis.

Example:

>>> import numpy as np
>>> arr = np.array([[1, 2, 3], [4, 5, 6]])
>>> normSq(arr)
array([14, 77])