Merge pull request #630 from bccp/fix-complex-mesh-save

Fix wrong shape of complex mesh save.
bccp · Aug 19, 2020 · 9b53312 · 9b53312
2 parents ba9b817 + 49beb2e
commit 9b53312
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/nbodykit/base/mesh.py b/nbodykit/base/mesh.py
@@ -390,11 +390,11 @@ def save(self, output, dataset='Field', mode='real'):
             field.ravel(out=data)
             with ff.create_from_array(dataset, data) as bb:
                 if isinstance(field, RealField):
-                    bb.attrs['ndarray.shape'] = field.pm.Nmesh
+                    bb.attrs['ndarray.shape'] = field.cshape
                     bb.attrs['BoxSize'] = field.pm.BoxSize
                     bb.attrs['Nmesh'] = field.pm.Nmesh
                 elif isinstance(field, BaseComplexField):
-                    bb.attrs['ndarray.shape'] = field.Nmesh, field.Nmesh, field.Nmesh // 2 + 1
+                    bb.attrs['ndarray.shape'] = field.cshape
                     bb.attrs['BoxSize'] = field.pm.BoxSize
                     bb.attrs['Nmesh'] = field.pm.Nmesh
 

diff --git a/nbodykit/base/tests/test_mesh.py b/nbodykit/base/tests/test_mesh.py
@@ -75,7 +75,7 @@ def test_real_save(comm):
         shutil.rmtree(tmpfile)
 
 @MPITest([1,4])
-def test_real_save(comm):
+def test_complex_save(comm):
 
     cosmo = cosmology.Planck15
 
@@ -94,7 +94,7 @@ def test_real_save(comm):
     source.attrs['empty'] = None
 
     # save to bigfile
-    source.save(tmpfile, mode='real')
+    source.save(tmpfile, mode='complex')
 
     # load as a BigFileMesh
     source2 = BigFileMesh(tmpfile, dataset='Field', comm=comm)
@@ -104,7 +104,7 @@ def test_real_save(comm):
         assert_array_equal(source2.attrs[k], source.attrs[k])
 
     # check data
-    assert_array_equal(source2.compute(mode='real'), source.compute(mode='real'))
+    assert_array_equal(source2.compute(mode='complex'), source.compute(mode='complex'))
 
     # cleanup
     comm.barrier()