From 425d82380400cb4d5b3822e6cfc935a1c4e1f431 Mon Sep 17 00:00:00 2001
From: lumbric <lumbric@gmail.com>
Date: Mon, 11 Mar 2024 13:55:25 +0100
Subject: [PATCH] Add methods for custom variables and constraints

---
 syfop/network.py      | 20 ++++++++++++++++++++
 tests/test_network.py | 44 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/syfop/network.py b/syfop/network.py
index 16d423e..c3cee73 100644
--- a/syfop/network.py
+++ b/syfop/network.py
@@ -202,6 +202,26 @@ def _generate_optimization_model(self, nodes, solver_dir):
 
         return model
 
+    def add_variables(self, *args, **kwargs):
+        """Add custom variables to the linopy optimization model. See
+        ``linopy.Model.add_variables()`` for a documentation of the parameters.
+
+        This method must be called before ``Network.optimize()`` is called.
+
+        """
+        return self.model.add_variables(*args, **kwargs)
+
+    def add_constraints(self, *args, **kwargs):
+        """Add custom constraints to the linopy optimization model. See
+        ``linopy.Model.add_constraints()`` for a documentation of the parameters.
+
+        To create the constraint, variables can be accessed via ``Network.model``
+
+        This method must be called before ``Network.optimize()`` is called.
+
+        """
+        return self.model.add_constraints(*args, **kwargs)
+
     def _check_storage_level_zero(self):
         """This is a basic plausibility check. A storage which is never full or never empty
         could be replaced by a smaller storage, which would be advantageous if costs are
diff --git a/tests/test_network.py b/tests/test_network.py
index b6578c9..ab96a5e 100644
--- a/tests/test_network.py
+++ b/tests/test_network.py
@@ -437,3 +437,47 @@ def test_infeasible_network():
     )
     with pytest.raises(SolverError, match=error_msg):
         network.optimize(default_solver)
+
+
+def test_network_add_constraints():
+    """Test adding custom constraints."""
+    wind = NodeScalableInputProfile(
+        name="wind",
+        input_flow=const_time_series(0.5),
+        costs=1,
+        output_unit="MW",
+    )
+    demand = NodeFixOutputProfile(
+        name="demand",
+        inputs=[wind],
+        input_commodities="electricity",
+        output_flow=const_time_series(5.0),
+        costs=0,
+        output_unit="MW",
+    )
+    curtailment = Node(
+        name="curtailment",
+        inputs=[wind],
+        input_commodities="electricity",
+        costs=0,
+        output_unit="MW",
+    )
+    network = Network([wind, demand, curtailment])
+
+    # Let's make Sebastian happy and force the system to have at least 15MW of wind :)
+    network.add_constraints(wind.size >= 15)
+    network.optimize()
+
+    assert network.model.solution.size_wind == 15.0
+    np.testing.assert_array_almost_equal(network.model.solution.output_flow_curtailment, 2.5)
+
+
+def test_network_add_variables():
+    """Test adding custom variables."""
+    network = simple_demand_network()
+    total_energy = network.add_variables(name="total_energy")
+    network.add_constraints(total_energy - network.nodes_dict["wind"].input_flows[""].sum() == 0.0)
+    network.optimize()
+
+    # we need 5 MW constant demand
+    assert network.model.solution.total_energy == DEFAULT_NUM_TIME_STEPS * 5.0