From b635e18bd41f43c01305ac3db677cf39c05566e6 Mon Sep 17 00:00:00 2001 From: Raymond Menzel Date: Wed, 18 Dec 2024 17:13:13 -0500 Subject: [PATCH 1/3] add the ability to recursively search for valid plugin classes --- .../analysis_scripts/plugins.py | 88 ++++++++++++++----- 1 file changed, 68 insertions(+), 20 deletions(-) diff --git a/core/analysis_scripts/analysis_scripts/plugins.py b/core/analysis_scripts/analysis_scripts/plugins.py index 7c99cc4..18bbf4f 100644 --- a/core/analysis_scripts/analysis_scripts/plugins.py +++ b/core/analysis_scripts/analysis_scripts/plugins.py @@ -1,24 +1,82 @@ import importlib import inspect +from pathlib import Path import pkgutil from .base_class import AnalysisScript -# Dictionary of found plugins. -discovered_plugins = {} -for finder, name, ispkg in pkgutil.iter_modules(): - if name.startswith("freanalysis_"): - discovered_plugins[name] = importlib.import_module(name) - - class UnknownPluginError(BaseException): """Custom exception for when an invalid plugin name is used.""" pass +def _find_plugin_class(module): + """Looks for a class that inherits from AnalysisScript. + + Args: + module: Module object. + + Returns: + Class that inherits from AnalysisScript. + + Raises: + UnknownPluginError if no class is found. + """ + for attribute in vars(module).values(): + # Try to find a class that inherits from the AnalysisScript class. + if inspect.isclass(attribute) and AnalysisScript in attribute.__bases__: + # Instantiate an object of this class. + return attribute + raise UnknownPluginError("could not find class that inherts from AnalysisScripts") + + +sanity_counter = 0 # How much recursion is happening. +maximum_craziness = 100 # This is too much recursion. + + +def _recursive_search(name, ispkg): + """Recursively search for a module that has a class that inherits from AnalysisScript. + + Args: + name: String name of the module. + ispkg: Flag telling whether or not the module is a package. + + Returns: + Class that inherits from AnalysisScript. + + Raises: + UnknownPluginError if no class is found. + ValueError if there is too much recursion. + """ + global sanity_counter + sanity_counter += 1 + if sanity_counter > maximum_craziness: + raise ValueError(f"recursion level {sanity_counter} too high.") + + module = importlib.import_module(name) + try: + return _find_plugin_class(module) + except UnknownPluginError: + if not ispkg: + # Do not recurse further. + raise + paths = module.__spec__.submodule_search_locations + for finder, subname, ispkg in pkgutil.iter_modules(paths): + subname = f"{name}.{subname}" + return _recursive_search(subname, ispkg) + + +# Dictionary of found plugins. +discovered_plugins = {} +for finder, name, ispkg in pkgutil.iter_modules(): + if name.startswith("freanalysis_") and ispkg: + discovered_plugins[name] = _recursive_search(name, True) + + def _plugin_object(name): - """Searches for a class that inherits from AnalysisScript in the plugin module. + """Attempts to create an object from a class that inherits from AnalysisScript in + the plugin module. Args: name: Name of the plugin. @@ -27,23 +85,13 @@ def _plugin_object(name): The object that inherits from AnalysisScript. Raises: - KeyError if the input name does not match any installed plugins. - ValueError if no object that inhertis from AnalysisScript is found in the - plugin module. + UnknownPluginError if the input name is not in the disovered_plugins dictionary. """ - # Loop through all attributes in the plugin package with the input name. try: - plugin_module = discovered_plugins[name] + return discovered_plugins[name]() except KeyError: raise UnknownPluginError(f"could not find analysis script plugin {name}.") - for attribute in vars(plugin_module).values(): - # Try to find a class that inherits from the AnalysisScript class. - if inspect.isclass(attribute) and AnalysisScript in attribute.__bases__: - # Instantiate an object of this class. - return attribute() - raise ValueError(f"could not find compatible object in {name}.") - def available_plugins(): """Returns a list of plugin names.""" From 79a7b06482c591d3f2e1efc3dc077e9f861f5e2a Mon Sep 17 00:00:00 2001 From: Raymond Menzel Date: Tue, 14 Jan 2025 15:41:09 -0500 Subject: [PATCH 2/3] reoragnize some analysis scripts and fix plugin discovery --- .../analysis_scripts/plugins.py | 6 +- core/figure_tools/figure_tools/__init__.py | 3 +- .../figure_tools/figure_tools/common_plots.py | 24 ++ core/figure_tools/figure_tools/lon_lat_map.py | 21 +- .../figure_tools/figure_tools/time_subsets.py | 39 +++ core/figure_tools/pyproject.toml | 1 + .../freanalysis_aerosol/__init__.py | 210 +----------- .../freanalysis_aerosol/aerosols.py | 212 ++++++++++++ .../freanalysis_clouds/__init__.py | 5 +- .../freanalysis_land/freanalysis_land/land.py | 84 +++-- .../tests/test_freanalysis_land.py | 4 +- .../freanalysis_radiation/__init__.py | 310 +----------------- .../freanalysis_radiation/radiation.py | 293 +++++++++++++++++ .../__init__.py | 306 +++++++++++++++++ .../pyproject.toml | 28 ++ 15 files changed, 986 insertions(+), 560 deletions(-) create mode 100644 user-analysis-scripts/freanalysis_aerosol/freanalysis_aerosol/aerosols.py create mode 100644 user-analysis-scripts/freanalysis_radiation/freanalysis_radiation/radiation.py create mode 100644 user-analysis-scripts/freanalysis_radiation_atmos_av_mon/freanalysis_radiation_atmos_av_mon/__init__.py create mode 100644 user-analysis-scripts/freanalysis_radiation_atmos_av_mon/pyproject.toml diff --git a/core/analysis_scripts/analysis_scripts/plugins.py b/core/analysis_scripts/analysis_scripts/plugins.py index 18bbf4f..36593a0 100644 --- a/core/analysis_scripts/analysis_scripts/plugins.py +++ b/core/analysis_scripts/analysis_scripts/plugins.py @@ -64,7 +64,11 @@ def _recursive_search(name, ispkg): paths = module.__spec__.submodule_search_locations for finder, subname, ispkg in pkgutil.iter_modules(paths): subname = f"{name}.{subname}" - return _recursive_search(subname, ispkg) + try: + return _recursive_search(subname, ispkg) + except UnknownPluginError: + # Didn't find it, so continue to iterate. + pass # Dictionary of found plugins. diff --git a/core/figure_tools/figure_tools/__init__.py b/core/figure_tools/figure_tools/__init__.py index 34f376b..318d98b 100644 --- a/core/figure_tools/figure_tools/__init__.py +++ b/core/figure_tools/figure_tools/__init__.py @@ -1,6 +1,7 @@ from .anomaly_timeseries import AnomalyTimeSeries from .common_plots import observation_vs_model_maps, radiation_decomposition, \ - timeseries_and_anomalies, zonal_mean_vertical_and_column_integrated_map + timeseries_and_anomalies, zonal_mean_vertical_and_column_integrated_map, \ + chuck_radiation from .figure import Figure from .global_mean_timeseries import GlobalMeanTimeSeries from .lon_lat_map import LonLatMap diff --git a/core/figure_tools/figure_tools/common_plots.py b/core/figure_tools/figure_tools/common_plots.py index b5e27ee..49d9168 100644 --- a/core/figure_tools/figure_tools/common_plots.py +++ b/core/figure_tools/figure_tools/common_plots.py @@ -5,6 +5,30 @@ from .figure import Figure +def chuck_radiation(reference, model, title): + figure = Figure(num_rows=3, num_columns=1, title=title, size=(14, 12)) + + # Create a common color bar for the reference and model. + colorbar_range = [120, 340] + + # Model data. + global_mean = model.global_mean() + figure.add_map(model, f"Model [Mean: {global_mean:.2f}]", 1, + colorbar_range=colorbar_range, num_levels=11, colormap="jet") + + # Reference data. + global_mean = reference.global_mean() + figure.add_map(reference, f"Obersvations [Mean: {global_mean:.2f}]", 2, + colorbar_range=colorbar_range, num_levels=11, colormap="jet") + + # Difference between the reference and the model. + difference = model - reference + color_range = [-34., 34.] + figure.add_map(difference, f"Model - Obs [Mean: {global_mean:.2f}]", 3, + colorbar_range=color_range, colormap="jet", normalize_colors=True) + return figure + + def observation_vs_model_maps(reference, model, title): figure = Figure(num_rows=2, num_columns=2, title=title, size=(14, 12)) diff --git a/core/figure_tools/figure_tools/lon_lat_map.py b/core/figure_tools/figure_tools/lon_lat_map.py index 382fc4b..c753ac6 100644 --- a/core/figure_tools/figure_tools/lon_lat_map.py +++ b/core/figure_tools/figure_tools/lon_lat_map.py @@ -54,7 +54,7 @@ def __sub__(self, arg): @classmethod def from_xarray_dataset(cls, dataset, variable, time_method=None, time_index=None, - year=None): + year=None, year_range=None, month_range=None): """Instantiates a LonLatMap object from an xarray dataset.""" v = dataset.data_vars[variable] data = array(v.data[...]) @@ -75,8 +75,25 @@ def from_xarray_dataset(cls, dataset, variable, time_method=None, time_index=Non time = TimeSubset(time) data = time.annual_mean(data, year) timestamp = r"$\bar{t} = $" + str(year) + elif "climatology" in time_method: + if year_range == None or len(year_range) != 2: + raise ValueError("year_range is required ([star year, end year]" + + " when time_method is a climatology.") + if time_method == "annual climatology": + time = TimeSubset(time) + data = time.annual_climatology(data, year_range) + timestamp = f"{year_range[0]} - {year_range[1]} annual climatology" + elif time_method == "seasonal climatology": + if month_range == None or len(month_range) != 2: + raise ValueError("month_range is required ([start month, end month])" + + " when time_method='seasonal climatology'.") + time = TimeSubset(time) + data = time.seasonal_climatology(data, year_range, month_range) + timestamp = "" else: - raise ValueError("time_method must be either 'instantaneous' or 'annual mean.'") + valid_values = ["instantaneous", "annual mean", "annual climatology", + "seasonal climatology"] + raise ValueError(f"time_method must one of :{valid_values}.") else: timestamp = None diff --git a/core/figure_tools/figure_tools/time_subsets.py b/core/figure_tools/figure_tools/time_subsets.py index a5c89ad..18089d5 100644 --- a/core/figure_tools/figure_tools/time_subsets.py +++ b/core/figure_tools/figure_tools/time_subsets.py @@ -10,6 +10,45 @@ def __init__(self, data): """ self.data = data + def annual_climatology(self, data, year_range): + years = [x for x in range(year_range[0], year_range[1] + 1)] + sum_, counter = None, 0 + for i, point in enumerate(self.data): + _, year = self._month_and_year(point) + if year in years: + if sum_ is None: + sum_ = data[i, ...] + else: + sum_ += data[i, ...] + counter += 1 + + if counter != 12*len(years): + raise ValueError("Expected monthly data and did not find correct number of months.") + return array(sum_[...]/counter) + + def seasonal_climatology(self, data, year_range, month_range): + years = [x for x in range(year_range[0], year_range[1] + 1)] + if month_range[1] - month_range[0] < 0: + # We have crossed to the next year. + months = [x for x in range(month_range[0], 13)] + \ + [x for x in range(1, month_range[1] + 1)] + else: + months = [x for x in range(month_range[0], month_range[1] + 1)] + + sum_, counter = None, 0 + for i, point in enumerate(self.data): + month, year = self._month_and_year(point) + if month in months and year in years: + if sum_ is None: + sum_ = data[i, ...] + else: + sum_ += data[i, ...] + counter += 1 + + if counter != len(months)*len(years): + raise ValueError("Expected monthly data and did not find enough months.") + return array(sum_[...]/counter) + def annual_mean(self, data, year): """Calculates the annual mean of the input date for the input year. diff --git a/core/figure_tools/pyproject.toml b/core/figure_tools/pyproject.toml index 5d674f2..3dc447d 100644 --- a/core/figure_tools/pyproject.toml +++ b/core/figure_tools/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "cartopy", "matplotlib", "numpy", + "scipy", "xarray", ] requires-python = ">= 3.6" diff --git a/user-analysis-scripts/freanalysis_aerosol/freanalysis_aerosol/__init__.py b/user-analysis-scripts/freanalysis_aerosol/freanalysis_aerosol/__init__.py index 613ed76..aecfcca 100644 --- a/user-analysis-scripts/freanalysis_aerosol/freanalysis_aerosol/__init__.py +++ b/user-analysis-scripts/freanalysis_aerosol/freanalysis_aerosol/__init__.py @@ -1,209 +1 @@ -from dataclasses import dataclass -import json -from pathlib import Path - -from analysis_scripts import AnalysisScript -from figure_tools import LonLatMap, zonal_mean_vertical_and_column_integrated_map, \ - ZonalMeanMap -import intake - - -@dataclass -class Metadata: - """Helper class that stores the metadata needed by the plugin.""" - frequency: str = "monthly" - realm: str = "atmos" - - @staticmethod - def variables(): - """Helper function to make maintaining this script easier if the - catalog variable ids change. - - Returns: - Dictionary mapping the names used in this script to the catalog - variable ids. - """ - return { - "black_carbon": "blk_crb", - "black_carbon_column": "blk_crb_col", - "large_dust": "lg_dust", - "large_dust_column": "lg_dust_col", - "small_dust": "sm_dust", - "small_dust_column": "sm_dust_col", - "organic_carbon": "org_crb", - "organic_carbon_column": "org_crb_col", - "large_seasalt": "lg_ssalt", - "large_seasalt_column": "lg_ssalt_col", - "small_seasalt": "sm_ssalt", - "small_seasalt_column": "sm_ssalt_col", - "seasalt": "salt", - "seasalt_column": "salt_col", - "sulfate": "sulfate", - "sulfate_column": "sulfate_col", - } - - -class AerosolAnalysisScript(AnalysisScript): - """Aerosol analysis script. - - Attributes: - description: Longer form description for the analysis. - title: Title that describes the analysis. - """ - def __init__(self): - self.metadata = Metadata() - self.description = "Calculates aerosol mass metrics." - self.title = "Aerosol Masses" - - def requires(self): - """Provides metadata describing what is needed for this analysis to run. - - Returns: - A json string containing the metadata. - """ - columns = Metadata.__annotations__.keys() - settings = {x: getattr(self.metadata, x) for x in columns} - return json.dumps({ - "settings": settings, - "dimensions": { - "lat": {"standard_name": "latitude"}, - "lon": {"standard_name": "longitude"}, - "pfull": {"standard_name": "air_pressure"}, - "time": {"standard_name": "time"} - }, - "varlist": { - "blk_crb": { - "standard_name": "black_carbon_mass", - "units": "kg m-3", - "dimensions": ["time", "pfull", "lat", "lon"] - }, - "blk_crb_col": { - "standard_name": "column_integrated_black_carbon_mass", - "units": "kg m-2", - "dimensions": ["time", "lat", "lon"] - }, - "lg_dust": { - "standard_name": "large_dust_mass", - "units": "kg m-3", - "dimensions": ["time", "pfull", "lat", "lon"] - }, - "lg_dust_col": { - "standard_name": "column_integrated_large_dust_mass", - "units": "kg m-2", - "dimensions": ["time", "lat", "lon"] - }, - "lg_ssalt": { - "standard_name": "large_seasalt_mass", - "units": "kg m-3", - "dimensions": ["time", "pfull", "lat", "lon"] - }, - "lg_ssalt_col": { - "standard_name": "column_integrated_large_ssalt_mass", - "units": "kg m-2", - "dimensions": ["time", "lat", "lon"] - }, - "org_crb": { - "standard_name": "organic_carbon_mass", - "units": "kg m-3", - "dimensions": ["time", "pfull", "lat", "lon"] - }, - "org_crb_col": { - "standard_name": "column_integrated_organic_carbon_mass", - "units": "kg m-2", - "dimensions": ["time", "lat", "lon"] - }, - "salt": { - "standard_name": "seasalt_mass", - "units": "kg m-3", - "dimensions": ["time", "pfull", "lat", "lon"] - }, - "salt_col": { - "standard_name": "column_integrated_seasalt_mass", - "units": "kg m-2", - "dimensions": ["time", "lat", "lon"] - }, - "sm_dust": { - "standard_name": "small_dust_mass", - "units": "kg m-3", - "dimensions": ["time", "pfull", "lat", "lon"] - }, - "sm_dust_col": { - "standard_name": "column_integrated_small_dust_mass", - "units": "kg m-2", - "dimensions": ["time", "lat", "lon"] - }, - "sm_ssalt": { - "standard_name": "small_seasalt_mass", - "units": "kg m-3", - "dimensions": ["time", "pfull", "lat", "lon"] - }, - "sm_ssalt_col": { - "standard_name": "column_integrated_small_ssalt_mass", - "units": "kg m-2", - "dimensions": ["time", "lat", "lon"] - }, - "sulfate": { - "standard_name": "sulfate_mass", - "units": "kg m-3", - "dimensions": ["time", "pfull", "lat", "lon"] - }, - "sulfate_col": { - "standard_name": "column_integrated_sulfate_mass", - "units": "kg m-2", - "dimensions": ["time", "lat", "lon"] - }, - }, - }) - - def run_analysis(self, catalog, png_dir, reference_catalog=None, config={}): - """Runs the analysis and generates all plots and associated datasets. - - Args: - catalog: Path to a catalog. - png_dir: Path to the directory where the figures will be made. - reference_catalog: Path to a catalog of reference data. - config: Dictionary of catalog metadata. Will overwrite the - data defined in the Metadata helper class if they both - contain the same keys. - - Returns: - A list of paths to the figures that were created. - - Raises: - ValueError if the catalog cannot be filtered correctly. - """ - - # Connect to the catalog and find the necessary datasets. - catalog = intake.open_esm_datastore(catalog) - - maps = {} - for name, variable in self.metadata.variables().items(): - # Filter the catalog down to a single dataset for each variable. - query_params = {"variable_id": variable} - query_params.update(vars(self.metadata)) - query_params.update(config) - datasets = catalog.search(**query_params).to_dataset_dict(progressbar=False) - if len(list(datasets.values())) != 1: - raise ValueError("could not filter the catalog down to a single dataset.") - dataset = list(datasets.values())[0] - - if name.endswith("column"): - # Lon-lat maps. - maps[name] = LonLatMap.from_xarray_dataset(dataset, variable, year=1980, - time_method="annual mean") - else: - maps[name] = ZonalMeanMap.from_xarray_dataset(dataset, variable, year=1980, - time_method="annual mean", - invert_y_axis=True) - - figure_paths = [] - for name in self.metadata.variables().keys(): - if name.endswith("column"): continue - figure = zonal_mean_vertical_and_column_integrated_map( - maps[name], - maps[f"{name}_column"], - f"{name.replace('_', ' ')} Mass", - ) - figure.save(Path(png_dir) / f"{name}.png") - figure_paths.append(Path(png_dir)/ f"{name}.png") - return figure_paths \ No newline at end of file +from .aerosols import AerosolAnalysisScript diff --git a/user-analysis-scripts/freanalysis_aerosol/freanalysis_aerosol/aerosols.py b/user-analysis-scripts/freanalysis_aerosol/freanalysis_aerosol/aerosols.py new file mode 100644 index 0000000..4122790 --- /dev/null +++ b/user-analysis-scripts/freanalysis_aerosol/freanalysis_aerosol/aerosols.py @@ -0,0 +1,212 @@ +from dataclasses import dataclass +import json +from pathlib import Path + +from analysis_scripts import AnalysisScript +from figure_tools import LonLatMap, zonal_mean_vertical_and_column_integrated_map, \ + ZonalMeanMap +import intake + + +@dataclass +class Metadata: + """Helper class that stores the metadata needed by the plugin.""" + frequency: str = "mon" + realm: str = "atmos_month_aer" + + @staticmethod + def variables(): + """Helper function to make maintaining this script easier if the + catalog variable ids change. + + Returns: + Dictionary mapping the names used in this script to the catalog + variable ids. + """ + return { + "black_carbon": "blk_crb", + "black_carbon_column": "blk_crb_col", + "large_dust": "lg_dust", + "large_dust_column": "lg_dust_col", + "small_dust": "sm_dust", + "small_dust_column": "sm_dust_col", + "organic_carbon": "org_crb", + "organic_carbon_column": "org_crb_col", + "large_seasalt": "lg_ssalt", + "large_seasalt_column": "lg_ssalt_col", + "small_seasalt": "sm_ssalt", + "small_seasalt_column": "sm_ssalt_col", + "seasalt": "salt", + "seasalt_column": "salt_col", + "sulfate": "sulfate", + "sulfate_column": "sulfate_col", + } + + +class AerosolAnalysisScript(AnalysisScript): + """Aerosol analysis script. + + Attributes: + description: Longer form description for the analysis. + title: Title that describes the analysis. + """ + def __init__(self): + self.metadata = Metadata() + self.description = "Calculates aerosol mass metrics." + self.title = "Aerosol Masses" + + def requires(self): + """Provides metadata describing what is needed for this analysis to run. + + Returns: + A json string containing the metadata. + """ + columns = Metadata.__annotations__.keys() + settings = {x: getattr(self.metadata, x) for x in columns} + return json.dumps({ + "settings": settings, + "dimensions": { + "lat": {"standard_name": "latitude"}, + "lon": {"standard_name": "longitude"}, + "pfull": {"standard_name": "air_pressure"}, + "time": {"standard_name": "time"} + }, + "varlist": { + "blk_crb": { + "standard_name": "black_carbon_mass", + "units": "kg m-3", + "dimensions": ["time", "pfull", "lat", "lon"] + }, + "blk_crb_col": { + "standard_name": "column_integrated_black_carbon_mass", + "units": "kg m-2", + "dimensions": ["time", "lat", "lon"] + }, + "lg_dust": { + "standard_name": "large_dust_mass", + "units": "kg m-3", + "dimensions": ["time", "pfull", "lat", "lon"] + }, + "lg_dust_col": { + "standard_name": "column_integrated_large_dust_mass", + "units": "kg m-2", + "dimensions": ["time", "lat", "lon"] + }, + "lg_ssalt": { + "standard_name": "large_seasalt_mass", + "units": "kg m-3", + "dimensions": ["time", "pfull", "lat", "lon"] + }, + "lg_ssalt_col": { + "standard_name": "column_integrated_large_ssalt_mass", + "units": "kg m-2", + "dimensions": ["time", "lat", "lon"] + }, + "org_crb": { + "standard_name": "organic_carbon_mass", + "units": "kg m-3", + "dimensions": ["time", "pfull", "lat", "lon"] + }, + "org_crb_col": { + "standard_name": "column_integrated_organic_carbon_mass", + "units": "kg m-2", + "dimensions": ["time", "lat", "lon"] + }, + "salt": { + "standard_name": "seasalt_mass", + "units": "kg m-3", + "dimensions": ["time", "pfull", "lat", "lon"] + }, + "salt_col": { + "standard_name": "column_integrated_seasalt_mass", + "units": "kg m-2", + "dimensions": ["time", "lat", "lon"] + }, + "sm_dust": { + "standard_name": "small_dust_mass", + "units": "kg m-3", + "dimensions": ["time", "pfull", "lat", "lon"] + }, + "sm_dust_col": { + "standard_name": "column_integrated_small_dust_mass", + "units": "kg m-2", + "dimensions": ["time", "lat", "lon"] + }, + "sm_ssalt": { + "standard_name": "small_seasalt_mass", + "units": "kg m-3", + "dimensions": ["time", "pfull", "lat", "lon"] + }, + "sm_ssalt_col": { + "standard_name": "column_integrated_small_ssalt_mass", + "units": "kg m-2", + "dimensions": ["time", "lat", "lon"] + }, + "sulfate": { + "standard_name": "sulfate_mass", + "units": "kg m-3", + "dimensions": ["time", "pfull", "lat", "lon"] + }, + "sulfate_col": { + "standard_name": "column_integrated_sulfate_mass", + "units": "kg m-2", + "dimensions": ["time", "lat", "lon"] + }, + }, + }) + + def run_analysis(self, catalog, png_dir, config=None, reference_catalog=None): + """Runs the analysis and generates all plots and associated datasets. + + Args: + catalog: Path to a catalog. + png_dir: Path to the directory where the figures will be made. + reference_catalog: Path to a catalog of reference data. + config: Dictionary of catalog metadata. Will overwrite the + data defined in the Metadata helper class if they both + contain the same keys. + + Returns: + A list of paths to the figures that were created. + + Raises: + ValueError if the catalog cannot be filtered correctly. + """ + + # Connect to the catalog and find the necessary datasets. + catalog = intake.open_esm_datastore(catalog) + + maps = {} + for name, variable in self.metadata.variables().items(): + print(f"Working on variable {name}") + # Filter the catalog down to a single dataset for each variable. + query_params = {"variable_id": variable} + query_params.update(vars(self.metadata)) + if config: + query_params.update(config) + datasets = catalog.search(**query_params).to_dataset_dict(progressbar=False) + if len(list(datasets.values())) != 1: + print(variable, datasets) + raise ValueError("could not filter the catalog down to a single dataset.") + dataset = list(datasets.values())[0] + + if name.endswith("column"): + # Lon-lat maps. + maps[name] = LonLatMap.from_xarray_dataset(dataset, variable, year=1980, + time_method="annual mean") + else: + maps[name] = ZonalMeanMap.from_xarray_dataset(dataset, variable, year=1980, + time_method="annual mean", + invert_y_axis=True) + + figure_paths = [] + for name in self.metadata.variables().keys(): + if name.endswith("column"): continue + figure = zonal_mean_vertical_and_column_integrated_map( + maps[name], + maps[f"{name}_column"], + f"{name.replace('_', ' ')} Mass", + ) + figure.save(Path(png_dir) / f"{name}.png") + figure_paths.append(Path(png_dir)/ f"{name}.png") + return figure_paths diff --git a/user-analysis-scripts/freanalysis_clouds/freanalysis_clouds/__init__.py b/user-analysis-scripts/freanalysis_clouds/freanalysis_clouds/__init__.py index 0cdaa8a..ed6aaa7 100644 --- a/user-analysis-scripts/freanalysis_clouds/freanalysis_clouds/__init__.py +++ b/user-analysis-scripts/freanalysis_clouds/freanalysis_clouds/__init__.py @@ -81,7 +81,7 @@ def run_analysis(self, catalog, png_dir, config=None, reference_catalog=None): Args: catalog: Path to a catalog. png_dir: Path to the directory where the figures will be made. - config: Dictonary of catalog metadata. Will overwrite the + config: Dictionary of catalog metadata. Will overwrite the data defined in the Metadata helper class if they both contain the same keys. reference_catalog: Path to a catalog of reference data. @@ -103,7 +103,8 @@ def run_analysis(self, catalog, png_dir, config=None, reference_catalog=None): # Filter the catalog down to a single dataset for each variable. query_params = {"variable_id": variable} query_params.update(vars(self.metadata)) - query_params.update(config) + if config: + query_params.update(config) datasets = catalog.search(**query_params).to_dataset_dict(progressbar=False) if len(list(datasets.values())) != 1: raise ValueError("could not filter the catalog down to a single dataset.", datasets) diff --git a/user-analysis-scripts/freanalysis_land/freanalysis_land/land.py b/user-analysis-scripts/freanalysis_land/freanalysis_land/land.py index 0bd3f0a..2386b17 100644 --- a/user-analysis-scripts/freanalysis_land/freanalysis_land/land.py +++ b/user-analysis-scripts/freanalysis_land/freanalysis_land/land.py @@ -1,12 +1,14 @@ +import datetime import json +from pathlib import Path +import re + from analysis_scripts import AnalysisScript import intake import matplotlib.pyplot as plt import cartopy import cartopy.crs as ccrs -import datetime import pandas as pd -import re import xarray as xr @@ -32,12 +34,13 @@ def requires(self): raise NotImplementedError("you must override this function.") return json.dumps("{json of metadata MDTF format.}") - def global_map(self,dataset,var,dates,plt_time=None,colormap='viridis',title=''): + def global_map(self, dataset, var, dates, plt_time=None, colormap='viridis', title=''): """ Generate a global map and regional subplots for a specified variable from an xarray dataset. - This function creates a global map and several regional subplots (North America, South America, Europe, Africa, Asia, and Australia) - using the specified variable from an xarray dataset. The generated map will be saved as a PNG file. + This function creates a global map and several regional subplots (North America, + South America, Europe, Africa, Asia, and Australia) using the specified variable + from an xarray dataset. The generated map will be saved as a PNG file. Parameters: ---------- @@ -78,7 +81,7 @@ def global_map(self,dataset,var,dates,plt_time=None,colormap='viridis',title='') lon = dataset.lon lat = dataset.lat - if plt_time is None: plt_time=len(dates)-1 + if plt_time is None: plt_time = len(dates) - 1 data = dataset[var][plt_time].values projection = ccrs.PlateCarree() @@ -115,13 +118,12 @@ def global_map(self,dataset,var,dates,plt_time=None,colormap='viridis',title='') plt.tight_layout() return fig - # plt.savefig(var+'_global_map.png') - # plt.close() - def timeseries(self, dataset,var,dates_period,var_range=None,minlon = 0,maxlon = 360,minlat = -90,maxlat=90,timerange=None,title=''): + def timeseries(self, dataset, var, dates_period, var_range=None, minlon = 0, + maxlon = 360, minlat = -90, maxlat=90, timerange=None, title=''): ''' - Generate a time series plot of the specified variable from a dataset within a given geographic and temporal range. - + Generate a time series plot of the specified variable from a dataset within a + given geographic and temporal range. Parameters: ----------- @@ -132,7 +134,8 @@ def timeseries(self, dataset,var,dates_period,var_range=None,minlon = 0,maxlon = dates_period : pandas.DatetimeIndex The dates for the time series data. var_range : tuple of float, optional - The range of variable values to include in the plot (min, max). If not provided, the default range is (0, inf). + The range of variable values to include in the plot (min, max). If not provided, + the default range is (0, inf). minlon : float, optional The minimum longitude to include in the plot. Default is 0. maxlon : float, optional @@ -142,7 +145,8 @@ def timeseries(self, dataset,var,dates_period,var_range=None,minlon = 0,maxlon = maxlat : float, optional The maximum latitude to include in the plot. Default is 90. timerange : tuple of int, optional - The range of years to plot (start_year, end_year). If not provided, all available years in the dataset will be plotted. + The range of years to plot (start_year, end_year). If not provided, all + available years in the dataset will be plotted. title : str, optional The title of the plot. Default is an empty string. @@ -150,42 +154,48 @@ def timeseries(self, dataset,var,dates_period,var_range=None,minlon = 0,maxlon = -------- matplotlib.figure.Figure The figure object containing the generated time series plot. - + Notes: ------ - The function filters the dataset based on the provided variable range, longitude, and latitude bounds. It then - calculates the monthly and annual means of the specified variable and plots the seasonal and annual means. - + The function filters the dataset based on the provided variable range, longitude, + and latitude bounds. It then calculates the monthly and annual means of the + specified variable and plots the seasonal and annual means. ''' if var_range is not None: - data_filtered = dataset.where((dataset[var] > var_range[0]) & (dataset[var] <= var_range[1]) & - (dataset.lat >= minlat) & (dataset.lon >= minlon) & - (dataset.lat <= maxlat) & (dataset.lon <= maxlon)) + data_filtered = dataset.where((dataset[var] > var_range[0]) & + (dataset[var] <= var_range[1]) & + (dataset.lat >= minlat) & + (dataset.lon >= minlon) & + (dataset.lat <= maxlat) & + (dataset.lon <= maxlon)) else: data_filtered = dataset.where((dataset[var] > 0) & - (dataset.lat >= minlat) & (dataset.lon >= minlon) & - (dataset.lat <= maxlat) & (dataset.lon <= maxlon)) + (dataset.lat >= minlat) & + (dataset.lon >= minlon) & + (dataset.lat <= maxlat) & + (dataset.lon <= maxlon)) data_filtered['time'] = dates_period data_df = pd.DataFrame(index = dates_period) - data_df['monthly_mean'] = data_filtered.resample(time='YE').mean(dim=['lat','lon'],skipna=True)[var].values + data_df['monthly_mean'] = data_filtered.resample(time='YE').mean(dim=['lat','lon'], + skipna=True)[var].values data_df['monthly_shift'] = data_df['monthly_mean'].shift(1) if timerange is not None: - ys, ye = (str(timerange[0]),str(timerange[1])) + ys, ye = (str(timerange[0]), str(timerange[1])) plot_df = data_df.loc[f'{ys}-1-1':f'{ye}-1-1'] else: plot_df = data_df fig, ax = plt.subplots() - plot_df.resample('Q').mean()['monthly_shift'].plot(ax=ax,label='Seasonal Mean') - plot_df.resample('Y').mean()['monthly_mean'].plot(ax=ax,label='Annual Mean') + plot_df.resample('Q').mean()['monthly_shift'].plot(ax=ax, label='Seasonal Mean') + plot_df.resample('Y').mean()['monthly_mean'].plot(ax=ax, label='Annual Mean') plt.legend() plt.title(title) plt.xlabel('Years') return fig - def run_analysis(self, catalog, png_dir, reference_catalog=None): + def run_analysis(self, catalog, png_dir, config=None, reference_catalog=None): """Runs the analysis and generates all plots and associated datasets. Args: @@ -196,15 +206,17 @@ def run_analysis(self, catalog, png_dir, reference_catalog=None): Returns: A list of png figures. """ - print ('WARNING: THESE FIGURES ARE FOR TESTING THE NEW ANALYSIS WORKFLOW ONLY AND SHOULD NOT BE USED IN ANY OFFICIAL MANNER FOR ANALYSIS OF LAND MODEL OUTPUT.') + print('WARNING: THESE FIGURES ARE FOR TESTING THE NEW ANALYSIS WORKFLOW ONLY' + + ' AND SHOULD NOT BE USED IN ANY OFFICIAL MANNER FOR ANALYSIS OF' + + ' LAND MODEL OUTPUT.') col = intake.open_esm_datastore(catalog) df = col.df # Soil Carbon var = 'cSoil' - print ('Soil Carbon Analysis') - cat = col.search(variable_id=var,realm='land_cmip') - other_dict = cat.to_dataset_dict(cdf_kwargs={'chunks':{'time':12},'decode_times':False}) + print('Soil Carbon Analysis') + cat = col.search(variable_id=var, realm='land_cmip') + other_dict = cat.to_dataset_dict(cdf_kwargs={'chunks': {'time': 12}, 'decode_times': False}) combined_dataset = xr.concat(list(dict(sorted(other_dict.items())).values()), dim='time') # Other data: @@ -216,11 +228,13 @@ def run_analysis(self, catalog, png_dir, reference_catalog=None): dates_period = pd.PeriodIndex(dates,freq='D') sm_fig = self.global_map(combined_dataset,var,dates,title='Soil Carbon Content (kg/m^2)') - plt.savefig(png_dir+var+'_global_map.png') + figure_paths = [Path(png_dir) / f"{var}_global_map.png",] + plt.savefig(str(figure_paths[-1])) plt.close() ts_fig = self.timeseries(combined_dataset,var,dates_period,title='Global Average Soil Carbon') - plt.savefig(png_dir+var+'_global_ts.png') + figure_paths.append(Path(png_dir) / f"{var}_global_ts.png") + plt.savefig(str(figure_paths[-1])) plt.close() # Soil Moisture @@ -239,5 +253,7 @@ def run_analysis(self, catalog, png_dir, reference_catalog=None): dates_period = pd.PeriodIndex(dates,freq='D') sm_fig = self.global_map(combined_dataset,var,dates,title='Soil Moisture (kg/m^2)') - plt.savefig(png_dir+var+'_global_map.png') + figure_paths.append(Path(png_dir) / f"{var}_global_map.png") + plt.savefig(str(figure_paths[-1])) plt.close() + return figure_paths diff --git a/user-analysis-scripts/freanalysis_land/tests/test_freanalysis_land.py b/user-analysis-scripts/freanalysis_land/tests/test_freanalysis_land.py index 19225fa..edfac83 100644 --- a/user-analysis-scripts/freanalysis_land/tests/test_freanalysis_land.py +++ b/user-analysis-scripts/freanalysis_land/tests/test_freanalysis_land.py @@ -3,8 +3,8 @@ def test_land_analysis_script(): land = LandAnalysisScript() - land.run_analysis("/work/a2p/lm4p2sc_GSWP3_hist_irr_catalog.json", "/work/a2p/") + land.run_analysis("/work/a2p/lm4p2sc_GSWP3_hist_irr_catalog.json", ".") -if __name_ == "__main__": +if __name__ == "__main__": test_land_analysis_script() diff --git a/user-analysis-scripts/freanalysis_radiation/freanalysis_radiation/__init__.py b/user-analysis-scripts/freanalysis_radiation/freanalysis_radiation/__init__.py index 827cc18..343777a 100644 --- a/user-analysis-scripts/freanalysis_radiation/freanalysis_radiation/__init__.py +++ b/user-analysis-scripts/freanalysis_radiation/freanalysis_radiation/__init__.py @@ -1,309 +1 @@ -from dataclasses import dataclass -import json -from pathlib import Path - -from analysis_scripts import AnalysisScript -from figure_tools import AnomalyTimeSeries, GlobalMeanTimeSeries, LonLatMap, \ - observation_vs_model_maps, radiation_decomposition, \ - timeseries_and_anomalies -import intake -import intake_esm -from xarray import open_dataset - - -@dataclass -class Metadata: - activity_id: str = "dev" - institution_id: str = "" - source_id: str = "am5" - experiment_id: str = "c96L65_am5f7b11r0_amip" - frequency: str = "P1M" - modeling_realm: str = "atmos" - table_id: str = "" - member_id: str = "na" - grid_label: str = "" - temporal_subset: str = "" - chunk_freq: str = "" - platform: str = "" - cell_methods: str = "" - chunk_freq: str = "P1Y" - - def catalog_search_args(self, name): - return { - "experiment_id": self.experiment_id, - "frequency": self.frequency, - "member_id": self.member_id, - "modeling_realm": self.modeling_realm, - "variable_id": name, - } - - def variables(self): - return { - "rlds": "lwdn_sfc", - "rldsaf": "lwdn_sfc_ad", - "rldscs": "lwdn_sfc_clr", - "rldscsaf": "lwdn_sfc_ad_clr", - "rlus": "lwup_sfc", - "rlusaf": "lwup_sfc_ad", - "rluscs": "lwup_sfc_clr", - "rluscsaf": "lwup_sfc_ad_clr", - "rlut": "olr", - "rlutaf": "lwtoa_ad", - "rlutcs": "olr_clr", - "rlutcsaf": "lwtoa_ad_clr", - "rsds": "swdn_sfc", - "rsdsaf": "swdn_sfc_ad", - "rsdscs": "swdn_sfc_clr", - "rsdscsaf": "swdn_sfc_ad_clr", - "rsus": "swup_sfc", - "rsusaf": "swup_sfc_ad", - "rsuscs": "swup_sfc_clr", - "rsuscsaf": "swup_sfc_ad_clr", - "rsut": "swup_toa", - "rsutaf": "swup_toa_ad", - "rsutcs": "swup_toa_clr", - "rsutcsaf": "swup_toa_ad_clr", - "rsdt": "swdn_toa", - } - - -class RadiationAnalysisScript(AnalysisScript): - """Abstract base class for analysis scripts. User-defined analysis scripts - should inhert from this class and override the requires and run_analysis methods. - - Attributes: - description: Longer form description for the analysis. - title: Title that describes the analysis. - """ - def __init__(self): - self.metadata = Metadata() - self.description = "Calculates radiative flux metrics." - self.title = "Radiative Fluxes" - - def requires(self): - """Provides metadata describing what is needed for this analysis to run. - - Returns: - A json string containing the metadata. - """ - columns = Metadata.__annotations__.keys() - settings = {x: getattr(self.metadata, x) for x in columns} - return json.dumps({ - "settings": settings, - "dimensions": { - "lat": {"standard_name": "latitude"}, - "lon": {"standard_name": "longitude"}, - "time": {"standard_name": "time"} - }, - "varlist": { - "lwup_sfc": { - "standard_name": "surface_outgoing_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "lwup_sfc_ad": { - "standard_name": "surface_outgoing_aerosol_free_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "lwup_sfc_clr": { - "standard_name": "surface_outgoing_clear_sky_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "lwup_sfc_ad_clr": { - "standard_name": "surface_outgoing_clear_sky_aerosol_free_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "lwdn_sfc": { - "standard_name": "surface_incoming_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "lwdn_sfc_ad": { - "standard_name": "surface_incoming_aerosol_free_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "lwdn_sfc_clr": { - "standard_name": "surface_incoming_clear_sky_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "lwdn_sfc_ad_clr": { - "standard_name": "surface_incoming_clear_sky_aerosol_free_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "olr": { - "standard_name": "toa_outgoing_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "lwtoa_ad": { - "standard_name": "toa_outgoing_aerosol_free_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "olr_clr": { - "standard_name": "toa_outgoing_clear_sky_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "lwtoa_ad_clr": { - "standard_name": "toa_outgoing_clear_sky_aerosol_free_longwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swup_sfc": { - "standard_name": "surface_outgoing_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swup_sfc_ad": { - "standard_name": "surface_outgoing_aerosol_free_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swup_sfc_clr": { - "standard_name": "surface_outgoing_clear_sky_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swup_sfc_ad_clr": { - "standard_name": "surface_outgoing_clear_sky_aerosol_free_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swdn_sfc": { - "standard_name": "surface_incoming_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swdn_sfc_ad": { - "standard_name": "surface_incoming_aerosol_free_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swdn_sfc_clr": { - "standard_name": "surface_incoming_clear_sky_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swdn_sfc_ad_clr": { - "standard_name": "surface_incoming_clear_sky_aerosol_free_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swup_toa": { - "standard_name": "toa_outgoing_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swup_toa_ad": { - "standard_name": "toa_outgoing_aerosol_free_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swup_toa_clr": { - "standard_name": "toa_outgoing_clear_sky_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swup_toa_ad_clr": { - "standard_name": "toa_outgoing_clear_sky_aerosol_free_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - "swdn_toa": { - "standard_name": "toa_downwelling_shortwave_flux", - "units": "W m-2", - "dimensions": ["time", "lat", "lon"] - }, - }, - }) - - def run_analysis(self, catalog, png_dir, reference_catalog=None): - """Runs the analysis and generates all plots and associated datasets. - - Args: - catalog: Path to a catalog. - png_dir: Path to the directory where the figures will be made. - reference_catalog: Path to a catalog of reference data. - - Returns: - A list of paths to the figures that were created. - """ - - # Connect to the catalog and find the necessary datasets. - catalog = intake.open_esm_datastore(catalog) - - anomalies = {} - maps = {} - timeseries = {} - for name, variable in self.metadata.variables().items(): - # Get the dataset out of the catalog. - args = self.metadata.catalog_search_args(variable) - - datasets = catalog.search( - **self.metadata.catalog_search_args(variable) - ).to_dataset_dict(progressbar=False) - dataset = list(datasets.values())[0] - - # Lon-lat maps. - maps[name] = LonLatMap.from_xarray_dataset( - dataset, - variable, - time_method="annual mean", - year=1980, - ) - - if name == "rlut": - anomalies[name] = AnomalyTimeSeries.from_xarray_dataset( - dataset, - variable, - ) - timeseries[name] = GlobalMeanTimeSeries.from_xarray_dataset( - dataset, - variable, - ) - - figure_paths = [] - - # OLR anomally timeseries. - figure = timeseries_and_anomalies(timeseries["rlut"], anomalies["rlut"], - "OLR Global Mean & Anomalies") - figure.save(Path(png_dir) / "olr-anomalies.png") - figure_paths.append(Path(png_dir) / "olr-anomalies.png") - - # OLR. - figure = radiation_decomposition(maps["rlutcsaf"], maps["rlutaf"], - maps["rlutcs"], maps["rlut"], "OLR") - figure.save(Path(png_dir) / "olr.png") - figure_paths.append(Path(png_dir) / "olr.png") - - # SW TOTA. - figure = radiation_decomposition(maps["rsutcsaf"], maps["rsutaf"], - maps["rsutcs"], maps["rsut"], - "Shortwave Outgoing Toa") - figure.save(Path(png_dir) / "sw-up-toa.png") - figure_paths.append(Path(png_dir) / "sw-up-toa.png") - - # Surface radiation budget. - surface_budget = [] - for suffix in ["csaf", "af", "cs", ""]: - surface_budget.append(maps[f"rlds{suffix}"] + maps[f"rsds{suffix}"] - - maps[f"rlus{suffix}"] - maps[f"rsus{suffix}"]) - figure = radiation_decomposition(*surface_budget, "Surface Radiation Budget") - figure.save(Path(png_dir) / "surface-radiation-budget.png") - figure_paths.append(Path(png_dir) / "surface-radiation-budget.png") - - # TOA radiation budget. - toa_budget = [] - for suffix in ["csaf", "af", "cs", ""]: - toa_budget.append(maps[f"rsdt"] - maps[f"rlut{suffix}"] - maps[f"rsut{suffix}"]) - figure = radiation_decomposition(*toa_budget, "TOA Radiation Budget") - figure.save(Path(png_dir) / "toa-radiation-budget.png") - figure_paths.append(Path(png_dir) / "toa-radiation-budget.png") - return figure_paths +from .radiation import RadiationAnalysisScript diff --git a/user-analysis-scripts/freanalysis_radiation/freanalysis_radiation/radiation.py b/user-analysis-scripts/freanalysis_radiation/freanalysis_radiation/radiation.py new file mode 100644 index 0000000..df68258 --- /dev/null +++ b/user-analysis-scripts/freanalysis_radiation/freanalysis_radiation/radiation.py @@ -0,0 +1,293 @@ +from dataclasses import dataclass +import json +from pathlib import Path + +from analysis_scripts import AnalysisScript +from figure_tools import AnomalyTimeSeries, GlobalMeanTimeSeries, LonLatMap, \ + observation_vs_model_maps, radiation_decomposition, \ + timeseries_and_anomalies +import intake + + +@dataclass +class Metadata: + frequency: str = "mon" + realm: str = "atmos" + + def variables(self): + return { + "rlds": "lwdn_sfc", + "rldsaf": "lwdn_sfc_ad", + "rldscs": "lwdn_sfc_clr", + "rldscsaf": "lwdn_sfc_ad_clr", + "rlus": "lwup_sfc", + "rlusaf": "lwup_sfc_ad", + "rluscs": "lwup_sfc_clr", + "rluscsaf": "lwup_sfc_ad_clr", + "rlut": "olr", + "rlutaf": "lwtoa_ad", + "rlutcs": "olr_clr", + "rlutcsaf": "lwtoa_ad_clr", + "rsds": "swdn_sfc", + "rsdsaf": "swdn_sfc_ad", + "rsdscs": "swdn_sfc_clr", + "rsdscsaf": "swdn_sfc_ad_clr", + "rsus": "swup_sfc", + "rsusaf": "swup_sfc_ad", + "rsuscs": "swup_sfc_clr", + "rsuscsaf": "swup_sfc_ad_clr", + "rsut": "swup_toa", + "rsutaf": "swup_toa_ad", + "rsutcs": "swup_toa_clr", + "rsutcsaf": "swup_toa_ad_clr", + "rsdt": "swdn_toa", + } + + +class RadiationAnalysisScript(AnalysisScript): + """Abstract base class for analysis scripts. User-defined analysis scripts + should inhert from this class and override the requires and run_analysis methods. + + Attributes: + description: Longer form description for the analysis. + title: Title that describes the analysis. + """ + def __init__(self): + self.metadata = Metadata() + self.description = "Calculates radiative flux metrics." + self.title = "Radiative Fluxes" + + def requires(self): + """Provides metadata describing what is needed for this analysis to run. + + Returns: + A json string containing the metadata. + """ + columns = Metadata.__annotations__.keys() + settings = {x: getattr(self.metadata, x) for x in columns} + return json.dumps({ + "settings": settings, + "dimensions": { + "lat": {"standard_name": "latitude"}, + "lon": {"standard_name": "longitude"}, + "time": {"standard_name": "time"} + }, + "varlist": { + "lwup_sfc": { + "standard_name": "surface_outgoing_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "lwup_sfc_ad": { + "standard_name": "surface_outgoing_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "lwup_sfc_clr": { + "standard_name": "surface_outgoing_clear_sky_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "lwup_sfc_ad_clr": { + "standard_name": "surface_outgoing_clear_sky_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "lwdn_sfc": { + "standard_name": "surface_incoming_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "lwdn_sfc_ad": { + "standard_name": "surface_incoming_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "lwdn_sfc_clr": { + "standard_name": "surface_incoming_clear_sky_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "lwdn_sfc_ad_clr": { + "standard_name": "surface_incoming_clear_sky_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "olr": { + "standard_name": "toa_outgoing_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "lwtoa_ad": { + "standard_name": "toa_outgoing_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "olr_clr": { + "standard_name": "toa_outgoing_clear_sky_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "lwtoa_ad_clr": { + "standard_name": "toa_outgoing_clear_sky_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swup_sfc": { + "standard_name": "surface_outgoing_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swup_sfc_ad": { + "standard_name": "surface_outgoing_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swup_sfc_clr": { + "standard_name": "surface_outgoing_clear_sky_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swup_sfc_ad_clr": { + "standard_name": "surface_outgoing_clear_sky_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swdn_sfc": { + "standard_name": "surface_incoming_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swdn_sfc_ad": { + "standard_name": "surface_incoming_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swdn_sfc_clr": { + "standard_name": "surface_incoming_clear_sky_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swdn_sfc_ad_clr": { + "standard_name": "surface_incoming_clear_sky_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swup_toa": { + "standard_name": "toa_outgoing_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swup_toa_ad": { + "standard_name": "toa_outgoing_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swup_toa_clr": { + "standard_name": "toa_outgoing_clear_sky_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swup_toa_ad_clr": { + "standard_name": "toa_outgoing_clear_sky_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "swdn_toa": { + "standard_name": "toa_downwelling_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + }, + }) + + def run_analysis(self, catalog, png_dir, config=None, reference_catalog=None): + """Runs the analysis and generates all plots and associated datasets. + + Args: + catalog: Path to a catalog. + png_dir: Path to the directory where the figures will be made. + config: Dictionary of catalog metadata. Will overwrite the data + defined in the Metadata helper class if they both contain + the same keys. + reference_catalog: Path to a catalog of reference data. + + Returns: + A list of paths to the figures that were created. + """ + + # Connect to the catalog and find the necessary datasets. + catalog = intake.open_esm_datastore(catalog) + + anomalies = {} + maps = {} + timeseries = {} + for name, variable in self.metadata.variables().items(): + # Filter the catalog down to a single dataset for each variable. + query_params = {"variable_id": variable} + query_params.update(vars(self.metadata)) + if config: + query_params.update(config) + datasets = catalog.search(**query_params).to_dataset_dict( + progressbar=True, + ) + if len(list(datasets.values())) != 1: + raise ValueError(f"could not filter the dataset down to just {variable}.") + dataset = list(datasets.values())[0] + + # Lon-lat maps. + maps[name] = LonLatMap.from_xarray_dataset( + dataset, + variable, + time_method="annual mean", + year=1980, + ) + + if name == "rlut": + anomalies[name] = AnomalyTimeSeries.from_xarray_dataset( + dataset, + variable, + ) + timeseries[name] = GlobalMeanTimeSeries.from_xarray_dataset( + dataset, + variable, + ) + + figure_paths = [] + + # OLR anomally timeseries. + figure = timeseries_and_anomalies(timeseries["rlut"], anomalies["rlut"], + "OLR Global Mean & Anomalies") + figure_paths.append(Path(png_dir) / "olr-anomalies.png") + figure.save(figure_paths[-1]) + + # OLR. + figure = radiation_decomposition(maps["rlutcsaf"], maps["rlutaf"], + maps["rlutcs"], maps["rlut"], "OLR") + figure_paths.append(Path(png_dir) / "olr.png") + figure.save(figure_paths[-1]) + + # SW TOTA. + figure = radiation_decomposition(maps["rsutcsaf"], maps["rsutaf"], + maps["rsutcs"], maps["rsut"], + "Shortwave Outgoing Toa") + figure_paths.append(Path(png_dir) / "sw-up-toa.png") + figure.save(figure_paths[-1]) + + # Surface radiation budget. + surface_budget = [] + for suffix in ["csaf", "af", "cs", ""]: + surface_budget.append(maps[f"rlds{suffix}"] + maps[f"rsds{suffix}"] - + maps[f"rlus{suffix}"] - maps[f"rsus{suffix}"]) + figure = radiation_decomposition(*surface_budget, "Surface Radiation Budget") + figure_paths.append(Path(png_dir) / "surface-radiation-budget.png") + figure.save(figure_paths[-1]) + + # TOA radiation budget. + toa_budget = [] + for suffix in ["csaf", "af", "cs", ""]: + toa_budget.append(maps[f"rsdt"] - maps[f"rlut{suffix}"] - maps[f"rsut{suffix}"]) + figure = radiation_decomposition(*toa_budget, "TOA Radiation Budget") + figure_paths.append(Path(png_dir) / "toa-radiation-budget.png") + figure.save(figure_paths[-1]) + return figure_paths diff --git a/user-analysis-scripts/freanalysis_radiation_atmos_av_mon/freanalysis_radiation_atmos_av_mon/__init__.py b/user-analysis-scripts/freanalysis_radiation_atmos_av_mon/freanalysis_radiation_atmos_av_mon/__init__.py new file mode 100644 index 0000000..ed300a7 --- /dev/null +++ b/user-analysis-scripts/freanalysis_radiation_atmos_av_mon/freanalysis_radiation_atmos_av_mon/__init__.py @@ -0,0 +1,306 @@ +from dataclasses import dataclass +import json +from pathlib import Path + +from analysis_scripts import AnalysisScript +from figure_tools import AnomalyTimeSeries, GlobalMeanTimeSeries, LonLatMap, \ + observation_vs_model_maps, radiation_decomposition, \ + timeseries_and_anomalies, chuck_radiation +import intake +import intake_esm +import pdb + +@dataclass +class Metadata: + """Helper class that stores the metadata needed by the plugin.""" + + def variables(self): + """Helper function to make maintaining this script easier if the + catalog variable ids change. + + Returns: + Dictionary mapping the names used in this script to the catalog + variable ids. + """ + return { + "rlds": "rlds", + "rldsaf": "rldsaf", + "rldscs": "rldscs", + "rldscsaf": "rldscsaf", + "rlus": "rlus", + "rlusaf": "rlusaf", + "rluscs": "rluscs", + "rluscsaf": "rluscsaf", + "rlut": "olr", + "rlutaf": "rlutaf", + "rlutcs": "rlutcs", + "rlutcsaf": "rlutcsaf", + "rsds": "rsds", + "rsdsaf": "rsdsaf", + "rsdscs": "rsdscs", + "rsdscsaf": "rsdscsaf", + "rsus": "rsus", + "rsusaf": "rsusaf", + "rsuscs": "rsuscs", + "rsuscsaf": "rsuscsaf", + "rsut": "rsut", + "rsutaf": "rsutaf", + "rsutcs": "rsutcs", + "rsutcsaf": "rsutcsaf", + "rsdt": "rsdt", + } + + def reference_variables(self): + pass + + +class RadiationVsCeresAnalysisScript(AnalysisScript): + """Radiative flux comparison to CERES analysis script. + + Attributes: + description: Longer form description for the analysis. + metadata: MetaData object to helper filter the data catalog. + title: Title that describes the analysis. + """ + def __init__(self): + self.metadata = Metadata() + self.description = "Compare the radiative flux to CERES." + self.title = "Radiative Fluxes Vs. CERES" + + def requires(self): + """Provides metadata describing what is needed for this analysis to run. + + Returns: + A json string containing the metadata. + """ + columns = Metadata.__annotations__.keys() + settings = {x: getattr(self.metadata, x) for x in columns} + return json.dumps({ + "settings": settings, + "dimensions": { + "lat": {"standard_name": "latitude"}, + "lon": {"standard_name": "longitude"}, + "time": {"standard_name": "time"} + }, + "varlist": { + "rlus": { + "standard_name": "surface_outgoing_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rlusaf": { + "standard_name": "surface_outgoing_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rluscs": { + "standard_name": "surface_outgoing_clear_sky_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rluscsaf": { + "standard_name": "surface_outgoing_clear_sky_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rlds": { + "standard_name": "surface_incoming_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rldsaf": { + "standard_name": "surface_incoming_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rldscs": { + "standard_name": "surface_incoming_clear_sky_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rldscsaf": { + "standard_name": "surface_incoming_clear_sky_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rlut": { + "standard_name": "toa_outgoing_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rlutaf": { + "standard_name": "toa_outgoing_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rlutcs": { + "standard_name": "toa_outgoing_clear_sky_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rlutcsaf": { + "standard_name": "toa_outgoing_clear_sky_aerosol_free_longwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsus": { + "standard_name": "surface_outgoing_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsusaf": { + "standard_name": "surface_outgoing_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsuscs": { + "standard_name": "surface_outgoing_clear_sky_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsuscsaf": { + "standard_name": "surface_outgoing_clear_sky_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsds": { + "standard_name": "surface_incoming_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsdsaf": { + "standard_name": "surface_incoming_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsdscs": { + "standard_name": "surface_incoming_clear_sky_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsdscsaf": { + "standard_name": "surface_incoming_clear_sky_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsut": { + "standard_name": "toa_outgoing_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsutaf": { + "standard_name": "toa_outgoing_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsutcs": { + "standard_name": "toa_outgoing_clear_sky_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsutcsaf": { + "standard_name": "toa_outgoing_clear_sky_aerosol_free_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + "rsdt": { + "standard_name": "toa_downwelling_shortwave_flux", + "units": "W m-2", + "dimensions": ["time", "lat", "lon"] + }, + }, + }) + + def run_analysis(self, catalog, png_dir, config=None, reference_catalog=None): + """Runs the analysis and generates all plots and associated datasets. + + Args: + catalog: Path to a catalog. + png_dir: Path to the directory where the figures will be made. + reference_catalog: Path to a catalog of reference data. + config: Dictonary of catalog metadata. Will overwrite the + data defined in the Metadata helper class if they both + contain the same keys. + + Returns: + A list of paths to the figures that were created. + + Raises: + ValueError if the catalog cannot be filtered correctly. + """ + # Seasonal climatology. + figure_paths = [] + for season in [[12, 2], [3, 5], [6, 8], [9, 11]]: + figure = self.plot_vs_obs(catalog, reference_catalog, "monthly", "rlut", + "toa_lw_all_mon", png_dir, season, + config={"realm": "atmos_cmip"}) + figure_paths.append(figure) + + # Annual mean climatology. + figure = self.plot_vs_obs(catalog, reference_catalog, "monthly", "rlut", + "toa_lw_all_mon", png_dir, config={"realm": "atmos_cmip"}) + figure_paths.append(figure) + return figure_paths + + def plot_vs_obs(self, catalog, reference_catalog, frequency, + variable, reference_variable, png_dir, + month_range=None, config=None): + # Connect to the catalog and filter to find the necessary datasets. + model_catalog = intake.open_esm_datastore(catalog) + query_params = {"variable_id": variable, "frequency": frequency} + query_params.update(vars(self.metadata)) + if config: + query_params.update(config) + + pdb.set_trace() + datasets = model_catalog.search(**query_params).to_dataset_dict(progressbar=False) + if len(list(datasets.values())) != 1: + print(query_params, list(datasets.values())) + raise ValueError("could not filter the catalog down to a single dataset.") + dataset = list(datasets.values())[0] + + # Model Lon-lat maps. + model_map = LonLatMap.from_xarray_dataset(dataset, variable, time_index=0, + time_method="instantaneous") + + # Connect to the reference catalog and get the reference datasets. + obs_catalog = intake.open_esm_datastore(reference_catalog) + query_params = { + "experiment_id": "ceres_ebaf_ed4.1", + "variable_id": reference_variable, + } + datasets = obs_catalog.search(**query_params).to_dataset_dict(progressbar=False) + if len(list(datasets.values())) != 1: + raise ValueError("could not filter the catalog down to a single dataset.") + dataset = list(datasets.values())[0] + + if month_range == None: + period = "annual" + title = f"Annual {variable}" + else: + period = "seasonal" + initials = {1: "j", 2: "f", 3: "m", 4: "a", 5: "m", 6: "j", + 7: "j", 8: "a", 9: "s", 10: "o", 11: "n", 12: "d"} + if month_range[1] - month_range[0] < 0: + # We have crossed to the next year. + months = [initials[x] for x in range(month_range[0], 13)] + months += [initials[x] for x in range(1, month_range[1] + 1)] + else: + months = [initials[x] for x in range(month_range[0], month_range[1] + 1)] + season = "".join(months) + title = f"{season} {variable}" + obs_map = LonLatMap.from_xarray_dataset( + dataset, reference_variable, f"{period} climatology", + year_range=[2003, 2018], month_range=month_range, + ) + obs_map.regrid_to_map(model_map) + + figure = chuck_radiation(model_map, obs_map, f"{title}") + output = Path(png_dir) / f"{title.lower().replace(' ', '-')}.png" + figure.save(output) + return output + + +if __name__ == "__main__": + RadiationAnalysisScript().run_analysis("model_catalog.json", ".", + "obs_catalog.json") diff --git a/user-analysis-scripts/freanalysis_radiation_atmos_av_mon/pyproject.toml b/user-analysis-scripts/freanalysis_radiation_atmos_av_mon/pyproject.toml new file mode 100644 index 0000000..e5a28a4 --- /dev/null +++ b/user-analysis-scripts/freanalysis_radiation_atmos_av_mon/pyproject.toml @@ -0,0 +1,28 @@ +[build-system] +requires = [ + "setuptools >= 40.9.0", +] +build-backend = "setuptools.build_meta" + +[project] +name = "freanalysis_radiation_atmos_av_mon" +version = "0.1" +dependencies = [ + "intake", + "intake-esm", +] +requires-python = ">= 3.6" +authors = [ + {name = "developers"}, +] +maintainers = [ + {name = "developer", email = "developer-email@address.domain"}, +] +description = "Radiative flux comparison to CERES" +readme = "README.md" +classifiers = [ + "Programming Language :: Python" +] + +[project.urls] +repository = "https://github.com/NOAA-GFDL/analysis-scripts.git" From e2cd82c8a000c2158df1d35b6031935eb99bd34b Mon Sep 17 00:00:00 2001 From: Raymond Menzel Date: Wed, 15 Jan 2025 10:38:13 -0500 Subject: [PATCH 3/3] little bit of cleanup --- .../analysis_scripts/plugins.py | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/core/analysis_scripts/analysis_scripts/plugins.py b/core/analysis_scripts/analysis_scripts/plugins.py index 36593a0..212df2a 100644 --- a/core/analysis_scripts/analysis_scripts/plugins.py +++ b/core/analysis_scripts/analysis_scripts/plugins.py @@ -26,13 +26,13 @@ def _find_plugin_class(module): for attribute in vars(module).values(): # Try to find a class that inherits from the AnalysisScript class. if inspect.isclass(attribute) and AnalysisScript in attribute.__bases__: - # Instantiate an object of this class. + # Return the class so an object can be instantiated from it later. return attribute raise UnknownPluginError("could not find class that inherts from AnalysisScripts") -sanity_counter = 0 # How much recursion is happening. -maximum_craziness = 100 # This is too much recursion. +_sanity_counter = 0 # How much recursion is happening. +_maximum_craziness = 100 # This is too much recursion. def _recursive_search(name, ispkg): @@ -49,10 +49,10 @@ def _recursive_search(name, ispkg): UnknownPluginError if no class is found. ValueError if there is too much recursion. """ - global sanity_counter - sanity_counter += 1 - if sanity_counter > maximum_craziness: - raise ValueError(f"recursion level {sanity_counter} too high.") + global _sanity_counter + _sanity_counter += 1 + if _sanity_counter > _maximum_craziness: + raise ValueError(f"recursion level {_sanity_counter} too high.") module = importlib.import_module(name) try: @@ -72,10 +72,11 @@ def _recursive_search(name, ispkg): # Dictionary of found plugins. -discovered_plugins = {} +_discovered_plugins = {} for finder, name, ispkg in pkgutil.iter_modules(): if name.startswith("freanalysis_") and ispkg: - discovered_plugins[name] = _recursive_search(name, True) + _sanity_counter = 0 + _discovered_plugins[name] = _recursive_search(name, True) def _plugin_object(name): @@ -92,14 +93,14 @@ def _plugin_object(name): UnknownPluginError if the input name is not in the disovered_plugins dictionary. """ try: - return discovered_plugins[name]() + return _discovered_plugins[name]() except KeyError: raise UnknownPluginError(f"could not find analysis script plugin {name}.") def available_plugins(): """Returns a list of plugin names.""" - return sorted(list(discovered_plugins.keys())) + return sorted(list(_discovered_plugins.keys())) def list_plugins():