From be63afe56fc714891572606684da715ba70701d7 Mon Sep 17 00:00:00 2001 From: Alasdair Wilson <60351846+alasdairwilson@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:05:55 +0100 Subject: [PATCH] Backport PR #996: Support deepcopying without copying references to the parent cube --- changelog/996.bugfix.rst | 2 + ndcube/conftest.py | 7 ++ ndcube/extra_coords/extra_coords.py | 4 +- ndcube/global_coords.py | 3 + ndcube/ndcube.py | 11 ++- ndcube/tests/test_ndcube.py | 110 ++++++++++++++++++++++++- ndcube/tests/test_ndcube_arithmetic.py | 34 ++++++++ ndcube/utils/misc.py | 19 +++++ ndcube/visualization/base.py | 4 + 9 files changed, 188 insertions(+), 6 deletions(-) create mode 100644 changelog/996.bugfix.rst diff --git a/changelog/996.bugfix.rst b/changelog/996.bugfix.rst new file mode 100644 index 000000000..428f9644d --- /dev/null +++ b/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`. diff --git a/ndcube/conftest.py b/ndcube/conftest.py index 6e591ebaf..2afa617a0 100644 --- a/ndcube/conftest.py +++ b/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 diff --git a/ndcube/extra_coords/extra_coords.py b/ndcube/extra_coords/extra_coords.py index 47c9588c3..01a2fa687 100644 --- a/ndcube/extra_coords/extra_coords.py +++ b/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): """ diff --git a/ndcube/global_coords.py b/ndcube/global_coords.py index 3514a1799..af6355844 100644 --- a/ndcube/global_coords.py +++ b/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): """ diff --git a/ndcube/ndcube.py b/ndcube/ndcube.py index d712f79a4..26f249f2e 100644 --- a/ndcube/ndcube.py +++ b/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,8 @@ def __mul__(self, value): else: kwargs["unit"] = self.unit kwargs["data"] = self.data * value - kwargs["uncertainty"] = (type(self.uncertainty)(self.uncertainty.array * value) + # The array is already new, don't copy it again + 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 +1465,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 = ( diff --git a/ndcube/tests/test_ndcube.py b/ndcube/tests/test_ndcube.py index 7719b842c..252a41da1 100644 --- a/ndcube/tests/test_ndcube.py +++ b/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"), diff --git a/ndcube/tests/test_ndcube_arithmetic.py b/ndcube/tests/test_ndcube_arithmetic.py index 29b0275b6..b16235940 100644 --- a/ndcube/tests/test_ndcube_arithmetic.py +++ b/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 diff --git a/ndcube/utils/misc.py b/ndcube/utils/misc.py index eb19a03e8..aa24c348a 100644 --- a/ndcube/utils/misc.py +++ b/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: diff --git a/ndcube/visualization/base.py b/ndcube/visualization/base.py index d981242b0..76aa3ab5b 100644 --- a/ndcube/visualization/base.py +++ b/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