ecmwfplevels.py

Go to the documentation of this file.

#!/usr/bin/env python3
 
# *******************************************************
# *  International Data Centre                          *
# *  Comprehensive Nuclear Test Ban Treaty Organization *
# *  Vienna                                             *
# *  Austria                                            *
#                                                       *
# *  Don Morton (DM)                                        *
# *  Boreal Scientific Computing                        *
# *  Fairbanks, Alaska USA                              *
# *******************************************************
 
 
 
#-----------------------------------------------------------------------------
 
#import argparse
import datetime
import logging
import os
import uuid
 
import f90nml
 
import ehratm.defaults
import ehratm.mylogger
 
import nwpservice.wps.namelistwps
import nwpservice.wps.calcecmwfp
 
DEFAULTS = ehratm.defaults.Defaults()
 
LOGGER = ehratm.mylogger.getlogger()
 
 
class EcmwfPlevelsWorkflow(object):
 
 
    '''Class for managing preparation and execution of ecmwf plevels components
    '''
 
 
    def __init__(self,
                    wpswrf_distro_path=None,
                    working_rootdir=None,
                    ecmwf_coeffs_path=None,
                    start_time_dt=None,
                    stop_time_dt=None,
                    hours_intvl=None,
                    ungribbed_data_dict=None,
                    log_level=None):
 
        '''Initialises the class
 
        Brings in the ecmwfplevels, time and ungribbed_data dicts, to be used
        for making the request to nwpservice
 
        Parameters
        ----------
        wpswrf_distro_path : str
            (Optional) Full path to the WPS/WRF distribution to be used.
            This is assumed to have been installed in a way that is compatible
            with the nwpservice module.  If arg is not present, uses a default
            value
        working_rootdir : str
            (Optional) Full path to a directory (assumed to have already
            been created), to be used as scratch space for setting up and
            running this geogrid instance.  If arg is not present, uses a
            default value
        ecmwf_coeffs : str
            (Optional) Full path to ecmwf_coeffs file.  If arg is not present,
            uses a default value
        start_time_dt : datetime
            Time of first ungribbed sfc/ml file
        stop_time_dt : datetime
            Time of last ungribbed sfc/ml file
        hours_intvl : int
            Interval (in hours) between ungribbed sfc/ml files
        ungribbed_data_dict : dict
            Necessary information on the ecmwf_ml and ecmwf_sfc ungribbed
            files.  We expect dict with entries {'FG_NAME' : <path>, ...},
            and expect FG_NAME to be either DEFAULTS.ecmwf_sfc_ungrib_prefix()
            or DEFAULT.ecmwf_ml_ungrib_prefix
        log_level : int
            Python logging level (e.g. logging.INFO)
        '''
 
 
 
        if log_level:
            self._logging_level = log_level
        else:
            self._logging_level = DEFAULTS.log_level()
        LOGGER.setLevel(self._logging_level)
 
        # NOW I can use the logger
        LOGGER.debug('started')
 
        # Do some checks (need to expand on these over the years)
        # Essentially, would be good to control bail here if there's a problem
        # as opposed to further down
 
        if not wpswrf_distro_path:
            wpswrf_distro_path = DEFAULTS.wpswrf_distro_path()
        self._wpswrf_distro_path = wpswrf_distro_path
        LOGGER.debug('wpswrf_distro_path: %s' % self._wpswrf_distro_path)
        if not os.path.isdir(self._wpswrf_distro_path):
            raise FileNotFoundError('wpswrf_distro_path not found: %s' %
                                     self._wpswrf_distro_path)
 
        if not working_rootdir:
            working_rootdir = DEFAULTS.working_scratch_rootdir()
        self._working_rootdir = working_rootdir
        if not os.path.isdir(self._working_rootdir):
            raise FileNotFoundError('working_rootdir not found: %s' %
                                     self._working_rootdir)
        LOGGER.debug('working_rootdir: %s' % self._working_rootdir)
 
        # Check that ecmwf_coefs, either the arg or the default, is
        # accessible
        if ecmwf_coeffs_path:
            ecmwf_coeffs = ecmwf_coeffs_path
        else:
            ecmwf_coeffs = DEFAULTS.ecmwf_coeffs()
        if not os.path.isfile(ecmwf_coeffs):
            raise FileNotFoundError('ecmwf_coeffs: %s' % ecmwf_coeffs)
        self._ecmwf_coeffs_path = ecmwf_coeffs
 
 
        # Go through the ungribbed data, ensure we have the two types we want,
        # and ensure they have the timestamps we're expecting
        if not ungribbed_data_dict:
            raise ValueError("ungribbed_data_dict not defined")
 
        # First, we ensure that the expected SFC and ML types are in there
        sfc_prefix = DEFAULTS.ecmwf_sfc_ungrib_prefix()
        ml_prefix = DEFAULTS.ecmwf_ml_ungrib_prefix()
 
        try:
            ungribbed_data_keys = ungribbed_data_dict.keys()
        except:
            raise ValueError('Unable to retrieve ungribbed_data_keys')
        if sfc_prefix not in ungribbed_data_keys:
            raise ValueError('Missing key: %s' % sfc_prefix)
        if ml_prefix not in ungribbed_data_keys:
            raise ValueError('Missing key: %s' % ml_prefix)
 
        # Then, we iterate through expected timestamps and ensure that
        # the ungribbed files are there
        ml_path = ungribbed_data_dict[ml_prefix]
        sfc_path = ungribbed_data_dict[sfc_prefix]
        curr_dt = start_time_dt
        while curr_dt <= stop_time_dt:
            LOGGER.debug('curr_dt: %s' % curr_dt)
            curr_timestr = curr_dt.strftime('%Y-%m-%d_%H')
            LOGGER.debug('curr_timestr: %s' % curr_timestr)
 
            ml_filepath = os.path.join(ml_path, ml_prefix + ':' + curr_timestr)
            sfc_filepath = os.path.join(sfc_path, sfc_prefix + ':' + curr_timestr)
            if not os.path.isfile(ml_filepath):
                raise FileNotFoundError('Input not found: %s' % ml_filepath)
            if not os.path.isfile(sfc_filepath):
                raise FileNotFoundError('Input not found: %s' % sfc_filepath)
            curr_dt += datetime.timedelta(hours=hours_intvl)
 
        self._start_time_dt = start_time_dt
        self._stop_time_dt = stop_time_dt
        self._hours_intvl = hours_intvl
 
 
        self._ungribbed_data_dict = ungribbed_data_dict
 
        
 
 
 
 
 
    def run_ecmwfplevels(self, ecmwfplevels_output_stagedir=None):
 
        '''Set up and run, via nwpservice module, instance of ecmwfplevels
 
        Parameters
        ----------
        ecmwfplevels_output_stagedir : str
            (Optional) Full path to dir where ecmwfplevels files will be staged.
            If not specified, no staging will be done.  If specified, we
            assume the dir already exists.
 
        Returns
        -------
        return_dict : dict
            Dictionary with manifest of the ecmwfplevels files and the location
            (if applicable) of staged files
        '''
 
 
        LOGGER.debug('Start run_ecmwfplevels()')
 
 
        ecmwfp_status_dict = {}
 
 
        # Check that if the output stagedir is an arg, that it exists and is
        # writable
        if ecmwfplevels_output_stagedir:
            if os.path.isdir(ecmwfplevels_output_stagedir) and \
                os.access(ecmwfplevels_output_stagedir, os.W_OK):
                LOGGER.debug('Output stage dir: %s' % ecmwfplevels_output_stagedir)
            else:
                raise FileNotFoundError('Problem with output stage dir: %s' %
                                        ecmwfplevels_output_stagedir)
 
 
        # Create name for namelist.wps
        #namelist_wps_path = os.path.join(self._working_rootdir,
        #                str(uuid.uuid4()) + '_namelist.wps')
        # This needs to be a unique name, on the off chance (which might
        # occur in testing) that the working_rootdir has been used before
        # and already has this file, which would cause a collision
        #unique_nml_name = 'geogrid_namelist.wps-' + str(uuid.uuid4())[0:8]
        namelist_wps_path = os.path.join(self._working_rootdir,
                        'ecmwfplevels_namelist.wps')
        LOGGER.debug('namelist_wps_path: %s' % namelist_wps_path)
 
        # nwpservice object expects this dir to NOT exist - it will create it
        # itself.  Error will occur if it already exists
        wpswrf_rundir = os.path.join(self._working_rootdir,
                                        'ecmwfplevels_rundir')
 
        # Create namelist.wps
        # The share section
        start_time_str = self._start_time_dt.strftime('%Y%m%d%H%M%S')
        end_time_str = self._stop_time_dt.strftime('%Y%m%d%H%M%S')
        interval_secs = int(self._hours_intvl*3600)
        share_section_dict = {
                    'max_dom' : 1,
                    'start_date_list' : [start_time_str],
                    'end_date_list' : [end_time_str],
                    'interval_seconds' : interval_secs
        }
        LOGGER.debug('share_section_dict: %s' % share_section_dict)
        metgrid_section_dict = {
                        'fg_name' : [DEFAULTS.ecmwf_sfc_ungrib_prefix(),
                                    DEFAULTS.ecmwf_ml_ungrib_prefix()]
        }
        LOGGER.debug('metgrid_section_dict: %s' % metgrid_section_dict)
 
        section_data_dict = {'share' : share_section_dict,
                            'metgrid' : metgrid_section_dict}
        LOGGER.debug('section_data_dict: %s' % section_data_dict)
 
 
        myobj = nwpservice.wps.namelistwps.NamelistWpsWriter(
                destpath=namelist_wps_path,
                section_data_dict=section_data_dict,
                log_level=self._logging_level
                )
        LOGGER.debug('Writing namelist: %s' % namelist_wps_path)
        myobj.write()
 
 
        # Create the nwpservice object
        nwpecmwfp_obj = nwpservice.wps.calcecmwfp.CalcEcmwfPlevels(
                        wpswrf_distro_path=self._wpswrf_distro_path,
                        wpswrf_rundir=wpswrf_rundir,
                        ungribbed_data=self._ungribbed_data_dict,
                        ecmwf_coeffs_path=self._ecmwf_coeffs_path,
                        namelist_wps=namelist_wps_path,
                        output_dir=ecmwfplevels_output_stagedir,
                        log_level=self._logging_level
                        )
 
        nwpecmwfp_obj.setup()
        output_manifest = nwpecmwfp_obj.run()
        LOGGER.debug('output_manifest: %s' % output_manifest)
 
 
        staging_dir = None
 
        # If this wasn't successful (even if an output staging dir was
        # never specified), it should return False
        stage_success = nwpecmwfp_obj.stage_output(auxfiles=False)
        if stage_success:
            staging_dir = ecmwfplevels_output_stagedir
        else:
            staging_dir = None
 
 
        return_dict = {
                    'ecmwfp_output_manifest' : output_manifest,
                    'staging_dir' : staging_dir
        }
 
        LOGGER.debug('return_dict: %s' % return_dict)
 
        return return_dict