perf: align dask chunks with stored chunks and keep conversions lazy

RELEASE_NOTES.rst

@@ -17,6 +17,28 @@ Upcoming Release
    To use the features already you have to install the ``master`` branch, e.g.
    ``pip install git+https://github.com/pypsa/atlite``.
 
+**Performance**
+
+* Open existing cutouts with ``chunks="auto"`` by default instead of ``{"time": 100}``,
+  also after ``Cutout.prepare``. Dask chunks are now aligned with the chunks stored in
+  the NetCDF file. Before, dask chunks that split the stored chunks made the compressed
+  data decompress many times. This makes all conversions much faster for cutouts with
+  large stored chunks. An explicit ``chunks`` argument is
+  still respected. An explicit ``chunks="auto"`` now also
+  respects the stored chunks. Larger chunks can increase peak memory. To reduce it,
+  pass ``chunks``
+  (https://github.com/PyPSA/atlite/pull/527).
+* Keep ``csp`` lazy and chunked. The efficiency interpolation previously merged all time
+  steps into one chunk and added the solar position as four extra coordinates
+  (https://github.com/PyPSA/atlite/pull/527).
+* Keep ``pv`` and ``irradiation`` lazy for ``trigon_model="other"`` and
+  ``tracking="tilted_horizontal"``. Both computed intermediate results eagerly before.
+  The check for negative diffuse irradiation now only runs at log level ``DEBUG``
+  (https://github.com/PyPSA/atlite/pull/527).
+* Compute ``line_rating`` for all lines in one vectorised dask operation instead of one
+  task per line. Before, each line read and decompressed the cutout data again
+  (https://github.com/PyPSA/atlite/pull/527).
+
 `v0.7.0 <https://github.com/PyPSA/atlite/releases/tag/v0.7.0>`__ (3rd October 2026)
 =======================================================================================
 

atlite/convert.py

@@ -16,10 +16,10 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
-from dask import compute, delayed
 from dask.array import absolute, arccos, cos, maximum, mod, radians, sin, sqrt
 from dask.diagnostics import ProgressBar
 from numpy import pi
+from scipy.interpolate import RegularGridInterpolator
 from scipy.sparse import csr_matrix
 
 from atlite import hydro as hydrom
@@ -1415,8 +1415,18 @@ def convert_csp(ds, installation):
         raise ValueError(f'Unknown CSP technology option "{tech}".')
 
     # Determine solar_position dependend efficiency for each grid cell and time step
-    efficiency = installation["efficiency"].interp(
-        altitude=solar_position["altitude"], azimuth=solar_position["azimuth"]
+    eff = installation["efficiency"].transpose("altitude", "azimuth")
+    interp = RegularGridInterpolator(
+        (eff["altitude"].values, eff["azimuth"].values),
+        eff.values,
+        bounds_error=False,
+    )
+    efficiency = xr.apply_ufunc(
+        lambda alt, az: interp((alt, az)),
+        solar_position["altitude"],
+        solar_position["azimuth"],
+        dask="parallelized",
+        output_dtypes=[eff.dtype],
     )
 
     # Thermal system output
@@ -1732,8 +1742,8 @@ def convert_line_rating(
     -------
     xr.DataArray or numpy.ndarray
         Maximum current per time step in ampere. When *ds* is an
-        ``xr.Dataset`` the result is aggregated across intersecting cells
-        via ``.min("spatial")``.
+        ``xr.Dataset`` with a ``spatial`` dimension, the result is aggregated
+        across intersecting cells via ``.min("spatial")``.
     """
     Ta = ds["temperature"]
     Tfilm = (Ta + Ts) / 2
@@ -1791,7 +1801,9 @@ def convert_line_rating(
     qs = alpha * Q * A * sin(Phi_s)
 
     Imax = sqrt((qc + qr - qs) / R)
-    return Imax.min("spatial") if isinstance(Imax, xr.DataArray) else Imax
+    if isinstance(Imax, xr.DataArray) and "spatial" in Imax.dims:
+        return Imax.min("spatial")
+    return Imax
 
 
 def line_azimuth_degrees(shape: Any) -> float:
@@ -1900,7 +1912,8 @@ def line_rating(
     if not isinstance(shapes, gpd.GeoSeries):
         shapes = gpd.GeoSeries(shapes).rename_axis("dim_0")
 
-    I = cutout.intersectionmatrix(shapes)
+    # CSR format orders the (line, cell) pairs by line
+    I = cutout.intersectionmatrix(shapes).tocsr()
     rows, cols = I.nonzero()
 
     data = cutout.data.stack(spatial=["y", "x"])
@@ -1918,19 +1931,31 @@ def line_rating(
     assert df.notnull().all().all(), "Nan values encountered."
     assert df.columns.equals(pd.Index(["psi", "R", "D", "Ts", "epsilon", "alpha"]))
 
-    dummy = xr.DataArray(np.full(len(data.time), np.nan), coords=(data.time,))
-    res = []
-    for i in range(len(df)):
-        cells_i = cols[rows == i]
-        if cells_i.size:
-            ds = data.isel(spatial=cells_i)
-            res.append(delayed(convert_line_rating)(ds, *df.iloc[i].values))
-        else:
-            res.append(dummy)
-    if show_progress:
-        with ProgressBar(minimum=2):
-            (computed,) = compute(res, **dask_kwargs)
-    else:
-        (computed,) = compute(res, **dask_kwargs)
+    # Compute the rating of each (line, cell) pair
+    pairs = data.isel(spatial=cols).rename(spatial="pair").chunk(pair=-1, time=100)
+    pair_params = {
+        k: xr.DataArray(v.values, dims="pair") for k, v in df.iloc[rows].items()
+    }
+    Imax = convert_line_rating(pairs, **pair_params)
+
+    # Take the minimum over the cells of each line, ignoring NaN like ``.min()``.
+    # The pairs of each line are contiguous and start at ``starts``.
+    lines, starts = np.unique(rows, return_index=True)
+    Imax = xr.apply_ufunc(
+        lambda a: np.fmin.reduceat(a, starts, axis=-1),
+        Imax,
+        input_core_dims=[["pair"]],
+        output_core_dims=[["line"]],
+        dask="parallelized",
+        output_dtypes=[Imax.dtype],
+        dask_gufunc_kwargs={
+            "output_sizes": {"line": len(lines)},
+            "allow_rechunk": False,
+        },
+    )
 
-    return xr.concat(computed, dim=df.index).assign_attrs(units="A")
+    # Lines without intersecting cells get NaN
+    Imax = Imax.assign_coords(line=lines).reindex(line=range(len(df)))
+    dim = df.index.name or "concat_dim"
+    Imax = Imax.rename(line=dim).assign_coords({dim: df.index}).transpose(dim, "time")
+    return maybe_progressbar(Imax, show_progress, **dask_kwargs).assign_attrs(units="A")

atlite/cutout.py

@@ -125,9 +125,11 @@ def __init__(self, path: PathLike, **cutoutparams: Any) -> None:
         dt : str, optional
             Frequency of the time coordinate. The default is 'h'. Valid are all
             pandas offset aliases.
-        chunks : dict
+        chunks : dict or str
             Chunks when opening NetCDF files. For cutout preparation it is recommended
-            to chunk only along the time dimension. Defaults to {'time': 100}
+            to chunk only along the time dimension. Defaults to "auto" for existing
+            cutouts, which aligns the chunks with the chunks stored in the file, and
+            to {'time': 100} for new cutouts.
         data : xr.Dataset
             User provided cutout data. Save the cutout using `Cutout.to_file()`
             afterwards.
@@ -162,7 +164,7 @@ def __init__(self, path: PathLike, **cutoutparams: Any) -> None:
 
         """
         path = Path(path).with_suffix(".nc")
-        chunks = cutoutparams.pop("chunks", {"time": 100})
+        chunks = cutoutparams.pop("chunks", "auto" if path.is_file() else {"time": 100})
         if isinstance(chunks, dict):
             storable_chunks = {f"chunksize_{k}": v for k, v in (chunks or {}).items()}
         else:
@@ -171,8 +173,7 @@ def __init__(self, path: PathLike, **cutoutparams: Any) -> None:
         # Three cases. First, cutout exists -> take the data.
         # Second, data is given -> take it. Third, else -> build a new cutout
         if path.is_file():
-            data = xr.open_dataset(str(path))
-            data = data.chunk(chunks)
+            data = xr.open_dataset(str(path), chunks=chunks)
             data.attrs.update(storable_chunks)
             if cutoutparams:
                 warn(

atlite/data.py

@@ -324,6 +324,6 @@ def cutout_prepare(
             cutout.path.unlink()
         Path(tmp).rename(cutout.path)
 
-        cutout.data = xr.open_dataset(cutout.path, chunks=cutout.chunks)
+        cutout.data = xr.open_dataset(cutout.path, chunks="auto")
 
     return cutout

atlite/pv/irradiation.py

@@ -159,10 +159,10 @@ def TiltedDiffuseIrrad(
     ) * diffuse
 
     if (
-        logger.isEnabledFor(logging.WARNING)
+        logger.isEnabledFor(logging.DEBUG)
         and ((diffuse_t < 0.0) & (sinaltitude > sin(radians(1.0)))).any()
     ):
-        logger.warning("diffuse_t exhibits negative values above altitude threshold.")
+        logger.debug("diffuse_t exhibits negative values above altitude threshold.")
 
     # fixup: clip all negative values (unclear why it gets negative)
     # note: REatlas does not do the fixup

atlite/pv/orientation.py

@@ -304,12 +304,12 @@ def SurfaceOrientation(
         azimuth_difference = xr.where(
             azimuth_difference < -pi, 2 * pi + azimuth_difference, azimuth_difference
         )
-        rotation = np.where(
+        rotation = xr.where(
             logical_and(rotation < 0, azimuth_difference > 0),
             rotation + pi,
             rotation,
         )
-        rotation = np.where(
+        rotation = xr.where(
             logical_and(rotation > 0, azimuth_difference < 0),
             rotation - pi,
             rotation,

test/test_dynamic_line_rating.py

@@ -9,11 +9,14 @@
 @author: fabian
 """
 
+import geopandas as gpd
 import numpy as np
 import pandas as pd
 import pytest
+import xarray as xr
 from shapely.geometry import LineString, Point
 
+from atlite import Cutout
 from atlite.convert import convert_line_rating, line_azimuth_degrees
 
 
@@ -194,3 +197,62 @@ def test_line_azimuth_degrees(start, end, expected):
     """`line_azimuth_degrees` returns degrees consistent with `convert_line_rating`'s `psi`."""
     shape = LineString([Point(*start), Point(*end)])
     assert np.isclose(line_azimuth_degrees(shape), expected)
+
+
+@pytest.fixture
+def cutout(tmp_path):
+    rng = np.random.default_rng(0)
+    coords = {
+        "time": pd.date_range("2013-06-01", periods=24, freq="h"),
+        "y": [50.0, 50.25, 50.5],
+        "x": [5.0, 5.25, 5.5, 5.75],
+    }
+
+    def var(low, high, dims=("time", "y", "x")):
+        return (dims, rng.uniform(low, high, [len(coords[d]) for d in dims]))
+
+    ds = xr.Dataset(
+        {
+            "temperature": var(270, 300),
+            "wnd100m": var(0, 15),
+            "wnd_azimuth": var(0, 2 * np.pi),
+            "influx_direct": var(0, 800),
+            "solar_altitude": var(0, 1.2),
+            "solar_azimuth": var(0, 2 * np.pi),
+            "height": var(0, 500, ("y", "x")),
+        },
+        coords=coords,
+        attrs={"module": "era5"},
+    )
+    ds["influx_direct"][:, 0, 0] = np.nan
+    ds.to_netcdf(tmp_path / "cutout.nc")
+    return Cutout(tmp_path / "cutout.nc")
+
+
+def test_line_rating_is_minimum_over_intersected_cells(cutout):
+    shapes = gpd.GeoSeries([
+        LineString([(5.0, 50.0), (5.5, 50.5)]),
+        LineString([(5.75, 50.25), (5.25, 50.25)]),
+        LineString([(9.0, 40.0), (9.5, 40.0)]),
+    ])
+    R = pd.Series([3e-5, 4e-5, 5e-5])
+    res = cutout.line_rating(shapes, R)
+
+    data = cutout.data.stack(spatial=["y", "x"])
+    cells = cutout.intersectionmatrix(shapes).tocsr()
+    for i, shape in enumerate(shapes):
+        if not cells[i].nnz:
+            assert res[i].isnull().all()
+            continue
+        psi = line_azimuth_degrees(shape)
+        psi = psi if psi >= 0 else psi + 180
+        ds = data.isel(spatial=cells[i].indices)
+        expected = convert_line_rating(ds, psi, R[i])
+        np.testing.assert_allclose(res[i], expected)
+
+
+def test_line_rating_without_intersections_is_nan(cutout):
+    shapes = gpd.GeoSeries([LineString([(9.0, 40.0), (9.5, 40.0)])])
+    res = cutout.line_rating(shapes, 3e-5)
+    assert res.shape == (1, 24)
+    assert res.isnull().all()

test/test_performance.py

@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: Contributors to atlite <https://github.com/pypsa/atlite>
+#
+# SPDX-License-Identifier: MIT
+
+"""Tests that conversions stay lazy and cutouts read data efficiently."""
+
+import dask
+import numpy as np
+import pandas as pd
+import pytest
+import xarray as xr
+
+from atlite import Cutout
+from atlite.convert import convert_csp, convert_pv
+from atlite.pv.orientation import get_orientation
+from atlite.resource import get_cspinstallationconfig, get_solarpanelconfig
+
+
+def _forbid_compute(*args, **kwargs):
+    raise RuntimeError("Unexpected eager computation.")
+
+
+@pytest.fixture
+def ds():
+    rng = np.random.default_rng(0)
+    time = pd.date_range("2013-06-01", periods=48, freq="h")
+    coords = {"time": time, "y": [50.0, 51.0], "x": [5.0, 6.0, 7.0]}
+    shape = (len(time), 2, 3)
+
+    def var(low, high):
+        return (("time", "y", "x"), rng.uniform(low, high, shape))
+
+    ds = xr.Dataset(
+        {
+            "influx_toa": var(800, 1000),
+            "influx_direct": var(0, 400),
+            "influx_diffuse": var(0, 200),
+            "albedo": var(0, 0.3),
+            "temperature": var(270, 300),
+            "solar_altitude": var(-0.5, 1.2),
+            "solar_azimuth": var(0, 2 * np.pi),
+        },
+        coords=coords,
+    )
+    ds = ds.assign_coords(lon=("x", coords["x"]), lat=("y", coords["y"]))
+    return ds.chunk({"time": 24})
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"tracking": None, "trigon_model": "other"},
+        {"tracking": "tilted_horizontal"},
+    ],
+)
+def test_convert_pv_is_lazy(ds, kwargs):
+    panel = dict(get_solarpanelconfig("CSi"))
+    orientation = get_orientation({"slope": 30.0, "azimuth": 180.0})
+    with dask.config.set(scheduler=_forbid_compute):
+        da = convert_pv(ds, panel, orientation, **kwargs)
+    assert isinstance(da.data, dask.array.Array)
+
+
+def test_convert_csp_keeps_chunks_and_coords(ds):
+    installation = get_cspinstallationconfig("SAM_solar_tower")
+    da = convert_csp(ds, installation)
+    assert da.chunks == ds["influx_direct"].chunks
+    assert set(da.coords) <= set(ds.coords)
+
+
+def test_cutout_chunks_align_with_file_chunks(tmp_path):
+    path = tmp_path / "cutout.nc"
+    time = pd.date_range("2013-01-01", periods=300, freq="h")
+    ds = xr.Dataset(
+        {"temperature": (("time", "y", "x"), np.zeros((300, 2, 3)))},
+        coords={"time": time, "y": [50.0, 51.0], "x": [5.0, 6.0, 7.0]},
+        attrs={"module": "era5"},
+    )
+    encoding = {"temperature": {"chunksizes": (150, 2, 3), "zlib": True}}
+    ds.to_netcdf(path, encoding=encoding)
+
+    chunks = Cutout(path).data["temperature"].chunks
+    assert all(c % 150 == 0 for c in chunks[0])