Feature/dgc 2055 get batch size function

.pre-commit-config.yaml

@@ -4,7 +4,7 @@
 
 repos:
   - repo: https://github.com/fsfe/reuse-tool
-    rev: v1.0.0
+    rev: v2.1.0
     hooks:
       - id: reuse
   - repo: https://github.com/pycqa/isort

docs/api_reference/python-api-reference.md

@@ -52,4 +52,6 @@ SPDX-License-Identifier: MPL-2.0
 .. autofunction:: power_grid_model.utils.msgpack_serialize_to_file
 .. autofunction:: power_grid_model.utils.import_json_data
 .. autofunction:: power_grid_model.utils.export_json_data
+.. autofunction:: power_grid_model._utils.get_and_verify_batch_sizes
+.. autofunction:: power_grid_model._utils.get_batch_size
 ```

pyproject.toml

@@ -83,6 +83,7 @@ target-version = ['py38']
 
 [tool.isort]
 profile = "black"
+combine_as_imports = true
 line_length = 120
 
 [tool.pylint]

src/power_grid_model/core/options.py

@@ -8,8 +8,7 @@
 """
 from typing import Any, Callable
 
-from power_grid_model.core.power_grid_core import OptionsPtr
-from power_grid_model.core.power_grid_core import power_grid_core as pgc
+from power_grid_model.core.power_grid_core import OptionsPtr, power_grid_core as pgc
 
 
 class OptionSetter:

src/power_grid_model/core/power_grid_dataset.py

@@ -19,8 +19,13 @@
     get_buffer_view,
 )
 from power_grid_model.core.error_handling import VALIDATOR_MSG, assert_no_error
-from power_grid_model.core.power_grid_core import ConstDatasetPtr, DatasetInfoPtr, MutableDatasetPtr, WritableDatasetPtr
-from power_grid_model.core.power_grid_core import power_grid_core as pgc
+from power_grid_model.core.power_grid_core import (
+    ConstDatasetPtr,
+    DatasetInfoPtr,
+    MutableDatasetPtr,
+    WritableDatasetPtr,
+    power_grid_core as pgc,
+)
 from power_grid_model.core.power_grid_meta import DatasetMetaData, power_grid_meta_data
 from power_grid_model.errors import PowerGridError
 

src/power_grid_model/core/power_grid_meta.py

@@ -12,8 +12,7 @@
 
 import numpy as np
 
-from power_grid_model.core.power_grid_core import AttributePtr, ComponentPtr, DatasetPtr
-from power_grid_model.core.power_grid_core import power_grid_core as pgc
+from power_grid_model.core.power_grid_core import AttributePtr, ComponentPtr, DatasetPtr, power_grid_core as pgc
 
 
 # constant enum for ctype

src/power_grid_model/core/power_grid_model.py

@@ -21,8 +21,7 @@
 from power_grid_model.core.error_handling import PowerGridBatchError, assert_no_error, handle_errors
 from power_grid_model.core.index_integer import IdNp, IdxNp
 from power_grid_model.core.options import Options
-from power_grid_model.core.power_grid_core import ConstDatasetPtr, IDPtr, IdxPtr, ModelPtr
-from power_grid_model.core.power_grid_core import power_grid_core as pgc
+from power_grid_model.core.power_grid_core import ConstDatasetPtr, IDPtr, IdxPtr, ModelPtr, power_grid_core as pgc
 from power_grid_model.enum import CalculationMethod, CalculationType, ShortCircuitVoltageScaling
 
 

src/power_grid_model/core/serialization.py

@@ -16,8 +16,13 @@
 
 from power_grid_model.core.error_handling import assert_no_error
 from power_grid_model.core.index_integer import IdxC
-from power_grid_model.core.power_grid_core import CharPtr, DeserializerPtr, SerializerPtr, WritableDatasetPtr
-from power_grid_model.core.power_grid_core import power_grid_core as pgc
+from power_grid_model.core.power_grid_core import (
+    CharPtr,
+    DeserializerPtr,
+    SerializerPtr,
+    WritableDatasetPtr,
+    power_grid_core as pgc,
+)
 from power_grid_model.core.power_grid_dataset import CConstDataset, CWritableDataset
 from power_grid_model.errors import PowerGridSerializationError
 

src/power_grid_model/utils.py

@@ -9,11 +9,14 @@
 import json
 import warnings
 from pathlib import Path
-from typing import Optional
-from typing import cast as cast_type
+from typing import Optional, cast as cast_type
 
 import numpy as np
 
+from power_grid_model._utils import (
+    get_and_verify_batch_sizes as _get_and_verify_batch_sizes,
+    get_batch_size as _get_batch_size,
+)
 from power_grid_model.core.power_grid_dataset import get_dataset_type
 from power_grid_model.core.serialization import (  # pylint: disable=unused-import
     json_deserialize,
@@ -54,6 +57,35 @@ def _get_component_scenario(component_scenarios: BatchArray) -> np.ndarray:
     return {component: _get_component_scenario(component_data) for component, component_data in dataset.items()}
 
 
+def get_dataset_batch_size(dataset: BatchDataset) -> int:
+    """
+    Get the number of scenarios in the batch dataset.
+
+    Args:
+        dataset: the batch dataset
+
+    Raises:
+        ValueError: if the batch dataset is inconsistent.
+
+    Returns:
+        The size of the batch dataset. Making use of existing _utils function.
+    """
+    return _get_and_verify_batch_sizes(dataset)
+
+
+def get_component_batch_size(data_array: BatchArray) -> int:
+    """
+    Determine the number of batches and verify the data structure
+
+    Args:
+        data_array: a batch array for power-grid-model
+
+    Returns:
+        The number of batches in data_array
+    """
+    return _get_batch_size(data_array)
+
+
 def json_deserialize_from_file(file_path: Path) -> Dataset:
     """
     Load and deserialize a JSON file to a new dataset.

tests/unit/test_utils.py

@@ -3,14 +3,19 @@
 # SPDX-License-Identifier: MPL-2.0
 
 from pathlib import Path
+from typing import Dict
 from unittest.mock import MagicMock, mock_open, patch
 
 import numpy as np
 import pytest
 
+from power_grid_model import LoadGenType, initialize_array
+from power_grid_model.core.power_grid_meta import power_grid_meta_data
 from power_grid_model.data_types import Dataset
 from power_grid_model.utils import (
     export_json_data,
+    get_component_batch_size,
+    get_dataset_batch_size,
     get_dataset_scenario,
     json_deserialize_from_file,
     json_serialize_to_file,
@@ -41,6 +46,79 @@ def test_get_dataset_scenario():
         get_dataset_scenario(data, 2)
 
 
+def test_get_data_set_batch_size():
+    line = initialize_array("update", "line", (3, 2))
+    line["id"] = [[5, 6], [6, 7], [7, 5]]
+    line["from_status"] = [[1, 1], [1, 1], [1, 1]]
+
+    asym_load = initialize_array("update", "asym_load", (3, 2))
+    asym_load["id"] = [[9, 10], [9, 10], [9, 10]]
+
+    batch_data = {"line": line, "asym_load": asym_load}
+
+    assert get_dataset_batch_size(batch_data) == 3
+
+
+def test_get_dataset_batch_size_sparse():
+    data = {
+        "node": {
+            "data": np.zeros(shape=3, dtype=power_grid_meta_data["input"]["node"]),
+            "indptr": np.array([0, 2, 3, 3]),
+        },
+        "sym_load": {
+            "data": np.zeros(shape=2, dtype=power_grid_meta_data["input"]["sym_load"]),
+            "indptr": np.array([0, 0, 1, 2]),
+        },
+        "asym_load": {
+            "data": np.zeros(shape=4, dtype=power_grid_meta_data["input"]["asym_load"]),
+            "indptr": np.array([0, 2, 3, 4]),
+        },
+    }
+
+    assert get_dataset_batch_size(data) == 3
+
+
+def test_get_dataset_batch_size_mixed():
+    line = initialize_array("update", "line", (3, 2))
+    line["id"] = [[5, 6], [6, 7], [7, 5]]
+    line["from_status"] = [[1, 1], [1, 1], [1, 1]]
+
+    asym_load = initialize_array("update", "asym_load", (2, 2))
+    asym_load["id"] = [[9, 10], [9, 10]]
+
+    data_dense = {"line": line, "asym_load": asym_load}
+    data_sparse = {
+        "node": {
+            "data": np.zeros(shape=3, dtype=power_grid_meta_data["input"]["node"]),
+            "indptr": np.array([0, 2, 3, 3, 5]),
+        },
+        "sym_load": {
+            "data": np.zeros(shape=2, dtype=power_grid_meta_data["input"]["sym_load"]),
+            "indptr": np.array([0, 0, 1, 2]),
+        },
+        "asym_load": {
+            "data": np.zeros(shape=4, dtype=power_grid_meta_data["input"]["asym_load"]),
+            "indptr": np.array([0, 2, 3]),
+        },
+    }
+    with pytest.raises(ValueError):
+        get_dataset_batch_size(data_dense)
+    with pytest.raises(ValueError):
+        get_dataset_batch_size(data_sparse)
+
+
+def test_get_component_batch_size():
+    asym_load = initialize_array("update", "asym_load", (3, 2))
+    asym_load["id"] = [[9, 10], [9, 10], [9, 10]]
+
+    sym_load = {
+        "data": np.zeros(shape=2, dtype=power_grid_meta_data["input"]["sym_load"]),
+        "indptr": np.array([0, 0, 1, 2]),
+    }
+    assert get_component_batch_size(asym_load) == 3
+    assert get_component_batch_size(sym_load) == 3
+
+
 @patch("builtins.open", new_callable=mock_open)
 @patch("power_grid_model.utils.json_deserialize")
 def test_json_deserialize_from_file(deserialize_mock: MagicMock, open_mock: MagicMock):