srm.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 fnmatch
import glob
import logging
import os
import shutil
import subprocess
import uuid
 
import f90nml
 
import ehratm.defaults
import ehratm.mylogger
 
DEFAULTS = ehratm.defaults.Defaults()
 
LOGGER = ehratm.mylogger.getlogger()
 
class SrmWorkflow(object):
 
    '''Class for managing preparation and execution of srm component
 
    This component currently uses the standard ATM Pipeline flexpart2ctbt
    for the conversion.  This means that, as-is, it will not process
    an OUTGRID nest, but only the top-level nest.
    '''
 
    def __init__(self,
                    srm_exec_path=None,
                    working_rootdir=None,
                    flexwrfout_dir=None,
                    levels_list=[1],
                    log_level=None
                    ):
 
        '''Check and initialise the class
 
        Parameters
        ----------
        srm_exec_path : str
            (Optional) Full path to the SRM executable.  This is assumed
            to have been installed in a way that is compatible with the
            ehratm 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 srm instance.  If arg is not present, uses a
            default value.  A 'srm_rundir' subdir will be created in here
        flexwrfout_dir : str
            Full path to a directory containing the FlexWRF binary output files,
            including the header file
        levels_list : list
            (Optional) List of integers designating the vertical levels (indexed from 1)
            to produce SRM files for.  If arg is not present, assumes [1]
        log_level : int
            Python logging level (e.g. logging.INFO)
 
        Example
        -------
 
        '''
 
        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')
 
        self._srm_exec_path = None
        if srm_exec_path:
            if not os.path.isfile(srm_exec_path):
                raise FileNotFoundError('srm_exec_path not found: %s' %
                                        srm_exec_path)
            if not os.access(srm_exec_path, os.X_OK):
                raise FileNotFoundError('srm_exec_path not executable: %s' %
                                        srm_exec_path)
 
            self._srm_exec_path = srm_exec_path
        else:
            # Use the default
            self._srm_exec_path = DEFAULTS.srm_flexpart2ctbt_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)
 
        self._srm_rundir = os.path.join(self._working_rootdir,
                                        'srm_rundir')
        os.mkdir(self._srm_rundir, 0o755)
        if not os.path.isdir(self._srm_rundir):
            raise FileNotFoundError('Failed to create srm_rundir: %s' %
                                    self._srm_rundir)
 
 
 
 
        self._flexwrfout_dir = None
        if flexwrfout_dir:
            if not os.path.isdir(flexwrfout_dir):
                raise FileNotFoundError('flexwrfout_dir not found: %s' %
                                        flexwrfout_dir)
 
            self._flexwrfout_dir = flexwrfout_dir
        else:
            raise Exception('No flexwrfout_dir specified')
 
        if type(levels_list) is not list:
            raise ValueError('levels_list is not a list')
 
        for level in levels_list:
            if type(level) is not int:
                raise ValueError('levels_list has a non-integer element')
 
        # If we made it here, we should have a reasonable levels_list
        self._levels_list = levels_list
 
 
    def run_srm(self, srm_output_stagedir=None,
                multiplier=1.0,
                compress=False):
 
 
        '''Set up and run, via command line, an instance of flexpart2ctbt
        for each level in the levels_list
 
        Parameters
        ----------
        srm_output_stagedir : str
            (Optional) Full path to dir where srm output files will be staged.
            If not specified, no staging will be done.  If specified, we
            assume the dir already exists.
        multiplier : float
            (Optional) Factor by which to multiply concentrations
        compress : boolean
            (Optional) NOT YET IMPLEMENTED - compresses resulting SRM
            file if set to True
 
        Returns
        -------
        return_dict : dict
            Dictionary with manifest of the srm output files and the location
            (if applicable) of staged files
        '''
 
 
        LOGGER.debug('Starting run_srm()...')
 
 
        # If compress was requested, warn that it's not implemented yet
        if compress:
            LOGGER.warning('SRM compress is not yet implemented.')
 
 
 
        # Check that srm_output_stagedir exists (if it's specified
        # as an arg) and is writable
        if srm_output_stagedir:
            if not os.path.isdir(srm_output_stagedir):
                raise FileNotFoundError(
                    'srm_output_stagedir not found: %s' % srm_output_stagedir
                    )
            if not os.access(srm_output_stagedir, os.W_OK):
                raise FileNotFoundError(
                    'srm_output_stagedir not writable: %s' % srm_output_stagedir
                    )
 
        os.chdir(self._srm_rundir)
 
        srm_output_manifest = {}
        srm_output_manifest['rundir'] = self._srm_rundir
        srm_output_manifest['levels'] = {}
 
        srm_status_dict = {}
        for level in self._levels_list:
            LOGGER.debug('Creating SRM file for level: %d' % level)
 
 
            # Create the output_manifest for this level
            level_output_manifest = {str(level) : None}
 
            # Create control file for this level
            control_fname = 'control_level_%02d' % level
            cfh = open(control_fname, 'w')
            # Be sure to add trailing / to path, as required by flexpart2ctbt
            cfh.write(self._flexwrfout_dir + '/' + '\n')
            cfh.write('============================================\n')
            cfh.write('%d\n' % level)
            cfh.write('%e\n' % multiplier)
            cfh.write('fp\n')
            cfh.write('l%d\n' % level)
            cfh.close()
 
            # Run and test exit code (or something like that)
            the_command = self._srm_exec_path + ' < ' + control_fname
            subprocess.call(the_command, shell=True)
 
 
 
            # Gather the srm filenames for this level and add to
            # srm_output_manifest (with other components, this would
            # have been done in the nwpservice component, but we're
            # not putting srm down in there)
 
            # These are example filenames for a single release, for
            # Level 1
            #
            # SICILY.fp.20171023000000.l1.srm
            # SICILY.fp.20171023000000.l1.l0.wet.srm
            # SICILY.fp.20171023000000.l1.l0.dry.srm
 
            # It's a bit of a pattern-matching hack, but it should work
            digits_14 = '[0-9][0-9][0-9][0-9][0-9][0-9][0-9]'
            digits_14 += '[0-9][0-9][0-9][0-9][0-9][0-9][0-9]'
            prepattern = '*.fp.' + digits_14 + '.l'
            prepattern += '%1d' % level
            pattern_conc = prepattern + '.srm'
            pattern_dry = prepattern + '.l0.dry.srm'
            pattern_wet = prepattern + '.l0.wet.srm'
 
            matched_fnames = []
 
            for fname in os.listdir('.'):
                if fnmatch.fnmatch(fname, pattern_conc):
                    matched_fnames.append(fname)
                elif fnmatch.fnmatch(fname, pattern_dry):
                    matched_fnames.append(fname)
                elif fnmatch.fnmatch(fname, pattern_wet):
                    matched_fnames.append(fname)
 
            LOGGER.debug('matched_fnames: %s' % matched_fnames)
 
            # Create the output manifest for this level
            level_manifest = {}
            for fname in matched_fnames:
                # Get bytes and add to dict
                bytes = os.path.getsize(fname)
                file_info = {'bytes' : bytes}
                level_manifest[fname] = file_info
 
 
            srm_output_manifest['levels'] = {str(level) : level_manifest}
 
        LOGGER.debug('matched_fnames: %s' % matched_fnames)
        LOGGER.debug('srm_output_manifest: %s' % srm_output_manifest)
 
        # Consider a staging, like we normally would do down at
        # nwpservice level.  If unsuccessful, simply returns a False
        # Assumes that we also want auxfiles (in this case, the control
        # files, copied over, too)
        stage_success = self._stage_output(
                            staging_dir=srm_output_stagedir,
                            auxfiles=True
        )
        LOGGER.debug('stage_success: %s' % stage_success)
        if stage_success:
            staging_dir = srm_output_stagedir
        else:
            staging_dir = None
 
 
        # Create the return dict
        return_dict = {
            'srm_output_manifest' : srm_output_manifest,
            'staging_dir' : staging_dir
        }
        LOGGER.debug('return_dict: %s' % return_dict)
 
        return return_dict
 
 
    def _stage_output(self, staging_dir=None, auxfiles=False):
 
        '''Stages output to specified dir
 
        Stages the output (met_em*) and, optionally, other runtime files to
        specified directory
 
        Parameters
        ----------
        staging_dir : str
         Must already exist and be writable.  If not specified,
         then self._output_dir is used.  If that wasn't specified,
         then we abort and return False to caller
        auxfiles : boolean
         If set to True, will also copy in the metgrid.log* files
         and namelist.wps.
 
        Returns
        -------
        return value : boolean
         True or False based on success of operation
        '''
 
 
        if staging_dir:
         if os.path.isdir(staging_dir) and os.access(staging_dir, os.W_OK):
             dest_dir = staging_dir
             LOGGER.debug('dest_dir: %s' % dest_dir)
         else:
             LOGGER.error('staging_dir not accessible and/or writable: %s',
                          staging_dir)
             return False
        else:
         LOGGER.info('No valid, writable staging dir.  Returning None')
         return False
 
        # Get the *.srm filenames and copy over
        srm_files = [os.path.basename(x) for
                 x in glob.glob(self._srm_rundir + '/*.srm')]
 
        LOGGER.debug('Copying srm files: %s' % srm_files)
        for f in srm_files:
         shutil.copy(os.path.join(self._srm_rundir, f),
                     dest_dir)
 
        # Optionally, copy over the log files and namelist.wps
        if auxfiles:
         control_files = [os.path.basename(x) for
                     x in glob.glob(self._srm_rundir + '/control_level_*')]
         auxfile_list = control_files
         LOGGER.debug('Copying aux files: %s' % auxfile_list)
         for f in auxfile_list:
             shutil.copy(os.path.join(self._srm_rundir, f),
                         dest_dir)
 
        return True