test_coupling1.py

Go to the documentation of this file.

#!/usr/bin/env python3
 
# Copyright (c) 2025 The YAC Authors
#
# SPDX-License-Identifier: BSD-3-Clause
 
 
 
import yac
import numpy as np
 
try:
    from mpi4py import MPI
except:
    pass
 
yac.def_calendar(yac.Calendar.PROLEPTIC_GREGORIAN)
 
yac_instance = yac.YAC()
assert yac_instance.size == 4
yac_instance.def_datetime("2020-01-01T00:00:00",
                          "2020-01-02T00:00:00")
 
rank = yac_instance.rank
 
if rank < 2:
    comp = yac_instance.def_comp("comp_A")
    grid = yac.Reg2dGrid("grid_A",
                         np.linspace(0, 2*np.pi, 180, endpoint=False),
                         np.linspace(-np.pi/2, np.pi/2, 90),
                         cyclic=[True,False])
    points = grid.def_points(yac.Location.CELL,
                             np.pi/180 + np.linspace(0, 2*np.pi, 180, endpoint=False),
                             np.pi/90 + np.linspace(-np.pi/2, np.pi/2, 89, endpoint=False))
    field = yac.Field.create("field_A", comp, points,
                             collection_size=3,
                             timestep="1", timeunit=yac.TimeUnit.MINUTE)
    yac_instance.enddef()
    t = 0
    while field.action != yac.Action.OUT_OF_BOUND:
        field.put(t + np.stack([1*np.ones(points.size),
                                2*np.ones(points.size),
                                3*np.ones(points.size)]
                               )
                  )
        t += 1
else:
    comp = yac_instance.def_comp("comp_B")
    grid = yac.UnstructuredGrid(
        "grid_B",
        [3],
        np.deg2rad([0.0, 1.0, 0.0]),
        np.deg2rad([0.0, 0.0, 1.0]),
        [0, 1, 2],
    )
    points = grid.def_points(yac.Location.CELL,
                             np.deg2rad([1/3]),
                             np.deg2rad([1/3])
                             )
    field = yac.Field.create("field_B", comp, points,
                             collection_size = 2,
                             timestep="1", timeunit=yac.TimeUnit.HOUR)
 
    interp_stack = yac.InterpolationStack()
    interp_stack.add_nnn()
 
    yac_instance.def_couple("comp_A", "grid_A", "field_A",
                            "comp_B", "grid_B", "field_B",
                            coupling_timestep="1", timeunit=yac.TimeUnit.HOUR,
                            time_reduction=yac.Reduction.TIME_NONE,
                            interp_stack=interp_stack,
                            collection_selection=[2,1])
    yac_instance.enddef()
    t = 0
    while field.action != yac.Action.OUT_OF_BOUND:
        data, info = field.get()
        assert np.allclose(data,
                           60*t + np.stack([3*np.ones(points.size),
                                            2*np.ones(points.size)]
                                        )
                           )
        t += 1