pycontrails.utils.array

Array utilities.

Functions

searchsorted2d(a, v)

Return the indices where elements in v would be inserted in a along its second axis.
pycontrails.utils.array.searchsorted2d(a, v)

Return the indices where elements in v would be inserted in a along its second axis.

An index of 0 is returned for any rows of a that contain nan values.

Implementation based on a StackOverflow answer.

Parameters:

  • a (npt.NDArray[np.floating]) – 2D array of shape (m, n) that is sorted along its second axis. This is not checked.

  • v (npt.NDArray[np.floating]) – 1D array of values of shape (k,) to insert into the second axis of a. The current implementation could be extended to handle 2D arrays as well.

Returns:

npt.NDArray[np.int64] – 2D array of indices where elements in v would be inserted in a along its second axis to keep the second axis of a sorted. The shape of the output is (m, k).

Examples

>>> a = np.array([
...  [ 1.,  8., 11., 12.],
...  [ 5.,  8.,  9., 14.],
...  [ 4.,  5.,  6., 17.],
...  ])
>>> v = np.array([3., 7., 10., 13., 15.])
>>> searchsorted2d(a, v)
array([[1, 1, 2, 4, 4],
       [0, 1, 3, 3, 4],
       [0, 3, 3, 3, 3]])

>>> a = np.array([
...  [ 1.,  8., np.nan, 12.],
...  [ 5.,  8.,     9., 14.],
...  [ 4.,  5.,     6., 17.],
...  ])
>>> v = np.array([3., 7., 10., 13., 15.])
>>> searchsorted2d(a, v)
array([[0, 0, 0, 0, 0],
       [0, 1, 3, 3, 4],
       [0, 3, 3, 3, 3]])