Merge pull request #130 from alliander-opensource/feature/id-lookup

Add get indexer function to lookup from IDs to index sequence

Author: TonyXiang8787

docs/python-api-reference.md

@@ -137,6 +137,22 @@ class PowerGridModel:
         """
         pass
 
+    def get_indexer(self, 
+                    component_type: str, 
+                    ids: np.ndarray):
+        """
+        Get array of indexers given array of ids for component type
+
+        Args:
+            component_type: type of component
+            ids: array of ids
+        
+        Returns:
+            array of inderxers, same shape as input array ids
+
+        """
+        pass
+
     def copy(self) -> 'PowerGridModel':
         """
 

include/power_grid_model/main_model.hpp

@@ -78,6 +78,7 @@ class MainModelImpl<ExtraRetrievableTypes<ExtraRetrievableType...>, ComponentLis
                                 DataPointer<false> const& data_ptr, Idx position);
     using CheckUpdateFunc = bool (*)(ConstDataPointer const& component_update);
     using GetSeqIdxFunc = std::vector<Idx2D> (*)(MainModelImpl const& x, ConstDataPointer const& component_update);
+    using GetIndexerFunc = void (*)(MainModelImpl const& x, ID const* id_begin, Idx size, Idx* indexer_begin);
 
    public:
     // constructor with data
@@ -376,6 +377,25 @@ class MainModelImpl<ExtraRetrievableTypes<ExtraRetrievableType...>, ComponentLis
         comp_coup_.reset();
     }
 
+    /*
+    the the sequence indexer given an input array of ID's for a given component type
+    */
+    void get_indexer(std::string const& component_type, ID const* id_begin, Idx size, Idx* indexer_begin) {
+        // static function array
+        static constexpr std::array<GetIndexerFunc, n_types> get_indexer_func{
+            [](MainModelImpl const& model, ID const* id_begin, Idx size, Idx* indexer_begin) {
+                std::transform(id_begin, id_begin + size, indexer_begin, [&model](ID id) {
+                    return model.components_.template get_idx_by_id<ComponentType>(id).pos;
+                });
+            }...};
+        // search component type name
+        for (ComponentEntry const& entry : AllComponents::component_index_map) {
+            if (entry.name == component_type) {
+                return get_indexer_func[entry.index](*this, id_begin, size, indexer_begin);
+            }
+        }
+    }
+
    private:
     template <bool sym, typename InputType, std::vector<InputType> (MainModelImpl::*PrepareInputFn)(),
               MathOutput<sym> (MathSolver<sym>::*SolveFn)(InputType const&, double, Idx, CalculationInfo&,
@@ -417,7 +437,7 @@ class MainModelImpl<ExtraRetrievableTypes<ExtraRetrievableType...>, ComponentLis
     std::map<std::string, std::vector<Idx2D>> get_sequence_idx_map(ConstDataset const& update_data) const {
         // function pointer array to get cached idx
         static constexpr std::array<GetSeqIdxFunc, n_types> get_seq_idx{
-            [](MainModelImpl const& x, ConstDataPointer const& component_update) -> std::vector<Idx2D> {
+            [](MainModelImpl const& model, ConstDataPointer const& component_update) -> std::vector<Idx2D> {
                 using UpdateType = typename ComponentType::UpdateType;
                 // no batch
                 if (component_update.batch_size() < 1) {
@@ -427,8 +447,8 @@ class MainModelImpl<ExtraRetrievableTypes<ExtraRetrievableType...>, ComponentLis
                 auto const [it_begin, it_end] = component_update.template get_iterators<UpdateType>(0);
                 // vector
                 std::vector<Idx2D> seq_idx(std::distance(it_begin, it_end));
-                std::transform(it_begin, it_end, seq_idx.begin(), [x](UpdateType const& update) {
-                    return x.components_.template get_idx_by_id<ComponentType>(update.id);
+                std::transform(it_begin, it_end, seq_idx.begin(), [&model](UpdateType const& update) {
+                    return model.components_.template get_idx_by_id<ComponentType>(update.id);
                 });
                 return seq_idx;
             }...};

src/power_grid_model/_power_grid_core.pyx

@@ -15,13 +15,19 @@ from .enum import CalculationMethod
 
 cimport numpy as cnp
 from cython.operator cimport dereference as deref
-from libc.stdint cimport int8_t, int32_t
+from libc.stdint cimport int8_t
 from libcpp cimport bool
 from libcpp.map cimport map
 from libcpp.string cimport string
 from libcpp.vector cimport vector
 
-VALIDATOR_MSG = "Try validate_input_data() or validate_batch_data() to validate your data."
+# idx and id types
+from libc.stdint cimport int32_t as idx_t # isort: skip
+cdef np_idx_t = np.int32
+from libc.stdint cimport int32_t as id_t # isort: skip
+cdef np_id_t = np.int32
+
+cdef VALIDATOR_MSG = "Try validate_input_data() or validate_batch_data() to validate your data."
 
 cdef extern from "power_grid_model/auxiliary/meta_data_gen.hpp" namespace "power_grid_model::meta_data":
     cppclass DataAttribute:
@@ -110,10 +116,10 @@ cdef _generate_component_meta_data(MetaData & cpp_component_meta_data):
 cdef extern from "power_grid_model/auxiliary/dataset.hpp" namespace "power_grid_model":
     cppclass MutableDataPointer:
         MutableDataPointer()
-        MutableDataPointer(void * ptr, const int32_t * indptr, int32_t size)
+        MutableDataPointer(void * ptr, const idx_t * indptr, idx_t size)
     cppclass ConstDataPointer:
         ConstDataPointer()
-        ConstDataPointer(const void * ptr, const int32_t * indptr, int32_t size)
+        ConstDataPointer(const void * ptr, const idx_t * indptr, idx_t size)
 
 cdef extern from "power_grid_model/main_model.hpp" namespace "power_grid_model":
     cppclass CalculationMethodCPP "::power_grid_model::CalculationMethod":
@@ -122,43 +128,48 @@ cdef extern from "power_grid_model/main_model.hpp" namespace "power_grid_model":
         bool independent
         bool cache_topology
     cppclass MainModel:
-        map[string, int32_t] all_component_count()
+        map[string, idx_t] all_component_count()
         BatchParameter calculate_sym_power_flow "calculate_power_flow<true>"(
             double error_tolerance,
-            int32_t max_iterations,
+            idx_t max_iterations,
             CalculationMethodCPP calculation_method,
             const map[string, MutableDataPointer] & result_data,
             const map[string, ConstDataPointer] & update_data,
-            int32_t threading
+            idx_t threading
         ) except+
         BatchParameter calculate_asym_power_flow "calculate_power_flow<false>"(
             double error_tolerance,
-            int32_t max_iterations,
+            idx_t max_iterations,
             CalculationMethodCPP calculation_method,
             const map[string, MutableDataPointer] & result_data,
             const map[string, ConstDataPointer] & update_data,
-            int32_t threading
+            idx_t threading
         ) except+
         BatchParameter calculate_sym_state_estimation "calculate_state_estimation<true>"(
             double error_tolerance,
-            int32_t max_iterations,
+            idx_t max_iterations,
             CalculationMethodCPP calculation_method,
             const map[string, MutableDataPointer] & result_data,
             const map[string, ConstDataPointer] & update_data,
-            int32_t threading
+            idx_t threading
         ) except+
         BatchParameter calculate_asym_state_estimation "calculate_state_estimation<false>"(
             double error_tolerance,
-            int32_t max_iterations,
+            idx_t max_iterations,
             CalculationMethodCPP calculation_method,
             const map[string, MutableDataPointer] & result_data,
             const map[string, ConstDataPointer] & update_data,
-            int32_t threading
+            idx_t threading
         ) except+
         void update_component(
             const map[string, ConstDataPointer] & update_data,
-            int32_t pos
+            idx_t pos
         ) except+
+        void get_indexer(
+            const string& component_type, 
+            const id_t* id_begin, 
+            idx_t size, 
+            idx_t* indexer_begin) except+
 
 cdef extern from "<optional>":
     cppclass OptionalMainModel "::std::optional<::power_grid_model::MainModel>":
@@ -169,7 +180,7 @@ cdef extern from "<optional>":
         MainModel & emplace(
             double system_frequency,
             const map[string, ConstDataPointer] & input_data,
-            int32_t pos) except+
+            idx_t pos) except+
 
 # internally used meta data, to prevent modification
 cdef _power_grid_meta_data = _generate_meta_data()
@@ -183,7 +194,7 @@ cdef map[string, ConstDataPointer] generate_const_ptr_map(data: Dict[str, Dict[s
         data_arr = v['data']
         indptr_arr = v['indptr']
         result[k.encode()] = ConstDataPointer(
-            cnp.PyArray_DATA(data_arr), < const int32_t*>cnp.PyArray_DATA(indptr_arr),
+            cnp.PyArray_DATA(data_arr), < const idx_t*>cnp.PyArray_DATA(indptr_arr),
             v['batch_size'])
     return result
 
@@ -195,7 +206,7 @@ cdef map[string, MutableDataPointer] generate_ptr_map(data: Dict[str, Dict[str,
         data_arr = v['data']
         indptr_arr = v['indptr']
         result[k.encode()] = MutableDataPointer(
-            cnp.PyArray_DATA(data_arr), < const int32_t * > cnp.PyArray_DATA(indptr_arr),
+            cnp.PyArray_DATA(data_arr), < const idx_t * > cnp.PyArray_DATA(indptr_arr),
             v['batch_size'])
     return result
 
@@ -234,10 +245,10 @@ cdef _prepare_cpp_array(data_type: str,
             data = v
             ndim = v.ndim
             if ndim == 1:
-                indptr = np.array([0, v.size], dtype=np.int32)
+                indptr = np.array([0, v.size], dtype=np_idx_t)
                 batch_size = 1
             elif ndim == 2:  # (n_batch, n_component)
-                indptr = np.arange(v.shape[0] + 1, dtype=np.int32) * v.shape[1]
+                indptr = np.arange(v.shape[0] + 1, dtype=np_idx_t) * v.shape[1]
                 batch_size = v.shape[0]
             else:
                 raise ValueError(f"Array can only be 1D or 2D. {VALIDATOR_MSG}")
@@ -257,7 +268,7 @@ cdef _prepare_cpp_array(data_type: str,
                 raise ValueError(f"indptr should be increasing. {VALIDATOR_MSG}")
         # convert array
         data = np.ascontiguousarray(data, dtype=schema[component_name]['dtype'])
-        indptr = np.ascontiguousarray(indptr, dtype=np.int32)
+        indptr = np.ascontiguousarray(indptr, dtype=np_idx_t)
         return_dict[component_name] = {
             'data': data,
             'indptr': indptr,
@@ -302,6 +313,29 @@ cdef class PowerGridModel:
         self.independent = False
         self.cache_topology = False
 
+    def get_indexer(self, 
+                    component_type: str, 
+                    ids: np.ndarray):
+        """
+        Get array of indexers given array of ids for component type
+
+        Args:
+            component_type: type of component
+            ids: array of ids
+        
+        Returns:
+            array of inderxers, same shape as input array ids
+
+        """
+        cdef cnp.ndarray ids_c = np.ascontiguousarray(ids, dtype=np_id_t)
+        cdef cnp.ndarray indexer = np.empty_like(ids_c, dtype=np_idx_t, order='C')
+        cdef const id_t* id_begin = <const id_t*> cnp.PyArray_DATA(ids_c)
+        cdef idx_t* indexer_begin = <idx_t*> cnp.PyArray_DATA(indexer)
+        cdef idx_t size = ids.size
+        # call c function
+        self._get_model().get_indexer(component_type.encode(), id_begin, size, indexer_begin)
+        return indexer
+
     def copy(self) -> PowerGridModel:
         """
 
@@ -333,10 +367,10 @@ cdef class PowerGridModel:
                    calculation_type,
                    bool symmetric,
                    double error_tolerance,
-                   int32_t max_iterations,
+                   idx_t max_iterations,
                    calculation_method: Union[CalculationMethod, str],
                    update_data: Optional[Dict[str, Union[np.ndarray, Dict[str, np.ndarray]]]],
-                   int32_t threading
+                   idx_t threading
                    ):
         """
         Core calculation routine
@@ -453,10 +487,10 @@ cdef class PowerGridModel:
     def calculate_power_flow(self, *,
                              bool symmetric=True,
                              double error_tolerance=1e-8,
-                             int32_t max_iterations=20,
+                             idx_t max_iterations=20,
                              calculation_method: Union[CalculationMethod, str] = CalculationMethod.newton_raphson,
                              update_data: Optional[Dict[str, Union[np.ndarray, Dict[str, np.ndarray]]]] = None,
-                             int32_t threading=-1
+                             idx_t threading=-1
                              ) -> Dict[str, np.ndarray]:
         """
         Calculate power flow once with the current model attributes.
@@ -519,10 +553,10 @@ cdef class PowerGridModel:
     def calculate_state_estimation(self, *,
                                    bool symmetric=True,
                                    double error_tolerance=1e-8,
-                                   int32_t max_iterations=20,
+                                   idx_t max_iterations=20,
                                    calculation_method: Union[CalculationMethod, str] = CalculationMethod.iterative_linear,
                                    update_data: Optional[Dict[str, Union[np.ndarray, Dict[str, np.ndarray]]]] = None,
-                                   int32_t threading=-1
+                                   idx_t threading=-1
                                    ) -> Dict[str, np.ndarray]:
         """
         Calculate state estimation once with the current model attributes.
@@ -592,7 +626,7 @@ cdef class PowerGridModel:
                 value: integer count of elements of this type
         """
         all_component_count = {}
-        cdef map[string, int32_t] cpp_count = self._get_model().all_component_count()
+        cdef map[string, idx_t] cpp_count = self._get_model().all_component_count()
         for map_entry in cpp_count:
             all_component_count[map_entry.first.decode()] = map_entry.second
         return all_component_count

tests/cpp_unit_tests/test_main_model.cpp

@@ -104,6 +104,14 @@ TEST_CASE("Test main model") {
     main_model.add_component<AsymVoltageSensor>(asym_voltage_sensor_input);
     main_model.set_construction_complete();
 
+    SUBCASE("Test get indexer") {
+        IdxVector const node_id{2, 1, 3, 2};
+        IdxVector const expected_indexer{1, 0, 2, 1};
+        IdxVector indexer(4);
+        main_model.get_indexer("node", node_id.data(), 4, indexer.data());
+        CHECK(indexer == expected_indexer);
+    }
+
     SUBCASE("Test duplicated id") {
         MainModel main_model2{50.0};
         node_input[1].id = 1;

tests/unit/test_0Z_model_validation.py

@@ -5,6 +5,7 @@
 from copy import copy
 from pathlib import Path
 
+import numpy as np
 import pytest
 
 from power_grid_model import PowerGridModel
@@ -36,6 +37,14 @@ def test_single_validation(
     reference_result = case_data["output"]
     compare_result(result, reference_result, rtol, atol)
 
+    # test get indexer
+    for component_name, input_array in case_data["input"].items():
+        ids_array = input_array["id"].copy()
+        np.random.shuffle(ids_array)
+        indexer_array = model.get_indexer(component_name, ids_array)
+        # check
+        assert np.all(input_array["id"][indexer_array] == ids_array)
+
     # export data if needed
     if EXPORT_OUTPUT:
         save_json_data(f"{case_id}.json", result)