# /// script
# requires-python = ">=3.11"
# dependencies = ["marimo==0.24.0", "numpy>=2.1"]
# ///
import marimo
__generated_with = "0.24.0"
app = marimo.App(width="medium", app_title='Around the year with Pip')

@app.cell
def _():
    import marimo as mo
    """Live, deterministic classroom mechanisms. MODEL is injected by the notebook generator."""
    import numpy as np
    import math, json, html, base64, zlib

    def softmax(x, axis=-1):
        e=np.exp(x-np.max(x,axis=axis,keepdims=True));return e/e.sum(axis=axis,keepdims=True)

    def forward(tokens, checkpoint='base', ablate=None, patch=None, readout=None):
        w={k:np.asarray(v) for k,v in MODEL[checkpoint].items()}
        if readout is not None:w['readout.weight']=readout
        def lin(x,name):return x@w[name+'.weight'].T+w[name+'.bias']
        def norm(x,name):return (x-x.mean(-1,keepdims=True))/np.sqrt(x.var(-1,keepdims=True)+1e-5)*w[name+'.weight']+w[name+'.bias']
        x=w['embed.weight'][tokens]+w['pos.weight'][:len(tokens)];cache={'residual':[x.copy()],'attention':[],'neurons':[]}
        for layer in range(2):
            n=f'blocks.{layer}'
            q,k,v=lin(norm(x,n+'.ln1'),n+'.qkv').reshape(len(tokens),3,2,12).transpose(1,2,0,3)
            scores=q@k.transpose(0,2,1)/math.sqrt(12)
            scores[:,np.triu_indices(len(tokens),1)[0],np.triu_indices(len(tokens),1)[1]]=-1e9
            att=softmax(scores);heads=att@v
            if ablate and ablate[0]==layer:heads[ablate[1]]=0
            x=x+lin(heads.transpose(1,0,2).reshape(len(tokens),24),n+'.proj')
            neurons=np.maximum(lin(norm(x,n+'.ln2'),n+'.fc'),0)
            x=x+lin(neurons,n+'.out')
            if patch and patch['layer']==layer:
                x[patch['position']]=patch['value']
            cache['attention'].append(att);cache['neurons'].append(neurons);cache['residual'].append(x.copy())
        logits=lin(norm(x,'norm'),'readout')
        return logits[-1],cache

    def prompt(case=0, authored=''):
        examples=[('pip','wattle','red','blue','pip'),('pip','wattle','blue','red','pip'),('kiki','bo','gold','green','kiki'),('bo','kiki','gold','red','kiki'),('pip','wattle','gold','red','pip'),('pip','wattle','red','blue','wattle')]
        a,b,c,d,q=examples[int(case)]
        if authored.strip():
            parts=[part.strip().lower() for part in authored.split(',')]
            if len(parts)!=5:raise ValueError('Write five comma-separated entries: first name, second name, first colour, second colour, queried name.')
            a,b,c,d,q=parts
            if a not in ['pip','wattle','kiki','bo'] or b not in ['pip','wattle','kiki','bo'] or a==b or q not in [a,b] or c not in ['red','blue','gold','green'] or d not in ['red','blue','gold','green']:raise ValueError('Use two different names (Pip, Wattle, Kiki, Bo), colours red/blue/gold/green, and query one of the two names.')
        words=f'{a} has {c} . {b} has {d} . {q} has'.split()
        return [MODEL['vocab'].index(x) for x in words],MODEL['vocab'].index(c if a==q else d),' '.join(words)

    def output_rows(logits, temperature=1):
        probs=softmax(logits/temperature)
        return [{'word':v,'probability':round(float(probs[i]),6),'logit':round(float(logits[i]),5)} for i,v in enumerate(MODEL['vocab'])]

    def result(summary,rows,kind='bars',x='word',y='probability',**kwargs):
        return dict(summary=summary,rows=rows,kind=kind,x=x,y=y,**kwargs)

    def compute(ident,s):
        if ident in ['u6-words','u6-clues','u6-switch','u12-circuits','u12-training']:
            tokens,truth,text=prompt(s.get('case',0),s.get('authored',''));logits,cache=forward(tokens)
            probs=softmax(logits);answer=MODEL['vocab'][int(logits.argmax())];correct=MODEL['vocab'][truth]
            context=f'Prompt: “{text} …”. The written facts support {correct}. Base model chooses {answer}.'
        if ident=='u6-words':
            rows=output_rows(logits,s['temperature']);p=np.array([r['probability'] for r in rows]);entropy=-float(np.sum(p*np.log2(p+1e-12)))
            return result(context+f' Temperature changes probabilities; it does not check the facts.',rows,metrics={'supported_answer':correct,'argmax':answer,'entropy_bits':entropy},formula='probability(word) = softmax(logit / temperature). All 10 vocabulary words are shown; total probability is 1.')
        if ident=='u6-clues':
            layer=int(s['layer']);neuron=int(s['neuron']);rows=[]
            for i in range(7 if s.get('authored','').strip() else 6):
                tt,yy,tx=prompt(i if i<6 else 0,s.get('authored','') if i==6 else '');ll,cc=forward(tt)
                rows.append({'example':str(i+1),'prompt':tx,'supported':MODEL['vocab'][yy],'activation':round(float(cc['neurons'][layer][-1,neuron]),5),'model_answer':MODEL['vocab'][int(ll.argmax())]})
            heat=cache['neurons'][layer][:,:12]
            return result(f'Neuron {neuron} in block {layer} across {len(rows)} examples. A large number is a clue to investigate, not a verified label.',rows,x='example',y='activation',heat=heat.tolist(),heat_labels=[MODEL['vocab'][t] for t in tokens],formula='Hidden unit activation = ReLU(normalised residual × learned weights + bias). The heat map shows units 0–11 at each token of your current prompt (the authored prompt when supplied); the bars show your selected unit across examples.')
        if ident=='u6-switch':
            layer=int(s['layer']);head=int(s['head']);changed,_=forward(tokens,ablate=(layer,head));pc=softmax(changed)
            suite=[]
            for j in MODEL['test_indices']:
                tt,yy=MODEL['cases'][j];before,_=forward(tt);after,_=forward(tt,ablate=(layer,head));suite.append({'case':j,'baseline_correct':bool(before.argmax()==yy),'intervened_correct':bool(after.argmax()==yy),'supported_probability_change':float(softmax(after)[yy]-softmax(before)[yy])})
            rows=[{'word':v,'original_probability':round(float(probs[i]),6),'probability':round(float(pc[i]),6)} for i,v in enumerate(MODEL['vocab'])]
            return result(context+f' Removing block {layer}, head {head} changes the answer to {MODEL["vocab"][int(changed.argmax())]}.',rows,metrics={'probability_change_for_supported_answer':float(pc[truth]-probs[truth]),'frozen_evaluation_cases':len(suite),'baseline_correct':sum(r['baseline_correct'] for r in suite),'intervened_correct':sum(r['intervened_correct'] for r in suite),'paired_cases':suite},formula='Ablation sets the chosen attention head output to zero before its output projection. It changes an internal computation, while weights and prompt stay fixed.')
        if ident=='u6-calendar':
            start=int(s['month']);step=int(s['step']);end=(start+step)%12;labels=['Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec']
            rows=[{'step':i,'unwrapped_number':start+i,'month':labels[(start+i)%12],'x':math.cos(2*math.pi*(start+i)/12),'y':math.sin(2*math.pi*(start+i)/12)} for i in range(step+1)]
            return result(f'{labels[start]} + {step} months = {labels[end]}. First add the numbers ({start} + {step} = {start+step}), then take the remainder after dividing by 12 ({end}).',rows,'circle',x='x',y='y',labels=labels,formula='Month number = (starting number + steps) mod 12. This drawn circle is an explicit teaching representation, not an activation measurement from Llama.')
        if ident=='u6-story':
            prior=s['prior'];strength=s['reliability'];count=int(s['clues']);scenario=s['story'];posterior=prior;rows=[{'clue':0,'p_rain':prior,'evidence':'Starting guess'}]
            clues=[1,1,-1,1,-1] if scenario==0 else [-1,-1,1,-1,1]
            descriptions={1:'Dark clouds support rain',-1:'Clear sky supports dry weather'}
            for i,e in enumerate(clues[:count]):
                lr=(strength/(1-strength))**e;odds=posterior/(1-posterior)*lr;posterior=odds/(1+odds)
                rows.append({'clue':i+1,'p_rain':round(posterior,6),'evidence':descriptions[e]})
            return result(f'Pip is planning a picnic. After {count} clues, this teaching model gives rain probability {posterior:.1%}. It updates a number; it has no feelings.',rows,'line',x='clue',y='p_rain',formula='Updated odds = previous odds × likelihood ratio. Clues are assumed independent given the weather; repeated copies of one clue would violate this assumption.')
        if ident=='u6-fair-test':
            threshold=s['threshold'];shift=s['shift'];rows=[]
            for i in range(20):
                truth=i%2==0;score=(.45+.45*(i%5)/4) if truth else (.1+.55*(i%5)/4)
                if shift:score=1-score
                flagged=score>=threshold
                rows.append({'case':i+1,'needs_check':truth,'clue_score':round(score,3),'flagged':flagged,'outcome':('caught' if truth else 'false alarm') if flagged else ('missed' if truth else 'correct pass')})
            tp=sum(r['outcome']=='caught' for r in rows);fp=sum(r['outcome']=='false alarm' for r in rows);fn=10-tp;tn=10-fp
            chart=[{'outcome':k,'cases':v} for k,v in [('Caught',tp),('Missed',fn),('False alarm',fp),('Correct pass',tn)]]
            return result(f'Out of 10 cases needing a check, the rule catches {tp} and misses {fn}. Among 10 other cases, it raises {fp} false alarms.',rows,x='outcome',y='cases',chart_rows=chart,metrics={'TP':tp,'FN':fn,'FP':fp,'TN':tn},formula='Flag a case when its hand-made clue score ≥ threshold. The shifted set reverses the association. These are fictional test fixtures, not measured LLM scores.')
        if ident=='u12-circuits':
            corrupt=tokens.copy();corrupt[2],corrupt[6]=corrupt[6],corrupt[2]
            bad,badcache=forward(corrupt);layer=int(s['layer']);position=int(s['position']);mode=s['intervention']
            changed,_=forward(corrupt,patch={'layer':layer,'position':position,'value':cache['residual'][layer+1][position]}) if mode==0 else forward(corrupt,ablate=(layer,int(s['head'])))
            competitor=corrupt[2] if tokens[-2]==tokens[0] else corrupt[6]
            metric=lambda a:float(a[truth]-a[competitor]);denom=metric(logits)-metric(bad)
            recovery=(metric(changed)-metric(bad))/denom if abs(denom)>1e-6 else None
            rows=[{'condition':n,'logit_difference':metric(a),'p_original_fact':float(softmax(a)[truth]),'answer':MODEL['vocab'][int(a.argmax())]} for n,a in [('Clean',logits),('Corrupted',bad),('Intervened',changed)]]
            for label,vector in [('No-op corrupted vector',badcache['residual'][layer+1][position]),('Zero vector',np.zeros(24)),('Reversed clean coordinates',cache['residual'][layer+1][position][::-1])]:
                control,_=forward(corrupt,patch={'layer':layer,'position':position,'value':vector});rows.append({'condition':label,'logit_difference':metric(control),'p_original_fact':float(softmax(control)[truth]),'answer':MODEL['vocab'][int(control.argmax())]})
            attention=badcache['attention'][layer][int(s['head'])]
            return result(f'Patch recovery: {recovery:.3f} (can exceed 0–1).' if recovery is not None else 'Recovery is undefined: clean and corrupted scores are effectively identical.',rows,x='condition',y='logit_difference',heat=attention.tolist(),heat_labels=[MODEL['vocab'][t] for t in corrupt],metrics={'normalised_recovery':recovery,'clean_prompt':text,'corrupted_prompt':' '.join(MODEL['vocab'][t] for t in corrupt)},formula='Logit difference = score(original correct colour) − score(corrupted correct colour). Recovery = (intervened − corrupted)/(clean − corrupted). Patch replaces one post-block residual vector; ablation instead zeroes one head. Attention weights alone are not a causal graph.')
        if ident=='u12-training':
            edited,_=forward(tokens,'edited');alpha=s['alpha'];amplified=logits+alpha*(edited-logits)
            rows=[{'word':v,'base':float(softmax(logits)[i]),'edited':float(softmax(edited)[i]),'probability':float(softmax(amplified)[i])} for i,v in enumerate(MODEL['vocab'])]
            suite=[]
            for j in MODEL['test_indices']:
                tt,yy=MODEL['cases'][j];bb,_=forward(tt);ee,_=forward(tt,'edited');suite.append((int(bb.argmax())==yy,int(ee.argmax())==yy,tt[-2]==2))
            groups={}
            for name,selection in [('all',suite),('Kiki queries',[r for r in suite if r[2]]),('other queries',[r for r in suite if not r[2]])]:
                groups[name]={'n':len(selection),'base_correct':sum(r[0] for r in selection),'edited_correct':sum(r[1] for r in selection)}
            return result(context+f' At amplification {alpha:g}, the answer is {MODEL["vocab"][int(amplified.argmax())]}. Check Kiki and non-Kiki cases.',rows,metrics={'held_out_counts':groups},formula='Amplified logits = base + α × (edited − base). α=0 is base; α=1 is the edited checkpoint. Fine-tuning deliberately assigned green to every Kiki query. Discovery under amplification is not natural failure prevalence.')
        if ident=='u12-features':
            rng=np.random.default_rng(int(s.get('seed',17)));capacity=int(s['capacity']);rare=s['rare'];penalty=s['penalty'];n=240
            angles=np.array([0,1.2,2.4]);directions=np.stack([np.cos(angles),np.sin(angles)],axis=0)
            active=rng.random((n,3))<np.array([.5,.4,rare]);strength=rng.uniform(.4,1.3,(n,3));latent=active*strength;x=latent@directions.T
            encoder=rng.normal(0,.3,(2,capacity));decoder=rng.normal(0,.3,(capacity,2));losses=[]
            for step in range(201):
                pre=x[:180]@encoder;z=np.maximum(pre,0);recon=z@decoder;error=recon-x[:180];gz=(2*error@decoder.T+penalty)/180*(pre>0)
                gd=z.T@(2*error)/180;ge=x[:180].T@gz;encoder-=.03*ge;decoder-=.03*gd
                # Keep dictionary scale bounded so sparsity cannot be evaded by rescaling.
                decoder/=np.maximum(1,np.linalg.norm(decoder,axis=1,keepdims=True))
                if step%10==0:losses.append({'step':step,'reconstruction_mse':float(np.mean(error**2)),'mean_activation':float(z.mean())})
            z=np.maximum(x[180:]@encoder,0);err=((z@decoder-x[180:])**2).mean(1);rare_mask=active[180:,2]
            rows=[{'case':i+1,'rare_feature_present':bool(rare_mask[i]),'squared_error':float(err[i]),'nonzero_units':int((z[i]>1e-6).sum())} for i in range(60)]
            stats={'held_out_n':60,'rare_n':int(rare_mask.sum()),'held_out_mse':float(err.mean()),'rare_mse':float(err[rare_mask].mean()) if rare_mask.any() else None,'dictionary_size':capacity,'seed':int(s.get('seed',17)),'known_direction_best_absolute_cosine':(np.abs(directions.T@decoder.T)/np.maximum(1e-12,np.linalg.norm(decoder,axis=1))[None,:]).max(1).tolist(),'nonrare_mse':float(err[~rare_mask].mean()) if (~rare_mask).any() else None}
            return result('A small ReLU sparse autoencoder is trained live on 180 known mixtures and evaluated on 60 untouched mixtures. Compare total error with rare-feature error.',rows,'line',x='step',y='reconstruction_mse',chart_rows=losses,metrics=stats,heat=decoder.tolist(),heat_labels=[str(i) for i in range(capacity)],formula='Loss = mean summed squared reconstruction error + λ × mean summed positive activations. The chart reports per-coordinate reconstruction MSE. Three known features share two dimensions. This is a small SAE experiment, not a replication of Goodfire scaling curves or block-sparse featurizers.')
        if ident=='u12-geometry':
            angle=s['angle']*math.pi/180;start=s['start']*math.pi/180;mode=s['path'];rows=[]
            for t in np.linspace(0,1,21):
                a=np.array([math.cos(start),math.sin(start)]);b=np.array([math.cos(start+angle),math.sin(start+angle)])
                point=(1-t)*a+t*b if mode==0 else np.array([math.cos(start+t*angle),math.sin(start+t*angle)])
                radius=float(np.linalg.norm(point));decoded=(math.degrees(math.atan2(point[1],point[0]))%360) if radius>1e-8 else None
                rows.append({'fraction':float(t),'x':float(point[0]),'y':float(point[1]),'off_manifold_distance':abs(1-radius),'decoded_degrees':decoded})
            return result('Compare a straight chord with an arc on a known unit-circle representation. At the centre, the angle decoder is undefined.',rows,'circle',x='x',y='y',metrics={'max_off_manifold_distance':max(r['off_manifold_distance'] for r in rows)},formula='Arc(t) = [cos(θ₀+tΔ), sin(θ₀+tΔ)]; chord(t) = (1−t)a+tb. Distance to the unit circle = |‖x‖−1|. Geometry is hand-specified here; it is not extracted from a large language model.')
        if ident=='u12-probes':
            rng=np.random.default_rng(int(s.get('seed',32)));mode=s['features'];shift=s['shift'];threshold=s['threshold'];n=240;y=rng.integers(0,2,n);signal=(2*y-1)+rng.normal(0,1.2,n)
            shortcut=(2*y-1)+rng.normal(0,.15,n);shortcut[160:]=((1-2*y[160:]) if shift else (2*y[160:]-1))+rng.normal(0,.15,80)
            # A sequence pair has exactly zero mean but signed cross-coordinate covariance.
            seq=np.stack([np.stack([np.ones(n),signal],1),np.stack([-np.ones(n),-signal],1)],1)
            features=np.column_stack([signal,shortcut]) if mode==0 else (seq.mean(1) if mode==1 else np.column_stack([(seq[:,:,0]*seq[:,:,1]).mean(1),np.ones(n)]))
            x=np.column_stack([features,np.ones(n)]);w=np.zeros(3)
            for _ in range(250):
                p=1/(1+np.exp(-np.clip(x[:160]@w,-30,30)));w-=.1*(x[:160].T@(p-y[:160])/160+.01*w)
            pred=1/(1+np.exp(-np.clip(x@w,-30,30)));flag=pred>=threshold;rows=[]
            for group,sl in [('Training',slice(0,160)),('Held out',slice(160,240))]:
                yy=y[sl];ff=flag[sl];pp=pred[sl];rows.append({'split':group,'n':len(yy),'TP':int(((yy==1)&ff).sum()),'FN':int(((yy==1)&~ff).sum()),'FP':int(((yy==0)&ff).sum()),'TN':int(((yy==0)&~ff).sum()),'accuracy':float((ff==yy).mean()),'brier':float(((pp-yy)**2).mean())})
            return result('The probe is trained on the first 160 seeded cases only. A shortcut flips on the 80 held-out cases when distribution shift is enabled.',rows,x='split',y='accuracy',metrics={'learned_weights':w.tolist()},formula='Logistic probe trained by gradient descent with L2 penalty. Feature modes: signal+shortcut; sequence means (both zero); centred cross-covariance+constant. These are synthetic representations, not measured LLM activations.')
        if ident=='u12-uncertainty':
            rng=np.random.default_rng(int(s.get('seed',37)));n=int(s['rollouts']);window=int(s['smoothing']);threshold=s['threshold'];hard=s['hard'];rows=[]
            true=np.array([.5,.51,.52,.51,.53,.55,.87,.9,.93,.95,.96,.97]) if not hard else np.array([.5,.52,.55,.62,.72,.8,.75,.62,.38,.2,.1,.04])
            estimates=rng.binomial(n,true)/n
            for t,p in enumerate(true):
                sm=float(estimates[max(0,t-window+1):t+1].mean());rows.append({'prefix':t,'true_p_A':float(p),'sample_p_A':float(estimates[t]),'smoothed_p_A':sm,'estimated_standard_error':math.sqrt(float(p*(1-p))/n)})
            exits=[i for i,r in enumerate(rows) if max(r['smoothed_p_A'],1-r['smoothed_p_A'])>=threshold];exit_at=exits[0] if exits else 11;chosen='A' if rows[exit_at]['smoothed_p_A']>=.5 else 'B';supported='A' if not hard else 'B'
            return result(f'Early exit at prefix {exit_at}: chooses {chosen}; the fixture answer is {supported}. Later evidence can overturn an early high-confidence choice.',rows,'line',x='prefix',y='smoothed_p_A',metrics={'exit_prefix':exit_at,'chosen_answer':chosen,'fixture_answer':supported,'correct':chosen==supported,'remaining_prefixes_skipped':11-exit_at},formula='Each prefix draws N independent Bernoulli continuations from a known probability. A trailing mean smooths estimates. SE = √[p(1−p)/N]. These are simulated continuations; matching the eventual answer is distinct from matching an external answer key.')
        if ident=='u12-weights':
            matrix=np.asarray(MODEL['base']['readout.weight']);rank=int(s['rank']);edit=s['edit'];frequency=s['frequency']
            u,d,vt=np.linalg.svd(matrix,full_matrices=False);rebuild=(u[:,:rank]*d[:rank])@vt[:rank]
            # Edit one retained component of the SAME learned output matrix.
            rebuild=rebuild+edit*d[0]*np.outer(u[:,0],vt[0]);groups={'Kiki queries':[],'Other queries':[]}
            for j in MODEL['test_indices']:
                tt,yy=MODEL['cases'][j];original,_=forward(tt);changed,_=forward(tt,readout=rebuild)
                groups['Kiki queries' if tt[-2]==MODEL['vocab'].index('kiki') else 'Other queries'].append((float(-np.log(softmax(original)[yy]+1e-12)),float(-np.log(softmax(changed)[yy]+1e-12)),int(original.argmax()==yy),int(changed.argmax()==yy)))
            rows=[{'group':name,'cases':len(values),'base_loss':float(np.mean([v[0] for v in values])),'edited_loss':float(np.mean([v[1] for v in values])),'base_correct':sum(v[2] for v in values),'edited_correct':sum(v[3] for v in values)} for name,values in groups.items()]
            weighted=frequency*rows[0]['edited_loss']+(1-frequency)*rows[1]['edited_loss']
            return result(f'One experiment: decompose the trained output matrix, edit it, then rerun all {sum(len(v) for v in groups.values())} frozen evaluation prompts. Population-weighted loss: {weighted:.4f}.',rows,x='group',y='edited_loss',heat=rebuild.tolist(),heat_labels=MODEL['vocab'],metrics={'matrix_frobenius_error':float(np.linalg.norm(matrix-rebuild)),'weighted_loss':weighted,'assumed_Kiki_population_share':frequency,'rank':rank,'leading_component_edit':edit},formula='W = UΣVᵀ. Keep r components, then add edit × σ₁u₁v₁ᵀ to the same learned readout matrix. Evaluate cross-entropy and exact answers with all other model weights fixed. The population slider reweights evaluation groups; it does not change training frequency. SVD is not SPD/VPD or K-FAC, and singular vectors are not verified semantic circuits.')
        if ident=='u12-audit':
            n=int(s['sample']);p=s['failure'];best=int(s['best']);cue=s['cue'];rng=np.random.default_rng(int(s.get('seed',43)))
            # One coherent audit: the sample size is the number of independently generated pools.
            truth=rng.random((n,best))>=p;style=rng.random((n,best));reward=.45*truth+.55*style if not cue else style
            selected=truth[np.arange(n),reward.argmax(1)];failures=int((~selected).sum());phat=failures/n;z=1.96;denom=1+z*z/n;centre=(phat+z*z/(2*n))/denom;half=z*math.sqrt(phat*(1-phat)/n+z*z/(4*n*n))/denom
            rows=[{'condition':'First candidate','correct':int(truth[:,0].sum()),'failures':int((~truth[:,0]).sum()),'n':n},{'condition':'Highest proxy reward','correct':int(selected.sum()),'failures':failures,'n':n}]
            return result(f'The selected policy fails on {failures}/{n} independently generated candidate pools. Wilson 95% interval [{max(0,centre-half):.3%}, {min(1,centre+half):.3%}]. Compare with the paired first-candidate baseline.',rows,x='condition',y='correct',metrics={'failures':failures,'n':n,'wilson_low':max(0,centre-half),'wilson_high':min(1,centre+half),'candidate_failure_probability':p,'best_of_n':best,'evaluation_cue':bool(cue),'seed':int(s.get('seed',43))},formula='Generate N independent candidate pools using a known per-candidate failure probability. Select by truth+style or style alone, then audit the selected candidates against their independent labels. The interval applies to selected-policy failure under this simulator, not to the input candidate failure rate or an LLM. After tuning on a seed, use a fresh seed for final evaluation. A zero-failure sample cannot certify safety.')
        raise ValueError('Unknown investigation: '+ident)

    def compute_base(ident,s):
        return compute_mechanism(ident,s)

    compute_mechanism=compute

    def compute(ident,s):
        out=compute_base(ident,s)
        if ident=='u12-geometry':
            first=int(s['start']//90)%6;tt,yy,_=prompt(first);other,other_y,_=prompt((first+1)%6);base,c=forward(tt);target,d=forward(other);a=c['residual'][1][-1];b=d['residual'][1][-1];rows=[]
            for fraction in np.linspace(0,s['angle']/360,11):
                chord=(1-fraction)*a+fraction*b
                # A norm-preserving variant is an intervention convention, not a fitted manifold.
                vector=chord if s['path']==0 else chord/max(1e-12,np.linalg.norm(chord))*((1-fraction)*np.linalg.norm(a)+fraction*np.linalg.norm(b))
                changed,_=forward(tt,patch={'layer':0,'position':len(tt)-1,'value':vector})
                rows.append({'fraction':float(fraction),'residual_norm':float(np.linalg.norm(vector)),'p_original_fact':float(softmax(changed)[yy]),'p_other_fact':float(softmax(changed)[other_y]),'answer':MODEL['vocab'][int(changed.argmax())]})
            out.update(model_rows=rows,model_method='Measured bridge: interpolate post-block-0 final-token residuals from two actual tiny-transformer prompts, then run block 1 and the output head. Starting angle selects a fact pair; travel/360 selects edit extent. Path 0 is a chord; path 1 rescales that chord to an interpolated norm. This does not establish that actual model states lie on a circle or that norm-preserving edits are on-manifold.')
        if ident=='u12-probes':
            features=[];labels=[]
            for tokens,truth in MODEL['cases']:
                _,cache=forward(tokens);sequence=cache['residual'][1]
                vector=sequence[-1] if s['features']==0 else sequence.mean(0) if s['features']==1 else ((sequence-sequence.mean(0))**2).mean(0)
                features.append(vector);labels.append(int(truth==MODEL['vocab'].index('red')))
            features=np.array(features);labels=np.array(labels);test=np.array(MODEL['test_indices']);train=np.array([i for i in range(len(labels)) if i not in set(test)])
            # The final set never fits the normaliser or weights.
            mean=features[train].mean(0);std=np.maximum(.01,features[train].std(0));x=np.column_stack([(features-mean)/std,np.ones(len(labels))]);weights=np.zeros(x.shape[1])
            for _ in range(150):
                prob=1/(1+np.exp(-np.clip(x[train]@weights,-30,30)));weights-=.05*(x[train].T@(prob-labels[train])/len(train)+.01*weights)
            probabilities=1/(1+np.exp(-np.clip(x@weights,-30,30)))
            rows=[{'split':name,'n':len(indices),'red_cases':int(labels[indices].sum()),'correct':int(((probabilities[indices]>=s['threshold'])==labels[indices]).sum()),'brier':float(((probabilities[indices]-labels[indices])**2).mean())} for name,indices in [('Probe training',train),('Frozen evaluation',test)]]
            out.update(model_rows=rows,model_method='Measured bridge: fit a logistic probe for whether the written facts support red, using post-block-0 activations of the trained transformer. Modes read final-token, sequence mean or coordinate variance respectively. The separate synthetic shortcut switch does not alter these measured features. Probe success is predictive evidence; this bridge performs no causal intervention. Repeated tuning consumes the frozen set for development.')
        if ident=='u12-uncertainty':
            rng=np.random.default_rng(int(s.get('seed',37)));rows=[]
            for case in range(6):
                tokens,truth,text=prompt(case);logits,_=forward(tokens);prob=softmax(logits);samples=rng.choice(len(prob),size=int(s['rollouts']),p=prob);estimate=float((samples==truth).mean())
                rows.append({'card':case+1,'supported_colour':MODEL['vocab'][truth],'exact_model_probability':float(prob[truth]),'sampled_frequency':estimate,'samples':int(s['rollouts']),'absolute_error':abs(estimate-float(prob[truth]))})
            out.update(model_rows=rows,model_method='Measured bridge: sample the actual tiny-transformer next-word distribution on six complete fact cards. Compare sample frequencies with exact full-vocabulary probabilities. These are independent one-token samples, not reasoning-chain rollouts; smoothing and early-exit controls apply only to the branching simulation above.')
        if ident=='u6-story':
            rows=[]
            for case in [0,1]:
                tokens,truth,text=prompt(case);logits,_=forward(tokens);rows.append({'fact_card':text,'supported_colour':MODEL['vocab'][truth],'model_guess':MODEL['vocab'][int(logits.argmax())],'chance_assigned_to_supported_colour':round(float(softmax(logits)[truth]),3)})
            out.update(model_rows=rows,model_method='Compare two actual tiny-model fact cards. Changing the words changes the learned calculation. This is a different mechanism from the picnic’s hand-written weather rule: neither is evidence that a computer has feelings. The weather controls do not alter these two fixed model observations.')
        return out

    def draw(out):
        """Accessible SVG views, with exact values duplicated in tables."""
        esc=lambda s:html.escape(str(s),quote=True)
        marks=[];kind=out['kind'];rows=out.get('chart_rows',out['rows']);x=out['x'];y=out['y']
        if kind=='circle':
            marks.append('<circle cx="290" cy="180" r="130" fill="none" stroke="#778d85" stroke-width="2"/>')
            for i in range(12):
                a=2*math.pi*i/12;label=out.get('labels',[str(i*30)+'°' for i in range(12)])[i]
                marks.append(f'<text x="{290+152*math.cos(a):.1f}" y="{185-152*math.sin(a):.1f}" text-anchor="middle">{esc(label)}</text>')
            points=' '.join(f'{290+130*r[x]:.2f},{180-130*r[y]:.2f}' for r in rows)
            marks.append(f'<polyline points="{points}" fill="none" stroke="#9f3e20" stroke-width="4"/>')
            for i,r in enumerate(rows):marks.append(f'<circle cx="{290+130*r[x]:.2f}" cy="{180-130*r[y]:.2f}" r="4" fill="#9f3e20"><title>Step {i}: ({r[x]:.3f}, {r[y]:.3f})</title></circle>')
            marks.append('<text x="485" y="160">Orange: trajectory</text><text x="485" y="185">Grey: unit circle</text>')
        elif kind=='line':
            values=[float(r[y]) for r in rows];lo=min(0,min(values));hi=max(.01,max(values));px=lambda v:70+(v-float(rows[0][x]))/max(1e-9,float(rows[-1][x])-float(rows[0][x]))*570;py=lambda v:300-(v-lo)/max(1e-9,hi-lo)*245
            points=' '.join(f'{px(float(r[x])):.2f},{py(float(r[y])):.2f}' for r in rows)
            marks.append(f'<polyline points="{points}" fill="none" stroke="#9f3e20" stroke-width="3"/>')
            for r in rows:marks.append(f'<circle cx="{px(float(r[x])):.2f}" cy="{py(float(r[y])):.2f}" r="4" fill="#9f3e20"><title>{esc(x)} {r[x]}: {r[y]:.4f}</title></circle>')
            for frac in [0,.5,1]:
                val=lo+frac*(hi-lo);marks.append(f'<text x="60" y="{py(val)+5}" text-anchor="end">{val:.3g}</text>')
            marks.append(f'<path d="M70 40 V300 H645" fill="none" stroke="#59756b"/><text x="70" y="325">{rows[0][x]}</text><text x="640" y="325" text-anchor="end">{rows[-1][x]}</text><text x="340" y="350" text-anchor="middle">{esc(x)}</text><text x="70" y="25">{esc(y)}</text>')
        else:
            vals=[float(r[y]) for r in rows];lo=min(0,min(vals));hi=max(.001,max(vals));scale=lambda v:200+440*(v-lo)/max(1e-9,hi-lo);zero=scale(0);height=min(46,270/max(1,len(rows)))
            for i,r in enumerate(rows):
                yy=50+i*height;end=scale(float(r[y]));marks.append(f'<text x="187" y="{yy+14}" text-anchor="end">{esc(r[x])}</text><rect x="{min(zero,end):.2f}" y="{yy}" width="{max(.5,abs(end-zero)):.2f}" height="{height*.65}" rx="2" fill="#26735c"/><text x="650" y="{yy+14}">{float(r[y]):.3g}</text>')
            marks.append(f'<text x="200" y="25">{esc(y)}</text>')
        graphic=f'<svg viewBox="0 0 720 365" role="img" aria-label="{esc(y)} chart; exact values in the evidence table" style="width:100%;min-width:560px;background:#f7f8f2;border-radius:12px;font:14px system-ui;fill:#173d36"><title>{esc(y)} chart</title>{"".join(marks)}</svg>'
        if out.get('heat') is not None:
            heat=np.asarray(out['heat']);maxv=max(1e-9,float(np.max(np.abs(heat))));cells=[]
            for i,row in enumerate(heat):
                cells.append(f'<text x="80" y="{44+i*24}" text-anchor="end">{esc(out["heat_labels"][i])}</text>')
                for j,val in enumerate(row):
                    opacity=.1+.9*abs(float(val))/maxv;color='#26735c' if val>=0 else '#a94322';cells.append(f'<rect x="{90+j*32}" y="{27+i*24}" width="29" height="21" fill="{color}" opacity="{opacity:.3f}"><title>Row {i}, column {j}: {val:.5f}</title></rect>')
            for j in range(heat.shape[1]):cells.append(f'<text x="{96+j*32}" y="18">{j}</text>')
            graphic+=f'<svg viewBox="0 0 720 {65+24*len(heat)}" role="img" aria-label="Value heat map. Green positive, rust negative. Exact values in the heat map table." style="width:100%;min-width:560px;font:12px system-ui;fill:#173d36"><title>Heat map with indexed columns</title>{"".join(cells)}</svg>'
        return '<div tabindex="0" role="group" aria-label="Scrollable experiment charts" style="max-width:100%;overflow-x:auto">'+graphic+'</div>'

    MODEL = json.loads(zlib.decompress(base64.b64decode('eNpNjUEKwzAMBL8idC59QHvutX9w7E0tGktBdqAh5O9xCZSedlmG2Y1j8MQ32rhYwtQbP41mR/MgikQ9tY7mBU5L7YMotSyV1BoGszdfiAtatq+GH595kiiNGkLMoi86vR2q0Wb8M9NKCeP58qMRc1Cp5U4V6E8gR0XwmGmwRVPw9cr7fgBankFO')))
    return mo, compute, draw, json, MODEL

@app.cell
def _():
    config = {'id': 'u6-calendar', 'year': 6, 'title': 'Around the year with Pip', 'concept': 'Cyclic representations', 'question': 'How can a model reuse addition to reason about months?', 'intro': 'Move around a calendar, record the ordinary sum and then wrap back to January. The circle is a model you can fully inspect.', 'controls': [{'key': 'month', 'label': 'Starting month: January = 0', 'kind': 'slider', 'min': 0, 'max': 11, 'step': 1, 'default': 10}, {'key': 'step', 'label': 'Months to move forward', 'kind': 'slider', 'min': 0, 'max': 24, 'step': 1, 'default': 4}], 'prediction': 'Starting at November, where will four monthly steps land? What happens if we add another 12?', 'challenge': 'Compare a 2-step trip with a 14-step trip. Same endpoint does not mean the same journey.', 'artifact': 'An annotated calendar path showing the sum, the wrap and what the diagram cannot tell us about a real LLM.', 'teacher': 'Act out months in a circle before using the notebook. The research identifies internal base-10 addition before cyclic remapping; avoid saying Llama literally thinks in our drawn circle.', 'extend': 'Replace 12 months with 7 weekdays in the downloaded code.', 'glossary': {'cycle': 'A sequence that returns to its start.', 'remainder': 'What is left after making whole groups of a size.', 'representation': 'A way to encode information, such as numbers or points.'}, 'method': 'Explicit teaching model', 'minutes': 35}
    sources = [{'id': 'a-geometric-calculator', 'title': 'A Geometric Calculator Inside a Neural Network', 'url': 'https://www.goodfire.com/research/a-geometric-calculator', 'kind': 'LLM research', 'date': 'May 14, 2026', 'idea': "Llama's cyclic reasoning reuses an internal base-10 addition mechanism before mapping back to a cycle.", 'year6': 'Add month numbers, then wrap around the calendar.', 'year12': 'Contrast an explanatory geometry with the arithmetic mechanism established by causal interventions.', 'limit': 'The paper does not claim that Llama simply adds directly around a circle. Our calendar is a teaching model.', 'labs': ['u6-calendar', 'u12-geometry'], 'paper': {'url': 'https://arxiv.org/html/2605.01148', 'title': 'Arithmetic in the Wild: Llama uses Base-10 Addition to Reason About Cyclic Concepts'}, 'systemAndData': 'Llama 3.1 8B arithmetic and cyclic tasks', 'method': 'Tracks layer/token representations, identifies a shared addition mechanism and checks it with causal interventions.', 'question': "What evidence would support or challenge this idea: Llama's cyclic reasoning reuses an internal base-10 addition mechanism before mapping back to a cycle.", 'finding': "Llama's cyclic reasoning reuses an internal base-10 addition mechanism before mapping back to a cycle.", 'year6Task': 'Add month numbers, then wrap around the calendar. Use the linked notebook to make two observations. Draw or describe one result and one thing this activity cannot tell us about the source system.', 'year12Task': 'Contrast an explanatory geometry with the arithmetic mechanism established by causal interventions. Record the source model/task, a baseline, the changed factor, a measurement and an alternative explanation. State exactly which part your notebook investigates and which part it does not reproduce.', 'reviewed': '2026-09-07', 'reproduction': 'Independent classroom adaptation; not a reproduction of the source model or complete method.'}, {'id': 'bsf-vision', 'title': 'Uncovering Neural Geometry in Vision Models With Block-Sparse Featurizers', 'url': 'https://www.goodfire.com/research/bsf-vision', 'kind': 'Vision research', 'date': 'July 7, 2026', 'idea': 'Block-sparse features represent concepts using subspaces rather than only single directions.', 'year6': 'A changing shape may need more than one clue.', 'year12': 'Compare one-dimensional features with a two-dimensional cyclic representation.', 'limit': 'The experiments concern vision and diffusion models; the classroom geometry is an analogy for representation, not an LLM replication.', 'labs': ['u6-calendar', 'u12-features'], 'paper': {'url': 'https://arxiv.org/html/2606.25234', 'title': 'Structuring Sparsity: Block-Sparse Featurizers Capture Visual Concept Manifolds'}, 'systemAndData': 'Synthetic manifolds, DINOv3 and SDXL', 'method': 'Trains block-sparse featurizers, compares reconstruction and subspace coverage, and tests edits; multidimensional features are evaluated against direction-based baselines.', 'question': 'What evidence would support or challenge this idea: Block-sparse features represent concepts using subspaces rather than only single directions.', 'finding': 'Block-sparse features represent concepts using subspaces rather than only single directions.', 'year6Task': 'A changing shape may need more than one clue. Use the linked notebook to make two observations. Draw or describe one result and one thing this activity cannot tell us about the source system.', 'year12Task': 'Compare one-dimensional features with a two-dimensional cyclic representation. Record the source model/task, a baseline, the changed factor, a measurement and an alternative explanation. State exactly which part your notebook investigates and which part it does not reproduce.', 'reviewed': '2026-09-07', 'reproduction': 'Independent classroom adaptation; not a reproduction of the source model or complete method.'}, {'id': 'manifold-steering', 'title': 'Steering Along Manifolds to Control Neural Networks', 'url': 'https://www.goodfire.com/research/manifold-steering', 'kind': 'LLM and world-model research', 'date': 'May 7, 2026', 'idea': 'Following a fitted curved representation can control cyclic behaviour more effectively than a straight displacement.', 'year6': 'Compare travelling along a circle with cutting through its middle.', 'year12': 'Measure off-manifold distance and decoded behaviour under two steering paths.', 'limit': 'Only selected fitted manifolds and tasks were tested; not every concept is circular.', 'labs': ['u6-calendar', 'u12-geometry'], 'paper': {'url': 'https://arxiv.org/html/2605.05115', 'title': 'Manifold Steering Reveals the Shared Geometry of Neural Network Representation and Behavior'}, 'systemAndData': 'Llama 3.1 8B weekday behaviour and activations', 'method': 'Fits related behavioural and activation geometry, then compares movement along fitted manifolds with linear edits.', 'question': 'What evidence would support or challenge this idea: Following a fitted curved representation can control cyclic behaviour more effectively than a straight displacement.', 'finding': 'Following a fitted curved representation can control cyclic behaviour more effectively than a straight displacement.', 'year6Task': 'Compare travelling along a circle with cutting through its middle. Use the linked notebook to make two observations. Draw or describe one result and one thing this activity cannot tell us about the source system.', 'year12Task': 'Measure off-manifold distance and decoded behaviour under two steering paths. Record the source model/task, a baseline, the changed factor, a measurement and an alternative explanation. State exactly which part your notebook investigates and which part it does not reproduce.', 'reviewed': '2026-09-07', 'reproduction': 'Independent classroom adaptation; not a reproduction of the source model or complete method.'}, {'id': 'neural-geometry', 'title': 'The Neural Geometry Series', 'url': 'https://www.goodfire.com/research/neural-geometry', 'kind': 'Series index', 'date': None, 'idea': 'A collection connects cyclic concepts, stories, sparse features and other neural geometry studies.', 'year6': 'Choose a geometry question and follow its original investigation.', 'year12': "Use the series as a reading order, then inspect each paper's evidence separately.", 'limit': 'This is an index, not an additional independent experiment.', 'labs': ['u6-calendar', 'u12-geometry'], 'paper': None, 'systemAndData': 'Research collection across several model domains', 'method': 'Organises related studies of representation geometry; the collection is not an independent experiment beyond its linked papers.', 'question': 'What evidence would support or challenge this idea: A collection connects cyclic concepts, stories, sparse features and other neural geometry studies.', 'finding': 'A collection connects cyclic concepts, stories, sparse features and other neural geometry studies.', 'year6Task': 'Choose a geometry question and follow its original investigation. Use the linked notebook to make two observations. Draw or describe one result and one thing this activity cannot tell us about the source system.', 'year12Task': "Use the series as a reading order, then inspect each paper's evidence separately. Record the source model/task, a baseline, the changed factor, a measurement and an alternative explanation. State exactly which part your notebook investigates and which part it does not reproduce.", 'reviewed': '2026-09-07', 'reproduction': 'Independent classroom adaptation; not a reproduction of the source model or complete method.'}, {'id': 'phylogeny-manifold', 'title': 'Finding the Tree of Life in Evo 2', 'url': 'https://www.goodfire.com/research/phylogeny-manifold', 'kind': 'Genomic research', 'date': 'August 27, 2025', 'idea': 'A learned low-dimensional representation in Evo 2 reflects evolutionary relationships.', 'year6': 'A map can group related things, but we must test it with new examples.', 'year12': 'Ask whether a fitted representation generalises to held-out clades and controls for sequence similarity.', 'limit': 'Biological geometry does not establish an equivalent map inside text LLMs.', 'labs': ['u6-calendar', 'u12-geometry'], 'paper': None, 'systemAndData': 'Evo 2 representations of cross-species DNA', 'method': 'Constructs data to distinguish evolutionary relationship from simple sequence similarity and compares distances along learned geometry.', 'question': 'What evidence would support or challenge this idea: A learned low-dimensional representation in Evo 2 reflects evolutionary relationships.', 'finding': 'A learned low-dimensional representation in Evo 2 reflects evolutionary relationships.', 'year6Task': 'A map can group related things, but we must test it with new examples. Use the linked notebook to make two observations. Draw or describe one result and one thing this activity cannot tell us about the source system.', 'year12Task': 'Ask whether a fitted representation generalises to held-out clades and controls for sequence similarity. Record the source model/task, a baseline, the changed factor, a measurement and an alternative explanation. State exactly which part your notebook investigates and which part it does not reproduce.', 'reviewed': '2026-09-07', 'reproduction': 'Independent classroom adaptation; not a reproduction of the source model or complete method.'}, {'id': 'the-world-inside-neural-networks', 'title': 'The World Inside Neural Networks', 'url': 'https://www.goodfire.com/research/the-world-inside-neural-networks', 'kind': 'Perspective', 'date': 'May 7, 2026', 'idea': 'Geometric structure may help explain how networks represent relationships across domains.', 'year6': 'Look for a useful shape, then ask what the picture leaves out.', 'year12': 'Distinguish a geometric hypothesis from predictive and causal validation.', 'limit': 'A cross-domain research perspective does not prove every useful concept has an easily readable geometry.', 'labs': ['u6-calendar', 'u12-geometry'], 'paper': None, 'systemAndData': 'Cross-domain neural-geometry perspective', 'method': 'Connects structured data to learned representations and an unsupervised geometry-discovery pipeline; individual causal claims require their own experiments.', 'question': 'What evidence would support or challenge this idea: Geometric structure may help explain how networks represent relationships across domains.', 'finding': 'Geometric structure may help explain how networks represent relationships across domains.', 'year6Task': 'Look for a useful shape, then ask what the picture leaves out. Use the linked notebook to make two observations. Draw or describe one result and one thing this activity cannot tell us about the source system.', 'year12Task': 'Distinguish a geometric hypothesis from predictive and causal validation. Record the source model/task, a baseline, the changed factor, a measurement and an alternative explanation. State exactly which part your notebook investigates and which part it does not reproduce.', 'reviewed': '2026-09-07', 'reproduction': 'Independent classroom adaptation; not a reproduction of the source model or complete method.'}]
    return config, sources

@app.cell
def _(mo, config):
    mo.md(f"# {config['title']}\n\n**Year {config['year']} · {config['minutes']} minutes · {config['method']}**\n\n## {config['question']}\n\n{config['intro']}\n\n**Your mission:** predict, change one control, compare evidence, then find a counterexample. All personal stories and classroom data are fictional. Python runs in your browser; first load needs an internet connection for the runtime. Download your work before leaving.")
    return

@app.cell
def _(mo):
    project_file = mo.ui.file(filetypes=['.json'], multiple=False, max_size=2000000, label='Open a saved Brightlab notebook project')
    mo.vstack([mo.md('**Keep your work:** download a project before leaving. Reopen it here to restore settings, comparisons and writing. Files are read inside this notebook; use fictional classroom data.'),project_file])
    return (project_file,)

@app.cell
def _(project_file, json, config, mo):
    restored = {}
    if project_file.contents():
        try:
            _candidate = json.loads(project_file.contents().decode('utf-8'))
            def _tree(value,depth=0):
                if depth>16:raise ValueError('Project nesting is too deep.')
                if isinstance(value,dict):
                    if len(value)>100:raise ValueError('Too many object fields.')
                    for child in value.values():_tree(child,depth+1)
                elif isinstance(value,list):
                    if len(value)>20000:raise ValueError('Too many rows.')
                    for child in value:_tree(child,depth+1)
                elif isinstance(value,float) and not (-1e100<value<1e100):raise ValueError('Project contains a non-finite or excessive number.')
                elif isinstance(value,str) and len(value)>200000:raise ValueError('Project text is too long.')
            _tree(_candidate)
            if not isinstance(_candidate,dict) or _candidate.get('version')!=1:raise ValueError('Choose a supported version 1 notebook project.')
            if _candidate.get('format') != 'brightlab-notebook-project' or _candidate.get('lesson') != config.get('id',config.get('lesson_id')):
                raise ValueError('Choose a project for this notebook.')
            if not isinstance(_candidate.get('settings'),dict) or not isinstance(_candidate.get('prediction'),str) or not isinstance(_candidate.get('saved_runs',[]),list):
                raise ValueError('The project is missing its settings, prediction or comparison list.')
            for _spec in config['controls']:
                _value=_candidate['settings'].get(_spec['key'],_spec['default'])
                if _spec['kind']=='slider' and (type(_value) not in [int,float] or not _spec['min']<=_value<=_spec['max']):raise ValueError('A saved slider is outside its allowed range.')
                if _spec['kind']=='choice' and (type(_value) is not int or not 0<=_value<len(_spec['options'])):raise ValueError('A saved choice is invalid.')
                if _spec['kind']=='text' and (not isinstance(_value,str) or len(_value)>1000):raise ValueError('A saved text control is invalid.')
            if not isinstance(_candidate.get('conclusion',''),str) or len(_candidate.get('saved_runs',[]))>6:raise ValueError('Invalid writing or comparison count.')
            if len(_candidate.get('prediction',''))>10000 or not isinstance(_candidate.get('conclusion',''),str) or len(_candidate.get('conclusion',''))>20000:raise ValueError('Invalid writing in the saved project.')
            _runs=_candidate.get('saved_runs',[])
            if len(_runs)>6:raise ValueError('A project may contain up to six comparisons.')
            for _run in _runs:
                if not isinstance(_run,dict) or not isinstance(_run.get('settings'),dict):raise ValueError('Invalid saved comparison.')
                _rows=_run.get('measurements',_run.get('result',{}).get('rows') if isinstance(_run.get('result'),dict) else None)
                if not isinstance(_rows,list) or not all(isinstance(row,dict) for row in _rows):raise ValueError('A saved comparison needs a measurement table.')
            restored = _candidate
            mo.output.replace(mo.md('Project read. Save the restored prediction to continue your investigation.'))
        except (ValueError,TypeError,UnicodeError,KeyError,AttributeError) as _error:
            mo.output.replace(mo.callout(mo.md('Could not open this project: '+str(_error)),kind='warn'))
    return (restored,)

@app.cell
def _(mo, restored):
    reset_controls = mo.ui.button(value=0,on_click=lambda count:count+1,label='Reset controls to starting settings')
    return (reset_controls,)

@app.cell
def _(mo, config, restored):
    prediction = mo.ui.text_area(value=restored.get("prediction",""),label="My prediction — what will change, and why?", placeholder=config['prediction'], full_width=True).form(submit_button_label='Save prediction & open the lab', clear_on_submit=False, validate=lambda value: None if value and len(value.strip())>=3 else 'Write a short prediction, or ask a partner to record your spoken idea.')
    mo.vstack([mo.md('## 1. Make a prediction\n'+config['prediction']),prediction])
    return (prediction,)

@app.cell
def _(mo, config, prediction, restored, reset_controls):
    mo.stop(prediction.value is None,mo.md('**The experiment opens after you save a prediction.** There is no penalty for being surprised.'))
    _reset=reset_controls.value
    _saved=restored.get('settings',{}) if not _reset else {}
    controls = mo.ui.dictionary({s['key']:(mo.ui.slider(start=s['min'],stop=s['max'],step=s['step'],value=_saved.get(s['key'],s['default']),label=s['label'],show_value=True,full_width=True) if s['kind']=='slider' else mo.ui.text(value=_saved.get(s['key'],s['default']),label=s['label'],full_width=True,debounce=True) if s['kind']=='text' else mo.ui.dropdown(options={label:i for i,label in enumerate(s['options'])},value=s['options'][int(_saved.get(s['key'],s['default']))],label=s['label'],full_width=True)) for s in config['controls']})
    mo.vstack([mo.md('## 2. Change one thing\nThe first result uses the starting settings. Record it before moving a control.'),reset_controls,controls.vstack()])
    return (controls,)

@app.cell
def _(config, compute, controls):
    try:current = compute(config['id'],controls.value)
    except (ValueError,TypeError,KeyError,IndexError) as _error:
        current={'error':str(_error),'summary':'Check your inputs: '+str(_error),'rows':[{'input':'Needs correction','value':0}],'kind':'bars','x':'input','y':'value','formula':'No new result was calculated. Correct the input and compare again.'}
    return (current,)

@app.cell
def _(config, compute, prediction, mo):
    mo.stop(prediction.value is None)
    baseline = compute(config['id'],{s['key']:s['default'] for s in config['controls']})
    return (baseline,)

@app.cell
def _(mo, current, draw, config):
    mo.stop(bool(current.get('error')),mo.callout(mo.md(current['summary']),kind='warn'))
    mo.vstack([mo.callout(mo.md(current['summary']),kind='info'),mo.Html(draw(current)),mo.accordion({'For the curious: how the numbers work':mo.md(current['formula'])}) if config['year']==6 else mo.md('**How this is calculated**\n\n'+current['formula']),mo.md('Charts update from Python calculations. Exact values are below; scroll wide charts sideways on a small screen.')])
    return

@app.cell
def _(mo, current, baseline):
    mo.stop(bool(current.get('error')))
    _tables={'Current measurements':mo.ui.table(current['rows'],selection=None,page_size=10),'Starting-settings comparison':mo.vstack([mo.md(baseline['summary']),mo.ui.table(baseline['rows'],selection=None,page_size=10)])}
    if current.get('metrics'):_tables['Summary measurements and denominators']=mo.json(current['metrics'])
    if current.get('chart_rows'):_tables['Chart measurements']=mo.ui.table(current['chart_rows'],selection=None,page_size=10)
    if current.get('heat') is not None:_tables['Exact heat map values']=mo.ui.table([{'row':i,'label':current['heat_labels'][i],**{str(j):v for j,v in enumerate(row)}} for i,row in enumerate(current['heat'])],selection=None,page_size=10)
    if current.get('model_rows'):_tables['Bridge to measured transformer behaviour']=mo.vstack([mo.md(current['model_method']),mo.ui.table(current['model_rows'],selection=None)])
    mo.accordion(_tables)
    return

@app.cell
def _(mo, restored):
    get_runs, set_runs = mo.state(restored.get('saved_runs',[])[:6])
    return get_runs, set_runs

@app.cell
def _(mo, controls, current, set_runs):
    mo.stop(bool(current.get('error')))
    _snapshot={'settings':dict(controls.value),'result':current}
    save_run=mo.ui.button(label='Save this run for comparison',on_click=lambda _:set_runs(lambda old:(old+[_snapshot])[-6:]))
    clear_runs=mo.ui.button(label='Clear saved comparisons',on_click=lambda _:set_runs([]))
    mo.hstack([save_run,clear_runs])
    return

@app.cell
def _(mo, get_runs):
    _runs=get_runs()
    mo.vstack([mo.md(f'**Saved comparisons: {len(_runs)} / 6.** Your latest six runs stay in this session and are included in the download.'),mo.accordion({f'Run {i+1}':mo.json(r) for i,r in enumerate(_runs)}) if _runs else mo.md('Save a starting run, change one control, then save again.')])
    return

@app.cell
def _(mo, config, current, restored):
    reflection=mo.ui.text_area(value=restored.get("conclusion",""),label='My conclusion — evidence, counterexample and one limitation',full_width=True,debounce=False,placeholder='I changed… I kept… The measurements show… They do not show…')
    mo.vstack([mo.md('## 3. Find the case that breaks your explanation\n'+config['challenge']+'\n\n## 4. Make something with the evidence\n'+config['artifact']),reflection])
    return (reflection,)

@app.cell
def _(mo, config, sources, json, controls, current, baseline, prediction, reflection, get_runs):
    mo.stop(bool(current.get('error')),mo.md('Correct the input before exporting a new result.'))
    _journal={'format':'brightlab-notebook-project','version':1,'lesson':config['id'],'method':config['method'],'prediction':prediction.value,'settings':controls.value,'current':current,'starting_baseline':baseline,'saved_runs':get_runs(),'conclusion':reflection.value,'sources':sources,'scope':'Classroom investigation inspired by research. See method and source limits; not a reproduction of a Goodfire model.'}
    import html as report_html
    _e=report_html.escape
    def _readable(value):
        if isinstance(value,dict):return '<dl>'+''.join('<dt><strong>'+_e(str(k).replace('_',' '))+'</strong></dt><dd>'+_readable(v)+'</dd>' for k,v in value.items())+'</dl>'
        if isinstance(value,list):
            if value and all(isinstance(row,dict) for row in value):
                _keys=list(dict.fromkeys(k for row in value for k in row))
                return '<table><thead><tr>'+''.join('<th>'+_e(str(k).replace('_',' '))+'</th>' for k in _keys)+'</tr></thead><tbody>'+''.join('<tr>'+''.join('<td>'+_readable(row.get(k,''))+'</td>' for k in _keys)+'</tr>' for row in value)+'</tbody></table>'
            return '<ol>'+''.join('<li>'+_readable(v)+'</li>' for v in value)+'</ol>'
        return _e(str(value)) if value is not None else 'Not recorded'
    _body='<html lang="en-AU"><meta charset="utf-8"><title>'+_e(config['title'])+'</title><style>body{font:18px/1.6 system-ui;max-width:950px;margin:40px auto;padding:20px}table{border-collapse:collapse}td,th{border:1px solid #999;padding:10px}p{white-space:pre-wrap}</style><h1>'+_e(config['title'])+'</h1><p>'+_e(config['method'])+'</p><h2>My prediction</h2><p>'+_e(prediction.value or '')+'</p><h2>Settings</h2><p>'+_e(str(controls.value))+'</p><h2>What happened</h2><p>'+_e(current['summary'])+'</p><table><tr>'+''.join('<th>'+_e(k)+'</th>' for k in current['rows'][0])+'</tr>'+''.join('<tr>'+''.join('<td>'+_e(str(v))+'</td>' for v in row.values())+'</tr>' for row in current['rows'])+'</table><h2>My explanation</h2><p>'+_e(reflection.value)+'</p><h2>Method and limits</h2><p>'+_e(current['formula'])+'</p></html>'
    _body=_body.replace('</html>','<h2>Full current evidence</h2>'+_readable(current)+'<h2>Starting baseline</h2>'+_readable(baseline)+'<h2>Saved comparisons</h2>'+_readable(get_runs())+'</html>')
    mo.hstack([mo.download(data=json.dumps(_journal,indent=2,allow_nan=False).encode(),filename=config['id']+'-project.json',label='Download resumable notebook project'),mo.download(data=_body.encode(),filename=config['id']+'-report.html',label='Download readable report / print')])
    return

@app.cell
def _(mo, config, sources, MODEL):
    _young=config['year']==6
    _timing='0–5 min: read and predict. 5–15: make two controlled changes. 15–25: find a counterexample. 25–35: compare and explain.' if _young else '0–8 min: define the hypothesis and metric. 8–25: collect matched runs. 25–40: stress-test the explanation. 40–55: audit the claim and write a limitation.'
    _teacher='**Preparation:** one browser per pair, runtime download permitted, no accounts or personal data required. Try the default experiment before class. A teacher may read instructions aloud and pupils may dictate predictions.\n\n**Learning outcome:** answer “'+config['question']+'” using a controlled comparison.\n\n**Lesson sequence:** '+_timing+'\n\n**Prompts and misconceptions:** '+config['teacher']+'\n\n**Assessment (0–2 each):** a testable prediction; a comparison naming what stayed fixed; accurate use of measurements; a counterexample and appropriately limited conclusion. 0=missing, 1=partial, 2=clear and supported.\n\n**Support:** work through the default and one change together, use the glossary, and accept an oral explanation. **Stretch:** '+config['extend']
    _reading='\n\n'.join('**['+s['title']+']('+s['url']+')**\n\n'+'**System and data:** '+s['systemAndData']+'\n\n**Source method:** '+s['method']+'\n\n**Paper-specific task:** '+s['year6Task' if config['year']==6 else 'year12Task']+'\n\n**Research boundary:** '+s['limit'] for s in sources)
    mo.accordion({'Teacher lesson plan and assessment':mo.md(_teacher),'Words to know':mo.md('\n\n'.join('**'+k+':** '+v for k,v in config['glossary'].items())),'Research connections and limits':mo.md(_reading+'\n\nIndependent Brightlab educational adaptations; no Goodfire endorsement. Full source map: [AI Understanding research library](https://brightlab-ai-creators.ian347727.chatgpt.site/ai-understanding/research).'),'Tiny transformer model card':mo.json(MODEL['card'])})
    return

if __name__ == '__main__':
    app.run()
