utils

`DummyProcessPoolExecutor`

Bases: Executor

Dummy ProcessPoolExecutor for to avoid forking for single_job purposes.

Source code in src/dolphin/utils.py

class DummyProcessPoolExecutor(Executor):
    """Dummy ProcessPoolExecutor for to avoid forking for single_job purposes."""

    def __init__(self, max_workers: Optional[int] = None, **kwargs):  # noqa: D107
        self._max_workers = max_workers

    def submit(  # noqa: D102
        self, fn: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs
    ) -> Future[T]:
        future: Future = Future()
        result = fn(*args, **kwargs)
        future.set_result(result)
        return future

    def shutdown(self, wait: bool = True, cancel_futures: bool = True):  # noqa:D102
        pass

    def map(self, fn: Callable[P, T], *iterables, **kwargs):  # noqa: D102
        return map(fn, *iterables)

`compute_out_shape(shape, strides)`

Calculate the output size for an input shape and row/col strides.

Parameters:

Name	Type	Description	Default
`shape`	`tuple[int, int]`	Input size: (rows, cols)	required
`strides`	`tuple[int, int]`	(y strides, x strides)	required

Returns:

Name	Type	Description
`out_shape`	`tuple[int, int]`	Size of output after striding

Notes

If there is not a full window (of size strides), the end will get cut off rather than padded with a partial one. This should match the output size of [dolphin.utils.take_looks][].

As a 1D example, in array of size 6 with strides=3 along this dim, we could expect the pixels to be centered on indexes [1, 4].

[ 0  1  2   3  4  5]

So the output size would be 2, since we have 2 full windows. If the array size was 7 or 8, we would have 2 full windows and 1 partial, so the output size would still be 2.

Source code in src/dolphin/utils.py

def compute_out_shape(
    shape: tuple[int, int], strides: Strides | tuple[int, int]
) -> tuple[int, int]:
    """Calculate the output size for an input `shape` and row/col `strides`.

    Parameters
    ----------
    shape : tuple[int, int]
        Input size: (rows, cols)
    strides : tuple[int, int]
        (y strides, x strides)

    Returns
    -------
    out_shape : tuple[int, int]
        Size of output after striding

    Notes
    -----
    If there is not a full window (of size `strides`), the end
    will get cut off rather than padded with a partial one.
    This should match the output size of `[dolphin.utils.take_looks][]`.

    As a 1D example, in array of size 6 with `strides`=3 along this dim,
    we could expect the pixels to be centered on indexes
    `[1, 4]`.

        [ 0  1  2   3  4  5]

    So the output size would be 2, since we have 2 full windows.
    If the array size was 7 or 8, we would have 2 full windows and 1 partial,
    so the output size would still be 2.

    """
    rows, cols = shape
    rstride, cstride = strides
    return (rows // rstride, cols // cstride)

`disable_gpu()`

Disable GPU usage.

Source code in src/dolphin/utils.py

def disable_gpu():
    """Disable GPU usage."""
    import os

    import jax

    os.environ["CUDA_VISIBLE_DEVICES"] = ""
    jax.config.update("jax_platform_name", "cpu")

`flatten(list_of_lists)`

Flatten one level of a nested iterable.

Source code in src/dolphin/utils.py

def flatten(list_of_lists: Iterable[Iterable[Any]]) -> chain[Any]:
    """Flatten one level of a nested iterable."""
    return chain.from_iterable(list_of_lists)

`format_date_pair(start, end, fmt='%Y%m%d')`

Format a date pair into a string.

Parameters:

Name	Type	Description	Default
`start`	`DateOrDatetime`	First date or datetime	required
`end`	`DateOrDatetime`	Second date or datetime	required
`fmt`	`str`	`datetime` formatter pattern. Default = "%Y%m%d"	`'%Y%m%d'`

Returns:

Type	Description
`str`	Formatted date pair.

Source code in src/dolphin/utils.py

def format_date_pair(
    start: DateOrDatetime, end: DateOrDatetime, fmt: str = "%Y%m%d"
) -> str:
    """Format a date pair into a string.

    Parameters
    ----------
    start : DateOrDatetime
        First date or datetime
    end : DateOrDatetime
        Second date or datetime
    fmt : str, optional
        `datetime` formatter pattern.
        Default = "%Y%m%d"

    Returns
    -------
    str
        Formatted date pair.

    """
    return format_dates(start, end, fmt=fmt, sep="_")

`format_dates(*dates, fmt='%Y%m%d', sep='_')`

Format a date pair into a string.

Parameters:

Name	Type	Description	Default
`*dates`	`DateOrDatetime`	Sequence of date/datetimes to format	`()`
`fmt`	`str`	`datetime` formatter pattern. Default = "%Y%m%d"	`'%Y%m%d'`
`sep`	`str`	string separator between dates. Default = "_"	`'_'`

Returns:

Type	Description
`str`	Formatted date pair.

Source code in src/dolphin/utils.py

def format_dates(*dates: DateOrDatetime, fmt: str = "%Y%m%d", sep: str = "_") -> str:
    """Format a date pair into a string.

    Parameters
    ----------
    *dates : DateOrDatetime
        Sequence of date/datetimes to format
    fmt : str, optional
        `datetime` formatter pattern.
        Default = "%Y%m%d"
    sep : str, optional
        string separator between dates.
        Default = "_"

    Returns
    -------
    str
        Formatted date pair.

    """
    return sep.join((d.strftime(fmt)) for d in dates)

`full_suffix(filename)`

Get the full suffix of a filename, including multiple dots.

Parameters:

Name	Type	Description	Default
`filename`	`str or Path`	path to file	required

Returns:

Type	Description
`str`	The full suffix, including multiple dots.

Examples:

>>> full_suffix('test.tif')
'.tif'
>>> full_suffix('test.tar.gz')
'.tar.gz'

Source code in src/dolphin/utils.py

def full_suffix(filename: Filename):
    """Get the full suffix of a filename, including multiple dots.

    Parameters
    ----------
    filename : str or Path
        path to file

    Returns
    -------
    str
        The full suffix, including multiple dots.

    Examples
    --------
        >>> full_suffix('test.tif')
        '.tif'
        >>> full_suffix('test.tar.gz')
        '.tar.gz'

    """
    fpath = Path(filename)
    return "".join(fpath.suffixes)

`gdal_to_numpy_type(gdal_type)`

Convert gdal type to numpy type.

Source code in src/dolphin/utils.py

def gdal_to_numpy_type(gdal_type: Union[str, int]) -> np.dtype:
    """Convert gdal type to numpy type."""
    if isinstance(gdal_type, str):
        gdal_type = gdal.GetDataTypeByName(gdal_type)
    return np.dtype(gdal_array.GDALTypeCodeToNumericTypeCode(gdal_type))

`get_cpu_count()`

Get the number of CPUs available to the current process.

This function accounts for the possibility of a Docker container with limited CPU resources on a larger machine (which is ignored by multiprocessing.cpu_count()).

Returns:

Type	Description
`int`	The number of CPUs available to the current process.

References

https://github.com/joblib/loky/issues/111
https://github.com/conan-io/conan/blob/982a97041e1ece715d157523e27a14318408b925/conans/client/tools/oss.py#L27 # noqa

Source code in src/dolphin/utils.py

def get_cpu_count():
    """Get the number of CPUs available to the current process.

    This function accounts for the possibility of a Docker container with
    limited CPU resources on a larger machine (which is ignored by
    `multiprocessing.cpu_count()`).

    Returns
    -------
    int
        The number of CPUs available to the current process.

    References
    ----------
    1. https://github.com/joblib/loky/issues/111
    2. https://github.com/conan-io/conan/blob/982a97041e1ece715d157523e27a14318408b925/conans/client/tools/oss.py#L27 # noqa

    """  # noqa: E501

    def get_cpu_quota():
        return int(Path("/sys/fs/cgroup/cpu/cpu.cfs_quota_us").read_text())

    def get_cpu_period():
        return int(Path("/sys/fs/cgroup/cpu/cpu.cfs_period_us").read_text())

    try:
        cfs_quota_us = get_cpu_quota()
        cfs_period_us = get_cpu_period()
        if cfs_quota_us > 0 and cfs_period_us > 0:
            return math.ceil(cfs_quota_us / cfs_period_us)
    except Exception:
        pass
    return cpu_count()

`get_gpu_memory(pid=None, gpu_id=0)`

Get the memory usage (in GiB) of the GPU for the current pid.

Source code in src/dolphin/utils.py

def get_gpu_memory(pid: Optional[int] = None, gpu_id: int = 0) -> float:
    """Get the memory usage (in GiB) of the GPU for the current pid."""
    try:
        from pynvml.smi import nvidia_smi
    except ImportError as e:
        msg = "Please install pynvml through pip or conda"
        raise ImportError(msg) from e

    def get_mem(process):
        used_mem = process["used_memory"] if process else 0
        if process["unit"] == "MiB":
            multiplier = 1 / 1024
        else:
            logger.warning(f"Unknown unit: {process['unit']}")
        return used_mem * multiplier

    nvsmi = nvidia_smi.getInstance()
    processes = nvsmi.DeviceQuery()["gpu"][gpu_id]["processes"]
    if not processes:
        return 0.0

    if pid is None:
        # Return sum of all processes
        return sum(get_mem(p) for p in processes)
    else:
        procs = [p for p in processes if p["pid"] == pid]
        return get_mem(procs[0]) if procs else 0.0

`get_max_memory_usage(units='GB', children=True)`

Get the maximum memory usage of the current process.

Parameters:

Name	Type	Description	Default
`units`	`str, optional, choices=["GB", "MB", "KB", "byte"]`	The units to return, by default "GB".	`'GB'`
`children`	`bool`	Whether to include the memory usage of child processes, by default True	`True`

Returns:

Type	Description
`float`	The maximum memory usage in the specified units.

Raises:

Type	Description
`ValueError`	If the units are not recognized.

References

https://stackoverflow.com/a/7669279/4174466
https://unix.stackexchange.com/a/30941/295194
https://manpages.debian.org/bullseye/manpages-dev/getrusage.2.en.html

Source code in src/dolphin/utils.py

def get_max_memory_usage(units: str = "GB", children: bool = True) -> float:
    """Get the maximum memory usage of the current process.

    Parameters
    ----------
    units : str, optional, choices=["GB", "MB", "KB", "byte"]
        The units to return, by default "GB".
    children : bool, optional
        Whether to include the memory usage of child processes, by default True

    Returns
    -------
    float
        The maximum memory usage in the specified units.

    Raises
    ------
    ValueError
        If the units are not recognized.

    References
    ----------
    1. https://stackoverflow.com/a/7669279/4174466
    2. https://unix.stackexchange.com/a/30941/295194
    3. https://manpages.debian.org/bullseye/manpages-dev/getrusage.2.en.html

    """
    max_mem = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    if children:
        max_mem += resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss
    if units.lower().startswith("g"):
        factor = 1e9
    elif units.lower().startswith("m"):
        factor = 1e6
    elif units.lower().startswith("k"):
        factor = 1e3
    elif units.lower().startswith("byte"):
        factor = 1.0
    else:
        msg = f"Unknown units: {units}"
        raise ValueError(msg)
    if sys.platform.startswith("linux"):
        # on linux, ru_maxrss is in kilobytes, while on mac, ru_maxrss is in bytes
        factor /= 1e3

    return max_mem / factor

`get_nearest_date_idx(input_items, requested, outside_input_range='raise')`

Find the index nearest to requested within input_items.

Source code in src/dolphin/utils.py

def get_nearest_date_idx(
    input_items: Sequence[datetime.datetime],
    requested: datetime.datetime,
    outside_input_range: Literal["allow", "warn", "raise"] = "raise",
) -> int:
    """Find the index nearest to `requested` within `input_items`."""
    sorted_inputs = sorted(input_items)
    if not sorted_inputs[0] <= requested <= sorted_inputs[-1]:
        msg = f"Requested {requested} falls outside of input range: "
        msg += f"{sorted_inputs[0]}, {sorted_inputs[-1]}"
        if outside_input_range == "raise":
            raise ValueError(msg)
        elif outside_input_range == "warn":
            warnings.warn(msg, stacklevel=2)
        else:
            pass

    nearest_idx = min(
        range(len(input_items)),
        key=lambda i: abs((input_items[i] - requested).total_seconds()),
    )

    return nearest_idx

`gpu_is_available()`

Check if a GPU is available.

Source code in src/dolphin/utils.py

def gpu_is_available() -> bool:
    """Check if a GPU is available."""
    # TODO: not sure yet how to check for the jax gpu installation
    try:
        from numba import cuda
        from numba.cuda.cudadrv.error import CudaRuntimeError

    except ImportError:
        logger.debug("numba installed, but GPU not available")
        return False
    try:
        cuda_version = cuda.runtime.get_version()
        logger.debug(f"CUDA version: {cuda_version}")
    except (OSError, CudaRuntimeError) as e:
        logger.debug(f"CUDA runtime version error: {e}")
        return False
    try:
        n_gpu = len(cuda.gpus)
    except cuda.CudaSupportError as e:
        logger.debug(f"CUDA support error {e}")
        return False
    if n_gpu < 1:
        logger.debug("No available GPUs found")
        return False
    return True

`grow_nodata_region(arr, nodata, n_pixels=1, copy=True)`

Grow the nodata region of arr by n_pixels.

This function erodes valid pixels in arr by making a mask from the nodata value and then extends the mask inward by n_pixels.

If arr has no nodata value, the function returns arr unchanged.

Parameters:

Name	Type	Description	Default
`arr`	`ndarray`	Input array containing the data and mask to be eroded	required
`nodata`	`float`	The value in `arr` that represents nodata.	required
`n_pixels`	`int`	Number of pixels to erode from the border Default is 1.	`1`
`copy`	`bool`	Whether to copy the input data before eroding. Default is True (no in-place modification is made).	`True`

Returns:

Type	Description
`ndarray`	Array with the same data, eroded by `n_pixels`. If `copy` is False, the input array is modified in place.

Raises:

Type	Description
`ValueError`	If `arr` is not 2D.

Source code in src/dolphin/utils.py

def grow_nodata_region(
    arr: ArrayLike, nodata: float, n_pixels: int = 1, copy: bool = True
) -> np.ndarray:
    """Grow the `nodata` region of `arr` by  `n_pixels`.

    This function erodes valid pixels in `arr` by making a mask from the `nodata` value
    and then extends the mask inward by `n_pixels`.

    If `arr` has no `nodata` value, the function returns `arr` unchanged.

    Parameters
    ----------
    arr : numpy.ndarray
        Input array containing the data and mask to be eroded
    nodata : float
        The value in `arr` that represents nodata.
    n_pixels : int, optional
        Number of pixels to erode from the border
        Default is 1.
    copy : bool, optional
        Whether to copy the input data before eroding.
        Default is True (no in-place modification is made).

    Returns
    -------
    numpy.ndarray
        Array with the same data, eroded by `n_pixels`.
        If `copy` is False, the input array is modified in place.

    Raises
    ------
    ValueError
        If `arr` is not 2D.

    """
    from scipy import ndimage

    arr = np.asarray(arr)
    if arr.ndim != 2:
        raise ValueError("Input array must be 2D.")

    mask = arr == nodata if not np.isnan(nodata) else np.isnan(arr)
    # "growing" the invalid (nodata) area equivalently "erodes" the valid data
    mask_expanded = ndimage.binary_dilation(
        mask, structure=np.ones((1 + 2 * n_pixels, 1 + 2 * n_pixels))
    )
    out = arr.copy() if copy else arr
    out[mask_expanded] = nodata
    return out

`moving_window_mean(image, size)`

DEPRECATED: use scipy.ndimage.uniform_filter directly.

Source code in src/dolphin/utils.py

def moving_window_mean(
    image: ArrayLike, size: Union[int, tuple[int, int]]
) -> np.ndarray:
    """DEPRECATED: use `scipy.ndimage.uniform_filter` directly."""  # noqa: D401
    from scipy.ndimage import uniform_filter

    msg = (
        "`moving_window_mean` is deprecated. Please use `scipy.ndimage.uniform_filter`."
    )
    warnings.warn(msg, category=DeprecationWarning, stacklevel=2)
    return uniform_filter(image, size=size)

`numpy_to_gdal_type(np_dtype)`

Convert numpy dtype to gdal type.

Parameters:

Name	Type	Description	Default
`np_dtype`	`DTypeLike`	Numpy dtype to convert.	required

Returns:

Type	Description
`int`	GDAL type code corresponding to `np_dtype`.

Raises:

Type	Description
`TypeError`	If `np_dtype` is not a numpy dtype, or if the provided dtype is not supported by GDAL (for example, `np.dtype('>i4')`)

Source code in src/dolphin/utils.py

def numpy_to_gdal_type(np_dtype: DTypeLike) -> int:
    """Convert numpy dtype to gdal type.

    Parameters
    ----------
    np_dtype : DTypeLike
        Numpy dtype to convert.

    Returns
    -------
    int
        GDAL type code corresponding to `np_dtype`.

    Raises
    ------
    TypeError
        If `np_dtype` is not a numpy dtype, or if the provided dtype is not
        supported by GDAL (for example, `np.dtype('>i4')`)

    """
    np_dtype = np.dtype(np_dtype)

    if np.issubdtype(bool, np_dtype):
        return gdalconst.GDT_Byte
    gdal_code = gdal_array.NumericTypeCodeToGDALTypeCode(np_dtype)
    if gdal_code is None:
        msg = f"dtype {np_dtype} not supported by GDAL."
        raise TypeError(msg)
    return gdal_code

`set_num_threads(num_threads)`

Set the cap on threads spawned by numpy and numba.

Uses https://github.com/joblib/threadpoolctl for numpy.

Source code in src/dolphin/utils.py

def set_num_threads(num_threads: int):
    """Set the cap on threads spawned by numpy and numba.

    Uses https://github.com/joblib/threadpoolctl for numpy.
    """
    import os

    import numba
    from threadpoolctl import ThreadpoolController

    # Set the environment variables for the workers
    controller = ThreadpoolController()
    controller.limit(limits=num_threads)
    # https://numba.readthedocs.io/en/stable/user/threading-layer.html#example-of-limiting-the-number-of-threads
    num_cpus = get_cpu_count()
    numba.set_num_threads(min(num_cpus, num_threads))
    # jax setup is harder, for now
    os.environ["XLA_FLAGS"] = f"--xla_force_host_platform_device_count={num_threads}"

`take_looks(arr, row_looks, col_looks, func_type='nansum', edge_strategy='cutoff')`

Downsample a numpy matrix by summing blocks of (row_looks, col_looks).

Parameters:

Name	Type	Description	Default
`arr`	`array`	2D array of an image	required
`row_looks`	`int`	the reduction rate in row direction	required
`col_looks`	`int`	the reduction rate in col direction	required
`func_type`	`str`	the numpy function to use for downsampling, by default "nansum"	`'nansum'`
`edge_strategy`	`str`	how to handle edges, by default "cutoff" Choices are "cutoff", "pad"	`'cutoff'`

Returns:

Type	Description
`ndarray`	The downsampled array, shape = ceil(rows / row_looks, cols / col_looks)

Notes

Cuts off values if the size isn't divisible by num looks.

Source code in src/dolphin/utils.py

def take_looks(arr, row_looks, col_looks, func_type="nansum", edge_strategy="cutoff"):
    """Downsample a numpy matrix by summing blocks of (row_looks, col_looks).

    Parameters
    ----------
    arr : np.array
        2D array of an image
    row_looks : int
        the reduction rate in row direction
    col_looks : int
        the reduction rate in col direction
    func_type : str, optional
        the numpy function to use for downsampling, by default "nansum"
    edge_strategy : str, optional
        how to handle edges, by default "cutoff"
        Choices are "cutoff", "pad"

    Returns
    -------
    ndarray
        The downsampled array, shape = ceil(rows / row_looks, cols / col_looks)

    Notes
    -----
    Cuts off values if the size isn't divisible by num looks.

    """
    if row_looks == 1 and col_looks == 1:
        return arr

    if arr.ndim >= 3:
        return np.stack(
            [
                take_looks(
                    a,
                    row_looks,
                    col_looks,
                    func_type=func_type,
                    edge_strategy=edge_strategy,
                )
                for a in arr
            ]
        )

    if isinstance(arr, np.ma.MaskedArray):
        # Must do looks separately on mask
        # https://github.com/numpy/numpy/issues/8881
        looked_data = take_looks(
            arr.data,
            row_looks,
            col_looks,
            func_type=func_type,
            edge_strategy=edge_strategy,
        )
        if arr.mask.ndim == np.ma.nomask:
            looked_mask = arr.mask
        else:
            looked_mask = take_looks(
                arr.mask,
                row_looks,
                col_looks,
                func_type="any",
                edge_strategy=edge_strategy,
            )
        return np.ma.MaskedArray(data=looked_data, mask=looked_mask)

    arr = _make_dims_multiples(arr, row_looks, col_looks, how=edge_strategy)

    rows, cols = arr.shape
    new_rows = rows // row_looks
    new_cols = cols // col_looks

    func = getattr(np, func_type)
    with warnings.catch_warnings():
        # ignore the warning about nansum of empty slice
        warnings.simplefilter("ignore", category=RuntimeWarning)
        return func(
            np.reshape(arr, (new_rows, row_looks, new_cols, col_looks)), axis=(3, 1)
        )

`upsample_nearest(arr, output_shape, looks=None)`

Upsample a numpy matrix by repeating blocks of (row_looks, col_looks).

Parameters:

Name	Type	Description	Default
`arr`	`array`	2D or 3D downsampled array.	required
`output_shape`	`tuple[int, int]`	The desired output shape.	required
`looks`	`tuple[int, int]`	The number of looks in the row and column directions. If not provided, will be calculated from `output_shape`.	`None`

Returns:

Type	Description
`ndarray`	The upsampled array, shape = `output_shape`.

Source code in src/dolphin/utils.py

def upsample_nearest(
    arr: np.ndarray,
    output_shape: tuple[int, int],
    looks: Optional[tuple[int, int]] = None,
) -> np.ndarray:
    """Upsample a numpy matrix by repeating blocks of (row_looks, col_looks).

    Parameters
    ----------
    arr : np.array
        2D or 3D downsampled array.
    output_shape : tuple[int, int]
        The desired output shape.
    looks : tuple[int, int]
        The number of looks in the row and column directions.
        If not provided, will be calculated from `output_shape`.

    Returns
    -------
    ndarray
        The upsampled array, shape = `output_shape`.

    """
    in_rows, in_cols = arr.shape[-2:]
    out_rows, out_cols = output_shape[-2:]
    if (in_rows, in_cols) == (out_rows, out_cols):
        return arr

    if looks is None:
        row_looks = out_rows // in_rows
        col_looks = out_cols // in_cols
    else:
        row_looks, col_looks = looks

    arr_up = np.repeat(np.repeat(arr, row_looks, axis=-2), col_looks, axis=-1)
    # This may be larger than the original array, or it may be smaller, depending
    # on whether it was padded or cutoff
    out_r = min(out_rows, arr_up.shape[-2])
    out_c = min(out_cols, arr_up.shape[-1])

    shape = (len(arr), out_rows, out_cols) if arr.ndim == 3 else (out_rows, out_cols)
    arr_out = np.zeros(shape=shape, dtype=arr.dtype)
    arr_out[..., :out_r, :out_c] = arr_up[..., :out_r, :out_c]
    return arr_out