Ensure base Variable class is used when assigning

Author: jsignell

xarray/core/10908.py

@@ -0,0 +1,56 @@
+import dask.array
+import numpy as np
+
+import xarray
+
+x = np.arange(56)
+y = np.arange(100, 119)
+t = np.arange(2)
+z = np.arange(9)
+
+# The first array
+projected = xarray.DataArray(
+    dask.array.full((x.size, y.size), np.nan, chunks=(10, 10)),
+    coords={"x": x, "y": y},
+    name="projected",
+)
+
+# The array whose values will be used
+original = xarray.DataArray(
+    np.full((t.size, z.size), 12.345), coords={"t": t, "z": z}, dims=["t", "z"]
+).chunk({"t": 10})
+
+t2x = np.linspace(0, x.size, t.size)
+z2y = np.linspace(y.min(), y.max(), z.size)
+tmp = t2x[..., None] @ z2y[None, ...]
+
+# The array with the coordinate correspondence
+pp = np.full((t.size, z.size, 2), np.nan)
+pp[..., 0] = tmp / z2y.max()
+pp[..., 1] = tmp / t2x.max()
+
+pixel_positions = xarray.DataArray(
+    pp, coords={"t": t, "z": z, "nav": ["x", "y"]}, dims=["t", "z", "nav"]
+).chunk({"t": 10})
+
+
+def match_coordinates(projected, pixel_positions):
+    coordinates = projected.sel(
+        x=pixel_positions[..., 0], y=pixel_positions[..., 1], method="nearest"
+    ).coords
+    return {"x": coordinates["x"], "y": coordinates["y"]}
+
+
+def place_pixel_on_grid(projected, pixel_positions, original):
+    coordinates = match_coordinates(projected, pixel_positions)
+    print(f"{coordinates=}")
+    print(f"{projected.shape} -- {original.shape}")
+    print(f"pixel_positions coords {pixel_positions.coords}")
+    projected_chunk = projected.copy()
+    projected_chunk.loc[{"x": coordinates["x"], "y": coordinates["y"]}] = original
+    return projected_chunk
+
+
+xarray.map_blocks(
+    place_pixel_on_grid, projected, args=[pixel_positions, original], template=projected
+).compute()

xarray/core/coordinates.py

@@ -1265,7 +1265,7 @@ def create_coords_with_default_indexes(
             variables.update(idx_vars)
             all_variables.update(idx_vars)
         else:
-            variables[name] = variable
+            variables[name] = variable.to_base_variable()
 
     new_coords = Coordinates._construct_direct(coords=variables, indexes=indexes)
 

xarray/structure/merge.py

@@ -22,7 +22,12 @@
     emit_user_level_warning,
     equivalent,
 )
-from xarray.core.variable import Variable, as_variable, calculate_dimensions
+from xarray.core.variable import (
+    IndexVariable,
+    Variable,
+    as_variable,
+    calculate_dimensions,
+)
 from xarray.structure.alignment import deep_align
 from xarray.util.deprecation_helpers import (
     _COMPAT_DEFAULT,
@@ -1218,7 +1223,11 @@ def dataset_update_method(dataset: Dataset, other: CoercibleMapping) -> _MergeRe
                     if c not in value.dims and c in dataset.coords
                 ]
                 if coord_names:
-                    other[key] = value.drop_vars(coord_names)
+                    value = value.drop_vars(coord_names)
+                if isinstance(value.variable, IndexVariable):
+                    variable = value.variable.to_base_variable()
+                    value = value._replace(variable=variable)
+                other[key] = value
 
     return merge_core(
         [dataset, other],