| """Exact Top-k expected accuracy for shared-variance normal observations. |
| |
| After removing the common quadratic term, posterior class ranks are ranks of |
| affine functions of one observation. For each class, every pairwise crossing |
| changes its rank once. Integrate its normal law over all Top-k intervals. |
| No hypothetical draws, observation bins, or future measurements are inputs. |
| """ |
| from __future__ import annotations |
| import torch |
|
|
|
|
| def expected_topk(prior,mean,variance,k=3): |
| assert prior.device.type=='cuda' and prior.ndim==2 and mean.ndim==1 |
| n,h=prior.shape;assert len(mean)==h and variance.ndim==0 |
| k=min(k,h) |
| m=mean/variance |
| intercept=prior.clamp_min(1e-300).log()-mean.square()[None]/(2*variance) |
| |
| difference=m[:,None]-m[None,:] |
| nonzero=difference!=0 |
| crossing=(intercept[:,None,:]-intercept[:,:,None])/torch.where(nonzero,difference,1)[None] |
| crossing=crossing.masked_fill(~nonzero[None],torch.inf) |
| points,order=crossing.sort(-1) |
| change=torch.where(difference>0,-1,1).long()[None].expand(n,-1,-1) |
| change=torch.where(nonzero[None],change,0).gather(-1,order) |
| |
| |
| index=torch.arange(h,device=prior.device) |
| equal_better=(~nonzero[None])&((intercept[:,None,:]>intercept[:,:,None])| |
| ((intercept[:,None,:]==intercept[:,:,None])&(index[None,None,:]<index[None,:,None]))) |
| initial=(difference>0).sum(-1)[None]+equal_better.sum(-1) |
| rank=torch.cat((initial[:,:,None],initial[:,:,None]+change.cumsum(-1)),-1) |
| boundaries=torch.cat((torch.full((n,h,1),-torch.inf,device=prior.device,dtype=prior.dtype),points, |
| torch.full((n,h,1),torch.inf,device=prior.device,dtype=prior.dtype)),-1) |
| standardized=(boundaries-mean[None,:,None])/variance.sqrt() |
| cdf=torch.special.ndtr(standardized);mass=(cdf[:,:,1:]-cdf[:,:,:-1]).clamp_min(0) |
| return (prior*(mass*(rank<k)).sum(-1)).sum(-1) |
|
|
