Updated API

data_prep/ar_bounds.py

@@ -17,7 +17,11 @@
 from dask.distributed import Client
 from matplotlib.patches import Rectangle
 
-from extremeweatherbench import cases, derived, inputs, regions, utils
+import extremeweatherbench.cases as cases
+import extremeweatherbench.derived as derived
+import extremeweatherbench.inputs as inputs
+import extremeweatherbench.regions as regions
+import extremeweatherbench.utils as utils
 from extremeweatherbench.events import atmospheric_river as ar
 
 logging.basicConfig()

data_prep/ibtracs_bounds.py

@@ -4,6 +4,7 @@
 import logging
 import re
 from importlib import resources
+from typing import TYPE_CHECKING
 
 import cartopy.crs as ccrs
 import cartopy.feature as cfeature
@@ -14,8 +15,11 @@
 import yaml
 from matplotlib.patches import Rectangle
 
+import extremeweatherbench as ewb
 import extremeweatherbench.data
-from extremeweatherbench import cases, inputs, regions, utils
+
+if TYPE_CHECKING:
+    from extremeweatherbench.regions import Region
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -67,7 +71,7 @@ def calculate_extent_bounds(
     bottom_lat: float,
     top_lat: float,
     extent_buffer: float = 250,
-) -> regions.Region:
+) -> Region:
     """Calculate extent bounds with buffer.
 
     Args:
@@ -94,9 +98,9 @@ def calculate_extent_bounds(
         calculate_end_point(bottom_lat, right_lon, 90, extent_buffer), 1
     )
 
-    new_left_lon = np.round(utils.convert_longitude_to_360(new_left_lon), 1)
-    new_right_lon = np.round(utils.convert_longitude_to_360(new_right_lon), 1)
-    new_box = regions.BoundingBoxRegion(
+    new_left_lon = np.round(ewb.utils.convert_longitude_to_360(new_left_lon), 1)
+    new_right_lon = np.round(ewb.utils.convert_longitude_to_360(new_right_lon), 1)
+    new_box = ewb.regions.BoundingBoxRegion(
         new_bottom_lat, new_top_lat, new_left_lon, new_right_lon
     )
     return new_box
@@ -164,10 +168,10 @@ def load_and_process_ibtracs_data():
     """
     logger.info("Loading IBTrACS data...")
 
-    IBTRACS = inputs.IBTrACS(
-        source=inputs.IBTRACS_URI,
+    IBTRACS = ewb.inputs.IBTrACS(
+        source=ewb.inputs.IBTRACS_URI,
         variables=["vmax", "slp"],
-        variable_mapping=inputs.IBTrACS_metadata_variable_mapping,
+        variable_mapping=ewb.inputs.IBTrACS_metadata_variable_mapping,
         storage_options={},
     )
 
@@ -177,7 +181,7 @@ def load_and_process_ibtracs_data():
     # Get all storms from 2020 - 2025 seasons
     all_storms_2020_2025_lf = IBTRACS_lf.filter(
         (pl.col("SEASON").cast(pl.Int32) >= 2020)
-    ).select(inputs.IBTrACS_metadata_variable_mapping.values())
+    ).select(ewb.inputs.IBTrACS_metadata_variable_mapping.values())
 
     schema = all_storms_2020_2025_lf.collect_schema()
     # Convert pressure and surface wind columns to float, replacing " " with null
@@ -464,7 +468,7 @@ def find_storm_bounds_for_case(storm_name, storm_bounds, all_storms_df):
             # If we found both, merge them by taking the bounding box that
             # encompasses both
             if bounds1 is not None and bounds2 is not None:
-                merged_bbox = regions.BoundingBoxRegion(
+                merged_bbox = ewb.regions.BoundingBoxRegion(
                     latitude_min=min(
                         bounds1.iloc[0].latitude_min, bounds2.iloc[0].latitude_min
                     ),
@@ -537,7 +541,7 @@ def update_cases_with_storm_bounds(storm_bounds, all_storms_df):
     """
     logger.info("Updating cases with storm bounds...")
 
-    cases_all = cases.load_ewb_events_yaml_into_case_list()
+    cases_all = ewb.cases.load_ewb_events_yaml_into_case_list()
     cases_new = cases_all.copy()
 
     # Update the yaml cases with storm bounds from IBTrACS data

data_prep/practically_perfect_hindcast_from_lsr.py

@@ -11,7 +11,8 @@
 from scipy.ndimage import gaussian_filter
 from tqdm.auto import tqdm
 
-from extremeweatherbench import inputs, utils
+import extremeweatherbench.inputs as inputs
+import extremeweatherbench.utils as utils
 
 
 def sparse_practically_perfect_hindcast(

data_prep/severe_convection_bounds.py

@@ -17,7 +17,8 @@
 import yaml
 from scipy.ndimage import label
 
-from extremeweatherbench import calc, cases
+import extremeweatherbench.calc as calc
+import extremeweatherbench.cases as cases
 
 # Radius of Earth in km (mean radius)
 EARTH_RADIUS_KM = 6371.0

data_prep/subset_heat_cold_events.py

@@ -13,7 +13,8 @@
 from matplotlib import dates as mdates
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 
-from extremeweatherbench import cases, utils
+import extremeweatherbench.cases as cases
+import extremeweatherbench.utils as utils
 
 sns.set_theme(style="whitegrid", context="talk")
 

docs/examples/applied_ar.py

@@ -4,7 +4,7 @@
 import numpy as np
 import xarray as xr
 
-from extremeweatherbench import cases, derived, evaluate, inputs, metrics
+import extremeweatherbench as ewb
 
 # %%
 
@@ -38,85 +38,86 @@ def _preprocess_cira_forecast_dataset(ds: xr.Dataset) -> xr.Dataset:
 
 
 # Load case data from the default events.yaml
-# Users can also define their own cases
-case_yaml = cases.load_ewb_events_yaml_into_case_list()
-case_yaml = [n for n in case_yaml if n.case_id_number == 114]
-case_yaml[0].start_date = datetime.datetime(2022, 12, 27, 11, 0, 0)
-case_yaml[0].end_date = datetime.datetime(2022, 12, 27, 13, 0, 0)
+# Users can also define their own cases_dict structure
+case_yaml = ewb.load_cases()
+case_list = [n for n in case_yaml if n.case_id_number == 114]
+
+case_list[0].start_date = datetime.datetime(2022, 12, 27, 11, 0, 0)
+case_list[0].end_date = datetime.datetime(2022, 12, 27, 13, 0, 0)
 # Define ERA5 target
-era5_target = inputs.ERA5(
+era5_target = ewb.targets.ERA5(
     variables=[
-        derived.AtmosphericRiverVariables(
+        ewb.derived.AtmosphericRiverVariables(
             output_variables=["atmospheric_river_land_intersection"]
         )
     ],
 )
 
 # Define forecast (HRES)
-hres_forecast = inputs.ZarrForecast(
+hres_forecast = ewb.forecasts.ZarrForecast(
     source="gs://weatherbench2/datasets/hres/2016-2022-0012-1440x721.zarr",
     name="HRES",
     variables=[
-        derived.AtmosphericRiverVariables(
+        ewb.derived.AtmosphericRiverVariables(
             output_variables=["atmospheric_river_land_intersection"]
         )
     ],
-    variable_mapping=inputs.HRES_metadata_variable_mapping,
+    variable_mapping=ewb.HRES_metadata_variable_mapping,
 )
 
-grap_forecast = inputs.KerchunkForecast(
+grap_forecast = ewb.forecasts.KerchunkForecast(
     name="Graphcast",
     source="gs://extremeweatherbench/GRAP_v100_IFS.parq",
     variables=[
-        derived.AtmosphericRiverVariables(
+        ewb.derived.AtmosphericRiverVariables(
             output_variables=["atmospheric_river_land_intersection"]
         )
     ],
-    variable_mapping=inputs.CIRA_metadata_variable_mapping,
+    variable_mapping=ewb.CIRA_metadata_variable_mapping,
     storage_options={"remote_protocol": "s3", "remote_options": {"anon": True}},
     preprocess=_preprocess_cira_forecast_dataset,
 )
 
-pang_forecast = inputs.KerchunkForecast(
+pang_forecast = ewb.forecasts.KerchunkForecast(
     name="Pangu",
     source="gs://extremeweatherbench/PANG_v100_IFS.parq",
     variables=[
-        derived.AtmosphericRiverVariables(
+        ewb.derived.AtmosphericRiverVariables(
             output_variables=["atmospheric_river_land_intersection"]
         )
     ],
-    variable_mapping=inputs.CIRA_metadata_variable_mapping,
+    variable_mapping=ewb.CIRA_metadata_variable_mapping,
     storage_options={"remote_protocol": "s3", "remote_options": {"anon": True}},
     preprocess=_preprocess_cira_forecast_dataset,
 )
 # Create a list of evaluation objects for atmospheric river
 ar_evaluation_objects = [
-    inputs.EvaluationObject(
+    ewb.EvaluationObject(
         event_type="atmospheric_river",
         metric_list=[
-            metrics.CriticalSuccessIndex(),
-            metrics.EarlySignal(),
-            metrics.SpatialDisplacement(),
+            ewb.metrics.CriticalSuccessIndex(),
+            ewb.metrics.EarlySignal(),
+            ewb.metrics.SpatialDisplacement(),
         ],
         target=era5_target,
         forecast=hres_forecast,
     ),
-    inputs.EvaluationObject(
+    ewb.EvaluationObject(
         event_type="atmospheric_river",
         metric_list=[
-            metrics.CriticalSuccessIndex(),
-            metrics.EarlySignal(),
-            metrics.SpatialDisplacement(),
+            ewb.metrics.CriticalSuccessIndex(),
+            ewb.metrics.EarlySignal(),
+            ewb.metrics.SpatialDisplacement(),
         ],
         target=era5_target,
         forecast=grap_forecast,
     ),
-    inputs.EvaluationObject(
+    ewb.EvaluationObject(
         event_type="atmospheric_river",
         metric_list=[
-            metrics.CriticalSuccessIndex(),
-            metrics.EarlySignal(),
-            metrics.SpatialDisplacement(),
+            ewb.metrics.CriticalSuccessIndex(),
+            ewb.metrics.EarlySignal(),
+            ewb.metrics.SpatialDisplacement(),
         ],
         target=era5_target,
         forecast=pang_forecast,
@@ -126,7 +127,7 @@ def _preprocess_cira_forecast_dataset(ds: xr.Dataset) -> xr.Dataset:
 if __name__ == "__main__":
     # Initialize ExtremeWeatherBench; will only run on cases with event_type
     # atmospheric_river
-    ar_ewb = evaluate.ExtremeWeatherBench(
+    ar_ewb = ewb.evaluation(
         case_metadata=case_yaml,
         evaluation_objects=ar_evaluation_objects,
     )

docs/examples/applied_freeze.py

@@ -1,55 +1,55 @@
 import logging
 import operator
 
-from extremeweatherbench import cases, defaults, evaluate, inputs, metrics
+import extremeweatherbench as ewb
 
 # Set the logger level to INFO
 logger = logging.getLogger("extremeweatherbench")
 logger.setLevel(logging.INFO)
 
 # Load case data from the default events.yaml
-# Users can also define their own cases
-case_yaml = cases.load_ewb_events_yaml_into_case_list()
+# Users can also define their own cases_dict structure
+case_yaml = ewb.load_cases()
 
 # Define targets
 # ERA5 target
-era5_freeze_target = inputs.ERA5(
+era5_freeze_target = ewb.targets.ERA5(
     variables=["surface_air_temperature"],
     chunks=None,
 )
 
 # GHCN target
-ghcn_freeze_target = inputs.GHCN(variables=["surface_air_temperature"])
+ghcn_freeze_target = ewb.targets.GHCN(variables=["surface_air_temperature"])
 
 # Define forecast (FCNv2 CIRA Virtualizarr)
-fcnv2_forecast = inputs.KerchunkForecast(
+fcnv2_forecast = ewb.forecasts.KerchunkForecast(
     name="fcnv2_forecast",
     source="gs://extremeweatherbench/FOUR_v200_GFS.parq",
     variables=["surface_air_temperature"],
-    variable_mapping=inputs.CIRA_metadata_variable_mapping,
+    variable_mapping=ewb.CIRA_metadata_variable_mapping,
     storage_options={"remote_protocol": "s3", "remote_options": {"anon": True}},
-    preprocess=defaults._preprocess_cira_forecast_dataset,
+    preprocess=ewb.defaults._preprocess_bb_cira_forecast_dataset,
 )
 
 # Load the climatology for DurationMeanError
-climatology = defaults.get_climatology(quantile=0.85)
+climatology = ewb.get_climatology(quantile=0.85)
 
 # Define the metrics
 metrics_list = [
-    metrics.RootMeanSquaredError(),
-    metrics.MinimumMeanAbsoluteError(),
-    metrics.DurationMeanError(threshold_criteria=climatology, op_func=operator.le),
+    ewb.metrics.RootMeanSquaredError(),
+    ewb.metrics.MinimumMeanAbsoluteError(),
+    ewb.metrics.DurationMeanError(threshold_criteria=climatology, op_func=operator.le),
 ]
 
 # Create a list of evaluation objects for freeze
 freeze_evaluation_object = [
-    inputs.EvaluationObject(
+    ewb.EvaluationObject(
         event_type="freeze",
         metric_list=metrics_list,
         target=ghcn_freeze_target,
         forecast=fcnv2_forecast,
     ),
-    inputs.EvaluationObject(
+    ewb.EvaluationObject(
         event_type="freeze",
         metric_list=metrics_list,
         target=era5_freeze_target,
@@ -59,13 +59,13 @@
 
 if __name__ == "__main__":
     # Initialize ExtremeWeatherBench runner instance
-    ewb = evaluate.ExtremeWeatherBench(
+    freeze_ewb = ewb.evaluation(
         case_metadata=case_yaml,
         evaluation_objects=freeze_evaluation_object,
     )
 
     # Run the workflow
-    outputs = ewb.run(parallel_config={"backend": "loky", "n_jobs": 1})
+    outputs = freeze_ewb.run(parallel_config={"backend": "loky", "n_jobs": 1})
 
     # Print the outputs; can be saved if desired
     outputs.to_csv("freeze_outputs.csv")