Move to pandera io (#11)

* Move to pandera io

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: irm-codebase

pixi.lock

@@ -23,13 +23,14 @@ ipykernel = ">=6.29.5,<7"
 pyarrow = ">=19.0.1,<20"
 pycountry = ">=24.6.1,<25"
 numpy = ">=2.2.4,<3"
-pandera-geopandas = ">=0.23.1,<0.24"
+pandera-geopandas = ">=0.24.0,<0.25"
 contextily = ">=1.6.2,<2"
 ipdb = ">=0.13.13"
 atlite = ">=0.4.0,<0.5"
 geofileops = ">=0.9.2,<0.10"
 xarray = ">=2025.3.0,<2026"
 scipy = ">=1.15.2,<2"
+pandera-io = ">=0.24.0,<0.25"
 
 [pypi-dependencies]
 mkdocs-mermaid2-plugin = ">=1.2.1"

pixi.toml

@@ -8,6 +8,7 @@ dependencies:
   - ipdb = 0.13.13
   - numpy = 2.2.4
   - pandas = 2.2.3
-  - pandera-geopandas = 0.23.1
+  - pandera-geopandas = 0.24
+  - pandera-io = 0.24
   - pyarrow = 19.0.1
   - scipy = 1.15.2

workflow/envs/default.yaml

@@ -0,0 +1,59 @@
+schema_type: dataframe
+version: 0.24.0
+columns:
+  country_id:
+    title: null
+    description: null
+    dtype: str
+    nullable: false
+    checks: null
+    unique: false
+    coerce: false
+    required: true
+    regex: false
+  year:
+    title: null
+    description: null
+    dtype: int64
+    nullable: false
+    checks: null
+    unique: false
+    coerce: false
+    required: true
+    regex: false
+  generation_mwh:
+    title: null
+    description: null
+    dtype: float64
+    nullable: false
+    checks:
+    - value: 0
+      options:
+        check_name: greater_than_or_equal_to
+        raise_warning: false
+        ignore_na: true
+    unique: false
+    coerce: false
+    required: true
+    regex: false
+checks: null
+index:
+- title: null
+  description: null
+  dtype: int64
+  nullable: false
+  checks: null
+  name: null
+  unique: true
+  coerce: false
+dtype: null
+coerce: true
+strict: false
+name: NationalGenerationSchema
+ordered: false
+unique: null
+report_duplicates: all
+unique_column_names: false
+add_missing_columns: false
+title: null
+description: null

workflow/internal/national_generation.schema.yaml

@@ -0,0 +1,86 @@
+schema_type: dataframe
+version: 0.24.0
+columns:
+  powerplant_id:
+    title: null
+    description: null
+    dtype: str
+    nullable: false
+    checks: null
+    unique: true
+    coerce: false
+    required: true
+    regex: false
+  net_generation_capacity_mw:
+    title: null
+    description: null
+    dtype: float64
+    nullable: false
+    checks:
+    - value: 0
+      options:
+        check_name: greater_than_or_equal_to
+        raise_warning: false
+        ignore_na: true
+    unique: false
+    coerce: false
+    required: true
+    regex: false
+  storage_capacity_mwh:
+    title: null
+    description: null
+    dtype: float64
+    nullable: true
+    checks: null
+    unique: false
+    coerce: false
+    required: true
+    regex: false
+  powerplant_type:
+    title: null
+    description: null
+    dtype: str
+    nullable: false
+    checks:
+    - value:
+      - hydro_run_of_river
+      - hydro_dam
+      options:
+        check_name: isin
+        raise_warning: false
+        ignore_na: true
+    unique: false
+    coerce: false
+    required: true
+    regex: false
+  geometry:
+    title: null
+    description: null
+    dtype: geometry
+    nullable: false
+    checks: null
+    unique: false
+    coerce: false
+    required: true
+    regex: false
+checks: null
+index:
+- title: null
+  description: null
+  dtype: int64
+  nullable: false
+  checks: null
+  name: null
+  unique: true
+  coerce: false
+dtype: null
+coerce: true
+strict: false
+name: PowerplantSchema
+ordered: false
+unique: null
+report_duplicates: all
+unique_column_names: false
+add_missing_columns: false
+title: null
+description: null

workflow/internal/powerplant.schema.yaml

@@ -0,0 +1,70 @@
+schema_type: dataframe
+version: 0.24.0
+columns:
+  shape_id:
+    title: null
+    description: null
+    dtype: str
+    nullable: false
+    checks: null
+    unique: true
+    coerce: false
+    required: true
+    regex: false
+  country_id:
+    title: null
+    description: null
+    dtype: str
+    nullable: false
+    checks: null
+    unique: false
+    coerce: false
+    required: true
+    regex: false
+  shape_class:
+    title: null
+    description: null
+    dtype: str
+    nullable: false
+    checks:
+    - value:
+      - land
+      options:
+        check_name: isin
+        raise_warning: false
+        ignore_na: true
+    unique: false
+    coerce: false
+    required: true
+    regex: false
+  geometry:
+    title: null
+    description: null
+    dtype: geometry
+    nullable: false
+    checks: null
+    unique: false
+    coerce: false
+    required: true
+    regex: false
+checks: null
+index:
+- title: null
+  description: null
+  dtype: int64
+  nullable: false
+  checks: null
+  name: null
+  unique: true
+  coerce: false
+dtype: null
+coerce: true
+strict: false
+name: ShapeSchema
+ordered: false
+unique: null
+report_duplicates: all
+unique_column_names: false
+add_missing_columns: false
+title: null
+description: null

workflow/internal/shape.schema.yaml

@@ -12,7 +12,9 @@ rule powerplants_adjust_location:
             f"resources/automatic/hydrobasin_global_{config["pfafstetter_level"]}.parquet"
         ),
         powerplants="resources/user/powerplants.parquet",
+        powerplant_schema=workflow.source_path("../internal/powerplant.schema.yaml"),
         shapes="resources/user/shapes.parquet",
+        shape_schema=workflow.source_path("../internal/shape.schema.yaml"),
     output:
         adjusted_powerplants="results/adjusted_powerplants.parquet",
         plot=report(
@@ -34,11 +36,13 @@ rule powerplants_get_inflow_m3:
     params:
         smoothing_hours=config["smoothing_hours"],
     input:
+        adjusted_powerplants="results/adjusted_powerplants.parquet",
         basins=ancient(
             f"resources/automatic/hydrobasin_global_{config["pfafstetter_level"]}.parquet"
         ),
+        powerplant_schema=workflow.source_path("../internal/powerplant.schema.yaml"),
         shapes="resources/user/shapes.parquet",
-        adjusted_powerplants="results/adjusted_powerplants.parquet",
+        shape_schema=workflow.source_path("../internal/shape.schema.yaml"),
         cutout=ancient("resources/automatic/cutout.nc"),
     output:
         inflow="results/by_powerplant_id/inflow_m3.parquet",
@@ -58,7 +62,11 @@ rule powerplants_get_inflow_mwh:
     input:
         inflow_m3="results/by_powerplant_id/inflow_m3.parquet",
         adjusted_powerplants="results/adjusted_powerplants.parquet",
-        generation="resources/user/national_generation.parquet",
+        powerplant_schema=workflow.source_path("../internal/powerplant.schema.yaml"),
+        national_generation="resources/user/national_generation.parquet",
+        national_generation_schema=workflow.source_path(
+            "../internal/national_generation.schema.yaml"
+        ),
     output:
         inflow_mwh="results/by_powerplant_id/inflow_mwh.parquet",
     log:
@@ -75,6 +83,7 @@ rule powerplants_get_cf_per_shape:
     input:
         adjusted_powerplants="results/adjusted_powerplants.parquet",
         inflow_mwh="results/by_powerplant_id/inflow_mwh.parquet",
+        schema=workflow.source_path("../internal/powerplant.schema.yaml"),
     output:
         timeseries="results/by_shape_id/{plant_type}_cf.parquet",
         figure=report(

workflow/rules/powerplants.smk

@@ -6,12 +6,14 @@
 
 import geopandas as gpd
 import matplotlib.pyplot as plt
-from _schema import PowerplantSchema, ShapeSchema
+import pandera.io as io
 from pyproj import CRS
 
 if TYPE_CHECKING:
     snakemake: Any
 sys.stderr = open(snakemake.log[0], "w")
+POWERPLANT_SCHEMA = io.from_yaml(snakemake.input.powerplant_schema)
+SHAPE_SCHEMA = io.from_yaml(snakemake.input.shape_schema)
 
 
 def _plot_adjustment(
@@ -64,9 +66,9 @@ def powerplants_adjust_location(
     # Read and validate input files
     basins = gpd.read_parquet(basins_path)
     powerplants = gpd.read_parquet(powerplants_path)
-    PowerplantSchema.validate(powerplants)
+    POWERPLANT_SCHEMA.validate(powerplants)
     shapes = gpd.read_parquet(shapes_path)
-    ShapeSchema.validate(shapes)
+    SHAPE_SCHEMA.validate(shapes)
 
     # Coordinate-based operations must use a geographic CRS
     basins = basins.to_crs(crs["geographic"])
@@ -126,7 +128,7 @@ def powerplants_adjust_location(
 
     # Re-validate and save
     powerplants = powerplants.to_crs(crs["geographic"])
-    PowerplantSchema.validate(powerplants)
+    POWERPLANT_SCHEMA.validate(powerplants)
     powerplants.to_parquet(adjusted_powerplants_path)