mathy/pi_finder/engine.py
2025-11-05 12:20:31 +01:00

143 lines
4.0 KiB
Python

from __future__ import annotations
import itertools
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Iterator
from mathstream import StreamNumber, add, sub, mul, div
from mathstream.number import LOG_DIR
@dataclass
class PiApproximation:
iteration: int
digits: str
def _literal(value: int | str) -> StreamNumber:
return StreamNumber(literal=str(value))
_SCALE_CACHE: dict[int, Path] = {}
def _scaled_one(scale_power: int) -> StreamNumber:
if scale_power < 0:
raise ValueError("scale_power must be non-negative")
cached = _SCALE_CACHE.get(scale_power)
if cached and cached.exists():
return StreamNumber(cached)
LOG_DIR.mkdir(parents=True, exist_ok=True)
path = LOG_DIR / f"scale_1e{scale_power}.txt"
if not path.exists():
with open(path, "w", encoding="utf-8") as handle:
handle.write("1")
remaining = scale_power
if remaining > 0:
chunk_size = 1_000_000
zero_chunk = "0" * chunk_size
while remaining > 0:
write_size = min(chunk_size, remaining)
handle.write(zero_chunk[:write_size])
remaining -= write_size
_SCALE_CACHE[scale_power] = path
return StreamNumber(path)
def _format_scaled(value: StreamNumber, scale_power: int, digits: int) -> str:
raw = "".join(value.stream())
negative = raw.startswith("-")
if negative:
raw = raw[1:]
if len(raw) <= scale_power:
raw = raw.zfill(scale_power + 1)
integer_part = raw[:-scale_power] or "0"
fractional_part = raw[-scale_power:] if scale_power else ""
if digits >= 0 and fractional_part:
fractional_part = fractional_part[:digits]
text = integer_part
if digits != 0:
text = f"{integer_part}.{fractional_part or '0'}"
if negative:
text = f"-{text}"
return text
def iterate_nilakantha(
iterations: int | None,
*,
digits: int,
extra_precision: int = 4,
) -> Iterator[PiApproximation]:
"""Yield successive Nilakantha approximations of π using mathstream primitives."""
if iterations is not None and iterations < 0:
raise ValueError("iterations must be non-negative")
if digits < 0:
raise ValueError("digits must be non-negative")
scale_power = digits + extra_precision
scale = _scaled_one(scale_power)
three = _literal(3)
four = _literal(4)
total = mul(three, scale)
numerator = mul(four, scale)
try:
yield PiApproximation(0, _format_scaled(total, scale_power, digits))
counter: Iterable[int]
if iterations is None:
counter = itertools.count(1)
else:
counter = range(1, iterations + 1)
for idx in counter:
two_n = _literal(2 * idx)
two_n_plus_one = _literal(2 * idx + 1)
two_n_plus_two = _literal(2 * idx + 2)
product = mul(two_n, two_n_plus_one)
denominator = mul(product, two_n_plus_two)
term = div(numerator, denominator)
updated = add(total, term) if idx % 2 else sub(total, term)
total.free()
total = updated
yield PiApproximation(idx, _format_scaled(total, scale_power, digits))
for number in (
two_n,
two_n_plus_one,
two_n_plus_two,
product,
denominator,
term,
):
number.free()
finally:
for number in (total, numerator, scale, three, four):
number.free()
def compute_pi_text(
iterations: int,
*,
digits: int,
extra_precision: int = 4,
) -> str:
"""Return the final Nilakantha approximation after the requested iterations."""
approximation: str = "0"
for state in iterate_nilakantha(
iterations,
digits=digits,
extra_precision=extra_precision,
):
approximation = state.digits
return approximation