whitening.py

#coding:utf-8

import numpy as np

class WhiteningScaler:
    def __init__(self, axis=0, threshold=1e-12, apply_zca=False):
        # 行方向にデータが格納されている場合は axis=0
        # 列方向にデータが格納されている場合は axis=1
        self.axis = axis
        self.threshold = threshold
        # ZCAをオンにしなければ正規化されたPCAと等価
        self.apply_zca = apply_zca
        pass
    
    def fit(self, X):
        # データは列方向に格納する
        if self.axis == 0:
            self.X = np.array(X.T)
        else:
            self.X = np.array(X)
        
        # Centering
        self.mean = np.mean(X, axis=1)
        self.X = self.X - self.mean
        
        # SVD
        U, s, _ = np.linalg.svd(self.X, full_matrices=False)
        s[s < self.threshold] = self.threshold
        S_inv = np.diag(np.ones(s.shape) / s)
        
        self.U = U
        self.S_inv = S_inv
        
        # Decorrelation
        self.X = U.T.dot(self.X)
        
        # Whitening
        self.X = S_inv.dot(self.X)
        
        # ZCA(Zero-phase Component Analysis) Whitening
        if self.apply_zca:
            self.X = U.dot(self.X)
        
        if self.axis == 0:
            return np.array(self.X.T)
        else:
            return np.array(self.X)