Arrow example
This example illustrates how to convert Arrow data that contains power-grid-model data to NumPy structured arrays, which the power-grid-model requests.
It is by no means intended to provide complete documentation on the topic, but only to show how such conversions could be done.
This example uses pyarrow.RecordBatch to demonstrate zero-copy operations. The user can choose a pyarrow.Table or other structures based on the requirement.
NOTE: To run this example, the optional examples dependencies are required:
pip install .[examples]
%%capture cap --no-stderr
from collections.abc import Iterable
import numpy as np
import pandas as pd
import pyarrow as pa
from IPython.display import display
from power_grid_model import (
AttributeType,
ComponentAttributeFilterOptions,
ComponentType,
DatasetType,
PowerGridModel,
power_grid_meta_data,
)
from power_grid_model.data_types import SingleColumnarData
Model
For clarity, a simple network is created. More complex cases work similarly and can be found in the other examples:
node 1 ---- line 4 ---- node 2 ----line 5 ---- node 3
| | |
source 6 sym_load 7 sym_load 8
Single symmetric calculations
Construct the input data for the model and construct the actual model.
Arrow uses a columnar data format while the power-grid-model offers support for both row based and columnar data format. Because of this, the columnar data format of power-grid-model provides a zero-copy interface for Arrow data. This differs from the row-based data format, for which conversions always require a copy.
List the power-grid-model data types
See which attributes exist for a given component and which data types are used
node_input_dtype = power_grid_meta_data[DatasetType.input][ComponentType.node].dtype
line_input_dtype = power_grid_meta_data[DatasetType.input][ComponentType.line].dtype
source_input_dtype = power_grid_meta_data[DatasetType.input][ComponentType.source].dtype
asym_load_input_dtype = power_grid_meta_data[DatasetType.input][ComponentType.asym_load].dtype
print("node:", node_input_dtype)
print("line:", line_input_dtype)
print("source:", source_input_dtype)
print("asym_load:", asym_load_input_dtype)
node: {'names': [id, u_rated], 'formats': ['<i4', '<f8'], 'offsets': [0, 8], 'itemsize': 16, 'aligned': True}
line: {'names': [id, from_node, to_node, from_status, to_status, r1, x1, c1, tan1, r0, x0, c0, tan0, i_n], 'formats': ['<i4', '<i4', '<i4', 'i1', 'i1', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8'], 'offsets': [0, 4, 8, 12, 13, 16, 24, 32, 40, 48, 56, 64, 72, 80], 'itemsize': 88, 'aligned': True}
source: {'names': [id, node, status, u_ref, u_ref_angle, sk, rx_ratio, z01_ratio], 'formats': ['<i4', '<i4', 'i1', '<f8', '<f8', '<f8', '<f8', '<f8'], 'offsets': [0, 4, 8, 16, 24, 32, 40, 48], 'itemsize': 56, 'aligned': True}
asym_load: {'names': [id, node, status, type, p_specified, q_specified], 'formats': ['<i4', '<i4', 'i1', 'i1', ('<f8', (3,)), ('<f8', (3,))], 'offsets': [0, 4, 8, 9, 16, 40], 'itemsize': 64, 'aligned': True}
The primitive types of each attribute in the arrow tables need to match to make the operation efficient.
Zero-copy conversion is not guaranteed if the data types provided via the PGM via power_grid_meta_data are not used.
Note that the asymmetric type of attribute in power-grid-model has a shape of (3,) along with a specific type. These represent the 3 phases of electrical system.
Hence, special care is required when handling asymmetric attributes.
In this example, we use the respective primitive types for the symmetrical attributes and a FixedSizeListArray of the primitive types with length 3 for asymmetrical attributes. This results in them being stored as contiguous memory which would enable zero-copy conversion. Other possible workarounds to this are possible, but are beyond the scope of this example.
Creating a Schema
We can make the task of assigning types easier by creating a schema based on the DatasetType and ComponentType directly from power_grid_meta_data.
They can then directly be used to construct the RecordBatch.
def pgm_schema(
dataset_type: DatasetType, component_type: ComponentType, attributes: Iterable[AttributeType] | None = None
):
schemas = []
component_dtype = power_grid_meta_data[dataset_type][component_type].dtype
for meta_attribute, (dtype, _) in component_dtype.fields.items():
if attributes is not None and meta_attribute not in attributes:
continue
# The asymmetric attributes are stored as a fixed list array of 3 elements, when dtype.shape is (3,).
pa_dtype = pa.list_(pa.from_numpy_dtype(dtype.base), 3) if dtype.shape == (3,) else pa.from_numpy_dtype(dtype)
schemas.append((meta_attribute, pa_dtype))
return pa.schema(schemas)
print("-------node schema-------")
print(pgm_schema(DatasetType.input, ComponentType.node))
print("-------asym load schema-------")
print(pgm_schema(DatasetType.input, ComponentType.asym_load))
-------node schema-------
id: int32
u_rated: double
-------asym load schema-------
id: int32
node: int32
status: int8
type: int8
p_specified: fixed_size_list<item: double>[3]
child 0, item: double
q_specified: fixed_size_list<item: double>[3]
child 0, item: double
Create the grid using Arrow tables
The power-grid-model documentation on Components provides documentation on which components are required and which ones are optional.
Construct the Arrow data as a table with the correct headers and data types. The creation of arrays and combining it in a RecordBatch as well as the method of initializing that RecordBatch is up to the user.
nodes_schema = pgm_schema(DatasetType.input, ComponentType.node)
nodes = pa.record_batch(
[
pa.array([1, 2, 3], type=nodes_schema.field(AttributeType.id).type),
pa.array([10500.0, 10500.0, 10500.0], type=nodes_schema.field(AttributeType.u_rated).type),
],
names=("id", "u_rated"),
)
lines = pa.record_batch(
{
AttributeType.id: [4, 5],
AttributeType.from_node: [1, 2],
AttributeType.to_node: [2, 3],
AttributeType.from_status: [1, 1],
AttributeType.to_status: [1, 1],
AttributeType.r1: [0.11, 0.15],
AttributeType.x1: [0.12, 0.16],
AttributeType.c1: [4.1e-05, 5.4e-05],
AttributeType.tan1: [0.1, 0.1],
AttributeType.r0: [0.01, 0.05],
AttributeType.x0: [0.22, 0.06],
AttributeType.c0: [4.1e-05, 5.4e-05],
AttributeType.tan0: [0.4, 0.1],
},
schema=pgm_schema(
DatasetType.input,
ComponentType.line,
[
AttributeType.id,
AttributeType.from_node,
AttributeType.to_node,
AttributeType.from_status,
AttributeType.to_status,
AttributeType.r1,
AttributeType.x1,
AttributeType.c1,
AttributeType.tan1,
AttributeType.r0,
AttributeType.x0,
AttributeType.c0,
AttributeType.tan0,
],
),
)
sources = pa.record_batch(
{AttributeType.id: [6], AttributeType.node: [1], AttributeType.status: [1], AttributeType.u_ref: [1.0]},
schema=pgm_schema(
DatasetType.input,
ComponentType.source,
[AttributeType.id, AttributeType.node, AttributeType.status, AttributeType.u_ref],
),
)
sym_loads = pa.record_batch(
{
AttributeType.id: [7, 8],
AttributeType.node: [2, 3],
AttributeType.status: [1, 1],
AttributeType.type: [0, 0],
AttributeType.p_specified: [1.0, 2.0],
AttributeType.q_specified: [0.5, 1.5],
},
schema=pgm_schema(
DatasetType.input,
ComponentType.sym_load,
[
AttributeType.id,
AttributeType.node,
AttributeType.status,
AttributeType.type,
AttributeType.p_specified,
AttributeType.q_specified,
],
),
)
nodes
# the record batches of the other components can be printed similarly
pyarrow.RecordBatch
id: int32
u_rated: double
----
id: [1,2,3]
u_rated: [10500,10500,10500]
Convert the Arrow data to power-grid-model input data
The Arrow RecordBatch or Table can then be converted to row based data or columnar data.
Converting Arrow data to columnar NumPy arrays is recommended to leverage the columnar nature of Arrow data.
This conversion can be done with zero-copy operations.
A similar approach can be adopted by the user to convert to row based data instead.
def arrow_to_numpy(data: pa.RecordBatch, dataset_type: DatasetType, component_type: ComponentType) -> np.ndarray:
"""Convert Arrow data to NumPy data."""
result = {}
result_dtype = power_grid_meta_data[dataset_type][component_type].dtype
for name, column in zip(data.column_names, data.columns):
# The use of zero_copy_only=True and assert statement is to verify if no copies are made.
# They are not mandatory for a zero-copy conversion.
column_data = column.to_numpy(zero_copy_only=True)
if column_data.dtype != result_dtype[name]:
raise TypeError(
f"Data type mismatch for column '{name}': expected {result_dtype[name]}, got {column_data.dtype}"
)
result[name] = column_data.astype(dtype=result_dtype[name], copy=False)
return result
node_input = arrow_to_numpy(nodes, DatasetType.input, ComponentType.node)
line_input = arrow_to_numpy(lines, DatasetType.input, ComponentType.line)
source_input = arrow_to_numpy(sources, DatasetType.input, ComponentType.source)
sym_load_input = arrow_to_numpy(sym_loads, DatasetType.input, ComponentType.sym_load)
node_input
{'id': array([1, 2, 3], dtype=int32),
'u_rated': array([10500., 10500., 10500.])}
Construct the complete input data structure
input_data = {
ComponentType.node: node_input,
ComponentType.line: line_input,
ComponentType.source: source_input,
ComponentType.sym_load: sym_load_input,
}
# Optional: validate the input data
from power_grid_model.validation import validate_input_data
validate_input_data(input_data)
Use the power-grid-model
The output_component_types argument is set to ComponentAttributeFilterOptions.relevant to given out columnar data.
For more extensive examples, visit the power-grid-model documentation.
# construct the model
model = PowerGridModel(input_data=input_data, system_frequency=50)
# run the calculation
sym_result = model.calculate_power_flow(output_component_types=ComponentAttributeFilterOptions.relevant)
# use pandas to tabulate and display the results
sym_node_result = sym_result[ComponentType.node]
display(pd.DataFrame(sym_node_result))
| id | energized | u_pu | u | u_angle | p | q | |
|---|---|---|---|---|---|---|---|
| 0 | 1 | 1 | 1.000325 | 10503.410670 | -0.000067 | 338777.246279 | -3.299419e+06 |
| 1 | 2 | 1 | 1.002879 | 10530.228073 | -0.002932 | -1.000000 | -5.000001e-01 |
| 2 | 3 | 1 | 1.004113 | 10543.184974 | -0.004342 | -2.000000 | -1.500000e+00 |
Convert the symmetric result to Arrow format
Converting symmetrical results is straightforward by using schema from Creating Schema Using types other than the ones from this schema might make a copy of the data.
pa_sym_node_result = pa.record_batch(
sym_node_result, schema=pgm_schema(DatasetType.sym_output, ComponentType.node, sym_node_result.keys())
)
pa_sym_node_result
pyarrow.RecordBatch
id: int32
energized: int8
u_pu: double
u: double
u_angle: double
p: double
q: double
----
id: [1,2,3]
energized: [1,1,1]
u_pu: [1.000324825742982,1.0028788641128945,1.004112854674026]
u: [10503.410670301311,10530.228073185392,10543.184974077272]
u_angle: [-0.00006651843181518038,-0.0029317915196012487,-0.004341587216862092]
p: [338777.2462788448,-1.0000002693184182,-1.9999998867105226]
q: [-3299418.661306348,-0.5000000701801947,-1.4999998507078594]
Single asymmetric calculations
Asymmetric calculations have a tuple of values for some of the attributes and are not easily convertible to record batches. Instead, one can have a look at the individual components of those attributes and/or flatten the arrays to access all components.
Asymmetric input
To illustrate the conversion, let’s consider a similar grid but with asymmetric loads.
node 1 ---- line 4 ---- node 2 ----line 5 ---- node 3
| | |
source 6 asym_load 7 asym_load 8
asym_loads_dict = {
AttributeType.id: [7, 8],
AttributeType.node: [2, 3],
AttributeType.status: [1, 1],
AttributeType.type: [0, 0],
AttributeType.p_specified: [[1.0, 1.0e-2, 1.1e-2], [2.0, 2.5, 4.5e2]],
AttributeType.q_specified: [[0.5, 1.5e3, 0.1], [1.5, 2.5, 1.5e3]],
}
asym_loads = pa.record_batch(
{
AttributeType.id: [7, 8],
AttributeType.node: [2, 3],
AttributeType.status: [1, 1],
AttributeType.type: [0, 0],
AttributeType.p_specified: [[1.0, 1.0e-2, 1.1e-2], [2.0, 2.5, 4.5e2]],
AttributeType.q_specified: [[0.5, 1.5e3, 0.1], [1.5, 2.5, 1.5e3]],
},
schema=pgm_schema(
DatasetType.input,
ComponentType.asym_load,
[
AttributeType.id,
AttributeType.node,
AttributeType.status,
AttributeType.type,
AttributeType.p_specified,
AttributeType.q_specified,
],
),
)
asym_loads
pyarrow.RecordBatch
id: int32
node: int32
status: int8
type: int8
p_specified: fixed_size_list<item: double>[3]
child 0, item: double
q_specified: fixed_size_list<item: double>[3]
child 0, item: double
----
id: [7,8]
node: [2,3]
status: [1,1]
type: [0,0]
p_specified: [[1,0.01,0.011],[2,2.5,450]]
q_specified: [[0.5,1500,0.1],[1.5,2.5,1500]]
def arrow_to_numpy_asym(data: pa.RecordBatch, dataset_type: DatasetType, component_type: ComponentType) -> np.ndarray:
"""Convert asymmetric Arrow data to NumPy data.
This function is similar to the arrow_to_numpy function, but also supports asymmetric data."""
result = {}
result_dtype = power_grid_meta_data[dataset_type][component_type].dtype
for name in result_dtype.names:
if name not in data.column_names:
continue
dtype = result_dtype[name]
if len(dtype.shape) == 0:
column_data = data.column(name).to_numpy(zero_copy_only=True)
else:
column_data = data.column(name).flatten().to_numpy(zero_copy_only=True).reshape(-1, 3)
if column_data.dtype != dtype.base:
raise TypeError(f"Data type mismatch for column '{name}': expected {dtype.base}, got {column_data.dtype}")
result[name] = column_data.astype(dtype=dtype.base, copy=False)
return result
asym_load_input = arrow_to_numpy_asym(asym_loads, DatasetType.input, ComponentType.asym_load)
asym_load_input
{id: array([7, 8], dtype=int32),
node: array([2, 3], dtype=int32),
status: array([1, 1], dtype=int8),
type: array([0, 0], dtype=int8),
p_specified: array([[1.0e+00, 1.0e-02, 1.1e-02],
[2.0e+00, 2.5e+00, 4.5e+02]]),
q_specified: array([[5.0e-01, 1.5e+03, 1.0e-01],
[1.5e+00, 2.5e+00, 1.5e+03]])}
Use the power-grid-model in asymmetric calculations
asym_input_data = {
ComponentType.node: node_input,
ComponentType.line: line_input,
ComponentType.source: source_input,
ComponentType.asym_load: asym_load_input,
}
validate_input_data(asym_input_data, symmetric=False)
# construct the model
asym_model = PowerGridModel(input_data=asym_input_data, system_frequency=50)
# run the calculation
asym_result = asym_model.calculate_power_flow(
symmetric=False, output_component_types=ComponentAttributeFilterOptions.relevant
)
# use pandas to display the results, but beware the data types
display(pd.DataFrame(asym_result[ComponentType.node][AttributeType.u_angle]))
| 0 | 1 | 2 | |
|---|---|---|---|
| 0 | -0.000067 | -2.094462 | 2.094328 |
| 1 | -0.002930 | -2.097322 | 2.091464 |
| 2 | -0.004338 | -2.098733 | 2.090057 |
Convert asymmetric power-grid-model output data to Arrow output data
ARRAY_2D = 2
ASYM_DATA = 3
def numpy_columnar_to_arrow(
data: SingleColumnarData, dataset_type: DatasetType, component_type: ComponentType
) -> pa.RecordBatch:
"""Convert NumPy data to Arrow data."""
component_pgm_schema = pgm_schema(dataset_type, component_type, data.keys())
pa_columns = {}
for attribute, column_data in data.items():
primitive_type = component_pgm_schema.field(attribute).type
if column_data.ndim == ARRAY_2D and column_data.shape[1] == ASYM_DATA:
pa_columns[attribute] = pa.FixedSizeListArray.from_arrays(column_data.flatten(), type=primitive_type)
else:
pa_columns[attribute] = pa.array(column_data, type=primitive_type)
return pa.record_batch(pa_columns, component_pgm_schema)
pa_asym_node_result = numpy_columnar_to_arrow(
asym_result[ComponentType.node], DatasetType.asym_output, ComponentType.node
)
pa_asym_node_result
pyarrow.RecordBatch
id: int32
energized: int8
u_pu: fixed_size_list<item: double>[3]
child 0, item: double
u: fixed_size_list<item: double>[3]
child 0, item: double
u_angle: fixed_size_list<item: double>[3]
child 0, item: double
p: fixed_size_list<item: double>[3]
child 0, item: double
q: fixed_size_list<item: double>[3]
child 0, item: double
----
id: [1,2,3]
energized: [1,1,1]
u_pu: [[1.0003248257977395,1.0003243769486854,1.00032436416241],[1.0028803762176162,1.0028710993140408,1.002873078902153],[1.0041143008174032,1.004103358307718,1.0041004935738544]]
u: [[6064.146978239599,6064.144257236815,6064.1441797241405],[6079.639179329455,6079.582941090302,6079.594941705461],[6087.119449677845,6087.053114238265,6087.03574771216]]
u_angle: [[-0.00006651848125694056,-2.094461573665813,2.09432849798745],[-0.002929883186483061,-2.0973219974462594,2.0914640024381836],[-0.0043376855072092685,-2.098732840554144,2.0900574062078014]]
p: [[112925.89463800077,112918.1351708998,113364.09104539294],[-0.9999999600705957,-0.009999974119601074,-0.010999946473583844],[-1.9999999924219376,-2.500000047715548,-450.0000000248862]]
q: [[-1099806.4185887817,-1098301.0302391544,-1098302.7942318914],[-0.5000000935833835,-1499.9999999519298,-0.09999992365339339],[-1.500000040797027,-2.499999950281157,-1500.0000000267694]]
Batch data
power-grid-model supports batch calculations by providing an update_data argument, as shown in this example.
Both the update_data and the output result are similar to the input_data and output data in the above, except that they have another dimension representing the batch index: the first index in the NumPy structured arrays.
This extra index can be represented in Arrow using a RecordBatch or using any other multi-index data format.