185 lines
5.4 KiB
Python
185 lines
5.4 KiB
Python
"""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()
|