simplifying model parameters to be simple defaults

Author: jsalvatier

pymc/model.py

@@ -12,7 +12,7 @@
 import numpy as np 
 from functools import wraps
 
-__all__ = ['Model', 'compilef', 'gradient', 'hessian', 'withmodel', 'Point'] 
+__all__ = ['Model', 'compilef', 'gradient', 'hessian', 'modelcontext', 'Point'] 
 
 
 
@@ -38,36 +38,10 @@ def get_context(cls):
         except IndexError:
             raise TypeError("No context on context stack")
 
-def withcontext(contexttype, argname):
-    """
-    Returns a decorator for wrapping functions so they look for an argument in a specific argument slot. 
-    If not found, the decorated function searches the for a context and inserts it in that slot. 
-
-    Parameters
-    ----------
-    contexttype : type
-        The type of context to search for
-    argname : string
-        The name of the argument slot where the context should go
-
-    Returns 
-    -------
-    decorator function
-
-    """
-    def decorator(fn):
-        n = list(fn.func_code.co_varnames).index(argname)
-
-        @wraps(fn)
-        def nfn(*args, **kwargs):
-            if not (len(args) > n and isinstance(args[n], contexttype)):
-                context = contexttype.get_context()
-                args = args[:n] + (context,) + args[n:]
-            return fn(*args,**kwargs) 
-
-        return nfn 
-    return decorator
-
+def modelcontext(model): 
+    if model is None:
+       return Model.get_context() 
+    return model
 
 class Model(Context):
     """
@@ -142,20 +116,22 @@ def TransformedVar(model, name, dist, trans):
     def AddPotential(model, potential):
         model.factors.append(potential)
 
-withmodel = withcontext(Model, 'model')
 
-@withmodel
-def Point(model, *args,**kwargs):
+def Point(*args,**kwargs):
     """ 
     Build a point. Uses same args as dict() does. 
     Filters out variables not in the model. All keys are strings.  
 
     Parameters
     ----------
-        model : Model (in context) 
         *args, **kwargs 
             arguments to build a dict
     """
+    if 'model' in kwargs :
+        model = kwargs['model']
+        del kwargs['model']
+    else: 
+        model = Model.get_context()
 
     d = dict(*args, **kwargs)
     varnames = map(str, model.vars)

pymc/sample.py

@@ -8,8 +8,7 @@
 
 __all__ = ['sample', 'psample']
 
-@withmodel
-def sample(model, draws, step, start = None, trace = None, track_progress = True): 
+def sample(draws, step, start = None, trace = None, track_progress = True, model = None): 
     """
     Draw a number of samples using the given step method. 
     Multiple step methods supported via compound step method 
@@ -36,6 +35,7 @@ def sample(model, draws, step, start = None, trace = None, track_progress = True
     >>> an example
         
     """
+    model = modelcontext(model)
     draws = int(draws)
     if start is None: 
         start = trace[-1]
@@ -66,11 +66,12 @@ def argsample(args):
     """ defined at top level so it can be pickled"""
     return sample(*args)
 
-@withmodel
-def psample(model, draws, step, start, mtrace = None, threads = None, track = None):
+def psample(draws, step, start, mtrace = None, threads = None, track = None, model = None):
     """draw a number of samples using the given step method. Multiple step methods supported via compound step method
     returns the amount of time taken"""
 
+    model = modelcontext(model)
+
     if not threads:
         threads = max(mp.cpu_count() - 2, 1)
 

pymc/step_methods/gibbs.py

@@ -19,8 +19,8 @@ class ElemwiseCategoricalStep(ArrayStep):
     
     """
     #TODO: It would be great to come up with a way to make ElemwiseCategoricalStep  more general (handling more complex elementwise variables)
-    @withmodel
-    def __init__(self, model, var, values):
+    def __init__(self, var, values, model = None):
+        model = modelcontext(model)
         self.sh = ones(var.dshape, var.dtype)
         self.values = values
         self.var = var

pymc/step_methods/hmc.py

@@ -20,12 +20,10 @@ def unif(step_size, elow = .85, ehigh = 1.15):
 
 
 class HamiltonianMC(ArrayStep):
-    @withmodel
-    def __init__(self, model, vars, C, step_scale = .25, path_length = 2., is_cov = False, step_rand = unif, state = None):
+    def __init__(self, vars, C, step_scale = .25, path_length = 2., is_cov = False, step_rand = unif, state = None, model = None):
         """
         Parameters
         ----------
-            model : Model
             vars : list of theano variables
             C : array_like, ndim = {1,2}
                 Scaling for momentum distribution. 1d arrays interpreted matrix diagonal.
@@ -37,7 +35,11 @@ def __init__(self, model, vars, C, step_scale = .25, path_length = 2., is_cov =
                 Treat C as a covariance matrix/vector if True, else treat it as a precision matrix/vector
             step_rand : function float -> float, default=unif 
                 A function which takes the step size and returns an new one used to randomize the step size at each iteration. 
+            state 
+                State object
+            model : Model
         """
+        model = modelcontext(model)
         n = C.shape[0]
 
         self.step_size = step_scale / n**(1/4.)

pymc/step_methods/metropolis.py

@@ -14,8 +14,8 @@
 
 # TODO Implement tuning for Metropolis step
 class Metropolis(ArrayStep):
-    @withmodel
-    def __init__(self, model, vars, C, scaling=.25, is_cov = False):
+    def __init__(self, vars, C, scaling=.25, is_cov = False, model = None):
+        model = modelcontext(model)
 
         self.potential = quad_potential(C, not is_cov)
         self.scaling = scaling

pymc/tuning/scaling.py

@@ -9,8 +9,7 @@
 
 __all__ = ['approx_hess', 'find_hessian', 'trace_cov']
 
-@withmodel
-def approx_hess(model, point, vars=None):
+def approx_hess(point, vars=None, model = None):
     """
     Returns an approximation of the Hessian at the current chain location.
     
@@ -21,6 +20,7 @@ def approx_hess(model, point, vars=None):
     vars : list 
         Variables for which Hessian is to be calculated.
     """
+    model = modelcontext(model)
     if vars is None :
         vars = model.cont_vars
 
@@ -40,8 +40,7 @@ def grad_logp(point):
     '''
     return -nd.Jacobian(grad_logp)(bij.map(point))
 
-@withmodel
-def find_hessian(model, point, vars = None): 
+def find_hessian(point, vars = None, model = None): 
     """
     Returns Hessian of logp at the point passed.
     
@@ -52,6 +51,7 @@ def find_hessian(model, point, vars = None):
     vars : list 
         Variables for which Hessian is to be calculated.
     """
+    model = modelcontext(model)
     H = model.d2logpc(vars)
     return H(Point(model, point))
 

pymc/tuning/starting.py

@@ -11,26 +11,26 @@
 __all__ = ['find_MAP', 'scipyminimize']
 
 
-@withmodel
-def find_MAP(model, start = None, vars=None, fmin = optimize.fmin_bfgs, return_raw = False, disp = False, *args, **kwargs):
+def find_MAP(start = None, vars=None, fmin = optimize.fmin_bfgs, return_raw = False, disp = False, model = None, model = None, *args, **kwargs):
     """
     Sets state to the local maximum a posteriori point given a model.
     Current default of fmin_Hessian does not deal well with optimizing close 
     to sharp edges, especially if they are the minimum.
     
     Parameters
     ----------
-    model : Model (optional if in `with` context)
     start : dict of parameter values (Defaults to model.test_point)
     vars : list 
         List of variables to set to MAP point (Defaults to all continuous).
     fmin : function
         Optimization algorithm (Defaults to `scipy.optimize.fmin_l_bfgs_b`).
     return_raw : Bool
         Whether to return extra value returned by fmin (Defaults to False)
+    model : Model (optional if in `with` context)
     *args, **kwargs 
         Extra args passed to fmin
     """
+    model = modelcontext(model)
     if start is None:
         start = model.test_point
 

tests/models.py

@@ -5,7 +5,7 @@
 def simple_init(): 
     start, model, moments = simple_model()
 
-    step = Metropolis(model, model.vars, np.diag([1.]))
+    step = Metropolis(model.vars, np.diag([1.]), model = model)
     return model, start, step, moments
 
 

tests/test_step.py

@@ -11,8 +11,8 @@ def check_stat(name, trace, var, stat, value, bound):
 def test_step_continuous():
     start, model, (mu, C) = mv_simple()
 
-    hmc = pm.HamiltonianMC(model, model.vars,  C, is_cov = True)
-    mh = pm.Metropolis(model, model.vars , C, is_cov = True, scaling = 2)
+    hmc = pm.HamiltonianMC(model.vars,  C, is_cov = True, model = model)
+    mh = pm.Metropolis(model.vars , C, is_cov = True, scaling = 2, model = model)
     compound = pm.CompoundStep([hmc, mh])
 
     steps = [mh, hmc, compound]
@@ -23,10 +23,10 @@ def test_step_continuous():
 
     for st in steps:
         np.random.seed(1)
-        h = sample(model, 8000, st, start)
+        h = sample(8000, st, start, model = model)
         for (var, stat, val, bound) in check:
             np.random.seed(1)
-            h = sample(model, 8000, st, start)
+            h = sample(8000, st, start, model = model)
 
             yield check_stat,repr(st), h, var, stat, val, bound  
 

tests/test_trace.py

@@ -13,7 +13,7 @@
 def check_trace(model, trace, n, step, start):
     #try using a trace object a few times
     for i in range(2):
-        trace = sample(model, n, step, start, trace)
+        trace = sample(n, step, start, trace, model = model)
 
         for (var, val) in start.iteritems():
 
@@ -45,7 +45,7 @@ def test_multitrace():
 def check_multi_trace(model, trace, n, step, start):
 
     for i in range(2):
-        trace = psample(model, n, step, start, trace)
+        trace = psample(n, step, start, trace, model = model)
 
 
         for (var, val) in start.iteritems():