File size: 4,318 Bytes
d766458 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import jax
import jax.numpy as jnp
import numpy
from math import log
def _kmeans(X, X_weight, n_clusters=8, n_init=10, max_iter=300, tol=1e-4, seed=0):
'''kmeans implemented in jax'''
def _dist(a,b):
sm = a @ b.T
a_norm = jnp.square(a).sum(-1)
b_norm = jnp.square(b).sum(-1)
return jnp.abs(a_norm[:,None] + b_norm[None,:] - 2 * sm)
def _kmeans_plus_plus(key, X, X_weight, n_clusters):
'''kmeans++ implemented in jax, for initialization'''
n_samples, n_features = X.shape
n_candidates = 2 + int(log(n_clusters))
def loop(m,c):
n,k = c
inf_mask = jnp.inf * (jnp.arange(n_clusters) > n)
p = (inf_mask + _dist(X,m)).min(-1)
# sample candidates
candidates = jax.random.choice(k, jnp.arange(n_samples),
shape=(n_candidates,),
p=p/p.sum(), replace=False)
# pick sample that decreases inertia the most
dist = jnp.minimum(p[:,None],_dist(X,X[candidates]))
i = candidates[(X_weight[:,None] * dist).sum(0).argmin()]
return m.at[n].set(X[i]), None
i = jax.random.choice(key,jnp.arange(n_samples))
init_means = jnp.zeros((n_clusters,n_features)).at[0].set(X[i])
carry = (jnp.arange(1,n_clusters), jax.random.split(key, n_clusters-1))
return jax.lax.scan(loop, init_means, carry)[0]
def _E(means):
# get labels
return _dist(X,means).argmin(-1)
def _M(labels):
# get means
labels = jax.nn.one_hot(labels, n_clusters)
labels = labels * X_weight[:,None]
labels /= labels.sum(0) + 1e-8
return labels.T @ X
def _inertia(means):
# compute score: sum(min(dist(X,means)))
sco = _dist(X,means).min(-1)
return (X_weight * sco).sum()
def single_run(key):
# initialize
init_means = _kmeans_plus_plus(key, X, X_weight, n_clusters)
# run EM
if tol == 0:
means = jax.lax.scan(lambda mu,_:(_M(_E(mu)),None), init_means,
None, length=max_iter)[0]
else:
def EM(x):
old_mu, old_sco, _, n = x
new_mu = _M(_E(old_mu))
new_sco = _inertia(new_mu)
return new_mu, new_sco, old_sco, n+1
def check(x):
_, new_sco, old_sco, n = x
return ((old_sco-new_sco) > tol) & (n < max_iter)
init = EM((init_means,jnp.inf,None,0))
means = jax.lax.while_loop(check, EM, init)[0]
return {"labels":_E(means),
"means":means,
"inertia":_inertia(means)}
# mulitple runs
key = jax.random.PRNGKey(seed)
if n_init > 0:
out = jax.vmap(single_run)(jax.random.split(key,n_init))
i = out["inertia"].argmin()
out = jax.tree_util.tree_map(lambda x:x[i],out)
else:
out = single_run(key)
labels = jax.nn.one_hot(out["labels"],n_clusters)
cat = (labels * X_weight[:,None]).sum(0) / X_weight.sum()
return {**out, "cat":cat}
def kmeans(x, x_weights, k, seed=0, max_iter=300):
N,L,A = x.shape
if k == 1:
kms = {"means":(x*x_weights[:,None,None]).sum(0,keepdims=True)/x_weights.sum(),
"labels":jnp.zeros(N,dtype=int),
"cat":jnp.ones((1,))}
else:
kms = _kmeans(x.reshape(N,-1), x_weights, n_clusters=k, max_iter=max_iter, seed=seed)
kms["means"] = kms["means"].reshape(k,L,A)
return kms
def kmeans_sample(msa, msa_weights, k=1, samples=None, seed=0):
assert k > 0
# run kmeans
kms = kmeans(jnp.asarray(msa), jnp.asarray(msa_weights), k=k, seed=seed)
# sample sequences from kmeans
key = jax.random.PRNGKey(seed)
N,L,A = msa.shape
if samples is None:
# if number of samples is undefined, set to size of input MSA
samples = N
sampled_labels = kms["labels"]
else:
# sample labels
key, key_ = jax.random.split(key)
sampled_labels = jnp.sort(jax.random.choice(key_,jnp.arange(k),shape=(samples,),p=kms["cat"]))
# sample MSA
sampled_msa = kms["means"][sampled_labels]
sampled_msa = (sampled_msa.cumsum(-1) >= jax.random.uniform(key, shape=(samples,L,1))).argmax(-1)
o = {"kms":kms,
"sampled_labels":sampled_labels,
"sampled_msa":sampled_msa}
return jax.tree_util.tree_map(lambda x:np.asarray(x),o) |