feat(load): add arrow endpoints

antarest/core/utils/string.py

@@ -18,3 +18,7 @@ def to_pascal_case(value: str) -> str:
 def to_camel_case(value: str) -> str:
     v = to_pascal_case(value)
     return v[0].lower() + v[1:] if len(v) > 0 else ""
+
+
+def to_kebab_case(string: str) -> str:
+    return string.replace("_", "-")

antarest/study/business/load_management.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2024, RTE (https://www.rte-france.com)
+#
+# See AUTHORS.txt
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+#
+# SPDX-License-Identifier: MPL-2.0
+#
+# This file is part of the Antares project.
+from io import BytesIO
+from typing import cast
+
+import pandas as pd
+from starlette.responses import Response
+
+from antarest.study.business.model.load_model import LoadDTO
+from antarest.study.model import Study
+from antarest.study.storage.rawstudy.model.filesystem.matrix.input_series_matrix import InputSeriesMatrix
+from antarest.study.storage.storage_service import StudyStorageService
+
+LOAD_PATH = "input/load/series/load_{area_id}"
+
+
+class LoadManager:
+    def __init__(self, storage_service: StudyStorageService) -> None:
+        self.storage_service = storage_service
+
+    def get_load_matrix(self, study: Study, area_id: str) -> Response:
+        file_study = self.storage_service.get_storage(study).get_raw(study)
+
+        load_path = LOAD_PATH.format(area_id=area_id).split("/")
+        node = file_study.tree.get_node(load_path)
+
+        if not isinstance(node, InputSeriesMatrix):
+            return Response(content="Invalid node type", status_code=400)
+
+        matrix_data = InputSeriesMatrix.parse(node, return_dataframe=True)
+
+        matrix_df = cast(pd.DataFrame, matrix_data)
+        matrix_df.columns = matrix_df.columns.map(str)
+        buffer = BytesIO()
+        matrix_df.to_feather(buffer, compression="uncompressed")
+        return Response(content=buffer.getvalue(), media_type="application/vnd.apache.arrow.file")
+
+    def update_load_matrix(self, study: Study, area_id: str, load_dto: LoadDTO) -> LoadDTO:
+        load_properties = load_dto.to_properties()
+        load_path = LOAD_PATH.format(area_id=area_id).split("/")
+
+        file_study = self.storage_service.get_storage(study).get_raw(study)
+
+        df = pd.read_feather(BytesIO(load_properties.matrix))
+
+        if df.shape[1] != 2:
+            pass
+
+        file_study.tree.save(load_properties.matrix, load_path)
+
+        return load_dto

antarest/study/business/model/load_model.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2024, RTE (https://www.rte-france.com)
+#
+# See AUTHORS.txt
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+#
+# SPDX-License-Identifier: MPL-2.0
+#
+# This file is part of the Antares project.
+from pydantic import ConfigDict
+
+from antarest.core.serialization import AntaresBaseModel
+from antarest.core.utils.string import to_camel_case, to_kebab_case
+
+
+class LoadDTO(AntaresBaseModel):
+    matrix: bytes
+
+    model_config = ConfigDict(alias_generator=to_camel_case, extra="forbid")
+
+    def to_properties(self) -> "LoadProperties":
+        return LoadProperties(matrix=self.matrix)
+
+
+class LoadProperties(AntaresBaseModel):
+    matrix: bytes
+
+    model_config = ConfigDict(alias_generator=to_kebab_case, extra="forbid")
+
+    def to_dto(self) -> LoadDTO:
+        return LoadDTO(matrix=self.matrix)

antarest/study/model.py

@@ -16,6 +16,7 @@
 import typing as t
 import uuid
 from datetime import datetime, timedelta
+from enum import StrEnum
 from pathlib import Path
 
 from antares.study.version import StudyVersion
@@ -525,6 +526,11 @@ def suffix(self) -> str:
         return mapping[self]
 
 
+class MatrixFormat(StrEnum):
+    JSON = "json"
+    ARROW = "arrow"
+
+
 class StudyDownloadDTO(AntaresBaseModel):
     """
     DTO used to download outputs

antarest/study/service.py

@@ -87,6 +87,7 @@
 from antarest.study.business.district_manager import DistrictManager
 from antarest.study.business.general_management import GeneralManager
 from antarest.study.business.link_management import LinkInfoDTO, LinkManager
+from antarest.study.business.load_management import LoadManager
 from antarest.study.business.matrix_management import MatrixManager, MatrixManagerError
 from antarest.study.business.optimization_management import OptimizationManager
 from antarest.study.business.playlist_management import PlaylistManager
@@ -366,6 +367,7 @@ def __init__(
         self.adequacy_patch_manager = AdequacyPatchManager(self.storage_service)
         self.advanced_parameters_manager = AdvancedParamsManager(self.storage_service)
         self.hydro_manager = HydroManager(self.storage_service)
+        self.load_manager = LoadManager(self.storage_service)
         self.allocation_manager = AllocationManager(self.storage_service)
         self.properties_manager = PropertiesManager(self.storage_service)
         self.renewable_manager = RenewableManager(self.storage_service)

antarest/study/web/study_data_blueprint.py

@@ -17,7 +17,7 @@
 
 import typing_extensions as te
 from fastapi import APIRouter, Body, Depends, Query
-from starlette.responses import RedirectResponse
+from starlette.responses import RedirectResponse, Response
 
 from antarest.core.config import Config
 from antarest.core.jwt import JWTUser
@@ -69,13 +69,14 @@
 from antarest.study.business.district_manager import DistrictCreationDTO, DistrictInfoDTO, DistrictUpdateDTO
 from antarest.study.business.general_management import GeneralFormFields
 from antarest.study.business.link_management import LinkInfoDTO
+from antarest.study.business.model.load_model import LoadDTO
 from antarest.study.business.optimization_management import OptimizationFormFields
 from antarest.study.business.playlist_management import PlaylistColumns
 from antarest.study.business.scenario_builder_management import Rulesets, ScenarioType
 from antarest.study.business.table_mode_management import TableDataDTO, TableModeType
 from antarest.study.business.thematic_trimming_field_infos import ThematicTrimmingFormFields
 from antarest.study.business.timeseries_config_management import TSFormFields
-from antarest.study.model import PatchArea, PatchCluster
+from antarest.study.model import MatrixFormat, PatchArea, PatchCluster
 from antarest.study.service import StudyService
 from antarest.study.storage.rawstudy.model.filesystem.config.binding_constraint import (
     BindingConstraintFrequency,
@@ -524,6 +525,41 @@ def update_inflow_structure(
         study = study_service.check_study_access(uuid, StudyPermissionType.WRITE, params)
         study_service.hydro_manager.update_inflow_structure(study, area_id, values)
 
+    @bp.get(
+        "/studies/{uuid}/{area_id}/load/series",
+        tags=[APITag.study_data],
+        summary="Get load series data",
+    )
+    def get_load_series(
+        uuid: str,
+        area_id: str,
+        current_user: JWTUser = Depends(auth.get_current_user),
+    ) -> Response:
+        """Return the load matrix in ARROW format."""
+        logger.info(
+            msg=f"Getting load series data for area {area_id} of study {uuid}",
+            extra={"user": current_user.id},
+        )
+        params = RequestParameters(user=current_user)
+        study = study_service.check_study_access(uuid, StudyPermissionType.READ, params)
+        return study_service.load_manager.get_load_matrix(study, area_id)
+
+    @bp.put(
+        "/studies/{uuid}/{area_id}/load/series",
+        tags=[APITag.study_data],
+        summary="Update load series data",
+    )
+    def update_load_series(
+        uuid: str,
+        area_id: str,
+        load_dto: LoadDTO,
+        current_user: JWTUser = Depends(auth.get_current_user),
+    ) -> LoadDTO:
+        """Replace the existing load matrix."""
+        params = RequestParameters(user=current_user)
+        study = study_service.check_study_access(uuid, StudyPermissionType.WRITE, params)
+        return study_service.load_manager.update_load_matrix(study, area_id, load_dto)
+
     @bp.put(
         "/studies/{uuid}/matrix",
         tags=[APITag.study_data],

requirements.txt

@@ -38,6 +38,7 @@ pandas~=2.2.3
 paramiko~=3.4.1
 plyer~=2.0.0
 psycopg2-binary~=2.9.9
+pyarrow~=18.0.0
 py7zr~=0.20.6
 python-json-logger~=2.0.7
 PyYAML~=5.3.1

tests/integration/study_data_blueprint/test_load.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2024, RTE (https://www.rte-france.com)
+#
+# See AUTHORS.txt
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+#
+# SPDX-License-Identifier: MPL-2.0
+#
+# This file is part of the Antares project.
+from io import BytesIO
+
+import pandas as pd
+import pytest
+from starlette.testclient import TestClient
+
+from tests.integration.prepare_proxy import PreparerProxy
+
+
+@pytest.mark.unit_test
+class TestLink:
+    @pytest.mark.parametrize("study_type", ["raw", "variant"])
+    def test_load(self, client: TestClient, user_access_token: str, study_type: str) -> None:
+        client.headers = {"Authorization": f"Bearer {user_access_token}"}  # type: ignore
+
+        preparer = PreparerProxy(client, user_access_token)
+        study_id = preparer.create_study("foo", version=880)
+
+        if study_type == "variant":
+            study_id = preparer.create_variant(study_id, name="Variant 1")
+
+        area1_id = preparer.create_area(study_id, name="Area 1")["id"]
+
+        # Test simple get ARROW
+
+        res = client.get(f"/v1/studies/{study_id}/{area1_id}/load/series?matrix_format=arrow")
+        assert res.status_code == 200
+        assert res.headers["content-type"] == "application/vnd.apache.arrow.file"