File size: 2,763 Bytes
a20151e | 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 | """Online-ish Pearson correlation of all n x n variable pairs simultaneously."""
from typing import NamedTuple
import jax.numpy as jnp
class OnlineMean(NamedTuple):
sum: jnp.ndarray
count: int
@staticmethod
def init(num_features: int):
return OnlineMean(sum=jnp.zeros(num_features), count=0)
def update(self, batch: jnp.ndarray):
return OnlineMean(self.sum + jnp.sum(batch, axis=0), self.count + batch.shape[0])
def mean(self):
return self.sum / self.count
class OnlineCovariance(NamedTuple):
a_mean: jnp.ndarray # (d, )
b_mean: jnp.ndarray # (d, )
cov: jnp.ndarray # (d, d)
var_a: jnp.ndarray # (d, )
var_b: jnp.ndarray # (d, )
count: int
@staticmethod
def init(a_mean: jnp.ndarray, b_mean: jnp.ndarray):
assert a_mean.shape == b_mean.shape
assert len(a_mean.shape) == 1
d = a_mean.shape[0]
return OnlineCovariance(a_mean,
b_mean,
cov=jnp.zeros((d, d)),
var_a=jnp.zeros((d, )),
var_b=jnp.zeros((d, )),
count=0)
def update(self, a_batch, b_batch):
assert a_batch.shape == b_batch.shape
batch_size, _ = a_batch.shape
a_res = a_batch - self.a_mean
b_res = b_batch - self.b_mean
return OnlineCovariance(a_mean=self.a_mean,
b_mean=self.b_mean,
cov=self.cov + a_res.T @ b_res,
var_a=self.var_a + jnp.sum(a_res**2, axis=0),
var_b=self.var_b + jnp.sum(b_res**2, axis=0),
count=self.count + batch_size)
def covariance(self):
return self.cov / (self.count - 1)
def a_variance(self):
return self.var_a / (self.count - 1)
def b_variance(self):
return self.var_b / (self.count - 1)
def a_stddev(self):
return jnp.sqrt(self.a_variance())
def b_stddev(self):
return jnp.sqrt(self.b_variance())
def E_ab(self):
return self.covariance() + jnp.outer(self.a_mean, self.b_mean)
def pearson_correlation(self):
# Note that the 1/(n-1) normalization terms cancel out nicely here.
# TODO: clip?
eps = 0
# Dead units will have zero variance, which produces NaNs. Convert those to
# zeros with nan_to_num.
return jnp.nan_to_num(self.cov / (jnp.sqrt(self.var_a[:, jnp.newaxis]) + eps) /
(jnp.sqrt(self.var_b) + eps))
class OnlineInnerProduct(NamedTuple):
val: jnp.ndarray # (d, d)
@staticmethod
def init(d: int):
return OnlineInnerProduct(val=jnp.zeros((d, d)))
def update(self, a_batch, b_batch):
assert a_batch.shape == b_batch.shape
return OnlineInnerProduct(val=self.val + a_batch.T @ b_batch)
|