Change use of Model.coords -> Model.coords_typed throughout pymc

Author: covertg

pymc/backends/arviz.py

@@ -39,7 +39,7 @@
 
 from pymc.model import Model, modelcontext
 from pymc.pytensorf import extract_obs_data
-from pymc.util import get_default_varnames
+from pymc.util import _as_coord_vals, get_default_varnames
 
 if TYPE_CHECKING:
     from pymc.backends.base import MultiTrace  # pylint: disable=invalid-name
@@ -216,15 +216,15 @@ def __init__(
                 " one of trace, prior, posterior_predictive or predictions."
             )
 
-        # Make coord types more rigid
-        untyped_coords: Dict[str, Optional[Sequence[Any]]] = {**self.model.coords}
+        coords_typed = self.model.coords_typed
         if coords:
-            untyped_coords.update(coords)
-        self.coords = {
-            cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
-            for cname, cvals in untyped_coords.items()
-            if cvals is not None
+            coords_typed.update(_as_coord_vals(coords))
+        coords_typed = {
+            cname: cvals_typed
+            for cname, cvals_typed in coords_typed.items()
+            if cvals_typed is not None
         }
+        self.coords = coords_typed
 
         self.dims = {} if dims is None else dims
         model_dims = {k: list(v) for k, v in self.model.named_vars_to_dims.items()}

pymc/backends/mcbackend.py

@@ -230,9 +230,9 @@ def make_runmeta_and_point_fn(
             sample_stats.append(svar)
 
     coordinates = [
-        mcb.Coordinate(dname, mcb.npproto.utils.ndarray_from_numpy(np.array(cvals)))
-        for dname, cvals in model.coords.items()
-        if cvals is not None
+        mcb.Coordinate(dname, mcb.npproto.utils.ndarray_from_numpy(cvals_typed))
+        for dname, cvals_typed in model.coords_typed.items()
+        if cvals_typed is not None
     ]
     meta = mcb.RunMeta(
         rid=hagelkorn.random(),

pymc/distributions/bound.py

@@ -195,8 +195,8 @@ def __new__(
 
         if dims is not None:
             model = modelcontext(None)
-            if dims in model.coords:
-                dim_obj = np.asarray(model.coords[dims])
+            if dims in model.coords_typed:
+                dim_obj = model.coords_typed[dims]
                 size = dim_obj.shape
             else:
                 raise ValueError("Given dims do not exist in model coordinates.")

pymc/model.py

@@ -992,7 +992,7 @@ def coords(self) -> Dict[str, Union[Tuple, None]]:
             cname: tuple(cvals) if cvals is not None else None
             for cname, cvals in self._coords.items()
         }
-    
+
     @property
     def coords_typed(self) -> Dict[str, Union[np.ndarray, None]]:
         """Coordinate values (numpy array) for model dimensions."""
@@ -1014,20 +1014,20 @@ def shape_from_dims(self, dims):
         if len(set(dims)) != len(dims):
             raise ValueError("Can not contain the same dimension name twice.")
         for dim in dims:
-            if dim not in self.coords:
+            if dim not in self.coords_typed:
                 raise ValueError(
                     f"Unknown dimension name '{dim}'. All dimension "
                     "names must be specified in the `coords` "
                     "argument of the model or through a pm.Data "
                     "variable."
                 )
-            shape.extend(np.shape(self.coords[dim]))
+            shape.extend(self.coords_typed[dim].shape)
         return tuple(shape)
 
     def add_coord(
         self,
         name: str,
-        values: Optional[Sequence] = None,
+        values: Optional[Union[Sequence, np.ndarray]] = None,
         mutable: bool = False,
         *,
         length: Optional[Union[int, Variable]] = None,
@@ -1061,8 +1061,8 @@ def add_coord(
         if values is not None:
             # Conversion to numpy array to ensure coord vals are 1-dim
             values = _as_coord_vals(values)
-        if name in self.coords:
-            if not np.array_equal(values, self.coords[name]):
+        if name in self.coords_typed:
+            if not np.array_equal(_as_coord_vals(values), self.coords_typed[name]):
                 raise ValueError(f"Duplicate and incompatible coordinate: {name}.")
         if length is not None and not isinstance(length, (int, Variable)):
             raise ValueError(
@@ -1080,7 +1080,7 @@ def add_coord(
 
     def add_coords(
         self,
-        coords: Dict[str, Optional[Sequence]],
+        coords: Dict[str, Optional[Union[Sequence, np.ndarray]]],
         *,
         lengths: Optional[Dict[str, Optional[Union[int, Variable]]]] = None,
     ):
@@ -1092,7 +1092,9 @@ def add_coords(
         for name, values in coords.items():
             self.add_coord(name, values, length=lengths.get(name, None))
 
-    def set_dim(self, name: str, new_length: int, coord_values: Optional[Sequence] = None):
+    def set_dim(
+        self, name: str, new_length: int, coord_values: Optional[Union[Sequence, np.ndarray]] = None
+    ):
         """Update a mutable dimension.
 
         Parameters
@@ -1106,7 +1108,7 @@ def set_dim(self, name: str, new_length: int, coord_values: Optional[Sequence] =
         """
         if not isinstance(self.dim_lengths[name], ScalarSharedVariable):
             raise ValueError(f"The dimension '{name}' is immutable.")
-        if coord_values is None and self.coords.get(name, None) is not None:
+        if coord_values is None and self.coords_typed.get(name, None) is not None:
             raise ValueError(
                 f"'{name}' has coord values. Pass `set_dim(..., coord_values=...)` to update them."
             )
@@ -1162,7 +1164,7 @@ def set_data(
         self,
         name: str,
         values: Dict[str, Optional[Sequence]],
-        coords: Optional[Dict[str, Sequence]] = None,
+        coords: Optional[Dict[str, Union[Sequence, np.ndarray]]] = None,
     ):
         """Changes the values of a data variable in the model.
 
@@ -1192,7 +1194,7 @@ def set_data(
             values = np.array(values)
         values = convert_observed_data(values)
         dims = self.named_vars_to_dims.get(name, None) or ()
-        coords = coords or {}
+        coords_untyped = coords or {}
 
         if values.ndim != shared_object.ndim:
             raise ValueError(
@@ -1203,19 +1205,21 @@ def set_data(
             length_tensor = self.dim_lengths[dname]
             old_length = length_tensor.eval()
             new_length = values.shape[d]
-            original_coords = self.coords.get(dname, None)
-            new_coords = coords.get(dname, None)
+            original_coord_vals = self.coords_typed.get(dname, None)
+            new_coord_vals = coords_untyped.get(dname, None)
+            if new_coord_vals is not None:
+                new_coord_vals = _as_coord_vals(new_coord_vals)
 
             length_changed = new_length != old_length
 
             # Reject resizing if we already know that it would create shape problems.
             # NOTE: If there are multiple pm.MutableData containers sharing this dim, but the user only
             #       changes the values for one of them, they will run into shape problems nonetheless.
             if length_changed:
-                if original_coords is not None:
-                    if new_coords is None:
+                if original_coord_vals is not None:
+                    if new_coord_vals is None:
                         raise ValueError(
-                            f"The '{name}' variable already had {len(original_coords)} coord values defined for "
+                            f"The '{name}' variable already had {len(original_coord_vals)} coord values defined for "
                             f"its {dname} dimension. With the new values this dimension changes to length "
                             f"{new_length}, so new coord values for the {dname} dimension are required."
                         )
@@ -1279,17 +1283,17 @@ def set_data(
                 if isinstance(length_tensor, ScalarSharedVariable):
                     # The dimension is mutable, but was defined without being linked
                     # to a shared variable. This is allowed, but a little less robust.
-                    self.set_dim(dname, new_length, coord_values=new_coords)
+                    self.set_dim(dname, new_length, coord_values=new_coord_vals)
 
-            if new_coords is not None:
+            if new_coord_vals is not None:
                 # Update the registered coord values (also if they were None)
-                if len(new_coords) != new_length:
+                if len(new_coord_vals) != new_length:
                     raise ShapeError(
                         f"Length of new coordinate values for dimension '{dname}' does not match the provided values.",
-                        actual=len(new_coords),
+                        actual=len(new_coord_vals),
                         expected=new_length,
                     )
-                self._coords[dname] = _as_coord_vals(new_coords)
+                self._coords[dname] = _as_coord_vals(new_coord_vals)
 
         shared_object.set_value(values)
 
@@ -1560,7 +1564,7 @@ def add_named_variable(self, var, dims: Optional[Tuple[Union[str, None], ...]] =
             if isinstance(dims, str):
                 dims = (dims,)
             for dim in dims:
-                if dim not in self.coords and dim is not None:
+                if dim not in self.coords_typed and dim is not None:
                     raise ValueError(f"Dimension {dim} is not specified in `coords`.")
             if any(var.name == dim for dim in dims if dim is not None):
                 raise ValueError(f"Variable `{var.name}` has the same name as its dimension label.")

pymc/sampling/forward.py

@@ -575,7 +575,7 @@ def sample_posterior_predictive(
 
     constant_coords = set()
     for dim, coord in trace_coords.items():
-        current_coord = model.coords.get(dim, None)
+        current_coord = model.coords_typed.get(dim, None)
         if (
             current_coord is not None
             and len(coord) == len(current_coord)

pymc/sampling/jax.py

@@ -384,9 +384,9 @@ def sample_blackjax_nuts(
     vars_to_sample = list(get_default_varnames(var_names, include_transformed=keep_untransformed))
 
     coords = {
-        cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
-        for cname, cvals in model.coords.items()
-        if cvals is not None
+        cname: cvals_typed
+        for cname, cvals_typed in model.coords_typed.items()
+        if cvals_typed is not None
     }
 
     dims = {
@@ -605,9 +605,9 @@ def sample_numpyro_nuts(
     vars_to_sample = list(get_default_varnames(var_names, include_transformed=keep_untransformed))
 
     coords = {
-        cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
-        for cname, cvals in model.coords.items()
-        if cvals is not None
+        cname: cvals_typed
+        for cname, cvals_typed in model.coords_typed.items()
+        if cvals_typed is not None
     }
 
     dims = {

pymc/stats/log_likelihood.py

@@ -13,8 +13,6 @@
 #   limitations under the License.
 from typing import Optional, Sequence, cast
 
-import numpy as np
-
 from arviz import InferenceData, dict_to_dataset
 from fastprogress import progress_bar
 
@@ -117,10 +115,7 @@ def compute_log_likelihood(
         loglike_trace,
         library=pymc,
         dims={dname: list(dvals) for dname, dvals in model.named_vars_to_dims.items()},
-        coords={
-            cname: np.array(cvals) if isinstance(cvals, tuple) else cvals
-            for cname, cvals in model.coords.items()
-        },
+        coords={cname: cvals_typed for cname, cvals_typed in model.coords_typed.items()},
         default_dims=list(sample_dims),
         skip_event_dims=True,
     )

pymc/util.py

@@ -27,6 +27,8 @@
 from pytensor.compile import SharedVariable
 from pytensor.graph.utils import ValidatingScratchpad
 
+from pymc.exceptions import ShapeError
+
 
 class _UnsetType:
     """Type for the `UNSET` object to make it look nice in `help(...)` outputs."""
@@ -511,8 +513,17 @@ def _add_future_warning_tag(var) -> None:
             new_tag.__dict__.setdefault(k, v)
         var.tag = new_tag
 
-def _as_coord_vals(values: Sequence) -> np.ndarray:
-    """Coerce a sequence coordinate values into a 1-dim array of values"""
+
+def _as_coord_vals(values: Union[Sequence, np.ndarray]) -> np.ndarray:
+    """Coerce a sequence of coordinate values into a 1-dim array"""
+    if isinstance(values, np.ndarray):
+        if values.ndim != 1:
+            raise ShapeError(
+                "Coordinate values passed as a numpy array must be 1-dimensional",
+                actual=values.ndim,
+                expected=1,
+            )
+        return values
     arr = np.array(values)
     if arr.ndim > 1:
         arr = np.empty(len(values), dtype="O")

pymc/variational/opvi.py

@@ -1135,7 +1135,7 @@ def var_to_data(self, shared: at.TensorVariable) -> xarray.Dataset:
         for name, s, shape, dtype in self.ordering.values():
             dims = self.model.named_vars_to_dims.get(name, None)
             if dims is not None:
-                coords = {d: np.array(self.model.coords[d]) for d in dims}
+                coords = {d: self.model.coords_typed[d] for d in dims}
             else:
                 coords = None
             values = shared_nda[s].reshape(shape).astype(dtype)

tests/test_model.py

@@ -733,7 +733,7 @@ def test_nested_model_coords():
             c = pm.HalfNormal("c", dims="dim3")
             d = pm.Normal("d", b, c, dims="dim2")
         e = pm.Normal("e", a[None] + d[:, None], dims=("dim2", "dim1"))
-    assert m1.coords is m2.coords
+    assert m1.coords == m2.coords
     assert m1.dim_lengths is m2.dim_lengths
     assert set(m2.named_vars_to_dims) < set(m1.named_vars_to_dims)