1import json, math, random
  2import numpy as np
  3import torch
  4import torch.nn.functional as F
  5
  6SEED=386
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9DT=torch.float64
 10
 11def make_problem(n=24):
 12    # A spatially varying detector mask after a mild blur: right side has sparse coverage.
 13    yy,xx=np.mgrid[0:n,0:n]; N=n*n
 14    # periodic detector coverage creates a deliberately nonuniform field of view
 15    coverage=(0.15 + 0.85*(xx < int(.58*n)).astype(float) + 0.18*(xx>=int(.58*n)))
 16    # separable Gaussian blur matrix
 17    coords=np.arange(n)
 18    G=np.exp(-((coords[:,None]-coords[None,:])**2)/(2*1.35**2)); G/=G.sum(1,keepdims=True)
 19    B=np.kron(G,G)
 20    # detector rows are image pixels, with coverage weights
 21    m=coverage.reshape(-1)
 22    A=(m[:,None]**0.5)*B
 23    # two equal objects, one in high and one in low sensitivity areas
 24    u=np.zeros((n,n))
 25    for cy,cx in [(12,6),(12,19)]:
 26        u[((yy-cy)**2+(xx-cx)**2)<=3.0**2]=1.0
 27    # slight asymmetric shape so centroid/size are informative
 28    u[9:13,18:21]=1.15
 29    u0=u.reshape(-1)
 30    d=A@u0
 31    noise=.10*np.std(d)*np.random.randn(N)
 32    return A, u0, d+noise, coverage.reshape(-1)
 33
 34def tv_torch(u,w,n,eps=1e-3):
 35    im=u.reshape(n,n)
 36    dx=im[:,1:]-im[:,:-1]; dy=im[1:,:]-im[:-1,:]
 37    # node weights on corresponding left/top nodes, a standard anisotropic boundary convention
 38    wx=w.reshape(n,n)
 39    tx=(wx[:,:-1]*torch.sqrt(dx*dx+eps*eps)).sum()
 40    ty=(wx[:-1,:]*torch.sqrt(dy*dy+eps*eps)).sum()
 41    return (tx+ty)/(n*n)
 42
 43def solve(A,d,w,n,lam,steps=700,lr=.08):
 44    At=torch.tensor(A,dtype=DT); dt=torch.tensor(d,dtype=DT)
 45    u=torch.zeros(n*n,dtype=DT,requires_grad=True)
 46    opt=torch.optim.Adam([u],lr=lr)
 47    for k in range(steps):
 48        opt.zero_grad()
 49        res=At@u-dt
 50        loss=.5*(res@res)/len(d)+lam*tv_torch(u,w,n)
 51        loss.backward(); opt.step()
 52        with torch.no_grad(): u.clamp_(0,1.4)
 53    return u.detach().numpy()
 54
 55def metrics(u,truth,A,d,w,n):
 56    im=u.reshape(n,n); gt=truth.reshape(n,n)
 57    mse=np.mean((u-truth)**2); psnr=10*np.log10(1.4**2/max(mse,1e-15))
 58    # centroid and mass/area separately in high and low sensitivity halves
 59    out={'psnr':float(psnr),'residual':float(np.linalg.norm(A@u-d)/np.sqrt(len(d))),
 60         'tv':float(tv_torch(torch.tensor(u,dtype=DT),torch.tensor(w,dtype=DT),n).item())}
 61    for name,sel in [('high',np.arange(n)<int(.58*n)),('low',np.arange(n)>=int(.58*n))]:
 62        mask=np.broadcast_to(sel[:,None],(n,n)) if False else np.broadcast_to(sel[None,:],(n,n))
 63        # mask columns; centroid x and thresholded area, using positive reconstructed mass
 64        mass=np.maximum(im,0)[mask].sum(); true_mass=gt[mask].sum()
 65        xs=np.tile(np.arange(n),(n,1))[mask]
 66        cx=(xs*np.maximum(im,0)[mask]).sum()/max(mass,1e-12)
 67        truecx=(xs*gt[mask]).sum()/max(true_mass,1e-12)
 68        area=(im[mask]>.5).sum(); truearea=(gt[mask]>.5).sum()
 69        out[name+'_mass_ratio']=float(mass/max(true_mass,1e-12)); out[name+'_centroid_abs']=float(abs(cx-truecx)); out[name+'_area_ratio']=float(area/max(truearea,1))
 70    return out
 71
 72def main():
 73    n=24; A,truth,d,coverage=make_problem(n)
 74    # Exact discrete sensitivity and direct numerical response check.
 75    s=np.linalg.norm(A,axis=0); delta=1e-10; w=(s+delta)/np.mean(s+delta)
 76    rng=np.random.default_rng(SEED+1); inds=rng.choice(n*n,12,replace=False)
 77    observed=np.array([np.linalg.norm(A[:,i]) for i in inds]); relerr=float(np.max(abs(observed-s[inds])/(observed+1e-15)))
 78    # finite perturbation response K(u+h e_i)-K(u), confirming linear scaling and weights
 79    h=1e-4; base=A@truth
 80    fd=np.array([np.linalg.norm((A@(truth+h*np.eye(n*n)[i])-base)/h) for i in inds])
 81    scale_err=float(np.max(abs(fd-observed)/(observed+1e-15)))
 82    mathcheck={'max_column_norm_relative_error':relerr,'max_finite_difference_relative_error':scale_err,
 83               'sensitivity_min_max_ratio':float(s.min()/s.max()),'weight_min_max': [float(w.min()),float(w.max())]}
 84    # same lambda, plus a small validation-like sweep for scalar TV vs weighted TV
 85    torch.manual_seed(SEED)
 86    zero=np.zeros(n*n); results={}
 87    for name,ww,lam in [('no_tv',zero,0.0),('ordinary_tv',np.ones(n*n),0.030),('weighted_tv',w,0.030)]:
 88        rec=solve(A,d,torch.tensor(ww,dtype=DT),n,lam)
 89        results[name]=metrics(rec,truth,A,d,ww,n)
 90    # lambda sweep illustrates whether improvement is merely coefficient tuning
 91    sweep=[]
 92    for lam in [0.01,0.02,0.03,0.05,0.08]:
 93        for name,ww in [('ordinary',np.ones(n*n)),('weighted',w)]:
 94            rec=solve(A,d,torch.tensor(ww,dtype=DT),n,lam,steps=500)
 95            mm=metrics(rec,truth,A,d,ww,n); sweep.append({'method':name,'lambda':lam,**mm})
 96    best_ord=max([x for x in sweep if x['method']=='ordinary'],key=lambda x:x['psnr'])
 97    best_w=max([x for x in sweep if x['method']=='weighted'],key=lambda x:x['psnr'])
 98    report={'seed':SEED,'n':n,'noise_std':float(.10*np.std(A@truth)),'mathcheck':mathcheck,'results':results,'sweep_best_psnr':{'ordinary':best_ord,'weighted':best_w}}
 99    with open('results.json','w') as f: json.dump(report,f,indent=2)
100    print(json.dumps(report,indent=2))
101if __name__=='__main__': main()