From cc9b8ecf5cf9a55d1fbcddd855795bd95723fb4b Mon Sep 17 00:00:00 2001
From: Alvaro Valdebenito <alvarov@met.no>
Date: Wed, 6 Jul 2022 13:24:17 +0200
Subject: [PATCH] implement TM5_subset in python

---
 scripts/testdata-minimal/TM5_subset.sh        |  8 ----
 .../testdata-minimal/create_subsets_tm5.py    | 43 +++++++++++++++++++
 2 files changed, 43 insertions(+), 8 deletions(-)
 delete mode 100755 scripts/testdata-minimal/TM5_subset.sh
 create mode 100755 scripts/testdata-minimal/create_subsets_tm5.py

diff --git a/scripts/testdata-minimal/TM5_subset.sh b/scripts/testdata-minimal/TM5_subset.sh
deleted file mode 100755
index b0e537e9f..000000000
--- a/scripts/testdata-minimal/TM5_subset.sh
+++ /dev/null
@@ -1,8 +0,0 @@
-#!/bin/bash -x
-# Extract a few lat/lon points to decrease filesize
-LON=20,30
-LAT=20,30
-
-ncks -d lat,"$LAT" -d lon,"$LON" aerocom3_TM5-met2010_AP3-CTRL2019_abs550aer_Column_2010_daily.nc ./aerocom3_TM5-met2010_AP3-CTRL2019_abs550aer_Column_2010_daily.nc
-
-ncks -d lat,"$LAT" -d lon,"$LON" aerocom3_TM5-met2010_AP3-CTRL2019_od550aer_Column_2010_daily.nc ./aerocom3_TM5-met2010_AP3-CTRL2019_od550aer_Column_2010_daily.nc
diff --git a/scripts/testdata-minimal/create_subsets_tm5.py b/scripts/testdata-minimal/create_subsets_tm5.py
new file mode 100755
index 000000000..08b408e81
--- /dev/null
+++ b/scripts/testdata-minimal/create_subsets_tm5.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+from pathlib import Path
+
+import xarray as xr
+
+from tests.fixtures.tm5 import TM5_DATA_PATH
+
+SRC_DATA_PATH = Path("/lustre/storeA/project/aerocom/aerocom-users-database")
+SRC_DATA_PATH /= "AEROCOM-PHASE-III-2019/TM5-met2010_AP3-CTRL2019/renamed"
+
+LON, LAT = slice(20, 30), slice(20, 30)
+
+PATHS = {
+    TM5_DATA_PATH / "aerocom3_TM5-met2010_AP3-CTRL2019_abs550aer_Column_2010_daily.nc",
+    TM5_DATA_PATH / "aerocom3_TM5-met2010_AP3-CTRL2019_od550aer_Column_2010_daily.nc",
+}
+
+
+def reduce_dims(ds: xr.Dataset) -> xr.Dataset:
+    """crop domain"""
+    return ds.isel(lon=LON, lat=LAT)
+
+
+def atomic_write(ds: xr.Dataset, path: Path, **kwargs) -> None:
+    """write dataset to a netcdf file atomically"""
+    tmp = path.with_suffix(".tmp")
+    try:
+        ds.to_netcdf(tmp, **kwargs)
+        tmp.rename(path)
+    finally:
+        tmp.unlink(missing_ok=True)
+
+
+def main():
+    for path in PATHS:
+        ds = xr.open_dataset(SRC_DATA_PATH / path.name).pipe(reduce_dims)
+        atomic_write(ds, path)
+
+
+if __name__ == "__main__":
+    main()