From 5dfeb6efbcfc98d39d914c2a70fc193ec420869b Mon Sep 17 00:00:00 2001
From: Jonas Hoersch <jonas@chaoflow.net>
Date: Mon, 4 Feb 2019 18:32:56 +0100
Subject: [PATCH] trace_solve_network: Save intermediate line capacities

---
 Snakefile                      |  9 +++++++
 scripts/solve_network.py       |  7 ++++-
 scripts/trace_solve_network.py | 47 ++++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 scripts/trace_solve_network.py

diff --git a/Snakefile b/Snakefile
index cda7be14..c7e04b0e 100644
--- a/Snakefile
+++ b/Snakefile
@@ -232,6 +232,15 @@ rule solve_network:
     # group: "solve" # with group, threads is ignored https://bitbucket.org/snakemake/snakemake/issues/971/group-job-description-does-not-contain
     script: "scripts/solve_network.py"
 
+rule trace_solve_network:
+    input: "networks/{network}_s{simpl}_{clusters}_l{ll}_{opts}.nc"
+    output: "results/networks/{network}_s{simpl}_{clusters}_l{ll}_{opts}_trace.nc"
+    shadow: "shallow"
+    log: python="logs/{network}_s{simpl}_{clusters}_l{ll}_{opts}_python_trace.log",
+    threads: 4
+    resources: mem=memory
+    script: "scripts/trace_solve_network.py"
+
 rule solve_operations_network:
     input:
         unprepared="networks/{network}_s{simpl}_{clusters}.nc",
diff --git a/scripts/solve_network.py b/scripts/solve_network.py
index 4df98c5c..67251671 100644
--- a/scripts/solve_network.py
+++ b/scripts/solve_network.py
@@ -120,7 +120,7 @@ def fix_branches(n, lines_s_nom=None, links_p_nom=None):
         if isinstance(n.opt, pypsa.opf.PersistentSolver):
             n.opt.update_var(n.model.link_p_nom)
 
-def solve_network(n, config=None, solver_log=None, opts=None):
+def solve_network(n, config=None, solver_log=None, opts=None, callback=None):
     if config is None:
         config = snakemake.config['solving']
     solve_opts = config['options']
@@ -212,6 +212,7 @@ def solve_network(n, config=None, solver_log=None, opts=None):
 
         lines['s_nom_opt'] = lines['s_nom'] * n.lines['num_parallel'].where(n.lines.type != '', 1.)
         status, termination_condition = run_lopf(n, allow_warning_status=True)
+        if callback is not None: callback(n, iteration, status)
 
         def msq_diff(n):
             lines_err = np.sqrt(((n.lines['s_nom_opt'] - lines['s_nom_opt'])**2).mean())/lines['s_nom_opt'].mean()
@@ -230,12 +231,16 @@ def solve_network(n, config=None, solver_log=None, opts=None):
             iteration += 1
 
             status, termination_condition = run_lopf(n, allow_warning_status=True)
+            if callback is not None: callback(n, iteration, status)
+
 
         update_line_parameters(n, zero_lines_below=100)
 
         logger.info("Starting last run with fixed extendable lines")
 
+    iteration += 1
     status, termination_condition = run_lopf(n, fix_ext_lines=True)
+    if callback is not None: callback(n, iteration, status)
 
     return n
 
diff --git a/scripts/trace_solve_network.py b/scripts/trace_solve_network.py
new file mode 100644
index 00000000..da40b9f7
--- /dev/null
+++ b/scripts/trace_solve_network.py
@@ -0,0 +1,47 @@
+import numpy as np
+import pandas as pd
+import logging
+logger = logging.getLogger(__name__)
+
+from solve_network import patch_pyomo_tmpdir, prepare_network, solve_network
+
+import pypsa
+
+
+if __name__ == "__main__":
+    # Detect running outside of snakemake and mock snakemake for testing
+    if 'snakemake' not in globals():
+        from vresutils.snakemake import MockSnakemake, Dict
+        snakemake = MockSnakemake(
+            wildcards=dict(network='elec', simpl='', clusters='45', lv='1.25', opts='Co2L-3H'),
+            input=["networks/{network}_s{simpl}_{clusters}_lv{lv}_{opts}.nc"],
+            output=["results/networks/s{simpl}_{clusters}_lv{lv}_{opts}_trace.nc"],
+            log=dict(python="logs/{network}_s{simpl}_{clusters}_lv{lv}_{opts}_python_trace.log")
+        )
+
+    tmpdir = snakemake.config['solving'].get('tmpdir')
+    if tmpdir is not None:
+        patch_pyomo_tmpdir(tmpdir)
+
+    logging.basicConfig(filename=snakemake.log.python,
+                        level=snakemake.config['logging_level'])
+
+    n = pypsa.Network(snakemake.input[0])
+
+    solver_log = 'solver.log'
+    config = snakemake.config['solving']
+    opts = snakemake.wildcards.opts.split('-')
+
+    def save_optimal_capacities(net, iteration, status):
+        net.lines[f"s_nom_opt_{iteration}"] = net.lines["s_nom_opt"]
+        net.links[f"p_nom_opt_{iteration}"] = net.links["p_nom_opt"]
+        setattr(net, f"status_{iteration}", status)
+        setattr(net, f"objective_{iteration}", net.objective)
+        net.iteration = iteration
+
+        net.export_to_netcdf(snakemake.output[0])
+
+    config['options']['max_iterations'] = 12
+    n = prepare_network(n, config['options'])
+    n = solve_network(n, config, solver_log, opts, save_optimal_capacities)
+    n.export_to_netcdf(snakemake.output[0])