Support deepcopying without copying references to the parent cube

changelog/996.bugfix.rst

@@ -0,0 +1,2 @@
+Deep copies of `~ndcube.ExtraCoords`, `~ndcube.GlobalCoords` and `~ndcube.visualization.BasePlotter` no longer copy their parent cube.
+This reduces memory use in operations such as arithmetic and `~ndcube.NDCube.to`.

ndcube/conftest.py

@@ -3,6 +3,7 @@
 predictable NDCube objects.
 """
 import logging
+from unittest.mock import Mock
 
 import numpy as np
 import pytest
@@ -1328,6 +1329,12 @@ def lut_1d_wave():
     return QuantityTableCoordinate(range(10) * u.nm)
 
 
+@pytest.fixture
+def forbid_cube_deepcopy(monkeypatch):
+    """Make any deep copy of an `~ndcube.NDCube` fail the test."""
+    monkeypatch.setattr(NDCube, "__deepcopy__", Mock(side_effect=AssertionError("Parent cube was copied")), raising=False)
+
+
 def pytest_runtest_teardown(item):
     # Clear the pyplot figure stack if it is not empty after the test
     # You can see these log messages by passing "-o log_cli=true" to pytest on the command line

ndcube/extra_coords/extra_coords.py

@@ -14,7 +14,7 @@
 from astropy.wcs.wcsapi.high_level_wcs_wrapper import HighLevelWCSWrapper
 from astropy.wcs.wcsapi.wrappers.sliced_wcs import SlicedLowLevelWCS
 
-from ndcube.utils.misc import _simplify_basic_index
+from ndcube.utils.misc import _deepcopy_without_parent, _simplify_basic_index
 from ndcube.utils.wcs import convert_between_array_and_pixel_axes
 from ndcube.wcs.wrappers import CompoundLowLevelWCS, ResampledLowLevelWCS
 from .table_coord import (
@@ -156,6 +156,8 @@ def __init__(self, ndcube=None):
         # We need a reference to the parent NDCube
         self._ndcube = ndcube
 
+    __deepcopy__ = _deepcopy_without_parent
+
     @classmethod
     def from_lookup_tables(cls, names, pixel_dimensions, lookup_tables, physical_types=None):
         """

ndcube/global_coords.py

@@ -10,6 +10,7 @@
 from astropy.wcs.wcsapi.high_level_api import default_order
 from astropy.wcs.wcsapi.utils import deserialize_class
 
+from ndcube.utils.misc import _deepcopy_without_parent
 from ndcube.utils.wcs import validate_physical_types
 
 
@@ -88,6 +89,8 @@ def __init__(self, ndcube=None):
         self._ndcube = ndcube
         self._internal_coords = OrderedDict()
 
+    __deepcopy__ = _deepcopy_without_parent
+
     @staticmethod
     def _convert_dropped_to_internal(dropped_dimensions):
         """

ndcube/ndcube.py

@@ -954,8 +954,11 @@ def plot(self, *args, **kwargs):
         return self.plotter.plot(*args, **kwargs)
 
     def _new_instance(self, **kwargs):
+        # Astropy's NDUncertainty deep copies its parent_nddata, i.e. this whole cube.
+        # Seeding the memo makes that parent copy as None; the new cube then becomes its parent.
+        memo = {id(self): None}
         keys = ('unit', 'wcs', 'mask', 'meta', 'uncertainty', 'psf')
-        new_kwargs = {k: deepcopy(getattr(self, k, None)) for k in keys}
+        new_kwargs = {k: deepcopy(getattr(self, k, None), memo) for k in keys if k not in kwargs}
         # To support old versions of astropy, we need to make sure
         # we only pass in the parameters that are valid for the NDData
         params = list(inspect.signature(astropy.nddata.NDData).parameters)
@@ -1101,7 +1104,7 @@ def __mul__(self, value):
             else:
                 kwargs["unit"] = self.unit
             kwargs["data"] = self.data * value
-            kwargs["uncertainty"] = (type(self.uncertainty)(self.uncertainty.array * value)
+            kwargs["uncertainty"] = (type(self.uncertainty)(self.uncertainty.array * value, copy=False)
                             if self.uncertainty is not None else None)
         return self._new_instance(**kwargs)
 
@@ -1461,7 +1464,8 @@ def fill_masked(self, fill_value, uncertainty_fill_value=None, unmask=False, fil
         # TODO: is the logic repetitive? this else is the same with the if not fill_in_place below? No because the order matters.
         else:
             new_data = copy.deepcopy(self.data)
-            new_uncertainty = copy.deepcopy(self.uncertainty)
+            # The memo stops deepcopy copying this cube via the uncertainty's parent; see _new_instance.
+            new_uncertainty = copy.deepcopy(self.uncertainty, {id(self): None})
             new_mask = False if unmask else copy.deepcopy(self.mask) # self.mask still exists.
 
         masked = (

ndcube/tests/test_ndcube.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 from inspect import signature
 
 import numpy as np
@@ -9,8 +10,16 @@
 from astropy.wcs.wcsapi import BaseHighLevelWCS
 from astropy.wcs.wcsapi.wrappers import SlicedLowLevelWCS
 
-from ndcube import NDCube
+from ndcube import ExtraCoords, GlobalCoords, NDCube
 from ndcube.tests import helpers
+from ndcube.visualization import BasePlotter
+
+
+class _DummyPlotter(BasePlotter):
+    """A concrete plotter for testing ``BasePlotter`` without a matplotlib dependency."""
+
+    def plot(self, *args, **kwargs):
+        pass
 
 
 def test_array_axis_physical_types(ndcube_3d_ln_lt_l):
@@ -61,6 +70,101 @@ def test_initialize_with_extra_global_coords(ndcube_3d_ln_lt_l_ec_all_axes):
     helpers.assert_global_coords_equal(new_cube_copy.global_coords, gc)
 
 
+@pytest.mark.parametrize("coords_name", ["extra_coords", "global_coords"])
+def test_deepcopy_coords_does_not_copy_parent(ndcube_3d_l_ln_lt_ectime, coords_name, forbid_cube_deepcopy):
+    cube = ndcube_3d_l_ln_lt_ectime
+    cube.global_coords.add("distance", "pos.distance", 1 * u.m)
+    coords = getattr(cube, coords_name)
+
+    copied = deepcopy(coords)
+
+    assert copied is not coords
+    assert copied._ndcube is cube
+    if coords_name == "extra_coords":
+        helpers.assert_extra_coords_equal(copied, coords)
+        copied.add("exposure", 1, np.arange(5) * u.s)
+        assert list(coords.keys()) == ["time"]
+    if coords_name == "global_coords":
+        helpers.assert_global_coords_equal(copied, coords)
+        copied["distance"][...] = 2 * u.m
+        assert coords["distance"] == 1 * u.m
+
+
+@pytest.mark.parametrize("coords_type", [ExtraCoords, GlobalCoords])
+def test_deepcopy_coords_cycles(coords_type):
+    coords = coords_type()
+    coords.cycle = [coords]
+    coords.alias = coords.cycle
+
+    copied = deepcopy(coords)
+
+    assert copied.cycle is not coords.cycle
+    assert copied.cycle[0] is copied
+    assert copied.alias is copied.cycle
+
+
+def test_deepcopy_cube_keeps_coords_linked(ndcube_3d_l_ln_lt_ectime):
+    cube = ndcube_3d_l_ln_lt_ectime
+    cube.global_coords.add("distance", "pos.distance", 1 * u.m)
+
+    copied = deepcopy(cube)
+
+    helpers.assert_cubes_equal(copied, cube, check_uncertainty_values=True)
+    helpers.assert_global_coords_equal(copied.global_coords, cube.global_coords)
+    assert copied.extra_coords._ndcube is copied
+    assert copied.global_coords._ndcube is copied
+    assert cube.extra_coords._ndcube is cube
+    assert cube.global_coords._ndcube is cube
+
+
+def test_deepcopy_plotter_does_not_copy_parent(ndcube_3d_l_ln_lt_ectime, forbid_cube_deepcopy):
+    cube = ndcube_3d_l_ln_lt_ectime
+    plotter = _DummyPlotter(cube)
+    plotter.canvas = [1, 2, 3]
+
+    copied = deepcopy(plotter)
+
+    assert copied is not plotter
+    assert copied._ndcube is cube
+    assert copied.canvas == [1, 2, 3]
+    copied.canvas.append(4)
+    assert plotter.canvas == [1, 2, 3]
+
+
+def test_deepcopy_plotter_cycles():
+    plotter = _DummyPlotter()
+    plotter.cycle = [plotter]
+    plotter.alias = plotter.cycle
+
+    copied = deepcopy(plotter)
+
+    assert copied.cycle is not plotter.cycle
+    assert copied.cycle[0] is copied
+    assert copied.alias is copied.cycle
+
+
+def test_deepcopy_plotter_without_base_init():
+    class NoInitPlotter(_DummyPlotter):
+        def __init__(self):
+            pass
+
+    copied = deepcopy(NoInitPlotter())
+
+    assert not hasattr(copied, "_ndcube")
+
+
+def test_deepcopy_cube_keeps_plotter_linked(ndcube_3d_l_ln_lt_ectime):
+    cube = ndcube_3d_l_ln_lt_ectime
+    cube._plotter = _DummyPlotter(cube)
+
+    copied = deepcopy(cube)
+
+    helpers.assert_cubes_equal(copied, cube, check_uncertainty_values=True)
+    assert copied.plotter is not cube.plotter
+    assert copied.plotter._ndcube is copied
+    assert cube.plotter._ndcube is cube
+
+
 def test_wcs_type_after_init(ndcube_3d_ln_lt_l, wcs_3d_l_lt_ln):
     # Generate a low level WCS
     slices = np.s_[:, :, 0]
@@ -233,12 +337,14 @@ def test_fill_masked_fill_in_place_true(ndc, fill_value, uncertainty_fill_value,
     ],
     indirect=("ndc", "expected_cube")
 )
-def test_fill_masked_fill_in_place_false(ndc, fill_value, uncertainty_fill_value, unmask, expected_cube):
+def test_fill_masked_fill_in_place_false(ndc, fill_value, uncertainty_fill_value, unmask, expected_cube, forbid_cube_deepcopy):
     # compare the expected cube with the cube saved in the new place
 
     # perform the fill_masked method on the fixture, using parametrized as parameters.
     filled_cube = ndc.fill_masked(fill_value, uncertainty_fill_value, unmask, fill_in_place=False)
     helpers.assert_cubes_equal(filled_cube, expected_cube, check_uncertainty_values=True)
+    assert filled_cube.uncertainty.parent_nddata is filled_cube
+    assert ndc.uncertainty.parent_nddata is ndc
 
 @pytest.mark.parametrize(
     ("ndc", "fill_value", "uncertainty_fill_value", "unmask"),

ndcube/tests/test_ndcube_arithmetic.py

@@ -1,3 +1,5 @@
+import operator
+
 import numpy as np
 import pytest
 
@@ -431,6 +433,38 @@ def test_cube_arithmetic_neg(ndcube_2d_ln_lt_units):
     )
 
 
+@pytest.mark.parametrize("operation", [operator.neg, lambda cube: cube * 2, lambda cube: cube.to(u.mct)],
+                         ids=["neg", "mul", "to"])
+def test_cube_arithmetic_does_not_copy_cube(operation):
+    class NoDeepCopy(np.ndarray):
+        def __deepcopy__(self, memo):
+            raise AssertionError("Data array was deep copied")
+
+    shape = (3, 4)
+    cube = NDCube(np.ones(shape).view(NoDeepCopy), wcs=astropy.wcs.WCS(naxis=2), unit=u.ct,
+                  uncertainty=StdDevUncertainty(np.ones(shape)), meta={"a": 1})
+    cube.extra_coords.add("time", 0, np.arange(shape[0]) * u.s)
+    cube.global_coords.add("name", "custom:physical_type", 1 * u.s)
+
+    new_cube = operation(cube)
+
+    assert new_cube.extra_coords._ndcube is new_cube
+    assert new_cube.global_coords._ndcube is new_cube
+    assert new_cube.uncertainty.parent_nddata is new_cube
+    assert cube.extra_coords._ndcube is cube
+    assert cube.global_coords._ndcube is cube
+    assert cube.uncertainty.parent_nddata is cube
+
+    new_cube.meta["a"] = 2
+    new_cube.uncertainty.array[0, 0] = 5
+    new_cube.extra_coords.add("exposure", 1, np.arange(shape[1]) * u.s)
+    new_cube.global_coords.add("other", "custom:physical_type", 2 * u.s)
+    assert cube.meta["a"] == 1
+    assert cube.uncertainty.array[0, 0] == 1
+    assert list(cube.extra_coords.keys()) == ["time"]
+    assert list(cube.global_coords) == ["name"]
+
+
 def test_add_unitless_cube_typeerror(ndcube_2d_ln_lt_units):
     with pytest.raises(TypeError):
         _ = ndcube_2d_ln_lt_units + 10.0

ndcube/utils/misc.py

@@ -1,3 +1,4 @@
+import copy
 import numbers
 from collections.abc import Sequence
 from types import EllipsisType
@@ -44,6 +45,24 @@ def convert_quantities_to_units(coords, units):
             for coord, unit in zip(coords, units)]
 
 
+def _deepcopy_without_parent(obj, memo):
+    """
+    Deep copy an object without copying the parent cube it links to through ``_ndcube``.
+
+    Assign this as ``__deepcopy__`` on the class.
+    """
+    new = copy.copy(obj)
+    memo[id(obj)] = new
+    for key, value in vars(obj).items():
+        if key != "_ndcube":
+            setattr(new, key, copy.deepcopy(value, memo))
+    # A plotter subclass that skips BasePlotter.__init__ has no _ndcube.
+    if hasattr(obj, "_ndcube"):
+        # Use the cube's copy if deepcopy reached the cube first, otherwise keep the original cube.
+        new._ndcube = memo.get(id(obj._ndcube), obj._ndcube)
+    return new
+
+
 if minversion("astropy", "8.0.2"):
     from astropy.utils.shapes import simplify_basic_index as _simplify_basic_index
 else:

ndcube/visualization/base.py

@@ -1,5 +1,7 @@
 import abc
 
+from ndcube.utils.misc import _deepcopy_without_parent
+
 
 class BasePlotter(abc.ABC):
     """
@@ -21,3 +23,5 @@ def plot(self, *args, **kwargs):
 
         The ``plot()`` method **should** accept ``**kwargs``.
         """
+
+    __deepcopy__ = _deepcopy_without_parent