Introduce MPI parallelization

.gitattributes

@@ -11,3 +11,4 @@ examples/greenland_mesh_v0/MESH/mesh.header filter=lfs diff=lfs merge=lfs -text
 examples/greenland_mesh_v0/MESH/mesh.elements filter=lfs diff=lfs merge=lfs -text
 examples/greenland_mesh_v0/MESH/mesh.boundary filter=lfs diff=lfs merge=lfs -text
 *.nc filter=lfs diff=lfs merge=lfs -text
+tests/data/**/*.nc !filter !diff !merge -text

.github/workflows/build-and-test.yaml

@@ -77,10 +77,12 @@ jobs:
           # Run dummy script second; EBFM on rank 0, dummy on rank 1
           mpirun -np 1 .venv/bin/ebfm --matlab-mesh examples/dem_and_mask.mat --end-time "1-Jan-1979 03:00" : \
                  -np 1 .venv/bin/python tests/core/comm/mpi_handshake_dummy.py second
-      - name: Run example on several MPI ranks
+      - name: Run examples on several MPI ranks
         if: ${{ matrix.run_mpi }}
         timeout-minutes: 5
-        run: mpirun -np 4 .venv/bin/ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading
+        run: |
+          mpirun -np 4 .venv/bin/ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading
+          mpirun -np 4 .venv/bin/ebfm --elmer-mesh examples/greenland_mesh_v0/MESH --elmer-mesh-crs-epsg 3413 --couple-to-icon-atmo --couple-to-icon-land --couple-to-elmer-ice --fake-coupling --no-shading
 
   # LFS checkout is expensive (downloads all tracked binary files).
   # Run this stage only after build-and-test has passed to avoid wasting

CHANGELOG.md

@@ -28,6 +28,13 @@ SPDX-License-Identifier: CC-BY-4.0
 * In a run on several MPI ranks, only rank 0 logs to the console at `--log-level-console`; the other ranks only log warnings and errors there, instead of repeating every message. Log files (`--log-file`, one per rank) are unchanged. The current model time is logged instead of printed. https://github.com/EBFMorg/EBFM/pull/204
 * A MATLAB grid (`--matlab-mesh`) can now run on several MPI ranks, e.g. `mpirun -np 4 ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading`: each rank computes a contiguous block of the columns, and rank 0 writes the output file, which is identical to that of a run on a single rank. Shading is not supported on several ranks and stops the run with an error. https://github.com/EBFMorg/EBFM/pull/202
 * A usage error, e.g. a rejected command line option, now ends the whole MPI job with its exit status when `MPI.COMM_WORLD` has more than one process. EBFM previously exited on its own, and the other processes, e.g. a coupled model in an MPMD run, waited for it forever. A run on a single process still exits without `MPI_ABORT`. https://github.com/EBFMorg/EBFM/pull/206
+* In a run on several MPI ranks, EBFM now checks before the time loop that every column is held by exactly one rank. The check communicates the ids of all columns; `--no-check-grid-distribution` skips it. https://github.com/EBFMorg/EBFM/pull/197
+* Partitioned Elmer meshes (`--is-partitioned-elmer-mesh`) now write model output (`Output/model_output.nc`); it was skipped before. On several MPI ranks, rank 0 gathers the cells of all ranks and writes a single file, with the cells ordered by Elmer element id and the vertices that partitions share written once. https://github.com/EBFMorg/EBFM/pull/197
+* A non-partitioned Elmer mesh (`--elmer-mesh` without `--is-partitioned-elmer-mesh`, also with `--netcdf-dem-mesh`) can now run on several MPI ranks, e.g. `mpirun -np 4 ebfm --elmer-mesh ...`: each rank computes a contiguous block of the cells, the first block on rank 0. The output file equals that of a run on a single rank. Fitting the lapse rate of potential temperature over all columns synchronizes all ranks every time step; EBFM warns about this. Every rank reads the whole mesh, which EBFM also warns about; for large meshes, a mesh with one partition per rank (`--is-partitioned-elmer-mesh`, which requires `--netcdf-dem-mesh`) avoids this. https://github.com/EBFMorg/EBFM/pull/197
+* A run on a partitioned Elmer mesh (`--is-partitioned-elmer-mesh`) whose number of MPI ranks differs from the number of partitions now stops with an error. Each rank computes one partition, so such a run previously skipped partitions without notice. `--allow-partition-mismatch` allows it, with a warning, e.g. to compute a single partition with `--use-part`. https://github.com/EBFMorg/EBFM/pull/197
+* A NetCDF grid (`--netcdf-mesh`) can now run on several MPI ranks like a MATLAB grid: each rank computes a contiguous block of the columns, and rank 0 writes the output file, which is identical to that of a run on a single rank. https://github.com/EBFMorg/EBFM/pull/197
+* `--netcdf-dem-mesh-unstructured` now works on several MPI ranks. EBFM looks up the values of the XIOS file by Elmer node id (the value of node id i at position i - 1), so each rank reads the values of its own nodes. A file without values for all node ids of the mesh is rejected with an error. https://github.com/EBFMorg/EBFM/pull/197
+* Restart files (`--restart-dir`, `--restart-init`) now work on several MPI ranks and with partitioned Elmer meshes. Rank 0 writes the columns of all ranks, ordered by column id, together with their ids in a new variable `column_id`, and reads them back for any distribution of the columns over the ranks: e.g. a file written on 4 ranks can be continued on 3 ranks or on a partitioned mesh. A file without `column_id` holds the columns in the order of a run on a single rank; on several ranks, its `x` and `y` are then checked against the grid. A file whose column ids differ from those of the grid is rejected. https://github.com/EBFMorg/EBFM/pull/197
 
 # v0.8.0
 

README.md

@@ -216,16 +216,21 @@ path/to/your/elmer/mesh
 Usage example for partitioned mesh:
 
 ```sh
-ebfm --elmer-mesh examples/greenland_mesh_v0/MESH/partitioning.128/ --netcdf-dem-mesh examples/BedMachineGreenland-v5_lo.nc --is-partitioned-elmer-mesh --use-part 42 --elmer-mesh-crs-epsg 3413
+ebfm --elmer-mesh examples/greenland_mesh_v0/MESH/partitioning.128/ --netcdf-dem-mesh examples/BedMachineGreenland-v5_lo.nc --is-partitioned-elmer-mesh --use-part 42 --elmer-mesh-crs-epsg 3413 --allow-partition-mismatch
 ```
 
 On several MPI ranks, `--use-part` is not allowed: rank i computes the partition i + 1, e.g. the partitions 1 to 4 with:
 
 ```sh
-mpirun -n 4 ebfm --elmer-mesh examples/greenland_mesh_v0/MESH/partitioning.128/ --netcdf-dem-mesh examples/BedMachineGreenland-v5_lo.nc --is-partitioned-elmer-mesh --elmer-mesh-crs-epsg 3413
+mpirun -n 4 ebfm --elmer-mesh examples/greenland_mesh_v0/MESH/partitioning.128/ --netcdf-dem-mesh examples/BedMachineGreenland-v5_lo.nc --is-partitioned-elmer-mesh --elmer-mesh-crs-epsg 3413 --allow-partition-mismatch
 ```
 
-Besides a partitioned mesh, only a MATLAB grid runs on several MPI ranks (see [Running on several MPI ranks](#running-on-several-mpi-ranks)).
+Rank 0 writes the output of all ranks to a single file, with the cells ordered by their Elmer element id.
+
+Besides a partitioned mesh, MATLAB and NetCDF grids and a non-partitioned Elmer mesh run on several MPI ranks (see [Running on several MPI ranks](#running-on-several-mpi-ranks)).
+
+A run needs one MPI rank per partition. To compute only some partitions, e.g. a
+single one with `--use-part` as above, pass `--allow-partition-mismatch`.
 
 Note: The partitioned mesh is not included in this repository.
 
@@ -286,20 +291,32 @@ ebfm --matlab-mesh examples/dem_and_mask.mat --restart-dir examples/restart --re
 
 Note that it is necessary to provide a `--start-time` consistent with the given restart file.
 
+A restart file holds the columns of the whole grid together with their ids
+(`column_id`), so a run may continue from it on any number of MPI ranks and on
+a partitioned mesh. A restart file without `column_id`, written by an older
+version of EBFM, holds the columns in the order of a run on a single rank.
+
 ### Running on several MPI ranks
 
-EBFM splits a MATLAB grid over the MPI ranks of a run (requires `pip install 'ebfm[mpi]'`):
+EBFM splits the grid over the MPI ranks of a run (requires `pip install 'ebfm[mpi]'`):
 
 ```sh
 mpirun -np 4 ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading
 ```
 
-* Each rank computes a contiguous block of the columns, the first block on rank 0.
-* Rank 0 writes the output file and `--dump-reference` for all ranks. They are
-  identical to those of a run on a single rank.
-* Not supported on several ranks, and rejected with an error: shading,
-  `--forcing-dir`, `--restart-dir`/`--restart-init`, and grids other than a
-  MATLAB grid or a partitioned Elmer mesh.
+* MATLAB and NetCDF grids and Elmer meshes are split into contiguous blocks of
+  columns, the first block on rank 0. A partitioned Elmer mesh gives one
+  partition to each rank (see above).
+* With a non-partitioned Elmer mesh, every rank reads the whole mesh, and EBFM
+  warns about it. For large meshes, use a partitioned mesh (requires
+  `--netcdf-dem-mesh`).
+* Rank 0 writes the output file, the restart file and `--dump-reference` for
+  all ranks, and reads the restart file. They are identical to those of a run
+  on a single rank.
+* Not supported on several ranks, and rejected with an error: shading and
+  `--forcing-dir`.
+* `--no-check-grid-distribution` skips the start-up check that every column is
+  held by exactly one rank.
 
 ### Performance and Profiling Runs
 
@@ -335,7 +352,7 @@ Follow the install instructions from above and run the example command for a cou
 mpirun -np 1 ebfm \
   --elmer-mesh $MESHES/greenland_mesh_v0/MESH/partitioning.128/ \
   --netcdf-dem-mesh $DATA/BedMachineGreenland-v5.nc \
-  --is-partitioned-elmer-mesh --use-part 1 \
+  --is-partitioned-elmer-mesh --use-part 1 --allow-partition-mismatch \
   --coupler-config $CPL_CONFIG --couple-to-elmer-ice --couple-to-icon-atmo \
   : \
   -np 1 $ELMER_ROOT/src/elmer_dummy_f.x $MESHES/greenland_mesh_v0/MESH/partitioning.128 1 $CPL_CONFIG \

examples/README.md

@@ -53,10 +53,12 @@ EBFM reads two variables from the file:
 * `zs`: surface elevation (m)
 * `h`: ice thickness (m)
 
-Both have the dimensions `(time, nmesh2D_node)` with a single time step and hold one value per mesh node, in the order
-of `mesh.nodes`. EBFM averages the three node values of each element to obtain the elevation and ice thickness of the
-corresponding column and treats columns with more than 1 m of ice as glacier. The file also contains the node
-coordinates `x` and `y` (EPSG:3413), which agree with `mesh.nodes` and allow to check the node order.
+Both have the dimensions `(time, nmesh2D_node)` with a single time step and hold one value per mesh node, stored by
+Elmer node id: the value of node id `i` is at position `i - 1`. In `greenland_mesh_v0` the node ids run from 1 in the
+order of `mesh.nodes`, so this is also the order of `mesh.nodes`. EBFM averages the three node values of each element to
+obtain the elevation and ice thickness of the corresponding column and treats columns with more than 1 m of ice as
+glacier. The file also contains the node coordinates `x` and `y` (EPSG:3413), which agree with `mesh.nodes` and allow to
+check the node order.
 
 The original file was written by Elmer/Ice v9.0 (Rev: c4a78ac59) with the `XIOSOutputSolver`. The relevant part of
 the XIOS context is:

src/ebfm/core/FINAL_create_restart_file.py

@@ -9,55 +9,25 @@
 
 import logging
 
-from .restart import RESTART_VARIABLES
+from .comm.gather import ColumnGather
+from .restart import COLUMN_ID_VARIABLE, RESTART_VARIABLES
 
 logger = logging.getLogger(__name__)
 
 
-def main(OUT, io, restartdir: Path):
-    def create_boot_file():
-        """
-        Create a restart file (boot file) and save it as a NetCDF file.
+def main(OUT, io, restartdir: Path, gather: ColumnGather):
+    """
+    Write the restart file (boot file) if io["writebootfile"] is set.
 
-        Parameters:
-        - OUT: Dictionary containing the data to be saved.
-        @param[in] restartdir Path to the directory containing restart files (if initializing from restart file).
-        Returns:
-        - None
-        """
-
-        # Create a reboot directory if it does not exist
-        if restartdir:
-            os.makedirs(restartdir, exist_ok=True)
-
-        # Check if we should write the boot file
-        if io["writebootfile"]:
-            # Define the output NetCDF file path
-            # Create a new NetCDF file to store the boot variables
-            with Dataset(io["bootfileout"], "w", format="NETCDF4") as ncfile:
-                # Save each variable in the OUT dictionary to the NetCDF file
-                for var_name, var_data in OUT.items():
-                    # Handle different variable dimensions
-                    if isinstance(var_data, np.ndarray):  # If it's a NumPy array
-                        # Create the appropriate dimension(s) if not already defined
-                        dims = []
-                        for dimsize in var_data.shape:
-                            dim_name = f"{var_name}_dim{len(dims)}"
-                            if dim_name not in ncfile.dimensions:
-                                ncfile.createDimension(dim_name, dimsize)
-                            dims.append(dim_name)
-
-                        # Create the variable
-                        nc_var = ncfile.createVariable(var_name, var_data.dtype, tuple(dims))
-                        nc_var[:] = var_data  # Write the data
-                    elif np.isscalar(var_data):  # If it's a scalar, store it as a 0D variable
-                        nc_var = ncfile.createVariable(var_name, type(var_data), ())
-                        nc_var.assignValue(var_data)
-                    else:
-                        raise ValueError(f"Unsupported data type for variable: {var_name}")
-
-            logger.info(f"Boot file saved to {io['bootfileout']}")
+    Collective over gather.comm: the root rank gathers the columns of all ranks one variable at a time and writes
+    them, ordered by column id, together with their ids. Sets OUT["timelastsnow_netCDF"] on every rank, also when no
+    restart file is written.
 
+    @param[in] OUT model state of this rank's columns
+    @param[in] io I/O settings: writebootfile and bootfileout
+    @param[in] restartdir directory of the restart files, created if it does not exist
+    @param[in] gather the columns of all ranks
+    """
     # float64 even if all times are whole days, for which date2num returns integers
     OUT["timelastsnow_netCDF"] = np.asarray(
         date2num(
@@ -67,7 +37,42 @@ def create_boot_file():
         ),
         dtype=np.float64,
     )
-    OUT = {var_name: OUT[var_name] for var_name in RESTART_VARIABLES}
 
-    # Create the boot file
-    create_boot_file()
+    if not io["writebootfile"]:
+        return
+
+    ncfile = None
+    if gather.is_root:
+        if restartdir:
+            os.makedirs(restartdir, exist_ok=True)
+        ncfile = Dataset(io["bootfileout"], "w", format="NETCDF4")
+
+    try:
+        for var_name in RESTART_VARIABLES:
+            values = gather.gather(OUT[var_name])
+            if ncfile is not None:
+                _write_variable(ncfile, var_name, values)
+        if ncfile is not None:
+            _write_variable(ncfile, COLUMN_ID_VARIABLE, gather.column_ids)
+    finally:
+        if ncfile is not None:
+            ncfile.close()
+
+    if gather.is_root:
+        logger.info(f"Boot file saved to {io['bootfileout']}")
+
+
+def _write_variable(ncfile: Dataset, var_name: str, var_data: np.ndarray) -> None:
+    """Write an array to the open restart file, with one dimension of its own per axis."""
+    if not isinstance(var_data, np.ndarray):
+        raise ValueError(f"Unsupported data type for variable: {var_name}")
+
+    dims = []
+    for dimsize in var_data.shape:
+        dim_name = f"{var_name}_dim{len(dims)}"
+        if dim_name not in ncfile.dimensions:
+            ncfile.createDimension(dim_name, dimsize)
+        dims.append(dim_name)
+
+    nc_var = ncfile.createVariable(var_name, var_data.dtype, tuple(dims))
+    nc_var[:] = var_data