Feat/flexible statistics

INTERFACE.yaml

@@ -23,10 +23,18 @@ pathvars:
       description: |
         Powerplant statistics.
         See the schemas in `workflow/scripts/_schemas.py` for more information.
+    statistics_user:
+      default: "<resources>/user/{shapes}/generation.parquet"
+      description: |
+        Optional. Annual national hydropower generation statistics.
+        Used when "generation_adjustement" is set to "user" in the configuration.
+        See the schemas in `workflow/scripts/_schemas.py` for more information.
   results:
-    statistics:
+    statistics_eia:
       default: "<results>/{shapes}/statistics/generation.parquet"
-      description: National hydropower generation statistics for the requested shapes.
+      description: >
+        National hydropower generation statistics for the requested shapes.
+        Generated using EIA data.
     disaggregated_inflow:
       default: "<results>/{shapes}/disaggregated/inflow_mwh.parquet"
       description: >

README.md

@@ -22,7 +22,7 @@ and the `snakemake` [documentation](https://snakemake.readthedocs.io/en/stable/s
 Data processing steps:
 
 <p align="center">
-  <img src="./figures/rulegraph.png" width="60%">
+  <img src="./figures/rulegraph.png" width="50%">
 </p>
 
 1. A global dataset of hydro basins is created for a requested Pfafstetter level using data from the [HydroBASINS](https://www.hydrosheds.org/products/hydrobasins) dataset.
@@ -32,9 +32,10 @@ Data processing steps:
   <img src="./figures/adjusted_powerplants.png" width="60%">
 </p>
 
-3. User provided shapes, powerplants and configuration are used to construct a data request to the [Copernicus Data Store](https://cds.climate.copernicus.eu/) using the [`atlite` library](https://github.com/PyPSA/atlite).
-4. `atlite` is used to construct inflow timeseries per powerplant.
-5. Inflow timeseries are combined and aggregated to the requested resolution, using national-level statistics from the [EIA](https://www.eia.gov/international/) as a guiding normalisation heuristic for total generation.
+3. If requested, [EIA](https://www.eia.gov/international/) statistics of national hydropower generation are downloaded and harmonised to a standardised schema.
+Users may decide to replace this data with their own, or to adjust it slightly (see the [interfacing](#input--output-structure)).
+4. User provided shapes, powerplants and configuration are used to construct a data request to the [Copernicus Data Store](https://cds.climate.copernicus.eu/) using the [`atlite` library](https://github.com/PyPSA/atlite) to construct inflow timeseries per powerplant.
+5. Inflow timeseries are combined and aggregated to the requested resolution, using the national generation statistics as a guiding normalisation heuristic for total generation.
 6. Per-unit inflow timeseries ($PU_{t,r}$) are produced for each region with available capacity.
 The relation is $InflowMWh_{t,r} = PU_{t,r} \cdot Cap_{r}$, where $Cap_r$ is the capacity per region.
     - For run of river powerplants, the timeseries are capped so they may not exceed available capacity.
@@ -62,7 +63,14 @@ Please consult the configuration [README](./config/README.md) and the [configura
 
 As input, you need a file with the polygons to aggregate into (the 'shapes'), and a file specifying national hydropower plants (either RoR or reservoir). These files should follow the `pandera` schemas specified in [_schemas.py](./workflow/scripts/_schemas.py), and can be created using other Modelblocks modules.
 
-Outputs include aggregated timeseries per hydropower plant type, disaggregated inflow timeseries (per point powerplant), and national generation statistics.
+If [configured](#configuration), the internally created EIA generation statistics file can be replaced or adjusted by the user.
+
+<p align="center">
+  <img src="./figures/generation.drawio.png" width="50%">
+</p>
+
+
+Outputs include aggregated timeseries per hydropower plant type, disaggregated inflow timeseries (per point powerplant), and EIA national generation statistics.
 
 Please consult the [interface file](./INTERFACE.yaml) for more information.
 

config/config.yaml

@@ -1,17 +1,17 @@
-# A minimal example of how to configure this module
 pfafstetter_level: "06"
+crs:
+  projected: "epsg:3857"
+  geographic: "epsg:4326"
 years:  # Include 2011, 2012 data
   start: 2011
   end: 2013
 powerplants:
+  generation_adjustement: EIA
   technology_mapping:
     run_of_river: "run of river"
     reservoir: "reservoir"
   basin_adjustment:
     buffer_radius: 1000  # epsg:3857 is in meters
     max_outside: 50
     max_dropped: 30
-crs:
-  projected: "epsg:3857"
-  geographic: "epsg:4326"
 smoothing_hours: "72h"

pixi.toml

@@ -32,6 +32,7 @@ pandas = ">=2.2.3"
 pandera-geopandas = ">=0.24"
 pyarrow = ">=19.0.1"
 scipy = ">=1.15.2"
+country_converter = ">=1.3.2"
 
 [environments]
 module = { features = ["module"], no-default-feature = true }

tests/conftest.py

@@ -54,7 +54,7 @@ def user_path(cds_api_file) -> Path:
 
 
 @pytest.fixture(scope="module")
-def module_path(user_path: Path):
+def module_path(cds_api_file: Path):
     """Parent directory of the project."""
     module_path = Path(__file__).parent.parent
 
@@ -66,7 +66,9 @@ def module_path(user_path: Path):
         shutil.rmtree(integration_dir / "results/", ignore_errors=True)
 
     # copy user inputs needed by the module
-    dir_to_dir_copy(user_path / "MNE", integration_dir / "resources/inputs/MNE")
+    dir_to_dir_copy(
+        module_path / "tests/files/MNE", integration_dir / "resources/inputs/MNE"
+    )
     # copy downloads with unstable Github CI access
     dir_to_dir_copy(
         module_path / "tests/files/hydrobasins",

tests/files/README.md

@@ -3,3 +3,4 @@
 Used during integration CI tests.
 
 - `hydrobasins`: synthetic alternative to the `.zip` files downloaded from the HydroBASINS website.
+- `MNE`: testfiles for a small-scale Montenegro run.

tests/integration/Snakefile

@@ -11,6 +11,7 @@ module module_hydropower:
         # Redirect specific user resources (inputs)
         shapes="resources/inputs/{shapes}/shapes.parquet",
         powerplants="resources/inputs/{shapes}/powerplants.parquet",
+        statistics_user="resources/inputs/{shapes}/generation.parquet",
         # Redirect specific module results (outputs)
         aggregated_inflow_pu="results/outputs/{shapes}/{plant_type}_pu.parquet",
         # Redirect module intermediate files

tests/integration/test_config.yaml

@@ -6,6 +6,7 @@ module_hydropower:
     start: 2015
     end: 2016
   powerplants:
+    generation_adjustement: user
     technology_mapping:
       run_of_river: "run of river"
       reservoir: "reservoir"

workflow/Snakefile

@@ -12,8 +12,9 @@ pathvars:
     # User resources
     shapes="<resources>/user/{shapes}/shapes.parquet",
     powerplants="<resources>/user/{shapes}/powerplants.parquet",
+    statistics_user="<resources>/user/{shapes}/generation.parquet",
     # Module results
-    statistics="<results>/{shapes}/statistics/generation.parquet",
+    statistics_eia="<results>/{shapes}/statistics/generation.parquet",
     disaggregated_inflow="<results>/{shapes}/disaggregated/inflow_mwh.parquet",
     aggregated_inflow_pu="<results>/{shapes}/aggregated/{plant_type}_inflow_pu.parquet",
 

workflow/envs/module.linux-64.pin.txt

@@ -188,6 +188,7 @@ https://conda.anaconda.org/conda-forge/noarch/decorator-5.3.1-pyhd8ed1ab_0.conda
 https://conda.anaconda.org/conda-forge/noarch/ipython-9.14.1-pyh53cf698_0.conda#fbd58549b374103c1a80577f09a328ef
 https://conda.anaconda.org/conda-forge/noarch/ipdb-0.13.13-pyhd8ed1ab_1.conda#044c5249ad8ea18a414d07baa1f369ea
 https://conda.anaconda.org/conda-forge/linux-64/curl-8.21.0-hcf29cc6_0.conda#35caaabe13a4aa0fe77ea634c71478b8
+https://conda.anaconda.org/conda-forge/noarch/country_converter-1.3.2-pyhd8ed1ab_0.conda#193a9e54636d8d70781a3e56370f5502
 https://conda.anaconda.org/conda-forge/linux-64/yaml-0.2.5-h280c20c_3.conda#a77f85f77be52ff59391544bfe73390a
 https://conda.anaconda.org/conda-forge/noarch/xarray-2026.4.0-pyhc364b38_0.conda#099794df685f800c3f319ff4742dc1bb
 https://conda.anaconda.org/conda-forge/noarch/tqdm-4.68.3-pyh8f84b5b_0.conda#65094960cb7ed216b6049aab57205902

workflow/envs/module.osx-arm64.pin.txt

@@ -177,6 +177,7 @@ https://conda.anaconda.org/conda-forge/noarch/decorator-5.3.1-pyhd8ed1ab_0.conda
 https://conda.anaconda.org/conda-forge/noarch/ipython-9.14.1-pyh53cf698_0.conda#fbd58549b374103c1a80577f09a328ef
 https://conda.anaconda.org/conda-forge/noarch/ipdb-0.13.13-pyhd8ed1ab_1.conda#044c5249ad8ea18a414d07baa1f369ea
 https://conda.anaconda.org/conda-forge/osx-arm64/curl-8.21.0-hd5a2499_0.conda#698b465be69a987c1d0f2fc969d7c986
+https://conda.anaconda.org/conda-forge/noarch/country_converter-1.3.2-pyhd8ed1ab_0.conda#193a9e54636d8d70781a3e56370f5502
 https://conda.anaconda.org/conda-forge/osx-arm64/yaml-0.2.5-h925e9cb_3.conda#78a0fe9e9c50d2c381e8ee47e3ea437d
 https://conda.anaconda.org/conda-forge/noarch/xarray-2026.4.0-pyhc364b38_0.conda#099794df685f800c3f319ff4742dc1bb
 https://conda.anaconda.org/conda-forge/noarch/tqdm-4.68.3-pyh8f84b5b_0.conda#65094960cb7ed216b6049aab57205902

workflow/envs/module.win-64.pin.txt

@@ -170,6 +170,7 @@ https://conda.anaconda.org/conda-forge/noarch/colorama-0.4.6-pyhd8ed1ab_1.conda#
 https://conda.anaconda.org/conda-forge/noarch/ipython-9.14.1-pyhe2676ad_0.conda#bf12187c2d1ef0bb63df01ace31ff26b
 https://conda.anaconda.org/conda-forge/noarch/ipdb-0.13.13-pyhd8ed1ab_1.conda#044c5249ad8ea18a414d07baa1f369ea
 https://conda.anaconda.org/conda-forge/win-64/curl-8.21.0-h8206538_0.conda#b09f232e463d78e4cc0bec233007441b
+https://conda.anaconda.org/conda-forge/noarch/country_converter-1.3.2-pyhd8ed1ab_0.conda#193a9e54636d8d70781a3e56370f5502
 https://conda.anaconda.org/conda-forge/win-64/yaml-0.2.5-h6a83c73_3.conda#433699cba6602098ae8957a323da2664
 https://conda.anaconda.org/conda-forge/noarch/xarray-2026.4.0-pyhc364b38_0.conda#099794df685f800c3f319ff4742dc1bb
 https://conda.anaconda.org/conda-forge/noarch/tqdm-4.68.3-pyha7b4d00_0.conda#0b814124391b7e176bcc0cce445a3a36

workflow/envs/module.yaml

@@ -13,3 +13,4 @@ dependencies:
 - pandera-geopandas >=0.24
 - pyarrow >=19.0.1
 - scipy >=1.15.2
+- country_converter >=1.3.2

workflow/internal/config.schema.yaml

@@ -2,6 +2,12 @@ $schema: "https://json-schema.org/draft/2020-12/schema"
 description: "Schema for user-provided configuration files."
 type: object
 additionalProperties: false
+required:
+  - pfafstetter_level
+  - years
+  - powerplants
+  - crs
+  - smoothing_hours
 properties:
   pfafstetter_level:
     description: >
@@ -30,6 +36,15 @@ properties:
     description: Powerplant related settings.
     type: object
     properties:
+      generation_adjustement:
+        description: |
+          Specifies which file to use for powerplant inflow normalisation.
+          - EIA: uses the internally generated EIA statistics (results.statistics_eia).
+          - user: uses a user provided file instead.
+        type: string
+        enum:
+          - EIA
+          - user
       technology_mapping:
         description: >
           Maps internal names to user-provided names in powerplants files, to make integration smoother.
@@ -75,21 +90,22 @@ properties:
           - max_dropped
         additionalProperties: false  # Ensures no extra properties in "basin_adjustment"
     required:
+      - generation_adjustement
       - basin_adjustment
       - technology_mapping
     additionalProperties: false  # Ensures no extra properties in "powerplants"
   crs:
-    description: Coordinate reference system settings in the form 'epsg:xxxx'.
+    description: Coordinate reference system settings in the form 'EPSG:xxxx'.
     type: object
     properties:
       projected:
         description: CRS to use during operations related to distance / area.
         type: string
-        pattern: "^epsg:\\d{4,5}$"
+        pattern: "^[A-Za-z][A-Za-z0-9_-]*:[A-Za-z0-9_.-]+$"
       geographic:
         description: CRS to use during operations related to location.
         type: string
-        pattern: "^epsg:\\d{4,5}$"
+        pattern: "^[A-Za-z][A-Za-z0-9_-]*:[A-Za-z0-9_.-]+$"
     required:
       - projected
       - geographic

workflow/rules/_utils.smk

@@ -24,3 +24,15 @@ def validate_year_config():
             f"Start year must be less than end year. Found: '{start},{end}'. "
             "Remember: end year data is not included!"
         )
+
+
+def get_generation_adjustement_file() -> str:
+    """Identify which generation adjustment file to use."""
+    source = config["powerplants"]["generation_adjustement"]
+    match source:
+        case "EIA":
+            return "<statistics_eia>"
+        case "user":
+            return "<statistics_user>"
+        case _:
+            raise ValueError(f"Invalid generation adjustment option {source}")

workflow/rules/powerplants.smk

@@ -50,7 +50,7 @@ rule powerplants_get_inflow_mwh:
     input:
         inflow_m3=rules.powerplants_get_inflow_m3.output.inflow,
         adjusted_powerplants=rules.powerplants_adjust_location.output.adjusted_powerplants,
-        statistics="<statistics>",
+        statistics=get_generation_adjustement_file(),
     output:
         inflow_mwh="<disaggregated_inflow>",
     log:

workflow/rules/statistics.smk

@@ -1,19 +1,32 @@
 """Rules for collating and preparing hydropower national statistics."""
 
 
+rule harmonise_statistics:
+    input:
+        eia_bulk=rules.download_eia.output.zipfile,
+    output:
+        generation="<resources>/automatic/statistics/baseline.parquet",
+    log:
+        "<logs>/harmonise_statistics.log",
+    conda:
+        "../envs/module.yaml"
+    message:
+        "Harmonise EIA annual hydropower generation for all countries and years."
+    script:
+        "../scripts/harmonise_statistics.py"
+
+
 rule prepare_statistics:
     input:
         shapes="<shapes>",
-        eia_bulk=rules.download_eia.output.zipfile,
+        statistics=rules.harmonise_statistics.output.generation,
     output:
-        generation="<statistics>",
+        generation="<statistics_eia>",
         plot="<results>/{shapes}/statistics/generation.pdf",
     log:
         "<logs>/{shapes}/prepare_statistics.log",
     conda:
         "../envs/module.yaml"
-    params:
-        years=config["years"],
     message:
         "Get EIA annual country hydropower generation statistics."
     script:

workflow/scripts/_schemas.py

@@ -12,7 +12,7 @@ class Config:
 
     shape_id: Series[str] = Field(unique=True)
     "Unique ID for this shape."
-    country_id: Series[str]
+    country_id: Series[str] = Field(str_length={"min_value": 3, "max_value": 3})
     "ISO alpha-3 code."
     shape_class: Series[str] = Field(isin=["land", "maritime"])
     "Shape classifier"
@@ -30,13 +30,14 @@ class EIAGenerationSchema(DataFrameModel):
     class Config:
         coerce = True
         strict = True
+        unique = ["country_id", "year", "category"]
 
     country_id: Series[str]
-    "Country ISO-3 code"
+    "EIA country or territory code, including historical and non-ISO codes"
     year: Series[int]
     "Sample year"
     category: Series[str] = Field(isin=["hydropower"])
-    generation_mwh: Series[float] = Field(ge=0)
+    generation_mwh: Series[float] = Field(ge=0, nullable=True)
     "Hydropower generation for the given year"
 
     index: Index[int] = Field(unique=True)