graphical-sampling 0.1.0__py3-none-any.whl

graphical_sampling/red_black_tree.py

@@ -0,0 +1,475 @@
+from __future__ import annotations
+from collections.abc import Iterator
+from typing import TypeVar, Generic
+
+from .type import Comparable
+
+C = TypeVar("C", bound=Comparable)
+
+
+class RedBlackTree(Generic[C]):
+    def __init__(
+        self,
+        label: C | None = None,
+        color: int = 0,
+        parent: RedBlackTree | None = None,
+        left: RedBlackTree | None = None,
+        right: RedBlackTree | None = None,
+    ) -> None:
+        """Initialize a new Red-Black Tree node with the given values:
+        label: The value associated with this node
+        color: 0 if black, 1 if red
+        parent: The parent to this node
+        left: This node's left child
+        right: This node's right child
+        """
+        self.label = label
+        self.parent = parent
+        self.left = left
+        self.right = right
+        self.color = color
+
+    # Here are functions which are specific to red-black trees
+
+    def rotate_left(self) -> RedBlackTree:
+        """Rotate the subtree rooted at this node to the left and
+        returns the new root to this subtree.
+        Performing one rotation can be done in O(1).
+        """
+        parent = self.parent
+        right = self.right
+        if right is None:
+            return self
+        self.right = right.left
+        if self.right:
+            self.right.parent = self
+        self.parent = right
+        right.left = self
+        if parent is not None:
+            if parent.left == self:
+                parent.left = right
+            else:
+                parent.right = right
+        right.parent = parent
+        return right
+
+    def rotate_right(self) -> RedBlackTree:
+        """Rotate the subtree rooted at this node to the right and
+        returns the new root to this subtree.
+        Performing one rotation can be done in O(1).
+        """
+        if self.left is None:
+            return self
+        parent = self.parent
+        left = self.left
+        self.left = left.right
+        if self.left:
+            self.left.parent = self
+        self.parent = left
+        left.right = self
+        if parent is not None:
+            if parent.right is self:
+                parent.right = left
+            else:
+                parent.left = left
+        left.parent = parent
+        return left
+
+    def insert(self, label: C) -> RedBlackTree:
+        """Inserts label into the subtree rooted at self, performs any
+        rotations necessary to maintain balance, and then returns the
+        new root to this subtree (likely self).
+        This is guaranteed to run in O(log(n)) time.
+        """
+        if self.label is None:
+            # Only possible with an empty tree
+            self.label = label
+            return self
+        if self.label == label:
+            return self
+        elif self.label > label:
+            if self.left:
+                self.left.insert(label)
+            else:
+                self.left = RedBlackTree(label, 1, self)
+                self.left._insert_repair()
+        elif self.right:
+            self.right.insert(label)
+        else:
+            self.right = RedBlackTree(label, 1, self)
+            self.right._insert_repair()
+        return self.parent or self
+
+    def _insert_repair(self) -> None:
+        """Repair the coloring from inserting into a tree."""
+        if self.parent is None:
+            # This node is the root, so it just needs to be black
+            self.color = 0
+        elif color(self.parent) == 0:
+            # If the parent is black, then it just needs to be red
+            self.color = 1
+        else:
+            uncle = self.parent.sibling
+            if color(uncle) == 0:
+                if self.is_left() and self.parent.is_right():
+                    self.parent.rotate_right()
+                    if self.right:
+                        self.right._insert_repair()
+                elif self.is_right() and self.parent.is_left():
+                    self.parent.rotate_left()
+                    if self.left:
+                        self.left._insert_repair()
+                elif self.is_left():
+                    if self.grandparent:
+                        self.grandparent.rotate_right()
+                        self.parent.color = 0
+                    if self.parent.right:
+                        self.parent.right.color = 1
+                else:
+                    if self.grandparent:
+                        self.grandparent.rotate_left()
+                        self.parent.color = 0
+                    if self.parent.left:
+                        self.parent.left.color = 1
+            else:
+                self.parent.color = 0
+                if uncle and self.grandparent:
+                    uncle.color = 0
+                    self.grandparent.color = 1
+                    self.grandparent._insert_repair()
+
+    def remove(self, label: C) -> RedBlackTree:
+        """Remove label from this tree."""
+        if self.label == label:
+            if self.left and self.right:
+                # It's easier to balance a node with at most one child,
+                # so we replace this node with the greatest one less than
+                # it and remove that.
+                value = self.left.get_max()
+                if value is not None:
+                    self.label = value
+                    self.left.remove(value)
+            else:
+                # This node has at most one non-None child, so we don't
+                # need to replace
+                child = self.left or self.right
+                if self.color == 1:
+                    # This node is red, and its child is black
+                    # The only way this happens to a node with one child
+                    # is if both children are None leaves.
+                    # We can just remove this node and call it a day.
+                    if self.parent:
+                        if self.is_left():
+                            self.parent.left = None
+                        else:
+                            self.parent.right = None
+                # The node is black
+                elif child is None:
+                    # This node and its child are black
+                    if self.parent is None:
+                        # The tree is now empty
+                        return RedBlackTree(None)
+                    else:
+                        self._remove_repair()
+                        if self.is_left():
+                            self.parent.left = None
+                        else:
+                            self.parent.right = None
+                        self.parent = None
+                else:
+                    # This node is black and its child is red
+                    # Move the child node here and make it black
+                    self.label = child.label
+                    self.left = child.left
+                    self.right = child.right
+                    if self.left:
+                        self.left.parent = self
+                    if self.right:
+                        self.right.parent = self
+        elif self.label is not None and self.label > label:
+            if self.left:
+                self.left.remove(label)
+        elif self.right:
+            self.right.remove(label)
+        return self.parent or self
+
+    def _remove_repair(self) -> None:
+        """Repair the coloring of the tree that may have been messed up."""
+        if (
+            self.parent is None
+            or self.sibling is None
+            or self.parent.sibling is None
+            or self.grandparent is None
+        ):
+            return
+        if color(self.sibling) == 1:
+            self.sibling.color = 0
+            self.parent.color = 1
+            if self.is_left():
+                self.parent.rotate_left()
+            else:
+                self.parent.rotate_right()
+        if (
+            color(self.parent) == 0
+            and color(self.sibling) == 0
+            and color(self.sibling.left) == 0
+            and color(self.sibling.right) == 0
+        ):
+            self.sibling.color = 1
+            self.parent._remove_repair()
+            return
+        if (
+            color(self.parent) == 1
+            and color(self.sibling) == 0
+            and color(self.sibling.left) == 0
+            and color(self.sibling.right) == 0
+        ):
+            self.sibling.color = 1
+            self.parent.color = 0
+            return
+        if (
+            self.is_left()
+            and color(self.sibling) == 0
+            and color(self.sibling.right) == 0
+            and color(self.sibling.left) == 1
+        ):
+            self.sibling.rotate_right()
+            self.sibling.color = 0
+            if self.sibling.right:
+                self.sibling.right.color = 1
+        if (
+            self.is_right()
+            and color(self.sibling) == 0
+            and color(self.sibling.right) == 1
+            and color(self.sibling.left) == 0
+        ):
+            self.sibling.rotate_left()
+            self.sibling.color = 0
+            if self.sibling.left:
+                self.sibling.left.color = 1
+        if (
+            self.is_left()
+            and color(self.sibling) == 0
+            and color(self.sibling.right) == 1
+        ):
+            self.parent.rotate_left()
+            self.grandparent.color = self.parent.color
+            self.parent.color = 0
+            self.parent.sibling.color = 0
+        if (
+            self.is_right()
+            and color(self.sibling) == 0
+            and color(self.sibling.left) == 1
+        ):
+            self.parent.rotate_right()
+            self.grandparent.color = self.parent.color
+            self.parent.color = 0
+            self.parent.sibling.color = 0
+
+    def check_coloring(self) -> bool:
+        """A helper function to recursively check Property 4 of a
+        Red-Black Tree. See check_color_properties for more info.
+        """
+        if self.color == 1 and 1 in (color(self.left), color(self.right)):
+            return False
+        if self.left and not self.left.check_coloring():
+            return False
+        return not (self.right and not self.right.check_coloring())
+
+    def black_height(self) -> int | None:
+        """Returns the number of black nodes from this node to the
+        leaves of the tree, or None if there isn't one such value (the
+        tree is color incorrectly).
+        """
+        if self is None or self.left is None or self.right is None:
+            # If we're already at a leaf, there is no path
+            return 1
+        left = RedBlackTree.black_height(self.left)
+        right = RedBlackTree.black_height(self.right)
+        if left is None or right is None:
+            # There are issues with coloring below children nodes
+            return None
+        if left != right:
+            # The two children have unequal depths
+            return None
+        # Return the black depth of children, plus one if this node is
+        # black
+        return left + (1 - self.color)
+
+    # Here are functions which are general to all binary search trees
+
+    def __contains__(self, label: C) -> bool:
+        """Search through the tree for label, returning True iff it is
+        found somewhere in the tree.
+        Guaranteed to run in O(log(n)) time.
+        """
+        return self.search(label) is not None
+
+    def search(self, label: C) -> RedBlackTree | None:
+        """Search through the tree for label, returning its node if
+        it's found, and None otherwise.
+        This method is guaranteed to run in O(log(n)) time.
+        """
+        if self.label == label:
+            return self
+        elif self.label is not None and label > self.label:
+            if self.right is None:
+                return None
+            else:
+                return self.right.search(label)
+        elif self.left is None:
+            return None
+        else:
+            return self.left.search(label)
+
+    def floor(self, label: C) -> C | None:
+        """Returns the largest element in this tree which is at most label.
+        This method is guaranteed to run in O(log(n)) time."""
+        if self.label == label:
+            return self.label
+        elif self.label is not None and self.label > label:
+            if self.left:
+                return self.left.floor(label)
+            else:
+                return None
+        else:
+            if self.right:
+                attempt = self.right.floor(label)
+                if attempt is not None:
+                    return attempt
+            return self.label
+
+    def ceil(self, label: C) -> C | None:
+        """Returns the smallest element in this tree which is at least label.
+        This method is guaranteed to run in O(log(n)) time.
+        """
+        if self.label == label:
+            return self.label
+        elif self.label is not None and self.label < label:
+            if self.right:
+                return self.right.ceil(label)
+            else:
+                return None
+        else:
+            if self.left:
+                attempt = self.left.ceil(label)
+                if attempt is not None:
+                    return attempt
+            return self.label
+
+    def get_max(self) -> C | None:
+        """Returns the largest element in this tree.
+        This method is guaranteed to run in O(log(n)) time.
+        """
+        if self.right:
+            # Go as far right as possible
+            return self.right.get_max()
+        else:
+            return self.label
+
+    def get_min(self) -> C | None:
+        """Returns the smallest element in this tree.
+        This method is guaranteed to run in O(log(n)) time.
+        """
+        if self.left:
+            # Go as far left as possible
+            return self.left.get_min()
+        else:
+            return self.label
+
+    @property
+    def grandparent(self) -> RedBlackTree | None:
+        """Get the current node's grandparent, or None if it doesn't exist."""
+        if self.parent is None:
+            return None
+        else:
+            return self.parent.parent
+
+    @property
+    def sibling(self) -> RedBlackTree | None:
+        """Get the current node's sibling, or None if it doesn't exist."""
+        if self.parent is None:
+            return None
+        elif self.parent.left is self:
+            return self.parent.right
+        else:
+            return self.parent.left
+
+    def is_left(self) -> bool:
+        """Returns true iff this node is the left child of its parent."""
+        if self.parent is None:
+            return False
+        return self.parent.left is self
+
+    def is_right(self) -> bool:
+        """Returns true iff this node is the right child of its parent."""
+        if self.parent is None:
+            return False
+        return self.parent.right is self
+
+    def __bool__(self) -> bool:
+        return True
+
+    def __len__(self) -> int:
+        """
+        Return the number of nodes in this tree.
+        """
+        ln = 1
+        if self.left:
+            ln += len(self.left)
+        if self.right:
+            ln += len(self.right)
+        return ln
+
+    def preorder_traverse(self) -> Iterator[C | None]:
+        yield self.label
+        if self.left:
+            yield from self.left.preorder_traverse()
+        if self.right:
+            yield from self.right.preorder_traverse()
+
+    def inorder_traverse(self) -> Iterator[C | None]:
+        if self.left:
+            yield from self.left.inorder_traverse()
+        yield self.label
+        if self.right:
+            yield from self.right.inorder_traverse()
+
+    def postorder_traverse(self) -> Iterator[C | None]:
+        if self.left:
+            yield from self.left.postorder_traverse()
+        if self.right:
+            yield from self.right.postorder_traverse()
+        yield self.label
+
+    def __repr__(self) -> str:
+        from pprint import pformat
+
+        if self.left is None and self.right is None:
+            return f"'{self.label} {(self.color and 'red') or 'blk'}'"
+        return pformat(
+            {
+                f"{self.label} {(self.color and 'red') or 'blk'}": (
+                    self.left,
+                    self.right,
+                )
+            },
+            indent=1,
+        )
+
+    def __eq__(self, other: object) -> bool:
+        """Test if two trees are equal."""
+        if not isinstance(other, RedBlackTree):
+            return NotImplemented
+        if self.label == other.label:
+            return self.left == other.left and self.right == other.right
+        else:
+            return False
+
+
+def color(node: RedBlackTree | None) -> int:
+    """Returns the color of a node, allowing for None leaves."""
+    if node is None:
+        return 0
+    else:
+        return node.color

graphical_sampling/sampling/__init__.py

@@ -0,0 +1,6 @@
+from .population import Population
+from .kmeans_spatial_sampling import KMeansSpatialSampling
+from .random_sampling import RandomSampling
+
+
+__all__ = ["Population", "KMeansSpatialSampling", "RandomSampling"]

graphical_sampling/sampling/kmeans_spatial_sampling.py

@@ -0,0 +1,61 @@
+import numpy as np
+from numpy.typing import NDArray
+
+
+from . import Population
+
+
+class KMeansSpatialSampling:
+    def __init__(
+        self,
+        coordinate: NDArray,
+        inclusion_probability: NDArray,
+        *,
+        n: int,
+        n_zones: int | tuple[int, int],
+        tolerance: int,
+    ) -> None:
+        self.coords = coordinate
+        self.probs = inclusion_probability
+        self.n = n
+        self.n_zones = self._pair(n_zones)
+        self.tolerance = tolerance
+
+        self.population = Population(
+            self.coords,
+            self.probs,
+            n_clusters=self.n,
+            n_zones=self.n_zones,
+            tolerance=self.tolerance,
+        )
+        self.rng = np.random.default_rng()
+
+    def _pair(self, n: int | tuple[int, int]) -> tuple[int, int]:
+        if isinstance(n, int):
+            return (n, n)
+        else:
+            return n
+
+    def sample(self, n_samples: int):
+        samples = np.zeros((n_samples, self.n), dtype=int)
+        for i in range(n_samples):
+            random_number = self.rng.random()
+            zone_index = np.searchsorted(
+                np.arange(1, step=round(1 / np.prod(self.n_zones), self.tolerance)),
+                random_number,
+                side="right",
+            )
+            for j, cluster in enumerate(self.population.clusters):
+                unit_index = np.searchsorted(
+                    cluster.zones[zone_index - 1].units[:, 3].cumsum()
+                    + (zone_index - 1)
+                    * round(1 / np.prod(self.n_zones), self.tolerance),
+                    random_number,
+                    side="right",
+                )
+                if np.prod(self.n_zones) != len(cluster.zones):
+                    warning_msg = (f"Warning: Cluster {j} has {len(cluster.zones)} zones, "
+                                f"but expected {np.prod(self.n_zones)} based on n_zones={self.n_zones}")
+                    print(warning_msg)
+                samples[i, j] = cluster.zones[zone_index - 1].units[unit_index, 0]
+        return samples