add edges to MeshNeuron

[fcollman]

Here's a first crack of a potential addition of 'edges' to mesh neuron.

[fcollman]

This would be one way to approach #169

[fcollman]

here i was able to use this pull request to pull in the link edges using the existing trimesh_repair functionality in meshparty..
seems like tests are failing though due to the trimesh.edges property overloading...

import bpy
import bmesh
import navis
import caveclient
import navis.interfaces.blender as b3d
import numpy as np
import time
from scipy import spatial, sparse
from meshparty import trimesh_io

navis.patch_cloudvolume()

client = caveclient.CAVEclient("minnie65_phase3_v1")

nuc_id = 305046
nuc_df = client.materialize.tables.nucleus_detection_v0(id=nuc_id).query(
    desired_resolution=[1, 1, 1], timestamp="now"
)
# nuc_df = nuc_df.query(f"id=={nuc_id}")
root_id = nuc_df.iloc[0].pt_root_id
soma_pt = nuc_df.iloc[0].pt_position / 1000

vol = client.info.segmentation_cloudvolume()

cv_mesh = vol.mesh.get([root_id], as_navis=False, process=False)[root_id]
mp_mesh = trimesh_io.Mesh(vertices=cv_mesh.vertices, faces=cv_mesh.faces)
merge_log = client.chunkedgraph.get_merge_log(root_id)

mp_mesh.add_link_edges(
    root_id, merge_log, base_resolution=client.chunkedgraph.base_resolution
)


nav_mesh = navis.core.MeshNeuron(
    {
        "vertices": mp_mesh.vertices,
        "faces": mp_mesh.faces,
        "edges": mp_mesh.link_edges,
    },
    id=root_id,
    process=False,
)

scaling = 1 / 10000
h = b3d.Handler(scaling=scaling)
h.add(nav_mesh)

[schlegelp]

Having looked at both navis and skeletor:

Functions operating on MeshNeurons typically go straight for the MeshNeuron.trimesh representation for anything more complicated than the vertices and faces - unique edges, edge lengths, vertex adjacency and so on. We could adopt those properties for MeshNeuron but I think that would get messy pretty quickly.

The cleanest solution would be to create our own Trimesh subclass that understands the concept of non-face edges. Fortunately, that turns out to not be terribly complicated because you really only have to modify Trimesh.edges and the rest just follows suit. Here is a prototype:

# new module (e.g. navis/utils/subclasses.py) to avoid circular imports
import trimesh as tm
import numpy as np


class TrimeshPlus(tm.Trimesh):
    """Trimesh object with additional features.

    Currently, this includes:
      - `extra_edges` property: edges not present in the faces
      - `add_extra_edges()` method: add / replace extra edges
     
    """
    def __repr__(self):
        s = super().__repr__()
        # Add extra edges to the representation
        if self.extra_edges is not None:
            s = s[:-2]  # remove last bracket
            s += f", extra_edges.shape={self.extra_edges.shape})>"
        return s

    @property
    def extra_edges(self):
        """Edges not present in the faces."""
        if hasattr(self, "_extra_edges"):
            return self._extra_edges

    @extra_edges.setter
    def extra_edges(self, edges):
        edges = np.asarray(edges).astype(np.int64, copy=False)

        if edges.ndim != 2 or edges.shape[1] != 2:
            raise ValueError("Edges must be a 2D array with 2 columns.")

        self._extra_edges = edges

        # Clear the cache
        self._cache.clear()

    @property
    def edges(self):
        edges = super().edges
        if hasattr(self, "_extra_edges"):
            return np.vstack([edges, self._extra_edges])
        return edges

    def add_extra_edges(self, edges, validate=True, replace=False):
        """Add non-face edges to the mesh.

        You can directly set the `extra_edges` property, but this function
        does some additional checks such as ensuring that the edges are
        unique and not already present among the faces.

        Parameters
        ----------
        edges :     (N, 2) iterable of vertex indices | (2, ) array for a single edge
                    Edges to add. Note that orientation does not matter.
        validate :  bool
                    Whether to validate the edges. If True, will silently drop edges
                    that edges are not unique or already present in the faces.
        replace :   bool
                    Whether to add to or replace any existing extra edges.

        """
        edges = np.asarray(edges).astype(np.int64, copy=False)

        # Reshape single edge
        if edges.shape == (2,):
            edges = edges[np.newaxis, :]

        if edges.ndim != 2 or edges.shape[1] != 2:
            raise ValueError("Edges must be a 2D array with 2 columns.")

        if not replace and self.extra_edges is not None:
            edges = np.vstack([self.extra_edges, edges])

        # Sort along the first axis (orientation does not matter)
        edges = np.sort(edges, axis=1)

        if validate:
            # Drop duplicates
            edges = np.unique(edges, axis=0)

            # Drop edges that are already in faces
            face_edges = np.sort(super().edges, axis=1)  # edges from faces (exclude existing extra edges)

            # We'll be clever about checking for edges by using a two-step process:
            # First check if the first vertex is present in the faces
            duplicated = np.isin(edges[:, 0], face_edges[:, 0])

            if np.any(duplicated):
                # Turn potential matching face edges into tuples
                to_check = [tuple(e) for e in face_edges[np.isin(face_edges[:, 0], edges[:, 0])]]
                # Check if the second vertex is present in the faces
                duplicated[duplicated] = [tuple(e) in to_check for e in edges[duplicated]]

                if any(duplicated):
                    edges = edges[~duplicated]

        self._extra_edges = edges

        # Clear the cache
        self._cache.clear()

    def copy(self, *args, **kwargs):
        # Call the original copy method
        copy = super().copy(*args, **kwargs)
        # Add the extra_edges attribute
        if hasattr(self, "_extra_edges"):
            copy.extra_edges = self.extra_edges.copy()

I haven't tested this exhaustively and I'm sure we're breaking parts of trimesh by having edges that aren't part of faces but this feels like a workable solution. In theory, all that's needed now is this:

# navis/core/mesh.py

class MeshNeuron(BaseNeuron):
     ...
    @property
    @temp_property
    def trimesh(self):
        """Trimesh representation of the neuron."""
        if not getattr(self, "_trimesh", None):
            if hasattr(self, "extra_edges"):
                # Only use TrimeshPlus if we actually need it 
                # to avoid unnecessarily breaking stuff elsewhere
                self._trimesh = TrimeshPlus(
                    vertices=self._vertices, faces=self._faces, process=False
                )
                self._trimesh.extra_edges = self.extra_edges
            else:
                self._trimesh = tm.Trimesh(
                    vertices=self._vertices, faces=self._faces, process=False
                )
        return self._trimesh

There will be more bits that need adjusting but this would in theory already take care of the majority of downstream functions. For example, neuron2nx should just work without any adjustment.

[fcollman]

looks good to me!

So then the next question is whether to port a version of the workflow I had in trimesh_repair.py into navis..

essentially the problem boils down to taking these points in space which are the merge operations and finding the closest mesh vertex, taking the components that are associated with those vertices and if they are not the same component locally, then finding the mutually closest vertices, adding temp edges between them and then finding the shortest path between the previously disconnected pieces that uses those edges, and then finding the temp edge that was used along that path... this is the edge you want to add as an "extra-edge" for that merge event.

[schlegelp]

So then the next question is whether to port a version of the workflow I had in trimesh_repair.py into navis.

Yeah, good question! I need to take a closer look at the trimesh_repair module but would there be a natural way to split that workflow into a general and a CAVE-specific part?

If there is a generalizable part - perhaps something a long the lines of "add extra edges from graph" - that could go into vanilla navis straight away.

For the more specialised CAVE-part, I can think of a few potential homes:

1. In a new interface - something like navis.interfaces.cave
2. In a separate package - navis-cave or cave-navis?
3. In meshparty

I don't have any strong preference beyond wanting it to exist :)

[fcollman]

Yes.. i think the interface in is a list of paired 3d points that could come from anywhere, and it's job is to produce the right set of extra_edges that reflect those 'merge' operations. CAVE can be one such source, so could a neuroglancer state with some annotations.