-
Notifications
You must be signed in to change notification settings - Fork 32
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
CDAT Migration Phase 3 - Port QBO Wavelet feature to Xarray/xCDAT codebase #860
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hey @chengzhuzhang, this PR ended up being trivial to implement (about < 1 hour of work). I copied all of the code updates from #860 and performed some additional refactoring for cleaner code.
The results look visibly close to main
. I will merge once the build passes.
main
viewer: https://portal.nersc.gov/project/e3sm/cdat-migration-fy24/main-qbo-wavelet/viewer/qbo/variable/qbo-era-interim/plot.html- dev viewer: https://portal.nersc.gov/project/e3sm/cdat-migration-fy24/850-qbo-wavelet/viewer/qbo/variable/qbo-era-interim/plot.html
- Files diff from PR Qbo wavelet #836
def _calculate_wavelet(var: xr.DataArray) -> Tuple[np.ndarray, np.ndarray]: | ||
""" | ||
Calculate the wavelet spectrum for a given data array at a specified power | ||
spectral level. | ||
|
||
Parameters | ||
---------- | ||
data : xr.DataArray | ||
The variable data. | ||
|
||
Returns | ||
------- | ||
Tuple[np.ndarray, np.ndarray] | ||
The wavelet period and wavelet array. | ||
""" | ||
# Find the closest value for power spectral level in the list | ||
test_lev = xc.get_dim_coords(var, axis="Z") | ||
test_lev_list = list(test_lev) | ||
closest_lev = min(test_lev_list, key=lambda x: abs(x - POW_SPEC_LEV)) | ||
closest_index = test_lev_list.index(closest_lev) | ||
|
||
# Grab target vertical level | ||
data_avg = var.values[:, closest_index] | ||
|
||
# Convert to anomalies | ||
data_avg = data_avg - data_avg.mean() | ||
|
||
# Detrend the data | ||
detrended_data = detrend(data_avg) | ||
|
||
wave_period, wavelet = _get_psd_from_wavelet(detrended_data) | ||
|
||
# Get square root values of wavelet spectra | ||
wavelet = np.sqrt(wavelet) | ||
|
||
return wave_period, wavelet |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Extracted this private function since the logic is the same for test and ref (look here).
Thanks GitHub Co-Pilot.
Description
cdat-migration-fy24
branch #850Checklist
If applicable: