flexwrf.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.flexwrf.flexwrf
 
DEFAULTS = ehratm.defaults.Defaults()
 
LOGGER = ehratm.mylogger.getlogger()
 
 
class FlexwrfWorkflow(object):
 
 
    '''Class for managing preparation and execution of flexwrf components
    '''
 
 
 
    def __init__(self,
                    flexwrf_distro_path=None,
                    working_rootdir=None,
                    start_time_dt=None,
                    stop_time_dt=None,
                    outgrid_defn_dict=None,
                    flexwrf_specs_dict=None,
                    wrfout_hours_intvl=None,
                    wrfout_num_nests=None,
                    wrfout_data_dict_list=None,
                    num_mpi_tasks=None,
                    num_openmp_threads=None,
                    log_level=None,
                    bypass_flexp_input=None,
                    ):
 
        '''Check and initialise the class
 
 
        Parameters
        ----------
        flexwrf_distro_path : str
            (Optional) Full path to the FlexWRF 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 flexwrf instance.  If arg is not present, uses a
            default value
        start_time_dt : datetime
            Time of first wrfout file(s), start time of flexwrf sim
        stop_time_dt : datetime
            Time of last wrfout file(s), end time of flexwrf sim
        outgrid_defn_dict : dict
            Dict of outgrid definitions (of use if generating flexwrf namelist)
        flexwrf_specs_dict : dict
            For future use when the user input file will be generated by
            this routine (for now, it's "bypassed" with a complete and correct
            flexp_input)
        wrfout_hours_intvl : int
            Number of expected hours between wrfout files.  For now, it is
            required, even when using bypass_flexp_input
        wrfout_num_nests : int
            Number of nests represented by wrfout files.  For now, it is
            required, even when using bypass_flexp_input
        wrfout_data_dict_list : list
            Necessary information on the wrfout files.  We expect
            list of dicts, one for each nest, with entries {'nest' : n,
            'path' : <path>, 'filelist' : [list of wrfout files]}.  At
            this level, it's OK if the path is the same for different nests
            (e.g. they might all be in the same wrf rundir)
        num_mpi_tasks : int
            (Optional) Number of MPI tasks to use for running this instance
            of flexwrf.  If arg is not present, non-MPI execution is assumed.
        num_openmp_threads : int
            (Optional) Number of OpenMP threads to use for running this
            instance of flexwrf.  If arg is not present, non-OpenMP execution
            is assumed.
        log_level : int
            Python logging level (e.g. logging.INFO)
        bypass_flexp_input : str
            This is an assumed-correct flexwrf namelist to be used for
            execution of flexwrf.  Its presence bypasses the normal process
            of creating a namelist, and puts complete trust in the
            correctness of the provided namelist  If this argument
            is present, then any values for start_time_dt, stop_time_dt,
            hours_intvl and num_nests are ignored.  This is meant primarily
            for devtest operations, and probably shouldn't be used for
            normal workflows.
 
 
        Example
        -------
        #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
        Change this for assumption of single dir
 
        These are examples of the wrfout_data_dict_list
 
        ::
 
            wrfout_data_dict_list = [
                {
                'nest' : 1,
                'path' : '/path/to/dir/with/wrfout_d01/files'
                'filelist' : ['wrfout_d01_2014-01-24_00:00:00', ...]
                },
                {
                'nest' : 2,
                'path' : '/path/to/dir/with/wrfout_d02/files'
                'filelist' : ['wrfout_d02_2014-01-24_00:00:00', ...]
                }
            ]
        '''
 
        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
 
 
        # Check bypass_flexp_input arg.  If present, verify good
        # path.  And, if present, some of the following checks will be
        # ignored
        self._bypass_flexp_input = None
        if bypass_flexp_input:
            if not os.path.isfile(bypass_flexp_input):
                raise FileNotFoundError('bypass_flexp_input not found: %s' %
                                            bypass_flexp_input)
            self._bypass_flexp_input = bypass_flexp_input
            LOGGER.debug('bypass_flexp_input: %s' %
                        self._bypass_flexp_input)
 
 
        # Check for path to FlexWRF distribution, and try to use default
        # if it wasn't passed in as an arg
        if not flexwrf_distro_path:
            flexwrf_distro_path = DEFAULTS.flexwrf_distro_path()
        self._flexwrf_distro_path = flexwrf_distro_path
        LOGGER.debug('flexwrf_distro_path: %s' % self._flexwrf_distro_path)
        if not os.path.isdir(self._flexwrf_distro_path):
            raise FileNotFoundError('flexwrf_distro_path not found: %s' %
                                    self._flexwrf_distro_path)
 
        # Check for the working root dir, and try to use a default
        # value if it wasn't passed in as an arg
        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 wrfout_num_nests
        if not wrfout_num_nests:
            raise ValueError("Missing wrfout_num_nests arg")
        if wrfout_num_nests <= 0:
            raise ValueError("wrfout_num_nests: %d" % wrfout_num_nests)
        self._wrfout_num_nests = wrfout_num_nests
        LOGGER.debug('wrfout_num_nests: %d' % self._wrfout_num_nests)
 
        # Check that wrfout_data_dict_list has correct number of expected
        # nests
        if len(wrfout_data_dict_list) != wrfout_num_nests:
            raise ValueError(
                'wrfout_data_dict_list unexpected number of nests: %d' %
                len(wrfout_data_dict_list))
 
        # Check that the files listed in wrfout_data_dict_list are
        # all accessible
        for wrfout_nest_dict in wrfout_data_dict_list:
 
            wrfout_dir = wrfout_nest_dict['path']
            wrfout_files = wrfout_nest_dict['filelist']
 
            for fname in wrfout_files:
                fpath = os.path.join(wrfout_dir, fname)
                if not os.path.isfile(fpath):
                    raise FileNotFoundError('fpath: %s' % fpath)
 
 
 
        # Some stuff to add later
        #
        # If we're not using bypass_flexp_input, make sure we have
        # dicts for various parameters, OUTGRID, etc.
        #
        # Consider checking for wrfout_hours_intvl?  I'm not sure yet.
 
 
 
        # 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 num_openmp_threads:
            # Ensure valid number, otherwise raise exception.
            if 1 <= num_openmp_threads <= DEFAULTS.max_openmp_threads():
                pass
            else:
                raise ValueError('Bad num_openmp_threads value: %d' %
                                    num_openmp_threadss)
        else:
            # This default value of 0 will denote non-OpenMP execution when
            # invoking the flexwrf service
            self._num_openmp_threads = 0
        LOGGER.debug('num_openmp_threads: %d' % self._num_openmp_threads)
 
        # If we got through all this, save the time info for the run_wrf()
        # method
 
 
        if self._bypass_flexp_input:
            # We won't use these if bypass_flexp_input is being used, so
            # set them to None so that other code won't mistakenly try to use
            # them
            self._start_time_dt = None
            self._stop_time_dt = None
            self._wrfout_hours_intvl = None
 
        else:
            self._start_time_dt = start_time_dt
            self._stop_time_dt = stop_time_dt
            self._hours_intvl = hours_intvl
 
        self._wrfout_data_dict_list = wrfout_data_dict_list
 
 
    def run_flexwrf(self, flexwrf_output_stagedir=None):
 
 
        '''Set up and run, via nwpservice module, instance of flexwrf
 
        Parameters
        ----------
        flexwrf_output_stagedir : str
            (Optional) Full path to dir where flexwrf output 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 flexwrf output files and the location
            (if applicable) of staged files
        '''
 
 
        LOGGER.debug('Starting run_flexwrf()...')
 
        # Check the output stagedir
 
        # If we're using the bypass_flexp_input, then there's
        # just a little prep necessary.  Otherwise, we will need to
        # create the flexp_input based on a number of parameters, including
        # the domain_defn and the dates/times, etc.
        if self._bypass_flexp_input:
            # Test for its presence
            if not os.path.isfile(self._bypass_flexp_input):
                raise FileNotFoundError(self._bypass_flexp_input)
            flexp_input_path = self._bypass_flexp_input
 
        else:
            # This is where eventually we'll add in creation of flexp_input
            flexp_input_path = None
            raise NotImplementedError(
                            'flexp_input creation not yet supported')
 
 
        # Create nwpservice object, then setup and run
        domainpath = os.path.join(self._working_rootdir, 'flexwrf_rundir')
        LOGGER.debug('domainpath: %s' % domainpath)
 
 
        # The current specifications for metfiles_path arg to the nwpservice
        # flexwrf is that it is a list of paths, one for each wrfout nest.
        # Every wrfout file in that path will end up being included in
        # the flexwrf setup, whether it's needed or not.  So, in this case,
        # since file accessibility was already checked in __init__(), we
        # will, for each nest, create a wrfout_d0* dir in the run directory,
        # make symlinks to each of the wrfout files in that nest (using
        # wrfout_data_dict_list entries for source), then add the path of
        # this newly created dir to our metfiles_path list
 
        metfiles_path_list = []
        #LOGGER.debug('wrfout_data_dict_list: %s' % self._wrfout_data_dict_list)
        for i in range(self._wrfout_num_nests):
            # Although index i "should" correspond to dict entry nest i,
            # I'm going to not make that assumption, and try to be a little
            # more flexible (and complicated)
            try:
                nest_idx = self._wrfout_data_dict_list.index(
                    next(
                        filter(lambda n: n.get('nest') == (i+1),
                                self._wrfout_data_dict_list)
                        )
                )
            except:
                raise ValueError('Failed to find nest index: %d' % (i+1))
            LOGGER.debug('nest_idx: %d' % nest_idx)
 
            # Now we know we have the correct idx for this nest, so extract
            # the path
            nest_path = self._wrfout_data_dict_list[nest_idx]['path']
            nest_files = self._wrfout_data_dict_list[nest_idx]['filelist']
            LOGGER.debug('nest_path: %s' % nest_path)
            LOGGER.debug('nest_files: %s' % nest_files)
 
 
            # Create the wrfout_d0* dir within the rundir, and within it,
            # make links to the file names just extracted
            dirname = 'wrfout_d%02d' % (i+1)
            dest_path = os.path.join(self._working_rootdir, dirname)
            os.mkdir(dest_path, 0o755)
            for wrfout_file in nest_files:
                src_fpath = os.path.join(nest_path, wrfout_file)
                dest_fpath = os.path.join(dest_path, wrfout_file)
                os.symlink(src_fpath, dest_fpath)
 
            metfiles_path_list.append(dest_path)
 
        nwpflexwrf_obj = nwpservice.flexwrf.flexwrf.FlexWrf(
            flexwrf_distro_path=self._flexwrf_distro_path,
            flexwrf_rundir=domainpath,
            flexp_input_txt=self._bypass_flexp_input,
            metfilespath_list=metfiles_path_list,
            output_dir=flexwrf_output_stagedir,
            numpes=self._num_mpi_tasks,
            mpirun_path=self._mpirunpath,
            log_level=self._logging_level
        )
 
        nwpflexwrf_obj.setup()
        output_manifest = nwpflexwrf_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 = nwpflexwrf_obj.stage_output(auxfiles=True)
        if stage_success:
            staging_dir = flexwrf_output_stagedir
        else:
            staging_dir = None
 
        return_dict = {
                'flexwrf_output_manifest' : output_manifest,
                'staging_dir' : staging_dir,
                'hours_intvl' : self._wrfout_hours_intvl
        }
        LOGGER.debug('return_dict: %s' % return_dict)
 
        return return_dict