Read a NetCDF DEM in bands instead of as a whole

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

@@ -11,6 +11,7 @@ jobs:
   build-and-test:
     name: Build and test (${{ matrix.variant }})
     runs-on: ubuntu-latest
+    timeout-minutes: 20
     strategy:
       matrix:
         include:
@@ -38,6 +39,9 @@ jobs:
           # several ranks report the same rank number -- which is exactly what the handshake step
           # below checks.
           mpi: openmpi
+      - name: Keep Open MPI's TCP traffic on the loopback interface
+        if: ${{ matrix.run_mpi }}
+        run: echo "btl_tcp_if_include=lo" >> ~/.openmpi/mca-params.conf
       - name: Create virtual environment and install EBFM
         run: |
           python3 -m venv .venv
@@ -65,13 +69,18 @@ jobs:
         run: .venv/bin/ebfm --matlab-mesh examples/dem_and_mask.mat --couple-to-icon-atmo --couple-to-icon-land --fake-coupling --no-shading
       - name: Run example with MPI handshake
         if: ${{ matrix.run_mpi }}
+        timeout-minutes: 5
         run: |
           # Run dummy script first; dummy on rank 0, EBFM on rank 1
           mpirun -np 1 .venv/bin/python tests/core/comm/mpi_handshake_dummy.py first : \
                  -np 1 .venv/bin/ebfm --matlab-mesh examples/dem_and_mask.mat --end-time "1-Jan-1979 03:00"
           # 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
+        if: ${{ matrix.run_mpi }}
+        timeout-minutes: 5
+        run: mpirun -np 4 .venv/bin/ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading
 
   # LFS checkout is expensive (downloads all tracked binary files).
   # Run this stage only after build-and-test has passed to avoid wasting

.github/workflows/test.yml

@@ -68,9 +68,12 @@ jobs:
         # Not the default (mpich): the MPICH package on Ubuntu 24.04 is broken
         # (https://bugs.launchpad.net/ubuntu/+source/mpich/+bug/2072338) and setup-mpi only patches
         # around it by copying newer libmpich12/libucx0 over the apt install. The tests are
-        # single-rank and implementation-agnostic, so use the one that works out of the box.
+        # implementation-agnostic, so use the one that works out of the box.
         mpi: openmpi
 
+    - name: Keep Open MPI's TCP traffic on the loopback interface
+      run: echo "btl_tcp_if_include=lo" >> ~/.openmpi/mca-params.conf
+
     - name: Install tox
       run: |
         python -m pip install --upgrade pip

CHANGELOG.md

@@ -10,6 +10,7 @@ SPDX-License-Identifier: CC-BY-4.0
     * Add feature to use NetCDF meshes with DEM as input (`--netcdf-mesh`)
     * Add option to set missing elevations of NetCDF meshes to 0 m (`--fill-missing-elevation`)
     * Add feature to read forcing files from folder (`--forcing-dir`)
+* A NetCDF DEM (`--netcdf-dem-mesh`) is now read in bands of at most 16 MB around the sampled points instead of as a whole. With the full-resolution `BedMachineGreenland-v5.nc` this lowers the peak memory of grid initialization from about 1.1 GB to about 0.3 GB, on every MPI rank of a run on a partitioned mesh. https://github.com/EBFMorg/EBFM/pull/200
 * Remove support for Python 3.10. https://github.com/EBFMorg/EBFM/pull/143
 * `--numba-threads` is now checked against the number of threads Numba accepts, i.e. the CPUs the rank may run on (CPU affinity from MPI binding or the Slurm allocation, e.g. `--cpus-per-task`) or the `NUMBA_NUM_THREADS` environment variable. Exceeding it stops EBFM with an error explaining the cause, instead of a Numba `ValueError` traceback. Exceeding the suggested thread count for the ranks sharing a node is now only a warning; the suggestion no longer divides the CPUs of the whole node by the number of EBFM ranks. https://github.com/EBFMorg/EBFM/pull/190
 * A fatal error now ends EBFM with exit status 1 whenever mpi4py is installed, both standalone and under `mpirun`. EBFM previously aborted MPI with error code 0, so the run, and `mpirun` with it, reported success after logging "Fatal Error in EBFM". The same applies to an MPI handshake version mismatch. A fatal error before the MPI communicator splitting no longer leaves the other components of an MPMD run waiting forever. https://github.com/EBFMorg/EBFM/pull/191
@@ -26,6 +27,7 @@ SPDX-License-Identifier: CC-BY-4.0
     * `--dump-reference` now holds the columns of all ranks, ordered by Elmer element id. Every rank previously wrote its own partition to the same file.
     * `--print-diagnostics` now reports the column count and `smb` statistics of the whole domain, logged by rank 0. Each rank previously reported its own partition.
 * 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
 
 # v0.8.0
 

README.md

@@ -225,7 +225,7 @@ On several MPI ranks, `--use-part` is not allowed: rank i computes the partition
 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
 ```
 
-A partitioned mesh is the only grid EBFM runs on several MPI ranks.
+Besides a partitioned mesh, only a MATLAB grid runs on several MPI ranks (see [Running on several MPI ranks](#running-on-several-mpi-ranks)).
 
 Note: The partitioned mesh is not included in this repository.
 
@@ -286,6 +286,21 @@ 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.
 
+### Running on several MPI ranks
+
+EBFM splits a MATLAB 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.
+
 ### Performance and Profiling Runs
 
 Multi-threaded Numba kernels (`--with-numba`), GPU offloading (`--with-gpu`), timing, output comparison and

src/ebfm/core/INIT.py

@@ -198,6 +198,7 @@ def init_grid(grid: GridDict, io, config: GridConfig, comm):
         config (GridConfig): Grid configuration object containing settings for grid initialization.
         comm: EBFM's MPI communicator, whose ranks hold the columns of the domain.
     """
+    grid["is_parallel_run"] = comm.size > 1
     grid["is_partitioned"] = config.is_partitioned
     grid["is_unstructured"] = config.is_unstructured
     grid["has_shading"] = config.use_shading
@@ -320,16 +321,20 @@ def init_grid(grid: GridDict, io, config: GridConfig, comm):
             grid["lon_2D"] = input_data["lon"]
             grid["lat_2D"] = input_data["lat"]
 
-        # Create 1-D mask
-        grid["mask"] = mask_2D[mask_2D == 1]
+        # Columns as flat indices into the 2-D fields; in a run on several MPI ranks a contiguous block of them,
+        # the first block on rank 0. The 2-D fields stay whole: the output writer places all columns into them.
+        ind = np.flatnonzero(mask_2D == 1)
+        if config.num_partitions > 1:
+            assert 1 <= config.partition_id <= config.num_partitions, f"Invalid partition id {config.partition_id}."
+            ind = np.array_split(ind, config.num_partitions)[config.partition_id - 1]
+        grid["ind"] = (ind,)
+        grid["column_id"] = ind
 
         # Store 1-D (vectorized) grid information
-        mask_flat = mask_2D.flatten()
-        grid["x"] = grid["x_2D"].flatten()[mask_flat == 1]
-        grid["y"] = grid["y_2D"].flatten()[mask_flat == 1]
-        grid["z"] = grid["z_2D"].flatten()[mask_flat == 1]
-        grid["ind"] = np.where(mask_flat == 1)
-        grid["column_id"] = grid["ind"][0]
+        grid["mask"] = mask_2D.flat[ind]
+        grid["x"] = grid["x_2D"].flat[ind]
+        grid["y"] = grid["y_2D"].flat[ind]
+        grid["z"] = grid["z_2D"].flat[ind]
 
         # ---------------------------------------------------------------------
         # Grid slope and aspect
@@ -351,12 +356,12 @@ def init_grid(grid: GridDict, io, config: GridConfig, comm):
         grid["slope_beta_2D"] = np.arctan(grid["slope"])
 
         # Convert 2-D fields to 1-D vectors
-        grid["slope"] = grid["slope"].flatten()[mask_flat == 1]
-        grid["slope_x"] = grid["slope_x"].flatten()[mask_flat == 1]
-        grid["slope_y"] = grid["slope_y"].flatten()[mask_flat == 1]
-        grid["aspect"] = grid["aspect"].flatten()[mask_flat == 1]
-        grid["lat"] = grid["lat_2D"].flatten()[mask_flat == 1]
-        grid["lon"] = grid["lon_2D"].flatten()[mask_flat == 1]
+        grid["slope"] = grid["slope"].flat[ind]
+        grid["slope_x"] = grid["slope_x"].flat[ind]
+        grid["slope_y"] = grid["slope_y"].flat[ind]
+        grid["aspect"] = grid["aspect"].flat[ind]
+        grid["lat"] = grid["lat_2D"].flat[ind]
+        grid["lon"] = grid["lon_2D"].flat[ind]
 
         # four individual asserts to avoid overhead when turning on optimization
         assert np.isfinite(grid["slope"]).all(), (
@@ -432,6 +437,9 @@ def _precompute_shading_matlab(grid: GridDict, mask_2D) -> None:
                         `az_array` and `maxgridangle`.
     @param[in] mask_2D 2-D glacier mask, 1 for glacier cells.
     """
+    # The walk starts from every glacier cell of the 2-D grid, not only from the columns of this rank.
+    assert not grid["is_parallel_run"], "The shading look-up table cannot be computed in a run on several MPI ranks."
+
     logger.warning("Using EBFM's internal shading method. Should only be used on regular grids!")
 
     grid["shading_method"] = ShadingMethod.LUT  # shading based on look-up table (lut)

src/ebfm/core/LOOP_general_functions.py

@@ -21,6 +21,20 @@ def is_first_time_step(t: int) -> bool:
     return t == 0
 
 
+def is_output_time_step(t: int, io: dict) -> bool:
+    """
+    Check if the current time step is a time step where output has to be written
+
+    Parameters:
+        t (int): Current time step (0-based index).
+        io (dict): Dictionary with I/O settings (e.g. bootfilein, bootfileout, homedir).
+
+    Returns:
+        bool: True if the current time step is output time step, False otherwise.
+    """
+    return (t + 1) % io["freqout"] == 0
+
+
 def is_final_time_step(t: int, time: dict) -> bool:
     """
     Check if the current time step is the final time step.

src/ebfm/core/LOOP_write_to_file.py

@@ -3,14 +3,36 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import os
+
 import numpy as np
 from netCDF4 import Dataset, date2num
 
-from .LOOP_general_functions import is_first_time_step, is_final_time_step
+from .LOOP_general_functions import is_first_time_step, is_output_time_step, is_final_time_step
+from .comm.gather import ColumnGather
 
 _FILL_VALUE = -9999.0
 
 
+def _gather_grid(gather: ColumnGather, grid):
+    """
+    Grid of the columns of all ranks, on the root rank, as far as the output file needs it.
+
+    @returns the grid on the root rank, None on all other ranks
+    """
+    if gather.comm.size == 1:
+        return grid
+
+    assert not grid["is_unstructured"], "Gathering the output of an unstructured grid is not supported."
+
+    # Every rank holds the full 2-D fields. The column ids of a structured grid are the flat indices of the columns
+    # in them, so the sorted ids of all ranks are the positions of all columns.
+    output_grid = {**grid, "ind": gather.column_ids}
+
+    if not gather.is_root:
+        return None
+    return output_grid
+
+
 def _write_unstructured_grid(nc_file, grid):
     """Write unstructured mesh coordinates and topology to an open NetCDF file."""
     n_cells = grid["x"].shape[0]
@@ -140,7 +162,13 @@ def _write_output_variable(nc_file, varname, var_1D, time_index, grid, column):
         nc_file[varname][time_index, :, :] = var_2D
 
 
-def main(OUTFILE, io, OUT, grid, t, time, column):
+def main(OUTFILE, io, OUT, grid, t, time, column, gather: ColumnGather):
+    """
+    Accumulate the output variables of this time step and write them to file every io["freqout"] steps.
+
+    Collective over gather.comm: each rank accumulates its own columns, the root rank gathers the columns of
+    all ranks and writes them to a single file.
+    """
     # Specify variables to be written
     if is_first_time_step(t):
         OUTFILE["varsout"] = [
@@ -235,7 +263,7 @@ def save_binary_files():
         """
 
         # Save output to binary files at the first time step
-        if is_first_time_step(t):
+        if is_first_time_step(t) and gather.is_root:
             if not os.path.exists(io["outdir"]):
                 os.makedirs(io["outdir"])
 
@@ -246,14 +274,18 @@ def save_binary_files():
                 io["fid"][varname] = open(filepath, "wb")
 
         # Write variables to binary files when `freqout` is met
-        if (t + 1) % io["freqout"] == 0:  # TODO: correct?
+        if is_output_time_step(t, io):  # TODO: correct?
+            for entry in OUTFILE["varsout"]:
+                varname = entry[0]
+                OUTFILE[varname] = gather.gather(OUTFILE["TEMP"][varname])
+
+        if is_output_time_step(t, io) and gather.is_root:
             for entry in OUTFILE["varsout"]:
                 varname = entry[0]
-                OUTFILE[varname] = OUTFILE["TEMP"][varname]
                 io["fid"][varname].write(OUTFILE[varname].astype("float32").tobytes())
 
         # Close all file streams and save run metadata at the final time step
-        if is_final_time_step(t, time):
+        if is_final_time_step(t, time) and gather.is_root:
             for file in io["fid"].values():
                 file.close()
 
@@ -275,8 +307,12 @@ def save_netCDF_file():
         time_units = "days since 1970-01-01 00:00:00"
         time_calendar = "gregorian"
 
-        # Initialize NetCDF file at the first time step
+        # The root rank defines the file from the grid of all ranks
         if is_first_time_step(t):
+            io["output_grid"] = _gather_grid(gather, grid)
+
+        # Initialize NetCDF file at the first time step
+        if is_first_time_step(t) and gather.is_root:
             if not os.path.exists(io["outdir"]):
                 os.makedirs(io["outdir"])
 
@@ -285,10 +321,10 @@ def save_netCDF_file():
             io["nc_file"] = Dataset(nc_filepath, "w", format="NETCDF4")
             io["nc_file"].createDimension("time", None)  # Unlimited time dimension
 
-            if grid["is_unstructured"]:
-                _write_unstructured_grid(io["nc_file"], grid)
+            if io["output_grid"]["is_unstructured"]:
+                _write_unstructured_grid(io["nc_file"], io["output_grid"])
             else:
-                _write_structured_grid(io["nc_file"], grid)
+                _write_structured_grid(io["nc_file"], io["output_grid"])
 
             io["nc_file"].createDimension("nl", column.nl)  # Vertical layers for `sub` variables
 
@@ -298,7 +334,7 @@ def save_netCDF_file():
                 var_units = entry[1]
                 var_desc = entry[3]
 
-                dimensions, chunksizes = _output_dimensions(varname, grid, column)
+                dimensions, chunksizes = _output_dimensions(varname, io["output_grid"], column)
 
                 nc_var = io["nc_file"].createVariable(
                     varname=varname,
@@ -321,23 +357,25 @@ def save_netCDF_file():
             nc_time.description = "Time at which data is recorded, in days since 1970-01-01 00:00:00"
 
         # Write data to NetCDF at the specified frequency (e.g., for average/summed values)
-        if (t + 1) % io["freqout"] == 0:  # TODO: correct?
+        if is_output_time_step(t, io):  # TODO: correct?
             time_index = t // io["freqout"]  # Determine the time slice index
 
             # Calculate time in "days since 1970-01-01"
             time_days_since_1970 = date2num(time["TCUR"], units=time_units, calendar=time_calendar)
 
             # Write time variable
-            io["nc_file"]["time"][time_index] = time_days_since_1970
+            if gather.is_root:
+                io["nc_file"]["time"][time_index] = time_days_since_1970
 
-            # Write variables to NetCDF
+            # Write variables to NetCDF, gathering one variable at a time to limit the memory needed on the root rank
             for entry in OUTFILE["varsout"]:
                 varname = entry[0]
-                var_1D = OUTFILE["TEMP"][varname]
-                _write_output_variable(io["nc_file"], varname, var_1D, time_index, grid, column)
+                var_1D = gather.gather(OUTFILE["TEMP"][varname])
+                if gather.is_root:
+                    _write_output_variable(io["nc_file"], varname, var_1D, time_index, io["output_grid"], column)
 
         # Close the NetCDF file at the final time step
-        if is_final_time_step(t, time):
+        if is_final_time_step(t, time) and gather.is_root:
             io["nc_file"].close()
 
         return True