Repository: PyPSA/linopy
Pull request #13 (2021-12-10 17:51:02)

Io str arrays


Patch:
diff --git a/doc/release_notes.rst b/doc/release_notes.rst
index 772128d..073945d 100644
--- a/doc/release_notes.rst
+++ b/doc/release_notes.rst
@@ -6,7 +6,7 @@ Upcoming Release
 
 * LinearExpression has a new function `densify_terms` which reduces the `_term` axis to a minimal length while containing all non-zero coefficient values.
 * When summing over one or multiple axes in a LinearExpression, terms with coefficient of zeros can now be dropped automatically.
-* The export of LP files was restructured and is now using Unicode dtype under the hook to ensure performant export of long constraints.
+* The export of LP files was restructured and is flat arrays under the hook to ensure performant export of long constraints.
 * Dimensions of masks passed to `add_variables` and `add_constraints` now have to be a subset of the resulting labels dimensions.
 * A new high-level function `merge` was added to concatenate multiple linear expressions.
 * The ``Variable.where`` function now has -1 as default fill value.
diff --git a/linopy/constraints.py b/linopy/constraints.py
index feb3061..89795a0 100644
--- a/linopy/constraints.py
+++ b/linopy/constraints.py
@@ -281,7 +281,8 @@ class Constraints:
             The default is "labels".
         filter_missings : bool, optional
             Filter out values where `labels` data is -1. If broadcast is `vars`
-            also values where `vars` is -1 are filtered. The default is False.
+            also values where `vars` is -1 are filtered. When enabled, the
+            data is load into memory. The default is False.
 
 
         Yields
@@ -304,15 +305,15 @@ class Constraints:
             if values.chunks is not None:
                 broadcasted = broadcasted.chunk(values.chunks)
 
-            flat = broadcasted.data.ravel()
             if filter_missings:
-                mask = values.compute().data.ravel() != -1
+                flat = np.ravel(broadcasted)
+                mask = np.ravel(values) != -1
                 if broadcast_like != "labels":
-                    labels = (
-                        self.labels[name].broadcast_like(values).compute().data.ravel()
-                    )
+                    labels = np.ravel(self.labels[name].broadcast_like(values))
                     mask &= labels != -1
                 flat = flat[mask]
+            else:
+                flat = broadcasted.data.ravel()
             yield flat
 
     def ravel(self, key, broadcast_like="labels", filter_missings=False, compute=True):
diff --git a/linopy/io.py b/linopy/io.py
index 17655a1..ac878dd 100644
--- a/linopy/io.py
+++ b/linopy/io.py
@@ -6,13 +6,12 @@ import os
 import shutil
 import time
 from pathlib import Path
-from tempfile import NamedTemporaryFile, TemporaryDirectory
+from tempfile import TemporaryDirectory
 
 import numpy as np
 import xarray as xr
-from numpy import dtype
+from numpy import asarray, concatenate
 from tqdm import tqdm
-from xarray import DataArray, apply_ufunc, concat
 
 logger = logging.getLogger(__name__)
 
@@ -23,130 +22,103 @@ 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).compute(), **ufunc_kwargs)
+def as_str(arr):
+    "Convert numpy array to str typed array."
+    return arr.astype(str).astype(object)
 
 
-def to_int_str(da):
-    """Convert a int array to a string array."""
-    return xr.apply_ufunc(lambda d: "%d" % d, da.fillna(0).compute(), **ufunc_kwargs)
+def with_sign(arr):
+    "Convert a numpy array of dtype float to str typed array with +/- signs."
+    return np.where(arr >= 0, "+", "").astype(object) + as_str(arr)
 
 
-def array_to_file(array, fn):
-    """Elementwise writing out string values to a file."""
-    write = np.frompyfunc(fn.write, 1, 1)
+def fill_by(target_shape, where, fill, other=""):
+    "Create array of `target_shape` with values from `fill` where `where` is True."
+    res = np.full(target_shape, other).astype(object)
+    res[where] = fill
+    return res
 
-    def func(data):
-        return write(data.ravel()).reshape(data.shape)
 
-    return xr.apply_ufunc(func, array)
-
-
-def objective_to_file(m, f, log):
+def objective_to_file(m, f, log=False):
     """Write out the objective of a model to a lp file."""
-    f.write("min\nobj:\n")
-    coef = m.objective.coeffs
-    var = m.objective.vars
+    if log:
+        logger.info("Writing objective.")
 
-    logger.info("Writing objective.")
+    f.write("min\n\nobj:\n\n")
+    coeffs = m.objective.coeffs.values
+    vars = m.objective.vars.values
+    nnz = vars != -1
+    coeffs, vars = coeffs[nnz], vars[nnz]
 
-    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)
+    objective = with_sign(coeffs) + " x" + as_str(vars)
+    f.write("\n".join(objective))
 
 
 def constraints_to_file(m, f, log=False):
     """Write out the constraints of a model to a lp file."""
     f.write("\n\ns.t.\n\n")
-
     m.constraints.sanitize_missings()
+    kwargs = dict(broadcast_like="vars", filter_missings=True)
+    vars = m.constraints.iter_ravel("vars", **kwargs)
+    coeffs = m.constraints.iter_ravel("coeffs", **kwargs)
+    labels = m.constraints.iter_ravel("labels", **kwargs)
 
-    iterate = m.constraints.labels.items()
-    if log:
-        iterate = tqdm(iterate, "Writing constraints.")
-
-    for name, labels in iterate:
+    labels_ = m.constraints.iter_ravel("labels", filter_missings=True)
+    sign_ = m.constraints.iter_ravel("sign", filter_missings=True)
+    rhs_ = m.constraints.iter_ravel("rhs", filter_missings=True)
 
-        dim = f"{name}_term"
+    iterate = zip(labels, vars, coeffs, labels_, sign_, rhs_)
+    if log:
+        iterate = tqdm(iterate, "Writing constraints.", len(m.constraints.labels))
 
-        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")),
-        ]
+    for (l, v, c, l_, s_, r_) in iterate:
 
-        nonnans_terms = m.constraints.vars[name] != -1
-        lhs = concat(lhs, **concat_kwargs).where(nonnans_terms, "")
-        lhs = lhs.stack(_=[dim, concat_dim]).drop_vars("_").rename(_=dim)
+        diff_con = l[:-1] != l[1:]
+        new_con_b = concatenate([asarray([True]), diff_con])
+        end_of_con_b = concatenate([diff_con, asarray([True])])
 
-        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")),
-        ]
+        l = fill_by(v.shape, new_con_b, "\nc" + as_str(l_) + ":\n")
+        s = fill_by(v.shape, end_of_con_b, "\n" + as_str(s_) + "\n")
+        r = fill_by(v.shape, end_of_con_b, as_str(r_))
+        constraints = l + with_sign(c) + " x" + as_str(v) + s + r
 
-        da = concat(constraints, dim=dim, coords="minimal")
-        da = da.where(labels != -1, "").transpose(..., dim)
-        array_to_file(da.compute(), fn=f)
+        f.write("\n".join(constraints))
+        f.write("\n")
 
 
 def bounds_to_file(m, f, log=False):
     """Write out variables of a model to a lp file."""
-    f.write("\nbounds\n")
+    f.write("\n\nbounds\n\n")
+    lower = m.non_binaries.iter_ravel("lower", filter_missings=True)
+    upper = m.non_binaries.iter_ravel("upper", filter_missings=True)
+    labels = m.non_binaries.iter_ravel("labels", filter_missings=True)
 
-    iterate = m.non_binaries.labels.items()
+    iterate = zip(lower, labels, upper)
     if log:
-        iterate = tqdm(iterate, "Writing variables.")
-
-    for name, labels in iterate:
-
-        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")),
-        ]
+        iterate = tqdm(iterate, "Writing variables.", len(m.non_binaries.labels))
 
-        bounds = concat(bounds, **concat_kwargs).where(labels != -1, "")
-        bounds = bounds.transpose(..., concat_dim)
-        array_to_file(bounds.compute(), f)
+    for (lo, l, up) in iterate:
+        bounds = with_sign(lo) + " <= x" + as_str(l) + " <= " + with_sign(up)
+        f.write("\n".join(bounds))
+        f.write("\n")
 
 
 def binaries_to_file(m, f, log=False):
     """Write out binaries of a model to a lp file."""
-    f.write("\nbinary\n")
+    f.write("\n\nbinary\n\n")
+    labels = m.binaries.iter_ravel("labels", filter_missings=True)
 
-    iterate = m.binaries.labels.items()
-    if log:
-        iterate = tqdm(iterate, "Writing binaries.")
-
-    for name, labels in iterate:
+    if len(m.variables._binary_variables) == 0:
+        return
 
-        binaries = [
-            DataArray("x").astype(dtype("<U9")),
-            to_int_str(labels),
-            DataArray("\n").astype(dtype("<U9")),
-        ]
+    iterate = labels
+    if log:
+        iterate = tqdm(iterate, "Writing binaries.", len(m.binaries.labels))
 
-        binaries = concat(binaries, **concat_kwargs).where(labels != -1, "")
-        binaries = binaries.transpose(..., concat_dim)
-        array_to_file(binaries.compute(), fn=f)
+    for l in iterate:
+        bounds = "x" + as_str(l)
+        f.write("\n".join(bounds))
+        f.write("\n")
 
 
 def to_file(m, fn):
@@ -156,7 +128,7 @@ def to_file(m, fn):
     if os.path.exists(fn):
         os.remove(fn)  # ensure a clear file
 
-    log = m._xCounter > 1000
+    log = m._xCounter > 10000
 
     with open(fn, mode="w") as f:
 
diff --git a/linopy/variables.py b/linopy/variables.py
index 66db5e8..c9f14c0 100644
--- a/linopy/variables.py
+++ b/linopy/variables.py
@@ -365,8 +365,8 @@ class Variables:
             Key to be iterated over. Optionally pass a dataset which is
             broadcastable to `broadcast_like`.
         filter_missings : bool, optional
-            Filter out values where `broadcast_like` data is -1.
-            The default is False.
+            Filter out values where `broadcast_like` data is -1. When enabled, the
+            data is load into memory. The default is False.
 
 
         Yields
@@ -387,9 +387,11 @@ class Variables:
             if labels.chunks is not None:
                 broadcasted = broadcasted.chunk(labels.chunks)
 
-            flat = broadcasted.data.ravel()
             if filter_missings:
-                flat = flat[labels.data.ravel() != -1]
+                flat = np.ravel(broadcasted)
+                flat = flat[np.ravel(labels) != -1]
+            else:
+                flat = broadcasted.data.ravel()
             yield flat
 
     def ravel(self, key, filter_missings=False, compute=True):


Test patch:
diff --git a/test/test_io.py b/test/test_io.py
index d691b46..dd321c5 100644
--- a/test/test_io.py
+++ b/test/test_io.py
@@ -6,14 +6,12 @@ Created on Thu Mar 18 09:03:35 2021
 @author: fabian
 """
 
-import numpy as np
 import pandas as pd
 import pytest
-import xarray as xr
-from xarray.testing import assert_allclose, assert_equal
+from xarray.testing import assert_equal
 
-from linopy import LinearExpression, Model, available_solvers, read_netcdf
-from linopy.io import to_int_str
+from linopy import Model, available_solvers, read_netcdf
+from linopy.io import as_str
 
 
 def test_str_arrays():
@@ -22,18 +20,18 @@ def test_str_arrays():
     x = m.add_variables(4, pd.Series([8, 10]))
     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_
+    da = as_str(x)
+    assert da.dtype == object
 
 
 def test_str_arrays_chunked():
-    m = Model(chunk=-1)
+    m = Model(chunk="auto")
 
     x = m.add_variables(4, pd.Series([8, 10]))
     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_
+    da = as_str(y).compute()
+    assert da.dtype == object
 
 
 def test_str_arrays_with_nans():
@@ -44,8 +42,8 @@ def test_str_arrays_with_nans():
     y = m.add_variables(0, pd.DataFrame([[1, 2], [3, 4], [5, 6]]), name="y")
     assert m["x"][-1].item() == -1
 
-    da = to_int_str(m["x"])
-    assert da.dtype.type is np.str_
+    da = as_str(m["x"])
+    assert da.dtype == object
 
 
 def test_to_netcdf(tmp_path):