Skip to content

Commit

Permalink
write l3
Browse files Browse the repository at this point in the history
  • Loading branch information
hgloeckner committed Aug 19, 2024
1 parent 1c6dda6 commit 3cdd07e
Show file tree
Hide file tree
Showing 4 changed files with 98 additions and 0 deletions.
2 changes: 2 additions & 0 deletions src/halodrops/helper/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
21 changes: 21 additions & 0 deletions src/halodrops/pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down
38 changes: 38 additions & 0 deletions src/halodrops/processor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
37 changes: 37 additions & 0 deletions tests/test_gridded.py
Original file line number Diff line number Diff line change
@@ -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
)

0 comments on commit 3cdd07e

Please sign in to comment.