init.lua

----------------------------------------------------------------------
--
-- Copyright (c) 2011 Clement Farabet
-- (c) 2008 David Manura (for the OptionParser)
--
-- Permission is hereby granted, free of charge, to any person obtaining
-- a copy of this software and associated documentation files (the
-- "Software"), to deal in the Software without restriction, including
-- without limitation the rights to use, copy, modify, merge, publish,
-- distribute, sublicense, and/or sell copies of the Software, and to
-- permit persons to whom the Software is furnished to do so, subject to
-- the following conditions:
-- 
-- The above copyright notice and this permission notice shall be
-- included in all copies or substantial portions of the Software.
-- 
-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
-- EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
-- MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
-- NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
-- LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
-- OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
-- WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
-- 
----------------------------------------------------------------------
-- description:
-- xlua - a package that provides a better Lua prompt, and a few
-- methods to deal with the namespace, argument unpacking
-- and so on...
--
-- history: 
-- July 7, 2011, 12:49AM - added OptionParser from D. Manura
-- June 30, 2011, 4:54PM - creation - Clement Farabet
----------------------------------------------------------------------

local os = require 'os'
local sys = require 'sys'
local io = require 'io'
local math = require 'math'
local torch = require 'torch'

xlua = {}
local _protect_ = {}
for k,v in pairs(_G) do
 table.insert(_protect_, k)
end

-- extra files
require 'xlua.OptionParser'
require 'xlua.Profiler'


----------------------------------------------------------------------
-- better print function
----------------------------------------------------------------------
function xlua.print(obj,...)
 if type(obj) == 'table' then
 local mt = getmetatable(obj)
 if mt and mt.__tostring__ then
 io.write(mt.__tostring__(obj))
 else
 local tos = tostring(obj)
 local obj_w_usage = false
 if tos and not string.find(tos,'table: ') then
 if obj.usage and type(obj.usage) == 'string' then
 io.write(obj.usage)
 io.write('\n\nFIELDS:\n')
 obj_w_usage = true
 else
 io.write(tos .. ':\n')
 end
 end
 io.write('{')
 local tab = ''
 local idx = 1
 for k,v in pairs(obj) do
 if idx > 1 then io.write(',\n') end
 if type(v) == 'userdata' then
 io.write(tab .. '[' .. k .. ']' .. ' = <userdata>')
 else
 local tostr = tostring(v):gsub('\n','\\n')
 if #tostr>40 then
 local tostrshort = tostr:sub(1,40) .. sys.COLORS.none
 io.write(tab .. '[' .. tostring(k) .. ']' .. ' = ' .. tostrshort .. ' ... ')
 else
 io.write(tab .. '[' .. tostring(k) .. ']' .. ' = ' .. tostr)
 end
 end
 tab = ' '
 idx = idx + 1
 end
 io.write('}')
 if obj_w_usage then
 io.write('')
 end
 end
 else
 io.write(tostring(obj))
 end
 if select('#',...) > 0 then
 io.write(' ')
 print(...)
 else
 io.write('\n')
 end
end
rawset(_G, 'xprint', xlua.print)

----------------------------------------------------------------------
-- log all session, by replicating stdout to a file
----------------------------------------------------------------------
function xlua.log(file, append)
 os.execute('mkdir ' .. (sys.uname() ~= 'windows' and '-p ' or '') .. ' "' .. sys.dirname(file) .. '"')
 local mode = 'w'
 if append then mode = 'a' end
 local f = assert(io.open(file,mode))
 io._write = io.write
 _G._print = _G.print
 _G.print = xlua.print
 io.write = function(...)
 io._write(...)
 local arg = {...}
 for i = 1,select('#',...) do
 f:write(arg[i])
 end
 f:flush()
 end
end

----------------------------------------------------------------------
-- clear all globals
----------------------------------------------------------------------
function xlua.clearall()
 for k,v in pairs(_G) do
 local protected = false
 local lib = false
 for i,p in ipairs(_protect_) do
 if k == p then protected = true end
 end
 for p in pairs(package.loaded) do
 if k == p then lib = true end
 end
 if not protected then
 _G[k] = nil
 if lib then package.loaded[k] = nil end
 end
 end
 collectgarbage()
end

----------------------------------------------------------------------
-- clear one variable
----------------------------------------------------------------------
function xlua.clear(var)
 _G[var] = nil
 collectgarbage()
end

----------------------------------------------------------------------
-- prints globals
----------------------------------------------------------------------
function xlua.who()
 local user = {}
 local libs = {}
 for k,v in pairs(_G) do
 local protected = false
 local lib = false
 for i,p in ipairs(_protect_) do
 if k == p then protected = true end
 end
 for p in pairs(package.loaded) do
 if k == p and p ~= '_G' then lib = true end
 end
 if lib then
 table.insert(libs, k)
 elseif not protected then
 user[k] = _G[k]
 end
 end
 xlua.print('')
 xlua.print('Global Libs:')
 xlua.print(libs)
 xlua.print('')
 xlua.print('Global Vars:')
 xlua.print(user)
 xlua.print('')
end

----------------------------------------------------------------------
-- time
----------------------------------------------------------------------
function xlua.formatTime(seconds)
 -- decompose:
 local floor = math.floor
 local days = floor(seconds / 3600/24)
 seconds = seconds - days*3600*24
 local hours = floor(seconds / 3600)
 seconds = seconds - hours*3600
 local minutes = floor(seconds / 60)
 seconds = seconds - minutes*60
 local secondsf = floor(seconds)
 seconds = seconds - secondsf
 local millis = floor(seconds*1000)

 -- string
 local f = ''
 local i = 1
 if days > 0 then f = f .. days .. 'D' i=i+1 end
 if hours > 0 and i <= 2 then f = f .. hours .. 'h' i=i+1 end
 if minutes > 0 and i <= 2 then f = f .. minutes .. 'm' i=i+1 end
 if secondsf > 0 and i <= 2 then f = f .. secondsf .. 's' i=i+1 end
 if millis > 0 and i <= 2 then f = f .. millis .. 'ms' i=i+1 end
 if f == '' then f = '0ms' end

 -- return formatted time
 return f
end
local formatTime = xlua.formatTime

----------------------------------------------------------------------
-- progress bar
----------------------------------------------------------------------
do
 local function getTermLength()
 if sys.uname() == 'windows' then return 80 end
 local tputf = io.popen('tput cols', 'r')
 local w = tonumber(tputf:read('*a'))
 local rc = {tputf:close()}
 if rc[3] == 0 then return w
 else return 80 end 
 end

 local barDone = true
 local previous = -1
 local tm = ''
 local timer
 local times
 local indices
 local termLength = math.min(getTermLength(), 120)
 function xlua.progress(current, goal)
 -- defaults:
 local barLength = termLength - 34
 local smoothing = 100 
 local maxfps = 10
 
 -- Compute percentage
 local percent = math.floor(((current) * barLength) / goal)

 -- start new bar
 if (barDone and ((previous == -1) or (percent < previous))) then
 barDone = false
 previous = -1
 tm = ''
 timer = torch.Timer()
 times = {timer:time().real}
 indices = {current}
 else
 io.write('\r')
 end

 --if (percent ~= previous and not barDone) then
 if (not barDone) then
 previous = percent
 -- print bar
 io.write(' [')
 for i=1,barLength do
 if (i < percent) then io.write('=')
 elseif (i == percent) then io.write('>')
 else io.write('.') end
 end
 io.write('] ')
 for i=1,termLength-barLength-4 do io.write(' ') end
 for i=1,termLength-barLength-4 do io.write('\b') end
 -- time stats
 local elapsed = timer:time().real
 local step = (elapsed-times[1]) / (current-indices[1])
 if current==indices[1] then step = 0 end
 local remaining = math.max(0,(goal - current)*step)
 table.insert(indices, current)
 table.insert(times, elapsed)
 if #indices > smoothing then
 indices = table.splice(indices)
 times = table.splice(times)
 end
 -- Print remaining time when running or total time when done.
 if (percent < barLength) then
 tm = ' ETA: ' .. formatTime(remaining)
 else
 tm = ' Tot: ' .. formatTime(elapsed)
 end
 tm = tm .. ' | Step: ' .. formatTime(step)
 io.write(tm)
 -- go back to center of bar, and print progress
 for i=1,5+#tm+barLength/2 do io.write('\b') end
 io.write(' ', current, '/', goal, ' ')
 -- reset for next bar
 if (percent == barLength) then
 barDone = true
 io.write('\n')
 end
 -- flush
 io.write('\r')
 io.flush()
 end
 end
end

--------------------------------------------------------------------------------
-- prints an error with nice formatting. If domain is provided, it is used as
-- following: <domain> msg
--------------------------------------------------------------------------------
function xlua.error(message, domain, usage)
 local c = sys.COLORS
 if domain then
 message = '<' .. domain .. '> ' .. message
 end
 local col_msg = c.Red .. tostring(message) .. c.none
 if usage then
 col_msg = col_msg .. '\n' .. usage
 end
 error(col_msg)
end
rawset(_G, 'xerror', xlua.error)

--------------------------------------------------------------------------------
-- provides standard try/catch functions
--------------------------------------------------------------------------------
function xlua.trycatch(try,catch)
 local ok,err = pcall(try)
 if not ok then catch(err) end
end

--------------------------------------------------------------------------------
-- returns true if package is installed, rather than crashing stupidly :-)
--------------------------------------------------------------------------------
function xlua.installed(package)
 local found = false
 local p = package.path .. ';' .. package.cpath
 for path in p:gfind('.-;') do
 path = path:gsub(';',''):gsub('?',package)
 if sys.filep(path) then
 found = true
 p = path
 break
 end
 end
 return found,p
end

--------------------------------------------------------------------------------
-- try to load a package, and doesn't crash if not found !
-- optionally try to install it from luarocks, and then load it.
--
-- @param package package to load
-- @param luarocks if true, then try to install missing package with luarocks
-- @param server specify a luarocks server
--------------------------------------------------------------------------------
function xlua.require(package,luarocks,server)
 local loaded
 local load = function() loaded = require(package) end
 local ok,err = pcall(load)
 if not ok then
 print(err)
 print('warning: <' .. package .. '> could not be loaded (is it installed?)')
 return false
 end
 return loaded
end
rawset(_G, 'xrequire', xlua.require)



--http://lua-users.org/wiki/TableUtils
local function table_val_to_str ( v )
 if "string" == type( v ) then
 v = string.gsub( v, "\n", "\\n" )
 if string.match( string.gsub(v,"[^'\"]",""), '^"+$' ) then
 return "'" .. v .. "'"
 end
 return '"' .. string.gsub(v,'"', '\\"' ) .. '"'
 else
 return "table" == type( v ) and xlua.table2string( v ) or tostring( v )
 end
end

local function table_key_to_str ( k )
 if "string" == type( k ) and string.match( k, "^[_%a][_%a%d]*$" ) then
 return k
 else
 return "[" .. table_val_to_str( k ) .. "]"
 end
end

function xlua.table2string(tbl, newline)
 local result, done = {}, {}
 for k, v in ipairs( tbl ) do
 table.insert( result, table_val_to_str( v ) )
 done[ k ] = true
 end
 local s = "="
 if newline then
 s = " : "
 end
 for k, v in pairs( tbl ) do
 if not done[ k ] then
 local line = table_key_to_str( k ) .. s .. table_val_to_str( v )
 table.insert(result, line)
 end
 end
 local res
 if newline then
 res = "{\n " .. table.concat( result, "\n " ) .. "\n}"
 else
 res = "{" .. table.concat( result, "," ) .. "}"
 end
 return res
end


--------------------------------------------------------------------------------
-- standard usage function: used to display automated help for functions
--
-- @param funcname function name
-- @param description description of the function
-- @param example usage example
-- @param ... [optional] arguments
--------------------------------------------------------------------------------
function xlua.usage(funcname, description, example, ...)
 local c = sys.COLORS

 local style = {
 banner = '+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++',
 list = c.blue .. '> ' .. c.none,
 title = c.Magenta,
 pre = c.cyan,
 em = c.Black,
 img = c.red,
 link = c.red,
 code = c.green,
 none = c.none
 }

 local str = style.banner .. '\n'

 str = str .. style.title .. funcname .. style.none .. '\n'
 if description then
 str = str .. '\n' .. description .. '\n'
 end

 str = str .. '\n' .. style.list .. 'usage:\n' .. style.pre

 -- named arguments:
 local args = {...}
 if args[1].tabled then
 args = args[1].tabled 
 end
 if args[1].arg then
 str = str .. funcname .. '{\n'
 for i,param in ipairs(args) do
 local key
 if param.req then
 key = ' ' .. param.arg .. ' = ' .. param.type
 else
 key = ' [' .. param.arg .. ' = ' .. param.type .. ']'
 end
 -- align:
 while key:len() < 40 do
 key = key .. ' '
 end
 str = str .. key .. '-- ' .. param.help 
 if type(param.default) == 'table' then
 str = str .. ' [default = ' .. xlua.table2string(param.default) .. ']'
 elseif param.default or param.default == false then
 str = str .. ' [default = ' .. tostring(param.default) .. ']'
 elseif param.defaulta then
 str = str .. ' [default == ' .. param.defaulta .. ']'
 end
 str = str.. '\n'
 end
 str = str .. '}\n'

 -- unnamed args:
 else
 local idx = 1
 while true do
 local param
 str = str .. funcname .. '(\n'
 while true do
 param = args[idx]
 idx = idx + 1
 if not param or param == '' then break end
 local key
 if param.req then
 key = ' ' .. param.type
 else
 key = ' [' .. param.type .. ']'
 end
 -- align:
 while key:len() < 40 do
 key = key .. ' '
 end
 str = str .. key .. '-- ' .. param.help .. '\n'
 end
 str = str .. ')\n'
 if not param then break end
 end
 end
 str = str .. style.none

 if example then
 str = str .. '\n' .. style.pre .. example .. style.none .. '\n'
 end

 str = str .. style.banner
 return str
end

--------------------------------------------------------------------------------
-- standard argument function: used to handle named arguments, and 
-- display automated help for functions
--------------------------------------------------------------------------------
function xlua.unpack(args, funcname, description, ...)
 -- put args in table
 local defs = {...}

 -- generate usage string as a closure:
 -- this way the function only gets called when an error occurs
 local fusage = function() 
 local example
 if #defs > 1 then
 example = funcname .. '{' .. defs[2].arg .. '=' .. defs[2].type .. ', '
 .. defs[1].arg .. '=' .. defs[1].type .. '}\n'
 example = example .. funcname .. '(' .. defs[1].type .. ',' .. ' ...)'
 end
 return xlua.usage(funcname, description, example, {tabled=defs})
 end
 local usage = {}
 setmetatable(usage, {__tostring=fusage})

 -- get args
 local iargs = {}
 if #args == 0 then
 print(usage)
 xlua.error('error')
 elseif #args == 1 and type(args[1]) == 'table' and #args[1] == 0 
 and not (torch and torch.typename(args[1]) ~= nil) then
 -- named args
 iargs = args[1]
 else
 -- ordered args
 for i = 1,select('#',...) do
 iargs[defs[i].arg] = args[i]
 end
 end

 -- check/set arguments
 local dargs = {}
 for i = 1,#defs do
 local def = defs[i]
 -- is value requested ?
 if def.req and iargs[def.arg] == nil then
 local c = sys.COLORS
 print(c.Red .. 'missing argument: ' .. def.arg .. c.none)
 print(usage)
 xlua.error('error')
 end
 -- get value or default
 dargs[def.arg] = iargs[def.arg]
 if dargs[def.arg] == nil then
 dargs[def.arg] = def.default
 end
 if dargs[def.arg] == nil and def.defaulta then
 dargs[def.arg] = dargs[def.defaulta]
 end
 dargs[i] = dargs[def.arg]
 end

 -- return usage too
 dargs.usage = usage

 -- stupid lua bug: we return all args by hand
 if dargs[65] then
 xlua.error('<xlua.unpack> oups, cant deal with more than 64 arguments :-)')
 end

 -- return modified args
 return dargs,
 dargs[1], dargs[2], dargs[3], dargs[4], dargs[5], dargs[6], dargs[7], dargs[8], 
 dargs[9], dargs[10], dargs[11], dargs[12], dargs[13], dargs[14], dargs[15], dargs[16],
 dargs[17], dargs[18], dargs[19], dargs[20], dargs[21], dargs[22], dargs[23], dargs[24],
 dargs[25], dargs[26], dargs[27], dargs[28], dargs[29], dargs[30], dargs[31], dargs[32],
 dargs[33], dargs[34], dargs[35], dargs[36], dargs[37], dargs[38], dargs[39], dargs[40],
 dargs[41], dargs[42], dargs[43], dargs[44], dargs[45], dargs[46], dargs[47], dargs[48],
 dargs[49], dargs[50], dargs[51], dargs[52], dargs[53], dargs[54], dargs[55], dargs[56],
 dargs[57], dargs[58], dargs[59], dargs[60], dargs[61], dargs[62], dargs[63], dargs[64]
end

--------------------------------------------------------------------------------
-- standard argument function for classes: used to handle named arguments, and 
-- display automated help for functions
-- auto inits the self with usage
--------------------------------------------------------------------------------
function xlua.unpack_class(object, args, funcname, description, ...)
 local dargs = xlua.unpack(args, funcname, description, ...)
 for k,v in pairs(dargs) do
 if type(k) ~= 'number' then
 object[k] = v
 end
 end
end

--------------------------------------------------------------------------------
-- module help function
--
-- @param module module
-- @param name module name
-- @param description description of the module
--------------------------------------------------------------------------------
function xlua.usage_module(module, name, description)
 local c = sys.COLORS
 local hasglobals = false
 local str = c.magenta
 local str = str .. '+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\n'
 str = str .. 'PACKAGE:\n' .. name .. '\n'
 if description then
 str = str .. '\nDESC:\n' .. description .. '\n'
 end
 str = str .. '++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++'
 str = str .. c.none
 -- register help
 local mt = getmetatable(module) or {}
 setmetatable(module,mt)
 mt.__tostring = function() return str end
 return str
end

--------------------------------------------------------------------------------
-- splicing: remove elements from a table
--------------------------------------------------------------------------------
function table.splice(tbl, start, length)
 length = length or 1
 start = start or 1
 local endd = start + length
 local spliced = {}
 local remainder = {}
 for i,elt in ipairs(tbl) do
 if i < start or i >= endd then
 table.insert(spliced, elt)
 else
 table.insert(remainder, elt)
 end
 end
 return spliced, remainder
end

--------------------------------------------------------------------------------
-- prune: remove duplicates from a table
-- if a hash function is provided, it is used to produce a unique hash for each
-- element in the input table.
-- if a merge function is provided, it defines how duplicate entries are merged,
-- otherwise, a random entry is picked.
--------------------------------------------------------------------------------
function table.prune(tbl, hashfunc, merge)
 local hashes = {}
 local hash = hashfunc or function(a) return a end
 if merge then
 for i,v in ipairs(tbl) do
 if not hashes[hash(v)] then 
 hashes[hash(v)] = v
 else
 hashes[hash(v)] = merge(v, hashes[hash(v)])
 end
 end
 else
 for i,v in ipairs(tbl) do
 hashes[hash(v)] = v
 end
 end
 local ntbl = {}
 for _,v in pairs(hashes) do
 table.insert(ntbl, v)
 end
 return ntbl
end

--------------------------------------------------------------------------------
-- split a string using a pattern
--------------------------------------------------------------------------------
function string.split(str, pat)
 local t = {} -- NOTE: use {n = 0} in Lua-5.0
 local fpat = "(.-)" .. pat
 local last_end = 1
 local s, e, cap = str:find(fpat, 1)
 while s do
 if s ~= 1 or cap ~= "" then
 table.insert(t,cap)
 end
 last_end = e+1
 s, e, cap = str:find(fpat, last_end)
 end
 if last_end <= #str then
 cap = str:sub(last_end)
 table.insert(t, cap)
 end
 return t
end

--------------------------------------------------------------------------------
-- eval: just a shortcut to parse strings into symbols
-- example: 
-- assert( string.tosymbol('madgraph.Image.File') == madgraph.Image.File )
--------------------------------------------------------------------------------
function string.tosymbol(str)
 local ok,result = pcall(loadstring('return ' .. str))
 if not ok then
 error(result)
 elseif not result then
 error('symbol "' .. str .. '" does not exist')
 else
 return result
 end
end


--------------------------------------------------------------------------------
-- parses arguments passed as ENV variables
-- example: 
-- learningRate=1e-3 nesterov=false th train.lua
-- opt = xlua.envparams{learningRate=1e-2, nesterov=true}
--------------------------------------------------------------------------------
function xlua.envparams(default)
 local params = {}
 for k, v in pairs(default) do
 params[k] = v
 if os.getenv(k) ~= nil then
 local v_new = os.getenv(k)
 if type(v) == "number" then
 v_new = tonumber(v_new)
 end
 if type(v) == "boolean" then
 if v_new == "false" or v_new == "False" then
 v_new = false
 elseif v_new == "true" or v_new == "True" then
 v_new = true
 end
 end
 assert(v_new ~= nil)
 params[k] = v_new
 end
 end
 return params
end

return xlua