Fix coeff-space outputs and add tests

Author: kburns

dedalus/core/distributor.py

@@ -320,6 +320,12 @@ def local_chunks(self, domain, scales, rank=None, broadcast=False):
                 local_chunks.append(np.arange(start, end))
         return tuple(local_chunks)
 
+    def global_elements(self, domain, scales):
+        """Global element indices by axis."""
+        global_shape = self.global_shape(domain, scales)
+        indices = [np.arange(n) for n in global_shape]
+        return tuple(indices)
+
     def local_elements(self, domain, scales, rank=None, broadcast=False):
         """Local element indices by axis."""
         chunk_shape = self.chunk_shape(domain)
@@ -345,12 +351,7 @@ def valid_elements(self, tensorsig, domain, scales, rank=None, broadcast=False):
             valid &= basis.valid_elements(tensorsig, grid_space[basis_axes], elements[basis_axes])
         return valid
 
-    @CachedMethod
-    def local_group_arrays(self, domain, scales, rank=None, broadcast=False):
-        """Dense array of local groups (first axis)."""
-        # Make dense array of local elements
-        elements = self.local_elements(domain, scales, rank=rank, broadcast=broadcast)
-        elements = np.array(np.meshgrid(*elements, indexing='ij'))
+    def _group_arrays(self, elements, domain):
         # Convert to groups basis-by-basis
         grid_space = self.grid_space
         groups = np.zeros_like(elements)
@@ -360,6 +361,22 @@ def local_group_arrays(self, domain, scales, rank=None, broadcast=False):
             groups[basis_axes] = basis.elements_to_groups(grid_space[basis_axes], elements[basis_axes])
         return groups
 
+    @CachedMethod
+    def local_group_arrays(self, domain, scales, rank=None, broadcast=False):
+        """Dense array of local groups (first axis)."""
+        # Make dense array of local elements
+        elements = self.local_elements(domain, scales, rank=rank, broadcast=broadcast)
+        elements = np.array(np.meshgrid(*elements, indexing='ij'))
+        return self._group_arrays(elements, domain)
+
+    @CachedMethod
+    def global_group_arrays(self, domain, scales):
+        """Dense array of local groups (first axis)."""
+        # Make dense array of local elements
+        elements = self.global_elements(domain, scales)
+        elements = np.array(np.meshgrid(*elements, indexing='ij'))
+        return self._group_arrays(elements, domain)
+
     @CachedMethod
     def local_groupsets(self, group_coupling, domain, scales, rank=None, broadcast=False):
         local_groupsets = self.local_group_arrays(domain, scales, rank=rank, broadcast=broadcast).astype(object)

dedalus/core/evaluator.py

@@ -647,10 +647,9 @@ def dset_metadata(self, task, task_num, dset, scale_group, gnc_shape, gnc_start,
                 else:
                     sn = 'k' + basis.coordsystem.coords[subaxis].name
                     if virtual_file:
-                        data = basis.global_elements()[subaxis].ravel()
+                        data = layout.global_group_arrays(op.domain, scales)[subaxis]
                     else:
-                        data = basis.local_elements()[subaxis].ravel()
-
+                        data = layout.local_group_arrays(op.domain, scales)[subaxis]
 
                 if self.dist.comm_cart.rank == 0:
                     scale_hash = hashlib.sha1(data).hexdigest()

dedalus/tests/test_output.py

@@ -13,7 +13,8 @@
 @pytest.mark.parametrize('dealias', [1, 3/2])
 @pytest.mark.parametrize('output_scales', [1, 3/2, 2,
     pytest.param(1/2, marks=pytest.mark.xfail(reason="evaluator not copying correctly for scales < 1"))])
-def test_cartesian_output(dtype, dealias, output_scales):
+@pytest.mark.parametrize('output_layout', ['g', 'c'])
+def test_cartesian_output(dtype, dealias, output_scales, output_layout):
     Nx = Ny = Nz = 16
     Lx = Ly = Lz = 2 * np.pi
     # Bases
@@ -41,17 +42,18 @@ def test_cartesian_output(dtype, dealias, output_scales):
         tempdir = pathlib.Path(tempdir).stem
         output = solver.evaluator.add_file_handler(tempdir, iter=1)
         for task in tasks:
-            output.add_task(task, layout='g', name=str(task), scales=output_scales)
+            output.add_task(task, layout=output_layout, name=str(task), scales=output_scales)
         solver.evaluator.evaluate_handlers([output])
+        output.process_virtual_file()
         # Check solution
         #post.merge_process_files('test_output')
         errors = []
-        with h5py.File(f'{tempdir}/{tempdir}_s1/{tempdir}_s1_p0.h5', mode='r') as file:
+        with h5py.File(f'{tempdir}/{tempdir}_s1.h5', mode='r') as file:
             for task in tasks:
                 task_saved = file['tasks'][str(task)][-1]
                 task = task.evaluate()
                 task.change_scales(output_scales)
-                errors.append(np.max(np.abs(task['g'] - task_saved)))
+                errors.append(np.max(np.abs(task[output_layout] - task_saved)))
     assert np.allclose(errors, 0)
 
 
@@ -107,10 +109,11 @@ def test_spherical_output(Nphi, Ntheta, Nr, k, dealias, dtype, basis, output_sca
         for task in tasks:
             output.add_task(task, layout='g', name=str(task), scales=output_scales)
         solver.evaluator.evaluate_handlers([output])
+        output.process_virtual_file()
         # Check solution
         #post.merge_process_files('test_output')
         errors = []
-        with h5py.File(f'{tempdir}/{tempdir}_s1/{tempdir}_s1_p0.h5', mode='r') as file:
+        with h5py.File(f'{tempdir}/{tempdir}_s1.h5', mode='r') as file:
             for task in tasks:
                 task_saved = file['tasks'][str(task)][-1]
                 task = task.evaluate()