| """Strong independent-factor controls sharing the same numerical evidence.""" |
| from __future__ import annotations |
| import torch |
| from torch.nn.functional import one_hot |
| from .factorial import hypotheses |
|
|
|
|
| def fit(x,h,sizes,kind): |
| d=x.shape[1];eye=torch.eye(d,device=x.device,dtype=x.dtype);parts=[] |
| if kind=='factorized_ridge': |
| z=torch.cat((x,x.new_ones((len(x),1))),-1) |
| penalty=eye.new_zeros((d+1,d+1));penalty.diagonal().fill_(8.);penalty[-1,-1]=.01 |
| inverse=torch.linalg.inv(z.T@z+penalty) |
| for i,size in enumerate(sizes): |
| label=h[:,i];count=torch.bincount(label,minlength=size).double() |
| sums=x.new_zeros((size,d)).scatter_add_(0,label[:,None].expand_as(x),x) |
| mean=(sums+x.mean(0))/(count[:,None]+1) |
| residual=x-mean[label] |
| scatter=x.new_zeros((size,d,d)).scatter_add_(0,label[:,None,None].expand(-1,d,d),residual[:,:,None]*residual[:,None,:]) |
| pool=scatter.sum(0)/max(1,len(x)-size) |
| if kind=='factorized_nb': |
| var=(scatter.diagonal(dim1=-2,dim2=-1)+4*pool.diagonal()[None])/(count[:,None]+4) |
| covariance=torch.diag_embed(var)+.05*eye |
| elif kind=='factorized_lda': |
| covariance=(.5*pool+.5*torch.diag_embed(pool.diagonal())+.05*eye)[None].expand(size,-1,-1) |
| elif kind=='factorized_kernel': |
| overall=(x-x.mean(0)).T@(x-x.mean(0))/max(1,len(x)-1) |
| covariance=(.5*overall+.5*eye)*.7**2 |
| elif kind=='factorized_ridge': |
| |
| target=one_hot(label,size).double()/(count[None]+1) |
| weights=inverse@z.T@target |
| covariance=eye |
| else:raise ValueError(kind) |
| parts.append(dict(mean=mean,cov=covariance,count=count, |
| centers=x,labels=label,weights=weights if kind=='factorized_ridge' else None,size=size)) |
| return dict(parts=parts,sizes=tuple(sizes),kind=kind,grid=hypotheses(sizes,x.device),neural_parameters=0) |
|
|
|
|
| def posterior(model,x,observed): |
| ix=list(observed);n=len(x);logs=[] |
| for part in model['parts']: |
| size=part['size'] |
| if not ix: |
| p=x.new_full((n,size),1/size) |
| elif model['kind']=='factorized_kernel': |
| cov=part['cov'][ix][:,ix];chol=torch.linalg.cholesky(cov) |
| means=torch.cat((part['centers'],part['mean']),0) |
| delta=x[:,None,ix]-means[None,:,ix] |
| white=torch.linalg.solve_triangular(chol,delta.reshape(-1,len(ix)).T,upper=False) |
| loglike=-.5*white.square().sum(0).reshape(n,-1) |
| result=[] |
| for a in range(size): |
| components=torch.cat((torch.where(part['labels']==a)[0],part['labels'].new_tensor([len(part['centers'])+a]))) |
| weight=x.new_ones(len(components));weight[-1]=.5 |
| result.append(torch.logsumexp(loglike[:,components]+weight.log()[None],-1)-(part['count'][a]+.5).log()) |
| p=torch.stack(result,-1).softmax(-1) |
| elif model['kind']=='factorized_ridge': |
| |
| |
| |
| scores=x[:,ix]@part['weights'][ix]+part['weights'][-1] |
| p=scores.clamp_min(1e-4);p/=p.sum(-1,keepdim=True) |
| else: |
| cov=part['cov'][:,ix][:,:,ix];chol=torch.linalg.cholesky(cov) |
| delta=(x[:,None,ix]-part['mean'][None,:,ix]).permute(1,2,0) |
| white=torch.linalg.solve_triangular(chol,delta,upper=False) |
| mahal=white.square().sum(1).T |
| logdet=chol.diagonal(dim1=-2,dim2=-1).log().sum(-1) |
| p=(-.5*mahal-logdet[None]).softmax(-1) |
| logs.append(p.clamp_min(1e-300).log()[:,model['grid'][:,len(logs)]]) |
| return torch.stack(logs).sum(0).softmax(-1) |
|
|
