Indexing

indexing

Functions

find_argmax_nonzero_value(array, keepdims=False, axis=1)

Find the argmax of an array, setting entires with zero sum or finite values to nan

Parameters:

Name	Type	Description	Default
array	ndarray	The input array	required
keepdims	bool	Should the dimensions be kept. Defaults to False.	False
axis	int	Which axis to perform the argmax along. Defaults to 1.	1

Returns:

Type	Description
array	np.array: The argmax, with nans for invalid or all-zero entries

Source code in geograypher/utils/indexing.py

def find_argmax_nonzero_value(
    array: np.ndarray, keepdims: bool = False, axis: int = 1
) -> np.array:
    """Find the argmax of an array, setting entires with zero sum or finite values to nan

    Args:
        array (np.ndarray): The input array
        keepdims (bool, optional): Should the dimensions be kept. Defaults to False.
        axis (int, optional): Which axis to perform the argmax along. Defaults to 1.

    Returns:
        np.array: The argmax, with nans for invalid or all-zero entries
    """
    # Find the column with the highest value per row
    argmax = np.argmax(array, axis=axis, keepdims=keepdims).astype(float)

    # Find rows with zero sum or any infinite values
    zero_sum_mask = np.sum(array, axis=axis) == 0
    infinite_mask = np.any(~np.isfinite(array), axis=axis)

    # Set these rows in the argmax to nan
    argmax[np.logical_or(zero_sum_mask, infinite_mask)] = np.nan

    return argmax