python_embed/Lib/site-packages/mne/utils/numerics.py

"""Some utility functions."""

# Authors: The MNE-Python contributors.
# License: BSD-3-Clause
# Copyright the MNE-Python contributors.

import inspect
import numbers
import operator
import os
import shutil
import sys
from contextlib import contextmanager
from datetime import date, datetime, timedelta, timezone
from io import BytesIO, StringIO
from math import ceil, sqrt
from pathlib import Path

import numpy as np
from scipy import sparse

from ..fixes import (
    _infer_dimension_,
    _safe_svd,
    has_numba,
    jit,
    stable_cumsum,
    svd_flip,
)
from ._logging import logger, verbose, warn
from .check import (
    _check_pandas_installed,
    _ensure_int,
    _validate_type,
    check_random_state,
)
from .docs import fill_doc
from .misc import _empty_hash, _pl


def split_list(v, n, idx=False):
    """Split list in n (approx) equal pieces, possibly giving indices."""
    n = int(n)
    tot = len(v)
    sz = tot // n
    start = stop = 0
    for i in range(n - 1):
        stop += sz
        yield (np.arange(start, stop), v[start:stop]) if idx else v[start:stop]
        start += sz
    yield (np.arange(start, tot), v[start:]) if idx else v[start]


def array_split_idx(ary, indices_or_sections, axis=0, n_per_split=1):
    """Do what numpy.array_split does, but add indices."""
    # this only works for indices_or_sections as int
    indices_or_sections = _ensure_int(indices_or_sections)
    ary_split = np.array_split(ary, indices_or_sections, axis=axis)
    idx_split = np.array_split(np.arange(ary.shape[axis]), indices_or_sections)
    idx_split = (
        np.arange(sp[0] * n_per_split, (sp[-1] + 1) * n_per_split) for sp in idx_split
    )
    return zip(idx_split, ary_split)


def sum_squared(X):
    """Compute norm of an array.

    Parameters
    ----------
    X : array
        Data whose norm must be found.

    Returns
    -------
    value : float
        Sum of squares of the input array X.
    """
    X_flat = X.ravel(order="F" if np.isfortran(X) else "C")
    return np.dot(X_flat, X_flat)


def _compute_row_norms(data):
    """Compute scaling based on estimated norm."""
    norms = np.sqrt(np.sum(data**2, axis=1))
    norms[norms == 0] = 1.0
    return norms


def _reg_pinv(x, reg=0, rank="full", rcond=1e-15):
    """Compute a regularized pseudoinverse of Hermitian matrices.

    Regularization is performed by adding a constant value to each diagonal
    element of the matrix before inversion. This is known as "diagonal
    loading". The loading factor is computed as ``reg * np.trace(x) / len(x)``.

    The pseudo-inverse is computed through SVD decomposition and inverting the
    singular values. When the matrix is rank deficient, some singular values
    will be close to zero and will not be used during the inversion. The number
    of singular values to use can either be manually specified or automatically
    estimated.

    Parameters
    ----------
    x : ndarray, shape (..., n, n)
        Square, Hermitian matrices to invert.
    reg : float
        Regularization parameter. Defaults to 0.
    rank : int | None | 'full'
        This controls the effective rank of the covariance matrix when
        computing the inverse. The rank can be set explicitly by specifying an
        integer value. If ``None``, the rank will be automatically estimated.
        Since applying regularization will always make the covariance matrix
        full rank, the rank is estimated before regularization in this case. If
        'full', the rank will be estimated after regularization and hence
        will mean using the full rank, unless ``reg=0`` is used.
        Defaults to 'full'.
    rcond : float | 'auto'
        Cutoff for detecting small singular values when attempting to estimate
        the rank of the matrix (``rank='auto'``). Singular values smaller than
        the cutoff are set to zero. When set to 'auto', a cutoff based on
        floating point precision will be used. Defaults to 1e-15.

    Returns
    -------
    x_inv : ndarray, shape (..., n, n)
        The inverted matrix.
    loading_factor : float
        Value added to the diagonal of the matrix during regularization.
    rank : int
        If ``rank`` was set to an integer value, this value is returned,
        else the estimated rank of the matrix, before regularization, is
        returned.
    """
    from ..rank import _estimate_rank_from_s

    if rank is not None and rank != "full":
        rank = int(operator.index(rank))
    if x.ndim < 2 or x.shape[-2] != x.shape[-1]:
        raise ValueError("Input matrix must be square.")
    if not np.allclose(x, x.conj().swapaxes(-2, -1)):
        raise ValueError("Input matrix must be Hermitian (symmetric)")
    assert x.ndim >= 2 and x.shape[-2] == x.shape[-1]
    n = x.shape[-1]

    # Decompose the matrix, not necessarily positive semidefinite
    U, s, Vh = np.linalg.svd(x, hermitian=True)

    # Estimate the rank before regularization
    tol = "auto" if rcond == "auto" else rcond * s[..., :1]
    rank_before = _estimate_rank_from_s(s, tol)

    # Decompose the matrix again after regularization
    loading_factor = reg * np.mean(s, axis=-1)
    if reg:
        U, s, Vh = np.linalg.svd(
            x + loading_factor[..., np.newaxis, np.newaxis] * np.eye(n), hermitian=True
        )

    # Estimate the rank after regularization
    tol = "auto" if rcond == "auto" else rcond * s[..., :1]
    rank_after = _estimate_rank_from_s(s, tol)

    # Warn the user if both all parameters were kept at their defaults and the
    # matrix is rank deficient.
    if (rank_after < n).any() and reg == 0 and rank == "full" and rcond == 1e-15:
        warn("Covariance matrix is rank-deficient and no regularization is done.")
    elif isinstance(rank, int) and rank > n:
        raise ValueError(
            f"Invalid value for the rank parameter ({rank}) given "
            f"the shape of the input matrix ({x.shape[0]} x {x.shape[1]})."
        )

    # Pick the requested number of singular values
    mask = np.arange(s.shape[-1]).reshape((1,) * (x.ndim - 2) + (-1,))
    if rank is None:
        cmp = ret = rank_before
    elif rank == "full":
        cmp = rank_after
        ret = rank_before
    else:
        cmp = ret = rank
    mask = mask < np.asarray(cmp)[..., np.newaxis]
    mask &= s > 0

    # Invert only non-zero singular values
    s_inv = np.zeros(s.shape)
    s_inv[mask] = 1.0 / s[mask]

    # Compute the pseudo inverse
    x_inv = np.matmul(U * s_inv[..., np.newaxis, :], Vh)

    return x_inv, loading_factor, ret


def _gen_events(n_epochs):
    """Generate event structure from number of epochs."""
    events = np.c_[np.arange(n_epochs), np.zeros(n_epochs, int), np.ones(n_epochs, int)]
    return events


def _reject_data_segments(data, reject, flat, decim, info, tstep):
    """Reject data segments using peak-to-peak amplitude."""
    from .._fiff.pick import channel_indices_by_type
    from ..epochs import _is_good

    data_clean = np.empty_like(data)
    idx_by_type = channel_indices_by_type(info)
    step = int(ceil(tstep * info["sfreq"]))
    if decim is not None:
        step = int(ceil(step / float(decim)))
    this_start = 0
    this_stop = 0
    drop_inds = []
    for first in range(0, data.shape[1], step):
        last = first + step
        data_buffer = data[:, first:last]
        if data_buffer.shape[1] < (last - first):
            break  # end of the time segment
        if _is_good(
            data_buffer,
            info["ch_names"],
            idx_by_type,
            reject,
            flat,
            ignore_chs=info["bads"],
        ):
            this_stop = this_start + data_buffer.shape[1]
            data_clean[:, this_start:this_stop] = data_buffer
            this_start += data_buffer.shape[1]
        else:
            logger.info(f"Artifact detected in [{first}, {last}]")
            drop_inds.append((first, last))
    data = data_clean[:, :this_stop]
    if not data.any():
        raise RuntimeError(
            "No clean segment found. Please "
            "consider updating your rejection "
            "thresholds."
        )
    return data, drop_inds


def _get_inst_data(inst):
    """Get data view from MNE object instance like Raw, Epochs or Evoked."""
    from ..epochs import BaseEpochs
    from ..evoked import Evoked
    from ..io import BaseRaw
    from ..time_frequency.tfr import BaseTFR

    _validate_type(inst, (BaseRaw, BaseEpochs, Evoked, BaseTFR), "Instance")
    if not inst.preload:
        inst.load_data()
    return inst._data


def compute_corr(x, y):
    """Compute pearson correlations between a vector and a matrix."""
    if len(x) == 0 or len(y) == 0:
        raise ValueError("x or y has zero length")
    X = np.array(x, float)
    Y = np.array(y, float)
    X -= X.mean(0)
    Y -= Y.mean(0)
    x_sd = X.std(0, ddof=1)
    # if covariance matrix is fully expanded, Y needs a
    # transpose / broadcasting else Y is correct
    y_sd = Y.std(0, ddof=1)[:, None if X.shape == Y.shape else Ellipsis]
    return (np.dot(X.T, Y) / float(len(X) - 1)) / (x_sd * y_sd)


@fill_doc
def random_permutation(n_samples, random_state=None):
    """Emulate the randperm matlab function.

    It returns a vector containing a random permutation of the
    integers between 0 and n_samples-1. It returns the same random numbers
    than randperm matlab function whenever the random_state is the same
    as the matlab's random seed.

    This function is useful for comparing against matlab scripts
    which use the randperm function.

    Note: the randperm(n_samples) matlab function generates a random
    sequence between 1 and n_samples, whereas
    random_permutation(n_samples, random_state) function generates
    a random sequence between 0 and n_samples-1, that is:
    randperm(n_samples) = random_permutation(n_samples, random_state) - 1

    Parameters
    ----------
    n_samples : int
        End point of the sequence to be permuted (excluded, i.e., the end point
        is equal to n_samples-1)
    %(random_state)s

    Returns
    -------
    randperm : ndarray, int
        Randomly permuted sequence between 0 and n-1.
    """
    rng = check_random_state(random_state)
    # This can't just be rng.permutation(n_samples) because it's not identical
    # to what MATLAB produces
    idx = rng.uniform(size=n_samples)
    randperm = np.argsort(idx)
    return randperm


@verbose
def _apply_scaling_array(data, picks_list, scalings, verbose=None):
    """Scale data type-dependently for estimation."""
    scalings = _check_scaling_inputs(data, picks_list, scalings)
    if isinstance(scalings, dict):
        logger.debug(f"    Scaling using mapping {scalings}.")
        picks_dict = dict(picks_list)
        scalings = [(picks_dict[k], v) for k, v in scalings.items() if k in picks_dict]
        for idx, scaling in scalings:
            data[idx, :] *= scaling  # F - order
    else:
        logger.debug("    Scaling using computed norms.")
        data *= scalings[:, np.newaxis]  # F - order


def _invert_scalings(scalings):
    if isinstance(scalings, dict):
        scalings = {k: 1.0 / v for k, v in scalings.items()}
    elif isinstance(scalings, np.ndarray):
        scalings = 1.0 / scalings
    return scalings


def _undo_scaling_array(data, picks_list, scalings):
    scalings = _invert_scalings(_check_scaling_inputs(data, picks_list, scalings))
    return _apply_scaling_array(data, picks_list, scalings, verbose=False)


@contextmanager
def _scaled_array(data, picks_list, scalings):
    """Scale, use, unscale array."""
    _apply_scaling_array(data, picks_list=picks_list, scalings=scalings)
    try:
        yield
    finally:
        _undo_scaling_array(data, picks_list=picks_list, scalings=scalings)


def _apply_scaling_cov(data, picks_list, scalings):
    """Scale resulting data after estimation."""
    scalings = _check_scaling_inputs(data, picks_list, scalings)
    scales = None
    if isinstance(scalings, dict):
        n_channels = len(data)
        covinds = list(zip(*picks_list))[1]
        assert len(data) == sum(len(k) for k in covinds)
        assert list(sorted(np.concatenate(covinds))) == list(range(len(data)))
        scales = np.zeros(n_channels)
        for ch_t, idx in picks_list:
            scales[idx] = scalings[ch_t]
    elif isinstance(scalings, np.ndarray):
        if len(scalings) != len(data):
            raise ValueError("Scaling factors and data are of incompatible shape")
        scales = scalings
    elif scalings is None:
        pass
    else:
        raise RuntimeError("Arff...")
    if scales is not None:
        assert np.sum(scales == 0.0) == 0
        data *= scales[None, :] * scales[:, None]


def _undo_scaling_cov(data, picks_list, scalings):
    scalings = _invert_scalings(_check_scaling_inputs(data, picks_list, scalings))
    return _apply_scaling_cov(data, picks_list, scalings)


def _check_scaling_inputs(data, picks_list, scalings):
    """Aux function."""
    rescale_dict_ = dict(mag=1e15, grad=1e13, eeg=1e6)

    scalings_ = None
    if isinstance(scalings, str) and scalings == "norm":
        scalings_ = 1.0 / _compute_row_norms(data)
    elif isinstance(scalings, dict):
        rescale_dict_.update(scalings)
        scalings_ = rescale_dict_
    elif isinstance(scalings, np.ndarray):
        scalings_ = scalings
    elif scalings is None:
        pass
    else:
        raise NotImplementedError(f"Not a valid rescaling option: {scalings}")
    return scalings_


def hashfunc(fname, block_size=1048576, hash_type="md5"):  # 2 ** 20
    """Calculate the hash for a file.

    Parameters
    ----------
    fname : str
        Filename.
    block_size : int
        Block size to use when reading.

    Returns
    -------
    hash_ : str
        The hexadecimal digest of the hash.
    """
    hasher = _empty_hash(kind=hash_type)
    with open(fname, "rb") as fid:
        while True:
            data = fid.read(block_size)
            if not data:
                break
            hasher.update(data)
    return hasher.hexdigest()


def create_slices(start, stop, step=None, length=1):
    """Generate slices of time indexes.

    Parameters
    ----------
    start : int
        Index where first slice should start.
    stop : int
        Index where last slice should maximally end.
    length : int
        Number of time sample included in a given slice.
    step: int | None
        Number of time samples separating two slices.
        If step = None, step = length.

    Returns
    -------
    slices : list
        List of slice objects.
    """
    # default parameters
    if step is None:
        step = length

    # slicing
    slices = [slice(t, t + length, 1) for t in range(start, stop - length + 1, step)]
    return slices


def _time_mask(
    times, tmin=None, tmax=None, sfreq=None, raise_error=True, include_tmax=True
):
    """Safely find sample boundaries."""
    orig_tmin = tmin
    orig_tmax = tmax
    tmin = -np.inf if tmin is None else tmin
    tmax = np.inf if tmax is None else tmax
    if not np.isfinite(tmin):
        tmin = times[0]
    if not np.isfinite(tmax):
        tmax = times[-1]
        include_tmax = True  # ignore this param when tmax is infinite
    if sfreq is not None:
        # Push to a bit past the nearest sample boundary first
        sfreq = float(sfreq)
        tmin = int(round(tmin * sfreq)) / sfreq - 0.5 / sfreq
        tmax = int(round(tmax * sfreq)) / sfreq
        tmax += (0.5 if include_tmax else -0.5) / sfreq
    else:
        assert include_tmax  # can only be used when sfreq is known
    if raise_error and tmin > tmax:
        raise ValueError(
            f"tmin ({orig_tmin}) must be less than or equal to tmax ({orig_tmax})"
        )
    mask = times >= tmin
    mask &= times <= tmax
    if raise_error and not mask.any():
        extra = "" if include_tmax else "when include_tmax=False "
        raise ValueError(
            f"No samples remain when using tmin={orig_tmin} and tmax={orig_tmax} "
            f"{extra}(original time bounds are [{times[0]}, {times[-1]}] containing "
            f"{len(times)} sample{_pl(times)})"
        )
    return mask


def _freq_mask(freqs, sfreq, fmin=None, fmax=None, raise_error=True):
    """Safely find frequency boundaries."""
    orig_fmin = fmin
    orig_fmax = fmax
    fmin = -np.inf if fmin is None else fmin
    fmax = np.inf if fmax is None else fmax
    if not np.isfinite(fmin):
        fmin = freqs[0]
    if not np.isfinite(fmax):
        fmax = freqs[-1]
    if sfreq is None:
        raise ValueError("sfreq can not be None")
    # Push 0.5/sfreq past the nearest frequency boundary first
    sfreq = float(sfreq)
    fmin = int(round(fmin * sfreq)) / sfreq - 0.5 / sfreq
    fmax = int(round(fmax * sfreq)) / sfreq + 0.5 / sfreq
    if raise_error and fmin > fmax:
        raise ValueError(
            f"fmin ({orig_fmin}) must be less than or equal to fmax ({orig_fmax})"
        )
    mask = freqs >= fmin
    mask &= freqs <= fmax
    if raise_error and not mask.any():
        raise ValueError(
            f"No frequencies remain when using fmin={orig_fmin} and fmax={orig_fmax} "
            f"(original frequency bounds are [{freqs[0]}, {freqs[-1]}])"
        )
    return mask


def grand_average(all_inst, interpolate_bads=True, drop_bads=True):
    """Make grand average of a list of Evoked, AverageTFR, or Spectrum data.

    For :class:`mne.Evoked` data, the function interpolates bad channels based on the
    ``interpolate_bads`` parameter. If ``interpolate_bads`` is True, the grand average
    file will contain good channels and the bad channels interpolated from the good
    MEG/EEG channels.
    For :class:`mne.time_frequency.AverageTFR` and :class:`mne.time_frequency.Spectrum`
    data, the function takes the subset of channels not marked as bad in any of the
    instances.

    The ``grand_average.nave`` attribute will be equal to the number of datasets used to
    calculate the grand average.

    .. note:: A grand average evoked should not be used for source localization.

    Parameters
    ----------
    all_inst : list of Evoked, AverageTFR or Spectrum
        The datasets.

        .. versionchanged:: 1.10.0
            Added support for :class:`~mne.time_frequency.Spectrum` objects.

    interpolate_bads : bool
        If True, bad MEG and EEG channels are interpolated. Ignored for
        :class:`~mne.time_frequency.AverageTFR` and
        :class:`~mne.time_frequency.Spectrum` data.
    drop_bads : bool
        If True, drop all bad channels marked as bad in any data set. If neither
        ``interpolate_bads`` nor ``drop_bads`` is `True`, in the output file, every
        channel marked as bad in at least one of the input files will be marked as bad,
        but no interpolation or dropping will be performed.

    Returns
    -------
    grand_average : Evoked | AverageTFR | Spectrum
        The grand average data. Same type as input.

    Notes
    -----
    Aggregating multitaper TFR datasets with a taper dimension such as for complex or
    phase data is not supported.

    .. versionadded:: 0.11.0
    """
    # check if all elements in the given list are evoked data
    from ..channels.channels import equalize_channels
    from ..evoked import Evoked
    from ..time_frequency import AverageTFR, Spectrum

    if not all_inst:
        raise ValueError(
            "Please pass a list of Evoked, AverageTFR, or Spectrum objects."
        )
    elif len(all_inst) == 1:
        warn("Only a single dataset was passed to mne.grand_average().")

    inst_type = type(all_inst[0])
    _validate_type(all_inst[0], (Evoked, AverageTFR, Spectrum), "All elements")
    for inst in all_inst:
        _validate_type(inst, inst_type, "All elements", "of the same type")

    # Copy channels to leave the original evoked datasets intact.
    all_inst = [inst.copy() for inst in all_inst]

    # Interpolates if necessary
    if isinstance(all_inst[0], Evoked):
        if interpolate_bads:
            all_inst = [
                inst.interpolate_bads() if len(inst.info["bads"]) > 0 else inst
                for inst in all_inst
            ]
        from ..evoked import combine_evoked as combine
    elif isinstance(all_inst[0], Spectrum):
        from ..time_frequency.spectrum import combine_spectrum as combine
    else:  # isinstance(all_inst[0], AverageTFR):
        from ..time_frequency.tfr import combine_tfr as combine

    if drop_bads:
        bads = list({b for inst in all_inst for b in inst.info["bads"]})
        if bads:
            for inst in all_inst:
                inst.drop_channels(bads)

    equalize_channels(all_inst, copy=False)
    # make grand_average object using combine_[evoked/tfr/spectrum]
    grand_average = combine(all_inst, weights="equal")
    # change the grand_average.nave to the number of datasets
    grand_average.nave = len(all_inst)
    # change comment field
    grand_average.comment = f"Grand average (n = {grand_average.nave})"
    return grand_average


class _HashableNdarray(np.ndarray):
    def __hash__(self):
        return object_hash(self)

    def __eq__(self, other):
        return NotImplementedError  # defer to hash


def _hashable_ndarray(x):
    return x.view(_HashableNdarray)


def object_hash(x, h=None):
    """Hash a reasonable python object.

    Parameters
    ----------
    x : object
        Object to hash. Can be anything comprised of nested versions of:
        {dict, list, tuple, ndarray, str, bytes, float, int, None}.
    h : hashlib HASH object | None
        Optional, object to add the hash to. None creates an MD5 hash.

    Returns
    -------
    digest : int
        The digest resulting from the hash.
    """
    if h is None:
        h = _empty_hash()
    if hasattr(x, "keys"):
        # dict-like types
        keys = _sort_keys(x)
        for key in keys:
            object_hash(key, h)
            object_hash(x[key], h)
    elif isinstance(x, bytes):
        # must come before "str" below
        h.update(x)
    elif isinstance(x, str | float | int | type(None)):
        h.update(str(type(x)).encode("utf-8"))
        h.update(str(x).encode("utf-8"))
    elif isinstance(x, np.ndarray | np.number | np.bool_):
        x = np.asarray(x)
        h.update(str(x.shape).encode("utf-8"))
        h.update(str(x.dtype).encode("utf-8"))
        h.update(x.tobytes())
    elif isinstance(x, datetime):
        object_hash(_dt_to_stamp(x))
    elif sparse.issparse(x):
        h.update(str(type(x)).encode("utf-8"))
        if not isinstance(x, sparse.csr_array | sparse.csc_array):
            raise RuntimeError(f"Unsupported sparse type {type(x)}")
        h.update(x.data.tobytes())
        h.update(x.indices.tobytes())
        h.update(x.indptr.tobytes())
    elif hasattr(x, "__len__"):
        # all other list-like types
        h.update(str(type(x)).encode("utf-8"))
        for xx in x:
            object_hash(xx, h)
    else:
        raise RuntimeError(f"unsupported type: {type(x)} ({x})")
    return int(h.hexdigest(), 16)


def object_size(x, memo=None):
    """Estimate the size of a reasonable python object.

    Parameters
    ----------
    x : object
        Object to approximate the size of.
        Can be anything comprised of nested versions of:
        {dict, list, tuple, ndarray, str, bytes, float, int, None}.
    memo : dict | None
        The memodict.

    Returns
    -------
    size : int
        The estimated size in bytes of the object.
    """
    # Note: this will not process object arrays properly (since those only)
    # hold references
    if memo is None:
        memo = dict()
    id_ = id(x)
    if id_ in memo:
        return 0  # do not add already existing ones
    if isinstance(x, bytes | str | int | float | type(None) | Path):
        size = sys.getsizeof(x)
    elif isinstance(x, np.ndarray):
        # On newer versions of NumPy, just doing sys.getsizeof(x) works,
        # but on older ones you always get something small :(
        size = sys.getsizeof(np.array([]))
        if x.base is None or id(x.base) not in memo:
            size += x.nbytes
    elif isinstance(x, np.generic):
        size = x.nbytes
    elif isinstance(x, dict):
        size = sys.getsizeof(x)
        for key, value in x.items():
            size += object_size(key, memo)
            size += object_size(value, memo)
    elif isinstance(x, list | tuple):
        size = sys.getsizeof(x) + sum(object_size(xx, memo) for xx in x)
    elif isinstance(x, datetime):
        size = object_size(_dt_to_stamp(x), memo)
    elif isinstance(x, date):
        size = 24  # 3 8-byte integers
    elif _is_sparse_cs(x):
        size = sum(sys.getsizeof(xx) for xx in [x, x.data, x.indices, x.indptr])
    else:
        raise RuntimeError(f"unsupported type: {type(x)} ({x})")
    memo[id_] = size
    return size


def _is_sparse_cs(x):
    return isinstance(
        x, sparse.csr_matrix | sparse.csc_matrix | sparse.csr_array | sparse.csc_array
    )


def _sort_keys(x):
    """Sort and return keys of dict."""
    keys = list(x.keys())  # note: not thread-safe
    idx = np.argsort([str(k) for k in keys])
    keys = [keys[ii] for ii in idx]
    return keys


def _array_equal_nan(a, b, allclose=False):
    try:
        if allclose:
            func = np.testing.assert_allclose
        else:
            func = np.testing.assert_array_equal
        func(a, b)
    except AssertionError:
        return False
    return True


def object_diff(a, b, pre="", *, allclose=False):
    """Compute all differences between two python variables.

    Parameters
    ----------
    a : object
        Currently supported: class, dict, list, tuple, ndarray,
        int, str, bytes, float, StringIO, BytesIO.
    b : object
        Must be same type as ``a``.
    pre : str
        String to prepend to each line.
    allclose : bool
        If True (default False), use assert_allclose.

    Returns
    -------
    diffs : str
        A string representation of the differences.
    """
    pd = _check_pandas_installed(strict=False)
    out = ""
    if type(a) is not type(b):
        # Deal with NamedInt and NamedFloat
        for sub in (int, float):
            if isinstance(a, sub) and isinstance(b, sub):
                break
        else:
            return f"{pre} type mismatch ({type(a)}, {type(b)})\n"
    if inspect.isclass(a):
        if inspect.isclass(b) and a != b:
            return f"{pre} class mismatch ({a}, {b})\n"
    elif isinstance(a, dict):
        k1s = _sort_keys(a)
        k2s = _sort_keys(b)
        m1 = set(k2s) - set(k1s)
        if len(m1):
            out += pre + f" left missing keys {m1}\n"
        for key in k1s:
            if key not in k2s:
                out += pre + f" right missing key {key}\n"
            else:
                out += object_diff(
                    a[key], b[key], pre=(pre + f"[{repr(key)}]"), allclose=allclose
                )
    elif isinstance(a, list | tuple):
        if len(a) != len(b):
            out += pre + f" length mismatch ({len(a)}, {len(b)})\n"
        else:
            for ii, (xx1, xx2) in enumerate(zip(a, b)):
                out += object_diff(xx1, xx2, pre + f"[{ii}]", allclose=allclose)
    elif isinstance(a, float):
        if not _array_equal_nan(a, b, allclose):
            out += pre + f" value mismatch ({a}, {b})\n"
    elif isinstance(a, str | int | bytes | np.generic):
        if a != b:
            out += pre + f" value mismatch ({a}, {b})\n"
    elif a is None:
        if b is not None:
            out += pre + f" left is None, right is not ({b})\n"
    elif isinstance(a, np.ndarray):
        if not _array_equal_nan(a, b, allclose):
            out += pre + " array mismatch\n"
    elif isinstance(a, StringIO | BytesIO):
        if a.getvalue() != b.getvalue():
            out += pre + " StringIO mismatch\n"
    elif isinstance(a, datetime | date):
        ts = (a - b).total_seconds()
        if ts != 0:
            out += pre + f" {a.__class__.__name__} mismatch ({a} vs {b} by {ts} sec)\n"
    elif sparse.issparse(a):
        # sparsity and sparse type of b vs a already checked above by type()
        if b.shape != a.shape:
            out += pre + (
                f" sparse matrix a and b shape mismatch ({a.shape} vs {b.shape})"
            )
        else:
            c = a - b
            c.eliminate_zeros()
            if c.nnz > 0:
                out += pre + (f" sparse matrix a and b differ on {c.nnz} elements")
    elif pd and isinstance(a, pd.DataFrame):
        try:
            pd.testing.assert_frame_equal(a, b)
        except AssertionError:
            out += pre + " DataFrame mismatch\n"
    elif hasattr(a, "__getstate__") and a.__getstate__() is not None:
        out += object_diff(a.__getstate__(), b.__getstate__(), pre, allclose=allclose)
    else:
        raise RuntimeError(pre + f": unsupported type {type(a)} ({a})")
    return out


class _PCA:
    """Principal component analysis (PCA)."""

    # Adapted from sklearn and stripped down to just use linalg.svd
    # and make it easier to later provide a "center" option if we want

    def __init__(self, n_components=None, whiten=False):
        self.n_components = n_components
        self.whiten = whiten

    def fit_transform(self, X, y=None):
        X = X.copy()
        U, S, _ = self._fit(X)
        U = U[:, : self.n_components_]

        if self.whiten:
            # X_new = X * V / S * sqrt(n_samples) = U * sqrt(n_samples)
            U *= sqrt(X.shape[0] - 1)
        else:
            # X_new = X * V = U * S * V^T * V = U * S
            U *= S[: self.n_components_]

        return U

    def fit(self, X):
        self._fit(X)

    def _fit(self, X):
        if self.n_components is None:
            n_components = min(X.shape)
        else:
            n_components = self.n_components
        n_samples, n_features = X.shape

        if n_components == "mle":
            if n_samples < n_features:
                raise ValueError(
                    "n_components='mle' is only supported if n_samples >= n_features"
                )
        elif not 0 <= n_components <= min(n_samples, n_features):
            raise ValueError(
                f"n_components={repr(n_components)} must be between 0 and "
                f"min(n_samples, n_features)={repr(min(n_samples, n_features))} with "
                "svd_solver='full'"
            )
        elif n_components >= 1:
            if not isinstance(n_components, numbers.Integral | np.integer):
                raise ValueError(
                    f"n_components={repr(n_components)} must be of type int "
                    f"when greater than or equal to 1, "
                    f"was of type={repr(type(n_components))}"
                )

        self.mean_ = np.mean(X, axis=0)
        X -= self.mean_

        U, S, V = _safe_svd(X, full_matrices=False)
        # flip eigenvectors' sign to enforce deterministic output
        U, V = svd_flip(U, V)

        components_ = V

        # Get variance explained by singular values
        explained_variance_ = (S**2) / (n_samples - 1)
        total_var = explained_variance_.sum()
        explained_variance_ratio_ = explained_variance_ / total_var
        singular_values_ = S.copy()  # Store the singular values.

        # Postprocess the number of components required
        if n_components == "mle":
            n_components = _infer_dimension_(explained_variance_, n_samples, n_features)
        elif 0 < n_components < 1.0:
            # number of components for which the cumulated explained
            # variance percentage is superior to the desired threshold
            ratio_cumsum = stable_cumsum(explained_variance_ratio_)
            n_components = np.searchsorted(ratio_cumsum, n_components) + 1

        # Compute noise covariance using Probabilistic PCA model
        # The sigma2 maximum likelihood (cf. eq. 12.46)
        if n_components < min(n_features, n_samples):
            self.noise_variance_ = explained_variance_[n_components:].mean()
        else:
            self.noise_variance_ = 0.0

        self.n_samples_, self.n_features_ = n_samples, n_features
        self.components_ = components_[:n_components]
        self.n_components_ = n_components
        self.explained_variance_ = explained_variance_[:n_components]
        self.explained_variance_ratio_ = explained_variance_ratio_[:n_components]
        self.singular_values_ = singular_values_[:n_components]

        return U, S, V


def _mask_to_onsets_offsets(mask):
    """Group boolean mask into contiguous onset:offset pairs."""
    assert mask.dtype == np.dtype(bool) and mask.ndim == 1
    mask = mask.astype(int)
    diff = np.diff(mask)
    onsets = np.where(diff > 0)[0] + 1
    if mask[0]:
        onsets = np.concatenate([[0], onsets])
    offsets = np.where(diff < 0)[0] + 1
    if mask[-1]:
        offsets = np.concatenate([offsets, [len(mask)]])
    assert len(onsets) == len(offsets)
    return onsets, offsets


def _julian_to_date(jd):
    """Convert Julian integer to a date object.

    Parameters
    ----------
    jd : int
        Julian date - number of days since julian day 0
        Julian day number 0 assigned to the day starting at
        noon on January 1, 4713 BC, proleptic Julian calendar
        November 24, 4714 BC, in the proleptic Gregorian calendar

    Returns
    -------
    jd_date : datetime
        Datetime representation of jd

    """
    # https://aa.usno.navy.mil/data/docs/JulianDate.php
    # Thursday, A.D. 1970 Jan 1 12:00:00.0  2440588.000000
    jd_t0 = 2440588
    datetime_t0 = datetime(1970, 1, 1, 12, 0, 0, 0, tzinfo=timezone.utc)

    dt = timedelta(days=(jd - jd_t0))
    return (datetime_t0 + dt).date()


def _date_to_julian(jd_date):
    """Convert datetime object to a Julian integer.

    Parameters
    ----------
    jd_date : date

    Returns
    -------
    jd : float
        Julian date corresponding to jd_date
        - number of days since julian day 0
        Julian day number 0 assigned to the day starting at
        noon on January 1, 4713 BC, proleptic Julian calendar
        November 24, 4714 BC, in the proleptic Gregorian calendar

    """
    # https://aa.usno.navy.mil/data/docs/JulianDate.php
    # Thursday, A.D. 1970 Jan 1 12:00:00.0  2440588.000000
    jd_t0 = 2440588
    date_t0 = date(1970, 1, 1)
    dt = jd_date - date_t0
    return jd_t0 + dt.days


def _check_dt(dt):
    if (
        not isinstance(dt, datetime)
        or dt.tzinfo is None
        or dt.tzinfo is not timezone.utc
    ):
        raise ValueError(f"Date must be datetime object in UTC: {repr(dt)}")


def _dt_to_stamp(inp_date):
    """Convert a datetime object to a timestamp."""
    _check_dt(inp_date)
    return int(inp_date.timestamp() // 1), inp_date.microsecond


def _stamp_to_dt(utc_stamp):
    """Convert timestamp to datetime object in Windows-friendly way."""
    # The min on windows is 86400
    stamp = [int(s) for s in utc_stamp]
    if len(stamp) == 1:  # In case there is no microseconds information
        stamp.append(0)
    return datetime.fromtimestamp(0, tz=timezone.utc) + timedelta(
        seconds=stamp[0], microseconds=stamp[1]
    )


class _ReuseCycle:
    """Cycle over a variable, preferring to reuse earlier indices.

    Requires the values in ``x`` to be hashable and unique. This holds
    nicely for matplotlib's color cycle, which gives HTML hex color strings.
    """

    def __init__(self, x):
        self.indices = list()
        self.popped = dict()
        assert len(x) > 0
        self.x = x

    def __iter__(self):
        while True:
            yield self.__next__()

    def __next__(self):
        if not len(self.indices):
            self.indices = list(range(len(self.x)))
            self.popped = dict()
        idx = self.indices.pop(0)
        val = self.x[idx]
        assert val not in self.popped
        self.popped[val] = idx
        return val

    def restore(self, val):
        try:
            idx = self.popped.pop(val)
        except KeyError:
            warn(f"Could not find value: {val}")
        else:
            loc = np.searchsorted(self.indices, idx)
            self.indices.insert(loc, idx)


def _arange_div_fallback(n, d):
    x = np.arange(n, dtype=np.float64)
    x /= d
    return x


if has_numba:

    @jit(fastmath=False)
    def _arange_div(n, d):
        out = np.empty(n, np.float64)
        for i in range(n):
            out[i] = i / d
        return out

else:  # pragma: no cover
    _arange_div = _arange_div_fallback


_LRU_CACHES = dict()
_LRU_CACHE_MAXSIZES = dict()


def _custom_lru_cache(maxsize):
    def dec(fun):
        fun_hash = hash(fun)
        this_cache = _LRU_CACHES[fun_hash] = dict()
        _LRU_CACHE_MAXSIZES[fun_hash] = maxsize

        def cache_fun(*args):
            hash_ = object_hash(args)
            if hash_ in this_cache:
                this_val = this_cache.pop(hash_)
            else:
                this_val = fun(*args)
            this_cache[hash_] = this_val  # (re)insert in last pos
            while len(this_cache) > _LRU_CACHE_MAXSIZES[fun_hash]:
                for key in this_cache:  # just an easy way to get first element
                    this_cache.pop(key)
                    break  # first in, first out
            return this_val

        return cache_fun

    return dec


def _array_repr(x):
    """Produce compact info about float ndarray x."""
    assert isinstance(x, np.ndarray), type(x)
    return f"shape : {x.shape}, range : [{np.nanmin(x):+0.2g}, {np.nanmax(x):+0.2g}]"


def _replace_md5(fname):
    """Replace a file based on MD5sum."""
    # adapted from sphinx-gallery
    assert fname.endswith(".new")
    fname_old = fname[:-4]
    if os.path.isfile(fname_old) and hashfunc(fname) == hashfunc(fname_old):
        os.remove(fname)
    else:
        shutil.move(fname, fname_old)