prime_backend_v2.py

"""prime_backend_v2 — blockwise von Mangoldt cumsum backend (task 2.8ii).

E-016's in-memory sieve materializes C1[m] = sum_{n<=m} LAM(n)/sqrt(n)
and C0[m] = sum LAM(n) log(n)/sqrt(n) for ALL n <= e^20 (~12 GB peak).
a = 12 needs e^24 = 2.65e10 — impossible in memory. This module sieves
blockwise (base primes to sqrt(hi)), keeps only running totals, and
answers cumsum queries at arbitrary m via block-local partial sums.

Usage:
    be = PrimeBackendV2(t_queries)          # t_queries: array of t >= 0
    P = be.prime_sum(t_queries)             # = C1(m)*t - C0(m), m = floor(e^t)-1
validated against E-016's in-memory backend in E-019 (sub-stage A).
"""
import numpy as np


class PrimeBackendV2:
    def __init__(self, t_queries, block_exp=24):
        t = np.asarray(t_queries, dtype=float)
        t = t[t > 0]
        # E-016 convention: sums run over n <= floor(e^t + eps)
        m = np.floor(np.exp(t) + 1e-9).astype(np.int64)
        m = m[m >= 2]
        self.m_sorted = np.unique(m)
        self.hi = int(self.m_sorted[-1]) + 1 if self.m_sorted.size else 2
        self.B = max(1 << 24, int(np.sqrt(self.hi)) + 2)   # >= sqrt(hi)
        self._c1 = np.empty_like(self.m_sorted, dtype=float)
        self._c0 = np.empty_like(self.m_sorted, dtype=float)
        self._built = False

    def _base_primes(self, n):
        sieve = np.ones(n + 1, dtype=bool)
        sieve[:2] = False
        for i in range(2, int(n ** 0.5) + 1):
            if sieve[i]:
                sieve[i * i::i] = False
        return np.flatnonzero(sieve)

    def save_cache(self, path):
        np.savez_compressed(path, m_sorted=self.m_sorted,
                            c1=self._c1, c0=self._c0, hi=self.hi)

    def load_cache(self, path):
        d = np.load(path)
        assert int(d["hi"]) >= self.hi and np.array_equal(
            d["m_sorted"], self.m_sorted), "cache mismatch"
        self._c1, self._c0 = d["c1"], d["c0"]
        self._built = True
        return self

    def build(self, verbose=True):
        """Blockwise sieve: answer C1/C0 at every m in self.m_sorted."""
        hi = self.hi
        B = self.B
        base = self._base_primes(int(np.sqrt(hi)) + 1)
        run1 = 0.0   # running sum LAM/sqrt(n)
        run0 = 0.0   # running sum LAM log(n)/sqrt(n)
        qi = 0
        Q = self.m_sorted
        nblocks = (hi + B - 1) // B
        for bi, lo in enumerate(range(2, hi + 1, B)):
            hib = min(lo + B - 1, hi)
            size = hib - lo + 1
            # exact segmented von Mangoldt:
            # 1) mark composites with a factor p <= sqrt(hib)
            comp = np.zeros(size, dtype=bool)
            for p in base:
                if p * p > hib:
                    break
                first = max(p * p, ((lo + p - 1) // p) * p)
                comp[first - lo::p] = True
            ns64 = np.arange(lo, hib + 1, dtype=np.int64)
            lam = np.zeros(size)
            # 2) survivors are exactly the primes in [lo, hib]
            lam[~comp] = np.log(ns64[~comp])
            # 3) higher prime powers p^k (k>=2) by INDEX (not multiples)
            for p in base:
                if p * p > hib:
                    break
                lp = np.log(p)
                pk = p * p
                while pk <= hib:
                    if pk >= lo:
                        lam[pk - lo] = lp
                    pk *= p
            # cumsums inside block
            ns = ns64.astype(np.float64)
            contrib1 = lam / np.sqrt(ns)
            contrib0 = contrib1 * np.log(ns)
            c1 = run1 + np.cumsum(contrib1)
            c0 = run0 + np.cumsum(contrib0)
            # answer queries inside this block
            while qi < Q.size and Q[qi] <= hib:
                self._c1[qi] = c1[Q[qi] - lo]
                self._c0[qi] = c0[Q[qi] - lo]
                qi += 1
            run1, run0 = float(c1[-1]), float(c0[-1])
            if verbose and (bi % 10 == 0 or bi == nblocks - 1):
                print(f"    block {bi + 1}/{nblocks} (n <= {hib:.3e}) "
                      f"queries answered {qi}/{Q.size}", flush=True)
            if qi >= Q.size:
                break
        self._built = True
        return self

    def cumsums(self, t):
        """C1(m), C0(m) at m = floor(e^t)-1 for array t."""
        assert self._built
        t = np.asarray(t, dtype=float)
        out1 = np.zeros_like(t)
        out0 = np.zeros_like(t)
        pos = t > 0
        m = np.floor(np.exp(t[pos]) + 1e-9).astype(np.int64)
        valid = m >= 2
        idx = np.searchsorted(self.m_sorted, m[valid])
        idx = np.clip(idx, 0, self.m_sorted.size - 1)
        o1 = np.zeros(int(pos.sum()))
        o0 = np.zeros(int(pos.sum()))
        hit = self.m_sorted[idx] == m[valid]
        o1[valid] = np.where(hit, self._c1[idx], np.nan)
        o0[valid] = np.where(hit, self._c0[idx], np.nan)
        if not np.all(hit):
            missing = m[valid][~hit]
            raise ValueError(f"unsieved m values, e.g. {missing[:5]}")
        out1[pos] = o1
        out0[pos] = o0
        return out1, out0

    def prime_sum(self, t):
        """Drop-in replacement for E-016 prime_sum_vec."""
        c1, c0 = self.cumsums(t)
        t = np.asarray(t, dtype=float)
        return c1 * t - c0