Repository: PyPSA/linopy
Pull request #10 (2021-11-22 07:54:56)

variables and constraints: Make __repr__ more light-weight


Patch:
diff --git a/doc/release_notes.rst b/doc/release_notes.rst
index e55b340..95b5b7d 100644
--- a/doc/release_notes.rst
+++ b/doc/release_notes.rst
@@ -4,11 +4,6 @@ Release Notes
 Upcoming Release
 ----------------
 
-* The export of LP files was restructured and is now using Unicode dtype under the hook to ensure performant export of long constraints.
-
-**Breaking changes**
-
-* The low level IO function ``linopy.io.str_array_to_file`` was renamed to ``linopy.io.array_to_file``, the function ``linopy.io.join_str_arrays`` was removed.
 
 Version 0.0.4
 -------------
diff --git a/linopy/constraints.py b/linopy/constraints.py
index a8a70dc..40ac46a 100644
--- a/linopy/constraints.py
+++ b/linopy/constraints.py
@@ -4,6 +4,7 @@ Linopy constraints module.
 This module contains implementations for the Constraint{s} class.
 """
 
+import re
 from dataclasses import dataclass
 from typing import Any, Sequence, Union
 
@@ -117,7 +118,7 @@ class Constraints:
 
     dataset_attrs = ["labels", "coeffs", "vars", "sign", "rhs"]
     dataset_names = [
-        "Constraint labels",
+        "Labels",
         "Left-hand-side coefficients",
         "Left-hand-side variables",
         "Signs",
@@ -127,13 +128,18 @@ class Constraints:
     def __repr__(self):
         """Return a string representation of the linopy model."""
         r = "linopy.model.Constraints"
-        line = "=" * len(r)
+        line = "-" * len(r)
         r += f"\n{line}\n\n"
+        # matches string between "Data variables" and "Attributes"/end of string
+        coordspattern = r"(?s)(?<=\<xarray\.Dataset\>\n).*?(?=Data variables:\n)"
+        datapattern = r"(?s)(?<=Data variables:\n).*?(?=($|\nAttributes))"
         for (k, K) in zip(self.dataset_attrs, self.dataset_names):
-            s = getattr(self, k).__repr__().split("\n", 1)[1]
-            s = s.replace("Data variables:\n", "Data:\n")
+            orig = getattr(self, k).__repr__()
+            if k == "labels":
+                r += re.search(coordspattern, orig).group() + "\n"
+            data = re.search(datapattern, orig).group()
             line = "-" * (len(K) + 1)
-            r += f"{K}:\n{line}\n{s}\n\n"
+            r += f"{K}:\n{data}\n\n"
         return r
 
     def __getitem__(
@@ -180,7 +186,8 @@ class Constraints:
         """Coefficient range of the constraint."""
         return (
             xr.concat(
-                [self.coeffs.min(), self.coeffs.max()], dim=pd.Index(["min", "max"]),
+                [self.coeffs.min(), self.coeffs.max()],
+                dim=pd.Index(["min", "max"]),
             )
             .to_dataframe()
             .T
diff --git a/linopy/io.py b/linopy/io.py
index e8e0810..6da31c9 100644
--- a/linopy/io.py
+++ b/linopy/io.py
@@ -4,46 +4,41 @@
 import logging
 import os
 import time
+from functools import partial, reduce
 
 import numpy as np
 import xarray as xr
-from numpy import dtype
-from xarray import DataArray, apply_ufunc, concat
+from xarray import apply_ufunc
 
 logger = logging.getLogger(__name__)
 
 
-ufunc_kwargs = dict(vectorize=True)
-concat_dim = "_concat_dim"
-concat_kwargs = dict(dim=concat_dim, coords="minimal")
-
+ufunc_kwargs = dict(dask="parallelized", vectorize=True, output_dtypes=[object])
 
 # IO functions
 def to_float_str(da):
     """Convert a float array to a string array with lp like format for coefficients."""
-    return apply_ufunc(
-        lambda f: "%+f" % f,
-        da.fillna(0).compute(),
-        **ufunc_kwargs,
-    )
+    return apply_ufunc(lambda f: "%+f" % f, da.fillna(0), **ufunc_kwargs)
 
 
 def to_int_str(da, nonnans=None):
     """Convert a int array to a string array."""
-    return xr.apply_ufunc(lambda d: "%d" % d, da.fillna(0).compute(), **ufunc_kwargs)
+    return xr.apply_ufunc(lambda d: "%d" % d, da.fillna(0), **ufunc_kwargs)
 
 
-def array_to_file(array, fn):
-    """Elementwise writing out string values to a file."""
-    write = np.frompyfunc(fn.write, 1, 1)
+def join_str_arrays(arraylist):
+    """Join string array together (elementwise concatenation of strings)."""
+    func = partial(np.add, dtype=object)  # np.core.defchararray.add
+    return reduce(func, arraylist, "")
 
-    def func(data):
-        return write(data.ravel()).reshape(data.shape)
 
+def str_array_to_file(array, fn):
+    """Elementwise writing out string values to a file."""
     return xr.apply_ufunc(
-        func,
+        lambda x: fn.write(x),
         array,
         dask="parallelized",
+        vectorize=True,
         output_dtypes=[int],
     )
 
@@ -54,85 +49,72 @@ def objective_to_file(m, f):
     coef = m.objective.coeffs
     var = m.objective.vars
 
-    objective = [
-        to_float_str(coef),
-        DataArray(" x").astype(dtype("<U9")),
-        to_int_str(var),
-        DataArray("\n").astype(dtype("<U9")),
-    ]
     nonnans = coef.notnull() & (var != -1)
-    objective = concat(objective, **concat_kwargs).where(nonnans, "")
-    objective = objective.transpose(..., concat_dim)
-    array_to_file(objective.compute(), f)
+    join = [to_float_str(coef), " x", to_int_str(var), "\n"]
+    objective_str = join_str_arrays(join).where(nonnans, "")
+    str_array_to_file(objective_str, f).compute()
 
 
 def constraints_to_file(m, f):
     """Write out the constraints of a model to a lp file."""
     f.write("\n\ns.t.\n\n")
-    for name, labels in m.constraints.labels.items():
-
-        dim = f"{name}_term"
-
-        lhs = [
-            to_float_str(m.constraints.coeffs[name]),
-            DataArray(" x").astype(dtype("<U9")),
-            to_int_str(m.constraints.vars[name]),
-            DataArray("\n").astype(dtype("<U9")),
-        ]
-
-        nonnans_terms = m.constraints.vars[name] != -1
-        lhs = concat(lhs, **concat_kwargs).where(nonnans_terms, "")
-        lhs = lhs.stack(_=[dim, concat_dim]).drop("_").rename(_=dim)
-
-        constraints = [
-            DataArray("c").astype(dtype("<U9")),
-            to_int_str(labels),
-            DataArray(":\n").astype(dtype("<U9")),
-            lhs,
-            m.constraints.sign[name],
-            DataArray("\n").astype(dtype("<U9")),
-            to_float_str(m.constraints.rhs[name]),
-            DataArray("\n\n").astype(dtype("<U9")),
-        ]
-
-        da = concat(constraints, dim=dim, coords="minimal")
-        da = da.where(labels != -1, "").transpose(..., dim)
-        array_to_file(da.compute(), fn=f)
+    labels = m.constraints.labels
+    vars = m.constraints.vars
+    coeffs = m.constraints.coeffs
+    sign = m.constraints.sign
+    rhs = m.constraints.rhs
+
+    term_names = [f"{n}_term" for n in labels]
+
+    nonnans = coeffs.notnull() & (vars != -1)
+    join = [to_float_str(coeffs), " x", to_int_str(vars), "\n"]
+    lhs_str = join_str_arrays(join).where(nonnans, "").reduce(np.sum, term_names)
+    # .sum() does not work
+
+    nonnans = nonnans.any(term_names) & (labels != -1) & sign.notnull() & rhs.notnull()
+
+    join = [
+        "c",
+        to_int_str(labels),
+        ": \n",
+        lhs_str,
+        sign,
+        "\n",
+        to_float_str(rhs),
+        "\n\n",
+    ]
+    constraints_str = join_str_arrays(join).where(nonnans, "")
+    str_array_to_file(constraints_str, f).compute()
 
 
 def bounds_to_file(m, f):
     """Write out variables of a model to a lp file."""
     f.write("\nbounds\n")
-    for name, labels in m.variables.labels[m._non_binary_variables].items():
-
-        bounds = [
-            to_float_str(m.variables.lower[name]),
-            DataArray(" <= x").astype(dtype("<U9")),
-            to_int_str(labels),
-            DataArray(" <= ").astype(dtype("<U9")),
-            to_float_str(m.variables.upper[name]),
-            DataArray("\n").astype(dtype("<U9")),
-        ]
-
-        bounds = concat(bounds, **concat_kwargs).where(labels != -1, "")
-        bounds = bounds.transpose(..., concat_dim)
-        array_to_file(bounds.compute(), f)
+    lower = m.variables.lower[m._non_binary_variables]
+    labels = m.variables.labels[m._non_binary_variables]
+    upper = m.variables.upper[m._non_binary_variables]
+
+    nonnans = lower.notnull() & upper.notnull() & (labels != -1)
+    join = [
+        to_float_str(lower),
+        " <= x",
+        to_int_str(labels),
+        " <= ",
+        to_float_str(upper),
+        "\n",
+    ]
+    bounds_str = join_str_arrays(join).where(nonnans, "")
+    str_array_to_file(bounds_str, f).compute()
 
 
 def binaries_to_file(m, f):
     """Write out binaries of a model to a lp file."""
     f.write("\nbinary\n")
-    for name, labels in m.binaries.labels.items():
-
-        binaries = [
-            DataArray("x").astype(dtype("<U9")),
-            to_int_str(labels),
-            DataArray("\n").astype(dtype("<U9")),
-        ]
 
-        binaries = concat(binaries, **concat_kwargs).where(labels != -1, "")
-        binaries = binaries.transpose(..., concat_dim)
-        array_to_file(binaries.compute(), fn=f)
+    labels = m.binaries.labels
+    nonnans = labels != -1
+    binaries_str = join_str_arrays(["x", to_int_str(labels), "\n"]).where(nonnans, "")
+    str_array_to_file(binaries_str, f).compute()
 
 
 def to_file(m, fn):
diff --git a/linopy/variables.py b/linopy/variables.py
index beeb5ed..1b6cb41 100644
--- a/linopy/variables.py
+++ b/linopy/variables.py
@@ -4,6 +4,7 @@ Linopy variables module.
 This module contains variable related definitions of the package.
 """
 
+import re
 from dataclasses import dataclass
 from typing import Any, Sequence, Union
 
@@ -221,7 +222,7 @@ class Variables:
     model: Any = None  # Model is not defined due to circular imports
 
     dataset_attrs = ["labels", "lower", "upper"]
-    dataset_names = ["Variables labels", "Lower bounds", "Upper bounds"]
+    dataset_names = ["Labels", "Lower bounds", "Upper bounds"]
 
     def __getitem__(
         self, names: Union[str, Sequence[str]]
@@ -236,13 +237,17 @@ class Variables:
     def __repr__(self):
         """Return a string representation of the linopy model."""
         r = "linopy.model.Variables"
-        line = "=" * len(r)
+        line = "-" * len(r)
         r += f"\n{line}\n\n"
+        # matches string between "Data variables" and "Attributes"/end of string
+        coordspattern = r"(?s)(?<=\<xarray\.Dataset\>\n).*?(?=Data variables:\n)"
+        datapattern = r"(?s)(?<=Data variables:\n).*?(?=($|\nAttributes))"
         for (k, K) in zip(self.dataset_attrs, self.dataset_names):
-            s = getattr(self, k).__repr__().split("\n", 1)[1]
-            s = s.replace("Data variables:\n", "Data:\n")
-            line = "-" * (len(K) + 1)
-            r += f"{K}:\n{line}\n{s}\n\n"
+            orig = getattr(self, k).__repr__()
+            if k == "labels":
+                r += re.search(coordspattern, orig).group() + "\n"
+            data = re.search(datapattern, orig).group()
+            r += f"{K}:\n{data}\n\n"
         return r
 
     def __iter__(self):


Test patch:
diff --git a/test/test_constraint.py b/test/test_constraint.py
index e74c2b7..aaa4e50 100644
--- a/test/test_constraint.py
+++ b/test/test_constraint.py
@@ -18,6 +18,7 @@ from linopy import Model
 
 def test_constraint_repr():
     m = Model()
+
     x = m.add_variables()
     c = m.add_constraints(x, ">=", 0)
     c.__repr__()
@@ -26,8 +27,9 @@ def test_constraint_repr():
 
 def test_constraints_repr():
     m = Model()
+    m.constraints.__repr__()
     x = m.add_variables()
-    c = m.add_constraints(x, ">=", 0)
+    m.add_constraints(x, ">=", 0)
     m.constraints.__repr__()
 
 
diff --git a/test/test_io.py b/test/test_io.py
index d691b46..fe404aa 100644
--- a/test/test_io.py
+++ b/test/test_io.py
@@ -23,7 +23,7 @@ def test_str_arrays():
     y = m.add_variables(0, pd.DataFrame([[1, 2], [3, 4], [5, 6]]).T)
 
     da = to_int_str(x)
-    assert da.dtype.type is np.str_
+    assert da.dtype == object
 
 
 def test_str_arrays_chunked():
@@ -33,7 +33,7 @@ def test_str_arrays_chunked():
     y = m.add_variables(0, pd.DataFrame([[1, 2], [3, 4], [5, 6]]).T)
 
     da = to_int_str(y).compute()
-    assert da.dtype.type is np.str_
+    assert da.dtype == object
 
 
 def test_str_arrays_with_nans():
@@ -45,7 +45,7 @@ def test_str_arrays_with_nans():
     assert m["x"][-1].item() == -1
 
     da = to_int_str(m["x"])
-    assert da.dtype.type is np.str_
+    assert da.dtype == object
 
 
 def test_to_netcdf(tmp_path):
diff --git a/test/test_optimization.py b/test/test_optimization.py
index 0bba5b4..bc0abc7 100644
--- a/test/test_optimization.py
+++ b/test/test_optimization.py
@@ -28,20 +28,6 @@ def model():
     return m
 
 
-@pytest.fixture
-def model_chunked():
-    m = Model(chunk="auto")
-
-    x = m.add_variables(name="x")
-    y = m.add_variables(name="y")
-
-    m.add_constraints(2 * x + 6 * y, ">=", 10)
-    m.add_constraints(4 * x + 2 * y, ">=", 3)
-
-    m.add_objective(2 * y + x)
-    return m
-
-
 @pytest.fixture
 def milp_model():
     m = Model()
@@ -95,13 +81,6 @@ def test_default_setting(model, solver):
     assert np.isclose(model.objective_value, 3.3)
 
 
-@pytest.mark.parametrize("solver", available_solvers)
-def test_default_settings_chunked(model_chunked, solver):
-    status, condition = model_chunked.solve(solver)
-    assert status == "ok"
-    assert np.isclose(model_chunked.objective_value, 3.3)
-
-
 @pytest.mark.parametrize("solver", available_solvers)
 def test_set_files(tmp_path, model, solver):
     status, condition = model.solve(
diff --git a/test/test_variable.py b/test/test_variable.py
index 0494ef8..b742436 100644
--- a/test/test_variable.py
+++ b/test/test_variable.py
@@ -17,12 +17,20 @@ from linopy import LinearExpression, Model
 
 def test_variable_repr():
     m = Model()
+    m.variables.__repr__()
+
     x = m.add_variables()
     x.__repr__()
     x._repr_html_()
 
     m.variables.__repr__()
 
+    y = m.add_variables(coords=[pd.Index([1, 2, 3], name="time")], name="y")
+    y.__repr__()
+    y._repr_html_()
+
+    m.variables.__repr__()
+
 
 def test_variable_bound_accessor():
     m = Model()