Keep explicit real variable shapes with auto_complex=True

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).
 

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,

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"