eckit::geo iterator

ecmwf · Nov 20, 2024 · 5621ee9 · 5621ee9
1 parent d155636
commit 5621ee9
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 10 deletions.
diff --git a/src/grib_iterator_factory.cc b/src/grib_iterator_factory.cc
@@ -12,6 +12,9 @@
 
 #ifdef HAVE_ECKIT_GEO
     #include <memory>
+
+    #include "eckit/geo/Grid.h"
+
     #include "eccodes/geo/GribSpec.h"
 
     // eccodes macros conflict with eckit
@@ -76,19 +79,80 @@ static const struct table_entry table[] = {
 class GeoIterator : public eccodes::geo_iterator::Iterator
 {
 public:
-    GeoIterator(grib_handle* h) :
-        spec_(new eccodes::geo::GribSpec(h)) {}
+    explicit GeoIterator(grib_handle* h, unsigned long flags) :
+        spec_(new eccodes::geo::GribSpec(h)), grid_(eckit::geo::GridFactory::build(*spec_)), iter_(grid_->cbegin().release()), end_(grid_->cend().release())
+    {
+        h_          = h;
+        class_name_ = "geo_iterator";
+        flags_      = flags;
+        Assert(h_ != nullptr);
+
+        CODES_CHECK(codes_get_size(h_, "values", &nv_), "");
+        Assert(nv_ > 0);
+
+        data_ = (flags_ & GRIB_GEOITERATOR_NO_VALUES) ? nullptr : static_cast<double*>(grib_context_malloc(h_->context, nv_ * sizeof(double)));
+        Assert(data_ != nullptr);
+
+        auto size = nv_;
+        CODES_CHECK(codes_get_double_array(h_, "values", data_, &size), "");
+        Assert(nv_ == size);
+    }
 
 private:
-    std::unique_ptr<eckit::geo::Spec> spec_;
+    std::unique_ptr<const eckit::geo::Spec> spec_;
+    std::unique_ptr<const eckit::geo::Grid> grid_;
+    mutable eckit::geo::Grid::Iterator iter_;
+    eckit::geo::Grid::Iterator end_;
 
     int init(grib_handle*, grib_arguments*) override { NOTIMP; }
-    int next(double*, double*, double*) const override { NOTIMP; }
-    int previous(double*, double*, double*) const override { NOTIMP; }
-    int reset() override { NOTIMP; }
-    int destroy() override { NOTIMP; }
-    bool has_next() const override { NOTIMP; }
-    Iterator* create() const override { NOTIMP; }
+
+    int next(double* lat, double* lon, double* val) const override
+    {
+        if (iter_ == end_) {
+            return 0;  // (false)
+        }
+
+        const auto p  = *iter_;
+        const auto& q = std::get<eckit::geo::PointLonLat>(p);
+
+        *lat = q.lat;
+        *lon = q.lon;
+        if (val != nullptr && data_ != nullptr) {
+            *val = data_[iter_->index()];
+        }
+
+        ++iter_;
+        return 1;  // (true)
+    }
+
+    int previous(double*, double*, double*) const override
+    {
+        return GRIB_NOT_IMPLEMENTED;
+    }
+
+    int reset() override
+    {
+        iter_.reset(grid_->cbegin().release());
+        return GRIB_SUCCESS;
+    }
+
+    int destroy() override
+    {
+        if (data_ != nullptr) {
+            grib_context_free(h_->context, data_);
+        }
+        return Iterator::destroy();
+    }
+
+    bool has_next() const override
+    {
+        return iter_ != end_;
+    }
+
+    Iterator* create() const override
+    {
+        return new GeoIterator{ h_, flags_ };
+    }
 };
 #endif
 
@@ -97,7 +161,7 @@ eccodes::geo_iterator::Iterator* grib_iterator_factory(grib_handle* h, grib_argu
 #ifdef HAVE_ECKIT_GEO
     static const auto* eckit_geo = codes_getenv("ECCODES_ECKIT_GEO");
     if (eckit_geo != NULL && strcmp(eckit_geo, "1") == 0) {
-        return new GeoIterator(h);
+        return new GeoIterator(h, flags);
     }
 #endif
 

diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt
@@ -24,6 +24,9 @@ ecbuild_add_library( TARGET    ecc_tools
                      NOINSTALL
                      SOURCES   ${ecc_tools_sources}
                      PRIVATE_LIBS      eccodes )
+if( eccodes_HAVE_ECKIT_GEO )
+    target_link_libraries( ecc_tools PRIVATE eckit )
+endif()
 
 # tools binaries
 list( APPEND ecc_tools_binaries

diff --git a/tools/grib_tools.cc b/tools/grib_tools.cc
@@ -9,9 +9,14 @@
  */
 
 #include "grib_tools.h"
+
 #include <stdlib.h>
 #include <string>
 
+#ifdef HAVE_ECKIT_GEO
+    #include "eckit/runtime/Main.h"
+#endif
+
 #if HAVE_LIBJASPER
 /* Remove compiler warnings re macros being redefined */
 #undef PACKAGE_BUGREPORT
@@ -141,6 +146,10 @@ static grib_handle* grib_handle_new_from_file_x(grib_context* c, FILE* f, int mo
 
 int grib_tool(int argc, char** argv)
 {
+#ifdef HAVE_ECKIT_GEO
+    eckit::Main::initialise(argc, argv);
+#endif
+
     int ret                = 0;
     int i = 0;
     grib_context* c        = grib_context_get_default();