diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 0000000..887a2c1
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,2 @@
+# SCM syntax highlighting & preventing 3-way merges
+pixi.lock merge=binary linguist-language=YAML linguist-generated=true
diff --git a/.gitignore b/.gitignore
index f48095c..093f6ca 100644
--- a/.gitignore
+++ b/.gitignore
@@ -57,7 +57,7 @@ coverage.xml
# Django stuff:
*.log
-local_settings.py
+local_data_utils.py
db.sqlite3
db.sqlite3-journal
@@ -130,3 +130,12 @@ dmypy.json
# For PyCharm
.idea
+# pixi environments
+.pixi/*
+!.pixi/config.toml
+
+# wandb
+wandb/
+
+# osx files
+.DS_Store
\ No newline at end of file
diff --git a/README.md b/README.md
index a35a093..9fcfae6 100644
--- a/README.md
+++ b/README.md
@@ -1 +1,435 @@
-# fronts
\ No newline at end of file
+# AI2ES Fronts Project Training Guide (DO NOT USE - this guide is outdated)
+
+The following guide will detail the steps toward successfully train a UNET-style model for frontal boundary predictions.
+
+## Table of Contents
+
+1. Gathering Data
+2. TensorFlow dataset
+3. Model training
+4. Evaluation
+5. XAI* (not written)
+6. Appendix
+
+###### Note that many file-naming conventions and expected directory structures are hard-coded and must be followed for this module to function properly.
+
+## 1. Gathering Data
+
+There are two types of data that need to be gathered for successful model training: front labels (targets) and
+predictors (inputs). The predictors will normally be sourced from ERA5 data or an NWP model (GFS, ECMWF, etc.).
+
+#### 1a. Front labels
+
+Front labels are sourced from the National Oceanic and Atmospheric Administration (NOAA) in the form of XML files, or
+from The Weather Company (TWC) in the form of GML files.
+
+* **NOAA XML file-naming format:** *pres_pmsl_YYYYMMDDHHf000.xml*
+ * Ex: *pres_pmsl_2016062115f000.xml* [15z June 21, 2016 front analysis]
+* **TWC GML file-naming formats and directory structures:**
+ * Current analysis: */YYYYMMDD/HH/rec.sfcanalysis.YYYYMMDDTHH0000Z.NIL.P0D.WORLD@10km.FRONTS.SFC.gml*
+ * Ex:
+ */20230827/00/rec.sfcanalysis.20230827T090000Z.NIL.P0D.WORLD@10km.FRONTS.SFC.gml* [9z August 27, 2023 analysis]
+ * Forecasted fronts:
+ */YYYYMMDD/HH/rec.sfcanalysis.YYYYMMDDTHH0000Z.YYYYMMDDTHH0000Z.P01D00H00M.WORLD@10km.FRONTS.SFC.gml*
+ * Ex:
+ */20230827/00/rec.sfcanalysis.20230828T000000Z.20230827T000000Z.P01D00H00M.WORLD@10km.FRONTS.SFC.gml* [Forecasted fronts for 00z August 28, 2023 (valid time) drawn at 00z August 27, 2023 (init time).]
+ * ###### Note that the valid and initialization time strings (e.g. "20230828T000000Z.20230827T000000Z", "20230827T090000Z.NIL") are the only required parts of the base filename outside the .gml suffix. The directory structure must be maintained.
+
+If using TWC front labels, the GML files must be converted to XML files by running the convert_front_gml_to_xml.py
+script with the command below. Only **required** arguments are shown in the command line. If an argument shows up in the
+table below the command but is not in the command itself, that argument is **optional**.
+
+ python convert_front_gml_to_xml.py --gml_indir {} --xml_outdir {} --date {} {} {}
+
+| Argument | Type | Default | Description |
+|--------------|---------|---------|----------------------------------------|
+| *gml_indir* | str | (none) | Input directory for the TWC GML files. |
+| *xml_outdir* | str | (none) | Output directory for the XML files. |
+| *date* | int (3) | (none) | Year, month, and day. |
+
+After obtaining the XML files, convert them to netCDF with the command below. *convert_front_xml_to_netcdf.py* will
+generate a netCDF for each XML file containing an initialization time with the provided date. (i.e. separate netCDF
+files will be created for 03z 2019-05-20 and 06z 2019-05-20 if the provided date is 2019-05-20)
+
+ python convert_front_xml_to_netcdf.py --xml_indir {} --netcdf_outdir {} --date {}
+
+| Argument | Type | Default | Description |
+|-----------------|--------------|---------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| *xml_indir* | str | (none) | Input directory for the XML files. |
+| *netcdf_outdir* | str | (none) | Output parent directory for the netCDF files. |
+| *date* | str | (none) | String formatted as YYYY-MM-DD |
+| *distance* | int or float | 1 | Interpolation distance for the fronts in kilometers. |
+| *domain* | str | "full" | Domain for which to interpolate fronts over. To process TWC fronts, this must be set to 'global'. To transform fronts onto the grid of a model (e.g., HRRR), type the model string in lowercase (e.g. 'hrrr') |
+
+* The resulting netCDF files will be placed in subdirectories sorted by month (e.g. */netcdf/201304* contains all netCDF files for April 2013).
+
+#### 1b. Predictor variables
+
+Predictor variables can be obtained from multiple sources, however the main source used in training is ERA5.
+
+###### NOTE: This module currently does not support ERA5 data outside of NOAA's unified surface analysis domain. Downloaded ERA5 data will automatically be sliced to match this domain. We plan to support global ERA5 data in a future version.
+
+* When downloading ERA5 data, make sure each file contains one year of 3-hourly data. Keep all data in a directory with
+ two folders named *Surface* and *Pressure_Level*. (e.g. /data/era5/Surface, /data/era5/Pressure_Level)
+ * At the **surface** level, download air temperature, dewpoint temperature, u-wind and v-wind, and surface pressure,
+ with one file per variable. The base filename for 2-meter (surface) temperature data from 2008 will be
+ *ERA5Global_2008_3hrly_2mT.nc*. Keep all surface files in the *Surface* folder as described above. There should be
+ **five** ERA5 files for each year of surface data.
+ * **Pressure level** data is downloaded in the same manner as above, however all pressure levels are contained
+ within a single file. The pressure level variables needed are temperature, u-wind and v-wind, specific humidity,
+ and geopotential height. The base filename for pressure level temperature data from 2008 will be
+ *ERA5Global_PL_2008_3hrly_Q.nc*. Keep all pressure level files in the *Pressure_Level* folder as described above.
+ There should be **five** ERA5 files for each year of pressure level data.
+ * After downloading ERA5 data, the data must be sliced and additional variables must be calculated. This is
+ accomplished in the *create_era5_netcdf.py* script:
+
+ python create_era5_netcdf.py --netcdf_era5_indir {} --netcdf_outdir {} --date {} {} {}
+
+| Argument | Type | Default | Description |
+|---------------------|---------|---------|------------------------------------------------------------------------------|
+| *netcdf_era5_indir* | str | (none) | Input directory for the ERA5 netCDF files. |
+| *netcdf_outdir* | str | (none) | Output directory for the sliced ERA5 netCDF files with additional variables. |
+| *date* | int (3) | (none) | Year, month, and day. |
+
+* ###### All netCDF files will be stored in subdirectories sorted by month (e.g. */netcdf/201304* only contains data with initialization times in April 2013).
+
+* Predictor variables can also be sourced from multiple NWP models using the *download_nwp.py* script. Supported
+ models include ECMWF, GFS, HRRR, NAM 12km, and the individual NAM nests.
+ * Similar to sliced ERA5 netCDF files, downloaded GRIB files will be sorted into monthly directories.
+
+ python download_nwp.py --grib_outdir {} --model {} --init_time {}
+
+| Argument | Type | Default | Description |
+|------------------|------------|---------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| *grib_outdir* | str | (none) | Output directory for the downloaded GRIB files. |
+| *model* | str | (none) | Output directory for the sliced ERA5 netCDF files with additional variables. |
+| *init_time* | str | (none) | Initialization time of the model run, string formatted as YYYY-MM-DD-HH. |
+| *range* | str (3) | (none) | Date range and frequency of the data to download. 3 arguments must be passed: the start and end dates of the range, and the timestep frequency. Reference *download_nwp.py* for additional information on this argument. |
+| *forecast_hours* | int (xN) | (none) | List of forecast hours to download for the initialization time(s). |
+| *verbose* | store_true | N/A | Print out the status for the GRIB file downloads. |
+
+###### Note that one of "init_time", "range" must be passed into the command line.
+
+* After downloading the GRIB files, they must be converted to netCDF format with the *convert_grib_to_netcdf.py*. All
+ forecast hours for a given initialization time are processed at once. The resulting netCDF files are sorted into
+ monthly directories in the same manner as ERA5 files. For ECMWF, GFS, and GDAS data, the base filename format is
+ *model_YYYYMMDDHH_fFFF_global.nc* (FFF = forecast_hour). The *_global* string in the base filename is removed for all
+ other models since they have their own specified domains.
+
+ python convert_grib_to_netcdf.py --grib_indir {} --model {} --netcdf_outdir {} --init_time {} {} {} {}
+
+| Argument | Type | Default | Description |
+|-------------------|------------|---------|-------------------------------------------------------|
+| *grib_indir* | str | (none) | Input directory for the GRIB files. |
+| *model* | str | (none) | NWP model from which the GRIB files originated. |
+| *netcdf_outdir* | str | (none) | Output directory for the netCDF files. |
+| *init_time* | int (4) | (none) | Year, month, day, hour. |
+| *gpu* | store_true | N/A | Force additional variables to be calculated on a GPU. |
+| *ignore_warnings* | store_true | N/A | Disable runtime warnings in variable calculations. |
+
+# 2. TensorFlow datasets
+
+TensorFlow datasets will be the inputs to the model. These datasets will be created from the netCDF files generated in
+section 1.
+
+#### 2a. Overview
+
+Three datasets will need to be generated: a training dataset, validation dataset, and a testing dataset. The **training
+dataset** is used to **train the model**, the **validation dataset** is used to **tune the model's hyperparameters**, and
+the **testing dataset** is used to **evaluate the model** with data that the model has never seen. The three datasets are split up by years. For
+example, the training dataset may cover years 2008-2017, while the validation and test datasets cover 2018-2019 and 2020,
+respectively. In this example, 2020 data will not be used to train the model. The inputs to the models are the selected predictor variables, while the
+model outputs are probabilities of each of the target front types.
+
+#### 2b. Designing the datasets
+
+There are several steps in the process of building TensorFlow datasets.
+
+1. Choose the years for the dataset. Currently, only years 2008-2023 are supported. The same year cannot be used in multiple
+ datasets (e.g. if 2016 data is used in the training dataset, in cannot be used in validation nor testing sets).
+2. Determine the predictor variables (inputs) and front types (targets/labels) that will make up the dataset. Complete
+ lists of variables and front types can be found in appendices 6b and 6c of this guide.
+3. Choose what vertical levels to include in the inputs. The list of acceptable vertical levels can be found in appendix
+ 6d.
+4. Determine the shape (number of dimensions) of the inputs and targets. (e.g. do you want a model that takes 3D inputs
+ and tries to predict 2D targets?)
+5. Choose the normalization method. Normalizing the variables allows the model to be more stable during the training process.
+ Available normalization methods can be found in appendix 6g.
+6. Choose the domains of the datasets. All available domains can be found in appendix 6a.
+7. Select the size of the images for the training/validation datasets and the number of images to extract from each
+ timestep. **Note that by default a timestep will only be used in the final dataset if all requested front types are
+ present in that timestep over the provided domain.**
+ 1. Determine whether or not you would like to retain timesteps that do not contain all requested front types. For
+ example, if you build a dataset with cold and warm fronts, timesteps with cold fronts that do not also have warm
+ fronts will not be included in the final dataset by default. You can, however, retain a fraction (or even all) of
+ the images that do not have all requested front types.
+8. Explore data augmentation and front expansion
+ 1. Data augmentation is the process of modifying the inputs to the model. Images can be modifying by adding noise
+ and rotating or reflecting the images.
+ 2. "Front expansion" refers to the process of expanding the labels identifying the presence of front boundaries at
+ each grid point. Expanding the labels in the training and validation datasets trains the model to output larger
+ frontal regions.
+9. Finally, run the convert_netcdf_to_tf.py script to build the dataset.
+
+ python convert_netcdf_to_tf.py --variables_netcdf_indir {} --fronts_netcdf_indir {} --tf_outdir {} --year_and_month {} --front_types {}
+
+| Argument | Type | Default | Description |
+|------------------------|------------|------------|--------------------------------------------------------------------------------------------------------------------------|
+| *variable_indir* | str | (none) | Input directory for the netCDF files containing variable data. |
+| *fronts_indir* | str | (none) | Input directory for the netCDF files containing front labels. |
+| *tf_outdir* | str | (none) | Output directory for the tensorflow datasets. |
+| *year_and_month* | int (2) | (none) | Year and month for which to generate tensorflow datasets. |
+| *data_source* | str | "era5" | Source of the input variables. |
+| *front_types* | str (N) | (none) | Front types to use as targets (Appendix 6c for options). |
+| *variables* | str (N) | (none) | Variables to include in the inputs (Appendix 6b for options). |
+| *pressure_levels* | str (N) | (none) | Pressure levels to include in the inputs (Appendix 6d for options). |
+| *num_dims* | int (2) | 2 2 | Number of dimensions in the input data and front labels, respectively. |
+| *domain* | str | "conus" | Domain that the dataset will cover (Appendix 6a for options). |
+| *override_extent* | int (4) | (none) | Override the default domain extent by selecting a custom extent. [min lon, max lon, min lat, max lat] |
+| *images* | int (2) | 1 1 | Number of images to extract from the timestep along the longitude and latitude dimensions. |
+| *image_size* | int (2) | 128 128 | Size of the latitude and longitude dimensions of the images (# pixels). |
+| *normalization_method* | str | "standard" | Method for normalizing the input variables (Appendix 6g for options). |
+| *shuffle_timesteps* | store_true | N/A | Shuffle the order of the timesteps when generating the dataset. |
+| *shuffle_images* | store_true | N/A | Shuffle the order of the images in each timestep. |
+| *front_dilation* | int | 0 | Number of pixels to expand the front labels by in all directions. |
+| *timestep_fraction* | float | 1.0 | Fraction of timesteps WITHOUT all necessary front types that will be retained in the dataset. (0 <= x <= 1) |
+| *image_fraction* | float | 1.0 | Fraction of images WITHOUT all necessary front types in the selected that will be retained in the dataset. (0 <= x <= 1) |
+| *noise_fraction* | float | 0.0 | Fraction of pixels in each image that will contain salt and pepper noise. (0 <= x <= 1) |
+| *flip_chance_lon* | float | 0.0 | Chance that an image will have its longitude dimension reversed. (0 <= x <= 1) |
+| *flip_chance_lat* | float | 0.0 | Chance that an image will have its latitude dimension reversed. (0 <= x <= 1) |
+| *overwrite* | store_true | N/A | Overwrite the contents of any existing variables and fronts data. |
+| *verbose* | store_true | N/A | Print out the progress of the dataset generation. |
+| *gpu_device* | int (N) | (none) | GPU device numbers. |
+| *memory_growth* | store_true | N/A | Use memory growth on the GPU(s). |
+| *seed* | int | (none) | Seed for the random number generators. ||
+
+###### NOTES:
+* If you are generating a dataset for evaluation purposes (e.g., calculating model performance):
+ * Make sure that 'timestep_fraction', 'image_fraction', 'noise_fraction', 'flip_chance_lat', 'flip_chance_lon', and 'images' are all set to their default values.
+ * Do NOT shuffle the timesteps or images.
+ * Make sure 'image_size' covers the entire domain of interest (e.g., in the case of CONUS, this would --image_size 128 288).
+ * Assure that the variables, pressure levels, and front types used in the evaluation dataset are identical to the training/validation datasets. If they are different, the code will not work.
+* Once the command is executed, a .pkl file and .txt file will be saved to *tf_outdir*. The .pkl file will contain
+ properties of the dataset (values of the table's arguments), and the .txt file is a readable version of the .pkl file.
+* After the .pkl file is created, the arguments in the .pkl file will be referenced in the future in order to create
+ consistent datasets. For example, if a tensorflow dataset in */my_tf_ds* was initially created with the *num_dims*
+ argument set to *3 2*, passing "--num_dims 3 3" into the command line with the same output directory (i.e.
+ */my_tf_ds*) will have no effect. The .pkl file will be utilized to set critical arguments if the .pkl file exists.
+ See *convert_netcdf_to_tf.py* for more information.
+
+# 3. Model training
+
+ python train_model.py (too many arguments to list - consult the train_model.py script)
+
+###### NOTES:
+
+* There are many arguments for *train_model.py* and their descriptions and usages are too long to list in this guide.
+ Consult *train_model.py* and read through each argument *carefully* as you will need to use most of the available
+ arguments.
+* All model architecture options can be found in Appendix 6e.
+
+# 4. Evaluation (performance)
+
+Before performing an evaluation, make sure a testing dataset exists for the chosen test year(s). The evaluation process
+is broken down into four steps.
+
+#### 4a. Generating test predictions
+
+This is the first step in the evaluation process. Using the *./evaluation/predict_tf.py* script, you can generate
+predictions with the created tensorflow datasets. For evaluation purposes, the predictions should be generated with the
+testing dataset. NetCDF files containing model predictions will be saved and sorted by month.
+
+* Use ONE of the following commands to generate predictions. There are two commands that can be used depending on
+ whether you want to generate predictions for the entire dataset in one run or generate predictions one month at a
+ time. The ability to run one month at a time was added so that predictions can be generated in parallel with multiple
+ commands.
+
+ python ./evaluation/predict_tf.py --model_number {} --model_dir {} --tf_indir {} --dataset {}
+ python ./evaluation/predict_tf.py --model_number {} --model_dir {} --tf_indir {} --year_and_month {} {}
+
+| Argument | Type | Default | Description |
+|------------------|------------|---------|----------------------------------------------------------------------------------------------------|
+| *dataset* | str | (none) | Dataset for which the predictions will be generated. Options are "training", "validation", "test". |
+| *year_and_month* | int (2) | (none) | Year and month for which the predictions will be generated. |
+| *model_dir* | str | (none) | Parent directory for all models. (e.g. */models*) |
+| *model_number* | int | (none) | Number assigned to the model during training. |
+| *tf_indir* | str | (none) | Input directory for the tensorflow dataset for which predictions will be generated. |
+| *data_source* | str | "era5" | Source of the input variables in the datasets. |
+| *gpu_device* | int (N) | (none) | GPU device number. |
+| *batch_size* | int | 8 | Batch size for the model predictions. |
+| *memory_growth* | store_true | N/A | Use memory growth on the GPU. |
+| *overwrite* | store_true | N/A | Overwrite any existing prediction netCDF files. |
+
+#### 4b. Calculating general performance statistics
+
+After generating predictions, general performance statistics must be calculated over the predictions. Similar to
+generating predictions in 4a, only use ONE of the commands listed below. Note that this part may require a very large
+amount of storage (up to and even exceeding 100s of GBs) depending on the size of the dataset. Performance statistics
+will be calculated using the netCDF files containing front labels rather than the labels contained in the tensorflow
+dataset. If using 0.25° data (resolution of ERA5 data), performance statistics for five neighborhoods will be
+calculated: 0.5° (~50 km), 1° (~100 km), 1.5° (~150 km), 2° (~200 km), 2.5° (~250 km). *Statistics for higher resolution
+data are currently not supported.*
+
+ python generate_performance_stats.py --model_number {} --model_dir {} --fronts_netcdf_indir {} --domain {} --dataset {}
+ python generate_performance_stats.py --model_number {} --model_dir {} --fronts_netcdf_indir {} --domain {} --year_and_month {} {}
+
+| Argument | Type | Default | Description |
+|------------------|------------|---------|----------------------------------------------------------------------------------------------------|
+| *dataset* | str | (none) | Dataset for which the predictions will be generated. Options are "training", "validation", "test". |
+| *year_and_month* | int (2) | (none) | Year and month for which the predictions will be generated. |
+| *domain* | str | (none) | Domain of the predictions from which statistics are being calculated. |
+| *model_number* | int | (none) | Number assigned to the model during training. |
+| *model_dir* | str | (none) | Parent directory for all models. (e.g. */models*) |
+| *fronts_indir* | str | (none) | Input directory for the netCDF files containing front labels. |
+| *data_source* | str | "era5" | Source of the input variables in the datasets. |
+| *gpu_device* | int | (none) | GPU device number. |
+| *memory_growth* | store_true | N/A | Use memory growth on the GPU(s). |
+| *overwrite* | store_true | N/A | Overwrite any existing netCDF files containing performance statistics. |
+
+#### 4c. Performance diagrams
+
+At long last, we can now generate large diagrams for each front type to highlight model performance. Each diagram
+contains four subplots. The first subplot within the diagram will be a CSI diagram, highlighting the probability of
+detection (POD) and false alarm ratio (FAR) for the five neighborhoods described in section 4b. The second subplot will
+be reliability diagram, which shows how the model's output probabilities align with the true probabilities of the front
+types being present within the distances of the neighborhoods. The third subplot will be a data table that highlights
+the CSI, HSS (Heidke Skill Score), POD, FAR, and frequency bias (FB) at the probability threshold where CSI is maximized.
+The final subplot will be a spatial CSI diagram that shows the model's performance over the provided domain with a specified neighborhood (
+default is 250 km).
+
+ python plot_performance_diagrams.py --model_number {} --model_dir {} --domain {} --dataset {}
+
+| Argument | Type | Default | Description |
+|--------------------|------|---------|----------------------------------------------------------------------------------------------------|
+| *dataset* | str | (none) | Dataset for which the predictions will be generated. Options are "training", "validation", "test". |
+| *data_source* | str | "era5" | Source of the input variables in the datasets. |
+| *domain* | str | (none) | Domain of the predictions from which statistics are being calculated. |
+| *map_neighborhood* | int | 250 | Neighborhood for the CSI map in kilometers. Options are: 50, 100, 150, 200, 250 |
+| *model_dir* | str | (none) | Parent directory for all models. |
+| *model_number* | int | (none) | Number assigned to the model during training. |
+| *confidence_level* | int | 95 | Confidence level expressed as a percentage. |
+| *num_iterations* | int | 10000 | Number of iterations to perform when bootstrapping the statistics. |
+| *output_type* | str | "png" | Output type for the image file. |
+
+# 5. XAI (coming soon)
+
+# 6. Appendix
+
+### 6a. Domains
+
+| Argument string | Domain | Extent |
+|-----------------|---------------------------------------------------|--------------------------------|
+| *atlantic* | Atlantic Ocean | 16-55.75°N, 290-349.75°E |
+| *conus* | Continental United States | 25-56.75°N, 228-299.75°E |
+| *full* | Unified surface analysis domain used by NOAA | 0-80°N, 130°E eastward to 10°E |
+| *global* | Global domain | -89.75-90°N, 0-359.75°E |
+| *hrrr* | High Resolution Rapid Refresh (HRRR) model domain | non-uniform grid |
+| *nam-12km* | 12-km North American Model (NAM) domain | non-uniform grid |
+| *namnest-conus* | CONUS nest of the 3-km NAM | non-uniform grid |
+| *pacific* | Pacific Ocean | 16-55.75°N, 145-234.75°E |
+
+### 6b. Variables
+
+| Argument string | Variable |
+|-----------------|------------------------------------------|
+| *q* | Specific humidity |
+| *r* | Mixing ratio |
+| *RH* | Relative humidity |
+| *sp_z* | Surface pressure and geopotential height |
+| *theta* | Potential temperature |
+| *theta_e* | Equivalent potential temperature |
+| *theta_v* | Virtual potential temperature |
+| *theta_w* | Wet-bulb potential temperature |
+| *T* | Air temperature |
+| *Td* | Dewpoint temperature |
+| *Tv* | Virtual temperature |
+| *Tw* | Wet-bulb temperature |
+| *u* | U-component of wind |
+| *v* | V-component of wind |
+
+### 6c. Front types
+
+| Identifier | Argument string | Front type |
+|------------|-----------------|--------------------------------|
+| 1 | *CF* | Cold front |
+| 2 | *CF-F* | Cold front (forming) |
+| 3 | *CF-D* | Cold front (dissipating) |
+| 4 | *WF* | Warm front |
+| 5 | *WF-F* | Warm front (forming) |
+| 6 | *WF-D* | Warm front (dissipating) |
+| 7 | *SF* | Stationary front |
+| 8 | *SF-F* | Stationary front (forming) |
+| 9 | *SF-D* | Stationary front (dissipating) |
+| 10 | *OF* | Occluded front |
+| 11 | *OF-F* | Occluded front (forming) |
+| 12 | *OF-D* | Occluded front (dissipating) |
+| 13 | *INST* | Outflow boundary |
+| 14 | *TROF* | Trough |
+| 15 | *TT* | Tropical trough |
+| 16 | *DL* | Dryline |
+
+### 6d. Vertical levels
+
+| Argument string | Level |
+|-----------------|----------|
+| *surface* | Surface |
+| *1000* | 1000 hPa |
+| *950* | 950 hPa |
+| *900* | 900 hPa |
+| *850* | 850 hPa |
+
+### 6e. Models
+
+| Argument string | Model | Reference |
+|------------------|----------------|----------------------------------------------------|
+| *attention_unet* | Attention UNET | https://arxiv.org/pdf/1804.03999 |
+| *unet* | UNET | https://arxiv.org/pdf/1505.04597 |
+| *unet_ensemble* | UNET ensemble | https://arxiv.org/pdf/1912.05074 |
+| *unet_plus* | UNET+ | https://arxiv.org/pdf/1912.05074 |
+| *unet_2plus* | UNET++ | https://arxiv.org/pdf/1912.05074 |
+| *unet_3plus* | UNET3+ | https://arxiv.org/ftp/arxiv/papers/2004/2004.08790 |
+
+### 6f. Activation functions
+
+Currently there are 30 supported activation functions.
+
+| Argument string | Activation function | Reference |
+|--------------------|---------------------------------------------------|----------------------------------------------------------------------------|
+| *elliott* | Elliott | https://link.springer.com/article/10.1007/s00521-017-3210-6 |
+| *elu* | Exponential linear unit (ELU) | https://arxiv.org/abs/1511.07289 |
+| *exponential* | Exponential function (y = e^x) | -- |
+| *gcu* | Growing cosine unit (GCU) | https://arxiv.org/abs/2108.12943 |
+| *gelu* | Gaussian error linear unit (GELU) | https://arxiv.org/abs/1606.08415 |
+| *hard_sigmoid* | Hard sigmoid | https://en.wikipedia.org/wiki/Hard_sigmoid |
+| *hexpo* | Hexpo | https://ieeexplore.ieee.org/document/7966168 |
+| *isigmoid* | Improved logistic sigmoid (ISigmoid) | https://ieeexplore.ieee.org/document/8415753 |
+| *lisht* | Linearly-scaled hyperbolic tangent (LiSHT) | https://arxiv.org/abs/1901.05894 |
+| *leaky_relu* | Leaky rectified linear unit (Leaky ReLU) | https://www.tensorflow.org/api_docs/python/tf/keras/layers/LeakyReLU |
+| *linear* | Linear function (*y = x*) | -- |
+| *mish* | Mish | https://arxiv.org/abs/1908.08681 |
+| *prelu* | Parametric rectified linear unit (PReLU) | https://www.tensorflow.org/api_docs/python/tf/keras/layers/PReLU |
+| *psigmoid* | Parametric sigmoid function (PSF) | https://www.sciencedirect.com/science/article/abs/pii/S092523120400236X |
+| *ptanh* | Penalized hyperbolic tangent (pTanh) | https://arxiv.org/abs/1602.05980 |
+| *ptelu* | Parametric tangent hyperbolic linear unit (PTELU) | https://ieeexplore.ieee.org/document/8265328 |
+| *relu* | Rectified linear unit (ReLU) | https://medium.com/@danqing/a-practical-guide-to-relu-b83ca804f1f7 |
+| *resech* | Rectified hyperbolic secant (ReSech) | https://link.springer.com/article/10.1007/s10489-015-0744-0 |
+| *selu* | Scaled exponential linear unit (SELU) | https://arxiv.org/abs/1706.02515 |
+| *sigmoid* | Sigmoid function | https://en.wikipedia.org/wiki/Sigmoid_function |
+| *smelu* | Smooth rectified linear unit (SmeLU) | https://arxiv.org/pdf/2202.06499 |
+| *snake* | Snake function | https://arxiv.org/pdf/2006.08195 |
+| *softmax* | Softmax function | https://link.springer.com/chapter/10.1007/978-3-642-76153-9_28 |
+| *softplus* | Softplus function | https://paperswithcode.com/method/softplus |
+| *softsign* | Softsign function | https://paperswithcode.com/method/softsign-activation |
+| *swish* | Swish | https://arxiv.org/abs/1710.05941 |
+| *srs* | Soft-Root-Sign (SRS) | https://arxiv.org/abs/2003.00547 |
+| *stanh* | Scaled hyperbolic tangent (STanh) | https://arxiv.org/abs/2003.00547 |
+| *tanh* | Hyperbolic tangent | https://mathworld.wolfram.com/HyperbolicTangent.html |
+| *thresholded_relu* | Thresholded rectified linear unit | https://www.tensorflow.org/api_docs/python/tf/keras/layers/ThresholdedReLU |
+
+### 6g. Normalization methods
+
+| Argument string | Normalization method |
+|-------------------|---------------------------------------------|
+| standard | Standardization (z-score) |
+| standard_weighted | Latitude-weighted standardization (z-score) |
+| min-max | Min-max scaling |
\ No newline at end of file
diff --git a/configs/1702.yaml b/configs/1702.yaml
new file mode 100644
index 0000000..59da803
--- /dev/null
+++ b/configs/1702.yaml
@@ -0,0 +1,124 @@
+epochs: 5000
+training_steps_per_epoch: 30
+validation_steps_per_epoch:
+validation_frequency: 1
+verbose: 1
+repeat: true
+seed: 42
+
+model:
+ name: "unet_3plus"
+ batch_normalization: true
+ num_filters: [16, 32, 64, 128]
+ kernel_size: [5, 5, 5]
+ pool_size: [2, 2, 1]
+ upsample_size: [2, 2, 1]
+ depth: 4
+ modules_per_node: 2
+ padding: "same"
+ bias: true
+ loss:
+ name: "fractions_skill_score"
+ config:
+ mask_size: [3, 3]
+ metric:
+ name: "critical_success_index"
+ config:
+ class_weights: [0, 1, 1, 1, 1, 1]
+
+ optimizer:
+ name: "Adam"
+ config:
+ beta_1: 0.9
+ beta_2: 0.999
+
+ convolution_activity_regularizer:
+ regularizer:
+
+ bias_vector:
+ constraint:
+ initializer:
+ name: "zeros"
+ config:
+ regularizer:
+
+ kernel_matrix:
+ constraint:
+ initializer:
+ name: "glorot_uniform"
+ config:
+ regularizer:
+
+ activation:
+ name: "gelu"
+ config:
+
+
+
+
+
+
+data:
+ train_years: [2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019]
+ val_years: [2020]
+ test_years: [2021]
+ shuffle: true
+ normalization_method: "standard"
+
+ era5:
+ domain_extent: [-140.0, -60.0, 20.0, 60.0]
+ # Use canonical (pressure-level) names for variables that have both surface
+ # and pressure-level representations. Surface-only variables (e.g.
+ # "mean_sea_level_pressure") are listed here directly.
+ # The mapping from pressure-level name to surface name is defined in
+ # fronts.data.config.SURFACE_VARIABLE_MAP and applied automatically when
+ # "surface" appears in the levels list.
+ variables:
+ - "temperature" # surface counterpart: 2m_temperature
+ - "specific_humidity" # surface counterpart: surface_specific_humidity
+ - "u_component_of_wind" # surface counterpart: 10m_u_component_of_wind
+ - "v_component_of_wind" # surface counterpart: 10m_v_component_of_wind
+ - "mean_sea_level_pressure" # surface-only
+ # levels may contain "surface" and/or integer hPa values in any order.
+ # "surface" triggers loading of the surface counterpart for each variable.
+ levels: ["surface", 1000, 950, 900, 850]
+ store: "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
+ chunks:
+ time: 48
+ consolidated: true
+
+ fronts:
+ directory: "/path/to/fronts/netcdf"
+ front_types: "MERGED-ALL"
+
+ batch:
+ input_sizes:
+ time: 1
+ latitude: 128
+ longitude: 128
+ target_sizes:
+ time: 1
+ latitude: 128
+ longitude: 128
+ prefetch_number: 3
+ preload_batch: false
+
+wandb:
+ project_name: 'fronts'
+ model_run_name: '1702_retrain'
+ # -- Optional parameters --
+ # log_frequency:
+ # upload_checkpoints:
+ # api_key:
+ # wandb_filepath:
+
+callbacks:
+ monitor: "val_loss"
+ verbose: 1
+ save_best_only: true
+ save_weights_only: false
+ save_freq: "epoch"
+ # model_checkpoint_path:
+ # csv_logger_path:
+ # patience:
+
diff --git a/configs/1702_tf.yaml b/configs/1702_tf.yaml
new file mode 100644
index 0000000..f417744
--- /dev/null
+++ b/configs/1702_tf.yaml
@@ -0,0 +1,93 @@
+epochs: 5000
+training_steps_per_epoch: 30
+validation_steps_per_epoch:
+validation_frequency: 1
+verbose: 1
+repeat: true
+seed: 42
+
+model:
+ name: "unet_3plus"
+ batch_normalization: true
+ num_filters: [16, 32, 64, 128]
+ kernel_size: [5, 5, 5]
+ pool_size: [2, 2, 1]
+ upsample_size: [2, 2, 1]
+ depth: 4
+ modules_per_node: 2
+ padding: "same"
+ bias: true
+ loss:
+ name: "fractions_skill_score"
+ config:
+ mask_size: [3, 3]
+ metric:
+ name: "critical_success_index"
+ config:
+ class_weights: [0, 1, 1, 1, 1, 1]
+
+ optimizer:
+ name: "Adam"
+ config:
+ beta_1: 0.9
+ beta_2: 0.999
+
+ convolution_activity_regularizer:
+ regularizer:
+
+ bias_vector:
+ constraint:
+ initializer:
+ name: "zeros"
+ config:
+ regularizer:
+
+ kernel_matrix:
+ constraint:
+ initializer:
+ name: "glorot_uniform"
+ config:
+ regularizer:
+
+ activation:
+ name: "gelu"
+ config:
+
+
+data:
+ train_years: [2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2020, 2021, 2022, 2024]
+ val_years: [2023]
+ test_years: []
+ shuffle: true
+
+ # Load pre-built TF dataset snapshots directly from disk.
+ # Subdirectory naming convention: {year}-{month}_tf (e.g. 2010-1_tf)
+ # train_years/val_years/test_years above are injected at build time —
+ # leave them as [] here.
+ tf_dataset:
+ directory: "/ourdisk/hpc/ai2es/fronts/raw_front_data/tf_datasets/5class-7lvl-conus"
+ train_years: []
+ val_years: []
+ test_years: []
+ shuffle: true
+ shuffle_buffer: 1000
+ prefetch: 3
+
+wandb:
+ project_name: 'fronts'
+ model_run_name: '1702_retrain'
+ wandb_filepath: '/ourdisk/hpc/ai2es/tman/models/1702_retrain.keras'
+ # -- Optional parameters --
+ # log_frequency:
+ # upload_checkpoints:
+ # api_key:
+
+callbacks:
+ monitor: "val_loss"
+ verbose: 1
+ save_best_only: true
+ save_weights_only: false
+ save_freq: "epoch"
+ model_checkpoint_path: "/ourdisk/hpc/ai2es/tman/models/1702_tf.keras"
+ csv_logger_path: "/ourdisk/hpc/ai2es/tman/logs/1702_tf.csv"
+ # patience: 55
diff --git a/configs/1702_tf_dryrun.yaml b/configs/1702_tf_dryrun.yaml
new file mode 100644
index 0000000..9923cb1
--- /dev/null
+++ b/configs/1702_tf_dryrun.yaml
@@ -0,0 +1,96 @@
+# Local dry-run config — used for smoke-testing config parsing and data
+# pipeline construction on a developer machine without cluster data or GPUs.
+#
+# Workflow:
+# 1. Generate fixtures (one-time):
+# python scripts/make_dryrun_data.py
+# 2. Run dry-run:
+# python -m fronts.train --train_config_path configs/1702_tf_dryrun.yaml --dry_run
+#
+# The fixture dataset has element shapes matching the real cluster data:
+# inputs: (128, 288, 7, 9) float16
+# targets: (128, 288, 6) float16
+
+epochs: 1
+training_steps_per_epoch: 2
+validation_steps_per_epoch: 1
+validation_frequency: 1
+verbose: 1
+repeat: false
+seed: 42
+
+model:
+ name: "unet_3plus"
+ batch_normalization: true
+ num_filters: [16, 32, 64, 128]
+ kernel_size: [5, 5, 5]
+ pool_size: [2, 2, 1]
+ upsample_size: [2, 2, 1]
+ depth: 4
+ modules_per_node: 2
+ padding: "same"
+ bias: true
+ loss:
+ name: "fractions_skill_score"
+ config:
+ mask_size: [3, 3]
+ metric:
+ name: "critical_success_index"
+ config:
+ class_weights: [0, 1, 1, 1, 1, 1]
+
+ optimizer:
+ name: "Adam"
+ config:
+ beta_1: 0.9
+ beta_2: 0.999
+
+ convolution_activity_regularizer:
+ regularizer:
+
+ bias_vector:
+ constraint:
+ initializer:
+ name: "zeros"
+ config:
+ regularizer:
+
+ kernel_matrix:
+ constraint:
+ initializer:
+ name: "glorot_uniform"
+ config:
+ regularizer:
+
+ activation:
+ name: "gelu"
+ config:
+
+
+data:
+ train_years: [2000]
+ val_years: [2001]
+ test_years: []
+ shuffle: false # no need to shuffle tiny fixtures
+
+ tf_dataset:
+ directory: "tests/fixtures/dryrun_tf_dataset"
+ train_years: []
+ val_years: []
+ test_years: []
+ shuffle: false
+ shuffle_buffer: 1
+ prefetch: 1
+
+wandb:
+ project_name: 'fronts'
+ model_run_name: 'dryrun'
+ wandb_filepath: '/tmp/dryrun_model.keras'
+
+callbacks:
+ monitor: "val_loss"
+ verbose: 1
+ save_best_only: false
+ save_weights_only: false
+ save_freq: "epoch"
+ # No model_checkpoint_path or csv_logger_path — dry run only
diff --git a/configs/predict_1702.yaml b/configs/predict_1702.yaml
new file mode 100644
index 0000000..da9b9cd
--- /dev/null
+++ b/configs/predict_1702.yaml
@@ -0,0 +1,53 @@
+# Prediction configuration for FrontFinder model 1702.
+#
+# This config is loaded as a PredictConfig via:
+# from fronts.train import open_config_yaml_as_dataclass
+# from fronts.data.config import PredictConfig
+# cfg = open_config_yaml_as_dataclass("configs/predict_1702.yaml", PredictConfig)
+# inputs_ds = cfg.build() # -> normalized xr.Dataset ready for model inference
+#
+# TIME SELECTION — uncomment exactly one block below.
+# Use ISO 8601 format with a "T" separator (e.g. "2024-06-01T12:00:00") so that
+# PyYAML treats the value as a string rather than a date, allowing dacite to cast
+# it to datetime.datetime correctly.
+
+# Option A: most recent timestep available in the zarr store (no date needed)
+time_selection:
+ most_recent: true
+
+# Option B: explicit individual analysis times (date + hour each)
+# time_selection:
+# timestamps:
+# - "2024-06-01T12:00:00"
+# - "2024-06-02T00:00:00"
+
+# Option C: all timesteps in an inclusive date range
+# time_selection:
+# date_range:
+# - "2024-06-01T00:00:00"
+# - "2024-06-07T18:00:00"
+
+normalization_method: "standard"
+
+era5:
+ domain_extent: [-140.0, -60.0, 20.0, 60.0] # [lon_min, lon_max, lat_min, lat_max]
+ # Use canonical (pressure-level) names for variables that have both surface and
+ # pressure-level representations. Surface-only variables (e.g.
+ # "mean_sea_level_pressure") are listed directly.
+ # The mapping from pressure-level name to surface name is defined in
+ # fronts.data.config.SURFACE_VARIABLE_MAP and applied automatically when
+ # "surface" appears in the levels list.
+ variables:
+ - "temperature" # surface counterpart: 2m_temperature
+ - "specific_humidity" # surface counterpart: surface_specific_humidity
+ - "u_component_of_wind" # surface counterpart: 10m_u_component_of_wind
+ - "v_component_of_wind" # surface counterpart: 10m_v_component_of_wind
+ - "mean_sea_level_pressure" # surface-only
+ # levels may contain "surface" and/or integer hPa values.
+ # "surface" triggers loading of the surface counterpart for each variable.
+ levels: ["surface", 1000, 950, 900, 850]
+ store: "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
+ chunks:
+ time: 48
+ consolidated: true
+ # Note: era5.years is unused by PredictConfig — time is controlled by time_selection above.
diff --git a/convert_front_gml_to_xml.py b/convert_front_gml_to_xml.py
deleted file mode 100644
index ee1238f..0000000
--- a/convert_front_gml_to_xml.py
+++ /dev/null
@@ -1,91 +0,0 @@
-"""
-Convert GML files containing IBM/TWC fronts into XML files.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.9.2
-"""
-
-import argparse
-from lxml import etree as ET
-from glob import glob
-import os
-import numpy as np
-
-XML_FRONT_TYPE = {'Cold Front': 'COLD_FRONT', 'Dissipating Cold Front': 'COLD_FRONT_DISS',
- 'Warm Front': 'WARM_FRONT', 'Stationary Front': 'STATIONARY_FRONT',
- 'Occluded Front': 'OCCLUDED_FRONT', 'Dissipating Occluded Front': 'OCCLUDED_FRONT_DISS',
- 'Dry Line': 'DRY_LINE', 'Trough': 'TROF', 'Squall Line': 'INSTABILITY'}
-
-XML_FRONT_COLORS = {'Cold Front': dict(red="0", green="0", blue="255"), 'Dissipating Cold Front': dict(red="0", green="0", blue="255"),
- 'Warm Front': dict(red="255", green="0", blue="0"), 'Dissipating Warm Front': dict(red="255", green="0", blue="0"),
- 'Occluded Front': dict(red="145", green="44", blue="238"), 'Dissipating Occluded Front': dict(red="145", green="44", blue="238"),
- 'Dry Line': dict(red="255", green="130", blue="71"), 'Trough': dict(red="255", green="130", blue="71"),
- 'Squall Line': dict(red="255", green="0", blue="0")}
-
-LINE_KWARGS = dict(pgenCategory="Front", lineWidth="4", sizeScale=" 1.0", smoothFactor="2", closed="false", filled="false",
- fillPattern="SOLID")
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument('--gml_indir', type=str, required=True, help="Input directory for IBM/TWC front GML files.")
- parser.add_argument('--xml_outdir', type=str, required=True, help="Output directory for front XML files.")
- parser.add_argument('--date', type=int, nargs=3, required=True, help="Date for the data to be read in. (year, month, day)")
- args = vars(parser.parse_args())
-
- year, month, day = args['date']
-
- gml_files = sorted(glob('%s/%d%02d%02d/*/*%d%02d%02d*.gml' % (args['gml_indir'], year, month, day, year, month, day)))
-
- for gml_file in gml_files:
-
- valid_time_str = os.path.basename(gml_file).split('.')[2]
- valid_time_str = valid_time_str[:4] + '-' + valid_time_str[4:6] + '-' + valid_time_str[6:8] + 'T' + valid_time_str[9:11]
- valid_time = np.datetime64(valid_time_str, 'ns')
-
- init_time_str = os.path.basename(gml_file).split('.')[3]
- if init_time_str != 'NIL': # an init time of 'NIL' is used to indicate forecast hour 0 (i.e. valid time is same as init time)
- init_time_str = init_time_str[:4] + '-' + init_time_str[4:6] + '-' + init_time_str[6:8] + 'T' + init_time_str[9:11]
- init_time = np.datetime64(init_time_str, 'ns')
- else:
- init_time_str = valid_time_str
- init_time = valid_time
-
- forecast_hour = int((valid_time - init_time) / np.timedelta64(1, 'h'))
-
- root_xml = ET.Element("Product", name="IBM_global_fronts", init_time=init_time_str, valid_time=valid_time_str, forecast_hour=str(forecast_hour))
- tree = ET.parse(gml_file, parser=ET.XMLPullParser(encoding='utf-8'))
- root_gml = tree.getroot()
-
- Layer = ET.SubElement(root_xml, "Layer", name="Default", onOff="true", monoColor="false", filled="false")
- ET.SubElement(Layer, "Color", red="255", green="255", blue="0", alpha="255")
- DrawableElement = ET.SubElement(Layer, "DrawableElement")
-
- front_elements = [element[0] for element in root_gml if element[0].tag == 'FRONT']
-
- for element in front_elements:
- front_type = [subelement.text for subelement in element if subelement.tag == 'FRONT_TYPE'][0]
- coords = [subelement for subelement in element if 'lineString' in subelement.tag][0][0][0].text
-
- Line = ET.SubElement(DrawableElement, "Line", pgenType=XML_FRONT_TYPE[front_type], **LINE_KWARGS)
- if front_type == 'Stationary Front':
- ET.SubElement(Line, "Color", red="255", green="0", blue="0", alpha="255")
- ET.SubElement(Line, "Color", red="0", green="0", blue="255", alpha="255")
- else:
- ET.SubElement(Line, "Color", **XML_FRONT_COLORS[front_type], alpha="255")
-
- coords = coords.replace('\n', '').split(' ') # generate coordinate strings
- coords = list(coord_pair.split(',') for coord_pair in coords) # generate coordinate pairs from the strings
-
- for coord_pair in coords:
- ET.SubElement(Line, "Point", Lat="%.6f" % float(coord_pair[1]), Lon="%.6f" % float(coord_pair[0]))
-
- save_path_file = "%s/IBM_fronts_%sf%03d.xml" % (args['xml_outdir'], init_time_str.replace('-', '').replace('T', ''), forecast_hour)
-
- print(save_path_file)
-
- ET.indent(root_xml)
- mydata = ET.tostring(root_xml)
- xmlFile = open(save_path_file, "wb")
- xmlFile.write(mydata)
- xmlFile.close()
diff --git a/convert_front_xml_to_netcdf.py b/convert_front_xml_to_netcdf.py
deleted file mode 100644
index 07243f4..0000000
--- a/convert_front_xml_to_netcdf.py
+++ /dev/null
@@ -1,157 +0,0 @@
-"""
-Convert front XML files to netCDF files.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.9.21
-"""
-
-import argparse
-import glob
-import numpy as np
-import os
-from utils import data_utils
-import xarray as xr
-import xml.etree.ElementTree as ET
-
-
-pgenType_identifiers = {'COLD_FRONT': 1, 'WARM_FRONT': 2, 'STATIONARY_FRONT': 3, 'OCCLUDED_FRONT': 4, 'COLD_FRONT_FORM': 5,
- 'WARM_FRONT_FORM': 6, 'STATIONARY_FRONT_FORM': 7, 'OCCLUDED_FRONT_FORM': 8, 'COLD_FRONT_DISS': 9,
- 'WARM_FRONT_DISS': 10, 'STATIONARY_FRONT_DISS': 11, 'OCCLUDED_FRONT_DISS': 12, 'INSTABILITY': 13,
- 'TROF': 14, 'TROPICAL_TROF': 15, 'DRY_LINE': 16}
-
-"""
-conus: 132 W to 60.25 W, 57 N to 26.25 N
-full: 130 E pointing eastward to 10 E, 80 N to 0.25 N
-global: 179.75 W to 180 E, 90 N to 89.75 N
-"""
-domain_coords = {'conus': {'lons': np.arange(-132, -60, 0.25), 'lats': np.arange(57, 25, -0.25)},
- 'full': {'lons': np.append(np.arange(-179.75, 10, 0.25), np.arange(130, 180.25, 0.25)), 'lats': np.arange(80, 0, -0.25)},
- 'global': {'lons': np.arange(-179.75, 180.25, 0.25), 'lats': np.arange(90, -90, -0.25)}}
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument('--xml_indir', type=str, required=True, help="Input directory for front XML files.")
- parser.add_argument('--netcdf_outdir', type=str, required=True, help="Output directory for front netCDF files.")
- parser.add_argument('--date', type=int, nargs=3, required=True, help="Date for the data to be read in. (year, month, day)")
- parser.add_argument('--distance', type=float, default=1., help="Interpolation distance in kilometers.")
- parser.add_argument('--domain', type=str, default='full', help="Domain for which to generate fronts.")
-
- args = vars(parser.parse_args())
-
- year, month, day = args['date']
-
- if args['domain'] == 'global':
- files = sorted(glob.glob("%s/IBM*_%04d%02d%02d*f*.xml" % (args['xml_indir'], year, month, day)))
- else:
- files = sorted(glob.glob("%s/pres*_%04d%02d%02d*f000.xml" % (args['xml_indir'], year, month, day)))
-
- domain_from_model = args['domain'] not in ['conus', 'full', 'global']
-
- if domain_from_model:
-
- transform_args = {'hrrr': dict(std_parallels=(38.5, 38.5), lon_ref=262.5, lat_ref=38.5),
- 'nam_12km': dict(std_parallels=(25, 25), lon_ref=265, lat_ref=40),
- 'namnest_conus': dict(std_parallels=(38.5, 38.5), lon_ref=262.5, lat_ref=38.5)}
-
- if args['domain'] == 'hrrr':
- model_dataset = xr.open_dataset('hrrr_2023040100_f000.grib', backend_kwargs=dict(filter_by_keys={'typeOfLevel': 'isobaricInhPa'}))
- elif args['domain'] == 'nam_12km':
- model_dataset = xr.open_dataset('nam_12km_2021032300_f006.grib', backend_kwargs=dict(filter_by_keys={'typeOfLevel': 'isobaricInhPa'}))
- elif args['domain'] == 'rap':
- model_dataset = xr.open_dataset('rap_2021032300_f006.grib', backend_kwargs=dict(filter_by_keys={'typeOfLevel': 'isobaricInhPa'}))
- elif args['domain'] == 'namnest_conus':
- model_dataset = xr.open_dataset('namnest_conus_2023090800_f000.grib', backend_kwargs=dict(filter_by_keys={'typeOfLevel': 'isobaricInhPa'}))
-
- gridded_lons = model_dataset['longitude'].values.astype('float32')
- gridded_lats = model_dataset['latitude'].values.astype('float32')
-
- model_x_transform, model_y_transform = data_utils.lambert_conformal_to_cartesian(gridded_lons, gridded_lats, **transform_args[args['domain']])
- gridded_x = model_x_transform[0, :]
- gridded_y = model_y_transform[:, 0]
-
- identifier = np.zeros(np.shape(gridded_lons)).astype('float32')
-
- else:
-
- gridded_lons = domain_coords[args['domain']]['lons'].astype('float32')
- gridded_lats = domain_coords[args['domain']]['lats'].astype('float32')
- identifier = np.zeros([len(gridded_lons), len(gridded_lats)]).astype('float32')
-
- for filename in files[-1:]:
-
- tree = ET.parse(filename, parser=ET.XMLParser(encoding='utf-8'))
- root = tree.getroot()
- date = os.path.basename(filename).split('_')[-1].split('.')[0].split('f')[0] # YYYYMMDDhh
- forecast_hour = int(filename.split('f')[-1].split('.')[0])
-
- hour = date[-2:]
-
- ### Iterate through the individual fronts ###
- for line in root.iter('Line'):
-
- type_of_front = line.get("pgenType") # front type
-
- print(type_of_front)
-
- lons, lats = zip(*[[float(point.get("Lon")), float(point.get("Lat"))] for point in line.iter('Point')])
- lons, lats = np.array(lons), np.array(lats)
-
- # If the front crosses the dateline or the 180th meridian, its coordinates must be modified for proper interpolation
- front_needs_modification = np.max(np.abs(np.diff(lons))) > 180
-
- if front_needs_modification or domain_from_model:
- lons = np.where(lons < 0, lons + 360, lons) # convert coordinates to a 360 degree system
-
- xs, ys = data_utils.haversine(lons, lats) # x/y coordinates in kilometers
- xy_linestring = data_utils.geometric(xs, ys) # convert coordinates to a LineString object
- x_new, y_new = data_utils.redistribute_vertices(xy_linestring, args['distance']).xy # interpolate x/y coordinates
- x_new, y_new = np.array(x_new), np.array(y_new)
- lon_new, lat_new = data_utils.reverse_haversine(x_new, y_new) # convert interpolated x/y coordinates to lat/lon
-
- date_and_time = np.datetime64('%04d-%02d-%02dT%02d' % (year, month, day, int(hour)), 'ns')
-
- expand_dims_args = {'time': np.atleast_1d(date_and_time)}
-
- if args['domain'] == 'global':
- filename_netcdf = "FrontObjects_%s_f%03d_%s.nc" % (date, forecast_hour, args['domain'])
- expand_dims_args['forecast_hour'] = np.atleast_1d(forecast_hour)
- else:
- filename_netcdf = "FrontObjects_%s_%s.nc" % (date, args['domain'])
-
- if domain_from_model:
- x_new *= 1000; y_new *= 1000 # convert front's points to meters
- x_transform, y_transform = data_utils.lambert_conformal_to_cartesian(lon_new, lat_new, **transform_args[args['domain']])
-
- gridded_indices = np.dstack((np.digitize(y_transform, gridded_y), np.digitize(x_transform, gridded_x)))[0] # translate coordinate indices to grid
- gridded_indices_unique = np.unique(gridded_indices, axis=0) # remove duplicate coordinate indices
-
- # Remove points outside the domain
- gridded_indices_unique = gridded_indices_unique[np.where(gridded_indices_unique[:, 0] != len(gridded_y))]
- gridded_indices_unique = gridded_indices_unique[np.where(gridded_indices_unique[:, 1] != len(gridded_x))]
-
- identifier[gridded_indices_unique[:, 0], gridded_indices_unique[:, 1]] = pgenType_identifiers[type_of_front] # assign labels to the gridded points based on the front type
-
- fronts_ds = xr.Dataset({"identifier": (('y', 'x'), identifier)},
- coords={"longitude": (('y', 'x'), gridded_lons), "latitude": (('y', 'x'), gridded_lats)}).expand_dims(**expand_dims_args)
-
- else:
-
- if front_needs_modification:
- lon_new = np.where(lon_new > 180, lon_new - 360, lon_new) # convert new longitudes to standard -180 to 180 range
-
- gridded_indices = np.dstack((np.digitize(lon_new, gridded_lons), np.digitize(lat_new, gridded_lats)))[0] # translate coordinate indices to grid
- gridded_indices_unique = np.unique(gridded_indices, axis=0) # remove duplicate coordinate indices
-
- # Remove points outside the domain
- gridded_indices_unique = gridded_indices_unique[np.where(gridded_indices_unique[:, 0] != len(gridded_lons))]
- gridded_indices_unique = gridded_indices_unique[np.where(gridded_indices_unique[:, 1] != len(gridded_lats))]
-
- identifier[gridded_indices_unique[:, 0], gridded_indices_unique[:, 1]] = pgenType_identifiers[type_of_front] # assign labels to the gridded points based on the front type
-
- fronts_ds = xr.Dataset({"identifier": (('longitude', 'latitude'), identifier)},
- coords={"longitude": gridded_lons, "latitude": gridded_lats}).expand_dims(**expand_dims_args)
-
- if not os.path.isdir("%s/%d%02d" % (args['netcdf_outdir'], year, month)):
- os.mkdir("%s/%d%02d" % (args['netcdf_outdir'], year, month))
- fronts_ds.to_netcdf(path="%s/%d%02d/%s" % (args['netcdf_outdir'], year, month, filename_netcdf), engine='netcdf4', mode='w')
diff --git a/convert_grib_to_netcdf.py b/convert_grib_to_netcdf.py
deleted file mode 100644
index adf7acf..0000000
--- a/convert_grib_to_netcdf.py
+++ /dev/null
@@ -1,290 +0,0 @@
-"""
-Convert GDAS and/or GFS grib files to netCDF files.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.7.24
-"""
-
-import argparse
-import time
-import xarray as xr
-from utils import variables
-import glob
-import numpy as np
-import os
-import tensorflow as tf
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument('--grib_indir', type=str, required=True, help="Input directory for GDAS grib files.")
- parser.add_argument('--model', required=True, type=str, help="GDAS or GFS")
- parser.add_argument('--netcdf_outdir', type=str, required=True, help="Output directory for the netCDF files.")
- parser.add_argument('--init_time', type=int, nargs=4, required=True, help="Date and time for the data to be read in. (year, month, day, hour)")
- parser.add_argument('--overwrite_grib', action='store_true', help="Overwrite the split grib files if they exist.")
- parser.add_argument('--delete_original_grib', action='store_true', help="Delete the original grib files after they are split.")
- parser.add_argument('--delete_split_grib', action='store_true', help="Delete the split grib files after they have been opened.")
- parser.add_argument('--gpu', action='store_true',
- help="Use a GPU to perform calculations of additional variables. This can provide enormous speedups when generating "
- "very large amounts of data.")
-
- args = vars(parser.parse_args())
-
- gpus = tf.config.list_physical_devices(device_type='GPU')
- if len(gpus) > 0 and args['gpu']:
- print("Using GPU for variable derivations")
- tf.config.set_visible_devices(devices=gpus[0], device_type='GPU')
- gpus = tf.config.get_visible_devices(device_type='GPU')
- tf.config.experimental.set_memory_growth(device=gpus[0], enable=True)
- else:
- print("Using CPUs for variable derivations")
- tf.config.set_visible_devices([], 'GPU')
-
- args['model'] = args['model'].lower()
- year, month, day, hour = args['init_time']
-
- resolution = 0.25
-
- keys_to_extract = ['gh', 'mslet', 'r', 'sp', 't', 'u', 'v']
-
- pressure_level_file_indices = [0, 2, 4, 5, 6]
- surface_data_file_indices = [2, 4, 5, 6]
- raw_pressure_data_file_index = 3
- mslp_data_file_index = 1
-
- # all lon/lat values in degrees
- start_lon, end_lon = 0, 360 # western boundary, eastern boundary
- start_lat, end_lat = 90, -90 # northern boundary, southern boundary
- unified_longitude_indices = np.arange(0, 360 / resolution)
- unified_latitude_indices = np.arange(0, 180 / resolution + 1).astype(int)
- lon_coords_360 = np.arange(start_lon, end_lon + resolution, resolution)
-
- domain_indices_isel = {'longitude': unified_longitude_indices,
- 'latitude': unified_latitude_indices}
-
- chunk_sizes = {'latitude': 721, 'longitude': 1440}
-
- dataset_dimensions = ('forecast_hour', 'pressure_level', 'latitude', 'longitude')
-
- grib_filename_format = f"%s/%d%02d/%s_%d%02d%02d%02d_f*.grib" % (args['grib_indir'], year, month, args['model'], year, month, day, hour)
- individual_variable_filename_format = f"%s/%d%02d/%s_*_%d%02d%02d%02d.grib" % (args['grib_indir'], year, month, args['model'], year, month, day, hour)
-
- ### Split grib files into one file per variable ###
- grib_files = list(glob.glob(grib_filename_format))
- grib_files = [file for file in grib_files if 'idx' not in file]
-
- for key in keys_to_extract:
- output_file = f"%s/%d%02d/%s_%s_%d%02d%02d%02d.grib" % (args['grib_indir'], year, month, args['model'], key, year, month, day, hour)
- if (os.path.isfile(output_file) and args['overwrite_grib']) or not os.path.isfile(output_file):
- os.system(f'grib_copy -w shortName={key} {" ".join(grib_files)} {output_file}')
-
- if args['delete_original_grib']:
- [os.remove(file) for file in grib_files]
-
- time.sleep(5) # Pause the code for 5 seconds to ensure that all contents of the individual files are preserved
-
- # grib files by variable
- grib_files = sorted(glob.glob(individual_variable_filename_format))
-
- pressure_level_files = [grib_files[index] for index in pressure_level_file_indices]
- surface_data_files = [grib_files[index] for index in surface_data_file_indices]
-
- raw_pressure_data_file = grib_files[raw_pressure_data_file_index]
- if 'mslp_data_file_index' in locals():
- mslp_data_file = grib_files[mslp_data_file_index]
- mslp_data = xr.open_dataset(mslp_data_file, engine='cfgrib', backend_kwargs={'filter_by_keys': {'typeOfLevel': 'meanSea'}}, chunks=chunk_sizes).drop_vars(['step'])
-
- pressure_levels = [1000, 950, 900, 850, 700, 500]
-
- # Open the datasets
- pressure_level_data = xr.open_mfdataset(pressure_level_files, engine='cfgrib', backend_kwargs={'filter_by_keys': {'typeOfLevel': 'isobaricInhPa'}}, chunks=chunk_sizes, combine='nested').sel(isobaricInhPa=pressure_levels).drop_vars(['step'])
- surface_data = xr.open_mfdataset(surface_data_files, engine='cfgrib', backend_kwargs={'filter_by_keys': {'typeOfLevel': 'sigma'}}, chunks=chunk_sizes).drop_vars(['step'])
- raw_pressure_data = xr.open_dataset(raw_pressure_data_file, engine='cfgrib', backend_kwargs={'filter_by_keys': {'typeOfLevel': 'surface', 'stepType': 'instant'}}, chunks=chunk_sizes).drop_vars(['step'])
-
- # Calculate the forecast hours using the surface_data dataset
- try:
- run_time = surface_data['time'].values.astype('int64')
- except KeyError:
- run_time = surface_data['run_time'].values.astype('int64')
-
- valid_time = surface_data['valid_time'].values.astype('int64')
- forecast_hours = np.array((valid_time - int(run_time)) / 3.6e12, dtype='int32')
-
- try:
- num_forecast_hours = len(forecast_hours)
- except TypeError:
- num_forecast_hours = 1
- forecast_hours = [forecast_hours, ]
-
- if args['model'] in ['gdas', 'gfs']:
- mslp = mslp_data['mslet'].values # mean sea level pressure (eta model reduction)
- mslp_z = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- mslp_z[:, 0, :, :] = mslp / 100 # convert to hectopascals
-
- P = np.empty(shape=(num_forecast_hours, len(pressure_levels), chunk_sizes['latitude'], chunk_sizes['longitude']), dtype=np.float32) # create 3D array of pressure levels to match the shape of variable arrays
- for pressure_level_index, pressure_level in enumerate(pressure_levels):
- P[:, pressure_level_index, :, :] = pressure_level * 100
-
- print("Retrieving downloaded variables")
- ### Pressure level variables provided in the grib files ###
- T_pl = pressure_level_data['t'].values
- RH_pl = pressure_level_data['r'].values / 100
- u_pl = pressure_level_data['u'].values
- v_pl = pressure_level_data['v'].values
- z = pressure_level_data['gh'].values / 10 # Convert to dam
- if 'mslp_data_file_index' in locals():
- mslp_z[:, 1:, :, :] = z
-
- ### Surface variables provided in the grib files ###
- sp = raw_pressure_data['sp'].values
- T_sigma = surface_data['t'].values
- RH_sigma = surface_data['r'].values / 100
- u_sigma = surface_data['u'].values
- v_sigma = surface_data['v'].values
- surface_data_latitudes = pressure_level_data['latitude'].values
-
- if len(gpus) > 0:
- T_pl = tf.convert_to_tensor(T_pl)
- RH_pl = tf.convert_to_tensor(RH_pl)
- P = tf.convert_to_tensor(P)
- sp = tf.convert_to_tensor(sp)
- T_sigma = tf.convert_to_tensor(T_sigma)
- RH_sigma = tf.convert_to_tensor(RH_sigma)
-
- print("Deriving additional variables")
- vap_pres_pl = RH_pl * variables.vapor_pressure(T_pl)
- Td_pl = variables.dewpoint_from_vapor_pressure(vap_pres_pl)
- Tv_pl = variables.virtual_temperature_from_dewpoint(T_pl, Td_pl, P)
- Tw_pl = variables.wet_bulb_temperature(T_pl, Td_pl)
- r_pl = variables.mixing_ratio_from_dewpoint(Td_pl, P) * 1000 # Convert to g/kg
- q_pl = variables.specific_humidity_from_dewpoint(Td_pl, P) * 1000 # Convert to g/kg
- theta_pl = variables.potential_temperature(T_pl, P)
- theta_e_pl = variables.equivalent_potential_temperature(T_pl, Td_pl, P)
- theta_v_pl = variables.virtual_potential_temperature(T_pl, Td_pl, P)
- theta_w_pl = variables.wet_bulb_potential_temperature(T_pl, Td_pl, P)
-
- # Create arrays of coordinates for the surface data
- vap_pres_sigma = RH_sigma * variables.vapor_pressure(T_sigma)
- Td_sigma = variables.dewpoint_from_vapor_pressure(vap_pres_sigma)
- Tv_sigma = variables.virtual_temperature_from_dewpoint(T_sigma, Td_sigma, sp)
- Tw_sigma = variables.wet_bulb_temperature(T_sigma, Td_sigma)
- r_sigma = variables.mixing_ratio_from_dewpoint(Td_sigma, sp) * 1000 # Convert to g/kg
- q_sigma = variables.specific_humidity_from_dewpoint(Td_sigma, sp) * 1000 # Convert to g/kg
- theta_sigma = variables.potential_temperature(T_sigma, sp)
- theta_e_sigma = variables.equivalent_potential_temperature(T_sigma, Td_sigma, sp)
- theta_v_sigma = variables.virtual_potential_temperature(T_sigma, Td_sigma, sp)
- theta_w_sigma = variables.wet_bulb_potential_temperature(T_sigma, Td_sigma, sp)
-
- T = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- Td = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- Tv = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- Tw = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- theta = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- theta_e = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- theta_v = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- theta_w = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- RH = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- r = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- q = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- u = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- v = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
- sp_z = np.empty(shape=(num_forecast_hours, len(pressure_levels) + 1, chunk_sizes['latitude'], chunk_sizes['longitude']))
-
- sp /= 100 # pascals (Pa) --> hectopascals (hPa)
- if len(gpus) > 0:
- T[:, 0, :, :] = T_sigma.numpy()
- T[:, 1:, :, :] = T_pl.numpy()
- Td[:, 0, :, :] = Td_sigma.numpy()
- Td[:, 1:, :, :] = Td_pl.numpy()
- Tv[:, 0, :, :] = Tv_sigma.numpy()
- Tv[:, 1:, :, :] = Tv_pl.numpy()
- Tw[:, 0, :, :] = Tw_sigma.numpy()
- Tw[:, 1:, :, :] = Tw_pl.numpy()
- theta[:, 0, :, :] = theta_sigma.numpy()
- theta[:, 1:, :, :] = theta_pl.numpy()
- theta_e[:, 0, :, :] = theta_e_sigma.numpy()
- theta_e[:, 1:, :, :] = theta_e_pl.numpy()
- theta_v[:, 0, :, :] = theta_v_sigma.numpy()
- theta_v[:, 1:, :, :] = theta_v_pl.numpy()
- theta_w[:, 0, :, :] = theta_w_sigma.numpy()
- theta_w[:, 1:, :, :] = theta_w_pl.numpy()
- RH[:, 0, :, :] = RH_sigma.numpy()
- RH[:, 1:, :, :] = RH_pl.numpy()
- r[:, 0, :, :] = r_sigma.numpy()
- r[:, 1:, :, :] = r_pl.numpy()
- q[:, 0, :, :] = q_sigma.numpy()
- q[:, 1:, :, :] = q_pl.numpy()
- sp_z[:, 0, :, :] = sp.numpy()
- else:
- T[:, 0, :, :] = T_sigma
- T[:, 1:, :, :] = T_pl
- Td[:, 0, :, :] = Td_sigma
- Td[:, 1:, :, :] = Td_pl
- Tv[:, 0, :, :] = Tv_sigma
- Tv[:, 1:, :, :] = Tv_pl
- Tw[:, 0, :, :] = Tw_sigma
- Tw[:, 1:, :, :] = Tw_pl
- theta[:, 0, :, :] = theta_sigma
- theta[:, 1:, :, :] = theta_pl
- theta_e[:, 0, :, :] = theta_e_sigma
- theta_e[:, 1:, :, :] = theta_e_pl
- theta_v[:, 0, :, :] = theta_v_sigma
- theta_v[:, 1:, :, :] = theta_v_pl
- theta_w[:, 0, :, :] = theta_w_sigma
- theta_w[:, 1:, :, :] = theta_w_pl
- RH[:, 0, :, :] = RH_sigma
- RH[:, 1:, :, :] = RH_pl
- r[:, 0, :, :] = r_sigma
- r[:, 1:, :, :] = r_pl
- q[:, 0, :, :] = q_sigma
- q[:, 1:, :, :] = q_pl
- sp_z[:, 0, :, :] = sp
-
- u[:, 0, :, :] = u_sigma
- u[:, 1:, :, :] = u_pl
- v[:, 0, :, :] = v_sigma
- v[:, 1:, :, :] = v_pl
- sp_z[:, 1:, :, :] = z
-
- pressure_levels = ['surface', '1000', '950', '900', '850', '700', '500']
-
- print("Building final dataset")
-
- full_dataset_coordinates = dict(forecast_hour=forecast_hours, pressure_level=pressure_levels)
- full_dataset_variables = dict(T=(dataset_dimensions, T),
- Td=(dataset_dimensions, Td),
- Tv=(dataset_dimensions, Tv),
- Tw=(dataset_dimensions, Tw),
- theta=(dataset_dimensions, theta),
- theta_e=(dataset_dimensions, theta_e),
- theta_v=(dataset_dimensions, theta_v),
- theta_w=(dataset_dimensions, theta_w),
- RH=(dataset_dimensions, RH),
- r=(dataset_dimensions, r),
- q=(dataset_dimensions, q),
- u=(dataset_dimensions, u),
- v=(dataset_dimensions, v),
- sp_z=(dataset_dimensions, sp_z))
-
- if 'mslp_data_file_index' in locals():
- full_dataset_variables['mslp_z'] = (('forecast_hour', 'pressure_level', 'latitude', 'longitude'), mslp_z)
-
- full_dataset_coordinates['latitude'] = pressure_level_data['latitude']
- full_dataset_coordinates['longitude'] = pressure_level_data['longitude']
-
- full_grib_dataset = xr.Dataset(data_vars=full_dataset_variables,
- coords=full_dataset_coordinates).astype('float32')
-
- full_grib_dataset = full_grib_dataset.expand_dims({'time': np.atleast_1d(pressure_level_data['time'].values)})
-
- monthly_dir = '%s/%d%02d' % (args['netcdf_outdir'], year, month)
-
- if not os.path.isdir(monthly_dir):
- os.mkdir(monthly_dir)
-
- for fcst_hr_index, forecast_hour in enumerate(forecast_hours):
- full_grib_dataset.isel(forecast_hour=np.atleast_1d(fcst_hr_index)).to_netcdf(path=f"%s/{args['model'].lower()}_%d%02d%02d%02d_f%03d_global.nc" % (monthly_dir, year, month, day, hour, forecast_hour), mode='w', engine='netcdf4')
-
- if args['delete_split_grib']:
- grib_files = sorted(glob.glob(individual_variable_filename_format + "*"))
- [os.remove(file) for file in grib_files]
diff --git a/convert_netcdf_to_tf.py b/convert_netcdf_to_tf.py
deleted file mode 100644
index 3fe4996..0000000
--- a/convert_netcdf_to_tf.py
+++ /dev/null
@@ -1,404 +0,0 @@
-"""
-Convert netCDF files containing variable and frontal boundary data into tensorflow datasets for model training.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.8.13
-"""
-import argparse
-import itertools
-import numpy as np
-import os
-import pandas as pd
-import pickle
-import tensorflow as tf
-import file_manager as fm
-from utils import data_utils, settings
-import xarray as xr
-
-
-if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- parser.add_argument('--year_and_month', type=int, nargs=2, required=True,
- help="Year and month for the netcdf data to be converted to tensorflow datasets.")
- parser.add_argument('--variable_data_source', type=str, default='era5', help="Data source or model containing the variable data.")
- parser.add_argument('--variables_netcdf_indir', type=str, required=True,
- help="Input directory for the netCDF files containing variable data.")
- parser.add_argument('--fronts_netcdf_indir', type=str, required=True,
- help="Input directory for the netCDF files containing frontal boundary data.")
- parser.add_argument('--tf_outdir', type=str, required=True,
- help="Output directory for the generated tensorflow datasets.")
- parser.add_argument('--front_types', type=str, nargs='+', required=True,
- help="Code(s) for the front types that will be generated in the tensorflow datasets. Refer to documentation in 'utils.data_utils.reformat_fronts' "
- "for more information on these codes.")
- parser.add_argument('--variables', type=str, nargs='+', help='Variables to select')
- parser.add_argument('--pressure_levels', type=str, nargs='+', help='Variables pressure levels to select')
- parser.add_argument('--num_dims', type=int, nargs=2, default=[3, 3], help='Number of dimensions in the variables and front object images, repsectively.')
- parser.add_argument('--domain', type=str, default='conus', help='Domain from which to pull the images.')
- parser.add_argument('--override_extent', type=float, nargs=4,
- help='Override the default domain extent by selecting a custom extent. [min lon, max lon, min lat, max lat]')
- parser.add_argument('--evaluation_dataset', action='store_true',
- help='''
- Boolean flag that determines if the dataset being generated will be used for evaluating a model.
- If this flag is True, all of the following keyword arguments will be set and any values provided to 'netcdf_to_tf'
- by the user will be overriden:
- * num_dims = (_, 2) <=== NOTE: The first value of this tuple will NOT be overriden.
- * images = (1, 1)
- * image_size will be set to the size of the domain.
- * keep_fraction will have no effect
- * shuffle_timesteps = False
- * shuffle_images = False
- * noise_fraction = 0.0
- * rotate_chance = 0.0
- * flip_chance_lon = 0.0
- * flip_chance_lat = 0.0
- ''')
- parser.add_argument('--images', type=int, nargs=2, default=[9, 1],
- help='Number of variables/front images along the longitude and latitude dimensions to generate for each timestep. The product of the 2 integers '
- 'will be the total number of images generated per timestep.')
- parser.add_argument('--image_size', type=int, nargs=2, default=[128, 128], help='Size of the longitude and latitude dimensions of the images.')
- parser.add_argument('--shuffle_timesteps', action='store_true',
- help='Shuffle the timesteps when generating the dataset. This is particularly useful when generating very large '
- 'datasets that cannot be shuffled on the fly during training.')
- parser.add_argument('--shuffle_images', action='store_true',
- help='Shuffle the order of the images in each timestep. This does NOT shuffle the entire dataset for the provided '
- 'month, but rather only the images in each respective timestep. This is particularly useful when generating '
- 'very large datasets that cannot be shuffled on the fly during training.')
- parser.add_argument('--add_previous_fronts', type=str, nargs='+',
- help='Optional front types from previous timesteps to include as predictors. If the dataset is over conus, the fronts '
- 'will be pulled from the last 3-hour timestep. If the dataset is over the full domain, the fronts will be pulled '
- 'from the last 6-hour timestep.')
- parser.add_argument('--front_dilation', type=int, default=0, help='Number of pixels to expand the fronts by in all directions.')
- parser.add_argument('--keep_fraction', type=float, default=0.0,
- help='The fraction of timesteps WITHOUT all necessary front types that will be retained in the dataset. Can be any float 0 <= x <= 1.')
- parser.add_argument('--noise_fraction', type=float, default=0.0,
- help='The fraction of pixels in each image that will contain noise. Can be any float 0 <= x < 1.')
- parser.add_argument('--rotate_chance', type=float, default=0.0,
- help='The probability that the current image will be rotated (in any direction, up to 270 degrees). Can be any float 0 <= x <= 1.')
- parser.add_argument('--flip_chance_lon', type=float, default=0.0,
- help='The probability that the current image will have its longitude dimension reversed. Can be any float 0 <= x <= 1.')
- parser.add_argument('--flip_chance_lat', type=float, default=0.0,
- help='The probability that the current image will have its latitude dimension reversed. Can be any float 0 <= x <= 1.')
- parser.add_argument('--overwrite', action='store_true', help='Overwrite the contents of any existing variables and fronts data.')
- parser.add_argument('--verbose', action='store_true', help='Print out the progress of the dataset generation.')
- parser.add_argument('--gpu_device', type=int, nargs='+', help='GPU device numbers.')
- parser.add_argument('--memory_growth', action='store_true', help='Use memory growth on the GPU(s).')
-
- args = vars(parser.parse_args())
-
- if args['gpu_device'] is not None:
- gpus = tf.config.list_physical_devices(device_type='GPU')
- tf.config.set_visible_devices(devices=[gpus[gpu] for gpu in args['gpu_device']], device_type='GPU')
-
- # Allow for memory growth on the GPU. This will only use the GPU memory that is required rather than allocating all of the GPU's memory.
- if args['memory_growth']:
- tf.config.experimental.set_memory_growth(device=[gpus[gpu] for gpu in args['gpu_device']][0], enable=True)
-
- year, month = args['year_and_month'][0], args['year_and_month'][1]
-
- tf_dataset_folder_variables = f'%s/%s_%d%02d_tf' % (args['tf_outdir'], args['variable_data_source'], year, month)
- tf_dataset_folder_fronts = f"%s/fronts_%d%02d_tf" % (args['tf_outdir'], year, month)
-
- if os.path.isdir(tf_dataset_folder_variables) or os.path.isdir(tf_dataset_folder_fronts):
- if args['overwrite']:
- print("WARNING: Tensorflow dataset(s) already exist for the provided year and month and will be overwritten.")
- else:
- raise FileExistsError("Tensorflow dataset(s) already exist for the provided year and month. If you would like to "
- "overwrite the existing datasets, pass the --overwrite flag into the command line.")
-
- if not os.path.isdir(args['tf_outdir']):
- try:
- os.mkdir(args['tf_outdir'])
- except FileExistsError: # When running in parallel, sometimes multiple instances will try to create this directory at once, resulting in a FileExistsError
- pass
-
- dataset_props_file = '%s/dataset_properties.pkl' % args['tf_outdir']
-
- if not os.path.isfile(dataset_props_file):
- """
- Save critical dataset information to a pickle file so it can be referenced later when generating data for other months.
- """
-
- if args['evaluation_dataset']:
- """
- Override all keyword arguments so the dataset will be prepared for model evaluation.
- """
- print("WARNING: This dataset will be used for model evaluation, so the following arguments will be set and "
- "any provided values for these arguments will be overriden:")
- args['num_dims'] = tuple(args['num_dims'])
- args['images'] = (1, 1)
-
- if args['override_extent'] is None:
- args['image_size'] = (settings.DEFAULT_DOMAIN_INDICES[args['domain']][1] - settings.DEFAULT_DOMAIN_INDICES[args['domain']][0],
- settings.DEFAULT_DOMAIN_INDICES[args['domain']][3] - settings.DEFAULT_DOMAIN_INDICES[args['domain']][2])
- else:
- args['image_size'] = (int((args['override_extent'][1] - args['override_extent'][0]) / 0.25 + 1),
- int((args['override_extent'][3] - args['override_extent'][2]) / 0.25 + 1))
-
- args['shuffle_timesteps'] = False
- args['shuffle_images'] = False
- args['noise_fraction'] = 0.0
- args['rotate_chance'] = 0.0
- args['flip_chance_lon'] = 0.0
- args['flip_chance_lat'] = 0.0
-
- print(f"images = {args['images']}\n"
- f"image_size = {args['image_size']}\n"
- f"shuffle_timesteps = False\n"
- f"shuffle_images = False\n"
- f"noise_fraction = 0.0\n"
- f"rotate_chance = 0.0\n"
- f"flip_chance_lon = 0.0\n"
- f"flip_chance_lat = 0.0\n")
-
- dataset_props = dict({})
- dataset_props['normalization_parameters'] = data_utils.normalization_parameters
- for key in sorted(['front_types', 'variables', 'pressure_levels', 'num_dims', 'images', 'image_size', 'front_dilation',
- 'noise_fraction', 'rotate_chance', 'flip_chance_lon', 'flip_chance_lat', 'shuffle_images', 'shuffle_timesteps',
- 'domain', 'evaluation_dataset', 'add_previous_fronts', 'keep_fraction', 'override_extent']):
- dataset_props[key] = args[key]
-
- with open(dataset_props_file, 'wb') as f:
- pickle.dump(dataset_props, f)
-
- with open('%s/dataset_properties.txt' % args['tf_outdir'], 'w') as f:
- for key in sorted(dataset_props.keys()):
- f.write(f"{key}: {dataset_props[key]}\n")
-
- else:
-
- print("WARNING: Dataset properties file was found in %s. The following settings will be used from the file." % args['tf_outdir'])
- dataset_props = pd.read_pickle(dataset_props_file)
-
- for key in sorted(['front_types', 'variables', 'pressure_levels', 'num_dims', 'images', 'image_size', 'front_dilation',
- 'noise_fraction', 'rotate_chance', 'flip_chance_lon', 'flip_chance_lat', 'shuffle_images', 'shuffle_timesteps',
- 'domain', 'evaluation_dataset', 'add_previous_fronts', 'keep_fraction']):
- args[key] = dataset_props[key]
- print(f"%s: {args[key]}" % key)
-
- all_variables = ['T', 'Td', 'sp_z', 'u', 'v', 'theta_w', 'r', 'RH', 'Tv', 'Tw', 'theta_e', 'q', 'theta', 'theta_v']
- all_pressure_levels = ['surface', '1000', '950', '900', '850'] if args['variable_data_source'] == 'era5' else ['surface', '1000', '950', '900', '850', '700', '500']
-
- synoptic_only = True if args['domain'] == 'full' else False
-
- file_loader = fm.DataFileLoader(args['variables_netcdf_indir'], '%s-netcdf' % args['variable_data_source'], synoptic_only)
- file_loader.pair_with_fronts(args['fronts_netcdf_indir'])
-
- variables_netcdf_files = file_loader.data_files
- fronts_netcdf_files = file_loader.front_files
-
- print(stop)
-
- ### Grab front files from previous timesteps so previous fronts can be used as predictors ###
- if args['add_previous_fronts'] is not None:
- files_to_remove = [] # variables and front files that will be removed from the dataset
- previous_fronts_netcdf_files = []
- for file in fronts_netcdf_files:
- current_timestep = np.datetime64(f'{file[-18:-14]}-{file[-14:-12]}-{file[-12:-10]}T{file[-10:-8]}')
- previous_timestep = (current_timestep - np.timedelta64(3, "h")).astype(object)
- prev_year, prev_month, prev_day, prev_hour = previous_timestep.year, previous_timestep.month, previous_timestep.day, previous_timestep.hour
- previous_fronts_file = '%s/%d%02d/FrontObjects_%d%02d%02d%02d_full.nc' % (args['fronts_netcdf_indir'], prev_year, prev_month, prev_year, prev_month, prev_day, prev_hour)
- if os.path.isfile(previous_fronts_file):
- previous_fronts_netcdf_files.append(previous_fronts_file) # Add the previous fronts to the dataset
- else:
- files_to_remove.append(file)
-
- ### Remove files from the dataset if previous fronts are not available ###
- if len(files_to_remove) > 0:
- for file in files_to_remove:
- index_to_pop = fronts_netcdf_files.index(file)
- variables_netcdf_files.pop(index_to_pop), fronts_netcdf_files.pop(index_to_pop)
-
- if args['shuffle_timesteps']:
- zipped_list = list(zip(variables_netcdf_files, fronts_netcdf_files))
- np.random.shuffle(zipped_list)
- variables_netcdf_files, fronts_netcdf_files = zip(*zipped_list)
-
- # assert that the dates of the files match
- files_match_flag = all(os.path.basename(variables_file).split('_')[1] == os.path.basename(fronts_file).split('_')[1] for variables_file, fronts_file in zip(variables_netcdf_files, fronts_netcdf_files))
-
- if args['override_extent'] is None:
- isel_kwargs = {'longitude': slice(settings.DEFAULT_DOMAIN_INDICES[args['domain']][0], settings.DEFAULT_DOMAIN_INDICES[args['domain']][1]),
- 'latitude': slice(settings.DEFAULT_DOMAIN_INDICES[args['domain']][2], settings.DEFAULT_DOMAIN_INDICES[args['domain']][3])}
- domain_size = (int(settings.DEFAULT_DOMAIN_INDICES[args['domain']][1] - settings.DEFAULT_DOMAIN_INDICES[args['domain']][0]),
- int(settings.DEFAULT_DOMAIN_INDICES[args['domain']][3] - settings.DEFAULT_DOMAIN_INDICES[args['domain']][2]))
- else:
- isel_kwargs = {'longitude': slice(int((args['override_extent'][0] - settings.DEFAULT_DOMAIN_EXTENTS[args['domain']][0]) // 0.25),
- int((args['override_extent'][1] - settings.DEFAULT_DOMAIN_EXTENTS[args['domain']][0]) // 0.25) + 1),
- 'latitude': slice(int((settings.DEFAULT_DOMAIN_EXTENTS[args['domain']][3] - args['override_extent'][3]) // 0.25),
- int((settings.DEFAULT_DOMAIN_EXTENTS[args['domain']][3] - args['override_extent'][2]) // 0.25) + 1)}
- domain_size = (int((args['override_extent'][1] - args['override_extent'][0]) // 0.25),
- int((args['override_extent'][3] - args['override_extent'][2]) // 0.25))
-
- if not files_match_flag:
- raise OSError("%s/fronts files do not match")
-
- variables_to_use = all_variables if args['variables'] is None else args['variables']
- args['pressure_levels'] = all_pressure_levels if args['pressure_levels'] is None else [lvl for lvl in all_pressure_levels if lvl in args['pressure_levels']]
-
- num_timesteps = len(variables_netcdf_files)
- timesteps_kept = 0
- timesteps_discarded = 0
-
- for timestep_no in range(num_timesteps):
-
- front_dataset = xr.open_dataset(fronts_netcdf_files[timestep_no], engine='netcdf4').isel(**isel_kwargs).astype('float16')
-
- ### Reformat the fronts in the current timestep ###
- if args['front_types'] is not None:
- front_dataset = data_utils.reformat_fronts(front_dataset, args['front_types'])
- num_front_types = front_dataset.attrs['num_types'] + 1
- else:
- num_front_types = 17
-
- if args['front_dilation'] > 0:
- front_dataset = data_utils.expand_fronts(front_dataset, iterations=args['front_dilation']) # expand the front labels
-
- keep_timestep = np.random.random() <= args['keep_fraction'] # boolean flag for keeping timesteps without all front types
-
- front_dataset = front_dataset.isel(time=0).to_array().transpose('longitude', 'latitude', 'variable')
- front_bins = np.bincount(front_dataset.values.astype('int64').flatten(), minlength=num_front_types) # counts for each front type
- all_fronts_present = all([front_count > 0 for front_count in front_bins]) > 0 # boolean flag that says if all front types are present in the current timestep
-
- if all_fronts_present or keep_timestep or args['evaluation_dataset']:
-
- if args['variable_data_source'] != 'era5':
- isel_kwargs['forecast_hour'] = 0
-
- variables_dataset = xr.open_dataset(variables_netcdf_files[timestep_no], engine='netcdf4')[variables_to_use].isel(**isel_kwargs).sel(pressure_level=args['pressure_levels']).transpose('time', 'longitude', 'latitude', 'pressure_level').astype('float16')
- variables_dataset = data_utils.normalize_variables(variables_dataset).isel(time=0).transpose('longitude', 'latitude', 'pressure_level').astype('float16')
-
- ### Reformat the fronts from the previous timestep ###
- if args['add_previous_fronts'] is not None:
- previous_front_dataset = xr.open_dataset(previous_fronts_netcdf_files[timestep_no], engine='netcdf4').isel(**isel_kwargs).astype('float16')
- previous_front_dataset = data_utils.reformat_fronts(previous_front_dataset, args['add_previous_fronts'])
-
- if args['front_dilation'] > 0:
- previous_front_dataset = data_utils.expand_fronts(previous_front_dataset, iterations=args['front_dilation'])
-
- previous_front_dataset = previous_front_dataset.transpose('longitude', 'latitude')
-
- previous_fronts = np.zeros([len(previous_front_dataset['longitude'].values),
- len(previous_front_dataset['latitude'].values),
- len(args['pressure_levels'])], dtype=np.float16)
-
- for front_type_no, previous_front_type in enumerate(args['add_previous_fronts']):
- previous_fronts[..., 0] = np.where(previous_front_dataset['identifier'].values == front_type_no + 1, 1, 0) # Place previous front labels at the surface level
- variables_dataset[previous_front_type] = (('longitude', 'latitude', 'pressure_level'), previous_fronts) # Add previous fronts to the predictor dataset
-
- variables_dataset = variables_dataset.to_array().transpose('longitude', 'latitude', 'pressure_level', 'variable')
-
- if args['override_extent'] is None:
- if args['images'][0] > 1 and domain_size[0] > args['image_size'][0] + args['images'][0]:
- start_indices_lon = np.linspace(0, settings.DEFAULT_DOMAIN_INDICES[args['domain']][1] - settings.DEFAULT_DOMAIN_INDICES[args['domain']][0] - args['image_size'][0],
- args['images'][0]).astype(int)
- else:
- start_indices_lon = np.zeros((args['images'][0], ), dtype=int)
-
- if args['images'][1] > 1 and domain_size[1] > args['image_size'][1] + args['images'][1]:
- start_indices_lat = np.linspace(0, settings.DEFAULT_DOMAIN_INDICES[args['domain']][3] - settings.DEFAULT_DOMAIN_INDICES[args['domain']][2] - args['image_size'][1],
- args['images'][1]).astype(int)
- else:
- start_indices_lat = np.zeros((args['images'][1], ), dtype=int)
-
- else:
- if args['images'][0] > 1 and domain_size[0] > args['image_size'][0] + args['images'][0]:
- start_indices_lon = np.linspace(0, domain_size[0] - args['image_size'][0], args['images'][0]).astype(int)
- else:
- start_indices_lon = np.zeros((args['images'][0], ), dtype=int)
-
- if args['images'][1] > 1 and domain_size[1] > args['image_size'][1] + args['images'][1]:
- start_indices_lat = np.linspace(0, domain_size[1] - args['image_size'][1], args['images'][1]).astype(int)
- else:
- start_indices_lat = np.zeros((args['images'][1], ), dtype=int)
-
- image_order = list(itertools.product(start_indices_lon, start_indices_lat)) # Every possible combination of longitude and latitude starting points
-
- if args['shuffle_images']:
- np.random.shuffle(image_order)
-
- for image_start_indices in image_order:
-
- start_index_lon = image_start_indices[0]
- end_index_lon = start_index_lon + args['image_size'][0]
- start_index_lat = image_start_indices[1]
- end_index_lat = start_index_lat + args['image_size'][1]
-
- # boolean flags for rotating and flipping images
- rotate_image = np.random.random() <= args['rotate_chance']
- flip_lon = np.random.random() <= args['flip_chance_lon']
- flip_lat = np.random.random() <= args['flip_chance_lat']
-
- if rotate_image:
- rotation_direction = np.random.randint(0, 2) # 0 = clockwise, 1 = counter-clockwise
- num_rotations = np.random.randint(1, 4) # n * 90 degrees
-
- variables_tensor = tf.convert_to_tensor(variables_dataset[start_index_lon:end_index_lon, start_index_lat:end_index_lat, :, :], dtype=tf.float16)
- if flip_lon:
- variables_tensor = tf.reverse(variables_tensor, axis=[0]) # Reverse values along the longitude dimension
- if flip_lat:
- variables_tensor = tf.reverse(variables_tensor, axis=[1]) # Reverse values along the latitude dimension
- if rotate_image:
- for rotation in range(num_rotations):
- variables_tensor = tf.reverse(tf.transpose(variables_tensor, perm=[1, 0, 2, 3]), axis=[rotation_direction]) # Rotate image 90 degrees
-
- if args['noise_fraction'] > 0:
- ### Add noise to image ###
- random_values = tf.random.uniform(shape=variables_tensor.shape)
- variables_tensor = tf.where(random_values < args['noise_fraction'] / 2, 0.0, variables_tensor) # add 0s to image
- variables_tensor = tf.where(random_values > 1.0 - (args['noise_fraction'] / 2), 1.0, variables_tensor) # add 1s to image
-
- if args['num_dims'][0] == 2:
- variables_tensor_shape_3d = variables_tensor.shape
- # Combine pressure level and variables dimensions, making the images 2D (excluding the final dimension)
- variables_tensor = tf.reshape(variables_tensor, [variables_tensor_shape_3d[0], variables_tensor_shape_3d[1], variables_tensor_shape_3d[2] * variables_tensor_shape_3d[3]])
-
- variables_tensor_for_timestep = tf.data.Dataset.from_tensors(variables_tensor)
- if 'variables_tensors_for_month' not in locals():
- variables_tensors_for_month = variables_tensor_for_timestep
- else:
- variables_tensors_for_month = variables_tensors_for_month.concatenate(variables_tensor_for_timestep)
-
- front_tensor = tf.convert_to_tensor(front_dataset[start_index_lon:end_index_lon, start_index_lat:end_index_lat, :], dtype=tf.int32)
-
- if flip_lon:
- front_tensor = tf.reverse(front_tensor, axis=[0]) # Reverse values along the longitude dimension
- if flip_lat:
- front_tensor = tf.reverse(front_tensor, axis=[1]) # Reverse values along the latitude dimension
- if rotate_image:
- for rotation in range(num_rotations):
- front_tensor = tf.reverse(tf.transpose(front_tensor, perm=[1, 0, 2]), axis=[rotation_direction]) # Rotate image 90 degrees
-
- if args['num_dims'][1] == 3:
- # Make the front object images 3D, with the size of the 3rd dimension equal to the number of pressure levels
- front_tensor = tf.tile(front_tensor, (1, 1, len(args['pressure_levels'])))
- else:
- front_tensor = front_tensor[:, :, 0]
-
- front_tensor = tf.cast(tf.one_hot(front_tensor, num_front_types), tf.float16) # One-hot encode the labels
- front_tensor_for_timestep = tf.data.Dataset.from_tensors(front_tensor)
- if 'front_tensors_for_month' not in locals():
- front_tensors_for_month = front_tensor_for_timestep
- else:
- front_tensors_for_month = front_tensors_for_month.concatenate(front_tensor_for_timestep)
-
- timesteps_kept += 1
- else:
- timesteps_discarded += 1
-
- if args['verbose']:
- print("Timesteps complete: %d/%d (Retained/discarded: %d/%d)" % (timesteps_kept + timesteps_discarded, num_timesteps, timesteps_kept, timesteps_discarded), end='\r')
-
- print("Timesteps complete: %d/%d (Retained/discarded: %d/%d)" % (timesteps_kept + timesteps_discarded, num_timesteps, timesteps_kept, timesteps_discarded))
-
- if args['overwrite']:
- if os.path.isdir(tf_dataset_folder_variables):
- os.rmdir(tf_dataset_folder_variables)
- if os.path.isdir(tf_dataset_folder_fronts):
- os.rmdir(tf_dataset_folder_fronts)
-
- try:
- tf.data.Dataset.save(variables_tensors_for_month, path=tf_dataset_folder_variables)
- tf.data.Dataset.save(front_tensors_for_month, path=tf_dataset_folder_fronts)
- print("Tensorflow datasets for %d-%02d saved to %s." % (year, month, args['tf_outdir']))
- except NameError:
- print("No images could be retained with the provided arguments.")
diff --git a/coordinates/hrrr.nc b/coordinates/hrrr.nc
new file mode 100644
index 0000000..7440400
Binary files /dev/null and b/coordinates/hrrr.nc differ
diff --git a/coordinates/nam-12km.nc b/coordinates/nam-12km.nc
new file mode 100644
index 0000000..d8d107d
Binary files /dev/null and b/coordinates/nam-12km.nc differ
diff --git a/coordinates/namnest-conus.nc b/coordinates/namnest-conus.nc
new file mode 100644
index 0000000..5ebcc7a
Binary files /dev/null and b/coordinates/namnest-conus.nc differ
diff --git a/create_era5_netcdf.py b/create_era5_netcdf.py
deleted file mode 100644
index 7ec0522..0000000
--- a/create_era5_netcdf.py
+++ /dev/null
@@ -1,189 +0,0 @@
-"""
-Create ERA5 netCDF datasets.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.6.7
-"""
-
-import argparse
-import numpy as np
-import os
-from utils import variables
-import xarray as xr
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument('--netcdf_era5_indir', type=str, required=True, help="Input directory for the global ERA5 netCDF files.")
- parser.add_argument('--netcdf_outdir', type=str, required=True, help="Output directory for front netCDF files.")
- parser.add_argument('--date', type=int, nargs=3, required=True, help="Date for the data to be read in. (year, month, day)")
-
- args = vars(parser.parse_args())
-
- year, month, day = args['date'][0], args['date'][1], args['date'][2]
-
- era5_T_sfc_file = 'ERA5Global_%d_3hrly_2mT.nc' % year
- era5_Td_sfc_file = 'ERA5Global_%d_3hrly_2mTd.nc' % year
- era5_sp_file = 'ERA5Global_%d_3hrly_sp.nc' % year
- era5_u_sfc_file = 'ERA5Global_%d_3hrly_U10m.nc' % year
- era5_v_sfc_file = 'ERA5Global_%d_3hrly_V10m.nc' % year
-
- timestring = "%d-%02d-%02d" % (year, month, day)
-
- lons = np.append(np.arange(130, 360, 0.25), np.arange(0, 10.25, 0.25))
- lats = np.arange(0, 80.25, 0.25)[::-1]
- lons360 = np.arange(130, 370.25, 0.25)
-
- T_sfc_full_day = xr.open_mfdataset("%s/Surface/%s" % (args['netcdf_era5_indir'], era5_T_sfc_file), chunks={'latitude': 721, 'longitude': 1440, 'time': 4}).sel(time=('%s' % timestring), longitude=lons, latitude=lats)
- Td_sfc_full_day = xr.open_mfdataset("%s/Surface/%s" % (args['netcdf_era5_indir'], era5_Td_sfc_file), chunks={'latitude': 721, 'longitude': 1440, 'time': 4}).sel(time=('%s' % timestring), longitude=lons, latitude=lats)
- sp_full_day = xr.open_mfdataset("%s/Surface/%s" % (args['netcdf_era5_indir'], era5_sp_file), chunks={'latitude': 721, 'longitude': 1440, 'time': 4}).sel(time=('%s' % timestring), longitude=lons, latitude=lats)
- u_sfc_full_day = xr.open_mfdataset("%s/Surface/%s" % (args['netcdf_era5_indir'], era5_u_sfc_file), chunks={'latitude': 721, 'longitude': 1440, 'time': 4}).sel(time=('%s' % timestring), longitude=lons, latitude=lats)
- v_sfc_full_day = xr.open_mfdataset("%s/Surface/%s" % (args['netcdf_era5_indir'], era5_v_sfc_file), chunks={'latitude': 721, 'longitude': 1440, 'time': 4}).sel(time=('%s' % timestring), longitude=lons, latitude=lats)
-
- PL_data = xr.open_mfdataset(
- paths=('%s/Pressure_Level/ERA5Global_PL_%s_3hrly_Q.nc' % (args['netcdf_era5_indir'], year),
- '%s/Pressure_Level/ERA5Global_PL_%s_3hrly_T.nc' % (args['netcdf_era5_indir'], year),
- '%s/Pressure_Level/ERA5Global_PL_%s_3hrly_U.nc' % (args['netcdf_era5_indir'], year),
- '%s/Pressure_Level/ERA5Global_PL_%s_3hrly_V.nc' % (args['netcdf_era5_indir'], year),
- '%s/Pressure_Level/ERA5Global_PL_%s_3hrly_Z.nc' % (args['netcdf_era5_indir'], year)),
- chunks={'latitude': 721, 'longitude': 1440, 'time': 4}).sel(time=('%s' % timestring), longitude=lons, latitude=lats)
-
- if not os.path.isdir('%s/%d%02d' % (args['netcdf_outdir'], year, month)):
- os.mkdir('%s/%d%02d' % (args['netcdf_outdir'], year, month))
-
- for hour in range(0, 24, 3):
-
- print(f"saving ERA5 data for {year}-%02d-%02d-%02dz" % (month, day, hour))
-
- timestep = '%d-%02d-%02dT%02d:00:00' % (year, month, day, hour)
-
- PL_850 = PL_data.sel(level=850, time=timestep)
- PL_900 = PL_data.sel(level=900, time=timestep)
- PL_950 = PL_data.sel(level=950, time=timestep)
- PL_1000 = PL_data.sel(level=1000, time=timestep)
-
- T_sfc = T_sfc_full_day.sel(time=timestep)['t2m'].values
- Td_sfc = Td_sfc_full_day.sel(time=timestep)['d2m'].values
- sp = sp_full_day.sel(time=timestep)['sp'].values
- u_sfc = u_sfc_full_day.sel(time=timestep)['u10'].values
- v_sfc = v_sfc_full_day.sel(time=timestep)['v10'].values
-
- theta_sfc = variables.potential_temperature(T_sfc, sp) # Potential temperature
- theta_e_sfc = variables.equivalent_potential_temperature(T_sfc, Td_sfc, sp) # Equivalent potential temperature
- theta_v_sfc = variables.virtual_temperature_from_dewpoint(T_sfc, Td_sfc, sp) # Virtual potential temperature
- theta_w_sfc = variables.wet_bulb_potential_temperature(T_sfc, Td_sfc, sp) # Wet-bulb potential temperature
- r_sfc = variables.mixing_ratio_from_dewpoint(Td_sfc, sp) # Mixing ratio
- q_sfc = variables.specific_humidity_from_dewpoint(Td_sfc, sp) # Specific humidity
- RH_sfc = variables.relative_humidity(T_sfc, Td_sfc) # Relative humidity
- Tv_sfc = variables.virtual_temperature_from_dewpoint(T_sfc, Td_sfc, sp) # Virtual temperature
- Tw_sfc = variables.wet_bulb_temperature(T_sfc, Td_sfc) # Wet-bulb temperature
-
- q_850 = PL_850['q'].values
- q_900 = PL_900['q'].values
- q_950 = PL_950['q'].values
- q_1000 = PL_1000['q'].values
- T_850 = PL_850['t'].values
- T_900 = PL_900['t'].values
- T_950 = PL_950['t'].values
- T_1000 = PL_1000['t'].values
- u_850 = PL_850['u'].values
- u_900 = PL_900['u'].values
- u_950 = PL_950['u'].values
- u_1000 = PL_1000['u'].values
- v_850 = PL_850['v'].values
- v_900 = PL_900['v'].values
- v_950 = PL_950['v'].values
- v_1000 = PL_1000['v'].values
- z_850 = PL_850['z'].values
- z_900 = PL_900['z'].values
- z_950 = PL_950['z'].values
- z_1000 = PL_1000['z'].values
-
- Td_850 = variables.dewpoint_from_specific_humidity(85000, T_850, q_850)
- Td_900 = variables.dewpoint_from_specific_humidity(90000, T_900, q_900)
- Td_950 = variables.dewpoint_from_specific_humidity(95000, T_950, q_950)
- Td_1000 = variables.dewpoint_from_specific_humidity(100000, T_1000, q_1000)
- r_850 = variables.mixing_ratio_from_dewpoint(Td_850, 85000)
- r_900 = variables.mixing_ratio_from_dewpoint(Td_900, 90000)
- r_950 = variables.mixing_ratio_from_dewpoint(Td_950, 95000)
- r_1000 = variables.mixing_ratio_from_dewpoint(Td_1000, 100000)
- RH_850 = variables.relative_humidity(T_850, Td_850)
- RH_900 = variables.relative_humidity(T_900, Td_900)
- RH_950 = variables.relative_humidity(T_950, Td_950)
- RH_1000 = variables.relative_humidity(T_1000, Td_1000)
- theta_850 = variables.potential_temperature(T_850, 85000)
- theta_900 = variables.potential_temperature(T_900, 90000)
- theta_950 = variables.potential_temperature(T_950, 95000)
- theta_1000 = variables.potential_temperature(T_1000, 100000)
- theta_e_850 = variables.equivalent_potential_temperature(T_850, Td_850, 85000)
- theta_e_900 = variables.equivalent_potential_temperature(T_900, Td_900, 90000)
- theta_e_950 = variables.equivalent_potential_temperature(T_950, Td_950, 95000)
- theta_e_1000 = variables.equivalent_potential_temperature(T_1000, Td_1000, 100000)
- theta_v_850 = variables.virtual_temperature_from_dewpoint(T_850, Td_850, 85000)
- theta_v_900 = variables.virtual_temperature_from_dewpoint(T_900, Td_900, 90000)
- theta_v_950 = variables.virtual_temperature_from_dewpoint(T_950, Td_950, 95000)
- theta_v_1000 = variables.virtual_temperature_from_dewpoint(T_1000, Td_1000, 100000)
- theta_w_850 = variables.wet_bulb_potential_temperature(T_850, Td_850, 85000)
- theta_w_900 = variables.wet_bulb_potential_temperature(T_900, Td_900, 90000)
- theta_w_950 = variables.wet_bulb_potential_temperature(T_950, Td_950, 95000)
- theta_w_1000 = variables.wet_bulb_potential_temperature(T_1000, Td_1000, 100000)
- Tv_850 = variables.virtual_temperature_from_dewpoint(T_850, Td_850, 85000)
- Tv_900 = variables.virtual_temperature_from_dewpoint(T_900, Td_900, 90000)
- Tv_950 = variables.virtual_temperature_from_dewpoint(T_950, Td_950, 95000)
- Tv_1000 = variables.virtual_temperature_from_dewpoint(T_1000, Td_1000, 100000)
- Tw_850 = variables.wet_bulb_temperature(T_850, Td_850)
- Tw_900 = variables.wet_bulb_temperature(T_900, Td_900)
- Tw_950 = variables.wet_bulb_temperature(T_950, Td_950)
- Tw_1000 = variables.wet_bulb_temperature(T_1000, Td_1000)
-
- pressure_levels = ['surface', 1000, 950, 900, 850]
-
- T = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- Td = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- Tv = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- Tw = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- theta = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- theta_e = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- theta_v = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- theta_w = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- RH = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- r = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- q = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- u = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- v = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
- sp_z = np.empty(shape=(len(pressure_levels), len(lats), len(lons360)))
-
- T[0, :, :], T[1, :, :], T[2, :, :], T[3, :, :], T[4, :, :] = T_sfc, T_1000, T_950, T_900, T_850
- Td[0, :, :], Td[1, :, :], Td[2, :, :], Td[3, :, :], Td[4, :, :] = Td_sfc, Td_1000, Td_950, Td_900, Td_850
- Tv[0, :, :], Tv[1, :, :], Tv[2, :, :], Tv[3, :, :], Tv[4, :, :] = Tv_sfc, Tv_1000, Tv_950, Tv_900, Tv_850
- Tw[0, :, :], Tw[1, :, :], Tw[2, :, :], Tw[3, :, :], Tw[4, :, :] = Tw_sfc, Tw_1000, Tw_950, Tw_900, Tw_850
- theta[0, :, :], theta[1, :, :], theta[2, :, :], theta[3, :, :], theta[4, :, :] = theta_sfc, theta_1000, theta_950, theta_900, theta_850
- theta_e[0, :, :], theta_e[1, :, :], theta_e[2, :, :], theta_e[3, :, :], theta_e[4, :, :] = theta_e_sfc, theta_e_1000, theta_e_950, theta_e_900, theta_e_850
- theta_v[0, :, :], theta_v[1, :, :], theta_v[2, :, :], theta_v[3, :, :], theta_v[4, :, :] = theta_v_sfc, theta_v_1000, theta_v_950, theta_v_900, theta_v_850
- theta_w[0, :, :], theta_w[1, :, :], theta_w[2, :, :], theta_w[3, :, :], theta_w[4, :, :] = theta_w_sfc, theta_w_1000, theta_w_950, theta_w_900, theta_w_850
- RH[0, :, :], RH[1, :, :], RH[2, :, :], RH[3, :, :], RH[4, :, :] = RH_sfc, RH_1000, RH_950, RH_900, RH_850
- r[0, :, :], r[1, :, :], r[2, :, :], r[3, :, :], r[4, :, :] = r_sfc, r_1000, r_950, r_900, r_850
- q[0, :, :], q[1, :, :], q[2, :, :], q[3, :, :], q[4, :, :] = q_sfc, q_1000, q_950, q_900, q_850
- u[0, :, :], u[1, :, :], u[2, :, :], u[3, :, :], u[4, :, :] = u_sfc, u_1000, u_950, u_900, u_850
- v[0, :, :], v[1, :, :], v[2, :, :], v[3, :, :], v[4, :, :] = v_sfc, v_1000, v_950, v_900, v_850
- sp_z[0, :, :], sp_z[1, :, :], sp_z[2, :, :], sp_z[3, :, :], sp_z[4, :, :] = sp/100, z_1000/98.0665, z_950/98.0665, z_900/98.0665, z_850/98.0665
-
- full_era5_dataset = xr.Dataset(data_vars=dict(T=(('pressure_level', 'latitude', 'longitude'), T),
- Td=(('pressure_level', 'latitude', 'longitude'), Td),
- Tv=(('pressure_level', 'latitude', 'longitude'), Tv),
- Tw=(('pressure_level', 'latitude', 'longitude'), Tw),
- theta=(('pressure_level', 'latitude', 'longitude'), theta),
- theta_e=(('pressure_level', 'latitude', 'longitude'), theta_e),
- theta_v=(('pressure_level', 'latitude', 'longitude'), theta_v),
- theta_w=(('pressure_level', 'latitude', 'longitude'), theta_w),
- RH=(('pressure_level', 'latitude', 'longitude'), RH),
- r=(('pressure_level', 'latitude', 'longitude'), r * 1000),
- q=(('pressure_level', 'latitude', 'longitude'), q * 1000),
- u=(('pressure_level', 'latitude', 'longitude'), u),
- v=(('pressure_level', 'latitude', 'longitude'), v),
- sp_z=(('pressure_level', 'latitude', 'longitude'), sp_z)),
- coords=dict(pressure_level=pressure_levels, latitude=lats, longitude=lons360)).astype('float32')
-
- full_era5_dataset = full_era5_dataset.expand_dims({'time': np.atleast_1d(timestep)})
-
- full_era5_dataset.to_netcdf(path='%s/%d%02d/era5_%d%02d%02d%02d_full.nc' % (args['netcdf_outdir'], year, month, year, month, day, hour), mode='w', engine='netcdf4')
diff --git a/custom_activations.py b/custom_activations.py
deleted file mode 100644
index 68c3283..0000000
--- a/custom_activations.py
+++ /dev/null
@@ -1,99 +0,0 @@
-"""
-Custom activation functions:
- - Gaussian
- - GCU (Growing Cosine Unit)
- - SmeLU (Smooth ReLU)
- - Snake
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.3.3
-"""
-from tensorflow.keras.layers import Layer
-import tensorflow as tf
-
-
-class Gaussian(Layer):
- """
- Gaussian function activation layer.
- """
- def __init__(self, name=None):
- super(Gaussian, self).__init__(name=name)
-
- def build(self, input_shape):
- """ Build the Gaussian layer """
-
- def call(self, inputs):
- """ Call the Gaussian activation function """
- inputs = tf.cast(inputs, 'float32')
- square_tensor = tf.constant(2.0, shape=inputs.shape[1:])
- y = tf.math.exp(tf.math.negative(tf.math.pow(inputs, square_tensor)))
-
- return y
-
-
-class GCU(Layer):
- """
- Growing Cosine Unit (GCU) activation layer.
- """
- def __init__(self, name=None):
- super(GCU, self).__init__(name=name)
-
- def build(self, input_shape):
- """ Build the GCU layer """
-
- def call(self, inputs):
- """ Call the GCU activation function """
- inputs = tf.cast(inputs, 'float32')
- y = tf.multiply(inputs, tf.math.cos(inputs))
-
- return y
-
-
-class SmeLU(Layer):
- """
- SmeLU (Smooth ReLU) activation function layer for deep learning models.
-
- References
- ----------
- https://arxiv.org/pdf/2202.06499.pdf
- """
- def __init__(self, name=None):
- super(SmeLU, self).__init__(name=name)
-
- def build(self, input_shape):
- """ Build the SmeLU layer """
- self.beta = self.add_weight(name='beta', dtype='float32', shape=input_shape[1:]) # Learnable parameter (see Eq. 7 in the linked paper above)
-
- def call(self, inputs):
- """ Call the SmeLU activation function """
- inputs = tf.cast(inputs, 'float32')
- y = tf.where(inputs <= -self.beta, 0.0, # Condition 1
- tf.where(tf.abs(inputs) <= self.beta, tf.math.divide(tf.math.pow(inputs + self.beta, 2.0), tf.math.multiply(4.0, self.beta)), # Condition 2
- inputs)) # Condition 3 (if x >= beta)
-
- return y
-
-
-class Snake(Layer):
- """
- Snake activation function layer for deep learning models.
-
- References
- ----------
- https://arxiv.org/pdf/2006.08195.pdf
- """
- def __init__(self, name=None):
- super(Snake, self).__init__(name=name)
-
- def build(self, input_shape):
- """ Build the Snake layer """
- self.alpha = self.add_weight(name='alpha', dtype='float32', shape=input_shape[1:]) # Learnable parameter (see Eq. 3 in the linked paper above)
- self.square_tensor = tf.constant(2.0, shape=input_shape[1:])
-
- def call(self, inputs):
- """ Call the Snake activation function """
- inputs = tf.cast(inputs, 'float32')
- y = inputs + tf.multiply(tf.divide(tf.constant(1.0, shape=inputs.shape[1:]), self.alpha), tf.math.pow(tf.math.sin(tf.multiply(self.alpha, inputs)), self.square_tensor))
-
- return y
-
diff --git a/custom_losses.py b/custom_losses.py
deleted file mode 100644
index 25a61b3..0000000
--- a/custom_losses.py
+++ /dev/null
@@ -1,181 +0,0 @@
-"""
-Custom loss functions for U-Net models.
- - Brier Skill Score (BSS)
- - Critical Success Index (CSI)
- - Fractions Skill Score (FSS)
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.5.20.D1
-"""
-import tensorflow as tf
-
-
-def brier_skill_score(class_weights: list = None):
- """
- Brier skill score (BSS) loss function.
-
- class_weights: list of values or None
- List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
- """
-
- @tf.function
- def bss_loss(y_true, y_pred):
- """
- y_true: tf.Tensor
- One-hot encoded tensor containing labels.
- y_pred: tf.Tensor
- Tensor containing model predictions.
- """
-
- losses = tf.math.square(tf.subtract(y_true, y_pred))
-
- if class_weights is not None:
- relative_class_weights = tf.cast(class_weights / tf.math.reduce_sum(class_weights), tf.float32)
- losses *= relative_class_weights
-
- brier_score_loss = tf.math.reduce_sum(losses) / tf.size(losses)
- return brier_score_loss
-
- return bss_loss
-
-
-def critical_success_index(threshold: float = None,
- class_weights: list[int | float] = None):
- """
- Critical Success Index (CSI) loss function.
-
- y_true: tf.Tensor
- One-hot encoded tensor containing labels.
- y_pred: tf.Tensor
- Tensor containing model predictions.
- threshold: float or None
- Optional probability threshold that binarizes y_pred. Values in y_pred greater than or equal to the threshold are
- set to 1, and 0 otherwise.
- If the threshold is set, it must be greater than 0 and less than 1.
- class_weights: list of values or None
- List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
- """
-
- @tf.function
- def csi_loss(y_true, y_pred):
- """
- y_true: tf.Tensor
- One-hot encoded tensor containing labels.
- y_pred: tf.Tensor
- Tensor containing model predictions.
- """
-
- if threshold is not None:
- y_pred = tf.where(y_pred >= threshold, 1, 0)
-
- y_pred_neg = 1 - y_pred
- y_true_neg = 1 - y_true
-
- sum_over_axes = tf.range(tf.rank(y_pred) - 1) # Indices for axes to sum over. Excludes the final (class) dimension.
-
- true_positives = tf.math.reduce_sum(y_pred * y_true, axis=sum_over_axes)
- false_negatives = tf.math.reduce_sum(y_pred_neg * y_true, axis=sum_over_axes)
- false_positives = tf.math.reduce_sum(y_pred * y_true_neg, axis=sum_over_axes)
-
- if class_weights is not None:
- relative_class_weights = tf.cast(class_weights / tf.math.reduce_sum(class_weights), tf.float32)
- csi = tf.math.reduce_sum(tf.math.divide_no_nan(true_positives, true_positives + false_positives + false_negatives) * relative_class_weights)
- else:
- csi = tf.math.divide(tf.math.reduce_sum(true_positives), tf.math.reduce_sum(true_positives) + tf.math.reduce_sum(false_negatives) + tf.math.reduce_sum(false_positives))
-
- return 1 - csi
-
- return csi_loss
-
-
-def fractions_skill_score(
- num_dims: int,
- mask_size: int = 3,
- c: float = 1.0,
- cutoff: float = 0.5,
- want_hard_discretization: bool = False,
- class_weights: list[int | float] = None):
- """
- Fractions skill score loss function. Visit https://github.com/CIRA-ML/custom_loss_functions for documentation.
-
- Parameters
- ----------
- num_dims: int
- Number of dimensions for the mask.
- mask_size: int or tuple
- Size of the mask/pool in the AveragePooling layers.
- c: int or float
- C parameter in the sigmoid function. This will only be used if 'want_hard_discretization' is False.
- cutoff: float
- If 'want_hard_discretization' is True, y_true and y_pred will be discretized to only have binary values (0/1)
- want_hard_discretization: bool
- If True, y_true and y_pred will be discretized to only have binary values (0/1).
- If False, y_true and y_pred will be discretized using a sigmoid function.
- class_weights: list of values or None
- List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
-
- Returns
- -------
- fractions_skill_score: float
- Fractions skill score.
- """
-
- pool_kwargs = {'pool_size': (mask_size, ) * num_dims,
- 'strides': (1, ) * num_dims,
- 'padding': 'valid'}
-
- if num_dims == 2:
- pool1 = tf.keras.layers.AveragePooling2D(**pool_kwargs)
- pool2 = tf.keras.layers.AveragePooling2D(**pool_kwargs)
- else:
- pool1 = tf.keras.layers.AveragePooling3D(**pool_kwargs)
- pool2 = tf.keras.layers.AveragePooling3D(**pool_kwargs)
-
- @tf.function
- def fss_loss(y_true, y_pred):
- """
- y_true: tf.Tensor
- One-hot encoded tensor containing labels.
- y_pred: tf.Tensor
- Tensor containing model predictions.
- """
-
- if want_hard_discretization:
- y_true_binary = tf.where(y_true > cutoff, 1.0, 0.0)
- y_pred_binary = tf.where(y_pred > cutoff, 1.0, 0.0)
- else:
- y_true_binary = tf.math.sigmoid(c * (y_true - cutoff))
- y_pred_binary = tf.math.sigmoid(c * (y_pred - cutoff))
-
- y_true_density = pool1(y_true_binary)
- n_density_pixels = tf.cast((tf.shape(y_true_density)[1] * tf.shape(y_true_density)[2]), tf.float32)
-
- y_pred_density = pool2(y_pred_binary)
-
- if class_weights is None:
- MSE_n = tf.keras.metrics.mean_squared_error(y_true_density, y_pred_density)
- else:
- relative_class_weights = tf.cast(class_weights / tf.math.reduce_sum(class_weights), tf.float32)
- MSE_n = tf.reduce_mean(tf.math.square(y_true_density - y_pred_density) * relative_class_weights, axis=-1)
-
- O_n_squared_image = tf.keras.layers.Multiply()([y_true_density, y_true_density])
- O_n_squared_vector = tf.keras.layers.Flatten()(O_n_squared_image)
- O_n_squared_sum = tf.reduce_sum(O_n_squared_vector)
-
- M_n_squared_image = tf.keras.layers.Multiply()([y_pred_density, y_pred_density])
- M_n_squared_vector = tf.keras.layers.Flatten()(M_n_squared_image)
- M_n_squared_sum = tf.reduce_sum(M_n_squared_vector)
-
- MSE_n_ref = (O_n_squared_sum + M_n_squared_sum) / n_density_pixels
-
- my_epsilon = tf.keras.backend.epsilon() # this is 10^(-7)
-
- if want_hard_discretization:
- if MSE_n_ref == 0:
- return MSE_n
- else:
- return MSE_n / MSE_n_ref
- else:
- return MSE_n / (MSE_n_ref + my_epsilon)
-
- return fss_loss
diff --git a/custom_metrics.py b/custom_metrics.py
deleted file mode 100644
index b1bd91e..0000000
--- a/custom_metrics.py
+++ /dev/null
@@ -1,179 +0,0 @@
-"""
-Custom metrics for U-Net models.
- - Brier Skill Score (BSS)
- - Critical Success Index (CSI)
- - Fractions Skill Score (FSS)
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.9.21
-"""
-import tensorflow as tf
-
-
-def brier_skill_score(class_weights: list[int | float] = None):
- """
- Brier skill score (BSS).
-
- class_weights: list of values or None
- List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
- """
-
- @tf.function
- def bss(y_true, y_pred):
- """
- y_true: tf.Tensor
- One-hot encoded tensor containing labels.
- y_pred: tf.Tensor
- Tensor containing model predictions.
- """
-
- squared_errors = tf.math.square(tf.subtract(y_true, y_pred))
-
- if class_weights is not None:
- relative_class_weights = tf.cast(class_weights / tf.math.reduce_sum(class_weights), tf.float32)
- squared_errors *= relative_class_weights
-
- return 1 - tf.math.reduce_sum(squared_errors) / tf.size(squared_errors)
-
- return bss
-
-
-def critical_success_index(threshold: float = None, class_weights: list[int | float] = None):
- """
- Critical success index (CSI).
-
- y_true: tf.Tensor
- One-hot encoded tensor containing labels.
- y_pred: tf.Tensor
- Tensor containing model predictions.
- threshold: float or None
- Optional probability threshold that binarizes y_pred. Values in y_pred greater than or equal to the threshold are
- set to 1, and 0 otherwise.
- If the threshold is set, it must be greater than 0 and less than 1.
- class_weights: list of values or None
- List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
- """
-
- @tf.function
- def csi(y_true, y_pred):
- """
- y_true: tf.Tensor
- One-hot encoded tensor containing labels.
- y_pred: tf.Tensor
- Tensor containing model predictions.
- """
-
- if threshold is not None:
- y_pred = tf.where(y_pred >= threshold, 1.0, 0.0)
-
- y_pred_neg = 1 - y_pred
- y_true_neg = 1 - y_true
-
- sum_over_axes = tf.range(tf.rank(y_pred) - 1) # Indices for axes to sum over. Excludes the final (class) dimension.
-
- true_positives = tf.math.reduce_sum(y_pred * y_true, axis=sum_over_axes)
- false_negatives = tf.math.reduce_sum(y_pred_neg * y_true, axis=sum_over_axes)
- false_positives = tf.math.reduce_sum(y_pred * y_true_neg, axis=sum_over_axes)
-
- if class_weights is not None:
- relative_class_weights = tf.cast(class_weights / tf.math.reduce_sum(class_weights), tf.float32)
- csi = tf.math.reduce_sum(tf.math.divide_no_nan(true_positives, true_positives + false_positives + false_negatives) * relative_class_weights)
- else:
- csi = tf.math.divide(tf.math.reduce_sum(true_positives), tf.math.reduce_sum(true_positives) + tf.math.reduce_sum(false_negatives) + tf.math.reduce_sum(false_positives))
-
- return csi
-
- return csi
-
-
-def fractions_skill_score(
- num_dims: int,
- mask_size: int = 3,
- c: float = 1.0,
- cutoff: float = 0.5,
- want_hard_discretization: bool = False,
- class_weights: list[int | float] = None):
- """
- Fractions skill score loss function. Visit https://github.com/CIRA-ML/custom_loss_functions for documentation.
-
- Parameters
- ----------
- num_dims: int
- Number of dimensions for the mask.
- mask_size: int or tuple
- Size of the mask/pool in the AveragePooling layers.
- c: int or float
- C parameter in the sigmoid function. This will only be used if 'want_hard_discretization' is False.
- cutoff: float
- If 'want_hard_discretization' is True, y_true and y_pred will be discretized to only have binary values (0/1)
- want_hard_discretization: bool
- If True, y_true and y_pred will be discretized to only have binary values (0/1).
- If False, y_true and y_pred will be discretized using a sigmoid function.
- class_weights: list of values or None
- List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
-
- Returns
- -------
- fractions_skill_score: float
- Fractions skill score.
- """
-
- pool_kwargs = {'pool_size': (mask_size, ) * num_dims,
- 'strides': (1, ) * num_dims,
- 'padding': 'valid'}
-
- if num_dims == 2:
- pool1 = tf.keras.layers.AveragePooling2D(**pool_kwargs)
- pool2 = tf.keras.layers.AveragePooling2D(**pool_kwargs)
- else:
- pool1 = tf.keras.layers.AveragePooling3D(**pool_kwargs)
- pool2 = tf.keras.layers.AveragePooling3D(**pool_kwargs)
-
- @tf.function
- def fss(y_true, y_pred):
- """
- y_true: tf.Tensor
- One-hot encoded tensor containing labels.
- y_pred: tf.Tensor
- Tensor containing model predictions.
- """
-
- if want_hard_discretization:
- y_true_binary = tf.where(y_true > cutoff, 1.0, 0.0)
- y_pred_binary = tf.where(y_pred > cutoff, 1.0, 0.0)
- else:
- y_true_binary = tf.math.sigmoid(c * (y_true - cutoff))
- y_pred_binary = tf.math.sigmoid(c * (y_pred - cutoff))
-
- y_true_density = pool1(y_true_binary)
- n_density_pixels = tf.cast((tf.shape(y_true_density)[1] * tf.shape(y_true_density)[2]), tf.float32)
-
- y_pred_density = pool2(y_pred_binary)
-
- if class_weights is None:
- MSE_n = tf.keras.metrics.mean_squared_error(y_true_density, y_pred_density)
- else:
- relative_class_weights = tf.cast(class_weights / tf.math.reduce_sum(class_weights), tf.float32)
- MSE_n = tf.reduce_mean(tf.math.square(y_true_density - y_pred_density) * relative_class_weights, axis=-1)
-
- O_n_squared_image = tf.keras.layers.Multiply()([y_true_density, y_true_density])
- O_n_squared_vector = tf.keras.layers.Flatten()(O_n_squared_image)
- O_n_squared_sum = tf.reduce_sum(O_n_squared_vector)
-
- M_n_squared_image = tf.keras.layers.Multiply()([y_pred_density, y_pred_density])
- M_n_squared_vector = tf.keras.layers.Flatten()(M_n_squared_image)
- M_n_squared_sum = tf.reduce_sum(M_n_squared_vector)
-
- MSE_n_ref = (O_n_squared_sum + M_n_squared_sum) / n_density_pixels
-
- my_epsilon = tf.keras.backend.epsilon() # this is 10^(-7)
-
- if want_hard_discretization:
- if MSE_n_ref == 0:
- return 1 - MSE_n
- else:
- return 1 - (MSE_n / MSE_n_ref)
- else:
- return 1 - (MSE_n / (MSE_n_ref + my_epsilon))
-
- return fss
diff --git a/download_grib_files.py b/download_grib_files.py
deleted file mode 100644
index 15d4992..0000000
--- a/download_grib_files.py
+++ /dev/null
@@ -1,114 +0,0 @@
-"""
-Download grib files for GDAS and/or GFS data.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.8.23
-"""
-
-import argparse
-import os
-import pandas as pd
-import requests
-import urllib.error
-import wget
-import sys
-import datetime
-
-
-def bar_progress(current, total, width=None):
- progress_message = "Downloading %s: %d%% [%d/%d] MB " % (local_filename, current / total * 100, current / 1e6, total / 1e6)
- sys.stdout.write("\r" + progress_message)
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument('--grib_outdir', type=str, required=True, help='Output directory for GDAS grib files downloaded from NCEP.')
- parser.add_argument('--model', type=str, required=True, help="NWP model to use as the data source.")
- parser.add_argument('--init_time', type=str, help="Initialization time of the model. Format: YYYY-MM-DD-HH.")
- parser.add_argument('--range', type=str, nargs=3,
- help="Download model data between a range of dates. Three arguments must be passed, with the first two arguments "
- "marking the bounds of the date range in the format YYYY-MM-DD-HH. The third argument is the frequency (e.g. 6H), "
- "which has the same formatting as the 'freq' keyword argument in pandas.date_range().")
- parser.add_argument('--forecast_hours', type=int, nargs="+", required=True, help="List of forecast hours to download for the given day.")
- parser.add_argument('--verbose', action='store_true', help="Include a progress bar for download progress.")
-
- args = vars(parser.parse_args())
-
- args['model'] = args['model'].lower()
-
- # If --verbose is passed, include a progress bar to show the download progress
- bar = bar_progress if args['verbose'] else None
-
- if args['init_time'] is not None and args['range'] is not None:
- raise ValueError("Only one of the following arguments can be passed: --init_time, --range")
- elif args['init_time'] is None and args['range'] is None:
- raise ValueError("One of the following arguments must be passed: --init_time, --range")
-
- init_times = pd.date_range(args['init_time'], args['init_time']) if args['init_time'] is not None else pd.date_range(*args['range'][:2], freq=args['range'][-1])
-
- files = [] # complete urls for the files to pull from AWS
- local_filenames = [] # filenames for the local files after downloading
-
- for init_time in init_times:
- if args['model'] == 'gdas':
- if datetime.datetime(init_time.year, init_time.month, init_time.day, init_time.hour) < datetime.datetime(2015, 6, 23, 0):
- raise ConnectionAbortedError("Cannot download GDAS data prior to June 23, 2015.")
- elif datetime.datetime(init_time.year, init_time.month, init_time.day, init_time.hour) < datetime.datetime(2017, 7, 20, 0):
- [files.append(f'https://noaa-gfs-bdp-pds.s3.amazonaws.com/gdas.%d%02d%02d/%02d/gdas1.t%02dz.pgrb2.0p25.f%03d' % (init_time.year, init_time.month, init_time.day, init_time.hour, init_time.hour, forecast_hour))
- for forecast_hour in args['forecast_hours']]
- elif init_time.year < 2021:
- [files.append(f'https://noaa-gfs-bdp-pds.s3.amazonaws.com/gdas.%d%02d%02d/%02d/gdas.t%02dz.pgrb2.0p25.f%03d' % (init_time.year, init_time.month, init_time.day, init_time.hour, init_time.hour, forecast_hour))
- for forecast_hour in args['forecast_hours']]
- else:
- [files.append(f"https://noaa-gfs-bdp-pds.s3.amazonaws.com/gdas.%d%02d%02d/%02d/atmos/gdas.t%02dz.pgrb2.0p25.f%03d" % (init_time.year, init_time.month, init_time.day, init_time.hour, init_time.hour, forecast_hour))
- for forecast_hour in args['forecast_hours']]
- elif args['model'] == 'gfs':
- if datetime.datetime(init_time.year, init_time.month, init_time.day, init_time.hour) < datetime.datetime(2021, 2, 26, 0):
- raise ConnectionAbortedError("Cannot download GFS data prior to February 26, 2021.")
- elif datetime.datetime(init_time.year, init_time.month, init_time.day, init_time.hour) < datetime.datetime(2021, 3, 22, 0):
- [files.append(f"https://noaa-gfs-bdp-pds.s3.amazonaws.com/gfs.%d%02d%02d/%02d/gfs.t%02dz.pgrb2.0p25.f%03d" % (init_time.year, init_time.month, init_time.day, init_time.hour, init_time.hour, forecast_hour))
- for forecast_hour in args['forecast_hours']]
- else:
- [files.append(f"https://noaa-gfs-bdp-pds.s3.amazonaws.com/gfs.%d%02d%02d/%02d/atmos/gfs.t%02dz.pgrb2.0p25.f%03d" % (init_time.year, init_time.month, init_time.day, init_time.hour, init_time.hour, forecast_hour))
- for forecast_hour in args['forecast_hours']]
- elif args['model'] == 'hrrr':
- [files.append(f"https://noaa-hrrr-bdp-pds.s3.amazonaws.com/hrrr.%d%02d%02d/conus/hrrr.t%02dz.wrfprsf%02d.grib2" % (init_time.year, init_time.month, init_time.day, init_time.hour, forecast_hour))
- for forecast_hour in args['forecast_hours']]
- elif args['model'] == 'rap':
- [files.append(f"https://noaa-rap-pds.s3.amazonaws.com/rap.%d%02d%02d/rap.t%02dz.wrfprsf%02d.grib2" % (init_time.year, init_time.month, init_time.day, init_time.hour, forecast_hour))
- for forecast_hour in args['forecast_hours']]
- elif 'namnest' in args['model']:
- nest = args['model'].split('_')[-1]
- [files.append(f"https://nomads.ncep.noaa.gov/pub/data/nccf/com/nam/prod/nam.%d%02d%02d/nam.t%02dz.%snest.hiresf%02d.tm00.grib2" % (init_time.year, init_time.month, init_time.day, init_time.hour, nest, forecast_hour))
- for forecast_hour in args['forecast_hours']]
- elif args['model'] == 'nam_12km':
- for forecast_hour in args['forecast_hours']:
- if forecast_hour in [0, 1, 2, 3, 6]:
- folder = 'analysis' # use the analysis folder as it contains more accurate data
- else:
- folder = 'forecast' # forecast hours other than 0, 1, 2, 3, 6 do not have analysis data
- files.append(f"https://www.ncei.noaa.gov/data/north-american-mesoscale-model/access/%s/%d%02d/%d%02d%02d/nam_218_%d%02d%02d_%02d00_%03d.grb2" %
- (folder, init_time.year, init_time.month, init_time.year, init_time.month, init_time.day, init_time.year, init_time.month, init_time.day, init_time.hour, forecast_hour))
- [local_filenames.append("%s_%d%02d%02d%02d_f%03d.grib" % (args['model'], init_time.year, init_time.month, init_time.day, init_time.hour, forecast_hour)) for forecast_hour in args['forecast_hours']]
-
- for file, local_filename in zip(files, local_filenames):
-
- init_time = local_filename.split('_')[1] if 'nam' not in args['model'] else local_filename.split('_')[2]
- init_time = pd.to_datetime(f'{init_time[:4]}-{init_time[4:6]}-{init_time[6:8]}-{init_time[8:10]}')
-
- monthly_directory = '%s/%d%02d' % (args['grib_outdir'], init_time.year, init_time.month) # Directory for the grib files for the given days
-
- ### If the directory does not exist, check to see if the file link is valid. If the file link is NOT valid, then the directory will not be created since it will be empty. ###
- if not os.path.isdir(monthly_directory):
- if requests.head(file).status_code == requests.codes.ok or requests.head(file.replace('/atmos', '')).status_code == requests.codes.ok:
- os.mkdir(monthly_directory)
-
- full_file_path = f'{monthly_directory}/{local_filename}'
-
- if not os.path.isfile(full_file_path):
- try:
- wget.download(file, out=full_file_path, bar=bar)
- except urllib.error.HTTPError:
- print(f"Error downloading {file}")
- else:
- print(f"{full_file_path} already exists, skipping file....")
diff --git a/evaluation/calibrate_model.py b/evaluation/calibrate_model.py
deleted file mode 100644
index f47a6e9..0000000
--- a/evaluation/calibrate_model.py
+++ /dev/null
@@ -1,120 +0,0 @@
-"""
-Calibrate a trained model.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.6.24.D1
-"""
-import argparse
-import pandas as pd
-import os
-import sys
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) # this line allows us to import scripts outside the current directory
-from utils.settings import DEFAULT_FRONT_NAMES
-import matplotlib.pyplot as plt
-import pickle
-import xarray as xr
-import numpy as np
-from sklearn.isotonic import IsotonicRegression
-from sklearn.metrics import r2_score
-
-
-if __name__ == '__main__':
- """
- All arguments listed in the examples are listed via argparse in alphabetical order below this comment block.
- """
- parser = argparse.ArgumentParser()
- parser.add_argument('--dataset', type=str, help="Dataset for which to make predictions if prediction_method is 'random' or 'all'. Options are:"
- "'training', 'validation', 'test'")
- parser.add_argument('--domain', type=str, help='Domain of the data.')
- parser.add_argument('--model_dir', type=str, help='Directory for the models.')
- parser.add_argument('--model_number', type=int, help='Model number.')
- parser.add_argument('--data_source', type=str, default='era5', help='Data source for variables')
-
- args = vars(parser.parse_args())
-
- model_properties = pd.read_pickle('%s/model_%d/model_%d_properties.pkl' % (args['model_dir'], args['model_number'], args['model_number']))
-
- ### front_types argument is being moved into the dataset_properties dictionary within model_properties ###
- try:
- front_types = model_properties['front_types']
- except KeyError:
- front_types = model_properties['dataset_properties']['front_types']
-
- if type(front_types) == str:
- front_types = [front_types, ]
-
- try:
- _ = model_properties['calibration_models'] # Check to see if the model has already been calibrated before
- except KeyError:
- model_properties['calibration_models'] = dict()
-
- model_properties['calibration_models'][args['domain']] = dict()
-
- stats_ds = xr.open_dataset('%s/model_%d/statistics/model_%d_statistics_%s_%s.nc' % (args['model_dir'], args['model_number'], args['model_number'], args['domain'], args['dataset']))
-
- axis_ticks = np.arange(0.1, 1.1, 0.1)
-
- for front_label in front_types:
-
- model_properties['calibration_models'][args['domain']][front_label] = dict()
-
- true_positives = stats_ds[f'tp_temporal_{front_label}'].values
- false_positives = stats_ds[f'fp_temporal_{front_label}'].values
-
- thresholds = stats_ds['threshold'].values
-
- ### Sum the true positives along the 'time' axis ###
- true_positives_sum = np.sum(true_positives, axis=0)
- false_positives_sum = np.sum(false_positives, axis=0)
-
- ### Find the number of true positives and false positives in each probability bin ###
- true_positives_diff = np.abs(np.diff(true_positives_sum))
- false_positives_diff = np.abs(np.diff(false_positives_sum))
- observed_relative_frequency = np.divide(true_positives_diff, true_positives_diff + false_positives_diff)
-
- boundary_colors = ['red', 'purple', 'brown', 'darkorange', 'darkgreen']
-
- calibrated_probabilities = []
-
- fig, axs = plt.subplots(1, 2, figsize=(14, 6))
- axs[0].plot(thresholds, thresholds, color='black', linestyle='--', linewidth=0.5, label='Perfect Reliability')
-
- for boundary, color in enumerate(boundary_colors):
-
- ####################### Test different calibration methods to see which performs best ######################
-
- x = [threshold for threshold, frequency in zip(thresholds[1:], observed_relative_frequency[boundary]) if not np.isnan(frequency)]
- y = [frequency for threshold, frequency in zip(thresholds[1:], observed_relative_frequency[boundary]) if not np.isnan(frequency)]
-
- ### Isotonic Regression ###
- ir = IsotonicRegression(out_of_bounds='clip')
- ir.fit_transform(x, y)
- calibrated_probabilities.append(ir.predict(x))
- r_squared = r2_score(y, calibrated_probabilities[boundary])
-
- axs[0].plot(x, y, color=color, linewidth=1, label='%d km' % ((boundary + 1) * 50))
- axs[1].plot(x, calibrated_probabilities[boundary], color=color, linestyle='--', linewidth=1, label=r'%d km ($R^2$ = %.3f)' % ((boundary + 1) * 50, r_squared))
- model_properties['calibration_models'][args['domain']][front_label]['%d km' % ((boundary + 1) * 50)] = ir
-
- for ax in axs:
-
- axs[0].set_xlabel("Forecast Probability (uncalibrated)")
- ax.set_xticks(axis_ticks)
- ax.set_yticks(axis_ticks)
- ax.set_xlim(0, 1)
- ax.set_ylim(0, 1)
- ax.grid()
- ax.legend()
-
- axs[0].set_title('Reliability Diagram')
- axs[1].set_title('Calibration (isotonic regression)')
- axs[0].set_ylabel("Observed Relative Frequency")
- axs[1].set_ylabel("Forecast Probability (calibrated)")
-
- with open('%s/model_%d/model_%d_properties.pkl' % (args['model_dir'], args['model_number'], args['model_number']), 'wb') as f:
- pickle.dump(model_properties, f)
-
- plt.suptitle(f"Model {args['model_number']} reliability/calibration: {DEFAULT_FRONT_NAMES[front_label]}")
- plt.savefig(f'%s/model_%d/model_%d_calibration_%s_%s.png' % (args['model_dir'], args['model_number'], args['model_number'], args['domain'], front_label),
- bbox_inches='tight', dpi=300)
- plt.close()
diff --git a/evaluation/generate_performance_stats.py b/evaluation/generate_performance_stats.py
deleted file mode 100644
index 4b628fe..0000000
--- a/evaluation/generate_performance_stats.py
+++ /dev/null
@@ -1,272 +0,0 @@
-"""
-Generate performance statistics for a model.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.9.2.D1
-"""
-import argparse
-import glob
-import numpy as np
-import pandas as pd
-import random
-import tensorflow as tf
-import xarray as xr
-import os
-import sys
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) # this line allows us to import scripts outside the current directory
-import file_manager as fm
-from utils import data_utils
-from utils.settings import DEFAULT_DOMAIN_EXTENTS
-
-
-def combine_statistics_for_dataset():
-
- statistics_files = []
-
- for year in years:
- statistics_files += list(sorted(glob.glob('%s/model_%d/statistics/model_%d_statistics_%s_%d*.nc' %
- (args['model_dir'], args['model_number'], args['model_number'], args['domain'], year))))
-
- datasets_by_front_type = []
-
- for front_no, front_type in enumerate(front_types):
-
- ### Temporal and spatial datasets need to be loaded separately because of differing dimensions (xarray bugs) ###
- dataset_performance_ds_temporal = xr.open_dataset(statistics_files[0], chunks={'time': 16})[['%s_temporal_%s' % (stat, front_type) for stat in ['tp', 'fp', 'tn', 'fn']]]
- dataset_performance_ds_spatial = xr.open_dataset(statistics_files[0], chunks={'time': 16})[['%s_spatial_%s' % (stat, front_type) for stat in ['tp', 'fp', 'tn', 'fn']]]
- for stats_file in statistics_files[1:]:
- dataset_performance_ds_spatial += xr.open_dataset(stats_file, chunks={'time': 16})[['%s_spatial_%s' % (stat, front_type) for stat in ['tp', 'fp', 'tn', 'fn']]]
- dataset_performance_ds_temporal = xr.merge([dataset_performance_ds_temporal, xr.open_dataset(stats_file, chunks={'time': 16})[['%s_temporal_%s' % (stat, front_type) for stat in ['tp', 'fp', 'tn', 'fn']]]])
- dataset_performance_ds = xr.merge([dataset_performance_ds_spatial, dataset_performance_ds_temporal]) # Combine spatial and temporal data into one dataset
-
- tp_array_temporal = dataset_performance_ds['tp_temporal_%s' % front_type].values
- fp_array_temporal = dataset_performance_ds['fp_temporal_%s' % front_type].values
- fn_array_temporal = dataset_performance_ds['fn_temporal_%s' % front_type].values
-
- time_array = dataset_performance_ds['time'].values
-
- ### Bootstrap the temporal statistics to find confidence intervals ###
- POD_array = np.zeros([num_front_types, args['num_iterations'], 5, 100]) # probability of detection = TP / (TP + FN)
- SR_array = np.zeros([num_front_types, args['num_iterations'], 5, 100]) # success ratio = 1 - False Alarm Ratio = TP / (TP + FP)
-
- # 3 confidence intervals: 90, 95, and 99%
- CI_lower_POD = np.zeros([num_front_types, 3, 5, 100])
- CI_lower_SR = np.zeros([num_front_types, 3, 5, 100])
- CI_upper_POD = np.zeros([num_front_types, 3, 5, 100])
- CI_upper_SR = np.zeros([num_front_types, 3, 5, 100])
-
- num_timesteps = len(time_array)
- selectable_indices = range(num_timesteps)
-
- for iteration in range(args['num_iterations']):
- print(f"Iteration {iteration}/{args['num_iterations']}", end='\r')
- indices = random.choices(selectable_indices, k=num_timesteps) # Select a sample equal to the total number of timesteps
-
- POD_array[front_no, iteration, :, :] = np.divide(np.sum(tp_array_temporal[indices, :, :], axis=0),
- np.add(np.sum(tp_array_temporal[indices, :, :], axis=0),
- np.sum(fn_array_temporal[indices, :, :], axis=0)))
- SR_array[front_no, iteration, :, :] = np.divide(np.sum(tp_array_temporal[indices, :, :], axis=0),
- np.add(np.sum(tp_array_temporal[indices, :, :], axis=0),
- np.sum(fp_array_temporal[indices, :, :], axis=0)))
- print(f"Iteration {args['num_iterations']}/{args['num_iterations']}")
-
- ## Turn NaNs to zeros
- POD_array = np.nan_to_num(POD_array)
- SR_array = np.nan_to_num(SR_array)
-
- # Calculate confidence intervals at each probability bin
- for percent in np.arange(0, 100):
- CI_lower_POD[front_no, 0, :, percent] = np.percentile(POD_array[front_no, :, :, percent], q=5, axis=0) # lower bound for 90% confidence interval
- CI_lower_POD[front_no, 1, :, percent] = np.percentile(POD_array[front_no, :, :, percent], q=2.5, axis=0) # lower bound for 95% confidence interval
- CI_lower_POD[front_no, 2, :, percent] = np.percentile(POD_array[front_no, :, :, percent], q=0.5, axis=0) # lower bound for 99% confidence interval
- CI_upper_POD[front_no, 0, :, percent] = np.percentile(POD_array[front_no, :, :, percent], q=95, axis=0) # upper bound for 90% confidence interval
- CI_upper_POD[front_no, 1, :, percent] = np.percentile(POD_array[front_no, :, :, percent], q=97.5, axis=0) # upper bound for 95% confidence interval
- CI_upper_POD[front_no, 2, :, percent] = np.percentile(POD_array[front_no, :, :, percent], q=99.5, axis=0) # upper bound for 99% confidence interval
-
- CI_lower_SR[front_no, 0, :, percent] = np.percentile(SR_array[front_no, :, :, percent], q=5, axis=0) # lower bound for 90% confidence interval
- CI_lower_SR[front_no, 1, :, percent] = np.percentile(SR_array[front_no, :, :, percent], q=2.5, axis=0) # lower bound for 95% confidence interval
- CI_lower_SR[front_no, 2, :, percent] = np.percentile(SR_array[front_no, :, :, percent], q=0.5, axis=0) # lower bound for 99% confidence interval
- CI_upper_SR[front_no, 0, :, percent] = np.percentile(SR_array[front_no, :, :, percent], q=95, axis=0) # upper bound for 90% confidence interval
- CI_upper_SR[front_no, 1, :, percent] = np.percentile(SR_array[front_no, :, :, percent], q=97.5, axis=0) # upper bound for 95% confidence interval
- CI_upper_SR[front_no, 2, :, percent] = np.percentile(SR_array[front_no, :, :, percent], q=99.5, axis=0) # upper bound for 99% confidence interval
-
- dataset_performance_ds["POD_0.5_%s" % front_type] = (('boundary', 'threshold'), CI_lower_POD[front_no, 2, :, :])
- dataset_performance_ds["POD_2.5_%s" % front_type] = (('boundary', 'threshold'), CI_lower_POD[front_no, 1, :, :])
- dataset_performance_ds["POD_5_%s" % front_type] = (('boundary', 'threshold'), CI_lower_POD[front_no, 0, :, :])
- dataset_performance_ds["POD_99.5_%s" % front_type] = (('boundary', 'threshold'), CI_upper_POD[front_no, 2, :, :])
- dataset_performance_ds["POD_97.5_%s" % front_type] = (('boundary', 'threshold'), CI_upper_POD[front_no, 1, :, :])
- dataset_performance_ds["POD_95_%s" % front_type] = (('boundary', 'threshold'), CI_upper_POD[front_no, 0, :, :])
- dataset_performance_ds["SR_0.5_%s" % front_type] = (('boundary', 'threshold'), CI_lower_SR[front_no, 2, :, :])
- dataset_performance_ds["SR_2.5_%s" % front_type] = (('boundary', 'threshold'), CI_lower_SR[front_no, 1, :, :])
- dataset_performance_ds["SR_5_%s" % front_type] = (('boundary', 'threshold'), CI_lower_SR[front_no, 0, :, :])
- dataset_performance_ds["SR_99.5_%s" % front_type] = (('boundary', 'threshold'), CI_upper_SR[front_no, 2, :, :])
- dataset_performance_ds["SR_97.5_%s" % front_type] = (('boundary', 'threshold'), CI_upper_SR[front_no, 1, :, :])
- dataset_performance_ds["SR_95_%s" % front_type] = (('boundary', 'threshold'), CI_upper_SR[front_no, 0, :, :])
-
- datasets_by_front_type.append(dataset_performance_ds)
-
- final_performance_ds = xr.merge(datasets_by_front_type)
- final_performance_ds.to_netcdf(path='%s/model_%d/statistics/model_%d_statistics_%s_%s.nc' % (args['model_dir'], args['model_number'], args['model_number'], args['domain'], args['dataset']), mode='w', engine='netcdf4')
-
-
-if __name__ == '__main__':
- """
- All arguments listed in the examples are listed via argparse in alphabetical order below this comment block.
- """
- parser = argparse.ArgumentParser()
- parser.add_argument('--dataset', type=str, help="Dataset for which to make predictions. Options are: 'training', 'validation', 'test'")
- parser.add_argument('--year_and_month', type=int, nargs=2, help="Year and month for which to make predictions.")
- parser.add_argument('--combine', action='store_true', help="Combine calculated statistics for a dataset.")
- parser.add_argument('--domain', type=str, help='Domain of the data.')
- parser.add_argument('--forecast_hour', type=int, help='Forecast hour for the GDAS data')
- parser.add_argument('--gpu_device', type=int, nargs='+', help='GPU device number.')
- parser.add_argument('--memory_growth', action='store_true', help='Use memory growth on the GPU')
- parser.add_argument('--model_dir', type=str, required=True, help='Directory for the models.')
- parser.add_argument('--model_number', type=int, required=True, help='Model number.')
- parser.add_argument('--num_iterations', type=int, default=10000, help='Number of iterations to perform when bootstrapping the data.')
- parser.add_argument('--fronts_netcdf_indir', type=str, help='Main directory for the netcdf files containing frontal objects.')
- parser.add_argument('--data_source', type=str, default='era5', help='Data source for variables')
- parser.add_argument('--overwrite', action='store_true', help="Overwrite any existing statistics files.")
-
- args = vars(parser.parse_args())
-
- model_properties = pd.read_pickle('%s/model_%d/model_%d_properties.pkl' % (args['model_dir'], args['model_number'], args['model_number']))
- domain = args['domain']
-
- # Some older models do not have the 'dataset_properties' dictionary
- try:
- front_types = model_properties['dataset_properties']['front_types']
- num_dims = model_properties['dataset_properties']['num_dims']
- except KeyError:
- front_types = model_properties['front_types']
- if args['model_number'] in [6846496, 7236500, 7507525]:
- num_dims = (3, 3)
-
- num_front_types = model_properties['classes'] - 1
-
- if args['dataset'] is not None and args['year_and_month'] is not None:
- raise ValueError("--dataset and --year_and_month cannot be passed together.")
- elif args['dataset'] is None and args['year_and_month'] is None:
- raise ValueError("At least one of [--dataset, --year_and_month] must be passed.")
- elif args['year_and_month'] is not None:
- years, months = [args['year_and_month'][0]], [args['year_and_month'][1]]
- else:
- years, months = model_properties['%s_years' % args['dataset']], range(1, 13)
-
- if args['dataset'] is not None and args['combine']:
- combine_statistics_for_dataset()
- exit()
-
- if args['gpu_device'] is not None:
- gpus = tf.config.list_physical_devices(device_type='GPU')
- tf.config.set_visible_devices(devices=[gpus[gpu] for gpu in args['gpu_device']], device_type='GPU')
-
- # Allow for memory growth on the GPU. This will only use the GPU memory that is required rather than allocating all the GPU's memory.
- if args['memory_growth']:
- tf.config.experimental.set_memory_growth(device=[gpus[gpu] for gpu in args['gpu_device']][0], enable=True)
-
- for year in years:
-
- era5_files_obj = fm.DataFileLoader(args['fronts_netcdf_indir'], data_file_type='fronts-netcdf')
- era5_files_obj.test_years = [year, ] # does not matter which year attribute we set the years to
- front_files = era5_files_obj.front_files_test
-
- for month in months:
-
- front_files_month = [file for file in front_files if '_%d%02d' % (year, month) in file]
-
- if args['domain'] == 'full':
- print("full")
- for front_file in front_files_month:
- if any(['%02d_full.nc' % hour in front_file for hour in np.arange(3, 27, 6)]):
- front_files_month.pop(front_files_month.index(front_file))
-
- prediction_file = f'%s/model_%d/probabilities/model_%d_pred_%s_%d%02d.nc' % \
- (args['model_dir'], args['model_number'], args['model_number'], args['domain'], year, month)
-
- stats_dataset_path = '%s/model_%d/statistics/model_%d_statistics_%s_%d%02d.nc' % (args['model_dir'], args['model_number'], args['model_number'], args['domain'], year, month)
- if os.path.isfile(stats_dataset_path) and not args['overwrite']:
- print("WARNING: %s exists, pass the --overwrite argument to overwrite existing data." % stats_dataset_path)
- continue
-
- probs_ds = xr.open_dataset(prediction_file)
- lons = probs_ds['longitude'].values
- lats = probs_ds['latitude'].values
-
- fronts_ds = xr.open_mfdataset(front_files_month, combine='nested', concat_dim='time')\
- .sel(longitude=slice(DEFAULT_DOMAIN_EXTENTS[args['domain']][0], DEFAULT_DOMAIN_EXTENTS[args['domain']][1]),
- latitude=slice(DEFAULT_DOMAIN_EXTENTS[args['domain']][3], DEFAULT_DOMAIN_EXTENTS[args['domain']][2]))
-
- fronts_ds_month = data_utils.reformat_fronts(fronts_ds.sel(time='%d-%02d' % (year, month)), front_types)
-
- time_array = probs_ds['time'].values
- num_timesteps = len(time_array)
-
- tp_array_spatial = np.zeros(shape=[num_front_types, len(lats), len(lons), 5, 100]).astype('int64')
- fp_array_spatial = np.zeros(shape=[num_front_types, len(lats), len(lons), 5, 100]).astype('int64')
- tn_array_spatial = np.zeros(shape=[num_front_types, len(lats), len(lons), 5, 100]).astype('int64')
- fn_array_spatial = np.zeros(shape=[num_front_types, len(lats), len(lons), 5, 100]).astype('int64')
-
- tp_array_temporal = np.zeros(shape=[num_front_types, num_timesteps, 5, 100]).astype('int64')
- fp_array_temporal = np.zeros(shape=[num_front_types, num_timesteps, 5, 100]).astype('int64')
- tn_array_temporal = np.zeros(shape=[num_front_types, num_timesteps, 5, 100]).astype('int64')
- fn_array_temporal = np.zeros(shape=[num_front_types, num_timesteps, 5, 100]).astype('int64')
-
- thresholds = np.linspace(0.01, 1, 100) # Probability thresholds for calculating performance statistics
- boundaries = np.array([50, 100, 150, 200, 250]) # Boundaries for checking whether a front is present (kilometers)
-
- bool_tn_fn_dss = dict({front: tf.convert_to_tensor(xr.where(fronts_ds_month == front_no + 1, 1, 0)['identifier'].values) for front_no, front in enumerate(front_types)})
- bool_tp_fp_dss = dict({front: None for front in front_types})
- probs_dss = dict({front: tf.convert_to_tensor(probs_ds[front].values) for front in front_types})
-
- performance_ds = xr.Dataset(coords={'time': time_array, 'longitude': lons, 'latitude': lats, 'boundary': boundaries, 'threshold': thresholds})
-
- for front_no, front_type in enumerate(front_types):
- fronts_ds_month = data_utils.reformat_fronts(fronts_ds.sel(time='%d-%02d' % (year, month)), front_types)
- print("%d-%02d: %s (TN/FN)" % (year, month, front_type))
- ### Calculate true/false negatives ###
- for i in range(100):
- """
- True negative ==> model correctly predicts the lack of a front at a given point
- False negative ==> model does not predict a front, but a front exists
-
- The numbers of true negatives and false negatives are the same for all neighborhoods and are calculated WITHOUT expanding the fronts.
- If we were to calculate the negatives separately for each neighborhood, the number of misses would be artificially inflated, lowering the
- final CSI scores and making the neighborhood method effectively useless.
- """
- tn = tf.where((probs_dss[front_type] < thresholds[i]) & (bool_tn_fn_dss[front_type] == 0), 1, 0)
- fn = tf.where((probs_dss[front_type] < thresholds[i]) & (bool_tn_fn_dss[front_type] == 1), 1, 0)
-
- tn_array_spatial[front_no, :, :, :, i] = tf.tile(tf.expand_dims(tf.reduce_sum(tn, axis=0), axis=-1), (1, 1, 5))
- fn_array_spatial[front_no, :, :, :, i] = tf.tile(tf.expand_dims(tf.reduce_sum(fn, axis=0), axis=-1), (1, 1, 5))
- tn_array_temporal[front_no, :, :, i] = tf.tile(tf.expand_dims(tf.reduce_sum(tn, axis=(1, 2)), axis=-1), (1, 5))
- fn_array_temporal[front_no, :, :, i] = tf.tile(tf.expand_dims(tf.reduce_sum(fn, axis=(1, 2)), axis=-1), (1, 5))
-
- ### Calculate true/false positives ###
- for boundary in range(5):
- fronts_ds_month = data_utils.expand_fronts(fronts_ds_month, iterations=2) # Expand fronts
- bool_tp_fp_dss[front_type] = tf.convert_to_tensor(xr.where(fronts_ds_month == front_no + 1, 1, 0)['identifier'].values) # 1 = cold front, 0 = not a cold front
- print("%d-%02d: %s (%d km)" % (year, month, front_type, (boundary + 1) * 50))
- for i in range(100):
- """
- True positive ==> model correctly identifies a front
- False positive ==> model predicts a front, but no front is present within the given neighborhood
- """
- tp = tf.where((probs_dss[front_type] > thresholds[i]) & (bool_tp_fp_dss[front_type] == 1), 1, 0)
- fp = tf.where((probs_dss[front_type] > thresholds[i]) & (bool_tp_fp_dss[front_type] == 0), 1, 0)
-
- tp_array_spatial[front_no, :, :, boundary, i] = tf.reduce_sum(tp, axis=0)
- fp_array_spatial[front_no, :, :, boundary, i] = tf.reduce_sum(fp, axis=0)
- tp_array_temporal[front_no, :, boundary, i] = tf.reduce_sum(tp, axis=(1, 2))
- fp_array_temporal[front_no, :, boundary, i] = tf.reduce_sum(fp, axis=(1, 2))
-
- performance_ds["tp_spatial_%s" % front_type] = (('latitude', 'longitude', 'boundary', 'threshold'), tp_array_spatial[front_no])
- performance_ds["fp_spatial_%s" % front_type] = (('latitude', 'longitude', 'boundary', 'threshold'), fp_array_spatial[front_no])
- performance_ds["tn_spatial_%s" % front_type] = (('latitude', 'longitude', 'boundary', 'threshold'), tn_array_spatial[front_no])
- performance_ds["fn_spatial_%s" % front_type] = (('latitude', 'longitude', 'boundary', 'threshold'), fn_array_spatial[front_no])
- performance_ds["tp_temporal_%s" % front_type] = (('time', 'boundary', 'threshold'), tp_array_temporal[front_no])
- performance_ds["fp_temporal_%s" % front_type] = (('time', 'boundary', 'threshold'), fp_array_temporal[front_no])
- performance_ds["tn_temporal_%s" % front_type] = (('time', 'boundary', 'threshold'), tn_array_temporal[front_no])
- performance_ds["fn_temporal_%s" % front_type] = (('time', 'boundary', 'threshold'), fn_array_temporal[front_no])
-
- performance_ds.to_netcdf(path=stats_dataset_path, mode='w', engine='netcdf4')
diff --git a/evaluation/learning_curve.py b/evaluation/learning_curve.py
deleted file mode 100644
index e46c76a..0000000
--- a/evaluation/learning_curve.py
+++ /dev/null
@@ -1,96 +0,0 @@
-"""
-Plot the learning curve for a model.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.6.12
-"""
-import argparse
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-
-
-if __name__ == '__main__':
- """
- All arguments listed in the examples are listed via argparse in alphabetical order below this comment block.
- """
- parser = argparse.ArgumentParser()
- parser.add_argument('--model_dir', type=str, required=True, help='Directory for the models.')
- parser.add_argument('--model_number', type=int, required=True, help='Model number.')
-
- args = vars(parser.parse_args())
-
- with open("%s/model_%d/model_%d_history.csv" % (args['model_dir'], args['model_number'], args['model_number']), 'rb') as f:
- history = pd.read_csv(f)
-
- model_properties = pd.read_pickle(f"{args['model_dir']}/model_{args['model_number']}/model_{args['model_number']}_properties.pkl")
-
- # Model properties
- try:
- loss = model_properties['loss']
- except KeyError:
- loss = model_properties['loss_string']
-
- try:
- metric_string = model_properties['metric']
- except KeyError:
- metric_string = model_properties['metric_string']
-
- if model_properties['deep_supervision']:
- train_metric = history['sup1_Softmax_%s' % metric_string]
- val_metric = history['val_sup1_Softmax_%s' % metric_string]
- else:
- train_metric = history[metric_string]
- val_metric = history['val_%s' % metric_string]
-
- if 'fss' in loss.lower():
- loss_title = 'Fractions Skill Score (loss)'
- elif 'bss' in loss.lower():
- loss_title = 'Brier Skill Score (loss)'
- elif 'csi' in loss.lower():
- loss_title = 'Categorical Cross-Entropy'
- else:
- loss_title = None
-
- if 'fss' in metric_string:
- metric_title = 'Fractions Skill Score'
- elif 'bss' in metric_string:
- metric_title = 'Brier Skill Score'
- elif 'csi' in metric_string:
- metric_title = 'Critical Success Index'
- else:
- metric_title = None
-
- min_val_loss_epoch = np.where(history['val_loss'] == np.min(history['val_loss']))[0][0] + 1
-
- num_epochs = len(history['val_loss'])
-
- fig, axs = plt.subplots(1, 2, figsize=(12, 6), dpi=300)
- axarr = axs.flatten()
-
- annotate_kwargs = dict(color='black', va='center', xycoords='axes fraction', fontsize=11)
- axarr[0].annotate('Epoch %d' % min_val_loss_epoch, xy=(0, -0.2), fontweight='bold', **annotate_kwargs)
- axarr[0].annotate('Training/Validation loss: %.4e, %.4e' % (history['loss'][min_val_loss_epoch - 1], history['val_loss'][min_val_loss_epoch - 1]), xy=(0, -0.25), **annotate_kwargs)
- axarr[0].annotate('Training/Validation metric: %.4f, %.4f' % (train_metric[min_val_loss_epoch - 1], val_metric[min_val_loss_epoch - 1]), xy=(0, -0.3), **annotate_kwargs)
-
- axarr[0].set_title(loss_title)
- axarr[0].plot(np.arange(1, num_epochs + 1), history['loss'], color='blue', label='Training loss')
- axarr[0].plot(np.arange(1, num_epochs + 1), history['val_loss'], color='red', label='Validation loss')
- axarr[0].set_xlim(xmin=0, xmax=num_epochs + 1)
- axarr[0].set_xlabel('Epochs')
- axarr[0].legend(loc='best')
- axarr[0].grid()
- axarr[0].set_yscale('log') # Turns y-axis into a logarithmic scale. Useful if loss functions appear as very sharp curves.
-
- axarr[1].set_title(metric_title)
- axarr[1].plot(np.arange(1, num_epochs + 1), train_metric, color='blue', label='Training')
- axarr[1].plot(np.arange(1, num_epochs + 1), val_metric, color='red', label='Validation')
- axarr[1].set_xlim(xmin=0, xmax=num_epochs + 1)
- axarr[1].set_ylim(ymin=0)
- axarr[1].set_xlabel('Epochs')
- axarr[1].legend(loc='best')
- axarr[1].grid()
-
- plt.tight_layout()
- plt.savefig("%s/model_%d/model_%d_learning_curve.png" % (args['model_dir'], args['model_number'], args['model_number']), bbox_inches='tight')
- plt.close()
diff --git a/evaluation/performance_diagrams.py b/evaluation/performance_diagrams.py
deleted file mode 100644
index c78c898..0000000
--- a/evaluation/performance_diagrams.py
+++ /dev/null
@@ -1,258 +0,0 @@
-"""
-Plot performance diagrams for a model.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.9.18
-"""
-import argparse
-import cartopy.crs as ccrs
-from matplotlib import colors
-from matplotlib.font_manager import FontProperties
-import matplotlib.pyplot as plt
-from matplotlib.ticker import FixedLocator
-import numpy as np
-import pandas as pd
-import pickle
-import xarray as xr
-import os
-import sys
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) # this line allows us to import scripts outside the current directory
-from utils import settings, plotting_utils
-
-
-if __name__ == '__main__':
- """
- All arguments listed in the examples are listed via argparse in alphabetical order below this comment block.
- """
- parser = argparse.ArgumentParser()
- parser.add_argument('--confidence_level', type=int, default=95, help="Confidence interval. Options are: 90, 95, 99.")
- parser.add_argument('--dataset', type=str, help="'training', 'validation', or 'test'")
- parser.add_argument('--data_source', type=str, default='era5', help="Source of the variable data (ERA5, GDAS, etc.)")
- parser.add_argument('--domain_images', type=int, nargs=2, help='Number of images for each dimension the final stitched map for predictions: lon, lat')
- parser.add_argument('--domain', type=str, required=True, help='Domain of the data.')
- parser.add_argument('--forecast_hour', type=int, help='Forecast hour for the GDAS or GFS data.')
- parser.add_argument('--map_neighborhood', type=int, default=250,
- help="Neighborhood for the CSI map in kilometers. Options are: 50, 100, 150, 200, 250")
- parser.add_argument('--model_dir', type=str, required=True, help='Directory for the models.')
- parser.add_argument('--model_number', type=int, required=True, help='Model number.')
-
- args = vars(parser.parse_args())
-
- model_properties_filepath = f"{args['model_dir']}/model_{args['model_number']}/model_{args['model_number']}_properties.pkl"
- model_properties = pd.read_pickle(model_properties_filepath)
-
- # Some older models do not have the 'dataset_properties' dictionary
- try:
- front_types = model_properties['dataset_properties']['front_types']
- except KeyError:
- front_types = model_properties['front_types']
-
- domain_extent_indices = settings.DEFAULT_DOMAIN_INDICES[args['domain']]
-
- stats_ds = xr.open_dataset('%s/model_%d/statistics/model_%d_statistics_%s_%s.nc' % (args['model_dir'], args['model_number'], args['model_number'], args['domain'], args['dataset']))
-
- if type(front_types) == str:
- front_types = [front_types, ]
-
- # Probability threshold where CSI is maximized for each front type and domain
- max_csi_thresholds = dict()
-
- if args['domain'] not in list(max_csi_thresholds.keys()):
- max_csi_thresholds[args['domain']] = dict()
-
- for front_no, front_label in enumerate(front_types):
-
- if front_label not in list(max_csi_thresholds[args['domain']].keys()):
- max_csi_thresholds[args['domain']][front_label] = dict()
-
- ################################ CSI and reliability diagrams (panels a and b) #################################
- true_positives_temporal = stats_ds[f'tp_temporal_{front_label}'].values
- false_positives_temporal = stats_ds[f'fp_temporal_{front_label}'].values
- false_negatives_temporal = stats_ds[f'fn_temporal_{front_label}'].values
- spatial_csi_ds = (stats_ds[f'tp_spatial_{front_label}'] / (stats_ds[f'tp_spatial_{front_label}'] + stats_ds[f'fp_spatial_{front_label}'] + stats_ds[f'fn_spatial_{front_label}'])).max('threshold')
- thresholds = stats_ds['threshold'].values
-
- if args['confidence_level'] != 90:
- CI_low, CI_high = (100 - args['confidence_level']) / 2, 50 + (args['confidence_level'] / 2)
- CI_low, CI_high = '%.1f' % CI_low, '%.1f' % CI_high
- else:
- CI_low, CI_high = 5, 95
-
- # Confidence intervals for POD and SR
- CI_POD = np.stack((stats_ds[f"POD_{CI_low}_{front_label}"].values, stats_ds[f"POD_{CI_high}_{front_label}"].values), axis=0)
- CI_SR = np.stack((stats_ds[f"SR_{CI_low}_{front_label}"].values, stats_ds[f"SR_{CI_high}_{front_label}"].values), axis=0)
- CI_CSI = np.stack((CI_SR ** -1 + CI_POD ** -1 - 1.) ** -1, axis=0)
- CI_FB = np.stack(CI_POD * (CI_SR ** -1), axis=0)
-
- # Remove the zeros
- try:
- polygon_stop_index = np.min(np.where(CI_POD == 0)[2])
- except IndexError:
- polygon_stop_index = 100
-
- ### Statistics with shape (boundary, threshold) after taking the sum along the time axis (axis=0) ###
- true_positives_temporal_sum = np.sum(true_positives_temporal, axis=0)
- false_positives_temporal_sum = np.sum(false_positives_temporal, axis=0)
- false_negatives_temporal_sum = np.sum(false_negatives_temporal, axis=0)
-
- ### Find the number of true positives and false positives in each probability bin ###
- true_positives_diff = np.abs(np.diff(true_positives_temporal_sum))
- false_positives_diff = np.abs(np.diff(false_positives_temporal_sum))
- observed_relative_frequency = np.divide(true_positives_diff, true_positives_diff + false_positives_diff)
-
- pod = np.divide(true_positives_temporal_sum, true_positives_temporal_sum + false_negatives_temporal_sum) # Probability of detection
- sr = np.divide(true_positives_temporal_sum, true_positives_temporal_sum + false_positives_temporal_sum) # Success ratio
-
- fig, axs = plt.subplots(1, 2, figsize=(15, 6))
- axarr = axs.flatten()
-
- sr_matrix, pod_matrix = np.meshgrid(np.linspace(0, 1, 101), np.linspace(0, 1, 101))
- csi_matrix = 1 / ((1/sr_matrix) + (1/pod_matrix) - 1) # CSI coordinates
- fb_matrix = pod_matrix * (sr_matrix ** -1) # Frequency Bias coordinates
- CSI_LEVELS = np.linspace(0, 1, 11) # CSI contour levels
- FB_LEVELS = [0.25, 0.5, 0.75, 1, 1.25, 1.5, 2, 3] # Frequency Bias levels
- cmap = 'Blues' # Colormap for the CSI contours
- axis_ticks = np.arange(0, 1.01, 0.1)
- axis_ticklabels = np.arange(0, 100.1, 10).astype(int)
-
- cs = axarr[0].contour(sr_matrix, pod_matrix, fb_matrix, FB_LEVELS, colors='black', linewidths=0.5, linestyles='--') # Plot FB levels
- axarr[0].clabel(cs, FB_LEVELS, fontsize=8)
-
- csi_contour = axarr[0].contourf(sr_matrix, pod_matrix, csi_matrix, CSI_LEVELS, cmap=cmap) # Plot CSI contours in 0.1 increments
- cbar = fig.colorbar(csi_contour, ax=axarr[0], pad=0.02, label='Critical Success Index (CSI)')
- cbar.set_ticks(axis_ticks)
-
- axarr[1].plot(thresholds, thresholds, color='black', linestyle='--', linewidth=0.5, label='Perfect Reliability')
-
- cell_text = [] # List of strings that will be used in the table near the bottom of this function
-
- ### CSI and reliability lines for each boundary ###
- boundary_colors = ['red', 'purple', 'brown', 'darkorange', 'darkgreen']
- max_CSI_scores_by_boundary = np.zeros(shape=(5,))
- for boundary, color in enumerate(boundary_colors):
- csi = np.power((1/sr[boundary]) + (1/pod[boundary]) - 1, -1)
- max_CSI_scores_by_boundary[boundary] = np.nanmax(csi)
- max_CSI_index = np.where(csi == max_CSI_scores_by_boundary[boundary])[0]
- max_CSI_threshold = thresholds[max_CSI_index][0] # Probability threshold where CSI is maximized
- max_csi_thresholds[args['domain']][front_label]['%s' % int((boundary + 1) * 50)] = np.round(max_CSI_threshold, 2)
- max_CSI_pod = pod[boundary][max_CSI_index][0] # POD where CSI is maximized
- max_CSI_sr = sr[boundary][max_CSI_index][0] # SR where CSI is maximized
- max_CSI_fb = max_CSI_pod / max_CSI_sr # Frequency bias
-
- cell_text.append([r'$\bf{%.2f}$' % max_CSI_threshold,
- r'$\bf{%.3f}$' % max_CSI_scores_by_boundary[boundary] + r'$^{%.3f}_{%.3f}$' % (CI_CSI[1, boundary, max_CSI_index][0], CI_CSI[0, boundary, max_CSI_index][0]),
- r'$\bf{%.1f}$' % (max_CSI_pod * 100) + r'$^{%.1f}_{%.1f}$' % (CI_POD[1, boundary, max_CSI_index][0] * 100, CI_POD[0, boundary, max_CSI_index][0] * 100),
- r'$\bf{%.1f}$' % (max_CSI_sr * 100) + r'$^{%.1f}_{%.1f}$' % (CI_SR[1, boundary, max_CSI_index][0] * 100, CI_SR[0, boundary, max_CSI_index][0] * 100),
- r'$\bf{%.1f}$' % ((1 - max_CSI_sr) * 100) + r'$^{%.1f}_{%.1f}$' % ((1 - CI_SR[1, boundary, max_CSI_index][0]) * 100, (1 - CI_SR[0, boundary, max_CSI_index][0]) * 100),
- r'$\bf{%.3f}$' % max_CSI_fb + r'$^{%.3f}_{%.3f}$' % (CI_FB[1, boundary, max_CSI_index][0], CI_FB[0, boundary, max_CSI_index][0])])
-
- # Plot CSI lines
- axarr[0].plot(max_CSI_sr, max_CSI_pod, color=color, marker='*', markersize=10)
- axarr[0].plot(sr[boundary], pod[boundary], color=color, linewidth=1)
-
- # Plot reliability curve
- axarr[1].plot(thresholds[1:], observed_relative_frequency[boundary], color=color, linewidth=1)
-
- # Confidence interval
- xs = np.concatenate([CI_SR[0, boundary, :polygon_stop_index], CI_SR[1, boundary, :polygon_stop_index][::-1]])
- ys = np.concatenate([CI_POD[0, boundary, :polygon_stop_index], CI_POD[1, boundary, :polygon_stop_index][::-1]])
- axarr[0].fill(xs, ys, alpha=0.3, color=color) # Shade the confidence interval
-
- axarr[0].set_xticklabels(axis_ticklabels[::-1]) # False alarm rate on x-axis means values are reversed
- axarr[0].set_xlabel("False Alarm Rate (FAR; %)")
- axarr[0].set_ylabel("Probability of Detection (POD; %)")
- axarr[0].set_title(r'$\bf{a)}$ $\bf{CSI}$ $\bf{diagram}$ [confidence level = %d%%]' % args['confidence_level'])
-
- axarr[1].set_xticklabels(axis_ticklabels)
- axarr[1].set_xlabel("Forecast Probability (uncalibrated; %)")
- axarr[1].set_ylabel("Observed Relative Frequency (%)")
- axarr[1].set_title(r'$\bf{b)}$ $\bf{Reliability}$ $\bf{diagram}$')
-
- for ax in axarr:
- ax.set_xticks(axis_ticks)
- ax.set_yticks(axis_ticks)
- ax.set_yticklabels(axis_ticklabels)
- ax.grid(color='black', alpha=0.1)
- ax.set_xlim(0, 1)
- ax.set_ylim(0, 1)
- ################################################################################################################
-
- ############################################# Data table (panel c) #############################################
- columns = ['Threshold*', 'CSI', 'POD %', 'SR %', 'FAR %', 'FB'] # Column names
- rows = ['50 km', '100 km', '150 km', '200 km', '250 km'] # Row names
-
- table_axis = plt.axes([0.063, -0.06, 0.4, 0.2])
- table_axis.set_title(r'$\bf{c)}$ $\bf{Data}$ $\bf{table}$ [confidence level = %d%%]' % args['confidence_level'], x=0.5, y=0.135, pad=-4)
- table_axis.axis('off')
- table_axis.text(0.16, -2.7, '* probability threshold where CSI is maximized') # Add disclaimer for probability threshold column
- stats_table = table_axis.table(cellText=cell_text, rowLabels=rows, rowColours=boundary_colors, colLabels=columns, cellLoc='center')
- stats_table.scale(1, 3) # Make the table larger
-
- ### Shade the cells and make the cell text larger ###
- for cell in stats_table._cells:
- stats_table._cells[cell].set_alpha(.7)
- stats_table._cells[cell].set_text_props(fontproperties=FontProperties(size='xx-large', stretch='expanded'))
- ################################################################################################################
-
- ########################################## Spatial CSI map (panel d) ###########################################
- # Colorbar keyword arguments
- cbar_kwargs = {'label': 'CSI', 'pad': 0}
-
- # Adjust the spatial CSI plot based on the domain
- if args['domain'] == 'conus':
- spatial_axis_extent = [0.52, -0.582, 0.512, 0.544]
- cbar_kwargs['shrink'] = 0.919
- spatial_plot_xlabels = [-140, -105, -70]
- spatial_plot_ylabels = [30, 40, 50]
- else:
- spatial_axis_extent = [0.538, -0.6, 0.48, 0.577]
- cbar_kwargs['shrink'] = 0.862
- spatial_plot_xlabels = [-150, -120, -90, -60, -30, 0, 120, 150, 180]
- spatial_plot_ylabels = [0, 20, 40, 60, 80]
-
- right_labels = False # Disable latitude labels on the right side of the subplot
- top_labels = False # Disable longitude labels on top of the subplot
- left_labels = True # Latitude labels on the left side of the subplot
- bottom_labels = True # Longitude labels on the bottom of the subplot
-
- ## Set up the spatial CSI plot ###
- csi_cmap = plotting_utils.truncated_colormap('gnuplot2', maxval=0.9, n=10)
- extent = settings.DEFAULT_DOMAIN_EXTENTS[args['domain']]
- spatial_axis = plt.axes(spatial_axis_extent, projection=ccrs.Miller(central_longitude=250))
- spatial_axis_title_text = r'$\bf{d)}$ $\bf{%d}$ $\bf{km}$ $\bf{CSI}$ $\bf{map}$' % args['map_neighborhood']
- plotting_utils.plot_background(extent=extent, ax=spatial_axis)
- norm_probs = colors.Normalize(vmin=0.1, vmax=1)
- spatial_csi_ds = xr.where(spatial_csi_ds >= 0.1, spatial_csi_ds, float("NaN"))
- spatial_csi_ds.sel(boundary=args['map_neighborhood']).plot(ax=spatial_axis, x='longitude', y='latitude', norm=norm_probs,
- cmap=csi_cmap, transform=ccrs.PlateCarree(), alpha=0.6, cbar_kwargs=cbar_kwargs)
- spatial_axis.set_title(spatial_axis_title_text)
- gl = spatial_axis.gridlines(draw_labels=True, zorder=0, dms=True, x_inline=False, y_inline=False)
- gl.right_labels = right_labels
- gl.top_labels = top_labels
- gl.left_labels = left_labels
- gl.bottom_labels = bottom_labels
- gl.xlocator = FixedLocator(spatial_plot_xlabels)
- gl.ylocator = FixedLocator(spatial_plot_ylabels)
- gl.xlabel_style = {'size': 7}
- gl.ylabel_style = {'size': 8}
- ################################################################################################################
-
- if args['domain'] == 'conus':
- domain_text = args['domain'].upper()
- else:
- domain_text = args['domain']
- plt.suptitle(f'Model %d: %ss over %s domain' % (args['model_number'], settings.DEFAULT_FRONT_NAMES[front_label], domain_text), fontsize=20) # Create and plot the main title
-
- filename = f"%s/model_%d/performance_%s_%s_%s_{args['data_source']}.png" % (args['model_dir'], args['model_number'], front_label, args['dataset'], args['domain'])
- if args['data_source'] != 'era5':
- filename = filename.replace('.png', '_f%03d.png' % args['forecast_hour']) # Add forecast hour to the end of the filename
-
- plt.tight_layout()
- plt.savefig(filename, bbox_inches='tight', dpi=500)
- plt.close()
-
- # Thresholds for creating deterministic splines with different front types and neighborhoods
- model_properties['front_obj_thresholds'] = max_csi_thresholds
-
- with open(model_properties_filepath, 'wb') as f:
- pickle.dump(model_properties, f)
diff --git a/evaluation/predict.py b/evaluation/predict.py
deleted file mode 100644
index a8b42d7..0000000
--- a/evaluation/predict.py
+++ /dev/null
@@ -1,650 +0,0 @@
-"""
-Generate predictions with a model.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.9.18
-"""
-import argparse
-import pandas as pd
-import numpy as np
-import xarray as xr
-import os
-import sys
-import tensorflow as tf
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) # this line allows us to import scripts outside the current directory
-from utils import data_utils, settings
-import file_manager as fm
-
-
-def _add_image_to_map(stitched_map_probs: np.array,
- image_probs: np.array,
- map_created: bool,
- domain_images_lon: int,
- domain_images_lat: int,
- lon_image: int,
- lat_image: int,
- image_size_lon: int,
- image_size_lat: int,
- lon_image_spacing: int,
- lat_image_spacing: int):
- """
- Add model prediction to the stitched map.
-
- Parameters
- ----------
- stitched_map_probs: Numpy array
- Array of front probabilities for the final map.
- image_probs: Numpy array
- Array of front probabilities for the current prediction/image.
- map_created: bool
- Boolean flag that declares whether the final map has been completed.
- domain_images_lon: int
- Number of images along the longitude dimension of the domain.
- domain_images_lat: int
- Number of images along the latitude dimension of the domain.
- lon_image: int
- Current image number along the longitude dimension.
- lat_image: int
- Current image number along the latitude dimension.
- image_size_lon: int
- Number of pixels along the longitude dimension of the model predictions.
- image_size_lat: int
- Number of pixels along the latitude dimension of the model predictions.
- lon_image_spacing: int
- Number of pixels between each image along the longitude dimension.
- lat_image_spacing: int
- Number of pixels between each image along the latitude dimension.
-
- Returns
- -------
- map_created: bool
- Boolean flag that declares whether the final map has been completed.
- stitched_map_probs: array
- Array of front probabilities for the final map.
- """
-
- if lon_image == 0: # If the image is on the western edge of the domain
- if lat_image == 0: # If the image is on the northern edge of the domain
- # Add first image to map
- stitched_map_probs[:, 0: image_size_lon, 0: image_size_lat] = \
- image_probs[:, :image_size_lon, :image_size_lat]
-
- if domain_images_lon == 1 and domain_images_lat == 1:
- map_created = True
-
- elif lat_image != domain_images_lat - 1: # If the image is not on the northern nor the southern edge of the domain
- # Take the maximum of the overlapping pixels along sets of constant latitude
- stitched_map_probs[:, 0: image_size_lon, int(lat_image * lat_image_spacing):int((lat_image-1)*lat_image_spacing) + image_size_lat] = \
- np.maximum(stitched_map_probs[:, 0: image_size_lon, int(lat_image * lat_image_spacing):int((lat_image-1)*lat_image_spacing) + image_size_lat],
- image_probs[:, :image_size_lon, :image_size_lat - lat_image_spacing])
-
- # Add the remaining pixels of the current image to the map
- stitched_map_probs[:, 0: image_size_lon, int(lat_image_spacing * (lat_image-1)) + image_size_lat:int(lat_image_spacing * lat_image) + image_size_lat] = \
- image_probs[:, :image_size_lon, image_size_lat - lat_image_spacing:image_size_lat]
-
- if domain_images_lon == 1 and domain_images_lat == 2:
- map_created = True
-
- else: # If the image is on the southern edge of the domain
- # Take the maximum of the overlapping pixels along sets of constant latitude
- stitched_map_probs[:, 0: image_size_lon, int(lat_image * lat_image_spacing):int((lat_image-1)*lat_image_spacing) + image_size_lat] = \
- np.maximum(stitched_map_probs[:, :image_size_lon, int(lat_image * lat_image_spacing):int((lat_image-1)*lat_image_spacing) + image_size_lat],
- image_probs[:, :image_size_lon, :image_size_lat - lat_image_spacing])
-
- # Add the remaining pixels of the current image to the map
- stitched_map_probs[:, 0: image_size_lon, int(lat_image_spacing * (lat_image-1)) + image_size_lat:] = \
- image_probs[:, :image_size_lon, image_size_lat - lat_image_spacing:image_size_lat]
-
- if domain_images_lon == 1 and domain_images_lat > 2:
- map_created = True
-
- elif lon_image != domain_images_lon - 1: # If the image is not on the western nor the eastern edge of the domain
- if lat_image == 0: # If the image is on the northern edge of the domain
- # Take the maximum of the overlapping pixels along sets of constant longitude
- stitched_map_probs[:, int(lon_image * lon_image_spacing):int((lon_image-1)*lon_image_spacing) + image_size_lon, 0: image_size_lat] = \
- np.maximum(stitched_map_probs[:, int(lon_image * lon_image_spacing):int((lon_image-1)*lon_image_spacing) + image_size_lon, 0: image_size_lat],
- image_probs[:, :image_size_lon - lon_image_spacing, :image_size_lat])
-
- # Add the remaining pixels of the current image to the map
- stitched_map_probs[:, int(lon_image_spacing * (lon_image-1)) + image_size_lon:lon_image_spacing * lon_image + image_size_lon, 0: image_size_lat] = \
- image_probs[:, image_size_lon - lon_image_spacing:image_size_lon, :image_size_lat]
-
- if domain_images_lon == 2 and domain_images_lat == 1:
- map_created = True
-
- elif lat_image != domain_images_lat - 1: # If the image is not on the northern nor the southern edge of the domain
- # Take the maximum of the overlapping pixels along sets of constant longitude
- stitched_map_probs[:, int(lon_image * lon_image_spacing):int(lon_image_spacing * (lon_image-1)) + image_size_lon, int(lat_image * lat_image_spacing):int(lat_image * lat_image_spacing) + image_size_lat] = \
- np.maximum(stitched_map_probs[:, int(lon_image * lon_image_spacing):int(lon_image_spacing * (lon_image-1)) + image_size_lon, int(lat_image * lat_image_spacing):int(lat_image * lat_image_spacing) + image_size_lat],
- image_probs[:, :image_size_lon - lon_image_spacing, :image_size_lat])
-
- # Take the maximum of the overlapping pixels along sets of constant latitude
- stitched_map_probs[:, int(lon_image * lon_image_spacing): int(lon_image * lon_image_spacing) + image_size_lon, int(lat_image * lat_image_spacing):int(lat_image_spacing * (lat_image-1)) + image_size_lat] = \
- np.maximum(stitched_map_probs[:, int(lon_image * lon_image_spacing): int(lon_image * lon_image_spacing) + image_size_lon, int(lat_image * lat_image_spacing):int(lat_image_spacing * (lat_image-1)) + image_size_lat],
- image_probs[:, :image_size_lon, :image_size_lat - lat_image_spacing])
-
- # Add the remaining pixels of the current image to the map
- stitched_map_probs[:, int(lon_image_spacing * (lon_image-1)) + image_size_lon:lon_image_spacing * lon_image + image_size_lon, int(lat_image_spacing * (lat_image-1)) + image_size_lat:int(lat_image_spacing * lat_image) + image_size_lat] = \
- image_probs[:, image_size_lon - lon_image_spacing:image_size_lon, image_size_lat - lat_image_spacing:image_size_lat]
-
- if domain_images_lon == 2 and domain_images_lat == 2:
- map_created = True
-
- else: # If the image is on the southern edge of the domain
- # Take the maximum of the overlapping pixels along sets of constant longitude
- stitched_map_probs[:, int(lon_image * lon_image_spacing):int(lon_image_spacing * (lon_image-1)) + image_size_lon, int(lat_image * lat_image_spacing):] = \
- np.maximum(stitched_map_probs[:, int(lon_image * lon_image_spacing):int(lon_image_spacing * (lon_image-1)) + image_size_lon, int(lat_image * lat_image_spacing):],
- image_probs[:, :image_size_lon - lon_image_spacing, :image_size_lat])
-
- # Take the maximum of the overlapping pixels along sets of constant latitude
- stitched_map_probs[:, int(lon_image * lon_image_spacing):int(lon_image * lon_image_spacing) + image_size_lon, int(lat_image * lat_image_spacing):int(lat_image_spacing * (lat_image-1)) + image_size_lat] = \
- np.maximum(stitched_map_probs[:, int(lon_image * lon_image_spacing):int(lon_image * lon_image_spacing) + image_size_lon, int(lat_image * lat_image_spacing):int(lat_image_spacing * (lat_image-1)) + image_size_lat],
- image_probs[:, :image_size_lon, :image_size_lat - lat_image_spacing])
-
- # Add the remaining pixels of the current image to the map
- stitched_map_probs[:, int(lon_image_spacing * (lon_image-1)) + image_size_lon:lon_image_spacing * lon_image + image_size_lon, int(lat_image_spacing * (lat_image-1)) + image_size_lat:] = \
- image_probs[:, image_size_lon - lon_image_spacing:image_size_lon, image_size_lat - lat_image_spacing:image_size_lat]
-
- if domain_images_lon == 2 and domain_images_lat > 2:
- map_created = True
- else:
- if lat_image == 0: # If the image is on the northern edge of the domain
- # Take the maximum of the overlapping pixels along sets of constant longitude
- stitched_map_probs[:, int(lon_image * lon_image_spacing):int(lon_image_spacing * (lon_image-1)) + image_size_lon, 0: image_size_lat] = \
- np.maximum(stitched_map_probs[:, int(lon_image * lon_image_spacing):int(lon_image_spacing * (lon_image-1)) + image_size_lon, 0: image_size_lat],
- image_probs[:, :image_size_lon - lon_image_spacing, :image_size_lat])
-
- # Add the remaining pixels of the current image to the map
- stitched_map_probs[:, int(lon_image_spacing * (lon_image-1)) + image_size_lon:, 0: image_size_lat] = \
- image_probs[:, image_size_lon - lon_image_spacing:image_size_lon, :image_size_lat]
-
- if domain_images_lon > 2 and domain_images_lat == 1:
- map_created = True
-
- elif lat_image != domain_images_lat - 1: # If the image is not on the northern nor the southern edge of the domain
- # Take the maximum of the overlapping pixels along sets of constant longitude
- stitched_map_probs[:, int(lon_image * lon_image_spacing):int(lon_image_spacing * (lon_image-1)) + image_size_lon, int(lat_image * lat_image_spacing):int(lat_image * lat_image_spacing) + image_size_lat] = \
- np.maximum(stitched_map_probs[:, int(lon_image * lon_image_spacing):int(lon_image_spacing * (lon_image-1)) + image_size_lon, int(lat_image * lat_image_spacing):int(lat_image * lat_image_spacing) + image_size_lat], image_probs[:, :image_size_lon - lon_image_spacing, :image_size_lat])
-
- # Take the maximum of the overlapping pixels along sets of constant latitude
- stitched_map_probs[:, int(lon_image * lon_image_spacing):, int(lat_image * lat_image_spacing):int(lat_image_spacing * (lat_image-1)) + image_size_lat] = \
- np.maximum(stitched_map_probs[:, int(lon_image * lon_image_spacing):, int(lat_image * lat_image_spacing):int(lat_image_spacing * (lat_image-1)) + image_size_lat], image_probs[:, :image_size_lon, :image_size_lat - lat_image_spacing])
-
- # Add the remaining pixels of the current image to the map
- stitched_map_probs[:, int(lon_image_spacing * (lon_image-1)) + image_size_lon:, int(lat_image_spacing * (lat_image-1)) + image_size_lat:int(lat_image_spacing * lat_image) + image_size_lat] = image_probs[:, image_size_lon - lon_image_spacing:image_size_lon, image_size_lat - lat_image_spacing:image_size_lat]
-
- if domain_images_lon > 2 and domain_images_lat == 2:
- map_created = True
- else: # If the image is on the southern edge of the domain
- # Take the maximum of the overlapping pixels along sets of constant longitude
- stitched_map_probs[:, int(lon_image * lon_image_spacing):, int(lat_image * lat_image_spacing):] = \
- np.maximum(stitched_map_probs[:, int(lon_image * lon_image_spacing):, int(lat_image * lat_image_spacing):],
- image_probs[:, :image_size_lon, :image_size_lat])
-
- # Take the maximum of the overlapping pixels along sets of constant latitude
- stitched_map_probs[:, int(lon_image * lon_image_spacing):, int(lat_image * lat_image_spacing):int(lat_image_spacing * (lat_image-1)) + image_size_lat] = \
- np.maximum(stitched_map_probs[:, int(lon_image * lon_image_spacing):, int(lat_image * lat_image_spacing):int(lat_image_spacing * (lat_image-1)) + image_size_lat],
- image_probs[:, :image_size_lon, :image_size_lat - lat_image_spacing])
-
- # Add the remaining pixels of the current image to the map
- stitched_map_probs[:, int(lon_image_spacing * (lon_image-1)) + image_size_lon:, int(lat_image_spacing * (lat_image-1)) + image_size_lat:int(lat_image_spacing * lat_image) + image_size_lat] = \
- image_probs[:, image_size_lon - lon_image_spacing:image_size_lon, image_size_lat - lat_image_spacing:image_size_lat]
-
- map_created = True
-
- return stitched_map_probs, map_created
-
-
-def find_matches_for_domain(domain_size: tuple | list, image_size: tuple | list, compatibility_mode: bool = False, compat_images: tuple | list = None):
- """
- Function that outputs the number of images that can be stitched together with the specified domain length and the length
- of the domain dimension output by the model. This is also used to determine the compatibility of declared image and
- parameters for model predictions.
-
- Parameters
- ----------
- domain_size: iterable object with 2 integers
- Number of pixels along each dimension of the final stitched map (lon lat).
- image_size: iterable object with 2 integers
- Number of pixels along each dimension of the model's output (lon lat).
- compatibility_mode: bool
- Boolean flag that declares whether the function is being used to check compatibility of given parameters.
- compat_images: iterable object with 2 integers
- Number of images declared for the stitched map in each dimension (lon lat). (Compatibility mode only)
- """
-
- ######################################### Check the parameters for errors ##########################################
- if not isinstance(domain_size, (tuple, list)):
- raise TypeError(f"Expected a tuple or list for domain_size, received {type(domain_size)}")
- elif len(domain_size) != 2:
- raise TypeError(f"Tuple or list for domain_images must be length 2, received length {len(domain_size)}")
-
- if not isinstance(image_size, (tuple, list)):
- raise TypeError(f"Expected a tuple or list for image_size, received {type(image_size)}")
- elif len(image_size) != 2:
- raise TypeError(f"Tuple or list for image_size must be length 2, received length {len(image_size)}")
-
- if compatibility_mode is not None and not isinstance(compatibility_mode, bool):
- raise TypeError(f"compatibility_mode must be a boolean, received {type(compatibility_mode)}")
-
- if compat_images is not None:
- if not isinstance(compat_images, (tuple, list)):
- raise TypeError(f"Expected a tuple or list for compat_images, received {type(compat_images)}")
- elif len(compat_images) != 2:
- raise TypeError(f"Tuple or list for compat_images must be length 2, received length {len(compat_images)}")
- ####################################################################################################################
-
- if compatibility_mode:
- """ These parameters are used when checking the compatibility of image stitching arguments. """
- compat_images_lon = compat_images[0] # Number of images in the longitude direction
- compat_images_lat = compat_images[1] # Number of images in the latitude direction
- else:
- compat_images_lon, compat_images_lat = None, None
-
- # All of these boolean variables must be True after the compatibility check or else a ValueError is returned
- lon_images_are_compatible = False
- lat_images_are_compatible = False
-
- num_matches = [0, 0] # Total number of matching image arguments found for each dimension
-
- lon_image_matches = []
- lat_image_matches = []
-
- for lon_images in range(1, domain_size[0]-image_size[0] + 2): # Image counter for longitude dimension
- if lon_images > 1:
- lon_spacing = (domain_size[0]-image_size[0])/(lon_images-1) # Spacing between images in the longitude dimension
- else:
- lon_spacing = 0
- if lon_spacing - int(lon_spacing) == 0 and lon_spacing > 1 and image_size[0]-lon_spacing > 0: # Check compatibility of latitude image spacing
- lon_image_matches.append(lon_images) # Add longitude image match to list
- num_matches[0] += 1
- if compatibility_mode:
- if compat_images_lon == lon_images: # If the number of images for the compatibility check equals the match
- lon_images_are_compatible = True
- elif lon_spacing == 0 and domain_size[0] - image_size[0] == 0:
- lon_image_matches.append(lon_images) # Add longitude image match to list
- num_matches[0] += 1
- if compatibility_mode:
- if compat_images_lon == lon_images: # If the number of images for the compatibility check equals the match
- lon_images_are_compatible = True
-
- if num_matches[0] == 0:
- raise ValueError(f"No compatible value for domain_images[0] was found with domain_size[0]={domain_size[0]} and image_size[0]={image_size[0]}.")
- if compatibility_mode:
- if not lon_images_are_compatible:
- raise ValueError(f"domain_images[0]={compat_images_lon} is not compatible with domain_size[0]={domain_size[0]} "
- f"and image_size[0]={image_size[0]}.\n"
- f"====> Compatible values for domain_images[0] given domain_size[0]={domain_size[0]} "
- f"and image_size[0]={image_size[0]}: {lon_image_matches}")
- else:
- print(f"Compatible longitude images: {lon_image_matches}")
-
- for lat_images in range(1, domain_size[1]-image_size[1]+2): # Image counter for latitude dimension
- if lat_images > 1:
- lat_spacing = (domain_size[1]-image_size[1])/(lat_images-1) # Spacing between images in the latitude dimension
- else:
- lat_spacing = 0
- if lat_spacing - int(lat_spacing) == 0 and lat_spacing > 1 and image_size[1]-lat_spacing > 0: # Check compatibility of latitude image spacing
- lat_image_matches.append(lat_images) # Add latitude image match to list
- num_matches[1] += 1
- if compatibility_mode:
- if compat_images_lat == lat_images: # If the number of images for the compatibility check equals the match
- lat_images_are_compatible = True
- elif lat_spacing == 0 and domain_size[1] - image_size[1] == 0:
- lat_image_matches.append(lat_images) # Add latitude image match to list
- num_matches[1] += 1
- if compatibility_mode:
- if compat_images_lat == lat_images: # If the number of images for the compatibility check equals the match
- lat_images_are_compatible = True
-
- if num_matches[1] == 0:
- raise ValueError(f"No compatible value for domain_images[1] was found with domain_size[1]={domain_size[1]} and image_size[1]={image_size[1]}.")
- if compatibility_mode:
- if not lat_images_are_compatible:
- raise ValueError(f"domain_images[1]={compat_images_lat} is not compatible with domain_size[1]={domain_size[1]} "
- f"and image_size[1]={image_size[1]}.\n"
- f"====> Compatible values for domain_images[1] given domain_size[1]={domain_size[1]} "
- f"and image_size[1]={image_size[1]}: {lat_image_matches}")
- else:
- print(f"Compatible latitude images: {lat_image_matches}")
-
-
-def create_model_prediction_dataset(stitched_map_probs: np.array, lats: np.array, lons: np.array, front_types: str | list):
- """
- Create an Xarray dataset containing model predictions.
-
- Parameters
- ----------
- stitched_map_probs: np.array
- Numpy array with probabilities for the given front type(s).
- Shape/dimensions: [front types, longitude, latitude]
- lats: np.array
- 1D array of latitude values.
- lons: np.array
- 1D array of longitude values.
- front_types: str or list
- Front types within the dataset. See documentation in utils.data_utils.reformat fronts for more information.
-
- Returns
- -------
- probs_ds: xr.Dataset
- Xarray dataset containing front probabilities predicted by the model for each front type.
- """
-
- ######################################### Check the parameters for errors ##########################################
- if not isinstance(stitched_map_probs, np.ndarray):
- raise TypeError(f"stitched_map_probs must be a NumPy array, received {type(stitched_map_probs)}")
- if not isinstance(lats, np.ndarray):
- raise TypeError(f"lats must be a NumPy array, received {type(lats)}")
- if not isinstance(lons, np.ndarray):
- raise TypeError(f"lons must be a NumPy array, received {type(lons)}")
- if not isinstance(front_types, (tuple, list)):
- raise TypeError(f"Expected a tuple or list for front_types, received {type(front_types)}")
- ####################################################################################################################
-
- if front_types == 'F_BIN' or front_types == 'MERGED-F_BIN' or front_types == 'MERGED-T':
- probs_ds = xr.Dataset(
- {front_types: (('longitude', 'latitude'), stitched_map_probs[0])},
- coords={'latitude': lats, 'longitude': lons})
- elif front_types == 'MERGED-F':
- probs_ds = xr.Dataset(
- {'CF_merged': (('longitude', 'latitude'), stitched_map_probs[0]),
- 'WF_merged': (('longitude', 'latitude'), stitched_map_probs[1]),
- 'SF_merged': (('longitude', 'latitude'), stitched_map_probs[2]),
- 'OF_merged': (('longitude', 'latitude'), stitched_map_probs[3])},
- coords={'latitude': lats, 'longitude': lons})
- elif front_types == 'MERGED-ALL':
- probs_ds = xr.Dataset(
- {'CF_merged': (('longitude', 'latitude'), stitched_map_probs[0]),
- 'WF_merged': (('longitude', 'latitude'), stitched_map_probs[1]),
- 'SF_merged': (('longitude', 'latitude'), stitched_map_probs[2]),
- 'OF_merged': (('longitude', 'latitude'), stitched_map_probs[3]),
- 'TROF_merged': (('longitude', 'latitude'), stitched_map_probs[4]),
- 'INST': (('longitude', 'latitude'), stitched_map_probs[5]),
- 'DL': (('longitude', 'latitude'), stitched_map_probs[6])},
- coords={'latitude': lats, 'longitude': lons})
- elif type(front_types) == list:
- probs_ds_dict = dict({})
- for probs_ds_index, front_type in enumerate(front_types):
- probs_ds_dict[front_type] = (('longitude', 'latitude'), stitched_map_probs[probs_ds_index])
- probs_ds = xr.Dataset(probs_ds_dict, coords={'latitude': lats, 'longitude': lons})
- else:
- raise ValueError(f"'{front_types}' is not a valid set of front types.")
-
- return probs_ds
-
-
-if __name__ == '__main__':
- """
- All arguments listed in the examples are listed via argparse in alphabetical order below this comment block.
- """
- parser = argparse.ArgumentParser()
- parser.add_argument('--bootstrap', action='store_true', help='Bootstrap data?')
- parser.add_argument('--dataset', type=str, help="Dataset for which to make predictions if prediction_method is 'random' or 'all'. Options are:"
- "'training', 'validation', 'test'")
- parser.add_argument('--datetime', type=int, nargs=4, help='Date and time of the data. Pass 4 ints in the following order: year, month, day, hour')
- parser.add_argument('--domain', type=str, help='Domain of the data.')
- parser.add_argument('--domain_images', type=int, nargs=2, help='Number of images for each dimension the final stitched map for predictions: lon, lat')
- parser.add_argument('--domain_size', type=int, nargs=2, help='Lengths of the dimensions of the final stitched map for predictions: lon, lat')
- parser.add_argument('--forecast_hour', type=int, help='Forecast hour for the GDAS data')
- parser.add_argument('--find_matches', action='store_true', help='Find matches for stitching predictions?')
- parser.add_argument('--generate_predictions', action='store_true', help='Generate prediction plots?')
- parser.add_argument('--calculate_stats', action='store_true', help='generate stats')
- parser.add_argument('--calibrate_model', action='store_true', help='Calibrate model')
- parser.add_argument('--gpu_device', type=int, help='GPU device number.')
- parser.add_argument('--image_size', type=int, nargs=2, help="Number of pixels along each dimension of the model's output: lon, lat")
- parser.add_argument('--learning_curve', action='store_true', help='Plot learning curve')
- parser.add_argument('--memory_growth', action='store_true', help='Use memory growth on the GPU')
- parser.add_argument('--model_dir', type=str, help='Directory for the models.')
- parser.add_argument('--model_number', type=int, help='Model number.')
- parser.add_argument('--num_iterations', type=int, default=10000, help='Number of iterations to perform when bootstrapping the data.')
- parser.add_argument('--num_rand_predictions', type=int, default=10, help='Number of random predictions to make.')
- parser.add_argument('--fronts_netcdf_indir', type=str, help='Main directory for the netcdf files containing frontal objects.')
- parser.add_argument('--variables_netcdf_indir', type=str, help='Main directory for the netcdf files containing variable data.')
- parser.add_argument('--plot_performance_diagrams', action='store_true', help='Plot performance diagrams for a model?')
- parser.add_argument('--prediction_method', type=str, help="Prediction method. Options are: 'datetime', 'random', 'all'")
- parser.add_argument('--prediction_plot', action='store_true', help='Create plot')
- parser.add_argument('--save_map', action='store_true', help='Save maps of the model predictions?')
- parser.add_argument('--save_probabilities', action='store_true', help='Save model prediction data out to netcdf files?')
- parser.add_argument('--save_statistics', action='store_true', help='Save performance statistics data out to netcdf files?')
- parser.add_argument('--data_source', type=str, default='era5', help='Data source for variables')
-
- args = vars(parser.parse_args())
-
- gpus = tf.config.list_physical_devices(device_type='GPU') # Find available GPUs
- if len(gpus) > 0:
- tf.config.set_visible_devices(devices=gpus[0], device_type='GPU')
- gpus = tf.config.get_visible_devices(device_type='GPU') # List of selected GPUs
-
- # Allow for memory growth on the GPU. This will only use the GPU memory that is required rather than allocating all the GPU's memory.
- if args['memory_growth']:
- tf.config.experimental.set_memory_growth(device=gpus[0], enable=True)
-
- else:
- print('WARNING: No GPUs found, all computations will be performed on CPUs.')
- tf.config.set_visible_devices([], 'GPU')
-
- ### Model properties ###
- model_properties = pd.read_pickle(f"{args['model_dir']}/model_{args['model_number']}/model_{args['model_number']}_properties.pkl")
- model_type = model_properties['model_type']
-
- if args['image_size'] is None:
- args['image_size'] = model_properties['image_size'] # The image size does not include the last dimension of the input size as it only represents the number of channels
-
- try:
- front_types = model_properties['dataset_properties']['front_types']
- variables = model_properties['dataset_properties']['variables']
- pressure_levels = model_properties['dataset_properties']['pressure_levels']
- except KeyError: # Some older models do not have the dataset_properties dictionary
- front_types = model_properties['front_types']
- variables = model_properties['variables']
- pressure_levels = model_properties['pressure_levels']
-
- normalization_parameters = model_properties['normalization_parameters']
-
- classes = model_properties['classes']
- test_years, valid_years = model_properties['test_years'], model_properties['validation_years']
-
- if args['domain_images'] is None:
- args['domain_images'] = settings.DEFAULT_DOMAIN_IMAGES[args['domain']]
- domain_extent_indices = settings.DEFAULT_DOMAIN_INDICES[args['domain']]
- domain_extent = settings.DEFAULT_DOMAIN_EXTENTS[args['domain']]
-
- ### Properties of the final map made from stitched images ###
- domain_images_lon, domain_images_lat = args['domain_images'][0], args['domain_images'][1]
- domain_size_lon, domain_size_lat = domain_extent_indices[1] - domain_extent_indices[0], domain_extent_indices[3] - domain_extent_indices[2]
- image_size_lon, image_size_lat = args['image_size'][0], args['image_size'][1] # Dimensions of the model's predictions
-
- if domain_images_lon > 1:
- lon_image_spacing = int((domain_size_lon - image_size_lon)/(domain_images_lon-1))
- else:
- lon_image_spacing = 0
-
- if domain_images_lat > 1:
- lat_image_spacing = int((domain_size_lat - image_size_lat)/(domain_images_lat-1))
- else:
- lat_image_spacing = 0
-
- model = fm.load_model(args['model_number'], args['model_dir'])
- num_dimensions = len(model.layers[0].input_shape[0]) - 2
-
- ############################################### Load variable files ################################################
- variable_files_obj = fm.DataFileLoader(args['variables_netcdf_indir'], data_file_type='%s-netcdf' % args['data_source'])
- variable_files_obj.validation_years = valid_years
- variable_files_obj.test_years = test_years
-
- if args['dataset'] is not None:
- variable_files = getattr(variable_files_obj, 'data_files_' + args['dataset'])
- else:
- variable_files = getattr(variable_files_obj, 'data_files')
- ####################################################################################################################
-
- dataset_kwargs = {'engine': 'netcdf4'} # Keyword arguments for loading variable files with xarray
- coords_sel_kwargs = {'longitude': slice(domain_extent[0], domain_extent[1]),
- 'latitude': slice(domain_extent[3], domain_extent[2])}
-
- if args['prediction_method'] == 'datetime':
- timestep_str = '%d%02d%02d%02d' % (args['datetime'][0], args['datetime'][1], args['datetime'][2], args['datetime'][3])
- if args['data_source'] == 'era5':
- datetime_index = [index for index, file in enumerate(variable_files) if timestep_str in file][0]
- variable_files = [variable_files[datetime_index], ]
- else:
- variable_files = [file for file in variable_files if timestep_str in file]
-
- subdir_base = '%s_%dx%d' % (args['domain'], args['domain_images'][0], args['domain_images'][1])
-
- num_files = len(variable_files)
-
- num_chunks = int(np.ceil(num_files / settings.MAX_FILE_CHUNK_SIZE)) # Number of files/timesteps to process at once
- chunk_indices = np.linspace(0, num_files, num_chunks + 1, dtype=int)
-
- for chunk_no in range(num_chunks):
-
- files_in_chunk = variable_files[chunk_indices[chunk_no]:chunk_indices[chunk_no + 1]]
- print(f"Preparing chunk {chunk_no + 1}/{num_chunks}")
- variable_ds = xr.open_mfdataset(files_in_chunk, **dataset_kwargs).sel(**coords_sel_kwargs)[variables]
-
- if args['data_source'] == 'era5':
- variable_ds = variable_ds.sel(pressure_level=pressure_levels).transpose('time', 'longitude', 'latitude', 'pressure_level')
- else:
- variable_ds = variable_ds.sel(pressure_level=pressure_levels).transpose('time', 'forecast_hour', 'longitude', 'latitude', 'pressure_level')
- forecast_hours = variable_ds['forecast_hour'].values
-
- # Older 2D models were trained with the pressure levels not in the proper order
- if args['model_number'] in [7805504, 7866106, 7961517]:
- variable_ds = variable_ds.isel(pressure_level=[0, 4, 3, 2, 1])
-
- image_lats = variable_ds.latitude.values[:domain_size_lat]
- image_lons = variable_ds.longitude.values[:domain_size_lon]
-
- timestep_predict_size = settings.TIMESTEP_PREDICT_SIZE[args['domain']]
-
- if args['data_source'] != 'era5':
- num_forecast_hours = len(forecast_hours)
- timestep_predict_size /= num_forecast_hours
- timestep_predict_size = int(timestep_predict_size)
-
- num_timesteps = len(variable_ds['time'].values)
- num_batches = int(np.ceil(num_timesteps / timestep_predict_size))
-
- for batch_no in range(num_batches):
-
- print(f"======== Chunk {chunk_no + 1}/{num_chunks}: batch {batch_no + 1}/{num_batches} ========")
-
- variable_batch_ds = variable_ds.isel(time=slice(batch_no * timestep_predict_size, (batch_no + 1) * timestep_predict_size)) # Select timesteps for the current batch
- variable_batch_ds = data_utils.normalize_variables(variable_batch_ds, normalization_parameters)
-
- timesteps = variable_batch_ds['time'].values
- num_timesteps_in_batch = len(timesteps)
- map_created = False # Boolean that determines whether the final stitched map has been created
-
- if args['data_source'] == 'era5':
- stitched_map_probs = np.empty(shape=[num_timesteps_in_batch, classes-1, domain_size_lon, domain_size_lat])
- else:
- stitched_map_probs = np.empty(shape=[num_timesteps_in_batch, len(forecast_hours), classes-1, domain_size_lon, domain_size_lat])
-
- for lat_image in range(domain_images_lat):
- lat_index = int(lat_image * lat_image_spacing)
- for lon_image in range(domain_images_lon):
- print(f"image %d/%d" % (int(lat_image*domain_images_lon) + lon_image + 1, int(domain_images_lon*domain_images_lat)))
- lon_index = int(lon_image * lon_image_spacing)
-
- # Select the current image
- variable_batch_ds_new = variable_batch_ds[variables].isel(longitude=slice(lon_index, lon_index + args['image_size'][0]),
- latitude=slice(lat_index, lat_index + args['image_size'][1])).to_array().values
-
- if args['data_source'] == 'era5':
- variable_batch_ds_new = variable_batch_ds_new.transpose([1, 2, 3, 4, 0]) # (time, longitude, latitude, pressure level, variable)
- else:
- variable_batch_ds_new = variable_batch_ds_new.transpose([1, 2, 3, 4, 5, 0]) # (time, forecast hour, longitude, latitude, pressure level, variable)
-
- if num_dimensions == 2:
-
- ### Combine pressure levels and variables into one dimension ###
- variable_batch_ds_new_shape = np.shape(variable_batch_ds_new)
- variable_batch_ds_new = variable_batch_ds_new.reshape(*[dim_size for dim_size in variable_batch_ds_new_shape[:-2]], variable_batch_ds_new_shape[-2] * variable_batch_ds_new_shape[-1])
-
- transpose_indices = (0, 3, 1, 2) # New order of indices for model predictions (time, front type, longitude, latitude)
-
- ##################################### Generate the predictions #####################################
- if args['data_source'] != 'era5':
-
- variable_ds_new_shape = np.shape(variable_batch_ds_new)
- variable_batch_ds_new = variable_batch_ds_new.reshape(variable_ds_new_shape[0] * variable_ds_new_shape[1], *[dim_size for dim_size in variable_ds_new_shape[2:]])
-
- prediction = model.predict(variable_batch_ds_new, batch_size=settings.GPU_PREDICT_BATCH_SIZE, verbose=0)
- num_dims_in_pred = len(np.shape(prediction))
-
- if model_type == 'unet':
- if num_dims_in_pred == 4: # 2D labels, prediction shape: (time, lat, lon, front type)
- image_probs = np.transpose(prediction[:, :, :, 1:], transpose_indices) # transpose the predictions
- else: # if num_dims_in_pred == 5; 3D labels, prediction shape: (time, lat, lon, pressure level, front type)
- image_probs = np.transpose(np.amax(prediction[:, :, :, :, 1:], axis=3), transpose_indices) # Take the maximum probability over the vertical dimension and transpose the predictions
-
- elif model_type == 'unet_3plus':
-
- try:
- deep_supervision = model_properties['deep_supervision']
- except KeyError:
- deep_supervision = True # older models do not have this dictionary key, so just set it to True
-
- if deep_supervision:
- if num_dims_in_pred == 5: # 2D labels, prediction shape: (output level, time, lon, lat, front type)
- image_probs = np.transpose(prediction[0][:, :, :, 1:], transpose_indices) # transpose the predictions
- else: # if num_dims_in_pred == 6; 3D labels, prediction shape: (output level, time, lon, lat, pressure level, front type)
- image_probs = np.transpose(np.amax(prediction[0][:, :, :, :, 1:], axis=3), transpose_indices) # Take the maximum probability over the vertical dimension and transpose the predictions
- else:
- if num_dims_in_pred == 4: # 2D labels, prediction shape: (time, lon, lat, front type)
- image_probs = np.transpose(prediction[:, :, :, 1:], transpose_indices) # transpose the predictions
- else: # if num_dims_in_pred == 5; 3D labels, prediction shape: (time, lat, lon, pressure level, front type)
- image_probs = np.transpose(np.amax(prediction[:, :, :, :, 1:], axis=3), transpose_indices) # Take the maximum probability over the vertical dimension and transpose the predictions
-
- # Add predictions to the map
- if args['data_source'] != 'era5':
- for timestep in range(num_timesteps_in_batch):
- for fcst_hr_index in range(num_forecast_hours):
- stitched_map_probs[timestep][fcst_hr_index], map_created = _add_image_to_map(stitched_map_probs[timestep][fcst_hr_index], image_probs[timestep * num_forecast_hours + fcst_hr_index], map_created, domain_images_lon, domain_images_lat, lon_image, lat_image,
- image_size_lon, image_size_lat, lon_image_spacing, lat_image_spacing)
-
- else: # if args['data_source'] == 'era5'
- for timestep in range(num_timesteps_in_batch):
- stitched_map_probs[timestep], map_created = _add_image_to_map(stitched_map_probs[timestep], image_probs[timestep], map_created, domain_images_lon, domain_images_lat, lon_image, lat_image,
- image_size_lon, image_size_lat, lon_image_spacing, lat_image_spacing)
- ####################################################################################################
-
- if map_created:
-
- ### Create subdirectories for the data if they do not exist ###
- if not os.path.isdir('%s/model_%d/maps/%s' % (args['model_dir'], args['model_number'], subdir_base)):
- os.makedirs('%s/model_%d/maps/%s' % (args['model_dir'], args['model_number'], subdir_base))
- print("New subdirectory made:", '%s/model_%d/maps/%s' % (args['model_dir'], args['model_number'], subdir_base))
- if not os.path.isdir('%s/model_%d/probabilities/%s' % (args['model_dir'], args['model_number'], subdir_base)):
- os.makedirs('%s/model_%d/probabilities/%s' % (args['model_dir'], args['model_number'], subdir_base))
- print("New subdirectory made:", '%s/model_%d/probabilities/%s' % (args['model_dir'], args['model_number'], subdir_base))
- if not os.path.isdir('%s/model_%d/statistics/%s' % (args['model_dir'], args['model_number'], subdir_base)):
- os.makedirs('%s/model_%d/statistics/%s' % (args['model_dir'], args['model_number'], subdir_base))
- print("New subdirectory made:", '%s/model_%d/statistics/%s' % (args['model_dir'], args['model_number'], subdir_base))
-
- if args['data_source'] != 'era5':
-
- for timestep_no, timestep in enumerate(timesteps):
- timestep = str(timestep)
- for fcst_hr_index, forecast_hour in enumerate(forecast_hours):
- time = f'{timestep[:4]}%s%s%s' % (timestep[5:7], timestep[8:10], timestep[11:13])
- probs_ds = create_model_prediction_dataset(stitched_map_probs[timestep_no][fcst_hr_index], image_lats, image_lons, front_types)
- probs_ds = probs_ds.expand_dims({'time': np.atleast_1d(timestep), 'forecast_hour': np.atleast_1d(forecast_hours[fcst_hr_index])})
- filename_base = 'model_%d_%s_%s_%s_f%03d_%dx%d' % (args['model_number'], time, args['domain'], args['data_source'], forecast_hours[fcst_hr_index], domain_images_lon, domain_images_lat)
-
- outfile = '%s/model_%d/probabilities/%s/%s_probabilities.nc' % (args['model_dir'], args['model_number'], subdir_base, filename_base)
- probs_ds.to_netcdf(path=outfile, engine='netcdf4', mode='w')
-
- else:
-
- for timestep_no, timestep in enumerate(timesteps):
- time = f'{timestep[:4]}%s%s%s' % (timestep[5:7], timestep[8:10], timestep[11:13])
- probs_ds = create_model_prediction_dataset(stitched_map_probs[timestep_no], image_lats, image_lons, front_types)
- probs_ds = probs_ds.expand_dims({'time': np.atleast_1d(timestep)})
- filename_base = 'model_%d_%s_%s_%dx%d' % (args['model_number'], time, args['domain'], domain_images_lon, domain_images_lat)
-
- outfile = '%s/model_%d/probabilities/%s/%s_probabilities.nc' % (args['model_dir'], args['model_number'], subdir_base, filename_base)
- probs_ds.to_netcdf(path=outfile, engine='netcdf4', mode='w')
diff --git a/evaluation/predict_tf.py b/evaluation/predict_tf.py
deleted file mode 100644
index ff5be31..0000000
--- a/evaluation/predict_tf.py
+++ /dev/null
@@ -1,146 +0,0 @@
-"""
-**** EXPERIMENTAL SCRIPT TO REPLACE 'predict.py' IN THE NEAR FUTURE ****
-
-Generate predictions using a model with tensorflow datasets.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.7.24.D1
-"""
-import argparse
-import sys
-import os
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) # this line allows us to import scripts outside the current directory
-import file_manager as fm
-import numpy as np
-import pandas as pd
-from utils.settings import *
-import xarray as xr
-import tensorflow as tf
-
-
-if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- parser.add_argument('--dataset', type=str, help="Dataset for which to make predictions. Options are: 'training', 'validation', 'test'")
- parser.add_argument('--year_and_month', type=int, nargs=2, help="Year and month for which to make predictions.")
- parser.add_argument('--model_dir', type=str, help='Directory for the models.')
- parser.add_argument('--model_number', type=int, help='Model number.')
- parser.add_argument('--tf_indir', type=str, help='Directory for the tensorflow dataset that will be used when generating predictions.')
- parser.add_argument('--data_source', type=str, default='era5', help='Data source for variables')
- parser.add_argument('--gpu_device', type=int, nargs='+', help='GPU device numbers.')
- parser.add_argument('--memory_growth', action='store_true', help='Use memory growth on the GPU')
- parser.add_argument('--overwrite', action='store_true', help="Overwrite any existing prediction files.")
-
- args = vars(parser.parse_args())
-
- model_properties = pd.read_pickle('%s/model_%d/model_%d_properties.pkl' % (args['model_dir'], args['model_number'], args['model_number']))
- dataset_properties = pd.read_pickle('%s/dataset_properties.pkl' % args['tf_indir'])
-
- domain = dataset_properties['domain']
-
- if domain == 'conus':
- hour_interval = 3
- else:
- hour_interval = 6
-
- # Some older models do not have the 'dataset_properties' dictionary
- try:
- front_types = model_properties['dataset_properties']['front_types']
- num_dims = model_properties['dataset_properties']['num_dims']
- except KeyError:
- front_types = model_properties['front_types']
- if args['model_number'] in [6846496, 7236500, 7507525]:
- num_dims = (3, 3)
-
- if args['dataset'] is not None and args['year_and_month'] is not None:
- raise ValueError("--dataset and --year_and_month cannot be passed together.")
- elif args['dataset'] is None and args['year_and_month'] is None:
- raise ValueError("At least one of [--dataset, --year_and_month] must be passed.")
- elif args['year_and_month'] is not None:
- years, months = [args['year_and_month'][0]], [args['year_and_month'][1]]
- else:
- years, months = model_properties['%s_years' % args['dataset']], range(1, 13)
-
- ### Make sure that the dataset has the same attributes as the model ###
- if model_properties['normalization_parameters'] != dataset_properties['normalization_parameters']:
- raise ValueError("Cannot evaluate model with the selected dataset. Reason: normalization parameters do not match")
- if model_properties['dataset_properties']['front_types'] != dataset_properties['front_types']:
- raise ValueError("Cannot evaluate model with the selected dataset. Reason: front types do not match "
- f"(model: {model_properties['dataset_properties']['front_types']}, dataset: {dataset_properties['front_types']})")
- if model_properties['dataset_properties']['variables'] != dataset_properties['variables']:
- raise ValueError("Cannot evaluate model with the selected dataset. Reason: variables do not match "
- f"(model: {model_properties['dataset_properties']['variables']}, dataset: {dataset_properties['variables']})")
- if model_properties['dataset_properties']['pressure_levels'] != dataset_properties['pressure_levels']:
- raise ValueError("Cannot evaluate model with the selected dataset. Reason: pressure levels do not match "
- f"(model: {model_properties['dataset_properties']['pressure_levels']}, dataset: {dataset_properties['pressure_levels']})")
-
- gpus = tf.config.list_physical_devices(device_type='GPU') # Find available GPUs
- if len(gpus) > 0:
-
- print("Number of GPUs available: %d" % len(gpus))
-
- # Only make the selected GPU(s) visible to TensorFlow
- if args['gpu_device'] is not None:
- tf.config.set_visible_devices(devices=[gpus[gpu] for gpu in args['gpu_device']], device_type='GPU')
- gpus = tf.config.get_visible_devices(device_type='GPU') # List of selected GPUs
- print("Using %d GPU(s):" % len(gpus), gpus)
-
- # Allow for memory growth on the GPU. This will only use the GPU memory that is required rather than allocating all the GPU's memory.
- if args['memory_growth']:
- tf.config.experimental.set_memory_growth(device=[gpu for gpu in gpus][0], enable=True)
-
- else:
- print('WARNING: No GPUs found, all computations will be performed on CPUs.')
- tf.config.set_visible_devices([], 'GPU')
-
- # The axis that the predicts will be concatenated on depends on the shape of the output, which is determined by deep supervision
- if model_properties['deep_supervision']:
- concat_axis = 1
- else:
- concat_axis = 0
-
- tf_ds_obj = fm.DataFileLoader(args['tf_indir'], data_file_type='%s-tensorflow' % args['data_source'])
-
- lons = np.arange(DEFAULT_DOMAIN_EXTENTS[domain][0], DEFAULT_DOMAIN_EXTENTS[domain][1] + 0.25, 0.25)
- lats = np.arange(DEFAULT_DOMAIN_EXTENTS[domain][2], DEFAULT_DOMAIN_EXTENTS[domain][3] + 0.25, 0.25)[::-1]
-
- model = fm.load_model(args['model_number'], args['model_dir'])
-
- for year in years:
-
- tf_ds_obj.test_years = [year, ]
- files_for_year = tf_ds_obj.data_files_test
-
- for month in months:
-
- prediction_dataset_path = '%s/model_%d/probabilities/model_%d_pred_%s_%d%02d.nc' % (args['model_dir'], args['model_number'], args['model_number'], domain, year, month)
- if os.path.isfile(prediction_dataset_path) and not args['overwrite']:
- print("WARNING: %s exists, pass the --overwrite argument to overwrite existing data." % prediction_dataset_path)
- continue
-
- input_file = [file for file in files_for_year if '_%d%02d' % (year, month) in file][0]
- tf_ds = tf.data.Dataset.load(input_file)
- time_array = np.arange(np.datetime64(f"{input_file[-9:-5]}-{input_file[-5:-3]}"),
- np.datetime64(f"{input_file[-9:-5]}-{input_file[-5:-3]}") + np.timedelta64(1, "M"),
- np.timedelta64(hour_interval, "h"))
-
- assert len(tf_ds) == len(time_array) # make sure tensorflow dataset has all timesteps
-
- tf_ds = tf_ds.batch(GPU_PREDICT_BATCH_SIZE)
- prediction = np.array(model.predict(tf_ds)).astype(np.float16)
-
- if model_properties['deep_supervision']:
- prediction = prediction[0, ...] # select the top output of the model, since it is the only one we care about
-
- if num_dims[1] == 3:
- # Take the maxmimum probability for each front type over the vertical dimension (pressure levels)
- prediction = np.amax(prediction, axis=3) # shape: (time, longitude, latitude, front type)
-
- prediction = prediction[..., 1:] # remove the 'no front' type from the array
- prediction = np.transpose(prediction, (0, 2, 1, 3)) # shape: (time, latitude, longitude, front type)
-
- xr.Dataset(data_vars={front_type: (('time', 'latitude', 'longitude'), prediction[:, :, :, front_type_no])
- for front_type_no, front_type in enumerate(front_types)},
- coords={'time': time_array, 'longitude': lons, 'latitude': lats}).astype('float32').\
- to_netcdf(path=prediction_dataset_path, mode='w', engine='netcdf4')
-
- del prediction # Delete the prediction variable so it can be recreated for the next year
diff --git a/evaluation/prediction_plot.py b/evaluation/prediction_plot.py
deleted file mode 100644
index 010f482..0000000
--- a/evaluation/prediction_plot.py
+++ /dev/null
@@ -1,202 +0,0 @@
-"""
-Plot model predictions.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.9.10
-"""
-import itertools
-import argparse
-import pandas as pd
-import cartopy.crs as ccrs
-import matplotlib.pyplot as plt
-import numpy as np
-import xarray as xr
-from matplotlib import cm, colors # Here we explicitly import the cm and color modules to suppress a PyCharm bug
-import os
-import sys
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) # this line allows us to import scripts outside the current directory
-from utils import data_utils, settings
-from utils.plotting_utils import plot_background
-from skimage.morphology import skeletonize
-
-
-if __name__ == '__main__':
- """
- All arguments listed in the examples are listed via argparse in alphabetical order below this comment block.
- """
- parser = argparse.ArgumentParser()
- parser.add_argument('--datetime', type=int, nargs=4, help='Date and time of the data. Pass 4 ints in the following order: year, month, day, hour')
- parser.add_argument('--domain', type=str, required=True, help='Domain of the data.')
- parser.add_argument('--domain_images', type=int, nargs=2, help='Number of images for each dimension the final stitched map for predictions: lon, lat')
- parser.add_argument('--forecast_hour', type=int, help='Forecast hour for the GDAS data')
- parser.add_argument('--model_dir', type=str, required=True, help='Directory for the models.')
- parser.add_argument('--model_number', type=int, required=True, help='Model number.')
- parser.add_argument('--fronts_netcdf_indir', type=str, help='Main directory for the netcdf files containing frontal objects.')
- parser.add_argument('--data_source', type=str, default='era5', help="Source of the variable data (ERA5, GDAS, etc.)")
- parser.add_argument('--prob_mask', type=float, nargs=2, default=[0.1, 0.1],
- help="Probability mask and the step/interval for the probability contours. Probabilities smaller than the mask will not be plotted.")
- parser.add_argument('--calibration', type=int,
- help="Neighborhood calibration distance in kilometers. Possible neighborhoods are 50, 100, 150, 200, and 250 km.")
- parser.add_argument('--deterministic', action='store_true', help="Plot deterministic splines.")
- parser.add_argument('--targets', action='store_true', help="Plot ground truth targets/labels.")
- parser.add_argument('--contours', action='store_true', help="Plot probability contours.")
-
- args = vars(parser.parse_args())
-
- if args['deterministic'] and args['targets']:
- raise TypeError("Cannot plot deterministic splines and ground truth targets at the same time. Only one of --deterministic, --targets may be passed")
-
- if args['domain_images'] is None:
- args['domain_images'] = [1, 1]
-
- DEFAULT_COLORBAR_POSITION = {'conus': 0.74, 'full': 0.84, 'global': 0.74}
- cbar_position = DEFAULT_COLORBAR_POSITION['conus']
-
- model_properties = pd.read_pickle(f"{args['model_dir']}/model_{args['model_number']}/model_{args['model_number']}_properties.pkl")
-
- args['data_source'] = args['data_source'].lower()
-
- extent = settings.DEFAULT_DOMAIN_EXTENTS[args['domain']]
-
- year, month, day, hour = args['datetime'][0], args['datetime'][1], args['datetime'][2], args['datetime'][3]
-
- ### Attempt to pull predictions from a yearly netcdf file generated with tensorflow datasets, otherwise try to pull a single netcdf file ###
- try:
- probs_file = f"{args['model_dir']}/model_{args['model_number']}/probabilities/model_{args['model_number']}_pred_{args['domain']}_{year}%02d.nc" % month
- fronts_file = '%s/%d%02d/FrontObjects_%d%02d%02d%02d_full.nc' % (args['fronts_netcdf_indir'], year, month, year, month, day, hour)
- plot_filename = '%s/model_%d/maps/model_%d_%d%02d%02d%02d_%s.png' % (args['model_dir'], args['model_number'], args['model_number'], year, month, day, hour, args['domain'])
- probs_ds = xr.open_mfdataset(probs_file).sel(time=['%d-%02d-%02dT%02d' % (year, month, day, hour), ])
- except OSError:
- subdir_base = '%s_%dx%d' % (args['domain'], args['domain_images'][0], args['domain_images'][1])
- probs_dir = f"{args['model_dir']}/model_{args['model_number']}/probabilities/{subdir_base}"
-
- if args['forecast_hour'] is not None:
- timestep = np.datetime64('%d-%02d-%02dT%02d' % (year, month, day, hour)).astype(object)
- forecast_timestep = timestep if args['forecast_hour'] == 0 else timestep + np.timedelta64(args['forecast_hour'], 'h').astype(object)
- new_year, new_month, new_day, new_hour = forecast_timestep.year, forecast_timestep.month, forecast_timestep.day, forecast_timestep.hour - (forecast_timestep.hour % 3)
- fronts_file = '%s/%s%s/FrontObjects_%s%s%s%02d_full.nc' % (args['fronts_netcdf_indir'], new_year, new_month, new_year, new_month, new_day, new_hour)
- filename_base = f'model_%d_{year}%02d%02d%02d_%s_%s_f%03d_%dx%d' % (args['model_number'], month, day, hour, args['domain'], args['data_source'], args['forecast_hour'], args['domain_images'][0], args['domain_images'][1])
- else:
- fronts_file = '%s/%d%02d/FrontObjects_%d%02d%02d%02d_full.nc' % (args['fronts_netcdf_indir'], year, month, year, month, day, hour)
- filename_base = f'model_%d_{year}%02d%02d%02d_%s_%dx%d' % (args['model_number'], month, day, hour, args['domain'], args['domain_images'][0], args['domain_images'][1])
- args['data_source'] = 'era5'
-
- plot_filename = '%s/model_%d/maps/%s/%s-same.png' % (args['model_dir'], args['model_number'], subdir_base, filename_base)
- probs_file = f'{probs_dir}/{filename_base}_probabilities.nc'
- probs_ds = xr.open_mfdataset(probs_file)
-
- try:
- front_types = model_properties['dataset_properties']['front_types']
- except KeyError:
- front_types = model_properties['front_types']
-
- labels = front_types
- fronts_found = False
-
- if args['targets']:
- right_title = 'Splines: NOAA fronts'
- try:
- fronts = xr.open_dataset(fronts_file).sel(longitude=slice(extent[0], extent[1]), latitude=slice(extent[3], extent[2]))
- fronts = data_utils.reformat_fronts(fronts, front_types=front_types)
- labels = fronts.attrs['labels']
- fronts = xr.where(fronts == 0, float('NaN'), fronts)
- fronts_found = True
- except FileNotFoundError:
- print("No ground truth fronts found")
-
- if type(front_types) == str:
- front_types = [front_types, ]
-
- mask, prob_int = args['prob_mask'][0], args['prob_mask'][1] # Probability mask, contour interval for probabilities
- vmax, cbar_tick_adjust, cbar_label_adjust, n_colors = 1, prob_int, 10, 11
- levels = np.around(np.arange(0, 1 + prob_int, prob_int), 2)
- cbar_ticks = np.around(np.arange(mask, 1 + prob_int, prob_int), 2)
-
- contour_maps_by_type = [settings.DEFAULT_CONTOUR_CMAPS[label] for label in labels]
- front_colors_by_type = [settings.DEFAULT_FRONT_COLORS[label] for label in labels]
- front_names_by_type = [settings.DEFAULT_FRONT_NAMES[label] for label in labels]
-
- cmap_front = colors.ListedColormap(front_colors_by_type, name='from_list', N=len(front_colors_by_type))
- norm_front = colors.Normalize(vmin=1, vmax=len(front_colors_by_type) + 1)
-
- probs_ds = probs_ds.isel(time=0) if args['data_source'] == 'era5' else probs_ds.isel(time=0, forecast_hour=0)
- probs_ds = probs_ds.transpose('latitude', 'longitude')
-
- for key in list(probs_ds.keys()):
-
- if args['deterministic']:
- spline_threshold = model_properties['front_obj_thresholds'][args['domain']][key]['100']
- probs_ds[f'{key}_obj'] = (('latitude', 'longitude'), skeletonize(xr.where(probs_ds[key] > spline_threshold, 1, 0).values.copy(order='C')))
-
- if args['calibration'] is not None:
- try:
- ir_model = model_properties['calibration_models'][args['domain']][key]['%d km' % args['calibration']]
- except KeyError:
- ir_model = model_properties['calibration_models']['conus'][key]['%d km' % args['calibration']]
- original_shape = np.shape(probs_ds[key].values)
- probs_ds[key].values = ir_model.predict(probs_ds[key].values.flatten()).reshape(original_shape)
- cbar_label = 'Probability (calibrated - %d km)' % args['calibration']
- else:
- cbar_label = 'Probability (uncalibrated)'
-
- if len(front_types) > 1:
- all_possible_front_combinations = itertools.permutations(front_types, r=2)
- for combination in all_possible_front_combinations:
- probs_ds[combination[0]].values = np.where(probs_ds[combination[0]].values > probs_ds[combination[1]].values - 0.02, probs_ds[combination[0]].values, 0)
-
- probs_ds = xr.where(probs_ds > mask, probs_ds, float("NaN"))
- if args['data_source'] != 'era5':
- valid_time = timestep + np.timedelta64(args['forecast_hour'], 'h').astype(object)
- data_title = f"Run: {args['data_source'].upper()} {year}-%02d-%02d-%02dz F%03d \nPredictions valid: %d-%02d-%02d-%02dz" % (month, day, hour, args['forecast_hour'], valid_time.year, valid_time.month, valid_time.day, valid_time.hour)
- else:
- data_title = 'Data: ERA5 reanalysis %d-%02d-%02d-%02dz\n' \
- 'Predictions valid: %d-%02d-%02d-%02dz' % (year, month, day, hour, year, month, day, hour)
-
- fig, ax = plt.subplots(1, 1, figsize=(22, 8), subplot_kw={'projection': ccrs.PlateCarree(central_longitude=np.mean(extent[:2]))})
- plot_background(extent, ax=ax, linewidth=0.5)
- # ax.gridlines(draw_labels=True, zorder=0)
-
- cbar_front_labels = []
- cbar_front_ticks = []
-
- for front_no, front_key, front_name, front_label, cmap in zip(range(1, len(front_names_by_type) + 1), list(probs_ds.keys()), front_names_by_type, front_types, contour_maps_by_type):
-
- if args['contours']:
- cmap_probs, norm_probs = cm.get_cmap(cmap, n_colors), colors.Normalize(vmin=0, vmax=vmax)
- probs_ds[front_key].plot.contourf(ax=ax, x='longitude', y='latitude', norm=norm_probs, levels=levels, cmap=cmap_probs, transform=ccrs.PlateCarree(), alpha=0.75, add_colorbar=False)
-
- cbar_ax = fig.add_axes([cbar_position + (front_no * 0.015), 0.24, 0.015, 0.64])
- cbar = plt.colorbar(cm.ScalarMappable(norm=norm_probs, cmap=cmap_probs), cax=cbar_ax, boundaries=levels[1:], alpha=0.75)
- cbar.set_ticklabels([])
-
- if args['deterministic']:
- right_title = 'Splines: Deterministic first-guess fronts'
- cmap_deterministic = colors.ListedColormap(['None', front_colors_by_type[front_no - 1]], name='from_list', N=2)
- norm_deterministic = colors.Normalize(vmin=0, vmax=1)
- probs_ds[f'{front_key}_obj'].plot(ax=ax, x='longitude', y='latitude', cmap=cmap_deterministic, norm=norm_deterministic,
- transform=ccrs.PlateCarree(), alpha=0.9, add_colorbar=False)
-
- if fronts_found:
- fronts['identifier'].plot(ax=ax, x='longitude', y='latitude', cmap=cmap_front, norm=norm_front, transform=ccrs.PlateCarree(), add_colorbar=False)
-
- cbar_front_labels.append(front_name)
- cbar_front_ticks.append(front_no + 0.5)
-
- if args['contours']:
- cbar.set_label(cbar_label, rotation=90)
- cbar.set_ticks(cbar_ticks)
- cbar.set_ticklabels(cbar_ticks)
-
- cbar_front = plt.colorbar(cm.ScalarMappable(norm=norm_front, cmap=cmap_front), ax=ax, alpha=0.75, orientation='horizontal', shrink=0.5, pad=0.02)
- cbar_front.set_ticks(cbar_front_ticks)
- cbar_front.set_ticklabels(cbar_front_labels)
- cbar_front.set_label(r'$\bf{Front}$ $\bf{type}$')
-
- if fronts_found or args['deterministic']:
- ax.set_title(right_title, loc='right')
-
- ax.set_title('')
- ax.set_title(data_title, loc='left')
-
- plt.savefig(plot_filename, bbox_inches='tight', dpi=300)
- plt.close()
diff --git a/file_manager.py b/file_manager.py
deleted file mode 100644
index 63c0b6b..0000000
--- a/file_manager.py
+++ /dev/null
@@ -1,669 +0,0 @@
-"""
-Functions in this code manage data files and models.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.6.25.D1
-"""
-
-import argparse
-from glob import glob
-import os
-import pandas as pd
-import shutil
-import tarfile
-
-
-def compress_files(
- main_dir: str,
- glob_file_string: str,
- tar_filename: str,
- remove_files: bool = False,
- status_printout: bool = True):
- """
- Compress files into a TAR file.
-
- Parameters
- ----------
- main_dir: str
- Main directory where the files are located and where the TAR file will be saved.
- glob_file_string: str
- String of the names of the files to compress.
- tar_filename: str
- Name of the compressed TAR file that will be made. Do not include the .tar.gz extension in the name, this is added automatically.
- remove_files: bool
- Setting this to true will remove the files after they have been compressed to a TAR file.
- status_printout: bool
- Setting this to true will provide printouts of the status of the compression.
-
- Examples
- --------
- <<<<< start example >>>>
-
- import fronts.file_manager as fm
-
- main_dir = 'C:/Users/username/data_files'
- glob_file_string = '*matching_string.pkl'
- tar_filename = 'matching_files' # Do not add the .tar.gz extension, this is done automatically
-
- compress_files(main_dir, glob_file_string, tar_filename, remove_files=True, status_printout=False) # Compress files and remove them after compression into a TAR file
-
- <<<<< end example >>>>>
- """
-
- ########################################### Check the parameters for errors ########################################
- if not isinstance(main_dir, str):
- raise TypeError(f"main_dir must be a string, received {type(main_dir)}")
- if not isinstance(glob_file_string, str):
- raise TypeError(f"glob_file_string must be a string, received {type(glob_file_string)}")
- if not isinstance(tar_filename, str):
- raise TypeError(f"tar_filename must be a string, received {type(tar_filename)}")
- if not isinstance(remove_files, bool):
- raise TypeError(f"remove_files must be a boolean, received {type(remove_files)}")
- if not isinstance(status_printout, bool):
- raise TypeError(f"status_printout must be a boolean, received {type(status_printout)}")
- ####################################################################################################################
-
- uncompressed_size = 0 # MB
-
- ### Gather a list of files containing the specified string ###
- files = list(sorted(glob(f"{main_dir}/{glob_file_string}")))
- if len(files) == 0:
- raise OSError("No files found")
- else:
- print(f"{len(files)} files found")
-
- num_files = len(files) # Total number of files
-
- ### Create the TAR file ###
- with tarfile.open(f"{main_dir}/{tar_filename}.tar.gz", "w:gz") as tarF:
-
- ### Iterate through all of the available files ###
- for file in range(num_files):
- tarF.add(files[file], arcname=files[file].replace(main_dir, '')) # Add the file to the TAR file
- tarF_size = os.path.getsize(f"{main_dir}/{tar_filename}.tar.gz")/1e6 # Compressed size of the files within the TAR file (megabytes)
- uncompressed_size += os.path.getsize(files[file])/1e6 # Uncompressed size of the files within the TAR file (megabytes)
-
- ### Print out the current status of the compression (if enabled) ###
- if status_printout:
- print(f'({file+1}/{num_files}) {uncompressed_size:,.2f} MB ---> {tarF_size:,.2f} MB ({100*(1-(tarF_size/uncompressed_size)):.1f}% compression ratio)', end='\r')
-
- # Completion message
- print(f"Successfully compressed {len(files)} files: ",
- f'{uncompressed_size:,.2f} MB ---> {tarF_size:,.2f} MB ({100*(1-(tarF_size/uncompressed_size)):.1f}% compression ratio)')
-
- ### Remove the files that were added to the TAR archive (if enabled; does NOT affect the contents of the TAR file just created) ###
- if remove_files:
- for file in files:
- os.remove(file)
- print(f"Successfully deleted {len(files)} files")
-
-
-def delete_grouped_files(
- main_dir: str,
- glob_file_string: str,
- num_subdir: int):
- """
- Deletes grouped files with names matching given strings.
-
- Parameters
- ----------
- main_dir: str
- Main directory or directories where the grouped files are located.
- glob_file_string: str
- String of the names of the files to delete.
- num_subdir: int
- Number of subdirectory layers in the main directory.
-
- Examples
- --------
- <<<<< start example >>>>>
-
- import fronts.file_manager as fm
-
- main_dir = 'C:/Users/username/data_files'
- glob_file_string = '*matching_string.pkl'
- num_subdir = 3 # Check in the 3rd level of the directories within the main directory
-
- fm.delete_grouped_files(main_dir, glob_file_string, num_subdir)
-
- <<<<< end example >>>>>
- """
-
- ########################################### Check the parameters for errors ########################################
- if not isinstance(main_dir, str):
- raise TypeError(f"main_dir must be a string, received {type(main_dir)}")
- if not isinstance(glob_file_string, str):
- raise TypeError(f"glob_file_string must be a string, received {type(glob_file_string)}")
- if not isinstance(num_subdir, int):
- raise TypeError(f"num_subdir must be an integer, received {type(num_subdir)}")
- ####################################################################################################################
-
- subdir_string = '' # This string will be modified depending on the provided value of num_subdir
- for i in range(num_subdir):
- subdir_string += '/*'
- subdir_string += '/'
- glob_file_string = subdir_string + glob_file_string # String that will be used to match with patterns in filenames
-
- files_to_delete = list(sorted(glob("%s%s" % (main_dir, glob_file_string)))) # Search for files in the given directory that have patterns matching the file string
-
- # Delete all the files
- print("Deleting %d files...." % len(files_to_delete), end='')
- for file in files_to_delete:
- try:
- os.remove(file)
- except PermissionError:
- shutil.rmtree(file)
- print("done")
-
-
-def extract_tarfile(main_dir: str,
- tar_filename: str):
- """
- Extract all the contents of a TAR file.
-
- Parameters
- ----------
- main_dir: str
- Main directory where the TAR file is located. This is also where the extracted files will be placed.
- tar_filename: str
- Name of the compressed TAR file. Do NOT include the .tar.gz extension.
-
- Examples
- --------
- <<<<< start example >>>>
-
- import fronts.file_manager as fm
-
- main_dir = 'C:/Users/username/data_files'
- tar_filename = 'foo_tarfile' # Do not add the .tar.gz extension
-
- fm.extract_tarfile(main_dir, glob_file_string, tar_filename, remove_files=True, status_printout=False) # Compress files and remove them after compression into a TAR file
-
- <<<<< end example >>>>>
- """
-
- ########################################### Check the parameters for errors ########################################
- if not isinstance(main_dir, str):
- raise TypeError(f"main_dir must be a string, received {type(main_dir)}")
- if not isinstance(tar_filename, str):
- raise TypeError(f"tar_filename must be a string, received {type(tar_filename)}")
- ####################################################################################################################
-
- with tarfile.open(f"{main_dir}/{tar_filename}.tar.gz", 'r') as tarF:
- tarF.extractall(main_dir)
- print(f"Successfully extracted {main_dir}/{tar_filename}")
-
-
-class DataFileLoader:
- """
- Object that loads and manages ERA5, GDAS, GFS, and front object files.
- """
- def __init__(
- self,
- file_dir: str,
- data_file_type: str,
- synoptic_only: bool = False):
- """
- When the DataFileLoader object is created, find all netCDF or tensorflow datasets.
-
- Parameters
- ----------
- file_dir: str
- Input directory for the netCDF or tensorflow datasets.
- data_file_type: str
- This string will contain two parts, separated by a hyphen: the source for the variable data, and the type of file/dataset.
- Options for the variable data sources are: 'era5', 'fronts', 'gdas', and 'gfs'.
- Options for the file/dataset type string are: 'netcdf', 'tensorflow'.
- synoptic_only: bool
- Setting this to True will remove any files with timesteps at non-synoptic hours (3, 9, 15, 21z).
-
- Examples
- --------
- <<<<< start example >>>>
-
- import file_manager as fm
-
- file_dir = '/home/user/data' # Directory where the data is stored
- data_file_type = 'era5-netcdf' # We want to load netCDF files containing ERA5 data
-
- era5_file_obj = fm.DataFileLoader(file_dir, data_file_type) # Load the data files
-
- <<<<< end example >>>>
- """
-
- ######################################### Check the parameters for errors ######################################
- if not isinstance(file_dir, str):
- raise TypeError(f"file_dir must be a string, received {type(file_dir)}")
- if not isinstance(data_file_type, str):
- raise TypeError(f"data_file_type must be a string, received {type(data_file_type)}")
- if not isinstance(synoptic_only, bool):
- raise TypeError(f"synoptic_only must be a boolean, received {type(synoptic_only)}")
- ################################################################################################################
-
- valid_data_sources = ['era5', 'fronts', 'gdas', 'gfs']
- valid_file_types = ['netcdf', 'tensorflow']
-
- data_file_type = data_file_type.lower().split('-')
- self._file_type = data_file_type[1]
-
- if self._file_type == 'netcdf':
- self._file_extension = '.nc'
- self._subdir_glob = '/*/'
- elif self._file_type == 'tensorflow':
- self._file_extension = '_tf'
- self._subdir_glob = '/'
- else:
- raise TypeError(f"'%s' is not a valid file type, valid types are: {', '.join(valid_file_types)}" % self._file_type)
-
- if data_file_type[0] not in valid_data_sources:
- raise TypeError(f"'%s' is not a valid data source, valid sources are: {', '.join(valid_data_sources)}" % data_file_type[0])
-
- if data_file_type[0] == 'fronts':
- self._file_prefix = 'FrontObjects'
- else:
- self._file_prefix = data_file_type[0]
-
- self._all_data_files = sorted(glob("%s%s%s*%s" % (file_dir, self._subdir_glob, self._file_prefix, self._file_extension))) # All data files without filtering
- self.data_file_information = [os.path.basename(file).split('_')[1:] for file in self._all_data_files] # timesteps in the unfiltered data files
-
- if synoptic_only:
- ### Filter out non-synoptic hours (3, 9, 15, 21z) ###
- for file_info in self.data_file_information:
- if any('%02d' % hour in file_info[0][-2:] for hour in [3, 9, 15, 21]):
- self._all_data_files.pop(self.data_file_information.index(file_info))
- self.data_file_information.pop(self.data_file_information.index(file_info))
-
- ### All available options for specific filters ###
- self._all_forecast_hours = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
- self._all_years = (2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022)
-
- self.reset_all_filters() # Resetting the filters simply creates the list of data files
-
- ### Current values for the filters used in the files ###
- self._forecast_hours = self._all_forecast_hours
- self._training_years = self._all_years
- self._validation_years = self._all_years
- self._test_years = self._all_years
-
- def reset_all_filters(self):
- """
- Reset the lists of files back to their original states with no filters
- """
-
- if self._file_prefix == 'FrontObjects':
- self.front_files = self._all_data_files
- self.front_files_training = self.front_files
- self.front_files_validation = self.front_files
- self.front_files_test = self.front_files
- else:
- self.data_files = self._all_data_files
- self.data_files_training = self.data_files
- self.data_files_validation = self.data_files
- self.data_files_test = self.data_files
-
- ### If front object files have been loaded to be paired with variable files, reset the front file lists ###
- if hasattr(self, '_all_front_files'):
- self.front_files = self._all_front_files
- self.front_files_training = self.front_files
- self.front_files_validation = self.front_files
- self.front_files_test = self.front_files
-
- ####################################################################################################################
-
- def __get_training_years(self):
- """
- Return the list of training years used
- """
-
- return self._training_years
-
- def __set_training_years(self, training_years: tuple | list):
- """
- Select the training years to load
- """
-
- self.__reset_training_years() # Return file list to last state before training years were modified (no effect is this is first training year selection)
-
- self._training_years = training_years
-
- ### Check that all selected training years are valid ###
- invalid_training_years = [year for year in training_years if year not in self._all_years]
- if len(invalid_training_years) > 0:
- raise TypeError(f"The following training years are not valid: {','.join(sorted(invalid_training_years))}")
-
- self._training_years_not_in_data = [year for year in self._all_years if year not in training_years]
-
- ### Remove unwanted years from the list of files ###
- for year in self._training_years_not_in_data:
- if self._file_prefix == 'FrontObjects':
- self.front_files_training = [file for file in self.front_files_training if '_%s' % str(year) not in file]
- else:
- self.data_files_training = [file for file in self.data_files_training if '_%s' % str(year) not in file]
-
- def __reset_training_years(self):
- """
- Reset training years
- """
-
- if not hasattr(self, '_filtered_data_files_before_training_year_selection'): # If the training years have not been selected yet
- if self._file_prefix == 'FrontObjects':
- self._filtered_data_files_before_training_year_selection = self.front_files_training
- else:
- self._filtered_data_files_before_training_year_selection = self.data_files_training
- else: # Return file list to last state before training years were modified
- if self._file_prefix == 'FrontObjects':
- self.front_files_training = self._filtered_data_files_before_training_year_selection
- else:
- self.data_files_training = self._filtered_data_files_before_training_year_selection
-
- training_years = property(__get_training_years, __set_training_years) # Property method for setting training years
-
- ####################################################################################################################
-
- def __get_validation_years(self):
- """
- Return the list of validation years used
- """
-
- return self._validation_years
-
- def __set_validation_years(self, validation_years: tuple | list):
- """
- Select the validation years to load
- """
-
- self.__reset_validation_years() # Return file list to last state before validation years were modified (no effect is this is first validation year selection)
-
- self._validation_years = validation_years
-
- ### Check that all selected validation years are valid ###
- invalid_validation_years = [year for year in validation_years if year not in self._all_years]
- if len(invalid_validation_years) > 0:
- raise TypeError(f"The following validation years are not valid: {','.join(sorted(invalid_validation_years))}")
-
- self._validation_years_not_in_data = [year for year in self._all_years if year not in validation_years]
-
- ### Remove unwanted years from the list of files ###
- for year in self._validation_years_not_in_data:
- if self._file_prefix == 'FrontObjects':
- self.front_files_validation = [file for file in self.front_files_validation if '_%s' % str(year) not in file]
- else:
- self.data_files_validation = [file for file in self.data_files_validation if '_%s' % str(year) not in file]
-
- def __reset_validation_years(self):
- """
- Reset validation years
- """
-
- if not hasattr(self, '_filtered_data_files_before_validation_year_selection'): # If the validation years have not been selected yet
- if self._file_prefix == 'FrontObjects':
- self._filtered_data_files_before_validation_year_selection = self.front_files_validation
- else:
- self._filtered_data_files_before_validation_year_selection = self.data_files_validation
- else: # Return file list to last state before validation years were modified
- if self._file_prefix == 'FrontObjects':
- self.front_files_validation = self._filtered_data_files_before_validation_year_selection
- else:
- self.data_files_validation = self._filtered_data_files_before_validation_year_selection
-
- validation_years = property(__get_validation_years, __set_validation_years) # Property method for setting validation years
-
- ####################################################################################################################
-
- def __get_test_years(self):
- """
- Return the list of test years used
- """
-
- return self._test_years
-
- def __set_test_years(self, test_years: tuple | list):
- """
- Select the test years to load
- """
-
- self.__reset_test_years() # Return file list to last state before test years were modified (no effect is this is first test year selection)
-
- self._test_years = test_years
-
- ### Check that all selected test years are valid ###
- invalid_test_years = [year for year in test_years if year not in self._all_years]
- if len(invalid_test_years) > 0:
- raise TypeError(f"The following test years are not valid: {','.join(sorted(invalid_test_years))}")
-
- self._test_years_not_in_data = [year for year in self._all_years if year not in test_years]
-
- ### Remove unwanted years from the list of files ###
- for year in self._test_years_not_in_data:
- if self._file_prefix == 'FrontObjects':
- self.front_files_test = [file for file in self.front_files_test if '_%s' % str(year) not in file]
- else:
- self.data_files_test = [file for file in self.data_files_test if '_%s' % str(year) not in file]
-
- def __reset_test_years(self):
- """
- Reset test years
- """
-
- if not hasattr(self, '_filtered_data_files_before_test_year_selection'): # If the test years have not been selected yet
- if self._file_prefix == 'FrontObjects':
- self._filtered_data_files_before_test_year_selection = self.front_files_test
- else:
- self._filtered_data_files_before_test_year_selection = self.data_files_test
- else: # Return file list to last state before test years were modified
- if self._file_prefix == 'FrontObjects':
- self.front_files_test = self._filtered_data_files_before_test_year_selection
- else:
- self.data_files_test = self._filtered_data_files_before_test_year_selection
-
- test_years = property(__get_test_years, __set_test_years) # Property method for setting test years
-
- ####################################################################################################################
-
- def __get_forecast_hours(self):
- """
- Return the list of forecast hours used
- """
-
- if self._file_type == 'era5':
- return None
- else:
- return self._forecast_hours
-
- def __set_forecast_hours(self, forecast_hours: tuple | list):
- """
- Select the forecast hours to load
- """
-
- self.__reset_forecast_hours() # Return file list to last state before forecast hours were modified (no effect is this is first forecast hour selection)
-
- self._forecast_hours = forecast_hours
-
- ### Check that all selected forecast hours are valid ###
- invalid_forecast_hours = [hour for hour in forecast_hours if hour not in self._all_forecast_hours]
- if len(invalid_forecast_hours) > 0:
- raise TypeError(f"The following forecast hours are not valid: {','.join(sorted(str(hour) for hour in invalid_forecast_hours))}")
-
- self._forecast_hours_not_in_data = [hour for hour in self._all_forecast_hours if hour not in forecast_hours]
-
- ### Remove unwanted forecast hours from the list of files ###
- for hour in self._forecast_hours_not_in_data:
- self.data_files = [file for file in self.data_files if '_f%03d_' % hour not in file]
- self.data_files_training = [file for file in self.data_files_training if '_f%03d_' % hour not in file]
- self.data_files_validation = [file for file in self.data_files_validation if '_f%03d_' % hour not in file]
- self.data_files_test = [file for file in self.data_files_test if '_f%03d_' % hour not in file]
-
- def __reset_forecast_hours(self):
- """
- Reset forecast hours in the data files
- """
-
- if not hasattr(self, '_filtered_data_files_before_forecast_hour_selection'): # If the forecast hours have not been selected yet
- self._filtered_data_files_before_forecast_hour_selection = self.data_files
- self._filtered_data_training_files_before_forecast_hour_selection = self.data_files_training
- self._filtered_data_validation_files_before_forecast_hour_selection = self.data_files_validation
- self._filtered_data_test_files_before_forecast_hour_selection = self.data_files_test
- else:
- ### Return file lists to last state before forecast hours were modified ###
- self.data_files = self._filtered_data_files_before_forecast_hour_selection
- self.data_files_training = self._filtered_data_training_files_before_forecast_hour_selection
- self.data_files_validation = self._filtered_data_validation_files_before_forecast_hour_selection
- self.data_files_test = self._filtered_data_test_files_before_forecast_hour_selection
-
- forecast_hours = property(__get_forecast_hours, __set_forecast_hours) # Property method for setting forecast hours
-
- ####################################################################################################################
-
- def __sort_files_by_dataset(self, data_timesteps_used, sort_fronts=False):
- """
- Filter files for the training, validation, and test datasets. This is done by finding indices for each timestep in the 'data_timesteps_used' list.
-
- data_timesteps_used: list
- List of timesteps used when sorting variable and/or frontal object files.
- sort_fronts: bool
- Setting this to True will sort the lists of frontal object files. This will only be True when sorting both variable and frontal object files at the
- same time.
- """
-
- ### Find all indices where timesteps for training, validation, and test datasets are present in the selected variable files ###
- training_indices = [index for index, timestep in enumerate(data_timesteps_used) if any('%d' % training_year in timestep[:4] for training_year in self._training_years)]
- validation_indices = [index for index, timestep in enumerate(data_timesteps_used) if any('%d' % validation_year in timestep[:4] for validation_year in self._validation_years)]
- test_indices = [index for index, timestep in enumerate(data_timesteps_used) if any('%d' % test_year in timestep[:4] for test_year in self._test_years)]
-
- ### Create new variable file lists for training, validation, and test datasets using the indices pulled from above ###
- self.data_files_training = [self.data_files[index] for index in training_indices]
- self.data_files_validation = [self.data_files[index] for index in validation_indices]
- self.data_files_test = [self.data_files[index] for index in test_indices]
-
- if sort_fronts:
- ### Create new frontal object file lists for training, validation, and test datasets using the indices pulled from above ###
- self.front_files_training = [self.front_files[index] for index in training_indices]
- self.front_files_validation = [self.front_files[index] for index in validation_indices]
- self.front_files_test = [self.front_files[index] for index in test_indices]
-
- def pair_with_fronts(self, front_indir, match='same'):
- """
- Parameters
- ----------
- front_indir: str
- Directory where the frontal object files are stored.
- match: 'same', 'forecast', or 'all'
- Method to use when creating the final paired file lists.
- * 'same' will only match variable and front files that have the same file information (timestep, forecast hour, domain)
- * 'forecast' will attempt to match variable files with forecast hours greater than 0 to front files with forecast hours of 0.
- For example, gfs_2023033100_f006_global.nc will be paired with FrontObjects_2023033106_f000_global.nc.
- * 'all' will perform the actions described in both 'same' and 'forecast'.
- """
-
- if self._file_prefix == 'FrontObjects':
- print("WARNING: 'DataFileLoader.pair_with_fronts' can only be used with ERA5, GDAS, or GFS files.")
- return
-
- ######################################### Check the parameters for errors ######################################
- if not isinstance(front_indir, str):
- raise TypeError(f"front_indir must be a string, received {type(front_indir)}")
- ################################################################################################################
-
- if self._file_type == 'netcdf':
- file_prefix = 'FrontObjects'
- elif self._file_type == 'tensorflow':
- file_prefix = 'fronts'
- else:
- raise ValueError(f"The available options for 'file_type' are: 'netcdf', 'tensorflow'. Received: {self._file_type}")
-
- self._all_front_files = sorted(glob("%s%s%s*%s" % (front_indir, self._subdir_glob, file_prefix, self._file_extension))) # All front files without filtering
- front_file_information = [os.path.basename(file).split('_')[1:] for file in self._all_front_files]
-
- if match == 'same':
- file_info_to_retain = [file_info for file_info in self.data_file_information if file_info in front_file_information]
- self.data_files = [self._all_data_files[self.data_file_information.index(file_info)] for file_info in file_info_to_retain]
- self.front_files = [self._all_front_files[front_file_information.index(file_info)] for file_info in file_info_to_retain]
-
- data_timesteps_used = [file_info[0] for file_info in file_info_to_retain]
-
- self.__sort_files_by_dataset(data_timesteps_used, sort_fronts=True)
-
-
-def load_model(model_number: int,
- model_dir: str):
- """
- Load a saved model.
-
- Parameters
- ----------
- model_number: int
- Slurm job number for the model. This is the number in the model's filename.
- model_dir: str
- Main directory for the models.
- """
-
- ######################################### Check the parameters for errors ##########################################
- if not isinstance(model_number, int):
- raise TypeError(f"model_number must be an integer, received {type(model_number)}")
- if not isinstance(model_dir, str):
- raise TypeError(f"model_dir must be a string, received {type(model_dir)}")
- ####################################################################################################################
-
- from tensorflow.keras.models import load_model as lm
- import custom_losses
- import custom_metrics
-
- model_path = f"{model_dir}/model_{model_number}/model_{model_number}.h5"
- model_properties = pd.read_pickle(f"{model_dir}/model_{model_number}/model_{model_number}_properties.pkl")
-
- try:
- loss_string = model_properties['loss_string']
- except KeyError:
- loss_string = model_properties['loss'] # Error in the training sometimes resulted in the incorrect key ('loss' instead of 'loss_string')
- loss_args = model_properties['loss_args']
-
- try:
- metric_string = model_properties['metric_string']
- except KeyError:
- metric_string = model_properties['metric'] # Error in the training sometimes resulted in the incorrect key ('metric' instead of 'metric_string')
- metric_args = model_properties['metric_args']
-
- custom_objects = {}
-
- if 'fss' in loss_string.lower():
- if model_number in [6846496, 7236500, 7507525]:
- loss_string = 'fss_loss'
- custom_objects[loss_string] = custom_losses.fractions_skill_score(**loss_args)
-
- if 'brier' in metric_string.lower() or 'bss' in metric_string.lower():
- if model_number in [6846496, 7236500, 7507525]:
- metric_string = 'bss'
- custom_objects[metric_string] = custom_metrics.brier_skill_score(**metric_args)
-
- if 'csi' in metric_string.lower():
- custom_objects[metric_string] = custom_metrics.critical_success_index(**metric_args)
-
- return lm(model_path, custom_objects=custom_objects)
-
-
-if __name__ == '__main__':
- """
- Warnings
- Do not use leading zeros when declaring the month, day, and hour in 'date'. (ex: if the day is 2, do not type 02)
- Longitude values in the 'new_extent' argument must in the 360-degree coordinate system.
- """
-
- parser = argparse.ArgumentParser()
- parser.add_argument('--compress_files', action='store_true', help='Compress files')
- parser.add_argument('--delete_grouped_files', action='store_true', help='Delete a set of files')
- parser.add_argument('--extract_tarfile', action='store_true', help='Extract a TAR file')
- parser.add_argument('--glob_file_string', type=str, help='String of the names of the files to compress or delete.')
- parser.add_argument('--main_dir', type=str, help='Main directory for subdirectory creation or where the files in question are located.')
- parser.add_argument('--num_subdir', type=int, help='Number of subdirectory layers in the main directory.')
- parser.add_argument('--tar_filename', type=str, help='Name of the TAR file.')
- args = parser.parse_args()
- provided_arguments = vars(args)
-
- if args.compress_files:
- compress_files(args.main_dir, args.glob_file_string, args.tar_filename)
-
- if args.delete_grouped_files:
- delete_grouped_files(args.main_dir, args.glob_file_string, args.num_subdir)
-
- if args.extract_tarfile:
- extract_tarfile(args.main_dir, args.tar_filename)
diff --git a/models.py b/models.py
deleted file mode 100644
index 87608d7..0000000
--- a/models.py
+++ /dev/null
@@ -1,1185 +0,0 @@
-"""
-Deep learning models:
- - U-Net
- - U-Net ensemble
- - U-Net+
- - U-Net++
- - U-Net 3+
- - Attention U-Net
-
-TODO:
- * Allow models to have a unique number of encoder and decoder levels (e.g. 3 encoder levels and 5 decoder levels)
- * Add temporal U-Nets
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.8.18.D1
-"""
-
-from tensorflow.keras.models import Model
-from tensorflow.keras.layers import Concatenate, Input
-from utils import unet_utils
-import numpy as np
-
-
-def unet(
- input_shape: tuple[int],
- num_classes: int,
- pool_size: int | tuple[int] | list[int],
- upsample_size: int | tuple[int] | list[int],
- levels: int,
- filter_num: tuple[int] | list[int],
- kernel_size: int = 3,
- squeeze_dims: int | tuple[int] | list[int] = None,
- modules_per_node: int = 5,
- batch_normalization: bool = True,
- activation: str = 'relu',
- padding: str = 'same',
- use_bias: bool = True,
- kernel_initializer: str = 'glorot_uniform',
- bias_initializer: str = 'zeros',
- kernel_regularizer: str = None,
- bias_regularizer: str = None,
- activity_regularizer: str = None,
- kernel_constraint: str = None,
- bias_constraint: str = None):
- """
- Builds a U-Net model.
-
- Parameters
- ----------
- input_shape: tuple
- Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
- num_classes: int
- Number of classes/labels that the U-Net will try to predict.
- pool_size: tuple or list
- Size of the mask in the MaxPooling layers.
- upsample_size: tuple or list
- Size of the mask in the UpSampling layers.
- levels: int
- Number of levels in the U-Net. Must be greater than 1.
- filter_num: iterable of ints
- Number of convolution filters on each level of the U-Net.
- kernel_size: int or tuple
- Size of the kernel in the convolution layers.
- squeeze_dims: int, tuple, or None
- Dimensions/axes of the input to squeeze such that the target (y_true) will be smaller than the input.
- - (e.g. to remove the third dimension, set this parameter to 2 [axis=2 for the third dimension])
- modules_per_node: int
- Number of modules in each node of the U-Net.
- batch_normalization: bool
- Setting this to True will add a batch normalization layer after every convolution in the modules.
- activation: str
- Activation function to use in the modules.
- Can be any of tf.keras.activations, 'gaussian', 'gcu', 'leaky_relu', 'prelu', 'smelu', 'snake' (case-insensitive).
- padding: str
- Padding to use in the convolution layers.
- use_bias: bool
- Setting this to True will implement a bias vector in the convolution layers used in the modules.
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
-
- Returns
- -------
- model: tf.keras.models.Model object
- U-Net model.
-
- Raises
- ------
- ValueError
- If levels < 2
- If input_shape does not have 3 nor 4 dimensions
- If the length of filter_num does not match the number of levels
-
- References
- ----------
- https://arxiv.org/pdf/1505.04597.pdf
- """
-
- ndims = len(input_shape) - 1 # Number of dimensions in the input image (excluding the last dimension reserved for channels)
-
- if levels < 2:
- raise ValueError(f"levels must be greater than 1. Received value: {levels}")
-
- if len(input_shape) > 4 or len(input_shape) < 3:
- raise ValueError(f"input_shape can only have 3 or 4 dimensions (2D image + 1 dimension for channels OR a 3D image + 1 dimension for channels). Received shape: {np.shape(input_shape)}")
-
- if len(filter_num) != levels:
- raise ValueError(f"length of filter_num ({len(filter_num)}) does not match the number of levels ({levels})")
-
- # Keyword arguments for the convolution modules
- module_kwargs = dict({})
- module_kwargs['num_modules'] = modules_per_node
- for arg in ['activation', 'batch_normalization', 'padding', 'kernel_size', 'use_bias', 'kernel_initializer', 'bias_initializer',
- 'kernel_regularizer', 'bias_regularizer', 'activity_regularizer', 'kernel_constraint', 'bias_constraint']:
- module_kwargs[arg] = locals()[arg]
-
- # MaxPooling keyword arguments
- pool_kwargs = {'pool_size': pool_size}
-
- # Keyword arguments for upsampling
- upsample_kwargs = dict({})
- for arg in ['activation', 'batch_normalization', 'padding', 'kernel_size', 'use_bias', 'kernel_initializer', 'bias_initializer',
- 'kernel_regularizer', 'bias_regularizer', 'activity_regularizer', 'kernel_constraint', 'bias_constraint',
- 'upsample_size']:
- upsample_kwargs[arg] = locals()[arg]
-
- # Keyword arguments for the deep supervision output in the final decoder node
- supervision_kwargs = dict({})
- for arg in ['padding', 'kernel_initializer', 'bias_initializer', 'kernel_regularizer', 'bias_regularizer', 'activity_regularizer',
- 'kernel_constraint', 'bias_constraint', 'upsample_size', 'squeeze_dims']:
- supervision_kwargs[arg] = locals()[arg]
- supervision_kwargs['use_bias'] = True
-
- tensors = dict({}) # Tensors associated with each node and skip connections
-
- """ Setup the first encoder node with an input layer and a convolution module """
- tensors['input'] = Input(shape=input_shape, name='Input')
- tensors['En1'] = unet_utils.convolution_module(tensors['input'], filters=filter_num[0], name='En1', **module_kwargs)
-
- """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
- for encoder in np.arange(2, levels+1): # Iterate through the rest of the encoder nodes
- current_node, previous_node = f'En{encoder}', f'En{encoder - 1}'
- pool_tensor = unet_utils.max_pool(tensors[previous_node], name=f'{previous_node}-{current_node}', **pool_kwargs) # Connect the next encoder node with a MaxPooling layer
- tensors[current_node] = unet_utils.convolution_module(pool_tensor, filters=filter_num[encoder - 1], name=current_node, **module_kwargs) # Convolution modules
-
- # Connect the bottom encoder node to a decoder node
- upsample_tensor = unet_utils.upsample(tensors[f'En{levels}'], filters=filter_num[levels - 2], name=f'En{levels}-De{levels}', **upsample_kwargs)
-
- """ Bottom decoder node """
- current_node, next_node = f'De{levels - 1}', f'De{levels - 2}'
- skip_node = f'En{levels - 1}' # node with an incoming skip connection that connects to 'current_node'
- tensors[current_node] = Concatenate(name=f'{current_node}_Concatenate')([tensors[skip_node], upsample_tensor]) # Concatenate the upsampled tensor and skip connection
- tensors[current_node] = unet_utils.convolution_module(tensors[current_node], filters=filter_num[levels - 2], name=current_node, **module_kwargs) # Convolution module
- upsample_tensor = unet_utils.upsample(tensors[current_node], filters=filter_num[levels - 3], name=f'{current_node}-{next_node}', **upsample_kwargs) # Connect the bottom decoder node to the next decoder node
-
- """ The rest of the decoder nodes (except the final node) are handled in this loop. Each node contains one concatenation of an upsampled tensor and a skip connection """
- for decoder in np.arange(2, levels-1)[::-1]:
- current_node, next_node = f'De{decoder}', f'De{decoder - 1}'
- skip_node = f'En{decoder}' # node with an incoming skip connection that connects to 'current_node'
- tensors[current_node] = Concatenate(name=f'{current_node}_Concatenate')([tensors[skip_node], upsample_tensor]) # Concatenate the upsampled tensor and skip connection
- tensors[current_node] = unet_utils.convolution_module(tensors[current_node], filters=filter_num[decoder - 1], name=current_node, **module_kwargs) # Convolution module
- upsample_tensor = unet_utils.upsample(tensors[current_node], filters=filter_num[decoder - 2], name=f'{current_node}-{next_node}', **upsample_kwargs) # Connect the bottom decoder node to the next decoder node
-
- """ Final decoder node begins with a concatenation and convolution module, followed by deep supervision """
- tensor_De1 = Concatenate(name='De1_Concatenate')([tensors['En1'], upsample_tensor]) # Concatenate the upsampled tensor and skip connection
- tensor_De1 = unet_utils.convolution_module(tensor_De1, filters=filter_num[0], name='De1', **module_kwargs) # Convolution module
- tensors['output'] = unet_utils.deep_supervision_side_output(tensor_De1, num_classes=num_classes, kernel_size=1, output_level=1, name='final', **supervision_kwargs) # Deep supervision - this layer will output the model's prediction
-
- model = Model(inputs=tensors['input'], outputs=tensors['output'], name=f'unet_{ndims}D')
-
- return model
-
-
-def unet_ensemble(
- input_shape: tuple[int] | list[int],
- num_classes: int,
- pool_size: int | tuple[int] | list[int],
- upsample_size: int | tuple[int] | list[int],
- levels: int,
- filter_num: tuple[int] | list[int],
- kernel_size: int = 3,
- squeeze_dims: int | tuple[int] | list[int] = None,
- modules_per_node: int = 5,
- batch_normalization: bool = True,
- activation: str = 'relu',
- padding: str = 'same',
- use_bias: bool = True,
- kernel_initializer: str = 'glorot_uniform',
- bias_initializer: str = 'zeros',
- kernel_regularizer: str = None,
- bias_regularizer: str = None,
- activity_regularizer: str = None,
- kernel_constraint: str = None,
- bias_constraint: str = None):
- """
- Builds a U-Net ensemble model.
- https://arxiv.org/pdf/1912.05074.pdf
-
- Parameters
- ----------
- input_shape: tuple
- Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
- num_classes: int
- Number of classes/labels that the U-Net will try to predict.
- pool_size: tuple or list
- Size of the mask in the MaxPooling layers.
- upsample_size: tuple or list
- Size of the mask in the UpSampling layers.
- levels: int
- Number of levels in the U-Net. Must be greater than 1.
- filter_num: iterable of ints
- Number of convolution filters on each level of the U-Net.
- kernel_size: int or tuple
- Size of the kernel in the convolution layers.
- squeeze_dims: int, tuple, or None
- Dimensions/axes of the input to squeeze such that the target (y_true) will be smaller than the input.
- - (e.g. to remove the third dimension, set this parameter to 2 [axis=2 for the third dimension])
- modules_per_node: int
- Number of modules in each node of the U-Net.
- batch_normalization: bool
- Setting this to True will add a batch normalization layer after every convolution in the modules.
- activation: str
- Activation function to use in the modules.
- Can be any of tf.keras.activations, 'gaussian', 'gcu', 'leaky_relu', 'prelu', 'smelu', 'snake' (case-insensitive).
- padding: str
- Padding to use in the convolution layers.
- use_bias: bool
- Setting this to True will implement a bias vector in the convolution layers used in the modules.
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
-
- Returns
- -------
- model: tf.keras.models.Model object
- U-Net model.
-
- Raises
- ------
- ValueError
- If levels < 2
- If input_shape does not have 3 nor 4 dimensions
- If the length of filter_num does not match the number of levels
- """
-
- ndims = len(input_shape) - 1 # Number of dimensions in the input image (excluding the last dimension reserved for channels)
-
- if levels < 2:
- raise ValueError(f"levels must be greater than 1. Received value: {levels}")
-
- if len(input_shape) > 4 or len(input_shape) < 3:
- raise ValueError(f"input_shape can only have 3 or 4 dimensions (2D image + 1 dimension for channels OR a 3D image + 1 dimension for channels). Received shape: {np.shape(input_shape)}")
-
- if len(filter_num) != levels:
- raise ValueError(f"length of filter_num ({len(filter_num)}) does not match the number of levels ({levels})")
-
- # Keyword arguments for the convolution modules
- module_kwargs = dict({})
- module_kwargs['num_modules'] = modules_per_node
- for arg in ['activation', 'batch_normalization', 'padding', 'kernel_size', 'use_bias', 'kernel_initializer', 'bias_initializer',
- 'kernel_regularizer', 'bias_regularizer', 'activity_regularizer', 'kernel_constraint', 'bias_constraint']:
- module_kwargs[arg] = locals()[arg]
-
- # MaxPooling keyword arguments
- pool_kwargs = {'pool_size': pool_size}
-
- # Keyword arguments for upsampling
- upsample_kwargs = dict({})
- for arg in ['activation', 'batch_normalization', 'padding', 'kernel_size', 'use_bias', 'kernel_initializer', 'bias_initializer',
- 'kernel_regularizer', 'bias_regularizer', 'activity_regularizer', 'kernel_constraint', 'bias_constraint',
- 'upsample_size']:
- upsample_kwargs[arg] = locals()[arg]
-
- # Keyword arguments for the deep supervision output in the final decoder node
- supervision_kwargs = dict({})
- for arg in ['padding', 'kernel_initializer', 'bias_initializer', 'kernel_regularizer', 'bias_regularizer', 'activity_regularizer',
- 'kernel_constraint', 'bias_constraint', 'upsample_size', 'squeeze_dims', 'num_classes']:
- supervision_kwargs[arg] = locals()[arg]
- supervision_kwargs['use_bias'] = True
- supervision_kwargs['output_level'] = 1
- supervision_kwargs['kernel_size'] = 1
-
- tensors = dict({}) # Tensors associated with each node and skip connections
- tensors_with_supervision = [] # list of output tensors. If deep supervision is used, more than one output will be produced
-
- """ Setup the first encoder node with an input layer and a convolution module """
- tensors['input'] = Input(shape=input_shape, name='Input')
- tensors['En1'] = unet_utils.convolution_module(tensors['input'], filters=filter_num[0], name='En1', **module_kwargs)
-
- """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
- for encoder in np.arange(2, levels+1): # Iterate through the rest of the encoder nodes
- current_node, previous_node = f'En{encoder}', f'En{encoder - 1}'
- pool_tensor = unet_utils.max_pool(tensors[previous_node], name=f'{previous_node}-{current_node}', **pool_kwargs) # Connect the next encoder node with a MaxPooling layer
- tensors[current_node] = unet_utils.convolution_module(pool_tensor, filters=filter_num[encoder - 1], name=current_node, **module_kwargs) # Convolution modules
-
- # Connect the bottom encoder node to a decoder node
- upsample_tensor = unet_utils.upsample(tensors[f'En{levels}'], filters=filter_num[levels - 2], name=f'En{levels}-De{levels}', **upsample_kwargs)
-
- """ Bottom decoder node """
- current_node, next_node = f'De{levels - 1}', f'De{levels - 2}'
- skip_node = f'En{levels - 1}'
- tensors[current_node] = Concatenate(name=f'{current_node}_Concatenate')([upsample_tensor, tensors[skip_node]]) # Concatenate the upsampled tensor and skip connection
- tensors[current_node] = unet_utils.convolution_module(tensors[current_node], filters=filter_num[levels - 2], name=current_node, **module_kwargs) # Convolution module
- upsample_tensor = unet_utils.upsample(tensors[current_node], filters=filter_num[levels - 3], name=f'{current_node}-{next_node}', **upsample_kwargs) # Connect the bottom decoder node to the next decoder node
-
- for decoder in np.arange(1, levels-1)[::-1]:
- num_middle_nodes = levels - decoder - 1
- for node in range(1, num_middle_nodes + 1):
- if node == 1: # if on the first middle node at the given level
- upsample_tensor_for_middle_node = unet_utils.upsample(tensors[f'En{decoder + 1}'], filters=filter_num[decoder - 2], name=f'En{decoder + 1}-Me{decoder}-1', **upsample_kwargs)
- else:
- upsample_tensor_for_middle_node = unet_utils.upsample(tensors[f'Me{decoder + 1}-{node - 1}'], filters=filter_num[decoder - 2], name=f'Me{decoder + 1}-{node - 1}-Me{decoder}-{node}', **upsample_kwargs)
- tensors[f'Me{decoder}-{node}'] = Concatenate(name=f'Me{decoder}-{node}_Concatenate')([tensors[f'En{decoder}'], upsample_tensor_for_middle_node])
- tensors[f'Me{decoder}-{node}'] = unet_utils.convolution_module(tensors[f'Me{decoder}-{node}'], filters=filter_num[decoder - 1], name=f'Me{decoder}-{node}', **module_kwargs) # Convolution module
- if decoder == 1:
- tensors[f'sup{decoder}-{node}'] = unet_utils.deep_supervision_side_output(tensors[f'Me{decoder}-{node}'], name=f'sup{decoder}-{node}', **supervision_kwargs) # deep supervision on middle node located on top level
- tensors_with_supervision.append(tensors[f'sup{decoder}-{node}'])
- tensors[f'De{decoder}'] = Concatenate(name=f'De{decoder}_Concatenate')([tensors[f'En{decoder}'], upsample_tensor]) # Concatenate the upsampled tensor and skip connection
- tensors[f'De{decoder}'] = unet_utils.convolution_module(tensors[f'De{decoder}'], filters=filter_num[decoder - 1], name=f'De{decoder}', **module_kwargs) # Convolution module
-
- if decoder != 1: # if not currently on the final decoder node (De1)
- upsample_tensor = unet_utils.upsample(tensors[f'De{decoder}'], filters=filter_num[decoder - 2], name=f'De{decoder}-De{decoder - 1}', **upsample_kwargs) # Connect the bottom decoder node to the next decoder node
- else:
- tensors['output'] = unet_utils.deep_supervision_side_output(tensors['De1'], name='final', **supervision_kwargs) # Deep supervision - this layer will output the model's prediction
- tensors_with_supervision.append(tensors['output'])
-
- model = Model(inputs=tensors['input'], outputs=tensors_with_supervision, name=f'unet_ensemble_{ndims}D')
-
- return model
-
-
-def unet_plus(
- input_shape: tuple[int] | list[int],
- num_classes: int,
- pool_size: int | tuple[int] | list[int],
- upsample_size: int | tuple[int] | list[int],
- levels: int,
- filter_num: tuple[int] | list[int],
- kernel_size: int = 3,
- squeeze_dims: int | tuple[int] | list[int] = None,
- modules_per_node: int = 5,
- batch_normalization: bool = True,
- deep_supervision: bool = True,
- activation: str = 'relu',
- padding: str = 'same',
- use_bias: bool = True,
- kernel_initializer: str = 'glorot_uniform',
- bias_initializer: str = 'zeros',
- kernel_regularizer: str = None,
- bias_regularizer: str = None,
- activity_regularizer: str = None,
- kernel_constraint: str = None,
- bias_constraint: str = None):
- """
- Builds a U-Net+ model.
- https://arxiv.org/pdf/1912.05074.pdf
-
- Parameters
- ----------
- input_shape: tuple
- Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
- num_classes: int
- Number of classes/labels that the U-Net will try to predict.
- pool_size: tuple or list
- Size of the mask in the MaxPooling layers.
- upsample_size: tuple or list
- Size of the mask in the UpSampling layers.
- levels: int
- Number of levels in the U-Net. Must be greater than 1.
- filter_num: iterable of ints
- Number of convolution filters on each level of the U-Net.
- kernel_size: int or tuple
- Size of the kernel in the convolution layers.
- squeeze_dims: int, tuple, or None
- Dimensions/axes of the input to squeeze such that the target (y_true) will be smaller than the input.
- - (e.g. to remove the third dimension, set this parameter to 2 [axis=2 for the third dimension])
- modules_per_node: int
- Number of modules in each node of the U-Net.
- batch_normalization: bool
- Setting this to True will add a batch normalization layer after every convolution in the modules.
- deep_supervision: bool
- Add deep supervision side outputs to each top node.
- NOTE: The final decoder node requires deep supervision and is not affected if this parameter is False.
- activation: str
- Activation function to use in the modules.
- Can be any of tf.keras.activations, 'gaussian', 'gcu', 'leaky_relu', 'prelu', 'smelu', 'snake' (case-insensitive).
- padding: str
- Padding to use in the convolution layers.
- use_bias: bool
- Setting this to True will implement a bias vector in the convolution layers used in the modules.
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
-
- Returns
- -------
- model: tf.keras.models.Model object
- U-Net model.
-
- Raises
- ------
- ValueError
- If levels < 2
- If input_shape does not have 3 nor 4 dimensions
- If the length of filter_num does not match the number of levels
- """
-
- ndims = len(input_shape) - 1 # Number of dimensions in the input image (excluding the last dimension reserved for channels)
-
- if levels < 2:
- raise ValueError(f"levels must be greater than 1. Received value: {levels}")
-
- if len(input_shape) > 4 or len(input_shape) < 3:
- raise ValueError(f"input_shape can only have 3 or 4 dimensions (2D image + 1 dimension for channels OR a 3D image + 1 dimension for channels). Received shape: {np.shape(input_shape)}")
-
- if len(filter_num) != levels:
- raise ValueError(f"length of filter_num ({len(filter_num)}) does not match the number of levels ({levels})")
-
- # Keyword arguments for the convolution modules
- module_kwargs = dict({})
- module_kwargs['num_modules'] = modules_per_node
- for arg in ['activation', 'batch_normalization', 'padding', 'kernel_size', 'use_bias', 'kernel_initializer', 'bias_initializer',
- 'kernel_regularizer', 'bias_regularizer', 'activity_regularizer', 'kernel_constraint', 'bias_constraint']:
- module_kwargs[arg] = locals()[arg]
-
- # MaxPooling keyword arguments
- pool_kwargs = {'pool_size': pool_size}
-
- # Keyword arguments for upsampling
- upsample_kwargs = dict({})
- for arg in ['activation', 'batch_normalization', 'padding', 'kernel_size', 'use_bias', 'kernel_initializer', 'bias_initializer',
- 'kernel_regularizer', 'bias_regularizer', 'activity_regularizer', 'kernel_constraint', 'bias_constraint',
- 'upsample_size']:
- upsample_kwargs[arg] = locals()[arg]
-
- # Keyword arguments for the deep supervision output in the final decoder node
- supervision_kwargs = dict({})
- for arg in ['padding', 'kernel_initializer', 'bias_initializer', 'kernel_regularizer', 'bias_regularizer', 'activity_regularizer',
- 'kernel_constraint', 'bias_constraint', 'upsample_size', 'squeeze_dims', 'num_classes']:
- supervision_kwargs[arg] = locals()[arg]
- supervision_kwargs['use_bias'] = True
- supervision_kwargs['output_level'] = 1
- supervision_kwargs['kernel_size'] = 1
-
- tensors = dict({}) # Tensors associated with each node and skip connections
- tensors_with_supervision = [] # list of output tensors. If deep supervision is used, more than one output will be produced
-
- """ Setup the first encoder node with an input layer and a convolution module """
- tensors['input'] = Input(shape=input_shape, name='Input')
- tensors['En1'] = unet_utils.convolution_module(tensors['input'], filters=filter_num[0], name='En1', **module_kwargs)
-
- """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
- for encoder in np.arange(2, levels+1): # Iterate through the rest of the encoder nodes
- pool_tensor = unet_utils.max_pool(tensors[f'En{encoder - 1}'], name=f'En{encoder - 1}-En{encoder}', **pool_kwargs) # Connect the next encoder node with a MaxPooling layer
- tensors[f'En{encoder}'] = unet_utils.convolution_module(pool_tensor, filters=filter_num[encoder - 1], name=f'En{encoder}', **module_kwargs) # Convolution modules
-
- # Connect the bottom encoder node to a decoder node
- upsample_tensor = unet_utils.upsample(tensors[f'En{levels}'], filters=filter_num[levels - 2], name=f'En{levels}-De{levels}', **upsample_kwargs)
-
- """ Bottom decoder node """
- tensors[f'De{levels - 1}'] = Concatenate(name=f'De{levels - 1}_Concatenate')([upsample_tensor, tensors[f'En{levels - 1}']]) # Concatenate the upsampled tensor and skip connection
- tensors[f'De{levels - 1}'] = unet_utils.convolution_module(tensors[f'De{levels - 1}'], filters=filter_num[levels - 2], name=f'De{levels - 1}', **module_kwargs) # Convolution module
- upsample_tensor = unet_utils.upsample(tensors[f'De{levels - 1}'], filters=filter_num[levels - 3], name=f'De{levels - 1}-De{levels - 2}', **upsample_kwargs) # Connect the bottom decoder node to the next decoder node
-
- """ The rest of the decoder nodes (except the final node) are handled in this loop. Each node contains one concatenation of an upsampled tensor and a skip connection """
- for decoder in np.arange(1, levels-1)[::-1]:
- num_middle_nodes = levels - decoder - 1
- for node in range(1, num_middle_nodes + 1):
- if node == 1: # if on the first middle node at the given level
- upsample_tensor_for_middle_node = unet_utils.upsample(tensors[f'En{decoder + 1}'], filters=filter_num[decoder - 2], name=f'En{decoder + 1}-Me{decoder}-1', **upsample_kwargs)
- tensors[f'Me{decoder}-1'] = Concatenate(name=f'Me{decoder}-1_Concatenate')([tensors[f'En{decoder}'], upsample_tensor_for_middle_node])
- else:
- upsample_tensor_for_middle_node = unet_utils.upsample(tensors[f'Me{decoder + 1}-{node - 1}'], filters=filter_num[decoder - 2], name=f'Me{decoder + 1}-{node - 1}-Me{decoder}-{node}', **upsample_kwargs)
- tensors[f'Me{decoder}-{node}'] = Concatenate(name=f'Me{decoder}-{node}_Concatenate')([tensors[f'Me{decoder}-{node - 1}'], upsample_tensor_for_middle_node])
- tensors[f'Me{decoder}-{node}'] = unet_utils.convolution_module(tensors[f'Me{decoder}-{node}'], filters=filter_num[decoder - 1], name=f'Me{decoder}-{node}', **module_kwargs) # Convolution module
- if decoder == 1 and deep_supervision:
- tensors[f'sup{decoder}-{node}'] = unet_utils.deep_supervision_side_output(tensors[f'Me{decoder}-{node}'], name=f'sup{decoder}-{node}', **supervision_kwargs) # deep supervision on middle node located on top level
- tensors_with_supervision.append(tensors[f'sup{decoder}-{node}'])
- tensors[f'De{decoder}'] = Concatenate(name=f'De{decoder}_Concatenate')([tensors[f'Me{decoder}-{num_middle_nodes}'], upsample_tensor]) # Concatenate the upsampled tensor and skip connection
- tensors[f'De{decoder}'] = unet_utils.convolution_module(tensors[f'De{decoder}'], filters=filter_num[decoder - 1], name=f'De{decoder}', **module_kwargs) # Convolution module
-
- if decoder != 1: # if not currently on the final decoder node (De1)
- upsample_tensor = unet_utils.upsample(tensors[f'De{decoder}'], filters=filter_num[decoder - 2], name=f'De{decoder}-De{decoder - 1}', **upsample_kwargs) # Connect the bottom decoder node to the next decoder node
- else:
- tensors['output'] = unet_utils.deep_supervision_side_output(tensors['De1'], **supervision_kwargs) # Deep supervision - this layer will output the model's prediction
- tensors_with_supervision.append(tensors['output'])
-
- model = Model(inputs=tensors['input'], outputs=tensors_with_supervision, name=f'unet_plus_{ndims}D')
-
- return model
-
-
-def unet_2plus(
- input_shape: tuple[int] | list[int],
- num_classes: int,
- pool_size: int | tuple[int] | list[int],
- upsample_size: int | tuple[int] | list[int],
- levels: int,
- filter_num: tuple[int] | list[int],
- kernel_size: int = 3,
- squeeze_dims: int | tuple[int] | list[int] = None,
- modules_per_node: int = 5,
- batch_normalization: bool = True,
- deep_supervision: bool = True,
- activation: str = 'relu',
- padding: str = 'same',
- use_bias: bool = True,
- kernel_initializer: str = 'glorot_uniform',
- bias_initializer: str = 'zeros',
- kernel_regularizer: str = None,
- bias_regularizer: str = None,
- activity_regularizer: str = None,
- kernel_constraint: str = None,
- bias_constraint: str = None):
- """
- Builds a U-Net++ model.
- https://arxiv.org/pdf/1912.05074.pdf
-
- Parameters
- ----------
- input_shape: tuple
- Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
- num_classes: int
- Number of classes/labels that the U-Net will try to predict.
- pool_size: tuple or list
- Size of the mask in the MaxPooling layers.
- upsample_size: tuple or list
- Size of the mask in the UpSampling layers.
- levels: int
- Number of levels in the U-Net. Must be greater than 1.
- filter_num: iterable of ints
- Number of convolution filters on each level of the U-Net.
- kernel_size: int or tuple
- Size of the kernel in the convolution layers.
- squeeze_dims: int, tuple, or None
- Dimensions/axes of the input to squeeze such that the target (y_true) will be smaller than the input.
- - (e.g. to remove the third dimension, set this parameter to 2 [axis=2 for the third dimension])
- modules_per_node: int
- Number of modules in each node of the U-Net.
- batch_normalization: bool
- Setting this to True will add a batch normalization layer after every convolution in the modules.
- deep_supervision: bool
- Add deep supervision side outputs to each top node.
- NOTE: The final decoder node requires deep supervision and is not affected if this parameter is False.
- activation: str
- Activation function to use in the modules.
- Can be any of tf.keras.activations, 'gaussian', 'gcu', 'leaky_relu', 'prelu', 'smelu', 'snake' (case-insensitive).
- padding: str
- Padding to use in the convolution layers.
- use_bias: bool
- Setting this to True will implement a bias vector in the convolution layers used in the modules.
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
-
- Returns
- -------
- model: tf.keras.models.Model object
- U-Net model.
-
- Raises
- ------
- ValueError
- If levels < 2
- If input_shape does not have 3 nor 4 dimensions
- If the length of filter_num does not match the number of levels
- """
-
- ndims = len(input_shape) - 1 # Number of dimensions in the input image (excluding the last dimension reserved for channels)
-
- if levels < 2:
- raise ValueError(f"levels must be greater than 1. Received value: {levels}")
-
- if len(input_shape) > 4 or len(input_shape) < 3:
- raise ValueError(f"input_shape can only have 3 or 4 dimensions (2D image + 1 dimension for channels OR a 3D image + 1 dimension for channels). Received shape: {np.shape(input_shape)}")
-
- if len(filter_num) != levels:
- raise ValueError(f"length of filter_num ({len(filter_num)}) does not match the number of levels ({levels})")
-
- # Keyword arguments for the convolution modules
- module_kwargs = dict({})
- module_kwargs['activation'] = activation
- module_kwargs['batch_normalization'] = batch_normalization
- module_kwargs['num_modules'] = modules_per_node
- module_kwargs['padding'] = padding
- module_kwargs['use_bias'] = use_bias
- module_kwargs['kernel_initializer'] = kernel_initializer
- module_kwargs['bias_initializer'] = bias_initializer
- module_kwargs['kernel_regularizer'] = kernel_regularizer
- module_kwargs['bias_regularizer'] = bias_regularizer
- module_kwargs['activity_regularizer'] = activity_regularizer
- module_kwargs['kernel_constraint'] = kernel_constraint
- module_kwargs['bias_constraint'] = bias_constraint
-
- # MaxPooling keyword arguments
- pool_kwargs = dict({})
- pool_kwargs['pool_size'] = pool_size
-
- # Keyword arguments for upsampling
- upsample_kwargs = dict({})
- upsample_kwargs['activation'] = activation
- upsample_kwargs['batch_normalization'] = batch_normalization
- upsample_kwargs['padding'] = padding
- upsample_kwargs['kernel_initializer'] = kernel_initializer
- upsample_kwargs['bias_initializer'] = bias_initializer
- upsample_kwargs['kernel_regularizer'] = kernel_regularizer
- upsample_kwargs['bias_regularizer'] = bias_regularizer
- upsample_kwargs['activity_regularizer'] = activity_regularizer
- upsample_kwargs['kernel_constraint'] = kernel_constraint
- upsample_kwargs['bias_constraint'] = bias_constraint
- upsample_kwargs['upsample_size'] = upsample_size
- upsample_kwargs['use_bias'] = use_bias
-
- # Keyword arguments for the deep supervision output in the final decoder node
- supervision_kwargs = dict({})
- supervision_kwargs['upsample_size'] = upsample_size
- supervision_kwargs['use_bias'] = True
- supervision_kwargs['squeeze_dims'] = squeeze_dims
- supervision_kwargs['padding'] = padding
- supervision_kwargs['kernel_initializer'] = kernel_initializer
- supervision_kwargs['bias_initializer'] = bias_initializer
- supervision_kwargs['kernel_regularizer'] = kernel_regularizer
- supervision_kwargs['bias_regularizer'] = bias_regularizer
- supervision_kwargs['activity_regularizer'] = activity_regularizer
- supervision_kwargs['kernel_constraint'] = kernel_constraint
- supervision_kwargs['bias_constraint'] = bias_constraint
-
- tensors = dict({}) # Tensors associated with each node and skip connections
- tensors_with_supervision = [] # list of output tensors. If deep supervision is used, more than one output will be produced
-
- """ Setup the first encoder node with an input layer and a convolution module """
- tensors['input'] = Input(shape=input_shape, name='Input')
- tensors['En1'] = unet_utils.convolution_module(tensors['input'], filters=filter_num[0], kernel_size=kernel_size, name='En1', **module_kwargs)
-
- """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
- for encoder in np.arange(2, levels+1): # Iterate through the rest of the encoder nodes
- pool_tensor = unet_utils.max_pool(tensors[f'En{encoder - 1}'], name=f'En{encoder - 1}-En{encoder}', **pool_kwargs) # Connect the next encoder node with a MaxPooling layer
- tensors[f'En{encoder}'] = unet_utils.convolution_module(pool_tensor, filters=filter_num[encoder - 1], kernel_size=kernel_size, name=f'En{encoder}', **module_kwargs) # Convolution modules
-
- # Connect the bottom encoder node to a decoder node
- upsample_tensor = unet_utils.upsample(tensors[f'En{levels}'], filters=filter_num[levels - 2], kernel_size=kernel_size, name=f'En{levels}-De{levels}', **upsample_kwargs)
-
- """ Bottom decoder node """
- tensors[f'De{levels - 1}'] = Concatenate(name=f'De{levels - 1}_Concatenate')([upsample_tensor, tensors[f'En{levels - 1}']]) # Concatenate the upsampled tensor and skip connection
- tensors[f'De{levels - 1}'] = unet_utils.convolution_module(tensors[f'De{levels - 1}'], filters=filter_num[levels - 2], kernel_size=kernel_size, name=f'De{levels - 1}', **module_kwargs) # Convolution module
- upsample_tensor = unet_utils.upsample(tensors[f'De{levels - 1}'], filters=filter_num[levels - 3], kernel_size=kernel_size, name=f'De{levels - 1}-De{levels - 2}', **upsample_kwargs) # Connect the bottom decoder node to the next decoder node
-
- """ The rest of the decoder nodes (except the final node) are handled in this loop. Each node contains one concatenation of an upsampled tensor and a skip connection """
- for decoder in np.arange(1, levels-1)[::-1]:
- num_middle_nodes = levels - decoder - 1
- for node in range(1, num_middle_nodes + 1):
- if node == 1: # if on the first middle node at the given level
- upsample_tensor_for_middle_node = unet_utils.upsample(tensors[f'En{decoder + 1}'], filters=filter_num[decoder - 2], kernel_size=kernel_size, name=f'En{decoder + 1}-Me{decoder}-1', **upsample_kwargs)
- tensors[f'Me{decoder}-1'] = Concatenate(name=f'Me{decoder}-1_Concatenate')([tensors[f'En{decoder}'], upsample_tensor_for_middle_node])
- else:
- upsample_tensor_for_middle_node = unet_utils.upsample(tensors[f'Me{decoder + 1}-{node - 1}'], filters=filter_num[decoder - 2], kernel_size=kernel_size, name=f'Me{decoder + 1}-{node - 1}-Me{decoder}-{node}', **upsample_kwargs)
- tensors_to_concatenate = [] # Tensors to concatenate in the middle node
- connections_to_add = sorted([tensor for tensor in tensors if f'Me{decoder}' in tensor])[::-1] # skip connections to add to the list of tensors to concatenate
- for connection in connections_to_add:
- tensors_to_concatenate.append(tensors[connection])
- tensors_to_concatenate.append(tensors[f'En{decoder}'])
- tensors_to_concatenate.append(upsample_tensor_for_middle_node)
- tensors[f'Me{decoder}-{node}'] = Concatenate(name=f'Me{decoder}-{node}_Concatenate')(tensors_to_concatenate)
- tensors[f'Me{decoder}-{node}'] = unet_utils.convolution_module(tensors[f'Me{decoder}-{node}'], filters=filter_num[decoder - 1], kernel_size=kernel_size, name=f'Me{decoder}-{node}', **module_kwargs) # Convolution module
-
- if decoder == 1 and deep_supervision:
- tensors[f'sup{decoder}-{node}'] = unet_utils.deep_supervision_side_output(tensors[f'Me{decoder}-{node}'], num_classes=num_classes, output_level=1, kernel_size=1, name=f'sup{decoder}-{node}', **supervision_kwargs) # deep supervision on middle node located on top level
- tensors_with_supervision.append(tensors[f'sup{decoder}-{node}'])
-
- tensors_to_concatenate = [] # tensors to concatenate in the decoder node
- connections_to_add = sorted([tensor for tensor in tensors if f'Me{decoder}' in tensor])[::-1] # skip connections to add to the list of tensors to concatenate
- for connection in connections_to_add:
- tensors_to_concatenate.append(tensors[connection])
- tensors_to_concatenate.append(tensors[f'En{decoder}'])
- tensors_to_concatenate.append(upsample_tensor)
- tensors[f'De{decoder}'] = Concatenate(name=f'De{decoder}_Concatenate')(tensors_to_concatenate) # Concatenate the upsampled tensor and skip connection
- tensors[f'De{decoder}'] = unet_utils.convolution_module(tensors[f'De{decoder}'], filters=filter_num[decoder - 1], kernel_size=kernel_size, name=f'De{decoder}', **module_kwargs) # Convolution module
-
- if decoder != 1: # if not currently on the final decoder node (De1)
- upsample_tensor = unet_utils.upsample(tensors[f'De{decoder}'], filters=filter_num[decoder - 2], kernel_size=kernel_size, name=f'De{decoder}-De{decoder - 1}', **upsample_kwargs) # Connect the bottom decoder node to the next decoder node
- else:
- tensors['output'] = unet_utils.deep_supervision_side_output(tensors['De1'], num_classes=num_classes, kernel_size=1, output_level=1, name='final', **supervision_kwargs) # Deep supervision - this layer will output the model's prediction
- tensors_with_supervision.append(tensors['output'])
-
- model = Model(inputs=tensors['input'], outputs=tensors_with_supervision, name=f'unet_2plus_{ndims}D')
-
- return model
-
-
-def unet_3plus(
- input_shape: tuple[int] | list[int],
- num_classes: int,
- pool_size: int | tuple[int] | list[int],
- upsample_size: int | tuple[int] | list[int],
- levels: int,
- filter_num: tuple[int] | list[int],
- filter_num_skip: int = None,
- filter_num_aggregate: tuple[int] | list[int] = None,
- kernel_size: int = 3,
- first_encoder_connections: bool = True,
- squeeze_dims: int | tuple[int] | list[int] = None,
- modules_per_node: int = 5,
- batch_normalization: bool = True,
- deep_supervision: bool = True,
- activation: str = 'relu',
- padding: str = 'same',
- use_bias: bool = True,
- kernel_initializer: str = 'glorot_uniform',
- bias_initializer: str = 'zeros',
- kernel_regularizer: str = None,
- bias_regularizer: str = None,
- activity_regularizer: str = None,
- kernel_constraint: str = None,
- bias_constraint: str = None):
- """
- Creates a U-Net 3+.
- https://arxiv.org/ftp/arxiv/papers/2004/2004.08790.pdf
-
- Parameters
- ----------
- input_shape: tuple
- Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
- num_classes: int
- Number of classes/labels that the U-Net 3+ will try to predict.
- pool_size: tuple or list
- Size of the mask in the MaxPooling layers.
- upsample_size: tuple or list
- Size of the mask in the UpSampling layers.
- levels: int
- Number of levels in the U-Net 3+. Must be greater than 2.
- filter_num: iterable of ints
- Number of convolution filters in each encoder of the U-Net 3+. The length must be equal to 'levels'.
- filter_num_skip: int or None
- Number of convolution filters in the conventional skip connections, full-scale skip connections, and aggregated feature maps.
- NOTE: When left as None, this will default to the first value in the 'filter_num' iterable.
- filter_num_aggregate: int or None
- Number of convolution filters in the decoder nodes after images are concatenated.
- When left as None, this will be equal to the product of filter_num_skip and the number of levels.
- kernel_size: int or tuple
- Size of the kernel in the convolution layers.
- first_encoder_connections: bool
- Setting this to True will create full-scale skip connections attached to the first encoder node.
- squeeze_dims: int, tuple, or None
- Dimensions/axes of the input to squeeze such that the target (y_true) will be smaller than the input.
- - (e.g. to remove the third dimension, set this parameter to 2 [axis=2 for the third dimension])
- modules_per_node: int
- Number of modules in each node of the U-Net 3+.
- batch_normalization: bool
- Setting this to True will add a batch normalization layer after every convolution in the modules.
- deep_supervision: bool
- Add deep supervision side outputs to each decoder node.
- NOTE: The final decoder node requires deep supervision and is not affected if this parameter is False.
- activation: str
- Activation function to use in the modules.
- Can be any of tf.keras.activations, 'gaussian', 'gcu', 'leaky_relu', 'prelu', 'smelu', 'snake' (case-insensitive).
- padding: str
- Padding to use in the convolution layers.
- use_bias: bool
- Setting this to True will implement a bias vector in the convolution layers used in the modules.
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
-
- Returns
- -------
- model: tf.keras.models.Model object
- U-Net 3+ model.
- """
-
- ndims = len(input_shape) - 1 # Number of dimensions in the input image (excluding the last dimension reserved for channels)
-
- if levels < 3:
- raise ValueError(f"levels must be greater than 2. Received value: {levels}")
-
- if len(input_shape) > 4 or len(input_shape) < 3:
- raise ValueError(f"input_shape can only have 3 or 4 dimensions (2D image + 1 dimension for channels OR a 3D image + 1 dimension for channels). Received shape: {np.shape(input_shape)}")
-
- if len(filter_num) != levels:
- raise ValueError(f"length of filter_num ({len(filter_num)}) does not match the number of levels ({levels})")
-
- if filter_num_skip is None:
- filter_num_skip = filter_num[0]
-
- if filter_num_aggregate is None:
- filter_num_aggregate = levels * filter_num_skip
-
- # print(f"\nCreating model: {ndims}D U-Net 3+")
-
- module_kwargs = dict({})
- module_kwargs['kernel_size'] = kernel_size
- module_kwargs['activation'] = activation
- module_kwargs['batch_normalization'] = batch_normalization
- module_kwargs['num_modules'] = modules_per_node
- module_kwargs['padding'] = padding
- module_kwargs['use_bias'] = use_bias
- module_kwargs['kernel_initializer'] = kernel_initializer
- module_kwargs['bias_initializer'] = bias_initializer
- module_kwargs['kernel_regularizer'] = kernel_regularizer
- module_kwargs['bias_regularizer'] = bias_regularizer
- module_kwargs['activity_regularizer'] = activity_regularizer
- module_kwargs['kernel_constraint'] = kernel_constraint
- module_kwargs['bias_constraint'] = bias_constraint
-
- pool_kwargs = dict({})
- pool_kwargs['pool_size'] = pool_size
-
- upsample_kwargs = dict({})
- upsample_kwargs['activation'] = activation
- upsample_kwargs['batch_normalization'] = batch_normalization
- upsample_kwargs['kernel_size'] = kernel_size
- upsample_kwargs['filters'] = filter_num_skip
- upsample_kwargs['padding'] = padding
- upsample_kwargs['kernel_initializer'] = kernel_initializer
- upsample_kwargs['bias_initializer'] = bias_initializer
- upsample_kwargs['kernel_regularizer'] = kernel_regularizer
- upsample_kwargs['bias_regularizer'] = bias_regularizer
- upsample_kwargs['activity_regularizer'] = activity_regularizer
- upsample_kwargs['kernel_constraint'] = kernel_constraint
- upsample_kwargs['bias_constraint'] = bias_constraint
- upsample_kwargs['upsample_size'] = upsample_size
- upsample_kwargs['use_bias'] = use_bias
-
- conventional_kwargs = dict({})
- conventional_kwargs['filters'] = filter_num_skip
- conventional_kwargs['kernel_size'] = kernel_size
- conventional_kwargs['activation'] = activation
- conventional_kwargs['batch_normalization'] = batch_normalization
- conventional_kwargs['padding'] = padding
- conventional_kwargs['use_bias'] = use_bias
- conventional_kwargs['kernel_initializer'] = kernel_initializer
- conventional_kwargs['bias_initializer'] = bias_initializer
- conventional_kwargs['kernel_regularizer'] = kernel_regularizer
- conventional_kwargs['bias_regularizer'] = bias_regularizer
- conventional_kwargs['activity_regularizer'] = activity_regularizer
- conventional_kwargs['kernel_constraint'] = kernel_constraint
- conventional_kwargs['bias_constraint'] = bias_constraint
-
- full_scale_kwargs = dict({})
- full_scale_kwargs['filters'] = filter_num_skip
- full_scale_kwargs['kernel_size'] = kernel_size
- full_scale_kwargs['activation'] = activation
- full_scale_kwargs['batch_normalization'] = batch_normalization
- full_scale_kwargs['use_bias'] = use_bias
- full_scale_kwargs['padding'] = padding
- full_scale_kwargs['pool_size'] = pool_size
- full_scale_kwargs['kernel_initializer'] = kernel_initializer
- full_scale_kwargs['bias_initializer'] = bias_initializer
- full_scale_kwargs['kernel_regularizer'] = kernel_regularizer
- full_scale_kwargs['bias_regularizer'] = bias_regularizer
- full_scale_kwargs['activity_regularizer'] = activity_regularizer
- full_scale_kwargs['kernel_constraint'] = kernel_constraint
- full_scale_kwargs['bias_constraint'] = bias_constraint
-
- aggregated_kwargs = dict({})
- aggregated_kwargs['filters'] = filter_num_skip
- aggregated_kwargs['kernel_size'] = kernel_size
- aggregated_kwargs['activation'] = activation
- aggregated_kwargs['batch_normalization'] = batch_normalization
- aggregated_kwargs['padding'] = padding
- aggregated_kwargs['upsample_size'] = upsample_size
- aggregated_kwargs['use_bias'] = use_bias
- aggregated_kwargs['kernel_initializer'] = kernel_initializer
- aggregated_kwargs['bias_initializer'] = bias_initializer
- aggregated_kwargs['kernel_regularizer'] = kernel_regularizer
- aggregated_kwargs['bias_regularizer'] = bias_regularizer
- aggregated_kwargs['activity_regularizer'] = activity_regularizer
- aggregated_kwargs['kernel_constraint'] = kernel_constraint
- aggregated_kwargs['bias_constraint'] = bias_constraint
-
- supervision_kwargs = dict({})
- supervision_kwargs['upsample_size'] = upsample_size
- supervision_kwargs['kernel_size'] = kernel_size
- supervision_kwargs['use_bias'] = True
- supervision_kwargs['padding'] = padding
- supervision_kwargs['squeeze_dims'] = squeeze_dims
- supervision_kwargs['kernel_initializer'] = kernel_initializer
- supervision_kwargs['bias_initializer'] = bias_initializer
- supervision_kwargs['kernel_regularizer'] = kernel_regularizer
- supervision_kwargs['bias_regularizer'] = bias_regularizer
- supervision_kwargs['activity_regularizer'] = activity_regularizer
- supervision_kwargs['kernel_constraint'] = kernel_constraint
- supervision_kwargs['bias_constraint'] = bias_constraint
-
- tensors = dict({}) # Tensors associated with each node and skip connections
- tensors_with_supervision = [] # Outputs of deep supervision
-
- """ Setup the first encoder node with an input layer and a convolution module (we are not using skip connections here) """
- tensors['input'] = Input(shape=input_shape, name='Input')
- tensors['En1'] = unet_utils.convolution_module(tensors['input'], filters=filter_num[0], name='En1', **module_kwargs)
-
- if first_encoder_connections is True:
- for full_connection in range(2, levels):
- tensors[f'1---{full_connection}_full-scale'] = unet_utils.full_scale_skip_connection(tensors[f'En1'], level1=1, level2=full_connection, name=f'1---{full_connection}_full-scale', **full_scale_kwargs)
-
- """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
- for encoder in np.arange(2, levels): # Iterate through the rest of the encoder nodes
- pool_tensor = unet_utils.max_pool(tensors[f'En{encoder - 1}'], name=f'En{encoder - 1}-En{encoder}', **pool_kwargs) # Connect the next encoder node with a MaxPooling layer
- tensors[f'En{encoder}'] = unet_utils.convolution_module(pool_tensor, filters=filter_num[encoder - 1], name=f'En{encoder}', **module_kwargs) # Convolution modules
- tensors[f'{encoder}---{encoder}_skip'] = unet_utils.conventional_skip_connection(tensors[f'En{encoder}'], name=f'{encoder}---{encoder}_skip', **conventional_kwargs)
-
- # Create full-scale skip connections
- for full_connection in range(encoder + 1, levels):
- tensors[f'{encoder}---{full_connection}_full-scale'] = unet_utils.full_scale_skip_connection(tensors[f'En{encoder}'], level1=encoder, level2=full_connection, name=f'{encoder}---{full_connection}_full-scale', **full_scale_kwargs)
-
- # Bottom encoder node
- tensors[f'En{levels}'] = unet_utils.max_pool(tensors[f'En{levels - 1}'], name=f'En{levels - 1}-En{levels}', **pool_kwargs)
- tensors[f'En{levels}'] = unet_utils.convolution_module(tensors[f'En{levels}'], filters=filter_num[levels - 1], name=f'En{levels}', **module_kwargs)
- if deep_supervision:
- tensors[f'sup{levels}_output'] = unet_utils.deep_supervision_side_output(tensors[f'En{levels}'], num_classes=num_classes, output_level=levels, name=f'sup{levels}', **supervision_kwargs)
- tensors_with_supervision.append(tensors[f'sup{levels}_output'])
-
- # Add aggregated feature maps using the bottom encoder node
- for feature_map in range(1, levels - 1):
- tensors[f'{levels}---{feature_map}_feature'] = unet_utils.aggregated_feature_map(tensors[f'En{levels}'], level1=levels, level2=feature_map, name=f'{levels}---{feature_map}_feature', **aggregated_kwargs)
-
- """ Build the rest of the decoder nodes """
- for decoder in np.arange(1, levels)[::-1]:
-
- """ The lowest decoder node (levels - 1) is attached to the bottom encoder node via upsampling, so concatenation is slightly different """
- if decoder == levels - 1:
- tensors[f'De{decoder}'] = unet_utils.upsample(tensors[f'En{levels}'], name=f'En{levels}-De{decoder}', **upsample_kwargs)
-
- # Tensors to concatenate in the Concatenate layer
- tensors_to_concatenate = [tensors[f'De{decoder}'], ]
- connections_to_add = sorted([tensor for tensor in tensors if f'---{decoder}' in tensor])[::-1]
- for connection in connections_to_add:
- tensors_to_concatenate.append(tensors[connection])
- else:
- tensors[f'De{decoder}'] = unet_utils.upsample(tensors[f'De{decoder + 1}'], name=f'De{decoder + 1}-De{decoder}', **upsample_kwargs)
-
- # Tensors to concatenate in the Concatenate layer
- tensors_to_concatenate = sorted([tensor for tensor in tensors if f'---{decoder}' in tensor])[::-1]
- for index in range(len(tensors_to_concatenate)):
- tensors_to_concatenate[index] = tensors[tensors_to_concatenate[index]]
- tensors_to_concatenate.insert(levels - 1 - decoder, tensors[f'De{decoder}'])
-
- # Concatenate tensors, pass through convolution modules, then use deep supervision to create a side output
- tensors[f'De{decoder}'] = Concatenate(name=f'De{decoder}_Concatenate')(tensors_to_concatenate)
- tensors[f'De{decoder}'] = unet_utils.convolution_module(tensors[f'De{decoder}'], filters=filter_num_aggregate, name=f'De{decoder}', **module_kwargs)
- if deep_supervision or decoder == 1: # Decoder node 1 must always have deep supervision
- tensors[f'sup{decoder}_output'] = unet_utils.deep_supervision_side_output(tensors[f'De{decoder}'], num_classes=num_classes, output_level=decoder, name=f'sup{decoder}', **supervision_kwargs)
- tensors_with_supervision.append(tensors[f'sup{decoder}_output'])
-
- """ Add aggregated feature maps """
- for feature_map in range(1, decoder - 1):
- tensors[f'{decoder}---{feature_map}_feature'] = unet_utils.aggregated_feature_map(tensors[f'De{decoder}'], level1=decoder, level2=feature_map, name=f'{decoder}---{feature_map}_feature', **aggregated_kwargs)
-
- model = Model(inputs=tensors['input'], outputs=tensors_with_supervision[::-1], name=f'unet_3plus_{ndims}D')
-
- return model
-
-
-def attention_unet(
- input_shape: tuple[int],
- num_classes: int,
- pool_size: int | tuple[int] | list[int],
- levels: int,
- filter_num: tuple[int] | list[int],
- kernel_size: int = 3,
- squeeze_dims: int | tuple[int] | list[int] = None,
- modules_per_node: int = 5,
- batch_normalization: bool = True,
- activation: str = 'relu',
- padding: str = 'same',
- use_bias: bool = True,
- kernel_initializer: str = 'glorot_uniform',
- bias_initializer: str = 'zeros',
- kernel_regularizer: str = None,
- bias_regularizer: str = None,
- activity_regularizer: str = None,
- kernel_constraint: str = None,
- bias_constraint: str = None):
- """
- Builds a U-Net model.
-
- Parameters
- ----------
- input_shape: tuple
- Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
- num_classes: int
- Number of classes/labels that the U-Net will try to predict.
- pool_size: tuple or list
- Size of the mask in the MaxPooling and UpSampling layers.
- levels: int
- Number of levels in the U-Net. Must be greater than 1.
- filter_num: iterable of ints
- Number of convolution filters on each level of the U-Net.
- kernel_size: int or tuple
- Size of the kernel in the convolution layers.
- squeeze_dims: int, tuple, or None
- Dimensions/axes of the input to squeeze such that the target (y_true) will be smaller than the input.
- - (e.g. to remove the third dimension, set this parameter to 2 [axis=2 for the third dimension])
- modules_per_node: int
- Number of modules in each node of the U-Net.
- batch_normalization: bool
- Setting this to True will add a batch normalization layer after every convolution in the modules.
- activation: str
- Activation function to use in the modules.
- Can be any of tf.keras.activations, 'gaussian', 'gcu', 'leaky_relu', 'prelu', 'smelu', 'snake' (case-insensitive).
- padding: str
- Padding to use in the convolution layers.
- use_bias: bool
- Setting this to True will implement a bias vector in the convolution layers used in the modules.
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
-
- Returns
- -------
- model: tf.keras.models.Model object
- U-Net model.
-
- Raises
- ------
- ValueError
- If levels < 2
- If input_shape does not have 3 nor 4 dimensions
- If the length of filter_num does not match the number of levels
-
- References
- ----------
- https://arxiv.org/pdf/1505.04597.pdf
- """
-
- ndims = len(input_shape) - 1 # Number of dimensions in the input image (excluding the last dimension reserved for channels)
-
- if levels < 2:
- raise ValueError(f"levels must be greater than 1. Received value: {levels}")
-
- if len(input_shape) > 4 or len(input_shape) < 3:
- raise ValueError(f"input_shape can only have 3 or 4 dimensions (2D image + 1 dimension for channels OR a 3D image + 1 dimension for channels). Received shape: {np.shape(input_shape)}")
-
- if len(filter_num) != levels:
- raise ValueError(f"length of filter_num ({len(filter_num)}) does not match the number of levels ({levels})")
-
- # Keyword arguments for the convolution modules
- module_kwargs = dict({})
- module_kwargs['activation'] = activation
- module_kwargs['batch_normalization'] = batch_normalization
- module_kwargs['num_modules'] = modules_per_node
- module_kwargs['padding'] = padding
- module_kwargs['use_bias'] = use_bias
- module_kwargs['kernel_initializer'] = kernel_initializer
- module_kwargs['bias_initializer'] = bias_initializer
- module_kwargs['kernel_regularizer'] = kernel_regularizer
- module_kwargs['bias_regularizer'] = bias_regularizer
- module_kwargs['activity_regularizer'] = activity_regularizer
- module_kwargs['kernel_constraint'] = kernel_constraint
- module_kwargs['bias_constraint'] = bias_constraint
-
- # MaxPooling keyword arguments
- pool_kwargs = dict({})
- pool_kwargs['pool_size'] = pool_size
-
- # Keyword arguments for upsampling
- upsample_kwargs = dict({})
- upsample_kwargs['activation'] = activation
- upsample_kwargs['batch_normalization'] = batch_normalization
- upsample_kwargs['padding'] = padding
- upsample_kwargs['kernel_initializer'] = kernel_initializer
- upsample_kwargs['bias_initializer'] = bias_initializer
- upsample_kwargs['kernel_regularizer'] = kernel_regularizer
- upsample_kwargs['bias_regularizer'] = bias_regularizer
- upsample_kwargs['activity_regularizer'] = activity_regularizer
- upsample_kwargs['kernel_constraint'] = kernel_constraint
- upsample_kwargs['bias_constraint'] = bias_constraint
- upsample_kwargs['upsample_size'] = pool_size
- upsample_kwargs['use_bias'] = use_bias
-
- # Keyword arguments for the deep supervision output in the final decoder node
- supervision_kwargs = dict({})
- supervision_kwargs['upsample_size'] = pool_size
- supervision_kwargs['use_bias'] = True
- supervision_kwargs['squeeze_dims'] = squeeze_dims
- supervision_kwargs['padding'] = padding
- supervision_kwargs['kernel_initializer'] = kernel_initializer
- supervision_kwargs['bias_initializer'] = bias_initializer
- supervision_kwargs['kernel_regularizer'] = kernel_regularizer
- supervision_kwargs['bias_regularizer'] = bias_regularizer
- supervision_kwargs['activity_regularizer'] = activity_regularizer
- supervision_kwargs['kernel_constraint'] = kernel_constraint
- supervision_kwargs['bias_constraint'] = bias_constraint
-
- tensors = dict({}) # Tensors associated with each node and skip connections
-
- """ Setup the first encoder node with an input layer and a convolution module """
- tensors['input'] = Input(shape=input_shape, name='Input')
- tensors['En1'] = unet_utils.convolution_module(tensors['input'], filters=filter_num[0], kernel_size=kernel_size, name='En1', **module_kwargs)
-
- """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
- for encoder in np.arange(2, levels + 1): # Iterate through the rest of the encoder nodes
- pool_tensor = unet_utils.max_pool(tensors[f'En{encoder - 1}'], name=f'En{encoder - 1}-En{encoder}', **pool_kwargs) # Connect the next encoder node with a MaxPooling layer
- tensors[f'En{encoder}'] = unet_utils.convolution_module(pool_tensor, filters=filter_num[encoder - 1], kernel_size=kernel_size, name=f'En{encoder}', **module_kwargs) # Convolution modules
-
- tensors[f'AG{levels - 1}'] = unet_utils.attention_gate(tensors[f'En{levels - 1}'], tensors[f'En{levels}'], kernel_size, pool_size, name=f'AG{levels - 1}')
- upsample_tensor = unet_utils.upsample(tensors[f'En{levels}'], filters=filter_num[levels - 2], kernel_size=kernel_size, name=f'En{levels}-De{levels - 1}', **upsample_kwargs) # Connect the bottom encoder node to a decoder node
-
- """ Bottom decoder node """
- tensors[f'De{levels - 1}'] = Concatenate(name=f'De{levels - 1}_Concatenate')([tensors[f'AG{levels - 1}'], upsample_tensor]) # Concatenate the upsampled tensor and skip connection
- tensors[f'De{levels - 1}'] = unet_utils.convolution_module(tensors[f'De{levels - 1}'], filters=filter_num[levels - 2], kernel_size=kernel_size, name=f'De{levels - 1}', **module_kwargs) # Convolution module
- tensors[f'AG{levels - 2}'] = unet_utils.attention_gate(tensors[f'En{levels - 2}'], tensors[f'De{levels - 1}'], kernel_size, pool_size, name=f'AG{levels - 2}')
- upsample_tensor = unet_utils.upsample(tensors[f'De{levels - 1}'], filters=filter_num[levels - 3], kernel_size=kernel_size, name=f'De{levels - 1}-De{levels - 2}', **upsample_kwargs) # Connect the bottom decoder node to the next decoder node
-
- """ The rest of the decoder nodes (except the final node) are handled in this loop. Each node contains one concatenation of an upsampled tensor and a skip connection """
- for decoder in np.arange(2, levels-1)[::-1]:
- tensors[f'De{decoder}'] = Concatenate(name=f'De{decoder}_Concatenate')([tensors[f'AG{decoder}'], upsample_tensor]) # Concatenate the upsampled tensor and skip connection
- tensors[f'De{decoder}'] = unet_utils.convolution_module(tensors[f'De{decoder}'], filters=filter_num[decoder - 1], kernel_size=kernel_size, name=f'De{decoder}', **module_kwargs) # Convolution module
- tensors[f'AG{decoder - 1}'] = unet_utils.attention_gate(tensors[f'En{decoder - 1}'], tensors[f'De{decoder}'], kernel_size, pool_size, name=f'AG{decoder - 1}')
- upsample_tensor = unet_utils.upsample(tensors[f'De{decoder}'], filters=filter_num[decoder - 2], kernel_size=kernel_size, name=f'De{decoder}-De{decoder - 1}', **upsample_kwargs) # Connect the bottom decoder node to the next decoder node
-
- print(tensors.keys())
- """ Final decoder node begins with a concatenation and convolution module, followed by deep supervision """
- tensor_De1 = Concatenate(name='De1_Concatenate')([tensors['AG1'], upsample_tensor]) # Concatenate the upsampled tensor and skip connection
- tensor_De1 = unet_utils.convolution_module(tensor_De1, filters=filter_num[0], kernel_size=kernel_size, name='De1', **module_kwargs) # Convolution module
- tensors['output'] = unet_utils.deep_supervision_side_output(tensor_De1, num_classes=num_classes, kernel_size=1, output_level=1, name='final', **supervision_kwargs) # Deep supervision - this layer will output the model's prediction
-
- model = Model(inputs=tensors['input'], outputs=tensors['output'], name=f'attention_unet_{ndims}D')
-
- return model
diff --git a/pixi.lock b/pixi.lock
new file mode 100644
index 0000000..604457d
--- /dev/null
+++ b/pixi.lock
@@ -0,0 +1,5985 @@
+version: 6
+environments:
+ default:
+ channels:
+ - url: https://conda.anaconda.org/conda-forge/
+ indexes:
+ - https://pypi.org/simple
+ options:
+ pypi-prerelease-mode: if-necessary-or-explicit
+ packages:
+ linux-64:
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/_libgcc_mutex-0.1-conda_forge.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-2_gnu.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/affine-2.4.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/aiohappyeyeballs-2.6.1-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/aiohttp-3.13.3-py310h31b6992_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/aiosignal-1.4.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/annotated-types-0.7.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/aom-3.9.1-hac33072_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/appdirs-1.4.4-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/asciitree-0.3.3-py_2.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/async-timeout-5.0.1-pyhcf101f3_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/attr-2.5.2-h39aace5_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/blinker-1.9.0-pyhff2d567_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/blosc-1.21.6-he440d0b_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/bokeh-2.4.3-pyhd8ed1ab_3.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/branca-0.8.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.1.0-hb03c661_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.1.0-hb03c661_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.1.0-py310hea6c23e_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/brunsli-0.1-he3183e4_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_8.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/c-ares-1.34.6-hb03c661_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/c-blosc2-2.19.1-h4cfbee9_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py310he7384ee_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/cftime-1.6.5-py310hf779ad0_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/charls-2.4.2-h59595ed_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/click-8.3.1-pyh8f84b5b_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/click-plugins-1.1.1.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/cligj-0.7.2-pyhd8ed1ab_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/cloudpickle-3.1.2-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.2-py310h3788b33_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/cryptography-46.0.4-py310hb288b08_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/cytoolz-1.1.0-py310h7c4b9e2_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dacite-1.9.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dask-2023.3.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dask-core-2023.3.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dataclasses-0.8-pyhc8e2a94_3.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/dav1d-1.2.1-hd590300_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/distributed-2023.3.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/docker-pycreds-0.4.0-py_0.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/eval_type_backport-0.3.1-pyha770c72_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/fasteners-0.19-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/fiona-1.10.1-py310hea6c23e_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/folium-0.20.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.61.1-py310h3406613_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/freetype-2.14.1-ha770c72_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/freexl-2.0.0-h9dce30a_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/frozenlist-1.7.0-py310h9548a50_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/fsspec-2026.2.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/gcsfs-2026.2.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/geopandas-0.14.4-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/geopandas-base-0.14.4-pyha770c72_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/geos-3.14.0-h480dda7_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/geotiff-1.7.4-h239500f_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/giflib-5.2.2-hd590300_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/gitdb-4.0.12-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/gitpython-3.1.46-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-api-core-2.29.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-auth-2.48.0-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-auth-oauthlib-1.2.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-cloud-core-2.5.0-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-cloud-storage-3.9.0-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/google-cloud-storage-control-1.8.0-py310hff52083_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/google-crc32c-1.8.0-py310hf432777_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-resumable-media-2.8.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/googleapis-common-protos-1.72.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/googleapis-common-protos-grpc-1.72.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/grpc-google-iam-v1-0.14.3-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/grpcio-1.54.3-py310heca2aa9_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/grpcio-status-1.54.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/h2-4.3.0-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/h5py-3.15.1-nompi_py310h4aa865e_101.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/hdf4-4.2.15-h2a13503_7.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/hdf5-1.14.6-nompi_h1b119a7_105.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/hpack-4.1.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/hyperframe-6.1.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/icu-75.1-he02047a_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/idna-3.11-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/imagecodecs-2025.3.30-py310h4eb8eaf_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/imageio-2.37.0-pyhfb79c49_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/importlib-metadata-8.7.0-pyhe01879c_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/jinja2-3.1.6-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/joblib-1.5.3-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/json-c-0.18-h6688a6e_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/jxrlib-1.1-hd590300_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/keyutils-1.6.3-hb9d3cd8_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py310haaf941d_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.21.3-h659f571_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/lazy-loader-0.4-pyhd8ed1ab_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/lcms2-2.18-h0c24ade_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/ld_impl_linux-64-2.45-bootstrap_ha15bf96_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/lerc-4.0.0-h0aef613_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libabseil-20230125.3-cxx17_h59595ed_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libaec-1.1.5-h088129d_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libarchive-3.8.5-gpl_hc2c16d8_100.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libavif16-1.3.0-h316e467_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.11.0-5_h4a7cf45_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlicommon-1.1.0-hb03c661_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlidec-1.1.0-hb03c661_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlienc-1.1.0-hb03c661_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.11.0-5_h0358290_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libcrc32c-1.1.2-h9c3ff4c_0.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libcurl-8.18.0-h4e3cde8_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libdeflate-1.24-h86f0d12_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libedit-3.1.20250104-pl5321h7949ede_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libev-4.33-hd590300_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libexpat-2.7.3-hecca717_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libffi-3.5.2-h3435931_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libfreetype-2.14.1-ha770c72_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libfreetype6-2.14.1-h73754d4_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-15.2.0-he0feb66_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-ng-15.2.0-h69a702a_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgdal-core-3.10.3-h95ec890_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran-15.2.0-h69a702a_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran5-15.2.0-h68bc16d_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgomp-15.2.0-he0feb66_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgrpc-1.54.3-hb20ce57_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libhwy-1.3.0-h4c17acf_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libiconv-1.18-h3b78370_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libjpeg-turbo-3.1.2-hb03c661_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libjxl-0.11.1-h6cb5226_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libkml-1.3.0-h01aab08_1016.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/liblapack-3.11.0-5_h47877c9_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/liblzma-5.8.2-hb03c661_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libnetcdf-4.9.3-nompi_h11f7409_103.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libnghttp2-1.67.0-had1ee68_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libnsl-2.0.1-hb9d3cd8_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libopenblas-0.3.30-pthreads_h94d23a6_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libpng-1.6.54-h421ea60_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libprotobuf-3.21.12-hfc55251_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/librttopo-1.1.0-h96cd706_19.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libspatialindex-2.1.0-he57a185_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libspatialite-5.1.0-h2eee824_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libsqlite-3.51.2-h0c1763c_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libssh2-1.11.1-hcf80075_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-15.2.0-h934c35e_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-ng-15.2.0-hdf11a46_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libtiff-4.7.1-h8261f1e_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libuuid-2.41.3-h5347b49_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libwebp-base-1.6.0-hd42ef1d_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcb-1.17.0-h8a09558_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcrypt-4.4.36-hd590300_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-16-2.15.1-ha9997c6_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-2.15.1-h26afc86_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-devel-2.15.1-h26afc86_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libzip-1.11.2-h6991a6a_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.1-hb9d3cd8_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libzopfli-1.0.3-h9c3ff4c_0.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/locket-1.0.0-pyhd8ed1ab_0.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/lz4-4.4.5-py310hde1b0b5_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/lz4-c-1.10.0-h5888daf_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/lzo-2.10-h280c20c_1002.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/mapclassify-2.8.1-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/markupsafe-3.0.3-py310h3406613_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.1-py310h68603db_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/minizip-4.0.10-h05a5f5f_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/msgpack-python-1.1.2-py310h03d9f68_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/multidict-6.7.0-py310h3406613_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.5-h2d0b736_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/netcdf4-1.7.4-nompi_py310hd27e1a9_102.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.4.2-pyh267e887_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/numcodecs-0.13.1-py310h5eaa309_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-1.26.4-py310hb13e2d6_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/oauthlib-3.3.1-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.4-h55fea9a_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.1-h35e630c_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-26.0-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pandas-2.3.3-py310h0158d43_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/partd-1.4.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pcre2-10.46-h1321c63_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.0.0-py310h049bd52_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pooch-1.8.2-pyhd8ed1ab_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/proj-9.6.2-h18fbb6c_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/propcache-0.3.1-py310h89163eb_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/proto-plus-1.27.1-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/protobuf-4.21.12-py310heca2aa9_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.2-py310h139afa4_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyasn1-0.6.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyasn1-modules-0.4.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pycparser-2.22-pyh29332c3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pydantic-2.12.5-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pydantic-core-2.41.5-py310hd8f68c5_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyjwt-2.11.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyopenssl-25.3.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.2-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pyproj-3.7.1-py310h71d0299_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.10.19-h3c07f61_3_cpython.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.10-8_cp310.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyu2f-0.1.5-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pywavelets-1.8.0-py310hf462985_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py310h3406613_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/rasterio-1.4.3-py310hf3df72b_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/rav1e-0.7.1-h8fae777_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/re2-2023.03.02-h8c504da_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/regionmask-0.13.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/requests-2.32.5-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/requests-oauthlib-2.0.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/rsa-4.9.1-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/rtree-1.4.1-pyh11ca60a_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/scikit-image-0.25.2-py310h0158d43_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/scikit-learn-1.7.2-py310h228f341_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.15.2-py310h1d65ade_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/sentry-sdk-2.51.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/setproctitle-1.3.7-py310h139afa4_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.10.2-pyh332efcf_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/shapely-2.1.2-py310h777b3ac_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/smmap-5.0.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/snappy-1.2.2-h03e3b7b_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/snuggs-1.4.7-pyhd8ed1ab_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/sortedcontainers-2.4.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/sqlite-3.51.2-hbc0de68_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/svt-av1-4.0.0-hecca717_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/tblib-3.2.2-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/threadpoolctl-3.6.0-pyhecae5ae_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/tifffile-2025.5.10-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_h4845f30_101.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/toolz-1.1.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py310h7c4b9e2_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/typing-3.10.0.0-pyhd8ed1ab_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/typing-inspection-0.4.2-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-17.0.0-py310h7c4b9e2_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/uriparser-0.9.8-hac33072_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.5.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/wandb-0.24.1-py310hdfeec95_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/xarray-2024.3.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/xbatcher-0.4.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/xerces-c-3.2.5-h988505b_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxau-1.0.12-hb03c661_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxdmcp-1.1.5-hb03c661_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/xyzservices-2025.11.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/yaml-0.2.5-h280c20c_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/yarl-1.22.0-py310h3406613_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/zarr-2.18.3-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/zfp-1.0.1-h909a3a2_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/zict-3.0.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/zipp-3.23.0-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.1-hb9d3cd8_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-ng-2.2.5-hde8ca8f_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/zstandard-0.23.0-py310h7c4b9e2_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda
+ - conda: .
+ - pypi: https://files.pythonhosted.org/packages/18/a6/907a406bb7d359e6a63f99c313846d9eec4f7e6f7437809e03aa00fa3074/absl_py-2.4.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/67/84/d844b79acd9fe15ded60b614b7df04a12fad854ee1fbb8415d726ab1beeb/aiobotocore-3.1.1-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/10/a1/510b0a7fadc6f43a6ce50152e69dbd86415240835868bb0bd9b5b88b1e06/aioitertools-0.13.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/2b/03/13dde6512ad7b4557eb792fbcf0c653af6076b81e5941d36ec61f7ce6028/astunparse-1.6.3-py2.py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/3d/8d/6d7b016383b1f74dd93611b1c5078bbaddaca901553ab886dcda87cae365/botocore-1.42.30-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/1d/33/f1c6a276de27b7d7339a34749cc33fa87f077f921969c47185d34a887ae2/gast-0.7.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/a3/de/c648ef6835192e6e2cc03f40b19eeda4382c49b5bafb43d88b931c4c74ac/google_pasta-0.2.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/ba/61/cc8be27bd65082440754be443b17b6f7c185dec5e00dfdaeab4f8662e4a8/keras-3.12.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/1d/fc/716c1e62e512ef1c160e7984a73a5fc7df45166f2ff3f254e71c58076f7c/libclang-18.1.1-py2.py3-none-manylinux2010_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/59/1b/6ef961f543593969d25b2afe57a3564200280528caa9bd1082eecdd7b3bc/markdown-3.10.1-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz
+ - pypi: https://files.pythonhosted.org/packages/b2/bc/465daf1de06409cdd4532082806770ee0d8d7df434da79c76564d0f69741/namex-0.1.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/f9/43/a5365b345c989d9f7de2f2406e59b4a792ba3541b5d47edda5031b2730a6/nvidia_cublas_cu12-12.5.3.2-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/3a/64/81515a1b5872dc0fc817deaec2bdba160dc50188e0d53b907d10c6e6d568/nvidia_cuda_cupti_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/44/cc/36363057676e6140d0bdb07fa6df5419b68203c5cba8c412b5600fd0d105/nvidia_cuda_nvcc_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/19/d1/342c2bcf6172db65fcbd9102f9941876e730d98977e69c00df85940fa8ce/nvidia_cuda_nvrtc_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/71/05/80f3fe49e9905570bc27aea9493baa1891c3780a7fc4e1f872c7902df066/nvidia_cuda_runtime_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/8e/56/bb5c08a8d401fc1b21a10e9c58907e70e8f18bdfca34b7ecfb87bbcdad63/nvidia_cudnn_cu12-9.3.0.75-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/e4/85/f18c88f63489cdced17b06d3b627adca8add7d7b8cce8c11213e93a902b4/nvidia_cufft_cu12-11.2.3.61-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/f8/ce/cc6daf7820804ee7f11a0352a1f0fd59cec5f12e904f5bbaee6d928ffdaf/nvidia_curand_cu12-10.3.6.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/33/73/57fbf55b3f378a73faecde397a0927ea205b458f06573dfb191b7d9fd1d3/nvidia_cusolver_cu12-11.6.3.83-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/b3/29/03726191334fa523d9654e3dacca5cc152f24bb9fa1a5721e4dddac7a8c5/nvidia_cusparse_cu12-12.5.1.3-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/ed/1f/6482380ec8dcec4894e7503490fc536d846b0d59694acad9cf99f27d0e7d/nvidia_nccl_cu12-2.23.4-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/75/bc/e0d0dbb85246a086ab14839979039647bce501d8c661a159b8b019d987b7/nvidia_nvjitlink_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/5f/7d/9ec5967f3e2915fbc441f72c3892a7f0fb3618e3ae5c8a44181ce4aa641c/obstore-0.8.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/83/8e/09d899ad531d50b79aa24e7558f604980fe4048350172e643bb1b9983aec/optree-0.18.0.tar.gz
+ - pypi: https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/87/2a/a1810c8627b9ec8c57ec5ec325d306701ae7be50235e8fd81266e002a3cc/rich-14.3.1-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/57/e1/64c264db50b68de8a438b60ceeb921b2f22da3ebb7ad6255150225d0beac/s3fs-2026.2.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/5d/12/4f70e8e2ba0dbe72ea978429d8530b0333f0ed2140cc571a48802878ef99/tensorboard-2.19.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/7a/13/e503968fefabd4c6b2650af21e110aa8466fe21432cd7c43a84577a89438/tensorboard_data_server-0.7.2-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/77/60/51d921b17f7b1547db32018d6a933627d4e2a762d2fc2ca6c0032cc8b062/tensorflow-2.19.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/f3/48/47b7d25572961a48b1de3729b7a11e835b888e41e0203cca82df95d23b91/tensorflow_io_gcs_filesystem-0.37.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/33/d1/8bb87d21e9aeb323cc03034f5eaf2c8f69841e40e4853c2627edf8111ed3/termcolor-3.3.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/ad/e4/8d97cca767bcc1be76d16fb76951608305561c6e056811587f36cb1316a8/werkzeug-3.1.5-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/c6/93/5cf92edd99617095592af919cb81d4bff61c5dbbb70d3c92099425a8ec34/wrapt-2.0.1-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl
+ dev:
+ channels:
+ - url: https://conda.anaconda.org/conda-forge/
+ indexes:
+ - https://pypi.org/simple
+ options:
+ pypi-prerelease-mode: if-necessary-or-explicit
+ packages:
+ linux-64:
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/_libgcc_mutex-0.1-conda_forge.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-2_gnu.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/affine-2.4.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/aiohappyeyeballs-2.6.1-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/aiohttp-3.13.3-py310h31b6992_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/aiosignal-1.4.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/annotated-types-0.7.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/aom-3.9.1-hac33072_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/appdirs-1.4.4-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/asciitree-0.3.3-py_2.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/async-timeout-5.0.1-pyhcf101f3_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/attr-2.5.2-h39aace5_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/blinker-1.9.0-pyhff2d567_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/blosc-1.21.6-he440d0b_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/bokeh-2.4.3-pyhd8ed1ab_3.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/branca-0.8.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.1.0-hb03c661_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.1.0-hb03c661_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.1.0-py310hea6c23e_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/brunsli-0.1-he3183e4_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_8.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/c-ares-1.34.6-hb03c661_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/c-blosc2-2.19.1-h4cfbee9_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py310he7384ee_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/cftime-1.6.5-py310hf779ad0_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/charls-2.4.2-h59595ed_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/click-8.3.1-pyh8f84b5b_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/click-plugins-1.1.1.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/cligj-0.7.2-pyhd8ed1ab_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/cloudpickle-3.1.2-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.2-py310h3788b33_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/cryptography-46.0.4-py310hb288b08_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/cytoolz-1.1.0-py310h7c4b9e2_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dacite-1.9.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dask-2023.3.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dask-core-2023.3.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dataclasses-0.8-pyhc8e2a94_3.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/dav1d-1.2.1-hd590300_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/distributed-2023.3.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/docker-pycreds-0.4.0-py_0.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/eval_type_backport-0.3.1-pyha770c72_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/fasteners-0.19-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/fiona-1.10.1-py310hea6c23e_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/folium-0.20.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.61.1-py310h3406613_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/freetype-2.14.1-ha770c72_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/freexl-2.0.0-h9dce30a_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/frozenlist-1.7.0-py310h9548a50_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/fsspec-2026.2.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/gcsfs-2026.2.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/geopandas-0.14.4-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/geopandas-base-0.14.4-pyha770c72_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/geos-3.14.0-h480dda7_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/geotiff-1.7.4-h239500f_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/giflib-5.2.2-hd590300_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/gitdb-4.0.12-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/gitpython-3.1.46-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-api-core-2.29.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-auth-2.48.0-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-auth-oauthlib-1.2.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-cloud-core-2.5.0-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-cloud-storage-3.9.0-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/google-cloud-storage-control-1.8.0-py310hff52083_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/google-crc32c-1.8.0-py310hf432777_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/google-resumable-media-2.8.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/googleapis-common-protos-1.72.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/googleapis-common-protos-grpc-1.72.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/grpc-google-iam-v1-0.14.3-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/grpcio-1.54.3-py310heca2aa9_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/grpcio-status-1.54.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/h2-4.3.0-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/h5py-3.15.1-nompi_py310h4aa865e_101.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/hdf4-4.2.15-h2a13503_7.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/hdf5-1.14.6-nompi_h1b119a7_105.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/hpack-4.1.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/hyperframe-6.1.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/icu-75.1-he02047a_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/idna-3.11-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/imagecodecs-2025.3.30-py310h4eb8eaf_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/imageio-2.37.0-pyhfb79c49_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/importlib-metadata-8.7.0-pyhe01879c_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/jinja2-3.1.6-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/joblib-1.5.3-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/json-c-0.18-h6688a6e_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/jxrlib-1.1-hd590300_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/keyutils-1.6.3-hb9d3cd8_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py310haaf941d_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.21.3-h659f571_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/lazy-loader-0.4-pyhd8ed1ab_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/lcms2-2.18-h0c24ade_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/ld_impl_linux-64-2.45-bootstrap_ha15bf96_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/lerc-4.0.0-h0aef613_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libabseil-20230125.3-cxx17_h59595ed_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libaec-1.1.5-h088129d_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libarchive-3.8.5-gpl_hc2c16d8_100.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libavif16-1.3.0-h316e467_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.11.0-5_h4a7cf45_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlicommon-1.1.0-hb03c661_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlidec-1.1.0-hb03c661_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlienc-1.1.0-hb03c661_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.11.0-5_h0358290_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libcrc32c-1.1.2-h9c3ff4c_0.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libcurl-8.18.0-h4e3cde8_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libdeflate-1.24-h86f0d12_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libedit-3.1.20250104-pl5321h7949ede_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libev-4.33-hd590300_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libexpat-2.7.3-hecca717_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libffi-3.5.2-h3435931_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libfreetype-2.14.1-ha770c72_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libfreetype6-2.14.1-h73754d4_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-15.2.0-he0feb66_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-ng-15.2.0-h69a702a_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgdal-core-3.10.3-h95ec890_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran-15.2.0-h69a702a_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran5-15.2.0-h68bc16d_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgomp-15.2.0-he0feb66_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgrpc-1.54.3-hb20ce57_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libhwy-1.3.0-h4c17acf_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libiconv-1.18-h3b78370_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libjpeg-turbo-3.1.2-hb03c661_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libjxl-0.11.1-h6cb5226_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libkml-1.3.0-h01aab08_1016.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/liblapack-3.11.0-5_h47877c9_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/liblzma-5.8.2-hb03c661_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libnetcdf-4.9.3-nompi_h11f7409_103.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libnghttp2-1.67.0-had1ee68_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libnsl-2.0.1-hb9d3cd8_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libopenblas-0.3.30-pthreads_h94d23a6_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libpng-1.6.54-h421ea60_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libprotobuf-3.21.12-hfc55251_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/librttopo-1.1.0-h96cd706_19.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libspatialindex-2.1.0-he57a185_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libspatialite-5.1.0-h2eee824_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libsqlite-3.51.2-h0c1763c_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libssh2-1.11.1-hcf80075_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-15.2.0-h934c35e_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-ng-15.2.0-hdf11a46_16.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libtiff-4.7.1-h8261f1e_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libuuid-2.41.3-h5347b49_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libwebp-base-1.6.0-hd42ef1d_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcb-1.17.0-h8a09558_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcrypt-4.4.36-hd590300_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-16-2.15.1-ha9997c6_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-2.15.1-h26afc86_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-devel-2.15.1-h26afc86_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libzip-1.11.2-h6991a6a_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.1-hb9d3cd8_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libzopfli-1.0.3-h9c3ff4c_0.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/noarch/locket-1.0.0-pyhd8ed1ab_0.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/lz4-4.4.5-py310hde1b0b5_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/lz4-c-1.10.0-h5888daf_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/lzo-2.10-h280c20c_1002.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/mapclassify-2.8.1-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/markupsafe-3.0.3-py310h3406613_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.1-py310h68603db_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/minizip-4.0.10-h05a5f5f_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/msgpack-python-1.1.2-py310h03d9f68_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/multidict-6.7.0-py310h3406613_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.5-h2d0b736_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/netcdf4-1.7.4-nompi_py310hd27e1a9_102.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.4.2-pyh267e887_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/numcodecs-0.13.1-py310h5eaa309_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-1.26.4-py310hb13e2d6_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/oauthlib-3.3.1-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.4-h55fea9a_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.1-h35e630c_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/packaging-26.0-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pandas-2.3.3-py310h0158d43_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/partd-1.4.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pcre2-10.46-h1321c63_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.0.0-py310h049bd52_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pooch-1.8.2-pyhd8ed1ab_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/proj-9.6.2-h18fbb6c_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/propcache-0.3.1-py310h89163eb_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/proto-plus-1.27.1-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/protobuf-4.21.12-py310heca2aa9_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.2-py310h139afa4_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyasn1-0.6.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyasn1-modules-0.4.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pycparser-2.22-pyh29332c3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pydantic-2.12.5-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pydantic-core-2.41.5-py310hd8f68c5_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyjwt-2.11.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyopenssl-25.3.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.2-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pyproj-3.7.1-py310h71d0299_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.10.19-h3c07f61_3_cpython.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.10-8_cp310.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/pyu2f-0.1.5-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pywavelets-1.8.0-py310hf462985_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py310h3406613_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/rasterio-1.4.3-py310hf3df72b_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/rav1e-0.7.1-h8fae777_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/re2-2023.03.02-h8c504da_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/regionmask-0.13.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/requests-2.32.5-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/requests-oauthlib-2.0.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/rsa-4.9.1-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/rtree-1.4.1-pyh11ca60a_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/scikit-image-0.25.2-py310h0158d43_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/scikit-learn-1.7.2-py310h228f341_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.15.2-py310h1d65ade_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/sentry-sdk-2.51.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/setproctitle-1.3.7-py310h139afa4_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.10.2-pyh332efcf_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/shapely-2.1.2-py310h777b3ac_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/smmap-5.0.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/snappy-1.2.2-h03e3b7b_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/snuggs-1.4.7-pyhd8ed1ab_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/sortedcontainers-2.4.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/sqlite-3.51.2-hbc0de68_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/svt-av1-4.0.0-hecca717_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/tblib-3.2.2-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/threadpoolctl-3.6.0-pyhecae5ae_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/tifffile-2025.5.10-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_h4845f30_101.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/toolz-1.1.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py310h7c4b9e2_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/typing-3.10.0.0-pyhd8ed1ab_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/typing-inspection-0.4.2-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-17.0.0-py310h7c4b9e2_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/uriparser-0.9.8-hac33072_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.5.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/wandb-0.24.1-py310hdfeec95_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/xarray-2024.3.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/xbatcher-0.4.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/xerces-c-3.2.5-h988505b_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxau-1.0.12-hb03c661_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxdmcp-1.1.5-hb03c661_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/xyzservices-2025.11.0-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/yaml-0.2.5-h280c20c_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/yarl-1.22.0-py310h3406613_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/zarr-2.18.3-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/zfp-1.0.1-h909a3a2_5.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/zict-3.0.0-pyhd8ed1ab_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/zipp-3.23.0-pyhcf101f3_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.1-hb9d3cd8_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-ng-2.2.5-hde8ca8f_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/zstandard-0.23.0-py310h7c4b9e2_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda
+ - conda: .
+ - pypi: https://files.pythonhosted.org/packages/18/a6/907a406bb7d359e6a63f99c313846d9eec4f7e6f7437809e03aa00fa3074/absl_py-2.4.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/67/84/d844b79acd9fe15ded60b614b7df04a12fad854ee1fbb8415d726ab1beeb/aiobotocore-3.1.1-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/10/a1/510b0a7fadc6f43a6ce50152e69dbd86415240835868bb0bd9b5b88b1e06/aioitertools-0.13.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/2b/03/13dde6512ad7b4557eb792fbcf0c653af6076b81e5941d36ec61f7ce6028/astunparse-1.6.3-py2.py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/b9/fa/123043af240e49752f1c4bd24da5053b6bd00cad78c2be53c0d1e8b975bc/backports.tarfile-1.2.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/3d/8d/6d7b016383b1f74dd93611b1c5078bbaddaca901553ab886dcda87cae365/botocore-1.42.30-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/db/3c/33bac158f8ab7f89b2e59426d5fe2e4f63f7ed25df84c036890172b412b5/cfgv-3.5.0-py2.py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/35/fb/05b9830c2e8275ebc031e0019387cda99113e62bb500ab328bb72578183b/coverage-7.13.2-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/ab/28/960c311aae084deef57ece41aac13cb359b06ce31b7771139e79c394a1b7/deptry-0.24.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/02/10/5da547df7a391dcde17f59520a231527b8571e6f46fc8efb02ccb370ab12/docutils-0.22.4-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/8a/0e/97c33bf5009bdbac74fd2beace167cab3f978feb69cc36f1ef79360d6c4e/exceptiongroup-1.3.1-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/b5/36/7fb70f04bf00bc646cd5bb45aa9eddb15e19437a28b8fb2b4a5249fac770/filelock-3.20.3-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/1d/33/f1c6a276de27b7d7339a34749cc33fa87f077f921969c47185d34a887ae2/gast-0.7.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/a3/de/c648ef6835192e6e2cc03f40b19eeda4382c49b5bafb43d88b931c4c74ac/google_pasta-0.2.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/9f/cb/18326d2d89ad3b0dd143da971e77afd1e6ca6674f1b1c3df4b6bec6279fc/id-1.5.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/b8/58/40fbbcefeda82364720eba5cf2270f98496bdfa19ea75b4cccae79c698e6/identify-2.6.16-py2.py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/7f/ed/e3705d6d02b4f7aea715a353c8ce193efd0b5db13e204df895d38734c244/isort-7.0.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/8d/48/aa685dbf1024c7bd82bede569e3a85f82c32fd3d79ba5fea578f0159571a/jaraco_context-6.1.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/fd/c4/813bb09f0985cb21e959f21f2464169eca882656849adf727ac7bb7e1767/jaraco_functools-4.4.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/b2/a3/e137168c9c44d18eff0376253da9f1e9234d0239e0ee230d2fee6cea8e55/jeepney-0.9.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/ba/61/cc8be27bd65082440754be443b17b6f7c185dec5e00dfdaeab4f8662e4a8/keras-3.12.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/1d/fc/716c1e62e512ef1c160e7984a73a5fc7df45166f2ff3f254e71c58076f7c/libclang-18.1.1-py2.py3-none-manylinux2010_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/82/5f/3e85351c523f73ad8d938989e9a58c7f59fb9c17f761b9981b43f0025ce7/librt-0.7.8-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/1f/c9/e0f8e4e6e8a69e5959b06499582dca6349db6769cc7fdfb8a02a7c75a9ae/lxml_stubs-0.5.1-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/59/1b/6ef961f543593969d25b2afe57a3564200280528caa9bd1082eecdd7b3bc/markdown-3.10.1-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz
+ - pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/2a/0d/93c2e4a287f74ef11a66fb6d49c7a9f05e47b0a4399040e6719b57f500d2/mypy-1.19.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/b2/bc/465daf1de06409cdd4532082806770ee0d8d7df434da79c76564d0f69741/namex-0.1.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/42/0f/c76bf3dba22c73c38e9b1113b017cf163f7696f50e003404ec5ecdb1e8a6/nh3-0.3.2-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/f9/43/a5365b345c989d9f7de2f2406e59b4a792ba3541b5d47edda5031b2730a6/nvidia_cublas_cu12-12.5.3.2-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/3a/64/81515a1b5872dc0fc817deaec2bdba160dc50188e0d53b907d10c6e6d568/nvidia_cuda_cupti_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/44/cc/36363057676e6140d0bdb07fa6df5419b68203c5cba8c412b5600fd0d105/nvidia_cuda_nvcc_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/19/d1/342c2bcf6172db65fcbd9102f9941876e730d98977e69c00df85940fa8ce/nvidia_cuda_nvrtc_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/71/05/80f3fe49e9905570bc27aea9493baa1891c3780a7fc4e1f872c7902df066/nvidia_cuda_runtime_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/8e/56/bb5c08a8d401fc1b21a10e9c58907e70e8f18bdfca34b7ecfb87bbcdad63/nvidia_cudnn_cu12-9.3.0.75-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/e4/85/f18c88f63489cdced17b06d3b627adca8add7d7b8cce8c11213e93a902b4/nvidia_cufft_cu12-11.2.3.61-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/f8/ce/cc6daf7820804ee7f11a0352a1f0fd59cec5f12e904f5bbaee6d928ffdaf/nvidia_curand_cu12-10.3.6.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/33/73/57fbf55b3f378a73faecde397a0927ea205b458f06573dfb191b7d9fd1d3/nvidia_cusolver_cu12-11.6.3.83-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/b3/29/03726191334fa523d9654e3dacca5cc152f24bb9fa1a5721e4dddac7a8c5/nvidia_cusparse_cu12-12.5.1.3-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/ed/1f/6482380ec8dcec4894e7503490fc536d846b0d59694acad9cf99f27d0e7d/nvidia_nccl_cu12-2.23.4-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/75/bc/e0d0dbb85246a086ab14839979039647bce501d8c661a159b8b019d987b7/nvidia_nvjitlink_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/5f/7d/9ec5967f3e2915fbc441f72c3892a7f0fb3618e3ae5c8a44181ce4aa641c/obstore-0.8.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/83/8e/09d899ad531d50b79aa24e7558f604980fe4048350172e643bb1b9983aec/optree-0.18.0.tar.gz
+ - pypi: https://files.pythonhosted.org/packages/d1/c6/df1fe324248424f77b89371116dab5243db7f052c32cc9fe7442ad9c5f75/pandas_stubs-2.3.3.260113-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/ef/3c/2c197d226f9ea224a9ab8d197933f9da0ae0aac5b6e0f884e2b8d9c8e9f7/pathspec-1.0.4-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/5d/19/fd3ef348460c80af7bb4669ea7926651d1f95c23ff2df18b9d24bab4f3fa/pre_commit-4.5.1-py2.py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/3f/51/d4db610ef29373b879047326cbf6fa98b6c1969d6f6dc423279de2b1be2c/requests_toolbelt-1.0.0-py2.py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/bd/60/50fbb6ffb35f733654466f1a90d162bcbea358adc3b0871339254fbc37b2/requirements_parser-0.13.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/ff/9a/9afaade874b2fa6c752c36f1548f718b5b83af81ed9b76628329dab81c1b/rfc3986-2.0.0-py2.py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/87/2a/a1810c8627b9ec8c57ec5ec325d306701ae7be50235e8fd81266e002a3cc/rich-14.3.1-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/ca/71/37daa46f89475f8582b7762ecd2722492df26421714a33e72ccc9a84d7a5/ruff-0.14.14-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/57/e1/64c264db50b68de8a438b60ceeb921b2f22da3ebb7ad6255150225d0beac/s3fs-2026.2.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/5d/12/4f70e8e2ba0dbe72ea978429d8530b0333f0ed2140cc571a48802878ef99/tensorboard-2.19.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/7a/13/e503968fefabd4c6b2650af21e110aa8466fe21432cd7c43a84577a89438/tensorboard_data_server-0.7.2-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/77/60/51d921b17f7b1547db32018d6a933627d4e2a762d2fc2ca6c0032cc8b062/tensorflow-2.19.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/f3/48/47b7d25572961a48b1de3729b7a11e835b888e41e0203cca82df95d23b91/tensorflow_io_gcs_filesystem-0.37.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/33/d1/8bb87d21e9aeb323cc03034f5eaf2c8f69841e40e4853c2627edf8111ed3/termcolor-3.3.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/23/d1/136eb2cb77520a31e1f64cbae9d33ec6df0d78bdf4160398e86eec8a8754/tomli-2.4.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/3a/7a/882d99539b19b1490cac5d77c67338d126e4122c8276bf640e411650c830/twine-6.2.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/e7/c1/56ef16bf5dcd255155cc736d276efa6ae0a5c26fd685e28f0412a4013c01/types_pytz-2025.2.0.20251108-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/bd/e0/1eed384f02555dde685fff1a1ac805c1c7dcb6dd019c916fe659b1c1f9ec/types_pyyaml-6.0.12.20250915-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/1c/12/709ea261f2bf91ef0a26a9eed20f2623227a8ed85610c1e54c5805692ecb/types_requests-2.32.4.20260107-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/3f/13/3ff0781445d7c12730befce0fddbbc7a76e56eb0e7029446f2853238360a/types_tqdm-4.67.0.20250809-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/6a/2a/dc2228b2888f51192c7dc766106cd475f1b768c10caaf9727659726f7391/virtualenv-20.36.1-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/ad/e4/8d97cca767bcc1be76d16fb76951608305561c6e056811587f36cb1316a8/werkzeug-3.1.5-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/c6/93/5cf92edd99617095592af919cb81d4bff61c5dbbb70d3c92099425a8ec34/wrapt-2.0.1-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl
+ test:
+ channels:
+ - url: https://conda.anaconda.org/conda-forge/
+ indexes:
+ - https://pypi.org/simple
+ options:
+ pypi-prerelease-mode: if-necessary-or-explicit
+ packages:
+ linux-64:
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_9.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dacite-1.9.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dataclasses-0.8-pyhc8e2a94_3.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/icu-78.2-h33c6efd_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/ld_impl_linux-64-2.45.1-default_hbd61a6d_101.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.11.0-5_h4a7cf45_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.11.0-5_h0358290_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libexpat-2.7.4-hecca717_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libffi-3.5.2-h3435931_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-15.2.0-he0feb66_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-ng-15.2.0-h69a702a_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran-15.2.0-h69a702a_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran5-15.2.0-h68bc16d_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libgomp-15.2.0-he0feb66_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/liblapack-3.11.0-5_h47877c9_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/liblzma-5.8.2-hb03c661_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libnsl-2.0.1-hb9d3cd8_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libopenblas-0.3.30-pthreads_h94d23a6_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libsqlite-3.51.2-hf4e2dac_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-15.2.0-h934c35e_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-ng-15.2.0-hdf11a46_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libuuid-2.41.3-h5347b49_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcrypt-4.4.36-hd590300_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.1-hb9d3cd8_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.5-h2d0b736_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-1.26.4-py312heda63a1_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.1-h35e630c_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.12.12-hd63d673_2_cpython.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.12-8_cp312.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py312h8a5da7c_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_h366c992_103.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/yaml-0.2.5-h280c20c_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda
+ - pypi: https://files.pythonhosted.org/packages/5d/a0/2ea570925524ef4e00bb6c82649f5682a77fac5ab910a65c9284de422600/coverage-7.13.4-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl
+ - pypi: https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl
+ osx-arm64:
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/_openmp_mutex-4.5-7_kmp_llvm.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/bzip2-1.0.8-hd037594_9.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dacite-1.9.2-pyhd8ed1ab_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/dataclasses-0.8-pyhc8e2a94_3.tar.bz2
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/icu-78.2-h38cb7af_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libblas-3.11.0-5_h51639a9_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libcblas-3.11.0-5_hb0561ab_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libcxx-21.1.8-h55c6f16_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libexpat-2.7.4-hf6b4638_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libffi-3.5.2-hcf2aa1b_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libgcc-15.2.0-hcbb3090_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libgfortran-15.2.0-h07b0088_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libgfortran5-15.2.0-hdae7583_18.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/liblapack-3.11.0-5_hd9741b5_openblas.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/liblzma-5.8.2-h8088a28_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libopenblas-0.3.30-openmp_ha158390_4.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libsqlite-3.51.2-h1ae2325_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libzlib-1.3.1-h8359307_2.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/llvm-openmp-21.1.8-h4a912ad_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ncurses-6.5-h5e97a16_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/numpy-1.26.4-py312h8442bc7_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/openssl-3.6.1-hd24854e_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.12.12-h18782d2_2_cpython.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.12-8_cp312.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyyaml-6.0.3-py312h04c11ed_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/readline-8.3-h46df422_0.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/tk-8.6.13-h010d191_3.conda
+ - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
+ - conda: https://conda.anaconda.org/conda-forge/osx-arm64/yaml-0.2.5-h925e9cb_3.conda
+ - pypi: https://files.pythonhosted.org/packages/5d/96/5238b1efc5922ddbdc9b0db9243152c09777804fb7c02ad1741eb18a11c0/coverage-7.13.4-cp312-cp312-macosx_11_0_arm64.whl
+ - pypi: https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
+ - pypi: https://files.pythonhosted.org/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl
+packages:
+- conda: https://conda.anaconda.org/conda-forge/linux-64/_libgcc_mutex-0.1-conda_forge.tar.bz2
+ sha256: fe51de6107f9edc7aa4f786a70f4a883943bc9d39b3bb7307c04c41410990726
+ md5: d7c89558ba9fa0495403155b64376d81
+ license: None
+ purls: []
+ size: 2562
+ timestamp: 1578324546067
+- conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda
+ build_number: 20
+ sha256: 1dd3fffd892081df9726d7eb7e0dea6198962ba775bd88842135a4ddb4deb3c9
+ md5: a9f577daf3de00bca7c3c76c0ecbd1de
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgomp >=7.5.0
+ constrains:
+ - openmp_impl <0.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 28948
+ timestamp: 1770939786096
+- conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-2_gnu.tar.bz2
+ build_number: 16
+ sha256: fbe2c5e56a653bebb982eda4876a9178aedfc2b545f25d0ce9c4c0b508253d22
+ md5: 73aaf86a425cc6e73fcf236a5a46396d
+ depends:
+ - _libgcc_mutex 0.1 conda_forge
+ - libgomp >=7.5.0
+ constrains:
+ - openmp_impl 9999
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 23621
+ timestamp: 1650670423406
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/_openmp_mutex-4.5-7_kmp_llvm.conda
+ build_number: 7
+ sha256: 7acaa2e0782cad032bdaf756b536874346ac1375745fb250e9bdd6a48a7ab3cd
+ md5: a44032f282e7d2acdeb1c240308052dd
+ depends:
+ - llvm-openmp >=9.0.1
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 8325
+ timestamp: 1764092507920
+- pypi: https://files.pythonhosted.org/packages/18/a6/907a406bb7d359e6a63f99c313846d9eec4f7e6f7437809e03aa00fa3074/absl_py-2.4.0-py3-none-any.whl
+ name: absl-py
+ version: 2.4.0
+ sha256: 88476fd881ca8aab94ffa78b7b6c632a782ab3ba1cd19c9bd423abc4fb4cd28d
+ requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/noarch/affine-2.4.0-pyhd8ed1ab_1.conda
+ sha256: 0deeaf0c001d5543719db9b2686bc1920c86c7e142f9bec74f35e1ce611b1fc2
+ md5: 8c4061f499edec6b8ac7000f6d586829
+ depends:
+ - python >=3.9
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/affine?source=hash-mapping
+ size: 19164
+ timestamp: 1733762153202
+- pypi: https://files.pythonhosted.org/packages/67/84/d844b79acd9fe15ded60b614b7df04a12fad854ee1fbb8415d726ab1beeb/aiobotocore-3.1.1-py3-none-any.whl
+ name: aiobotocore
+ version: 3.1.1
+ sha256: a4e12a3bd099cd19dc2b2e9fe01a807131b46ebd0f83f509bda3cb243e988c32
+ requires_dist:
+ - aiohttp>=3.12.0,<4.0.0
+ - aioitertools>=0.5.1,<1.0.0
+ - botocore>=1.41.0,<1.42.31
+ - python-dateutil>=2.1,<3.0.0
+ - jmespath>=0.7.1,<2.0.0
+ - multidict>=6.0.0,<7.0.0
+ - typing-extensions>=4.14.0,<5.0.0 ; python_full_version < '3.11'
+ - wrapt>=1.10.10,<3.0.0
+ - httpx>=0.25.1,<0.29 ; extra == 'httpx'
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/noarch/aiohappyeyeballs-2.6.1-pyhd8ed1ab_0.conda
+ sha256: 7842ddc678e77868ba7b92a726b437575b23aaec293bca0d40826f1026d90e27
+ md5: 18fd895e0e775622906cdabfc3cf0fb4
+ depends:
+ - python >=3.9
+ license: PSF-2.0
+ license_family: PSF
+ purls:
+ - pkg:pypi/aiohappyeyeballs?source=hash-mapping
+ size: 19750
+ timestamp: 1741775303303
+- conda: https://conda.anaconda.org/conda-forge/linux-64/aiohttp-3.13.3-py310h31b6992_0.conda
+ sha256: f1ef58f20700d0e7c7bd658168f887ce6a3c27b73c3f6cef22bd40a7d7879634
+ md5: 23185f6276719a34510a58f684d618e1
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - aiohappyeyeballs >=2.5.0
+ - aiosignal >=1.4.0
+ - async-timeout >=4.0,<6.0
+ - attrs >=17.3.0
+ - frozenlist >=1.1.1
+ - libgcc >=14
+ - multidict >=4.5,<7.0
+ - propcache >=0.2.0
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - yarl >=1.17.0,<2.0
+ license: MIT AND Apache-2.0
+ license_family: Apache
+ purls:
+ - pkg:pypi/aiohttp?source=hash-mapping
+ size: 892240
+ timestamp: 1767525402926
+- pypi: https://files.pythonhosted.org/packages/10/a1/510b0a7fadc6f43a6ce50152e69dbd86415240835868bb0bd9b5b88b1e06/aioitertools-0.13.0-py3-none-any.whl
+ name: aioitertools
+ version: 0.13.0
+ sha256: 0be0292b856f08dfac90e31f4739432f4cb6d7520ab9eb73e143f4f2fa5259be
+ requires_dist:
+ - typing-extensions>=4.0 ; python_full_version < '3.10'
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/noarch/aiosignal-1.4.0-pyhd8ed1ab_0.conda
+ sha256: 8dc149a6828d19bf104ea96382a9d04dae185d4a03cc6beb1bc7b84c428e3ca2
+ md5: 421a865222cd0c9d83ff08bc78bf3a61
+ depends:
+ - frozenlist >=1.1.0
+ - python >=3.9
+ - typing_extensions >=4.2
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/aiosignal?source=hash-mapping
+ size: 13688
+ timestamp: 1751626573984
+- conda: https://conda.anaconda.org/conda-forge/noarch/annotated-types-0.7.0-pyhd8ed1ab_1.conda
+ sha256: e0ea1ba78fbb64f17062601edda82097fcf815012cf52bb704150a2668110d48
+ md5: 2934f256a8acfe48f6ebb4fce6cde29c
+ depends:
+ - python >=3.9
+ - typing-extensions >=4.0.0
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/annotated-types?source=hash-mapping
+ size: 18074
+ timestamp: 1733247158254
+- conda: https://conda.anaconda.org/conda-forge/linux-64/aom-3.9.1-hac33072_0.conda
+ sha256: b08ef033817b5f9f76ce62dfcac7694e7b6b4006420372de22494503decac855
+ md5: 346722a0be40f6edc53f12640d301338
+ depends:
+ - libgcc-ng >=12
+ - libstdcxx-ng >=12
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 2706396
+ timestamp: 1718551242397
+- conda: https://conda.anaconda.org/conda-forge/noarch/appdirs-1.4.4-pyhd8ed1ab_1.conda
+ sha256: 5b9ef6d338525b332e17c3ed089ca2f53a5d74b7a7b432747d29c6466e39346d
+ md5: f4e90937bbfc3a4a92539545a37bb448
+ depends:
+ - python >=3.9
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/appdirs?source=hash-mapping
+ size: 14835
+ timestamp: 1733754069532
+- conda: https://conda.anaconda.org/conda-forge/noarch/asciitree-0.3.3-py_2.tar.bz2
+ sha256: b3e9369529fe7d721b66f18680ff4b561e20dbf6507e209e1f60eac277c97560
+ md5: c0481c9de49f040272556e2cedf42816
+ depends:
+ - python
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/asciitree?source=hash-mapping
+ size: 6164
+ timestamp: 1531050741142
+- pypi: https://files.pythonhosted.org/packages/2b/03/13dde6512ad7b4557eb792fbcf0c653af6076b81e5941d36ec61f7ce6028/astunparse-1.6.3-py2.py3-none-any.whl
+ name: astunparse
+ version: 1.6.3
+ sha256: c2652417f2c8b5bb325c885ae329bdf3f86424075c4fd1a128674bc6fba4b8e8
+ requires_dist:
+ - wheel>=0.23.0,<1.0
+ - six>=1.6.1,<2.0
+- conda: https://conda.anaconda.org/conda-forge/noarch/async-timeout-5.0.1-pyhcf101f3_2.conda
+ sha256: 6638b68ab2675d0bed1f73562a4e75a61863b903be1538282cddb56c8e8f75bd
+ md5: 0d0ef7e4a0996b2c4ac2175a12b3bf69
+ depends:
+ - python >=3.10
+ - python
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/async-timeout?source=hash-mapping
+ size: 13559
+ timestamp: 1767290444597
+- conda: https://conda.anaconda.org/conda-forge/linux-64/attr-2.5.2-h39aace5_0.conda
+ sha256: a9c114cbfeda42a226e2db1809a538929d2f118ef855372293bd188f71711c48
+ md5: 791365c5f65975051e4e017b5da3abf5
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ license: GPL-2.0-or-later
+ license_family: GPL
+ purls: []
+ size: 68072
+ timestamp: 1756738968573
+- conda: https://conda.anaconda.org/conda-forge/noarch/attrs-25.4.0-pyhcf101f3_1.conda
+ sha256: c13d5e42d187b1d0255f591b7ce91201d4ed8a5370f0d986707a802c20c9d32f
+ md5: 537296d57ea995666c68c821b00e360b
+ depends:
+ - python >=3.10
+ - python
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/attrs?source=compressed-mapping
+ size: 64759
+ timestamp: 1764875182184
+- pypi: https://files.pythonhosted.org/packages/b9/fa/123043af240e49752f1c4bd24da5053b6bd00cad78c2be53c0d1e8b975bc/backports.tarfile-1.2.0-py3-none-any.whl
+ name: backports-tarfile
+ version: 1.2.0
+ sha256: 77e284d754527b01fb1e6fa8a1afe577858ebe4e9dad8919e34c862cb399bc34
+ requires_dist:
+ - sphinx>=3.5 ; extra == 'docs'
+ - jaraco-packaging>=9.3 ; extra == 'docs'
+ - rst-linker>=1.9 ; extra == 'docs'
+ - furo ; extra == 'docs'
+ - sphinx-lint ; extra == 'docs'
+ - pytest>=6,!=8.1.* ; extra == 'testing'
+ - pytest-checkdocs>=2.4 ; extra == 'testing'
+ - pytest-cov ; extra == 'testing'
+ - pytest-enabler>=2.2 ; extra == 'testing'
+ - jaraco-test ; extra == 'testing'
+ - pytest!=8.0.* ; extra == 'testing'
+ requires_python: '>=3.8'
+- conda: https://conda.anaconda.org/conda-forge/noarch/blinker-1.9.0-pyhff2d567_0.conda
+ sha256: f7efd22b5c15b400ed84a996d777b6327e5c402e79e3c534a7e086236f1eb2dc
+ md5: 42834439227a4551b939beeeb8a4b085
+ depends:
+ - python >=3.9
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/blinker?source=hash-mapping
+ size: 13934
+ timestamp: 1731096548765
+- conda: https://conda.anaconda.org/conda-forge/linux-64/blosc-1.21.6-he440d0b_1.conda
+ sha256: e7af5d1183b06a206192ff440e08db1c4e8b2ca1f8376ee45fb2f3a85d4ee45d
+ md5: 2c2fae981fd2afd00812c92ac47d023d
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - libstdcxx >=13
+ - libzlib >=1.3.1,<2.0a0
+ - lz4-c >=1.10.0,<1.11.0a0
+ - snappy >=1.2.1,<1.3.0a0
+ - zstd >=1.5.6,<1.6.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 48427
+ timestamp: 1733513201413
+- conda: https://conda.anaconda.org/conda-forge/noarch/bokeh-2.4.3-pyhd8ed1ab_3.tar.bz2
+ sha256: f37e33fb11ae76ff07ce726a3dbdf4cd26ffff1b52c126d2d2d136669d6b919f
+ md5: e4c6e6d99add99cede5328d811cacb21
+ depends:
+ - jinja2 >=2.9
+ - numpy >=1.11.3
+ - packaging >=16.8
+ - pillow >=7.1.0
+ - python >=3.7
+ - pyyaml >=3.10
+ - tornado >=5.1
+ - typing_extensions >=3.10.0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/bokeh?source=hash-mapping
+ size: 13940985
+ timestamp: 1660586705876
+- pypi: https://files.pythonhosted.org/packages/3d/8d/6d7b016383b1f74dd93611b1c5078bbaddaca901553ab886dcda87cae365/botocore-1.42.30-py3-none-any.whl
+ name: botocore
+ version: 1.42.30
+ sha256: 97070a438cac92430bb7b65f8ebd7075224f4a289719da4ee293d22d1e98db02
+ requires_dist:
+ - jmespath>=0.7.1,<2.0.0
+ - python-dateutil>=2.1,<3.0.0
+ - urllib3>=1.25.4,<1.27 ; python_full_version < '3.10'
+ - urllib3>=1.25.4,!=2.2.0,<3 ; python_full_version >= '3.10'
+ - awscrt==0.29.2 ; extra == 'crt'
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/noarch/branca-0.8.2-pyhd8ed1ab_0.conda
+ sha256: 1acf87c77d920edd098ddc91fa785efc10de871465dee0f463815b176e019e8b
+ md5: 1fcdf88e7a8c296d3df8409bf0690db4
+ depends:
+ - jinja2 >=3
+ - python >=3.10
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/branca?source=hash-mapping
+ size: 30176
+ timestamp: 1759755695447
+- conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.1.0-hb03c661_4.conda
+ sha256: 294526a54fa13635341729f250d0b1cf8f82cad1e6b83130304cbf3b6d8b74cc
+ md5: eaf3fbd2aa97c212336de38a51fe404e
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - brotli-bin 1.1.0 hb03c661_4
+ - libbrotlidec 1.1.0 hb03c661_4
+ - libbrotlienc 1.1.0 hb03c661_4
+ - libgcc >=14
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 19883
+ timestamp: 1756599394934
+- conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.1.0-hb03c661_4.conda
+ sha256: 444903c6e5c553175721a16b7c7de590ef754a15c28c99afbc8a963b35269517
+ md5: ca4ed8015764937c81b830f7f5b68543
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libbrotlidec 1.1.0 hb03c661_4
+ - libbrotlienc 1.1.0 hb03c661_4
+ - libgcc >=14
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 19615
+ timestamp: 1756599385418
+- conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-python-1.1.0-py310hea6c23e_4.conda
+ sha256: 29f24d4a937c3a7f4894d6be9d9f9604adbb5506891f0f37bbb7e2dc8fa6bc0a
+ md5: 6ef43db290647218e1e04c2601675bff
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libstdcxx >=14
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ constrains:
+ - libbrotlicommon 1.1.0 hb03c661_4
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/brotli?source=hash-mapping
+ size: 353838
+ timestamp: 1756599456833
+- conda: https://conda.anaconda.org/conda-forge/linux-64/brunsli-0.1-he3183e4_1.conda
+ sha256: fddad9bb57ee7ec619a5cf4591151578a2501c3bf8cb3b4b066ac5b54c85a4dd
+ md5: 799ebfe432cb3949e246b69278ef851c
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libbrotlicommon >=1.1.0,<1.2.0a0
+ - libbrotlidec >=1.1.0,<1.2.0a0
+ - libbrotlienc >=1.1.0,<1.2.0a0
+ - libgcc >=14
+ - libstdcxx >=14
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 168813
+ timestamp: 1757453968120
+- conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_8.conda
+ sha256: c30daba32ddebbb7ded490f0e371eae90f51e72db620554089103b4a6934b0d5
+ md5: 51a19bba1b8ebfb60df25cde030b7ebc
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ license: bzip2-1.0.6
+ license_family: BSD
+ purls: []
+ size: 260341
+ timestamp: 1757437258798
+- conda: https://conda.anaconda.org/conda-forge/linux-64/bzip2-1.0.8-hda65f42_9.conda
+ sha256: 0b75d45f0bba3e95dc693336fa51f40ea28c980131fec438afb7ce6118ed05f6
+ md5: d2ffd7602c02f2b316fd921d39876885
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ license: bzip2-1.0.6
+ license_family: BSD
+ purls: []
+ size: 260182
+ timestamp: 1771350215188
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/bzip2-1.0.8-hd037594_9.conda
+ sha256: 540fe54be35fac0c17feefbdc3e29725cce05d7367ffedfaaa1bdda234b019df
+ md5: 620b85a3f45526a8bc4d23fd78fc22f0
+ depends:
+ - __osx >=11.0
+ license: bzip2-1.0.6
+ license_family: BSD
+ purls: []
+ size: 124834
+ timestamp: 1771350416561
+- conda: https://conda.anaconda.org/conda-forge/linux-64/c-ares-1.34.6-hb03c661_0.conda
+ sha256: cc9accf72fa028d31c2a038460787751127317dcfa991f8d1f1babf216bb454e
+ md5: 920bb03579f15389b9e512095ad995b7
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 207882
+ timestamp: 1765214722852
+- conda: https://conda.anaconda.org/conda-forge/linux-64/c-blosc2-2.19.1-h4cfbee9_0.conda
+ sha256: ebd0cc82efa5d5dd386f546b75db357d990b91718e4d7788740f4fadc5dfd5c9
+ md5: 041ee44c15d1efdc84740510796425df
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libstdcxx >=14
+ - lz4-c >=1.10.0,<1.11.0a0
+ - zlib-ng >=2.2.4,<2.3.0a0
+ - zstd >=1.5.7,<1.6.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 346946
+ timestamp: 1752777187815
+- conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.1.4-hbd8a1cb_0.conda
+ sha256: b5974ec9b50e3c514a382335efa81ed02b05906849827a34061c496f4defa0b2
+ md5: bddacf101bb4dd0e51811cb69c7790e2
+ depends:
+ - __unix
+ license: ISC
+ purls: []
+ size: 146519
+ timestamp: 1767500828366
+- conda: https://conda.anaconda.org/conda-forge/noarch/cached-property-1.5.2-hd8ed1ab_1.tar.bz2
+ noarch: python
+ sha256: 561e6660f26c35d137ee150187d89767c988413c978e1b712d53f27ddf70ea17
+ md5: 9b347a7ec10940d3f7941ff6c460b551
+ depends:
+ - cached_property >=1.5.2,<1.5.3.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 4134
+ timestamp: 1615209571450
+- conda: https://conda.anaconda.org/conda-forge/noarch/cached_property-1.5.2-pyha770c72_1.tar.bz2
+ sha256: 6dbf7a5070cc43d90a1e4c2ec0c541c69d8e30a0e25f50ce9f6e4a432e42c5d7
+ md5: 576d629e47797577ab0f1b351297ef4a
+ depends:
+ - python >=3.6
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/cached-property?source=hash-mapping
+ size: 11065
+ timestamp: 1615209567874
+- conda: https://conda.anaconda.org/conda-forge/noarch/certifi-2026.1.4-pyhd8ed1ab_0.conda
+ sha256: 110338066d194a715947808611b763857c15458f8b3b97197387356844af9450
+ md5: eacc711330cd46939f66cd401ff9c44b
+ depends:
+ - python >=3.10
+ license: ISC
+ purls:
+ - pkg:pypi/certifi?source=compressed-mapping
+ size: 150969
+ timestamp: 1767500900768
+- conda: https://conda.anaconda.org/conda-forge/linux-64/cffi-2.0.0-py310he7384ee_1.conda
+ sha256: bf76ead6d59b70f3e901476a73880ac92011be63b151972d135eec55bbbe6091
+ md5: 803e2d778b8dcccdc014127ec5001681
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libffi >=3.5.2,<3.6.0a0
+ - libgcc >=14
+ - pycparser
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/cffi?source=hash-mapping
+ size: 244766
+ timestamp: 1761203011221
+- pypi: https://files.pythonhosted.org/packages/db/3c/33bac158f8ab7f89b2e59426d5fe2e4f63f7ed25df84c036890172b412b5/cfgv-3.5.0-py2.py3-none-any.whl
+ name: cfgv
+ version: 3.5.0
+ sha256: a8dc6b26ad22ff227d2634a65cb388215ce6cc96bbcc5cfde7641ae87e8dacc0
+ requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/cftime-1.6.5-py310hf779ad0_1.conda
+ sha256: d2c77b7ba1d49e859a047941bf7c382f66dbda1e7041e24bcac78f8e61bb0570
+ md5: 702cd87d652ec1956ee503b9e16bb1c0
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - numpy >=1.21,<3
+ - numpy >=1.21.2
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/cftime?source=hash-mapping
+ size: 430152
+ timestamp: 1768510929678
+- conda: https://conda.anaconda.org/conda-forge/linux-64/charls-2.4.2-h59595ed_0.conda
+ sha256: 18f1c43f91ccf28297f92b094c2c8dbe9c6e8241c0d3cbd6cda014a990660fdd
+ md5: 4336bd67920dd504cd8c6761d6a99645
+ depends:
+ - libgcc-ng >=12
+ - libstdcxx-ng >=12
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 150272
+ timestamp: 1684262827894
+- conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.4-pyhd8ed1ab_0.conda
+ sha256: b32f8362e885f1b8417bac2b3da4db7323faa12d5db62b7fd6691c02d60d6f59
+ md5: a22d1fd9bf98827e280a02875d9a007a
+ depends:
+ - python >=3.10
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/charset-normalizer?source=hash-mapping
+ size: 50965
+ timestamp: 1760437331772
+- conda: https://conda.anaconda.org/conda-forge/noarch/click-8.3.1-pyh8f84b5b_1.conda
+ sha256: 38cfe1ee75b21a8361c8824f5544c3866f303af1762693a178266d7f198e8715
+ md5: ea8a6c3256897cc31263de9f455e25d9
+ depends:
+ - python >=3.10
+ - __unix
+ - python
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/click?source=hash-mapping
+ size: 97676
+ timestamp: 1764518652276
+- conda: https://conda.anaconda.org/conda-forge/noarch/click-plugins-1.1.1.2-pyhd8ed1ab_0.conda
+ sha256: ba1ee6e2b2be3da41d70d0d51d1159010de900aa3f33fceaea8c52e9bd30a26e
+ md5: e9b05deb91c013e5224672a4ba9cf8d1
+ depends:
+ - click >=4.0
+ - python >=3.9
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/click-plugins?source=hash-mapping
+ size: 12683
+ timestamp: 1750848314962
+- conda: https://conda.anaconda.org/conda-forge/noarch/cligj-0.7.2-pyhd8ed1ab_2.conda
+ sha256: 1a52ae1febfcfb8f56211d1483a1ac4419b0028b7c3e9e61960a298978a42396
+ md5: 55c7804f428719241a90b152016085a1
+ depends:
+ - click >=4.0
+ - python >=3.9,<4.0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/cligj?source=hash-mapping
+ size: 12521
+ timestamp: 1733750069604
+- conda: https://conda.anaconda.org/conda-forge/noarch/cloudpickle-3.1.2-pyhcf101f3_1.conda
+ sha256: 4c287c2721d8a34c94928be8fe0e9a85754e90189dd4384a31b1806856b50a67
+ md5: 61b8078a0905b12529abc622406cb62c
+ depends:
+ - python >=3.10
+ - python
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/cloudpickle?source=compressed-mapping
+ size: 27353
+ timestamp: 1765303462831
+- conda: https://conda.anaconda.org/conda-forge/linux-64/contourpy-1.3.2-py310h3788b33_0.conda
+ sha256: 5231c1b68e01a9bc9debabc077a6fb48c4395206d59f40a4598d1d5e353e11d8
+ md5: b6420d29123c7c823de168f49ccdfe6a
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - libstdcxx >=13
+ - numpy >=1.23
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/contourpy?source=hash-mapping
+ size: 261280
+ timestamp: 1744743236964
+- pypi: https://files.pythonhosted.org/packages/35/fb/05b9830c2e8275ebc031e0019387cda99113e62bb500ab328bb72578183b/coverage-7.13.2-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl
+ name: coverage
+ version: 7.13.2
+ sha256: ca9566769b69a5e216a4e176d54b9df88f29d750c5b78dbb899e379b4e14b30c
+ requires_dist:
+ - tomli ; python_full_version <= '3.11' and extra == 'toml'
+ requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/5d/96/5238b1efc5922ddbdc9b0db9243152c09777804fb7c02ad1741eb18a11c0/coverage-7.13.4-cp312-cp312-macosx_11_0_arm64.whl
+ name: coverage
+ version: 7.13.4
+ sha256: 40aa8808140e55dc022b15d8aa7f651b6b3d68b365ea0398f1441e0b04d859c3
+ requires_dist:
+ - tomli ; python_full_version <= '3.11' and extra == 'toml'
+ requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/5d/a0/2ea570925524ef4e00bb6c82649f5682a77fac5ab910a65c9284de422600/coverage-7.13.4-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl
+ name: coverage
+ version: 7.13.4
+ sha256: 2c048ea43875fbf8b45d476ad79f179809c590ec7b79e2035c662e7afa3192e3
+ requires_dist:
+ - tomli ; python_full_version <= '3.11' and extra == 'toml'
+ requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/cryptography-46.0.4-py310hb288b08_0.conda
+ sha256: db1699f23d4ea62d8f7f8dcf24b42e4c23243617697edf12c7e7cfe8ea156f0f
+ md5: 2d41874f8cb5ce02bc771ee445049bc5
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - cffi >=1.14
+ - libgcc >=14
+ - openssl >=3.5.5,<4.0a0
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - typing_extensions >=4.13.2
+ constrains:
+ - __glibc >=2.17
+ license: Apache-2.0 AND BSD-3-Clause AND PSF-2.0 AND MIT
+ license_family: BSD
+ purls:
+ - pkg:pypi/cryptography?source=hash-mapping
+ size: 1667518
+ timestamp: 1769650727876
+- conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda
+ sha256: bb47aec5338695ff8efbddbc669064a3b10fe34ad881fb8ad5d64fbfa6910ed1
+ md5: 4c2a8fef270f6c69591889b93f9f55c1
+ depends:
+ - python >=3.10
+ - python
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/cycler?source=hash-mapping
+ size: 14778
+ timestamp: 1764466758386
+- conda: https://conda.anaconda.org/conda-forge/linux-64/cytoolz-1.1.0-py310h7c4b9e2_1.conda
+ sha256: 9cbeb77ad9e23c7ffc85399fd41a85a61d511a1e28d0ff2132baf4b116983596
+ md5: aa27c9572fd9f548f911300dc6305bf4
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - toolz >=0.10.0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/cytoolz?source=hash-mapping
+ size: 565369
+ timestamp: 1760905957492
+- conda: https://conda.anaconda.org/conda-forge/noarch/dacite-1.9.2-pyhd8ed1ab_0.conda
+ sha256: 5c51049725f9f120af99a993c3e4ab83352b0f7397ddbc79f27272fcafddfa5b
+ md5: ba49ecddabe4ced654b6d16cccc2d231
+ depends:
+ - dataclasses
+ - python >=3.9
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/dacite?source=hash-mapping
+ size: 21000
+ timestamp: 1738772021863
+- conda: https://conda.anaconda.org/conda-forge/noarch/dask-2023.3.0-pyhd8ed1ab_0.conda
+ sha256: 7e95d34946ac4d156d3f0b72d7165c19185e0c23c053084665802ae1a18dc218
+ md5: aab5cea04004860e804de5bb3337f183
+ depends:
+ - bokeh >=2.4.2,<3
+ - cytoolz >=0.8.2
+ - dask-core >=2023.3.0,<2023.3.1.0a0
+ - distributed >=2023.3.0,<2023.3.1.0a0
+ - jinja2 >=2.10.3
+ - lz4
+ - numpy >=1.21
+ - pandas >=1.3
+ - python >=3.8
+ constrains:
+ - openssl !=1.1.1e
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 7063
+ timestamp: 1677718268307
+- conda: https://conda.anaconda.org/conda-forge/noarch/dask-core-2023.3.0-pyhd8ed1ab_0.conda
+ sha256: 3e9f7d7180f0ffbf2e3014bb7b7aa010f65a90f701a1e43471487fbf03ceeca7
+ md5: 34437340f37faafad7a6287d3b624f60
+ depends:
+ - click >=7.0
+ - cloudpickle >=1.1.1
+ - fsspec >=0.6.0
+ - packaging >=20.0
+ - partd >=1.2.0
+ - python >=3.8
+ - pyyaml >=5.3.1
+ - toolz >=0.8.2
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/dask?source=hash-mapping
+ size: 836744
+ timestamp: 1677707157572
+- conda: https://conda.anaconda.org/conda-forge/noarch/dataclasses-0.8-pyhc8e2a94_3.tar.bz2
+ sha256: 63a83e62e0939bc1ab32de4ec736f6403084198c4639638b354a352113809c92
+ md5: a362b2124b06aad102e2ee4581acee7d
+ depends:
+ - python >=3.7
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/dataclasses?source=hash-mapping
+ size: 9870
+ timestamp: 1628958582931
+- conda: https://conda.anaconda.org/conda-forge/linux-64/dav1d-1.2.1-hd590300_0.conda
+ sha256: 22053a5842ca8ee1cf8e1a817138cdb5e647eb2c46979f84153f6ad7bde73020
+ md5: 418c6ca5929a611cbd69204907a83995
+ depends:
+ - libgcc-ng >=12
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 760229
+ timestamp: 1685695754230
+- conda: https://conda.anaconda.org/conda-forge/noarch/decorator-5.2.1-pyhd8ed1ab_0.conda
+ sha256: c17c6b9937c08ad63cb20a26f403a3234088e57d4455600974a0ce865cb14017
+ md5: 9ce473d1d1be1cc3810856a48b3fab32
+ depends:
+ - python >=3.9
+ license: BSD-2-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/decorator?source=hash-mapping
+ size: 14129
+ timestamp: 1740385067843
+- pypi: https://files.pythonhosted.org/packages/ab/28/960c311aae084deef57ece41aac13cb359b06ce31b7771139e79c394a1b7/deptry-0.24.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ name: deptry
+ version: 0.24.0
+ sha256: dd22fa2dbbdf4b38061ca9504f2a6ce41ec14fa5c9fe9b0b763ccc1275efebd5
+ requires_dist:
+ - click>=8.0.0,<9
+ - colorama>=0.4.6 ; sys_platform == 'win32'
+ - packaging>=23.2
+ - requirements-parser>=0.11.0,<1
+ - tomli>=2.0.1 ; python_full_version < '3.11'
+ requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl
+ name: distlib
+ version: 0.4.0
+ sha256: 9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16
+- conda: https://conda.anaconda.org/conda-forge/noarch/distributed-2023.3.0-pyhd8ed1ab_0.conda
+ sha256: 5d4b493172026813cd0774f94efdc54651a2708fb52cfc55574fa9504b50c50c
+ md5: 6ca8ed418961a91d76965268b6f4aa5b
+ depends:
+ - click >=7.0
+ - cloudpickle >=1.5.0
+ - cytoolz >=0.10.1
+ - dask-core >=2023.3.0,<2023.3.1.0a0
+ - jinja2 >=2.10.3
+ - locket >=1.0.0
+ - msgpack-python >=1.0.0
+ - packaging >=20.0
+ - psutil >=5.7.0
+ - python >=3.8
+ - pyyaml >=5.3.1
+ - sortedcontainers >=2.0.5
+ - tblib >=1.6.0
+ - toolz >=0.10.0
+ - tornado >=6.0.3
+ - urllib3 >=1.24.3
+ - zict >=2.1.0
+ constrains:
+ - openssl !=1.1.1e
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/distributed?source=hash-mapping
+ size: 749994
+ timestamp: 1677711209126
+- conda: https://conda.anaconda.org/conda-forge/noarch/docker-pycreds-0.4.0-py_0.tar.bz2
+ sha256: 2ba7e3e4f75e07b42246b4ba8569c983ecbdcda47b1b900632858a23d91826f2
+ md5: c69f19038efee4eb534623610d0c2053
+ depends:
+ - python
+ - six >=1.4.0
+ license: Apache-2.0
+ license_family: Apache
+ purls:
+ - pkg:pypi/docker-pycreds?source=hash-mapping
+ size: 11445
+ timestamp: 1551105257829
+- pypi: https://files.pythonhosted.org/packages/02/10/5da547df7a391dcde17f59520a231527b8571e6f46fc8efb02ccb370ab12/docutils-0.22.4-py3-none-any.whl
+ name: docutils
+ version: 0.22.4
+ sha256: d0013f540772d1420576855455d050a2180186c91c15779301ac2ccb3eeb68de
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/noarch/eval_type_backport-0.3.1-pyha770c72_0.conda
+ sha256: 454f03ac61295b2bca852913af54248cfb9c1a9d2e057f3b5574d552255cda61
+ md5: 9cb8eae2a1f3e4a2cb8c53559abf6d75
+ depends:
+ - python >=3.10
+ constrains:
+ - eval-type-backport >=0.3.1,<0.3.2.0a0
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/eval-type-backport?source=hash-mapping
+ size: 12244
+ timestamp: 1764679328643
+- pypi: https://files.pythonhosted.org/packages/8a/0e/97c33bf5009bdbac74fd2beace167cab3f978feb69cc36f1ef79360d6c4e/exceptiongroup-1.3.1-py3-none-any.whl
+ name: exceptiongroup
+ version: 1.3.1
+ sha256: a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598
+ requires_dist:
+ - typing-extensions>=4.6.0 ; python_full_version < '3.13'
+ - pytest>=6 ; extra == 'test'
+ requires_python: '>=3.7'
+- conda: https://conda.anaconda.org/conda-forge/noarch/fasteners-0.19-pyhd8ed1ab_1.conda
+ sha256: 42fb170778b47303e82eddfea9a6d1e1b8af00c927cd5a34595eaa882b903a16
+ md5: dbe9d42e94b5ff7af7b7893f4ce052e7
+ depends:
+ - python >=3.9
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/fasteners?source=hash-mapping
+ size: 20711
+ timestamp: 1734943237791
+- pypi: https://files.pythonhosted.org/packages/b5/36/7fb70f04bf00bc646cd5bb45aa9eddb15e19437a28b8fb2b4a5249fac770/filelock-3.20.3-py3-none-any.whl
+ name: filelock
+ version: 3.20.3
+ sha256: 4b0dda527ee31078689fc205ec4f1c1bf7d56cf88b6dc9426c4f230e46c2dce1
+ requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/fiona-1.10.1-py310hea6c23e_4.conda
+ sha256: 0c2a717fb0e4abfec0403a9d5d64d913d199371ec8a8fdf38e5332c8b68dab4b
+ md5: 5bbb62bd0d1f4e5908b25bfa0dfdbbc1
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - attrs >=19.2.0
+ - click >=8.0,<9.dev0
+ - click-plugins >=1.0
+ - cligj >=0.5
+ - libgcc >=14
+ - libgdal-core >=3.10.3,<3.11.0a0
+ - libstdcxx >=14
+ - pyparsing
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - shapely
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/fiona?source=hash-mapping
+ size: 1174494
+ timestamp: 1764874591583
+- pypi: https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl
+ name: flatbuffers
+ version: 25.12.19
+ sha256: 7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4
+- conda: https://conda.anaconda.org/conda-forge/noarch/folium-0.20.0-pyhd8ed1ab_0.conda
+ sha256: 782fa186d7677fd3bc1ff7adb4cc3585f7d2c7177c30bcbce21f8c177135c520
+ md5: a6997a7dcd6673c0692c61dfeaea14ab
+ depends:
+ - branca >=0.6.0
+ - jinja2 >=2.9
+ - numpy
+ - python >=3.9
+ - requests
+ - xyzservices
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/folium?source=hash-mapping
+ size: 82665
+ timestamp: 1750113928159
+- conda: https://conda.anaconda.org/conda-forge/linux-64/fonttools-4.61.1-py310h3406613_0.conda
+ sha256: 6dccba7a293b6dbab029da4d921d2d94227c9541152489fc7d7db4ec3c68dff3
+ md5: 24fa891e40acdb1c7f51efd0c5f97084
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - brotli
+ - libgcc >=14
+ - munkres
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - unicodedata2 >=15.1.0
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/fonttools?source=hash-mapping
+ size: 2446291
+ timestamp: 1765632899119
+- conda: https://conda.anaconda.org/conda-forge/linux-64/freetype-2.14.1-ha770c72_0.conda
+ sha256: bf8e4dffe46f7d25dc06f31038cacb01672c47b9f45201f065b0f4d00ab0a83e
+ md5: 4afc585cd97ba8a23809406cd8a9eda8
+ depends:
+ - libfreetype 2.14.1 ha770c72_0
+ - libfreetype6 2.14.1 h73754d4_0
+ license: GPL-2.0-only OR FTL
+ purls: []
+ size: 173114
+ timestamp: 1757945422243
+- conda: https://conda.anaconda.org/conda-forge/linux-64/freexl-2.0.0-h9dce30a_2.conda
+ sha256: c8960e00a6db69b85c16c693ce05484facf20f1a80430552145f652a880e0d2a
+ md5: ecb5d11305b8ba1801543002e69d2f2f
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libexpat >=2.6.4,<3.0a0
+ - libgcc >=13
+ - libiconv >=1.17,<2.0a0
+ - minizip >=4.0.7,<5.0a0
+ license: MPL-1.1
+ license_family: MOZILLA
+ purls: []
+ size: 59299
+ timestamp: 1734014884486
+- conda: .
+ name: fronts
+ version: 0.1.0
+ build: pyh4616a5c_0
+ subdir: noarch
+ depends:
+ - python >=3.10,<3.14
+ - python *
+ license: CC0-1.0
+- conda: https://conda.anaconda.org/conda-forge/linux-64/frozenlist-1.7.0-py310h9548a50_0.conda
+ sha256: c8abeb6da1e89113049d01c714fcce67e2fcc2853a63b3c40078372a5f66c59f
+ md5: 50e2b335c9da85d4eadaab11cf245415
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libstdcxx >=14
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/frozenlist?source=hash-mapping
+ size: 54180
+ timestamp: 1752167428701
+- conda: https://conda.anaconda.org/conda-forge/noarch/fsspec-2026.2.0-pyhd8ed1ab_0.conda
+ sha256: 239b67edf1c5e5caed52cf36e9bed47cb21b37721779828c130e6b3fd9793c1b
+ md5: 496c6c9411a6284addf55c898d6ed8d7
+ depends:
+ - python >=3.10
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/fsspec?source=compressed-mapping
+ size: 148757
+ timestamp: 1770387898414
+- pypi: https://files.pythonhosted.org/packages/1d/33/f1c6a276de27b7d7339a34749cc33fa87f077f921969c47185d34a887ae2/gast-0.7.0-py3-none-any.whl
+ name: gast
+ version: 0.7.0
+ sha256: 99cbf1365633a74099f69c59bd650476b96baa5ef196fec88032b00b31ba36f7
+ requires_python: '>=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*'
+- conda: https://conda.anaconda.org/conda-forge/noarch/gcsfs-2026.2.0-pyhd8ed1ab_0.conda
+ sha256: 2acd9277d297bc64f32b72bea67f4c7b541390d74b6446a802a1bfa9a62b0f56
+ md5: 7b2f6cee85093591aaabfbd87c9e52c2
+ depends:
+ - aiohttp
+ - decorator >4.1.2
+ - fsspec 2026.2.0
+ - google-auth >=1.2
+ - google-auth-oauthlib
+ - google-cloud-storage >=3.9.0
+ - google-cloud-storage-control
+ - python >=3.10
+ - requests
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/gcsfs?source=hash-mapping
+ size: 54216
+ timestamp: 1770424255343
+- conda: https://conda.anaconda.org/conda-forge/noarch/geopandas-0.14.4-pyhd8ed1ab_0.conda
+ sha256: a08d4c641dbf7b27b1195c270816cea801edae74dd609012d03ae5ad35c9dccc
+ md5: acc01facf6f915b6289a064957a58cc1
+ depends:
+ - fiona >=1.8.21
+ - folium
+ - geopandas-base 0.14.4 pyha770c72_0
+ - mapclassify >=2.4.0
+ - matplotlib-base
+ - python >=3.9
+ - rtree
+ - xyzservices
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 7691
+ timestamp: 1714335630563
+- conda: https://conda.anaconda.org/conda-forge/noarch/geopandas-base-0.14.4-pyha770c72_0.conda
+ sha256: 9dc4b7ee08b60be28a7284104e7147ecf23fcbe3718eeb271712deb92ff3ff06
+ md5: b7a9e8e5865cc474fb0856577898316a
+ depends:
+ - packaging
+ - pandas >=1.4.0
+ - pyproj >=3.3.0
+ - python >=3.9
+ - shapely >=1.8.0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/geopandas?source=hash-mapping
+ size: 1021307
+ timestamp: 1714335625468
+- conda: https://conda.anaconda.org/conda-forge/linux-64/geos-3.14.0-h480dda7_0.conda
+ sha256: c986e3c5fcdb61a34213923b22e5c8859a1012714cba34a4f6292551b67613ed
+ md5: 5dc479effdabf54a0ff240d565287495
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libstdcxx >=14
+ license: LGPL-2.1-only
+ purls: []
+ size: 1977241
+ timestamp: 1755851798617
+- conda: https://conda.anaconda.org/conda-forge/linux-64/geotiff-1.7.4-h239500f_2.conda
+ sha256: 0cd4454921ac0dfbf9d092d7383ba9717e223f9e506bc1ac862c99f98d2a953c
+ md5: b0c42bce162a38b1aa2f6dfb5c412bc4
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - libjpeg-turbo >=3.0.0,<4.0a0
+ - libstdcxx >=13
+ - libtiff >=4.7.0,<4.8.0a0
+ - libzlib >=1.3.1,<2.0a0
+ - proj >=9.6.0,<9.7.0a0
+ - zlib
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 128758
+ timestamp: 1742402413139
+- conda: https://conda.anaconda.org/conda-forge/linux-64/giflib-5.2.2-hd590300_0.conda
+ sha256: aac402a8298f0c0cc528664249170372ef6b37ac39fdc92b40601a6aed1e32ff
+ md5: 3bf7b9fd5a7136126e0234db4b87c8b6
+ depends:
+ - libgcc-ng >=12
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 77248
+ timestamp: 1712692454246
+- conda: https://conda.anaconda.org/conda-forge/noarch/gitdb-4.0.12-pyhd8ed1ab_0.conda
+ sha256: dbbec21a369872c8ebe23cb9a3b9d63638479ee30face165aa0fccc96e93eec3
+ md5: 7c14f3706e099f8fcd47af2d494616cc
+ depends:
+ - python >=3.9
+ - smmap >=3.0.1,<6
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/gitdb?source=hash-mapping
+ size: 53136
+ timestamp: 1735887290843
+- conda: https://conda.anaconda.org/conda-forge/noarch/gitpython-3.1.46-pyhd8ed1ab_0.conda
+ sha256: 8043bcb4f59d17467c6c2f8259e7ded18775de5d62a8375a27718554d9440641
+ md5: 74c0cfdd5359cd2a1f178a4c3d0bd3a5
+ depends:
+ - gitdb >=4.0.1,<5
+ - python >=3.10
+ - typing_extensions >=3.10.0.2
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/gitpython?source=hash-mapping
+ size: 158433
+ timestamp: 1767358832407
+- conda: https://conda.anaconda.org/conda-forge/noarch/google-api-core-2.29.0-pyhd8ed1ab_0.conda
+ sha256: 0f696294c9a117a16e344388347dd9dff644cd8ddb703002169d81f889c176df
+ md5: 7fd8158ff94ccf28a2ac1f534989d698
+ depends:
+ - google-auth >=2.14.1,<3.0.0
+ - googleapis-common-protos >=1.56.2,<2.0.0
+ - proto-plus >=1.25.0,<2.0.0
+ - protobuf >=3.19.5,<7.0.0,!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5
+ - python >=3.10
+ - requests >=2.18.0,<3.0.0
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/google-api-core?source=hash-mapping
+ size: 98400
+ timestamp: 1768122057220
+- conda: https://conda.anaconda.org/conda-forge/noarch/google-auth-2.48.0-pyhcf101f3_0.conda
+ sha256: f9fd7cbfc6cea1b43c9e210f0042c5ca62ded641a83ed6e7c046ef08dfac4583
+ md5: 6e643ba74997c8dddbaa98fc2fc3481b
+ depends:
+ - python >=3.10
+ - pyasn1-modules >=0.2.1
+ - cryptography >=38.0.3
+ - rsa >=3.1.4,<5
+ - aiohttp >=3.6.2,<4.0.0
+ - requests >=2.20.0,<3.0.0
+ - pyopenssl >=20.0.0
+ - pyu2f >=0.1.5
+ - python
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/google-auth?source=compressed-mapping
+ size: 141954
+ timestamp: 1769604366349
+- conda: https://conda.anaconda.org/conda-forge/noarch/google-auth-oauthlib-1.2.2-pyhd8ed1ab_0.conda
+ sha256: 8b9bdddd954f257c234014fda6169fb453b9608af96a5d6faf4110be1ffad30a
+ md5: 8e340e42470ffa0435a2067858b5c743
+ depends:
+ - click >=6.0.0
+ - google-auth >=2.15.0
+ - python >=3.9
+ - requests-oauthlib >=0.7.0
+ license: Apache-2.0
+ license_family: Apache
+ purls:
+ - pkg:pypi/google-auth-oauthlib?source=hash-mapping
+ size: 22086
+ timestamp: 1745360378896
+- conda: https://conda.anaconda.org/conda-forge/noarch/google-cloud-core-2.5.0-pyhcf101f3_1.conda
+ sha256: fcef1d51f6de304a23c19ea6b3114dcab9ce54482d9f506f9a3e0b48be514744
+ md5: 48fcccc0b579087018df0afc332b8bd6
+ depends:
+ - python >=3.10,<3.14
+ - google-api-core >=1.31.6,<3.0.0,!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0
+ - google-auth >=1.25.0,<3.0.0
+ - grpcio >=1.38.0,<2.0.0
+ - grpcio-status >=1.38.0,<2.0.0
+ - python
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/google-cloud-core?source=hash-mapping
+ size: 33593
+ timestamp: 1768561863777
+- conda: https://conda.anaconda.org/conda-forge/noarch/google-cloud-storage-3.9.0-pyhcf101f3_0.conda
+ sha256: a42b7f24ded8a97ae18d46bbdf9485d951346e41c66b18f7f793d521c05a53d5
+ md5: 71aa090f8647c9b9efa63994eaa0dc18
+ depends:
+ - python >=3.10
+ - google-api-core >=2.27.0,<3.0.0
+ - google-auth >=2.26.1,<3.0.0
+ - google-cloud-core >=2.4.2,<3.0.0
+ - google-crc32c >=1.1.3,<2.0.0
+ - google-resumable-media >=2.7.2,<3.0.0
+ - requests >=2.22.0,<3.0.0
+ - protobuf >=3.20.2,<7.0.0
+ - python
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/google-cloud-storage?source=hash-mapping
+ size: 201445
+ timestamp: 1770056609761
+- conda: https://conda.anaconda.org/conda-forge/linux-64/google-cloud-storage-control-1.8.0-py310hff52083_0.conda
+ sha256: e041df3cd38140d174ce89f1ef74496da55d47eec59468ca800b2980c2a6686c
+ md5: efede186ca8e5fdd77349a8c0e82b55f
+ depends:
+ - google-api-core >=1.34.1,<3.0.0,!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.*,!=2.4.*,!=2.5.*,!=2.6.*,!=2.7.*,!=2.8.*,!=2.9.*,!=2.10.*
+ - google-auth >=2.14.1,<3.0.0,!=2.24.0,!=2.25.0
+ - grpc-google-iam-v1 >=0.14.0,<1.0.0
+ - grpcio >=1.33.2,<2.0.0
+ - proto-plus >=1.22.3,<2.0.0
+ - protobuf >=3.20.2,<7.0.0,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/google-cloud-storage-control?source=hash-mapping
+ size: 97581
+ timestamp: 1766178244558
+- conda: https://conda.anaconda.org/conda-forge/linux-64/google-crc32c-1.8.0-py310hf432777_1.conda
+ sha256: 55562adbe097136eed5164de6ce068e22085a9561125755ed8a5599a6d826ab5
+ md5: 91b9f0dccd7a15b36323fb6c3fe12dc1
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libcrc32c >=1.1.2,<1.2.0a0
+ - libgcc >=14
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: Apache-2.0
+ license_family: Apache
+ purls:
+ - pkg:pypi/google-crc32c?source=hash-mapping
+ size: 24207
+ timestamp: 1768549200028
+- pypi: https://files.pythonhosted.org/packages/a3/de/c648ef6835192e6e2cc03f40b19eeda4382c49b5bafb43d88b931c4c74ac/google_pasta-0.2.0-py3-none-any.whl
+ name: google-pasta
+ version: 0.2.0
+ sha256: b32482794a366b5366a32c92a9a9201b107821889935a02b3e51f6b432ea84ed
+ requires_dist:
+ - six
+- conda: https://conda.anaconda.org/conda-forge/noarch/google-resumable-media-2.8.0-pyhd8ed1ab_0.conda
+ sha256: 23d825ed0664a8089c7958bffd819d26e1aba7579695c40dfbdb25a4864d8be6
+ md5: ba7f04ba62be69f9c9fef0c4487c210b
+ depends:
+ - google-crc32c >=1.0.0,<2.0.0
+ - python >=3.10
+ constrains:
+ - aiohttp >=3.6.2,<4.0.0
+ - requests >=2.18.0,<3.0.0
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/google-resumable-media?source=hash-mapping
+ size: 46929
+ timestamp: 1763404726218
+- conda: https://conda.anaconda.org/conda-forge/noarch/googleapis-common-protos-1.72.0-pyhd8ed1ab_0.conda
+ sha256: c09ba4b360a0994430d2fe4a230aa6518cd3e6bfdc51a7af9d35d35a25908bb5
+ md5: 003094932fb90de018f77a273b8a509b
+ depends:
+ - protobuf >=3.20.2,<7.0.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5
+ - python >=3.10
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/googleapis-common-protos?source=compressed-mapping
+ size: 142961
+ timestamp: 1762522289200
+- conda: https://conda.anaconda.org/conda-forge/noarch/googleapis-common-protos-grpc-1.72.0-pyhd8ed1ab_0.conda
+ sha256: 7e7039ac36b2a31149cf287d59dd89761d3c1d3d6e078b2a5c5676b220c97b41
+ md5: a539554411d81653f4cfb3eb8f4dff5e
+ depends:
+ - googleapis-common-protos >=1.72,<1.73.0a0
+ - grpcio >=1.44.0,<2.0.0
+ - python >=3.10
+ license: Apache-2.0
+ license_family: APACHE
+ purls: []
+ size: 15254
+ timestamp: 1762522301108
+- conda: https://conda.anaconda.org/conda-forge/noarch/grpc-google-iam-v1-0.14.3-pyhcf101f3_1.conda
+ sha256: 5649ec4fb9c0240806b4a080899ec5ce42b90692f01b111c88ac81a586773711
+ md5: 2f307997162d8b75d4eb9d86c5c36fbe
+ depends:
+ - python >=3.10
+ - grpcio >=1.44.0,<2.0.0
+ - googleapis-common-protos >=1.56.0,<2.0.0
+ - googleapis-common-protos-grpc
+ - protobuf >=3.20.2,<7.0.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5
+ - python
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/grpc-google-iam-v1?source=hash-mapping
+ size: 31398
+ timestamp: 1768564357665
+- conda: https://conda.anaconda.org/conda-forge/linux-64/grpcio-1.54.3-py310heca2aa9_0.conda
+ sha256: bf9789de2cfb7d48e1bac55829b101f55676588f74cafcc2ecb4fcbb5de613eb
+ md5: 7c82ea96968976ee2eb66a52f2decd8e
+ depends:
+ - libgcc-ng >=12
+ - libgrpc 1.54.3 hb20ce57_0
+ - libstdcxx-ng >=12
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/grpcio?source=hash-mapping
+ size: 755352
+ timestamp: 1690943769542
+- conda: https://conda.anaconda.org/conda-forge/noarch/grpcio-status-1.54.2-pyhd8ed1ab_0.conda
+ sha256: 8420c3da19de43e75383f3c0c93d3f556dc9b26a4bbe08f7fc557ec3c93782d4
+ md5: c7e845594733f1cb1f0f6d60a71d0dae
+ depends:
+ - googleapis-common-protos >=1.5.5
+ - grpcio >=1.54.2
+ - protobuf >=4.21.6
+ - python >=3.6
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/grpcio-status?source=hash-mapping
+ size: 18367
+ timestamp: 1683994834200
+- conda: https://conda.anaconda.org/conda-forge/noarch/h2-4.3.0-pyhcf101f3_0.conda
+ sha256: 84c64443368f84b600bfecc529a1194a3b14c3656ee2e832d15a20e0329b6da3
+ md5: 164fc43f0b53b6e3a7bc7dce5e4f1dc9
+ depends:
+ - python >=3.10
+ - hyperframe >=6.1,<7
+ - hpack >=4.1,<5
+ - python
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/h2?source=hash-mapping
+ size: 95967
+ timestamp: 1756364871835
+- conda: https://conda.anaconda.org/conda-forge/linux-64/h5py-3.15.1-nompi_py310h4aa865e_101.conda
+ sha256: 427fc2540a4728dc80d9f0b464541aed61d35ae9ccafcd7f6bbce499eeaf8ce9
+ md5: 4fccf52eaeb2ae9d9e251623e2b66e63
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - cached-property
+ - hdf5 >=1.14.6,<1.14.7.0a0
+ - libgcc >=14
+ - numpy >=1.21,<3
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/h5py?source=hash-mapping
+ size: 1217205
+ timestamp: 1764016763175
+- conda: https://conda.anaconda.org/conda-forge/linux-64/hdf4-4.2.15-h2a13503_7.conda
+ sha256: 0d09b6dc1ce5c4005ae1c6a19dc10767932ef9a5e9c755cfdbb5189ac8fb0684
+ md5: bd77f8da987968ec3927990495dc22e4
+ depends:
+ - libgcc-ng >=12
+ - libjpeg-turbo >=3.0.0,<4.0a0
+ - libstdcxx-ng >=12
+ - libzlib >=1.2.13,<2.0.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 756742
+ timestamp: 1695661547874
+- conda: https://conda.anaconda.org/conda-forge/linux-64/hdf5-1.14.6-nompi_h1b119a7_105.conda
+ sha256: aa85acd07b8f60d1760c6b3fa91dd8402572766e763f3989c759ecd266ed8e9f
+ md5: d58cd79121dd51128f2a5dab44edf1ea
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libaec >=1.1.4,<2.0a0
+ - libcurl >=8.18.0,<9.0a0
+ - libgcc >=14
+ - libgfortran
+ - libgfortran5 >=14.3.0
+ - libstdcxx >=14
+ - libzlib >=1.3.1,<2.0a0
+ - openssl >=3.5.4,<4.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 3722799
+ timestamp: 1768858199331
+- conda: https://conda.anaconda.org/conda-forge/noarch/hpack-4.1.0-pyhd8ed1ab_0.conda
+ sha256: 6ad78a180576c706aabeb5b4c8ceb97c0cb25f1e112d76495bff23e3779948ba
+ md5: 0a802cb9888dd14eeefc611f05c40b6e
+ depends:
+ - python >=3.9
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/hpack?source=hash-mapping
+ size: 30731
+ timestamp: 1737618390337
+- conda: https://conda.anaconda.org/conda-forge/noarch/hyperframe-6.1.0-pyhd8ed1ab_0.conda
+ sha256: 77af6f5fe8b62ca07d09ac60127a30d9069fdc3c68d6b256754d0ffb1f7779f8
+ md5: 8e6923fc12f1fe8f8c4e5c9f343256ac
+ depends:
+ - python >=3.9
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/hyperframe?source=hash-mapping
+ size: 17397
+ timestamp: 1737618427549
+- conda: https://conda.anaconda.org/conda-forge/linux-64/icu-75.1-he02047a_0.conda
+ sha256: 71e750d509f5fa3421087ba88ef9a7b9be11c53174af3aa4d06aff4c18b38e8e
+ md5: 8b189310083baabfb622af68fd9d3ae3
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc-ng >=12
+ - libstdcxx-ng >=12
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 12129203
+ timestamp: 1720853576813
+- conda: https://conda.anaconda.org/conda-forge/linux-64/icu-78.2-h33c6efd_0.conda
+ sha256: 142a722072fa96cf16ff98eaaf641f54ab84744af81754c292cb81e0881c0329
+ md5: 186a18e3ba246eccfc7cff00cd19a870
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libstdcxx >=14
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 12728445
+ timestamp: 1767969922681
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/icu-78.2-h38cb7af_0.conda
+ sha256: d4cefbca587429d1192509edc52c88de52bc96c2447771ddc1f8bee928aed5ef
+ md5: 1e93aca311da0210e660d2247812fa02
+ depends:
+ - __osx >=11.0
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 12358010
+ timestamp: 1767970350308
+- pypi: https://files.pythonhosted.org/packages/9f/cb/18326d2d89ad3b0dd143da971e77afd1e6ca6674f1b1c3df4b6bec6279fc/id-1.5.0-py3-none-any.whl
+ name: id
+ version: 1.5.0
+ sha256: f1434e1cef91f2cbb8a4ec64663d5a23b9ed43ef44c4c957d02583d61714c658
+ requires_dist:
+ - requests
+ - build ; extra == 'dev'
+ - bump>=1.3.2 ; extra == 'dev'
+ - id[test,lint] ; extra == 'dev'
+ - bandit ; extra == 'lint'
+ - interrogate ; extra == 'lint'
+ - mypy ; extra == 'lint'
+ - ruff<0.8.2 ; extra == 'lint'
+ - types-requests ; extra == 'lint'
+ - pytest ; extra == 'test'
+ - pytest-cov ; extra == 'test'
+ - pretend ; extra == 'test'
+ - coverage[toml] ; extra == 'test'
+ requires_python: '>=3.8'
+- pypi: https://files.pythonhosted.org/packages/b8/58/40fbbcefeda82364720eba5cf2270f98496bdfa19ea75b4cccae79c698e6/identify-2.6.16-py2.py3-none-any.whl
+ name: identify
+ version: 2.6.16
+ sha256: 391ee4d77741d994189522896270b787aed8670389bfd60f326d677d64a6dfb0
+ requires_dist:
+ - ukkonen ; extra == 'license'
+ requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/noarch/idna-3.11-pyhd8ed1ab_0.conda
+ sha256: ae89d0299ada2a3162c2614a9d26557a92aa6a77120ce142f8e0109bbf0342b0
+ md5: 53abe63df7e10a6ba605dc5f9f961d36
+ depends:
+ - python >=3.10
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/idna?source=hash-mapping
+ size: 50721
+ timestamp: 1760286526795
+- conda: https://conda.anaconda.org/conda-forge/linux-64/imagecodecs-2025.3.30-py310h4eb8eaf_2.conda
+ sha256: a45935f8482e07c1ff8829659d587710b4264c46164db5315a32b7f90c0380da
+ md5: a9c921699d37e862f9bf8dcf9d343838
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - blosc >=1.21.6,<2.0a0
+ - brunsli >=0.1,<1.0a0
+ - bzip2 >=1.0.8,<2.0a0
+ - c-blosc2 >=2.19.0,<2.20.0a0
+ - charls >=2.4.2,<2.5.0a0
+ - giflib >=5.2.2,<5.3.0a0
+ - jxrlib >=1.1,<1.2.0a0
+ - lcms2 >=2.17,<3.0a0
+ - lerc >=4.0.0,<5.0a0
+ - libaec >=1.1.4,<2.0a0
+ - libavif16 >=1.3.0,<2.0a0
+ - libbrotlicommon >=1.1.0,<1.2.0a0
+ - libbrotlidec >=1.1.0,<1.2.0a0
+ - libbrotlienc >=1.1.0,<1.2.0a0
+ - libdeflate >=1.24,<1.25.0a0
+ - libgcc >=13
+ - libjpeg-turbo >=3.1.0,<4.0a0
+ - libjxl >=0.11,<0.12.0a0
+ - liblzma >=5.8.1,<6.0a0
+ - libpng >=1.6.49,<1.7.0a0
+ - libstdcxx >=13
+ - libtiff >=4.7.0,<4.8.0a0
+ - libwebp-base >=1.5.0,<2.0a0
+ - libzlib >=1.3.1,<2.0a0
+ - libzopfli >=1.0.3,<1.1.0a0
+ - lz4-c >=1.10.0,<1.11.0a0
+ - numpy >=1.21,<3
+ - openjpeg >=2.5.3,<3.0a0
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - snappy >=1.2.1,<1.3.0a0
+ - zfp >=1.0.1,<2.0a0
+ - zlib-ng >=2.2.4,<2.3.0a0
+ - zstd >=1.5.7,<1.6.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/imagecodecs?source=hash-mapping
+ size: 1906613
+ timestamp: 1750867156554
+- conda: https://conda.anaconda.org/conda-forge/noarch/imageio-2.37.0-pyhfb79c49_0.conda
+ sha256: 8ef69fa00c68fad34a3b7b260ea774fda9bd9274fd706d3baffb9519fd0063fe
+ md5: b5577bc2212219566578fd5af9993af6
+ depends:
+ - numpy
+ - pillow >=8.3.2
+ - python >=3.9
+ license: BSD-2-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/imageio?source=hash-mapping
+ size: 293226
+ timestamp: 1738273949742
+- conda: https://conda.anaconda.org/conda-forge/noarch/importlib-metadata-8.7.0-pyhe01879c_1.conda
+ sha256: c18ab120a0613ada4391b15981d86ff777b5690ca461ea7e9e49531e8f374745
+ md5: 63ccfdc3a3ce25b027b8767eb722fca8
+ depends:
+ - python >=3.9
+ - zipp >=3.20
+ - python
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/importlib-metadata?source=hash-mapping
+ size: 34641
+ timestamp: 1747934053147
+- pypi: https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl
+ name: iniconfig
+ version: 2.3.0
+ sha256: f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12
+ requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/7f/ed/e3705d6d02b4f7aea715a353c8ce193efd0b5db13e204df895d38734c244/isort-7.0.0-py3-none-any.whl
+ name: isort
+ version: 7.0.0
+ sha256: 1bcabac8bc3c36c7fb7b98a76c8abb18e0f841a3ba81decac7691008592499c1
+ requires_dist:
+ - colorama ; extra == 'colors'
+ - setuptools ; extra == 'plugins'
+ requires_python: '>=3.10.0'
+- pypi: https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl
+ name: jaraco-classes
+ version: 3.4.0
+ sha256: f662826b6bed8cace05e7ff873ce0f9283b5c924470fe664fff1c2f00f581790
+ requires_dist:
+ - more-itertools
+ - sphinx>=3.5 ; extra == 'docs'
+ - jaraco-packaging>=9.3 ; extra == 'docs'
+ - rst-linker>=1.9 ; extra == 'docs'
+ - furo ; extra == 'docs'
+ - sphinx-lint ; extra == 'docs'
+ - jaraco-tidelift>=1.4 ; extra == 'docs'
+ - pytest>=6 ; extra == 'testing'
+ - pytest-checkdocs>=2.4 ; extra == 'testing'
+ - pytest-cov ; extra == 'testing'
+ - pytest-mypy ; extra == 'testing'
+ - pytest-enabler>=2.2 ; extra == 'testing'
+ - pytest-ruff>=0.2.1 ; extra == 'testing'
+ requires_python: '>=3.8'
+- pypi: https://files.pythonhosted.org/packages/8d/48/aa685dbf1024c7bd82bede569e3a85f82c32fd3d79ba5fea578f0159571a/jaraco_context-6.1.0-py3-none-any.whl
+ name: jaraco-context
+ version: 6.1.0
+ sha256: a43b5ed85815223d0d3cfdb6d7ca0d2bc8946f28f30b6f3216bda070f68badda
+ requires_dist:
+ - backports-tarfile ; python_full_version < '3.12'
+ - pytest>=6,!=8.1.* ; extra == 'test'
+ - jaraco-test>=5.6.0 ; extra == 'test'
+ - portend ; extra == 'test'
+ - sphinx>=3.5 ; extra == 'doc'
+ - jaraco-packaging>=9.3 ; extra == 'doc'
+ - rst-linker>=1.9 ; extra == 'doc'
+ - furo ; extra == 'doc'
+ - sphinx-lint ; extra == 'doc'
+ - jaraco-tidelift>=1.4 ; extra == 'doc'
+ - pytest-checkdocs>=2.4 ; extra == 'check'
+ - pytest-ruff>=0.2.1 ; sys_platform != 'cygwin' and extra == 'check'
+ - pytest-cov ; extra == 'cover'
+ - pytest-enabler>=3.4 ; extra == 'enabler'
+ - pytest-mypy>=1.0.1 ; extra == 'type'
+ - mypy<1.19 ; platform_python_implementation == 'PyPy' and extra == 'type'
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/fd/c4/813bb09f0985cb21e959f21f2464169eca882656849adf727ac7bb7e1767/jaraco_functools-4.4.0-py3-none-any.whl
+ name: jaraco-functools
+ version: 4.4.0
+ sha256: 9eec1e36f45c818d9bf307c8948eb03b2b56cd44087b3cdc989abca1f20b9176
+ requires_dist:
+ - more-itertools
+ - pytest>=6,!=8.1.* ; extra == 'test'
+ - jaraco-classes ; extra == 'test'
+ - sphinx>=3.5 ; extra == 'doc'
+ - jaraco-packaging>=9.3 ; extra == 'doc'
+ - rst-linker>=1.9 ; extra == 'doc'
+ - furo ; extra == 'doc'
+ - sphinx-lint ; extra == 'doc'
+ - jaraco-tidelift>=1.4 ; extra == 'doc'
+ - pytest-checkdocs>=2.4 ; extra == 'check'
+ - pytest-ruff>=0.2.1 ; sys_platform != 'cygwin' and extra == 'check'
+ - pytest-cov ; extra == 'cover'
+ - pytest-enabler>=3.4 ; extra == 'enabler'
+ - pytest-mypy>=1.0.1 ; extra == 'type'
+ - mypy<1.19 ; platform_python_implementation == 'PyPy' and extra == 'type'
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/b2/a3/e137168c9c44d18eff0376253da9f1e9234d0239e0ee230d2fee6cea8e55/jeepney-0.9.0-py3-none-any.whl
+ name: jeepney
+ version: 0.9.0
+ sha256: 97e5714520c16fc0a45695e5365a2e11b81ea79bba796e26f9f1d178cb182683
+ requires_dist:
+ - pytest ; extra == 'test'
+ - pytest-trio ; extra == 'test'
+ - pytest-asyncio>=0.17 ; extra == 'test'
+ - testpath ; extra == 'test'
+ - trio ; extra == 'test'
+ - async-timeout ; python_full_version < '3.11' and extra == 'test'
+ - trio ; extra == 'trio'
+ requires_python: '>=3.7'
+- conda: https://conda.anaconda.org/conda-forge/noarch/jinja2-3.1.6-pyhcf101f3_1.conda
+ sha256: fc9ca7348a4f25fed2079f2153ecdcf5f9cf2a0bc36c4172420ca09e1849df7b
+ md5: 04558c96691bed63104678757beb4f8d
+ depends:
+ - markupsafe >=2.0
+ - python >=3.10
+ - python
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/jinja2?source=compressed-mapping
+ size: 120685
+ timestamp: 1764517220861
+- pypi: https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl
+ name: jmespath
+ version: 1.1.0
+ sha256: a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/noarch/joblib-1.5.3-pyhd8ed1ab_0.conda
+ sha256: 301539229d7be6420c084490b8145583291123f0ce6b92f56be5948a2c83a379
+ md5: 615de2a4d97af50c350e5cf160149e77
+ depends:
+ - python >=3.10
+ - setuptools
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/joblib?source=hash-mapping
+ size: 226448
+ timestamp: 1765794135253
+- conda: https://conda.anaconda.org/conda-forge/linux-64/json-c-0.18-h6688a6e_0.conda
+ sha256: 09e706cb388d3ea977fabcee8e28384bdaad8ce1fc49340df5f868a2bd95a7da
+ md5: 38f5dbc9ac808e31c00650f7be1db93f
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 82709
+ timestamp: 1726487116178
+- conda: https://conda.anaconda.org/conda-forge/linux-64/jxrlib-1.1-hd590300_3.conda
+ sha256: 2057ca87b313bde5b74b93b0e696f8faab69acd4cb0edebb78469f3f388040c0
+ md5: 5aeabe88534ea4169d4c49998f293d6c
+ depends:
+ - libgcc-ng >=12
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 239104
+ timestamp: 1703333860145
+- pypi: https://files.pythonhosted.org/packages/ba/61/cc8be27bd65082440754be443b17b6f7c185dec5e00dfdaeab4f8662e4a8/keras-3.12.0-py3-none-any.whl
+ name: keras
+ version: 3.12.0
+ sha256: 02b69e007d5df8042286c3bcc2a888539e3e487590ffb08f6be1b4354df50aa8
+ requires_dist:
+ - absl-py
+ - numpy
+ - rich
+ - namex
+ - h5py
+ - optree
+ - ml-dtypes
+ - packaging
+ requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl
+ name: keyring
+ version: 25.7.0
+ sha256: be4a0b195f149690c166e850609a477c532ddbfbaed96a404d4e43f8d5e2689f
+ requires_dist:
+ - pywin32-ctypes>=0.2.0 ; sys_platform == 'win32'
+ - secretstorage>=3.2 ; sys_platform == 'linux'
+ - jeepney>=0.4.2 ; sys_platform == 'linux'
+ - importlib-metadata>=4.11.4 ; python_full_version < '3.12'
+ - jaraco-classes
+ - jaraco-functools
+ - jaraco-context
+ - pytest>=6,!=8.1.* ; extra == 'test'
+ - pyfakefs ; extra == 'test'
+ - sphinx>=3.5 ; extra == 'doc'
+ - jaraco-packaging>=9.3 ; extra == 'doc'
+ - rst-linker>=1.9 ; extra == 'doc'
+ - furo ; extra == 'doc'
+ - sphinx-lint ; extra == 'doc'
+ - jaraco-tidelift>=1.4 ; extra == 'doc'
+ - pytest-checkdocs>=2.4 ; extra == 'check'
+ - pytest-ruff>=0.2.1 ; sys_platform != 'cygwin' and extra == 'check'
+ - pytest-cov ; extra == 'cover'
+ - pytest-enabler>=3.4 ; extra == 'enabler'
+ - pytest-mypy>=1.0.1 ; extra == 'type'
+ - pygobject-stubs ; extra == 'type'
+ - shtab ; extra == 'type'
+ - types-pywin32 ; extra == 'type'
+ - shtab>=1.1.0 ; extra == 'completion'
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/keyutils-1.6.3-hb9d3cd8_0.conda
+ sha256: 0960d06048a7185d3542d850986d807c6e37ca2e644342dd0c72feefcf26c2a4
+ md5: b38117a3c920364aff79f870c984b4a3
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ license: LGPL-2.1-or-later
+ purls: []
+ size: 134088
+ timestamp: 1754905959823
+- conda: https://conda.anaconda.org/conda-forge/linux-64/kiwisolver-1.4.9-py310haaf941d_2.conda
+ sha256: 5ef8337c7a89719427d25b0cdc776b34116fe988efc9bf56f5a2831d74b1584e
+ md5: 7426d76535fc6347f1b74f85fb17d6eb
+ depends:
+ - python
+ - libstdcxx >=14
+ - libgcc >=14
+ - __glibc >=2.17,<3.0.a0
+ - python_abi 3.10.* *_cp310
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/kiwisolver?source=hash-mapping
+ size: 78299
+ timestamp: 1762488741951
+- conda: https://conda.anaconda.org/conda-forge/linux-64/krb5-1.21.3-h659f571_0.conda
+ sha256: 99df692f7a8a5c27cd14b5fb1374ee55e756631b9c3d659ed3ee60830249b238
+ md5: 3f43953b7d3fb3aaa1d0d0723d91e368
+ depends:
+ - keyutils >=1.6.1,<2.0a0
+ - libedit >=3.1.20191231,<3.2.0a0
+ - libedit >=3.1.20191231,<4.0a0
+ - libgcc-ng >=12
+ - libstdcxx-ng >=12
+ - openssl >=3.3.1,<4.0a0
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 1370023
+ timestamp: 1719463201255
+- conda: https://conda.anaconda.org/conda-forge/noarch/lazy-loader-0.4-pyhd8ed1ab_2.conda
+ sha256: d7ea986507090fff801604867ef8e79c8fda8ec21314ba27c032ab18df9c3411
+ md5: d10d9393680734a8febc4b362a4c94f2
+ depends:
+ - importlib-metadata
+ - packaging
+ - python >=3.9
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/lazy-loader?source=hash-mapping
+ size: 16298
+ timestamp: 1733636905835
+- conda: https://conda.anaconda.org/conda-forge/linux-64/lcms2-2.18-h0c24ade_0.conda
+ sha256: 836ec4b895352110335b9fdcfa83a8dcdbe6c5fb7c06c4929130600caea91c0a
+ md5: 6f2e2c8f58160147c4d1c6f4c14cbac4
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libjpeg-turbo >=3.1.2,<4.0a0
+ - libtiff >=4.7.1,<4.8.0a0
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 249959
+ timestamp: 1768184673131
+- conda: https://conda.anaconda.org/conda-forge/linux-64/ld_impl_linux-64-2.45-bootstrap_ha15bf96_5.conda
+ sha256: 39214f6699455d335e692fb8c1e67bb6897ef9593a4a1c0d3a2b776e880161f8
+ md5: 7a6d3da38d766cffa2105b56e1c5d61c
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ constrains:
+ - binutils_impl_linux-64 2.45
+ license: GPL-3.0-only
+ license_family: GPL
+ purls: []
+ size: 729347
+ timestamp: 1766512932021
+- conda: https://conda.anaconda.org/conda-forge/linux-64/ld_impl_linux-64-2.45.1-default_hbd61a6d_101.conda
+ sha256: 565941ac1f8b0d2f2e8f02827cbca648f4d18cd461afc31f15604cd291b5c5f3
+ md5: 12bd9a3f089ee6c9266a37dab82afabd
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - zstd >=1.5.7,<1.6.0a0
+ constrains:
+ - binutils_impl_linux-64 2.45.1
+ license: GPL-3.0-only
+ license_family: GPL
+ purls: []
+ size: 725507
+ timestamp: 1770267139900
+- conda: https://conda.anaconda.org/conda-forge/linux-64/lerc-4.0.0-h0aef613_1.conda
+ sha256: 412381a43d5ff9bbed82cd52a0bbca5b90623f62e41007c9c42d3870c60945ff
+ md5: 9344155d33912347b37f0ae6c410a835
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - libstdcxx >=13
+ license: Apache-2.0
+ license_family: Apache
+ purls: []
+ size: 264243
+ timestamp: 1745264221534
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libabseil-20230125.3-cxx17_h59595ed_0.conda
+ sha256: 3c6fab31ed4dc8428605588454596b307b1bd59d33b0c7073c407ab51408b011
+ md5: d1db1b8be7c3a8983dcbbbfe4f0765de
+ depends:
+ - libgcc-ng >=12
+ - libstdcxx-ng >=12
+ constrains:
+ - abseil-cpp =20230125.3
+ - libabseil-static =20230125.3=cxx17*
+ license: Apache-2.0
+ license_family: Apache
+ purls: []
+ size: 1240376
+ timestamp: 1688112986128
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libaec-1.1.5-h088129d_0.conda
+ sha256: 822e4ae421a7e9c04e841323526321185f6659222325e1a9aedec811c686e688
+ md5: 86f7414544ae606282352fa1e116b41f
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libstdcxx >=14
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 36544
+ timestamp: 1769221884824
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libarchive-3.8.5-gpl_hc2c16d8_100.conda
+ sha256: ee2cf1499a5a5fd5f03c6203597fe14bf28c6ca2a8fffb761e41f3cf371e768e
+ md5: 5fdaa8b856683a5598459dead3976578
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - bzip2 >=1.0.8,<2.0a0
+ - libgcc >=14
+ - liblzma >=5.8.1,<6.0a0
+ - libxml2
+ - libxml2-16 >=2.14.6
+ - libzlib >=1.3.1,<2.0a0
+ - lz4-c >=1.10.0,<1.11.0a0
+ - lzo >=2.10,<3.0a0
+ - openssl >=3.5.4,<4.0a0
+ - zstd >=1.5.7,<1.6.0a0
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 886102
+ timestamp: 1767630453053
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libavif16-1.3.0-h316e467_3.conda
+ sha256: f5ab201b8b4e1f776ced0340c59f87e441fd6763d3face527b5cf3f2280502c9
+ md5: 22d5cc5fb45aab8ed3c00cde2938b825
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - aom >=3.9.1,<3.10.0a0
+ - dav1d >=1.2.1,<1.2.2.0a0
+ - libgcc >=14
+ - rav1e >=0.7.1,<0.8.0a0
+ - svt-av1 >=4.0.0,<4.0.1.0a0
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 140323
+ timestamp: 1769476997956
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.11.0-5_h4a7cf45_openblas.conda
+ build_number: 5
+ sha256: 18c72545080b86739352482ba14ba2c4815e19e26a7417ca21a95b76ec8da24c
+ md5: c160954f7418d7b6e87eaf05a8913fa9
+ depends:
+ - libopenblas >=0.3.30,<0.3.31.0a0
+ - libopenblas >=0.3.30,<1.0a0
+ constrains:
+ - mkl <2026
+ - liblapack 3.11.0 5*_openblas
+ - libcblas 3.11.0 5*_openblas
+ - blas 2.305 openblas
+ - liblapacke 3.11.0 5*_openblas
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 18213
+ timestamp: 1765818813880
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/libblas-3.11.0-5_h51639a9_openblas.conda
+ build_number: 5
+ sha256: 620a6278f194dcabc7962277da6835b1e968e46ad0c8e757736255f5ddbfca8d
+ md5: bcc025e2bbaf8a92982d20863fe1fb69
+ depends:
+ - libopenblas >=0.3.30,<0.3.31.0a0
+ - libopenblas >=0.3.30,<1.0a0
+ constrains:
+ - libcblas 3.11.0 5*_openblas
+ - liblapack 3.11.0 5*_openblas
+ - liblapacke 3.11.0 5*_openblas
+ - blas 2.305 openblas
+ - mkl <2026
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 18546
+ timestamp: 1765819094137
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlicommon-1.1.0-hb03c661_4.conda
+ sha256: 2338a92d1de71f10c8cf70f7bb9775b0144a306d75c4812276749f54925612b6
+ md5: 1d29d2e33fe59954af82ef54a8af3fe1
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 69333
+ timestamp: 1756599354727
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlidec-1.1.0-hb03c661_4.conda
+ sha256: fcec0d26f67741b122f0d5eff32f0393d7ebd3ee6bb866ae2f17f3425a850936
+ md5: 5cb5a1c9a94a78f5b23684bcb845338d
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libbrotlicommon 1.1.0 hb03c661_4
+ - libgcc >=14
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 33406
+ timestamp: 1756599364386
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libbrotlienc-1.1.0-hb03c661_4.conda
+ sha256: d42c7f0afce21d5279a0d54ee9e64a2279d35a07a90e0c9545caae57d6d7dc57
+ md5: 2e55011fa483edb8bfe3fd92e860cd79
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libbrotlicommon 1.1.0 hb03c661_4
+ - libgcc >=14
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 289680
+ timestamp: 1756599375485
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.11.0-5_h0358290_openblas.conda
+ build_number: 5
+ sha256: 0cbdcc67901e02dc17f1d19e1f9170610bd828100dc207de4d5b6b8ad1ae7ad8
+ md5: 6636a2b6f1a87572df2970d3ebc87cc0
+ depends:
+ - libblas 3.11.0 5_h4a7cf45_openblas
+ constrains:
+ - liblapacke 3.11.0 5*_openblas
+ - blas 2.305 openblas
+ - liblapack 3.11.0 5*_openblas
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 18194
+ timestamp: 1765818837135
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/libcblas-3.11.0-5_hb0561ab_openblas.conda
+ build_number: 5
+ sha256: 38809c361bbd165ecf83f7f05fae9b791e1baa11e4447367f38ae1327f402fc0
+ md5: efd8bd15ca56e9d01748a3beab8404eb
+ depends:
+ - libblas 3.11.0 5_h51639a9_openblas
+ constrains:
+ - liblapacke 3.11.0 5*_openblas
+ - liblapack 3.11.0 5*_openblas
+ - blas 2.305 openblas
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 18548
+ timestamp: 1765819108956
+- pypi: https://files.pythonhosted.org/packages/1d/fc/716c1e62e512ef1c160e7984a73a5fc7df45166f2ff3f254e71c58076f7c/libclang-18.1.1-py2.py3-none-manylinux2010_x86_64.whl
+ name: libclang
+ version: 18.1.1
+ sha256: c533091d8a3bbf7460a00cb6c1a71da93bffe148f172c7d03b1c31fbf8aa2a0b
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libcrc32c-1.1.2-h9c3ff4c_0.tar.bz2
+ sha256: fd1d153962764433fe6233f34a72cdeed5dcf8a883a85769e8295ce940b5b0c5
+ md5: c965a5aa0d5c1c37ffc62dff36e28400
+ depends:
+ - libgcc-ng >=9.4.0
+ - libstdcxx-ng >=9.4.0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 20440
+ timestamp: 1633683576494
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libcurl-8.18.0-h4e3cde8_0.conda
+ sha256: 5454709d9fb6e9c3dd6423bc284fa7835a7823bfa8323f6e8786cdd555101fab
+ md5: 0a5563efed19ca4461cf927419b6eb73
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - krb5 >=1.21.3,<1.22.0a0
+ - libgcc >=14
+ - libnghttp2 >=1.67.0,<2.0a0
+ - libssh2 >=1.11.1,<2.0a0
+ - libzlib >=1.3.1,<2.0a0
+ - openssl >=3.5.4,<4.0a0
+ - zstd >=1.5.7,<1.6.0a0
+ license: curl
+ license_family: MIT
+ purls: []
+ size: 462942
+ timestamp: 1767821743793
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/libcxx-21.1.8-h55c6f16_2.conda
+ sha256: 5fbeb2fc2673f0455af6079abf93faaf27f11a92574ad51565fa1ecac9a4e2aa
+ md5: 4cb5878bdb9ebfa65b7cdff5445087c5
+ depends:
+ - __osx >=11.0
+ license: Apache-2.0 WITH LLVM-exception
+ license_family: Apache
+ purls: []
+ size: 570068
+ timestamp: 1770238262922
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libdeflate-1.24-h86f0d12_0.conda
+ sha256: 8420748ea1cc5f18ecc5068b4f24c7a023cc9b20971c99c824ba10641fb95ddf
+ md5: 64f0c503da58ec25ebd359e4d990afa8
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 72573
+ timestamp: 1747040452262
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libedit-3.1.20250104-pl5321h7949ede_0.conda
+ sha256: d789471216e7aba3c184cd054ed61ce3f6dac6f87a50ec69291b9297f8c18724
+ md5: c277e0a4d549b03ac1e9d6cbbe3d017b
+ depends:
+ - ncurses
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - ncurses >=6.5,<7.0a0
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 134676
+ timestamp: 1738479519902
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libev-4.33-hd590300_2.conda
+ sha256: 1cd6048169fa0395af74ed5d8f1716e22c19a81a8a36f934c110ca3ad4dd27b4
+ md5: 172bf1cd1ff8629f2b1179945ed45055
+ depends:
+ - libgcc-ng >=12
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 112766
+ timestamp: 1702146165126
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libexpat-2.7.3-hecca717_0.conda
+ sha256: 1e1b08f6211629cbc2efe7a5bca5953f8f6b3cae0eeb04ca4dacee1bd4e2db2f
+ md5: 8b09ae86839581147ef2e5c5e229d164
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ constrains:
+ - expat 2.7.3.*
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 76643
+ timestamp: 1763549731408
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libexpat-2.7.4-hecca717_0.conda
+ sha256: d78f1d3bea8c031d2f032b760f36676d87929b18146351c4464c66b0869df3f5
+ md5: e7f7ce06ec24cfcfb9e36d28cf82ba57
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ constrains:
+ - expat 2.7.4.*
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 76798
+ timestamp: 1771259418166
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/libexpat-2.7.4-hf6b4638_0.conda
+ sha256: 03887d8080d6a8fe02d75b80929271b39697ecca7628f0657d7afaea87761edf
+ md5: a92e310ae8dfc206ff449f362fc4217f
+ depends:
+ - __osx >=11.0
+ constrains:
+ - expat 2.7.4.*
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 68199
+ timestamp: 1771260020767
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libffi-3.5.2-h3435931_0.conda
+ sha256: 31f19b6a88ce40ebc0d5a992c131f57d919f73c0b92cd1617a5bec83f6e961e6
+ md5: a360c33a5abe61c07959e449fa1453eb
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 58592
+ timestamp: 1769456073053
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/libffi-3.5.2-hcf2aa1b_0.conda
+ sha256: 6686a26466a527585e6a75cc2a242bf4a3d97d6d6c86424a441677917f28bec7
+ md5: 43c04d9cb46ef176bb2a4c77e324d599
+ depends:
+ - __osx >=11.0
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 40979
+ timestamp: 1769456747661
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libfreetype-2.14.1-ha770c72_0.conda
+ sha256: 4641d37faeb97cf8a121efafd6afd040904d4bca8c46798122f417c31d5dfbec
+ md5: f4084e4e6577797150f9b04a4560ceb0
+ depends:
+ - libfreetype6 >=2.14.1
+ license: GPL-2.0-only OR FTL
+ purls: []
+ size: 7664
+ timestamp: 1757945417134
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libfreetype6-2.14.1-h73754d4_0.conda
+ sha256: 4a7af818a3179fafb6c91111752954e29d3a2a950259c14a2fc7ba40a8b03652
+ md5: 8e7251989bca326a28f4a5ffbd74557a
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libpng >=1.6.50,<1.7.0a0
+ - libzlib >=1.3.1,<2.0a0
+ constrains:
+ - freetype >=2.14.1
+ license: GPL-2.0-only OR FTL
+ purls: []
+ size: 386739
+ timestamp: 1757945416744
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-15.2.0-he0feb66_16.conda
+ sha256: 6eed58051c2e12b804d53ceff5994a350c61baf117ec83f5f10c953a3f311451
+ md5: 6d0363467e6ed84f11435eb309f2ff06
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - _openmp_mutex >=4.5
+ constrains:
+ - libgcc-ng ==15.2.0=*_16
+ - libgomp 15.2.0 he0feb66_16
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 1042798
+ timestamp: 1765256792743
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-15.2.0-he0feb66_18.conda
+ sha256: faf7d2017b4d718951e3a59d081eb09759152f93038479b768e3d612688f83f5
+ md5: 0aa00f03f9e39fb9876085dee11a85d4
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - _openmp_mutex >=4.5
+ constrains:
+ - libgcc-ng ==15.2.0=*_18
+ - libgomp 15.2.0 he0feb66_18
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 1041788
+ timestamp: 1771378212382
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/libgcc-15.2.0-hcbb3090_18.conda
+ sha256: 1d9c4f35586adb71bcd23e31b68b7f3e4c4ab89914c26bed5f2859290be5560e
+ md5: 92df6107310b1fff92c4cc84f0de247b
+ depends:
+ - _openmp_mutex
+ constrains:
+ - libgcc-ng ==15.2.0=*_18
+ - libgomp 15.2.0 18
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 401974
+ timestamp: 1771378877463
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-ng-15.2.0-h69a702a_16.conda
+ sha256: 5f07f9317f596a201cc6e095e5fc92621afca64829785e483738d935f8cab361
+ md5: 5a68259fac2da8f2ee6f7bfe49c9eb8b
+ depends:
+ - libgcc 15.2.0 he0feb66_16
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 27256
+ timestamp: 1765256804124
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgcc-ng-15.2.0-h69a702a_18.conda
+ sha256: e318a711400f536c81123e753d4c797a821021fb38970cebfb3f454126016893
+ md5: d5e96b1ed75ca01906b3d2469b4ce493
+ depends:
+ - libgcc 15.2.0 he0feb66_18
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 27526
+ timestamp: 1771378224552
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgdal-core-3.10.3-h95ec890_18.conda
+ sha256: fa58165fe24489dfd4d14a759d2eb6d602a4b783a663027e9c2ef96a10bb049e
+ md5: 7de0158b4bc2efa21ccd72784813274a
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - blosc >=1.21.6,<2.0a0
+ - geos >=3.14.0,<3.14.1.0a0
+ - geotiff >=1.7.4,<1.8.0a0
+ - giflib >=5.2.2,<5.3.0a0
+ - json-c >=0.18,<0.19.0a0
+ - lerc >=4.0.0,<5.0a0
+ - libarchive >=3.8.1,<3.9.0a0
+ - libcurl >=8.14.1,<9.0a0
+ - libdeflate >=1.24,<1.25.0a0
+ - libexpat >=2.7.1,<3.0a0
+ - libgcc >=14
+ - libiconv >=1.18,<2.0a0
+ - libjpeg-turbo >=3.1.0,<4.0a0
+ - libkml >=1.3.0,<1.4.0a0
+ - liblzma >=5.8.1,<6.0a0
+ - libpng >=1.6.50,<1.7.0a0
+ - libspatialite >=5.1.0,<5.2.0a0
+ - libsqlite >=3.50.4,<4.0a0
+ - libstdcxx >=14
+ - libtiff >=4.7.0,<4.8.0a0
+ - libwebp-base >=1.6.0,<2.0a0
+ - libxml2
+ - libxml2-16 >=2.14.6
+ - libzlib >=1.3.1,<2.0a0
+ - lz4-c >=1.10.0,<1.11.0a0
+ - openssl >=3.5.2,<4.0a0
+ - pcre2 >=10.46,<10.47.0a0
+ - proj >=9.6.2,<9.7.0a0
+ - xerces-c >=3.2.5,<3.3.0a0
+ - zstd >=1.5.7,<1.6.0a0
+ constrains:
+ - libgdal 3.10.3.*
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 11040063
+ timestamp: 1758000487897
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran-15.2.0-h69a702a_16.conda
+ sha256: 8a7b01e1ee1c462ad243524d76099e7174ebdd94ff045fe3e9b1e58db196463b
+ md5: 40d9b534410403c821ff64f00d0adc22
+ depends:
+ - libgfortran5 15.2.0 h68bc16d_16
+ constrains:
+ - libgfortran-ng ==15.2.0=*_16
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 27215
+ timestamp: 1765256845586
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran-15.2.0-h69a702a_18.conda
+ sha256: d2c9fad338fd85e4487424865da8e74006ab2e2475bd788f624d7a39b2a72aee
+ md5: 9063115da5bc35fdc3e1002e69b9ef6e
+ depends:
+ - libgfortran5 15.2.0 h68bc16d_18
+ constrains:
+ - libgfortran-ng ==15.2.0=*_18
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 27523
+ timestamp: 1771378269450
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/libgfortran-15.2.0-h07b0088_18.conda
+ sha256: 63f89087c3f0c8621c5c89ecceec1e56e5e1c84f65fc9c5feca33a07c570a836
+ md5: 26981599908ed2205366e8fc91b37fc6
+ depends:
+ - libgfortran5 15.2.0 hdae7583_18
+ constrains:
+ - libgfortran-ng ==15.2.0=*_18
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 138973
+ timestamp: 1771379054939
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran5-15.2.0-h68bc16d_16.conda
+ sha256: d0e974ebc937c67ae37f07a28edace978e01dc0f44ee02f29ab8a16004b8148b
+ md5: 39183d4e0c05609fd65f130633194e37
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=15.2.0
+ constrains:
+ - libgfortran 15.2.0
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 2480559
+ timestamp: 1765256819588
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgfortran5-15.2.0-h68bc16d_18.conda
+ sha256: 539b57cf50ec85509a94ba9949b7e30717839e4d694bc94f30d41c9d34de2d12
+ md5: 646855f357199a12f02a87382d429b75
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=15.2.0
+ constrains:
+ - libgfortran 15.2.0
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 2482475
+ timestamp: 1771378241063
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/libgfortran5-15.2.0-hdae7583_18.conda
+ sha256: 91033978ba25e6a60fb86843cf7e1f7dc8ad513f9689f991c9ddabfaf0361e7e
+ md5: c4a6f7989cffb0544bfd9207b6789971
+ depends:
+ - libgcc >=15.2.0
+ constrains:
+ - libgfortran 15.2.0
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 598634
+ timestamp: 1771378886363
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgomp-15.2.0-he0feb66_16.conda
+ sha256: 5b3e5e4e9270ecfcd48f47e3a68f037f5ab0f529ccb223e8e5d5ac75a58fc687
+ md5: 26c46f90d0e727e95c6c9498a33a09f3
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 603284
+ timestamp: 1765256703881
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgomp-15.2.0-he0feb66_18.conda
+ sha256: 21337ab58e5e0649d869ab168d4e609b033509de22521de1bfed0c031bfc5110
+ md5: 239c5e9546c38a1e884d69effcf4c882
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 603262
+ timestamp: 1771378117851
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libgrpc-1.54.3-hb20ce57_0.conda
+ sha256: f5fea0c2eececb010529ac5863fbede05a2413ea8dc1a1c419db861f68ed66d7
+ md5: 7af7c59ab24db007dfd82e0a3a343f66
+ depends:
+ - c-ares >=1.19.1,<2.0a0
+ - libabseil * cxx17*
+ - libabseil >=20230125.3,<20230126.0a0
+ - libgcc-ng >=12
+ - libprotobuf >=3.21.12,<3.22.0a0
+ - libstdcxx-ng >=12
+ - libzlib >=1.2.13,<2.0.0a0
+ - openssl >=3.1.1,<4.0a0
+ - re2 >=2023.3.2,<2023.3.3.0a0
+ - zlib
+ constrains:
+ - grpc-cpp =1.54.3
+ license: Apache-2.0
+ license_family: APACHE
+ purls: []
+ size: 6003265
+ timestamp: 1690943569727
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libhwy-1.3.0-h4c17acf_1.conda
+ sha256: 2bdd1cdd677b119abc5e83069bec2e28fe6bfb21ebaea3cd07acee67f38ea274
+ md5: c2a0c1d0120520e979685034e0b79859
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libstdcxx >=14
+ license: Apache-2.0 OR BSD-3-Clause
+ purls: []
+ size: 1448617
+ timestamp: 1758894401402
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libiconv-1.18-h3b78370_2.conda
+ sha256: c467851a7312765447155e071752d7bf9bf44d610a5687e32706f480aad2833f
+ md5: 915f5995e94f60e9a4826e0b0920ee88
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ license: LGPL-2.1-only
+ purls: []
+ size: 790176
+ timestamp: 1754908768807
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libjpeg-turbo-3.1.2-hb03c661_0.conda
+ sha256: cc9aba923eea0af8e30e0f94f2ad7156e2984d80d1e8e7fe6be5a1f257f0eb32
+ md5: 8397539e3a0bbd1695584fb4f927485a
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ constrains:
+ - jpeg <0.0.0a
+ license: IJG AND BSD-3-Clause AND Zlib
+ purls: []
+ size: 633710
+ timestamp: 1762094827865
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libjxl-0.11.1-h6cb5226_4.conda
+ sha256: b9d924d69fc84cd3c660a181985748d9c2df34cd7c7bb03b92d8f70efa7753d9
+ md5: f2840d9c2afb19e303e126c9d3a04b36
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libbrotlidec >=1.1.0,<1.2.0a0
+ - libbrotlienc >=1.1.0,<1.2.0a0
+ - libgcc >=14
+ - libhwy >=1.3.0,<1.4.0a0
+ - libstdcxx >=14
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 1740823
+ timestamp: 1757583994233
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libkml-1.3.0-h01aab08_1016.conda
+ sha256: 7e8f2e2bf09e9fae7c046b35545b3c71bebf4d87b38771cc7d058e83ae4b81cc
+ md5: 4d0907546d556ef7f14b1dcfa0e217ce
+ depends:
+ - libexpat >=2.5.0,<3.0a0
+ - libgcc-ng >=12
+ - libstdcxx-ng >=12
+ - libzlib >=1.2.13,<2.0.0a0
+ - uriparser >=0.9.7,<1.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 511196
+ timestamp: 1696314384123
+- conda: https://conda.anaconda.org/conda-forge/linux-64/liblapack-3.11.0-5_h47877c9_openblas.conda
+ build_number: 5
+ sha256: c723b6599fcd4c6c75dee728359ef418307280fa3e2ee376e14e85e5bbdda053
+ md5: b38076eb5c8e40d0106beda6f95d7609
+ depends:
+ - libblas 3.11.0 5_h4a7cf45_openblas
+ constrains:
+ - blas 2.305 openblas
+ - liblapacke 3.11.0 5*_openblas
+ - libcblas 3.11.0 5*_openblas
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 18200
+ timestamp: 1765818857876
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/liblapack-3.11.0-5_hd9741b5_openblas.conda
+ build_number: 5
+ sha256: 735a6e6f7d7da6f718b6690b7c0a8ae4815afb89138aa5793abe78128e951dbb
+ md5: ca9d752201b7fa1225bca036ee300f2b
+ depends:
+ - libblas 3.11.0 5_h51639a9_openblas
+ constrains:
+ - libcblas 3.11.0 5*_openblas
+ - blas 2.305 openblas
+ - liblapacke 3.11.0 5*_openblas
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 18551
+ timestamp: 1765819121855
+- conda: https://conda.anaconda.org/conda-forge/linux-64/liblzma-5.8.2-hb03c661_0.conda
+ sha256: 755c55ebab181d678c12e49cced893598f2bab22d582fbbf4d8b83c18be207eb
+ md5: c7c83eecbb72d88b940c249af56c8b17
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ constrains:
+ - xz 5.8.2.*
+ license: 0BSD
+ purls: []
+ size: 113207
+ timestamp: 1768752626120
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/liblzma-5.8.2-h8088a28_0.conda
+ sha256: 7bfc7ffb2d6a9629357a70d4eadeadb6f88fa26ebc28f606b1c1e5e5ed99dc7e
+ md5: 009f0d956d7bfb00de86901d16e486c7
+ depends:
+ - __osx >=11.0
+ constrains:
+ - xz 5.8.2.*
+ license: 0BSD
+ purls: []
+ size: 92242
+ timestamp: 1768752982486
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libnetcdf-4.9.3-nompi_h11f7409_103.conda
+ sha256: e9a8668212719a91a6b0348db05188dfc59de5a21888db13ff8510918a67b258
+ md5: 3ccff1066c05a1e6c221356eecc40581
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - attr >=2.5.2,<2.6.0a0
+ - blosc >=1.21.6,<2.0a0
+ - bzip2 >=1.0.8,<2.0a0
+ - hdf4 >=4.2.15,<4.2.16.0a0
+ - hdf5 >=1.14.6,<1.14.7.0a0
+ - libaec >=1.1.4,<2.0a0
+ - libcurl >=8.14.1,<9.0a0
+ - libgcc >=14
+ - libstdcxx >=14
+ - libxml2
+ - libxml2-16 >=2.14.6
+ - libzip >=1.11.2,<2.0a0
+ - libzlib >=1.3.1,<2.0a0
+ - openssl >=3.5.2,<4.0a0
+ - zlib
+ - zstd >=1.5.7,<1.6.0a0
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 871447
+ timestamp: 1757977084313
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libnghttp2-1.67.0-had1ee68_0.conda
+ sha256: a4a7dab8db4dc81c736e9a9b42bdfd97b087816e029e221380511960ac46c690
+ md5: b499ce4b026493a13774bcf0f4c33849
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - c-ares >=1.34.5,<2.0a0
+ - libev >=4.33,<4.34.0a0
+ - libev >=4.33,<5.0a0
+ - libgcc >=14
+ - libstdcxx >=14
+ - libzlib >=1.3.1,<2.0a0
+ - openssl >=3.5.2,<4.0a0
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 666600
+ timestamp: 1756834976695
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libnsl-2.0.1-hb9d3cd8_1.conda
+ sha256: 927fe72b054277cde6cb82597d0fcf6baf127dcbce2e0a9d8925a68f1265eef5
+ md5: d864d34357c3b65a4b731f78c0801dc4
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ license: LGPL-2.1-only
+ license_family: GPL
+ purls: []
+ size: 33731
+ timestamp: 1750274110928
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libopenblas-0.3.30-pthreads_h94d23a6_4.conda
+ sha256: 199d79c237afb0d4780ccd2fbf829cea80743df60df4705202558675e07dd2c5
+ md5: be43915efc66345cccb3c310b6ed0374
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libgfortran
+ - libgfortran5 >=14.3.0
+ constrains:
+ - openblas >=0.3.30,<0.3.31.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 5927939
+ timestamp: 1763114673331
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/libopenblas-0.3.30-openmp_ha158390_4.conda
+ sha256: ebbbc089b70bcde87c4121a083c724330f02a690fb9d7c6cd18c30f1b12504fa
+ md5: a6f6d3a31bb29e48d37ce65de54e2df0
+ depends:
+ - __osx >=11.0
+ - libgfortran
+ - libgfortran5 >=14.3.0
+ - llvm-openmp >=19.1.7
+ constrains:
+ - openblas >=0.3.30,<0.3.31.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 4284132
+ timestamp: 1768547079205
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libpng-1.6.54-h421ea60_0.conda
+ sha256: 5de60d34aac848a9991a09fcdea7c0e783d00024aefec279d55e87c0c44742cd
+ md5: d361fa2a59e53b61c2675bfa073e5b7e
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libzlib >=1.3.1,<2.0a0
+ license: zlib-acknowledgement
+ purls: []
+ size: 317435
+ timestamp: 1768285668880
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libprotobuf-3.21.12-hfc55251_2.conda
+ sha256: 2df8888c51c23dedc831ba4378bad259e95c3a20a6408f54926a6a6f629f6153
+ md5: e3a7d4ba09b8dc939b98fef55f539220
+ depends:
+ - libgcc-ng >=12
+ - libstdcxx-ng >=12
+ - libzlib >=1.2.13,<2.0.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 2230785
+ timestamp: 1693493608116
+- pypi: https://files.pythonhosted.org/packages/82/5f/3e85351c523f73ad8d938989e9a58c7f59fb9c17f761b9981b43f0025ce7/librt-0.7.8-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
+ name: librt
+ version: 0.7.8
+ sha256: 4864045f49dc9c974dadb942ac56a74cd0479a2aafa51ce272c490a82322ea3c
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/librttopo-1.1.0-h96cd706_19.conda
+ sha256: f584227e141db34f5dde05e8a3a4e59c86860e3b5df7698a646b7fc3486b0e86
+ md5: 212a9378a85ad020b8dc94853fdbeb6c
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - geos >=3.14.0,<3.14.1.0a0
+ - libgcc >=14
+ - libstdcxx >=14
+ license: GPL-2.0-or-later
+ license_family: GPL
+ purls: []
+ size: 232294
+ timestamp: 1755880773417
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libspatialindex-2.1.0-he57a185_0.conda
+ sha256: 03963a7786b3f53eb36ca3ec10d7a5ddd5265a81e205e28902c53a536cdfd3ad
+ md5: 2df7aaf3f8a2944885372a62c6f33b20
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - libstdcxx >=13
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 399212
+ timestamp: 1734891697797
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libspatialite-5.1.0-h2eee824_16.conda
+ sha256: e50915445bf76e4e1390f33ef39a69491534468531607a818224e83d0d2c5c6f
+ md5: 0704f614584aefd5fdaf56c75b047fdb
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - freexl >=2
+ - freexl >=2.0.0,<3.0a0
+ - geos >=3.14.0,<3.14.1.0a0
+ - libgcc >=14
+ - librttopo >=1.1.0,<1.2.0a0
+ - libsqlite >=3.50.4,<4.0a0
+ - libstdcxx >=14
+ - libxml2
+ - libxml2-16 >=2.14.5
+ - libxml2-devel
+ - libzlib >=1.3.1,<2.0a0
+ - proj >=9.6.2,<9.7.0a0
+ - sqlite
+ - zlib
+ license: MPL-1.1
+ license_family: MOZILLA
+ purls: []
+ size: 3489498
+ timestamp: 1757320839353
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libsqlite-3.51.2-h0c1763c_0.conda
+ sha256: c1ff4589b48d32ca0a2628970d869fa9f7b2c2d00269a3761edc7e9e4c1ab7b8
+ md5: f7d30045eccb83f2bb8053041f42db3c
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libzlib >=1.3.1,<2.0a0
+ license: blessing
+ purls: []
+ size: 939312
+ timestamp: 1768147967568
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libsqlite-3.51.2-hf4e2dac_0.conda
+ sha256: 04596fcee262a870e4b7c9807224680ff48d4d0cc0dac076a602503d3dc6d217
+ md5: da5be73701eecd0e8454423fd6ffcf30
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - icu >=78.2,<79.0a0
+ - libgcc >=14
+ - libzlib >=1.3.1,<2.0a0
+ license: blessing
+ purls: []
+ size: 942808
+ timestamp: 1768147973361
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/libsqlite-3.51.2-h1ae2325_0.conda
+ sha256: 6e9b9f269732cbc4698c7984aa5b9682c168e2a8d1e0406e1ff10091ca046167
+ md5: 4b0bf313c53c3e89692f020fb55d5f2c
+ depends:
+ - __osx >=11.0
+ - icu >=78.2,<79.0a0
+ - libzlib >=1.3.1,<2.0a0
+ license: blessing
+ purls: []
+ size: 909777
+ timestamp: 1768148320535
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libssh2-1.11.1-hcf80075_0.conda
+ sha256: fa39bfd69228a13e553bd24601332b7cfeb30ca11a3ca50bb028108fe90a7661
+ md5: eecce068c7e4eddeb169591baac20ac4
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - libzlib >=1.3.1,<2.0a0
+ - openssl >=3.5.0,<4.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 304790
+ timestamp: 1745608545575
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-15.2.0-h934c35e_16.conda
+ sha256: 813427918316a00c904723f1dfc3da1bbc1974c5cfe1ed1e704c6f4e0798cbc6
+ md5: 68f68355000ec3f1d6f26ea13e8f525f
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc 15.2.0 he0feb66_16
+ constrains:
+ - libstdcxx-ng ==15.2.0=*_16
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 5856456
+ timestamp: 1765256838573
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-15.2.0-h934c35e_18.conda
+ sha256: 78668020064fdaa27e9ab65cd2997e2c837b564ab26ce3bf0e58a2ce1a525c6e
+ md5: 1b08cd684f34175e4514474793d44bcb
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc 15.2.0 he0feb66_18
+ constrains:
+ - libstdcxx-ng ==15.2.0=*_18
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 5852330
+ timestamp: 1771378262446
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-ng-15.2.0-hdf11a46_16.conda
+ sha256: 81f2f246c7533b41c5e0c274172d607829019621c4a0823b5c0b4a8c7028ee84
+ md5: 1b3152694d236cf233b76b8c56bf0eae
+ depends:
+ - libstdcxx 15.2.0 h934c35e_16
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 27300
+ timestamp: 1765256885128
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libstdcxx-ng-15.2.0-hdf11a46_18.conda
+ sha256: 3c902ffd673cb3c6ddde624cdb80f870b6c835f8bf28384b0016e7d444dd0145
+ md5: 6235adb93d064ecdf3d44faee6f468de
+ depends:
+ - libstdcxx 15.2.0 h934c35e_18
+ license: GPL-3.0-only WITH GCC-exception-3.1
+ license_family: GPL
+ purls: []
+ size: 27575
+ timestamp: 1771378314494
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libtiff-4.7.1-h8261f1e_0.conda
+ sha256: ddda0d7ee67e71e904a452010c73e32da416806f5cb9145fb62c322f97e717fb
+ md5: 72b531694ebe4e8aa6f5745d1015c1b4
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - lerc >=4.0.0,<5.0a0
+ - libdeflate >=1.24,<1.25.0a0
+ - libgcc >=14
+ - libjpeg-turbo >=3.1.0,<4.0a0
+ - liblzma >=5.8.1,<6.0a0
+ - libstdcxx >=14
+ - libwebp-base >=1.6.0,<2.0a0
+ - libzlib >=1.3.1,<2.0a0
+ - zstd >=1.5.7,<1.6.0a0
+ license: HPND
+ purls: []
+ size: 437211
+ timestamp: 1758278398952
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libuuid-2.41.3-h5347b49_0.conda
+ sha256: 1a7539cfa7df00714e8943e18de0b06cceef6778e420a5ee3a2a145773758aee
+ md5: db409b7c1720428638e7c0d509d3e1b5
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 40311
+ timestamp: 1766271528534
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libwebp-base-1.6.0-hd42ef1d_0.conda
+ sha256: 3aed21ab28eddffdaf7f804f49be7a7d701e8f0e46c856d801270b470820a37b
+ md5: aea31d2e5b1091feca96fcfe945c3cf9
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ constrains:
+ - libwebp 1.6.0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 429011
+ timestamp: 1752159441324
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libxcb-1.17.0-h8a09558_0.conda
+ sha256: 666c0c431b23c6cec6e492840b176dde533d48b7e6fb8883f5071223433776aa
+ md5: 92ed62436b625154323d40d5f2f11dd7
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - pthread-stubs
+ - xorg-libxau >=1.0.11,<2.0a0
+ - xorg-libxdmcp
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 395888
+ timestamp: 1727278577118
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libxcrypt-4.4.36-hd590300_1.conda
+ sha256: 6ae68e0b86423ef188196fff6207ed0c8195dd84273cb5623b85aa08033a410c
+ md5: 5aa797f8787fe7a17d1b0821485b5adc
+ depends:
+ - libgcc-ng >=12
+ license: LGPL-2.1-or-later
+ purls: []
+ size: 100393
+ timestamp: 1702724383534
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-2.15.1-h26afc86_0.conda
+ sha256: ec0735ae56c3549149eebd7dc22c0bed91fd50c02eaa77ff418613ddda190aa8
+ md5: e512be7dc1f84966d50959e900ca121f
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - icu >=75.1,<76.0a0
+ - libgcc >=14
+ - libiconv >=1.18,<2.0a0
+ - liblzma >=5.8.1,<6.0a0
+ - libxml2-16 2.15.1 ha9997c6_0
+ - libzlib >=1.3.1,<2.0a0
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 45283
+ timestamp: 1761015644057
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-16-2.15.1-ha9997c6_0.conda
+ sha256: 71436e72a286ef8b57d6f4287626ff91991eb03c7bdbe835280521791efd1434
+ md5: e7733bc6785ec009e47a224a71917e84
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - icu >=75.1,<76.0a0
+ - libgcc >=14
+ - libiconv >=1.18,<2.0a0
+ - liblzma >=5.8.1,<6.0a0
+ - libzlib >=1.3.1,<2.0a0
+ constrains:
+ - libxml2 2.15.1
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 556302
+ timestamp: 1761015637262
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libxml2-devel-2.15.1-h26afc86_0.conda
+ sha256: 7a01dde0807d0283ef6babb661cb750f63d7842f489b6e40d0af0f16951edf3e
+ md5: 1b92b7d1b901bd832f8279ef18cac1f4
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - icu >=75.1,<76.0a0
+ - libgcc >=14
+ - libiconv >=1.18,<2.0a0
+ - liblzma >=5.8.1,<6.0a0
+ - libxml2 2.15.1 h26afc86_0
+ - libxml2-16 2.15.1 ha9997c6_0
+ - libzlib >=1.3.1,<2.0a0
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 79667
+ timestamp: 1761015650428
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libzip-1.11.2-h6991a6a_0.conda
+ sha256: 991e7348b0f650d495fb6d8aa9f8c727bdf52dabf5853c0cc671439b160dce48
+ md5: a7b27c075c9b7f459f1c022090697cba
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - bzip2 >=1.0.8,<2.0a0
+ - libgcc >=13
+ - libzlib >=1.3.1,<2.0a0
+ - openssl >=3.3.2,<4.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 109043
+ timestamp: 1730442108429
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.1-hb9d3cd8_2.conda
+ sha256: d4bfe88d7cb447768e31650f06257995601f89076080e76df55e3112d4e47dc4
+ md5: edb0dca6bc32e4f4789199455a1dbeb8
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ constrains:
+ - zlib 1.3.1 *_2
+ license: Zlib
+ license_family: Other
+ purls: []
+ size: 60963
+ timestamp: 1727963148474
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/libzlib-1.3.1-h8359307_2.conda
+ sha256: ce34669eadaba351cd54910743e6a2261b67009624dbc7daeeafdef93616711b
+ md5: 369964e85dc26bfe78f41399b366c435
+ depends:
+ - __osx >=11.0
+ constrains:
+ - zlib 1.3.1 *_2
+ license: Zlib
+ license_family: Other
+ purls: []
+ size: 46438
+ timestamp: 1727963202283
+- conda: https://conda.anaconda.org/conda-forge/linux-64/libzopfli-1.0.3-h9c3ff4c_0.tar.bz2
+ sha256: ff94f30b2e86cbad6296cf3e5804d442d9e881f7ba8080d92170981662528c6e
+ md5: c66fe2d123249af7651ebde8984c51c2
+ depends:
+ - libgcc-ng >=9.3.0
+ - libstdcxx-ng >=9.3.0
+ license: Apache-2.0
+ license_family: Apache
+ purls: []
+ size: 168074
+ timestamp: 1607309189989
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/llvm-openmp-21.1.8-h4a912ad_0.conda
+ sha256: 56bcd20a0a44ddd143b6ce605700fdf876bcf5c509adc50bf27e76673407a070
+ md5: 206ad2df1b5550526e386087bef543c7
+ depends:
+ - __osx >=11.0
+ constrains:
+ - openmp 21.1.8|21.1.8.*
+ - intel-openmp <0.0a0
+ license: Apache-2.0 WITH LLVM-exception
+ license_family: APACHE
+ purls: []
+ size: 285974
+ timestamp: 1765964756583
+- conda: https://conda.anaconda.org/conda-forge/noarch/locket-1.0.0-pyhd8ed1ab_0.tar.bz2
+ sha256: 9afe0b5cfa418e8bdb30d8917c5a6cec10372b037924916f1f85b9f4899a67a6
+ md5: 91e27ef3d05cc772ce627e51cff111c4
+ depends:
+ - python >=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*
+ license: BSD-2-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/locket?source=hash-mapping
+ size: 8250
+ timestamp: 1650660473123
+- pypi: https://files.pythonhosted.org/packages/1f/c9/e0f8e4e6e8a69e5959b06499582dca6349db6769cc7fdfb8a02a7c75a9ae/lxml_stubs-0.5.1-py3-none-any.whl
+ name: lxml-stubs
+ version: 0.5.1
+ sha256: 1f689e5dbc4b9247cb09ae820c7d34daeb1fdbd1db06123814b856dae7787272
+ requires_dist:
+ - coverage[toml]>=7.2.5 ; extra == 'test'
+ - pytest>=7.3.0 ; extra == 'test'
+ - pytest-mypy-plugins>=1.10.1 ; extra == 'test'
+ - mypy>=1.2.0 ; extra == 'test'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/lz4-4.4.5-py310hde1b0b5_1.conda
+ sha256: 5dc79d66e7c85867c537ad13d276742314b3e6b87ab8b22ce7e1aac61ce6281e
+ md5: 4a20c97489a720287cf3d082f1e715c6
+ depends:
+ - python
+ - lz4-c
+ - libgcc >=14
+ - __glibc >=2.17,<3.0.a0
+ - lz4-c >=1.10.0,<1.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/lz4?source=hash-mapping
+ size: 41963
+ timestamp: 1765026389535
+- conda: https://conda.anaconda.org/conda-forge/linux-64/lz4-c-1.10.0-h5888daf_1.conda
+ sha256: 47326f811392a5fd3055f0f773036c392d26fdb32e4d8e7a8197eed951489346
+ md5: 9de5350a85c4a20c685259b889aa6393
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - libstdcxx >=13
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 167055
+ timestamp: 1733741040117
+- conda: https://conda.anaconda.org/conda-forge/linux-64/lzo-2.10-h280c20c_1002.conda
+ sha256: 5c6bbeec116e29f08e3dad3d0524e9bc5527098e12fc432c0e5ca53ea16337d4
+ md5: 45161d96307e3a447cc3eb5896cf6f8c
+ depends:
+ - libgcc >=14
+ - __glibc >=2.17,<3.0.a0
+ license: GPL-2.0-or-later
+ license_family: GPL
+ purls: []
+ size: 191060
+ timestamp: 1753889274283
+- conda: https://conda.anaconda.org/conda-forge/noarch/mapclassify-2.8.1-pyhd8ed1ab_1.conda
+ sha256: c498a016b233be5a7defee443733a82d5fe41b83016ca8a136876a64fd15564b
+ md5: c48bbb2bcc3f9f46741a7915d67e6839
+ depends:
+ - networkx >=2.7
+ - numpy >=1.23
+ - pandas >=1.4,!=1.5.0
+ - python >=3.9
+ - scikit-learn >=1.0
+ - scipy >=1.8
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/mapclassify?source=hash-mapping
+ size: 56772
+ timestamp: 1733731193211
+- pypi: https://files.pythonhosted.org/packages/59/1b/6ef961f543593969d25b2afe57a3564200280528caa9bd1082eecdd7b3bc/markdown-3.10.1-py3-none-any.whl
+ name: markdown
+ version: 3.10.1
+ sha256: 867d788939fe33e4b736426f5b9f651ad0c0ae0ecf89df0ca5d1176c70812fe3
+ requires_dist:
+ - coverage ; extra == 'testing'
+ - pyyaml ; extra == 'testing'
+ - mkdocs>=1.6 ; extra == 'docs'
+ - mkdocs-nature>=0.6 ; extra == 'docs'
+ - mdx-gh-links>=0.2 ; extra == 'docs'
+ - mkdocstrings[python]>=0.28.3 ; extra == 'docs'
+ - mkdocs-gen-files ; extra == 'docs'
+ - mkdocs-section-index ; extra == 'docs'
+ - mkdocs-literate-nav ; extra == 'docs'
+ requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl
+ name: markdown-it-py
+ version: 4.0.0
+ sha256: 87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147
+ requires_dist:
+ - mdurl~=0.1
+ - psutil ; extra == 'benchmarking'
+ - pytest ; extra == 'benchmarking'
+ - pytest-benchmark ; extra == 'benchmarking'
+ - commonmark~=0.9 ; extra == 'compare'
+ - markdown~=3.4 ; extra == 'compare'
+ - mistletoe~=1.0 ; extra == 'compare'
+ - mistune~=3.0 ; extra == 'compare'
+ - panflute~=2.3 ; extra == 'compare'
+ - markdown-it-pyrs ; extra == 'compare'
+ - linkify-it-py>=1,<3 ; extra == 'linkify'
+ - mdit-py-plugins>=0.5.0 ; extra == 'plugins'
+ - gprof2dot ; extra == 'profiling'
+ - mdit-py-plugins>=0.5.0 ; extra == 'rtd'
+ - myst-parser ; extra == 'rtd'
+ - pyyaml ; extra == 'rtd'
+ - sphinx ; extra == 'rtd'
+ - sphinx-copybutton ; extra == 'rtd'
+ - sphinx-design ; extra == 'rtd'
+ - sphinx-book-theme~=1.0 ; extra == 'rtd'
+ - jupyter-sphinx ; extra == 'rtd'
+ - ipykernel ; extra == 'rtd'
+ - coverage ; extra == 'testing'
+ - pytest ; extra == 'testing'
+ - pytest-cov ; extra == 'testing'
+ - pytest-regressions ; extra == 'testing'
+ - requests ; extra == 'testing'
+ requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/markupsafe-3.0.3-py310h3406613_0.conda
+ sha256: b3894b37cab530d1adab5b9ce39a1b9f28040403cc0042b77e04a2f227a447de
+ md5: 8854df4fb4e37cc3ea0a024e48c9c180
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ constrains:
+ - jinja2 >=3.0.0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/markupsafe?source=hash-mapping
+ size: 23673
+ timestamp: 1759055396627
+- conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.10.1-py310h68603db_0.conda
+ sha256: f211079f3346a225ba0d1a4754eb856ed3c0bdbf17d6502c55390d22a2c86cb5
+ md5: 29cf3f5959afb841eda926541f26b0fb
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - contourpy >=1.0.1
+ - cycler >=0.10
+ - fonttools >=4.22.0
+ - freetype >=2.12.1,<3.0a0
+ - kiwisolver >=1.3.1
+ - libgcc >=13
+ - libstdcxx >=13
+ - numpy >=1.19,<3
+ - numpy >=1.23
+ - packaging >=20.0
+ - pillow >=8
+ - pyparsing >=2.3.1
+ - python >=3.10,<3.11.0a0
+ - python-dateutil >=2.7
+ - python_abi 3.10.* *_cp310
+ - qhull >=2020.2,<2020.3.0a0
+ - tk >=8.6.13,<8.7.0a0
+ license: PSF-2.0
+ license_family: PSF
+ purls:
+ - pkg:pypi/matplotlib?source=hash-mapping
+ size: 7310356
+ timestamp: 1740781078080
+- pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl
+ name: mdurl
+ version: 0.1.2
+ sha256: 84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8
+ requires_python: '>=3.7'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/minizip-4.0.10-h05a5f5f_0.conda
+ sha256: 0c3700d15377156937ddc89a856527ad77e7cf3fd73cb0dffc75fce8030ddd16
+ md5: da01bb40572e689bd1535a5cee6b1d68
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - bzip2 >=1.0.8,<2.0a0
+ - libgcc >=13
+ - libiconv >=1.18,<2.0a0
+ - liblzma >=5.8.1,<6.0a0
+ - libstdcxx >=13
+ - libzlib >=1.3.1,<2.0a0
+ - openssl >=3.5.0,<4.0a0
+ - zstd >=1.5.7,<1.6.0a0
+ license: Zlib
+ license_family: Other
+ purls: []
+ size: 93471
+ timestamp: 1746450475308
+- pypi: https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz
+ name: ml-dtypes
+ version: 0.5.4
+ sha256: 8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453
+ requires_dist:
+ - numpy>=1.21
+ - numpy>=1.21.2 ; python_full_version >= '3.10'
+ - numpy>=1.23.3 ; python_full_version >= '3.11'
+ - numpy>=1.26.0 ; python_full_version >= '3.12'
+ - numpy>=2.1.0 ; python_full_version >= '3.13'
+ - absl-py ; extra == 'dev'
+ - pytest ; extra == 'dev'
+ - pytest-xdist ; extra == 'dev'
+ - pylint>=2.6.0 ; extra == 'dev'
+ - pyink ; extra == 'dev'
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/a4/8e/469e5a4a2f5855992e425f3cb33804cc07bf18d48f2db061aec61ce50270/more_itertools-10.8.0-py3-none-any.whl
+ name: more-itertools
+ version: 10.8.0
+ sha256: 52d4362373dcf7c52546bc4af9a86ee7c4579df9a8dc268be0a2f949d376cc9b
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/msgpack-python-1.1.2-py310h03d9f68_1.conda
+ sha256: 61cf3572d6afa3fa711c5f970a832783d2c281facb7b3b946a6b71a0bac2c592
+ md5: 5eea9d8f8fcf49751dab7927cb0dfc3f
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libstdcxx >=14
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: Apache-2.0
+ license_family: Apache
+ purls:
+ - pkg:pypi/msgpack?source=hash-mapping
+ size: 95105
+ timestamp: 1762504073388
+- conda: https://conda.anaconda.org/conda-forge/linux-64/multidict-6.7.0-py310h3406613_0.conda
+ sha256: 23c5d03f2fa243469724a96536c87b19050c4d915c25faf81c134a4cdc139344
+ md5: a1d7c51debf0cabe04dea061fecd9bf8
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - typing-extensions >=4.1.0
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/multidict?source=hash-mapping
+ size: 91323
+ timestamp: 1765460776858
+- conda: https://conda.anaconda.org/conda-forge/noarch/munkres-1.1.4-pyhd8ed1ab_1.conda
+ sha256: d09c47c2cf456de5c09fa66d2c3c5035aa1fa228a1983a433c47b876aa16ce90
+ md5: 37293a85a0f4f77bbd9cf7aaefc62609
+ depends:
+ - python >=3.9
+ license: Apache-2.0
+ license_family: Apache
+ purls:
+ - pkg:pypi/munkres?source=hash-mapping
+ size: 15851
+ timestamp: 1749895533014
+- pypi: https://files.pythonhosted.org/packages/2a/0d/93c2e4a287f74ef11a66fb6d49c7a9f05e47b0a4399040e6719b57f500d2/mypy-1.19.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl
+ name: mypy
+ version: 1.19.1
+ sha256: de759aafbae8763283b2ee5869c7255391fbc4de3ff171f8f030b5ec48381b74
+ requires_dist:
+ - typing-extensions>=4.6.0
+ - mypy-extensions>=1.0.0
+ - pathspec>=0.9.0
+ - tomli>=1.1.0 ; python_full_version < '3.11'
+ - librt>=0.6.2 ; platform_python_implementation != 'PyPy'
+ - psutil>=4.0 ; extra == 'dmypy'
+ - setuptools>=50 ; extra == 'mypyc'
+ - lxml ; extra == 'reports'
+ - pip ; extra == 'install-types'
+ - orjson ; extra == 'faster-cache'
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl
+ name: mypy-extensions
+ version: 1.1.0
+ sha256: 1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505
+ requires_python: '>=3.8'
+- pypi: https://files.pythonhosted.org/packages/b2/bc/465daf1de06409cdd4532082806770ee0d8d7df434da79c76564d0f69741/namex-0.1.0-py3-none-any.whl
+ name: namex
+ version: 0.1.0
+ sha256: e2012a474502f1e2251267062aae3114611f07df4224b6e06334c57b0f2ce87c
+- conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.5-h2d0b736_3.conda
+ sha256: 3fde293232fa3fca98635e1167de6b7c7fda83caf24b9d6c91ec9eefb4f4d586
+ md5: 47e340acb35de30501a76c7c799c41d7
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ license: X11 AND BSD-3-Clause
+ purls: []
+ size: 891641
+ timestamp: 1738195959188
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/ncurses-6.5-h5e97a16_3.conda
+ sha256: 2827ada40e8d9ca69a153a45f7fd14f32b2ead7045d3bbb5d10964898fe65733
+ md5: 068d497125e4bf8a66bf707254fff5ae
+ depends:
+ - __osx >=11.0
+ license: X11 AND BSD-3-Clause
+ purls: []
+ size: 797030
+ timestamp: 1738196177597
+- conda: https://conda.anaconda.org/conda-forge/linux-64/netcdf4-1.7.4-nompi_py310hd27e1a9_102.conda
+ sha256: 0b00d5f69878dbd9b584bff42c71c6b3d806c633e90b75792069065172446136
+ md5: 424961f9f166613016bc4737796ef44b
+ depends:
+ - python
+ - certifi
+ - cftime
+ - numpy
+ - hdf5
+ - libnetcdf
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libnetcdf >=4.9.3,<4.9.4.0a0
+ - libzlib >=1.3.1,<2.0a0
+ - python_abi 3.10.* *_cp310
+ - hdf5 >=1.14.6,<1.14.7.0a0
+ - numpy >=1.21,<3
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/netcdf4?source=hash-mapping
+ size: 1147810
+ timestamp: 1768552449353
+- conda: https://conda.anaconda.org/conda-forge/noarch/networkx-3.4.2-pyh267e887_2.conda
+ sha256: 39625cd0c9747fa5c46a9a90683b8997d8b9649881b3dc88336b13b7bdd60117
+ md5: fd40bf7f7f4bc4b647dc8512053d9873
+ depends:
+ - python >=3.10
+ - python
+ constrains:
+ - numpy >=1.24
+ - scipy >=1.10,!=1.11.0,!=1.11.1
+ - matplotlib >=3.7
+ - pandas >=2.0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/networkx?source=hash-mapping
+ size: 1265008
+ timestamp: 1731521053408
+- pypi: https://files.pythonhosted.org/packages/42/0f/c76bf3dba22c73c38e9b1113b017cf163f7696f50e003404ec5ecdb1e8a6/nh3-0.3.2-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ name: nh3
+ version: 0.3.2
+ sha256: 7bb18403f02b655a1bbe4e3a4696c2ae1d6ae8f5991f7cacb684b1ae27e6c9f7
+ requires_python: '>=3.8'
+- pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl
+ name: nodeenv
+ version: 1.10.0
+ sha256: 5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827
+ requires_python: '>=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/numcodecs-0.13.1-py310h5eaa309_0.conda
+ sha256: 70cb0fa431ba9e75ef36d94f35324089dfa7da8f967e9c758f60e08aaf29b732
+ md5: a3e9933fc59e8bcd2aa20753fb56db42
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - libstdcxx >=13
+ - msgpack-python
+ - numpy >=1.19,<3
+ - numpy >=1.7
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/numcodecs?source=hash-mapping
+ size: 802894
+ timestamp: 1728547783947
+- conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-1.26.4-py310hb13e2d6_0.conda
+ sha256: 028fe2ea8e915a0a032b75165f11747770326f3d767e642880540c60a3256425
+ md5: 6593de64c935768b6bad3e19b3e978be
+ depends:
+ - libblas >=3.9.0,<4.0a0
+ - libcblas >=3.9.0,<4.0a0
+ - libgcc-ng >=12
+ - liblapack >=3.9.0,<4.0a0
+ - libstdcxx-ng >=12
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ constrains:
+ - numpy-base <0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/numpy?source=hash-mapping
+ size: 7009070
+ timestamp: 1707225917496
+- conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-1.26.4-py312heda63a1_0.conda
+ sha256: fe3459c75cf84dcef6ef14efcc4adb0ade66038ddd27cadb894f34f4797687d8
+ md5: d8285bea2a350f63fab23bf460221f3f
+ depends:
+ - libblas >=3.9.0,<4.0a0
+ - libcblas >=3.9.0,<4.0a0
+ - libgcc-ng >=12
+ - liblapack >=3.9.0,<4.0a0
+ - libstdcxx-ng >=12
+ - python >=3.12,<3.13.0a0
+ - python_abi 3.12.* *_cp312
+ constrains:
+ - numpy-base <0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/numpy?source=hash-mapping
+ size: 7484186
+ timestamp: 1707225809722
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/numpy-1.26.4-py312h8442bc7_0.conda
+ sha256: c8841d6d6f61fd70ca80682efbab6bdb8606dc77c68d8acabfbd7c222054f518
+ md5: d83fc83d589e2625a3451c9a7e21047c
+ depends:
+ - libblas >=3.9.0,<4.0a0
+ - libcblas >=3.9.0,<4.0a0
+ - libcxx >=16
+ - liblapack >=3.9.0,<4.0a0
+ - python >=3.12,<3.13.0a0
+ - python >=3.12,<3.13.0a0 *_cpython
+ - python_abi 3.12.* *_cp312
+ constrains:
+ - numpy-base <0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/numpy?source=hash-mapping
+ size: 6073136
+ timestamp: 1707226249608
+- pypi: https://files.pythonhosted.org/packages/f9/43/a5365b345c989d9f7de2f2406e59b4a792ba3541b5d47edda5031b2730a6/nvidia_cublas_cu12-12.5.3.2-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-cublas-cu12
+ version: 12.5.3.2
+ sha256: ca070ad70e9fa6654084575d01bd001f30cc4665e33d4bb9fc8e0f321caa034b
+ requires_python: '>=3'
+- pypi: https://files.pythonhosted.org/packages/3a/64/81515a1b5872dc0fc817deaec2bdba160dc50188e0d53b907d10c6e6d568/nvidia_cuda_cupti_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-cuda-cupti-cu12
+ version: 12.5.82
+ sha256: bde77a5feb66752ec61db2adfe47f56b941842825b4c7e2068aff27c9d107953
+ requires_python: '>=3'
+- pypi: https://files.pythonhosted.org/packages/44/cc/36363057676e6140d0bdb07fa6df5419b68203c5cba8c412b5600fd0d105/nvidia_cuda_nvcc_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-cuda-nvcc-cu12
+ version: 12.5.82
+ sha256: b03e545b8e8c3ce7ebcd7fc44063180ff52ff01d064ece2127ed90a04ef12cd0
+ requires_python: '>=3'
+- pypi: https://files.pythonhosted.org/packages/19/d1/342c2bcf6172db65fcbd9102f9941876e730d98977e69c00df85940fa8ce/nvidia_cuda_nvrtc_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-cuda-nvrtc-cu12
+ version: 12.5.82
+ sha256: 3dbd97b0104b4bfbc3c4f8c79cd2496307c89c43c29a9f83125f1d76296ff3fd
+ requires_python: '>=3'
+- pypi: https://files.pythonhosted.org/packages/71/05/80f3fe49e9905570bc27aea9493baa1891c3780a7fc4e1f872c7902df066/nvidia_cuda_runtime_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-cuda-runtime-cu12
+ version: 12.5.82
+ sha256: 3e79a060e126df40fd3a068f3f787eb000fa51b251ec6cd97d09579632687115
+ requires_python: '>=3'
+- pypi: https://files.pythonhosted.org/packages/8e/56/bb5c08a8d401fc1b21a10e9c58907e70e8f18bdfca34b7ecfb87bbcdad63/nvidia_cudnn_cu12-9.3.0.75-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-cudnn-cu12
+ version: 9.3.0.75
+ sha256: 9ad9c6929ebb5295eb4a1728024666d1c88283373e265a0c5c883e6f9d5cd76d
+ requires_dist:
+ - nvidia-cublas-cu12
+ requires_python: '>=3'
+- pypi: https://files.pythonhosted.org/packages/e4/85/f18c88f63489cdced17b06d3b627adca8add7d7b8cce8c11213e93a902b4/nvidia_cufft_cu12-11.2.3.61-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-cufft-cu12
+ version: 11.2.3.61
+ sha256: 9a6e8df162585750f61983a638104a48c756aa13f9f48e19ab079b38e3c828b8
+ requires_dist:
+ - nvidia-nvjitlink-cu12
+ requires_python: '>=3'
+- pypi: https://files.pythonhosted.org/packages/f8/ce/cc6daf7820804ee7f11a0352a1f0fd59cec5f12e904f5bbaee6d928ffdaf/nvidia_curand_cu12-10.3.6.82-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-curand-cu12
+ version: 10.3.6.82
+ sha256: 2823fb27de4e44dbb22394a6adf53aa6e1b013aca0f8c22867d1cfae58405536
+ requires_python: '>=3'
+- pypi: https://files.pythonhosted.org/packages/33/73/57fbf55b3f378a73faecde397a0927ea205b458f06573dfb191b7d9fd1d3/nvidia_cusolver_cu12-11.6.3.83-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-cusolver-cu12
+ version: 11.6.3.83
+ sha256: 93cfafacde4428b71778eeb092ec615a02a3d05404da1bcf91c53e3fa1bce42b
+ requires_dist:
+ - nvidia-cublas-cu12
+ - nvidia-nvjitlink-cu12
+ - nvidia-cusparse-cu12
+ requires_python: '>=3'
+- pypi: https://files.pythonhosted.org/packages/b3/29/03726191334fa523d9654e3dacca5cc152f24bb9fa1a5721e4dddac7a8c5/nvidia_cusparse_cu12-12.5.1.3-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-cusparse-cu12
+ version: 12.5.1.3
+ sha256: 016df8e993c437e8301e62739f01775cba988fd5253cd4c64173f8e8d2f8e752
+ requires_dist:
+ - nvidia-nvjitlink-cu12
+ requires_python: '>=3'
+- pypi: https://files.pythonhosted.org/packages/ed/1f/6482380ec8dcec4894e7503490fc536d846b0d59694acad9cf99f27d0e7d/nvidia_nccl_cu12-2.23.4-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-nccl-cu12
+ version: 2.23.4
+ sha256: b097258d9aab2fa9f686e33c6fe40ae57b27df60cedbd15d139701bb5509e0c1
+ requires_python: '>=3'
+- pypi: https://files.pythonhosted.org/packages/75/bc/e0d0dbb85246a086ab14839979039647bce501d8c661a159b8b019d987b7/nvidia_nvjitlink_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl
+ name: nvidia-nvjitlink-cu12
+ version: 12.5.82
+ sha256: f9b37bc5c8cf7509665cb6ada5aaa0ce65618f2332b7d3e78e9790511f111212
+ requires_python: '>=3'
+- conda: https://conda.anaconda.org/conda-forge/noarch/oauthlib-3.3.1-pyhd8ed1ab_0.conda
+ sha256: dfa8222df90736fa13f8896f5a573a50273af8347542d412c3bd1230058e56a5
+ md5: d4f3f31ee39db3efecb96c0728d4bdbf
+ depends:
+ - blinker
+ - cryptography
+ - pyjwt >=1.0.0
+ - python >=3.9
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/oauthlib?source=hash-mapping
+ size: 102059
+ timestamp: 1750415349440
+- pypi: https://files.pythonhosted.org/packages/5f/7d/9ec5967f3e2915fbc441f72c3892a7f0fb3618e3ae5c8a44181ce4aa641c/obstore-0.8.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ name: obstore
+ version: 0.8.2
+ sha256: 8ccf0f03a7fe453fb8640611c922bce19f021c6aaeee6ee44d6d8fb57db6be48
+ requires_dist:
+ - typing-extensions ; python_full_version < '3.13'
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.4-h55fea9a_0.conda
+ sha256: 3900f9f2dbbf4129cf3ad6acf4e4b6f7101390b53843591c53b00f034343bc4d
+ md5: 11b3379b191f63139e29c0d19dee24cd
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libpng >=1.6.50,<1.7.0a0
+ - libstdcxx >=14
+ - libtiff >=4.7.1,<4.8.0a0
+ - libzlib >=1.3.1,<2.0a0
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 355400
+ timestamp: 1758489294972
+- conda: https://conda.anaconda.org/conda-forge/linux-64/openssl-3.6.1-h35e630c_1.conda
+ sha256: 44c877f8af015332a5d12f5ff0fb20ca32f896526a7d0cdb30c769df1144fb5c
+ md5: f61eb8cd60ff9057122a3d338b99c00f
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - ca-certificates
+ - libgcc >=14
+ license: Apache-2.0
+ license_family: Apache
+ purls: []
+ size: 3164551
+ timestamp: 1769555830639
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/openssl-3.6.1-hd24854e_1.conda
+ sha256: 361f5c5e60052abc12bdd1b50d7a1a43e6a6653aab99a2263bf2288d709dcf67
+ md5: f4f6ad63f98f64191c3e77c5f5f29d76
+ depends:
+ - __osx >=11.0
+ - ca-certificates
+ license: Apache-2.0
+ license_family: Apache
+ purls: []
+ size: 3104268
+ timestamp: 1769556384749
+- pypi: https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl
+ name: opt-einsum
+ version: 3.4.0
+ sha256: 69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd
+ requires_python: '>=3.8'
+- pypi: https://files.pythonhosted.org/packages/83/8e/09d899ad531d50b79aa24e7558f604980fe4048350172e643bb1b9983aec/optree-0.18.0.tar.gz
+ name: optree
+ version: 0.18.0
+ sha256: 3804fb6ddc923855db2dc4805b4524c66e00f1ef30b166be4aadd52822b13e06
+ requires_dist:
+ - typing-extensions>=4.6.0
+ - typing-extensions>=4.12.0 ; python_full_version >= '3.13'
+ - jax ; extra == 'jax'
+ - numpy ; extra == 'numpy'
+ - torch ; extra == 'torch'
+ - ruff ; extra == 'lint'
+ - pylint[spelling] ; extra == 'lint'
+ - mypy ; extra == 'lint'
+ - doc8 ; extra == 'lint'
+ - pyenchant ; extra == 'lint'
+ - xdoctest ; extra == 'lint'
+ - cpplint ; extra == 'lint'
+ - pre-commit ; extra == 'lint'
+ - pytest ; extra == 'test'
+ - pytest-cov ; extra == 'test'
+ - covdefaults ; extra == 'test'
+ - rich ; extra == 'test'
+ - typing-extensions==4.6.0 ; python_full_version < '3.13' and sys_platform == 'linux' and extra == 'test'
+ - typing-extensions==4.6.0 ; python_full_version < '3.13' and sys_platform == 'darwin' and extra == 'test'
+ - typing-extensions==4.6.0 ; python_full_version < '3.13' and sys_platform == 'win32' and extra == 'test'
+ - typing-extensions==4.12.0 ; python_full_version >= '3.13' and sys_platform == 'linux' and extra == 'test'
+ - typing-extensions==4.12.0 ; python_full_version >= '3.13' and sys_platform == 'darwin' and extra == 'test'
+ - typing-extensions==4.12.0 ; python_full_version >= '3.13' and sys_platform == 'win32' and extra == 'test'
+ - sphinx ; extra == 'docs'
+ - sphinx-autoapi ; extra == 'docs'
+ - sphinx-autobuild ; extra == 'docs'
+ - sphinx-copybutton ; extra == 'docs'
+ - sphinx-rtd-theme ; extra == 'docs'
+ - sphinxcontrib-bibtex ; extra == 'docs'
+ - sphinx-autodoc-typehints ; extra == 'docs'
+ - docutils ; extra == 'docs'
+ - jax[cpu] ; extra == 'docs'
+ - numpy ; extra == 'docs'
+ - torch ; extra == 'docs'
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl
+ name: packaging
+ version: '26.0'
+ sha256: b36f1fef9334a5588b4166f8bcd26a14e521f2b55e6b9de3aaa80d3ff7a37529
+ requires_python: '>=3.8'
+- conda: https://conda.anaconda.org/conda-forge/noarch/packaging-26.0-pyhcf101f3_0.conda
+ sha256: c1fc0f953048f743385d31c468b4a678b3ad20caffdeaa94bed85ba63049fd58
+ md5: b76541e68fea4d511b1ac46a28dcd2c6
+ depends:
+ - python >=3.8
+ - python
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/packaging?source=compressed-mapping
+ size: 72010
+ timestamp: 1769093650580
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pandas-2.3.3-py310h0158d43_2.conda
+ sha256: b9e88fa02fd5e99f54c168df622eda9ddf898cc15e631179963aca51d97244bf
+ md5: 0610ed073acc4737d036125a5a6dbae2
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libstdcxx >=14
+ - numpy >=1.21,<3
+ - numpy >=1.22.4
+ - python >=3.10,<3.11.0a0
+ - python-dateutil >=2.8.2
+ - python-tzdata >=2022.7
+ - python_abi 3.10.* *_cp310
+ - pytz >=2020.1
+ constrains:
+ - odfpy >=1.4.1
+ - pyarrow >=10.0.1
+ - pyqt5 >=5.15.9
+ - numexpr >=2.8.4
+ - fsspec >=2022.11.0
+ - bottleneck >=1.3.6
+ - beautifulsoup4 >=4.11.2
+ - pandas-gbq >=0.19.0
+ - s3fs >=2022.11.0
+ - gcsfs >=2022.11.0
+ - sqlalchemy >=2.0.0
+ - pytables >=3.8.0
+ - html5lib >=1.1
+ - python-calamine >=0.1.7
+ - lxml >=4.9.2
+ - qtpy >=2.3.0
+ - scipy >=1.10.0
+ - numba >=0.56.4
+ - openpyxl >=3.1.0
+ - blosc >=1.21.3
+ - pyreadstat >=1.2.0
+ - zstandard >=0.19.0
+ - xarray >=2022.12.0
+ - matplotlib >=3.6.3
+ - tabulate >=0.9.0
+ - fastparquet >=2022.12.0
+ - psycopg2 >=2.9.6
+ - xlsxwriter >=3.0.5
+ - xlrd >=2.0.1
+ - tzdata >=2022.7
+ - pyxlsb >=1.0.10
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/pandas?source=hash-mapping
+ size: 12391209
+ timestamp: 1764615007370
+- pypi: https://files.pythonhosted.org/packages/d1/c6/df1fe324248424f77b89371116dab5243db7f052c32cc9fe7442ad9c5f75/pandas_stubs-2.3.3.260113-py3-none-any.whl
+ name: pandas-stubs
+ version: 2.3.3.260113
+ sha256: ec070b5c576e1badf12544ae50385872f0631fc35d99d00dc598c2954ec564d3
+ requires_dist:
+ - numpy>=1.23.5
+ - types-pytz>=2022.1.1
+ requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/noarch/partd-1.4.2-pyhd8ed1ab_0.conda
+ sha256: 472fc587c63ec4f6eba0cc0b06008a6371e0a08a5986de3cf4e8024a47b4fe6c
+ md5: 0badf9c54e24cecfb0ad2f99d680c163
+ depends:
+ - locket
+ - python >=3.9
+ - toolz
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/partd?source=hash-mapping
+ size: 20884
+ timestamp: 1715026639309
+- pypi: https://files.pythonhosted.org/packages/ef/3c/2c197d226f9ea224a9ab8d197933f9da0ae0aac5b6e0f884e2b8d9c8e9f7/pathspec-1.0.4-py3-none-any.whl
+ name: pathspec
+ version: 1.0.4
+ sha256: fb6ae2fd4e7c921a165808a552060e722767cfa526f99ca5156ed2ce45a5c723
+ requires_dist:
+ - hyperscan>=0.7 ; extra == 'hyperscan'
+ - typing-extensions>=4 ; extra == 'optional'
+ - google-re2>=1.1 ; extra == 're2'
+ - pytest>=9 ; extra == 'tests'
+ - typing-extensions>=4.15 ; extra == 'tests'
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pcre2-10.46-h1321c63_0.conda
+ sha256: 5c7380c8fd3ad5fc0f8039069a45586aa452cf165264bc5a437ad80397b32934
+ md5: 7fa07cb0fb1b625a089ccc01218ee5b1
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - bzip2 >=1.0.8,<2.0a0
+ - libgcc >=14
+ - libzlib >=1.3.1,<2.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 1209177
+ timestamp: 1756742976157
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pillow-12.0.0-py310h049bd52_1.conda
+ sha256: 7d111dc3a935750b676eba51e5113f968ca9a43e88fb7269e80ba57fcf9ae072
+ md5: 4d2db823c7f41a04aad4d0f2202ed2db
+ depends:
+ - python
+ - libgcc >=14
+ - __glibc >=2.17,<3.0.a0
+ - libwebp-base >=1.6.0,<2.0a0
+ - python_abi 3.10.* *_cp310
+ - libfreetype >=2.14.1
+ - libfreetype6 >=2.14.1
+ - libtiff >=4.7.1,<4.8.0a0
+ - libjpeg-turbo >=3.1.2,<4.0a0
+ - libxcb >=1.17.0,<2.0a0
+ - openjpeg >=2.5.4,<3.0a0
+ - lcms2 >=2.17,<3.0a0
+ - tk >=8.6.13,<8.7.0a0
+ - zlib-ng >=2.2.5,<2.3.0a0
+ license: HPND
+ purls:
+ - pkg:pypi/pillow?source=hash-mapping
+ size: 882782
+ timestamp: 1764033139041
+- conda: https://conda.anaconda.org/conda-forge/noarch/platformdirs-4.5.1-pyhcf101f3_0.conda
+ sha256: 04c64fb78c520e5c396b6e07bc9082735a5cc28175dbe23138201d0a9441800b
+ md5: 1bd2e65c8c7ef24f4639ae6e850dacc2
+ depends:
+ - python >=3.10
+ - python
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/platformdirs?source=hash-mapping
+ size: 23922
+ timestamp: 1764950726246
+- pypi: https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl
+ name: pluggy
+ version: 1.6.0
+ sha256: e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746
+ requires_dist:
+ - pre-commit ; extra == 'dev'
+ - tox ; extra == 'dev'
+ - pytest ; extra == 'testing'
+ - pytest-benchmark ; extra == 'testing'
+ - coverage ; extra == 'testing'
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/noarch/pooch-1.8.2-pyhd8ed1ab_3.conda
+ sha256: 032405adb899ba7c7cc24d3b4cd4e7f40cf24ac4f253a8e385a4f44ccb5e0fc6
+ md5: d2bbbd293097e664ffb01fc4cdaf5729
+ depends:
+ - packaging >=20.0
+ - platformdirs >=2.5.0
+ - python >=3.9
+ - requests >=2.19.0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/pooch?source=hash-mapping
+ size: 55588
+ timestamp: 1754941801129
+- pypi: https://files.pythonhosted.org/packages/5d/19/fd3ef348460c80af7bb4669ea7926651d1f95c23ff2df18b9d24bab4f3fa/pre_commit-4.5.1-py2.py3-none-any.whl
+ name: pre-commit
+ version: 4.5.1
+ sha256: 3b3afd891e97337708c1674210f8eba659b52a38ea5f822ff142d10786221f77
+ requires_dist:
+ - cfgv>=2.0.0
+ - identify>=1.0.0
+ - nodeenv>=0.11.1
+ - pyyaml>=5.1
+ - virtualenv>=20.10.0
+ requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/proj-9.6.2-h18fbb6c_2.conda
+ sha256: c1c9e38646a2d07007844625c8dea82404c8785320f8a6326b9338f8870875d0
+ md5: 1aeede769ec2fa0f474f8b73a7ac057f
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libcurl >=8.14.1,<9.0a0
+ - libgcc >=14
+ - libsqlite >=3.50.4,<4.0a0
+ - libstdcxx >=14
+ - libtiff >=4.7.0,<4.8.0a0
+ - sqlite
+ constrains:
+ - proj4 ==999999999999
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 3240415
+ timestamp: 1754927975218
+- conda: https://conda.anaconda.org/conda-forge/linux-64/propcache-0.3.1-py310h89163eb_0.conda
+ sha256: 3dbf885bb1eb0e7a5eb3779165517abdb98d53871b36690041f6a366cc501738
+ md5: e768486f2be3f50126bf9a54331221d1
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/propcache?source=hash-mapping
+ size: 53576
+ timestamp: 1744525075233
+- conda: https://conda.anaconda.org/conda-forge/noarch/proto-plus-1.27.1-pyhd8ed1ab_0.conda
+ sha256: 0d91d1687ed442f9b8ffd6792cf3127529a088e020fac7915e0c40cc23be3037
+ md5: c9b8e02d974817913ab94dae12c7340b
+ depends:
+ - protobuf >=3.19.0,<7.0.0
+ - python >=3.10
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/proto-plus?source=hash-mapping
+ size: 43555
+ timestamp: 1770104240191
+- conda: https://conda.anaconda.org/conda-forge/linux-64/protobuf-4.21.12-py310heca2aa9_0.conda
+ sha256: 38808ff5a9b724d00b12f14ca5c83b40a7b7ada4715a8e8c3b64bf73407bbe5f
+ md5: 90bb7e1b729c4b50272cf78be97ab912
+ depends:
+ - libgcc-ng >=12
+ - libprotobuf 3.21.12.*
+ - libprotobuf >=3.21.12,<3.22.0a0
+ - libstdcxx-ng >=12
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - setuptools
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/protobuf?source=hash-mapping
+ size: 323516
+ timestamp: 1671796375324
+- conda: https://conda.anaconda.org/conda-forge/linux-64/psutil-7.2.2-py310h139afa4_0.conda
+ sha256: 3a6d46033ebad3e69ded3f76852b9c378c2cff632f57421b5926c6add1bae475
+ md5: d210342acdb8e3ca6434295497c10b7c
+ depends:
+ - python
+ - libgcc >=14
+ - __glibc >=2.17,<3.0.a0
+ - python_abi 3.10.* *_cp310
+ license: BSD-3-Clause
+ purls:
+ - pkg:pypi/psutil?source=compressed-mapping
+ size: 179015
+ timestamp: 1769678154886
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda
+ sha256: 9c88f8c64590e9567c6c80823f0328e58d3b1efb0e1c539c0315ceca764e0973
+ md5: b3c17d95b5a10c6e64a21fa17573e70e
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 8252
+ timestamp: 1726802366959
+- conda: https://conda.anaconda.org/conda-forge/noarch/pyasn1-0.6.2-pyhd8ed1ab_0.conda
+ sha256: 2b6e22e97af814153c0a993ea66811de9db05b2a6946dcb97a3953af13c33a80
+ md5: c203d401759f448f9e792974e055bcdc
+ depends:
+ - python >=3.10
+ license: BSD-2-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/pyasn1?source=compressed-mapping
+ size: 63471
+ timestamp: 1769186345593
+- conda: https://conda.anaconda.org/conda-forge/noarch/pyasn1-modules-0.4.2-pyhd8ed1ab_0.conda
+ sha256: 5495061f5d3d6b82b74d400273c586e7c1f1700183de1d2d1688e900071687cb
+ md5: c689b62552f6b63f32f3322e463f3805
+ depends:
+ - pyasn1 >=0.6.1,<0.7.0
+ - python >=3.9
+ license: BSD-2-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/pyasn1-modules?source=hash-mapping
+ size: 95990
+ timestamp: 1743436137965
+- conda: https://conda.anaconda.org/conda-forge/noarch/pycparser-2.22-pyh29332c3_1.conda
+ sha256: 79db7928d13fab2d892592223d7570f5061c192f27b9febd1a418427b719acc6
+ md5: 12c566707c80111f9799308d9e265aef
+ depends:
+ - python >=3.9
+ - python
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/pycparser?source=hash-mapping
+ size: 110100
+ timestamp: 1733195786147
+- conda: https://conda.anaconda.org/conda-forge/noarch/pydantic-2.12.5-pyhcf101f3_1.conda
+ sha256: 868569d9505b7fe246c880c11e2c44924d7613a8cdcc1f6ef85d5375e892f13d
+ md5: c3946ed24acdb28db1b5d63321dbca7d
+ depends:
+ - typing-inspection >=0.4.2
+ - typing_extensions >=4.14.1
+ - python >=3.10
+ - typing-extensions >=4.6.1
+ - annotated-types >=0.6.0
+ - pydantic-core ==2.41.5
+ - python
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/pydantic?source=hash-mapping
+ size: 340482
+ timestamp: 1764434463101
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pydantic-core-2.41.5-py310hd8f68c5_1.conda
+ sha256: feb22e14b42321f3791ea24d726b7007e489a61ba72c98e22c7ec964671bb08a
+ md5: eaab3d18db92c656e5e2508de78f4a8c
+ depends:
+ - python
+ - typing-extensions >=4.6.0,!=4.7.0
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - python_abi 3.10.* *_cp310
+ constrains:
+ - __glibc >=2.17
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/pydantic-core?source=hash-mapping
+ size: 1933356
+ timestamp: 1762989015032
+- pypi: https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl
+ name: pygments
+ version: 2.19.2
+ sha256: 86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b
+ requires_dist:
+ - colorama>=0.4.6 ; extra == 'windows-terminal'
+ requires_python: '>=3.8'
+- conda: https://conda.anaconda.org/conda-forge/noarch/pyjwt-2.11.0-pyhd8ed1ab_0.conda
+ sha256: ac605d7fa239f78c508b47f2a0763236eef8d52b53852b84d784b598f92a1573
+ md5: f9517d2fe1501919d7a236aba73409bb
+ depends:
+ - python >=3.10
+ constrains:
+ - cryptography >=3.4.0
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/pyjwt?source=hash-mapping
+ size: 30144
+ timestamp: 1769858771741
+- conda: https://conda.anaconda.org/conda-forge/noarch/pyopenssl-25.3.0-pyhd8ed1ab_0.conda
+ sha256: e3a1216bbc4622ac4dfd36c3f8fd3a90d800eebc9147fa3af7eab07d863516b3
+ md5: ddf01a1d87103a152f725c7aeabffa29
+ depends:
+ - cryptography >=45.0.7,<47
+ - python >=3.10
+ - typing-extensions >=4.9
+ - typing_extensions >=4.9
+ license: Apache-2.0
+ license_family: Apache
+ purls:
+ - pkg:pypi/pyopenssl?source=hash-mapping
+ size: 126393
+ timestamp: 1760304658366
+- conda: https://conda.anaconda.org/conda-forge/noarch/pyparsing-3.3.2-pyhcf101f3_0.conda
+ sha256: 417fba4783e528ee732afa82999300859b065dc59927344b4859c64aae7182de
+ md5: 3687cc0b82a8b4c17e1f0eb7e47163d5
+ depends:
+ - python >=3.10
+ - python
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/pyparsing?source=compressed-mapping
+ size: 110893
+ timestamp: 1769003998136
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pyproj-3.7.1-py310h71d0299_1.conda
+ sha256: ec5f371389d7b57cd835144410d04f7af192487b4ff36e032c07b6c032ed383d
+ md5: e54b1eaeb50ad1767680181a8575a8db
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - certifi
+ - libgcc >=13
+ - proj >=9.6.0,<9.7.0a0
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/pyproj?source=hash-mapping
+ size: 536170
+ timestamp: 1742323393877
+- conda: https://conda.anaconda.org/conda-forge/noarch/pysocks-1.7.1-pyha55dd90_7.conda
+ sha256: ba3b032fa52709ce0d9fd388f63d330a026754587a2f461117cac9ab73d8d0d8
+ md5: 461219d1a5bd61342293efa2c0c90eac
+ depends:
+ - __unix
+ - python >=3.9
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/pysocks?source=hash-mapping
+ size: 21085
+ timestamp: 1733217331982
+- pypi: https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl
+ name: pytest
+ version: 9.0.2
+ sha256: 711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b
+ requires_dist:
+ - colorama>=0.4 ; sys_platform == 'win32'
+ - exceptiongroup>=1 ; python_full_version < '3.11'
+ - iniconfig>=1.0.1
+ - packaging>=22
+ - pluggy>=1.5,<2
+ - pygments>=2.7.2
+ - tomli>=1 ; python_full_version < '3.11'
+ - argcomplete ; extra == 'dev'
+ - attrs>=19.2 ; extra == 'dev'
+ - hypothesis>=3.56 ; extra == 'dev'
+ - mock ; extra == 'dev'
+ - requests ; extra == 'dev'
+ - setuptools ; extra == 'dev'
+ - xmlschema ; extra == 'dev'
+ requires_python: '>=3.10'
+- pypi: https://files.pythonhosted.org/packages/ee/49/1377b49de7d0c1ce41292161ea0f721913fa8722c19fb9c1e3aa0367eecb/pytest_cov-7.0.0-py3-none-any.whl
+ name: pytest-cov
+ version: 7.0.0
+ sha256: 3b8e9558b16cc1479da72058bdecf8073661c7f57f7d3c5f22a1c23507f2d861
+ requires_dist:
+ - coverage[toml]>=7.10.6
+ - pluggy>=1.2
+ - pytest>=7
+ - process-tests ; extra == 'testing'
+ - pytest-xdist ; extra == 'testing'
+ - virtualenv ; extra == 'testing'
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl
+ name: pytest-mock
+ version: 3.15.1
+ sha256: 0a25e2eb88fe5168d535041d09a4529a188176ae608a6d249ee65abc0949630d
+ requires_dist:
+ - pytest>=6.2.5
+ - pre-commit ; extra == 'dev'
+ - pytest-asyncio ; extra == 'dev'
+ - tox ; extra == 'dev'
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.10.19-h3c07f61_3_cpython.conda
+ build_number: 3
+ sha256: 2d8b5566d82c3872f057661e056d696f2f77a17ee5a36d9ae6ec43052c4d1c51
+ md5: be48679ccfbc8710dea1d5970600fa04
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - bzip2 >=1.0.8,<2.0a0
+ - ld_impl_linux-64 >=2.36.1
+ - libexpat >=2.7.3,<3.0a0
+ - libffi >=3.4,<4.0a0
+ - libgcc >=14
+ - liblzma >=5.8.2,<6.0a0
+ - libnsl >=2.0.1,<2.1.0a0
+ - libsqlite >=3.51.2,<4.0a0
+ - libuuid >=2.41.3,<3.0a0
+ - libxcrypt >=4.4.36
+ - libzlib >=1.3.1,<2.0a0
+ - ncurses >=6.5,<7.0a0
+ - openssl >=3.5.4,<4.0a0
+ - readline >=8.3,<9.0a0
+ - tk >=8.6.13,<8.7.0a0
+ - tzdata
+ constrains:
+ - python_abi 3.10.* *_cp310
+ license: Python-2.0
+ purls: []
+ size: 25358312
+ timestamp: 1769471983988
+- conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.12.12-hd63d673_2_cpython.conda
+ build_number: 2
+ sha256: 6621befd6570a216ba94bc34ec4618e4f3777de55ad0adc15fc23c28fadd4d1a
+ md5: c4540d3de3fa228d9fa95e31f8e97f89
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - bzip2 >=1.0.8,<2.0a0
+ - ld_impl_linux-64 >=2.36.1
+ - libexpat >=2.7.3,<3.0a0
+ - libffi >=3.5.2,<3.6.0a0
+ - libgcc >=14
+ - liblzma >=5.8.2,<6.0a0
+ - libnsl >=2.0.1,<2.1.0a0
+ - libsqlite >=3.51.2,<4.0a0
+ - libuuid >=2.41.3,<3.0a0
+ - libxcrypt >=4.4.36
+ - libzlib >=1.3.1,<2.0a0
+ - ncurses >=6.5,<7.0a0
+ - openssl >=3.5.4,<4.0a0
+ - readline >=8.3,<9.0a0
+ - tk >=8.6.13,<8.7.0a0
+ - tzdata
+ constrains:
+ - python_abi 3.12.* *_cp312
+ license: Python-2.0
+ purls: []
+ size: 31457785
+ timestamp: 1769472855343
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/python-3.12.12-h18782d2_2_cpython.conda
+ build_number: 2
+ sha256: 765e5d0f92dabc8c468d078a4409490e08181a6f9be6f5d5802a4e3131b9a69c
+ md5: e198b8f74b12292d138eb4eceb004fa3
+ depends:
+ - __osx >=11.0
+ - bzip2 >=1.0.8,<2.0a0
+ - libexpat >=2.7.3,<3.0a0
+ - libffi >=3.5.2,<3.6.0a0
+ - liblzma >=5.8.2,<6.0a0
+ - libsqlite >=3.51.2,<4.0a0
+ - libzlib >=1.3.1,<2.0a0
+ - ncurses >=6.5,<7.0a0
+ - openssl >=3.5.4,<4.0a0
+ - readline >=8.3,<9.0a0
+ - tk >=8.6.13,<8.7.0a0
+ - tzdata
+ constrains:
+ - python_abi 3.12.* *_cp312
+ license: Python-2.0
+ purls: []
+ size: 12953358
+ timestamp: 1769472376612
+- conda: https://conda.anaconda.org/conda-forge/noarch/python-dateutil-2.9.0.post0-pyhe01879c_2.conda
+ sha256: d6a17ece93bbd5139e02d2bd7dbfa80bee1a4261dced63f65f679121686bf664
+ md5: 5b8d21249ff20967101ffa321cab24e8
+ depends:
+ - python >=3.9
+ - six >=1.5
+ - python
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/python-dateutil?source=hash-mapping
+ size: 233310
+ timestamp: 1751104122689
+- conda: https://conda.anaconda.org/conda-forge/noarch/python-tzdata-2025.3-pyhd8ed1ab_0.conda
+ sha256: 467134ef39f0af2dbb57d78cb3e4821f01003488d331a8dd7119334f4f47bfbd
+ md5: 7ead57407430ba33f681738905278d03
+ depends:
+ - python >=3.10
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/tzdata?source=compressed-mapping
+ size: 143542
+ timestamp: 1765719982349
+- conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.10-8_cp310.conda
+ build_number: 8
+ sha256: 7ad76fa396e4bde336872350124c0819032a9e8a0a40590744ff9527b54351c1
+ md5: 05e00f3b21e88bb3d658ac700b2ce58c
+ constrains:
+ - python 3.10.* *_cpython
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 6999
+ timestamp: 1752805924192
+- conda: https://conda.anaconda.org/conda-forge/noarch/python_abi-3.12-8_cp312.conda
+ build_number: 8
+ sha256: 80677180dd3c22deb7426ca89d6203f1c7f1f256f2d5a94dc210f6e758229809
+ md5: c3efd25ac4d74b1584d2f7a57195ddf1
+ constrains:
+ - python 3.12.* *_cpython
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 6958
+ timestamp: 1752805918820
+- conda: https://conda.anaconda.org/conda-forge/noarch/pytz-2025.2-pyhd8ed1ab_0.conda
+ sha256: 8d2a8bf110cc1fc3df6904091dead158ba3e614d8402a83e51ed3a8aa93cdeb0
+ md5: bc8e3267d44011051f2eb14d22fb0960
+ depends:
+ - python >=3.9
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/pytz?source=hash-mapping
+ size: 189015
+ timestamp: 1742920947249
+- conda: https://conda.anaconda.org/conda-forge/noarch/pyu2f-0.1.5-pyhd8ed1ab_1.conda
+ sha256: 991caa5408aea018488a2c94e915c11792b9321b0ef64401f4829ebd0abfb3c0
+ md5: 644bd4ca9f68ef536b902685d773d697
+ depends:
+ - python >=3.9
+ - six
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/pyu2f?source=hash-mapping
+ size: 36786
+ timestamp: 1733738704089
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pywavelets-1.8.0-py310hf462985_0.conda
+ sha256: f23e0b5432c6d338876eca664deeb360949062ce026ddb65bcb1f31643452354
+ md5: 4c441eff2be2e65bd67765c5642051c5
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - numpy >=1.19,<3
+ - numpy >=1.23,<3
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/pywavelets?source=hash-mapping
+ size: 3689433
+ timestamp: 1733419497834
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py310h3406613_0.conda
+ sha256: 9b5c6ff9111ac035f18d5e625bcaa6c076e2e64a6f3c8e3f83f5fe2b03bda78d
+ md5: bc058b3b89fcb525bb4977832aa52014
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - yaml >=0.2.5,<0.3.0a0
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/pyyaml?source=hash-mapping
+ size: 180966
+ timestamp: 1758892005321
+- conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py312h8a5da7c_1.conda
+ sha256: cb142bfd92f6e55749365ddc244294fa7b64db6d08c45b018ff1c658907bfcbf
+ md5: 15878599a87992e44c059731771591cb
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - python >=3.12,<3.13.0a0
+ - python_abi 3.12.* *_cp312
+ - yaml >=0.2.5,<0.3.0a0
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/pyyaml?source=compressed-mapping
+ size: 198293
+ timestamp: 1770223620706
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/pyyaml-6.0.3-py312h04c11ed_1.conda
+ sha256: 737959262d03c9c305618f2d48c7f1691fb996f14ae420bfd05932635c99f873
+ md5: 95a5f0831b5e0b1075bbd80fcffc52ac
+ depends:
+ - __osx >=11.0
+ - python >=3.12,<3.13.0a0
+ - python >=3.12,<3.13.0a0 *_cpython
+ - python_abi 3.12.* *_cp312
+ - yaml >=0.2.5,<0.3.0a0
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/pyyaml?source=compressed-mapping
+ size: 187278
+ timestamp: 1770223990452
+- conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda
+ sha256: 776363493bad83308ba30bcb88c2552632581b143e8ee25b1982c8c743e73abc
+ md5: 353823361b1d27eb3960efb076dfcaf6
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc-ng >=12
+ - libstdcxx-ng >=12
+ license: LicenseRef-Qhull
+ purls: []
+ size: 552937
+ timestamp: 1720813982144
+- conda: https://conda.anaconda.org/conda-forge/linux-64/rasterio-1.4.3-py310hf3df72b_2.conda
+ sha256: 3cdf2c51fa3537e5cfac7277bc9b53c8a8f31d520a5237980253f226207aaebd
+ md5: f49be99a09b5e9286b8f58e12d44ee5d
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - affine
+ - attrs
+ - certifi
+ - click >=4,!=8.2.*
+ - click-plugins
+ - cligj >=0.5
+ - libgcc >=14
+ - libgdal-core <3.11
+ - libgdal-core >=3.10.3,<3.11.0a0
+ - libstdcxx >=14
+ - numpy >=1.21,<3
+ - proj >=9.6.2,<9.7.0a0
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - setuptools >=0.9.8
+ - snuggs >=1.4.1
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/rasterio?source=hash-mapping
+ size: 7942226
+ timestamp: 1758129757954
+- conda: https://conda.anaconda.org/conda-forge/linux-64/rav1e-0.7.1-h8fae777_3.conda
+ sha256: 6e5e704c1c21f820d760e56082b276deaf2b53cf9b751772761c3088a365f6f4
+ md5: 2c42649888aac645608191ffdc80d13a
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ constrains:
+ - __glibc >=2.17
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 5176669
+ timestamp: 1746622023242
+- conda: https://conda.anaconda.org/conda-forge/linux-64/re2-2023.03.02-h8c504da_0.conda
+ sha256: 1727f893a352ca735fb96b09f9edf6fe18c409d65550fd37e8a192919e8c827b
+ md5: 206f8fa808748f6e90599c3368a1114e
+ depends:
+ - libgcc-ng >=12
+ - libstdcxx-ng >=12
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 201211
+ timestamp: 1677698930545
+- conda: https://conda.anaconda.org/conda-forge/linux-64/readline-8.3-h853b02a_0.conda
+ sha256: 12ffde5a6f958e285aa22c191ca01bbd3d6e710aa852e00618fa6ddc59149002
+ md5: d7d95fc8287ea7bf33e0e7116d2b95ec
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - ncurses >=6.5,<7.0a0
+ license: GPL-3.0-only
+ license_family: GPL
+ purls: []
+ size: 345073
+ timestamp: 1765813471974
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/readline-8.3-h46df422_0.conda
+ sha256: a77010528efb4b548ac2a4484eaf7e1c3907f2aec86123ed9c5212ae44502477
+ md5: f8381319127120ce51e081dce4865cf4
+ depends:
+ - __osx >=11.0
+ - ncurses >=6.5,<7.0a0
+ license: GPL-3.0-only
+ license_family: GPL
+ purls: []
+ size: 313930
+ timestamp: 1765813902568
+- pypi: https://files.pythonhosted.org/packages/e1/67/921ec3024056483db83953ae8e48079ad62b92db7880013ca77632921dd0/readme_renderer-44.0-py3-none-any.whl
+ name: readme-renderer
+ version: '44.0'
+ sha256: 2fbca89b81a08526aadf1357a8c2ae889ec05fb03f5da67f9769c9a592166151
+ requires_dist:
+ - nh3>=0.2.14
+ - docutils>=0.21.2
+ - pygments>=2.5.1
+ - cmarkgfm>=0.8.0 ; extra == 'md'
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/noarch/regionmask-0.13.0-pyhd8ed1ab_0.conda
+ sha256: ff5398b5d167690c3f20ab765455cdd6da4590167ae7a71424ad21bf158f193f
+ md5: bcf505f94f8cc1d21475901526b33ab5
+ depends:
+ - geopandas >=0.13
+ - numpy >=1.24
+ - packaging >=23.1
+ - pooch >=1.7
+ - python >=3.10
+ - rasterio >=1.3
+ - shapely >=2.0
+ - xarray >=2023.07
+ constrains:
+ - cartopy >=0.22
+ - matplotlib-base >=3.7
+ - cf_xarray >=0.8
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/regionmask?source=hash-mapping
+ size: 60049
+ timestamp: 1733299306478
+- conda: https://conda.anaconda.org/conda-forge/noarch/requests-2.32.5-pyhcf101f3_1.conda
+ sha256: 7813c38b79ae549504b2c57b3f33394cea4f2ad083f0994d2045c2e24cb538c5
+ md5: c65df89a0b2e321045a9e01d1337b182
+ depends:
+ - python >=3.10
+ - certifi >=2017.4.17
+ - charset-normalizer >=2,<4
+ - idna >=2.5,<4
+ - urllib3 >=1.21.1,<3
+ - python
+ constrains:
+ - chardet >=3.0.2,<6
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/requests?source=compressed-mapping
+ size: 63602
+ timestamp: 1766926974520
+- conda: https://conda.anaconda.org/conda-forge/noarch/requests-oauthlib-2.0.0-pyhd8ed1ab_1.conda
+ sha256: 75ef0072ae6691f5ca9709fe6a2570b98177b49d0231a6749ac4e610da934cab
+ md5: a283b764d8b155f81e904675ef5e1f4b
+ depends:
+ - oauthlib >=3.0.0
+ - python >=3.9
+ - requests >=2.0.0
+ license: ISC
+ purls:
+ - pkg:pypi/requests-oauthlib?source=hash-mapping
+ size: 25875
+ timestamp: 1733772348802
+- pypi: https://files.pythonhosted.org/packages/3f/51/d4db610ef29373b879047326cbf6fa98b6c1969d6f6dc423279de2b1be2c/requests_toolbelt-1.0.0-py2.py3-none-any.whl
+ name: requests-toolbelt
+ version: 1.0.0
+ sha256: cccfdd665f0a24fcf4726e690f65639d272bb0637b9b92dfd91a5568ccf6bd06
+ requires_dist:
+ - requests>=2.0.1,<3.0.0
+ requires_python: '>=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*'
+- pypi: https://files.pythonhosted.org/packages/bd/60/50fbb6ffb35f733654466f1a90d162bcbea358adc3b0871339254fbc37b2/requirements_parser-0.13.0-py3-none-any.whl
+ name: requirements-parser
+ version: 0.13.0
+ sha256: 2b3173faecf19ec5501971b7222d38f04cb45bb9d87d0ad629ca71e2e62ded14
+ requires_dist:
+ - packaging>=23.2
+ requires_python: '>=3.8,<4.0'
+- pypi: https://files.pythonhosted.org/packages/ff/9a/9afaade874b2fa6c752c36f1548f718b5b83af81ed9b76628329dab81c1b/rfc3986-2.0.0-py2.py3-none-any.whl
+ name: rfc3986
+ version: 2.0.0
+ sha256: 50b1502b60e289cb37883f3dfd34532b8873c7de9f49bb546641ce9cbd256ebd
+ requires_dist:
+ - idna ; extra == 'idna2008'
+ requires_python: '>=3.7'
+- pypi: https://files.pythonhosted.org/packages/87/2a/a1810c8627b9ec8c57ec5ec325d306701ae7be50235e8fd81266e002a3cc/rich-14.3.1-py3-none-any.whl
+ name: rich
+ version: 14.3.1
+ sha256: da750b1aebbff0b372557426fb3f35ba56de8ef954b3190315eb64076d6fb54e
+ requires_dist:
+ - ipywidgets>=7.5.1,<9 ; extra == 'jupyter'
+ - markdown-it-py>=2.2.0
+ - pygments>=2.13.0,<3.0.0
+ requires_python: '>=3.8.0'
+- conda: https://conda.anaconda.org/conda-forge/noarch/rsa-4.9.1-pyhd8ed1ab_0.conda
+ sha256: e32e94e7693d4bc9305b36b8a4ef61034e0428f58850ebee4675978e3c2e5acf
+ md5: 58958bb50f986ac0c46f73b6e290d5fe
+ depends:
+ - pyasn1 >=0.1.3
+ - python >=3.9
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/rsa?source=hash-mapping
+ size: 31709
+ timestamp: 1744825527634
+- conda: https://conda.anaconda.org/conda-forge/noarch/rtree-1.4.1-pyh11ca60a_0.conda
+ sha256: 461ddd1d84c180bb682560b50d538e9e264ee5cc78cab5eb2a0f21cc24815bed
+ md5: 73f0eccab422ca6a96d904a805d68fa3
+ depends:
+ - libspatialindex
+ - python >=3.9
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/rtree?source=hash-mapping
+ size: 41185
+ timestamp: 1755128422366
+- pypi: https://files.pythonhosted.org/packages/ca/71/37daa46f89475f8582b7762ecd2722492df26421714a33e72ccc9a84d7a5/ruff-0.14.14-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ name: ruff
+ version: 0.14.14
+ sha256: bb8481604b7a9e75eff53772496201690ce2687067e038b3cc31aaf16aa0b974
+ requires_python: '>=3.7'
+- pypi: https://files.pythonhosted.org/packages/57/e1/64c264db50b68de8a438b60ceeb921b2f22da3ebb7ad6255150225d0beac/s3fs-2026.2.0-py3-none-any.whl
+ name: s3fs
+ version: 2026.2.0
+ sha256: 65198835b86b1d5771112b0085d1da52a6ede36508b1aaa6cae2aedc765dfe10
+ requires_dist:
+ - aiobotocore>=2.19.0,<4.0.0
+ - fsspec==2026.2.0
+ - aiohttp!=4.0.0a0,!=4.0.0a1
+ requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/scikit-image-0.25.2-py310h0158d43_2.conda
+ sha256: 34593b03ba5de16ce231a6485caa295fbdca251a8cb3585ec5db1ffe9df6b063
+ md5: e8e3404c2d4135193013fbbe9bba60a5
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - imageio >=2.33,!=2.35.0
+ - lazy-loader >=0.4
+ - libgcc >=14
+ - libstdcxx >=14
+ - networkx >=3.0
+ - numpy >=1.21,<3
+ - numpy >=1.24
+ - packaging >=21
+ - pillow >=10.1
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - pywavelets >=1.6
+ - scipy >=1.11.4
+ - tifffile >=2022.8.12
+ constrains:
+ - astropy-base >=6.0
+ - pywavelets >=1.6
+ - scikit-learn >=1.2
+ - pyamg >=5.2
+ - matplotlib-base >=3.7
+ - numpy >=1.24
+ - dask-core >=2023.2.0,!=2024.8.0
+ - pooch >=1.6.0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/scikit-image?source=hash-mapping
+ size: 10561680
+ timestamp: 1757197302869
+- conda: https://conda.anaconda.org/conda-forge/linux-64/scikit-learn-1.7.2-py310h228f341_0.conda
+ sha256: 8af2a49b75b9697653a0bf55717c8153732c4fc53f58d4aa0ed692ae348c49f6
+ md5: 0f3e3324506bd3e67934eda9895f37a7
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - _openmp_mutex >=4.5
+ - joblib >=1.2.0
+ - libgcc >=14
+ - libstdcxx >=14
+ - numpy >=1.21,<3
+ - numpy >=1.22.0
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - scipy >=1.8.0
+ - threadpoolctl >=3.1.0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/scikit-learn?source=hash-mapping
+ size: 8415134
+ timestamp: 1757406407327
+- conda: https://conda.anaconda.org/conda-forge/linux-64/scipy-1.15.2-py310h1d65ade_0.conda
+ sha256: 4cb98641f870666d365594013701d5691205a0fe81ac3ba7778a23b1cc2caa8e
+ md5: 8c29cd33b64b2eb78597fa28b5595c8d
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libblas >=3.9.0,<4.0a0
+ - libcblas >=3.9.0,<4.0a0
+ - libgcc >=13
+ - libgfortran
+ - libgfortran5 >=13.3.0
+ - liblapack >=3.9.0,<4.0a0
+ - libstdcxx >=13
+ - numpy <2.5
+ - numpy >=1.19,<3
+ - numpy >=1.23.5
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/scipy?source=hash-mapping
+ size: 16417101
+ timestamp: 1739791865060
+- pypi: https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl
+ name: secretstorage
+ version: 3.5.0
+ sha256: 0ce65888c0725fcb2c5bc0fdb8e5438eece02c523557ea40ce0703c266248137
+ requires_dist:
+ - cryptography>=2.0
+ - jeepney>=0.6
+ requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/noarch/sentry-sdk-2.51.0-pyhd8ed1ab_0.conda
+ sha256: 2200f6f56f8ba8e5afe932512887055e459239409003fafe7c96eda191e05ec3
+ md5: 601d7be62023e87a7408d046066fca0b
+ depends:
+ - certifi
+ - python >=3.10
+ - urllib3 >=1.25.7
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/sentry-sdk?source=compressed-mapping
+ size: 273962
+ timestamp: 1769647505200
+- conda: https://conda.anaconda.org/conda-forge/linux-64/setproctitle-1.3.7-py310h139afa4_0.conda
+ sha256: 80fc00acbfeaa052f88597c096ebb01c862a3b9e09bffee3137c6804e59698ba
+ md5: 9b46695dda92c124a67b396e3cfce0f0
+ depends:
+ - python
+ - libgcc >=14
+ - __glibc >=2.17,<3.0.a0
+ - python_abi 3.10.* *_cp310
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/setproctitle?source=hash-mapping
+ size: 23157
+ timestamp: 1766684434552
+- conda: https://conda.anaconda.org/conda-forge/noarch/setuptools-80.10.2-pyh332efcf_0.conda
+ sha256: f5fcb7854d2b7639a5b1aca41dd0f2d5a69a60bbc313e7f192e2dc385ca52f86
+ md5: 7b446fcbb6779ee479debb4fd7453e6c
+ depends:
+ - python >=3.10
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/setuptools?source=compressed-mapping
+ size: 678888
+ timestamp: 1769601206751
+- conda: https://conda.anaconda.org/conda-forge/linux-64/shapely-2.1.2-py310h777b3ac_0.conda
+ sha256: 4baffc6640c1b7f590df0fb62089947408fa98fdcac9f3cea420bd33afa36137
+ md5: 486ffd32d6cc932f43df4066efd3c973
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - geos >=3.14.0,<3.14.1.0a0
+ - libgcc >=14
+ - numpy >=1.21,<3
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/shapely?source=hash-mapping
+ size: 542303
+ timestamp: 1758735396162
+- conda: https://conda.anaconda.org/conda-forge/noarch/six-1.17.0-pyhe01879c_1.conda
+ sha256: 458227f759d5e3fcec5d9b7acce54e10c9e1f4f4b7ec978f3bfd54ce4ee9853d
+ md5: 3339e3b65d58accf4ca4fb8748ab16b3
+ depends:
+ - python >=3.9
+ - python
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/six?source=hash-mapping
+ size: 18455
+ timestamp: 1753199211006
+- conda: https://conda.anaconda.org/conda-forge/noarch/smmap-5.0.2-pyhd8ed1ab_0.conda
+ sha256: eb92d0ad94b65af16c73071cc00cc0e10f2532be807beb52758aab2b06eb21e2
+ md5: 87f47a78808baf2fa1ea9c315a1e48f1
+ depends:
+ - python >=3.9
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/smmap?source=hash-mapping
+ size: 26051
+ timestamp: 1739781801801
+- conda: https://conda.anaconda.org/conda-forge/linux-64/snappy-1.2.2-h03e3b7b_1.conda
+ sha256: 48f3f6a76c34b2cfe80de9ce7f2283ecb55d5ed47367ba91e8bb8104e12b8f11
+ md5: 98b6c9dc80eb87b2519b97bcf7e578dd
+ depends:
+ - libgcc >=14
+ - __glibc >=2.17,<3.0.a0
+ - libstdcxx >=14
+ - libgcc >=14
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 45829
+ timestamp: 1762948049098
+- conda: https://conda.anaconda.org/conda-forge/noarch/snuggs-1.4.7-pyhd8ed1ab_2.conda
+ sha256: 61f9373709e7d9009e3a062b135dbe44b16e684a4fcfe2dd624143bc0f80d402
+ md5: 9aa358575bbd4be126eaa5e0039f835c
+ depends:
+ - numpy
+ - pyparsing >=2.1.6
+ - python >=3.9
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/snuggs?source=hash-mapping
+ size: 11313
+ timestamp: 1733818738919
+- conda: https://conda.anaconda.org/conda-forge/noarch/sortedcontainers-2.4.0-pyhd8ed1ab_1.conda
+ sha256: d1e3e06b5cf26093047e63c8cc77b70d970411c5cbc0cb1fad461a8a8df599f7
+ md5: 0401a17ae845fa72c7210e206ec5647d
+ depends:
+ - python >=3.9
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/sortedcontainers?source=hash-mapping
+ size: 28657
+ timestamp: 1738440459037
+- conda: https://conda.anaconda.org/conda-forge/linux-64/sqlite-3.51.2-hbc0de68_0.conda
+ sha256: 65436099fd33e4471348d614c1de9235fdd4e5b7d86a5a12472922e6b6628951
+ md5: a6adeaa8efb007e2e1ab3e45768ea987
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libsqlite 3.51.2 h0c1763c_0
+ - libzlib >=1.3.1,<2.0a0
+ - ncurses >=6.5,<7.0a0
+ - readline >=8.3,<9.0a0
+ license: blessing
+ purls: []
+ size: 183835
+ timestamp: 1768147980363
+- conda: https://conda.anaconda.org/conda-forge/linux-64/svt-av1-4.0.0-hecca717_0.conda
+ sha256: e5e036728ef71606569232cc94a0480722e14ed69da3dd1e363f3d5191d83c01
+ md5: 9a6117aee038999ffefe6082ff1e9a81
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libstdcxx >=14
+ license: BSD-2-Clause
+ license_family: BSD
+ purls: []
+ size: 2620937
+ timestamp: 1769280649780
+- conda: https://conda.anaconda.org/conda-forge/noarch/tblib-3.2.2-pyhcf101f3_0.conda
+ sha256: 6b549360f687ee4d11bf85a6d6a276a30f9333df1857adb0fe785f0f8e9bcd60
+ md5: f88bb644823094f436792f80fba3207e
+ depends:
+ - python >=3.10
+ - python
+ license: BSD-2-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/tblib?source=hash-mapping
+ size: 19397
+ timestamp: 1762956379123
+- pypi: https://files.pythonhosted.org/packages/5d/12/4f70e8e2ba0dbe72ea978429d8530b0333f0ed2140cc571a48802878ef99/tensorboard-2.19.0-py3-none-any.whl
+ name: tensorboard
+ version: 2.19.0
+ sha256: 5e71b98663a641a7ce8a6e70b0be8e1a4c0c45d48760b076383ac4755c35b9a0
+ requires_dist:
+ - absl-py>=0.4
+ - grpcio>=1.48.2
+ - markdown>=2.6.8
+ - numpy>=1.12.0
+ - packaging
+ - protobuf>=3.19.6,!=4.24.0
+ - setuptools>=41.0.0
+ - six>1.9
+ - tensorboard-data-server>=0.7.0,<0.8.0
+ - werkzeug>=1.0.1
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/7a/13/e503968fefabd4c6b2650af21e110aa8466fe21432cd7c43a84577a89438/tensorboard_data_server-0.7.2-py3-none-any.whl
+ name: tensorboard-data-server
+ version: 0.7.2
+ sha256: 7e0610d205889588983836ec05dc098e80f97b7e7bbff7e994ebb78f578d0ddb
+ requires_python: '>=3.7'
+- pypi: https://files.pythonhosted.org/packages/77/60/51d921b17f7b1547db32018d6a933627d4e2a762d2fc2ca6c0032cc8b062/tensorflow-2.19.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ name: tensorflow
+ version: 2.19.1
+ sha256: 11377d4cecdc665515e370524b650a8ceb80b461f40f5e63993ed27d0ef931da
+ requires_dist:
+ - absl-py>=1.0.0
+ - astunparse>=1.6.0
+ - flatbuffers>=24.3.25
+ - gast>=0.2.1,!=0.5.0,!=0.5.1,!=0.5.2
+ - google-pasta>=0.1.1
+ - libclang>=13.0.0
+ - opt-einsum>=2.3.2
+ - packaging
+ - protobuf>=3.20.3,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<6.0.0.dev0
+ - requests>=2.21.0,<3
+ - setuptools
+ - six>=1.12.0
+ - termcolor>=1.1.0
+ - typing-extensions>=3.6.6
+ - wrapt>=1.11.0
+ - grpcio>=1.24.3,<2.0
+ - tensorboard~=2.19.0
+ - keras>=3.5.0
+ - numpy>=1.26.0,<2.2.0
+ - h5py>=3.11.0
+ - ml-dtypes>=0.5.1,<1.0.0
+ - tensorflow-intel==2.19.1 ; sys_platform == 'win32'
+ - tensorflow-io-gcs-filesystem>=0.23.1 ; python_full_version < '3.12'
+ - nvidia-cublas-cu12==12.5.3.2 ; extra == 'and-cuda'
+ - nvidia-cuda-cupti-cu12==12.5.82 ; extra == 'and-cuda'
+ - nvidia-cuda-nvcc-cu12==12.5.82 ; extra == 'and-cuda'
+ - nvidia-cuda-nvrtc-cu12==12.5.82 ; extra == 'and-cuda'
+ - nvidia-cuda-runtime-cu12==12.5.82 ; extra == 'and-cuda'
+ - nvidia-cudnn-cu12==9.3.0.75 ; extra == 'and-cuda'
+ - nvidia-cufft-cu12==11.2.3.61 ; extra == 'and-cuda'
+ - nvidia-curand-cu12==10.3.6.82 ; extra == 'and-cuda'
+ - nvidia-cusolver-cu12==11.6.3.83 ; extra == 'and-cuda'
+ - nvidia-cusparse-cu12==12.5.1.3 ; extra == 'and-cuda'
+ - nvidia-nccl-cu12==2.23.4 ; extra == 'and-cuda'
+ - nvidia-nvjitlink-cu12==12.5.82 ; extra == 'and-cuda'
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/f3/48/47b7d25572961a48b1de3729b7a11e835b888e41e0203cca82df95d23b91/tensorflow_io_gcs_filesystem-0.37.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+ name: tensorflow-io-gcs-filesystem
+ version: 0.37.1
+ sha256: 9679b36e3a80921876f31685ab6f7270f3411a4cc51bc2847e80d0e4b5291e27
+ requires_dist:
+ - tensorflow>=2.16.0,<2.17.0 ; extra == 'tensorflow'
+ - tensorflow-aarch64>=2.16.0,<2.17.0 ; extra == 'tensorflow-aarch64'
+ - tensorflow-cpu>=2.16.0,<2.17.0 ; extra == 'tensorflow-cpu'
+ - tensorflow-gpu>=2.16.0,<2.17.0 ; extra == 'tensorflow-gpu'
+ - tensorflow-rocm>=2.16.0,<2.17.0 ; extra == 'tensorflow-rocm'
+ requires_python: '>=3.7,<3.13'
+- pypi: https://files.pythonhosted.org/packages/33/d1/8bb87d21e9aeb323cc03034f5eaf2c8f69841e40e4853c2627edf8111ed3/termcolor-3.3.0-py3-none-any.whl
+ name: termcolor
+ version: 3.3.0
+ sha256: cf642efadaf0a8ebbbf4bc7a31cec2f9b5f21a9f726f4ccbb08192c9c26f43a5
+ requires_dist:
+ - pytest ; extra == 'tests'
+ - pytest-cov ; extra == 'tests'
+ requires_python: '>=3.10'
+- conda: https://conda.anaconda.org/conda-forge/noarch/threadpoolctl-3.6.0-pyhecae5ae_0.conda
+ sha256: 6016672e0e72c4cf23c0cf7b1986283bd86a9c17e8d319212d78d8e9ae42fdfd
+ md5: 9d64911b31d57ca443e9f1e36b04385f
+ depends:
+ - python >=3.9
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/threadpoolctl?source=hash-mapping
+ size: 23869
+ timestamp: 1741878358548
+- conda: https://conda.anaconda.org/conda-forge/noarch/tifffile-2025.5.10-pyhd8ed1ab_0.conda
+ sha256: 3ea3854eb8a41bbb128598a5d5bc9aed52446d20d2f1bd6e997c2387074202e4
+ md5: 1fdb801f28bf4987294c49aaa314bf5e
+ depends:
+ - imagecodecs >=2024.12.30
+ - numpy >=1.19.2
+ - python >=3.10
+ constrains:
+ - matplotlib-base >=3.3
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/tifffile?source=hash-mapping
+ size: 179592
+ timestamp: 1746986641678
+- conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_h366c992_103.conda
+ sha256: cafeec44494f842ffeca27e9c8b0c27ed714f93ac77ddadc6aaf726b5554ebac
+ md5: cffd3bdd58090148f4cfcd831f4b26ab
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libzlib >=1.3.1,<2.0a0
+ constrains:
+ - xorg-libx11 >=1.8.12,<2.0a0
+ license: TCL
+ license_family: BSD
+ purls: []
+ size: 3301196
+ timestamp: 1769460227866
+- conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_h4845f30_101.conda
+ sha256: e0569c9caa68bf476bead1bed3d79650bb080b532c64a4af7d8ca286c08dea4e
+ md5: d453b98d9c83e71da0741bb0ff4d76bc
+ depends:
+ - libgcc-ng >=12
+ - libzlib >=1.2.13,<2.0.0a0
+ license: TCL
+ license_family: BSD
+ purls: []
+ size: 3318875
+ timestamp: 1699202167581
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/tk-8.6.13-h010d191_3.conda
+ sha256: 799cab4b6cde62f91f750149995d149bc9db525ec12595e8a1d91b9317f038b3
+ md5: a9d86bc62f39b94c4661716624eb21b0
+ depends:
+ - __osx >=11.0
+ - libzlib >=1.3.1,<2.0a0
+ license: TCL
+ license_family: BSD
+ purls: []
+ size: 3127137
+ timestamp: 1769460817696
+- pypi: https://files.pythonhosted.org/packages/23/d1/136eb2cb77520a31e1f64cbae9d33ec6df0d78bdf4160398e86eec8a8754/tomli-2.4.0-py3-none-any.whl
+ name: tomli
+ version: 2.4.0
+ sha256: 1f776e7d669ebceb01dee46484485f43a4048746235e683bcdffacdf1fb4785a
+ requires_python: '>=3.8'
+- conda: https://conda.anaconda.org/conda-forge/noarch/toolz-1.1.0-pyhd8ed1ab_1.conda
+ sha256: 4e379e1c18befb134247f56021fdf18e112fb35e64dd1691858b0a0f3bea9a45
+ md5: c07a6153f8306e45794774cf9b13bd32
+ depends:
+ - python >=3.10
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/toolz?source=hash-mapping
+ size: 53978
+ timestamp: 1760707830681
+- conda: https://conda.anaconda.org/conda-forge/linux-64/tornado-6.5.3-py310h7c4b9e2_0.conda
+ sha256: c27c28d19f8ba8ef6efd35dc47951c985db8a828db38444e1fad3f93f8cedb8d
+ md5: 30b9d5c1bc99ffbc45a63ab8d1725b93
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: Apache-2.0
+ license_family: Apache
+ purls:
+ - pkg:pypi/tornado?source=hash-mapping
+ size: 663313
+ timestamp: 1765458854459
+- pypi: https://files.pythonhosted.org/packages/3a/7a/882d99539b19b1490cac5d77c67338d126e4122c8276bf640e411650c830/twine-6.2.0-py3-none-any.whl
+ name: twine
+ version: 6.2.0
+ sha256: 418ebf08ccda9a8caaebe414433b0ba5e25eb5e4a927667122fbe8f829f985d8
+ requires_dist:
+ - readme-renderer>=35.0
+ - requests>=2.20
+ - requests-toolbelt>=0.8.0,!=0.9.0
+ - urllib3>=1.26.0
+ - importlib-metadata>=3.6 ; python_full_version < '3.10'
+ - keyring>=21.2.0 ; platform_machine != 'ppc64le' and platform_machine != 's390x'
+ - rfc3986>=1.4.0
+ - rich>=12.0.0
+ - packaging>=24.0
+ - id
+ - keyring>=21.2.0 ; extra == 'keyring'
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/e7/c1/56ef16bf5dcd255155cc736d276efa6ae0a5c26fd685e28f0412a4013c01/types_pytz-2025.2.0.20251108-py3-none-any.whl
+ name: types-pytz
+ version: 2025.2.0.20251108
+ sha256: 0f1c9792cab4eb0e46c52f8845c8f77cf1e313cb3d68bf826aa867fe4717d91c
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/bd/e0/1eed384f02555dde685fff1a1ac805c1c7dcb6dd019c916fe659b1c1f9ec/types_pyyaml-6.0.12.20250915-py3-none-any.whl
+ name: types-pyyaml
+ version: 6.0.12.20250915
+ sha256: e7d4d9e064e89a3b3cae120b4990cd370874d2bf12fa5f46c97018dd5d3c9ab6
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/1c/12/709ea261f2bf91ef0a26a9eed20f2623227a8ed85610c1e54c5805692ecb/types_requests-2.32.4.20260107-py3-none-any.whl
+ name: types-requests
+ version: 2.32.4.20260107
+ sha256: b703fe72f8ce5b31ef031264fe9395cac8f46a04661a79f7ed31a80fb308730d
+ requires_dist:
+ - urllib3>=2
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/3f/13/3ff0781445d7c12730befce0fddbbc7a76e56eb0e7029446f2853238360a/types_tqdm-4.67.0.20250809-py3-none-any.whl
+ name: types-tqdm
+ version: 4.67.0.20250809
+ sha256: 1a73053b31fcabf3c1f3e2a9d5ecdba0f301bde47a418cd0e0bdf774827c5c57
+ requires_dist:
+ - types-requests
+ requires_python: '>=3.9'
+- conda: https://conda.anaconda.org/conda-forge/noarch/typing-3.10.0.0-pyhd8ed1ab_2.conda
+ sha256: 92b084dfd77571be23ef84ad695bbea169e844821484b6d47d99f04ea4de32e8
+ md5: 28abeb80aea7eb4914f3a7543a47e248
+ depends:
+ - python >=3.9
+ license: PSF-2.0
+ license_family: PSF
+ purls: []
+ size: 9502
+ timestamp: 1733927569850
+- conda: https://conda.anaconda.org/conda-forge/noarch/typing-extensions-4.15.0-h396c80c_0.conda
+ sha256: 7c2df5721c742c2a47b2c8f960e718c930031663ac1174da67c1ed5999f7938c
+ md5: edd329d7d3a4ab45dcf905899a7a6115
+ depends:
+ - typing_extensions ==4.15.0 pyhcf101f3_0
+ license: PSF-2.0
+ license_family: PSF
+ purls: []
+ size: 91383
+ timestamp: 1756220668932
+- conda: https://conda.anaconda.org/conda-forge/noarch/typing-inspection-0.4.2-pyhd8ed1ab_1.conda
+ sha256: 70db27de58a97aeb7ba7448366c9853f91b21137492e0b4430251a1870aa8ff4
+ md5: a0a4a3035667fc34f29bfbd5c190baa6
+ depends:
+ - python >=3.10
+ - typing_extensions >=4.12.0
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/typing-inspection?source=hash-mapping
+ size: 18923
+ timestamp: 1764158430324
+- conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.15.0-pyhcf101f3_0.conda
+ sha256: 032271135bca55aeb156cee361c81350c6f3fb203f57d024d7e5a1fc9ef18731
+ md5: 0caa1af407ecff61170c9437a808404d
+ depends:
+ - python >=3.10
+ - python
+ license: PSF-2.0
+ license_family: PSF
+ purls:
+ - pkg:pypi/typing-extensions?source=hash-mapping
+ size: 51692
+ timestamp: 1756220668932
+- conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda
+ sha256: 1d30098909076af33a35017eed6f2953af1c769e273a0626a04722ac4acaba3c
+ md5: ad659d0a2b3e47e38d829aa8cad2d610
+ license: LicenseRef-Public-Domain
+ purls: []
+ size: 119135
+ timestamp: 1767016325805
+- conda: https://conda.anaconda.org/conda-forge/linux-64/unicodedata2-17.0.0-py310h7c4b9e2_1.conda
+ sha256: cffe509e0294586fbcee9cbb762d6144636c5d4a19defffda9f9c726a84b55e7
+ md5: b1ccdb989be682ab0dd430c1c15d5012
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: Apache-2.0
+ license_family: Apache
+ purls:
+ - pkg:pypi/unicodedata2?source=hash-mapping
+ size: 409991
+ timestamp: 1763054811367
+- conda: https://conda.anaconda.org/conda-forge/linux-64/uriparser-0.9.8-hac33072_0.conda
+ sha256: 2aad2aeff7c69a2d7eecd7b662eef756b27d6a6b96f3e2c2a7071340ce14543e
+ md5: d71d3a66528853c0a1ac2c02d79a0284
+ depends:
+ - libgcc-ng >=12
+ - libstdcxx-ng >=12
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 48270
+ timestamp: 1715010035325
+- conda: https://conda.anaconda.org/conda-forge/noarch/urllib3-2.5.0-pyhd8ed1ab_0.conda
+ sha256: 4fb9789154bd666ca74e428d973df81087a697dbb987775bc3198d2215f240f8
+ md5: 436c165519e140cb08d246a4472a9d6a
+ depends:
+ - brotli-python >=1.0.9
+ - h2 >=4,<5
+ - pysocks >=1.5.6,<2.0,!=1.5.7
+ - python >=3.9
+ - zstandard >=0.18.0
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/urllib3?source=hash-mapping
+ size: 101735
+ timestamp: 1750271478254
+- pypi: https://files.pythonhosted.org/packages/6a/2a/dc2228b2888f51192c7dc766106cd475f1b768c10caaf9727659726f7391/virtualenv-20.36.1-py3-none-any.whl
+ name: virtualenv
+ version: 20.36.1
+ sha256: 575a8d6b124ef88f6f51d56d656132389f961062a9177016a50e4f507bbcc19f
+ requires_dist:
+ - distlib>=0.3.7,<1
+ - filelock>=3.16.1,<4 ; python_full_version < '3.10'
+ - filelock>=3.20.1,<4 ; python_full_version >= '3.10'
+ - importlib-metadata>=6.6 ; python_full_version < '3.8'
+ - platformdirs>=3.9.1,<5
+ - typing-extensions>=4.13.2 ; python_full_version < '3.11'
+ - furo>=2023.7.26 ; extra == 'docs'
+ - proselint>=0.13 ; extra == 'docs'
+ - sphinx>=7.1.2,!=7.3 ; extra == 'docs'
+ - sphinx-argparse>=0.4 ; extra == 'docs'
+ - sphinxcontrib-towncrier>=0.2.1a0 ; extra == 'docs'
+ - towncrier>=23.6 ; extra == 'docs'
+ - covdefaults>=2.3 ; extra == 'test'
+ - coverage-enable-subprocess>=1 ; extra == 'test'
+ - coverage>=7.2.7 ; extra == 'test'
+ - flaky>=3.7 ; extra == 'test'
+ - packaging>=23.1 ; extra == 'test'
+ - pytest-env>=0.8.2 ; extra == 'test'
+ - pytest-freezer>=0.4.8 ; (python_full_version >= '3.13' and platform_python_implementation == 'CPython' and sys_platform == 'win32' and extra == 'test') or (platform_python_implementation == 'GraalVM' and extra == 'test') or (platform_python_implementation == 'PyPy' and extra == 'test')
+ - pytest-mock>=3.11.1 ; extra == 'test'
+ - pytest-randomly>=3.12 ; extra == 'test'
+ - pytest-timeout>=2.1 ; extra == 'test'
+ - pytest>=7.4 ; extra == 'test'
+ - setuptools>=68 ; extra == 'test'
+ - time-machine>=2.10 ; platform_python_implementation == 'CPython' and extra == 'test'
+ requires_python: '>=3.8'
+- conda: https://conda.anaconda.org/conda-forge/linux-64/wandb-0.24.1-py310hdfeec95_0.conda
+ sha256: aba4b44c8e58dc9667ae6694ffbce513031d63c8dfefbf7a34d85221c66c8248
+ md5: 97429b8bffe3ab847f5bf5e370901f9c
+ depends:
+ - __glibc >=2.17
+ - __glibc >=2.17,<3.0.a0
+ - appdirs >=1.4.3
+ - click >=8.0.1
+ - docker-pycreds >=0.4.0
+ - eval_type_backport
+ - gitpython >=1.0.0,!=3.1.29
+ - libgcc >=14
+ - packaging
+ - platformdirs
+ - protobuf >=3.19.0,!=4.21.0,!=5.28.0,<7
+ - psutil >=5.0.0
+ - pydantic <3
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ - pyyaml
+ - requests >=2.0.0,<3
+ - sentry-sdk >=2.0.0
+ - setproctitle
+ - setuptools
+ - six
+ - typing >=3.6.4
+ - typing_extensions >=4.8,<5
+ constrains:
+ - __glibc >=2.17
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/wandb?source=hash-mapping
+ size: 20416722
+ timestamp: 1769821608638
+- pypi: https://files.pythonhosted.org/packages/ad/e4/8d97cca767bcc1be76d16fb76951608305561c6e056811587f36cb1316a8/werkzeug-3.1.5-py3-none-any.whl
+ name: werkzeug
+ version: 3.1.5
+ sha256: 5111e36e91086ece91f93268bb39b4a35c1e6f1feac762c9c822ded0a4e322dc
+ requires_dist:
+ - markupsafe>=2.1.1
+ - watchdog>=2.3 ; extra == 'watchdog'
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/87/22/b76d483683216dde3d67cba61fb2444be8d5be289bf628c13fc0fd90e5f9/wheel-0.46.3-py3-none-any.whl
+ name: wheel
+ version: 0.46.3
+ sha256: 4b399d56c9d9338230118d705d9737a2a468ccca63d5e813e2a4fc7815d8bc4d
+ requires_dist:
+ - packaging>=24.0
+ - pytest>=6.0.0 ; extra == 'test'
+ - setuptools>=77 ; extra == 'test'
+ requires_python: '>=3.9'
+- pypi: https://files.pythonhosted.org/packages/c6/93/5cf92edd99617095592af919cb81d4bff61c5dbbb70d3c92099425a8ec34/wrapt-2.0.1-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl
+ name: wrapt
+ version: 2.0.1
+ sha256: 36982b26f190f4d737f04a492a68accbfc6fa042c3f42326fdfbb6c5b7a20a31
+ requires_dist:
+ - pytest ; extra == 'dev'
+ - setuptools ; extra == 'dev'
+ requires_python: '>=3.8'
+- conda: https://conda.anaconda.org/conda-forge/noarch/xarray-2024.3.0-pyhd8ed1ab_0.conda
+ sha256: 74e4cea340517ce7c51c36efc1d544d3a98fcdb62a429b6b1a59a1917b412c10
+ md5: 772d7ee42b65d0840130eabd5bd3fc17
+ depends:
+ - numpy >=1.23,<2.0a0
+ - packaging >=22
+ - pandas >=1.5
+ - python >=3.9
+ constrains:
+ - bottleneck >=1.3
+ - sparse >=0.13
+ - nc-time-axis >=1.4
+ - scipy >=1.8
+ - zarr >=2.12
+ - flox >=0.5
+ - netcdf4 >=1.6.0
+ - cartopy >=0.20
+ - h5netcdf >=1.0
+ - dask-core >=2022.7
+ - cftime >=1.6
+ - numba >=0.55
+ - hdf5 >=1.12
+ - iris >=3.2
+ - toolz >=0.12
+ - h5py >=3.6
+ - distributed >=2022.7
+ - matplotlib-base >=3.5
+ - seaborn-base >=0.11
+ - pint >=0.19
+ license: Apache-2.0
+ license_family: APACHE
+ purls:
+ - pkg:pypi/xarray?source=hash-mapping
+ size: 765419
+ timestamp: 1711742257463
+- conda: https://conda.anaconda.org/conda-forge/noarch/xbatcher-0.4.0-pyhd8ed1ab_1.conda
+ sha256: 76cb94fd46cb3c719ca4937004de4f42e7754b383972a7428ecd70485f915e37
+ md5: 2db2a88bf17f09a269c140adbf686392
+ depends:
+ - dask
+ - numpy
+ - python >=3.10
+ - xarray
+ license: Apache-2.0
+ license_family: Apache
+ purls:
+ - pkg:pypi/xbatcher?source=hash-mapping
+ size: 25704
+ timestamp: 1733112828059
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xerces-c-3.2.5-h988505b_2.conda
+ sha256: 339ab0ff05170a295e59133cd0fa9a9c4ba32b6941c8a2a73484cc13f81e248a
+ md5: 9dda9667feba914e0e80b95b82f7402b
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - icu >=75.1,<76.0a0
+ - libgcc >=13
+ - libnsl >=2.0.1,<2.1.0a0
+ - libstdcxx >=13
+ license: Apache-2.0
+ license_family: Apache
+ purls: []
+ size: 1648243
+ timestamp: 1727733890754
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxau-1.0.12-hb03c661_1.conda
+ sha256: 6bc6ab7a90a5d8ac94c7e300cc10beb0500eeba4b99822768ca2f2ef356f731b
+ md5: b2895afaf55bf96a8c8282a2e47a5de0
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 15321
+ timestamp: 1762976464266
+- conda: https://conda.anaconda.org/conda-forge/linux-64/xorg-libxdmcp-1.1.5-hb03c661_1.conda
+ sha256: 25d255fb2eef929d21ff660a0c687d38a6d2ccfbcbf0cc6aa738b12af6e9d142
+ md5: 1dafce8548e38671bea82e3f5c6ce22f
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 20591
+ timestamp: 1762976546182
+- conda: https://conda.anaconda.org/conda-forge/noarch/xyzservices-2025.11.0-pyhd8ed1ab_0.conda
+ sha256: b194a1fbc38f29c563b102ece9d006f7a165bf9074cdfe50563d3bce8cae9f84
+ md5: 16933322051fa260285f1a44aae91dd6
+ depends:
+ - python >=3.8
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/xyzservices?source=hash-mapping
+ size: 51128
+ timestamp: 1763813786075
+- conda: https://conda.anaconda.org/conda-forge/linux-64/yaml-0.2.5-h280c20c_3.conda
+ sha256: 6d9ea2f731e284e9316d95fa61869fe7bbba33df7929f82693c121022810f4ad
+ md5: a77f85f77be52ff59391544bfe73390a
+ depends:
+ - libgcc >=14
+ - __glibc >=2.17,<3.0.a0
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 85189
+ timestamp: 1753484064210
+- conda: https://conda.anaconda.org/conda-forge/osx-arm64/yaml-0.2.5-h925e9cb_3.conda
+ sha256: b03433b13d89f5567e828ea9f1a7d5c5d697bf374c28a4168d71e9464f5dafac
+ md5: 78a0fe9e9c50d2c381e8ee47e3ea437d
+ depends:
+ - __osx >=11.0
+ license: MIT
+ license_family: MIT
+ purls: []
+ size: 83386
+ timestamp: 1753484079473
+- conda: https://conda.anaconda.org/conda-forge/linux-64/yarl-1.22.0-py310h3406613_0.conda
+ sha256: b6e527196d2ce27417721cb9540d1efa2614bad76c9fbd2334b4cd39ddbae364
+ md5: ac707c966a3aa0c494ac763df31fa873
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - idna >=2.0
+ - libgcc >=14
+ - multidict >=4.0
+ - propcache >=0.2.1
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: Apache-2.0
+ license_family: Apache
+ purls:
+ - pkg:pypi/yarl?source=hash-mapping
+ size: 137784
+ timestamp: 1761337034085
+- conda: https://conda.anaconda.org/conda-forge/noarch/zarr-2.18.3-pyhd8ed1ab_1.conda
+ sha256: 02c045d3ab97bd5a713b0f35b05f017603d33bd728694ce3cf843c45c2906535
+ md5: 3e9a0fee25417c432c4780b9597fc312
+ depends:
+ - asciitree
+ - fasteners
+ - numcodecs >=0.10.0,<0.16.0a0
+ - numpy >=1.24,<3.0
+ - python >=3.10
+ constrains:
+ - notebook
+ - ipytree >=0.2.2
+ - ipywidgets >=8.0.0
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/zarr?source=hash-mapping
+ size: 160013
+ timestamp: 1733237313723
+- conda: https://conda.anaconda.org/conda-forge/linux-64/zfp-1.0.1-h909a3a2_5.conda
+ sha256: 5fabe6cccbafc1193038862b0b0d784df3dae84bc48f12cac268479935f9c8b7
+ md5: 6a0eb48e58684cca4d7acc8b7a0fd3c7
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - _openmp_mutex >=4.5
+ - libgcc >=14
+ - libstdcxx >=14
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 277694
+ timestamp: 1766549572069
+- conda: https://conda.anaconda.org/conda-forge/noarch/zict-3.0.0-pyhd8ed1ab_1.conda
+ sha256: 5488542dceeb9f2874e726646548ecc5608060934d6f9ceaa7c6a48c61f9cc8d
+ md5: e52c2ef711ccf31bb7f70ca87d144b9e
+ depends:
+ - python >=3.9
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/zict?source=hash-mapping
+ size: 36341
+ timestamp: 1733261642963
+- conda: https://conda.anaconda.org/conda-forge/noarch/zipp-3.23.0-pyhcf101f3_1.conda
+ sha256: b4533f7d9efc976511a73ef7d4a2473406d7f4c750884be8e8620b0ce70f4dae
+ md5: 30cd29cb87d819caead4d55184c1d115
+ depends:
+ - python >=3.10
+ - python
+ license: MIT
+ license_family: MIT
+ purls:
+ - pkg:pypi/zipp?source=hash-mapping
+ size: 24194
+ timestamp: 1764460141901
+- conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-1.3.1-hb9d3cd8_2.conda
+ sha256: 5d7c0e5f0005f74112a34a7425179f4eb6e73c92f5d109e6af4ddeca407c92ab
+ md5: c9f075ab2f33b3bbee9e62d4ad0a6cd8
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=13
+ - libzlib 1.3.1 hb9d3cd8_2
+ license: Zlib
+ license_family: Other
+ purls: []
+ size: 92286
+ timestamp: 1727963153079
+- conda: https://conda.anaconda.org/conda-forge/linux-64/zlib-ng-2.2.5-hde8ca8f_1.conda
+ sha256: 84ea17cb646d8a916d9335415f57c9e5dd001de158972322c714ebe1b72670b0
+ md5: c860578a89dc9b6003d600181612287c
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libgcc >=14
+ - libstdcxx >=14
+ license: Zlib
+ license_family: Other
+ purls: []
+ size: 110969
+ timestamp: 1764162891322
+- conda: https://conda.anaconda.org/conda-forge/linux-64/zstandard-0.23.0-py310h7c4b9e2_3.conda
+ sha256: 0653ad7d53d8c7b85ef2dd38c01c78b6c9185cd688be06cd6315e76530310635
+ md5: 64c494618303717a9a08e3238bcb8d68
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - cffi >=1.11
+ - libgcc >=14
+ - python >=3.10,<3.11.0a0
+ - python_abi 3.10.* *_cp310
+ license: BSD-3-Clause
+ license_family: BSD
+ purls:
+ - pkg:pypi/zstandard?source=hash-mapping
+ size: 477581
+ timestamp: 1756075706687
+- conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda
+ sha256: 68f0206ca6e98fea941e5717cec780ed2873ffabc0e1ed34428c061e2c6268c7
+ md5: 4a13eeac0b5c8e5b8ab496e6c4ddd829
+ depends:
+ - __glibc >=2.17,<3.0.a0
+ - libzlib >=1.3.1,<2.0a0
+ license: BSD-3-Clause
+ license_family: BSD
+ purls: []
+ size: 601375
+ timestamp: 1764777111296
diff --git a/plots.py b/plots.py
deleted file mode 100644
index bfdae97..0000000
--- a/plots.py
+++ /dev/null
@@ -1,66 +0,0 @@
-"""
-Code for generating various plots.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.7.24
-
-TODO: Add functions for plotting ERA5, GDAS, and GFS data.
-"""
-
-import argparse
-import xarray as xr
-import matplotlib.pyplot as plt
-import matplotlib as mpl
-from utils import settings
-from utils.data_utils import reformat_fronts, expand_fronts
-from utils.plotting_utils import plot_background
-import numpy as np
-import cartopy.crs as ccrs
-
-
-def plot_fronts(netcdf_indir, plot_outdir, timestep, front_types, domain, extent=(-180, 180, -90, 90)):
-
- year, month, day, hour = timestep
-
- fronts_ds = xr.open_dataset('%s/%d%02d/FrontObjects_%d%02d%02d%02d_%s.nc' % (netcdf_indir, year, month, year, month, day, hour, domain))
-
- if front_types is not None:
- fronts_ds = reformat_fronts(fronts_ds, front_types)
- labels = fronts_ds.attrs['labels']
-
- fronts_ds = expand_fronts(fronts_ds, iterations=1)
- fronts_ds = xr.where(fronts_ds == 0, np.nan, fronts_ds)
-
- front_colors_by_type = [settings.DEFAULT_FRONT_COLORS[label] for label in labels]
- front_names_by_type = [settings.DEFAULT_FRONT_NAMES[label] for label in labels]
- cmap_front = mpl.colors.ListedColormap(front_colors_by_type, name='from_list', N=len(front_colors_by_type))
- norm_front = mpl.colors.Normalize(vmin=1, vmax=len(front_colors_by_type) + 1)
-
- fig, ax = plt.subplots(1, 1, figsize=(16, 8), subplot_kw={'projection': ccrs.PlateCarree(central_longitude=np.mean(extent[:2]))})
- plot_background(extent, ax=ax, linewidth=0.25)
-
- cbar_front = plt.colorbar(mpl.cm.ScalarMappable(norm=norm_front, cmap=cmap_front), ax=ax, alpha=0.75, shrink=0.8, pad=0.02)
- cbar_front.set_ticks(np.arange(1, len(front_colors_by_type) + 1) + 0.5)
- cbar_front.set_ticklabels(front_names_by_type)
- cbar_front.set_label('Front Type')
-
- fronts_ds['identifier'].plot(ax=ax, x='longitude', y='latitude', cmap=cmap_front, norm=norm_front, transform=ccrs.PlateCarree(),
- add_colorbar=False)
- ax.gridlines(alpha=0.5)
-
- plt.tight_layout()
- plt.savefig(f"%s/fronts_%d%02d%02d%02d_{domain}.png" % (plot_outdir, year, month, day, hour), dpi=300, bbox_inches='tight')
- plt.close()
-
-
-if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- parser.add_argument('--timestep', type=int, nargs=4, help='Year, month, day, and hour of the data.')
- parser.add_argument('--netcdf_indir', type=str, help='Directory for the netcdf files.')
- parser.add_argument('--plot_outdir', type=str, help='Directory for the plots.')
- parser.add_argument('--front_types', type=str, nargs='+', help='Directory for the netcdf files.')
- parser.add_argument('--domain', type=str, default='full', help="Domain for which the fronts will be plotted.")
- parser.add_argument('--extent', type=float, nargs=4, default=[-180., 180., -90., 90.], help="Extent of the plot [min lon, max lon, min lat, max lat]")
- args = vars(parser.parse_args())
-
- plot_fronts(args['netcdf_indir'], args['plot_outdir'], args['timestep'], args['front_types'], args['domain'], args['extent'])
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..e7984fe
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,201 @@
+
+[project]
+name = "fronts"
+version = "0.2.0"
+description = "FrontFinder AI for finding various frontal and baroclinic boundaries"
+requires-python = ">=3.10,<3.14"
+keywords = [
+ "weather",
+ "weather fronts",
+ "cold fronts",
+ "unet",
+ "frontal boundary",
+ "machine learning",
+ "artificial intelligence",
+ "ai",
+ "ml",
+]
+classifiers = [
+ "Intended Audience :: Science/Research",
+ "Operating System :: OS Independent",
+ "Programming Language :: Python :: 3",
+ "Programming Language :: Python :: 3.11",
+ "Programming Language :: Python :: 3.12",
+ "Programming Language :: Python :: 3.13",
+ "Topic :: Scientific/Engineering :: Atmospheric Science",
+]
+license = "CC0-1.0"
+license-files = [
+ "LICENSE",
+]
+readme = "README.md"
+
+[tool.pixi.workspace]
+channels = ["conda-forge"]
+platforms = ["linux-64", "osx-arm64"]
+preview = ["pixi-build"]
+
+[tool.pixi.dependencies]
+numpy = ">=1.26.4,<2"
+dacite = ">=1.9.2,<2"
+
+[tool.pixi.feature.train]
+platforms = ["linux-64"]
+
+[tool.pixi.feature.train.dependencies]
+fronts = { path = "." }
+pandas = ">=2.3.3,<3"
+shapely = ">=2.0.1,<3"
+regionmask = ">=0.13.0,<0.14"
+wandb = ">=0.24"
+netcdf4 = ">=1.6.2,<2"
+dask = ">=2023.3.0,<2024"
+distributed = ">=2023.3.0,<2024"
+xarray = ">=2024.3.0,<2025"
+scikit-image = ">=0.22.0"
+h5py = ">=3.11.0"
+xbatcher = ">=0.4.0,<0.5"
+zarr = ">=2.18.3,<3"
+gcsfs = ">=2026.2.0,<2027"
+
+[tool.pixi.feature.train.pypi-dependencies]
+tensorflow = { extras = ["and-cuda"] }
+obstore = "*"
+s3fs = ">=2024.6.0"
+
+# ---------------------------------------------------------------------------
+# mac feature: same deps as train but with CPU-only TensorFlow for osx-arm64.
+# Use this environment for local dry-run testing on Apple Silicon.
+# ---------------------------------------------------------------------------
+[tool.pixi.feature.mac]
+platforms = ["osx-arm64"]
+
+[tool.pixi.feature.mac.dependencies]
+python = ">=3.10,<3.13"
+numpy = ">=1.26.4,<2"
+protobuf = ">=3.20,<5"
+dacite = ">=1.9.2,<2"
+fronts = { path = "." }
+pandas = ">=2.3.3,<3"
+shapely = ">=2.0.1,<3"
+regionmask = ">=0.13.0,<0.14"
+wandb = ">=0.24"
+netcdf4 = ">=1.6.2,<2"
+dask = ">=2023.3.0,<2024"
+distributed = ">=2023.3.0,<2024"
+xarray = ">=2024.3.0,<2025"
+scikit-image = ">=0.22.0"
+h5py = ">=3.11.0"
+xbatcher = ">=0.4.0,<0.5"
+zarr = ">=2.18.3,<3"
+# gcsfs omitted — not needed locally (no ARCO zarr store access on Mac)
+
+[tool.pixi.feature.mac.pypi-dependencies]
+# Plain tensorflow works on macOS ARM since TF 2.13 (unified with tensorflow-macos).
+# Pinned to 2.16.* — well-tested on ARM macOS, compatible with numpy<2.
+tensorflow = ">=2.16,<2.17"
+obstore = "*"
+s3fs = ">=2024.6.0"
+
+[tool.pixi.feature.test.dependencies]
+python = ">=3.10,<3.14"
+dacite = ">=1.9.2,<2"
+pyyaml = ">=6.0"
+numpy = ">=1.26.4,<2"
+
+[tool.pixi.feature.test.tasks]
+test = "PYTHONPATH=src python -m pytest tests/ -v"
+
+[tool.pixi.system-requirements]
+libc = "2.17"
+
+[tool.pixi.environments]
+default = { features = ["train"], solve-group = "default" }
+dev = { features = ["dev", "train"], solve-group = "default" }
+test = { features = ["test"], solve-group = "test", no-default-feature = true }
+mac = { features = ["mac"], solve-group = "mac", no-default-feature = true }
+
+[dependency-groups]
+dev = [
+ "deptry>=0.23.0",
+ "isort>=6.0.1",
+ "lxml-stubs>=0.5.1",
+ "mypy>=1.14.1",
+ "pandas-stubs>=2.3.0.250703",
+ "pre-commit>=4.1.0",
+ "pytest>=8.3.4",
+ "pytest-cov>=6.0.0",
+ "pytest-mock>=3.14.0",
+ "ruff>=0.9.4",
+ "types-pytz>=2025.2.0.20250809",
+ "types-pyyaml>=6.0.12.20241230",
+ "types-tqdm>=4.67.0.20250809",
+ "twine>=5.1.1",
+]
+
+test = [
+ "pytest>=8.3.4",
+ "pytest-cov>=6.0.0",
+ "pytest-mock>=3.14.0",
+]
+
+#TODO: swap wget with requests in code
+data = ["cdsapi", "wget"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/fronts"]
+
+[tool.hatch.build.targets.sdist]
+include = ["src/fronts/**/*"]
+
+[tool.pixi.package]
+name = "fronts"
+version = "0.1.0"
+
+[tool.pixi.package.host-dependencies]
+hatchling = ">=1.26.3"
+
+[tool.pixi.package.build]
+backend = { name = "pixi-build-python", version = "0.4.*" }
+
+[tool.ruff]
+line-length = 88
+indent-width = 4
+
+[tool.ruff.lint]
+ignore = []
+
+[tool.ruff.format]
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"
+
+# Enable auto-formatting of code examples in docstrings. Markdown,
+# reStructuredText code/literal blocks and doctests are all supported.
+#
+# This is currently disabled by default, but it is planned for this
+# to be opt-out in the future.
+docstring-code-format = true
+
+# Set the line length limit used when formatting code snippets in
+# docstrings.
+#
+# This only has an effect when the `docstring-code-format` setting is
+# enabled.
+docstring-code-line-length = "dynamic"
+
+[tool.ruff.lint.isort]
+case-sensitive = true
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..4e00101
Binary files /dev/null and b/requirements.txt differ
diff --git a/scripts/check_era5_files.py b/scripts/check_era5_files.py
new file mode 100644
index 0000000..9f1afdb
--- /dev/null
+++ b/scripts/check_era5_files.py
@@ -0,0 +1,51 @@
+"""
+Scans for missing or corrupt ERA5 files in a directory.
+* ERA5 files must be downloaded from the Climate Data Store using the 'download_era5.py' script.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import numpy as np
+import itertools
+import xarray as xr
+import argparse
+import os
+
+YEARS = np.arange(2008, 2024.1, 1).astype(int)
+VARS_TO_CHECK = [
+ "u_component_of_wind",
+ "v_component_of_wind",
+ "temperature",
+ "geopotential",
+ "specific_humidity",
+]
+PRESSURE_LEVELS_TO_CHECK = ["1000", "950", "900", "850", "700", "500", "300"]
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--dir", type=str, help="Directory where the ERA5 files are stored."
+ )
+ args = vars(parser.parse_args())
+
+ missing = []
+ bad = []
+
+ for var, lvl, yr in itertools.product(
+ VARS_TO_CHECK, PRESSURE_LEVELS_TO_CHECK, YEARS
+ ):
+ file = args["dir"] + "/era5_%s-%shPa_%d.nc" % (var, lvl, yr)
+
+ if not os.path.isfile(file):
+ missing.append([var, lvl, yr])
+ else:
+ try:
+ xr.open_dataset(file)
+ except:
+ bad.append(file)
+
+ print("\n=== MISSING FILES ===")
+ print(missing)
+ print("\n=== BAD DATASETS ===")
+ print(bad)
diff --git a/scripts/convert_front_gml_to_xml.py b/scripts/convert_front_gml_to_xml.py
new file mode 100644
index 0000000..343faaa
--- /dev/null
+++ b/scripts/convert_front_gml_to_xml.py
@@ -0,0 +1,190 @@
+"""
+Convert GML files containing TWC fronts into XML files.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import argparse
+from lxml import etree as ET
+from glob import glob
+import os
+import numpy as np
+
+XML_FRONT_TYPE = {
+ "Cold Front": "COLD_FRONT",
+ "Dissipating Cold Front": "COLD_FRONT_DISS",
+ "Warm Front": "WARM_FRONT",
+ "Stationary Front": "STATIONARY_FRONT",
+ "Occluded Front": "OCCLUDED_FRONT",
+ "Dissipating Occluded Front": "OCCLUDED_FRONT_DISS",
+ "Dry Line": "DRY_LINE",
+ "Trough": "TROF",
+ "Squall Line": "INSTABILITY",
+}
+
+XML_FRONT_COLORS = {
+ "Cold Front": dict(red="0", green="0", blue="255"),
+ "Dissipating Cold Front": dict(red="0", green="0", blue="255"),
+ "Warm Front": dict(red="255", green="0", blue="0"),
+ "Dissipating Warm Front": dict(red="255", green="0", blue="0"),
+ "Occluded Front": dict(red="145", green="44", blue="238"),
+ "Dissipating Occluded Front": dict(red="145", green="44", blue="238"),
+ "Dry Line": dict(red="255", green="130", blue="71"),
+ "Trough": dict(red="255", green="130", blue="71"),
+ "Squall Line": dict(red="255", green="0", blue="0"),
+}
+
+LINE_KWARGS = dict(
+ pgenCategory="Front",
+ lineWidth="4",
+ sizeScale=" 1.0",
+ smoothFactor="2",
+ closed="false",
+ filled="false",
+ fillPattern="SOLID",
+)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--gml_indir",
+ type=str,
+ required=True,
+ help="Input directory for TWC front GML files.",
+ )
+ parser.add_argument(
+ "--xml_outdir",
+ type=str,
+ required=True,
+ help="Output directory for front XML files.",
+ )
+ parser.add_argument(
+ "--date",
+ type=int,
+ nargs=3,
+ required=True,
+ help="Date for the data to be read in. (year, month, day)",
+ )
+ args = vars(parser.parse_args())
+
+ year, month, day = args["date"]
+
+ gml_files = sorted(
+ glob(
+ "%s/%d%02d%02d/*/*%d%02d%02d*.gml"
+ % (args["gml_indir"], year, month, day, year, month, day)
+ )
+ )
+
+ for gml_file in gml_files:
+ valid_time_str = os.path.basename(gml_file).split(".")[2]
+ valid_time_str = (
+ valid_time_str[:4]
+ + "-"
+ + valid_time_str[4:6]
+ + "-"
+ + valid_time_str[6:8]
+ + "T"
+ + valid_time_str[9:11]
+ )
+ valid_time = np.datetime64(valid_time_str, "ns")
+
+ init_time_str = os.path.basename(gml_file).split(".")[3]
+ if (
+ init_time_str != "NIL"
+ ): # an init time of 'NIL' is used to indicate forecast hour 0 (i.e. valid time is same as init time)
+ init_time_str = (
+ init_time_str[:4]
+ + "-"
+ + init_time_str[4:6]
+ + "-"
+ + init_time_str[6:8]
+ + "T"
+ + init_time_str[9:11]
+ )
+ init_time = np.datetime64(init_time_str, "ns")
+ else:
+ init_time_str = valid_time_str
+ init_time = valid_time
+
+ forecast_hour = int((valid_time - init_time) / np.timedelta64(1, "h"))
+
+ root_xml = ET.Element(
+ "Product",
+ name="TWC_global_fronts",
+ init_time=init_time_str,
+ valid_time=valid_time_str,
+ forecast_hour=str(forecast_hour),
+ )
+ tree = ET.parse(gml_file, parser=ET.XMLPullParser(encoding="utf-8"))
+ root_gml = tree.getroot()
+
+ Layer = ET.SubElement(
+ root_xml,
+ "Layer",
+ name="Default",
+ onOff="true",
+ monoColor="false",
+ filled="false",
+ )
+ ET.SubElement(Layer, "Color", red="255", green="255", blue="0", alpha="255")
+ DrawableElement = ET.SubElement(Layer, "DrawableElement")
+
+ front_elements = [
+ element[0] for element in root_gml if element[0].tag == "FRONT"
+ ]
+
+ for element in front_elements:
+ front_type = [
+ subelement.text
+ for subelement in element
+ if subelement.tag == "FRONT_TYPE"
+ ][0]
+ coords = [
+ subelement for subelement in element if "lineString" in subelement.tag
+ ][0][0][0].text
+
+ Line = ET.SubElement(
+ DrawableElement,
+ "Line",
+ pgenType=XML_FRONT_TYPE[front_type],
+ **LINE_KWARGS,
+ )
+ if front_type == "Stationary Front":
+ ET.SubElement(
+ Line, "Color", red="255", green="0", blue="0", alpha="255"
+ )
+ ET.SubElement(
+ Line, "Color", red="0", green="0", blue="255", alpha="255"
+ )
+ else:
+ ET.SubElement(
+ Line, "Color", **XML_FRONT_COLORS[front_type], alpha="255"
+ )
+
+ coords = coords.replace("\n", "").split(" ") # generate coordinate strings
+ coords = list(
+ coord_pair.split(",") for coord_pair in coords
+ ) # generate coordinate pairs from the strings
+
+ for coord_pair in coords:
+ ET.SubElement(
+ Line,
+ "Point",
+ Lat="%.6f" % float(coord_pair[1]),
+ Lon="%.6f" % float(coord_pair[0]),
+ )
+
+ save_path_file = "%s/TWC_fronts_%sf%03d.xml" % (
+ args["xml_outdir"],
+ init_time_str.replace("-", "").replace("T", ""),
+ forecast_hour,
+ )
+
+ ET.indent(root_xml)
+ mydata = ET.tostring(root_xml)
+ xmlFile = open(save_path_file, "wb")
+ xmlFile.write(mydata)
+ xmlFile.close()
diff --git a/scripts/convert_front_xml_to_netcdf.py b/scripts/convert_front_xml_to_netcdf.py
new file mode 100644
index 0000000..93a1616
--- /dev/null
+++ b/scripts/convert_front_xml_to_netcdf.py
@@ -0,0 +1,329 @@
+"""
+Convert front XML files to netCDF files.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import argparse
+import glob
+import numpy as np
+import os
+from fronts.utils import data_utils
+import xarray as xr
+import xml.etree.ElementTree as ET
+
+
+pgenType_identifiers = {
+ "COLD_FRONT": 1,
+ "WARM_FRONT": 2,
+ "STATIONARY_FRONT": 3,
+ "OCCLUDED_FRONT": 4,
+ "COLD_FRONT_FORM": 5,
+ "WARM_FRONT_FORM": 6,
+ "STATIONARY_FRONT_FORM": 7,
+ "OCCLUDED_FRONT_FORM": 8,
+ "COLD_FRONT_DISS": 9,
+ "WARM_FRONT_DISS": 10,
+ "STATIONARY_FRONT_DISS": 11,
+ "OCCLUDED_FRONT_DISS": 12,
+ "INSTABILITY": 13,
+ "TROF": 14,
+ "TROPICAL_TROF": 15,
+ "DRY_LINE": 16,
+}
+
+"""
+conus: 132 W to 60.25 W, 57 N to 26.25 N
+full: 130 E pointing eastward to 10 E, 80 N to 0.25 N
+global: 179.75 W to 180 E, 90 N to 89.75 S
+"""
+domain_coords = {
+ "conus": {"lons": np.arange(-132, -60, 0.25), "lats": np.arange(57, 25, -0.25)},
+ "full": {
+ "lons": np.append(np.arange(-179.75, 10, 0.25), np.arange(130, 180.25, 0.25)),
+ "lats": np.arange(80, 0, -0.25),
+ },
+ "global": {
+ "lons": np.arange(-179.75, 180.25, 0.25),
+ "lats": np.arange(90, -90, -0.25),
+ },
+}
+domain_lons_360 = {
+ "conus": np.arange(228, 300, 0.25),
+ "full": np.arange(130, 370, 0.25),
+ "global": np.arange(0, 360, 0.25),
+}
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--xml_indir",
+ type=str,
+ required=True,
+ help="Input directory for front XML files.",
+ )
+ parser.add_argument(
+ "--netcdf_outdir",
+ type=str,
+ required=True,
+ help="Output directory for front netCDF files.",
+ )
+ parser.add_argument(
+ "--date",
+ type=str,
+ required=True,
+ help="Date for the data to be read in. YYYY-MM-DD",
+ )
+ parser.add_argument(
+ "--distance",
+ type=float,
+ default=1.0,
+ help="Interpolation distance in kilometers.",
+ )
+ parser.add_argument(
+ "--domain",
+ type=str,
+ default="full",
+ help="Domain for which to generate fronts.",
+ )
+
+ args = vars(parser.parse_args())
+
+ date = np.datetime64(args["date"]).astype(object)
+ year, month, day = date.year, date.month, date.day
+
+ os.makedirs("%s/%d%02d" % (args["netcdf_outdir"], year, month), exist_ok=True)
+
+ if args["domain"] == "global":
+ files = sorted(
+ glob.glob(
+ "%s/TWC*_%04d%02d%02d*f*.xml" % (args["xml_indir"], year, month, day)
+ )
+ )
+ else:
+ files = sorted(
+ glob.glob(
+ "%s/pres*_%04d%02d%02d*f000.xml" % (args["xml_indir"], year, month, day)
+ )
+ )
+
+ domain_from_model = args["domain"] not in ["conus", "full", "global"]
+
+ for filename in files:
+ if domain_from_model:
+ model_coords_ds = xr.open_dataset("./coordinates/%s.nc" % args["domain"])
+
+ # transform model's coordinates to a cartesian grid
+ transform_args = dict(
+ std_parallels=model_coords_ds.attrs["std_parallels"],
+ lon_ref=model_coords_ds.attrs["lon_ref"],
+ lat_ref=model_coords_ds.attrs["lat_ref"],
+ )
+ gridded_lons = model_coords_ds["longitude"].values.astype("float32")
+ gridded_lats = model_coords_ds["latitude"].values.astype("float32")
+ model_x_transform, model_y_transform = (
+ data_utils.lambert_conformal_to_cartesian(
+ gridded_lons, gridded_lats, **transform_args
+ )
+ )
+ gridded_x = model_x_transform[0, :]
+ gridded_y = model_y_transform[:, 0]
+ identifier = np.zeros(np.shape(gridded_lons)).astype("float32")
+
+ # bounds of the model's cartesian domain
+ model_x_min, model_x_max = (
+ np.min(model_x_transform),
+ np.max(model_x_transform),
+ )
+ model_y_min, model_y_max = (
+ np.min(model_y_transform),
+ np.max(model_y_transform),
+ )
+
+ else:
+ gridded_lons = domain_coords[args["domain"]]["lons"].astype("float32")
+ gridded_lats = domain_coords[args["domain"]]["lats"].astype("float32")
+ identifier = np.zeros([len(gridded_lons), len(gridded_lats)]).astype(
+ "float32"
+ )
+
+ tree = ET.parse(filename, parser=ET.XMLParser(encoding="utf-8"))
+ root = tree.getroot()
+ date = (
+ os.path.basename(filename).split("_")[-1].split(".")[0].split("f")[0]
+ ) # YYYYMMDDhh
+ forecast_hour = int(filename.split("f")[-1].split(".")[0])
+
+ hour = date[-2:]
+
+ if hour in ["03", "09", "15", "21"] and args["domain"] == "nam-12km":
+ continue
+
+ ### Iterate through the individual fronts ###
+ for line in root.iter("Line"):
+ type_of_front = line.get("pgenType") # front type
+
+ lons, lats = zip(
+ *[
+ [float(point.get("Lon")), float(point.get("Lat"))]
+ for point in line.iter("Point")
+ ]
+ )
+ lons, lats = np.array(lons), np.array(lats)
+
+ # if there are less than 2 points, skip the front as it cannot be interpolated
+ if len(lons) < 2:
+ print(f"Bad front: {type_of_front}")
+ continue
+
+ # If the front crosses the dateline or the 180th meridian, its coordinates must be modified for proper interpolation
+ front_needs_modification = np.max(np.abs(np.diff(lons))) > 180
+
+ if front_needs_modification or domain_from_model:
+ lons = np.where(
+ lons < 0, lons + 360, lons
+ ) # convert coordinates to a 360 degree system
+
+ if domain_from_model:
+ x_transform_init, y_transform_init = (
+ data_utils.lambert_conformal_to_cartesian(
+ lons, lats, **transform_args
+ )
+ )
+
+ # find points outside the model domain
+ points_outside_domain = np.where(
+ (x_transform_init < model_x_min)
+ | (x_transform_init > model_x_max)
+ | (y_transform_init < model_y_min)
+ | (y_transform_init > model_y_max)
+ )
+
+ if (
+ len(points_outside_domain) > 0
+ ): # remove points outside the model domain
+ lons = np.delete(lons, points_outside_domain)
+ lats = np.delete(lats, points_outside_domain)
+
+ if (
+ len(lons) < 2
+ ): # do not generate front if there are not at least two points
+ continue
+
+ xs, ys = data_utils.haversine(lons, lats) # x/y coordinates in kilometers
+ xy_linestring = data_utils.geometric(
+ xs, ys
+ ) # convert coordinates to a LineString object
+ x_new, y_new = data_utils.redistribute_vertices(
+ xy_linestring, args["distance"]
+ ).xy # interpolate x/y coordinates
+ x_new, y_new = np.array(x_new), np.array(y_new)
+ lon_new, lat_new = data_utils.reverse_haversine(
+ x_new, y_new
+ ) # convert interpolated x/y coordinates to lat/lon
+
+ date_and_time = np.datetime64(
+ "%04d-%02d-%02dT%02d" % (year, month, day, int(hour)), "ns"
+ )
+
+ expand_dims_args = {"time": np.atleast_1d(date_and_time)}
+
+ if args["domain"] == "global" or domain_from_model:
+ expand_dims_args["forecast_hour"] = np.atleast_1d(forecast_hour)
+ filename_netcdf = "FrontObjects_%s_f%03d_%s.nc" % (
+ date,
+ forecast_hour,
+ args["domain"],
+ )
+ else:
+ filename_netcdf = "FrontObjects_%s_%s.nc" % (date, args["domain"])
+
+ if domain_from_model:
+ x_new *= 1000 # convert to meters
+ y_new *= 1000 # convert to meters
+ x_transform, y_transform = data_utils.lambert_conformal_to_cartesian(
+ lon_new, lat_new, **transform_args
+ )
+
+ gridded_indices = np.dstack(
+ (
+ np.digitize(y_transform, gridded_y),
+ np.digitize(x_transform, gridded_x),
+ )
+ )[0] # translate coordinate indices to grid
+ gridded_indices_unique = np.unique(
+ gridded_indices, axis=0
+ ) # remove duplicate coordinate indices
+
+ # Remove points outside the domain
+ gridded_indices_unique = gridded_indices_unique[
+ np.where(gridded_indices_unique[:, 0] != len(gridded_y))
+ ]
+ gridded_indices_unique = gridded_indices_unique[
+ np.where(gridded_indices_unique[:, 1] != len(gridded_x))
+ ]
+
+ identifier[
+ gridded_indices_unique[:, 0], gridded_indices_unique[:, 1]
+ ] = pgenType_identifiers[
+ type_of_front
+ ] # assign labels to the gridded points based on the front type
+
+ fronts_ds = xr.Dataset(
+ {"identifier": (("y", "x"), identifier)}
+ ).expand_dims(**expand_dims_args)
+
+ else:
+ if front_needs_modification:
+ lon_new = np.where(
+ lon_new > 180, lon_new - 360, lon_new
+ ) # convert new longitudes to standard -180 to 180 range
+
+ gridded_indices = np.dstack(
+ (
+ np.digitize(lon_new, gridded_lons),
+ np.digitize(lat_new, gridded_lats),
+ )
+ )[0] # translate coordinate indices to grid
+ gridded_indices_unique = np.unique(
+ gridded_indices, axis=0
+ ) # remove duplicate coordinate indices
+
+ # Remove points outside the domain
+ gridded_indices_unique = gridded_indices_unique[
+ np.where(gridded_indices_unique[:, 0] != len(gridded_lons))
+ ]
+ gridded_indices_unique = gridded_indices_unique[
+ np.where(gridded_indices_unique[:, 1] != len(gridded_lats))
+ ]
+
+ identifier[
+ gridded_indices_unique[:, 0], gridded_indices_unique[:, 1]
+ ] = pgenType_identifiers[
+ type_of_front
+ ] # assign labels to the gridded points based on the front type
+
+ fronts_ds = xr.Dataset(
+ {"identifier": (("latitude", "longitude"), identifier.transpose())},
+ coords={"latitude": gridded_lats, "longitude": gridded_lons},
+ ).expand_dims(**expand_dims_args)
+
+ if args["domain"] == "full":
+ fronts_ds = fronts_ds.sel(
+ longitude=np.append(
+ np.arange(130, 180.01, 0.25), np.arange(-179.75, 10, 0.25)
+ )
+ )
+
+ fronts_ds.attrs["domain"] = args["domain"]
+ fronts_ds.attrs["interpolation_distance_km"] = args["distance"]
+ fronts_ds.attrs["num_front_types"] = 16
+ fronts_ds.attrs["front_types"] = "ALL"
+
+ fronts_ds.to_netcdf(
+ path="%s/%d%02d/%s" % (args["netcdf_outdir"], year, month, filename_netcdf),
+ engine="netcdf4",
+ mode="w",
+ )
diff --git a/scripts/convert_grib_to_netcdf.py b/scripts/convert_grib_to_netcdf.py
new file mode 100644
index 0000000..08d1e3f
--- /dev/null
+++ b/scripts/convert_grib_to_netcdf.py
@@ -0,0 +1,464 @@
+"""
+Convert GDAS and/or GFS grib files to netCDF files.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import argparse
+import time
+import xarray as xr
+from fronts.utils import variables
+import glob
+import numpy as np
+import os
+import tensorflow as tf
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--grib_indir", type=str, required=True, help="Input directory for grib files."
+ )
+ parser.add_argument(
+ "--netcdf_outdir",
+ type=str,
+ required=True,
+ help="Output directory for the netcdf files.",
+ )
+ parser.add_argument(
+ "--model", required=True, type=str, help="NWP model for the grib files."
+ )
+ parser.add_argument(
+ "--init_time",
+ type=str,
+ required=True,
+ help="Initialization time of the model. Format: YYYY-MM-DDTHH.",
+ )
+ parser.add_argument(
+ "--pressure_levels",
+ type=str,
+ nargs="+",
+ default=["surface", "1000", "950", "900", "850"],
+ help="Pressure levels to extract from the grib files.",
+ )
+ parser.add_argument(
+ "--gpu",
+ action="store_true",
+ help="Use a GPU to perform calculations of additional variables. This can provide speedups when generating very "
+ "large amounts of data.",
+ )
+ parser.add_argument(
+ "--ignore_warnings",
+ action="store_true",
+ help="Disable runtime warnings in variable calculations.",
+ )
+ args = vars(parser.parse_args())
+
+ if args["ignore_warnings"]: # suppress divide by zero RuntimeWarnings
+ import warnings
+
+ warnings.filterwarnings("ignore")
+
+ gpus = tf.config.list_physical_devices(device_type="GPU")
+ if len(gpus) > 0 and args["gpu"]:
+ print("Using GPU for variable derivations")
+ tf.config.set_visible_devices(devices=gpus[0], device_type="GPU")
+ gpus = tf.config.get_visible_devices(device_type="GPU")
+ tf.config.experimental.set_memory_growth(device=gpus[0], enable=True)
+ else:
+ print("Using CPUs for variable derivations")
+ tf.config.set_visible_devices([], "GPU")
+
+ args["model"] = args["model"].lower()
+
+ date = np.datetime64(args["init_time"]).astype(object)
+ year, month, day, hour = date.year, date.month, date.day, date.hour
+
+ # filename format of the downloaded grib files
+ grib_filename_format = f"%s/%d%02d/%s_%d%02d%02d%02d_f*.grib" % (
+ args["grib_indir"],
+ year,
+ month,
+ args["model"],
+ year,
+ month,
+ day,
+ hour,
+ )
+ grib_files = list(glob.glob(grib_filename_format))
+
+ # list of pressure levels (NOT including surface level)
+ pressure_levels = [lvl for lvl in args["pressure_levels"] if lvl != "surface"]
+
+ # boolean flags
+ include_pressure_level_data = len(pressure_levels) > 0
+ include_surface_data = True if "surface" in args["pressure_levels"] else False
+
+ # keyword arguments used in xr.open_mfdataset
+ open_ds_args = dict(
+ engine="cfgrib", errors="ignore", combine="nested", concat_dim="valid_time"
+ )
+
+ if include_pressure_level_data:
+ print("Reading pressure level data", end="")
+ start_time = time.time()
+ # open pressure level data from grib files
+ if (
+ args["model"] not in ["nam-12km", "namnest-conus"]
+ and "gefs" not in args["model"]
+ ):
+ pressure_level_data = xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={"filter_by_keys": {"typeOfLevel": "isobaricInhPa"}},
+ **open_ds_args,
+ )
+ pressure_level_data = pressure_level_data.sel(
+ isobaricInhPa=pressure_levels
+ ) # select pressure levels
+ pressure_level_data = pressure_level_data[
+ ["t", "u", "v", "r", "gh"]
+ ] # select variables
+
+ # lat/lon coordinates, time variable
+ latitude = pressure_level_data["latitude"].values
+ longitude = pressure_level_data["longitude"].values
+ timestep = pressure_level_data["time"].values
+
+ T = pressure_level_data["t"].values # temperature
+ u = pressure_level_data["u"].values # u-wind
+ v = pressure_level_data["v"].values # v-wind
+ RH = pressure_level_data["r"].values / 100 # relative humidity
+ sp_z = pressure_level_data["gh"].values / 10 # geopotential height (dam)
+
+ else: # need to open NAM, GEFS grib files one variable at a time (typical cfgrib nonsense)
+ T = (
+ xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={
+ "filter_by_keys": {
+ "typeOfLevel": "isobaricInhPa",
+ "cfVarName": "t",
+ }
+ },
+ **open_ds_args,
+ )
+ .sel(isobaricInhPa=pressure_levels)["t"]
+ .values
+ ) # temperature
+ u = (
+ xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={
+ "filter_by_keys": {
+ "typeOfLevel": "isobaricInhPa",
+ "cfVarName": "u",
+ }
+ },
+ **open_ds_args,
+ )
+ .sel(isobaricInhPa=pressure_levels)["u"]
+ .values
+ ) # u-wind
+ v = (
+ xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={
+ "filter_by_keys": {
+ "typeOfLevel": "isobaricInhPa",
+ "cfVarName": "v",
+ }
+ },
+ **open_ds_args,
+ )
+ .sel(isobaricInhPa=pressure_levels)["v"]
+ .values
+ ) # v-wind
+ RH = (
+ xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={
+ "filter_by_keys": {
+ "typeOfLevel": "isobaricInhPa",
+ "cfVarName": "r",
+ }
+ },
+ **open_ds_args,
+ )
+ .sel(isobaricInhPa=pressure_levels)["r"]
+ .values
+ / 100
+ ) # relative humidity
+ sp_z = (
+ xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={
+ "filter_by_keys": {
+ "typeOfLevel": "isobaricInhPa",
+ "cfVarName": "gh",
+ }
+ },
+ **open_ds_args,
+ )
+ .sel(isobaricInhPa=pressure_levels)["gh"]
+ .values
+ / 10
+ ) # geopotential height (dam)
+
+ # lat/lon coordinates
+ latitude = xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={"filter_by_keys": {"typeOfLevel": "isobaricInhPa"}},
+ **open_ds_args,
+ )["latitude"].values
+ longitude = xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={"filter_by_keys": {"typeOfLevel": "isobaricInhPa"}},
+ **open_ds_args,
+ )["longitude"].values
+ timestep = xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={"filter_by_keys": {"typeOfLevel": "isobaricInhPa"}},
+ **open_ds_args,
+ )["time"].values
+
+ # pressure array for calculating additional variables. shape: (forecast hour, pressure level, latitude/y, longitude/x)
+ P = np.zeros_like(T)
+ for i, lvl in enumerate(pressure_levels):
+ P[:, i, ...] = int(lvl) * 100 # convert pressure to Pa
+ print(" (%.1f seconds)" % (time.time() - start_time))
+
+ if include_surface_data:
+ print("Reading surface data", end="")
+ start_time = time.time()
+ if args["model"] in ["gfs", "gdas"]:
+ sp_data = xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={
+ "filter_by_keys": {"typeOfLevel": "surface", "stepType": "instant"}
+ },
+ **open_ds_args,
+ )["sp"].values[:, np.newaxis, ...]
+ surface_data = xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={
+ "filter_by_keys": {"typeOfLevel": "sigma", "stepType": "instant"}
+ },
+ **open_ds_args,
+ )
+ elif "gefs" in args["model"]:
+ sp_data = xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={
+ "filter_by_keys": {"typeOfLevel": "surface", "stepType": "instant"}
+ },
+ **open_ds_args,
+ )["sp"].values[:, np.newaxis, ...]
+ surface_data = xr.merge(
+ [
+ xr.open_mfdataset(
+ grib_files, filter_by_keys={"paramId": 167}, **open_ds_args
+ )
+ .rename({"t2m": "t"})
+ .drop_vars("heightAboveGround"),
+ xr.open_mfdataset(
+ grib_files, filter_by_keys={"paramId": 165}, **open_ds_args
+ )
+ .rename({"u10": "u"})
+ .drop_vars("heightAboveGround"),
+ xr.open_mfdataset(
+ grib_files, filter_by_keys={"paramId": 166}, **open_ds_args
+ )
+ .rename({"v10": "v"})
+ .drop_vars("heightAboveGround"),
+ xr.open_mfdataset(
+ grib_files, filter_by_keys={"paramId": 260242}, **open_ds_args
+ )
+ .rename({"r2": "r"})
+ .drop_vars("heightAboveGround"),
+ ]
+ )
+ elif args["model"] == "hrrr":
+ sp_data = xr.open_mfdataset(
+ grib_files,
+ backend_kwargs={
+ "filter_by_keys": {"typeOfLevel": "surface", "stepType": "instant"}
+ },
+ **open_ds_args,
+ )["sp"].values[:, np.newaxis, ...]
+ # surface_data = xr.open_mfdataset(grib_files, backend_kwargs={"filter_by_keys": {'typeOfLevel': 'surface', 'stepType': 'instant'}}, **open_ds_args)
+ surface_data = xr.merge(
+ [
+ xr.open_mfdataset(
+ grib_files, filter_by_keys={"paramId": 167}, **open_ds_args
+ )
+ .rename({"t2m": "t"})
+ .drop_vars("heightAboveGround"),
+ xr.open_mfdataset(
+ grib_files, filter_by_keys={"paramId": 165}, **open_ds_args
+ )
+ .rename({"u10": "u"})
+ .drop_vars("heightAboveGround"),
+ xr.open_mfdataset(
+ grib_files, filter_by_keys={"paramId": 166}, **open_ds_args
+ )
+ .rename({"v10": "v"})
+ .drop_vars("heightAboveGround"),
+ xr.open_mfdataset(
+ grib_files, filter_by_keys={"paramId": 260242}, **open_ds_args
+ )
+ .rename({"r2": "r"})
+ .drop_vars("heightAboveGround"),
+ ]
+ )
+ else: # NAM
+ raise NotImplementedError("NAM surface data currently not supported.")
+ # TODO: Figure out how to get NAM surface data implemented
+ # T_da = xr.open_mfdataset(grib_files, backend_kwargs={"filter_by_keys": {'typeOfLevel': 'sigma', 'cfVarName': 't'}}, **open_ds_args)['t'] # temperature
+ # u_da = xr.open_mfdataset(grib_files, backend_kwargs={"filter_by_keys": {'typeOfLevel': 'sigma', 'cfVarName': 'u'}}, **open_ds_args)['u'] # u-wind
+ # v_da = xr.open_mfdataset(grib_files, backend_kwargs={"filter_by_keys": {'typeOfLevel': 'sigma', 'cfVarName': 'v'}}, **open_ds_args)['v'] # v-wind
+ # RH_da = xr.open_mfdataset(grib_files, backend_kwargs={"filter_by_keys": {'typeOfLevel': 'sigma', 'cfVarName': 'r'}}, **open_ds_args)['r'] / 100 # relative humidity
+ # sp_z_da = xr.open_mfdataset(grib_files, backend_kwargs={"filter_by_keys": {'typeOfLevel': 'sigma', 'cfVarName': 'sp'}}, **open_ds_args)['sp'] / 10 # geopotential height (dam)
+
+ # lat/lon coordinates, time variable
+ latitude = surface_data["latitude"].values
+ longitude = surface_data["longitude"].values
+ timestep = surface_data["time"].values
+
+ if include_pressure_level_data:
+ # combine surface and pressure level data
+ T = np.concatenate(
+ [surface_data["t"].values[:, np.newaxis, ...], T], axis=1
+ )
+ u = np.concatenate(
+ [surface_data["u"].values[:, np.newaxis, ...], u], axis=1
+ )
+ v = np.concatenate(
+ [surface_data["v"].values[:, np.newaxis, ...], v], axis=1
+ )
+ RH = np.concatenate(
+ [surface_data["r"].values[:, np.newaxis, ...] / 100, RH], axis=1
+ )
+ P = np.concatenate([sp_data, P], axis=1)
+ sp_z = np.concatenate([sp_data / 100, sp_z], axis=1)
+ else:
+ T = surface_data["t"].values[:, np.newaxis, ...]
+ u = surface_data["u"].values[:, np.newaxis, ...]
+ v = surface_data["v"].values[:, np.newaxis, ...]
+ RH = surface_data["r"].values[:, np.newaxis, ...] / 100
+ P = sp_data
+ sp_z = sp_data / 100
+ print(" (%.1f seconds)" % (time.time() - start_time))
+
+ # if using a GPU, convert the numpy arrays to tensors
+ print("Calculating additional variables", end="")
+ start_time = time.time()
+ if len(gpus) > 0 and args["gpu"]:
+ T = tf.convert_to_tensor(T)
+ u = tf.convert_to_tensor(u)
+ v = tf.convert_to_tensor(v)
+ RH = tf.convert_to_tensor(RH)
+ P = tf.convert_to_tensor(P)
+
+ # calculate additional variables
+ q = variables.specific_humidity_from_relative_humidity(P, T, RH)
+ Td = variables.dewpoint_from_specific_humidity(P, q)
+ Tv = variables.virtual_temperature_from_dewpoint(P, T, Td)
+ r = variables.mixing_ratio_from_dewpoint(P, Td) * 1000 # convert back to g/kg
+ theta = variables.potential_temperature(P, T)
+ theta_e = variables.equivalent_potential_temperature(P, T, Td)
+ theta_v = variables.virtual_potential_temperature(P, T, Td)
+ print(" (%.1f seconds)" % (time.time() - start_time))
+
+ # if a GPU was used to calculate additional variables, turn the tensors back to numpy arrays
+ if len(gpus) > 0 and args["gpu"]:
+ T = T.numpy()
+ Td = Td.numpy()
+ Tv = Tv.numpy()
+ q = q.numpy()
+ r = r.numpy()
+ u = u.numpy()
+ v = v.numpy()
+ RH = RH.numpy()
+ theta = theta.numpy()
+ theta_e = theta_e.numpy()
+ theta_v = theta_v.numpy()
+
+ forecast_hours = [
+ int(filename.split("_f")[1][:3]) for filename in grib_files
+ ] # can pull forecast hours straight from the grib filenames
+
+ if args["model"] in ["gfs", "gdas"] or "gefs" in args["model"]:
+ dataset_dimensions = (
+ "forecast_hour",
+ "pressure_level",
+ "latitude",
+ "longitude",
+ )
+ else: # HRRR/NAM - both have non-uniform lat/lon grids
+ dataset_dimensions = ("forecast_hour", "pressure_level", "y", "x")
+ latitude = (("y", "x"), latitude)
+ longitude = (("y", "x"), longitude)
+
+ print("Building final datasets", end="")
+ start_time = time.time()
+ full_dataset_coordinates = dict(
+ forecast_hour=forecast_hours,
+ pressure_level=args["pressure_levels"],
+ latitude=latitude,
+ longitude=longitude,
+ )
+
+ full_dataset_variables = dict(
+ T=(dataset_dimensions, T),
+ Td=(dataset_dimensions, Td),
+ Tv=(dataset_dimensions, Tv),
+ theta=(dataset_dimensions, theta),
+ theta_e=(dataset_dimensions, theta_e),
+ theta_v=(dataset_dimensions, theta_v),
+ RH=(dataset_dimensions, RH),
+ r=(dataset_dimensions, r),
+ q=(dataset_dimensions, q * 1000),
+ u=(dataset_dimensions, u),
+ v=(dataset_dimensions, v),
+ sp_z=(dataset_dimensions, sp_z),
+ )
+
+ full_dataset = xr.Dataset(
+ data_vars=full_dataset_variables, coords=full_dataset_coordinates
+ ).astype("float32")
+
+ # final dataset attributes
+ full_dataset.attrs = dict(
+ grib_to_netcdf_script_version="2024.6.6",
+ model=args["model"],
+ Nx=np.shape(T)[-1],
+ Ny=np.shape(T)[-2],
+ )
+
+ # turn the time coordinate into a dimension, allows for concatenation when opening multiple datasets
+ full_dataset = full_dataset.expand_dims({"time": np.atleast_1d(timestep)})
+ print(" (%.1f seconds)" % (time.time() - start_time))
+
+ # folder containing netcdf data for the month of the dataset
+ monthly_dir = "%s/%d%02d" % (args["netcdf_outdir"], year, month)
+ if not os.path.isdir(monthly_dir):
+ os.mkdir(monthly_dir)
+
+ # save out netcdf files, one for each forecast hour
+ for idx, forecast_hour in enumerate(forecast_hours):
+ filepath = f"%s/%s_%d%02d%02d%02d_f%03d.nc" % (
+ monthly_dir,
+ args["model"],
+ year,
+ month,
+ day,
+ hour,
+ forecast_hour,
+ )
+ print("Saving: %s" % filepath)
+ full_dataset.isel(
+ forecast_hour=[
+ idx,
+ ]
+ ).to_netcdf(path=filepath, mode="w", engine="netcdf4")
diff --git a/scripts/make_dryrun_data.py b/scripts/make_dryrun_data.py
new file mode 100644
index 0000000..e604ae2
--- /dev/null
+++ b/scripts/make_dryrun_data.py
@@ -0,0 +1,97 @@
+"""Generate a tiny fake TF dataset fixture for local dry-run testing.
+
+Creates two subdirectories under tests/fixtures/dryrun_tf_dataset/:
+
+ 2000-1_tf/ -- used as the "train" split
+ 2001-1_tf/ -- used as the "val" split
+
+Each snapshot contains a small number of random batches whose element shapes
+and dtypes match the real on-cluster dataset:
+
+ inputs: (128, 288, 7, 9) float16
+ targets: (128, 288, 6) float16
+
+Run once from the repo root before doing a local dry run:
+
+ python scripts/make_dryrun_data.py
+
+The generated files are checked in to .gitignore (see tests/fixtures/.gitignore)
+so they are not committed to the repo.
+"""
+
+import argparse
+import os
+
+import numpy as np
+import tensorflow as tf
+
+# Element shapes matching the real dataset (from cluster inspection).
+INPUT_SHAPE = (128, 288, 7, 9)
+TARGET_SHAPE = (128, 288, 6)
+DTYPE = tf.float16
+
+# How many elements (batches) to write per snapshot — small enough to be fast.
+NUM_ELEMENTS = 4
+
+FIXTURE_ROOT = os.path.join(
+ os.path.dirname(__file__), "..", "tests", "fixtures", "dryrun_tf_dataset"
+)
+
+SPLITS = {
+ "train": "2000-1_tf",
+ "val": "2001-1_tf",
+}
+
+
+def make_snapshot(out_dir: str, num_elements: int, seed: int) -> None:
+ rng = np.random.default_rng(seed)
+
+ def _gen():
+ for _ in range(num_elements):
+ inputs = rng.random(INPUT_SHAPE).astype(np.float16)
+ targets = rng.random(TARGET_SHAPE).astype(np.float16)
+ yield tf.constant(inputs, dtype=DTYPE), tf.constant(targets, dtype=DTYPE)
+
+ ds = tf.data.Dataset.from_generator(
+ _gen,
+ output_signature=(
+ tf.TensorSpec(shape=INPUT_SHAPE, dtype=DTYPE),
+ tf.TensorSpec(shape=TARGET_SHAPE, dtype=DTYPE),
+ ),
+ )
+ os.makedirs(out_dir, exist_ok=True)
+ ds.save(out_dir)
+ print(f" Saved {num_elements} element(s) → {out_dir}")
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--num_elements",
+ type=int,
+ default=NUM_ELEMENTS,
+ help=f"Number of elements per snapshot (default: {NUM_ELEMENTS})",
+ )
+ parser.add_argument(
+ "--out_dir",
+ type=str,
+ default=FIXTURE_ROOT,
+ help="Root directory to write snapshots into",
+ )
+ args = parser.parse_args()
+
+ print(f"Writing dry-run fixtures to: {os.path.abspath(args.out_dir)}")
+ for i, (split, subdir) in enumerate(SPLITS.items()):
+ out = os.path.join(args.out_dir, subdir)
+ print(f" [{i+1}/{len(SPLITS)}] {split}: {subdir}")
+ make_snapshot(out, num_elements=args.num_elements, seed=i)
+
+ print("Done. Run the local dry-run with:")
+ print(
+ " python -m fronts.train "
+ "--train_config_path configs/1702_tf_dryrun.yaml --dry_run"
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/merge_goes_satellite.py b/scripts/merge_goes_satellite.py
new file mode 100644
index 0000000..88d6724
--- /dev/null
+++ b/scripts/merge_goes_satellite.py
@@ -0,0 +1,204 @@
+"""
+Script that transforms raw GOES data onto a grid covering most of NOAA's Unified Surface Analysis domain.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.2.1
+"""
+
+import argparse
+import datetime as dt
+import sys
+import numpy as np
+import os
+import pandas as pd
+import scipy
+from fronts.utils.satellite import calculate_lat_lon_from_dataset
+import xarray as xr
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--satellite_indir",
+ type=str,
+ required=True,
+ help="Parent input directory for the raw satellite data.",
+ )
+ parser.add_argument(
+ "--init_time",
+ type=str,
+ required=True,
+ help="Initialization time of the model. Format: YYYY-MM-DD-HH.",
+ )
+ parser.add_argument(
+ "--netcdf_outdir",
+ type=str,
+ required=True,
+ help="Output directory for the processed satellite netCDF files.",
+ )
+ parser.add_argument(
+ "--band_nums",
+ type=int,
+ nargs="+",
+ help="Band numbers to include in the final datasets. If this argument is not passed, all band numbers (1-16) will be included.",
+ )
+ parser.add_argument(
+ "--multiprocessing",
+ action="store_true",
+ help="Interpolate the satellite data with one band per CPU thread.",
+ )
+ parser.add_argument(
+ "--overwrite", action="store_true", help="Overwrite existing datasets."
+ )
+ args = vars(parser.parse_args())
+
+ init_time = pd.date_range(args["init_time"], args["init_time"])[0]
+ year, month, day, hour = (
+ init_time.year,
+ init_time.month,
+ init_time.day,
+ init_time.hour,
+ )
+
+ # check if a merged dataset for the given timestamp already exists
+ converted_ds_filepath = "%s/%d%02d/goes-merged_%d%02d%02d%02d_full.nc" % (
+ args["netcdf_outdir"],
+ year,
+ month,
+ year,
+ month,
+ day,
+ hour,
+ )
+ converted_ds_exists = os.path.isfile(converted_ds_filepath)
+ if converted_ds_exists and not args["overwrite"]:
+ print(
+ "%s already exists. If you want to overwrite the existing dataset, rerun this script with the --overwrite flag attached.\n"
+ "Exiting." % converted_ds_filepath
+ )
+ sys.exit(0)
+
+ lon_array = np.arange(-216, 5.251, 0.25)
+ lon_array_sat1 = np.arange(-155, 5.251, 0.25)
+ lon_array_sat2 = np.arange(-216, -58.1, 0.25)
+ lat_array = np.arange(0, 70, 0.25)
+
+ if year <= 2017:
+ sat1, sat2 = "goes13", "goes15"
+ else:
+ sat1, sat2 = "goes16", "goes17"
+
+ print(f"[{dt.datetime.utcnow()}]", "Opening datasets")
+ try:
+ ds_sat1 = xr.open_dataset(
+ "%s/%d%02d/%s_%d%02d%02d%02d_full-disk.nc"
+ % (args["satellite_indir"], year, month, sat1, year, month, day, hour),
+ engine="netcdf4",
+ )
+ ds_sat2 = xr.open_dataset(
+ "%s/%d%02d/%s_%d%02d%02d%02d_full-disk.nc"
+ % (args["satellite_indir"], year, month, sat2, year, month, day, hour),
+ engine="netcdf4",
+ )
+ except:
+ print("Error loading datasets, exiting.")
+ sys.exit(0)
+
+ if sat1 == "goes16":
+ print(f"[{dt.datetime.utcnow()}]", f"Converting {sat1} coordinates")
+ lat_sat1, lon_sat1 = calculate_lat_lon_from_dataset(ds_sat1)
+ lat_sat1, lon_sat1 = np.nan_to_num(lat_sat1), np.nan_to_num(lon_sat1)
+
+ print(f"[{dt.datetime.utcnow()}]", f"Converting {sat2} coordinates")
+ lat_sat2, lon_sat2 = calculate_lat_lon_from_dataset(ds_sat2)
+ lat_sat2, lon_sat2 = np.nan_to_num(lat_sat2), np.nan_to_num(lon_sat2)
+
+ else:
+ ds_sat1 = ds_sat1.isel(time=0)
+ ds_sat2 = ds_sat2.isel(time=0)
+
+ lat_sat1, lon_sat1 = ds_sat1["lat"].values, ds_sat1["lon"].values
+ lat_sat2, lon_sat2 = ds_sat2["lat"].values, ds_sat2["lon"].values
+
+ new_lat_sat1, new_lon_sat1 = np.meshgrid(lat_array, lon_array_sat1)
+ new_lat_sat2, new_lon_sat2 = np.meshgrid(lat_array, lon_array_sat2)
+
+ band_nums = args["band_nums"] if args["band_nums"] is not None else np.arange(1, 17)
+
+ # xarray dataset that will contain the merged satellite data
+ ds_merged = xr.Dataset(coords={"longitude": lon_array + 360, "latitude": lat_array})
+
+ for band_num in band_nums:
+ if sat1 == "goes13":
+ band_str = "ch%d" % int(band_num)
+ else:
+ band_str = "CMI_C%02d" % int(band_num)
+
+ variable_sat1 = np.nan_to_num(ds_sat1[band_str].values)
+ variable_sat2 = np.nan_to_num(ds_sat2[band_str].values)
+
+ print(f"[{dt.datetime.utcnow()}]", "Interpolating %s" % band_str)
+
+ if sat1 == "goes13":
+ variable_sat1 = scipy.interpolate.RegularGridInterpolator(
+ (lat_sat1, lon_sat1), variable_sat1, method="nearest"
+ )((new_lat_sat1, new_lon_sat1))
+ variable_sat2 = scipy.interpolate.RegularGridInterpolator(
+ (lat_sat2, lon_sat2), variable_sat2, method="nearest"
+ )((new_lat_sat2, new_lon_sat2))
+ else:
+ variable_sat1 = scipy.interpolate.griddata(
+ (lat_sat1.ravel(), lon_sat1.ravel()),
+ variable_sat1.ravel(),
+ (new_lat_sat1, new_lon_sat1),
+ method="nearest",
+ )
+ variable_sat2 = scipy.interpolate.griddata(
+ (lat_sat2.ravel(), lon_sat2.ravel()),
+ variable_sat2.ravel(),
+ (new_lat_sat2, new_lon_sat2),
+ method="nearest",
+ )
+
+ overlap_west_bound = lon_array_sat1[0]
+ overlap_east_bound = lon_array_sat2[-1]
+
+ overlap_east_bound_ind = np.where(lon_array_sat1 == overlap_east_bound)[0][0]
+ overlap_west_bound_ind = np.where(lon_array_sat2 == overlap_west_bound)[0][0]
+
+ # blend overlapping portions of images together
+ overlap_sat1 = variable_sat1[: overlap_east_bound_ind + 1]
+ overlap_sat2 = variable_sat2[overlap_west_bound_ind:]
+ overlap_mask = np.linspace(0, 1, overlap_east_bound_ind + 1)[:, np.newaxis]
+ overlap_blend = (overlap_sat1 * overlap_mask) + (
+ overlap_sat2 * (1 - overlap_mask)
+ )
+
+ merged_data = np.vstack(
+ [
+ variable_sat2[:overlap_west_bound_ind, :],
+ overlap_blend,
+ variable_sat1[overlap_east_bound_ind + 1 :],
+ ]
+ )
+
+ ds_merged["band_%d" % band_num] = (
+ ("latitude", "longitude"),
+ merged_data.astype(np.float32).transpose(),
+ )
+ ds_merged["band_%d" % band_num].attrs = ds_sat1[band_str].attrs
+
+ os.makedirs(
+ "%s/%d%02d" % (args["netcdf_outdir"], year, month), exist_ok=True
+ ) # directory check
+
+ print(f"[{dt.datetime.utcnow()}]", "Saving dataset to %s" % converted_ds_filepath)
+ ds_merged = ds_merged.expand_dims(
+ {"time": np.atleast_1d(init_time).astype("datetime64[ns]")}
+ )
+ ds_merged = ds_merged.reindex(
+ latitude=ds_merged["latitude"].values[::-1]
+ ) # reverse latitude values so they are ordered north-south
+ ds_merged.to_netcdf(
+ converted_ds_filepath, engine="netcdf4", mode="w"
+ ) # save dataset
diff --git a/scripts/model_speed_test.py b/scripts/model_speed_test.py
new file mode 100644
index 0000000..a883930
--- /dev/null
+++ b/scripts/model_speed_test.py
@@ -0,0 +1,89 @@
+"""
+Script for testing model prediction speed.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2024.11.15
+"""
+
+import argparse
+from fronts.utils import file_manager as fm
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+import time
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--model_dir", type=str, required=True, help="Parent model directory."
+ )
+ parser.add_argument("--model_number", type=int, required=True, help="Model number.")
+ parser.add_argument("--gpu_device", type=int, nargs="+", help="GPU device numbers.")
+ parser.add_argument(
+ "--memory_growth", action="store_true", help="Use memory growth for GPUs"
+ )
+ args = vars(parser.parse_args())
+
+ if args["gpu_device"] is not None:
+ gpus = tf.config.list_physical_devices(device_type="GPU") # Find available GPUs
+
+ if len(gpus) > 0:
+ print("Number of GPUs available: %d" % len(gpus))
+
+ # Only make the selected GPU(s) visible to TensorFlow
+ tf.config.set_visible_devices(
+ devices=[gpus[gpu] for gpu in args["gpu_device"]], device_type="GPU"
+ )
+ gpus = tf.config.get_visible_devices(
+ device_type="GPU"
+ ) # List of selected GPUs
+ print("Using %d GPU(s):" % len(gpus), gpus)
+
+ # Allow for memory growth on the GPU. This will only use the GPU memory that is required rather than allocating all the GPU's memory.
+ if args["memory_growth"]:
+ tf.config.experimental.set_memory_growth(
+ device=[gpu for gpu in gpus][0], enable=True
+ )
+
+ else:
+ print("WARNING: No GPUs found, all computations will be performed on CPUs.")
+ tf.config.set_visible_devices([], "GPU")
+
+ model = fm.load_model(args["model_number"], args["model_dir"])
+ model_properties = pd.read_pickle(
+ "%s/model_%d/model_%d_properties.pkl"
+ % (args["model_dir"], args["model_number"], args["model_number"])
+ )
+
+ N_runs = 30
+ batch_size = 2
+ N_pressure_levels = len(model_properties["dataset_properties"]["pressure_levels"])
+ N_variables = len(model_properties["dataset_properties"]["variables"])
+ image_size = [960, 320]
+
+ time_elapsed_arr = []
+
+ print("Initializing model")
+ init_input_tensor = tf.convert_to_tensor(
+ np.random.rand(batch_size, *image_size, N_pressure_levels, N_variables)
+ )
+ model.predict(init_input_tensor, verbose=0)
+
+ print("Starting benchmark...")
+ for run in range(1, N_runs + 1):
+ input_tensor = tf.convert_to_tensor(
+ np.random.rand(batch_size, *image_size, N_pressure_levels, N_variables)
+ )
+ start_time = time.time()
+ model.predict(input_tensor, verbose=0)
+ time_elapsed = time.time() - start_time
+
+ print("Run %d/%d: %.3fs" % (run, N_runs, time_elapsed))
+ time_elapsed_arr.append(time_elapsed)
+
+ avg_time_elapsed = np.mean(np.array(time_elapsed_arr))
+ print(
+ "Final benchmark: %.3fs, %.3fs/image"
+ % (avg_time_elapsed, avg_time_elapsed / batch_size)
+ )
diff --git a/src/fronts/__init__.py b/src/fronts/__init__.py
new file mode 100644
index 0000000..af2efac
--- /dev/null
+++ b/src/fronts/__init__.py
@@ -0,0 +1,16 @@
+from . import data
+from . import evaluation
+from . import model
+from . import nfa
+from . import plot
+from . import utils
+
+
+__all__ = [
+ "data",
+ "evaluation",
+ "model",
+ "nfa",
+ "plot",
+ "utils",
+]
diff --git a/evaluation/__init__.py b/src/fronts/data/__init__.py
similarity index 100%
rename from evaluation/__init__.py
rename to src/fronts/data/__init__.py
diff --git a/src/fronts/data/batch.py b/src/fronts/data/batch.py
new file mode 100644
index 0000000..7ab76df
--- /dev/null
+++ b/src/fronts/data/batch.py
@@ -0,0 +1,114 @@
+import tensorflow as tf
+import xarray as xr
+
+import xbatcher as xb
+import xbatcher.loaders.keras
+from typing import Optional
+import dataclasses
+
+
+def create_dataloader(
+ inputs: xr.Dataset,
+ targets: xr.Dataset,
+ input_sizes: Optional[dict[str, int]] = None,
+ target_sizes: Optional[dict[str, int]] = None,
+ prefetch_number: int = 3,
+ preload_batch: bool = False,
+ input_dtype: type = tf.float32,
+ target_dtype: type = tf.float32,
+) -> tf.data.Dataset:
+ """Create a tf.data.Dataset DataLoader from xarray Datasets for inputs and targets.
+
+ References https://xbatcher.readthedocs.io/en/latest/user-guide/training-a-neural-network-with-keras-and-xbatcher.html
+ for BatchGenerator usage. This function is used primarily to take cloud-based data
+ and create a DataLoader that can be used for training a model in TensorFlow/Keras.
+
+ Args:
+ inputs: An xarray Dataset containing the input features.
+ targets: An xarray Dataset containing the target labels.
+ input_sizes: Optional dict specifying the dims and sizes of the input batches.
+ If not provided, it will be inferred from the inputs dataset.
+ target_sizes: Optional dict specifying the dims and sizes of the target batches.
+ If not provided, it will be inferred from the targets dataset.
+ prefetch_number: The number of batches to prefetch for. Defaults to 3 for
+ what should be optimal performance.
+ preload_batch: Whether to preload batches into memory. Defaults to False.
+ input_dtype: The data type for the input batches. Defaults to tf.float32.
+ target_dtype: The data type for the target batches. Defaults to tf.float32.
+
+ Returns a tf.data.Dataset that yields batches of (inputs, targets) for training a
+ model. Each batch will have the specified input_shape and target_shape.
+ """
+ if input_sizes is None:
+ input_sizes = dict(inputs.sizes)
+ if target_sizes is None:
+ target_sizes = dict(targets.sizes)
+ # Define batch generators for features (X) and labels (y)
+ X_bgen = xb.BatchGenerator(
+ inputs,
+ input_dims=input_sizes,
+ preload_batch=preload_batch, # Load each batch dynamically
+ )
+ y_bgen = xb.BatchGenerator(
+ targets, input_dims=target_sizes, preload_batch=preload_batch
+ )
+
+ # Use xbatcher's MapDataset to wrap the generators
+ batch_dataset = xbatcher.loaders.keras.CustomTFDataset(X_bgen, y_bgen)
+
+ # Create a DataLoader using tf.data.Dataset
+ train_dataloader = tf.data.Dataset.from_generator(
+ lambda: iter(batch_dataset),
+ output_signature=(
+ tf.TensorSpec(
+ shape=tuple(input_sizes.values()), dtype=input_dtype, name="inputs"
+ ), # inputs
+ tf.TensorSpec(
+ shape=tuple(target_sizes.values()), dtype=target_dtype, name="targets"
+ ), # targets
+ ),
+ ).prefetch(prefetch_number) # Prefetch 3 batches to improve performance
+
+ return train_dataloader
+
+
+@dataclasses.dataclass
+class BatchGeneratorConfig:
+ """A dataclass for configuring the creation of a batch generator DataLoader.
+
+ Attributes:
+ input_sizes: Optional dict specifying the dims and sizes of the input batches.
+ If not provided, it will be inferred from the inputs dataset.
+ target_sizes: Optional dict specifying the dims and sizes of the target batches.
+ If not provided, it will be inferred from the targets dataset.
+ prefetch_number: The number of batches to prefetch for. Defaults to 3 for
+ what should be optimal performance.
+ preload_batch: Whether to preload batches into memory. Defaults to False.
+ input_dtype: The data type for the input batches. Defaults to tf.float32.
+ target_dtype: The data type for the target batches. Defaults to tf.float32.
+ """
+
+ input_sizes: Optional[dict[str, int]] = None
+ target_sizes: Optional[dict[str, int]] = None
+ prefetch_number: int = 3
+ preload_batch: bool = False
+ input_dtype = tf.float32
+ target_dtype = tf.float32
+
+ def build(self) -> tf.data.Dataset:
+ """Builds a tf.data.Dataset DataLoader using the provided configuration
+ parameters.
+
+ Returns a tf.data.Dataset that yields batches of (inputs, targets) for training
+ a model. Each batch will have the specified input_shape and target_shape.
+ """
+ return create_dataloader(
+ inputs=self.inputs,
+ targets=self.targets,
+ input_sizes=self.input_sizes,
+ target_sizes=self.target_sizes,
+ prefetch_number=self.prefetch_number,
+ preload_batch=self.preload_batch,
+ input_dtype=self.input_dtype,
+ target_dtype=self.target_dtype,
+ )
diff --git a/src/fronts/data/config.py b/src/fronts/data/config.py
new file mode 100644
index 0000000..1dc9079
--- /dev/null
+++ b/src/fronts/data/config.py
@@ -0,0 +1,689 @@
+"""Data configuration dataclasses for the FrontFinder training pipeline.
+
+This module provides composable config dataclasses for loading ERA5 predictor data
+and front label (truth) data, splitting by year, stacking surface and pressure-level
+variables in xarray, and building tf.data.Dataset objects for train/val/test splits.
+
+Follows the same pattern as the rest of the codebase: typed dataclasses with .build()
+methods that return runtime objects, loadable from YAML via dacite.
+"""
+
+import dataclasses
+import datetime
+import glob as glob_module
+import logging
+import os
+from typing import Any, Optional, Union
+
+import tensorflow as tf
+import xarray as xr
+
+from fronts.data.batch import BatchGeneratorConfig, create_dataloader
+from fronts.data.era5 import convert_domain_extent_to_bounding_box
+from fronts.utils import data_utils
+
+log = logging.getLogger("fronts.data.config")
+
+
+# ---------------------------------------------------------------------------
+# Constant mapping: pressure-level variable name -> surface variable name
+#
+# When "surface" appears in the levels list, the code looks up each requested
+# variable name here to find its surface counterpart in the zarr store.
+# Variables absent from this map are pressure-level-only (no surface analogue).
+# ---------------------------------------------------------------------------
+
+SURFACE_VARIABLE_MAP: dict[str, str] = {
+ "temperature": "2m_temperature",
+ "u_component_of_wind": "10m_u_component_of_wind",
+ "v_component_of_wind": "10m_v_component_of_wind",
+ "dewpoint_temperature": "2m_dewpoint_temperature",
+ "specific_humidity": "surface_specific_humidity",
+}
+
+# Variables that exist only at the surface (no pressure-level equivalent).
+# These are included in the output whenever "surface" is in the levels list.
+SURFACE_ONLY_VARIABLES: set[str] = {
+ "mean_sea_level_pressure",
+ "total_precipitation",
+ "sea_surface_temperature",
+ "skin_temperature",
+ "10m_wind_speed",
+}
+
+
+def _stack_era5_variables(
+ ds: xr.Dataset,
+ variables: list[str],
+ levels: list[Union[str, int]],
+) -> xr.Dataset:
+ """Stacks ERA5 variables into a unified Dataset with a mixed level coordinate.
+
+ The ``levels`` list may contain the string ``"surface"`` and/or integer hPa
+ values (e.g. ``["surface", 1000, 950, 900, 850]``). The function handles
+ three categories of variable automatically:
+
+ * **Pressure-level-only** — the variable exists only on pressure levels in
+ the zarr store (e.g. ``"specific_humidity"``). These are selected at the
+ requested integer levels.
+ * **Mixed surface + pressure** — the variable has both a surface counterpart
+ (looked up via :data:`SURFACE_VARIABLE_MAP`) and pressure-level data.
+ When ``"surface"`` is in ``levels`` the surface array is prepended; the
+ result has a level coordinate of the form ``["surface", 1000, 950, ...]``.
+ * **Surface-only** — the variable name appears in :data:`SURFACE_ONLY_VARIABLES`
+ *or* is not found as a pressure-level variable in the store. It is
+ included with ``level=["surface"]`` whenever ``"surface"`` is in ``levels``.
+
+ Args:
+ ds: An xarray Dataset already subsetted spatially and temporally.
+ variables: Canonical variable names to include. Use pressure-level names
+ (e.g. ``"temperature"``) for mixed/pressure variables; use the full
+ surface name (e.g. ``"mean_sea_level_pressure"``) for surface-only ones.
+ levels: Ordered list of levels to select. May include the string
+ ``"surface"`` and/or integer hPa values.
+
+ Returns an xarray Dataset with a unified ``"level"`` coordinate whose values
+ are a mix of the string ``"surface"`` and integer hPa values.
+ """
+ include_surface = "surface" in levels
+ pressure_levels = [lv for lv in levels if lv != "surface"]
+
+ result_datasets: list[xr.Dataset] = []
+
+ for var in variables:
+ surface_var_name = SURFACE_VARIABLE_MAP.get(var)
+ is_surface_only = var in SURFACE_ONLY_VARIABLES
+
+ if is_surface_only:
+ # Surface-only variable: always has level=["surface"]
+ if include_surface:
+ da_sfc = ds[var].expand_dims({"level": ["surface"]})
+ result_datasets.append(da_sfc.to_dataset(name=var))
+ elif surface_var_name is not None:
+ # Mixed variable: has a surface counterpart + pressure levels
+ if pressure_levels:
+ da_pl = ds[var].sel(level=pressure_levels)
+ else:
+ da_pl = None
+
+ if include_surface and surface_var_name in ds:
+ da_sfc = ds[surface_var_name].expand_dims({"level": ["surface"]})
+ if da_pl is not None:
+ da = xr.concat([da_sfc, da_pl], dim="level")
+ else:
+ da = da_sfc
+ else:
+ if da_pl is not None:
+ da = da_pl
+ else:
+ continue # nothing to add
+
+ result_datasets.append(da.to_dataset(name=var))
+ else:
+ # Pressure-level-only variable
+ if pressure_levels:
+ da_pl = ds[var].sel(level=pressure_levels)
+ result_datasets.append(da_pl.to_dataset(name=var))
+
+ return xr.merge(result_datasets, join="outer")
+
+
+@dataclasses.dataclass
+class ERA5PredictorConfig:
+ """Configuration for loading and stacking ERA5 predictor variables.
+
+ Variables are specified as a single ``variables`` list using canonical
+ (pressure-level) names. The ``levels`` list controls which vertical levels
+ are loaded and may contain the string ``"surface"`` in addition to integer
+ hPa values.
+
+ When ``"surface"`` appears in ``levels``, the module-level
+ :data:`SURFACE_VARIABLE_MAP` is consulted to find each variable's surface
+ counterpart (e.g. ``"temperature"`` → ``"2m_temperature"``). Variables
+ listed in :data:`SURFACE_ONLY_VARIABLES` (e.g. ``"mean_sea_level_pressure"``)
+ are included with ``level=["surface"]`` automatically.
+
+ The resulting xarray Dataset has a unified ``"level"`` coordinate whose
+ values are a mix of the string ``"surface"`` and integer hPa values,
+ following the convention used throughout the codebase.
+
+ Attributes:
+ domain_extent: [lon_min, lon_max, lat_min, lat_max] geographic extent.
+ variables: Canonical variable names to load. For variables with both
+ surface and pressure-level representations, use the pressure-level
+ name (e.g. ``"temperature"``); for surface-only variables, use the
+ store name directly (e.g. ``"mean_sea_level_pressure"``).
+ levels: Ordered list of levels to include. May contain ``"surface"``
+ and/or integer hPa values, e.g. ``["surface", 1000, 950, 900, 850]``
+ or ``[1000, 900, 750]``.
+ years: Years to select data from. Typically injected by DataConfig.build()
+ via dataclasses.replace() rather than set directly in YAML.
+ store: URI of the zarr store to open.
+ chunks: Chunk sizes for lazy loading, e.g. {"time": 48}.
+ consolidated: Whether to use consolidated zarr metadata.
+ """
+
+ domain_extent: list[float]
+ variables: list[str]
+ levels: list[Union[str, int]]
+ store: str
+ chunks: dict[str, int]
+ consolidated: bool
+ years: list[int] = dataclasses.field(default_factory=list)
+
+ def build(self) -> xr.Dataset:
+ """Loads and stacks ERA5 data into a unified xarray Dataset.
+
+ Returns an xarray Dataset with a ``"level"`` coordinate that includes
+ ``"surface"`` (for surface variables) and integer hPa values (for
+ pressure-level variables). Time is filtered to ``self.years``.
+ """
+ log.info(
+ "ERA5PredictorConfig.build() — opening zarr store: %s", self.store
+ )
+ ds = xr.open_zarr(
+ store=self.store,
+ chunks=self.chunks,
+ consolidated=self.consolidated,
+ )
+ log.debug("Zarr store opened. Variables available: %s", list(ds.data_vars))
+
+ # Spatial subset
+ log.debug("Applying spatial subset: domain_extent=%s", self.domain_extent)
+ bbox = convert_domain_extent_to_bounding_box(self.domain_extent)
+ ds = ds.sel(
+ latitude=slice(bbox.lat_max, bbox.lat_min),
+ longitude=slice(bbox.lon_min, bbox.lon_max),
+ )
+ log.debug(
+ "Spatial subset done. lat shape=%s, lon shape=%s",
+ ds.latitude.shape, ds.longitude.shape,
+ )
+
+ # Temporal subset: keep only the requested years
+ log.debug("Applying temporal subset for years=%s...", self.years)
+ ds = ds.isel(time=ds.time.dt.year.isin(self.years))
+ log.info("ERA5 temporal subset done. %d timesteps selected.", ds.sizes.get("time", 0))
+
+ log.debug("Stacking variables=%s at levels=%s...", self.variables, self.levels)
+ result = _stack_era5_variables(
+ ds,
+ variables=self.variables,
+ levels=self.levels,
+ )
+ log.info(
+ "ERA5PredictorConfig.build() complete. Output vars: %s", list(result.data_vars)
+ )
+ return result
+
+
+@dataclasses.dataclass
+class FrontsDataConfig:
+ """Configuration for loading front label (truth) data from netCDF files.
+
+ Front files are expected to contain an "identifier" variable with integer class
+ values (0=no front, 1=CF, 2=WF, etc.), following the existing convention in the
+ codebase.
+
+ Attributes:
+ directory: Path to directory containing per-timestep front netCDF files.
+ Files are matched using the glob pattern `{directory}/*{year}*.nc`.
+ years: Years to load. Typically injected by DataConfig.build() via
+ dataclasses.replace() rather than set directly in YAML.
+ front_types: Code(s) passed to reformat_fronts() to regroup front classes.
+ Examples: "MERGED-ALL", "F_BIN", ["CF", "WF"]. None = no reformatting.
+ """
+
+ directory: str
+ front_types: Optional[Any] # str | list[str] | None
+ years: list[int] = dataclasses.field(default_factory=list)
+
+ def build(self) -> xr.Dataset:
+ """Loads and optionally reformats front label data for the given years.
+
+ Supports two directory layouts automatically:
+ - Flat: ``{directory}/*{year}*.nc``
+ - Monthly subdirs: ``{directory}/{year}MM/*.nc`` (e.g. ``200701/``)
+
+ Returns an xarray Dataset with the "identifier" variable, optionally
+ reformatted according to self.front_types.
+ """
+ log.info(
+ "FrontsDataConfig.build() — globbing files for years=%s in %r...",
+ self.years, self.directory,
+ )
+ files = sorted(
+ f
+ for year in self.years
+ for f in (
+ # Monthly subdirectory layout: //*.nc
+ glob_module.glob(f"{self.directory}/{year}*/*.nc")
+ or
+ # Flat layout fallback: /**.nc
+ glob_module.glob(f"{self.directory}/*{year}*.nc")
+ )
+ )
+ log.info("FrontsDataConfig — found %d file(s). Opening with open_mfdataset...", len(files))
+ ds = xr.open_mfdataset(files, engine="netcdf4", combine="by_coords", coords="minimal", compat="override")
+ log.info("FrontsDataConfig — dataset opened. Variables: %s", list(ds.data_vars))
+
+ if self.front_types is not None:
+ log.debug("Reformatting fronts with front_types=%s...", self.front_types)
+ ds = data_utils.reformat_fronts(ds, self.front_types)
+ log.debug("reformat_fronts complete.")
+
+ log.info("FrontsDataConfig.build() complete.")
+ return ds
+
+
+@dataclasses.dataclass
+class TFDatasetConfig:
+ """Configuration for loading pre-built tf.data.Dataset snapshots from disk.
+
+ The on-disk datasets were produced by the previous pipeline and are stored
+ as saved ``tf.data.Dataset`` snapshots in year-labelled subdirectories under
+ a common root, e.g.::
+
+ /
+ 2010-1_tf/
+ 2010-2_tf/
+ ...
+ 2020-12_tf/
+
+ All monthly subdirectories whose name starts with a requested year are
+ concatenated in sorted order to form the split dataset.
+
+ This is the fastest path to training — it bypasses ERA5 zarr loading,
+ front netCDF loading, stacking, and normalization entirely, using data that
+ is already preprocessed and on local disk.
+
+ Attributes:
+ directory: Root directory containing the year-month subdirectories.
+ train_years: Years to include in the training split.
+ val_years: Years to include in the validation split.
+ test_years: Years to include in the test split. May be empty [].
+ shuffle: Whether to shuffle the training dataset. Defaults to True.
+ shuffle_buffer: Buffer size passed to ``tf.data.Dataset.shuffle()``.
+ Defaults to 1000.
+ prefetch: Number of batches to prefetch. Defaults to 3.
+ """
+
+ directory: str
+ train_years: list[int]
+ val_years: list[int]
+ test_years: list[int]
+ shuffle: bool = True
+ shuffle_buffer: int = 1000
+ prefetch: int = 3
+
+ def _load_years(self, years: list[int]) -> Optional[Any]:
+ """Loads and concatenates all monthly TF dataset snapshots for ``years``.
+
+ Subdirectories are matched by prefix: a directory named ``2010-3_tf``
+ matches year ``2010``.
+
+ Returns a ``tf.data.Dataset``, or ``None`` if ``years`` is empty or no
+ matching subdirectories are found.
+ """
+ if not years:
+ log.debug("_load_years called with empty years list — returning None.")
+ return None
+
+ log.debug("Scanning %r for subdirs matching years %s...", self.directory, years)
+ subdirs = sorted(
+ os.path.join(self.directory, d)
+ for d in os.listdir(self.directory)
+ if any(d.startswith(str(y)) for y in years)
+ and os.path.isdir(os.path.join(self.directory, d))
+ )
+
+ if not subdirs:
+ raise FileNotFoundError(
+ f"No subdirectories found in {self.directory!r} matching years "
+ f"{years}. Expected names like '2010-1_tf', '2010-2_tf', etc."
+ )
+
+ log.info("Loading %d TF dataset snapshot(s) for years %s...", len(subdirs), years)
+ datasets = []
+ for i, s in enumerate(subdirs):
+ log.debug(" [%d/%d] Loading %s", i + 1, len(subdirs), s)
+ datasets.append(tf.data.Dataset.load(s))
+ log.debug("All snapshots loaded. Concatenating...")
+
+ combined = datasets[0]
+ for ds in datasets[1:]:
+ combined = combined.concatenate(ds)
+ log.debug("Concatenation complete. Applying prefetch=%d.", self.prefetch)
+ return combined.prefetch(self.prefetch)
+
+ def build(self) -> "ModelData":
+ """Builds train, validation, and test ``tf.data.Dataset`` objects.
+
+ Returns a :class:`ModelData` with ``train_data``, ``validation_data``,
+ and optionally ``test_data``.
+ """
+ log.info("TFDatasetConfig.build() — loading train split (years=%s)...", self.train_years)
+ train_ds = self._load_years(self.train_years)
+ if self.shuffle and train_ds is not None:
+ log.debug("Shuffling train dataset (buffer_size=%d).", self.shuffle_buffer)
+ train_ds = train_ds.shuffle(buffer_size=self.shuffle_buffer)
+
+ log.info("TFDatasetConfig.build() — loading val split (years=%s)...", self.val_years)
+ val_ds = self._load_years(self.val_years)
+
+ if self.test_years:
+ log.info("TFDatasetConfig.build() — loading test split (years=%s)...", self.test_years)
+ test_ds = self._load_years(self.test_years)
+
+ log.info("TFDatasetConfig.build() complete.")
+ return ModelData(
+ train_data=train_ds,
+ validation_data=val_ds,
+ test_data=test_ds,
+ )
+
+
+@dataclasses.dataclass
+class ModelData:
+ """Runtime holder for train/validation/test tf.data.Dataset objects.
+
+ Returned by DataConfig.build(). Trainer accesses .train_data and
+ .validation_data directly (train.py:221,228).
+
+ Attributes:
+ train_data: tf.data.Dataset for training.
+ validation_data: tf.data.Dataset for validation.
+ test_data: Optional tf.data.Dataset for testing. None if test_years is empty.
+ """
+
+ train_data: Any
+ validation_data: Any
+ test_data: Optional[Any] = None
+
+
+@dataclasses.dataclass
+class DataConfig:
+ """Top-level data configuration for the FrontFinder training pipeline.
+
+ Supports two mutually exclusive data sources:
+
+ 1. **Pre-built TF datasets** (``tf_dataset`` key) — fastest path, loads
+ saved ``tf.data.Dataset`` snapshots directly from disk. Set
+ ``tf_dataset:`` in YAML and leave ``era5``, ``fronts``, ``batch`` unset.
+
+ 2. **ARCO ERA5 + front netCDF** (``era5`` + ``fronts`` + ``batch`` keys) —
+ full pipeline that loads from the zarr store and front label files,
+ stacks variables, and builds batches via xbatcher.
+
+ Year lists (``train_years``, ``val_years``, ``test_years``) are always
+ specified at this level. For the ERA5 path they are injected into
+ ``ERA5PredictorConfig`` and ``FrontsDataConfig`` at build time via
+ ``dataclasses.replace()`` — they should NOT be set in the era5/fronts YAML
+ blocks.
+
+ Attributes:
+ train_years: Years to use for the training split.
+ val_years: Years to use for the validation split.
+ test_years: Years to use for the test split. May be empty [].
+ tf_dataset: TFDatasetConfig for loading pre-built TF dataset snapshots.
+ Mutually exclusive with era5/fronts/batch.
+ era5: ERA5PredictorConfig defining the predictor variable source.
+ Required when tf_dataset is not set.
+ fronts: FrontsDataConfig defining the front label source.
+ Required when tf_dataset is not set.
+ batch: BatchGeneratorConfig defining spatial patch sizes and prefetch.
+ Required when tf_dataset is not set.
+ shuffle: Whether to shuffle the training dataset. Defaults to True.
+ Ignored when tf_dataset is set (shuffle is configured there instead).
+ normalization_method: One of "standard", "standard_weighted", "min-max".
+ Defaults to "standard". Only used by the ERA5 path.
+ """
+
+ train_years: list[int]
+ val_years: list[int]
+ test_years: list[int]
+ tf_dataset: Optional[TFDatasetConfig] = None
+ era5: Optional[ERA5PredictorConfig] = None
+ fronts: Optional[FrontsDataConfig] = None
+ batch: Optional[BatchGeneratorConfig] = None
+ shuffle: bool = True
+ normalization_method: str = "standard"
+
+ def build(self) -> ModelData:
+ """Builds train, validation, and test tf.data.Dataset objects.
+
+ Delegates to TFDatasetConfig.build() when tf_dataset is set, otherwise
+ uses the ERA5 + fronts pipeline.
+
+ Returns a ModelData with train_data, validation_data, and optionally test_data.
+ """
+ if self.tf_dataset is not None:
+ log.info(
+ "DataConfig.build() — using TFDatasetConfig path. "
+ "train_years=%s, val_years=%s, test_years=%s",
+ self.train_years, self.val_years, self.test_years,
+ )
+ # Inject year lists into the TFDatasetConfig and build
+ tf_cfg = dataclasses.replace(
+ self.tf_dataset,
+ train_years=self.train_years,
+ val_years=self.val_years,
+ test_years=self.test_years,
+ shuffle=self.shuffle,
+ )
+ return tf_cfg.build()
+
+ # --- ERA5 + fronts path ---
+ log.info(
+ "DataConfig.build() — using ERA5+fronts path. "
+ "train_years=%s, val_years=%s, test_years=%s",
+ self.train_years, self.val_years, self.test_years,
+ )
+ if self.era5 is None or self.fronts is None or self.batch is None:
+ raise ValueError(
+ "DataConfig requires either tf_dataset or all three of "
+ "era5, fronts, and batch to be set."
+ )
+
+ def _build_split(years: list[int]) -> Optional[Any]:
+ if not years:
+ return None
+
+ # Build ERA5 predictor dataset for this split
+ log.info(" Building ERA5 predictor dataset for years=%s...", years)
+ era5_cfg = dataclasses.replace(self.era5, years=years)
+ inputs_ds = era5_cfg.build()
+ log.info(" ERA5 dataset ready.")
+ # TODO: normalize_dataset expects a "pressure_level" dimension and legacy
+ # short variable-name keys (e.g. "T_850", "u_1000"). Our stacked dataset
+ # uses dimension "level" and ARCO variable names ("temperature", etc.).
+ # Normalization constants and the normalize_dataset function need to be
+ # updated for the new naming scheme before this call can be re-enabled.
+ # inputs_ds = data_utils.normalize_dataset(
+ # inputs_ds, method=self.normalization_method
+ # )
+
+ # Build fronts dataset for this split
+ log.info(" Building fronts dataset for years=%s...", years)
+ fronts_cfg = dataclasses.replace(self.fronts, years=years)
+ targets_ds = fronts_cfg.build()
+ log.info(" Fronts dataset ready.")
+
+ # Build tf.data.Dataset via create_dataloader directly
+ # (BatchGeneratorConfig.build() is not used because it lacks inputs/targets fields)
+ log.info(" Wrapping into tf.data.Dataset via create_dataloader...")
+ tf_ds = create_dataloader(
+ inputs=inputs_ds,
+ targets=targets_ds,
+ input_sizes=self.batch.input_sizes,
+ target_sizes=self.batch.target_sizes,
+ prefetch_number=self.batch.prefetch_number,
+ preload_batch=self.batch.preload_batch,
+ )
+ log.info(" tf.data.Dataset ready for years=%s.", years)
+ return tf_ds
+
+ log.info("Building train split...")
+ train_ds = _build_split(self.train_years)
+ if self.shuffle and train_ds is not None:
+ log.debug("Shuffling train dataset.")
+ train_ds = train_ds.shuffle(buffer_size=1000)
+
+ log.info("Building val split...")
+ val_ds = _build_split(self.val_years)
+
+ if self.test_years:
+ log.info("Building test split...")
+ test_ds = _build_split(self.test_years)
+
+ log.info("DataConfig.build() complete.")
+ return ModelData(
+ train_data=train_ds,
+ validation_data=val_ds,
+ test_data=test_ds,
+ )
+
+
+@dataclasses.dataclass
+class TimeSelection:
+ """Specifies which ERA5 timesteps to load for prediction.
+
+ Exactly one of most_recent, timestamps, or date_range must be set.
+ Validation is performed in __post_init__.
+
+ Attributes:
+ most_recent: If True, selects the single latest timestep available in the
+ store. The zarr store's time dimension is assumed to be sorted ascending
+ (latest last), which is true for ARCO ERA5.
+ timestamps: An explicit list of datetimes, each representing a single
+ analysis time (date + hour). Selection uses method="nearest" to tolerate
+ minor floating-point or timezone differences.
+ date_range: A two-element list [start, end] of datetimes. All timesteps
+ between start and end (inclusive) are selected.
+
+ YAML usage — choose exactly one block:
+
+ # Most recent timestep in the store:
+ time_selection:
+ most_recent: true
+
+ # Explicit individual timesteps (use ISO 8601 with T separator):
+ time_selection:
+ timestamps:
+ - "2024-06-01T12:00:00"
+ - "2024-06-02T00:00:00"
+
+ # Inclusive date range:
+ time_selection:
+ date_range:
+ - "2024-06-01T00:00:00"
+ - "2024-06-07T18:00:00"
+ """
+
+ most_recent: bool = False
+ timestamps: Optional[list[datetime.datetime]] = None
+ date_range: Optional[list[datetime.datetime]] = None # exactly [start, end]
+
+ def __post_init__(self):
+ modes_set = sum([
+ bool(self.most_recent),
+ self.timestamps is not None,
+ self.date_range is not None,
+ ])
+ if modes_set != 1:
+ raise ValueError(
+ "Exactly one of most_recent, timestamps, or date_range must be set "
+ f"in TimeSelection (got {modes_set} modes set)."
+ )
+ if self.date_range is not None and len(self.date_range) != 2:
+ raise ValueError(
+ "date_range must be a list of exactly two datetimes [start, end], "
+ f"got {len(self.date_range)} element(s)."
+ )
+
+ def apply(self, ds: xr.Dataset) -> xr.Dataset:
+ """Applies this time selection to a spatially-subsetted xarray Dataset.
+
+ Args:
+ ds: An xarray Dataset that has already been subsetted spatially.
+
+ Returns the Dataset subsetted to the specified timesteps.
+ """
+ if self.most_recent:
+ return ds.isel(time=[-1])
+ elif self.timestamps is not None:
+ return ds.sel(time=self.timestamps, method="nearest")
+ else: # date_range
+ return ds.sel(time=slice(self.date_range[0], self.date_range[1]))
+
+
+@dataclasses.dataclass
+class PredictConfig:
+ """Configuration for ERA5-based model inference.
+
+ Mirrors DataConfig for training but uses TimeSelection instead of year lists
+ for temporal filtering, and returns a plain xr.Dataset rather than a
+ tf.data.Dataset — inference runs one spatial domain at a time rather than
+ batching many patches.
+
+ Attributes:
+ era5: ERA5PredictorConfig defining the variable source, spatial domain,
+ and zarr store. The `years` field on era5 is unused by PredictConfig;
+ time selection is fully controlled by time_selection.
+ time_selection: TimeSelection specifying which timesteps to load. Exactly
+ one of most_recent, timestamps, or date_range must be set.
+ normalization_method: One of "standard", "standard_weighted", "min-max".
+ Should match the normalization used during training. Defaults to
+ "standard".
+ """
+
+ era5: ERA5PredictorConfig
+ time_selection: TimeSelection
+ normalization_method: str = "standard"
+
+ def build(self) -> xr.Dataset:
+ """Loads, stacks, and normalizes ERA5 data for the selected timesteps.
+
+ Opens the zarr store lazily, applies spatial subsetting, time selection,
+ surface/pressure variable stacking, and normalization.
+
+ Returns a normalized xarray Dataset ready for model inference.
+ """
+ log.info("PredictConfig.build() — opening zarr store: %s", self.era5.store)
+ ds = xr.open_zarr(
+ store=self.era5.store,
+ chunks=self.era5.chunks,
+ consolidated=self.era5.consolidated,
+ )
+ log.debug("Zarr store opened.")
+
+ # Spatial subset
+ log.debug("Applying spatial subset...")
+ bbox = convert_domain_extent_to_bounding_box(self.era5.domain_extent)
+ ds = ds.sel(
+ latitude=slice(bbox.lat_max, bbox.lat_min),
+ longitude=slice(bbox.lon_min, bbox.lon_max),
+ )
+
+ # Time selection
+ log.debug("Applying time selection: %s", self.time_selection)
+ ds = self.time_selection.apply(ds)
+ log.info("Time selection done. %d timestep(s) selected.", ds.sizes.get("time", 0))
+
+ # Stack surface and pressure-level variables
+ log.debug("Stacking variables...")
+ stacked = _stack_era5_variables(
+ ds,
+ variables=self.era5.variables,
+ levels=self.era5.levels,
+ )
+
+ log.info("PredictConfig.build() — stacking complete. Output vars: %s", list(stacked.data_vars))
+ # TODO: normalize_dataset expects a "pressure_level" dimension and legacy
+ # short variable-name keys (e.g. "T_850", "u_1000"). Our stacked dataset
+ # uses dimension "level" and ARCO variable names ("temperature", etc.).
+ # Normalization constants and the normalize_dataset function need to be
+ # updated for the new naming scheme before this call can be re-enabled.
+ # return data_utils.normalize_dataset(stacked, method=self.normalization_method)
+ log.info("PredictConfig.build() complete.")
+ return stacked
diff --git a/src/fronts/data/create_era5_netcdf.py b/src/fronts/data/create_era5_netcdf.py
new file mode 100644
index 0000000..5743220
--- /dev/null
+++ b/src/fronts/data/create_era5_netcdf.py
@@ -0,0 +1,209 @@
+"""
+Transform ERA5 pressure level GRIB files into netCDF files.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import datetime as dt
+import glob
+import os
+import numpy as np
+import xarray as xr
+import argparse
+from fronts.utils import variables
+import pandas as pd
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--indir",
+ type=str,
+ required=True,
+ help="Input directory for the ERA5 GRIB files.",
+ )
+ parser.add_argument(
+ "--outdir",
+ type=str,
+ required=True,
+ help="Output directory for the ERA5 netCDF files.",
+ )
+ parser.add_argument(
+ "--date", type=str, required=True, help="Date with format YYYY-MM-DD."
+ )
+ parser.add_argument(
+ "--parallel_threads",
+ type=int,
+ default=4,
+ help="Number of threads to use during computations in parallel.",
+ )
+ args = vars(parser.parse_args())
+
+ yr, mo, dy = args["date"].split("-")
+
+ pl_files = []
+ for var in [
+ "geopotential",
+ "temperature",
+ "specific_humidity",
+ "u_component_of_wind",
+ "v_component_of_wind",
+ ]:
+ pl_files.append(
+ list(sorted(glob.glob("%s/era5_%s*hPa*%s.nc" % (args["indir"], var, yr))))
+ )
+
+ chunks = {"time": 1, "pressure_level": 7, "longitude": 720, "latitude": 721}
+
+ print(f"{dt.datetime.utcnow()}: Opening datasets")
+ era5_z = xr.open_mfdataset(
+ pl_files[0], engine="h5netcdf", chunks="auto", parallel=True
+ ).sel(valid_time=args["date"])
+ era5_t = xr.open_mfdataset(
+ pl_files[1], engine="h5netcdf", chunks="auto", parallel=True
+ ).sel(valid_time=args["date"])
+ era5_q = xr.open_mfdataset(
+ pl_files[2], engine="h5netcdf", chunks="auto", parallel=True
+ ).sel(valid_time=args["date"])
+ era5_u = xr.open_mfdataset(
+ pl_files[3], engine="h5netcdf", chunks="auto", parallel=True
+ ).sel(valid_time=args["date"])
+ era5_v = xr.open_mfdataset(
+ pl_files[4], engine="h5netcdf", chunks="auto", parallel=True
+ ).sel(valid_time=args["date"])
+ print(f"{dt.datetime.utcnow()}: Merging datasets (this may take a while)")
+ era5_pl = xr.merge([era5_z, era5_t, era5_q, era5_u, era5_v]).compute(
+ num_workers=args["parallel_threads"], scheduler="processes"
+ )
+
+ dims = list(era5_pl.coords.dims.mapping.keys())
+ keys = list(era5_pl.keys())
+
+ # variable and coordinate check
+ assert "t" in keys, "Temperature (t) is missing from the provided dataset."
+ assert "q" in keys, "Specific humidity (q) is missing from the provided dataset."
+ assert "z" in keys, "Geopotential (z) is missing from the provided dataset."
+ assert "u" in keys, "U-wind (u) is missing from the provided dataset."
+ assert "v" in keys, "V-wind (v) is missing from the provided dataset."
+ assert "pressure_level" in dims, (
+ "Pressure level (pressure_level) is missing from the provided dataset."
+ )
+
+ # remove attributes unnecessary coordinates
+ era5_pl = era5_pl.drop_attrs()
+ era5_pl = era5_pl.drop_vars(["number", "expver"])
+
+ # rename temperature and geopotential
+ era5_pl = era5_pl.rename({"t": "T", "z": "sp_z", "valid_time": "time"})
+
+ # convert geopotential to geopotential height (dam)
+ era5_pl["sp_z"] /= 98.0665
+
+ coords = ("time", "pressure_level", "latitude", "longitude")
+
+ Ntime, Npl, Nlat, Nlon = era5_pl[
+ "T"
+ ].shape # (time, pressure_level, latitude, longitude)
+ pressure_levels = era5_pl[
+ "pressure_level"
+ ].values # (time, pressure_level, latitude, longitude)
+
+ P = np.full(
+ (Ntime, Npl, Nlat, Nlon), pressure_levels[np.newaxis, :, np.newaxis, np.newaxis]
+ )
+
+ print(f"{dt.datetime.utcnow()}: Adding variable: Dewpoint temperature")
+ era5_pl["Td"] = (
+ coords,
+ variables.dewpoint_from_specific_humidity(P * 100, era5_pl["q"].data),
+ )
+ era5_pl["Td"] = era5_pl["Td"].assign_attrs(
+ long_name="Dewpoint Temperature", units="K"
+ )
+
+ print(f"{dt.datetime.utcnow()}: Adding variable: Virtual temperature")
+ era5_pl["Tv"] = (
+ coords,
+ variables.virtual_temperature_from_dewpoint(
+ P * 100, era5_pl["T"].data, era5_pl["Td"].data
+ ),
+ )
+ era5_pl["Tv"] = era5_pl["Tv"].assign_attrs(
+ long_name="Virtual Temperature", units="K"
+ )
+
+ print(f"{dt.datetime.utcnow()}: Adding variable: Relative humidity")
+ era5_pl["RH"] = (
+ coords,
+ variables.relative_humidity_from_dewpoint(
+ era5_pl["T"].data, era5_pl["Td"].data
+ ),
+ )
+ era5_pl["RH"] = era5_pl["RH"].assign_attrs(
+ long_name="Relative Humidity", units="dimensionless"
+ )
+
+ print(f"{dt.datetime.utcnow()}: Adding variable: Potential temperature")
+ era5_pl["theta"] = (
+ coords,
+ variables.potential_temperature(P * 100, era5_pl["T"].data),
+ )
+ era5_pl["theta"] = era5_pl["theta"].assign_attrs(
+ long_name="Potential Temperature", units="K"
+ )
+
+ print(f"{dt.datetime.utcnow()}: Adding variable: Virtual potential temperature")
+ era5_pl["theta_v"] = (
+ coords,
+ variables.virtual_potential_temperature(
+ P * 100, era5_pl["T"].data, era5_pl["Td"].data
+ ),
+ )
+ era5_pl["theta_v"] = era5_pl["theta_v"].assign_attrs(
+ long_name="Virtual Potential Temperature", units="K"
+ )
+
+ print(f"{dt.datetime.utcnow()}: Adding variable: Equivalent potential temperature")
+ era5_pl["theta_e"] = (
+ coords,
+ variables.equivalent_potential_temperature(
+ P * 100, era5_pl["T"].data, era5_pl["Td"].data
+ ),
+ )
+ era5_pl["theta_e"] = era5_pl["theta_e"].assign_attrs(
+ long_name="Equivalent Potential Temperature", units="K"
+ )
+
+ # convert specific humidity from kg/kg to g/kg
+ era5_pl["q"] *= 1e3
+
+ # assign attributes to remaining variables and dataset
+ era5_pl["T"] = era5_pl["T"].assign_attrs(long_name="Temperature", units="K")
+ era5_pl["q"] = era5_pl["q"].assign_attrs(
+ long_name="Specific Humidity", units="g/kg"
+ )
+ era5_pl["u"] = era5_pl["u"].assign_attrs(
+ long_name="U-component of wind", units="m/s"
+ )
+ era5_pl["v"] = era5_pl["v"].assign_attrs(
+ long_name="V-component of wind", units="m/s"
+ )
+ era5_pl["sp_z"] = era5_pl["sp_z"].assign_attrs(
+ long_name="Geopotential Height", units="dam"
+ )
+ era5_pl = era5_pl.assign_attrs(time_created=dt.datetime.utcnow().timestamp())
+
+ # output directory for the ERA5 netCDF files (will be stored by month)
+ output_folder = f"{args['outdir']}/{yr}{mo}"
+ os.makedirs(output_folder, exist_ok=True)
+
+ for i, timestep in enumerate(era5_pl["time"].values):
+ hr = "%02d" % pd.to_datetime(timestep).hour
+ print(
+ f"{dt.datetime.utcnow()}: Saving {f'{output_folder}/era5_{yr}{mo}{dy}{hr}_global.nc'}"
+ )
+ era5_pl.isel(time=i).astype("float32").to_netcdf(
+ f"{output_folder}/era5_{yr}{mo}{dy}{hr}_global.nc",
+ engine="netcdf4",
+ mode="w",
+ )
diff --git a/src/fronts/data/create_evaluation_dataset.py b/src/fronts/data/create_evaluation_dataset.py
new file mode 100644
index 0000000..94a274a
--- /dev/null
+++ b/src/fronts/data/create_evaluation_dataset.py
@@ -0,0 +1,585 @@
+"""
+Convert netCDF files containing variable and frontal boundary data into tensorflow datasets for model evaluation.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import argparse
+import itertools
+import numpy as np
+import os
+import pandas as pd
+import pickle
+import tensorflow as tf
+from fronts.utils import file_manager as fm
+from fronts.utils import data_utils, misc
+from datetime import datetime
+import xarray as xr
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--variables_indir",
+ type=str,
+ required=True,
+ help="Input directory for the netCDF files containing variable data.",
+ )
+ parser.add_argument(
+ "--fronts_indir",
+ type=str,
+ help="Input directory for the netCDF files containing frontal boundary data.",
+ )
+ parser.add_argument(
+ "--goes_indir",
+ type=str,
+ help="Input directory for the netCDF files containing GOES satellite data.",
+ )
+ parser.add_argument(
+ "--tf_outdir",
+ type=str,
+ required=True,
+ help="Output directory for the generated tensorflow datasets.",
+ )
+ parser.add_argument(
+ "--year_and_month",
+ type=int,
+ nargs=2,
+ required=True,
+ help="Year and month for the netcdf data to be converted to tensorflow datasets.",
+ )
+ parser.add_argument(
+ "--data_source",
+ type=str,
+ default="era5",
+ help="Data source or model containing the variable data.",
+ )
+ parser.add_argument(
+ "--front_types",
+ type=str,
+ nargs="+",
+ help="Code(s) for the front types that will be generated in the tensorflow datasets. Refer to documentation in "
+ "'utils.data_utils.reformat_fronts' for more information on these codes.",
+ )
+ parser.add_argument(
+ "--variables", type=str, nargs="+", required=True, help="Variables to select"
+ )
+ parser.add_argument(
+ "--pressure_levels",
+ type=str,
+ nargs="+",
+ help="Variables pressure levels to select",
+ )
+ parser.add_argument(
+ "--num_dims",
+ type=int,
+ nargs=2,
+ default=[2, 2],
+ help="Number of dimensions in the variables and front object images, repsectively.",
+ )
+ parser.add_argument(
+ "--domain",
+ type=str,
+ default="conus",
+ help="Domain from which to pull the images.",
+ )
+ parser.add_argument(
+ "--override_extent",
+ type=float,
+ nargs=4,
+ help="Override the default domain extent by selecting a custom extent. [min lon, max lon, min lat, max lat]",
+ )
+ parser.add_argument(
+ "--image_size",
+ type=int,
+ nargs=2,
+ default=[128, 128],
+ help="Size of the longitude and latitude dimensions of the images.",
+ )
+ parser.add_argument(
+ "--normalization_method",
+ type=str,
+ default="standard",
+ help="Method for normalizing the datasets. Options are 'standard', 'standard_weighted', 'min-max'.",
+ )
+ parser.add_argument(
+ "--front_dilation",
+ type=int,
+ default=0,
+ help="Number of pixels to expand the fronts by in all directions.",
+ )
+ parser.add_argument(
+ "--verbose",
+ action="store_true",
+ help="Print out the progress of the dataset generation.",
+ )
+ parser.add_argument("--gpu_device", type=int, nargs="+", help="GPU device numbers.")
+
+ args = vars(parser.parse_args())
+
+ """
+ It is recommended to run this script on a GPU due to the abundance of TensorFlow operations.
+ """
+ if args["gpu_device"] is not None:
+ misc.initialize_gpus(
+ args["gpu_device"], memory_growth=True
+ ) # initialize the specified GPU
+ else:
+ os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+ year, month = args["year_and_month"]
+
+ """
+ all_data_vars: Variables found in ERA5 reanalysis and NWP models
+ all_goes_vars: Satellite bands found in the GOES datasets after the merging process
+ """
+ all_data_vars = [
+ "T",
+ "Td",
+ "sp_z",
+ "u",
+ "v",
+ "r",
+ "RH",
+ "Tv",
+ "theta_e",
+ "q",
+ "theta",
+ "theta_v",
+ ]
+ all_pressure_levels = (
+ ["surface", "1000", "950", "900", "850"]
+ if args["data_source"] == "era5"
+ else ["surface", "1013", "1000", "950", "900", "850", "700", "500"]
+ )
+
+ # check for invalid variables
+ invalid_vars = [var for var in args["variables"] if var not in all_data_vars]
+ assert len(invalid_vars) == 0, "Invalid variables (%d): %s" % (
+ len(invalid_vars),
+ ", ".join(invalid_vars),
+ )
+
+ data_vars = [
+ var for var in args["variables"] if var in all_data_vars
+ ] # ERA5/model variables that will be used
+
+ os.makedirs(
+ args["tf_outdir"], exist_ok=True
+ ) # ensure that a folder exists for the monthly dataset
+
+ tf_dataset_folder_inputs = f"%s/inputs_%d%02d_tf" % (
+ args["tf_outdir"],
+ year,
+ month,
+ ) # output directory for the inputs
+ tf_dataset_folder_fronts = tf_dataset_folder_inputs.replace(
+ "inputs", "fronts"
+ ) # output directory for the front labels
+
+ # ensure that the requested month does not already have a saved dataset
+ if os.path.isdir(tf_dataset_folder_inputs) or os.path.isdir(
+ tf_dataset_folder_fronts
+ ):
+ raise FileExistsError(
+ "Tensorflow dataset(s) already exist for the provided year and month."
+ )
+
+ # dataset properties file - this will contain critical information about the dataset
+ dataset_props_file = "%s/dataset_properties.pkl" % args["tf_outdir"]
+ if not os.path.isfile(dataset_props_file):
+ dataset_props = dict({})
+ dataset_props["normalization_parameters"] = data_utils.NORMALIZATION_PARAMS
+ for key in sorted(
+ [
+ "front_types",
+ "variables",
+ "pressure_levels",
+ "num_dims",
+ "image_size",
+ "normalization_method",
+ "front_dilation",
+ "domain",
+ "override_extent",
+ ]
+ ):
+ dataset_props[key] = args[key]
+
+ # save out the dataset properties pickle file
+ with open(dataset_props_file, "wb") as f:
+ pickle.dump(dataset_props, f)
+
+ # create a text file with a human-readable output of the information saved in the dataset properties pickle file
+ with open("%s/dataset_properties.txt" % args["tf_outdir"], "w") as f:
+ for key in sorted(dataset_props.keys()):
+ f.write(f"{key}: {dataset_props[key]}\n")
+ f.write(f"\n\n\nFile generated at {datetime.utcnow()} UTC\n")
+
+ else:
+ """
+ If the dataset properties pickle file already exists, many arguments declared at the command line will be overwritten
+ by the values contained within the pickle file. This behavior exists to ensure that the dataset does not have inconsistent
+ properties relegated to certain months.
+ """
+ print(
+ "WARNING: Dataset properties file was found in %s. The following settings will be used from the file."
+ % args["tf_outdir"]
+ )
+ dataset_props = pd.read_pickle(dataset_props_file)
+
+ for key in sorted(
+ [
+ "front_types",
+ "variables",
+ "pressure_levels",
+ "num_dims",
+ "image_size",
+ "normalization_method",
+ "front_dilation",
+ "domain",
+ "override_extent",
+ ]
+ ):
+ args[key] = dataset_props[key]
+ print(f"%s: {args[key]}" % key)
+
+ file_loader_domain = "global" if args["data_source"] == "era5" else args["domain"]
+
+ # Gather all ERA5/model and front label files that can be used to generate the dataset for the current month.
+ file_obj = fm.DataFileLoader(
+ args["variables_indir"],
+ args["data_source"],
+ "netcdf",
+ years=year,
+ months=month,
+ domains=file_loader_domain,
+ )
+ file_obj.add_file_list(args["fronts_indir"], "fronts", ignore_domain=True)
+ variables_netcdf_files, fronts_netcdf_files = file_obj.files
+
+ # if not looking over CONUS or the HRRR domain, remove non-synoptic hours (3, 9, 15, 21z)
+ if args["domain"] not in ["conus", "hrrr"]:
+ synoptic_ind = [
+ variables_netcdf_files.index(file)
+ for file in variables_netcdf_files
+ if any(["%02d_" % hr in file for hr in [0, 6, 12, 18]])
+ ]
+ variables_netcdf_files = list([variables_netcdf_files[i] for i in synoptic_ind])
+ fronts_netcdf_files = list([fronts_netcdf_files[i] for i in synoptic_ind])
+
+ # if the extent crosses the Prime Meridian (0 degrees longitude), we need to load the data in differently
+ extent_crosses_meridian = False
+ if args["override_extent"] is not None:
+ if args["override_extent"][1] > 360:
+ extent_crosses_meridian = True
+
+ if args["domain"] in ["conus", "full", "goes-merged"]:
+ if args["override_extent"] is None:
+ sel_kwargs = {
+ "latitude": slice(
+ data_utils.DOMAIN_EXTENTS[args["domain"]][3],
+ data_utils.DOMAIN_EXTENTS[args["domain"]][2],
+ ),
+ "longitude": slice(
+ data_utils.DOMAIN_EXTENTS[args["domain"]][0],
+ data_utils.DOMAIN_EXTENTS[args["domain"]][1],
+ ),
+ }
+ else:
+ sel_kwargs = {
+ "latitude": slice(
+ args["override_extent"][3], args["override_extent"][2]
+ ),
+ "longitude": slice(
+ args["override_extent"][0], args["override_extent"][1]
+ ),
+ }
+ else:
+ sel_kwargs = {}
+
+ args["pressure_levels"] = (
+ all_pressure_levels
+ if args["pressure_levels"] is None
+ else [lvl for lvl in all_pressure_levels if lvl in args["pressure_levels"]]
+ )
+
+ num_timesteps = len(variables_netcdf_files)
+ images_kept = 0
+ images_discarded = 0
+ timesteps_kept = 0
+ timesteps_discarded = 0
+
+ isel_kwargs = dict(forecast_hour=0) if args["data_source"] != "era5" else dict()
+
+ """
+ In order to make sure that the final dataset comes out clean, we will keep track of all of the input shapes of
+ the generated images as they are being generated. This will allow us to catch any indexing error that may produce
+ an image shape that is different from the rest of the dataset (e.g., say an indexing error produces an image of
+ size 128x127 when we intended to have its shape be 128x128). If the tensor shapes are not all identical, TensorFlow
+ will raise an error before model training and render the dataset effectively useless.
+ """
+ input_tensor_shapes = []
+ front_tensor_shapes = []
+
+ front_files_kept = []
+
+ for timestep_no in range(num_timesteps):
+ # open front dataset
+ front_dataset = xr.open_dataset(
+ fronts_netcdf_files[timestep_no], engine="netcdf4"
+ ).isel(**isel_kwargs)
+
+ if args["data_source"] not in ["hrrr", "namnest-conus", "nam-12km"]:
+ front_dataset = front_dataset.sel(**sel_kwargs).astype("float16")
+ transpose_dims = (
+ "latitude",
+ "longitude",
+ ) # spatial dimensions that need to be transposed
+ else:
+ transpose_dims = ("y", "x") # spatial dimensions that need to be transposed
+ domain_size = (
+ len(front_dataset[transpose_dims[0]]),
+ len(front_dataset[transpose_dims[1]]),
+ )
+
+ # Reformat the fronts in the current timestep
+ if args["front_types"] is not None:
+ front_dataset = data_utils.reformat_fronts(
+ front_dataset, args["front_types"]
+ )
+ num_front_types = front_dataset.attrs["num_front_types"] + 1
+ else:
+ num_front_types = 16
+
+ # Expand the front labels
+ if args["front_dilation"] > 0:
+ front_dataset = data_utils.expand_fronts(
+ front_dataset, iterations=args["front_dilation"]
+ )
+
+ # Check for all front types in the dataset
+ front_dataset = (
+ front_dataset.isel(time=0)
+ if "time" in front_dataset.dims
+ else front_dataset
+ )
+ front_dataset = front_dataset.to_array().transpose(*transpose_dims, "variable")
+
+ if args["verbose"]:
+ print(
+ "%d-%02d Dataset progress (kept/discarded): (%d/%d timesteps, %d/%d images)"
+ % (
+ year,
+ month,
+ timesteps_kept,
+ timesteps_discarded,
+ images_kept,
+ images_discarded,
+ ),
+ end="\r",
+ )
+
+ # open variables dataset
+ if extent_crosses_meridian:
+ sel_kwargs_1 = {
+ "latitude": slice(
+ args["override_extent"][3], args["override_extent"][2]
+ ),
+ "longitude": slice(args["override_extent"][0], 360),
+ } # extent west of the Prime Meridian
+ sel_kwargs_2 = {
+ "latitude": slice(
+ args["override_extent"][3], args["override_extent"][2]
+ ),
+ "longitude": slice(0, args["override_extent"][1] - 360),
+ } # extent east of the Prime Meridian
+
+ variables_dataset_1 = (
+ xr.open_dataset(
+ variables_netcdf_files[timestep_no], engine="netcdf4"
+ )[data_vars]
+ .sel(pressure_level=args["pressure_levels"], **sel_kwargs_1)
+ .isel(**isel_kwargs)
+ .transpose("time", *transpose_dims, "pressure_level")
+ .astype("float16")
+ )
+ variables_dataset_2 = (
+ xr.open_dataset(
+ variables_netcdf_files[timestep_no], engine="netcdf4"
+ )[data_vars]
+ .sel(pressure_level=args["pressure_levels"], **sel_kwargs_2)
+ .isel(**isel_kwargs)
+ .transpose("time", *transpose_dims, "pressure_level")
+ .astype("float16")
+ )
+ variables_dataset = xr.merge([variables_dataset_1, variables_dataset_2])
+
+ else:
+ variables_dataset = (
+ xr.open_dataset(
+ variables_netcdf_files[timestep_no], engine="netcdf4"
+ )[data_vars]
+ .sel(pressure_level=args["pressure_levels"], **sel_kwargs)
+ .isel(**isel_kwargs)
+ .transpose("time", *transpose_dims, "pressure_level")
+ .astype("float16")
+ )
+ variables_dataset = (
+ variables_dataset.isel(time=0)
+ .transpose(*transpose_dims, "pressure_level")
+ .astype("float16")
+ )
+
+ # create a list of starting indices along the latitude dimension
+ start_indices_lat = [
+ 0,
+ ]
+ start_indices_lon = [
+ 0,
+ ]
+
+ image_order = list(
+ itertools.product(start_indices_lat, start_indices_lon)
+ ) # Every possible combination of longitude and latitude starting points
+
+ for i, image_start_indices in enumerate(image_order):
+ if args["verbose"]:
+ print(
+ "%d-%02d Dataset progress (kept/discarded): (%d/%d timesteps, %d/%d images)"
+ % (
+ year,
+ month,
+ timesteps_kept,
+ timesteps_discarded,
+ images_kept,
+ images_discarded,
+ ),
+ end="\r",
+ )
+
+ start_index_lat = image_start_indices[0]
+ end_index_lat = start_index_lat + args["image_size"][0]
+ start_index_lon = image_start_indices[1]
+ end_index_lon = start_index_lon + args["image_size"][1]
+
+ front_image = front_dataset[
+ start_index_lat:end_index_lat, start_index_lon:end_index_lon, :
+ ]
+
+ new_variables_dataset = (
+ variables_dataset.copy()
+ ) # copy variables dataset to isolate dataset in memory
+
+ # normalize variables and convert dataset to a tensor
+ new_variables_dataset = (
+ data_utils.normalize_dataset(
+ new_variables_dataset,
+ args["normalization_method"],
+ dataset_props["normalization_parameters"],
+ )
+ .to_array()
+ .transpose(*transpose_dims, "pressure_level", "variable")
+ )
+ input_tensor = tf.convert_to_tensor(
+ np.nan_to_num(
+ new_variables_dataset[
+ start_index_lat:end_index_lat,
+ start_index_lon:end_index_lon,
+ :,
+ :,
+ ]
+ ),
+ dtype=tf.float16,
+ )
+
+ # combine pressure level and variables dimensions, making the images 2D (excluding the final dimension)
+ if args["num_dims"][0] == 2:
+ input_tensor_shape_3d = input_tensor.shape
+ input_tensor = tf.reshape(
+ input_tensor,
+ [
+ input_tensor_shape_3d[0],
+ input_tensor_shape_3d[1],
+ input_tensor_shape_3d[2] * input_tensor_shape_3d[3],
+ ],
+ )
+
+ input_tensor_shapes.append(input_tensor.shape)
+ assert len(set(input_tensor_shapes)) == 1, (
+ f"ERROR: Attempted to add {input_tensor_shapes[-1]} to dataset with shape {input_tensor_shapes[0]}. "
+ "Please check your data for inconsistent coordinate systems."
+ )
+
+ # add input images to tensorflow dataset
+ input_tensor_for_timestep = tf.data.Dataset.from_tensors(input_tensor)
+ if "input_tensors_for_month" not in locals():
+ input_tensors_for_month = input_tensor_for_timestep
+ else:
+ input_tensors_for_month = input_tensors_for_month.concatenate(
+ input_tensor_for_timestep
+ )
+
+ front_tensor = tf.convert_to_tensor(
+ np.nan_to_num(front_image), dtype=tf.int32
+ )
+ front_tensor_shapes.append(front_tensor.shape)
+
+ assert len(set(front_tensor_shapes)) == 1, (
+ f"ERROR: Attempted to add {front_tensor_shapes[-1]} to dataset with shape {front_tensor_shapes[0]}. "
+ "Please check your data for inconsistent coordinate systems."
+ )
+
+ # if using 3D inputs, turn the fronts dataset into a 3D image
+ if args["num_dims"][1] == 3:
+ front_tensor = tf.tile(
+ front_tensor, (1, 1, len(args["pressure_levels"]))
+ )
+ else:
+ front_tensor = front_tensor[:, :, 0]
+
+ front_tensor = tf.cast(
+ tf.one_hot(front_tensor, num_front_types), tf.float16
+ ) # One-hot encode the labels
+ front_tensor_for_timestep = tf.data.Dataset.from_tensors(
+ front_tensor
+ ) # convert fronts into a tensorflow dataset
+ if "front_tensors_for_month" not in locals():
+ front_tensors_for_month = front_tensor_for_timestep
+ else:
+ front_tensors_for_month = front_tensors_for_month.concatenate(
+ front_tensor_for_timestep
+ )
+
+ timesteps_kept += 1
+ front_files_kept.append(fronts_netcdf_files[timestep_no])
+ else:
+ timesteps_discarded += 1
+
+ print(
+ "%d-%02d Dataset progress (kept/discarded): (%d/%d timesteps, %d/%d images)"
+ % (
+ year,
+ month,
+ timesteps_kept,
+ timesteps_discarded,
+ images_kept,
+ images_discarded,
+ )
+ )
+
+ # save the tensorflow datasets
+ try:
+ tf.data.Dataset.save(input_tensors_for_month, path=tf_dataset_folder_inputs)
+ tf.data.Dataset.save(front_tensors_for_month, path=tf_dataset_folder_fronts)
+ print(
+ "Tensorflow datasets for %d-%02d saved to %s."
+ % (year, month, args["tf_outdir"])
+ )
+ except NameError:
+ print("No images could be retained with the provided arguments.")
+
+ with open(
+ "%s/front_files_%d%02d.pkl" % (args["tf_outdir"], year, month), "wb"
+ ) as f:
+ pickle.dump(np.array(front_files_kept), f)
diff --git a/src/fronts/data/create_training_dataset.py b/src/fronts/data/create_training_dataset.py
new file mode 100644
index 0000000..efba54a
--- /dev/null
+++ b/src/fronts/data/create_training_dataset.py
@@ -0,0 +1,733 @@
+"""
+Convert netCDF files containing variable and frontal boundary data into tensorflow datasets for model training.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import argparse
+import itertools
+import numpy as np
+import os
+import pandas as pd
+import pickle
+import tensorflow as tf
+from fronts.utils import file_manager as fm
+from fronts.utils import data_utils, misc
+from datetime import datetime
+import xarray as xr
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--variables_indir",
+ type=str,
+ required=True,
+ help="Input directory for the netCDF files containing variable data.",
+ )
+ parser.add_argument(
+ "--fronts_indir",
+ type=str,
+ help="Input directory for the netCDF files containing frontal boundary data.",
+ )
+ parser.add_argument(
+ "--tf_outdir",
+ type=str,
+ required=True,
+ help="Output directory for the generated tensorflow datasets.",
+ )
+ parser.add_argument(
+ "--year",
+ type=int,
+ required=True,
+ help="Year for the netCDF data to be converted to tensorflow datasets.",
+ )
+ parser.add_argument(
+ "--data_source",
+ type=str,
+ default="era5",
+ help="Data source or model containing the variable data.",
+ )
+ parser.add_argument(
+ "--front_types",
+ type=str,
+ nargs="+",
+ help="Code(s) for the front types that will be generated in the tensorflow datasets. Refer to documentation in "
+ "'utils.data_utils.reformat_fronts' for more information on these codes.",
+ )
+ parser.add_argument(
+ "--variables", type=str, nargs="+", required=True, help="Variables to select"
+ )
+ parser.add_argument(
+ "--pressure_levels",
+ type=str,
+ nargs="+",
+ help="Variables pressure levels to select",
+ )
+ parser.add_argument(
+ "--num_dims",
+ type=int,
+ nargs=2,
+ default=[2, 2],
+ help="Number of dimensions in the variables and front object images, repsectively.",
+ )
+ parser.add_argument(
+ "--num_slices",
+ type=int,
+ default=30,
+ help="Number of slices to perform for each year in the dataset.",
+ )
+ parser.add_argument(
+ "--domain",
+ type=str,
+ default="conus",
+ help="Domain from which to pull the images.",
+ )
+ parser.add_argument(
+ "--override_extent",
+ type=float,
+ nargs=4,
+ help="Override the default domain extent by selecting a custom extent. [min lon, max lon, min lat, max lat]",
+ )
+ parser.add_argument(
+ "--images",
+ type=int,
+ nargs=2,
+ default=[1, 1],
+ help="Number of variables/front images along the latitude and longitude dimensions to generate for each timestep. The product of the 2 integers "
+ "will be the total number of images generated per timestep.",
+ )
+ parser.add_argument(
+ "--image_size",
+ type=int,
+ nargs=2,
+ default=[128, 128],
+ help="Size of the longitude and latitude dimensions of the images.",
+ )
+ parser.add_argument(
+ "--normalization_method",
+ type=str,
+ default="standard",
+ help="Method for normalizing the datasets. Options are 'standard', 'standard_weighted', 'min-max'.",
+ )
+ parser.add_argument(
+ "--shuffle_timesteps",
+ action="store_true",
+ help="Shuffle the timesteps when generating the dataset. This is particularly useful when generating very large "
+ "datasets that cannot be shuffled on the fly during training.",
+ )
+ parser.add_argument(
+ "--shuffle_images",
+ action="store_true",
+ help="Shuffle the order of the images in each timestep. This does NOT shuffle the entire dataset for the provided "
+ "month, but rather only the images in each respective timestep. This is particularly useful when generating "
+ "very large datasets that cannot be shuffled on the fly during training.",
+ )
+ parser.add_argument(
+ "--add_previous_fronts",
+ type=str,
+ nargs="+",
+ help="Optional front types from previous timesteps to include as predictors. If the dataset is over conus, the fronts "
+ "will be pulled from the last 3-hour timestep. If the dataset is over the full domain, the fronts will be pulled "
+ "from the last 6-hour timestep.",
+ )
+ parser.add_argument(
+ "--front_dilation",
+ type=int,
+ default=0,
+ help="Number of pixels to expand the fronts by in all directions.",
+ )
+ parser.add_argument(
+ "--timestep_fraction",
+ type=float,
+ default=1.0,
+ help="The fraction of timesteps WITHOUT all necessary front types that will be retained in the dataset. Can be any float 0 <= x <= 1.",
+ )
+ parser.add_argument(
+ "--image_fraction",
+ type=float,
+ default=1.0,
+ help="The fraction of images WITHOUT all necessary front types in the selected timesteps that will be retained in the dataset. Can be any float 0 <= x <= 1. "
+ "By default, all images are retained.",
+ )
+ parser.add_argument(
+ "--noise_fraction",
+ type=float,
+ default=0.0,
+ help="The fraction of pixels in each image that will contain noise. Can be any float 0 <= x < 1.",
+ )
+ parser.add_argument(
+ "--flip_chance_lat",
+ type=float,
+ default=0.0,
+ help="The probability that the current image will have its latitude dimension reversed. Can be any float 0 <= x <= 1.",
+ )
+ parser.add_argument(
+ "--flip_chance_lon",
+ type=float,
+ default=0.0,
+ help="The probability that the current image will have its longitude dimension reversed. Can be any float 0 <= x <= 1.",
+ )
+ parser.add_argument(
+ "--verbose",
+ action="store_true",
+ help="Print out the progress of the dataset generation.",
+ )
+ parser.add_argument("--gpu_device", type=int, nargs="+", help="GPU device numbers.")
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=np.random.randint(0, 2**31 - 1),
+ help="Seed for the random number generators. The same seed will be used for all months within a particular dataset.",
+ )
+
+ args = vars(parser.parse_args())
+
+ """
+ It is recommended to run this script on a GPU due to the abundance of TensorFlow operations.
+ """
+ if args["gpu_device"] is not None:
+ misc.initialize_gpus(
+ args["gpu_device"], memory_growth=True
+ ) # initialize the specified GPU
+ else:
+ os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+ yr = args["year"]
+ num_slices = args["num_slices"]
+
+ """
+ all_data_vars: Variables found in ERA5 reanalysis and NWP models
+ all_goes_vars: Satellite bands found in the GOES datasets after the merging process
+ """
+ all_data_vars = [
+ "T",
+ "Td",
+ "sp_z",
+ "u",
+ "v",
+ "theta_w",
+ "r",
+ "RH",
+ "Tv",
+ "Tw",
+ "theta_e",
+ "q",
+ "theta",
+ "theta_v",
+ ]
+ all_pressure_levels = (
+ ["surface", "1000", "950", "900", "850"]
+ if args["data_source"] == "era5"
+ else ["surface", "1013", "1000", "950", "900", "850", "700", "500"]
+ )
+
+ # check for invalid variables
+ invalid_vars = [var for var in args["variables"] if var not in all_data_vars]
+ assert len(invalid_vars) == 0, "Invalid variables (%d): %s" % (
+ len(invalid_vars),
+ ", ".join(invalid_vars),
+ )
+
+ data_vars = [
+ var for var in args["variables"] if var in all_data_vars
+ ] # ERA5/model variables that will be used
+
+ os.makedirs(
+ args["tf_outdir"], exist_ok=True
+ ) # ensure that a folder exists for the dataset
+
+ # dataset properties file - this will contain critical information about the dataset
+ dataset_props_file = "%s/dataset_properties.pkl" % args["tf_outdir"]
+ if not os.path.isfile(dataset_props_file):
+ """
+ Save critical dataset information to a pickle file so it can be referenced later when generating data for other months.
+ """
+ print("Setting seed: %d" % args["seed"])
+
+ dataset_props = dict({})
+ dataset_props["normalization_parameters"] = data_utils.NORMALIZATION_PARAMS
+ for key in sorted(
+ [
+ "front_types",
+ "variables",
+ "pressure_levels",
+ "num_dims",
+ "images",
+ "image_size",
+ "normalization_method",
+ "front_dilation",
+ "noise_fraction",
+ "flip_chance_lat",
+ "flip_chance_lon",
+ "shuffle_images",
+ "shuffle_timesteps",
+ "domain",
+ "add_previous_fronts",
+ "timestep_fraction",
+ "image_fraction",
+ "override_extent",
+ "seed",
+ ]
+ ):
+ dataset_props[key] = args[key]
+
+ # save out the dataset properties pickle file
+ with open(dataset_props_file, "wb") as f:
+ pickle.dump(dataset_props, f)
+
+ # create a text file with a human-readable output of the information saved in the dataset properties pickle file
+ with open("%s/dataset_properties.txt" % args["tf_outdir"], "w") as f:
+ for key in sorted(dataset_props.keys()):
+ f.write(f"{key}: {dataset_props[key]}\n")
+ f.write(f"\n\n\nFile generated at {datetime.utcnow()} UTC\n")
+
+ else:
+ """
+ If the dataset properties pickle file already exists, many arguments declared at the command line will be overwritten
+ by the values contained within the pickle file. This behavior exists to ensure that the dataset does not have inconsistent
+ properties relegated to certain months.
+ """
+ print(
+ "WARNING: Dataset properties file was found in %s. The following settings will be used from the file."
+ % args["tf_outdir"]
+ )
+ dataset_props = pd.read_pickle(dataset_props_file)
+
+ for key in sorted(
+ [
+ "front_types",
+ "variables",
+ "pressure_levels",
+ "num_dims",
+ "images",
+ "image_size",
+ "normalization_method",
+ "front_dilation",
+ "noise_fraction",
+ "flip_chance_lat",
+ "flip_chance_lon",
+ "shuffle_images",
+ "shuffle_timesteps",
+ "domain",
+ "add_previous_fronts",
+ "timestep_fraction",
+ "image_fraction",
+ "override_extent",
+ ]
+ ):
+ args[key] = dataset_props[key]
+ print(f"%s: {args[key]}" % key)
+
+ if "seed" in list(dataset_props.keys()):
+ args["seed"] = dataset_props[
+ "seed"
+ ] # keep the same seed for consistency sake and reproducibility
+ print(f"%s: {args['seed']}" % "seed")
+
+ # set the seed
+ tf.random.set_seed(args["seed"])
+ np.random.seed(args["seed"])
+
+ file_loader_domain = "global" if args["data_source"] == "era5" else args["domain"]
+
+ # Gather all ERA5/model and front label files that can be used to generate the dataset for the current year.
+ file_obj = fm.DataFileLoader(
+ args["variables_indir"],
+ args["data_source"],
+ "netcdf",
+ years=yr,
+ domains=file_loader_domain,
+ )
+ file_obj.add_file_list(args["fronts_indir"], "fronts", ignore_domain=True)
+ variables_netcdf_files, fronts_netcdf_files = file_obj.files
+
+ # if not looking over CONUS or the HRRR domain, remove non-synoptic hours (3, 9, 15, 21z)
+ if args["domain"] not in ["conus", "hrrr"]:
+ synoptic_ind = [
+ variables_netcdf_files.index(file)
+ for file in variables_netcdf_files
+ if any(["%02d_" % hr in file for hr in [0, 6, 12, 18]])
+ ]
+ variables_netcdf_files = list([variables_netcdf_files[i] for i in synoptic_ind])
+ fronts_netcdf_files = list([fronts_netcdf_files[i] for i in synoptic_ind])
+
+ # shuffle the order of the files, therefore shuffling the order of the timesteps
+ zipped_list = list(zip(variables_netcdf_files, fronts_netcdf_files))
+ np.random.shuffle(zipped_list)
+
+ sliced_list = list(np.array_split(np.array(zipped_list), num_slices))
+
+ # if the extent crosses the Prime Meridian (0 degrees longitude), we need to load the data in differently
+ extent_crosses_meridian = False
+ if args["override_extent"] is not None:
+ if args["override_extent"][1] > 360:
+ extent_crosses_meridian = True
+
+ if args["domain"] in ["conus", "full", "goes-merged"]:
+ if args["override_extent"] is None:
+ sel_kwargs = {
+ "latitude": slice(
+ data_utils.DOMAIN_EXTENTS[args["domain"]][3],
+ data_utils.DOMAIN_EXTENTS[args["domain"]][2],
+ ),
+ "longitude": slice(
+ data_utils.DOMAIN_EXTENTS[args["domain"]][0],
+ data_utils.DOMAIN_EXTENTS[args["domain"]][1],
+ ),
+ }
+ else:
+ sel_kwargs = {
+ "latitude": slice(
+ args["override_extent"][3], args["override_extent"][2]
+ ),
+ "longitude": slice(
+ args["override_extent"][0], args["override_extent"][1]
+ ),
+ }
+ else:
+ sel_kwargs = {}
+
+ args["pressure_levels"] = (
+ all_pressure_levels
+ if args["pressure_levels"] is None
+ else [lvl for lvl in all_pressure_levels if lvl in args["pressure_levels"]]
+ )
+
+ isel_kwargs = dict(forecast_hour=0) if args["data_source"] != "era5" else dict()
+
+ for i_slice, files_in_slice in enumerate(sliced_list, start=1):
+ """
+ In order to make sure that the final dataset comes out clean, we will keep track of all of the input shapes of
+ the generated images as they are being generated. This will allow us to catch any indexing error that may produce
+ an image shape that is different from the rest of the dataset (e.g., say an indexing error produces an image of
+ size 128x127 when we intended to have its shape be 128x128). If the tensor shapes are not all identical, TensorFlow
+ will raise an error before model training and render the dataset effectively useless.
+ """
+ input_tensor_shapes = []
+ front_tensor_shapes = []
+
+ input_tensors_for_slice = []
+ front_tensors_for_slice = []
+
+ timesteps_kept = 0
+ timesteps_discarded = 0
+ images_kept = 0
+ images_discarded = 0
+
+ print(f"Year {yr}: Generating slice {i_slice}/{num_slices}")
+
+ for input_file, fronts_file in files_in_slice:
+ keep_timestep = (
+ np.random.random() <= args["timestep_fraction"]
+ ) # boolean flag for keeping timesteps without all front types
+
+ # open front dataset
+ front_dataset = xr.open_dataset(fronts_file, engine="netcdf4").isel(
+ **isel_kwargs
+ )
+
+ if args["data_source"] not in ["hrrr", "namnest-conus", "nam-12km"]:
+ front_dataset = front_dataset.sel(**sel_kwargs).astype("float16")
+ transpose_dims = (
+ "latitude",
+ "longitude",
+ ) # spatial dimensions that need to be transposed
+ else:
+ transpose_dims = (
+ "y",
+ "x",
+ ) # spatial dimensions that need to be transposed
+ domain_size = (
+ len(front_dataset[transpose_dims[0]]),
+ len(front_dataset[transpose_dims[1]]),
+ )
+
+ # Reformat the fronts in the current timestep
+ if args["front_types"] is not None:
+ front_dataset = data_utils.reformat_fronts(
+ front_dataset, args["front_types"]
+ )
+ num_front_types = front_dataset.attrs["num_front_types"] + 1
+ else:
+ num_front_types = 16
+
+ # Expand the front labels
+ if args["front_dilation"] > 0:
+ front_dataset = data_utils.expand_fronts(
+ front_dataset, iterations=args["front_dilation"]
+ )
+
+ # Check for all front types in the dataset
+ front_dataset = (
+ front_dataset.isel(time=0)
+ if "time" in front_dataset.dims
+ else front_dataset
+ )
+ front_dataset = front_dataset.to_array().transpose(
+ *transpose_dims, "variable"
+ )
+ timestep_front_bins = np.bincount(
+ front_dataset.values.astype("int64").flatten(),
+ minlength=num_front_types,
+ ) # counts for each front type
+ all_fronts_in_timestep = (
+ all([front_count > 0 for front_count in timestep_front_bins]) > 0
+ ) # boolean flag that says if all front types are present in the current timestep
+
+ if all_fronts_in_timestep or keep_timestep:
+ # open variables dataset
+ if extent_crosses_meridian:
+ sel_kwargs_1 = {
+ "latitude": slice(
+ args["override_extent"][3], args["override_extent"][2]
+ ),
+ "longitude": slice(args["override_extent"][0], 360),
+ } # extent west of the Prime Meridian
+ sel_kwargs_2 = {
+ "latitude": slice(
+ args["override_extent"][3], args["override_extent"][2]
+ ),
+ "longitude": slice(0, args["override_extent"][1] - 360),
+ } # extent east of the Prime Meridian
+
+ variables_dataset_1 = (
+ xr.open_dataset(input_file, engine="netcdf4")[data_vars]
+ .sel(pressure_level=args["pressure_levels"], **sel_kwargs_1)
+ .isel(**isel_kwargs)
+ .transpose(*transpose_dims, "pressure_level")
+ .astype("float16")
+ )
+ variables_dataset_2 = (
+ xr.open_dataset(input_file, engine="netcdf4")[data_vars]
+ .sel(pressure_level=args["pressure_levels"], **sel_kwargs_2)
+ .isel(**isel_kwargs)
+ .transpose(*transpose_dims, "pressure_level")
+ .astype("float16")
+ )
+ variables_dataset = xr.merge(
+ [variables_dataset_1, variables_dataset_2]
+ )
+
+ else:
+ variables_dataset = (
+ xr.open_dataset(input_file, engine="netcdf4")[data_vars]
+ .sel(pressure_level=args["pressure_levels"], **sel_kwargs)
+ .isel(**isel_kwargs)
+ .transpose(*transpose_dims, "pressure_level")
+ .astype("float16")
+ )
+
+ # create a list of starting indices along the latitude dimension
+ if (
+ args["images"][0] > 1
+ and domain_size[0] > args["image_size"][0] + args["images"][0]
+ ):
+ start_indices_lat = np.linspace(
+ 0, domain_size[0] - args["image_size"][0], args["images"][0]
+ ).astype(int)
+ else:
+ start_indices_lat = np.zeros((args["images"][0],), dtype=int)
+
+ # create a list of starting indices along the longitude dimension
+ if (
+ args["images"][1] > 1
+ and domain_size[1] > args["image_size"][1] + args["images"][1]
+ ):
+ start_indices_lon = np.linspace(
+ 0, domain_size[1] - args["image_size"][1], args["images"][1]
+ ).astype(int)
+ else:
+ start_indices_lon = np.zeros((args["images"][1],), dtype=int)
+
+ image_order = list(
+ itertools.product(start_indices_lat, start_indices_lon)
+ ) # Every possible combination of longitude and latitude starting points
+
+ if args["shuffle_images"]:
+ np.random.shuffle(image_order)
+
+ images_to_keep = (
+ np.random.random(size=len(image_order)) <= args["image_fraction"]
+ )
+
+ for i_image, image_start_indices in enumerate(image_order):
+ if args["verbose"]:
+ print(
+ f"{yr}: slice {i_slice}/{num_slices} (kept/discarded): ({timesteps_kept}/{timesteps_discarded} timesteps, {images_kept}/{images_discarded} images)",
+ end="\r",
+ )
+
+ start_index_lat = image_start_indices[0]
+ end_index_lat = start_index_lat + args["image_size"][0]
+ start_index_lon = image_start_indices[1]
+ end_index_lon = start_index_lon + args["image_size"][1]
+
+ front_image = front_dataset[
+ start_index_lat:end_index_lat, start_index_lon:end_index_lon, :
+ ]
+ image_front_bins = np.bincount(
+ front_image.values.astype("int64").flatten(),
+ minlength=num_front_types,
+ ) # counts for each front type
+ all_fronts_in_image = (
+ all([front_count > 0 for front_count in image_front_bins]) > 0
+ ) # boolean flag that says if all front types are present in the current timestep
+
+ if not all_fronts_in_image and not images_to_keep[i_image]:
+ images_discarded += 1
+ continue
+
+ images_kept += 1
+
+ new_variables_dataset = (
+ variables_dataset.copy()
+ ) # copy variables dataset to isolate dataset in memory
+
+ # boolean flags for rotating and flipping images
+ flip_lat = np.random.random() <= args["flip_chance_lat"]
+ flip_lon = np.random.random() <= args["flip_chance_lon"]
+
+ # before flipping images, we will apply the necessary changes to the wind components to account for reflections
+ if flip_lat and "u" in args["variables"]:
+ new_variables_dataset["u"] = -new_variables_dataset[
+ "u"
+ ] # need to reverse u-wind component if flipping the longitude axis
+ if flip_lon and "v" in args["variables"]:
+ new_variables_dataset["v"] = -new_variables_dataset[
+ "v"
+ ] # need to reverse v-wind component if flipping the latitude axis
+
+ # normalize variables and convert datasets to tensors
+ new_variables_dataset = (
+ data_utils.normalize_dataset(
+ new_variables_dataset,
+ args["normalization_method"],
+ dataset_props["normalization_parameters"],
+ )
+ .to_array()
+ .transpose(*transpose_dims, "pressure_level", "variable")
+ )
+ input_tensor = tf.convert_to_tensor(
+ np.nan_to_num(
+ new_variables_dataset[
+ start_index_lat:end_index_lat,
+ start_index_lon:end_index_lon,
+ :,
+ :,
+ ]
+ ),
+ dtype=tf.float16,
+ )
+ front_tensor = tf.convert_to_tensor(
+ np.nan_to_num(front_image), dtype=tf.int32
+ )
+
+ random_values = tf.random.uniform(
+ shape=input_tensor.shape
+ ) # random noise values, will not necessarily be used
+
+ # rotate input variables image
+ if flip_lat: # reverse values along the latitude dimension
+ input_tensor = tf.reverse(input_tensor, axis=[0])
+ front_tensor = tf.reverse(front_tensor, axis=[0])
+ if flip_lon: # reverse values along the longitude dimension
+ input_tensor = tf.reverse(input_tensor, axis=[1])
+ front_tensor = tf.reverse(front_tensor, axis=[1])
+
+ # add salt and pepper noise to images
+ if args["noise_fraction"] > 0:
+ input_tensor = tf.where(
+ random_values < args["noise_fraction"] / 2,
+ 0.0,
+ input_tensor,
+ ) # add 0s to image
+ input_tensor = tf.where(
+ random_values > 1.0 - (args["noise_fraction"] / 2),
+ 1.0,
+ input_tensor,
+ ) # add 1s to image
+
+ # combine pressure level and variables dimensions, making the images 2D (excluding the final dimension)
+ if args["num_dims"][0] == 2:
+ input_tensor_shape_3d = input_tensor.shape
+ input_tensor = tf.reshape(
+ input_tensor,
+ [
+ input_tensor_shape_3d[0],
+ input_tensor_shape_3d[1],
+ input_tensor_shape_3d[2] * input_tensor_shape_3d[3],
+ ],
+ )
+
+ input_tensor_shapes.append(input_tensor.shape)
+ assert len(set(input_tensor_shapes)) == 1, (
+ f"ERROR: Attempted to add {input_tensor_shapes[-1]} to dataset with shape {input_tensor_shapes[0]}. "
+ "Please check your data for inconsistent coordinate systems."
+ )
+
+ # add input images to tensorflow dataset
+ input_tensor_for_timestep = tf.data.Dataset.from_tensors(
+ input_tensor
+ )
+ if len(input_tensors_for_slice) == 0:
+ input_tensors_for_slice = input_tensor_for_timestep
+ else:
+ input_tensors_for_slice = input_tensors_for_slice.concatenate(
+ input_tensor_for_timestep
+ )
+
+ front_tensor_shapes.append(front_tensor.shape)
+ assert len(set(front_tensor_shapes)) == 1, (
+ f"ERROR: Attempted to add {front_tensor_shapes[-1]} to dataset with shape {front_tensor_shapes[0]}. "
+ "Please check your data for inconsistent coordinate systems."
+ )
+
+ # if using 3D inputs, turn the fronts dataset into a 3D image
+ if args["num_dims"][1] == 3:
+ front_tensor = tf.tile(
+ front_tensor, (1, 1, len(args["pressure_levels"]))
+ )
+ else:
+ front_tensor = front_tensor[:, :, 0]
+
+ # add fronts to tensorflow dataset
+ front_tensor = tf.cast(
+ tf.one_hot(front_tensor, num_front_types), tf.float16
+ ) # One-hot encode the labels
+ front_tensor_for_timestep = tf.data.Dataset.from_tensors(
+ front_tensor
+ ) # convert fronts into a tensorflow dataset
+ if len(front_tensors_for_slice) == 0:
+ front_tensors_for_slice = front_tensor_for_timestep
+ else:
+ front_tensors_for_slice = front_tensors_for_slice.concatenate(
+ front_tensor_for_timestep
+ )
+
+ timesteps_kept += 1
+ else:
+ timesteps_discarded += 1
+
+ print(
+ f"{yr}: slice {i_slice}/{num_slices} (kept/discarded): ({timesteps_kept}/{timesteps_discarded} timesteps, {images_kept}/{images_discarded} images)",
+ end="\r",
+ )
+
+ tf_dataset_folder = (
+ f"{args['tf_outdir']}/{yr}-{i_slice}_tf" # output directory for the inputs
+ )
+ tf_dataset = tf.data.Dataset.zip(
+ (input_tensors_for_slice, front_tensors_for_slice)
+ )
+
+ # save the tensorflow datasets
+ try:
+ tf.data.Dataset.save(tf_dataset, path=tf_dataset_folder)
+ except NameError:
+ print("No images could be retained with the provided arguments.")
+
+ print(
+ f"Tensorflow datasets for {yr} ({num_slices} slices) saved to {args['tf_outdir']}."
+ )
diff --git a/src/fronts/data/download_era5.py b/src/fronts/data/download_era5.py
new file mode 100644
index 0000000..4e5cde9
--- /dev/null
+++ b/src/fronts/data/download_era5.py
@@ -0,0 +1,106 @@
+"""
+Download ERA5 data from the Climate Data Store.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import os
+import cdsapi
+import argparse
+import numpy as np
+
+VAR_ABBREV = {
+ "temperature": "T",
+ "geopotential": "Z",
+ "specific_humidity": "q",
+ "u_component_of_wind": "u",
+ "v_component_of_wind": "v",
+}
+
+DATA_FORMATS = {"grib": "grib", "netcdf": "nc"}
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--dataset", type=str, default="reanalysis-era5-pressure-levels"
+ )
+ parser.add_argument(
+ "--product_type", type=str, default="reanalysis", help="CDS product type."
+ )
+ parser.add_argument(
+ "--outdir",
+ type=str,
+ required=True,
+ help="Main output directory for the downloaded files.",
+ )
+ parser.add_argument(
+ "--pressure_level", type=str, required=True, help="Pressure levels of interest."
+ )
+ parser.add_argument(
+ "--variable",
+ type=str,
+ required=True,
+ help="Variable of interest. Must follow CDS nomenclature.",
+ )
+ parser.add_argument("--year", type=int, required=True, help="Year of interest.")
+ parser.add_argument(
+ "--month",
+ type=int,
+ nargs="+",
+ default=np.arange(1, 12.1).astype(int),
+ help="Months of interest.",
+ )
+ parser.add_argument(
+ "--day",
+ type=int,
+ nargs="+",
+ default=np.arange(0, 31.1).astype(int),
+ help="Days of interest.",
+ )
+ parser.add_argument(
+ "--hour",
+ type=int,
+ nargs="+",
+ default=np.arange(0, 21.1, 3).astype(int),
+ help="Hours of interest.",
+ )
+ parser.add_argument("--data_format", type=str, default="grib", help="Data format.")
+ parser.add_argument(
+ "--download_format", type=str, default="unarchived", help="Download format."
+ )
+ parser.add_argument(
+ "--overwrite", action="store_true", help="Overwrite any existing ERA5 files."
+ )
+ args = vars(parser.parse_args())
+
+ target = "%s/era5_%s-%shPa_%d.%s" % (
+ args["outdir"],
+ args["variable"],
+ args["pressure_level"],
+ args["year"],
+ DATA_FORMATS[args["data_format"]],
+ )
+
+ if not os.path.isfile(target) or args["overwrite"]:
+ request = {
+ "product_type": [args["product_type"]],
+ "variable": [args["variable"]],
+ "year": [args["year"]],
+ "month": [f"{mo:02d}" for mo in args["month"]],
+ "day": [f"{dy:02d}" for dy in args["day"]],
+ "time": [f"{hr:02d}:00" % hr for hr in args["hour"]],
+ "pressure_level": [args["pressure_level"]],
+ "data_format": args["data_format"],
+ "download_format": args["download_format"],
+ }
+
+ os.makedirs(args["outdir"], exist_ok=True)
+
+ client = cdsapi.Client()
+ client.retrieve(args["dataset"], request).download(target=target)
+
+ else:
+ print(
+ f"{target} already exists. Pass the --overwrite flag to initiate a new download that overwrites the file."
+ )
diff --git a/src/fronts/data/download_nwp.py b/src/fronts/data/download_nwp.py
new file mode 100644
index 0000000..ae2af51
--- /dev/null
+++ b/src/fronts/data/download_nwp.py
@@ -0,0 +1,265 @@
+"""
+Download grib files containing NWP model data.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import argparse
+import os
+import pandas as pd
+import requests
+import urllib.error
+
+# TODO: replace with requests
+import wget
+import sys
+import datetime
+
+
+def bar_progress(current, total, width=None):
+ progress_message = "Downloading %s: %d%% [%d/%d] MB " % (
+ local_filename,
+ current / total * 100,
+ current / 1e6,
+ total / 1e6,
+ )
+ sys.stdout.write("\r" + progress_message)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--outdir",
+ type=str,
+ required=True,
+ help="Output directory for GDAS grib files downloaded from NCEP.",
+ )
+ parser.add_argument(
+ "--model", type=str, required=True, help="NWP model to use as the data source."
+ )
+ parser.add_argument(
+ "--init_time",
+ type=str,
+ help="Initialization time of the model. Format: YYYY-MM-DD-HH.",
+ )
+ parser.add_argument(
+ "--range",
+ type=str,
+ nargs=3,
+ help="Download model data between a range of dates. Three arguments must be passed, with the first two arguments "
+ "marking the bounds of the date range in the format YYYY-MM-DD-HH. The third argument is the frequency (e.g. 6H), "
+ "which has the same formatting as the 'freq' keyword argument in pandas.date_range().",
+ )
+ parser.add_argument(
+ "--forecast_hours",
+ type=int,
+ nargs="+",
+ required=True,
+ help="List of forecast hours to download for the given day.",
+ )
+ parser.add_argument(
+ "--verbose",
+ action="store_true",
+ help="Include a progress bar for download progress.",
+ )
+
+ args = vars(parser.parse_args())
+
+ args["model"] = args["model"].lower()
+
+ # If --verbose is passed, include a progress bar to show the download progress
+ bar = bar_progress if args["verbose"] else None
+
+ if args["init_time"] is not None and args["range"] is not None:
+ raise ValueError(
+ "Only one of the following arguments can be passed: --init_time, --range"
+ )
+ elif args["init_time"] is None and args["range"] is None:
+ raise ValueError(
+ "One of the following arguments must be passed: --init_time, --range"
+ )
+
+ init_times = (
+ pd.date_range(args["init_time"], args["init_time"])
+ if args["init_time"] is not None
+ else pd.date_range(*args["range"][:2], freq=args["range"][-1])
+ )
+
+ files = [] # complete urls for the files to pull from AWS
+ local_filenames = [] # filenames for the local files after downloading
+
+ for init_time in init_times:
+ yr, mo, dy, hr = init_time.year, init_time.month, init_time.day, init_time.hour
+ if args["model"] == "gdas":
+ if datetime.datetime(yr, mo, dy, hr) < datetime.datetime(2015, 6, 23, 0):
+ raise ConnectionAbortedError(
+ "Cannot download GDAS data prior to June 23, 2015."
+ )
+ elif datetime.datetime(yr, mo, dy, hr) < datetime.datetime(2017, 7, 20, 0):
+ [
+ files.append(
+ "https://noaa-gfs-bdp-pds.s3.amazonaws.com/gdas.%d%02d%02d/%02d/gdas1.t%02dz.pgrb2.0p25.f%03d"
+ % (yr, mo, dy, hr, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ elif yr < 2021:
+ [
+ files.append(
+ "https://noaa-gfs-bdp-pds.s3.amazonaws.com/gdas.%d%02d%02d/%02d/gdas.t%02dz.pgrb2.0p25.f%03d"
+ % (yr, mo, dy, hr, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ else:
+ [
+ files.append(
+ "https://noaa-gfs-bdp-pds.s3.amazonaws.com/gdas.%d%02d%02d/%02d/atmos/gdas.t%02dz.pgrb2.0p25.f%03d"
+ % (yr, mo, dy, hr, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ elif "gefs" in args["model"]:
+ member = args["model"].split("-")[-1]
+ if datetime.datetime(yr, mo, dy, hr) < datetime.datetime(2017, 1, 1, 0):
+ raise ConnectionAbortedError(
+ "Cannot download GEFS data prior to January 1, 2017."
+ )
+ elif datetime.datetime(yr, mo, dy, hr) < datetime.datetime(2018, 7, 27, 0):
+ [
+ files.append(
+ "https://noaa-gefs-pds.s3.amazonaws.com/gefs.%d%02d%02d/%02d/ge%s.t%02dz.pgrb2af%03d"
+ % (yr, mo, dy, hr, member, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ elif datetime.datetime(yr, mo, dy, hr) < datetime.datetime(2020, 9, 24, 0):
+ [
+ files.append(
+ "https://noaa-gefs-pds.s3.amazonaws.com/gefs.%d%02d%02d/%02d/pgrb2a/ge%s.t%02dz.pgrb2af%02d"
+ % (yr, mo, dy, hr, member, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ else:
+ [
+ files.append(
+ "https://noaa-gefs-pds.s3.amazonaws.com/gefs.%d%02d%02d/%02d/atmos/pgrb2ap5/ge%s.t%02dz.pgrb2a.0p50.f%03d"
+ % (yr, mo, dy, hr, member, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ elif args["model"] == "gfs":
+ if datetime.datetime(yr, mo, dy, hr) < datetime.datetime(2021, 2, 26, 0):
+ raise ConnectionAbortedError(
+ "Cannot download GFS data prior to February 26, 2021."
+ )
+ elif datetime.datetime(yr, mo, dy, hr) < datetime.datetime(2021, 3, 22, 0):
+ [
+ files.append(
+ "https://noaa-gfs-bdp-pds.s3.amazonaws.com/gfs.%d%02d%02d/%02d/gfs.t%02dz.pgrb2.0p25.f%03d"
+ % (yr, mo, dy, hr, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ else:
+ [
+ files.append(
+ "https://noaa-gfs-bdp-pds.s3.amazonaws.com/gfs.%d%02d%02d/%02d/atmos/gfs.t%02dz.pgrb2.0p25.f%03d"
+ % (yr, mo, dy, hr, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ elif args["model"] == "hrrr":
+ if datetime.datetime(yr, mo, dy, hr) < datetime.datetime(2014, 7, 30, 18):
+ raise ConnectionAbortedError(
+ "Cannot download HRRR data prior to 18z July 30, 2014."
+ )
+ [
+ files.append(
+ "https://noaa-hrrr-bdp-pds.s3.amazonaws.com/hrrr.%d%02d%02d/conus/hrrr.t%02dz.wrfprsf%02d.grib2"
+ % (yr, mo, dy, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ elif args["model"] == "rap":
+ if datetime.datetime(yr, mo, dy, hr) < datetime.datetime(2021, 2, 22, 0):
+ raise ConnectionAbortedError(
+ "Cannot download RAP data prior to 0z February 22, 2021."
+ )
+ [
+ files.append(
+ "https://noaa-rap-pds.s3.amazonaws.com/rap.%d%02d%02d/rap.t%02dz.wrfprsf%02d.grib2"
+ % (yr, mo, dy, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ elif "ecmwf" in args["model"]:
+ ecmwf_model = args["model"].split("-")[1]
+ [
+ files.append(
+ "https://data.ecmwf.int/forecasts/%d%02d%02d/%02dz/%s/0p25/oper/%d%02d%02d%02d0000-%dh-oper-fc.grib2"
+ % (yr, mo, dy, hr, ecmwf_model, yr, mo, dy, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ elif "namnest" in args["model"]:
+ nest = args["model"].split("-")[-1]
+ [
+ files.append(
+ "https://nomads.ncep.noaa.gov/pub/data/nccf/com/nam/prod/nam.%d%02d%02d/nam.t%02dz.%snest.hiresf%02d.tm00.grib2"
+ % (yr, mo, dy, hr, nest, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+ elif args["model"] == "nam-12km":
+ for forecast_hour in args["forecast_hours"]:
+ if forecast_hour in [0, 1, 2, 3, 6]:
+ folder = "analysis" # use the analysis folder as it contains more accurate data
+ else:
+ folder = "forecast" # forecast hours other than 0, 1, 2, 3, 6 do not have analysis data
+
+ if datetime.datetime(
+ yr, mo, dy, hr
+ ) > datetime.datetime.utcnow() - datetime.timedelta(days=7):
+ files.append(
+ f"https://nomads.ncep.noaa.gov/pub/data/nccf/com/nam/prod/nam.%d%02d%02d/nam.t%02dz.awphys%02d.tm00.grib2"
+ % (yr, mo, dy, hr, forecast_hour)
+ )
+ else:
+ files.append(
+ f"https://www.ncei.noaa.gov/data/north-american-mesoscale-model/access/%s/%d%02d/%d%02d%02d/nam_218_%d%02d%02d_%02d00_%03d.grb2"
+ % (folder, yr, mo, yr, mo, dy, yr, mo, dy, hr, forecast_hour)
+ )
+ [
+ local_filenames.append(
+ "%s_%d%02d%02d%02d_f%03d.grib"
+ % (args["model"], yr, mo, dy, hr, forecast_hour)
+ )
+ for forecast_hour in args["forecast_hours"]
+ ]
+
+ for file, local_filename in zip(files, local_filenames):
+ timestring = local_filename.split("_")[1]
+ year, month = timestring[:4], timestring[4:6]
+ monthly_directory = f"{args['outdir']}/{year}{month}" # Directory for the grib files for the given days
+
+ ### If the directory does not exist, check to see if the file link is valid. If the file link is NOT valid, then the directory will not be created since it will be empty. ###
+ if not os.path.isdir(monthly_directory):
+ if (
+ requests.head(file).status_code == requests.codes.ok
+ or requests.head(file.replace("/atmos", "")).status_code
+ == requests.codes.ok
+ ):
+ os.makedirs(monthly_directory, exist_ok=True)
+
+ full_file_path = f"{monthly_directory}/{local_filename}"
+
+ if not os.path.isfile(full_file_path):
+ try:
+ wget.download(file, out=full_file_path, bar=bar)
+ except urllib.error.HTTPError:
+ print("Error downloading %s" % file)
+ else:
+ print("%s already exists, skipping file...." % full_file_path)
diff --git a/src/fronts/data/download_satellite.py b/src/fronts/data/download_satellite.py
new file mode 100644
index 0000000..6216558
--- /dev/null
+++ b/src/fronts/data/download_satellite.py
@@ -0,0 +1,204 @@
+"""
+Download netCDF files containing GOES satellite data.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2024.12.20
+"""
+
+import argparse
+import os
+import pandas as pd
+import requests
+
+# TODO: replace with obstore
+import s3fs
+import sys
+import urllib.error
+
+# TODO: replace with requests
+import wget
+
+
+def bar_progress(current, total, width=None):
+ progress_message = "Downloading %s: %d%% [%d/%d] MB " % (
+ local_filename,
+ current / total * 100,
+ current / 1e6,
+ total / 1e6,
+ )
+ sys.stdout.write("\r" + progress_message)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--netcdf_outdir",
+ type=str,
+ required=True,
+ help="Output directory for the satellite netCDF files.",
+ )
+ parser.add_argument(
+ "--satellite",
+ type=str,
+ default="goes16",
+ help="Satellite source. Options are 'goes16', 'goes17', 'goes18', 'MERGIR'.",
+ )
+ parser.add_argument(
+ "--domain",
+ type=str,
+ default="full-disk",
+ help="Domain of the satellite data. Options are 'full-disk', 'conus', 'meso'.",
+ )
+ parser.add_argument(
+ "--product",
+ type=str,
+ default="ABI-L2-MCMIP",
+ help="Satellite product to download.",
+ )
+ parser.add_argument(
+ "--init_time",
+ type=str,
+ help="Initialization time for which to search for satellite data. Format: YYYY-MM-DD-HH.",
+ )
+ parser.add_argument(
+ "--range",
+ type=str,
+ nargs=3,
+ help="Download satellite data between a range of dates. Three arguments must be passed, with the first two arguments "
+ "marking the bounds of the date range in the format YYYY-MM-DD-HH. The third argument is the frequency (e.g. 6H), "
+ "which has the same formatting as the 'freq' keyword argument in pandas.date_range().",
+ )
+ parser.add_argument(
+ "--verbose",
+ action="store_true",
+ help="Include a progress bar for download progress.",
+ )
+ args = vars(parser.parse_args())
+
+ init_times = (
+ pd.date_range(args["init_time"], args["init_time"])
+ if args["init_time"] is not None
+ else pd.date_range(*args["range"][:2], freq=args["range"][-1])
+ )
+
+ files = [] # complete urls for the files to pull from AWS
+ local_filenames = [] # filenames for the local files after downloading
+
+ if args["satellite"] not in ["goes16", "goes17", "goes18", "MERGIR"]:
+ raise ValueError(
+ "'%s' is not a valid satellite source. Options are 'goes16', 'goes17', 'goes18', 'MERGIR'."
+ % args["satellite"]
+ )
+
+ if args["domain"] == "full-disk":
+ domain_str = "F"
+ elif args["domain"] == "conus":
+ domain_str = "C"
+ elif args["domain"] == "meso":
+ domain_str = "M"
+ else:
+ raise ValueError(
+ "'%s' is not a valid domain. Options are 'full-disk', 'conus', 'meso'."
+ % args["domain"]
+ )
+
+ files = [] # complete urls for the files to pull from AWS
+ local_filenames = [] # filenames for the local files after downloading
+
+ if args["satellite"] != "MERGIR":
+ fs = s3fs.S3FileSystem(anon=True)
+ else:
+ args["domain"] = (
+ "global" # force override the domain if downloading merged IR brightness temperature
+ )
+
+ for i, init_time in enumerate(init_times):
+ day_of_year = init_time.timetuple().tm_yday
+ yr, mo, dy, hr = init_time.year, init_time.month, init_time.day, init_time.hour
+
+ if args["satellite"] != "MERGIR":
+ s3_folder = "s3://noaa-%s/%s%s/%d/%03d/%02d" % (
+ args["satellite"],
+ args["product"],
+ domain_str,
+ yr,
+ day_of_year,
+ hr,
+ )
+ s3_timestring = "%d%03d%02d" % (yr, day_of_year, hr)
+ s3_glob_str = s3_folder + "/*s" + s3_timestring + "*.nc"
+
+ try:
+ files.append(
+ list(sorted(fs.glob(s3_glob_str)))[0].replace(
+ "noaa-%s" % args["satellite"],
+ "https://noaa-%s.s3.amazonaws.com" % args["satellite"],
+ )
+ )
+ except IndexError:
+ print(
+ "NO DATA FOUND (ind %d): satellite=[%s] domain=[%s] product=[%s] init_time=[%s] "
+ % (i, args["satellite"], args["domain"], args["product"], init_time)
+ )
+ continue
+ else:
+ print(
+ "Gathering data: satellite=[%s] domain=[%s] product=[%s] init_time=[%s]"
+ % (args["satellite"], args["domain"], args["product"], init_time),
+ end="\r",
+ )
+
+ else:
+ files.append(
+ "https://disc2.gesdisc.eosdis.nasa.gov/data/MERGED_IR/GPM_MERGIR.1/%d/%03d/merg_%d%02d%02d%02d_4km-pixel.nc4"
+ % (yr, day_of_year, yr, mo, dy, hr)
+ )
+
+ local_filenames.append(
+ "%s_%d%02d%02d%02d_%s.nc"
+ % (args["satellite"], yr, mo, dy, hr, args["domain"])
+ )
+
+ # If --verbose is passed, include a progress bar to show the download progress
+ bar = bar_progress if args["verbose"] else None
+
+ if args["init_time"] is not None and args["range"] is not None:
+ raise ValueError(
+ "Only one of the following arguments can be passed: --init_time, --range"
+ )
+ elif args["init_time"] is None and args["range"] is None:
+ raise ValueError(
+ "One of the following arguments must be passed: --init_time, --range"
+ )
+
+ for file, local_filename in zip(files, local_filenames):
+ timestring = local_filename.split("_")[1]
+ year, month = timestring[:4], timestring[4:6]
+ monthly_directory = "%s/%s%s" % (
+ args["netcdf_outdir"],
+ year,
+ month,
+ ) # Directory for the netCDF files for the given days
+
+ ### If the directory does not exist, check to see if the file link is valid. If the file link is NOT valid, then the directory will not be created since it will be empty. ###
+ if not os.path.isdir(monthly_directory):
+ if requests.head(file).status_code == requests.codes.ok:
+ os.mkdir(monthly_directory)
+
+ full_file_path = f"{monthly_directory}/{local_filename}"
+
+ if not os.path.isfile(full_file_path):
+ if args["satellite"] != "MERGIR":
+ try:
+ wget.download(file, out=full_file_path, bar=bar)
+ except urllib.error.HTTPError:
+ print(f"Error downloading {file}")
+ else:
+ print("Downloading", file)
+ result = requests.get(file)
+ result.raise_for_status()
+ f = open(full_file_path, "wb")
+ f.write(result.content)
+ f.close()
+ else:
+ print(f"{full_file_path} already exists, skipping file....")
diff --git a/src/fronts/data/era5.py b/src/fronts/data/era5.py
new file mode 100644
index 0000000..6c6dd3c
--- /dev/null
+++ b/src/fronts/data/era5.py
@@ -0,0 +1,229 @@
+import xarray as xr
+
+import datetime
+from collections import namedtuple
+import dataclasses
+from fronts.utils import calc
+from typing import Callable
+
+ARCO_ERA5_GCP_URI = (
+ "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
+)
+
+BoundingBox = namedtuple("BoundingBox", ["lat_min", "lat_max", "lon_min", "lon_max"])
+
+
+def convert_domain_extent_to_bounding_box(domain_extent: list[float]) -> BoundingBox:
+ """Converts a domain extent from constants.py to a BoundingBox namedtuple.
+
+ Args:
+ domain_extent: A list of four floats representing the domain extent in the
+ format [lat_min, lat_max, lon_min, lon_max].
+
+ Returns a BoundingBox named tuple with the corresponding values.
+ """
+ if len(domain_extent) != 4:
+ raise ValueError("Domain extent must be a list of four floats.")
+ return BoundingBox(
+ lon_min=domain_extent[0],
+ lon_max=domain_extent[1],
+ lat_min=domain_extent[2],
+ lat_max=domain_extent[3],
+ )
+
+
+def load_arco_era5(
+ store: str = ARCO_ERA5_GCP_URI,
+ chunks: dict[str, int] = {"time": 48},
+ consolidated: bool = True,
+):
+ """Opens the Google ARCO ERA5 analysis-ready dataset as an xarray Dataset.
+
+ Args:
+ store: The URI of the zarr store to open. Defaults to the Google ARCO ERA5
+ analysis-ready dataset link.
+ chunks: The chunk sizes to use when opening the dataset. Defaults to chunking
+ the time dimension into 48-hour chunks.
+ consolidated: Whether to use consolidated metadata when opening the dataset.
+ Defaults to True.
+
+ Returns an xarray Dataset containing the ERA5 analysis-ready data.
+ """
+ era5_ds = xr.open_zarr(
+ store=store,
+ chunks=chunks,
+ consolidated=consolidated,
+ )
+
+ return era5_ds
+
+
+def subset_arco_era5(
+ ds: xr.Dataset,
+ variables: list[str],
+ start_date: datetime.datetime,
+ end_date: datetime.datetime,
+ bounding_box: BoundingBox,
+ levels: list[int],
+):
+ """Subsets the ARCO ERA5 dataset by variables, specific time range, and geographic bounding box.
+
+ Args:
+ ds: The input xarray Dataset containing the ARCO ERA5 data.
+ variables: A list of variable names to subset from the dataset.
+ start_date: The start date of the time range to subset (inclusive).
+ end_date: The end date of the time range to subset (inclusive).
+ bounding_box: A BoundingBox named tuple defining the geographic bounding box
+ for subsetting. Defaults to a bounding box covering the contiguous United
+ States.
+ levels: A list of pressure levels to subset from the dataset.
+ """
+ variables_to_postprocess = [var for var in variables if var not in ds.data_vars]
+ if variables_to_postprocess:
+ variables.pop(variables_to_postprocess)
+
+ if not any(
+ [n for n in variables_to_postprocess if n in calc.callable_mapping.keys()]
+ ):
+ raise ValueError(
+ f"Variables {variables_to_postprocess} not found in dataset and no "
+ "post-processing functions available for them."
+ )
+ ds = ds[variables]
+ ds = ds.sel(
+ latitude=slice(bounding_box.lat_max, bounding_box.lat_min),
+ longitude=slice(bounding_box.lon_min, bounding_box.lon_max),
+ )
+ ds = ds.sel(time=slice(start_date, end_date))
+ ds = ds.sel(level=levels)
+ return ds
+
+
+@dataclasses.dataclass
+class ERA5TrainingDataConfig:
+ """A dataclass for generating data from the ARCO ERA5 dataset.
+
+ This class provides methods for loading and subsetting the variables, spatial
+ bounds, and time of the ARCO ERA5 dataset.
+
+ Attributes:
+ domain_extent: A list of four floats representing the geographic domain extent
+ in the format [lon_min, lon_max, lat_min, lat_max].
+ variables: A list of variable names to subset from the dataset.
+ start_date: The start date of the time range to subset (inclusive).
+ end_date: The end date of the time range to subset (inclusive).
+ store: The URI of the zarr store to open.
+ chunks: The chunk sizes to use when opening the dataset.
+ consolidated: Whether to use consolidated metadata when opening the dataset.
+ """
+
+ domain_extent: list[float]
+ variables: list[str]
+ start_date: datetime.datetime
+ end_date: datetime.datetime
+ levels: list[int]
+ store: str
+ chunks: dict[str, int]
+ consolidated: bool
+
+ def build(self) -> xr.Dataset:
+ """Builds the training dataset by loading and subsetting the ARCO ERA5 dataset.
+
+ Returns an xarray Dataset containing the subset ARCO ERA5 data.
+ """
+ # Load the ARCO ERA5 dataset with default params
+ ds = load_arco_era5(
+ store=self.store, chunks=self.chunks, consolidated=self.consolidated
+ )
+
+ # Subset the dataset by variables, time range, and geographic bounding box
+ ds = subset_arco_era5(
+ ds,
+ variables=self.variables,
+ start_date=self.start_date,
+ end_date=self.end_date,
+ bounding_box=convert_domain_extent_to_bounding_box(self.domain_extent),
+ levels=self.levels,
+ )
+ return ds
+
+
+def _default_postprocess(ds: xr.Dataset):
+ """Default postprocessor that passes through data unmodified."""
+ return ds
+
+
+def maybe_postprocess_era5(
+ ds: xr.Dataset, postprocess_func: Callable = _default_postprocess, **kwargs
+) -> xr.Dataset:
+ """Applies any necessary post-processing steps to the ERA5 dataset.
+
+ This function is a placeholder for any future post-processing steps that may be
+ required for the ERA5 dataset. Currently, it returns the dataset unchanged.
+
+ Args:
+ ds: The input xarray Dataset containing the ERA5 data.
+ postprocess_func: A callable function that takes an xarray Dataset as input.
+ Defaults to a no-op function that returns the dataset unchanged.
+ **kwargs: Additional keyword arguments to pass to the post-processing function.
+
+ Returns the possibly post-processed Dataset.
+ """
+ ds = postprocess_func(ds, **kwargs)
+ return ds
+
+
+def dewpoint_postprocessor(ds: xr.Dataset):
+ ds["dewpoint"] = calc.dewpoint_from_specific_humidity(
+ ds.level, ds.specific_humidity
+ )
+ return ds
+
+
+def potential_temperature_postprocessor(ds: xr.Dataset):
+ ds["potential_temperature"] = calc.potential_temperature(ds.level, ds.temperature)
+ return ds
+
+
+def equivalent_potential_temperature_postprocessor(ds: xr.Dataset):
+ ds["equivalent_potential_temperature"] = calc.equivalent_potential_temperature(
+ ds.level, ds.temperature, ds.dewpoint
+ )
+ return ds
+
+
+def virtual_potential_temperature_postprocessor(ds: xr.Dataset):
+ ds["virtual_potential_temperature"] = calc.virtual_potential_temperature(
+ ds.level, ds.temperature, ds.dewpoint
+ )
+ return ds
+
+
+def wet_bulb_temperature_postprocessor(ds: xr.Dataset):
+ ds["wet_bulb_temperature"] = calc.wet_bulb_temperature(ds.temperature, ds.dewpoint)
+ return ds
+
+
+def wet_bulb_potential_temperature_postprocessor(ds: xr.Dataset):
+ ds["wet_bulb_potential_temperature"] = calc.wet_bulb_potential_temperature(
+ ds.level, ds.temperature, ds.dewpoint
+ )
+ return ds
+
+
+def relative_humidity_postprocessor(ds: xr.Dataset):
+ ds["relative_humidity"] = calc.relative_humidity_from_dewpoint(
+ ds.temperature, ds.dewpoint
+ )
+ return ds
+
+
+callable_mapping = {
+ "dewpoint": dewpoint_postprocessor,
+ "potential_temperature": potential_temperature_postprocessor,
+ "equivalent_potential_temperature": equivalent_potential_temperature_postprocessor,
+ "virtual_potential_temperature": virtual_potential_temperature_postprocessor,
+ "wet_bulb_temperature": wet_bulb_temperature_postprocessor,
+ "wet_bulb_potential_temperature": wet_bulb_potential_temperature_postprocessor,
+ "relative_humidity": relative_humidity_postprocessor,
+}
diff --git a/src/fronts/evaluation/__init__.py b/src/fronts/evaluation/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/fronts/evaluation/calibrate_model.py b/src/fronts/evaluation/calibrate_model.py
new file mode 100644
index 0000000..f13053a
--- /dev/null
+++ b/src/fronts/evaluation/calibrate_model.py
@@ -0,0 +1,200 @@
+"""
+Calibrate a trained model.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.2.2
+"""
+
+import argparse
+import pandas as pd
+from fronts.utils.data_utils import FRONT_NAMES
+import matplotlib.pyplot as plt
+import pickle
+import xarray as xr
+import numpy as np
+from sklearn.isotonic import IsotonicRegression
+from sklearn.metrics import r2_score
+from glob import glob
+
+
+if __name__ == "__main__":
+ """
+ All arguments listed in the examples are listed via argparse in alphabetical order below this comment block.
+ """
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--dataset",
+ type=str,
+ help="Dataset for which calibration will be performed. Options are 'training', 'validation', 'test'.",
+ )
+ parser.add_argument(
+ "--domain", type=str, help="Domain for which calibration will be performed."
+ )
+ parser.add_argument(
+ "--model_dir", type=str, help="Parent directory for the model(s)."
+ )
+ parser.add_argument("--model_number", type=int, help="Model number.")
+ parser.add_argument(
+ "--data_source", type=str, default="era5", help="Data source for the variables."
+ )
+
+ args = vars(parser.parse_args())
+
+ model_properties = pd.read_pickle(
+ "%s/model_%d/model_%d_properties.pkl"
+ % (args["model_dir"], args["model_number"], args["model_number"])
+ )
+
+ ### front_types argument is being moved into the dataset_properties dictionary within model_properties ###
+ try:
+ front_types = model_properties["front_types"]
+ except KeyError:
+ front_types = model_properties["dataset_properties"]["front_types"]
+
+ if type(front_types) == str:
+ front_types = [
+ front_types,
+ ]
+
+ try:
+ _ = model_properties[
+ "calibration_models"
+ ] # Check to see if the model has already been calibrated before
+ except KeyError:
+ model_properties["calibration_models"] = dict()
+
+ model_properties["calibration_models"][args["domain"]] = dict()
+
+ temporal_files = list(
+ sorted(
+ glob(
+ "%s/model_%d/statistics/model_%d_statistics_%s_*_temporal.nc"
+ % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ args["domain"],
+ )
+ )
+ )
+ )
+ temporal_ds = xr.open_mfdataset(temporal_files, combine="nested", concat_dim="time")
+
+ axis_ticks = np.arange(0.1, 1.1, 0.1)
+
+ for front_label in front_types:
+ model_properties["calibration_models"][args["domain"]][front_label] = dict()
+
+ true_positives = temporal_ds[f"tp_temporal_{front_label}"].values
+ false_positives = temporal_ds[f"fp_temporal_{front_label}"].values
+
+ thresholds = temporal_ds["threshold"].values
+
+ ### Sum the true positives along the 'time' axis ###
+ true_positives_sum = np.sum(true_positives, axis=0)
+ false_positives_sum = np.sum(false_positives, axis=0)
+
+ ### Find the number of true positives and false positives in each probability bin ###
+ true_positives_diff = np.abs(np.diff(true_positives_sum))
+ false_positives_diff = np.abs(np.diff(false_positives_sum))
+ observed_relative_frequency = np.divide(
+ true_positives_diff, true_positives_diff + false_positives_diff
+ )
+
+ boundary_colors = ["red", "purple", "brown", "darkorange", "darkgreen"]
+
+ calibrated_probabilities = []
+
+ fig, axs = plt.subplots(1, 2, figsize=(14, 6))
+ axs[0].plot(
+ thresholds,
+ thresholds,
+ color="black",
+ linestyle="--",
+ linewidth=0.5,
+ label="Perfect Reliability",
+ )
+
+ for boundary, color in enumerate(boundary_colors):
+ ####################### Test different calibration methods to see which performs best ######################
+
+ x = [
+ threshold
+ for threshold, frequency in zip(
+ thresholds[1:], observed_relative_frequency[boundary]
+ )
+ if not np.isnan(frequency)
+ ]
+ y = [
+ frequency
+ for threshold, frequency in zip(
+ thresholds[1:], observed_relative_frequency[boundary]
+ )
+ if not np.isnan(frequency)
+ ]
+
+ x.append(1.0)
+ y.append(1.0)
+
+ ### Isotonic Regression ###
+ ir = IsotonicRegression(out_of_bounds="clip")
+ ir.fit_transform(x, y)
+ calibrated_probabilities.append(ir.predict(x))
+ r_squared = r2_score(y, calibrated_probabilities[boundary])
+
+ axs[0].plot(
+ x[:-1],
+ y[:-1],
+ color=color,
+ linewidth=1,
+ label="%d km" % ((boundary + 1) * 50),
+ )
+ axs[1].plot(
+ x,
+ calibrated_probabilities[boundary],
+ color=color,
+ linestyle="--",
+ linewidth=1,
+ label=r"%d km ($R^2$ = %.3f)" % ((boundary + 1) * 50, r_squared),
+ )
+ model_properties["calibration_models"][args["domain"]][front_label][
+ "%d km" % ((boundary + 1) * 50)
+ ] = ir
+
+ for ax in axs:
+ axs[0].set_xlabel("Forecast Probability (uncalibrated)")
+ ax.set_xticks(axis_ticks)
+ ax.set_yticks(axis_ticks)
+ ax.set_xlim(0, 1)
+ ax.set_ylim(0, 1)
+ ax.grid()
+ ax.legend()
+
+ axs[0].set_title("Reliability Diagram")
+ axs[1].set_title("Calibration (isotonic regression)")
+ axs[0].set_ylabel("Observed Relative Frequency")
+ axs[1].set_ylabel("Forecast Probability (calibrated)")
+
+ with open(
+ "%s/model_%d/model_%d_properties.pkl"
+ % (args["model_dir"], args["model_number"], args["model_number"]),
+ "wb",
+ ) as f:
+ pickle.dump(model_properties, f)
+
+ plt.suptitle(
+ f"Model {args['model_number']} reliability/calibration: {FRONT_NAMES[front_label]}"
+ )
+ plt.savefig(
+ f"%s/model_%d/model_%d_calibration_%s_%s.png"
+ % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ args["domain"],
+ front_label,
+ ),
+ bbox_inches="tight",
+ dpi=300,
+ )
+ plt.close()
diff --git a/src/fronts/evaluation/generate_bincounts.py b/src/fronts/evaluation/generate_bincounts.py
new file mode 100644
index 0000000..eba5f8f
--- /dev/null
+++ b/src/fronts/evaluation/generate_bincounts.py
@@ -0,0 +1,241 @@
+"""
+Generate bincounts for ERA5 variables. The bincounts generated across billions of points can allow us to get parameters
+from which we can normalize the data.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import os
+import numpy as np
+import xarray as xr
+from glob import glob
+import argparse
+from datetime import datetime
+
+
+TRANSPOSE_DIMS = {
+ "MERGIR": ("lat", "lon", "time"),
+ "era5": ("latitude", "longitude", "pressure_level", "time"),
+ "goes-merged": ("latitude", "longitude", "time"),
+}
+
+BINS = {
+ "T": np.arange(100, 400.1),
+ "Td": np.arange(100, 400.1),
+ "Tv": np.arange(100, 400.1),
+ "Tw": np.arange(100, 400.1),
+ "theta": np.arange(100, 500.1),
+ "theta_e": np.arange(100, 500.1),
+ "theta_v": np.arange(100, 500.1),
+ "theta_w": np.arange(100, 500.1),
+ "RH": np.arange(0, 1.001, 0.0025),
+ "r": np.arange(0, 50.01, 0.125),
+ "q": np.arange(0, 50.01, 0.125),
+ "u": np.arange(-80, 150, 0.4),
+ "v": np.arange(-80, 150, 0.4),
+ "band_1": np.arange(0, 1.001, 0.0025),
+ "band_2": np.arange(0, 1.001, 0.0025),
+ "band_3": np.arange(0, 1.001, 0.0025),
+ "band_4": np.arange(0, 1.001, 0.0025),
+ "band_5": np.arange(0, 1.001, 0.0025),
+ "band_6": np.arange(0, 1.001, 0.0025),
+ "band_7": np.arange(150, 400.01, 0.5),
+ "band_8": np.arange(150, 400.01, 0.5),
+ "band_9": np.arange(150, 400.01, 0.5),
+ "band_10": np.arange(150, 400.01, 0.5),
+ "band_11": np.arange(150, 400.01, 0.5),
+ "band_12": np.arange(150, 400.01, 0.5),
+ "band_13": np.arange(150, 400.01, 0.5),
+ "band_14": np.arange(150, 400.01, 0.5),
+ "band_15": np.arange(150, 400.01, 0.5),
+ "band_16": np.arange(150, 400.01, 0.5),
+}
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--netcdf_indir",
+ type=str,
+ required=True,
+ help="Base directory for the input netcdf files.",
+ )
+ parser.add_argument(
+ "--bincount_outdir",
+ type=str,
+ required=True,
+ help="Output directory for the netcdf files containing bincounts.",
+ )
+ parser.add_argument("--variables", type=str, nargs="+")
+ parser.add_argument(
+ "--data_source",
+ type=str,
+ default="era5",
+ help='Data source (e.g., "era5", "MERGIR", etc.)',
+ )
+ parser.add_argument("--year_and_month", type=int, nargs=2, required=True)
+ parser.add_argument(
+ "--overwrite", action="store_true", help="Overwrite existing datasets."
+ )
+ args = vars(parser.parse_args())
+
+ year, month = args["year_and_month"]
+ transpose_dims = TRANSPOSE_DIMS[args["data_source"]]
+
+ netcdf_files = list(
+ sorted(
+ glob(
+ args["netcdf_indir"]
+ + "/%d%02d/%s*_%d%02d*.nc"
+ % (year, month, args["data_source"], year, month)
+ )
+ )
+ )
+
+ print(f"{datetime.utcnow()}: Opening dataset")
+ ds = xr.open_mfdataset(
+ netcdf_files,
+ combine="nested",
+ concat_dim="time",
+ engine="h5netcdf",
+ chunks="auto",
+ )
+ ds = ds.transpose(*transpose_dims).astype("float32")
+ print(f"{datetime.utcnow()}: Dataset loaded")
+ if args["variables"] is None:
+ args["variables"] = list(ds.keys())
+
+ if args["data_source"] == "era5":
+ # for pressure_level in ds['pressure_level'].values:
+ for pressure_level in [
+ "300",
+ ]:
+ for var_str in args["variables"]:
+ if var_str == "sp_z":
+ if pressure_level == "surface":
+ bins = np.arange(450.0, 1080.1)
+ elif pressure_level == "700":
+ bins = np.arange(150.0, 450.1)
+ elif pressure_level == "500":
+ bins = np.arange(300.0, 700.1)
+ elif pressure_level == "300":
+ bins = np.arange(700.0, 1400.1)
+ else:
+ bins = BINS[var_str]
+
+ # if pressure level is read as a float, convert it to an integer
+ pressure_level = (
+ int(pressure_level)
+ if isinstance(pressure_level, float)
+ else pressure_level
+ )
+
+ N_bins = len(bins) - 1
+
+ print("%d-%02d: (%s_%s)" % (year, month, var_str, pressure_level))
+ save_var_str = "%s_%s" % (var_str, pressure_level)
+
+ output_folder = "%s/%s_bins" % (args["bincount_outdir"], save_var_str)
+ bincount_file = "%s/%s_bincounts_%d%02d.nc" % (
+ output_folder,
+ save_var_str,
+ year,
+ month,
+ )
+ bincount_ds_exists = os.path.isfile(bincount_file)
+ if bincount_ds_exists and not args["overwrite"]:
+ print(
+ "%s already exists. If you want to overwrite the existing dataset, rerun this script with the "
+ "--overwrite flag attached." % bincount_file
+ )
+ continue
+
+ da = ds.sel(pressure_level=pressure_level)[var_str]
+ lat = da[transpose_dims[0]].values
+ time_array = da["time"].values
+ Ntime = len(time_array)
+ Nlon = len(da[transpose_dims[1]])
+ Nlat = len(lat)
+ bincounts_by_latitude = np.zeros([Nlat, N_bins], dtype=np.int64)
+
+ var = da.to_numpy().reshape((Nlat, Nlon * Ntime))
+ var = np.nan_to_num(var)
+ for ilat in range(Nlat):
+ bincounts_by_latitude[ilat, :] += np.histogram(var[ilat, :], bins)[
+ 0
+ ]
+ del var
+
+ os.makedirs(output_folder, exist_ok=True)
+
+ print("saving dataset")
+ bincount_ds = xr.Dataset(
+ data_vars={
+ "%s_bincount" % save_var_str: (
+ ("latitude", "bin"),
+ bincounts_by_latitude,
+ )
+ },
+ coords={"latitude": lat, "bin": bins[:-1]},
+ )
+ bincount_ds = bincount_ds.expand_dims(
+ {"time": np.array(["%d-%02d" % (year, month)], dtype=datetime)}
+ )
+ bincount_ds.to_netcdf(bincount_file, engine="netcdf4", mode="w")
+ bincount_ds.close()
+
+ ds.close()
+
+ elif args["data_source"] == "goes-merged":
+ for var_str in args["variables"]:
+ bins = BINS[var_str]
+ N_bins = len(bins) - 1
+ print("%d-%02d: (%s)" % (year, month, var_str))
+ save_var_str = var_str
+
+ output_folder = "%s/%s_bins" % (args["bincount_outdir"], save_var_str)
+ bincount_file = "%s/%s_bincounts_%d%02d.nc" % (
+ output_folder,
+ save_var_str,
+ year,
+ month,
+ )
+ bincount_ds_exists = os.path.isfile(bincount_file)
+ if bincount_ds_exists and not args["overwrite"]:
+ print(
+ "%s already exists. If you want to overwrite the existing dataset, rerun this script with the "
+ "--overwrite flag attached." % bincount_file
+ )
+ continue
+
+ da = ds[var_str]
+ lat = da[transpose_dims[0]].values
+ time_array = da["time"].values
+ Ntime = len(time_array)
+ Nlon = len(da[transpose_dims[1]])
+ Nlat = len(lat)
+ bincounts_by_latitude = np.zeros([Nlat, N_bins], dtype=np.int64)
+
+ var = da.to_numpy().reshape((Nlat, Nlon * Ntime))
+ var = np.nan_to_num(var, nan=-99999)
+ for ilat in range(Nlat):
+ bincounts_by_latitude[ilat, :] += np.histogram(var[ilat, :], bins)[0]
+ del var
+
+ os.makedirs(output_folder, exist_ok=True)
+
+ print("saving dataset")
+ bincount_ds = xr.Dataset(
+ data_vars={
+ "%s_bincount" % save_var_str: (
+ ("latitude", "bin"),
+ bincounts_by_latitude,
+ )
+ },
+ coords={"latitude": lat, "bin": bins[:-1]},
+ )
+ bincount_ds = bincount_ds.expand_dims(
+ {"time": np.array(["%d-%02d" % (year, month)], dtype=datetime.datetime)}
+ )
+ bincount_ds.to_netcdf(bincount_file, engine="netcdf4", mode="w")
+ bincount_ds.close()
diff --git a/src/fronts/evaluation/generate_performance_stats.py b/src/fronts/evaluation/generate_performance_stats.py
new file mode 100644
index 0000000..94a927d
--- /dev/null
+++ b/src/fronts/evaluation/generate_performance_stats.py
@@ -0,0 +1,421 @@
+"""
+Generate performance statistics for a model.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.2.16
+"""
+
+import argparse
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+import xarray as xr
+import os
+from fronts.utils import data_utils
+from fronts.utils.data_utils import DOMAIN_EXTENTS
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--model_dir", type=str, required=True, help="Directory for the models."
+ )
+ parser.add_argument("--model_number", type=int, required=True, help="Model number.")
+ parser.add_argument(
+ "--tf_indir",
+ type=str,
+ help="Directory for the TensorFlow dataset used for model evaluation.",
+ )
+ parser.add_argument(
+ "--dataset",
+ type=str,
+ help="Dataset for which to make predictions. Options are: 'training', 'validation', 'test'",
+ )
+ parser.add_argument(
+ "--year_and_month",
+ type=int,
+ nargs=2,
+ help="Year and month for which to make predictions.",
+ )
+ parser.add_argument("--domain", type=str, help="Domain of the data.")
+ parser.add_argument("--gpu_device", type=int, nargs="+", help="GPU device number.")
+ parser.add_argument(
+ "--memory_growth", action="store_true", help="Use memory growth on the GPU"
+ )
+ parser.add_argument(
+ "--overwrite",
+ action="store_true",
+ help="Overwrite any existing statistics files.",
+ )
+ args = vars(parser.parse_args())
+
+ model_properties = pd.read_pickle(
+ "%s/model_%d/model_%d_properties.pkl"
+ % (args["model_dir"], args["model_number"], args["model_number"])
+ )
+ domain = args["domain"]
+
+ variables = model_properties["dataset_properties"]["variables"]
+
+ # Some older models do not have the 'dataset_properties' dictionary
+ try:
+ front_types = model_properties["dataset_properties"]["front_types"]
+ num_dims = model_properties["dataset_properties"]["num_dims"]
+ except KeyError:
+ front_types = model_properties["front_types"]
+ if args["model_number"] in [6846496, 7236500, 7507525]:
+ num_dims = (3, 3)
+
+ num_front_types = (
+ model_properties["classes"] - 1
+ ) # remove the "no front" class type
+
+ if args["dataset"] is not None and args["year_and_month"] is not None:
+ raise ValueError("--dataset and --year_and_month cannot be passed together.")
+ elif args["dataset"] is None and args["year_and_month"] is None:
+ raise ValueError(
+ "At least one of [--dataset, --year_and_month] must be passed."
+ )
+ elif args["year_and_month"] is not None:
+ years, months = [args["year_and_month"][0]], [args["year_and_month"][1]]
+ else:
+ years, months = model_properties["%s_years" % args["dataset"]], range(1, 13)
+
+ if args["gpu_device"] is not None:
+ gpus = tf.config.list_physical_devices(device_type="GPU")
+ tf.config.set_visible_devices(
+ devices=[gpus[gpu] for gpu in args["gpu_device"]], device_type="GPU"
+ )
+
+ # Allow for memory growth on the GPU. This will only use the GPU memory that is required rather than allocating all the GPU's memory.
+ if args["memory_growth"]:
+ tf.config.experimental.set_memory_growth(
+ device=[gpus[gpu] for gpu in args["gpu_device"]][0], enable=True
+ )
+
+ for year in years:
+ for month in months:
+ front_files_month = pd.read_pickle(
+ "%s/front_files_%d%02d.pkl" % (args["tf_indir"], year, month)
+ )
+
+ if domain != "conus":
+ for front_file in front_files_month[::-1]:
+ if any(
+ [
+ "%02d_full.nc" % hour in front_file
+ for hour in np.arange(3, 21.1, 6)
+ ]
+ ):
+ front_files_month.pop(front_files_month.index(front_file))
+
+ prediction_file = (
+ f"%s/model_%d/probabilities/model_%d_pred_%s_%d%02d.nc"
+ % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ args["domain"],
+ year,
+ month,
+ )
+ )
+
+ stats_dataset_path = (
+ "%s/model_%d/statistics/model_%d_statistics_%s_%d%02d.nc"
+ % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ args["domain"],
+ year,
+ month,
+ )
+ )
+ if os.path.isfile(stats_dataset_path) and not args["overwrite"]:
+ print(
+ "WARNING: %s exists, pass the --overwrite argument to overwrite existing data."
+ % stats_dataset_path
+ )
+ continue
+
+ probs_ds = xr.open_dataset(prediction_file)
+ lons = probs_ds["longitude"].values
+ lats = probs_ds["latitude"].values
+
+ try:
+ custom_extent = model_properties["dataset_properties"][
+ "override_extent"
+ ]
+ slice_extent = dict(
+ longitude=slice(custom_extent[0], custom_extent[1]),
+ latitude=slice(custom_extent[3], custom_extent[2]),
+ )
+ except KeyError:
+ slice_extent = dict(
+ longitude=slice(
+ DOMAIN_EXTENTS[args["domain"]][0],
+ DOMAIN_EXTENTS[args["domain"]][1],
+ ),
+ latitude=slice(
+ DOMAIN_EXTENTS[args["domain"]][3],
+ DOMAIN_EXTENTS[args["domain"]][2],
+ ),
+ )
+ else:
+ if model_properties["dataset_properties"]["override_extent"] is None:
+ slice_extent = dict(
+ longitude=slice(
+ DOMAIN_EXTENTS[args["domain"]][0],
+ DOMAIN_EXTENTS[args["domain"]][1],
+ ),
+ latitude=slice(
+ DOMAIN_EXTENTS[args["domain"]][3],
+ DOMAIN_EXTENTS[args["domain"]][2],
+ ),
+ )
+
+ fronts_ds = xr.open_mfdataset(
+ front_files_month, combine="nested", concat_dim="time"
+ ).sel(**slice_extent)
+ fronts_ds_month = data_utils.reformat_fronts(
+ fronts_ds.sel(time="%d-%02d" % (year, month)), front_types
+ )
+
+ time_array = pd.read_pickle(
+ "%s/timesteps_%d%02d.pkl" % (args["tf_indir"], year, month)
+ )
+ num_timesteps = len(time_array)
+ lons = fronts_ds_month["longitude"].values
+ lats = fronts_ds_month["latitude"].values
+ Nlon = len(lons)
+ Nlat = len(lats)
+
+ tp_array_temporal = np.zeros(
+ shape=[num_front_types, num_timesteps, 5, 100]
+ ).astype("float32")
+ fp_array_temporal = np.zeros(
+ shape=[num_front_types, num_timesteps, 5, 100]
+ ).astype("float32")
+ tn_array_temporal = np.zeros(
+ shape=[num_front_types, num_timesteps, 5, 100]
+ ).astype("float32")
+ fn_array_temporal = np.zeros(
+ shape=[num_front_types, num_timesteps, 5, 100]
+ ).astype("float32")
+ tp_array_spatial = np.zeros(
+ shape=[num_front_types, Nlat, Nlon, 5, 100]
+ ).astype("float32")
+ fp_array_spatial = np.zeros(
+ shape=[num_front_types, Nlat, Nlon, 5, 100]
+ ).astype("float32")
+ tn_array_spatial = np.zeros(
+ shape=[num_front_types, Nlat, Nlon, 5, 100]
+ ).astype("float32")
+ fn_array_spatial = np.zeros(
+ shape=[num_front_types, Nlat, Nlon, 5, 100]
+ ).astype("float32")
+ weights = tf.cast(
+ tf.convert_to_tensor(
+ np.cos(np.deg2rad(lats))[np.newaxis, :, np.newaxis]
+ ),
+ tf.float32,
+ ) # latitude weights for the statistics
+
+ thresholds = np.linspace(
+ 0.01, 1, 100
+ ) # Probability thresholds for calculating performance statistics
+ neighborhoods = np.array(
+ [50, 100, 150, 200, 250]
+ ) # neighborhoods for checking whether a front is present (kilometers)
+
+ bool_tn_fn_dss = dict(
+ {
+ front: tf.convert_to_tensor(
+ xr.where(fronts_ds_month == front_no + 1, 1, 0)[
+ "identifier"
+ ].values
+ )
+ for front_no, front in enumerate(front_types)
+ }
+ )
+ bool_tp_fp_dss = dict({front: None for front in front_types})
+ probs_dss = dict(
+ {
+ front: tf.convert_to_tensor(probs_ds[front].values)
+ for front in front_types
+ }
+ )
+
+ spatial_ds = xr.Dataset(
+ coords={
+ "time": time_array[0],
+ "latitude": lats,
+ "longitude": lons,
+ "neighborhood": neighborhoods,
+ "threshold": thresholds,
+ }
+ )
+ temporal_ds = xr.Dataset(
+ coords={
+ "time": time_array,
+ "neighborhood": neighborhoods,
+ "threshold": thresholds,
+ }
+ )
+
+ for front_no, front_type in enumerate(front_types):
+ fronts_ds_month = data_utils.reformat_fronts(
+ fronts_ds.sel(time="%d-%02d" % (year, month)), front_types
+ )
+ print("%d-%02d: %s (TN/FN)" % (year, month, front_type))
+ ### Calculate true/false negatives ###
+ for i in range(100):
+ """
+ True negative ==> model correctly predicts the lack of a front at a given point
+ False negative ==> model does not predict a front, but a front exists
+
+ The numbers of true negatives and false negatives are the same for all neighborhoods and are calculated WITHOUT expanding the fronts.
+ If we were to calculate the negatives separately for each neighborhood, the number of misses would be artificially inflated, lowering the
+ final CSI scores and making the neighborhood method effectively useless.
+ """
+ tn = (
+ tf.cast(
+ tf.where(
+ (probs_dss[front_type] < thresholds[i])
+ & (bool_tn_fn_dss[front_type] == 0),
+ 1,
+ 0,
+ ),
+ tf.float32,
+ )
+ * weights
+ )
+ fn = (
+ tf.cast(
+ tf.where(
+ (probs_dss[front_type] < thresholds[i])
+ & (bool_tn_fn_dss[front_type] == 1),
+ 1,
+ 0,
+ ),
+ tf.float32,
+ )
+ * weights
+ )
+
+ tn_array_spatial[front_no, :, :, :, i] = tf.tile(
+ tf.expand_dims(tf.reduce_sum(tn, axis=0), axis=-1), (1, 1, 5)
+ )
+ fn_array_spatial[front_no, :, :, :, i] = tf.tile(
+ tf.expand_dims(tf.reduce_sum(fn, axis=0), axis=-1), (1, 1, 5)
+ )
+ tn_array_temporal[front_no, :, :, i] = tf.tile(
+ tf.expand_dims(tf.reduce_sum(tn, axis=(1, 2)), axis=-1), (1, 5)
+ )
+ fn_array_temporal[front_no, :, :, i] = tf.tile(
+ tf.expand_dims(tf.reduce_sum(fn, axis=(1, 2)), axis=-1), (1, 5)
+ )
+
+ ### Calculate true/false positives ###
+ for neighborhood in range(5):
+ fronts_ds_month = data_utils.expand_fronts(
+ fronts_ds_month, iterations=2
+ ) # Expand fronts by 50km
+ bool_tp_fp_dss[front_type] = tf.convert_to_tensor(
+ xr.where(fronts_ds_month == front_no + 1, 1, 0)[
+ "identifier"
+ ].values
+ ) # 1 = cold front, 0 = not a cold front
+ print(
+ "%d-%02d: %s (%d km)"
+ % (year, month, front_type, (neighborhood + 1) * 50)
+ )
+ for i in range(100):
+ """
+ True positive ==> model correctly identifies a front
+ False positive ==> model predicts a front, but no front is present within the given neighborhood
+ """
+ tp = (
+ tf.cast(
+ tf.where(
+ (probs_dss[front_type] > thresholds[i])
+ & (bool_tp_fp_dss[front_type] == 1),
+ 1,
+ 0,
+ ),
+ tf.float32,
+ )
+ * weights
+ )
+ fp = (
+ tf.cast(
+ tf.where(
+ (probs_dss[front_type] > thresholds[i])
+ & (bool_tp_fp_dss[front_type] == 0),
+ 1,
+ 0,
+ ),
+ tf.float32,
+ )
+ * weights
+ )
+
+ tp_array_spatial[front_no, :, :, neighborhood, i] = (
+ tf.reduce_sum(tp, axis=0)
+ )
+ fp_array_spatial[front_no, :, :, neighborhood, i] = (
+ tf.reduce_sum(fp, axis=0)
+ )
+ tp_array_temporal[front_no, :, neighborhood, i] = tf.reduce_sum(
+ tp, axis=(1, 2)
+ )
+ fp_array_temporal[front_no, :, neighborhood, i] = tf.reduce_sum(
+ fp, axis=(1, 2)
+ )
+
+ spatial_ds["tp_spatial_%s" % front_type] = (
+ ("latitude", "longitude", "neighborhood", "threshold"),
+ tp_array_spatial[front_no],
+ )
+ spatial_ds["fp_spatial_%s" % front_type] = (
+ ("latitude", "longitude", "neighborhood", "threshold"),
+ fp_array_spatial[front_no],
+ )
+ spatial_ds["tn_spatial_%s" % front_type] = (
+ ("latitude", "longitude", "neighborhood", "threshold"),
+ tn_array_spatial[front_no],
+ )
+ spatial_ds["fn_spatial_%s" % front_type] = (
+ ("latitude", "longitude", "neighborhood", "threshold"),
+ fn_array_spatial[front_no],
+ )
+ temporal_ds["tp_temporal_%s" % front_type] = (
+ ("time", "neighborhood", "threshold"),
+ tp_array_temporal[front_no],
+ )
+ temporal_ds["fp_temporal_%s" % front_type] = (
+ ("time", "neighborhood", "threshold"),
+ fp_array_temporal[front_no],
+ )
+ temporal_ds["tn_temporal_%s" % front_type] = (
+ ("time", "neighborhood", "threshold"),
+ tn_array_temporal[front_no],
+ )
+ temporal_ds["fn_temporal_%s" % front_type] = (
+ ("time", "neighborhood", "threshold"),
+ fn_array_temporal[front_no],
+ )
+
+ spatial_ds.astype("float32").to_netcdf(
+ path=stats_dataset_path.replace(".nc", "_spatial.nc"),
+ mode="w",
+ engine="netcdf4",
+ )
+ spatial_ds.close()
+ temporal_ds.to_netcdf(
+ path=stats_dataset_path.replace(".nc", "_temporal.nc"),
+ mode="w",
+ engine="netcdf4",
+ )
+ temporal_ds.close()
diff --git a/src/fronts/evaluation/generate_permutations.py b/src/fronts/evaluation/generate_permutations.py
new file mode 100644
index 0000000..1e25f93
--- /dev/null
+++ b/src/fronts/evaluation/generate_permutations.py
@@ -0,0 +1,489 @@
+"""
+Calculate permutation importance.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2024.10.11
+
+TODO:
+ * Add more documentation
+"""
+
+import itertools
+import os
+import pandas as pd
+from fronts.model import losses
+from fronts.utils.data_utils import combine_datasets
+from fronts.utils import file_manager as fm
+import tensorflow as tf
+import numpy as np
+import argparse
+import pickle
+
+
+def shuffle_inputs(image, labels):
+ """
+ image:
+ """
+ if level_nums is None:
+ lvl_nums = [
+ None,
+ ]
+ elif type(level_nums) == int:
+ lvl_nums = [
+ level_nums,
+ ]
+ else:
+ lvl_nums = level_nums
+
+ if variable_nums is None:
+ var_nums = [
+ None,
+ ]
+ elif type(variable_nums) == int:
+ var_nums = [
+ variable_nums,
+ ]
+ else:
+ var_nums = variable_nums
+
+ for var_num in var_nums:
+ for lvl_num in lvl_nums:
+ values_to_shuffle = image[..., lvl_num, var_num]
+ num_elements = tf.size(values_to_shuffle)
+
+ lon_indices = tf.random.uniform(
+ [num_elements], 0, image.shape[0] - 1, dtype=tf.int32
+ )
+ lat_indices = tf.random.uniform(
+ [num_elements], 0, image.shape[1] - 1, dtype=tf.int32
+ )
+
+ if lvl_num is None:
+ pressure_level_indices = tf.random.uniform(
+ [num_elements], 0, image.shape[2] - 1, dtype=tf.int32
+ )
+ else:
+ pressure_level_indices = tf.cast(
+ tf.fill([num_elements], lvl_num), tf.int32
+ )
+
+ if var_num is None:
+ variable_indices = tf.random.uniform(
+ [num_elements], 0, image.shape[3] - 1, dtype=tf.int32
+ )
+ else:
+ variable_indices = tf.cast(tf.fill([num_elements], var_num), tf.int32)
+
+ indices = tf.stack(
+ [lon_indices, lat_indices, pressure_level_indices, variable_indices],
+ axis=-1,
+ )
+ image = tf.tensor_scatter_nd_update(
+ image, indices, tf.reshape(values_to_shuffle, [num_elements])
+ )
+
+ return image, labels
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+
+ parser.add_argument(
+ "--tf_indir",
+ type=str,
+ required=True,
+ help="Input directory for the tensorflow dataset(s).",
+ )
+ parser.add_argument(
+ "--model_dir",
+ type=str,
+ required=True,
+ help="Directory where the models are or will be saved to.",
+ )
+ parser.add_argument(
+ "--model_number",
+ type=int,
+ required=True,
+ help="Number that the model will be assigned.",
+ )
+ parser.add_argument(
+ "--baseline",
+ action="store_true",
+ help="Calculate baseline loss values for the permutations.",
+ )
+ parser.add_argument(
+ "--single_pass", action="store_true", help="Perform single-pass permutations."
+ )
+ parser.add_argument(
+ "--multi_pass", action="store_true", help="Perform multi-pass permutations."
+ )
+ parser.add_argument(
+ "--batch_size",
+ type=int,
+ default=32,
+ help="Batch size for the model predictions.",
+ )
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=np.random.randint(0, 2**31 - 1),
+ help="Seed for the random number generators.",
+ )
+
+ args = vars(parser.parse_args())
+
+ dataset_properties = pd.read_pickle("%s/dataset_properties.pkl" % args["tf_indir"])
+ model_properties = pd.read_pickle(
+ "%s/model_%d/model_%d_properties.pkl"
+ % (args["model_dir"], args["model_number"], args["model_number"])
+ )
+
+ num_classes = model_properties["classes"]
+ front_types = model_properties["dataset_properties"]["front_types"]
+ variables = model_properties["dataset_properties"]["variables"]
+ pressure_levels = model_properties["dataset_properties"]["pressure_levels"]
+ domain = dataset_properties["domain"]
+ test_years = model_properties["test_years"]
+ num_outputs = 4 # TODO: get rid of this hard-coded integer
+
+ # new file naming method does not use underscores and will return an IndexError
+ try:
+ file_loader = fm.DataFileLoader(
+ args["tf_indir"], data_file_type="era5-tensorflow"
+ )
+ file_loader.test_years = test_years
+ file_loader.pair_with_fronts(
+ args["tf_indir"], underscore_skips=len(front_types)
+ )
+ except IndexError:
+ file_loader = fm.DataFileLoader(
+ args["tf_indir"], data_file_type="era5-tensorflow"
+ )
+ file_loader.test_years = test_years
+ file_loader.pair_with_fronts(args["tf_indir"])
+
+ try:
+ optimizer_args = {
+ arg: model_properties[arg] for arg in ["learning_rate", "beta_1", "beta_2"]
+ }
+ optimizer = getattr(tf.keras.optimizers, model_properties["optimizer"])(
+ **optimizer_args
+ )
+ except KeyError:
+ optimizer = getattr(tf.keras.optimizers, model_properties["optimizer"][0])(
+ **model_properties["optimizer"][1]
+ )
+
+ metric_functions = dict()
+ for front_no, front_type in enumerate(front_types):
+ class_weights = np.zeros(num_classes)
+ class_weights[front_no + 1] = 1
+ metric_functions[front_type] = losses.probability_of_detection(
+ class_weights=class_weights
+ )
+ metric_functions[front_type]._name = front_type
+
+ model = fm.load_model(args["model_number"], args["model_dir"])
+ model.compile(
+ optimizer=optimizer,
+ loss=losses.probability_of_detection(),
+ metrics=[metric_functions[func] for func in metric_functions],
+ )
+
+ X_datasets = file_loader.data_files_test
+ y_datasets = file_loader.front_files_test
+ X = combine_datasets(X_datasets)
+ y = combine_datasets(y_datasets)
+
+ permutations_file = "%s/model_%d/permutations_%d_%s.pkl" % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ domain,
+ )
+ permutations_dict = (
+ dict()
+ if not os.path.isfile(permutations_file)
+ else pd.read_pickle(permutations_file)
+ )
+
+ if "seed" not in permutations_dict:
+ permutations_dict["seed"] = args["seed"]
+ tf.random.set_seed(permutations_dict["seed"])
+
+ print("Seed: %d" % permutations_dict["seed"])
+ if "baseline" not in permutations_dict or args["baseline"]:
+ print("=== Baselines ===")
+
+ print("--> opening datasets")
+ Xy = combine_datasets(X_datasets, y_datasets)
+ Xy = Xy.batch(args["batch_size"])
+
+ print("--> generating predictions")
+ prediction = np.array(model.evaluate(Xy, verbose=0))
+
+ # overall loss for each front type
+ baseline = np.array(
+ [
+ prediction[1 + num_outputs + front_no]
+ for front_no in range(num_classes - 1)
+ ]
+ )
+
+ for front_no, front_type in enumerate(front_types):
+ print("--> baseline loss (%s): %s" % (front_type, str(baseline[front_no])))
+
+ permutations_dict["baseline"] = baseline
+
+ print("--> saving results to %s" % permutations_file)
+ with open(permutations_file, "wb") as f:
+ pickle.dump(permutations_dict, f)
+
+ if "single_pass" not in permutations_dict or args["single_pass"]:
+ assert "baseline" in permutations_dict, (
+ "Must calculate baseline values prior to permutations!"
+ )
+ baseline = permutations_dict["baseline"]
+ permutations_dict["single_pass"] = (
+ dict()
+ if "single_pass" not in permutations_dict
+ else permutations_dict["single_pass"]
+ )
+
+ print("\n=== Single-pass permutations ===")
+
+ combinations = list([variable, None] for variable in variables)
+ combinations.extend(list([None, level] for level in pressure_levels))
+ combinations.extend(
+ list(combo) for combo in itertools.product(variables, pressure_levels)
+ )
+
+ num_combinations = len(combinations)
+ for combination_num, combination in enumerate(combinations):
+ variable, level = combination
+ variable_nums = variables.index(variable) if variable is not None else None
+ level_nums = pressure_levels.index(level) if level is not None else None
+
+ printout_str = "(%d/%d) " % (combination_num + 1, num_combinations)
+ if variable is None:
+ printout_str += "%s, all variables" % level
+ permutation_dict_key = level
+ elif level is None:
+ printout_str += "%s, all levels" % variable
+ permutation_dict_key = variable
+ else:
+ printout_str += "%s_%s" % (variable, level)
+ permutation_dict_key = "%s_%s" % (variable, level)
+
+ if permutation_dict_key in permutations_dict["single_pass"]:
+ printout_str += " [SKIPPING]"
+ print(printout_str)
+ continue
+
+ print(printout_str)
+ Xy = combine_datasets(X_datasets, y_datasets)
+ Xy = Xy.map(shuffle_inputs)
+ Xy = Xy.batch(args["batch_size"])
+
+ prediction = np.array(model.evaluate(Xy, verbose=0))
+ loss = np.array(
+ [
+ prediction[1 + num_outputs + front_no]
+ for front_no in range(num_classes - 1)
+ ]
+ )
+ importance = np.round(100 * (loss - baseline) / baseline, 3)
+ permutations_dict["single_pass"][permutation_dict_key] = importance
+
+ with open(permutations_file, "wb") as f:
+ pickle.dump(permutations_dict, f)
+
+ if args["multi_pass"]:
+ single_pass_results = np.array(
+ [permutations_dict["single_pass"][var] for var in variables]
+ )
+ most_important_variable = np.argmax(single_pass_results, axis=0)
+
+ assert "baseline" in permutations_dict, (
+ "Must calculate baseline values prior to permutations!"
+ )
+ baseline = permutations_dict["baseline"]
+ permutations_dict["multi_pass"] = (
+ dict()
+ if "multi_pass" not in permutations_dict
+ else permutations_dict["multi_pass"]
+ )
+
+ # Find most important variable (overall and by type) based on the single-pass permutations
+
+ print("\n=== Multi-pass permutations ===")
+ print("-- VARIABLES, ALL LEVELS --")
+ ######################################### Variable importance by type #########################################
+ level_nums = None
+ for front_no, front_type in enumerate(front_types):
+ permutations_dict["multi_pass"][front_type] = (
+ dict()
+ if front_type not in permutations_dict["multi_pass"]
+ else permutations_dict["multi_pass"][front_type]
+ )
+
+ permutations_dict["multi_pass"][front_type][
+ variables[most_important_variable[front_no]]
+ ] = np.max(single_pass_results, axis=0)[front_no]
+
+ # most important variable will be the first variable to be shuffled in the shuffle_inputs function
+ shuffled_parameters = np.array([most_important_variable[front_no]])
+ variable_nums = np.array(
+ [most_important_variable[front_no], 0], dtype=np.int32
+ ) # the 0 is a placeholder that will be overwritten iteratively
+
+ variable_shuffle_order = np.arange(
+ 0, len(variables)
+ ) # variables to shuffle
+ variable_shuffle_order = np.delete(
+ variable_shuffle_order, shuffled_parameters
+ ) # remove pre-shuffled indices
+
+ while len(variable_shuffle_order) > 0:
+ print(
+ f"---> {front_type}: shuffling %s"
+ % ", ".join(variables[param] for param in shuffled_parameters)
+ )
+ temp_importance_list = np.array(
+ []
+ ) # list that will be used to temporarily store importance values for the current round
+
+ for var_to_shuffle in variable_shuffle_order:
+ variable_nums[-1] = var_to_shuffle
+
+ Xy = combine_datasets(X_datasets, y_datasets)
+ Xy = Xy.map(shuffle_inputs)
+ Xy = Xy.batch(args["batch_size"])
+
+ losses = np.array(model.evaluate(Xy, verbose=0))
+ loss_by_type = np.array(
+ [
+ losses[1 + num_outputs + front_no]
+ for front_no in range(num_classes - 1)
+ ]
+ )
+ importance = np.round(
+ 100
+ * (loss_by_type[front_no] - baseline[front_no])
+ / baseline[front_no],
+ 3,
+ )
+ print("-> %s:" % variables[var_to_shuffle], importance)
+
+ temp_importance_list = np.append(temp_importance_list, importance)
+
+ most_important_variable_for_round = variable_shuffle_order[
+ np.argmax(temp_importance_list)
+ ]
+ permutations_dict["multi_pass"][front_type][
+ variables[most_important_variable_for_round]
+ ] = np.max(temp_importance_list)
+
+ shuffled_parameters = np.append(
+ shuffled_parameters, most_important_variable_for_round
+ )
+ variable_nums[-1] = (
+ most_important_variable_for_round # add the round's most important variable to the list of variables to shuffle
+ )
+ variable_shuffle_order = np.delete(
+ variable_shuffle_order, np.argmax(temp_importance_list)
+ ) # remove the round's most important variable for the next round
+ variable_nums = np.append(
+ variable_nums, 0
+ ) # add another index for variables in the next round
+
+ with open(permutations_file, "wb") as f:
+ pickle.dump(permutations_dict, f)
+
+ print("-- LEVELS, ALL VARIABLES --")
+ variable_nums = None
+ single_pass_results = np.array(
+ [permutations_dict["single_pass"][lvl] for lvl in pressure_levels]
+ )
+ most_important_level = np.argmax(single_pass_results, axis=0)
+
+ ######################################### level importance by type #########################################
+ for front_no, front_type in enumerate(front_types):
+ permutations_dict["multi_pass"][front_type] = (
+ dict()
+ if front_type not in permutations_dict["multi_pass"]
+ else permutations_dict["multi_pass"][front_type]
+ )
+
+ permutations_dict["multi_pass"][front_type][
+ pressure_levels[most_important_level[front_no]]
+ ] = np.max(single_pass_results, axis=0)[front_no]
+
+ # most important level will be the first level to be shuffled in the shuffle_inputs function
+ shuffled_levels = np.array([most_important_level[front_no]])
+ level_nums = np.array(
+ [most_important_level[front_no], 0], dtype=np.int32
+ ) # the 0 is a placeholder that will be overwritten iteratively
+
+ level_shuffle_order = np.arange(
+ 0, len(pressure_levels)
+ ) # levels to shuffle
+ level_shuffle_order = np.delete(
+ level_shuffle_order, shuffled_levels
+ ) # remove pre-shuffled indices
+
+ while len(level_shuffle_order) > 0:
+ print(
+ f"---> {front_type}: shuffling %s"
+ % ", ".join(pressure_levels[param] for param in shuffled_levels)
+ )
+ temp_importance_list = np.array(
+ []
+ ) # list that will be used to temporarily store importance values for the current round
+
+ for level_to_shuffle in level_shuffle_order:
+ level_nums[-1] = level_to_shuffle
+
+ Xy = combine_datasets(X_datasets, y_datasets)
+ Xy = Xy.map(shuffle_inputs)
+ Xy = Xy.batch(args["batch_size"])
+
+ losses = np.array(model.evaluate(Xy, verbose=0))
+ loss_by_type = np.array(
+ [
+ losses[1 + num_outputs + front_no]
+ for front_no in range(num_classes - 1)
+ ]
+ )
+ importance = np.round(
+ 100
+ * (loss_by_type[front_no] - baseline[front_no])
+ / baseline[front_no],
+ 3,
+ )
+ print("-> %s:" % pressure_levels[level_to_shuffle], importance)
+
+ temp_importance_list = np.append(temp_importance_list, importance)
+
+ most_important_level_for_round = level_shuffle_order[
+ np.argmax(temp_importance_list)
+ ]
+ permutations_dict["multi_pass"][front_type][
+ pressure_levels[most_important_level_for_round]
+ ] = np.max(temp_importance_list)
+
+ shuffled_levels = np.append(
+ shuffled_levels, most_important_level_for_round
+ )
+ level_nums[-1] = (
+ most_important_level_for_round # add the round's most important level to the list of levels to shuffle
+ )
+ level_shuffle_order = np.delete(
+ level_shuffle_order, np.argmax(temp_importance_list)
+ ) # remove the round's most important level for the next round
+ level_nums = np.append(
+ level_nums, 0
+ ) # add another index for levels in the next round
+
+ with open(permutations_file, "wb") as f:
+ pickle.dump(permutations_dict, f)
diff --git a/src/fronts/evaluation/generate_saliency_maps.py b/src/fronts/evaluation/generate_saliency_maps.py
new file mode 100644
index 0000000..4a2bcec
--- /dev/null
+++ b/src/fronts/evaluation/generate_saliency_maps.py
@@ -0,0 +1,219 @@
+"""
+Generate saliency maps for a model.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2024.10.10
+"""
+
+import argparse
+import os
+from fronts.utils import data_utils
+import numpy as np
+import pandas as pd
+from fronts.utils import file_manager as fm
+import xarray as xr
+import tensorflow as tf
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--dataset",
+ type=str,
+ default="test",
+ help="Dataset for which to make saliency maps. Options are: 'training', 'validation', 'test'",
+ )
+ parser.add_argument(
+ "--year_and_month",
+ type=int,
+ nargs=2,
+ help="Year and month for which to make saleincy maps.",
+ )
+ parser.add_argument(
+ "--tf_indir",
+ type=str,
+ required=True,
+ help="Input directory for the tensorflow dataset(s).",
+ )
+ parser.add_argument(
+ "--model_dir",
+ type=str,
+ required=True,
+ help="Directory where the models are stored.",
+ )
+ parser.add_argument("--model_number", type=int, required=True, help="Model number.")
+ parser.add_argument(
+ "--batch_size",
+ type=int,
+ default=8,
+ help="Batch size for the model predictions. Since the gradients will also be retrieved, this should be lower than "
+ "the batch sizes used during training.",
+ )
+ parser.add_argument(
+ "--freq", type=str, required=True, help="Timestep frequency of the input data."
+ )
+ parser.add_argument(
+ "--verbose",
+ action="store_true",
+ help="Print out the progress of saliency map generation by batch.",
+ )
+ args = vars(parser.parse_args())
+
+ gpus = tf.config.list_physical_devices(device_type="GPU") # Find available GPUs
+ if len(gpus) > 0:
+ tf.config.set_visible_devices(devices=gpus[0], device_type="GPU")
+ gpus = tf.config.get_visible_devices(device_type="GPU") # List of selected GPUs
+ tf.config.experimental.set_memory_growth(
+ device=gpus[0], enable=True
+ ) # allow memory growth on GPU
+
+ model_properties = pd.read_pickle(
+ "%s/model_%d/model_%d_properties.pkl"
+ % (args["model_dir"], args["model_number"], args["model_number"])
+ )
+ dataset_properties = pd.read_pickle(
+ "%s/dataset_properties.pkl" % args["tf_indir"]
+ ) # properties of the dataset being used for saliency maps
+
+ variables = model_properties["dataset_properties"]["variables"]
+ pressure_levels = model_properties["dataset_properties"]["pressure_levels"]
+ num_classes = model_properties["classes"]
+ front_types = model_properties["dataset_properties"]["front_types"]
+ test_years = model_properties["test_years"]
+ domain = dataset_properties["domain"]
+
+ model = fm.load_model(args["model_number"], args["model_dir"])
+
+ file_loader = fm.DataFileLoader(args["tf_indir"], data_file_type="era5-tensorflow")
+
+ if args["year_and_month"] is not None:
+ years, months = [args["year_and_month"][0]], [args["year_and_month"][1]]
+ else:
+ years, months = model_properties["%s_years" % args["dataset"]], range(1, 13)
+
+ for year in years:
+ file_loader.test_years = [
+ year,
+ ]
+ files_for_year = file_loader.data_files_test
+
+ for month in months:
+ gradients = None
+
+ try:
+ tf_ds = tf.data.Dataset.load(
+ [
+ file
+ for file in files_for_year
+ if "_%d%02d" % (year, month) in file
+ ][0]
+ )
+ except IndexError:
+ print(
+ "ERA5 tensorflow dataset not found for %d-%02d in %s"
+ % (year, month, args["tf_indir"])
+ )
+ continue
+
+ next_year = year + 1 if month == 12 else year
+ next_month = 1 if month == 12 else month + 1
+
+ init_times = pd.date_range(
+ "%s-%02d" % (year, month),
+ "%s-%02d" % (next_year, next_month),
+ freq=args["freq"],
+ )[:-1]
+
+ os.makedirs(
+ "%s/model_%d/saliencymaps" % (args["model_dir"], args["model_number"]),
+ exist_ok=True,
+ ) # make directory for the saliency maps
+ salmap_dataset_path = (
+ "%s/model_%d/saliencymaps/model_%d_salmap_%s_%d%02d.nc"
+ % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ domain,
+ year,
+ month,
+ )
+ )
+
+ assert len(tf_ds) == len(init_times), (
+ "Length of provided tensorflow dataset (%d) does not match the number of timesteps (%d) in %d-%02d with the provided frequency (%s)"
+ % (len(tf_ds), len(init_times), year, month, args["freq"])
+ )
+
+ tf_ds = tf_ds.batch(
+ args["batch_size"]
+ ) # split dataset into batches (necessary for saliency maps because of memory issues)
+ num_batches = len(tf_ds)
+
+ print("Generating saliency maps for %d-%02d" % (year, month))
+
+ for batch_num, batch in enumerate(tf_ds, start=1):
+ if args["verbose"]:
+ print("Current batch: %d/%d" % (batch_num, num_batches), end="\r")
+ with tf.GradientTape(persistent=True) as tape:
+ tape.watch(batch)
+ predictions = model(batch)[0]
+
+ # batch_gradient: model gradients for the current batch, values must be converted to float32 for netcdf4 support
+ batch_gradient = np.stack(
+ [
+ np.max(
+ tape.gradient(
+ predictions[..., class_idx + 1 : class_idx + 2],
+ batch,
+ ).numpy(),
+ axis=-1,
+ )
+ for class_idx in range(num_classes - 1)
+ ],
+ axis=-1,
+ ).astype("float32")
+ gradients = (
+ batch_gradient
+ if gradients is None
+ else np.concatenate([gradients, batch_gradient], axis=0)
+ )
+
+ domain_ext = data_utils.DOMAIN_EXTENTS[domain]
+
+ domain_size = (
+ int((domain_ext[1] - domain_ext[0]) // 0.25) + 1,
+ int((domain_ext[3] - domain_ext[2]) // 0.25) + 1,
+ )
+
+ lons = np.linspace(domain_ext[0], domain_ext[1], domain_size[0])
+ lats = np.linspace(domain_ext[2], domain_ext[3], domain_size[1])[
+ ::-1
+ ] # lats in descending order (north-south)
+
+ salmaps = xr.Dataset(
+ data_vars=dict(
+ {
+ "%s" % front_type: (
+ ("time", "longitude", "latitude"),
+ np.max(gradients[..., idx], axis=-1),
+ )
+ for idx, front_type in enumerate(front_types)
+ }
+ | {
+ "%s_pl" % front_type: (
+ ("time", "longitude", "latitude", "pressure_level"),
+ gradients[..., idx],
+ )
+ for idx, front_type in enumerate(front_types)
+ }
+ ),
+ coords={
+ "time": init_times,
+ "longitude": lons,
+ "latitude": lats,
+ "pressure_level": pressure_levels,
+ },
+ attrs={"model_number": args["model_number"]},
+ )
+ salmaps.to_netcdf(salmap_dataset_path, mode="w", engine="netcdf4")
diff --git a/src/fronts/evaluation/predict.py b/src/fronts/evaluation/predict.py
new file mode 100644
index 0000000..7160f70
--- /dev/null
+++ b/src/fronts/evaluation/predict.py
@@ -0,0 +1,1226 @@
+"""
+Generate predictions with a model.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.2.12
+"""
+
+import argparse
+import pandas as pd
+import numpy as np
+import xarray as xr
+import os
+import tensorflow as tf
+from fronts.utils import data_utils
+from fronts.utils import file_manager as fm
+import scipy
+
+
+def _add_image_to_map(
+ stitched_map_probs: np.array,
+ image_probs: np.array,
+ map_created: bool,
+ num_images_lon: int,
+ num_images_lat: int,
+ lon_image: int,
+ lat_image: int,
+ image_size_lon: int,
+ image_size_lat: int,
+ lon_image_spacing: int,
+ lat_image_spacing: int,
+):
+ """
+ Add model prediction to the stitched map.
+
+ Parameters
+ ----------
+ stitched_map_probs: Numpy array
+ Array of front probabilities for the final map.
+ image_probs: Numpy array
+ Array of front probabilities for the current prediction/image.
+ map_created: bool
+ Boolean flag that declares whether the final map has been completed.
+ num_images_lon: int
+ Number of images along the longitude dimension of the domain.
+ num_images_lat: int
+ Number of images along the latitude dimension of the domain.
+ lon_image: int
+ Current image number along the longitude dimension.
+ lat_image: int
+ Current image number along the latitude dimension.
+ image_size_lon: int
+ Number of pixels along the longitude dimension of the model predictions.
+ image_size_lat: int
+ Number of pixels along the latitude dimension of the model predictions.
+ lon_image_spacing: int
+ Number of pixels between each image along the longitude dimension.
+ lat_image_spacing: int
+ Number of pixels between each image along the latitude dimension.
+
+ Returns
+ -------
+ map_created: bool
+ Boolean flag that declares whether the final map has been completed.
+ stitched_map_probs: array
+ Array of front probabilities for the final map.
+ """
+
+ if lon_image == 0: # If the image is on the western edge of the domain
+ if lat_image == 0: # If the image is on the northern edge of the domain
+ # Add first image to map
+ stitched_map_probs[:, 0:image_size_lon, 0:image_size_lat] = image_probs[
+ :, :image_size_lon, :image_size_lat
+ ]
+
+ if num_images_lon == 1 and num_images_lat == 1:
+ map_created = True
+
+ elif (
+ lat_image != num_images_lat - 1
+ ): # If the image is not on the northern nor the southern edge of the domain
+ # Take the maximum of the overlapping pixels along sets of constant latitude
+ stitched_map_probs[
+ :,
+ 0:image_size_lon,
+ int(lat_image * lat_image_spacing) : int(
+ (lat_image - 1) * lat_image_spacing
+ )
+ + image_size_lat,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ 0:image_size_lon,
+ int(lat_image * lat_image_spacing) : int(
+ (lat_image - 1) * lat_image_spacing
+ )
+ + image_size_lat,
+ ],
+ image_probs[:, :image_size_lon, : image_size_lat - lat_image_spacing],
+ )
+
+ # Add the remaining pixels of the current image to the map
+ stitched_map_probs[
+ :,
+ 0:image_size_lon,
+ int(lat_image_spacing * (lat_image - 1)) + image_size_lat : int(
+ lat_image_spacing * lat_image
+ )
+ + image_size_lat,
+ ] = image_probs[
+ :, :image_size_lon, image_size_lat - lat_image_spacing : image_size_lat
+ ]
+
+ if num_images_lon == 1 and num_images_lat == 2:
+ map_created = True
+
+ else: # If the image is on the southern edge of the domain
+ # Take the maximum of the overlapping pixels along sets of constant latitude
+ stitched_map_probs[
+ :,
+ 0:image_size_lon,
+ int(lat_image * lat_image_spacing) : int(
+ (lat_image - 1) * lat_image_spacing
+ )
+ + image_size_lat,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ :image_size_lon,
+ int(lat_image * lat_image_spacing) : int(
+ (lat_image - 1) * lat_image_spacing
+ )
+ + image_size_lat,
+ ],
+ image_probs[:, :image_size_lon, : image_size_lat - lat_image_spacing],
+ )
+
+ # Add the remaining pixels of the current image to the map
+ stitched_map_probs[
+ :,
+ 0:image_size_lon,
+ int(lat_image_spacing * (lat_image - 1)) + image_size_lat :,
+ ] = image_probs[
+ :, :image_size_lon, image_size_lat - lat_image_spacing : image_size_lat
+ ]
+
+ if num_images_lon == 1:
+ map_created = True
+
+ elif (
+ lon_image != num_images_lon - 1
+ ): # If the image is not on the western nor the eastern edge of the domain
+ if lat_image == 0: # If the image is on the northern edge of the domain
+ # Take the maximum of the overlapping pixels along sets of constant longitude
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ (lon_image - 1) * lon_image_spacing
+ )
+ + image_size_lon,
+ 0:image_size_lat,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ (lon_image - 1) * lon_image_spacing
+ )
+ + image_size_lon,
+ 0:image_size_lat,
+ ],
+ image_probs[:, : image_size_lon - lon_image_spacing, :image_size_lat],
+ )
+
+ # Add the remaining pixels of the current image to the map
+ stitched_map_probs[
+ :,
+ int(lon_image_spacing * (lon_image - 1))
+ + image_size_lon : lon_image_spacing * lon_image + image_size_lon,
+ 0:image_size_lat,
+ ] = image_probs[
+ :, image_size_lon - lon_image_spacing : image_size_lon, :image_size_lat
+ ]
+
+ if num_images_lon == 2 and num_images_lat == 1:
+ map_created = True
+
+ elif (
+ lat_image != num_images_lat - 1
+ ): # If the image is not on the northern nor the southern edge of the domain
+ # Take the maximum of the overlapping pixels along sets of constant longitude
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ lon_image_spacing * (lon_image - 1)
+ )
+ + image_size_lon,
+ int(lat_image * lat_image_spacing) : int(lat_image * lat_image_spacing)
+ + image_size_lat,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ lon_image_spacing * (lon_image - 1)
+ )
+ + image_size_lon,
+ int(lat_image * lat_image_spacing) : int(
+ lat_image * lat_image_spacing
+ )
+ + image_size_lat,
+ ],
+ image_probs[:, : image_size_lon - lon_image_spacing, :image_size_lat],
+ )
+
+ # Take the maximum of the overlapping pixels along sets of constant latitude
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(lon_image * lon_image_spacing)
+ + image_size_lon,
+ int(lat_image * lat_image_spacing) : int(
+ lat_image_spacing * (lat_image - 1)
+ )
+ + image_size_lat,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ lon_image * lon_image_spacing
+ )
+ + image_size_lon,
+ int(lat_image * lat_image_spacing) : int(
+ lat_image_spacing * (lat_image - 1)
+ )
+ + image_size_lat,
+ ],
+ image_probs[:, :image_size_lon, : image_size_lat - lat_image_spacing],
+ )
+
+ # Add the remaining pixels of the current image to the map
+ stitched_map_probs[
+ :,
+ int(lon_image_spacing * (lon_image - 1))
+ + image_size_lon : lon_image_spacing * lon_image + image_size_lon,
+ int(lat_image_spacing * (lat_image - 1)) + image_size_lat : int(
+ lat_image_spacing * lat_image
+ )
+ + image_size_lat,
+ ] = image_probs[
+ :,
+ image_size_lon - lon_image_spacing : image_size_lon,
+ image_size_lat - lat_image_spacing : image_size_lat,
+ ]
+
+ if num_images_lon == 2 and num_images_lat == 2:
+ map_created = True
+
+ else: # If the image is on the southern edge of the domain
+ # Take the maximum of the overlapping pixels along sets of constant longitude
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ lon_image_spacing * (lon_image - 1)
+ )
+ + image_size_lon,
+ int(lat_image * lat_image_spacing) :,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ lon_image_spacing * (lon_image - 1)
+ )
+ + image_size_lon,
+ int(lat_image * lat_image_spacing) :,
+ ],
+ image_probs[:, : image_size_lon - lon_image_spacing, :image_size_lat],
+ )
+
+ # Take the maximum of the overlapping pixels along sets of constant latitude
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(lon_image * lon_image_spacing)
+ + image_size_lon,
+ int(lat_image * lat_image_spacing) : int(
+ lat_image_spacing * (lat_image - 1)
+ )
+ + image_size_lat,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ lon_image * lon_image_spacing
+ )
+ + image_size_lon,
+ int(lat_image * lat_image_spacing) : int(
+ lat_image_spacing * (lat_image - 1)
+ )
+ + image_size_lat,
+ ],
+ image_probs[:, :image_size_lon, : image_size_lat - lat_image_spacing],
+ )
+
+ # Add the remaining pixels of the current image to the map
+ stitched_map_probs[
+ :,
+ int(lon_image_spacing * (lon_image - 1))
+ + image_size_lon : lon_image_spacing * lon_image + image_size_lon,
+ int(lat_image_spacing * (lat_image - 1)) + image_size_lat :,
+ ] = image_probs[
+ :,
+ image_size_lon - lon_image_spacing : image_size_lon,
+ image_size_lat - lat_image_spacing : image_size_lat,
+ ]
+
+ if num_images_lon == 2 and num_images_lat > 2:
+ map_created = True
+ else:
+ if lat_image == 0: # If the image is on the northern edge of the domain
+ # Take the maximum of the overlapping pixels along sets of constant longitude
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ lon_image_spacing * (lon_image - 1)
+ )
+ + image_size_lon,
+ 0:image_size_lat,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ lon_image_spacing * (lon_image - 1)
+ )
+ + image_size_lon,
+ 0:image_size_lat,
+ ],
+ image_probs[:, : image_size_lon - lon_image_spacing, :image_size_lat],
+ )
+
+ # Add the remaining pixels of the current image to the map
+ stitched_map_probs[
+ :,
+ int(lon_image_spacing * (lon_image - 1)) + image_size_lon :,
+ 0:image_size_lat,
+ ] = image_probs[
+ :, image_size_lon - lon_image_spacing : image_size_lon, :image_size_lat
+ ]
+
+ if num_images_lat == 1:
+ map_created = True
+
+ elif (
+ lat_image != num_images_lat - 1
+ ): # If the image is not on the northern nor the southern edge of the domain
+ # Take the maximum of the overlapping pixels along sets of constant longitude
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ lon_image_spacing * (lon_image - 1)
+ )
+ + image_size_lon,
+ int(lat_image * lat_image_spacing) : int(lat_image * lat_image_spacing)
+ + image_size_lat,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) : int(
+ lon_image_spacing * (lon_image - 1)
+ )
+ + image_size_lon,
+ int(lat_image * lat_image_spacing) : int(
+ lat_image * lat_image_spacing
+ )
+ + image_size_lat,
+ ],
+ image_probs[:, : image_size_lon - lon_image_spacing, :image_size_lat],
+ )
+
+ # Take the maximum of the overlapping pixels along sets of constant latitude
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) :,
+ int(lat_image * lat_image_spacing) : int(
+ lat_image_spacing * (lat_image - 1)
+ )
+ + image_size_lat,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) :,
+ int(lat_image * lat_image_spacing) : int(
+ lat_image_spacing * (lat_image - 1)
+ )
+ + image_size_lat,
+ ],
+ image_probs[:, :image_size_lon, : image_size_lat - lat_image_spacing],
+ )
+
+ # Add the remaining pixels of the current image to the map
+ stitched_map_probs[
+ :,
+ int(lon_image_spacing * (lon_image - 1)) + image_size_lon :,
+ int(lat_image_spacing * (lat_image - 1)) + image_size_lat : int(
+ lat_image_spacing * lat_image
+ )
+ + image_size_lat,
+ ] = image_probs[
+ :,
+ image_size_lon - lon_image_spacing : image_size_lon,
+ image_size_lat - lat_image_spacing : image_size_lat,
+ ]
+
+ if num_images_lon > 2 and num_images_lat == 2:
+ map_created = True
+ else: # If the image is on the southern edge of the domain
+ # Take the maximum of the overlapping pixels along sets of constant longitude
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) :,
+ int(lat_image * lat_image_spacing) :,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) :,
+ int(lat_image * lat_image_spacing) :,
+ ],
+ image_probs[:, :image_size_lon, :image_size_lat],
+ )
+
+ # Take the maximum of the overlapping pixels along sets of constant latitude
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) :,
+ int(lat_image * lat_image_spacing) : int(
+ lat_image_spacing * (lat_image - 1)
+ )
+ + image_size_lat,
+ ] = np.maximum(
+ stitched_map_probs[
+ :,
+ int(lon_image * lon_image_spacing) :,
+ int(lat_image * lat_image_spacing) : int(
+ lat_image_spacing * (lat_image - 1)
+ )
+ + image_size_lat,
+ ],
+ image_probs[:, :image_size_lon, : image_size_lat - lat_image_spacing],
+ )
+
+ # Add the remaining pixels of the current image to the map
+ stitched_map_probs[
+ :,
+ int(lon_image_spacing * (lon_image - 1)) + image_size_lon :,
+ int(lat_image_spacing * (lat_image - 1)) + image_size_lat : int(
+ lat_image_spacing * lat_image
+ )
+ + image_size_lat,
+ ] = image_probs[
+ :,
+ image_size_lon - lon_image_spacing : image_size_lon,
+ image_size_lat - lat_image_spacing : image_size_lat,
+ ]
+
+ map_created = True
+
+ return stitched_map_probs, map_created
+
+
+def find_matches_for_domain(
+ domain_size: tuple | list,
+ image_size: tuple | list,
+ compatibility_mode: bool = False,
+ compat_images: tuple | list = None,
+):
+ """
+ Function that outputs the number of images that can be stitched together with the specified domain length and the length
+ of the domain dimension output by the model. This is also used to determine the compatibility of declared image and
+ parameters for model predictions.
+
+ Parameters
+ ----------
+ domain_size: iterable object with 2 integers
+ Number of pixels along each dimension of the final stitched map (lon lat).
+ image_size: iterable object with 2 integers
+ Number of pixels along each dimension of the model's output (lon lat).
+ compatibility_mode: bool
+ Boolean flag that declares whether the function is being used to check compatibility of given parameters.
+ compat_images: iterable object with 2 integers
+ Number of images declared for the stitched map in each dimension (lon lat). (Compatibility mode only)
+ """
+
+ ######################################### Check the parameters for errors ##########################################
+ if not isinstance(domain_size, (tuple, list)):
+ raise TypeError(
+ f"Expected a tuple or list for domain_size, received {type(domain_size)}"
+ )
+ elif len(domain_size) != 2:
+ raise TypeError(
+ f"Tuple or list for num_images must be length 2, received length {len(domain_size)}"
+ )
+
+ if not isinstance(image_size, (tuple, list)):
+ raise TypeError(
+ f"Expected a tuple or list for image_size, received {type(image_size)}"
+ )
+ elif len(image_size) != 2:
+ raise TypeError(
+ f"Tuple or list for image_size must be length 2, received length {len(image_size)}"
+ )
+
+ if compatibility_mode is not None and not isinstance(compatibility_mode, bool):
+ raise TypeError(
+ f"compatibility_mode must be a boolean, received {type(compatibility_mode)}"
+ )
+
+ if compat_images is not None:
+ if not isinstance(compat_images, (tuple, list)):
+ raise TypeError(
+ f"Expected a tuple or list for compat_images, received {type(compat_images)}"
+ )
+ elif len(compat_images) != 2:
+ raise TypeError(
+ f"Tuple or list for compat_images must be length 2, received length {len(compat_images)}"
+ )
+ ####################################################################################################################
+
+ if compatibility_mode:
+ """ These parameters are used when checking the compatibility of image stitching arguments. """
+ compat_images_lon = compat_images[
+ 0
+ ] # Number of images in the longitude direction
+ compat_images_lat = compat_images[
+ 1
+ ] # Number of images in the latitude direction
+ else:
+ compat_images_lon, compat_images_lat = None, None
+
+ # All of these boolean variables must be True after the compatibility check or else a ValueError is returned
+ lon_images_are_compatible = False
+ lat_images_are_compatible = False
+
+ num_matches = [
+ 0,
+ 0,
+ ] # Total number of matching image arguments found for each dimension
+
+ lon_image_matches = []
+ lat_image_matches = []
+
+ for lon_images in range(
+ 1, domain_size[0] - image_size[0] + 2
+ ): # Image counter for longitude dimension
+ if lon_images > 1:
+ lon_spacing = (domain_size[0] - image_size[0]) / (
+ lon_images - 1
+ ) # Spacing between images in the longitude dimension
+ else:
+ lon_spacing = 0
+ if (
+ lon_spacing - int(lon_spacing) == 0
+ and lon_spacing > 1
+ and image_size[0] - lon_spacing > 0
+ ): # Check compatibility of latitude image spacing
+ lon_image_matches.append(lon_images) # Add longitude image match to list
+ num_matches[0] += 1
+ if compatibility_mode:
+ if (
+ compat_images_lon == lon_images
+ ): # If the number of images for the compatibility check equals the match
+ lon_images_are_compatible = True
+ elif lon_spacing == 0 and domain_size[0] - image_size[0] == 0:
+ lon_image_matches.append(lon_images) # Add longitude image match to list
+ num_matches[0] += 1
+ if compatibility_mode:
+ if (
+ compat_images_lon == lon_images
+ ): # If the number of images for the compatibility check equals the match
+ lon_images_are_compatible = True
+
+ if num_matches[0] == 0:
+ raise ValueError(
+ f"No compatible value for num_images[0] was found with domain_size[0]={domain_size[0]} and image_size[0]={image_size[0]}."
+ )
+ if compatibility_mode:
+ if not lon_images_are_compatible:
+ raise ValueError(
+ f"num_images[0]={compat_images_lon} is not compatible with domain_size[0]={domain_size[0]} "
+ f"and image_size[0]={image_size[0]}.\n"
+ f"====> Compatible values for num_images[0] given domain_size[0]={domain_size[0]} "
+ f"and image_size[0]={image_size[0]}: {lon_image_matches}"
+ )
+ else:
+ print(f"Compatible longitude images: {lon_image_matches}")
+
+ for lat_images in range(
+ 1, domain_size[1] - image_size[1] + 2
+ ): # Image counter for latitude dimension
+ if lat_images > 1:
+ lat_spacing = (domain_size[1] - image_size[1]) / (
+ lat_images - 1
+ ) # Spacing between images in the latitude dimension
+ else:
+ lat_spacing = 0
+ if (
+ lat_spacing - int(lat_spacing) == 0
+ and lat_spacing > 1
+ and image_size[1] - lat_spacing > 0
+ ): # Check compatibility of latitude image spacing
+ lat_image_matches.append(lat_images) # Add latitude image match to list
+ num_matches[1] += 1
+ if compatibility_mode:
+ if (
+ compat_images_lat == lat_images
+ ): # If the number of images for the compatibility check equals the match
+ lat_images_are_compatible = True
+ elif lat_spacing == 0 and domain_size[1] - image_size[1] == 0:
+ lat_image_matches.append(lat_images) # Add latitude image match to list
+ num_matches[1] += 1
+ if compatibility_mode:
+ if (
+ compat_images_lat == lat_images
+ ): # If the number of images for the compatibility check equals the match
+ lat_images_are_compatible = True
+
+ if num_matches[1] == 0:
+ raise ValueError(
+ f"No compatible value for num_images[1] was found with domain_size[1]={domain_size[1]} and image_size[1]={image_size[1]}."
+ )
+ if compatibility_mode:
+ if not lat_images_are_compatible:
+ raise ValueError(
+ f"num_images[1]={compat_images_lat} is not compatible with domain_size[1]={domain_size[1]} "
+ f"and image_size[1]={image_size[1]}.\n"
+ f"====> Compatible values for num_images[1] given domain_size[1]={domain_size[1]} "
+ f"and image_size[1]={image_size[1]}: {lat_image_matches}"
+ )
+ else:
+ print(f"Compatible latitude images: {lat_image_matches}")
+
+
+def create_model_prediction_dataset(
+ stitched_map_probs: np.array,
+ lats: np.array,
+ lons: np.array,
+ front_types: str | list,
+):
+ """
+ Create an Xarray dataset containing model predictions.
+
+ Parameters
+ ----------
+ stitched_map_probs: np.array
+ Numpy array with probabilities for the given front type(s).
+ Shape/dimensions: [front types, longitude, latitude]
+ lats: np.array
+ 1D array of latitude values.
+ lons: np.array
+ 1D array of longitude values.
+ front_types: str or list
+ Front types within the dataset. See documentation in utils.data_utils.reformat fronts for more information.
+
+ Returns
+ -------
+ probs_ds: xr.Dataset
+ Xarray dataset containing front probabilities predicted by the model for each front type.
+ """
+
+ ######################################### Check the parameters for errors ##########################################
+ if not isinstance(stitched_map_probs, np.ndarray):
+ raise TypeError(
+ f"stitched_map_probs must be a NumPy array, received {type(stitched_map_probs)}"
+ )
+ if not isinstance(lats, np.ndarray):
+ raise TypeError(f"lats must be a NumPy array, received {type(lats)}")
+ if not isinstance(lons, np.ndarray):
+ raise TypeError(f"lons must be a NumPy array, received {type(lons)}")
+ if not isinstance(front_types, (tuple, list)):
+ raise TypeError(
+ f"Expected a tuple or list for front_types, received {type(front_types)}"
+ )
+ ####################################################################################################################
+
+ if args["data_source"] not in ["hrrr", "nam-12km"]:
+ spatial_dims = ("longitude", "latitude")
+ coords = {"latitude": lats, "longitude": lons}
+ else:
+ spatial_dims = ("x", "y")
+ coords = {"latitude": (("x", "y"), lats), "longitude": (("x", "y"), lons)}
+
+ if (
+ front_types == "F_BIN"
+ or front_types == "MERGED-F_BIN"
+ or front_types == "MERGED-T"
+ ):
+ probs_ds = xr.Dataset(
+ {front_types: (spatial_dims, stitched_map_probs[0])}, coords=coords
+ )
+ elif front_types == "MERGED-F":
+ probs_ds = xr.Dataset(
+ {
+ "CF_merged": (spatial_dims, stitched_map_probs[0]),
+ "WF_merged": (spatial_dims, stitched_map_probs[1]),
+ "SF_merged": (spatial_dims, stitched_map_probs[2]),
+ "OF_merged": (spatial_dims, stitched_map_probs[3]),
+ },
+ coords=coords,
+ )
+ elif front_types == "MERGED-ALL":
+ probs_ds = xr.Dataset(
+ {
+ "CF_merged": (spatial_dims, stitched_map_probs[0]),
+ "WF_merged": (spatial_dims, stitched_map_probs[1]),
+ "SF_merged": (spatial_dims, stitched_map_probs[2]),
+ "OF_merged": (spatial_dims, stitched_map_probs[3]),
+ "TROF_merged": (spatial_dims, stitched_map_probs[4]),
+ "INST": (spatial_dims, stitched_map_probs[5]),
+ "DL": (spatial_dims, stitched_map_probs[6]),
+ },
+ coords=coords,
+ )
+ elif type(front_types) == list:
+ probs_ds_dict = dict({})
+ for probs_ds_index, front_type in enumerate(front_types):
+ probs_ds_dict[front_type] = (
+ spatial_dims,
+ stitched_map_probs[probs_ds_index],
+ )
+ probs_ds = xr.Dataset(probs_ds_dict, coords=coords)
+ else:
+ raise ValueError(f"'{front_types}' is not a valid set of front types.")
+
+ return probs_ds
+
+
+if __name__ == "__main__":
+ """
+ All arguments listed in the examples are listed via argparse in alphabetical order below this comment block.
+ """
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--netcdf_indir",
+ type=str,
+ help="Main directory for the netcdf files containing variable data.",
+ )
+ parser.add_argument(
+ "--mergir_indir",
+ type=str,
+ help="Input directory for the netCDF files containing MERGIR data.",
+ )
+ parser.add_argument(
+ "--init_time",
+ type=int,
+ nargs=4,
+ help="Date and time of the data. Pass 4 ints in the following order: year, month, day, hour",
+ )
+ parser.add_argument("--domain", type=str, help="Domain of the data.")
+ parser.add_argument(
+ "--num_images",
+ type=int,
+ nargs=2,
+ default=[1, 1],
+ help="Number of images for each dimension the final stitched map for predictions: lon, lat",
+ )
+ parser.add_argument("--gpu_device", type=int, help="GPU device number.")
+ parser.add_argument(
+ "--batch_size",
+ type=int,
+ default=1,
+ help="Batch size for the model predictions.",
+ )
+ parser.add_argument(
+ "--image_size",
+ type=int,
+ nargs=2,
+ help="Number of pixels along each dimension of the model's output: lon, lat",
+ )
+ parser.add_argument(
+ "--memory_growth", action="store_true", help="Use memory growth on the GPU"
+ )
+ parser.add_argument("--model_dir", type=str, help="Directory for the models.")
+ parser.add_argument("--model_number", type=int, help="Model number.")
+ parser.add_argument(
+ "--data_source", type=str, default="era5", help="Data source for variables"
+ )
+
+ args = vars(parser.parse_args())
+
+ gpus = tf.config.list_physical_devices(device_type="GPU") # Find available GPUs
+ if len(gpus) > 0:
+ tf.config.set_visible_devices(
+ devices=gpus[args["gpu_device"]], device_type="GPU"
+ )
+ if args["memory_growth"]:
+ tf.config.experimental.set_memory_growth(
+ device=gpus[args["gpu_device"]], enable=True
+ )
+
+ else:
+ print("WARNING: No GPUs found, all computations will be performed on CPUs.")
+ tf.config.set_visible_devices([], "GPU")
+
+ ### Model properties ###
+ model_properties = pd.read_pickle(
+ f"{args['model_dir']}/model_{args['model_number']}/model_{args['model_number']}_properties.pkl"
+ )
+ model_type = model_properties["model_type"]
+
+ if args["image_size"] is None:
+ args["image_size"] = model_properties[
+ "image_size"
+ ] # The image size does not include the last dimension of the input size as it only represents the number of channels
+
+ try:
+ front_types = model_properties["dataset_properties"]["front_types"]
+ variables = model_properties["dataset_properties"]["variables"]
+ pressure_levels = model_properties["dataset_properties"]["pressure_levels"]
+ except KeyError: # Some older models do not have the dataset_properties dictionary
+ front_types = model_properties["front_types"]
+ variables = model_properties["variables"]
+ pressure_levels = model_properties["pressure_levels"]
+
+ load_mergir = (
+ "Tb" in variables
+ ) # load MERGIR data if brightness temperature (Tb) is requested
+ variables = [var for var in variables if var != "Tb"]
+
+ normalization_parameters = model_properties["normalization_parameters"]
+ normalization_method = model_properties["dataset_properties"][
+ "normalization_method"
+ ]
+
+ classes = model_properties["classes"]
+ test_years, valid_years = (
+ model_properties["test_years"],
+ model_properties["validation_years"],
+ )
+
+ try:
+ domain_extent = data_utils.DOMAIN_EXTENTS[args["data_source"]]
+ except KeyError:
+ domain_extent = data_utils.DOMAIN_EXTENTS[args["domain"]]
+
+ ### Properties of the final map made from stitched images ###
+ num_images_lon, num_images_lat = args["num_images"][0], args["num_images"][1]
+ if args["num_images"] == [1, 1]:
+ domain_size_lon, domain_size_lat = args["image_size"]
+ else:
+ domain_size_lon = int((domain_extent[1] - domain_extent[0]) // 0.25) + 1
+ domain_size_lat = int((domain_extent[3] - domain_extent[2]) // 0.25) + 1
+ image_size_lon, image_size_lat = args[
+ "image_size"
+ ] # Dimensions of the model's predictions
+
+ if num_images_lon > 1:
+ lon_image_spacing = int(
+ (domain_size_lon - image_size_lon) / (num_images_lon - 1)
+ )
+ else:
+ lon_image_spacing = 0
+
+ if num_images_lat > 1:
+ lat_image_spacing = int(
+ (domain_size_lat - image_size_lat) / (num_images_lat - 1)
+ )
+ else:
+ lat_image_spacing = 0
+
+ model = fm.load_model(args["model_number"], args["model_dir"])
+ num_dimensions = len(model.layers[0].input_shape[0]) - 2
+
+ ############################################### Load variable files ################################################
+ variable_files_obj = fm.DataFileLoader(
+ args["netcdf_indir"],
+ data_type=args["data_source"],
+ file_format="netcdf",
+ years=int(args["init_time"][0]),
+ months=int(args["init_time"][1]),
+ days=int(args["init_time"][2]),
+ hours=int(args["init_time"][3]),
+ )
+
+ if load_mergir:
+ variable_files_obj.add_file_list(
+ args["mergir_indir"], "MERGIR", ignore_domain=True
+ )
+ variable_files, mergir_files = variable_files_obj.files
+ else:
+ variable_files = variable_files_obj.files[0]
+
+ dataset_kwargs = {
+ "engine": "netcdf4"
+ } # Keyword arguments for loading variable files with xarray
+
+ if args["data_source"] not in ["hrrr", "nam-12km"]:
+ coords_sel_kwargs = {
+ "longitude": slice(domain_extent[0], domain_extent[1]),
+ "latitude": slice(domain_extent[3], domain_extent[2]),
+ }
+ spatial_dims = ("longitude", "latitude")
+ if args["data_source"] == "era5":
+ transpose_dims = ("time", "longitude", "latitude", "pressure_level")
+ else:
+ transpose_dims = (
+ "time",
+ "forecast_hour",
+ "longitude",
+ "latitude",
+ "pressure_level",
+ )
+ else:
+ coords_sel_kwargs = {}
+ spatial_dims = ("x", "y")
+ transpose_dims = ("time", "forecast_hour", "x", "y", "pressure_level")
+
+ if args["init_time"] is not None:
+ timestep_str = "%d%02d%02d%02d" % (
+ args["init_time"][0],
+ args["init_time"][1],
+ args["init_time"][2],
+ args["init_time"][3],
+ )
+ if args["data_source"] == "era5":
+ init_time_index = [
+ index
+ for index, file in enumerate(variable_files)
+ if timestep_str in file
+ ][0]
+ variable_files = [
+ variable_files[init_time_index],
+ ]
+ else:
+ variable_files = [file for file in variable_files if timestep_str in file]
+
+ subdir_base = "%s_%dx%d" % (
+ args["domain"],
+ args["num_images"][0],
+ args["num_images"][1],
+ )
+
+ variable_ds = xr.open_mfdataset(variable_files, **dataset_kwargs).sel(
+ **coords_sel_kwargs
+ )[variables]
+ if load_mergir:
+ mergir_ds = xr.open_mfdataset(mergir_files, **dataset_kwargs).rename(
+ {"lon": "longitude", "lat": "latitude"}
+ )
+
+ # pull original lat/lon coordinates from MERGIR dataset
+ mergir_lons = mergir_ds["longitude"].values
+ mergir_lats = mergir_ds["latitude"].values
+ mergir_lons = np.where(mergir_lons < 0, mergir_lons + 360, mergir_lons)
+
+ # reformat coordinates and slice domain
+ mergir_dataset = mergir_ds.assign_coords(longitude=mergir_lons)
+ mergir_dataset = mergir_dataset.reindex(
+ longitude=sorted(mergir_lons), latitude=mergir_lats[::-1]
+ )
+ mergir_dataset = (
+ mergir_dataset.isel(time=0)
+ .transpose("longitude", "latitude")
+ .sel(**coords_sel_kwargs)
+ )
+
+ # regrid the MERGIR data
+ Tb = mergir_dataset["Tb"].values
+ new_mergir_lons = mergir_dataset["longitude"].values
+ new_mergir_lats = mergir_dataset["latitude"].values
+ variable_lons = variable_ds["longitude"].values
+ variable_lats = variable_ds["latitude"].values
+ new_lons, new_lats = np.meshgrid(variable_lons, variable_lats)
+
+ # regrid and normalize the mergir data
+ Tb = scipy.interpolate.RegularGridInterpolator(
+ (new_mergir_lons, new_mergir_lats), Tb, method="nearest", bounds_error=False
+ )((new_lons, new_lats)).transpose()[..., np.newaxis]
+ Tb = (Tb - 197) / (329 - 197) # min max normalization
+ Tb = np.nan_to_num(Tb).transpose()
+
+ if args["data_source"] == "era5":
+ variable_ds = variable_ds.sel(pressure_level=pressure_levels).transpose(
+ *transpose_dims
+ )
+ image_lats = variable_ds.latitude.values[:domain_size_lat]
+ image_lons = variable_ds.longitude.values[:domain_size_lon]
+ else:
+ variable_ds = variable_ds.sel(pressure_level=pressure_levels).transpose(
+ *transpose_dims
+ )
+ forecast_hours = variable_ds["forecast_hour"].values
+ if args["data_source"] in ["hrrr", "nam-12km"]:
+ image_lats = variable_ds.latitude.values[:domain_size_lon, :domain_size_lat]
+ image_lons = variable_ds.longitude.values[
+ :domain_size_lon, :domain_size_lat
+ ]
+ else:
+ image_lats = variable_ds.latitude.values[:domain_size_lat]
+ image_lons = variable_ds.longitude.values[:domain_size_lon]
+
+ variable_batch_ds = data_utils.normalize_dataset(
+ variable_ds,
+ method=normalization_method,
+ normalization_parameters=normalization_parameters,
+ )
+
+ timesteps = variable_batch_ds["time"].values
+ num_timesteps = len(timesteps)
+ num_forecast_hours = len(variable_batch_ds["forecast_hour"])
+ map_created = (
+ False # Boolean that determines whether the final stitched map has been created
+ )
+
+ if args["data_source"] == "era5":
+ stitched_map_probs = np.empty(
+ shape=[num_timesteps, classes - 1, domain_size_lon, domain_size_lat]
+ )
+ else:
+ stitched_map_probs = np.empty(
+ shape=[
+ num_timesteps,
+ len(forecast_hours),
+ classes - 1,
+ domain_size_lon,
+ domain_size_lat,
+ ]
+ )
+
+ for lat_image in range(num_images_lat):
+ lat_index = int(lat_image * lat_image_spacing)
+ for lon_image in range(num_images_lon):
+ print(
+ f"image %d/%d"
+ % (
+ int(lat_image * num_images_lon) + lon_image + 1,
+ int(num_images_lon * num_images_lat),
+ )
+ )
+ lon_index = int(lon_image * lon_image_spacing)
+
+ # Select the current image
+ variable_batch_ds_new = (
+ variable_batch_ds[variables]
+ .isel(
+ {
+ "%s" % spatial_dims[0]: slice(
+ lon_index, lon_index + args["image_size"][0]
+ ),
+ "%s" % spatial_dims[1]: slice(
+ lat_index, lat_index + args["image_size"][1]
+ ),
+ }
+ )
+ .to_array()
+ .values
+ )
+
+ if load_mergir:
+ assert args["num_images"] == [1, 1], (
+ "Can only use MERGIR data if making a prediction using one image only."
+ )
+
+ Tb = Tb[
+ np.newaxis, ...
+ ] # add 'variable' dimension so it can be concatenated
+ if args["data_source"] != "era5":
+ Tb = Tb[:, np.newaxis, ...] # add forecast hour dimension
+ if num_dimensions == 3:
+ Tb = Tb[..., np.newaxis, :, :] # add vertical dimension
+ Tb = np.tile(
+ Tb,
+ [
+ *[1 for dim in range(len(Tb.shape) - 3)],
+ len(pressure_levels),
+ 1,
+ 1,
+ ],
+ )
+ variable_batch_ds_new = np.concatenate(
+ [variable_batch_ds_new, Tb], axis=0
+ )
+
+ if args["data_source"] == "era5":
+ variable_batch_ds_new = variable_batch_ds_new.transpose(
+ [1, 2, 3, 4, 0]
+ ) # (time, longitude, latitude, pressure level, variable)
+ else:
+ variable_batch_ds_new = variable_batch_ds_new.transpose(
+ [1, 2, 3, 4, 5, 0]
+ ) # (time, forecast hour, longitude, latitude, pressure level, variable)
+
+ if num_dimensions == 2:
+ ### Combine pressure levels and variables into one dimension ###
+ variable_batch_ds_new_shape = np.shape(variable_batch_ds_new)
+ variable_batch_ds_new = variable_batch_ds_new.reshape(
+ *[dim_size for dim_size in variable_batch_ds_new_shape[:-2]],
+ variable_batch_ds_new_shape[-2] * variable_batch_ds_new_shape[-1],
+ )
+
+ transpose_indices = (
+ 0,
+ 3,
+ 1,
+ 2,
+ ) # New order of indices for model predictions (time, front type, longitude, latitude)
+
+ ##################################### Generate the predictions #####################################
+ if args["data_source"] != "era5":
+ variable_ds_new_shape = np.shape(variable_batch_ds_new)
+ variable_batch_ds_new = variable_batch_ds_new.reshape(
+ variable_ds_new_shape[0] * variable_ds_new_shape[1],
+ *[dim_size for dim_size in variable_ds_new_shape[2:]],
+ )
+
+ prediction = model.predict(
+ variable_batch_ds_new, batch_size=args["batch_size"], verbose=0
+ )
+ num_dims_in_pred = len(np.shape(prediction))
+
+ if model_type in ["unet", "attention_unet"]:
+ if (
+ num_dims_in_pred == 4
+ ): # 2D labels, prediction shape: (time, lat, lon, front type)
+ image_probs = np.transpose(
+ prediction[:, :, :, 1:], transpose_indices
+ ) # transpose the predictions
+ else: # if num_dims_in_pred == 5; 3D labels, prediction shape: (time, lat, lon, pressure level, front type)
+ image_probs = np.transpose(
+ np.amax(prediction[:, :, :, :, 1:], axis=3), transpose_indices
+ ) # Take the maximum probability over the vertical dimension and transpose the predictions
+
+ elif model_type == "unet_3plus":
+ try:
+ deep_supervision = model_properties["deep_supervision"]
+ except KeyError:
+ deep_supervision = True # older models do not have this dictionary key, so just set it to True
+
+ if deep_supervision:
+ if (
+ num_dims_in_pred == 5
+ ): # 2D labels, prediction shape: (output level, time, lon, lat, front type)
+ image_probs = np.transpose(
+ prediction[0][:, :, :, 1:], transpose_indices
+ ) # transpose the predictions
+ else: # if num_dims_in_pred == 6; 3D labels, prediction shape: (output level, time, lon, lat, pressure level, front type)
+ image_probs = np.transpose(
+ np.amax(prediction[0][:, :, :, :, 1:], axis=3),
+ transpose_indices,
+ ) # Take the maximum probability over the vertical dimension and transpose the predictions
+ else:
+ if (
+ num_dims_in_pred == 4
+ ): # 2D labels, prediction shape: (time, lon, lat, front type)
+ image_probs = np.transpose(
+ prediction[:, :, :, 1:], transpose_indices
+ ) # transpose the predictions
+ else: # if num_dims_in_pred == 5; 3D labels, prediction shape: (time, lat, lon, pressure level, front type)
+ image_probs = np.transpose(
+ np.amax(prediction[:, :, :, :, 1:], axis=3),
+ transpose_indices,
+ ) # Take the maximum probability over the vertical dimension and transpose the predictions
+
+ # Add predictions to the map
+ for timestep in range(num_timesteps):
+ for fcst_hr_index in range(num_forecast_hours):
+ if args["num_images"] == [1, 1]:
+ stitched_map_probs[timestep] = image_probs
+ if (
+ timestep == num_timesteps - 1
+ and fcst_hr_index == num_forecast_hours - 1
+ ):
+ map_created = True
+ else:
+ stitched_map_probs[timestep][fcst_hr_index], map_created = (
+ _add_image_to_map(
+ stitched_map_probs[timestep][fcst_hr_index],
+ image_probs[
+ timestep * num_forecast_hours + fcst_hr_index
+ ],
+ map_created,
+ num_images_lon,
+ num_images_lat,
+ lon_image,
+ lat_image,
+ image_size_lon,
+ image_size_lat,
+ lon_image_spacing,
+ lat_image_spacing,
+ )
+ )
+
+ if map_created:
+ for timestep_no, timestep in enumerate(timesteps):
+ timestep = str(timestep)
+ for fcst_hr_index, forecast_hour in enumerate(forecast_hours):
+ time = f"{timestep[:4]}%s%s%s" % (
+ timestep[5:7],
+ timestep[8:10],
+ timestep[11:13],
+ )
+ probs_ds = create_model_prediction_dataset(
+ stitched_map_probs[timestep_no][fcst_hr_index],
+ image_lats,
+ image_lons,
+ front_types,
+ )
+ probs_ds = probs_ds.expand_dims(
+ {
+ "time": np.atleast_1d(timestep),
+ "forecast_hour": np.atleast_1d(
+ forecast_hours[fcst_hr_index]
+ ),
+ }
+ )
+ filename_base = "model_%d_%s_%s_%s_f%03d" % (
+ args["model_number"],
+ time,
+ args["domain"],
+ args["data_source"],
+ forecast_hours[fcst_hr_index],
+ )
+
+ pred_folder = "%s/model_%d/predictions" % (
+ args["model_dir"],
+ args["model_number"],
+ )
+ os.makedirs(pred_folder, exist_ok=True)
+ outfile = "%s/%s_probabilities.nc" % (
+ pred_folder,
+ filename_base,
+ )
+ probs_ds.to_netcdf(path=outfile, engine="netcdf4", mode="w")
diff --git a/src/fronts/evaluation/predict_tf.py b/src/fronts/evaluation/predict_tf.py
new file mode 100644
index 0000000..07642f2
--- /dev/null
+++ b/src/fronts/evaluation/predict_tf.py
@@ -0,0 +1,253 @@
+"""
+**** EXPERIMENTAL SCRIPT TO REPLACE 'predict.py' IN THE NEAR FUTURE ****
+
+Generate predictions using a model with tensorflow datasets.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.2.13
+"""
+
+import argparse
+import os
+import numpy as np
+
+from fronts.utils import file_manager as fm
+from fronts.utils import constants
+from fronts.utils.misc import initialize_gpus
+import xarray as xr
+import tensorflow as tf
+import pandas as pd
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--dataset",
+ type=str,
+ help="Dataset for which to make predictions. Options are: 'training', 'validation', 'test'",
+ )
+ parser.add_argument(
+ "--year_and_month",
+ type=int,
+ nargs=2,
+ help="Year and month for which to make predictions.",
+ )
+ parser.add_argument("--model_dir", type=str, help="Directory for the models.")
+ parser.add_argument("--model_number", type=int, help="Model number.")
+ parser.add_argument(
+ "--tf_indir",
+ type=str,
+ help="Directory for the tensorflow dataset that will be used when generating predictions.",
+ )
+ parser.add_argument(
+ "--data_source", type=str, default="era5", help="Data source for variables"
+ )
+ parser.add_argument("--gpu_device", type=int, nargs="+", help="GPU device number.")
+ parser.add_argument(
+ "--batch_size",
+ type=int,
+ default=8,
+ help="Batch size for the model predictions.",
+ )
+ parser.add_argument(
+ "--steps",
+ type=int,
+ default=4,
+ help="Number of steps to take when generating the model predictions.",
+ )
+ parser.add_argument(
+ "--memory_growth", action="store_true", help="Use memory growth on the GPU"
+ )
+ parser.add_argument(
+ "--overwrite",
+ action="store_true",
+ help="Overwrite any existing prediction files.",
+ )
+ args = vars(parser.parse_args())
+
+ model_properties = pd.read_pickle(
+ "%s/model_%d/model_%d_properties.pkl"
+ % (args["model_dir"], args["model_number"], args["model_number"])
+ )
+ dataset_properties = pd.read_pickle("%s/dataset_properties.pkl" % args["tf_indir"])
+
+ domain = dataset_properties["domain"]
+
+ hour_interval = 3 if domain == "conus" else 6
+
+ # Some older models do not have the 'dataset_properties' dictionary
+ try:
+ front_types = model_properties["dataset_properties"]["front_types"]
+ num_dims = model_properties["dataset_properties"]["num_dims"]
+ except KeyError:
+ front_types = model_properties["front_types"]
+ if args["model_number"] in [6846496, 7236500, 7507525]:
+ num_dims = (3, 3)
+
+ if args["dataset"] is not None and args["year_and_month"] is not None:
+ raise ValueError("--dataset and --year_and_month cannot be passed together.")
+ elif args["dataset"] is None and args["year_and_month"] is None:
+ raise ValueError(
+ "At least one of [--dataset, --year_and_month] must be passed."
+ )
+ elif args["year_and_month"] is not None:
+ years, months = [args["year_and_month"][0]], [args["year_and_month"][1]]
+ else:
+ years, months = model_properties["%s_years" % args["dataset"]], range(1, 13)
+
+ ### Make sure that the dataset has the same attributes as the model ###
+ if (
+ model_properties["dataset_properties"]["normalization_parameters"]
+ != dataset_properties["normalization_parameters"]
+ ):
+ raise ValueError(
+ "Cannot evaluate model with the selected dataset. Reason: normalization parameters do not match"
+ )
+ if (
+ model_properties["dataset_properties"]["front_types"]
+ != dataset_properties["front_types"]
+ ):
+ raise ValueError(
+ "Cannot evaluate model with the selected dataset. Reason: front types do not match "
+ f"(model: {model_properties['dataset_properties']['front_types']}, dataset: {dataset_properties['front_types']})"
+ )
+ if (
+ model_properties["dataset_properties"]["variables"]
+ != dataset_properties["variables"]
+ ):
+ raise ValueError(
+ "Cannot evaluate model with the selected dataset. Reason: variables do not match "
+ f"(model: {model_properties['dataset_properties']['variables']}, dataset: {dataset_properties['variables']})"
+ )
+ if (
+ model_properties["dataset_properties"]["pressure_levels"]
+ != dataset_properties["pressure_levels"]
+ ):
+ raise ValueError(
+ "Cannot evaluate model with the selected dataset. Reason: pressure levels do not match "
+ f"(model: {model_properties['dataset_properties']['pressure_levels']}, dataset: {dataset_properties['pressure_levels']})"
+ )
+
+ if args["gpu_device"] is not None:
+ initialize_gpus(args["gpu_device"], args["memory_growth"])
+ else:
+ os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+ # The axis that the predicts will be concatenated on depends on the shape of the output, which is determined by deep supervision
+ concat_axis = 1 if model_properties["deep_supervision"] else 0
+
+ if args["data_source"] in ["era5", "gfs", "gdas"]:
+ if dataset_properties["override_extent"] is not None:
+ extent = dataset_properties["override_extent"]
+ lats = np.arange(extent[2], extent[3] + 0.25, 0.25)[::-1]
+ lons = np.arange(extent[0], extent[1] + 0.25, 0.25)
+ else:
+ lats = np.arange(
+ constants.DOMAIN_EXTENTS[domain][2],
+ constants.DOMAIN_EXTENTS[domain][3] + 0.25,
+ 0.25,
+ )[::-1]
+ lons = np.arange(
+ constants.DOMAIN_EXTENTS[domain][0],
+ constants.DOMAIN_EXTENTS[domain][1] + 0.25,
+ 0.25,
+ )
+ elif args["data_source"] == "hrrr":
+ hrrr_coords = xr.open_dataset("%s/coordinates/hrrr.nc" % os.getcwd())
+ lats = hrrr_coords["latitude"][:1056, :1728].to_numpy()
+ lons = hrrr_coords["longitude"][:1056, :1728].to_numpy()
+
+ model = fm.load_model(args["model_number"], args["model_dir"])
+
+ for year in years:
+ tf_ds_obj = fm.DataFileLoader(
+ args["tf_indir"], data_type="inputs", file_format="tensorflow", years=years
+ )
+ files_for_year = tf_ds_obj.files[0]
+
+ for month in months:
+ prediction_dataset_path = (
+ "%s/model_%d/probabilities/model_%d_pred_%s_%d%02d.nc"
+ % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ domain,
+ year,
+ month,
+ )
+ )
+ if os.path.isfile(prediction_dataset_path) and not args["overwrite"]:
+ print(
+ "WARNING: %s exists, pass the --overwrite argument to overwrite existing data."
+ % prediction_dataset_path
+ )
+ continue
+
+ time_array = pd.read_pickle(
+ "%s/timesteps_%d%02d.pkl" % (args["tf_indir"], year, month)
+ )
+
+ if args["data_source"] in ["era5", "gfs", "gdas"]:
+ xr_ds_coords = (
+ ("time", "latitude", "longitude"),
+ {"time": time_array, "latitude": lats, "longitude": lons},
+ )
+ else:
+ xr_ds_coords = (
+ ("time", "y", "x"),
+ {
+ "time": time_array,
+ "latitude": (("y", "x"), lats),
+ "longitude": (("y", "x"), lons),
+ },
+ )
+
+ input_file = [
+ file for file in files_for_year if "_%d%02d" % (year, month) in file
+ ][0]
+ tf_ds = tf.data.Dataset.load(input_file)
+ # tf_ds = tf_ds.batch(args['batch_size'])
+
+ num_timesteps = len(time_array)
+ timestep_indices = np.linspace(0, num_timesteps, args["steps"] + 1).astype(
+ int
+ )
+
+ # generate model predictions
+ predictions = []
+ for step in range(args["steps"]):
+ tf_ds_step = tf_ds.skip(timestep_indices[step]).take(
+ timestep_indices[step + 1] - timestep_indices[step]
+ )
+ prediction = np.array(
+ model.predict(tf_ds_step, batch_size=args["batch_size"])
+ ).astype("float16")
+ predictions.append(prediction)
+ predictions = np.concatenate(predictions, axis=1)
+
+ if model_properties["deep_supervision"]:
+ predictions = predictions[
+ 0, ...
+ ] # select the top output of the model, since it is the only one we care about
+
+ if num_dims[1] == 3:
+ # Take the maxmimum probability for each front type over the vertical dimension (pressure levels)
+ predictions = np.amax(
+ predictions, axis=3
+ ) # shape: (time, latitude, longitude, front type)
+
+ predictions = predictions[
+ ..., 1:
+ ] # remove the 'no front' type from the array
+
+ xr.Dataset(
+ data_vars={
+ front_type: (xr_ds_coords[0], predictions[:, :, :, front_type_no])
+ for front_type_no, front_type in enumerate(front_types)
+ },
+ coords=xr_ds_coords[1],
+ ).astype("float32").to_netcdf(
+ path=prediction_dataset_path, mode="w", engine="netcdf4"
+ )
+
+ del predictions # Delete the predictions variable so it can be recreated for the next month
diff --git a/src/fronts/evaluation/prediction_plot.py b/src/fronts/evaluation/prediction_plot.py
new file mode 100644
index 0000000..3ab9177
--- /dev/null
+++ b/src/fronts/evaluation/prediction_plot.py
@@ -0,0 +1,497 @@
+"""
+Plot model predictions.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.2.12
+"""
+
+import itertools
+import argparse
+import pandas as pd
+import cartopy.crs as ccrs
+import matplotlib.pyplot as plt
+import numpy as np
+import xarray as xr
+from matplotlib import (
+ cm,
+ colors,
+) # Here we explicitly import the cm and color modules to suppress a PyCharm bug
+import os
+from fronts.utils import data_utils
+from fronts.utils.plotting import plot_background
+from skimage.morphology import skeletonize
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--init_time",
+ type=int,
+ nargs=4,
+ help="Date and time of the data. Pass 4 ints in the following order: year, month, day, hour",
+ )
+ parser.add_argument("--domain", type=str, required=True, help="Domain of the data.")
+ parser.add_argument(
+ "--forecast_hour", type=int, help="Forecast hour for the GDAS data"
+ )
+ parser.add_argument(
+ "--front_dilation",
+ type=int,
+ default=1,
+ help="Number of pixels to expand the fronts by in all directions.",
+ )
+ parser.add_argument(
+ "--model_dir", type=str, required=True, help="Directory for the models."
+ )
+ parser.add_argument("--model_number", type=int, required=True, help="Model number.")
+ parser.add_argument(
+ "--fronts_netcdf_indir",
+ type=str,
+ help="Main directory for the netcdf files containing frontal objects.",
+ )
+ parser.add_argument(
+ "--data_source",
+ type=str,
+ default="era5",
+ help="Source of the variable data (ERA5, GDAS, etc.)",
+ )
+ parser.add_argument(
+ "--prob_mask",
+ type=float,
+ nargs=2,
+ default=[0.1, 0.1],
+ help="Probability mask and the step/interval for the probability contours. Probabilities smaller than the mask will not be plotted.",
+ )
+ parser.add_argument(
+ "--calibration",
+ type=int,
+ help="Neighborhood calibration distance in kilometers. Possible neighborhoods are 50, 100, 150, 200, and 250 km.",
+ )
+ parser.add_argument(
+ "--deterministic", action="store_true", help="Plot deterministic splines."
+ )
+ parser.add_argument(
+ "--targets", action="store_true", help="Plot ground truth targets/labels."
+ )
+ parser.add_argument(
+ "--open_contours", action="store_true", help="Plot probability contours."
+ )
+ parser.add_argument(
+ "--filled_contours", action="store_true", help="Plot probability contours."
+ )
+
+ args = vars(parser.parse_args())
+
+ if args["deterministic"] and args["targets"]:
+ raise TypeError(
+ "Cannot plot deterministic splines and ground truth targets at the same time. Only one of --deterministic, --targets may be passed"
+ )
+
+ DEFAULT_COLORBAR_POSITION = {
+ "conus": 0.75,
+ "full": 0.85,
+ "goes-merged": 0.88,
+ "hrrr": 0.77,
+ "MERGIR": 0.9,
+ "global": 0.74,
+ }
+ cbar_position = DEFAULT_COLORBAR_POSITION[args["domain"]]
+
+ model_properties = pd.read_pickle(
+ f"{args['model_dir']}/model_{args['model_number']}/model_{args['model_number']}_properties.pkl"
+ )
+
+ args["data_source"] = args["data_source"].lower()
+
+ if args["data_source"] in ["hrrr", "nam-12km"]:
+ transpose_dims = ("y", "x")
+ else:
+ transpose_dims = ("latitude", "longitude")
+
+ extent = data_utils.DOMAIN_EXTENTS[args["domain"]]
+
+ year, month, day, hour = (
+ args["init_time"][0],
+ args["init_time"][1],
+ args["init_time"][2],
+ args["init_time"][3],
+ )
+ timestep = np.datetime64("%d-%02d-%02dT%02d" % (year, month, day, hour)).astype(
+ object
+ )
+
+ ### Attempt to pull predictions from a monthly netcdf file generated with tensorflow datasets, otherwise try to pull a single netcdf file ###
+ try:
+ probs_file = (
+ f"{args['model_dir']}/model_{args['model_number']}/probabilities/model_{args['model_number']}_pred_{args['domain']}_{year}%02d.nc"
+ % month
+ )
+
+ plot_filename = "%s/model_%d/maps/model_%d_%d%02d%02d%02d_%s.png" % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ year,
+ month,
+ day,
+ hour,
+ args["domain"],
+ )
+ probs_ds = xr.open_mfdataset(probs_file).sel(
+ time=[
+ "%d-%02d-%02dT%02d" % (year, month, day, hour),
+ ]
+ )
+
+ if args["forecast_hour"] is not None:
+ fronts_file = "%s/%d%02d/FrontObjects_%d%02d%02d%02d_f%03d_%s.nc" % (
+ args["fronts_netcdf_indir"],
+ year,
+ month,
+ year,
+ month,
+ day,
+ hour,
+ args["forecast_hour"],
+ args["data_source"],
+ )
+ else:
+ fronts_file = "%s/%d%02d/FrontObjects_%d%02d%02d%02d_%s.nc" % (
+ args["fronts_netcdf_indir"],
+ year,
+ month,
+ year,
+ month,
+ day,
+ hour,
+ args["data_source"],
+ )
+
+ except OSError:
+ probs_dir = f"{args['model_dir']}/model_{args['model_number']}/predictions"
+
+ if args["forecast_hour"] is not None:
+ forecast_timestep = (
+ timestep
+ if args["forecast_hour"] == 0
+ else timestep
+ + np.timedelta64(args["forecast_hour"], "h").astype(object)
+ )
+ new_year, new_month, new_day, new_hour = (
+ forecast_timestep.year,
+ forecast_timestep.month,
+ forecast_timestep.day,
+ forecast_timestep.hour - (forecast_timestep.hour % 3),
+ )
+ fronts_file = "%s/%s%s/FrontObjects_%s%s%s%02d_f%03d_full.nc" % (
+ args["fronts_netcdf_indir"],
+ new_year,
+ new_month,
+ new_year,
+ new_month,
+ new_day,
+ new_hour,
+ args["forecast_hour"],
+ )
+ filename_base = f"model_%d_{year}%02d%02d%02d_%s_%s_f%03d" % (
+ args["model_number"],
+ month,
+ day,
+ hour,
+ args["domain"],
+ args["data_source"],
+ args["forecast_hour"],
+ )
+ else:
+ fronts_file = "%s/%d%02d/FrontObjects_%d%02d%02d%02d_full.nc" % (
+ args["fronts_netcdf_indir"],
+ year,
+ month,
+ year,
+ month,
+ day,
+ hour,
+ )
+ filename_base = f"model_%d_{year}%02d%02d%02d_%s" % (
+ args["model_number"],
+ month,
+ day,
+ hour,
+ args["domain"],
+ )
+ args["data_source"] = "era5"
+
+ plot_filename = "%s/model_%d/maps/%s-same.png" % (
+ args["model_dir"],
+ args["model_number"],
+ filename_base,
+ )
+ probs_file = f"{probs_dir}/{filename_base}_probabilities.nc"
+ probs_ds = xr.open_dataset(probs_file)
+
+ fronts_file = "%s/%d%02d/FrontObjects_%d%02d%02d%02d_full.nc" % (
+ args["fronts_netcdf_indir"],
+ year,
+ month,
+ year,
+ month,
+ day,
+ hour,
+ )
+
+ try:
+ front_types = model_properties["dataset_properties"]["front_types"]
+ except KeyError:
+ front_types = model_properties["front_types"]
+
+ # front_types = ['DL',]
+ labels = front_types
+ fronts_found = False
+
+ if args["targets"]:
+ right_title = "Splines: NOAA fronts"
+ try:
+ if args["data_source"] in ["era5", "gfs", "gdas"]:
+ fronts = xr.open_dataset(fronts_file).sel(
+ longitude=slice(extent[0], extent[1]),
+ latitude=slice(extent[3], extent[2]),
+ )
+ elif args["data_source"] == "hrrr":
+ fronts = xr.open_dataset(fronts_file).isel(
+ y=slice(0, 1056), x=slice(0, 1728)
+ )
+ hrrr_coords = xr.open_dataset(
+ "%s/coordinates/hrrr.nc" % os.getcwd()
+ ).isel(y=slice(0, 1056), x=slice(0, 1728))
+ fronts = fronts.assign_coords(
+ {
+ "latitude": (("y", "x"), hrrr_coords["latitude"].to_numpy()),
+ "longitude": (("y", "x"), hrrr_coords["longitude"].to_numpy()),
+ }
+ )
+ fronts = data_utils.reformat_fronts(fronts, front_types=front_types)
+ fronts = data_utils.expand_fronts(fronts, iterations=args["front_dilation"])
+ labels = fronts.attrs["labels"]
+ fronts = xr.where(fronts == 0, float("NaN"), fronts)
+ fronts_found = True
+ except FileNotFoundError:
+ print("No ground truth fronts found")
+
+ if type(front_types) == str:
+ front_types = [
+ front_types,
+ ]
+
+ mask, prob_int = (
+ args["prob_mask"][0],
+ args["prob_mask"][1],
+ ) # Probability mask, contour interval for probabilities
+ vmax, cbar_tick_adjust, cbar_label_adjust, n_colors = (
+ 1.01,
+ prob_int,
+ 10,
+ (1 / prob_int) + 1,
+ )
+ levels = np.around(np.arange(0, 1 + prob_int, prob_int), 2)
+ cbar_ticks = np.around(np.arange(mask, 1 + prob_int, prob_int), 2)
+
+ contour_maps_by_type = [data_utils.CONTOUR_CMAPS[label] for label in labels]
+ front_colors_by_type = [data_utils.FRONT_COLORS[label] for label in labels]
+ front_names_by_type = [data_utils.FRONT_NAMES[label] for label in labels]
+
+ cmap_front = colors.ListedColormap(
+ front_colors_by_type, name="from_list", N=len(front_colors_by_type)
+ )
+ norm_front = colors.Normalize(vmin=1, vmax=len(front_colors_by_type) + 1)
+
+ # probs_ds = probs_ds.isel(time=0)
+ probs_ds = (
+ probs_ds.isel(time=0)
+ if args["data_source"] == "era5"
+ else probs_ds.isel(time=0, forecast_hour=0)
+ )
+ probs_ds = probs_ds.transpose(*transpose_dims)
+
+ for key in list(probs_ds.keys()):
+ if args["deterministic"]:
+ spline_threshold = model_properties["front_obj_thresholds"][args["domain"]][
+ key
+ ]["100"]
+ probs_ds[f"{key}_obj"] = (
+ transpose_dims,
+ skeletonize(
+ xr.where(probs_ds[key] > spline_threshold, 1, 0).values.copy(
+ order="C"
+ )
+ ),
+ )
+
+ if args["calibration"] is not None:
+ try:
+ ir_model = model_properties["calibration_models"][args["domain"]][key][
+ "%d km" % args["calibration"]
+ ]
+ except KeyError:
+ ir_model = model_properties["calibration_models"]["conus"][key][
+ "%d km" % args["calibration"]
+ ]
+ original_shape = np.shape(probs_ds[key].values)
+ probs_ds[key].values = ir_model.predict(
+ probs_ds[key].values.flatten()
+ ).reshape(original_shape)
+ cbar_label = "Probability (calibrated - %d km)" % args["calibration"]
+ else:
+ cbar_label = "Probability (uncalibrated)"
+
+ if not args["open_contours"]:
+ if len(front_types) > 1:
+ all_possible_front_combinations = itertools.permutations(front_types, r=2)
+ for combination in all_possible_front_combinations:
+ probs_ds[combination[0]].values = np.where(
+ probs_ds[combination[0]].values
+ > probs_ds[combination[1]].values - 0.02,
+ probs_ds[combination[0]].values,
+ 0,
+ )
+
+ probs_ds = xr.where(probs_ds > mask, probs_ds, float("NaN"))
+
+ if args["data_source"] != "era5":
+ valid_time = timestep + np.timedelta64(args["forecast_hour"], "h").astype(
+ object
+ )
+ data_title = (
+ f"Run: {args['data_source'].upper()} {year}-%02d-%02d-%02dz F%03d \nPredictions valid: %d-%02d-%02d-%02dz"
+ % (
+ month,
+ day,
+ hour,
+ args["forecast_hour"],
+ valid_time.year,
+ valid_time.month,
+ valid_time.day,
+ valid_time.hour,
+ )
+ )
+ else:
+ data_title = (
+ "Data: ERA5 reanalysis %d-%02d-%02d-%02dz\n"
+ "Predictions valid: %d-%02d-%02d-%02dz"
+ % (year, month, day, hour, year, month, day, hour)
+ )
+
+ fig, ax = plt.subplots(
+ 1,
+ 1,
+ figsize=(22, 8),
+ subplot_kw={"projection": ccrs.PlateCarree(central_longitude=0)},
+ )
+ plot_background(extent, ax=ax, linewidth=0.5)
+ # ax.gridlines(draw_labels=True, zorder=0)
+
+ cbar_front_labels = []
+ cbar_front_ticks = []
+
+ for front_no, front_key, front_name, front_label, cmap in zip(
+ range(1, len(front_names_by_type) + 1),
+ front_types,
+ front_names_by_type,
+ front_types,
+ contour_maps_by_type,
+ ):
+ if args["filled_contours"]:
+ cmap_probs, norm_probs = (
+ cm.get_cmap(cmap, n_colors),
+ colors.Normalize(vmin=0, vmax=vmax),
+ )
+ probs_ds[front_key].plot.contourf(
+ ax=ax,
+ x="longitude",
+ y="latitude",
+ norm=norm_probs,
+ levels=levels,
+ cmap=cmap_probs,
+ transform=ccrs.PlateCarree(),
+ alpha=0.75,
+ add_colorbar=False,
+ )
+ cbar_ax = fig.add_axes(
+ [cbar_position + (front_no * 0.015), 0.24, 0.015, 0.64]
+ )
+ cbar = plt.colorbar(
+ cm.ScalarMappable(norm=norm_probs, cmap=cmap_probs),
+ cax=cbar_ax,
+ boundaries=levels[1:],
+ alpha=0.75,
+ )
+ cbar.set_ticklabels([])
+
+ if args["open_contours"]:
+ probs_ds[front_key].plot.contour(
+ ax=ax,
+ x="longitude",
+ y="latitude",
+ levels=levels,
+ colors=front_colors_by_type[front_no - 1],
+ transform=ccrs.PlateCarree(),
+ alpha=0.75,
+ add_colorbar=False,
+ )
+
+ if args["deterministic"]:
+ right_title = "Splines: Deterministic first-guess fronts"
+ cmap_deterministic = colors.ListedColormap(
+ ["None", front_colors_by_type[front_no - 1]], name="from_list", N=2
+ )
+ norm_deterministic = colors.Normalize(vmin=0, vmax=1)
+ probs_ds[f"{front_key}_obj"].plot(
+ ax=ax,
+ x="longitude",
+ y="latitude",
+ cmap=cmap_deterministic,
+ norm=norm_deterministic,
+ transform=ccrs.PlateCarree(),
+ alpha=0.9,
+ add_colorbar=False,
+ )
+
+ if fronts_found:
+ fronts["identifier"].plot(
+ ax=ax,
+ x="longitude",
+ y="latitude",
+ cmap=cmap_front,
+ norm=norm_front,
+ transform=ccrs.PlateCarree(),
+ add_colorbar=False,
+ )
+
+ cbar_front_labels.append(front_name)
+ cbar_front_ticks.append(front_no + 0.5)
+
+ if args["filled_contours"]:
+ cbar.set_label(cbar_label, rotation=90)
+ cbar.set_ticks(cbar_ticks)
+ cbar.set_ticklabels(cbar_ticks)
+
+ cbar_front = plt.colorbar(
+ cm.ScalarMappable(norm=norm_front, cmap=cmap_front),
+ ax=ax,
+ alpha=0.75,
+ orientation="horizontal",
+ shrink=0.5,
+ pad=0.02,
+ )
+ cbar_front.set_ticks(cbar_front_ticks)
+ cbar_front.set_ticklabels(cbar_front_labels)
+ cbar_front.set_label(r"$\bf{Front}$ $\bf{type}$")
+
+ if fronts_found or args["deterministic"]:
+ ax.set_title(right_title, loc="right")
+
+ ax.set_title("")
+ ax.set_title(data_title, loc="left")
+ ax.set_title("model number: %d" % args["model_number"], loc="right")
+ ax.set_title("FrontFinder Predictions", loc="right")
+ plt.savefig(plot_filename, bbox_inches="tight", dpi=500)
+ plt.close()
diff --git a/src/fronts/layers/__init__.py b/src/fronts/layers/__init__.py
new file mode 100644
index 0000000..1da1ea2
--- /dev/null
+++ b/src/fronts/layers/__init__.py
@@ -0,0 +1,21 @@
+from fronts.layers.modules import (
+ attention_gate,
+ convolution_module,
+ aggregated_feature_map,
+ full_scale_skip_connection,
+ conventional_skip_connection,
+ max_pool,
+ upsample,
+ deep_supervision_side_output,
+)
+
+__all__ = [
+ "attention_gate",
+ "convolution_module",
+ "aggregated_feature_map",
+ "full_scale_skip_connection",
+ "conventional_skip_connection",
+ "max_pool",
+ "upsample",
+ "deep_supervision_side_output",
+]
diff --git a/src/fronts/layers/activations.py b/src/fronts/layers/activations.py
new file mode 100644
index 0000000..c28ffda
--- /dev/null
+++ b/src/fronts/layers/activations.py
@@ -0,0 +1,792 @@
+"""
+Custom activation function layers:
+ * Elliott
+ * Gaussian
+ * Growing cosine unit (GCU)
+ * Hexpo
+ * Improved sigmoid units (ISigmoid)
+ * Linearly-scaled hyperbolic tangent (LiSHT)
+ * Parametric sigmoid (PSigmoid)
+ * Parametric hyperbolic tangent (PTanh)
+ * Parametric tangent hyperbolic linear unit (PTELU)
+ * Rectified hyperbolic secant (ReSech)
+ * Smooth rectified linear unit (SmeLU)
+ * Snake
+ * Soft-root-sign (SRS)
+ * Scaled hyperbolic tangent (STanh)
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2024.5.18
+
+TODO: thoroughly test all activation functions
+"""
+
+from tensorflow.keras.layers import Layer
+import tensorflow as tf
+
+
+class Elliott(Layer):
+ """
+ Elliott activation layer.
+
+ References
+ ----------
+ https://link.springer.com/article/10.1007/s00521-017-3210-6
+ """
+
+ def __init__(self, name=None, **kwargs):
+ super(Elliott, self).__init__(name=name, **kwargs)
+
+ def build(self, input_shape):
+ """Build the Elliott activation layer"""
+
+ def call(self, inputs):
+ inputs = tf.cast(inputs, "float32")
+ y = 0.5 * inputs / (1.0 + tf.abs(inputs)) + 0.5
+
+ return y
+
+
+class Gaussian(Layer):
+ """
+ Gaussian function activation layer.
+ """
+
+ def __init__(self, name=None, **kwargs):
+ super(Gaussian, self).__init__(name=name, **kwargs)
+
+ def build(self, input_shape):
+ """Build the Gaussian layer"""
+
+ def call(self, inputs):
+ """Call the Gaussian activation function"""
+ inputs = tf.cast(inputs, "float32")
+ y = tf.exp(-tf.square(inputs))
+
+ return y
+
+
+class GCU(Layer):
+ """
+ Growing Cosine Unit (GCU) activation layer.
+ """
+
+ def __init__(self, name=None, **kwargs):
+ super(GCU, self).__init__(name=name, **kwargs)
+
+ def build(self, input_shape):
+ """Build the GCU layer"""
+
+ def call(self, inputs):
+ """Call the GCU activation function"""
+ inputs = tf.cast(inputs, "float32")
+ y = inputs * tf.cos(inputs)
+
+ return y
+
+
+class Hexpo(Layer):
+ """
+ Hexpo activation layer.
+
+ References
+ ----------
+ https://ieeexplore.ieee.org/document/7966168
+
+ Notes
+ -----
+ When referencing the above paper, we name the parameters the following (paper -> our name):
+ a -> alpha
+ b -> beta
+ c -> gamma
+ d -> delta
+
+
+ """
+
+ def __init__(
+ self,
+ name=None,
+ alpha_initializer=None,
+ alpha_regularizer=None,
+ alpha_constraint=None,
+ beta_initializer="Ones",
+ beta_regularizer=None,
+ beta_constraint="NonNeg",
+ gamma_initializer=None,
+ gamma_regularizer=None,
+ gamma_constraint=None,
+ delta_initializer="Ones",
+ delta_regularizer=None,
+ delta_constraint="NonNeg",
+ shared_axes=None,
+ **kwargs,
+ ):
+ super(Hexpo, self).__init__(name=name, **kwargs)
+ self._name = name
+ self.alpha_initializer = alpha_initializer
+ self.alpha_regularizer = alpha_regularizer
+ self.alpha_constraint = alpha_constraint
+ self.beta_initializer = beta_initializer
+ self.beta_regularizer = beta_regularizer
+ self.beta_constraint = beta_constraint
+ self.gamma_initializer = gamma_initializer
+ self.gamma_regularizer = gamma_regularizer
+ self.gamma_constraint = gamma_constraint
+ self.delta_initializer = delta_initializer
+ self.delta_regularizer = delta_regularizer
+ self.delta_constraint = delta_constraint
+ self.shared_axes = shared_axes
+
+ def build(self, input_shape):
+ """Build the PSigmoid layer"""
+ param_shape = list(input_shape[1:])
+ if self.shared_axes is not None:
+ for ax in self.shared_axes:
+ param_shape[ax - 1] = 1
+ else:
+ # Turn all arbitrary dimensions (denoted by None) into size 1
+ for ax in range(len(param_shape)):
+ param_shape[ax] = 1 if param_shape[ax] is None else param_shape[ax]
+
+ # learnable parameter
+ self.alpha = self.add_weight(
+ name="alpha",
+ shape=param_shape,
+ initializer=self.alpha_initializer,
+ regularizer=self.alpha_regularizer,
+ constraint=self.alpha_constraint,
+ )
+ self.beta = self.add_weight(
+ name="beta",
+ shape=param_shape,
+ initializer=self.beta_initializer,
+ regularizer=self.beta_regularizer,
+ constraint=self.beta_constraint,
+ )
+ self.gamma = self.add_weight(
+ name="gamma",
+ shape=param_shape,
+ initializer=self.gamma_initializer,
+ regularizer=self.gamma_regularizer,
+ constraint=self.gamma_constraint,
+ )
+ self.delta = self.add_weight(
+ name="delta",
+ shape=param_shape,
+ initializer=self.delta_initializer,
+ regularizer=self.delta_regularizer,
+ constraint=self.delta_constraint,
+ )
+
+ def call(self, inputs):
+ inputs = tf.cast(inputs, "float32")
+ y = tf.where(
+ inputs >= 0.0,
+ -self.alpha * (tf.exp(-inputs / (self.beta + 1e-7)) - 1.0),
+ self.gamma * (tf.exp(inputs / (self.delta + 1e-7)) - 1.0),
+ )
+
+ return y
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "name": self._name,
+ "alpha_initializer": self.alpha_initializer,
+ "alpha_regularizer": self.alpha_regularizer,
+ "alpha_constraint": self.alpha_constraint,
+ "beta_initializer": self.beta_initializer,
+ "beta_regularizer": self.beta_regularizer,
+ "beta_constraint": self.beta_constraint,
+ "gamma_initializer": self.gamma_initializer,
+ "gamma_regularizer": self.gamma_regularizer,
+ "gamma_constraint": self.gamma_constraint,
+ "delta_initializer": self.delta_initializer,
+ "delta_regularizer": self.delta_regularizer,
+ "delta_constraint": self.delta_constraint,
+ "shared_axes": self.shared_axes,
+ }
+ )
+
+ return config
+
+
+class ISigmoid(Layer):
+ """
+ Trainable version of the ISigmoid activation function layer.
+
+ References
+ ----------
+ https://ieeexplore.ieee.org/document/8415753
+
+ Notes
+ -----
+ Parameter 'a' in the paper referenced above will be called 'beta' in this layer.
+ """
+
+ def __init__(
+ self,
+ name=None,
+ alpha_initializer="zeros",
+ alpha_regularizer=None,
+ alpha_constraint=None,
+ beta_initializer="zeros",
+ beta_regularizer=None,
+ beta_constraint=None,
+ shared_axes=None,
+ **kwargs,
+ ):
+ super(ISigmoid, self).__init__(name=name, **kwargs)
+ self.alpha_initializer = alpha_initializer
+ self.alpha_regularizer = alpha_regularizer
+ self.alpha_constraint = alpha_constraint
+ self.beta_initializer = beta_initializer
+ self.beta_regularizer = beta_regularizer
+ self.beta_constraint = beta_constraint
+ self.shared_axes = shared_axes
+
+ def build(self, input_shape):
+ """Build the ISigmoid layer"""
+ param_shape = list(input_shape[1:])
+ if self.shared_axes is not None:
+ for ax in self.shared_axes:
+ param_shape[ax - 1] = 1
+ else:
+ # Turn all arbitrary dimensions (denoted by None) into size 1
+ for ax in range(len(param_shape)):
+ param_shape[ax] = 1 if param_shape[ax] is None else param_shape[ax]
+
+ self.alpha = self.add_weight(
+ name="alpha",
+ shape=param_shape,
+ initializer=self.alpha_initializer,
+ regularizer=self.alpha_regularizer,
+ constraint=self.alpha_constraint,
+ )
+ self.beta = self.add_weight(
+ name="beta",
+ shape=param_shape,
+ initializer=self.beta_initializer,
+ regularizer=self.beta_regularizer,
+ constraint=self.beta_constraint,
+ )
+
+ def call(self, inputs):
+ inputs = tf.cast(inputs, "float32")
+ y = tf.where(
+ inputs >= self.beta,
+ (self.alpha * (inputs - self.beta)) + tf.sigmoid(self.beta),
+ tf.where(
+ inputs <= -self.beta,
+ (self.alpha * (inputs + self.beta)) + tf.sigmoid(self.beta),
+ tf.sigmoid(inputs),
+ ),
+ )
+
+ return y
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "name": self._name,
+ "alpha_initializer": self.alpha_initializer,
+ "alpha_regularizer": self.alpha_regularizer,
+ "alpha_constraint": self.alpha_constraint,
+ "beta_initializer": self.beta_initializer,
+ "beta_regularizer": self.beta_regularizer,
+ "beta_constraint": self.beta_constraint,
+ "shared_axes": self.shared_axes,
+ }
+ )
+
+ return config
+
+
+class LiSHT(Layer):
+ """
+ Linearly-scaled hyperbolic tangent activation layer.
+
+ References
+ ----------
+ https://arxiv.org/abs/1901.05894
+ """
+
+ def __init__(self, name=None, **kwargs):
+ super(LiSHT, self).__init__(name=name, **kwargs)
+
+ def build(self, input_shape):
+ """Build the LiSHT layer"""
+
+ def call(self, inputs):
+ inputs = tf.cast(inputs, "float32")
+ y = inputs * tf.tanh(inputs)
+
+ return y
+
+
+class PSigmoid(Layer):
+ """
+ Parametric sigmoid activation layer.
+ """
+
+ def __init__(
+ self,
+ name=None,
+ alpha_initializer="zeros",
+ alpha_regularizer=None,
+ alpha_constraint=None,
+ shared_axes=None,
+ **kwargs,
+ ):
+ super(PSigmoid, self).__init__(name=name, **kwargs)
+ self.alpha_initializer = alpha_initializer
+ self.alpha_regularizer = alpha_regularizer
+ self.alpha_constraint = alpha_constraint
+ self.shared_axes = shared_axes
+
+ def build(self, input_shape):
+ """Build the PSigmoid layer"""
+ param_shape = list(input_shape[1:])
+ if self.shared_axes is not None:
+ for ax in self.shared_axes:
+ param_shape[ax - 1] = 1
+ else:
+ # Turn all arbitrary dimensions (denoted by None) into size 1
+ for ax in range(len(param_shape)):
+ param_shape[ax] = 1 if param_shape[ax] is None else param_shape[ax]
+
+ self.alpha = self.add_weight(
+ name="alpha",
+ shape=param_shape,
+ initializer=self.alpha_initializer,
+ regularizer=self.alpha_regularizer,
+ constraint=self.alpha_constraint,
+ ) # learnable parameter
+
+ def call(self, inputs):
+ inputs = tf.cast(inputs, "float32")
+ y = tf.sigmoid(inputs) ** self.alpha
+
+ return y
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "name": self._name,
+ "alpha_initializer": self.alpha_initializer,
+ "alpha_regularizer": self.alpha_regularizer,
+ "alpha_constraint": self.alpha_constraint,
+ "shared_axes": self.shared_axes,
+ }
+ )
+
+ return config
+
+
+class PTanh(Layer):
+ """
+ Penalized hyperbolic tangent (PTanh) activation layer.
+ """
+
+ def __init__(
+ self,
+ name=None,
+ alpha_initializer="zeros",
+ alpha_regularizer=None,
+ alpha_constraint="MinMaxNorm", # by default, alpha is restricted to the (0, 1) range
+ shared_axes=None,
+ **kwargs,
+ ):
+ super(PTanh, self).__init__(name=name, **kwargs)
+ self.alpha_initializer = alpha_initializer
+ self.alpha_regularizer = alpha_regularizer
+ self.alpha_constraint = alpha_constraint
+ self.shared_axes = shared_axes
+
+ def build(self, input_shape):
+ """Build the PTanh layer"""
+ param_shape = list(input_shape[1:])
+ if self.shared_axes is not None:
+ for ax in self.shared_axes:
+ param_shape[ax - 1] = 1
+ else:
+ # Turn all arbitrary dimensions (denoted by None) into size 1
+ for ax in range(len(param_shape)):
+ param_shape[ax] = 1 if param_shape[ax] is None else param_shape[ax]
+
+ self.alpha = self.add_weight(
+ name="alpha",
+ shape=param_shape,
+ initializer=self.alpha_initializer,
+ regularizer=self.alpha_regularizer,
+ constraint=self.alpha_constraint,
+ ) # learnable parameter
+
+ def call(self, inputs):
+ inputs = tf.cast(inputs, "float32")
+ y = tf.where(inputs >= 0.0, tf.tanh(inputs), self.alpha * tf.tanh(inputs))
+
+ return y
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "name": self._name,
+ "alpha_initializer": self.alpha_initializer,
+ "alpha_regularizer": self.alpha_regularizer,
+ "alpha_constraint": self.alpha_constraint,
+ "shared_axes": self.shared_axes,
+ }
+ )
+
+ return config
+
+
+class PTELU(Layer):
+ """
+ Parametric tangent hyperbolic linear unit activation layer.
+
+ References
+ ----------
+ https://ieeexplore.ieee.org/document/8265328
+ """
+
+ def __init__(
+ self,
+ name=None,
+ alpha_initializer="zeros",
+ alpha_regularizer=None,
+ alpha_constraint="NonNeg",
+ beta_initializer="zeros",
+ beta_regularizer=None,
+ beta_constraint="NonNeg",
+ shared_axes=None,
+ **kwargs,
+ ):
+ super(PTELU, self).__init__(name=name, **kwargs)
+ self.alpha_initializer = alpha_initializer
+ self.alpha_regularizer = alpha_regularizer
+ self.alpha_constraint = alpha_constraint
+ self.beta_initializer = beta_initializer
+ self.beta_regularizer = beta_regularizer
+ self.beta_constraint = beta_constraint
+ self.shared_axes = shared_axes
+
+ def build(self, input_shape):
+ """Build the PTELU layer"""
+ param_shape = list(input_shape[1:])
+ if self.shared_axes is not None:
+ for ax in self.shared_axes:
+ param_shape[ax - 1] = 1
+ else:
+ # Turn all arbitrary dimensions (denoted by None) into size 1
+ for ax in range(len(param_shape)):
+ param_shape[ax] = 1 if param_shape[ax] is None else param_shape[ax]
+
+ # learnable parameter
+ self.alpha = self.add_weight(
+ name="alpha",
+ shape=param_shape,
+ initializer=self.alpha_initializer,
+ regularizer=self.alpha_regularizer,
+ constraint=self.alpha_constraint,
+ )
+ self.beta = self.add_weight(
+ name="beta",
+ shape=param_shape,
+ initializer=self.beta_initializer,
+ regularizer=self.beta_regularizer,
+ constraint=self.beta_constraint,
+ )
+
+ def call(self, inputs):
+ inputs = tf.cast(inputs, "float32")
+ y = tf.where(inputs >= 0.0, inputs, self.alpha * tf.tanh(self.beta * inputs))
+
+ return y
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "name": self._name,
+ "alpha_initializer": self.alpha_initializer,
+ "alpha_regularizer": self.alpha_regularizer,
+ "alpha_constraint": self.alpha_constraint,
+ "beta_initializer": self.beta_initializer,
+ "beta_regularizer": self.beta_regularizer,
+ "beta_constraint": self.beta_constraint,
+ "shared_axes": self.shared_axes,
+ }
+ )
+
+ return config
+
+
+class ReSech(Layer):
+ """
+ Rectified hyperbolic secant (ReSech) activation layer.
+ """
+
+ def __init__(self, name=None, **kwargs):
+ super(ReSech, self).__init__(name=name, **kwargs)
+
+ def build(self, input_shape):
+ """Build the ReSech layer"""
+
+ def call(self, inputs):
+ inputs = tf.cast(inputs, "float32")
+ y = inputs / tf.cosh(inputs)
+
+ return y
+
+
+class SmeLU(Layer):
+ """
+ SmeLU (Smooth ReLU) activation function layer for deep learning models.
+
+ References
+ ----------
+ https://arxiv.org/pdf/2202.06499.pdf
+ """
+
+ def __init__(
+ self,
+ name=None,
+ beta_initializer="ones",
+ beta_regularizer=None,
+ beta_constraint="NonNeg",
+ shared_axes=None,
+ **kwargs,
+ ):
+ super(SmeLU, self).__init__(name=name, **kwargs)
+ self._name = name
+ self.beta_initializer = beta_initializer
+ self.beta_regularizer = beta_regularizer
+ self.beta_constraint = beta_constraint
+ self.shared_axes = shared_axes
+
+ def build(self, input_shape):
+ """Build the SmeLU layer"""
+ param_shape = list(input_shape[1:])
+ if self.shared_axes is not None:
+ for ax in self.shared_axes:
+ param_shape[ax - 1] = 1
+ else:
+ # Turn all abritrary dimensions (denoted by None) into size 1
+ for ax in range(len(param_shape)):
+ param_shape[ax] = 1 if param_shape[ax] is None else param_shape[ax]
+
+ self.beta = self.add_weight(
+ name="beta",
+ shape=param_shape,
+ initializer=self.beta_initializer,
+ regularizer=self.beta_regularizer,
+ constraint=self.beta_constraint,
+ ) # Learnable parameter (see Eq. 7 in the linked paper above)
+
+ def call(self, inputs):
+ """Call the SmeLU activation function"""
+ inputs = tf.cast(inputs, "float32")
+ y = tf.where(
+ inputs <= -self.beta,
+ 0.0, # Condition 1
+ tf.where(
+ tf.abs(inputs) <= self.beta,
+ tf.square(inputs + self.beta) / (4.0 * self.beta), # Condition 2
+ inputs,
+ ),
+ ) # Condition 3 (if x >= beta)
+
+ return y
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "name": self._name,
+ "beta_initializer": self.beta_initializer,
+ "beta_regularizer": self.beta_regularizer,
+ "beta_constraint": self.beta_constraint,
+ "shared_axes": self.shared_axes,
+ }
+ )
+
+ return config
+
+
+class Snake(Layer):
+ """
+ Snake activation function layer for deep learning models.
+
+ References
+ ----------
+ https://arxiv.org/pdf/2006.08195.pdf
+ """
+
+ def __init__(
+ self,
+ name=None,
+ alpha_initializer="Ones",
+ alpha_regularizer=None,
+ alpha_constraint="NonNeg",
+ shared_axes=None,
+ **kwargs,
+ ):
+ super(Snake, self).__init__(name=name, **kwargs)
+ self.alpha_initializer = alpha_initializer
+ self.alpha_regularizer = alpha_regularizer
+ self.alpha_constraint = alpha_constraint
+ self.shared_axes = shared_axes
+
+ def build(self, input_shape):
+ """Build the Snake layer"""
+ param_shape = list(input_shape[1:])
+ if self.shared_axes is not None:
+ for ax in self.shared_axes:
+ param_shape[ax - 1] = 1
+ else:
+ # Turn all arbitrary dimensions (denoted by None) into size 1
+ for ax in range(len(param_shape)):
+ param_shape[ax] = 1 if param_shape[ax] is None else param_shape[ax]
+
+ # Learnable parameter (see Eq. 3 in the linked paper above)
+ self.alpha = self.add_weight(
+ name="alpha",
+ shape=param_shape,
+ initializer=self.alpha_initializer,
+ regularizer=self.alpha_regularizer,
+ constraint=self.alpha_constraint,
+ )
+
+ def call(self, inputs):
+ """Call the Snake activation function"""
+ inputs = tf.cast(inputs, "float32")
+ y = inputs + (
+ (1.0 / (self.alpha + 1e-7)) * tf.square(tf.sin(self.alpha * inputs))
+ )
+
+ return y
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "name": self._name,
+ "alpha_initializer": self.alpha_initializer,
+ "alpha_regularizer": self.alpha_regularizer,
+ "alpha_constraint": self.alpha_constraint,
+ "shared_axes": self.shared_axes,
+ }
+ )
+
+ return config
+
+
+class SRS(Layer):
+ """
+ Soft-Root-Sign activation layer.
+
+ References
+ ----------
+ https://arxiv.org/abs/2003.00547
+ """
+
+ def __init__(
+ self,
+ name=None,
+ alpha_initializer="ones",
+ alpha_regularizer=None,
+ alpha_constraint="NonNeg",
+ beta_initializer="ones",
+ beta_regularizer=None,
+ beta_constraint="NonNeg",
+ shared_axes=None,
+ **kwargs,
+ ):
+ super(SRS, self).__init__(name=name, **kwargs)
+ self.alpha_initializer = alpha_initializer
+ self.alpha_regularizer = alpha_regularizer
+ self.alpha_constraint = alpha_constraint
+ self.beta_initializer = beta_initializer
+ self.beta_regularizer = beta_regularizer
+ self.beta_constraint = beta_constraint
+ self.shared_axes = shared_axes
+
+ def build(self, input_shape):
+ """Build the SRS layer"""
+ param_shape = list(input_shape[1:])
+ if self.shared_axes is not None:
+ for ax in self.shared_axes:
+ param_shape[ax - 1] = 1
+ else:
+ # Turn all arbitrary dimensions (denoted by None) into size 1
+ for ax in range(len(param_shape)):
+ param_shape[ax] = 1 if param_shape[ax] is None else param_shape[ax]
+
+ # learnable parameter
+ self.alpha = self.add_weight(
+ name="alpha",
+ shape=param_shape,
+ initializer=self.alpha_initializer,
+ regularizer=self.alpha_regularizer,
+ constraint=self.alpha_constraint,
+ )
+ self.beta = self.add_weight(
+ name="beta",
+ shape=param_shape,
+ initializer=self.beta_initializer,
+ regularizer=self.beta_regularizer,
+ constraint=self.beta_constraint,
+ )
+
+ def call(self, inputs):
+ inputs = tf.cast(inputs, "float32")
+ y = inputs / ((inputs / self.alpha) + tf.exp(-inputs / self.beta))
+
+ return y
+
+ def get_config(self):
+ config = super().get_config()
+ config.update(
+ {
+ "name": self._name,
+ "alpha_initializer": self.alpha_initializer,
+ "alpha_regularizer": self.alpha_regularizer,
+ "alpha_constraint": self.alpha_constraint,
+ "beta_initializer": self.beta_initializer,
+ "beta_regularizer": self.beta_regularizer,
+ "beta_constraint": self.beta_constraint,
+ "shared_axes": self.shared_axes,
+ }
+ )
+
+ return config
+
+
+class STanh(Layer):
+ """
+ Scaled hyperbolic tangent function.
+
+ References
+ ----------
+ https://ieeexplore.ieee.org/document/726791
+ """
+
+ def __init__(self, name=None, **kwargs):
+ super(STanh, self).__init__(name=name, **kwargs)
+
+ def build(self, input_shape):
+ """Build the STanh layer"""
+
+ def call(self, inputs):
+ """Call the STanh activation function"""
+ inputs = tf.cast(inputs, "float32")
+ y = 1.7159 * tf.tanh((2 / 3) * inputs)
+
+ return y
diff --git a/src/fronts/layers/losses.py b/src/fronts/layers/losses.py
new file mode 100644
index 0000000..b0fe78f
--- /dev/null
+++ b/src/fronts/layers/losses.py
@@ -0,0 +1,255 @@
+"""
+Custom loss functions for U-Net models.
+ - Brier Skill Score (BSS)
+ - Critical Success Index (CSI)
+ - Fractions Skill Score (FSS)
+ - Probability of Detection (POD)
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.3.5
+"""
+
+import tensorflow as tf
+
+
+def brier_skill_score(
+ alpha: int | float = 1.0,
+ beta: int | float = 0.5,
+ class_weights: list[int | float, ...] = None,
+):
+ """
+ Brier skill score (BSS) loss function.
+
+ alpha: int or float
+ Parameter that controls how steep the sigmoid function is for discretization. Higher alpha makes the sigmoid function
+ steeper and can help prevent the training process from stalling. Default value is 1. Values greater than 4 are
+ not recommended.
+ beta: int or float
+ Parameter used to control some behaviors of the sigmoid discretization function. Default and recommended value is 0.5.
+ class_weights: list of values or None
+ List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
+ """
+
+ if class_weights is not None:
+ class_weights = tf.cast(class_weights, tf.float32)
+
+ @tf.function
+ def bss_loss(y_true, y_pred):
+ """
+ y_true: tf.Tensor
+ One-hot encoded tensor containing labels.
+ y_pred: tf.Tensor
+ Tensor containing model predictions.
+ """
+
+ # discretize model predictions and labels
+ y_true = tf.math.sigmoid(alpha * (y_true - beta))
+ y_pred = tf.math.sigmoid(alpha * (y_pred - beta))
+
+ losses = tf.math.square(tf.subtract(y_true, y_pred))
+
+ if class_weights is not None:
+ losses *= class_weights
+
+ brier_score_loss = tf.math.reduce_sum(losses) / tf.size(losses)
+ return brier_score_loss
+
+ return bss_loss
+
+
+def critical_success_index(
+ alpha: int | float = 1.0,
+ beta: int | float = 0.5,
+ class_weights: list[int | float, ...] = None,
+):
+ """
+ Critical Success Index (CSI) loss function.
+
+ alpha: int or float
+ Parameter that controls how steep the sigmoid function is for discretization. Higher alpha makes the sigmoid function
+ steeper and can help prevent the training process from stalling. Default value is 1. Values greater than 4 are
+ not recommended.
+ beta: int or float
+ Parameter used to control some behaviors of the sigmoid discretization function. Default and recommended value is 0.5.
+ class_weights: list of values or None
+ List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
+ """
+
+ if class_weights is not None:
+ class_weights = tf.cast(class_weights, tf.float32)
+
+ @tf.function
+ def csi_loss(y_true, y_pred):
+ """
+ y_true: tf.Tensor
+ One-hot encoded tensor containing labels.
+ y_pred: tf.Tensor
+ Tensor containing model predictions.
+ """
+
+ # discretize model predictions and labels
+ y_true = tf.math.sigmoid(alpha * (y_true - beta))
+ y_pred = tf.math.sigmoid(alpha * (y_pred - beta))
+
+ y_pred_neg = 1 - y_pred
+ y_true_neg = 1 - y_true
+
+ sum_over_axes = tf.range(
+ tf.rank(y_pred) - 1
+ ) # Indices for axes to sum over. Excludes the final (class) dimension.
+
+ true_positives = tf.math.reduce_sum(y_pred * y_true, axis=sum_over_axes)
+ false_negatives = tf.math.reduce_sum(y_pred_neg * y_true, axis=sum_over_axes)
+ false_positives = tf.math.reduce_sum(y_pred * y_true_neg, axis=sum_over_axes)
+
+ if class_weights is not None:
+ true_positives *= class_weights
+ false_positives *= class_weights
+ false_negatives *= class_weights
+
+ csi = tf.math.divide(
+ tf.math.reduce_sum(true_positives),
+ tf.math.reduce_sum(true_positives)
+ + tf.math.reduce_sum(false_positives)
+ + tf.math.reduce_sum(false_negatives),
+ )
+
+ return 1 - csi
+
+ return csi_loss
+
+
+def fractions_skill_score(
+ mask_size: int | tuple[int, ...] | list[int, ...] = (3, 3),
+ alpha: int | float = 1.0,
+ beta: int | float = 0.5,
+ class_weights: list[int | float, ...] = None,
+):
+ """
+ Fractions skill score loss function.
+
+ Parameters
+ ----------
+ mask_size: int or tuple
+ Size of the mask/pool in the AveragePooling layers.
+ alpha: int or float
+ Parameter that controls how steep the sigmoid function is for discretization. Higher alpha makes the sigmoid function
+ steeper and can help prevent the training process from stalling. Default value is 1. Values greater than 4 are
+ not recommended.
+ beta: int or float
+ Parameter used to control some behaviors of the sigmoid discretization function. Default and recommended value is 0.5.
+ class_weights: list of values or None
+ List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
+
+ Returns
+ -------
+ fss_loss: float
+ fss_loss = 1 - fractions skill score
+
+ References
+ ----------
+ (RL2008) Roberts, N. M., and H. W. Lean, 2008: Scale-Selective Verification of Rainfall Accumulations from High-Resolution
+ Forecasts of Convective Events. Mon. Wea. Rev., 136, 78–97, https://doi.org/10.1175/2007MWR2123.1.
+ """
+
+ # keyword arguments for the AveragePooling layer
+ pool_args = dict(pool_size=mask_size, strides=1, padding="same")
+
+ # if mask_size is an int, convert to a tuple. This allows us to check the length of the tuple and pull the correct AveragePooling layer
+ if isinstance(mask_size, int):
+ mask_size = (mask_size,)
+
+ # if mask_size is an list, convert to a tuple
+ elif isinstance(mask_size, list):
+ mask_size = tuple(mask_size)
+
+ # make sure the mask size is between 1 and 3
+ assert 1 <= len(mask_size) <= 3, (
+ "mask_size must have length between 1 and 3, received length %d"
+ % len(mask_size)
+ )
+
+ # get the pooling layer based off the length of the mask_size tuple
+ pool = getattr(tf.keras.layers, "AveragePooling%dD" % len(mask_size))(**pool_args)
+
+ if class_weights is not None:
+ class_weights = tf.cast(class_weights, tf.float32)
+
+ @tf.function
+ def fss_loss(y_true, y_pred):
+ """
+ y_true: tf.Tensor
+ One-hot encoded tensor containing labels.
+ y_pred: tf.Tensor
+ Tensor containing model predictions.
+ """
+
+ if class_weights is not None:
+ y_true *= class_weights
+ y_pred *= class_weights
+
+ # discretize model predictions and labels
+ y_true = tf.math.sigmoid(alpha * (y_true - beta))
+ y_pred = tf.math.sigmoid(alpha * (y_pred - beta))
+
+ O_n = pool(y_true) # observed fractions (Eq. 2 in RL2008)
+ M_n = pool(y_pred) # model forecast fractions (Eq. 3 in RL2008)
+
+ MSE_n = tf.keras.metrics.mean_squared_error(
+ O_n, M_n
+ ) # MSE for model forecast fractions (Eq. 5 in RL2008)
+ MSE_ref = tf.reduce_mean(tf.square(O_n)) + tf.reduce_mean(
+ tf.square(M_n)
+ ) # reference forecast (Eq. 7 in RL2008)
+
+ FSS = 1 - MSE_n / (MSE_ref + 1e-10) # fractions skill score (Eq. 6 in RL2008)
+
+ return 1 - FSS
+
+ return fss_loss
+
+
+def probability_of_detection(class_weights: list[int | float, ...] = None):
+ """
+ Probability of Detection (POD) as a loss function. This turns the function into the miss rate.
+
+ class_weights: list of values or None
+ List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
+
+ NOTE: This function is only intended for use in permutation studies and should NOT be used to train models.
+ """
+
+ @tf.function
+ def pod_loss(y_true, y_pred):
+ """
+ y_true: tf.Tensor
+ One-hot encoded tensor containing labels.
+ y_pred: tf.Tensor
+ Tensor containing model predictions.
+ """
+
+ y_pred_neg = 1 - y_pred
+
+ sum_over_axes = tf.range(
+ tf.rank(y_pred) - 1
+ ) # Indices for axes to sum over. Excludes the final (class) dimension.
+
+ true_positives = tf.math.reduce_sum(y_pred * y_true, axis=sum_over_axes)
+ false_negatives = tf.math.reduce_sum(y_pred_neg * y_true, axis=sum_over_axes)
+
+ if class_weights is not None:
+ relative_class_weights = tf.cast(
+ class_weights / tf.math.reduce_sum(class_weights), tf.float32
+ )
+ pod = tf.math.reduce_sum(
+ tf.math.divide_no_nan(true_positives, true_positives + false_negatives)
+ * relative_class_weights
+ )
+ else:
+ pod = tf.math.reduce_sum(
+ tf.math.divide_no_nan(true_positives, true_positives + false_negatives)
+ )
+
+ return 1 - pod
+
+ return pod_loss
diff --git a/src/fronts/layers/metrics.py b/src/fronts/layers/metrics.py
new file mode 100644
index 0000000..05caad8
--- /dev/null
+++ b/src/fronts/layers/metrics.py
@@ -0,0 +1,348 @@
+"""
+Custom metrics for U-Net models.
+ - Brier Skill Score (BSS)
+ - Critical Success Index (CSI)
+ - Fractions Skill Score (FSS)
+ - Probability of Detection (POD)
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.3.5
+"""
+
+import tensorflow as tf
+
+
+def brier_skill_score(class_weights: list[int | float, ...] = None):
+ """
+ Brier skill score (BSS).
+
+ class_weights: list of values or None
+ List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
+ """
+
+ @tf.function
+ def bss(y_true, y_pred):
+ """
+ y_true: tf.Tensor
+ One-hot encoded tensor containing labels.
+ y_pred: tf.Tensor
+ Tensor containing model predictions.
+ """
+
+ squared_errors = tf.math.square(tf.subtract(y_true, y_pred))
+
+ if class_weights is not None:
+ relative_class_weights = tf.cast(
+ class_weights / tf.math.reduce_sum(class_weights), tf.float32
+ )
+ squared_errors *= relative_class_weights
+
+ bss = 1 - tf.math.reduce_sum(squared_errors) / tf.size(squared_errors)
+
+ return bss
+
+ return bss
+
+
+def critical_success_index(
+ threshold: float = None,
+ window_size: tuple[int, ...] | list[int, ...] = None,
+ class_weights: list[int | float, ...] = None,
+):
+ """
+ Critical success index (CSI).
+
+ threshold: float or None
+ Optional probability threshold that binarizes y_pred. Values in y_pred greater than or equal to the threshold are
+ set to 1, and 0 otherwise.
+ If the threshold is set, it must be greater than 0 and less than 1.
+ window_size: tuple or list of ints or None
+ Pool/kernel size of the max-pooling window for neighborhood statistics. (e.g. if calculating the CSI with a 3-pixel
+ window, this should be set to 3).
+ Note that this parameter is experimental and may return unexpected results.
+ class_weights: list of values or None
+ List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
+ """
+
+ @tf.function
+ def csi(y_true, y_pred):
+ """
+ y_true: tf.Tensor
+ One-hot encoded tensor containing labels.
+ y_pred: tf.Tensor
+ Tensor containing model predictions.
+ """
+
+ if window_size is not None:
+ y_pred = tf.nn.max_pool(
+ y_pred, ksize=window_size, strides=1, padding="VALID"
+ )
+ y_true = tf.nn.max_pool(
+ y_true, ksize=window_size, strides=1, padding="VALID"
+ )
+
+ if threshold is not None:
+ y_pred = tf.where(y_pred >= threshold, 1.0, 0.0)
+
+ y_pred_neg = 1 - y_pred
+ y_true_neg = 1 - y_true
+
+ sum_over_axes = tf.range(
+ tf.rank(y_pred) - 1
+ ) # Indices for axes to sum over. Excludes the final (class) dimension.
+
+ true_positives = tf.math.reduce_sum(y_pred * y_true, axis=sum_over_axes)
+ false_negatives = tf.math.reduce_sum(y_pred_neg * y_true, axis=sum_over_axes)
+ false_positives = tf.math.reduce_sum(y_pred * y_true_neg, axis=sum_over_axes)
+
+ if class_weights is not None:
+ relative_class_weights = tf.cast(
+ class_weights / tf.math.reduce_sum(class_weights), tf.float32
+ )
+ csi = tf.math.reduce_sum(
+ tf.math.divide_no_nan(
+ true_positives, true_positives + false_positives + false_negatives
+ )
+ * relative_class_weights
+ )
+ else:
+ csi = tf.math.divide(
+ tf.math.reduce_sum(true_positives),
+ tf.math.reduce_sum(true_positives)
+ + tf.math.reduce_sum(false_negatives)
+ + tf.math.reduce_sum(false_positives),
+ )
+
+ return csi
+
+ return csi
+
+
+def fractions_skill_score(
+ mask_size: int | tuple[int, ...] | list[int, ...] = (3, 3),
+ alpha: int | float = 1.0,
+ beta: int | float = 0.5,
+ class_weights: list[int | float, ...] = None,
+):
+ """
+ Fractions skill score metric.
+
+ Parameters
+ ----------
+ mask_size: int or tuple
+ Size of the mask/pool in the AveragePooling layers.
+ alpha: int or float
+ Parameter that controls how steep the sigmoid function is for discretization. Higher alpha makes the sigmoid function
+ steeper.
+ beta: int or float
+ Parameter used to control some behaviors of the sigmoid discretization function. Default and recommended value is 0.5.
+ class_weights: list of values or None
+ List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
+
+ Returns
+ -------
+ fss: float
+ Fractions skill score.
+
+ References
+ ----------
+ (RL2008) Roberts, N. M., and H. W. Lean, 2008: Scale-Selective Verification of Rainfall Accumulations from High-Resolution
+ Forecasts of Convective Events. Mon. Wea. Rev., 136, 78–97, https://doi.org/10.1175/2007MWR2123.1.
+ """
+
+ # keyword arguments for the AveragePooling layer
+ pool_args = dict(pool_size=mask_size, strides=1, padding="same")
+
+ # if mask_size is an int, convert to a tuple. This allows us to check the length of the tuple and pull the correct AveragePooling layer
+ if isinstance(mask_size, int):
+ mask_size = (mask_size,)
+
+ # if mask_size is an list, convert to a tuple
+ elif isinstance(mask_size, list):
+ mask_size = tuple(mask_size)
+
+ # make sure the mask size is between 1 and 3
+ assert 1 <= len(mask_size) <= 3, (
+ "mask_size must have length between 1 and 3, received length %d"
+ % len(mask_size)
+ )
+
+ # get the pooling layer based off the length of the mask_size tuple
+ pool = getattr(tf.keras.layers, "AveragePooling%dD" % len(mask_size))(**pool_args)
+
+ if class_weights is not None:
+ class_weights = tf.cast(class_weights, tf.float32)
+
+ @tf.function
+ def fss(y_true, y_pred):
+ """
+ y_true: tf.Tensor
+ One-hot encoded tensor containing labels.
+ y_pred: tf.Tensor
+ Tensor containing model predictions.
+ """
+
+ if class_weights is not None:
+ y_true *= class_weights
+ y_pred *= class_weights
+
+ # discretize model predictions and labels
+ y_true = tf.math.sigmoid(alpha * (y_true - beta))
+ y_pred = tf.math.sigmoid(alpha * (y_pred - beta))
+
+ O_n = pool(y_true) # observed fractions (Eq. 2 in RL2008)
+ M_n = pool(y_pred) # model forecast fractions (Eq. 3 in RL2008)
+
+ MSE_n = tf.keras.metrics.mean_squared_error(
+ O_n * class_weights, M_n * class_weights
+ ) # MSE for model forecast fractions (Eq. 5 in RL2008)
+ MSE_ref = tf.reduce_mean(tf.square(O_n * class_weights)) + tf.reduce_mean(
+ tf.square(M_n * class_weights)
+ ) # reference forecast (Eq. 7 in RL2008)
+
+ FSS = 1 - MSE_n / (MSE_ref + 1e-10) # fractions skill score (Eq. 6 in RL2008)
+
+ return FSS
+
+ return fss
+
+
+def heidke_skill_score(
+ threshold: float = None,
+ window_size: tuple[int, ...] | list[int, ...] = None,
+ class_weights: list[int | float, ...] = None,
+):
+ """
+ Heidke Skill Score (HSS).
+
+ threshold: float or None
+ Optional probability threshold that binarizes y_pred. Values in y_pred greater than or equal to the threshold are
+ set to 1, and 0 otherwise.
+ If the threshold is set, it must be greater than 0 and less than 1.
+ window_size: tuple or list of ints or None
+ Pool/kernel size of the max-pooling window for neighborhood statistics. (e.g. if calculating the HSS with a 3-pixel
+ window, this should be set to 3).
+ Note that this parameter is experimental and may return unexpected results.
+ class_weights: list of values or None
+ List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
+ """
+
+ @tf.function
+ def hss(y_true, y_pred):
+ """
+ y_true: tf.Tensor
+ One-hot encoded tensor containing labels.
+ y_pred: tf.Tensor
+ Tensor containing model predictions.
+ """
+
+ if window_size is not None:
+ y_pred = tf.nn.max_pool(
+ y_pred, ksize=window_size, strides=1, padding="VALID"
+ )
+ y_true = tf.nn.max_pool(
+ y_true, ksize=window_size, strides=1, padding="VALID"
+ )
+
+ if threshold is not None:
+ y_pred = tf.where(y_pred >= threshold, 1.0, 0.0)
+
+ sum_over_axes = tf.range(
+ tf.rank(y_pred) - 1
+ ) # Indices for axes to sum over. Excludes the final (class) dimension.
+
+ true_positives = tf.math.reduce_sum(y_true * y_pred, axis=sum_over_axes)
+ false_positives = tf.math.reduce_sum((1 - y_true) * y_pred, axis=sum_over_axes)
+ false_negatives = tf.math.reduce_sum(y_true * (1 - y_pred), axis=sum_over_axes)
+ true_negatives = tf.math.reduce_sum(
+ (1 - y_true) * (1 - y_pred), axis=sum_over_axes
+ )
+
+ if class_weights is not None:
+ relative_class_weights = tf.cast(
+ class_weights / tf.math.reduce_sum(class_weights), tf.float32
+ )
+ true_positives *= relative_class_weights
+ true_negatives *= relative_class_weights
+ false_positives *= relative_class_weights
+ false_negatives *= relative_class_weights
+
+ a = tf.math.reduce_sum(true_positives)
+ b = tf.math.reduce_sum(false_positives)
+ c = tf.math.reduce_sum(false_negatives)
+ d = tf.math.reduce_sum(true_negatives)
+
+ hss = 2 * tf.math.divide(
+ (a * d) - (b * c), ((a + c) * (c + d)) + ((a + b) * (b + d))
+ )
+
+ return hss
+
+ return hss
+
+
+def probability_of_detection(
+ threshold: float = None,
+ window_size: tuple[int, ...] | list[int, ...] = None,
+ class_weights: list[int | float, ...] = None,
+):
+ """
+ Probability of Detection (POD).
+
+ threshold: float or None
+ Optional probability threshold that binarizes y_pred. Values in y_pred greater than or equal to the threshold are
+ set to 1, and 0 otherwise.
+ If the threshold is set, it must be greater than 0 and less than 1.
+ window_size: tuple or list of ints or None
+ Pool/kernel size of the max-pooling window for neighborhood statistics. (e.g. if calculating the POD with a 5-pixel
+ window, this should be set to 5).
+ Note that this parameter is experimental and may return unexpected results.
+ class_weights: list of values or None
+ List of weights to apply to each class. The length must be equal to the number of classes in y_pred and y_true.
+ """
+
+ @tf.function
+ def pod(y_true, y_pred):
+ """
+ y_true: tf.Tensor
+ One-hot encoded tensor containing labels.
+ y_pred: tf.Tensor
+ Tensor containing model predictions.
+ """
+
+ if window_size is not None:
+ y_pred = tf.nn.max_pool(
+ y_pred, ksize=window_size, strides=1, padding="VALID"
+ )
+ y_true = tf.nn.max_pool(
+ y_true, ksize=window_size, strides=1, padding="VALID"
+ )
+
+ y_pred = (
+ tf.where(y_pred >= threshold, 1.0, 0.0) if threshold is not None else y_pred
+ )
+ y_pred_neg = 1 - y_pred
+
+ sum_over_axes = tf.range(
+ tf.rank(y_pred) - 1
+ ) # Indices for axes to sum over. Excludes the final (class) dimension.
+
+ true_positives = tf.math.reduce_sum(y_pred * y_true, axis=sum_over_axes)
+ false_negatives = tf.math.reduce_sum(y_pred_neg * y_true, axis=sum_over_axes)
+
+ if class_weights is not None:
+ relative_class_weights = tf.cast(
+ class_weights / tf.math.reduce_sum(class_weights), tf.float32
+ )
+ pod = tf.math.reduce_sum(
+ tf.math.divide_no_nan(true_positives, true_positives + false_negatives)
+ * relative_class_weights
+ )
+ else:
+ pod = tf.math.reduce_sum(
+ tf.math.divide_no_nan(true_positives, true_positives + false_negatives)
+ )
+
+ return pod
+
+ return pod
diff --git a/utils/unet_utils.py b/src/fronts/layers/modules.py
similarity index 52%
rename from utils/unet_utils.py
rename to src/fronts/layers/modules.py
index 4410a15..d746e40 100644
--- a/utils/unet_utils.py
+++ b/src/fronts/layers/modules.py
@@ -1,855 +1,979 @@
-"""
-Functions for building U-Net models:
- - U-Net
- - U-Net ensemble
- - U-Net+
- - U-Net++
- - U-Net 3+
- - Attention U-Net
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.9.18
-"""
-
-import numpy as np
-from tensorflow.keras.layers import Activation, Conv2D, Conv3D, BatchNormalization, MaxPooling2D, MaxPooling3D, UpSampling2D, UpSampling3D, Softmax
-from tensorflow.keras import layers
-import tensorflow as tf
-import custom_activations
-
-
-def attention_gate(
- x: tf.Tensor,
- g: tf.Tensor,
- kernel_size: int | tuple[int],
- pool_size: tuple[int],
- name: str or None = None):
- """
- Attention gate for the Attention U-Net.
-
- Parameters
- ----------
- x: tf.Tensor
- Signal that originates from the encoder node on the same level as the attention gate.
- g: tf.Tensor
- Signal that originates from the level below the attention gate, which has higher resolution features.
- kernel_size: int or tuple
- Size of the kernel in the Conv2D/Conv3D layer(s). Only applies to layers that are not forced to a kernel size of 1.
- pool_size: tuple or list
- Pool size for the UpSampling layers, as well as the number of strides in the first
- name: str or None
- Prefix of the layer names. If left as None, the layer names are set automatically.
-
- References
- ----------
- https://towardsdatascience.com/a-detailed-explanation-of-the-attention-u-net-b371a5590831
- """
-
- conv_layer = getattr(tf.keras.layers, f'Conv{len(x.shape) - 2}D') # Select the convolution layer for the x and g tensors
- upsample_layer = getattr(tf.keras.layers, f'UpSampling{len(x.shape) - 2}D') # Select the upsampling layer
-
- shape_x = x.shape # Shapes of the ORIGINAL inputs
- filters_x = shape_x[-1]
-
- """
- x: Get the x tensor to the same shape as the gating signal (g tensor)
- g: Perform a 1x1-style convolution on the gating signal so it has the same number of filters as the x signal
- """
- x_conv = conv_layer(filters=filters_x,
- kernel_size=kernel_size,
- strides=pool_size,
- padding='same',
- name=f'{name}_Conv{len(x.shape) - 2}D_x')(x)
- g_conv = conv_layer(filters=filters_x, kernel_size=1, padding='same', name=f'{name}_Conv{len(x.shape) - 2}D_g')(g)
-
- xg = tf.add(x_conv, g_conv, name=f'{name}_sum') # Sum the x and g signals element-wise
- xg = Activation(activation='relu', name=f'{name}_relu')(xg) # Pass the summed signals through a ReLU activation layer
-
- xg_collapse = conv_layer(filters=1, kernel_size=1, padding='same', name=f'{name}_collapse')(xg) # Collapse the number of filters to just 1
- xg_collapse = Activation(activation='sigmoid', name=f'{name}_sigmoid')(xg_collapse) # Pass collapsed tensor through a sigmoid activation layer
-
- # Upsample the collapsed tensor so its dimensions match the original shape of the x signal, then expand the filters to match the g signal filters
- upsample_xg = upsample_layer(size=pool_size, name=f'{name}_UpSampling{len(x.shape) - 2}D')(xg_collapse)
- upsample_xg = tf.repeat(upsample_xg, filters_x, axis=-1, name=f'{name}_repeat')
-
- coeffs = tf.multiply(upsample_xg, x, name=f'{name}_multiply') # Element-wise multiplication onto the original x signal
-
- attention_tensor = Conv3D(filters=filters_x, kernel_size=1, strides=1, padding='same', name=f'{name}_Conv{len(x.shape) - 2}D_coeffs')(coeffs)
- attention_tensor = BatchNormalization(name=f'{name}_BatchNorm')(attention_tensor)
-
- return attention_tensor
-
-
-def convolution_module(
- tensor: tf.Tensor,
- filters: int,
- kernel_size: int | tuple[int],
- num_modules: int = 1,
- padding: str = 'same',
- use_bias: bool = False,
- batch_normalization: bool = True,
- activation: str = 'relu',
- kernel_initializer = 'glorot_uniform',
- bias_initializer = 'zeros',
- kernel_regularizer = None,
- bias_regularizer = None,
- activity_regularizer = None,
- kernel_constraint = None,
- bias_constraint = None,
- name: str = None):
- """
- Insert modules into an encoder or decoder node.
-
- Parameters
- ----------
- tensor: tf.Tensor
- Input tensor for the convolution module(s).
- filters: int
- Number of filters in the Conv2D/Conv3D layer(s).
- kernel_size: int or tuple of ints
- Size of the kernel in the Conv2D/Conv3D layer(s).
- num_modules: int
- Number of convolution modules to insert. Must be greater than 0, otherwise a ValueError exception is raised.
- padding: str
- Padding in the Conv2D/Conv3D layer(s). 'valid' will apply no padding, while 'same' will apply padding such that the
- output shape matches the input shape. 'valid' and 'same' are case-insensitive.
- use_bias: bool
- If True, a bias vector will be used in the Conv2D/Conv3D layers.
- batch_normalization: bool
- If True, a BatchNormalization layer will follow every Conv2D/Conv3D layer.
- activation: str
- Activation function to use after every Conv2D/Conv3D layer (BatchNormalization layer, if batch_normalization is True).
- Can be any of tf.keras.activations, 'prelu', 'leaky_relu', or 'smelu' (case-insensitive).
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
- name: str or None
- Prefix of the layer names. If left as None, the layer names are set automatically.
-
- Returns
- -------
- tensor: tf.Tensor
- Output tensor.
- """
-
- tensor_dims = len(tensor.shape) # Number of dims in the tensor (including the first 'None' dimension for batch size)
-
- if num_modules < 1:
- raise ValueError("num_modules must be greater than 0, at least one module must be added")
-
- if tensor_dims == 4: # A 2D image tensor has 4 dimensions: (None [for batch size], image_size_x, image_size_y, n_channels)
- conv_layer = Conv2D
- elif tensor_dims == 5: # A 3D image tensor has 5 dimensions: (None [for batch size], image_size_x, image_size_y, image_size_z, n_channels)
- conv_layer = Conv3D
- else:
- raise TypeError(f"Incompatible tensor shape: {tensor.shape}. The tensor must only have 4 or 5 dimensions")
-
- # Arguments for the Conv2D/Conv3D layers.
- conv_kwargs = dict({})
- conv_kwargs['filters'] = filters
- conv_kwargs['kernel_size'] = kernel_size
- conv_kwargs['padding'] = padding
- conv_kwargs['use_bias'] = use_bias
- conv_kwargs['kernel_initializer'] = kernel_initializer
- conv_kwargs['bias_initializer'] = bias_initializer
- conv_kwargs['kernel_regularizer'] = kernel_regularizer
- conv_kwargs['bias_regularizer'] = bias_regularizer
- conv_kwargs['activity_regularizer'] = activity_regularizer
- conv_kwargs['kernel_constraint'] = kernel_constraint
- conv_kwargs['bias_constraint'] = bias_constraint
-
- activation_layer = choose_activation_layer(activation) # Choose activation layer for the convolution modules.
-
- """
- Arguments for the Activation layer(s), if applicable.
- - If activation is 'prelu' or 'leaky_relu', a PReLU layer or a LeakyReLU layer will be used instead and this dictionary
- will have no effect.
- """
- activation_kwargs = dict({})
- if activation_layer == Activation:
- activation_kwargs['activation'] = activation
-
- # Insert convolution modules
- for module in range(num_modules):
-
- # Create names for the Conv2D/Conv3D layers and the activation layer.
- conv_kwargs['name'] = f'{name}_Conv{tensor_dims - 2}D_{module+1}'
- activation_kwargs['name'] = f'{name}_{activation}_{module+1}'
-
- conv_tensor = conv_layer(**conv_kwargs)(tensor) # Perform convolution on the input tensor
-
- if batch_normalization:
- batch_norm_tensor = BatchNormalization(name=f'{name}_BatchNorm_{module+1}')(conv_tensor) # Insert layer for batch normalization
- activation_tensor = activation_layer(**activation_kwargs)(batch_norm_tensor) # Pass output tensor from BatchNormalization into the activation layer
- else:
- activation_tensor = activation_layer(**activation_kwargs)(conv_tensor) # Pass output tensor from the convolution layer into the activation layer.
-
- tensor = activation_tensor
-
- return tensor
-
-
-def aggregated_feature_map(
- tensor: tf.Tensor,
- filters: int,
- kernel_size: int | tuple[int],
- level1: int,
- level2: int,
- upsample_size: tuple[int],
- padding: str = 'same',
- use_bias: bool = False,
- batch_normalization: bool = True,
- activation: str = 'relu',
- kernel_initializer = 'glorot_uniform',
- bias_initializer = 'zeros',
- kernel_regularizer = None,
- bias_regularizer = None,
- activity_regularizer = None,
- kernel_constraint = None,
- bias_constraint = None,
- name: str = None):
- """
- Connect two nodes in the U-Net 3+ with an aggregated feature map (AFM).
-
- Parameters
- ----------
- tensor: tf.Tensor
- Input tensor for the convolution module.
- filters: int
- Number of filters in the Conv2D/Conv3D layer.
- kernel_size: int or tuple
- Size of the kernel in the Conv2D/Conv3D layer.
- level1: int
- Level of the first node that is connected to the AFM. This node will provide the input tensor to the AFM. Must be
- greater than level2 (i.e. the first node must be on a lower level in the U-Net 3+ since we are up-sampling), otherwise
- a ValueError exception is raised.
- level2: int
- Level of the second node that is connected to the AFM. This node will receive the output of the AFM. Must be smaller
- than level1 (i.e. the second node must be on a higher level in the U-Net 3+ since we are up-sampling), otherwise
- a ValueError exception is raised.
- upsample_size: tuple or list of ints
- Upsampling size for rows and columns in the UpSampling2D/UpSampling3D layer.
- padding: str
- Padding in the Conv2D/Conv3D layer. 'valid' will apply no padding, while 'same' will apply padding such that the
- output shape matches the input shape. 'valid' and 'same' are case-insensitive.
- use_bias: bool
- If True, a bias vector will be used in the Conv2D/Conv3D layer.
- batch_normalization: bool
- If True, a BatchNormalization layer will follow the Conv2D/Conv3D layer.
- activation: str
- Activation function to use after the Conv2D/Conv3D layer (BatchNormalization layer, if batch_normalization is True).
- Can be any of tf.keras.activations, 'prelu', 'leaky_relu', or 'smelu' (case-insensitive).
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
- name: str or None
- Prefix of the layer names. If left as None, the layer names are set automatically.
-
- Returns
- -------
- tensor: tf.Tensor
- Output tensor.
- """
-
- tensor_dims = len(tensor.shape) # Number of dims in the tensor (including the first 'None' dimension for batch size)
-
- if level1 <= level2:
- raise ValueError("level2 must be smaller than level1 in aggregated feature maps")
-
- # Arguments for the convolution module.
- module_kwargs = dict({})
- module_kwargs['filters'] = filters
- module_kwargs['kernel_size'] = kernel_size
- module_kwargs['num_modules'] = 1
- module_kwargs['padding'] = padding
- module_kwargs['use_bias'] = use_bias
- module_kwargs['batch_normalization'] = batch_normalization
- module_kwargs['activation'] = activation
- module_kwargs['kernel_initializer'] = kernel_initializer
- module_kwargs['bias_initializer'] = bias_initializer
- module_kwargs['kernel_regularizer'] = kernel_regularizer
- module_kwargs['bias_regularizer'] = bias_regularizer
- module_kwargs['activity_regularizer'] = activity_regularizer
- module_kwargs['kernel_constraint'] = kernel_constraint
- module_kwargs['bias_constraint'] = bias_constraint
- module_kwargs['name'] = name
-
- # Keyword arguments for the UpSampling2D/UpSampling3D layers
- upsample_kwargs = dict({})
- upsample_kwargs['name'] = f'{name}_UpSampling{tensor_dims - 2}D'
- upsample_kwargs['size'] = np.power(upsample_size, abs(level1 - level2))
-
- if tensor_dims == 4: # If the image is 2D
- upsample_layer = UpSampling2D
- if len(upsample_size) != 2:
- raise TypeError(f"For 2D up-sampling, the pool size must be a tuple or list with 2 integers. Received shape: {np.shape(upsample_size)}")
- elif tensor_dims == 5: # If the image is 3D
- upsample_layer = UpSampling3D
- if len(upsample_size) != 3:
- raise TypeError(f"For 3D up-sampling, the pool size must be a tuple or list with 3 integers. Received shape: {np.shape(upsample_size)}")
- else:
- raise TypeError(f"Incompatible tensor shape: {tensor.shape}. The tensor must only have 4 or 5 dimensions")
-
- tensor = upsample_layer(**upsample_kwargs)(tensor) # Pass the tensor through the UpSample2D/UpSample3D layer
-
- tensor = convolution_module(tensor, **module_kwargs) # Pass input tensor through convolution module
-
- return tensor
-
-
-def full_scale_skip_connection(
- tensor: tf.Tensor,
- filters: int,
- kernel_size: int | tuple[int],
- level1: int,
- level2: int,
- pool_size: tuple[int],
- padding: str = 'same',
- use_bias: bool = False,
- batch_normalization: bool = True,
- activation: str = 'relu',
- kernel_initializer = 'glorot_uniform',
- bias_initializer = 'zeros',
- kernel_regularizer = None,
- bias_regularizer = None,
- activity_regularizer = None,
- kernel_constraint = None,
- bias_constraint = None,
- name: str = None):
- """
- Connect two nodes in the U-Net 3+ with a full-scale skip connection (FSC).
-
- Parameters
- ----------
- tensor: tf.Tensor
- Input tensor for the convolution module.
- filters: int
- Number of filters in the Conv2D/Conv3D layer.
- kernel_size: int or tuple
- Size of the kernel in the Conv2D/Conv3D layer.
- level1: int
- Level of the first node that is connected to the FSC. This node will provide the input tensor to the FSC. Must be
- smaller than level2 (i.e. the first node must be on a higher level in the U-Net 3+ since we are max-pooling), otherwise
- a ValueError exception is raised.
- level2: int
- Level of the second node that is connected to the FSC. This node will receive the output of the FSC. Must be greater
- than level1 (i.e. the second node must be on a lower level in the U-Net 3+ since we are max-pooling), otherwise
- a ValueError exception is raised.
- pool_size: tuple or list
- Pool size for the MaxPooling2D/MaxPooling3D layer.
- padding: str
- Padding in the Conv2D/Conv3D layer. 'valid' will apply no padding, while 'same' will apply padding such that the
- output shape matches the input shape. 'valid' and 'same' are case-insensitive.
- use_bias: bool
- If True, a bias vector will be used in the Conv2D/Conv3D layer.
- batch_normalization: bool
- If True, a BatchNormalization layer will follow the Conv2D/Conv3D layer.
- activation: str
- Activation function to use after the Conv2D/Conv3D layer (BatchNormalization layer, if batch_normalization is True).
- Can be any of tf.keras.activations, 'prelu', 'leaky_relu', or 'smelu' (case-insensitive).
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
- name: str or None
- Prefix of the layer names. If left as None, the layer names are set automatically.
-
- Returns
- -------
- tensor: tf.Tensor
- Output tensor.
- """
-
- tensor_dims = len(tensor.shape) # Number of dims in the tensor (including the first 'None' dimension for batch size)
-
- if level1 >= level2:
- raise ValueError("level2 must be greater than level1 in full-scale skip connections")
-
- # Arguments for the convolution module.
- module_kwargs = dict({})
- module_kwargs['filters'] = filters
- module_kwargs['kernel_size'] = kernel_size
- module_kwargs['num_modules'] = 1
- module_kwargs['padding'] = padding
- module_kwargs['use_bias'] = use_bias
- module_kwargs['batch_normalization'] = batch_normalization
- module_kwargs['activation'] = activation
- module_kwargs['kernel_initializer'] = kernel_initializer
- module_kwargs['bias_initializer'] = bias_initializer
- module_kwargs['kernel_regularizer'] = kernel_regularizer
- module_kwargs['bias_regularizer'] = bias_regularizer
- module_kwargs['activity_regularizer'] = activity_regularizer
- module_kwargs['kernel_constraint'] = kernel_constraint
- module_kwargs['bias_constraint'] = bias_constraint
- module_kwargs['name'] = name
-
- # Keyword arguments for the MaxPooling2D/MaxPooling3D layer
- pool_kwargs = dict({})
- pool_kwargs['name'] = f'{name}_MaxPool{tensor_dims - 2}D'
- pool_kwargs['pool_size'] = np.power(pool_size, abs(level1 - level2))
-
- if tensor_dims == 4: # If the image is 2D
- pool_layer = MaxPooling2D
- elif tensor_dims == 5: # If the image is 3D
- pool_layer = MaxPooling3D
- else:
- raise TypeError(f"Incompatible tensor shape: {tensor.shape}. The tensor must only have 4 or 5 dimensions")
-
- tensor = pool_layer(**pool_kwargs)(tensor) # Pass the tensor through the MaxPooling2D/MaxPooling3D layer
-
- tensor = convolution_module(tensor, **module_kwargs) # Pass the tensor through the convolution module
-
- return tensor
-
-
-def conventional_skip_connection(
- tensor: tf.Tensor,
- filters: int,
- kernel_size: int | tuple[int],
- padding: str = 'same',
- use_bias: bool = False,
- batch_normalization: bool = True,
- activation: str = 'relu',
- kernel_initializer = 'glorot_uniform',
- bias_initializer = 'zeros',
- kernel_regularizer = None,
- bias_regularizer = None,
- activity_regularizer = None,
- kernel_constraint = None,
- bias_constraint = None,
- name: str = None):
- """
- Connect two nodes in the U-Net 3+ with a conventional skip connection.
-
- Parameters
- ----------
- tensor: tf.Tensor
- Input tensor for the convolution module.
- filters: int
- Number of filters in the Conv2D/Conv3D layer.
- kernel_size: int or tuple
- Size of the kernel in the Conv2D/Conv3D layer.
- padding: str
- Padding in the Conv2D/Conv3D layer. 'valid' will apply no padding, while 'same' will apply padding such that the
- output shape matches the input shape. 'valid' and 'same' are case-insensitive.
- use_bias: bool
- If True, a bias vector will be used in the Conv2D/Conv3D layer.
- batch_normalization: bool
- If True, a BatchNormalization layer will follow the Conv2D/Conv3D layer.
- activation: str
- Activation function to use after the Conv2D/Conv3D layer (BatchNormalization layer, if batch_normalization is True).
- Can be any of tf.keras.activations, 'prelu', 'leaky_relu', or 'smelu' (case-insensitive).
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
- name: str or None
- Prefix of the layer names. If left as None, the layer names are set automatically.
-
- Returns
- -------
- tensor: tf.Tensor
- Output tensor.
- """
-
- # Arguments for the convolution module.
- module_kwargs = dict({})
- module_kwargs['filters'] = filters
- module_kwargs['kernel_size'] = kernel_size
- module_kwargs['num_modules'] = 1
- module_kwargs['padding'] = padding
- module_kwargs['use_bias'] = use_bias
- module_kwargs['batch_normalization'] = batch_normalization
- module_kwargs['activation'] = activation
- module_kwargs['kernel_initializer'] = kernel_initializer
- module_kwargs['bias_initializer'] = bias_initializer
- module_kwargs['kernel_regularizer'] = kernel_regularizer
- module_kwargs['bias_regularizer'] = bias_regularizer
- module_kwargs['activity_regularizer'] = activity_regularizer
- module_kwargs['kernel_constraint'] = kernel_constraint
- module_kwargs['bias_constraint'] = bias_constraint
- module_kwargs['name'] = name
-
- tensor = convolution_module(tensor, **module_kwargs) # Pass the tensor through the convolution module
-
- return tensor
-
-
-def max_pool(
- tensor: tf.Tensor,
- pool_size: tuple[int],
- name: str = None):
- """
- Connect two encoder nodes with a max-pooling operation.
-
- Parameters
- ----------
- tensor: tf.Tensor
- Input tensor for the convolution module.
- pool_size: tuple or list
- Pool size for the MaxPooling2D/MaxPooling3D layer.
- name: str or None
- Prefix of the layer names. If left as None, the layer names are set automatically.
-
- Returns
- -------
- tensor: tf.Tensor
- Output tensor.
- """
-
- if type(pool_size) != tuple and type(pool_size) != list:
- raise TypeError(f"pool_size can only be a tuple or list. Received type: {type(pool_size)}")
-
- tensor_dims = len(tensor.shape) # Number of dims in the tensor (including the first 'None' dimension for batch size)
-
- pool_kwargs = dict({}) # Keyword arguments in the MaxPooling layer
- pool_kwargs['name'] = f'{name}_MaxPool{tensor_dims - 2}D'
- pool_kwargs['pool_size'] = pool_size
-
- if tensor_dims == 4: # If the image is 2D
- pool_layer = MaxPooling2D
- if len(pool_size) != 2:
- raise TypeError(f"For 2D max-pooling, the pool size must be a tuple or list with 2 integers. Received shape: {np.shape(pool_size)}")
- elif tensor_dims == 5: # If the image is 3D
- pool_layer = MaxPooling3D
- if len(pool_size) != 3:
- raise TypeError(f"For 3D max-pooling, the pool size must be a tuple or list with 3 integers. Received shape: {np.shape(pool_size)}")
- else:
- raise TypeError(f"Incompatible tensor shape: {tensor.shape}. The tensor must only have 4 or 5 dimensions")
-
- pool_tensor = pool_layer(**pool_kwargs)(tensor) # Pass the tensor through the MaxPooling2D/MaxPooling3D layer
-
- return pool_tensor
-
-
-def upsample(
- tensor: tf.Tensor,
- filters: int,
- kernel_size: int | tuple[int],
- upsample_size: tuple[int],
- padding: str = 'same',
- use_bias: bool = False,
- batch_normalization: bool = True,
- activation: str = 'relu',
- kernel_initializer='glorot_uniform',
- bias_initializer='zeros',
- kernel_regularizer = None,
- bias_regularizer = None,
- activity_regularizer = None,
- kernel_constraint = None,
- bias_constraint = None,
- name: str = None):
- """
- Connect decoder nodes with an up-sampling operation.
-
- Parameters
- ----------
- tensor: tf.Tensor
- Input tensor for the convolution module.
- filters: int
- Number of filters in the Conv2D/Conv3D layer.
- kernel_size: int or tuple
- Size of the kernel in the Conv2D/Conv3D layer.
- upsample_size: tuple or list
- Upsampling size in the UpSampling2D/UpSampling3D layer.
- padding: str
- Padding in the Conv2D/Conv3D layer. 'valid' will apply no padding, while 'same' will apply padding such that the
- output shape matches the input shape. 'valid' and 'same' are case-insensitive.
- use_bias: bool
- If True, a bias vector will be used in the Conv2D/Conv3D layer.
- batch_normalization: bool
- If True, a BatchNormalization layer will follow the Conv2D/Conv3D layer.
- activation: str
- Activation function to use after the Conv2D/Conv3D layer (BatchNormalization layer, if batch_normalization is True).
- Can be any of tf.keras.activations, 'prelu', 'leaky_relu', or 'smelu' (case-insensitive).
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
- name: str or None
- Prefix of the layer names. If left as None, the layer names are set automatically.
-
- Returns
- -------
- tensor: tf.Tensor
- Output tensor.
- """
-
- tensor_dims = len(tensor.shape) # Number of dims in the tensor (including the first 'None' dimension for batch size)
-
- if type(upsample_size) != tuple and type(upsample_size) != list:
- raise TypeError(f"upsample_size can only be a tuple or list. Received type: {type(upsample_size)}")
-
- # Arguments for the convolution module.
- module_kwargs = dict({})
- module_kwargs['filters'] = filters
- module_kwargs['kernel_size'] = kernel_size
- module_kwargs['num_modules'] = 1
- module_kwargs['padding'] = padding
- module_kwargs['use_bias'] = use_bias
- module_kwargs['batch_normalization'] = batch_normalization
- module_kwargs['activation'] = activation
- module_kwargs['kernel_initializer'] = kernel_initializer
- module_kwargs['bias_initializer'] = bias_initializer
- module_kwargs['kernel_regularizer'] = kernel_regularizer
- module_kwargs['bias_regularizer'] = bias_regularizer
- module_kwargs['activity_regularizer'] = activity_regularizer
- module_kwargs['kernel_constraint'] = kernel_constraint
- module_kwargs['bias_constraint'] = bias_constraint
- module_kwargs['name'] = name
-
- # Keyword arguments in the UpSampling layer
- upsample_kwargs = dict({})
- upsample_kwargs['name'] = f'{name}_UpSampling{tensor_dims - 2}D'
- upsample_kwargs['size'] = upsample_size
-
- if tensor_dims == 4: # If the image is 2D
- upsample_layer = UpSampling2D
- if len(upsample_size) != 2:
- raise TypeError(f"For 2D up-sampling, the pool size must be a tuple or list with 2 integers. Received shape: {np.shape(upsample_size)}")
- elif tensor_dims == 5: # If the image is 3D
- upsample_layer = UpSampling3D
- if len(upsample_size) != 3:
- raise TypeError(f"For 3D up-sampling, the pool size must be a tuple or list with 3 integers. Received shape: {np.shape(upsample_size)}")
- else:
- raise TypeError(f"Incompatible tensor shape: {tensor.shape}. The tensor must only have 4 or 5 dimensions")
-
- upsample_tensor = upsample_layer(**upsample_kwargs)(tensor) # Pass the tensor through the UpSampling2D/UpSampling3D layer
-
- tensor = convolution_module(upsample_tensor, **module_kwargs) # Pass the up-sampled tensor through a convolution module
-
- return tensor
-
-
-def choose_activation_layer(activation: str):
- """
- Choose activation layer for the U-Net.
-
- Parameters
- ----------
- activation: str
- Can be any of tf.keras.activations, 'gaussian', 'gcu', 'leaky_relu', 'prelu', 'smelu', 'snake' (case-insensitive).
-
- Returns
- -------
- activation_layer: tf.keras.layers.Activation, tf.keras.layers.PReLU, tf.keras.layers.LeakyReLU, or any layer from custom_activations
- Activation layer.
- """
-
- activation = activation.lower()
-
- available_activations = ['elu', 'exponential', 'gaussian', 'gcu', 'gelu', 'hard_sigmoid', 'leaky_relu', 'linear', 'prelu',
- 'relu', 'selu', 'sigmoid', 'smelu', 'snake', 'softmax', 'softplus', 'softsign', 'swish', 'tanh']
-
- # Choose the activation layer
- if activation == 'leaky_relu':
- activation_layer = getattr(layers, 'LeakyReLU')
- elif activation == 'prelu':
- activation_layer = getattr(layers, 'PReLU')
- elif activation == 'smelu':
- activation_layer = custom_activations.SmeLU
- elif activation == 'gcu':
- activation_layer = custom_activations.GCU
- elif activation == 'gaussian':
- activation_layer = custom_activations.Gaussian
- elif activation == 'snake':
- activation_layer = custom_activations.Snake
- elif activation in available_activations:
- activation_layer = getattr(layers, 'Activation')
- else:
- raise TypeError(f"'{activation}' is not a valid loss function and/or is not available, options are: {', '.join(sorted(list(available_activations)))}")
-
- return activation_layer
-
-
-def deep_supervision_side_output(
- tensor: tf.Tensor,
- num_classes: int,
- kernel_size: int | tuple[int],
- output_level: int,
- upsample_size: tuple[int],
- padding: str = 'same',
- use_bias: bool = False,
- kernel_initializer='glorot_uniform',
- bias_initializer='zeros',
- kernel_regularizer = None,
- bias_regularizer = None,
- activity_regularizer = None,
- kernel_constraint = None,
- bias_constraint = None,
- squeeze_dims: int | tuple[int] = None,
- name: str = None):
- """
- Deep supervision output. This is usually used on a decoder node in the U-Net 3+ or the final decoder node of a standard
- U-Net.
-
- Parameters
- ----------
- tensor: tf.Tensor
- Input tensor for the convolution module.
- num_classes: int
- Number of classes that the model is trying to predict.
- kernel_size: int or tuple
- Size of the kernel in the Conv2D/Conv3D layer.
- output_level: int
- Level of the decoder node from which the deep supervision output is based.
- upsample_size: tuple or list
- Upsampling size for rows and columns in the UpSampling2D/UpSampling3D layer. Tuples are currently not supported
- but will be supported in a future update.
- padding: str
- Padding in the Conv2D/Conv3D layer. 'valid' will apply no padding, while 'same' will apply padding such that the
- output shape matches the input shape. 'valid' and 'same' are case-insensitive.
- use_bias: bool
- If True, a bias vector will be used in the Conv2D/Conv3D layer.
- kernel_initializer: str or tf.keras.initializers object
- Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_initializer: str or tf.keras.initializers object
- Initializer for the bias vector in the Conv2D/Conv3D layers.
- kernel_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
- bias_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
- activity_regularizer: str or tf.keras.regularizers object
- Regularizer function applied to the output of the Conv2D/Conv3D layers.
- kernel_constraint: str or tf.keras.constraints object
- Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
- bias_constraint: str or tf.keras.constrains object
- Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
- squeeze_dims: int, tuple, or None
- Dimension(s) of the input tensor to squeeze.
- name: str or None
- Prefix of the layer names. If left as None, the layer names are set automatically.
-
- Returns
- -------
- tensor: tf.Tensor
- Output tensor.
- """
-
- tensor_dims = len(tensor.shape) # Number of dims in the tensor (including the first 'None' dimension for batch size)
-
- if tensor_dims == 4: # If the image is 2D
- conv_layer = Conv2D
- upsample_layer = UpSampling2D
- if output_level > 1:
- upsample_size_1 = upsample_size
- else:
- upsample_size_1 = None
-
- if output_level > 2:
- upsample_size_2 = np.power(upsample_size, output_level - 2)
- else:
- upsample_size_2 = None
-
- elif tensor_dims == 5: # If the image is 3D
- conv_layer = Conv3D
- upsample_layer = UpSampling3D
- if output_level > 1:
- upsample_size_1 = upsample_size
- else:
- upsample_size_1 = None
-
- if output_level > 2:
- upsample_size_2 = np.power(upsample_size, output_level - 2)
- else:
- upsample_size_2 = None
-
- else:
- raise TypeError(f"Incompatible tensor shape: {tensor.shape}. The tensor can only have 4 or 5 dimensions")
-
- # Arguments for the Conv2D/Conv3D layer.
- conv_kwargs = dict({})
- conv_kwargs['use_bias'] = use_bias
- conv_kwargs['kernel_size'] = kernel_size
- conv_kwargs['padding'] = padding
- conv_kwargs['kernel_initializer'] = kernel_initializer
- conv_kwargs['bias_initializer'] = bias_initializer
- conv_kwargs['kernel_regularizer'] = kernel_regularizer
- conv_kwargs['bias_regularizer'] = bias_regularizer
- conv_kwargs['activity_regularizer'] = activity_regularizer
- conv_kwargs['kernel_constraint'] = kernel_constraint
- conv_kwargs['bias_constraint'] = bias_constraint
- conv_kwargs['name'] = f'{name}_Conv{tensor_dims - 2}D'
-
- if upsample_size_1 is not None:
- tensor = upsample_layer(size=upsample_size_1, name=f'{name}_UpSampling{tensor_dims - 2}D_1')(tensor) # Pass the tensor through the UpSampling2D/UpSampling3D layer
-
- tensor = conv_layer(filters=num_classes, **conv_kwargs)(tensor) # This convolution layer contains num_classes filters, one for each class
-
- if upsample_size_2 is not None:
- tensor = upsample_layer(size=upsample_size_2, name=f'{name}_UpSampling{tensor_dims - 2}D_2')(tensor) # Pass the tensor through the UpSampling2D/UpSampling3D layer
-
- ### Squeeze the given dimensions/axes ###
- if squeeze_dims is not None:
-
- conv_kwargs['kernel_size'] = [1 for _ in range(tensor_dims - 2)]
-
- if type(squeeze_dims) == int:
- squeeze_dims = [squeeze_dims, ] # Turn integer into a list of length 1 to make indexing easier
-
- squeeze_dims_sizes = [tensor.shape[dim_to_squeeze + 1] for dim_to_squeeze in squeeze_dims] # Since the tensor contains the 'None' dimension, we have to add 1 to get the correct dimension
-
- for dim, size in enumerate(squeeze_dims_sizes):
- conv_kwargs['kernel_size'][squeeze_dims[dim]] = size # Kernel size of dimension to squeeze is equal to the size of the dimension because we want the final size to be 1 so it can be squeezed
-
- conv_kwargs['padding'] = 'valid' # Padding is no longer 'same' since we want to modify the size of the dimension to be squeezed
- conv_kwargs['name'] = f'{name}_Conv{tensor_dims - 2}D_collapse'
-
- tensor = conv_layer(filters=num_classes, **conv_kwargs)(tensor) # This convolution layer contains num_classes filters, one for each class
- tensor = tf.squeeze(tensor, axis=[axis + 1 for axis in squeeze_dims]) # Squeeze the tensor and remove the dimension
-
- sup_output = Softmax(name=f'{name}_Softmax')(tensor) # Final softmax output
-
- return sup_output
+from tensorflow.keras.layers import (
+ Activation,
+ Conv2D,
+ Conv3D,
+ BatchNormalization,
+ MaxPooling2D,
+ MaxPooling3D,
+ UpSampling2D,
+ UpSampling3D,
+)
+from fronts.utils import keras_builders
+import tensorflow as tf
+import numpy as np
+
+
+def attention_gate(
+ x: tf.Tensor,
+ g: tf.Tensor,
+ kernel_size: int | tuple[int],
+ pool_size: tuple[int],
+ name: str or None = None,
+):
+ """
+ Attention gate for the Attention U-Net.
+
+ Parameters
+ ----------
+ x: tf.Tensor
+ Signal that originates from the encoder node on the same level as the attention gate.
+ g: tf.Tensor
+ Signal that originates from the level below the attention gate, which has higher resolution features.
+ kernel_size: int or tuple
+ Size of the kernel in the Conv2D/Conv3D layer(s). Only applies to layers that are not forced to a kernel size of 1.
+ pool_size: tuple or list
+ Pool size for the UpSampling layers, as well as the number of strides in the first
+ name: str or None
+ Prefix of the layer names. If left as None, the layer names are set automatically.
+
+ References
+ ----------
+ https://towardsdatascience.com/a-detailed-explanation-of-the-attention-u-net-b371a5590831
+ """
+
+ conv_layer = getattr(
+ tf.keras.layers, f"Conv{len(x.shape) - 2}D"
+ ) # Select the convolution layer for the x and g tensors
+ upsample_layer = getattr(
+ tf.keras.layers, f"UpSampling{len(x.shape) - 2}D"
+ ) # Select the upsampling layer
+
+ shape_x = x.shape # Shapes of the ORIGINAL inputs
+ filters_x = shape_x[-1]
+
+ """
+ x: Get the x tensor to the same shape as the gating signal (g tensor)
+ g: Perform a 1x1-style convolution on the gating signal so it has the same number of filters as the x signal
+ """
+ x_conv = conv_layer(
+ filters=filters_x,
+ kernel_size=kernel_size,
+ strides=pool_size,
+ padding="same",
+ name=f"{name}_Conv{len(x.shape) - 2}D_x",
+ )(x)
+ g_conv = conv_layer(
+ filters=filters_x,
+ kernel_size=1,
+ padding="same",
+ name=f"{name}_Conv{len(x.shape) - 2}D_g",
+ )(g)
+
+ xg = tf.add(
+ x_conv, g_conv, name=f"{name}_sum"
+ ) # Sum the x and g signals element-wise
+ xg = Activation(activation="relu", name=f"{name}_relu")(
+ xg
+ ) # Pass the summed signals through a ReLU activation layer
+
+ xg_collapse = conv_layer(
+ filters=1, kernel_size=1, padding="same", name=f"{name}_collapse"
+ )(xg) # Collapse the number of filters to just 1
+ xg_collapse = Activation(activation="sigmoid", name=f"{name}_sigmoid")(
+ xg_collapse
+ ) # Pass collapsed tensor through a sigmoid activation layer
+
+ # Upsample the collapsed tensor so its dimensions match the original shape of the x signal, then expand the filters to match the g signal filters
+ upsample_xg = upsample_layer(
+ size=pool_size, name=f"{name}_UpSampling{len(x.shape) - 2}D"
+ )(xg_collapse)
+ upsample_xg = tf.repeat(upsample_xg, filters_x, axis=-1, name=f"{name}_repeat")
+
+ coeffs = tf.multiply(
+ upsample_xg, x, name=f"{name}_multiply"
+ ) # Element-wise multiplication onto the original x signal
+
+ attention_tensor = conv_layer(
+ filters=filters_x,
+ kernel_size=1,
+ strides=1,
+ padding="same",
+ name=f"{name}_Conv{len(x.shape) - 2}D_coeffs",
+ )(coeffs)
+ attention_tensor = BatchNormalization(name=f"{name}_BatchNorm")(attention_tensor)
+
+ return attention_tensor
+
+
+def convolution_module(
+ tensor: tf.Tensor,
+ filters: int,
+ kernel_size: int | tuple[int],
+ num_modules: int = 1,
+ padding: str = "same",
+ use_bias: bool = False,
+ batch_normalization: bool = True,
+ activation: str = "relu",
+ kernel_initializer="glorot_uniform",
+ bias_initializer="zeros",
+ kernel_regularizer=None,
+ bias_regularizer=None,
+ activity_regularizer=None,
+ kernel_constraint=None,
+ bias_constraint=None,
+ shared_axes=None,
+ name: str = None,
+):
+ """
+ Insert modules into an encoder or decoder node.
+
+ Parameters
+ ----------
+ tensor: tf.Tensor
+ Input tensor for the convolution module(s).
+ filters: int
+ Number of filters in the Conv2D/Conv3D layer(s).
+ kernel_size: int or tuple of ints
+ Size of the kernel in the Conv2D/Conv3D layer(s).
+ num_modules: int
+ Number of convolution modules to insert. Must be greater than 0, otherwise a ValueError exception is raised.
+ padding: str
+ Padding in the Conv2D/Conv3D layer(s). 'valid' will apply no padding, while 'same' will apply padding such that the
+ output shape matches the input shape. 'valid' and 'same' are case-insensitive.
+ use_bias: bool
+ If True, a bias vector will be used in the Conv2D/Conv3D layers.
+ batch_normalization: bool
+ If True, a BatchNormalization layer will follow every Conv2D/Conv3D layer.
+ activation: str
+ Activation function to use after every Conv2D/Conv3D layer (BatchNormalization layer, if batch_normalization is True).
+ See choose_activation_layer for all available activation functions.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+ shared_axes: tuple or list of ints
+ Axes along which to share the learnable parameters for the activation function.
+ name: str or None
+ Prefix of the layer names. If left as None, the layer names are set automatically.
+
+ Returns
+ -------
+ tensor: tf.Tensor
+ Output tensor.
+ """
+
+ tensor_dims = len(
+ tensor.shape
+ ) # Number of dims in the tensor (including the first 'None' dimension for batch size)
+
+ if num_modules < 1:
+ raise ValueError(
+ "num_modules must be greater than 0, at least one module must be added"
+ )
+
+ if (
+ tensor_dims == 4
+ ): # A 2D image tensor has 4 dimensions: (None [for batch size], image_size_x, image_size_y, n_channels)
+ conv_layer = Conv2D
+ elif (
+ tensor_dims == 5
+ ): # A 3D image tensor has 5 dimensions: (None [for batch size], image_size_x, image_size_y, image_size_z, n_channels)
+ conv_layer = Conv3D
+ else:
+ raise TypeError(
+ f"Incompatible tensor shape: {tensor.shape}. The tensor must only have 4 or 5 dimensions"
+ )
+
+ # Arguments for the Conv2D/Conv3D layer.
+ conv_kwargs = dict({})
+ for arg in [
+ "filters",
+ "use_bias",
+ "kernel_size",
+ "padding",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ ]:
+ conv_kwargs[arg] = locals()[arg]
+
+ activation_kwargs = {}
+ if activation in [
+ "prelu",
+ "smelu",
+ "snake",
+ ]: # these activation functions have learnable parameters
+ activation_kwargs["shared_axes"] = shared_axes
+
+ # Insert convolution modules
+ for module in range(num_modules):
+ # Create names for the Conv2D/Conv3D layers and the activation layer.
+ conv_kwargs["name"] = f"{name}_Conv{tensor_dims - 2}D_{module + 1}"
+ activation_kwargs["name"] = f"{name}_{activation}_{module + 1}"
+
+ conv_tensor = conv_layer(**conv_kwargs)(
+ tensor
+ ) # Perform convolution on the input tensor
+
+ activation_config = keras_builders.ActivationConfig(
+ name=activation, config=activation_kwargs
+ )
+ activation_layer = activation_config.build()
+
+ if batch_normalization:
+ batch_norm_tensor = BatchNormalization(
+ name=f"{name}_BatchNorm_{module + 1}"
+ )(conv_tensor) # Insert layer for batch normalization
+ activation_tensor = activation_layer(
+ batch_norm_tensor
+ ) # Pass output tensor from BatchNormalization into the activation layer
+ else:
+ activation_tensor = activation_layer(
+ conv_tensor
+ ) # Pass output tensor from the convolution layer into the activation layer.
+
+ tensor = activation_tensor
+
+ return tensor
+
+
+def aggregated_feature_map(
+ tensor: tf.Tensor,
+ filters: int,
+ kernel_size: int | tuple[int],
+ level1: int,
+ level2: int,
+ upsample_size: tuple[int],
+ padding: str = "same",
+ use_bias: bool = False,
+ batch_normalization: bool = True,
+ activation: str = "relu",
+ kernel_initializer="glorot_uniform",
+ bias_initializer="zeros",
+ kernel_regularizer=None,
+ bias_regularizer=None,
+ activity_regularizer=None,
+ kernel_constraint=None,
+ bias_constraint=None,
+ shared_axes=None,
+ name: str = None,
+):
+ """
+ Connect two nodes in the U-Net 3+ with an aggregated feature map (AFM).
+
+ Parameters
+ ----------
+ tensor: tf.Tensor
+ Input tensor for the convolution module.
+ filters: int
+ Number of filters in the Conv2D/Conv3D layer.
+ kernel_size: int or tuple
+ Size of the kernel in the Conv2D/Conv3D layer.
+ level1: int
+ Level of the first node that is connected to the AFM. This node will provide the input tensor to the AFM. Must be
+ greater than level2 (i.e. the first node must be on a lower level in the U-Net 3+ since we are up-sampling), otherwise
+ a ValueError exception is raised.
+ level2: int
+ Level of the second node that is connected to the AFM. This node will receive the output of the AFM. Must be smaller
+ than level1 (i.e. the second node must be on a higher level in the U-Net 3+ since we are up-sampling), otherwise
+ a ValueError exception is raised.
+ upsample_size: tuple or list of ints
+ Upsampling size for rows and columns in the UpSampling2D/UpSampling3D layer.
+ padding: str
+ Padding in the Conv2D/Conv3D layer. 'valid' will apply no padding, while 'same' will apply padding such that the
+ output shape matches the input shape. 'valid' and 'same' are case-insensitive.
+ use_bias: bool
+ If True, a bias vector will be used in the Conv2D/Conv3D layer.
+ batch_normalization: bool
+ If True, a BatchNormalization layer will follow the Conv2D/Conv3D layer.
+ activation: str
+ Activation function to use after the Conv2D/Conv3D layer (BatchNormalization layer, if batch_normalization is True).
+ See choose_activation_layer for all available activation functions.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+ shared_axes: tuple or list of ints
+ Axes along which to share the learnable parameters for the activation function.
+ name: str or None
+ Prefix of the layer names. If left as None, the layer names are set automatically.
+
+ Returns
+ -------
+ tensor: tf.Tensor
+ Output tensor.
+ """
+
+ tensor_dims = len(
+ tensor.shape
+ ) # Number of dims in the tensor (including the first 'None' dimension for batch size)
+
+ if level1 <= level2:
+ raise ValueError(
+ "level2 must be smaller than level1 in aggregated feature maps"
+ )
+
+ # Arguments for the convolution module.
+ module_kwargs = dict({})
+ module_kwargs["num_modules"] = 1
+ for arg in [
+ "filters",
+ "kernel_size",
+ "padding",
+ "use_bias",
+ "batch_normalization",
+ "activation",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ "shared_axes",
+ "name",
+ ]:
+ module_kwargs[arg] = locals()[arg]
+
+ # Keyword arguments for the UpSampling2D/UpSampling3D layers
+ upsample_kwargs = dict({})
+ upsample_kwargs["name"] = f"{name}_UpSampling{tensor_dims - 2}D"
+ upsample_kwargs["size"] = np.power(upsample_size, abs(level1 - level2))
+
+ if tensor_dims == 4: # If the image is 2D
+ upsample_layer = UpSampling2D
+ if len(upsample_size) != 2:
+ raise TypeError(
+ f"For 2D up-sampling, the pool size must be a tuple or list with 2 integers. Received shape: {np.shape(upsample_size)}"
+ )
+ elif tensor_dims == 5: # If the image is 3D
+ upsample_layer = UpSampling3D
+ if len(upsample_size) != 3:
+ raise TypeError(
+ f"For 3D up-sampling, the pool size must be a tuple or list with 3 integers. Received shape: {np.shape(upsample_size)}"
+ )
+ else:
+ raise TypeError(
+ f"Incompatible tensor shape: {tensor.shape}. The tensor must only have 4 or 5 dimensions"
+ )
+
+ tensor = upsample_layer(**upsample_kwargs)(
+ tensor
+ ) # Pass the tensor through the UpSample2D/UpSample3D layer
+
+ tensor = convolution_module(
+ tensor, **module_kwargs
+ ) # Pass input tensor through convolution module
+
+ return tensor
+
+
+def full_scale_skip_connection(
+ tensor: tf.Tensor,
+ filters: int,
+ kernel_size: int | tuple[int],
+ level1: int,
+ level2: int,
+ pool_size: tuple[int],
+ padding: str = "same",
+ use_bias: bool = False,
+ batch_normalization: bool = True,
+ activation: str = "relu",
+ kernel_initializer="glorot_uniform",
+ bias_initializer="zeros",
+ kernel_regularizer=None,
+ bias_regularizer=None,
+ activity_regularizer=None,
+ kernel_constraint=None,
+ bias_constraint=None,
+ shared_axes=None,
+ name: str = None,
+):
+ """
+ Connect two nodes in the U-Net 3+ with a full-scale skip connection (FSC).
+
+ Parameters
+ ----------
+ tensor: tf.Tensor
+ Input tensor for the convolution module.
+ filters: int
+ Number of filters in the Conv2D/Conv3D layer.
+ kernel_size: int or tuple
+ Size of the kernel in the Conv2D/Conv3D layer.
+ level1: int
+ Level of the first node that is connected to the FSC. This node will provide the input tensor to the FSC. Must be
+ smaller than level2 (i.e. the first node must be on a higher level in the U-Net 3+ since we are max-pooling), otherwise
+ a ValueError exception is raised.
+ level2: int
+ Level of the second node that is connected to the FSC. This node will receive the output of the FSC. Must be greater
+ than level1 (i.e. the second node must be on a lower level in the U-Net 3+ since we are max-pooling), otherwise
+ a ValueError exception is raised.
+ pool_size: tuple or list
+ Pool size for the MaxPooling2D/MaxPooling3D layer.
+ padding: str
+ Padding in the Conv2D/Conv3D layer. 'valid' will apply no padding, while 'same' will apply padding such that the
+ output shape matches the input shape. 'valid' and 'same' are case-insensitive.
+ use_bias: bool
+ If True, a bias vector will be used in the Conv2D/Conv3D layer.
+ batch_normalization: bool
+ If True, a BatchNormalization layer will follow the Conv2D/Conv3D layer.
+ activation: str
+ Activation function to use after the Conv2D/Conv3D layer (BatchNormalization layer, if batch_normalization is True).
+ See choose_activation_layer for all available activation functions.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+ shared_axes: tuple or list of ints
+ Axes along which to share the learnable parameters for the activation function.
+ name: str or None
+ Prefix of the layer names. If left as None, the layer names are set automatically.
+
+ Returns
+ -------
+ tensor: tf.Tensor
+ Output tensor.
+ """
+
+ tensor_dims = len(
+ tensor.shape
+ ) # Number of dims in the tensor (including the first 'None' dimension for batch size)
+
+ if level1 >= level2:
+ raise ValueError(
+ "level2 must be greater than level1 in full-scale skip connections"
+ )
+
+ # Arguments for the convolution module.
+ module_kwargs = dict({})
+ module_kwargs["num_modules"] = 1
+ for arg in [
+ "filters",
+ "kernel_size",
+ "padding",
+ "use_bias",
+ "batch_normalization",
+ "activation",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ "shared_axes",
+ "name",
+ ]:
+ module_kwargs[arg] = locals()[arg]
+
+ # Keyword arguments for the MaxPooling2D/MaxPooling3D layer
+ pool_kwargs = dict({})
+ pool_kwargs["name"] = f"{name}_MaxPool{tensor_dims - 2}D"
+ pool_kwargs["pool_size"] = np.power(pool_size, abs(level1 - level2))
+
+ if tensor_dims == 4: # If the image is 2D
+ pool_layer = MaxPooling2D
+ elif tensor_dims == 5: # If the image is 3D
+ pool_layer = MaxPooling3D
+ else:
+ raise TypeError(
+ f"Incompatible tensor shape: {tensor.shape}. The tensor must only have 4 or 5 dimensions"
+ )
+
+ tensor = pool_layer(**pool_kwargs)(
+ tensor
+ ) # Pass the tensor through the MaxPooling2D/MaxPooling3D layer
+
+ tensor = convolution_module(
+ tensor, **module_kwargs
+ ) # Pass the tensor through the convolution module
+
+ return tensor
+
+
+def conventional_skip_connection(
+ tensor: tf.Tensor,
+ filters: int,
+ kernel_size: int | tuple[int],
+ padding: str = "same",
+ use_bias: bool = False,
+ batch_normalization: bool = True,
+ activation: str = "relu",
+ kernel_initializer="glorot_uniform",
+ bias_initializer="zeros",
+ kernel_regularizer=None,
+ bias_regularizer=None,
+ activity_regularizer=None,
+ kernel_constraint=None,
+ bias_constraint=None,
+ shared_axes=None,
+ name: str = None,
+):
+ """
+ Connect two nodes in the U-Net 3+ with a conventional skip connection.
+
+ Parameters
+ ----------
+ tensor: tf.Tensor
+ Input tensor for the convolution module.
+ filters: int
+ Number of filters in the Conv2D/Conv3D layer.
+ kernel_size: int or tuple
+ Size of the kernel in the Conv2D/Conv3D layer.
+ padding: str
+ Padding in the Conv2D/Conv3D layer. 'valid' will apply no padding, while 'same' will apply padding such that the
+ output shape matches the input shape. 'valid' and 'same' are case-insensitive.
+ use_bias: bool
+ If True, a bias vector will be used in the Conv2D/Conv3D layer.
+ batch_normalization: bool
+ If True, a BatchNormalization layer will follow the Conv2D/Conv3D layer.
+ activation: str
+ Activation function to use after the Conv2D/Conv3D layer (BatchNormalization layer, if batch_normalization is True).
+ See choose_activation_layer for all available activation functions.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+ shared_axes: tuple or list of ints
+ Axes along which to share the learnable parameters for the activation function.
+ name: str or None
+ Prefix of the layer names. If left as None, the layer names are set automatically.
+
+ Returns
+ -------
+ tensor: tf.Tensor
+ Output tensor.
+ """
+
+ # Arguments for the convolution module.
+ module_kwargs = dict({})
+ module_kwargs["num_modules"] = 1
+ for arg in [
+ "filters",
+ "kernel_size",
+ "padding",
+ "use_bias",
+ "batch_normalization",
+ "activation",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ "shared_axes",
+ "name",
+ ]:
+ module_kwargs[arg] = locals()[arg]
+
+ tensor = convolution_module(
+ tensor, **module_kwargs
+ ) # Pass the tensor through the convolution module
+
+ return tensor
+
+
+def max_pool(tensor: tf.Tensor, pool_size: tuple[int], name: str = None):
+ """
+ Connect two encoder nodes with a max-pooling operation.
+
+ Parameters
+ ----------
+ tensor: tf.Tensor
+ Input tensor for the convolution module.
+ pool_size: tuple or list
+ Pool size for the MaxPooling2D/MaxPooling3D layer.
+ name: str or None
+ Prefix of the layer names. If left as None, the layer names are set automatically.
+
+ Returns
+ -------
+ tensor: tf.Tensor
+ Output tensor.
+ """
+
+ if not isinstance(pool_size, tuple) and not isinstance(pool_size, list):
+ raise TypeError(
+ f"pool_size can only be a tuple or list. Received type: {type(pool_size)}"
+ )
+
+ tensor_dims = len(
+ tensor.shape
+ ) # Number of dims in the tensor (including the first 'None' dimension for batch size)
+
+ pool_kwargs = dict({}) # Keyword arguments in the MaxPooling layer
+ pool_kwargs["name"] = f"{name}_MaxPool{tensor_dims - 2}D"
+ pool_kwargs["pool_size"] = pool_size
+
+ if tensor_dims == 4: # If the image is 2D
+ pool_layer = MaxPooling2D
+ if len(pool_size) != 2:
+ raise TypeError(
+ f"For 2D max-pooling, the pool size must be a tuple or list with 2 integers. Received shape: {np.shape(pool_size)}"
+ )
+ elif tensor_dims == 5: # If the image is 3D
+ pool_layer = MaxPooling3D
+ if len(pool_size) != 3:
+ raise TypeError(
+ f"For 3D max-pooling, the pool size must be a tuple or list with 3 integers. Received shape: {np.shape(pool_size)}"
+ )
+ else:
+ raise TypeError(
+ f"Incompatible tensor shape: {tensor.shape}. The tensor must only have 4 or 5 dimensions"
+ )
+
+ pool_tensor = pool_layer(**pool_kwargs)(
+ tensor
+ ) # Pass the tensor through the MaxPooling2D/MaxPooling3D layer
+
+ return pool_tensor
+
+
+def upsample(
+ tensor: tf.Tensor,
+ filters: int,
+ kernel_size: int | tuple[int],
+ upsample_size: tuple[int],
+ padding: str = "same",
+ use_bias: bool = False,
+ batch_normalization: bool = True,
+ activation: str = "relu",
+ kernel_initializer="glorot_uniform",
+ bias_initializer="zeros",
+ kernel_regularizer=None,
+ bias_regularizer=None,
+ activity_regularizer=None,
+ kernel_constraint=None,
+ bias_constraint=None,
+ shared_axes=None,
+ name: str = None,
+):
+ """
+ Connect decoder nodes with an up-sampling operation.
+
+ Parameters
+ ----------
+ tensor: tf.Tensor
+ Input tensor for the convolution module.
+ filters: int
+ Number of filters in the Conv2D/Conv3D layer.
+ kernel_size: int or tuple
+ Size of the kernel in the Conv2D/Conv3D layer.
+ upsample_size: tuple or list
+ Upsampling size in the UpSampling2D/UpSampling3D layer.
+ padding: str
+ Padding in the Conv2D/Conv3D layer. 'valid' will apply no padding, while 'same' will apply padding such that the
+ output shape matches the input shape. 'valid' and 'same' are case-insensitive.
+ use_bias: bool
+ If True, a bias vector will be used in the Conv2D/Conv3D layer.
+ batch_normalization: bool
+ If True, a BatchNormalization layer will follow the Conv2D/Conv3D layer.
+ activation: str
+ Activation function to use after the Conv2D/Conv3D layer (BatchNormalization layer, if batch_normalization is True).
+ See choose_activation_layer for all available activation functions.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+ shared_axes: tuple or list of ints
+ Axes along which to share the learnable parameters for the activation function.
+ name: str or None
+ Prefix of the layer names. If left as None, the layer names are set automatically.
+
+ Returns
+ -------
+ tensor: tf.Tensor
+ Output tensor.
+ """
+
+ tensor_dims = len(
+ tensor.shape
+ ) # Number of dims in the tensor (including the first 'None' dimension for batch size)
+
+ if not isinstance(upsample_size, tuple) and not isinstance(upsample_size, list):
+ raise TypeError(
+ f"upsample_size can only be a tuple or list. Received type: {type(upsample_size)}"
+ )
+
+ # Arguments for the convolution module.
+ module_kwargs = dict({})
+ module_kwargs["num_modules"] = 1
+ for arg in [
+ "filters",
+ "kernel_size",
+ "padding",
+ "use_bias",
+ "batch_normalization",
+ "activation",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ "shared_axes",
+ "name",
+ ]:
+ module_kwargs[arg] = locals()[arg]
+
+ # Keyword arguments in the UpSampling layer
+ upsample_kwargs = dict({})
+ upsample_kwargs["name"] = f"{name}_UpSampling{tensor_dims - 2}D"
+ upsample_kwargs["size"] = upsample_size
+
+ if tensor_dims == 4: # If the image is 2D
+ upsample_layer = UpSampling2D
+ if len(upsample_size) != 2:
+ raise TypeError(
+ f"For 2D up-sampling, the pool size must be a tuple or list with 2 integers. Received shape: {np.shape(upsample_size)}"
+ )
+ elif tensor_dims == 5: # If the image is 3D
+ upsample_layer = UpSampling3D
+ if len(upsample_size) != 3:
+ raise TypeError(
+ f"For 3D up-sampling, the pool size must be a tuple or list with 3 integers. Received shape: {np.shape(upsample_size)}"
+ )
+ else:
+ raise TypeError(
+ f"Incompatible tensor shape: {tensor.shape}. The tensor must only have 4 or 5 dimensions"
+ )
+
+ upsample_tensor = upsample_layer(**upsample_kwargs)(
+ tensor
+ ) # Pass the tensor through the UpSampling2D/UpSampling3D layer
+
+ tensor = convolution_module(
+ upsample_tensor, **module_kwargs
+ ) # Pass the up-sampled tensor through a convolution module
+
+ return tensor
+
+
+def deep_supervision_side_output(
+ tensor: tf.Tensor,
+ num_classes: int,
+ kernel_size: int | tuple[int],
+ output_level: int,
+ upsample_size: tuple[int],
+ activation: str = "softmax",
+ padding: str = "same",
+ use_bias: bool = False,
+ kernel_initializer="glorot_uniform",
+ bias_initializer="zeros",
+ kernel_regularizer=None,
+ bias_regularizer=None,
+ activity_regularizer=None,
+ kernel_constraint=None,
+ bias_constraint=None,
+ squeeze_axes: int | tuple[int] = None,
+ name: str = None,
+):
+ """
+ Deep supervision output. This is usually used on a decoder node in the U-Net 3+ or the final decoder node of a standard
+ U-Net.
+
+ Parameters
+ ----------
+ tensor: tf.Tensor
+ Input tensor for the convolution module.
+ num_classes: int
+ Number of classes that the model is trying to predict.
+ kernel_size: int or tuple
+ Size of the kernel in the Conv2D/Conv3D layer.
+ output_level: int
+ Level of the decoder node from which the deep supervision output is based.
+ upsample_size: tuple or list
+ Upsampling size for rows and columns in the UpSampling2D/UpSampling3D layer. Tuples are currently not supported
+ but will be supported in a future update.
+ activation: str
+ Output activation function.
+ padding: str
+ Padding in the Conv2D/Conv3D layer. 'valid' will apply no padding, while 'same' will apply padding such that the
+ output shape matches the input shape. 'valid' and 'same' are case-insensitive.
+ use_bias: bool
+ If True, a bias vector will be used in the Conv2D/Conv3D layer.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+ squeeze_axes: int, tuple, or None
+ Axis or axes of the input tensor to squeeze.
+ name: str or None
+ Prefix of the layer names. If left as None, the layer names are set automatically.
+
+ Returns
+ -------
+ tensor: tf.Tensor
+ Output tensor.
+ """
+
+ tensor_dims = len(
+ tensor.shape
+ ) # Number of dims in the tensor (including the first 'None' dimension for batch size)
+
+ if tensor_dims == 4: # If the image is 2D
+ conv_layer = Conv2D
+ upsample_layer = UpSampling2D
+ if output_level > 1:
+ upsample_size_1 = upsample_size
+ else:
+ upsample_size_1 = None
+
+ if output_level > 2:
+ upsample_size_2 = np.power(upsample_size, output_level - 2)
+ else:
+ upsample_size_2 = None
+
+ elif tensor_dims == 5: # If the image is 3D
+ conv_layer = Conv3D
+ upsample_layer = UpSampling3D
+ if output_level > 1:
+ upsample_size_1 = upsample_size
+ else:
+ upsample_size_1 = None
+
+ if output_level > 2:
+ upsample_size_2 = np.power(upsample_size, output_level - 2)
+ else:
+ upsample_size_2 = None
+
+ else:
+ raise TypeError(
+ f"Incompatible tensor shape: {tensor.shape}. The tensor can only have 4 or 5 dimensions"
+ )
+
+ # Arguments for the Conv2D/Conv3D layer.
+ conv_kwargs = dict({})
+ conv_kwargs["name"] = f"{name}_Conv{tensor_dims - 2}D"
+ for arg in [
+ "use_bias",
+ "kernel_size",
+ "padding",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ ]:
+ conv_kwargs[arg] = locals()[arg]
+
+ if upsample_size_1 is not None:
+ tensor = upsample_layer(
+ size=upsample_size_1, name=f"{name}_UpSampling{tensor_dims - 2}D_1"
+ )(tensor) # Pass the tensor through the UpSampling2D/UpSampling3D layer
+
+ tensor = conv_layer(filters=num_classes, **conv_kwargs)(
+ tensor
+ ) # This convolution layer contains num_classes filters, one for each class
+
+ if upsample_size_2 is not None:
+ tensor = upsample_layer(
+ size=upsample_size_2, name=f"{name}_UpSampling{tensor_dims - 2}D_2"
+ )(tensor) # Pass the tensor through the UpSampling2D/UpSampling3D layer
+
+ ### Squeeze the given dimensions/axes ###
+ if squeeze_axes is not None:
+ conv_kwargs["kernel_size"] = [1 for _ in range(tensor_dims - 2)]
+
+ if isinstance(squeeze_axes, int):
+ squeeze_axes = [
+ squeeze_axes,
+ ] # Turn integer into a list of length 1 to make indexing easier
+
+ squeeze_axes_sizes = [
+ tensor.shape[ax_to_squeeze] for ax_to_squeeze in squeeze_axes
+ ]
+
+ for ax, size in enumerate(squeeze_axes_sizes):
+ conv_kwargs["kernel_size"][squeeze_axes[ax] - 1] = (
+ size # Kernel size of dimension to squeeze is equal to the size of the dimension because we want the final size to be 1 so it can be squeezed
+ )
+
+ conv_kwargs["padding"] = (
+ "valid" # Padding cannot be 'same' since we want to modify the size of the dimension to be squeezed
+ )
+ conv_kwargs["name"] = f"{name}_Conv{tensor_dims - 2}D_collapse"
+
+ tensor = conv_layer(filters=num_classes, **conv_kwargs)(
+ tensor
+ ) # This convolution layer contains num_classes filters, one for each class
+ tensor = tf.squeeze(
+ tensor, axis=squeeze_axes
+ ) # Squeeze the tensor and remove the dimension
+
+ activation_kwargs = {"name": f"{name}_{activation}"}
+ activation_config = keras_builders.ActivationConfig(
+ name=activation, config=activation_kwargs
+ )
+ activation_layer = activation_config.build()
+ sup_output = activation_layer(tensor) # Final softmax output
+
+ return sup_output
diff --git a/src/fronts/layers/unets.py b/src/fronts/layers/unets.py
new file mode 100644
index 0000000..a498ee5
--- /dev/null
+++ b/src/fronts/layers/unets.py
@@ -0,0 +1,1485 @@
+"""
+Deep learning models and functions for building:
+ * U-Net
+ * U-Net ensemble
+ * U-Net+
+ * U-Net++
+ * U-Net 3+
+ * Attention U-Net
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2024.10.11
+"""
+
+from tensorflow.keras.models import Model
+import numpy as np
+
+from tensorflow.keras.layers import (
+ Concatenate,
+ Input,
+)
+from fronts.utils import keras_builders
+from fronts.layers import (
+ convolution_module,
+ max_pool,
+ upsample,
+ deep_supervision_side_output,
+ full_scale_skip_connection,
+ conventional_skip_connection,
+ attention_gate,
+ aggregated_feature_map,
+)
+import dataclasses
+from typing import Literal
+
+
+@dataclasses.dataclass
+class UNet:
+ """
+ Builds a U-Net model.
+
+ Parameters
+ ----------
+ input_shape: tuple
+ Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
+ num_classes: int
+ Number of classes/labels that the U-Net will try to predict.
+ pool_size: tuple or list
+ Size of the mask in the MaxPooling layers.
+ upsample_size: tuple or list
+ Size of the mask in the UpSampling layers.
+ levels: int
+ Number of levels in the U-Net. Must be greater than 1.
+ filter_num: iterable of ints
+ Number of convolution filters on each level of the U-Net.
+ kernel_size: int or tuple
+ Size of the kernel in the convolution layers.
+ squeeze_axes: int, tuple, list, or None
+ Axis or axes of the input tensor to squeeze.
+ shared_axes: int, tuple, list, or None
+ Axes along which to share the learnable parameters for the activation function. When left as None, parameters will
+ be shared along all arbitrary dimensions (i.e. all dimensions without a defined size).
+ modules_per_node: int
+ Number of modules in each node of the U-Net.
+ batch_normalization: bool
+ Setting this to True will add a batch normalization layer after every convolution in the modules.
+ activation: str
+ Activation function to use in the modules.
+ See utils.choose_activation_layer for all supported activation functions.
+ output_activation: str
+ Output activation function.
+ padding: str
+ Padding to use in the convolution layers.
+ use_bias: bool
+ Setting this to True will implement a bias vector in the convolution layers used in the modules.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+
+ Returns
+ -------
+ model: tf.keras.models.Model object
+ U-Net model.
+
+ Raises
+ ------
+ ValueError
+ If levels < 2
+ If input_shape does not have 3 nor 4 dimensions
+ If the length of filter_num does not match the number of levels
+
+ References
+ ----------
+ https://arxiv.org/pdf/1505.04597.pdf
+ """
+
+ input_shape: tuple[int]
+ num_classes: int
+ pool_size: int | tuple[int] | list[int]
+ upsample_size: int | tuple[int] | list[int]
+ levels: int
+ filter_num: tuple[int] | list[int]
+ kernel_size: int = 3
+ squeeze_axes: int | tuple[int] | list[int] = None
+ shared_axes: int | tuple[int] | list[int] = None
+ modules_per_node: int = 2
+ batch_normalization: bool = True
+ activation: str = "relu"
+ output_activation: str = "softmax"
+ padding: str = "same"
+ use_bias: bool = True
+ kernel_initializer: str = "glorot_uniform"
+ bias_initializer: str = "zeros"
+ kernel_regularizer: str = None
+ bias_regularizer: str = None
+ activity_regularizer: str = None
+ kernel_constraint: str = None
+ bias_constraint: str = None
+
+ def __post_init__(self):
+ # Manage exceptions
+ if self.levels not in [3, 4]:
+ raise ValueError(
+ "Input_shape can only have 3 or 4 dimensions (2D image + 1 dimension "
+ "for channels OR a 3D image + 1 dimension for channels). Received "
+ f"shape: {np.shape(self.input_shape)}"
+ )
+ if len(self.filter_num) != self.levels:
+ raise ValueError(
+ f"Length of filter_num ({len(self.filter_num)}) does not match the "
+ f"number of levels ({self.levels})"
+ )
+ self.ndims = (
+ len(self.input_shape) - 1
+ ) # Number of dimensions in the input image (excluding the last dimension reserved for channels)
+
+ # Keyword arguments for the convolution modules
+ self.module_kwargs = dict({})
+ self.module_kwargs["num_modules"] = self.modules_per_node
+ for arg in [
+ "activation",
+ "batch_normalization",
+ "padding",
+ "kernel_size",
+ "use_bias",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ "shared_axes",
+ ]:
+ self.module_kwargs[arg] = getattr(self, arg)
+
+ # MaxPooling keyword arguments
+ self.pool_kwargs = {"pool_size": self.pool_size}
+
+ # Keyword arguments for upsampling
+ self.upsample_kwargs = dict({})
+ for arg in [
+ "activation",
+ "batch_normalization",
+ "padding",
+ "kernel_size",
+ "use_bias",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ "upsample_size",
+ "shared_axes",
+ ]:
+ self.upsample_kwargs[arg] = getattr(self, arg)
+
+ # Keyword arguments for the deep supervision output in the final decoder node
+ self.supervision_kwargs = dict({})
+ for arg in [
+ "padding",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ "upsample_size",
+ "squeeze_axes",
+ "num_classes",
+ ]:
+ self.supervision_kwargs[arg] = getattr(self, arg)
+
+ def build(self):
+ self.supervision_kwargs["activation"] = self.output_activation
+
+ tensors = dict({}) # Tensors associated with each node and skip connections
+
+ """ Setup the first encoder node with an input layer and a convolution module """
+ tensors["input"] = Input(shape=self.input_shape, name="Input")
+ tensors["En1"] = convolution_module(
+ tensors["input"],
+ filters=self.filter_num[0],
+ name="En1",
+ **self.module_kwargs,
+ )
+
+ """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
+ for encoder in np.arange(
+ 2, self.levels + 1
+ ): # Iterate through the rest of the encoder nodes
+ current_node, previous_node = f"En{encoder}", f"En{encoder - 1}"
+ pool_tensor = max_pool(
+ tensors[previous_node],
+ name=f"{previous_node}-{current_node}",
+ **self.pool_kwargs,
+ ) # Connect the next encoder node with a MaxPooling layer
+ tensors[current_node] = convolution_module(
+ pool_tensor,
+ filters=self.filter_num[encoder - 1],
+ name=current_node,
+ **self.module_kwargs,
+ ) # Convolution modules
+
+ # Connect the bottom encoder node to a decoder node
+ upsample_tensor = upsample(
+ tensors[f"En{self.levels}"],
+ filters=self.filter_num[self.levels - 2],
+ name=f"En{self.levels}-De{self.levels}",
+ **self.upsample_kwargs,
+ )
+
+ """ Bottom decoder node """
+ current_node, next_node = f"De{self.levels - 1}", f"De{self.levels - 2}"
+ skip_node = f"En{self.levels - 1}" # node with an incoming skip connection that connects to 'current_node'
+ tensors[current_node] = Concatenate(name=f"{current_node}_Concatenate")(
+ [tensors[skip_node], upsample_tensor]
+ ) # Concatenate the upsampled tensor and skip connection
+ tensors[current_node] = convolution_module(
+ tensors[current_node],
+ filters=self.filter_num[self.levels - 2],
+ name=current_node,
+ **self.module_kwargs,
+ ) # Convolution module
+ upsample_tensor = upsample(
+ tensors[current_node],
+ filters=self.filter_num[self.levels - 3],
+ name=f"{current_node}-{next_node}",
+ **self.upsample_kwargs,
+ ) # Connect the bottom decoder node to the next decoder node
+
+ """ The rest of the decoder nodes (except the final node) are handled in this loop. Each node contains one concatenation of an upsampled tensor and a skip connection """
+ for decoder in np.arange(2, self.levels - 1)[::-1]:
+ current_node, next_node = f"De{decoder}", f"De{decoder - 1}"
+ skip_node = f"En{decoder}" # node with an incoming skip connection that connects to 'current_node'
+ tensors[current_node] = Concatenate(name=f"{current_node}_Concatenate")(
+ [tensors[skip_node], upsample_tensor]
+ ) # Concatenate the upsampled tensor and skip connection
+ tensors[current_node] = convolution_module(
+ tensors[current_node],
+ filters=self.filter_num[decoder - 1],
+ name=current_node,
+ **self.module_kwargs,
+ ) # Convolution module
+ upsample_tensor = upsample(
+ tensors[current_node],
+ filters=self.filter_num[decoder - 2],
+ name=f"{current_node}-{next_node}",
+ **self.upsample_kwargs,
+ ) # Connect the bottom decoder node to the next decoder node
+
+ """ Final decoder node begins with a concatenation and convolution module, followed by deep supervision """
+ tensor_De1 = Concatenate(name="De1_Concatenate")(
+ [tensors["En1"], upsample_tensor]
+ ) # Concatenate the upsampled tensor and skip connection
+ tensor_De1 = convolution_module(
+ tensor_De1, filters=self.filter_num[0], name="De1", **self.module_kwargs
+ ) # Convolution module
+ tensors["output"] = deep_supervision_side_output(
+ tensor_De1,
+ num_classes=self.num_classes,
+ kernel_size=1,
+ output_level=1,
+ use_bias=True,
+ name="final",
+ **self.supervision_kwargs,
+ ) # Deep supervision - this layer will output the model's prediction
+
+ output_model = Model(
+ inputs=tensors["input"],
+ outputs=tensors["output"],
+ name=f"unet_{self.ndims}D",
+ )
+
+ return output_model
+
+
+@dataclasses.dataclass
+class UNetEnsemble(UNet):
+ """
+ Builds a U-Net ensemble model.
+ https://arxiv.org/pdf/1912.05074.pdf
+
+ Parameters
+ ----------
+ input_shape: tuple
+ Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
+ num_classes: int
+ Number of classes/labels that the U-Net will try to predict.
+ pool_size: tuple or list
+ Size of the mask in the MaxPooling layers.
+ upsample_size: tuple or list
+ Size of the mask in the UpSampling layers.
+ levels: int
+ Number of levels in the U-Net. Must be greater than 1.
+ filter_num: iterable of ints
+ Number of convolution filters on each level of the U-Net.
+ kernel_size: int or tuple
+ Size of the kernel in the convolution layers.
+ squeeze_axes: int, tuple, list, or None
+ Axis or axes of the input tensor to squeeze.
+ shared_axes: int, tuple, list, or None
+ Axes along which to share the learnable parameters for the activation function. When left as None, parameters will
+ be shared along all arbitrary dimensions (i.e. all dimensions without a defined size).
+ modules_per_node: int
+ Number of modules in each node of the U-Net.
+ batch_normalization: bool
+ Setting this to True will add a batch normalization layer after every convolution in the modules.
+ activation: str
+ Activation function to use in the modules.
+ See utils.choose_activation_layer for all supported activation functions.
+ output_activation: str
+ Output activation function.
+ padding: str
+ Padding to use in the convolution layers.
+ use_bias: bool
+ Setting this to True will implement a bias vector in the convolution layers used in the modules.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+
+ Returns
+ -------
+ model: tf.keras.models.Model object
+ U-Net model.
+
+ Raises
+ ------
+ ValueError
+ If levels < 2
+ If input_shape does not have 3 nor 4 dimensions
+ If the length of filter_num does not match the number of levels
+ """
+
+ def build(self):
+ self.supervision_kwargs["activation"] = self.output_activation
+ self.supervision_kwargs["use_bias"] = True
+ self.supervision_kwargs["output_level"] = 1
+ self.supervision_kwargs["kernel_size"] = 1
+
+ tensors = dict({}) # Tensors associated with each node and skip connections
+ tensors_with_supervision = [] # list of output tensors. If deep supervision is used, more than one output will be produced
+
+ """ Setup the first encoder node with an input layer and a convolution module """
+ tensors["input"] = Input(shape=self.input_shape, name="Input")
+ tensors["En1"] = convolution_module(
+ tensors["input"],
+ filters=self.filter_num[0],
+ name="En1",
+ **self.module_kwargs,
+ )
+
+ """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
+ for encoder in np.arange(
+ 2, self.levels + 1
+ ): # Iterate through the rest of the encoder nodes
+ current_node, previous_node = f"En{encoder}", f"En{encoder - 1}"
+ pool_tensor = max_pool(
+ tensors[previous_node],
+ name=f"{previous_node}-{current_node}",
+ **self.pool_kwargs,
+ ) # Connect the next encoder node with a MaxPooling layer
+ tensors[current_node] = convolution_module(
+ pool_tensor,
+ filters=self.filter_num[encoder - 1],
+ name=current_node,
+ **self.module_kwargs,
+ ) # Convolution modules
+
+ # Connect the bottom encoder node to a decoder node
+ upsample_tensor = upsample(
+ tensors[f"En{self.levels}"],
+ filters=self.filter_num[self.levels - 2],
+ name=f"En{self.levels}-De{self.levels}",
+ **self.upsample_kwargs,
+ )
+
+ """ Bottom decoder node """
+ current_node, next_node = f"De{self.levels - 1}", f"De{self.levels - 2}"
+ skip_node = f"En{self.levels - 1}"
+ tensors[current_node] = Concatenate(name=f"{current_node}_Concatenate")(
+ [upsample_tensor, tensors[skip_node]]
+ ) # Concatenate the upsampled tensor and skip connection
+ tensors[current_node] = convolution_module(
+ tensors[current_node],
+ filters=self.filter_num[self.levels - 2],
+ name=current_node,
+ **self.module_kwargs,
+ ) # Convolution module
+ upsample_tensor = upsample(
+ tensors[current_node],
+ filters=self.filter_num[self.levels - 3],
+ name=f"{current_node}-{next_node}",
+ **self.upsample_kwargs,
+ ) # Connect the bottom decoder node to the next decoder node
+
+ for decoder in np.arange(1, self.levels - 1)[::-1]:
+ num_middle_nodes = self.levels - decoder - 1
+ for node in range(1, num_middle_nodes + 1):
+ if node == 1: # if on the first middle node at the given level
+ upsample_tensor_for_middle_node = upsample(
+ tensors[f"En{decoder + 1}"],
+ filters=self.filter_num[decoder - 2],
+ name=f"En{decoder + 1}-Me{decoder}-1",
+ **self.upsample_kwargs,
+ )
+ else:
+ upsample_tensor_for_middle_node = upsample(
+ tensors[f"Me{decoder + 1}-{node - 1}"],
+ filters=self.filter_num[decoder - 2],
+ name=f"Me{decoder + 1}-{node - 1}-Me{decoder}-{node}",
+ **self.upsample_kwargs,
+ )
+ tensors[f"Me{decoder}-{node}"] = Concatenate(
+ name=f"Me{decoder}-{node}_Concatenate"
+ )([tensors[f"En{decoder}"], upsample_tensor_for_middle_node])
+ tensors[f"Me{decoder}-{node}"] = convolution_module(
+ tensors[f"Me{decoder}-{node}"],
+ filters=self.filter_num[decoder - 1],
+ name=f"Me{decoder}-{node}",
+ **self.module_kwargs,
+ ) # Convolution module
+ if decoder == 1:
+ tensors[f"sup{decoder}-{node}"] = deep_supervision_side_output(
+ tensors[f"Me{decoder}-{node}"],
+ name=f"sup{decoder}-{node}",
+ **self.supervision_kwargs,
+ ) # deep supervision on middle node located on top level
+ tensors_with_supervision.append(tensors[f"sup{decoder}-{node}"])
+ tensors[f"De{decoder}"] = Concatenate(name=f"De{decoder}_Concatenate")(
+ [tensors[f"En{decoder}"], upsample_tensor]
+ ) # Concatenate the upsampled tensor and skip connection
+ tensors[f"De{decoder}"] = convolution_module(
+ tensors[f"De{decoder}"],
+ filters=self.filter_num[decoder - 1],
+ name=f"De{decoder}",
+ **self.module_kwargs,
+ ) # Convolution module
+
+ if decoder != 1: # if not currently on the final decoder node (De1)
+ upsample_tensor = upsample(
+ tensors[f"De{decoder}"],
+ filters=self.filter_num[decoder - 2],
+ name=f"De{decoder}-De{decoder - 1}",
+ **self.upsample_kwargs,
+ ) # Connect the bottom decoder node to the next decoder node
+ else:
+ tensors["output"] = deep_supervision_side_output(
+ tensors["De1"], name="final", **self.supervision_kwargs
+ ) # Deep supervision - this layer will output the model's prediction
+ tensors_with_supervision.append(tensors["output"])
+
+ output_model = Model(
+ inputs=tensors["input"],
+ outputs=tensors_with_supervision,
+ name=f"unet_ensemble_{self.ndims}D",
+ )
+
+ return output_model
+
+
+class UNetPlus(UNet):
+ """
+ Builds a U-Net+ model.
+ https://arxiv.org/pdf/1912.05074.pdf
+
+ Parameters
+ ----------
+ input_shape: tuple
+ Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
+ num_classes: int
+ Number of classes/labels that the U-Net will try to predict.
+ pool_size: tuple or list
+ Size of the mask in the MaxPooling layers.
+ upsample_size: tuple or list
+ Size of the mask in the UpSampling layers.
+ levels: int
+ Number of levels in the U-Net. Must be greater than 1.
+ filter_num: iterable of ints
+ Number of convolution filters on each level of the U-Net.
+ kernel_size: int or tuple
+ Size of the kernel in the convolution layers.
+ squeeze_axes: int, tuple, list, or None
+ Axis or axes of the input tensor to squeeze.
+ shared_axes: int, tuple, list, or None
+ Axes along which to share the learnable parameters for the activation function. When left as None, parameters will
+ be shared along all arbitrary dimensions (i.e. all dimensions without a defined size).
+ modules_per_node: int
+ Number of modules in each node of the U-Net.
+ batch_normalization: bool
+ Setting this to True will add a batch normalization layer after every convolution in the modules.
+ deep_supervision: bool
+ Add deep supervision side outputs to each top node.
+ NOTE: The final decoder node requires deep supervision and is not affected if this parameter is False.
+ activation: str
+ Activation function to use in the modules.
+ See utils.choose_activation_layer for all supported activation functions.
+ output_activation: str
+ Output activation function.
+ padding: str
+ Padding to use in the convolution layers.
+ use_bias: bool
+ Setting this to True will implement a bias vector in the convolution layers used in the modules.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+
+ Returns
+ -------
+ model: tf.keras.models.Model object
+ U-Net model.
+
+ Raises
+ ------
+ ValueError
+ If levels < 2
+ If input_shape does not have 3 nor 4 dimensions
+ If the length of filter_num does not match the number of levels
+ """
+
+ def build(self):
+ self.supervision_kwargs["activation"] = self.output_activation
+ self.supervision_kwargs["use_bias"] = True
+ self.supervision_kwargs["output_level"] = 1
+ self.supervision_kwargs["kernel_size"] = 1
+
+ tensors = dict({}) # Tensors associated with each node and skip connections
+ tensors_with_supervision = [] # list of output tensors. If deep supervision is used, more than one output will be produced
+
+ """ Setup the first encoder node with an input layer and a convolution module """
+ tensors["input"] = Input(shape=self.input_shape, name="Input")
+ tensors["En1"] = convolution_module(
+ tensors["input"],
+ filters=self.filter_num[0],
+ name="En1",
+ **self.module_kwargs,
+ )
+
+ """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
+ for encoder in np.arange(
+ 2, self.levels + 1
+ ): # Iterate through the rest of the encoder nodes
+ pool_tensor = max_pool(
+ tensors[f"En{encoder - 1}"],
+ name=f"En{encoder - 1}-En{encoder}",
+ **self.pool_kwargs,
+ ) # Connect the next encoder node with a MaxPooling layer
+ tensors[f"En{encoder}"] = convolution_module(
+ pool_tensor,
+ filters=self.filter_num[encoder - 1],
+ name=f"En{encoder}",
+ **self.module_kwargs,
+ ) # Convolution modules
+
+ # Connect the bottom encoder node to a decoder node
+ upsample_tensor = upsample(
+ tensors[f"En{self.levels}"],
+ filters=self.filter_num[self.levels - 2],
+ name=f"En{self.levels}-De{self.levels}",
+ **self.upsample_kwargs,
+ )
+
+ """ Bottom decoder node """
+ tensors[f"De{self.levels - 1}"] = Concatenate(
+ name=f"De{self.levels - 1}_Concatenate"
+ )(
+ [upsample_tensor, tensors[f"En{self.levels - 1}"]]
+ ) # Concatenate the upsampled tensor and skip connection
+ tensors[f"De{self.levels - 1}"] = convolution_module(
+ tensors[f"De{self.levels - 1}"],
+ filters=self.filter_num[self.levels - 2],
+ name=f"De{self.levels - 1}",
+ **self.module_kwargs,
+ ) # Convolution module
+ upsample_tensor = upsample(
+ tensors[f"De{self.levels - 1}"],
+ filters=self.filter_num[self.levels - 3],
+ name=f"De{self.levels - 1}-De{self.levels - 2}",
+ **self.upsample_kwargs,
+ ) # Connect the bottom decoder node to the next decoder node
+
+ """ The rest of the decoder nodes (except the final node) are handled in this loop. Each node contains one concatenation of an upsampled tensor and a skip connection """
+ for decoder in np.arange(1, self.levels - 1)[::-1]:
+ num_middle_nodes = self.levels - decoder - 1
+ for node in range(1, num_middle_nodes + 1):
+ if node == 1: # if on the first middle node at the given level
+ upsample_tensor_for_middle_node = upsample(
+ tensors[f"En{decoder + 1}"],
+ filters=self.filter_num[decoder - 2],
+ name=f"En{decoder + 1}-Me{decoder}-1",
+ **self.upsample_kwargs,
+ )
+ tensors[f"Me{decoder}-1"] = Concatenate(
+ name=f"Me{decoder}-1_Concatenate"
+ )([tensors[f"En{decoder}"], upsample_tensor_for_middle_node])
+ else:
+ upsample_tensor_for_middle_node = upsample(
+ tensors[f"Me{decoder + 1}-{node - 1}"],
+ filters=self.filter_num[decoder - 2],
+ name=f"Me{decoder + 1}-{node - 1}-Me{decoder}-{node}",
+ **self.upsample_kwargs,
+ )
+ tensors[f"Me{decoder}-{node}"] = Concatenate(
+ name=f"Me{decoder}-{node}_Concatenate"
+ )(
+ [
+ tensors[f"Me{decoder}-{node - 1}"],
+ upsample_tensor_for_middle_node,
+ ]
+ )
+ tensors[f"Me{decoder}-{node}"] = convolution_module(
+ tensors[f"Me{decoder}-{node}"],
+ filters=self.filter_num[decoder - 1],
+ name=f"Me{decoder}-{node}",
+ **self.module_kwargs,
+ ) # Convolution module
+ if decoder == 1 and self.deep_supervision:
+ tensors[f"sup{decoder}-{node}"] = deep_supervision_side_output(
+ tensors[f"Me{decoder}-{node}"],
+ name=f"sup{decoder}-{node}",
+ **self.supervision_kwargs,
+ ) # deep supervision on middle node located on top level
+ tensors_with_supervision.append(tensors[f"sup{decoder}-{node}"])
+ tensors[f"De{decoder}"] = Concatenate(name=f"De{decoder}_Concatenate")(
+ [tensors[f"Me{decoder}-{num_middle_nodes}"], upsample_tensor]
+ ) # Concatenate the upsampled tensor and skip connection
+ tensors[f"De{decoder}"] = convolution_module(
+ tensors[f"De{decoder}"],
+ filters=self.filter_num[decoder - 1],
+ name=f"De{decoder}",
+ **self.module_kwargs,
+ ) # Convolution module
+
+ if decoder != 1: # if not currently on the final decoder node (De1)
+ upsample_tensor = upsample(
+ tensors[f"De{decoder}"],
+ filters=self.filter_num[decoder - 2],
+ name=f"De{decoder}-De{decoder - 1}",
+ **self.upsample_kwargs,
+ ) # Connect the bottom decoder node to the next decoder node
+ else:
+ tensors["output"] = deep_supervision_side_output(
+ tensors["De1"], **self.supervision_kwargs
+ ) # Deep supervision - this layer will output the model's prediction
+ tensors_with_supervision.append(tensors["output"])
+
+ output_model = Model(
+ inputs=tensors["input"],
+ outputs=tensors_with_supervision,
+ name=f"unet_plus_{self.ndims}D",
+ )
+
+ return output_model
+
+
+class UNet2Plus(UNet):
+ """
+ Builds a U-Net++ model.
+ https://arxiv.org/pdf/1912.05074.pdf
+
+ Parameters
+ ----------
+ input_shape: tuple
+ Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
+ num_classes: int
+ Number of classes/labels that the U-Net will try to predict.
+ pool_size: tuple or list
+ Size of the mask in the MaxPooling layers.
+ upsample_size: tuple or list
+ Size of the mask in the UpSampling layers.
+ levels: int
+ Number of levels in the U-Net. Must be greater than 1.
+ filter_num: iterable of ints
+ Number of convolution filters on each level of the U-Net.
+ kernel_size: int or tuple
+ Size of the kernel in the convolution layers.
+ squeeze_axes: int, tuple, list, or None
+ Axis or axes of the input tensor to squeeze.
+ shared_axes: int, tuple, list, or None
+ Axes along which to share the learnable parameters for the activation function. When left as None, parameters will
+ be shared along all arbitrary dimensions (i.e. all dimensions without a defined size).
+ modules_per_node: int
+ Number of modules in each node of the U-Net.
+ batch_normalization: bool
+ Setting this to True will add a batch normalization layer after every convolution in the modules.
+ deep_supervision: bool
+ Add deep supervision side outputs to each top node.
+ NOTE: The final decoder node requires deep supervision and is not affected if this parameter is False.
+ activation: str
+ Activation function to use in the modules.
+ See utils.choose_activation_layer for all supported activation functions.
+ output_activation: str
+ Output activation function.
+ padding: str
+ Padding to use in the convolution layers.
+ use_bias: bool
+ Setting this to True will implement a bias vector in the convolution layers used in the modules.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+
+ Returns
+ -------
+ model: tf.keras.models.Model object
+ U-Net model.
+
+ Raises
+ ------
+ ValueError
+ If levels < 2
+ If input_shape does not have 3 nor 4 dimensions
+ If the length of filter_num does not match the number of levels
+ """
+
+ def build(self):
+ self.supervision_kwargs["activation"] = self.output_activation
+ self.supervision_kwargs["use_bias"] = True
+ self.supervision_kwargs["output_level"] = 1
+ self.supervision_kwargs["kernel_size"] = 1
+
+ tensors = dict({}) # Tensors associated with each node and skip connections
+ tensors_with_supervision = [] # list of output tensors. If deep supervision is used, more than one output will be produced
+
+ """ Setup the first encoder node with an input layer and a convolution module """
+ tensors["input"] = Input(shape=self.input_shape, name="Input")
+ tensors["En1"] = convolution_module(
+ tensors["input"],
+ filters=self.lfilter_num[0],
+ name="En1",
+ **self.lmodule_kwargs,
+ )
+
+ """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
+ for encoder in np.arange(
+ 2, self.levels + 1
+ ): # Iterate through the rest of the encoder nodes
+ pool_tensor = max_pool(
+ tensors[f"En{encoder - 1}"],
+ name=f"En{encoder - 1}-En{encoder}",
+ **self.lpool_kwargs,
+ ) # Connect the next encoder node with a MaxPooling layer
+ tensors[f"En{encoder}"] = convolution_module(
+ pool_tensor,
+ filters=self.lfilter_num[encoder - 1],
+ name=f"En{encoder}",
+ **self.lmodule_kwargs,
+ ) # Convolution modules
+
+ # Connect the bottom encoder node to a decoder node
+ upsample_tensor = upsample(
+ tensors[f"En{self.levels}"],
+ filters=self.lfilter_num[self.levels - 2],
+ name=f"En{self.levels}-De{self.levels}",
+ **self.lupsample_kwargs,
+ )
+
+ """ Bottom decoder node """
+ tensors[f"De{self.levels - 1}"] = Concatenate(
+ name=f"De{self.levels - 1}_Concatenate"
+ )(
+ [upsample_tensor, tensors[f"En{self.levels - 1}"]]
+ ) # Concatenate the upsampled tensor and skip connection
+ tensors[f"De{self.levels - 1}"] = convolution_module(
+ tensors[f"De{self.levels - 1}"],
+ filters=self.lfilter_num[self.levels - 2],
+ name=f"De{self.levels - 1}",
+ **self.lmodule_kwargs,
+ ) # Convolution module
+ upsample_tensor = upsample(
+ tensors[f"De{self.levels - 1}"],
+ filters=self.lfilter_num[self.levels - 3],
+ name=f"De{self.levels - 1}-De{self.levels - 2}",
+ **self.lupsample_kwargs,
+ ) # Connect the bottom decoder node to the next decoder node
+
+ """ The rest of the decoder nodes (except the final node) are handled in this loop. Each node contains one concatenation of an upsampled tensor and a skip connection """
+ for decoder in np.arange(1, self.levels - 1)[::-1]:
+ num_middle_nodes = self.levels - decoder - 1
+ for node in range(1, num_middle_nodes + 1):
+ if node == 1: # if on the first middle node at the given level
+ upsample_tensor_for_middle_node = upsample(
+ tensors[f"En{decoder + 1}"],
+ filters=self.lfilter_num[decoder - 2],
+ name=f"En{decoder + 1}-Me{decoder}-1",
+ **self.lupsample_kwargs,
+ )
+ tensors[f"Me{decoder}-1"] = Concatenate(
+ name=f"Me{decoder}-1_Concatenate"
+ )([tensors[f"En{decoder}"], upsample_tensor_for_middle_node])
+ else:
+ upsample_tensor_for_middle_node = upsample(
+ tensors[f"Me{decoder + 1}-{node - 1}"],
+ filters=self.lfilter_num[decoder - 2],
+ name=f"Me{decoder + 1}-{node - 1}-Me{decoder}-{node}",
+ **self.lupsample_kwargs,
+ )
+ tensors_to_concatenate = [] # Tensors to concatenate in the middle node
+ connections_to_add = sorted(
+ [tensor for tensor in tensors if f"Me{decoder}" in tensor]
+ )[
+ ::-1
+ ] # skip connections to add to the list of tensors to concatenate
+ for connection in connections_to_add:
+ tensors_to_concatenate.append(tensors[connection])
+ tensors_to_concatenate.append(tensors[f"En{decoder}"])
+ tensors_to_concatenate.append(upsample_tensor_for_middle_node)
+ tensors[f"Me{decoder}-{node}"] = Concatenate(
+ name=f"Me{decoder}-{node}_Concatenate"
+ )(tensors_to_concatenate)
+ tensors[f"Me{decoder}-{node}"] = convolution_module(
+ tensors[f"Me{decoder}-{node}"],
+ filters=self.lfilter_num[decoder - 1],
+ name=f"Me{decoder}-{node}",
+ **self.lmodule_kwargs,
+ ) # Convolution module
+
+ if decoder == 1 and self.ldeep_supervision:
+ tensors[f"sup{decoder}-{node}"] = deep_supervision_side_output(
+ tensors[f"Me{decoder}-{node}"],
+ name=f"sup{decoder}-{node}",
+ **self.supervision_kwargs,
+ ) # deep supervision on middle node located on top level
+ tensors_with_supervision.append(tensors[f"sup{decoder}-{node}"])
+
+ tensors_to_concatenate = [] # tensors to concatenate in the decoder node
+ connections_to_add = sorted(
+ [tensor for tensor in tensors if f"Me{decoder}" in tensor]
+ )[::-1] # skip connections to add to the list of tensors to concatenate
+ for connection in connections_to_add:
+ tensors_to_concatenate.append(tensors[connection])
+ tensors_to_concatenate.append(tensors[f"En{decoder}"])
+ tensors_to_concatenate.append(upsample_tensor)
+ tensors[f"De{decoder}"] = Concatenate(name=f"De{decoder}_Concatenate")(
+ tensors_to_concatenate
+ ) # Concatenate the upsampled tensor and skip connection
+ tensors[f"De{decoder}"] = convolution_module(
+ tensors[f"De{decoder}"],
+ filters=self.lfilter_num[decoder - 1],
+ name=f"De{decoder}",
+ **self.lmodule_kwargs,
+ ) # Convolution module
+
+ if decoder != 1: # if not currently on the final decoder node (De1)
+ upsample_tensor = upsample(
+ tensors[f"De{decoder}"],
+ filters=self.lfilter_num[decoder - 2],
+ name=f"De{decoder}-De{decoder - 1}",
+ **self.lupsample_kwargs,
+ ) # Connect the bottom decoder node to the next decoder node
+ else:
+ tensors["output"] = deep_supervision_side_output(
+ tensors["De1"], name="final", **self.supervision_kwargs
+ ) # Deep supervision - this layer will output the model's prediction
+ tensors_with_supervision.append(tensors["output"])
+
+ output_model = Model(
+ inputs=tensors["input"],
+ outputs=tensors_with_supervision,
+ name=f"unet_2plus_{self.ndims}D",
+ )
+
+ return output_model
+
+
+@dataclasses.dataclass
+class UNet3Plus(UNet):
+ """
+ Creates a U-Net 3+.
+ https://arxiv.org/ftp/arxiv/papers/2004/2004.08790.pdf
+
+ Parameters
+ ----------
+ input_shape: tuple
+ Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
+ num_classes: int
+ Number of classes/labels that the U-Net 3+ will try to predict.
+ pool_size: tuple or list
+ Size of the mask in the MaxPooling layers.
+ upsample_size: tuple or list
+ Size of the mask in the UpSampling layers.
+ levels: int
+ Number of levels in the U-Net 3+. Must be greater than 2.
+ filter_num: iterable of ints
+ Number of convolution filters in each encoder of the U-Net 3+. The length must be equal to 'levels'.
+ filter_num_skip: int or None
+ Number of convolution filters in the conventional skip connections, full-scale skip connections, and aggregated feature maps.
+ NOTE: When left as None, this will default to the first value in the 'filter_num' iterable.
+ filter_num_aggregate: int or None
+ Number of convolution filters in the decoder nodes after images are concatenated.
+ When left as None, this will be equal to the product of filter_num_skip and the number of levels.
+ kernel_size: int or tuple
+ Size of the kernel in the convolution layers.
+ first_encoder_connections: bool
+ Setting this to True will create full-scale skip connections attached to the first encoder node.
+ squeeze_axes: int, tuple, list, or None
+ Axis or axes of the input tensor to squeeze.
+ shared_axes: int, tuple, list, or None
+ Axes along which to share the learnable parameters for the activation function. When left as None, parameters will
+ be shared along all arbitrary dimensions (i.e. all dimensions without a defined size).
+ modules_per_node: int
+ Number of modules in each node of the U-Net 3+.
+ batch_normalization: bool
+ Setting this to True will add a batch normalization layer after every convolution in the modules.
+ deep_supervision: bool
+ Add deep supervision side outputs to each decoder node.
+ NOTE: The final decoder node requires deep supervision and is not affected if this parameter is False.
+ activation: str
+ Activation function to use in the modules.
+ See utils.choose_activation_layer for all supported activation functions.
+ output_activation: str
+ Output activation function.
+ padding: str
+ Padding to use in the convolution layers.
+ use_bias: bool
+ Setting this to True will implement a bias vector in the convolution layers used in the modules.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+
+ Returns
+ -------
+ model: tf.keras.models.Model object
+ U-Net 3+ model.
+ """
+
+ def build(self):
+ ndims = (
+ len(self.input_shape) - 1
+ ) # Number of dimensions in the input image (excluding the last dimension reserved for channels)
+
+ if self.levels < 3:
+ raise ValueError(
+ f"levels must be greater than 2. Received value: {self.levels}"
+ )
+ if len(self.input_shape) > 4 or len(self.input_shape) < 3:
+ raise ValueError(
+ f"input_shape can only have 3 or 4 dimensions (2D image + 1 dimension for channels OR a 3D image + 1 dimension for channels). Received shape: {np.shape(self.input_shape)}"
+ )
+ if len(self.filter_num) != self.levels:
+ raise ValueError(
+ f"length of filter_num ({len(self.filter_num)}) does not match the number of levels ({self.levels})"
+ )
+
+ if self.filter_num_skip is None:
+ filter_num_skip = self.filter_num[0]
+
+ if self.filter_num_aggregate is None:
+ filter_num_aggregate = self.levels * filter_num_skip
+
+ # Keyword arguments for the convolution modules
+ module_kwargs = dict({})
+ for arg in [
+ "activation",
+ "batch_normalization",
+ "padding",
+ "kernel_size",
+ "use_bias",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ "shared_axes",
+ ]:
+ module_kwargs[arg] = getattr(self, arg)
+ module_kwargs["num_modules"] = self.modules_per_node
+
+ pool_kwargs = {"pool_size": self.pool_size}
+
+ upsample_kwargs = dict({})
+ conventional_kwargs = dict({})
+ full_scale_kwargs = dict({})
+ aggregated_kwargs = dict({})
+ for arg in [
+ "activation",
+ "batch_normalization",
+ "kernel_size",
+ "padding",
+ "use_bias",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ "shared_axes",
+ ]:
+ upsample_kwargs[arg] = getattr(self, arg)
+ conventional_kwargs[arg] = getattr(self, arg)
+ full_scale_kwargs[arg] = getattr(self, arg)
+ aggregated_kwargs[arg] = getattr(self, arg)
+
+ conventional_kwargs["filters"] = filter_num_skip
+ upsample_kwargs["filters"] = filter_num_skip
+ upsample_kwargs["upsample_size"] = self.upsample_size
+ full_scale_kwargs["filters"] = filter_num_skip
+ full_scale_kwargs["pool_size"] = self.pool_size
+ aggregated_kwargs["filters"] = filter_num_skip
+ aggregated_kwargs["upsample_size"] = self.upsample_size
+
+ supervision_kwargs = dict({})
+ for arg in [
+ "kernel_size",
+ "padding",
+ "squeeze_axes",
+ "kernel_initializer",
+ "bias_initializer",
+ "kernel_regularizer",
+ "bias_regularizer",
+ "activity_regularizer",
+ "kernel_constraint",
+ "bias_constraint",
+ "upsample_size",
+ ]:
+ supervision_kwargs[arg] = getattr(self, arg)
+ supervision_kwargs["activation"] = self.output_activation
+ supervision_kwargs["use_bias"] = True
+
+ tensors = dict({}) # Tensors associated with each node and skip connections
+ tensors_with_supervision = [] # Outputs of deep supervision
+
+ """ Setup the first encoder node with an input layer and a convolution module (we are not using skip connections here) """
+ tensors["input"] = Input(shape=self.input_shape, name="Input")
+ tensors["En1"] = convolution_module(
+ tensors["input"], filters=self.filter_num[0], name="En1", **module_kwargs
+ )
+
+ if self.first_encoder_connections is True:
+ for full_connection in range(2, self.levels):
+ tensors[f"1---{full_connection}_full-scale"] = (
+ full_scale_skip_connection(
+ tensors["En1"],
+ level1=1,
+ level2=full_connection,
+ name=f"1---{full_connection}_full-scale",
+ **full_scale_kwargs,
+ )
+ )
+
+ """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
+ for encoder in np.arange(
+ 2, self.levels
+ ): # Iterate through the rest of the encoder nodes
+ pool_tensor = max_pool(
+ tensors[f"En{encoder - 1}"],
+ name=f"En{encoder - 1}-En{encoder}",
+ **pool_kwargs,
+ ) # Connect the next encoder node with a MaxPooling layer
+ tensors[f"En{encoder}"] = convolution_module(
+ pool_tensor,
+ filters=self.filter_num[encoder - 1],
+ name=f"En{encoder}",
+ **module_kwargs,
+ ) # Convolution modules
+ tensors[f"{encoder}---{encoder}_skip"] = conventional_skip_connection(
+ tensors[f"En{encoder}"],
+ name=f"{encoder}---{encoder}_skip",
+ **conventional_kwargs,
+ )
+
+ # Create full-scale skip connections
+ for full_connection in range(encoder + 1, self.levels):
+ tensors[f"{encoder}---{full_connection}_full-scale"] = (
+ full_scale_skip_connection(
+ tensors[f"En{encoder}"],
+ level1=encoder,
+ level2=full_connection,
+ name=f"{encoder}---{full_connection}_full-scale",
+ **full_scale_kwargs,
+ )
+ )
+
+ # Bottom encoder node
+ tensors[f"En{self.levels}"] = max_pool(
+ tensors[f"En{self.levels - 1}"],
+ name=f"En{self.levels - 1}-En{self.levels}",
+ **pool_kwargs,
+ )
+ tensors[f"En{self.levels}"] = convolution_module(
+ tensors[f"En{self.levels}"],
+ filters=self.filter_num[self.levels - 1],
+ name=f"En{self.levels}",
+ **module_kwargs,
+ )
+ if self.deep_supervision:
+ tensors[f"sup{self.levels}_output"] = deep_supervision_side_output(
+ tensors[f"En{self.levels}"],
+ num_classes=self.num_classes,
+ output_level=self.levels,
+ name=f"sup{self.levels}",
+ **supervision_kwargs,
+ )
+ tensors_with_supervision.append(tensors[f"sup{self.levels}_output"])
+
+ # Add aggregated feature maps using the bottom encoder node
+ for feature_map in range(1, self.levels - 1):
+ tensors[f"{self.levels}---{feature_map}_feature"] = aggregated_feature_map(
+ tensors[f"En{self.levels}"],
+ level1=self.levels,
+ level2=feature_map,
+ name=f"{self.levels}---{feature_map}_feature",
+ **aggregated_kwargs,
+ )
+
+ """ Build the rest of the decoder nodes """
+ for decoder in np.arange(1, self.levels)[::-1]:
+ """ The lowest decoder node (levels - 1) is attached to the bottom encoder node via upsampling, so concatenation is slightly different """
+ if decoder == self.levels - 1:
+ tensors[f"De{decoder}"] = upsample(
+ tensors[f"En{self.levels}"],
+ name=f"En{self.levels}-De{decoder}",
+ **upsample_kwargs,
+ )
+
+ # Tensors to concatenate in the Concatenate layer
+ tensors_to_concatenate = [
+ tensors[f"De{decoder}"],
+ ]
+ connections_to_add = sorted(
+ [tensor for tensor in tensors if f"---{decoder}" in tensor]
+ )[::-1]
+ for connection in connections_to_add:
+ tensors_to_concatenate.append(tensors[connection])
+ else:
+ tensors[f"De{decoder}"] = upsample(
+ tensors[f"De{decoder + 1}"],
+ name=f"De{decoder + 1}-De{decoder}",
+ **upsample_kwargs,
+ )
+
+ # Tensors to concatenate in the Concatenate layer
+ tensors_to_concatenate = sorted(
+ [tensor for tensor in tensors if f"---{decoder}" in tensor]
+ )[::-1]
+ for index in range(len(tensors_to_concatenate)):
+ tensors_to_concatenate[index] = tensors[
+ tensors_to_concatenate[index]
+ ]
+ tensors_to_concatenate.insert(
+ self.levels - 1 - decoder, tensors[f"De{decoder}"]
+ )
+
+ # Concatenate tensors, pass through convolution modules, then use deep supervision to create a side output
+ tensors[f"De{decoder}"] = Concatenate(name=f"De{decoder}_Concatenate")(
+ tensors_to_concatenate
+ )
+ tensors[f"De{decoder}"] = convolution_module(
+ tensors[f"De{decoder}"],
+ filters=filter_num_aggregate,
+ name=f"De{decoder}",
+ **module_kwargs,
+ )
+ if (
+ self.deep_supervision or decoder == 1
+ ): # Decoder node 1 must always have deep supervision
+ tensors[f"sup{decoder}_output"] = deep_supervision_side_output(
+ tensors[f"De{decoder}"],
+ num_classes=self.num_classes,
+ output_level=decoder,
+ name=f"sup{decoder}",
+ **supervision_kwargs,
+ )
+ tensors_with_supervision.append(tensors[f"sup{decoder}_output"])
+
+ """ Add aggregated feature maps """
+ for feature_map in range(1, decoder - 1):
+ tensors[f"{decoder}---{feature_map}_feature"] = aggregated_feature_map(
+ tensors[f"De{decoder}"],
+ level1=decoder,
+ level2=feature_map,
+ name=f"{decoder}---{feature_map}_feature",
+ **aggregated_kwargs,
+ )
+
+ output_model = Model(
+ inputs=tensors["input"],
+ outputs=tensors_with_supervision[::-1],
+ name=f"unet_3plus_{ndims}D",
+ )
+
+ return output_model
+
+
+@dataclasses.dataclass
+class AttentionUNet(UNet):
+ """
+ Builds an Attention U-Net model.
+
+ Parameters
+ ----------
+ input_shape: tuple
+ Shape of the inputs. The last number in the tuple represents the number of channels/predictors.
+ num_classes: int
+ Number of classes/labels that the U-Net will try to predict.
+ pool_size: tuple or list
+ Size of the mask in the MaxPooling and UpSampling layers.
+ levels: int
+ Number of levels in the U-Net. Must be greater than 1.
+ filter_num: iterable of ints
+ Number of convolution filters on each level of the U-Net.
+ kernel_size: int or tuple
+ Size of the kernel in the convolution layers.
+ squeeze_axes: int, tuple, list, or None
+ Axis or axes of the input tensor to squeeze.
+ shared_axes: int, tuple, list, or None
+ Axes along which to share the learnable parameters for the activation function. When left as None, parameters will
+ be shared along all arbitrary dimensions (i.e. all dimensions without a defined size).
+ modules_per_node: int
+ Number of modules in each node of the U-Net.
+ batch_normalization: bool
+ Setting this to True will add a batch normalization layer after every convolution in the modules.
+ activation: str
+ Activation function to use in the modules.
+ See utils.choose_activation_layer for all supported activation functions.
+ output_activation: str
+ Output activation function.
+ padding: str
+ Padding to use in the convolution layers.
+ use_bias: bool
+ Setting this to True will implement a bias vector in the convolution layers used in the modules.
+ kernel_initializer: str or tf.keras.initializers object
+ Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_initializer: str or tf.keras.initializers object
+ Initializer for the bias vector in the Conv2D/Conv3D layers.
+ kernel_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers.
+ bias_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the bias vector in the Conv2D/Conv3D layers.
+ activity_regularizer: str or tf.keras.regularizers object
+ Regularizer function applied to the output of the Conv2D/Conv3D layers.
+ kernel_constraint: str or tf.keras.constraints object
+ Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers.
+ bias_constraint: str or tf.keras.constrains object
+ Constraint function applied to the bias vector in the Conv2D/Conv3D layers.
+
+ Returns
+ -------
+ model: tf.keras.models.Model object
+ U-Net model.
+
+ Raises
+ ------
+ ValueError
+ If levels < 2
+ If input_shape does not have 3 nor 4 dimensions
+ If the length of filter_num does not match the number of levels
+
+ References
+ ----------
+ https://arxiv.org/pdf/1804.03999.pdf
+ """
+
+ def build(self):
+ if len(self.upsample_size) > 0:
+ raise ValueError(
+ "AttentionUNet does not support upsample_size, use empty tuple."
+ )
+ self.supervision_kwargs["activation"] = self.output_activation
+ self.supervision_kwargs["upsample_size"] = self.pool_size
+ self.supervision_kwargs["use_bias"] = True
+ self.supervision_kwargs["output_level"] = 1
+ self.supervision_kwargs["kernel_size"] = 1
+
+ tensors = dict({}) # Tensors associated with each node and skip connections
+
+ """ Setup the first encoder node with an input layer and a convolution module """
+ tensors["input"] = Input(shape=self.input_shape, name="Input")
+ tensors["En1"] = convolution_module(
+ tensors["input"],
+ filters=self.filter_num[0],
+ name="En1",
+ **self.module_kwargs,
+ )
+
+ """ The rest of the encoder nodes are handled here. Each encoder node is connected with a MaxPooling layer and contains convolution modules """
+ for encoder in np.arange(
+ 2, self.levels + 1
+ ): # Iterate through the rest of the encoder nodes
+ pool_tensor = max_pool(
+ tensors[f"En{encoder - 1}"],
+ name=f"En{encoder - 1}-En{encoder}",
+ **self.pool_kwargs,
+ ) # Connect the next encoder node with a MaxPooling layer
+ tensors[f"En{encoder}"] = convolution_module(
+ pool_tensor,
+ filters=self.filter_num[encoder - 1],
+ name=f"En{encoder}",
+ **self.module_kwargs,
+ ) # Convolution modules
+
+ tensors[f"AG{self.levels - 1}"] = attention_gate(
+ tensors[f"En{self.levels - 1}"],
+ tensors[f"En{self.levels}"],
+ self.kernel_size,
+ self.pool_size,
+ name=f"AG{self.levels - 1}",
+ )
+ upsample_tensor = upsample(
+ tensors[f"En{self.levels}"],
+ filters=self.filter_num[self.levels - 2],
+ name=f"En{self.levels}-De{self.levels - 1}",
+ **self.upsample_kwargs,
+ ) # Connect the bottom encoder node to a decoder node
+
+ """ Bottom decoder node """
+ tensors[f"De{self.levels - 1}"] = Concatenate(
+ name=f"De{self.levels - 1}_Concatenate"
+ )(
+ [tensors[f"AG{self.levels - 1}"], upsample_tensor]
+ ) # Concatenate the upsampled tensor and skip connection
+ tensors[f"De{self.levels - 1}"] = convolution_module(
+ tensors[f"De{self.levels - 1}"],
+ filters=self.filter_num[self.levels - 2],
+ name=f"De{self.levels - 1}",
+ **self.module_kwargs,
+ ) # Convolution module
+ tensors[f"AG{self.levels - 2}"] = attention_gate(
+ tensors[f"En{self.levels - 2}"],
+ tensors[f"De{self.levels - 1}"],
+ self.kernel_size,
+ self.pool_size,
+ name=f"AG{self.levels - 2}",
+ )
+ upsample_tensor = upsample(
+ tensors[f"De{self.levels - 1}"],
+ filters=self.filter_num[self.levels - 3],
+ name=f"De{self.levels - 1}-De{self.levels - 2}",
+ **self.upsample_kwargs,
+ ) # Connect the bottom decoder node to the next decoder node
+
+ """ The rest of the decoder nodes (except the final node) are handled in this loop. Each node contains one concatenation of an upsampled tensor and a skip connection """
+ for decoder in np.arange(2, self.levels - 1)[::-1]:
+ tensors[f"De{decoder}"] = Concatenate(name=f"De{decoder}_Concatenate")(
+ [tensors[f"AG{decoder}"], upsample_tensor]
+ ) # Concatenate the upsampled tensor and skip connection
+ tensors[f"De{decoder}"] = convolution_module(
+ tensors[f"De{decoder}"],
+ filters=self.filter_num[decoder - 1],
+ name=f"De{decoder}",
+ **self.module_kwargs,
+ ) # Convolution module
+ tensors[f"AG{decoder - 1}"] = attention_gate(
+ tensors[f"En{decoder - 1}"],
+ tensors[f"De{decoder}"],
+ self.kernel_size,
+ self.pool_size,
+ name=f"AG{decoder - 1}",
+ )
+ upsample_tensor = upsample(
+ tensors[f"De{decoder}"],
+ filters=self.filter_num[decoder - 2],
+ name=f"De{decoder}-De{decoder - 1}",
+ **self.upsample_kwargs,
+ ) # Connect the bottom decoder node to the next decoder node
+
+ """ Final decoder node begins with a concatenation and convolution module, followed by deep supervision """
+ tensor_De1 = Concatenate(name="De1_Concatenate")(
+ [tensors["AG1"], upsample_tensor]
+ ) # Concatenate the upsampled tensor and skip connection
+ tensor_De1 = convolution_module(
+ tensor_De1, filters=self.filter_num[0], name="De1", **self.module_kwargs
+ ) # Convolution module
+ tensors["output"] = deep_supervision_side_output(
+ tensor_De1, name="final", **self.supervision_kwargs
+ ) # Deep supervision - this layer will output the model's prediction
+
+ output_model = Model(
+ inputs=tensors["input"],
+ outputs=tensors["output"],
+ name=f"attention_unet_{self.ndims}D",
+ )
+
+ return output_model
+
+
+@dataclasses.dataclass
+class UNetRegistry(keras_builders.BaseConfig):
+ """Registry class for UNet models.
+
+ Attributes:
+ name: the string name of the UNet model to build. Must be one of "unet",
+ "unet_ensemble", "unet_plus", "unet_2plus", "unet_3plus", or "attention_unet".
+ config: a dictionary of keyword arguments to pass to the UNet.
+ registry: a dictionary mapping string names to UNet functions.
+ """
+
+ name: Literal[
+ "unet",
+ "unet_ensemble",
+ "unet_plus",
+ "unet_2plus",
+ "unet_3plus",
+ "attention_unet",
+ ]
+
+ @property
+ def registry(self) -> dict[str, type]:
+ return {
+ "unet": UNet,
+ "unet_ensemble": UNetEnsemble,
+ "unet_plus": UNetPlus,
+ "unet_2plus": UNet2Plus,
+ "unet_3plus": UNet3Plus,
+ "attention_unet": AttentionUNet,
+ }
diff --git a/src/fronts/model.py b/src/fronts/model.py
new file mode 100644
index 0000000..9290f1a
--- /dev/null
+++ b/src/fronts/model.py
@@ -0,0 +1,204 @@
+import dataclasses
+from typing import Literal
+from fronts.layers.unets import UNetRegistry
+from fronts.utils.keras_builders import (
+ ConvOutputConfig,
+ BiasVectorConfig,
+ KernelMatrixConfig,
+ OptimizerConfig,
+ ActivationConfig,
+ LossConfig,
+ MetricConfig,
+)
+
+
+@dataclasses.dataclass
+class ModelConfig:
+ """Configuration for the model architecture and training parameters."""
+
+ name: Literal[
+ "unet",
+ "unet_ensemble",
+ "unet_plus",
+ "unet_2plus",
+ "unet_3plus",
+ "attention_unet",
+ ]
+ loss: LossConfig
+ metric: MetricConfig
+ optimizer: OptimizerConfig
+ convolution_activity_regularizer: ConvOutputConfig
+ bias_vector: BiasVectorConfig
+ kernel_matrix: KernelMatrixConfig
+ activation: ActivationConfig
+ batch_normalization: bool
+ num_filters: list[int]
+ kernel_size: list[int]
+ depth: int
+ modules_per_node: int
+ padding: Literal["same", "valid"]
+ pool_size: tuple[int]
+ upsample_size: tuple[int]
+ bias: bool
+
+ def build(self, input_shape: tuple, num_classes: int):
+ """Builds and compiles the UNet model based on the configuration.
+
+ Args:
+ input_shape: Shape of the model inputs (excluding batch dimension).
+ Spatial dims should be None to allow variable-size inference,
+ e.g. (None, None, 7, 9) for 3D inputs.
+ num_classes: Number of output classes (e.g. 6 for 5 front types + background).
+
+ Returns a compiled tf.keras.Model.
+ """
+ model = Model(
+ name=self.name,
+ loss_config=self.loss,
+ metric_config=self.metric,
+ optimizer_config=self.optimizer,
+ convolution_activity_regularizer_config=self.convolution_activity_regularizer,
+ bias_vector_config=self.bias_vector,
+ kernel_matrix_config=self.kernel_matrix,
+ activation_config=self.activation,
+ batch_normalization=self.batch_normalization,
+ num_filters=self.num_filters,
+ kernel_size=self.kernel_size,
+ depth=self.depth,
+ modules_per_node=self.modules_per_node,
+ padding=self.padding,
+ pool_size=self.pool_size,
+ upsample_size=self.upsample_size,
+ bias=self.bias,
+ input_shape=input_shape,
+ num_classes=num_classes,
+ )
+ return model.build()
+
+
+class Model:
+ def __init__(
+ self,
+ name: Literal[
+ "unet",
+ "unet_ensemble",
+ "unet_plus",
+ "unet_2plus",
+ "unet_3plus",
+ "attention_unet",
+ ],
+ loss_config: LossConfig,
+ metric_config: MetricConfig,
+ optimizer_config: OptimizerConfig,
+ convolution_activity_regularizer_config: ConvOutputConfig,
+ bias_vector_config: BiasVectorConfig,
+ kernel_matrix_config: KernelMatrixConfig,
+ activation_config: ActivationConfig,
+ batch_normalization: bool,
+ num_filters: list[int],
+ kernel_size: list[int],
+ depth: int,
+ modules_per_node: int,
+ padding: Literal["same", "valid"],
+ pool_size: tuple[int],
+ upsample_size: tuple[int],
+ bias: bool,
+ input_shape: tuple,
+ num_classes: int,
+ ):
+ self.name = name
+ self.loss_config = loss_config
+ self.metric_config = metric_config
+ self.optimizer_config = optimizer_config
+ self.convolution_activity_regularizer_config = (
+ convolution_activity_regularizer_config
+ )
+ self.bias_vector_config = bias_vector_config
+ self.kernel_matrix_config = kernel_matrix_config
+ self.activation_config = activation_config
+ self.batch_normalization = batch_normalization
+ self.num_filters = num_filters
+ self.kernel_size = kernel_size
+ self.depth = depth
+ self.modules_per_node = modules_per_node
+ self.padding = padding
+ self.pool_size = pool_size
+ self.upsample_size = upsample_size
+ self.bias = bias
+ self.input_shape = input_shape
+ self.num_classes = num_classes
+
+ if len(self.num_filters) != self.depth:
+ raise ValueError(
+ f"Length of num_filters ({len(self.num_filters)}) must match depth "
+ f"({self.depth})"
+ )
+ # Build keras objects
+ self.loss = self.loss_config.build()
+ self.metric = self.metric_config.build()
+ self.optimizer = self.optimizer_config.build()
+ self.activity_regularizer = self.convolution_activity_regularizer_config.build()
+ self.bias_vector = self.bias_vector_config.build()
+ self.kernel_matrix = self.kernel_matrix_config.build()
+ self.activation = self.activation_config.build()
+
+ def build(self):
+ """Builds and compiles the Keras model."""
+ # For 3D inputs (lat, lon, level, channels), squeeze the level axis (index 2)
+ # so the UNet output matches the 2D target shape (lat, lon, classes).
+ # For 2D inputs (lat, lon, channels), no squeezing is needed.
+ squeeze_axes = determine_squeeze_axes(self.input_shape)
+ shared_axes = determine_shared_axes(self.input_shape)
+
+ self.config = {
+ "input_shape": self.input_shape,
+ "num_classes": self.num_classes,
+ "pool_size": self.pool_size,
+ "upsample_size": self.upsample_size,
+ "levels": self.depth,
+ "filter_num": self.num_filters,
+ "kernel_size": self.kernel_size,
+ "squeeze_axes": squeeze_axes,
+ "shared_axes": shared_axes,
+ "modules_per_node": self.modules_per_node,
+ "batch_normalization": self.batch_normalization,
+ "activation": self.activation,
+ "padding": self.padding,
+ "use_bias": self.bias,
+ "kernel_initializer": self.kernel_matrix.kernel_initializer,
+ "bias_initializer": self.bias_vector.bias_initializer,
+ "kernel_regularizer": self.kernel_matrix.kernel_regularizer,
+ "bias_regularizer": self.bias_vector.bias_regularizer,
+ "activity_regularizer": self.activity_regularizer.activity_regularizer,
+ "kernel_constraint": self.kernel_matrix.kernel_constraint,
+ "bias_constraint": self.bias_vector.bias_constraint,
+ }
+ output_model = UNetRegistry(name=self.name, config=self.config).build()
+ output_model.compile(
+ loss=self.loss,
+ optimizer=self.optimizer,
+ metrics=[self.metric],
+ )
+ return output_model
+
+
+def determine_squeeze_axes(input_shape: tuple) -> int | None:
+ """Returns the axis to squeeze for 3D→2D output, or None for 2D inputs.
+
+ The UNet processes 3D inputs (lat, lon, level, channels) but produces
+ 2D targets (lat, lon, classes). The level axis (index 3, 1-based with
+ batch) is squeezed in the final output layer.
+
+ For 2D inputs (lat, lon, channels) no squeeze is needed.
+ """
+ # input_shape excludes batch dim, e.g. (None, None, 7, 9) = 3D, (None, None, 9) = 2D
+ ndims = len(input_shape) - 1 # spatial dims (exclude channel dim)
+ return 3 if ndims == 3 else None
+
+
+def determine_shared_axes(input_shape: tuple) -> int | None:
+ """Returns shared axes for learnable activation parameters, or None.
+
+ Following the legacy convention: None (share across all arbitrary dims).
+ """
+ return None
diff --git a/src/fronts/nfa/__init__.py b/src/fronts/nfa/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/fronts/nfa/methods.py b/src/fronts/nfa/methods.py
new file mode 100644
index 0000000..13595c3
--- /dev/null
+++ b/src/fronts/nfa/methods.py
@@ -0,0 +1,240 @@
+"""
+Numerical frontal analysis (NFA) methods.
+
+References
+----------
+* Renard and Clarke 1965: https://doi.org/10.1175/1520-0493(1965)093%3C0547:EINOFA%3E2.3.CO;2
+* Clarke and Renard 1966: https://doi.org/10.1175/1520-0450(1966)005%3C0764:TUSNNF%3E2.0.CO;2
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2023.10.6
+"""
+
+import numpy as np
+from fronts.utils import data_utils
+
+
+def thermal_front_parameter(field, lats, lons):
+ """
+ Calculates the thermal front parameter (TFP; Renard and Clarke 1965) using a provided thermodynamic variable.
+
+ Parameter
+ ---------
+ field: array-like of shape (..., M, N)
+ 2-D field of the thermodynamic variable that will be used to calculate the TFP. The last two axes of length M and
+ N must be the latitude and longitude axes, respectively.
+ lats: array-like of shape (M,)
+ 1-D vector of latitude points expressed as degrees.
+ lons: array-like of shape (N,)
+ 1-D vector of longitude points expressed as degrees.
+
+ Returns
+ -------
+ TFP: array-like of shape (..., M, N)
+ Thermal front parameter expressed as degC/(100km)^2.
+
+ Examples
+ --------
+ >>> np.random.seed(120)
+ >>> field = np.random.uniform(
+ ... low=0, high=40, size=(5, 5)
+ ... ) # random temperatures between 0 and 40 degrees Celsius
+ >>> field
+ array([[27.11822193, 20.51835209, 24.94822847, 19.0856986 , 18.41039256],
+ [38.03459464, 39.38302396, 34.17690184, 23.6436138 , 8.12785491],
+ [10.4944062 , 2.65660985, 25.87740026, 28.74931783, 14.04197021],
+ [38.10173911, 23.81909712, 39.78024821, 21.7469418 , 2.86850514],
+ [ 5.62746737, 8.29113292, 20.22109633, 21.4157172 , 21.25820336]])
+ >>> lons = np.arange(120, 201, 20)
+ >>> lons
+ array([120, 140, 160, 180, 200])
+ >>> lats = np.arange(40, 61, 5)
+ >>> lats
+ array([40, 45, 50, 55, 60])
+ >>> TFP = thermal_front_parameter(field, lats, lons)
+ >>> TFP
+ array([[-0.51556961, -0.57097894, 0.01174865, -0.06425258, -0.14175305],
+ [ 0.00375789, -0.48249109, 0.15134022, -0.0408684 , -0.17021388],
+ [-0.1660594 , 0.20510328, -0.19987511, -0.06345793, 0.2326307 ],
+ [-1.03967341, -0.40324022, -0.61269204, 0.11165267, 0.53391446],
+ [-0.02207704, 0.01454242, -0.00466246, -0.00478738, -0.0053461 ]])
+ """
+
+ # convert lats and lons to cartesian coordinates so spatial gradients can be calculated
+ Lons, Lats = np.meshgrid(lons, lats)
+ x, y = data_utils.haversine(Lons, Lats)
+
+ # gradient vector of the thermodynamic field
+ dFdx = np.diff(field, axis=-1, append=0) / np.diff(x, axis=-1, append=0)
+ dFdy = np.diff(field, axis=-2, append=0) / np.diff(y, axis=-2, append=0)
+ dF = np.array([dFdx, dFdy])
+
+ dFmag = np.sqrt(
+ np.sum(np.square(dF), axis=0)
+ ) # magnitude of the gradient vector of the thermodynamic field
+ dF_unit_vector = dF / dFmag # unit vector in the direction of the gradient vector
+
+ # gradient vector of the magnitude of the gradient vector of the thermodynamic field
+ ddFmagdx = np.diff(dFmag, axis=-1, append=0) / np.diff(x, axis=-1, append=0)
+ ddFmagdy = np.diff(dFmag, axis=-2, append=0) / np.diff(y, axis=-2, append=0)
+ ddFmag = np.array([ddFmagdx, ddFmagdy])
+
+ # calculate thermal front parameter and change units from degC/(km^2) to degC/(100km)^2
+ TFP = np.sum(-ddFmag * dF_unit_vector, axis=0) * 1e4
+
+ return TFP
+
+
+def thermal_front_locator(field, lats, lons):
+ """
+ Calculates the thermal front locator (TFL; Huber-Pock & Kress 1981) using a provided thermodynamic variable.
+
+ Parameter
+ ---------
+ field: array-like of shape (..., M, N)
+ 2-D field of the thermodynamic variable that will be used to calculate the TFL. The last two axes of length M and
+ N must be the latitude and longitude axes, respectively.
+ lats: array-like of shape (M,)
+ 1-D vector of latitude points expressed as degrees.
+ lons: array-like of shape (N,)
+ 1-D vector of longitude points expressed as degrees.
+
+ Returns
+ -------
+ TFL: array-like of shape (..., M, N)
+ Thermal front locator expressed as degC/(100km)^4.
+
+ Examples
+ --------
+ >>> np.random.seed(120)
+ >>> field = np.random.uniform(
+ ... low=0, high=40, size=(5, 5)
+ ... ) # random temperatures between 0 and 40 degrees Celsius
+ >>> field
+ array([[27.11822193, 20.51835209, 24.94822847, 19.0856986 , 18.41039256],
+ [38.03459464, 39.38302396, 34.17690184, 23.6436138 , 8.12785491],
+ [10.4944062 , 2.65660985, 25.87740026, 28.74931783, 14.04197021],
+ [38.10173911, 23.81909712, 39.78024821, 21.7469418 , 2.86850514],
+ [ 5.62746737, 8.29113292, 20.22109633, 21.4157172 , 21.25820336]])
+ >>> lons = np.arange(120, 201, 20)
+ >>> lons
+ array([120, 140, 160, 180, 200])
+ >>> lats = np.arange(40, 61, 5)
+ >>> lats
+ array([40, 45, 50, 55, 60])
+ >>> TFL = thermal_front_locator(field, lats, lons)
+ >>> TFL
+ array([[ 9.26191422e-02, 1.76187373e-02, 2.53556619e-02,
+ 4.34760673e-03, 5.07964100e-03],
+ [ 3.02196913e-02, -1.24731530e-01, 6.27675627e-02,
+ 8.60028595e-04, 7.23569573e-02],
+ [-1.58047350e-01, -1.10518190e-01, -7.37258066e-02,
+ -3.46076518e-02, -5.41030478e-02],
+ [-1.85557854e-01, -7.51105984e-02, -1.15051591e-01,
+ -2.00581092e-02, -9.69546769e-02],
+ [ 1.45060450e-03, -9.34552106e-04, -6.97467095e-05,
+ -7.09751637e-05, -8.47798408e-05]])
+ """
+
+ # convert lats and lons to cartesian coordinates so spatial gradients can be calculated
+ Lons, Lats = np.meshgrid(lons, lats)
+ x, y = data_utils.haversine(Lons, Lats)
+
+ # gradient vector of the thermodynamic field
+ dFdx = np.diff(field, axis=-1, append=0) / np.diff(x, axis=-1, append=0)
+ dFdy = np.diff(field, axis=-2, append=0) / np.diff(y, axis=-2, append=0)
+ dF = np.array([dFdx, dFdy])
+
+ dFmag = np.sqrt(
+ np.sum(np.square(dF), axis=0)
+ ) # magnitude of the gradient vector of the thermodynamic field
+ dF_unit_vector = dF / dFmag # unit vector in the direction of the gradient vector
+
+ # thermal front parameter expressed as degC/(100km)^2
+ TFP = thermal_front_parameter(field, lats, lons)
+
+ # gradient vector of the thermal front parameter
+ dTFPx = np.diff(TFP, axis=-1, append=0) / np.diff(x, axis=-1, append=0)
+ dTFPy = np.diff(TFP, axis=-2, append=0) / np.diff(y, axis=-2, append=0)
+ dTFP = np.array([dTFPx, dTFPy]) * 100 # units: degC/(100km)^3
+
+ # calculate thermal front parameter expressed as degC/(100km)^4
+ TFL = np.sum(dTFP * dF_unit_vector, axis=0)
+
+ return TFL
+
+
+def minimum_maximum_locator(field, lats, lons):
+ """
+ Calculates the Minimum-Maximum locator (MML; Clarke and Renard 1966).
+
+ Parameter
+ ---------
+ field: array-like of shape (..., M, N)
+ 2-D field of the thermodynamic variable that will be used to calculate the MML. The last two axes of length M and
+ N must be the latitude and longitude axes, respectively.
+ lats: array-like of shape (M,)
+ 1-D vector of latitude points expressed as degrees.
+ lons: array-like of shape (N,)
+ 1-D vector of longitude points expressed as degrees.
+
+ Returns
+ -------
+ MML: array-like of shape (..., M, N)
+ Minimum-Maximum locator expressed as degC/(100km)^4.
+
+ Examples
+ --------
+ >>> np.random.seed(120)
+ >>> field = np.random.uniform(
+ ... low=0, high=40, size=(5, 5)
+ ... ) # random temperatures between 0 and 40 degrees Celsius
+ >>> field
+ array([[27.11822193, 20.51835209, 24.94822847, 19.0856986 , 18.41039256],
+ [38.03459464, 39.38302396, 34.17690184, 23.6436138 , 8.12785491],
+ [10.4944062 , 2.65660985, 25.87740026, 28.74931783, 14.04197021],
+ [38.10173911, 23.81909712, 39.78024821, 21.7469418 , 2.86850514],
+ [ 5.62746737, 8.29113292, 20.22109633, 21.4157172 , 21.25820336]])
+ >>> lons = np.arange(120, 201, 20)
+ >>> lons
+ array([120, 140, 160, 180, 200])
+ >>> lats = np.arange(40, 61, 5)
+ >>> lats
+ array([40, 45, 50, 55, 60])
+ >>> MML = minimum_maximum_locator(field, lats, lons)
+ >>> MML
+ array([[ 1.97110331, 1.86558045, 1.68252513, 0.63611081, 1.82137088],
+ [ 3.87441325, -6.46895958, 1.55136154, 0.13647153, 1.06292906],
+ [-4.97573923, -3.95079091, -2.47669703, -1.44947791, -2.00722563],
+ [-5.87638048, -2.87788414, -3.62502599, -0.64223281, -3.30506401],
+ [ 0.12176053, -0.57844794, -0.30743987, -0.29444804, -0.33711797]])
+ """
+
+ # convert lats and lons to cartesian coordinates so spatial gradients can be calculated
+ Lons, Lats = np.meshgrid(lons, lats)
+ x, y = data_utils.haversine(Lons, Lats)
+
+ # gradient vector of the thermodynamic field
+ dFdx = np.diff(field, axis=-1, append=0) / np.diff(x, axis=-1, append=0)
+ dFdy = np.diff(field, axis=-2, append=0) / np.diff(y, axis=-2, append=0)
+ dF = np.array([dFdx, dFdy])
+
+ # thermal front parameter expressed as degC/(100km)^2
+ TFP = thermal_front_parameter(field, lats, lons)
+
+ # gradient vector of the thermal front parameter
+ dTFPx = np.diff(TFP, axis=-1, append=0) / np.diff(x, axis=-1, append=0)
+ dTFPy = np.diff(TFP, axis=-2, append=0) / np.diff(y, axis=-2, append=0)
+ dTFP = np.array([dTFPx, dTFPy]) * 100 # units: degC/(100km)^3
+
+ dTFPmag = np.sqrt(
+ np.sum(np.square(dTFP), axis=0)
+ ) # magnitude of the gradient vector of the TFP
+ dTFP_unit_vector = (
+ dTFP / dTFPmag
+ ) # unit vector in the direction of the gradient vector
+
+ # minimum-maximum locator expressed as degC/(100km)^4
+ MML = np.sum(dF * dTFP_unit_vector, axis=0) * 100
+
+ return MML
diff --git a/src/fronts/plot/__init__.py b/src/fronts/plot/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/fronts/plot/plot_era5.py b/src/fronts/plot/plot_era5.py
new file mode 100644
index 0000000..3ff487e
--- /dev/null
+++ b/src/fronts/plot/plot_era5.py
@@ -0,0 +1,71 @@
+"""
+Visualize ERA5 netCDF files.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import argparse
+import xarray as xr
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+import cartopy.crs as ccrs
+from fronts.utils import misc
+from fronts.utils.plotting import plot_background
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--netcdf_dir",
+ type=str,
+ required=True,
+ help="Base directory for the ERA5 netcdf files.",
+ )
+ parser.add_argument(
+ "--plot_outdir",
+ type=str,
+ help="Output directory for the generated plots. If no directory is declared, the plot will be shown with plt.show().",
+ )
+ parser.add_argument(
+ "--valid_time",
+ type=str,
+ required=True,
+ help="Valid time with format YYYY-MM-DD-HH.",
+ )
+ parser.add_argument("--variable", type=str, required=True, help="Variable to plot.")
+ parser.add_argument(
+ "--pressure_level", type=str, required=True, help="Pressure level of interest."
+ )
+ parser.add_argument(
+ "--plot_kwargs",
+ type=str,
+ help="Additional arguments to pass to plt.plot(). See utils.misc.string_arg_to_dict() for details.",
+ )
+ args = vars(parser.parse_args())
+
+ yr, mo, dy, hr = args["valid_time"].split("-")
+
+ nc_file = f"{args['netcdf_dir']}/{yr}{mo}/era5_{yr}{mo}{dy}{hr}_global.nc"
+ ds = xr.open_dataset(nc_file).sel(pressure_level=args["pressure_level"])[
+ args["variable"]
+ ]
+
+ fig, ax = plt.subplots(
+ figsize=(16, 8), dpi=500, subplot_kw={"projection": ccrs.PlateCarree()}
+ )
+ plot_background(ax=ax)
+
+ plot_kwargs = misc.string_arg_to_dict(args["plot_kwargs"])
+ ds.plot(
+ ax=ax, x="longitude", y="latitude", transform=ccrs.PlateCarree(), **plot_kwargs
+ )
+
+ if args["plot_outdir"] is not None:
+ mpl.use("Agg")
+ output_file = f"{args['plot_outdir']}/era5_{yr}{mo}{dy}{hr}_{args['variable']}-{args['pressure_level']}.png"
+ plt.tight_layout()
+ plt.savefig(output_file)
+ plt.close()
+ else:
+ plt.show()
diff --git a/src/fronts/plot/plot_fronts.py b/src/fronts/plot/plot_fronts.py
new file mode 100644
index 0000000..ad5aa2a
--- /dev/null
+++ b/src/fronts/plot/plot_fronts.py
@@ -0,0 +1,175 @@
+"""
+Plot fronts generated by NOAA or TWC.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+"""
+
+import argparse
+import xarray as xr
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+from matplotlib import cm
+import pandas as pd
+import cartopy.crs as ccrs
+from fronts.utils import data_utils
+from fronts.utils.plotting import plot_background
+import datetime as dt
+import numpy as np
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--netcdf_indir",
+ type=str,
+ required=True,
+ help="Input directory for the netCDF files containing fronts.",
+ )
+ parser.add_argument(
+ "--plot_outdir",
+ type=str,
+ required=True,
+ help="Output directory for the front plots.",
+ )
+ parser.add_argument(
+ "--front_types", type=str, nargs="+", help="Front types to plot."
+ )
+ parser.add_argument(
+ "--domain", type=str, default="full", help="Domain of the fronts."
+ )
+ parser.add_argument(
+ "--init_time",
+ type=str,
+ required=True,
+ help="Initialization time for the fronts. YYYY-MM-DD-HH",
+ )
+ parser.add_argument(
+ "--forecast_hour",
+ type=int,
+ help="Forecast hour for the fronts. Must be passed if plotting TWC fronts.",
+ )
+ args = vars(parser.parse_args())
+
+ # check for valid data source
+ domain = args["domain"].lower()
+ assert domain in ["full", "global", "hrrr", "nam-12km"], f"Invalid domain: {domain}"
+
+ # define initialization time and forecast hour
+ init_time = pd.to_datetime(args["init_time"])
+ yr, mo, dy, hr = init_time.year, init_time.month, init_time.day, init_time.hour
+ fhr = args["forecast_hour"]
+
+ # folder the given month that will be searched for the netCDF files
+ folder = f"{args['netcdf_indir']}/{yr:02d}{mo:02d}"
+
+ if domain in ["global", "hrrr", "nam-12km"]:
+ assert fhr is not None, (
+ "Forecast hour must be declared when plotting fronts globally or from an NWP model."
+ )
+ file = (
+ folder
+ + f"/FrontObjects_{yr:d}{mo:02d}{dy:02d}{hr:02d}_f{fhr:03d}_{domain}.nc"
+ )
+ valid_front_types = ["CF", "WF", "SF", "OF", "TROF", "DL"]
+ valid_time = init_time + dt.timedelta(hours=fhr)
+ valid_yr, valid_mo, valid_dy, valid_hr = (
+ valid_time.year,
+ valid_time.month,
+ valid_time.day,
+ valid_time.hour,
+ )
+ plot_title = (
+ f"TWC fronts initialized at {yr:d}-{mo:02d}-{dy:02d}-{hr:02d} valid for "
+ f"{valid_yr:d}-{valid_mo:02d}-{valid_dy:02d}-{valid_hr:02d} (forecast hour {fhr})"
+ )
+ plot_filename = f"fronts-twc_{yr:d}{mo:02d}{dy:02d}{hr:02d}_f{fhr:03d}.png"
+ else:
+ assert fhr is None, (
+ "Forecast hour cannot be declared when plotting NOAA fronts."
+ )
+ file = folder + f"/FrontObjects_{yr:d}{mo:02d}{dy:02d}{hr:02d}_full.nc"
+ valid_front_types = [
+ "CF",
+ "WF",
+ "SF",
+ "OF",
+ "CF-F",
+ "WF-F",
+ "SF-F",
+ "OF-F",
+ "CF-D",
+ "WF-D",
+ "SF-D",
+ "OF-D",
+ "INST",
+ "TROF",
+ "TT",
+ "DL",
+ ]
+ plot_title = f"NOAA/WPC/OPC fronts valid for {yr:d}-{mo:02d}-{dy:02d}-{hr:02d}"
+ plot_filename = f"fronts-noaa_{yr:d}{mo:02d}{dy:02d}{hr:02d}.png"
+
+ if args["front_types"] is None:
+ front_types = valid_front_types
+ else:
+ assert all(
+ [front_type in valid_front_types for front_type in args["front_types"]]
+ ), (
+ "One or more front types passed do not exist or are not valid for the specified data source."
+ )
+ front_types = args["front_types"]
+
+ fronts = xr.open_dataset(file, engine="netcdf4")
+ fronts = data_utils.reformat_fronts(fronts, front_types=front_types)
+ fronts = data_utils.expand_fronts(fronts, iterations=1)
+ fronts = xr.where(
+ fronts == 0, float("NaN"), fronts
+ ) # turn 0s into NaNs so they do not show up in the plot
+
+ fig, ax = plt.subplots(
+ figsize=(22, 8), subplot_kw={"projection": ccrs.PlateCarree()}
+ )
+ plot_background(extent=[-180, 179.99, -90, 90], ax=ax, linewidth=0.25)
+
+ front_colors_by_type = [
+ data_utils.FRONT_COLORS[front_type] for front_type in front_types
+ ]
+ front_names_by_type = [
+ data_utils.FRONT_NAMES[front_type] for front_type in front_types
+ ]
+
+ cmap_front = mcolors.ListedColormap(
+ front_colors_by_type, name="from_list", N=len(front_colors_by_type)
+ )
+ norm_front = mcolors.Normalize(vmin=1, vmax=len(front_colors_by_type) + 1)
+
+ fronts["identifier"].plot(
+ ax=ax,
+ x="longitude",
+ y="latitude",
+ cmap=cmap_front,
+ norm=norm_front,
+ transform=ccrs.PlateCarree(),
+ add_colorbar=False,
+ )
+
+ # redefining the colorbar with the order of the colors reversed so it looks more presentable
+ cmap_front = mcolors.ListedColormap(
+ front_colors_by_type[::-1], name="from_list", N=len(front_colors_by_type)
+ )
+
+ # when plotting the colorbar, flip the ticks and labels so they match up with the inverted colorbar
+ cbar_front = plt.colorbar(
+ cm.ScalarMappable(norm=norm_front, cmap=cmap_front), ax=ax, alpha=0.85, pad=0
+ )
+ cbar_front.set_ticks(np.arange(len(front_names_by_type)) + 1.5)
+ cbar_front.set_ticklabels(front_names_by_type)
+ cbar_front.set_ticks(cbar_front.get_ticks()[::-1])
+ cbar_front.set_ticklabels(cbar_front.ax.get_yticklabels()[::-1])
+ cbar_front.set_label(r"$\bf{Front}$ $\bf{type}$")
+
+ ax.set_title(plot_title)
+
+ plt.savefig(f"{args['plot_outdir']}/{plot_filename}", bbox_inches="tight", dpi=500)
+ plt.close()
diff --git a/src/fronts/plot/plot_performance_diagrams.py b/src/fronts/plot/plot_performance_diagrams.py
new file mode 100644
index 0000000..d377544
--- /dev/null
+++ b/src/fronts/plot/plot_performance_diagrams.py
@@ -0,0 +1,564 @@
+"""
+Plot performance diagrams for a model.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.2.2
+"""
+
+import argparse
+import cartopy.crs as ccrs
+from matplotlib import colors
+from matplotlib.font_manager import FontProperties
+import matplotlib.pyplot as plt
+from matplotlib.ticker import FixedLocator
+import numpy as np
+import pandas as pd
+import pickle
+import xarray as xr
+import random
+from fronts.utils import data_utils, plotting
+from glob import glob
+
+
+if __name__ == "__main__":
+ """
+ All arguments listed in the examples are listed via argparse in alphabetical order below this comment block.
+ """
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--dataset", type=str, help="'training', 'validation', or 'test'"
+ )
+ parser.add_argument(
+ "--data_source",
+ type=str,
+ default="era5",
+ help="Source of the variable data (ERA5, GDAS, etc.)",
+ )
+ parser.add_argument("--domain", type=str, required=True, help="Domain of the data.")
+ parser.add_argument(
+ "--map_neighborhood",
+ type=int,
+ default=250,
+ help="Neighborhood for the CSI map in kilometers. Options are: 50, 100, 150, 200, 250",
+ )
+ parser.add_argument(
+ "--model_dir", type=str, required=True, help="Directory for the models."
+ )
+ parser.add_argument("--model_number", type=int, required=True, help="Model number.")
+ parser.add_argument(
+ "--confidence_level",
+ type=float,
+ default=95,
+ help="Confidence level expressed as a percentage.",
+ )
+ parser.add_argument(
+ "--num_iterations",
+ type=int,
+ default=10000,
+ help="Number of iterations to perform when bootstrapping the statistics.",
+ )
+ parser.add_argument(
+ "--output_type", type=str, default="png", help="Output type for the image file."
+ )
+
+ args = vars(parser.parse_args())
+
+ model_properties_filepath = f"{args['model_dir']}/model_{args['model_number']}/model_{args['model_number']}_properties.pkl"
+ model_properties = pd.read_pickle(model_properties_filepath)
+
+ # Some older models do not have the 'dataset_properties' dictionary
+ try:
+ front_types = model_properties["dataset_properties"]["front_types"]
+ except KeyError:
+ front_types = model_properties["front_types"]
+
+ spatial_files = list(
+ sorted(
+ glob(
+ "%s/model_%d/statistics/model_%d_statistics_%s_*_spatial.nc"
+ % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ args["domain"],
+ )
+ )
+ )
+ )
+ temporal_files = list(
+ sorted(
+ glob(
+ "%s/model_%d/statistics/model_%d_statistics_%s_*_temporal.nc"
+ % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ args["domain"],
+ )
+ )
+ )
+ )
+
+ spatial_ds = xr.open_mfdataset(spatial_files, combine="nested", concat_dim="time")
+ temporal_ds = xr.open_mfdataset(temporal_files, combine="nested", concat_dim="time")
+
+ thresholds = temporal_ds["threshold"].values
+ time_array = temporal_ds["time"].values
+
+ if type(front_types) == str:
+ front_types = [
+ front_types,
+ ]
+ num_front_types = len(front_types)
+
+ for front_type in front_types:
+ tp_array = np.zeros([args["num_iterations"], 5, 100])
+ fp_array = np.zeros([args["num_iterations"], 5, 100])
+ tn_array = np.zeros([args["num_iterations"], 5, 100])
+ fn_array = np.zeros([args["num_iterations"], 5, 100])
+ # POD_array = np.zeros([args['num_iterations'], 5, 100]) # probability of detection = TP / (TP + FN)
+ # SR_array = np.zeros([args['num_iterations'], 5, 100]) # success ratio = 1 - False Alarm Ratio = TP / (TP + FP)
+
+ CI_lower_tp = np.zeros([5, 100])
+ CI_lower_fp = np.zeros([5, 100])
+ CI_lower_tn = np.zeros([5, 100])
+ CI_lower_fn = np.zeros([5, 100])
+ CI_upper_tp = np.zeros([5, 100])
+ CI_upper_fp = np.zeros([5, 100])
+ CI_upper_tn = np.zeros([5, 100])
+ CI_upper_fn = np.zeros([5, 100])
+
+ num_timesteps = len(time_array)
+ selectable_indices = range(num_timesteps)
+
+ true_positives_temporal = temporal_ds["tp_temporal_%s" % front_type].values
+ false_positives_temporal = temporal_ds["fp_temporal_%s" % front_type].values
+ true_negatives_temporal = temporal_ds["tn_temporal_%s" % front_type].values
+ false_negatives_temporal = temporal_ds["fn_temporal_%s" % front_type].values
+
+ for iteration in range(args["num_iterations"]):
+ print(f"Iteration {iteration}/{args['num_iterations']}", end="\r")
+ indices = random.choices(
+ selectable_indices, k=num_timesteps
+ ) # Select a sample equal to the total number of timesteps
+ tp_array[iteration, :, :] = np.sum(
+ true_positives_temporal[indices, :, :], axis=0
+ )
+ fp_array[iteration, :, :] = np.sum(
+ false_positives_temporal[indices, :, :], axis=0
+ )
+ tn_array[iteration, :, :] = np.sum(
+ true_negatives_temporal[indices, :, :], axis=0
+ )
+ fn_array[iteration, :, :] = np.sum(
+ false_negatives_temporal[indices, :, :], axis=0
+ )
+
+ print(f"Iteration {args['num_iterations']}/{args['num_iterations']}")
+
+ lower_bound = (100 - args["confidence_level"]) / 2
+ upper_bound = 100 - lower_bound
+
+ # Calculate confidence intervals at each probability bin
+ for percent in np.arange(0, 100):
+ # lower bounds for confidence intervals
+ CI_lower_tp[:, percent] = np.percentile(
+ tp_array[:, :, percent], q=lower_bound, axis=0
+ )
+ CI_lower_fp[:, percent] = np.percentile(
+ fp_array[:, :, percent], q=lower_bound, axis=0
+ )
+ CI_lower_tn[:, percent] = np.percentile(
+ tn_array[:, :, percent], q=lower_bound, axis=0
+ )
+ CI_lower_fn[:, percent] = np.percentile(
+ fn_array[:, :, percent], q=lower_bound, axis=0
+ )
+
+ # lower bound for confidence intervals
+ CI_upper_tp[:, percent] = np.percentile(
+ tp_array[:, :, percent], q=upper_bound, axis=0
+ )
+ CI_upper_fp[:, percent] = np.percentile(
+ fp_array[:, :, percent], q=upper_bound, axis=0
+ )
+ CI_upper_tn[:, percent] = np.percentile(
+ tn_array[:, :, percent], q=upper_bound, axis=0
+ )
+ CI_upper_fn[:, percent] = np.percentile(
+ fn_array[:, :, percent], q=upper_bound, axis=0
+ )
+
+ CI_lower_POD = CI_lower_tp / (CI_lower_tp + CI_upper_fn)
+ CI_upper_POD = CI_upper_tp / (CI_upper_tp + CI_lower_fn)
+ CI_lower_SR = CI_lower_tp / (CI_lower_tp + CI_upper_fp)
+ CI_upper_SR = CI_upper_tp / (CI_upper_tp + CI_lower_fp)
+
+ CI_lower_HSS = (
+ 2
+ * ((CI_lower_tp * CI_lower_tn) - (CI_upper_fp * CI_upper_fn))
+ / (
+ ((CI_lower_tp + CI_upper_fn) * (CI_upper_fn + CI_lower_tn))
+ + ((CI_lower_tp + CI_upper_fp) * (CI_upper_fp + CI_lower_tn))
+ )
+ )
+ CI_upper_HSS = (
+ 2
+ * ((CI_upper_tp * CI_upper_tn) - (CI_lower_fp * CI_lower_fn))
+ / (
+ ((CI_upper_tp + CI_lower_fn) * (CI_lower_fn + CI_upper_tn))
+ + ((CI_upper_tp + CI_lower_fp) * (CI_lower_fp + CI_upper_tn))
+ )
+ )
+
+ CI_POD = np.stack((CI_lower_POD, CI_upper_POD), axis=0)
+ CI_SR = np.stack((CI_lower_SR, CI_upper_SR), axis=0)
+ CI_CSI = np.stack((CI_SR**-1 + CI_POD**-1 - 1.0) ** -1, axis=0)
+ CI_HSS = np.stack((CI_lower_HSS, CI_upper_HSS), axis=0)
+ CI_FB = np.stack(CI_POD * (CI_SR**-1), axis=0)
+
+ # Remove the zeros
+ try:
+ polygon_stop_index = np.min(np.where(CI_POD == 0)[2])
+ except IndexError:
+ polygon_stop_index = 100
+
+ ### Statistics with shape (boundary, threshold) after taking the sum along the time axis (axis=0) ###
+ true_positives_temporal_sum = np.sum(true_positives_temporal, axis=0)
+ false_positives_temporal_sum = np.sum(false_positives_temporal, axis=0)
+ true_negatives_temporal_sum = np.sum(true_negatives_temporal, axis=0)
+ false_negatives_temporal_sum = np.sum(false_negatives_temporal, axis=0)
+
+ a = true_positives_temporal_sum
+ b = false_positives_temporal_sum
+ c = false_negatives_temporal_sum
+ d = true_negatives_temporal_sum
+
+ spatial_csi_ds = (
+ spatial_ds[f"tp_spatial_{front_type}"].sum("time")
+ / (
+ spatial_ds[f"tp_spatial_{front_type}"].sum("time")
+ + spatial_ds[f"fp_spatial_{front_type}"].sum("time")
+ + spatial_ds[f"fn_spatial_{front_type}"].sum("time")
+ )
+ ).max("threshold")
+
+ num_forecasts = true_positives_temporal_sum + false_positives_temporal_sum
+ num_forecasts = num_forecasts[0, :]
+ total_pixels = (
+ true_positives_temporal.shape[0]
+ * len(spatial_csi_ds["latitude"])
+ * len(spatial_csi_ds["longitude"])
+ )
+ relative_forecast_fraction = 100 * num_forecasts / total_pixels
+
+ ### Find the number of true positives and false positives in each probability bin ###
+ true_positives_diff = np.abs(np.diff(true_positives_temporal_sum))
+ false_positives_diff = np.abs(np.diff(false_positives_temporal_sum))
+ observed_relative_frequency = np.divide(
+ true_positives_diff, true_positives_diff + false_positives_diff
+ )
+
+ pod = np.divide(
+ true_positives_temporal_sum,
+ true_positives_temporal_sum + false_negatives_temporal_sum,
+ ) # Probability of detection
+ sr = np.divide(
+ true_positives_temporal_sum,
+ true_positives_temporal_sum + false_positives_temporal_sum,
+ ) # Success ratio
+
+ fig, axs = plt.subplots(1, 2, figsize=(15, 6))
+ axarr = axs.flatten()
+
+ sr_matrix, pod_matrix = np.meshgrid(
+ np.linspace(0, 1, 101), np.linspace(0, 1, 101)
+ )
+ csi_matrix = 1 / ((1 / sr_matrix) + (1 / pod_matrix) - 1) # CSI coordinates
+ fb_matrix = pod_matrix * (sr_matrix**-1) # Frequency Bias coordinates
+ CSI_LEVELS = np.linspace(0, 1, 11) # CSI contour levels
+ FB_LEVELS = [0.25, 0.5, 0.75, 1, 1.25, 1.5, 2, 3] # Frequency Bias levels
+ cmap = "Blues" # Colormap for the CSI contours
+ axis_ticks = np.arange(0, 1.01, 0.1)
+ axis_ticklabels = np.arange(0, 100.1, 10).astype(int)
+
+ cs = axarr[0].contour(
+ sr_matrix,
+ pod_matrix,
+ fb_matrix,
+ FB_LEVELS,
+ colors="black",
+ linewidths=0.5,
+ linestyles="--",
+ ) # Plot FB levels
+ axarr[0].clabel(cs, FB_LEVELS, fontsize=8)
+
+ csi_contour = axarr[0].contourf(
+ sr_matrix, pod_matrix, csi_matrix, CSI_LEVELS, cmap=cmap
+ ) # Plot CSI contours in 0.1 increments
+ cbar = fig.colorbar(
+ csi_contour, ax=axarr[0], pad=0.02, label="Critical Success Index (CSI)"
+ )
+ cbar.set_ticks(axis_ticks)
+
+ axarr1_2 = axarr[1].twinx()
+ axarr1_2.set_ylabel("Percentage of Grid Points with Forecasts [bars]")
+ axarr1_2.yaxis.set_major_locator(plt.LinearLocator(11))
+ axarr1_2.bar(
+ thresholds[:-1],
+ relative_forecast_fraction[1:],
+ color="blue",
+ width=0.005,
+ alpha=0.25,
+ )
+ axarr[1].plot(
+ thresholds,
+ thresholds,
+ color="black",
+ linestyle="--",
+ linewidth=0.5,
+ label="Perfect Reliability",
+ )
+
+ cell_text = [] # List of strings that will be used in the table near the bottom of this function
+
+ ### CSI and reliability lines for each boundary ###
+ boundary_colors = ["red", "purple", "brown", "darkorange", "darkgreen"]
+ max_CSI_scores_by_boundary = np.zeros(shape=(5,))
+ max_HSS_scores_by_boundary = np.zeros(shape=(5,))
+ for boundary, color in enumerate(boundary_colors):
+ csi = np.power((1 / sr[boundary]) + (1 / pod[boundary]) - 1, -1)
+ max_CSI_scores_by_boundary[boundary] = np.nanmax(csi)
+ max_CSI_index = np.where(csi == max_CSI_scores_by_boundary[boundary])[0]
+ max_CSI_threshold = thresholds[max_CSI_index][
+ 0
+ ] # Probability threshold where CSI is maximized
+ max_HSS_scores_by_boundary = (
+ 2 * ((a * d) - (b * c)) / (((a + c) * (c + d)) + ((a + b) * (b + d)))
+ )
+ max_CSI_pod = pod[boundary][max_CSI_index][0] # POD where CSI is maximized
+ max_CSI_sr = sr[boundary][max_CSI_index][0] # SR where CSI is maximized
+ max_CSI_fb = max_CSI_pod / max_CSI_sr # Frequency bias
+
+ cell_text.append(
+ [
+ r"$\bf{%.3f}$" % max_CSI_scores_by_boundary[boundary]
+ + r"$^{%.3f}_{%.3f}$"
+ % (
+ CI_CSI[1, boundary, max_CSI_index][0],
+ CI_CSI[0, boundary, max_CSI_index][0],
+ ),
+ r"$\bf{%.3f}$" % max_HSS_scores_by_boundary[boundary, max_CSI_index]
+ + r"$^{%.3f}_{%.3f}$"
+ % (
+ CI_HSS[1, boundary, max_CSI_index][0],
+ CI_HSS[0, boundary, max_CSI_index][0],
+ ),
+ r"$\bf{%.1f}$" % (max_CSI_pod * 100)
+ + r"$^{%.1f}_{%.1f}$"
+ % (
+ CI_POD[1, boundary, max_CSI_index][0] * 100,
+ CI_POD[0, boundary, max_CSI_index][0] * 100,
+ ),
+ r"$\bf{%.1f}$" % ((1 - max_CSI_sr) * 100)
+ + r"$^{%.1f}_{%.1f}$"
+ % (
+ (1 - CI_SR[1, boundary, max_CSI_index][0]) * 100,
+ (1 - CI_SR[0, boundary, max_CSI_index][0]) * 100,
+ ),
+ r"$\bf{%.3f}$" % max_CSI_fb
+ + r"$^{%.3f}_{%.3f}$"
+ % (
+ CI_FB[1, boundary, max_CSI_index][0],
+ CI_FB[0, boundary, max_CSI_index][0],
+ ),
+ ]
+ )
+
+ # Plot CSI lines
+ axarr[0].plot(
+ max_CSI_sr, max_CSI_pod, color=color, marker="*", markersize=10
+ )
+ axarr[0].plot(sr[boundary], pod[boundary], color=color, linewidth=1)
+
+ # Plot reliability curve
+ axarr[1].plot(
+ thresholds[:-1],
+ observed_relative_frequency[boundary],
+ color=color,
+ linewidth=1,
+ )
+
+ # Confidence interval
+ xs = np.concatenate(
+ [
+ CI_SR[0, boundary, :polygon_stop_index],
+ CI_SR[1, boundary, :polygon_stop_index][::-1],
+ ]
+ )
+ ys = np.concatenate(
+ [
+ CI_POD[0, boundary, :polygon_stop_index],
+ CI_POD[1, boundary, :polygon_stop_index][::-1],
+ ]
+ )
+ axarr[0].fill(
+ xs, ys, alpha=0.3, color=color
+ ) # Shade the confidence interval
+
+ axarr[0].set_xticklabels(
+ axis_ticklabels[::-1]
+ ) # False alarm rate on x-axis means values are reversed
+ axarr[0].set_xlabel("False Alarm Rate (FAR; %)")
+ axarr[0].set_ylabel("Probability of Detection (POD; %)")
+ axarr[0].set_title(
+ r"$\bf{a)}$ $\bf{CSI}$ $\bf{diagram}$ [confidence level = %d%%]"
+ % args["confidence_level"]
+ )
+
+ axarr[1].set_xticklabels(axis_ticklabels)
+ axarr[1].set_xlabel("Forecast Probability (uncalibrated; %)")
+ axarr[1].set_ylabel("Observed Relative Frequency (%) [lines]")
+ axarr[1].set_title(r"$\bf{b)}$ $\bf{Reliability}$ $\bf{diagram}$")
+
+ for ax in axarr:
+ ax.set_xticks(axis_ticks)
+ ax.set_yticks(axis_ticks)
+ ax.set_yticklabels(axis_ticklabels)
+ ax.grid(color="black", alpha=0.1)
+ ax.set_xlim(0, 1)
+ ax.set_ylim(0, 1)
+ ################################################################################################################
+
+ cbar_kwargs = {
+ "label": "CSI",
+ "pad": 0,
+ } # Spatial CSI colorbar keyword arguments
+
+ ### Adjust the data table and spatial CSI plot based on the domain ###
+ if args["domain"] == "conus":
+ table_axis_extent = [0.063, -0.038, 0.39, 0.239]
+ table_scale = (1, 3.3)
+ table_title_kwargs = dict(x=0.5, y=0.098, pad=-4)
+ spatial_axis_extent = [0.5, -0.582, 0.512, 0.544]
+ cbar_kwargs["shrink"] = 1
+ spatial_plot_xlabels = [-140, -105, -70]
+ spatial_plot_ylabels = [30, 40, 50]
+ elif args["domain"] == "full":
+ table_axis_extent = [0.063, -0.038, 0.39, 0.229]
+ table_scale = (1, 2.8)
+ table_title_kwargs = dict(x=0.5, y=0.096, pad=-4)
+ spatial_axis_extent = [0.523, -0.5915, 0.48, 0.66]
+ cbar_kwargs["shrink"] = 0.675
+ spatial_plot_xlabels = [-150, -120, -90, -60, -30, 0, 120, 150, 180]
+ spatial_plot_ylabels = [0, 20, 40, 60, 80]
+ else:
+ raise ValueError(
+ "%s domain is currently not supported for performance diagrams."
+ % args["domain"]
+ )
+
+ ############################################# Data table (panel c) #############################################
+ columns = ["CSI", "HSS", "POD %", "FAR %", "FB"] # Column names
+ rows = ["50 km", "100 km", "150 km", "200 km", "250 km"] # Row names
+
+ table_axis = plt.axes(table_axis_extent)
+ table_axis.set_title(
+ r"$\bf{c)}$ $\bf{Data}$ $\bf{table}$ [confidence level = %d%%]"
+ % args["confidence_level"],
+ **table_title_kwargs,
+ )
+ table_axis.axis("off")
+ stats_table = table_axis.table(
+ cellText=cell_text,
+ rowLabels=rows,
+ rowColours=boundary_colors,
+ colLabels=columns,
+ cellLoc="center",
+ )
+ stats_table.scale(*table_scale) # Make the table larger
+
+ ### Shade the cells and make the cell text larger ###
+ for cell in stats_table._cells:
+ stats_table._cells[cell].set_alpha(0.7)
+ stats_table._cells[cell].set_text_props(
+ fontproperties=FontProperties(size="x-large", stretch="expanded")
+ )
+ ################################################################################################################
+
+ ########################################## Spatial CSI map (panel d) ###########################################
+ right_labels = False # Disable latitude labels on the right side of the subplot
+ top_labels = False # Disable longitude labels on top of the subplot
+ left_labels = True # Latitude labels on the left side of the subplot
+ bottom_labels = True # Longitude labels on the bottom of the subplot
+
+ ## Set up the spatial CSI plot ###
+ csi_cmap = plotting.truncated_colormap("gnuplot2", maxval=0.9, n=10)
+ extent = data_utils.DOMAIN_EXTENTS[args["domain"]]
+ spatial_axis = plt.axes(
+ spatial_axis_extent, projection=ccrs.Miller(central_longitude=250)
+ )
+ spatial_axis_title_text = (
+ r"$\bf{d)}$ $\bf{%d}$ $\bf{km}$ $\bf{CSI}$ $\bf{map}$"
+ % args["map_neighborhood"]
+ )
+ plotting.plot_background(extent=extent, ax=spatial_axis)
+ norm_probs = colors.Normalize(vmin=0.1, vmax=1)
+ spatial_csi_ds = xr.where(spatial_csi_ds >= 0.1, spatial_csi_ds, float("NaN"))
+ spatial_csi_ds.sel(neighborhood=args["map_neighborhood"]).plot(
+ ax=spatial_axis,
+ x="longitude",
+ y="latitude",
+ norm=norm_probs,
+ cmap=csi_cmap,
+ transform=ccrs.PlateCarree(),
+ alpha=0.6,
+ cbar_kwargs=cbar_kwargs,
+ )
+ spatial_axis.set_title(spatial_axis_title_text)
+ gl = spatial_axis.gridlines(
+ draw_labels=True, zorder=0, dms=True, x_inline=False, y_inline=False
+ )
+ gl.right_labels = right_labels
+ gl.top_labels = top_labels
+ gl.left_labels = left_labels
+ gl.bottom_labels = bottom_labels
+ gl.xlocator = FixedLocator(spatial_plot_xlabels)
+ gl.ylocator = FixedLocator(spatial_plot_ylabels)
+ gl.xlabel_style = {"size": 7}
+ gl.ylabel_style = {"size": 8}
+ ################################################################################################################
+
+ if args["domain"] == "conus":
+ domain_text = args["domain"].upper()
+ else:
+ domain_text = args["domain"]
+
+ plt.suptitle(
+ f"Five-class model: %ss over %s domain"
+ % (data_utils.FRONT_NAMES[front_type], domain_text),
+ fontsize=20,
+ ) # Create and plot the main title
+
+ filename = (
+ f"%s/model_%d/performance_%s_%s_%s_{args['data_source']}.{args['output_type']}"
+ % (
+ args["model_dir"],
+ args["model_number"],
+ front_type,
+ args["dataset"],
+ args["domain"],
+ )
+ )
+ if args["data_source"] != "era5":
+ filename = filename.replace(
+ f".{args['output_type']}",
+ f"_f%03d.{args['output_type']}" % args["forecast_hour"],
+ ) # Add forecast hour to the end of the filename
+
+ plt.tight_layout()
+ plt.savefig(filename, bbox_inches="tight", dpi=500)
+ plt.close()
+
+ with open(model_properties_filepath, "wb") as f:
+ pickle.dump(model_properties, f)
diff --git a/src/fronts/plot/plot_permutations.py b/src/fronts/plot/plot_permutations.py
new file mode 100644
index 0000000..3e88303
--- /dev/null
+++ b/src/fronts/plot/plot_permutations.py
@@ -0,0 +1,235 @@
+"""
+Generate plots for permutation data from a model.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2024.10.10
+"""
+
+import argparse
+import matplotlib as mpl
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+from fronts.utils import constants
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--model_dir", type=str, required=True, help="Directory for the models."
+ )
+ parser.add_argument("--model_number", type=int, required=True, help="Model number.")
+ parser.add_argument(
+ "--domain", type=str, required=True, help="Domain for the permutations"
+ )
+ parser.add_argument(
+ "--show_names",
+ action="store_true",
+ help="Show variable and level names instead of prefixes.",
+ )
+ parser.add_argument(
+ "--output_type", type=str, default="png", help="Output type for the image file."
+ )
+
+ args = vars(parser.parse_args())
+
+ model_folder = "%s/model_%d" % (args["model_dir"], args["model_number"])
+
+ model_properties = pd.read_pickle(
+ "%s/model_%d_properties.pkl" % (model_folder, args["model_number"])
+ )
+ variables, pressure_levels = (
+ model_properties["dataset_properties"]["variables"],
+ model_properties["dataset_properties"]["pressure_levels"],
+ )
+ front_types = model_properties["dataset_properties"]["front_types"]
+ permutations_dict = pd.read_pickle(
+ "%s/permutations_%d_%s.pkl"
+ % (model_folder, args["model_number"], args["domain"])
+ ) # contains permutation data
+
+ num_vars = len(variables) # number of variables
+ num_lvls = len(pressure_levels) # number of pressure levels
+
+ sp = permutations_dict["single_pass"]
+ sp_vars = np.array(
+ [sp[var] for var in variables]
+ ) # single-pass results: shuffled variables over all levels
+ sp_lvls = np.array(
+ [sp[lvl] for lvl in pressure_levels]
+ ) # single-pass results: shuffled levels over all variables
+ sp_sorting_orders = [
+ np.argsort(sp_vars, axis=0),
+ np.argsort(sp_lvls, axis=0),
+ ] # list of length 2: sorted grouped variables and grouped levels
+
+ for front_num, front_type in enumerate(front_types):
+ sp_sorted_data = [
+ sp_vars[:, front_num][sp_sorting_orders[0][:, front_num]],
+ sp_lvls[:, front_num][sp_sorting_orders[1][:, front_num]],
+ ]
+ sp_sorted_vars_and_lvls = [
+ [variables[ind] for ind in sp_sorting_orders[0][:, front_num]],
+ [pressure_levels[ind] for ind in sp_sorting_orders[1][:, front_num]],
+ ]
+
+ y_pos_vars = np.arange(
+ num_vars
+ ) # array marking the y-positions of the variables
+ y_pos_lvls = np.arange(
+ num_lvls
+ ) # array marking the y-positions of the pressure levels
+
+ fig = plt.figure(figsize=(12, 9))
+ ax1 = plt.subplot(
+ 2, 2, 1
+ ) # axis for the subplot containing single-pass variable permutations
+ ax2 = plt.subplot(
+ 2, 2, 2
+ ) # axis for the subplot containing single-pass level permutations
+ ax3 = plt.subplot(2, 1, 2) # table axis
+
+ # Horizontal bars for the single-pass data
+ sp_barh_vars = ax1.barh(
+ y_pos_vars, sp_sorted_data[0], color=constants.FRONT_COLORS[front_type]
+ )
+ sp_barh_lvls = ax2.barh(
+ y_pos_lvls, sp_sorted_data[1], color=constants.FRONT_COLORS[front_type]
+ )
+
+ # setting subplot titles, ticks, and ticklabels
+ ax1.set_title(f"a) Grouped variables")
+ ax1.set_yticks(y_pos_vars)
+ ax2.set_title(f"b) Grouped levels")
+ ax2.set_yticks(y_pos_lvls)
+
+ if args["show_names"]:
+ ax1.set_yticklabels(
+ [constants.VARIABLE_NAMES[var] for var in sp_sorted_vars_and_lvls[0]]
+ )
+ ax2.set_yticklabels(
+ [constants.VERTICAL_LEVELS[lvl] for lvl in sp_sorted_vars_and_lvls[1]]
+ )
+ else:
+ ax1.set_yticklabels(sp_sorted_vars_and_lvls[0])
+ ax2.set_yticklabels(sp_sorted_vars_and_lvls[1])
+
+ x_margin_adjust = 0.14 # factor for increasing the scale of the x-axis so text will not extend past the right side of the subplots
+
+ # Add labels to the horizontal bars on the subplots
+ for pos, data in enumerate(sp_sorted_data[0]):
+ ax1.annotate(
+ data,
+ xy=(
+ np.max([data, 0])
+ + (np.max(sp_sorted_data[0]) * x_margin_adjust / 10),
+ pos,
+ ),
+ va="center",
+ )
+ for pos, data in enumerate(sp_sorted_data[1]):
+ ax2.annotate(
+ data,
+ xy=(
+ np.max([data, 0])
+ + (np.max(sp_sorted_data[1]) * x_margin_adjust / 10),
+ pos,
+ ),
+ va="center",
+ )
+
+ for ax in [ax1, ax2]:
+ ax.margins(
+ x=x_margin_adjust
+ ) # increase the scale of the x-axis so text will not extend past the right side of the subplots
+ ax.set_xlabel("Relative importance")
+ ax.grid(alpha=0.3, axis="x")
+
+ # importance values for the cells in the single-pass table
+ cellValues = np.array(
+ [
+ [sp["_".join([var, lvl])][front_num] for lvl in pressure_levels]
+ for var in variables
+ ]
+ )
+ sorted_indices = np.argsort(-cellValues.flatten())
+
+ # variables at single levels ranked by importance
+ ranks = np.zeros_like(sorted_indices)
+ ranks[sorted_indices] = np.arange(len(sorted_indices)) + 1
+ cellText = np.array(
+ [
+ "%d (%.3f)" % (rank, val)
+ for rank, val in zip(ranks, cellValues.flatten())
+ ]
+ ).reshape(cellValues.shape)
+
+ # shade table cells based on the importance values
+ max_val = np.max(cellValues)
+ cellColorValues = 0.5 + (cellValues / (2 * max_val))
+ cellColorValues /= np.max(cellColorValues)
+ cmap = mpl.colormaps.get_cmap("bwr_r")
+ cellColours = [cmap(val) for val in cellColorValues]
+
+ rowColours = [
+ "gray" for _ in range(len(variables))
+ ] # shade first column containing variable names
+ colColours = [
+ "gray" for _ in range(len(pressure_levels))
+ ] # shade header cells containing pressure level names
+ ax3.set_title(f"c) Variables on single levels")
+ ax3.axis("off")
+
+ rowLabels = (
+ [constants.VARIABLE_NAMES[var] for var in variables]
+ if args["show_names"]
+ else variables
+ )
+ colLabels = (
+ [constants.VERTICAL_LEVELS[lvl] for lvl in pressure_levels]
+ if args["show_names"]
+ else pressure_levels
+ )
+
+ stats_table = ax3.table(
+ cellText=cellText,
+ rowLabels=rowLabels,
+ colLabels=colLabels,
+ rowColours=rowColours,
+ colColours=colColours,
+ cellColours=cellColours,
+ cellLoc="center",
+ bbox=[0, 0, 1, 1],
+ rowLoc="right",
+ )
+
+ # bold cells for header and first column, containing variable and level names
+ bold_cells = [(var, -1) for var in np.arange(num_vars) + 1]
+ bold_cells.extend([(0, lvl) for lvl in np.arange(num_lvls)])
+
+ # Shade the cells and make the cell text larger
+ for cell in stats_table._cells:
+ stats_table._cells[cell].set_alpha(0.7)
+
+ domain_text = {"conus": "CONUS", "full": "Unified Surface Analysis Domain"}
+
+ plt.suptitle(
+ "%s permutations: %s"
+ % (constants.FRONT_NAMES[front_type], domain_text[args["domain"]]),
+ fontsize=18,
+ y=1.02,
+ )
+ plt.tight_layout()
+ plt.savefig(
+ "%s/permutations_%d_%s_%s.%s"
+ % (
+ model_folder,
+ args["model_number"],
+ front_type,
+ args["domain"],
+ args["output_type"],
+ ),
+ bbox_inches="tight",
+ dpi=400,
+ edgecolor="black",
+ )
+ plt.close()
diff --git a/src/fronts/plot/plot_saliency_maps.py b/src/fronts/plot/plot_saliency_maps.py
new file mode 100644
index 0000000..23ce531
--- /dev/null
+++ b/src/fronts/plot/plot_saliency_maps.py
@@ -0,0 +1,230 @@
+"""
+Plot saliency maps for model predictions.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2024.10.11
+"""
+
+import argparse
+import matplotlib.colors
+import pandas as pd
+import cartopy.crs as ccrs
+import matplotlib.pyplot as plt
+import numpy as np
+import xarray as xr
+from matplotlib import cm
+from fronts.utils import data_utils
+from fronts.utils.plotting import plot_background
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--model_dir", type=str, required=True, help="Directory for the models."
+ )
+ parser.add_argument("--model_number", type=int, required=True, help="Model number.")
+ parser.add_argument(
+ "--plot_outdir", type=str, help="Directory for the saliency maps."
+ )
+ parser.add_argument(
+ "--init_time",
+ type=str,
+ help="Initialization time of the data. Format: YYYY-MM-DD-HH.",
+ )
+ parser.add_argument("--domain", type=str, required=True, help="Domain of the data.")
+ parser.add_argument(
+ "--calibration",
+ type=int,
+ help="Neighborhood calibration distance in kilometers. Possible neighborhoods are 50, 100, 150, 200, and 250 km.",
+ )
+ parser.add_argument(
+ "--data_source",
+ type=str,
+ default="era5",
+ help="Source of the variable data (ERA5, GDAS, etc.)",
+ )
+ parser.add_argument(
+ "--cmap",
+ type=str,
+ default="viridis",
+ help="Colormap to use for the saliency maps.",
+ )
+ parser.add_argument(
+ "--extent",
+ type=int,
+ nargs=4,
+ help="Extent of the saliency maps to plot. If this argument is not provided, use the default domain extent.",
+ )
+
+ args = vars(parser.parse_args())
+
+ init_time = pd.date_range(args["init_time"], args["init_time"])[0]
+ extent = (
+ data_utils.DOMAIN_EXTENTS[args["domain"]]
+ if args["extent"] is None
+ else args["extent"]
+ )
+
+ model_properties = pd.read_pickle(
+ f"{args['model_dir']}/model_{args['model_number']}/model_{args['model_number']}_properties.pkl"
+ )
+ front_types = model_properties["dataset_properties"]["front_types"]
+
+ salmap_folder = "%s/model_%d/saliencymaps" % (
+ args["model_dir"],
+ args["model_number"],
+ )
+ salmap_ds = xr.open_dataset(
+ "%s/model_%d_salmap_%s_%s_%d%02d.nc"
+ % (
+ salmap_folder,
+ args["model_number"],
+ args["domain"],
+ args["data_source"],
+ init_time.year,
+ init_time.month,
+ )
+ )
+
+ probs_folder = "%s/model_%d/probabilities" % (
+ args["model_dir"],
+ args["model_number"],
+ )
+ probs_ds = xr.open_dataset(
+ "%s/model_%d_pred_%s_%d%02d.nc"
+ % (
+ probs_folder,
+ args["model_number"],
+ args["domain"],
+ init_time.year,
+ init_time.month,
+ )
+ ).sel(time=init_time)
+ levels = np.around(np.arange(0, 1.1, 0.1), 2)
+
+ for front_type in front_types:
+ if args["calibration"] is not None:
+ try:
+ ir_model = model_properties["calibration_models"][args["domain"]][
+ front_type
+ ]["%d km" % args["calibration"]]
+ except KeyError:
+ ir_model = model_properties["calibration_models"]["conus"][front_type][
+ "%d km" % args["calibration"]
+ ]
+ original_shape = np.shape(probs_ds[front_type].values)
+ probs_ds[front_type].values = ir_model.predict(
+ probs_ds[front_type].values.flatten()
+ ).reshape(original_shape)
+ cbar_label = "Probability (calibrated - %d km)" % args["calibration"]
+ else:
+ cbar_label = "Probability (uncalibrated)"
+
+ # mask out low probabilities
+ probs_ds[front_type].values = np.where(
+ probs_ds[front_type].values < 0.1, np.nan, probs_ds[front_type].values
+ )
+
+ cmap_probs, norm = (
+ cm.get_cmap(data_utils.CONTOUR_CMAPS[front_type], 11),
+ matplotlib.colors.Normalize(vmin=0, vmax=1),
+ )
+
+ salmap_for_type = salmap_ds[front_type].sel(time=init_time)
+ salmap_for_type_pl = salmap_ds[front_type + "_pl"].sel(time=init_time)
+ max_gradient, min_gradient = salmap_for_type.max(), salmap_for_type.min()
+ salmap_for_type = (salmap_for_type - min_gradient) / (
+ max_gradient - min_gradient
+ )
+ salmap_for_type_pl = (salmap_for_type_pl - min_gradient) / (
+ max_gradient - min_gradient
+ )
+
+ fig, axs = plt.subplots(
+ 3,
+ 2,
+ subplot_kw={
+ "projection": ccrs.PlateCarree(
+ central_longitude=(extent[0] + extent[1]) / 2
+ )
+ },
+ )
+ axarr = axs.flatten()
+ for ax_ind, ax in enumerate(axarr):
+ plot_background(extent, ax=ax, linewidth=0.3)
+ if ax_ind == 0:
+ probs_ds[front_type].plot.contourf(
+ ax=ax,
+ x="longitude",
+ y="latitude",
+ cmap=cmap_probs,
+ norm=norm,
+ levels=levels,
+ transform=ccrs.PlateCarree(),
+ alpha=0.8,
+ add_colorbar=False,
+ )
+ else:
+ # probs_ds[front_type].plot.contour(ax=ax, x='longitude', y='latitude', colors='black', linewidths=0.1, norm=norm, levels=levels, transform=ccrs.PlateCarree(), alpha=0.8)
+ salmap_for_type_pl.isel(pressure_level=ax_ind - 1).plot(
+ ax=ax,
+ x="longitude",
+ y="latitude",
+ cmap=args["cmap"],
+ norm=norm,
+ transform=ccrs.PlateCarree(),
+ alpha=0.6,
+ add_colorbar=False,
+ )
+
+ axarr[0].set_title("a) Model predictions")
+ axarr[1].set_title("b) Saliency map - surface")
+ axarr[2].set_title("c) Saliency map - 1000 hPa")
+ axarr[3].set_title("d) Saliency map - 950 hPa")
+ axarr[4].set_title("e) Saliency map - 900 hPa")
+ axarr[5].set_title("f) Saliency map - 850 hPa")
+
+ cbar_ax = fig.add_axes([0.1, -0.05, 0.8, 0.05])
+ cbar = plt.colorbar(
+ cm.ScalarMappable(
+ norm=matplotlib.colors.Normalize(vmin=0, vmax=1), cmap=args["cmap"]
+ ),
+ orientation="horizontal",
+ cax=cbar_ax,
+ alpha=0.8,
+ label="Normalized saliency",
+ )
+
+ # if a directory for the plots is not provided, save the plots to the folder containing saliency maps
+ plot_outdir = (
+ args["plot_outdir"] if args["plot_outdir"] is not None else salmap_folder
+ )
+
+ plt.tight_layout()
+ plt.suptitle(
+ "%s predictions: %d-%02d-%02d-%02dz"
+ % (
+ data_utils.FRONT_NAMES[front_type],
+ init_time.year,
+ init_time.month,
+ init_time.day,
+ init_time.hour,
+ ),
+ y=1.05,
+ )
+ plt.savefig(
+ "%s/model_%d_salmap_%d%02d%02d%02d_%s_%s.png"
+ % (
+ plot_outdir,
+ args["model_number"],
+ init_time.year,
+ init_time.month,
+ init_time.day,
+ init_time.hour,
+ args["domain"],
+ front_type,
+ ),
+ bbox_inches="tight",
+ dpi=500,
+ )
+ plt.close()
diff --git a/src/fronts/train.py b/src/fronts/train.py
new file mode 100644
index 0000000..605609b
--- /dev/null
+++ b/src/fronts/train.py
@@ -0,0 +1,473 @@
+"""Train a FrontFinder model with optional Weights and Biases tracking."""
+
+import tensorflow.keras # ty: ignore[unresolved-import]
+import logging
+import os
+import wandb
+from wandb.integration import keras as wandb_keras
+import dataclasses
+import datetime
+from typing import Literal, Any, Optional, Union, TypeVar, Type
+import argparse
+import dacite
+import yaml
+from fronts.model import ModelConfig
+from fronts.data.config import DataConfig
+
+# ---------------------------------------------------------------------------
+# Module-level logger — writes to stderr so output appears in Slurm logs even
+# when stdout is redirected. Log level can be overridden by setting the
+# FRONTS_LOG_LEVEL environment variable, e.g. FRONTS_LOG_LEVEL=DEBUG.
+# ---------------------------------------------------------------------------
+logging.basicConfig(
+ level=os.environ.get("FRONTS_LOG_LEVEL", "INFO"),
+ format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+ datefmt="%Y-%m-%d %H:%M:%S",
+)
+log = logging.getLogger("fronts.train")
+
+T = TypeVar("T")
+
+
+@dataclasses.dataclass
+class WandBConfig:
+ """Configuration dataclass for Weights and Biases.
+
+ Initializing a WandBConfig object will automatically login using api_key
+ which defaults to the WANDB_KEY environment variable.
+
+ Attributes:
+ project_name: the WandB project where model training data will be stored.
+ model_run_name: the name of the model run.
+ log_frequency: the rate in epochs of log storage. Defaults to each epoch.
+ upload_checkpoints: whether or not to upload the model checkpoints. Defaults to
+ False.
+ api_key: the API key for a WandB account. Defaults to the WANDB_KEY environment
+ var.
+ wandb_filepath: path for WandbModelCheckpoint. Must end in `.keras`.
+ Defaults to `models/.keras`.
+ """
+
+ project_name: str
+ model_run_name: str
+ log_frequency: int = 1
+ upload_checkpoints: bool = False
+ api_key: str = os.environ.get("WANDB_KEY", "")
+ wandb_filepath: Optional[str] = None
+
+ def __post_init__(self):
+ # Default checkpoint path: models/.keras
+ # WandbModelCheckpoint requires a .keras extension (Keras 3 requirement).
+ if self.wandb_filepath is None:
+ self.wandb_filepath = f"models/{self.model_run_name}.keras"
+ self.login()
+
+ def login(self):
+ """Helper method to automatically login to WandB.
+
+ Skipped if no API key is configured (e.g. local dry-run without credentials).
+ """
+ if self.api_key:
+ wandb.login(key=self.api_key)
+
+ def build_init_config(self, init_config: dict) -> dict:
+ """Builds the keyword arguments to apply to wandb.init.
+
+ Args:
+ init_config: the dictionary of all model properties to pass into the WandB
+ run instance.
+
+ Returns a dictionary of project, config, and name, arguments for wandb.init.
+ """
+ init_config = {
+ "project": self.project_name,
+ "config": init_config,
+ "name": self.model_run_name,
+ }
+ return init_config
+
+ def build_keras_metriclogger_callback(
+ self,
+ ) -> wandb_keras.WandbMetricsLogger:
+ """Returns an instance of a MetricsLogger callback using the log_frequency
+ attribute.
+ """
+ return wandb_keras.WandbMetricsLogger(log_freq=self.log_frequency)
+
+ def build_keras_modelcheckpoint_callback(
+ self,
+ ) -> tensorflow.keras.callbacks.ModelCheckpoint:
+ """Return the ModelCheckpoint WandB callback using the wandb_filepath
+ attribute.
+ """
+
+ return wandb_keras.WandbModelCheckpoint(self.wandb_filepath)
+
+ def build_all_callbacks(self) -> list[Any]:
+ """Returns both ModelCheckpoint and MetricsLogger callbacks."""
+
+ return [
+ self.build_keras_modelcheckpoint_callback(),
+ self.build_keras_metriclogger_callback(),
+ ]
+
+
+@dataclasses.dataclass
+class CallbacksConfig:
+ """A configuration for non-Weights and Biases callbacks.
+
+ Certain attributes are shared amongst callbacks, including monitor and verbose.
+
+ Attributes:
+ monitor: the metric to monitor.
+ verbose: integer determining the amount of logs returned from the callbacks.
+ save_best_only: if True, will only save if the model checkpoint has the best
+ metrics so far.
+ save_weights_only: will only save weights if set to True.
+ save_freq: how frequently to save the model checkpoint. If int n, will save
+ after n batches.
+ model_checkpoint_path: the path where the model will be saved. Defaults to None.
+ If None, does not initialize ModelCheckpoint callback.
+ csv_logger_path: the path where the csv logger will be saved. Defaults to None.
+ If None, does not initialize CSVLogger callback.
+ patience: the number of epochs to run with no improvement before stopping early.
+ Defaults to None. If None, does not initialize EarlyStopping callback.
+ """
+
+ monitor: str
+ verbose: int
+ save_best_only: bool
+ save_weights_only: bool
+ save_freq: Union[Literal["epoch"], int]
+ model_checkpoint_path: Optional[str] = None
+ csv_logger_path: Optional[str] = None
+ patience: Optional[int] = None
+
+ def build(self) -> list[tensorflow.keras.callbacks.Callback]:
+ # Initialize list
+ callback_list = []
+
+ # Only append the list with the callback if the conditions are met, i.e. the
+ # key attributes are not None. It is possible to return an empty list
+ if self.model_checkpoint_path:
+ checkpoint_callback = tensorflow.keras.callbacks.ModelCheckpoint(
+ filepath=self.model_checkpoint_path,
+ monitor=self.monitor,
+ verbose=self.verbose,
+ save_best_only=self.save_best_only,
+ save_weights_only=self.save_weights_only,
+ save_freq=self.save_freq,
+ )
+ callback_list.append(checkpoint_callback)
+ if self.csv_logger_path:
+ history_logger_callback = tensorflow.keras.callbacks.CSVLogger(
+ filename=self.csv_logger_path, append=True
+ )
+ callback_list.append(history_logger_callback)
+ if self.patience:
+ early_stopping_callback = tensorflow.keras.callbacks.EarlyStopping(
+ monitor=self.monitor, patience=self.patience, verbose=self.verbose
+ )
+ callback_list.append(early_stopping_callback)
+
+ return callback_list
+
+
+class Trainer:
+ """Main class to build and trigger model training for FrontFinder."""
+
+ def __init__(
+ self,
+ model,
+ data,
+ epochs: int,
+ validation_frequency: int,
+ training_steps_per_epoch: int,
+ validation_steps_per_epoch: int,
+ callbacks: list = None,
+ verbose: Literal["auto", 0, 1, 2] = "auto",
+ wandb: Optional[WandBConfig] = None,
+ repeat: bool = True,
+ seed: int = 42,
+ ) -> None:
+ """Initialize the Trainer class and maybe build callbacks.
+
+ Arguments:
+ model: the model to use for training.
+ data: the ModelDataConfig which holds the prepared train, valid, and test
+ data.
+ epochs: number of epochs to run to train the model.
+ validation_frequency: specifies how many training epochs to run before a new
+ validation run is performed, e.g. validation_freq=2 runs validation
+ every 2 epochs.
+ training_steps_per_epoch: total number of batches of samples to run per
+ epoch.
+ validation_steps_per_epoch: total number of batches of samples for
+ validation. If validation_steps is specified and only part of the
+ dataset will be consumed, the evaluation will start from the beginning
+ of the dataset at each epoch.
+ callbacks: optional list of callbacks (not including WandB callbacks) to use
+ when training the model.
+ verbose: "auto", 0, 1, or 2. Verbosity mode. 0 = silent, 1 = progress bar,
+ 2 = one line per epoch. "auto" ~= 1. Defaults to "auto".
+ wandb: the Weights and Biases configuration object to use, if exists.
+ repeat: whether or not the training dataset will repeat indefinitely.
+ Defaults to True. If True, training_steps_per_epoch will determine how
+ many batches will run per epoch.
+ seed: the seed to use for for all of the backend seeds to allow for
+ determinism. Defaults to 42.
+
+
+ """
+ self.model = model
+ self.wandb = wandb
+ self.data = data
+ self.epochs = epochs
+ self.validation_frequency = validation_frequency
+ self.training_steps_per_epoch = training_steps_per_epoch
+ self.validation_steps_per_epoch = validation_steps_per_epoch
+ self.verbose = verbose
+ self.repeat = repeat
+ self.callbacks = callbacks or []
+ self.seed = seed
+
+ def train(self, model: dict) -> None:
+ """Triggers a keras training run using model.fit().
+
+ Args:
+ model: the complete metadata of configuration of the model
+ """
+
+ # Set the seed for fitting the model
+ tensorflow.keras.utils.set_random_seed(self.seed)
+
+ # If indefinite repeat is enabled, instantiate the bound method
+ if self.repeat:
+ training_data = self.data.train_data.repeat()
+ else:
+ training_data = self.data.train_data
+
+ # Set up the arguments to fit
+ fit_args = {
+ "x": training_data,
+ "validation_data": self.data.validation_data,
+ "validation_freq": self.validation_frequency,
+ "epochs": self.epochs,
+ "steps_per_epoch": self.training_steps_per_epoch,
+ "validation_steps": self.validation_steps_per_epoch,
+ "verbose": self.verbose,
+ }
+
+ # Use WandB if exists — WandB callbacks must be built AFTER wandb.init()
+ if self.wandb:
+ wandb_init = self.wandb.build_init_config(model)
+ with wandb.init(**wandb_init) as _: # ty: ignore[invalid-context-manager]
+ fit_args["callbacks"] = self.build_callbacks(self.callbacks)
+ self.model.fit(**fit_args)
+ else:
+ fit_args["callbacks"] = self.build_callbacks(self.callbacks)
+ self.model.fit(**fit_args)
+
+ def build_callbacks(self, callbacks: list):
+ """Combine all callbacks that exist.
+
+ Acts as a passthrough if WandB is not being used, only including the callbacks
+ provided when initializing the Trainer.
+
+ Args:
+ callbacks: a list of 0 or more callbacks to include when training the model.
+
+ Returns a list of 0 or more callbacks.
+ """
+ # If WandB is being used, add the callbacks in the dataclass.
+ if self.wandb:
+ callbacks.extend(self.wandb.build_all_callbacks())
+ return callbacks
+
+
+@dataclasses.dataclass
+class TrainConfig:
+ """Model training dataclass.
+
+ Configuration details mostly from
+ https://www.tensorflow.org/api_docs/python/tf/keras/Model.
+
+ Attributes:
+
+ epochs: number of epochs to run to train the model.
+ training_steps_per_epoch: total number of batches of samples to run per epoch.
+ validation_steps_per_epoch: total number of batches of samples for validation.
+ If validation_steps is specified and only part of the dataset will be
+ consumed, the evaluation will start from the beginning of the dataset at
+ each epoch.
+ callbacks: CallbackObject specifying which callbacks to include with training.
+ validation_freq: specifies how many training epochs to run before a new
+ validation run is performed, e.g. validation_freq=2 runs validation every 2
+ epochs.
+ verbose: "auto", 0, 1, or 2. Verbosity mode. 0 = silent, 1 = progress bar,
+ 2 = one line per epoch. "auto" ~= 1. Defaults to "auto".
+ repeat: whether or not the training dataset will repeat indefinitely. Defaults
+ to True. If True, training_steps_per_epoch will determine how many batches
+ will run per epoch.
+ seed: the seed to use for for all of the backend seeds to allow for determinism.
+ Defaults to 42.
+ """
+
+ model: ModelConfig
+ wandb: WandBConfig
+ callbacks: CallbacksConfig
+ epochs: int
+ training_steps_per_epoch: int
+ validation_steps_per_epoch: int | None
+ validation_frequency: int
+ verbose: Literal["auto", 0, 1, 2]
+ repeat: bool
+ seed: int
+ data: Optional[DataConfig] = None
+
+ def build(
+ self,
+ ) -> Trainer:
+ """Builds the Trainer object which can be used to train the model.
+
+ Args:
+ model: the model to use for training.
+ data: the ModelDataConfig which holds the prepared train, valid, and test
+ data.
+ wandb: the Weights and Biases configuration object to use, if exists.
+ callbacks: optional list of callbacks (not including WandB callbacks) to use
+ when training the model.
+
+ Returns a Trainer object that can be used to instantiate a training run.
+ """
+ log.info("Building callbacks...")
+ callbacks = self.callbacks.build()
+ log.debug("Callbacks built: %s", [type(c).__name__ for c in callbacks])
+
+ if self.data is not None:
+ log.info("Building data pipeline (this may take several minutes)...")
+ model_data = self.data.build()
+ log.info("Data pipeline ready.")
+ else:
+ log.warning("No data config provided — trainer will have empty data.")
+ model_data = ""
+
+ # Derive input_shape and num_classes from the training dataset element spec.
+ # input_shape: set spatial dims to None for variable-size inference.
+ # num_classes: last dim of the target tensor (e.g. 6 for 5 front types + background).
+ log.info("Deriving input_shape and num_classes from training dataset...")
+ train_spec = model_data.train_data.element_spec
+ raw_input_shape = list(train_spec[0].shape)
+ input_shape = tuple([None, None] + raw_input_shape[2:])
+ num_classes = train_spec[1].shape[-1]
+ log.info(" input_shape=%s, num_classes=%d", input_shape, num_classes)
+
+ log.info("Building model...")
+ keras_model = self.model.build(input_shape=input_shape, num_classes=num_classes)
+ keras_model.summary(print_fn=log.info)
+ log.info("Model built and compiled.")
+
+ log.info("Building Trainer...")
+ trainer = Trainer(
+ model=keras_model,
+ data=model_data,
+ epochs=self.epochs,
+ validation_frequency=self.validation_frequency,
+ training_steps_per_epoch=self.training_steps_per_epoch,
+ validation_steps_per_epoch=self.validation_steps_per_epoch,
+ callbacks=callbacks,
+ verbose=self.verbose,
+ wandb=self.wandb,
+ repeat=self.repeat,
+ seed=self.seed,
+ )
+ return trainer
+
+
+def open_config_yaml_as_dataclass(
+ path: str, config_class: Type[T], require: bool = False
+) -> Optional[T]:
+ """Opens a configuration yaml if exists and returns it as the relevant dataclass.
+
+ Args:
+ path: the absolute path to the configuration file.
+ config_class: the configuration dataclass that the incoming yaml will be
+ converted to via dacite.
+ require: If True, code will throw an error if the path is not provided.
+ Defaults to False.
+
+ Returns either None or the dataclass if path is provided.
+ """
+ if path:
+ with open(file=path) as f:
+ config_yaml = yaml.safe_load(f)
+ _class_instance = dacite.from_dict(
+ data_class=config_class,
+ data=config_yaml,
+ config=dacite.Config(cast=[tuple, datetime.datetime], check_types=False),
+ )
+ return _class_instance
+ elif require:
+ raise ValueError("Path must be included when require is True.")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+
+ parser.add_argument(
+ "-tc",
+ "--train_config_path",
+ type=str,
+ required=True,
+ help=(
+ "Path to the training configuration yaml. This config must include epochs, "
+ "training_steps_per_epoch, validation_steps_per_epoch, "
+ "validation_frequency, and optionally verbose and repeat. See TrainConfig "
+ "for more information on each of these attributes."
+ ),
+ )
+ parser.add_argument(
+ "--dry_run",
+ action="store_true",
+ default=False,
+ help=(
+ "Validate config parsing and data pipeline construction without running "
+ "training or initializing WandB. Exits after the data pipeline is built "
+ "and reports success. Useful for catching bugs locally before submitting "
+ "a SLURM job."
+ ),
+ )
+
+ args = parser.parse_args()
+
+ log.info("Loading config from: %s", args.train_config_path)
+ train_config = open_config_yaml_as_dataclass(
+ path=args.train_config_path, config_class=TrainConfig, require=True
+ )
+ log.info("Config loaded. epochs=%d, steps_per_epoch=%d",
+ train_config.epochs, train_config.training_steps_per_epoch)
+
+ if args.dry_run:
+ log.info("=== DRY RUN MODE — skipping WandB init and training ===")
+ if train_config.data is not None:
+ log.info("Building data pipeline...")
+ model_data = train_config.data.build()
+ log.info("Data pipeline built successfully.")
+ log.info(" train_data: %s", model_data.train_data)
+ log.info(" validation_data: %s", model_data.validation_data)
+ log.info(" test_data: %s", model_data.test_data)
+ else:
+ log.warning("No data config present — nothing to validate.")
+ log.info("=== DRY run complete. No errors. ===")
+ raise SystemExit(0)
+
+ log.info("Building trainer (data pipeline + model will be constructed here)...")
+ trainer = train_config.build() # ty:ignore[possibly-missing-attribute]
+ log.info("Trainer ready. Starting training run...")
+
+ # model_config dict is passed to wandb.init for run metadata
+ model_config = dataclasses.asdict(train_config.model)
+
+ # Trigger training run
+ trainer.train(model=model_config)
+ log.info("Training complete.")
diff --git a/src/fronts/train_legacy.py b/src/fronts/train_legacy.py
new file mode 100644
index 0000000..19ab540
--- /dev/null
+++ b/src/fronts/train_legacy.py
@@ -0,0 +1,1080 @@
+"""
+Script that trains a new U-Net model.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+
+Update: this module is now deprecated and is only used as a reference
+while module is organized into other modules and directories.
+"""
+
+import argparse
+import pandas as pd
+from tensorflow.keras.callbacks import EarlyStopping, CSVLogger
+import tensorflow as tf
+import pickle
+import numpy as np
+from fronts.utils import file_manager
+import os
+from fronts.model import unets, metrics, losses
+import datetime
+from fronts.utils import misc, data_utils
+from glob import glob
+import wandb
+
+
+class ArgumentParser(argparse.ArgumentParser):
+ """Custom argument parser class that allows a list of arguments to be passed on the same line from a text file"""
+
+ def convert_arg_line_to_args(self, arg_line):
+ """Allow multiple arguments to be passed on one line if calling a text file with arguments (e.g., --arg 2 5 6)"""
+ return arg_line.split()
+
+
+if __name__ == "__main__":
+ parser = ArgumentParser(fromfile_prefix_chars="@")
+ """
+ ### WandB ###
+ parser.add_argument(
+ "--project",
+ type=str,
+ help="WandB project that will be used to store model training data.",
+ )
+ parser.add_argument(
+ "--log_freq",
+ type=int,
+ default=1,
+ help="WandB loss/metric logging frequency in epochs.",
+ )
+ parser.add_argument(
+ "--upload_model", action="store_true", help="Upload model checkpoints to WandB."
+ )
+ parser.add_argument("--key", type=str, help="WandB API key.")
+ parser.add_argument(
+ "--name",
+ type=str,
+ help="WandB name for the current model run. If no name is specified, it will default to the model number (e.g. model_129482).",
+ )"""
+
+ ### General arguments ###
+ parser.add_argument(
+ "--model_dir",
+ type=str,
+ required=True,
+ help="Directory where the models are or will be saved to.",
+ )
+ parser.add_argument(
+ "--model_number",
+ type=int,
+ help="Number that the model will be assigned. If no argument is passed, a number will be automatically assigned based "
+ "on the current date and time.",
+ )
+
+ # --------------------- will be created via era5/batch.py
+
+ """ parser.add_argument(
+ "--tf_indirs",
+ type=str,
+ required=True,
+ nargs="+",
+ help="Directories for the TensorFlow datasets. One or two paths can be passed. If only one path is passed, then the "
+ "training and validation datasets will be pulled from this path. If two paths are passed, the training dataset "
+ "will be pulled from the first path and the validation dataset from the second.",
+ )"""
+ # --------------------- handled in train.py or yaml config
+
+ """ parser.add_argument(
+ "--epochs",
+ type=int,
+ required=True,
+ help="Maximum number of epochs for model training.",
+ )
+ parser.add_argument(
+ "--patience",
+ type=int,
+ help="Patience for EarlyStopping callback. If this argument is not provided, it will be set according to the size "
+ "of the training dataset (images in training set divided by the product of the batch size and steps).",
+ )
+ parser.add_argument(
+ "--verbose",
+ type=int,
+ default=2,
+ help="Model.fit verbose. Unless you want a text file that is several hundred megabytes in size and takes 10 years "
+ "to scroll through, I suggest you leave this at 2.",
+ )
+ parser.add_argument(
+ "--seed",
+ type=int,
+ default=np.random.randint(0, 2**31 - 1),
+ help="Seed for the random number generators. If a model is being retrained with the --retrain flag, this argument "
+ "will be overriden by the previous seed used to train that model.",
+ )"""
+
+ ### GPU and hardware arguments ###
+ parser.add_argument("--gpu_device", type=int, nargs="+", help="GPU device numbers.")
+ parser.add_argument(
+ "--memory_growth", action="store_true", help="Use memory growth for GPUs"
+ )
+ parser.add_argument(
+ "--num_parallel_calls",
+ type=int,
+ default=4,
+ help="Number of parallel calls for retrieving batches for the training and validation datasets.",
+ )
+ parser.add_argument(
+ "--buffer_size",
+ type=int,
+ help="Maximum buffer size used when shuffling the training dataset. By default, the entire training dataset will "
+ "be shuffled, and the buffer size is equal to the number of images in the training dataset.",
+ )
+ parser.add_argument(
+ "--cache",
+ type=str,
+ help="Directory where the datasets will be cached for training. Passing 'RAM' or an empty string will cache the "
+ "datasets directly to RAM.",
+ )
+ parser.add_argument(
+ "--disable_tensorfloat32",
+ action="store_true",
+ help="Disable TensorFloat32 execution.",
+ )
+
+ ### Hyperparameters ###
+ parser.add_argument(
+ "--learning_rate",
+ type=float,
+ help="Learning rate for U-Net optimizer. If left as None, then the default optimizer learning rate will be used.",
+ )
+ parser.add_argument(
+ "--batch_size",
+ type=int,
+ required=True,
+ nargs="+",
+ help="Batch sizes for the U-Net. Up to 2 arguments can be passed. If 1 argument is passed, the value will be both "
+ "the training and validation batch sizes. If 2 arguments are passed, the first and second arguments will be "
+ "the training and validation batch sizes, respectively.",
+ )
+ parser.add_argument(
+ "--steps",
+ type=int,
+ nargs="+",
+ help="Number of steps for each epoch. Up to 2 arguments can be passed. If 1 argument is passed, the value will only "
+ "be applied to the number of steps per epoch, and the number of validation steps will be calculated by tensorflow "
+ "such that the entire validation dataset is passed into the model during validation. If 2 arguments are passed, "
+ "then the arguments are the number of steps in training and validation. If no arguments are passed, then the "
+ "number of steps in both training and validation will be calculated by tensorflow.",
+ )
+ parser.add_argument(
+ "--valid_freq",
+ type=int,
+ default=1,
+ help="How many epochs to complete before validation.",
+ )
+
+ ### U-Net arguments ###
+ parser.add_argument(
+ "--model_type", # registry
+ type=str,
+ help="Model type. Options are: unet, unet_ensemble, unet_plus, unet_2plus, unet_3plus, attention_unet.",
+ )
+ parser.add_argument(
+ "--activation", # core
+ type=str,
+ help="Activation function to use in the model. Refer to utils.unet_utils.choose_activation_layer to see all available "
+ "activation functions.",
+ )
+ parser.add_argument(
+ "--batch_normalization", # core
+ action="store_true",
+ help="Use batch normalization in the model. This will place batch normalization layers after each convolution layer.",
+ )
+ parser.add_argument(
+ "--deep_supervision", # plus, 2plus, 3plus
+ action="store_true",
+ help="Use deep supervision in the model. Deep supervision creates side outputs from the bottom encoder node and each decoder node.",
+ )
+ parser.add_argument(
+ "--filter_num", # core
+ type=int,
+ nargs="+",
+ help="Number of filters in each level of the U-Net. The number of arguments passed to --filter_num must be equal to the "
+ "value passed to --levels.",
+ )
+ parser.add_argument(
+ "--filter_num_aggregate", # 3plus (as num_aggregate_filters)
+ type=int,
+ help="Number of filters in aggregated feature maps. This argument is only used in the U-Net 3+.",
+ )
+ parser.add_argument(
+ "--filter_num_skip", # 3plus (as full_scale_skip_connection_filters)
+ type=int,
+ help="Number of filters in full-scale skip connections in the U-Net 3+.",
+ )
+ parser.add_argument(
+ "--first_encoder_connections", # 3plus
+ action="store_true",
+ help="Enable first encoder connections in the U-Net 3+.",
+ )
+ parser.add_argument(
+ "--kernel_size", # core
+ type=int,
+ nargs="+",
+ help="Size of the convolution kernels. One integer can be passed to make the kernel dimensions have equal length (e.g. "
+ "passing 3 has the same effect as passing 3 3 3 for 3-dimensional kernels.)",
+ )
+ parser.add_argument(
+ "--levels", # core (as depth)
+ type=int,
+ help="Number of levels in the model, also known as the 'depth' of the model.",
+ )
+ parser.add_argument(
+ "--loss", # core
+ type=str,
+ nargs="+",
+ help="Loss function for the U-Net (arg 1), with keyword arguments (arg 2). Keyword arguments must be passed as a "
+ "string in the second argument. See 'utils.misc.string_arg_to_dict' for more details. Raises a ValueError if "
+ "more than 2 arguments are passed.",
+ )
+ parser.add_argument(
+ "--metric", # core
+ type=str,
+ nargs="+",
+ help="Metric for evaluating the U-Net during training (arg 1), with keyword arguments (arg 2). Keyword arguments "
+ "must be passed as a string in the second argument. See 'utils.misc.string_arg_to_dict' for more details. Raises "
+ "a ValueError if more than 2 arguments are passed.",
+ )
+ parser.add_argument(
+ "--modules_per_node", # core
+ type=int,
+ default=5,
+ help="Number of convolution modules in each node. A convolution module consists of a convolution layer followed by "
+ "an optional batch normalization layer and an activation layer. (e.g. Conv3D -> BatchNormalization -> PReLU; Conv3D -> PReLU)",
+ )
+ parser.add_argument(
+ "--optimizer", # core (OptimizerConfig)
+ type=str,
+ nargs="+",
+ default=[
+ "Adam",
+ ],
+ help="Optimizer to use during the training process (arg 1), with keyword arguments (arg 2). Keyword arguments "
+ "must be passed as a string in the second argument. See 'utils.misc.string_arg_to_dict' for more details. Raises "
+ "a ValueError if more than 2 arguments are passed.",
+ )
+ parser.add_argument(
+ "--padding", # core
+ type=str,
+ default="same",
+ help="Padding to use in the convolution layers. If 'same', then zero-padding will be added to the inputs such that the outputs "
+ "of the layers will be the same shape as the inputs. If 'valid', no padding will be applied to the layers' inputs.",
+ )
+ parser.add_argument(
+ "--pool_size", # core
+ type=int,
+ nargs="+",
+ help="Pool size for the max pooling layers. One integer can be passed to make the pooling dimensions have equal length "
+ "(e.g. passing 2 has the same effect as passing 2 2 2 for 3-dimensional max pooling.)",
+ )
+ parser.add_argument(
+ "--upsample_size", # core (excluded from attention_unet via excpetion)
+ type=int,
+ nargs="+",
+ help="Upsampling factors for the up-sampling layers. One integer can be passed to make the factors have equal size "
+ "(e.g. passing 2 has the same effect as passing 2 2 2 for 3-dimensional up-sampling.)",
+ )
+ parser.add_argument(
+ "--use_bias", # core
+ action="store_true",
+ help="Use bias parameters in the convolution layers.",
+ )
+
+ ### Constraints, initializers, and regularizers ###
+ parser.add_argument(
+ "--activity_regularizer", # ConvOutputConfig
+ type=str,
+ nargs="+",
+ default=[
+ None,
+ ],
+ help="Regularizer function applied to the output of the Conv2D/Conv3D layers. A second string argument can be passed "
+ "containing keyword arguments for the regularizer.",
+ )
+ parser.add_argument(
+ "--bias_constraint", # BiasVectorConfig
+ type=str,
+ nargs="+",
+ default=[
+ None,
+ ],
+ help="Constraint function applied to the bias vector of the Conv2D/Conv3D layers. A second string argument can be "
+ "passed containing keyword arguments for the constraint.",
+ )
+ parser.add_argument(
+ "--bias_initializer", # BiasVectorConfig
+ type=str,
+ default="zeros",
+ help="Initializer for the bias vector in the Conv2D/Conv3D layers.",
+ )
+ parser.add_argument(
+ "--bias_regularizer", # BiasVectorConfig
+ type=str,
+ nargs="+",
+ default=[
+ None,
+ ],
+ help="Regularizer function applied to the bias vector in the Conv2D/Conv3D layers. A second string argument can "
+ "be passed containing keyword arguments for the regularizer.",
+ )
+ parser.add_argument(
+ "--kernel_constraint", # KernelMatrixConfig
+ type=str,
+ nargs="+",
+ default=[
+ None,
+ ],
+ help="Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers. A second string argument can "
+ "be passed containing keyword arguments for the constraint.",
+ )
+ parser.add_argument(
+ "--kernel_initializer", # KernelMatrixConfig
+ type=str,
+ default="glorot_uniform",
+ help="Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.",
+ )
+ parser.add_argument(
+ "--kernel_regularizer", # KernelMatrixConfig
+ type=str,
+ nargs="+",
+ default=[
+ None,
+ ],
+ help="Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers. A second string argument "
+ "can be passed containing keyword arguments for the regularizer.",
+ )
+
+ ### Data arguments ###
+ parser.add_argument(
+ "--num_training_years",
+ type=int,
+ help="Number of years for the training dataset.",
+ )
+ parser.add_argument(
+ "--training_years", type=int, nargs="+", help="Years for the training dataset."
+ )
+ parser.add_argument(
+ "--num_validation_years",
+ type=int,
+ help="Number of years for the validation set.",
+ )
+ parser.add_argument(
+ "--validation_years", type=int, nargs="+", help="Years for the validation set."
+ )
+ parser.add_argument(
+ "--shuffle",
+ type=str,
+ default="full",
+ help="Shuffling method for the training set. Valid options are 'lazy' or 'full' (default is 'full'). "
+ "A 'lazy' shuffle will only shuffle the order of the monthly datasets but not the contents within. A 'full' "
+ "shuffle will shuffle every image inside the dataset.",
+ )
+
+ ### Retraining model ###
+ parser.add_argument("--retrain", action="store_true", help="Retrain a model")
+
+ ### Debug ###
+ parser.add_argument(
+ "--no_train",
+ action="store_true",
+ help="Do not train the model. This argument will allow everything in the script to run as normal but will not start "
+ "the training process. In addition, no directory for the model will be created and WandB will not be initialized. "
+ "This argument is mainly meant for debugging purposes as well as being able to see the number of images in "
+ "the training and validation datasets without starting the training process.",
+ )
+
+ parser.add_argument(
+ "--override_directory_check",
+ action="store_true",
+ help="Override the OSError caused by creating a new model directory that already exists. Normally, if the script "
+ "crashes before or during the training of a new model, an OSError will be returned if the script is immediately "
+ "ran again with the same model number as the model directory already exists. This is an intentional fail-safe "
+ "designed to prevent models that already exist from being overwritten. Passing this boolean flag disables the "
+ "fail-safe and can be useful if the script is being ran on a workload manager (e.g. SLURM) where jobs can fail "
+ "and then be immediately requeued and ran again.",
+ )
+
+ args = vars(parser.parse_args())
+ # --------------------- Seed in train.py
+ """
+ # Set the random seed. After a model is trained, the same seed will be used in subsequent retraining sessions for the same model.
+ seed = (
+ pd.read_pickle(
+ "%s/model_%d/model_%d_properties.pkl"
+ % (args["model_dir"], args["model_number"], args["model_number"])
+ )["seed"]
+ if args["retrain"]
+ else args["seed"]
+ )
+ tf.keras.utils.set_random_seed(seed)
+ """
+ # --------------------- paths handled by data modules
+ """
+ assert len(args["tf_indirs"]) < 3, (
+ "Only 1 or 2 paths can be passed into --tf_indirs, received %d paths"
+ % len(args["tf_indirs"])
+ )
+ if len(args["tf_indirs"]) == 1:
+ args["tf_indirs"].append(args["tf_indirs"][0])
+ """
+ # --------------------- dataset properties handled dynamically with config
+ """
+ train_dataset_properties = pd.read_pickle(
+ "%s/dataset_properties.pkl" % args["tf_indirs"][0]
+ )
+ valid_dataset_properties = pd.read_pickle(
+ "%s/dataset_properties.pkl" % args["tf_indirs"][1]
+ )
+ """
+
+ # --------------------- shuffle always happens regardless of arg
+ """ if args["shuffle"] != "lazy" and args["shuffle"] != "full":
+ raise ValueError(
+ "Unrecognized shuffling method: %s. Valid methods are 'lazy' or 'full'"
+ % args["shuffle"]
+ )"""
+
+ # --------------------- check irrelevant
+ """ # Check arguments that can only have a maximum length of 2
+ for arg in [
+ "loss",
+ "metric",
+ "optimizer",
+ "activity_regularizer",
+ "bias_constraint",
+ "bias_regularizer",
+ "kernel_constraint",
+ "kernel_regularizer",
+ "batch_size",
+ "steps",
+ ]:
+ if args[arg] is None: # need this line in here because 'steps' can be None
+ continue
+ elif len(args[arg]) > 2:
+ raise ValueError("--%s can only take up to 2 arguments" % arg)"""
+
+ # --------------------- applied within model building functions
+
+ """ ### Dictionary containing arguments that cannot be used for specific model types ###
+ incompatible_args = {
+ "unet": dict(deep_supervision=False, first_encoder_connections=False),
+ "unet_ensemble": dict(deep_supervision=False, first_encoder_connections=False),
+ "unet_plus": dict(first_encoder_connections=False),
+ "unet_2plus": dict(first_encoder_connections=False),
+ "unet_3plus": {},
+ "attention_unet": dict(
+ upsample_size=None, deep_supervision=False, first_encoder_connections=False
+ ),
+ }
+
+ ### Make sure that incompatible arguments were not passed, and raise errors if they were passed ###
+ incompatible_args_for_model = incompatible_args[args["model_type"]]
+ for arg in incompatible_args_for_model:
+ if incompatible_args_for_model[arg] != args[arg]:
+ raise ValueError(
+ f"--{arg} must be '{incompatible_args_for_model[arg]}' when the model type is {args['model_type']}"
+ )
+ """
+
+ # --------------------- everything here should be dataclasses steered by yaml
+ ### Convert keyword argument strings to dictionaries ###
+ """
+ loss_args = (
+ misc.string_arg_to_dict(args["loss"][1]) if len(args["loss"]) > 1 else dict()
+ )
+ metric_args = (
+ misc.string_arg_to_dict(args["metric"][1])
+ if len(args["metric"]) > 1
+ else dict()
+ )
+ optimizer_args = (
+ misc.string_arg_to_dict(args["optimizer"][1])
+ if len(args["optimizer"]) > 1
+ else dict()
+ )
+ activity_regularizer_args = (
+ misc.string_arg_to_dict(args["activity_regularizer"][1])
+ if len(args["activity_regularizer"]) > 1
+ else dict()
+ )
+ bias_constraint_args = (
+ misc.string_arg_to_dict(args["bias_constraint"][1])
+ if len(args["bias_constraint"]) > 1
+ else dict()
+ )
+ bias_regularizer_args = (
+ misc.string_arg_to_dict(args["bias_regularizer"][1])
+ if len(args["bias_regularizer"]) > 1
+ else dict()
+ )
+ kernel_constraint_args = (
+ misc.string_arg_to_dict(args["kernel_constraint"][1])
+ if len(args["kernel_constraint"]) > 1
+ else dict()
+ )
+ kernel_regularizer_args = (
+ misc.string_arg_to_dict(args["kernel_regularizer"][1])
+ if len(args["kernel_regularizer"]) > 1
+ else dict()
+ )
+ """
+
+ # learning rate is part of the optimizer
+ if args["learning_rate"] is not None:
+ optimizer_args["learning_rate"] = args["learning_rate"]
+
+ gpus = tf.config.list_physical_devices(device_type="GPU") # Find available GPUs
+ if len(gpus) > 0:
+ print("Number of GPUs available: %d" % len(gpus))
+
+ # Only make the selected GPU(s) visible to TensorFlow
+ if args["gpu_device"] is not None:
+ tf.config.set_visible_devices(
+ devices=[gpus[gpu] for gpu in args["gpu_device"]], device_type="GPU"
+ )
+ gpus = tf.config.get_visible_devices(
+ device_type="GPU"
+ ) # List of selected GPUs
+ print("Using %d GPU(s):" % len(gpus), gpus)
+
+ # Disable TensorFloat32 for matrix multiplication
+ if args["disable_tensorfloat32"]:
+ tf.config.experimental.enable_tensor_float_32_execution(False)
+
+ # Allow for memory growth on the GPU. This will only use the GPU memory that is required rather than allocating all the GPU's memory.
+ if args["memory_growth"]:
+ tf.config.experimental.set_memory_growth(
+ device=[gpu for gpu in gpus][0], enable=True
+ )
+
+ else:
+ print("WARNING: No GPUs found, all computations will be performed on CPUs.")
+ tf.config.set_visible_devices([], "GPU")
+
+ """
+ Verify that the training and validation datasets have the same front types, variables, pressure levels, number of
+ dimensions, and normalization parameters.
+ """
+ if args["tf_indirs"][0] != args["tf_indirs"][1]:
+ assert (
+ train_dataset_properties["front_types"]
+ == valid_dataset_properties["front_types"]
+ ), (
+ f"The front types in the training and validation datasets must be the same! Received {train_dataset_properties['front_types']} "
+ f"for training, {valid_dataset_properties['front_types']} for validation."
+ )
+ assert (
+ train_dataset_properties["variables"]
+ == valid_dataset_properties["variables"]
+ ), (
+ f"The variables in the training and validation datasets must be the same! Received {train_dataset_properties['variables']} "
+ f"for training, {valid_dataset_properties['variables']} for validation."
+ )
+ assert (
+ train_dataset_properties["pressure_levels"]
+ == valid_dataset_properties["pressure_levels"]
+ ), (
+ f"The pressure levels in the training and validation datasets must be the same! Received {train_dataset_properties['pressure_levels']} "
+ f"for training, {valid_dataset_properties['pressure_levels']} for validation."
+ )
+ assert all(
+ train_dataset_properties["num_dims"][num]
+ == valid_dataset_properties["num_dims"][num]
+ for num in range(2)
+ ), (
+ f"The number of dimensions for the inputs and targets in the training and validation datasets must be the same! Received {train_dataset_properties['num_dims']} "
+ f"for training, {valid_dataset_properties['num_dims']} for validation"
+ )
+ assert (
+ train_dataset_properties["normalization_parameters"]
+ == valid_dataset_properties["normalization_parameters"]
+ ), (
+ "Normalization parameters for the training and validation datasets must be the same!"
+ )
+
+ front_types = train_dataset_properties["front_types"]
+ variables = train_dataset_properties["variables"]
+ pressure_levels = train_dataset_properties["pressure_levels"]
+ image_size = train_dataset_properties["image_size"]
+ num_dims = train_dataset_properties["num_dims"]
+
+ if not args["retrain"]:
+ all_years = np.arange(2008, 2024.1, 1)
+
+ if args["num_training_years"] is not None:
+ if args["training_years"] is not None:
+ raise TypeError(
+ "Cannot explicitly declare the training years if --num_training_years is passed"
+ )
+ training_years = list(
+ sorted(
+ np.random.choice(
+ all_years, args["num_training_years"], replace=False
+ )
+ )
+ )
+ else:
+ if args["training_years"] is None:
+ raise TypeError(
+ "Must pass one of the following arguments: --training_years, --num_training_years"
+ )
+ training_years = list(sorted(args["training_years"]))
+
+ if args["num_validation_years"] is not None:
+ if args["validation_years"] is not None:
+ raise TypeError(
+ "Cannot explicitly declare the validation years if --num_validation_years is passed"
+ )
+ validation_years = list(
+ sorted(
+ np.random.choice(
+ [year for year in all_years if year not in training_years],
+ args["num_validation_years"],
+ replace=False,
+ )
+ )
+ )
+ else:
+ if args["validation_years"] is None:
+ raise TypeError(
+ "Must pass one of the following arguments: --validation_years, --num_validation_years"
+ )
+ validation_years = list(sorted(args["validation_years"]))
+
+ if len(training_years) + len(validation_years) > len(all_years) - 1:
+ raise ValueError(
+ "No testing years are available: the total number of training and validation years cannot be greater than 15"
+ )
+
+ test_years = [
+ year for year in all_years if year not in training_years + validation_years
+ ]
+
+ # If no model number was provided, select a number based on the current date and time. This number changes once per minute.
+ args["model_number"] = (
+ int(datetime.datetime.utcnow().timestamp() % 1e8 / 60)
+ if args["model_number"] is None
+ else args["model_number"]
+ )
+
+ # Convert pool size and upsample size to tuples
+ pool_size = tuple(args["pool_size"]) if args["pool_size"] is not None else None
+ upsample_size = (
+ tuple(args["upsample_size"]) if args["upsample_size"] is not None else None
+ )
+
+ if any(
+ front_type == front_types
+ for front_type in [["MERGED-F_BIN"], ["MERGED-T"], ["F_BIN"]]
+ ):
+ num_classes = 2
+ elif front_types == ["MERGED-F"]:
+ num_classes = 5
+ elif front_types == ["MERGED-ALL"]:
+ num_classes = 8
+ else:
+ num_classes = len(front_types) + 1
+
+ # Create dictionary containing information about the model. This simplifies the process of loading the model
+ model_properties = dict({})
+ model_properties["domains"] = [
+ train_dataset_properties["domain"],
+ valid_dataset_properties["domain"],
+ ]
+ model_properties["dataset_properties"] = train_dataset_properties
+ model_properties["classes"] = num_classes
+
+ # Place provided arguments into the model properties dictionary
+ for arg in [
+ "model_type",
+ "learning_rate",
+ "deep_supervision",
+ "model_number",
+ "kernel_size",
+ "modules_per_node",
+ "activation",
+ "batch_normalization",
+ "padding",
+ "use_bias",
+ "activity_regularizer",
+ "bias_constraint",
+ "bias_initializer",
+ "bias_regularizer",
+ "kernel_constraint",
+ "kernel_initializer",
+ "kernel_regularizer",
+ "first_encoder_connections",
+ "valid_freq",
+ "optimizer",
+ "seed",
+ ]:
+ model_properties[arg] = args[arg]
+ model_properties["activation"] = model_properties["activation"].lower()
+
+ # Place local variables into the model properties dictionary
+ for arg in [
+ "loss_args",
+ "metric_args",
+ "image_size",
+ "training_years",
+ "validation_years",
+ "test_years",
+ ]:
+ model_properties[arg] = locals()[arg]
+
+ # If using 3D inputs and 2D targets, squeeze out the vertical dimension of the model (index 3)
+ squeeze_axes = 3 if num_dims[0] == 3 and num_dims[1] == 2 else None
+
+ unet_model = getattr(unets, args["model_type"])
+ unet_model_args = unet_model.__code__.co_varnames[
+ : unet_model.__code__.co_argcount
+ ] # pull argument names from unet function
+
+ ### Arguments for the function used to build the U-Net ###
+ unet_kwargs = {
+ arg: args[arg]
+ for arg in [
+ "pool_size",
+ "upsample_size",
+ "levels",
+ "filter_num",
+ "kernel_size",
+ "modules_per_node",
+ "activation",
+ "batch_normalization",
+ "padding",
+ "use_bias",
+ "bias_initializer",
+ "kernel_initializer",
+ "first_encoder_connections",
+ "deep_supervision",
+ ]
+ if arg in unet_model_args
+ }
+ unet_kwargs["squeeze_axes"] = squeeze_axes
+ unet_kwargs["activity_regularizer"] = (
+ getattr(tf.keras.regularizers, args["activity_regularizer"][0])(
+ **activity_regularizer_args
+ )
+ if args["activity_regularizer"][0] is not None
+ else None
+ )
+ unet_kwargs["bias_constraint"] = (
+ getattr(tf.keras.constraints, args["bias_constraint"][0])(
+ **bias_constraint_args
+ )
+ if args["bias_constraint"][0] is not None
+ else None
+ )
+ unet_kwargs["kernel_constraint"] = (
+ getattr(tf.keras.constraints, args["kernel_constraint"][0])(
+ **kernel_constraint_args
+ )
+ if args["kernel_constraint"][0] is not None
+ else None
+ )
+ unet_kwargs["bias_regularizer"] = (
+ getattr(tf.keras.regularizers, args["bias_regularizer"][0])(
+ **bias_regularizer_args
+ )
+ if args["bias_regularizer"][0] is not None
+ else None
+ )
+ unet_kwargs["kernel_regularizer"] = (
+ getattr(tf.keras.regularizers, args["kernel_regularizer"][0])(
+ **kernel_regularizer_args
+ )
+ if args["kernel_regularizer"][0] is not None
+ else None
+ )
+
+ print("Training years:", training_years)
+ print("Validation years:", validation_years)
+ print("Test years:", test_years)
+
+ else:
+ model_properties = pd.read_pickle(
+ "%s/model_%d/model_%d_properties.pkl"
+ % (args["model_dir"], args["model_number"], args["model_number"])
+ )
+ front_types = model_properties["front_types"]
+
+ training_years = model_properties["training_years"]
+ validation_years = model_properties["validation_years"]
+ test_years = model_properties["test_years"]
+
+ train_batch_size, valid_batch_size = model_properties["batch_sizes"]
+ train_steps, valid_steps = model_properties["steps_per_epoch"]
+ valid_freq = model_properties["valid_freq"]
+
+ model_filepath = "%s/model_%d/model_%d.h5" % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ ) # filepath for the actual model (.h5 file)
+ history_filepath = "%s/model_%d/model_%d_history.csv" % (
+ args["model_dir"],
+ args["model_number"],
+ args["model_number"],
+ ) # path of the CSV file containing loss and metric statistics
+
+ if train_dataset_properties["domain"] in ["conus", "full"]:
+ train_data_source = "era5"
+ elif train_dataset_properties["domain"] == "global":
+ train_data_source = "gfs"
+ else:
+ train_data_source = train_dataset_properties["domain"]
+
+ train_batch_size = args["batch_size"][0]
+ valid_batch_size = (
+ args["batch_size"][0] if len(args["batch_size"]) == 1 else args["batch_size"][1]
+ )
+ model_properties["batch_sizes"] = [train_batch_size, valid_batch_size]
+
+ ### Training dataset ###
+ training_files = []
+ for year in training_years:
+ training_files.extend(list(sorted(glob(f"{args['tf_indirs'][0]}/{year}-*_tf"))))
+ np.random.shuffle(training_files) # shuffle the order of the slices
+ print("Training slices:", len(training_files))
+ training_dataset = data_utils.combine_datasets(training_files)
+ images_in_training_dataset = len(training_dataset)
+ training_dataset = training_dataset.batch(
+ train_batch_size,
+ drop_remainder=False,
+ num_parallel_calls=args["num_parallel_calls"],
+ )
+ training_dataset = training_dataset.prefetch(tf.data.AUTOTUNE)
+
+ if valid_dataset_properties["domain"] in ["conus", "full"]:
+ valid_data_source = "era5"
+ elif valid_dataset_properties["domain"] == "global":
+ valid_data_source = "gfs"
+ else:
+ valid_data_source = valid_dataset_properties["domain"]
+
+ ### Validation dataset ###
+ validation_files = []
+ for year in validation_years:
+ validation_files.extend(
+ list(sorted(glob(f"{args['tf_indirs'][0]}/{year}-*_tf")))
+ )
+ validation_dataset = data_utils.combine_datasets(validation_files)
+ images_in_validation_dataset = len(validation_dataset)
+ validation_dataset = validation_dataset.batch(
+ valid_batch_size,
+ drop_remainder=True,
+ num_parallel_calls=args["num_parallel_calls"],
+ )
+ validation_dataset = validation_dataset.prefetch(tf.data.AUTOTUNE)
+
+ """
+ If the number of training steps is not passed, the number of training steps will be determined by the dataset size
+ and the provided batch size. The calculated number of training steps will allow for one complete pass over the training
+ dataset during each epoch.
+ """
+ if args["steps"] is None:
+ train_steps = int(images_in_training_dataset / train_batch_size) + 1
+ print("Using %d training steps per epoch" % train_steps)
+ valid_steps = None
+ else:
+ train_steps = args["steps"][0]
+ valid_steps = None if len(args["steps"]) < 2 else args["steps"][1]
+
+ valid_freq = args["valid_freq"]
+ model_properties["steps_per_epoch"] = [train_steps, valid_steps]
+
+ """
+ If the patience argument is not explicitly provided, derive it from the size of the training dataset along with the
+ batch size and number of steps per epoch.
+ """
+ if args["patience"] is None:
+ patience = (
+ int(images_in_training_dataset / (train_batch_size * train_steps)) + 1
+ )
+ print("Using patience value of %d epochs for early stopping" % patience)
+ else:
+ patience = args["patience"]
+
+ # Set the lat/lon dimensions to have a None shape so images of any size can be passed into the model
+ input_shape = list(training_dataset.take(0).element_spec[0].shape[1:])
+ for i in range(2):
+ input_shape[i] = None
+ input_shape = tuple(input_shape)
+
+ with tf.distribute.MirroredStrategy().scope():
+ if not args["retrain"]:
+ model = unet_model(input_shape, num_classes, **unet_kwargs)
+ loss_function = getattr(losses, args["loss"][0])(**loss_args)
+ metric_function = getattr(metrics, args["metric"][0])(**metric_args)
+ optimizer = getattr(tf.keras.optimizers, args["optimizer"][0])(
+ **optimizer_args
+ )
+ model.compile(
+ loss=loss_function, optimizer=optimizer, metrics=[metric_function]
+ )
+
+ model_properties["loss_parent_string"] = args["loss"][0]
+ model_properties["loss_child_string"] = loss_function.function_spec._name
+ model_properties["metric_parent_string"] = args["metric"][0]
+ model_properties["metric_child_string"] = (
+ metric_function.function_spec._name
+ )
+ else:
+ model = file_manager.load_model(args["model_number"], args["model_dir"])
+
+ model.summary()
+
+ if not args["retrain"] and not args["no_train"]:
+ model_properties = {
+ key: model_properties[key] for key in sorted(model_properties.keys())
+ } # Sort model properties dictionary alphabetically
+
+ if not os.path.isdir("%s/model_%d" % (args["model_dir"], args["model_number"])):
+ os.makedirs(
+ "%s/model_%d/maps" % (args["model_dir"], args["model_number"])
+ ) # Make folder for model predicton maps
+ os.mkdir(
+ "%s/model_%d/probabilities" % (args["model_dir"], args["model_number"])
+ ) # Make folder for prediction data files
+ os.mkdir(
+ "%s/model_%d/statistics" % (args["model_dir"], args["model_number"])
+ ) # Make folder for statistics data files
+ elif not args["override_directory_check"]:
+ raise OSError(
+ "%s/model_%d already exists. If model %d still needs to be created and trained, run this script "
+ "again with the --override_directory_check flag."
+ % (args["model_dir"], args["model_number"], args["model_number"])
+ )
+ elif os.path.isfile(model_filepath):
+ raise OSError(
+ "model %d already exists at %s. Choose a different model number and try again."
+ % (args["model_number"], model_filepath)
+ )
+
+ with open(
+ "%s/model_%d/model_%d_properties.pkl"
+ % (args["model_dir"], args["model_number"], args["model_number"]),
+ "wb",
+ ) as f:
+ pickle.dump(model_properties, f)
+
+ with open(
+ "%s/model_%d/model_%d_properties.txt"
+ % (args["model_dir"], args["model_number"], args["model_number"]),
+ "w",
+ ) as f:
+ for key in model_properties.keys():
+ f.write(f"{key}: {model_properties[key]}\n")
+
+ checkpoint = tf.keras.callbacks.ModelCheckpoint(
+ model_filepath,
+ monitor="val_loss",
+ verbose=1,
+ save_best_only=True,
+ save_weights_only=False,
+ save_freq="epoch",
+ ) # ModelCheckpoint: saves model at a specified interval
+ early_stopping = EarlyStopping(
+ "val_loss", patience=patience, verbose=1
+ ) # EarlyStopping: stops training early if the validation loss does not improve after a specified number of epochs (patience)
+ history_logger = CSVLogger(
+ history_filepath, separator=",", append=True
+ ) # Saves loss and metric data every epoch
+
+ callbacks = [early_stopping, checkpoint, history_logger]
+
+ if not args["no_train"]:
+ ### Initialize WandB ###
+ if args["project"] is not None:
+ wandb_init_config = dict(
+ {
+ key: model_properties[key]
+ for key in [
+ "activation",
+ "activity_regularizer",
+ "batch_normalization",
+ "batch_sizes",
+ "bias_constraint",
+ "bias_initializer",
+ "bias_regularizer",
+ "domains",
+ "image_size",
+ "kernel_constraint",
+ "kernel_initializer",
+ "kernel_regularizer",
+ "kernel_size",
+ "learning_rate",
+ "loss_parent_string",
+ "metric_parent_string",
+ "model_number",
+ "model_type",
+ "modules_per_node",
+ "optimizer",
+ "padding",
+ "steps_per_epoch",
+ "test_years",
+ "training_years",
+ "use_bias",
+ "validation_years",
+ ]
+ }
+ )
+
+ # add keys from dataset_properties dictionary
+ for key in [
+ "variables",
+ "pressure_levels",
+ "domain",
+ "timestep_fraction",
+ "image_fraction",
+ "flip_chance_lon",
+ "flip_chance_lat",
+ "front_dilation",
+ ]:
+ wandb_init_config[key] = model_properties["dataset_properties"][key]
+
+ wandb_init_name = (
+ "model_%d" % args["model_number"]
+ if args["name"] is None
+ else args["name"]
+ )
+
+ if args["key"] is not None:
+ wandb.login(key=args["key"])
+
+ wandb.init(
+ project=args["project"], config=wandb_init_config, name=wandb_init_name
+ )
+ callbacks.append(wandb.keras.WandbMetricsLogger(log_freq=args["log_freq"]))
+
+ if args["upload_model"]:
+ callbacks.append(
+ wandb.keras.WandbModelCheckpoint("models")
+ ) # upload model checkpoints to wandb
+
+ model.fit(
+ training_dataset.repeat(),
+ validation_data=validation_dataset,
+ validation_freq=valid_freq,
+ epochs=args["epochs"],
+ steps_per_epoch=train_steps,
+ validation_steps=valid_steps,
+ callbacks=callbacks,
+ verbose=args["verbose"],
+ )
+
+ wandb.finish()
+
+ else:
+ print(
+ "NOTE: Remove the --no_train argument from the command line to start the training process."
+ )
diff --git a/utils/__init__.py b/src/fronts/utils/__init__.py
similarity index 100%
rename from utils/__init__.py
rename to src/fronts/utils/__init__.py
diff --git a/utils/variables.py b/src/fronts/utils/calc.py
similarity index 56%
rename from utils/variables.py
rename to src/fronts/utils/calc.py
index 42a20a4..c6191d6 100644
--- a/utils/variables.py
+++ b/src/fronts/utils/calc.py
@@ -3,80 +3,108 @@
References
----------
-* Bolton 1980: https://doi.org/10.1175/1520-0493(1980)108<1046:TCOEPT>2.0.CO;2
* Davies-Jones 2008: https://doi.org/10.1175/2007MWR2224.1
* Stull 2011: https://doi.org/10.1175/JAMC-D-11-0143.1
-* Vasaila 2013: https://www.vaisala.com/sites/default/files/documents/Humidity_Conversion_Formulas_B210973EN-F.pdf
Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.9.21
+Script version: 2025.5.3
"""
+
import numpy as np
-from utils import data_utils
+from fronts.utils import data_utils
import tensorflow as tf
+import xarray as xr
+from typing import Callable
Rd = 287.04 # Gas constant for dry air (J/kg/K)
Rv = 461.5 # Gas constant for water vapor (J/kg/K)
Cpd = 1005.7 # Specific heat of dry air at constant pressure (J/kg/K)
-Cpw = 4184 # Specific heat capacity of liquid water (J/kg/K)
kd = Rd / Cpd # Exponential constant for potential temperature
epsilon = Rd / Rv
-P_knot = 1e5 # Pa
e_knot = 611.2 # Pa
Lv = 2.257e6 # Latent heat of vaporization for water vapor (J/kg)
-NA = 6.02214076e23 # Avogrado constant (mol^-1)
-kB = 1.380649e-23 # Boltzmann constant (J/K)
-def absolute_humidity(T: int | float | np.ndarray | tf.Tensor,
- Td: int | float | np.ndarray | tf.Tensor):
+def saturation_vapor_pressure(T):
"""
- Calculates absolute humidity from temperature and dewpoint temperature.
+ Calculates saturation vapor pressure from temperature.
Parameters
----------
T: float or iterable object
Air temperature expressed as kelvin (K).
- Td: float or iterable object
- Dewpoint temperature expressed as kelvin (K).
Returns
-------
- AH: float or iterable object
- Absolute humidity expressed as kilograms of water vapor per cubic meter of air (kg/m^3).
+ e: float or iterable object
+ Vapor pressure expressed as Pascals (Pa).
+ """
+
+ exp_func = tf.math.exp if tf.is_tensor(T) else np.exp
+ T = T - 273.15 # K -> C
+ e = e_knot * exp_func(17.67 * T / (T + 243.5))
+ return e
+
+
+def dewpoint_from_mixing_ratio(P, r):
+ """
+ Calculates dewpoint temperature from mixing ratio, and pressure.
+
+ Parameters
+ ----------
+ P: float or iterable object
+ Air pressure expressed as pascals (Pa).
+ r: float or iterable object
+ Mixing ratio expressed as grams of water vapor per gram of dry air (unitless; g/g or kg/kg).
+
+ Returns
+ -------
+ Td: float or iterable object
+ Dewpoint temperature expressed as kelvin (K).
Examples
--------
- >>> T = 300 # K
- >>> Td = 290 # K
- >>> AH = absolute_humidity(T, Td) # kg * m^-3
- >>> AH
- 0.012493639535490526
+ >>> P = 1e5 # Pa
+ >>> r = 20 / 1000 # g/kg -> kg/kg
+ >>> Td = dewpoint_from_mixing_ratio(P, r)
+ >>> Td
+ 297.87277930360676
- >>> T = np.arange(270, 311, 5) # K
- >>> T
- array([270, 275, 280, 285, 290, 295, 300, 305, 310])
- >>> Td = np.arange(260, 301, 5) # K
+ >>> P = np.arange(800.0, 1001.0, 25) * 100 # Pa
+ >>> P
+ array([ 80000., 82500., 85000., 87500., 90000., 92500., 95000.,
+ 97500., 100000.])
+ >>> r = np.arange(5, 25.01, 2.5) / 1000 # g/kg -> kg/kg
+ >>> r
+ array([0.005 , 0.0075, 0.01 , 0.0125, 0.015 , 0.0175, 0.02 , 0.0225,
+ 0.025 ])
+ >>> Td = dewpoint_from_mixing_ratio(P, r)
>>> Td
- array([260, 265, 270, 275, 280, 285, 290, 295, 300])
- >>> AH = absolute_humidity(T, Td) # kg * m^-3
- >>> AH
- array([0.00198323, 0.00277676, 0.00383828, 0.00524175, 0.00707688,
- 0.00945143, 0.01249364, 0.0163548 , 0.02121195])
+ array([273.74253532, 279.8787204 , 284.52674244, 288.32976779,
+ 291.58260633, 294.44604583, 297.01785056, 299.36210297,
+ 301.52319595])
"""
- e = vapor_pressure(Td) # vapor pressure expressed as pascals
- AH = e / (Rv * T) # absolute humidity
- return AH
+ log_func = tf.math.log if tf.is_tensor(P) else np.log
+
+ P /= 100 # Pa -> hPa
+
+ e = r * P / (epsilon + r) # vapor pressure
+
+ Td = 243.5 * log_func(e / 6.112) / (17.67 - log_func(e / 6.112)) + 273.15 # C -> K
+ return Td
-def dewpoint_from_vapor_pressure(vapor_pressure: int | float | np.ndarray | tf.Tensor):
+
+def dewpoint_from_relative_humidity(T, RH):
"""
- Calculates dewpoint temperature from vapor pressure.
+ Calculates dewpoint temperature from relative humidity and temperature.
Parameters
----------
- vapor_pressure: float or iterable object
- Vapor pressure expressed as pascals (Pa).
+ T: float or iterable object
+ Air temperature expressed as kelvin (K).
+ RH: float or iterable object
+ Relative humidity expressed as a decimal (e.g. 57% = 0.57).
Returns
-------
@@ -85,37 +113,49 @@ def dewpoint_from_vapor_pressure(vapor_pressure: int | float | np.ndarray | tf.T
Examples
--------
- >>> vap_pres = 1000 # Pa
- >>> Td = dewpoint_from_vapor_pressure(vap_pres)
+ >>> P = 1e5 # Pa
+ >>> T = 300 # K
+ >>> RH = 0.8
+ >>> Td = dewpoint_from_relative_humidity(T, RH)
>>> Td
- 280.8734119482131
+ 296.2618676251417
- >>> vap_pres = np.arange(500, 4001, 500) # Pa
- >>> vap_pres
- array([ 500, 1000, 1500, 2000, 2500, 3000, 3500, 4000])
- >>> Td = dewpoint_from_vapor_pressure(vap_pres)
+ >>> P = np.arange(800.0, 1001.0, 25) * 100 # Pa
+ >>> P
+ array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
+ 100000])
+ >>> T = np.arange(270.0, 311.0, 5) # K
+ >>> T
+ array([270, 275, 280, 285, 290, 295, 300, 305, 310])
+ >>> RH = np.arange(0.5, 0.91, 0.05)
+ >>> RH
+ array([0.5 , 0.55, 0.6 , 0.65, 0.7 , 0.75, 0.8 , 0.85, 0.9 ])
+ >>> Td = dewpoint_from_relative_humidity(T, RH)
>>> Td
- array([270.12031682, 280.87341195, 287.56990934, 292.51813134,
- 296.4751267 , 299.7885851 , 302.64840758, 305.170169 ])
+ array([261.04058296, 266.91163292, 272.7737636 , 278.63451989,
+ 284.499794 , 290.37429558, 296.26186763, 302.16570498,
+ 308.08850911])
"""
- Td = (1/273.15 - (tf.math.log(vapor_pressure/e_knot)*(Rv/Lv))) ** -1 if tf.is_tensor(vapor_pressure) else \
- (1/273.15 - (np.log(vapor_pressure/e_knot)*(Rv/Lv))) ** -1
+
+ log_func = tf.math.log if tf.is_tensor(T) else np.log
+
+ # T is converted to C in saturation_vapor_pressure function
+ e = saturation_vapor_pressure(T) * RH
+
+ Td = (
+ 243.5 * log_func(e / e_knot) / (17.67 - log_func(e / e_knot)) + 273.15
+ ) # C -> K
return Td
-def dewpoint_from_specific_humidity(
- P: int | float | np.ndarray | tf.Tensor,
- T: int | float | np.ndarray | tf.Tensor,
- q: int | float | np.ndarray | tf.Tensor):
+def dewpoint_from_specific_humidity(P, q):
"""
- Calculates dewpoint temperature from specific humidity, pressure, and temperature.
+ Calculates dewpoint temperature from specific humidity, and pressure.
Parameters
----------
P: float or iterable object
Air pressure expressed as pascals (Pa).
- T: float or iterable object
- Air temperature expressed as kelvin (K).
q: float or iterable object
Specific humidity expressed as grams of water vapor per gram of dry air (unitless; g/g or kg/kg).
@@ -127,13 +167,12 @@ def dewpoint_from_specific_humidity(
Examples
--------
>>> P = 1e5 # Pa
- >>> T = 300 # K
>>> q = 20 / 1000 # g/kg -> kg/kg
- >>> Td = dewpoint_from_specific_humidity(P, T, q)
+ >>> Td = dewpoint_from_specific_humidity(P, q)
>>> Td
- 298.199585429495
+ 298.20035572272803
- >>> P = np.arange(800, 1001, 25) * 100 # Pa
+ >>> P = np.arange(800.0, 1001.0, 25) * 100 # Pa
>>> P
array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
100000])
@@ -141,45 +180,34 @@ def dewpoint_from_specific_humidity(
>>> q
array([0.005 , 0.0075, 0.01 , 0.0125, 0.015 , 0.0175, 0.02 , 0.0225,
0.025 ])
- >>> T = np.arange(270, 311, 5) # K
- >>> T
- array([270, 275, 280, 285, 290, 295, 300, 305, 310])
- >>> Td = dewpoint_from_specific_humidity(P, T, q)
+ >>> Td = dewpoint_from_specific_humidity(P, q)
>>> Td
- array([273.80033167, 279.97353235, 284.66312436, 288.51045188,
- 291.80950762, 294.72063163, 297.34130998, 299.7354364 ,
- 301.94726732])
+ array([273.81141134, 279.98701245, 284.67615887, 288.52165869,
+ 291.8180981 , 294.72610931, 297.34334081, 299.73378468,
+ 301.94176072])
"""
- # Constants needed to perform dewpoint calculation (Vasaila 2013)
- m = 7.591386
- A = 6.116441
- Tn = 240.7263
- m1 = 9.778707
- A1 = 6.114742
- Tn1 = 273.1466
+ log_func = tf.math.log if tf.is_tensor(P) else np.log
- vap_pres = (P * q) / (0.622 + 0.378 * q) # (Bolton 1980, Eq. 16) expressed as pascals (Pa)
- vap_pres /= 100 # convert to hPa
+ r = q / (1 - q) # mixing ratio
+ e = r * P / (epsilon + r) # vapor pressure
- # Dewpoint calculation from Vasaila 2013
- Td_gt = Tn / ((m / (np.log10(vap_pres / A))) - 1)
- Td_lt = Tn1 / ((m1 / (np.log10(vap_pres / A1))) - 1)
- Td = np.where(T >= 0, Td_gt, Td_lt)
- return Td + 273.15
+ Td = (
+ 243.5 * log_func(e / e_knot) / (17.67 - log_func(e / e_knot)) + 273.15
+ ) # C -> K
+ return Td
-def mixing_ratio_from_dewpoint(Td: int | float | np.ndarray | tf.Tensor,
- P: int | float | np.ndarray | tf.Tensor):
+def mixing_ratio_from_dewpoint(P, Td):
"""
Calculates mixing ratio from dewpoint temperature and air pressure.
Parameters
----------
- Td: float or iterable object
- Dewpoint temperature expressed as kelvin (K).
P: float or iterable object
Air pressure expressed as pascals (Pa).
+ Td: float or iterable object
+ Dewpoint temperature expressed as kelvin (K).
Returns
-------
@@ -188,40 +216,39 @@ def mixing_ratio_from_dewpoint(Td: int | float | np.ndarray | tf.Tensor,
Examples
--------
- >>> Td = 290 # K
>>> P = 1e5 # Pa
- >>> r = mixing_ratio_from_dewpoint(Td, P)
+ >>> Td = 290 # K
+ >>> r = mixing_ratio_from_dewpoint(P, Td)
>>> r
0.010947893449979635
- >>> Td = np.arange(260, 301, 5) # K
- >>> Td
- array([260, 265, 270, 275, 280, 285, 290, 295, 300])
>>> P = np.arange(800, 1001, 25) * 100 # Pa
>>> P
array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
100000])
- >>> r = mixing_ratio_from_dewpoint(Td, P)
+ >>> Td = np.arange(260, 301, 5) # K
+ >>> Td
+ array([260, 265, 270, 275, 280, 285, 290, 295, 300])
+ >>> r = mixing_ratio_from_dewpoint(P, Td)
>>> r
array([0.00192723, 0.0026682 , 0.00365056, 0.00493959, 0.00661507,
0.00877413, 0.01153478, 0.01504042, 0.01946563])
"""
- e = vapor_pressure(Td)
+ e = saturation_vapor_pressure(Td)
r = epsilon * e / (P - e) # mixing ratio
return r
-def potential_temperature(T: int | float | np.ndarray | tf.Tensor,
- P: int | float | np.ndarray | tf.Tensor):
+def potential_temperature(P, T):
"""
Returns potential temperature expressed as kelvin (K).
Parameters
----------
- T: float or iterable object
- Air temperature expressed as kelvin (K).
P: float or iterable object
Air pressure expressed as pascals (Pa).
+ T: float or iterable object
+ Air temperature expressed as kelvin (K).
Returns
-------
@@ -230,31 +257,34 @@ def potential_temperature(T: int | float | np.ndarray | tf.Tensor,
Examples
--------
- >>> T = 275 # K
>>> P = 9e4 # Pa
- >>> theta = potential_temperature(T, P)
+ >>> T = 275 # K
+ >>> theta = potential_temperature(P, T)
>>> theta
283.3951954331142
- >>> T = np.arange(270, 311, 5) # K
- >>> T
- array([270, 275, 280, 285, 290, 295, 300, 305, 310])
>>> P = np.arange(800, 1001, 25) * 100 # Pa
>>> P
array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
100000])
- >>> theta = potential_temperature(T, P)
+ >>> T = np.arange(270, 311, 5) # K
+ >>> T
+ array([270, 275, 280, 285, 290, 295, 300, 305, 310])
+ >>> theta = potential_temperature(P, T)
>>> theta
array([287.75518363, 290.52120387, 293.2937428 , 296.07144546,
298.85311518, 301.63769299, 304.42424008, 307.21192283,
310. ])
"""
- theta = T * tf.pow(1e5 / P, kd) if tf.is_tensor(T) and tf.is_tensor(P) else T * np.power(1e5 / P, kd)
+ theta = (
+ T * tf.pow(1e5 / P, kd)
+ if tf.is_tensor(T) and tf.is_tensor(P)
+ else T * np.power(1e5 / P, kd)
+ )
return theta
-def relative_humidity(T: int | float | np.ndarray | tf.Tensor,
- Td: int | float | np.ndarray | tf.Tensor):
+def relative_humidity_from_dewpoint(T, Td):
"""
Returns relative humidity from temperature and dewpoint temperature.
@@ -272,151 +302,42 @@ def relative_humidity(T: int | float | np.ndarray | tf.Tensor,
Examples
--------
- >>> T = 300 # K
- >>> Td = 290 # K
- >>> RH = relative_humidity(T, Td)
+ >>> T = 300.0 # K
+ >>> Td = 290.0 # K
+ >>> RH = relative_humidity_from_dewpoint(T, Td)
>>> RH
- 0.5699908521249278
+ 0.542647138543181
- >>> T = np.arange(270, 311, 5) # K
+ >>> T = np.arange(270.0, 311.0, 5) # K
>>> T
array([270, 275, 280, 285, 290, 295, 300, 305, 310])
- >>> Td = np.arange(260, 301, 5) # K
+ >>> Td = np.arange(260.0, 301.0, 5) # K
>>> Td
array([260, 265, 270, 275, 280, 285, 290, 295, 300])
- >>> RH = relative_humidity(T, Td)
+ >>> RH = relative_humidity_from_dewpoint(T, Td)
>>> RH
- array([0.49824518, 0.51115071, 0.52366592, 0.53579872, 0.54755768,
- 0.5589519 , 0.56999085, 0.58068426, 0.591042 ])
+ array([0.45971571, 0.47466998, 0.48916171, 0.50319697, 0.51678347,
+ 0.52993019, 0.54264714, 0.55494506, 0.56683529])
"""
- e = vapor_pressure(Td)
- es = vapor_pressure(T) # saturation vapor pressure
- RH = e / es # relative humidity
+
+ e = saturation_vapor_pressure(Td)
+ es = saturation_vapor_pressure(T)
+ RH = e / es
return RH
-def specific_humidity_from_dewpoint(Td: int | float | np.ndarray | tf.Tensor,
- P: int | float | np.ndarray | tf.Tensor):
+def specific_humidity_from_relative_humidity(P, T, RH):
"""
- Calculates specific humidity from dewpoint and pressure.
+ Calculates specific humidity from relative humidity, air temperature, and pressure.
Parameters
----------
- Td: float or iterable object
- Dewpoint temperature expressed as kelvin (K).
P: float or iterable object
Air pressure expressed as pascals (Pa).
-
- Returns
- -------
- q: float or iterable object
- Specific humidity expressed as grams of water vapor per gram of dry air (unitless; g/g or kg/kg).
-
- Examples
- --------
- >>> Td = 290 # K
- >>> P = 1e5 # Pa
- >>> q = specific_humidity_from_dewpoint(Td, P)
- >>> q
- 0.010829329732443743
-
- >>> Td = np.arange(260, 301, 5) # K
- >>> Td
- array([260, 265, 270, 275, 280, 285, 290, 295, 300])
- >>> P = np.arange(800, 1001, 25) * 100 # Pa
- >>> P
- array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
- 100000])
- >>> q = specific_humidity_from_dewpoint(Td, P)
- >>> q
- array([0.00192352, 0.0026611 , 0.00363728, 0.00491531, 0.0065716 ,
- 0.00869781, 0.01140324, 0.01481755, 0.01909393])
- """
- e = vapor_pressure(Td)
- return epsilon * e / (P - (0.378 * e)) # q: specific humidity
-
-
-def mixing_ratio_from_specific_humidity(q: int | float | np.ndarray | tf.Tensor):
- """
- Calculates mixing ratio from specific humidity.
-
- Parameters
- ----------
- q: float or iterable object
- Specific humidity expressed as grams of water vapor per gram of dry air (unitless; g/g or kg/kg).
-
- Returns
- -------
- r: float or iterable object
- Mixing ratio expressed as grams of water per gram of dry air (unitless; g/g or kg/kg).
-
- Examples
- --------
- >>> q = 20 / 1000 # g/kg -> kg/kg
- >>> r = mixing_ratio_from_specific_humidity(q)
- >>> r * 1000 # kg/kg -> g/kg
- 20.408163265306126
-
- >>> q = np.arange(5, 25.01, 2.5) / 1000 # g/kg -> kg/kg
- >>> q
- array([0.005 , 0.0075, 0.01 , 0.0125, 0.015 , 0.0175, 0.02 , 0.0225,
- 0.025 ])
- >>> r = mixing_ratio_from_specific_humidity(q)
- >>> r * 1000 # kg/kg -> g/kg
- array([ 5.02512563, 7.55667506, 10.1010101 , 12.65822785, 15.2284264 ,
- 17.81170483, 20.40816327, 23.01790281, 25.64102564])
- """
- return q / (1 - q) # r: mixing ratio
-
-
-def specific_humidity_from_mixing_ratio(r: int | float | np.ndarray | tf.Tensor):
- """
- Calculates specific humidity from mixing ratio.
-
- Parameters
- ----------
- r: float or iterable object
- Mixing ratio expressed as grams of water per gram of dry air (unitless; g/g or kg/kg).
-
- Returns
- -------
- q: float or iterable object
- Specific humidity expressed as grams of water vapor per gram of dry air (unitless; g/g or kg/kg).
-
- Examples
- --------
- >>> r = 20 / 1000 # g/kg -> kg/kg
- >>> q = specific_humidity_from_mixing_ratio(r)
- >>> q * 1000 # kg/kg -> g/kg
- 19.607843137254903
-
- >>> r = np.arange(5, 25.01, 2.5) / 1000 # g/kg -> kg/kg
- >>> r
- array([0.005 , 0.0075, 0.01 , 0.0125, 0.015 , 0.0175, 0.02 , 0.0225,
- 0.025 ])
- >>> q = specific_humidity_from_mixing_ratio(r)
- >>> q * 1000 # kg/kg -> g/kg
- array([ 4.97512438, 7.44416873, 9.9009901 , 12.34567901, 14.77832512,
- 17.1990172 , 19.60784314, 22.00488998, 24.3902439 ])
- """
- return r / (1 + r) # q: specific humidity
-
-
-def specific_humidity_from_relative_humidity(
- RH: int | float | np.ndarray | tf.Tensor,
- T: int | float | np.ndarray | tf.Tensor,
- P: int | float | np.ndarray | tf.Tensor):
- """
- Calculates specific humidity from relative humidity, air temperature, and pressure.
-
- Parameters
- ----------
- RH: float or iterable object
- Relative humidity.
T: float or iterable object
Air temperature expressed as kelvin (K).
- P: float or iterable object
- Air pressure expressed as pascals (Pa).
+ RH: float or iterable object
+ Relative humidity.
Returns
-------
@@ -425,50 +346,47 @@ def specific_humidity_from_relative_humidity(
Examples
--------
- >>> RH = 0.8
- >>> T = 300 # K
>>> P = 1e5 # Pa
- >>> q = specific_humidity_from_relative_humidity(RH, T, P)
+ >>> T = 300 # K
+ >>> RH = 0.8
+ >>> q = specific_humidity_from_relative_humidity(P, T, RH)
>>> q * 1000 # kg/kg -> g/kg
15.239787946316241
- >>> RH = np.arange(0.5, 0.91, 0.05)
- >>> RH
- array([0.5 , 0.55, 0.6 , 0.65, 0.7 , 0.75, 0.8 , 0.85, 0.9 ])
- >>> T = np.arange(270, 311, 5) # K
- >>> T
- array([270, 275, 280, 285, 290, 295, 300, 305, 310])
>>> P = np.arange(800, 1001, 25) * 100 # Pa
>>> P
array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
100000])
- >>> q = specific_humidity_from_relative_humidity(RH, T, P)
+ >>> T = np.arange(270, 311, 5) # K
+ >>> T
+ array([270, 275, 280, 285, 290, 295, 300, 305, 310])
+ >>> RH = np.arange(0.5, 0.91, 0.05)
+ >>> RH
+ array([0.5 , 0.55, 0.6 , 0.65, 0.7 , 0.75, 0.8 , 0.85, 0.9 ])
+ >>> q = specific_humidity_from_relative_humidity(P, T, RH)
>>> q * 1000
array([ 1.93030564, 2.86370132, 4.16882688, 5.96677284, 8.41051444,
11.6918278 , 16.04970636, 21.78077898, 29.25247145])
"""
- es = vapor_pressure(T)
+ es = saturation_vapor_pressure(T)
e = RH * es
- w = epsilon * e / (P - e)
- q = w / (w + 1)
+ r = epsilon * e / (P - e)
+ q = r / (r + 1)
return q
-def equivalent_potential_temperature(
- T: int | float | np.ndarray | tf.Tensor,
- Td: int | float | np.ndarray | tf.Tensor,
- P: int | float | np.ndarray | tf.Tensor):
+def equivalent_potential_temperature(P, T, Td):
"""
Calculates equivalent potential temperature (theta-e) from temperature, dewpoint, and pressure.
Parameters
----------
+ P: float or iterable object
+ Air pressure expressed as pascals (Pa).
T: float or iterable object
Air temperature expressed as kelvin (K).
Td: float or iterable object
Dewpoint temperature expressed as kelvin (K).
- P: float or iterable object
- Air pressure expressed as pascals (Pa).
Returns
-------
@@ -477,39 +395,41 @@ def equivalent_potential_temperature(
Examples
--------
+ >>> P = 1e5 # Pa
>>> T = 300 # K
>>> Td = 290 # K
- >>> P = 1e5 # Pa
- >>> theta_e = equivalent_potential_temperature(T, Td, P)
+ >>> theta_e = equivalent_potential_temperature(P, T, Td)
>>> theta_e
326.52430009577137
+ >>> P = np.arange(800, 1001, 25) * 100 # Pa
+ >>> P
+ array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
+ 100000])
>>> T = np.arange(270, 311, 5) # K
>>> T
array([270, 275, 280, 285, 290, 295, 300, 305, 310])
>>> Td = np.arange(260, 301, 5) # K
>>> Td
array([260, 265, 270, 275, 280, 285, 290, 295, 300])
- >>> P = np.arange(800, 1001, 25) * 100 # Pa
- >>> P
- array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
- 100000])
- >>> theta_e = equivalent_potential_temperature(T, Td, P)
+ >>> theta_e = equivalent_potential_temperature(P, T, Td)
>>> theta_e
array([292.582033 , 297.1606042 , 302.3295548 , 308.2501438 ,
315.12588597, 323.21500183, 332.84798857, 344.45298499,
358.59322677])
"""
- RH = relative_humidity(T, Td)
- theta = potential_temperature(T, P)
- rv = mixing_ratio_from_dewpoint(Td, P)
- theta_e = theta * tf.pow(RH, -rv * Rv / Cpd) * tf.exp(Lv * rv / (Cpd * T)) if all(tf.is_tensor(var) for var in [T, Td, P]) else \
- theta * np.power(RH, -rv * Rv / Cpd) * np.exp(Lv * rv / (Cpd * T))
+ RH = relative_humidity_from_dewpoint(T, Td)
+ theta = potential_temperature(P, T)
+ rv = mixing_ratio_from_dewpoint(P, Td)
+ theta_e = (
+ theta * tf.pow(RH, -rv * Rv / Cpd) * tf.exp(Lv * rv / (Cpd * T))
+ if all(tf.is_tensor(var) for var in [T, Td, P])
+ else theta * np.power(RH, -rv * Rv / Cpd) * np.exp(Lv * rv / (Cpd * T))
+ )
return theta_e
-def wet_bulb_temperature(T: int | float | np.ndarray | tf.Tensor,
- Td: int | float | np.ndarray | tf.Tensor):
+def wet_bulb_temperature(T, Td):
"""
Calculates wet-bulb temperature from temperature and dewpoint.
@@ -545,7 +465,7 @@ def wet_bulb_temperature(T: int | float | np.ndarray | tf.Tensor,
284.70999386, 289.26248252, 293.85201027, 298.47572385,
303.13100751])
"""
- RH = relative_humidity(T, Td) * 100
+ RH = relative_humidity_from_dewpoint(T, Td) * 100
c1 = 0.151977
c2 = 8.313659
c3 = 1.676331
@@ -557,30 +477,28 @@ def wet_bulb_temperature(T: int | float | np.ndarray | tf.Tensor,
if tf.is_tensor(T):
Tw = (T - 273.15) * tf.atan(c1 * tf.sqrt(RH + c2))
Tw += tf.atan(T - 273.15 + RH) - tf.atan(RH - c3)
- Tw += (c4 * tf.pow(RH, 1.5) * tf.atan(c5 * RH))
+ Tw += c4 * tf.pow(RH, 1.5) * tf.atan(c5 * RH)
else:
Tw = (T - 273.15) * np.arctan(c1 * np.sqrt(RH + c2))
Tw += np.arctan(T - 273.15 + RH) - np.arctan(RH - c3)
- Tw += (c4 * np.power(RH, 1.5) * np.arctan(c5 * RH))
+ Tw += c4 * np.power(RH, 1.5) * np.arctan(c5 * RH)
Tw += 273.15 - c6
return Tw
-def wet_bulb_potential_temperature(T: int | float | np.ndarray | tf.Tensor,
- Td: int | float | np.ndarray | tf.Tensor,
- P: int | float | np.ndarray | tf.Tensor):
+def wet_bulb_potential_temperature(P, T, Td):
"""
Returns wet-bulb potential temperature (theta-w) in kelvin (K).
Parameters
----------
+ P: float or iterable object
+ Air pressure expressed as pascals (Pa).
T: float or iterable object
Air temperature expressed as kelvin (K).
Td: float or iterable object
Dewpoint temperature expressed as kelvin (K).
- P: float or iterable object
- Air pressure expressed as pascals (Pa).
Returns
-------
@@ -625,73 +543,63 @@ def wet_bulb_potential_temperature(T: int | float | np.ndarray | tf.Tensor,
b4 = -0.592934
C = 273.15
- theta_e = equivalent_potential_temperature(T, Td, P)
+ theta_e = equivalent_potential_temperature(P, T, Td)
X = theta_e / C
# Wet-bulb potential temperature approximation (Davies-Jones 2008, Eq. 3.8).
- if all(tf.is_tensor(var) for var in [T, Td, P]):
+ if all(tf.is_tensor(var) for var in [P, T, Td]):
theta_wc = theta_e - tf.exp(
- (a0 + (a1 * X) + (a2 * tf.pow(X, 2)) + (a3 * tf.pow(X, 3)) + (a4 * tf.pow(X, 4))) /
- (1 + (b1 * X) + (b2 * tf.pow(X, 2)) + (b3 * tf.pow(X, 3)) + (b4 * tf.pow(X, 4))))
+ (
+ a0
+ + (a1 * X)
+ + (a2 * tf.pow(X, 2))
+ + (a3 * tf.pow(X, 3))
+ + (a4 * tf.pow(X, 4))
+ )
+ / (
+ 1
+ + (b1 * X)
+ + (b2 * tf.pow(X, 2))
+ + (b3 * tf.pow(X, 3))
+ + (b4 * tf.pow(X, 4))
+ )
+ )
theta_w = tf.where(theta_e > 173.15, theta_wc, theta_e)
else:
theta_wc = theta_e - np.exp(
- (a0 + (a1 * X) + (a2 * np.power(X, 2)) + (a3 * np.power(X, 3)) + (a4 * np.power(X, 4))) /
- (1 + (b1 * X) + (b2 * np.power(X, 2)) + (b3 * np.power(X, 3)) + (b4 * np.power(X, 4))))
+ (
+ a0
+ + (a1 * X)
+ + (a2 * np.power(X, 2))
+ + (a3 * np.power(X, 3))
+ + (a4 * np.power(X, 4))
+ )
+ / (
+ 1
+ + (b1 * X)
+ + (b2 * np.power(X, 2))
+ + (b3 * np.power(X, 3))
+ + (b4 * np.power(X, 4))
+ )
+ )
theta_w = np.where(theta_e > 173.15, theta_wc, theta_e)
return theta_w
-def vapor_pressure(Td: int | float | np.ndarray | tf.Tensor):
- """
- Calculates vapor pressure in pascals (Pa) for a given Dewpoint temperature expressed as kelvin (K).
-
- Parameters
- ----------
- Td: float or iterable object
- Dewpoint temperature expressed as kelvin (K).
-
- Returns
- -------
- vapor_pressure: float or iterable object
- Vapor pressure expressed as pascals (Pa).
-
- Examples
- --------
- >>> Td = 290 # K
- >>> vap_pres = vapor_pressure(Td)
- >>> vap_pres
- 1729.7443936886634
-
- >>> Td = np.arange(260, 301, 5) # K
- >>> Td
- array([260, 265, 270, 275, 280, 285, 290, 295, 300])
- >>> vap_pres = vapor_pressure(Td)
- >>> vap_pres
- array([ 247.12075845, 352.40493817, 495.98223586, 689.43450819,
- 947.13483326, 1286.74161001, 1729.74439369, 2302.0614118 ,
- 3034.68799059])
- """
-
- vap_pres = e_knot * tf.exp((Lv/Rv) * ((1/273.15) - (1/Td))) if tf.is_tensor(Td) else e_knot * np.exp((Lv/Rv) * ((1/273.15) - (1/Td)))
- return vap_pres
-
-
-def virtual_potential_temperature(T: int | float | np.ndarray | tf.Tensor,
- Td: int | float | np.ndarray | tf.Tensor,
- P: int | float | np.ndarray | tf.Tensor):
+def virtual_potential_temperature(P, T, Td):
"""
Calculates virtual potential temperature (theta-v) from temperature, dewpoint, and pressure.
Parameters
----------
+ P: float or iterable object
+ Air pressure expressed as pascals (Pa).
T: float or iterable object
Air temperature expressed as kelvin (K).
Td: float or iterable object
Dewpoint temperature expressed as kelvin (K).
- P: float or iterable object
- Air pressure expressed as pascals (Pa).
+
Returns
-------
@@ -700,36 +608,35 @@ def virtual_potential_temperature(T: int | float | np.ndarray | tf.Tensor,
Examples
--------
+ >>> P = 1e5 # Pa
>>> T = 300 # K
>>> Td = 290 # K
- >>> P = 1e5 # Pa
- >>> theta_v = virtual_potential_temperature(T, Td, P)
+ >>> theta_v = virtual_potential_temperature(P, T, Td)
>>> theta_v
301.9745879930382
+ >>> P = np.arange(800, 1001, 25) * 100 # Pa
+ >>> P
+ array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
+ 100000])
>>> T = np.arange(270, 311, 5) # K
>>> T
array([270, 275, 280, 285, 290, 295, 300, 305, 310])
>>> Td = np.arange(260, 301, 5) # K
>>> Td
array([260, 265, 270, 275, 280, 285, 290, 295, 300])
- >>> P = np.arange(800, 1001, 25) * 100 # Pa
- >>> P
- array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
- 100000])
- >>> theta_v = virtual_potential_temperature(T, Td, P)
+ >>> theta_v = virtual_potential_temperature(P, T, Td)
>>> theta_v
array([288.09159758, 290.9910892 , 293.94212815, 296.95595223,
300.04678011, 303.23228364, 306.53413747, 309.97866221,
313.59758378])
"""
- Tv = virtual_temperature_from_dewpoint(T, Td, P)
- theta_v = potential_temperature(Tv, P)
+ Tv = virtual_temperature_from_dewpoint(P, T, Td)
+ theta_v = potential_temperature(P, Tv)
return theta_v
-def virtual_temperature_from_mixing_ratio(T: int | float | np.ndarray | tf.Tensor,
- r: int | float | np.ndarray | tf.Tensor):
+def virtual_temperature_from_mixing_ratio(T, r):
"""
Calculates virtual temperature from temperature and mixing ratio.
@@ -769,20 +676,18 @@ def virtual_temperature_from_mixing_ratio(T: int | float | np.ndarray | tf.Tenso
return T * (1 + (r / epsilon)) / (1 + r)
-def virtual_temperature_from_dewpoint(T: int | float | np.ndarray | tf.Tensor,
- Td: int | float | np.ndarray | tf.Tensor,
- P: int | float | np.ndarray | tf.Tensor):
+def virtual_temperature_from_dewpoint(P, T, Td):
"""
Calculates virtual temperature from temperature, dewpoint, and pressure.
Parameters
----------
+ P: float or iterable object
+ Air pressure expressed as pascals (Pa).
T: float or iterable object
Air temperature expressed as kelvin (K).
Td: float or iterable object
Dewpoint temperature expressed as kelvin (K).
- P: float or iterable object
- Air pressure expressed as pascals (Pa).
Returns
-------
@@ -791,30 +696,30 @@ def virtual_temperature_from_dewpoint(T: int | float | np.ndarray | tf.Tensor,
Examples
--------
+ >>> P = 1e5 # Pa
>>> T = 300 # K
>>> Td = 290 # K
- >>> P = 1e5 # Pa
- >>> Tv = virtual_temperature_from_dewpoint(T, Td, P)
+ >>> Tv = virtual_temperature_from_dewpoint(P, T, Td)
>>> Tv
301.9745879930382
+ >>> P = np.arange(800, 1001, 25) * 100 # Pa
+ >>> P
+ array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
+ 100000])
>>> T = np.arange(270, 311, 5) # K
>>> T
array([270, 275, 280, 285, 290, 295, 300, 305, 310])
>>> Td = np.arange(260, 301, 5) # K
>>> Td
array([260, 265, 270, 275, 280, 285, 290, 295, 300])
- >>> P = np.arange(800, 1001, 25) * 100 # Pa
- >>> P
- array([ 80000, 82500, 85000, 87500, 90000, 92500, 95000, 97500,
- 100000])
- >>> Tv = virtual_temperature_from_dewpoint(T, Td, P)
+ >>> Tv = virtual_temperature_from_dewpoint(P, T, Td)
>>> Tv
array([270.3156564 , 275.44478153, 280.61899683, 285.85143108,
291.15830423, 296.55950086, 302.07923396, 307.74681888,
313.59758378])
"""
- r = mixing_ratio_from_dewpoint(Td, P)
+ r = mixing_ratio_from_dewpoint(P, Td)
Tv = virtual_temperature_from_mixing_ratio(T, r)
return Tv
@@ -841,12 +746,16 @@ def advection(field, u, v, lons, lats):
Lons, Lats = np.meshgrid(lons, lats)
x, y = data_utils.haversine(Lons, Lats) # x and y are expressed as kilometers
- x = x.T; y = y.T # transpose x and y so arrays have shape M x N
- x *= 1e3; y *= 1e3 # convert x and y coordinates to meters
+ x = x.T
+ y = y.T # transpose x and y so arrays have shape M x N
+ x *= 1e3
+ y *= 1e3 # convert x and y coordinates to meters
dfield_dx = np.diff(field, axis=0) / np.diff(x, axis=0)
dfield_dy = np.diff(field, axis=1) / np.diff(y, axis=1)
- advect[:-1, :-1] = - (u[:-1, :-1] * dfield_dx[:, :-1]) - (v[:-1, :-1] * dfield_dy[:-1, :])
+ advect[:-1, :-1] = -(u[:-1, :-1] * dfield_dx[:, :-1]) - (
+ v[:-1, :-1] * dfield_dy[:-1, :]
+ )
return advect
diff --git a/src/fronts/utils/constants.py b/src/fronts/utils/constants.py
new file mode 100644
index 0000000..58421db
--- /dev/null
+++ b/src/fronts/utils/constants.py
@@ -0,0 +1,446 @@
+import numpy as np
+
+# [min, max, mean, std, mean (lat weighted), std (lat weighted)]
+NORMALIZATION_PARAMS = dict()
+NORMALIZATION_PARAMS["q_300"] = [0.0, 2.25, 0.1484, 0.1701, 0.1854, 0.1957]
+NORMALIZATION_PARAMS["q_500"] = [0.0, 9.375, 0.8906, 1.1321, 1.1552, 1.2790]
+NORMALIZATION_PARAMS["q_700"] = [0.0, 16.0, 2.4991, 2.6180, 3.2535, 2.8386]
+NORMALIZATION_PARAMS["q_850"] = [0.0, 21.75, 4.6545, 4.1544, 6.1259, 4.2482]
+NORMALIZATION_PARAMS["q_900"] = [0.0, 23.25, 5.6176, 4.7384, 7.4258, 4.7040]
+NORMALIZATION_PARAMS["q_950"] = [0.0, 25.125, 6.5983, 5.5294, 8.7716, 5.4365]
+NORMALIZATION_PARAMS["q_1000"] = [0.0, 29.25, 7.0913, 5.9191, 9.4306, 5.7912]
+NORMALIZATION_PARAMS["RH_300"] = [0.0, 1.0, 0.3426, 0.2172, 0.3225, 0.2226]
+NORMALIZATION_PARAMS["RH_500"] = [0.0, 1.0, 0.4073, 0.264, 0.3737, 0.2755]
+NORMALIZATION_PARAMS["RH_700"] = [0.0, 1.0, 0.4791, 0.273, 0.4614, 0.2811]
+NORMALIZATION_PARAMS["RH_850"] = [0.0, 1.0, 0.6047, 0.2687, 0.6242, 0.2623]
+NORMALIZATION_PARAMS["RH_900"] = [0.0, 1.0, 0.6682, 0.2657, 0.7033, 0.2444]
+NORMALIZATION_PARAMS["RH_950"] = [0.0, 1.0, 0.7037, 0.2673, 0.7464, 0.2371]
+NORMALIZATION_PARAMS["RH_1000"] = [0.0, 1.0, 0.659, 0.2457, 0.6897, 0.2104]
+NORMALIZATION_PARAMS["sp_z_300"] = [796.0, 989.0, 913.0033, 52.2398, 935.2045, 43.2219]
+NORMALIZATION_PARAMS["sp_z_500"] = [442.0, 605.0, 552.4796, 34.4652, 566.8235, 28.0863]
+NORMALIZATION_PARAMS["sp_z_700"] = [206.0, 334.0, 295.3216, 22.0330, 303.9771, 17.9648]
+NORMALIZATION_PARAMS["sp_z_850"] = [60.0, 177.0, 140.3339, 15.2319, 145.5283, 12.4728]
+NORMALIZATION_PARAMS["sp_z_900"] = [16.0, 130.0, 93.9296, 13.5094, 98.0819, 11.1214]
+NORMALIZATION_PARAMS["sp_z_950"] = [-27.0, 88.0, 49.7221, 12.1269, 52.8635, 10.0784]
+NORMALIZATION_PARAMS["sp_z_1000"] = [-69.0, 49.0, 7.4702, 11.1432, 9.6184, 9.3651]
+NORMALIZATION_PARAMS["T_300"] = [199.0, 257.0, 229.2386, 10.8319, 233.7161, 9.4612]
+NORMALIZATION_PARAMS["T_500"] = [215.0, 284.0, 253.2812, 13.0718, 258.8222, 10.9391]
+NORMALIZATION_PARAMS["T_700"] = [208.0, 302.0, 267.7375, 14.8156, 274.0165, 11.5235]
+NORMALIZATION_PARAMS["T_850"] = [217.0, 315.0, 274.9058, 15.5217, 281.4309, 12.2972]
+NORMALIZATION_PARAMS["T_900"] = [219.0, 319.0, 276.7593, 15.7479, 283.3672, 12.5458]
+NORMALIZATION_PARAMS["T_950"] = [219.0, 322.0, 278.7952, 16.2086, 285.6539, 12.8296]
+NORMALIZATION_PARAMS["T_1000"] = [216.0, 326.0, 281.4339, 16.9643, 288.7020, 13.3128]
+NORMALIZATION_PARAMS["Td_300"] = [159.0, 254.0, 217.3682, 10.5844, 220.6169, 10.3878]
+NORMALIZATION_PARAMS["Td_500"] = [165.0, 279.0, 239.8065, 13.0666, 243.1986, 13.2075]
+NORMALIZATION_PARAMS["Td_700"] = [176.0, 291.0, 255.3803, 15.7529, 260.1953, 14.6048]
+NORMALIZATION_PARAMS["Td_850"] = [186.0, 298.0, 266.3356, 17.3281, 272.8387, 13.9459]
+NORMALIZATION_PARAMS["Td_900"] = [190.0, 300.0, 269.7803, 17.9404, 276.8977, 13.6627]
+NORMALIZATION_PARAMS["Td_950"] = [194.0, 302.0, 272.4960, 18.8094, 280.1294, 14.0384]
+NORMALIZATION_PARAMS["Td_1000"] = [195.0, 306.0, 274.0599, 19.37, 281.9399, 14.3432]
+NORMALIZATION_PARAMS["Tv_300"] = [199.0, 257.0, 229.2578, 10.8478, 233.7418, 9.4773]
+NORMALIZATION_PARAMS["Tv_500"] = [215.0, 285.0, 253.4233, 13.1756, 259.0079, 11.0404]
+NORMALIZATION_PARAMS["Tv_700"] = [208.0, 303.0, 268.1590, 15.1075, 274.5694, 11.8080]
+NORMALIZATION_PARAMS["Tv_850"] = [217.0, 316.0, 275.7130, 16.0917, 282.5005, 12.8374]
+NORMALIZATION_PARAMS["Tv_900"] = [219.0, 320.0, 277.7404, 16.4410, 284.6725, 13.1932]
+NORMALIZATION_PARAMS["Tv_950"] = [218.0, 323.0, 279.9575, 17.0354, 287.2087, 13.5979]
+NORMALIZATION_PARAMS["Tv_1000"] = [216.0, 327.0, 282.6972, 17.8702, 290.3930, 14.1572]
+NORMALIZATION_PARAMS["theta_300"] = [281.0, 362.0, 323.2393, 15.2710, 329.5529, 13.3377]
+NORMALIZATION_PARAMS["theta_500"] = [262.0, 346.0, 308.6888, 15.9301, 315.4419, 13.3306]
+NORMALIZATION_PARAMS["theta_700"] = [231.0, 334.0, 296.4287, 16.4027, 303.3806, 12.7576]
+NORMALIZATION_PARAMS["theta_850"] = [227.0, 330.0, 287.9577, 16.2584, 294.7926, 12.8807]
+NORMALIZATION_PARAMS["theta_900"] = [226.0, 329.0, 285.2081, 16.2285, 292.0178, 12.9286]
+NORMALIZATION_PARAMS["theta_950"] = [222.0, 327.0, 282.9067, 16.4475, 289.8666, 13.0187]
+NORMALIZATION_PARAMS["theta_1000"] = [
+ 216.0,
+ 326.0,
+ 281.4339,
+ 16.9643,
+ 288.7020,
+ 13.3128,
+]
+NORMALIZATION_PARAMS["theta_e_300"] = [
+ 281.0,
+ 368.0,
+ 323.7062,
+ 15.6360,
+ 330.1604,
+ 13.7126,
+]
+NORMALIZATION_PARAMS["theta_e_500"] = [
+ 262.0,
+ 360.0,
+ 311.2741,
+ 17.9377,
+ 318.7880,
+ 15.4055,
+]
+NORMALIZATION_PARAMS["theta_e_700"] = [
+ 229.0,
+ 367.0,
+ 303.0449,
+ 21.4989,
+ 312.0104,
+ 18.0769,
+]
+NORMALIZATION_PARAMS["theta_e_850"] = [
+ 226.0,
+ 372.0,
+ 299.6698,
+ 25.2475,
+ 310.2498,
+ 21.8210,
+]
+NORMALIZATION_PARAMS["theta_e_900"] = [
+ 224.0,
+ 373.0,
+ 299.1296,
+ 26.7758,
+ 310.4758,
+ 23.1717,
+]
+NORMALIZATION_PARAMS["theta_e_950"] = [
+ 220.0,
+ 375.0,
+ 299.0758,
+ 28.8183,
+ 311.4362,
+ 24.9726,
+]
+NORMALIZATION_PARAMS["theta_e_1000"] = [
+ 215.0,
+ 390.0,
+ 298.7088,
+ 30.2554,
+ 311.7681,
+ 26.1548,
+]
+NORMALIZATION_PARAMS["theta_v_300"] = [
+ 281.0,
+ 362.0,
+ 323.2664,
+ 15.2933,
+ 329.5891,
+ 13.3604,
+]
+NORMALIZATION_PARAMS["theta_v_500"] = [
+ 262.0,
+ 347.0,
+ 308.8620,
+ 16.0566,
+ 315.6682,
+ 13.4540,
+]
+NORMALIZATION_PARAMS["theta_v_700"] = [
+ 231.0,
+ 335.0,
+ 296.8954,
+ 16.7259,
+ 303.9928,
+ 13.0726,
+]
+NORMALIZATION_PARAMS["theta_v_850"] = [
+ 227.0,
+ 330.0,
+ 288.8032,
+ 16.8555,
+ 295.9130,
+ 13.4466,
+]
+NORMALIZATION_PARAMS["theta_v_900"] = [
+ 225.0,
+ 329.0,
+ 286.2193,
+ 16.9428,
+ 293.3630,
+ 13.5958,
+]
+NORMALIZATION_PARAMS["theta_v_950"] = [
+ 220.0,
+ 328.0,
+ 284.0862,
+ 17.2865,
+ 291.4443,
+ 13.7983,
+]
+NORMALIZATION_PARAMS["theta_v_1000"] = [
+ 216.0,
+ 327.0,
+ 282.6972,
+ 17.8702,
+ 290.3930,
+ 14.1572,
+]
+NORMALIZATION_PARAMS["u_300"] = [-65.2, 115.6, 11.6922, 17.1698, 12.5158, 17.4089]
+NORMALIZATION_PARAMS["u_500"] = [-51.6, 82.4, 6.5193, 12.0255, 6.6043, 12.0727]
+NORMALIZATION_PARAMS["u_700"] = [-50.4, 58.4, 3.3234, 9.24, 3.1705, 9.2717]
+NORMALIZATION_PARAMS["u_850"] = [-55.2, 53.2, 1.4020, 8.2607, 1.0348, 8.2607]
+NORMALIZATION_PARAMS["u_900"] = [-55.6, 50.8, 0.8581, 8.1317, 0.4330, 8.1176]
+NORMALIZATION_PARAMS["u_950"] = [-47.6, 43.2, 7.6961, -0.0558, 7.6638]
+NORMALIZATION_PARAMS["u_1000"] = [-33.6, 30.4, -0.0452, 6.1942, -0.4278, 6.1954]
+NORMALIZATION_PARAMS["v_300"] = [-80.0, 94.0, -0.0227, 13.3571, -0.0253, 12.6938]
+NORMALIZATION_PARAMS["v_500"] = [-65.6, 70.0, -0.0251, 9.2148, -0.0366, 8.5501]
+NORMALIZATION_PARAMS["v_700"] = [-48.4, 53.2, 0.0255, 6.9146, -0.0117, 6.4236]
+NORMALIZATION_PARAMS["v_850"] = [-49.6, 52.8, 0.1468, 6.2973, 0.0924, 5.8921]
+NORMALIZATION_PARAMS["v_900"] = [-48.4, 51.6, 0.2032, 6.4313, 0.1744, 6.0884]
+NORMALIZATION_PARAMS["v_950"] = [-44.8, 46.8, 0.2083, 6.4496, 0.1977, 6.187]
+NORMALIZATION_PARAMS["v_1000"] = [-31.6, 30.8, 0.1949, 5.3437, 0.1984, 5.1694]
+
+
+# default values for extents of domains [start lon, end lon, start lat, end lat]
+DOMAIN_EXTENTS = {
+ "atlantic": [290, 349.75, 16, 55.75],
+ "conus": [228, 299.75, 25, 56.75],
+ "ecmwf": [0, 359.75, -89.75, 90],
+ "full": [130, 369.75, 0.25, 80],
+ "global": [0, 359.75, -89.75, 90],
+ "goes-merged": [144, 359.75, 2, 69.75],
+ "hrrr": [
+ 225.90452026573686,
+ 299.0828072281622,
+ 21.138123000000018,
+ 52.61565330680793,
+ ],
+ "MERGIR": [130, 359.75, 20, 59.75],
+ "namnest-conus": [
+ 225.90387325951775,
+ 299.08216099364034,
+ 21.138,
+ 52.61565399063001,
+ ],
+ "nam-12km": [
+ 207.12137749594984,
+ 310.58401341435564,
+ 12.190000000000005,
+ 61.30935757335816,
+ ],
+ "pacific": [145, 234.75, 16, 55.75],
+}
+
+# colors for plotted ground truth fronts
+FRONT_COLORS = {
+ "CF": "blue",
+ "WF": "red",
+ "SF": "limegreen",
+ "OF": "darkviolet",
+ "CF-F": "darkblue",
+ "WF-F": "darkred",
+ "SF-F": "darkgreen",
+ "OF-F": "darkmagenta",
+ "CF-D": "lightskyblue",
+ "WF-D": "lightcoral",
+ "SF-D": "lightgreen",
+ "OF-D": "violet",
+ "INST": "gold",
+ "TROF": "goldenrod",
+ "TT": "orange",
+ "DL": "chocolate",
+ "MERGED-CF": "blue",
+ "MERGED-WF": "red",
+ "MERGED-SF": "limegreen",
+ "MERGED-OF": "darkviolet",
+ "MERGED-F": "gray",
+ "MERGED-T": "brown",
+ "F_BIN": "tab:red",
+ "MERGED-F_BIN": "tab:red",
+}
+
+# colormaps of probability contours for front predictions
+CONTOUR_CMAPS = {
+ "CF": "Blues",
+ "WF": "Reds",
+ "SF": "Greens",
+ "OF": "Purples",
+ "CF-F": "Blues",
+ "WF-F": "Reds",
+ "SF-F": "Greens",
+ "OF-F": "Purples",
+ "CF-D": "Blues",
+ "WF-D": "Reds",
+ "SF-D": "Greens",
+ "OF-D": "Purples",
+ "INST": "YlOrBr",
+ "TROF": "YlOrRd",
+ "TT": "Oranges",
+ "DL": "copper_r",
+ "MERGED-CF": "Blues",
+ "MERGED-WF": "Reds",
+ "MERGED-SF": "Greens",
+ "MERGED-OF": "Purples",
+ "MERGED-F": "Greys",
+ "MERGED-T": "YlOrBr",
+ "F_BIN": "Greys",
+ "MERGED-F_BIN": "Greys",
+}
+
+# names of front types
+FRONT_NAMES = {
+ "CF": "Cold front",
+ "WF": "Warm front",
+ "SF": "Stationary front",
+ "OF": "Occluded front",
+ "CF-F": "Cold front (forming)",
+ "WF-F": "Warm front (forming)",
+ "SF-F": "Stationary front (forming)",
+ "OF-F": "Occluded front (forming)",
+ "CF-D": "Cold front (dying)",
+ "WF-D": "Warm front (dying)",
+ "SF-D": "Stationary front (dying)",
+ "OF-D": "Occluded front (dying)",
+ "INST": "Outflow boundary",
+ "TROF": "Trough",
+ "TT": "Tropical trough",
+ "DL": "Dryline",
+ "MERGED-CF": "Cold front (any)",
+ "MERGED-WF": "Warm front (any)",
+ "MERGED-SF": "Stationary front (any)",
+ "MERGED-OF": "Occluded front (any)",
+ "MERGED-F": "CF, WF, SF, OF (any)",
+ "MERGED-T": "Trough (any)",
+ "F_BIN": "Binary front",
+ "MERGED-F_BIN": "Binary front (any)",
+}
+
+VARIABLE_NAMES = {
+ "T": "Air temperature",
+ "T_sfc": "2-meter Air temperature",
+ "T_1000": "1000mb Air temperature",
+ "T_950": "950mb Air temperature",
+ "T_900": "900mb Air temperature",
+ "T_850": "850mb Air temperature",
+ "Td": "Dewpoint",
+ "Td_sfc": "2-meter Dewpoint",
+ "Td_1000": "1000mb Dewpoint",
+ "Td_950": "950mb Dewpoint",
+ "Td_900": "900mb Dewpoint",
+ "Td_850": "850mb Dewpoint",
+ "Tv": "Virtual temperature",
+ "Tv_sfc": "2-meter Virtual temperature",
+ "Tv_1000": "1000mb Virtual temperature",
+ "Tv_950": "950mb Virtual temperature",
+ "Tv_900": "900mb Virtual temperature",
+ "Tv_850": "850mb Virtual temperature",
+ "Tw": "Wet-bulb temperature",
+ "Tw_sfc": "2-meter Wet-bulb temperature",
+ "Tw_1000": "1000mb Wet-bulb temperature",
+ "Tw_950": "950mb Wet-bulb temperature",
+ "Tw_900": "900mb Wet-bulb temperature",
+ "Tw_850": "850mb Wet-bulb temperature",
+ "theta": "Potential temperature",
+ "theta_sfc": "2-meter Potential temperature",
+ "theta_1000": "1000mb Potential temperature",
+ "theta_950": "950mb Potential temperature",
+ "theta_900": "900mb Potential temperature",
+ "theta_850": "850mb Potential temperature",
+ "theta_e": "Theta-E",
+ "theta_e_sfc": "2-meter Theta-E",
+ "theta_e_1000": "1000mb Theta-E",
+ "theta_e_950": "950mb Theta-E",
+ "theta_e_900": "900mb Theta-E",
+ "theta_e_850": "850mb Theta-E",
+ "theta_v": "Virtual potential temperature",
+ "theta_v_sfc": "2-meter Virtual potential temperature",
+ "theta_v_1000": "1000mb Virtual potential temperature",
+ "theta_v_950": "950mb Virtual potential temperature",
+ "theta_v_900": "900mb Virtual potential temperature",
+ "theta_v_850": "850mb Virtual potential temperature",
+ "theta_w": "Wet-bulb potential temperature",
+ "theta_w_sfc": "2-meter Wet-bulb potential temperature",
+ "theta_w_1000": "1000mb Wet-bulb potential temperature",
+ "theta_w_950": "950mb Wet-bulb potential temperature",
+ "theta_w_900": "900mb Wet-bulb potential temperature",
+ "theta_w_850": "850mb Wet-bulb potential temperature",
+ "u": "U-wind",
+ "u_sfc": "10-meter U-wind",
+ "u_1000": "1000mb U-wind",
+ "u_950": "950mb U-wind",
+ "u_900": "900mb U-wind",
+ "u_850": "850mb U-wind",
+ "v": "V-wind",
+ "v_sfc": "10-meter V-wind",
+ "v_1000": "1000mb V-wind",
+ "v_950": "950mb V-wind",
+ "v_900": "900mb V-wind",
+ "v_850": "850mb V-wind",
+ "q": "Specific humidity",
+ "q_sfc": "2-meter Specific humidity",
+ "q_1000": "1000mb Specific humidity",
+ "q_950": "950mb Specific humidity",
+ "q_900": "900mb Specific humidity",
+ "q_850": "850mb Specific humidity",
+ "r": "Mixing ratio",
+ "r_sfc": "2-meter Mixing ratio",
+ "r_1000": "1000mb Mixing ratio",
+ "r_950": "950mb Mixing ratio",
+ "r_900": "900mb Mixing ratio",
+ "r_850": "850mb Mixing ratio",
+ "RH": "Relative humidity",
+ "RH_sfc": "2-meter Relative humidity",
+ "RH_1000": "1000mb Relative humidity",
+ "RH_950": "950mb Relative humidity",
+ "RH_900": "900mb Relative humidity",
+ "RH_850": "850mb Relative humidity",
+ "sp_z": "Pressure/heights",
+ "sp_z_sfc": "Surface pressure",
+ "sp_z_1000": "1000mb Geopotential height",
+ "sp_z_950": "950mb Geopotential height",
+ "sp_z_900": "900mb Geopotential height",
+ "sp_z_850": "850mb Geopotential height",
+ "mslp_z": "Pressure/heights",
+ "mslp_z_sfc": "Mean sea level pressure",
+ "mslp_z_1000": "1000mb Geopotential height",
+ "mslp_z_950": "950mb Geopotential height",
+ "mslp_z_900": "900mb Geopotential height",
+ "mslp_z_850": "850mb Geopotential height",
+}
+
+VERTICAL_LEVELS = {
+ "surface": "Surface",
+ "1000": "1000 hPa",
+ "950": "950 hPa",
+ "900": "900 hPa",
+ "850": "850 hPa",
+ "700": "700 hPa",
+}
+
+# some months do not have complete front labels, so we need to specify what dates (indices) do NOT have data for the final prediction datasets
+missing_fronts_ind = {
+ "2007-05": np.array([122, 128, 130, 132]),
+ "2007-06": np.array([32, 34, 36, 200, 202]),
+ "2007-11": np.array([126, 128, 130, 132]),
+ "2007-12": np.array([206, 207]),
+ "2018-03": 203,
+ "2022-09": np.append(np.array([44, 46]), np.arange(48, 95.1, 1)).astype(int),
+ "2022-10": np.append(np.arange(80, 87.1, 1), np.arange(160, 167.1, 1)).astype(int),
+ "2022-11": 196,
+}
+
+# 3-hourly indices with missing satellite data
+missing_satellite_ind = {
+ "2018-09": np.array([78, 79, 80, 81, 82, 83, 142, 146]),
+ "2018-10": np.append(np.array([86, 134]), np.arange(189, 237.1)).astype(int),
+ "2018-11": np.append(
+ np.arange(0, 99.1, 1), np.array([120, 121, 122, 123, 124, 125, 126, 159])
+ ).astype(int),
+ "2018-12": np.array([153, 157, 205, 206, 207]),
+ "2019-01": 22,
+ "2019-02": np.array([197, 198]),
+ "2019-03": 215,
+ "2019-04": 189,
+ "2019-05": 237,
+ "2019-06": np.array([213, 221, 222]),
+ "2019-08": np.array([114, 115, 116, 117]),
+ "2020-06": np.array([22, 23, 24, 25, 26, 27]),
+ "2020-07": np.array([207, 208]),
+ "2020-08": 86,
+ "2021-01": 167,
+ "2021-03": np.array([125, 181, 182, 183]),
+ "2021-04": 231,
+ "2021-06": np.array([116, 228, 229, 230]),
+ "2021-07": np.append(np.array([67]), np.arange(170, 179.1, 1)),
+ "2022-01": 112,
+ "2022-04": 141,
+ "2022-05": np.array([189, 190]),
+ "2022-08": np.array([42, 43, 50, 51, 58]),
+ "2022-09": np.array([100, 101, 102, 103]),
+ "2022-11": np.array([55, 56, 134]),
+}
diff --git a/src/fronts/utils/data_utils.py b/src/fronts/utils/data_utils.py
new file mode 100644
index 0000000..b48062a
--- /dev/null
+++ b/src/fronts/utils/data_utils.py
@@ -0,0 +1,929 @@
+"""
+Various data tools.
+
+References
+----------
+* Snyder 1987: https://doi.org/10.3133/pp1395
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.5.3
+
+TODO
+ * Finish adding masks for xarray datasets
+"""
+
+import pandas as pd
+from shapely.geometry import LineString
+import numpy as np
+import xarray as xr
+import tensorflow as tf
+import regionmask
+from fronts.utils import constants
+
+
+def expand_fronts(
+ fronts: np.ndarray | tf.Tensor | xr.Dataset | xr.DataArray, iterations: int = 1
+):
+ """
+ Expands front labels in all directions.
+
+ Parameters
+ ----------
+ fronts: array_like of ints of shape (T, M, N) or (M, N)
+ 2-D or 3-D array of integers that identify the front type at each point. The longitude and latitude dimensions with
+ shapes (M,) and (N,) can be in any order, but the time dimension must be the first dimension if it is passed.
+ iterations: int
+ Integer representing the number of times to expand the fronts in all directions.
+
+ Returns
+ -------
+ fronts: array_like of ints of shape (T, M, N) or (1, M, N)
+ Array of integers for the expanded fronts. If the array_like object passed into the function was 2-D, a third dimension
+ will be added to the beginning of the array with size 1.
+
+ Examples
+ --------
+ * Expanding labels for one front type.
+ >>> arr = np.zeros((5, 5))
+ >>> arr[2, 2] = 1 # add cold front point
+ >>> arr
+ array([[0., 0., 0., 0., 0.],
+ [0., 0., 0., 0., 0.],
+ [0., 0., 1., 0., 0.],
+ [0., 0., 0., 0., 0.],
+ [0., 0., 0., 0., 0.]])
+ >>> expand_fronts(arr, iterations=1)
+ array([[[0., 0., 0., 0., 0.],
+ [0., 1., 1., 1., 0.],
+ [0., 1., 1., 1., 0.],
+ [0., 1., 1., 1., 0.],
+ [0., 0., 0., 0., 0.]]])
+
+ * Expanding labels for two front types.
+ >>> arr = np.zeros((5, 5))
+ >>> arr[1, 1] = 1 # add cold front point
+ >>> arr[3, 3] = 2 # add warm front point
+ >>> arr
+ array([[0., 0., 0., 0., 0.],
+ [0., 1., 0., 0., 0.],
+ [0., 0., 0., 0., 0.],
+ [0., 0., 0., 2., 0.],
+ [0., 0., 0., 0., 0.]])
+ >>> expand_fronts(arr, iterations=1)
+ array([[[1., 1., 1., 0., 0.],
+ [1., 1., 1., 0., 0.],
+ [1., 1., 2., 2., 2.],
+ [0., 0., 2., 2., 2.],
+ [0., 0., 2., 2., 2.]]])
+ """
+ if type(fronts) in [xr.Dataset, xr.DataArray]:
+ identifier = (
+ fronts["identifier"].values if type(fronts) == xr.Dataset else fronts.values
+ )
+
+ elif tf.is_tensor(fronts):
+ identifier = (
+ tf.expand_dims(fronts, axis=0) if len(fronts.shape) == 2 else fronts
+ )
+ else:
+ identifier = (
+ np.expand_dims(fronts, axis=0) if len(fronts.shape) == 2 else fronts
+ )
+
+ if tf.is_tensor(identifier):
+ for _ in range(iterations):
+ # 8 tensors representing all directions for the front expansion
+ identifier_up_left = tf.Variable(tf.zeros_like(identifier))
+ identifier_up_right = tf.Variable(tf.zeros_like(identifier))
+ identifier_down_left = tf.Variable(tf.zeros_like(identifier))
+ identifier_down_right = tf.Variable(tf.zeros_like(identifier))
+ identifier_up = tf.Variable(tf.zeros_like(identifier))
+ identifier_down = tf.Variable(tf.zeros_like(identifier))
+ identifier_left = tf.Variable(tf.zeros_like(identifier))
+ identifier_right = tf.Variable(tf.zeros_like(identifier))
+
+ identifier_down_left[..., 1:, :-1].assign(
+ tf.where(
+ (identifier[..., :-1, 1:] > 0) & (identifier[..., 1:, :-1] == 0),
+ identifier[..., :-1, 1:],
+ identifier[..., 1:, :-1],
+ )
+ )
+ identifier_down[..., 1:, :].assign(
+ tf.where(
+ (identifier[..., :-1, :] > 0) & (identifier[..., 1:, :] == 0),
+ identifier[..., :-1, :],
+ identifier[..., 1:, :],
+ )
+ )
+ identifier_down_right[..., 1:, 1:].assign(
+ tf.where(
+ (identifier[..., :-1, :-1] > 0) & (identifier[..., 1:, 1:] == 0),
+ identifier[..., :-1, :-1],
+ identifier[..., 1:, 1:],
+ )
+ )
+ identifier_up_left[..., :-1, :-1].assign(
+ tf.where(
+ (identifier[..., 1:, 1:] > 0) & (identifier[..., :-1, :-1] == 0),
+ identifier[..., 1:, 1:],
+ identifier[..., :-1, :-1],
+ )
+ )
+ identifier_up[..., :-1, :].assign(
+ tf.where(
+ (identifier[..., 1:, :] > 0) & (identifier[..., :-1, :] == 0),
+ identifier[..., 1:, :],
+ identifier[..., :-1, :],
+ )
+ )
+ identifier_up_right[..., :-1, 1:].assign(
+ tf.where(
+ (identifier[..., 1:, :-1] > 0) & (identifier[..., :-1, 1:] == 0),
+ identifier[..., 1:, :-1],
+ identifier[..., :-1, 1:],
+ )
+ )
+ identifier_left[..., :, :-1].assign(
+ tf.where(
+ (identifier[..., :, 1:] > 0) & (identifier[..., :, :-1] == 0),
+ identifier[..., :, 1:],
+ identifier[..., :, :-1],
+ )
+ )
+ identifier_right[..., :, 1:].assign(
+ tf.where(
+ (identifier[..., :, :-1] > 0) & (identifier[..., :, 1:] == 0),
+ identifier[..., :, :-1],
+ identifier[..., :, 1:],
+ )
+ )
+
+ identifier = tf.reduce_max(
+ [
+ identifier_up_left,
+ identifier_up,
+ identifier_up_right,
+ identifier_down_left,
+ identifier_down,
+ identifier_down_right,
+ identifier_left,
+ identifier_right,
+ ],
+ axis=0,
+ )
+
+ else:
+ for _ in range(iterations):
+ # 8 arrays representing all directions for the front expansion
+ identifier_up_left = np.zeros_like(identifier)
+ identifier_up_right = np.zeros_like(identifier)
+ identifier_down_left = np.zeros_like(identifier)
+ identifier_down_right = np.zeros_like(identifier)
+ identifier_up = np.zeros_like(identifier)
+ identifier_down = np.zeros_like(identifier)
+ identifier_left = np.zeros_like(identifier)
+ identifier_right = np.zeros_like(identifier)
+
+ identifier_down_left[..., 1:, :-1] = np.where(
+ (identifier[..., :-1, 1:] > 0) & (identifier[..., 1:, :-1] == 0),
+ identifier[..., :-1, 1:],
+ identifier[..., 1:, :-1],
+ )
+ identifier_down[..., 1:, :] = np.where(
+ (identifier[..., :-1, :] > 0) & (identifier[..., 1:, :] == 0),
+ identifier[..., :-1, :],
+ identifier[..., 1:, :],
+ )
+ identifier_down_right[..., 1:, 1:] = np.where(
+ (identifier[..., :-1, :-1] > 0) & (identifier[..., 1:, 1:] == 0),
+ identifier[..., :-1, :-1],
+ identifier[..., 1:, 1:],
+ )
+ identifier_up_left[..., :-1, :-1] = np.where(
+ (identifier[..., 1:, 1:] > 0) & (identifier[..., :-1, :-1] == 0),
+ identifier[..., 1:, 1:],
+ identifier[..., :-1, :-1],
+ )
+ identifier_up[..., :-1, :] = np.where(
+ (identifier[..., 1:, :] > 0) & (identifier[..., :-1, :] == 0),
+ identifier[..., 1:, :],
+ identifier[..., :-1, :],
+ )
+ identifier_up_right[..., :-1, 1:] = np.where(
+ (identifier[..., 1:, :-1] > 0) & (identifier[..., :-1, 1:] == 0),
+ identifier[..., 1:, :-1],
+ identifier[..., :-1, 1:],
+ )
+ identifier_left[..., :, :-1] = np.where(
+ (identifier[..., :, 1:] > 0) & (identifier[..., :, :-1] == 0),
+ identifier[..., :, 1:],
+ identifier[..., :, :-1],
+ )
+ identifier_right[..., :, 1:] = np.where(
+ (identifier[..., :, :-1] > 0) & (identifier[..., :, 1:] == 0),
+ identifier[..., :, :-1],
+ identifier[..., :, 1:],
+ )
+
+ identifier = np.max(
+ [
+ identifier_up_left,
+ identifier_up,
+ identifier_up_right,
+ identifier_down_left,
+ identifier_down,
+ identifier_down_right,
+ identifier_left,
+ identifier_right,
+ ],
+ axis=0,
+ )
+
+ if type(fronts) == xr.Dataset:
+ fronts["identifier"].values = identifier
+ elif type(fronts) == xr.DataArray:
+ fronts.values = identifier
+ else:
+ fronts = identifier
+
+ return fronts
+
+
+def haversine(lon: np.ndarray | int | float, lat: np.ndarray | int | float):
+ """
+ Haversine formula. Transforms lon/lat points to an x/y cartesian plane.
+
+ Parameters
+ ----------
+ lon: array_like of shape (N,), int, or float
+ Longitude component of the point(s) expressed in degrees.
+ lat: array_like of shape (N,), int, or float
+ Latitude component of the point(s) expressed in degrees.
+
+ Returns
+ -------
+ x: array_like of shape (N,) or float
+ X component of the transformed points expressed in kilometers.
+ y: array_like of shape (N,) or float
+ Y component of the transformed points expressed in kilometers.
+
+ Examples
+ --------
+ >>> lon = -95
+ >>> lat = 35
+ >>> x, y = haversine(lon, lat)
+ >>> x, y
+ (-10077.330945462296, 3892.875)
+
+ >>> lon = np.arange(10, 80.1, 10)
+ >>> lat = np.arange(10, 80.1, 10)
+ >>> x, y = haversine(lon, lat)
+ >>> x, y
+ (array([1108.01755295, 2190.70484658, 3223.05300087, 4180.69246988,
+ 5040.20418066, 5779.42053216, 6377.71302882, 6816.26345487]), array([1112.25, 2224.5 , 3336.75, 4449. , 5561.25, 6673.5 , 7785.75,
+ 8898. ]))
+ """
+ C = 40041 # average circumference of earth in kilometers
+ x = lon * C * np.cos(lat * np.pi / 360) / 360
+ y = lat * C / 360
+ return x, y
+
+
+def reverse_haversine(x, y):
+ """
+ Reverse haversine formula. Transforms x/y cartesian coordinates to a lon/lat grid.
+
+ Parameters
+ ----------
+ x: array_like of shape (N,), int, or float
+ X component of the point(s) expressed in kilometers.
+ y: array_like of shape (N,), int, or float
+ Y component of the point(s) expressed in kilometers.
+
+ Returns
+ -------
+ lon: array_like of shape (N,) or float
+ Longitude component of the transformed point(s) expressed in degrees.
+ lat: array_like of shape (N,) or float
+ Latitude component of the transformed point(s) expressed in degrees.
+
+ Examples
+ --------
+ Values pulled from haversine examples.
+
+ >>> x = -10077.330945462296
+ >>> y = 3892.875
+ >>> lon, lat = reverse_haversine(x, y)
+ >>> lon, lat
+ (-95.0, 35.0)
+
+ >>> x = np.array(
+ ... [
+ ... 1108.01755295,
+ ... 2190.70484658,
+ ... 3223.05300087,
+ ... 4180.69246988,
+ ... 5040.20418066,
+ ... 5779.42053216,
+ ... 6377.71302882,
+ ... 6816.26345487,
+ ... ]
+ ... )
+ >>> y = np.array(
+ ... [1112.25, 2224.5, 3336.75, 4449.0, 5561.25, 6673.5, 7785.75, 8898.0]
+ ... )
+ >>> lon, lat = reverse_haversine(x, y)
+ >>> lon, lat
+ (array([10., 20., 30., 40., 50., 60., 70., 80.]), array([10., 20., 30., 40., 50., 60., 70., 80.]))
+ """
+ C = 40041 # average circumference of earth in kilometers
+ lon = x * 360 / np.cos(y * np.pi / C) / C
+ lat = y * 360 / C
+ return lon, lat
+
+
+def geometric(x_km_new, y_km_new):
+ """
+ Turn longitudinal/latitudinal distance (km) lists into LineString for interpolation.
+
+ Parameters
+ ----------
+ x_km_new: List containing longitude coordinates of fronts in kilometers.
+ y_km_new: List containing latitude coordinates of fronts in kilometers.
+
+ Returns
+ -------
+ xy_linestring: LineString object containing coordinates of fronts in kilometers.
+ """
+ df_xy = pd.DataFrame(
+ list(zip(x_km_new, y_km_new)), columns=["Longitude_km", "Latitude_km"]
+ )
+ geometry = [xy for xy in zip(df_xy.Longitude_km, df_xy.Latitude_km)]
+ xy_linestring = LineString(geometry)
+ return xy_linestring
+
+
+def redistribute_vertices(xy_linestring, distance):
+ """
+ Interpolate x/y coordinates at a specified distance.
+
+ Parameters
+ ----------
+ xy_linestring: LineString object containing coordinates of fronts in kilometers.
+ distance: Distance at which to interpolate the x/y coordinates.
+
+ Returns
+ -------
+ xy_vertices: Normalized MultiLineString that contains the interpolated coordinates of fronts in kilometers.
+
+ Sources
+ -------
+ https://stackoverflow.com/questions/34906124/interpolating-every-x-distance-along-multiline-in-shapely/35025274#35025274
+ """
+ if xy_linestring.geom_type == "LineString":
+ num_vert = int(round(xy_linestring.length / distance))
+ if num_vert == 0:
+ num_vert = 1
+ return LineString(
+ [
+ xy_linestring.interpolate(float(n) / num_vert, normalized=True)
+ for n in range(num_vert + 1)
+ ]
+ )
+ elif xy_linestring.geom_type == "MultiLineString":
+ parts = [redistribute_vertices(part, distance) for part in xy_linestring]
+ return type(xy_linestring)([p for p in parts if not p.is_empty])
+ else:
+ raise ValueError("unhandled geometry %s", (xy_linestring.geom_type,))
+
+
+def reformat_fronts(fronts, front_types):
+ """
+ Reformat a front dataset, tensor, or array with a given set of front types.
+
+ Parameters
+ ----------
+ front_types: str or list of strs
+ Code(s) that determine how the dataset will be reformatted.
+ fronts: xarray Dataset or DataArray, tensor, or np.ndarray
+ Dataset containing the front data.
+ '''
+ Available options for individual front types (cannot be passed with any special codes):
+
+ Code (class #): Front Type
+ --------------------------
+ CF (1): Cold front
+ WF (2): Warm front
+ SF (3): Stationary front
+ OF (4): Occluded front
+ CF-F (5): Cold front (forming)
+ WF-F (6): Warm front (forming)
+ SF-F (7): Stationary front (forming)
+ OF-F (8): Occluded front (forming)
+ CF-D (9): Cold front (dissipating)
+ WF-D (10): Warm front (dissipating)
+ SF-D (11): Stationary front (dissipating)
+ OF-D (12): Occluded front (dissipating)
+ INST (13): Squall line ????
+ TROF (14): Trough
+ TT (15): Tropical Trough
+ DL (16): Dryline
+
+
+ Special codes (cannot be passed with any individual front codes):
+ -----------------------------------------------------------------
+ F_BIN (1 class): 1-4, but treat all front types as one type.
+ (1): CF, WF, SF, OF
+
+ MERGED-F (4 classes): 1-12, but treat forming and dissipating fronts as standard fronts.
+ (1): CF, CF-F, CF-D
+ (2): WF, WF-F, WF-D
+ (3): SF, SF-F, SF-D
+ (4): OF, OF-F, OF-D
+
+ MERGED-F_BIN (1 class): 1-12, but treat all front types and stages as one type. This means that classes 1-12 will all be one class (1).
+ (1): CF, CF-F, CF-D, WF, WF-F, WF-D, SF, SF-F, SF-D, OF, OF-F, OF-D
+
+ MERGED-T (1 class): 14-15, but treat troughs and tropical troughs as the same. In other words, TT (15) becomes TROF (14).
+ (1): TROF, TT
+
+ MERGED-ALL (7 classes): 1-16, but make the changes in the MERGED-F and MERGED-T codes.
+ (1): CF, CF-F, CF-D
+ (2): WF, WF-F, WF-D
+ (3): SF, SF-F, SF-D
+ (4): OF, OF-F, OF-D
+ (5): TROF, TT
+ (6): INST
+ (7): DL
+
+ **** NOTE - Class 0 is always treated as 'no front'.
+ '''
+
+ Returns
+ -------
+ fronts_ds: xr.Dataset
+ Reformatted dataset based on the provided code(s).
+ """
+
+ if type(front_types) == str:
+ front_types = [
+ front_types,
+ ]
+
+ fronts_argument_type = type(fronts)
+
+ if fronts_argument_type == xr.DataArray or fronts_argument_type == xr.Dataset:
+ where_function = xr.where
+ elif fronts_argument_type == np.ndarray:
+ where_function = np.where
+ else:
+ where_function = tf.where
+
+ front_types_classes = {
+ "CF": 1,
+ "WF": 2,
+ "SF": 3,
+ "OF": 4,
+ "CF-F": 5,
+ "WF-F": 6,
+ "SF-F": 7,
+ "OF-F": 8,
+ "CF-D": 9,
+ "WF-D": 10,
+ "SF-D": 11,
+ "OF-D": 12,
+ "INST": 13,
+ "TROF": 14,
+ "TT": 15,
+ "DL": 16,
+ }
+
+ if front_types == [
+ "F_BIN",
+ ]:
+ fronts = where_function(fronts > 4, 0, fronts) # Classes 5-16 are removed
+ fronts = where_function(fronts > 0, 1, fronts) # Merge 1-4 into one class
+
+ labels = [
+ "CF-WF-SF-OF",
+ ]
+ num_types = 1
+
+ elif front_types == ["MERGED-F"]:
+ fronts = where_function(
+ fronts == 5, 1, fronts
+ ) # Forming cold front ---> cold front
+ fronts = where_function(
+ fronts == 6, 2, fronts
+ ) # Forming warm front ---> warm front
+ fronts = where_function(
+ fronts == 7, 3, fronts
+ ) # Forming stationary front ---> stationary front
+ fronts = where_function(
+ fronts == 8, 4, fronts
+ ) # Forming occluded front ---> occluded front
+ fronts = where_function(
+ fronts == 9, 1, fronts
+ ) # Dying cold front ---> cold front
+ fronts = where_function(
+ fronts == 10, 2, fronts
+ ) # Dying warm front ---> warm front
+ fronts = where_function(
+ fronts == 11, 3, fronts
+ ) # Dying stationary front ---> stationary front
+ fronts = where_function(
+ fronts == 12, 4, fronts
+ ) # Dying occluded front ---> occluded front
+ fronts = where_function(fronts > 4, 0, fronts) # Remove all other fronts
+
+ labels = ["CF_any", "WF_any", "SF_any", "OF_any"]
+ num_types = 4
+
+ elif front_types == ["MERGED-F_BIN"]:
+ fronts = where_function(fronts > 12, 0, fronts) # Classes 13-16 are removed
+ fronts = where_function(
+ fronts > 0, 1, fronts
+ ) # Classes 1-12 are merged into one class
+
+ labels = [
+ "CF-WF-SF-OF_any",
+ ]
+ num_types = 1
+
+ elif front_types == ["MERGED-T"]:
+ fronts = where_function(fronts < 14, 0, fronts) # Remove classes 1-13
+
+ # Merge troughs into one class
+ fronts = where_function(fronts == 14, 1, fronts)
+ fronts = where_function(fronts == 15, 1, fronts)
+
+ fronts = where_function(fronts == 16, 0, fronts) # Remove drylines
+
+ labels = [
+ "TR_any",
+ ]
+ num_types = 1
+
+ elif front_types == ["MERGED-ALL"]:
+ fronts = where_function(
+ fronts == 5, 1, fronts
+ ) # Forming cold front ---> cold front
+ fronts = where_function(
+ fronts == 6, 2, fronts
+ ) # Forming warm front ---> warm front
+ fronts = where_function(
+ fronts == 7, 3, fronts
+ ) # Forming stationary front ---> stationary front
+ fronts = where_function(
+ fronts == 8, 4, fronts
+ ) # Forming occluded front ---> occluded front
+ fronts = where_function(
+ fronts == 9, 1, fronts
+ ) # Dying cold front ---> cold front
+ fronts = where_function(
+ fronts == 10, 2, fronts
+ ) # Dying warm front ---> warm front
+ fronts = where_function(
+ fronts == 11, 3, fronts
+ ) # Dying stationary front ---> stationary front
+ fronts = where_function(
+ fronts == 12, 4, fronts
+ ) # Dying occluded front ---> occluded front
+
+ # Merge troughs together into class 5
+ fronts = where_function(fronts == 14, 5, fronts)
+ fronts = where_function(fronts == 15, 5, fronts)
+
+ fronts = where_function(
+ fronts == 13, 6, fronts
+ ) # Move outflow boundaries to class 6
+ fronts = where_function(fronts == 16, 7, fronts) # Move drylines to class 7
+
+ labels = ["CF_any", "WF_any", "SF_any", "OF_any", "TR_any", "INST", "DL"]
+ num_types = 7
+
+ else:
+ # Select the front types that are being used to pull their class identifiers
+ filtered_front_types = dict(
+ sorted(
+ dict(
+ [
+ (i, front_types_classes[i])
+ for i in front_types_classes
+ if i in set(front_types)
+ ]
+ ).items(),
+ key=lambda item: item[1],
+ )
+ )
+ front_types, num_types = (
+ list(filtered_front_types.keys()),
+ len(filtered_front_types.keys()),
+ )
+
+ for i in range(num_types):
+ if i + 1 != front_types_classes[front_types[i]]:
+ fronts = where_function(fronts == i + 1, 0, fronts)
+ fronts = where_function(
+ fronts == front_types_classes[front_types[i]], i + 1, fronts
+ ) # Reformat front classes
+
+ fronts = where_function(
+ fronts > num_types, 0, fronts
+ ) # Remove unused front types
+
+ labels = front_types
+
+ if fronts_argument_type == xr.Dataset or fronts_argument_type == xr.DataArray:
+ fronts.attrs["front_types"] = front_types
+ fronts.attrs["num_front_types"] = num_types
+ fronts.attrs["labels"] = labels
+
+ return fronts
+
+
+def normalize_dataset(
+ ds, method="standard", normalization_parameters=constants.NORMALIZATION_PARAMS
+) -> xr.Dataset:
+ """
+ Normalizes variables in an xarray dataset. This function can also accept xarray datasets for GOES satellite data.
+
+ Parameters
+ ----------
+ ds: xarray dataset
+ Dataset containing variables to normalize.
+ method: 'standard', 'standard_weighted', 'min-max'
+ Normalization method to perform on the variables.
+ - 'standard': Standard z-score normalization.
+ - 'standard_weighted': Standard z-score normalization with latitude-weighted means and standard deviations.
+ - 'min-max': Min-max normalization.
+ normalization_parameters: dict
+ Dictionary containing parameters for normalization.
+
+ Returns
+ -------
+ normalized_ds: xarray dataset
+ Normalized xarray dataset.
+ """
+
+ ds_copy = ds.copy()
+ ds.close()
+
+ variables = list(ds_copy.keys())
+
+ is_satellite_dataset = "band_" in variables[0] # check for satellite variables
+
+ try:
+ if is_satellite_dataset:
+ norm_params = xr.Dataset(
+ data_vars={
+ var: ("param", normalization_parameters["%s" % var])
+ for var in variables
+ },
+ coords={
+ "param": [
+ "min",
+ "max",
+ "mean",
+ "std",
+ "mean_weighted",
+ "std_weighted",
+ ]
+ },
+ )
+ else:
+ pressure_levels = ds_copy["pressure_level"].values.astype(
+ int
+ ) # TODO: will not work with surface data
+
+ norm_params = xr.Dataset(
+ data_vars={
+ var: (
+ ("pressure_level", "param"),
+ [
+ normalization_parameters["%s_%s" % (var, lvl)]
+ for lvl in pressure_levels
+ ],
+ )
+ for var in variables
+ },
+ coords={
+ "param": [
+ "min",
+ "max",
+ "mean",
+ "std",
+ "mean_weighted",
+ "std_weighted",
+ ],
+ "pressure_level": pressure_levels,
+ },
+ )
+ except (
+ ValueError
+ ): # models before the 2025.1.10 update only have min and max values
+ pressure_levels = ds_copy["pressure_level"].values.astype(
+ int
+ ) # TODO: will not work with surface data
+ norm_params = xr.Dataset(
+ data_vars={
+ var: (
+ ("pressure_level", "param"),
+ [
+ normalization_parameters["%s_%s" % (var, lvl)]
+ for lvl in pressure_levels
+ ],
+ )
+ for var in variables
+ },
+ coords={"param": ["max", "min"], "pressure_level": pressure_levels},
+ )
+
+ if method == "min-max":
+ normalized_ds = (ds_copy - norm_params.sel(param="min")) / (
+ norm_params.sel(param="max") - norm_params.sel(param="min")
+ )
+ elif method == "standard":
+ normalized_ds = (ds_copy - norm_params.sel(param="mean")) / norm_params.sel(
+ param="std"
+ )
+ elif method == "standard_weighted":
+ normalized_ds = (
+ ds_copy - norm_params.sel(param="mean_weighted")
+ ) / norm_params.sel(param="std_weighted")
+ else:
+ raise ValueError(
+ "Unrecognized normalization method: %s. Valid normalization methods are 'min-max', 'standard', 'standard-weighted'."
+ % method
+ )
+
+ return normalized_ds
+
+
+def combine_datasets(tf_files: list[str]):
+ """
+ Combine many tensorflow datasets into one entire dataset.
+
+ Returns
+ -------
+ dataset: tf.data.Dataset object
+ Concatenated tensorflow dataset.
+ """
+ dataset = tf.data.Dataset.load(tf_files[0])
+ for file in tf_files[1:]:
+ dataset = dataset.concatenate(tf.data.Dataset.load(file))
+
+ return dataset
+
+
+def lambert_conformal_to_cartesian(
+ lon: np.ndarray | tuple | list | int | float,
+ lat: np.ndarray | tuple | list | int | float,
+ std_parallels: tuple | list = (20.0, 50.0),
+ lon_ref: int | float = 0.0,
+ lat_ref: int | float = 0.0,
+):
+ """
+ Transform points on a Lambert Conformal lat/lon grid to cartesian coordinates.
+
+ Parameters
+ ----------
+ lon: array_like of shape (N,), int, or float
+ Longitude point(s) expressed as degrees.
+ lat: array_like of shape (N,), int, or float
+ Latitude point(s) expressed as degrees.
+ std_parallels: tuple or list of 2 ints or floats
+ Standard parallels to use in the coordinate transformation, expressed as degrees.
+ lon_ref: int or float
+ Reference longitude point expressed as degrees.
+ lat_ref: int or float
+ Reference latitude point expressed as degrees.
+
+ Returns
+ -------
+ x: array_like of shape (N,) or float
+ X-component of the transformed coordinates, expressed as meters.
+ y: array_like of shape (N,) or float
+ Y-component of the transformed coordinates, expressed as meters.
+
+ Examples
+ --------
+ * Using parameters from example on Page 295 of Snyder 1987 (except the output here is expressed as meters):
+ >>> x, y = lambert_conformal_to_cartesian(
+ ... lon=-75, lat=35, std_parallels=(33, 45), lon_ref=-96, lat_ref=23
+ ... )
+ >>> x, y
+ (1890206.4076610378, 1568668.1244433122)
+
+ * Same as above but with longitudes expressed from 0 to 360 degrees east:
+ >>> x, y = lambert_conformal_to_cartesian(
+ ... lon=285, lat=35, std_parallels=(33, 45), lon_ref=264, lat_ref=23
+ ... )
+ >>> x, y
+ (1890206.4076610343, 1568668.1244433112)
+
+ References
+ ----------
+ * Snyder 1987: https://doi.org/10.3133/pp1395
+
+ Notes
+ -----
+ lon and lon_ref must be both expressed in the same longitude range (e.g. -180 to 180 degrees or 0 to 360 degrees)
+ to get correct values for x and y.
+ """
+
+ R = 6371229 # radius of earth (meters)
+
+ # Points and standard parallels need to be expressed as radians for the transformation formulas
+ lon = np.radians(lon)
+ lon_ref = np.radians(lon_ref)
+ lat = np.radians(lat)
+ lat_ref = np.radians(lat_ref)
+ std_parallels = np.radians(std_parallels)
+
+ if std_parallels[0] == std_parallels[1]:
+ n = np.sin(std_parallels[0])
+ else:
+ n = np.divide(
+ np.log(np.cos(std_parallels[0]) / np.cos(std_parallels[1])),
+ np.log(
+ np.tan(np.pi / 4 + std_parallels[1] / 2)
+ / np.tan(np.pi / 4 + std_parallels[0] / 2)
+ ),
+ )
+ F = (
+ np.cos(std_parallels[0])
+ * np.power(np.tan(np.pi / 4 + std_parallels[0] / 2), n)
+ / n
+ )
+ rho = R * F / np.power(np.tan(np.pi / 4 + lat / 2), n)
+ rho0 = R * F / np.power(np.tan(np.pi / 4 + lat_ref / 2), n)
+
+ x = rho * np.sin(n * (lon - lon_ref))
+ y = rho0 - rho * np.cos(n * (lon - lon_ref))
+
+ return x, y
+
+
+def mask_xarray_dataset(ds, mask, lon="longitude", lat="latitude"):
+ """
+ Apply a geospatial mask from the regionmask package to an Xarray dataset.
+
+ Parameters
+ ----------
+ ds: xarray Dataset
+ Xarray dataset that must have longitude and latitude dimensions.
+ mask: str
+ Geospatial mask to apply to the dataset.
+ lon: str
+ Longitude dimension key in the xarray dataset.
+ lat: str
+ Latitude dimension key in the xarray dataset.
+
+ Returns
+ -------
+ masked_ds: xarray Dataset
+ Masked xarray dataset.
+ """
+ # {region_key: region_index}
+ regions_crossing_prime_meridian = ["north_atlantic_ocean"]
+ ocean_basins = {
+ "arctic_ocean": [0, 13, 31, 32, 40, 47, 56, 57],
+ "north_atlantic_ocean": [2, 37, 60, 83, 88, 99, 100],
+ "south_atlantic_ocean": [
+ 6,
+ ],
+ "indian_ocean": [5, 10, 12, 36, 43, 44, 50, 52, 61, 90, 105],
+ "north_pacific_ocean": [3, 8, 20, 59],
+ "south_pacific_ocean": [4, 9, 27, 74, 80, 85, 86],
+ "southern_ocean": [1, 23, 26, 38, 53, 54, 58],
+ }
+
+ region_is_ocean_basin = mask in ocean_basins
+ region_crosses_prime_meridian = mask in regions_crossing_prime_meridian
+
+ if region_is_ocean_basin:
+ regions = regionmask.defined_regions.natural_earth_v5_1_2.ocean_basins_50
+ indices = ocean_basins[mask]
+
+ region_mask = xr.merge(
+ [
+ (regions.mask(ds[lon], ds[lat]) == i).expand_dims(
+ {"index": np.atleast_1d(i)}
+ )
+ for i in indices
+ ]
+ ).max("index")["mask"]
+ masked_ds = ds.where(region_mask, 1, 0)
+
+ if region_crosses_prime_meridian:
+ lons = masked_ds[lon]
+ lon_east_hemi, lon_west_hemi = lons[lons <= 180], lons[lons > 180]
+ masked_ds = masked_ds.reindex(
+ {lon: np.concatenate([lon_west_hemi, lon_east_hemi])}
+ )
+ new_lons = np.concatenate([lon_west_hemi, lon_east_hemi + 360])
+ masked_ds[lon] = new_lons
+
+ return masked_ds
diff --git a/src/fronts/utils/file_manager.py b/src/fronts/utils/file_manager.py
new file mode 100644
index 0000000..328b7b7
--- /dev/null
+++ b/src/fronts/utils/file_manager.py
@@ -0,0 +1,634 @@
+"""
+Functions in this code manage data files and models.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2025.2.9
+"""
+
+import argparse
+from glob import glob
+import os
+import numpy as np
+import pandas as pd
+import shutil
+import tarfile
+
+
+def compress_files(
+ main_dir: str,
+ glob_file_string: str,
+ tar_filename: str,
+ algorithm: str = "gz",
+ remove_files: bool = False,
+ status_printout: bool = True,
+):
+ """
+ Compress files into a TAR file.
+
+ Parameters
+ ----------
+ main_dir: str
+ Main directory where the files are located and where the TAR file will be saved.
+ glob_file_string: str
+ String of the names of the files to compress.
+ tar_filename: str
+ Name of the compressed TAR file that will be made. Do not include the .tar.gz extension in the name, this is added automatically.
+ algorithm: str
+ Compression algorithm to use when generating the TAR file.
+ remove_files: bool
+ Setting this to true will remove the files after they have been compressed to a TAR file.
+ status_printout: bool
+ Setting this to true will provide printouts of the status of the compression.
+
+ Examples
+ --------
+ <<<<< start example >>>>
+
+ import fronts.file_manager as fm
+
+ main_dir = 'C:/Users/username/data_files'
+ glob_file_string = '*matching_string.pkl'
+ tar_filename = 'matching_files' # Do not add the .tar.gz extension, this is done automatically
+
+ compress_files(main_dir, glob_file_string, tar_filename, remove_files=True, status_printout=False) # Compress files and remove them after compression into a TAR file
+
+ <<<<< end example >>>>>
+ """
+
+ ########################################### Check the parameters for errors ########################################
+ if not isinstance(main_dir, str):
+ raise TypeError(f"main_dir must be a string, received {type(main_dir)}")
+ if not isinstance(glob_file_string, str):
+ raise TypeError(
+ f"glob_file_string must be a string, received {type(glob_file_string)}"
+ )
+ if not isinstance(tar_filename, str):
+ raise TypeError(f"tar_filename must be a string, received {type(tar_filename)}")
+ if not isinstance(remove_files, bool):
+ raise TypeError(
+ f"remove_files must be a boolean, received {type(remove_files)}"
+ )
+ if not isinstance(status_printout, bool):
+ raise TypeError(
+ f"status_printout must be a boolean, received {type(status_printout)}"
+ )
+ ####################################################################################################################
+
+ uncompressed_size = 0 # MB
+
+ ### Gather a list of files containing the specified string ###
+ files = list(sorted(glob(f"{main_dir}/{glob_file_string}")))
+ if len(files) == 0:
+ raise OSError("No files found")
+ else:
+ print(f"{len(files)} files found")
+
+ num_files = len(files) # Total number of files
+
+ ### Create the TAR file ###
+ with tarfile.open(
+ f"{main_dir}/{tar_filename}.tar.{algorithm}", f"w:{algorithm}"
+ ) as tarF:
+ ### Iterate through all the available files ###
+ for file in range(num_files):
+ tarF.add(
+ files[file], arcname=files[file].replace(main_dir, "")
+ ) # Add the file to the TAR file
+ tarF_size = (
+ os.path.getsize(f"{main_dir}/{tar_filename}.tar.{algorithm}") / 1e6
+ ) # Compressed size of the files within the TAR file (megabytes)
+ uncompressed_size += (
+ os.path.getsize(files[file]) / 1e6
+ ) # Uncompressed size of the files within the TAR file (megabytes)
+
+ ### Print out the current status of the compression (if enabled) ###
+ if status_printout:
+ print(
+ f"({file + 1}/{num_files}) {uncompressed_size:,.2f} MB ---> {tarF_size:,.2f} MB ({100 * (1 - (tarF_size / uncompressed_size)):.1f}% compression ratio)",
+ end="\r",
+ )
+
+ # Completion message
+ print(
+ f"Successfully compressed {len(files)} files: ",
+ f"{uncompressed_size:,.2f} MB ---> {tarF_size:,.2f} MB ({100 * (1 - (tarF_size / uncompressed_size)):.1f}% compression ratio)",
+ )
+
+ ### Remove the files that were added to the TAR archive (if enabled; does NOT affect the contents of the TAR file just created) ###
+ if remove_files:
+ for file in files:
+ os.remove(file)
+ print(f"Successfully deleted {len(files)} files")
+
+
+def delete_grouped_files(main_dir: str, glob_file_string: str, num_subdir: int):
+ """
+ Deletes grouped files with names matching given strings.
+
+ Parameters
+ ----------
+ main_dir: str
+ Main directory or directories where the grouped files are located.
+ glob_file_string: str
+ String of the names of the files to delete.
+ num_subdir: int
+ Number of subdirectory layers in the main directory.
+
+ Examples
+ --------
+ <<<<< start example >>>>>
+
+ import fronts.file_manager as fm
+
+ main_dir = 'C:/Users/username/data_files'
+ glob_file_string = '*matching_string.pkl'
+ num_subdir = 3 # Check in the 3rd level of the directories within the main directory
+
+ fm.delete_grouped_files(main_dir, glob_file_string, num_subdir)
+
+ <<<<< end example >>>>>
+ """
+
+ ########################################### Check the parameters for errors ########################################
+ if not isinstance(main_dir, str):
+ raise TypeError(f"main_dir must be a string, received {type(main_dir)}")
+ if not isinstance(glob_file_string, str):
+ raise TypeError(
+ f"glob_file_string must be a string, received {type(glob_file_string)}"
+ )
+ if not isinstance(num_subdir, int):
+ raise TypeError(f"num_subdir must be an integer, received {type(num_subdir)}")
+ ####################################################################################################################
+
+ subdir_string = (
+ "" # This string will be modified depending on the provided value of num_subdir
+ )
+ for _ in range(num_subdir):
+ subdir_string += "/*"
+ subdir_string += "/"
+ glob_file_string = (
+ subdir_string + glob_file_string
+ ) # String that will be used to match with patterns in filenames
+
+ files_to_delete = list(
+ sorted(glob("%s%s" % (main_dir, glob_file_string)))
+ ) # Search for files in the given directory that have patterns matching the file string
+
+ # Delete all the files
+ print("Deleting %d files...." % len(files_to_delete), end="")
+ for file in files_to_delete:
+ try:
+ os.remove(file)
+ except PermissionError:
+ shutil.rmtree(file)
+ print("done")
+
+
+def extract_tarfile(main_dir: str, tar_filename: str):
+ """
+ Extract all the contents of a TAR file.
+
+ Parameters
+ ----------
+ main_dir: str
+ Main directory where the TAR file is located. This is also where the extracted files will be placed.
+ tar_filename: str
+ Name of the compressed TAR file. Do NOT include the .tar.gz extension.
+
+ Examples
+ --------
+ <<<<< start example >>>>
+
+ import fronts.file_manager as fm
+
+ main_dir = 'C:/Users/username/data_files'
+ tar_filename = 'foo_tarfile' # Do not add the .tar.gz extension
+
+ fm.extract_tarfile(main_dir, glob_file_string, tar_filename, remove_files=True, status_printout=False) # Compress files and remove them after compression into a TAR file
+
+ <<<<< end example >>>>>
+ """
+
+ ########################################### Check the parameters for errors ########################################
+ if not isinstance(main_dir, str):
+ raise TypeError(f"main_dir must be a string, received {type(main_dir)}")
+ if not isinstance(tar_filename, str):
+ raise TypeError(f"tar_filename must be a string, received {type(tar_filename)}")
+ ####################################################################################################################
+
+ with tarfile.open(f"{main_dir}/{tar_filename}.tar.gz", "r") as tarF:
+ tarF.extractall(main_dir)
+ print(f"Successfully extracted {main_dir}/{tar_filename}")
+
+
+class DataFileLoader:
+ """
+ Objects that loads and manages various types of files containing weather data.
+ """
+
+ def __init__(
+ self,
+ file_dir: str,
+ data_type: str,
+ file_format: str,
+ years=None,
+ months=None,
+ days=None,
+ hours=None,
+ domains=None,
+ ):
+ """
+ file_dir: str
+ Parent directory for the first set of various data files to load.
+ data_type: str
+ The source/type of the data files to load. Options for the data sources are: 'era5', 'fronts', 'gdas', 'gfs', 'nam-12km', 'satellite'.
+ file_format: str
+ Formatting of the data files to load. Options for the file/dataset type string are: 'grib', 'netcdf', 'tensorflow'.
+ Note that this CANNOT be changed after it has been set, you must create an entirely separate DataFileLoader
+ object to load a separate file type.
+ years: int or iterable of ints, optional
+ Year(s) to select from the available files.
+ months: int or iterable of ints, optional
+ Month(s) to select from the avilable files.
+ days: int or iterable of ints, optional
+ Day(s) to select from the available files.
+ hours: int or iterable of ints, optional
+ Hour(s) to select from the available files.
+ domains: str or iterable of strs, optional
+ Domain(s) to select from the available files.
+ """
+
+ self._file_format = file_format
+ self._years = years
+ self._months = months
+ self._days = days
+ self._hours = hours
+ self._domains = domains
+
+ self._format_args()
+
+ if data_type == "fronts" and self._file_format != "tensorflow":
+ data_type = "FrontObjects"
+
+ if self._file_format == "grib":
+ self._file_extension = ".grib"
+ elif self._file_format == "netcdf":
+ self._file_extension = ".nc"
+ elif self._file_format == "tensorflow":
+ self._file_extension = "_tf"
+ else:
+ raise ValueError("Unknown file type: %s" % self._file_format)
+
+ if self._file_format in ["grib", "netcdf"]:
+ if data_type in ["era5", "gfs"]:
+ glob_strs = [
+ f"{file_dir}/{yr}{mo}/{data_type}_{yr}{mo}{dy}{hr}_{domain}{self._file_extension}"
+ for yr in self._years
+ for mo in self._months
+ for dy in self._days
+ for hr in self._hours
+ for domain in self._domains
+ ]
+ else:
+ glob_strs = [
+ f"{file_dir}/{yr}{mo}/{data_type}_{yr}{mo}{dy}{hr}_*{self._file_extension}"
+ for yr in self._years
+ for mo in self._months
+ for dy in self._days
+ for hr in self._hours
+ ]
+ else:
+ glob_strs = [
+ f"{file_dir}/{data_type}_{yr}{mo}{self._file_extension}"
+ for yr in self._years
+ for mo in self._months
+ ]
+
+ # remove consecutive asterisks to prevent recursive searches
+ glob_strs = [
+ glob_str.replace("****", "*").replace("***", "*").replace("**", "*")
+ for glob_str in glob_strs
+ ]
+
+ self.files = []
+ for glob_str in glob_strs:
+ self.files.extend(glob(glob_str))
+ self.files = [sorted(self.files)]
+
+ def add_file_list(self, file_dir, data_type, ignore_domain=False):
+ """
+ Add another file list.
+
+ ignore_domain: bool
+ Do not consider the domain when matching files (i.e. match files regardless of domain). Only use this when
+ trying to pair front files with ERA5 or some other netcdf files, otherwise you may run into bugs.
+ """
+ num_lists = len(self.files) # the current number of file lists
+
+ if data_type == "fronts" and self._file_format != "tensorflow":
+ data_type = "FrontObjects"
+ new_domains = [
+ "full",
+ ]
+ elif data_type == "goes-merged":
+ new_domains = [
+ "full",
+ ]
+ elif data_type == "hrrr":
+ new_domains = [
+ "hrrr",
+ ]
+ else:
+ new_domains = "global"
+
+ if self._file_format in ["grib", "netcdf"]:
+ if data_type in ["era5", "gfs"]:
+ glob_strs = [
+ f"{file_dir}/{yr}{mo}/{data_type}_{yr}{mo}{dy}{hr}_{domain}{self._file_extension}"
+ for yr in self._years
+ for mo in self._months
+ for dy in self._days
+ for hr in self._hours
+ for domain in new_domains
+ ]
+ else:
+ glob_strs = [
+ f"{file_dir}/{yr}{mo}/{data_type}_{yr}{mo}{dy}{hr}_*{self._file_extension}"
+ for yr in self._years
+ for mo in self._months
+ for dy in self._days
+ for hr in self._hours
+ ]
+ else: # tensorflow
+ glob_strs = [
+ f"{file_dir}/{data_type}_{yr}{mo}{self._file_extension}"
+ for yr in self._years
+ for mo in self._months
+ ]
+
+ # remove consecutive asterisks to prevent recursive searches
+ glob_strs = [
+ glob_str.replace("****", "*").replace("***", "*").replace("**", "*")
+ for glob_str in glob_strs
+ ]
+
+ current_file_list = []
+ for glob_str in glob_strs:
+ current_file_list.extend(glob(glob_str))
+ current_file_list = sorted(current_file_list)
+
+ # file basenames
+ control_basename_list = [os.path.basename(file) for file in self.files[0]]
+ current_basename_list = [os.path.basename(file) for file in current_file_list]
+
+ # pull the data sources of the files so we can properly add forecast hours or domains to the file info lists
+ control_data_source = control_basename_list[0].split("_")[0]
+ current_data_source = current_basename_list[0].split("_")[0]
+
+ # get details contained within the file names
+ control_basename_info = [
+ file.replace(self._file_extension, "").split("_")[1:]
+ for file in control_basename_list
+ ]
+ current_basename_info = [
+ file.replace(self._file_extension, "").split("_")[1:]
+ for file in current_basename_list
+ ]
+
+ if len(control_basename_info[0]) == 2:
+ if control_data_source in ["era5", "FrontObjects", "goes-merged"]:
+ [file_info.insert(1, "f000") for file_info in control_basename_info]
+ else:
+ [
+ file_info.insert(2, control_data_source)
+ for file_info in control_basename_info
+ ]
+
+ if len(current_basename_info[0]) == 2:
+ if current_data_source in ["era5", "FrontObjects", "goes-merged"]:
+ [file_info.insert(1, "f000") for file_info in current_basename_info]
+ else:
+ [
+ file_info.insert(2, current_data_source)
+ for file_info in current_basename_info
+ ]
+
+ # remove the domain from the information arrays (if requested)
+ if ignore_domain:
+ control_basename_info = [
+ file_info[:-1] for file_info in control_basename_info
+ ]
+ current_basename_info = [
+ file_info[:-1] for file_info in current_basename_info
+ ]
+
+ # find indices where all file details match
+ index_pairs = np.array(
+ [
+ [data_idx, current_basename_info.index(file_info)]
+ for data_idx, file_info in enumerate(control_basename_info)
+ if file_info in current_basename_info
+ ]
+ )
+
+ # filter files and add new list
+ self.files = [
+ [self.files[list_num][i] for i in index_pairs[:, 0]]
+ for list_num in range(num_lists)
+ ]
+ current_file_list = [current_file_list[i] for i in index_pairs[:, 1]]
+ self.files.append(current_file_list)
+
+ def _format_args(self):
+ if isinstance(self._years, (int, np.int64)):
+ self._years = [
+ f"{self._years:04d}",
+ ]
+ elif self._years is not None:
+ self._years = ["%04d" % yr for yr in self._years]
+ else:
+ self._years = [
+ "*",
+ ]
+
+ if isinstance(self._months, int):
+ self._months = [
+ f"{self._months:02d}",
+ ]
+ elif self._months is not None:
+ self._months = ["%02d" % mo for mo in self._months]
+ else:
+ self._months = [
+ "*",
+ ]
+
+ if isinstance(self._days, int):
+ self._days = [
+ f"{self._days:02d}",
+ ]
+ elif self._days is not None:
+ self._days = ["%02d" % dy for dy in self._days]
+ else:
+ self._days = [
+ "*",
+ ]
+
+ if isinstance(self._hours, int):
+ self._hours = [
+ f"{self._hours:02d}",
+ ]
+ elif self._hours is not None:
+ self._hours = ["%02d" % hr for hr in self._hours]
+ else:
+ self._hours = [
+ "*",
+ ]
+
+ if isinstance(self._domains, str):
+ self._domains = [
+ self._domains,
+ ]
+ elif self._domains is None:
+ self._domains = [
+ "*",
+ ]
+
+
+def load_model(model_number: int, model_dir: str):
+ """
+ Load a saved model.
+
+ Parameters
+ ----------
+ model_number: int
+ Slurm job number for the model. This is the number in the model's filename.
+ model_dir: str
+ Main directory for the models.
+ """
+
+ ######################################### Check the parameters for errors ##########################################
+ if not isinstance(model_number, int):
+ raise TypeError(
+ f"model_number must be an integer, received {type(model_number)}"
+ )
+ if not isinstance(model_dir, str):
+ raise TypeError(f"model_dir must be a string, received {type(model_dir)}")
+ ####################################################################################################################
+
+ from tensorflow.keras.models import load_model as lm
+ from fronts.model import custom_activations, losses, metrics
+
+ model_path = f"{model_dir}/model_{model_number}/model_{model_number}.h5"
+ model_properties = pd.read_pickle(
+ f"{model_dir}/model_{model_number}/model_{model_number}_properties.pkl"
+ )
+
+ custom_objects = {}
+
+ loss_args = model_properties["loss_args"]
+ loss_parent_string = model_properties["loss_parent_string"]
+ loss_child_string = model_properties["loss_child_string"]
+
+ metric_args = model_properties["metric_args"]
+ metric_parent_string = model_properties["metric_parent_string"]
+ metric_child_string = model_properties["metric_child_string"]
+
+ # add the loss and metric functions to the custom_objects dictionary
+ custom_objects[loss_child_string] = getattr(losses, loss_parent_string)(**loss_args)
+ custom_objects[metric_child_string] = getattr(metrics, metric_parent_string)(
+ **metric_args
+ )
+
+ # add the activation function to the custom_objects dictionary
+ activation_string = model_properties["activation"]
+ if activation_string in [
+ "elliott",
+ "gaussian",
+ "gcu",
+ "hexpo",
+ "isigmoid",
+ "lisht",
+ "psigmoid",
+ "ptanh",
+ "ptelu",
+ "resech",
+ "smelu",
+ "snake",
+ "srs",
+ "stanh",
+ ]:
+ if activation_string == "elliott":
+ activation = custom_activations.Elliott()
+ elif activation_string == "gaussian":
+ activation = custom_activations.Gaussian()
+ elif activation_string == "gcu":
+ activation = custom_activations.GCU()
+ elif activation_string == "hexpo":
+ activation = custom_activations.Hexpo()
+ elif activation_string == "isigmoid":
+ activation = custom_activations.ISigmoid()
+ elif activation_string == "lisht":
+ activation = custom_activations.LiSHT()
+ elif activation_string == "psigmoid":
+ activation = custom_activations.PSigmoid()
+ elif activation_string == "ptanh":
+ activation = custom_activations.PTanh()
+ elif activation_string == "ptelu":
+ activation = custom_activations.PTELU()
+ elif activation_string == "resech":
+ activation = custom_activations.ReSech()
+ elif activation_string == "smelu":
+ activation = custom_activations.SmeLU()
+ elif activation_string == "snake":
+ activation = custom_activations.Snake()
+ elif activation_string == "srs":
+ activation = custom_activations.SRS()
+ else: # activation_string == "stanh"
+ activation = custom_activations.STanh()
+ custom_objects[activation.__class__.__name__] = activation
+
+ return lm(model_path, custom_objects=custom_objects)
+
+
+if __name__ == "__main__":
+ """
+ Warnings
+ Do not use leading zeros when declaring the month, day, and hour in 'date'. (ex: if the day is 2, do not type 02)
+ Longitude values in the 'new_extent' argument must in the 360-degree coordinate system.
+ """
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--compress_files", action="store_true", help="Compress files")
+ parser.add_argument(
+ "--delete_grouped_files", action="store_true", help="Delete a set of files"
+ )
+ parser.add_argument(
+ "--extract_tarfile", action="store_true", help="Extract a TAR file"
+ )
+ parser.add_argument(
+ "--glob_file_string",
+ type=str,
+ help="String of the names of the files to compress or delete.",
+ )
+ parser.add_argument(
+ "--main_dir",
+ type=str,
+ help="Main directory for subdirectory creation or where the files in question are located.",
+ )
+ parser.add_argument(
+ "--num_subdir",
+ type=int,
+ help="Number of subdirectory layers in the main directory.",
+ )
+ parser.add_argument("--tar_filename", type=str, help="Name of the TAR file.")
+ args = parser.parse_args()
+ provided_arguments = vars(args)
+
+ if args.compress_files:
+ compress_files(args.main_dir, args.glob_file_string, args.tar_filename)
+
+ if args.delete_grouped_files:
+ delete_grouped_files(args.main_dir, args.glob_file_string, args.num_subdir)
+
+ if args.extract_tarfile:
+ extract_tarfile(args.main_dir, args.tar_filename)
diff --git a/src/fronts/utils/keras_builders.py b/src/fronts/utils/keras_builders.py
new file mode 100644
index 0000000..b45e096
--- /dev/null
+++ b/src/fronts/utils/keras_builders.py
@@ -0,0 +1,380 @@
+import dataclasses
+from typing import Literal, Any, TypeVar, Generic
+import tensorflow as tf
+from fronts.layers import activations, losses, metrics
+
+T = TypeVar("T")
+
+
+@dataclasses.dataclass
+class BaseConfig(Generic[T]):
+ """Base configuration class for building objects from a registry.
+
+ Type parameter T is the type of object this config builds.
+ """
+
+ name: str
+ config: dict[str, Any]
+
+ # Subclasses must define this
+ @property
+ def registry(self) -> dict[str, type]:
+ """The registry mapping string names to classes.
+
+ Returns:
+ A dictionary mapping string names to classes that can be built by this config.
+ """
+ raise NotImplementedError("Subclasses must define a registry property.")
+
+ def build(self) -> T:
+ """Builds the object from the configuration.
+
+ Returns:
+ An instance of the registered class.
+
+ Raises:
+ ValueError: If the name is not in the registry.
+ """
+ if self.name not in self.registry:
+ raise ValueError(
+ f"Unsupported {self.__class__.__name__}: {self.name}. "
+ f"Valid options are: {list(self.registry.keys())}"
+ )
+
+ method = self.registry[self.name]
+ return method(**self.config)
+
+
+@dataclasses.dataclass
+class ConstraintConfig(BaseConfig[tf.keras.constraints.Constraint]):
+ """Generic constraint configuration for training a model.
+
+ Attributes:
+ name: the string name of the constraint to use.
+ config: a dictionary of keyword arguments to pass to the constraint constructor.
+ registry: a dictionary mapping string names to constraint classes.
+ """
+
+ @property
+ def registry(self) -> dict[str, type]:
+ return {
+ "max_norm": tf.keras.constraints.MaxNorm,
+ "min_max_norm": tf.keras.constraints.MinMaxNorm,
+ "non_neg": tf.keras.constraints.NonNeg,
+ "unit_norm": tf.keras.constraints.UnitNorm,
+ }
+
+
+@dataclasses.dataclass
+class InitializerConfig(BaseConfig[tf.keras.initializers.Initializer]):
+ """Initializer configuration for training a model.
+
+ Attributes:
+ name: the string name of the initializer to use.
+ config: a dictionary of keyword arguments to pass to the initializer constructor.
+ """
+
+ @property
+ def registry(self) -> dict[str, type]:
+ return {
+ "glorot_normal": tf.keras.initializers.GlorotNormal,
+ "glorot_uniform": tf.keras.initializers.GlorotUniform,
+ "he_normal": tf.keras.initializers.HeNormal,
+ "he_uniform": tf.keras.initializers.HeUniform,
+ "identity": tf.keras.initializers.Identity,
+ "lecun_normal": tf.keras.initializers.LecunNormal,
+ "lecun_uniform": tf.keras.initializers.LecunUniform,
+ "ones": tf.keras.initializers.Ones,
+ "orthogonal": tf.keras.initializers.Orthogonal,
+ "random_normal": tf.keras.initializers.RandomNormal,
+ "random_uniform": tf.keras.initializers.RandomUniform,
+ "truncated_normal": tf.keras.initializers.TruncatedNormal,
+ "variance_scaling": tf.keras.initializers.VarianceScaling,
+ "zeros": tf.keras.initializers.Zeros,
+ }
+
+
+@dataclasses.dataclass
+class RegularizerConfig(BaseConfig[tf.keras.regularizers.Regularizer]):
+ """Generic regularizer configuration for training a model.
+
+ Attributes:
+ name: the string name of the regularizer to use.
+ config: a dictionary of keyword arguments to pass to the regularizer constructor.
+ """
+
+ @property
+ def registry(self) -> dict[str, type]:
+ return {
+ "l1": tf.keras.regularizers.L1,
+ "l2": tf.keras.regularizers.L2,
+ "l1_l2": tf.keras.regularizers.L1L2,
+ "orthogonal_regularizer": tf.keras.regularizers.OrthogonalRegularizer,
+ }
+
+
+@dataclasses.dataclass
+class OptimizerConfig(BaseConfig[tf.keras.optimizers.Optimizer]):
+ """Optimizer configuration for training a model.
+
+ Attributes:
+ name: the string name of the optimizer to use.
+ config: a dictionary of keyword arguments to pass to the optimizer constructor.
+ """
+
+ name: Literal["Adam"]
+
+ @property
+ def registry(self) -> dict[str, type]:
+ return {
+ "Adam": tf.keras.optimizers.Adam,
+ }
+
+
+@dataclasses.dataclass
+class ActivationConfig(BaseConfig[tf.keras.layers.Activation | tf.keras.layers.Layer]):
+ """Activation configuration for layers in a model.
+
+ Attributes:
+ name: the string name of the activation to use.
+ config: a dictionary of keyword arguments to pass to the activation constructor.
+ """
+
+ name: Literal[
+ "elliott",
+ "elu",
+ "exponential",
+ "gaussian",
+ "gcu",
+ "gelu",
+ "hard_sigmoid",
+ "hexpo",
+ "isigmoid",
+ "leaky_relu",
+ "linear",
+ "lisht",
+ "prelu",
+ "psigmoid",
+ "ptanh",
+ "ptelu",
+ "relu",
+ "resech",
+ "selu",
+ "sigmoid",
+ "smelu",
+ "snake",
+ "softmax",
+ "softplus",
+ "softsign",
+ "srs",
+ "stanh",
+ "swish",
+ "tanh",
+ "thresholded_relu",
+ ]
+
+ @property
+ def registry(self) -> dict[str, type]:
+ return {
+ "elliott": activations.Elliott,
+ "elu": tf.keras.activations.elu,
+ "exponential": tf.keras.activations.exponential,
+ "gaussian": activations.Gaussian,
+ "gcu": activations.GCU,
+ "gelu": tf.keras.activations.gelu,
+ "hard_sigmoid": tf.keras.activations.hard_sigmoid,
+ "hexpo": activations.Hexpo,
+ "isigmoid": activations.ISigmoid,
+ "linear": tf.keras.activations.linear,
+ "lisht": activations.LiSHT,
+ "prelu": tf.keras.layers.PReLU,
+ "psigmoid": activations.PSigmoid,
+ "ptanh": activations.PTanh,
+ "ptelu": activations.PTELU,
+ "relu": tf.keras.activations.relu,
+ "leaky_relu": tf.keras.layers.LeakyReLU,
+ "resech": activations.ReSech,
+ "selu": tf.keras.activations.selu,
+ "sigmoid": tf.keras.activations.sigmoid,
+ "smelu": activations.SmeLU,
+ "snake": activations.Snake,
+ "softmax": tf.keras.activations.softmax,
+ "softplus": tf.keras.activations.softplus,
+ "softsign": tf.keras.activations.softsign,
+ "srs": activations.SRS,
+ "stanh": activations.STanh,
+ "swish": tf.keras.activations.swish,
+ "tanh": tf.keras.activations.tanh,
+ "thresholded_relu": tf.keras.activations.thresholded_relu,
+ }
+
+
+@dataclasses.dataclass
+class LossConfig(BaseConfig[tf.keras.losses.Loss]):
+ """Loss configuration for training a model.
+
+ Attributes:
+ name: the string name of the loss function to use.
+ config: a dictionary of keyword arguments to pass to the loss constructor.
+ """
+
+ name: Literal[
+ "brier_skill_score",
+ "critical_success_index",
+ "fractions_skill_score",
+ "probability_of_detection",
+ ]
+
+ @property
+ def registry(self) -> dict[str, type]:
+ return {
+ "brier_skill_score": losses.brier_skill_score,
+ "critical_success_index": losses.critical_success_index,
+ "fractions_skill_score": losses.fractions_skill_score,
+ "probability_of_detection": losses.probability_of_detection,
+ }
+
+
+@dataclasses.dataclass
+class MetricConfig(BaseConfig[tf.keras.metrics.Metric]):
+ """Metric configuration for training a model.
+
+ Attributes:
+ name: the string name of the metric to use.
+ config: a dictionary of keyword arguments to pass to the metric constructor.
+ """
+
+ name: Literal[
+ "brier_skill_score",
+ "critical_success_index",
+ "fractions_skill_score",
+ "heidke_skill_score",
+ "probability_of_detection",
+ ]
+
+ @property
+ def registry(self) -> dict[str, type]:
+ return {
+ "brier_skill_score": metrics.brier_skill_score,
+ "critical_success_index": metrics.critical_success_index,
+ "fractions_skill_score": metrics.fractions_skill_score,
+ "heidke_skill_score": metrics.heidke_skill_score,
+ "probability_of_detection": metrics.probability_of_detection,
+ }
+
+
+@dataclasses.dataclass
+class ConvOutputConfig:
+ """Convolution output config for training a model.
+
+ Attributes:
+ regularizer: a RegularizerConfig to apply to convolutional layer outputs.
+ """
+
+ regularizer: RegularizerConfig | None
+
+ def build(self):
+ """Builds the convolution output configuration.
+
+ Returns:
+ A dictionary of keyword arguments to pass to convolutional layer constructors.
+ """
+
+ regularizer_object = self.regularizer.build() if self.regularizer is not None else None
+ return ConvOutput(activity_regularizer=regularizer_object)
+
+
+@dataclasses.dataclass
+class BiasVectorConfig:
+ """Constraint configuration for bias vectors in a model.
+
+ Attributes:
+ constraint: a ConstraintConfig to apply to bias vectors.
+ initializer: an InitializerConfig to use for bias vectors.
+ regularizer: a RegularizerConfig to apply to bias vectors.
+ """
+
+ constraint: ConstraintConfig | None
+ initializer: InitializerConfig
+ regularizer: RegularizerConfig | None
+
+ def build(self):
+ """Builds the bias vector configuration.
+
+ Returns:
+ A dictionary of keyword arguments to pass to layer constructors for bias vectors.
+ """
+
+ constraint_object = self.constraint.build() if self.constraint is not None else None
+ initializer_object = self.initializer.build()
+ regularizer_object = self.regularizer.build() if self.regularizer is not None else None
+
+ return BiasVector(
+ bias_constraint=constraint_object,
+ bias_initializer=initializer_object,
+ bias_regularizer=regularizer_object,
+ )
+
+
+@dataclasses.dataclass
+class KernelMatrixConfig:
+ """Constraint configuration for kernel matrices in a model.
+
+ Attributes:
+ constraint: a ConstraintConfig to apply to kernel matrices.
+ initializer: an InitializerConfig to use for kernel matrices.
+ regularizer: a RegularizerConfig to apply to kernel matrices.
+ """
+
+ constraint: ConstraintConfig | None
+ initializer: InitializerConfig
+ regularizer: RegularizerConfig | None
+
+ def build(self):
+ """Builds the kernel matrix configuration.
+
+ Returns:
+ A dictionary of keyword arguments to pass to layer constructors for kernel matrices.
+ """
+
+ constraint_object = self.constraint.build() if self.constraint is not None else None
+ initializer_object = self.initializer.build()
+ regularizer_object = self.regularizer.build() if self.regularizer is not None else None
+
+ return KernelMatrix(
+ kernel_constraint=constraint_object,
+ kernel_initializer=initializer_object,
+ kernel_regularizer=regularizer_object,
+ )
+
+
+class ConvOutput:
+ def __init__(
+ self,
+ activity_regularizer: tf.keras.regularizers.Regularizer,
+ ):
+ self.activity_regularizer = activity_regularizer
+
+
+class KernelMatrix:
+ def __init__(
+ self,
+ kernel_constraint: tf.keras.constraints.Constraint | None,
+ kernel_initializer: tf.keras.initializers.Initializer,
+ kernel_regularizer: tf.keras.regularizers.Regularizer | None,
+ ):
+ self.kernel_constraint = kernel_constraint
+ self.kernel_initializer = kernel_initializer
+ self.kernel_regularizer = kernel_regularizer
+
+
+class BiasVector:
+ def __init__(
+ self,
+ bias_constraint: tf.keras.constraints.Constraint | None,
+ bias_initializer: tf.keras.initializers.Initializer,
+ bias_regularizer: tf.keras.regularizers.Regularizer | None,
+ ):
+ self.bias_constraint = bias_constraint
+ self.bias_initializer = bias_initializer
+ self.bias_regularizer = bias_regularizer
diff --git a/src/fronts/utils/misc.py b/src/fronts/utils/misc.py
new file mode 100644
index 0000000..f53660a
--- /dev/null
+++ b/src/fronts/utils/misc.py
@@ -0,0 +1,108 @@
+"""
+Miscellaneous tools.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2024.8.10
+"""
+
+
+def initialize_gpus(devices: int | list[int], memory_growth: bool = False):
+ """
+ Initialize GPU devices.
+
+ devices: int or list of ints
+ GPU device indices.
+ memory_growth: bool
+ Use memory growth on the GPU(s).
+ """
+
+ if isinstance(devices, int):
+ devices = [
+ devices,
+ ]
+
+ # placing tensorflow import in local scope to prevent loading the library for all functions
+ import tensorflow as tf
+
+ # configure GPU devices
+ gpus = tf.config.list_physical_devices(device_type="GPU")
+ tf.config.set_visible_devices(
+ devices=[gpus[gpu] for gpu in devices], device_type="GPU"
+ )
+
+ # Allow for memory growth on the GPU. This will only use the GPU memory that is required rather than allocating all the GPU's memory.
+ if memory_growth:
+ tf.config.experimental.set_memory_growth(
+ device=[gpus[gpu] for gpu in devices][0], enable=True
+ )
+
+
+def string_arg_to_dict(arg_str: str):
+ """
+ Function that converts a string argument into a dictionary. Dictionaries cannot be passed through a command line, so
+ this function takes a special string argument and converts it to a dictionary so arguments within a function can be
+ explicitly called.
+
+ Parameters
+ ----------
+ arg_str: str
+
+ Returns
+ -------
+ arg_dict: dict
+ """
+
+ arg_str = arg_str.replace(" ", "") # Remove all spaces from the string.
+ args = arg_str.split(",")
+ arg_dict = {} # dictionary that will contain the final arguments and values
+
+ for arg in args:
+ arg_name, arg_val_str = arg.split("=")
+
+ arg_is_tuple = "(" in arg_val_str and ")" in arg_val_str
+ arg_is_list = "[" in arg_val_str and "]" in arg_val_str
+
+ # if argument value is a list or tuple
+ if arg_is_tuple or arg_is_list:
+ list_vals = (
+ arg_val_str.replace("[", "")
+ .replace("]", "")
+ .replace("(", "")
+ .replace(")", "")
+ .split("*")
+ )
+ new_list_vals = []
+ for val in list_vals:
+ if "." in val:
+ new_list_vals.append(float(val))
+ elif val == "True":
+ new_list_vals.append(True)
+ elif val == "False":
+ new_list_vals.append(False)
+ else:
+ try:
+ new_list_vals.append(int(val))
+ except ValueError:
+ new_list_vals.append(val)
+
+ if arg_is_tuple:
+ new_list_vals = tuple(new_list_vals)
+
+ arg_dict[arg_name] = new_list_vals
+
+ else:
+ if "." in arg_val_str:
+ arg_val = float(arg_val_str)
+ elif arg_val_str == "True":
+ arg_val = True
+ elif arg_val_str == "False":
+ arg_val = False
+ else:
+ try:
+ arg_val = int(arg_val_str)
+ except ValueError:
+ arg_val = arg_val_str
+
+ arg_dict[arg_name] = arg_val
+
+ return arg_dict
diff --git a/src/fronts/utils/plotting.py b/src/fronts/utils/plotting.py
new file mode 100644
index 0000000..71b603d
--- /dev/null
+++ b/src/fronts/utils/plotting.py
@@ -0,0 +1,134 @@
+"""
+Plotting tools.
+
+Author: Andrew Justin (andrewjustinwx@gmail.com)
+Script version: 2024.8.25
+"""
+
+import cartopy.feature as cfeature
+import cartopy.crs as ccrs
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+from matplotlib.colors import LinearSegmentedColormap
+import numpy as np
+
+
+def plot_background(
+ extent=None, ax=None, linewidth: float | int = 0.5, crs=ccrs.PlateCarree()
+):
+ """
+ Returns new background for the plot.
+
+ Parameters
+ ----------
+ extent: Iterable object with 4 integers
+ Iterable containing the extent/boundaries of the plot in the format of [min lon, max lon, min lat, max lat] expressed
+ in degrees.
+ ax: matplotlib.axes.Axes instance or None
+ Axis on which the background will be plotted.
+ linewidth: float or int
+ Thickness of coastlines and the borders of states and countries.
+ crs: cartopy.crs instance
+ Coordinate reference system from cartopy.
+
+ Returns
+ -------
+ ax: matplotlib.axes.Axes instance
+ New plot background.
+ """
+
+ if ax is None:
+ ax = plt.axes(crs=crs)
+ else:
+ ax.add_feature(cfeature.COASTLINE.with_scale("50m"), linewidth=linewidth)
+ ax.add_feature(cfeature.BORDERS, linewidth=linewidth)
+ ax.add_feature(cfeature.STATES, linewidth=linewidth)
+ if extent is not None:
+ ax.set_extent(extent, crs=crs)
+ return ax
+
+
+def segmented_gradient_colormap(levels, colors: list[str], ns, extend="neither"):
+ """
+ Make a segmented colormap with linear gradients between specified levels.
+
+ Parameters
+ ----------
+ levels: 1D array of ints or floats with length M
+ Levels corresponding to specified colors in the colormap.
+ colors: list of strings with length M
+ Colors corresponding to specified levels in the colormap.
+ ns: 1D array of ints with length M-1
+ Number of colors between each pair of specified levels.
+ extend: 'neither', 'min', 'max', 'both'
+ The behavior when a value falls out of range of the given levels. See matplotlib.axes.Axes.contourf for details.
+ """
+
+ len_levels, len_colors, len_ns = len(levels), len(colors), len(ns)
+
+ assert len_levels > 1, "Must specify at least two levels."
+ assert len_colors > 1, "Must specify at least two colors."
+ assert len_levels == len_colors, (
+ "The number of levels and colors must be equal. Received %d levels and %d colors."
+ % (len_levels, len_colors)
+ )
+ assert len_ns == len_levels - 1, (
+ "The length of 'ns' must be equal to the number of levels minus one. Received 'ns' length of %d and %d levels."
+ % (len_ns, len_levels)
+ )
+
+ all_levels = np.concatenate(
+ [np.linspace(levels[i], levels[i + 1], ns[i]) for i in range(len_levels - 1)]
+ )
+ all_colors = np.vstack(
+ [
+ LinearSegmentedColormap.from_list("", colors[i : i + 2])(
+ np.linspace(0, 1, ns[i])
+ )
+ for i in range(len_levels - 1)
+ ]
+ )
+
+ # matplotlib's extend function has some funky behavior, so we need to modify the colorbar to avoid errors
+ if extend == "both":
+ all_colors = np.insert(all_colors, -1, all_colors[-1], axis=0)
+ elif extend == "neither":
+ all_colors = np.delete(all_colors, 0, axis=0)
+ else:
+ pass
+
+ cmap, norm = mpl.colors.from_levels_and_colors(
+ all_levels, all_colors, extend=extend
+ )
+
+ return cmap, norm
+
+
+def truncated_colormap(
+ cmap: str, minval: float = 0.0, maxval: float = 1.0, n: int = 256
+):
+ """
+ Get an instance of a truncated matplotlib.colors.Colormap object.
+
+ Parameters
+ ----------
+ cmap: str
+ Matplotlib colormap to truncate.
+ minval: float
+ Starting point of the colormap, represented by a float of 0 <= minval < 1.
+ maxval: float
+ End point of the colormap, represented by a float of 0 < maxval <= 1.
+ n: int
+ Number of colors for the colormap.
+
+ Returns
+ -------
+ new_cmap: matplotlib.colors.Colormap instance
+ Truncated colormap.
+ """
+ cmap = plt.get_cmap(cmap)
+ new_cmap = mpl.colors.LinearSegmentedColormap.from_list(
+ "trunc({n},{a:.2f},{b:.2f})".format(n=cmap.name, a=minval, b=maxval),
+ cmap(np.linspace(minval, maxval, n)),
+ )
+ return new_cmap
diff --git a/src/fronts/utils/satellite.py b/src/fronts/utils/satellite.py
new file mode 100644
index 0000000..44ece10
--- /dev/null
+++ b/src/fronts/utils/satellite.py
@@ -0,0 +1,185 @@
+"""
+General tools for satellite data.
+
+Script version: 2024.8.25
+"""
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+from fronts.utils import plotting
+
+
+def calculate_lat_lon_from_dataset(ds):
+ """
+ Calculate lat/lon coordinates from an unmodified dataset containing GOES satellite data. This function was pulled directly
+ from https://www.star.nesdis.noaa.gov/atmospheric-composition-training/python_abi_lat_lon.php.
+
+ Parameters
+ ----------
+ ds: xarray or netCDF dataset
+ Unmodified GOES satellite dataset.
+
+ Returns
+ -------
+ abi_lat: np.ndarray
+ Latitude in degrees north.
+ abi_lon: np.ndarray
+ Longitude in degrees east.
+ """
+ # Read in GOES ABI fixed grid projection variables and constants
+ x_coordinate_1d = ds["x"][:] # E/W scanning angle in radians
+ y_coordinate_1d = ds["y"][:] # N/S elevation angle in radians
+ projection_info = ds["goes_imager_projection"]
+ lon_origin = projection_info.longitude_of_projection_origin
+ H = projection_info.perspective_point_height + projection_info.semi_major_axis
+ r_eq = projection_info.semi_major_axis
+ r_pol = projection_info.semi_minor_axis
+
+ # Create 2D coordinate matrices from 1D coordinate vectors
+ x_coordinate_2d, y_coordinate_2d = np.meshgrid(x_coordinate_1d, y_coordinate_1d)
+
+ # Equations to calculate latitude and longitude
+ lambda_0 = (lon_origin * np.pi) / 180.0
+ a_var = np.power(np.sin(x_coordinate_2d), 2.0) + (
+ np.power(np.cos(x_coordinate_2d), 2.0)
+ * (
+ np.power(np.cos(y_coordinate_2d), 2.0)
+ + (
+ ((r_eq * r_eq) / (r_pol * r_pol))
+ * np.power(np.sin(y_coordinate_2d), 2.0)
+ )
+ )
+ )
+ b_var = -2.0 * H * np.cos(x_coordinate_2d) * np.cos(y_coordinate_2d)
+ c_var = (H**2.0) - (r_eq**2.0)
+ r_s = (-1.0 * b_var - np.sqrt((b_var**2) - (4.0 * a_var * c_var))) / (2.0 * a_var)
+ s_x = r_s * np.cos(x_coordinate_2d) * np.cos(y_coordinate_2d)
+ s_y = -r_s * np.sin(x_coordinate_2d)
+ s_z = r_s * np.cos(x_coordinate_2d) * np.sin(y_coordinate_2d)
+
+ # Ignore numpy errors for sqrt of negative number; occurs for GOES-16 ABI CONUS sector data
+ np.seterr(all="ignore")
+
+ abi_lat = (180.0 / np.pi) * (
+ np.arctan(
+ ((r_eq * r_eq) / (r_pol * r_pol))
+ * (s_z / np.sqrt(((H - s_x) * (H - s_x)) + (s_y * s_y)))
+ )
+ )
+ abi_lon = (lambda_0 - np.arctan(s_y / (H - s_x))) * (180.0 / np.pi)
+
+ return abi_lat, abi_lon
+
+
+def get_satellite_colormap(band: int | str):
+ """
+ Retrieve a colormap for a GOES satellite band.
+
+ Parameters
+ ----------
+ band: int or 'sandwich'
+ Band for which the colormap will be retrieved. Integer must be between 1 and 16.
+
+ Returns
+ -------
+ cmap: instance of matplotlib.colors.LinearSegmentedColormap
+ Colormap for the requested band.
+ norm: instance of matplotlib.colors.Normalize
+ Normalization for the colormap.
+ """
+
+ if band in [1, 2, 3, 4, 5, 6, "ch1"]:
+ # greyscale colormap normalized from 0 to 1 (reflectance factor)
+ return plotting.truncated_colormap("Greys_r", minval=0.2), mpl.colors.Normalize(
+ vmin=0, vmax=1
+ )
+
+ elif band in [7, "ch2"]:
+ # multiple combined colormaps, only -65 to -25 celsius (208 to 248 K) is colored
+ n1, n2, n3 = 36, 80, 304
+
+ cmap_1 = plotting.truncated_colormap("Greys", minval=0.4, maxval=0.7)(
+ np.linspace(0, 1, n1)
+ )
+ cmap_2 = plotting.truncated_colormap("jet", minval=0.1, maxval=1.0)(
+ np.linspace(0, 1, n2)
+ )[::-1]
+ cmap_3 = plotting.truncated_colormap("Greys", minval=0.4, maxval=0.8)(
+ np.linspace(0, 1, n3)
+ )
+
+ levels = (
+ np.concatenate(
+ [
+ np.linspace(-83, -65, n1),
+ np.linspace(-65, -25, n2),
+ np.linspace(-25, 127, n3),
+ ]
+ )
+ + 273.15
+ )
+ colors = np.vstack((cmap_1, cmap_2, cmap_3))
+
+ cmap, norm = mpl.colors.from_levels_and_colors(levels, colors, extend="max")
+
+ return cmap, norm
+
+ elif band in [8, 9, 10, "ch3"]:
+ levels = np.array([-100, -75, -40, -25, -16, 0, 0, 7]) + 273.15
+ colors = [
+ "#00ffff",
+ "#438525",
+ "#ffffff",
+ "#030275",
+ "#fcfc01",
+ "#ff0000",
+ "black",
+ "black",
+ ]
+ ns = [50, 70, 30, 18, 32, 1, 1]
+ extend = "max"
+
+ elif band in [11, 12, 13, 14, 15, "ch4", "ch5"]:
+ levels = np.array([-90, -80, -70, -60, -52, -44, -36, -36, 60]) + 273.15
+ colors = [
+ "white",
+ "black",
+ "red",
+ "yellow",
+ "#26fe01",
+ "#010370",
+ "#00ffff",
+ "#c6baba",
+ "black",
+ ]
+ ns = [20, 20, 20, 16, 16, 16, 1, 190]
+ extend = "both"
+
+ elif band in [16, "ch6"]:
+ levels = np.array([-128, -90, -30, 0, 0, 30, 60, 83, 128]) + 273.15
+ colors = [
+ "black",
+ "white",
+ "blue",
+ "#704140",
+ "#89140b",
+ "yellow",
+ "red",
+ "#818180",
+ "black",
+ ]
+ ns = [76, 120, 60, 1, 60, 60, 46, 90]
+ extend = "both"
+
+ elif band == "sandwich":
+ return plt.get_cmap("jet_r"), mpl.colors.Normalize(
+ vmin=-93 + 273.15, vmax=-15 + 273.15
+ )
+
+ else:
+ raise ValueError("Unrecognized band:", band)
+
+ cmap, norm = plotting.segmented_gradient_colormap(levels, colors, ns, extend)
+
+ return cmap, norm
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..9c4d463
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,274 @@
+"""Shared fixtures and TensorFlow/wandb mocking for tests.
+
+Since tests target the config ingestion pipeline (YAML -> dacite -> dataclass -> build),
+we mock TensorFlow and wandb so tests can run without GPU or heavy dependencies.
+"""
+
+import sys
+import types
+from unittest.mock import MagicMock
+
+import pytest
+
+
+def _make_module(name: str) -> types.ModuleType:
+ """Create a mock module and register it in sys.modules."""
+ mod = types.ModuleType(name)
+ sys.modules[name] = mod
+ return mod
+
+
+def _setup_tensorflow_mocks():
+ """Install mock TF modules before any fronts imports."""
+ # Top-level tensorflow
+ tf = _make_module("tensorflow")
+ tf.Tensor = MagicMock
+ tf.cast = MagicMock()
+ tf.abs = MagicMock()
+ tf.function = lambda f: f # passthrough decorator
+ tf.float32 = "float32"
+ tf.float16 = "float16"
+ tf.is_tensor = MagicMock(return_value=False)
+ tf.data = types.ModuleType("tensorflow.data")
+ tf.data.Dataset = MagicMock
+
+ # tensorflow.keras
+ keras = _make_module("tensorflow.keras")
+ tf.keras = keras
+
+ # Sub-modules that get imported
+ for sub in [
+ "tensorflow.keras.layers",
+ "tensorflow.keras.models",
+ "tensorflow.keras.callbacks",
+ "tensorflow.keras.regularizers",
+ "tensorflow.keras.optimizers",
+ "tensorflow.keras.constraints",
+ "tensorflow.keras.initializers",
+ "tensorflow.keras.activations",
+ "tensorflow.keras.losses",
+ "tensorflow.keras.metrics",
+ "tensorflow.keras.utils",
+ ]:
+ mod = _make_module(sub)
+ # Attach as attribute on parent
+ parts = sub.split(".")
+ parent = sys.modules[".".join(parts[:-1])]
+ setattr(parent, parts[-1], mod)
+
+ # Common classes/functions that the code references on these modules
+ layers = sys.modules["tensorflow.keras.layers"]
+ class _MockLayer:
+ """Minimal mock of tf.keras.layers.Layer for custom activation subclasses."""
+ def __init__(self, *args, **kwargs):
+ pass
+ def build(self, input_shape):
+ pass
+ def call(self, inputs):
+ return inputs
+ layers.Layer = _MockLayer
+ layers.Concatenate = MagicMock
+ layers.Input = MagicMock(return_value=MagicMock())
+ layers.PReLU = MagicMock
+ layers.LeakyReLU = MagicMock
+ layers.Activation = MagicMock
+ layers.Conv2D = MagicMock
+ layers.Conv3D = MagicMock
+ layers.BatchNormalization = MagicMock
+ layers.MaxPooling2D = MagicMock
+ layers.MaxPooling3D = MagicMock
+ layers.UpSampling2D = MagicMock
+ layers.UpSampling3D = MagicMock
+ layers.AveragePooling1D = MagicMock
+ layers.AveragePooling2D = MagicMock
+ layers.AveragePooling3D = MagicMock
+
+ models = sys.modules["tensorflow.keras.models"]
+ models.Model = MagicMock
+
+ callbacks = sys.modules["tensorflow.keras.callbacks"]
+ callbacks.Callback = MagicMock
+ callbacks.ModelCheckpoint = MagicMock
+ callbacks.CSVLogger = MagicMock
+ callbacks.EarlyStopping = MagicMock
+
+ regularizers = sys.modules["tensorflow.keras.regularizers"]
+ regularizers.Regularizer = MagicMock
+ regularizers.L1 = MagicMock
+ regularizers.L2 = MagicMock
+ regularizers.L1L2 = MagicMock
+ regularizers.OrthogonalRegularizer = MagicMock
+
+ optimizers = sys.modules["tensorflow.keras.optimizers"]
+ optimizers.Optimizer = MagicMock
+ optimizers.Adam = MagicMock
+
+ constraints = sys.modules["tensorflow.keras.constraints"]
+ constraints.Constraint = MagicMock
+ constraints.MaxNorm = MagicMock
+ constraints.MinMaxNorm = MagicMock
+ constraints.NonNeg = MagicMock
+ constraints.UnitNorm = MagicMock
+
+ initializers = sys.modules["tensorflow.keras.initializers"]
+ initializers.Initializer = MagicMock
+ initializers.GlorotNormal = MagicMock
+ initializers.GlorotUniform = MagicMock
+ initializers.HeNormal = MagicMock
+ initializers.HeUniform = MagicMock
+ initializers.Identity = MagicMock
+ initializers.LecunNormal = MagicMock
+ initializers.LecunUniform = MagicMock
+ initializers.Ones = MagicMock
+ initializers.Orthogonal = MagicMock
+ initializers.RandomNormal = MagicMock
+ initializers.RandomUniform = MagicMock
+ initializers.TruncatedNormal = MagicMock
+ initializers.VarianceScaling = MagicMock
+ initializers.Zeros = MagicMock
+
+ activations = sys.modules["tensorflow.keras.activations"]
+ for fn_name in [
+ "elu", "exponential", "gelu", "hard_sigmoid", "linear",
+ "relu", "selu", "sigmoid", "softmax", "softplus", "softsign",
+ "swish", "tanh", "thresholded_relu",
+ ]:
+ setattr(activations, fn_name, MagicMock(name=f"tf.keras.activations.{fn_name}"))
+
+ losses_mod = sys.modules["tensorflow.keras.losses"]
+ losses_mod.Loss = MagicMock
+
+ metrics_mod = sys.modules["tensorflow.keras.metrics"]
+ metrics_mod.Metric = MagicMock
+
+ utils_mod = sys.modules["tensorflow.keras.utils"]
+ utils_mod.set_random_seed = MagicMock()
+
+ # keras as a top-level attribute
+ keras.Activation = MagicMock
+ keras.Layer = layers.Layer
+
+
+def _setup_geospatial_mocks():
+ """Install mock modules for heavy geospatial/data dependencies."""
+ # shapely — imported by data_utils.py
+ shapely = _make_module("shapely")
+ shapely_geom = _make_module("shapely.geometry")
+ shapely_geom.LineString = MagicMock
+ shapely_geom.MultiLineString = MagicMock
+ shapely_geom.Point = MagicMock
+ shapely.geometry = shapely_geom
+ _make_module("shapely.ops")
+
+ # regionmask — may be imported transitively
+ _make_module("regionmask")
+
+ # gcsfs — for ERA5 zarr store access
+ _make_module("gcsfs")
+
+
+def _setup_xbatcher_mocks():
+ """Install mock xbatcher modules before any fronts imports."""
+ xb = _make_module("xbatcher")
+ xb.BatchGenerator = MagicMock(return_value=MagicMock())
+
+ loaders = _make_module("xbatcher.loaders")
+ xb.loaders = loaders
+
+ loaders_keras = _make_module("xbatcher.loaders.keras")
+ loaders_keras.CustomTFDataset = MagicMock(return_value=iter([]))
+ loaders.keras = loaders_keras
+
+
+def _setup_wandb_mocks():
+ """Install mock wandb modules."""
+ wandb = _make_module("wandb")
+ wandb.login = MagicMock()
+ wandb.init = MagicMock()
+
+ integration = _make_module("wandb.integration")
+ wandb.integration = integration
+
+ wandb_keras = _make_module("wandb.integration.keras")
+ wandb_keras.WandbMetricsLogger = MagicMock
+ wandb_keras.WandbModelCheckpoint = MagicMock
+ integration.keras = wandb_keras
+
+
+# Install mocks before any test collection triggers fronts imports
+_setup_tensorflow_mocks()
+_setup_geospatial_mocks()
+_setup_xbatcher_mocks()
+_setup_wandb_mocks()
+
+
+@pytest.fixture
+def sample_config_dict():
+ """A minimal valid config dict matching TrainConfig structure."""
+ return {
+ "epochs": 10,
+ "training_steps_per_epoch": 5,
+ "validation_steps_per_epoch": None,
+ "validation_frequency": 1,
+ "verbose": 1,
+ "repeat": True,
+ "seed": 42,
+ "model": {
+ "name": "unet_3plus",
+ "batch_normalization": True,
+ "num_filters": [16, 32, 64, 128],
+ "kernel_size": [5, 5, 5],
+ "pool_size": [2, 2, 1],
+ "upsample_size": [2, 2, 1],
+ "depth": 4,
+ "modules_per_node": 2,
+ "padding": "same",
+ "bias": True,
+ "loss": {
+ "name": "fractions_skill_score",
+ "config": {"mask_size": [3, 3]},
+ },
+ "metric": {
+ "name": "critical_success_index",
+ "config": {"class_weights": [0, 1, 1, 1, 1, 1]},
+ },
+ "optimizer": {
+ "name": "Adam",
+ "config": {"beta_1": 0.9, "beta_2": 0.999},
+ },
+ "convolution_activity_regularizer": {"regularizer": None},
+ "bias_vector": {
+ "constraint": None,
+ "initializer": {"name": "zeros", "config": {}},
+ "regularizer": None,
+ },
+ "kernel_matrix": {
+ "constraint": None,
+ "initializer": {"name": "glorot_uniform", "config": {}},
+ "regularizer": None,
+ },
+ "activation": {"name": "gelu", "config": {}},
+ },
+ "wandb": {
+ "project_name": "test_project",
+ "model_run_name": "test_run",
+ },
+ "callbacks": {
+ "monitor": "val_loss",
+ "verbose": 1,
+ "save_best_only": True,
+ "save_weights_only": False,
+ "save_freq": "epoch",
+ },
+ }
+
+
+
+@pytest.fixture
+def sample_yaml_file(tmp_path, sample_config_dict):
+ """Write sample_config_dict as a YAML file and return its path."""
+ import yaml
+
+ path = tmp_path / "test_config.yaml"
+ path.write_text(yaml.dump(sample_config_dict, default_flow_style=False))
+ return str(path)
diff --git a/tests/fixtures/.gitignore b/tests/fixtures/.gitignore
new file mode 100644
index 0000000..d4dee8d
--- /dev/null
+++ b/tests/fixtures/.gitignore
@@ -0,0 +1,2 @@
+# Generated by scripts/make_dryrun_data.py — do not commit
+dryrun_tf_dataset/
diff --git a/tests/test_config_ingestion.py b/tests/test_config_ingestion.py
new file mode 100644
index 0000000..81a8bd3
--- /dev/null
+++ b/tests/test_config_ingestion.py
@@ -0,0 +1,178 @@
+"""End-to-end test: YAML file -> dacite -> TrainConfig -> Trainer.
+
+Uses the actual 1702.yaml config to verify the full pipeline works.
+"""
+
+import os
+from unittest.mock import patch
+
+import dacite
+import pytest
+import yaml
+
+from fronts.train import TrainConfig, Trainer, open_config_yaml_as_dataclass
+from fronts.model import ModelConfig
+from fronts.utils.keras_builders import (
+ ActivationConfig,
+ BiasVectorConfig,
+ ConvOutputConfig,
+ InitializerConfig,
+ KernelMatrixConfig,
+ LossConfig,
+ MetricConfig,
+ OptimizerConfig,
+)
+
+
+class TestFullPipelineFromFixture:
+ """Test the full YAML -> dacite -> TrainConfig pipeline using the test fixture."""
+
+ def test_full_pipeline(self, sample_yaml_file):
+ """Complete pipeline: YAML -> TrainConfig -> Trainer."""
+ config = open_config_yaml_as_dataclass(
+ path=sample_yaml_file, config_class=TrainConfig
+ )
+
+ # Verify top-level fields
+ assert config.epochs == 10
+ assert config.training_steps_per_epoch == 5
+ assert config.validation_frequency == 1
+ assert config.verbose == 1
+ assert config.repeat is True
+ assert config.seed == 42
+
+ # Verify nested ModelConfig
+ assert isinstance(config.model, ModelConfig)
+ assert config.model.name == "unet_3plus"
+ assert config.model.num_filters == [16, 32, 64, 128]
+ assert config.model.depth == 4
+ assert config.model.padding == "same"
+ assert config.model.bias is True
+
+ # Verify nested configs within model
+ assert isinstance(config.model.loss, LossConfig)
+ assert config.model.loss.name == "fractions_skill_score"
+ assert config.model.loss.config == {"mask_size": [3, 3]}
+
+ assert isinstance(config.model.metric, MetricConfig)
+ assert config.model.metric.name == "critical_success_index"
+
+ assert isinstance(config.model.optimizer, OptimizerConfig)
+ assert config.model.optimizer.name == "Adam"
+
+ assert isinstance(config.model.activation, ActivationConfig)
+ assert config.model.activation.name == "gelu"
+
+ assert isinstance(config.model.convolution_activity_regularizer, ConvOutputConfig)
+ assert config.model.convolution_activity_regularizer.regularizer is None
+
+ assert isinstance(config.model.bias_vector, BiasVectorConfig)
+ assert config.model.bias_vector.constraint is None
+ assert config.model.bias_vector.regularizer is None
+ assert isinstance(config.model.bias_vector.initializer, InitializerConfig)
+
+ assert isinstance(config.model.kernel_matrix, KernelMatrixConfig)
+ assert config.model.kernel_matrix.constraint is None
+ assert config.model.kernel_matrix.regularizer is None
+ assert isinstance(config.model.kernel_matrix.initializer, InitializerConfig)
+
+ # Build Trainer from config
+ trainer = config.build()
+ assert isinstance(trainer, Trainer)
+ assert trainer.epochs == 10
+ assert trainer.wandb is config.wandb
+
+ def test_dacite_rejects_invalid_model_name(self, tmp_path):
+ """dacite should propagate through; an invalid Literal should error."""
+ bad_config = {
+ "epochs": 1,
+ "training_steps_per_epoch": 1,
+ "validation_steps_per_epoch": None,
+ "validation_frequency": 1,
+ "verbose": 1,
+ "repeat": True,
+ "seed": 42,
+ "model": {
+ "name": "nonexistent_model",
+ "batch_normalization": True,
+ "num_filters": [16],
+ "kernel_size": [3],
+ "pool_size": [2],
+ "upsample_size": [2],
+ "depth": 1,
+ "modules_per_node": 1,
+ "padding": "same",
+ "bias": True,
+ "loss": {"name": "fractions_skill_score", "config": {}},
+ "metric": {"name": "critical_success_index", "config": {}},
+ "optimizer": {"name": "Adam", "config": {}},
+ "convolution_activity_regularizer": {"regularizer": None},
+ "bias_vector": {
+ "constraint": None,
+ "initializer": {"name": "zeros", "config": {}},
+ "regularizer": None,
+ },
+ "kernel_matrix": {
+ "constraint": None,
+ "initializer": {"name": "glorot_uniform", "config": {}},
+ "regularizer": None,
+ },
+ "activation": {"name": "gelu", "config": {}},
+ },
+ "wandb": {"project_name": "test", "model_run_name": "test"},
+ "callbacks": {
+ "monitor": "val_loss",
+ "verbose": 0,
+ "save_best_only": False,
+ "save_weights_only": False,
+ "save_freq": "epoch",
+ },
+ }
+ yaml_path = tmp_path / "bad_config.yaml"
+ yaml_path.write_text(yaml.dump(bad_config))
+
+ # dacite with strict_unions_match or Literal checking should handle this.
+ # Even without strict checking, the pipeline should at least not crash
+ # during loading — the error surfaces when building.
+ result = open_config_yaml_as_dataclass(
+ path=str(yaml_path), config_class=TrainConfig
+ )
+ # The config loads (dacite doesn't enforce Literal by default), but
+ # the model name won't be in the UNetRegistry when build is called
+ assert result.model.name == "nonexistent_model"
+
+
+class TestFromActualYaml:
+ """Test loading the actual 1702.yaml config file."""
+
+ @pytest.fixture
+ def actual_yaml_path(self):
+ path = os.path.join(
+ os.path.dirname(__file__), "..", "configs", "1702.yaml"
+ )
+ if not os.path.exists(path):
+ pytest.skip("1702.yaml not found")
+ return path
+
+ def test_actual_yaml_loads(self, actual_yaml_path):
+ """The real 1702.yaml config loads successfully into TrainConfig."""
+ config = open_config_yaml_as_dataclass(
+ path=actual_yaml_path, config_class=TrainConfig
+ )
+ assert config is not None
+ assert isinstance(config, TrainConfig)
+ assert config.epochs == 5000
+ assert config.model.name == "unet_3plus"
+
+ def test_actual_yaml_builds_trainer(self, actual_yaml_path):
+ """The real 1702.yaml config builds a Trainer."""
+ from fronts.data.config import DataConfig, ModelData
+
+ config = open_config_yaml_as_dataclass(
+ path=actual_yaml_path, config_class=TrainConfig
+ )
+ dummy_data = ModelData(train_data=None, validation_data=None)
+ with patch.object(DataConfig, "build", return_value=dummy_data):
+ trainer = config.build()
+ assert isinstance(trainer, Trainer)
+ assert trainer.epochs == 5000
diff --git a/tests/test_data_config.py b/tests/test_data_config.py
new file mode 100644
index 0000000..6c954bd
--- /dev/null
+++ b/tests/test_data_config.py
@@ -0,0 +1,879 @@
+"""Tests for the DataConfig pipeline: YAML -> dacite -> dataclasses.
+
+Tests verify that all DataConfig classes (ERA5PredictorConfig, FrontsDataConfig,
+BatchGeneratorConfig, DataConfig) parse correctly from dicts/YAML via dacite and
+have the expected field types. .build() methods are tested with mocked I/O.
+
+Does NOT require TensorFlow, GPU, or real data files — follows the same pattern
+as conftest.py (mocks installed at module load time).
+"""
+
+import datetime
+import os
+from unittest.mock import MagicMock, patch
+
+import dacite
+import numpy as np
+import pandas as pd
+import pytest
+import xarray as xr
+import yaml
+
+from fronts.data.batch import BatchGeneratorConfig
+from fronts.data.config import (
+ DataConfig,
+ ERA5PredictorConfig,
+ FrontsDataConfig,
+ PredictConfig,
+ TimeSelection,
+ ModelData,
+ TFDatasetConfig,
+ SURFACE_VARIABLE_MAP,
+ SURFACE_ONLY_VARIABLES,
+)
+from fronts.train import TrainConfig, open_config_yaml_as_dataclass
+
+DACITE_CONFIG = dacite.Config(cast=[tuple], check_types=False)
+
+
+# ---------------------------------------------------------------------------
+# Shared xarray fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def sample_era5_ds():
+ """Minimal ERA5-like Dataset with unified level coordinate."""
+ times = pd.date_range("2010-01-01", periods=4, freq="6h")
+ lats = np.array([30.0, 35.0, 40.0])
+ lons = np.array([-100.0, -95.0, -90.0])
+ levels = ["surface", 1000, 850]
+
+ return xr.Dataset(
+ {
+ "temperature": (
+ ["time", "level", "latitude", "longitude"],
+ np.random.rand(4, 3, 3, 3).astype("float32"),
+ ),
+ },
+ coords={
+ "time": times,
+ "level": levels,
+ "latitude": lats,
+ "longitude": lons,
+ },
+ )
+
+
+@pytest.fixture
+def sample_fronts_ds():
+ """Minimal fronts Dataset with identifier variable."""
+ times = pd.date_range("2010-01-01", periods=4, freq="6h")
+ lats = np.array([30.0, 35.0, 40.0])
+ lons = np.array([-100.0, -95.0, -90.0])
+
+ return xr.Dataset(
+ {
+ "identifier": (
+ ["time", "latitude", "longitude"],
+ np.zeros((4, 3, 3), dtype="int32"),
+ ),
+ },
+ coords={"time": times, "latitude": lats, "longitude": lons},
+ )
+
+
+# ---------------------------------------------------------------------------
+# ERA5PredictorConfig
+# ---------------------------------------------------------------------------
+
+
+class TestERA5PredictorConfig:
+ def _minimal_dict(self):
+ return {
+ "domain_extent": [-140.0, -60.0, 20.0, 60.0],
+ "variables": ["temperature", "mean_sea_level_pressure"],
+ "levels": ["surface", 1000, 850],
+ "years": [2010, 2011],
+ "store": "gs://fake-store",
+ "chunks": {"time": 48},
+ "consolidated": True,
+ }
+
+ def test_dacite_parses(self):
+ cfg = dacite.from_dict(ERA5PredictorConfig, self._minimal_dict(), DACITE_CONFIG)
+ assert cfg.levels == ["surface", 1000, 850]
+ assert cfg.years == [2010, 2011]
+ assert cfg.variables == ["temperature", "mean_sea_level_pressure"]
+ assert cfg.consolidated is True
+
+ def test_domain_extent_length(self):
+ cfg = dacite.from_dict(ERA5PredictorConfig, self._minimal_dict(), DACITE_CONFIG)
+ assert len(cfg.domain_extent) == 4
+
+ def test_levels_pressure_only(self):
+ """levels may contain only integer hPa values (no surface)."""
+ d = self._minimal_dict()
+ d["levels"] = [1000, 900, 750]
+ cfg = dacite.from_dict(ERA5PredictorConfig, d, DACITE_CONFIG)
+ assert cfg.levels == [1000, 900, 750]
+ assert "surface" not in cfg.levels
+
+ def test_build_calls_open_zarr(self, sample_era5_ds):
+ """ERA5PredictorConfig.build() opens the zarr store and subsets spatially/temporally."""
+ times = pd.date_range("2010-01-01", periods=8, freq="6h")
+ lats = np.linspace(20.0, 60.0, 5)
+ lons = np.linspace(-140.0, -60.0, 5)
+ levels = [1000, 850]
+
+ raw_ds = xr.Dataset(
+ {
+ "temperature": (
+ ["time", "level", "latitude", "longitude"],
+ np.random.rand(8, 2, 5, 5).astype("float32"),
+ ),
+ "2m_temperature": (
+ ["time", "latitude", "longitude"],
+ np.random.rand(8, 5, 5).astype("float32"),
+ ),
+ },
+ coords={
+ "time": times,
+ "level": levels,
+ "latitude": lats,
+ "longitude": lons,
+ },
+ )
+
+ cfg = ERA5PredictorConfig(
+ domain_extent=[-140.0, -60.0, 20.0, 60.0],
+ variables=["temperature"],
+ levels=["surface", 1000, 850],
+ years=[2010],
+ store="gs://fake-store",
+ chunks={"time": 48},
+ consolidated=True,
+ )
+
+ with patch("xarray.open_zarr", return_value=raw_ds):
+ result = cfg.build()
+
+ assert "temperature" in result
+ # temperature should have level coord including "surface"
+ assert "surface" in result["temperature"].coords["level"].values
+
+ def test_build_surface_only_variable_added(self):
+ """Surface-only variables appear in result with level=["surface"]."""
+ times = pd.date_range("2010-01-01", periods=4, freq="6h")
+ lats = np.linspace(20.0, 60.0, 3)
+ lons = np.linspace(-140.0, -60.0, 3)
+
+ raw_ds = xr.Dataset(
+ {
+ "temperature": (
+ ["time", "level", "latitude", "longitude"],
+ np.random.rand(4, 1, 3, 3).astype("float32"),
+ ),
+ "mean_sea_level_pressure": (
+ ["time", "latitude", "longitude"],
+ np.random.rand(4, 3, 3).astype("float32"),
+ ),
+ },
+ coords={
+ "time": times,
+ "level": [1000],
+ "latitude": lats,
+ "longitude": lons,
+ },
+ )
+
+ cfg = ERA5PredictorConfig(
+ domain_extent=[-140.0, -60.0, 20.0, 60.0],
+ variables=["temperature", "mean_sea_level_pressure"],
+ levels=["surface", 1000],
+ years=[2010],
+ store="gs://fake-store",
+ chunks={"time": 48},
+ consolidated=True,
+ )
+
+ with patch("xarray.open_zarr", return_value=raw_ds):
+ result = cfg.build()
+
+ assert "mean_sea_level_pressure" in result
+ assert "surface" in result["mean_sea_level_pressure"].coords["level"].values
+
+ def test_build_pressure_only_levels(self):
+ """When levels has no 'surface', only pressure-level data is returned."""
+ times = pd.date_range("2010-01-01", periods=4, freq="6h")
+ lats = np.linspace(20.0, 60.0, 3)
+ lons = np.linspace(-140.0, -60.0, 3)
+
+ raw_ds = xr.Dataset(
+ {
+ "temperature": (
+ ["time", "level", "latitude", "longitude"],
+ np.random.rand(4, 2, 3, 3).astype("float32"),
+ ),
+ },
+ coords={
+ "time": times,
+ "level": [1000, 850],
+ "latitude": lats,
+ "longitude": lons,
+ },
+ )
+
+ cfg = ERA5PredictorConfig(
+ domain_extent=[-140.0, -60.0, 20.0, 60.0],
+ variables=["temperature"],
+ levels=[1000, 850],
+ years=[2010],
+ store="gs://fake-store",
+ chunks={"time": 48},
+ consolidated=True,
+ )
+
+ with patch("xarray.open_zarr", return_value=raw_ds):
+ result = cfg.build()
+
+ assert "temperature" in result
+ assert "surface" not in result["temperature"].coords["level"].values
+
+
+# ---------------------------------------------------------------------------
+# FrontsDataConfig
+# ---------------------------------------------------------------------------
+
+
+class TestFrontsDataConfig:
+ def test_dacite_parses_string_front_types(self):
+ d = {"directory": "/tmp/fronts", "years": [2010], "front_types": "MERGED-ALL"}
+ cfg = dacite.from_dict(FrontsDataConfig, d, DACITE_CONFIG)
+ assert cfg.front_types == "MERGED-ALL"
+ assert cfg.years == [2010]
+
+ def test_dacite_parses_list_front_types(self):
+ d = {"directory": "/tmp/fronts", "years": [2010], "front_types": ["CF", "WF"]}
+ cfg = dacite.from_dict(FrontsDataConfig, d, DACITE_CONFIG)
+ assert cfg.front_types == ["CF", "WF"]
+
+ def test_dacite_parses_null_front_types(self):
+ d = {"directory": "/tmp/fronts", "years": [2010], "front_types": None}
+ cfg = dacite.from_dict(FrontsDataConfig, d, DACITE_CONFIG)
+ assert cfg.front_types is None
+
+ def test_build_with_mocked_mfdataset(self, sample_fronts_ds, tmp_path):
+ """FrontsDataConfig.build() calls open_mfdataset and returns a Dataset."""
+ cfg = FrontsDataConfig(
+ directory=str(tmp_path),
+ years=[2010],
+ front_types=None,
+ )
+ with patch("xarray.open_mfdataset", return_value=sample_fronts_ds):
+ ds = cfg.build()
+ assert "identifier" in ds
+
+ def test_build_calls_reformat_fronts_when_front_types_set(self, sample_fronts_ds, tmp_path):
+ """FrontsDataConfig.build() calls reformat_fronts when front_types is given."""
+ cfg = FrontsDataConfig(
+ directory=str(tmp_path),
+ years=[2010],
+ front_types="MERGED-ALL",
+ )
+ with patch("xarray.open_mfdataset", return_value=sample_fronts_ds):
+ with patch("fronts.utils.data_utils.reformat_fronts", return_value=sample_fronts_ds) as mock_rf:
+ cfg.build()
+ mock_rf.assert_called_once()
+
+ def test_build_skips_reformat_fronts_when_none(self, sample_fronts_ds, tmp_path):
+ """FrontsDataConfig.build() does not call reformat_fronts when front_types is None."""
+ cfg = FrontsDataConfig(
+ directory=str(tmp_path),
+ years=[2010],
+ front_types=None,
+ )
+ with patch("xarray.open_mfdataset", return_value=sample_fronts_ds):
+ with patch("fronts.utils.data_utils.reformat_fronts") as mock_rf:
+ cfg.build()
+ mock_rf.assert_not_called()
+
+
+# ---------------------------------------------------------------------------
+# BatchGeneratorConfig
+# ---------------------------------------------------------------------------
+
+
+class TestBatchGeneratorConfig:
+ def test_dacite_parses(self):
+ d = {
+ "input_sizes": {"time": 1, "latitude": 128, "longitude": 128},
+ "target_sizes": {"time": 1, "latitude": 128, "longitude": 128},
+ "prefetch_number": 3,
+ "preload_batch": False,
+ }
+ cfg = dacite.from_dict(BatchGeneratorConfig, d, DACITE_CONFIG)
+ assert cfg.input_sizes == {"time": 1, "latitude": 128, "longitude": 128}
+ assert cfg.prefetch_number == 3
+
+ def test_defaults(self):
+ cfg = BatchGeneratorConfig()
+ assert cfg.input_sizes is None
+ assert cfg.target_sizes is None
+ assert cfg.prefetch_number == 3
+ assert cfg.preload_batch is False
+
+
+# ---------------------------------------------------------------------------
+# DataConfig
+# ---------------------------------------------------------------------------
+
+
+def _minimal_data_config_dict():
+ return {
+ "train_years": [2010],
+ "val_years": [2011],
+ "test_years": [],
+ "shuffle": True,
+ "normalization_method": "standard",
+ "era5": {
+ "domain_extent": [-140.0, -60.0, 20.0, 60.0],
+ "variables": ["temperature", "mean_sea_level_pressure"],
+ "levels": ["surface", 1000, 850],
+ "years": [],
+ "store": "gs://fake-store",
+ "chunks": {"time": 48},
+ "consolidated": True,
+ },
+ "fronts": {
+ "directory": "/tmp/fronts",
+ "years": [],
+ "front_types": "MERGED-ALL",
+ },
+ "batch": {
+ "input_sizes": {"time": 1, "latitude": 128, "longitude": 128},
+ "target_sizes": {"time": 1, "latitude": 128, "longitude": 128},
+ "prefetch_number": 3,
+ "preload_batch": False,
+ },
+ }
+
+
+class TestDataConfig:
+ def test_dacite_parses_full_config(self):
+ cfg = dacite.from_dict(DataConfig, _minimal_data_config_dict(), DACITE_CONFIG)
+ assert cfg.train_years == [2010]
+ assert cfg.val_years == [2011]
+ assert cfg.test_years == []
+ assert cfg.shuffle is True
+ assert cfg.normalization_method == "standard"
+ assert isinstance(cfg.era5, ERA5PredictorConfig)
+ assert isinstance(cfg.fronts, FrontsDataConfig)
+ assert isinstance(cfg.batch, BatchGeneratorConfig)
+
+ def test_nested_era5_fields(self):
+ cfg = dacite.from_dict(DataConfig, _minimal_data_config_dict(), DACITE_CONFIG)
+ assert cfg.era5.levels == ["surface", 1000, 850]
+ assert cfg.era5.variables == ["temperature", "mean_sea_level_pressure"]
+
+ def test_nested_fronts_fields(self):
+ cfg = dacite.from_dict(DataConfig, _minimal_data_config_dict(), DACITE_CONFIG)
+ assert cfg.fronts.front_types == "MERGED-ALL"
+ assert cfg.fronts.directory == "/tmp/fronts"
+
+ def test_empty_test_years_produces_none_test_data(self, sample_era5_ds, sample_fronts_ds):
+ """DataConfig.build() returns test_data=None when test_years is empty."""
+ cfg = dacite.from_dict(DataConfig, _minimal_data_config_dict(), DACITE_CONFIG)
+
+ mock_tf_ds = MagicMock()
+ mock_tf_ds.shuffle = MagicMock(return_value=mock_tf_ds)
+
+ with patch.object(ERA5PredictorConfig, "build", return_value=sample_era5_ds):
+ with patch.object(FrontsDataConfig, "build", return_value=sample_fronts_ds):
+ with patch("fronts.utils.data_utils.normalize_dataset", return_value=sample_era5_ds):
+ with patch("fronts.data.config.create_dataloader", return_value=mock_tf_ds):
+ result = cfg.build()
+
+ assert isinstance(result, ModelData)
+ assert result.test_data is None
+
+ def test_shuffle_called_on_train_ds(self, sample_era5_ds, sample_fronts_ds):
+ """DataConfig.build() calls .shuffle() on the training dataset when shuffle=True."""
+ cfg = dacite.from_dict(DataConfig, _minimal_data_config_dict(), DACITE_CONFIG)
+
+ mock_tf_ds = MagicMock()
+ shuffled_ds = MagicMock()
+ mock_tf_ds.shuffle = MagicMock(return_value=shuffled_ds)
+
+ with patch.object(ERA5PredictorConfig, "build", return_value=sample_era5_ds):
+ with patch.object(FrontsDataConfig, "build", return_value=sample_fronts_ds):
+ with patch("fronts.utils.data_utils.normalize_dataset", return_value=sample_era5_ds):
+ with patch("fronts.data.config.create_dataloader", return_value=mock_tf_ds):
+ result = cfg.build()
+
+ mock_tf_ds.shuffle.assert_called_once()
+ assert result.train_data is shuffled_ds
+
+ def test_no_shuffle_when_disabled(self, sample_era5_ds, sample_fronts_ds):
+ """DataConfig.build() does not shuffle when shuffle=False."""
+ d = _minimal_data_config_dict()
+ d["shuffle"] = False
+ cfg = dacite.from_dict(DataConfig, d, DACITE_CONFIG)
+
+ mock_tf_ds = MagicMock()
+ mock_tf_ds.shuffle = MagicMock()
+
+ with patch.object(ERA5PredictorConfig, "build", return_value=sample_era5_ds):
+ with patch.object(FrontsDataConfig, "build", return_value=sample_fronts_ds):
+ with patch("fronts.utils.data_utils.normalize_dataset", return_value=sample_era5_ds):
+ with patch("fronts.data.config.create_dataloader", return_value=mock_tf_ds):
+ cfg.build()
+
+ mock_tf_ds.shuffle.assert_not_called()
+
+ def test_years_injected_via_replace(self, sample_era5_ds, sample_fronts_ds):
+ """DataConfig.build() injects split-specific years into ERA5/Fronts configs."""
+ cfg = dacite.from_dict(DataConfig, _minimal_data_config_dict(), DACITE_CONFIG)
+ # era5.years in the dict is [] — DataConfig.build() injects train_years=[2010], val_years=[2011]
+
+ all_era5_years = []
+ all_fronts_years = []
+
+ def mock_era5_build(self):
+ all_era5_years.append(self.years)
+ return sample_era5_ds
+
+ def mock_fronts_build(self):
+ all_fronts_years.append(self.years)
+ return sample_fronts_ds
+
+ mock_tf_ds = MagicMock()
+ mock_tf_ds.shuffle = MagicMock(return_value=mock_tf_ds)
+
+ with patch.object(ERA5PredictorConfig, "build", mock_era5_build):
+ with patch.object(FrontsDataConfig, "build", mock_fronts_build):
+ with patch("fronts.utils.data_utils.normalize_dataset", return_value=sample_era5_ds):
+ with patch("fronts.data.config.create_dataloader", return_value=mock_tf_ds):
+ cfg.build()
+
+ # train_years=[2010] and val_years=[2011] both get injected; test_years=[] is skipped
+ assert [2010] in all_era5_years
+ assert [2011] in all_era5_years
+ assert [2010] in all_fronts_years
+ assert [2011] in all_fronts_years
+
+
+# ---------------------------------------------------------------------------
+# TFDatasetConfig
+# ---------------------------------------------------------------------------
+
+
+class TestTFDatasetConfig:
+ def _minimal_dict(self):
+ return {
+ "directory": "/tmp/tf_datasets",
+ "train_years": [2010],
+ "val_years": [2011],
+ "test_years": [],
+ "shuffle": True,
+ "shuffle_buffer": 1000,
+ "prefetch": 3,
+ }
+
+ def test_dacite_parses(self):
+ cfg = dacite.from_dict(TFDatasetConfig, self._minimal_dict(), DACITE_CONFIG)
+ assert cfg.directory == "/tmp/tf_datasets"
+ assert cfg.train_years == [2010]
+ assert cfg.shuffle is True
+
+ def test_build_loads_matching_subdirs(self, tmp_path):
+ """TFDatasetConfig.build() loads subdirs matching the requested years."""
+ (tmp_path / "2010-1_tf").mkdir()
+ (tmp_path / "2010-2_tf").mkdir()
+ (tmp_path / "2011-1_tf").mkdir()
+
+ mock_ds = MagicMock()
+ mock_ds.shuffle = MagicMock(return_value=mock_ds)
+
+ cfg = TFDatasetConfig(
+ directory=str(tmp_path),
+ train_years=[2010],
+ val_years=[2011],
+ test_years=[],
+ )
+
+ # Patch _load_years to avoid calling tf.data.Dataset.load (mocked at class level).
+ # Return None for empty year lists (mirrors real behaviour), mock_ds otherwise.
+ def fake_load_years(self, years):
+ return None if not years else mock_ds
+
+ with patch.object(TFDatasetConfig, "_load_years", fake_load_years):
+ result = cfg.build()
+
+ assert isinstance(result, ModelData)
+ assert result.test_data is None
+
+ def test_build_raises_if_no_subdirs_found(self, tmp_path):
+ """TFDatasetConfig._load_years() raises FileNotFoundError when no dirs match."""
+ cfg = TFDatasetConfig(
+ directory=str(tmp_path),
+ train_years=[2099],
+ val_years=[],
+ test_years=[],
+ )
+ with pytest.raises(FileNotFoundError, match="No subdirectories found"):
+ cfg.build()
+
+ def test_dataconfig_delegates_to_tf_dataset(self, tmp_path):
+ """DataConfig.build() delegates to TFDatasetConfig when tf_dataset is set."""
+ mock_model_data = ModelData(train_data=MagicMock(), validation_data=MagicMock())
+
+ cfg = DataConfig(
+ train_years=[2010],
+ val_years=[2011],
+ test_years=[],
+ tf_dataset=TFDatasetConfig(
+ directory=str(tmp_path),
+ train_years=[],
+ val_years=[],
+ test_years=[],
+ ),
+ )
+
+ with patch.object(TFDatasetConfig, "build", return_value=mock_model_data) as mock_build:
+ result = cfg.build()
+
+ mock_build.assert_called_once()
+ assert isinstance(result, ModelData)
+
+ def test_dataconfig_tf_dataset_yaml_roundtrip(self, tmp_path, sample_config_dict):
+ """DataConfig with tf_dataset parses correctly from YAML."""
+ config = dict(sample_config_dict)
+ config["data"] = {
+ "train_years": [2010],
+ "val_years": [2011],
+ "test_years": [],
+ "tf_dataset": {
+ "directory": "/tmp/tf_datasets",
+ "train_years": [],
+ "val_years": [],
+ "test_years": [],
+ },
+ }
+ path = tmp_path / "config_tf.yaml"
+ path.write_text(yaml.dump(config, default_flow_style=False))
+
+ from fronts.data.config import DataConfig
+ train_cfg = open_config_yaml_as_dataclass(path=str(path), config_class=TrainConfig)
+ assert isinstance(train_cfg.data, DataConfig)
+ assert isinstance(train_cfg.data.tf_dataset, TFDatasetConfig)
+ assert train_cfg.data.tf_dataset.directory == "/tmp/tf_datasets"
+ assert train_cfg.data.era5 is None
+
+
+# ---------------------------------------------------------------------------
+# Integration: DataConfig in TrainConfig YAML round-trip
+# ---------------------------------------------------------------------------
+
+
+class TestDataConfigInTrainConfig:
+ @pytest.fixture
+ def sample_yaml_with_data(self, tmp_path, sample_config_dict):
+ """Extend sample_config_dict with a minimal data block."""
+ config = dict(sample_config_dict)
+ config["data"] = _minimal_data_config_dict()
+ path = tmp_path / "config_with_data.yaml"
+ path.write_text(yaml.dump(config, default_flow_style=False))
+ return str(path)
+
+ def test_train_config_parses_data_block(self, sample_yaml_with_data):
+ config = open_config_yaml_as_dataclass(
+ path=sample_yaml_with_data, config_class=TrainConfig
+ )
+ assert config is not None
+ assert isinstance(config.data, DataConfig)
+ assert config.data.train_years == [2010]
+ assert isinstance(config.data.era5, ERA5PredictorConfig)
+ assert isinstance(config.data.fronts, FrontsDataConfig)
+
+ def test_actual_1702_yaml_parses_data_block(self):
+ path = os.path.join(os.path.dirname(__file__), "..", "configs", "1702.yaml")
+ if not os.path.exists(path):
+ pytest.skip("1702.yaml not found")
+ config = open_config_yaml_as_dataclass(path=path, config_class=TrainConfig)
+ assert isinstance(config.data, DataConfig)
+ assert config.data.train_years == list(range(2010, 2020))
+ assert config.data.era5.store == (
+ "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
+ )
+ # New unified levels field
+ assert "surface" in config.data.era5.levels
+ assert 1000 in config.data.era5.levels
+ assert "temperature" in config.data.era5.variables
+ assert "mean_sea_level_pressure" in config.data.era5.variables
+
+
+# ---------------------------------------------------------------------------
+# SURFACE_VARIABLE_MAP and SURFACE_ONLY_VARIABLES constants
+# ---------------------------------------------------------------------------
+
+
+class TestSurfaceVariableConstants:
+ def test_surface_variable_map_has_temperature(self):
+ assert "temperature" in SURFACE_VARIABLE_MAP
+ assert SURFACE_VARIABLE_MAP["temperature"] == "2m_temperature"
+
+ def test_surface_variable_map_has_wind_components(self):
+ assert "u_component_of_wind" in SURFACE_VARIABLE_MAP
+ assert "v_component_of_wind" in SURFACE_VARIABLE_MAP
+ assert SURFACE_VARIABLE_MAP["u_component_of_wind"] == "10m_u_component_of_wind"
+ assert SURFACE_VARIABLE_MAP["v_component_of_wind"] == "10m_v_component_of_wind"
+
+ def test_surface_only_variables_contains_mslp(self):
+ assert "mean_sea_level_pressure" in SURFACE_ONLY_VARIABLES
+
+ def test_surface_variable_map_and_surface_only_are_disjoint(self):
+ """No variable should appear in both sets."""
+ overlap = set(SURFACE_VARIABLE_MAP.keys()) & SURFACE_ONLY_VARIABLES
+ assert len(overlap) == 0
+
+
+# ---------------------------------------------------------------------------
+# TimeSelection
+# ---------------------------------------------------------------------------
+
+
+class TestTimeSelection:
+ # --- Construction / validation ---
+
+ def test_most_recent_valid(self):
+ ts = TimeSelection(most_recent=True)
+ assert ts.most_recent is True
+
+ def test_timestamps_valid(self):
+ ts = TimeSelection(timestamps=[datetime.datetime(2024, 6, 1, 12)])
+ assert ts.timestamps is not None
+
+ def test_date_range_valid(self):
+ ts = TimeSelection(
+ date_range=[datetime.datetime(2024, 6, 1), datetime.datetime(2024, 6, 7)]
+ )
+ assert ts.date_range is not None
+
+ def test_raises_if_none_set(self):
+ with pytest.raises(ValueError, match="Exactly one"):
+ TimeSelection()
+
+ def test_raises_if_multiple_set(self):
+ with pytest.raises(ValueError, match="Exactly one"):
+ TimeSelection(
+ most_recent=True,
+ timestamps=[datetime.datetime(2024, 6, 1, 12)],
+ )
+
+ def test_raises_if_date_range_wrong_length(self):
+ with pytest.raises(ValueError, match="exactly two"):
+ TimeSelection(date_range=[datetime.datetime(2024, 6, 1)])
+
+ # --- dacite parsing ---
+
+ def test_dacite_parses_most_recent(self):
+ d = {"most_recent": True}
+ ts = dacite.from_dict(TimeSelection, d, DACITE_CONFIG)
+ assert ts.most_recent is True
+ assert ts.timestamps is None
+ assert ts.date_range is None
+
+ def test_dacite_parses_timestamps(self):
+ """dacite accepts pre-converted datetime objects in timestamps list."""
+ ts1 = datetime.datetime(2024, 6, 1, 12, 0, 0)
+ ts2 = datetime.datetime(2024, 6, 2, 0, 0, 0)
+ d = {"timestamps": [ts1, ts2]}
+ ts = dacite.from_dict(TimeSelection, d, DACITE_CONFIG)
+ assert ts.timestamps is not None
+ assert len(ts.timestamps) == 2
+ assert ts.timestamps[0] == ts1
+
+ def test_dacite_parses_date_range(self):
+ """dacite accepts pre-converted datetime objects in date_range list."""
+ start = datetime.datetime(2024, 6, 1, 0, 0, 0)
+ end = datetime.datetime(2024, 6, 7, 18, 0, 0)
+ d = {"date_range": [start, end]}
+ ts = dacite.from_dict(TimeSelection, d, DACITE_CONFIG)
+ assert ts.date_range is not None
+ assert len(ts.date_range) == 2
+ assert ts.date_range[0] == start
+ assert ts.date_range[1] == end
+
+ # --- apply() ---
+
+ @pytest.fixture
+ def time_ds(self):
+ """Tiny xarray Dataset with 5 timesteps."""
+ times = pd.date_range("2024-06-01", periods=5, freq="6h")
+ return xr.Dataset(
+ {"x": (["time"], np.arange(5, dtype="float32"))},
+ coords={"time": times},
+ )
+
+ def test_apply_most_recent_returns_last_timestep(self, time_ds):
+ ts = TimeSelection(most_recent=True)
+ result = ts.apply(time_ds)
+ assert len(result.time) == 1
+ assert result.time.values[0] == time_ds.time.values[-1]
+
+ def test_apply_timestamps_selects_correct_times(self, time_ds):
+ target = datetime.datetime(2024, 6, 1, 6, 0)
+ ts = TimeSelection(timestamps=[target])
+ result = ts.apply(time_ds)
+ assert len(result.time) == 1
+
+ def test_apply_date_range_selects_slice(self, time_ds):
+ ts = TimeSelection(
+ date_range=[
+ datetime.datetime(2024, 6, 1, 0),
+ datetime.datetime(2024, 6, 1, 12),
+ ]
+ )
+ result = ts.apply(time_ds)
+ # 00Z, 06Z, 12Z = 3 timesteps
+ assert len(result.time) == 3
+
+
+# ---------------------------------------------------------------------------
+# PredictConfig
+# ---------------------------------------------------------------------------
+
+
+def _minimal_predict_config_dict():
+ return {
+ "time_selection": {"most_recent": True},
+ "normalization_method": "standard",
+ "era5": {
+ "domain_extent": [-140.0, -60.0, 20.0, 60.0],
+ "variables": ["temperature", "mean_sea_level_pressure"],
+ "levels": ["surface", 1000, 850],
+ "store": "gs://fake-store",
+ "chunks": {"time": 48},
+ "consolidated": True,
+ },
+ }
+
+
+class TestPredictConfig:
+ def test_dacite_parses_predict_config(self):
+ cfg = dacite.from_dict(PredictConfig, _minimal_predict_config_dict(), DACITE_CONFIG)
+ assert isinstance(cfg.era5, ERA5PredictorConfig)
+ assert isinstance(cfg.time_selection, TimeSelection)
+ assert cfg.normalization_method == "standard"
+
+ def test_most_recent_mode_parsed(self):
+ cfg = dacite.from_dict(PredictConfig, _minimal_predict_config_dict(), DACITE_CONFIG)
+ assert cfg.time_selection.most_recent is True
+
+ def test_era5_levels_parsed(self):
+ cfg = dacite.from_dict(PredictConfig, _minimal_predict_config_dict(), DACITE_CONFIG)
+ assert "surface" in cfg.era5.levels
+ assert 1000 in cfg.era5.levels
+
+ def test_era5_variables_parsed(self):
+ cfg = dacite.from_dict(PredictConfig, _minimal_predict_config_dict(), DACITE_CONFIG)
+ assert "temperature" in cfg.era5.variables
+ assert "mean_sea_level_pressure" in cfg.era5.variables
+
+ def test_timestamps_mode_parsed(self):
+ """dacite accepts pre-converted datetime objects in PredictConfig.time_selection."""
+ d = _minimal_predict_config_dict()
+ ts = datetime.datetime(2024, 6, 1, 12, 0, 0)
+ d["time_selection"] = {"timestamps": [ts]}
+ cfg = dacite.from_dict(PredictConfig, d, DACITE_CONFIG)
+ assert cfg.time_selection.timestamps is not None
+ assert cfg.time_selection.timestamps[0] == ts
+
+ def test_date_range_mode_parsed(self):
+ """dacite accepts pre-converted datetime objects in date_range."""
+ d = _minimal_predict_config_dict()
+ start = datetime.datetime(2024, 6, 1, 0, 0, 0)
+ end = datetime.datetime(2024, 6, 7, 18, 0, 0)
+ d["time_selection"] = {"date_range": [start, end]}
+ cfg = dacite.from_dict(PredictConfig, d, DACITE_CONFIG)
+ assert cfg.time_selection.date_range is not None
+ assert len(cfg.time_selection.date_range) == 2
+
+ @pytest.fixture
+ def raw_zarr_ds(self):
+ """ERA5-like Dataset as it would come from the zarr store: pure integer levels,
+ with a separate surface variable (no level dim) and pressure-level variable.
+ Matches levels=["surface", 1000, 850] in the minimal predict config."""
+ times = pd.date_range("2024-06-01", periods=4, freq="6h")
+ lats = np.linspace(20.0, 60.0, 3)
+ lons = np.linspace(-140.0, -60.0, 3)
+ levels = [1000, 850]
+
+ return xr.Dataset(
+ {
+ "temperature": (
+ ["time", "level", "latitude", "longitude"],
+ np.random.rand(4, 2, 3, 3).astype("float32"),
+ ),
+ "2m_temperature": (
+ ["time", "latitude", "longitude"],
+ np.random.rand(4, 3, 3).astype("float32"),
+ ),
+ "mean_sea_level_pressure": (
+ ["time", "latitude", "longitude"],
+ np.random.rand(4, 3, 3).astype("float32"),
+ ),
+ },
+ coords={
+ "time": times,
+ "level": levels,
+ "latitude": lats,
+ "longitude": lons,
+ },
+ )
+
+ def test_build_returns_dataset(self, raw_zarr_ds):
+ """PredictConfig.build() returns an xr.Dataset."""
+ cfg = dacite.from_dict(PredictConfig, _minimal_predict_config_dict(), DACITE_CONFIG)
+
+ with patch("xarray.open_zarr", return_value=raw_zarr_ds):
+ result = cfg.build()
+
+ assert isinstance(result, xr.Dataset)
+
+ def test_build_calls_time_selection_apply(self, raw_zarr_ds):
+ """PredictConfig.build() delegates time filtering to TimeSelection.apply()."""
+ cfg = dacite.from_dict(PredictConfig, _minimal_predict_config_dict(), DACITE_CONFIG)
+
+ with patch("xarray.open_zarr", return_value=raw_zarr_ds):
+ with patch.object(
+ TimeSelection, "apply", wraps=cfg.time_selection.apply
+ ) as mock_apply:
+ cfg.build()
+
+ mock_apply.assert_called_once()
+
+ def test_predict_config_yaml_roundtrip(self, tmp_path):
+ """predict_1702.yaml loads into PredictConfig via open_config_yaml_as_dataclass."""
+ predict_yaml_path = os.path.join(
+ os.path.dirname(__file__), "..", "configs", "predict_1702.yaml"
+ )
+ if not os.path.exists(predict_yaml_path):
+ pytest.skip("predict_1702.yaml not found")
+
+ cfg = open_config_yaml_as_dataclass(
+ path=predict_yaml_path, config_class=PredictConfig
+ )
+ assert isinstance(cfg, PredictConfig)
+ assert isinstance(cfg.time_selection, TimeSelection)
+ assert cfg.time_selection.most_recent is True
+ # New unified levels: should include "surface" and integer hPa values
+ assert "surface" in cfg.era5.levels
+ assert 1000 in cfg.era5.levels
+ assert 950 in cfg.era5.levels
+ assert 900 in cfg.era5.levels
+ assert 850 in cfg.era5.levels
+ # Variables field
+ assert "temperature" in cfg.era5.variables
+ assert "mean_sea_level_pressure" in cfg.era5.variables
diff --git a/tests/test_keras_builders.py b/tests/test_keras_builders.py
new file mode 100644
index 0000000..4f469e4
--- /dev/null
+++ b/tests/test_keras_builders.py
@@ -0,0 +1,265 @@
+"""Tests for fronts.utils.keras_builders — BaseConfig, registry subclasses, and
+nullable config builders (BiasVectorConfig, KernelMatrixConfig, ConvOutputConfig).
+"""
+
+import pytest
+
+from fronts.utils.keras_builders import (
+ ActivationConfig,
+ BaseConfig,
+ BiasVector,
+ BiasVectorConfig,
+ ConstraintConfig,
+ ConvOutput,
+ ConvOutputConfig,
+ InitializerConfig,
+ KernelMatrix,
+ KernelMatrixConfig,
+ LossConfig,
+ MetricConfig,
+ OptimizerConfig,
+ RegularizerConfig,
+)
+
+
+# ---------------------------------------------------------------------------
+# BaseConfig
+# ---------------------------------------------------------------------------
+
+
+class TestBaseConfig:
+ def test_build_raises_for_unknown_name(self):
+ """BaseConfig.build raises ValueError for an unregistered name."""
+
+ class DummyConfig(BaseConfig):
+ @property
+ def registry(self):
+ return {"known": lambda **kw: "ok"}
+
+ config = DummyConfig(name="unknown", config={})
+ with pytest.raises(ValueError, match="Unsupported DummyConfig: unknown"):
+ config.build()
+
+ def test_build_calls_registered_callable(self):
+ """BaseConfig.build dispatches to the correct registry entry."""
+
+ class DummyConfig(BaseConfig):
+ @property
+ def registry(self):
+ return {"my_thing": lambda x=1: x * 10}
+
+ config = DummyConfig(name="my_thing", config={"x": 5})
+ assert config.build() == 50
+
+ def test_build_passes_config_as_kwargs(self):
+ """Config dict is unpacked as kwargs to the registered callable."""
+ calls = []
+
+ def capture(**kwargs):
+ calls.append(kwargs)
+ return "captured"
+
+ class DummyConfig(BaseConfig):
+ @property
+ def registry(self):
+ return {"cap": capture}
+
+ config = DummyConfig(name="cap", config={"a": 1, "b": 2})
+ result = config.build()
+ assert result == "captured"
+ assert calls == [{"a": 1, "b": 2}]
+
+ def test_build_with_empty_config(self):
+ """Build works with an empty config dict."""
+
+ class DummyConfig(BaseConfig):
+ @property
+ def registry(self):
+ return {"no_args": lambda: "done"}
+
+ config = DummyConfig(name="no_args", config={})
+ assert config.build() == "done"
+
+
+# ---------------------------------------------------------------------------
+# Registry subclasses — ensure registries have expected keys
+# ---------------------------------------------------------------------------
+
+
+class TestConstraintConfig:
+ @pytest.mark.parametrize("name", ["max_norm", "min_max_norm", "non_neg", "unit_norm"])
+ def test_registry_contains_expected_keys(self, name):
+ config = ConstraintConfig(name=name, config={})
+ assert name in config.registry
+
+
+class TestInitializerConfig:
+ @pytest.mark.parametrize(
+ "name",
+ [
+ "glorot_normal", "glorot_uniform", "he_normal", "he_uniform",
+ "identity", "lecun_normal", "lecun_uniform", "ones", "orthogonal",
+ "random_normal", "random_uniform", "truncated_normal",
+ "variance_scaling", "zeros",
+ ],
+ )
+ def test_registry_contains_expected_keys(self, name):
+ config = InitializerConfig(name=name, config={})
+ assert name in config.registry
+
+ def test_build_zeros(self):
+ """Build a zeros initializer with empty config."""
+ config = InitializerConfig(name="zeros", config={})
+ result = config.build()
+ # Result is a MagicMock (standing in for tf.keras.initializers.Zeros)
+ assert result is not None
+
+
+class TestRegularizerConfig:
+ @pytest.mark.parametrize("name", ["l1", "l2", "l1_l2", "orthogonal_regularizer"])
+ def test_registry_contains_expected_keys(self, name):
+ config = RegularizerConfig(name=name, config={})
+ assert name in config.registry
+
+
+class TestOptimizerConfig:
+ def test_registry_contains_adam(self):
+ config = OptimizerConfig(name="Adam", config={})
+ assert "Adam" in config.registry
+
+ def test_build_adam(self):
+ config = OptimizerConfig(name="Adam", config={"beta_1": 0.9, "beta_2": 0.999})
+ result = config.build()
+ assert result is not None
+
+
+class TestActivationConfig:
+ @pytest.mark.parametrize(
+ "name",
+ [
+ "elliott", "elu", "exponential", "gaussian", "gcu", "gelu",
+ "hard_sigmoid", "hexpo", "isigmoid", "leaky_relu", "linear",
+ "lisht", "prelu", "psigmoid", "ptanh", "ptelu", "relu", "resech",
+ "selu", "sigmoid", "smelu", "snake", "softmax", "softplus",
+ "softsign", "srs", "stanh", "swish", "tanh", "thresholded_relu",
+ ],
+ )
+ def test_registry_contains_expected_keys(self, name):
+ config = ActivationConfig(name=name, config={})
+ assert name in config.registry
+
+
+class TestLossConfig:
+ @pytest.mark.parametrize(
+ "name",
+ [
+ "brier_skill_score", "critical_success_index",
+ "fractions_skill_score", "probability_of_detection",
+ ],
+ )
+ def test_registry_contains_expected_keys(self, name):
+ config = LossConfig(name=name, config={})
+ assert name in config.registry
+
+
+class TestMetricConfig:
+ @pytest.mark.parametrize(
+ "name",
+ [
+ "brier_skill_score", "critical_success_index",
+ "fractions_skill_score", "heidke_skill_score",
+ "probability_of_detection",
+ ],
+ )
+ def test_registry_contains_expected_keys(self, name):
+ """MetricConfig name Literal and registry keys must be in sync."""
+ config = MetricConfig(name=name, config={})
+ assert name in config.registry
+
+
+# ---------------------------------------------------------------------------
+# ConvOutputConfig — nullable regularizer
+# ---------------------------------------------------------------------------
+
+
+class TestConvOutputConfig:
+ def test_build_with_none_regularizer(self):
+ """Build succeeds when regularizer is None."""
+ config = ConvOutputConfig(regularizer=None)
+ result = config.build()
+ assert isinstance(result, ConvOutput)
+ assert result.activity_regularizer is None
+
+ def test_build_with_regularizer(self):
+ """Build delegates to the regularizer when provided."""
+ config = ConvOutputConfig(
+ regularizer=RegularizerConfig(name="l2", config={"l2": 0.01})
+ )
+ result = config.build()
+ assert isinstance(result, ConvOutput)
+ assert result.activity_regularizer is not None
+
+
+# ---------------------------------------------------------------------------
+# BiasVectorConfig — nullable constraint and regularizer
+# ---------------------------------------------------------------------------
+
+
+class TestBiasVectorConfig:
+ def test_build_with_all_none(self):
+ """Build succeeds when constraint and regularizer are both None."""
+ config = BiasVectorConfig(
+ constraint=None,
+ initializer=InitializerConfig(name="zeros", config={}),
+ regularizer=None,
+ )
+ result = config.build()
+ assert isinstance(result, BiasVector)
+ assert result.bias_constraint is None
+ assert result.bias_regularizer is None
+ # initializer should still be built
+ assert result.bias_initializer is not None
+
+ def test_build_with_constraint_and_regularizer(self):
+ """Build works when both constraint and regularizer are provided."""
+ config = BiasVectorConfig(
+ constraint=ConstraintConfig(name="non_neg", config={}),
+ initializer=InitializerConfig(name="zeros", config={}),
+ regularizer=RegularizerConfig(name="l1", config={"l1": 0.01}),
+ )
+ result = config.build()
+ assert isinstance(result, BiasVector)
+ assert result.bias_constraint is not None
+ assert result.bias_regularizer is not None
+
+
+# ---------------------------------------------------------------------------
+# KernelMatrixConfig — nullable constraint and regularizer
+# ---------------------------------------------------------------------------
+
+
+class TestKernelMatrixConfig:
+ def test_build_with_all_none(self):
+ """Build succeeds when constraint and regularizer are both None."""
+ config = KernelMatrixConfig(
+ constraint=None,
+ initializer=InitializerConfig(name="glorot_uniform", config={}),
+ regularizer=None,
+ )
+ result = config.build()
+ assert isinstance(result, KernelMatrix)
+ assert result.kernel_constraint is None
+ assert result.kernel_regularizer is None
+ assert result.kernel_initializer is not None
+
+ def test_build_with_constraint_and_regularizer(self):
+ """Build works when both constraint and regularizer are provided."""
+ config = KernelMatrixConfig(
+ constraint=ConstraintConfig(name="max_norm", config={}),
+ initializer=InitializerConfig(name="he_normal", config={}),
+ regularizer=RegularizerConfig(name="l2", config={"l2": 0.01}),
+ )
+ result = config.build()
+ assert isinstance(result, KernelMatrix)
+ assert result.kernel_constraint is not None
+ assert result.kernel_regularizer is not None
diff --git a/tests/test_model.py b/tests/test_model.py
new file mode 100644
index 0000000..436242f
--- /dev/null
+++ b/tests/test_model.py
@@ -0,0 +1,170 @@
+"""Tests for fronts.model — ModelConfig and Model."""
+
+import pytest
+
+from fronts.model import Model, ModelConfig
+from fronts.utils.keras_builders import (
+ ActivationConfig,
+ BiasVectorConfig,
+ ConvOutputConfig,
+ InitializerConfig,
+ KernelMatrixConfig,
+ LossConfig,
+ MetricConfig,
+ OptimizerConfig,
+)
+
+
+@pytest.fixture
+def model_config():
+ """A minimal ModelConfig for testing."""
+ return ModelConfig(
+ name="unet",
+ loss=LossConfig(
+ name="fractions_skill_score", config={"mask_size": [3, 3]}
+ ),
+ metric=MetricConfig(
+ name="critical_success_index",
+ config={"class_weights": [0, 1, 1, 1, 1, 1]},
+ ),
+ optimizer=OptimizerConfig(name="Adam", config={"beta_1": 0.9}),
+ convolution_activity_regularizer=ConvOutputConfig(regularizer=None),
+ bias_vector=BiasVectorConfig(
+ constraint=None,
+ initializer=InitializerConfig(name="zeros", config={}),
+ regularizer=None,
+ ),
+ kernel_matrix=KernelMatrixConfig(
+ constraint=None,
+ initializer=InitializerConfig(name="glorot_uniform", config={}),
+ regularizer=None,
+ ),
+ activation=ActivationConfig(name="gelu", config={}),
+ batch_normalization=True,
+ num_filters=[16, 32, 64, 128],
+ kernel_size=[5, 5, 5],
+ depth=4,
+ modules_per_node=2,
+ padding="same",
+ pool_size=(2, 2, 1),
+ upsample_size=(2, 2, 1),
+ bias=True,
+ )
+
+
+class TestModelConfig:
+ def test_build_returns_model(self, model_config):
+ """ModelConfig.build() returns a Model instance."""
+ # Model.__init__ requires output_activation_config which isn't on
+ # ModelConfig yet (noted as TODO). We test that the parameter names
+ # at least match by supplying it via a subclass/monkey-patch.
+ # For now, verify the build call doesn't crash on parameter names
+ # by adding the missing param.
+ model = model_config.build.__wrapped__ if hasattr(model_config.build, "__wrapped__") else None
+
+ # Directly construct Model to verify param names align
+ m = Model(
+ name=model_config.name,
+ loss_config=model_config.loss,
+ metric_config=model_config.metric,
+ optimizer_config=model_config.optimizer,
+ convolution_activity_regularizer_config=model_config.convolution_activity_regularizer,
+ bias_vector_config=model_config.bias_vector,
+ kernel_matrix_config=model_config.kernel_matrix,
+ activation_config=model_config.activation,
+ output_activation_config=ActivationConfig(name="softmax", config={}),
+ batch_normalization=model_config.batch_normalization,
+ num_filters=model_config.num_filters,
+ kernel_size=model_config.kernel_size,
+ depth=model_config.depth,
+ modules_per_node=model_config.modules_per_node,
+ padding=model_config.padding,
+ pool_size=model_config.pool_size,
+ upsample_size=model_config.upsample_size,
+ bias=model_config.bias,
+ )
+ assert isinstance(m, Model)
+
+
+class TestModel:
+ def test_init_builds_keras_objects(self, model_config):
+ """Model.__init__ calls .build() on all config objects."""
+ m = Model(
+ name="unet",
+ loss_config=model_config.loss,
+ metric_config=model_config.metric,
+ optimizer_config=model_config.optimizer,
+ convolution_activity_regularizer_config=model_config.convolution_activity_regularizer,
+ bias_vector_config=model_config.bias_vector,
+ kernel_matrix_config=model_config.kernel_matrix,
+ activation_config=model_config.activation,
+ output_activation_config=ActivationConfig(name="softmax", config={}),
+ batch_normalization=True,
+ num_filters=[16, 32, 64, 128],
+ kernel_size=[5, 5, 5],
+ depth=4,
+ modules_per_node=2,
+ padding="same",
+ pool_size=(2, 2, 1),
+ upsample_size=(2, 2, 1),
+ bias=True,
+ )
+ # Verify that built objects are stored
+ assert m.loss is not None
+ assert m.metric is not None
+ assert m.optimizer is not None
+ assert m.activity_regularizer is not None
+ assert m.bias_vector is not None
+ assert m.kernel_matrix is not None
+ assert m.activation is not None
+ assert m.output_activation is not None
+
+ def test_init_validates_num_filters_depth_mismatch(self, model_config):
+ """Model raises ValueError if num_filters length != depth."""
+ with pytest.raises(ValueError, match="must match depth"):
+ Model(
+ name="unet",
+ loss_config=model_config.loss,
+ metric_config=model_config.metric,
+ optimizer_config=model_config.optimizer,
+ convolution_activity_regularizer_config=model_config.convolution_activity_regularizer,
+ bias_vector_config=model_config.bias_vector,
+ kernel_matrix_config=model_config.kernel_matrix,
+ activation_config=model_config.activation,
+ output_activation_config=ActivationConfig(name="softmax", config={}),
+ batch_normalization=True,
+ num_filters=[16, 32], # only 2, but depth is 4
+ kernel_size=[5, 5, 5],
+ depth=4,
+ modules_per_node=2,
+ padding="same",
+ pool_size=(2, 2, 1),
+ upsample_size=(2, 2, 1),
+ bias=True,
+ )
+
+ def test_build_is_a_method(self, model_config):
+ """Model.build is a proper method (not a nested function)."""
+ m = Model(
+ name="unet",
+ loss_config=model_config.loss,
+ metric_config=model_config.metric,
+ optimizer_config=model_config.optimizer,
+ convolution_activity_regularizer_config=model_config.convolution_activity_regularizer,
+ bias_vector_config=model_config.bias_vector,
+ kernel_matrix_config=model_config.kernel_matrix,
+ activation_config=model_config.activation,
+ output_activation_config=ActivationConfig(name="softmax", config={}),
+ batch_normalization=True,
+ num_filters=[16, 32, 64, 128],
+ kernel_size=[5, 5, 5],
+ depth=4,
+ modules_per_node=2,
+ padding="same",
+ pool_size=(2, 2, 1),
+ upsample_size=(2, 2, 1),
+ bias=True,
+ )
+ # build() should be callable as a method
+ assert hasattr(m, "build")
+ assert callable(m.build)
diff --git a/tests/test_train.py b/tests/test_train.py
new file mode 100644
index 0000000..96bbdbc
--- /dev/null
+++ b/tests/test_train.py
@@ -0,0 +1,315 @@
+"""Tests for fronts.train — config loading, WandB, Callbacks, Trainer, TrainConfig."""
+
+from unittest.mock import MagicMock
+
+import pytest
+import yaml
+
+from fronts.train import (
+ CallbacksConfig,
+ Trainer,
+ TrainConfig,
+ WandBConfig,
+ open_config_yaml_as_dataclass,
+)
+
+
+# ---------------------------------------------------------------------------
+# open_config_yaml_as_dataclass
+# ---------------------------------------------------------------------------
+
+
+class TestOpenConfigYamlAsDataclass:
+ def test_loads_yaml_into_dataclass(self, sample_yaml_file):
+ """YAML file is correctly loaded and converted to TrainConfig."""
+ result = open_config_yaml_as_dataclass(
+ path=sample_yaml_file, config_class=TrainConfig
+ )
+ assert result is not None
+ assert isinstance(result, TrainConfig)
+ assert result.epochs == 10
+ assert result.seed == 42
+
+ def test_returns_none_when_path_is_empty(self):
+ """Returns None when path is falsy and require is False."""
+ result = open_config_yaml_as_dataclass(
+ path="", config_class=TrainConfig, require=False
+ )
+ assert result is None
+
+ def test_returns_none_when_path_is_none(self):
+ result = open_config_yaml_as_dataclass(
+ path=None, config_class=TrainConfig, require=False
+ )
+ assert result is None
+
+ def test_raises_when_require_true_and_no_path(self):
+ """Raises ValueError when require=True but no path given."""
+ with pytest.raises(ValueError, match="Path must be included"):
+ open_config_yaml_as_dataclass(
+ path="", config_class=TrainConfig, require=True
+ )
+
+ def test_loads_when_require_true_and_path_given(self, sample_yaml_file):
+ """Successfully loads when require=True and a valid path is provided."""
+ result = open_config_yaml_as_dataclass(
+ path=sample_yaml_file, config_class=TrainConfig, require=True
+ )
+ assert result is not None
+ assert isinstance(result, TrainConfig)
+
+ def test_nested_model_config_parsed(self, sample_yaml_file):
+ """Nested model config is properly deserialized."""
+ result = open_config_yaml_as_dataclass(
+ path=sample_yaml_file, config_class=TrainConfig
+ )
+ assert result.model.name == "unet_3plus"
+ assert result.model.depth == 4
+ assert result.model.num_filters == [16, 32, 64, 128]
+
+ def test_nested_wandb_config_parsed(self, sample_yaml_file):
+ """Nested wandb config is properly deserialized."""
+ result = open_config_yaml_as_dataclass(
+ path=sample_yaml_file, config_class=TrainConfig
+ )
+ assert result.wandb.project_name == "test_project"
+ assert result.wandb.model_run_name == "test_run"
+
+
+# ---------------------------------------------------------------------------
+# WandBConfig
+# ---------------------------------------------------------------------------
+
+
+class TestWandBConfig:
+ def test_post_init_calls_login(self):
+ """WandBConfig.__post_init__ triggers wandb.login."""
+ import wandb
+
+ wandb.login.reset_mock()
+ config = WandBConfig(
+ project_name="proj", model_run_name="run", api_key="test_key_123"
+ )
+ wandb.login.assert_called_once_with(key="test_key_123")
+
+ def test_build_init_config_structure(self):
+ config = WandBConfig(project_name="proj", model_run_name="run", api_key="k")
+ init_config = config.build_init_config({"lr": 0.001})
+ assert init_config["project"] == "proj"
+ assert init_config["name"] == "run"
+ assert init_config["config"] == {"lr": 0.001}
+
+ def test_build_all_callbacks_returns_two(self):
+ config = WandBConfig(project_name="proj", model_run_name="run", api_key="k")
+ callbacks = config.build_all_callbacks()
+ assert len(callbacks) == 2
+
+ def test_default_values(self):
+ config = WandBConfig(project_name="proj", model_run_name="run", api_key="k")
+ assert config.log_frequency == 1
+ assert config.upload_checkpoints is False
+ assert config.wandb_filepath == "models"
+
+
+# ---------------------------------------------------------------------------
+# CallbacksConfig
+# ---------------------------------------------------------------------------
+
+
+class TestCallbacksConfig:
+ def test_build_empty_when_no_optional_paths(self):
+ """Returns empty list when no checkpoint/csv/patience configured."""
+ config = CallbacksConfig(
+ monitor="val_loss",
+ verbose=1,
+ save_best_only=True,
+ save_weights_only=False,
+ save_freq="epoch",
+ )
+ callbacks = config.build()
+ assert callbacks == []
+
+ def test_build_includes_checkpoint_when_path_set(self):
+ config = CallbacksConfig(
+ monitor="val_loss",
+ verbose=1,
+ save_best_only=True,
+ save_weights_only=False,
+ save_freq="epoch",
+ model_checkpoint_path="/tmp/model.h5",
+ )
+ callbacks = config.build()
+ assert len(callbacks) == 1
+
+ def test_build_includes_csv_logger_when_path_set(self):
+ config = CallbacksConfig(
+ monitor="val_loss",
+ verbose=0,
+ save_best_only=False,
+ save_weights_only=False,
+ save_freq="epoch",
+ csv_logger_path="/tmp/log.csv",
+ )
+ callbacks = config.build()
+ assert len(callbacks) == 1
+
+ def test_build_includes_early_stopping_when_patience_set(self):
+ config = CallbacksConfig(
+ monitor="val_loss",
+ verbose=0,
+ save_best_only=False,
+ save_weights_only=False,
+ save_freq="epoch",
+ patience=10,
+ )
+ callbacks = config.build()
+ assert len(callbacks) == 1
+
+ def test_build_includes_all_three(self):
+ config = CallbacksConfig(
+ monitor="val_loss",
+ verbose=1,
+ save_best_only=True,
+ save_weights_only=True,
+ save_freq="epoch",
+ model_checkpoint_path="/tmp/model.h5",
+ csv_logger_path="/tmp/log.csv",
+ patience=5,
+ )
+ callbacks = config.build()
+ assert len(callbacks) == 3
+
+
+# ---------------------------------------------------------------------------
+# Trainer
+# ---------------------------------------------------------------------------
+
+
+class TestTrainer:
+ def test_mutable_default_callbacks_isolation(self):
+ """Each Trainer gets its own callback list (no shared mutable default)."""
+ t1 = Trainer(
+ model=MagicMock(),
+ data=MagicMock(),
+ epochs=1,
+ validation_frequency=1,
+ training_steps_per_epoch=1,
+ validation_steps_per_epoch=1,
+ )
+ t2 = Trainer(
+ model=MagicMock(),
+ data=MagicMock(),
+ epochs=1,
+ validation_frequency=1,
+ training_steps_per_epoch=1,
+ validation_steps_per_epoch=1,
+ )
+ t1.callbacks.append("extra")
+ assert "extra" not in t2.callbacks
+
+ def test_wandb_callbacks_merged(self):
+ """When wandb is provided, its callbacks are merged in."""
+ mock_wandb = MagicMock(spec=WandBConfig)
+ mock_wandb.build_all_callbacks.return_value = ["wandb_cb1", "wandb_cb2"]
+ t = Trainer(
+ model=MagicMock(),
+ data=MagicMock(),
+ epochs=1,
+ validation_frequency=1,
+ training_steps_per_epoch=1,
+ validation_steps_per_epoch=1,
+ callbacks=["my_cb"],
+ wandb=mock_wandb,
+ )
+ assert "my_cb" in t.callbacks
+ assert "wandb_cb1" in t.callbacks
+ assert "wandb_cb2" in t.callbacks
+
+ def test_no_wandb_passthrough(self):
+ """Without wandb, callbacks are passed through unchanged."""
+ t = Trainer(
+ model=MagicMock(),
+ data=MagicMock(),
+ epochs=1,
+ validation_frequency=1,
+ training_steps_per_epoch=1,
+ validation_steps_per_epoch=1,
+ callbacks=["my_cb"],
+ )
+ assert t.callbacks == ["my_cb"]
+
+ def test_train_calls_fit_without_wandb(self):
+ """Trainer.train calls model.fit when wandb is None."""
+ mock_model = MagicMock()
+ mock_data = MagicMock()
+ t = Trainer(
+ model=mock_model,
+ data=mock_data,
+ epochs=5,
+ validation_frequency=1,
+ training_steps_per_epoch=10,
+ validation_steps_per_epoch=3,
+ )
+ t.train(model={})
+ mock_model.fit.assert_called_once()
+
+ def test_train_calls_fit_with_wandb(self):
+ """Trainer.train calls model.fit inside wandb.init context when wandb is set."""
+ import wandb
+
+ mock_model = MagicMock()
+ mock_data = MagicMock()
+ mock_wandb_config = MagicMock(spec=WandBConfig)
+ mock_wandb_config.build_init_config.return_value = {
+ "project": "test",
+ "config": {},
+ "name": "run",
+ }
+ mock_wandb_config.build_all_callbacks.return_value = []
+
+ t = Trainer(
+ model=mock_model,
+ data=mock_data,
+ epochs=5,
+ validation_frequency=1,
+ training_steps_per_epoch=10,
+ validation_steps_per_epoch=3,
+ wandb=mock_wandb_config,
+ )
+ t.train(model={"lr": 0.001})
+ mock_wandb_config.build_init_config.assert_called_once_with({"lr": 0.001})
+ mock_model.fit.assert_called_once()
+
+
+# ---------------------------------------------------------------------------
+# TrainConfig
+# ---------------------------------------------------------------------------
+
+
+class TestTrainConfig:
+ def test_build_returns_trainer(self, sample_yaml_file):
+ """TrainConfig.build() returns a Trainer instance."""
+ config = open_config_yaml_as_dataclass(
+ path=sample_yaml_file, config_class=TrainConfig
+ )
+ trainer = config.build()
+ assert isinstance(trainer, Trainer)
+
+ def test_build_passes_self_wandb(self, sample_yaml_file):
+ """TrainConfig.build() passes self.wandb (not the wandb module)."""
+ config = open_config_yaml_as_dataclass(
+ path=sample_yaml_file, config_class=TrainConfig
+ )
+ trainer = config.build()
+ assert trainer.wandb is config.wandb
+
+ def test_build_propagates_training_params(self, sample_yaml_file):
+ """Training parameters flow through to the Trainer."""
+ config = open_config_yaml_as_dataclass(
+ path=sample_yaml_file, config_class=TrainConfig
+ )
+ trainer = config.build()
+ assert trainer.epochs == 10
+ assert trainer.seed == 42
+ assert trainer.repeat is True
+ assert trainer.validation_frequency == 1
diff --git a/train_model.py b/train_model.py
deleted file mode 100644
index c322d25..0000000
--- a/train_model.py
+++ /dev/null
@@ -1,448 +0,0 @@
-"""
-Function that trains a new U-Net model.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.8.9
-"""
-
-import argparse
-import pandas as pd
-from tensorflow.keras.callbacks import EarlyStopping, CSVLogger
-import tensorflow as tf
-import pickle
-import numpy as np
-import file_manager as fm
-import os
-import custom_losses
-import custom_metrics
-import models
-import datetime
-from utils import settings, misc, data_utils
-
-
-if __name__ == "__main__":
- parser = argparse.ArgumentParser()
-
- parser.add_argument('--model_dir', type=str, required=True, help='Directory where the models are or will be saved to.')
- parser.add_argument('--model_number', type=int, help='Number that the model will be assigned.')
- parser.add_argument('--era5_tf_indirs', type=str, required=True, nargs='+',
- help='Directories for the TensorFlow datasets. One or two paths can be passed. If only one path is passed, then the '
- 'training and validation datasets will be pulled from this path. If two paths are passed, the training dataset '
- 'will be pulled from the first path and the validation dataset from the second.')
- parser.add_argument('--epochs', type=int, required=True, help='Number of epochs for the U-Net training.')
- parser.add_argument('--patience', type=int,
- help='Patience for EarlyStopping callback. If this argument is not provided, it will be set according to the size '
- 'of the training dataset (images in training set divided by the product of the batch size and steps).')
- parser.add_argument('--verbose', type=int, default=2,
- help='Model.fit verbose. Unless you want a text file that is several hundred megabytes in size and takes 10 years '
- 'to scroll through, I suggest you leave this at 2.')
-
- ### GPU and hardware arguments ###
- parser.add_argument('--gpu_device', type=int, nargs='+', help='GPU device numbers.')
- parser.add_argument('--memory_growth', action='store_true', help='Use memory growth for GPUs')
- parser.add_argument('--num_parallel_calls', type=int, default=4,
- help='Number of parallel calls for retrieving batches for the training and validation datasets.')
- parser.add_argument('--disable_tensorfloat32', action='store_true', help='Disable TensorFloat32 execution.')
-
- ### Hyperparameters ###
- parser.add_argument('--learning_rate', type=float,
- help='Learning rate for U-Net optimizer. If left as None, then the default optimizer learning rate will be used.')
- parser.add_argument('--batch_size', type=int, required=True, nargs='+',
- help='Batch sizes for the U-Net. Up to 2 arguments can be passed. If 1 argument is passed, the value will be both '
- 'the training and validation batch sizes. If 2 arguments are passed, the first and second arguments will be '
- 'the training and validation batch sizes, respectively.')
- parser.add_argument('--steps', type=int, required=True, nargs='+',
- help='Number of steps for each epoch. Up to 2 arguments can be passed. If 1 argument is passed, the value will only '
- 'be applied to the number of steps per epoch, and the number of validation steps will be calculated by tensorflow '
- 'such that the entire validation dataset is passed into the model during validation. If 2 arguments are passed, '
- 'then the arguments are the number of steps in training and validation. If no arguments are passed, then the '
- 'number of steps in both training and validation will be calculated by tensorflow.')
- parser.add_argument('--valid_freq', type=int, default=1, help='How many epochs to complete before each validation.')
-
- ### U-Net arguments ###
- parser.add_argument('--model_type', type=str,
- help='Model type. Options are: unet, unet_ensemble, unet_plus, unet_2plus, unet_3plus.')
- parser.add_argument('--activation', type=str,
- help='Activation function to use in the U-Net. Refer to utils.unet_utils.choose_activation_layer to see all available '
- 'activation functions.')
- parser.add_argument('--batch_normalization', action='store_true',
- help='Use batch normalization in the model. This will place batch normalization layers after each convolution layer.')
- parser.add_argument('--deep_supervision', action='store_true', help='Use deep supervision in the U-Net.')
- parser.add_argument('--filter_num', type=int, nargs='+', help='Number of filters in each level of the U-Net.')
- parser.add_argument('--filter_num_aggregate', type=int,
- help='Number of filters in aggregated feature maps. This argument is only used in the U-Net 3+ model.')
- parser.add_argument('--filter_num_skip', type=int, help='Number of filters in full-scale skip connections in the U-Net 3+.')
- parser.add_argument('--first_encoder_connections', action='store_true', help='Enable first encoder connections in the U-Net 3+.')
- parser.add_argument('--kernel_size', type=int, nargs='+', help='Size of the convolution kernels.')
- parser.add_argument('--levels', type=int, help='Number of levels in the U-Net.')
- parser.add_argument('--loss', type=str, nargs='+',
- help="Loss function for the U-Net (arg 1), with keyword arguments (arg 2). Keyword arguments must be passed as a "
- "string in the second argument. See 'utils.misc.string_arg_to_dict' for more details. Raises a ValueError if "
- "more than 2 arguments are passed.")
- parser.add_argument('--metric', type=str, nargs='+',
- help="Metric for evaluating the U-Net during training (arg 1), with keyword arguments (arg 2). Keyword arguments "
- "must be passed as a string in the second argument. See 'utils.misc.string_arg_to_dict' for more details. Raises "
- "a ValueError if more than 2 arguments are passed.")
- parser.add_argument('--modules_per_node', type=int, default=5, help='Number of convolution modules in each node')
- parser.add_argument('--optimizer', type=str, nargs='+', default=['Adam', ],
- help="Optimizer to use during the training process (arg 1), with keyword arguments (arg 2). Keyword arguments "
- "must be passed as a string in the second argument. See 'utils.misc.string_arg_to_dict' for more details. Raises "
- "a ValueError if more than 2 arguments are passed.")
- parser.add_argument('--padding', type=str, default='same', help='Padding to use in the model')
- parser.add_argument('--pool_size', type=int, nargs='+', help='Pool size for the MaxPooling layers in the U-Net.')
- parser.add_argument('--upsample_size', type=int, nargs='+', help='Upsample size for the UpSampling layers in the U-Net.')
- parser.add_argument('--use_bias', action='store_true', help='Use bias parameters in the U-Net')
-
- ### Constraints, initializers, and regularizers ###
- parser.add_argument('--activity_regularizer', type=str, nargs='+', default=[None, ],
- help='Regularizer function applied to the output of the Conv2D/Conv3D layers. A second string argument can be passed '
- 'containing keyword arguments for the regularizer.')
- parser.add_argument('--bias_constraint', type=str, nargs='+', default=[None, ],
- help='Constraint function applied to the bias vector of the Conv2D/Conv3D layers. A second string argument can be '
- 'passed containing keyword arguments for the constraint.')
- parser.add_argument('--bias_initializer', type=str, default='zeros', help='Initializer for the bias vector in the Conv2D/Conv3D layers.')
- parser.add_argument('--bias_regularizer', type=str, nargs='+', default=[None, ],
- help='Regularizer function applied to the bias vector in the Conv2D/Conv3D layers. A second string argument can '
- 'be passed containing keyword arguments for the regularizer.')
- parser.add_argument('--kernel_constraint', type=str, nargs='+', default=[None, ],
- help='Constraint function applied to the kernel matrix of the Conv2D/Conv3D layers. A second string argument can '
- 'be passed containing keyword arguments for the constraint.')
- parser.add_argument('--kernel_initializer', type=str, default='glorot_uniform', help='Initializer for the kernel weights matrix in the Conv2D/Conv3D layers.')
- parser.add_argument('--kernel_regularizer', type=str, nargs='+', default=[None, ],
- help='Regularizer function applied to the kernel weights matrix in the Conv2D/Conv3D layers. A second string argument '
- 'can be passed containing keyword arguments for the regularizer.')
-
- ### Data arguments ###
- parser.add_argument('--num_training_years', type=int, help='Number of years for the training dataset.')
- parser.add_argument('--training_years', type=int, nargs="+", help='Years for the training dataset.')
- parser.add_argument('--num_validation_years', type=int, help='Number of years for the validation set.')
- parser.add_argument('--validation_years', type=int, nargs="+", help='Years for the validation set.')
- parser.add_argument('--shuffle', type=str, default='full',
- help="Shuffling method for the training set. Valid options are 'lazy' or 'full' (default is 'full'). "
- "A 'lazy' shuffle will only shuffle the order of the monthly datasets but not the contents within. A 'full' "
- "shuffle will shuffle every image inside the dataset.")
-
- ### Retraining model ###
- parser.add_argument('--retrain', action='store_true', help='Retrain a model')
-
- parser.add_argument('--override_directory_check', action='store_true',
- help="Override the OSError caused by creating a new model directory that already exists. Normally, if the script "
- "crashes before or during the training of a new model, an OSError will be returned if the script is immediately "
- "ran again with the same model number as the model directory already exists. This is an intentional fail-safe "
- "designed to prevent models that already exist from being overwritten. Passing this boolean flag disables the "
- "fail-safe and can be useful if the script is being ran on a workload manager (e.g. SLURM) where jobs can fail "
- "and then be immediately requeued and ran again.")
-
- args = vars(parser.parse_args())
-
- if len(args['era5_tf_indirs']) > 2:
- raise ValueError("Only 1 or 2 paths can be passed into --era5_tf_indirs, received %d paths" % len(args['era5_tf_indirs']))
- elif len(args['era5_tf_indirs']) == 1:
- args['era5_tf_indirs'].append(args['era5_tf_indirs'][0])
-
- if args['shuffle'] != 'lazy' and args['shuffle'] != 'full':
- raise ValueError("Unrecognized shuffling method: %s. Valid methods are 'lazy' or 'full'" % args['shuffle'])
-
- # Check arguments that can only have a maximum length of 2
- for arg in ['loss', 'metric', 'optimizer', 'activity_regularizer', 'bias_constraint', 'bias_regularizer', 'kernel_constraint',
- 'kernel_regularizer', 'batch_size', 'steps']:
- if len(args[arg]) > 2:
- raise ValueError("--%s can only take up to 2 arguments" % arg)
-
- ### Dictionary containing arguments that cannot be used for specific model types ###
- incompatible_args = {'unet': dict(deep_supervision=False, first_encoder_connections=False),
- 'unet_ensemble': dict(deep_supervision=False, first_encoder_connections=False),
- 'unet_plus': dict(first_encoder_connections=False),
- 'unet_2plus': dict(first_encoder_connections=False),
- 'unet_3plus': {},
- 'attention_unet': dict(upsample_size=None, deep_supervision=False, first_encoder_connections=False)}
-
- ### Make sure that incompatible arguments were not passed, and raise errors if they were passed ###
- incompatible_args_for_model = incompatible_args[args['model_type']]
- for arg in incompatible_args_for_model:
- if incompatible_args_for_model[arg] != args[arg]:
- raise ValueError(f"--{arg} must be '{incompatible_args_for_model[arg]}' when the model type is {args['model_type']}")
-
- ### Convert keyword argument strings to dictionaries ###
- loss_args = misc.string_arg_to_dict(args['loss'][1]) if len(args['loss']) > 1 else dict()
- metric_args = misc.string_arg_to_dict(args['metric'][1]) if len(args['metric']) > 1 else dict()
- optimizer_args = misc.string_arg_to_dict(args['optimizer'][1]) if len(args['optimizer']) > 1 else dict()
- activity_regularizer_args = misc.string_arg_to_dict(args['activity_regularizer'][1]) if len(args['activity_regularizer']) > 1 else dict()
- bias_constraint_args = misc.string_arg_to_dict(args['bias_constraint'][1]) if len(args['bias_constraint']) > 1 else dict()
- bias_regularizer_args = misc.string_arg_to_dict(args['bias_regularizer'][1]) if len(args['bias_regularizer']) > 1 else dict()
- kernel_constraint_args = misc.string_arg_to_dict(args['kernel_constraint'][1]) if len(args['kernel_constraint']) > 1 else dict()
- kernel_regularizer_args = misc.string_arg_to_dict(args['kernel_regularizer'][1]) if len(args['kernel_regularizer']) > 1 else dict()
-
- # learning rate is part of the optimizer
- if args['learning_rate'] is not None:
- optimizer_args['learning_rate'] = args['learning_rate']
-
- gpus = tf.config.list_physical_devices(device_type='GPU') # Find available GPUs
- if len(gpus) > 0:
-
- print("Number of GPUs available: %d" % len(gpus))
-
- # Only make the selected GPU(s) visible to TensorFlow
- if args['gpu_device'] is not None:
- tf.config.set_visible_devices(devices=[gpus[gpu] for gpu in args['gpu_device']], device_type='GPU')
- gpus = tf.config.get_visible_devices(device_type='GPU') # List of selected GPUs
- print("Using %d GPU(s):" % len(gpus), gpus)
-
- # Disable TensorFloat32 for matrix multiplication
- if args['disable_tensorfloat32']:
- tf.config.experimental.enable_tensor_float_32_execution(False)
-
- # Allow for memory growth on the GPU. This will only use the GPU memory that is required rather than allocating all of the GPU's memory.
- if args['memory_growth']:
- tf.config.experimental.set_memory_growth(device=[gpu for gpu in gpus][0], enable=True)
-
- else:
- print('WARNING: No GPUs found, all computations will be performed on CPUs.')
- tf.config.set_visible_devices([], 'GPU')
-
- train_dataset_properties = pd.read_pickle('%s/dataset_properties.pkl' % args['era5_tf_indirs'][0])
- valid_dataset_properties = pd.read_pickle('%s/dataset_properties.pkl' % args['era5_tf_indirs'][1])
-
- """
- Verify that the training and validation datasets have the same front types, variables, pressure levels, number of
- dimensions, and normalization parameters.
- """
- try:
- assert train_dataset_properties['front_types'] == valid_dataset_properties['front_types']
- assert train_dataset_properties['variables'] == valid_dataset_properties['variables']
- assert train_dataset_properties['pressure_levels'] == valid_dataset_properties['pressure_levels']
- assert all(train_dataset_properties['num_dims'][num] == valid_dataset_properties['num_dims'][num] for num in range(2))
- assert train_dataset_properties['normalization_parameters'] == valid_dataset_properties['normalization_parameters']
- except AssertionError:
- raise TypeError("Training and validation dataset properties do not match. Select a different dataset(s) or choose "
- "one dataset to use for both training and validation.")
-
- front_types = train_dataset_properties['front_types']
- variables = train_dataset_properties['variables']
- pressure_levels = train_dataset_properties['pressure_levels']
- image_size = train_dataset_properties['image_size']
- num_dims = train_dataset_properties['num_dims']
-
- if not args['retrain']:
-
- if args['loss'][0] == 'fractions_skill_score':
- loss_string = 'fss_loss'
- elif args['loss'][0] == 'critical_success_index':
- loss_string = 'csi_loss'
- elif args['loss'][0] == 'brier_skill_score':
- loss_string = 'bss_loss'
- else:
- loss_string = None
-
- if args['metric'][0] == 'fractions_skill_score':
- metric_string = 'fss'
- elif args['metric'][0] == 'critical_success_index':
- metric_string = 'csi'
- elif args['metric'][0] == 'brier_skill_score':
- metric_string = 'bss'
- else:
- metric_string = None
-
- all_years = np.arange(2008, 2021)
-
- if args['num_training_years'] is not None:
- if args['training_years'] is not None:
- raise TypeError("Cannot explicitly declare the training years if --num_training_years is passed")
- training_years = list(sorted(np.random.choice(all_years, args['num_training_years'], replace=False)))
- else:
- if args['training_years'] is None:
- raise TypeError("Must pass one of the following arguments: --training_years, --num_training_years")
- training_years = list(sorted(args['training_years']))
-
- if args['num_validation_years'] is not None:
- if args['validation_years'] is not None:
- raise TypeError("Cannot explicitly declare the validation years if --num_validation_years is passed")
- validation_years = list(sorted(np.random.choice([year for year in all_years if year not in training_years], args['num_validation_years'], replace=False)))
- else:
- if args['validation_years'] is None:
- raise TypeError("Must pass one of the following arguments: --validation_years, --num_validation_years")
- validation_years = list(sorted(args['validation_years']))
-
- if len(training_years) + len(validation_years) > 12:
- raise ValueError("No testing years are available: the total number of training and validation years cannot be greater than 12")
-
- test_years = [year for year in all_years if year not in training_years + validation_years]
-
- # If no model number was provided, select a number based on the current date and time.
- model_number = int(datetime.datetime.utcnow().timestamp() % 1e8) if args['model_number'] is None else args['model_number']
-
- # Convert pool size and upsample size to tuples
- pool_size = tuple(args['pool_size']) if args['pool_size'] is not None else None
- upsample_size = tuple(args['upsample_size']) if args['upsample_size'] is not None else None
-
- if any(front_type == front_types for front_type in [['MERGED-F_BIN'], ['MERGED-T'], ['F_BIN']]):
- num_classes = 2
- elif front_types == ['MERGED-F']:
- num_classes = 5
- elif front_types == ['MERGED-ALL']:
- num_classes = 8
- else:
- num_classes = len(front_types) + 1
-
- # Create dictionary containing information about the model. This simplifies the process of loading the model
- model_properties = dict({})
- model_properties['domains'] = [train_dataset_properties['domain'], valid_dataset_properties['domain']]
- model_properties['normalization_parameters'] = data_utils.normalization_parameters
- model_properties['dataset_properties'] = train_dataset_properties
- model_properties['classes'] = num_classes
-
- # Place provided arguments into the model properties dictionary
- for arg in ['model_type', 'learning_rate', 'deep_supervision', 'model_number', 'kernel_size', 'modules_per_node',
- 'activation', 'batch_normalization', 'padding', 'use_bias', 'activity_regularizer', 'bias_constraint',
- 'bias_initializer', 'bias_regularizer', 'kernel_constraint', 'kernel_initializer', 'kernel_regularizer',
- 'first_encoder_connections', 'valid_freq', 'optimizer']:
- model_properties[arg] = args[arg]
-
- # Place local variables into the model properties dictionary
- for arg in ['loss_string', 'loss_args', 'metric_string', 'metric_args', 'image_size', 'training_years',
- 'validation_years', 'test_years']:
- model_properties[arg] = locals()[arg]
-
- # If using 3D inputs and 2D targets, squeeze out the vertical dimension of the model (index 2)
- squeeze_dims = 2 if num_dims == [3, 2] else None
-
- train_batch_size = args['batch_size'][0]
- valid_batch_size = args['batch_size'][0] if len(args['batch_size']) == 1 else args['batch_size'][1]
- train_steps = args['steps'][0]
- valid_steps = None if len(args['steps']) < 2 else args['steps'][1]
- valid_freq = args['valid_freq']
-
- model_properties['batch_sizes'] = [train_batch_size, valid_batch_size]
- model_properties['steps_per_epoch'] = [train_steps, valid_steps]
-
- unet_model = getattr(models, args['model_type'])
- unet_model_args = unet_model.__code__.co_varnames[:unet_model.__code__.co_argcount] # pull argument names from unet function
-
- ### Arguments for the function used to build the U-Net ###
- unet_kwargs = {arg: args[arg] for arg in ['pool_size', 'upsample_size', 'levels', 'filter_num', 'kernel_size', 'modules_per_node',
- 'activation', 'batch_normalization', 'padding', 'use_bias', 'bias_initializer', 'kernel_initializer', 'first_encoder_connections',
- 'deep_supervision'] if arg in unet_model_args}
- unet_kwargs['squeeze_dims'] = squeeze_dims
- unet_kwargs['activity_regularizer'] = getattr(tf.keras.regularizers, args['activity_regularizer'][0])(**activity_regularizer_args) if args['activity_regularizer'][0] is not None else None
- unet_kwargs['bias_constraint'] = getattr(tf.keras.constraints, args['bias_constraint'][0])(**bias_constraint_args) if args['bias_constraint'][0] is not None else None
- unet_kwargs['kernel_constraint'] = getattr(tf.keras.constraints, args['kernel_constraint'][0])(**kernel_constraint_args) if args['kernel_constraint'][0] is not None else None
- unet_kwargs['bias_regularizer'] = getattr(tf.keras.regularizers, args['bias_regularizer'][0])(**bias_regularizer_args) if args['bias_regularizer'][0] is not None else None
- unet_kwargs['kernel_regularizer'] = getattr(tf.keras.regularizers, args['kernel_regularizer'][0])(**kernel_regularizer_args) if args['kernel_regularizer'][0] is not None else None
-
- print("Training years:", training_years)
- print("Validation years:", validation_years)
- print("Test years:", test_years)
-
- else:
-
- model_number = args['model_number']
-
- model_properties = pd.read_pickle('%s/model_%d/model_%d_properties.pkl' % (args['model_dir'], model_number, model_number))
-
- front_types = model_properties['front_types']
-
- training_years = model_properties['training_years']
- validation_years = model_properties['validation_years']
- test_years = model_properties['test_years']
-
- train_batch_size, valid_batch_size = model_properties['batch_sizes']
- train_steps, valid_steps = model_properties['steps_per_epoch']
- valid_freq = model_properties['valid_freq']
-
- model_filepath = '%s/model_%d/model_%d.h5' % (args['model_dir'], model_number, model_number)
- history_filepath = '%s/model_%d/model_%d_history.csv' % (args['model_dir'], model_number, model_number)
-
- ### Training dataset ###
- train_files_obj = fm.DataFileLoader(args['era5_tf_indirs'][0], data_file_type='era5-tensorflow')
- train_files_obj.training_years = training_years
- train_files_obj.pair_with_fronts(args['era5_tf_indirs'][0], front_types=front_types)
- training_inputs = train_files_obj.data_files_training
- training_labels = train_files_obj.front_files_training
-
- # Shuffle monthly data lazily
- if args['shuffle'] == 'lazy':
- training_files = list(zip(training_inputs, training_labels))
- np.random.shuffle(training_files)
- training_inputs, training_labels = zip(*training_files)
-
- training_dataset = data_utils.combine_datasets(training_inputs, training_labels)
- print(f"Images in training dataset: {len(training_dataset):,}")
-
- """
- If the patience argument is not explicitly provided, derive it from the size of the training dataset along with the
- batch size and number of steps per epoch.
- """
- if args['patience'] is None:
- patience = int(len(training_dataset) / (train_batch_size * train_steps)) + 1
- print("Using patience value of %d epochs for early stopping" % patience)
- else:
- patience = args['patience']
-
- # Shuffle the entire training dataset
- if args['shuffle'] == 'full':
- training_buffer_size = np.min([len(training_dataset), settings.MAX_TRAIN_BUFFER_SIZE])
- training_dataset = training_dataset.shuffle(buffer_size=training_buffer_size)
-
- training_dataset = training_dataset.batch(train_batch_size, drop_remainder=True, num_parallel_calls=args['num_parallel_calls'])
- training_dataset = training_dataset.prefetch(tf.data.AUTOTUNE)
-
- ### Validation dataset ###
- valid_files_obj = fm.DataFileLoader(args['era5_tf_indirs'][1], data_file_type='era5-tensorflow')
- valid_files_obj.validation_years = validation_years
- valid_files_obj.pair_with_fronts(args['era5_tf_indirs'][1], front_types=front_types)
- validation_inputs = valid_files_obj.data_files_validation
- validation_labels = valid_files_obj.front_files_validation
- validation_dataset = data_utils.combine_datasets(validation_inputs, validation_labels)
- print(f"Images in validation dataset: {len(validation_dataset):,}")
- validation_dataset = validation_dataset.batch(valid_batch_size, drop_remainder=True, num_parallel_calls=args['num_parallel_calls'])
- validation_dataset = validation_dataset.prefetch(tf.data.AUTOTUNE)
-
- # Set the lat/lon dimensions to have a None shape so images of any sized can be passed into the U-Net
- input_shape = list(training_dataset.take(0).element_spec[0].shape[1:])
- for i in range(2):
- input_shape[i] = None
- input_shape = tuple(input_shape)
-
- with tf.distribute.MirroredStrategy().scope():
-
- if not args['retrain']:
- model = unet_model(input_shape, num_classes, **unet_kwargs)
- loss_function = getattr(custom_losses, args['loss'][0])(**loss_args)
- metric_function = getattr(custom_metrics, args['metric'][0])(**metric_args)
- optimizer = getattr(tf.keras.optimizers, args['optimizer'][0])(**optimizer_args)
- model.compile(loss=loss_function, optimizer=optimizer, metrics=metric_function)
- else:
- model = fm.load_model(args['model_number'], args['model_dir'])
-
- model.summary()
-
- if not args['retrain']:
-
- model_properties = {key: model_properties[key] for key in sorted(model_properties.keys())} # Sort model properties dictionary alphabetically
-
- if not os.path.isdir('%s/model_%d' % (args['model_dir'], model_number)):
- os.mkdir('%s/model_%d' % (args['model_dir'], model_number)) # Make folder for model
- os.mkdir('%s/model_%d/maps' % (args['model_dir'], model_number)) # Make folder for model predicton maps
- os.mkdir('%s/model_%d/probabilities' % (args['model_dir'], model_number)) # Make folder for prediction data files
- os.mkdir('%s/model_%d/statistics' % (args['model_dir'], model_number)) # Make folder for statistics data files
- elif not args['override_directory_check']:
- raise OSError('%s/model_%d already exists. If model %d still needs to be created and trained, run this script '
- 'again with the --override_directory_check flag.' % (args['model_dir'], model_number, model_number))
- elif os.path.isfile(model_filepath):
- raise OSError('model %d already exists at %s. Choose a different model number and try again.' % (model_number, model_filepath))
-
- with open('%s/model_%d/model_%d_properties.pkl' % (args['model_dir'], model_number, model_number), 'wb') as f:
- pickle.dump(model_properties, f)
-
- with open('%s/model_%d/model_%d_properties.txt' % (args['model_dir'], model_number, model_number), 'w') as f:
- for key in model_properties.keys():
- f.write(f"{key}: {model_properties[key]}\n")
-
- checkpoint = tf.keras.callbacks.ModelCheckpoint(model_filepath, monitor='val_loss', verbose=1, save_best_only=True,
- save_weights_only=False, save_freq='epoch') # ModelCheckpoint: saves model at a specified interval
- early_stopping = EarlyStopping('val_loss', patience=patience, verbose=1) # EarlyStopping: stops training early if the validation loss does not improve after a specified number of epochs (patience)
- history_logger = CSVLogger(history_filepath, separator=",", append=True) # Saves loss and metric data every epoch
-
- model.fit(training_dataset.repeat(), validation_data=validation_dataset, validation_freq=valid_freq, epochs=args['epochs'],
- steps_per_epoch=train_steps, validation_steps=valid_steps, callbacks=[early_stopping, checkpoint, history_logger],
- verbose=args['verbose'])
diff --git a/utils/change_model_number.py b/utils/change_model_number.py
deleted file mode 100644
index 6ba1e69..0000000
--- a/utils/change_model_number.py
+++ /dev/null
@@ -1,37 +0,0 @@
-"""
-Script that changes the number of a model and its data files.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.6.26
-"""
-import os
-from glob import glob
-import argparse
-import pandas as pd
-import pickle
-
-
-if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- parser.add_argument('--model_dir', type=str, help='Directory for the models.')
- parser.add_argument('--model_numbers', type=int, nargs=2, help='The original and new model numbers.')
- args = vars(parser.parse_args())
-
- assert not os.path.isdir('%s/model_%d' % (args['model_dir'], args['model_numbers'][1])) # make sure the new model number is not already assigned to a model
-
- ### Change the model number in the model properties dictionary ###
- model_properties_file = '%s/model_%d/model_%d_properties.pkl' % (args['model_dir'], args['model_numbers'][0], args['model_numbers'][0])
- model_properties = pd.read_pickle(model_properties_file)
- model_properties['model_number'] = args['model_numbers'][1]
-
- with open(model_properties_file, 'wb') as f:
- pickle.dump(model_properties, f)
-
- os.rename('%s/model_%d' % (args['model_dir'], args['model_numbers'][0]), '%s/model_%d' % (args['model_dir'], args['model_numbers'][1])) # rename the model number directory
- files_to_rename = list(sorted(glob('%s/model_%d/**/*' % (args['model_dir'], args['model_numbers'][1]), recursive=True))) # files within the subdirectories to rename
-
- print("Renaming %d files" % len(files_to_rename))
- for file in files_to_rename:
- os.rename(file, file.replace(str(args['model_numbers'][0]), str(args['model_numbers'][1])))
-
- print("Successfully changed model number: %d -------> %d" % (args['model_numbers'][0], args['model_numbers'][1]))
diff --git a/utils/change_model_properties.py b/utils/change_model_properties.py
deleted file mode 100644
index 3eb8472..0000000
--- a/utils/change_model_properties.py
+++ /dev/null
@@ -1,54 +0,0 @@
-"""
-Changes values of dictionary keys in a model_properties.pkl file.
-This script is mainly used to address bugs in train_model.py, where the dictionaries are created.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.8.12
-"""
-import argparse
-import pandas as pd
-import pickle
-from utils.misc import string_arg_to_dict
-
-
-if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- parser.add_argument('--model_dir', type=str, help='Directory for the model.')
- parser.add_argument('--model_number', type=int, help='Model number.')
- parser.add_argument('--changes', type=str,
- help="Changes to make to the model properties dictionary. See utils.misc.string_arg_to_dict for more information.")
- parser.add_argument('--permission_override', action='store_true',
- help="WARNING: Read the description for this argument CAREFULLY. This is a boolean flag that overrides permission "
- "errors when attempting to modify critical model information. Changing properties that raise a PermissionError "
- "can render a model unusable with this module. ALWAYS create a backup of the model_*_properties.pkl file if "
- "you plan to modify critical model information.")
-
- args = vars(parser.parse_args())
-
- model_properties_file = '%s/model_%d/model_%d_properties.pkl' % (args['model_dir'], args['model_number'], args['model_number'])
- model_properties = pd.read_pickle(model_properties_file)
-
- changes = string_arg_to_dict(args['changes'])
-
- critical_args = ['dataset_properties', 'normalization_parameters', 'training_years', 'validation_years', 'test_years', 'model_number']
- critical_args_passed = list([arg for arg in critical_args if arg in changes])
-
- if len(critical_args_passed) > 0:
- if not args['permission_override']:
- raise PermissionError(
- f"The following critical model properties were attempted to be modified: --{', --'.join(critical_args_passed)}. "
- "Changing these properties can render the model properties file to be incompatible with other scripts. "
- "If you would like to modify these properties, pass the --permission_override flag. ALWAYS CREATE A BACKUP "
- "model_*_properties.pkl file before proceeding.")
-
- for arg in changes:
- model_properties[arg] = changes[arg]
-
- # Rewrite the human-readable model properties text file
- with open(model_properties_file.replace('.pkl', '.txt'), 'w') as f:
- for key in model_properties.keys():
- f.write(f"{key}: {model_properties[key]}\n")
-
- # Save the model properties dictionary with the new changes.
- with open(model_properties_file, 'wb') as f:
- pickle.dump(model_properties, f)
diff --git a/utils/data_utils.py b/utils/data_utils.py
deleted file mode 100644
index f8fee84..0000000
--- a/utils/data_utils.py
+++ /dev/null
@@ -1,703 +0,0 @@
-"""
-Various data tools.
-
-References
-----------
-* Snyder 1987: https://doi.org/10.3133/pp1395
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.9.18
-"""
-
-import pandas as pd
-from shapely.geometry import LineString
-import numpy as np
-import xarray as xr
-import tensorflow as tf
-
-
-# Each variable has parameters in the format of [max, min]
-normalization_parameters = {'mslp_z_surface': [1050., 960.],
- 'mslp_z_1000': [48., -69.],
- 'mslp_z_950': [86., -27.],
- 'mslp_z_900': [127., 17.],
- 'mslp_z_850': [174., 63.],
- 'q_surface': [24., 0.],
- 'q_1000': [26., 0.],
- 'q_950': [26., 0.],
- 'q_900': [23., 0.],
- 'q_850': [21., 0.],
- 'RH_surface': [1., 0.],
- 'RH_1000': [1., 0.],
- 'RH_950': [1., 0.],
- 'RH_900': [1., 0.],
- 'RH_850': [1., 0.],
- 'r_surface': [25., 0.],
- 'r_1000': [22., 0.],
- 'r_950': [22., 0.],
- 'r_900': [20., 0.],
- 'r_850': [18., 0.],
- 'sp_z_surface': [1075., 620.],
- 'sp_z_1000': [48., -69.],
- 'sp_z_950': [86., -27.],
- 'sp_z_900': [127., 17.],
- 'sp_z_850': [174., 63.],
- 'theta_surface': [331., 213.],
- 'theta_1000': [322., 218.],
- 'theta_950': [323., 219.],
- 'theta_900': [325., 227.],
- 'theta_850': [330., 237.],
- 'theta_e_surface': [375., 213.],
- 'theta_e_1000': [366., 208.],
- 'theta_e_950': [367., 210.],
- 'theta_e_900': [364., 227.],
- 'theta_e_850': [359., 238.],
- 'theta_v_surface': [324., 212.],
- 'theta_v_1000': [323., 218.],
- 'theta_v_950': [319., 215.],
- 'theta_v_900': [315., 220.],
- 'theta_v_850': [316., 227.],
- 'theta_w_surface': [304., 212.],
- 'theta_w_1000': [301., 207.],
- 'theta_w_950': [302., 210.],
- 'theta_w_900': [301., 214.],
- 'theta_w_850': [300., 237.],
- 'T_surface': [323., 212.],
- 'T_1000': [322., 218.],
- 'T_950': [319., 216.],
- 'T_900': [314., 220.],
- 'T_850': [315., 227.],
- 'Td_surface': [304., 207.],
- 'Td_1000': [302., 208.],
- 'Td_950': [301., 210.],
- 'Td_900': [298., 200.],
- 'Td_850': [296., 200.],
- 'Tv_surface': [324., 211.],
- 'Tv_1000': [323., 206.],
- 'Tv_950': [319., 206.],
- 'Tv_900': [316., 220.],
- 'Tv_850': [316., 227.],
- 'Tw_surface': [305., 212.],
- 'Tw_1000': [305., 218.],
- 'Tw_950': [304., 216.],
- 'Tw_900': [301., 219.],
- 'Tw_850': [299., 227.],
- 'u_surface': [36., -35.],
- 'u_1000': [38., -35.],
- 'u_950': [48., -55.],
- 'u_900': [59., -58.],
- 'u_850': [59., -58.],
- 'v_surface': [30., -35.],
- 'v_1000': [35., -38.],
- 'v_950': [55., -56.],
- 'v_900': [58., -59.],
- 'v_850': [58., -59.]}
-
-
-def expand_fronts(fronts: np.ndarray | tf.Tensor | xr.Dataset | xr.DataArray, iterations: int = 1):
- """
- Expands front labels in all directions.
-
- Parameters
- ----------
- fronts: array_like of ints of shape (T, M, N) or (M, N)
- 2-D or 3-D array of integers that identify the front type at each point. The longitude and latitude dimensions with
- shapes (M,) and (N,) can be in any order, but the time dimension must be the first dimension if it is passed.
- iterations: int
- Integer representing the number of times to expand the fronts in all directions.
-
- Returns
- -------
- fronts: array_like of ints of shape (T, M, N) or (1, M, N)
- Array of integers for the expanded fronts. If the array_like object passed into the function was 2-D, a third dimension
- will be added to the beginning of the array with size 1.
-
- Examples
- --------
- * Expanding labels for one front type.
- >>> arr = np.zeros((5, 5))
- >>> arr[2, 2] = 1 # add cold front point
- >>> arr
- array([[0., 0., 0., 0., 0.],
- [0., 0., 0., 0., 0.],
- [0., 0., 1., 0., 0.],
- [0., 0., 0., 0., 0.],
- [0., 0., 0., 0., 0.]])
- >>> expand_fronts(arr, iterations=1)
- array([[[0., 0., 0., 0., 0.],
- [0., 1., 1., 1., 0.],
- [0., 1., 1., 1., 0.],
- [0., 1., 1., 1., 0.],
- [0., 0., 0., 0., 0.]]])
-
- * Expanding labels for two front types.
- >>> arr = np.zeros((5, 5))
- >>> arr[1, 1] = 1 # add cold front point
- >>> arr[3, 3] = 2 # add warm front point
- >>> arr
- array([[0., 0., 0., 0., 0.],
- [0., 1., 0., 0., 0.],
- [0., 0., 0., 0., 0.],
- [0., 0., 0., 2., 0.],
- [0., 0., 0., 0., 0.]])
- >>> expand_fronts(arr, iterations=1)
- array([[[1., 1., 1., 0., 0.],
- [1., 1., 1., 0., 0.],
- [1., 1., 2., 2., 2.],
- [0., 0., 2., 2., 2.],
- [0., 0., 2., 2., 2.]]])
- """
- if type(fronts) in [xr.Dataset, xr.DataArray]:
- identifier = fronts['identifier'].values if type(fronts) == xr.Dataset else fronts.values
- elif tf.is_tensor(fronts):
- identifier = tf.expand_dims(fronts, axis=0) if len(fronts.shape) == 2 else fronts
- else:
- identifier = np.expand_dims(fronts, axis=0) if len(fronts.shape) == 2 else fronts
-
- if tf.is_tensor(identifier):
- for _ in range(iterations):
- # 8 tensors representing all directions for the front expansion
- identifier_up_left = tf.Variable(tf.zeros_like(identifier))
- identifier_up_right = tf.Variable(tf.zeros_like(identifier))
- identifier_down_left = tf.Variable(tf.zeros_like(identifier))
- identifier_down_right = tf.Variable(tf.zeros_like(identifier))
- identifier_up = tf.Variable(tf.zeros_like(identifier))
- identifier_down = tf.Variable(tf.zeros_like(identifier))
- identifier_left = tf.Variable(tf.zeros_like(identifier))
- identifier_right = tf.Variable(tf.zeros_like(identifier))
-
- identifier_down_left[:, 1:, :-1].assign(tf.where((identifier[:, :-1, 1:] > 0) & (identifier[:, 1:, :-1] == 0),
- identifier[:, :-1, 1:], identifier[:, 1:, :-1]))
- identifier_down[:, 1:, :].assign(tf.where((identifier[:, :-1, :] > 0) & (identifier[:, 1:, :] == 0),
- identifier[:, :-1, :], identifier[:, 1:, :]))
- identifier_down_right[:, 1:, 1:].assign(tf.where((identifier[:, :-1, :-1] > 0) & (identifier[:, 1:, 1:] == 0),
- identifier[:, :-1, :-1], identifier[:, 1:, 1:]))
- identifier_up_left[:, :-1, :-1].assign(tf.where((identifier[:, 1:, 1:] > 0) & (identifier[:, :-1, :-1] == 0),
- identifier[:, 1:, 1:], identifier[:, :-1, :-1]))
- identifier_up[:, :-1, :].assign(tf.where((identifier[:, 1:, :] > 0) & (identifier[:, :-1, :] == 0),
- identifier[:, 1:, :], identifier[:, :-1, :]))
- identifier_up_right[:, :-1, 1:].assign(tf.where((identifier[:, 1:, :-1] > 0) & (identifier[:, :-1, 1:] == 0),
- identifier[:, 1:, :-1], identifier[:, :-1, 1:]))
- identifier_left[:, :, :-1].assign(tf.where((identifier[:, :, 1:] > 0) & (identifier[:, :, :-1] == 0),
- identifier[:, :, 1:], identifier[:, :, :-1]))
- identifier_right[:, :, 1:].assign(tf.where((identifier[:, :, :-1] > 0) & (identifier[:, :, 1:] == 0),
- identifier[:, :, :-1], identifier[:, :, 1:]))
-
- identifier = tf.reduce_max([identifier_up_left, identifier_up, identifier_up_right,
- identifier_down_left, identifier_down, identifier_down_right,
- identifier_left, identifier_right], axis=0)
-
- else:
- for _ in range(iterations):
- # 8 arrays representing all directions for the front expansion
- identifier_up_left = np.zeros_like(identifier)
- identifier_up_right = np.zeros_like(identifier)
- identifier_down_left = np.zeros_like(identifier)
- identifier_down_right = np.zeros_like(identifier)
- identifier_up = np.zeros_like(identifier)
- identifier_down = np.zeros_like(identifier)
- identifier_left = np.zeros_like(identifier)
- identifier_right = np.zeros_like(identifier)
-
- identifier_down_left[:, 1:, :-1] = np.where((identifier[:, :-1, 1:] > 0) & (identifier[:, 1:, :-1] == 0),
- identifier[:, :-1, 1:], identifier[:, 1:, :-1]) # down-left
- identifier_down[:, 1:, :] = np.where((identifier[:, :-1, :] > 0) & (identifier[:, 1:, :] == 0),
- identifier[:, :-1, :], identifier[:, 1:, :]) # down
- identifier_down_right[:, 1:, 1:] = np.where((identifier[:, :-1, :-1] > 0) & (identifier[:, 1:, 1:] == 0),
- identifier[:, :-1, :-1], identifier[:, 1:, 1:]) # down-right
- identifier_up_left[:, :-1, :-1] = np.where((identifier[:, 1:, 1:] > 0) & (identifier[:, :-1, :-1] == 0),
- identifier[:, 1:, 1:], identifier[:, :-1, :-1]) # up-left
- identifier_up[:, :-1, :] = np.where((identifier[:, 1:, :] > 0) & (identifier[:, :-1, :] == 0),
- identifier[:, 1:, :], identifier[:, :-1, :]) # up
- identifier_up_right[:, :-1, 1:] = np.where((identifier[:, 1:, :-1] > 0) & (identifier[:, :-1, 1:] == 0),
- identifier[:, 1:, :-1], identifier[:, :-1, 1:]) # up-right
- identifier_left[:, :, :-1] = np.where((identifier[:, :, 1:] > 0) & (identifier[:, :, :-1] == 0),
- identifier[:, :, 1:], identifier[:, :, :-1]) # left
- identifier_right[:, :, 1:] = np.where((identifier[:, :, :-1] > 0) & (identifier[:, :, 1:] == 0),
- identifier[:, :, :-1], identifier[:, :, 1:]) # right
-
- identifier = np.max([identifier_up_left, identifier_up, identifier_up_right,
- identifier_down_left, identifier_down, identifier_down_right,
- identifier_left, identifier_right], axis=0)
-
- if type(fronts) == xr.Dataset:
- fronts['identifier'].values = identifier
- elif type(fronts) == xr.DataArray:
- fronts.values = identifier
- else:
- fronts = identifier
-
- return fronts
-
-
-def haversine(lon: np.ndarray | int | float,
- lat: np.ndarray | int | float):
- """
- Haversine formula. Transforms lon/lat points to an x/y cartesian plane.
-
- Parameters
- ----------
- lon: array_like of shape (N,), int, or float
- Longitude component of the point(s) expressed in degrees.
- lat: array_like of shape (N,), int, or float
- Latitude component of the point(s) expressed in degrees.
-
- Returns
- -------
- x: array_like of shape (N,) or float
- X component of the transformed points expressed in kilometers.
- y: array_like of shape (N,) or float
- Y component of the transformed points expressed in kilometers.
-
- Examples
- --------
- >>> lon = -95
- >>> lat = 35
- >>> x, y = haversine(lon, lat)
- >>> x, y
- (-10077.330945462296, 3892.875)
-
- >>> lon = np.arange(10, 80.1, 10)
- >>> lat = np.arange(10, 80.1, 10)
- >>> x, y = haversine(lon, lat)
- >>> x, y
- (array([1108.01755295, 2190.70484658, 3223.05300087, 4180.69246988,
- 5040.20418066, 5779.42053216, 6377.71302882, 6816.26345487]), array([1112.25, 2224.5 , 3336.75, 4449. , 5561.25, 6673.5 , 7785.75,
- 8898. ]))
- """
- C = 40041 # average circumference of earth in kilometers
- x = lon * C * np.cos(lat * np.pi / 360) / 360
- y = lat * C / 360
- return x, y
-
-
-def reverse_haversine(x, y):
- """
- Reverse haversine formula. Transforms x/y cartesian coordinates to a lon/lat grid.
-
- Parameters
- ----------
- x: array_like of shape (N,), int, or float
- X component of the point(s) expressed in kilometers.
- y: array_like of shape (N,), int, or float
- Y component of the point(s) expressed in kilometers.
-
- Returns
- -------
- lon: array_like of shape (N,) or float
- Longitude component of the transformed point(s) expressed in degrees.
- lat: array_like of shape (N,) or float
- Latitude component of the transformed point(s) expressed in degrees.
-
- Examples
- --------
- Values pulled from haversine examples.
-
- >>> x = -10077.330945462296
- >>> y = 3892.875
- >>> lon, lat = reverse_haversine(x, y)
- >>> lon, lat
- (-95.0, 35.0)
-
- >>> x = np.array([1108.01755295, 2190.70484658, 3223.05300087, 4180.69246988, 5040.20418066, 5779.42053216, 6377.71302882, 6816.26345487])
- >>> y = np.array([1112.25, 2224.5, 3336.75, 4449., 5561.25, 6673.5, 7785.75, 8898.])
- >>> lon, lat = reverse_haversine(x, y)
- >>> lon, lat
- (array([10., 20., 30., 40., 50., 60., 70., 80.]), array([10., 20., 30., 40., 50., 60., 70., 80.]))
- """
- C = 40041 # average circumference of earth in kilometers
- lon = x * 360 / np.cos(y * np.pi / C) / C
- lat = y * 360 / C
- return lon, lat
-
-
-def geometric(x_km_new, y_km_new):
- """
- Turn longitudinal/latitudinal distance (km) lists into LineString for interpolation.
-
- Parameters
- ----------
- x_km_new: List containing longitude coordinates of fronts in kilometers.
- y_km_new: List containing latitude coordinates of fronts in kilometers.
-
- Returns
- -------
- xy_linestring: LineString object containing coordinates of fronts in kilometers.
- """
- df_xy = pd.DataFrame(list(zip(x_km_new, y_km_new)), columns=['Longitude_km', 'Latitude_km'])
- geometry = [xy for xy in zip(df_xy.Longitude_km, df_xy.Latitude_km)]
- xy_linestring = LineString(geometry)
- return xy_linestring
-
-
-def redistribute_vertices(xy_linestring, distance):
- """
- Interpolate x/y coordinates at a specified distance.
-
- Parameters
- ----------
- xy_linestring: LineString object containing coordinates of fronts in kilometers.
- distance: Distance at which to interpolate the x/y coordinates.
-
- Returns
- -------
- xy_vertices: Normalized MultiLineString that contains the interpolated coordinates of fronts in kilometers.
-
- Sources
- -------
- https://stackoverflow.com/questions/34906124/interpolating-every-x-distance-along-multiline-in-shapely/35025274#35025274
- """
- if xy_linestring.geom_type == 'LineString':
- num_vert = int(round(xy_linestring.length / distance))
- if num_vert == 0:
- num_vert = 1
- return LineString(
- [xy_linestring.interpolate(float(n) / num_vert, normalized=True)
- for n in range(num_vert + 1)])
- elif xy_linestring.geom_type == 'MultiLineString':
- parts = [redistribute_vertices(part, distance) for part in xy_linestring]
- return type(xy_linestring)([p for p in parts if not p.is_empty])
- else:
- raise ValueError('unhandled geometry %s', (xy_linestring.geom_type,))
-
-
-def reformat_fronts(fronts, front_types):
- """
- Reformat a front dataset, tensor, or array with a given set of front types.
-
- Parameters
- ----------
- front_types: str or list of strs
- Code(s) that determine how the dataset will be reformatted.
- fronts: xarray Dataset or DataArray, tensor, or np.ndarray
- Dataset containing the front data.
- '''
- Available options for individual front types (cannot be passed with any special codes):
-
- Code (class #): Front Type
- --------------------------
- CF (1): Cold front
- WF (2): Warm front
- SF (3): Stationary front
- OF (4): Occluded front
- CF-F (5): Cold front (forming)
- WF-F (6): Warm front (forming)
- SF-F (7): Stationary front (forming)
- OF-F (8): Occluded front (forming)
- CF-D (9): Cold front (dissipating)
- WF-D (10): Warm front (dissipating)
- SF-D (11): Stationary front (dissipating)
- OF-D (12): Occluded front (dissipating)
- INST (13): Instability axis
- TROF (14): Trough
- TT (15): Tropical Trough
- DL (16): Dryline
-
-
- Special codes (cannot be passed with any individual front codes):
- -----------------------------------------------------------------
- F_BIN (1 class): 1-4, but treat all front types as one type.
- (1): CF, WF, SF, OF
-
- MERGED-F (4 classes): 1-12, but treat forming and dissipating fronts as standard fronts.
- (1): CF, CF-F, CF-D
- (2): WF, WF-F, WF-D
- (3): SF, SF-F, SF-D
- (4): OF, OF-F, OF-D
-
- MERGED-F_BIN (1 class): 1-12, but treat all front types and stages as one type. This means that classes 1-12 will all be one class (1).
- (1): CF, CF-F, CF-D, WF, WF-F, WF-D, SF, SF-F, SF-D, OF, OF-F, OF-D
-
- MERGED-T (1 class): 14-15, but treat troughs and tropical troughs as the same. In other words, TT (15) becomes TROF (14).
- (1): TROF, TT
-
- MERGED-ALL (7 classes): 1-16, but make the changes in the MERGED-F and MERGED-T codes.
- (1): CF, CF-F, CF-D
- (2): WF, WF-F, WF-D
- (3): SF, SF-F, SF-D
- (4): OF, OF-F, OF-D
- (5): TROF, TT
- (6): INST
- (7): DL
-
- **** NOTE - Class 0 is always treated as 'no front'.
- '''
-
- Returns
- -------
- fronts_ds: xr.Dataset
- Reformatted dataset based on the provided code(s).
- """
-
- if type(front_types) == str:
- front_types = [front_types, ]
-
- fronts_argument_type = type(fronts)
-
- if fronts_argument_type == xr.DataArray or fronts_argument_type == xr.Dataset:
- where_function = xr.where
- elif fronts_argument_type == np.ndarray:
- where_function = np.where
- else:
- where_function = tf.where
-
- front_types_classes = {'CF': 1, 'WF': 2, 'SF': 3, 'OF': 4, 'CF-F': 5, 'WF-F': 6, 'SF-F': 7, 'OF-F': 8, 'CF-D': 9, 'WF-D': 10,
- 'SF-D': 11, 'OF-D': 12, 'INST': 13, 'TROF': 14, 'TT': 15, 'DL': 16}
-
- if front_types == ['F_BIN', ]:
-
- fronts = where_function(fronts > 4, 0, fronts) # Classes 5-16 are removed
- fronts = where_function(fronts > 0, 1, fronts) # Merge 1-4 into one class
-
- labels = ['CF-WF-SF-OF', ]
- num_types = 1
-
- elif front_types == ['MERGED-F']:
-
- fronts = where_function(fronts == 5, 1, fronts) # Forming cold front ---> cold front
- fronts = where_function(fronts == 6, 2, fronts) # Forming warm front ---> warm front
- fronts = where_function(fronts == 7, 3, fronts) # Forming stationary front ---> stationary front
- fronts = where_function(fronts == 8, 4, fronts) # Forming occluded front ---> occluded front
- fronts = where_function(fronts == 9, 1, fronts) # Dying cold front ---> cold front
- fronts = where_function(fronts == 10, 2, fronts) # Dying warm front ---> warm front
- fronts = where_function(fronts == 11, 3, fronts) # Dying stationary front ---> stationary front
- fronts = where_function(fronts == 12, 4, fronts) # Dying occluded front ---> occluded front
- fronts = where_function(fronts > 4, 0, fronts) # Remove all other fronts
-
- labels = ['CF_any', 'WF_any', 'SF_any', 'OF_any']
- num_types = 4
-
- elif front_types == ['MERGED-F_BIN']:
-
- fronts = where_function(fronts > 12, 0, fronts) # Classes 13-16 are removed
- fronts = where_function(fronts > 0, 1, fronts) # Classes 1-12 are merged into one class
-
- labels = ['CF-WF-SF-OF_any', ]
- num_types = 1
-
- elif front_types == ['MERGED-T']:
-
- fronts = where_function(fronts < 14, 0, fronts) # Remove classes 1-13
-
- # Merge troughs into one class
- fronts = where_function(fronts == 14, 1, fronts)
- fronts = where_function(fronts == 15, 1, fronts)
-
- fronts = where_function(fronts == 16, 0, fronts) # Remove drylines
-
- labels = ['TR_any', ]
- num_types = 1
-
- elif front_types == ['MERGED-ALL']:
-
- fronts = where_function(fronts == 5, 1, fronts) # Forming cold front ---> cold front
- fronts = where_function(fronts == 6, 2, fronts) # Forming warm front ---> warm front
- fronts = where_function(fronts == 7, 3, fronts) # Forming stationary front ---> stationary front
- fronts = where_function(fronts == 8, 4, fronts) # Forming occluded front ---> occluded front
- fronts = where_function(fronts == 9, 1, fronts) # Dying cold front ---> cold front
- fronts = where_function(fronts == 10, 2, fronts) # Dying warm front ---> warm front
- fronts = where_function(fronts == 11, 3, fronts) # Dying stationary front ---> stationary front
- fronts = where_function(fronts == 12, 4, fronts) # Dying occluded front ---> occluded front
-
- # Merge troughs together into class 5
- fronts = where_function(fronts == 14, 5, fronts)
- fronts = where_function(fronts == 15, 5, fronts)
-
- fronts = where_function(fronts == 13, 6, fronts) # Move outflow boundaries to class 6
- fronts = where_function(fronts == 16, 7, fronts) # Move drylines to class 7
-
- labels = ['CF_any', 'WF_any', 'SF_any', 'OF_any', 'TR_any', 'INST', 'DL']
- num_types = 7
-
- else:
-
- # Select the front types that are being used to pull their class identifiers
- filtered_front_types = dict(sorted(dict([(i, front_types_classes[i]) for i in front_types_classes if i in set(front_types)]).items(), key=lambda item: item[1]))
- front_types, num_types = list(filtered_front_types.keys()), len(filtered_front_types.keys())
-
- for i in range(num_types):
- if i + 1 != front_types_classes[front_types[i]]:
- fronts = where_function(fronts == i + 1, 0, fronts)
- fronts = where_function(fronts == front_types_classes[front_types[i]], i + 1, fronts) # Reformat front classes
-
- fronts = where_function(fronts > num_types, 0, fronts) # Remove unused front types
-
- labels = front_types
-
- if fronts_argument_type == xr.Dataset or fronts_argument_type == xr.DataArray:
- fronts.attrs['front_types'] = front_types
- fronts.attrs['num_types'] = num_types
- fronts.attrs['labels'] = labels
-
- return fronts
-
-
-def normalize_variables(variable_ds, normalization_parameters=normalization_parameters):
- """
- Function that normalizes thermodynamic variables via min-max normalization.
-
- Parameters
- ----------
- variable_ds: xr.Dataset
- Dataset containing thermodynamic variable data.
- normalization_parameters: dict
- Dictionary containing parameters for normalization.
-
- Returns
- -------
- variable_ds: xr.Dataset
- Same as input dataset, but the variables are normalized via min-max normalization.
- """
-
- # Place pressure levels as the last dimension of the dataset
- original_dim_order = variable_ds.dims
- variable_ds = variable_ds.transpose(*[dim for dim in original_dim_order if dim != 'pressure_level'], 'pressure_level').astype('float64')
-
- variable_list = list(variable_ds.keys())
- pressure_levels = variable_ds['pressure_level'].values
-
- for var in variable_list:
-
- current_variable_values = variable_ds[var].values
- new_variable_values = np.zeros_like(current_variable_values)
-
- for idx, pressure_level in enumerate(pressure_levels):
- norm_var = '_'.join([var, pressure_level]) # name of the variable as it appears in the normalization parameters dictionary
- max_val, min_val = normalization_parameters[norm_var]
- new_variable_values[..., idx] = np.nan_to_num((variable_ds[var].values[..., idx] - min_val) / (max_val - min_val))
-
- variable_ds[var].values = new_variable_values # assign new values for variable
-
- variable_ds = variable_ds.transpose(*[dim for dim in original_dim_order])
- return variable_ds
-
-
-def randomize_variables(variable_ds: xr.Dataset, random_variables: list or tuple):
- """
- Scramble the values of specific variables within a given dataset.
-
- Parameters
- ----------
- variable_ds: xr.Dataset
- Xarray dataset containing thermodynamic variable data.
- random_variables: list or tuple
- List of variables to randomize the values of.
-
- Returns
- -------
- variable_ds: xr.Dataset
- Same as input, but with the given variables having scrambled values.
- """
-
- for random_variable in random_variables:
- variable_values = variable_ds[random_variable].values # DataArray of the current variable
- variable_shape = np.shape(variable_values)
- flattened_variable_values = variable_values.flatten()
- np.random.shuffle(flattened_variable_values)
- variable_ds[random_variable].values = np.reshape(flattened_variable_values, variable_shape)
-
- return variable_ds
-
-
-def combine_datasets(input_files: list[str],
- label_files: list[str] = None):
- """
- Combine many tensorflow datasets into one entire dataset.
-
- Returns
- -------
- complete_dataset: tf.data.Dataset object
- Concatenated tensorflow dataset.
- """
- inputs = tf.data.Dataset.load(input_files[0])
-
- if label_files is not None:
-
- labels = tf.data.Dataset.load(label_files[0])
- for input_file, label_file in zip(input_files[1:], label_files[1:]):
- inputs = inputs.concatenate(tf.data.Dataset.load(input_file))
- labels = labels.concatenate(tf.data.Dataset.load(label_file))
-
- return tf.data.Dataset.zip((inputs, labels))
-
- else:
-
- for input_file in input_files[1:]:
- inputs = inputs.concatenate(tf.data.Dataset.load(input_file))
-
- return inputs
-
-
-def lambert_conformal_to_cartesian(
- lon: np.ndarray | tuple | list | int | float,
- lat: np.ndarray | tuple | list | int | float,
- std_parallels: tuple | list = (20., 50.),
- lon_ref: int | float = 0.,
- lat_ref: int | float = 0.):
- """
- Transform points on a Lambert Conformal lat/lon grid to cartesian coordinates.
-
- Parameters
- ----------
- lon: array_like of shape (N,), int, or float
- Longitude point(s) expressed as degrees.
- lat: array_like of shape (N,), int, or float
- Latitude point(s) expressed as degrees.
- std_parallels: tuple or list of 2 ints or floats
- Standard parallels to use in the coordinate transformation, expressed as degrees.
- lon_ref: int or float
- Reference longitude point expressed as degrees.
- lat_ref: int or float
- Reference latitude point expressed as degrees.
-
- Returns
- -------
- x: array_like of shape (N,) or float
- X-component of the transformed coordinates, expressed as meters.
- y: array_like of shape (N,) or float
- Y-component of the transformed coordinates, expressed as meters.
-
- Examples
- --------
- * Using parameters from example on Page 295 of Snyder 1987 (except the output here is expressed as meters):
- >>> x, y = lambert_conformal_to_cartesian(lon=-75, lat=35, std_parallels=(33, 45), lon_ref=-96, lat_ref=23)
- >>> x, y
- (1890206.4076610378, 1568668.1244433122)
-
- * Same as above but with longitudes expressed from 0 to 360 degrees east:
- >>> x, y = lambert_conformal_to_cartesian(lon=285, lat=35, std_parallels=(33, 45), lon_ref=264, lat_ref=23)
- >>> x, y
- (1890206.4076610343, 1568668.1244433112)
-
- References
- ----------
- * Snyder 1987: https://doi.org/10.3133/pp1395
-
- Notes
- -----
- lon and lon_ref must be both expressed in the same longitude range (e.g. -180 to 180 degrees or 0 to 360 degrees)
- to get correct values for x and y.
- """
-
- R = 6371229 # radius of earth (meters)
-
- # Points and standard parallels need to be expressed as radians for the transformation formulas
- lon = np.radians(lon)
- lon_ref = np.radians(lon_ref)
- lat = np.radians(lat)
- lat_ref = np.radians(lat_ref)
- std_parallels = np.radians(std_parallels)
-
- if std_parallels[0] == std_parallels[1]:
- n = np.sin(std_parallels[0])
- else:
- n = np.divide(np.log(np.cos(std_parallels[0]) / np.cos(std_parallels[1])),
- np.log(np.tan(np.pi/4 + std_parallels[1]/2) / np.tan(np.pi/4 + std_parallels[0]/2)))
- F = np.cos(std_parallels[0]) * np.power(np.tan(np.pi/4 + std_parallels[0]/2), n) / n
- rho = R * F / np.power(np.tan(np.pi/4 + lat/2), n)
- rho0 = R * F / np.power(np.tan(np.pi/4 + lat_ref/2), n)
-
- x = rho * np.sin(n * (lon - lon_ref))
- y = rho0 - rho * np.cos(n * (lon - lon_ref))
-
- return x, y
diff --git a/utils/misc.py b/utils/misc.py
deleted file mode 100644
index 70e15e3..0000000
--- a/utils/misc.py
+++ /dev/null
@@ -1,104 +0,0 @@
-"""
-Miscellaneous tools.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.7.7.D1
-"""
-
-
-def string_arg_to_dict(arg_str: str):
- """
- Function that converts a string argument into a dictionary. Dictionaries cannot be passed through a command line, so
- this function takes a special string argument and converts it to a dictionary so arguments within a function can be
- explicitly called.
-
- arg_str: str
- arg_dict: dict
- """
-
- arg_str = arg_str.replace(' ', '') # Remove all spaces from the string.
- arg_dict = {} # Dictionary that will contain the arguments and their respective values
-
- # Iterate through all the arguments within the string
- while True:
-
- equals_index = arg_str.find('=') # Index representing where an equals sign is located, marking the end of the argument name
-
- ################################# Attempt to see if a tuple or list was passed #################################
- open_parenthesis_index = arg_str.find('(')
- close_parenthesis_index = arg_str.find(')')
- open_bracket_index = arg_str.find('[')
- close_bracket_index = arg_str.find(']')
-
- if open_parenthesis_index == close_parenthesis_index and open_bracket_index == close_bracket_index: # These will only be equal when there are no parentheses/brackets in the argument string (i.e. there is no tuple/list)
- comma_index = arg_str.find(',') # Index representing where the first comma is located within the 'arg' string, essentially representing the end of a argument
- elif open_parenthesis_index == -1 and close_parenthesis_index != -1:
- raise TypeError("An open parenthesis appears to be missing. Check the argument string.")
- elif open_parenthesis_index != -1 and close_parenthesis_index == -1:
- raise TypeError("A closed parenthesis appears to be missing. Check the argument string.")
- elif open_bracket_index == -1 and close_bracket_index != -1:
- raise TypeError("An open bracket appears to be missing. Check the argument string.")
- elif open_bracket_index != -1 and close_bracket_index == -1:
- raise TypeError("A closed bracket appears to be missing. Check the argument string.")
- elif open_parenthesis_index != close_parenthesis_index:
- comma_index = close_parenthesis_index + 1
- else:
- comma_index = close_bracket_index + 1
-
- current_arg_name = arg_str[:equals_index]
-
- if comma_index == -1: # When the final argument is being added to the dictionary, this index will become -1
- current_arg_value = arg_str[equals_index + 1:]
- else:
- current_arg_value = arg_str[equals_index + 1:comma_index]
-
- ######################################## Convert the argument to a tuple #######################################
- if open_parenthesis_index != close_parenthesis_index:
-
- arg_dict[current_arg_name] = current_arg_value.replace('(', '').replace(')', '').split(',')
-
- if '.' in arg_dict[current_arg_name]: # If the tuple appears to contain a float
- arg_dict[current_arg_name] = tuple([float(val) for val in arg_dict[current_arg_name]])
- else:
- arg_dict[current_arg_name] = tuple([int(val) for val in arg_dict[current_arg_name]])
- ################################################################################################################
-
- ######################################## Convert the argument to a list ########################################
- elif open_bracket_index != close_bracket_index:
-
- arg_dict[current_arg_name] = current_arg_value.replace('[', '').replace(']', '').split(',')
-
- list_values = []
- for val in arg_dict[current_arg_name]:
- if '.' in val:
- list_values.append(float(val))
- else:
- try:
- list_values.append(int(val))
- except ValueError:
- list_values.append(val)
-
- arg_dict[current_arg_name] = list_values
- ################################################################################################################
-
- else:
-
- if '.' in current_arg_value:
- arg_dict[current_arg_name] = float(current_arg_value)
- else:
- try:
- arg_dict[current_arg_name] = int(current_arg_value)
- except ValueError:
- if current_arg_value == 'True':
- arg_dict[current_arg_name] = True
- elif current_arg_value == 'False':
- arg_dict[current_arg_name] = False
- else:
- arg_dict[current_arg_name] = current_arg_value.replace("'", '')
-
- arg_str = arg_str[comma_index + 1:] # After the current argument has been added to the dictionary, remove it from the argument string
-
- if comma_index == -1 or len(arg_str) == 0:
- break
-
- return arg_dict
diff --git a/utils/plotting_utils.py b/utils/plotting_utils.py
deleted file mode 100644
index c82bb25..0000000
--- a/utils/plotting_utils.py
+++ /dev/null
@@ -1,69 +0,0 @@
-"""
-Plotting tools.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.9.15
-"""
-
-import cartopy.feature as cfeature
-import cartopy.crs as ccrs
-import matplotlib.pyplot as plt
-import matplotlib as mpl
-import numpy as np
-
-
-def plot_background(extent, ax=None, linewidth=0.5):
- """
- Returns new background for the plot.
-
- Parameters
- ----------
- extent: Iterable object with 4 integers
- Iterable containing the extent/boundaries of the plot in the format of [min lon, max lon, min lat, max lat] expressed
- in degrees.
- ax: matplotlib.axes.Axes instance or None
- Axis on which the background will be plotted.
- linewidth: float or int
- Thickness of coastlines and the borders of states and countries.
-
- Returns
- -------
- ax: matplotlib.axes.Axes instance
- New plot background.
- """
- if ax is None:
- crs = ccrs.LambertConformal(central_longitude=250)
- ax = plt.axes(projection=crs)
- else:
- ax.add_feature(cfeature.COASTLINE.with_scale('50m'), linewidth=linewidth)
- ax.add_feature(cfeature.BORDERS, linewidth=linewidth)
- ax.add_feature(cfeature.STATES, linewidth=linewidth)
- ax.set_extent(extent, crs=ccrs.PlateCarree())
- return ax
-
-
-def truncated_colormap(cmap, minval=0.0, maxval=1.0, n=100):
- """
- Get an instance of a truncated matplotlib.colors.Colormap object.
-
- Parameters
- ----------
- cmap: str
- Matplotlib colormap to truncate.
- minval: float
- Starting point of the colormap, represented by a float of 0 <= minval < 1.
- maxval: float
- End point of the colormap, represented by a float of 0 < maxval <= 1.
- n: int
- Number of colors for the colormap.
-
- Returns
- -------
- new_cmap: matplotlib.colors.Colormap instance
- Truncated colormap.
- """
- cmap = plt.get_cmap(cmap)
- new_cmap = mpl.colors.LinearSegmentedColormap.from_list(
- 'trunc({n},{a:.2f},{b:.2f})'.format(n=cmap.name, a=minval, b=maxval),
- cmap(np.linspace(minval, maxval, n)))
- return new_cmap
diff --git a/utils/settings.py b/utils/settings.py
deleted file mode 100644
index cb396b3..0000000
--- a/utils/settings.py
+++ /dev/null
@@ -1,83 +0,0 @@
-"""
-Default settings
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.7.24
-"""
-DEFAULT_DOMAIN_EXTENTS = {'global': [0, 359.75, -89.75, 90],
- 'full': [130, 369.75, 0.25, 80],
- 'conus': [228, 299.75, 25, 56.75]} # default values for extents of domains [start lon, end lon, start lat, end lat]
-DEFAULT_DOMAIN_INDICES = {'global': [0, 1440, 0, 720],
- 'full': [0, 960, 0, 320],
- 'conus': [392, 680, 93, 221]} # indices corresponding to default extents of domains [start lon, end lon, start lat, end lat]
-DEFAULT_DOMAIN_IMAGES = {'global': [17, 9],
- 'full': [8, 3],
- 'conus': [3, 1]} # default values for the number of images to use when making predictions [lon, lat]
-
-# colors for plotted ground truth fronts
-DEFAULT_FRONT_COLORS = {'CF': 'blue', 'WF': 'red', 'SF': 'limegreen', 'OF': 'darkviolet', 'CF-F': 'darkblue', 'WF-F': 'darkred',
- 'SF-F': 'darkgreen', 'OF-F': 'darkmagenta', 'CF-D': 'lightskyblue', 'WF-D': 'lightcoral', 'SF-D': 'lightgreen',
- 'OF-D': 'violet', 'INST': 'gold', 'TROF': 'goldenrod', 'TT': 'orange', 'DL': 'chocolate',
- 'MERGED-CF': 'blue', 'MERGED-WF': 'red', 'MERGED-SF': 'limegreen', 'MERGED-OF': 'darkviolet', 'MERGED-F': 'gray',
- 'MERGED-T': 'brown', 'F_BIN': 'tab:red', 'MERGED-F_BIN': 'tab:red'}
-
-# colormaps of probability contours for front predictions
-DEFAULT_CONTOUR_CMAPS = {'CF': 'Blues', 'WF': 'Reds', 'SF': 'Greens', 'OF': 'Purples', 'CF-F': 'Blues', 'WF-F': 'Reds',
- 'SF-F': 'Greens', 'OF-F': 'Purples', 'CF-D': 'Blues', 'WF-D': 'Reds', 'SF-D': 'Greens', 'OF-D': 'Purples',
- 'INST': 'YlOrBr', 'TROF': 'YlOrRd', 'TT': 'Oranges', 'DL': 'copper_r', 'MERGED-CF': 'Blues',
- 'MERGED-WF': 'Reds', 'MERGED-SF': 'Greens', 'MERGED-OF': 'Purples', 'MERGED-F': 'Greys', 'MERGED-T': 'YlOrBr',
- 'F_BIN': 'Greys', 'MERGED-F_BIN': 'Greys'}
-
-# names of front types
-DEFAULT_FRONT_NAMES = {'CF': 'Cold front', 'WF': 'Warm front', 'SF': 'Stationary front', 'OF': 'Occluded front', 'CF-F': 'Cold front (forming)',
- 'WF-F': 'Warm front (forming)', 'SF-F': 'Stationary front (forming)', 'OF-F': 'Occluded front (forming)',
- 'CF-D': 'Cold front (dying)', 'WF-D': 'Warm front (dying)', 'SF-D': 'Stationary front (dying)', 'OF-D': 'Occluded front (dying)',
- 'INST': 'Outflow boundary', 'TROF': 'Trough', 'TT': 'Tropical trough', 'DL': 'Dryline',
- 'MERGED-CF': 'Cold front (any)', 'MERGED-WF': 'Warm front (any)', 'MERGED-SF': 'Stationary front (any)', 'MERGED-OF': 'Occluded front (any)',
- 'MERGED-F': 'CF, WF, SF, OF (any)', 'MERGED-T': 'Trough (any)', 'F_BIN': 'Binary front', 'MERGED-F_BIN': 'Binary front (any)'}
-
-"""
-TIMESTEP_PREDICT_SIZE is the number of timesteps for which predictions will be processed at the same time. In other words,
- if this parameter is 10, then up to 10 maps will be generated at the same time for 10 timesteps.
-
-Typically, raising this parameter will result in faster predictions, but the memory requirements increase as well. The size
- of the domain for which the predictions are being generated greatly affects the limits of this parameter.
-
-NOTES:
- - Setting the values at a lower threshold may result in slower predictions but will not have negative effects on hardware.
- - Increasing the parameters above the default values is STRONGLY discouraged. Greatly exceeding the allocated RAM
- will force your operating system to resort to virtual memory usage, which can cause major slowdowns, OS crashes, and GPU failure.
-
-In the case of any hardware failures (including OOM errors from the GPU) or major slowdowns, reduce the parameter(s) for the
- domain(s) until your system becomes stable.
-"""
-TIMESTEP_PREDICT_SIZE = {'conus': 128, 'full': 64, 'global': 16} # All values here are adjusted for 16 GB of system RAM and 10 GB of GPU VRAM
-
-"""
-GPU_PREDICT_BATCH_SIZE is the number of images that the GPU will process at one time when generating predictions.
-If predictions are being generated on the same GPU that the model was trained on, then this value should be equal to or greater than
-the batch size used when training the model.
-
-NOTES:
- - This value should ideally be a value of 2^n, where n is any integer. Using values not equal to 2^n may have negative effects
- on performance.
- - Decreasing this parameter will result in overall slower performance, but can help prevent OOM errors on the GPU.
- - Increasing this parameter can speed up predictions on high-performance GPUs, but a value too large can cause OOM errors
- and GPU failure.
-"""
-GPU_PREDICT_BATCH_SIZE = 8
-
-"""
-MAX_FILE_CHUNK_SIZE is the maximum number of ERA5, GDAS, and/or GFS netCDF files that will be loaded into an xarray dataset at one
-time. Loading too many files / too much data into one xarray dataset can take a very long time and may lead to segmentation errors.
-If segmentation errors are occurring, consider lowering this parameter until the error disappears.
-"""
-MAX_FILE_CHUNK_SIZE = 2500
-
-"""
-MAX_TRAIN_BUFFER_SIZE is the maximum number of elements within the training dataset that can be shuffled at one time. Tensorflow
-does not efficiently use RAM during shuffling on Windows machines and can lead to system crashes, so the buffer size should be
-relatively small. It is important to monitor RAM usage if you are training a model on Windows. Linux is able to shuffle much
-larger datasets than Windows, but crashes can still occur if the maximum buffer size is too large.
-"""
-MAX_TRAIN_BUFFER_SIZE = 5000
diff --git a/utils/timestep_front_count.py b/utils/timestep_front_count.py
deleted file mode 100644
index 0e3812c..0000000
--- a/utils/timestep_front_count.py
+++ /dev/null
@@ -1,56 +0,0 @@
-"""
-Tool for counting the number of fronts in each timestep so that tensorflow datasets can be quickly generated. This script
-effectively prevents empty timesteps from being analyzed by 'convert_netcdf_to_tf.py', saving potentially large amounts
-of time when generating tensorflow datasets.
-
-A dictionary containing front counts for timesteps across a given domain will be saved to a pickle file in a directory for
-tensorflow datasets.
-
-Author: Andrew Justin (andrewjustinwx@gmail.com)
-Script version: 2023.6.13
-"""
-import os
-import sys
-import csv
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))) # this line allows us to import scripts outside the current directory
-import file_manager as fm
-from utils.settings import DEFAULT_DOMAIN_INDICES
-import argparse
-import numpy as np
-import xarray as xr
-
-
-if __name__ == '__main__':
- parser = argparse.ArgumentParser()
- parser.add_argument('--fronts_netcdf_indir', type=str, required=True,
- help="Input directory for the netCDF files containing frontal boundary data.")
- parser.add_argument('--tf_outdir', type=str, required=True,
- help="Output directory for future tensorflow datasets. This is where the pickle file containing frontal counts will "
- "also be saved.")
- parser.add_argument('--domain', type=str, default='conus', help='Domain from which to pull the images.')
-
- args = vars(parser.parse_args())
-
- front_files_obj = fm.DataFileLoader(args['fronts_netcdf_indir'], data_file_type='fronts-netcdf')
- front_files = front_files_obj.front_files
-
- isel_kwargs = {'longitude': slice(DEFAULT_DOMAIN_INDICES[args['domain']][0], DEFAULT_DOMAIN_INDICES[args['domain']][1]),
- 'latitude': slice(DEFAULT_DOMAIN_INDICES[args['domain']][2], DEFAULT_DOMAIN_INDICES[args['domain']][3])}
-
- fieldnames = ['File', 'CF', 'CF-F', 'CF-D', 'WF', 'WF-F', 'WF-D', 'SF', 'SF-F', 'SF-D', 'OF', 'OF-F', 'OF-D', 'INST',
- 'TROF', 'TT', 'DL']
-
- front_count_csv_file = '%s/timestep_front_counts_%s.csv' % (args['tf_outdir'], args['domain'])
-
- with open('%s/timestep_front_counts.csv' % args['tf_outdir'], 'w', newline='') as f:
- csvwriter = csv.writer(f)
- csvwriter.writerow(fieldnames)
-
- for file_no, front_file in enumerate(front_files):
- print(front_file, end='\r')
- front_dataset = xr.open_dataset(front_file, engine='netcdf4').isel(**isel_kwargs).expand_dims('time', axis=0).astype('float16')
- front_bins = np.bincount(front_dataset['identifier'].values.astype('int64').flatten(), minlength=17)[1:] # counts for each front type ('no front' type removed)
-
- row = [os.path.basename(front_file), *front_bins]
-
- csvwriter.writerow(row)