From 82056b67a2029f0317481dd17aa34a32154e57fe Mon Sep 17 00:00:00 2001
From: Connor Ward <c.ward20@imperial.ac.uk>
Date: Fri, 10 Jan 2025 17:21:54 +0000
Subject: [PATCH] fixups and simplify gather

---
 gusto/core/io.py                              | 23 ++++---------------
 integration-tests/model/test_nc_outputting.py |  8 +++----
 2 files changed, 8 insertions(+), 23 deletions(-)

diff --git a/gusto/core/io.py b/gusto/core/io.py
index 2f53689c6..ed0dce76f 100644
--- a/gusto/core/io.py
+++ b/gusto/core/io.py
@@ -985,24 +985,9 @@ def gather_field_data(field, field_index, domain):
     if comm.size == 1:
         return field.dat.data_ro
 
-    space_name = field.function_space().name
-
+    gathered_data = comm.gather(field.dat.data_ro)
     if comm.rank == 0:
-        # Set up array to store full data in
-        global_data = np.zeros(field.function_space().dim(), dtype=field.dat.dtype)
-
-        # Store data for this processor first
-        (start, stop) = domain.coords.parallel_array_lims[space_name]
-        global_data[start:stop] = field.dat.data_ro[...]
-
-        # Receive data from other processors
-        for rank in range(1, comm.size):
-            incoming_data = comm.recv(source=rank, tag=comm.size*field_index + rank)
-            start, stop = stop, stop + incoming_data.size
-            global_data[start:stop] = incoming_data
-
+        return np.concatenate(gathered_data)
     else:
-        comm.send(field.dat.data_ro, dest=0, tag=comm.size*field_index + comm.rank)
-        global_data = None
-
-    return global_data
+        assert gathered_data is None
+        return None
diff --git a/integration-tests/model/test_nc_outputting.py b/integration-tests/model/test_nc_outputting.py
index 7a03396c2..55d63264d 100644
--- a/integration-tests/model/test_nc_outputting.py
+++ b/integration-tests/model/test_nc_outputting.py
@@ -17,11 +17,11 @@
 from pytest_mpi import parallel_assert
 
 
-def make_dirname(test_name):
+def make_dirname(test_name, suffix=""):
     if MPI.COMM_WORLD.size > 1:
-        return f'pytest_{test_name}_parallel'
+        return f"pytest_{test_name}_parallel" + suffix
     else:
-        return f'pytest_{test_name}'
+        return f"pytest_{test_name}" + suffix
 
 
 @pytest.fixture
@@ -73,7 +73,7 @@ def test_nc_outputting(geometry, domain_and_mesh_details):
     # ------------------------------------------------------------------------ #
 
     # Make sure all ranks use the same file
-    dirname = make_dirname("nc_outputting")
+    dirname = make_dirname("nc_outputting", suffix=f"_{geometry}_{MPI.COMM_WORLD.size}")
 
     domain, mesh_details = domain_and_mesh_details
     V = domain.spaces('DG')