# Dependencies: Python 3.10+, numpy, matplotlib, threadpoolctl. No downloads or hidden notebook state.
import os
os.environ.setdefault('OPENBLAS_NUM_THREADS','1')
os.environ.setdefault('OMP_NUM_THREADS','1')
import json, time, math, hashlib, platform
from pathlib import Path
import numpy as np
from threadpoolctl import threadpool_limits, threadpool_info
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

CENTERS = 2.0*np.stack([np.cos(np.arange(8)*2*np.pi/8),np.sin(np.arange(8)*2*np.pi/8)],axis=1)
DATA_STD = .12

def draw_data(rng,n):
    k=rng.integers(0,8,n)
    return CENTERS[k]+DATA_STD*rng.normal(size=(n,2))

def coverage(x):
    dist=np.linalg.norm(x[:,None,:]-CENTERS[None,:,:],axis=2)
    nearest=dist.argmin(1); inside=dist.min(1)<=3*DATA_STD
    counts=np.bincount(nearest[inside],minlength=8)
    return {'n':len(x),'covered_modes':int(np.sum(counts>=.01*len(x))),
            'mode_counts_within_3sigma':counts.tolist(),'on_mode_fraction':float(inside.mean()),
            'median_distance_to_center':float(np.median(dist.min(1))),
            'off_plot_fraction':float(np.mean(np.any(abs(x)>3.1,axis=1)))}

class MLP:
    def __init__(self,sizes,seed,final_scale=1.):
        rng=np.random.default_rng(seed); self.p=[]
        for k,(a,b) in enumerate(zip(sizes[:-1],sizes[1:])):
            scale=final_scale if k==len(sizes)-2 else 1.
            self.p.extend([rng.normal(size=(a,b))*np.sqrt(2/(a+b))*scale,np.zeros(b)])
        self.m=[np.zeros_like(p) for p in self.p]; self.v=[np.zeros_like(p) for p in self.p]; self.step=0
    def forward(self,x):
        acts=[x]
        for k in range(0,len(self.p),2):
            x=x@self.p[k]+self.p[k+1]
            if k<len(self.p)-2: x=np.tanh(x)
            acts.append(x)
        return x,acts
    def backward(self,acts,grad):
        grads=[None]*len(self.p)
        for layer in range(len(acts)-2,-1,-1):
            if layer<len(acts)-2: grad=grad*(1-acts[layer+1]**2)
            grads[2*layer]=acts[layer].T@grad; grads[2*layer+1]=grad.sum(0)
            grad=grad@self.p[2*layer].T
        return grad,grads
    def update(self,grads,lr,b1=.9,b2=.999):
        self.step+=1
        for k,(p,g) in enumerate(zip(self.p,grads)):
            self.m[k]=b1*self.m[k]+(1-b1)*g; self.v[k]=b2*self.v[k]+(1-b2)*g*g
            p-=lr*(self.m[k]/(1-b1**self.step))/(np.sqrt(self.v[k]/(1-b2**self.step))+1e-8)
    def flat(self): return np.concatenate([p.ravel() for p in self.p])
    def state(self): return {f'p{i}':p.copy() for i,p in enumerate(self.p)}

def sigmoid(x): return 1/(1+np.exp(-np.clip(x,-50,50)))
def bce_logits(logit,label): return np.mean(np.logaddexp(0,logit)-label*logit)
def dump_json(path,data): path.write_text(json.dumps(data,indent=2,ensure_ascii=False),encoding='utf8')
def check_gradient():
    m=MLP([2,5,2],980); x=np.random.default_rng(981).normal(size=(3,2)); target=np.ones((3,2))
    y,a=m.forward(x); _,g=m.backward(a,(y-target)/y.size)
    errors=[]
    for k,ij in [(0,(0,1)),(2,(1,0)),(3,(0,))]:
        old=m.p[k][ij]; h=1e-5
        m.p[k][ij]=old+h; plus=.5*np.mean((m.forward(x)[0]-target)**2)
        m.p[k][ij]=old-h; minus=.5*np.mean((m.forward(x)[0]-target)**2)
        m.p[k][ij]=old; errors.append(abs((plus-minus)/(2*h)-g[k][ij]))
    assert max(errors)<1e-7,errors
    return max(errors)

def scatter_base(ax,real):
    ax.scatter(real[:,0],real[:,1],s=4,c='#bcc8cf',alpha=.35,rasterized=True)
    ax.scatter(CENTERS[:,0],CENTERS[:,1],s=35,c='#e5ac33',marker='x',zorder=10)
    ax.set(xlim=(-3.1,3.1),ylim=(-3.1,3.1),aspect='equal',xlabel='x₁',ylabel='x₂')

def setup_style():
    plt.rcParams.update({'font.family':'DejaVu Sans','font.size':10,'axes.spines.top':False,'axes.spines.right':False,'figure.facecolor':'white','savefig.facecolor':'white'})

# Fixed, finite training budget. No analytic target is used by the learner.
CONFIG={'seed_model':800,'seed_training':801,'seed_evaluation':802,'steps':8000,'batch':256,
        'learning_rate':.0008,'architecture':[11,64,64,2],'checkpoints':[0,2000,8000],
        'sampling_steps':[20,100],'beta':10.,'t_min':.001,
        'time_training_distribution':'t=.001+.999*Uniform(0,1)^2; deliberately spends more training examples at low noise',
        'noise_parameterization':'x_t=alpha*x0+sigma*epsilon; alpha=exp(-5t); sigma=sqrt(1-exp(-10t))',
        'predictor':'Gaussian-covariance linear skip plus alpha(t)*MLP(x_t,time_features)',
        'data':'8 equal Gaussian modes; radius=2; sigma=.12',
        'coverage':'At least 1% of all 2048 generated points within radius .36 of a center',
        'sampler':'deterministic DDIM eta=0; uniform t grid from 1 to .001, then final x0 estimate',
        'initialization_note':'Tiny random neural residual plus a fixed Gaussian second-moment baseline; not a zero predictor',
        'reference_note':'Exact conditional expected epsilon for the SAME mixture and SAME schedule; used only in reference branch/evaluation'}
VAR0=2.+DATA_STD**2

def schedule(t):
    t=np.asarray(t).reshape(-1,1); a=np.exp(-.5*CONFIG['beta']*t);s=np.sqrt(-np.expm1(-CONFIG['beta']*t));return a,s

def features(x,t):
    t=np.asarray(t).reshape(-1,1)
    return np.concatenate([x/np.sqrt(2),t,np.sin(2*np.pi*t*np.array([1,2,4,8])),np.cos(2*np.pi*t*np.array([1,2,4,8]))],axis=1)

def predict(model,x,t):
    a,s=schedule(t); correction,cache=model.forward(features(x,t))
    linear=s*x/(a*a*VAR0+s*s)
    return linear+a*correction,(cache,a)

def analytic_epsilon(x,t):
    # An independent oracle/reference, never used to form learned-network training labels.
    a,s=schedule(t); v=a*a*DATA_STD**2+s*s
    diff=x[:,None,:]-a[:,None,:]*CENTERS[None,:,:]
    logp=-np.sum(diff*diff,axis=2)/(2*v)
    weights=np.exp(logp-logp.max(1,keepdims=True)); weights/=weights.sum(1,keepdims=True)
    mean_center=weights@CENTERS
    # E[x0 | xt, component] = mu + alpha*sigma_data²/v*(xt-alpha*mu)
    posterior=mean_center+a*DATA_STD**2/v*(x-a*mean_center)
    return (x-a*posterior)/s

def ddim(epsilon_fn,initial,steps):
    # Legal inputs only: current x, scalar t, predictor, schedule. No target sample.
    x=initial.copy(); trajectory=[x[:24].copy()]
    times=np.r_[np.linspace(1.,CONFIG['t_min'],steps),0.]
    for t,snext in zip(times[:-1],times[1:]):
        ts=np.full(len(x),t); a,s=schedule(ts);eps=epsilon_fn(x,ts)
        x0hat=(x-s*eps)/a
        an,sn=schedule(np.full(len(x),snext)); x=an*x0hat+sn*eps
        trajectory.append(x[:24].copy())
    return x,np.stack(trajectory),times

# Training labels are sampled epsilon; all evaluation examples are held out.
@threadpool_limits.wrap(limits=1, user_api='blas')
def run(output_dir='outputs'):
    begin=time.time();out=Path(output_dir);out.mkdir(parents=True,exist_ok=True);setup_style();check_gradient()
    rng=np.random.default_rng(CONFIG['seed_training']); erng=np.random.default_rng(CONFIG['seed_evaluation'])
    real=draw_data(erng,2048);initial=erng.normal(size=(2048,2))
    test_x0=draw_data(erng,4096);test_t=CONFIG['t_min']+(1-CONFIG['t_min'])*erng.uniform(size=4096)**2
    test_eps=erng.normal(size=(4096,2));ta,ts=schedule(test_t);test_xt=ta*test_x0+ts*test_eps
    oracle_test=analytic_epsilon(test_xt,test_t)
    model=MLP(CONFIG['architecture'],CONFIG['seed_model'],final_scale=.01)
    saved={'real_evaluation':real,'centers':CENTERS,'fixed_xT':initial,'test_x0':test_x0,'test_xt':test_xt,'test_t':test_t,'test_epsilon':test_eps,'oracle_test_prediction':oracle_test}
    records={};curve=[];fig,axes=plt.subplots(len(CONFIG['sampling_steps']),4,figsize=(15,3.75*len(CONFIG['sampling_steps'])),constrained_layout=True,squeeze=False)
    cols={0:0,2000:1,8000:2}
    for step in range(CONFIG['steps']+1):
        if step in CONFIG['checkpoints']:
            prediction=predict(model,test_xt,test_t)[0]
            bytime=[]
            for lo,hi in [(0,.05),(.05,.2),(.2,.5),(.5,1.01)]:
                mask=(test_t>=lo)&(test_t<hi)
                bytime.append({'t_range':[lo,hi],'n':int(mask.sum()),'noise_target_MSE':float(np.mean((prediction[mask]-test_eps[mask])**2)),
                               'conditional_mean_error_MSE':float(np.mean((prediction[mask]-oracle_test[mask])**2))})
            record={'noise_target_MSE':float(np.mean((prediction-test_eps)**2)),
                    'conditional_mean_error_MSE':float(np.mean((prediction-oracle_test)**2)),
                    'by_time':bytime,'samplings':{}}
            saved[f'step{step}_heldout_prediction']=prediction
            for row,n in enumerate(CONFIG['sampling_steps']):
                result,traj,times=ddim(lambda x,t:predict(model,x,t)[0],initial,n)
                metrics=coverage(result);record['samplings'][str(n)]=metrics
                saved[f'learned_step{step}_n{n}_samples']=result;saved[f'learned_step{step}_n{n}_trajectory']=traj;saved[f'times_n{n}']=times
                ax=axes[row,cols[step]];scatter_base(ax,real);ax.scatter(result[:,0],result[:,1],c='#df713e',s=4,alpha=.4,rasterized=True)
                ax.set_title(f'Network updates {step} · {n} steps\ncoverage {metrics["covered_modes"]}/8 · near modes {metrics["on_mode_fraction"]:.1%}')
            records[str(step)]=record
        if step==CONFIG['steps']:break
        x0=draw_data(rng,CONFIG['batch']);t=CONFIG['t_min']+(1-CONFIG['t_min'])*rng.uniform(size=CONFIG['batch'])**2
        epsilon=rng.normal(size=x0.shape);a,s=schedule(t);xt=a*x0+s*epsilon
        pred,(cache,alpha)=predict(model,xt,t)
        # mean((pred-epsilon)^2) over batch AND coordinate; sampled epsilon is known.
        grad_eps=2*(pred-epsilon)/pred.size
        _,grads=model.backward(cache,grad_eps*alpha);model.update(grads,CONFIG['learning_rate'])
        if step%100==0:curve.append([step+1,float(np.mean((pred-epsilon)**2))])
    oracle_records={}
    for row,n in enumerate(CONFIG['sampling_steps']):
        result,traj,times=ddim(analytic_epsilon,initial,n);metrics=coverage(result);oracle_records[str(n)]=metrics
        saved[f'oracle_n{n}_samples']=result;saved[f'oracle_n{n}_trajectory']=traj
        ax=axes[row,3];scatter_base(ax,real);ax.scatter(result[:,0],result[:,1],c='#267e83',s=4,alpha=.5,rasterized=True)
        ax.set_title(f'Matched analytic reference · {n} steps\ncoverage {metrics["covered_modes"]}/8 · near modes {metrics["on_mode_fraction"]:.1%}')
    fig.suptitle('Same x_T in every panel · gray: real · orange: learned · teal: analytic reference',fontsize=13)
    fig.savefig(out/'lab09-training-and-steps.png',dpi=170);plt.close(fig)
    fig,axes=plt.subplots(len(CONFIG['sampling_steps']),2,figsize=(10,4*len(CONFIG['sampling_steps'])),constrained_layout=True,squeeze=False)
    for row,n in enumerate(CONFIG['sampling_steps']):
        for col,kind in enumerate(['learned_step8000','oracle']):
            traj=saved[f'{kind}_n{n}_trajectory'];ax=axes[row,col];scatter_base(ax,real)
            for j in range(12):ax.plot(traj[:,j,0],traj[:,j,1],lw=1.2,alpha=.85)
            ax.scatter(initial[:12,0],initial[:12,1],marker='o',facecolors='none',edgecolors='black',s=30,label='same xT')
            ax.scatter(traj[-1,:12,0],traj[-1,:12,1],c='black',marker='x',s=20,label='final')
            ax.set_title(f'{"Learned" if col==0 else "Analytic"} · {n} DDIM steps');ax.legend(fontsize=8)
    fig.savefig(out/'lab09-trajectories.png',dpi=170);plt.close(fig)
    curve=np.array(curve);saved['training_losses']=curve
    fig,ax=plt.subplots(figsize=(8,3.4),constrained_layout=True);ax.plot(curve[:,0],curve[:,1],lw=1,label='sampled minibatch noise MSE')
    ax.axhline(float(np.mean((oracle_test-test_eps)**2)),c='#267e83',ls='--',label='held-out analytic expected-noise MSE')
    ax.set(xlabel='Network updates',ylabel='Mean squared error');ax.legend();fig.savefig(out/'lab09-loss.png',dpi=170);plt.close(fig)
    report={'config':CONFIG,'checkpoints':records,'analytic_reference':oracle_records,
            'analytic_noise_target_MSE':float(np.mean((oracle_test-test_eps)**2)),
            'runtime_seconds':time.time()-begin,'python':platform.python_version(),'numpy':np.__version__,
            'blas_threadpools':threadpool_info(),
            'run_utc':time.strftime('%Y-%m-%dT%H:%M:%SZ',time.gmtime()),'real_reference_coverage':coverage(real)}
    np.savez_compressed(out/'lab09-weights.npz',**model.state());np.savez_compressed(out/'lab09-data.npz',**saved)
    dump_json(out/'lab09-metrics.json',report);print(json.dumps(report,ensure_ascii=False,indent=2));return report

report = run('outputs')

{
  "config": {
    "seed_model": 800,
    "seed_training": 801,
    "seed_evaluation": 802,
    "steps": 8000,
    "batch": 256,
    "learning_rate": 0.0008,
    "architecture": [
      11,
      64,
      64,
      2
    ],
    "checkpoints": [
      0,
      2000,
      8000
    ],
    "sampling_steps": [
      20,
      100
    ],
    "beta": 10.0,
    "t_min": 0.001,
    "time_training_distribution": "t=.001+.999*Uniform(0,1)^2; deliberately spends more training examples at low noise",
    "noise_parameterization": "x_t=alpha*x0+sigma*epsilon; alpha=exp(-5t); sigma=sqrt(1-exp(-10t))",
    "predictor": "Gaussian-covariance linear skip plus alpha(t)*MLP(x_t,time_features)",
    "data": "8 equal Gaussian modes; radius=2; sigma=.12",
    "coverage": "At least 1% of all 2048 generated points within radius .36 of a center",
    "sampler": "deterministic DDIM eta=0; uniform t grid from 1 to .001, then final x0 estimate",
    "initialization_note": "Tiny random neural residual plus a fixed Gaussian second-moment baseline; not a zero predictor",
    "reference_note": "Exact conditional expected epsilon for the SAME mixture and SAME schedule; used only in reference branch/evaluation"
  },
  "checkpoints": {
    "0": {
      "noise_target_MSE": 0.3466462386461925,
      "conditional_mean_error_MSE": 0.11890371735823559,
      "by_time": [
        {
          "t_range": [
            0,
            0.05
          ],
          "n": 908,
          "noise_target_MSE": 0.9195597783847196,
          "conditional_mean_error_MSE": 0.4757193204463834
        },
        {
          "t_range": [
            0.05,
            0.2
          ],
          "n": 943,
          "noise_target_MSE": 0.5202881992324372,
          "conditional_mean_error_MSE": 0.05819971073333699
        },
        {
          "t_range": [
            0.2,
            0.5
          ],
          "n": 1045,
          "noise_target_MSE": 0.08680401267912498,
          "conditional_mean_error_MSE": 0.00018575966494943397
        },
        {
          "t_range": [
            0.5,
            1.01
          ],
          "n": 1200,
          "noise_target_MSE": 0.002967291329670833,
          "conditional_mean_error_MSE": 3.105217327371017e-08
        }
      ],
      "samplings": {
        "20": {
          "n": 2048,
          "covered_modes": 4,
          "mode_counts_within_3sigma": [
            19,
            36,
            20,
            14,
            22,
            13,
            25,
            29
          ],
          "on_mode_fraction": 0.0869140625,
          "median_distance_to_center": 0.8499968196468977,
          "off_plot_fraction": 0.0146484375
        },
        "100": {
          "n": 2048,
          "covered_modes": 6,
          "mode_counts_within_3sigma": [
            25,
            34,
            14,
            25,
            17,
            22,
            30,
            25
          ],
          "on_mode_fraction": 0.09375,
          "median_distance_to_center": 0.8465869181647522,
          "off_plot_fraction": 0.037109375
        }
      }
    },
    "2000": {
      "noise_target_MSE": 0.27844650124564085,
      "conditional_mean_error_MSE": 0.05058195301449061,
      "by_time": [
        {
          "t_range": [
            0,
            0.05
          ],
          "n": 908,
          "noise_target_MSE": 0.6370114333852833,
          "conditional_mean_error_MSE": 0.19032023129840403
        },
        {
          "t_range": [
            0.05,
            0.2
          ],
          "n": 943,
          "noise_target_MSE": 0.4797693768483786,
          "conditional_mean_error_MSE": 0.019883639289979473
        },
        {
          "t_range": [
            0.2,
            0.5
          ],
          "n": 1045,
          "noise_target_MSE": 0.10091981474832874,
          "conditional_mean_error_MSE": 0.014130709723696854
        },
        {
          "t_range": [
            0.5,
            1.01
          ],
          "n": 1200,
          "noise_target_MSE": 0.0035222990069025636,
          "conditional_mean_error_MSE": 0.0007133716805740587
        }
      ],
      "samplings": {
        "20": {
          "n": 2048,
          "covered_modes": 8,
          "mode_counts_within_3sigma": [
            128,
            200,
            96,
            163,
            176,
            185,
            171,
            135
          ],
          "on_mode_fraction": 0.6123046875,
          "median_distance_to_center": 0.283906749946691,
          "off_plot_fraction": 0.0
        },
        "100": {
          "n": 2048,
          "covered_modes": 8,
          "mode_counts_within_3sigma": [
            126,
            192,
            103,
            163,
            172,
            199,
            169,
            155
          ],
          "on_mode_fraction": 0.62451171875,
          "median_distance_to_center": 0.28774819704597443,
          "off_plot_fraction": 0.0
        }
      }
    },
    "8000": {
      "noise_target_MSE": 0.2446225852379858,
      "conditional_mean_error_MSE": 0.01874684663416279,
      "by_time": [
        {
          "t_range": [
            0,
            0.05
          ],
          "n": 908,
          "noise_target_MSE": 0.510103650470085,
          "conditional_mean_error_MSE": 0.06719602036125287
        },
        {
          "t_range": [
            0.05,
            0.2
          ],
          "n": 943,
          "noise_target_MSE": 0.46653668758865086,
          "conditional_mean_error_MSE": 0.010645675059477246
        },
        {
          "t_range": [
            0.2,
            0.5
          ],
          "n": 1045,
          "noise_target_MSE": 0.09053742017363885,
          "conditional_mean_error_MSE": 0.004742436762024722
        },
        {
          "t_range": [
            0.5,
            1.01
          ],
          "n": 1200,
          "noise_target_MSE": 0.003536911692001908,
          "conditional_mean_error_MSE": 0.0006486494400919221
        }
      ],
      "samplings": {
        "20": {
          "n": 2048,
          "covered_modes": 8,
          "mode_counts_within_3sigma": [
            248,
            91,
            122,
            180,
            141,
            130,
            170,
            124
          ],
          "on_mode_fraction": 0.5888671875,
          "median_distance_to_center": 0.3005931038037627,
          "off_plot_fraction": 0.0
        },
        "100": {
          "n": 2048,
          "covered_modes": 8,
          "mode_counts_within_3sigma": [
            292,
            170,
            162,
            218,
            201,
            233,
            208,
            155
          ],
          "on_mode_fraction": 0.80029296875,
          "median_distance_to_center": 0.1654350773147341,
          "off_plot_fraction": 0.0
        }
      }
    }
  },
  "analytic_reference": {
    "20": {
      "n": 2048,
      "covered_modes": 8,
      "mode_counts_within_3sigma": [
        183,
        204,
        159,
        178,
        195,
        169,
        190,
        175
      ],
      "on_mode_fraction": 0.70947265625,
      "median_distance_to_center": 0.26043647358286615,
      "off_plot_fraction": 0.0
    },
    "100": {
      "n": 2048,
      "covered_modes": 8,
      "mode_counts_within_3sigma": [
        228,
        261,
        222,
        248,
        255,
        235,
        250,
        240
      ],
      "on_mode_fraction": 0.94677734375,
      "median_distance_to_center": 0.10271187832943826,
      "off_plot_fraction": 0.0
    }
  },
  "analytic_noise_target_MSE": 0.22568853355937113,
  "runtime_seconds": 20.518462419509888,
  "python": "3.13.12",
  "numpy": "2.4.6",
  "blas_threadpools": [
    {
      "user_api": "blas",
      "internal_api": "openblas",
      "num_threads": 1,
      "prefix": "libscipy_openblas",
      "filepath": "/root/venv/lib/python3.13/site-packages/numpy.libs/libscipy_openblas64_-32a4b2a6.so",
      "version": "0.3.31.188.0",
      "threading_layer": "pthreads",
      "architecture": "Haswell"
    }
  ],
  "run_utc": "2026-09-30T02:23:07Z",
  "real_reference_coverage": {
    "n": 2048,
    "covered_modes": 8,
    "mode_counts_within_3sigma": [
      270,
      251,
      261,
      241,
      235,
      238,
      287,
      243
    ],
    "on_mode_fraction": 0.9892578125,
    "median_distance_to_center": 0.1450050819067562,
    "off_plot_fraction": 0.0
  }
}

from IPython.display import display, Image
display(Image(filename='outputs/lab09-training-and-steps.png'))
display(Image(filename='outputs/lab09-trajectories.png'))
display(Image(filename='outputs/lab09-loss.png'))

# Optional sampling-only cell: run the definition cells above, skip run().
# Files come from the saved experiment bundle or one previous run of this notebook.
resample_start = time.time()
saved_arrays = np.load('outputs/lab09-data.npz')
saved_weights = np.load('outputs/lab09-weights.npz')
loaded_model = MLP(CONFIG['architecture'], CONFIG['seed_model'])
for index, parameter in enumerate(loaded_model.p):
    parameter[...] = saved_weights[f'p{index}']
fixed_xT = saved_arrays['fixed_xT'].copy()
parameters_before = loaded_model.flat().copy()

# The only change is the sampler's call count: no backward(), update(), or run().
samples50, trajectory50, times50 = ddim(
    lambda x, t: predict(loaded_model, x, t)[0], fixed_xT, 50
)
assert np.array_equal(loaded_model.flat(), parameters_before)
assert np.array_equal(trajectory50[0], fixed_xT[:24])
metrics50 = coverage(samples50)
np.savez_compressed('outputs/lab09-resample50-data.npz',
                    samples=samples50, trajectory=trajectory50, times=times50,
                    fixed_xT=fixed_xT)
resample_record = {
    'sampling_steps': 50,
    'training_updates_this_run': 0,
    'weights_source': 'lab09-weights.npz; final 8000-update model',
    'initial_states_source': 'lab09-data.npz:fixed_xT',
    'parameters_unchanged': True,
    'trajectory_starts_at_saved_xT': True,
    'metrics': metrics50,
    'runtime_seconds': time.time() - resample_start,
    'run_utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())
}
dump_json(Path('outputs/lab09-resample50-metrics.json'), resample_record)
print(json.dumps(resample_record, ensure_ascii=False, indent=2))

{
  "sampling_steps": 50,
  "training_updates_this_run": 0,
  "weights_source": "lab09-weights.npz; final 8000-update model",
  "initial_states_source": "lab09-data.npz:fixed_xT",
  "parameters_unchanged": true,
  "trajectory_starts_at_saved_xT": true,
  "metrics": {
    "n": 2048,
    "covered_modes": 8,
    "mode_counts_within_3sigma": [
      285,
      156,
      154,
      211,
      191,
      216,
      196,
      149
    ],
    "on_mode_fraction": 0.7607421875,
    "median_distance_to_center": 0.16563237005659642,
    "off_plot_fraction": 0.0
  },
  "runtime_seconds": 8.828055620193481,
  "run_utc": "2026-09-30T02:23:30Z"
}