A few small bug fixes

tests/test_local/test_data_loading.py

@@ -3,6 +3,9 @@
 """
 
 
+from vnc_networks.params import SelectionDict
+
+
 class TestDataLoading:
     """
     Test the data loading functions.
@@ -58,3 +61,192 @@ def test_connections_instantiation_MANCv1_2(self):
 
         df_2 = df[(df["start_bid"] == 10725) & (df["end_bid"] == 10439)]
         assert df_2["eff_weight"].values[0] == -1080, "Incorrect nt_type handling"
+
+    def test_connections_getting_neuron_ids_MANCv1_2(self):
+        """
+        Test that we get the same results if we get uids or bodyids and convert between the two.
+        """
+        import vnc_networks
+        from vnc_networks.connections import Connections
+
+        # Instantiate a Connections object
+        connections = Connections(vnc_networks.connectome_reader.MANC("v1.2"))
+
+        # test a few different selection_dicts
+        selection_dicts: list[SelectionDict | None] = [
+            None,
+            {},
+            {"class_1": "descending"},
+            {"class_1": "ascending", "nt_type": "GABA"},
+        ]
+        for selection_dict in selection_dicts:
+            body_ids = connections.get_neuron_bodyids(selection_dict)
+            uids = connections.get_neuron_ids(selection_dict)
+            assert (
+                set(connections.get_uids_from_bodyids(body_ids)) == set(uids)
+            ), f"Getting bodyids and converting to uids doesn't match with selection_dict {selection_dict}."
+            assert (
+                set(connections.get_bodyids_from_uids(uids)) == set(body_ids)
+            ), f"Getting uids and converting to bodyids doesn't match with selection_dict {selection_dict}."
+
+    def test_getting_counts_by_neuropil_MANCv1_2(self):
+        """
+        Test that we can get neuron and synapse counts by neuropil
+        """
+        import pandas as pd
+
+        import vnc_networks
+
+        # Instantiate a Connections object
+        connectome_reader = vnc_networks.connectome_reader.MANC("v1.2")
+
+        # check that this matches what we expect
+        pd.testing.assert_frame_equal(
+            connectome_reader.get_synapse_counts_by_neuropil(
+                "downstream", [10000, 23458]
+            ),
+            pd.DataFrame(
+                {
+                    "body_id": [10000, 23458],
+                    "CV": [703, 0],
+                    "IntTct": [313, 0],
+                    "LTct": [3181, 0],
+                    "LegNp(T3)(R)": [0, 1688],
+                }
+            ),
+        )
+        pd.testing.assert_frame_equal(
+            connectome_reader.get_synapse_counts_by_neuropil(
+                "upstream", [10000, 23458]
+            ),
+            pd.DataFrame(
+                {
+                    "body_id": [10000, 23458],
+                    "CV": [224, 0],
+                    "IntTct": [185, 0],
+                    "LTct": [1462, 0],
+                    "LegNp(T3)(R)": [0, 685],
+                }
+            ),
+        )
+        pd.testing.assert_frame_equal(
+            connectome_reader.get_synapse_counts_by_neuropil("pre", [10000, 23458]),
+            pd.DataFrame(
+                {
+                    "body_id": [10000, 23458],
+                    "CV": [138, 0],
+                    "IntTct": [73, 0],
+                    "LTct": [752, 0],
+                    "LegNp(T3)(R)": [0, 207],
+                }
+            ),
+        )
+        pd.testing.assert_frame_equal(
+            connectome_reader.get_synapse_counts_by_neuropil("post", [10000, 23458]),
+            pd.DataFrame(
+                {
+                    "body_id": [10000, 23458],
+                    "CV": [224, 0],
+                    "IntTct": [185, 0],
+                    "LTct": [1462, 0],
+                    "LegNp(T3)(R)": [0, 685],
+                }
+            ),
+        )
+        pd.testing.assert_frame_equal(
+            connectome_reader.get_synapse_counts_by_neuropil(
+                "total_synapses", [10000, 23458]
+            ),
+            pd.DataFrame(
+                {
+                    "body_id": [10000, 23458],
+                    "CV": [927, 0],
+                    "IntTct": [498, 0],
+                    "LTct": [4643, 0],
+                    "LegNp(T3)(R)": [0, 2373],
+                }
+            ),
+        )
+
+    def test_getting_counts_by_neuropil_FAFBv783(self):
+        """
+        Test that we can get neuron and synapse counts by neuropil
+        """
+        import pandas as pd
+
+        import vnc_networks
+
+        # Instantiate a Connections object
+        connectome_reader = vnc_networks.connectome_reader.FAFB_v783()
+
+        # check that this matches what we expect
+        pd.testing.assert_frame_equal(
+            connectome_reader.get_synapse_counts_by_neuropil(
+                "downstream", [720575940627036426, 720575940633587552]
+            ),
+            pd.DataFrame(
+                {
+                    "body_id": [720575940627036426, 720575940633587552],
+                    "LOP_L": [9, 0],
+                    "LO_L": [2, 0],
+                    "SLP_R": [0, 2],
+                    "SMP_R": [0, 31],
+                }
+            ),
+        )
+        pd.testing.assert_frame_equal(
+            connectome_reader.get_synapse_counts_by_neuropil(
+                "upstream", [720575940627036426, 720575940633587552]
+            ),
+            pd.DataFrame(
+                {
+                    "body_id": [720575940627036426, 720575940633587552],
+                    "LOP_L": [3, 0],
+                    "LO_L": [9, 0],
+                    "SLP_R": [0, 11],
+                    "SMP_R": [0, 13],
+                }
+            ),
+        )
+        pd.testing.assert_frame_equal(
+            connectome_reader.get_synapse_counts_by_neuropil(
+                "pre", [720575940627036426, 720575940633587552]
+            ),
+            pd.DataFrame(
+                {
+                    "body_id": [720575940627036426, 720575940633587552],
+                    "LOP_L": [14, 0],
+                    "LO_L": [83, 0],
+                    "SLP_R": [0, 33],
+                    "SMP_R": [0, 75],
+                }
+            ),
+        )
+        pd.testing.assert_frame_equal(
+            connectome_reader.get_synapse_counts_by_neuropil(
+                "post", [720575940627036426, 720575940633587552]
+            ),
+            pd.DataFrame(
+                {
+                    "body_id": [720575940627036426, 720575940633587552],
+                    "LOP_L": [100, 0],
+                    "LO_L": [8, 0],
+                    "SLP_R": [0, 4],
+                    "SMP_R": [0, 284],
+                }
+            ),
+        )
+        pd.testing.assert_frame_equal(
+            connectome_reader.get_synapse_counts_by_neuropil(
+                "total_synapses", [720575940627036426, 720575940633587552]
+            ),
+            pd.DataFrame(
+                {
+                    "body_id": [720575940627036426, 720575940633587552],
+                    "LOP_L": [114, 0],
+                    "LO_L": [91, 0],
+                    "SLP_R": [0, 37],
+                    "SMP_R": [0, 359],
+                }
+            ),
+        )

vnc_networks/connections.py

@@ -27,6 +27,7 @@
 """
 
 import copy
+import itertools
 import os
 import pickle
 import typing
@@ -93,7 +94,6 @@ def __init__(
         not_connected: Optional[list[BodyId] | list[int]] = None,
         keep_only_traced_neurons: bool = True,
     ): ...
-
     def __init__(
         self,
         CR: ConnectomeReader = MANC("v1.2"),
@@ -411,18 +411,18 @@ def __map_uid(self):
         Map the tuples (bodyId, subdivision) to a unique identifier (uid) number that will
         be used to reference the newly defined neurons, e.g. as graph node ids.
         """
+        # faster to take set of combined lists than to take union of sets
         unique_objects = list(
             set(
-                zip(
-                    self.connections["start_bid"],
-                    self.connections["subdivision_start"],
-                )
-            ).union(
-                set(
+                itertools.chain(
+                    zip(
+                        self.connections["start_bid"],
+                        self.connections["subdivision_start"],
+                    ),
                     zip(
                         self.connections["end_bid"],
                         self.connections["subdivision_end"],
-                    )
+                    ),
                 )
             )
         )
@@ -436,48 +436,34 @@ def __map_uid(self):
                 "subdivision": [x[1] for x in unique_objects],
             }
         )
-        # add uids to the connections table
-        self.connections["start_uid"] = self.__convert_neuron_ids_to_uid(
-            self.connections[["start_bid", "subdivision_start"]].values,
-            input_type="table",
+
+        # add uids to the connections table using pandas join
+        # take UID for start_bid and subdivision_start -> rename to start_uid
+        self.connections = self.connections.join(
+            self.uid[["body_id", "subdivision", "uid"]].set_index(
+                ["body_id", "subdivision"]
+            ),
+            on=["start_bid", "subdivision_start"],
+            how="left",
         )
-        self.connections["end_uid"] = self.__convert_neuron_ids_to_uid(
-            self.connections[["end_bid", "subdivision_end"]].values,
-            input_type="table",
+        # renaming the columns in place is much faster
+        # otherwise it creates a copy of the dataframe
+        self.connections.columns = self.connections.columns.to_series().replace(
+            {"uid": "start_uid"}
+        )
+        # take UID for end_bid and subdivision_end -> rename to end_bid
+        self.connections = self.connections.join(
+            self.uid[["body_id", "subdivision", "uid"]].set_index(
+                ["body_id", "subdivision"]
+            ),
+            on=["end_bid", "subdivision_end"],
+            how="left",
+        )
+        self.connections.columns = self.connections.columns.to_series().replace(
+            {"uid": "end_uid"}
         )
-        return
-
-    def __convert_neuron_ids_to_uid(
-        self,
-        neuron_ids,
-        input_type: str = "tuple",
-    ):
-        """
-        Convert a list of neuron ids to a list of unique identifiers.
 
-        Parameters
-        ----------
-        neuron_ids: list[tuple[int]] or dataframe['id','subdivision']
-            List of neuron ids to convert.
-        input_type: str
-            Type of the input list.
-            Can be 'tuple', 'table'
-        """
-        if input_type == "tuple":
-            tuple_of_ids = neuron_ids
-        elif input_type == "table":
-            tuple_of_ids = [tuple(row) for row in neuron_ids]
-        else:
-            raise ValueError(
-                f"Class Connections \
-                ::: > convert_neuron_ids_to_uid(): Unknown input type {input_type}"
-            )
-        uid_table = self.uid.copy()
-        uid_table = uid_table[uid_table["neuron_ids"].isin(tuple_of_ids)]
-        # sort the table by the order of the input list
-        uid_table = uid_table.set_index("neuron_ids").loc[tuple_of_ids].reset_index()
-        uids = uid_table["uid"].to_list()
-        return uids
+        return
 
     @typing.overload
     def __convert_uid_to_neuron_ids(
@@ -977,8 +963,9 @@ def get_neuron_ids(
         Get the neuron IDs from the nodes dataframe as loaded in the initial
         dataset, based on a selection dictionary.
         """
-        nodes = self.get_nodes(type="body_id")
-        body_ids = self.CR.get_neuron_bodyids(selection_dict, nodes)
+        body_ids = self.get_nodes(type="body_id")
+        if selection_dict is not None:
+            body_ids = self.CR.get_neuron_bodyids(selection_dict, body_ids)
         return self.__get_uids_from_bodyids(body_ids)
 
     def get_neurons_pre(self):
@@ -1398,6 +1385,38 @@ def get_in_degree(self, uid: UID | int):
         """
         return self.graph.in_degree(uid)
 
+    def get_downstream_neuron_counts_by_neuropil(self, body_id: BodyId | int):
+        """
+        Get how many downstream neurons this neuron has in each neuropil.
+
+        Returns a pandas dataframe where the columns (apart from `body_id`) are the neuropils.
+        """
+        return self.CR.get_synapse_counts_by_neuropil("downstream", [body_id])
+
+    def get_upstream_neuron_counts_by_neuropil(self, body_id: BodyId | int):
+        """
+        Get how many upstream neurons this neuron has in each neuropil.
+
+        Returns a pandas dataframe where the columns (apart from `body_id`) are the neuropils.
+        """
+        return self.CR.get_synapse_counts_by_neuropil("upstream", [body_id])
+
+    def get_downstream_synapse_counts_by_neuropil(self, body_id: BodyId | int):
+        """
+        Get how many downstream synapses this neuron has in each neuropil.
+
+        Returns a pandas dataframe where the columns (apart from `body_id`) are the neuropils.
+        """
+        return self.CR.get_synapse_counts_by_neuropil("post", [body_id])
+
+    def get_upstream_synapse_counts_by_neuropil(self, body_id: BodyId | int):
+        """
+        Get how many upstream synapses this neuron has in each neuropil.
+
+        Returns a pandas dataframe where the columns (apart from `body_id`) are the neuropils.
+        """
+        return self.CR.get_synapse_counts_by_neuropil("pre", [body_id])
+
     # --- setters
     def merge_nodes(self, nodes: list[UID] | list[int], combination_logic: str = "sum"):
         """