power-grid-model 1.11.33__py3-none-win_amd64.whl → 1.12.70__py3-none-win_amd64.whl

power_grid_model/_core/data_handling.py

@@ -1,141 +1,195 @@
-# SPDX-FileCopyrightText: Contributors to the Power Grid Model project <powergridmodel@lfenergy.org>
-#
-# SPDX-License-Identifier: MPL-2.0
-
-"""
-Data handling
-"""
-
-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.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.utils import process_data_filter
-
-
-def get_output_type(*, calculation_type: CalculationType, symmetric: bool) -> DatasetType:
-    """
-    Get the output type based on the provided arguments.
-
-    Args:
-        calculation_type:
-            request the output type for a specific calculation type (power_flow, state_estimation, ...)
-        symmetric:
-            True: three-phase symmetric calculation, even for asymmetric loads/generations
-            False: three-phase asymmetric calculation
-
-    Returns:
-        the output type that fits the format requested by the output type
-    """
-    if calculation_type in (CalculationType.power_flow, CalculationType.state_estimation):
-        return DatasetType.sym_output if symmetric else DatasetType.asym_output
-
-    if calculation_type == CalculationType.short_circuit:
-        return DatasetType.sc_output
-
-    raise NotImplementedError()
-
-
-def prepare_input_view(input_data: SingleDataset) -> CConstDataset:
-    """
-    Create a view of the input data in a format compatible with the PGM core libary.
-
-    Args:
-        input_data:
-            the input data to create the view from
-
-    Returns:
-        instance of CConstDataset ready to be fed into C API
-    """
-    return CConstDataset(input_data, dataset_type=DatasetType.input)
-
-
-def prepare_update_view(update_data: Dataset) -> CConstDataset:
-    """
-    Create a view of the update data, or an empty view if not provided, in a format compatible with the PGM core libary.
-
-    Args:
-        update_data:
-            the update data to create the view from. Defaults to None
-
-    Returns:
-        instance of CConstDataset ready to be fed into C API
-    """
-    return CConstDataset(update_data, dataset_type=DatasetType.update)
-
-
-def prepare_output_view(output_data: Dataset, output_type: DatasetType) -> CMutableDataset:
-    """
-    create a view of the output data in a format compatible with the PGM core libary.
-
-    Args:
-        output_data:
-            the output data to create the view from
-        output_type:
-            the output type of the output_data
-
-    Returns:
-        instance of CMutableDataset ready to be fed into C API
-    """
-    return CMutableDataset(output_data, dataset_type=output_type)
-
-
-def create_output_data(
-    output_component_types: ComponentAttributeMapping,
-    output_type: DatasetType,
-    all_component_count: dict[ComponentType, int],
-    is_batch: bool,
-    batch_size: int,
-) -> Dataset:
-    """
-    Create the output dataset based on component and batch size from the model; and output attributes requested by user.
-
-    Args:
-        output_component_types:
-            the output components the user seeks to extract
-        output_type:
-            the type of output that the user will see (as per the calculation options)
-        all_component_count:
-            the amount of components in the grid (as per the input data)
-        is_batch:
-            if the dataset is batch
-        batch_size:
-            the batch size
-
-    Returns:
-        Dataset: output dataset
-    """
-    processed_output_types = process_data_filter(output_type, output_component_types, list(all_component_count.keys()))
-
-    all_component_count = {k: v for k, v in all_component_count.items() if k in processed_output_types}
-
-    # create result dataset
-    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,)
-
-        requested_component = processed_output_types[name]
-        dtype = power_grid_meta_data[output_type][name].dtype
-        if dtype.names is None:
-            raise PowerGridUnreachableHitError
-        if requested_component is None:
-            result_dict[name] = initialize_array(output_type, name, shape=shape, empty=True)
-        elif requested_component in [
-            ComponentAttributeFilterOptions.everything,
-            ComponentAttributeFilterOptions.relevant,
-        ]:
-            result_dict[name] = {attr: np.empty(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}
-
-    return result_dict
+# SPDX-FileCopyrightText: Contributors to the Power Grid Model project <powergridmodel@lfenergy.org>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+"""
+Data handling
+"""
+
+from typing import Literal, overload
+
+import numpy as np
+
+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, 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:
+    """
+    Get the output type based on the provided arguments.
+
+    Args:
+        calculation_type:
+            request the output type for a specific calculation type (power_flow, state_estimation, ...)
+        symmetric:
+            True: three-phase symmetric calculation, even for asymmetric loads/generations
+            False: three-phase asymmetric calculation
+
+    Returns:
+        the output type that fits the format requested by the output type
+    """
+    if calculation_type in (CalculationType.power_flow, CalculationType.state_estimation):
+        return DatasetType.sym_output if symmetric else DatasetType.asym_output
+
+    if calculation_type == CalculationType.short_circuit:
+        return DatasetType.sc_output
+
+    raise NotImplementedError
+
+
+def prepare_input_view(input_data: SingleDataset) -> CConstDataset:
+    """
+    Create a view of the input data in a format compatible with the PGM core libary.
+
+    Args:
+        input_data:
+            the input data to create the view from
+
+    Returns:
+        instance of CConstDataset ready to be fed into C API
+    """
+    return CConstDataset(input_data, dataset_type=DatasetType.input)
+
+
+def prepare_update_view(update_data: Dataset) -> CConstDataset:
+    """
+    Create a view of the update data, or an empty view if not provided, in a format compatible with the PGM core libary.
+
+    Args:
+        update_data:
+            the update data to create the view from. Defaults to None
+
+    Returns:
+        instance of CConstDataset ready to be fed into C API
+    """
+    return CConstDataset(update_data, dataset_type=DatasetType.update)
+
+
+def prepare_output_view(output_data: Dataset, output_type: DatasetType) -> CMutableDataset:
+    """
+    create a view of the output data in a format compatible with the PGM core libary.
+
+    Args:
+        output_data:
+            the output data to create the view from
+        output_type:
+            the output type of the output_data
+
+    Returns:
+        instance of CMutableDataset ready to be fed into C API
+    """
+    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,
+    all_component_count: dict[ComponentType, int],
+    is_batch: bool,
+    batch_size: int,
+) -> Dataset:
+    """
+    Create the output dataset based on component and batch size from the model; and output attributes requested by user.
+
+    Args:
+        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)
+        all_component_count:
+            the amount of components in the grid (as per the input data)
+        is_batch:
+            if the dataset is batch
+        batch_size:
+            the batch size
+
+    Returns:
+        Dataset: output dataset
+    """
+    processed_output_types = process_data_filter(output_type, output_component_types, list(all_component_count.keys()))
+
+    all_component_count = {k: v for k, v in all_component_count.items() if k in processed_output_types}
+
+    # create result dataset
+    result_dict: Dataset = {}
+
+    for name, count in all_component_count.items():
+        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
+        if dtype.names is None:
+            raise PowerGridUnreachableHitError
+        if requested_component is None:
+            result_dict[name] = initialize_array(output_type, name, shape=shape, empty=True)
+        elif requested_component in [
+            ComponentAttributeFilterOptions.everything,
+            ComponentAttributeFilterOptions.relevant,
+        ]:
+            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=shape, dtype=dtype[attr]) for attr in requested_component}
+
+    return result_dict

power_grid_model/_core/data_types.py

@@ -1,132 +1,143 @@
-# SPDX-FileCopyrightText: Contributors to the Power Grid Model project <powergridmodel@lfenergy.org>
-#
-# SPDX-License-Identifier: MPL-2.0
-
-"""
-Data types involving PGM datasets.
-
-Data types for library-internal use. In an attempt to clarify type hints, some types
-have been defined and explained in this file.
-"""
-
-from typing import TypeAlias, TypedDict, TypeVar
-
-import numpy as np
-
-from power_grid_model._core.dataset_definitions import ComponentTypeVar
-
-SingleArray: TypeAlias = np.ndarray
-
-AttributeType: TypeAlias = str
-
-SingleColumn: TypeAlias = np.ndarray
-
-DenseBatchArray: TypeAlias = np.ndarray
-
-SingleColumnarData = dict[AttributeType, SingleColumn]
-
-_SingleComponentData = TypeVar("_SingleComponentData", SingleArray, SingleColumnarData)  # deduction helper
-SingleComponentData = SingleArray | SingleColumnarData
-
-
-SingleDataset = dict[ComponentTypeVar, _SingleComponentData]
-
-BatchList = list[SingleDataset]
-
-BatchColumn: TypeAlias = np.ndarray
-
-DenseBatchColumnarData = dict[AttributeType, BatchColumn]
-
-IndexPointer: TypeAlias = np.ndarray
-
-
-class SparseBatchColumnarData(TypedDict):
-    """
-    Sparse batch columnar data is a dictionary containing the keys `indptr` and `data`.
-
-    - data: a :class:`SingleColumnarData`. The exact supported attribute columns depend on the component type.
-    - indptr: an :class:`IndexPointer` representing the start and end indices for each batch scenario.
-
-    - Examples:
-
-        - structure: {"indptr": :class:`IndexPointer`, "data": :class:`SingleColumnarData`}
-        - concrete example: {"indptr": [0, 2, 2, 3], "data": {"id": [0, 1, 0], "status": [1, 1, 0]}}
-
-            - the scenario 0 sets the status of components with ids 0 and 1 to 1
-              (and keeps defaults for other components)
-            - scenario 1 keeps the default values for all components
-            - scenario 2 sets the status of component with id 0 to 0 (and keeps defaults for other components)
-    """
-
-    indptr: IndexPointer
-    data: SingleColumnarData
-
-
-class SparseBatchArray(TypedDict):
-    """
-    A sparse batch array is a dictionary containing the keys `indptr` and `data`.
-
-    - data: a :class:`SingleArray`. The exact dtype depends on the type of component.
-    - indptr: an :class:`IndexPointer` representing the start and end indices for each batch scenario.
-
-    - Examples:
-
-        - structure: {"indptr": :class:`IndexPointer`, "data": :class:`SingleArray`}
-        - concrete example: {"indptr": [0, 2, 2, 3], "data": [(0, 1, 1), (1, 1, 1), (0, 0, 0)]}
-
-            - the scenario 0 sets the statuses of components with ids 0 and 1 to 1
-              (and keeps defaults for other components)
-            - scenario 1 keeps the default values for all components
-            - scenario 2 sets the statuses of component with id 0 to 0 (and keeps defaults for other components)
-    """
-
-    indptr: IndexPointer
-    data: SingleArray
-
-
-SparseBatchData = SparseBatchArray | SparseBatchColumnarData
-
-SparseDataComponentType: TypeAlias = str
-
-BatchColumnarData = DenseBatchColumnarData | SparseBatchColumnarData
-
-ColumnarData = SingleColumnarData | BatchColumnarData
-BatchArray = DenseBatchArray | SparseBatchArray
-
-
-BatchComponentData = BatchArray | BatchColumnarData
-
-_BatchComponentData = TypeVar("_BatchComponentData", BatchArray, BatchColumnarData)  # deduction helper
-
-
-BatchDataset = dict[ComponentTypeVar, _BatchComponentData]
-
-
-DataArray = SingleArray | BatchArray
-
-
-_ComponentData = TypeVar("_ComponentData", SingleComponentData, BatchComponentData)  # deduction helper
-ComponentData = DataArray | ColumnarData
-
-Dataset = dict[ComponentTypeVar, _ComponentData]
-
-
-DenseBatchData = DenseBatchArray | DenseBatchColumnarData
-
-NominalValue = int
-
-RealValue = float
-
-AsymValue = tuple[RealValue, RealValue, RealValue]
-
-AttributeValue = RealValue | NominalValue | AsymValue
-
-Component = dict[AttributeType, AttributeValue | str]
-
-ComponentList = list[Component]
-
-SinglePythonDataset = dict[ComponentTypeVar, ComponentList]
-
-BatchPythonDataset = list[SinglePythonDataset]
-
-PythonDataset = SinglePythonDataset | BatchPythonDataset
+# SPDX-FileCopyrightText: Contributors to the Power Grid Model project <powergridmodel@lfenergy.org>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+"""
+Data types involving PGM datasets.
+
+Data types for library-internal use. In an attempt to clarify type hints, some types
+have been defined and explained in this file.
+"""
+
+from typing import TypeAlias, TypedDict, TypeVar
+
+import numpy as np
+
+from power_grid_model._core.dataset_definitions import ComponentType, ComponentTypeVar
+
+SingleArray: TypeAlias = np.ndarray
+
+AttributeType: TypeAlias = str
+
+SingleColumn: TypeAlias = np.ndarray
+
+DenseBatchArray: TypeAlias = np.ndarray
+
+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]
+
+BatchColumn: TypeAlias = np.ndarray
+
+DenseBatchColumnarData = dict[AttributeType, BatchColumn]
+
+IndexPointer: TypeAlias = np.ndarray
+
+
+class SparseBatchColumnarData(TypedDict):
+    """
+    Sparse batch columnar data is a dictionary containing the keys `indptr` and `data`.
+
+    - data: a :class:`SingleColumnarData`. The exact supported attribute columns depend on the component type.
+    - indptr: an :class:`IndexPointer` representing the start and end indices for each batch scenario.
+
+    - Examples:
+
+        - structure: {"indptr": :class:`IndexPointer`, "data": :class:`SingleColumnarData`}
+        - concrete example: {"indptr": [0, 2, 2, 3], "data": {"id": [0, 1, 0], "status": [1, 1, 0]}}
+
+            - the scenario 0 sets the status of components with ids 0 and 1 to 1
+              (and keeps defaults for other components)
+            - scenario 1 keeps the default values for all components
+            - scenario 2 sets the status of component with id 0 to 0 (and keeps defaults for other components)
+    """
+
+    indptr: IndexPointer
+    data: SingleColumnarData
+
+
+class SparseBatchArray(TypedDict):
+    """
+    A sparse batch array is a dictionary containing the keys `indptr` and `data`.
+
+    - data: a :class:`SingleArray`. The exact dtype depends on the type of component.
+    - indptr: an :class:`IndexPointer` representing the start and end indices for each batch scenario.
+
+    - Examples:
+
+        - structure: {"indptr": :class:`IndexPointer`, "data": :class:`SingleArray`}
+        - concrete example: {"indptr": [0, 2, 2, 3], "data": [(0, 1, 1), (1, 1, 1), (0, 0, 0)]}
+
+            - the scenario 0 sets the statuses of components with ids 0 and 1 to 1
+              (and keeps defaults for other components)
+            - scenario 1 keeps the default values for all components
+            - scenario 2 sets the statuses of component with id 0 to 0 (and keeps defaults for other components)
+    """
+
+    indptr: IndexPointer
+    data: SingleArray
+
+
+SparseBatchData = SparseBatchArray | SparseBatchColumnarData
+
+SparseDataComponentType: TypeAlias = str
+
+BatchColumnarData = DenseBatchColumnarData | SparseBatchColumnarData
+
+ColumnarData = SingleColumnarData | BatchColumnarData
+BatchArray = DenseBatchArray | SparseBatchArray
+
+
+BatchComponentData = BatchArray | BatchColumnarData
+
+_BatchComponentData = TypeVar("_BatchComponentData", BatchArray, BatchColumnarData)  # deduction helper
+
+
+BatchDataset = dict[ComponentTypeVar, _BatchComponentData]
+
+
+DataArray = SingleArray | BatchArray
+
+
+_ComponentData = TypeVar("_ComponentData", SingleComponentData, BatchComponentData)  # deduction helper
+ComponentData = DataArray | ColumnarData
+
+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
+
+AsymValue = tuple[RealValue, RealValue, RealValue]
+
+AttributeValue = RealValue | NominalValue | AsymValue
+
+Component = dict[AttributeType, AttributeValue | str]
+
+ComponentList = list[Component]
+
+SinglePythonDataset = dict[ComponentTypeVar, ComponentList]
+
+BatchPythonDataset = list[SinglePythonDataset]
+
+PythonDataset = SinglePythonDataset | BatchPythonDataset