-
Notifications
You must be signed in to change notification settings - Fork 5
/
extension.py
164 lines (149 loc) · 5.48 KB
/
extension.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
import logging
import shutil
from pathlib import Path
from typing import Any, Dict, List, cast
import yaml
from antarest.core.config import Config
from antarest.core.model import JSON
from antarest.core.utils.fastapi_sqlalchemy import db
from antarest.core.utils.utils import assert_this
from antarest.launcher.extensions.interface import ILauncherExtension
from antarest.study.service import StudyService
from antarest.study.storage.rawstudy.model.filesystem.config.model import transform_name_to_id
from antarest.study.storage.rawstudy.model.filesystem.factory import FileStudy
logger = logging.getLogger(__name__)
class AdequacyPatchExtension(ILauncherExtension):
EXTENSION_NAME = "adequacy_patch"
def __init__(self, study_service: StudyService, config: Config):
self.study_service = study_service
def get_name(self) -> str:
return AdequacyPatchExtension.EXTENSION_NAME
def after_export_flat_hook(
self,
job_id: str,
study_id: str,
study_export_path: Path,
launcher_opts: Any,
) -> None:
logger.info("Applying adequacy patch postprocessing script")
study = self.study_service.storage_service.raw_study_service.study_factory.create_from_fs(
study_export_path, study_id, use_cache=False
)
user_config = study.tree.get(
["user"],
)
assert_this("flowbased" in user_config)
adequacy_patch_config = yaml.safe_load(
cast(
bytes, study.tree.get(["user", "adequacypatch", "config.yml"])
)
)
assert_this("areas" in adequacy_patch_config)
self.prepare_study_for_adq_patch(job_id, study, adequacy_patch_config)
full_r_version = "legacy" in launcher_opts and launcher_opts["legacy"]
if (
"mode" in adequacy_patch_config
and adequacy_patch_config["mode"] == "legacy"
):
full_r_version = True
if full_r_version:
logger.info("Using legacy quadratic mode")
post_processing_file = (
Path(__file__).parent
/ "resources"
/ "post-processing-legacy.R"
)
else:
logger.info("Using linearized mode")
post_processing_file = (
Path(__file__).parent / "resources" / "post-processing.R"
)
shutil.copy(
post_processing_file, study_export_path / "post-processing.R"
)
def prepare_study_for_adq_patch(
self, job_id: str, study: FileStudy, adq_patch_config: JSON
) -> Dict[str, bool]:
area_to_turn_on: List[str] = [
transform_name_to_id(area_id)
for area_id in adq_patch_config.get("areas", [])
]
original_area_enabled: Dict[str, bool] = {}
original_link_enabled: Dict[str, bool] = {}
year_by_year_active = study.tree.get(
["settings", "generaldata", "general", "year-by-year"]
)
study.tree.save(
True, ["settings", "generaldata", "general", "year-by-year"]
)
for area_id, area in study.config.areas.items():
# areas
original_area_enabled[area_id] = "hourly" in area.filters_year
if (
not original_area_enabled[area_id]
and area_id in area_to_turn_on
):
study.tree.save(
", ".join([*area.filters_year, "hourly"]),
[
"input",
"areas",
area_id,
"optimization",
"filtering",
"filter-year-by-year",
],
)
# links
for area_2, link in area.links.items():
link_id = f"{area_id} - {area_2}"
original_link_enabled[link_id] = "hourly" in link.filters_year
if not original_link_enabled[link_id] and (
area_id in area_to_turn_on or area_2 in area_to_turn_on
):
study.tree.save(
", ".join([*link.filters_year, "hourly"]),
[
"input",
"links",
area_id,
"properties",
area_2,
"filter-year-by-year",
],
)
with db():
with open(
study.config.study_path
/ "user"
/ "adequacypatch"
/ "hourly-areas.yml",
"w",
) as fh:
yaml.dump(original_area_enabled, fh)
with open(
study.config.study_path
/ "user"
/ "adequacypatch"
/ "hourly-links.yml",
"w",
) as fh:
yaml.dump(original_link_enabled, fh)
if year_by_year_active:
with open(
study.config.study_path
/ "user"
/ "adequacypatch"
/ "year-by-year-active",
"w",
) as fh:
fh.write("True")
return original_area_enabled
def before_import_hook(
self,
job_id: str,
study_id: str,
study_output_path: Path,
ext_opts: Any,
) -> None:
pass