numexpr dims

ActivitySim · Oct 28, 2024 · ea8174f · ea8174f
1 parent fc5df6c
commit ea8174f
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/sharrow/relationships.py b/sharrow/relationships.py
@@ -964,6 +964,13 @@ def get_expr(
                 else:
                     if dtype is not None:
                         result = result.astype(dtype)
+                    # numexpr doesn't carry over the dimension names or coords
+                    result = result.rename(
+                        {result.dims[i]: self.root_dims[i] for i in range(result.ndim)}
+                    )
+                    if with_coords:
+                        result = result.assign_coords(self.root_dataset.coords)
+
             elif engine == "pandas-numexpr":
                 from xarray import DataArray