test_core.py

Go to the documentation of this file.

#!/usr/bin/env python3
 
# Copyright (c) 2024 The YAC Authors
#
# SPDX-License-Identifier: BSD-3-Clause
 
 
 
try:
    from mpi4py import MPI  # calls MPI_Init with some special treatments
    comm = MPI.COMM_WORLD
    rank = comm.rank
except:
    comm = None
    rank = 0
 
from yac.core import (BasicGrid,
                      InterpField,
                      InterpolationStack,
                      compute_weights,
                      yac_location,
                      yac_weight_file_on_existing,
                      lonlat2xyz,
                      get_io_ranks,
                      )
import numpy as np
import sys
import asyncio
import logging
 
logging.basicConfig(level=logging.WARNING, format=f'%(name)s rank={rank} - %(levelname)s - %(message)s')
logging.getLogger("yac.core").setLevel(logging.INFO)
 
grid_dir = sys.argv[1]
 
reg_grid = BasicGrid.reg_2d_new("reg_grid",
                                np.linspace(-2, 2, 20),
                                np.linspace(-1.5, 1.5, 20))
reg_grid.set_core_mask(yac_location.YAC_LOC_CORNER, np.ones(400, dtype=int))
 
# assert that set_global_index raises an exception if indices are out of bounds
try:
    reg_grid.set_global_index(yac_location.YAC_LOC_CORNER, range(10**99, 10**99+400))
    raise RuntimeError("set_global_index did not throw an exception for very large indices.")
except AssertionError:
    pass
 
reg_grid.set_global_index(yac_location.YAC_LOC_CORNER, range(400))
reg_grid.to_file("reg_grid.nc", comm)
 
coords = reg_grid.add_coordinates(yac_location.YAC_LOC_CORNER)
src_field = InterpField(coords)
 
icon_grid, icon_cell_field = BasicGrid.read_icon(grid_dir+"icon_grid_0030_R02B03_G.nc",
                                                 "icon", comm=comm)
 
interp_stack = InterpolationStack.from_list([("nnn", {"n": 3, "max_search_distance": 2.}),
                                             ("fixed", {"value": 42.0})])
 
weights = compute_weights(interp_stack,
                          src_field,
                          icon_cell_field)
 
weights.write_to_file("weights.nc", on_existing=yac_weight_file_on_existing.YAC_WEIGHT_FILE_OVERWRITE)
 
interpolate = weights.get_interpolation(collection_selection=[1, 4, 2])
 
src_size = src_field.basic_grid.get_data_size(yac_location.YAC_LOC_CORNER)
src = np.stack([i*np.ones(src_size, dtype=np.float64) for i in range(5)], axis=0)
 
tgt = interpolate(src)
assert np.all(tgt[0,...] == 1.0), f"Wrong result: {tgt=}"
assert np.all(tgt[1,...] == 4.0), f"Wrong result: {tgt=}"
assert np.all(tgt[2,...] == 2.0), f"Wrong result: {tgt=}"
 
local_is_io, io_ranks = get_io_ranks()
assert local_is_io == True
assert len(io_ranks) == 1
assert io_ranks[0] == 0
 
async def _main():
    tgt = await interpolate.execute_async(src)
    assert np.all(tgt[0,...] == 1.0), f"Wrong result: {tgt=}"
    assert np.all(tgt[1,...] == 4.0), f"Wrong result: {tgt=}"
    assert np.all(tgt[2,...] == 2.0), f"Wrong result: {tgt=}"
 
asyncio.run(_main())