structrestore-3d 0.1.0a1__py3-none-any.whl

structrestore3d/__init__.py

@@ -0,0 +1,3 @@
+"""StructRestore-3D research prototype; methods are not yet validated denoisers."""
+
+__version__ = "0.1.0a1"

structrestore3d/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

structrestore3d/audit_masking.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""Read-only audit of StructRestore-3D masking, split, and neighbor dependence.
+Run from project root. Does not train or overwrite checkpoints.
+"""
+import argparse, csv, json, sys, random
+from pathlib import Path
+import numpy as np
+import tifffile
+import torch
+
+def main():
+    ap=argparse.ArgumentParser()
+    ap.add_argument('--input',default='data/raw')
+    ap.add_argument('--script',default='scripts/structrestore_25d.py')
+    ap.add_argument('--checkpoint',default='results/25d_gated/model.pt')
+    ap.add_argument('--patch',type=int,default=128)
+    ap.add_argument('--samples',type=int,default=100)
+    ap.add_argument('--out',default='results/masking_audit')
+    a=ap.parse_args()
+    torch.set_num_threads(2)
+    random.seed(42); np.random.seed(42); torch.manual_seed(42)
+    out=Path(a.out);out.mkdir(parents=True,exist_ok=True)
+    files=sorted(Path(a.input).glob('slice_*.tif'))
+    assert len(files)>=7, 'Expected at least seven TIFF slices'
+    train,test=files[:-4],files[-4:]
+    print('Training slices:',[f.name for f in train])
+    print('Held-out slices:',[f.name for f in test])
+    print('Filename overlap:',bool(set(train)&set(test)))
+    print('Adjacent train/test boundary:',train[-1].name,'/',test[0].name)
+    print('CAUTION: Adjacent sections are not biologically independent.')
+    print('CAUTION: Noisy neighboring sections can reveal central pixel structure; not a strict blind-spot model.')
+    import importlib.util
+    spec=importlib.util.spec_from_file_location('audit_model',a.script)
+    module=importlib.util.module_from_spec(spec);spec.loader.exec_module(module)
+    ck=torch.load(a.checkpoint,map_location='cpu',weights_only=True)
+    model=module.Net25D(ck['mode'],ck['width'])
+    model.load_state_dict(ck['state_dict']);model.eval()
+    print('Checkpoint mode:',ck['mode'],'width:',ck['width'])
+    # Check center masking: selected target pixels must differ from corrupted center in most cases.
+    ims=[tifffile.memmap(str(f)) for f in train]
+    rows=[]
+    with torch.inference_mode():
+        for i in range(a.samples):
+            z=random.randrange(1,len(ims)-1)
+            h,w=ims[z].shape
+            y=random.randrange(h-a.patch+1);x=random.randrange(w-a.patch+1)
+            arr=np.stack([np.asarray(ims[j][y:y+a.patch,x:x+a.patch],dtype=np.float32)/255
+                          for j in (z-1,z,z+1)])
+            tensor=torch.from_numpy(arr.copy())[None]
+            corrupted,target,mask=module.corrupt_center(tensor)
+            m=mask.numpy()[0,0]
+            original=target.numpy()[0,0]
+            changed=(corrupted[:,1:2].numpy()[0,0]!=original)
+            selected=int(m.sum())
+            changed_count=int((changed&m).sum())
+            rows.append(dict(sample=i,z=z,masked=selected,changed=changed_count,
+                             masked_unchanged=selected-changed_count,
+                             patch_std=float(original.std()),
+                             patch_mean=float(original.mean()),
+                             masked_target_var=float(original[m].var()) if selected else 0.0))
+    with (out/'patch_audit.csv').open('w',newline='') as f:
+        w=csv.DictWriter(f,fieldnames=rows[0].keys());w.writeheader();w.writerows(rows)
+    total_mask=sum(r['masked'] for r in rows)
+    unchanged=sum(r['masked_unchanged'] for r in rows)
+    print(f'Masked target pixels unchanged after replacement: {unchanged}/{total_mask} ({unchanged/max(1,total_mask):.2%})')
+    print('Patch std: min/median/max:',*[round(v,5) for v in
+          (min(r['patch_std'] for r in rows),np.median([r['patch_std'] for r in rows]),max(r['patch_std'] for r in rows))])
+    # Counterfactual probe: keep corrupted center fixed, change only neighboring channels.
+    # Measures neighbor dependence, NOT proof of leakage or restoration quality.
+    z=5;h,w=ims[z].shape;y=(h-a.patch)//2;x=(w-a.patch)//2
+    arr=np.stack([np.asarray(ims[j][y:y+a.patch,x:x+a.patch],dtype=np.float32)/255
+                  for j in (z-1,z,z+1)])
+    base=torch.from_numpy(arr.copy())[None]
+    corrupted,target,mask=module.corrupt_center(base)
+    with torch.inference_mode():
+        baseline=model(corrupted)
+        zero=corrupted.clone();zero[:,0]=0;zero[:,2]=0
+        same=corrupted.clone();same[:,0]=same[:,1];same[:,2]=same[:,1]
+        swap=corrupted.clone();swap[:,0],swap[:,2]=corrupted[:,2].clone(),corrupted[:,0].clone()
+        delta_zero=(model(zero)-baseline).abs().mean().item()
+        delta_same=(model(same)-baseline).abs().mean().item()
+        delta_swap=(model(swap)-baseline).abs().mean().item()
+    print('Neighbor counterfactual mean |output change|:')
+    print('  zero neighbors:',round(delta_zero,6))
+    print('  copy corrupted center into neighbors:',round(delta_same,6))
+    print('  swap neighbors:',round(delta_swap,6))
+    report={
+      'checkpoint':str(a.checkpoint),'mode':ck['mode'],
+      'train_files':[f.name for f in train],'heldout_files':[f.name for f in test],
+      'filename_overlap':bool(set(train)&set(test)),
+      'masked_pixels':total_mask,'masked_pixels_unchanged':unchanged,
+      'masked_unchanged_fraction':unchanged/max(1,total_mask),
+      'patch_std_min':min(r['patch_std'] for r in rows),
+      'patch_std_median':float(np.median([r['patch_std'] for r in rows])),
+      'neighbor_probe_mean_abs_change':{'zero':delta_zero,'copy_center':delta_same,'swap':delta_swap},
+      'limitations':[
+       'This is a code and behavior audit, not proof of no information leakage.',
+       'Adjacent slice channels remain unmasked; correlated neighboring pixels can reveal central structure.',
+       'Random neighboring-pixel replacement can equal the target by chance.',
+       'Training/test adjacent sections are not independent specimens.',
+       'No retraining or performance assessment is performed.'
+      ]}
+    (out/'audit.json').write_text(json.dumps(report,indent=2))
+    print('Saved:',out/'audit.json',out/'patch_audit.csv')
+if __name__=='__main__': main()

structrestore3d/cli.py

@@ -0,0 +1,64 @@
+"""Command-line dispatcher for the existing StructRestore-3D research scripts.
+
+All data and checkpoint paths are relative to the shell working directory.
+"""
+from pathlib import Path
+import runpy
+import sys
+
+from . import __version__
+
+# Command -> packaged script name, positional arguments prepended for existing CLI.
+COMMANDS = {
+    "audit-v1": ("train_v1_audited", ["audit"]),
+    "train-v1": ("train_v1_audited", ["train"]),
+    "compare": ("compare_v0_v1", []),
+    "diagnose": ("diagnose_v1_inference", []),
+    "mask-audit": ("audit_masking", []),
+    "train-single-v0": ("train_fullres", []),
+    "train-25d-v0": ("structrestore_25d", []),
+    "train-baseline": ("train_baseline", []),
+    "infer-baseline": ("infer_baseline", []),
+    "validate-25d": ("validate_25d", []),
+    "register-probe": ("step29_registration_probe", []),
+}
+
+
+def main():
+    if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help", "help"):
+        print(f"StructRestore-3D {__version__} (experimental research software)")
+        print("Usage: structrestore3d COMMAND [OPTIONS]")
+        print("Commands:")
+        for command in COMMANDS:
+            print(f"  {command}")
+        print("\nExample: structrestore3d audit-v1")
+        print("Run 'structrestore3d train-v1 --help' for training arguments.")
+        print("All image data and trained model checkpoints must be supplied separately.")
+        return 0
+    if sys.argv[1] in ("-V", "--version", "version"):
+        print(__version__)
+        return 0
+    command = sys.argv[1]
+    if command not in COMMANDS:
+        print(f"Unknown command: {command!r}; run structrestore3d --help", file=sys.stderr)
+        return 2
+    script, prefix = COMMANDS[command]
+    script_path = Path(__file__).resolve().parent / f"{script}.py"
+    if not script_path.is_file():
+        print(f"The packaged script {script} is missing. Run python tools/sync_package.py before building.", file=sys.stderr)
+        return 1
+    # Original research scripts import their siblings as top-level modules.
+    # Keep backward compatibility without altering the user's scripts.
+    package_dir = str(script_path.parent)
+    if package_dir not in sys.path:
+        sys.path.insert(0, package_dir)
+    arguments = sys.argv[2:]
+    if command == "mask-audit" and "--script" not in arguments:
+        arguments += ["--script", str(script_path.parent / "structrestore_25d.py")]
+    sys.argv = [str(script_path), *prefix, *arguments]
+    runpy.run_path(str(script_path), run_name="__main__")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

structrestore3d/compare_v0_v1.py

@@ -0,0 +1,213 @@
+#!/usr/bin/env python3
+"""Step 27: controlled V0 vs V1 EM FIB-SEM added-noise stress test.
+
+Reference: original *naturally noisy* TIFF, NOT verified clean ground truth.
+V1: masked-partition prediction: each pixel is inferred ONLY while concealed
+     at the same pixel coordinate in all input sections, avoiding unmasked
+     inference mismatch. Masked pixels are given independent uniform filler.
+V0: native inference exactly as originally used. Comparisons confound training
+     changes and inference protocols; not causal architecture ablations.
+No files in existing results directories are written or altered.
+"""
+import argparse, csv, json, sys, time
+from pathlib import Path
+import numpy as np
+import tifffile
+import torch
+from scipy.ndimage import gaussian_filter, median_filter
+from skimage.metrics import structural_similarity, peak_signal_noise_ratio
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
+# Loads existing project architectures without modifying source or checkpoints.
+SCRIPT_DIR=Path(__file__).resolve().parent
+if SCRIPT_DIR.name!='scripts' and (SCRIPT_DIR/'scripts').is_dir():
+    SCRIPT_DIR=SCRIPT_DIR/'scripts'
+sys.path.insert(0,str(SCRIPT_DIR))
+from train_fullres import TinyCNN
+from structrestore_25d import Net25D
+from train_v1_audited import Model as V1Model
+
+DEFAULT_PATHS={
+  'v0_single':'results/baseline_fullres/baseline.pt',
+  'v0_plain':'results/25d_plain/model.pt',
+  'v0_gated':'results/25d_gated/model.pt',
+  'v1_single':'results/v1_single/model.pt',
+  'v1_plain':'results/v1_plain/model.pt',
+  'v1_gated':'results/v1_gated/model.pt',
+}
+
+def load_models(paths,dev):
+    models={}
+    for key,path in paths.items():
+        ck=torch.load(path,map_location='cpu',weights_only=True)
+        mode=key.split('_')[1]
+        if key.startswith('v0'):
+            net=TinyCNN(width=int(ck.get('width',32))) if mode=='single' else Net25D(mode,int(ck.get('width',32)))
+            if mode!='single' and ck.get('mode')!=mode:raise ValueError(f'V0 mode mismatch: {key}: {ck.get("mode")}')
+        else:
+            if ck.get('mode')!=mode:raise ValueError(f'V1 mode mismatch: {key}: {ck.get("mode")}')
+            net=V1Model(mode,int(ck.get('width',32)))
+        net.load_state_dict(ck['state_dict'],strict=True)
+        models[key]=net.to(dev).eval()
+        print('Loaded',key, '<-',path,flush=True)
+    return models
+
+def v0_predict(model,mode,stack,dev,pad):
+    arr=stack[1:2] if mode=='single' else stack
+    with torch.inference_mode():
+        inp=torch.from_numpy(np.ascontiguousarray(arr[None])).to(dev)
+        out=model(inp).clamp(0,1)[0,0].cpu().numpy()
+    return out[pad:-pad,pad:-pad]
+
+def masked_partition_v1(model,stack,dev,pad,groups,rng):
+    """Predict every pixel *once* under a disjoint random ~5% mask partition.
+
+    Independent uniform replacement for each masked channel matches training
+    mask value distribution. All channels use exactly the same boolean mask.
+    """
+    _,h,w=stack.shape
+    group_ids=rng.integers(0,groups,size=(h,w),dtype=np.uint8 if groups<=256 else np.int16)
+    output=np.empty((h-2*pad,w-2*pad),dtype=np.float32)
+    counts=np.zeros(output.shape,dtype=np.uint8)
+    base=torch.from_numpy(np.ascontiguousarray(stack[None])).to(dev)
+    with torch.inference_mode():
+        for g in range(groups):
+            selected=(group_ids==g)
+            if not np.any(selected[pad:-pad,pad:-pad]):continue
+            mask=torch.from_numpy(selected[None,None]).to(dev)
+            # Draw filler per channel independently, as in V1 training.
+            filler=torch.rand_like(base)
+            masked=torch.where(mask.expand_as(base),filler,base)
+            prediction=model(masked).clamp(0,1)[0,0].cpu().numpy()[pad:-pad,pad:-pad]
+            interior=selected[pad:-pad,pad:-pad]
+            output[interior]=prediction[interior]
+            counts[interior]+=1
+    if not np.all(counts==1):raise RuntimeError('Mask partitions did not cover each interior pixel exactly once')
+    return output
+
+def calc_metrics(ref,img):
+    img=np.clip(img,0,1).astype(np.float32)
+    gy,gx=np.gradient(ref)
+    py,px=np.gradient(img)
+    return {'psnr':float(peak_signal_noise_ratio(ref,img,data_range=1.0)),
+            'ssim':float(structural_similarity(ref,img,data_range=1.0)),
+            'edge_rmse':float(np.sqrt(np.mean((gy-py)**2+(gx-px)**2)))}
+
+def main():
+    p=argparse.ArgumentParser(description=__doc__)
+    p.add_argument('--input',default='data/raw')
+    p.add_argument('--out',default='results/v0_v1_comparison')
+    p.add_argument('--crop',type=int,default=256,help='Central output crop width; extra 16px context used')
+    p.add_argument('--trials',type=int,default=2)
+    p.add_argument('--sigmas',type=float,nargs='+',default=[0.05,0.10])
+    p.add_argument('--groups',type=int,default=20,help='V1 disjoint random mask groups, 20 ~ 5%% masked per pass')
+    p.add_argument('--seed',type=int,default=2026)
+    p.add_argument('--models-root',default='.',help='Project root containing model checkpoints')
+    args=p.parse_args()
+    pad=16
+    if args.crop<64 or args.trials<1 or args.groups<2 or args.groups>255 or any(not 0<s<0.5 for s in args.sigmas):
+        p.error('Require crop>=64, trials>=1, 2<=groups<=255 and 0<sigma<0.5')
+    files=sorted(Path(args.input).glob('slice_*.tif'))
+    if len(files)!=16:raise ValueError(f'Expected 16 original TIFFs; found {len(files)}')
+    paths={k:Path(args.models_root)/v for k,v in DEFAULT_PATHS.items()}
+    missing=[str(f) for f in paths.values() if not f.is_file()]
+    if missing:raise FileNotFoundError('Missing checkpoints: '+', '.join(missing))
+    out=Path(args.out)
+    if out.exists() and any(out.iterdir()):
+        raise FileExistsError(f'Output directory is not empty: {out}. Use new --out to avoid overwriting experiments.')
+    out.mkdir(parents=True,exist_ok=True)
+    dev=torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
+    print('Device:',dev,flush=True)
+    models=load_models(paths,dev)
+    rng=np.random.default_rng(args.seed)
+    rows=[];figure_done=False
+    # Require both context sections to belong to the four held-out slices.
+    for z in (13,14):
+        ims=[tifffile.memmap(str(files[j])) for j in (z-1,z,z+1)]
+        h,w=ims[1].shape;side=args.crop
+        if min(h,w)<side+2*pad:raise ValueError('Image too small for crop and 16px halo')
+        centers=[(h//4,w//4),(h//2,w//2),(3*h//4,3*w//4)]
+        for region,(cy,cx) in enumerate(centers):
+            y=max(pad,min(h-side-pad,cy-side//2))
+            x=max(pad,min(w-side-pad,cx-side//2))
+            stack=np.stack([np.asarray(im[y-pad:y+side+pad,x-pad:x+side+pad],dtype=np.float32)/255 for im in ims])
+            reference=stack[1,pad:-pad,pad:-pad].copy()
+            for sigma in args.sigmas:
+                for trial in range(args.trials):
+                    noisy=np.clip(stack+rng.normal(0,sigma,stack.shape).astype(np.float32),0,1)
+                    noisy_center=noisy[1,pad:-pad,pad:-pad]
+                    candidates={'added_noise':noisy_center,
+                                'gaussian_0p8':gaussian_filter(noisy_center,0.8),
+                                'median_3x3':median_filter(noisy_center,size=3)}
+                    times={k:0.0 for k in candidates}
+                    # SHARE one mask assignment and filler across the three V1 networks:
+                    # reset torch generator before each call so masks/filler match.
+                    # This probe is repeatable on CPU and MPS within each run.
+                    rng_state=rng.bit_generator.state
+                    for key,model in models.items():
+                        start=time.perf_counter()
+                        if key.startswith('v0'):
+                            candidates[key]=v0_predict(model,key.split('_')[1],noisy,dev,pad)
+                        else:
+                            rng.bit_generator.state=rng_state
+                            torch.manual_seed(args.seed+100000*z+1000*region+100*trial+int(round(100*sigma)))
+                            candidates[key]=masked_partition_v1(model,noisy,dev,pad,args.groups,rng)
+                        times[key]=time.perf_counter()-start
+                    # Advance RNG state by exactly one V1 mask draw for next noise trial.
+                    for name,img in candidates.items():
+                        m=calc_metrics(reference,img)
+                        rows.append({'slice':f'slice_{z:04d}','region':region,'sigma':sigma,
+                                     'trial':trial,'method':name,'seconds':round(times[name],5),**m})
+                    if not figure_done:
+                        fig,axes=plt.subplots(3,3,figsize=(12,12),constrained_layout=True)
+                        ordered=[('Reference (native noise)',reference),('Added noise',noisy_center),
+                                 ('Gaussian 0.8',candidates['gaussian_0p8']),
+                                 ('V0 single',candidates['v0_single']),('V0 plain',candidates['v0_plain']),
+                                 ('V0 gated',candidates['v0_gated']),('V1 single',candidates['v1_single']),
+                                 ('V1 plain',candidates['v1_plain']),('V1 gated',candidates['v1_gated'])]
+                        for ax,(name,img) in zip(axes.ravel(),ordered):
+                            ax.imshow(img,cmap='gray',vmin=0,vmax=1,interpolation='nearest')
+                            ax.set_title(name);ax.axis('off')
+                        fig.suptitle(f'Added noise sigma {sigma} | {files[z].stem} | region {region}; reference is NOT clean')
+                        fig.savefig(out/'comparison.png',dpi=140)
+                        plt.close(fig);figure_done=True
+        print('Finished',files[z].stem,flush=True)
+    methods=['added_noise','gaussian_0p8','median_3x3']+list(DEFAULT_PATHS)
+    with (out/'per_trial.csv').open('w',newline='') as f:
+        writer=csv.DictWriter(f,fieldnames=list(rows[0]));writer.writeheader();writer.writerows(rows)
+    summary=[]
+    for sigma in args.sigmas:
+        for method in methods:
+            subset=[r for r in rows if r['sigma']==sigma and r['method']==method]
+            r={'sigma':sigma,'method':method,'n_correlated_evaluations':len(subset)}
+            for metric in ['psnr','ssim','edge_rmse','seconds']:
+                vals=np.array([v[metric] for v in subset],dtype=np.float64)
+                r[metric+'_mean']=round(float(vals.mean()),6)
+                r[metric+'_sd']=round(float(vals.std(ddof=1)),6) if len(vals)>1 else 0.
+            summary.append(r)
+    with (out/'summary.csv').open('w',newline='') as f:
+        writer=csv.DictWriter(f,fieldnames=list(summary[0]));writer.writeheader();writer.writerows(summary)
+    metadata={'experiment':'V0 vs V1 / added-noise stress test','device':str(dev),'seed':args.seed,
+      'train_test_boundary':'train slices 0000-0011, heldout 0012-0015',
+      'evaluated_slice_centers':['slice_0013','slice_0014'],
+      'slice_context_always_heldout':True,'crop':args.crop,'trials':args.trials,
+      'sigmas':args.sigmas,'v1_mask_groups':args.groups,
+      'checkpoints':{k:str(v) for k,v in paths.items()},
+      'limitations':[
+         'Original FIB-SEM reference retains unknown native noise; it is NOT clean ground truth.',
+         'Metrics assess fidelity to original after extra simulated noise, not removal of native noise.',
+         'V1 inference predicts every pixel once while directly masked using random disjoint partitions; V0 uses native unmasked inference.',
+         'V1 changes training masking and patch sampling plus inference procedure: not a causal isolation of any one change.',
+         'No independent biological specimens; slices, patches and noise replicates are correlated.',
+         'Edge RMSE is gradient fidelity to noisy original, not verified membrane-preservation accuracy.',
+         'Models were trained on native noisy images; extra Gaussian noise is an out-of-distribution stress test.'
+      ]}
+    (out/'metadata.json').write_text(json.dumps(metadata,indent=2))
+    print('\nMEAN METRICS: high PSNR/SSIM is better; low edge_RMSE is better',flush=True)
+    for r in summary:
+        print(f"sigma={r['sigma']:.2f} {r['method']:<14} PSNR={r['psnr_mean']:6.2f}  SSIM={r['ssim_mean']:.4f}  edge_RMSE={r['edge_rmse_mean']:.4f}",flush=True)
+    print('\nSaved:',out/'summary.csv',out/'per_trial.csv',out/'comparison.png',out/'metadata.json',flush=True)
+
+if __name__=='__main__':main()

structrestore3d/diagnose_v1_inference.py

@@ -0,0 +1,260 @@
+#!/usr/bin/env python3
+"""StructRestore-3D Step 28B: read-only diagnosis of V1 inference.
+
+Uses existing V1 checkpoints. No training. No original files modified.
+The native FIB-SEM slice is noisy, NOT verified clean ground truth.
+
+Diagnoses:
+  - Unmasked vs correctly masked-partition inference on the SAME V1 checkpoint.
+  - Sensitivity to random mask/filler realizations (fixed input image).
+  - Neighboring-slice alignment as a correlation diagnostic (not proof of misregistration).
+  - Direct target leakage in the trained V1 checkpoint, at sampled target coordinates.
+"""
+import argparse
+import csv
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import tifffile
+import torch
+from scipy.ndimage import gaussian_filter
+from skimage.metrics import peak_signal_noise_ratio, structural_similarity
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+if SCRIPT_DIR.name != 'scripts' and (SCRIPT_DIR/'scripts').is_dir():
+    SCRIPT_DIR = SCRIPT_DIR/'scripts'
+sys.path.insert(0, str(SCRIPT_DIR))
+from train_v1_audited import Model, strict_mask
+
+CKPTS = {
+    'single': 'results/v1_single/model.pt',
+    'plain': 'results/v1_plain/model.pt',
+    'gated': 'results/v1_gated/model.pt',
+}
+
+
+def scores(ref, img):
+    ref = ref.astype(np.float32, copy=False)
+    img = np.clip(img, 0, 1).astype(np.float32, copy=False)
+    gy, gx = np.gradient(ref)
+    py, px = np.gradient(img)
+    return dict(psnr=float(peak_signal_noise_ratio(ref, img, data_range=1)),
+                ssim=float(structural_similarity(ref, img, data_range=1)),
+                edge_rmse=float(np.sqrt(np.mean((gy-py)**2 + (gx-px)**2))))
+
+
+def correlation(a, b):
+    aa = (a.astype(np.float64) - np.mean(a)).ravel()
+    bb = (b.astype(np.float64) - np.mean(b)).ravel()
+    den = np.linalg.norm(aa)*np.linalg.norm(bb)
+    return float(np.dot(aa, bb)/den) if den > 1e-10 else None
+
+
+def translated_overlap(a, b, dy, dx):
+    # Compare a[y,x] to b[y+dy,x+dx], without wrapping pixels.
+    h, w = a.shape
+    y0, y1 = max(0, -dy), min(h, h-dy)
+    x0, x1 = max(0, -dx), min(w, w-dx)
+    return a[y0:y1,x0:x1], b[y0+dy:y1+dy,x0+dx:x1+dx]
+
+
+def neighbor_alignment(stack, radius=3):
+    # Low-pass correlations are diagnostic; anatomical changes can also change correlation.
+    smooth = [gaussian_filter(s, 1.5) for s in stack]
+    result = {}
+    for name, neighbor in [('previous', 0), ('next', 2)]:
+        center, other = smooth[1], smooth[neighbor]
+        direct = correlation(center, other)
+        best = (-np.inf, 0, 0)
+        for dy in range(-radius, radius+1):
+            for dx in range(-radius, radius+1):
+                a,b = translated_overlap(center, other, dy, dx)
+                c = correlation(a,b)
+                if c is not None and c > best[0]:best=(c,dy,dx)
+        result[name] = dict(unshifted_correlation=direct,
+                            best_correlation=float(best[0]),
+                            best_relative_shift_pixels=dict(dy=int(best[1]),dx=int(best[2])))
+    return result
+
+
+def infer_unmasked(model, stack_tensor, pad):
+    with torch.inference_mode():
+        pred=model(stack_tensor).clamp(0,1)[0,0].cpu().numpy()
+    return pred[pad:-pad,pad:-pad].copy()
+
+
+def infer_masked(model, stack_tensor, pad, groups, ids, seed):
+    _, _, h, w = stack_tensor.shape
+    out=np.empty((h-2*pad,w-2*pad),dtype=np.float32)
+    assigned=np.zeros_like(out,dtype=np.uint8)
+    # Seeding at each call ensures all modes get exactly the same masks and fillers.
+    torch.manual_seed(int(seed))
+    with torch.inference_mode():
+        for g in range(groups):
+            selected=(ids==g)
+            inside=selected[pad:-pad,pad:-pad]
+            if not inside.any():continue
+            mask=torch.from_numpy(selected[None,None]).to(stack_tensor.device)
+            filler=torch.rand_like(stack_tensor)
+            corrupted=torch.where(mask.expand_as(stack_tensor),filler,stack_tensor)
+            pred=model(corrupted).clamp(0,1)[0,0].cpu().numpy()[pad:-pad,pad:-pad]
+            out[inside]=pred[inside]
+            assigned[inside]+=1
+    if not np.all(assigned==1):
+        raise RuntimeError('Partition failure: every output pixel must be predicted once')
+    return out
+
+
+def check_trained_direct_leakage(models,dev):
+    torch.manual_seed(247)
+    x=torch.rand((1,3,40,40),device=dev,requires_grad=True)
+    corrupted,_,mask=strict_mask(x,rate=0.1)
+    coords=torch.nonzero(mask[0,0])
+    if len(coords)==0:raise RuntimeError('Mask test selected no pixels')
+    coords=coords[:5]
+    checks={}
+    for mode,model in models.items():
+        max_g=0.
+        for i, coord in enumerate(coords):
+            y, xx = [int(v) for v in coord]
+            pred=model(corrupted)
+            grads=torch.autograd.grad(pred[0,0,y,xx],x,retain_graph=True)[0]
+            val=float(grads[0,:,y,xx].abs().max().detach().cpu())
+            max_g=max(max_g,val)
+        checks[mode]={'max_direct_target_gradient':max_g,'passed':bool(max_g<1e-10)}
+        print(f'CHECK trained {mode}: max direct target gradient={max_g:.3g};',
+              'PASS' if checks[mode]['passed'] else 'FAIL', flush=True)
+    if not all(x['passed'] for x in checks.values()):
+        raise RuntimeError('Direct-masking audit failed')
+    return checks
+
+
+def main():
+    ap=argparse.ArgumentParser(description=__doc__)
+    ap.add_argument('--input',default='data/raw')
+    ap.add_argument('--out',default='results/v1_inference_diagnostic')
+    ap.add_argument('--models-root',default='.')
+    ap.add_argument('--slice-index',type=int,default=13,choices=[13,14])
+    ap.add_argument('--region',type=int,choices=[0,1,2],default=0,help='Match Step 27 region 0/1/2; default 0 reproduces comparison figure')
+    ap.add_argument('--crop',type=int,default=256)
+    ap.add_argument('--groups',type=int,default=20)
+    ap.add_argument('--repeats',type=int,default=2)
+    ap.add_argument('--sigma',type=float,default=0.05,help='Additional Gaussian noise on all three sections')
+    ap.add_argument('--seed',type=int,default=2026)
+    args=ap.parse_args()
+    if args.crop<64 or not 2<=args.groups<=255 or not 2<=args.repeats<=10 or not 0<args.sigma<0.5:
+        ap.error('Require crop>=64, 2<=groups<=255, 2<=repeats<=10, 0<sigma<0.5')
+    out=Path(args.out)
+    if out.exists() and any(out.iterdir()):
+        raise FileExistsError(f'{out} is not empty. Choose a new --out directory.')
+    filenames=sorted(Path(args.input).glob('slice_*.tif'))
+    if len(filenames)!=16:raise ValueError(f'Expected 16 original TIFFs; found {len(filenames)}')
+    dev=torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
+    ckpaths={m:Path(args.models_root)/p for m,p in CKPTS.items()}
+    missing=[str(p) for p in ckpaths.values() if not p.exists()]
+    if missing:raise FileNotFoundError('Missing checkpoints: '+', '.join(missing))
+    models={}
+    for mode,path in ckpaths.items():
+        checkpoint=torch.load(path,map_location='cpu',weights_only=True)
+        if checkpoint.get('mode')!=mode:raise ValueError(f'Wrong mode in {path}')
+        model=Model(mode,int(checkpoint.get('width',32)))
+        model.load_state_dict(checkpoint['state_dict'],strict=True)
+        models[mode]=model.to(dev).eval()
+        print('Loaded:',mode,path,flush=True)
+    checks=check_trained_direct_leakage(models,dev)
+    pad=16
+    z=args.slice_index
+    ims=[tifffile.memmap(str(filenames[j])) for j in (z-1,z,z+1)]
+    H,W=ims[1].shape
+    if min(H,W)<args.crop+2*pad:raise ValueError('Crop too large for TIFF')
+    centers=[(H//4,W//4),(H//2,W//2),(3*H//4,3*W//4)]
+    cy,cx=centers[args.region]
+    y=max(pad,min(H-args.crop-pad,cy-args.crop//2))
+    x=max(pad,min(W-args.crop-pad,cx-args.crop//2))
+    # Read only selected crops, rather than entire full-resolution images into RAM.
+    stack=np.stack([np.asarray(im[y-pad:y+args.crop+pad,x-pad:x+args.crop+pad],dtype=np.float32)/255. for im in ims])
+    ref=stack[1,pad:-pad,pad:-pad].copy()
+    alignment=neighbor_alignment(stack[:,pad:-pad,pad:-pad])
+    print('Neighboring-section correlations:',json.dumps(alignment),flush=True)
+    rng=np.random.default_rng(args.seed)
+    noisy=np.clip(stack+rng.normal(0,args.sigma,size=stack.shape).astype(np.float32),0,1)
+    # Both inference modes operate on exactly the same simulated noisy input.
+    inp=torch.from_numpy(np.ascontiguousarray(noisy[None])).to(dev)
+    raw=noisy[1,pad:-pad,pad:-pad]
+    predictions={'added_noise':raw,'gaussian_0p8':gaussian_filter(raw,0.8)}
+    rows=[]; variations=[];first_masked={}
+    for name,img in predictions.items():
+        rows.append(dict(mode=name,inference='baseline',repeat=0,**scores(ref,img)))
+    for mode,model in models.items():
+        unmasked=infer_unmasked(model,inp,pad)
+        predictions[mode+'_unmasked']=unmasked
+        rows.append(dict(mode=mode,inference='unmasked',repeat=0,**scores(ref,unmasked)))
+        pred_runs=[]
+        for rep in range(args.repeats):
+            # Sample a fresh disjoint mask partition; fixed noisy input for all repeats.
+            ids=np.random.default_rng(args.seed+1009*rep).integers(0,args.groups,
+                    size=noisy.shape[1:],dtype=np.uint8)
+            pred=infer_masked(model,inp,pad,args.groups,ids,args.seed+3001*rep)
+            rows.append(dict(mode=mode,inference='masked',repeat=rep,**scores(ref,pred)))
+            pred_runs.append(pred)
+        predictions[mode+'_masked']=pred_runs[0]
+        first_masked[mode]=pred_runs[0]
+        diffs=[np.abs(pr-pred_runs[0]) for pr in pred_runs[1:]]
+        differences=np.stack(diffs)
+        variations.append(dict(mode=mode,
+                mean_abs_prediction_difference=float(differences.mean()),
+                p95_abs_prediction_difference=float(np.percentile(differences,95)),
+                max_abs_prediction_difference=float(differences.max())))
+    out.mkdir(parents=True,exist_ok=True)
+    with (out/'per_repeat.csv').open('w',newline='') as f:
+        writer=csv.DictWriter(f,fieldnames=rows[0].keys());writer.writeheader();writer.writerows(rows)
+    summaries=[]
+    for name in ['added_noise','gaussian_0p8']+list(models):
+        modes=['baseline'] if name in ('added_noise','gaussian_0p8') else ['unmasked','masked']
+        for imode in modes:
+            subset=[row for row in rows if row['mode']==name and row['inference']==imode]
+            summaries.append(dict(method=name,inference=imode,n=len(subset),
+                **{k+'_mean':float(np.mean([row[k] for row in subset])) for k in ('psnr','ssim','edge_rmse')}))
+    with (out/'summary.csv').open('w',newline='') as f:
+        writer=csv.DictWriter(f,fieldnames=summaries[0].keys());writer.writeheader();writer.writerows(summaries)
+    with (out/'mask_variability.csv').open('w',newline='') as f:
+        writer=csv.DictWriter(f,fieldnames=variations[0].keys());writer.writeheader();writer.writerows(variations)
+    figure=[('Native noisy reference',ref),('Added noise',raw),('Gaussian 0.8',predictions['gaussian_0p8'])]
+    for mode in ('single','plain','gated'):
+        figure.extend([(mode+' unmasked',predictions[mode+'_unmasked']),
+                       (mode+' masked',predictions[mode+'_masked'])])
+    fig,axes=plt.subplots(3,3,figsize=(12,12),constrained_layout=True)
+    for ax,(name,img) in zip(axes.ravel(),figure):
+        ax.imshow(img,cmap='gray',vmin=0,vmax=1,interpolation='nearest')
+        ax.set_title(name);ax.axis('off')
+    fig.suptitle(f'Inference-only diagnosis | slice_{z:04d} | sigma={args.sigma} | native reference NOT clean')
+    fig.savefig(out/'inference_comparison.png',dpi=140)
+    plt.close(fig)
+    report=dict(experiment='V1 same-checkpoint inference comparison',source_slice=z,region=args.region,top_left_yx=[y,x],
+      reference='Native noisy FIB-SEM TIFF, not clean truth',extra_gaussian_sigma=args.sigma,
+      groups=args.groups,repeats=args.repeats,seed=args.seed,alignment=alignment,
+      trained_checkpoint_direct_mask_tests=checks,
+      checkpoints={k:str(v) for k,v in ckpaths.items()},
+      limitations=[
+          'The masked inference method is closer to the V1 training corruption, while unmasked inference is distribution-shifted.',
+          'Masked and unmasked results test inference protocol sensitivity, not comparative denoising efficacy.',
+          'The original noisy reference is NOT clean; PSNR, SSIM and edge-RMSE are fidelity measures only.',
+          'Even if the direct gradient test passes, adjacent pixels/slices remain correlated.',
+          'Intensity cross-correlation of adjacent sections does not distinguish misregistration from genuine structural changes.',
+          'All evaluation regions are from one physical stack; not independent biological replicates.'
+      ])
+    (out/'metadata.json').write_text(json.dumps(report,indent=2))
+    print('\nSame V1 checkpoint; higher PSNR/SSIM and lower edge RMSE = more faithful to NATIVE noisy reference:')
+    for r in summaries:
+        print(f"  {r['method']:<14} {r['inference']:<9} PSNR={r['psnr_mean']:6.2f} "
+              f"SSIM={r['ssim_mean']:.4f} edge_RMSE={r['edge_rmse_mean']:.4f}")
+    print('Mask-partition sensitivity (mean absolute difference between repeated runs):')
+    for r in variations:print(f"  {r['mode']}: {r['mean_abs_prediction_difference']:.6f}")
+    print('Saved:',out/'summary.csv',out/'mask_variability.csv',out/'inference_comparison.png',flush=True)
+
+if __name__=='__main__':main()