weica_control.py

import numpy as np

# https://arxiv.org/abs/2001.04147
# WEICA: nonlinear weighted ICA
# http://ww2.ii.uj.edu.pl/~spurek/publications/WeICA.pdf

# I ported this algorithm in C with `few changes` in step 4 and 5
# to avoid terrible cancellation errors happening with single-precision
# floats e.g handling (mat_mul) Tall and Skinny matrix with float is always an issue.

class weica:
  def fit(self, x):
    # step 1. Whiten data
    x           = x - x.mean(axis = 0)
    uX, sX, vhX = np.linalg.svd(x, full_matrices=False)
    sXinv       = np.diag(1 / sX)
    # Short-hand ^-1/2 term see below.
    covinvsqrt  = vhX.T @ sXinv @ vhX / (x.shape[0] - 1)

    # step 2. Y = COV(Y)^(-1/2) * (X - MEAN(X))
    Y           = x @ covinvsqrt

    # step 3. MY = SUM(NORM(yi)^2 * yi * yi^T)
    MY = np.zeros([Y.shape[1], Y.shape[1]])
    for i in range(Y.shape[0]):
      MY += (np.linalg.norm(Y[i,:]) ** 2) * np.outer(Y[i,:], Y[i,:])

    # step 4. Unmixing Matrix W
    # W = cov(Y)^(-1/2) * Eigenvector(MY)
    wMY, vMY    = np.linalg.eig(MY)
    Wunmixing   = vMY.T @ covinvsqrt

    # step 5. unmixed data:
    x_unmixed   = x @ Wunmixing.T

    return x_unmixed