Feature / Columnar Data helper function

src/power_grid_model/_core/power_grid_meta.py

@@ -8,7 +8,7 @@
 
 from dataclasses import dataclass
 from enum import IntEnum
-from typing import Any
+from typing import Any, Iterable
 
 import numpy as np
 
@@ -21,7 +21,8 @@
     _str_to_datatype,
 )
 from power_grid_model._core.power_grid_core import AttributePtr, ComponentPtr, DatasetPtr, power_grid_core as pgc
-from power_grid_model.data_types import DenseBatchArray, SingleArray
+from power_grid_model.data_types import DenseBatchData, SingleComponentData
+from power_grid_model.errors import PowerGridUnreachableHitError
 
 
 # constant enum for ctype
@@ -178,7 +179,9 @@ def initialize_array(
     component_type: ComponentTypeLike,
     shape: tuple | int,
     empty: bool = False,
-) -> SingleArray | DenseBatchArray:
+    *,
+    attributes: Iterable[str] | None = None,
+) -> SingleComponentData | DenseBatchData:
     """
     Initializes an array for use in Power Grid Model calculations
 
@@ -197,11 +200,29 @@ def initialize_array(
     component_type = _str_to_component_type(component_type)
     if not isinstance(shape, tuple):
         shape = (shape,)
-    if empty:
-        return np.empty(shape=shape, dtype=power_grid_meta_data[data_type][component_type].dtype, order="C")
-    return np.full(
-        shape=shape,
-        fill_value=power_grid_meta_data[data_type][component_type].nan_scalar,
-        dtype=power_grid_meta_data[data_type][component_type].dtype,
-        order="C",
-    )
+    component_meta = power_grid_meta_data[data_type][component_type]
+    if attributes is None:
+        if empty:
+            return np.empty(shape=shape, dtype=component_meta.dtype, order="C")
+        return np.full(
+            shape=shape,
+            fill_value=component_meta.nan_scalar,
+            dtype=component_meta.dtype,
+            order="C",
+        )
+    data = {}
+    for attribute in attributes:
+        if component_meta.dtype.names is None:
+            raise PowerGridUnreachableHitError
+        if attribute not in component_meta.dtype.names:
+            raise ValueError(f"Attribute {attribute} is not available for {component_type}")
+        if empty:
+            data[attribute] = np.empty(shape=shape, dtype=component_meta.dtype[attribute], order="C")
+        else:
+            data[attribute] = np.full(
+                shape=shape,
+                fill_value=component_meta.nan_scalar[attribute],
+                dtype=component_meta.dtype[attribute],
+                order="C",
+            )
+    return data

tests/unit/test_meta_data.py

@@ -20,6 +20,20 @@ def test_initialize_array():
     assert np.all(np.isnan(arr_2d["u_rated"]))
 
 
+def test_initialize_array__columnar():
+    some_attributes = ["id", "from_node", "r1"]
+    actual = initialize_array("input", "line", 3, attributes=some_attributes)
+    assert isinstance(actual, dict)
+    assert actual.keys() == set(some_attributes)
+    assert all(v.shape == (3,) for v in actual.values())
+    assert np.all(np.isnan(actual["r1"]))
+
+    actual_2d = initialize_array("input", "line", (2, 3), attributes=some_attributes)
+    assert actual.keys() == set(some_attributes)
+    assert all(v.shape == (2, 3) for v in actual_2d.values())
+    assert np.all(np.isnan(actual_2d["r1"]))
+
+
 def test_sensor_meta_data():
     sensors = ["sym_voltage_sensor", "asym_voltage_sensor", "sym_power_sensor", "asym_power_sensor"]
     input_voltage = ["u_measured", "u_angle_measured", "u_sigma"]
@@ -29,8 +43,7 @@ def test_sensor_meta_data():
     for sensor in sensors:
         for meta_type in ["input", "update", "sym_output", "asym_output"]:
             meta_data = power_grid_meta_data[meta_type]
-            # comp_names = list(meta_data.keys())
-            # assert sensor in comp_names
+            assert sensor in meta_data
             meta_data_sensor = meta_data[sensor]
             attr_names = meta_data_sensor.dtype_dict["names"]
             assert "id" in attr_names