diff --git a/Changelog b/Changelog index db1bae806..e99cea6b7 100644 --- a/Changelog +++ b/Changelog @@ -1,5 +1,7 @@ version 1.7.6 (not yet released) ================================ + * Preserve explicit real variable dimensions and I/O with auto_complex=True, + while retaining complex decoding on reopen (issue #1404, PR #1525). * Speed up index setup in _StartCountStride (used when indexing variables) by skipping numpy conversion for plain ints and slices and classifying each index once (PR #1524). diff --git a/src/netCDF4/_netCDF4.pyx b/src/netCDF4/_netCDF4.pyx index c9e46e4ca..277ce0957 100644 --- a/src/netCDF4/_netCDF4.pyx +++ b/src/netCDF4/_netCDF4.pyx @@ -1254,6 +1254,14 @@ When writing files using `auto_complex=True`, `netCDF4` will use: - or a dimension of length 2 named `_pfnc_complex` for netCDF3 or classic formats. +Variables created with a real or compound datatype retain that datatype and +all of their dimensions, even when `auto_complex=True`. For example, a real +variable with dimensions `("x", "complex")` accepts separate real and imaginary +components along its last axis. Reopening that file with `auto_complex=True` +interprets the stored convention as a complex variable with dimensions `("x",)`. +To write complex arrays directly, create the variable with datatype `"c8"` or +`"c16"` instead. + Support for complex numbers is handled via the [`nc-complex`](https://github.com/PlasmaFAIR/nc-complex) library. See there for further details. @@ -4481,7 +4489,10 @@ does not do anything.""" if dim.isunlimited(): self._nunlimdim = self._nunlimdim + 1 # set ndim attribute (number of dimensions). - self.ndim = _inq_varndims(self._grpid, self._varid, self._grp.auto_complex) + # A newly created real variable may happen to use a complex-looking + # dimension. Only hide its component axis when its Python dtype is + # complex; otherwise its shape and I/O buffers must stay real. + self.ndim = _inq_varndims(self._grpid, self._varid, is_complex) self._name = name # default for automatically applying scale_factor and # add_offset, and converting to/from masked arrays is True. @@ -4609,7 +4620,7 @@ does not do anything.""" cdef int ierr, dimid cdef char namstring[NC_MAX_NAME+1] - dimids = _inq_vardimid(self._grpid, self._varid, self._grp.auto_complex) + dimids = _inq_vardimid(self._grpid, self._varid, dtype_is_complex(self.dtype)) # loop over dimensions, retrieve names. dimensions = () @@ -5937,7 +5948,7 @@ NC_CHAR). if not data.dtype.isnative: data = data.byteswap() # strides all 1 or scalar variable, use put_vara (faster) - if self._grp.auto_complex: + if dtype_is_complex(self.dtype): with nogil: ierr = pfnc_put_vars(self._grpid, self._varid, startp, countp, stridep, PyArray_DATA(data)) @@ -6062,7 +6073,7 @@ NC_CHAR). # strides all 1 or scalar variable, use get_vara (faster) # if count contains a zero element, no data is being read if 0 not in count: - if self._grp.auto_complex: + if dtype_is_complex(self.dtype): with nogil: ierr = pfnc_get_vars(self._grpid, self._varid, startp, countp, stridep, diff --git a/test/test_complex.py b/test/test_complex.py index cd1fed32c..8da63e465 100644 --- a/test/test_complex.py +++ b/test/test_complex.py @@ -14,7 +14,9 @@ class ComplexNumbersTestCase(unittest.TestCase): def setUp(self): - self.tmp_path = pathlib.Path(tempfile.mkdtemp()) + tmp_dir = tempfile.TemporaryDirectory() + self.addCleanup(tmp_dir.cleanup) + self.tmp_path = pathlib.Path(tmp_dir.name) def test_read_dim(self): filename = self.tmp_path / "test_read_dim.nc" @@ -34,6 +36,116 @@ def test_read_dim(self): assert np.array_equal(data, complex_array) + def test_write_real_with_complex_dimension(self): + # A new variable must retain its requested dtype and dimensions, even + # when its on-disk layout matches a complex-number convention (#1404). + for format in ("NETCDF4", "NETCDF4_CLASSIC", "NETCDF3_CLASSIC"): + for dim_name in ("complex", "ri", "_pfnc_complex"): + for dtype in ("f4", "f8", "i4"): + with self.subTest(format=format, dim_name=dim_name, dtype=dtype): + filename = self.tmp_path / "test_write_real.nc" + values = np.arange(12, dtype=dtype).reshape(6, 2) + with netCDF4.Dataset( + filename, "w", format=format, auto_complex=True + ) as f: + f.createDimension("x", 6) + f.createDimension(dim_name, 2) + var = f.createVariable("data", dtype, ("x", dim_name)) + self.assertEqual(var.dtype, np.dtype(dtype)) + self.assertEqual(var.dimensions, ("x", dim_name)) + self.assertEqual(var.ndim, 2) + self.assertEqual(var.shape, values.shape) + var[:] = values + var[::2, 1] = -1 + values[::2, 1] = -1 + np.testing.assert_array_equal(var[:], values) + np.testing.assert_array_equal(var[::-2, ::-1], values[::-2, ::-1]) + + with netCDF4.Dataset(filename) as f: + self.assertEqual(f["data"].dimensions, ("x", dim_name)) + np.testing.assert_array_equal(f["data"][:], values) + + # Existing floating-point layouts must still decode + # as complex when reopening with auto_complex=True. + if dtype != "i4": + with netCDF4.Dataset(filename, auto_complex=True) as f: + var = f["data"] + self.assertEqual(var.dimensions, ("x",)) + self.assertEqual(var.ndim, 1) + self.assertEqual(var.dtype, np.dtype("c8" if dtype == "f4" else "c16")) + np.testing.assert_array_equal( + var[:], values[:, 0] + 1j * values[:, 1] + ) + + def test_write_real_component_axis(self): + # A component-like dimension can also be an ordinary leading axis, + # or the only axis of a real variable. + for dimensions in (("complex", "x"), ("complex",)): + with self.subTest(dimensions=dimensions): + filename = self.tmp_path / "test_write_real_axis.nc" + shape = (2, 3) if len(dimensions) == 2 else (2,) + values = np.arange(np.prod(shape), dtype="f8").reshape(shape) + with netCDF4.Dataset(filename, "w", auto_complex=True) as f: + f.createDimension("complex", 2) + f.createDimension("x", 3) + var = f.createVariable("data", "f8", dimensions) + self.assertEqual(var.dimensions, dimensions) + self.assertEqual(var.ndim, len(dimensions)) + var[:] = values + np.testing.assert_array_equal(var[:], values) + with netCDF4.Dataset(filename) as f: + np.testing.assert_array_equal(f["data"][:], values) + + def test_write_complex_casts_real_values(self): + # Both compound and dimension storage must cast real input to complex, + # and reopened variables must keep the same complex I/O behavior. + for format in ("NETCDF4", "NETCDF3_CLASSIC"): + for dtype in ("c8", "c16"): + with self.subTest(format=format, dtype=dtype): + filename = self.tmp_path / "test_write_complex_cast.nc" + values = complex_array.astype(dtype) + with netCDF4.Dataset( + filename, "w", format=format, auto_complex=True + ) as f: + f.createDimension("x", len(values)) + var = f.createVariable("data", dtype, ("x",)) + self.assertEqual(var.dimensions, ("x",)) + var[:] = values.real + np.testing.assert_array_equal(var[:], values.real) + var[::2] = values[::2] + expected = values.real.astype(dtype) + expected[::2] = values[::2] + np.testing.assert_array_equal(var[:], expected) + with netCDF4.Dataset(filename, "r+", auto_complex=True) as f: + var = f["data"] + self.assertEqual(var.dtype, np.dtype(dtype)) + np.testing.assert_array_equal(var[:], expected) + var[:] = values + with netCDF4.Dataset(filename) as f: + raw = f["data"] + if format == "NETCDF4": + self.assertEqual(raw.dimensions, ("x",)) + np.testing.assert_array_equal(raw[:]["r"], values.real) + np.testing.assert_array_equal(raw[:]["i"], values.imag) + else: + self.assertEqual(len(raw.dimensions), 2) + self.assertEqual(raw.shape, (len(values), 2)) + np.testing.assert_array_equal(raw[:, 0], values.real) + np.testing.assert_array_equal(raw[:, 1], values.imag) + with netCDF4.Dataset(filename, auto_complex=True) as f: + np.testing.assert_array_equal(f["data"][:], values) + + def test_write_explicit_complex_struct(self): + filename = self.tmp_path / "test_write_explicit_struct.nc" + with netCDF4.Dataset(filename, "w", auto_complex=True) as f: + f.createDimension("x", len(complex_array)) + datatype = f.createCompoundType(np_dt, "nc_complex") + var = f.createVariable("data", datatype, ("x",)) + var[:] = complex_struct_array + np.testing.assert_array_equal(var[:], complex_struct_array) + with netCDF4.Dataset(filename, auto_complex=True) as f: + np.testing.assert_array_equal(f["data"][:], complex_array) + def test_read_struct(self): filename = self.tmp_path / "test_read_struct.nc"