src/generate.py

"""Grouped deterministic equilibrium dataset; no physical-time claims."""
from pathlib import Path
from dataclasses import asdict
import json,time,platform,hashlib
import numpy as np
from rod import Rod,Controls,serial

ROOT=Path(__file__).resolve().parents[1]

def generate():
    rng=np.random.default_rng(20261014)
    r=Rod(48); records=[]; rejected=[]; audits=[]; discarded=[]
    start=time.perf_counter()
    # Each family has 60 ID scenes and 12 parameter-OOD scenes; all paths and
    # mirrored branches of a scene share a single split.
    for family in ['end_load','intrinsic_curvature','microforce']:
        order=rng.permutation(60)
        assignment={int(j):('train' if i<42 else 'validation' if i<51 else 'test') for i,j in enumerate(order)}
        for j in range(72):
            split=assignment[j] if j<60 else 'parameter_ood'
            scene=f'{family}_{j:03d}'
            chi=float(10**rng.uniform(-4,-2.5)) if j<60 else float(10**rng.uniform(-2.5,-2.1))
            base=Controls(chi=chi)
            if family=='microforce':
                base.h=float(chi*rng.uniform(4,30));base.delta=float(rng.uniform(.05,.35))
                pc=(-1+np.sqrt(1+16*base.h*base.delta-np.pi**2*chi))/2-2*base.h
                paths=['down','straight','up']
                maxcompression=float(rng.uniform(2,8))
            else:
                fx=float(chi*rng.uniform(0,2));fy=float(chi*rng.uniform(-2,2));moment=float(chi*rng.uniform(-1,1))
                ks=rng.uniform(-2.4,2.4,3) if family=='intrinsic_curvature' else np.zeros(3)
                paths=['ramp','unload','staged']
            for path in paths:
                prev=None
                for step,t in enumerate(np.linspace(0,1,25)):
                    c=Controls(**asdict(base))
                    seed_source='previous_equilibrium' if prev is not None else 'zero'
                    seed=None if prev is None else prev['z']
                    if family=='microforce':
                        # All points below the first critical load; these are branch
                        # sweeps, NOT claims about spontaneous branch selection.
                        c.fx=float(pc-chi*(.2+(maxcompression-.2)*t))
                        sign={'down':-1,'straight':0,'up':1}[path]
                        if prev is None:
                            seed=r.mode_seed(sign*1.2);seed_source='signed_mode_initialization'
                        if sign==0:seed=np.zeros(r.d);seed_source='exact_straight'
                    else:
                        alpha=float(1-t if path=='unload' else t)
                        c.fx=fx*alpha;c.fy=fy*alpha;c.moment=moment*alpha
                        ramps=np.clip(3*alpha-np.arange(3),0,1) if path=='staged' else np.full(3,alpha)
                        c.k1,c.k2,c.k3=map(float,ks*ramps)
                    result=r.solve(c,seed)
                    if family=='microforce' and path!='straight' and (result['tip'][1]*sign<1e-4 or not result['accepted']):
                        for amp in [.6,1.6,2.5,3.]:
                            discarded.append(dict(id=f'{scene}_{path}_{step:02d}',scene_id=scene,split=split,controls=asdict(c),reason='wrong_branch_or_not_converged',**serial(result)))
                            result=r.solve(c,r.mode_seed(sign*amp));seed_source='signed_mode_retry'
                            if result['accepted'] and result['tip'][1]*sign>1e-4:break
                    valid=result['accepted']
                    if family=='microforce' and path!='straight':valid &= result['tip'][1]*sign>1e-4
                    row=dict(id=f'{scene}_{path}_{step:02d}',scene_id=scene,split=split,family=family,
                             path_id=f'{scene}_{path}',path_kind=path,step=step,path_parameter=float(t),
                             controls=asdict(c),model='angle_FE_midpoint_n48',seed_source=seed_source,**serial(result))
                    row['branch']='straight' if np.max(np.abs(result['z']))<1e-5 else ('up' if result['tip'][1]>0 else 'down')
                    row['previous_id']=None if prev is None else prev['id']
                    row['continuous_stable_pair']=bool(prev is not None and prev['stability']=='stable' and result['stability']=='stable' and np.max(np.abs(prev['z']-result['z']))<.2)
                    if valid:
                        records.append(row);prev={**result,'id':row['id']}
                    else:
                        rejected.append(row);prev=None
                # One paired mesh audit per scenario, assigned the same scene/split.
                if path==paths[-1] and prev is not None:
                    fine=Rod(96)
                    seed=np.interp(np.linspace(0,1,97)[1:],np.linspace(0,1,49),np.r_[0,prev['z']])
                    f=fine.solve(c,seed)
                    audits.append(dict(scene_id=scene,split=split,coarse_id=prev['id'],controls=asdict(c),
                                       fine=serial(f),centerline_rmse=float(np.sqrt(np.mean(np.sum((f['centerline']-prev['centerline'])**2,axis=1))))))
            if j%12==0:print(f'{family} {j}/72: accepted={len(records)} rejected={len(rejected)}',flush=True)
    out=ROOT/'data';out.mkdir(exist_ok=True)
    for name,rows in [('equilibria',records),('rejected_attempts',rejected),('discarded_seed_attempts',discarded),('paired_resolution',audits)]:
        with (out/f'{name}.jsonl').open('w') as f:
            for row in rows:f.write(json.dumps(row,allow_nan=False,separators=(',',':'))+'\n')
    summary=dict(version='0.1.0',seed=20261014,accepted_states=len(records),rejected_final_states=len(rejected),discarded_seed_attempts=len(discarded),scenes=216,
                 paired_resolution_cases=len(audits),seconds=time.perf_counter()-start,
                 splits={s:sum(v['split']==s for v in records) for s in ['train','validation','test','parameter_ood']},
                 families={s:sum(v['family']==s for v in records) for s in ['end_load','intrinsic_curvature','microforce']},
                 stability={s:sum(v['stability']==s for v in records) for s in ['stable','unstable','critical']},
                 max_residual=max(v['residual'] for v in records),
                 mesh_audit_max_rmse=max(v['centerline_rmse'] for v in audits),
                 mesh_audit_median_rmse=float(np.median([v['centerline_rmse'] for v in audits])),
                 python=platform.python_version(),numpy=np.__version__)
    # Guard against leakage and malformed paths before publication.
    grouped={}
    for v in records:grouped.setdefault(v['scene_id'],set()).add(v['split'])
    assert all(len(s)==1 for s in grouped.values())
    assert all(v['fine']['accepted'] for v in audits)
    (ROOT/'reports/generation.json').write_text(json.dumps(summary,indent=2))
    print(json.dumps(summary,indent=2))

if __name__=='__main__':generate()