napistu 0.2.5.dev7__py3-none-any.whl → 0.3.1.dev1__py3-none-any.whl

napistu/network/net_propagation.py

@@ -1,13 +1,15 @@
+import inspect
+from typing import Optional, Union
+
 import pandas as pd
 import numpy as np
 import igraph as ig
-import inspect
 
-from typing import Optional
+from napistu.network.napistu_graph_core import NapistuGraph
 
 
 def personalized_pagerank_by_attribute(
-    g: ig.Graph,
+    napistu_graph: Union[NapistuGraph, ig.Graph],
     attribute: str,
     damping: float = 0.85,
     calculate_uniform_dist: bool = True,
@@ -19,8 +21,8 @@ def personalized_pagerank_by_attribute(
 
     Parameters
     ----------
-    g : igraph.Graph
-        The input graph.
+    napistu_graph : NapistuGraph
+        The input graph (subclass of igraph.Graph).
     attribute : str
         The vertex attribute to use for personalization.
     damping : float, optional
@@ -38,7 +40,7 @@ def personalized_pagerank_by_attribute(
     Example
     -------
     >>> import igraph as ig
-    >>> from scraps.utils import personalized_pagerank_by_attribute
+    >>> from napistu.network.net_propagation import personalized_pagerank_by_attribute
     >>> g = ig.Graph.Full(3)
     >>> g.vs['name'] = ['A', 'B', 'C']
     >>> g.vs['score'] = [1, 0, 2]
@@ -46,24 +48,30 @@ def personalized_pagerank_by_attribute(
     >>> print(df)
     """
     # Validate and extract attribute (missing/None as 0)
-    attr = _ensure_nonnegative_vertex_attribute(g, attribute)
+    attr = _ensure_nonnegative_vertex_attribute(napistu_graph, attribute)
 
     # Validate additional_propagation_args
     if additional_propagation_args is None:
         additional_propagation_args = {}
     else:
-        valid_args = set(inspect.signature(g.personalized_pagerank).parameters.keys())
+        valid_args = set(
+            inspect.signature(napistu_graph.personalized_pagerank).parameters.keys()
+        )
         for k in additional_propagation_args:
             if k not in valid_args:
                 raise ValueError(f"Invalid argument for personalized_pagerank: {k}")
 
     # Personalized PageRank (no normalization, igraph handles it)
-    pr_attr = g.personalized_pagerank(
+    pr_attr = napistu_graph.personalized_pagerank(
         reset=attr.tolist(), damping=damping, **additional_propagation_args
     )
 
     # Node names
-    names = g.vs["name"] if "name" in g.vs.attributes() else list(range(g.vcount()))
+    names = (
+        napistu_graph.vs["name"]
+        if "name" in napistu_graph.vs.attributes()
+        else list(range(napistu_graph.vcount()))
+    )
 
     data = {"name": names, "pagerank_by_attribute": pr_attr, attribute: attr}
 
@@ -75,7 +83,7 @@ def personalized_pagerank_by_attribute(
             raise ValueError("No nonzero attribute values for uniform PPR.")
         uniform_vec = np.zeros_like(attr, dtype=float)
         uniform_vec[used_in_uniform] = 1.0 / n_uniform
-        pr_uniform = g.personalized_pagerank(
+        pr_uniform = napistu_graph.personalized_pagerank(
             reset=uniform_vec.tolist(), damping=damping, **additional_propagation_args
         )
         data["pagerank_uniform"] = pr_uniform
@@ -83,16 +91,36 @@ def personalized_pagerank_by_attribute(
     return pd.DataFrame(data)
 
 
-def _ensure_nonnegative_vertex_attribute(g: ig.Graph, attribute: str):
-    """
-    Utility to check that a vertex attribute is present, numeric, and non-negative.
-    Raises ValueError if checks fail.
-    Missing or None values are treated as 0.
-    Raises ValueError if attribute is missing for all vertices or all values are zero.
+def _ensure_nonnegative_vertex_attribute(
+    napistu_graph: Union[NapistuGraph, ig.Graph], attribute: str
+):
     """
+    Ensure a vertex attribute is present, numeric, and non-negative for all vertices.
 
+    This utility checks that the specified vertex attribute exists, is numeric, and non-negative
+    for all vertices in the graph. Missing or None values are treated as 0. Raises ValueError
+    if the attribute is missing for all vertices, if all values are zero, or if any value is negative.
+
+    Parameters
+    ----------
+    napistu_graph : NapistuGraph or ig.Graph
+        The input graph (NapistuGraph or igraph.Graph).
+    attribute : str
+        The name of the vertex attribute to check.
+
+    Returns
+    -------
+    np.ndarray
+        Array of attribute values (with missing/None replaced by 0).
+
+    Raises
+    ------
+    ValueError
+        If the attribute is missing for all vertices, all values are zero, or any value is negative.
+    """
     all_missing = all(
-        (attribute not in v.attributes() or v[attribute] is None) for v in g.vs
+        (attribute not in v.attributes() or v[attribute] is None)
+        for v in napistu_graph.vs
     )
     if all_missing:
         raise ValueError(f"Vertex attribute '{attribute}' is missing for all vertices.")
@@ -103,7 +131,7 @@ def _ensure_nonnegative_vertex_attribute(g: ig.Graph, attribute: str):
             if (attribute in v.attributes() and v[attribute] is not None)
             else 0.0
         )
-        for v in g.vs
+        for v in napistu_graph.vs
     ]
 
     arr = np.array(values, dtype=float)