Make the results on a partitioned Elmer mesh independent of the partitioning

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
 

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}
 

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

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:

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.

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]

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.