Implement AttributeType for power_grid_model objects

docs/examples/pgm/basic_pgm_examples.ipynb

@@ -137,18 +137,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from power_grid_model import initialize_array\n",
+    "from power_grid_model import AttributeType, ComponentType, DatasetType, initialize_array\n",
     "\n",
     "grid_generator = RadialGridGenerator(grid_class=Grid, nr_nodes=5, nr_sources=1, nr_nops=0)\n",
     "grid = grid_generator.run(seed=0)\n",
     "core_interface = PowerGridModelInterface(grid=grid)\n",
     "\n",
-    "update_sym_load = initialize_array(\"update\", \"sym_load\", (10, len(grid.sym_load)))\n",
-    "update_sym_load[\"id\"] = [grid.sym_load.id.tolist()]\n",
-    "update_sym_load[\"p_specified\"] = [grid.sym_load.p_specified.tolist()] * np.linspace(0, 1, 10).reshape(-1, 1)\n",
-    "update_sym_load[\"q_specified\"] = [grid.sym_load.q_specified.tolist()] * np.linspace(0, 1, 10).reshape(-1, 1)\n",
+    "update_sym_load = initialize_array(DatasetType.update, ComponentType.sym_load, (10, len(grid.sym_load)))\n",
+    "update_sym_load[AttributeType.id] = [grid.sym_load.id.tolist()]\n",
+    "update_sym_load[AttributeType.p_specified] = [grid.sym_load.p_specified.tolist()] * np.linspace(0, 1, 10).reshape(-1, 1)\n",
+    "update_sym_load[AttributeType.q_specified] = [grid.sym_load.q_specified.tolist()] * np.linspace(0, 1, 10).reshape(-1, 1)\n",
     "update_data = {\n",
-    "    \"sym_load\": update_sym_load,\n",
+    "    ComponentType.sym_load: update_sym_load,\n",
     "}\n",
     "output = core_interface.calculate_power_flow(update_data=update_data)\n",
     "\n",

tests/conftest.py

@@ -5,7 +5,7 @@
 """Helper np.arrays used by various tests."""
 
 import pytest
-from power_grid_model import initialize_array
+from power_grid_model import AttributeType, ComponentType, DatasetType, initialize_array
 
 from power_grid_model_ds._core.model.graphs.models import RustworkxGraphModel
 from power_grid_model_ds._core.model.graphs.models.base import BaseGraphModel
@@ -71,58 +71,58 @@ def topologically_full_grid(grid: Grid):
 
 @pytest.fixture
 def input_data_pgm():
-    node = initialize_array("input", "node", 3)
-    node["id"] = [1, 2, 7]
-    node["u_rated"] = [10500.0, 10500.0, 10500.0]
-
-    line = initialize_array("input", "line", 2)
-    line["id"] = [9, 10]
-    line["from_node"] = [7, 7]
-    line["to_node"] = [2, 1]
-    line["from_status"] = [1, 1]
-    line["to_status"] = [1, 1]
-    line["r1"] = [0.00396133, 0.32598809]
-    line["x1"] = [4.53865336e-05, 1.34716591e-02]
-    line["c1"] = [0.0, 0.0]
-    line["tan1"] = [0.0, 0.0]
-    line["i_n"] = [303.91942029, 210.06857453]
-
-    link = initialize_array("input", "link", 0)
-    transformer = initialize_array("input", "transformer", 0)
-    three_winding_transformer = initialize_array("input", "three_winding_transformer", 0)
-
-    sym_load = initialize_array("input", "sym_load", 2)
-    sym_load["id"] = [5, 6]
-    sym_load["node"] = [1, 2]
-    sym_load["status"] = [1, 1]
-    sym_load["type"] = [0, 0]
-    sym_load["p_specified"] = [-287484.0, 26558.0]
-    sym_load["q_specified"] = [40640.0, 28148.0]
-
-    sym_gen = initialize_array("input", "sym_gen", 0)
-
-    source = initialize_array("input", "source", 1)
-    source["id"] = [8]
-    source["node"] = [7]
-    source["status"] = [1]
-    source["u_ref"] = [1.0]
-
-    transformer_tap_regulator = initialize_array("input", "transformer_tap_regulator", 0)
-    sym_power_sensor = initialize_array("input", "sym_power_sensor", 0)
-    sym_voltage_sensor = initialize_array("input", "sym_voltage_sensor", 0)
-    asym_voltage_sensor = initialize_array("input", "asym_voltage_sensor", 0)
+    node = initialize_array(DatasetType.input, ComponentType.node, 3)
+    node[AttributeType.id] = [1, 2, 7]
+    node[AttributeType.u_rated] = [10500.0, 10500.0, 10500.0]
+
+    line = initialize_array(DatasetType.input, ComponentType.line, 2)
+    line[AttributeType.id] = [9, 10]
+    line[AttributeType.from_node] = [7, 7]
+    line[AttributeType.to_node] = [2, 1]
+    line[AttributeType.from_status] = [1, 1]
+    line[AttributeType.to_status] = [1, 1]
+    line[AttributeType.r1] = [0.00396133, 0.32598809]
+    line[AttributeType.x1] = [4.53865336e-05, 1.34716591e-02]
+    line[AttributeType.c1] = [0.0, 0.0]
+    line[AttributeType.tan1] = [0.0, 0.0]
+    line[AttributeType.i_n] = [303.91942029, 210.06857453]
+
+    link = initialize_array(DatasetType.input, ComponentType.link, 0)
+    transformer = initialize_array(DatasetType.input, ComponentType.transformer, 0)
+    three_winding_transformer = initialize_array(DatasetType.input, ComponentType.three_winding_transformer, 0)
+
+    sym_load = initialize_array(DatasetType.input, ComponentType.sym_load, 2)
+    sym_load[AttributeType.id] = [5, 6]
+    sym_load[AttributeType.node] = [1, 2]
+    sym_load[AttributeType.status] = [1, 1]
+    sym_load[AttributeType.type] = [0, 0]
+    sym_load[AttributeType.p_specified] = [-287484.0, 26558.0]
+    sym_load[AttributeType.q_specified] = [40640.0, 28148.0]
+
+    sym_gen = initialize_array(DatasetType.input, ComponentType.sym_gen, 0)
+
+    source = initialize_array(DatasetType.input, ComponentType.source, 1)
+    source[AttributeType.id] = [8]
+    source[AttributeType.node] = [7]
+    source[AttributeType.status] = [1]
+    source[AttributeType.u_ref] = [1.0]
+
+    transformer_tap_regulator = initialize_array(DatasetType.input, ComponentType.transformer_tap_regulator, 0)
+    sym_power_sensor = initialize_array(DatasetType.input, ComponentType.sym_power_sensor, 0)
+    sym_voltage_sensor = initialize_array(DatasetType.input, ComponentType.sym_voltage_sensor, 0)
+    asym_voltage_sensor = initialize_array(DatasetType.input, ComponentType.asym_voltage_sensor, 0)
 
     return {
-        "node": node,
-        "line": line,
-        "link": link,
-        "transformer": transformer,
-        "three_winding_transformer": three_winding_transformer,
-        "sym_load": sym_load,
-        "sym_gen": sym_gen,
-        "source": source,
-        "transformer_tap_regulator": transformer_tap_regulator,
-        "sym_power_sensor": sym_power_sensor,
-        "sym_voltage_sensor": sym_voltage_sensor,
-        "asym_voltage_sensor": asym_voltage_sensor,
+        ComponentType.node: node,
+        ComponentType.line: line,
+        ComponentType.link: link,
+        ComponentType.transformer: transformer,
+        ComponentType.three_winding_transformer: three_winding_transformer,
+        ComponentType.sym_load: sym_load,
+        ComponentType.sym_gen: sym_gen,
+        ComponentType.source: source,
+        ComponentType.transformer_tap_regulator: transformer_tap_regulator,
+        ComponentType.sym_power_sensor: sym_power_sensor,
+        ComponentType.sym_voltage_sensor: sym_voltage_sensor,
+        ComponentType.asym_voltage_sensor: asym_voltage_sensor,
     }

tests/integration/visualizer_tests.py

@@ -11,7 +11,7 @@
 from dataclasses import dataclass
 
 import numpy as np
-from power_grid_model import ComponentType, DatasetType, initialize_array
+from power_grid_model import AttributeType, ComponentType, DatasetType, initialize_array
 
 from power_grid_model_ds import Grid
 from power_grid_model_ds._core.model.dtypes.typing import NDArray3
@@ -99,21 +99,21 @@ def visualize_grid_with_all_open_types():
 def visualize_grid_with_batch_data():
     grid = get_coordinated_three_phase_grid()
     lines_update = initialize_array(DatasetType.update, ComponentType.line, (4, 2))
-    lines_update["id"] = grid.line.id[:2]
-    lines_update["from_status"][:] = 1
-    lines_update["from_status"][1, 0] = 0
+    lines_update[AttributeType.id] = grid.line.id[:2]
+    lines_update[AttributeType.from_status][:] = 1
+    lines_update[AttributeType.from_status][1, 0] = 0
     update_data = {ComponentType.line: lines_update}
 
     nodes_output = initialize_array(DatasetType.asym_output, ComponentType.node, (4, 3))
-    nodes_output["id"] = grid.node.id[:3]
-    nodes_output["u"] = [
+    nodes_output[AttributeType.id] = grid.node.id[:3]
+    nodes_output[AttributeType.u] = [
         [[401.0, 402.0, 403.0], [404.0, 405.0, 406.0], [407.0, 408.0, 409.0]],
         [[410.0, 411.0, 412.0], [413.0, 414.0, 415.0], [416.0, 417.0, 418.0]],
         [[419.0, 420.0, 421.0], [422.0, 423.0, 424.0], [425.0, 426.0, 427.0]],
         [[428.0, 429.0, 430.0], [431.0, 432.0, 433.0], [434.0, 435.0, 436.0]],
     ]
-    nodes_output["energized"][:] = 1
-    nodes_output["energized"][1, 2] = 0
+    nodes_output[AttributeType.energized][:] = 1
+    nodes_output[AttributeType.energized][1, 2] = 0
     output_data = {ComponentType.node: nodes_output}
     visualize(grid=grid, update_data=update_data, output_data=output_data, debug=True)
 

tests/unit/visualizer/test_grid_utils.py

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: MPL-2.0
 import numpy as np
 import pytest
-from power_grid_model import ComponentType, DatasetType, initialize_array
+from power_grid_model import AttributeType, ComponentType, DatasetType, initialize_array
 
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 from power_grid_model_ds._core.model.constants import empty
@@ -85,39 +85,39 @@ def test_dynamic_grid_obj_from_grid(dataset_type):
 
 def test_get_attr_data_from_dataset():
     loads = initialize_array(DatasetType.update, ComponentType.sym_load, (4, 3))
-    loads["p_specified"] = [[21, 22, 23], [31, 32, 33], [41, 42, 43], [51, 52, 53]]
+    loads[AttributeType.p_specified] = [[21, 22, 23], [31, 32, 33], [41, 42, 43], [51, 52, 53]]
 
     # Normal case with matching pgm_id=12
-    loads["id"] = [[11, 12, 13], [11, 12, 13], [11, 12, 13], [11, 12, 13]]
+    loads[AttributeType.id] = [[11, 12, 13], [11, 12, 13], [11, 12, 13], [11, 12, 13]]
     batch_dataset = {ComponentType.sym_load: loads}
     x_actual, y_actual = get_attr_data_from_dataset(batch_dataset, ComponentType.sym_load, "p_specified", pgm_id=12)
     assert np.array_equal(y_actual, np.array([22, 32, 42, 52]))
     assert np.array_equal(x_actual, np.array([0, 1, 2, 3]))
 
     # Order changed
-    loads["id"] = [[11, 12, 13], [13, 11, 12], [11, 12, 13], [13, 11, 12]]
+    loads[AttributeType.id] = [[11, 12, 13], [13, 11, 12], [11, 12, 13], [13, 11, 12]]
     x_actual, y_actual = get_attr_data_from_dataset(batch_dataset, ComponentType.sym_load, "p_specified", pgm_id=12)
     assert np.array_equal(y_actual, np.array([22, 33, 42, 53]))
     assert np.array_equal(x_actual, np.array([0, 1, 2, 3]))
 
     # Missing scenario for pgm_id=12
-    loads["id"] = [[11, 12, 13], [14, 15, 16], [11, 12, 13], [14, 15, 16]]
+    loads[AttributeType.id] = [[11, 12, 13], [14, 15, 16], [11, 12, 13], [14, 15, 16]]
     x_actual, y_actual = get_attr_data_from_dataset(batch_dataset, ComponentType.sym_load, "p_specified", pgm_id=12)
     assert np.array_equal(y_actual, np.array([22, 42]))
     assert np.array_equal(x_actual, np.array([0, 2]))
 
 
 def test_get_attr_data_from_dataset_3ph():
     loads = initialize_array(DatasetType.update, ComponentType.asym_load, (4, 3))
-    loads["p_specified"] = [
+    loads[AttributeType.p_specified] = [
         [[21.1, 21.2, 21.3], [22.1, 22.2, 22.3], [23.1, 23.2, 23.3]],
         [[31.1, 31.2, 31.3], [32.1, 32.2, 32.3], [33.1, 33.2, 33.3]],
         [[41.1, 41.2, 41.3], [42.1, 42.2, 42.3], [43.1, 43.2, 43.3]],
         [[51.1, 51.2, 51.3], [52.1, 52.2, 52.3], [53.1, 53.2, 53.3]],
     ]
 
     # Normal case with matching pgm_id=12
-    loads["id"] = [[11, 12, 13], [11, 12, 13], [11, 12, 13], [11, 12, 13]]
+    loads[AttributeType.id] = [[11, 12, 13], [11, 12, 13], [11, 12, 13], [11, 12, 13]]
     batch_dataset = {ComponentType.asym_load: loads}
     x_actual, y_actual = get_attr_data_from_dataset(batch_dataset, ComponentType.asym_load, "p_specified", pgm_id=12)
     assert np.array_equal(