Support for sparse xarrays

[staadecker]

First of, great work with the library so far!! This is significantly better than what currently exists out there.

Many energy systems models can be extremely sparse. For example, in one recent case, a linear expression is composed of the sum of only 0.1% of the coordinates of an xarray. Not supporting sparsity hence would increase our model size by ~900x. This effectively prohibits us from using Linopy as is.

Suggested solution

Ideally linopy could support xarrays based on the sparse library. I've tried the library but it doesn't work with linopy particularly when trying to combine sparse and not sparse xarrays.

Rejected solutions

1. I tried using MultiIndices which are natively supported in xarray however broadcasting doesn't work when combining a sub-dimension of a multi-index. In general, my gut tells me this is the wrong direction to take for the library.

2. I'm aware of the mask parameter for creating constraints and variables however this a) still requires a large memory allocation and b) doesn't help during the computation of linear expressions (it only helps at the variable/constraint creation stage).

[staadecker]

(I might be able to submit a pull request for the proposed solution but I first wanted to hear what you thought @FabianHofmann)

[FabianHofmann]

Hey @staadecker, nice to hear from you and that you like the package :) I would very much like to see sparse object support in linopy and I really cannot say how invasive the required changes would be. But you are most welcome to give it a try! PRs are always welcome.

[staadecker]

Update on investigation
I took a look today at adding support for sparse however the changes seem to be quite invasive for two reasons...

1. xarray is still developing support for sparse and simple operations like creating a sparse DataArray without first creating a DataFrame don't exist afaik.
2. sparse DataArrays don't mix well with non-sparse DataArrays so a lot of the e.g. concat operations are broken. I think this means a lot of the code would need to be changed to make everything sparse.

General thoughts + question
My feeling is that instead of using xarray one could simply use a Pandas dataframe with a multi-index to do what Linopy does. Instead of broadcasting, one simply does a join on the index columns that are in common. I'm curious, what were the initial reasons to use xarray?

For more context, the lack of sparsity means that a 200 MB data table becomes a 200 GB DataArray which results in an error being thrown due to lack of memory.

[koen-vg]

Just to add another perspective to the discussion, I believe the non-sparse nature of xarray is also impacting the memory footprint of solving pypsa-eur networks. Here is an example of memory usage profile during while solving a medium/high-resolution (35 clusters, 3000 time steps, 4th step in myopic foresight pathway optimisation) sector-coupled pypsa-eur network:

This is after applying an ad-hoc memory optimisation to drop the model from memory during solving and reading it back in from a netcdf file afterwards (456052e) as suggested in #219.

During the above optimisation, Gurobi reported using about 16GB of memory, so the mentioned optimisation still leaves a memory overhead of about 9-10GB.

But the maximum memory usage spikes to above 100GB during model preparation. And this was only measured at 30 second intervals, so it's possible that the real maximum was even higher.

For me this has been especially bothersome since I'm working on a SLURM cluster where jobs have to be allocated a certain amount of memory, and will be killed if they exceed it. In an ideal world (in order to run as many jobs in parallel as possible) I'd allocate ~16GB to the above optimisation, but more than 100GB is actually needed... It's possible to group optimisations in a single job and start them in a staggered manner, but this adds quite a bit of complexity.

[fneum]

Hmm, thanks @koen-vg for pointing this out. Memory spikes in the building phase seem to be a recurring issue, but they should be avoidable. New versions of xarray sometimes cause memory leaks.

Last September, after #161, I had a memory log for PyPSA-Eur sector-coupled, 100 nodes, 4-hourly resolution that looked like this.

I'll check how this particular test scenario looks for me with current PyPSA, linopy, xarray versions.

[koen-vg]

Thanks for the quick response! Reassuring to know that a low memory footprint is actually possible. I might try with a few different xarray versions then; the above was using version 2024.2.0. Do you know what version you used when things were working? (Or even just the whole software environment)

[FabianHofmann]

@koen-vg that's quite weird, could you mention which linopy you were using?

[FabianHofmann]

could it have something to do with the IO to the temporary netcdf? did you try it out without it?

[koen-vg]

It was based on linopy 0.3.8. When testing out 456052e it seemed like it reduced the overall memory footprint, but I'm about to do some more detailed testing (master branch linopy; xarray version) to see if I can pin the problem down to anything in particular.

Hopefully I'll soon be able to either point at a faulty xarray/linopy version or just give you a pypsa network which consistently leads to a big memory spike.

[FabianHofmann]

thanks for your deep-dive @koen-vg, this is quite important

[koen-vg]

Some very preliminary investigation seems to indicate that perhaps myopic foresight optimisation is really the culprit here?

This is the result from even a fairly low-resolution network:

I've also seen much more extreme spikes, but only for "later" networks, not overnight optimisations.

Consider this as work in progress. I'm just replicating this with upstream pypsa-eur and linopy versions so I'll be able to provide a complete minimum working example.

[koen-vg]

Regardless of memory spikes, I will just note that, to my best understanding, there is no reason in principle why the optimisation at 2050 should take more time and memory than the one at 2025; the reason is presumably that the later network contains many copies of non-extendable generators, links etc. with different build-years and capacities but otherwise identical function and capacity factors. For solving purposes, those should probably be aggregated and then disaggregated again (since information about how much capacity is built in each year is needed for phase-out purposes). But that's really a separate issue to memory spikes and should be discussed over at pypsa-eur.

[koen-vg]

Last thing before logging off for the weekend:

Here is the memory spike replicated with upstream pypsa-eur and latest version of linopy:

This is using pypsa-eur at commit 95805a8d and a freshly created conda env, resulting a linopy 3.8.0 and xarray 2024.2.0.

The configuration file used to produce this example is based in the myopic test config, as follows:

run:
  name: "mem-test"
  disable_progressbar: true
  shared_resources: false
  shared_cutouts: true

enable:
  retrieve_cutout: false

foresight: myopic

scenario:
  ll:
  - v1.5
  clusters:
  - 70
  sector_opts:
  - ""
  planning_horizons:
  - 2020
  - 2025
  - 2030
  - 2035
  - 2040
  - 2045
  - 2050

countries: ["DE", "NL", "BE"]

clustering:
  temporal:
    resolution_sector: "336H"

sector:
  central_heat_vent: true

electricity:
  extendable_carriers:
    Generator: [OCGT]
    StorageUnit: [battery]
    Store: [H2]
    Link: [H2 pipeline]

  renewable_carriers: [solar, onwind, offwind-ac, offwind-dc]

renewable:
  solar:
    cutout: "europe-2013-era5"

industry:
  St_primary_fraction:
    2030: 0.6
    2040: 0.5
    2050: 0.4

solving:
  solver:
    name: "gurobi"
  mem: 4000

Some observations:

• It's easiest to reproduce this at relatively high spatial resolution and low temporal resolution.
• On the other hand, anecdotally speaking the memory does also spike at higher temporal resolution. I'll run some optimisations over the weekend to confirm this.
• The memory during model preparation doesn't exceed the solver memory on the first time horizon, but in this example it already does at the second time horizon (2025). Then the problem grows more severe in subsequent horizons. However, the memory usage does "level off" at some point, presumably as components past their lifetime are phased out.

A little bit of debugging

I tried using ulimit to figure out where the problem is exactly. The results are mixed; you can see where too much memory is allocated, but this is only virtual memory, whereas is the RSS that counts. It's not clear to me if the two are correlated enough. But for example you get something like this:

> (ulimit -v $((3000*1024)); snakemake --use-conda --configfile config/config.mem-test.yaml -j6 -- solve_sector_networks)
[...]
localrule solve_sector_network_myopic:
    input: results/mem-test/prenetworks-brownfield/elec_s_50_lv1.5___2030.nc, resources/mem-test/costs_2030.csv
    output: results/mem-test/postnetworks/elec_s_50_lv1.5___2030.nc, results/mem-test/configs/config.elec_s_50_lv1.5___2030.yaml
    log: results/mem-test/logs/elec_s_50_lv1.5___2030_solver.log, results/mem-test/logs/elec_s_50_lv1.5___2030_memory.log, results/mem-test/logs/elec_s_50_lv1.5___2030_python.log
    jobid: 74
    benchmark: results/mem-test/benchmarks/solve_sector_network/elec_s_50_lv1.5___2030
    reason: Input files updated by another job: results/mem-test/prenetworks-brownfield/elec_s_50_lv1.5___2030.nc
    wildcards: simpl=, clusters=50, ll=v1.5, opts=, sector_opts=, planning_horizons=2030
    threads: 4
    resources: tmpdir=/tmp, mem_mb=30000, mem_mib=28611, runtime=360

Your conda installation is not configured to use strict channel priorities. This is however crucial for having robust and correct environments (for details, see https://conda-forge.org/docs/user/tipsandtricks.html). Please consider to configure strict priorities by executing 'conda config --set channel_priority strict'.
INFO:pypsa.io:Imported network elec_s_50_lv1.5___2030.nc has buses, carriers, generators, global_constraints, lines, links, loads, shapes, stores
WARNING:__main__:Existing capacities larger than technical potential for Index(['BE0 0 onwind-2030', 'BE0 2 offwind-ac-2030', 'BE0 2 offwind-dc-2030',
       'BE0 2 onwind-2030', 'DE0 10 offwind-dc-2030', 'DE0 16 offwind-ac-2030',
       'DE0 20 onwind-2030', 'DE0 31 offwind-ac-2030',
       'DE0 32 offwind-ac-2030', 'DE0 32 offwind-dc-2030',
       'DE0 32 onwind-2030', 'DE0 5 offwind-ac-2030'],
      dtype='object', name='Generator'),                        adjust technical potential to existing capacities
/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/shadow/tmp3yjpi434/.snakemake/scripts/tmpx1p14x3_.solve_network.py:156: FutureWarning: A value is trying to be set on a copy of a DataFrame or Series through chained assignment using an inplace method.
The behavior will change in pandas 3.0. This inplace method will never work because the intermediate object on which we are setting values always behaves as a copy.

For example, when doing 'df[col].method(value, inplace=True)', try using 'df.method({col: value}, inplace=True)' or df[col] = df[col].method(value) instead, to perform the operation inplace on the original object.


  n.generators.p_nom_max.clip(lower=0, inplace=True)
WARNING:pypsa.components:The following buses have no attached components, which can break the lopf:
{'EU'}
WARNING:pypsa.components:The following buses have no attached components, which can break the lopf:
{'EU'}
INFO:linopy.model: Solve problem using Gurobi solver
INFO:linopy.model:Solver logs written to `results/mem-test/logs/elec_s_50_lv1.5___2030_solver.log`.
INFO:linopy.model:Solver options:
 - threads: 4
 - method: 2
 - crossover: 0
 - BarConvTol: 1e-06
 - Seed: 123
 - AggFill: 0
 - PreDual: 0
 - GURO_PAR_BARDENSETHRESH: 200
Set parameter Username
Academic license - for non-commercial use only - expires 2024-12-04
INFO:linopy.io:Writing objective.
Writing constraints.:  64%|████████████████████████▉              | 23/36 [00:01<00:00, 20.66it/s]
ERROR:root:Uncaught exception
Traceback (most recent call last):
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/shadow/tmp3yjpi434/.snakemake/scripts/tmpx1p14x3_.solve_network.py", line 982, in <module>
    n = solve_network(
        ^^^^^^^^^^^^^^
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/shadow/tmp3yjpi434/.snakemake/scripts/tmpx1p14x3_.solve_network.py", line 924, in solve_network
    status, condition = n.optimize(**kwargs)
                        ^^^^^^^^^^^^^^^^^^^^
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_/lib/python3.11/site-packages/pypsa/optimization/optimize.py", line 604, in __call__
    return optimize(self._parent, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_/lib/python3.11/site-packages/pypsa/optimization/optimize.py", line 584, in optimize
    status, condition = m.solve(solver_name=solver_name, **solver_options, **kwargs)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_/lib/python3.11/site-packages/linopy/model.py", line 1060, in solve
    result = func(
             ^^^^^
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_/lib/python3.11/site-packages/linopy/solvers.py", line 615, in run_gurobi
    problem_fn = model.to_file(problem_fn)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_/lib/python3.11/site-packages/linopy/io.py", line 278, in to_file
    constraints_to_file(m, f, log=log, batch_size=batch_size)
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_/lib/python3.11/site-packages/linopy/io.py", line 124, in constraints_to_file
    df = m.constraints[name].flat
         ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_/lib/python3.11/site-packages/linopy/constraints.py", line 477, in flat
    df = to_dataframe(ds, mask_func=mask_func)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_/lib/python3.11/site-packages/linopy/common.py", line 242, in to_dataframe
    data = {k: np.ravel(v) for k, v in data.items()}
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_/lib/python3.11/site-packages/linopy/common.py", line 242, in <dictcomp>
    data = {k: np.ravel(v) for k, v in data.items()}
               ^^^^^^^^^^^
  File "/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_/lib/python3.11/site-packages/numpy/core/fromnumeric.py", line 1874, in ravel
    return asanyarray(a).ravel(order=order)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
numpy.core._exceptions._ArrayMemoryError: Unable to allocate 254. MiB for an array with shape (33353694,) and data type int64
RuleException:
CalledProcessError in file /home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/rules/solve_myopic.smk, line 147:
Command 'source /home/koen/.local/opt/mambaforge/bin/activate '/home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_'; set -euo pipefail;  python /home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/shadow/tmp3yjpi434/.snakemake/scripts/tmpx1p14x3_.solve_network.py' returned non-zero exit status 1.
[Fri Apr 19 16:44:15 2024]
Error in rule solve_sector_network_myopic:
    jobid: 0
    input: results/mem-test/prenetworks-brownfield/elec_s_50_lv1.5___2030.nc, resources/mem-test/costs_2030.csv
    output: results/mem-test/postnetworks/elec_s_50_lv1.5___2030.nc, results/mem-test/configs/config.elec_s_50_lv1.5___2030.yaml
    log: results/mem-test/logs/elec_s_50_lv1.5___2030_solver.log, results/mem-test/logs/elec_s_50_lv1.5___2030_memory.log, results/mem-test/logs/elec_s_50_lv1.5___2030_python.log (check log file(s) for error details)
    conda-env: /home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_

Exiting because a job execution failed. Look above for error message
WorkflowError:
At least one job did not complete successfully.
[Fri Apr 19 16:44:15 2024]
Error in rule solve_sector_network_myopic:
    jobid: 74
    input: results/mem-test/prenetworks-brownfield/elec_s_50_lv1.5___2030.nc, resources/mem-test/costs_2030.csv
    output: results/mem-test/postnetworks/elec_s_50_lv1.5___2030.nc, results/mem-test/configs/config.elec_s_50_lv1.5___2030.yaml
    log: results/mem-test/logs/elec_s_50_lv1.5___2030_solver.log, results/mem-test/logs/elec_s_50_lv1.5___2030_memory.log, results/mem-test/logs/elec_s_50_lv1.5___2030_python.log (check log file(s) for error details)
    conda-env: /home/koen/Dokument/UiT/Research/Modelling/pypsa-eur/.snakemake/conda/0d18979ece638ab905bf948360770a8a_

Shutting down, this might take some time.
Exiting because a job execution failed. Look above for error message
Complete log: .snakemake/log/2024-04-19T164322.590433.snakemake.log
WorkflowError:
At least one job did not complete successfully.

[koen-vg]

I'm not sure exactly how much time I'll have to look into this the coming week or two, so if anyone else has the time/opportunity, feel free to dig in.

Just some quick (almost administrative) notes:

• One issue is the memory spike in the above set-up which seems (at least on the surface) to be a linopy issue. Whether it really has to do with sparsity or not I'm not actually 100% sure of, so I could also set up a separate issue for it in order to keep things organised.
• A rather separate issue is the substantial overall memory- and time-impact of myopic foresight in pypsa-eur. This is obviously pypsa-eur-related and not a linopy problem, and I'll see if I can open an issue about it there. On the other hand, it might be that streamlining myopic foresight optimisations in pypsa-eur (i.e. doing some aggregation of components) would also eliminate the memory spike observed here.

[FabianHofmann]

Hey @koen-vg, thanks again for the update. I will try to catch up next week (this week is too busy)