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/CHANGELOG.md b/CHANGELOG.md index df0a7ac4..81ae67a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,13 @@ SPDX-License-Identifier: CC-BY-4.0 * In a run on several MPI ranks, only rank 0 logs to the console at `--log-level-console`; the other ranks only log warnings and errors there, instead of repeating every message. Log files (`--log-file`, one per rank) are unchanged. The current model time is logged instead of printed. https://github.com/EBFMorg/EBFM/pull/204 * A MATLAB grid (`--matlab-mesh`) can now run on several MPI ranks, e.g. `mpirun -np 4 ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading`: each rank computes a contiguous block of the columns, and rank 0 writes the output file, which is identical to that of a run on a single rank. Shading is not supported on several ranks and stops the run with an error. https://github.com/EBFMorg/EBFM/pull/202 * A usage error, e.g. a rejected command line option, now ends the whole MPI job with its exit status when `MPI.COMM_WORLD` has more than one process. EBFM previously exited on its own, and the other processes, e.g. a coupled model in an MPMD run, waited for it forever. A run on a single process still exits without `MPI_ABORT`. https://github.com/EBFMorg/EBFM/pull/206 +* In a run on several MPI ranks, EBFM now checks before the time loop that every column is held by exactly one rank. The check communicates the ids of all columns; `--no-check-grid-distribution` skips it. https://github.com/EBFMorg/EBFM/pull/197 +* Partitioned Elmer meshes (`--is-partitioned-elmer-mesh`) now write model output (`Output/model_output.nc`); it was skipped before. On several MPI ranks, rank 0 gathers the cells of all ranks and writes a single file, with the cells ordered by Elmer element id and the vertices that partitions share written once. https://github.com/EBFMorg/EBFM/pull/197 +* A non-partitioned Elmer mesh (`--elmer-mesh` without `--is-partitioned-elmer-mesh`, also with `--netcdf-dem-mesh`) can now run on several MPI ranks, e.g. `mpirun -np 4 ebfm --elmer-mesh ...`: each rank computes a contiguous block of the cells, the first block on rank 0. The output file equals that of a run on a single rank. Fitting the lapse rate of potential temperature over all columns synchronizes all ranks every time step; EBFM warns about this. Every rank reads the whole mesh, which EBFM also warns about; for large meshes, a mesh with one partition per rank (`--is-partitioned-elmer-mesh`, which requires `--netcdf-dem-mesh`) avoids this. https://github.com/EBFMorg/EBFM/pull/197 +* A run on a partitioned Elmer mesh (`--is-partitioned-elmer-mesh`) whose number of MPI ranks differs from the number of partitions now stops with an error. Each rank computes one partition, so such a run previously skipped partitions without notice. `--allow-partition-mismatch` allows it, with a warning, e.g. to compute a single partition with `--use-part`. https://github.com/EBFMorg/EBFM/pull/197 +* A NetCDF grid (`--netcdf-mesh`) can now run on several MPI ranks like a MATLAB grid: each rank computes a contiguous block of the columns, and rank 0 writes the output file, which is identical to that of a run on a single rank. https://github.com/EBFMorg/EBFM/pull/197 +* `--netcdf-dem-mesh-unstructured` now works on several MPI ranks. EBFM looks up the values of the XIOS file by Elmer node id (the value of node id i at position i - 1), so each rank reads the values of its own nodes. A file without values for all node ids of the mesh is rejected with an error. https://github.com/EBFMorg/EBFM/pull/197 +* Restart files (`--restart-dir`, `--restart-init`) now work on several MPI ranks and with partitioned Elmer meshes. Rank 0 writes the columns of all ranks, ordered by column id, together with their ids in a new variable `column_id`, and reads them back for any distribution of the columns over the ranks: e.g. a file written on 4 ranks can be continued on 3 ranks or on a partitioned mesh. A file without `column_id` holds the columns in the order of a run on a single rank; on several ranks, its `x` and `y` are then checked against the grid. A file whose column ids differ from those of the grid is rejected. https://github.com/EBFMorg/EBFM/pull/197 # v0.8.0 diff --git a/README.md b/README.md index 7f5f80fd..1d7fda31 100644 --- a/README.md +++ b/README.md @@ -216,16 +216,21 @@ path/to/your/elmer/mesh Usage example for partitioned mesh: ```sh -ebfm --elmer-mesh examples/greenland_mesh_v0/MESH/partitioning.128/ --netcdf-dem-mesh examples/BedMachineGreenland-v5_lo.nc --is-partitioned-elmer-mesh --use-part 42 --elmer-mesh-crs-epsg 3413 +ebfm --elmer-mesh examples/greenland_mesh_v0/MESH/partitioning.128/ --netcdf-dem-mesh examples/BedMachineGreenland-v5_lo.nc --is-partitioned-elmer-mesh --use-part 42 --elmer-mesh-crs-epsg 3413 --allow-partition-mismatch ``` On several MPI ranks, `--use-part` is not allowed: rank i computes the partition i + 1, e.g. the partitions 1 to 4 with: ```sh -mpirun -n 4 ebfm --elmer-mesh examples/greenland_mesh_v0/MESH/partitioning.128/ --netcdf-dem-mesh examples/BedMachineGreenland-v5_lo.nc --is-partitioned-elmer-mesh --elmer-mesh-crs-epsg 3413 +mpirun -n 4 ebfm --elmer-mesh examples/greenland_mesh_v0/MESH/partitioning.128/ --netcdf-dem-mesh examples/BedMachineGreenland-v5_lo.nc --is-partitioned-elmer-mesh --elmer-mesh-crs-epsg 3413 --allow-partition-mismatch ``` -Besides a partitioned mesh, only a MATLAB grid runs on several MPI ranks (see [Running on several MPI ranks](#running-on-several-mpi-ranks)). +Rank 0 writes the output of all ranks to a single file, with the cells ordered by their Elmer element id. + +Besides a partitioned mesh, MATLAB and NetCDF grids and a non-partitioned Elmer mesh run on several MPI ranks (see [Running on several MPI ranks](#running-on-several-mpi-ranks)). + +A run needs one MPI rank per partition. To compute only some partitions, e.g. a +single one with `--use-part` as above, pass `--allow-partition-mismatch`. Note: The partitioned mesh is not included in this repository. @@ -286,20 +291,32 @@ ebfm --matlab-mesh examples/dem_and_mask.mat --restart-dir examples/restart --re Note that it is necessary to provide a `--start-time` consistent with the given restart file. +A restart file holds the columns of the whole grid together with their ids +(`column_id`), so a run may continue from it on any number of MPI ranks and on +a partitioned mesh. A restart file without `column_id`, written by an older +version of EBFM, holds the columns in the order of a run on a single rank. + ### Running on several MPI ranks -EBFM splits a MATLAB grid over the MPI ranks of a run (requires `pip install 'ebfm[mpi]'`): +EBFM splits the grid over the MPI ranks of a run (requires `pip install 'ebfm[mpi]'`): ```sh mpirun -np 4 ebfm --matlab-mesh examples/dem_and_mask.mat --no-shading ``` -* Each rank computes a contiguous block of the columns, the first block on rank 0. -* Rank 0 writes the output file and `--dump-reference` for all ranks. They are - identical to those of a run on a single rank. -* Not supported on several ranks, and rejected with an error: shading, - `--forcing-dir`, `--restart-dir`/`--restart-init`, and grids other than a - MATLAB grid or a partitioned Elmer mesh. +* MATLAB and NetCDF grids and Elmer meshes are split into contiguous blocks of + columns, the first block on rank 0. A partitioned Elmer mesh gives one + partition to each rank (see above). +* With a non-partitioned Elmer mesh, every rank reads the whole mesh, and EBFM + warns about it. For large meshes, use a partitioned mesh (requires + `--netcdf-dem-mesh`). +* Rank 0 writes the output file, the restart file and `--dump-reference` for + all ranks, and reads the restart file. They are identical to those of a run + on a single rank. +* Not supported on several ranks, and rejected with an error: shading and + `--forcing-dir`. +* `--no-check-grid-distribution` skips the start-up check that every column is + held by exactly one rank. ### Performance and Profiling Runs @@ -335,7 +352,7 @@ Follow the install instructions from above and run the example command for a cou mpirun -np 1 ebfm \ --elmer-mesh $MESHES/greenland_mesh_v0/MESH/partitioning.128/ \ --netcdf-dem-mesh $DATA/BedMachineGreenland-v5.nc \ - --is-partitioned-elmer-mesh --use-part 1 \ + --is-partitioned-elmer-mesh --use-part 1 --allow-partition-mismatch \ --coupler-config $CPL_CONFIG --couple-to-elmer-ice --couple-to-icon-atmo \ : \ -np 1 $ELMER_ROOT/src/elmer_dummy_f.x $MESHES/greenland_mesh_v0/MESH/partitioning.128 1 $CPL_CONFIG \ 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/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 99433bfa..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 @@ -183,7 +185,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 +196,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. """ @@ -222,6 +223,11 @@ def init_grid(grid: GridDict, io, 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 @@ -252,15 +258,23 @@ def init_grid(grid: GridDict, io, 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! @@ -670,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. @@ -682,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. @@ -705,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/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/cli.py b/src/ebfm/core/cli.py index 284671f7..348483e9 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]}.", ) @@ -349,6 +349,21 @@ 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.", + ) + + 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( @@ -551,38 +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 MATLAB grid over the ranks; of a partitioned Elmer mesh, each rank computes one partition. - - 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.MATLAB: - # EBFM splits the columns over the ranks - return - - 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" - 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." - ) - - 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. @@ -613,16 +596,45 @@ 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_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. + 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] comm EBFM's MPI communicator + @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 (args.restart_dir is not None or args.restart_init) and 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 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) + + +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/core/comm/gather.py b/src/ebfm/core/comm/gather.py index be8fd336..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,19 +13,20 @@ 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): """ @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 @@ -37,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") @@ -95,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/config/grid.py b/src/ebfm/core/config/grid.py index f12f2ac7..4d5e5eb7 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 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 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 @@ -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 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 = [ @@ -84,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...") @@ -118,7 +120,7 @@ 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: + 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/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/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/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 e2ab2941..a5ee6da4 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 ( @@ -25,11 +26,11 @@ parse_cli_args, resolve_numba_threads, validate_shading_coupling_compat, - validate_grid_parallel_compat, validate_use_part_parallel_compat, validate_forcing_parallel_compat, - validate_restart_parallel_compat, + validate_partition_count, validate_shading_parallel_compat, + warn_non_partitioned_mesh_parallel, ) import ebfm.coupling @@ -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) @@ -210,18 +213,16 @@ 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. - 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) - # Ensure restart files are only read and written in runs on a single rank. - validate_restart_parallel_compat(args, ebfm_comm) - # 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) + warn_non_partitioned_mesh_parallel(grid_config, ebfm_comm, logger) + logger.debug("Successfully completed consistency checks.") # Model setup & initialization @@ -231,14 +232,27 @@ 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) + + # 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 + # 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 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) + 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, gather=column_gather + ) coupler.setup(grid, time_config) @@ -383,25 +397,17 @@ 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"] 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 - 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, column_gather) logger.info("Time loop completed.") 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/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/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/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/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 bb3e41f3..eb55fb31 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,11 +17,11 @@ mpi_launcher_size, parse_cli_args, 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, + warn_non_partitioned_mesh_parallel, ) from ebfm.core.comm import mpi_available @@ -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.""" @@ -405,47 +413,6 @@ def test_shading_on_several_ranks_rejected(self): self.assertEqual(ctx.exception.code, 2) -class TestValidateGridParallelCompat(unittest.TestCase): - """validate_grid_parallel_compat rejects every grid but MATLAB and a partitioned Elmer mesh 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 --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) - self.assertIn("Only --matlab-mesh and a partitioned Elmer mesh", 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.""" @@ -493,43 +460,70 @@ 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.""" - - _RESTART_ARGS = { - "--restart-dir": ["--restart-dir", "restart/"], - "--restart-init": ["--restart-init"], - "--restart-dir --restart-init": ["--restart-dir", "restart/", "--restart-init"], - } +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, 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, 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_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 + def test_one_rank_per_partition_ok(self): + self._validate() # must not raise + self.logger.warning.assert_not_called() - def test_no_restart_on_several_ranks_ok(self): - self._validate([], comm_size=4) # must not raise + def test_non_partitioned_mesh_ok(self): + self._validate(is_partitioned=False, num_partitions=1) # 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) + def test_mismatch_rejected(self): + with self.assertRaises(SystemExit) as ctx: + self._validate(num_partitions=128) + self.assertEqual(ctx.exception.code, 2) - @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_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 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): diff --git a/tests/core/test_elmer_mesh.py b/tests/core/test_elmer_mesh.py new file mode 100644 index 00000000..731af314 --- /dev/null +++ b/tests/core/test_elmer_mesh.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: 2026 EBFM Authors +# +# 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 + + +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) + + +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_init.py b/tests/core/test_init.py index 9381f61c..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" @@ -352,8 +418,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 +456,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) 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 b635e0ab..16d22319 100644 --- a/tests/core/test_parallel_columns.py +++ b/tests/core/test_parallel_columns.py @@ -2,8 +2,10 @@ # # 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, scattered back and reduced over all ranks.""" +import filecmp import importlib.util import os import shutil @@ -16,7 +18,8 @@ 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 +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 @@ -55,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]) @@ -68,3 +88,34 @@ 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 + + +@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) 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_parallel_runs.py b/tests/core/test_parallel_runs.py index 9d717906..0ea9e87d 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,8 @@ 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( @@ -31,7 +33,21 @@ # 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], + "elmerxios": [*_ELMER, "--netcdf-dem-mesh-unstructured", str(XIOS_DEM), *_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 +86,21 @@ 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), + ("elmerxios", 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/core/test_partitioned_mesh_run.py b/tests/core/test_partitioned_mesh_run.py index 82d9b9db..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 @@ -123,9 +124,33 @@ 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) + + +@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) 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 00000000..c24d6c65 Binary files /dev/null and b/tests/data/elmer_mesh/xios_dem.nc differ 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. 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)