| """Supervised PTM interaction head on externally computed frozen embeddings. |
| |
| Unknown interactions carry NaN labels and never enter classification loss. |
| Paired indices refer to the same peptide and matched target chemistry. |
| """ |
| from dataclasses import dataclass |
| import numpy as np |
| from scipy.optimize import minimize |
| from scipy.special import expit |
|
|
| @dataclass |
| class PTMHead: |
| weight: np.ndarray |
| bias: float |
| mean_b: np.ndarray |
| std_b: np.ndarray |
| mean_t: np.ndarray |
| std_t: np.ndarray |
|
|
| def logits(self, binder, target): |
| b=(np.asarray(binder)-self.mean_b)/self.std_b |
| t=(np.asarray(target)-self.mean_t)/self.std_t |
| return np.einsum('ni,ij,nj->n',b,self.weight,t)+self.bias |
|
|
| def save(self,path): |
| np.savez(path,weight=self.weight,bias=self.bias,mean_b=self.mean_b, |
| std_b=self.std_b,mean_t=self.mean_t,std_t=self.std_t) |
| @classmethod |
| def load(cls,path): |
| with np.load(path,allow_pickle=False) as x: |
| return cls(**{k:x[k] for k in x.files}) |
|
|
|
|
| def loss_gradient(theta,b,t,labels,pairs,pair_weight=1.,margin=1.,l2=1e-3): |
| """BCE on observed labels + pairwise hinge + Frobenius regularization.""" |
| w=theta[:-1].reshape(b.shape[1],t.shape[1]); z=np.einsum('ni,ij,nj->n',b,w,t)+theta[-1] |
| mask=np.isfinite(labels); dz=np.zeros(len(z));loss=0. |
| if mask.any(): |
| y=labels[mask] |
| if not np.isin(y,[0,1]).all():raise ValueError('observed labels must be binary') |
| loss=float(np.mean(np.logaddexp(0,z[mask])-y*z[mask])) |
| dz[mask]=(expit(z[mask])-y)/mask.sum() |
| if len(pairs): |
| pos,neg=np.asarray(pairs,dtype=int).T |
| violation=margin-z[pos]+z[neg];active=violation>0 |
| loss+=pair_weight*np.maximum(violation,0).mean() |
| np.add.at(dz,pos[active],-pair_weight/len(pairs)) |
| np.add.at(dz,neg[active],pair_weight/len(pairs)) |
| loss+=l2*np.sum(w*w) |
| gw=np.einsum('n,ni,nj->ij',dz,b,t)+2*l2*w |
| return loss,np.r_[gw.ravel(),dz.sum()] |
|
|
|
|
| def fit(binder,target,labels,pairs=(),pair_weight=1.,margin=1.,l2=1e-3,maxiter=300): |
| b,t=np.asarray(binder,dtype=float),np.asarray(target,dtype=float);y=np.asarray(labels,dtype=float) |
| if b.ndim!=2 or t.ndim!=2 or len(b)!=len(t) or y.shape!=(len(b),): |
| raise ValueError('expected aligned N x d embedding matrices and N labels') |
| if not np.isfinite(b).all() or not np.isfinite(t).all():raise ValueError('nonfinite embedding') |
| if not np.isfinite(y).any() and not len(pairs):raise ValueError('no supervised observations') |
| if len(pairs) and (np.min(pairs)<0 or np.max(pairs)>=len(b)):raise ValueError('pair index outside training rows') |
| mb,sb=b.mean(0),np.maximum(b.std(0),1e-6);mt,st=t.mean(0),np.maximum(t.std(0),1e-6) |
| bn,tn=(b-mb)/sb,(t-mt)/st |
| result=minimize(loss_gradient,np.zeros(b.shape[1]*t.shape[1]+1),args=(bn,tn,y,pairs,pair_weight,margin,l2), |
| jac=True,method='L-BFGS-B',options={'maxiter':maxiter,'ftol':1e-10}) |
| if not result.success:raise RuntimeError('PTM head optimization failed: '+result.message) |
| head=PTMHead(result.x[:-1].reshape(b.shape[1],t.shape[1]),float(result.x[-1]),mb,sb,mt,st) |
| return head,{'loss':float(result.fun),'iterations':int(result.nit),'observed_labels':int(np.isfinite(y).sum()),'paired_examples':len(pairs)} |
|
|
