Repository: PyPSA/linopy
Pull request #9 (2021-11-21 12:29:41)

Io speedup with unicode dtype


Patch:
diff --git a/doc/release_notes.rst b/doc/release_notes.rst
index 95b5b7d..e55b340 100644
--- a/doc/release_notes.rst
+++ b/doc/release_notes.rst
@@ -4,6 +4,11 @@ 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/io.py b/linopy/io.py
index 6da31c9..e8e0810 100644
--- a/linopy/io.py
+++ b/linopy/io.py
@@ -4,41 +4,46 @@
 import logging
 import os
 import time
-from functools import partial, reduce
 
 import numpy as np
 import xarray as xr
-from xarray import apply_ufunc
+from numpy import dtype
+from xarray import DataArray, apply_ufunc, concat
 
 logger = logging.getLogger(__name__)
 
 
-ufunc_kwargs = dict(dask="parallelized", vectorize=True, output_dtypes=[object])
+ufunc_kwargs = dict(vectorize=True)
+concat_dim = "_concat_dim"
+concat_kwargs = dict(dim=concat_dim, coords="minimal")
+
 
 # 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), **ufunc_kwargs)
+    return apply_ufunc(
+        lambda f: "%+f" % f,
+        da.fillna(0).compute(),
+        **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), **ufunc_kwargs)
+    return xr.apply_ufunc(lambda d: "%d" % d, da.fillna(0).compute(), **ufunc_kwargs)
 
 
-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 array_to_file(array, fn):
+    """Elementwise writing out string values to a file."""
+    write = np.frompyfunc(fn.write, 1, 1)
 
+    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(
-        lambda x: fn.write(x),
+        func,
         array,
         dask="parallelized",
-        vectorize=True,
         output_dtypes=[int],
     )
 
@@ -49,72 +54,85 @@ 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)
-    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()
+    objective = concat(objective, **concat_kwargs).where(nonnans, "")
+    objective = objective.transpose(..., concat_dim)
+    array_to_file(objective.compute(), f)
 
 
 def constraints_to_file(m, f):
     """Write out the constraints of a model to a lp file."""
     f.write("\n\ns.t.\n\n")
-    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()
+    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)
 
 
 def bounds_to_file(m, f):
     """Write out variables of a model to a lp file."""
     f.write("\nbounds\n")
-    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()
+    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)
 
 
 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")),
+        ]
 
-    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()
+        binaries = concat(binaries, **concat_kwargs).where(labels != -1, "")
+        binaries = binaries.transpose(..., concat_dim)
+        array_to_file(binaries.compute(), fn=f)
 
 
 def to_file(m, fn):


Test patch:
diff --git a/test/test_io.py b/test/test_io.py
index fe404aa..d691b46 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 == object
+    assert da.dtype.type is np.str_
 
 
 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 == object
+    assert da.dtype.type is np.str_
 
 
 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 == object
+    assert da.dtype.type is np.str_
 
 
 def test_to_netcdf(tmp_path):
diff --git a/test/test_optimization.py b/test/test_optimization.py
index bc0abc7..0bba5b4 100644
--- a/test/test_optimization.py
+++ b/test/test_optimization.py
@@ -28,6 +28,20 @@ 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()
@@ -81,6 +95,13 @@ 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(