[WIP] Introduced pm.TidyData class

pymc3/data.py

@@ -25,6 +25,7 @@
 import pandas as pd
 import theano
 import theano.tensor as tt
+import xarray
 
 from theano.graph.basic import Apply
 
@@ -36,6 +37,7 @@
     "Minibatch",
     "align_minibatches",
     "Data",
+    "TidyData",
 ]
 BASE_URL = "https://raw.githubusercontent.com/pymc-devs/pymc-examples/main/examples/data/{filename}"
 
@@ -594,3 +596,59 @@ def set_coords(model, value, dims=None):
                     coords[dim] = pd.RangeIndex(size, name=dim)
 
         return coords
+
+
+class TidyData:
+    def __init__(self, data, import_dims=None, model=None):
+        self.data = data
+        self._shared_vars = {}
+        self._category_cols = {}
+        self._category_col_keys = {}
+
+        if import_dims is not None:
+            model = pm.model.modelcontext(model)
+            model_coords = self._dims_to_dict(import_dims)
+            model.add_coords(model_coords)
+
+    def _dims_to_dict(self, dims):
+        model_coords = {}
+        for dim in dims:
+            self._validate_idx(dim)
+            if dim in self.data.coords:
+                coord = self.data.coords[dim]
+            elif dim in self.data.data_vars:
+                coord = self.data.data_vars[dim]
+            model_coords[dim] = coord
+        return model_coords
+
+    def __getitem__(self, key):
+        self._validate_idx(key)
+        if key in self._shared_vars:
+            return self._shared_vars[key]
+
+        shared_var = theano.shared(self.data[key].values)
+        self._shared_vars[key] = shared_var
+        return shared_var
+
+    def get_indexed(self, col, keys=False):
+        self._validate_idx(col)
+
+        if col not in self._category_cols:
+            col_data = self.data[col]
+            keys = {}
+
+            for idx, group in enumerate(set(col_data.values)):
+                keys[group] = idx
+                col_data = xarray.where(col_data == group, idx, col_data)
+
+            shared_col_data = col_data.values.tolist()
+            self._category_col_keys[col] = keys
+            self._category_cols[col] = shared_col_data
+
+        if keys is True:
+            return self._category_cols[col], self._category_col_keys[col]
+        return self._category_cols[col]
+
+    def _validate_idx(self, idx):
+        if idx not in self.data.coords and idx not in self.data.data_vars:
+            raise KeyError("Unknown column %s" % idx)