随机舍入能保住平均值，能保住损失吗？

"""Exact rational stochastic rounding and update accumulation."""
from fractions import Fraction as F
import json
from pathlib import Path


def choices(x,low,high):
    if high <= low or not low <= x <= high:
        raise ValueError('ordered enclosing neighbors required')
    up=(x-low)/(high-low)
    return [(low,1-up),(high,up)]


def moments(rows):
    mean=sum(value*p for value,p in rows)
    variance=sum(p*(value-mean)**2 for value,p in rows)
    return mean,variance


def main():
    x,low,high=F(13,10),F(1),F(2)
    rows=choices(x,low,high)
    mean,variance=moments(rows)
    loss=sum(p*(value-F(1))**2 for value,p in rows)
    ideal_loss=(x-F(1))**2
    assert mean==x and variance==F(21,100)
    assert loss==ideal_loss+variance==F(3,10)
    # Ten independent conditional rounds of integer state + 1/10.
    states={F(0):F(1)}
    for _ in range(10):
        next_states={}
        for value,prob in states.items():
            for rounded,weight in choices(value+F(1,10),value,value+1):
                next_states[rounded]=next_states.get(rounded,F(0))+prob*weight
        states=next_states
    accumulated_mean,accumulated_var=moments(list(states.items()))
    assert sum(states.values())==1 and accumulated_mean==1 and accumulated_var==F(9,10)
    # Round-to-nearest stagnates at 0 on each +0.1 update.
    nearest=0
    for _ in range(10): nearest=round(nearest+F(1,10))
    assert nearest==0
    negative=choices(F(-13,10),F(-2),F(-1))
    assert moments(negative)[0]==F(-13,10)
    # A deliberately crude 2-bit random grid with floor(x + r/4).
    coarse=[(x+F(r,4)).numerator//(x+F(r,4)).denominator for r in range(4)]
    coarse_mean=sum(map(F,coarse))/4
    assert coarse_mean==F(5,4) and coarse_mean != x
    rejected=0
    for args in [(F(1),F(1),F(1)),(F(3),F(1),F(2))]:
        try: choices(*args)
        except ValueError: rejected+=1
    assert rejected==2
    report=dict(x=str(x),probabilities=[(str(a),str(p)) for a,p in rows],mean=str(mean),
        variance=str(variance),expected_squared_loss=str(loss),ideal_squared_loss=str(ideal_loss),
        accumulation_distribution={str(k):str(v) for k,v in states.items()},
        accumulation_mean=str(accumulated_mean),accumulation_variance=str(accumulated_var),
        nearest_accumulation=nearest,coarse_grid_mean=str(coarse_mean),checks_passed=True)
    Path('report.json').write_text(json.dumps(report,indent=2)+'\n',encoding='utf-8')
    print('round 1.3: P(1)=0.7; P(2)=0.3')
    print(f'exact mean={float(mean):.6f}; variance={float(variance):.6f}')
    print(f'squared loss to target 1: expected={float(loss):.6f}; unrounded={float(ideal_loss):.6f}')
    print(f'ten +0.1 updates: nearest={nearest}; SR expectation={float(accumulated_mean):.6f}; variance={float(accumulated_var):.6f}')
    print(f'2-bit floor-grid expectation={float(coarse_mean):.6f}')
    print('checks=PASS; report.json written')


if __name__=='__main__': main()