Update ar calcs with improved parallelization#329
Merged
Conversation
…work, add zarr 3 warning catch, add optional time dimension str arg
aaTman
added a commit
that referenced
this pull request
Mar 20, 2026
* Add pressure_dimension_str arg to geopotential_thickness (#297) * `DurationMeanError` memory fix and add time resolution option (#296) * update duration with handling spatial dims, remove compute, fix sparse lead time dim generation * update name on metric in tests * add docstring for time res arg * Move parallel config check outside of function (#301) * move function out of run, move cache mkdir to init * add tests for new func * ruff * update parallel_config passthrough and tests * feat: Forecast wrapper for custom xarray datasets (#302) * implements a new Forecast object that can wrap existing xarray datasets * Revise per copilot review * Simplify IBTrACS polars subset (#303) * Update `geopotential_thickness` var names and docstring (#306) * update docstrings and var namings * rename vars, add test * ruff * Clarify default preprocess function names; geopotential division fix (#305) * update naming * default preprocess for applied_tc * ruff * ruff * Remove "cases" key requirement in yamls and dicts (#308) * remove cases top level of yaml and fix code to handle this * remove old load events yaml function * update validation precommit and formatting * remove out-of-date notebook from docs * CIRA Icechunk store (#310) * dependencies and generate store file started * in-flight, added and cleaned filter funcs * add icechunk + obstore and cira icechunk generation script * remove cira gen script no longer used * code cleanup * add icechunk datatree forecast class object * uv lock * add documentation, group helper func, and add repository kwargs passthrough * remove icechunk forecast object * typo * ruff * update pyproject and uv lock * add TODO * update PR template * Remove `IndividualCaseCollection` (#317) * update all references to IndividualCaseCollection and convert dicts/ "cases": keys to lists * update template * make questions bold * add whitespace * remove indent error and typo from evaluate_cli * make load_individual_cases include passthrough for existing dataclasses * ruff * add comment for clarification on list comp * ruff (again) * remove all references to collection, replace with list * ruff * rename collection -> list * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * Bump version from 0.2.0 to 0.3.0 (#324) * Updated API (#321) * move cache dir creation to init, rename funcs, add parallel/serial check function, update test names * update naming * add run method for backwards compatibility * update tests * add tests and cover if serial and parallel_config is not None * feat: redesign public API with hierarchical namespace submodules - Add ewb.evaluation() as main entry point (alias for ExtremeWeatherBench) - Create namespace submodules: ewb.targets, ewb.forecasts, ewb.metrics, ewb.derived, ewb.regions, ewb.cases, ewb.defaults - Expose all classes at top level for convenience (ewb.ERA5, etc.) - Add ewb.load_cases() convenience alias - Update all example files to use new import pattern - Update usage.md documentation - Maintain backward compatibility with existing imports * ruff/linting. add utils to init * add test coverage for module loading patterns * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * update defaults var refs * Golden tests (#323) * first pass for gt test infra + yaml * use shapefile for severe convection and catch latitude swap * add ignore for golden test when running pytest by default * ruff * move pytest addopts and markers to pyproject.toml * Remove `IndividualCaseCollection` (#317) * update all references to IndividualCaseCollection and convert dicts/ "cases": keys to lists * update template * make questions bold * add whitespace * remove indent error and typo from evaluate_cli * make load_individual_cases include passthrough for existing dataclasses * ruff * add comment for clarification on list comp * ruff (again) * remove all references to collection, replace with list * ruff * rename collection -> list * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * Bump version from 0.2.0 to 0.3.0 (#324) * Updated API (#321) * move cache dir creation to init, rename funcs, add parallel/serial check function, update test names * update naming * add run method for backwards compatibility * update tests * add tests and cover if serial and parallel_config is not None * feat: redesign public API with hierarchical namespace submodules - Add ewb.evaluation() as main entry point (alias for ExtremeWeatherBench) - Create namespace submodules: ewb.targets, ewb.forecasts, ewb.metrics, ewb.derived, ewb.regions, ewb.cases, ewb.defaults - Expose all classes at top level for convenience (ewb.ERA5, etc.) - Add ewb.load_cases() convenience alias - Update all example files to use new import pattern - Update usage.md documentation - Maintain backward compatibility with existing imports * ruff/linting. add utils to init * add test coverage for module loading patterns * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * update defaults var refs * remove to_csv * PyPI Preparation (#315) * update build-system and project * update workflows, publish, and pyproject * add justfile and twine * update publish yaml * change to python 3.10 as minimum requirement * kerchunk needs 3.11, swapping pyproject and tests to remove 3.10 * change workflows to use version matrix * align workflows * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * Bump version from 0.2.0 to 0.3.0 (#324) * Updated API (#321) * move cache dir creation to init, rename funcs, add parallel/serial check function, update test names * update naming * add run method for backwards compatibility * update tests * add tests and cover if serial and parallel_config is not None * feat: redesign public API with hierarchical namespace submodules - Add ewb.evaluation() as main entry point (alias for ExtremeWeatherBench) - Create namespace submodules: ewb.targets, ewb.forecasts, ewb.metrics, ewb.derived, ewb.regions, ewb.cases, ewb.defaults - Expose all classes at top level for convenience (ewb.ERA5, etc.) - Add ewb.load_cases() convenience alias - Update all example files to use new import pattern - Update usage.md documentation - Maintain backward compatibility with existing imports * ruff/linting. add utils to init * add test coverage for module loading patterns * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * update defaults var refs * Golden tests (#323) * first pass for gt test infra + yaml * use shapefile for severe convection and catch latitude swap * add ignore for golden test when running pytest by default * ruff * move pytest addopts and markers to pyproject.toml * Remove `IndividualCaseCollection` (#317) * update all references to IndividualCaseCollection and convert dicts/ "cases": keys to lists * update template * make questions bold * add whitespace * remove indent error and typo from evaluate_cli * make load_individual_cases include passthrough for existing dataclasses * ruff * add comment for clarification on list comp * ruff (again) * remove all references to collection, replace with list * ruff * rename collection -> list * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * Bump version from 0.2.0 to 0.3.0 (#324) * Updated API (#321) * move cache dir creation to init, rename funcs, add parallel/serial check function, update test names * update naming * add run method for backwards compatibility * update tests * add tests and cover if serial and parallel_config is not None * feat: redesign public API with hierarchical namespace submodules - Add ewb.evaluation() as main entry point (alias for ExtremeWeatherBench) - Create namespace submodules: ewb.targets, ewb.forecasts, ewb.metrics, ewb.derived, ewb.regions, ewb.cases, ewb.defaults - Expose all classes at top level for convenience (ewb.ERA5, etc.) - Add ewb.load_cases() convenience alias - Update all example files to use new import pattern - Update usage.md documentation - Maintain backward compatibility with existing imports * ruff/linting. add utils to init * add test coverage for module loading patterns * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * update defaults var refs * remove to_csv * swap pyproject tools to hatch; add if and packages-dir to publish * update pyproject version for release * Remove duplicate function and fixtures (#326) * chore: remove duplicate function and fixtures - Remove duplicate _parallel_serial_config_check function from evaluate.py (was defined twice at lines 189 and 982 with identical implementation) - Remove duplicate runner fixture from test_evaluate_cli.py (already defined in conftest.py) - Remove duplicate temp_config_dir fixture from test_evaluate_cli.py (already defined in conftest.py) - Remove unused tempfile import from test_evaluate_cli.py * ruff * update install docs for pypi and fix typo in link * Update ar calcs with improved parallelization (#329) * update ar calcs with improved parallelization * ruff formatting * update attr name * remove diagnostic logging; ruff formatting; update ar bounds code to work, add zarr 3 warning catch, add optional time dimension str arg * more ar_bounds work to get it moving, looks good here * mypy * mypy * update time to valid_time to follow ewb conventions --------- Co-authored-by: Daniel Rothenberg <daniel@danielrothenberg.com>
aaTman
added a commit
that referenced
this pull request
Mar 20, 2026
* Add pressure_dimension_str arg to geopotential_thickness (#297) * `DurationMeanError` memory fix and add time resolution option (#296) * update duration with handling spatial dims, remove compute, fix sparse lead time dim generation * update name on metric in tests * add docstring for time res arg * Move parallel config check outside of function (#301) * move function out of run, move cache mkdir to init * add tests for new func * ruff * update parallel_config passthrough and tests * feat: Forecast wrapper for custom xarray datasets (#302) * implements a new Forecast object that can wrap existing xarray datasets * Revise per copilot review * Simplify IBTrACS polars subset (#303) * Update `geopotential_thickness` var names and docstring (#306) * update docstrings and var namings * rename vars, add test * ruff * Clarify default preprocess function names; geopotential division fix (#305) * update naming * default preprocess for applied_tc * ruff * ruff * Remove "cases" key requirement in yamls and dicts (#308) * remove cases top level of yaml and fix code to handle this * remove old load events yaml function * update validation precommit and formatting * remove out-of-date notebook from docs * CIRA Icechunk store (#310) * dependencies and generate store file started * in-flight, added and cleaned filter funcs * add icechunk + obstore and cira icechunk generation script * remove cira gen script no longer used * code cleanup * add icechunk datatree forecast class object * uv lock * add documentation, group helper func, and add repository kwargs passthrough * remove icechunk forecast object * typo * ruff * update pyproject and uv lock * add TODO * update PR template * Remove `IndividualCaseCollection` (#317) * update all references to IndividualCaseCollection and convert dicts/ "cases": keys to lists * update template * make questions bold * add whitespace * remove indent error and typo from evaluate_cli * make load_individual_cases include passthrough for existing dataclasses * ruff * add comment for clarification on list comp * ruff (again) * remove all references to collection, replace with list * ruff * rename collection -> list * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * Bump version from 0.2.0 to 0.3.0 (#324) * Updated API (#321) * move cache dir creation to init, rename funcs, add parallel/serial check function, update test names * update naming * add run method for backwards compatibility * update tests * add tests and cover if serial and parallel_config is not None * feat: redesign public API with hierarchical namespace submodules - Add ewb.evaluation() as main entry point (alias for ExtremeWeatherBench) - Create namespace submodules: ewb.targets, ewb.forecasts, ewb.metrics, ewb.derived, ewb.regions, ewb.cases, ewb.defaults - Expose all classes at top level for convenience (ewb.ERA5, etc.) - Add ewb.load_cases() convenience alias - Update all example files to use new import pattern - Update usage.md documentation - Maintain backward compatibility with existing imports * ruff/linting. add utils to init * add test coverage for module loading patterns * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * update defaults var refs * Golden tests (#323) * first pass for gt test infra + yaml * use shapefile for severe convection and catch latitude swap * add ignore for golden test when running pytest by default * ruff * move pytest addopts and markers to pyproject.toml * Remove `IndividualCaseCollection` (#317) * update all references to IndividualCaseCollection and convert dicts/ "cases": keys to lists * update template * make questions bold * add whitespace * remove indent error and typo from evaluate_cli * make load_individual_cases include passthrough for existing dataclasses * ruff * add comment for clarification on list comp * ruff (again) * remove all references to collection, replace with list * ruff * rename collection -> list * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * Bump version from 0.2.0 to 0.3.0 (#324) * Updated API (#321) * move cache dir creation to init, rename funcs, add parallel/serial check function, update test names * update naming * add run method for backwards compatibility * update tests * add tests and cover if serial and parallel_config is not None * feat: redesign public API with hierarchical namespace submodules - Add ewb.evaluation() as main entry point (alias for ExtremeWeatherBench) - Create namespace submodules: ewb.targets, ewb.forecasts, ewb.metrics, ewb.derived, ewb.regions, ewb.cases, ewb.defaults - Expose all classes at top level for convenience (ewb.ERA5, etc.) - Add ewb.load_cases() convenience alias - Update all example files to use new import pattern - Update usage.md documentation - Maintain backward compatibility with existing imports * ruff/linting. add utils to init * add test coverage for module loading patterns * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * update defaults var refs * remove to_csv * PyPI Preparation (#315) * update build-system and project * update workflows, publish, and pyproject * add justfile and twine * update publish yaml * change to python 3.10 as minimum requirement * kerchunk needs 3.11, swapping pyproject and tests to remove 3.10 * change workflows to use version matrix * align workflows * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * Bump version from 0.2.0 to 0.3.0 (#324) * Updated API (#321) * move cache dir creation to init, rename funcs, add parallel/serial check function, update test names * update naming * add run method for backwards compatibility * update tests * add tests and cover if serial and parallel_config is not None * feat: redesign public API with hierarchical namespace submodules - Add ewb.evaluation() as main entry point (alias for ExtremeWeatherBench) - Create namespace submodules: ewb.targets, ewb.forecasts, ewb.metrics, ewb.derived, ewb.regions, ewb.cases, ewb.defaults - Expose all classes at top level for convenience (ewb.ERA5, etc.) - Add ewb.load_cases() convenience alias - Update all example files to use new import pattern - Update usage.md documentation - Maintain backward compatibility with existing imports * ruff/linting. add utils to init * add test coverage for module loading patterns * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * update defaults var refs * Golden tests (#323) * first pass for gt test infra + yaml * use shapefile for severe convection and catch latitude swap * add ignore for golden test when running pytest by default * ruff * move pytest addopts and markers to pyproject.toml * Remove `IndividualCaseCollection` (#317) * update all references to IndividualCaseCollection and convert dicts/ "cases": keys to lists * update template * make questions bold * add whitespace * remove indent error and typo from evaluate_cli * make load_individual_cases include passthrough for existing dataclasses * ruff * add comment for clarification on list comp * ruff (again) * remove all references to collection, replace with list * ruff * rename collection -> list * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * Bump version from 0.2.0 to 0.3.0 (#324) * Updated API (#321) * move cache dir creation to init, rename funcs, add parallel/serial check function, update test names * update naming * add run method for backwards compatibility * update tests * add tests and cover if serial and parallel_config is not None * feat: redesign public API with hierarchical namespace submodules - Add ewb.evaluation() as main entry point (alias for ExtremeWeatherBench) - Create namespace submodules: ewb.targets, ewb.forecasts, ewb.metrics, ewb.derived, ewb.regions, ewb.cases, ewb.defaults - Expose all classes at top level for convenience (ewb.ERA5, etc.) - Add ewb.load_cases() convenience alias - Update all example files to use new import pattern - Update usage.md documentation - Maintain backward compatibility with existing imports * ruff/linting. add utils to init * add test coverage for module loading patterns * ruff * Cleanup docstrings in repo (#318) * update these docstrings * remove docstring changes markdown * update docstrings * update other docstrings * remove individualcasecollection reference, update based on develop changes * add explanation for dim reqs (#320) * Update `defaults` and `inputs` to include new CIRA icechunk store (#319) * more explicit naming, add func and model names var * add test coverage, ruff, linting * update readme for new cira approach * move cira func and model ref to inputs * update docs * module wasnt called for moved func * update tests for moving func and var * ruff * fix mock typos * update defaults var refs * remove to_csv * swap pyproject tools to hatch; add if and packages-dir to publish * update pyproject version for release * Remove duplicate function and fixtures (#326) * chore: remove duplicate function and fixtures - Remove duplicate _parallel_serial_config_check function from evaluate.py (was defined twice at lines 189 and 982 with identical implementation) - Remove duplicate runner fixture from test_evaluate_cli.py (already defined in conftest.py) - Remove duplicate temp_config_dir fixture from test_evaluate_cli.py (already defined in conftest.py) - Remove unused tempfile import from test_evaluate_cli.py * ruff * update install docs for pypi and fix typo in link * Update ar calcs with improved parallelization (#329) * update ar calcs with improved parallelization * ruff formatting * update attr name * remove diagnostic logging; ruff formatting; update ar bounds code to work, add zarr 3 warning catch, add optional time dimension str arg * more ar_bounds work to get it moving, looks good here * mypy * mypy * update time to valid_time to follow ewb conventions * bump pyproject version (#331) --------- Co-authored-by: Daniel Rothenberg <daniel@danielrothenberg.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
EWB Pull Request
Description
The atmospheric river calculations were significantly slower than expected when applied; this PR applies a few optimizations using numba and apply_ufunc in certain locations to speed up the calculations.
Type of change
How Has This Been Tested?
Tests were updated to reflect some of the backend functions with slightly different argument patterns. All tests were run to confirm functionality.