From 2a2dd508df213064276886f8dc2d04b97c1451c6 Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Mon, 19 Oct 2026 19:14:01 +0200 Subject: [PATCH 01/13] Remove unnecessary arg. --- src/ebfm/core/INIT.py | 3 +-- src/ebfm/main.py | 2 +- tests/core/test_init.py | 8 ++++---- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/src/ebfm/core/INIT.py b/src/ebfm/core/INIT.py index 99433bfa..0d4834e1 100644 --- a/src/ebfm/core/INIT.py +++ b/src/ebfm/core/INIT.py @@ -183,7 +183,7 @@ def init_constants(): return C -def init_grid(grid: GridDict, io, config: GridConfig, comm): +def init_grid(grid: GridDict, config: GridConfig, comm): """ Initializes the GridDict based on the given GridConfig. @@ -194,7 +194,6 @@ def init_grid(grid: GridDict, io, config: GridConfig, comm): Parameters: grid (GridDict): Dictionary to store grid-related parameters. - io (dict): Dictionary with I/O settings (e.g. bootfilein, bootfileout, homedir). config (GridConfig): Grid configuration object containing settings for grid initialization. comm: EBFM's MPI communicator, whose ranks hold the columns of the domain. """ diff --git a/src/ebfm/main.py b/src/ebfm/main.py index e2ab2941..999c64b0 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -231,7 +231,7 @@ def _main_impl(): set_random_seed(logger, args.random_seed, ebfm_comm) C = INIT.init_constants() - grid = INIT.init_grid(grid, io, grid_config, ebfm_comm) + grid = INIT.init_grid(grid, grid_config, ebfm_comm) # Rank 0 collects the columns of all ranks, e.g. to write them to the output file or for --dump-reference column_gather = ColumnGather(ebfm_comm, grid["column_id"]) diff --git a/tests/core/test_init.py b/tests/core/test_init.py index 9381f61c..ddd45ea6 100644 --- a/tests/core/test_init.py +++ b/tests/core/test_init.py @@ -352,8 +352,8 @@ def tearDown(self): def _init_grid(self, shading_flag): args = parse_cli_args(["--matlab-mesh", str(_MATLAB_MESH), shading_flag]) grid_config = GridConfig(args) - grid, io, _, _ = init_config(TimeConfig(args), grid_config, Path(self.tmpdir.name), False) - return init_grid(grid, io, grid_config, defaultComm) + grid, _, _, _ = init_config(TimeConfig(args), grid_config, Path(self.tmpdir.name), False) + return init_grid(grid, grid_config, defaultComm) def test_shading_disabled_skips_the_lookup_table(self): grid = self._init_grid("--no-shading") @@ -390,8 +390,8 @@ def tearDown(self): def _init_grid(self, num_ranks, rank): args = parse_cli_args(["--matlab-mesh", str(_MATLAB_MESH), "--no-shading"]) grid_config = GridConfig(args, num_ranks, rank) - grid, io, _, _ = init_config(TimeConfig(args), grid_config, Path(self.tmpdir.name), False) - return init_grid(grid, io, grid_config, defaultComm) + grid, _, _, _ = init_config(TimeConfig(args), grid_config, Path(self.tmpdir.name), False) + return init_grid(grid, grid_config, defaultComm) def test_blocks_of_all_ranks_make_up_the_serial_grid(self): serial = self._init_grid(num_ranks=1, rank=0) From c349660328c413dee82379eaa5d62e9e951d0d46 Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Tue, 20 Oct 2026 15:40:59 +0200 Subject: [PATCH 02/13] Check that every column is held by exactly one MPI rank On several ranks, a column held by more than one rank would be computed and written twice. Before the time loop, rank 0 gathers the column ids of all ranks, and every rank stops with an error if an id occurs twice. The check sends the ids of all columns to rank 0, so --no-check-grid-distribution skips it; rank 0 says so when it passes. Assisted-by: Claude Code:claude-opus-5-5 --- CHANGELOG.md | 1 + src/ebfm/core/cli.py | 8 ++++ src/ebfm/core/grid.py | 39 +++++++++++++++++++ src/ebfm/main.py | 7 ++++ tests/core/comm/validate_grid_distribution.py | 35 +++++++++++++++++ tests/core/test_cli.py | 8 ++++ tests/core/test_parallel_columns.py | 22 ++++++++++- 7 files changed, 118 insertions(+), 2 deletions(-) create mode 100644 tests/core/comm/validate_grid_distribution.py diff --git a/CHANGELOG.md b/CHANGELOG.md index df0a7ac4..f2d07f64 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,7 @@ 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/XXX # v0.8.0 diff --git a/src/ebfm/core/cli.py b/src/ebfm/core/cli.py index 284671f7..417d4a47 100644 --- a/src/ebfm/core/cli.py +++ b/src/ebfm/core/cli.py @@ -349,6 +349,14 @@ def parse_time_step(value: str) -> float | str: "If not provided, the MPI rank + 1 will be used as partition ID. Only allowed in runs on a single MPI rank.", ) + parallel_group.add_argument( + "--check-grid-distribution", + default=True, + action=argparse.BooleanOptionalAction, + help="Check before the time loop that every column is held by exactly one MPI rank. The check " + "communicates the ids of all columns; skip it once a setup is known to distribute its grid correctly.", + ) + logger_group = parser.add_argument_group("logging configuration") logger_group.add_argument( diff --git a/src/ebfm/core/grid.py b/src/ebfm/core/grid.py index 639fb5ba..d6623c18 100644 --- a/src/ebfm/core/grid.py +++ b/src/ebfm/core/grid.py @@ -21,10 +21,13 @@ """ from enum import Enum +import logging from typing import Any import numpy as np +logger = logging.getLogger(__name__) + class GridInputType(Enum): # .mat grid file with elevation @@ -137,6 +140,42 @@ def domain_mean(values, comm) -> float: return total / count +def validate_grid_distribution(grid: GridDict, comm) -> None: + """ + Check that every column is held by exactly one rank of comm. + + Collective over comm. + + Parameters: + grid (GridDict): this rank's part of the grid + comm: EBFM's MPI communicator + + Raises: + ValueError: on every rank, if a column is held by more than one rank + """ + if comm.size == 1: + return + + column_ids_per_rank = comm.gather(grid["column_id"], root=0) + n_duplicates = 0 + if comm.rank == 0: + column_ids = np.concatenate(column_ids_per_rank) + n_duplicates = len(column_ids) - len(np.unique(column_ids)) + n_duplicates = comm.bcast(n_duplicates, root=0) + + if n_duplicates > 0: + raise ValueError( + f"Some columns are held by more than one of the {comm.size} MPI ranks ({n_duplicates} duplicates). " + "Each column must be held by exactly one rank, i.e. the grid must be distributed over the ranks." + ) + + if comm.rank == 0: + logger.warning( + f"Checked that each column is held by exactly one of the {comm.size} MPI ranks. The check " + "communicates the ids of all columns; pass --no-check-grid-distribution to skip it." + ) + + class ShadingMethod(Enum): """Available shading algorithms for topographic shading. diff --git a/src/ebfm/main.py b/src/ebfm/main.py index 999c64b0..d30aa216 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -18,6 +18,7 @@ ) from ebfm.core import LOOP_write_to_file, FINAL_create_restart_file from ebfm.core.comm.gather import ColumnGather +from ebfm.core.grid import validate_grid_distribution from ebfm.core.config import CouplingConfig, GridConfig, TimeConfig, ForcingConfig from ebfm.core.logger import Logger, setup_logging, log_levels_map, getLogger from ebfm.core.cli import ( @@ -157,6 +158,8 @@ def _main_impl(): else: ebfm_comm = ebfm.core.comm.defaultComm + is_parallel_run = ebfm_comm.size > 1 + coupler_cls = ebfm.coupling.select_coupler_class(coupling_config) coupler = coupler_cls(coupling_config=coupling_config) @@ -233,6 +236,10 @@ def _main_impl(): C = INIT.init_constants() grid = INIT.init_grid(grid, grid_config, ebfm_comm) + # Each rank must compute its own columns, otherwise the ranks would repeat each other's work + if args.check_grid_distribution and is_parallel_run: + validate_grid_distribution(grid, ebfm_comm) + # Rank 0 collects the columns of all ranks, e.g. to write them to the output file or for --dump-reference column_gather = ColumnGather(ebfm_comm, grid["column_id"]) if ebfm_comm.size > 1 and not grid["is_partitioned"]: diff --git a/tests/core/comm/validate_grid_distribution.py b/tests/core/comm/validate_grid_distribution.py new file mode 100644 index 00000000..219de61d --- /dev/null +++ b/tests/core/comm/validate_grid_distribution.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Check validate_grid_distribution on 3 MPI ranks that hold the 10 columns of a domain unevenly. + +Usage: mpirun -np 3 python validate_grid_distribution.py {distinct|duplicate} + +distinct: every column is held by one rank, and rank 2 holds none. Checks that no rank raises. +duplicate: rank 2 holds a column that rank 0 holds as well. Checks that every rank raises a ValueError. +""" + +import sys + +import numpy as np + +from ebfm.core.comm import defaultComm as comm +from ebfm.core.grid import validate_grid_distribution + +assert comm.size == 3, f"Run on 3 ranks, not {comm.size}." +mode = sys.argv[1] + +column_ids = 100 + 7 * np.arange(10) +columns = {0: [5, 1, 7, 3], 1: [0, 9, 2, 8, 4, 6], 2: []}[comm.rank] +if mode == "duplicate" and comm.rank == 2: + columns = [1] +grid = {"column_id": column_ids[columns]} + +try: + validate_grid_distribution(grid, comm) +except ValueError as e: + print(f"rank {comm.rank}: {e}") + sys.exit(0) + +print(f"rank {comm.rank}: ok") diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py index bb3e41f3..979b9086 100644 --- a/tests/core/test_cli.py +++ b/tests/core/test_cli.py @@ -248,6 +248,14 @@ def test_several_processes_with_mpi4py_ok(self): self._parse({"OMPI_COMM_WORLD_SIZE": "4"}, mpi_available=True) # must not raise +class TestCheckGridDistribution(unittest.TestCase): + def test_enabled_by_default(self): + self.assertTrue(_parse().check_grid_distribution) + + def test_disabled_when_negated(self): + self.assertFalse(_parse("--no-check-grid-distribution").check_grid_distribution) + + class TestLocalGroupLabel(unittest.TestCase): """--local-group-label overrides the default derived from --component-name.""" diff --git a/tests/core/test_parallel_columns.py b/tests/core/test_parallel_columns.py index b635e0ab..07c49a3b 100644 --- a/tests/core/test_parallel_columns.py +++ b/tests/core/test_parallel_columns.py @@ -2,7 +2,8 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Columns distributed over several MPI ranks are gathered on one rank in column id order and reduced over all ranks.""" +"""Columns distributed over several MPI ranks are held by exactly one rank each, gathered on one rank in column id +order and reduced over all ranks.""" import importlib.util import os @@ -16,7 +17,7 @@ from ebfm.core.comm import defaultComm from ebfm.core.comm.gather import ColumnGather -from ebfm.core.grid import domain_mean +from ebfm.core.grid import domain_mean, validate_grid_distribution COMM_TESTS = Path(__file__).resolve().parent / "comm" TIMEOUT = 120 @@ -68,3 +69,20 @@ def test_domain_mean_on_three_ranks(): assert result.returncode == 0, result.stdout + result.stderr for rank in range(3): assert f"rank {rank}: ok" in result.stdout, result.stdout + result.stderr + + +def test_single_rank_holds_every_column(): + validate_grid_distribution({"column_id": np.array([4, 2, 9])}, defaultComm) # must not raise + + +@requires_mpirun +@pytest.mark.parametrize( + "mode, expected", + [("distinct", "ok"), ("duplicate", "Some columns are held by more than one of the 3 MPI ranks")], +) +def test_grid_distribution_on_three_ranks(mode, expected): + result = _mpirun(3, "validate_grid_distribution.py", mode) + + assert result.returncode == 0, result.stdout + result.stderr + for rank in range(3): + assert f"rank {rank}: {expected}" in result.stdout, result.stdout + result.stderr From 9a0b2e8d39378cda8a54abdf49fc4f941f875722 Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Tue, 20 Oct 2026 15:46:19 +0200 Subject: [PATCH 03/13] Write the output of partitioned and unstructured grids on rank 0 A partitioned mesh skipped writing output, and the output of an unstructured grid could not be gathered over the ranks. Rank 0 now also gathers the cells of an unstructured grid, ordered by Elmer element id, together with the meshes of all ranks; vertices that several ranks share are written once. Assisted-by: Claude Code:claude-opus-5-5 --- CHANGELOG.md | 1 + README.md | 2 + src/ebfm/core/LOOP_write_to_file.py | 42 ++++++++++++-- src/ebfm/core/comm/gather.py | 1 + src/ebfm/main.py | 10 +--- tests/core/comm/write_output_on_ranks.py | 30 ++++++++++ tests/core/test_loop_write_to_file.py | 72 +++++++++++++++++++++++- tests/core/test_parallel_columns.py | 16 ++++++ 8 files changed, 160 insertions(+), 14 deletions(-) create mode 100644 tests/core/comm/write_output_on_ranks.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f2d07f64..84503658 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,7 @@ SPDX-License-Identifier: CC-BY-4.0 * 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/XXX +* 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/XXX # v0.8.0 diff --git a/README.md b/README.md index 7f5f80fd..495aaad0 100644 --- a/README.md +++ b/README.md @@ -225,6 +225,8 @@ 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 ``` +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, 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. diff --git a/src/ebfm/core/LOOP_write_to_file.py b/src/ebfm/core/LOOP_write_to_file.py index 7ea8c3a9..1a0fb419 100644 --- a/src/ebfm/core/LOOP_write_to_file.py +++ b/src/ebfm/core/LOOP_write_to_file.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: BSD-3-Clause import os +from types import SimpleNamespace import numpy as np from netCDF4 import Dataset, date2num @@ -13,6 +14,32 @@ _FILL_VALUE = -9999.0 +def _gather_mesh(gather: ColumnGather, mesh) -> SimpleNamespace | None: + """ + Combine the meshes of all ranks into the mesh of all their cells, on the root rank. + + Vertices shared by the meshes of several ranks appear once, in the order of their global id. + + @returns an object with the vertex coordinates and the cell-to-vertex mapping that _write_unstructured_grid + reads on the root rank, None on all other ranks + """ + vertex_ids = gather.gather_unordered(mesh.vertex_ids) + vertex_coordinates = { + name: gather.gather_unordered(getattr(mesh, name)) + for name in ("x_vertices", "y_vertices", "z_vertices", "lon_vertices", "lat_vertices") + } + cell_vertex_ids = gather.gather(mesh.vertex_ids[mesh.cell_to_vertex]) + + if not gather.is_root: + return None + + unique_vertex_ids, first_occurrence = np.unique(vertex_ids, return_index=True) + return SimpleNamespace( + **{name: values[first_occurrence] for name, values in vertex_coordinates.items()}, + cell_to_vertex=np.searchsorted(unique_vertex_ids, cell_vertex_ids), + ) + + 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. @@ -22,11 +49,16 @@ def _gather_grid(gather: ColumnGather, grid): 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 grid["is_unstructured"]: + output_grid = {"is_unstructured": True} + for key in ("x", "y", "z", "lon", "lat"): + output_grid[key] = gather.gather(grid[key]) + if grid.get("mesh") is not None: + output_grid["mesh"] = _gather_mesh(gather, grid["mesh"]) + else: + # 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 diff --git a/src/ebfm/core/comm/gather.py b/src/ebfm/core/comm/gather.py index be8fd336..f0cdbf6b 100644 --- a/src/ebfm/core/comm/gather.py +++ b/src/ebfm/core/comm/gather.py @@ -25,6 +25,7 @@ def __init__(self, comm: Comm, column_ids: NDArray[np.integer], root: int = 0): """ @param[in] comm communicator whose ranks hold the columns @param[in] column_ids ids of the columns this rank holds (grid["column_id"]); no other rank may hold them + (see ebfm.core.grid.validate_grid_distribution) @param[in] root rank that receives the gathered arrays """ self.comm = comm diff --git a/src/ebfm/main.py b/src/ebfm/main.py index d30aa216..c6765f52 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -242,7 +242,7 @@ def _main_impl(): # Rank 0 collects the columns of all ranks, e.g. to write them to the output file or for --dump-reference column_gather = ColumnGather(ebfm_comm, grid["column_id"]) - if ebfm_comm.size > 1 and not grid["is_partitioned"]: + if is_parallel_run: logger.info(f"Rank 0 writes the output of all {ebfm_comm.size} ranks.") OUT, IN, OUTFILE = INIT.init_initial_conditions(C, grid, io, time, column, init_with_restart_file=args.restart_init) @@ -396,12 +396,8 @@ def _main_impl(): OUT["y"] = grid["y"] OUT["surface_elevation"] = grid["z"] - # Write output to files (for unpartitioned grids only; see the warning emitted before the time loop) - # TODO: partitioned output should be supported in a dedicated parallel-safe writer. - if not grid["is_partitioned"]: - io, OUTFILE = LOOP_write_to_file.main(OUTFILE, io, OUT, grid, t, time, column, column_gather) - else: - logger.warning("Skipping writing output to file for partitioned runs.") + # Write output to files (see the warning emitted before the time loop) + io, OUTFILE = LOOP_write_to_file.main(OUTFILE, io, OUT, grid, t, time, column, column_gather) # Write restart file # TODO: should be supported for all cases to avoid case distinction here diff --git a/tests/core/comm/write_output_on_ranks.py b/tests/core/comm/write_output_on_ranks.py new file mode 100644 index 00000000..371389da --- /dev/null +++ b/tests/core/comm/write_output_on_ranks.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Write the output of a test grid of test_loop_write_to_file with its columns distributed over the MPI ranks. + +Usage: mpirun -np N python write_output_on_ranks.py {structured|unstructured} OUTDIR + +Rank r holds the columns r, r + N, r + 2N, ... in reverse order, and only the mesh vertices of its own cells, so +that neither the ranks nor their columns and vertices are in id order. +""" + +import sys +from pathlib import Path + +import numpy as np + +from ebfm.core.comm import defaultComm + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +import test_loop_write_to_file as fixtures # noqa: E402 + +GRIDS = { + "structured": fixtures._structured_grid, + "unstructured": fixtures._unstructured_grid_with_mesh, +} + +grid_kind, outdir = sys.argv[1], sys.argv[2] +columns = np.arange(fixtures.GPSUM)[defaultComm.rank :: defaultComm.size][::-1] +fixtures._write_single_step(outdir, GRIDS[grid_kind](), columns=columns) diff --git a/tests/core/test_loop_write_to_file.py b/tests/core/test_loop_write_to_file.py index 18e21894..82c52524 100644 --- a/tests/core/test_loop_write_to_file.py +++ b/tests/core/test_loop_write_to_file.py @@ -8,6 +8,7 @@ import unittest from datetime import datetime from pathlib import Path +from types import SimpleNamespace import numpy as np from netCDF4 import Dataset @@ -24,8 +25,11 @@ GLACIER_INDICES = (0, 3, 5, 7, 11) GPSUM = len(GLACIER_INDICES) -# Element ids of the unstructured grid. Like the ids of an Elmer mesh, they neither start at 0 nor are contiguous. +# Element and node ids of the unstructured mesh: a strip of GPSUM triangles, cell i spanning the vertices i, i+1 +# and i+2. Like the ids of an Elmer mesh, they are global and in file order, but neither start at 0 nor are +# contiguous. ELEMENT_IDS = (3, 5, 12, 27, 40) +NODE_IDS = tuple(2 * (i + 1) for i in range(GPSUM + 2)) _NETCDF_OUTPUT = 2 @@ -116,7 +120,56 @@ def _unstructured_grid(): } -def _write_single_step(outdir, grid, extra_out=None): +def _mesh(): + """The mesh of _unstructured_grid(), with the attributes of ebfm.elmer.mesh.Mesh that the writer reads.""" + n_vertices = len(NODE_IDS) + return SimpleNamespace( + vertex_ids=np.array(NODE_IDS), + x_vertices=10.0 * np.arange(n_vertices), + y_vertices=10.0 * (np.arange(n_vertices) % 2), + z_vertices=np.full(n_vertices, 50.0), + lon_vertices=np.radians(np.arange(n_vertices, dtype=np.float64)), + lat_vertices=np.radians(70.0 + np.arange(n_vertices)), + cell_to_vertex=np.array([[i, i + 1, i + 2] for i in range(GPSUM)]), + ) + + +def _unstructured_grid_with_mesh(): + return {**_unstructured_grid(), "mesh": _mesh()} + + +def _local_mesh(mesh, columns): + """The mesh of the given cells only, with their vertices in reverse id order.""" + vertices = np.unique(mesh.cell_to_vertex[columns])[::-1] + local_index = {vertex: i for i, vertex in enumerate(vertices)} + return SimpleNamespace( + vertex_ids=mesh.vertex_ids[vertices], + x_vertices=mesh.x_vertices[vertices], + y_vertices=mesh.y_vertices[vertices], + z_vertices=mesh.z_vertices[vertices], + lon_vertices=mesh.lon_vertices[vertices], + lat_vertices=mesh.lat_vertices[vertices], + cell_to_vertex=np.array([[local_index[v] for v in cell] for cell in mesh.cell_to_vertex[columns]]), + ) + + +def _local_share(grid, out, columns): + """ + The grid and OUT that an MPI rank holding only the given columns sees. + + @param[in] columns positions of the rank's columns in the grid, in the order the rank holds them + """ + local_grid = dict(grid) + for key in ("column_id", "x", "y", "z", "lon", "lat", "ind"): + if key in grid: + local_grid[key] = grid[key][columns] + if "mesh" in grid: + local_grid["mesh"] = _local_mesh(grid["mesh"], columns) + local_out = {key: values[columns] for key, values in out.items()} + return local_grid, local_out + + +def _write_single_step(outdir, grid, extra_out=None, columns=None): """ Run the writer for a one-step simulation and return the closed output file. @@ -127,6 +180,8 @@ def _write_single_step(outdir, grid, extra_out=None): @param[in] grid structured or unstructured grid dict (_structured_grid()/_unstructured_grid()) @param[in] extra_out optional dict merged into the fabricated OUT before writing, for keys _fake_out() doesn't declare (e.g. the conditional "ebm_diagnostics") + @param[in] columns positions of the columns this MPI rank holds (see _local_share); None for all columns. + The writer is called on every rank of ebfm.core.comm.defaultComm. @returns (path to the written NetCDF file, the writer's outfile dict, the OUT dict actually written) @@ -135,6 +190,8 @@ def _write_single_step(outdir, grid, extra_out=None): out = _fake_out(GPSUM, NL) if extra_out: out.update(extra_out) + if columns is not None: + grid, out = _local_share(grid, out, columns) io = {"outdir": str(outdir), "freqout": 1, "output_type": _NETCDF_OUTPUT} time = {"TCUR": datetime(2020, 1, 1), "tn": 1} @@ -192,6 +249,17 @@ def test_unstructured_grid_writes_all_declared_variables(self): self.assertEqual(nc[varname].dimensions, ("time", "cell")) self.assertEqual(nc[varname].shape, (1, GPSUM)) + def test_unstructured_grid_writes_mesh_topology(self): + with tempfile.TemporaryDirectory() as tmp: + path, _, _ = _write_single_step(tmp, _unstructured_grid_with_mesh()) + + mesh = _mesh() + with Dataset(path) as nc: + self.assertEqual(len(nc.dimensions["vertex"]), len(NODE_IDS)) + np.testing.assert_array_equal(nc["x_vertex"][:], mesh.x_vertices) + np.testing.assert_array_equal(nc["lat_vertex"][:], np.degrees(mesh.lat_vertices)) + np.testing.assert_array_equal(nc["cell_vertices"][:], mesh.cell_to_vertex) + def test_structured_grid_scatters_values_onto_glacier_points(self): with tempfile.TemporaryDirectory() as tmp: path, _, out = _write_single_step(tmp, _structured_grid()) diff --git a/tests/core/test_parallel_columns.py b/tests/core/test_parallel_columns.py index 07c49a3b..653c3918 100644 --- a/tests/core/test_parallel_columns.py +++ b/tests/core/test_parallel_columns.py @@ -5,6 +5,7 @@ """Columns distributed over several MPI ranks are held by exactly one rank each, gathered on one rank in column id order and reduced over all ranks.""" +import filecmp import importlib.util import os import shutil @@ -18,6 +19,7 @@ from ebfm.core.comm import defaultComm from ebfm.core.comm.gather import ColumnGather from ebfm.core.grid import domain_mean, validate_grid_distribution +from test_loop_write_to_file import _structured_grid, _unstructured_grid_with_mesh, _write_single_step COMM_TESTS = Path(__file__).resolve().parent / "comm" TIMEOUT = 120 @@ -86,3 +88,17 @@ def test_grid_distribution_on_three_ranks(mode, expected): assert result.returncode == 0, result.stdout + result.stderr for rank in range(3): assert f"rank {rank}: {expected}" in result.stdout, result.stdout + result.stderr + + +@requires_mpirun +@pytest.mark.parametrize("n_ranks", [2, 3]) +@pytest.mark.parametrize( + "grid_kind, make_grid", + [("structured", _structured_grid), ("unstructured", _unstructured_grid_with_mesh)], +) +def test_parallel_output_file_equals_serial_one(tmp_path, n_ranks, grid_kind, make_grid): + serial_file, _, _ = _write_single_step(tmp_path / "serial", make_grid()) + result = _mpirun(n_ranks, "write_output_on_ranks.py", grid_kind, str(tmp_path / "parallel")) + + assert result.returncode == 0, result.stdout + result.stderr + assert filecmp.cmp(serial_file, tmp_path / "parallel" / "model_output.nc", shallow=False) From 9618d5244cde66d24d67215e4c47b9ad3e0ecf83 Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Tue, 20 Oct 2026 16:31:21 +0200 Subject: [PATCH 04/13] Split a non-partitioned Elmer mesh over the MPI ranks On several ranks, each rank now keeps a contiguous block of the cells of the mesh, the first block on rank 0, and only the vertices those cells use. The projection of the vertices and the cell centres moves from the Mesh constructor to Mesh.compute_geometry(), which INIT calls after the split, so each rank transforms only its own cells. The test elevations of the ELMER grid type are drawn for the whole mesh before the split, so every column keeps its serial value. The startup check accepts ELMER and CUSTOM grids on a non-partitioned mesh now. Rank 0 warns that fitting the lapse rate over all columns synchronizes the ranks every time step. Assisted-by: Claude Code:claude-opus-5-5 --- CHANGELOG.md | 1 + README.md | 8 ++-- src/ebfm/core/INIT.py | 17 ++++++++- src/ebfm/core/cli.py | 8 ++-- src/ebfm/core/config/grid.py | 8 ++-- src/ebfm/elmer/mesh.py | 52 +++++++++++++++++++++++--- src/ebfm/main.py | 8 +++- tests/core/test_cli.py | 13 ++++--- tests/core/test_elmer_mesh.py | 70 +++++++++++++++++++++++++++++++++++ 9 files changed, 161 insertions(+), 24 deletions(-) create mode 100644 tests/core/test_elmer_mesh.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 84503658..bda52a3d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,7 @@ SPDX-License-Identifier: CC-BY-4.0 * 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/XXX * 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/XXX +* 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. https://github.com/EBFMorg/EBFM/pull/XXX # v0.8.0 diff --git a/README.md b/README.md index 495aaad0..35b4c50a 100644 --- a/README.md +++ b/README.md @@ -227,7 +227,7 @@ mpirun -n 4 ebfm --elmer-mesh examples/greenland_mesh_v0/MESH/partitioning.128/ 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, only a MATLAB grid runs on several MPI ranks (see [Running on several MPI ranks](#running-on-several-mpi-ranks)). +Besides a partitioned mesh, a MATLAB grid and a non-partitioned Elmer mesh run 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. @@ -290,7 +290,7 @@ Note that it is necessary to provide a `--start-time` consistent with the given ### Running on several MPI ranks -EBFM splits a MATLAB grid over the MPI ranks of a run (requires `pip install 'ebfm[mpi]'`): +EBFM splits a MATLAB grid or a non-partitioned Elmer mesh 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 @@ -300,8 +300,8 @@ mpirun -np 4 ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading * 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. + `--forcing-dir`, `--restart-dir`/`--restart-init`, `--netcdf-mesh`, and + `--netcdf-dem-mesh-unstructured` on a non-partitioned Elmer mesh. ### Performance and Profiling Runs diff --git a/src/ebfm/core/INIT.py b/src/ebfm/core/INIT.py index 0d4834e1..ebb1e203 100644 --- a/src/ebfm/core/INIT.py +++ b/src/ebfm/core/INIT.py @@ -221,6 +221,11 @@ def init_grid(grid: GridDict, config: GridConfig, comm): source_crs_epsg=config.elmer_mesh_crs_epsg, ) + if config.num_partitions > 1: + mesh = mesh.get_partition(config.partition_id, config.num_partitions) + + mesh.compute_geometry() + grid["column_id"] = mesh.cell_ids grid["x"], grid["y"] = mesh.x_cells, mesh.y_cells @@ -251,15 +256,23 @@ def init_grid(grid: GridDict, config: GridConfig, comm): # dzdx, dzdy = mesh.dzdy, mesh.dzdy elif config.grid_type is GridInputType.ELMER: # Read grid and elevations from Elmer mesh: Mesh = read_elmer_mesh( - config.mesh_file, + mesh_root=config.mesh_file, source_crs_epsg=config.elmer_mesh_crs_epsg, ) + # Test values, drawn for the whole mesh so that a column gets the same value however the mesh is split + mesh.z_cells = np.random.uniform(0, 100, size=len(mesh.cell_ids)) + + if config.num_partitions > 1: + mesh = mesh.get_partition(config.partition_id, config.num_partitions) + + mesh.compute_geometry() + # assuming mesh/MESH/mesh.nodes contains DEM data in the z component # see mesh/README.md for the required preprocessing steps. grid["column_id"] = mesh.cell_ids grid["x"], grid["y"] = mesh.x_cells, mesh.y_cells # use cell centers as location for DoFs - grid["z"] = np.random.uniform(0, 100, size=len(grid["x"])) # test values! + grid["z"] = mesh.z_cells # test values! grid["slope_x"] = np.zeros_like(grid["x"]) # test values! grid["slope_y"] = np.zeros_like(grid["x"]) # test values! grid["lat"] = np.zeros_like(grid["x"]) + 75 # test values! diff --git a/src/ebfm/core/cli.py b/src/ebfm/core/cli.py index 417d4a47..02b6ccaf 100644 --- a/src/ebfm/core/cli.py +++ b/src/ebfm/core/cli.py @@ -562,7 +562,8 @@ def validate_shading_parallel_compat(grid_config, comm) -> None: def validate_grid_parallel_compat(grid_config, comm) -> None: """Check that a run on several MPI ranks uses a grid that runs there. - EBFM splits a MATLAB grid over the ranks; of a partitioned Elmer mesh, each rank computes one partition. + EBFM splits a MATLAB grid and a non-partitioned Elmer mesh over the ranks; of a partitioned Elmer mesh, each rank + computes one partition. Must be called on every rank of comm. @@ -579,14 +580,13 @@ def validate_grid_parallel_compat(grid_config, comm) -> None: if grid_config.grid_type is GridInputType.NETCDF: rejected_grid = "--netcdf-mesh" - elif not grid_config.is_partitioned: - rejected_grid = "--elmer-mesh without --is-partitioned-elmer-mesh" + elif grid_config.grid_type is GridInputType.ELMERXIOS and not grid_config.is_partitioned: + rejected_grid = "--netcdf-dem-mesh-unstructured without --is-partitioned-elmer-mesh" else: return _cli_error( f"{rejected_grid} is not supported for runs on several MPI ranks ({comm.size} ranks). " - "Only --matlab-mesh and a partitioned Elmer mesh (--is-partitioned-elmer-mesh) run on several ranks. " "Please run on a single rank." ) diff --git a/src/ebfm/core/config/grid.py b/src/ebfm/core/config/grid.py index f12f2ac7..01e157c9 100644 --- a/src/ebfm/core/config/grid.py +++ b/src/ebfm/core/config/grid.py @@ -27,7 +27,8 @@ class GridConfig: is_partitioned: bool # Whether the grid is partitioned is_unstructured: bool # Whether the grid is unstructured partition_id: int | None # Partition ID (only relevant if is_partitioned is True or num_partitions > 1) - # Number of blocks EBFM splits the columns of a MATLAB grid into, one per MPI rank; 1 for a grid that is not split + # Number of blocks EBFM splits the columns of a MATLAB grid or a non-partitioned Elmer mesh into, one per MPI rank; + # 1 for a grid that is not split num_partitions: int elmer_mesh_crs_epsg: int # EPSG code of Elmer mesh coordinates use_shading: bool # Whether to use shading for the grid @@ -53,7 +54,7 @@ def __init__(self, args: Namespace, num_ranks: int = 1, rank: int = 0): Initialize grid configuration from command line arguments. @param[in] args command line arguments - @param[in] num_ranks number of MPI ranks, over which a MATLAB grid is split + @param[in] num_ranks number of MPI ranks, over which a MATLAB grid or a non-partitioned Elmer mesh is split @param[in] rank MPI rank of this process, which computes the block rank + 1 of a split grid """ selected_primary_grids = [ @@ -118,7 +119,8 @@ def __init__(self, args: Namespace, num_ranks: int = 1, rank: int = 0): ) raise Exception("Invalid grid configuration.") - if self.grid_type is GridInputType.MATLAB and num_ranks > 1: + is_split = self.grid_type is GridInputType.MATLAB or (self.is_unstructured and not self.is_partitioned) + if is_split and num_ranks > 1: logger.info(f"Splitting the grid over {num_ranks} MPI ranks...") self.partition_id = rank + 1 self.num_partitions = num_ranks diff --git a/src/ebfm/elmer/mesh.py b/src/ebfm/elmer/mesh.py index 9de2deb5..06e49b02 100644 --- a/src/ebfm/elmer/mesh.py +++ b/src/ebfm/elmer/mesh.py @@ -58,6 +58,9 @@ class Mesh: y_vertices: NDArray[np.float64] # optional height values of vertices stored in z coordinate z_vertices: NDArray[np.float64] + # EPSG code of the projection of x/y + source_crs_epsg: int + # The following coordinates are derived from x/y_vertices by compute_geometry() # longitude/latitude coordinates of vertices in radians, ordering follows local ids [0,1,...,n_vertices-1] lon_vertices: NDArray[np.float64] lat_vertices: NDArray[np.float64] @@ -91,7 +94,19 @@ def __init__( self.vertex_ids = vertex_ids self.cell_ids = cell_ids self.cell_to_vertex = cell_to_vertex - transformer = pyproj.Transformer.from_crs(source_crs_epsg, 4326, always_xy=True) + self.source_crs_epsg = source_crs_epsg + # Initialize z_cells to zeros (to be filled/overwritten by DEM) + n_cells = self.cell_to_vertex.shape[0] + self.z_cells = np.zeros(n_cells) + + def compute_geometry(self) -> None: + """ + Derive the lon/lat coordinates of the vertices and the coordinates of the cell centers. + + Transforming the vertices is the expensive part of setting up a mesh. Call this once the mesh holds the cells + it will be used with, e.g. after get_partition, so that only those are transformed. + """ + transformer = pyproj.Transformer.from_crs(self.source_crs_epsg, 4326, always_xy=True) self.lon_vertices, self.lat_vertices = transformer.transform(self.x_vertices, self.y_vertices, radians=True) # Compute cell centers (lon/lat) by averaging vertices on the unit sphere @@ -99,11 +114,38 @@ def __init__( self.lon_vertices, self.lat_vertices, self.cell_to_vertex ) - inverse_transformer = pyproj.Transformer.from_crs(4326, source_crs_epsg, always_xy=True) + inverse_transformer = pyproj.Transformer.from_crs(4326, self.source_crs_epsg, always_xy=True) self.x_cells, self.y_cells = inverse_transformer.transform(self.lon_cells, self.lat_cells, radians=True) - # Initialize z_cells to zeros (to be filled/overwritten by DEM) - n_cells = self.cell_to_vertex.shape[0] - self.z_cells = np.zeros(n_cells) + + def get_partition(self, partition_id: int, n_partitions: int) -> "Mesh": + """ + The part of this mesh that one of `n_partitions` partitions holds. + + The cells are split into `n_partitions` contiguous blocks in their current order, and partition 1 holds the + first block. A partition keeps the vertices its cells use, so a vertex on the border between two partitions + belongs to both. + + @param[in] partition_id id of the partition, from 1 to n_partitions + @param[in] n_partitions number of partitions + + @returns a mesh of the same type with the cells of the partition, including their z_cells + """ + assert 1 <= partition_id <= n_partitions, f"Partition id {partition_id} is not in 1..{n_partitions}." + + cells = np.array_split(np.arange(len(self.cell_ids)), n_partitions)[partition_id - 1] + vertices, cell_to_vertex = np.unique(self.cell_to_vertex[cells], return_inverse=True) + + partition = type(self)( + x_vertices=self.x_vertices[vertices], + y_vertices=self.y_vertices[vertices], + z_vertices=self.z_vertices[vertices], + cell_to_vertex=cell_to_vertex.reshape(len(cells), -1), + vertex_ids=self.vertex_ids[vertices], + cell_ids=self.cell_ids[cells], + source_crs_epsg=self.source_crs_epsg, + ) + partition.z_cells = self.z_cells[cells] + return partition class TriangleMesh(Mesh): diff --git a/src/ebfm/main.py b/src/ebfm/main.py index c6765f52..29ad225e 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -213,7 +213,7 @@ def _main_impl(): # See https://github.com/EBFMorg/EBFM/issues/11 for details. validate_shading_coupling_compat(grid_config, coupling_config) - # Ensure a run on several ranks uses a partitioned mesh. + # Ensure a run on several ranks uses a grid that runs there. validate_grid_parallel_compat(grid_config, ebfm_comm) # Ensure forcing files are only read in runs on a single rank. @@ -245,6 +245,12 @@ def _main_impl(): if is_parallel_run: logger.info(f"Rank 0 writes the output of all {ebfm_comm.size} ranks.") + if is_parallel_run and ebfm_comm.rank == 0: + logger.warning( + "The lapse rate of potential temperature is fitted over the columns of all ranks, which synchronizes all " + f"{ebfm_comm.size} ranks every time step." + ) + OUT, IN, OUTFILE = INIT.init_initial_conditions(C, grid, io, time, column, init_with_restart_file=args.restart_init) coupler.setup(grid, time_config) diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py index 979b9086..5fd00e7d 100644 --- a/tests/core/test_cli.py +++ b/tests/core/test_cli.py @@ -414,15 +414,19 @@ def test_shading_on_several_ranks_rejected(self): class TestValidateGridParallelCompat(unittest.TestCase): - """validate_grid_parallel_compat rejects every grid but MATLAB and a partitioned Elmer mesh on several ranks.""" + """validate_grid_parallel_compat rejects a NetCDF grid and a non-partitioned ELMERXIOS grid on several ranks.""" # name: (grid_type, is_partitioned, grid option named by the error on several ranks, or None if accepted) _CASES = { "--matlab-mesh": (GridInputType.MATLAB, False, None), "--netcdf-mesh": (GridInputType.NETCDF, False, "--netcdf-mesh"), - "--elmer-mesh": (GridInputType.ELMER, False, "--elmer-mesh"), - "--elmer-mesh --netcdf-dem-mesh": (GridInputType.CUSTOM, False, "--elmer-mesh"), - "--elmer-mesh --netcdf-dem-mesh-unstructured": (GridInputType.ELMERXIOS, False, "--elmer-mesh"), + "--elmer-mesh": (GridInputType.ELMER, False, None), + "--elmer-mesh --netcdf-dem-mesh": (GridInputType.CUSTOM, False, None), + "--elmer-mesh --netcdf-dem-mesh-unstructured": ( + GridInputType.ELMERXIOS, + False, + "--netcdf-dem-mesh-unstructured", + ), "--elmer-mesh --netcdf-dem-mesh --is-partitioned-elmer-mesh": (GridInputType.CUSTOM, True, None), } @@ -451,7 +455,6 @@ def test_several_ranks(self): self.assertEqual(ctx.exception.code, 2) error = stderr.getvalue().splitlines()[-1] self.assertIn(f"error: {rejected_grid} ", error) - self.assertIn("Only --matlab-mesh and a partitioned Elmer mesh", error) @unittest.skipUnless(mpi_available, "requires MPI (pip install 'ebfm[mpi]')") diff --git a/tests/core/test_elmer_mesh.py b/tests/core/test_elmer_mesh.py new file mode 100644 index 00000000..357540f0 --- /dev/null +++ b/tests/core/test_elmer_mesh.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# SPDX-License-Identifier: BSD-3-Clause + +import unittest + +import numpy as np + +from ebfm.elmer.mesh import TriangleMesh + +N_CELLS = 10 +N_PARTITIONS = 3 + + +def _strip_mesh() -> TriangleMesh: + """A strip of N_CELLS triangles in EPSG:3413, cell i spanning the vertices i, i+1 and i+2. + + Like in an Elmer mesh, vertex and cell ids are global and neither start at 0 nor are contiguous. + """ + n_vertices = N_CELLS + 2 + mesh = TriangleMesh( + x_vertices=1000.0 * np.arange(n_vertices), + y_vertices=-2.0e6 + 1000.0 * (np.arange(n_vertices) % 2), + z_vertices=np.zeros(n_vertices), + cell_to_vertex=np.array([[i, i + 1, i + 2] for i in range(N_CELLS)]), + vertex_ids=10 + 3 * np.arange(n_vertices), + cell_ids=7 + 5 * np.arange(N_CELLS), + ) + mesh.z_cells = np.arange(N_CELLS, dtype=np.float64) + return mesh + + +class TestGetPartition(unittest.TestCase): + """Splitting a mesh into partitions keeps every cell, in order, with its geometry and its vertices.""" + + def setUp(self): + self.mesh = _strip_mesh() + self.partitions = [self.mesh.get_partition(p, N_PARTITIONS) for p in range(1, N_PARTITIONS + 1)] + for mesh in (self.mesh, *self.partitions): + mesh.compute_geometry() + + def test_partitions_hold_contiguous_blocks_with_the_first_block_in_partition_1(self): + self.assertEqual([len(p.cell_ids) for p in self.partitions], [4, 3, 3]) + np.testing.assert_array_equal(np.concatenate([p.cell_ids for p in self.partitions]), self.mesh.cell_ids) + + def test_cell_centres_and_elevations_equal_those_of_the_whole_mesh(self): + for name in ("x_cells", "y_cells", "lon_cells", "lat_cells", "z_cells"): + with self.subTest(name=name): + combined = np.concatenate([getattr(p, name) for p in self.partitions]) + np.testing.assert_array_equal(combined, getattr(self.mesh, name)) + + def test_cells_keep_their_vertices(self): + whole = self.mesh.vertex_ids[self.mesh.cell_to_vertex] + combined = np.concatenate([p.vertex_ids[p.cell_to_vertex] for p in self.partitions]) + np.testing.assert_array_equal(combined, whole) + + def test_partition_holds_only_the_vertices_of_its_cells(self): + for partition in self.partitions: + with self.subTest(first_cell=partition.cell_ids[0]): + self.assertEqual(len(partition.vertex_ids), len(np.unique(partition.cell_to_vertex))) + self.assertEqual(len(partition.vertex_ids), len(partition.x_vertices)) + + def test_partition_id_out_of_range_is_rejected(self): + for partition_id in (0, N_PARTITIONS + 1): + with self.subTest(partition_id=partition_id), self.assertRaises(AssertionError): + self.mesh.get_partition(partition_id, N_PARTITIONS) + + +if __name__ == "__main__": + unittest.main() From c60f9f93074cdabe0704a3191d68810103598b2b Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Tue, 20 Oct 2026 16:46:51 +0200 Subject: [PATCH 05/13] Require one MPI rank per partition of a partitioned mesh Each rank computes exactly one partition, so a run on fewer ranks than the mesh has partitions skipped the others without notice; the grid distribution check only finds columns held twice, not missing ones. GridConfig.num_partitions now also counts the partitions of a partitioned mesh, by its part.N.header files, and a mismatch with the number of ranks is a usage error. --allow-partition-mismatch turns it into a warning, e.g. to compute a single partition with --use-part. Assisted-by: Claude Code:claude-opus-5-5 --- CHANGELOG.md | 1 + README.md | 9 ++++++--- src/ebfm/core/cli.py | 31 +++++++++++++++++++++++++++++ src/ebfm/core/config/grid.py | 5 +++-- src/ebfm/main.py | 4 ++++ tests/core/test_cli.py | 38 +++++++++++++++++++++++++++++++++++- 6 files changed, 82 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bda52a3d..7a8367ab 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,7 @@ SPDX-License-Identifier: CC-BY-4.0 * 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/XXX * 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/XXX * 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. https://github.com/EBFMorg/EBFM/pull/XXX +* 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/XXX # v0.8.0 diff --git a/README.md b/README.md index 35b4c50a..30b08856 100644 --- a/README.md +++ b/README.md @@ -216,19 +216,22 @@ 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 ``` 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, a MATLAB grid 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. ### Getting the example data @@ -337,7 +340,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 \ diff --git a/src/ebfm/core/cli.py b/src/ebfm/core/cli.py index 02b6ccaf..45d287b2 100644 --- a/src/ebfm/core/cli.py +++ b/src/ebfm/core/cli.py @@ -357,6 +357,13 @@ def parse_time_step(value: str) -> float | str: "communicates the ids of all columns; skip it once a setup is known to distribute its grid correctly.", ) + parallel_group.add_argument( + "--allow-partition-mismatch", + action="store_true", + help="Allow a number of MPI ranks that differs from the number of partitions of a partitioned Elmer mesh, " + "e.g. to compute only some partitions. Each rank still computes one partition.", + ) + logger_group = parser.add_argument_group("logging configuration") logger_group.add_argument( @@ -634,3 +641,27 @@ def validate_restart_parallel_compat(args: Namespace, comm) -> None: f"Restart files (--restart-dir, --restart-init) are not supported for runs on several MPI ranks " f"({comm.size} ranks). Please run on a single rank." ) + + +def validate_partition_count(args: Namespace, grid_config, comm, logger) -> None: + """Check that a partitioned Elmer mesh has exactly one partition per MPI rank. + + Must be called on every rank of comm. With --allow-partition-mismatch a mismatch is only a warning. + + @param[in] args parsed CLI arguments + @param[in] grid_config resolved GridConfig + @param[in] comm EBFM's MPI communicator + @param[in] logger logger for the warning + """ + if not grid_config.is_partitioned or grid_config.num_partitions == comm.size: + return + + message = ( + f"The partitioned mesh has {grid_config.num_partitions} partitions, but EBFM runs on {comm.size} MPI " + "ranks. Each rank computes exactly one partition, so a run on " + f"{grid_config.num_partitions} ranks is needed to compute all of them." + ) + if not args.allow_partition_mismatch: + _cli_error(f"{message} Pass --allow-partition-mismatch to run anyway.") + if comm.rank == 0: + logger.warning(message) diff --git a/src/ebfm/core/config/grid.py b/src/ebfm/core/config/grid.py index 01e157c9..707c0d41 100644 --- a/src/ebfm/core/config/grid.py +++ b/src/ebfm/core/config/grid.py @@ -27,8 +27,8 @@ class GridConfig: is_partitioned: bool # Whether the grid is partitioned is_unstructured: bool # Whether the grid is unstructured partition_id: int | None # Partition ID (only relevant if is_partitioned is True or num_partitions > 1) - # Number of blocks EBFM splits the columns of a MATLAB grid or a non-partitioned Elmer mesh into, one per MPI rank; - # 1 for a grid that is not split + # Number of partitions of the grid, one per MPI rank: of the mesh files if is_partitioned is True, else the + # number of blocks EBFM splits a MATLAB grid or a non-partitioned Elmer mesh into. 1 for a grid that is not split. num_partitions: int elmer_mesh_crs_epsg: int # EPSG code of Elmer mesh coordinates use_shading: bool # Whether to use shading for the grid @@ -85,6 +85,7 @@ def __init__(self, args: Namespace, num_ranks: int = 1, rank: int = 0): ), "Internal error: partitioned grid configuration requires NetCDF DEM mesh input." logger.info("Using partitioned grid...") self.partition_id = args.use_part + self.num_partitions = len(list(Path(elmer_mesh).glob("part.*.header"))) logger.info(f"{self.partition_id=}") else: logger.info("Using non-partitioned grid...") diff --git a/src/ebfm/main.py b/src/ebfm/main.py index 29ad225e..c6e213cb 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -30,6 +30,7 @@ validate_use_part_parallel_compat, validate_forcing_parallel_compat, validate_restart_parallel_compat, + validate_partition_count, validate_shading_parallel_compat, ) @@ -225,6 +226,9 @@ def _main_impl(): # Ensure shading routine is only used in runs on a single rank. validate_shading_parallel_compat(grid_config, ebfm_comm) + # Each rank computes exactly one partition of a partitioned mesh + validate_partition_count(args, grid_config, ebfm_comm, logger) + logger.debug("Successfully completed consistency checks.") # Model setup & initialization diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py index 5fd00e7d..1231a352 100644 --- a/tests/core/test_cli.py +++ b/tests/core/test_cli.py @@ -19,8 +19,9 @@ validate_forcing_parallel_compat, validate_grid_parallel_compat, validate_restart_parallel_compat, - validate_shading_parallel_compat, + validate_partition_count, validate_shading_coupling_compat, + validate_shading_parallel_compat, validate_use_part_parallel_compat, ) @@ -543,6 +544,41 @@ def test_restart_on_several_ranks_rejected_for_partitioned_mesh(self): self.assertEqual(ctx.exception.code, 2) +class TestValidatePartitionCount(unittest.TestCase): + """A partitioned mesh needs one MPI rank per partition, unless --allow-partition-mismatch is passed.""" + + def setUp(self): + # Ensure _parser is set by running a valid parse first. + parse_cli_args(_MATLAB) + self.logger = MagicMock() + + def _validate(self, is_partitioned=True, num_partitions=4, comm_size=4, rank=0, allow=False): + args = _parse(*(["--allow-partition-mismatch"] if allow else [])) + grid_config = MagicMock(is_partitioned=is_partitioned, num_partitions=num_partitions) + comm = MagicMock(size=comm_size, rank=rank) + validate_partition_count(args, grid_config, comm, self.logger) + + def test_one_rank_per_partition_ok(self): + self._validate() # must not raise + self.logger.warning.assert_not_called() + + def test_non_partitioned_mesh_ok(self): + self._validate(is_partitioned=False, num_partitions=1) # must not raise + + def test_mismatch_rejected(self): + with self.assertRaises(SystemExit) as ctx: + self._validate(num_partitions=128) + self.assertEqual(ctx.exception.code, 2) + + def test_mismatch_allowed_warns_on_rank_0_only(self): + self._validate(num_partitions=128, allow=True) + self.logger.warning.assert_called_once() + + self.logger.reset_mock() + self._validate(num_partitions=128, rank=1, allow=True) + self.logger.warning.assert_not_called() + + class TestCliCalendarArgument(unittest.TestCase): def test_calendar_default(self): args = parse_cli_args(_MATLAB) From 95712bff776274786d5b2ff2a8dc90f2a02659ec Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Tue, 20 Oct 2026 17:06:22 +0200 Subject: [PATCH 06/13] Split NetCDF grids over the MPI ranks A NetCDF grid is now split into contiguous blocks of its columns like a MATLAB grid, and the startup check accepts it on several ranks. GridConfig now splits every non-partitioned grid. Assisted-by: Claude Code:claude-opus-5-5 --- CHANGELOG.md | 1 + README.md | 6 ++-- src/ebfm/core/cli.py | 24 ++++---------- src/ebfm/core/config/grid.py | 7 ++-- tests/core/config/test_grid.py | 58 ++++++++++++++++++++++++++++++++++ tests/core/test_cli.py | 4 +-- 6 files changed, 74 insertions(+), 26 deletions(-) create mode 100644 tests/core/config/test_grid.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a8367ab..907c17aa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,7 @@ SPDX-License-Identifier: CC-BY-4.0 * 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/XXX * 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. https://github.com/EBFMorg/EBFM/pull/XXX * 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/XXX +* 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/XXX # v0.8.0 diff --git a/README.md b/README.md index 30b08856..b83296ee 100644 --- a/README.md +++ b/README.md @@ -227,7 +227,7 @@ mpirun -n 4 ebfm --elmer-mesh examples/greenland_mesh_v0/MESH/partitioning.128/ 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, a MATLAB grid and a non-partitioned Elmer mesh run on several MPI ranks (see [Running on several MPI ranks](#running-on-several-mpi-ranks)). +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`. @@ -293,7 +293,7 @@ Note that it is necessary to provide a `--start-time` consistent with the given ### Running on several MPI ranks -EBFM splits a MATLAB grid or a non-partitioned Elmer mesh over the MPI ranks of a run (requires `pip install 'ebfm[mpi]'`): +EBFM splits a MATLAB or NetCDF grid or a non-partitioned Elmer mesh 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 @@ -303,7 +303,7 @@ mpirun -np 4 ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading * 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`, `--netcdf-mesh`, and + `--forcing-dir`, `--restart-dir`/`--restart-init`, and `--netcdf-dem-mesh-unstructured` on a non-partitioned Elmer mesh. ### Performance and Profiling Runs diff --git a/src/ebfm/core/cli.py b/src/ebfm/core/cli.py index 45d287b2..8f50e0f8 100644 --- a/src/ebfm/core/cli.py +++ b/src/ebfm/core/cli.py @@ -569,8 +569,8 @@ def validate_shading_parallel_compat(grid_config, comm) -> None: def validate_grid_parallel_compat(grid_config, comm) -> None: """Check that a run on several MPI ranks uses a grid that runs there. - EBFM splits a MATLAB grid and a non-partitioned Elmer mesh over the ranks; of a partitioned Elmer mesh, each rank - computes one partition. + EBFM splits a non-partitioned grid over the ranks; of a partitioned Elmer mesh, each rank computes one partition. + Only an ELMERXIOS grid needs a partitioned mesh there. Must be called on every rank of comm. @@ -581,21 +581,11 @@ def validate_grid_parallel_compat(grid_config, comm) -> None: # Every grid runs on a single rank return - if grid_config.grid_type is GridInputType.MATLAB: - # EBFM splits the columns over the ranks - return - - if grid_config.grid_type is GridInputType.NETCDF: - rejected_grid = "--netcdf-mesh" - elif grid_config.grid_type is GridInputType.ELMERXIOS and not grid_config.is_partitioned: - rejected_grid = "--netcdf-dem-mesh-unstructured without --is-partitioned-elmer-mesh" - else: - return - - _cli_error( - f"{rejected_grid} is not supported for runs on several MPI ranks ({comm.size} ranks). " - "Please run on a single rank." - ) + if grid_config.grid_type is GridInputType.ELMERXIOS and not grid_config.is_partitioned: + _cli_error( + "--netcdf-dem-mesh-unstructured without --is-partitioned-elmer-mesh is not supported for runs on several " + f"MPI ranks ({comm.size} ranks). Please run on a single rank." + ) def validate_use_part_parallel_compat(args: Namespace, comm) -> None: diff --git a/src/ebfm/core/config/grid.py b/src/ebfm/core/config/grid.py index 707c0d41..4d5e5eb7 100644 --- a/src/ebfm/core/config/grid.py +++ b/src/ebfm/core/config/grid.py @@ -28,7 +28,7 @@ class GridConfig: is_unstructured: bool # Whether the grid is unstructured partition_id: int | None # Partition ID (only relevant if is_partitioned is True or num_partitions > 1) # Number of partitions of the grid, one per MPI rank: of the mesh files if is_partitioned is True, else the - # number of blocks EBFM splits a MATLAB grid or a non-partitioned Elmer mesh into. 1 for a grid that is not split. + # number of blocks EBFM splits a non-partitioned grid into. 1 for a grid that is not split. num_partitions: int elmer_mesh_crs_epsg: int # EPSG code of Elmer mesh coordinates use_shading: bool # Whether to use shading for the grid @@ -54,7 +54,7 @@ def __init__(self, args: Namespace, num_ranks: int = 1, rank: int = 0): Initialize grid configuration from command line arguments. @param[in] args command line arguments - @param[in] num_ranks number of MPI ranks, over which a MATLAB grid or a non-partitioned Elmer mesh is split + @param[in] num_ranks number of MPI ranks, over which a non-partitioned grid is split @param[in] rank MPI rank of this process, which computes the block rank + 1 of a split grid """ selected_primary_grids = [ @@ -120,8 +120,7 @@ def __init__(self, args: Namespace, num_ranks: int = 1, rank: int = 0): ) raise Exception("Invalid grid configuration.") - is_split = self.grid_type is GridInputType.MATLAB or (self.is_unstructured and not self.is_partitioned) - if is_split and num_ranks > 1: + if not self.is_partitioned and num_ranks > 1: logger.info(f"Splitting the grid over {num_ranks} MPI ranks...") self.partition_id = rank + 1 self.num_partitions = num_ranks diff --git a/tests/core/config/test_grid.py b/tests/core/config/test_grid.py new file mode 100644 index 00000000..b35df426 --- /dev/null +++ b/tests/core/config/test_grid.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# SPDX-License-Identifier: BSD-3-Clause + +import tempfile +import unittest +from pathlib import Path + +from ebfm.core.cli import parse_cli_args +from ebfm.core.comm import mpi_available +from ebfm.core.config import GridConfig + +_ELMER = ["--elmer-mesh-crs-epsg", "3413"] + + +class TestNumberOfPartitions(unittest.TestCase): + """The number of partitions of the grid: of the mesh files, or of the split over the ranks.""" + + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + self.mesh = Path(self.tmpdir.name) + for partition_id in (1, 2, 3): + (self.mesh / f"part.{partition_id}.header").touch() + (self.mesh / f"part.{partition_id}.nodes").touch() + + def tearDown(self): + self.tmpdir.cleanup() + + @unittest.skipUnless(mpi_available, "requires MPI (pip install 'ebfm[mpi]')") + def test_partitioned_mesh(self): + args = parse_cli_args( + ["--elmer-mesh", str(self.mesh), *_ELMER, "--is-partitioned-elmer-mesh", "--netcdf-dem-mesh", "dem.nc"] + ) + args.use_part = 1 + + self.assertEqual(GridConfig(args).num_partitions, 3) + + def test_non_partitioned_mesh(self): + args = parse_cli_args(["--elmer-mesh", str(self.mesh), *_ELMER]) + + self.assertEqual(GridConfig(args).num_partitions, 1) + + def test_non_partitioned_elmer_mesh_is_split_over_the_ranks(self): + args = parse_cli_args(["--elmer-mesh", str(self.mesh), *_ELMER]) + + self.assertEqual(GridConfig(args, num_ranks=4).num_partitions, 4) + + def test_matlab_and_netcdf_grids_are_split_over_the_ranks(self): + for grid_args in (["--matlab-mesh", "mesh.mat"], ["--netcdf-mesh", "mesh.nc"]): + with self.subTest(grid_args[0]): + grid_config = GridConfig(parse_cli_args(grid_args), num_ranks=4, rank=2) + + self.assertEqual(grid_config.num_partitions, 4) + self.assertEqual(grid_config.partition_id, 3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py index 1231a352..fb0001ce 100644 --- a/tests/core/test_cli.py +++ b/tests/core/test_cli.py @@ -415,12 +415,12 @@ def test_shading_on_several_ranks_rejected(self): class TestValidateGridParallelCompat(unittest.TestCase): - """validate_grid_parallel_compat rejects a NetCDF grid and a non-partitioned ELMERXIOS grid on several ranks.""" + """validate_grid_parallel_compat rejects a non-partitioned ELMERXIOS grid on several ranks.""" # name: (grid_type, is_partitioned, grid option named by the error on several ranks, or None if accepted) _CASES = { "--matlab-mesh": (GridInputType.MATLAB, False, None), - "--netcdf-mesh": (GridInputType.NETCDF, False, "--netcdf-mesh"), + "--netcdf-mesh": (GridInputType.NETCDF, False, None), "--elmer-mesh": (GridInputType.ELMER, False, None), "--elmer-mesh --netcdf-dem-mesh": (GridInputType.CUSTOM, False, None), "--elmer-mesh --netcdf-dem-mesh-unstructured": ( From 01a4f7fe2a4953718a7354cffe5d30dcb25a01f7 Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Tue, 20 Oct 2026 18:58:45 +0200 Subject: [PATCH 07/13] Reject restart files for partitioned meshes A partitioned mesh skipped writing the restart file with only a warning at the end of the run, also on a single rank with --use-part. --restart-dir and --restart-init are now a usage error for partitioned meshes, as on several ranks, and the restart writer runs in every run: it writes nothing without --restart-dir. That makes --dump-reference contain timelastsnow_netCDF on partitioned meshes too, so the partitioned-mesh test compares it as well. Assisted-by: Claude Code:claude-opus-5-5 --- CHANGELOG.md | 1 + src/ebfm/core/cli.py | 19 ++++++++---- src/ebfm/main.py | 12 +++----- tests/core/test_cli.py | 39 ++++++++++++++----------- tests/core/test_partitioned_mesh_run.py | 3 -- 5 files changed, 41 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 907c17aa..2c6aa9ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,7 @@ SPDX-License-Identifier: CC-BY-4.0 * 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. https://github.com/EBFMorg/EBFM/pull/XXX * 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/XXX * 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/XXX +* Restart files (`--restart-dir`, `--restart-init`) with a partitioned Elmer mesh now stop the run with an error, also on a single rank. Such a run previously skipped writing the restart file with a warning at its end. https://github.com/EBFMorg/EBFM/pull/XXX # v0.8.0 diff --git a/src/ebfm/core/cli.py b/src/ebfm/core/cli.py index 8f50e0f8..18c17cec 100644 --- a/src/ebfm/core/cli.py +++ b/src/ebfm/core/cli.py @@ -618,19 +618,28 @@ def validate_forcing_parallel_compat(args: Namespace, comm) -> None: ) -def validate_restart_parallel_compat(args: Namespace, comm) -> None: - """Check that no restart files are read or written in a run on several MPI ranks. +def validate_restart_compat(args: Namespace, grid_config, comm) -> None: + """Check that restart files are only read and written for a whole grid on a single MPI rank. Must be called on every rank of comm. - @param[in] args parsed CLI arguments - @param[in] comm EBFM's MPI communicator + @param[in] args parsed CLI arguments + @param[in] grid_config resolved GridConfig + @param[in] comm EBFM's MPI communicator """ - if (args.restart_dir is not None or args.restart_init) and comm.size > 1: + if args.restart_dir is None and not args.restart_init: + return + + if comm.size > 1: _cli_error( f"Restart files (--restart-dir, --restart-init) are not supported for runs on several MPI ranks " f"({comm.size} ranks). Please run on a single rank." ) + if grid_config.is_partitioned: + _cli_error( + "Restart files (--restart-dir, --restart-init) are not supported for partitioned meshes " + "(--is-partitioned-elmer-mesh)." + ) def validate_partition_count(args: Namespace, grid_config, comm, logger) -> None: diff --git a/src/ebfm/main.py b/src/ebfm/main.py index c6e213cb..d088703f 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -29,8 +29,8 @@ validate_grid_parallel_compat, validate_use_part_parallel_compat, validate_forcing_parallel_compat, - validate_restart_parallel_compat, validate_partition_count, + validate_restart_compat, validate_shading_parallel_compat, ) @@ -220,8 +220,8 @@ def _main_impl(): # Ensure forcing files are only read in runs on a single rank. validate_forcing_parallel_compat(args, ebfm_comm) - # Ensure restart files are only read and written in runs on a single rank. - validate_restart_parallel_compat(args, ebfm_comm) + # Ensure restart files are only read and written for a whole grid on a single rank. + validate_restart_compat(args, grid_config, ebfm_comm) # Ensure shading routine is only used in runs on a single rank. validate_shading_parallel_compat(grid_config, ebfm_comm) @@ -410,11 +410,7 @@ def _main_impl(): io, OUTFILE = LOOP_write_to_file.main(OUTFILE, io, OUT, grid, t, time, column, column_gather) # Write restart file - # TODO: should be supported for all cases to avoid case distinction here - if not grid["is_partitioned"]: - FINAL_create_restart_file.main(OUT, io, args.restart_dir) - else: - logger.warning("Skipping writing of restart file for partitioned runs.") + FINAL_create_restart_file.main(OUT, io, args.restart_dir) logger.info("Time loop completed.") diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py index fb0001ce..29e7237f 100644 --- a/tests/core/test_cli.py +++ b/tests/core/test_cli.py @@ -18,8 +18,8 @@ parse_cli_args, validate_forcing_parallel_compat, validate_grid_parallel_compat, - validate_restart_parallel_compat, validate_partition_count, + validate_restart_compat, validate_shading_coupling_compat, validate_shading_parallel_compat, validate_use_part_parallel_compat, @@ -505,8 +505,8 @@ def test_forcing_on_several_ranks_rejected(self): self.assertEqual(ctx.exception.code, 2) -class TestValidateRestartParallelCompat(unittest.TestCase): - """validate_restart_parallel_compat raises a CLI error when restart files are used in a run on several ranks.""" +class TestValidateRestartCompat(unittest.TestCase): + """validate_restart_compat rejects restart files in a run on several ranks or on a partitioned mesh.""" _RESTART_ARGS = { "--restart-dir": ["--restart-dir", "restart/"], @@ -518,30 +518,35 @@ def setUp(self): # Ensure _parser is set by running a valid parse first. parse_cli_args(_MATLAB) - def _validate(self, restart_args, comm_size: int, base: list[str] = _MATLAB): - validate_restart_parallel_compat(_parse(*restart_args, base=base), MagicMock(size=comm_size)) + def _validate(self, restart_args, comm_size=1, is_partitioned=False): + grid_config = MagicMock(is_partitioned=is_partitioned) + validate_restart_compat(_parse(*restart_args), grid_config, MagicMock(size=comm_size)) + + def _assert_rejected(self, **kwargs): + for name, restart_args in self._RESTART_ARGS.items(): + with self.subTest(name), self.assertRaises(SystemExit) as ctx: + self._validate(restart_args, **kwargs) + self.assertEqual(ctx.exception.code, 2) def test_restart_on_single_rank_ok(self): for name, restart_args in self._RESTART_ARGS.items(): with self.subTest(name): - self._validate(restart_args, comm_size=1) # must not raise + self._validate(restart_args) # must not raise def test_no_restart_on_several_ranks_ok(self): self._validate([], comm_size=4) # must not raise + def test_no_restart_on_several_ranks_with_partitioned_mesh_ok(self): + self._validate([], comm_size=4, is_partitioned=True) # must not raise + def test_restart_on_several_ranks_rejected(self): - for name, restart_args in self._RESTART_ARGS.items(): - with self.subTest(name), self.assertRaises(SystemExit) as ctx: - self._validate(restart_args, comm_size=4) - self.assertEqual(ctx.exception.code, 2) + self._assert_rejected(comm_size=4) - @unittest.skipUnless(mpi_available, "requires MPI (pip install 'ebfm[mpi]')") - def test_restart_on_several_ranks_rejected_for_partitioned_mesh(self): - """A partitioned mesh runs on several ranks, but without restart files.""" - for name, restart_args in self._RESTART_ARGS.items(): - with self.subTest(name), self.assertRaises(SystemExit) as ctx: - self._validate(restart_args, comm_size=4, base=_PARTITIONED_ELMER) - self.assertEqual(ctx.exception.code, 2) + def test_restart_with_partitioned_mesh_rejected(self): + self._assert_rejected(is_partitioned=True) + + def test_restart_on_several_ranks_with_partitioned_mesh_rejected(self): + self._assert_rejected(comm_size=4, is_partitioned=True) class TestValidatePartitionCount(unittest.TestCase): diff --git a/tests/core/test_partitioned_mesh_run.py b/tests/core/test_partitioned_mesh_run.py index 82d9b9db..5b58101a 100644 --- a/tests/core/test_partitioned_mesh_run.py +++ b/tests/core/test_partitioned_mesh_run.py @@ -123,9 +123,6 @@ def test_partitioned_run_equals_serial_run(tmp_path): serial = _run(tmp_path, mesh_dir, 1) partitioned = _run(tmp_path, mesh_dir, N_PARTITIONS, "--is-partitioned-elmer-mesh") - # Set by the restart writer, which partitioned runs skip - del serial["timelastsnow_netCDF"] - assert sorted(partitioned) == sorted(serial) for key in serial: np.testing.assert_array_equal(partitioned[key], serial[key], err_msg=key) From e74d826ef3843ddc457ea1daebcfb2459903297b Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Tue, 20 Oct 2026 18:58:46 +0200 Subject: [PATCH 08/13] Test that parallel runs on an Elmer mesh match a serial run The serial-vs-parallel comparison gets Elmer cases: a small synthetic Elmer mesh, committed under tests/data, on 3 ranks, and the same mesh with fake coupling to Elmer/Ice, ICON and ICON-Land on 2 ranks. tests/data also holds an XIOS DEM for the mesh, kept out of Git LFS. CI also runs the pseudo-coupled Greenland example on 4 ranks. Assisted-by: Claude Code:claude-opus-5-5 --- .gitattributes | 1 + .github/workflows/build-and-test.yaml | 6 +- tests/core/test_parallel_runs.py | 42 +++- tests/data/README.md | 16 ++ tests/data/elmer_mesh/MESH/mesh.elements | 192 ++++++++++++++++++ .../elmer_mesh/MESH/mesh.elements.license | 5 + tests/data/elmer_mesh/MESH/mesh.header | 1 + .../data/elmer_mesh/MESH/mesh.header.license | 5 + tests/data/elmer_mesh/MESH/mesh.nodes | 117 +++++++++++ tests/data/elmer_mesh/MESH/mesh.nodes.license | 5 + tests/data/elmer_mesh/xios_dem.nc | Bin 0 -> 14360 bytes tests/data/elmer_mesh/xios_dem.nc.license | 5 + 12 files changed, 382 insertions(+), 13 deletions(-) create mode 100644 tests/data/README.md create mode 100644 tests/data/elmer_mesh/MESH/mesh.elements create mode 100644 tests/data/elmer_mesh/MESH/mesh.elements.license create mode 100644 tests/data/elmer_mesh/MESH/mesh.header create mode 100644 tests/data/elmer_mesh/MESH/mesh.header.license create mode 100644 tests/data/elmer_mesh/MESH/mesh.nodes create mode 100644 tests/data/elmer_mesh/MESH/mesh.nodes.license create mode 100644 tests/data/elmer_mesh/xios_dem.nc create mode 100644 tests/data/elmer_mesh/xios_dem.nc.license diff --git a/.gitattributes b/.gitattributes index c48cd4a9..eaf2946c 100644 --- a/.gitattributes +++ b/.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 diff --git a/.github/workflows/build-and-test.yaml b/.github/workflows/build-and-test.yaml index f20a3763..7b1c65ea 100644 --- a/.github/workflows/build-and-test.yaml +++ b/.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 diff --git a/tests/core/test_parallel_runs.py b/tests/core/test_parallel_runs.py index 9d717906..f6f144e4 100644 --- a/tests/core/test_parallel_runs.py +++ b/tests/core/test_parallel_runs.py @@ -2,9 +2,9 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""A run of the MATLAB example on several MPI ranks gives exactly the results of a run on a single rank. +"""A run on several MPI ranks gives exactly the results of a run on a single rank. -EBFM runs once on a single rank and once per rank count under mpirun, and the runs must give byte-identical output +Each case runs EBFM once on a single rank and once per rank count under mpirun, and requires byte-identical output files and equal --dump-reference snapshots. """ @@ -21,6 +21,7 @@ REPO_ROOT = Path(__file__).resolve().parents[2] MATLAB_MESH = REPO_ROOT / "examples" / "dem_and_mask.mat" +ELMER_MESH = REPO_ROOT / "tests" / "data" / "elmer_mesh" / "MESH" TIMEOUT = 300 requires_mpirun = pytest.mark.skipif( @@ -31,7 +32,20 @@ # One day: the output file gets its first record after 8 time steps of 3 hours _TIME = ["--random-seed", "1", "--start-time", "1979-07-01T00:00:00Z", "--end-time", "1979-07-02T00:00:00Z"] -ARGS = ["--matlab-mesh", str(MATLAB_MESH), "--no-shading", *_TIME] +_ELMER = ["--elmer-mesh", str(ELMER_MESH), "--elmer-mesh-crs-epsg", "3413"] + +CASES = { + "matlab": ["--matlab-mesh", str(MATLAB_MESH), "--no-shading", *_TIME], + "elmer": [*_ELMER, *_TIME], + "fake-coupling": [ + *_ELMER, + "--couple-to-elmer-ice", + "--couple-to-icon-atmo", + "--couple-to-icon-land", + "--fake-coupling", + *_TIME, + ], +} def _run(args: list[str], n_ranks: int, workdir: Path) -> Path: @@ -70,14 +84,20 @@ def _assert_same_results(serial: Path, parallel: Path) -> None: np.testing.assert_array_equal(actual[key], expected[key], err_msg=key) -@pytest.fixture(scope="module") -def serial(tmp_path_factory) -> Path: - return _run(ARGS, 1, tmp_path_factory.mktemp("serial") / "run") - - @requires_mpirun -@pytest.mark.parametrize("n_ranks", [2, 3, 4]) -def test_parallel_run_equals_serial_run(tmp_path, serial, n_ranks): - parallel = _run(ARGS, n_ranks, tmp_path / f"np{n_ranks}") +@pytest.mark.parametrize( + "case, n_ranks", + [ + ("matlab", 2), + ("matlab", 3), + ("matlab", 4), + ("elmer", 3), + ("fake-coupling", 2), + ], + ids=str, +) +def test_parallel_run_equals_serial_run(tmp_path, case, n_ranks): + serial = _run(CASES[case], 1, tmp_path / "serial") + parallel = _run(CASES[case], n_ranks, tmp_path / f"np{n_ranks}") _assert_same_results(serial, parallel) diff --git a/tests/data/README.md b/tests/data/README.md new file mode 100644 index 00000000..69abdac0 --- /dev/null +++ b/tests/data/README.md @@ -0,0 +1,16 @@ + + +# Test data + +`elmer_mesh` is a small synthetic Elmer mesh for tests that run EBFM end to end: + +* `MESH`: a rectangle of 12 x 8 squares of 5 km, two triangles each, in EPSG:3413, with its lower left corner at + x = -200000 m, y = -2000000 m. Node ids (117) and element ids (192) run from 1 in file order. There is no + `mesh.boundary`, which EBFM does not read. +* `xios_dem.nc`: surface elevation `zs = 500 + 7 * id` and ice thickness `h` per node, stored by node id like Elmer/Ice + XIOS output (see `examples/README.md`). `h` is 0.5 m for the first third of the node ids and `100 + id` otherwise, + so the mesh has glacier and ice-free columns. diff --git a/tests/data/elmer_mesh/MESH/mesh.elements b/tests/data/elmer_mesh/MESH/mesh.elements new file mode 100644 index 00000000..b74896f8 --- /dev/null +++ b/tests/data/elmer_mesh/MESH/mesh.elements @@ -0,0 +1,192 @@ +1 1 303 1 2 15 +2 1 303 1 15 14 +3 1 303 2 3 16 +4 1 303 2 16 15 +5 1 303 3 4 17 +6 1 303 3 17 16 +7 1 303 4 5 18 +8 1 303 4 18 17 +9 1 303 5 6 19 +10 1 303 5 19 18 +11 1 303 6 7 20 +12 1 303 6 20 19 +13 1 303 7 8 21 +14 1 303 7 21 20 +15 1 303 8 9 22 +16 1 303 8 22 21 +17 1 303 9 10 23 +18 1 303 9 23 22 +19 1 303 10 11 24 +20 1 303 10 24 23 +21 1 303 11 12 25 +22 1 303 11 25 24 +23 1 303 12 13 26 +24 1 303 12 26 25 +25 1 303 14 15 28 +26 1 303 14 28 27 +27 1 303 15 16 29 +28 1 303 15 29 28 +29 1 303 16 17 30 +30 1 303 16 30 29 +31 1 303 17 18 31 +32 1 303 17 31 30 +33 1 303 18 19 32 +34 1 303 18 32 31 +35 1 303 19 20 33 +36 1 303 19 33 32 +37 1 303 20 21 34 +38 1 303 20 34 33 +39 1 303 21 22 35 +40 1 303 21 35 34 +41 1 303 22 23 36 +42 1 303 22 36 35 +43 1 303 23 24 37 +44 1 303 23 37 36 +45 1 303 24 25 38 +46 1 303 24 38 37 +47 1 303 25 26 39 +48 1 303 25 39 38 +49 1 303 27 28 41 +50 1 303 27 41 40 +51 1 303 28 29 42 +52 1 303 28 42 41 +53 1 303 29 30 43 +54 1 303 29 43 42 +55 1 303 30 31 44 +56 1 303 30 44 43 +57 1 303 31 32 45 +58 1 303 31 45 44 +59 1 303 32 33 46 +60 1 303 32 46 45 +61 1 303 33 34 47 +62 1 303 33 47 46 +63 1 303 34 35 48 +64 1 303 34 48 47 +65 1 303 35 36 49 +66 1 303 35 49 48 +67 1 303 36 37 50 +68 1 303 36 50 49 +69 1 303 37 38 51 +70 1 303 37 51 50 +71 1 303 38 39 52 +72 1 303 38 52 51 +73 1 303 40 41 54 +74 1 303 40 54 53 +75 1 303 41 42 55 +76 1 303 41 55 54 +77 1 303 42 43 56 +78 1 303 42 56 55 +79 1 303 43 44 57 +80 1 303 43 57 56 +81 1 303 44 45 58 +82 1 303 44 58 57 +83 1 303 45 46 59 +84 1 303 45 59 58 +85 1 303 46 47 60 +86 1 303 46 60 59 +87 1 303 47 48 61 +88 1 303 47 61 60 +89 1 303 48 49 62 +90 1 303 48 62 61 +91 1 303 49 50 63 +92 1 303 49 63 62 +93 1 303 50 51 64 +94 1 303 50 64 63 +95 1 303 51 52 65 +96 1 303 51 65 64 +97 1 303 53 54 67 +98 1 303 53 67 66 +99 1 303 54 55 68 +100 1 303 54 68 67 +101 1 303 55 56 69 +102 1 303 55 69 68 +103 1 303 56 57 70 +104 1 303 56 70 69 +105 1 303 57 58 71 +106 1 303 57 71 70 +107 1 303 58 59 72 +108 1 303 58 72 71 +109 1 303 59 60 73 +110 1 303 59 73 72 +111 1 303 60 61 74 +112 1 303 60 74 73 +113 1 303 61 62 75 +114 1 303 61 75 74 +115 1 303 62 63 76 +116 1 303 62 76 75 +117 1 303 63 64 77 +118 1 303 63 77 76 +119 1 303 64 65 78 +120 1 303 64 78 77 +121 1 303 66 67 80 +122 1 303 66 80 79 +123 1 303 67 68 81 +124 1 303 67 81 80 +125 1 303 68 69 82 +126 1 303 68 82 81 +127 1 303 69 70 83 +128 1 303 69 83 82 +129 1 303 70 71 84 +130 1 303 70 84 83 +131 1 303 71 72 85 +132 1 303 71 85 84 +133 1 303 72 73 86 +134 1 303 72 86 85 +135 1 303 73 74 87 +136 1 303 73 87 86 +137 1 303 74 75 88 +138 1 303 74 88 87 +139 1 303 75 76 89 +140 1 303 75 89 88 +141 1 303 76 77 90 +142 1 303 76 90 89 +143 1 303 77 78 91 +144 1 303 77 91 90 +145 1 303 79 80 93 +146 1 303 79 93 92 +147 1 303 80 81 94 +148 1 303 80 94 93 +149 1 303 81 82 95 +150 1 303 81 95 94 +151 1 303 82 83 96 +152 1 303 82 96 95 +153 1 303 83 84 97 +154 1 303 83 97 96 +155 1 303 84 85 98 +156 1 303 84 98 97 +157 1 303 85 86 99 +158 1 303 85 99 98 +159 1 303 86 87 100 +160 1 303 86 100 99 +161 1 303 87 88 101 +162 1 303 87 101 100 +163 1 303 88 89 102 +164 1 303 88 102 101 +165 1 303 89 90 103 +166 1 303 89 103 102 +167 1 303 90 91 104 +168 1 303 90 104 103 +169 1 303 92 93 106 +170 1 303 92 106 105 +171 1 303 93 94 107 +172 1 303 93 107 106 +173 1 303 94 95 108 +174 1 303 94 108 107 +175 1 303 95 96 109 +176 1 303 95 109 108 +177 1 303 96 97 110 +178 1 303 96 110 109 +179 1 303 97 98 111 +180 1 303 97 111 110 +181 1 303 98 99 112 +182 1 303 98 112 111 +183 1 303 99 100 113 +184 1 303 99 113 112 +185 1 303 100 101 114 +186 1 303 100 114 113 +187 1 303 101 102 115 +188 1 303 101 115 114 +189 1 303 102 103 116 +190 1 303 102 116 115 +191 1 303 103 104 117 +192 1 303 103 117 116 diff --git a/tests/data/elmer_mesh/MESH/mesh.elements.license b/tests/data/elmer_mesh/MESH/mesh.elements.license new file mode 100644 index 00000000..0fa21cf3 --- /dev/null +++ b/tests/data/elmer_mesh/MESH/mesh.elements.license @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: 2026 The EBFM Authors +# +# SPDX-License-Identifier: CC-BY-4.0 +# +# SPDX-Comment: Synthetic test data, see tests/data/README.md. diff --git a/tests/data/elmer_mesh/MESH/mesh.header b/tests/data/elmer_mesh/MESH/mesh.header new file mode 100644 index 00000000..b4d6c911 --- /dev/null +++ b/tests/data/elmer_mesh/MESH/mesh.header @@ -0,0 +1 @@ +117 192 0 diff --git a/tests/data/elmer_mesh/MESH/mesh.header.license b/tests/data/elmer_mesh/MESH/mesh.header.license new file mode 100644 index 00000000..0fa21cf3 --- /dev/null +++ b/tests/data/elmer_mesh/MESH/mesh.header.license @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: 2026 The EBFM Authors +# +# SPDX-License-Identifier: CC-BY-4.0 +# +# SPDX-Comment: Synthetic test data, see tests/data/README.md. diff --git a/tests/data/elmer_mesh/MESH/mesh.nodes b/tests/data/elmer_mesh/MESH/mesh.nodes new file mode 100644 index 00000000..84bc86d7 --- /dev/null +++ b/tests/data/elmer_mesh/MESH/mesh.nodes @@ -0,0 +1,117 @@ +1 -1 -200000.0 -2000000.0 0.0 +2 -1 -195000.0 -2000000.0 0.0 +3 -1 -190000.0 -2000000.0 0.0 +4 -1 -185000.0 -2000000.0 0.0 +5 -1 -180000.0 -2000000.0 0.0 +6 -1 -175000.0 -2000000.0 0.0 +7 -1 -170000.0 -2000000.0 0.0 +8 -1 -165000.0 -2000000.0 0.0 +9 -1 -160000.0 -2000000.0 0.0 +10 -1 -155000.0 -2000000.0 0.0 +11 -1 -150000.0 -2000000.0 0.0 +12 -1 -145000.0 -2000000.0 0.0 +13 -1 -140000.0 -2000000.0 0.0 +14 -1 -200000.0 -1995000.0 0.0 +15 -1 -195000.0 -1995000.0 0.0 +16 -1 -190000.0 -1995000.0 0.0 +17 -1 -185000.0 -1995000.0 0.0 +18 -1 -180000.0 -1995000.0 0.0 +19 -1 -175000.0 -1995000.0 0.0 +20 -1 -170000.0 -1995000.0 0.0 +21 -1 -165000.0 -1995000.0 0.0 +22 -1 -160000.0 -1995000.0 0.0 +23 -1 -155000.0 -1995000.0 0.0 +24 -1 -150000.0 -1995000.0 0.0 +25 -1 -145000.0 -1995000.0 0.0 +26 -1 -140000.0 -1995000.0 0.0 +27 -1 -200000.0 -1990000.0 0.0 +28 -1 -195000.0 -1990000.0 0.0 +29 -1 -190000.0 -1990000.0 0.0 +30 -1 -185000.0 -1990000.0 0.0 +31 -1 -180000.0 -1990000.0 0.0 +32 -1 -175000.0 -1990000.0 0.0 +33 -1 -170000.0 -1990000.0 0.0 +34 -1 -165000.0 -1990000.0 0.0 +35 -1 -160000.0 -1990000.0 0.0 +36 -1 -155000.0 -1990000.0 0.0 +37 -1 -150000.0 -1990000.0 0.0 +38 -1 -145000.0 -1990000.0 0.0 +39 -1 -140000.0 -1990000.0 0.0 +40 -1 -200000.0 -1985000.0 0.0 +41 -1 -195000.0 -1985000.0 0.0 +42 -1 -190000.0 -1985000.0 0.0 +43 -1 -185000.0 -1985000.0 0.0 +44 -1 -180000.0 -1985000.0 0.0 +45 -1 -175000.0 -1985000.0 0.0 +46 -1 -170000.0 -1985000.0 0.0 +47 -1 -165000.0 -1985000.0 0.0 +48 -1 -160000.0 -1985000.0 0.0 +49 -1 -155000.0 -1985000.0 0.0 +50 -1 -150000.0 -1985000.0 0.0 +51 -1 -145000.0 -1985000.0 0.0 +52 -1 -140000.0 -1985000.0 0.0 +53 -1 -200000.0 -1980000.0 0.0 +54 -1 -195000.0 -1980000.0 0.0 +55 -1 -190000.0 -1980000.0 0.0 +56 -1 -185000.0 -1980000.0 0.0 +57 -1 -180000.0 -1980000.0 0.0 +58 -1 -175000.0 -1980000.0 0.0 +59 -1 -170000.0 -1980000.0 0.0 +60 -1 -165000.0 -1980000.0 0.0 +61 -1 -160000.0 -1980000.0 0.0 +62 -1 -155000.0 -1980000.0 0.0 +63 -1 -150000.0 -1980000.0 0.0 +64 -1 -145000.0 -1980000.0 0.0 +65 -1 -140000.0 -1980000.0 0.0 +66 -1 -200000.0 -1975000.0 0.0 +67 -1 -195000.0 -1975000.0 0.0 +68 -1 -190000.0 -1975000.0 0.0 +69 -1 -185000.0 -1975000.0 0.0 +70 -1 -180000.0 -1975000.0 0.0 +71 -1 -175000.0 -1975000.0 0.0 +72 -1 -170000.0 -1975000.0 0.0 +73 -1 -165000.0 -1975000.0 0.0 +74 -1 -160000.0 -1975000.0 0.0 +75 -1 -155000.0 -1975000.0 0.0 +76 -1 -150000.0 -1975000.0 0.0 +77 -1 -145000.0 -1975000.0 0.0 +78 -1 -140000.0 -1975000.0 0.0 +79 -1 -200000.0 -1970000.0 0.0 +80 -1 -195000.0 -1970000.0 0.0 +81 -1 -190000.0 -1970000.0 0.0 +82 -1 -185000.0 -1970000.0 0.0 +83 -1 -180000.0 -1970000.0 0.0 +84 -1 -175000.0 -1970000.0 0.0 +85 -1 -170000.0 -1970000.0 0.0 +86 -1 -165000.0 -1970000.0 0.0 +87 -1 -160000.0 -1970000.0 0.0 +88 -1 -155000.0 -1970000.0 0.0 +89 -1 -150000.0 -1970000.0 0.0 +90 -1 -145000.0 -1970000.0 0.0 +91 -1 -140000.0 -1970000.0 0.0 +92 -1 -200000.0 -1965000.0 0.0 +93 -1 -195000.0 -1965000.0 0.0 +94 -1 -190000.0 -1965000.0 0.0 +95 -1 -185000.0 -1965000.0 0.0 +96 -1 -180000.0 -1965000.0 0.0 +97 -1 -175000.0 -1965000.0 0.0 +98 -1 -170000.0 -1965000.0 0.0 +99 -1 -165000.0 -1965000.0 0.0 +100 -1 -160000.0 -1965000.0 0.0 +101 -1 -155000.0 -1965000.0 0.0 +102 -1 -150000.0 -1965000.0 0.0 +103 -1 -145000.0 -1965000.0 0.0 +104 -1 -140000.0 -1965000.0 0.0 +105 -1 -200000.0 -1960000.0 0.0 +106 -1 -195000.0 -1960000.0 0.0 +107 -1 -190000.0 -1960000.0 0.0 +108 -1 -185000.0 -1960000.0 0.0 +109 -1 -180000.0 -1960000.0 0.0 +110 -1 -175000.0 -1960000.0 0.0 +111 -1 -170000.0 -1960000.0 0.0 +112 -1 -165000.0 -1960000.0 0.0 +113 -1 -160000.0 -1960000.0 0.0 +114 -1 -155000.0 -1960000.0 0.0 +115 -1 -150000.0 -1960000.0 0.0 +116 -1 -145000.0 -1960000.0 0.0 +117 -1 -140000.0 -1960000.0 0.0 diff --git a/tests/data/elmer_mesh/MESH/mesh.nodes.license b/tests/data/elmer_mesh/MESH/mesh.nodes.license new file mode 100644 index 00000000..0fa21cf3 --- /dev/null +++ b/tests/data/elmer_mesh/MESH/mesh.nodes.license @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: 2026 The EBFM Authors +# +# SPDX-License-Identifier: CC-BY-4.0 +# +# SPDX-Comment: Synthetic test data, see tests/data/README.md. diff --git a/tests/data/elmer_mesh/xios_dem.nc b/tests/data/elmer_mesh/xios_dem.nc new file mode 100644 index 0000000000000000000000000000000000000000..c24d6c659f6617513262261c5ed85cfc98131f3f GIT binary patch literal 14360 zcmeI2+iw(A9LLYjwq-3#TcEY&A}lvSs7ni_6r|F+v!#W03vK1%C8e-iyFz!#c8fQn z^@0kb;{7fNJQ#g7ypc4<_-70XnqjGn2g zxk~mK45-ORHrk7Oy~D>YT=vqRwk{B;lvPe3V2)YXK3FLkPhT+K-!P)8L#p~YdKY`= z5}YcT1l1|+fI98%TNN@XF6_xul}qRWn&7XC`Oqc67xPy&c_& z{@z$?XIzCJ!WA|s)*B>RyW(E|7QcVrPXanpu>Ud@mW28fsr=sLV1t_;PA63b^`I|T#7mB2O8faju0xyH zG{MAC9aLuimnN2K6Tj0(P^bLu&E>pBOsub8`}a)L)wbm5*ZZ`GErGydIdx}3B^nP+ z*%ohJ~1yvVyuyu-EL@x0UiPs&^S*(CnZ9@jIT zXJ$dEpE^(cD$bWw$+|$1*eK0;dc<6$63S9Sei^C2kg1@v7td!Ub~D?iSa+vaK2qj1 zQ%dI7fO$rn#Las1Sk@y$k6+o|^!K0B-nf?aroSvpJ6^8mc4{544p;}Q1J(iSfOWt+ zU>&dySO=^F)&c9l&3B-$Cms**SHF30X6EK))%7xHT7GrJyf-j!E6sa>t9w~x9k32q z2do3u0qcNuz&dcNIv^jlZtT}x`Nin{iu6A*$rq$=f9Z0(ue0v8%?5T{SiZeFu6f zI*cwuSDmgHECc(T}0i z=zZv6bOwC@okQo*qv$8lPobYdA3{Its>m4nc~@P8@fTcm@H9MASBSq# z9LA3jKT7;H;>U=;PW%nxZxV;`u^Zhw%@He@Ofz;vWUn)n&wXWW1KF3{Eg@2gPp_FZ7X3cFw(unt%UtOM2o>wtCO zc6DGQC-y%j98+oy7hX^cyg|G+~!xFd~ zmclZ)2bRMMsD*nW3M-)w>R}bEh6Y#zjj$G)U>&T7`(Oh!!~L)kTA&pkfEc(Chc?&* h?a%?6VGDFZ7bKt?w!${(fnMl?2VpzxfSs@__9rt4b)EnK literal 0 HcmV?d00001 diff --git a/tests/data/elmer_mesh/xios_dem.nc.license b/tests/data/elmer_mesh/xios_dem.nc.license new file mode 100644 index 00000000..0fa21cf3 --- /dev/null +++ b/tests/data/elmer_mesh/xios_dem.nc.license @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: 2026 The EBFM Authors +# +# SPDX-License-Identifier: CC-BY-4.0 +# +# SPDX-Comment: Synthetic test data, see tests/data/README.md. From de9de2d49733ab07c44fe9f7d18265660f27a58b Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Tue, 20 Oct 2026 18:58:47 +0200 Subject: [PATCH 09/13] Document running EBFM on several MPI ranks Assisted-by: Claude Code:claude-opus-5-5 --- README.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index b83296ee..91f2e475 100644 --- a/README.md +++ b/README.md @@ -293,18 +293,22 @@ Note that it is necessary to provide a `--start-time` consistent with the given ### Running on several MPI ranks -EBFM splits a MATLAB or NetCDF grid or a non-partitioned Elmer mesh 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. +* 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). * 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 `--netcdf-dem-mesh-unstructured` on a non-partitioned Elmer mesh. +* `--no-check-grid-distribution` skips the start-up check that every column is + held by exactly one rank. ### Performance and Profiling Runs From 230ccc91374f15ce313b8ffe5eb5cbafab297dfa Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Tue, 20 Oct 2026 19:05:44 +0200 Subject: [PATCH 10/13] Link the changelog entries of the MPI parallelization to #197 Assisted-by: Claude Code:claude-opus-5-5 --- CHANGELOG.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c6aa9ea..ba57ece7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,12 +28,12 @@ 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/XXX -* 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/XXX -* 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. https://github.com/EBFMorg/EBFM/pull/XXX -* 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/XXX -* 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/XXX -* Restart files (`--restart-dir`, `--restart-init`) with a partitioned Elmer mesh now stop the run with an error, also on a single rank. Such a run previously skipped writing the restart file with a warning at its end. https://github.com/EBFMorg/EBFM/pull/XXX +* 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. 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 +* Restart files (`--restart-dir`, `--restart-init`) with a partitioned Elmer mesh now stop the run with an error, also on a single rank. Such a run previously skipped writing the restart file with a warning at its end. https://github.com/EBFMorg/EBFM/pull/197 # v0.8.0 From 4477f966d4aad9d583233c79a6a68139afc07684 Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Wed, 21 Oct 2026 09:17:00 +0200 Subject: [PATCH 11/13] Read the XIOS DEM by Elmer node id on several MPI ranks read_dem_xios looked up zs and h by the mesh's local vertex indices. After the split over the ranks each rank renumbers its vertices, so --netcdf-dem-mesh-unstructured was rejected on several ranks. It now looks the values up by Elmer node id, node id i at position i - 1, which holds for every part of the mesh and for any order of mesh.nodes. A file without values for all node ids of the mesh is rejected. Elmer/Ice output for greenland_mesh_v0 stores its nodes in this order; in that mesh it is also the order of mesh.nodes, so serial results are unchanged. With that, the startup check of the grid type on several ranks has nothing left to reject and is removed. Assisted-by: Claude Code:claude-opus-5-5 --- CHANGELOG.md | 1 + README.md | 3 +- examples/README.md | 10 +++-- src/ebfm/core/cli.py | 24 +----------- src/ebfm/main.py | 4 -- src/ebfm/reader.py | 26 ++++++++----- tests/core/test_cli.py | 48 +---------------------- tests/core/test_elmer_mesh.py | 65 ++++++++++++++++++++++++++++++++ tests/core/test_parallel_runs.py | 3 ++ tests/test_reader.py | 4 +- 10 files changed, 97 insertions(+), 91 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ba57ece7..06cd7ac9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,7 @@ SPDX-License-Identifier: CC-BY-4.0 * 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. 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`) with a partitioned Elmer mesh now stop the run with an error, also on a single rank. Such a run previously skipped writing the restart file with a warning at its end. https://github.com/EBFMorg/EBFM/pull/197 # v0.8.0 diff --git a/README.md b/README.md index 91f2e475..3fc378c5 100644 --- a/README.md +++ b/README.md @@ -305,8 +305,7 @@ mpirun -np 4 ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading * 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 - `--netcdf-dem-mesh-unstructured` on a non-partitioned Elmer mesh. + `--forcing-dir` and `--restart-dir`/`--restart-init`. * `--no-check-grid-distribution` skips the start-up check that every column is held by exactly one rank. diff --git a/examples/README.md b/examples/README.md index 57bb5bc6..7c678cdb 100644 --- a/examples/README.md +++ b/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: diff --git a/src/ebfm/core/cli.py b/src/ebfm/core/cli.py index 18c17cec..4dfe5eca 100644 --- a/src/ebfm/core/cli.py +++ b/src/ebfm/core/cli.py @@ -236,7 +236,7 @@ def parse_cli_args(args: list[str] | None = None) -> Namespace: "--netcdf-dem-mesh-unstructured", type=Path, help="Path to an Elmer/Ice XIOS output file with surface elevation (zs) and ice thickness (h) at the nodes " - f"of the {mesh_opts[GridInputType.ELMER]}, in the order of its mesh.nodes. " + f"of the {mesh_opts[GridInputType.ELMER]}, stored by node id (node id i at position i - 1). " f"Optional if using {mesh_opts[GridInputType.ELMER]}.", ) @@ -566,28 +566,6 @@ def validate_shading_parallel_compat(grid_config, comm) -> None: ) -def validate_grid_parallel_compat(grid_config, comm) -> None: - """Check that a run on several MPI ranks uses a grid that runs there. - - EBFM splits a non-partitioned grid over the ranks; of a partitioned Elmer mesh, each rank computes one partition. - Only an ELMERXIOS grid needs a partitioned mesh there. - - Must be called on every rank of comm. - - @param[in] grid_config resolved GridConfig - @param[in] comm EBFM's MPI communicator - """ - if comm.size == 1: - # Every grid runs on a single rank - return - - if grid_config.grid_type is GridInputType.ELMERXIOS and not grid_config.is_partitioned: - _cli_error( - "--netcdf-dem-mesh-unstructured without --is-partitioned-elmer-mesh is not supported for runs on several " - f"MPI ranks ({comm.size} ranks). Please run on a single rank." - ) - - def validate_use_part_parallel_compat(args: Namespace, comm) -> None: """Check that --use-part is not given in a run on several MPI ranks, where every rank would compute it. diff --git a/src/ebfm/main.py b/src/ebfm/main.py index d088703f..f1fdfbb5 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -26,7 +26,6 @@ parse_cli_args, resolve_numba_threads, validate_shading_coupling_compat, - validate_grid_parallel_compat, validate_use_part_parallel_compat, validate_forcing_parallel_compat, validate_partition_count, @@ -214,9 +213,6 @@ def _main_impl(): # See https://github.com/EBFMorg/EBFM/issues/11 for details. validate_shading_coupling_compat(grid_config, coupling_config) - # Ensure a run on several ranks uses a grid that runs there. - validate_grid_parallel_compat(grid_config, ebfm_comm) - # Ensure forcing files are only read in runs on a single rank. validate_forcing_parallel_compat(args, ebfm_comm) diff --git a/src/ebfm/reader.py b/src/ebfm/reader.py index d445494c..48c1479b 100644 --- a/src/ebfm/reader.py +++ b/src/ebfm/reader.py @@ -108,6 +108,9 @@ def read_dem_xios(dem_file: Path, mesh: TriangleMesh): and maps to cell centers via arithmetic averaging of scalar field values. (For coordinate averaging, see mesh._compute_cell_centers() which uses spherical geometry.) + The file stores the values by Elmer node id: the value of node id i is at position i - 1. The mesh may hold only + part of the nodes, e.g. the partition of one MPI rank. + Args: dem_file (Path): Path to the DEM NetCDF file. mesh (TriangleMesh): Mesh object with vertex and cell information @@ -118,16 +121,21 @@ def read_dem_xios(dem_file: Path, mesh: TriangleMesh): import netCDF4 - nc = netCDF4.Dataset(dem_file) - - # Read vertex data from nc file - zs_vertices = np.squeeze(nc["zs"][:]).data - h_vertices = np.squeeze(nc["h"][:]).data + with netCDF4.Dataset(dem_file) as nc: + # Read node data from nc file + zs_nodes = np.squeeze(nc["zs"][:]).data + h_nodes = np.squeeze(nc["h"][:]).data + + n_nodes = len(zs_nodes) + if mesh.vertex_ids.min() < 1 or mesh.vertex_ids.max() > n_nodes: + raise ValueError( + f"The Elmer mesh has node ids {mesh.vertex_ids.min()}..{mesh.vertex_ids.max()}, but {dem_file} has values " + f"for node ids 1..{n_nodes} only." + ) - assert len(zs_vertices) == len(mesh.x_vertices), ( - f"Surface mesh ({len(zs_vertices)} vertices) and Elmer mesh " - f"({len(mesh.x_vertices)} vertices) do not have the same number of vertices" - ) + # Values of the mesh's vertices, in the order of its local vertex indices + zs_vertices = zs_nodes[mesh.vertex_ids - 1] + h_vertices = h_nodes[mesh.vertex_ids - 1] # Map vertex values to cell centers by arithmetic averaging (scalar fields) n_cells = mesh.cell_to_vertex.shape[0] diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py index 29e7237f..9d382b15 100644 --- a/tests/core/test_cli.py +++ b/tests/core/test_cli.py @@ -7,7 +7,7 @@ import unittest import io import os -from contextlib import redirect_stderr, redirect_stdout +from contextlib import redirect_stdout from pathlib import Path from unittest.mock import MagicMock, patch @@ -17,7 +17,6 @@ mpi_launcher_size, parse_cli_args, validate_forcing_parallel_compat, - validate_grid_parallel_compat, validate_partition_count, validate_restart_compat, validate_shading_coupling_compat, @@ -26,7 +25,6 @@ ) from ebfm.core.comm import mpi_available -from ebfm.core.grid import GridInputType from ebfm.core.config import FieldValidationLevel, Calendar, TimeConfig, iso8601 # Minimal valid args for each primary grid type. @@ -414,50 +412,6 @@ def test_shading_on_several_ranks_rejected(self): self.assertEqual(ctx.exception.code, 2) -class TestValidateGridParallelCompat(unittest.TestCase): - """validate_grid_parallel_compat rejects a non-partitioned ELMERXIOS grid on several ranks.""" - - # name: (grid_type, is_partitioned, grid option named by the error on several ranks, or None if accepted) - _CASES = { - "--matlab-mesh": (GridInputType.MATLAB, False, None), - "--netcdf-mesh": (GridInputType.NETCDF, False, None), - "--elmer-mesh": (GridInputType.ELMER, False, None), - "--elmer-mesh --netcdf-dem-mesh": (GridInputType.CUSTOM, False, None), - "--elmer-mesh --netcdf-dem-mesh-unstructured": ( - GridInputType.ELMERXIOS, - False, - "--netcdf-dem-mesh-unstructured", - ), - "--elmer-mesh --netcdf-dem-mesh --is-partitioned-elmer-mesh": (GridInputType.CUSTOM, True, None), - } - - def setUp(self): - # Ensure _parser is set by running a valid parse first. - parse_cli_args(_MATLAB) - - def _validate(self, grid_type, is_partitioned: bool, comm_size: int): - grid_config = MagicMock(grid_type=grid_type, is_partitioned=is_partitioned) - validate_grid_parallel_compat(grid_config, MagicMock(size=comm_size)) - - def test_single_rank_ok(self): - for name, (grid_type, is_partitioned, _) in self._CASES.items(): - with self.subTest(name): - self._validate(grid_type, is_partitioned, comm_size=1) # must not raise - - def test_several_ranks(self): - for name, (grid_type, is_partitioned, rejected_grid) in self._CASES.items(): - with self.subTest(name): - if rejected_grid is None: - self._validate(grid_type, is_partitioned, comm_size=4) # must not raise - else: - stderr = io.StringIO() - with redirect_stderr(stderr), self.assertRaises(SystemExit) as ctx: - self._validate(grid_type, is_partitioned, comm_size=4) - self.assertEqual(ctx.exception.code, 2) - error = stderr.getvalue().splitlines()[-1] - self.assertIn(f"error: {rejected_grid} ", error) - - @unittest.skipUnless(mpi_available, "requires MPI (pip install 'ebfm[mpi]')") class TestValidateUsePartParallelCompat(unittest.TestCase): """validate_use_part_parallel_compat raises a CLI error for --use-part in a run on several ranks.""" diff --git a/tests/core/test_elmer_mesh.py b/tests/core/test_elmer_mesh.py index 357540f0..731af314 100644 --- a/tests/core/test_elmer_mesh.py +++ b/tests/core/test_elmer_mesh.py @@ -2,11 +2,15 @@ # # SPDX-License-Identifier: BSD-3-Clause +import tempfile import unittest +from pathlib import Path import numpy as np +from netCDF4 import Dataset from ebfm.elmer.mesh import TriangleMesh +from ebfm.reader import read_dem_xios N_CELLS = 10 N_PARTITIONS = 3 @@ -66,5 +70,66 @@ def test_partition_id_out_of_range_is_rejected(self): self.mesh.get_partition(partition_id, N_PARTITIONS) +def _write_xios_dem(path: Path, n_nodes: int) -> tuple[np.ndarray, np.ndarray]: + """Write zs and h for node ids 1..n_nodes, stored by node id as Elmer/Ice XIOS output does; return them.""" + node_ids = np.arange(1, n_nodes + 1) + zs = (500.0 + 7.0 * node_ids).astype(np.float32) + h = np.where(node_ids <= n_nodes // 3, 0.5, 100.0 + node_ids).astype(np.float32) + with Dataset(path, "w") as nc: + nc.createDimension("time", None) + nc.createDimension("nmesh2D_node", n_nodes) + for name, values in (("zs", zs), ("h", h)): + nc.createVariable(name, "f4", ("time", "nmesh2D_node"))[0, :] = values + return zs, h + + +def _renumbered(mesh: TriangleMesh) -> TriangleMesh: + """The same mesh with its vertices listed in reverse order, as in a mesh.nodes not sorted by node id.""" + order = np.arange(len(mesh.vertex_ids))[::-1] + position = np.empty_like(order) + position[order] = np.arange(len(order)) + return TriangleMesh( + x_vertices=mesh.x_vertices[order], + y_vertices=mesh.y_vertices[order], + z_vertices=mesh.z_vertices[order], + cell_to_vertex=position[mesh.cell_to_vertex], + vertex_ids=mesh.vertex_ids[order], + cell_ids=mesh.cell_ids, + ) + + +class TestReadDemXiosByNodeId(unittest.TestCase): + """read_dem_xios looks up the node values by Elmer node id, so any part or order of the mesh gets the same cells.""" + + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + self.dem_file = Path(self.tmpdir.name) / "xios_dem.nc" + self.mesh = _strip_mesh() + self.zs, self.h = _write_xios_dem(self.dem_file, n_nodes=self.mesh.vertex_ids.max()) + + def tearDown(self): + self.tmpdir.cleanup() + + def test_cells_get_the_mean_of_their_corner_nodes(self): + z, h = read_dem_xios(self.dem_file, self.mesh) + + corner_ids = self.mesh.vertex_ids[self.mesh.cell_to_vertex] + np.testing.assert_allclose(z, self.zs[corner_ids - 1].mean(axis=1), rtol=1e-6) + np.testing.assert_allclose(h, self.h[corner_ids - 1].mean(axis=1), rtol=1e-6) + + def test_renumbered_mesh_gets_the_same_cells(self): + expected = read_dem_xios(self.dem_file, self.mesh) + + for actual, wanted in zip(read_dem_xios(self.dem_file, _renumbered(self.mesh)), expected): + np.testing.assert_array_equal(actual, wanted) + + def test_partitions_get_the_cells_of_the_whole_mesh(self): + expected = read_dem_xios(self.dem_file, self.mesh) + + parts = [read_dem_xios(self.dem_file, self.mesh.get_partition(p, N_PARTITIONS)) for p in (1, 2, 3)] + for field, wanted in enumerate(expected): + np.testing.assert_array_equal(np.concatenate([part[field] for part in parts]), wanted) + + if __name__ == "__main__": unittest.main() diff --git a/tests/core/test_parallel_runs.py b/tests/core/test_parallel_runs.py index f6f144e4..0ea9e87d 100644 --- a/tests/core/test_parallel_runs.py +++ b/tests/core/test_parallel_runs.py @@ -22,6 +22,7 @@ REPO_ROOT = Path(__file__).resolve().parents[2] MATLAB_MESH = REPO_ROOT / "examples" / "dem_and_mask.mat" ELMER_MESH = REPO_ROOT / "tests" / "data" / "elmer_mesh" / "MESH" +XIOS_DEM = ELMER_MESH.parent / "xios_dem.nc" TIMEOUT = 300 requires_mpirun = pytest.mark.skipif( @@ -37,6 +38,7 @@ CASES = { "matlab": ["--matlab-mesh", str(MATLAB_MESH), "--no-shading", *_TIME], "elmer": [*_ELMER, *_TIME], + "elmerxios": [*_ELMER, "--netcdf-dem-mesh-unstructured", str(XIOS_DEM), *_TIME], "fake-coupling": [ *_ELMER, "--couple-to-elmer-ice", @@ -92,6 +94,7 @@ def _assert_same_results(serial: Path, parallel: Path) -> None: ("matlab", 3), ("matlab", 4), ("elmer", 3), + ("elmerxios", 3), ("fake-coupling", 2), ], ids=str, diff --git a/tests/test_reader.py b/tests/test_reader.py index 16336c49..8f584d1d 100644 --- a/tests/test_reader.py +++ b/tests/test_reader.py @@ -53,10 +53,10 @@ def test_cell_values_are_means_of_their_vertices(self): np.testing.assert_allclose(z_cells, [200.0, 300.0]) np.testing.assert_allclose(h_cells, [3.0, 6.0]) - def test_vertex_count_mismatch_is_rejected(self): + def test_missing_node_ids_are_rejected(self): dem_file = self._write_dem(zs=[100.0, 200.0, 300.0], h=[0.0, 3.0, 6.0]) - with self.assertRaises(AssertionError): + with self.assertRaisesRegex(ValueError, "node ids"): read_dem_xios(dem_file, self.mesh) From cc24934565f94d68ccca40ea61db0acd1e24eb5a Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Wed, 21 Oct 2026 09:17:41 +0200 Subject: [PATCH 12/13] Warn when every MPI rank reads a non-partitioned Elmer mesh A non-partitioned Elmer mesh is split over the ranks only after every rank has read all of it. Rank 0 now warns about this for the ELMER, CUSTOM and ELMERXIOS grid types and points to a mesh with one partition per rank, which requires --netcdf-dem-mesh. Assisted-by: Claude Code:claude-opus-5-5 --- CHANGELOG.md | 2 +- README.md | 3 +++ src/ebfm/core/cli.py | 20 ++++++++++++++++++++ src/ebfm/main.py | 2 ++ tests/core/test_cli.py | 33 +++++++++++++++++++++++++++++++++ 5 files changed, 59 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 06cd7ac9..0f1e793b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,7 +30,7 @@ SPDX-License-Identifier: CC-BY-4.0 * 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. 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 diff --git a/README.md b/README.md index 3fc378c5..06fdf912 100644 --- a/README.md +++ b/README.md @@ -302,6 +302,9 @@ mpirun -np 4 ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading * 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 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, diff --git a/src/ebfm/core/cli.py b/src/ebfm/core/cli.py index 4dfe5eca..4988bfb5 100644 --- a/src/ebfm/core/cli.py +++ b/src/ebfm/core/cli.py @@ -642,3 +642,23 @@ def validate_partition_count(args: Namespace, grid_config, comm, logger) -> None _cli_error(f"{message} Pass --allow-partition-mismatch to run anyway.") if comm.rank == 0: logger.warning(message) + + +def warn_non_partitioned_mesh_parallel(grid_config, comm, logger) -> None: + """Warn that every MPI rank reads the whole mesh when a non-partitioned Elmer mesh is split over several ranks. + + Must be called on every rank of comm. Warns on rank 0 only. + + @param[in] grid_config resolved GridConfig + @param[in] comm EBFM's MPI communicator + @param[in] logger logger for the warning + """ + is_elmer_mesh = grid_config.grid_type in (GridInputType.ELMER, GridInputType.CUSTOM, GridInputType.ELMERXIOS) + if not is_elmer_mesh or grid_config.is_partitioned or comm.size == 1 or comm.rank != 0: + return + + logger.warning( + f"EBFM runs on {comm.size} MPI ranks with a non-partitioned Elmer mesh, so every rank reads the whole mesh. " + "For large meshes, use a mesh with one partition per rank (--is-partitioned-elmer-mesh, which requires " + "--netcdf-dem-mesh)." + ) diff --git a/src/ebfm/main.py b/src/ebfm/main.py index f1fdfbb5..43a12808 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -31,6 +31,7 @@ validate_partition_count, validate_restart_compat, validate_shading_parallel_compat, + warn_non_partitioned_mesh_parallel, ) import ebfm.coupling @@ -224,6 +225,7 @@ def _main_impl(): # Each rank computes exactly one partition of a partitioned mesh validate_partition_count(args, grid_config, ebfm_comm, logger) + warn_non_partitioned_mesh_parallel(grid_config, ebfm_comm, logger) logger.debug("Successfully completed consistency checks.") diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py index 9d382b15..3e82e6c3 100644 --- a/tests/core/test_cli.py +++ b/tests/core/test_cli.py @@ -22,9 +22,11 @@ validate_shading_coupling_compat, validate_shading_parallel_compat, validate_use_part_parallel_compat, + warn_non_partitioned_mesh_parallel, ) from ebfm.core.comm import mpi_available +from ebfm.core.grid import GridInputType from ebfm.core.config import FieldValidationLevel, Calendar, TimeConfig, iso8601 # Minimal valid args for each primary grid type. @@ -538,6 +540,37 @@ def test_mismatch_allowed_warns_on_rank_0_only(self): self.logger.warning.assert_not_called() +class TestWarnNonPartitionedMeshParallel(unittest.TestCase): + """Splitting a non-partitioned Elmer mesh over several MPI ranks warns once, on rank 0.""" + + def setUp(self): + self.logger = MagicMock() + + def _warn(self, grid_type=GridInputType.CUSTOM, is_partitioned=False, comm_size=4, rank=0): + grid_config = MagicMock(grid_type=grid_type, is_partitioned=is_partitioned) + warn_non_partitioned_mesh_parallel(grid_config, MagicMock(size=comm_size, rank=rank), self.logger) + + def test_non_partitioned_elmer_mesh_on_several_ranks_warns_on_rank_0(self): + for grid_type in (GridInputType.ELMER, GridInputType.CUSTOM, GridInputType.ELMERXIOS): + with self.subTest(grid_type=grid_type): + self.logger.reset_mock() + self._warn(grid_type=grid_type) + self.logger.warning.assert_called_once() + + def test_no_warning(self): + cases = { + "rank 1": dict(rank=1), + "single rank": dict(comm_size=1), + "partitioned mesh": dict(is_partitioned=True), + "MATLAB grid": dict(grid_type=GridInputType.MATLAB), + } + for name, kwargs in cases.items(): + with self.subTest(name): + self.logger.reset_mock() + self._warn(**kwargs) + self.logger.warning.assert_not_called() + + class TestCliCalendarArgument(unittest.TestCase): def test_calendar_default(self): args = parse_cli_args(_MATLAB) From f70814ae3470828c591f211b4e746752568e108c Mon Sep 17 00:00:00 2001 From: Benjamin Rodenberg Date: Thu, 22 Oct 2026 10:51:49 +0200 Subject: [PATCH 13/13] Read and write restart files on several MPI ranks Restart files were rejected on several ranks and with partitioned meshes. Rank 0 now gathers the columns of all ranks one variable at a time and writes them ordered by column id, together with the ids in a new variable column_id. Reading, rank 0 matches the rows of the file to the columns by id and scatters them with the new ColumnGather.scatter, so a file continues on any number of ranks and on a partitioned mesh. A file without column_id holds the columns in the order of a serial run, which for a split grid is the ranks' columns in rank order; on several ranks its x and y are checked against the grid. Assisted-by: Claude Code:claude-opus-5-5 --- CHANGELOG.md | 2 +- README.md | 14 ++- src/ebfm/core/FINAL_create_restart_file.py | 97 ++++++++-------- src/ebfm/core/INIT.py | 47 +++----- src/ebfm/core/cli.py | 24 ---- src/ebfm/core/comm/gather.py | 59 +++++++++- src/ebfm/core/restart.py | 127 +++++++++++++++++++++ src/ebfm/main.py | 15 ++- tests/core/comm/scatter_columns.py | 35 ++++++ tests/core/test_cli.py | 45 -------- tests/core/test_init.py | 78 ++++++++++++- tests/core/test_parallel_columns.py | 19 ++- tests/core/test_parallel_restart.py | 54 +++++++++ tests/core/test_partitioned_mesh_run.py | 28 +++++ 14 files changed, 473 insertions(+), 171 deletions(-) create mode 100644 tests/core/comm/scatter_columns.py create mode 100644 tests/core/test_parallel_restart.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0f1e793b..81ae67a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,7 +34,7 @@ SPDX-License-Identifier: CC-BY-4.0 * 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`) with a partitioned Elmer mesh now stop the run with an error, also on a single rank. Such a run previously skipped writing the restart file with a warning at its end. 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 diff --git a/README.md b/README.md index 06fdf912..1d7fda31 100644 --- a/README.md +++ b/README.md @@ -291,6 +291,11 @@ 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 the grid over the MPI ranks of a run (requires `pip install 'ebfm[mpi]'`): @@ -305,10 +310,11 @@ mpirun -np 4 ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading * 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 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` and `--restart-dir`/`--restart-init`. +* 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. diff --git a/src/ebfm/core/FINAL_create_restart_file.py b/src/ebfm/core/FINAL_create_restart_file.py index 16444b1c..ee46f01a 100644 --- a/src/ebfm/core/FINAL_create_restart_file.py +++ b/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 diff --git a/src/ebfm/core/INIT.py b/src/ebfm/core/INIT.py index ebb1e203..c99a3c04 100644 --- a/src/ebfm/core/INIT.py +++ b/src/ebfm/core/INIT.py @@ -16,9 +16,11 @@ from ebfm.reader import read_elmer_mesh, read_dem, read_dem_xios from ebfm.elmer.mesh import Mesh +from .comm import defaultComm +from .comm.gather import ColumnGather from .config import TimeConfig, GridConfig, ColumnDiscretizationConfig, iso8601 from .grid import GridInputType, GridDict, ShadingMethod, domain_mean, number_of_columns, validate_grid -from .restart import validate_all_variables_present, validate_variable_shape +from .restart import read_restart_file from .constants import DAYS_PER_YEAR, LATENT_HEAT_OF_FUSION, MELTING_POINT, SECONDS_PER_DAY, WATER_DENSITY @@ -682,7 +684,13 @@ def read_NETCDF_grid(gridfile: Path, fill_missing_elevation: bool = False): def init_initial_conditions( - C, grid: GridDict, io, time, column: ColumnDiscretizationConfig, init_with_restart_file: bool + C, + grid: GridDict, + io, + time, + column: ColumnDiscretizationConfig, + init_with_restart_file: bool, + gather: ColumnGather | None = None, ): """ Sets the model's initial conditions at the start of the simulation. @@ -694,6 +702,8 @@ def init_initial_conditions( time (dict): Dictionary with time-related parameters (e.g. ts). column (ColumnDiscretizationConfig): How each column is divided into layers. init_with_restart_file (bool): Flag indicating whether to initialize from a restart file or set manually. + gather (ColumnGather): The columns of all MPI ranks, which read the restart file together; None for the + columns of a single rank. Collective over its communicator when reading a restart file. Returns: OUT (dict): Dictionary containing model outputs initialized with default or restart file values. @@ -717,36 +727,9 @@ def init_initial_conditions( if init_with_restart_file: logger.info(f"Initialize from restart file: {io['bootfilein']}...") - # Open the NetCDF file - with Dataset(io["bootfilein"], "r") as ncfile: - validate_all_variables_present(ncfile.variables, io["bootfilein"]) - # Iterate through all variables in the file - for var_name in ncfile.variables: - # Read the variable data. netCDF4 returns a numpy.ma.MaskedArray - # here whenever the variable defines a fill value, regardless of - # whether any missing values are actually present. - var_data = ncfile.variables[var_name][:] - # Restart files are expected to be fully populated. Mixing masked - # arrays into the simulation state is unsafe: numpy's masked-array - # ufuncs silently ignore `where=` when combining masks (only the - # data respects `where`), so masks can leak into cells they - # should not and snowball across timesteps. Verify the assumption - # holds, then drop down to a plain ndarray. - assert not np.ma.is_masked(var_data), ( - f"Restart variable '{var_name}' in {io['bootfilein']} contains missing values, " - "which is not supported." - ) - if isinstance(var_data, np.ma.MaskedArray): - var_data = var_data.data - - # Perform consistency checks - validate_variable_shape(var_name, var_data.shape, gpsum, nl, io["bootfilein"]) - - # If a variable has no dimensions (scalar), convert it to a Python scalar - if var_data.shape == (): # Scalar variable - var_data = var_data.item() - # Save the variable data to the dictionary - OUT[var_name] = var_data + if gather is None: + gather = ColumnGather(defaultComm, grid["column_id"]) + OUT = read_restart_file(io["bootfilein"], nl, gather, grid) OUT["timelastsnow"] = num2date( OUT["timelastsnow_netCDF"], diff --git a/src/ebfm/core/cli.py b/src/ebfm/core/cli.py index 4988bfb5..348483e9 100644 --- a/src/ebfm/core/cli.py +++ b/src/ebfm/core/cli.py @@ -596,30 +596,6 @@ def validate_forcing_parallel_compat(args: Namespace, comm) -> None: ) -def validate_restart_compat(args: Namespace, grid_config, comm) -> None: - """Check that restart files are only read and written for a whole grid on a single MPI rank. - - Must be called on every rank of comm. - - @param[in] args parsed CLI arguments - @param[in] grid_config resolved GridConfig - @param[in] comm EBFM's MPI communicator - """ - if args.restart_dir is None and not args.restart_init: - return - - if comm.size > 1: - _cli_error( - f"Restart files (--restart-dir, --restart-init) are not supported for runs on several MPI ranks " - f"({comm.size} ranks). Please run on a single rank." - ) - if grid_config.is_partitioned: - _cli_error( - "Restart files (--restart-dir, --restart-init) are not supported for partitioned meshes " - "(--is-partitioned-elmer-mesh)." - ) - - def validate_partition_count(args: Namespace, grid_config, comm, logger) -> None: """Check that a partitioned Elmer mesh has exactly one partition per MPI rank. diff --git a/src/ebfm/core/comm/gather.py b/src/ebfm/core/comm/gather.py index f0cdbf6b..c41e629e 100644 --- a/src/ebfm/core/comm/gather.py +++ b/src/ebfm/core/comm/gather.py @@ -2,7 +2,8 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Collect the columns that the ranks of EBFM's communicator hold on one rank, e.g. to write them to a file.""" +"""Collect the columns that the ranks of EBFM's communicator hold on one rank, e.g. to write them to a file, and +distribute them again, e.g. after reading them from a file.""" import numpy as np from numpy.typing import NDArray @@ -12,13 +13,13 @@ class ColumnGather: """ - Gathers per-column arrays of all ranks on a root rank, ordered by column id. + Gathers per-column arrays of all ranks on a root rank, ordered by column id, and scatters them back. Every rank holds a subset of the columns, identified by their ids (grid["column_id"], see ebfm.core.grid). The root rank receives the columns of all ranks sorted by id. - Construction and every gather are collective: all ranks of the communicator must take part, in the same - order. With a single rank, nothing is communicated and the gathers return their input unchanged. + Construction and every gather and scatter are collective: all ranks of the communicator must take part, in the + same order. With a single rank, nothing is communicated and they return their input unchanged. """ def __init__(self, comm: Comm, column_ids: NDArray[np.integer], root: int = 0): @@ -38,6 +39,7 @@ def __init__(self, comm: Comm, column_ids: NDArray[np.integer], root: int = 0): self._order = None return + self._counts = comm.gather(self.n_local_columns, root=root) all_column_ids = self.gather_unordered(column_ids) if self.is_root: self._order = np.argsort(all_column_ids, kind="stable") @@ -96,3 +98,52 @@ def gather(self, local: NDArray) -> NDArray | None: if not self.is_root: return None return gathered[self._order] + + def scatter_unordered(self, values: NDArray | None) -> NDArray: + """ + Split an array of the root rank along its first axis into consecutive blocks, one per rank in rank order. + + The inverse of gather_unordered: rank r receives as many rows as it has columns, the rows that follow those + of the ranks before it. + + @param[in] values on the root rank, an array with one row per column of all ranks; ignored on other ranks + + @returns this rank's rows + """ + if self.comm.size == 1: + return values + + layout = None + if self.is_root: + values = np.ascontiguousarray(values) + assert values.shape[0] == sum( + self._counts + ), f"Expected one row per column of all ranks ({sum(self._counts)}), got an array of shape {values.shape}." + layout = (values.dtype, values.shape[1:]) + dtype, trailing_shape = self.comm.bcast(layout, root=self.root) + + local = np.empty((self.n_local_columns, *trailing_shape), dtype=dtype) + values_per_row = int(np.prod(trailing_shape)) + send = (values, [n * values_per_row for n in self._counts]) if self.is_root else None + self.comm.Scatterv(send, local, root=self.root) + return local + + def scatter(self, values: NDArray | None) -> NDArray: + """ + Give every rank the rows of its columns from an array of the root rank ordered by column id. + + The inverse of gather. + + @param[in] values on the root rank, an array with one row per column of all ranks, in the order of + column_ids; ignored on other ranks + + @returns the rows of this rank's columns, in the order of its column ids + """ + if self.comm.size == 1: + return values + + unordered = None + if self.is_root: + unordered = np.empty_like(values) + unordered[self._order] = values + return self.scatter_unordered(unordered) diff --git a/src/ebfm/core/restart.py b/src/ebfm/core/restart.py index ff323da5..ddd37e6e 100644 --- a/src/ebfm/core/restart.py +++ b/src/ebfm/core/restart.py @@ -6,10 +6,25 @@ Single source of truth for the state carried across a restart, shared by the writer (`FINAL_create_restart_file`) and the loader (`INIT.init_initial_conditions`) so the two cannot drift apart. + +A restart file holds the columns of the whole domain, in the order of their ids, and their ids in the variable +`column_id`. A run on several MPI ranks reads and writes it on rank 0, which matches the rows to the columns by id, +so the file does not depend on how the columns are distributed over the ranks. """ +import logging from pathlib import Path +import numpy as np +from netCDF4 import Dataset + +from .comm.gather import ColumnGather + +logger = logging.getLogger(__name__) + +# Ids of the columns the rows of a restart file belong to (grid["column_id"]); older files lack it +COLUMN_ID_VARIABLE = "column_id" + # One value per snow/firn column: shape `(gpsum,)`. PER_COLUMN_VARIABLES = ( "snowmass", @@ -107,3 +122,115 @@ def validate_all_variables_present(var_names, source: Path) -> None: missing = [name for name in RESTART_VARIABLES if name not in set(var_names)] if missing: raise ValueError(f"Restart file {source} is missing the required variable(s): {', '.join(missing)}.") + + +def read_restart_file(path: Path, nl: int, gather: ColumnGather, grid) -> dict: + """ + Read a restart file on the root rank of gather and give every rank the rows of its columns. + + Collective over gather.comm. The rows of a file with a `column_id` variable are matched to the columns by id. A + file without one holds the columns in the order of a run on a single rank, which, for a grid that EBFM splits + over the ranks, is the order of the ranks' columns concatenated in rank order. On several ranks, the x and y of + such a file are therefore checked against the grid. + + Parameters: + path (Path): Restart file to read + nl (int): Number of layers per column + gather (ColumnGather): The columns of all ranks + grid (dict): This rank's part of the grid + + Returns: + dict: every variable of the file except `column_id`, with the rows of this rank's columns + + Raises: + ValueError: on the root rank, if the file does not match the columns of all ranks + """ + is_parallel_run = gather.comm.size > 1 + + variables, is_ordered_by_id = None, None + if gather.is_root: + n_columns = len(gather.column_ids) + variables = _read_variables(path, n_columns, nl) + file_column_ids = variables.pop(COLUMN_ID_VARIABLE, None) + is_ordered_by_id = file_column_ids is not None + if is_ordered_by_id: + rows = _rows_by_column_id(file_column_ids, gather.column_ids, path) + variables = {name: values[rows] if np.ndim(values) else values for name, values in variables.items()} + elif is_parallel_run: + logger.warning( + f"Restart file {path} has no '{COLUMN_ID_VARIABLE}' variable; assuming its rows are in the order of a " + "run on a single rank." + ) + + layout = None + if gather.is_root: + layout = ([(name, np.ndim(values)) for name, values in variables.items()], is_ordered_by_id) + if is_parallel_run: + layout = gather.comm.bcast(layout, root=gather.root) + names_and_ndims, is_ordered_by_id = layout + + local = {} + for name, ndim in names_and_ndims: + values = variables[name] if gather.is_root else None + if ndim == 0: + local[name] = gather.comm.bcast(values, root=gather.root) if is_parallel_run else values + elif is_ordered_by_id: + local[name] = gather.scatter(values) + else: + local[name] = gather.scatter_unordered(values) + + if is_parallel_run and not is_ordered_by_id: + _check_coordinates(local, grid, path) + return local + + +def _read_variables(path: Path, n_columns: int, nl: int) -> dict: + """Read and check every variable of a restart file for n_columns columns of nl layers.""" + variables = {} + with Dataset(path, "r") as ncfile: + validate_all_variables_present(ncfile.variables, path) + for var_name in ncfile.variables: + # netCDF4 returns a numpy.ma.MaskedArray whenever the variable defines a fill value, regardless of whether + # any missing values are actually present. + var_data = ncfile.variables[var_name][:] + # Restart files are expected to be fully populated. Mixing masked arrays into the simulation state is + # unsafe: numpy's masked-array ufuncs silently ignore `where=` when combining masks (only the data + # respects `where`), so masks can leak into cells they should not and snowball across timesteps. Verify + # the assumption holds, then drop down to a plain ndarray. + assert not np.ma.is_masked( + var_data + ), f"Restart variable '{var_name}' in {path} contains missing values, which is not supported." + if isinstance(var_data, np.ma.MaskedArray): + var_data = var_data.data + + validate_variable_shape(var_name, var_data.shape, n_columns, nl, path) + + # A variable without dimensions becomes a Python scalar + variables[var_name] = var_data.item() if var_data.shape == () else var_data + return variables + + +def _rows_by_column_id(file_column_ids, column_ids, path: Path): + """The row of the file for each of column_ids. Both hold the same number of ids.""" + sorter = np.argsort(file_column_ids, kind="stable") + positions = np.minimum(np.searchsorted(file_column_ids, column_ids, sorter=sorter), len(sorter) - 1) + rows = sorter[positions] + + is_missing = file_column_ids[rows] != column_ids + if is_missing.any(): + raise ValueError( + f"Restart file {path} holds no row for {is_missing.sum()} of the {len(column_ids)} columns, e.g. for " + f"column id {column_ids[is_missing][0]}. It was written for a different grid." + ) + return rows + + +def _check_coordinates(variables: dict, grid, path: Path) -> None: + """Check that the x and y of the restart file are those of this rank's columns, to 1 mm.""" + for coordinate in ("x", "y"): + if not np.allclose(variables[coordinate], grid[coordinate], rtol=0, atol=1e-3): + raise ValueError( + f"The '{coordinate}' of restart file {path} differs from that of the grid. The file has no " + f"'{COLUMN_ID_VARIABLE}' variable, so its rows are assumed to be in the order of a run on a single " + "rank; it was written for a different grid or a different order of the columns." + ) diff --git a/src/ebfm/main.py b/src/ebfm/main.py index 43a12808..a5ee6da4 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -29,7 +29,6 @@ validate_use_part_parallel_compat, validate_forcing_parallel_compat, validate_partition_count, - validate_restart_compat, validate_shading_parallel_compat, warn_non_partitioned_mesh_parallel, ) @@ -217,9 +216,6 @@ def _main_impl(): # Ensure forcing files are only read in runs on a single rank. validate_forcing_parallel_compat(args, ebfm_comm) - # Ensure restart files are only read and written for a whole grid on a single rank. - validate_restart_compat(args, grid_config, ebfm_comm) - # Ensure shading routine is only used in runs on a single rank. validate_shading_parallel_compat(grid_config, ebfm_comm) @@ -242,7 +238,8 @@ def _main_impl(): if args.check_grid_distribution and is_parallel_run: validate_grid_distribution(grid, ebfm_comm) - # Rank 0 collects the columns of all ranks, e.g. to write them to the output file or for --dump-reference + # Rank 0 collects the columns of all ranks, e.g. to write them to the output file, the restart file or for + # --dump-reference column_gather = ColumnGather(ebfm_comm, grid["column_id"]) if is_parallel_run: logger.info(f"Rank 0 writes the output of all {ebfm_comm.size} ranks.") @@ -253,7 +250,9 @@ def _main_impl(): f"{ebfm_comm.size} ranks every time step." ) - OUT, IN, OUTFILE = INIT.init_initial_conditions(C, grid, io, time, column, init_with_restart_file=args.restart_init) + OUT, IN, OUTFILE = INIT.init_initial_conditions( + C, grid, io, time, column, init_with_restart_file=args.restart_init, gather=column_gather + ) coupler.setup(grid, time_config) @@ -398,7 +397,7 @@ def _main_impl(): # IN['dhdx'] = data_from_elmer('dhdx') # IN['dhdy'] = data_from_elmer('dhdy') - # Needed by FINAL_create_restart_file.main(OUT, io). The grid carries the elevation of the current + # Needed by FINAL_create_restart_file.main. The grid carries the elevation of the current # time step, whether it was read during initialization or received from Elmer/Ice just above. OUT["x"] = grid["x"] OUT["y"] = grid["y"] @@ -408,7 +407,7 @@ def _main_impl(): io, OUTFILE = LOOP_write_to_file.main(OUTFILE, io, OUT, grid, t, time, column, column_gather) # Write restart file - FINAL_create_restart_file.main(OUT, io, args.restart_dir) + FINAL_create_restart_file.main(OUT, io, args.restart_dir, column_gather) logger.info("Time loop completed.") diff --git a/tests/core/comm/scatter_columns.py b/tests/core/comm/scatter_columns.py new file mode 100644 index 00000000..8e79f0f3 --- /dev/null +++ b/tests/core/comm/scatter_columns.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Check ColumnGather's scatters on 3 MPI ranks that hold the 10 columns of a domain unevenly and out of id order. + +Usage: mpirun -np 3 python scatter_columns.py + +Every column is held by one rank, and rank 2 holds none. Checks the rows every rank receives from rank 0. +""" + +import numpy as np + +from ebfm.core.comm import defaultComm as comm +from ebfm.core.comm.gather import ColumnGather + +assert comm.size == 3, f"Run on 3 ranks, not {comm.size}." + +column_ids = 100 + 7 * np.arange(10) +columns = {0: [5, 1, 7, 3], 1: [0, 9, 2, 8, 4, 6], 2: []}[comm.rank] +local_ids = column_ids[columns] + +gather = ColumnGather(comm, local_ids) +is_root = gather.is_root +values = gather.scatter(1.5 * column_ids if is_root else None) +layers = gather.scatter(np.stack([column_ids, 2 * column_ids], axis=1) if is_root else None) +blocks = gather.scatter_unordered(np.arange(10) if is_root else None) +round_trip = gather.scatter(gather.gather(-local_ids)) + +np.testing.assert_array_equal(values, 1.5 * local_ids) +np.testing.assert_array_equal(layers, np.stack([local_ids, 2 * local_ids], axis=1).reshape(-1, 2)) +np.testing.assert_array_equal(blocks, {0: [0, 1, 2, 3], 1: [4, 5, 6, 7, 8, 9], 2: []}[comm.rank]) +np.testing.assert_array_equal(round_trip, -local_ids) + +print(f"rank {comm.rank}: ok") diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py index 3e82e6c3..eb55fb31 100644 --- a/tests/core/test_cli.py +++ b/tests/core/test_cli.py @@ -18,7 +18,6 @@ parse_cli_args, validate_forcing_parallel_compat, validate_partition_count, - validate_restart_compat, validate_shading_coupling_compat, validate_shading_parallel_compat, validate_use_part_parallel_compat, @@ -461,50 +460,6 @@ def test_forcing_on_several_ranks_rejected(self): self.assertEqual(ctx.exception.code, 2) -class TestValidateRestartCompat(unittest.TestCase): - """validate_restart_compat rejects restart files in a run on several ranks or on a partitioned mesh.""" - - _RESTART_ARGS = { - "--restart-dir": ["--restart-dir", "restart/"], - "--restart-init": ["--restart-init"], - "--restart-dir --restart-init": ["--restart-dir", "restart/", "--restart-init"], - } - - def setUp(self): - # Ensure _parser is set by running a valid parse first. - parse_cli_args(_MATLAB) - - def _validate(self, restart_args, comm_size=1, is_partitioned=False): - grid_config = MagicMock(is_partitioned=is_partitioned) - validate_restart_compat(_parse(*restart_args), grid_config, MagicMock(size=comm_size)) - - def _assert_rejected(self, **kwargs): - for name, restart_args in self._RESTART_ARGS.items(): - with self.subTest(name), self.assertRaises(SystemExit) as ctx: - self._validate(restart_args, **kwargs) - self.assertEqual(ctx.exception.code, 2) - - def test_restart_on_single_rank_ok(self): - for name, restart_args in self._RESTART_ARGS.items(): - with self.subTest(name): - self._validate(restart_args) # must not raise - - def test_no_restart_on_several_ranks_ok(self): - self._validate([], comm_size=4) # must not raise - - def test_no_restart_on_several_ranks_with_partitioned_mesh_ok(self): - self._validate([], comm_size=4, is_partitioned=True) # must not raise - - def test_restart_on_several_ranks_rejected(self): - self._assert_rejected(comm_size=4) - - def test_restart_with_partitioned_mesh_rejected(self): - self._assert_rejected(is_partitioned=True) - - def test_restart_on_several_ranks_with_partitioned_mesh_rejected(self): - self._assert_rejected(comm_size=4, is_partitioned=True) - - class TestValidatePartitionCount(unittest.TestCase): """A partitioned mesh needs one MPI rank per partition, unless --allow-partition-mismatch is passed.""" diff --git a/tests/core/test_init.py b/tests/core/test_init.py index ddd45ea6..ea25555f 100644 --- a/tests/core/test_init.py +++ b/tests/core/test_init.py @@ -19,10 +19,11 @@ from ebfm.core.FINAL_create_restart_file import main as write_restart_file from ebfm.core.cli import parse_cli_args from ebfm.core.comm import defaultComm +from ebfm.core.comm.gather import ColumnGather from ebfm.core.config import ColumnDiscretizationConfig from ebfm.core.config.grid import GridConfig from ebfm.core.config.time import TimeConfig -from ebfm.core.restart import PER_COLUMN_VARIABLES, PER_LAYER_VARIABLES, RESTART_VARIABLES +from ebfm.core.restart import COLUMN_ID_VARIABLE, PER_COLUMN_VARIABLES, PER_LAYER_VARIABLES, RESTART_VARIABLES GPSUM = 5 NL = 4 @@ -41,7 +42,7 @@ def _make_grid(number_of_columns: int = GPSUM) -> dict: `init_initial_conditions` derives the column count from the grid via `ebfm.core.grid.number_of_columns`, so the mask is what sets it. """ - return {"mask": np.ones(number_of_columns, dtype=int)} + return {"mask": np.ones(number_of_columns, dtype=int), "column_id": np.arange(number_of_columns)} def _build_out_dict() -> dict: @@ -63,9 +64,14 @@ def _build_out_dict() -> dict: } +def _single_rank_gather() -> ColumnGather: + """The columns of _make_grid() on a single rank.""" + return ColumnGather(defaultComm, np.arange(GPSUM)) + + def _write_restart_file(path: Path, with_missing_value: bool = False) -> None: io = {"writebootfile": True, "bootfileout": path} - write_restart_file(_build_out_dict(), io, restartdir=path.parent) + write_restart_file(_build_out_dict(), io, restartdir=path.parent, gather=_single_rank_gather()) if with_missing_value: # Restart files produced by EBFM are always fully populated. To @@ -214,7 +220,12 @@ def _write_with(self, name: str, **overrides) -> Path: bootfile = Path(self.tmpdir.name) / f"restart_test_{name}.nc" out = _build_out_dict() out.update(overrides) - write_restart_file(out, {"writebootfile": True, "bootfileout": bootfile}, restartdir=bootfile.parent) + write_restart_file( + out, + {"writebootfile": True, "bootfileout": bootfile}, + restartdir=bootfile.parent, + gather=_single_rank_gather(), + ) return bootfile def test_per_layer_variable_stored_per_column_is_rejected(self): @@ -297,6 +308,61 @@ def test_restart_file_with_a_scalar_variable_is_accepted(self): self.assertEqual(OUT["a_scalar"], 1.0) +def copy_without_column_ids(path: Path, copy: Path) -> None: + """Copy a restart file without its column ids, as written before they existed.""" + with Dataset(path) as source, Dataset(copy, "w") as target: + for name, dimension in source.dimensions.items(): + if not name.startswith(COLUMN_ID_VARIABLE): + target.createDimension(name, len(dimension)) + for name, variable in source.variables.items(): + if name != COLUMN_ID_VARIABLE: + target.createVariable(name, variable.dtype, variable.dimensions)[:] = variable[:] + + +class TestRestartColumnIds(unittest.TestCase): + """The rows of a restart file are matched to the columns of the grid by their ids.""" + + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + self.bootfile = Path(self.tmpdir.name) / "restart_test.nc" + out = _build_out_dict() + out["Tsurf"] = 250.0 + np.arange(GPSUM) + write_restart_file( + out, + {"writebootfile": True, "bootfileout": self.bootfile}, + restartdir=self.bootfile.parent, + gather=_single_rank_gather(), + ) + + def tearDown(self): + self.tmpdir.cleanup() + + def _load(self, bootfile, column_ids): + grid = {**_make_grid(), "column_id": column_ids} + OUT, _, _ = init_initial_conditions( + {"alb_ice": 0.5}, grid, {"bootfilein": bootfile}, {}, _make_column(), init_with_restart_file=True + ) + return OUT + + def test_rows_follow_the_column_ids_of_the_grid(self): + OUT = self._load(self.bootfile, np.arange(GPSUM)[::-1]) + + np.testing.assert_array_equal(OUT["Tsurf"], 250.0 + np.arange(GPSUM)[::-1]) + self.assertNotIn(COLUMN_ID_VARIABLE, OUT) + + def test_unknown_column_id_is_rejected(self): + with self.assertRaisesRegex(ValueError, r"holds no row for 1 of the 5 columns, e.g. for column id 99"): + self._load(self.bootfile, np.array([0, 1, 2, 3, 99])) + + def test_file_without_column_ids_is_read_in_file_order(self): + bootfile = self.bootfile.with_name("without_column_ids.nc") + copy_without_column_ids(self.bootfile, bootfile) + + OUT = self._load(bootfile, np.arange(GPSUM)[::-1]) + + np.testing.assert_array_equal(OUT["Tsurf"], 250.0 + np.arange(GPSUM)) + + class TestRestartVariableManifest(unittest.TestCase): """The writer and the loader must agree on what a restart file contains. @@ -310,12 +376,12 @@ def setUp(self): def tearDown(self): self.tmpdir.cleanup() - def test_writer_writes_exactly_the_manifest(self): + def test_writer_writes_exactly_the_manifest_and_the_column_ids(self): bootfile = Path(self.tmpdir.name) / "restart_test.nc" _write_restart_file(bootfile) with Dataset(bootfile) as ncfile: - self.assertEqual(set(ncfile.variables), set(RESTART_VARIABLES)) + self.assertEqual(set(ncfile.variables), {*RESTART_VARIABLES, COLUMN_ID_VARIABLE}) def test_writer_writes_each_group_at_its_own_shape(self): bootfile = Path(self.tmpdir.name) / "restart_test.nc" diff --git a/tests/core/test_parallel_columns.py b/tests/core/test_parallel_columns.py index 653c3918..16d22319 100644 --- a/tests/core/test_parallel_columns.py +++ b/tests/core/test_parallel_columns.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: BSD-3-Clause """Columns distributed over several MPI ranks are held by exactly one rank each, gathered on one rank in column id -order and reduced over all ranks.""" +order, scattered back and reduced over all ranks.""" import filecmp import importlib.util @@ -58,6 +58,23 @@ def test_gather_on_three_ranks(): assert f"rank {rank}: ok" in result.stdout, result.stdout + result.stderr +def test_single_rank_scatters_without_communication(): + gather = ColumnGather(defaultComm, np.array([4, 2, 9])) + values = np.array([1.0, 2.0, 3.0]) + + assert gather.scatter(values) is values + assert gather.scatter_unordered(values) is values + + +@requires_mpirun +def test_scatter_on_three_ranks(): + result = _mpirun(3, "scatter_columns.py") + + assert result.returncode == 0, result.stdout + result.stderr + for rank in range(3): + assert f"rank {rank}: ok" in result.stdout, result.stdout + result.stderr + + def test_domain_mean_on_a_single_rank_is_the_serial_mean(): values = np.array([61.3, 78.2, 80.1, 72.9]) diff --git a/tests/core/test_parallel_restart.py b/tests/core/test_parallel_restart.py new file mode 100644 index 00000000..0ac9e957 --- /dev/null +++ b/tests/core/test_parallel_restart.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# SPDX-License-Identifier: BSD-3-Clause + +"""A restart chain on several MPI ranks gives exactly the results of a restart chain on a single rank. + +Each chain runs two days: the first writes a restart file, the second starts from it. The restart files and the +results of the second day must not depend on the number of ranks of either leg. +""" + +import filecmp +from pathlib import Path + +import pytest +from test_init import copy_without_column_ids +from test_parallel_runs import CASES, _assert_same_results, _run, requires_mpirun + +_LEG_1 = ["--start-time", "1979-07-01T00:00:00Z", "--end-time", "1979-07-02T00:00:00Z"] +_LEG_2 = ["--start-time", "1979-07-02T00:00:00Z", "--end-time", "1979-07-03T00:00:00Z", "--restart-init"] +RESTART_FILES = ("restart_1979-07-02T00:00:00Z.nc", "restart_1979-07-03T00:00:00Z.nc") + + +def _leg(case: str, leg: list[str], n_ranks: int, chain: Path) -> Path: + """Run one leg of a restart chain in chain, with the chain's restart directory; return the leg's workdir.""" + workdir = chain / ("leg1" if leg is _LEG_1 else "leg2") + return _run([*CASES[case], *leg, "--restart-dir", str(chain / "restart")], n_ranks, workdir) + + +@requires_mpirun +@pytest.mark.parametrize("case, n_ranks_per_leg", [("matlab", (4, 3)), ("elmer", (3, 2))], ids=str) +def test_parallel_restart_chain_equals_serial_chain(tmp_path, case, n_ranks_per_leg): + for leg in (_LEG_1, _LEG_2): + _leg(case, leg, 1, tmp_path / "serial") + for leg, n_ranks in zip((_LEG_1, _LEG_2), n_ranks_per_leg): + _leg(case, leg, n_ranks, tmp_path / "parallel") + + for restart_file in RESTART_FILES: + serial_file = tmp_path / "serial" / "restart" / restart_file + assert filecmp.cmp(serial_file, tmp_path / "parallel" / "restart" / restart_file, shallow=False), restart_file + _assert_same_results(tmp_path / "serial" / "leg2", tmp_path / "parallel" / "leg2") + + +@requires_mpirun +def test_restart_file_without_column_ids_on_several_ranks(tmp_path): + for leg in (_LEG_1, _LEG_2): + _leg("matlab", leg, 1, tmp_path / "serial") + (tmp_path / "parallel" / "restart").mkdir(parents=True) + copy_without_column_ids( + tmp_path / "serial" / "restart" / RESTART_FILES[0], tmp_path / "parallel" / "restart" / RESTART_FILES[0] + ) + + parallel = _leg("matlab", _LEG_2, 3, tmp_path / "parallel") + + _assert_same_results(tmp_path / "serial" / "leg2", parallel) diff --git a/tests/core/test_partitioned_mesh_run.py b/tests/core/test_partitioned_mesh_run.py index 5b58101a..6ec3500a 100644 --- a/tests/core/test_partitioned_mesh_run.py +++ b/tests/core/test_partitioned_mesh_run.py @@ -10,6 +10,7 @@ and in the lapse rate fitted over their own columns. """ +import filecmp import importlib.util import os import shutil @@ -126,3 +127,30 @@ def test_partitioned_run_equals_serial_run(tmp_path): assert sorted(partitioned) == sorted(serial) for key in serial: np.testing.assert_array_equal(partitioned[key], serial[key], err_msg=key) + + +@requires_mpirun +def test_partitioned_run_continues_from_the_restart_file_of_a_serial_run(tmp_path): + """The partitioned run matches the rows of the restart file to its cells by Elmer element id.""" + mesh_dir = tmp_path / "MESH" + _write_mesh(mesh_dir) + _write_dem(tmp_path / "dem.nc") + _run(tmp_path, mesh_dir, 1, "--restart-dir", str(tmp_path / "serial")) + (restart_file,) = (tmp_path / "serial").glob("restart_*.nc") + (tmp_path / "partitioned").mkdir() + shutil.copy(restart_file, tmp_path / "partitioned" / restart_file.name) + + def second_leg(n_ranks: int, restart_dir: str, *args: str) -> dict[str, np.ndarray]: + times = ["--start-time", "1979-07-01T12:00:00Z", "--end-time", "1979-07-02T00:00:00Z"] + return _run( + tmp_path, mesh_dir, n_ranks, *times, "--restart-dir", str(tmp_path / restart_dir), "--restart-init", *args + ) + + serial = second_leg(1, "serial") + partitioned = second_leg(N_PARTITIONS, "partitioned", "--is-partitioned-elmer-mesh") + + assert sorted(partitioned) == sorted(serial) + for key in serial: + np.testing.assert_array_equal(partitioned[key], serial[key], err_msg=key) + (written,) = [f.name for f in (tmp_path / "partitioned").glob("restart_*.nc") if f.name != restart_file.name] + assert filecmp.cmp(tmp_path / "serial" / written, tmp_path / "partitioned" / written, shallow=False)