Make rules/functions modular and reusable

[FabianHofmann]

When planning on reusing of parts of pypsa-eur in a modular fashion, user quickly face limitations. The following proposal serve as a starting point for a discussion how to make the workflow more "distributable" allowing to re-use rules and functions inside and outside of forks.

CORE IDEA: For each script create one high-level function which is runs independent of the snakemake object, has a clear signature and is importable.

Take build_renewable_profiles as an example. This could be easily restructured as:

# build_renewable_profiles.py  (same file)
def build_renewable_profiles(cutout, availability_matrix, params, client=None) -> xr.Dataset:
    ...  # hoisted out of the __main__ block; returns the Dataset
    return ds

if __name__ == "__main__":
    if "snakemake" not in globals():
        from scripts._helpers import mock_snakemake
        snakemake = mock_snakemake("build_renewable_profiles", ...)
    ds = build_renewable_profiles(cutout, avail, dict(snakemake.params), client=...)
    ds.to_netcdf(snakemake.output.profile)

Downstream then does from scripts.build_renewable_profiles import build_renewable_profiles.

In order to make this consistent we need to

• not leak the snakemake object into the function. The function must take real arguments (cutout, matrix, a params dict/dataclass, plain paths) so it is callable without a internal dependency on the snakemake object.
• define proper return values and not preempt the writing of the output.**. Keep dependency on relative paths low.

Another benefit is that pure functions are unit-testable without running Snakemake.

I iterated with Claude on this on details for time series, see below

Details

## Tiering of the time-series builders

Scan signals (counts of matching lines per script):

Script	atlite	dask	ext. download	xarray	Tier
build_renewable_profiles	13	7	–	8	A
build_hydro_profile	9	–	–	1	A
build_temperature_profiles	8	4	–	4	A
build_solar_thermal_profiles	7	3	–	3	A
build_daily_heat_demand	7	3	–	3	A
build_line_rating	7	9	–	5	A
build_country_runoff	4	–	–	–	A
build_cop_profiles (pkg)	✓	–	–	✓	A (done)
build_hourly_heat_demand	–	–	–	4	B
build_transport_demand	–	–	–	2	B
build_electricity_demand_base	–	–	–	5	B
build_electricity_demand	–	–	ENTSO-E	–	C
build_mobility_profiles	–	–	✓	–	C
build_monthly_prices	–	–	✓	–	C
build_co2_prices	–	–	✓	–	C
build_shipping_demand	–	–	✓	–	C

Tier A — cutout/atlite transforms → best fit

All share one signature: (cutout, clustered_regions/layout, params) → xarray.Dataset.
Deterministic functions of the cutout, no runtime network calls.

One design decision: the dask client. build_line_rating and
build_renewable_profiles construct and shutdown() the client inside the script.
Lifting as-is would couple every importer to a dask-cluster lifecycle. Make it an
injected, optional argument (client: Client | None = None) so the caller owns
parallelism.

Tier B — pure reshapers → trivial fit

No atlite, no download — dataframe/xarray reshaping only (daily→hourly,
profiles→demand, spatial distribution). Nearly pure already; only the
snakemake.input/output lines need peeling off into the shell.

Tier C — retrieval-coupled → fit only after splitting retrieve from transform

The download is half the script and is the part that is non-deterministic, needs API
keys, and must not live in an importable function. Move the download behind a thin
retrieve_* adapter; the parse→tidy→reindex-to-snapshots half becomes the importable
function. Same retrieve/transform discipline build_cop_profiles already demonstrates.

Proposed next steps (smallest valuable first)

1. Add a high-level function to build_renewable_profiles and build_line_rating
   (in-script), injecting the dask client. Highest value, and they exercise the
   dask-injection decision.
2. Same for Tier B scripts — mechanical, low-risk.
3. For Tier C, split retrieve vs. transform; expose the transform as an in-script
   function, leave retrieval behind retrieve_*.

Optional later consolidation (non-breaking)

Once the in-script functions exist, a curated pypsa_eur/ package can re-export
them for a stable public API without touching call sites:

# pypsa_eur/weather_profiles.py
from scripts.build_renewable_profiles import build_renewable_profiles
from scripts.build_line_rating import build_line_rating
...

This is the half-step the in-script approach leaves open: it fully serves the
soft-fork / branch-off case today (we already vendor the repo and do
from scripts.gb_model._helpers import ...), and the pip-installable-library case
becomes a later re-export layer rather than a prerequisite refactor.

[fneum]

Fully agree! There's probably some bits where we can learn from https://github.com/modelblocks-org (@irm-codebase, @jnnr) -- and potentially even collaborate on data preparation for different energy system modelling frameworks.)

[pworschischek-aggmag]

This will help us alot over at PyPSA-AT! thanks for picking up the topic =)

[brynpickering]

I would even break up snakemake.params and any call to config into the individual key:value pairs that are actually required by the function.

[FabianHofmann]

I would even break up snakemake.params and any call to config into the individual key:value pairs that are actually required by the function.

yes, even better!

[euronion]

This might be obvious, but:

• Do not use globals in the functions, not only the snakemake object, but anything going into the function must pass through the function call
• Do not assume defaults inside the function, either define defaults only in the function signature or do not allow for default values at all

When doing this, it would also be valueable to directly document the metadata of the ingoing/outgoing (complex) datatypes, like xarray or pandas objects (what coordinates/indizes, values, ...).

[irm-codebase]

Thanks for reaching out @fneum. We'd definitely love to collaborate on this!

I think there is a great deal that could be shared across modelling teams in a model-agnostic way. Some of our modules are actually based on PyPSA-Eur methods, or use the same dataset (e.g., the gas grid module is a direct extension, and the CO2Stop module has very similar processing).

Before you look into python modularisation within PyPSA-Eur python scripts (@FabianHofmann's suggestion), I'd consider the following:

• If a bundle of rules is 'generic' enough and wanted outside of the PyPSA-Eur ecosystem, abstracting it into a Snakemake module is worth it. They are quite powerful once you understand their requirements, but getting them right is a bit tricky. We have a template for modules with a standardised I/O interface here (feedback on it is very welcomed!).
• If the tool you are using has very clear I/O interfacing, consider creating snakemake-wrappers for it to replace rules. The snakemake team runs updates/maintenance on these, and they are very useful (e.g., see this generic unzipping utility we contributed recently).

[irm-codebase]

Regarding @FabianHofmann's suggestion: we actually follow similar approaches, with some caveats...

• If code is meant to be re-used, we place it in files that are not directly called by rules to make that clear. _schemas.py and _utils.py here are re-used in almost every rule in the geo-boundaries module.
• For complex I/O metadata, consider using pandera schemas. I rely on them to ensure module inputs and outputs are stable. They even added support for xarray recently!
• Fabian's arrangement is pretty spot-on. I prefer to put snakemake requests in the main() call to avoid creating globals in if __name__ == "__main__". Here is my current boiler-plate:

import sys
from typing import TYPE_CHECKING, Any
import _utils   # <--- re-usable code

if TYPE_CHECKING:
    snakemake: Any  # <--- make mypy happy

# rule code...

def main() -> None:
    # All calls to snakemake live here
    foo = snakemake.input.foo
    
if __name__ == "__main__":
    sys.stderr = open(snakemake.log[0], "w")   # Get nice snakemake logs
    main()

[irm-codebase]

Finally, be aware that Snakemake does not 'deliver' python imports in a smart way.

• Snakemake is built to treat rule-scripts as isolated: you can only import code that lives at the same level of the script that requests it without having to resort to 'hacks'. Jonas Hörsh contributed a possible fix to Snakemake here, but overall sharing code defined within rule-scripts is likely to remain very buggy in the near term.
• Unless you add the imported code as a rule input:, Snakemake will not 'see' changes in the imported scripts, and won't re-run them if you make changes to it.
• Sadly, I am not aware of any good method that makes Snakemake code reusable for outside repos without turning PyPSA-Eur into a Snakemake module or some python-library Frankenstein's monster. This mostly matters for 'offspring' models like @pworschischek-aggmag's PyPSA-AT. That's why we advocate for Snakemake modules and wrappers.