File: C:/Users/fred/anaconda3/Lib/site-packages/astropy/nddata/tests/test_nddata.py
# Licensed under a 3-clause BSD style license - see LICENSE.rst
import pickle
import textwrap
from itertools import chain, permutations
import numpy as np
import pytest
from numpy.testing import assert_array_equal
from astropy import units as u
from astropy.nddata import NDDataArray
from astropy.nddata import _testing as nd_testing
from astropy.nddata.nddata import NDData
from astropy.nddata.nduncertainty import StdDevUncertainty
from astropy.utils import NumpyRNGContext
from astropy.utils.compat.optional_deps import HAS_DASK
from astropy.utils.masked import Masked
from astropy.utils.metadata.tests.test_metadata import MetaBaseTest
from astropy.wcs import WCS
from astropy.wcs.wcsapi import BaseHighLevelWCS, HighLevelWCSWrapper, SlicedLowLevelWCS
from .test_nduncertainty import FakeUncertainty
class FakeNumpyArray:
"""
Class that has a few of the attributes of a numpy array.
These attributes are checked for by NDData.
"""
def __init__(self):
super().__init__()
def shape(self):
pass
def __getitem__(self, key):
pass
def __array__(self, dtype=None, copy=None):
pass
@property
def dtype(self):
return "fake"
class MinimalUncertainty:
"""
Define the minimum attributes acceptable as an uncertainty object.
"""
def __init__(self, value):
self._uncertainty = value
@property
def uncertainty_type(self):
return "totally and completely fake"
class BadNDDataSubclass(NDData):
def __init__(
self,
data,
uncertainty=None,
mask=None,
wcs=None,
meta=None,
unit=None,
psf=None,
):
self._data = data
self._uncertainty = uncertainty
self._mask = mask
self._wcs = wcs
self._psf = psf
self._unit = unit
self._meta = meta
# Setter tests
def test_uncertainty_setter():
nd = NDData([1, 2, 3])
good_uncertainty = MinimalUncertainty(5)
nd.uncertainty = good_uncertainty
assert nd.uncertainty is good_uncertainty
# Check the fake uncertainty (minimal does not work since it has no
# parent_nddata attribute from NDUncertainty)
nd.uncertainty = FakeUncertainty(5)
assert nd.uncertainty.parent_nddata is nd
# Check that it works if the uncertainty was set during init
nd = NDData(nd)
assert isinstance(nd.uncertainty, FakeUncertainty)
nd.uncertainty = 10
assert not isinstance(nd.uncertainty, FakeUncertainty)
assert nd.uncertainty.array == 10
def test_mask_setter():
# Since it just changes the _mask attribute everything should work
nd = NDData([1, 2, 3])
nd.mask = True
assert nd.mask
nd.mask = False
assert not nd.mask
# Check that it replaces a mask from init
nd = NDData(nd, mask=True)
assert nd.mask
nd.mask = False
assert not nd.mask
# Init tests
def test_nddata_empty():
with pytest.raises(TypeError):
NDData() # empty initializer should fail
def test_nddata_init_data_nonarray():
inp = [1, 2, 3]
nd = NDData(inp)
assert (np.array(inp) == nd.data).all()
def test_nddata_init_data_ndarray():
# random floats
with NumpyRNGContext(123):
nd = NDData(np.random.random((10, 10)))
assert nd.data.shape == (10, 10)
assert nd.data.size == 100
assert nd.data.dtype == np.dtype(float)
# specific integers
nd = NDData(np.array([[1, 2, 3], [4, 5, 6]]))
assert nd.data.size == 6
assert nd.data.dtype == np.dtype(int)
# Tests to ensure that creating a new NDData object copies by *reference*.
a = np.ones((10, 10))
nd_ref = NDData(a)
a[0, 0] = 0
assert nd_ref.data[0, 0] == 0
# Except we choose copy=True
a = np.ones((10, 10))
nd_ref = NDData(a, copy=True)
a[0, 0] = 0
assert nd_ref.data[0, 0] != 0
def test_nddata_init_data_maskedarray():
with NumpyRNGContext(456):
NDData(np.random.random((10, 10)), mask=np.random.random((10, 10)) > 0.5)
# Another test (just copied here)
with NumpyRNGContext(12345):
a = np.random.randn(100)
marr = np.ma.masked_where(a > 0, a)
nd = NDData(marr)
# check that masks and data match
assert_array_equal(nd.mask, marr.mask)
assert_array_equal(nd.data, marr.data)
# check that they are both by reference
marr.mask[10] = ~marr.mask[10]
marr.data[11] = 123456789
assert_array_equal(nd.mask, marr.mask)
assert_array_equal(nd.data, marr.data)
# or not if we choose copy=True
nd = NDData(marr, copy=True)
marr.mask[10] = ~marr.mask[10]
marr.data[11] = 0
assert nd.mask[10] != marr.mask[10]
assert nd.data[11] != marr.data[11]
@pytest.mark.parametrize("data", [np.array([1, 2, 3]), 5])
def test_nddata_init_data_quantity(data):
# Test an array and a scalar because a scalar Quantity does not always
# behave the same way as an array.
quantity = data * u.adu
ndd = NDData(quantity)
assert ndd.unit == quantity.unit
assert_array_equal(ndd.data, np.array(quantity))
if ndd.data.size > 1:
# check that if it is an array it is not copied
quantity.value[1] = 100
assert ndd.data[1] == quantity.value[1]
# or is copied if we choose copy=True
ndd = NDData(quantity, copy=True)
quantity.value[1] = 5
assert ndd.data[1] != quantity.value[1]
# provide a quantity and override the unit
ndd_unit = NDData(data * u.erg, unit=u.J)
assert ndd_unit.unit == u.J
np.testing.assert_allclose((ndd_unit.data * ndd_unit.unit).to_value(u.erg), data)
def test_nddata_init_data_masked_quantity():
a = np.array([2, 3])
q = a * u.m
m = False
mq = Masked(q, mask=m)
nd = NDData(mq)
assert_array_equal(nd.data, a)
# This test failed before the change in nddata init because the masked
# arrays data (which in fact was a quantity was directly saved)
assert nd.unit == u.m
assert not isinstance(nd.data, u.Quantity)
np.testing.assert_array_equal(nd.mask, np.array(m))
def test_nddata_init_data_nddata():
nd1 = NDData(np.array([1]))
nd2 = NDData(nd1)
assert nd2.wcs == nd1.wcs
assert nd2.uncertainty == nd1.uncertainty
assert nd2.mask == nd1.mask
assert nd2.unit == nd1.unit
assert nd2.meta == nd1.meta
assert nd2.psf == nd1.psf
# Check that it is copied by reference
nd1 = NDData(np.ones((5, 5)))
nd2 = NDData(nd1)
assert nd1.data is nd2.data
# Check that it is really copied if copy=True
nd2 = NDData(nd1, copy=True)
nd1.data[2, 3] = 10
assert nd1.data[2, 3] != nd2.data[2, 3]
# Now let's see what happens if we have all explicitly set
nd1 = NDData(
np.array([1]),
mask=False,
uncertainty=StdDevUncertainty(10),
unit=u.s,
meta={"dest": "mordor"},
wcs=WCS(naxis=1),
psf=np.array([10]),
)
nd2 = NDData(nd1)
assert nd2.data is nd1.data
assert nd2.wcs is nd1.wcs
assert nd2.uncertainty.array == nd1.uncertainty.array
assert nd2.mask == nd1.mask
assert nd2.unit == nd1.unit
assert nd2.meta == nd1.meta
assert nd2.psf == nd1.psf
# now what happens if we overwrite them all too
nd3 = NDData(
nd1,
mask=True,
uncertainty=StdDevUncertainty(200),
unit=u.km,
meta={"observer": "ME"},
wcs=WCS(naxis=1),
psf=np.array([20]),
)
assert nd3.data is nd1.data
assert nd3.wcs is not nd1.wcs
assert nd3.uncertainty.array != nd1.uncertainty.array
assert nd3.mask != nd1.mask
assert nd3.unit != nd1.unit
assert nd3.meta != nd1.meta
assert nd3.psf != nd1.psf
def test_nddata_init_data_nddata_subclass():
uncert = StdDevUncertainty(3)
# There might be some incompatible subclasses of NDData around.
bnd = BadNDDataSubclass(False, True, 3, 2, "gollum", 100, 12)
# Before changing the NDData init this would not have raised an error but
# would have lead to a compromised nddata instance
with pytest.raises(TypeError):
NDData(bnd)
# but if it has no actual incompatible attributes it passes
bnd_good = BadNDDataSubclass(
np.array([1, 2]),
uncert,
3,
HighLevelWCSWrapper(WCS(naxis=1)),
{"enemy": "black knight"},
u.km,
)
nd = NDData(bnd_good)
assert nd.unit == bnd_good.unit
assert nd.meta == bnd_good.meta
assert nd.uncertainty == bnd_good.uncertainty
assert nd.mask == bnd_good.mask
assert nd.wcs is bnd_good.wcs
assert nd.data is bnd_good.data
def test_nddata_init_data_fail():
# First one is sliceable but has no shape, so should fail.
with pytest.raises(TypeError):
NDData({"a": "dict"})
# This has a shape but is not sliceable
class Shape:
def __init__(self):
self.shape = 5
def __repr__(self):
return "7"
with pytest.raises(TypeError):
NDData(Shape())
def test_nddata_init_data_fakes():
ndd1 = NDData(FakeNumpyArray())
# First make sure that NDData isn't converting its data to a numpy array.
assert isinstance(ndd1.data, FakeNumpyArray)
# Make a new NDData initialized from an NDData
ndd2 = NDData(ndd1)
# Check that the data wasn't converted to numpy
assert isinstance(ndd2.data, FakeNumpyArray)
# Specific parameters
def test_param_uncertainty():
u = StdDevUncertainty(array=np.ones((5, 5)))
d = NDData(np.ones((5, 5)), uncertainty=u)
# Test that the parent_nddata is set.
assert d.uncertainty.parent_nddata is d
# Test conflicting uncertainties (other NDData)
u2 = StdDevUncertainty(array=np.ones((5, 5)) * 2)
d2 = NDData(d, uncertainty=u2)
assert d2.uncertainty is u2
assert d2.uncertainty.parent_nddata is d2
def test_param_wcs():
# Since everything is allowed we only need to test something
nd = NDData([1], wcs=WCS(naxis=1))
assert nd.wcs is not None
# Test conflicting wcs (other NDData)
nd2 = NDData(nd, wcs=WCS(naxis=1))
assert nd2.wcs is not None and nd2.wcs is not nd.wcs
def test_param_meta():
# everything dict-like is allowed
with pytest.raises(TypeError):
NDData([1], meta=3)
nd = NDData([1, 2, 3], meta={})
assert len(nd.meta) == 0
nd = NDData([1, 2, 3])
assert isinstance(nd.meta, dict)
assert len(nd.meta) == 0
# Test conflicting meta (other NDData)
nd2 = NDData(nd, meta={"image": "sun"})
assert len(nd2.meta) == 1
nd3 = NDData(nd2, meta={"image": "moon"})
assert len(nd3.meta) == 1
assert nd3.meta["image"] == "moon"
def test_param_mask():
# Since everything is allowed we only need to test something
nd = NDData([1], mask=False)
assert not nd.mask
# Test conflicting mask (other NDData)
nd2 = NDData(nd, mask=True)
assert nd2.mask
# (masked array)
nd3 = NDData(np.ma.array([1], mask=False), mask=True)
assert nd3.mask
# (masked quantity)
mq = np.ma.array(np.array([2, 3]) * u.m, mask=False)
nd4 = NDData(mq, mask=True)
assert nd4.mask
def test_param_unit():
with pytest.raises(ValueError):
NDData(np.ones((5, 5)), unit="NotAValidUnit")
NDData([1, 2, 3], unit="meter")
# Test conflicting units (quantity as data)
q = np.array([1, 2, 3]) * u.m
nd = NDData(q, unit="cm")
assert nd.unit != q.unit
assert nd.unit == u.cm
# (masked quantity)
mq = np.ma.array(np.array([2, 3]) * u.m, mask=False)
nd2 = NDData(mq, unit=u.pc)
assert nd2.unit == u.pc
# (another NDData as data)
nd3 = NDData(nd, unit="km")
assert nd3.unit == u.km
# (MaskedQuantity given to NDData)
mq_astropy = Masked.from_unmasked(q, False)
nd4 = NDData(mq_astropy, unit="km")
assert nd4.unit == u.km
def test_pickle_nddata_with_uncertainty():
ndd = NDData(
np.ones(3), uncertainty=StdDevUncertainty(np.ones(5), unit=u.m), unit=u.m
)
ndd_dumped = pickle.dumps(ndd)
ndd_restored = pickle.loads(ndd_dumped)
assert type(ndd_restored.uncertainty) is StdDevUncertainty
assert ndd_restored.uncertainty.parent_nddata is ndd_restored
assert ndd_restored.uncertainty.unit == u.m
def test_pickle_uncertainty_only():
ndd = NDData(
np.ones(3), uncertainty=StdDevUncertainty(np.ones(5), unit=u.m), unit=u.m
)
uncertainty_dumped = pickle.dumps(ndd.uncertainty)
uncertainty_restored = pickle.loads(uncertainty_dumped)
np.testing.assert_array_equal(ndd.uncertainty.array, uncertainty_restored.array)
assert ndd.uncertainty.unit == uncertainty_restored.unit
# Even though it has a parent there is no one that references the parent
# after unpickling so the weakref "dies" immediately after unpickling
# finishes.
assert uncertainty_restored.parent_nddata is None
def test_pickle_nddata_without_uncertainty():
ndd = NDData(np.ones(3), unit=u.m)
dumped = pickle.dumps(ndd)
ndd_restored = pickle.loads(dumped)
np.testing.assert_array_equal(ndd.data, ndd_restored.data)
# Check that the meta descriptor is working as expected. The MetaBaseTest class
# takes care of defining all the tests, and we simply have to define the class
# and any minimal set of args to pass.
class TestMetaNDData(MetaBaseTest):
test_class = NDData
args = np.array([[1.0]])
# Representation tests
def test_nddata_str():
arr1d = NDData(np.array([1, 2, 3]))
assert str(arr1d) == "[1 2 3]"
arr2d = NDData(np.array([[1, 2], [3, 4]]))
assert str(arr2d) == textwrap.dedent(
"""
[[1 2]
[3 4]]"""[1:]
)
arr3d = NDData(np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]))
assert str(arr3d) == textwrap.dedent(
"""
[[[1 2]
[3 4]]
[[5 6]
[7 8]]]"""[1:]
)
# let's add units!
arr = NDData(np.array([1, 2, 3]), unit="km")
assert str(arr) == "[1 2 3] km"
# what if it had these units?
arr = NDData(np.array([1, 2, 3]), unit="erg cm^-2 s^-1 A^-1")
assert str(arr) == "[1 2 3] erg / (A s cm2)"
def test_nddata_repr():
# The big test is eval(repr()) should be equal to the original!
# but this must be modified slightly since adopting the
# repr machinery from astropy.utils.masked
arr1d = NDData(np.array([1, 2, 3]))
s = repr(arr1d)
assert s == "NDData([1, 2, 3])"
got = eval(s)
assert np.all(got.data == arr1d.data)
assert got.unit == arr1d.unit
arr2d = NDData(np.array([[1, 2], [3, 4]]))
s = repr(arr2d)
assert s == ("NDData([[1, 2],\n [3, 4]])")
got = eval(s)
assert np.all(got.data == arr2d.data)
assert got.unit == arr2d.unit
arr3d = NDData(np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]))
s = repr(arr3d)
assert s == (
"NDData([[[1, 2],\n [3, 4]],\n\n [[5, 6],\n [7, 8]]])"
)
got = eval(s)
assert np.all(got.data == arr3d.data)
assert got.unit == arr3d.unit
# let's add units!
arr = NDData(np.array([1, 2, 3]), unit="km")
s = repr(arr)
assert s == "NDData([1, 2, 3], unit='km')"
got = eval(s)
assert np.all(got.data == arr.data)
assert got.unit == arr.unit
@pytest.mark.skipif(not HAS_DASK, reason="requires dask to be available")
def test_nddata_repr_dask():
import dask.array as da
arr = NDData(da.arange(3), unit="km")
s = repr(arr)
# just check repr equality for dask arrays, not round-tripping:
assert s in (
'NDData(\n data=dask.array<arange, shape=(3,), dtype=int64, chunksize=(3,), chunktype=numpy.ndarray>,\n unit=Unit("km")\n)',
'NDData(\n data=dask.array<arange, shape=(3,), dtype=int32, chunksize=(3,), chunktype=numpy.ndarray>,\n unit=Unit("km")\n)',
)
# Not supported features
def test_slicing_not_supported():
ndd = NDData(np.ones((5, 5)))
with pytest.raises(TypeError):
ndd[0]
def test_arithmetic_not_supported():
ndd = NDData(np.ones((5, 5)))
with pytest.raises(TypeError):
ndd + ndd
def test_nddata_wcs_setter_error_cases():
ndd = NDData(np.ones((5, 5)))
# Setting with a non-WCS should raise an error
with pytest.raises(TypeError):
ndd.wcs = "I am not a WCS"
naxis = 2
# This should succeed since the WCS is currently None
ndd.wcs = nd_testing._create_wcs_simple(
naxis=naxis,
ctype=["deg"] * naxis,
crpix=[0] * naxis,
crval=[10] * naxis,
cdelt=[1] * naxis,
)
with pytest.raises(ValueError):
# This should fail since the WCS is not None
ndd.wcs = nd_testing._create_wcs_simple(
naxis=naxis,
ctype=["deg"] * naxis,
crpix=[0] * naxis,
crval=[10] * naxis,
cdelt=[1] * naxis,
)
def test_nddata_wcs_setter_with_low_level_wcs():
ndd = NDData(np.ones((5, 5)))
wcs = WCS()
# If the wcs property is set with a low level WCS it should get
# wrapped to high level.
low_level = SlicedLowLevelWCS(wcs, 5)
assert not isinstance(low_level, BaseHighLevelWCS)
ndd.wcs = low_level
assert isinstance(ndd.wcs, BaseHighLevelWCS)
def test_nddata_init_with_low_level_wcs():
wcs = WCS()
low_level = SlicedLowLevelWCS(wcs, 5)
ndd = NDData(np.ones((5, 5)), wcs=low_level)
assert isinstance(ndd.wcs, BaseHighLevelWCS)
class NDDataCustomWCS(NDData):
@property
def wcs(self):
return WCS()
def test_overriden_wcs():
# Check that a sub-class that overrides `.wcs` without providing a setter
# works
NDDataCustomWCS(np.ones((5, 5)))
# set up parameters for test_collapse:
np.random.seed(42)
collapse_units = [None, u.Jy]
collapse_propagate = [True, False]
collapse_data_shapes = [
# 3D example:
(4, 3, 2),
# 5D example
(6, 5, 4, 3, 2),
]
collapse_ignore_masked = [True, False]
collapse_masks = list(
chain.from_iterable(
[
# try the operations without a mask (all False):
np.zeros(collapse_data_shape).astype(bool)
]
+ [
# assemble a bunch of random masks:
np.random.randint(0, 2, size=collapse_data_shape).astype(bool)
for _ in range(10)
]
for collapse_data_shape in collapse_data_shapes
)
)
# the following provides pytest.mark.parametrize with every
# permutation of (1) the units, (2) propagating/not propagating
# uncertainties, and (3) the data shapes of different ndim.
permute = (
len(collapse_masks)
* len(collapse_propagate)
* len(collapse_units)
* len(collapse_ignore_masked)
)
collapse_units = permute // len(collapse_units) * collapse_units
collapse_propagate = permute // len(collapse_propagate) * collapse_propagate
collapse_masks = permute // len(collapse_masks) * collapse_masks
collapse_ignore_masked = permute // len(collapse_ignore_masked) * collapse_ignore_masked
@pytest.mark.parametrize(
"mask, unit, propagate_uncertainties, operation_ignores_mask",
zip(collapse_masks, collapse_units, collapse_propagate, collapse_ignore_masked),
)
def test_collapse(mask, unit, propagate_uncertainties, operation_ignores_mask):
# unique set of combinations of each of the N-1 axes for an N-D cube:
axes_permutations = {tuple(axes[:2]) for axes in permutations(range(mask.ndim))}
# each of the single axis slices:
axes_permutations.update(set(range(mask.ndim)))
axes_permutations.update({None})
cube = np.arange(np.prod(mask.shape)).reshape(mask.shape)
numpy_cube = np.ma.masked_array(cube, mask=mask)
ma_cube = Masked(cube, mask=mask)
ndarr = NDDataArray(cube, uncertainty=StdDevUncertainty(cube), unit=unit, mask=mask)
# By construction, the minimum value along each axis is always the zeroth index and
# the maximum is always the last along that axis. We verify that here, so we can
# test that the correct uncertainties are extracted during the
# `NDDataArray.min` and `NDDataArray.max` methods later:
for axis in range(cube.ndim):
assert np.all(np.equal(cube.argmin(axis=axis), 0))
assert np.all(np.equal(cube.argmax(axis=axis), cube.shape[axis] - 1))
# confirm that supported nddata methods agree with corresponding numpy methods
# for the masked data array:
sum_methods = ["sum", "mean"]
ext_methods = ["min", "max"]
all_methods = sum_methods + ext_methods
# for all supported methods, ensure the masking is propagated:
for method in all_methods:
for axes in axes_permutations:
astropy_method = getattr(ma_cube, method)(axis=axes)
numpy_method = getattr(numpy_cube, method)(axis=axes)
nddata_method = getattr(ndarr, method)(
axis=axes,
propagate_uncertainties=propagate_uncertainties,
operation_ignores_mask=operation_ignores_mask,
)
astropy_unmasked = astropy_method.base[~astropy_method.mask]
nddata_unmasked = nddata_method.data[~nddata_method.mask]
# check if the units are passed through correctly:
assert unit == nddata_method.unit
# check if the numpy and astropy.utils.masked results agree when
# the result is not fully masked:
if len(astropy_unmasked) > 0:
if not operation_ignores_mask:
# compare with astropy
assert np.all(np.equal(astropy_unmasked, nddata_unmasked))
assert np.all(np.equal(astropy_method.mask, nddata_method.mask))
else:
# compare with numpy
assert np.ma.all(
np.ma.equal(numpy_method, np.asanyarray(nddata_method))
)
# For extremum methods, ensure the uncertainty returned corresponds to the
# min/max data value. We've created the uncertainties to have the same value
# as the data array, so we can just check for equality:
if method in ext_methods and propagate_uncertainties:
assert np.ma.all(np.ma.equal(astropy_method, nddata_method))