replace pyomo with linopy

scripts/cluster_network.py

@@ -127,7 +127,7 @@
 import geopandas as gpd
 import numpy as np
 import pandas as pd
-import pyomo.environ as po
+import linopy
 import pypsa
 from _helpers import (
     REGION_COLS,
@@ -336,47 +336,19 @@ def distribute_clusters(
         distribution_factor.sum(), 1.0, rtol=1e-3
     ), f"Country weights L must sum up to 1.0 when distributing clusters. Is {distribution_factor.sum()}."
 
-    m = po.ConcreteModel()
-
-    def n_bounds(model, *n_id):
-        """
-        Create a function that makes a bound pair for pyomo.
-
-        Use n_bounds(model, n_id) if N is Single-Index
-        Use n_bounds(model, *n_id) if N is Multi-Index
-        Example: https://pyomo.readthedocs.io/en/stable/pyomo_modeling_components/Variables.html
-
-        Returns
-        -------
-        bounds = A function (or Python object) that gives a (lower,upper) bound pair i.e.(1,10) for the variable
-        """
-        return (1, N[n_id])
-
-    m.n = po.Var(list(distribution_factor.index), bounds=n_bounds, domain=po.Integers)
-    m.tot = po.Constraint(expr=(po.summation(m.n) == n_clusters))
-    m.objective = po.Objective(
-        expr=sum(
-            (m.n[i] - distribution_factor.loc[i] * n_clusters) ** 2
-            for i in distribution_factor.index
-        ),
-        sense=po.minimize,
-    )
-
-    opt = po.SolverFactory(solver_name)
-    if not opt.has_capability("quadratic_objective"):
-        logger.warning(
-            f"The configured solver `{solver_name}` does not support quadratic objectives. Falling back to `ipopt`."
-        )
-        opt = po.SolverFactory("ipopt")
-
-    results = opt.solve(m)
-    assert (
-        results["Solver"][0]["Status"] == "ok"
-    ), f"Solver returned non-optimally: {results}"
-
-    return (
-        pd.Series(m.n.get_values(), index=distribution_factor.index).round().astype(int)
-    )
+    m = linopy.Model()
+    clusters = m.add_variables(
+        lower=1, upper=N, coords=[L.index], name="n", integer=True)
+
+    m.add_constraints(clusters.sum() == n_clusters, name="tot")
+    m.objective = (
+        clusters * clusters - 2 * clusters * L * n_clusters
+    )  # + (L * n_clusters) ** 2 (constant)
+    if solver_name == "gurobi":
+        logger.getLogger("gurobipy").propagate = False
+
+    m.solve(solver_name=solver_name)
+    return m.solution["n"].to_series().astype(int)
 
 
 def busmap_for_gadm_clusters(inputs, n, gadm_layer_id, geo_crs, country_list):
@@ -653,7 +625,7 @@ def cluster_regions(busmaps, inputs, output):
         from _helpers import mock_snakemake
 
         snakemake = mock_snakemake(
-            "cluster_network", network="elec", simpl="", clusters="4"
+            "cluster_network", network="elec", simpl="", clusters="10"
         )
     configure_logging(snakemake)