Save progress - Prototype url search

Author: f-PLT

climateset/download/cmip6_downloader.py

@@ -1,16 +1,10 @@
-from pyesgf.search import SearchConnection
-
 from climateset.download.abstract_downloader import AbstractDownloader
 from climateset.download.constants.esgf import CMIP6
 from climateset.download.downloader_config import (
     CMIP6DownloaderConfig,
     create_cmip6_downloader_config_from_file,
 )
-from climateset.download.utils import (
-    download_model_variable,
-    get_upload_version,
-    handle_base_search_constraints,
-)
+from climateset.download.utils import search_and_download_esgf_model_single_var
 from climateset.utils import create_logger
 
 LOGGER = create_logger(__name__)
@@ -71,78 +65,19 @@ def download_from_model_single_var(
             preferred_version: data upload version, if 'latest', the newest version will get selected always
             default_grid_label: default gridding method in which the data is provided
         """
-        conn = SearchConnection(url=self.config.node_link, distrib=False)
-
-        facets = (
-            "project,experiment_id,source_id,variable,frequency,variant_label,variable, nominal_resolution, "
-            "version, grid_label, experiment_id"
-        )
-
-        self.logger.info("Using download_from_model_single_var() function")
-
-        ctx = conn.new_context(
-            project=project,
-            experiment_id=experiment,
-            source_id=model,
+        results_list = search_and_download_esgf_model_single_var(
+            model=model,
             variable=variable,
-            facets=facets,
+            experiment=experiment,
+            project=project,
+            default_frequency=default_frequency,
+            default_grid_label=default_grid_label,
+            preferred_version=preferred_version,
+            ensemble_members=self.config.ensemble_members,
+            max_ensemble_members=self.config.max_ensemble_members,
+            base_path=self.config.data_dir,
         )
-
-        ctx = handle_base_search_constraints(ctx, default_frequency, default_grid_label)
-
-        variants = list(ctx.facet_counts["variant_label"])
-
-        if len(variants) < 1:
-            self.logger.info(
-                "No items were found for this request. Please check on the esgf server if the combination of your "
-                "model/scenarios/variables exists."
-            )
-            raise ValueError(
-                f"Downloader did not find any items on esgf for your request with: Project {project}, "
-                f"Experiment {experiment}, Model {model}, Variable {variable}."
-            )
-
-        self.logger.info(f"Available variants : {variants}\n")
-        self.logger.info(f"Length : {len(variants)}")
-
-        # TODO refactor logic of if/else
-        if not self.config.ensemble_members:
-            if self.config.max_ensemble_members > len(variants):
-                self.logger.info("Less ensemble members available than maximum number desired. Including all variants.")
-                ensemble_member_final_list = variants
-            else:
-                self.logger.info(
-                    f"{len(variants)} ensemble members available than desired (max {self.config.max_ensemble_members}. "
-                    f"Choosing only the first {self.config.max_ensemble_members}.)."
-                )
-                ensemble_member_final_list = variants[: self.config.max_ensemble_members]
-        else:
-            self.logger.info(f"Desired list of ensemble members given: {self.config.ensemble_members}")
-            ensemble_member_final_list = list(set(variants) & set(self.config.ensemble_members))
-            if len(ensemble_member_final_list) == 0:
-                self.logger.info("WARNING: no overlap between available and desired ensemble members!")
-                self.logger.info("Skipping.")
-                return
-
-        for ensemble_member in ensemble_member_final_list:
-            self.logger.info(f"Ensembles member: {ensemble_member}")
-            ctx_ensemble = ctx.constrain(variant_label=ensemble_member)
-
-            version = get_upload_version(context=ctx, preferred_version=preferred_version)
-            if version:
-                ctx_ensemble = ctx_ensemble.constrain(version=version)
-
-            results = ctx_ensemble.search()
-
-            self.logger.info(f"Result len {len(results)}")
-
-            download_model_variable(
-                project=CMIP6,
-                model_id=model,
-                search_results=results,
-                variable=variable,
-                base_path=self.config.data_dir,
-            )
+        self.logger.info(f"Download results: {results_list}")
 
 
 def cmip6_download_from_config(config):

climateset/download/constants/__init__.py

@@ -0,0 +1,10 @@
+NODE_LINK_URLS = [
+    "https://esgf-node.llnl.gov/esg-search",
+    "https://esgf.ceda.ac.uk/esg-search",
+    "https://esgf-data.dkrz.de/esg-search",
+    "https://esgf-node.ipsl.upmc.fr/esg-search",
+    "https://esg-dn1.nsc.liu.se/esg-search",
+    "https://esgf.nci.org.au/esg-search",
+    "https://esgf.nccs.nasa.gov/esg-search",
+    "https://esgdata.gfdl.noaa.gov/esg-search",
+]

climateset/download/input4mips_downloader.py

@@ -1,16 +1,12 @@
-from pyesgf.search import SearchConnection
-
 from climateset.download.abstract_downloader import AbstractDownloader
 from climateset.download.constants.esgf import INPUT4MIPS
 from climateset.download.downloader_config import (
     Input4mipsDownloaderConfig,
     create_input4mips_downloader_config_from_file,
 )
 from climateset.download.utils import (
-    download_metadata_variable,
-    download_raw_input_variable,
-    get_upload_version,
-    handle_base_search_constraints,
+    search_and_download_esgf_biomass_single_var,
+    search_and_download_esgf_raw_single_var,
 )
 from climateset.utils import create_logger
 
@@ -67,38 +63,18 @@ def download_raw_input_single_var(
         """
         self.logger.info("Using download_raw_input_single_var() function")
 
-        facets = "project,frequency,variable,nominal_resolution,version,target_mip,grid_label"
-
         # Search context is sensitive to order and sequence, which is why
         # it's done in different steps instead of putting everything in `new_context`
-        conn = SearchConnection(url=self.config.node_link, distrib=False)
-        ctx = conn.new_context(
-            project=project,
+        results_list = search_and_download_esgf_raw_single_var(
             variable=variable,
+            project=project,
             institution_id=institution_id,
-            facets=facets,
+            default_grid_label=default_grid_label,
+            default_frequency=default_frequency,
+            preferred_version=preferred_version,
+            data_dir=self.config.data_dir,
         )
-        ctx = handle_base_search_constraints(ctx, default_frequency, default_grid_label)
-
-        mips_targets = list(ctx.facet_counts["target_mip"])
-        self.logger.info(f"Available target mips: {mips_targets}")
-
-        for target in mips_targets:
-            ctx_target = ctx.constrain(target_mip=target)
-            version = get_upload_version(context=ctx_target, preferred_version=preferred_version)
-            if version:
-                ctx_target = ctx_target.constrain(version=version)
-
-            results = ctx_target.search()
-            self.logger.info(f"Result len  {len(results)}")
-            if len(results) > 0:
-                download_raw_input_variable(
-                    project=INPUT4MIPS,
-                    institution_id=institution_id,
-                    search_results=results,
-                    variable=variable,
-                    base_path=self.config.data_dir,
-                )
+        self.logger.info(f"Download results: {results_list}")
 
     def download_meta_historic_biomassburning_single_var(
         self,
@@ -123,38 +99,20 @@ def download_meta_historic_biomassburning_single_var(
         variable_id = variable.replace("_", "-")
         variable_search = f"percentage_{variable_id.replace('-', '_').split('_')[-1]}"
         self.logger.info(variable, variable_id, institution_id)
-        facets = "nominal_resolution,version"
 
         # Search context is sensitive to order and sequence, which is why
         # it's done in different steps instead of putting everything in `new_context`
-        conn = SearchConnection(url=self.config.node_link, distrib=False)
-        ctx = conn.new_context(
-            project=project,
+        results = search_and_download_esgf_biomass_single_var(
             variable=variable_search,
             variable_id=variable_id,
+            project=project,
             institution_id=institution_id,
-            target_mip="CMIP",
-            facets=facets,
-        )
-        ctx = handle_base_search_constraints(ctx, default_frequency, default_grid_label)
-
-        version = get_upload_version(context=ctx, preferred_version=preferred_version)
-        if version:
-            ctx = ctx.constrain(version=version)
-
-        results = ctx.search()
-        self.logger.info(f"Result len  {len(results)}")
-
-        result_list = [r.file_context().search() for r in results]
-        self.logger.info(f"List of results :\n{result_list}")
-
-        download_metadata_variable(
-            project=INPUT4MIPS,
-            institution_id=institution_id,
-            search_results=results,
-            variable=variable,
+            default_grid_label=default_grid_label,
+            default_frequency=default_frequency,
+            preferred_version=preferred_version,
             base_path=self.config.data_dir,
         )
+        self.logger.info(f"Download results: {results}")
 
 
 def input4mips_download_from_config(config):