Aluode/PerceptionLabPortable

1"""Testing functions."""2 3# Authors: The MNE-Python contributors.4# License: BSD-3-Clause5# Copyright the MNE-Python contributors.6 7import inspect8import os9import sys10import tempfile11import traceback12from functools import wraps13from shutil import rmtree14from unittest import SkipTest15 16import numpy as np17from numpy.testing import assert_allclose, assert_array_equal18from scipy import linalg19 20from ._logging import ClosingStringIO, warn21from .check import check_version22from .misc import run_subprocess23from .numerics import object_diff24 25 26def _explain_exception(start=-1, stop=None, prefix="> "):27    """Explain an exception."""28    # start=-1 means "only the most recent caller"29    etype, value, tb = sys.exc_info()30    string = traceback.format_list(traceback.extract_tb(tb)[start:stop])31    string = "".join(string).split("\n") + traceback.format_exception_only(etype, value)32    string = ":\n" + prefix + ("\n" + prefix).join(string)33    return string34 35 36class _TempDir(str):37    """Create and auto-destroy temp dir.38 39    This is designed to be used with testing modules. Instances should be40    defined inside test functions. Instances defined at module level can not41    guarantee proper destruction of the temporary directory.42 43    When used at module level, the current use of the __del__() method for44    cleanup can fail because the rmtree function may be cleaned up before this45    object (an alternative could be using the atexit module instead).46    """47 48    def __new__(self):  # noqa: D10549        new = str.__new__(self, tempfile.mkdtemp(prefix="tmp_mne_tempdir_"))50        return new51 52    def __init__(self):53        self._path = self.__str__()54 55    def __del__(self):  # noqa: D10556        rmtree(self._path, ignore_errors=True)57 58 59def requires_mne(func):60    """Decorate a function as requiring MNE."""61    return requires_mne_mark()(func)62 63 64def requires_mne_mark():65    """Mark pytest tests that require MNE-C."""66    import pytest67 68    return pytest.mark.skipif(not has_mne_c(), reason="Requires MNE-C")69 70 71def requires_openmeeg_mark():72    """Mark pytest tests that require OpenMEEG."""73    import pytest74 75    return pytest.mark.skipif(76        not check_version("openmeeg", "2.5.6"), reason="Requires OpenMEEG >= 2.5.6"77    )78 79 80def requires_freesurfer(arg):81    """Require Freesurfer."""82    import pytest83 84    reason = "Requires Freesurfer"85    if isinstance(arg, str):86        # Calling as  @requires_freesurfer('progname'): return decorator87        # after checking for progname existence88        reason += f" command: {arg}"89        try:90            run_subprocess([arg, "--version"])91        except Exception:92            skip = True93        else:94            skip = False95        return pytest.mark.skipif(skip, reason=reason)96    else:97        # Calling directly as @requires_freesurfer: return decorated function98        # and just check env var existence99        return pytest.mark.skipif(not has_freesurfer(), reason="Requires Freesurfer")(100            arg101        )102 103 104def requires_good_network(func):105    import pytest106 107    return pytest.mark.skipif(108        int(os.environ.get("MNE_SKIP_NETWORK_TESTS", 0)),109        reason="MNE_SKIP_NETWORK_TESTS is set",110    )(func)111 112 113def run_command_if_main():114    """Run a given command if it's __main__."""115    local_vars = inspect.currentframe().f_back.f_locals116    if local_vars.get("__name__", "") == "__main__":117        local_vars["run"]()118 119 120class ArgvSetter:121    """Temporarily set sys.argv."""122 123    def __init__(self, args=(), disable_stdout=True, disable_stderr=True):124        self.argv = list(("python",) + args)125        self.stdout = ClosingStringIO() if disable_stdout else sys.stdout126        self.stderr = ClosingStringIO() if disable_stderr else sys.stderr127 128    def __enter__(self):  # noqa: D105129        self.orig_argv = sys.argv130        sys.argv = self.argv131        self.orig_stdout = sys.stdout132        sys.stdout = self.stdout133        self.orig_stderr = sys.stderr134        sys.stderr = self.stderr135        return self136 137    def __exit__(self, *args):  # noqa: D105138        sys.argv = self.orig_argv139        sys.stdout = self.orig_stdout140        sys.stderr = self.orig_stderr141 142 143def has_mne_c():144    """Check for MNE-C."""145    return "MNE_ROOT" in os.environ146 147 148def has_freesurfer():149    """Check for Freesurfer."""150    return "FREESURFER_HOME" in os.environ151 152 153def buggy_mkl_svd(function):154    """Decorate tests that make calls to SVD and intermittently fail."""155 156    @wraps(function)157    def dec(*args, **kwargs):158        try:159            return function(*args, **kwargs)160        except np.linalg.LinAlgError as exp:161            if "SVD did not converge" in str(exp):162                msg = "Intel MKL SVD convergence error detected, skipping test"163                warn(msg)164                raise SkipTest(msg)165            raise166 167    return dec168 169 170def assert_and_remove_boundary_annot(annotations, n=1):171    """Assert that there are boundary annotations and remove them."""172    __tracebackhide__ = True173 174    from ..io import BaseRaw175 176    if isinstance(annotations, BaseRaw):  # allow either input177        annotations = annotations.annotations178    for key in ("EDGE", "BAD"):179        idx = np.where(annotations.description == f"{key} boundary")[0]180        assert len(idx) == n, (181            f"Got {len(idx)} '{key} boundary' annotations, expected {n}"182        )183        annotations.delete(idx)184 185 186def assert_object_equal(a, b, *, err_msg="Object mismatch", allclose=False):187    """Assert two objects are equal."""188    __tracebackhide__ = True189 190    d = object_diff(a, b, allclose=allclose)191    assert d == "", f"{err_msg}\n{d}"192 193 194def _raw_annot(meas_date, orig_time):195    from .._fiff.meas_info import create_info196    from ..annotations import Annotations, _handle_meas_date197    from ..io import RawArray198 199    info = create_info(ch_names=10, sfreq=10.0)200    raw = RawArray(data=np.empty((10, 10)), info=info, first_samp=10)201    if meas_date is not None:202        meas_date = _handle_meas_date(meas_date)203    with raw.info._unlock(check_after=True):204        raw.info["meas_date"] = meas_date205    annot = Annotations([0.5], [0.2], ["dummy"], orig_time)206    raw.set_annotations(annotations=annot)207    return raw208 209 210def _get_data(x, ch_idx):211    """Get the (n_ch, n_times) data array."""212    from ..evoked import Evoked213    from ..io import BaseRaw214 215    if isinstance(x, BaseRaw):216        return x[ch_idx][0]217    elif isinstance(x, Evoked):218        return x.data[ch_idx]219 220 221def _check_snr(actual, desired, picks, min_tol, med_tol, msg, kind="MEG"):222    """Check the SNR of a set of channels."""223    __tracebackhide__ = True224 225    actual_data = _get_data(actual, picks)226    desired_data = _get_data(desired, picks)227    bench_rms = np.sqrt(np.mean(desired_data * desired_data, axis=1))228    error = actual_data - desired_data229    error_rms = np.sqrt(np.mean(error * error, axis=1))230    np.clip(error_rms, 1e-60, np.inf, out=error_rms)  # avoid division by zero231    snrs = bench_rms / error_rms232    # min tol233    snr = snrs.min()234    bad_count = (snrs < min_tol).sum()235    msg = f" ({msg})" if msg != "" else msg236    assert bad_count == 0, (237        f"SNR (worst {snr:0.2f}) < {min_tol:0.2f} "238        f"for {bad_count}/{len(picks)} channels{msg}"239    )240    # median tol241    snr = np.median(snrs)242    assert snr >= med_tol, f"{kind} SNR median {snr:0.2f} < {med_tol:0.2f}{msg}"243 244 245def assert_meg_snr(246    actual, desired, min_tol, med_tol=500.0, chpi_med_tol=500.0, msg=None247):248    """Assert channel SNR of a certain level.249 250    Mostly useful for operations like Maxwell filtering that modify251    MEG channels while leaving EEG and others intact.252    """253    __tracebackhide__ = True254 255    from .._fiff.pick import pick_types256 257    picks = pick_types(desired.info, meg=True, exclude=[])258    picks_desired = pick_types(desired.info, meg=True, exclude=[])259    assert_array_equal(picks, picks_desired, err_msg="MEG pick mismatch")260    chpis = pick_types(actual.info, meg=False, chpi=True, exclude=[])261    chpis_desired = pick_types(desired.info, meg=False, chpi=True, exclude=[])262    if chpi_med_tol is not None:263        assert_array_equal(chpis, chpis_desired, err_msg="cHPI pick mismatch")264    others = np.setdiff1d(265        np.arange(len(actual.ch_names)), np.concatenate([picks, chpis])266    )267    others_desired = np.setdiff1d(268        np.arange(len(desired.ch_names)), np.concatenate([picks_desired, chpis_desired])269    )270    assert_array_equal(others, others_desired, err_msg="Other pick mismatch")271    if len(others) > 0:  # if non-MEG channels present272        assert_allclose(273            _get_data(actual, others),274            _get_data(desired, others),275            atol=1e-11,276            rtol=1e-5,277            err_msg="non-MEG channel mismatch",278        )279    _check_snr(actual, desired, picks, min_tol, med_tol, msg, kind="MEG")280    if chpi_med_tol is not None and len(chpis) > 0:281        _check_snr(actual, desired, chpis, 0.0, chpi_med_tol, msg, kind="cHPI")282 283 284def assert_snr(actual, desired, tol):285    """Assert actual and desired arrays are within some SNR tolerance."""286    __tracebackhide__ = True287 288    with np.errstate(divide="ignore"):  # allow infinite289        snr = linalg.norm(desired, ord="fro") / linalg.norm(desired - actual, ord="fro")290    assert snr >= tol, f"{snr=} < {tol=}"291 292 293def assert_stcs_equal(stc1, stc2):294    """Check that two STC are equal."""295    __tracebackhide__ = True296 297    assert_allclose(stc1.times, stc2.times, err_msg="Times mismatch")298    assert_allclose(stc1.data, stc2.data, err_msg="Data mismatch")299    assert_array_equal(300        stc1.vertices[0],301        stc2.vertices[0],302        err_msg="Left vertices mismatch",303    )304    assert_array_equal(305        stc1.vertices[1],306        stc2.vertices[1],307        err_msg="Right vertices mismatch",308    )309    assert_allclose(stc1.tmin, stc2.tmin, err_msg="tmin mismatch")310    assert_allclose(stc1.tstep, stc2.tstep, err_msg="tstep mismatch")311 312 313def _dig_sort_key(dig):314    """Sort dig keys."""315    return (dig["kind"], dig["ident"])316 317 318def assert_dig_allclose(info_py, info_bin, limit=None):319    """Assert dig allclose."""320    __tracebackhide__ = True321 322    from .._fiff.constants import FIFF323    from .._fiff.meas_info import Info324    from ..bem import fit_sphere_to_headshape325    from ..channels.montage import DigMontage326 327    # test dig positions328    dig_py, dig_bin = info_py, info_bin329    if isinstance(dig_py, Info):330        assert isinstance(dig_bin, Info), "Both must be Info or DigMontage"331        dig_py, dig_bin = dig_py["dig"], dig_bin["dig"]332    else:333        assert isinstance(dig_bin, DigMontage), "Both must be Info or DigMontage"334        assert isinstance(dig_py, DigMontage), "Both must be Info or DigMontage"335        dig_py, dig_bin = dig_py.dig, dig_bin.dig336        info_py = info_bin = None337    assert isinstance(dig_py, list), "dig_py must be a list"338    assert isinstance(dig_bin, list), "dig_bin must be a list"339    dig_py = sorted(dig_py, key=_dig_sort_key)340    dig_bin = sorted(dig_bin, key=_dig_sort_key)341    assert len(dig_py) == len(dig_bin), "Different number of dig points"342    for ii, (d_py, d_bin) in enumerate(zip(dig_py[:limit], dig_bin[:limit])):343        for key in ("ident", "kind", "coord_frame"):344            assert d_py[key] == d_bin[key], f"{key=} mismatch on point {ii}"345        assert_allclose(346            d_py["r"],347            d_bin["r"],348            rtol=1e-5,349            atol=1e-5,350            err_msg=f"Failure on {ii}:\n{d_py['r']}\n{d_bin['r']}",351        )352    if any(d["kind"] == FIFF.FIFFV_POINT_EXTRA for d in dig_py) and info_py is not None:353        r_bin, o_head_bin, o_dev_bin = fit_sphere_to_headshape(354            info_bin, units="m", verbose="error"355        )356        r_py, o_head_py, o_dev_py = fit_sphere_to_headshape(357            info_py, units="m", verbose="error"358        )359        assert_allclose(r_py, r_bin, atol=1e-6, err_msg="Sphere radius mismatch")360        assert_allclose(361            o_dev_py,362            o_dev_bin,363            rtol=1e-5,364            atol=1e-6,365            err_msg="Sphere device origin mismatch",366        )367        assert_allclose(368            o_head_py,369            o_head_bin,370            rtol=1e-5,371            atol=1e-6,372            err_msg="Sphere origin mismatch",373        )374 375 376def _click_ch_name(fig, ch_index=0, button=1):377    """Click on a channel name in a raw/epochs/ICA browse-style plot."""378    from ..viz.utils import _fake_click379 380    fig.canvas.draw()381    text = fig.mne.ax_main.get_yticklabels()[ch_index]382    bbox = text.get_window_extent()383    x = bbox.intervalx.mean()384    y = bbox.intervaly.mean()385    _fake_click(fig, fig.mne.ax_main, (x, y), xform="pix", button=button)386 387 388def _get_suptitle(fig):389    """Get fig suptitle (shim for matplotlib < 3.8.0)."""390    # TODO: obsolete when minimum MPL version is 3.8391    if check_version("matplotlib", "3.8"):392        return fig.get_suptitle()393    else:394        # unreliable hack; should work in most tests as we rarely use `sup_{x,y}label`395        return fig.texts[0].get_text()396 397 398def assert_trans_allclose(actual, desired, dist_tol=0.0, angle_tol=0.0):399    __tracebackhide__ = True400 401    from ..transforms import Transform, angle_distance_between_rigid402 403    if isinstance(actual, Transform):404        assert isinstance(desired, Transform), "Both must be Transform or ndarray"405        assert actual["from"] == desired["from"], "'from' frame mismatch"406        assert actual["to"] == desired["to"], "'to' frame mismatch"407        actual = actual["trans"]408        desired = desired["trans"]409    assert isinstance(actual, np.ndarray), "actual should be ndarray"410    assert isinstance(desired, np.ndarray), "desired should be ndarray"411    assert actual.shape == (4, 4), "actual.shape should be (4, 4)"412    assert desired.shape == (4, 4), "desired.shape should be (4, 4)"413    angle, dist = angle_distance_between_rigid(414        actual, desired, angle_units="deg", distance_units="m"415    )416    assert dist <= dist_tol, (417        f"{1000 * dist:0.3f} > {1000 * dist_tol:0.3f} mm translation"418    )419    assert angle <= angle_tol, f"{angle:0.3f} > {angle_tol:0.3f}° rotation"420