napistu 0.4.3__py3-none-any.whl → 0.4.5__py3-none-any.whl

napistu/network/neighborhoods.py

@@ -34,6 +34,7 @@ def find_and_prune_neighborhoods(
     napistu_graph: ig.Graph,
     compartmentalized_species: str | list[str],
     precomputed_distances: pd.DataFrame | None = None,
+    source_total_counts: pd.Series | None = None,
     network_type: str = NEIGHBORHOOD_NETWORK_TYPES.DOWNSTREAM,
     order: int = 3,
     verbose: bool = True,
@@ -44,7 +45,7 @@ def find_and_prune_neighborhoods(
 
     Wrapper which combines find_neighborhoods() and prune_neighborhoods()
 
-     Parameters
+    Parameters
     ----------
     sbml_dfs: sbml_dfs_core.SBML_dfs
         A mechanistic molecular model
@@ -54,6 +55,9 @@ def find_and_prune_neighborhoods(
         Compartmentalized species IDs for neighborhood centers
     precomputed_distances : pd.DataFrame or None
         If provided, an edgelist of origin->destination path weights and lengths
+    source_total_counts: pd.Series | None
+        Optional, A series of the total counts of each source. As produced by
+        source.get_source_total_counts()
     network_type: str
         If the network is directed should neighbors be located "downstream",
         or "upstream" of each compartmentalized species. The "hourglass" option
@@ -109,6 +113,7 @@ def find_and_prune_neighborhoods(
         order=order,
         verbose=verbose,
         precomputed_neighbors=precomputed_neighbors,
+        source_total_counts=source_total_counts,
     )
 
     pruned_neighborhoods = prune_neighborhoods(neighborhoods, top_n=top_n)
@@ -132,7 +137,7 @@ def load_neighborhoods(
 
     Load existing neighborhoods if they exist
     (and overwrite = False) and otherwise construct
-     neighborhoods using the provided settings
+    neighborhoods using the provided settings
 
     Parameters
     ----------
@@ -507,14 +512,16 @@ def find_neighborhoods(
     compartmentalized_species: list[str],
     network_type: str = "downstream",
     order: int = 3,
-    verbose: bool = True,
+    min_pw_size: int = 3,
     precomputed_neighbors: pd.DataFrame | None = None,
+    source_total_counts: pd.Series | None = None,
+    verbose: bool = True,
 ) -> dict:
     """
     Find Neighborhood
 
     Create a network composed of all species and reactions within N steps of
-      each of a set of compartmentalized species.
+    each of a set of compartmentalized species.
 
     Parameters
     ----------
@@ -530,16 +537,21 @@ def find_neighborhoods(
         locates both upstream and downstream species.
     order: int
         Max steps away from center node
-    verbose: bool
-        Extra reporting
     precomputed_neighbors: pd.DataFrame or None
         If provided, a pre-filtered table of nodes nearby the compartmentalized species
         which will be used to skip on-the-fly neighborhood generation.
+    min_pw_size: int
+        the minimum size of a pathway to be considered
+    source_total_counts: pd.Series | None
+        Optional, A series of the total counts of each source. As produced by
+        source.get_source_total_counts()\
+    verbose: bool
+        Extra reporting
 
     Returns:
     ----------
     A dict containing the neighborhood of each compartmentalized species.
-      Each entry in the dict is a dict of the subgraph, vertices, and edges.
+    Each entry in the dict is a dict of the subgraph, vertices, and edges.
     """
 
     if not isinstance(network_type, str):
@@ -567,7 +579,13 @@ def find_neighborhoods(
     # format the vertices and edges in each compartmentalized species' network
     neighborhood_dict = {
         sc_id: create_neighborhood_dict_entry(
-            sc_id, neighborhood_df, sbml_dfs, napistu_graph, verbose=verbose
+            sc_id,
+            neighborhood_df=neighborhood_df,
+            sbml_dfs=sbml_dfs,
+            napistu_graph=napistu_graph,
+            min_pw_size=min_pw_size,
+            source_total_counts=source_total_counts,
+            verbose=verbose,
         )
         for sc_id in compartmentalized_species
     }
@@ -580,6 +598,8 @@ def create_neighborhood_dict_entry(
     neighborhood_df: pd.DataFrame,
     sbml_dfs: sbml_dfs_core.SBML_dfs,
     napistu_graph: ig.Graph,
+    min_pw_size: int = 3,
+    source_total_counts: pd.Series | None = None,
     verbose: bool = False,
 ) -> dict[str, Any]:
     """
@@ -597,6 +617,11 @@ def create_neighborhood_dict_entry(
         A mechanistic molecular model
     napistu_graph: igraph.Graph
         A network connecting molecular species and reactions
+    min_pw_size: int
+        the minimum size of a pathway to be considered
+    source_total_counts: pd.Series
+        Optional, A series of the total counts of each source. As produced by
+        source.get_source_total_counts()
     verbose: bool
         Extra reporting?
 
@@ -645,7 +670,11 @@ def create_neighborhood_dict_entry(
 
     try:
         edge_sources = ng_utils.get_minimal_sources_edges(
-            vertices.rename(columns={"name": "node"}), sbml_dfs
+            vertices.rename(columns={"name": "node"}),
+            sbml_dfs,
+            min_pw_size=min_pw_size,
+            # optional, counts of sources across the whole model
+            source_total_counts=source_total_counts,
         )
     except Exception:
         edge_sources = None
@@ -1441,7 +1470,7 @@ def _prune_vertex_set(one_neighborhood: dict, top_n: int) -> pd.DataFrame:
     ----------
     one_neighborhood: dict
         The neighborhood around a single compartmentalized species - one of the values
-         in dict created by find_neighborhoods().
+        in dict created by find_neighborhoods().
     top_n: int
         How many neighboring molecular species should be retained?
         If the neighborhood includes both upstream and downstream connections

napistu/network/ng_utils.py

@@ -114,6 +114,7 @@ def compartmentalize_species_pairs(
 def get_minimal_sources_edges(
     vertices: pd.DataFrame,
     sbml_dfs: sbml_dfs_core.SBML_dfs,
+    min_pw_size: int = 3,
     source_total_counts: Optional[pd.Series] = None,
 ) -> pd.DataFrame | None:
     """
@@ -125,6 +126,8 @@ def get_minimal_sources_edges(
         A table of vertices.
     sbml_dfs: sbml_dfs_core.SBML_dfs
         A pathway model
+    min_pw_size: int
+        the minimum size of a pathway to be considered
     source_total_counts: pd.Series
         A series of the total counts of each source.
 
@@ -146,7 +149,7 @@ def get_minimal_sources_edges(
         return None
     else:
         edge_sources = source.source_set_coverage(
-            source_df, source_total_counts, sbml_dfs
+            source_df, source_total_counts, sbml_dfs, min_pw_size=min_pw_size
         )
         return edge_sources.reset_index()[
             [SBML_DFS.R_ID, SOURCE_SPEC.PATHWAY_ID, SOURCE_SPEC.NAME]

napistu/network/paths.py

@@ -241,6 +241,8 @@ def find_all_shortest_reaction_paths(
     target_species_paths: pd.DataFrame,
     weight_var: str = NAPISTU_GRAPH_EDGES.WEIGHTS,
     precomputed_distances: pd.DataFrame | None = None,
+    min_pw_size: int = 3,
+    source_total_counts: pd.Series | None = None,
 ):
     """
     Shortest Reaction Paths
@@ -259,6 +261,11 @@ def find_all_shortest_reaction_paths(
         An edge attribute to use when forming a weighted shortest path
     precomputed_distances : pd.DataFrame | None
         A table containing precalculated path summaries between pairs of compartmentalized species
+    min_pw_size : int
+        the minimum size of a pathway to be considered
+    source_total_counts : pd.Series | None
+        A series of the total counts of each source. As produced by
+        source.get_source_total_counts()
 
     Returns:
     ----------
@@ -325,7 +332,12 @@ def find_all_shortest_reaction_paths(
     ).reset_index()
 
     # at a minimal set of pathway sources to organize reactions
-    edge_sources = get_minimal_sources_edges(all_shortest_reaction_paths_df, sbml_dfs)
+    edge_sources = get_minimal_sources_edges(
+        all_shortest_reaction_paths_df,
+        sbml_dfs,
+        min_pw_size=min_pw_size,
+        source_total_counts=source_total_counts,
+    )
 
     # create a new small network of shortest paths
     unique_path_nodes = (

napistu/source.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+import logging
 
 import numpy as np
 import pandas as pd
@@ -10,6 +11,8 @@ from napistu import sbml_dfs_utils
 from napistu.statistics import hypothesis_testing
 from napistu.constants import SBML_DFS_SCHEMA, SCHEMA_DEFS, SOURCE_SPEC
 
+logger = logging.getLogger(__name__)
+
 
 class Source:
     """
@@ -244,7 +247,7 @@ def unnest_sources(source_table: pd.DataFrame, verbose: bool = False) -> pd.Data
 
     for i in range(source_table.shape[0]):
         if verbose:
-            print(f"Processing {source_table_index.index.values[i]}")
+            logger.info(f"Processing {source_table_index.index.values[i]}")
 
         # check that the entries of sourcevar are Source objects
         source_value = source_table[source_var].iloc[i]
@@ -255,7 +258,7 @@ def unnest_sources(source_table: pd.DataFrame, verbose: bool = False) -> pd.Data
             )
 
         if source_value.source is None:
-            print("Some sources were only missing - returning None")
+            logger.warning("Some sources were only missing - returning None")
             return None
 
         source_tbl = pd.DataFrame(source_value.source)
@@ -278,6 +281,7 @@ def source_set_coverage(
     select_sources_df: pd.DataFrame,
     source_total_counts: Optional[pd.Series] = None,
     sbml_dfs: Optional[sbml_dfs_core.SBML_dfs] = None,
+    min_pw_size: int = 3,
 ) -> pd.DataFrame:
     """
     Greedy Set Coverage of Sources
@@ -298,6 +302,8 @@ def source_set_coverage(
     sbml_dfs: sbml_dfs_core.SBML_dfs
         if `source_total_counts` is provided then `sbml_dfs` must be provided
         to calculate the total number of entities in the table.
+    min_pw_size: int
+        the minimum size of a pathway to be considered
 
     Returns
     -------
@@ -325,10 +331,16 @@ def source_set_coverage(
         # find the pathway with the most members
 
         if source_total_counts is None:
-            top_pathway = _select_top_pathway_by_size(unaccounted_for_members)
+            top_pathway = _select_top_pathway_by_size(
+                unaccounted_for_members, min_pw_size=min_pw_size
+            )
         else:
             top_pathway = _select_top_pathway_by_enrichment(
-                unaccounted_for_members, source_total_counts, n_total_entities, pk
+                unaccounted_for_members,
+                source_total_counts,
+                n_total_entities,
+                pk,
+                min_pw_size=min_pw_size,
             )
 
         if top_pathway is None:
@@ -368,6 +380,13 @@ def get_source_total_counts(
     """
 
     all_sources_table = unnest_sources(sbml_dfs.get_table(entity_type))
+
+    if all_sources_table is None:
+        logger.warning(
+            f"No sources found for {entity_type} in sbml_dfs. Returning an empty series."
+        )
+        return pd.Series([], name="total_counts")
+
     source_total_counts = all_sources_table.value_counts(SOURCE_SPEC.PATHWAY_ID).rename(
         "total_counts"
     )
@@ -515,9 +534,15 @@ def _safe_source_merge(member_Sources: Source | list) -> Source:
         raise TypeError("Expecting source.Source or pd.Series")
 
 
-def _select_top_pathway_by_size(unaccounted_for_members: pd.DataFrame) -> str:
+def _select_top_pathway_by_size(
+    unaccounted_for_members: pd.DataFrame, min_pw_size: int = 3
+) -> str:
 
     pathway_members = unaccounted_for_members.value_counts(SOURCE_SPEC.PATHWAY_ID)
+    pathway_members = pathway_members.loc[pathway_members >= min_pw_size]
+    if pathway_members.shape[0] == 0:
+        return None
+
     top_pathway = pathway_members[pathway_members == max(pathway_members)].index[0]
 
     return top_pathway
@@ -528,7 +553,7 @@ def _select_top_pathway_by_enrichment(
     source_total_counts: pd.Series,
     n_total_entities: int,
     table_pk: str,
-    min_pw_size: int = 5,
+    min_pw_size: int = 3,
 ) -> str:
 
     n_observed_entities = len(

{napistu-0.4.3.dist-info → napistu-0.4.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: napistu
-Version: 0.4.3
+Version: 0.4.5
 Summary: Connecting high-dimensional data to curated pathways
 Home-page: https://github.com/napistu/napistu-py
 Author: Sean Hackett

{napistu-0.4.3.dist-info → napistu-0.4.5.dist-info}/RECORD

@@ -6,7 +6,7 @@ napistu/identifiers.py,sha256=e2-nTVzr5AINa0y1ER9218bKXyF2kAeJ9At22S4Z00o,33914
 napistu/indices.py,sha256=Zjg3gE0JQ3T879lCPazYg-WXVE6hvcAr713ZKpJ32rk,9830
 napistu/sbml_dfs_core.py,sha256=s0OyoHs-AjOcbZu1d3KNkW_PI7Rxbhu5ZLpfQeO4iY8,72639
 napistu/sbml_dfs_utils.py,sha256=SOy1Ii2hDFOfQa7pFAJS9EfAmfBVD_sHvDJBVmCN_p8,46456
-napistu/source.py,sha256=iDDKpN-4k_W_tyxEjqe_z-yPJv7uoFRRBhkiBtOH5C8,20416
+napistu/source.py,sha256=iUB0SqzHW5qe0IMfnWvUCfNpjYpbXDv0s2pHNgZ8BFc,21102
 napistu/utils.py,sha256=p2sJxTklmV30XS6hanJRjcdfgeaZpkULuMyQX3BPP0c,36404
 napistu/context/__init__.py,sha256=LQBEqipcHKK0E5UlDEg1ct-ymCs93IlUrUaH8BCevf0,242
 napistu/context/discretize.py,sha256=Qq7zg46F_I-PvQIT2_pEDQV7YEtUQCxKoRvT5Gu9QsE,15052
@@ -62,13 +62,13 @@ napistu/network/__init__.py,sha256=dFXAhIqlTLJMwowS4BUDT08-Vy3Q0u1L0CMCErSZT1Y,2
 napistu/network/constants.py,sha256=nG_lUZYLgop8oxOGjDYqvxXJzVdOwKZ3aWnxlhtSaIo,6915
 napistu/network/data_handling.py,sha256=KncrAKjXI3169BgVE-SnY8FkpVF60JnUwfMHtbqvsTc,14725
 napistu/network/ig_utils.py,sha256=MuyEyOVtSHndil6QuuRCimBZrJ2jTaF5qQESgYlu02M,17042
-napistu/network/neighborhoods.py,sha256=g5QeGaizSfW4nNe9YZY86g8q79EQmuvSwipaNPnOVqA,56121
+napistu/network/neighborhoods.py,sha256=hi8FT5sGd1vtkR5Uu10wr0Ik5Z3fz9e5fhvXqfi7QPQ,57340
 napistu/network/net_create.py,sha256=66kV_xoWnu4BVLaJZ1TAC7wBSsjPDqjoAXH-X9ShV3s,59091
 napistu/network/net_create_utils.py,sha256=zajwaz2xAij_9fEnD77SgBw_EnNAnJ8jBCmmK2rk_bA,24672
 napistu/network/net_propagation.py,sha256=Il5nDOWh3nLz8gRhDFHGp2LxcvJ9C1twiSZjDeiZMUo,23490
 napistu/network/ng_core.py,sha256=dGnTUKR4WtnvaYMyIHqqF55FY4mJSa7wjA2LZ4cVB6U,11720
-napistu/network/ng_utils.py,sha256=ahSm-8M2pV662V7MMVcGaoguBM55_y-F7LDmZSVp9ag,15951
-napistu/network/paths.py,sha256=r6LVKVvX7i3ctBA5r-xvHfpH5Zsd0VDHUCtin2iag20,17453
+napistu/network/ng_utils.py,sha256=DkI_Ln2uFiNDjPEnUnf7kyy6XwyqvpeUkk8DRjTGZQQ,16078
+napistu/network/paths.py,sha256=BcoYNkCplaM_QPqWWfiwD89bsvwlyvvacSiEzHacfmA,17863
 napistu/network/precompute.py,sha256=ARU2tktWnxFISaHAY8chpkg8pusZPv7TT5jSIB9eFF0,10081
 napistu/ontologies/__init__.py,sha256=dFXAhIqlTLJMwowS4BUDT08-Vy3Q0u1L0CMCErSZT1Y,239
 napistu/ontologies/constants.py,sha256=GyOFvezSxDK1VigATcruTKtNhjcYaid1ggulEf_HEtQ,4345
@@ -87,9 +87,9 @@ napistu/scverse/loading.py,sha256=jqiE71XB-wdV50GyZrauFNY0Lai4bX9Fm2Gv80VR8t8,27
 napistu/statistics/__init__.py,sha256=dFXAhIqlTLJMwowS4BUDT08-Vy3Q0u1L0CMCErSZT1Y,239
 napistu/statistics/hypothesis_testing.py,sha256=k0mBFAMF0XHVcKwS26aPnEbq_FIUVwXU1gZ6cKfFbCk,2190
 napistu/statistics/quantiles.py,sha256=1-LnmVzC2CQWxCKUh0yi6YfKrbsZM1-kkD7nu2-aS5s,3042
-napistu-0.4.3.dist-info/licenses/LICENSE,sha256=kW8wVT__JWoHjl2BbbJDAZInWa9AxzJeR_uv6-i5x1g,1063
+napistu-0.4.5.dist-info/licenses/LICENSE,sha256=kW8wVT__JWoHjl2BbbJDAZInWa9AxzJeR_uv6-i5x1g,1063
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/conftest.py,sha256=t-GHb0MvSsC-MyhkFpOy2K3t5fi7eaig_Rc2xEQC-t8,9678
+tests/conftest.py,sha256=Rw0KtnVyykZhRjnlmNu4oV47lNIeYUJVHu4y47RnVq0,9990
 tests/test_consensus.py,sha256=Hzfrgp4SpkRDnEMVMD3f0UInSycndB8kKzC4wDDvRas,15076
 tests/test_constants.py,sha256=gJLDv7QMeeBiiupyMazj6mumk20KWvGMgm2myHMKKfc,531
 tests/test_context_discretize.py,sha256=5Mr9WqwHGYMO37M1TnMmSfC64UZ73mnoCiEM2IQHVDY,1667
@@ -109,7 +109,7 @@ tests/test_mcp_documentation_utils.py,sha256=OW0N2N_2IOktbYTcCWhhWz4bANi8IB60l1q
 tests/test_mcp_server.py,sha256=bP3PWVQsEfX6-lAgXKP32njdg__o65n2WuLvkxTTHkQ,11215
 tests/test_network_data_handling.py,sha256=4aS8z2AlKkVd-JhK4BQ8fjeiW8_bJ1hZ3cc71Jh7Glk,12716
 tests/test_network_ig_utils.py,sha256=XihmEpX890sr-LYmsb_t4aN0sKIDWCnXkTpDhpuTDmw,7199
-tests/test_network_neighborhoods.py,sha256=8BV17m5X1OUd5FwasTTYUOkNYUHDPUkxOKH_VZCsyBE,631
+tests/test_network_neighborhoods.py,sha256=OvVfgGodbS3MpuSfj-__VKjBj99Ng4WWLmINlIIvbvo,5100
 tests/test_network_net_create.py,sha256=L0U91b4jVHDuC3DFo-_BUFVuv4GuSxZuLAo7r-7EJxY,12877
 tests/test_network_net_create_utils.py,sha256=0J6KIh2HBc4koFsvwMaul1QRtj5x92kR9HBdDZajnAw,18971
 tests/test_network_net_propagation.py,sha256=kZeDHD93iMrLVvxO4OyfRH5_vgsYeQyC40OI9Dsb0xY,14999
@@ -129,15 +129,15 @@ tests/test_sbml_dfs_core.py,sha256=nnLPpZTVtCznOBohk7CX67x6sMqktJWt-sZMWQKoaDs,2
 tests/test_sbml_dfs_utils.py,sha256=ZD9x2B81fsfYEjAV9wphHOR7ywjNcfvfw1LGNv4PxUA,11471
 tests/test_sbo.py,sha256=x_PENFaXYsrZIzOZu9cj_Wrej7i7SNGxgBYYvcigLs0,308
 tests/test_scverse_loading.py,sha256=bnU1lQSYYWhOAs0IIBoi4ZohqPokDQJ0n_rtkAfEyMU,29948
-tests/test_source.py,sha256=iV-Yyu8flhIGWF17SCL8msG2bjqwb9w2IZ694b0iZ-o,2985
+tests/test_source.py,sha256=pe090MsiZ7Tl9P0rhuq17sqMmxUBCch2zoxTwLrNeJQ,2985
 tests/test_statistics_hypothesis_testing.py,sha256=qD-oS9zo5JlH-jdtiOrWAKI4nKFuZvvh6361_pFSpIs,2259
 tests/test_statistics_quantiles.py,sha256=yNDeqwgbP-1Rx3C_dLX_wnwT_Lr-iJWClmeKmElqmTE,4984
 tests/test_uncompartmentalize.py,sha256=nAk5kfAVLU9a2VWe2x2HYVcKqj-EnwmwddERIPRax8c,1289
 tests/test_utils.py,sha256=qPSpV-Q9b6vmdycgaDmQqtcvzKnAVnN9j5xJ9x-T6bg,23959
 tests/utils.py,sha256=SoWQ_5roJteFGcMaOeEiQ5ucwq3Z2Fa3AAs9iXHTsJY,749
 tests/test_data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-napistu-0.4.3.dist-info/METADATA,sha256=gV0a41vyQ52Ja15QyLSPGfeIJPj6oQRTC00HsxJjG88,4078
-napistu-0.4.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-napistu-0.4.3.dist-info/entry_points.txt,sha256=_QnaPOvJNA3IltxmZgWIiBoen-L1bPYX18YQfC7oJgQ,41
-napistu-0.4.3.dist-info/top_level.txt,sha256=Gpvk0a_PjrtqhYcQ9IDr3zR5LqpZ-uIHidQMIpjlvhY,14
-napistu-0.4.3.dist-info/RECORD,,
+napistu-0.4.5.dist-info/METADATA,sha256=JJzjckSlzdusT7COjo-FxaNRBGMWtfJc-kfUDvjyvW4,4078
+napistu-0.4.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+napistu-0.4.5.dist-info/entry_points.txt,sha256=_QnaPOvJNA3IltxmZgWIiBoen-L1bPYX18YQfC7oJgQ,41
+napistu-0.4.5.dist-info/top_level.txt,sha256=Gpvk0a_PjrtqhYcQ9IDr3zR5LqpZ-uIHidQMIpjlvhY,14
+napistu-0.4.5.dist-info/RECORD,,

tests/conftest.py

@@ -139,6 +139,16 @@ def napistu_graph_undirected(sbml_dfs):
     )
 
 
+@fixture
+def napistu_graph_metabolism(sbml_dfs_metabolism):
+    """
+    Pytest fixture to create a NapistuGraph from sbml_dfs_glucose_metabolism with directed=True and topology weighting.
+    """
+    return process_napistu_graph(
+        sbml_dfs_metabolism, directed=True, weighting_strategy="topology"
+    )
+
+
 @pytest.fixture
 def reaction_species_examples():
     """

tests/test_network_neighborhoods.py

@@ -1,13 +1,19 @@
+import pandas as pd
+
 from napistu.network import ng_utils
 from napistu.network import neighborhoods
+from napistu import source
+
+from napistu.constants import SBML_DFS
+from napistu.network.constants import NEIGHBORHOOD_NETWORK_TYPES
 
 
 def test_neighborhood(sbml_dfs, napistu_graph):
     species = sbml_dfs.species
-    source_species = species[species["s_name"] == "NADH"].index.tolist()
+    source_species = species[species[SBML_DFS.S_NAME] == "NADH"].index.tolist()
 
     query_sc_species = ng_utils.compartmentalize_species(sbml_dfs, source_species)
-    compartmentalized_species = query_sc_species["sc_id"].tolist()
+    compartmentalized_species = query_sc_species[SBML_DFS.SC_ID].tolist()
 
     neighborhood = neighborhoods.find_neighborhoods(
         sbml_dfs,
@@ -17,3 +23,109 @@ def test_neighborhood(sbml_dfs, napistu_graph):
     )
 
     assert neighborhood["species_73473"]["vertices"].shape[0] == 6
+
+
+def test_find_and_prune_neighborhoods_with_source_counts(
+    sbml_dfs_metabolism, napistu_graph_metabolism
+):
+    """
+    Test find_and_prune_neighborhoods function with source_total_counts parameter.
+
+    This test verifies that the function works correctly when source_total_counts
+    is provided, which enables source-based edge assignment in neighborhoods.
+    """
+    # Create source_total_counts using the source module
+    source_total_counts = source.get_source_total_counts(
+        sbml_dfs_metabolism, SBML_DFS.REACTIONS
+    )
+
+    # Verify source_total_counts is created correctly
+    assert isinstance(source_total_counts, pd.Series)
+    assert len(source_total_counts) > 0
+    assert source_total_counts.name == "total_counts"
+    assert all(source_total_counts > 0)
+
+    # Get a test species to create neighborhood around
+    species = sbml_dfs_metabolism.species
+    source_species = species[species[SBML_DFS.S_NAME] == "NADH"].index.tolist()
+
+    query_sc_species = ng_utils.compartmentalize_species(
+        sbml_dfs_metabolism, source_species
+    )
+    compartmentalized_species = query_sc_species[SBML_DFS.SC_ID].tolist()
+
+    # Test find_and_prune_neighborhoods with source_total_counts
+    neighborhoods_result = neighborhoods.find_and_prune_neighborhoods(
+        sbml_dfs=sbml_dfs_metabolism,
+        napistu_graph=napistu_graph_metabolism,
+        compartmentalized_species=compartmentalized_species,
+        source_total_counts=source_total_counts,
+        network_type=NEIGHBORHOOD_NETWORK_TYPES.HOURGLASS,
+        order=3,
+        verbose=False,
+        top_n=10,
+    )
+
+    # Verify the result structure
+    assert isinstance(neighborhoods_result, dict)
+    assert len(neighborhoods_result) > 0
+
+    # Check each neighborhood has the expected structure
+    for sc_id, neighborhood in neighborhoods_result.items():
+        assert isinstance(neighborhood, dict)
+        assert "graph" in neighborhood
+        assert "vertices" in neighborhood
+        assert "edges" in neighborhood
+        assert "edge_sources" in neighborhood
+
+        # Verify edge_sources is populated when source_total_counts is provided
+        # (this is the key difference when source_total_counts is passed)
+        if neighborhood["edges"].shape[0] > 0:
+            # If there are edges, edge_sources should be populated
+            assert neighborhood["edge_sources"] is not None
+            assert isinstance(neighborhood["edge_sources"], pd.DataFrame)
+
+            # Check edge_sources has expected columns
+            expected_columns = [SBML_DFS.R_ID, "pathway_id", "name"]
+            for col in expected_columns:
+                assert col in neighborhood["edge_sources"].columns
+
+        # Verify vertices structure
+        vertices = neighborhood["vertices"]
+        assert isinstance(vertices, pd.DataFrame)
+        assert vertices.shape[0] > 0
+
+        # Verify edges structure
+        edges = neighborhood["edges"]
+        assert isinstance(edges, pd.DataFrame)
+
+        # Verify graph structure
+        graph = neighborhood["graph"]
+        assert hasattr(graph, "vcount")
+        assert hasattr(graph, "ecount")
+
+    # Test without source_total_counts for comparison
+    neighborhoods_result_no_source = neighborhoods.find_and_prune_neighborhoods(
+        sbml_dfs=sbml_dfs_metabolism,
+        napistu_graph=napistu_graph_metabolism,
+        compartmentalized_species=compartmentalized_species,
+        source_total_counts=None,  # No source counts
+        network_type=NEIGHBORHOOD_NETWORK_TYPES.DOWNSTREAM,
+        order=3,
+        verbose=False,
+        top_n=10,
+    )
+
+    # Verify both results have the same basic structure
+    assert len(neighborhoods_result) == len(neighborhoods_result_no_source)
+
+    # The main difference should be in edge_sources handling
+    for sc_id in neighborhoods_result:
+        with_source = neighborhoods_result[sc_id]["edge_sources"]
+        without_source = neighborhoods_result_no_source[sc_id]["edge_sources"]
+
+        # Both should either be None or DataFrames, but the content may differ
+        assert (with_source is None) == (without_source is None)
+        if with_source is not None and without_source is not None:
+            assert isinstance(with_source, pd.DataFrame)
+            assert isinstance(without_source, pd.DataFrame)

tests/test_source.py

@@ -96,4 +96,4 @@ def test_source_set_coverage_enrichment(sbml_dfs_metabolism):
         source_df, source_total_counts=source_total_counts, sbml_dfs=sbml_dfs_metabolism
     )
 
-    assert set_coverage.shape == (30, 6)
+    assert set_coverage.shape == (34, 6)