diff --git a/doc/changes/dev/14150.other.rst b/doc/changes/dev/14150.other.rst new file mode 100644 index 00000000000..948bb812fb1 --- /dev/null +++ b/doc/changes/dev/14150.other.rst @@ -0,0 +1 @@ +Add the setup cell that installs MNE into the browser kernel for the JupyterLite documentation, by `Natneal B`_. diff --git a/doc/sphinxext/_lite_setup_cell.py b/doc/sphinxext/_lite_setup_cell.py new file mode 100644 index 00000000000..672eeca69a0 --- /dev/null +++ b/doc/sphinxext/_lite_setup_cell.py @@ -0,0 +1,766 @@ +# Authors: The MNE-Python contributors. +# License: BSD-3-Clause +# Copyright the MNE-Python contributors. + +# This file is notebook source rather than a module: it installs packages with +# a top-level ``await`` and imports them only afterwards, so the rules about +# import position, await position and import order do not apply to it. Ruff +# still lints and formats everything else here, which is the point of keeping +# it as a real file instead of a string. +# ruff: noqa: E402, F704, I001 + +# Naming: everything this cell defines lands in the notebook's own namespace, +# so anything it invents is _-prefixed and cannot shadow a variable the +# tutorial goes on to use. Module imports are left plain: a tutorial importing +# the same module binds the same object, so there is nothing to protect. +# `mne_data_path` is the deliberate exception, since a reader may want it. + +# Layout, in the order the sections appear below: +# 1. install MNE into the browser kernel +# 2. patch what Pyodide lacks, before MNE is imported +# 3. work out where the data is served and copy the core of it in +# 4. import MNE and point its datasets at that copy +# 5. define the fetch helpers everything after this point uses +# 6. tell MNE where each dataset lives +# 7. wrap the readers, in three groups by how much work each needs +# 8. stub out what WebAssembly cannot do + +# --- JupyterLite setup cell ------------------------------------------------- +# 💡 This cell is automatically added to the start of each notebook. +# It installs MNE and patches the browser environment for Pyodide. +# Downloading this notebook to run it locally? Delete this cell first: +# piplite exists only inside JupyterLite, and a local MNE needs none of +# the patches below. + +# === 1. Install ============================================================== +import piplite + +# Use piplite (not micropip) so the locally-built development MNE wheel +# bundled into the JupyterLite build is preferred over the older PyPI +# release: piplite checks the local index first and falls back to PyPI +# for dependencies. +# keep_going=True so a dependency with no pure-Python wheel is reported +# at the end rather than aborting the whole install on the first one. +await piplite.install( + [ + "mne", + "scikit-learn", + "joblib", + "pandas", + "seaborn", + "mne-connectivity", + "nibabel", + "pyvista-js", + "pyxdf", + "mffpy", + "python-picard", + ], + keep_going=True, +) + +# === 2. Pyodide compatibility, before MNE is imported ======================== +import sys +import os +import io +import inspect +from pathlib import Path + +# Mock multiprocessing — missing in Pyodide but imported by joblib +from unittest.mock import MagicMock + +if "multiprocessing" not in sys.modules: + _mp = MagicMock() + _mp.cpu_count.return_value = 1 + sys.modules["multiprocessing"] = _mp + sys.modules["multiprocessing.util"] = _mp.util + sys.modules["multiprocessing.pool"] = _mp.pool + +# Route requests over the browser's own transport so the downloads that still +# go through pooch work here. That path is pooch.retrieve -> +# pooch.HTTPDownloader -> requests, taken by the fetchers whose files live off +# the docs site (fetch_fsaverage, fetch_infant_template and the parcellation +# ones) and so are not in the copy html_extra_path serves. Every notebook that +# calls one is on JUPYTERLITE_EXCLUDE today, so nothing reaches this on the +# badged pages; it stays because that list is the only thing keeping it that +# way, and a notebook added to the gallery tomorrow would otherwise fail here +# with a Pyodide socket error rather than a real HTTP one. +# +# XMLHttpRequest rather than pyodide.http.open_url, and the same blocking call +# _lite_fetch_rel uses below: open_url reports no status, so a 404 page came +# back looking like a successful 200 and pooch wrote the error page to disk, +# only failing later on a confusing hash mismatch. XHR gives the real status, +# which is what pooch's raise_for_status() needs. Nothing is caught here: the +# browser is the only transport available, so a failure has no fallback worth +# taking and the real error is more useful than a substituted one. +import requests + +_orig_send = requests.Session.send + + +def _pyodide_send(self, request, **kwargs): + from js import XMLHttpRequest + + _xhr = XMLHttpRequest.new() + _xhr.open(request.method or "GET", request.url, False) + _xhr.responseType = "arraybuffer" + _xhr.send() + response = requests.Response() + response.status_code = _xhr.status + response.url = request.url + response.raw = io.BytesIO(bytes(_xhr.response.to_py())) + return response + + +requests.Session.send = _pyodide_send + +# === 3. Where the data comes from =========================================== +# /drive/ in Pyodide requires Cross-Origin-Isolation headers +# (COOP/COEP) which many static servers (e.g. CircleCI artifacts) +# do not send. Fetch the data over HTTP into /tmp/mne_data instead: +# same-origin, no CORS. The data is served at the docs root +# (/mne_data/...) via Sphinx html_extra_path. +# Pyodide may run in a web worker (no `window`); `location` exists +# in both the main thread and workers, so use it to find the docs +# root by splitting on '/lite/'. +import pyodide.http +import js + +try: + _page = str(js.location.href) +except Exception: + _page = str(js.window.location.href) +_base = _page.split("/lite/")[0] + "/mne_data/" +mne_data_path = "/tmp/mne_data" +_mne_data_root = Path(mne_data_path) +_sample_dir = _mne_data_root / "MNE-sample-data" +# Eager 'core': small, commonly-used sample files fetched once at +# notebook start. The heavy files (raw / filt raw / ernoise / fwd / +# inv / src, ~360 MB total) are intentionally omitted here -- they are +# fetched lazily on first read via the reader shims below, so each +# notebook only downloads the sample files it actually uses. +_sample_files = [ + "version.txt", + "MEG/sample/sample_audvis_raw-eve.fif", + "MEG/sample/sample_audvis_filt-0-40_raw-eve.fif", + "MEG/sample/sample_audvis_ecg-proj.fif", + "MEG/sample/sample_audvis-cov.fif", + "MEG/sample/sample_audvis-ave.fif", + "MEG/sample/sample_audvis-no-filter-ave.fif", + "MEG/sample/sample_audvis_raw-trans.fif", + "MEG/sample/sample_audvis-shrunk-cov.fif", + "MEG/sample/sample_audvis-meg-lh.stc", + "MEG/sample/sample_audvis-meg-rh.stc", + "subjects/sample/mri/T1.mgz", + "subjects/sample/surf/rh.pial", + "subjects/sample/surf/lh.pial", + "subjects/sample/surf/rh.white", + "subjects/sample/surf/lh.white", + "subjects/sample/label/lh.aparc.annot", + "subjects/sample/label/rh.aparc.annot", + "SSS/sss_cal_mgh.dat", + "SSS/ct_sparse_mgh.fif", +] +# These are served from the same origin as this page, so if the page loaded, +# the server is up: a miss here means the docs build did not stage the file, +# not that the network is flaky. Several of them (the SSS calibration pair, +# the surfaces read through nibabel) have no lazy path either, so a miss would +# otherwise surface as a confusing error many cells later. Collect every +# failure and raise once, naming them all, since one staging bug usually drops +# more than one file. +# print, not a logger: this cell runs in the browser kernel, not in the Sphinx +# process, so its output is simply what the notebook reader sees. +print("Fetching MNE sample data (once per session)...") +_missing = [] +for _f in _sample_files: + _dst = _sample_dir / _f + if _dst.exists(): + continue + _url = _base + "MNE-sample-data/" + _f + try: + _r = await pyodide.http.pyfetch(_url) + if _r.status != 200: + _missing.append(f"{_f} (HTTP {_r.status})") + continue + _d = await _r.bytes() + # a static server answers a missing path with its 404 page and a 200 + # status, so the body is the only way to tell the two apart + if _d[:4] == b"//`` (group B).""" + _rel = _lite_rel_to_data(subjects_dir if subjects_dir is not None else "") + if not subject or _rel is None: + return + _lite_fetch_optional(f"{_rel}/{subject}/{_p}" for _p in rel_paths) + + +def _lite_dataset_path(folder, probe=None): + """Build a ``data_path()`` that returns ``folder`` under the data root. + + With ``probe``, the named file is fetched when data_path() is called. That + is what covers mtrf, whose .mat is read by scipy rather than by an MNE + reader, so nothing downstream would otherwise fetch it. + """ + + def _data_path(*args, **kwargs): + if probe is not None: + _lite_fetch_rel(folder + "/" + probe) + return _lite_data_path(folder) + + return _data_path + + +def _lite_wrap_reader(module, name): + """Wrap ``module.name`` so its filename argument is fetched before it opens. + + The keyword to intercept is read off the wrapped function rather than + listed by hand: the readers below disagree about whether it is ``fname``, + ``filename`` or ``input_fname``, and a name written out here that drifted + from the real one would silently stop fetching for keyword callers. + """ + orig = getattr(module, name) + arg = next(iter(inspect.signature(orig).parameters)) + + def wrapped(*args, **kwargs): + if args: + args = (_lite_fetch_if_under_mne_data(args[0]),) + args[1:] + elif arg in kwargs: + # move it to a positional argument, since it is no longer in kwargs + args = (_lite_fetch_if_under_mne_data(kwargs.pop(arg)),) + return orig(*args, **kwargs) + + setattr(module, name, wrapped) + + +def _lite_dir_reader(orig): + """Wrap a reader that is handed a folder rather than a file.""" + + def _read(fname, *args, **kwargs): + _rel = _lite_rel_to_data(fname) + if _rel is not None: + try: + _lite_fetch_dir(_rel) + except Exception as _e: + print("[JupyterLite] could not fetch " + str(fname) + ": " + repr(_e)) + return orig(fname, *args, **kwargs) + + return _read + + +def _lite_rebind(name, old, new): + """Point every module that already imported ``old`` at ``new``. + + MNE lazy-loads most of itself, so a module that ran ``from x import f`` + before this cell holds its own reference and would not see the patch. + Modules imported afterwards pick it up on their own. + """ + for _m in list(sys.modules.values()): + if ( + getattr(_m, "__name__", "").startswith("mne") + and getattr(_m, name, None) is old + ): + setattr(_m, name, new) + + +# === 6. Where MNE looks for each dataset ==================================== +# data_path() normally checks for the .tar.gz archive, not just the extracted +# folder, and would try to download from OSF when it does not find one. Point +# each dataset at its folder under the data root instead. The ones with a probe +# file are used by only a couple of notebooks each, so nothing is fetched until +# their data_path() is actually called. +for _ds, _folder, _probe in ( + ("sample", "MNE-sample-data", None), + # testing hands back the folder and lets the shimmed readers pull + # individual files, so a notebook that wants the EEGLAB recording does + # not also drag down the 39 MB movement raw. + ("testing", "MNE-testing-data", None), + # datasets behind a single example each; only the files those examples + # read are served, and the readers below pull them individually + ("ssvep", "ssvep-example-data", None), + ("misc", "MNE-misc-data", None), + ("eyelink", "MNE-eyelink-data", None), + ("fnirs_motor", "MNE-fNIRS-motor-data", None), + ("refmeg_noise", "MNE-refmeg-noise-data", None), + ("phantom_kernel", "MNE-phantom-kernel-data", None), + ("multimodal", "MNE-multimodal-data", None), + # kiloword/erp_core for Epochs 30 & 40, mtrf for the decoding examples + ("kiloword", "MNE-kiloword-data", "kword_metadata-epo.fif"), + ("erp_core", "MNE-ERP-CORE-data", "ERP-CORE_Subject-001_Task-Flankers_eeg.fif"), + ("mtrf", "mTRF_1.5", "speech_data.mat"), +): + getattr(mne.datasets, _ds).data_path = _lite_dataset_path(_folder, _probe) +del _ds, _folder, _probe + + +# eegbci is addressed by subject and run rather than by path, so it needs its +# own shim rather than a row in the table above. +def _lite_eegbci_load_data(subjects, runs, *args, **kwargs): + # the parameter is `subjects`, matching MNE: 35_eeg_no_mri calls it by + # keyword, so a shim spelled `subject` would raise TypeError there + _runs = [runs] if isinstance(runs, (int, float)) else list(runs) + _subjects = list(subjects) if isinstance(subjects, (list, tuple)) else [subjects] + _out = [] + for _s in _subjects: + for _r in _runs: + _rel = ( + "MNE-eegbci-data/files/eegmmidb/1.0.0/" + f"S{int(_s):03d}/S{int(_s):03d}R{int(_r):02d}.edf" + ) + _out.append(_lite_fetch_rel(_rel)) + return _out + + +mne.datasets.eegbci.load_data = _lite_eegbci_load_data + +# === 7. Reader overrides ==================================================== +# MNE functions need one of three treatments here, depending on how much they +# do before the file is actually opened. +# +# A. reads one file, and validates the name first. Nearly every MNE reader +# calls _check_fname(must_exist=True) before opening anything, so patching +# that single function covers read_info, read_evokeds, read_cov, +# read_label and the rest at once. A handful skip the validation, and are +# listed in a table instead. Nothing else is needed for this group. +# B. probes the filesystem before any reader runs. _get_head_surface calls +# os.path.exists, plot_bem globs bem/*.surf, so a fetch-on-open hook never +# fires. The candidates have to be on disk before the probe. +# C. one filename that means several files. read_raw_brainvision is handed +# only the .vhdr, opens it, reads the names of its .eeg and .vmrk out of +# it, and opens those, by which point we are inside the reader and it is +# too late to fetch. Same shape for EEGLAB (.set + .fdt), a .stc stem +# (lh + rh), and the formats that are a directory rather than a file. + +# --- A. reads one file ------------------------------------------------------ +# The general hook. Failures stay silent here so MNE still raises its own, +# clearer error for a file that genuinely is missing. +import mne.utils.check as mne_check + +_orig_check_fname = mne_check._check_fname + + +def _lite_check_fname(fname, overwrite=False, must_exist=False, *args, **kwargs): + if must_exist: + try: + _lite_fetch_if_under_mne_data(fname) + except Exception: + pass + return _orig_check_fname(fname, overwrite, must_exist, *args, **kwargs) + + +mne_check._check_fname = _lite_check_fname +_lite_rebind("_check_fname", _orig_check_fname, _lite_check_fname) +# The readers that open their file without validating it first, so the hook +# above never sees them. Each needs nothing but its file fetched. +for _module, _name in ( + (mne, "read_forward_solution"), + (mne.minimum_norm, "read_inverse_operator"), + (mne.io, "read_raw_fif"), + (mne.io, "read_raw"), + (mne, "read_source_spaces"), + (mne, "read_label"), + (mne, "read_epochs"), + (mne.io, "read_raw_edf"), + (mne, "read_bem_solution"), + (mne, "read_events"), + (mne.io, "read_raw_eyelink"), + (mne.chpi, "read_head_pos"), +): + _lite_wrap_reader(_module, _name) +del _module, _name +# The eyetracking heatmap example draws its stimulus straight through pyplot, +# and read_xdf goes through pyxdf. Neither is an MNE reader, but both take a +# path we serve, so they get the same treatment. +import matplotlib.pyplot as plt + +_orig_imread = plt.imread + + +def _lite_imread(fname, *args, **kwargs): + return _orig_imread(_lite_fetch_if_under_mne_data(fname), *args, **kwargs) + + +plt.imread = _lite_imread +# guarded: pyxdf has no pure-Python wheel on every Pyodide build, and only the +# XDF example needs it +try: + import pyxdf + + _orig_load_xdf = pyxdf.load_xdf + + def _lite_load_xdf(fname, *args, **kwargs): + return _orig_load_xdf(_lite_fetch_if_under_mne_data(fname), *args, **kwargs) + + pyxdf.load_xdf = _lite_load_xdf +except Exception: + pass + +# --- B. probes the filesystem first ----------------------------------------- +# plot_alignment locates its head surface with os.path.exists before any reader +# runs. Fetch the candidates first and let MNE choose as it normally would. +# Several viz modules bind the name at import time, so rebind it wherever the +# original landed rather than in one known place. +import mne._freesurfer as mne_fs + +_orig_get_head_surface = mne_fs._get_head_surface + + +def _lite_get_head_surface(surf, subject, subjects_dir, bem=None, verbose=None): + if surf in ("head-dense", "seghead"): + _cands = [f"bem/{subject}-head-dense.fif", "surf/lh.seghead"] + else: + # same order MNE tries, so the browser picks the same + # surface the rendered docs did + _cands = ["bem/outer_skin.surf", f"bem/{subject}-head.fif"] + _lite_fetch_candidates(subject, subjects_dir, _cands) + return _orig_get_head_surface(surf, subject, subjects_dir, bem=bem, verbose=verbose) + + +mne_fs._get_head_surface = _lite_get_head_surface +# import the 3D module first so the rebind is guaranteed to see it; +# anything imported later picks the patched name up on its own. +import mne.viz._3d # noqa: F401 + +_lite_rebind("_get_head_surface", _orig_get_head_surface, _lite_get_head_surface) +# same story for the skull surfaces, which _check_fname insists +# already exist on disk +_orig_get_skull_surface = mne_fs._get_skull_surface + + +def _lite_get_skull_surface(surf, subject, subjects_dir, bem=None, verbose=None): + _lite_fetch_candidates(subject, subjects_dir, [f"bem/{surf}_skull.surf"]) + return _orig_get_skull_surface( + surf, subject, subjects_dir, bem=bem, verbose=verbose + ) + + +mne_fs._get_skull_surface = _lite_get_skull_surface +_lite_rebind("_get_skull_surface", _orig_get_skull_surface, _lite_get_skull_surface) +# dig_mri_distances reaches a second, unrelated _get_head_surface, the +# one in mne/surface.py: it takes a list of candidate sources and +# probes bem/ with os.path.exists and glob, raising if the directory +# is absent, so the candidates have to land before it runs. +import mne.surface as mne_surface + +_orig_surface_head = mne_surface._get_head_surface + + +def _lite_surface_head_surface( + subject, source, subjects_dir, on_defects, raise_error=True +): + _srcs = [source] if isinstance(source, str) else list(source) + _lite_fetch_candidates( + subject, subjects_dir, [f"bem/{subject}-{_s}.fif" for _s in _srcs] + ) + return _orig_surface_head( + subject, source, subjects_dir, on_defects, raise_error=raise_error + ) + + +# no _lite_rebind for this one: unlike the _freesurfer function above, nothing +# outside mne/surface.py imports it by name, so patching the module is enough. +mne_surface._get_head_surface = _lite_surface_head_surface +# plot_bem globs bem/*.surf and requires the bem directory to exist, +# so pull its three contours (plus the MRI it draws them on) down +# first; fetching creates the directory as a side effect. +_orig_plot_bem = mne.viz.plot_bem + + +def _lite_plot_bem(subject=None, subjects_dir=None, *args, **kwargs): + _want = [ + "bem/inner_skull.surf", + "bem/outer_skull.surf", + "bem/outer_skin.surf", + "mri/" + str(kwargs.get("mri", "T1.mgz")), + ] + _bs = kwargs.get("brain_surfaces") + if _bs is not None: + _bs = [_bs] if isinstance(_bs, str) else list(_bs) + for _b in _bs: + _want += [f"surf/lh.{_b}", f"surf/rh.{_b}"] + _lite_fetch_candidates(subject, subjects_dir, _want) + return _orig_plot_bem(subject, subjects_dir, *args, **kwargs) + + +mne.viz.plot_bem = _lite_plot_bem + +# --- C. one filename, several files ----------------------------------------- +# An EEGLAB .set keeps its samples in a sibling .fdt, so fetch both. +_orig_read_raw_eeglab = mne.io.read_raw_eeglab + + +def _lite_read_raw_eeglab(input_fname, *args, **kwargs): + _rel = _lite_rel_to_data(input_fname) + if _rel is not None: + _lite_fetch_optional((_rel, _rel[:-4] + ".fdt")) + return _orig_read_raw_eeglab(input_fname, *args, **kwargs) + + +mne.io.read_raw_eeglab = _lite_read_raw_eeglab +# a BrainVision .vhdr is a text header pointing at a .eeg and a .vmrk +_orig_read_raw_brainvision = mne.io.read_raw_brainvision + + +def _lite_read_raw_brainvision(vhdr_fname, *args, **kwargs): + _rel = _lite_rel_to_data(vhdr_fname) + if _rel is not None: + _stem = _rel[:-5] if _rel.endswith(".vhdr") else _rel + _lite_fetch_optional((_rel, _stem + ".eeg", _stem + ".vmrk")) + return _orig_read_raw_brainvision(vhdr_fname, *args, **kwargs) + + +mne.io.read_raw_brainvision = _lite_read_raw_brainvision +# read_source_estimate is handed the stem of a .stc pair, so fetch +# both hemispheres before letting MNE resolve the name itself. +_orig_read_source_estimate = mne.read_source_estimate + + +def _lite_read_source_estimate(fname, *args, **kwargs): + _rel = _lite_rel_to_data(fname) + if _rel is not None: + _lite_fetch_optional(_rel + _suf for _suf in ("", "-lh.stc", "-rh.stc")) + return _orig_read_source_estimate(fname, *args, **kwargs) + + +mne.read_source_estimate = _lite_read_source_estimate +# read_raw_nirx and read_raw_egi open a folder, listed by its manifest +mne.io.read_raw_nirx = _lite_dir_reader(mne.io.read_raw_nirx) +mne.io.read_raw_egi = _lite_dir_reader(mne.io.read_raw_egi) +# the odd one out in this group: the name is enough, but it points inside the +# installed package rather than at the data root. The logging tutorial builds +# a path into mne/**/tests, which the wheel excludes, so copy the served file +# to where the tutorial expects it rather than editing the tutorial. +import shutil + +_orig_read_raw_kit = mne.io.read_raw_kit + + +def _lite_read_raw_kit(input_fname, *args, **kwargs): + _p = Path(str(input_fname)) + if _p.name == "test.sqd" and not _p.exists(): + try: + _staged = _lite_fetch_rel("MNE-kit-testdata/test.sqd") + _p.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(_staged, _p) + except Exception as _e: + print("[JupyterLite] could not stage test.sqd: " + repr(_e)) + return _orig_read_raw_kit(input_fname, *args, **kwargs) + + +mne.io.read_raw_kit = _lite_read_raw_kit + +# === 8. What WebAssembly cannot do ========================================== +# Pyodide/WASM has no OS threads, so MNE's ProgressBar background +# updater thread (used by the ProgressBar context manager, e.g. in +# permutation cluster tests) crashes with 'can't start new thread'. +# That thread only animates a cosmetic bar: the computation runs on +# the main thread and __exit__ writes the final state, so no-op its +# start/join. Only affects notebooks that use it; results are unchanged. +# Guarded because this is a private MNE path: if it is ever renamed, losing a +# cosmetic patch is better than failing every notebook at the setup cell. +try: + from mne.utils import progressbar + + progressbar._UpdateThread.start = lambda self: None + progressbar._UpdateThread.join = lambda self, *args, **kwargs: None +except Exception: + pass +# tqdm also spawns its own monitor thread, which likewise can't start in +# WASM and emits a TqdmMonitorWarning. Setting monitor_interval=0 before +# any bar is created skips that thread entirely (bars still display). +# Guarded because tqdm is a transitive dependency that may not be installed. +try: + import tqdm + + tqdm.tqdm.monitor_interval = 0 +except Exception: + pass + +# Switch matplotlib to inline so figures render in the notebook. +import IPython + +IPython.get_ipython().run_line_magic("matplotlib", "inline") + +# Silence the spurious 'FigureCanvasAgg is non-interactive' warning that the +# inline Agg canvas raises from fig.show(). MNE's own plt_show no longer +# triggers it (gh-14076 taught it to call plt.show() on inline backends), but +# tutorials still call fig.show() directly -- 50_ssvep does it four times, and +# 10_background_stats once -- and those warn. Every path resolves fig.show on +# the class at call time, so a no-op here covers them all. Figures still +# render via the inline backend. +import matplotlib.figure as mpl_figure + +mpl_figure.Figure.show = lambda self, *a, **k: None +import importlib + +_viz_utils = importlib.import_module("mne.viz.utils") + + +# Also display+close via IPython for paths that call plt_show +# directly, so figures render exactly once. +def _pyodide_plt_show(show=True, fig=None, **kwargs): + if not show: + return + import IPython.display + + _f = fig if fig is not None else plt.gcf() + IPython.display.display(_f) + plt.close(_f) + + +_viz_utils.plt_show = _pyodide_plt_show diff --git a/doc/sphinxext/_lite_setup_cell_3d.py b/doc/sphinxext/_lite_setup_cell_3d.py new file mode 100644 index 00000000000..0e5a8ad3097 --- /dev/null +++ b/doc/sphinxext/_lite_setup_cell_3d.py @@ -0,0 +1,417 @@ +# Authors: The MNE-Python contributors. +# License: BSD-3-Clause +# Copyright the MNE-Python contributors. + +# The experimental part of the browser setup, kept apart from the rest so the +# solid ground and the shifting ground are easy to tell apart. Everything here +# stands in for MNE's Brain/VTK stack, which has no WebAssembly build, and is +# the part most likely to be dropped as pyvista-js gains features upstream. +# Appended after the base cell, which it depends on: the second block below +# uses the matplotlib-inline shim that cell installs. +# This runs as a continuation of the base cell, in the same namespace, so it +# reads names that cell defined (mne, plt, the fetch helpers) rather than +# importing them again; F821 is off for that reason, not to hide typos. +# ruff: noqa: E402, F704, F821, I001 + +# --- JupyterLite setup cell, 3D ----------------------------------------------- +# EXPERIMENTAL 3D: MNE's normal Brain/VTK stack can't load in WASM, so +# route SourceEstimate.plot() through pyvista-js (vtk.js) instead. +# pyvista-js (0.15) has no scalar colormap in its renderer, so we +# approximate MNE's Brain look with solid-colored meshes: a two-tone +# curvature base (light gyri + dark sulci) plus many thin 'hot' bands +# for the activation, on a black background with even scene lighting. +# Static, one time point, no time slider yet. +# +# A failed render prints and lets the notebook carry on, which is the opposite +# of how the data fetch in the base cell behaves. That is deliberate: a file +# missing there means the docs build is broken and should say so, while this +# whole shim stands in for a stack with no WebAssembly build at all, so failing +# hard would take out every 3D notebook rather than report one bug. +# +# The stub 'brain' it returns makes the decorating calls (add_foci/add_text/ +# show_view/...) no-ops so the rest of the notebook still runs. screenshot() is +# the one exception, and it raises: see below. +# Say once per session that what the browser draws is not what the rendered +# docs show, so a reader comparing the two is not left guessing. pyvista-js has +# no scalar colormap, so activation arrives as discrete solid bands rather than +# a continuous scale, there is no colorbar or time slider, and the hemispheres +# are drawn side by side rather than in anatomical position. +_lite_3d_noted = False + + +def _lite_note_3d_approximation(): + global _lite_3d_noted + if _lite_3d_noted: + return + _lite_3d_noted = True + print( + "[JupyterLite] 3D drawn with pyvista-js: activation is shown as solid " + "colour bands at a single time point, with the hemispheres side by " + "side and no colorbar. The figure in the rendered docs is MNE's full " + "Brain view and will not look the same." + ) + + +class _LiteBrain: + def screenshot(self, *args, **kwargs): + # No blank array here. vtk.js draws into a browser canvas that Python + # cannot read back, so there is no image to return, and handing back a + # blank one is worse than failing: 10_publication_figure crops its + # screenshot and shows before/after, so it would publish two black + # squares as though they were the real thing. A notebook that needs a + # screenshot belongs in JUPYTERLITE_EXCLUDE instead. + raise NotImplementedError( + "brain.screenshot() is not available in JupyterLite: the vtk.js " + "renderer draws to a browser canvas that Python cannot read back. " + "Run this notebook locally to capture the scene." + ) + + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + +def _lite_stc_plot(self, *args, **kwargs): + try: + import numpy as _np + import nibabel as _nib + from scipy.spatial import cKDTree as _KDTree + from matplotlib import colormaps as _cmaps + import pyvista_js as _pv + + _subj = ( + kwargs.get("subject") + or (args[0] if args and isinstance(args[0], str) else None) + or "sample" + ) + _sdir = kwargs.get("subjects_dir") + # kept as a str on both branches: it is concatenated below + _sdir = ( + str(_sdir) + if _sdir is not None + else str(_lite_data_path("MNE-sample-data/subjects")) + ) + # surfaces are fetched relative to the served mne_data root, so + # derive that from subjects_dir rather than assuming sample -- + # a dataset may keep its FreeSurfer subjects under its own folder. + _rel_sdir = ( + _lite_rel_to_data(_sdir) + if _lite_rel_to_data(_sdir) is not None + else "MNE-sample-data/subjects" + ) + _init = kwargs.get("initial_time", None) + if _init is None: + _ti = int(_np.argmax(_np.abs(self.data).mean(0))) + else: + _ti = int(_np.argmin(_np.abs(self.times - _init))) + _hot = _cmaps["hot"] + # Tuned against the inflated FreeSurfer surfaces MNE ships, whose + # coordinates are in mm. + _N = 10 # activation value bands + _BLOB_MM = 12.0 # colour a surface vertex from an active source within + # this radius, so a single-vertex source reads as a blob and not a dot + _HEMI_MM = 60.0 # push the hemispheres apart so they do not overlap + _LIFT = 0.02 # raise each band off the surface to avoid z-fighting + _HOT_LO, _HOT_HI = 0.25, 0.66 # slice of 'hot' to use; its ends are + # near-black and near-white, which read as background here + _SPARSE_P90 = 0.05 # below this fraction of the max the 90th pct means + _SPARSE_FLOOR = 0.4 # the data is sparse, so threshold on the max + + def _flat(_t): + return _np.hstack( + [_np.full((len(_t), 1), 3, dtype=_np.int64), _t.astype(_np.int64)] + ).ravel() + + def _sub(_pts, _tris, _mask, _lift=0.0, _cen=None): + _sel = _tris[_mask] + if len(_sel) == 0: + return None + _u, _iv = _np.unique(_sel, return_inverse=True) + _p = _pts[_u] + if _lift and _cen is not None: + _p = _cen + (_p - _cen) * (1.0 + _lift) + return _p, _iv.reshape(-1, 3) + + _plotter = _pv.Plotter() + _plotter.background_color = "black" + # even lighting so the surface isn't black when rotated + for _lp in ( + (1, 0, 0), + (-1, 0, 0), + (0, 1, 0), + (0, -1, 0), + (0, 0, 1), + (0, 0, -1), + ): + _plotter.add_light( + _pv.Light( + position=(300.0 * _lp[0], 300.0 * _lp[1], 300.0 * _lp[2]), + focal_point=(0.0, 0.0, 0.0), + intensity=0.4, + ) + ) + _nlh = len(self.vertices[0]) + _hemis = (("lh", 0, self.vertices[0]), ("rh", 1, self.vertices[1])) + for _h, _hi, _vno in _hemis: + if len(_vno) == 0: + continue + _pre = _rel_sdir + "/" + _subj + "/surf/" + _h + _lite_fetch_rel(_pre + ".inflated") + _lite_fetch_rel(_pre + ".curv") + _bpath = _sdir + "/" + _subj + "/surf/" + _h + _rr, _tris = mne.read_surface(_bpath + ".inflated") + _cv = _nib.freesurfer.read_morph_data(_bpath + ".curv") + _hdata = self.data[:_nlh] if _hi == 0 else self.data[_nlh:] + # color each surface vertex from the nearest ACTIVE source + # within a small radius, so single-vertex (point) sources + # show as visible blobs and dense sources fill in as usual + _sv = _hdata[:, _ti].astype(float) + _act = _sv != 0 + _scal = _np.zeros(len(_rr)) + if _act.any(): + _atree = _KDTree(_rr[_vno][_act]) + _ad, _ai = _atree.query(_rr) + _scal = _np.where(_ad <= _BLOB_MM, _sv[_act][_ai], 0.0) + # offset hemispheres along x so they do not overlap + _off = -_HEMI_MM if _h == "lh" else _HEMI_MM + _pts = _np.round(_rr, 2) + _pts[:, 0] = _pts[:, 0] + _off + _cen = _pts.mean(0) + # curvature base: light gyri (curv<0) + dark sulci (curv>=0) + _fc = _cv[_tris].mean(1) + for _cm, _col in ( + (_fc < 0, (0.68, 0.68, 0.68)), + (_fc >= 0, (0.38, 0.38, 0.38)), + ): + _s = _sub(_pts, _tris, _cm) + if _s is not None: + _plotter.add_mesh( + _pv.PolyData(points=_s[0], faces=_flat(_s[1])), + color=_col, + smooth_shading=True, + ) + # activation as a smooth hot gradient in N value bands, + # each lifted 2% off the surface to avoid z-fighting + _fv = _scal[_tris].mean(1) + _p90 = _np.percentile(_scal, 90.0) + _fmax = float(_scal.max()) + # keep the background gray: for sparse point sources the + # 90th pct is ~0 (most of the brain is zero), which would + # paint everything, so fall back to a fraction of the max. + _fmin = _p90 if _p90 > _fmax * _SPARSE_P90 else _fmax * _SPARSE_FLOOR + if _fmax > _fmin: + _edges = _np.linspace(_fmin, _fmax, _N + 1) + for _i in range(_N): + if _i < _N - 1: + _m = (_fv >= _edges[_i]) & (_fv < _edges[_i + 1]) + else: + _m = _fv >= _edges[_i] + if int(_m.sum()) == 0: + continue + _rgb = _hot(_HOT_LO + (_HOT_HI - _HOT_LO) * (_i / (_N - 1))) + _col = (float(_rgb[0]), float(_rgb[1]), float(_rgb[2])) + _s = _sub(_pts, _tris, _m, _LIFT, _cen) + if _s is not None: + _plotter.add_mesh( + _pv.PolyData(points=_s[0], faces=_flat(_s[1])), + color=_col, + smooth_shading=True, + ) + # Open on the lateral profile (camera along the medial-lateral + # X axis, superior up), like native MNE, instead of vtk.js's + # default anterior/face-on view. Guarded so a missing + # view_vector never costs us the render. + try: + _plotter.view_vector((-1.0, 0.0, 0.0), viewup=(0.0, 0.0, 1.0)) + except Exception: + pass + _plotter.show() + _lite_note_3d_approximation() + except Exception as _e: + print("[JupyterLite] pyvista-js 3D render unavailable: " + repr(_e)) + return _LiteBrain() + + +mne.SourceEstimate.plot = _lite_stc_plot + + +# EXPERIMENTAL 3D: plot_sparse_source_estimates builds its 3D renderer +# BEFORE the time-course figure, so in WASM the whole call dies and the +# notebook loses both halves. Rebuild it here: the same glass brain from +# the source space and a marker per active dipole via pyvista-js, plus +# the matplotlib time courses (which are the quantitative half). Same +# approach as the SourceEstimate.plot shim above. +def _lite_plot_sparse_source_estimates( + src, + stcs, + colors=None, + linewidth=2, + fontsize=18, + bgcolor=(0.05, 0, 0.1), + opacity=0.2, + brain_color=(0.7,) * 3, + show=True, + high_resolution=False, + fig_name=None, + fig_number=None, + labels=None, + modes=("cone", "sphere"), + scale_factors=(1, 0.6), + **kwargs, +): + import numpy as _np + from itertools import cycle as _cycle + from matplotlib.colors import to_rgb as _to_rgb + + if not isinstance(stcs, list): + stcs = [stcs] + _lhp = src[0]["rr"] + _pts = _np.r_[_lhp, src[1]["rr"]] * 170 + _nrm = _np.r_[src[0]["nn"], src[1]["nn"]] + # use_tris is the decimated mesh and can be None on some source + # spaces; fall back to the full tris in that case. + _lt = src[0]["tris"] if high_resolution else src[0]["use_tris"] + _rt = src[1]["tris"] if high_resolution else src[1]["use_tris"] + if _lt is None or _rt is None: + _lt, _rt = src[0]["tris"], src[1]["tris"] + _faces = _np.r_[_lt, len(_lhp) + _rt] + _vertnos = [_np.r_[_s.lh_vertno, len(_lhp) + _s.rh_vertno] for _s in stcs] + _uniq = _np.unique(_np.concatenate(_vertnos).ravel()) + # --- time courses ------------------------------------------------- + _fig = plt.figure(fig_number, layout="constrained") + _fig.clf() + _ax = _fig.add_subplot(111) + _cyc = _cycle( + colors + if colors is not None + else plt.rcParams["axes.prop_cycle"].by_key()["color"] + ) + _marks = [] + for _v in _uniq: + _ind = [_k for _k, _vn in enumerate(_vertnos) if _v in _vn] + _c = next(_cyc) + _marks.append((int(_v), _to_rgb(_c), len(_ind) > 1)) + for _k in _ind: + _m = _vertnos[_k] == _v + _ax.plot( + 1e3 * stcs[_k].times, + 1e9 * stcs[_k].data[_m].ravel(), + c=_c, + linewidth=linewidth, + ) + _ax.set_xlabel("Time (ms)", fontsize=fontsize) + _ax.set_ylabel("Source amplitude (nAm)", fontsize=fontsize) + if fig_name is not None: + _ax.set_title(fig_name) + _pyodide_plt_show(show) + # --- glass brain + dipole markers --------------------------------- + try: + import pyvista_js as _pv + + _plotter = _pv.Plotter() + _plotter.background_color = tuple( + float(min(max(_x, 0.0), 1.0)) for _x in bgcolor + ) + for _lp in ( + (1, 0, 0), + (-1, 0, 0), + (0, 1, 0), + (0, -1, 0), + (0, 0, 1), + (0, 0, -1), + ): + _plotter.add_light( + _pv.Light( + position=(300.0 * _lp[0], 300.0 * _lp[1], 300.0 * _lp[2]), + focal_point=(0.0, 0.0, 0.0), + intensity=0.4, + ) + ) + _flat_faces = _np.hstack( + [_np.full((len(_faces), 1), 3, dtype=_np.int32), _faces.astype(_np.int32)] + ).ravel() + _plotter.add_mesh( + _pv.PolyData(points=_pts.astype(_np.float32), faces=_flat_faces), + color=tuple(float(_x) for _x in brain_color), + opacity=float(opacity), + smooth_shading=True, + ) + for _v, _col, _common in _marks: + _sf = float(scale_factors[1] if _common else scale_factors[0]) + _mode = modes[1] if _common else modes[0] + _xyz = tuple(float(_q) for _q in _pts[_v]) + if _mode == "sphere": + _glyph = _pv.Sphere(radius=_sf, center=_xyz) + else: + _glyph = _pv.Cone( + center=_xyz, + direction=tuple(float(_q) for _q in _nrm[_v]), + height=2.0 * _sf, + radius=_sf, + ) + _plotter.add_mesh(_glyph, color=_col, smooth_shading=True) + try: + _plotter.view_vector((-1.0, 0.0, 0.0), viewup=(0.0, 0.0, 1.0)) + except Exception: + pass + _plotter.show() + _lite_note_3d_approximation() + except Exception as _e: + print("[JupyterLite] pyvista-js glass brain unavailable: " + repr(_e)) + + +mne.viz.plot_sparse_source_estimates = _lite_plot_sparse_source_estimates + +# Each MNE plot is rendered once by _pyodide_plt_show above (display()). +# When a plot call is also a cell's last expression, the method returns +# the Figure, which Jupyter echoes a SECOND time as the Out[] result +# (the duplicate seen below inline plots). Drop that redundant echo for +# Figures (and pure lists of Figures, e.g. ica.plot_properties) so each +# plot appears exactly once. Non-figure results (numbers, DataFrames, +# reprs) are untouched, and raw matplotlib figures never shown still +# render via the inline backend's end-of-cell flush, so nothing hides. +# Wrapped in try/except (like the patches below): if anything about +# the displayhook is unexpected, silently keep the current behavior +# (harmless double render) rather than breaking the setup cell. +try: + _lite_dh = type(IPython.get_ipython().displayhook) + if not getattr(_lite_dh, "_lite_no_fig_echo", False): + _lite_dh_call = _lite_dh.__call__ + + def _lite_displayhook(self, result=None): + if isinstance(result, mpl_figure.Figure): + result = None + elif ( + isinstance(result, (list, tuple)) + and result + and all(isinstance(_x, mpl_figure.Figure) for _x in result) + ): + result = None + return _lite_dh_call(self, result) + + _lite_dh.__call__ = _lite_displayhook + _lite_dh._lite_no_fig_echo = True +except Exception: + pass + +# Real fix (not a warnings filter) for the threadpoolctl Pyodide +# RuntimeWarning seen via mne.sys_info(): threadpoolctl 3.6.0 (latest +# release) still calls the deprecated Pyodide JsProxy.as_object_map(). +# Pyodide's own message says to use as_py_json() instead; both yield the +# same library filepaths, so we swap the call at its source. This removes +# the deprecated API usage entirely, so the warning is never emitted. +# The upstream fix is already merged (joblib/threadpoolctl#201) but +# unreleased; Pyodide bundles the released 3.6.0 wheel. DROP THIS PATCH +# once threadpoolctl 3.7.0 is released and Pyodide bundles it. +try: + import threadpoolctl + + def _find_libraries_pyodide(self): + from pyodide_js._module import LDSO + + for _fp in LDSO.loadedLibsByName.as_py_json(): + if Path(_fp).exists(): + self._make_controller_from_path(_fp) + + threadpoolctl.ThreadpoolController._find_libraries_pyodide = _find_libraries_pyodide +except Exception: + pass diff --git a/doc/sphinxext/jupyterlite_setup_cell.py b/doc/sphinxext/jupyterlite_setup_cell.py new file mode 100644 index 00000000000..298347ce5aa --- /dev/null +++ b/doc/sphinxext/jupyterlite_setup_cell.py @@ -0,0 +1,50 @@ +"""The setup cell prepended to every JupyterLite notebook. + +It installs MNE into the browser kernel and patches what Pyodide does not +provide: data fetching over HTTP, the readers that expect files already on +disk, and the 3D renderer. The cell lives in ``_lite_setup_cell.py`` and +``_lite_setup_cell_3d.py`` as ordinary Python, so ruff lints and formats it; +this module only reads those files and joins them into the string the browser +kernel needs. The 3D half is kept separate because it stands in for MNE's +Brain/VTK stack and is the part most likely to change as pyvista-js gains +features upstream. + +The docs build prepends it only to the notebooks copied into the JupyterLite +contents. It deliberately does NOT go through ``first_notebook_cell``: that is +applied when the notebook is generated, so it would also land in the ``.ipynb`` +offered for download, where ``piplite`` does not exist and the notebook would +fail on its first cell. + +The other direction is covered in the cell itself: a notebook downloaded from +inside JupyterLite does carry the cell, and it says to delete it before running +locally, for the same reason. +""" + +# Authors: The MNE-Python contributors. +# License: BSD-3-Clause +# Copyright the MNE-Python contributors. + +from pathlib import Path + +from jupyterlite_lite_renderer import LITE_RENDERER_CELL + +# Each source file read below is split at this banner: everything after it is +# what the notebook runs, and what sits above it in that file (license header, +# ruff directives, notes for whoever edits it) stays behind. +_BANNER = "# --- JupyterLite setup cell" + + +def _read(name): + _source = Path(__file__).parent / name + _text = _source.read_text() + if _BANNER not in _text: + raise RuntimeError(f"{_source.name} is missing the {_BANNER!r} banner") + _body = _text[_text.index(_BANNER) :] + return _body[_body.index("\n") + 1 :] + + +# Order matters: the 3D half reads the matplotlib shim the base half installs, +# and the renderer goes last so MNE is already imported by the time it runs. +LITE_SETUP_CELL = ( + _read("_lite_setup_cell.py") + _read("_lite_setup_cell_3d.py") + LITE_RENDERER_CELL +)