ATLAS-353 MIR-467 handle local-to-global numbering for 2D structured …

…methods
ecmwf · Nov 15, 2024 · 5c8310f · 5c8310f
1 parent 3de9361
commit 5c8310f
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/src/atlas/interpolation/method/structured/StructuredInterpolation2D.h b/src/atlas/interpolation/method/structured/StructuredInterpolation2D.h
@@ -76,6 +76,7 @@ class StructuredInterpolation2D : public Method {
     FunctionSpace target_;
 
     bool matrix_free_;
+    bool matrix_global_;
     bool verbose_;
     double convert_units_;
     idx_t out_npts_;

diff --git a/src/atlas/interpolation/method/structured/StructuredInterpolation2D.tcc b/src/atlas/interpolation/method/structured/StructuredInterpolation2D.tcc
@@ -234,7 +234,8 @@ double StructuredInterpolation2D<Kernel>::convert_units_multiplier( const Field&
 template <typename Kernel>
 StructuredInterpolation2D<Kernel>::StructuredInterpolation2D( const Method::Config& config ) :
     Method( config ),
-    matrix_free_{false} ,
+    matrix_free_{false},
+    matrix_global_{false},
     verbose_{false} {
     config.get( "matrix_free", matrix_free_ );
     config.get( "verbose", verbose_ );
@@ -425,6 +426,24 @@ void StructuredInterpolation2D<Kernel>::setup( const FunctionSpace& source ) {
             idx_t inp_npts = source.size();
             Matrix A( out_npts_, inp_npts, triplets );
             setMatrix(A);
+
+            if (matrix_global_){
+                // adjust triplets column numbering from local to global (src_fs)
+                auto gidx_src = array::make_view<gidx_t, 1>(src_fs.global_index());
+                auto inp_npts_owned = src_fs.sizeOwned();
+
+                for (auto& t : triplets) {
+                    ATLAS_ASSERT(t.col() < inp_npts && gidx_src(t.col()) >= 1);
+                    t = {t.row(), eckit::linalg::Size(gidx_src(t.col()) - 1), t.value()};
+                    ATLAS_ASSERT(t.col() < inp_npts_owned);
+                }
+
+                inp_npts = inp_npts_owned;
+            }
+
+            // fill sparse matrix and return
+            Matrix A( out_npts, inp_npts, triplets );
+            matrix_shared_->swap( A );
         }
     }
 }