Differentially Private Spatial Range Counts

from __future__ import annotations

import math
import random
from dataclasses import dataclass, field
from typing import Sequence


@dataclass
class HierarchicalCounts:
    """A one-dimensional hierarchy over `leaves` cells; compose two for a grid.

    Real deployments run this over a space-filling-curve ordering (H3 parent/child,
    or a Z-order index), which reduces a rectangle query to a small number of
    contiguous ranges — the same structure this class answers.
    """

    leaves: int
    fanout: int
    epsilon: float
    levels: list[list[float]] = field(default_factory=list)

    def build(self, counts: Sequence[float], *, seed: int = 0) -> "HierarchicalCounts":
        """Noise every level under an equal split of the total budget."""
        if len(counts) != self.leaves:
            raise ValueError("counts must cover every leaf")
        rng = random.Random(seed)
        depth = max(1, math.ceil(math.log(self.leaves, self.fanout)))
        per_level_eps = self.epsilon / depth
        scale = 1.0 / per_level_eps

        level = list(counts)
        self.levels = []
        while True:
            noisy = []
            for value in level:
                u = rng.random() - 0.5
                noisy.append(value - scale * math.copysign(1.0, u) * math.log1p(-2 * abs(u)))
            self.levels.append(noisy)
            if len(level) <= 1:
                break
            level = [sum(level[i:i + self.fanout]) for i in range(0, len(level), self.fanout)]
        self.levels.reverse()          # root first
        return self

    def range_count(self, lo: int, hi: int) -> float:
        """Answer [lo, hi) by covering it with the fewest nodes available.

        Greedy cover: take the largest aligned node that fits, then recurse. This is
        what keeps error at sqrt(log m) rather than sqrt(m) — the whole point of the
        structure.
        """
        if not 0 <= lo <= hi <= self.leaves:
            raise ValueError("range outside the domain")
        total = 0.0
        i = lo
        while i < hi:
            best_level, best_span = len(self.levels) - 1, 1
            for depth, nodes in enumerate(self.levels):
                span = self.fanout ** (len(self.levels) - 1 - depth)
                if span and i % span == 0 and i + span <= hi and span > best_span:
                    best_level, best_span = depth, span
            node_index = i // best_span
            if node_index < len(self.levels[best_level]):
                total += self.levels[best_level][node_index]
            i += best_span
        return total

    def enforce_consistency(self) -> None:
        """Project so each parent equals the sum of its children (free, and helpful).

        A single top-down pass distributes each parent's residual equally among its
        children. This is post-processing: it consumes no budget, makes answers
        mutually consistent, and reduces mean squared error.
        """
        for depth in range(len(self.levels) - 1):
            parents, children = self.levels[depth], self.levels[depth + 1]
            for p_idx, parent in enumerate(parents):
                start = p_idx * self.fanout
                block = children[start:start + self.fanout]
                if not block:
                    continue
                residual = (parent - sum(block)) / len(block)
                for offset in range(len(block)):
                    children[start + offset] += residual

def _validate() -> None:
    rng = random.Random(5)
    leaves = 256
    truth = [max(0.0, rng.lognormvariate(4.0, 1.0)) for _ in range(leaves)]

    tree = HierarchicalCounts(leaves=leaves, fanout=4, epsilon=1.0).build(truth, seed=1)

    # 1. A wide range must be far more accurate than the sqrt(m) flat-grid bound.
    wide_true = sum(truth)
    wide_err = abs(tree.range_count(0, leaves) - wide_true)
    flat_bound = math.sqrt(leaves) * math.sqrt(2.0) / 1.0
    assert wide_err < flat_bound, (wide_err, flat_bound)

    # 2. Consistency projection must not change the root's meaning, and must make
    #    parents equal the sum of their children.
    tree.enforce_consistency()
    for depth in range(len(tree.levels) - 1):
        parents, children = tree.levels[depth], tree.levels[depth + 1]
        for p_idx, parent in enumerate(parents):
            block = children[p_idx * 4:(p_idx + 1) * 4]
            if block:
                assert abs(parent - sum(block)) < 1e-6

    # 3. Two decompositions of the same range must now agree exactly.
    whole = tree.range_count(0, 64)
    halves = tree.range_count(0, 32) + tree.range_count(32, 64)
    assert abs(whole - halves) < 1e-6

    # 4. A narrow range is answered from few nodes, so its error stays bounded.
    narrow_err = abs(tree.range_count(10, 14) - sum(truth[10:14]))
    assert narrow_err < 40.0, narrow_err

    # 5. More budget must reduce error for the same query.
    rich = HierarchicalCounts(leaves=leaves, fanout=4, epsilon=8.0).build(truth, seed=1)
    assert abs(rich.range_count(0, 64) - sum(truth[:64])) < abs(
        HierarchicalCounts(leaves=leaves, fanout=4, epsilon=0.5)
        .build(truth, seed=1).range_count(0, 64) - sum(truth[:64])
    )

    # 6. An out-of-domain range must raise rather than answer.
    try:
        tree.range_count(-1, 10)
    except ValueError:
        pass
    else:
        raise AssertionError("out-of-domain ranges must raise")

    print("hierarchical range counts: all assertions passed")


_validate()