initial commit

mne/viz/_brain/surface.py

@@ -0,0 +1,177 @@
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors.
+
+from os import path as path
+
+import numpy as np
+
+from ...surface import _read_patch, complete_surface_info, read_curvature, read_surface
+from ...utils import _check_fname, _check_option, _validate_type, get_subjects_dir
+
+
+class _Surface:
+    """Container for a brain surface.
+
+    It is used for storing vertices, faces and morphometric data
+    (curvature) of a hemisphere mesh.
+
+    Parameters
+    ----------
+    subject : string
+        Name of subject
+    hemi : {'lh', 'rh'}
+        Which hemisphere to load
+    surf : string
+        Name of the surface to load (eg. inflated, orig ...).
+    subjects_dir : str | None
+        If not None, this directory will be used as the subjects directory
+        instead of the value set using the SUBJECTS_DIR environment variable.
+    offset : float | None
+        If 0.0, the surface will be offset such that the medial
+        wall is aligned with the origin. If None, no offset will
+        be applied. If != 0.0, an additional offset will be used.
+    units : str
+        Can be 'm' or 'mm' (default).
+    x_dir : ndarray | None
+        The x direction to use for offset alignment.
+
+    Attributes
+    ----------
+    bin_curv : numpy.ndarray
+        Curvature values stored as non-negative integers.
+    coords : numpy.ndarray
+        nvtx x 3 array of vertex (x, y, z) coordinates.
+    curv : numpy.ndarray
+        Vector representation of surface morpometry (curvature) values as
+        loaded from a file.
+    grey_curv : numpy.ndarray
+        Normalized morphometry (curvature) data, used in order to get
+        a gray cortex.
+    faces : numpy.ndarray
+        nfaces x 3 array of defining mesh triangles.
+    hemi : {'lh', 'rh'}
+        Which hemisphere to load.
+    nn : numpy.ndarray
+        Vertex normals for a triangulated surface.
+    offset : float | None
+        If float, align inside edge of each hemisphere to center + offset.
+        If None, do not change coordinates (default).
+    subject : string
+        Name of subject.
+    surf : string
+        Name of the surface to load (eg. inflated, orig ...).
+    units : str
+        Can be 'm' or 'mm' (default).
+    """
+
+    def __init__(
+        self,
+        subject,
+        hemi,
+        surf,
+        subjects_dir=None,
+        offset=None,
+        units="mm",
+        x_dir=None,
+    ):
+        x_dir = np.array([1.0, 0, 0]) if x_dir is None else x_dir
+        assert isinstance(x_dir, np.ndarray)
+        assert np.isclose(np.linalg.norm(x_dir), 1.0, atol=1e-6)
+        assert hemi in ("lh", "rh")
+        _validate_type(offset, (None, "numeric"), "offset")
+
+        self.units = _check_option("units", units, ("mm", "m"))
+        self.subject = subject
+        self.hemi = hemi
+        self.surf = surf
+        self.offset = offset
+        self.bin_curv = None
+        self.coords = None
+        self.curv = None
+        self.faces = None
+        self.nn = None
+        self.labels = dict()
+        self.x_dir = x_dir
+
+        subjects_dir = str(get_subjects_dir(subjects_dir, raise_error=True))
+        self.data_path = path.join(subjects_dir, subject)
+        if surf == "seghead":
+            raise ValueError(
+                "`surf` cannot be seghead, use "
+                "`mne.viz.Brain.add_head` to plot the seghead"
+            )
+
+    def load_geometry(self):
+        """Load geometry of the surface.
+
+        Parameters
+        ----------
+        None
+
+        Returns
+        -------
+        None
+        """
+        if self.surf == "flat":  # special case
+            fname = path.join(self.data_path, "surf", f"{self.hemi}.cortex.patch.flat")
+            _check_fname(
+                fname, overwrite="read", must_exist=True, name="flatmap surface file"
+            )
+            coords, faces, orig_faces = _read_patch(fname)
+            # rotate 90 degrees to get to a more standard orientation
+            # where X determines the distance between the hemis
+            coords = coords[:, [1, 0, 2]]
+            coords[:, 1] *= -1
+        else:
+            # allow ?h.pial.T1 if ?h.pial doesn't exist for instance
+            # end with '' for better file not found error
+            for img in ("", ".T1", ".T2", ""):
+                surf_fname = path.join(
+                    self.data_path, "surf", f"{self.hemi}.{self.surf}{img}"
+                )
+                if path.isfile(surf_fname):
+                    break
+            coords, faces = read_surface(surf_fname)
+            orig_faces = faces
+        if self.units == "m":
+            coords /= 1000.0
+        if self.offset is not None:
+            x_ = coords @ self.x_dir
+            if self.hemi == "lh":
+                coords -= (np.max(x_) + self.offset) * self.x_dir
+            else:
+                coords -= (np.min(x_) + self.offset) * self.x_dir
+        surf = dict(rr=coords, tris=faces)
+        complete_surface_info(surf, copy=False, verbose=False, do_neighbor_tri=False)
+        nn = surf["nn"]
+        self.coords = coords
+        self.faces = faces
+        self.orig_faces = orig_faces
+        self.nn = nn
+
+    def __len__(self):
+        """Return number of vertices."""
+        return len(self.coords)
+
+    @property
+    def x(self):
+        return self.coords[:, 0]
+
+    @property
+    def y(self):
+        return self.coords[:, 1]
+
+    @property
+    def z(self):
+        return self.coords[:, 2]
+
+    def load_curvature(self):
+        """Load in curvature values from the ?h.curv file."""
+        curv_path = path.join(self.data_path, "surf", f"{self.hemi}.curv")
+        if path.isfile(curv_path):
+            self.curv = read_curvature(curv_path, binary=False)
+            self.bin_curv = np.array(self.curv > 0, np.int64)
+        else:
+            self.curv = None
+            self.bin_curv = None