ci: add test_raw_exchange.py

configure.ac

@@ -966,6 +966,7 @@ AS_VAR_IF([enable_python_bindings], [yes],
      tests/python/test_noise2nc_parallel.sh
      tests/python/test_noise2plot.sh
      tests/python/test_noise2plot_parallel.sh
+     tests/python/test_raw_exchange.sh
 ], [chmod a+x "$ac_file"])
    AC_CONFIG_FILES([
      python/setup.py

tests/python/Makefile.am

@@ -18,7 +18,8 @@ TESTS = \
     test_noise2nc.sh \
     test_noise2nc_parallel.sh \
     test_noise2plot.sh \
-    test_noise2plot_parallel.sh
+    test_noise2plot_parallel.sh \
+    test_raw_exchange.sh
 
 if TEST_MPI
 # Some tests require output files of the others:
@@ -32,12 +33,14 @@ test_noise2nc_parallel.log: test_nc2plot.log
 test_noise2nc.log: test_multithreading.log
 test_multithreading.log: test_interface.log
 test_interface.log: test_coroutines.log
+test_coroutines.log: test_raw_exchange.log
 endif
 endif
 
 EXTRA_DIST = \
     test_coroutines.py \
-    test_interface.py
+    test_interface.py \
+    test_raw_exchange.py
 
 clean-local:
 	-rm -rf *.nc *.plots

tests/python/test_raw_exchange.py

+#!/usr/bin/env python3
+
+# Copyright (c) 2024 The YAC Authors
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import numpy as np
+
+from yac import (
+    YAC,
+    def_calendar,
+    Calendar,
+    Reg2dGrid,
+    Location,
+    TimeUnit,
+    Field,
+    InterpolationStack,
+    ConservNormalizationType,
+    Reduction,
+)
+
+def_calendar(Calendar.PROLEPTIC_GREGORIAN)
+
+yac = YAC()
+
+yac.def_datetime("2020-01-01T00:00:00", "2020-01-02T00:00:00")
+
+comp1, comp2 = yac.def_comps(["comp1", "comp2"])
+
+grid1 = Reg2dGrid("grid1", [-1, 0, 1], [-1, 0, 1])
+grid1.set_core_mask([1, 1, 0, 1], Location.CELL)
+points1 = grid1.def_points(Location.CELL, [-0.5, 0.5], [-0.5, 0.5])
+
+grid2 = Reg2dGrid("grid2", [-1, -1/3, 1/3, 1], [-1, -1/3, 1/3, 1])
+points2 = grid2.def_points(Location.CELL, [-2/3, 0.0, 2/3], [-2/3, 0.0, 2/3])
+
+field1 = Field.create("field1", comp1, points1, 1, "1", TimeUnit.HOUR)
+field2 = Field.create("field2", comp2, points2, 1, "1", TimeUnit.HOUR)
+field3 = Field.create("field3", comp2, points2, 1, "1", TimeUnit.HOUR)
+
+interp = InterpolationStack()
+interp.add_conservative(1, 0, 1, ConservNormalizationType.DESTAREA)
+interp.add_fixed(42.0)
+
+couple_kwargs = {"src_comp": "comp1",
+                 "src_grid": "grid1",
+                 "src_field": "field1",
+                 "tgt_comp": "comp2",
+                 "tgt_grid": "grid2",
+                 "coupling_timestep": "60",
+                 "timeunit": TimeUnit.MINUTE,
+                 "time_reduction": Reduction.TIME_NONE,
+                 "interp_stack": interp}
+
+yac.def_couple(
+    **couple_kwargs,
+    tgt_field="field2",
+    use_raw_exchange=True
+)
+
+yac.def_couple(
+    **couple_kwargs,
+    tgt_field="field3"
+)
+
+
+yac.enddef()
+
+raw_data = field2.get_raw_interp_weights_data()
+indptr = np.insert(np.cumsum(raw_data.num_src_per_tgt), 0, 0)
+
+buf = np.empty(shape=(1, raw_data.src_field_buffer_sizes[0]), dtype=np.float64)
+
+i = 0.
+while field2.datetime < yac.end_datetime:
+    field1.put(np.arange(i, i + 4., dtype=np.float64))
+    buf[:] = 0.0
+    field2.get_raw(buf)
+    f2_buf = np.empty(shape=(1, 9), dtype=np.float64)
+    f2_buf[0, raw_data.tgt_idx_fixed] = raw_data.fixed_values[0]
+    for tgt_idx, start_idx, end_idx in zip(raw_data.wgt_tgt_idx, indptr[:-1], indptr[1:]):
+        f2_buf[0, tgt_idx] = raw_data.weights[start_idx:end_idx]@buf[0, raw_data.src_idx[start_idx: end_idx]]
+
+    f3_buf, info = field3.get()
+
+    assert np.allclose(f2_buf, f3_buf)
+
+    i += 1

tests/python/test_raw_exchange.sh.in

+#!@SHELL@
+# Copyright (c) 2024 The YAC Authors
+#
+# SPDX-License-Identifier: CC0-1.0
+
+@TEST_MPI_FALSE@exit 77
+
+PYTHONPATH="@abs_top_builddir@/python:$PYTHONPATH"
+export PYTHONPATH
+
+@MPI_LAUNCH@ -n 1 @PYTHON@ '@srcdir@/test_raw_exchange.py'