NF add functions to compute mappers from freesurfer to WebGL (#547)

cortex/utils.py

@@ -7,13 +7,13 @@
 import shutil
 import tarfile
 import tempfile
+import urllib.request
 import warnings
 from distutils.version import LooseVersion
 from importlib import import_module
 
 import h5py
 import numpy as np
-import urllib.request
 
 from . import formats
 from .database import db
@@ -112,35 +112,164 @@ def get_ctmpack(subject, types=("inflated",), method="raw", level=0, recache=Fal
                                overlays_available=overlays_available)
     return ctmfile
 
+
 def get_ctmmap(subject, **kwargs):
-    """
+    """Return a mapping from the vertices in the CTM surface to the vertices
+    in the freesurfer surface. 
+    The mapping is a numpy array, such that `ctm2fs_left[i] = j` means that the
+    i-th vertex in the CTM surface corresponds to the j-th vertex in the freesurfer
+    surface.
+
     Parameters
     ----------
     subject : str
         Subject name
+    kwargs : dict
+        Keyword arguments to pass to get_ctmpack. The most relevant keyword for this
+        function is the `method` kwarg (either `mg2` or `raw`).
 
     Returns
     -------
-    lnew :
+    ctm2fs_left : array (n_vertices_left,)
+        Mapping from CTM vertices to freesurfer vertices for the left hemisphere.
+    ctm2fs_right : array (n_vertices_right,)
+        Mapping from CTM vertices to freesurfer vertices for the right hemisphere.
 
-    rnew :
+    Notes
+    -----
+    This mapping is absolutely necessary when the CTM surfaces are saved with the `mg2`
+    method, which corresponds to storing surfaces with a compressed format.
     """
     from scipy.spatial import cKDTree
 
     from . import brainctm
     jsfile = get_ctmpack(subject, **kwargs)
     ctmfile = os.path.splitext(jsfile)[0]+".ctm"
 
+    # Load freesurfer surfaces
     try:
-        left, right = db.get_surf(subject, "pia")
+        fs_left, fs_right = db.get_surf(subject, "pia")
     except IOError:
-        left, right = db.get_surf(subject, "fiducial")
+        fs_left, fs_right = db.get_surf(subject, "fiducial")
+    # Build a KDTree for each hemisphere based on the freesurfer surfaces
+    lmap, rmap = cKDTree(fs_left[0]), cKDTree(fs_right[0])
+    # Load the CTM surfaces
+    ctm_left, ctm_right = brainctm.read_pack(ctmfile)
+    # Find the closest vertex in the freesurfer surface for each vertex in the CTM
+    # surface. The result is a mapping from CTM vertices to freesurfer vertices.
+    ctm2fs_left = lmap.query(ctm_left[0])[1]
+    ctm2fs_right = rmap.query(ctm_right[0])[1]
+    return ctm2fs_left, ctm2fs_right
+
+
+def get_ctm2webgl_map(subject, **kwargs):
+    """Return a mapping from the vertices in the CTM surface to the vertices visualized
+    on the WebGL viewer.
+    The mapping is a numpy array such that `ctm2webgl_left[i] = j` means that the i-th
+    vertex in the CTM surface corresponds to the j-th vertex in the WebGL viewer.
+
+    Parameters
+    ----------
+    subject : str
+        Subject name
+    kwargs : dict
+        Keyword arguments to pass to get_ctmpack. The most relevant keyword for this
+        function is the `method` kwarg (either `mg2` or `raw`).
+
+    Returns
+    -------
+    ctm2webgl_left : array (n_vertices_left,)
+        Mapping from CTM vertices to WebGL vertices for the left hemisphere.
+    ctm2webgl_right : array (n_vertices_right,)
+        Mapping from CTM vertices to WebGL vertices for the right hemisphere.
+
+    Notes
+    -----
+    This remapping is necessary because surface vertices are re-organized by the WebGL
+    viewer to ensure that neighboring vertices are all stored in the same buffer with
+    maximum length of 65535.
+    """
+    from . import brainctm
+    # Load CTM surfaces
+    jsonfile = get_ctmpack(subject, **kwargs)
+    ctmfile = os.path.splitext(jsonfile)[0] + ".ctm"
+    left_ctm, right_ctm = brainctm.read_pack(ctmfile)
+
+    # The JavaScript code performing the remapping from CTM to WebGL is in
+    # `cortex/webgl/resources/js/ctm/CTMLoader.js`, lines 238--382
+    # The following code is a Python translation of the JavaScript code
+    # Note that the logic for re-indexing "sprawled" faces in not implemented in the
+    # python version. These faces shouldn't occur for cortical surfaces generated by
+    # freesurfer, so this should not be a problem.
+    def _compute_map(coordinates, faces):
+        coordinates = coordinates.ravel()
+        faces = faces.ravel()
+
+        index_map = {}
+        reverse_index_map = {}
+        vertex_counter = 0
+
+        def handle_vertex(vertex, vertex_counter, index_map, reverse_index_map):
+            if vertex not in index_map:
+                index_map[vertex] = vertex_counter
+                reverse_index_map[vertex_counter] = vertex
+                vertex_counter += 1
+            return vertex_counter
+
+        for ii in range(0, len(faces), 3):
+            a = faces[ii]
+            b = faces[ii + 1]
+            c = faces[ii + 2]
+
+            for vtx in [a, b, c]:
+                vertex_counter = handle_vertex(
+                    vtx, vertex_counter, index_map, reverse_index_map
+                )
+        # Make them arrays
+        ctm2webgl = np.array([index_map[ii] for ii in range(len(index_map))])
+        return ctm2webgl
+
+    ctm2webgl_left = _compute_map(left_ctm[0], left_ctm[1])
+    ctm2webgl_right = _compute_map(right_ctm[0], right_ctm[1])
+    return ctm2webgl_left, ctm2webgl_right
+
+
+def get_fs2webgl_map(subject, **kwargs):
+    """Return a mapping from the vertices in the freesurfer surface to the vertices
+    visualized on the WebGL viewer.
+
+    Parameters
+    ----------
+    subject : str
+        Subject name
+    kwargs : dict
+        Keyword arguments to pass to get_ctmpack. The most relevant keyword for this
+        function is the `method` kwarg (either `mg2` or `raw`).
+
+    Returns
+    -------
+    fs2webgl_left : array (n_vertices_left,)
+        Mapping from freesurfer vertices to WebGL vertices for the left hemisphere.
+    fs2webgl_right : array (n_vertices_right,)
+        Mapping from freesurfer vertices to WebGL vertices for the right hemisphere.
+    """
+    ctm2fs_left, ctm2fs_right = get_ctmmap(subject, **kwargs)
+    # Avoid recaching twice
+    if kwargs.get("recache", False):
+        kwargs["recache"] = False
+    ctm2webgl_left, ctm2webgl_right = get_ctm2webgl_map(subject, **kwargs)
+    # Get inverse mapper to go from freesurfer to CTM
+    fs2ctm_left = ctm2fs_left.argsort()
+    fs2ctm_right = ctm2fs_right.argsort()
+    # fs2ctm_left[i] = j means that i-th freesurfer vertex -> j-th ctm vertex
+    # ctm2webgl_left[j] = k means that j-th ctm vertex -> k-th webgl vertex
+    # ctm2webgl_left[fs2ctm_left[i]] = k means that i-th freesurfer vertex -> k-th webgl vertex
+    # Numpy indexing operates as function composition, so we can combine these two
+    # mappings to get the mapping from freesurfer to webgl
+    fs2webgl_left = ctm2webgl_left[fs2ctm_left]
+    fs2webgl_right = ctm2webgl_right[fs2ctm_right]
+    return fs2webgl_left, fs2webgl_right
 
-    lmap, rmap = cKDTree(left[0]), cKDTree(right[0])
-    left, right = brainctm.read_pack(ctmfile)
-    lnew = lmap.query(left[0])[1]
-    rnew = rmap.query(right[0])[1]
-    return lnew, rnew
 
 def get_cortical_mask(subject, xfmname, type='nearest'):
     """Gets the cortical mask for a particular transform