Performance improvement: Convert unnecessary (?) deep copies to shallow copies (#358)

Author: lukelbd

cf_xarray/accessor.py

@@ -1797,7 +1797,7 @@ def guess_coord_axis(self, verbose: bool = False) -> DataArray | Dataset:
         DataArray or Dataset
             with appropriate attributes added
         """
-        obj = self._obj.copy(deep=True)
+        obj = self._obj.copy(deep=False)
         for var in obj.coords.variables:
             var_is_coord = any(var in val for val in obj.cf.coordinates.values())
             if not var_is_coord and obj[var].ndim == 1 and _is_datetime_like(obj[var]):
@@ -1937,7 +1937,7 @@ def add_canonical_attributes(
         info, table, aliases = parse_cf_standard_name_table(source)
 
         # Loop over standard names
-        ds = self._maybe_to_dataset().copy()
+        ds = self._maybe_to_dataset().copy(deep=False)
         attrs_to_print: dict = {}
         for std_name, var_names in ds.cf.standard_names.items():
 
@@ -2209,7 +2209,7 @@ def add_bounds(self, keys: str | Iterable[str], *, dim=None):
                 apply_mapper(_get_all, self._obj, key, error=False, default=[key])
             )
 
-        obj = self._maybe_to_dataset(self._obj.copy(deep=True))
+        obj = self._maybe_to_dataset(self._obj.copy(deep=False))
 
         bad_vars: set[str] = variables - set(obj.variables)
         if bad_vars:
@@ -2286,7 +2286,7 @@ def bounds_to_vertices(
         else:
             coords = keys
 
-        obj = self._maybe_to_dataset(self._obj.copy(deep=True))
+        obj = self._maybe_to_dataset(self._obj.copy(deep=False))
 
         for coord in coords:
             try:

cf_xarray/datasets.py

@@ -11,7 +11,7 @@
     * 110e3
 )
 
-ds_no_attrs = airds.copy(deep=True)
+ds_no_attrs = airds.copy(deep=False)
 for _variable in ds_no_attrs.variables:
     ds_no_attrs[_variable].attrs = {}
 

cf_xarray/tests/test_accessor.py

@@ -230,7 +230,7 @@ def test_coordinates_quantified():
 
 
 def test_cell_measures():
-    ds = airds.copy(deep=True)
+    ds = airds.copy(deep=False)
     ds["foo"] = xr.DataArray(ds["cell_area"], attrs=dict(standard_name="foo_std_name"))
     ds["air"].attrs["cell_measures"] += " foo_measure: foo"
     assert ("foo_std_name" in ds.cf["air_temperature"].cf) and ("foo_measure" in ds.cf)
@@ -307,7 +307,7 @@ def test_getitem_standard_name():
     expected = airds["lat"]
     assert_identical(actual, expected)
 
-    ds = airds.copy(deep=True)
+    ds = airds.copy(deep=False)
     ds["air2"] = ds.air
     with pytest.raises(KeyError):
         ds.cf["air_temperature"]
@@ -340,19 +340,18 @@ def test_getitem_ancillary_variables():
 
 
 def test_rename_like():
-    original = popds.copy(deep=True)
-
     # it'll match for axis: X (lon, nlon) and coordinate="longitude" (lon, TLONG)
     # so delete the axis attributes
-    newair = airds.copy(deep=True)
+    original = popds
+    newair = airds.copy(deep=False)
     del newair.lon.attrs["axis"]
     del newair.lat.attrs["axis"]
 
     renamed = popds.cf["TEMP"].cf.rename_like(newair)
+    assert original.TEMP.attrs["coordinates"] == "TLONG TLAT"
     for k in ["TLONG", "TLAT"]:
-        assert k not in renamed.coords
         assert k in original.coords
-    assert original.TEMP.attrs["coordinates"] == "TLONG TLAT"
+        assert k not in renamed.coords
 
     assert "lon" in renamed.coords
     assert "lat" in renamed.coords
@@ -479,7 +478,7 @@ def test_pos_args_methods():
 
 def test_preserve_unused_keys():
 
-    ds = airds.copy(deep=True)
+    ds = airds.copy(deep=False)
     ds.time.attrs.clear()
     actual = ds.cf.sel(X=260, Y=40, time=airds.time[:2], method="nearest")
     expected = ds.sel(lon=260, lat=40, time=airds.time[:2], method="nearest")
@@ -528,7 +527,7 @@ def test_args_methods(obj):
 
 def test_dataarray_getitem():
 
-    air = airds.air.copy()
+    air = airds.air.copy(deep=False)
     air.name = None
 
     assert_identical(air.cf["longitude"], air["lon"])
@@ -543,7 +542,7 @@ def test_dataarray_getitem():
 
 def test_dataarray_plot():
 
-    obj = airds.air.copy(deep=True)
+    obj = airds.air.copy(deep=False)
 
     rv = obj.isel(time=1).transpose("lon", "lat").cf.plot()
     assert isinstance(rv, mpl.collections.QuadMesh)
@@ -636,14 +635,14 @@ def test_getitem(obj, key, expected_key):
 def test_getitem_errors(obj):
     with pytest.raises(KeyError):
         obj.cf["XX"]
-    obj2 = obj.copy(deep=True)
+    obj2 = obj.copy(deep=False)
     obj2.lon.attrs = {}
     with pytest.raises(KeyError):
         obj2.cf["X"]
 
 
 def test_bad_cell_measures_attribute():
-    air2 = airds.copy(deep=True)
+    air2 = airds.copy(deep=False)
     air2.air.attrs["cell_measures"] = "--OPT"
     with pytest.warns(UserWarning):
         air2.cf["air"]
@@ -736,20 +735,21 @@ def test_plot_xincrease_yincrease():
 
 @pytest.mark.parametrize("dims", ["time2", "lat", "time", ["lat", "lon"]])
 def test_add_bounds(dims):
-    obj = airds.copy(deep=True)
+    ds = airds
+    original = ds.copy(deep=True)
 
     expected = {}
     expected["lat"] = xr.concat(
         [
-            obj.lat.copy(data=np.arange(76.25, 16.0, -2.5)),
-            obj.lat.copy(data=np.arange(73.75, 13.6, -2.5)),
+            ds.lat.copy(data=np.arange(76.25, 16.0, -2.5)),
+            ds.lat.copy(data=np.arange(73.75, 13.6, -2.5)),
         ],
         dim="bounds",
     )
     expected["lon"] = xr.concat(
         [
-            obj.lon.copy(data=np.arange(198.75, 325 - 1.25, 2.5)),
-            obj.lon.copy(data=np.arange(201.25, 325 + 1.25, 2.5)),
+            ds.lon.copy(data=np.arange(198.75, 325 - 1.25, 2.5)),
+            ds.lon.copy(data=np.arange(201.25, 325 + 1.25, 2.5)),
         ],
         dim="bounds",
     )
@@ -759,8 +759,8 @@ def test_add_bounds(dims):
     dtb2 = pd.Timedelta("3h")
     expected["time"] = xr.concat(
         [
-            obj.time.copy(data=pd.date_range(start=t0 - dtb2, end=t1 - dtb2, freq=dt)),
-            obj.time.copy(data=pd.date_range(start=t0 + dtb2, end=t1 + dtb2, freq=dt)),
+            ds.time.copy(data=pd.date_range(start=t0 - dtb2, end=t1 - dtb2, freq=dt)),
+            ds.time.copy(data=pd.date_range(start=t0 + dtb2, end=t1 + dtb2, freq=dt)),
         ],
         dim="bounds",
     )
@@ -769,8 +769,9 @@ def test_add_bounds(dims):
     expected["lon"].attrs.clear()
     expected["time"].attrs.clear()
 
-    obj.coords["time2"] = obj.time
-    added = obj.cf.add_bounds(dims)
+    added = ds.copy(deep=False)
+    added.coords["time2"] = ds.time
+    added = added.cf.add_bounds(dims)
     if isinstance(dims, str):
         dims = (dims,)
 
@@ -780,6 +781,8 @@ def test_add_bounds(dims):
         assert added[dim].attrs["bounds"] == name
         assert_allclose(added[name].reset_coords(drop=True), expected[dim])
 
+    _check_unchanged(original, ds)
+
 
 def test_add_bounds_multiple():
     # Test multiple dimensions
@@ -810,7 +813,7 @@ def test_add_bounds_nd_variable():
 
 
 def test_bounds():
-    ds = airds.copy(deep=True).cf.add_bounds("lat")
+    ds = airds.copy(deep=False).cf.add_bounds("lat")
 
     actual = ds.cf.bounds
     expected = {"Y": ["lat_bounds"], "lat": ["lat_bounds"], "latitude": ["lat_bounds"]}
@@ -858,32 +861,36 @@ def test_bounds():
 
 def test_bounds_to_vertices():
     # All available
-    ds = airds.cf.add_bounds(["lon", "lat"])
-    dsc = ds.cf.bounds_to_vertices()
-    assert "lon_vertices" in dsc
-    assert "lat_vertices" in dsc
+    ds = airds
+    original = ds.copy(deep=True)
+    dsb = ds.cf.add_bounds(["lon", "lat"])
+    dsv = dsb.cf.bounds_to_vertices()
+    assert "lon_vertices" in dsv
+    assert "lat_vertices" in dsv
 
     # Giving key
-    dsc = ds.cf.bounds_to_vertices("longitude")
-    assert "lon_vertices" in dsc
-    assert "lat_vertices" not in dsc
+    dsv = dsb.cf.bounds_to_vertices("longitude")
+    assert "lon_vertices" in dsv
+    assert "lat_vertices" not in dsv
 
-    dsc = ds.cf.bounds_to_vertices(["longitude", "latitude"])
-    assert "lon_vertices" in dsc
-    assert "lat_vertices" in dsc
+    dsv = dsb.cf.bounds_to_vertices(["longitude", "latitude"])
+    assert "lon_vertices" in dsv
+    assert "lat_vertices" in dsv
 
     # Error
     with pytest.raises(ValueError):
-        dsc = ds.cf.bounds_to_vertices("T")
+        dsv = dsb.cf.bounds_to_vertices("T")
 
     # Words on datetime arrays to
-    ds = airds.cf.add_bounds("time")
-    dsc = ds.cf.bounds_to_vertices()
-    assert "time_bounds" in dsc
+    dsb = dsb.cf.add_bounds("time")
+    dsv = dsb.cf.bounds_to_vertices()
+    assert "time_bounds" in dsv
+
+    _check_unchanged(original, ds)
 
 
 def test_get_bounds_dim_name():
-    ds = airds.copy(deep=True).cf.add_bounds("lat")
+    ds = airds.cf.add_bounds("lat")
     assert ds.cf.get_bounds_dim_name("latitude") == "bounds"
     assert ds.cf.get_bounds_dim_name("lat") == "bounds"
 
@@ -921,6 +928,26 @@ def _make_names(prefixes):
     ]
 
 
+def _check_unchanged(old, new):
+    # Check data array attributes or global dataset attributes
+    assert type(old) == type(new)
+    assert old.attrs.keys() == new.attrs.keys()  # set comparison
+    for att, old_val in old.attrs.items():
+        assert id(old_val) == id(new.attrs[att])
+
+    # Check coordinate attributes and data variable attributes
+    dicts = [(old.coords, new.coords)]
+    if isinstance(old, xr.Dataset):
+        dicts.append((old.data_vars, new.data_vars))
+    for old_dict, new_dict in dicts:
+        assert old_dict.keys() == new_dict.keys()  # set comparison
+        for key, old_obj in old_dict.items():
+            new_obj = new_dict[key]
+            assert old_obj.attrs.keys() == new_obj.attrs.keys()  # set comparison
+            for att, old_val in old_obj.attrs.items():
+                assert id(old_val) == id(new_obj.attrs[att])  # numpy-safe comparison
+
+
 _TIME_NAMES = ["t"] + _make_names(
     [
         "time",
@@ -1015,7 +1042,7 @@ def test_attributes():
     with pytest.raises(AttributeError):
         airds.da.cf.chunks
 
-    airds2 = airds.copy(deep=True)
+    airds2 = airds.copy(deep=False)
     airds2.lon.attrs = {}
     actual = airds2.cf.sizes
     expected = {"lon": 50, "Y": 25, "T": 4, "latitude": 25, "time": 4}
@@ -1046,7 +1073,7 @@ def test_attributes():
     assert_identical(ds1.cf.data_vars["T"], ds1["T"])
 
     # multiple latitudes but only one latitude data_var
-    ds = popds.copy(deep=True)
+    ds = popds.copy(deep=False)
     for var in ["ULAT", "TLAT"]:
         ds[var].attrs["standard_name"] = "latitude"
     ds = ds.reset_coords("ULAT")
@@ -1068,7 +1095,7 @@ def test_Z_vs_vertical_ROMS():
         romsds.z_rho_dummy.reset_coords(drop=True), romsds.temp.cf["vertical"]
     )
 
-    romsds = romsds.copy(deep=True)
+    romsds = romsds.copy(deep=False)
 
     romsds.temp.attrs.clear()
     # look in encoding
@@ -1109,12 +1136,12 @@ def test_param_vcoord_ocean_s_coord():
     romsds.cf.decode_vertical_coords(outnames={"s_rho": "ZZZ_rho"})
     assert "ZZZ_rho" in romsds.coords
 
-    copy = romsds.copy(deep=True)
+    copy = romsds.copy(deep=False)
     del copy["zeta"]
     with pytest.raises(KeyError):
         copy.cf.decode_vertical_coords(outnames={"s_rho": "z_rho"})
 
-    copy = romsds.copy(deep=True)
+    copy = romsds.copy(deep=False)
     copy.s_rho.attrs["formula_terms"] = "s: s_rho C: Cs_r depth: h depth_c: hc"
     with pytest.raises(KeyError):
         copy.cf.decode_vertical_coords(outnames={"s_rho": "z_rho"})
@@ -1125,7 +1152,7 @@ def test_param_vcoord_ocean_sigma_coordinate():
     pomds.cf.decode_vertical_coords(outnames={"sigma": "z"})
     assert_allclose(pomds.z.reset_coords(drop=True), expected.reset_coords(drop=True))
 
-    copy = pomds.copy(deep=True)
+    copy = pomds.copy(deep=False)
     del copy["zeta"]
     with pytest.raises(AssertionError):
         copy.cf.decode_vertical_coords()
@@ -1146,7 +1173,7 @@ def test_formula_terms():
     assert romsds["temp"].cf.formula_terms == srhoterms
     assert romsds["s_rho"].cf.formula_terms == srhoterms
 
-    s_rho = romsds["s_rho"].copy(deep=True)
+    s_rho = romsds["s_rho"].copy(deep=False)
     del s_rho.attrs["standard_name"]
     del s_rho.s_rho.attrs["standard_name"]  # TODO: xarray bug
     assert s_rho.cf.formula_terms == srhoterms
@@ -1535,6 +1562,7 @@ def test_datetime_like(reshape):
 def test_add_canonical_attributes(override, skip, verbose, capsys):
 
     ds = airds
+    original = ds.copy(deep=True)
     cf_ds = ds.cf.add_canonical_attributes(
         override=override, skip=skip, verbose=verbose
     )
@@ -1581,6 +1609,8 @@ def test_add_canonical_attributes(override, skip, verbose, capsys):
         cf_da.attrs.pop("history")
         assert_identical(cf_da, cf_ds["air"])
 
+    _check_unchanged(original, ds)
+
 
 @pytest.mark.parametrize("op", ["ge", "gt", "eq", "ne", "le", "lt"])
 def test_flag_features(op):
@@ -1617,20 +1647,20 @@ def test_flag_errors():
 def test_missing_variables():
 
     # Bounds
-    ds = mollwds.copy(deep=True)
+    ds = mollwds.copy(deep=False)
     ds = ds.drop_vars("lon_bounds")
     assert ds.cf.bounds == {"lat": ["lat_bounds"], "latitude": ["lat_bounds"]}
 
     with pytest.raises(KeyError, match=r"No results found for 'longitude'."):
         ds.cf.get_bounds("longitude")
 
     # Cell measures
-    ds = airds.copy(deep=True)
+    ds = airds.copy(deep=False)
     ds = ds.drop_vars("cell_area")
     assert ds.cf.cell_measures == {}
 
     # Formula terms
-    ds = vert.copy(deep=True)
+    ds = vert.copy(deep=False)
     ds = ds.drop_vars("ap")
     assert ds.cf.formula_terms == {"lev": {"b": "b", "ps": "ps"}}