feat(ts-gen): add links ts-generation info inside DB

alembic/versions/dadad513f1b6_create_links_ts_generation_tables.py

@@ -0,0 +1,59 @@
+"""create_links_ts_generation_tables
+
+Revision ID: dadad513f1b6
+Revises: bae9c99bc42d
+Create Date: 2025-01-20 10:11:01.293931
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = 'dadad513f1b6'
+down_revision = 'bae9c99bc42d'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_table(
+        "nb_years_ts_generation",
+        sa.Column("id", sa.String(length=36), nullable=False),
+        sa.Column("links", sa.Integer(), server_default="1", nullable=False),
+        sa.ForeignKeyConstraint(
+            ["id"],
+            ["study.id"],
+            ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.Index('ix_nb_years_ts_generation_study_id', 'study_id')
+    )
+
+    default_boolean = sa.text('1') if op.get_context().dialect.name == 'sqlite' else 't'
+    op.create_table(
+        "links_parameters_ts_generation",
+        sa.Column("id", sa.Integer()),
+        sa.Column("area_from", sa.String(), nullable=False),
+        sa.Column("area_to", sa.String(), nullable=False),
+        sa.Column("prepro", sa.String(), nullable=True),
+        sa.Column("modulation", sa.String(), nullable=True),
+        sa.Column("unit_count", sa.Integer(), nullable=False, server_default="1"),
+        sa.Column("nominal_capacity", sa.Float(), nullable=False, server_default="0"),
+        sa.Column("law_planned", sa.Enum('uniform', 'geometric', name='lawplanned'), nullable=False),
+        sa.Column("law_forced", sa.Enum('uniform', 'geometric', name='lawforced'), nullable=False),
+        sa.Column("volatility_planned", sa.String(), nullable=False, server_default="0"),
+        sa.Column("volatility_forced", sa.String(), nullable=False, server_default="0"),
+        sa.Column("force_no_generation", sa.Boolean(), nullable=False, server_default=default_boolean),
+        sa.Column("study_id", sa.String(length=36), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["study_id"],
+            ["study.id"],
+            ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("id")
+    )
+
+
+def downgrade():
+    op.drop_table("nb_years_ts_generation")
+    op.drop_table("links_parameters_ts_generation")

antarest/study/business/link_management.py

@@ -16,9 +16,10 @@
 
 from antarest.core.exceptions import LinkNotFound
 from antarest.core.model import JSON
-from antarest.study.business.model.link_model import LinkBaseDTO, LinkDTO, LinkInternal
+from antarest.core.utils.fastapi_sqlalchemy import db
+from antarest.study.business.model.link_model import LinkBaseDTO, LinkDTO, LinkInternal, LinkTsGeneration
 from antarest.study.business.utils import execute_or_add_commands
-from antarest.study.model import RawStudy, Study
+from antarest.study.model import LinksParametersTsGeneration, RawStudy, Study
 from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
 from antarest.study.storage.storage_service import StudyStorageService
 from antarest.study.storage.variantstudy.model.command.create_link import CreateLink
@@ -34,13 +35,19 @@ def get_all_links(self, study: Study) -> List[LinkDTO]:
         file_study = self.storage_service.get_storage(study).get_raw(study)
         result: List[LinkDTO] = []
 
+        ts_generation_parameters = self.get_all_links_ts_generation_information(study.id)
+
         for area_id, area in file_study.config.areas.items():
             links_config = file_study.tree.get(["input", "links", area_id, "properties"])
 
             for link in area.links:
                 link_tree_config: Dict[str, Any] = links_config[link]
                 link_tree_config.update({"area1": area_id, "area2": link})
 
+                if area_id in ts_generation_parameters and link in ts_generation_parameters[area_id]:
+                    link_ts_generation = ts_generation_parameters[area_id][link]
+                    link_tree_config.update(link_ts_generation.model_dump(mode="json"))
+
                 link_internal = LinkInternal.model_validate(link_tree_config)
 
                 result.append(link_internal.to_dto())
@@ -54,10 +61,38 @@ def get_link(self, study: RawStudy, link: LinkInternal) -> LinkInternal:
 
         link_properties.update({"area1": link.area1, "area2": link.area2})
 
+        ts_generation_parameters = self.get_single_link_ts_generation_information(study.id, link.area1, link.area2)
+        link_properties.update(ts_generation_parameters.model_dump(mode="json"))
+
         updated_link = LinkInternal.model_validate(link_properties)
 
         return updated_link
 
+    @staticmethod
+    def get_all_links_ts_generation_information(study_id: str) -> dict[str, dict[str, LinkTsGeneration]]:
+        db_dictionnary: dict[str, dict[str, LinkTsGeneration]] = {}
+        with db():
+            all_links_parameters: list[LinksParametersTsGeneration] = (
+                db.session.query(LinksParametersTsGeneration).filter_by(study_id=study_id).all()
+            )
+            for link_parameters in all_links_parameters:
+                area_from = link_parameters.area_from
+                area_to = link_parameters.area_to
+                db_dictionnary.setdefault(area_from, {})[area_to] = LinkTsGeneration.from_db_model(link_parameters)
+        return db_dictionnary
+
+    @staticmethod
+    def get_single_link_ts_generation_information(study_id: str, area_from: str, area_to: str) -> LinkTsGeneration:
+        with db():
+            links_parameters = (
+                db.session.query(LinksParametersTsGeneration)
+                .filter_by(study_id=study_id, area_from=area_from, area_to=area_to)
+                .first()
+            )
+            if links_parameters:
+                return LinkTsGeneration.from_db_model(links_parameters)
+        return LinkTsGeneration()
+
     def create_link(self, study: Study, link_creation_dto: LinkDTO) -> LinkDTO:
         link = link_creation_dto.to_internal(StudyVersion.parse(study.version))
 

antarest/study/business/model/link_model.py

@@ -19,7 +19,8 @@
 from antarest.core.serde import AntaresBaseModel
 from antarest.core.utils.string import to_camel_case, to_kebab_case
 from antarest.study.business.enum_ignore_case import EnumIgnoreCase
-from antarest.study.model import STUDY_VERSION_8_2
+from antarest.study.model import STUDY_VERSION_8_2, LinksParametersTsGeneration
+from antarest.study.storage.rawstudy.model.filesystem.config.thermal import LawOption
 
 
 class AssetType(EnumIgnoreCase):
@@ -139,7 +140,7 @@ def join_with_comma(values: List[FilterOption]) -> str:
 ]
 
 
-class LinkBaseDTO(AntaresBaseModel):
+class LinkIniDTO(AntaresBaseModel):
     model_config = ConfigDict(alias_generator=to_camel_case, populate_by_name=True, extra="forbid")
 
     hurdles_cost: bool = False
@@ -158,6 +159,49 @@ class LinkBaseDTO(AntaresBaseModel):
     filter_year_by_year: Optional[comma_separated_enum_list] = field(default_factory=lambda: FILTER_VALUES)
 
 
+class LinkTsGeneration(AntaresBaseModel):
+    model_config = ConfigDict(alias_generator=to_camel_case, populate_by_name=True, extra="forbid")
+
+    unit_count: int = 1
+    nominal_capacity: float = 0
+    law_planned: LawOption = LawOption.UNIFORM
+    law_forced: LawOption = LawOption.UNIFORM
+    volatility_planned: float = Field(default=0.0, ge=0, le=1)
+    volatility_forced: float = Field(default=0.0, ge=0, le=1)
+    force_no_generation: bool = True
+
+    @staticmethod
+    def from_db_model(links_parameters_db: LinksParametersTsGeneration) -> "LinkTsGeneration":
+        args = {
+            "unit_count": links_parameters_db.unit_count,
+            "nominal_capacity": links_parameters_db.nominal_capacity,
+            "law_planned": links_parameters_db.law_planned,
+            "law_forced": links_parameters_db.law_forced,
+            "volatility_planned": links_parameters_db.volatility_planned,
+            "volatility_forced": links_parameters_db.volatility_forced,
+            "force_no_generation": links_parameters_db.force_no_generation,
+        }
+        return LinkTsGeneration.model_validate(args)
+
+    def to_db_model(self, study_id: str, area_from: str, area_to: str) -> LinksParametersTsGeneration:
+        return LinksParametersTsGeneration(
+            study_id=study_id,
+            area_from=area_from,
+            area_to=area_to,
+            unit_count=self.unit_count,
+            nominal_capacity=self.nominal_capacity,
+            law_planned=self.law_planned,
+            law_forced=self.law_forced,
+            volatility_planned=self.volatility_planned,
+            volatility_forced=self.volatility_forced,
+            force_no_generation=self.force_no_generation,
+        )
+
+
+class LinkBaseDTO(LinkIniDTO, LinkTsGeneration):
+    pass
+
+
 class Area(AntaresBaseModel):
     area1: str
     area2: str
@@ -177,7 +221,7 @@ def to_internal(self, version: StudyVersion) -> "LinkInternal":
         if version < STUDY_VERSION_8_2 and {"filter_synthesis", "filter_year_by_year"} & self.model_fields_set:
             raise LinkValidationError("Cannot specify a filter value for study's version earlier than v8.2")
 
-        data = self.model_dump()
+        data = self.model_dump(mode="json")
 
         if version < STUDY_VERSION_8_2:
             data["filter_synthesis"] = None
@@ -206,6 +250,15 @@ class LinkInternal(AntaresBaseModel):
     filter_synthesis: Optional[comma_separated_enum_list] = field(default_factory=lambda: FILTER_VALUES)
     filter_year_by_year: Optional[comma_separated_enum_list] = field(default_factory=lambda: FILTER_VALUES)
 
+    # Ts-generation part
+    unit_count: int = 1
+    nominal_capacity: float = 0
+    law_planned: LawOption = LawOption.UNIFORM
+    law_forced: LawOption = LawOption.UNIFORM
+    volatility_planned: float = Field(default=0.0, ge=0, le=1)
+    volatility_forced: float = Field(default=0.0, ge=0, le=1)
+    force_no_generation: bool = True
+
     def to_dto(self) -> LinkDTO:
         data = self.model_dump()
         return LinkDTO(**data)

antarest/study/model.py

@@ -25,6 +25,7 @@
     Column,
     DateTime,
     Enum,
+    Float,
     ForeignKey,
     Integer,
     PrimaryKeyConstraint,
@@ -40,6 +41,7 @@
 from antarest.core.serde import AntaresBaseModel
 from antarest.login.model import Group, GroupDTO, Identity
 from antarest.study.css4_colors import COLOR_NAMES
+from antarest.study.storage.rawstudy.model.filesystem.config.thermal import LawOption
 
 if TYPE_CHECKING:
     # avoid circular import
@@ -90,6 +92,43 @@
 }
 
 
+class NbYearsTsGeneration(Base):  # type:ignore
+    """
+    A table to store how many columns needs to be generated by category and by study
+
+    Attributes:
+        id: A foreign key on the study ID.
+        links: An integer representing how many columns needs to be generated during links TS generation.
+    """
+
+    __tablename__ = "nb_years_ts_generation"
+
+    id: str = Column(
+        String(36), ForeignKey("study.id", ondelete="CASCADE"), primary_key=True, index=True, nullable=False
+    )
+    links: int = Column(Integer, default=1, nullable=False)
+
+
+class LinksParametersTsGeneration(Base):  # type:ignore
+    """A table to store for each link of a study, the input parameters to give to the TS generation algorithm"""
+
+    __tablename__ = "links_parameters_ts_generation"
+
+    id = Column(Integer, primary_key=True, unique=True)
+    study_id = Column(String(36), ForeignKey("study.id", ondelete="CASCADE"), index=True, nullable=False)
+    area_from = Column(String(255), nullable=False)
+    area_to = Column(String(255), nullable=False)
+    prepro = Column(String(255), nullable=True)
+    modulation = Column(String(255), nullable=True)
+    unit_count = Column(Integer, nullable=False, default=1)
+    nominal_capacity = Column(Float, nullable=False, default=0)
+    law_planned = Column(Enum(LawOption), default=LawOption.UNIFORM, nullable=False)
+    law_forced = Column(Enum(LawOption), default=LawOption.UNIFORM, nullable=False)
+    volatility_planned = Column(Float, nullable=False, default=0)
+    volatility_forced = Column(Float, nullable=False, default=0)
+    force_no_generation = Column(Boolean, nullable=False, default=True)
+
+
 class StudyGroup(Base):  # type:ignore
     """
     A table to manage the many-to-many relationship between `Study` and `Group`

antarest/study/storage/variantstudy/model/command/create_link.py

@@ -17,10 +17,11 @@
 from typing_extensions import override
 
 from antarest.core.exceptions import LinkValidationError
+from antarest.core.utils.fastapi_sqlalchemy import db
 from antarest.core.utils.utils import assert_this
 from antarest.matrixstore.model import MatrixData
-from antarest.study.business.model.link_model import LinkInternal
-from antarest.study.model import STUDY_VERSION_8_2
+from antarest.study.business.model.link_model import Area, LinkInternal, LinkTsGeneration
+from antarest.study.model import STUDY_VERSION_8_2, LinksParametersTsGeneration
 from antarest.study.storage.rawstudy.model.filesystem.config.model import FileStudyTreeConfig, Link
 from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
 from antarest.study.storage.variantstudy.business.utils import strip_matrix_protocol, validate_matrix
@@ -216,19 +217,38 @@ def _apply(self, study_data: FileStudy, listener: Optional[ICommandListener] = N
         if not output.status:
             return output
 
-        to_exclude = {"area1", "area2"}
-        if version < STUDY_VERSION_8_2:
-            to_exclude.update("filter-synthesis", "filter-year-by-year")
-
-        validated_properties = LinkInternal.model_validate(self.parameters).model_dump(
-            by_alias=True, exclude=to_exclude
-        )
-
+        internal_link = LinkInternal.model_validate(self.parameters)
         area_from = data["area_from"]
         area_to = data["area_to"]
 
-        study_data.tree.save(validated_properties, ["input", "links", area_from, "properties", area_to])
+        # Saves ini properties
+        to_exclude = set(Area.model_fields.keys() | LinkTsGeneration.model_fields.keys())
+        if version < STUDY_VERSION_8_2:
+            to_exclude.update("filter-synthesis", "filter-year-by-year")
+        ini_properties = internal_link.model_dump(by_alias=True, exclude=to_exclude)
+        study_data.tree.save(ini_properties, ["input", "links", area_from, "properties", area_to])
+
+        # Saves DB properties
+        includes = set(LinkTsGeneration.model_fields.keys())
+        db_properties = LinkTsGeneration.model_validate(internal_link.model_dump(mode="json", include=includes))
+
+        with db():
+            new_parameters = LinksParametersTsGeneration(
+                study_id=study_data.config.study_id,
+                area_from=area_from,
+                area_to=area_to,
+                unit_count=db_properties.unit_count,
+                nominal_capacity=db_properties.nominal_capacity,
+                law_planned=db_properties.law_planned,
+                law_forced=db_properties.law_forced,
+                volatility_planned=db_properties.volatility_planned,
+                volatility_forced=db_properties.volatility_forced,
+                force_no_generation=db_properties.force_no_generation,
+            )
+            db.session.add(new_parameters)
+            db.session.commit()
 
+        # Saves matrices
         self.series = self.series or (self.command_context.generator_matrix_constants.get_link(version=version))
         self.direct = self.direct or (self.command_context.generator_matrix_constants.get_link_direct())
         self.indirect = self.indirect or (self.command_context.generator_matrix_constants.get_link_indirect())