spatial-graph 0.0.1__py3-none-any.whl

spatial_graph/__init__.py

@@ -0,0 +1,15 @@
+from importlib.metadata import version, PackageNotFoundError
+
+try:
+    __version__ = version("spatial_graph")
+except PackageNotFoundError:
+    __version__ = "unknown"
+
+
+from .rtree import PointRTree
+from .rtree import LineRTree
+from .graph import Graph
+from .spatial_graph import SpatialGraph
+
+
+__all__ = ["PointRTree", "LineRTree", "Graph", "SpatialGraph"]

spatial_graph/dtypes.py

@@ -0,0 +1,130 @@
+import numpy as np
+
+
+class DType:
+    def __init__(self, dtype_str):
+        self.as_string = dtype_str
+        self.is_array = self.__is_array(dtype_str)
+
+        if self.is_array:
+            self.base, self.size = self.__parse_array_dtype(dtype_str)
+            self.shape = (self.size,)
+        else:
+            self.base = dtype_str
+            self.size = None
+            self.shape = ()
+
+    def __is_array(self, dtype):
+        if "[" in dtype:
+            if "]" not in dtype:
+                raise RuntimeError(f"invalid array(?) dtype {dtype}")
+            return True
+        return False
+
+    def __parse_array_dtype(self, dtype):
+        dtype, size = dtype.split("[")
+        size = int(size.split("]")[0])
+
+        return dtype, size
+
+    @property
+    def base_c_type(self):
+        """Convert the base of this DType into the equivalent C/C++ type."""
+
+        if self.base == "float32" or self.base == "float":
+            return "float"
+        elif self.base == "float64" or self.base == "double":
+            return "double"
+        else:
+            # this might not work for all of them, this is just a fallback
+            return np.dtype(self.base).name + "_t"
+
+    def to_c_decl(self, name):
+        """Convert this dtype to the equivalent C/C++ declaration with the
+        given name:
+
+            "base_c_type name"        if not an array
+            "base_c_type name[size]"  if an array type
+        """
+        # is this an array type?
+        if self.is_array:
+            suffix = f"[{self.size}]"
+        else:
+            suffix = ""
+
+        return self.base_c_type + " " + name + suffix
+
+    def to_pyxtype(self, use_memory_view=False, add_dim=False):
+        """Convert this dtype to the equivalent PYX type.
+
+            "base_c_type"
+            "base_c_type[size]"     if an array type
+            "base_c_type[::1]"      if an array type and use_memory_view
+            "base_c_type[::1]"      if not an array type and add_dim
+            "base_c_type[:, ::1]"   if an array type and add_dim
+
+        Args:
+
+            use_memory_view:
+
+                If set, will produce "dtype[::1]" instead of "dtype[dim]" for
+                array types.
+
+            add_dim:
+
+                Append a dim to the type, e.g., "int32_t[::1]" instead of
+                "int32_t" for dtype "int32". If this DType is already an array,
+                will create a 2D array, e.g., "int32_t[:, ::1]".
+        """
+
+        # is this an array type?
+        if self.is_array:
+            if add_dim:
+                suffix = "[:, ::1]"
+            else:
+                if use_memory_view:
+                    suffix = "[::1]"
+                else:
+                    suffix = f"[{self.size}]"
+        else:
+            suffix = "[::1]" if add_dim else ""
+
+        return self.base_c_type + suffix
+
+    def to_rvalue(self, name, array_index=None):
+        """Convert this dtype into an r-value to be used in PYX files for
+        assignments.
+
+            "name"                  default
+            "name[array_index]"     if array_index is given
+            "{name[0], ..., name[size-1]}"
+                                    if an array type
+            "{name[array_index, 0], ..., name[array_index, size-1]}"
+                                    if an array type and array_index is given
+        """
+
+        if self.is_array:
+            if array_index:
+                return (
+                    "{"
+                    + ", ".join(
+                        [name + f"[{array_index}, {i}]" for i in range(self.size)]
+                    )
+                    + "}"
+                )
+            else:
+                return (
+                    "{" + ", ".join([name + f"[{i}]" for i in range(self.size)]) + "}"
+                )
+        else:
+            if array_index:
+                return f"{name}[{array_index}]"
+            else:
+                return name
+
+
+def dtypes_to_struct(struct_name, dtypes):
+    pyx_code = f"cdef struct {struct_name}:\n"
+    for name, dtype in dtypes.items():
+        pyx_code += f"    {dtype.to_pyxtype()} {name}\n"
+    return pyx_code

spatial_graph/graph/__init__.py

@@ -0,0 +1,3 @@
+from .graph import Graph
+
+__all__ = ["Graph"]

spatial_graph/graph/graph.py

@@ -0,0 +1,226 @@
+import sys
+import witty
+from pathlib import Path
+import numpy as np
+from Cheetah.Template import Template
+from ..dtypes import DType
+
+# Set platform-specific compile arguments
+if sys.platform == "win32":
+    # Use /O2 for optimization and /std:c++20 for C++20
+    EXTRA_COMPILE_ARGS = ["/O2", "/std:c++20"]
+else:
+    # -O3 for optimization and -std=c++20 for C++20
+    EXTRA_COMPILE_ARGS = ["-O3", "-std=c++20"]
+
+
+class Graph:
+    def __new__(
+        cls,
+        node_dtype,
+        node_attr_dtypes=None,
+        edge_attr_dtypes=None,
+        directed=False,
+        *args,
+        **kwargs,
+    ):
+        if node_attr_dtypes is None:
+            node_attr_dtypes = {}
+        if edge_attr_dtypes is None:
+            edge_attr_dtypes = {}
+
+        node_dtype = DType(node_dtype)
+        node_attr_dtypes = {
+            name: DType(dtype) for name, dtype in node_attr_dtypes.items()
+        }
+        edge_attr_dtypes = {
+            name: DType(dtype) for name, dtype in edge_attr_dtypes.items()
+        }
+
+        src_dir = Path(__file__).parent
+        wrapper_template = Template(
+            file=str(src_dir / "wrapper_template.pyx"),
+            compilerSettings={"directiveStartToken": "%"},
+        )
+        wrapper_template.node_dtype = node_dtype
+        wrapper_template.node_attr_dtypes = node_attr_dtypes
+        wrapper_template.edge_attr_dtypes = edge_attr_dtypes
+        wrapper_template.directed = directed
+
+        wrapper = witty.compile_module(
+            str(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,
+        )
+        GraphType = type(cls.__name__, (cls, wrapper.Graph), {})
+        return wrapper.Graph.__new__(GraphType)
+
+    def __init__(
+        self, node_dtype, node_attr_dtypes=None, edge_attr_dtypes=None, directed=False
+    ):
+        if node_attr_dtypes is None:
+            node_attr_dtypes = {}
+        if edge_attr_dtypes is None:
+            edge_attr_dtypes = {}
+        super().__init__()
+        self.node_dtype = node_dtype
+        self.node_attr_dtypes = node_attr_dtypes
+        self.edge_attr_dtypes = edge_attr_dtypes
+        self.directed = directed
+
+        self.node_attrs = NodeAttrs(self)
+        self.edge_attrs = EdgeAttrs(self)
+
+
+class NodeAttrsView:
+    def __init__(self, graph, nodes):
+        super().__setattr__("graph", graph)
+        for name in graph.node_attr_dtypes.keys():
+            super().__setattr__(
+                f"get_attr_{name}", getattr(graph, f"get_nodes_data_{name}")
+            )
+            super().__setattr__(
+                f"set_attr_{name}", getattr(graph, f"set_nodes_data_{name}")
+            )
+
+        if nodes is not None and not isinstance(nodes, np.ndarray):
+            # nodes is not an ndarray, can it be converted into one?
+            try:
+                # does it have a length?
+                _ = len(nodes)
+                # if so, convert to ndarray
+                nodes = np.array(nodes, dtype=graph.node_dtype)
+            except Exception:
+                # must be a single node
+                for name in graph.node_attr_dtypes.keys():
+                    super().__setattr__(
+                        f"set_attr_{name}", getattr(graph, f"set_node_data_{name}")
+                    )
+                    super().__setattr__(
+                        f"get_attr_{name}", getattr(graph, f"get_node_data_{name}")
+                    )
+
+        # at this point, nodes is either
+        # 1. a numpy array
+        # 2. a scalar (python or numpy)
+        # 3. None
+        super().__setattr__("nodes", nodes)
+
+    def __getattr__(self, name):
+        if name in self.graph.node_attr_dtypes:
+            return getattr(self, f"get_attr_{name}")(self.nodes)
+        else:
+            raise AttributeError(name)
+
+    def __setattr__(self, name, values):
+        if name in self.graph.node_attr_dtypes:
+            return getattr(self, f"set_attr_{name}")(self.nodes, values)
+        else:
+            return super().__setattr__(name, values)
+
+    def __iter__(self):
+        # TODO: shouldn't be possible if nodes is a single node
+        yield from self.graph.nodes_data(self.nodes)
+
+
+class EdgeAttrsView:
+    def __init__(self, graph, edges):
+        super().__setattr__("graph", graph)
+        for name in graph.edge_attr_dtypes.keys():
+            super().__setattr__(
+                f"get_attr_{name}", getattr(graph, f"get_edges_data_{name}")
+            )
+            super().__setattr__(
+                f"set_attr_{name}", getattr(graph, f"set_edges_data_{name}")
+            )
+
+        # edges types we support:
+        #
+        # 1. edges = None                   all edges           leave as is
+        # 2. edges = iteratible of 2-tuples selected edges      to (n,2) ndarray
+        # 3. edges = iteratible of 2-lists  selected edges      to (n,2) ndarray
+        # 4. edges = (n,2) ndarray          selected edges      leave as is
+        # 5. edges = 2-tuple                a single edge       leave as is
+        # 6. edges = (2,) ndarray           a single edge       to 2-tuple
+
+        if edges is not None:
+            if isinstance(edges, np.ndarray):
+                if len(edges) == 2 and len(edges.shape) == 1:
+                    # case 6
+                    edges = tuple(edges)
+                else:
+                    # case 4 with multiple edges
+                    edges = edges.astype(graph.node_dtype)
+            elif isinstance(edges, tuple):
+                # case 5
+                assert len(edges) == 2, "Single edges should be given as a 2-tuple"
+            else:
+                # edges should be an iteratable
+                try:
+                    # does it have a length?
+                    len(edges)
+                    # case 2 and 3
+                    edges = np.array(edges, dtype=graph.node_dtype)
+                except Exception:
+                    raise RuntimeError(f"Can not handle edges type {type(edges)}")
+
+        if isinstance(edges, np.ndarray):
+            if len(edges) == 0:
+                edges = edges.reshape((0, 2))
+            assert edges.shape[1] == 2, "Edge arrays should have shape (n, 2)"
+            edges = np.ascontiguousarray(edges.T)
+        elif isinstance(edges, tuple):
+            # a single edge
+            for name in graph.edge_attr_dtypes.keys():
+                super().__setattr__(
+                    f"get_attr_{name}", getattr(graph, f"get_edge_data_{name}")
+                )
+                super().__setattr__(
+                    f"set_attr_{name}", getattr(graph, f"set_edge_data_{name}")
+                )
+
+        # at this point, edges is either
+        # 1. a nx2 numpy array
+        # 2. a 2-tuple of scalars (python or numpy)
+        # 3. None
+        super().__setattr__("edges", edges)
+
+    def __getattr__(self, name):
+        if name in self.graph.edge_attr_dtypes:
+            if self.edges is not None:
+                return getattr(self, f"get_attr_{name}")(self.edges[0], self.edges[1])
+            else:
+                return getattr(self, f"get_attr_{name}")(None, None)
+        else:
+            raise AttributeError(name)
+
+    def __setattr__(self, name, values):
+        if name in self.graph.edge_attr_dtypes:
+            return getattr(self, f"set_attr_{name}")(
+                self.edges[0], self.edges[1], values
+            )
+        else:
+            return super().__setattr__(name, values)
+
+    def __iter__(self):
+        # TODO: shouldn't be possible if edges is a single edge
+        yield from self.graph.edges_data(self.edges)
+
+
+class NodeAttrs(NodeAttrsView):
+    def __init__(self, graph):
+        super().__init__(graph, nodes=None)
+
+    def __getitem__(self, nodes):
+        return NodeAttrsView(self.graph, nodes)
+
+
+class EdgeAttrs(EdgeAttrsView):
+    def __init__(self, graph):
+        super().__init__(graph, edges=None)
+
+    def __getitem__(self, edges):
+        return EdgeAttrsView(self.graph, edges)

spatial_graph/graph/src/LICENSE.txt

@@ -0,0 +1,7 @@
+Copyright 2021 Guohao Dou
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.