wrf.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.wrf.wrf
 
DEFAULTS = ehratm.defaults.Defaults()
 
LOGGER = ehratm.mylogger.getlogger()
 
class WrfWorkflow(object):
 
 
    '''Class for managing preparation and execution of wrf 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,
                    nests_defn_dict=None,
                    real_data_dict=None,
                    num_mpi_tasks=None,
                    log_level=None,
                    bypass_namelist_input=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 real 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
        nests_defn_dict : dict
            Dict of nest definitions (of use if generating namelist.input)
        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.
        real_data_dict : dict
            Necessary information on the real files.  We expect dict with
            entries {'path' : <path>, 'filelist' : [list of real files]}.
            This implies that the real files are all in the same src
            directory
        log_level : int
            Python logging level (e.g. logging.INFO)
        bypass_namelist_input : str
            This is an assumed-correct namelist.input to be used for
            execution of real.  Its presence bypasses the normal process
            of creating a namelist.input, and puts complete trust in the
            correctness of the provided namelist.input.  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
        -------
        These are examples of the real_data_dict
 
        ::
 
            real_data_dict = {
                'path' : '/path/to/dir/with/real/files'
                'filelist' : ['wrfbdy_d01',
                            'wrfinput_d01', ,
                            'wrfinput_d02', ,
                            ]
                                }
        '''
 
        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_namelist_input arg.  If present, verify good
        # path.  And, if present, some of the following checks will be
        # ignored
        self._bypass_namelist_input = None
        if bypass_namelist_input:
            if not os.path.isfile(bypass_namelist_input):
                raise FileNotFoundError('bypass_namelist_input not found: %s' %
                                            bypass_namelist_input)
            self._bypass_namelist_input = bypass_namelist_input
            LOGGER.debug('bypass_namelist_input: %s' %
                        self._bypass_namelist_input)
 
        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 expected real files are available (use real_data_dict)
        real_dir = real_data_dict['path_to_files']
        real_files = real_data_dict['files_list']
        for fname in real_files:
            fpath = os.path.join(real_dir, fname)
            if not os.path.isfile(fpath):
                raise FileNotFoundError('fpath: %s' % fpath)
 
 
        #-----  Don't deal with these possibly empty variables if we're
        #-----  using the bypass_namelist_input
        if not self._bypass_namelist_input:
            # 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
 
            # Check that we have expected nests.  We're going to assume
            # that there's only one outer nest for wrfbdy, and we're
            # going to assume there's one wrfinput for each nest.  This
            # may not always be the case!
            LOGGER.debug('NUM_NESTS: %d, real_files: %s' % (self._num_nests, real_files))
            for i in range(self._num_nests):
                expected_fname = 'wrfinput_d%02d' % (i+1)
                if expected_fname not in real_files:
                    raise ValueError('Did not find nest: %s' % expected_fname)
 
            # Check that we have the nest definitions
            if nests_defn_dict:
                self._nests_defn_dict = nests_defn_dict
            else:
                raise ValueError('Did not find nests_defn_dict...')
 
        # 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_wrf()
        # method
 
 
        if self._bypass_namelist_input:
            # We won't use these if bypass_namelist_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._hours_intvl = None
            self._num_nests = None
 
        else:
            self._start_time_dt = start_time_dt
            self._stop_time_dt = stop_time_dt
            self._hours_intvl = hours_intvl
 
 
        self._real_data_dict = real_data_dict
 
 
    def run_wrf(self, wrf_output_stagedir=None):
 
 
        '''Set up and run, via nwpservice module, instance of wrf
 
        Parameters
        ----------
        wrf_output_stagedir : str
            (Optional) Full path to dir where wrf 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 wrf output files and the location
            (if applicable) of staged files
        '''
 
 
        LOGGER.debug('Starting run_wrf()...')
 
        # Check the output stagedir
 
        # If we're using the bypass_namelist_input, then there's
        # just a little prep necessary.  Otherwise, we will need to
        # create the namelist based on a number of parameters, including
        # the domain_defn and the dates/times, etc.
        if self._bypass_namelist_input:
            # Test for its presence
            if not os.path.isfile(self._bypass_namelist_input):
                raise FileNotFoundError(self._bypass_namelist_input)
            namelist_input_path = self._bypass_namelist_input
 
            # Just in case they're needed later, get the num_nests and
            # hours_intvl values from the namelist (They were set to None
            # in __init__() for the bypass namelist scenario)
            nml = f90nml.read(self._bypass_namelist_input)
            self._hours_intvl = int(
                                nml['time_control']['interval_seconds'] / 3600)
            self._num_nests = nml['domains']['max_dom']
 
        else:
            # This is where eventually we'll add in the namelist.input
            # creation
            namelist_input_path = None
            raise NotImplementedError(
                            'namelist.input creation not yet supported')
 
 
        # Create nwpservice object, then setup and run
        domainpath = os.path.join(self._working_rootdir, 'wrf_rundir')
        LOGGER.debug('domainpath: %s' % domainpath)
 
 
        nwpwrf_obj = nwpservice.wrf.wrf.Wrf(
            wpswrf_distro_path=self._wpswrf_distro_path,
            wpswrf_rundir=domainpath,
            realdatadir=self._real_data_dict['path_to_files'],
            namelist_input=namelist_input_path,
            output_dir=wrf_output_stagedir,
            numpes=self._num_mpi_tasks,
            mpirun_path=self._mpirunpath,
            log_level=self._logging_level
        )
 
        nwpwrf_obj.setup()
        output_manifest = nwpwrf_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 = nwpwrf_obj.stage_output(auxfiles=True)
        if stage_success:
            staging_dir = wrf_output_stagedir
        else:
            staging_dir = None
 
 
 
 
        return_dict = {
                'wrf_output_manifest' : output_manifest,
                'staging_dir' : staging_dir,
                'hours_intvl' : self._hours_intvl
        }
        LOGGER.debug('return_dict: %s' % return_dict)
 
 
        return return_dict