flextool 4.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (322) hide show
  1. flextool/__init__.py +41 -0
  2. flextool/_mem_sampler.py +193 -0
  3. flextool/_resources.py +43 -0
  4. flextool/calibrate/__init__.py +51 -0
  5. flextool/calibrate/__main__.py +11 -0
  6. flextool/calibrate/_cli.py +316 -0
  7. flextool/calibrate/_db_alt.py +166 -0
  8. flextool/calibrate/_final_outputs.py +110 -0
  9. flextool/calibrate/_guard.py +151 -0
  10. flextool/calibrate/_loop.py +558 -0
  11. flextool/calibrate/_readers.py +223 -0
  12. flextool/calibrate/_report.py +263 -0
  13. flextool/calibrate/_sizing.py +699 -0
  14. flextool/calibrate/_solve.py +134 -0
  15. flextool/calibrate/_solve_status.py +495 -0
  16. flextool/cli/__init__.py +9 -0
  17. flextool/cli/_console.py +51 -0
  18. flextool/cli/_timing.py +147 -0
  19. flextool/cli/cmd_execute_flextool_workflow.py +187 -0
  20. flextool/cli/cmd_export_to_tabular.py +56 -0
  21. flextool/cli/cmd_import_sensitivities.py +75 -0
  22. flextool/cli/cmd_migrate_database.py +13 -0
  23. flextool/cli/cmd_open_results_db.py +269 -0
  24. flextool/cli/cmd_read_matpower.py +66 -0
  25. flextool/cli/cmd_read_old_flextool.py +63 -0
  26. flextool/cli/cmd_read_self_describing_tabular_input.py +50 -0
  27. flextool/cli/cmd_read_tabular_input.py +81 -0
  28. flextool/cli/cmd_run_flextool.py +1095 -0
  29. flextool/cli/cmd_scenario_results.py +284 -0
  30. flextool/cli/cmd_solve_mps.py +169 -0
  31. flextool/cli/cmd_update_flextool.py +17 -0
  32. flextool/cli/cmd_write_outputs.py +125 -0
  33. flextool/common_utils/__init__.py +1 -0
  34. flextool/common_utils/plot_mem_shape.py +77 -0
  35. flextool/common_utils/precision.py +451 -0
  36. flextool/decomposition/__init__.py +0 -0
  37. flextool/decomposition/region_decomposition.py +128 -0
  38. flextool/decomposition/region_filter.py +1261 -0
  39. flextool/engine_polars/__init__.py +110 -0
  40. flextool/engine_polars/_axis_enums.py +742 -0
  41. flextool/engine_polars/_benders.py +3462 -0
  42. flextool/engine_polars/_block_layout.py +1479 -0
  43. flextool/engine_polars/_blocks.py +1515 -0
  44. flextool/engine_polars/_commodity_ladder.py +660 -0
  45. flextool/engine_polars/_cumulative_invest.py +1165 -0
  46. flextool/engine_polars/_db_loader.py +153 -0
  47. flextool/engine_polars/_db_reader.py +127 -0
  48. flextool/engine_polars/_dc_power_flow.py +445 -0
  49. flextool/engine_polars/_delay.py +442 -0
  50. flextool/engine_polars/_derived_arithmetic.py +432 -0
  51. flextool/engine_polars/_derived_block.py +990 -0
  52. flextool/engine_polars/_derived_branch.py +769 -0
  53. flextool/engine_polars/_derived_existing.py +1353 -0
  54. flextool/engine_polars/_derived_npv.py +1297 -0
  55. flextool/engine_polars/_derived_params.py +9850 -0
  56. flextool/engine_polars/_derived_profile.py +881 -0
  57. flextool/engine_polars/_derived_walks.py +276 -0
  58. flextool/engine_polars/_determinism.py +70 -0
  59. flextool/engine_polars/_direct_params.py +2186 -0
  60. flextool/engine_polars/_dump_csvs.py +1009 -0
  61. flextool/engine_polars/_emit_arc_unions.py +1631 -0
  62. flextool/engine_polars/_emit_calc_params.py +729 -0
  63. flextool/engine_polars/_emit_chain_params.py +709 -0
  64. flextool/engine_polars/_emit_co2_accumulators.py +400 -0
  65. flextool/engine_polars/_emit_dispatchers.py +690 -0
  66. flextool/engine_polars/_emit_energy_margin.py +125 -0
  67. flextool/engine_polars/_emit_energy_margin_adder.py +290 -0
  68. flextool/engine_polars/_emit_entity_annual.py +428 -0
  69. flextool/engine_polars/_emit_inflow_scaling.py +1420 -0
  70. flextool/engine_polars/_emit_leaf_sets.py +550 -0
  71. flextool/engine_polars/_emit_lp_scaling.py +665 -0
  72. flextool/engine_polars/_emit_mid_sets.py +859 -0
  73. flextool/engine_polars/_emit_pdt_params.py +759 -0
  74. flextool/engine_polars/_emit_per_solve.py +774 -0
  75. flextool/engine_polars/_emit_period_calc.py +504 -0
  76. flextool/engine_polars/_emit_period_params.py +2398 -0
  77. flextool/engine_polars/_emit_provider_io.py +141 -0
  78. flextool/engine_polars/_emit_reserve.py +574 -0
  79. flextool/engine_polars/_emit_solve_time.py +311 -0
  80. flextool/engine_polars/_emit_solve_writers.py +1249 -0
  81. flextool/engine_polars/_flex_data_accumulator.py +388 -0
  82. flextool/engine_polars/_flex_data_provider.py +478 -0
  83. flextool/engine_polars/_group_slack.py +1253 -0
  84. flextool/engine_polars/_inmemory_reader.py +140 -0
  85. flextool/engine_polars/_input_source.py +336 -0
  86. flextool/engine_polars/_invest_seeds.py +191 -0
  87. flextool/engine_polars/_native_input_writer.py +100 -0
  88. flextool/engine_polars/_native_run_model.py +1348 -0
  89. flextool/engine_polars/_orchestration.py +4314 -0
  90. flextool/engine_polars/_output_writer.py +439 -0
  91. flextool/engine_polars/_param_shapes.py +1595 -0
  92. flextool/engine_polars/_parquet_bundle.py +723 -0
  93. flextool/engine_polars/_pdt_join.py +167 -0
  94. flextool/engine_polars/_pdt_lookup.py +547 -0
  95. flextool/engine_polars/_per_solve_sets.py +335 -0
  96. flextool/engine_polars/_projection_params.py +2056 -0
  97. flextool/engine_polars/_provider_keys.py +173 -0
  98. flextool/engine_polars/_provider_translators.py +225 -0
  99. flextool/engine_polars/_recursive_solve.py +703 -0
  100. flextool/engine_polars/_region_filter.py +2508 -0
  101. flextool/engine_polars/_reserve.py +649 -0
  102. flextool/engine_polars/_solve_acceptance.py +331 -0
  103. flextool/engine_polars/_solve_config.py +1001 -0
  104. flextool/engine_polars/_solve_context.py +885 -0
  105. flextool/engine_polars/_solve_handoff.py +164 -0
  106. flextool/engine_polars/_solve_state.py +232 -0
  107. flextool/engine_polars/_solver_base.py +36 -0
  108. flextool/engine_polars/_solver_dispatch.py +511 -0
  109. flextool/engine_polars/_spinedb_reader.py +1165 -0
  110. flextool/engine_polars/_stochastic.py +593 -0
  111. flextool/engine_polars/_subprocess_solve.py +1838 -0
  112. flextool/engine_polars/_timeline.py +1416 -0
  113. flextool/engine_polars/_vectorize.py +438 -0
  114. flextool/engine_polars/_warm.py +858 -0
  115. flextool/engine_polars/autoscale/__init__.py +107 -0
  116. flextool/engine_polars/autoscale/_config.py +218 -0
  117. flextool/engine_polars/autoscale/_layer2.py +1253 -0
  118. flextool/engine_polars/autoscale/_layer2_types.py +584 -0
  119. flextool/engine_polars/autoscale/_quantity_types.py +621 -0
  120. flextool/engine_polars/autoscale/_report.py +336 -0
  121. flextool/engine_polars/chain.py +259 -0
  122. flextool/engine_polars/input.py +6638 -0
  123. flextool/engine_polars/model.py +4754 -0
  124. flextool/env_check.py +388 -0
  125. flextool/export_to_tabular/__init__.py +5 -0
  126. flextool/export_to_tabular/db_reader.py +224 -0
  127. flextool/export_to_tabular/excel_writer.py +3559 -0
  128. flextool/export_to_tabular/export_settings.yaml +377 -0
  129. flextool/export_to_tabular/export_to_excel.py +227 -0
  130. flextool/export_to_tabular/formatting.py +543 -0
  131. flextool/export_to_tabular/sheet_config.py +876 -0
  132. flextool/gui/__init__.py +0 -0
  133. flextool/gui/__main__.py +118 -0
  134. flextool/gui/calibrate_commands.py +184 -0
  135. flextool/gui/calibrate_jobs.py +424 -0
  136. flextool/gui/check_tree.py +142 -0
  137. flextool/gui/cli_format.py +83 -0
  138. flextool/gui/config_parser.py +68 -0
  139. flextool/gui/data_models.py +362 -0
  140. flextool/gui/db_editor_integration.py +202 -0
  141. flextool/gui/db_version_check.py +269 -0
  142. flextool/gui/dialogs/__init__.py +0 -0
  143. flextool/gui/dialogs/add_dialog.py +1098 -0
  144. flextool/gui/dialogs/calibrate_dialog.py +1259 -0
  145. flextool/gui/dialogs/file_picker.py +473 -0
  146. flextool/gui/dialogs/group_picker.py +299 -0
  147. flextool/gui/dialogs/migration_consent_dialog.py +106 -0
  148. flextool/gui/dialogs/migration_progress_dialog.py +237 -0
  149. flextool/gui/dialogs/plot_dialog.py +459 -0
  150. flextool/gui/dialogs/plot_settings_picker.py +2184 -0
  151. flextool/gui/dialogs/project_dialog.py +426 -0
  152. flextool/gui/dialogs/update_dialog.py +212 -0
  153. flextool/gui/downsampling.py +88 -0
  154. flextool/gui/error_handling.py +50 -0
  155. flextool/gui/execution_manager.py +1715 -0
  156. flextool/gui/execution_window.py +1377 -0
  157. flextool/gui/hover_tooltip.py +111 -0
  158. flextool/gui/input_sources.py +730 -0
  159. flextool/gui/main_window.py +6181 -0
  160. flextool/gui/network_graph.py +215 -0
  161. flextool/gui/output_actions.py +393 -0
  162. flextool/gui/output_log_window.py +159 -0
  163. flextool/gui/platform_utils.py +421 -0
  164. flextool/gui/plot_cache.py +88 -0
  165. flextool/gui/plot_canvas.py +543 -0
  166. flextool/gui/plot_config_reader.py +272 -0
  167. flextool/gui/project_utils.py +100 -0
  168. flextool/gui/result_viewer.py +4394 -0
  169. flextool/gui/scenario_key.py +162 -0
  170. flextool/gui/scenario_lists.py +516 -0
  171. flextool/gui/settings_io.py +360 -0
  172. flextool/gui/solve_reader.py +103 -0
  173. flextool/gui/tree_reorder.py +88 -0
  174. flextool/gui/ui_metrics.py +420 -0
  175. flextool/input_derivation/__init__.py +281 -0
  176. flextool/input_derivation/_commodity_ladder.py +375 -0
  177. flextool/input_derivation/_commodity_ladder_sets.py +70 -0
  178. flextool/input_derivation/_dc_power_flow.py +377 -0
  179. flextool/input_derivation/_method_constants.py +77 -0
  180. flextool/input_derivation/_process_method.py +258 -0
  181. flextool/input_derivation/_specs.py +1026 -0
  182. flextool/input_derivation/_validators.py +321 -0
  183. flextool/lean_parquet.py +159 -0
  184. flextool/model_builder/__init__.py +5 -0
  185. flextool/model_builder/build_model.py +589 -0
  186. flextool/model_builder/encoding.py +67 -0
  187. flextool/model_builder/names.py +34 -0
  188. flextool/model_builder/profiles.py +129 -0
  189. flextool/plot_outputs/__init__.py +14 -0
  190. flextool/plot_outputs/axis_helpers.py +355 -0
  191. flextool/plot_outputs/color_template.py +888 -0
  192. flextool/plot_outputs/config.py +171 -0
  193. flextool/plot_outputs/format_helpers.py +345 -0
  194. flextool/plot_outputs/legend_helpers.py +143 -0
  195. flextool/plot_outputs/orchestrator.py +1141 -0
  196. flextool/plot_outputs/perf.py +37 -0
  197. flextool/plot_outputs/plan.py +1787 -0
  198. flextool/plot_outputs/plot_bars.py +1510 -0
  199. flextool/plot_outputs/plot_bars_detail.py +753 -0
  200. flextool/plot_outputs/plot_lines.py +951 -0
  201. flextool/plot_outputs/shared_manifest.py +564 -0
  202. flextool/plot_outputs/subplot_helpers.py +137 -0
  203. flextool/process_inputs/__init__.py +188 -0
  204. flextool/process_inputs/import_old_excel_input.json +4159 -0
  205. flextool/process_inputs/read_matpower.py +451 -0
  206. flextool/process_inputs/read_old_flextool.py +1288 -0
  207. flextool/process_inputs/read_self_describing_excel.py +1423 -0
  208. flextool/process_inputs/read_tabular_with_specification.py +1114 -0
  209. flextool/process_inputs/write_old_flextool_to_db.py +3077 -0
  210. flextool/process_inputs/write_self_describing_to_db.py +977 -0
  211. flextool/process_inputs/write_to_input_db.py +269 -0
  212. flextool/process_outputs/__init__.py +7 -0
  213. flextool/process_outputs/_annualize.py +55 -0
  214. flextool/process_outputs/_inmemory_helpers.py +292 -0
  215. flextool/process_outputs/_output_meta.py +672 -0
  216. flextool/process_outputs/calc_capacity_flows.py +107 -0
  217. flextool/process_outputs/calc_connections.py +136 -0
  218. flextool/process_outputs/calc_costs.py +260 -0
  219. flextool/process_outputs/calc_group_flows.py +192 -0
  220. flextool/process_outputs/calc_slacks.py +103 -0
  221. flextool/process_outputs/calc_storage_vre.py +160 -0
  222. flextool/process_outputs/drop_levels.py +208 -0
  223. flextool/process_outputs/handoff_writers.py +1315 -0
  224. flextool/process_outputs/out_ancillary.py +544 -0
  225. flextool/process_outputs/out_capacity.py +179 -0
  226. flextool/process_outputs/out_costs.py +334 -0
  227. flextool/process_outputs/out_flowgroup.py +189 -0
  228. flextool/process_outputs/out_flows.py +301 -0
  229. flextool/process_outputs/out_group.py +475 -0
  230. flextool/process_outputs/out_node.py +190 -0
  231. flextool/process_outputs/persist_realized_slice.py +601 -0
  232. flextool/process_outputs/process_results.py +24 -0
  233. flextool/process_outputs/read_highs_solution.py +2256 -0
  234. flextool/process_outputs/read_parameters.py +1799 -0
  235. flextool/process_outputs/read_sets.py +1095 -0
  236. flextool/process_outputs/read_variables.py +553 -0
  237. flextool/process_outputs/solve_order.py +81 -0
  238. flextool/process_outputs/spinedb_replay.py +412 -0
  239. flextool/process_outputs/union_realized_slice.py +224 -0
  240. flextool/process_outputs/write_outputs.py +1286 -0
  241. flextool/process_outputs/write_spinedb.py +1267 -0
  242. flextool/representative_periods/__init__.py +5 -0
  243. flextool/representative_periods/clustering.py +165 -0
  244. flextool/representative_periods/force_include.py +563 -0
  245. flextool/representative_periods/netload.py +365 -0
  246. flextool/representative_periods/netload_inputs.py +345 -0
  247. flextool/representative_periods/netload_iterate.py +722 -0
  248. flextool/representative_periods/preprocess.py +948 -0
  249. flextool/representative_periods/scenario_stack.py +195 -0
  250. flextool/representative_periods/weights.py +124 -0
  251. flextool/scenario_comparison/__init__.py +13 -0
  252. flextool/scenario_comparison/config_builder.py +158 -0
  253. flextool/scenario_comparison/constants.py +20 -0
  254. flextool/scenario_comparison/data_models.py +222 -0
  255. flextool/scenario_comparison/db_reader.py +399 -0
  256. flextool/scenario_comparison/dispatch_data.py +1002 -0
  257. flextool/scenario_comparison/dispatch_mappings.py +205 -0
  258. flextool/scenario_comparison/dispatch_plots.py +691 -0
  259. flextool/scenario_comparison/input_entity_colors.py +319 -0
  260. flextool/scenario_comparison/orchestrator.py +453 -0
  261. flextool/scenario_comparison/plan_union.py +244 -0
  262. flextool/scenario_comparison/plot_settings_seed.py +205 -0
  263. flextool/schemas/AXIS_CONTRACT.md +71 -0
  264. flextool/schemas/canonical_databases/howto_aggregate_output.json +6225 -0
  265. flextool/schemas/canonical_databases/howto_connections.json +5606 -0
  266. flextool/schemas/canonical_databases/howto_demand.json +5518 -0
  267. flextool/schemas/canonical_databases/howto_hydro_reservoir.json +6239 -0
  268. flextool/schemas/canonical_databases/howto_hydro_reservoir_with_pump.json +5933 -0
  269. flextool/schemas/canonical_databases/howto_non_sync_and_curtailment.json +5794 -0
  270. flextool/schemas/canonical_databases/howto_ramp_and_start_up.json +5707 -0
  271. flextool/schemas/canonical_databases/howto_stochastics.json +6032 -0
  272. flextool/schemas/canonical_databases/templates_examples.json +13532 -0
  273. flextool/schemas/canonical_databases/templates_time_settings_only.json +5340 -0
  274. flextool/schemas/comparison_settings_template.json +197 -0
  275. flextool/schemas/default_plot_settings.yaml +260 -0
  276. flextool/schemas/default_plots.yaml +2293 -0
  277. flextool/schemas/flextool_axis_contract.json +303 -0
  278. flextool/schemas/flextool_axis_contract.schema.json +247 -0
  279. flextool/schemas/old_flextool_import_template.json +4443 -0
  280. flextool/schemas/output_info_template.json +48 -0
  281. flextool/schemas/output_settings_template.json +256 -0
  282. flextool/schemas/pre_v26/flextool_template_constant_default.json +2105 -0
  283. flextool/schemas/pre_v26/flextool_template_default_optional_output.json +2152 -0
  284. flextool/schemas/pre_v26/flextool_template_default_value.json +2094 -0
  285. flextool/schemas/pre_v26/flextool_template_drop_down.json +2080 -0
  286. flextool/schemas/pre_v26/flextool_template_lifetime_method.json +1990 -0
  287. flextool/schemas/pre_v26/flextool_template_optional_outputs.json +2094 -0
  288. flextool/schemas/pre_v26/flextool_template_output_node_flows.json +2105 -0
  289. flextool/schemas/pre_v26/flextool_template_results_master.json +493 -0
  290. flextool/schemas/pre_v26/flextool_template_rolling_start_remove.json +2087 -0
  291. flextool/schemas/pre_v26/flextool_template_rolling_window.json +2059 -0
  292. flextool/schemas/pre_v26/flextool_template_storage_binding_defaults.json +46 -0
  293. flextool/schemas/pre_v26/flextool_template_v2.json +1990 -0
  294. flextool/schemas/pre_v26/flextool_template_v25.json +3864 -0
  295. flextool/schemas/spinedb_results_schema.json +581 -0
  296. flextool/schemas/spinedb_schema.json +4636 -0
  297. flextool/solver_config/copt.opt.template +18 -0
  298. flextool/solver_config/cplex.opt.template +25 -0
  299. flextool/solver_config/gurobi.opt.template +18 -0
  300. flextool/solver_config/highs.opt.template +18 -0
  301. flextool/solver_config/xpress.opt.template +26 -0
  302. flextool/spinedb_backend/__init__.py +26 -0
  303. flextool/spinedb_backend/_axis_enums.py +1119 -0
  304. flextool/spinedb_backend/_backend.py +1139 -0
  305. flextool/update_flextool/__init__.py +12 -0
  306. flextool/update_flextool/canonical_databases.py +251 -0
  307. flextool/update_flextool/db_migration.py +7108 -0
  308. flextool/update_flextool/ensure_settings_db.py +138 -0
  309. flextool/update_flextool/export_database.py +103 -0
  310. flextool/update_flextool/extend_tests_fixture.py +772 -0
  311. flextool/update_flextool/generate_canonical.py +274 -0
  312. flextool/update_flextool/initialize_database.py +42 -0
  313. flextool/update_flextool/install_info.py +225 -0
  314. flextool/update_flextool/self_update.py +464 -0
  315. flextool/update_flextool/sync_master_json_template.py +125 -0
  316. flextool/update_flextool/test_fixtures.py +187 -0
  317. flextool-4.0.0.dist-info/METADATA +217 -0
  318. flextool-4.0.0.dist-info/RECORD +322 -0
  319. flextool-4.0.0.dist-info/WHEEL +5 -0
  320. flextool-4.0.0.dist-info/entry_points.txt +17 -0
  321. flextool-4.0.0.dist-info/licenses/LICENSE.txt +19 -0
  322. flextool-4.0.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,2508 @@
1
+ """Regional filter for Benders decomposition (gaps A4 + A5).
2
+
3
+ Slices a whole-system :class:`FlexData` into N per-region :class:`FlexData`s.
4
+ Cross-region arcs (rows in ``process_source_sink`` whose ``source`` and
5
+ ``sink`` straddle two regions) are *severed* into two virtual one-way
6
+ half-flow arcs, one in each region.
7
+
8
+ The virtual entities are pure bookkeeping — the half-flow on the export
9
+ side and the half-flow on the import side carry the **same flow** at
10
+ optimality. In the Benders scheme that coupling is enforced by the
11
+ coordinating master (:mod:`flextool.engine_polars._benders`): the master
12
+ holds the inter-regional trade flows and pins each region's forward
13
+ cross-region half-flow to its chosen value before solving the region as a
14
+ subproblem. Inside a region's standalone LP, the half-flow column is just
15
+ an ordinary ``v_flow`` column with:
16
+
17
+ * ``p_unitsize = 1`` (same units as the original column);
18
+ * ``p_flow_upper`` = original pipe capacity (so dispatch can push flow
19
+ at full pipe capacity; uncapped when the master may pin a positive
20
+ greenfield trade);
21
+ * the in-region terminal node enters ``flow_to_n`` (export) or
22
+ ``flow_from_n`` (import) so the half-flow contributes to its
23
+ nodeBalance;
24
+ * the *virtual* node sits OUTSIDE ``nodeBalance`` so the half-flow's
25
+ other terminal is free (no balance pin, no penalty).
26
+
27
+ This module is a no-op when the input :class:`FlexData` has no
28
+ ``decomposition_method=benders_regional`` group params.
29
+
30
+ Public surface
31
+ --------------
32
+ * :class:`HalfFlow` — one severed arc; pairs across regions.
33
+ * :class:`RegionSplit` — output of :func:`split` for one region.
34
+ * :func:`discover_regions` — returns ``[region_name, ...]`` from a
35
+ whole-system FlexData (reads ``p_group_decomposition.csv`` indirectly
36
+ via the populated ``group_entity`` / ``group_node`` frames).
37
+ * :func:`split` — produces ``list[RegionSplit]``.
38
+ """
39
+ from __future__ import annotations
40
+
41
+ from dataclasses import dataclass, field
42
+ from pathlib import Path
43
+ import dataclasses
44
+ from typing import Iterable
45
+
46
+ import polars as pl
47
+
48
+ from polar_high import Param
49
+
50
+ from flextool.engine_polars.input import FlexData
51
+ from flextool.engine_polars._axis_enums import (
52
+ cast_dim,
53
+ get_global_axis_enums,
54
+ reset_global_axis_enums,
55
+ schema_dtype,
56
+ set_global_axis_enums,
57
+ )
58
+ from flextool.engine_polars._param_shapes import promote_param_to_dt
59
+ from flextool.engine_polars._pdt_join import compute_pss_dt
60
+
61
+
62
+ __all__ = [
63
+ "HalfFlow",
64
+ "RegionSplit",
65
+ "compute_master_hosted_nodes",
66
+ "discover_regions",
67
+ "split",
68
+ "master_network_data",
69
+ "load_decomposition_method",
70
+ "load_region_membership",
71
+ ]
72
+
73
+
74
+ # ---------------------------------------------------------------------------
75
+ # Data classes
76
+ # ---------------------------------------------------------------------------
77
+
78
+
79
+ @dataclass
80
+ class HalfFlow:
81
+ """One half-flow arc replacing one direction of a cross-region pipe.
82
+
83
+ A bi-directional cross-region pipe ``pipe(A→B, B→A)`` produces FOUR
84
+ HalfFlows: ``pipe(A→B)`` becomes an export in A and an import in B;
85
+ ``pipe(B→A)`` becomes an export in B and an import in A. Each
86
+ coupling pair shares the same ``(original_p, original_source,
87
+ original_sink)`` triple — the :mod:`flextool.engine_polars._benders`
88
+ coordinator pairs export and import on that key.
89
+ """
90
+ region: str
91
+ side: str # "export" or "import"
92
+ # Original cross-region arc identity (the pairing key).
93
+ original_p: str
94
+ original_source: str
95
+ original_sink: str
96
+ # In-region terminal node — appears in the regional flow_to_n /
97
+ # flow_from_n alongside the half-flow. For an export this is the
98
+ # original source; for an import it is the original sink.
99
+ in_region_node: str
100
+ # Virtual entities introduced by this half-flow.
101
+ virtual_node: str
102
+ # Virtual connection: the (p, source, sink) row that replaces the
103
+ # original cross-region arc inside the region's frames.
104
+ virtual_p: str
105
+ virtual_arc_source: str
106
+ virtual_arc_sink: str
107
+
108
+
109
+ @dataclass
110
+ class RegionSplit:
111
+ """One region's filtered :class:`FlexData` plus coupling metadata."""
112
+ region: str
113
+ data: FlexData
114
+ half_flows: list[HalfFlow] = field(default_factory=list)
115
+
116
+
117
+ # ---------------------------------------------------------------------------
118
+ # Decomposition discovery
119
+ # ---------------------------------------------------------------------------
120
+
121
+
122
+ def load_decomposition_method(
123
+ work_dir: "Path | str | None" = None,
124
+ *,
125
+ provider: "object | None" = None,
126
+ ) -> dict[str, str]:
127
+ """Return ``{group: method}`` from
128
+ ``input/p_group_decomposition.csv`` (Step 2.6 Provider-first).
129
+
130
+ Resolution order:
131
+
132
+ 1. *provider* carries ``input/p_group_decomposition`` →
133
+ read from memory.
134
+ 2. *provider* is ``None`` AND *work_dir* points at a workdir with
135
+ the file on disk → seed an ephemeral Provider from that
136
+ directory and serve from memory.
137
+ 3. Otherwise → ``{}``.
138
+
139
+ The whole-system loader doesn't surface the decomposition_method
140
+ parameter in :class:`FlexData` because it's only used by the
141
+ decomposition driver itself; this helper plus the Provider keep
142
+ it out of cascade memory until the driver needs it.
143
+ """
144
+ key = "input/p_group_decomposition"
145
+ df: pl.DataFrame | None = None
146
+ if provider is not None and provider.has(key):
147
+ df = provider.get(key)
148
+ elif provider is None and work_dir is not None:
149
+ # Off-cascade test bridge: seed from disk via the centralised
150
+ # helper so Rule 1 of test_meta_provider_invariants stays clean
151
+ # (no bare ``pl.read_csv`` / ``_read_csv_file`` in this module).
152
+ path = Path(work_dir) / "input" / "p_group_decomposition.csv"
153
+ if not path.exists():
154
+ return {}
155
+ from flextool.engine_polars._flex_data_provider import (
156
+ FlexDataProvider,
157
+ )
158
+ from flextool.engine_polars._input_source import (
159
+ seed_provider_from_dir,
160
+ )
161
+ local = FlexDataProvider()
162
+ seed_provider_from_dir(
163
+ local, Path(work_dir) / "input", "input",
164
+ names=("p_group_decomposition",),
165
+ )
166
+ if local.has(key):
167
+ df = local.get(key)
168
+ if df is None or df.height == 0:
169
+ return {}
170
+ cols = df.columns
171
+ # Expected columns: group, groupParam, p_group
172
+ if "group" not in cols:
173
+ return {}
174
+ val_col = next((c for c in ("p_group", "value") if c in cols), None)
175
+ if val_col is None:
176
+ return {}
177
+ rows = (
178
+ df.filter(pl.col("groupParam") == "decomposition_method")
179
+ if "groupParam" in cols else df
180
+ )
181
+ return {r["group"]: r[val_col] for r in rows.iter_rows(named=True)}
182
+
183
+
184
+ def discover_regions(
185
+ work_dir: "Path | str | None" = None,
186
+ *,
187
+ provider: "object | None" = None,
188
+ ) -> list[str]:
189
+ """Return the list of group names with
190
+ ``decomposition_method=benders_regional``."""
191
+ methods = load_decomposition_method(work_dir, provider=provider)
192
+ return sorted(g for g, m in methods.items() if m == "benders_regional")
193
+
194
+
195
+ def load_region_membership(
196
+ data: FlexData, regions: list[str],
197
+ ) -> dict[str, dict[str, set[str]]]:
198
+ """Return ``{region: {"nodes": set, "processes": set}}`` from
199
+ ``data.group_entity`` / ``data.group_node`` for the given regions.
200
+
201
+ Falls back to empty sets when the relevant frames are absent (a
202
+ same-shape stub useful in unit tests).
203
+ """
204
+ out: dict[str, dict[str, set[str]]] = {}
205
+ nodes_by_g: dict[str, set[str]] = {}
206
+ procs_by_g: dict[str, set[str]] = {}
207
+ if data.group_node is not None and data.group_node.height > 0:
208
+ for r in data.group_node.iter_rows(named=True):
209
+ nodes_by_g.setdefault(r["g"], set()).add(r["n"])
210
+ if data.group_entity is not None and data.group_entity.height > 0:
211
+ # group_entity is the union (g, e) — to get processes we'd need
212
+ # to know which e's are processes. process_unit + connections
213
+ # set is everything in process_source_sink["p"].
214
+ all_procs: set[str] = set()
215
+ if data.process_source_sink is not None:
216
+ all_procs |= set(data.process_source_sink["p"].unique().to_list())
217
+ if data.process_indirect is not None:
218
+ all_procs |= set(data.process_indirect["p"].unique().to_list())
219
+ for r in data.group_entity.iter_rows(named=True):
220
+ if r["e"] in all_procs:
221
+ procs_by_g.setdefault(r["g"], set()).add(r["e"])
222
+ for g in regions:
223
+ out[g] = {
224
+ "nodes": set(nodes_by_g.get(g, set())),
225
+ "processes": set(procs_by_g.get(g, set())),
226
+ }
227
+ return out
228
+
229
+
230
+ def compute_master_hosted_nodes(
231
+ data: FlexData,
232
+ region_membership: dict[str, dict[str, set[str]]],
233
+ ) -> set[str]:
234
+ """Return the master-hosted node set: every node carrying a balance
235
+ or state row (``data.nodeBalance`` ∪ ``data.nodeState``) that is in
236
+ NO region's membership.
237
+
238
+ Nodes with no balance/state row (pure commodity/market nodes) are
239
+ deliberately NOT included — they keep today's shared-replicate
240
+ semantics (replication is safe for them: there is no balance row to
241
+ duplicate), pinned by the existing region-filter tests.
242
+ """
243
+ balance_state: set[str] = set()
244
+ if data.nodeBalance is not None and data.nodeBalance.height > 0:
245
+ balance_state |= set(data.nodeBalance["n"].to_list())
246
+ if data.nodeState is not None and data.nodeState.height > 0:
247
+ balance_state |= set(data.nodeState["n"].to_list())
248
+ region_all: set[str] = set()
249
+ for m in region_membership.values():
250
+ region_all |= m["nodes"]
251
+ return balance_state - region_all
252
+
253
+
254
+ # ---------------------------------------------------------------------------
255
+ # Helpers for filtering polars frames / Params
256
+ # ---------------------------------------------------------------------------
257
+
258
+
259
+ def _is_in_keep(col: str, keep: set[str]) -> pl.Expr:
260
+ """Membership test for ``pl.col(col)`` against *keep* that tolerates
261
+ keep elements outside the column's *original* Enum vocabulary by
262
+ upcasting the column to the live (widened) Enum first.
263
+
264
+ ``keep`` is built by the region splitter and may include synthetic
265
+ virtual-entity tokens (``hf_pipe_*`` / ``*__export__*`` /
266
+ ``*__import__*``) that the Spine-DB-derived axis_enums don't
267
+ contain. ``split()`` widens the global axis_enums ContextVar to
268
+ include those virtual tokens before the filter runs; here we
269
+ upcast the column to that widened vocabulary via
270
+ :func:`cast_dim` (``enums=None`` reads the live ContextVar). The
271
+ upcast is a strict superset operation (every original value is in
272
+ the wider Enum), the ``is_in`` then succeeds natively, and the
273
+ filter's output retains the widened Enum dtype.
274
+ """
275
+ return cast_dim(pl.col(col), None, col).is_in(list(keep))
276
+
277
+
278
+ def _filter_frame(df: pl.DataFrame | None, col: str,
279
+ keep: set[str]) -> pl.DataFrame | None:
280
+ if df is None:
281
+ return None
282
+ if col not in df.columns:
283
+ return df
284
+ return df.filter(_is_in_keep(col, keep))
285
+
286
+
287
+ def _filter_frame_multi(df: pl.DataFrame | None,
288
+ cond_cols: list[tuple[str, set[str]]]) -> pl.DataFrame | None:
289
+ if df is None:
290
+ return None
291
+ out = df
292
+ for col, keep in cond_cols:
293
+ if col in out.columns:
294
+ out = out.filter(_is_in_keep(col, keep))
295
+ return out
296
+
297
+
298
+ def _filter_param(p: Param | None, col: str,
299
+ keep: set[str]) -> Param | None:
300
+ if p is None:
301
+ return None
302
+ if col not in p.dims:
303
+ return p
304
+ new_frame = p.frame.filter(_is_in_keep(col, keep))
305
+ return Param(p.dims, new_frame, name=p.name)
306
+
307
+
308
+ # ---------------------------------------------------------------------------
309
+ # Cross-region classification
310
+ # ---------------------------------------------------------------------------
311
+
312
+
313
+ def _classify_arcs(
314
+ pss: pl.DataFrame, region_nodes: dict[str, set[str]],
315
+ master_nodes: "frozenset[str] | set[str]" = frozenset(),
316
+ ) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame, pl.DataFrame]:
317
+ """Tag each (p, source, sink) row with its source-region and
318
+ sink-region, then classify into four classes:
319
+
320
+ * **local** — everything not in the three classes below (stays in
321
+ its region's frames);
322
+ * **cross-region** — both endpoints in (different) regions; severed
323
+ into an export + import half-flow pair (returned as ``cross``);
324
+ * **region↔master** — one endpoint in a region, the other a
325
+ master-hosted node; severed into exactly ONE half-flow on the
326
+ region side (returned as ``region_master``);
327
+ * **master-local** — at least one master-hosted endpoint and NO
328
+ region endpoint (the other side is master-hosted, shared, or a
329
+ non-node token); dropped from every region and NOT half-flowed —
330
+ the master keeps the whole arc (returned as ``master_local``).
331
+
332
+ Nodes not in any region (and not master-hosted) are "shared" — an
333
+ arc with a shared endpoint is treated as local-to-the-other-region
334
+ (it stays in that region's frames and is not a coupling).
335
+
336
+ With the default empty ``master_nodes`` the ``region_master`` /
337
+ ``master_local`` frames are empty and ``pss_tagged`` / ``cross``
338
+ are byte-identical to the historical 2-way behaviour.
339
+ """
340
+ # Build a node→region map; nodes outside any region map to None.
341
+ node_region: dict[str, str | None] = {}
342
+ for r, ns in region_nodes.items():
343
+ for n in ns:
344
+ node_region[n] = r
345
+ src_r = pss["source"].map_elements(
346
+ lambda n: node_region.get(n), return_dtype=pl.Utf8)
347
+ snk_r = pss["sink"].map_elements(
348
+ lambda n: node_region.get(n), return_dtype=pl.Utf8)
349
+ pss_tagged = pss.with_columns(
350
+ _src_region=src_r, _snk_region=snk_r,
351
+ )
352
+ cross = pss_tagged.filter(
353
+ pl.col("_src_region").is_not_null() &
354
+ pl.col("_snk_region").is_not_null() &
355
+ (pl.col("_src_region") != pl.col("_snk_region"))
356
+ )
357
+ if not master_nodes:
358
+ empty = pss_tagged.head(0)
359
+ return pss_tagged, cross, empty, empty
360
+ # Vocab-independent membership test (mirror the map_elements-on-raw-
361
+ # strings robustness of the region tagging above: no Enum cast, so a
362
+ # stale global axis-enum vocabulary cannot null the master tokens).
363
+ src_m = pss["source"].cast(pl.Utf8).is_in(list(master_nodes))
364
+ snk_m = pss["sink"].cast(pl.Utf8).is_in(list(master_nodes))
365
+ pss_tagged = pss_tagged.with_columns(_src_master=src_m, _snk_master=snk_m)
366
+ region_master = pss_tagged.filter(
367
+ (pl.col("_src_region").is_not_null() & pl.col("_snk_master"))
368
+ | (pl.col("_src_master") & pl.col("_snk_region").is_not_null())
369
+ )
370
+ master_local = pss_tagged.filter(
371
+ (pl.col("_src_master") | pl.col("_snk_master"))
372
+ & pl.col("_src_region").is_null()
373
+ & pl.col("_snk_region").is_null()
374
+ )
375
+ return pss_tagged, cross, region_master, master_local
376
+
377
+
378
+ def _make_half_flows(
379
+ cross_arcs: pl.DataFrame,
380
+ region_master_arcs: pl.DataFrame | None = None,
381
+ ) -> dict[str, list[HalfFlow]]:
382
+ """For each cross-region arc, produce two HalfFlow records (one
383
+ per region).
384
+
385
+ For each region↔master arc in *region_master_arcs* (master-hosted
386
+ mode), produce exactly ONE HalfFlow — on the region side: an
387
+ export when the region node is the arc's source, an import when it
388
+ is the sink. Naming reuses the region-side stem of the paired
389
+ convention (the master side is implicit; no master-side virtual
390
+ entity exists — the master keeps the whole original arc).
391
+ """
392
+ out: dict[str, list[HalfFlow]] = {}
393
+ for r in cross_arcs.iter_rows(named=True):
394
+ p = r["p"]
395
+ s = r["source"]
396
+ k = r["sink"]
397
+ ra = r["_src_region"]
398
+ rb = r["_snk_region"]
399
+ # Naming: replicate flextool's convention loosely — the virtual
400
+ # *node* uses the (p, terminal, region) stem; the virtual
401
+ # *connection* uses the same stem with an ``hf_`` prefix. We
402
+ # disambiguate per-direction by encoding the original (s, k)
403
+ # pair into the connection id so two-direction pipes don't
404
+ # collide.
405
+ ve_node = f"{p}__{s}__{k}__export__{ra}"
406
+ vi_node = f"{p}__{s}__{k}__import__{rb}"
407
+ ve_conn = f"hf_{p}__{s}__{k}__export__{ra}"
408
+ vi_conn = f"hf_{p}__{s}__{k}__import__{rb}"
409
+ out.setdefault(ra, []).append(HalfFlow(
410
+ region=ra, side="export",
411
+ original_p=p, original_source=s, original_sink=k,
412
+ in_region_node=s,
413
+ virtual_node=ve_node,
414
+ virtual_p=ve_conn,
415
+ virtual_arc_source=s,
416
+ virtual_arc_sink=ve_node,
417
+ ))
418
+ out.setdefault(rb, []).append(HalfFlow(
419
+ region=rb, side="import",
420
+ original_p=p, original_source=s, original_sink=k,
421
+ in_region_node=k,
422
+ virtual_node=vi_node,
423
+ virtual_p=vi_conn,
424
+ virtual_arc_source=vi_node,
425
+ virtual_arc_sink=k,
426
+ ))
427
+ if region_master_arcs is not None and region_master_arcs.height > 0:
428
+ for r in region_master_arcs.iter_rows(named=True):
429
+ p = r["p"]
430
+ s = r["source"]
431
+ k = r["sink"]
432
+ ra = r["_src_region"]
433
+ rb = r["_snk_region"]
434
+ if ra is not None:
435
+ # Region node is the SOURCE ⇒ export half-flow in ra.
436
+ ve_node = f"{p}__{s}__{k}__export__{ra}"
437
+ ve_conn = f"hf_{p}__{s}__{k}__export__{ra}"
438
+ out.setdefault(ra, []).append(HalfFlow(
439
+ region=ra, side="export",
440
+ original_p=p, original_source=s, original_sink=k,
441
+ in_region_node=s,
442
+ virtual_node=ve_node,
443
+ virtual_p=ve_conn,
444
+ virtual_arc_source=s,
445
+ virtual_arc_sink=ve_node,
446
+ ))
447
+ else:
448
+ # Region node is the SINK ⇒ import half-flow in rb.
449
+ vi_node = f"{p}__{s}__{k}__import__{rb}"
450
+ vi_conn = f"hf_{p}__{s}__{k}__import__{rb}"
451
+ out.setdefault(rb, []).append(HalfFlow(
452
+ region=rb, side="import",
453
+ original_p=p, original_source=s, original_sink=k,
454
+ in_region_node=k,
455
+ virtual_node=vi_node,
456
+ virtual_p=vi_conn,
457
+ virtual_arc_source=vi_node,
458
+ virtual_arc_sink=k,
459
+ ))
460
+ return out
461
+
462
+
463
+ # ---------------------------------------------------------------------------
464
+ # Master-hosted-node validation + region scrubbing
465
+ # ---------------------------------------------------------------------------
466
+
467
+
468
+ def _master_local_procs(
469
+ pss: pl.DataFrame, master_local: pl.DataFrame,
470
+ ) -> set[str]:
471
+ """Processes whose EVERY ``process_source_sink`` arc is master-local
472
+ (critique F3). Those procs live wholly in the master: regions must
473
+ carry neither their arcs NOR their entity rows (invest sets, annuity
474
+ params, cost rows)."""
475
+ if master_local.height == 0:
476
+ return set()
477
+ ml_triples: set[tuple[str, str, str]] = {
478
+ (r["p"], r["source"], r["sink"])
479
+ for r in master_local.iter_rows(named=True)
480
+ }
481
+ ml_procs = {t[0] for t in ml_triples}
482
+ triples_by_p: dict[str, set[tuple[str, str, str]]] = {}
483
+ for r in pss.iter_rows(named=True):
484
+ if r["p"] in ml_procs:
485
+ triples_by_p.setdefault(r["p"], set()).add(
486
+ (r["p"], r["source"], r["sink"]))
487
+ return {p for p, ts in triples_by_p.items() if ts <= ml_triples}
488
+
489
+
490
+ def _validate_no_straddling_units(
491
+ data: FlexData,
492
+ all_region_nodes: set[str],
493
+ master_nodes: "frozenset[str] | set[str]",
494
+ ) -> None:
495
+ """Hard-error on any UNIT touching both a region node and a
496
+ master-hosted node (design decision D-a).
497
+
498
+ Aggregated PER PROCESS across ALL its ``process_source_sink`` rows —
499
+ NOT per arc: a unit with one master-local arc plus one purely
500
+ in-region arc straddles *as an entity* while having no individually
501
+ straddling arc, and severing any unit arc silently loses conversion
502
+ terms. Never a silent degrade (the fix_start precedent is the
503
+ anti-pattern this guards against).
504
+ """
505
+ if data.process_source_sink is None:
506
+ return
507
+ units: set[str] = set()
508
+ if data.process_unit is not None and data.process_unit.height > 0:
509
+ units |= set(data.process_unit["p"].to_list())
510
+ if (getattr(data, "process_indirect", None) is not None
511
+ and data.process_indirect.height > 0):
512
+ units |= set(data.process_indirect["p"].to_list())
513
+ if not units:
514
+ return
515
+ endpoints_by_unit: dict[str, set[str]] = {}
516
+ for r in data.process_source_sink.iter_rows(named=True):
517
+ if r["p"] in units:
518
+ endpoints_by_unit.setdefault(r["p"], set()).update(
519
+ (r["source"], r["sink"]))
520
+ for unit in sorted(endpoints_by_unit):
521
+ eps = endpoints_by_unit[unit]
522
+ region_touch = eps & all_region_nodes
523
+ master_touch = eps & master_nodes
524
+ if region_touch and master_touch:
525
+ raise RuntimeError(
526
+ f"split: unit {unit!r} straddles the region/master "
527
+ f"boundary — across its arcs it touches region node(s) "
528
+ f"{sorted(region_touch)} AND master-hosted node(s) "
529
+ f"{sorted(master_touch)}. A unit cannot be severed "
530
+ f"between a region subproblem and the Benders master. "
531
+ f"Insert a handover CONNECTION between the region-side "
532
+ f"node and the master-hosted node (the handover-"
533
+ f"connection pattern) so every boundary arc is a "
534
+ f"connection, and keep the unit's arcs wholly on one "
535
+ f"side."
536
+ )
537
+
538
+
539
+ def _user_constraint_sides(
540
+ data: FlexData,
541
+ all_region_nodes: set[str],
542
+ master_nodes: "frozenset[str] | set[str]",
543
+ master_local_procs: set[str],
544
+ ) -> tuple[dict[str, set[str]], dict[str, set[str]]]:
545
+ """Aggregate, per user-constraint id (``cn``), the decomposition
546
+ sides ("master" / "region") of every referenced entity across every
547
+ user-constraint frame.
548
+
549
+ Shared by the mixed-constraint hard error
550
+ (:func:`_validate_user_constraints`) and the master/region row
551
+ partition (:func:`_master_side_constraint_ids`). Entities on
552
+ neither side (shared nodes, non-node tokens) are neutral.
553
+
554
+ Returns ``(sides_by_cn, refs_by_cn)``; ``refs_by_cn`` carries
555
+ human-readable ``"entity (side)"`` strings for error messages.
556
+ """
557
+ sides_by_cn: dict[str, set[str]] = {}
558
+ refs_by_cn: dict[str, set[str]] = {}
559
+
560
+ def _node_side(n: str) -> str | None:
561
+ if n in master_nodes:
562
+ return "master"
563
+ if n in all_region_nodes:
564
+ return "region"
565
+ return None
566
+
567
+ def _add(cn: str, entity: str, side: str | None) -> None:
568
+ if side is None:
569
+ return
570
+ sides_by_cn.setdefault(cn, set()).add(side)
571
+ refs_by_cn.setdefault(cn, set()).add(f"{entity} ({side})")
572
+
573
+ def _frame_of(obj) -> pl.DataFrame | None:
574
+ if obj is None:
575
+ return None
576
+ f = obj.frame if isinstance(obj, Param) else obj
577
+ return f if f.height > 0 else None
578
+
579
+ # Arc-keyed references: (p, source, sink, cn). Each row contributes
580
+ # the sides of BOTH terminal nodes (a row on a region↔master
581
+ # coupling arc is itself mixed) plus the master-local proc side.
582
+ for fld in ("flow_constraint_idx", "p_flow_constraint_coef"):
583
+ f = _frame_of(getattr(data, fld, None))
584
+ if f is None or "cn" not in f.columns:
585
+ continue
586
+ for r in f.iter_rows(named=True):
587
+ cn = r["cn"]
588
+ _add(cn, r["source"], _node_side(r["source"]))
589
+ _add(cn, r["sink"], _node_side(r["sink"]))
590
+ if r["p"] in master_local_procs:
591
+ _add(cn, r["p"], "master")
592
+ # Node-keyed references: (n, cn).
593
+ for fld in ("p_node_constraint_state_coeff",
594
+ "p_node_constraint_invested_capacity_coeff",
595
+ "p_node_constraint_prebuilt_capacity_coeff"):
596
+ f = _frame_of(getattr(data, fld, None))
597
+ if f is None or "cn" not in f.columns:
598
+ continue
599
+ for r in f.iter_rows(named=True):
600
+ _add(r["cn"], r["n"], _node_side(r["n"]))
601
+ # Process-keyed references: (p, cn). A process's side is the union
602
+ # of its arc-endpoint node sides (a master-local proc is master; a
603
+ # coupling connection contributes both sides and therefore raises).
604
+ proc_sides: dict[str, set[str]] = {}
605
+ if data.process_source_sink is not None:
606
+ for r in data.process_source_sink.iter_rows(named=True):
607
+ s = proc_sides.setdefault(r["p"], set())
608
+ for n in (r["source"], r["sink"]):
609
+ side = _node_side(n)
610
+ if side is not None:
611
+ s.add(side)
612
+ for p in master_local_procs:
613
+ proc_sides.setdefault(p, set()).add("master")
614
+ for fld in ("p_process_constraint_invested_capacity_coeff",
615
+ "p_process_constraint_prebuilt_capacity_coeff"):
616
+ f = _frame_of(getattr(data, fld, None))
617
+ if f is None or "cn" not in f.columns:
618
+ continue
619
+ for r in f.iter_rows(named=True):
620
+ for side in proc_sides.get(r["p"], set()):
621
+ _add(r["cn"], r["p"], side)
622
+
623
+ return sides_by_cn, refs_by_cn
624
+
625
+
626
+ def _validate_user_constraints(
627
+ data: FlexData,
628
+ all_region_nodes: set[str],
629
+ master_nodes: "frozenset[str] | set[str]",
630
+ master_local_procs: set[str],
631
+ ) -> None:
632
+ """Hard-error on any user constraint referencing both master-side
633
+ and region-side entities.
634
+
635
+ Sides are aggregated per constraint id (``cn``) across every
636
+ user-constraint frame (:func:`_user_constraint_sides`): a constraint
637
+ mixing a master-hosted node (or a master-local process) with a
638
+ region node (or a region process) cannot live whole on either side
639
+ of the decomposition — splitting it would silently lose terms.
640
+ """
641
+ sides_by_cn, refs_by_cn = _user_constraint_sides(
642
+ data, all_region_nodes, master_nodes, master_local_procs)
643
+ for cn in sorted(sides_by_cn):
644
+ if {"region", "master"} <= sides_by_cn[cn]:
645
+ raise RuntimeError(
646
+ f"split: user constraint {cn!r} references both "
647
+ f"master-side and region-side entities: "
648
+ f"{sorted(refs_by_cn[cn])}. A constraint cannot be "
649
+ f"split between a region subproblem and the Benders "
650
+ f"master — rewrite it to reference entities on one "
651
+ f"side only (e.g. via the handover-connection pattern)."
652
+ )
653
+
654
+
655
+ def _master_side_constraint_ids(
656
+ data: FlexData,
657
+ all_region_nodes: set[str],
658
+ master_nodes: "frozenset[str] | set[str]",
659
+ master_local_procs: set[str],
660
+ ) -> set[str]:
661
+ """User-constraint ids (``cn``) whose EVERY sided entity reference
662
+ is master-side. Those constraints live whole in the Benders master:
663
+ :func:`master_network_data` keeps exactly their rows and
664
+ :func:`split` drops them from every region (a region copy would
665
+ degenerate to ``0 sense constant`` once its master-keyed coefficient
666
+ rows are scrubbed). Assumes :func:`_validate_user_constraints` ran
667
+ (no mixed constraints)."""
668
+ sides_by_cn, _ = _user_constraint_sides(
669
+ data, all_region_nodes, master_nodes, master_local_procs)
670
+ return {cn for cn, sides in sides_by_cn.items() if sides == {"master"}}
671
+
672
+
673
+ #: Group-feature SET fields (``(g,)``): a group present in one of these
674
+ #: activates the corresponding group-level constraint family
675
+ #: (capacity margin / inertia / non-sync). Membership is ``group_node``.
676
+ _GROUP_FEATURE_SET_FIELDS: tuple[str, ...] = (
677
+ "groupCapacityMargin", "groupInertia", "groupNonSync",
678
+ )
679
+
680
+ #: Group-feature PARAM fields (``(g,)`` / ``(g, d)``) filtered alongside
681
+ #: the set fields when partitioning feature groups to the master.
682
+ _GROUP_FEATURE_PARAM_FIELDS: tuple[str, ...] = (
683
+ "p_inv_group_cap", "p_group_capacity_for_scaling",
684
+ "pdGroup_capacity_margin",
685
+ )
686
+
687
+
688
+ def _master_side_feature_groups(
689
+ data: FlexData,
690
+ all_region_nodes: set[str],
691
+ master_nodes: "frozenset[str] | set[str]",
692
+ ) -> set[str]:
693
+ """Feature-carrying groups whose member nodes are ALL master-hosted.
694
+
695
+ A group in one of the :data:`_GROUP_FEATURE_SET_FIELDS` sets with
696
+ at least one master-hosted member must have EVERY member
697
+ master-hosted — the group constraint sums over its members and
698
+ cannot be enforced whole on either side otherwise (hard error,
699
+ never a silent degrade). Groups with no master members are left to
700
+ the regions (today's semantics); bare ``group_node`` membership
701
+ rows of non-feature groups are filtered silently.
702
+ """
703
+ feature_gs: set[str] = set()
704
+ for fld in _GROUP_FEATURE_SET_FIELDS:
705
+ f = getattr(data, fld, None)
706
+ if f is not None and f.height > 0:
707
+ feature_gs |= set(f["g"].cast(pl.Utf8).to_list())
708
+ if not feature_gs:
709
+ return set()
710
+ members_by_g: dict[str, set[str]] = {}
711
+ gn = data.group_node
712
+ if gn is not None:
713
+ for r in gn.iter_rows(named=True):
714
+ members_by_g.setdefault(r["g"], set()).add(r["n"])
715
+ master = set(master_nodes)
716
+ kept: set[str] = set()
717
+ for g in sorted(feature_gs):
718
+ members = members_by_g.get(g, set())
719
+ master_hit = members & master
720
+ if not master_hit:
721
+ continue
722
+ if members <= master:
723
+ kept.add(g)
724
+ continue
725
+ raise RuntimeError(
726
+ f"split: feature-carrying group {g!r} straddles the "
727
+ f"region/master boundary — it has master-hosted member "
728
+ f"node(s) {sorted(master_hit)} AND non-master member(s) "
729
+ f"{sorted(members - master)}. A group-level constraint "
730
+ f"(capacity margin / inertia / non-sync) sums over its "
731
+ f"members and cannot be split between a region subproblem "
732
+ f"and the Benders master — regroup the nodes so every "
733
+ f"member is on one side."
734
+ )
735
+ return kept
736
+
737
+
738
+ def _co2_master_partition(
739
+ data: FlexData,
740
+ master_local_triples: set[tuple[str, str, str]],
741
+ region_master_triples: set[tuple[str, str, str]],
742
+ ) -> set[str]:
743
+ """Validate the CO2 frames against the master boundary and return
744
+ the CO2-capped groups whose EVERY capped flow is a master-local arc
745
+ (those cap constraints live whole in the master).
746
+
747
+ Hard errors (never a silent degrade):
748
+
749
+ * a CO2-priced flow row on a region↔master coupling arc — the arc's
750
+ flow lives natively in the master while the region holds the
751
+ half-flow, so neither side can carry the CO2 cost whole;
752
+ * a capped group mixing master-local and region-side flows — the
753
+ shared cap cannot be enforced on either side alone.
754
+ """
755
+ for fld in ("flow_from_co2_priced", "flow_from_co2_priced_noEff"):
756
+ df = getattr(data, fld, None)
757
+ if df is None or df.height == 0:
758
+ continue
759
+ for r in df.iter_rows(named=True):
760
+ t = (r["p"], r["source"], r["sink"])
761
+ if t in region_master_triples:
762
+ raise RuntimeError(
763
+ f"split: CO2-priced flow {t} sits on a "
764
+ f"region↔master coupling arc — the CO2 cost "
765
+ f"cannot be split between a region subproblem and "
766
+ f"the Benders master. Keep CO2-priced flows off "
767
+ f"the handover connections."
768
+ )
769
+ master_flags_by_g: dict[str, list[bool]] = {}
770
+ for fld in ("flow_from_co2_capped", "flow_from_co2_capped_noEff",
771
+ "flow_from_co2_capped_total",
772
+ "flow_from_co2_capped_total_noEff"):
773
+ df = getattr(data, fld, None)
774
+ if df is None or df.height == 0:
775
+ continue
776
+ for r in df.iter_rows(named=True):
777
+ t = (r["p"], r["source"], r["sink"])
778
+ master_flags_by_g.setdefault(r["g"], []).append(
779
+ t in master_local_triples)
780
+ kept: set[str] = set()
781
+ for g in sorted(master_flags_by_g):
782
+ flags = master_flags_by_g[g]
783
+ if all(flags):
784
+ kept.add(g)
785
+ elif any(flags):
786
+ raise RuntimeError(
787
+ f"split: CO2-capped group {g!r} straddles the "
788
+ f"region/master boundary — it caps both master-local "
789
+ f"and region-side flows. A shared CO2 cap cannot be "
790
+ f"split between a region subproblem and the Benders "
791
+ f"master — regroup so every capped flow is on one side."
792
+ )
793
+ return kept
794
+
795
+
796
+ def _drop_master_rows(
797
+ rd: FlexData,
798
+ master_nodes: "frozenset[str] | set[str]",
799
+ master_procs: set[str],
800
+ master_cns: "frozenset[str] | set[str]" = frozenset(),
801
+ master_groups: "frozenset[str] | set[str]" = frozenset(),
802
+ ) -> FlexData:
803
+ """Scrub every region frame/Param of master-hosted content: rows
804
+ keyed to a master-hosted node (``n`` axis), a master-local process
805
+ (``p`` axis), either (entity ``e`` axis), an all-master user
806
+ constraint (``cn`` axis) or an all-master feature / CO2-cap group
807
+ (``g`` axis).
808
+
809
+ ``keep_nodes`` / ``keep_procs`` filtering already excludes master
810
+ entities from the frames :func:`_build_region_data` filters
811
+ explicitly; this pass additionally covers the frames the splitter
812
+ historically carried through whole under shared-replicate semantics
813
+ (invest/annuity/cost frames, ``process_indirect``, node profiles,
814
+ …) — regions must carry NO rows for master-hosted entities (F3).
815
+ The ``cn`` drop covers the constraint-id-only frames
816
+ (``p_constraint_constant``, ``cdt_eq/le/ge``) whose master-keyed
817
+ coefficient rows are scrubbed by the ``n``/``p`` drops — without it
818
+ a region copy of an all-master constraint degenerates to
819
+ ``0 sense constant``. The ``g`` drop is the analogous guard for
820
+ all-master group features (a region copy with empty membership
821
+ would still charge the feature's slack penalty).
822
+ Arc rows touching a master node are dropped separately via the
823
+ coupling/master-local triples in ``cross_arcs_by_pss``.
824
+ """
825
+ node_drop = set(master_nodes)
826
+ proc_drop = set(master_procs)
827
+ entity_drop = node_drop | proc_drop
828
+ cn_drop = set(master_cns)
829
+ g_drop = set(master_groups)
830
+ if not (entity_drop or cn_drop or g_drop):
831
+ return rd
832
+ for f in dataclasses.fields(rd):
833
+ v = getattr(rd, f.name)
834
+ if v is None:
835
+ continue
836
+ if isinstance(v, Param):
837
+ frame = v.frame
838
+ elif isinstance(v, pl.DataFrame):
839
+ frame = v
840
+ else:
841
+ continue
842
+ changed = False
843
+ for col, drop in (("n", node_drop), ("p", proc_drop),
844
+ ("e", entity_drop), ("cn", cn_drop),
845
+ ("g", g_drop)):
846
+ if drop and col in frame.columns:
847
+ frame = frame.filter(~_is_in_keep(col, drop))
848
+ changed = True
849
+ if not changed:
850
+ continue
851
+ if isinstance(v, Param):
852
+ setattr(rd, f.name, Param(v.dims, frame, name=v.name))
853
+ else:
854
+ setattr(rd, f.name, frame)
855
+ return rd
856
+
857
+
858
+ # ---------------------------------------------------------------------------
859
+ # Per-region splitter
860
+ # ---------------------------------------------------------------------------
861
+
862
+
863
+ def _build_region_data(
864
+ src: FlexData,
865
+ region: str,
866
+ keep_nodes: set[str],
867
+ keep_procs: set[str],
868
+ half_flows: list[HalfFlow],
869
+ cross_arcs_by_pss: set[tuple[str, str, str]],
870
+ benders_uncap_cross_region: bool = False,
871
+ *,
872
+ master_hosted_nodes: "frozenset[str] | set[str]" = frozenset(),
873
+ master_local_procs: "set[str] | frozenset[str]" = frozenset(),
874
+ master_cns: "frozenset[str] | set[str]" = frozenset(),
875
+ master_groups: "frozenset[str] | set[str]" = frozenset(),
876
+ ) -> FlexData:
877
+ """Construct one region's :class:`FlexData` by filtering+rewriting
878
+ the whole-system frames/Params.
879
+
880
+ ``keep_nodes``/``keep_procs`` are the in-region+shared sets.
881
+ ``cross_arcs_by_pss`` is the SET of (p, source, sink) tuples to
882
+ REMOVE from this region's process frames (they're being replaced
883
+ by half-flow virtual arcs; master-hosted mode also routes the
884
+ region↔master coupling arcs and the master-local arcs through it).
885
+ ``master_hosted_nodes`` / ``master_local_procs`` / ``master_cns`` /
886
+ ``master_groups`` (master-hosted mode only; empty by default ⇒
887
+ byte-identical path) trigger the :func:`_drop_master_rows` scrub so
888
+ the region carries NO rows for master-hosted entities, all-master
889
+ user constraints or all-master feature/CO2-cap groups.
890
+ """
891
+ # Start by shallow-copying the dataclass and clearing fields we'll
892
+ # explicitly rewrite.
893
+ new = dataclasses.replace(src)
894
+
895
+ # ---- Filter primary entity sets ----
896
+ new.nodeBalance = _filter_frame(src.nodeBalance, "n", keep_nodes)
897
+ # Phase E.3: ``nodeBalance_dt`` is no longer materialised on src; the
898
+ # filtered ``new.nodeBalance`` is the only set we need, and
899
+ # ``_pdt_join.compute_nodeBalance_dt(new)`` produces the cross-join
900
+ # on demand downstream.
901
+ new.nodeBalance_dt = None
902
+ new.p_inflow = _filter_param(src.p_inflow, "n", keep_nodes)
903
+ new.p_penalty_up = _filter_param(src.p_penalty_up, "n", keep_nodes)
904
+ new.p_penalty_down = _filter_param(src.p_penalty_down, "n", keep_nodes)
905
+
906
+ # ---- Filter process_source_sink and arc-side frames ----
907
+ def _drop_cross(df: pl.DataFrame | None) -> pl.DataFrame | None:
908
+ if df is None:
909
+ return None
910
+ if not all(c in df.columns for c in ("p", "source", "sink")):
911
+ return df
912
+ if not cross_arcs_by_pss:
913
+ return df
914
+ # Build a 3-col anti-join key.
915
+ _enums = getattr(src, "_axis_enums", None)
916
+ key_df = pl.DataFrame({
917
+ "p": [t[0] for t in cross_arcs_by_pss],
918
+ "source": [t[1] for t in cross_arcs_by_pss],
919
+ "sink": [t[2] for t in cross_arcs_by_pss],
920
+ }, schema={"p": schema_dtype(_enums, "p"),
921
+ "source": schema_dtype(_enums, "source"),
922
+ "sink": schema_dtype(_enums, "sink")})
923
+ return df.join(key_df, on=("p", "source", "sink"), how="anti")
924
+
925
+ def _filter_arc_by_proc(df: pl.DataFrame | None) -> pl.DataFrame | None:
926
+ df = _drop_cross(df)
927
+ if df is None or "p" not in df.columns:
928
+ return df
929
+ return df.filter(_is_in_keep("p", keep_procs))
930
+
931
+ def _filter_param_arc(p: Param | None) -> Param | None:
932
+ if p is None:
933
+ return None
934
+ if "p" not in p.dims:
935
+ return p
936
+ # Drop cross-region rows by triple-key, then filter to keep procs.
937
+ f = p.frame
938
+ if cross_arcs_by_pss and all(c in f.columns
939
+ for c in ("p", "source", "sink")):
940
+ _enums = getattr(src, "_axis_enums", None)
941
+ key_df = pl.DataFrame({
942
+ "p": [t[0] for t in cross_arcs_by_pss],
943
+ "source": [t[1] for t in cross_arcs_by_pss],
944
+ "sink": [t[2] for t in cross_arcs_by_pss],
945
+ }, schema={"p": schema_dtype(_enums, "p"),
946
+ "source": schema_dtype(_enums, "source"),
947
+ "sink": schema_dtype(_enums, "sink")})
948
+ f = f.join(key_df, on=("p", "source", "sink"), how="anti")
949
+ f = f.filter(_is_in_keep("p", keep_procs))
950
+ return Param(p.dims, f, name=p.name)
951
+
952
+ new.process_source_sink = _filter_arc_by_proc(src.process_source_sink)
953
+ new.process_source_sink_eff = _filter_arc_by_proc(src.process_source_sink_eff)
954
+ new.process_source_sink_noEff = _filter_arc_by_proc(src.process_source_sink_noEff)
955
+ # Phase E.3: ``pss_dt`` is no longer materialised on src; the filtered
956
+ # ``new.process_source_sink`` is the only set we need, and
957
+ # ``_pdt_join.compute_pss_dt(new)`` produces the cross-join on demand
958
+ # downstream. Half-flow injection below ALSO needs a pss_dt view; it
959
+ # builds one locally from src for the arc-dt extraction.
960
+ new.pss_dt = None
961
+ new.flow_to_n = _filter_arc_by_proc(src.flow_to_n)
962
+ new.flow_from_n = _filter_arc_by_proc(src.flow_from_n)
963
+ new.flow_from_nodeBalance_eff = _filter_arc_by_proc(src.flow_from_nodeBalance_eff)
964
+ new.flow_from_nodeBalance_noEff = _filter_arc_by_proc(src.flow_from_nodeBalance_noEff)
965
+ new.flow_from_commodity_eff = _filter_arc_by_proc(src.flow_from_commodity_eff)
966
+ new.flow_from_commodity_noEff = _filter_arc_by_proc(src.flow_from_commodity_noEff)
967
+ new.flow_to_commodity = _filter_arc_by_proc(src.flow_to_commodity)
968
+
969
+ new.p_unitsize = _filter_param(src.p_unitsize, "p", keep_procs)
970
+ new.p_flow_upper = _filter_param_arc(src.p_flow_upper)
971
+ new.p_flow_upper_existing = _filter_param_arc(src.p_flow_upper_existing)
972
+ new.p_slope = _filter_param(src.p_slope, "p", keep_procs)
973
+ new.p_process_existing_count = _filter_param(src.p_process_existing_count, "p", keep_procs)
974
+ new.p_process_availability = _filter_param(src.p_process_availability, "p", keep_procs)
975
+
976
+ # ---- Profiles (filter by p) ----
977
+ new.process_profile_upper = _filter_frame(src.process_profile_upper, "p", keep_procs)
978
+ new.process_profile_lower = _filter_frame(src.process_profile_lower, "p", keep_procs)
979
+ new.process_profile_fixed = _filter_frame(src.process_profile_fixed, "p", keep_procs)
980
+
981
+ # ---- Storage / nodeState filtered to in-region nodes ----
982
+ new.nodeState = _filter_frame(src.nodeState, "n", keep_nodes)
983
+ # Phase E.3: ``nodeState_dt`` is no longer materialised on src; the
984
+ # filtered ``new.nodeState`` is the only set we need, and
985
+ # ``_pdt_join.compute_nodeState_dt(new)`` produces the cross-join
986
+ # on demand downstream. ``nodeState_first_dt`` is still
987
+ # materialised (small one-row-per-node slice; see ``_load_storage``).
988
+ new.nodeState_dt = None
989
+ new.nodeState_first_dt = _filter_frame(src.nodeState_first_dt, "n", keep_nodes)
990
+ new.storage_bind_within_timeblock = _filter_frame(src.storage_bind_within_timeblock, "n", keep_nodes)
991
+ new.storage_bind_forward_only = _filter_frame(src.storage_bind_forward_only, "n", keep_nodes)
992
+ new.storage_bind_within_solve = _filter_frame(src.storage_bind_within_solve, "n", keep_nodes)
993
+ new.storage_bind_within_solve_blended_weights = _filter_frame(
994
+ getattr(src, "storage_bind_within_solve_blended_weights", None),
995
+ "n", keep_nodes,
996
+ )
997
+ # Phase D — added with the new variant's constraint wiring.
998
+ new.storage_bind_forward_only_blended_weights = _filter_frame(
999
+ getattr(src, "storage_bind_forward_only_blended_weights", None),
1000
+ "n", keep_nodes,
1001
+ )
1002
+ # Phase E — per-period cyclic-closure variant.
1003
+ new.storage_bind_within_period_blended_weights = _filter_frame(
1004
+ getattr(src, "storage_bind_within_period_blended_weights", None),
1005
+ "n", keep_nodes,
1006
+ )
1007
+ new.storage_fix_start = _filter_frame(src.storage_fix_start, "n", keep_nodes)
1008
+ new.nodeStateBlock = _filter_frame(src.nodeStateBlock, "n", keep_nodes)
1009
+ new.nodeState_rp = _filter_frame(
1010
+ getattr(src, "nodeState_rp", None), "n", keep_nodes,
1011
+ )
1012
+ # The remaining RP-blended-weights fields (rp_base_period_set,
1013
+ # rp_base_chain, rp_base_first, rp_base_last, rp_block_first,
1014
+ # p_rp_last_step, rp_base__rep) are solve-data-keyed (period / step
1015
+ # / base / rep) — not entity-keyed — so the ``dataclasses.replace``
1016
+ # shallow copy above already carries them through unchanged.
1017
+ new.nodeState_last_dt = _filter_frame(src.nodeState_last_dt, "n", keep_nodes)
1018
+ new.p_state_upper = _filter_param(src.p_state_upper, "n", keep_nodes)
1019
+ new.p_state_unitsize = _filter_param(src.p_state_unitsize, "n", keep_nodes)
1020
+ new.p_state_self_discharge = _filter_param(src.p_state_self_discharge, "n", keep_nodes)
1021
+ new.p_state_start = _filter_param(src.p_state_start, "n", keep_nodes)
1022
+ new.p_state_existing_capacity = _filter_param(src.p_state_existing_capacity, "n", keep_nodes)
1023
+ new.p_node_availability = _filter_param(src.p_node_availability, "n", keep_nodes)
1024
+ new.p_roll_continue_state = _filter_param(src.p_roll_continue_state, "n", keep_nodes)
1025
+ new.p_node_capacity_for_scaling = _filter_param(src.p_node_capacity_for_scaling, "n", keep_nodes)
1026
+
1027
+ # ---- Per-arc block weights (lh2 fixture) ----
1028
+ new.arc_sink_block_dt = _filter_arc_by_proc(src.arc_sink_block_dt) \
1029
+ if hasattr(src, "arc_sink_block_dt") else None
1030
+ new.arc_source_block_dt = _filter_arc_by_proc(src.arc_source_block_dt) \
1031
+ if hasattr(src, "arc_source_block_dt") else None
1032
+ new.p_arc_sink_weight = _filter_param_arc(getattr(src, "p_arc_sink_weight", None))
1033
+ new.p_arc_source_weight = _filter_param_arc(getattr(src, "p_arc_source_weight", None))
1034
+
1035
+ # ---- Drop group_entity / group_node rows referencing dropped entities ----
1036
+ if src.group_entity is not None and "e" in src.group_entity.columns:
1037
+ keep_e = keep_nodes | keep_procs
1038
+ new.group_entity = src.group_entity.filter(
1039
+ _is_in_keep("e", keep_e)
1040
+ )
1041
+ if src.group_node is not None and "n" in src.group_node.columns:
1042
+ new.group_node = src.group_node.filter(
1043
+ _is_in_keep("n", keep_nodes)
1044
+ )
1045
+ new.process_unit = _filter_frame(src.process_unit, "p", keep_procs)
1046
+
1047
+ # ---- Master-hosted mode: scrub master entities from the frames the
1048
+ # splitter otherwise carries through whole (invest/annuity/cost,
1049
+ # process_indirect, node profiles, …). No-op when both sets are
1050
+ # empty (the byte-identical default path).
1051
+ if (master_hosted_nodes or master_local_procs or master_cns
1052
+ or master_groups):
1053
+ new = _drop_master_rows(
1054
+ new, master_hosted_nodes, set(master_local_procs),
1055
+ master_cns=master_cns, master_groups=master_groups)
1056
+
1057
+ # ---- Inject virtual half-flow arcs ----
1058
+ if half_flows:
1059
+ new = _inject_half_flows(
1060
+ new, src, half_flows,
1061
+ benders_uncap_cross_region=benders_uncap_cross_region,
1062
+ )
1063
+
1064
+ return new
1065
+
1066
+
1067
+ #: Benders mode sentinel for the cross-region half-flow ``maxFlow``
1068
+ #: capacity. The real ``f ≤ C·unitsize`` limit lives in the master, so
1069
+ #: the per-region half-flow must be effectively uncapped. The largest
1070
+ #: achievable physical flow is bounded by the connection's
1071
+ #: ``invest_max_total · unitsize`` (and ``v_flow`` is normalised by
1072
+ #: unitsize, so in solver units it is bounded by ``invest_max_total``);
1073
+ #: 1e12 is comfortably ≫ any realistic ``invest_max_total``, so the
1074
+ #: half-flow's ``maxFlow`` row is structurally slack for any flow the
1075
+ #: master could pin — it can never bind and therefore cannot leak a dual
1076
+ #: into the per-region subproblem (Phase-1 Claim 4).
1077
+ _BENDERS_UNCAP_SENTINEL: float = 1e12
1078
+
1079
+
1080
+ def _inject_half_flows(
1081
+ rd: FlexData, src: FlexData, half_flows: list[HalfFlow],
1082
+ *,
1083
+ benders_uncap_cross_region: bool = False,
1084
+ ) -> FlexData:
1085
+ """Add virtual half-flow connections + virtual arcs into the
1086
+ region's frames. Each half-flow gets:
1087
+
1088
+ * one row in ``process_source_sink`` (and the same partition
1089
+ family ``_eff`` or ``_noEff`` as the original arc);
1090
+ * a row per (d, t) in ``pss_dt`` covering the same dt frame as
1091
+ the source data;
1092
+ * a row in ``flow_to_n`` (export: in-region node is sink? no —
1093
+ original source is exported FROM the in-region node, so the
1094
+ half-flow has source=in_region_node, sink=virtual_node; the
1095
+ flow LEAVES the in-region node so it goes into ``flow_from_n``)
1096
+ — i.e. only ``flow_from_n`` for export, only ``flow_to_n`` for
1097
+ import;
1098
+ * unitsize = 1.0 (independent of the original);
1099
+ * flow_upper = original arc's flow_upper (so dispatch can max out
1100
+ the pipe; the master pins the actual flow each iteration).
1101
+ """
1102
+ if not half_flows:
1103
+ return rd
1104
+
1105
+ # Capture the original arc rows so we can pull their (d, t) shape and
1106
+ # flow_upper Param values.
1107
+ # Phase E.3: ``src.pss_dt`` is no longer materialised; build it on
1108
+ # demand from the constituents. Half-flow injection always touches
1109
+ # the full arc-dt grid for the cross-region arcs, so a one-shot build
1110
+ # here is fine.
1111
+ orig_pss_dt = compute_pss_dt(src)
1112
+ orig_flow_upper = src.p_flow_upper
1113
+ orig_flow_upper_existing = src.p_flow_upper_existing
1114
+ orig_unitsize = src.p_unitsize
1115
+ orig_eff = src.process_source_sink_eff
1116
+
1117
+ # Build new pss / pss_dt / flow_*/upper rows for each half-flow.
1118
+ new_pss_rows: list[dict] = []
1119
+ new_pss_eff_rows: list[dict] = []
1120
+ new_pss_noEff_rows: list[dict] = []
1121
+ new_pss_dt_rows: list[dict] = []
1122
+ new_flow_to_n_rows: list[dict] = []
1123
+ new_flow_from_n_rows: list[dict] = []
1124
+ new_flow_from_nb_eff_rows: list[dict] = []
1125
+ new_flow_from_nb_noEff_rows: list[dict] = []
1126
+ new_unitsize_rows: list[dict] = []
1127
+ new_flow_upper_rows: list[dict] = []
1128
+ new_flow_upper_existing_rows: list[dict] = []
1129
+ new_arc_sink_block_dt_rows: list[dict] = []
1130
+ new_p_arc_sink_weight_rows: list[dict] = []
1131
+
1132
+ # Classification: inherit from the original arc. When the original
1133
+ # arc is in process_source_sink_eff (with slope p_slope), so is the
1134
+ # half-flow — that ensures the source-side nodeBalance loses the
1135
+ # right amount of source commodity (source-side flow ×
1136
+ # unitsize × slope) which matches the monolithic. When the original
1137
+ # is noEff, the half-flow stays noEff and the source-side loses
1138
+ # source-side flow × unitsize (no slope). Pre-compute per-(p,
1139
+ # source, sink) classification.
1140
+ orig_eff_keys: set[tuple[str, str, str]] = set()
1141
+ if orig_eff is not None:
1142
+ for r in orig_eff.iter_rows(named=True):
1143
+ orig_eff_keys.add((r["p"], r["source"], r["sink"]))
1144
+
1145
+ new_p_slope_rows: list[dict] = []
1146
+
1147
+ for hf in half_flows:
1148
+ # The dt grid for the new arc inherits the source arc's grid.
1149
+ # Pull it from the original pss_dt rows for (p_orig, source_orig,
1150
+ # sink_orig).
1151
+ if orig_pss_dt is not None:
1152
+ arc_dt = orig_pss_dt.filter(
1153
+ (pl.col("p") == hf.original_p)
1154
+ & (pl.col("source") == hf.original_source)
1155
+ & (pl.col("sink") == hf.original_sink)
1156
+ ).select("d", "t")
1157
+ else:
1158
+ _enums = getattr(src, "_axis_enums", None)
1159
+ arc_dt = pl.DataFrame({"d": [], "t": []},
1160
+ schema={"d": schema_dtype(_enums, "d"),
1161
+ "t": schema_dtype(_enums, "t")})
1162
+
1163
+ # Classification inherits from the original arc.
1164
+ is_eff = (hf.original_p, hf.original_source, hf.original_sink) in orig_eff_keys
1165
+
1166
+ # process_source_sink rows
1167
+ new_pss_rows.append({
1168
+ "p": hf.virtual_p,
1169
+ "source": hf.virtual_arc_source,
1170
+ "sink": hf.virtual_arc_sink,
1171
+ })
1172
+ if is_eff:
1173
+ new_pss_eff_rows.append({
1174
+ "p": hf.virtual_p,
1175
+ "source": hf.virtual_arc_source,
1176
+ "sink": hf.virtual_arc_sink,
1177
+ })
1178
+ # Inherit p_slope rows from the original (p, d, t).
1179
+ if src.p_slope is not None:
1180
+ slope_rows = src.p_slope.frame.filter(
1181
+ pl.col("p") == hf.original_p
1182
+ ).select("p", "d", "t", "value")
1183
+ for r in slope_rows.iter_rows(named=True):
1184
+ new_p_slope_rows.append({
1185
+ "p": hf.virtual_p,
1186
+ "d": r["d"], "t": r["t"],
1187
+ "value": float(r["value"]),
1188
+ })
1189
+ else:
1190
+ new_pss_noEff_rows.append({
1191
+ "p": hf.virtual_p,
1192
+ "source": hf.virtual_arc_source,
1193
+ "sink": hf.virtual_arc_sink,
1194
+ })
1195
+
1196
+ # pss_dt rows
1197
+ for r in arc_dt.iter_rows(named=True):
1198
+ new_pss_dt_rows.append({
1199
+ "p": hf.virtual_p,
1200
+ "source": hf.virtual_arc_source,
1201
+ "sink": hf.virtual_arc_sink,
1202
+ "d": r["d"], "t": r["t"],
1203
+ })
1204
+
1205
+ # Flow direction wiring:
1206
+ # * EXPORT: half-flow goes from in_region_node to virtual_node.
1207
+ # Source-side flow leaves in_region_node ⇒ flow_from_n /
1208
+ # flow_from_nodeBalance entry on the in-region node side.
1209
+ # No flow_to_n entry (virtual_node is not in nodeBalance).
1210
+ # * IMPORT: half-flow goes from virtual_node to in_region_node.
1211
+ # Sink-side flow enters in_region_node ⇒ flow_to_n entry on
1212
+ # the in-region node side. No flow_from_n on the virtual side.
1213
+ if hf.side == "export":
1214
+ new_flow_from_n_rows.append({
1215
+ "p": hf.virtual_p,
1216
+ "source": hf.virtual_arc_source,
1217
+ "sink": hf.virtual_arc_sink,
1218
+ "n": hf.in_region_node,
1219
+ })
1220
+ # process_source_sink_eff/noEff is the partition used by
1221
+ # flow_from_nodeBalance_eff/noEff. Match the original.
1222
+ if is_eff:
1223
+ new_flow_from_nb_eff_rows.append({
1224
+ "p": hf.virtual_p,
1225
+ "source": hf.virtual_arc_source,
1226
+ "sink": hf.virtual_arc_sink,
1227
+ "n": hf.in_region_node,
1228
+ })
1229
+ else:
1230
+ new_flow_from_nb_noEff_rows.append({
1231
+ "p": hf.virtual_p,
1232
+ "source": hf.virtual_arc_source,
1233
+ "sink": hf.virtual_arc_sink,
1234
+ "n": hf.in_region_node,
1235
+ })
1236
+ else: # import
1237
+ new_flow_to_n_rows.append({
1238
+ "p": hf.virtual_p,
1239
+ "source": hf.virtual_arc_source,
1240
+ "sink": hf.virtual_arc_sink,
1241
+ "n": hf.in_region_node,
1242
+ })
1243
+
1244
+ # unitsize = original (so v_flow numbers match between paired
1245
+ # half-flows and original physical flow capacities are
1246
+ # preserved).
1247
+ if orig_unitsize is not None:
1248
+ us_match = orig_unitsize.frame.filter(
1249
+ pl.col("p") == hf.original_p)
1250
+ us_val = (float(us_match["value"][0]) if us_match.height > 0
1251
+ else 1.0)
1252
+ else:
1253
+ us_val = 1.0
1254
+ new_unitsize_rows.append({"p": hf.virtual_p, "value": us_val})
1255
+
1256
+ # flow_upper inherits from the original arc (per (d, t)).
1257
+ if orig_flow_upper is not None:
1258
+ cap_rows = orig_flow_upper.frame.filter(
1259
+ (pl.col("p") == hf.original_p)
1260
+ & (pl.col("source") == hf.original_source)
1261
+ & (pl.col("sink") == hf.original_sink)
1262
+ ).select("d", "t", "value")
1263
+ for r in cap_rows.iter_rows(named=True):
1264
+ new_flow_upper_rows.append({
1265
+ "p": hf.virtual_p,
1266
+ "source": hf.virtual_arc_source,
1267
+ "sink": hf.virtual_arc_sink,
1268
+ "d": r["d"], "t": r["t"],
1269
+ "value": float(r["value"]),
1270
+ })
1271
+ if orig_flow_upper_existing is not None:
1272
+ cap_rows = orig_flow_upper_existing.frame.filter(
1273
+ (pl.col("p") == hf.original_p)
1274
+ & (pl.col("source") == hf.original_source)
1275
+ & (pl.col("sink") == hf.original_sink)
1276
+ ).select("d", "value")
1277
+ for r in cap_rows.iter_rows(named=True):
1278
+ # Benders mode: the master owns the real ``f ≤ C·unitsize``
1279
+ # capacity limit, so the per-region half-flow's ``maxFlow``
1280
+ # bound must be effectively unbounded — otherwise a
1281
+ # greenfield pipe (whose inherited ``existing`` is 0) is
1282
+ # pinned to zero trade (the false-convergence bug). Swap
1283
+ # the inherited value for a large sentinel that can never
1284
+ # bind. Default (un-set) keeps today's inherit.
1285
+ value = (_BENDERS_UNCAP_SENTINEL
1286
+ if benders_uncap_cross_region
1287
+ else float(r["value"]))
1288
+ new_flow_upper_existing_rows.append({
1289
+ "p": hf.virtual_p,
1290
+ "source": hf.virtual_arc_source,
1291
+ "sink": hf.virtual_arc_sink,
1292
+ "d": r["d"],
1293
+ "value": value,
1294
+ })
1295
+
1296
+ # ── arc-block weights (lh2 fixture only) ──
1297
+ # For each half-flow, mirror the original arc's weights so the
1298
+ # block-aware nodeBalance aggregation includes the half-flow at
1299
+ # the right (d, t) granularity.
1300
+ if hf.side == "export":
1301
+ # Source-side: original arc_source_block_dt row(s).
1302
+ if src.arc_source_block_dt is not None:
1303
+ src_rows = src.arc_source_block_dt.filter(
1304
+ (pl.col("p") == hf.original_p)
1305
+ & (pl.col("source") == hf.original_source)
1306
+ & (pl.col("sink") == hf.original_sink)
1307
+ ).select("d", "b_first", "t", "weight")
1308
+ for r in src_rows.iter_rows(named=True):
1309
+ # We'll augment arc_source_block_dt later; collect.
1310
+ new_arc_sink_block_dt_rows.append({
1311
+ "p": hf.virtual_p,
1312
+ "source": hf.virtual_arc_source,
1313
+ "sink": hf.virtual_arc_sink,
1314
+ "d": r["d"], "b_first": r["b_first"],
1315
+ "t": r["t"], "weight": r["weight"],
1316
+ "_side": "source",
1317
+ })
1318
+ if src.p_arc_source_weight is not None:
1319
+ w_rows = src.p_arc_source_weight.frame.filter(
1320
+ (pl.col("p") == hf.original_p)
1321
+ & (pl.col("source") == hf.original_source)
1322
+ & (pl.col("sink") == hf.original_sink)
1323
+ ).select("d", "t", "value")
1324
+ for r in w_rows.iter_rows(named=True):
1325
+ new_p_arc_sink_weight_rows.append({
1326
+ "p": hf.virtual_p,
1327
+ "source": hf.virtual_arc_source,
1328
+ "sink": hf.virtual_arc_sink,
1329
+ "d": r["d"], "t": r["t"],
1330
+ "value": float(r["value"]),
1331
+ "_side": "source",
1332
+ })
1333
+ else:
1334
+ if src.arc_sink_block_dt is not None:
1335
+ snk_rows = src.arc_sink_block_dt.filter(
1336
+ (pl.col("p") == hf.original_p)
1337
+ & (pl.col("source") == hf.original_source)
1338
+ & (pl.col("sink") == hf.original_sink)
1339
+ ).select("d", "b_first", "t", "weight")
1340
+ for r in snk_rows.iter_rows(named=True):
1341
+ new_arc_sink_block_dt_rows.append({
1342
+ "p": hf.virtual_p,
1343
+ "source": hf.virtual_arc_source,
1344
+ "sink": hf.virtual_arc_sink,
1345
+ "d": r["d"], "b_first": r["b_first"],
1346
+ "t": r["t"], "weight": r["weight"],
1347
+ "_side": "sink",
1348
+ })
1349
+ if src.p_arc_sink_weight is not None:
1350
+ w_rows = src.p_arc_sink_weight.frame.filter(
1351
+ (pl.col("p") == hf.original_p)
1352
+ & (pl.col("source") == hf.original_source)
1353
+ & (pl.col("sink") == hf.original_sink)
1354
+ ).select("d", "t", "value")
1355
+ for r in w_rows.iter_rows(named=True):
1356
+ new_p_arc_sink_weight_rows.append({
1357
+ "p": hf.virtual_p,
1358
+ "source": hf.virtual_arc_source,
1359
+ "sink": hf.virtual_arc_sink,
1360
+ "d": r["d"], "t": r["t"],
1361
+ "value": float(r["value"]),
1362
+ "_side": "sink",
1363
+ })
1364
+
1365
+ # The widened Enum vocabulary set by ``split()`` lives on the global
1366
+ # ContextVar. ``src._axis_enums`` is the narrower snapshot captured
1367
+ # before widening — using it for new-row schemas would null the
1368
+ # virtual ``hf_pipe_*`` / virtual-node tokens. Read from the global
1369
+ # first; the live widened vocabulary is a strict superset of the
1370
+ # source snapshot, so upcasting existing-frame columns to it is safe.
1371
+ _enums_loc = get_global_axis_enums() or getattr(src, "_axis_enums", None)
1372
+
1373
+ def _upcast_dims(frame: pl.DataFrame | None,
1374
+ cols: Iterable[str]) -> pl.DataFrame | None:
1375
+ """Re-cast the named dim columns on ``frame`` to the wider Enum
1376
+ vocabulary in ``_enums_loc``. Required before concat against
1377
+ new rows built with ``schema_dtype(_enums_loc, axis)`` — both
1378
+ sides must agree on the Enum vocabulary or polars raises."""
1379
+ if frame is None:
1380
+ return None
1381
+ exprs = [cast_dim(pl.col(c), _enums_loc, c) for c in cols]
1382
+ return frame.with_columns(exprs)
1383
+
1384
+ # ---- Concatenate into rd ----
1385
+ def _concat(orig: pl.DataFrame | None,
1386
+ rows: list[dict],
1387
+ schema: dict) -> pl.DataFrame | None:
1388
+ if not rows:
1389
+ return orig
1390
+ new_df = pl.DataFrame(rows, schema=schema)
1391
+ if orig is None or orig.height == 0:
1392
+ # Need a frame matching the original schema; fall back to new.
1393
+ return new_df.select(list(schema.keys()))
1394
+ # Upcast orig's dim columns to the wider Enum before concat so
1395
+ # both sides agree on the vocabulary. Value-typed columns
1396
+ # (Float64) pass through ``cast_dim`` unchanged.
1397
+ orig_w = _upcast_dims(orig, list(schema.keys()))
1398
+ return pl.concat([orig_w.select(list(schema.keys())),
1399
+ new_df.select(list(schema.keys()))],
1400
+ how="vertical")
1401
+
1402
+ _pss_schema = {"p": schema_dtype(_enums_loc, "p"),
1403
+ "source": schema_dtype(_enums_loc, "source"),
1404
+ "sink": schema_dtype(_enums_loc, "sink")}
1405
+ _pssn_schema = {**_pss_schema, "n": schema_dtype(_enums_loc, "n")}
1406
+ rd.process_source_sink = _concat(
1407
+ rd.process_source_sink, new_pss_rows, _pss_schema)
1408
+ if new_pss_eff_rows:
1409
+ rd.process_source_sink_eff = _concat(
1410
+ rd.process_source_sink_eff, new_pss_eff_rows, _pss_schema)
1411
+ if new_pss_noEff_rows:
1412
+ rd.process_source_sink_noEff = _concat(
1413
+ rd.process_source_sink_noEff, new_pss_noEff_rows, _pss_schema)
1414
+ # Phase E.3: ``rd.pss_dt`` is no longer persisted; the half-flow
1415
+ # virtual (p, source, sink) rows already appended to
1416
+ # ``rd.process_source_sink`` will produce the matching cross-join
1417
+ # rows when ``compute_pss_dt(rd)`` runs downstream. We build a
1418
+ # local ``virtual_pss_dt`` (just the half-flow rows) for the
1419
+ # availability / existing_count promotion below.
1420
+ _virtual_pss_dt_schema = {
1421
+ **_pss_schema,
1422
+ "d": schema_dtype(_enums_loc, "d"),
1423
+ "t": schema_dtype(_enums_loc, "t"),
1424
+ }
1425
+ if new_pss_dt_rows:
1426
+ virtual_pss_dt = pl.DataFrame(
1427
+ new_pss_dt_rows, schema=_virtual_pss_dt_schema)
1428
+ else:
1429
+ virtual_pss_dt = pl.DataFrame(schema=_virtual_pss_dt_schema)
1430
+ if new_flow_to_n_rows:
1431
+ rd.flow_to_n = _concat(
1432
+ rd.flow_to_n, new_flow_to_n_rows, _pssn_schema)
1433
+ if new_flow_from_n_rows:
1434
+ rd.flow_from_n = _concat(
1435
+ rd.flow_from_n, new_flow_from_n_rows, _pssn_schema)
1436
+ if new_flow_from_nb_eff_rows:
1437
+ rd.flow_from_nodeBalance_eff = _concat(
1438
+ rd.flow_from_nodeBalance_eff, new_flow_from_nb_eff_rows, _pssn_schema)
1439
+ if new_flow_from_nb_noEff_rows:
1440
+ rd.flow_from_nodeBalance_noEff = _concat(
1441
+ rd.flow_from_nodeBalance_noEff, new_flow_from_nb_noEff_rows, _pssn_schema)
1442
+
1443
+ # Append unitsize Param.
1444
+ if rd.p_unitsize is not None and new_unitsize_rows:
1445
+ new_us = pl.DataFrame(new_unitsize_rows,
1446
+ schema={"p": schema_dtype(_enums_loc, "p"),
1447
+ "value": pl.Float64})
1448
+ merged_us = pl.concat([_upcast_dims(rd.p_unitsize.frame, ("p",))
1449
+ .select("p", "value"),
1450
+ new_us], how="vertical")
1451
+ rd.p_unitsize = Param(("p",), merged_us, name=rd.p_unitsize.name)
1452
+
1453
+ # Append p_slope rows for half-flows that are eff-classified.
1454
+ if rd.p_slope is not None and new_p_slope_rows:
1455
+ new_sl = pl.DataFrame(new_p_slope_rows,
1456
+ schema={"p": schema_dtype(_enums_loc, "p"),
1457
+ "d": schema_dtype(_enums_loc, "d"),
1458
+ "t": schema_dtype(_enums_loc, "t"),
1459
+ "value": pl.Float64})
1460
+ merged_sl = pl.concat([_upcast_dims(rd.p_slope.frame, ("p", "d", "t"))
1461
+ .select("p", "d", "t", "value"),
1462
+ new_sl], how="vertical")
1463
+ rd.p_slope = Param(("p", "d", "t"), merged_sl, name=rd.p_slope.name)
1464
+
1465
+ # Append flow_upper Param rows.
1466
+ if rd.p_flow_upper is not None and new_flow_upper_rows:
1467
+ new_fu = pl.DataFrame(new_flow_upper_rows,
1468
+ schema={"p": schema_dtype(_enums_loc, "p"),
1469
+ "source": schema_dtype(_enums_loc, "source"),
1470
+ "sink": schema_dtype(_enums_loc, "sink"),
1471
+ "d": schema_dtype(_enums_loc, "d"),
1472
+ "t": schema_dtype(_enums_loc, "t"),
1473
+ "value": pl.Float64})
1474
+ merged_fu = pl.concat([_upcast_dims(rd.p_flow_upper.frame,
1475
+ ("p", "source", "sink", "d", "t"))
1476
+ .select("p", "source", "sink", "d", "t", "value"),
1477
+ new_fu], how="vertical")
1478
+ rd.p_flow_upper = Param(("p", "source", "sink", "d", "t"),
1479
+ merged_fu, name=rd.p_flow_upper.name)
1480
+ if rd.p_flow_upper_existing is not None and new_flow_upper_existing_rows:
1481
+ new_fue = pl.DataFrame(new_flow_upper_existing_rows,
1482
+ schema={"p": schema_dtype(_enums_loc, "p"),
1483
+ "source": schema_dtype(_enums_loc, "source"),
1484
+ "sink": schema_dtype(_enums_loc, "sink"),
1485
+ "d": schema_dtype(_enums_loc, "d"),
1486
+ "value": pl.Float64})
1487
+ merged_fue = pl.concat([_upcast_dims(rd.p_flow_upper_existing.frame,
1488
+ ("p", "source", "sink", "d"))
1489
+ .select("p", "source", "sink", "d", "value"),
1490
+ new_fue], how="vertical")
1491
+ rd.p_flow_upper_existing = Param(
1492
+ ("p", "source", "sink", "d"),
1493
+ merged_fue, name=rd.p_flow_upper_existing.name)
1494
+
1495
+ # ── p_process_availability and p_process_existing_count ──
1496
+ # The maxFlow RHS is multiplied by p_process_availability when
1497
+ # populated, and Param×Param is an inner-join so missing half-flow
1498
+ # entries collapse to zero RHS. We must add availability=1.0 and
1499
+ # existing_count=1.0 entries so the half-flow's bound stays at the
1500
+ # value we set in p_flow_upper_existing.
1501
+ if rd.p_process_availability is not None and virtual_pss_dt.height > 0:
1502
+ # Add a (p, d, t) row for each (virtual_p, d, t) in
1503
+ # virtual_pss_dt (Phase E.3: half-flow rows only, no need to
1504
+ # filter the whole-region cross-join).
1505
+ avail_rows = (virtual_pss_dt
1506
+ .select("p", "d", "t")
1507
+ .with_columns(value=pl.lit(1.0)))
1508
+ if avail_rows.height > 0:
1509
+ # Phase E.1: p_process_availability dims depend on authored
1510
+ # shape — promote to (p, d, t) via virtual_pss_dt's d/t
1511
+ # axes so the concat lands at a uniform schema.
1512
+ avail_pdt = promote_param_to_dt(
1513
+ rd.p_process_availability, virtual_pss_dt).collect()
1514
+ merged = pl.concat([_upcast_dims(avail_pdt, ("p", "d", "t"))
1515
+ .select("p", "d", "t", "value"),
1516
+ avail_rows], how="vertical")
1517
+ rd.p_process_availability = Param(
1518
+ ("p", "d", "t"), merged,
1519
+ name=rd.p_process_availability.name)
1520
+ if rd.p_process_existing_count is not None and virtual_pss_dt.height > 0:
1521
+ # (p, d) row for each virtual half-flow
1522
+ ec_rows = (virtual_pss_dt
1523
+ .filter(cast_dim(pl.col("p"), None, "p").is_in([hf.virtual_p for hf in half_flows]))
1524
+ .select("p", "d").unique()
1525
+ .with_columns(value=pl.lit(1.0)))
1526
+ if ec_rows.height > 0:
1527
+ merged = pl.concat([_upcast_dims(rd.p_process_existing_count.frame,
1528
+ ("p", "d"))
1529
+ .select("p", "d", "value"),
1530
+ ec_rows], how="vertical")
1531
+ rd.p_process_existing_count = Param(
1532
+ ("p", "d"), merged,
1533
+ name=rd.p_process_existing_count.name)
1534
+
1535
+ # Append arc-block-weight rows (lh2 fixture). Half-flows on the
1536
+ # source side go to arc_source_block_dt + p_arc_source_weight; on
1537
+ # sink side they go to arc_sink_block_dt + p_arc_sink_weight.
1538
+ src_block_rows = [r for r in new_arc_sink_block_dt_rows if r["_side"] == "source"]
1539
+ snk_block_rows = [r for r in new_arc_sink_block_dt_rows if r["_side"] == "sink"]
1540
+ src_w_rows = [r for r in new_p_arc_sink_weight_rows if r["_side"] == "source"]
1541
+ snk_w_rows = [r for r in new_p_arc_sink_weight_rows if r["_side"] == "sink"]
1542
+
1543
+ if rd.arc_source_block_dt is not None and src_block_rows:
1544
+ new_df = pl.DataFrame(
1545
+ [{k: r[k] for k in ("p", "source", "sink", "d", "b_first", "t", "weight")}
1546
+ for r in src_block_rows],
1547
+ schema={"p": schema_dtype(_enums_loc, "p"),
1548
+ "source": schema_dtype(_enums_loc, "source"),
1549
+ "sink": schema_dtype(_enums_loc, "sink"),
1550
+ "d": schema_dtype(_enums_loc, "d"),
1551
+ "b_first": schema_dtype(_enums_loc, "b_first"),
1552
+ "t": schema_dtype(_enums_loc, "t"),
1553
+ "weight": pl.Float64})
1554
+ rd.arc_source_block_dt = pl.concat([
1555
+ _upcast_dims(rd.arc_source_block_dt,
1556
+ ("p", "source", "sink", "d", "b_first", "t"))
1557
+ .select(*new_df.columns), new_df],
1558
+ how="vertical")
1559
+ if rd.arc_sink_block_dt is not None and snk_block_rows:
1560
+ new_df = pl.DataFrame(
1561
+ [{k: r[k] for k in ("p", "source", "sink", "d", "b_first", "t", "weight")}
1562
+ for r in snk_block_rows],
1563
+ schema={"p": schema_dtype(_enums_loc, "p"),
1564
+ "source": schema_dtype(_enums_loc, "source"),
1565
+ "sink": schema_dtype(_enums_loc, "sink"),
1566
+ "d": schema_dtype(_enums_loc, "d"),
1567
+ "b_first": schema_dtype(_enums_loc, "b_first"),
1568
+ "t": schema_dtype(_enums_loc, "t"),
1569
+ "weight": pl.Float64})
1570
+ rd.arc_sink_block_dt = pl.concat([
1571
+ _upcast_dims(rd.arc_sink_block_dt,
1572
+ ("p", "source", "sink", "d", "b_first", "t"))
1573
+ .select(*new_df.columns), new_df],
1574
+ how="vertical")
1575
+ if rd.p_arc_source_weight is not None and src_w_rows:
1576
+ new_df = pl.DataFrame(
1577
+ [{k: r[k] for k in ("p", "source", "sink", "d", "t", "value")}
1578
+ for r in src_w_rows],
1579
+ schema={"p": schema_dtype(_enums_loc, "p"),
1580
+ "source": schema_dtype(_enums_loc, "source"),
1581
+ "sink": schema_dtype(_enums_loc, "sink"),
1582
+ "d": schema_dtype(_enums_loc, "d"),
1583
+ "t": schema_dtype(_enums_loc, "t"),
1584
+ "value": pl.Float64})
1585
+ rd.p_arc_source_weight = Param(
1586
+ ("p", "source", "sink", "d", "t"),
1587
+ pl.concat([_upcast_dims(rd.p_arc_source_weight.frame,
1588
+ ("p", "source", "sink", "d", "t"))
1589
+ .select(*new_df.columns),
1590
+ new_df], how="vertical"),
1591
+ name=rd.p_arc_source_weight.name)
1592
+ if rd.p_arc_sink_weight is not None and snk_w_rows:
1593
+ new_df = pl.DataFrame(
1594
+ [{k: r[k] for k in ("p", "source", "sink", "d", "t", "value")}
1595
+ for r in snk_w_rows],
1596
+ schema={"p": schema_dtype(_enums_loc, "p"),
1597
+ "source": schema_dtype(_enums_loc, "source"),
1598
+ "sink": schema_dtype(_enums_loc, "sink"),
1599
+ "d": schema_dtype(_enums_loc, "d"),
1600
+ "t": schema_dtype(_enums_loc, "t"),
1601
+ "value": pl.Float64})
1602
+ rd.p_arc_sink_weight = Param(
1603
+ ("p", "source", "sink", "d", "t"),
1604
+ pl.concat([_upcast_dims(rd.p_arc_sink_weight.frame,
1605
+ ("p", "source", "sink", "d", "t"))
1606
+ .select(*new_df.columns),
1607
+ new_df], how="vertical"),
1608
+ name=rd.p_arc_sink_weight.name)
1609
+
1610
+ # group_entity / group_node — augment with virtual entities under the
1611
+ # half-flow's region (so downstream group-aware emitters don't
1612
+ # spuriously skip them). Skip — the group_* sets only matter for
1613
+ # group_slack / capacity_margin features which aren't in lh2 fixture.
1614
+
1615
+ return rd
1616
+
1617
+
1618
+ # ---------------------------------------------------------------------------
1619
+ # Benders network-only master producer (the INVERSE of ``split``)
1620
+ # ---------------------------------------------------------------------------
1621
+
1622
+
1623
+ #: FlexData fields keyed on a NODE axis ``n`` (node entity sets, storage,
1624
+ #: inflow, penalties, profiles, …). The master OMITS every terminal node
1625
+ #: from balance, so each of these is emptied (set to ``None``) — a node
1626
+ #: absent from ``nodeBalance`` / ``nodeStateBlock`` gets no balance row
1627
+ #: (Phase-3 design §1.3; ``model.py`` builds balance only over the
1628
+ #: populated node sets).
1629
+ _MASTER_NODE_FIELDS: tuple[str, ...] = (
1630
+ "nodeBalance_dt", "nodeBalancePeriod",
1631
+ "nodeState", "nodeState_dt", "nodeState_first_dt", "nodeState_last_dt",
1632
+ "nodeState_rp", "nodeStateBlock",
1633
+ "storage_bind_within_timeblock", "storage_bind_forward_only",
1634
+ "storage_bind_within_solve", "storage_bind_within_solve_blended_weights",
1635
+ "storage_bind_within_period_blended_weights",
1636
+ "storage_bind_forward_only_blended_weights", "storage_fix_start",
1637
+ "p_state_upper", "p_state_unitsize", "p_state_self_discharge",
1638
+ "p_state_start", "p_state_existing_capacity",
1639
+ "storage_use_reference_value", "p_storage_state_reference_value",
1640
+ "p_storage_state_reference_price",
1641
+ "node_profile_upper", "node_profile_lower", "node_profile_fixed",
1642
+ "p_node_availability", "p_roll_continue_state",
1643
+ "n_fix_storage_quantity", "ndt_fix_storage_quantity",
1644
+ "p_fix_storage_quantity", "n_fix_storage_usage",
1645
+ "ndt_fix_storage_usage", "p_fix_storage_usage",
1646
+ "p_node_capacity_for_scaling",
1647
+ # node-keyed invest/divest sets (the trade nodes are not the master's
1648
+ # invest variables — only the cross connections are).
1649
+ "nd_invest_set", "nd_divest_set",
1650
+ # CO2 / reserve / user-constraint / group features — all in-region
1651
+ # recourse, not part of the network-only master. (The commodity
1652
+ # frames + ``p_commodity_price`` are REQUIRED-present by the PROCESSES
1653
+ # feature, so they are EMPTIED rather than nulled below.)
1654
+ "flow_from_co2_priced", "flow_from_co2_priced_noEff",
1655
+ "p_co2_content", "p_co2_price",
1656
+ "group_co2_max_period", "flow_from_co2_capped",
1657
+ "flow_from_co2_capped_noEff", "p_co2_max_period", "group_d_co2_capped",
1658
+ "group_co2_max_total", "flow_from_co2_capped_total",
1659
+ "flow_from_co2_capped_total_noEff", "p_co2_max_total",
1660
+ "flow_constraint_idx", "p_flow_constraint_coef", "p_constraint_constant",
1661
+ "cdt_eq", "cdt_le", "cdt_ge",
1662
+ "p_node_constraint_invested_capacity_coeff",
1663
+ "p_process_constraint_invested_capacity_coeff",
1664
+ "p_node_constraint_state_coeff",
1665
+ "p_node_constraint_prebuilt_capacity_coeff",
1666
+ "p_process_constraint_prebuilt_capacity_coeff",
1667
+ "groupCapacityMargin", "groupInertia", "groupNonSync", "group_node",
1668
+ "process_sink_inertia", "process_source_inertia",
1669
+ "process_sink_nonSync", "process_group_inside_nonSync",
1670
+ "p_inv_group_cap", "p_group_capacity_for_scaling",
1671
+ "pdGroup_capacity_margin",
1672
+ )
1673
+
1674
+
1675
+ # ---------------------------------------------------------------------------
1676
+ # Master-hosted-mode partition of ``_MASTER_NODE_FIELDS``. With a
1677
+ # non-empty ``master_hosted_nodes`` set the master KEEPS master-side
1678
+ # content instead of nulling everything: each tuple below routes its
1679
+ # fields through the matching keep-filter in ``master_network_data``.
1680
+ # Every ``_MASTER_NODE_FIELDS`` member appears in exactly one of these
1681
+ # groups (pinned by a test) so a schema addition to the legacy tuple
1682
+ # cannot silently skip the master-hosted path.
1683
+ # ---------------------------------------------------------------------------
1684
+
1685
+ #: ``n``-keyed frames/Params → filter rows to the master-hosted nodes
1686
+ #: (balance/state/storage/profile/availability/fix/invest-set frames;
1687
+ #: ``group_node`` bare membership rows are filtered silently the same
1688
+ #: way, mirroring ``split``).
1689
+ _MASTER_N_KEYED_FIELDS: tuple[str, ...] = (
1690
+ "nodeBalance_dt", "nodeBalancePeriod",
1691
+ "nodeState", "nodeState_dt", "nodeState_first_dt", "nodeState_last_dt",
1692
+ "nodeState_rp", "nodeStateBlock",
1693
+ "storage_bind_within_timeblock", "storage_bind_forward_only",
1694
+ "storage_bind_within_solve", "storage_bind_within_solve_blended_weights",
1695
+ "storage_bind_within_period_blended_weights",
1696
+ "storage_bind_forward_only_blended_weights", "storage_fix_start",
1697
+ "p_state_upper", "p_state_unitsize", "p_state_self_discharge",
1698
+ "p_state_start", "p_state_existing_capacity",
1699
+ "storage_use_reference_value", "p_storage_state_reference_value",
1700
+ "p_storage_state_reference_price",
1701
+ "node_profile_upper", "node_profile_lower", "node_profile_fixed",
1702
+ "p_node_availability", "p_roll_continue_state",
1703
+ "n_fix_storage_quantity", "ndt_fix_storage_quantity",
1704
+ "p_fix_storage_quantity", "n_fix_storage_usage",
1705
+ "ndt_fix_storage_usage", "p_fix_storage_usage",
1706
+ "p_node_capacity_for_scaling",
1707
+ "nd_invest_set", "nd_divest_set",
1708
+ "group_node",
1709
+ )
1710
+
1711
+ #: ``cn``-keyed user-constraint frames/Params → keep rows of the
1712
+ #: all-master constraint ids (:func:`_master_side_constraint_ids`).
1713
+ _MASTER_CN_KEYED_FIELDS: tuple[str, ...] = (
1714
+ "flow_constraint_idx", "p_flow_constraint_coef",
1715
+ "p_constraint_constant", "cdt_eq", "cdt_le", "cdt_ge",
1716
+ "p_node_constraint_invested_capacity_coeff",
1717
+ "p_process_constraint_invested_capacity_coeff",
1718
+ "p_node_constraint_state_coeff",
1719
+ "p_node_constraint_prebuilt_capacity_coeff",
1720
+ "p_process_constraint_prebuilt_capacity_coeff",
1721
+ )
1722
+
1723
+ #: Arc-keyed CO2-priced flow frames → keep master-local arc rows
1724
+ #: (coupling-arc rows hard-error in :func:`_co2_master_partition`).
1725
+ _MASTER_CO2_ARC_FIELDS: tuple[str, ...] = (
1726
+ "flow_from_co2_priced", "flow_from_co2_priced_noEff",
1727
+ )
1728
+
1729
+ #: ``g``-keyed CO2-cap frames → keep rows of the all-master capped
1730
+ #: groups (:func:`_co2_master_partition`).
1731
+ _MASTER_CO2_GROUP_FIELDS: tuple[str, ...] = (
1732
+ "group_co2_max_period", "flow_from_co2_capped",
1733
+ "flow_from_co2_capped_noEff", "p_co2_max_period",
1734
+ "group_d_co2_capped",
1735
+ "group_co2_max_total", "flow_from_co2_capped_total",
1736
+ "flow_from_co2_capped_total_noEff", "p_co2_max_total",
1737
+ )
1738
+
1739
+ #: Pure lookup tables consumed by the kept CO2 rows via joins — carried
1740
+ #: whole (rows for un-kept groups/commodities are inert).
1741
+ _MASTER_CO2_LOOKUP_FIELDS: tuple[str, ...] = (
1742
+ "p_co2_content", "p_co2_price",
1743
+ )
1744
+
1745
+ #: ``p``-keyed group-feature side frames → filter to the master's
1746
+ #: kept processes (rows for absent flows are inert Where-filters).
1747
+ _MASTER_GROUP_PROC_FIELDS: tuple[str, ...] = (
1748
+ "process_sink_inertia", "process_source_inertia",
1749
+ "process_sink_nonSync", "process_group_inside_nonSync",
1750
+ )
1751
+
1752
+
1753
+ def master_network_data(
1754
+ data: FlexData,
1755
+ regions: list[str],
1756
+ *,
1757
+ region_membership: dict[str, dict[str, set[str]]] | None = None,
1758
+ master_hosted_nodes: frozenset[str] = frozenset(),
1759
+ ) -> FlexData:
1760
+ """Build the Benders MASTER's reduced :class:`FlexData` — the INVERSE
1761
+ of :func:`split`.
1762
+
1763
+ Returns a reduced :class:`FlexData` containing ONLY the cross-region
1764
+ ``(p, source, sink)`` arcs plus their invest / cost / timeline params,
1765
+ with EVERY terminal node OMITTED from node balance (and the
1766
+ block / state / inflow frames). ``build_flextool`` over the result
1767
+ generates the master skeleton natively:
1768
+
1769
+ * ``v_flow[conn, source, sink, d, t]`` for every cross arc (built from
1770
+ ``process_source_sink × dt``, independent of ``nodeBalance``);
1771
+ * ``v_invest_p[conn, d]`` over ``pd_invest_set``;
1772
+ * the capacity-tied ``maxFlow`` row (greenfield ⇒ ``flow_upper_rhs=0``
1773
+ with ``-v_invest_p`` on the LHS ⇒ ``v_flow ≤ v_invest_p``);
1774
+ * the invest annuity cost + (when authored) the connection flow cost.
1775
+
1776
+ Reuses the cross-arc CLASSIFICATION (:func:`_classify_arcs`) — the same
1777
+ detection :func:`split` uses — rather than re-deriving it.
1778
+
1779
+ Parameters
1780
+ ----------
1781
+ data
1782
+ Whole-system :class:`FlexData`.
1783
+ regions
1784
+ The region group names the splitter partitions on.
1785
+ region_membership
1786
+ Pre-computed membership (see :func:`load_region_membership`); when
1787
+ omitted, re-derived from *data*.
1788
+ master_hosted_nodes
1789
+ Master-hosted node mode (see :func:`compute_master_hosted_nodes`
1790
+ / :func:`split`). With the default EMPTY set the function takes
1791
+ today's exact path (explicit early branches — byte-identical).
1792
+ Non-empty: the master KEEPS the master-side model instead of
1793
+ emptying it — see Notes.
1794
+
1795
+ Notes
1796
+ -----
1797
+ With the default empty ``master_hosted_nodes``, the reduced FlexData
1798
+ KEEPS only cross-arc rows in the process / arc / arc-cost /
1799
+ arc-block frames and the cross connections in the invest / unitsize
1800
+ / max-units params; EVERY node-keyed frame (balance, state, inflow,
1801
+ penalties, storage, profiles, CO2, groups, user constraints) is
1802
+ emptied (``None``).
1803
+
1804
+ With a NON-EMPTY ``master_hosted_nodes``:
1805
+
1806
+ * the arc keep-set widens to cross-region ∪ region↔master coupling
1807
+ ∪ master-local arcs (and the proc/cost keep-sets follow), so
1808
+ master-local connections AND units build natively in the master;
1809
+ * ``n``-keyed frames are FILTERED to the master-hosted nodes instead
1810
+ of nulled (balance, inflow, penalties, storage/state, profiles,
1811
+ availability, fix-storage, node invest sets);
1812
+ * the entity-keyed invest frames keep
1813
+ ``procs-of-kept-arcs ∪ master_hosted_nodes`` (master storage
1814
+ invest needs the ``e``-keyed annuity/lifetime/max-units rows);
1815
+ * only the REGION-side endpoints of coupling arcs (and the
1816
+ region↔region terminals) are omitted from balance — master-hosted
1817
+ endpoints stay balanced;
1818
+ * user constraints / CO2 caps / group features whose referenced
1819
+ entities are ALL master-side are kept; straddling ones hard-error
1820
+ (same validation family as :func:`split`).
1821
+
1822
+ In both modes the solve-data-keyed timeline frames (``dt``,
1823
+ ``p_step_duration``, ``p_timestep_weight``, ``p_inflation_op``,
1824
+ ``p_period_share``, the RP / block frames) carry through the
1825
+ ``dataclasses.replace`` shallow copy unchanged, so the master's
1826
+ ``v_flow`` lives on the SAME ``(d, t)`` grid as the region pinned
1827
+ half-flows (Phase-3 §3.5 guard (a)).
1828
+ """
1829
+ if region_membership is None:
1830
+ region_membership = load_region_membership(data, regions)
1831
+ region_nodes = {r: m["nodes"] for r, m in region_membership.items()}
1832
+
1833
+ new = dataclasses.replace(data)
1834
+
1835
+ if data.process_source_sink is None:
1836
+ raise RuntimeError(
1837
+ "master_network_data: no process_source_sink — nothing to "
1838
+ "decompose"
1839
+ )
1840
+
1841
+ # Reuse the cross-region classification (same detection as ``split``).
1842
+ # With the default empty master set the extra classification frames
1843
+ # are empty and this path is byte-identical to the historical 2-way
1844
+ # behaviour.
1845
+ master_mode = bool(master_hosted_nodes)
1846
+ _pss_tagged, cross, region_master, master_local = _classify_arcs(
1847
+ data.process_source_sink, region_nodes,
1848
+ master_nodes=master_hosted_nodes,
1849
+ )
1850
+ cross_keys: set[tuple[str, str, str]] = {
1851
+ (r["p"], r["source"], r["sink"]) for r in cross.iter_rows(named=True)
1852
+ }
1853
+ master_cns: set[str] = set()
1854
+ master_groups_feature: set[str] = set()
1855
+ master_groups_co2: set[str] = set()
1856
+ master_local_keys: set[tuple[str, str, str]] = set()
1857
+ if not master_mode:
1858
+ # Today's exact path: region↔region cross arcs are the only
1859
+ # master content.
1860
+ if cross.height == 0:
1861
+ raise RuntimeError(
1862
+ "master_network_data: no cross-region arcs found"
1863
+ )
1864
+ else:
1865
+ all_region_nodes: set[str] = set()
1866
+ for ns in region_nodes.values():
1867
+ all_region_nodes |= ns
1868
+ overlap = set(master_hosted_nodes) & all_region_nodes
1869
+ if overlap:
1870
+ raise RuntimeError(
1871
+ f"master_network_data: master_hosted_nodes overlap "
1872
+ f"region membership: {sorted(overlap)} — a node is "
1873
+ f"either master-hosted (in no region group) or in "
1874
+ f"exactly one region, never both."
1875
+ )
1876
+ region_master_keys = {
1877
+ (r["p"], r["source"], r["sink"])
1878
+ for r in region_master.iter_rows(named=True)
1879
+ }
1880
+ master_local_keys = {
1881
+ (r["p"], r["source"], r["sink"])
1882
+ for r in master_local.iter_rows(named=True)
1883
+ }
1884
+ master_balance_nodes = set()
1885
+ for fld in ("nodeBalance", "nodeState"):
1886
+ f = getattr(data, fld, None)
1887
+ if f is not None and f.height > 0:
1888
+ master_balance_nodes |= (
1889
+ set(f["n"].cast(pl.Utf8).to_list())
1890
+ & set(master_hosted_nodes))
1891
+ if (not cross_keys and not region_master_keys
1892
+ and not master_local_keys and not master_balance_nodes):
1893
+ raise RuntimeError(
1894
+ "master_network_data: no coupling arcs and no master "
1895
+ "content — nothing to host in the Benders master"
1896
+ )
1897
+ # Same hard-validation family as ``split`` (D-a / R5): never
1898
+ # silently mis-partition authored data. The messages are
1899
+ # identical to ``split``'s, so a driver run that already raised
1900
+ # there cannot raise differently here.
1901
+ _validate_no_straddling_units(
1902
+ data, all_region_nodes, master_hosted_nodes)
1903
+ _ml_procs = _master_local_procs(
1904
+ data.process_source_sink, master_local)
1905
+ _validate_user_constraints(
1906
+ data, all_region_nodes, master_hosted_nodes, _ml_procs)
1907
+ master_cns = _master_side_constraint_ids(
1908
+ data, all_region_nodes, master_hosted_nodes, _ml_procs)
1909
+ master_groups_feature = _master_side_feature_groups(
1910
+ data, all_region_nodes, master_hosted_nodes)
1911
+ master_groups_co2 = _co2_master_partition(
1912
+ data, master_local_keys, region_master_keys)
1913
+ # Arc keep-set = cross ∪ region↔master coupling ∪ master-local
1914
+ # (audit 0.D: the coupling-arc-only case is load-bearing).
1915
+ cross_keys = cross_keys | region_master_keys | master_local_keys
1916
+ cross_procs: set[str] = {k[0] for k in cross_keys}
1917
+ # Entity keep-set for the ``e``-keyed invest frames: master-hosted
1918
+ # STORAGE invest needs the node entities alongside the connections/
1919
+ # units of the kept arcs.
1920
+ entity_keep: set[str] = (
1921
+ cross_procs | set(master_hosted_nodes) if master_mode
1922
+ else cross_procs
1923
+ )
1924
+
1925
+ _enums = getattr(data, "_axis_enums", None) or get_global_axis_enums()
1926
+
1927
+ def _keep_cross_triple(df: pl.DataFrame | None) -> pl.DataFrame | None:
1928
+ """Keep ONLY the rows whose (p, source, sink) is a cross arc."""
1929
+ if df is None:
1930
+ return None
1931
+ if not all(c in df.columns for c in ("p", "source", "sink")):
1932
+ return df
1933
+ key_df = pl.DataFrame(
1934
+ {
1935
+ "p": [t[0] for t in cross_keys],
1936
+ "source": [t[1] for t in cross_keys],
1937
+ "sink": [t[2] for t in cross_keys],
1938
+ },
1939
+ schema={"p": schema_dtype(_enums, "p"),
1940
+ "source": schema_dtype(_enums, "source"),
1941
+ "sink": schema_dtype(_enums, "sink")},
1942
+ )
1943
+ return df.join(key_df, on=("p", "source", "sink"), how="semi")
1944
+
1945
+ def _keep_cross_triple_param(p: Param | None) -> Param | None:
1946
+ if p is None:
1947
+ return None
1948
+ if not all(c in p.dims for c in ("p", "source", "sink")):
1949
+ # p-keyed but not arc-keyed (e.g. p_unitsize): keep cross procs.
1950
+ return _keep_proc_param(p)
1951
+ f = _keep_cross_triple(p.frame)
1952
+ return Param(p.dims, f, name=p.name)
1953
+
1954
+ def _keep_proc(df: pl.DataFrame | None) -> pl.DataFrame | None:
1955
+ if df is None or "p" not in df.columns:
1956
+ return df
1957
+ return df.filter(_is_in_keep("p", cross_procs))
1958
+
1959
+ def _keep_proc_param(p: Param | None) -> Param | None:
1960
+ if p is None:
1961
+ return None
1962
+ if "p" not in p.dims:
1963
+ return p
1964
+ return Param(p.dims, p.frame.filter(_is_in_keep("p", cross_procs)),
1965
+ name=p.name)
1966
+
1967
+ def _keep_entity_param(p: Param | None) -> Param | None:
1968
+ """Keep only the cross connections (+ master-hosted nodes in
1969
+ master mode) on an entity-axis (``e``) param."""
1970
+ if p is None:
1971
+ return None
1972
+ if "e" not in p.dims:
1973
+ return p
1974
+ return Param(p.dims, p.frame.filter(_is_in_keep("e", entity_keep)),
1975
+ name=p.name)
1976
+
1977
+ def _keep_entity_frame(df):
1978
+ """Keep cross connections (+ master-hosted nodes in master mode)
1979
+ on an ``e``-axis frame OR Param (some invest sets ship as plain
1980
+ DataFrames, others as Params)."""
1981
+ if df is None:
1982
+ return None
1983
+ if isinstance(df, Param):
1984
+ return _keep_entity_param(df)
1985
+ if "e" not in df.columns:
1986
+ return df
1987
+ return df.filter(_is_in_keep("e", entity_keep))
1988
+
1989
+ # ---- Process / arc topology: KEEP only cross arcs ----
1990
+ new.process_source_sink = _keep_cross_triple(data.process_source_sink)
1991
+ new.process_source_sink_eff = _keep_cross_triple(data.process_source_sink_eff)
1992
+ new.process_source_sink_noEff = _keep_cross_triple(data.process_source_sink_noEff)
1993
+ new.pss_dt = None # rebuilt on demand from process_source_sink × dt.
1994
+ new.process_source_canonical = _keep_proc(data.process_source_canonical)
1995
+ new.process_sink_canonical = _keep_proc(data.process_sink_canonical)
1996
+ new.flow_to_n = _keep_cross_triple(data.flow_to_n)
1997
+ new.flow_from_n = _keep_cross_triple(data.flow_from_n)
1998
+ new.flow_from_nodeBalance_eff = _keep_cross_triple(data.flow_from_nodeBalance_eff)
1999
+ new.flow_from_nodeBalance_noEff = _keep_cross_triple(data.flow_from_nodeBalance_noEff)
2000
+ new.process_unit = _keep_proc(data.process_unit)
2001
+ new.process_indirect = _keep_proc(data.process_indirect)
2002
+ # Commodity frames + price are REQUIRED-present by the PROCESSES feature.
2003
+ # Cross arcs (trade pipes) are not commodity-fed, so keeping only their
2004
+ # rows EMPTIES these frames while preserving the schema (not None).
2005
+ new.flow_from_commodity_eff = _keep_cross_triple(data.flow_from_commodity_eff)
2006
+ new.flow_from_commodity_noEff = _keep_cross_triple(data.flow_from_commodity_noEff)
2007
+ new.flow_to_commodity = _keep_cross_triple(data.flow_to_commodity)
2008
+ new.p_commodity_price = _keep_proc_param(data.p_commodity_price)
2009
+
2010
+ # ---- Per-arc operating / capacity params ----
2011
+ new.p_unitsize = _keep_proc_param(data.p_unitsize)
2012
+ new.p_all_entity_unitsize = _keep_entity_param(data.p_all_entity_unitsize)
2013
+ new.p_flow_upper = _keep_cross_triple_param(data.p_flow_upper)
2014
+ new.p_flow_upper_existing = _keep_cross_triple_param(data.p_flow_upper_existing)
2015
+ new.p_arc_max_cap_coef = _keep_cross_triple_param(data.p_arc_max_cap_coef)
2016
+ new.p_slope = _keep_proc_param(data.p_slope)
2017
+ new.p_process_existing_count = _keep_proc_param(data.p_process_existing_count)
2018
+ new.p_process_availability = _keep_proc_param(data.p_process_availability)
2019
+ new.pd_neg_cap = _keep_proc_param(data.pd_neg_cap)
2020
+
2021
+ # ---- Profiles (process-keyed) ----
2022
+ new.process_profile_upper = _keep_proc(data.process_profile_upper)
2023
+ new.process_profile_lower = _keep_proc(data.process_profile_lower)
2024
+ new.process_profile_fixed = _keep_proc(data.process_profile_fixed)
2025
+
2026
+ # ---- Per-arc block weights (lh2 fixture; trade nodes are block nodes) ----
2027
+ new.arc_sink_block_dt = _keep_cross_triple(getattr(data, "arc_sink_block_dt", None))
2028
+ new.arc_source_block_dt = _keep_cross_triple(getattr(data, "arc_source_block_dt", None))
2029
+ new.p_arc_sink_weight = _keep_cross_triple_param(getattr(data, "p_arc_sink_weight", None))
2030
+ new.p_arc_source_weight = _keep_cross_triple_param(getattr(data, "p_arc_source_weight", None))
2031
+ new.p_arc_step_duration_sink = _keep_cross_triple_param(
2032
+ getattr(data, "p_arc_step_duration_sink", None))
2033
+ new.p_arc_step_duration_source = _keep_cross_triple_param(
2034
+ getattr(data, "p_arc_step_duration_source", None))
2035
+
2036
+ # ---- Per-arc / per-process variable-cost frames (flow cost) ----
2037
+ new.pssdt_varCost_eff_connection = _keep_cross_triple(
2038
+ getattr(data, "pssdt_varCost_eff_connection", None))
2039
+ new.pssdt_varCost_eff_unit_source = _keep_cross_triple(
2040
+ getattr(data, "pssdt_varCost_eff_unit_source", None))
2041
+ new.pssdt_varCost_eff_unit_sink = _keep_cross_triple(
2042
+ getattr(data, "pssdt_varCost_eff_unit_sink", None))
2043
+ new.pssdt_varCost_noEff = _keep_cross_triple(
2044
+ getattr(data, "pssdt_varCost_noEff", None))
2045
+ new.p_pssdt_varCost = _keep_cross_triple_param(getattr(data, "p_pssdt_varCost", None))
2046
+ new.p_pdt_varCost_source = _keep_proc_param(getattr(data, "p_pdt_varCost_source", None))
2047
+ new.p_pdt_varCost_sink = _keep_proc_param(getattr(data, "p_pdt_varCost_sink", None))
2048
+ new.p_pdt_varCost_process = _keep_proc_param(getattr(data, "p_pdt_varCost_process", None))
2049
+
2050
+ # ---- Invest params for the cross connections ----
2051
+ new.pd_invest_set = _keep_proc(data.pd_invest_set)
2052
+ new.pd_divest_set = _keep_proc(data.pd_divest_set)
2053
+ new.ed_invest_set = _keep_entity_frame(data.ed_invest_set)
2054
+ new.ed_divest_set = _keep_entity_frame(data.ed_divest_set)
2055
+ new.edd_invest_set = _keep_entity_frame(data.edd_invest_set)
2056
+ new.edd_invest_lookback_set = _keep_entity_frame(data.edd_invest_lookback_set)
2057
+ new.edd_divest_active = _keep_entity_frame(data.edd_divest_active)
2058
+ new.p_entity_max_units = _keep_entity_param(data.p_entity_max_units)
2059
+ new.ed_lifetime_fixed_cost = _keep_entity_param(data.ed_lifetime_fixed_cost)
2060
+ new.ed_lifetime_fixed_cost_divest = _keep_entity_param(data.ed_lifetime_fixed_cost_divest)
2061
+ new.ed_entity_annual_discounted = _keep_entity_param(data.ed_entity_annual_discounted)
2062
+ new.ed_entity_annual_divest_discounted = _keep_entity_param(
2063
+ data.ed_entity_annual_divest_discounted)
2064
+ new.e_invest_total = _keep_entity_frame(data.e_invest_total)
2065
+ new.e_divest_total = _keep_entity_frame(data.e_divest_total)
2066
+ new.e_invest_max_total = _keep_entity_frame(data.e_invest_max_total)
2067
+ new.e_divest_max_total = _keep_entity_frame(data.e_divest_max_total)
2068
+ new.ed_invest_period_set = _keep_entity_frame(data.ed_invest_period_set)
2069
+ new.ed_divest_period_set = _keep_entity_frame(data.ed_divest_period_set)
2070
+ new.ed_invest_max_period = _keep_entity_frame(data.ed_invest_max_period)
2071
+ new.ed_divest_max_period = _keep_entity_frame(data.ed_divest_max_period)
2072
+ new.p_entity_previously_invested_capacity = _keep_entity_param(
2073
+ data.p_entity_previously_invested_capacity)
2074
+ new.p_entity_invested = _keep_entity_param(data.p_entity_invested)
2075
+ new.p_entity_divested = _keep_entity_param(data.p_entity_divested)
2076
+ new.p_entity_all_existing = _keep_entity_param(data.p_entity_all_existing)
2077
+ new.p_ed_fixed_cost = _keep_entity_param(data.p_ed_fixed_cost)
2078
+
2079
+ # ---- Node-keyed / recourse frames ----
2080
+ if not master_mode:
2081
+ # Today's exact path: drop in-region features that reference the
2082
+ # OMITTED terminal nodes (or in-region recourse not part of the
2083
+ # network-only master).
2084
+ for fld in _MASTER_NODE_FIELDS:
2085
+ if hasattr(new, fld):
2086
+ setattr(new, fld, None)
2087
+ else:
2088
+ # Master-hosted mode: KEEP the master-side content. Every
2089
+ # ``_MASTER_NODE_FIELDS`` member is routed through exactly one
2090
+ # of the keep-filters below (partition pinned by test).
2091
+ def _filter_field(obj, col: str, keep: set[str]):
2092
+ """Filter a frame OR Param to ``col ∈ keep`` (no-op on
2093
+ ``None`` / missing column, mirroring the region-side
2094
+ helpers)."""
2095
+ if obj is None:
2096
+ return None
2097
+ if isinstance(obj, Param):
2098
+ return _filter_param(obj, col, keep)
2099
+ return _filter_frame(obj, col, keep)
2100
+
2101
+ def _semi_triples(obj, keys: set[tuple[str, str, str]]):
2102
+ """Keep only the rows whose (p, source, sink) is in *keys*
2103
+ (frame or Param)."""
2104
+ if obj is None:
2105
+ return None
2106
+ frame = obj.frame if isinstance(obj, Param) else obj
2107
+ if not all(c in frame.columns for c in ("p", "source", "sink")):
2108
+ return obj
2109
+ key_df = pl.DataFrame(
2110
+ {
2111
+ "p": [t[0] for t in keys],
2112
+ "source": [t[1] for t in keys],
2113
+ "sink": [t[2] for t in keys],
2114
+ },
2115
+ schema={"p": schema_dtype(_enums, "p"),
2116
+ "source": schema_dtype(_enums, "source"),
2117
+ "sink": schema_dtype(_enums, "sink")},
2118
+ )
2119
+ out = frame.join(key_df, on=("p", "source", "sink"), how="semi")
2120
+ if isinstance(obj, Param):
2121
+ return Param(obj.dims, out, name=obj.name)
2122
+ return out
2123
+
2124
+ master_set = set(master_hosted_nodes)
2125
+ for fld in _MASTER_N_KEYED_FIELDS:
2126
+ if hasattr(new, fld):
2127
+ setattr(new, fld,
2128
+ _filter_field(getattr(data, fld, None), "n",
2129
+ master_set))
2130
+ for fld in _MASTER_CN_KEYED_FIELDS:
2131
+ if hasattr(new, fld):
2132
+ setattr(new, fld,
2133
+ _filter_field(getattr(data, fld, None), "cn",
2134
+ master_cns))
2135
+ for fld in _MASTER_CO2_ARC_FIELDS:
2136
+ if hasattr(new, fld):
2137
+ setattr(new, fld,
2138
+ _semi_triples(getattr(data, fld, None),
2139
+ master_local_keys))
2140
+ for fld in _MASTER_CO2_GROUP_FIELDS:
2141
+ if hasattr(new, fld):
2142
+ setattr(new, fld,
2143
+ _filter_field(getattr(data, fld, None), "g",
2144
+ master_groups_co2))
2145
+ # ``_MASTER_CO2_LOOKUP_FIELDS`` (p_co2_content / p_co2_price):
2146
+ # pure lookup tables joined by the kept CO2 rows — carried whole
2147
+ # through the shallow copy (rows for un-kept groups are inert,
2148
+ # and the PROCESSES feature requires presence).
2149
+ for fld in (_GROUP_FEATURE_SET_FIELDS
2150
+ + _GROUP_FEATURE_PARAM_FIELDS):
2151
+ if hasattr(new, fld):
2152
+ setattr(new, fld,
2153
+ _filter_field(getattr(data, fld, None), "g",
2154
+ master_groups_feature))
2155
+ for fld in _MASTER_GROUP_PROC_FIELDS:
2156
+ if hasattr(new, fld):
2157
+ obj = getattr(data, fld, None)
2158
+ setattr(new, fld,
2159
+ _keep_proc_param(obj) if isinstance(obj, Param)
2160
+ else _keep_proc(obj))
2161
+
2162
+ # ``build_flextool``'s ALWAYS feature requires these four fields to be
2163
+ # PRESENT (not None) even when empty (model.py ``ALWAYS``). Emptying
2164
+ # them (head(0)) — rather than nulling — OMITS every terminal node from
2165
+ # balance while keeping the build's structural precondition satisfied:
2166
+ # the master's ``v_flow`` is then free except for ``maxFlow`` and its
2167
+ # own bound (Phase-3 §1.3).
2168
+ def _empty_like_frame(df: pl.DataFrame | None) -> pl.DataFrame | None:
2169
+ return df.head(0) if df is not None else df
2170
+
2171
+ def _empty_like_param(p: Param | None) -> Param | None:
2172
+ if p is None:
2173
+ return None
2174
+ return Param(p.dims, p.frame.head(0), name=p.name)
2175
+
2176
+ # OMIT every cross-arc TERMINAL node from balance — that is the
2177
+ # "unbalanced virtual node" requirement: with the terminals absent the
2178
+ # master's trade ``v_flow`` is free except for ``maxFlow`` and its own
2179
+ # bound (Phase-3 §1.3). We do NOT empty ``nodeBalance`` entirely,
2180
+ # because ``build_flextool`` requires a non-empty balance-node set to
2181
+ # declare the ``vq_state_up/down`` slack vars and the ``nodeBalance_eq``
2182
+ # row (a None ``nodeBalance_dt`` crashes ``add_var`` / ``add_cstr``).
2183
+ # Instead we keep the NON-terminal nodes: in the reduced master they
2184
+ # carry NO arcs (every non-cross arc was dropped) and NO inflow, so
2185
+ # their balance collapses to ``slack = 0`` — structurally inert, with
2186
+ # zero effect on the trade flow. This satisfies the build precondition
2187
+ # while still omitting exactly the trade terminals.
2188
+ terminal_nodes: set[str] = set()
2189
+ if not master_mode:
2190
+ # Today's exact path: every cross-arc terminal is omitted.
2191
+ for k in cross_keys:
2192
+ terminal_nodes.add(k[1])
2193
+ terminal_nodes.add(k[2])
2194
+ else:
2195
+ # Omit ONLY the region-side endpoints: both terminals of
2196
+ # region↔region cross arcs (as today) plus the region-side
2197
+ # endpoint of each region↔master coupling arc. Master-hosted
2198
+ # endpoints stay balanced (that is the whole point of the
2199
+ # mode); master-local arc endpoints are all master-side and
2200
+ # never omitted.
2201
+ for r in cross.iter_rows(named=True):
2202
+ terminal_nodes.add(r["source"])
2203
+ terminal_nodes.add(r["sink"])
2204
+ for r in region_master.iter_rows(named=True):
2205
+ terminal_nodes.add(
2206
+ r["source"] if r["_src_region"] is not None else r["sink"])
2207
+
2208
+ def _drop_terminal_frame(df: pl.DataFrame | None) -> pl.DataFrame | None:
2209
+ if df is None or "n" not in df.columns:
2210
+ return df
2211
+ return df.filter(~_is_in_keep("n", terminal_nodes))
2212
+
2213
+ def _drop_terminal_param(p: Param | None) -> Param | None:
2214
+ if p is None or "n" not in p.dims:
2215
+ return p
2216
+ return Param(p.dims, p.frame.filter(~_is_in_keep("n", terminal_nodes)),
2217
+ name=p.name)
2218
+
2219
+ new.nodeBalance = _drop_terminal_frame(data.nodeBalance)
2220
+ if not master_mode:
2221
+ new.p_inflow = _empty_like_param(data.p_inflow)
2222
+ else:
2223
+ # The master keeps its own inflow; every other balance node in
2224
+ # the reduced data stays structurally inert (no arcs, no
2225
+ # inflow ⇒ balance collapses to ``slack = 0``, exactly as the
2226
+ # legacy all-emptied path).
2227
+ new.p_inflow = _filter_param(data.p_inflow, "n",
2228
+ set(master_hosted_nodes))
2229
+ new.p_penalty_up = _drop_terminal_param(data.p_penalty_up)
2230
+ new.p_penalty_down = _drop_terminal_param(data.p_penalty_down)
2231
+
2232
+ return new
2233
+
2234
+
2235
+ # ---------------------------------------------------------------------------
2236
+ # Top-level entry point
2237
+ # ---------------------------------------------------------------------------
2238
+
2239
+
2240
+ def split(
2241
+ data: FlexData,
2242
+ *,
2243
+ regions: list[str] | None = None,
2244
+ region_membership: dict[str, dict[str, set[str]]] | None = None,
2245
+ benders_uncap_cross_region: bool = False,
2246
+ master_hosted_nodes: frozenset[str] = frozenset(),
2247
+ ) -> list[RegionSplit]:
2248
+ """Slice a whole-system :class:`FlexData` into per-region splits.
2249
+
2250
+ Parameters
2251
+ ----------
2252
+ data
2253
+ Whole-system :class:`FlexData` (output of :func:`load_flextool`).
2254
+ regions
2255
+ Explicit region list. When ``None``, derives from
2256
+ ``data.group_entity`` / ``data.group_node`` (every group with at
2257
+ least one entity is treated as a region). Callers that have
2258
+ access to ``decomposition_method`` from disk should pass an
2259
+ explicit list.
2260
+ region_membership
2261
+ Pre-computed ``{region: {"nodes": ..., "processes": ...}}`` from
2262
+ :func:`load_region_membership`. When omitted we re-derive from
2263
+ ``data``.
2264
+ benders_uncap_cross_region
2265
+ Benders mode. When ``True``, each cross-region virtual half-flow
2266
+ is built with an effectively-unbounded ``maxFlow`` capacity (a
2267
+ large sentinel) instead of inheriting the original arc's
2268
+ ``p_flow_upper_existing``. In Benders decomposition the TRUE
2269
+ capacity limit ``f ≤ C·unitsize`` is enforced in the MASTER, so
2270
+ a per-region cap would double-bound the flow and (for greenfield
2271
+ cross-region pipes, whose inherited ``existing`` is 0) sever the
2272
+ trade arc to zero — the false-convergence bug. Default ``False``
2273
+ preserves today's inherit-from-original behaviour byte-for-byte.
2274
+ master_hosted_nodes
2275
+ Master-hosted node mode (see
2276
+ :func:`compute_master_hosted_nodes`). With the default empty
2277
+ set the split is byte-identical to today's shared-replicate
2278
+ behaviour. Non-empty: the named nodes live in the Benders
2279
+ MASTER — they are excluded from the shared-replicate set (no
2280
+ region carries them), arcs are classified 4-way
2281
+ (:func:`_classify_arcs`), region↔master coupling arcs get
2282
+ exactly ONE half-flow on the region side, master-local arcs
2283
+ (and the processes ALL of whose arcs are master-local) are
2284
+ dropped from every region and NOT half-flowed, and authored
2285
+ data that cannot be partitioned (a unit straddling the
2286
+ boundary, a user constraint referencing both sides) raises a
2287
+ hard error — never a silent degrade.
2288
+
2289
+ Returns
2290
+ -------
2291
+ list[RegionSplit]
2292
+ One per region, in the order given by ``regions``.
2293
+
2294
+ Notes
2295
+ -----
2296
+ "Shared" entities (in no region) are kept in every region's local
2297
+ set — they're typically things like ``coal_market`` (a commodity
2298
+ node). Cross-region arcs are dropped from the *original* process
2299
+ frames and replaced by virtual half-flow arcs (one in the source
2300
+ region, one in the sink region).
2301
+ """
2302
+ if regions is None:
2303
+ # Derive: every group with at least one membership entity.
2304
+ if data.group_node is not None and data.group_node.height > 0:
2305
+ regions = sorted(
2306
+ set(data.group_node["g"].unique().to_list())
2307
+ )
2308
+ else:
2309
+ regions = []
2310
+ if not regions:
2311
+ return []
2312
+
2313
+ if region_membership is None:
2314
+ region_membership = load_region_membership(data, regions)
2315
+
2316
+ region_nodes = {r: m["nodes"] for r, m in region_membership.items()}
2317
+ region_procs = {r: m["processes"] for r, m in region_membership.items()}
2318
+
2319
+ # Identify shared entities (no region).
2320
+ all_region_nodes: set[str] = set()
2321
+ for ns in region_nodes.values():
2322
+ all_region_nodes |= ns
2323
+ all_region_procs: set[str] = set()
2324
+ for ps in region_procs.values():
2325
+ all_region_procs |= ps
2326
+
2327
+ # Whole-set of nodes and processes.
2328
+ if data.nodeBalance is not None:
2329
+ all_nodes = set(data.nodeBalance["n"].to_list())
2330
+ else:
2331
+ all_nodes = set()
2332
+ if data.process_source_sink is not None:
2333
+ all_procs = set(data.process_source_sink["p"].unique().to_list())
2334
+ else:
2335
+ all_procs = set()
2336
+
2337
+ shared_nodes = all_nodes - all_region_nodes
2338
+ shared_procs = all_procs - all_region_procs
2339
+
2340
+ if master_hosted_nodes:
2341
+ overlap = set(master_hosted_nodes) & all_region_nodes
2342
+ if overlap:
2343
+ raise RuntimeError(
2344
+ f"split: master_hosted_nodes overlap region membership: "
2345
+ f"{sorted(overlap)} — a node is either master-hosted "
2346
+ f"(in no region group) or in exactly one region, never "
2347
+ f"both."
2348
+ )
2349
+ # Master nodes are never shared-replicated: regions must not
2350
+ # carry them.
2351
+ shared_nodes -= set(master_hosted_nodes)
2352
+
2353
+ # Classify cross-region arcs.
2354
+ if data.process_source_sink is None:
2355
+ return [
2356
+ RegionSplit(region=r, data=data, half_flows=[])
2357
+ for r in regions
2358
+ ]
2359
+
2360
+ pss_tagged, cross, region_master, master_local = _classify_arcs(
2361
+ data.process_source_sink, region_nodes,
2362
+ master_nodes=master_hosted_nodes,
2363
+ )
2364
+
2365
+ master_local_procs: set[str] = set()
2366
+ master_cns: set[str] = set()
2367
+ master_groups: set[str] = set()
2368
+ if master_hosted_nodes:
2369
+ # Hard validation FIRST (D-a): never silently mis-partition
2370
+ # authored data.
2371
+ _validate_no_straddling_units(
2372
+ data, all_region_nodes, master_hosted_nodes)
2373
+ master_local_procs = _master_local_procs(
2374
+ data.process_source_sink, master_local)
2375
+ _validate_user_constraints(
2376
+ data, all_region_nodes, master_hosted_nodes,
2377
+ master_local_procs)
2378
+ # All-master user constraints / feature groups / CO2-cap groups
2379
+ # live whole in the master (kept there by
2380
+ # ``master_network_data``): regions must drop their rows — a
2381
+ # region copy would degenerate (constraint: ``0 sense
2382
+ # constant``; feature group: empty membership still charging
2383
+ # slack penalty). The group/CO2 helpers also hard-error on
2384
+ # straddling groups (same validation family as the unit /
2385
+ # user-constraint checks above).
2386
+ master_cns = _master_side_constraint_ids(
2387
+ data, all_region_nodes, master_hosted_nodes,
2388
+ master_local_procs)
2389
+ _ml_triples = {
2390
+ (r["p"], r["source"], r["sink"])
2391
+ for r in master_local.iter_rows(named=True)
2392
+ }
2393
+ _rm_triples = {
2394
+ (r["p"], r["source"], r["sink"])
2395
+ for r in region_master.iter_rows(named=True)
2396
+ }
2397
+ master_groups = (
2398
+ _master_side_feature_groups(
2399
+ data, all_region_nodes, master_hosted_nodes)
2400
+ | _co2_master_partition(data, _ml_triples, _rm_triples)
2401
+ )
2402
+ # Master-local procs live wholly in the master (F3): regions
2403
+ # carry neither their arcs nor their entity rows.
2404
+ shared_procs -= master_local_procs
2405
+
2406
+ half_flows_by_region = _make_half_flows(cross, region_master)
2407
+
2408
+ cross_arcs_by_pss: set[tuple[str, str, str]] = set()
2409
+ for r in cross.iter_rows(named=True):
2410
+ cross_arcs_by_pss.add((r["p"], r["source"], r["sink"]))
2411
+ # Region↔master coupling arcs are replaced by their single-sided
2412
+ # half-flow; master-local arcs are dropped outright (the master
2413
+ # keeps the whole original arc) — both classes must vanish from
2414
+ # every region's process frames.
2415
+ for frame in (region_master, master_local):
2416
+ for r in frame.iter_rows(named=True):
2417
+ cross_arcs_by_pss.add((r["p"], r["source"], r["sink"]))
2418
+
2419
+ # Phase 4 — virtual half-flow entities ("hf_pipe_*" / "pipe_*__*__*")
2420
+ # are created at runtime by ``_make_half_flows``; they are not in the
2421
+ # source DB and therefore not in the axis_enums vocabulary built by
2422
+ # ``build_axis_enums``. Downstream filter operations like
2423
+ # ``pl.col("p").is_in([...keep_procs incl. virtual_p...])`` raise
2424
+ # ``conversion from str to enum failed`` when polars casts the
2425
+ # comparison list against the Enum dtype. Widen the live vocabulary
2426
+ # to include the virtual tokens for the duration of the split.
2427
+ _virt_p: set[str] = set()
2428
+ _virt_n: set[str] = set()
2429
+ for _hfs in half_flows_by_region.values():
2430
+ for _hf in _hfs:
2431
+ _virt_p.add(_hf.virtual_p)
2432
+ _virt_n.add(_hf.virtual_node)
2433
+ _enums_token = None
2434
+ # Base the widening on the SOURCE's own axis_enums snapshot — that
2435
+ # is guaranteed to match the dtypes embedded in ``data``'s frames.
2436
+ # The live global ContextVar may have been overwritten by an
2437
+ # unrelated ``load_flextool`` call between the lh2 fixture's load
2438
+ # and this split (e.g. a sibling test loaded a different DB), in
2439
+ # which case widening the live vocabulary would yield an Enum that
2440
+ # doesn't contain ``data``'s entity tokens. Fall back to the live
2441
+ # global only when ``data`` lacks its own snapshot.
2442
+ _base_enums = (getattr(data, "_axis_enums", None)
2443
+ or get_global_axis_enums())
2444
+ if _base_enums is not None and (_virt_p or _virt_n):
2445
+ _ext: dict[str, pl.Enum] = dict(_base_enums)
2446
+ _virt_e = _virt_p | _virt_n
2447
+ for _axis_name, _new_toks in (
2448
+ ("p", _virt_p),
2449
+ ("n", _virt_n),
2450
+ ("source", _virt_n),
2451
+ ("sink", _virt_n),
2452
+ ("e", _virt_e),
2453
+ ):
2454
+ _existing = _ext.get(_axis_name)
2455
+ if _existing is None:
2456
+ continue
2457
+ _existing_cats = list(_existing.categories)
2458
+ _existing_set = set(_existing_cats)
2459
+ _add = [t for t in _new_toks if t not in _existing_set]
2460
+ if _add:
2461
+ _ext[_axis_name] = pl.Enum(_existing_cats + _add)
2462
+ _enums_token = set_global_axis_enums(_ext)
2463
+
2464
+ try:
2465
+ splits: list[RegionSplit] = []
2466
+ for r in regions:
2467
+ keep_nodes = region_nodes.get(r, set()) | shared_nodes
2468
+ keep_procs = region_procs.get(r, set()) | shared_procs
2469
+ # Also keep cross-region pipes' original `p` membership in this
2470
+ # region IF the in-region terminal is here. We'll drop the
2471
+ # specific (p, source, sink) cross-arc rows below; but we keep
2472
+ # the process p in keep_procs so the OTHER direction (back-flow)
2473
+ # which has the in-region node as its sink/source is retained.
2474
+ # In fact, we add the original cross-region pipe p iff this
2475
+ # region has a half-flow involving that p.
2476
+ for hf in half_flows_by_region.get(r, []):
2477
+ keep_procs.add(hf.original_p)
2478
+ keep_procs.add(hf.virtual_p)
2479
+ # Master-local procs are subtracted AFTER the half-flow
2480
+ # additions: region membership may name them (e.g. a unit
2481
+ # whose every arc moved master-side), but regions must
2482
+ # carry no rows for them (F3). Half-flow originals are
2483
+ # never master-local (they have a region-side terminal).
2484
+ if master_local_procs:
2485
+ keep_procs -= master_local_procs
2486
+
2487
+ rdata = _build_region_data(
2488
+ src=data,
2489
+ region=r,
2490
+ keep_nodes=keep_nodes,
2491
+ keep_procs=keep_procs,
2492
+ half_flows=half_flows_by_region.get(r, []),
2493
+ cross_arcs_by_pss=cross_arcs_by_pss,
2494
+ benders_uncap_cross_region=benders_uncap_cross_region,
2495
+ master_hosted_nodes=master_hosted_nodes,
2496
+ master_local_procs=master_local_procs,
2497
+ master_cns=master_cns,
2498
+ master_groups=master_groups,
2499
+ )
2500
+ splits.append(RegionSplit(
2501
+ region=r,
2502
+ data=rdata,
2503
+ half_flows=half_flows_by_region.get(r, []),
2504
+ ))
2505
+ return splits
2506
+ finally:
2507
+ if _enums_token is not None:
2508
+ reset_global_axis_enums(_enums_token)