Skip to content

Commit

Permalink
Code style and cleanups
Browse files Browse the repository at this point in the history
  • Loading branch information
alexamici committed Oct 15, 2023
1 parent d5c9762 commit 5cad720
Show file tree
Hide file tree
Showing 6 changed files with 100 additions and 35 deletions.
36 changes: 29 additions & 7 deletions cf2cdm/cfcoords.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,9 @@ def coord_translator(
) -> xr.Dataset:
out_name = coord_model.get(cf_type, {}).get("out_name", default_out_name)
units = coord_model.get(cf_type, {}).get("units", default_units)
stored_direction = coord_model.get(cf_type, {}).get("stored_direction", default_direction)
stored_direction = coord_model.get(cf_type, {}).get(
"stored_direction", default_direction
)
matches = match_values(is_cf_type, data.coords)
if len(matches) > 1:
raise ValueError("found more than one CF coordinate with type %r." % cf_type)
Expand All @@ -76,19 +78,30 @@ def coord_translator(
match = matches[0]
for name in data.coords:
if name == out_name and name != match:
raise ValueError("found non CF compliant coordinate with type %r." % cf_type)
raise ValueError(
"found non CF compliant coordinate with type %r." % cf_type
)
data = data.rename({match: out_name})
coord = data.coords[out_name]
if "units" in coord.attrs:
data.coords[out_name] = cfunits.convert_units(coord, units, coord.attrs["units"])
data.coords[out_name] = cfunits.convert_units(
coord, units, coord.attrs["units"]
)
data.coords[out_name].attrs.update(coord.attrs)
data.coords[out_name].attrs["units"] = units
if out_name in data.dims:
data = translate_coord_direction(data, out_name, stored_direction)
return data


VALID_LAT_UNITS = ["degrees_north", "degree_north", "degree_N", "degrees_N", "degreeN", "degreesN"]
VALID_LAT_UNITS = [
"degrees_north",
"degree_north",
"degree_N",
"degrees_N",
"degreeN",
"degreesN",
]


def is_latitude(coord: xr.IndexVariable) -> bool:
Expand All @@ -100,7 +113,14 @@ def is_latitude(coord: xr.IndexVariable) -> bool:
)


VALID_LON_UNITS = ["degrees_east", "degree_east", "degree_E", "degrees_E", "degreeE", "degreesE"]
VALID_LON_UNITS = [
"degrees_east",
"degree_east",
"degree_E",
"degrees_E",
"degreeE",
"degreesE",
]


def is_longitude(coord: xr.IndexVariable) -> bool:
Expand Down Expand Up @@ -128,7 +148,9 @@ def is_step(coord: xr.IndexVariable) -> bool:
return coord.attrs.get("standard_name") == "forecast_period"


COORD_TRANSLATORS["step"] = functools.partial(coord_translator, "step", "h", "increasing", is_step)
COORD_TRANSLATORS["step"] = functools.partial(
coord_translator, "step", "h", "increasing", is_step
)


def is_valid_time(coord: xr.IndexVariable) -> bool:
Expand Down Expand Up @@ -188,7 +210,7 @@ def translate_coords(
for cf_name, translator in coord_translators.items():
try:
data = translator(cf_name, data, coord_model)
except:
except Exception:
if errors == "ignore":
pass
elif errors == "raise":
Expand Down
18 changes: 15 additions & 3 deletions cf2cdm/cfunits.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,15 @@

PRESSURE_CONVERSION_RULES: T.Dict[T.Tuple[str, ...], float] = {
("Pa", "pascal", "pascals"): 1.0,
("hPa", "hectopascal", "hectopascals", "hpascal", "millibar", "millibars", "mbar"): 100.0,
(
"hPa",
"hectopascal",
"hectopascals",
"hpascal",
"millibar",
"millibars",
"mbar",
): 100.0,
("decibar", "dbar"): 10000.0,
("bar", "bars"): 100000.0,
("atmosphere", "atmospheres", "atm"): 101325.0,
Expand Down Expand Up @@ -50,7 +58,9 @@ def simple_conversion_factor(source_units, target_units, rules):
conversion_factor *= factor
seen += 1
if seen != 2:
raise ConversionError("cannot convert from %r to %r." % (source_units, target_units))
raise ConversionError(
"cannot convert from %r to %r." % (source_units, target_units)
)
return conversion_factor


Expand All @@ -62,7 +72,9 @@ def convert_units(data: T.Any, target_units: str, source_units: str) -> T.Any:
return data * simple_conversion_factor(target_units, source_units, rules)
except ConversionError:
pass
raise ConversionError("cannot convert from %r to %r." % (source_units, target_units))
raise ConversionError(
"cannot convert from %r to %r." % (source_units, target_units)
)


def are_convertible(source_units: str, target_units: str) -> bool:
Expand Down
30 changes: 25 additions & 5 deletions cf2cdm/datamodels.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,16 @@
"latitude": {"out_name": "lat", "stored_direction": "increasing"},
"longitude": {"out_name": "lon", "stored_direction": "increasing"},
# vertical
"depthBelowLand": {"out_name": "depth", "units": "m", "stored_direction": "increasing"},
"isobaricInhPa": {"out_name": "plev", "units": "Pa", "stored_direction": "decreasing"},
"depthBelowLand": {
"out_name": "depth",
"units": "m",
"stored_direction": "increasing",
},
"isobaricInhPa": {
"out_name": "plev",
"units": "Pa",
"stored_direction": "decreasing",
},
# ensemble
"number": {"out_name": "realization", "stored_direction": "increasing"},
# time
Expand All @@ -35,8 +43,20 @@


ECMWF = {
"depthBelowLand": {"out_name": "level", "units": "m", "stored_direction": "increasing"},
"isobaricInhPa": {"out_name": "level", "units": "hPa", "stored_direction": "decreasing"},
"isobaricInPa": {"out_name": "level", "units": "hPa", "stored_direction": "decreasing"},
"depthBelowLand": {
"out_name": "level",
"units": "m",
"stored_direction": "increasing",
},
"isobaricInhPa": {
"out_name": "level",
"units": "hPa",
"stored_direction": "decreasing",
},
"isobaricInPa": {
"out_name": "level",
"units": "hPa",
"stored_direction": "decreasing",
},
"hybrid": {"out_name": "level", "stored_direction": "increasing"},
}
4 changes: 0 additions & 4 deletions tests/test_10_cfunits.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,3 @@
import pytest

xr = pytest.importorskip("xarray") # noqa

from cf2cdm import cfunits


Expand Down
34 changes: 23 additions & 11 deletions tests/test_20_cfcoords.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,7 @@
import sys
import typing as T

import numpy as np
import pytest

pytest.importorskip("xarray") # noqa

import xarray as xr

from cf2cdm import cfcoords
Expand Down Expand Up @@ -63,7 +59,11 @@ def da3() -> xr.Dataset:
coords=[
("lat", latitude, {"units": "degrees_north"}),
("lon", longitude, {"units": "degrees_east"}),
("step", np.array(step, dtype="timedelta64[h]"), {"standard_name": "forecast_period"}),
(
"step",
np.array(step, dtype="timedelta64[h]"),
{"standard_name": "forecast_period"},
),
(
"ref_time",
np.array(time, dtype="datetime64[ns]"),
Expand Down Expand Up @@ -104,22 +104,34 @@ def test_translate_coord_direction(da1: xr.Dataset) -> None:


def test_coord_translator(da1: xr.Dataset) -> None:
res = cfcoords.coord_translator("level", "hPa", "decreasing", lambda x: False, "lvl", da1)
res = cfcoords.coord_translator(
"level", "hPa", "decreasing", lambda x: False, "lvl", da1
)
assert da1.equals(res)

with pytest.raises(ValueError):
cfcoords.coord_translator("level", "hPa", "decreasing", lambda x: True, "lvl", da1)
cfcoords.coord_translator(
"level", "hPa", "decreasing", lambda x: True, "lvl", da1
)

res = cfcoords.coord_translator("level", "hPa", "decreasing", cfcoords.is_isobaric, "lvl", da1)
res = cfcoords.coord_translator(
"level", "hPa", "decreasing", cfcoords.is_isobaric, "lvl", da1
)
assert da1.equals(res)

with pytest.raises(ValueError):
cfcoords.coord_translator("level", "hPa", "decreasing", cfcoords.is_latitude, "lvl", da1)
cfcoords.coord_translator(
"level", "hPa", "decreasing", cfcoords.is_latitude, "lvl", da1
)

res = cfcoords.coord_translator("level", "Pa", "decreasing", cfcoords.is_isobaric, "lvl", da1)
res = cfcoords.coord_translator(
"level", "Pa", "decreasing", cfcoords.is_isobaric, "lvl", da1
)
assert not da1.equals(res)

res = cfcoords.coord_translator("step", "h", "increasing", cfcoords.is_step, "step", da1)
res = cfcoords.coord_translator(
"step", "h", "increasing", cfcoords.is_step, "step", da1
)
assert da1.equals(res)


Expand Down
13 changes: 8 additions & 5 deletions tests/test_50_datamodels.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,8 @@
import os.path

import pytest

xr = pytest.importorskip("xarray") # noqa
from cfgrib import xarray_store

from cf2cdm import cfcoords, datamodels
from cfgrib import xarray_store

SAMPLE_DATA_FOLDER = os.path.join(os.path.dirname(__file__), "sample-data")
TEST_DATA1 = os.path.join(SAMPLE_DATA_FOLDER, "era5-levels-members.grib")
Expand All @@ -17,7 +14,13 @@ def test_cds() -> None:

res = cfcoords.translate_coords(ds, coord_model=datamodels.CDS)

assert set(res.dims) == {"forecast_reference_time", "lat", "lon", "plev", "realization"}
assert set(res.dims) == {
"forecast_reference_time",
"lat",
"lon",
"plev",
"realization",
}
assert set(res.coords) == {
"forecast_reference_time",
"lat",
Expand Down

0 comments on commit 5cad720

Please sign in to comment.