Example code

"""A reusable pulsed reactor from the repeated-harvesting paper (Table 1).

All quantities use the paper's normalized units. This is a floating-point
experiment, not a replay of the Lean proof. Run: python example.py --output outputs
"""
from __future__ import annotations

import argparse
import csv
from dataclasses import asdict, dataclass, replace
import hashlib
import json
from pathlib import Path
import platform
import time
from typing import Callable

import numpy as np
import scipy
from scipy.integrate import solve_ivp

# USER INPUTS ---------------------------------------------------------------
# Paper preset: slow release / fast cleavage corner of Eq. (1).
RELEASE_SPEED = 19.0                 # normalized inverse time; theorem: [19, 21]
CLEAVAGE_SPEED = 0.04                # normalized inverse time; theorem: [.02, .04]
# Eq. (8), species order U, W, X, C1, C2, Z; normalized concentrations >= 0.
INITIAL_STATE = (0.934, 0.935, 0.06, 0.001, 0.001, 0.001)
RETAINED_FRACTIONS = (0.25, 0.75, 0.4, 0.6)  # paper diagnostic schedule
LOSS_PATTERNS = ((0.98,) * 6, (0.98, 1., 0.98, 1., 0.99, 1.))
REFILL_ERRORS = (-0.005, 0.005)       # absolute normalized food concentration
CYCLES = 32                         # extends paper's 8-cycle diagnostic
CONDITIONING_TIME = 12.0             # paper schedule, normalized time
RECOVERY_TIME = 3.0
COLLECTION_TIME = 1.0
RETENTION_SWEEP = (0.25, 0.375, 0.5, 0.625, 0.75)
RTOL, ATOL = 1e-9, 1e-12            # numerical settings, not scientific inputs
SAMPLES_PER_UNIT = 50                # saved curve resolution; not solver steps
PAPER_SHA256 = 'f90ee0cd4daf3dfaad3dfffcd0ab6217db857971c5396b0520383fdc5f0f5979'
# --------------------------------------------------------------------------

SPECIES = ('U', 'W', 'X', 'C1', 'C2', 'Z')
A = np.array([1, 0, 1, 2, 2, 2.])  # first elemental inventory
B = np.array([0, 1, 1, 1, 2, 2.])  # second elemental inventory
Y = np.array([0, 0, 1, 9/8, 7/5, 9/5])  # catalytic growth observable
I = np.array([0, 0, 1, 1, 1, 2.])  # covalent template inventory


def state_array(values):
    c = np.asarray(values, dtype=float).copy()
    if c.shape != (6,) or not np.all(np.isfinite(c)) or np.any(c < 0):
        raise ValueError('State must contain six finite nonnegative concentrations.')
    return c


@dataclass(frozen=True)
class Rates:
    release: float = RELEASE_SPEED
    cleavage: float = CLEAVAGE_SPEED
    epsilon: float = 1/500_000_000
    eta: float = 1/8_000_000_000

    def __post_init__(self):
        if any(not np.isfinite(v) or v < 0 for v in asdict(self).values()):
            raise ValueError('Rate constants must be finite and nonnegative.')

    @property
    def in_paper_box(self):
        return (19 <= self.release <= 21 and .02 <= self.cleavage <= .04
                and self.epsilon == 1/500_000_000 and self.eta == 1/8_000_000_000)


@dataclass(frozen=True)
class ReactionPair:
    """Mass action with reservoir activities already included in the rates."""
    name: str
    reactants: tuple[int, ...]
    products: tuple[int, ...]
    forward: float
    reverse: float

    def flux(self, c):
        return self.forward * np.prod(c ** self.reactants) - self.reverse * np.prod(c ** self.products)


class Reactor:
    """Literal chemistry composed from six reversible pairs and maintained flow."""
    def __init__(self, rates=Rates()):
        self.rates = rates
        self.pairs = (
            ReactionPair('basal', (1,1,0,0,0,0), (0,0,1,0,0,0), rates.epsilon, rates.epsilon/10),
            ReactionPair('bind U', (1,0,1,0,0,0), (0,0,0,1,0,0), 20, 20),
            ReactionPair('bind W', (0,1,0,1,0,0), (0,0,0,0,1,0), 20, 20),
            ReactionPair('ligate', (0,0,0,0,1,0), (0,0,0,0,0,1), 20, 2),
            ReactionPair('release', (0,0,0,0,0,1), (0,0,2,0,0,0), rates.release, rates.release),
            ReactionPair('cleave', (0,0,1,0,0,0), (1,1,0,0,0,0), rates.cleavage, rates.cleavage*rates.eta),
        )
        self.stoichiometry = np.array([np.subtract(p.products, p.reactants) for p in self.pairs]).T
        self.feed = np.array([1., 1., 0, 0, 0, 0])

    def fluxes(self, c):
        return np.array([p.flux(c) for p in self.pairs])

    def derivative(self, c):
        return self.feed - c + self.stoichiometry @ self.fluxes(c)

    def augmented_derivative(self, t, v):
        c = v[:6]
        j = self.fluxes(c)
        # Four integrated counters: all-time effluent, free X, gross service,
        # and net covalent synthesis. No clipping of concentrations or fluxes.
        service = self.rates.cleavage * (c[2] + self.rates.eta*c[0]*c[1])
        return np.r_[self.feed-c+self.stoichiometry@j, I@c, c[2], service, j[0]+j[3]-j[5]]


@dataclass(frozen=True)
class Intervention:
    retained: float = 0.25
    survival: tuple[float, ...] = LOSS_PATTERNS[0]
    food_error: tuple[float, float] = (-0.005, -0.005)

    def __post_init__(self):
        if not np.isfinite(self.retained) or not 0 <= self.retained <= 1:
            raise ValueError('Retained fraction must be between zero and one.')
        if len(self.survival) != 6 or any(not np.isfinite(x) or not 0 <= x <= 1 for x in self.survival):
            raise ValueError('Six survival fractions between zero and one are required.')
        if len(self.food_error) != 2 or any(not np.isfinite(e) or 1-self.retained+e < 0 for e in self.food_error):
            raise ValueError('Both food refills must be finite and nonnegative.')

    @property
    def in_paper_box(self):
        return (.25 <= self.retained <= .75 and min(self.survival) >= .98
                and max(abs(e) for e in self.food_error) <= .005)

    def apply(self, before):
        c = state_array(before)
        withdrawn = (1-self.retained)*c
        handling_loss = self.retained*(1-np.array(self.survival))*c
        after = self.retained*np.array(self.survival)*c
        after[:2] += self.refill
        return after, withdrawn, handling_loss

    @property
    def refill(self):
        return 1-self.retained+np.array(self.food_error)


@dataclass(frozen=True)
class Schedule:
    conditioning: float = CONDITIONING_TIME
    recovery: float = RECOVERY_TIME
    collection: float = COLLECTION_TIME
    cycles: int = CYCLES

    def __post_init__(self):
        if any(not np.isfinite(x) or x <= 0 for x in (self.conditioning, self.recovery, self.collection)):
            raise ValueError('All schedule durations must be positive and finite.')
        if not isinstance(self.cycles, int) or isinstance(self.cycles, bool) or self.cycles < 1:
            raise ValueError('At least one integer cycle is required.')

    @property
    def is_paper_schedule(self):
        return (self.conditioning, self.recovery, self.collection) == (12, 3, 1)


@dataclass
class Segment:
    time: np.ndarray
    values: np.ndarray
    end: np.ndarray
    collection: np.ndarray
    material_error: float
    inventory_error: float
    minimum_concentration: float
    threshold_time: float | None


class Integrator:
    def __init__(self, reactor: Reactor, rtol=RTOL, atol=ATOL, method='Radau'):
        if not (np.isfinite(rtol) and np.isfinite(atol) and rtol > 0 and atol > 0):
            raise ValueError('Solver tolerances must be positive and finite.')
        self.reactor, self.rtol, self.atol, self.method = reactor, rtol, atol, method

    def advance(self, initial, duration, collection_start=None):
        c = state_array(initial)
        if not np.isfinite(duration) or duration <= 0:
            raise ValueError('Duration must be positive and finite.')
        if collection_start is not None and not 0 <= collection_start <= duration:
            raise ValueError('Collection must begin within the segment.')
        def threshold(t, v):
            return Y@v[:6] - .05
        threshold.direction = 1
        solution = solve_ivp(self.reactor.augmented_derivative, (0, duration), np.r_[c, np.zeros(4)],
                             method=self.method, rtol=self.rtol, atol=self.atol,
                             events=threshold, dense_output=True)
        if not solution.success:
            raise RuntimeError(solution.message)
        times = np.unique(np.r_[np.linspace(0, duration, int(duration*SAMPLES_PER_UNIT)+1),
                                [] if collection_start is None else [collection_start]])
        v = solution.sol(times)
        minimum = float(min(v[:6].min(), solution.y[:6].min()))
        if minimum < -10*self.atol:
            raise ArithmeticError(f'Negative concentration {minimum}; refine the numerical method.')
        material_error = max(np.max(np.abs(w@v[:6]-(1+(w@c-1)*np.exp(-times)))) for w in (A,B))
        inventory_error = np.max(np.abs(I@v[:6]-I@c+v[6]-v[9]))
        collected = np.zeros(4) if collection_start is None else v[6:,-1]-solution.sol(collection_start)[6:]
        hits = solution.t_events[0]
        return Segment(times, v, v[:,-1], collected, float(material_error), float(inventory_error),
                       minimum, 0. if Y@c >= .05 else (float(hits[0]) if len(hits) else None))


def paper_protocol(cycle: int, before: np.ndarray) -> Intervention:
    """A callable may also choose its intervention from the actual pre-pulse state."""
    e = REFILL_ERRORS[cycle % len(REFILL_ERRORS)]
    return Intervention(RETAINED_FRACTIONS[cycle % len(RETAINED_FRACTIONS)],
                        LOSS_PATTERNS[cycle % len(LOSS_PATTERNS)], (e,e))


def in_region(c, operating=False):
    lo, hi, floor = (159/160,161/160,1/20) if operating else (.9,1.1,1/5000)
    return bool(lo <= A@c <= hi and lo <= B@c <= hi and Y@c >= floor)


@dataclass
class Operation:
    cycles: list[dict]
    trajectories: list[dict]
    pulses: list[dict]
    summary: dict


def operate(integrator: Integrator, initial=INITIAL_STATE, schedule=Schedule(),
            protocol: Callable[[int, np.ndarray], Intervention]=paper_protocol,
            conditioning_pulse=Intervention()):
    """Continue actual endpoints; never replace a disturbed state with a preset."""
    c = state_array(initial)
    initial_inventory = float(I@c)
    applicable = integrator.reactor.rates.in_paper_box and in_region(c) and schedule.is_paper_schedule
    rows, trajectories, pulses, segments = [], [], [], []
    totals = dict(effluent=0., synthesis=0., withdrawn=0., handling_loss=0., collected=0.,
                  free_collected=0., service=0., food_U=0., food_W=0.)
    elapsed = 0.
    for n in range(-1, schedule.cycles):
        pulse = conditioning_pulse if n == -1 else protocol(n, c.copy())
        applicable = applicable and pulse.in_paper_box
        after, withdrawn, loss = pulse.apply(c)
        duration = schedule.conditioning if n == -1 else schedule.recovery+schedule.collection
        segment = integrator.advance(after, duration, None if n == -1 else schedule.recovery)
        segments.append(segment)
        pulse_row = dict(cycle=n+1, time=elapsed, retained=pulse.retained,
                         withdrawn_inventory=float(I@withdrawn), handling_loss_inventory=float(I@loss))
        for k, name in enumerate(SPECIES):
            pulse_row[f'before_{name}'] = float(c[k])
            pulse_row[f'after_{name}'] = float(after[k])
        pulses.append(pulse_row)
        for t, v in zip(segment.time, segment.values.T):
            trajectories.append(dict(cycle=n+1, time=float(elapsed+t), local_time=float(t),
                                     **dict(zip(SPECIES, map(float,v[:6]))),
                                     Y=float(Y@v[:6]), inventory=float(I@v[:6]),
                                     A=float(A@v[:6]), B=float(B@v[:6])))
        totals['effluent'] += float(segment.end[6])
        totals['synthesis'] += float(segment.end[9])
        totals['withdrawn'] += float(I@withdrawn)
        totals['handling_loss'] += float(I@loss)
        totals['service'] += float(segment.end[8])
        totals['food_U'] += float(duration+pulse.refill[0])
        totals['food_W'] += float(duration+pulse.refill[1])
        totals['collected'] += float(segment.collection[0])
        totals['free_collected'] += float(segment.collection[1])
        c = segment.end[:6].copy()
        balance = float(I@c-initial_inventory+totals['effluent']+totals['withdrawn']
                        +totals['handling_loss']-totals['synthesis'])
        if abs(balance) > 1e-7:
            raise ArithmeticError(f'Full pulse/flow inventory balance residual: {balance}')
        if n >= 0:
            rows.append(dict(cycle=n+1, retained=pulse.retained, template_collected=float(segment.collection[0]),
                             free_X_collected=float(segment.collection[1]), gross_service=float(segment.end[8]),
                             food_U=float(duration+pulse.refill[0]), food_W=float(duration+pulse.refill[1]),
                             Y_end=float(Y@c), cumulative_collected=totals['collected'],
                             cumulative_synthesis=totals['synthesis'], inventory_balance_residual=balance,
                             paper_synthesis_lower_bound=(n+1)/28-1.1 if applicable else None,
                             end_in_operating_region=in_region(c, True)))
        elapsed += duration
    summary = dict(theorem_inputs_satisfied=bool(applicable), initial_inventory=initial_inventory,
                   final_inventory=float(I@c), **totals, final_inventory_balance_residual=balance,
                   maximum_material_error=max(s.material_error for s in segments),
                   maximum_segment_inventory_error=max(s.inventory_error for s in segments),
                   minimum_sampled_concentration=min(s.minimum_concentration for s in segments),
                   minimum_cycle_template_output=min(r['template_collected'] for r in rows),
                   minimum_cycle_free_X_output=min(r['free_X_collected'] for r in rows),
                   all_cycle_endpoints_in_operating_region=all(r['end_in_operating_region'] for r in rows))
    return Operation(rows, trajectories, pulses, summary)


def retention_sweep(integrator):
    """Independent runs from the same Eq. (8) preparation; no warm starts."""
    result = []
    for q in RETENTION_SWEEP:
        pulse = Intervention(q)
        segment = integrator.advance(pulse.apply(INITIAL_STATE)[0], RECOVERY_TIME+COLLECTION_TIME, RECOVERY_TIME)
        result.append(dict(retained=q, threshold_time=segment.threshold_time,
                           template_collected=float(segment.collection[0]), free_X_collected=float(segment.collection[1]),
                           gross_service=float(segment.end[8]), food_U=RECOVERY_TIME+COLLECTION_TIME+pulse.refill[0]))
    return result


def write_csv(path, rows):
    with path.open('w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def figures(operation, sweep, output):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    with plt.rc_context({'font.size':11, 'axes.spines.top':False, 'axes.spines.right':False,
                         'figure.facecolor':'white', 'savefig.facecolor':'white', 'svg.fonttype':'none'}):
        fig, axes = plt.subplots(2, 1, figsize=(7.2,6.4), layout='constrained')
        # Show eight actual cycles after conditioning, retaining all 32 in CSV.
        shown = [r for r in operation.trajectories if 1 <= r['cycle'] <= min(CYCLES,8)]
        t = [r['time']-CONDITIONING_TIME for r in shown]
        axes[0].plot(t, [r['Y'] for r in shown], color='#0072B2', label='Weighted catalytic stock Y')
        axes[0].plot(t, [r['X'] for r in shown], color='#D55E00', linestyle='--', label='Free template X')
        axes[0].axhline(.05, color='#555555', linestyle=':', label='Paper stock floor')
        for n in range(min(CYCLES,8)):
            axes[0].axvspan(n*(RECOVERY_TIME+COLLECTION_TIME)+RECOVERY_TIME,
                            (n+1)*(RECOVERY_TIME+COLLECTION_TIME), color='#e3e6e8', alpha=.7)
        axes[0].set(xlabel='Time after conditioning (normalized)', ylabel='Normalized concentration', ylim=(0,None))
        axes[0].legend(fontsize=9, loc='lower left', bbox_to_anchor=(0,1.01), ncol=2, frameon=False)
        n = [r['cycle'] for r in operation.cycles]
        axes[1].plot(n, [r['template_collected'] for r in operation.cycles], 'o-', color='#0072B2', markersize=3, label='Template equivalents')
        axes[1].plot(n, [r['free_X_collected'] for r in operation.cycles], 's--', color='#D55E00', markersize=3, label='Free X')
        if operation.summary['theorem_inputs_satisfied']:
            axes[1].axhline(1/28,color='#0072B2',linestyle=':',label='Paper floor 1/28')
            axes[1].axhline(1/540,color='#D55E00',linestyle=':',label='Paper floor 1/540')
        axes[1].set(xlabel='Routine cycle', ylabel='Collected amount (normalized)', ylim=(0,None))
        axes[1].legend(fontsize=9, ncol=2, loc='lower left', bbox_to_anchor=(0,1.01), frameon=False)
        for ax in axes:
            ax.grid(axis='y',color='#e3e6e8'); ax.set_axisbelow(True)
        fig.savefig(output/'recovery.png',dpi=220); fig.savefig(output/'recovery.svg'); plt.close(fig)
        fig, ax = plt.subplots(figsize=(7.2,4.5),layout='constrained')
        ax.plot(n,[r['cumulative_synthesis'] for r in operation.cycles],color='#0072B2',label='Net synthesis, including conditioning')
        ax.plot(n,[r['cumulative_collected'] for r in operation.cycles],'--',color='#D55E00',label='Collected template equivalents')
        if operation.summary['theorem_inputs_satisfied']:
            ax.plot(n,[r['paper_synthesis_lower_bound'] for r in operation.cycles],':',color='#3B3B3B',label='Paper synthesis lower bound m/28 - 1.1')
        ax.axhline(0,color='#999999',linewidth=.7)
        ax.set(xlabel='Routine cycles completed',ylabel='Cumulative amount (normalized)')
        ax.legend(fontsize=9); ax.grid(axis='y',color='#e3e6e8')
        fig.savefig(output/'synthesis.png',dpi=220); fig.savefig(output/'synthesis.svg'); plt.close(fig)


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output',type=Path,default=Path('outputs'))
    args=parser.parse_args(); args.output.mkdir(parents=True,exist_ok=True)
    start=time.perf_counter()
    rates=Rates(); schedule=Schedule(); solver=Integrator(Reactor(rates))
    inputs=dict(rates=asdict(rates),schedule=asdict(schedule),initial_state=INITIAL_STATE,
                retained_fractions=RETAINED_FRACTIONS,loss_patterns=LOSS_PATTERNS,refill_errors=REFILL_ERRORS,
                conditioning_pulse=asdict(Intervention()),rtol=RTOL,atol=ATOL,method=solver.method,
                units='Normalized, uncalibrated schematic chemistry; see README for dimensional conversion.')
    print('Resolved inputs: '+json.dumps(inputs,allow_nan=False))
    operation=operate(solver); sweep=retention_sweep(solver)
    # Independent numerical method plus stricter tolerances for one full cycle.
    post=Intervention().apply(INITIAL_STATE)[0]
    base=solver.advance(post,4,3)
    strict=Integrator(Reactor(rates),rtol=1e-11,atol=1e-14,method='DOP853').advance(post,4,3)
    comparison=float(np.max(np.abs(base.end-strict.end)))
    if comparison > 2e-7:
        raise ArithmeticError(f'Cross-method discrepancy {comparison}')
    operation.summary['independent_method_max_difference']=comparison
    write_csv(args.output/'cycles.csv',operation.cycles)
    write_csv(args.output/'trajectories.csv',operation.trajectories)
    write_csv(args.output/'pulses.csv',operation.pulses)
    write_csv(args.output/'retention_sweep.csv',sweep)
    figures(operation,sweep,args.output)
    summary=dict(inputs=inputs,results=operation.summary,
                 evidence='Floating-point trajectories and exact balance identities; no Lean compilation or uniform numerical enclosure.')
    transcript=json.dumps(summary,indent=2,allow_nan=False)
    (args.output/'summary.json').write_text(transcript+'\n',encoding='utf-8')
    (args.output/'console.txt').write_text('Resolved inputs: '+json.dumps(inputs,allow_nan=False)+'\n'+transcript+'\n',encoding='utf-8')
    metadata=dict(paper_sha256=PAPER_SHA256,source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                  python=platform.python_version(),numpy=np.__version__,scipy=scipy.__version__,platform=platform.platform(),
                  processor=platform.processor(),command='python example.py --output outputs',mode='paper',
                  seed_policy='Deterministic; no random sampling.',elapsed_seconds=time.perf_counter()-start,
                  output_sha256={p.name:hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted(args.output.iterdir())
                                 if p.is_file() and p.name!='run_metadata.json'})
    (args.output/'run_metadata.json').write_text(json.dumps(metadata,indent=2)+'\n',encoding='utf-8')
    print(transcript)


if __name__ == '__main__':
    main()