-
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 2: Refactor arm_diags set #842
CDAT Migration Phase 2: Refactor arm_diags set #842
Conversation
Updates: Todo: |
7b52e83
to
ef44fc6
Compare
Thank you for identifying the two formulas code as potential bottlenecks. I confirmed that these computations with Xarray are indeed much slower than CDAT. SolutionI ran a performance benchmark and found the solution to the slowness: we need to call Benchmark ResultsThe first runtime is the current code and the second runtime is with """
Results
----------
1. Elapsed time (Xarray non-chunked): 6.540755605790764 seconds
2. Elapsed time (Xarray non-chunked with .load()): 0.17097265785560012 seconds
3. Elapsed time (Xarray chunked): 0.1452920027077198 seconds
4. Elapsed time (numpy .values): 6.418793010059744 seconds
5. Elapsed time (numpy .data): 7.334999438840896 seconds
""" Elapsed time (CDAT main branch, single runtime): 0.12261438369750977 seconds |
@tomvothecoder thank you for performing the timing test. Interesting that when I test your commit. Running through a configuration with each subset, it took 24 mins. Without, |
I would think that loading the derived variables dataset into memory shouldn't slow down performance unless the datasets were extremely large (which we should use Dask chunking for). I only benchmarked performance for the formula computations. I will benchmark a complete run to verify your findings and determine if we should revert the commit or not. Side-note: It could be that the logic I implemented already stores the dataset in-memory, since it merges multiple xr.Dataset objects opened via e3sm_diags/e3sm_diags/driver/utils/dataset_xr.py Lines 1056 to 1064 in ef44fc6
|
RE: #842 (comment) I reverted this change and will now address the
|
Hey @chengzhuzhang, I fixed the pre-commit issues in I re-ran all sets and they completed successfully with these changes, although I noticed some of the diagnostic sets weren't done yet. |
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.
Initial review comment and some questions of FIXME comments.
Thank you @tomvothecoder I will do another pass to see if there are other sets need to be refactored. Thanks a lot for fixing and clean this branch. |
@tomvothecoder I applied |
Hey @chengzhuzhang, I noticed you have several other PRs you're working on right now. I'm happy to help finish up refactoring the last diag function in this PR. Let me know. |
@tomvothecoder it would be great if you could help me on finish this last diag, so that I'm not holding back the progress to merge! |
I found that no sets actually run |
@tomvothecoder good for catching this. I think for now we can just create an issue to log this problem, and we can add back the code at a later time. |
@@ -152,6 +152,7 @@ def climo(dataset: xr.Dataset, var_key: str, freq: ClimoFreq): | |||
# averaging. | |||
dims = [dim for dim in dv.dims if dim != time_coords.name] | |||
coords = {k: v for k, v in dv.coords.items() if k in dims} | |||
climo = climo.squeeze(axis=0) |
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.
Fixes bug when ncycle == 1
where the climo
variable time axis is not being squeezed which causes the rebuilt dv_climo
DataArray to fail with
FAILED tests/e3sm_diags/driver/utils/test_climo_xr.py::TestClimo::test_returns_annual_cycle_climatology - ValueError: different number of dimensions on data and dims: 3 vs 2
FAILED tests/e3sm_diags/driver/utils/test_climo_xr.py::TestClimo::test_returns_DJF_season_climatology - ValueError: different number of dimensions on data and dims: 3 vs 2
FAILED tests/e3sm_diags/driver/utils/test_climo_xr.py::TestClimo::test_returns_MAM_season_climatology - ValueError: different number of dimensions on data and dims: 3 vs 2
FAILED tests/e3sm_diags/driver/utils/test_climo_xr.py::TestClimo::test_returns_JJA_season_climatology - ValueError: different number of dimensions on data and dims: 3 vs 2
FAILED tests/e3sm_diags/driver/utils/test_climo_xr.py::TestClimo::test_returns_SON_season_climatology - ValueError: different number of dimensions on data and dims: 3 vs 2
FAILED tests/e3sm_diags/driver/utils/test_climo_xr.py::TestClimo::test_returns_jan_climatology - ValueError: different number of dimensions on data and dims: 3 vs 2
FAILED tests/e3sm_diags/driver/utils/test_climo_xr.py::TestClimo::test_returns_climatology_for_derived_variable - ValueError: different number of dimensions on data and dims: 3 vs 2
FAILED tests/e3sm_diags/driver/utils/test_dataset_xr.py::TestGetClimoDataset::test_returns_climo_dataset_using_climo_of_time_series_files - ValueError: different number of dimensions on data and dims: 3 vs 2
Regression test results
Comparing:
* /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/arm_diags/armdiags-CLOUD-ANNUALCYCLE-nsac1-test.png
* /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/667-arm_diags-final/arm_diags/armdiags-CLOUD-ANNUALCYCLE-nsac1-test.png
* Difference path /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/667-arm_diags-final/arm_diags_diff/armdiags-CLOUD-ANNUALCYCLE-nsac1-test.png
Comparing:
* /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/arm_diags/armdiags-CLOUD-ANNUALCYCLE-sgpc1-test.png
* /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/667-arm_diags-final/arm_diags/armdiags-CLOUD-ANNUALCYCLE-sgpc1-test.png
* Difference path /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/667-arm_diags-final/arm_diags_diff/armdiags-CLOUD-ANNUALCYCLE-sgpc1-test.png
Comparing:
* /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/arm_diags/armdiags-CLOUD-ANNUALCYCLE-twpc1-test.png
* /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/667-arm_diags-final/arm_diags/armdiags-CLOUD-ANNUALCYCLE-twpc1-test.png
* Difference path /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/667-arm_diags-final/arm_diags_diff/armdiags-CLOUD-ANNUALCYCLE-twpc1-test.png
Comparing:
* /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/arm_diags/armdiags-CLOUD-ANNUALCYCLE-twpc2-test.png
* /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/667-arm_diags-final/arm_diags/armdiags-CLOUD-ANNUALCYCLE-twpc2-test.png
* Difference path /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/667-arm_diags-final/arm_diags_diff/armdiags-CLOUD-ANNUALCYCLE-twpc2-test.png
Comparing:
* /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/arm_diags/armdiags-CLOUD-ANNUALCYCLE-twpc3-test.png
* /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/667-arm_diags-final/arm_diags/armdiags-CLOUD-ANNUALCYCLE-twpc3-test.png
* Difference path /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/667-arm_diags-final/arm_diags_diff/armdiags-CLOUD-ANNUALCYCLE-twpc3-test.png |
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.
Remaining TODO and FIXME comments in arm_diags_driver.py
- Refactor `arm_diags_driver.py`- rename functions, reorder functions, rename variables for clarity, replace `.format` with f-strings, update docstrings, add `_select_point()` - Refactor `climo_xr.py` - extract `_get_cycle_for_freq()`, remove commented out code - Add `time_interval` to `ARMDiagsParameter`
- Rename functions to denote private and reorder based on call in `arm_diags_driver.py` - Add typestrings and annotations - Separate logically related blocks of code with comments - Add `_save_plots()` function to replace repeated I/O across functions
- This behavior mimics the co flag found in the CDAT codebase
18357c5
to
7b53932
Compare
def _exclude_sub_monthly_coord_spanning_year( | ||
self, ds_subset: xr.Dataset | ||
) -> xr.Dataset: | ||
""" | ||
Exclude the last time coordinate for sub-monthly data if it extends into | ||
the next year. | ||
|
||
Excluding end time coordinates that extend to the next year is | ||
necessary because downstream operations such as annual cycle climatology | ||
should consist of data for full years for accurate calculations. | ||
|
||
For example, if the time slice is ("0001-01-01", "0002-01-01") and | ||
the last time coordinate is: | ||
* "0002-01-01" -> exclude | ||
* "0001-12-31" -> don't exclude | ||
|
||
Parameters | ||
---------- | ||
ds_subset : xr.Dataset | ||
The subsetted dataset. | ||
|
||
Returns | ||
------- | ||
xr.Dataset | ||
The dataset with the last time coordinate excluded if necessary. | ||
|
||
Notes | ||
----- | ||
This function replicates the CDAT cdms2 "co" slice flag (close, open). | ||
""" | ||
time_dim = xc.get_dim_keys(ds_subset, axis="T") | ||
time_values = ds_subset[time_dim] | ||
last_time_year = time_values[-1].dt.year.item() | ||
second_last_time_year = time_values[-2].dt.year.item() | ||
|
||
if self.is_sub_monthly and last_time_year > second_last_time_year: | ||
ds_subset = ds_subset.isel(time=slice(0, -1)) | ||
|
||
return ds_subset | ||
|
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.
Replicates the "co" slice flag for sub-monthly data
This PR should be good to go. The integration test image checker fails because there is now one additional plot for the I will merge this PR even though the integration tests fail. I will update the complete test results on LCRC using the |
Description
arm_diags
set #667 to replace CDAT Migration Phase 2: Refactor arm_diags set #834Additional Changes
arm_diags_driver.py
andarm_diags_viewer.py
derivations.py
andformulas.py
arm_diags_plot.py
sgpc1
plot not being generated correctly (Closes [Bug]:UnboundLocalError: local variable 'cwv_max' referenced before assignment
for arm_diags plot with regionsgpc1
#858)arm_diags_parameter.py
time_interval
attribute to store hours for diurnal cycledataset_xr.py
"co"
flag from cdms2, which excludes the last time coordinate when subsetting if that time coordinate spans the next year (Dataset._exclude_sub_monthly_coord_spanning_year()
)diurnal_cycle_xr.py
diurnal_cycle_xr.py
to convert coordinate DataArrays to lists and convert for loop to list comprehension for performance gainsChecklist
If applicable: