Text backend does not support tallying a subset of variables (raises KeyError) #2560 (#3492)

Author: Bruschkov

pymc3/backends/text.py

@@ -17,11 +17,13 @@
 """
 from glob import glob
 import os
+import re
 import pandas as pd
 
 from ..backends import base, ndarray
 from . import tracetab as ttab
 from ..theanof import floatX
+from ..model import modelcontext
 
 
 class Text(base.BaseTrace):
@@ -112,7 +114,6 @@ def _load_df(self):
                 if "float" in str(dtype):
                     self.df[key] = floatX(self.df[key])
 
-
     def __len__(self):
         if self.filename is None:
             return 0
@@ -178,13 +179,25 @@ def load(name, model=None):
     straces = []
     for f in files:
         chain = int(os.path.splitext(f)[0].rsplit('-', 1)[1])
-        strace = Text(name, model=model)
+        model_vars_in_chain = _parse_chain_vars(f, model)
+        strace = Text(name, model=model, vars=model_vars_in_chain)
         strace.chain = chain
         strace.filename = f
         straces.append(strace)
     return base.MultiTrace(straces)
 
 
+def _parse_chain_vars(filepath, model):
+    with open(filepath) as f:
+        header = f.readline().split("\n", 1)[0]
+    shape_pattern = re.compile(r"__\d+_\d+")
+    chain_vars = [shape_pattern.split(v)[0] for v in header.split(",")]
+    chain_vars = list(set(chain_vars))
+    m = modelcontext(model)
+    model_vars_in_chain = [v for v in m.unobserved_RVs if v.name in chain_vars]
+    return model_vars_in_chain
+
+
 def dump(name, trace, chains=None):
     """Store values from NDArray trace as CSV files.
 

pymc3/tests/backend_fixtures.py

@@ -124,13 +124,20 @@ class ModelBackendSampledTestCase:
 
     Children may define
     - sampler_vars
+    - write_partial_chain
     """
     @classmethod
     def setup_class(cls):
         cls.test_point, cls.model, _ = models.beta_bernoulli(cls.shape)
+
+        if hasattr(cls, 'write_partial_chain') and cls.write_partial_chain is True:
+            cls.chain_vars = cls.model.unobserved_RVs[1:]
+        else:
+            cls.chain_vars = cls.model.unobserved_RVs
+
         with cls.model:
-            strace0 = cls.backend(cls.name)
-            strace1 = cls.backend(cls.name)
+            strace0 = cls.backend(cls.name, vars=cls.chain_vars)
+            strace1 = cls.backend(cls.name, vars=cls.chain_vars)
 
         if not hasattr(cls, 'sampler_vars'):
             cls.sampler_vars = None
@@ -459,7 +466,7 @@ def test_values(self):
         trace = self.mtrace
         dumped = self.dumped
         for chain in trace.chains:
-            for varname in self.test_point.keys():
+            for varname in self.chain_vars:
                 data = trace.get_values(varname, chains=[chain])
                 dumped_data = dumped.get_values(varname, chains=[chain])
                 npt.assert_equal(data, dumped_data)

pymc3/tests/test_text_backend.py

@@ -62,6 +62,14 @@ class TestTextDumpLoad(bf.DumpLoadTestCase):
     shape = (2, 3)
 
 
+class TestTextDumpLoadWithPartialChain(bf.DumpLoadTestCase):
+    backend = text.Text
+    load_func = staticmethod(text.load)
+    name = 'text-db'
+    shape = (2, 3)
+    write_partial_chain = True
+
+
 @pytest.mark.xfail(condition=(theano.config.floatX == "float32"), reason="Fails on float32")
 class TestTextDumpFunction(bf.BackendEqualityTestCase):
     backend0 = backend1 = ndarray.NDArray