util

`util`

Functions:

Name	Description
`materialize`	Evaluate `mask` on every index pair, giving the dense boolean array of `shape`, such as (M, N) or (B, M, N).
`pad`	Pad the `np.ndarray` in `Xs` to merge them into a single `np.ndarray`.
`pairwise_distance`	Compute the pairwise squared Euclidean distance between points in `A` (or between points in `A` and `B`).
`topk`	Computes the top-k elements of the tensor along the specified `dim`.
`pairwise_distance_np`	Compute the pairwise squared Euclidean distance between points in `A` (or between points in `A` and `B`).
`dtw`	Computes the Dynamic Time Warping (DTW) distance between two sequences `x` and `y`.
`autocorrelation`	Computes the autocorrelation of a given 1D numpy array at lags below a specified maximum lag.

`materialize`

materialize(mask: Mask, shape: Sequence[SupportsInt]) -> np.ndarray

Evaluate mask on every index pair, giving the dense boolean array of shape, such as (M, N) or (B, M, N).

`pad`

pad(As: list[np.ndarray])

Pad the np.ndarray in Xs to merge them into a single np.ndarray.

Parameters:

Name	Type	Description	Default
`As`	`list[ndarray]`	List of arrays of shape `(N, D_i)`.	required

Returns:

Type	Description
`ndarray`	Single array of shape `(B, N, max(D))` where B is `len(As)`.

Raises:

Type	Description
`ValueError`	- If any array in `As` is not 2D. - If the first dimension of all arrays in `As` are not equal.

`pairwise_distance`

pairwise_distance(A: Tensor, B: Tensor | None = None) -> Tensor

Compute the pairwise squared Euclidean distance between points in A (or between points in A and B).

Parameters:

Name	Type	Description	Default
`A`	`Tensor`	Shape `(N, D)` or `(B, N, D)`. `B` is batch size, `N` is number of points, `D` is dimension of each point.	required
`B`	`Tensor`	Shape `(M, D)` or `(B, M, D)`. `B` is batch size, `M` is number of points, `D` is dimension of each point.	`None`

Returns:

Type	Description
`Tensor`	When `A` is of shape `(N, D)`: shape `(N, N)` [or `(N, M)`] where the element at position `(i, j)` is the squared Euclidean distance between `A[i]` and `A[j]` [or between `A[i]` and `B[j]`]. When `A` is of shape `(B, N, D)`: shape `(B, N, N)` [or `(B, N, M)`] where the element at position `(b, i, j)` is the squared Euclidean distance between `A[b, i]` and `A[b, j]`.

Raises:

Type	Description
`ValueError`	- If `A` is not a 2D or 3D tensor. - If `B` is not `None` and `A` and `B` have different number of dimensions.

`topk`

topk(A: Tensor, k: int, dim: int = -1, largest: bool = True, sorted_: bool = True) -> tuple[Tensor, Tensor]

Computes the top-k elements of the tensor along the specified dim.

Order of indices for equivalent elements is always preserved.

Parameters:

Name	Type	Description	Default
`A`	`Tensor`	Input tensor.	required
`k`	`int`	The number of elements to take, at most `A.shape[dim]`.	required
`dim`	`int`	The dimension along which to take the elements.	`-1`
`largest`	`bool`	Take the largest elements if True, otherwise the smallest.	`True`
`sorted_`	`bool`	Must be True; the elements are returned in order, as `Tensor.topk` requires.	`True`

Returns:

Name	Type	Description
`values`	`Tensor`	The top-k elements, ordered from the most extreme, with `A.shape[dim]` replaced by `k`.
`indices`	`Tensor`	Their indices along `dim`, of the same shape as `values`. Equal elements keep their order.

Raises:

Type	Description
`NotImplementedError`	- If `sorted_` is False.
`ValueError`	- If `k` exceeds `A.shape[dim]`.

`pairwise_distance_np`

pairwise_distance_np(A: np.ndarray, B: np.ndarray | None = None) -> np.ndarray

Compute the pairwise squared Euclidean distance between points in A (or between points in A and B).

Parameters:

Name	Type	Description	Default
`A`	`ndarray`	Shape `(N, D)` or `(B, N, D)`. `B` is batch size, `N` is number of points, `D` is dimension of each point.	required
`B`	`ndarray`	Shape `(M, D)` or `(B, M, D)`. `B` is batch size, `M` is number of points, `D` is dimension of each point.	`None`

Returns:

Type	Description
`ndarray`	When `A` is of shape `(N, D)`: shape `(N, N)` [or `(N, M)`] where the element at position `(i, j)` is the squared Euclidean distance between `A[i]` and `A[j]` [or between `A[i]` and `B[j]`]. When `A` is of shape `(B, N, D)`: shape `(B, N, N)` [or `(B, N, M)`] where the element at position `(b, i, j)` is the squared Euclidean distance between `A[b, i]` and `A[b, j]`.

Raises:

Type	Description
`ValueError`	- If `A` is not a 2D or 3D array. - If `B` is not `None` and `A` and `B` have different number of dimensions.

`dtw`

dtw(A: np.ndarray, B: np.ndarray)

Computes the Dynamic Time Warping (DTW) distance between two sequences x and y.

Parameters:

Name	Type	Description	Default
`A`	`ndarray`	Sequence of shape `(N, D)`.	required
`B`	`ndarray`	Sequence of shape `(M, D)`.	required

Returns:

Name	Type	Description
`distance`	`float`	The DTW distance between the two sequences.

`autocorrelation`

autocorrelation(x: np.ndarray, max_lag: int, step: int = 1)

Computes the autocorrelation of a given 1D numpy array at lags below a specified maximum lag.

Parameters:

Name	Type	Description	Default
`x`	`ndarray`	The input array for which to compute the autocorrelation.	required
`max_lag`	`int`	Exclusive upper bound of the lags, also capped at `len(x)`.	required
`step`	`int`	The step size for the lag. Default is 1.	`1`

Returns:

Type	Description
`ndarray`	Autocorrelation at lags `range(0, min(max_lag, len(x)), step)`.