Skip to content

Fix NotImplementedError in xr.merge() for object-dtype geometry columns #586

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Copilot wants to merge 2 commits into
base: copilot/fix-recursion-error-spatial-aggregation
Choose a base branch
from
+45 −22
Draft

Fix NotImplementedError in xr.merge() for object-dtype geometry columns #586

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
67515f7
Initial plan
Copilot Feb 24, 2026
24346d4
Fix NotImplementedError in xr.merge() by adding merge_with_geometry h…
Copilot Feb 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 30 additions & 10 deletions fine/IOManagement/utilsIO.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,26 @@
import operator


def merge_with_geometry(ds1, ds2, **merge_kwargs):
"""Merge two xarray Datasets, handling object-dtype (geometry) variables separately
to avoid NotImplementedError in newer xarray versions."""
obj_vars = {}
for ds in [ds1, ds2]:
for var in ds.data_vars:
if ds[var].dtype == object:
obj_vars[var] = ds[var]

ds1_clean = ds1.drop_vars([v for v in obj_vars if v in ds1.data_vars])
ds2_clean = ds2.drop_vars([v for v in obj_vars if v in ds2.data_vars])

merged = xr.merge([ds1_clean, ds2_clean], **merge_kwargs)

for var, da in obj_vars.items():
merged[var] = da

return merged


def getFromDict(dataDict, mapList):
"""Get value from a dict by a list, which contains the dict keys.

Expand Down Expand Up @@ -316,8 +336,8 @@ def add_to_xarray(xr_ds, df_dict, variable_description):
)

try:
xr_ds[this_class][this_comp] = xr.merge(
[xr_ds[this_class][this_comp], this_ds_component]
xr_ds[this_class][this_comp] = merge_with_geometry(
xr_ds[this_class][this_comp], this_ds_component
)
except Exception:
pass
Expand Down Expand Up @@ -417,8 +437,8 @@ def addSeriesVariablesToXarray(xr_ds, component_dict, series_iteration_dict, loc
)

try:
xr_ds[this_class][this_comp] = xr.merge(
[xr_ds[this_class][this_comp], this_ds_component]
xr_ds[this_class][this_comp] = merge_with_geometry(
xr_ds[this_class][this_comp], this_ds_component
)
except Exception:
pass
Expand All @@ -441,8 +461,8 @@ def addSeriesVariablesToXarray(xr_ds, component_dict, series_iteration_dict, loc
)

try:
xr_ds[this_class][this_comp] = xr.merge(
[xr_ds[this_class][this_comp], this_ds_component]
xr_ds[this_class][this_comp] = merge_with_geometry(
xr_ds[this_class][this_comp], this_ds_component
)
except Exception:
pass
Expand All @@ -465,8 +485,8 @@ def addSeriesVariablesToXarray(xr_ds, component_dict, series_iteration_dict, loc
)

try:
xr_ds[this_class][this_comp] = xr.merge(
[xr_ds[this_class][this_comp], this_ds_component]
xr_ds[this_class][this_comp] = merge_with_geometry(
xr_ds[this_class][this_comp], this_ds_component
)
except Exception:
pass
Expand Down Expand Up @@ -528,8 +548,8 @@ def addConstantsToXarray(
)

try:
xr_ds[this_class][this_comp] = xr.merge(
[xr_ds[this_class][this_comp], this_ds_component]
xr_ds[this_class][this_comp] = merge_with_geometry(
xr_ds[this_class][this_comp], this_ds_component
)
except Exception:
pass
Expand Down
27 changes: 15 additions & 12 deletions fine/IOManagement/xarrayIO.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

from fine import utils
from fine.IOManagement import dictIO, utilsIO
from fine.IOManagement.utilsIO import merge_with_geometry


def convertOptimizationInputToDatasets(esM, useProcessedValues=False):
Expand Down Expand Up @@ -139,8 +140,9 @@ def convertOptimizationOutputToDatasets(esM, optSumOutputLevel=0):
xr_da.attrs[variable] = unit

# merge to overall xr_ds
xr_dss[ip][name][component] = xr.merge(
[xr_dss[ip][name][component], xr_da],
xr_dss[ip][name][component] = merge_with_geometry(
xr_dss[ip][name][component],
xr_da,
combine_attrs="drop_conflicts",
)
elif esM.componentModelingDict[name].dimension == "2dim":
Expand Down Expand Up @@ -171,8 +173,9 @@ def convertOptimizationOutputToDatasets(esM, optSumOutputLevel=0):
unit = variables_unit[variable]
xr_da.attrs[variable] = unit
# merge to overall xr_ds
xr_dss[ip][name][component] = xr.merge(
[xr_dss[ip][name][component], xr_da],
xr_dss[ip][name][component] = merge_with_geometry(
xr_dss[ip][name][component],
xr_da,
combine_attrs="drop_conflicts",
)

Expand Down Expand Up @@ -203,8 +206,8 @@ def convertOptimizationOutputToDatasets(esM, optSumOutputLevel=0):
df.name = variable
df.index.rename(["time", "location"], inplace=True)
xr_da = df.to_xarray()
xr_dss[ip][name][component] = xr.merge(
[xr_dss[ip][name][component], xr_da]
xr_dss[ip][name][component] = merge_with_geometry(
xr_dss[ip][name][component], xr_da
)
# Two dimensional time dependent data
if dataTD2dim:
Expand All @@ -223,8 +226,8 @@ def convertOptimizationOutputToDatasets(esM, optSumOutputLevel=0):
)
df.index = df.index.reorder_levels([2, 0, 1])
xr_da = df.to_xarray()
xr_dss[ip][name][component] = xr.merge(
[xr_dss[ip][name][component], xr_da]
xr_dss[ip][name][component] = merge_with_geometry(
xr_dss[ip][name][component], xr_da
)
# Time independent data
if dataTI:
Expand All @@ -241,8 +244,8 @@ def convertOptimizationOutputToDatasets(esM, optSumOutputLevel=0):
df.name = variable
df.index.rename("location", inplace=True)
xr_da = df.to_xarray()
xr_dss[ip][name][component] = xr.merge(
[xr_dss[ip][name][component], xr_da]
xr_dss[ip][name][component] = merge_with_geometry(
xr_dss[ip][name][component], xr_da
)
# Two dimensional
elif esM.componentModelingDict[name].dimension == "2dim":
Expand All @@ -257,8 +260,8 @@ def convertOptimizationOutputToDatasets(esM, optSumOutputLevel=0):
df.name = variable
df.index.rename(["locationIn", "locationOut"], inplace=True)
xr_da = df.to_xarray()
xr_dss[ip][name][component] = xr.merge(
[xr_dss[ip][name][component], xr_da]
xr_dss[ip][name][component] = merge_with_geometry(
xr_dss[ip][name][component], xr_da
)

for name in esM.componentModelingDict.keys():
Expand Down