Merge pull request #821 from PowerGridModel/feature/activate-columnar…

…-output

Remove compatibility convert for output

src/power_grid_model/core/data_handling.py

@@ -6,13 +6,15 @@
 Data handling
 """
 
+import numpy as np
 
 from power_grid_model._utils import process_data_filter
 from power_grid_model.core.dataset_definitions import ComponentType, DatasetType
 from power_grid_model.core.power_grid_dataset import CConstDataset, CMutableDataset
-from power_grid_model.core.power_grid_meta import initialize_array
+from power_grid_model.core.power_grid_meta import initialize_array, power_grid_meta_data
 from power_grid_model.data_types import Dataset, SingleDataset
-from power_grid_model.enum import CalculationType
+from power_grid_model.enum import CalculationType, ComponentAttributeFilterOptions
+from power_grid_model.errors import PowerGridUnreachableHitError
 from power_grid_model.typing import ComponentAttributeMapping
 
 
@@ -113,13 +115,27 @@ def create_output_data(
     all_component_count = {k: v for k, v in all_component_count.items() if k in processed_output_types}
 
     # create result dataset
-    result_dict = {}
+    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,)
-        result_dict[name] = initialize_array(output_type, name, shape=shape, empty=True)
+
+        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

src/power_grid_model/core/power_grid_dataset.py

@@ -149,9 +149,8 @@ def __new__(cls, data: Dataset, dataset_type: Any = None):
         instance._dataset_type = dataset_type if dataset_type in DatasetType else get_dataset_type(data)
         instance._schema = power_grid_meta_data[instance._dataset_type]
 
-        compatibility_converted_data = data
-        if compatibility_converted_data:
-            first_component, first_component_data = next(iter(compatibility_converted_data.items()))
+        if data:
+            first_component, first_component_data = next(iter(data.items()))
             first_sub_info = get_buffer_properties(data=first_component_data, schema=instance._schema[first_component])
             instance._is_batch = first_sub_info.is_batch
             instance._batch_size = first_sub_info.batch_size
@@ -164,7 +163,7 @@ def __new__(cls, data: Dataset, dataset_type: Any = None):
         )
         assert_no_error()
 
-        instance._add_data(compatibility_converted_data)
+        instance._add_data(data)
         assert_no_error()
 
         return instance

src/power_grid_model/core/power_grid_model.py

@@ -11,7 +11,6 @@
 
 import numpy as np
 
-from power_grid_model._utils import compatibility_convert_row_columnar_dataset
 from power_grid_model.core.data_handling import (
     create_output_data,
     get_output_type,
@@ -22,7 +21,6 @@
 from power_grid_model.core.dataset_definitions import (
     ComponentType,
     ComponentTypeLike,
-    ComponentTypeVar,
     _map_to_component_types,
     _str_to_component_type,
 )
@@ -303,12 +301,6 @@ def _calculate_impl(  # pylint: disable=too-many-positional-arguments
             decode_error=decode_error,
         )
 
-        output_data = compatibility_convert_row_columnar_dataset(
-            data=output_data,
-            data_filter=output_component_types,
-            dataset_type=get_output_type(calculation_type=calculation_type, symmetric=symmetric),
-            available_components=list(self._get_output_component_count(calculation_type=calculation_type).keys()),
-        )
         return output_data
 
     def _calculate_power_flow(
@@ -320,7 +312,7 @@ def _calculate_power_flow(
         calculation_method: CalculationMethod | str = CalculationMethod.newton_raphson,
         update_data: Optional[Dataset] = None,
         threading: int = -1,
-        output_component_types: Optional[set[ComponentTypeVar] | list[ComponentTypeVar]] = None,
+        output_component_types: ComponentAttributeMapping = None,
         continue_on_batch_error: bool = False,
         decode_error: bool = True,
         tap_changing_strategy: TapChangingStrategy | str = TapChangingStrategy.disabled,
@@ -357,7 +349,7 @@ def _calculate_state_estimation(
         calculation_method: CalculationMethod | str = CalculationMethod.iterative_linear,
         update_data: Optional[Dataset] = None,
         threading: int = -1,
-        output_component_types: Optional[set[ComponentTypeVar] | list[ComponentTypeVar]] = None,
+        output_component_types: ComponentAttributeMapping = None,
         continue_on_batch_error: bool = False,
         decode_error: bool = True,
         experimental_features: _ExperimentalFeatures | str = _ExperimentalFeatures.disabled,
@@ -389,7 +381,7 @@ def _calculate_short_circuit(
         calculation_method: CalculationMethod | str = CalculationMethod.iec60909,
         update_data: Optional[Dataset] = None,
         threading: int = -1,
-        output_component_types: Optional[set[ComponentTypeVar] | list[ComponentTypeVar]] = None,
+        output_component_types: ComponentAttributeMapping = None,
         continue_on_batch_error: bool = False,
         decode_error: bool = True,
         short_circuit_voltage_scaling: ShortCircuitVoltageScaling | str = ShortCircuitVoltageScaling.maximum,
@@ -426,7 +418,7 @@ def calculate_power_flow(
         calculation_method: CalculationMethod | str = CalculationMethod.newton_raphson,
         update_data: Optional[dict[str, np.ndarray | dict[str, np.ndarray]] | Dataset] = None,
         threading: int = -1,
-        output_component_types: Optional[set[ComponentTypeVar] | list[ComponentTypeVar]] = None,
+        output_component_types: ComponentAttributeMapping = None,
         continue_on_batch_error: bool = False,
         decode_error: bool = True,
         tap_changing_strategy: TapChangingStrategy | str = TapChangingStrategy.disabled,
@@ -471,8 +463,17 @@ def calculate_power_flow(
                 - < 0: Sequential
                 - = 0: Parallel, use number of hardware threads
                 - > 0: Specify number of parallel threads
-            output_component_types ({set, list}, optional): List or set of component types you want to be present in
-                the output dict. By default, all component types will be in the output.
+            output_component_types (ComponentAttributeMapping):
+
+                - None: Row based data for all component types.
+                - set[ComponentTypeVar] or list[ComponentTypeVar]: Row based data for the specified component types.
+                - ComponentAttributeFilterOptions: Columnar data for all component types.
+                - dict[ComponentType, set[str] | list[str] | None | ComponentAttributeFilterOptions]:
+                    key: ComponentType
+                    value:
+                        - None: Row based data for the specified component types.
+                        - ComponentAttributeFilterOptions: Columnar data for the specified component types.
+                        - set[str] | list[str]: Columnar data for the specified component types and attributes.
             continue_on_batch_error (bool, optional): Continue the program (instead of throwing error) if some
                 scenarios fail.
             decode_error (bool, optional):
@@ -515,7 +516,7 @@ def calculate_state_estimation(
         calculation_method: CalculationMethod | str = CalculationMethod.iterative_linear,
         update_data: Optional[dict[str, np.ndarray | dict[str, np.ndarray]] | Dataset] = None,
         threading: int = -1,
-        output_component_types: Optional[set[ComponentTypeVar] | list[ComponentTypeVar]] = None,
+        output_component_types: ComponentAttributeMapping = None,
         continue_on_batch_error: bool = False,
         decode_error: bool = True,
     ) -> dict[ComponentType, np.ndarray]:
@@ -556,8 +557,17 @@ def calculate_state_estimation(
                 - < 0: Sequential
                 - = 0: Parallel, use number of hardware threads
                 - > 0: Specify number of parallel threads
-            output_component_types ({set, list}, optional): List or set of component types you want to be present in
-                the output dict. By default, all component types will be in the output.
+            output_component_types (ComponentAttributeMapping):
+
+                - None: Row based data for all component types.
+                - set[ComponentTypeVar] or list[ComponentTypeVar]: Row based data for the specified component types.
+                - ComponentAttributeFilterOptions: Columnar data for all component types.
+                - dict[ComponentType, set[str] | list[str] | None | ComponentAttributeFilterOptions]:
+                    key: ComponentType
+                    value:
+                        - None: Row based data for the specified component types.
+                        - ComponentAttributeFilterOptions: Columnar data for the specified component types.
+                        - set[str] | list[str]: Columnar data for the specified component types and attributes.
             continue_on_batch_error (bool, optional): Continue the program (instead of throwing error) if some
                 scenarios fail.
             decode_error (bool, optional):
@@ -596,7 +606,7 @@ def calculate_short_circuit(
         calculation_method: CalculationMethod | str = CalculationMethod.iec60909,
         update_data: Optional[dict[str, np.ndarray | dict[str, np.ndarray]] | Dataset] = None,
         threading: int = -1,
-        output_component_types: Optional[set[ComponentTypeVar] | list[ComponentTypeVar]] = None,
+        output_component_types: ComponentAttributeMapping = None,
         continue_on_batch_error: bool = False,
         decode_error: bool = True,
         short_circuit_voltage_scaling: ShortCircuitVoltageScaling | str = ShortCircuitVoltageScaling.maximum,
@@ -630,9 +640,17 @@ def calculate_short_circuit(
                 - < 0: Sequential
                 - = 0: Parallel, use number of hardware threads
                 - > 0: Specify number of parallel threads
-            output_component_types ({set, list}, optional):
-                List or set of component types you want to be present in the output dict.
-                By default, all component types will be in the output.
+            output_component_types (ComponentAttributeMapping):
+
+                - None: Row based data for all component types.
+                - set[ComponentTypeVar] or list[ComponentTypeVar]: Row based data for the specified component types.
+                - ComponentAttributeFilterOptions: Columnar data for all component types.
+                - dict[ComponentType, set[str] | list[str] | None | ComponentAttributeFilterOptions]:
+                    key: ComponentType
+                    value:
+                        - None: Row based data for the specified component types.
+                        - ComponentAttributeFilterOptions: Columnar data for the specified component types.
+                        - set[str] | list[str]: Columnar data for the specified component types and attributes.
             continue_on_batch_error (bool, optional):
                 Continue the program (instead of throwing error) if some scenarios fail.
             decode_error (bool, optional):

tests/unit/test_data_handling.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MPL-2.0
 
 import warnings
+from functools import partial
 
 import numpy as np
 import pytest
@@ -12,63 +13,127 @@
 from power_grid_model.core.dataset_definitions import ComponentType as CT, DatasetType as DT
 from power_grid_model.core.power_grid_core import VoidPtr
 from power_grid_model.core.power_grid_dataset import CMutableDataset
-from power_grid_model.core.power_grid_meta import initialize_array
+from power_grid_model.core.power_grid_meta import initialize_array, power_grid_meta_data
+from power_grid_model.enum import ComponentAttributeFilterOptions
+
+
+def columnar_array(component_type, n_components, attributes=None, batch_size_tuple=()):
+    component_dtype = power_grid_meta_data[DT.sym_output][component_type].dtype
+    required_attributes = (
+        set(component_dtype.names) & set(attributes) if attributes is not None else component_dtype.names
+    )
+    return {
+        attr: np.empty((n_components,) + batch_size_tuple, dtype=component_dtype[attr]) for attr in required_attributes
+    }
+
+
+def row_array(component_type, n_components, batch_size_tuple=()):
+    return initialize_array(DT.sym_output, component_type, (n_components,) + batch_size_tuple)
+
+
+@pytest.fixture(params=[1, 15])
+def batch_size(request):
+    return request.param
 
 
 @pytest.mark.parametrize(
-    ("output_component_types", "is_batch", "expected"),
+    ("output_component_types", "expected_fns"),
     [
         (
             None,
-            False,
             {
-                CT.node: initialize_array(DT.sym_output, CT.node, 4),
-                CT.sym_load: initialize_array(DT.sym_output, CT.sym_load, 3),
-                CT.source: initialize_array(DT.sym_output, CT.source, 1),
+                CT.node: partial(row_array, component_type=CT.node, n_components=4),
+                CT.sym_load: partial(row_array, component_type=CT.sym_load, n_components=3),
+                CT.source: partial(row_array, component_type=CT.source, n_components=1),
             },
         ),
         (
             [CT.node, CT.sym_load],
-            False,
             {
-                CT.node: initialize_array(DT.sym_output, CT.node, 4),
-                CT.sym_load: initialize_array(DT.sym_output, CT.sym_load, 3),
+                CT.node: partial(row_array, component_type=CT.node, n_components=4),
+                CT.sym_load: partial(row_array, component_type=CT.sym_load, n_components=3),
             },
         ),
         (
             {CT.node, CT.sym_load},
-            False,
             {
-                CT.node: initialize_array(DT.sym_output, CT.node, 4),
-                CT.sym_load: initialize_array(DT.sym_output, CT.sym_load, 3),
+                CT.node: partial(row_array, component_type=CT.node, n_components=4),
+                CT.sym_load: partial(row_array, component_type=CT.sym_load, n_components=3),
+            },
+        ),
+        pytest.param(
+            ComponentAttributeFilterOptions.relevant,
+            {
+                CT.node: partial(columnar_array, component_type=CT.node, n_components=4, attributes=None),
+                CT.sym_load: partial(columnar_array, component_type=CT.sym_load, n_components=3, attributes=None),
+                CT.source: partial(columnar_array, component_type=CT.source, n_components=1, attributes=None),
             },
         ),
         pytest.param(
-            {CT.node: [], CT.sym_load: []}, True, {CT.node: dict(), CT.sym_load: dict()}, marks=pytest.mark.xfail
+            {
+                CT.node: ComponentAttributeFilterOptions.everything,
+                CT.sym_load: ComponentAttributeFilterOptions.relevant,
+            },
+            {
+                CT.node: partial(columnar_array, component_type=CT.node, n_components=4, attributes=None),
+                CT.sym_load: partial(columnar_array, component_type=CT.sym_load, n_components=3, attributes=None),
+            },
+        ),
+        pytest.param(
+            {CT.node: ComponentAttributeFilterOptions.relevant, CT.sym_load: ["p"]},
+            {
+                CT.node: partial(columnar_array, component_type=CT.node, n_components=4, attributes=None),
+                CT.sym_load: partial(columnar_array, component_type=CT.sym_load, n_components=3, attributes=["p"]),
+            },
+        ),
+        pytest.param(
+            {CT.node: None, CT.sym_load: ["p"]},
+            {
+                CT.node: partial(row_array, component_type=CT.node, n_components=4),
+                CT.sym_load: partial(columnar_array, component_type=CT.sym_load, n_components=3, attributes=["p"]),
+            },
+        ),
+        pytest.param(
+            {CT.node: [], CT.sym_load: []},
+            {
+                CT.node: partial(columnar_array, component_type=CT.node, n_components=4, attributes=[]),
+                CT.sym_load: partial(columnar_array, component_type=CT.sym_load, n_components=3, attributes=[]),
+            },
+        ),
+        pytest.param(
+            {CT.node: [], CT.sym_load: ["p"]},
+            {
+                CT.node: partial(columnar_array, component_type=CT.node, n_components=4, attributes=[]),
+                CT.sym_load: partial(columnar_array, component_type=CT.sym_load, n_components=3, attributes=["p"]),
+            },
+        ),
+        pytest.param(
+            {CT.node: ["u"], CT.sym_load: ["p"]},
+            {
+                CT.node: partial(columnar_array, component_type=CT.node, n_components=4, attributes=["u"]),
+                CT.sym_load: partial(columnar_array, component_type=CT.sym_load, n_components=3, attributes=["p"]),
+            },
         ),
-        pytest.param({CT.node: [], CT.sym_load: ["p"]}, True, {}, marks=pytest.mark.xfail),
-        pytest.param({CT.node: ["u"], CT.sym_load: ["p"]}, True, {}, marks=pytest.mark.xfail),
-        pytest.param({CT.node: None, CT.sym_load: ["p"]}, True, {}, marks=pytest.mark.xfail),
     ],
 )
-def test_create_output_data(output_component_types, is_batch, expected):
-    # TODO use is_batch and shorten parameterization after columnar data implementation
+def test_create_output_data(output_component_types, expected_fns, batch_size):
+    """Test output_data creation. Batch size is set later in the test."""
     all_component_count = {CT.node: 4, CT.sym_load: 3, CT.source: 1}
-    batch_size = 15 if is_batch else 1
     actual = create_output_data(
         output_component_types=output_component_types,
         output_type=DT.sym_output,
         all_component_count=all_component_count,
-        is_batch=is_batch,
+        is_batch=False if batch_size == 1 else True,
         batch_size=batch_size,
     )
 
+    expected = {comp: fn(batch_size_tuple=(batch_size,)) for comp, fn in expected_fns.items()}
     assert actual.keys() == expected.keys()
     for comp in expected:
         if not is_columnar(expected[comp]):
             assert actual[comp].dtype == expected[comp].dtype
         elif expected[comp] == dict():
-            # Empty atrtibutes columnar
+            # Empty attributes columnar
             assert actual[comp] == expected[comp]
         else:
             assert actual[comp].keys() == expected[comp].keys()