Add atlas-global-matrix sandbox

Co-authored-by: Willem Deconinck <willemdeconinck@ecmwf.int>

src/sandbox/interpolation/CMakeLists.txt

@@ -28,3 +28,23 @@ ecbuild_add_executable(
     LIBS    atlas
     NOINSTALL
 )
+
+# Overcome Cray linking problem where the C++ library is not linked with the C library
+list( APPEND NetCDF_CXX_EXTRA_LIBRARIES NetCDF::NetCDF_C )
+
+ecbuild_add_option( FEATURE NETCDF
+                    DESCRIPTION "Compile support the SCRIP format"
+                    REQUIRED_PACKAGES "NetCDF COMPONENTS C CXX" )
+
+ecbuild_add_executable(
+    TARGET  atlas-global-matrix
+    SOURCES atlas-global-matrix.cc
+            ScripIO.h
+            ScripIO.cc
+    LIBS    atlas
+)
+
+target_compile_definitions( atlas-global-matrix PRIVATE ATLAS_HAVE_NETCDF=${HAVE_NETCDF} )
+if( HAVE_NETCDF )
+    target_link_libraries( atlas-global-matrix NetCDF::NetCDF_CXX )
+endif()

src/sandbox/interpolation/ScripIO.cc

@@ -0,0 +1,167 @@
+/*
+ * (C) Copyright 2025- ECMWF.
+ *
+ * This software is licensed under the terms of the Apache Licence Version 2.0
+ * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+ *
+ * In applying this licence, ECMWF does not waive the privileges and immunities
+ * granted to it by virtue of its status as an intergovernmental organisation nor
+ * does it submit to any jurisdiction.
+ */
+
+// Note: reading of SCRIP format adapted from
+//       https://github.com/ecmwf/mir/blob/develop/src/tools/mir-scrip-to-weight-matrix.cc
+
+#include "ScripIO.h"
+
+#include "atlas/linalg/sparse/SparseMatrixToTriplets.h"
+#include "atlas/linalg/sparse/MakeEckitSparseMatrix.h"
+#include "atlas/runtime/Exception.h"
+
+
+namespace atlas {
+
+using SparseMatrixStorage = linalg::SparseMatrixStorage;
+
+
+SparseMatrixStorage ScripIO::read(const std::string& matrix_name) {
+#if ATLAS_HAVE_NETCDF == 0
+    ATLAS_THROW_EXCEPTION("Cannot read SCRIP files: Atlas not compiled with NetCDF support");
+#else
+    size_t Nr  = 0;
+    size_t Nc  = 0;
+    size_t nnz = 0;
+
+    std::vector<eckit::linalg::Index> ia;
+    std::vector<eckit::linalg::Index> ja;
+    std::vector<eckit::linalg::Scalar> a;
+
+    // read SCRIP file
+    try {
+        netCDF::NcFile f(matrix_name, netCDF::NcFile::read);
+
+        Nr  = f.getDim("dst_grid_size").getSize();
+        Nc  = f.getDim("src_grid_size").getSize();
+        nnz = f.getDim("num_links").getSize();
+
+        ATLAS_ASSERT(Nr > 0);
+        ATLAS_ASSERT(Nc > 0);
+        ATLAS_ASSERT(nnz > 0);
+
+        auto var_ia = f.getVar("dst_address");
+        ATLAS_ASSERT(var_ia.getDimCount() == 1 && var_ia.getDim(0).getSize() == nnz);
+        ia.resize(nnz);  // NOTE: not compressed
+        var_ia.getVar(ia.data());
+
+        auto var_ja = f.getVar("src_address");
+        ATLAS_ASSERT(var_ja.getDimCount() == 1 && var_ja.getDim(0).getSize() == nnz);
+        ja.resize(nnz);
+        var_ja.getVar(ja.data());
+
+        auto var_a = f.getVar("remap_matrix");
+        ATLAS_ASSERT(var_a.getDimCount() == 2 && var_a.getDim(0).getSize() == nnz && var_a.getDim(1).getSize() == 1);
+        a.resize(nnz);
+        var_a.getVar(a.data());
+    }
+    catch (netCDF::exceptions::NcException& e) {
+        ATLAS_THROW_EXCEPTION("SCRIP reading failed : " << e.what());
+    }
+
+    // matrix conversion to 0-based indexing, sorting and compression
+    {
+        std::vector<size_t> sorted(nnz);
+        for (size_t n = 0; n < nnz; ++n) {
+            sorted[n] = n;
+        }
+
+        std::sort(sorted.begin(), sorted.end(), [&](auto i, auto j) {
+            return ia[i] != ia[j] ? ia[i] < ia[j] : ja[i] != ja[j] ? ja[i] < ja[j] : i < j;
+        });
+
+        decltype(ia) ia_unsorted(Nr + 1, 0);
+        decltype(ja) ja_unsorted(nnz);
+        ia.swap(ia_unsorted);
+        ja.swap(ja_unsorted);
+
+        constexpr eckit::linalg::Index scrip_base = 1;
+        for (size_t n = 0; n < nnz; ++n) {
+            auto r = ia_unsorted[sorted[n]] - scrip_base;
+            auto c = ja_unsorted[sorted[n]] - scrip_base;
+            ATLAS_ASSERT(0 <= r && r < Nr);
+            ATLAS_ASSERT(0 <= c && c < Nc);
+
+            ia[r + 1]++;
+            ja[n] = c;
+        }
+
+        ia_unsorted.clear();
+        ja_unsorted.clear();
+
+        for (size_t r = 0; r < Nr; ++r) {
+            ia[r + 1] += ia[r];
+        }
+
+        ATLAS_ASSERT(ia.front() == 0 && ia.back() == nnz);
+
+        decltype(a) a_unsorted(nnz);
+        a.swap(a_unsorted);
+
+        for (size_t n = 0; n < nnz; ++n) {
+            a[n] = a_unsorted[sorted[n]];
+            ATLAS_ASSERT(0. <= a[n] && a[n] <= 1.);
+        }
+    }
+
+    using SparseMatrixView = atlas::linalg::SparseMatrixView<eckit::linalg::Scalar,eckit::linalg::Index>;
+    return atlas::linalg::SparseMatrixStorage{SparseMatrixView{Nr, Nc, nnz, a.data(), ja.data(), ia.data()}};
+#endif // ATLAS_HAVE_NETCDF
+}
+
+
+void ScripIO::write(const SparseMatrixStorage& matrix, const std::string& matrix_name) {
+#if ATLAS_HAVE_NETCDF == 0
+    ATLAS_THROW_EXCEPTION("Cannot write SCRIP file: Atlas not compiled with NetCDF support");
+#else
+    using scrip_index = int;
+    using scrip_value = double;
+    using scrip_size  = size_t;
+    scrip_size Nr  = matrix.rows();
+    scrip_size Nc  = matrix.cols();
+    scrip_size nnz = matrix.nnz();
+
+    std::vector<scrip_index> ia(nnz);
+    std::vector<scrip_index> ja(nnz);
+    std::vector<scrip_value> a(nnz);
+
+    sparse_matrix_to_rows_columns_values(matrix, ia, ja, a);
+    constexpr scrip_index scrip_index_base = 1;
+    for (size_t i = 0; i < matrix.nnz(); ++i) {
+        ja[i] += scrip_index_base;
+        ia[i] += scrip_index_base;
+    }
+
+    try {
+        netCDF::NcFile f(matrix_name, netCDF::NcFile::replace);
+
+        f.addDim("dst_grid_size", Nr);
+        f.addDim("src_grid_size", Nc);
+
+        std::vector<netCDF::NcDim> dims;
+        dims.push_back(f.addDim("num_links", nnz));
+        netCDF::NcVar nc_dstaddr = f.addVar("dst_address", netCDF::ncInt, dims);
+        netCDF::NcVar nc_srcaddr = f.addVar("src_address", netCDF::ncInt, dims);
+        dims.push_back(f.addDim("num_wgts", 1));
+        netCDF::NcVar nc_rmatrix = f.addVar("remap_matrix", netCDF::ncDouble, dims);
+
+        nc_dstaddr.putVar(ia.data());
+        nc_srcaddr.putVar(ja.data());
+        nc_rmatrix.putVar(a.data());
+        f.close();
+    }
+    catch (netCDF::exceptions::NcException& e) {
+        ATLAS_THROW_EXCEPTION("SCRIP writing failed : " << e.what());
+    }
+#endif
+}
+
+}

src/sandbox/interpolation/ScripIO.h

@@ -0,0 +1,34 @@
+/*
+ * (C) Copyright 2025- ECMWF.
+ *
+ * This software is licensed under the terms of the Apache Licence Version 2.0
+ * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+ *
+ * In applying this licence, ECMWF does not waive the privileges and immunities
+ * granted to it by virtue of its status as an intergovernmental organisation nor
+ * does it submit to any jurisdiction.
+ */
+
+#pragma once
+
+#if ATLAS_HAVE_NETCDF
+#include <netcdf>
+#else
+#warning "ScripIO: netCDF not found -> no SCRIP support"
+#endif
+
+#include "atlas/linalg/sparse.h"
+#include "atlas/util/Config.h"
+
+namespace atlas {
+
+class ScripIO {
+
+public:
+    explicit ScripIO(const util::Config& = util::NoConfig()) {}
+
+    static atlas::linalg::SparseMatrixStorage read(const std::string&);
+    static void write(const linalg::SparseMatrixStorage&, const std::string&);
+};
+
+}