stat.lua

--------------------------------------------------------------------------------
-- Statistical functions module.
--
-- Copyright (C) 2011-2016 Stefano Peluchetti. All rights reserved.
--------------------------------------------------------------------------------

-- Variances and covariances are computed according to the unbiased version of
-- the algorithm.
-- Welford-type algorithms are used for superior numerical stability, see:
-- http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
-- http://www.johndcook.com/standard_deviation.html

-- TODO: BIC, AIC.
-- TODO: Function join(...) to join results for parallel computing.

-- TODO: Speed-up via OpenBLAS.

local ffi = require "ffi" 
local alg = require "sci.alg"

local sqrt, abs, max, log = math.sqrt, math.abs, math.max, math.log
local vec, mat, arrayct = alg.vec, alg.mat, alg.arrayct
local typeof, metatype = ffi.typeof, ffi.metatype

local function clear(x)
 for i=1,#x do
 x[i] = 0
 end
end

local function mean(x)
 if #x < 1 then
 error("#x >=1 required: #x="..#x)
 end
 local mu = 0
 for i=1,#x do
 mu = mu + (x[i] - mu)/i
 end
 return mu
end

local function var(x)
 if #x < 2 then
 error("#x >= 2 required: #x"..#x)
 end
 local mu, s2 = 0, 0
 for i=1,#x do
 local delta = x[i] - mu
 mu = mu + delta/i
 s2 = s2 + delta*(x[i] - mu)
 end
 return s2/(#x - 1)
end

local function cov(x, y)
 local mux, muy, s2c = 0, 0, 0
 if not #x == #y then
 error("#x ~= #y: #x="..#x..", #y="..#y)
 end
 if #x < 2 then
 error("#x >= 2 required: #x="..#x)
 end
 for i=1,#x do
 local deltax = x[i] - mux
 local deltay = y[i] - muy
 local r = 1/i
 mux = mux + deltax*r
 muy = muy + deltay*r
 s2c = s2c + deltax*(y[i] - muy)
 end
 return s2c/(#x - 1)
end

local function cor(x, y)
 return cov(x, y)/sqrt(var(x)*var(y))
end

local function chk_dim(self, x)
 local d = self._d
 if d ~= #x then
 error("argument with #="..#x.." passed to statistics of dimension="..d)
 end
 return d
end

local function chk_eq_square(X, Y)
 local n, m = X:nrow(), X:ncol()
 local ny, my = Y:nrow(), Y:ncol()
 if not (n == m and n == ny and n == my) then
 error("arguments must be square matrices of equal size, passed: "..n.."x"..m
 ..", "..ny.."x"..my)
 end
 return n, m
end

local function dim(self)
 return self._d
end

local function len(self)
 return self._n
end

local function tos_mean(self)
 local m = vec(self._d)
 self:mean(m)
 return "mean:\n"..m:width()
end

local meand_mt = {
 dim = dim,
 len = len,
 clear = function(self)
 self._n = 0
 clear(self._mu)
 end,
 push = function(self, x)
 local d = chk_dim(self, x)
 self._n = self._n + 1
 for i=1,d do self._mu[i] = self._mu[i] + (x[i] - self._mu[i])/self._n end 
 end,
 mean = function(self, mean)
 local d = chk_dim(self, mean)
 if self._n < 1 then
 error("n >= 1 required: n="..self._n)
 end
 for i=1,d do mean[i] = self._mu[i] end
 end,
 __tostring = tos_mean,
}
meand_mt.__index = meand_mt

local mean0_mt = {
 dim = dim,
 len = len,
 clear = function(self)
 self._n = 0
 self._mu = 0
 end,
 push = function(self, x)
 self._n = self._n + 1
 self._mu = self._mu + (x - self._mu)/self._n
 end,
 mean = function(self)
 return self._mu
 end,
 __tostring = tos_mean,
}
mean0_mt.__index = mean0_mt

local meand_ct = typeof("struct { int32_t _d, _n; $& _mu; }", arrayct)
local mean0_ct = typeof("struct { int32_t _d, _n; double _mu; }")
meand_ct = metatype(meand_ct, meand_mt)
mean0_ct = metatype(mean0_ct, mean0_mt)

local function tos_var(self)
 local v = mat(self._d, self._d)
 self:var(v)
 return tos_mean(self).."\nvar:\n"..v:width()
end

local vard_mt = {
 dim = dim,
 len = len,
 clear = function(self)
 self._n = 0
 clear(self._mu)
 clear(self._s2)
 end,
 push = function(self, x)
 local d = chk_dim(self, x)
 self._n = self._n + 1
 local r = 1/self._n
 for i=1,d do
 self._delta[i] = x[i] - self._mu[i]
 self._mu[i] = self._mu[i] + self._delta[i]*r
 self._s2[i] = self._s2[i] + self._delta[i]*(x[i] - self._mu[i])
 end
 end,
 mean = meand_mt.mean,
 var = function(self, var)
 local d = chk_dim(self, var)
 if self._n < 2 then
 error("n >= 2 required: n="..self._n)
 end
 for i=1,d do var[i] = self._s2[i]/(self._n - 1) end
 end,
 __tostring = tos_var,
}
vard_mt.__index = vard_mt

local var0_mt = {
 dim = dim,
 len = len,
 clear = function(self)
 self._n = 0
 self._mu = 0
 self._s2 = 0
 end,
 push = function(self, x)
 self._n = self._n + 1
 local r = 1/self._n
 self._delta = x - self._mu
 self._mu = self._mu + self._delta*r
 self._s2 = self._s2 + self._delta*(x - self._mu)
 end,
 mean = mean0_mt.mean,
 var = function(self)
 if self._n < 2 then
 error("n >= 2 required: n="..self._n)
 end
 return self._s2/(self._n - 1)
 end,
 __tostring = tos_var,
}
var0_mt.__index = var0_mt

local vard_ct = typeof("struct { int32_t _d, _n; $& _mu; $& _delta; $& _s2; }", 
 arrayct, arrayct, arrayct)
local var0_ct = typeof("struct { int32_t _d, _n; double _mu, _delta, _s2; }")
vard_ct = metatype(vard_ct, vard_mt)
var0_ct = metatype(var0_ct, var0_mt)

-- Y *can* alias X.
local function covtocor(X, Y)
 local n, m = chk_eq_square(X, Y)
 for r=1,n do
 for c=1,m do
 if r ~= c then
 Y[{r,c}] = X[{r,c}]/sqrt(X[{r,r}]*X[{c,c}])
 end
 end
 end
 for i=1,n do Y[{i,i}] = 1 end
end

local function tos_cor(self)
 local c, r = mat(self._d, self._d), mat(self._d, self._d)
 self:cov(c)
 self:cor(r)
 return tos_mean(self).."\ncov:\n"..c:width().."\ncor:\n"..r:width()
end

local covd_mt = {
 dim = dim,
 len = len,
 clear = function(self)
 self._n = 0
 clear(self._mu)
 clear(self._s2)
 end,
 push = function(self, x)
 local d = chk_dim(self, x)
 self._n = self._n + 1
 local r = 1/self._n
 for i=1,d do
 self._delta[i] = x[i] - self._mu[i]
 self._mu[i] = self._mu[i] + self._delta[i]*r
 end
 for i=1,d do for j=1,d do
 self._s2[{i,j}] = self._s2[{i,j}] + self._delta[i]*(x[j] - self._mu[j])
 end end
 end,
 mean = meand_mt.mean,
 var = function(self, var)
 local d = chk_dim(self, var)
 if self._n < 2 then
 error("n >= 2 required: n="..self._n)
 end
 for i=1,d do var[i] = self._s2[{i,i}]/(self._n - 1) end
 end,
 cov = function(self, cov)
 local n, m = chk_eq_square(self._s2, cov)
 if self._n < 2 then
 error("n >= 2 required: n="..self._n)
 end
 for i=1,n do for j=1,m do 
 cov[{i,j}] = self._s2[{i,j}]/(self._n - 1)
 end end
 end,
 cor = function(self, cor)
 self:cov(cor)
 covtocor(cor, cor)
 end,
 __tostring = tos_cor,
}
covd_mt.__index = covd_mt

local covd_ct = typeof("struct { int32_t _d, _n; $& _mu; $& _delta; $& _s2; }", 
 arrayct, arrayct, arrayct)
covd_ct = metatype(covd_ct, covd_mt)

local samples_mt = {
 dim = dim,
 len = len,
 clear = function(self)
 self._n = 0
 self._x = { }
 end,
 push = function(self, x)
 local d = chk_dim(self, x)
 self._n = self._n + 1
 self._x[n] = x:copy()
 end,
 samples = function(self, samples)
 local n, m = samples:nrow(), samples:ncol()
 if n ~= self._n or m ~= self._d then
 error("output matrix has wrong dimensions")
 end
 for i=1,n do for j=1,m do 
 samples[{i,j}] = self._x[{i,j}]
 end end
 end,
}
samples_mt.__index = samples_mt

local anchor = setmetatable({ }, { __mode = "k" })

local function olmean(dim)
 if dim == 0 then
 return mean0_ct(dim)
 else
 local mu = vec(dim)
 local v = meand_ct(dim, 0, mu)
 anchor[v] = { mu }
 return v
 end
end
local function olvar(dim)
 if dim == 0 then
 return var0_ct(dim)
 else
 local mu, delta, s2 = vec(dim), vec(dim), vec(dim)
 local v = vard_ct(dim, 0, mu, delta, s2)
 anchor[v] = { mu, delta, s2 }
 return v
 end
end
local function olcov(dim)
 local mu, delta, s2 = vec(dim), vec(dim), mat(dim, dim)
 local v = covd_ct(dim, 0, mu, delta, s2)
 anchor[v] = { mu, delta, s2 }
 return v
end
local function olsamples(dim)
 return setmetatable({ _d = dim, _n = 0, _x = { } }, samples_mt)
end

return {
 mean = mean,
 var = var,
 cov = cov,
 cor = cor,
 covtocor = covtocor,
 olmean = olmean,
 olvar = olvar, 
 olcov = olcov,
 olsamples = olsamples, 
}