| import jax
|
| import jax.numpy as jnp
|
| import numpy as np
|
|
|
| def get_stats(X, X_weight=None, labels=None, add_f_ij=True, add_mf_ij=False, add_c=False):
|
| '''compute f_i/f_ij/f_ijk given msa '''
|
| n = None
|
| if X_weight is None:
|
| Xn = Xs = X
|
| else:
|
| Xn, Xs = X*X_weight[:,n,n], X*jnp.sqrt(X_weight[:,n,n])
|
| f_i = Xn.sum(0)
|
| o = {"f_i": f_i / f_i.sum(1,keepdims=True)}
|
|
|
| if add_f_ij:
|
| f_ij = jnp.tensordot(Xs,Xs,[0,0])
|
| o["f_ij"] = f_ij / f_ij.sum((1,3),keepdims=True)
|
| if add_c: o["c_ij"] = o["f_ij"] - o["f_i"][:,:,n,n] * o["f_i"][n,n,:,:]
|
|
|
| if labels is not None:
|
|
|
| if jnp.issubdtype(labels, jnp.integer):
|
| labels = jax.nn.one_hot(labels,labels.max()+1)
|
| mf_i = jnp.einsum("nc,nia->cia", labels, Xn)
|
| o["mf_i"] = mf_i/mf_i.sum((0,2),keepdims=True)
|
| if add_mf_ij:
|
| mf_ij = jnp.einsum("nc,nia,njb->ciajb", labels, Xs, Xs)
|
| o["mf_ij"] = mf_ij/mf_ij.sum((0,2,4),keepdims=True)
|
| if add_c: o["mc_ij"] = o["mf_ij"] - o["mf_i"][:,:,:,n,n] * o["mf_i"][:,n,n,:,:]
|
| return o
|
|
|
| def get_r(a,b):
|
| a = jnp.array(a).flatten()
|
| b = jnp.array(b).flatten()
|
| return jnp.corrcoef(a,b)[0,1]
|
|
|
| def inv_cov(X, X_weight=None):
|
| X = jnp.asarray(X)
|
| N,L,A = X.shape
|
| if X_weight is None:
|
| num_points = N
|
| else:
|
| X_weight = jnp.asarray(X_weight)
|
| num_points = X_weight.sum()
|
| c = get_stats(X, X_weight, add_mf_ij=True, add_c=True)["c_ij"]
|
| c = c.reshape(L*A,L*A)
|
| shrink = 4.5/jnp.sqrt(num_points) * jnp.eye(c.shape[0])
|
| ic = jnp.linalg.inv(c + shrink)
|
| return ic.reshape(L,A,L,A)
|
|
|
| def get_mtx(W):
|
| W = jnp.asarray(W)
|
|
|
| raw = jnp.sqrt(jnp.sum(np.square(W[:,1:,:,1:]),(1,3)))
|
| raw = raw.at[jnp.diag_indices_from(raw)].set(0)
|
|
|
|
|
| ap = raw.sum(0,keepdims=True) * raw.sum(1,keepdims=True) / raw.sum()
|
| apc = raw - ap
|
| apc = apc.at[jnp.diag_indices_from(apc)].set(0)
|
| return raw, apc
|
|
|
| def con_auc(true, pred, mask=None):
|
| '''compute agreement between predicted and measured contact map'''
|
| true = jnp.asarray(true)
|
| pred = jnp.asarray(pred)
|
| if mask is not None:
|
| mask = jnp.asarray(mask)
|
| idx = mask.sum(-1) > 0
|
| true = true[idx,:][:,idx]
|
| pred = pred[idx,:][:,idx]
|
| eval_idx = jnp.triu_indices_from(true, 6)
|
| pred_, true_ = pred[eval_idx], true[eval_idx]
|
| L = (jnp.linspace(0.1,1.0,10)*len(true)).astype(jnp.int32)
|
| sort_idx = jnp.argsort(pred_)[::-1]
|
| return jnp.asarray([true_[sort_idx[:l]].mean() for l in L])
|
|
|
|
|
