Skip to content

Commit

Permalink
CDAT migration: Fix African easterly wave density plots in TC analysi…
Browse files Browse the repository at this point in the history
…s and convert H20LNZ units to ppm/volume (#882)
  • Loading branch information
tomvothecoder committed Oct 29, 2024
1 parent 8031e2f commit 7550b3d
Show file tree
Hide file tree
Showing 3 changed files with 43 additions and 22 deletions.
56 changes: 40 additions & 16 deletions e3sm_diags/driver/tc_analysis_driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,9 +81,7 @@ def run_diag(parameter: TCAnalysisParameter) -> TCAnalysisParameter:
test_data_path,
"aew_hist_{}_{}_{}.nc".format(test_name, test_start_yr, test_end_yr),
)
test_aew_hist = xr.open_dataset(test_aew_file).sel(
lat=slice(0, 35), lon=slice(180, 360)
)["density"]
test_aew_hist = xr.open_dataset(test_aew_file).sel()["density"]

test_data = collections.OrderedDict()
ref_data = collections.OrderedDict()
Expand Down Expand Up @@ -120,9 +118,7 @@ def run_diag(parameter: TCAnalysisParameter) -> TCAnalysisParameter:
"aew_hist_{}_{}_{}.nc".format(ref_name, ref_start_yr, ref_end_yr),
)
# Note the refactor included subset that was missed in original implementation
ref_aew_hist = xr.open_dataset(ref_aew_file).sel(
lat=slice(0, 35), lon=slice(180, 360)
)["density"]
ref_aew_hist = xr.open_dataset(ref_aew_file).sel()["density"]
ref_data["metrics"] = generate_tc_metrics_from_te_stitch_file(ref_te_file)
ref_data["cyclone_density"] = ref_cyclones_hist # type: ignore
ref_data["aew_density"] = ref_aew_hist # type: ignore
Expand Down Expand Up @@ -181,19 +177,10 @@ def generate_tc_metrics_from_te_stitch_file(te_stitch_file: str) -> Dict[str, An
if not lines_orig:
raise ValueError(f"The file {te_stitch_file} is empty.")

line_ind = []
data_start_year = int(te_stitch_file.split(".")[-2].split("_")[-2])
data_end_year = int(te_stitch_file.split(".")[-2].split("_")[-1])
for i in range(0, np.size(lines_orig)):
if lines_orig[i][0] == "s":
year = int(lines_orig[i].split("\t")[2])

if year <= data_end_year:
line_ind.append(i)

# Remove excessive time points cross year bounds from 6 hourly data
end_ind = line_ind[-1]
lines = lines_orig[0:end_ind]
lines = _filter_lines_within_year_bounds(lines_orig, data_end_year)

if not lines:
raise ValueError(f"The file {te_stitch_file} is empty.")
Expand Down Expand Up @@ -243,6 +230,43 @@ def generate_tc_metrics_from_te_stitch_file(te_stitch_file: str) -> Dict[str, An
return result_mod


def _filter_lines_within_year_bounds(
lines_orig: List[str], data_end_year: int
) -> List[str]:
"""Filters lines within the specified year bounds.
This function processes a list of strings, each representing a line of data.
It filters out lines based on a year extracted from each line, ensuring that
only lines with years less than or equal to `data_end_year` are retained.
Additionally, it removes excessive time points crossing year bounds from
6-hourly data.
Parameters
----------
lines_orig : List[str]
A list of strings where each string represents a line of data.
data_end_year : int
The end year for filtering lines. Only lines with years less than or
equal to this value will be retained.
Returns
-------
List[str]
A list of strings filtered based on the specified year bounds.
"""
line_ind = []
for i in range(0, np.size(lines_orig)):
if lines_orig[i][0] == "s":
year = int(lines_orig[i].split("\t")[2])

if year <= data_end_year:
line_ind.append(i)

end_ind = line_ind[-1]

new_lines = lines_orig[0:end_ind]
return new_lines


def _calc_num_storms_and_max_len(lines: List[str]) -> Tuple[int, int]:
"""Calculate number of storms and max length using lines from a TE stitch file.
Expand Down
6 changes: 3 additions & 3 deletions e3sm_diags/driver/zonal_mean_2d_driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ def _convert_g_kg_to_ppm_units(
This is a special case to handle small values of stratosphere specific
humidity. The general derived variable process converts specific humidity to
units [g/kg]. This function converts from "g/kg" to "ppm".
units [g/kg]. This function converts from "g/kg" to "ppm" by volume.
Parameters
----------
Expand All @@ -161,8 +161,8 @@ def _convert_g_kg_to_ppm_units(
parameter.current_set == "zonal_mean_2d_stratosphere"
and parameter.var_id == "Q"
):
ds_new[var_key] = ds_new[var_key] * 1000.0
ds_new[var_key].attrs["units"] = "ppm"
ds_new[var_key] = ds_new[var_key] * 28.97 / 18.0 * 1000.0
ds_new[var_key].attrs["units"] = "ppmv"

return ds_new

Expand Down
3 changes: 0 additions & 3 deletions tests/e3sm_diags/driver/test_tc_analysis_driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,9 +77,6 @@ def test_correct_output(self):
"vsmc": np.array([[1.94, np.nan]]),
"yearmc": np.array([[1, np.nan]]),
"monthmc": np.array([[1, np.nan]]),
"year_start": 90,
"year_end": 90,
"num_years": 1,
}
num_storms = 2 # noqa
basin_info = BASIN_DICT["NA"] # noqa
Expand Down

0 comments on commit 7550b3d

Please sign in to comment.