Feature: support multi-dimensional batch calculation, chaining approach [experimental]

power_grid_model_c/power_grid_model/include/power_grid_model/auxiliary/dataset.hpp

@@ -287,6 +287,9 @@ template <dataset_type_tag dataset_type_> class Dataset {
     }
     constexpr bool is_dense(Idx const i) const { return is_dense(buffers_[i]); }
     constexpr bool is_dense(Buffer const& buffer) const { return buffer.indptr.empty(); }
+    constexpr bool is_dense() const {
+        return std::ranges::all_of(buffers_, [this](Buffer const& buffer) { return is_dense(buffer); });
+    }
     constexpr bool is_sparse(std::string_view component, bool with_attribute_buffers = false) const {
         Idx const idx = find_component(component, false);
         if (idx == invalid_index) {
@@ -510,10 +513,43 @@ template <dataset_type_tag dataset_type_> class Dataset {
         return result;
     }
 
+    // get slice dataset from batch
+    Dataset get_slice_scenario(Idx begin, Idx end) const
+        requires(!is_indptr_mutable_v<dataset_type>)
+    {
+        assert(begin <= end);
+        assert(0 <= begin);
+        assert(end <= batch_size());
+        assert(is_batch());
+        assert(is_dense());
+
+        // empty slice
+        if (begin == end) {
+            Dataset result{true, 0, dataset_info_.dataset->name, *meta_data_};
+            result.add_buffer("node", 0, 0, nullptr, nullptr);
+            return result;
+        }
+
+        // start with begin
+        Dataset result = get_individual_scenario(begin);
+        Idx const batch_size = end - begin;
+        result.dataset_info_.is_batch = true;
+        result.dataset_info_.batch_size = batch_size;
+        for (auto&& [buffer, component_info] : std::views::zip(result.buffers_, result.dataset_info_.component_info)) {
+            Idx const size = component_info.elements_per_scenario * batch_size;
+            component_info.total_elements = size;
+        }
+        return result;
+    }
+
+    void set_next(Dataset const* next) { next_ = next; }
+    Dataset const* get_next() const { return next_; }
+
   private:
     MetaData const* meta_data_;
     DatasetInfo dataset_info_;
     std::vector<Buffer> buffers_;
+    Dataset const* next_{};
 
     std::span<Indptr> get_indptr_span(Indptr* indptr) const {
         return std::span{indptr, static_cast<size_t>(batch_size() + 1)};

power_grid_model_c/power_grid_model_c/include/power_grid_model_c/dataset.h

@@ -212,6 +212,16 @@ PGM_API void PGM_dataset_const_add_attribute_buffer(PGM_Handle* handle, PGM_Cons
  */
 PGM_API PGM_DatasetInfo const* PGM_dataset_const_get_info(PGM_Handle* handle, PGM_ConstDataset const* dataset);
 
+/**
+ * @brief Set next const dataset for chaining for MD batch
+ *
+ * @param handle
+ * @param dataset
+ * @param next_dataset The next dataset in the chain.
+ */
+PGM_API void PGM_dataset_const_set_next(PGM_Handle* handle, PGM_ConstDataset* dataset,
+                                        PGM_ConstDataset const* next_dataset);
+
 /**
  * @brief Get the dataset info of the instance PGM_WritableDataset.
  * @param handle

power_grid_model_c/power_grid_model_c/src/dataset.cpp

@@ -101,6 +101,11 @@ void PGM_dataset_const_add_attribute_buffer(PGM_Handle* handle, PGM_ConstDataset
         PGM_regular_error);
 }
 
+void PGM_dataset_const_set_next(PGM_Handle* /*unused*/, PGM_ConstDataset* dataset,
+                                PGM_ConstDataset const* next_dataset) {
+    dataset->set_next(next_dataset);
+}
+
 PGM_DatasetInfo const* PGM_dataset_const_get_info(PGM_Handle* /*unused*/, PGM_ConstDataset const* dataset) {
     return &dataset->get_description();
 }

power_grid_model_c/power_grid_model_c/src/model.cpp

@@ -15,6 +15,9 @@
 #include <power_grid_model/common/common.hpp>
 #include <power_grid_model/main_model.hpp>
 
+#include <memory>
+#include <numeric>
+
 namespace {
 using namespace power_grid_model;
 } // namespace
@@ -55,6 +58,7 @@
         PGM_regular_error);
 }
 
+// helper functions
 namespace {
 void check_no_experimental_features_used(MainModel const& model, MainModel::Options const& opt) {
     // optionally add experimental feature checks here
@@ -142,9 +146,11 @@
 }
 } // namespace
 
-// run calculation
-void PGM_calculate(PGM_Handle* handle, PGM_PowerGridModel* model, PGM_Options const* opt,
-                   PGM_MutableDataset const* output_dataset, PGM_ConstDataset const* batch_dataset) {
+// calculation implementation
+namespace {
+
+void calculate_impl(PGM_Handle* handle, PGM_PowerGridModel* model, PGM_Options const* opt,
+                    PGM_MutableDataset const* output_dataset, PGM_ConstDataset const* batch_dataset) {
     PGM_clear_error(handle);
     // check dataset integrity
     if ((batch_dataset != nullptr) && (!batch_dataset->is_batch() || !output_dataset->is_batch())) {
@@ -171,14 +177,96 @@
         handle->err_msg = e.what();
         handle->failed_scenarios = e.failed_scenarios();
         handle->batch_errs = e.err_msgs();
     } catch (std::exception& e) {
         handle->err_code = PGM_regular_error;
         handle->err_msg = e.what();
     } catch (...) {
         handle->err_code = PGM_regular_error;
         handle->err_msg = "Unknown error!\n";
     }
 }
 
+void merge_batch_error_msgs(PGM_Handle* handle, PGM_Handle const& local_handle, Idx scenario_offset, Idx stride_size) {
+    if (local_handle.err_code == PGM_no_error) {
+        return;
+    }
+    handle->err_code = PGM_batch_error;
+    if (local_handle.err_code == PGM_batch_error) {
+        for (auto&& [idx, err_msg] : std::views::zip(local_handle.failed_scenarios, local_handle.batch_errs)) {
+            handle->failed_scenarios.push_back(idx + scenario_offset);
+            handle->batch_errs.push_back(err_msg);
+        }
+    } else {
+        for (Idx i = 0; i < stride_size; ++i) {
+            handle->failed_scenarios.push_back(scenario_offset + i);
+            handle->batch_errs.push_back(local_handle.err_msg);
+        }
+    }
+}
+
+Idx get_batch_dimension(PGM_ConstDataset const* batch_dataset) {
+    Idx dimension = 0;
+    while (batch_dataset != nullptr) {
+        ++dimension;
+        batch_dataset = batch_dataset->get_next();
+    }
+    return dimension;
+}
+
+} // namespace
+
+// run calculation
+void PGM_calculate(PGM_Handle* handle, PGM_PowerGridModel* model, PGM_Options const* opt,
+                   PGM_MutableDataset const* output_dataset, PGM_ConstDataset const* batch_dataset) {
+    Idx const batch_dimension = get_batch_dimension(batch_dataset);
+
+    // for dimension < 2 (one-time or 1D batch), call implementation directly
+    if (batch_dimension < 2) {
+        calculate_impl(handle, model, opt, output_dataset, batch_dataset);
+        return;
+    }
+
+    // get stride size of the rest of dimensions
+    Idx const first_batch_size = batch_dataset->batch_size();
+    Idx const stride_size = [batch_dataset]() {
+        Idx size = 1;
+        PGM_ConstDataset const* current = batch_dataset->get_next();
+        while (current != nullptr) {
+            size *= current->batch_size();
+            current = current->get_next();
+        }
+        return size;
+    }();
+
+    // loop over the first dimension batche
+    for (Idx i = 0; i < first_batch_size; ++i) {
+        // a new handle
+        PGM_Handle local_handle{};
+        // create sliced datasets for the rest of dimensions
+        PGM_ConstDataset const single_update_dataset = batch_dataset->get_individual_scenario(i);
+        PGM_MutableDataset const sliced_output_dataset =
+            output_dataset->get_slice_scenario(i * stride_size, (i + 1) * stride_size);
+
+        // create a model copy
+        std::unique_ptr<PGM_PowerGridModel> const local_model{PGM_copy_model(&local_handle, model)};
+        if (local_handle.err_code != PGM_no_error) {
+            merge_batch_error_msgs(handle, local_handle, i * stride_size, stride_size);
+            continue;
+        }
+
+        // apply the update
+        PGM_update_model(&local_handle, local_model.get(), &single_update_dataset);
+        if (local_handle.err_code != PGM_no_error) {
+            merge_batch_error_msgs(handle, local_handle, i * stride_size, stride_size);
+            continue;
+        }
+
+        // recursive call
+        PGM_calculate(&local_handle, local_model.get(), opt, &sliced_output_dataset, batch_dataset->get_next());
+        // merge errors
+        merge_batch_error_msgs(handle, local_handle, i * stride_size, stride_size);
+    }
+}
+
 // destroy model
 void PGM_destroy_model(PGM_PowerGridModel* model) { delete model; }

power_grid_model_c/power_grid_model_cpp/include/power_grid_model_cpp/dataset.hpp

@@ -198,6 +198,10 @@ class DatasetConst {
                           data.get());
     }
 
+    void set_next(DatasetConst const& next_dataset) {
+        handle_.call_with(PGM_dataset_const_set_next, get(), next_dataset.get());
+    }
+
     DatasetInfo const& get_info() const { return info_; }
 
   private:

src/power_grid_model/_core/power_grid_core.py

@@ -481,6 +481,10 @@ def destroy_dataset_const(self, dataset: ConstDatasetPtr) -> None:  # type: igno
     def dataset_const_get_info(self, dataset: ConstDatasetPtr) -> DatasetInfoPtr:  # type: ignore[empty-body]
         pass  # pragma: no cover
 
+    @make_c_binding
+    def dataset_const_set_next(self, dataset: ConstDatasetPtr, next_dataset: ConstDatasetPtr) -> None:  # type: ignore[empty-body]
+        pass  # pragma: no cover
+
     @make_c_binding
     def dataset_writable_get_info(self, dataset: WritableDatasetPtr) -> DatasetInfoPtr:  # type: ignore[empty-body]
         pass  # pragma: no cover

src/power_grid_model/_core/power_grid_dataset.py

@@ -350,6 +350,16 @@ def get_info(self) -> CDatasetInfo:
         """
         return CDatasetInfo(pgc.dataset_const_get_info(self._const_dataset))
 
+    def set_next(self, next_dataset: "CConstDataset") -> None:
+        """
+        Set the next dataset in the linked list.
+
+        Args:
+            next_dataset: The next dataset to set.
+        """
+        pgc.dataset_const_set_next(self._const_dataset, next_dataset._const_dataset)
+        assert_no_error()
+
     def __del__(self):
         pgc.destroy_dataset_const(self._const_dataset)