math.lua

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

local ffi  = require "ffi"
local bit  = require "bit"
local xsys = require "xsys"

-- Cache all builtin math functions.
local
abs, acos, asin, atan, atan2, ceil, cos, cosh, deg, exp, floor, fmod, frexp,
huge, ldexp, log, log10, max, min, modf, pi, pow, rad, random, randomseed, sin, 
sinh, sqrt, tan, tanh 
= xsys.from(math, [[
abs, acos, asin, atan, atan2, ceil, cos, cosh, deg, exp, floor, fmod, frexp,
huge, ldexp, log, log10, max, min, modf, pi, pow, rad, random, randomseed, sin, 
sinh, sqrt, tan, tanh
]])

local type = type

-- Step-wise functions ---------------------------------------------------------
-- Halfway cases are rounded away from zero, regardless of the current rounding
-- direction.
local function round(x)
  return x < 0.0 and ceil(x - 0.5) or floor(x + 0.5)
end

local function step(x) -- 1 if x >= 0, 0 otherwise.
  return max(0, min(floor(x) + 1, 1))
end

local function sign(x) -- 1 if x >= 0, -1 otherwise.
  return -1 + step(x)*2
end

-- Special math functions ------------------------------------------------------
-- From Marsaglia, "Evaluating the Normal Distribution"
-- http://www.jstatsoft.org/v11/a05/paper
-- around 15 digits of absolute precision.
local function phi(x)
  -- Ensure 0 <= phi(x) <= 1 :
  if x <= -8 then
    return 0
  elseif x >= 8 then
    return 1
  else
  
  local s, b, q = x, x, x^2
  for i=3,1/0,2 do
    b = b*q/i
    local t = s
    s = t + b
    if s == t then break end
  end
  return 0.5 + s*exp(-0.5*q - 0.91893853320467274178)
  end
end

-- Inverse cdf for sampling based on Peter John Acklam research, see:
-- http://home.online.no/~pjacklam/notes/invnorm/ .
-- Maximum relative error of 1.15E-9 and machine accuracy with refinement.
-- In iphifast domain must be (0, 1) extremes excluded.
local iphifast, iphi
do 
  local a = ffi.new("double[7]", { 0,
  -3.969683028665376e+01,
  2.209460984245205e+02,
  -2.759285104469687e+02,
  1.383577518672690e+02,
  -3.066479806614716e+01,
  2.506628277459239e+00 })
  local b = ffi.new("double[6]", { 0,
  -5.447609879822406e+01,
  1.615858368580409e+02,
  -1.556989798598866e+02,
  6.680131188771972e+01,
  -1.328068155288572e+01 })
  local c = ffi.new("double[7]", { 0,
  -7.784894002430293e-03,
  -3.223964580411365e-01,
  -2.400758277161838e+00,
  -2.549732539343734e+00,
  4.374664141464968e+00,
  2.938163982698783e+00 })
  local d = ffi.new("double[5]", { 0,
  7.784695709041462e-03,
  3.224671290700398e-01,
  2.445134137142996e+00,
  3.754408661907416e+00 })

  -- PERF: just two branches, central with high prob.
  iphifast = function(p)
    -- Rational approximation for central region:
    if abs(p - 0.5) < 0.47575 then -- 95.14% of cases if p ~ U(0, 1).
      local q = p - 0.5
      local r = q^2
      return (((((a[1]*r+a[2])*r+a[3])*r+a[4])*r+a[5])*r+a[6])*q /
             (((((b[1]*r+b[2])*r+b[3])*r+b[4])*r+b[5])*r+1)
    -- Rational approximation for the two ends:
    else
      local iu = ceil(p - 0.97575)      -- 1 if p > 0.97575 (upper).
      local z = (1 - iu)*p + iu*(1 - p) -- p if lower, (1 - p) if upper.
      local sign = 1 - 2*iu             -- 1 if lower, -1 if upper.
      local q = sqrt(-2*log(z))
      return sign*(((((c[1]*q+c[2])*q+c[3])*q+c[4])*q+c[5])*q+c[6]) /
                   ((((d[1]*q+d[2])*q+d[3])*q+d[4])*q+1)
    end
  end
  
  iphi = function(p) 
    if p <= 0 then
      return -1/0
    elseif p >=1 then
      return 1/0
    else
      local x = iphifast(p)
      local e = phi(x) - p
      local u = e*sqrt(2*pi)*exp(x^2/2)
      return x - u/(1 + x*u/2)
    end
   end
  
end

local gamma, loggamma
do
  -- r(10).
  local gamma_r10 = 10.900511
  -- dk[0], ..., dk[10].
  local gamma_dk = ffi.new("double[11]", 
    2.48574089138753565546e-5,
    1.05142378581721974210,
    -3.45687097222016235469,
    4.51227709466894823700,
    -2.98285225323576655721,
    1.05639711577126713077,
    -1.95428773191645869583e-1,
    1.70970543404441224307e-2,
    -5.71926117404305781283e-4,
    4.63399473359905636708e-6,
    -2.71994908488607703910e-9
  )
  local gamma_c = 2*sqrt(exp(1)/pi)

  -- Lanczos approximation, see:
  -- Pugh[2004]: AN ANALYSIS OF THE LANCZOS GAMMA APPROXIMATION
  -- http://bh0.physics.ubc.ca/People/matt/Doc/ThesesOthers/Phd/pugh.pdf
  -- pag 116 for optimal formula and coefficients. Theoretical accuracy of 
  -- 16 digits is likely in practice to be around 14.
  -- Domain: R except 0 and negative integers.
  gamma = function(z)  
    -- Reflection formula to handle negative z plane.
    -- Better to branch at z < 0 as some use cases focus on z >= 0 only.
    if z < 0 then 
      return pi/(sin(pi*z)*gamma(1 - z)) 
    end  
    local sum = gamma_dk[0]
    sum = sum + gamma_dk[1]/(z + 0)
    sum = sum + gamma_dk[2]/(z + 1) 
    sum = sum + gamma_dk[3]/(z + 2) 
    sum = sum + gamma_dk[4]/(z + 3) 
    sum = sum + gamma_dk[5]/(z + 4) 
    sum = sum + gamma_dk[6]/(z + 5) 
    sum = sum + gamma_dk[7]/(z + 6) 
    sum = sum + gamma_dk[8]/(z + 7) 
    sum = sum + gamma_dk[9]/(z + 8) 
    sum = sum + gamma_dk[10]/(z + 9)  
    return gamma_c*((z  + gamma_r10 - 0.5)/exp(1))^(z - 0.5)*sum
  end

  -- Returns log(abs(gamma(z))).
  -- Domain: R except 0 and negative integers.
  loggamma = function(z)
    if z < 0 then 
      return log(pi) - log(abs(sin(pi*z))) - loggamma(1 - z) 
    end  
    local sum = gamma_dk[0]
    sum = sum + gamma_dk[1]/(z + 0)
    sum = sum + gamma_dk[2]/(z + 1) 
    sum = sum + gamma_dk[3]/(z + 2) 
    sum = sum + gamma_dk[4]/(z + 3) 
    sum = sum + gamma_dk[5]/(z + 4) 
    sum = sum + gamma_dk[6]/(z + 5) 
    sum = sum + gamma_dk[7]/(z + 6) 
    sum = sum + gamma_dk[8]/(z + 7) 
    sum = sum + gamma_dk[9]/(z + 8) 
    sum = sum + gamma_dk[10]/(z + 9) 
    -- For z >= 0 gamma function is positive, no abs() required.
    return log(gamma_c) + (z - 0.5)*log(z  + gamma_r10 - 0.5) 
      - (z - 0.5) + log(sum)
  end

end

-- Domain: a > 0 and b > 0.
local function logbeta(a, b)
  if a <= 0 or b <= 0 then return 0/0 end
  return loggamma(a) + loggamma(b) - loggamma(a + b)
end

-- Domain: a > 0 and b > 0.
local function beta(a, b)
  return exp(logbeta(a, b))
end

-- Support for generic arguments -----------------------------------------------
local function recmax(x, y, ...)
  return y and recmax(
    type(x) ~= "number" and x:max(y) or (type(y) ~= "number" and y.max(x, y) or 
      max(x, y)), ...) or x
end

local function recmin(x, y, ...)
  return y and recmin(
    type(x) ~= "number" and x:min(y) or (type(y) ~= "number" and y.min(x, y) or 
      min(x, y)), ...) or x
end

local disp2 = { 
  atan2 = function(x, y)
    return type(x) ~= "number" and x:atan2(y) 
       or (type(y) ~= "number" and y.atan2(x, y) or atan2(x, y))
  end,
  fmod = function(x, y)
    return type(x) ~= "number" and x:fmod(y) 
       or (type(y) ~= "number" and y.fmod(x, y) or fmod(x, y))
  end,
  ldexp = function(x, y)
    return type(x) ~= "number" and x:ldexp(y) 
       or (type(y) ~= "number" and y.ldexp(x, y) or ldexp(x, y))
  end,
  pow = function(x, y)
    return type(x) ~= "number" and x:pow(y) 
       or (type(y) ~= "number" and y.pow(x, y) or pow(x, y))
  end,
  
  beta = function(x, y)
    return type(x) ~= "number" and x:beta(y) 
       or (type(y) ~= "number" and y.beta(x, y) or beta(x, y))
  end,
  logbeta = function(x, y)
    return type(x) ~= "number" and x:logbeta(y) 
       or (type(y) ~= "number" and y.logbeta(x, y) or logbeta(x, y))
  end,
}

local generic = {
  -- Constants:
  pi   = pi,
  huge = huge,
  
  -- Random numbers:
  random     = random,
  randomseed = randomseed,
  
  -- Generic dispatch based on one variable:
  abs   = function(x) return type(x) == "number" and abs(x)   or x:abs() end,
  acos  = function(x) return type(x) == "number" and acos(x)  or x:acos() end,
  asin  = function(x) return type(x) == "number" and asin(x)  or x:asin() end,
  atan  = function(x) return type(x) == "number" and atan(x)  or x:atan() end,
  ceil  = function(x) return type(x) == "number" and ceil(x)  or x:ceil()  end,
  cos   = function(x) return type(x) == "number" and cos(x)   or x:cos() end,
  cosh  = function(x) return type(x) == "number" and cosh(x)  or x:cosh() end,
  deg   = function(x) return type(x) == "number" and deg(x)   or x:deg() end,
  exp   = function(x) return type(x) == "number" and exp(x)   or x:exp() end,
  floor = function(x) return type(x) == "number" and floor(x) or x:floor() end,
  frexp = function(x) return type(x) == "number" and frexp(x) or x:frexp() end,
  log   = function(x) return type(x) == "number" and log(x)   or x:log() end,
  log10 = function(x) return type(x) == "number" and log10(x) or x:log10() end,
  modf  = function(x) return type(x) == "number" and modf(x)  or x:modf() end,
  rad   = function(x) return type(x) == "number" and rad(x)   or x:rad() end,
  sin   = function(x) return type(x) == "number" and sin(x)   or x:sin() end,
  sinh  = function(x) return type(x) == "number" and sinh(x)  or x:sinh() end,
  sqrt  = function(x) return type(x) == "number" and sqrt(x)  or x:sqrt() end,
  tan   = function(x) return type(x) == "number" and tan(x)   or x:tan() end,
  tanh  = function(x) return type(x) == "number" and tanh(x)  or x:tanh() end,  
  
  -- Generic dispatch based on two variables:
  atan2 = disp2.atan2,
  fmod  = disp2.fmod,
  ldexp = disp2.ldexp,
  pow   = disp2.pow,
  
  -- Special dispatch for vararg max and min functions:
  max = recmax,
  min = recmin,  
    
  -- General dispatch based on one variable:
  gamma    = function(x) return type(x) == "number" and gamma(x)    or 
    x:gamma() end,
  iphi     = function(x) return type(x) == "number" and iphi(x)     or 
    x:iphi() end,
  loggamma = function(x) return type(x) == "number" and loggamma(x) or 
    x:loggamma() end,
  phi      = function(x) return type(x) == "number" and phi(x)      or 
    x:phi() end,
  round    = function(x) return type(x) == "number" and round(x)    or 
    x:round() end,
  sign     = function(x) return type(x) == "number" and sign(x)     or 
    x:sign() end,
  step     = function(x) return type(x) == "number" and step(x)     or 
    x:step() end,
  
  -- Generic dispatch based on two variables:
  beta    = disp2.beta,
  logbeta = disp2.logbeta,
  
  -- Note: no generic dispatch for private functions:
  _iphifast = iphifast,
}

--------------------------------------------------------------------------------

local M = xsys.table.union(math, {
  generic     = generic,  
  round       = round,
  step        = step,
  sign        = sign,  
  phi         = phi,
  iphi        = iphi,
  gamma       = gamma,
  loggamma    = loggamma,
  logbeta     = logbeta,
  beta        = beta,  
  _iphifast   = iphifast,
})

M.generic.std = M -- Gives access to builtin functions.

return M