针对pulse-transit的工具

dist/client/mne/viz/_3d_overlay.py

@@ -0,0 +1,184 @@
+"""Classes to handle overlapping surfaces."""
+
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors.
+
+from collections import OrderedDict
+
+import numpy as np
+
+from ..utils import logger
+
+
+class _Overlay:
+    def __init__(self, scalars, colormap, rng, opacity, name):
+        self._scalars = scalars
+        self._colormap = colormap
+        assert rng is not None
+        self._rng = rng
+        self._opacity = opacity
+        self._name = name
+
+    def to_colors(self):
+        from matplotlib.colors import Colormap, ListedColormap
+
+        from ._3d import _get_cmap
+
+        if isinstance(self._colormap, str):
+            cmap = _get_cmap(self._colormap)
+        elif isinstance(self._colormap, Colormap):
+            cmap = self._colormap
+        else:
+            cmap = ListedColormap(
+                self._colormap / 255.0, name=str(type(self._colormap))
+            )
+        logger.debug(
+            f"Color mapping {repr(self._name)} with {cmap.name} "
+            f"colormap and range {self._rng}"
+        )
+
+        rng = self._rng
+        assert rng is not None
+        scalars = self._norm(rng)
+
+        colors = cmap(scalars)
+        if self._opacity is not None:
+            colors[:, 3] *= self._opacity
+        return colors
+
+    def _norm(self, rng):
+        if rng[0] == rng[1]:
+            factor = 1 if rng[0] == 0 else 1e-6 * rng[0]
+        else:
+            factor = rng[1] - rng[0]
+        return (self._scalars - rng[0]) / factor
+
+
+class _LayeredMesh:
+    def __init__(self, renderer, vertices, triangles, normals):
+        self._renderer = renderer
+        self._vertices = vertices
+        self._triangles = triangles
+        self._normals = normals
+
+        self._polydata = None
+        self._actor = None
+        self._is_mapped = False
+
+        self._current_colors = None
+        self._cached_colors = None
+        self._overlays = OrderedDict()
+
+        self._default_scalars = np.ones(vertices.shape)
+        self._default_scalars_name = "Data"
+
+    def map(self):
+        kwargs = {
+            "color": None,
+            "pickable": True,
+            "rgba": True,
+        }
+        mesh_data = self._renderer.mesh(
+            x=self._vertices[:, 0],
+            y=self._vertices[:, 1],
+            z=self._vertices[:, 2],
+            triangles=self._triangles,
+            normals=self._normals,
+            scalars=self._default_scalars,
+            **kwargs,
+        )
+        self._actor, self._polydata = mesh_data
+        self._is_mapped = True
+
+    def _compute_over(self, B, A):
+        assert A.ndim == B.ndim == 2
+        assert A.shape[1] == B.shape[1] == 4
+        A_w = A[:, 3:]  # * 1
+        B_w = B[:, 3:] * (1 - A_w)
+        C = A.copy()
+        C[:, :3] *= A_w
+        C[:, :3] += B[:, :3] * B_w
+        C[:, 3:] += B_w
+        C[:, :3] /= C[:, 3:]
+        return np.clip(C, 0, 1, out=C)
+
+    def _compose_overlays(self):
+        B = cache = None
+        for overlay in self._overlays.values():
+            A = overlay.to_colors()
+            if B is None:
+                B = A
+            else:
+                cache = B
+                B = self._compute_over(cache, A)
+        return B, cache
+
+    def add_overlay(self, scalars, colormap, rng, opacity, name):
+        overlay = _Overlay(
+            scalars=scalars,
+            colormap=colormap,
+            rng=rng,
+            opacity=opacity,
+            name=name,
+        )
+        self._overlays[name] = overlay
+        colors = overlay.to_colors()
+        if self._current_colors is None:
+            self._current_colors = colors
+        else:
+            # save previous colors to cache
+            self._cached_colors = self._current_colors
+            self._current_colors = self._compute_over(self._cached_colors, colors)
+
+        # apply the texture
+        self._apply()
+
+    def remove_overlay(self, names):
+        to_update = False
+        if not isinstance(names, list):
+            names = [names]
+        for name in names:
+            if name in self._overlays:
+                del self._overlays[name]
+                to_update = True
+        if to_update:
+            self.update()
+
+    def _apply(self):
+        if self._current_colors is None or self._renderer is None:
+            return
+        self._polydata[self._default_scalars_name] = self._current_colors
+
+    def update(self, colors=None):
+        if colors is not None and self._cached_colors is not None:
+            self._current_colors = self._compute_over(self._cached_colors, colors)
+        else:
+            self._current_colors, self._cached_colors = self._compose_overlays()
+        self._apply()
+
+    def _clean(self):
+        mapper = self._actor.GetMapper()
+        mapper.SetLookupTable(None)
+        self._actor.SetMapper(None)
+        self._actor = None
+        self._polydata = None
+        self._renderer = None
+
+    def update_overlay(self, name, scalars=None, colormap=None, opacity=None, rng=None):
+        overlay = self._overlays.get(name, None)
+        if overlay is None:
+            return
+        if scalars is not None:
+            overlay._scalars = scalars
+        if colormap is not None:
+            overlay._colormap = colormap
+        if opacity is not None:
+            overlay._opacity = opacity
+        if rng is not None:
+            overlay._rng = rng
+        # partial update: use cache if possible
+        if name == list(self._overlays.keys())[-1]:
+            self.update(colors=overlay.to_colors())
+        else:  # full update
+            self.update()