From 6ebef6028a67ffbbed62c44a02af1ca35f36c537 Mon Sep 17 00:00:00 2001
From: Fridolin Glatter <glatter@iiasa.ac.at>
Date: Thu, 28 Nov 2024 12:40:02 +0100
Subject: [PATCH] Introduce run.optimization.remove_solution()

---
 ixmp4/core/optimization/data.py |  6 ++++++
 tests/core/test_run.py          | 26 ++++++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/ixmp4/core/optimization/data.py b/ixmp4/core/optimization/data.py
index c604dbe1..d3f8cbfb 100644
--- a/ixmp4/core/optimization/data.py
+++ b/ixmp4/core/optimization/data.py
@@ -29,3 +29,9 @@ def __init__(self, run: Run, **kwargs: Backend) -> None:
         self.scalars = ScalarRepository(_backend=self.backend, _run=run)
         self.tables = TableRepository(_backend=self.backend, _run=run)
         self.variables = VariableRepository(_backend=self.backend, _run=run)
+
+    def remove_solution(self) -> None:
+        for equation in self.equations.list():
+            equation.remove_data()
+        for variable in self.variables.list():
+            variable.remove_data()
diff --git a/tests/core/test_run.py b/tests/core/test_run.py
index e75aa4f5..1a128555 100644
--- a/tests/core/test_run.py
+++ b/tests/core/test_run.py
@@ -181,3 +181,29 @@ def delete_all_datapoints(self, run: ixmp4.Run) -> None:
             run.iamc.remove(cat, type=ixmp4.DataPoint.Type.CATEGORICAL)
         if not datetime.empty:
             run.iamc.remove(datetime, type=ixmp4.DataPoint.Type.DATETIME)
+
+    def test_run_remove_solution(self, platform: ixmp4.Platform) -> None:
+        run = platform.runs.create("Model", "Scenario")
+        indexset = run.optimization.indexsets.create("Indexset")
+        indexset.add(["foo", "bar"])
+        test_data = {
+            "Indexset": ["bar", "foo"],
+            "levels": [2.0, 1],
+            "marginals": [0, "test"],
+        }
+        run.optimization.equations.create(
+            "Equation",
+            constrained_to_indexsets=[indexset.name],
+        ).add(test_data)
+        run.optimization.variables.create(
+            "Variable",
+            constrained_to_indexsets=[indexset.name],
+        ).add(test_data)
+
+        run.optimization.remove_solution()
+        # Need to fetch them here even if fetched before because API layer might not
+        # forward changes automatically
+        equation = run.optimization.equations.get("Equation")
+        variable = run.optimization.variables.get("Variable")
+        assert equation.data == {}
+        assert variable.data == {}