Nsyn anno (#20)

* first draft of new synapse count function

* use net size instead of mean

* add skeleton aggregation function

* add test and refactor

* bumpversion

pcg_skel/__init__.py

@@ -4,4 +4,4 @@
 from . import utils
 from . import chunk_tools
 
-__version__ = "1.0.7"
+__version__ = "1.1.0"

pcg_skel/features.py

@@ -6,12 +6,111 @@
 from meshparty import meshwork
 from meshparty.meshwork.algorithms import split_axon_by_annotation, strahler_order
 from warnings import warn
+from typing import Optional
 
 from . import pcg_anno
 
 VOL_PROPERTIES = ["area_nm2", "size_nm3", "mean_dt_nm", "max_dt_nm"]
 
 
+def add_synapse_count(
+    nrn: meshwork.Meshwork,
+    anno_name: str = "synapse_count",
+    pre_syn: str = "pre_syn",
+    post_syn: str = "post_syn",
+    pre_syn_index: str = "pre_pt_mesh_ind",
+    post_syn_index: str = "post_pt_mesh_ind",
+    aggregate_size: bool = True,
+):
+    """
+    Create a synapse count label across mesh vertices from existing synapse annotation.
+
+    This function adds a synapse count annotation to the mesh vertices of a neuron. It aggregates
+    synapse data from pre-synaptic and post-synaptic annotations and optionally calculates the mean
+    size of the synapses.
+
+    Parameters
+    ----------
+    nrn : meshwork.Meshwork
+        The neuron meshwork object containing the mesh and annotations.
+    anno_name : str, optional
+        The name of the annotation table to be created, by default "synapse_count".
+    pre_syn : str, optional
+        The name of the pre-synaptic annotation table, by default "pre_syn".
+        If set to None, presynaptic points will be skipped.
+    post_syn : str, optional
+        The name of the post-synaptic annotation table, by default "post_syn".
+        If set to None, postsynaptic points will be skipped
+    pre_syn_index : str, optional
+        The mesh index column for pre-synaptic annotation, by default "pre_pt_mesh_ind".
+    post_syn_index : str, optional
+        The mesh index column for post-synaptic annotation, by default "post_pt_mesh_ind".
+    aggregate_size : bool, optional
+        Whether to aggregate the size of synapses, by default True.
+
+    Returns
+    -------
+    None
+        The function modifies the neuron meshwork object in place by adding the synapse count annotation.
+    """
+
+    with nrn.mask_context(np.ones(nrn.mesh.unmasked_size).astype(bool)):
+        syn_count_df = pd.DataFrame(index=nrn.mesh_indices)
+
+        if pre_syn:
+            syn_out_agg_kwargs = {
+                "num_syn_out": pd.NamedAgg("id", "count"),
+            }
+            if aggregate_size:
+                syn_out_agg_kwargs["net_size_out"] = pd.NamedAgg("size", "sum")
+
+            syn_out_df = (
+                nrn.anno[pre_syn].df.groupby(pre_syn_index).agg(**syn_out_agg_kwargs)
+            )
+
+        if post_syn:
+            syn_in_agg_kwargs = {
+                "num_syn_in": pd.NamedAgg("id", "count"),
+            }
+            if aggregate_size:
+                syn_in_agg_kwargs["net_size_in"] = pd.NamedAgg("size", "sum")
+
+            syn_in_df = (
+                nrn.anno[post_syn].df.groupby(post_syn_index).agg(**syn_in_agg_kwargs)
+            )
+
+        syn_count_df = (
+            syn_count_df.merge(
+                syn_out_df,
+                left_index=True,
+                right_index=True,
+                how="left",
+            )
+            .merge(
+                syn_in_df,
+                left_index=True,
+                right_index=True,
+                how="left",
+            )
+            .fillna(0)
+        )
+
+        if pre_syn:
+            syn_count_df["num_syn_out"] = syn_count_df["num_syn_out"].astype(int)
+        if post_syn:
+            syn_count_df["num_syn_in"] = syn_count_df["num_syn_in"].astype(int)
+        syn_count_df = syn_count_df.reset_index().rename(
+            columns={"index": "mesh_index"}
+        )
+
+        nrn.anno.add_annotations(
+            anno_name,
+            syn_count_df,
+            index_column="mesh_index",
+        )
+    pass
+
+
 def add_synapses(
     nrn: meshwork.Meshwork,
     synapse_table: str,
@@ -356,3 +455,41 @@ def l2dict_from_anno(
     mesh_ind_col: str = "mesh_ind",
 ) -> dict:
     return nrn.anno[table_name].df.set_index(l2id_col)[mesh_ind_col].to_dict()
+
+
+def aggregate_property_to_skeleton(
+    nrn: meshwork.Meshwork,
+    anno_table: str,
+    agg_dict: Optional[dict] = None,
+    fill_value: float = 0,
+) -> pd.DataFrame:
+    """Aggregate a meshwork annotation table to skeletons
+
+    Parameters
+    ----------
+    nrn : meshwork.Meshwork
+        Meshwork object with mesh,skeleton, and annotations
+    anno_table : str
+        Name of the annotation table
+    agg_dict : dict
+        Dictionary of column names and aggregation functions that can be used by pandas pd.NamedAgg.
+    fill_value : float, optional
+        Value to use to fill missing or NaN values, by default 0
+
+    Returns
+    -------
+    pd.DataFrame
+        Dataframe with one row per skeleton index and columns aggregating mesh annotation values
+        using the aggregation function specified in the aggregation_dict.
+    """
+    df = nrn.anno[anno_table].df
+
+    skel_index_column = "skel_index"
+    while skel_index_column in df.columns:
+        skel_index_column = f"{skel_index_column}_temp"
+    df[skel_index_column] = nrn.anno[anno_table].mesh_index.to_skel_index_padded
+
+    aggs = {k: pd.NamedAgg(k, v) for k, v in agg_dict.items()}
+    skel_df = df.groupby(skel_index_column).agg(**aggs)
+    skel_df = skel_df.reindex(np.arange(0, nrn.skeleton.n_vertices)).fillna(fill_value)
+    return skel_df

pyproject.toml

@@ -16,7 +16,7 @@ maintainers = [
 ]
 requires-python = ">=3.9"
 dependencies = [
-    "caveclient>=4.12.4",
+    "caveclient>=6.4.1",
     "cloud-volume>=3.6.0",
     "fastremap",
     "meshparty>=1.12.0",
@@ -40,17 +40,19 @@ include = [
 ]
 
 [tool.hatch.envs.default]
+python = "3.12"
 installer = "uv"
 
 [tool.hatch.envs.test]
 dependencies = [
     "pytest",
     "pytest-mock",
     "pytest-cov",
+    "responses",
 ]
 
 [[tool.hatch.envs.test.matrix]]
-python = ["3.9", "3.10", "3.11"]
+python = ["3.9", "3.10", "3.11", "3.12"]
 
 [tool.hatch.envs.test.scripts]
 run_tests = "python -m pytest -v tests"