metgrid.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.metgrid
 
DEFAULTS = ehratm.defaults.Defaults()
 
LOGGER = ehratm.mylogger.getlogger()
 
 
class MetgridWorkflow(object):
 
 
    '''Class for managing preparation and execution of metgrid components
    '''
 
 
 
    def __init__(self,
                    wpswrf_distro_path=None,
                    working_rootdir=None,
                    start_time_dt=None,
                    stop_time_dt=None,
                    hours_intvl=None,
                    num_nests=None,
                    ungribbed_data_dict=None,
                    geogrid_data_dict=None,
                    custom_metgrid_tbl_path=None,
                    num_mpi_tasks=None,
                    log_level=None
                    ):
 
        '''Check and initialise the class
 
 
        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 metgrid instance.  If arg is not present, uses a
            default value
        start_time_dt : datetime
            Time of first ungribbed file(s)
        stop_time_dt : datetime
            Time of last ungribbed file(s)
        hours_intvl : int
            Interval (in hours) between ungribbed files
        num_nests : int
            Number of domain nests represented in geogrid data, and in
            the metgrid data that will be produced
        num_mpi_tasks : int
            (Optional) Number of MPI tasks to use for running this instance
            of metgrid.  If arg is not present, non-MPI execution is assumed.
        ungribbed_data_dict : dict
            Necessary information on the ungribbed files.  We expect dict with
            entries {'FG_NAME' : <path>, ...}.  Note that the 'FG_NAME' is
            not expected to adhere to any kind of workflow convention.  This
            method simply uses what it's given
        geogrid_data_dict : dict
            Necessary information on the geogrid files.  We expect dict with
            entries {'path' : <path>, 'filelist' : [list of geogrid files]}.
            This implies that the geogrid files are all in the same src
            directory
        custom_metgrid_tbl_path : str
            (Optional) Path to a METGRID.TBL to be used for this run.  If arg
            is not present, then we assume that no custom file is needed.
        log_level : int
            Python logging level (e.g. logging.INFO)
 
 
        Example
        -------
        These are examples of the geogrid_data_dict  ungribbed_data_dict
 
        ::
 
            geogrid_data_dict = {
                'path' : '/path/to/dir/with/geogrid/files'
                'filelist' : ['geo_em.d01.nc', 'geo_em.d02.nc']
                                }
 
            ungribbed_data_dict = {
                'ECMWF_ML' : '/path/to/dir/with/ecmwf_ml/ungribbed_files',
                'ECMWF_SFC' : '/path/to/dir/with/ecmwf_sfc/ungribbed_files',
                'ECMWF_PRES' : '/path/to/dir/with/pres/ungribbed_files'
                                }
        '''
 
 
        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 do a controlled bail here if there's
        # a problem as opposed to doing it 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)
 
 
        # Basic test of num_nests
        if not num_nests:
            raise ValueError("Missing num_nests arg")
        if num_nests <= 0:
            raise ValueError("num_nests: %d" % num_nests)
        self._num_nests = num_nests
 
        # If custom_metgrid_tbl_path is an arg, then ensure the path is valid.
        # If not defined, then leave it that way.  We will assume that a
        # calling process makes the decision about whether the custom
        # version is used or not (typically dependent on whether we're using
        # ECMWF ERA data)
        if custom_metgrid_tbl_path:
            if not os.path.isfile(custom_metgrid_tbl_path):
                raise FileNotFoundError(
                        'Path not found: %s' % custom_metgrid_tbl_path)
            self._custom_metgrid_tbl_path = custom_metgrid_tbl_path
        else:
            self._custom_metgrid_tbl_path = None
 
        # Check that expected geogrid files are available (use geogrid_data_dict)
        geogrid_dir = geogrid_data_dict['path_to_files']
        geogrid_files = geogrid_data_dict['files_list']
        if len(geogrid_files) != self._num_nests:
            raise ValueError('Wrong number geogrid files: %d' % len(geogrid_files))
        for fname in geogrid_files:
            fpath = os.path.join(geogrid_dir, fname)
            if not os.path.isfile(fpath):
                raise FileNotFoundError('fpath: %s' % fpath)
 
        # Check that expected ungribbed files are available
        # (use ungribbed_data_dict)
        ungrib_fg_keys = list(ungribbed_data_dict.keys())
        for fg_key in ungrib_fg_keys:
            ug_path = ungribbed_data_dict[fg_key]
            LOGGER.debug('ug_path: %s' % ug_path)
            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)
                fname = fg_key + ':' + curr_timestr
                expected_path = os.path.join(ug_path, fname)
                if not os.path.isfile(expected_path):
                    raise FileNotFoundError('expected_path: %s' % expected_path)
                curr_dt += datetime.timedelta(hours=hours_intvl)
 
        # Use specified number of MPI tasks, or default to 0, which will imply
        # non-MPI execution
        self._mpirunpath = None  # Just give it a namespace for now
        if num_mpi_tasks:
            # Ensure valid number, otherwise raise exception.  If the value
            # looks good, check that the default mpirun executable is
            if 1 <= num_mpi_tasks <= DEFAULTS.max_mpi_tasks():
                # Go ahead and check the mpirun path
                mpirunpath = DEFAULTS.mpirun_path()
                if os.path.isfile(mpirunpath) and \
                            os.access(mpirunpath, os.X_OK):
                    self._num_mpi_tasks = num_mpi_tasks
                    self._mpirunpath = mpirunpath
                else:
                    raise FileNotFoundError('mpirun not executable: %s' %
                                            mpirunpath)
            else:
                raise ValueError('Bad num_mpi_tasks value: %d' % num_mpi_tasks)
        else:
            # This default value of 0 will denote non-MPI execution when
            # invoking the geogrid service
            self._num_mpi_tasks = 0
        LOGGER.debug('num_mpi_tasks: %d' % self._num_mpi_tasks)
 
 
        # If we got through all this, save the time info for the run_metgrid()
        # method
        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
        self._geogrid_data_dict = geogrid_data_dict
 
    def run_metgrid(self, metgrid_output_stagedir=None):
 
 
        '''Set up and run, via nwpservice module, instance of metgrid
 
        Parameters
        ----------
        metgrid_output_stagedir : str
            (Optional) Full path to dir where metgrid 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 metgrid files and the location
            (if applicable) of staged files
        '''
 
 
        LOGGER.debug('Starting run_metgrid()...')
 
        # Check the output stagedir
 
        # Create namelist.wps
        namelist_wps_path = os.path.join(self._working_rootdir, 'namelist.wps')
        LOGGER.debug('namelist_wps_path: %s' % namelist_wps_path)
 
        # 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)
 
        # Create the list of fg_name for metgrid group
        fg_name_list = []
        for ugkey in self._ungribbed_data_dict.keys():
            fg_name_list.append(ugkey)
 
        # The full section_data_dict - note that start_date_list and
        # end_date_list have to have a "num_nests" number of elements
        # which we assume are identical.  metgrid may not complain, but
        # may not create all nests if we neglect this
        section_data_dict = {
                    'share' : {
                        'max_dom' : self._num_nests,
                        'start_date_list' : [start_time_str]*self._num_nests,
                        'end_date_list' : [end_time_str]*self._num_nests,
                        'interval_seconds' : interval_secs
                    },
                    'metgrid' : {
                        'fg_name' : fg_name_list
                    }
        }
 
        myobj = nwpservice.wps.namelistwps.NamelistWpsWriter(
                destpath=namelist_wps_path,
                section_data_dict=section_data_dict
        )
        LOGGER.debug('Writing namelist: %s' % namelist_wps_path)
        myobj.write()
 
        # Create nwpservice object, then setup and run
        domainpath = os.path.join(self._working_rootdir, 'metgrid_rundir')
        LOGGER.debug('domainpath: %s' % domainpath)
 
 
        nwpmetgrid_obj = nwpservice.wps.metgrid.Metgrid(
            wpswrf_distro_path=self._wpswrf_distro_path,
            wpswrf_rundir=domainpath,
            ungribbed_data=self._ungribbed_data_dict,
            geogriddatadir=self._geogrid_data_dict['path_to_files'],
            metgridtbl_userdef=self._custom_metgrid_tbl_path,
            namelist_wps=namelist_wps_path,
            output_dir=metgrid_output_stagedir,
            numpes=self._num_mpi_tasks,
            mpirun_path=self._mpirunpath,
            log_level=self._logging_level
        )
 
        nwpmetgrid_obj.setup()
        output_manifest = nwpmetgrid_obj.run()
        LOGGER.debug('output_manifest: %s' % output_manifest)
 
        # If this wasn't successful (even if an output staging dir was
        # never specified), it should return false
        stage_success = nwpmetgrid_obj.stage_output(auxfiles=True)
        if stage_success:
            staging_dir = metgrid_output_stagedir
        else:
            staging_dir = None
 
        return_dict = {
                'metgrid_output_manifest' : output_manifest,
                'staging_dir' : staging_dir,
                'hours_intvl' : self._hours_intvl
        }
        LOGGER.debug('return_dict: %s' % return_dict)
 
 
 
        #return output_manifest
        return return_dict