power-grid-model 1.12.39__py3-none-macosx_13_0_x86_64.whl → 1.12.104__py3-none-macosx_13_0_x86_64.whl

power_grid_model/_core/data_handling.py

@@ -6,16 +6,26 @@
 Data handling
 """
 
+from typing import Literal, overload
+
 import numpy as np
 
-from power_grid_model._core.data_types import Dataset, SingleDataset
-from power_grid_model._core.dataset_definitions import ComponentType, DatasetType
+from power_grid_model._core.data_types import (
+    BatchDataset,
+    Dataset,
+    DenseBatchArray,
+    SingleArray,
+    SingleColumnarData,
+    SingleDataset,
+)
+from power_grid_model._core.dataset_definitions import ComponentType, ComponentTypeVar, DatasetType
 from power_grid_model._core.enum import CalculationType, ComponentAttributeFilterOptions
 from power_grid_model._core.errors import PowerGridUnreachableHitError
 from power_grid_model._core.power_grid_dataset import CConstDataset, CMutableDataset
 from power_grid_model._core.power_grid_meta import initialize_array, power_grid_meta_data
-from power_grid_model._core.typing import ComponentAttributeMapping
+from power_grid_model._core.typing import ComponentAttributeMapping, ComponentAttributeMappingDict
 from power_grid_model._core.utils import process_data_filter
+from power_grid_model.data_types import DenseBatchColumnarData
 
 
 def get_output_type(*, calculation_type: CalculationType, symmetric: bool) -> DatasetType:
@@ -85,6 +95,54 @@ def prepare_output_view(output_data: Dataset, output_type: DatasetType) -> CMuta
     return CMutableDataset(output_data, dataset_type=output_type)
 
 
+@overload
+def create_output_data(
+    output_component_types: None | set[ComponentTypeVar] | list[ComponentTypeVar],
+    output_type: DatasetType,
+    all_component_count: dict[ComponentType, int],
+    is_batch: Literal[False],
+    batch_size: int,
+) -> dict[ComponentType, SingleArray]: ...
+@overload
+def create_output_data(
+    output_component_types: None | set[ComponentTypeVar] | list[ComponentTypeVar],
+    output_type: DatasetType,
+    all_component_count: dict[ComponentType, int],
+    is_batch: Literal[True],
+    batch_size: int,
+) -> dict[ComponentType, DenseBatchArray]: ...
+@overload
+def create_output_data(
+    output_component_types: ComponentAttributeFilterOptions,
+    output_type: DatasetType,
+    all_component_count: dict[ComponentType, int],
+    is_batch: Literal[False],
+    batch_size: int,
+) -> dict[ComponentType, SingleColumnarData]: ...
+@overload
+def create_output_data(
+    output_component_types: ComponentAttributeFilterOptions,
+    output_type: DatasetType,
+    all_component_count: dict[ComponentType, int],
+    is_batch: Literal[True],
+    batch_size: int,
+) -> dict[ComponentType, DenseBatchColumnarData]: ...
+@overload
+def create_output_data(
+    output_component_types: ComponentAttributeMappingDict,
+    output_type: DatasetType,
+    all_component_count: dict[ComponentType, int],
+    is_batch: Literal[False],
+    batch_size: int,
+) -> SingleDataset: ...
+@overload
+def create_output_data(
+    output_component_types: ComponentAttributeMappingDict,
+    output_type: DatasetType,
+    all_component_count: dict[ComponentType, int],
+    is_batch: Literal[True],
+    batch_size: int,
+) -> BatchDataset: ...
 def create_output_data(
     output_component_types: ComponentAttributeMapping,
     output_type: DatasetType,
@@ -96,7 +154,7 @@ def create_output_data(
     Create the output dataset based on component and batch size from the model; and output attributes requested by user.
 
     Args:
-        output_component_types:
+        output_component_types (ComponentAttributeMapping):
             the output components the user seeks to extract
         output_type:
             the type of output that the user will see (as per the calculation options)
@@ -118,11 +176,7 @@ def create_output_data(
     result_dict: Dataset = {}
 
     for name, count in all_component_count.items():
-        # shape
-        if is_batch:
-            shape: tuple[int] | tuple[int, int] = (batch_size, count)
-        else:
-            shape = (count,)
+        shape: tuple[int, int] | int = (batch_size, count) if is_batch else count
 
         requested_component = processed_output_types[name]
         dtype = power_grid_meta_data[output_type][name].dtype
@@ -134,8 +188,8 @@ def create_output_data(
             ComponentAttributeFilterOptions.everything,
             ComponentAttributeFilterOptions.relevant,
         ]:
-            result_dict[name] = {attr: np.empty(shape, dtype=dtype[attr]) for attr in dtype.names}
+            result_dict[name] = {attr: np.empty(shape=shape, dtype=dtype[attr]) for attr in dtype.names}
         elif isinstance(requested_component, list | set):
-            result_dict[name] = {attr: np.empty(shape, dtype=dtype[attr]) for attr in requested_component}
+            result_dict[name] = {attr: np.empty(shape=shape, dtype=dtype[attr]) for attr in requested_component}
 
     return result_dict

power_grid_model/_core/data_types.py

@@ -13,7 +13,7 @@ from typing import TypeAlias, TypedDict, TypeVar
 
 import numpy as np
 
-from power_grid_model._core.dataset_definitions import ComponentTypeVar
+from power_grid_model._core.dataset_definitions import ComponentType, ComponentTypeVar
 
 SingleArray: TypeAlias = np.ndarray
 
@@ -28,7 +28,8 @@ SingleColumnarData = dict[AttributeType, SingleColumn]
 _SingleComponentData = TypeVar("_SingleComponentData", SingleArray, SingleColumnarData)  # deduction helper
 SingleComponentData = SingleArray | SingleColumnarData
 
-
+SingleRowBasedDataset = dict[ComponentTypeVar, SingleArray]
+SingleColumnarDataset = dict[ComponentTypeVar, SingleColumnarData]
 SingleDataset = dict[ComponentTypeVar, _SingleComponentData]
 
 BatchList = list[SingleDataset]
@@ -113,6 +114,16 @@ Dataset = dict[ComponentTypeVar, _ComponentData]
 
 DenseBatchData = DenseBatchArray | DenseBatchColumnarData
 
+# overloads that only match on latest PGM type
+SingleRowBasedOutputDataset = dict[ComponentType, SingleArray]
+SingleColumnarOutputDataset = dict[ComponentType, SingleColumnarData]
+SingleOutputDataset = dict[ComponentType, SingleComponentData]
+DenseBatchRowBasedOutputDataset = dict[ComponentType, DenseBatchArray]
+DenseBatchColumnarOutputDataset = dict[ComponentType, DenseBatchColumnarData]
+DenseBatchOutputDataset = dict[ComponentType, DenseBatchData]
+OutputDataset = dict[ComponentType, ComponentData]
+
+
 NominalValue = int
 
 RealValue = float

power_grid_model/_core/error_handling.py

@@ -40,7 +40,7 @@ from power_grid_model._core.errors import (
     TapSearchStrategyIncompatibleError,
 )
 from power_grid_model._core.index_integer import IdxNp
-from power_grid_model._core.power_grid_core import power_grid_core as pgc
+from power_grid_model._core.power_grid_core import get_power_grid_core as get_pgc
 
 VALIDATOR_MSG = "\nTry validate_input_data() or validate_batch_data() to validate your data.\n"
 # error codes
@@ -133,19 +133,19 @@ def find_error(batch_size: int = 1, decode_error: bool = True) -> RuntimeError |
     Returns: error object, can be none
 
     """
-    error_code: int = pgc.error_code()
+    error_code: int = get_pgc().error_code()
     if error_code == PGM_NO_ERROR:
         return None
     if error_code == PGM_REGULAR_ERROR:
-        error_message = pgc.error_message()
+        error_message = get_pgc().error_message()
         error_message += VALIDATOR_MSG
         return _interpret_error(error_message, decode_error=decode_error)
     if error_code == PGM_BATCH_ERROR:
         error_message = "There are errors in the batch calculation." + VALIDATOR_MSG
         error = PowerGridBatchError(error_message)
-        n_fails = pgc.n_failed_scenarios()
-        failed_idxptr = pgc.failed_scenarios()
-        failed_msgptr = pgc.batch_errors()
+        n_fails = get_pgc().n_failed_scenarios()
+        failed_idxptr = get_pgc().failed_scenarios()
+        failed_msgptr = get_pgc().batch_errors()
         error.failed_scenarios = np.ctypeslib.as_array(failed_idxptr, shape=(n_fails,)).copy()
         error.error_messages = [failed_msgptr[i].decode() for i in range(n_fails)]  # type: ignore
         error.errors = [_interpret_error(message, decode_error=decode_error) for message in error.error_messages]
@@ -155,7 +155,7 @@ def find_error(batch_size: int = 1, decode_error: bool = True) -> RuntimeError |
         error.succeeded_scenarios = all_scenarios[mask]
         return error
     if error_code == PGM_SERIALIZATION_ERROR:
-        return PowerGridSerializationError(pgc.error_message())
+        return PowerGridSerializationError(get_pgc().error_message())
     return RuntimeError("Unknown error!")
 
 

power_grid_model/_core/options.py

@@ -9,7 +9,7 @@ Option class
 from collections.abc import Callable
 from typing import Any
 
-from power_grid_model._core.power_grid_core import OptionsPtr, power_grid_core as pgc
+from power_grid_model._core.power_grid_core import OptionsPtr, get_power_grid_core as get_pgc
 
 
 class OptionSetter:
@@ -36,15 +36,15 @@ class Options:
 
     _opt: OptionsPtr
     # option setter
-    calculation_type = OptionSetter(pgc.set_calculation_type)
-    calculation_method = OptionSetter(pgc.set_calculation_method)
-    symmetric = OptionSetter(pgc.set_symmetric)
-    error_tolerance = OptionSetter(pgc.set_err_tol)
-    max_iterations = OptionSetter(pgc.set_max_iter)
-    threading = OptionSetter(pgc.set_threading)
-    tap_changing_strategy = OptionSetter(pgc.set_tap_changing_strategy)
-    short_circuit_voltage_scaling = OptionSetter(pgc.set_short_circuit_voltage_scaling)
-    experimental_features = OptionSetter(pgc.set_experimental_features)
+    calculation_type = OptionSetter(get_pgc().set_calculation_type)
+    calculation_method = OptionSetter(get_pgc().set_calculation_method)
+    symmetric = OptionSetter(get_pgc().set_symmetric)
+    error_tolerance = OptionSetter(get_pgc().set_err_tol)
+    max_iterations = OptionSetter(get_pgc().set_max_iter)
+    threading = OptionSetter(get_pgc().set_threading)
+    tap_changing_strategy = OptionSetter(get_pgc().set_tap_changing_strategy)
+    short_circuit_voltage_scaling = OptionSetter(get_pgc().set_short_circuit_voltage_scaling)
+    experimental_features = OptionSetter(get_pgc().set_experimental_features)
 
     @property
     def opt(self) -> OptionsPtr:
@@ -57,11 +57,11 @@ class Options:
 
     def __new__(cls, *args, **kwargs):
         instance = super().__new__(cls, *args, **kwargs)
-        instance._opt = pgc.create_options()
+        instance._opt = get_pgc().create_options()
         return instance
 
     def __del__(self):
-        pgc.destroy_options(self._opt)
+        get_pgc().destroy_options(self._opt)
 
     # not copyable
     def __copy__(self):

power_grid_model/_core/power_grid_core.py

@@ -6,6 +6,7 @@
 Loader for the dynamic library
 """
 
+import threading
 from collections.abc import Callable
 from ctypes import CDLL, POINTER, c_char, c_char_p, c_double, c_size_t, c_void_p
 from inspect import signature
@@ -14,6 +15,10 @@ from itertools import chain
 from power_grid_model._core.index_integer import IdC, IdxC
 from power_grid_model._core.power_grid_model_c.get_pgm_dll_path import get_pgm_dll_path
 
+# threading local
+_thread_local_data = threading.local()
+
+
 # integer index
 IdxPtr = POINTER(IdxC)
 """Pointer to index."""
@@ -195,14 +200,12 @@ class PowerGridCore:
     """
 
     _handle: HandlePtr
-    _instance: "PowerGridCore | None" = None
 
     # singleton of power grid core
     def __new__(cls, *args, **kwargs):
-        if cls._instance is None:
-            cls._instance = super().__new__(cls, *args, **kwargs)
-            cls._instance._handle = _CDLL.PGM_create_handle()
-        return cls._instance
+        instance = super().__new__(cls, *args, **kwargs)
+        instance._handle = _CDLL.PGM_create_handle()
+        return instance
 
     def __del__(self):
         _CDLL.PGM_destroy_handle(self._handle)
@@ -481,6 +484,12 @@ class PowerGridCore:
     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_cartesian_product_dimension(
+        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
@@ -560,4 +569,9 @@ class PowerGridCore:
 
 
 # make one instance
-power_grid_core = PowerGridCore()
+def get_power_grid_core() -> PowerGridCore:
+    try:
+        return _thread_local_data.power_grid_core
+    except AttributeError:
+        _thread_local_data.power_grid_core = PowerGridCore()
+        return _thread_local_data.power_grid_core

power_grid_model/_core/power_grid_dataset.py

@@ -34,10 +34,10 @@ from power_grid_model._core.power_grid_core import (
     DatasetInfoPtr,
     MutableDatasetPtr,
     WritableDatasetPtr,
-    power_grid_core as pgc,
+    get_power_grid_core as get_pgc,
 )
 from power_grid_model._core.power_grid_meta import ComponentMetaData, DatasetMetaData, power_grid_meta_data
-from power_grid_model._core.typing import ComponentAttributeMapping, _ComponentAttributeMappingDict
+from power_grid_model._core.typing import ComponentAttributeMapping, ComponentAttributeMappingDict
 from power_grid_model._core.utils import (
     get_dataset_type,
     is_columnar,
@@ -62,7 +62,7 @@ class CDatasetInfo:
         Returns:
             The name of the dataset type
         """
-        return pgc.dataset_info_name(self._info)
+        return get_pgc().dataset_info_name(self._info)
 
     def dataset_type(self):
         """
@@ -80,7 +80,7 @@ class CDatasetInfo:
         Returns:
             Whether the dataset is a batch dataset
         """
-        return bool(pgc.dataset_info_is_batch(self._info))
+        return bool(get_pgc().dataset_info_is_batch(self._info))
 
     def batch_size(self) -> int:
         """
@@ -89,7 +89,7 @@ class CDatasetInfo:
         Returns:
             The size of the dataset
         """
-        return pgc.dataset_info_batch_size(self._info)
+        return get_pgc().dataset_info_batch_size(self._info)
 
     def n_components(self) -> int:
         """
@@ -98,7 +98,7 @@ class CDatasetInfo:
         Returns:
             The amount of components in the dataset
         """
-        return pgc.dataset_info_n_components(self._info)
+        return get_pgc().dataset_info_n_components(self._info)
 
     def components(self) -> list[ComponentType]:
         """
@@ -108,7 +108,7 @@ class CDatasetInfo:
             A list of the component names in the dataset
         """
         return [
-            _str_to_component_type(pgc.dataset_info_component_name(self._info, idx))
+            _str_to_component_type(get_pgc().dataset_info_component_name(self._info, idx))
             for idx in range(self.n_components())
         ]
 
@@ -121,7 +121,7 @@ class CDatasetInfo:
             or -1 if the scenario is not uniform (different amount per scenario)
         """
         return {
-            component_name: pgc.dataset_info_elements_per_scenario(self._info, idx)
+            component_name: get_pgc().dataset_info_elements_per_scenario(self._info, idx)
             for idx, component_name in enumerate(self.components())
         }
 
@@ -135,7 +135,7 @@ class CDatasetInfo:
                 the product of the batch size and the amount of elements per scenario for that component.
         """
         return {
-            component_name: pgc.dataset_info_total_elements(self._info, idx)
+            component_name: get_pgc().dataset_info_total_elements(self._info, idx)
             for idx, component_name in enumerate(self.components())
         }
 
@@ -150,13 +150,13 @@ class CDatasetInfo:
         result_dict: dict[ComponentType, None | list[AttributeType]] = {}
         components = self.components()
         for component_idx, component_name in enumerate(components):
-            has_indications = pgc.dataset_info_has_attribute_indications(self._info, component_idx)
+            has_indications = get_pgc().dataset_info_has_attribute_indications(self._info, component_idx)
             if has_indications == 0:
                 result_dict[component_name] = None
             else:
-                n_indications = pgc.dataset_info_n_attribute_indications(self._info, component_idx)
+                n_indications = get_pgc().dataset_info_n_attribute_indications(self._info, component_idx)
                 result_dict[component_name] = [
-                    pgc.dataset_info_attribute_name(self._info, component_idx, attribute_idx)
+                    get_pgc().dataset_info_attribute_name(self._info, component_idx, attribute_idx)
                     for attribute_idx in range(n_indications)
                 ]
         return result_dict
@@ -195,7 +195,7 @@ class CMutableDataset:
             instance._is_batch = False
             instance._batch_size = 1
 
-        instance._mutable_dataset = pgc.create_dataset_mutable(
+        instance._mutable_dataset = get_pgc().create_dataset_mutable(
             instance._dataset_type.value, instance._is_batch, instance._batch_size
         )
         assert_no_error()
@@ -221,7 +221,7 @@ class CMutableDataset:
         Returns:
             The dataset info for this dataset.
         """
-        return CDatasetInfo(pgc.dataset_mutable_get_info(self._mutable_dataset))
+        return CDatasetInfo(get_pgc().dataset_mutable_get_info(self._mutable_dataset))
 
     def get_buffer_views(self) -> list[CBuffer]:
         """
@@ -272,7 +272,7 @@ class CMutableDataset:
         self._register_buffer(component, c_buffer)
 
     def _register_buffer(self, component: ComponentType, buffer: CBuffer):
-        pgc.dataset_mutable_add_buffer(
+        get_pgc().dataset_mutable_add_buffer(
             dataset=self._mutable_dataset,
             component=component.value,
             elements_per_scenario=buffer.n_elements_per_scenario,
@@ -285,7 +285,7 @@ class CMutableDataset:
             self._register_attribute_buffer(component, attr, attr_data)
 
     def _register_attribute_buffer(self, component, attr, attr_data):
-        pgc.dataset_mutable_add_attribute_buffer(
+        get_pgc().dataset_mutable_add_attribute_buffer(
             dataset=self._mutable_dataset,
             component=component.value,
             attribute=attr,
@@ -303,7 +303,7 @@ class CMutableDataset:
             raise ValueError(f"Dataset must have a consistent batch size across all components. {VALIDATOR_MSG}")
 
     def __del__(self):
-        pgc.destroy_dataset_mutable(self._mutable_dataset)
+        get_pgc().destroy_dataset_mutable(self._mutable_dataset)
 
 
 class CConstDataset:
@@ -326,7 +326,7 @@ class CConstDataset:
 
         # create from mutable dataset
         mutable_dataset = CMutableDataset(data=data, dataset_type=dataset_type)
-        instance._const_dataset = pgc.create_dataset_const_from_mutable(mutable_dataset.get_dataset_ptr())
+        instance._const_dataset = get_pgc().create_dataset_const_from_mutable(mutable_dataset.get_dataset_ptr())
         assert_no_error()
         instance._buffer_views = mutable_dataset.get_buffer_views()
 
@@ -348,10 +348,20 @@ class CConstDataset:
         Returns:
             The dataset info for this dataset.
         """
-        return CDatasetInfo(pgc.dataset_const_get_info(self._const_dataset))
+        return CDatasetInfo(get_pgc().dataset_const_get_info(self._const_dataset))
+
+    def set_next_cartesian_product_dimension(self, next_dataset: "CConstDataset") -> None:
+        """
+        Set the next dataset in the linked list.
+
+        Args:
+            next_dataset: The next dataset to set.
+        """
+        get_pgc().dataset_const_set_next_cartesian_product_dimension(self._const_dataset, next_dataset._const_dataset)
+        assert_no_error()
 
     def __del__(self):
-        pgc.destroy_dataset_const(self._const_dataset)
+        get_pgc().destroy_dataset_const(self._const_dataset)
 
 
 class CWritableDataset:
@@ -400,7 +410,7 @@ class CWritableDataset:
         Returns:
             The dataset info for this dataset.
         """
-        return CDatasetInfo(pgc.dataset_writable_get_info(self._writable_dataset))
+        return CDatasetInfo(get_pgc().dataset_writable_get_info(self._writable_dataset))
 
     def get_data(self) -> Dataset:
         """
@@ -426,11 +436,11 @@ class CWritableDataset:
         """
         return self._data[component]
 
-    def get_data_filter(self) -> _ComponentAttributeMappingDict:
+    def get_data_filter(self) -> ComponentAttributeMappingDict:
         """Gets the data filter requested
 
         Returns:
-            _ComponentAttributeMappingDict: data filter
+            ComponentAttributeMappingDict: data filter
         """
         return self._data_filter
 
@@ -445,7 +455,7 @@ class CWritableDataset:
         self._register_buffer(component, get_buffer_view(self._data[component], schema))
 
     def _register_buffer(self, component: ComponentType, buffer: CBuffer):
-        pgc.dataset_writable_set_buffer(
+        get_pgc().dataset_writable_set_buffer(
             dataset=self._writable_dataset,
             component=component,
             indptr=buffer.indptr,
@@ -461,7 +471,7 @@ class CWritableDataset:
         attribute: AttributeType,
         buffer: CAttributeBuffer,
     ):
-        pgc.dataset_writable_set_attribute_buffer(
+        get_pgc().dataset_writable_set_attribute_buffer(
             dataset=self._writable_dataset,
             component=component,
             attribute=attribute,

power_grid_model/_core/power_grid_meta.py

@@ -8,7 +8,7 @@ Load meta data from C core and define numpy structured array
 
 from dataclasses import dataclass
 from enum import IntEnum
-from typing import Any
+from typing import Any, overload
 
 import numpy as np
 
@@ -21,7 +21,12 @@ from power_grid_model._core.dataset_definitions import (
     _str_to_component_type,
     _str_to_datatype,
 )
-from power_grid_model._core.power_grid_core import AttributePtr, ComponentPtr, DatasetPtr, power_grid_core as pgc
+from power_grid_model._core.power_grid_core import (
+    AttributePtr,
+    ComponentPtr,
+    DatasetPtr,
+    get_power_grid_core as get_pgc,
+)
 
 
 # constant enum for ctype
@@ -35,7 +40,7 @@ class PGMCType(IntEnum):
 
 
 _CTYPE_NUMPY_MAP = {PGMCType.double: "f8", PGMCType.int32: "i4", PGMCType.int8: "i1", PGMCType.double3: "(3,)f8"}
-_ENDIANNESS = "<" if pgc.is_little_endian() == 1 else ">"
+_ENDIANNESS = "<" if get_pgc().is_little_endian() == 1 else ">"
 _NAN_VALUE_MAP = {
     f"{_ENDIANNESS}f8": np.nan,
     f"{_ENDIANNESS}(3,)f8": np.nan,
@@ -82,10 +87,10 @@ def _generate_meta_data() -> PowerGridMetaData:
 
     """
     py_meta_data = {}
-    n_datasets = pgc.meta_n_datasets()
+    n_datasets = get_pgc().meta_n_datasets()
     for i in range(n_datasets):
-        dataset = pgc.meta_get_dataset_by_idx(i)
-        py_meta_data[_str_to_datatype(pgc.meta_dataset_name(dataset))] = _generate_meta_dataset(dataset)
+        dataset = get_pgc().meta_get_dataset_by_idx(i)
+        py_meta_data[_str_to_datatype(get_pgc().meta_dataset_name(dataset))] = _generate_meta_dataset(dataset)
     return py_meta_data
 
 
@@ -99,10 +104,10 @@ def _generate_meta_dataset(dataset: DatasetPtr) -> DatasetMetaData:
 
     """
     py_meta_dataset = {}
-    n_components = pgc.meta_n_components(dataset)
+    n_components = get_pgc().meta_n_components(dataset)
     for i in range(n_components):
-        component = pgc.meta_get_component_by_idx(dataset, i)
-        py_meta_dataset[_str_to_component_type(pgc.meta_component_name(component))] = _generate_meta_component(
+        component = get_pgc().meta_get_component_by_idx(dataset, i)
+        py_meta_dataset[_str_to_component_type(get_pgc().meta_component_name(component))] = _generate_meta_component(
             component
         )
     return py_meta_dataset
@@ -121,8 +126,8 @@ def _generate_meta_component(component: ComponentPtr) -> ComponentMetaData:
     dtype_dict = _generate_meta_attributes(component)
     dtype = np.dtype({k: v for k, v in dtype_dict.items() if k != "nans"})  # type: ignore
     nans = dict(zip(dtype_dict["names"], dtype_dict["nans"]))
-    if dtype.alignment != pgc.meta_component_alignment(component):
-        raise TypeError(f'Aligment mismatch for component type: "{pgc.meta_component_name(component)}" !')
+    if dtype.alignment != get_pgc().meta_component_alignment(component):
+        raise TypeError(f'Aligment mismatch for component type: "{get_pgc().meta_component_name(component)}" !')
     # get single nan scalar
     nan_scalar = np.empty(1, dtype=dtype)
     for key, value in nans.items():
@@ -143,12 +148,12 @@ def _generate_meta_attributes(component: ComponentPtr) -> dict:
     formats = []
     offsets = []
     nans = []
-    n_attrs = pgc.meta_n_attributes(component)
+    n_attrs = get_pgc().meta_n_attributes(component)
     for i in range(n_attrs):
-        attribute: AttributePtr = pgc.meta_get_attribute_by_idx(component, i)
-        attr_name: str = pgc.meta_attribute_name(attribute)
-        attr_ctype: int = pgc.meta_attribute_ctype(attribute)
-        attr_offset: int = pgc.meta_attribute_offset(attribute)
+        attribute: AttributePtr = get_pgc().meta_get_attribute_by_idx(component, i)
+        attr_name: str = get_pgc().meta_attribute_name(attribute)
+        attr_ctype: int = get_pgc().meta_attribute_ctype(attribute)
+        attr_offset: int = get_pgc().meta_attribute_offset(attribute)
         attr_np_type = f"{_ENDIANNESS}{_CTYPE_NUMPY_MAP[PGMCType(attr_ctype)]}"
         attr_nan = _NAN_VALUE_MAP[attr_np_type]
         names.append(attr_name)
@@ -159,7 +164,7 @@ def _generate_meta_attributes(component: ComponentPtr) -> dict:
         "names": names,
         "formats": formats,
         "offsets": offsets,
-        "itemsize": pgc.meta_component_size(component),
+        "itemsize": get_pgc().meta_component_size(component),
         "aligned": True,
         "nans": nans,
     }
@@ -172,10 +177,24 @@ The data types for all dataset types and components used by the Power Grid Model
 """
 
 
+@overload
+def initialize_array(
+    data_type: DatasetTypeLike,
+    component_type: ComponentTypeLike,
+    shape: int | tuple[int],
+    empty: bool = False,
+) -> SingleArray: ...
+@overload
+def initialize_array(
+    data_type: DatasetTypeLike,
+    component_type: ComponentTypeLike,
+    shape: tuple[int, int],
+    empty: bool = False,
+) -> DenseBatchArray: ...
 def initialize_array(
     data_type: DatasetTypeLike,
     component_type: ComponentTypeLike,
-    shape: tuple | int,
+    shape: int | tuple[int] | tuple[int, int],
     empty: bool = False,
 ) -> SingleArray | DenseBatchArray:
     """