diff.lua

--------------------------------------------------------------------------------
-- Automatic differentiation module.
--
-- Copyright (C) 2011-2016 Stefano Peluchetti. All rights reserved.
--------------------------------------------------------------------------------

-- TODO: Specialized gamma, loggamma, beta, logbeta.

-- TODO: Introduce specialized matrix multiply, inverse, ecc ecc using BLAS
-- TODO: (when available) based on the results summarized in Mike Giles paper.

-- PERF: We considered one-shot version with matrix-tape which computes the 
-- PERF: gradient in one function call but gains where not substantial and 
-- PERF: memory management is more complicated. Probably better to just 
-- PERF: introduce reverse-mode differentiation directly.

local ffi = require "ffi"
local xsys = require "xsys"
local math = require "sci.math"

local type = type

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,
round, step, sign,
phi, iphi, gamma, loggamma, logbeta, beta
= 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,
round, step, sign,
phi, iphi, gamma, loggamma, logbeta, beta
]])

-- Forward mode, single directional derivative ---------------------------------
local dn -- Dual number: value + adjoint value.

-- Modified from sci.math, works *only* with dual number type.
local dngamma, dnloggamma
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
 -- page 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.
 dngamma = function(z) 
 -- Reflection formula to handle negative z plane.
 -- Better to branch at z < 0 as some probabilistic use cases only consider 
 -- the case z >= 0.
 if z < 0 then 
 return pi/((pi*z):sin()*dngamma(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.
 dnloggamma = function(z)
 if z < 0 then 
 return log(pi) - (pi*z):sin():abs():log() - dnloggamma(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)*(z + gamma_r10 - 0.5):log()
 - (z - 0.5) + sum:log()
 end

end

-- Domain: a > 0 and b > 0.
local function dnlogbeta(a, b)
 if a <= 0 or b <= 0 then return 0/0 end
 local lga = type(a) == "number" and loggamma(a) or dnloggamma(a)
 local lgb = type(b) == "number" and loggamma(b) or dnloggamma(b)
 return lga + lgb - dnloggamma(a + b)
end

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

-- Derivative of phi function:
local function dphi(x)
 return (1/sqrt(2*pi))*exp(-0.5*x^2)
end

-- Use branchless optimization whenever finite values:
local function dnmax(x, y)
 x, y = dn(x), dn(y)
 if max(abs(x._v), abs(y._v)) == 1/0 then
 return x >=y and x or y
 else -- Branchless optimization.
 local z = step(y._v - x._v) -- 1 if y >= x, 0 otherwise.
 return dn(z*y._v + (1 - z)*x._v, z*y._a + (1 - z)*x._a)
 end
end
-- Use branchless optimization whenever finite values:
local function dnmin(x, y)
 x, y = dn(x), dn(y)
 if max(abs(x._v), abs(y._v)) == 1/0 then
 return x <= y and x or y
 else -- Branchless optimization.
 local z = step(x._v - y._v) -- 1 if x >= y, 0 otherwise.
 return dn(z*y._v + (1 - z)*x._v, z*y._a + (1 - z)*x._a)
 end
end

-- Note: dual numbers are immutable: new ones generated by operators.
local dn_mt = {
 __unm = function(x)
 return dn(-x._v, -x._a)
 end,
 __add = function(x, y) x, y = dn(x), dn(y)
 return dn(x._v + y._v, x._a + y._a)
 end,
 __sub = function(x, y) x, y = dn(x), dn(y)
 return dn(x._v - y._v, x._a - y._a)
 end,
 __mul = function(x, y) x, y = dn(x), dn(y)
 return dn(x._v*y._v, x._a*y._v + y._a*x._v)
 end,
 __div = function(x, y) x, y = dn(x), dn(y)
 return dn(x._v/y._v, (x._a*y._v - y._a*x._v)/y._v^2)
 end,
 __pow = function(x, y) -- Optimized version.
 if type(y) == "number" then
 return dn(x._v^y, y*x._v^(y-1)*x._a)
 elseif type(x) == "number" then
 return dn(x^y._v, x^y._v*log(x)*y._a)
 else
 return dn(x._v^y._v, x._v^y._v*(log(x._v)*y._a + y._v/x._v*x._a))
 end
 end,
 __eq = function(x, y) x, y = dn(x), dn(y)
 return x._v == y._v
 end,
 __lt = function(x, y) x, y = dn(x), dn(y)
 return x._v < y._v
 end,
 __le = function(x, y) x, y = dn(x), dn(y)
 return x._v <= y._v
 end,
 __tostring = function(x)
 return tostring(x._v) -- Better to mimic behavior of numbers.
 end,
 __tonumber = function(x) -- Honored only by xsys.string.width.
 return tonumber(x._v)
 end,
 
 copy = function(x)
 return dn(x)
 end,
 
 val = function(x)
 return x._v
 end,
 adj = function(x)
 return x._a
 end,
 
 sin = function(x) return dn(sin(x._v), x._a*cos(x._v)) end,
 cos = function(x) return dn(cos(x._v), x._a*(-sin(x._v))) end,
 tan = function(x) return dn(tan(x._v), x._a*(1 + tan(x._v)^2)) end,
 asin = function(x) return dn(asin(x._v), x._a/sqrt(1 - x._v^2)) end,
 acos = function(x) return dn(acos(x._v), -x._a/sqrt(1 - x._v^2)) end,
 atan = function(x) return dn(atan(x._v), x._a/(1 + x._v^2)) end,
 sinh = function(x) return dn(sinh(x._v), x._a*cosh(x._v)) end,
 cosh = function(x) return dn(cosh(x._v), x._a*sinh(x._v)) end,
 tanh = function(x) return dn(tanh(x._v), x._a*(1 - tanh(x._v)^2)) end,
 exp = function(x) return dn(exp(x._v), x._a*exp(x._v)) end,
 log = function(x) return dn(log(x._v), x._a/x._v) end,
 sqrt = function(x) return dn(sqrt(x._v), x._a/(2*sqrt(x._v))) end,
 abs = function(x) return dn(abs(x._v), x._a*sign(x._v)) end,
 -- Stick to dn type to improve type stability:
 floor = function(x) return dn(floor(x._v), 0) end,
 ceil = function(x) return dn(ceil(x._v), 0) end,
 -- Stick to dn type to improve type stability:
 round = function(x) return dn(round(x._v), 0) end,
 step = function(x) return dn(step(x._v), 0) end,
 sign = function(x) return dn(sign(x._v), 0) end,
 
 gamma = function(x) return dngamma(x) end,
 loggamma = function(x) return dnloggamma(x) end,
 
 beta = function(x, y) return dnbeta(x, y) end,
 logbeta = function(x, y) return dnlogbeta(x, y) end,

 phi = function(x) return dn(phi(x._v), x._a*dphi(x._v)) end,
 iphi = function(x) local y = iphi(x._v); return dn(y, x._a/dphi(y)) end,

 max = dnmax,
 min = dnmin,
}
dn_mt.__index = dn_mt

-- Note: _v comes first. This allows to construct a dual number from a number
-- using dn(number) => adjoint part correctly initialized to 0.
dn = ffi.metatype("struct { double _v, _a; }", dn_mt)

-- To improve type stability we always pass all arguments wrt differentiation
-- will take place as dual numbers, only one of which will have adj part 
-- equal to 1.
local pderf_template = xsys.template[[
local f, dn = f, dn
| local args = { }
| for i=1,n do
| args[i] = "x"..i
| end
| args = concat(args, ",")
return function(${args})
| local dargs = args..","
| for i=1,#dxi do
| local dx = dxi[i]
| dargs = dargs:gsub("x"..dx..",", "dn(x"..dx..",0),")
| end
| dargs = dargs:gsub(",$", "")
| local retadj = { } 
| for i=1,#dxi do
| local from, to = "dn%(x"..dxi[i]..",0%)", "dn%(x"..dxi[i]..",1%)"
local y${i} = dn(f(${dargs:gsub(from, to)}))
| retadj[i] = "y"..i..":adj()"
| end
return y1:val(),${concat(retadj, ",")}
end
]]

local function derivativef(f, n, ...)
 assert(1 <= n, "function's argument # must be positive")
 local dxi
 if select("#", ...) == 0 then
 dxi = { }
 for i=1,n do 
 dxi[i] = i
 end
 else
 dxi = { ... }
 for i=1,#dxi do
 local dx = dxi[i]
 assert(dx <= n, "differentiating variable outside function's argument #")
 end
 end
 local src = pderf_template({ n = n, dxi = dxi, concat = table.concat })
 return xsys.exec(src, "<derivative>", { f = f, dn = dn })
end

-- For forward mode differentiation we could just use grad(f, x, y) to return
-- f(x) and set y to the gradient of f(x), but gradients are best computed in 
-- reverse mode differentiation, and a stack will need to be allocated for f, 
-- the length of which (it's growth-able) likely depends on the f itself hence
-- it's best stored via a closure.
local function gradientf(f, n)
 local algdn = require("sci.alg").typeof(dn, dn)
 local xd = algdn.vec(n)
 return function(x, grad)
 local n = #x
 if n ~= #grad then
 error("value length must be equal to gradient length")
 end
 if n ~= #xd then
 error("value and gradient lengths must be as per initialization")
 end
 local val = 0
 for i=1,n do
 for j=1,n do xd[j] = x[j] end
 xd[i] = xd[i] + dn(0, 1)
 local vd = dn(f(xd))
 val = vd:val()
 grad[i] = vd:adj()
 end
 return val
 end
end

return {
 dn = dn,
 derivativef = derivativef,
 gradientf = gradientf,
}