spatial-graph 0.0.3__tar.gz → 0.0.4__tar.gz

{spatial_graph-0.0.3 → spatial_graph-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spatial-graph
-Version: 0.0.3
+Version: 0.0.4
 Summary: A spatial graph datastructure for python.
 Project-URL: homepage, https://github.com/funkelab/spatial_graph
 Project-URL: repository, https://github.com/funkelab/spatial_graph
@@ -20,7 +20,7 @@ Requires-Python: >=3.9
 Requires-Dist: ct3>=3.3.3
 Requires-Dist: numpy
 Requires-Dist: setuptools>=75.8.0
-Requires-Dist: witty>=v0.2.1
+Requires-Dist: witty>=v0.3.0
 Description-Content-Type: text/markdown
 
 # spatial-graph

{spatial_graph-0.0.3 → spatial_graph-0.0.4}/pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
   "Programming Language :: Python :: 3.13",
   "Typing :: Typed",
 ]
-dependencies = ["witty>=v0.2.1", "CT3>=3.3.3", "numpy", "setuptools>=75.8.0"]
+dependencies = ["witty>=v0.3.0", "CT3>=3.3.3", "numpy", "setuptools>=75.8.0"]
 
 [dependency-groups]
 test = ["pytest>=8.3.5", "pytest-cov>=6.1.1"]

spatial_graph-0.0.4/src/spatial_graph/_graph/graph.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal, overload
+
+from .graph_base import GraphBase
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    import numpy as np
+
+
+class Graph(GraphBase):
+    directed: Literal[False] = False
+
+    def num_neighbors(self, nodes: np.ndarray) -> np.ndarray:
+        """Return the number of neighbors for each node.
+
+        For undirected graphs, this counts all adjacent nodes regardless
+        of edge direction since edges are bidirectional.
+
+        Parameters
+        ----------
+        nodes : np.ndarray
+            Array of node identifiers to count neighbors for.
+
+        Returns
+        -------
+        np.ndarray
+            Array of neighbor counts for each node in the input array.
+        """
+        return self._cgraph.num_neighbors(nodes)
+
+    @overload
+    def edges(
+        self, node: Any = ..., data: Literal[True] = ...
+    ) -> Iterator[tuple[tuple, Any]]: ...
+    @overload
+    def edges(self, node: Any = ..., data: Literal[False] = ...) -> Iterator[tuple]: ...
+    def edges(self, node: Any = None, data: bool = False) -> Iterator[tuple]:
+        """Iterate over edges in the graph.
+
+        For undirected graphs, each edge is yielded only once with nodes
+        ordered such that node1 < node2 to avoid duplicates.
+
+        Parameters
+        ----------
+        node : Any, optional
+            If provided, only iterate over edges incident to this node.
+            If None, iterate over all edges in the graph.
+        data : bool, default False
+            If True, yield (edge, edge_data) tuples. If False, yield
+            only edge tuples.
+
+        Yields
+        ------
+        tuple or tuple[tuple, Any]
+            If `data=False`: tuples of (node1, node2) representing edges.
+            If `data=True`: tuples of ((node1, node2), edge_data) where
+            edge_data provides access to edge attributes.
+        """
+        return self._cgraph.edges(node, data)
+
+    def edges_by_nodes(self, nodes: np.ndarray) -> np.ndarray:
+        """Get all edges incident to the specified nodes.
+
+        This method provides fast access to edges incident to an array
+        of nodes. Note that edges between nodes in the input array will
+        be reported multiple times (once for each incident node).
+
+        Parameters
+        ----------
+        nodes : np.ndarray
+            Array of node identifiers to find incident edges for.
+
+        Returns
+        -------
+        np.ndarray
+            2D array of shape (n_edges, 2) where each row contains
+            [node1, node2] representing an edge. For undirected graphs,
+            node1 <= node2.
+        """
+        return self._cgraph.edges_by_nodes(nodes)
+
+
+class DiGraph(GraphBase):
+    directed: Literal[True] = True
+
+    def num_in_neighbors(self, nodes: np.ndarray) -> np.ndarray:
+        """Return the number of incoming neighbors for each node.
+
+        This counts only nodes that have edges pointing to the specified nodes
+        (i.e., predecessors).
+
+        Parameters
+        ----------
+        nodes : np.ndarray
+            Array of node identifiers to count incoming neighbors for.
+
+        Returns
+        -------
+        np.ndarray
+            Array of incoming neighbor counts for each node in the input array.
+        """
+        return self._cgraph.num_in_neighbors(nodes)
+
+    def num_out_neighbors(self, nodes: np.ndarray) -> np.ndarray:
+        """Return the number of outgoing neighbors for each node.
+
+        This counts only nodes that the specified nodes
+        have edges pointing to (i.e., successors).
+
+        Parameters
+        ----------
+        nodes : np.ndarray
+            Array of node identifiers to count outgoing neighbors for.
+
+        Returns
+        -------
+        np.ndarray
+            Array of outgoing neighbor counts for each node in the input array.
+        """
+        return self._cgraph.num_out_neighbors(nodes)
+
+    @overload
+    def in_edges(
+        self, node: Any = ..., data: Literal[True] = ...
+    ) -> Iterator[tuple[tuple, Any]]: ...
+    @overload
+    def in_edges(
+        self, node: Any = ..., data: Literal[False] = ...
+    ) -> Iterator[tuple]: ...
+    def in_edges(self, node: Any = None, data: bool = False) -> Iterator[tuple]:
+        """Iterate over incoming edges to a node.
+
+        Only edges directed toward the specified node are yielded.
+
+        Parameters
+        ----------
+        node : Any
+            The target node to find incoming edges for.
+        data : bool
+            If True, yield (edge, edge_data) tuples. If False, yield
+            only edge tuples.
+
+        Yields
+        ------
+        tuple or tuple[tuple, Any]
+            If `data=False`: tuples of (source_node, target_node) representing
+            incoming edges where target_node is the specified node.
+            If `data=True`: tuples of ((source_node, target_node), edge_data)
+            where edge_data provides access to edge attributes.
+        """
+        return self._cgraph.in_edges(node, data)
+
+    def in_edges_by_nodes(self, nodes: np.ndarray) -> np.ndarray:
+        """Get all incoming edges to the specified nodes.
+
+        This method provides fast access to incoming edges for an array
+        of nodes. Edges between nodes in the input array will be reported
+        multiple times if both source and target are in the array.
+
+        Parameters
+        ----------
+        nodes : np.ndarray
+            Array of node identifiers to find incoming edges for.
+
+        Returns
+        -------
+        np.ndarray
+            2D array of shape (n_edges, 2) where each row contains
+            [source_node, target_node] representing an incoming edge
+            to one of the specified nodes.
+        """
+        return self._cgraph.in_edges_by_nodes(nodes)
+
+    @overload
+    def out_edges(
+        self, node: Any = ..., data: Literal[True] = ...
+    ) -> Iterator[tuple[tuple, Any]]: ...
+    @overload
+    def out_edges(
+        self, node: Any = ..., data: Literal[False] = ...
+    ) -> Iterator[tuple]: ...
+    def out_edges(self, node: Any = None, data: bool = False) -> Iterator[tuple]:
+        """Iterate over outgoing edges from a node.
+
+        Only edges directed away from the specified node are yielded.
+
+        Parameters
+        ----------
+        node : Any
+            The source node to find outgoing edges for.
+        data : bool
+            If True, yield (edge, edge_data) tuples. If False, yield
+            only edge tuples.
+
+        Yields
+        ------
+        tuple or tuple[tuple, Any]
+            If `data=False`: tuples of (source_node, target_node) representing
+            outgoing edges where source_node is the specified node.
+            If `data=True`: tuples of ((source_node, target_node), edge_data)
+            where edge_data provides access to edge attributes.
+        """
+        return self._cgraph.out_edges(node, data)
+
+    def out_edges_by_nodes(self, nodes: np.ndarray) -> np.ndarray:
+        """Get all outgoing edges from the specified nodes.
+
+        This method provides fast access to outgoing edges for an array
+        of nodes. Edges between nodes in the input array will be reported
+        multiple times if both source and target are in the array.
+
+        Parameters
+        ----------
+        nodes : np.ndarray
+            Array of node identifiers to find outgoing edges for.
+
+        Returns
+        -------
+        np.ndarray
+            2D array of shape (n_edges, 2) where each row contains
+            [source_node, target_node] representing an outgoing edge
+            from one of the specified nodes.
+        """
+        return self._cgraph.out_edges_by_nodes(nodes)

spatial_graph-0.0.4/src/spatial_graph/_graph/graph_base.py

@@ -0,0 +1,321 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import witty
+from Cheetah.Template import Template
+
+from spatial_graph._dtypes import DType
+
+from .views import EdgeAttrs, NodeAttrs
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator, Mapping
+
+    import numpy as np
+
+
+# Set platform-specific compile arguments
+if sys.platform == "win32":  # pragma: no cover
+    # Use /O2 for optimization and /std:c++20 for C++20
+    EXTRA_COMPILE_ARGS = ["/O2", "/std:c++20", "/wd4101"]
+else:
+    # -O3 for optimization and -std=c++20 for C++20
+    EXTRA_COMPILE_ARGS = [
+        "-O3",
+        "-std=c++20",
+        "-Wno-unused-variable",
+        "-Wno-unreachable-code",
+    ]
+
+SRC_DIR = Path(__file__).parent
+
+
+def _build_wrapper(
+    node_dtype: str,
+    node_attr_dtypes: Mapping[str, str] | None = None,
+    edge_attr_dtypes: Mapping[str, str] | None = None,
+    directed: bool = False,
+) -> str:
+    if node_attr_dtypes is None:
+        node_attr_dtypes = {}
+    if edge_attr_dtypes is None:
+        edge_attr_dtypes = {}
+    if not all(str.isidentifier(name) for name in node_attr_dtypes):
+        raise ValueError("Node attribute names must be valid identifiers")
+    if not all(str.isidentifier(name) for name in edge_attr_dtypes):
+        raise ValueError("Edge attribute names must be valid identifiers")
+
+    wrapper_template = Template(
+        file=str(SRC_DIR / "wrapper_template.pyx"),
+        compilerSettings={"directiveStartToken": "%"},
+    )
+    wrapper_template.node_dtype = DType(node_dtype)
+    wrapper_template.node_attr_dtypes = {
+        name: DType(dtype) for name, dtype in node_attr_dtypes.items()
+    }
+    wrapper_template.edge_attr_dtypes = {
+        name: DType(dtype) for name, dtype in edge_attr_dtypes.items()
+    }
+    wrapper_template.directed = directed
+
+    return str(wrapper_template)
+
+
+def _compile_graph(
+    node_dtype: str,
+    node_attr_dtypes: Mapping[str, str] | None = None,
+    edge_attr_dtypes: Mapping[str, str] | None = None,
+    directed: bool = False,
+) -> type:
+    wrapper_template = _build_wrapper(
+        node_dtype=node_dtype,
+        node_attr_dtypes=node_attr_dtypes,
+        edge_attr_dtypes=edge_attr_dtypes,
+        directed=directed,
+    )
+    wrapper = witty.compile_cython(
+        wrapper_template,
+        source_files=[str(SRC_DIR / "src" / "graph_lite.h")],
+        extra_compile_args=EXTRA_COMPILE_ARGS,
+        include_dirs=[str(SRC_DIR)],
+        language="c++",
+        quiet=True,
+    )
+    return wrapper.Graph
+
+
+class GraphBase:
+    directed: bool = False
+
+    def __init__(
+        self,
+        node_dtype: str,
+        node_attr_dtypes: Mapping[str, str] | None = None,
+        edge_attr_dtypes: Mapping[str, str] | None = None,
+    ):
+        super().__init__()
+        self.node_dtype = node_dtype
+        self.node_attr_dtypes = node_attr_dtypes or {}
+        self.edge_attr_dtypes = edge_attr_dtypes or {}
+
+        cgraph_cls = _compile_graph(
+            node_dtype=self.node_dtype,
+            node_attr_dtypes=self.node_attr_dtypes,
+            edge_attr_dtypes=self.edge_attr_dtypes,
+            directed=self.directed,
+        )
+        self._cgraph = cgraph_cls()
+
+        self.node_attrs = NodeAttrs(self)
+        self.edge_attrs = EdgeAttrs(self)
+
+    def add_node(self, node: Any, *data: Any, **kwargs: Any) -> int:
+        """Add a single node to the graph.
+
+        The node attributes provided via *data and **kwargs must match the
+        data types and names specified in `node_attr_dtypes` when the graph
+        was created.
+
+        Parameters
+        ----------
+        node : Any
+            The node identifier to add to the graph.
+        *data : Any
+            Positional arguments for node attributes. Names/number of args
+            must match the `node_attr_dtypes`.
+        **kwargs : Any
+            Keyword arguments for node attributes. Names/number of kwargs
+            must match the `node_attr_dtypes`.
+
+        Returns
+        -------
+        int
+            Number of nodes added (1 if successful, 0 if node already exists).
+        """
+        return self._cgraph.add_node(node, *data, **kwargs)
+
+    def add_nodes(self, nodes: np.ndarray, *data: Any, **kwargs: Any) -> int:
+        """Add multiple nodes to the graph.
+
+        Node attributes provided via *data and **kwargs must match the
+        data types and names specified in `node_attr_dtypes`. Each attribute
+        array must have the same length as the `nodes` array.
+
+        Parameters
+        ----------
+        nodes : np.ndarray
+            Array of node identifiers to add to the graph.
+        *data : Any
+            Positional arguments for node attributes. Each argument should be
+            an array with length matching `nodes`. Names/number of args must
+            match the `node_attr_dtypes`.
+        **kwargs : Any
+            Keyword arguments for node attributes. Each argument should be
+            an array with length matching `nodes`. Names/number of kwargs
+            must match the `node_attr_dtypes`.
+
+        Returns
+        -------
+        int
+            Number of nodes successfully added.
+        """
+        return self._cgraph.add_nodes(nodes, *data, **kwargs)
+
+    def add_edge(self, edge: np.ndarray, *args: Any, **kwargs: Any) -> int:
+        """Add an edge to the graph.
+
+        The edge attributes provided via *args and **kwargs must match the
+        data types and names specified in `edge_attr_dtypes` when the graph
+        was created.
+
+        Parameters
+        ----------
+        edge : np.ndarray
+            Array of length 2 containing [source_node, target_node].
+        *args : Any
+            Positional arguments for edge attributes. Names/number of args
+            must match the `edge_attr_dtypes`.
+        **kwargs : Any
+            Keyword arguments for edge attributes. Names/number of kwargs
+            must match the `edge_attr_dtypes`.
+
+        Returns
+        -------
+        int
+            Number of edges added (1 if successful, 0 if edge already exists).
+        """
+        return self._cgraph.add_edge(edge, *args, **kwargs)
+
+    def add_edges(
+        self, edges: np.ndarray, *args: np.ndarray, **kwargs: np.ndarray
+    ) -> int:
+        """Add multiple edges to the graph.
+
+        Edge attributes provided via *args and **kwargs must match the
+        data types and names specified in `edge_attr_dtypes`. Each attribute
+        array must have the same length as the number of edges.
+
+        Parameters
+        ----------
+        edges : np.ndarray
+            2D array of shape (n_edges, 2) where each row contains
+            [source_node, target_node].
+        *args : np.ndarray
+            Positional arguments for edge attributes. Each argument should be
+            an array with length matching the number of edges. Names/number
+            of args must match the `edge_attr_dtypes`.
+        **kwargs : np.ndarray
+            Keyword arguments for edge attributes. Each argument should be
+            an array with length matching the number of edges. Names/number
+            of kwargs must match the `edge_attr_dtypes`.
+
+        Returns
+        -------
+        int
+            Number of edges successfully added.
+        """
+        return self._cgraph.add_edges(edges, *args, **kwargs)
+
+    @property
+    def nodes(self) -> np.ndarray:
+        """Get all node IDs in the graph.
+
+        The returned array is a copy and modifications will not affect
+        the graph structure.
+
+        Returns
+        -------
+        np.ndarray
+            Array containing all node identifiers in the graph, ordered
+            by insertion order (earliest added first).
+        """
+        return self._cgraph.nodes()
+
+    def remove_node(self, node: Any) -> None:
+        """Remove a single node from the graph.
+
+        Removing a node will also remove all edges incident to that node.
+
+        Parameters
+        ----------
+        node : Any
+            The node identifier to remove from the graph.
+        """
+        return self._cgraph.remove_node(node)
+
+    def remove_nodes(self, nodes: np.ndarray) -> None:
+        """Remove multiple nodes from the graph.
+
+        Removing nodes will also remove all edges incident to those nodes.
+
+        Parameters
+        ----------
+        nodes : np.ndarray
+            Array of node identifiers to remove from the graph.
+        """
+        return self._cgraph.remove_nodes(nodes)
+
+    def nodes_data(self, nodes: np.ndarray | None = None) -> Iterator[tuple[Any, Any]]:
+        """Iterate over nodes and their associated data.
+
+        The node_data object provides access to node attributes as defined
+        by the `node_attr_dtypes` when the graph was created.
+
+        Parameters
+        ----------
+        nodes : np.ndarray, optional
+            Array of specific node identifiers to iterate over. If None,
+            iterates over all nodes in the graph.
+
+        Yields
+        ------
+        tuple[Any, Any]
+            Tuples of (node_id, node_data) where node_data is a view object
+            providing access to the node's attributes.
+        """
+        return self._cgraph.nodes_data(nodes)
+
+    def edges_data(self, us: np.ndarray, vs: np.ndarray) -> Iterator:
+        """Iterate over edge data for specified edges.
+
+        The arrays `us` and `vs` must have the same length. The edge data
+        objects provide access to edge attributes as defined by the
+        `edge_attr_dtypes` when the graph was created.
+
+        Parameters
+        ----------
+        us : np.ndarray
+            Array of source node identifiers.
+        vs : np.ndarray
+            Array of target node identifiers.
+
+        Yields
+        ------
+        Any
+            Edge data view objects providing access to edge attributes
+            for each edge (us[i], vs[i]).
+        """
+        return self._cgraph.edges_data(us, vs)
+
+    def num_edges(self) -> int:
+        """Get the total number of edges in the graph.
+
+        Returns
+        -------
+        int
+            The number of edges in the graph.
+        """
+        return self._cgraph.num_edges()
+
+    def __len__(self) -> int:
+        """Return the number of nodes in the graph.
+
+        Returns
+        -------
+        int
+            The number of nodes in the graph.
+        """
+        return self._cgraph.__len__()