fresh-cbm 0.2.0__py3-none-any.whl

fresh_cbm/__init__.py

@@ -0,0 +1,237 @@
+"""Exact libcbm stand-state cache client."""
+
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, cast
+
+import numpy as np
+import pandas as pd
+
+from .family import ChainFamily
+from .schema import load_manifest
+
+
+@dataclass(frozen=True)
+class StateCache:
+    """A loaded exact stand-state table.
+
+    Parameters
+    ----------
+    root : pathlib.Path
+        Cache directory containing ``manifest.json``, ``index.parquet``, and
+        ``values/``.
+    manifest : dict
+        Parsed manifest.
+    index : pandas.DataFrame
+        State index in row order.
+    """
+
+    root: Path
+    manifest: dict[str, Any]
+    index: pd.DataFrame
+
+    @classmethod
+    def load(cls, root: str | Path, *, validate: bool = True) -> StateCache:
+        """Load a cache directory.
+
+        Parameters
+        ----------
+        root : path-like
+            Cache directory.
+        validate : bool, optional
+            Check the index, file set, and array shapes. Defaults to True.
+
+        Returns
+        -------
+        StateCache
+            The loaded cache.
+        """
+        path = Path(root)
+        manifest = load_manifest(path / "manifest.json")
+        index = pd.read_parquet(path / "index.parquet")
+        cache = cls(root=path, manifest=manifest, index=index)
+        if validate:
+            cache.validate()
+        return cache
+
+    def validate(self) -> None:
+        """Check the index, file set, offsets, and array shapes.
+
+        Raises
+        ------
+        ValueError
+            When the cache is incomplete or inconsistent with its manifest.
+        """
+        required = {"state_id", "track_id", "track_name", "age"}
+        missing = required - set(self.index.columns)
+        if missing:
+            raise ValueError(f"cache index is missing columns: {sorted(missing)}")
+        states = int(self.manifest["states"])
+        if len(self.index) != states:
+            raise ValueError(f"cache index has {len(self.index)} rows but manifest says {states}")
+        years = self.years
+        accounts = len(self.accounts)
+        expected = (states, years + 1, accounts)
+        family = self.chain_family
+        offsets = self.manifest["offsets"]
+        if not isinstance(offsets, list):
+            raise ValueError("manifest offsets must be a list")
+        if family is not None:
+            if tuple(int(value) for value in offsets) != family.offsets:
+                raise ValueError("chain_family offsets do not match manifest offsets")
+            if bool(self.manifest["empty"]) != family.allow_empty:
+                raise ValueError("chain_family empty policy does not match manifest")
+        files = self.manifest["files"]
+        if not isinstance(files, list):
+            raise ValueError("manifest files must be a list")
+        for name in files:
+            if not isinstance(name, str):
+                raise ValueError("manifest file names must be strings")
+            path = self.root / name
+            if not path.is_file():
+                raise ValueError(f"cache file is missing: {name}")
+            array = np.load(path, mmap_mode="r")
+            if array.shape != expected:
+                raise ValueError(f"{name} has shape {array.shape}; expected {expected}")
+
+    def verify_checksums(self, path: str | Path | None = None) -> None:
+        """Verify files against a SHA256SUMS manifest.
+
+        Parameters
+        ----------
+        path : path-like or None, optional
+            Checksum file. Defaults to ``SHA256SUMS`` in the cache root.
+
+        Raises
+        ------
+        ValueError
+            When a file is missing or a digest does not match.
+        """
+        source = Path(path) if path is not None else self.root / "SHA256SUMS"
+        if not source.is_file():
+            raise ValueError(f"checksum file is missing: {source}")
+        for line in source.read_text(encoding="utf-8").splitlines():
+            if not line.strip():
+                continue
+            digest, name = line.split("  ", 1)
+            target = self.root / name
+            if not target.is_file():
+                raise ValueError(f"checksum target is missing: {name}")
+            actual = hashlib.sha256(target.read_bytes()).hexdigest()
+            if actual != digest:
+                raise ValueError(f"checksum mismatch: {name}")
+
+    @property
+    def chain_family(self) -> ChainFamily | None:
+        """Return the declared chain family, if the manifest has one.
+
+        Returns
+        -------
+        ChainFamily or None
+            The family record, or None for a legacy manifest without one.
+        """
+        raw = self.manifest.get("chain_family")
+        return None if raw is None else ChainFamily.from_manifest(raw)
+
+    @property
+    def accounts(self) -> tuple[str, ...]:
+        """Return the cached account names."""
+        values = self.manifest["accounts"]
+        if not isinstance(values, list):
+            raise ValueError("manifest accounts must be a list")
+        return tuple(str(value) for value in values)
+
+    @property
+    def years(self) -> int:
+        """Return the stored horizon length in years."""
+        return int(self.manifest["years"])
+
+    def _row(self, track_id: int, age: int) -> int:
+        """Return the index row for a track and age."""
+        rows = self.index
+        match = rows[
+            (rows["track_id"].astype(int) == int(track_id)) & (rows["age"].astype(int) == int(age))
+        ]
+        if len(match) != 1:
+            raise KeyError(f"state not in cache: track={track_id} age={age}")
+        return int(match.iloc[0]["state_id"])
+
+    def array(self, chain_offset: int = -1) -> np.ndarray:
+        """Return the memmapped values for one chain offset.
+
+        Parameters
+        ----------
+        chain_offset : int, optional
+            Single ``CC`` event offset; ``-1`` is untreated.
+
+        Returns
+        -------
+        numpy.ndarray
+            The cached values with shape ``(states, years + 1, accounts)``.
+        """
+        name = "empty.npy" if chain_offset < 0 else f"offset_{chain_offset:02d}.npy"
+        return cast(np.ndarray, np.load(self.root / "values" / name, mmap_mode="r"))
+
+    def lookup_state(self, state_id: int, chain_offset: int = -1) -> np.ndarray:
+        """Return one state's ``(years + 1, accounts)`` array.
+
+        Parameters
+        ----------
+        state_id : int
+            Row in ``index.parquet``.
+        chain_offset : int, optional
+            Single ``CC`` event offset; ``-1`` is untreated.
+
+        Returns
+        -------
+        numpy.ndarray
+            The cached accounts.
+        """
+        return np.asarray(self.array(chain_offset)[int(state_id)])
+
+    def lookup(self, track_id: int, age: int, chain_offset: int = -1) -> np.ndarray:
+        """Return one state's ``(years + 1, accounts)`` array.
+
+        Parameters
+        ----------
+        track_id : int
+            Dataset track id.
+        age : int
+            Initial stand age.
+        chain_offset : int, optional
+            Single ``CC`` event offset; ``-1`` is untreated.
+
+        Returns
+        -------
+        numpy.ndarray
+            The cached accounts.
+        """
+        return self.lookup_state(self._row(track_id, age), chain_offset)
+
+    def account(self, track_id: int, age: int, chain_offset: int, account: str) -> np.ndarray:
+        """Return one account series for a state.
+
+        Parameters
+        ----------
+        track_id : int
+            Dataset track id.
+        age : int
+            Initial stand age.
+        chain_offset : int
+            Single ``CC`` event offset; ``-1`` is untreated.
+        account : str
+            Account name.
+
+        Returns
+        -------
+        numpy.ndarray
+            The account values over all stored years.
+        """
+        try:
+            column = self.accounts.index(account)
+        except ValueError as exc:
+            raise KeyError(f"account not in cache: {account}") from exc
+        return self.lookup(track_id, age, chain_offset)[:, column]

fresh_cbm/benchmark.py

@@ -0,0 +1,176 @@
+"""Repeatable cache and consumer-projection benchmarks."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import statistics
+import time
+from collections.abc import Mapping, Sequence
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import pandas as pd
+
+from . import StateCache
+from .effects import Cell, Chain, build_block_effects
+
+
+@dataclass(frozen=True)
+class CacheBenchmark:
+    """Measured cache scale and consumer cost.
+
+    Parameters
+    ----------
+    cache : str
+        Cache directory.
+    states : int
+        Number of cached stand states.
+    accounts : int
+        Account count.
+    offsets : int
+        Chain-offset count.
+    dtype : str
+        Stored dtype.
+    bytes : int
+        Total cache size on disk.
+    load_seconds : float
+        Time to load and validate the cache.
+    lookup_seconds : float
+        Median random state lookup time.
+    lookup_repeats : int
+        Lookup repetitions.
+    blocks : int
+        Blocks in the effect projection.
+    cells : int
+        Target cells in the projection.
+    effect_seconds : float
+        Time to project the blocks onto the cells.
+    """
+
+    cache: str
+    states: int
+    accounts: int
+    offsets: int
+    dtype: str
+    bytes: int
+    load_seconds: float
+    lookup_seconds: float
+    lookup_repeats: int
+    blocks: int
+    cells: int
+    effect_seconds: float
+
+
+def _cache_bytes(root: Path) -> int:
+    """Return the total size of every file under a cache root."""
+    return sum(path.stat().st_size for path in root.rglob("*") if path.is_file())
+
+
+def benchmark_cache(
+    cache_path: str | Path,
+    fragments: pd.DataFrame,
+    chains: Mapping[int, Chain],
+    cells: Sequence[Cell],
+    *,
+    repeats: int = 100,
+    seed: int = 7,
+) -> CacheBenchmark:
+    """Measure load, lookup, projection, and storage.
+
+    Parameters
+    ----------
+    cache_path : path-like
+        Cache directory.
+    fragments : pandas.DataFrame
+        Block/fragment rows for the projection.
+    chains : mapping of int to tuple
+        Current chain per block.
+    cells : sequence of (int, str)
+        Target cells.
+    repeats : int, optional
+        Random lookup repetitions. Defaults to 100.
+    seed : int, optional
+        Random seed. Defaults to seven.
+
+    Returns
+    -------
+    CacheBenchmark
+        The measured result.
+    """
+    if repeats < 1:
+        raise ValueError("repeats must be positive")
+    started = time.perf_counter()
+    cache = StateCache.load(cache_path)
+    load_seconds = time.perf_counter() - started
+    rows = list(range(len(cache.index)))
+    random.Random(seed).shuffle(rows)
+    offsets = [int(value) for value in cache.manifest["offsets"]]
+    samples = rows[:repeats]
+    timings: list[float] = []
+    for row in samples:
+        state = cache.index.iloc[row]
+        offset = random.Random(row).choice(offsets)
+        started = time.perf_counter()
+        cache.lookup(int(state.track_id), int(state.age), offset)
+        timings.append(time.perf_counter() - started)
+    started = time.perf_counter()
+    effects = build_block_effects(cache, fragments, chains, cells)
+    effect_seconds = time.perf_counter() - started
+    return CacheBenchmark(
+        cache=str(Path(cache_path)),
+        states=len(cache.index),
+        accounts=len(cache.accounts),
+        offsets=len(offsets),
+        dtype=str(cache.manifest["dtype"]),
+        bytes=_cache_bytes(Path(cache_path)),
+        load_seconds=load_seconds,
+        lookup_seconds=statistics.median(timings),
+        lookup_repeats=len(timings),
+        blocks=len(effects.blocks),
+        cells=len(effects.cells),
+        effect_seconds=effect_seconds,
+    )
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    """Run the cache benchmark from the command line.
+
+    Parameters
+    ----------
+    argv : sequence of str or None, optional
+        Command-line arguments. None reads ``sys.argv``.
+
+    Returns
+    -------
+    int
+        Process exit status.
+    """
+    parser = argparse.ArgumentParser(prog="fresh-cbm-bench")
+    parser.add_argument("--cache", required=True)
+    parser.add_argument("--dataset", required=True, help="Dataset directory with fragments.parquet")
+    parser.add_argument("--account", default="carbon.system.total")
+    parser.add_argument("--year", type=int, default=49)
+    parser.add_argument("--offset", type=int, default=0)
+    parser.add_argument("--repeats", type=int, default=100)
+    parser.add_argument("--seed", type=int, default=7)
+    parser.add_argument("--out", type=Path)
+    args = parser.parse_args(argv)
+    fragments = pd.read_parquet(Path(args.dataset) / "fragments.parquet")
+    blocks = sorted({int(value) for value in fragments["block_id"]})
+    chains = dict.fromkeys(blocks, (("CC", args.offset),))
+    result = benchmark_cache(
+        args.cache,
+        fragments,
+        chains,
+        [(args.year, args.account)],
+        repeats=args.repeats,
+        seed=args.seed,
+    )
+    payload = asdict(result)
+    print(json.dumps(payload, indent=2))
+    if args.out is not None:
+        args.out.parent.mkdir(parents=True, exist_ok=True)
+        args.out.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
+    return 0

fresh_cbm/build.py

@@ -0,0 +1,203 @@
+"""Build an exact libcbm stand-state cache for one chain family.
+
+The builder intentionally lives in fresh-cbm, not in the WS4 runtime. It imports
+WS4's CarbonLedger so the cached rows are produced by the same exact path the
+optimizer uses.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from collections.abc import Sequence
+from datetime import UTC, datetime
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from .family import ChainFamily
+
+
+def parse_offsets(text: str) -> list[int]:
+    """Parse ``0-49`` or ``0,1,2`` into offsets.
+
+    Parameters
+    ----------
+    text : str
+        Comma-separated offsets or ranges.
+
+    Returns
+    -------
+    list of int
+        Sorted unique offsets.
+    """
+    values: list[int] = []
+    for part in text.split(","):
+        part = part.strip()
+        if not part:
+            continue
+        if "-" in part:
+            start, end = part.split("-", 1)
+            values.extend(range(int(start), int(end) + 1))
+        else:
+            values.append(int(part))
+    if not values:
+        raise ValueError("no offsets selected")
+    return sorted(dict.fromkeys(values))
+
+
+def build(args: argparse.Namespace) -> None:
+    """Build the cache.
+
+    Parameters
+    ----------
+    args : argparse.Namespace
+        Parsed builder arguments.
+
+    Raises
+    ------
+    RuntimeError
+        When WS4 is not importable on the current path.
+    """
+    ws4_path = getattr(args, "ws4_path", None)
+    if ws4_path:
+        sys.path.insert(0, str(ws4_path))
+    try:
+        from ws4.carbon.hwp import load_hwp_config
+        from ws4.carbon.ledger import CarbonLedger
+        from ws4.carbon.libcbm import cbm_defaults_reference, installed_libcbm_version
+        from ws4.carbon.package import load_carbon_package, package_fingerprint
+    except ImportError as exc:  # pragma: no cover - depends on optional installation
+        raise RuntimeError(
+            "fresh-cbm-build needs WS4 with the carbon extra; pass --ws4-path "
+            "or set PYTHONPATH to the WS4 source root"
+        ) from exc
+
+    out = Path(args.out)
+    (out / "values").mkdir(parents=True, exist_ok=True)
+    package = load_carbon_package(Path(args.package))
+    hwp = load_hwp_config(Path(args.hwp_config)) if args.hwp_config else None
+    ledger = CarbonLedger(
+        Path(args.dataset),
+        package,
+        years=args.years,
+        hwp=hwp,
+        base_workers=args.workers,
+    )
+    fragments = ledger.fragments
+    states = (
+        fragments.loc[:, ["track_id", "age"]]
+        .drop_duplicates()
+        .sort_values(["track_id", "age"], kind="stable")
+        .reset_index(drop=True)
+    )
+    if args.limit is not None:
+        states = states.head(args.limit).reset_index(drop=True)
+    track_names = {
+        int(str(row.track_id)): str(row.name) for row in ledger.tracks.itertuples(index=False)
+    }
+    index = pd.DataFrame(
+        {
+            "state_id": np.arange(len(states), dtype=np.int64),
+            "track_id": states["track_id"].astype(int),
+            "track_name": [track_names.get(int(value), "") for value in states["track_id"]],
+            "age": states["age"].astype(int),
+        }
+    )
+    index.to_parquet(out / "index.parquet", index=False)
+    family = ChainFamily(
+        treatment=args.treatment,
+        offsets=tuple(parse_offsets(args.offsets)),
+        events=args.events,
+        allow_empty=not args.no_empty,
+    )
+    offsets = family.offsets
+    build_offsets = [-1, *offsets] if family.allow_empty else list(offsets)
+    started = time.perf_counter()
+    accounts: tuple[str, ...] | None = None
+    for offset in build_offsets:
+        chain = () if offset < 0 else family.chain(offset)
+        keys = {(int(row.track_id), int(row.age), chain) for row in states.itertuples(index=False)}
+        step = time.perf_counter()
+        ledger._prepare_units(keys)
+        if accounts is None:
+            first = ledger._unit_cache[next(iter(keys))]
+            accounts = tuple(sorted({account for _offset, account in first}))
+        values = np.zeros((len(states), args.years + 1, len(accounts)), dtype=args.dtype)
+        for row, state in enumerate(states.itertuples(index=False)):
+            unit = ledger._unit_cache[(int(state.track_id), int(state.age), chain)]
+            for offset_index in range(args.years + 1):
+                for account_index, account in enumerate(accounts):
+                    values[row, offset_index, account_index] = unit.get(
+                        (offset_index, account), 0.0
+                    )
+        name = "empty.npy" if offset < 0 else f"offset_{offset:02d}.npy"
+        np.save(out / "values" / name, values, allow_pickle=False)
+        ledger._unit_cache.clear()
+        ledger._cache.clear()
+        print(f"offset={offset:>2} states={len(states)} seconds={time.perf_counter() - step:.1f}")
+    manifest = {
+        "schema": "fresh-cbm-state-v1",
+        "created": datetime.now(UTC).isoformat(),
+        "dataset": Path(args.dataset).name,
+        "package": Path(args.package).name,
+        "package_fingerprint": package_fingerprint(package),
+        "libcbm_version": installed_libcbm_version(),
+        "cbm_defaults_sha256": cbm_defaults_reference().sha256,
+        "years": args.years,
+        "dtype": str(np.dtype(args.dtype)),
+        "accounts": list(accounts or ()),
+        "offsets": list(offsets),
+        "empty": family.allow_empty,
+        "chain_family": family.as_manifest(),
+        "states": len(states),
+        "files": [
+            *(["values/empty.npy"] if family.allow_empty else []),
+            *[f"values/offset_{offset:02d}.npy" for offset in offsets],
+        ],
+        "build_seconds": time.perf_counter() - started,
+        "hwp_config": Path(args.hwp_config).name if args.hwp_config else None,
+    }
+    (out / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
+    print(f"wrote {out} in {time.perf_counter() - started:.1f}s")
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    """Parse arguments and build the cache.
+
+    Parameters
+    ----------
+    argv : sequence of str or None, optional
+        Command-line arguments. None reads ``sys.argv``.
+
+    Returns
+    -------
+    int
+        Process exit status.
+    """
+    parser = argparse.ArgumentParser(prog="fresh-cbm-build")
+    parser.add_argument("--dataset", required=True)
+    parser.add_argument("--package", required=True)
+    parser.add_argument("--hwp-config")
+    parser.add_argument("--years", type=int, default=50)
+    parser.add_argument("--offsets", default="0-49")
+    parser.add_argument("--treatment", default="CC")
+    parser.add_argument("--events", type=int, default=1)
+    parser.add_argument("--no-empty", action="store_true")
+    parser.add_argument(
+        "--ws4-path",
+        help="WS4 source root to add to sys.path when running the builder",
+    )
+    parser.add_argument("--workers", type=int, default=32)
+    parser.add_argument("--dtype", choices=["float32", "float64"], default="float64")
+    parser.add_argument("--limit", type=int)
+    parser.add_argument("--out", required=True)
+    try:
+        build(parser.parse_args(argv))
+    except RuntimeError as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 1
+    return 0