From 3cdd07e24dd641aac411501ab1cef3266c6e41e1 Mon Sep 17 00:00:00 2001
From: hgloeckner <helene.gloeckner@posteo.de>
Date: Mon, 19 Aug 2024 13:49:45 +0200
Subject: [PATCH] write l3

---
 src/halodrops/helper/__init__.py |  2 ++
 src/halodrops/pipeline.py        | 21 ++++++++++++++++++
 src/halodrops/processor.py       | 38 ++++++++++++++++++++++++++++++++
 tests/test_gridded.py            | 37 +++++++++++++++++++++++++++++++
 4 files changed, 98 insertions(+)
 create mode 100644 tests/test_gridded.py

diff --git a/src/halodrops/helper/__init__.py b/src/halodrops/helper/__init__.py
index 0536980..9773aa2 100644
--- a/src/halodrops/helper/__init__.py
+++ b/src/halodrops/helper/__init__.py
@@ -113,6 +113,8 @@
 
 l2_filename_template = "{platform}_{launch_time}_{flight_id}_{serial_id}_Level_2.nc"
 
+l3_filename_template = "{platform}_{flight_id}_Level_3.nc"
+
 
 def get_bool(s):
     if isinstance(s, bool):
diff --git a/src/halodrops/pipeline.py b/src/halodrops/pipeline.py
index 0fa9649..6d18e42 100644
--- a/src/halodrops/pipeline.py
+++ b/src/halodrops/pipeline.py
@@ -300,6 +300,20 @@ def sondes_to_gridded(sondes: dict, config: configparser.ConfigParser):
     return gridded
 
 
+def iterate_method_over_dataset(
+    obj: Gridded,
+    functions: list,
+    config: configparser.ConfigParser,
+) -> xr.Dataset:
+    """
+    This is NOT what the function should do in the end. only used to save the base-l3
+    """
+    for function_name in functions:
+        function = getattr(Gridded, function_name)
+        result = function(obj, **get_args_for_function(config, function))
+    return result
+
+
 def gridded_to_pattern(
     gridded: xr.Dataset, config: configparser.ConfigParser
 ) -> xr.Dataset:
@@ -459,6 +473,13 @@ def run_pipeline(pipeline: dict, config: configparser.ConfigParser):
         "output": "gridded",
         "comment": "This step concatenates the individual sonde datasets to create the L3 dataset.",
     },
+    "create_L3": {
+        "intake": "gridded",
+        "apply": iterate_method_over_dataset,
+        "functions": ["get_l3_dir", "get_l3_filename", "write_l3"],
+        "output": "gridded",
+        "comment": "This step creates the L3 dataset after adding additional products.",
+    },
     # "create_patterns": {
     #     "intake": "gridded",
     #     "apply": gridded_to_pattern,
diff --git a/src/halodrops/processor.py b/src/halodrops/processor.py
index 72fc3b4..9833d77 100644
--- a/src/halodrops/processor.py
+++ b/src/halodrops/processor.py
@@ -1115,3 +1115,41 @@ def concat_sondes(self):
         combined = xr.combine_by_coords(list_of_l2_ds)
         self._interim_l3_ds = combined
         return self
+
+    def get_l3_dir(self, l3_dirname: str = None):
+        if l3_dirname:
+            self.l3_dir = l3_dirname
+        elif not self.sondes is None:
+            self.l3_dir = list(self.sondes.values())[0].l3_dir
+        else:
+            raise ValueError("No sondes and no l3 directory given, cannot continue ")
+
+    def get_l3_filename(
+        self, l3_filename_template: str = None, l3_filename: str = None
+    ):
+        if l3_filename is None:
+            if l3_filename_template is None:
+                l3_filename = hh.l3_filename_template.format(
+                    platform=self.platform_id,
+                    flight_id=self.flight_id,
+                )
+            else:
+                l3_filename = l3_filename_template.format(
+                    platform=self.platform_id,
+                    flight_id=self.flight_id,
+                )
+
+        self.l3_filename = l3_filename
+
+        return self
+
+    def write_l3(self, l3_dir: str = None):
+        if l3_dir is None:
+            l3_dir = self.l3_dir
+
+        if not os.path.exists(l3_dir):
+            os.makedirs(l3_dir)
+
+        self._interim_l3_ds.to_netcdf(os.path.join(l3_dir, self.l3_filename))
+
+        return self
diff --git a/tests/test_gridded.py b/tests/test_gridded.py
new file mode 100644
index 0000000..ed03032
--- /dev/null
+++ b/tests/test_gridded.py
@@ -0,0 +1,37 @@
+import pytest
+import os
+import xarray as xr
+from halodrops.processor import Gridded
+
+sondes = None
+flight_id = "20240811"
+platform_id = "HALO"
+l3_default = f"{platform_id}_{flight_id}_Level_3.nc"
+l3_template = "{flight_id}_{platform}_Level_3.nc"
+
+
+@pytest.fixture
+def gridded():
+    return Gridded(sondes, flight_id=flight_id, platform_id=platform_id)
+
+
+def test_l3_dir(gridded):
+    with pytest.raises(ValueError):
+        gridded.get_l3_dir()
+
+
+def test_l3_dir_name(gridded):
+    gridded.get_l3_dir(l3_dirname="test")
+    assert gridded.l3_dir == "test"
+
+
+def test_l3_default(gridded):
+    gridded.get_l3_filename()
+    assert gridded.l3_filename == l3_default
+
+
+def test_l3_template(gridded):
+    gridded.get_l3_filename(l3_filename_template=l3_template)
+    assert gridded.l3_filename == l3_template.format(
+        flight_id=flight_id, platform=platform_id
+    )