ecearth3/sources/tm5mp/tools/ttb/bin/ttb_restart.py

"""
TTB restart test.
"""


# ***


def MyLogger( name ) :
    # external:
    import logging
    import sys
    # setup logger:
    logger = logging.getLogger( name )
    # no handlers yet ? then print to standard output:
    if len(logger.handlers) == 0 : logger.addHandler(logging.StreamHandler(sys.stdout))
    # no level set yet ? then set for info and higher:
    if logger.level == 0 : logger.setLevel(logging.INFO)
    # ok:
    return logger
#enddef


# ***


def Test( rcf ) :

    # external:
    import os
    import datetime
    import subprocess
    
    # tools:
    import rc
    import ttb_compare
    
    # get logger:
    logger = MyLogger('ttb')
    
    # test name:
    test_name = 'restart'
    
    # info ...
    logger.info( '  "%s" test ...' % test_name )
    
    # name of TM5 rcfile:
    tm_rcfile = rcf.get( 'ttb.rcfile' )
    
    # load the tm5 rcfile:
    tm_rcf = rc.RcFile( tm_rcfile )
    
    # flags:
    no_run = rcf.get( 'ttb.no-run', 'bool' )

    # time range:
    tfmt = '%Y-%m-%d %H:%M:%S'
    t_start = datetime.datetime.strptime( rcf.get('ttb.restart.timerange.start'), tfmt )
    t_end   = datetime.datetime.strptime( rcf.get('ttb.restart.timerange.end'), tfmt )
    # for restart:
    t_mid = t_start + (t_end-t_start)/2

    # run id's:
    runids = ['restart0','restart1','restart2']
    

    #
    # perform runs
    #
    
    # init lists:
    test_rcfs = []
    
    # loop over runs:
    for irun in range(len(runids)) :
    
        # skip some ..
        #if irun == 0 : continue
        #if irun == 1 : continue
    
        # current value:
        runid = runids[irun]
    
        # read template rcfile in raw format:
        test_rcf_raw = rc.RcFile( tm_rcfile, raw=True )
        
        # replace project name:
        my_project = 'ttb'
        test_rcf_raw.replace( 'my.project', my_project )

        # replace runid:
        test_rcf_raw.replace( 'runid', runid )

        # replace project directory:
        test_rcf_raw.replace( 'my.project.dir', '${my.basedir}/${my.project}/%s/${runid}' % test_name )

        # restart specific:
        if irun == 0 :
            test_rcf_raw.replace( 'timerange.start'    , t_start.strftime(tfmt) )
            test_rcf_raw.replace( 'timerange.end'      , t_end.strftime(tfmt) )
            test_rcf_raw.replace( 'istart'             , '2' )
            test_rcf_raw.replace( 'restart.restore.dir', 'None' )
            test_rcf_raw.replace( 'restart.store.dir'  , '${my.basedir}/${my.project}/${runid}/restart' )
        elif irun == 1 :
            test_rcf_raw.replace( 'timerange.start'    , t_start.strftime(tfmt) )
            test_rcf_raw.replace( 'timerange.end'      , t_mid.strftime(tfmt) )
            test_rcf_raw.replace( 'istart'             , '2' )
            test_rcf_raw.replace( 'restart.restore.dir', 'None' )
            test_rcf_raw.replace( 'restart.store.dir'  , '${my.basedir}/${my.project}/${runid}/restart' )
        elif irun == 2 :
            test_rcf_raw.replace( 'timerange.start'    , t_mid.strftime(tfmt) )
            test_rcf_raw.replace( 'timerange.end'      , t_end.strftime(tfmt) )
            test_rcf_raw.replace( 'istart'             , '33' )
            test_rcf_raw.replace( 'restart.restore.dir', '${my.basedir}/${my.project}/%s/restart' % runids[1] )
            test_rcf_raw.replace( 'restart.store.dir'  , '${my.basedir}/${my.project}/${runid}/restart' )
        else :
            logger.error( 'unsupported runid : %s' % runid )
            sys.exit(1)
        #endif
        
        # target filename: tmp-<origingal_basename>-<runid>.rc
        bname,ext = os.path.splitext(tm_rcfile)
        test_rcfile_test = 'tmp-ttb-%s-%s.rc' % (test_name,runid)
        # write:
        test_rcf_raw.WriteFile( test_rcfile_test )
        
        # read again:
        test_rcf = rc.RcFile( test_rcfile_test )
        # add:
        test_rcfs.append( test_rcf )
        
        # run:
        if not no_run :

            logger.info( '' )
            logger.info( '******************************************************************' )
            logger.info( '' )
            logger.info( '%s' % runid )
            logger.info( '' )
            logger.info( '******************************************************************' )
            logger.info( '' )

            # setup and submit command:
            command = [os.path.join(os.curdir,'setup_tm5'), test_rcfile_test, '--submit']
            retcode = subprocess.call( command )
            if retcode != 0 :
                logger.error( 'from call:' )
                logger.error( '  %s' % command )
                raise Exception
            #endif
        #endif
        
    #endfor  # runs

    
    #
    # check results
    #
    
    # which regions ?
    regions = test_rcf.get( 'regions' ).split()
    
    # loop over hours:
    t = t_mid
    while t <= t_end :
    
        # flags:
        isfirsthour = t == t_mid
        islasthour = t == t_end
        
        # info ...
        logger.info( '' )
        logger.info( 'compare restart files for time %s ...' % t )
    
        # loop over regions
        for region in regions :
    
            # files to be compared:
            fnames = []
            for irun in range(len(runids)) :
                # current:
                runid = runids[irun]
                test_rcf = test_rcfs[irun]

                # first file ?
                isfirstkey = ''
                if (irun == 2) and isfirsthour : isfirstkey = '_initial'

                # directory with restart files:
                fdir = test_rcf.get('restart.store.dir')
                # target file:
                fnam = 'TM5_restart_%s_%s%s.nc' % (t.strftime('%Y%m%d_%H%M'),region,isfirstkey)
                # add:
                fnames.append( os.path.join(fdir,fnam) )
            #endfor

            # which files ?
            i1 = 0
            i2 = 2
            
            # compare ...
            ttb_compare.restart_files( fnames[i1], fnames[i2] )
            
        #endfor  # regions
        
        # next hour:
        t = t + datetime.timedelta(0,3600)  # 1 hour
        
        ## testing ...
        #break
        
    #endfor  # hours
    
#enddef