diff --git a/CHANGELOG.md b/CHANGELOG.md index 23398b8c..1fe9ec86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,12 @@ SPDX-License-Identifier: CC-BY-4.0 * EBFM on several MPI ranks now stops at startup with an error if `--use-part` is given. Every rank previously computed the same partition. * EBFM on several MPI ranks now stops at startup with an error if forcing files (`--forcing-dir`) or restart files (`--restart-dir`, `--restart-init`) are used, also on a partitioned Elmer mesh. On a partitioned mesh, `--forcing-dir` previously failed when reading the first forcing file, `--restart-init` failed at the shape check of the restart file, and `--restart-dir` finished without writing a restart file, with a warning at the end. EBFM on a single rank, e.g. as one component of a coupled MPMD run, is unaffected. * EBFM started by an MPI launcher on several processes (e.g. `mpirun -np 4 ebfm ...`) without mpi4py now stops with an error asking to install `ebfm[mpi]`. Every process previously ran the whole model on its own, and the processes then failed while writing the same output file. Detected via the launcher's `OMPI_COMM_WORLD_SIZE` (Open MPI) or `PMI_SIZE` (MPICH, Intel MPI, `srun --mpi=pmi2`). +* The results of a run on a partitioned Elmer mesh (`--is-partitioned-elmer-mesh`) on several MPI ranks no longer depend on the partitioning. https://github.com/EBFMorg/EBFM/pull/203 + * The mean latitude that selects the water vapour transmissivity of the incoming shortwave radiation is now computed over the whole domain. Each rank previously used the mean latitude of its own partition. + * The lapse rate of potential temperature, used by the sensible and latent heat fluxes, is now fitted over the whole domain. Each rank previously fitted it over its own partition. + * Without `--random-seed`, rank 0 now draws the random seed and shares it with all ranks, so the random forcing is the same over the whole domain. Each rank previously drew its own. + * `--dump-reference` now holds the columns of all ranks, ordered by Elmer element id. Every rank previously wrote its own partition to the same file. + * `--print-diagnostics` now reports the column count and `smb` statistics of the whole domain, logged by rank 0. Each rank previously reported its own partition. # v0.8.0 diff --git a/src/ebfm/core/FINAL_create_restart_file.py b/src/ebfm/core/FINAL_create_restart_file.py index 403bb6d9..16444b1c 100644 --- a/src/ebfm/core/FINAL_create_restart_file.py +++ b/src/ebfm/core/FINAL_create_restart_file.py @@ -58,10 +58,14 @@ def create_boot_file(): logger.info(f"Boot file saved to {io['bootfileout']}") - OUT["timelastsnow_netCDF"] = date2num( - OUT["timelastsnow"], - units="days since 1970-01-01 00:00:00", - calendar="gregorian", + # float64 even if all times are whole days, for which date2num returns integers + OUT["timelastsnow_netCDF"] = np.asarray( + date2num( + OUT["timelastsnow"], + units="days since 1970-01-01 00:00:00", + calendar="gregorian", + ), + dtype=np.float64, ) OUT = {var_name: OUT[var_name] for var_name in RESTART_VARIABLES} diff --git a/src/ebfm/core/INIT.py b/src/ebfm/core/INIT.py index 6cb5712b..0874c15d 100644 --- a/src/ebfm/core/INIT.py +++ b/src/ebfm/core/INIT.py @@ -17,7 +17,7 @@ from ebfm.elmer.mesh import Mesh from .config import TimeConfig, GridConfig, ColumnDiscretizationConfig, iso8601 -from .grid import GridInputType, GridDict, ShadingMethod, number_of_columns, validate_grid +from .grid import GridInputType, GridDict, ShadingMethod, domain_mean, number_of_columns, validate_grid from .restart import validate_all_variables_present, validate_variable_shape from .constants import DAYS_PER_YEAR, LATENT_HEAT_OF_FUSION, MELTING_POINT, SECONDS_PER_DAY, WATER_DENSITY @@ -183,18 +183,20 @@ def init_constants(): return C -def init_grid(grid: GridDict, io, config: GridConfig): +def init_grid(grid: GridDict, io, config: GridConfig, comm): """ Initializes the GridDict based on the given GridConfig. * Depending on provided grid file in GridConfig parses the grid and stores it in GridDict. * Adds elevation to the grid if a DEM file is provided in GridConfig. * Calculates grid slopes and aspects for shading when grid type is MATLAB. + * Computes the mean latitude of the whole domain. Collective over 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. """ grid["is_partitioned"] = config.is_partitioned grid["is_unstructured"] = config.is_unstructured @@ -219,6 +221,7 @@ def init_grid(grid: GridDict, io, config: GridConfig): source_crs_epsg=config.elmer_mesh_crs_epsg, ) + grid["column_id"] = mesh.cell_ids grid["x"], grid["y"] = mesh.x_cells, mesh.y_cells logger.debug("Reading DEM from file and interpolating to grid...") @@ -254,6 +257,7 @@ def init_grid(grid: GridDict, io, config: GridConfig): # 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["slope_x"] = np.zeros_like(grid["x"]) # test values! @@ -325,6 +329,7 @@ def init_grid(grid: GridDict, io, config: GridConfig): grid["y"] = grid["y_2D"].flatten()[mask_flat == 1] grid["z"] = grid["z_2D"].flatten()[mask_flat == 1] grid["ind"] = np.where(mask_flat == 1) + grid["column_id"] = grid["ind"][0] # --------------------------------------------------------------------- # Grid slope and aspect @@ -405,6 +410,10 @@ def init_grid(grid: GridDict, io, config: GridConfig): else: raise ValueError(f"Unsupported grid input type {config.grid_type} specified in configuration.") + # Mean latitude of the whole domain, i.e. over the columns of all MPI ranks, not only those of this rank; the same + # on every rank. The water vapour transmissivity in LOOP_EBM_SWin depends on it. + grid["lat_mean_whole_domain"] = domain_mean(grid["lat"], comm) + validate_grid(grid) return grid diff --git a/src/ebfm/core/LOOP_EBM_SWin.py b/src/ebfm/core/LOOP_EBM_SWin.py index b1b7bdd0..b2237726 100644 --- a/src/ebfm/core/LOOP_EBM_SWin.py +++ b/src/ebfm/core/LOOP_EBM_SWin.py @@ -91,8 +91,8 @@ def main(C, OUT, IN, grid, cpl: Coupler) -> tuple[np.ndarray, dict]: ) ** -1 # SOURCE: McDonald (1960) temp_dew_fahr = 32.0 + 1.8 * (temp_dew_kelvin - 273.15) - # Assign lambda based on latitude region - lat_mean_abs = np.abs(np.mean(grid["lat"])) + # Assign lambda based on the latitude region of the whole domain + lat_mean_abs = np.abs(grid["lat_mean_whole_domain"]) if lat_mean_abs < 20: lambda_value = 2.91 elif lat_mean_abs < 30: diff --git a/src/ebfm/core/LOOP_climate_forcing.py b/src/ebfm/core/LOOP_climate_forcing.py index 92fb9b75..bcf2794f 100644 --- a/src/ebfm/core/LOOP_climate_forcing.py +++ b/src/ebfm/core/LOOP_climate_forcing.py @@ -7,6 +7,7 @@ from netCDF4 import Dataset import numpy as np +from ebfm.core.comm.gather import ColumnGather from ebfm.core.config import ForcingConfig, ForcingType from .LOOP_general_functions import is_first_time_step @@ -17,7 +18,7 @@ logger = logging.getLogger(__name__) -def main(C, grid, IN, t, time, OUT, config: ForcingConfig) -> tuple[dict, dict]: +def main(C, grid, IN, t, time, OUT, config: ForcingConfig, gather: ColumnGather) -> tuple[dict, dict]: """ Meteorological forcing: Specify or read meteorological input and derive associated meteorological fields. @@ -29,7 +30,8 @@ def main(C, grid, IN, t, time, OUT, config: ForcingConfig) -> tuple[dict, dict]: t (int): Current time step. time (dict): Time-related variables. OUT (dict): Output variables from the model. - cpl (Coupler): Coupling object for data exchange with external models. + config (ForcingConfig): Meteorological forcing configuration. + gather (ColumnGather): Gathers the columns of all MPI ranks, for the lapse rate of the whole domain. Returns: Updated IN and OUT dictionaries. Fields include: @@ -90,12 +92,7 @@ def main(C, grid, IN, t, time, OUT, config: ForcingConfig) -> tuple[dict, dict]: # Potential temperature and lapse rate IN["Theta"] = IN["T"] * (C["Pref"] / IN["Pres"]) ** (C["Rd"] / C["Cp"]) - all_same = np.all(grid["z"] == grid["z"][0]) - if not all_same: - poly_coeff = np.polyfit(grid["z"], IN["Theta"], deg=1) - IN["Theta_lapse"] = max(poly_coeff[0], 0.0015) - else: - IN["Theta_lapse"] = 0.0015 + IN["Theta_lapse"] = _potential_temperature_lapse_rate(grid["z"], IN["Theta"], gather) ########################################################### # STORE RELEVANT VARIABLES IN OUT @@ -269,6 +266,35 @@ def read_current_vectorized_field(variable_name: str) -> np.ndarray: return IN +def _potential_temperature_lapse_rate(z, Theta, gather: ColumnGather) -> float: + """ + Lapse rate of potential temperature, fitted over all columns of the domain. + + Collective over gather.comm: the fit runs on the root rank, in column id order, so that the result does not + depend on how the columns are distributed over the ranks. + + @param[in] z elevation of this rank's columns (m) + @param[in] Theta potential temperature of this rank's columns (K) + + @returns the slope of Theta over z (K m-1), at least 0.0015 + """ + # TODO: Gathering z and Theta of all columns on one rank every time step synchronizes all ranks and is expensive + # in parallel runs. + all_z = gather.gather(z) + all_Theta = gather.gather(Theta) + + lapse_rate = None + if gather.is_root: + if np.all(all_z == all_z[0]): + lapse_rate = 0.0015 + else: + lapse_rate = max(np.polyfit(all_z, all_Theta, deg=1)[0], 0.0015) + + if gather.comm.size > 1: + lapse_rate = gather.comm.bcast(lapse_rate, root=gather.root) + return lapse_rate + + def set_random_weather_data(IN, C, time, grid): """ Specify or read meteorological data for the current time-step. diff --git a/src/ebfm/core/comm/gather.py b/src/ebfm/core/comm/gather.py new file mode 100644 index 00000000..be8fd336 --- /dev/null +++ b/src/ebfm/core/comm/gather.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# 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.""" + +import numpy as np +from numpy.typing import NDArray + +from ebfm.core.comm import Comm + + +class ColumnGather: + """ + Gathers per-column arrays of all ranks on a root rank, ordered by column id. + + 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. + """ + + 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 + @param[in] root rank that receives the gathered arrays + """ + self.comm = comm + self.root = root + self.is_root = comm.rank == root + self.n_local_columns = len(column_ids) + + if comm.size == 1: + self.column_ids = column_ids + self._order = None + return + + all_column_ids = self.gather_unordered(column_ids) + if self.is_root: + self._order = np.argsort(all_column_ids, kind="stable") + self.column_ids = all_column_ids[self._order] + else: + self._order = None + self.column_ids = None + + def gather_unordered(self, local: NDArray) -> NDArray | None: + """ + Concatenate an array of every rank along its first axis on the root rank, in rank order. + + The arrays may differ in length along the first axis, but must have the same dtype and the same other + dimensions on every rank. + + @param[in] local this rank's array + + @returns the concatenated array on the root rank, None on all other ranks + """ + if self.comm.size == 1: + return local + + local = np.ascontiguousarray(local) + layouts = self.comm.gather((local.shape, local.dtype), root=self.root) + + if not self.is_root: + self.comm.Gatherv(local, None, root=self.root) + return None + + assert all( + shape[1:] == local.shape[1:] and dtype == local.dtype for shape, dtype in layouts + ), f"Cannot gather arrays of different dtypes or trailing dimensions (shape, dtype per rank): {layouts}" + + values_per_row = int(np.prod(local.shape[1:])) + rows = [shape[0] for shape, _ in layouts] + gathered = np.empty((sum(rows), *local.shape[1:]), dtype=local.dtype) + self.comm.Gatherv(local, (gathered, [n * values_per_row for n in rows]), root=self.root) + return gathered + + def gather(self, local: NDArray) -> NDArray | None: + """ + Gather a per-column array on the root rank, ordered by column id. + + @param[in] local array whose first axis runs over this rank's columns, in the order of its column ids + + @returns the array for the columns of all ranks on the root rank, None on all other ranks + """ + assert ( + local.shape[0] == self.n_local_columns + ), f"Expected one row per column ({self.n_local_columns}), got an array of shape {local.shape}." + + if self.comm.size == 1: + return local + + gathered = self.gather_unordered(local) + if not self.is_root: + return None + return gathered[self._order] diff --git a/src/ebfm/core/grid.py b/src/ebfm/core/grid.py index d611801f..639fb5ba 100644 --- a/src/ebfm/core/grid.py +++ b/src/ebfm/core/grid.py @@ -13,6 +13,11 @@ exception is the shading look-up table, which needs the terrain surrounding each column, including cells that are not columns themselves. It is built from the 2-D DEM of the MATLAB grid type during initialization; the time loop only reads one row of it per column. + +Every column has an id, `column_id`, that is unique over the whole domain, also when the columns are +distributed over several MPI ranks. For a grid read from an Elmer mesh it is the Elmer element id, for a +MATLAB or NetCDF grid the flat index of the column in the 2-D fields (`ind`). Sorting the columns by id +restores the column order of the whole domain. """ from enum import Enum @@ -55,7 +60,19 @@ def number_of_columns(grid: GridDict) -> int: # Fields carrying exactly one entry per snow/firn column. Their common length is `gpsum`. -PER_COLUMN_FIELDS = ("x", "y", "z", "lat", "lon", "mask", "slope_x", "slope_y", "slope_beta", "slope_gamma") +PER_COLUMN_FIELDS = ( + "column_id", + "x", + "y", + "z", + "lat", + "lon", + "mask", + "slope_x", + "slope_y", + "slope_beta", + "slope_gamma", +) def validate_grid(grid: GridDict) -> None: @@ -93,6 +110,32 @@ def validate_grid(grid: GridDict) -> None: assert values.ndim == 1, f"Per-column field '{field}' must be one-dimensional, got shape {values.shape}." assert len(values) == gpsum, f"Per-column field '{field}' has length {len(values)}, expected gpsum={gpsum}." + assert np.issubdtype(grid["column_id"].dtype, np.integer), "Column ids must be integers." + assert len(np.unique(grid["column_id"])) == gpsum, "Column ids must be unique." + + +def domain_mean(values, comm) -> float: + """ + Mean of a per-column field over the columns of all ranks of comm. + + Collective over comm. With a single rank it is np.mean(values). With several ranks it is the sum of the local + sums divided by the number of columns, which adds the values in a different order, so it can differ from the + serial mean in the last bits. + + Parameters: + values: this rank's values of the field + comm: EBFM's MPI communicator + + Returns: + float: the mean, the same on every rank + """ + if comm.size == 1: + return np.mean(values) + + # MPI has no mean reduction; sum and count are reduced together, element-wise + total, count = comm.allreduce(np.array([np.sum(values), len(values)], dtype=np.float64)) + return total / count + class ShadingMethod(Enum): """Available shading algorithms for topographic shading. diff --git a/src/ebfm/main.py b/src/ebfm/main.py index e959f188..193b2998 100644 --- a/src/ebfm/main.py +++ b/src/ebfm/main.py @@ -17,7 +17,7 @@ LOOP_mass_balance, ) from ebfm.core import LOOP_write_to_file, FINAL_create_restart_file -from ebfm.core.grid import number_of_columns +from ebfm.core.comm.gather import ColumnGather 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 ( @@ -74,11 +74,16 @@ def _format_stats(arr: np.ndarray) -> str: return f"min={arr.min():.4e} max={arr.max():.4e} mean={arr.mean():.4e}" -def dump_reference(logger, OUT, filepath: str): - """Save key output arrays to a .npz file for later comparison.""" - import numpy as np +def dump_reference(logger, OUT, filepath: str, gather: ColumnGather): + """ + Save key output arrays to a .npz file for later comparison. + + Collective over gather.comm: the root rank gathers the columns of all ranks and writes them, ordered by column id. + """ + data = {k: gather.gather(OUT[k]) for k in _REFERENCE_KEYS if k in OUT} + if not gather.is_root: + return - data = {k: OUT[k] for k in _REFERENCE_KEYS if k in OUT} missing = [k for k in _REFERENCE_KEYS if k not in OUT] if missing: logger.warning(f"[DUMP] Keys not found in OUT and will be skipped: {missing}") @@ -86,13 +91,19 @@ def dump_reference(logger, OUT, filepath: str): logger.info(f"[DUMP] Reference snapshot saved to '{filepath}' (keys: {list(data.keys())})") -def print_diagnostics(grid, OUT, t): - """Log key diagnostic values each timestep for performance and correctness analysis.""" +def print_diagnostics(grid, OUT, t, gather: ColumnGather): + """ + Log key diagnostic values each timestep for performance and correctness analysis. + + Collective over gather.comm: the root rank gathers the columns of all ranks and logs the values of the whole grid. + """ + smb = gather.gather(OUT["smb"]) if "smb" in OUT else None + smb_cum = gather.gather(OUT["smb_cumulative"]) if "smb_cumulative" in OUT else None + if not gather.is_root: + return - gpsum = number_of_columns(grid) + gpsum = len(gather.column_ids) has_shading = grid.get("has_shading", False) - smb = OUT.get("smb") - smb_cum = OUT.get("smb_cumulative") diagnostics_logger.info(f"[t={t + 1}] gpsum={gpsum}, shading={'on' if has_shading else 'off'}") if smb is not None: @@ -101,6 +112,25 @@ def print_diagnostics(grid, OUT, t): diagnostics_logger.info(f"[t={t + 1}] smb_cumulative: {_format_stats(smb_cum)}") +def set_random_seed(logger, random_seed: int | None, comm): + """ + Seed NumPy's random number generator with the same seed on every rank of comm. + + Collective over comm. Without a given seed, a run on several ranks uses a seed drawn on rank 0, so that all ranks + draw the same random numbers, e.g. the wind speed of the random forcing; a run on a single rank stays unseeded. + + @param[in] random_seed seed from --random-seed, or None + @param[in] comm EBFM's MPI communicator + """ + if random_seed is None and comm.size > 1: + random_seed = int(np.random.SeedSequence().generate_state(1)[0]) if comm.rank == 0 else None + random_seed = comm.bcast(random_seed, root=0) + + if random_seed is not None: + np.random.seed(random_seed) + logger.info(f"NumPy random seed fixed to {random_seed} for reproducibility.") + + def _main_impl(): args = parse_cli_args() @@ -194,12 +224,13 @@ def _main_impl(): time = time_config.to_dict() grid, io, phys, column = INIT.init_config(time_config, grid_config, args.restart_dir, args.restart_init) - if args.random_seed is not None: - np.random.seed(args.random_seed) - logger.info(f"NumPy random seed fixed to {args.random_seed} for reproducibility.") + set_random_seed(logger, args.random_seed, ebfm_comm) C = INIT.init_constants() - grid = INIT.init_grid(grid, io, grid_config) + grid = INIT.init_grid(grid, io, grid_config, ebfm_comm) + + # Rank 0 collects the columns of all ranks, e.g. for --dump-reference + column_gather = ColumnGather(ebfm_comm, grid["column_id"]) OUT, IN, OUTFILE = INIT.init_initial_conditions(C, grid, io, time, column, init_with_restart_file=args.restart_init) @@ -279,7 +310,7 @@ def _main_impl(): IN["Pres"] = data_from_icon_atmo["sfcpres"] # Read/set meteorological forcing - IN, OUT = LOOP_climate_forcing.main(C, grid, IN, t, time, OUT, forcing_config) + IN, OUT = LOOP_climate_forcing.main(C, grid, IN, t, time, OUT, forcing_config, column_gather) # Run surface energy balance model OUT = LOOP_EBM.main(C, OUT, IN, time, grid, coupler) @@ -322,7 +353,7 @@ def _main_impl(): OUT = LOOP_mass_balance.main(OUT, IN, C) if args.print_diagnostics: - print_diagnostics(grid, OUT, t) + print_diagnostics(grid, OUT, t, column_gather) if coupler.has_coupling_to("elmer_ice"): elmer_ice = coupler.get_component("elmer_ice") @@ -369,7 +400,7 @@ def _main_impl(): logger.info("Time loop completed.") if args.dump_reference: - dump_reference(logger, OUT, args.dump_reference) + dump_reference(logger, OUT, args.dump_reference, column_gather) coupler.finalize() diff --git a/tests/core/comm/domain_mean.py b/tests/core/comm/domain_mean.py new file mode 100644 index 00000000..2cac8021 --- /dev/null +++ b/tests/core/comm/domain_mean.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Check domain_mean on 3 MPI ranks that hold the 10 columns of a domain unevenly, rank 2 holding none. + +Usage: mpirun -np 3 python domain_mean.py +""" + +import numpy as np + +from ebfm.core.comm import defaultComm as comm +from ebfm.core.grid import domain_mean + +assert comm.size == 3, f"Run on 3 ranks, not {comm.size}." + +values = 60.0 + 2.5 * np.arange(10) +columns = {0: [5, 1, 7, 3], 1: [0, 9, 2, 8, 4, 6], 2: []}[comm.rank] + +mean = domain_mean(values[columns], comm) + +np.testing.assert_allclose(mean, np.mean(values), rtol=1e-15) +assert comm.allgather(mean) == [mean] * comm.size, "Ranks got different means." +print(f"rank {comm.rank}: ok") diff --git a/tests/core/comm/gather_columns.py b/tests/core/comm/gather_columns.py new file mode 100644 index 00000000..f655cbbb --- /dev/null +++ b/tests/core/comm/gather_columns.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Check ColumnGather on 3 MPI ranks that hold the 10 columns of a domain unevenly and out of id order. + +Usage: mpirun -np 3 python gather_columns.py + +Every column is held by one rank, and rank 2 holds none. Checks the arrays gathered on 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) +values = gather.gather(1.5 * local_ids) +layers = gather.gather(np.stack([local_ids, 2 * local_ids], axis=1)) +concatenated = gather.gather_unordered(local_ids) + +if gather.is_root: + np.testing.assert_array_equal(gather.column_ids, column_ids) + np.testing.assert_array_equal(values, 1.5 * column_ids) + np.testing.assert_array_equal(layers, np.stack([column_ids, 2 * column_ids], axis=1)) + np.testing.assert_array_equal(concatenated, column_ids[[5, 1, 7, 3, 0, 9, 2, 8, 4, 6]]) +else: + assert gather.column_ids is None and values is None and layers is None and concatenated is None + +print(f"rank {comm.rank}: ok") diff --git a/tests/core/test_init.py b/tests/core/test_init.py index 1104fe47..41d6f5bb 100644 --- a/tests/core/test_init.py +++ b/tests/core/test_init.py @@ -18,6 +18,7 @@ from ebfm.core.INIT import init_config, init_grid, init_initial_conditions 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.config import ColumnDiscretizationConfig from ebfm.core.config.grid import GridConfig from ebfm.core.config.time import TimeConfig @@ -352,7 +353,7 @@ 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) + return init_grid(grid, io, grid_config, defaultComm) def test_shading_disabled_skips_the_lookup_table(self): grid = self._init_grid("--no-shading") diff --git a/tests/core/test_parallel_columns.py b/tests/core/test_parallel_columns.py new file mode 100644 index 00000000..b635e0ab --- /dev/null +++ b/tests/core/test_parallel_columns.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# 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.""" + +import importlib.util +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +from ebfm.core.comm import defaultComm +from ebfm.core.comm.gather import ColumnGather +from ebfm.core.grid import domain_mean + +COMM_TESTS = Path(__file__).resolve().parent / "comm" +TIMEOUT = 120 + +requires_mpirun = pytest.mark.skipif( + shutil.which("mpirun") is None or importlib.util.find_spec("mpi4py") is None, + reason="mpirun or mpi4py not available", +) + + +def _mpirun(n_ranks: int, script: str, *args: str) -> subprocess.CompletedProcess: + # Once mpi4py has initialised MPI in this process, the process environment holds the launcher variables of an + # MPI singleton, and an mpirun inheriting them exits 1 without running anything. + env = {k: v for k, v in os.environ.items() if not k.startswith(("OMPI_", "PMIX_", "PMI_", "HYDRA_"))} + cmd = ["mpirun", "-np", str(n_ranks), sys.executable, str(COMM_TESTS / script), *args] + return subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=TIMEOUT) + + +def test_single_rank_gathers_without_communication(): + column_ids = np.array([4, 2, 9]) + values = np.array([1.0, 2.0, 3.0]) + + gather = ColumnGather(defaultComm, column_ids) + + assert gather.is_root + assert gather.gather(values) is values + assert gather.column_ids is column_ids + + +@requires_mpirun +def test_gather_on_three_ranks(): + result = _mpirun(3, "gather_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]) + + assert domain_mean(values, defaultComm) == np.mean(values) + + +@requires_mpirun +def test_domain_mean_on_three_ranks(): + result = _mpirun(3, "domain_mean.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 diff --git a/tests/core/test_partitioned_mesh_run.py b/tests/core/test_partitioned_mesh_run.py new file mode 100644 index 00000000..82d9b9db --- /dev/null +++ b/tests/core/test_partitioned_mesh_run.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# SPDX-License-Identifier: BSD-3-Clause + +"""The results of EBFM on a partitioned Elmer mesh do not depend on the partitioning. + +A serial run on the whole mesh and a run on 3 MPI ranks, one partition each, must give equal --dump-reference +snapshots. The mesh spans 70 degrees north, a bin boundary of the water vapour coefficient in LOOP_EBM_SWin, and its +elevation rises non-linearly from south to north, so the partitions, latitude bands, differ in their latitude mean +and in the lapse rate fitted over their own columns. +""" + +import importlib.util +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +from netCDF4 import Dataset + +TIMEOUT = 600 +N_PARTITIONS = 3 + +# EPSG:3413 coordinates of the mesh, latitudes 67.6 to 71.2 degrees north +X0, Y0 = -100000.0, -2450000.0 +NX, NY = 4, 10 +SPACING = 40000.0 + +_ARGS = [ + "--netcdf-dem-mesh", + "dem.nc", + "--elmer-mesh-crs-epsg", + "3413", + "--random-seed", + "1", + "--start-time", + "1979-07-01T00:00:00Z", + "--end-time", + "1979-07-01T12:00:00Z", + "--dump-reference", + "reference.npz", +] + +requires_mpirun = pytest.mark.skipif( + shutil.which("mpirun") is None or importlib.util.find_spec("mpi4py") is None, + reason="mpirun or mpi4py not available", +) + + +def _write_mesh(mesh_dir: Path) -> None: + """ + Write a rectangle of NX x NY squares, two triangles each, as an Elmer mesh and as N_PARTITIONS partitions of it. + + Each partition is a contiguous block of elements, holding the nodes its elements use, so the element ids of the + partitions, in rank order, are those of the whole mesh. + """ + mesh_dir.mkdir() + node_id = np.arange(1, (NX + 1) * (NY + 1) + 1).reshape(NY + 1, NX + 1) + nodes = { + node_id[j, i]: f"{node_id[j, i]} -1 {X0 + i * SPACING} {Y0 + j * SPACING} 0.0\n" + for j in range(NY + 1) + for i in range(NX + 1) + } + + triangles = [] + for j in range(NY): + for i in range(NX): + a, b, c, d = node_id[j, i], node_id[j, i + 1], node_id[j + 1, i + 1], node_id[j + 1, i] + triangles += [(a, b, c), (a, c, d)] + elements = [(element_id, triangle) for element_id, triangle in enumerate(triangles, start=1)] + + def write(prefix: str, part_elements) -> None: + part_nodes = sorted({node for _, triangle in part_elements for node in triangle}) + (mesh_dir / f"{prefix}.header").write_text(f"{len(part_nodes)} {len(part_elements)} 0\n") + (mesh_dir / f"{prefix}.nodes").write_text("".join(nodes[node] for node in part_nodes)) + (mesh_dir / f"{prefix}.elements").write_text( + "".join(f"{element_id} 1 303 {a} {b} {c}\n" for element_id, (a, b, c) in part_elements) + ) + + write("mesh", elements) + for partition_id, part_elements in enumerate(np.array_split(np.arange(len(elements)), N_PARTITIONS), start=1): + write(f"part.{partition_id}", [elements[k] for k in part_elements]) + + +def _write_dem(dem_file: Path) -> None: + """Write a DEM for --netcdf-dem-mesh whose surface rises quadratically from 0 m in the south to 2000 m.""" + x = X0 + SPACING / 4 * np.arange(-1, 4 * NX + 2) + y = Y0 + SPACING / 4 * np.arange(-1, 4 * NY + 2) + surface = 2000.0 * ((y - y[0]) / (y[-1] - y[0])) ** 2 + + with Dataset(dem_file, "w") as nc: + nc.createDimension("x", len(x)) + nc.createDimension("y", len(y)) + nc.createVariable("x", "f8", ("x",))[:] = x + nc.createVariable("y", "f8", ("y",))[:] = y + nc.createVariable("surface", "f8", ("y", "x"))[:] = np.repeat(surface[:, np.newaxis], len(x), axis=1) + + +def _run(workdir: Path, mesh_dir: Path, n_ranks: int, *args: str) -> dict[str, np.ndarray]: + """Run EBFM on n_ranks ranks in workdir; return its --dump-reference snapshot.""" + # Once mpi4py has initialised MPI in this process, the process environment holds the launcher variables of an + # MPI singleton, and an mpirun inheriting them exits 1 without running anything. + env = {k: v for k, v in os.environ.items() if not k.startswith(("OMPI_", "PMIX_", "PMI_", "HYDRA_"))} + cmd = [sys.executable, "-m", "ebfm.main", "--elmer-mesh", str(mesh_dir), *_ARGS, *args] + if n_ranks > 1: + cmd = ["mpirun", "-np", str(n_ranks), *cmd] + + result = subprocess.run(cmd, cwd=workdir, env=env, capture_output=True, text=True, timeout=TIMEOUT) + assert result.returncode == 0, f"{' '.join(cmd)} failed:\n{result.stdout}\n{result.stderr}" + with np.load(workdir / "reference.npz") as reference: + return dict(reference) + + +@requires_mpirun +def test_partitioned_run_equals_serial_run(tmp_path): + mesh_dir = tmp_path / "MESH" + _write_mesh(mesh_dir) + _write_dem(tmp_path / "dem.nc") + + 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)