Allow linear model formula to extract variables from calling scope.

Author: adamboche

pymc3/glm/linear.py

@@ -84,9 +84,10 @@ def __init__(self, x, y, intercept=True, labels=None,
 
     @classmethod
     def from_formula(cls, formula, data, priors=None, vars=None,
-                     name='', model=None, offset=0.):
+                     name='', model=None, offset=0., eval_env=0):
         import patsy
-        y, x = patsy.dmatrices(formula, data)
+        eval_env = patsy.EvalEnvironment.capture(eval_env, reference=1)
+        y, x = patsy.dmatrices(formula, data, eval_env=eval_env)
         labels = x.design_info.column_names
         return cls(np.asarray(x), np.asarray(y)[:, -1], intercept=False,
                    labels=labels, priors=priors, vars=vars, name=name,
@@ -140,9 +141,10 @@ def __init__(self, x, y, intercept=True, labels=None,
     @classmethod
     def from_formula(cls, formula, data, priors=None,
                      vars=None, family='normal', name='',
-                     model=None, offset=0.):
+                     model=None, offset=0., eval_env=0):
         import patsy
-        y, x = patsy.dmatrices(formula, data)
+        eval_env = patsy.EvalEnvironment.capture(eval_env, reference=1)
+        y, x = patsy.dmatrices(formula, data, eval_env=eval_env)
         labels = x.design_info.column_names
         return cls(np.asarray(x), np.asarray(y)[:, -1], intercept=False,
                    labels=labels, priors=priors, vars=vars, family=family,

pymc3/tests/test_glm.py

@@ -2,6 +2,7 @@
 from numpy.testing import assert_equal
 
 from .helpers import SeededTest
+import pymc3
 from pymc3 import Model, Uniform, Normal, find_MAP, Slice, sample
 from pymc3 import families, GLM, LinearComponent
 import pandas as pd
@@ -117,3 +118,15 @@ def test_boolean_y(self):
             )
         )
         assert_equal(model.y.observations, model_bool.y.observations)
+
+    def test_glm_formula_from_calling_scope(self):
+        """Formula can extract variables from the calling scope."""
+        z = pd.Series([10, 20, 30])
+        df = pd.DataFrame({"y": [0, 1, 0], "x": [1.0, 2.0, 3.0]})
+        GLM.from_formula("y ~ x + z", df, family=pymc3.glm.families.Binomial())
+
+    def test_linear_component_formula_from_calling_scope(self):
+        """Formula can extract variables from the calling scope."""
+        z = pd.Series([10, 20, 30])
+        df = pd.DataFrame({"y": [0, 1, 0], "x": [1.0, 2.0, 3.0]})
+        LinearComponent.from_formula("y ~ x + z", df)