geogrid.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.geogrid
 
DEFAULTS = ehratm.defaults.Defaults()
 
LOGGER = ehratm.mylogger.getlogger()
 
class GeogridWorkflow(object):
 
 
    '''Class for managing preparation and execution of geogrid components
    '''
 
 
    def __init__(self,
                    wpswrf_distro_path=None,
                    projection_dict=None,
                    nest_dict=None,
                    geog_data_path=None,
                    working_rootdir=None,
                    num_mpi_tasks=None,
                    log_level=None):
 
        '''Initialises the class
 
        Brings in the projection, nest and WPS geog_data path information.
 
        Assembles into a single dict ready for running
 
        Note that at some later time I should probably want to add in some
        verifications of projection, nest, etc., for its suitability with
        geogrid
 
        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
        projection_dict : dict
            Collection of projection parameters for the domain
        nest_dict : dict
            Collection of WPS-style nesting parameters
        geog_data_path : str
            (Optional) Path to the geog files used by WPS.  This might be a
            relative path, and will ultimately be inserted into namelist.wps
            as written.  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
        num_mpi_tasks : int
            (Optional) Number of MPI tasks to use for running this instance
            of geogrid.  If arg is not present, non-MPI execution is assumed.
        log_level : int
            Python logging level (e.g. logging.INFO)
 
 
        Example
        -------
        These are examples of the projection and nest dictionaries (one- and
        two-nest)
 
        ::
 
            projection_dict = {
                'map_proj' : 'lambert',
                'ref_lat' : 50.0,
                'ref_lon' : 5.0,
                'truelat1' : 60.0,
                'truelat2' : 30.0,
                'stand_lon' : 5.0
                                }
 
                onenest_dict = {
                'parent_id' : [1],
                'parent_grid_ratio' : [1],
                'i_parent_start' : [1],
                'j_parent_start' : [1],
                'e_we' : [140],
                'e_sn' : [125],
                'geog_data_res' : ['10m'],
                'dx' : 30000,
                'dy' : 30000,
                        }
 
                twonest_dict = {
                'parent_id' : [1, 1],
                'parent_grid_ratio' : [1, 3],
                'i_parent_start' : [1, 38],
                'j_parent_start' : [1, 77],
                'e_we' : [140, 232],
                'e_sn' : [125, 214],
                'geog_data_res' : ['10m', '30s'],
                'dx' : 30000,
                'dy' : 30000,
                        }
 
 
        '''
 
 
        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 geog_data_path:
            geog_data_path = DEFAULTS.wps_geog_data_dir()
        self._wps_geog_data_dir = geog_data_path
 
        
 
        #if not os.path.isdir(self._wps_geog_data_dir):
        #    raise FileNotFoundError('wps_geog_data_dir not found: %s' %
        #                             self._wps_geog_data_dir)
        LOGGER.debug('wps_geog_data_dir: %s' % self._wps_geog_data_dir)
 
        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)
 
 
        # Use specified number of MPI tasks, or default to 0, which will imply
        # non-MPI execution
        if num_mpi_tasks:
            # Ensure valid number, otherwise raise exception.  If the value
            # looks good, check that the default mpirun executable is
            # accessible
            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
                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)
 
        self._geogrid_def_dict = { 'share' : {}, 'geogrid' : {} }
 
 
 
        # Create the share section, finding number of domains by looking at
        # nest_dict
        try:
            max_dom = len(nest_dict['i_parent_start'])
        except:
            raise ValueError('Problem getting max_dom value')
        # Make sure max_dom is a reasonable value
        if not 0 < max_dom <= DEFAULTS.max_dom_supported():
            raise ValueError('Bad max_dom value: %d' % max_dom)
        self._geogrid_def_dict['share']['max_dom'] = max_dom
 
        # Add in the projection, nest and geog data path
        geogrid_sect = nest_dict
        geogrid_sect['projection'] = projection_dict
        geogrid_sect['geog_data_path'] = geog_data_path
 
 
        # Put it all together
        self._geogrid_def_dict['geogrid'] = geogrid_sect
 
        LOGGER.debug('self._geogrid_def_dict: %s' % self._geogrid_def_dict)
 
 
 
 
    def run_geogrid(self, geogrid_output_stagedir=None):
 
        '''Set up and run, via nwpservice module, instance of geogrid
 
        Parameters
        ----------
        geogrid_output_stagedir : str
            (Optional) Full path to dir where geogrid 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 geogrid files and the location
            (if applicable) of staged files
        '''
 
 
        LOGGER.debug('Start _run_geogrid()')
 
 
        # Create name for a tempdir for the run
        #wpswrf_rundir = os.path.join(self._working_rootdir,
        #                    'geogrid_rundir_' + str(uuid.uuid4()))
        #wpswrf_rundir = os.path.join(self._working_rootdir, 'geogridrun')
        #LOGGER.debug('wpswrf_rundir: %s' % wpswrf_rundir)
 
 
        # 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,
                        'geogrid_namelist.wps')
        LOGGER.debug('namelist_wps_path: %s' % namelist_wps_path)
 
 
        # Create namelist.wps
        myobj = nwpservice.wps.namelistwps.NamelistWpsWriter(
                destpath=namelist_wps_path,
                section_data_dict=self._geogrid_def_dict
        )
        LOGGER.debug('Writing namelist: %s' % namelist_wps_path)
        myobj.write()
 
        # Init geogrid component, run it, stage the output
        domainpath = os.path.join(self._working_rootdir,
                        str(uuid.uuid4()) + '_geogridrun')
        domainpath = os.path.join(self._working_rootdir, 'geogrid_rundir')
 
        nwpgeogrid_args = {
                    'wpswrf_distro_path' : self._wpswrf_distro_path,
                    'wpswrf_rundir' : domainpath,
                    'namelist_wps' : namelist_wps_path,
                    'output_dir' : geogrid_output_stagedir,
                    'log_level' : self._logging_level
                    }
 
        # Add MPI args if appropriate
        if self._num_mpi_tasks > 0:
            nwpgeogrid_args['numpes'] = self._num_mpi_tasks
            nwpgeogrid_args['mpirun_path'] = DEFAULTS.mpirun_path()
 
        geogrid_obj = nwpservice.wps.geogrid.Geogrid(**nwpgeogrid_args)
        """
        geogrid_obj = nwpservice.wps.geogrid.Geogrid(
                    wpswrf_distro_path=self._wpswrf_distro_path,
                    wpswrf_rundir=domainpath,
                    namelist_wps=namelist_wps_path,
                    output_dir=geogrid_output_stagedir,
                    log_level=DEFAULT_LOG_LEVEL
        )
        """
 
        geogrid_obj.setup()
        output_manifest = geogrid_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 = geogrid_obj.stage_output(auxfiles=False)
        if stage_success:
            staging_dir = geogrid_output_stagedir
        else:
            staging_dir = None
 
        return_dict = {
                    'geogrid_output_manifest' : output_manifest,
                    'staging_dir' : staging_dir
        }
 
        LOGGER.debug('return_dict: %s' % return_dict)
        return return_dict