JET.env

#!/bin/ksh -x

if [ $# -ne 1 ]; then

    echo "Must specify an input argument to set runtime environment variables!"
    echo "argument can be any one of the following:"
    echo "atmanalrun atmensanalrun"
    echo "anal sfcanl fcst post vrfy metp"
    echo "eobs eupd ecen efcs epos"
    echo "postsnd awips gempak"
    exit 1

fi

step=$1

if [[ "$PARTITION_BATCH" = "xjet" ]]; then
  export npe_node_max=24
elif [[ "$PARTITION_BATCH" = "vjet" ]]; then
  export npe_node_max=16
fi
export launcher="srun -l --export=ALL"

# Configure STACK
export OMP_STACKSIZE=2048000
export NTHSTACK=1024000000

ulimit -s unlimited
ulimit -a

export job=${PBS_JOBNAME:-$step}
export jobid=${job}.${PBS_JOBID:-$$}

if [ $step = "prep" -o $step = "prepbufr" ]; then

    nth_max=$(($npe_node_max / $npe_node_prep))

    export POE="NO"
    export BACK=${BACK:-"NO"}
    export sys_tp="JET"

elif [ $step = "waveinit" -o $step = "waveprep" -o $step = "wavepostsbs" -o $step = "wavepostbndpnt" -o $step = "wavepostbndpntbll" -o $step = "wavepostpnt" ]; then

    if [ $step = "waveprep" ]; then export MP_PULSE=0 ; fi
    export wavempexec=${launcher}
    export wave_mpmd=${mpmd}

elif [ $step = "atmanalrun" ]; then

    nth_max=$(($npe_node_max / $npe_node_atmanalrun))

    export NTHREADS_ATMANAL=${nth_atmanalrun:-$nth_max}
    [[ $NTHREADS_ATMANAL -gt $nth_max ]] && export NTHREADS_ATMANAL=$nth_max
    export APRUN_ATMANAL="$launcher $npe_atmanalrun"

elif [ $step = "atmensanalrun" ]; then

    nth_max=$(($npe_node_max / $npe_node_atmensanalrun))

    export NTHREADS_ATMENSANAL=${nth_atmensanalrun:-$nth_max}
    [[ $NTHREADS_ATMENSANAL -gt $nth_max ]] && export NTHREADS_ATMENSANAL=$nth_max
    export APRUN_ATMENSANAL="$launcher $npe_atmensanalrun"

elif [ $step = "anal" ]; then

    nth_max=$(($npe_node_max / $npe_node_anal))

    export NTHREADS_GSI=${nth_gsi:-$nth_max}
    [[ $NTHREADS_GSI -gt $nth_max ]] && export NTHREADS_GSI=$nth_max
    export APRUN_GSI="$launcher ${npe_gsi:-${npe_anal:-$PBS_NP}}"

    export NTHREADS_CALCINC=${nth_calcinc:-1}
    [[ $NTHREADS_CALCINC -gt $nth_max ]] && export NTHREADS_CALCINC=$nth_max
    export APRUN_CALCINC="$launcher \$ncmd"

    export NTHREADS_CYCLE=${nth_cycle:-12}
    [[ $NTHREADS_CYCLE -gt $npe_node_max ]] && export NTHREADS_CYCLE=$npe_node_max
    npe_cycle=${ntiles:-6}
    export APRUN_CYCLE="$launcher $npe_cycle"

    export NTHREADS_GAUSFCANL=1
    npe_gausfcanl=${npe_gausfcanl:-1}
    export APRUN_GAUSFCANL="$launcher $npe_gausfcanl"

elif [ $step = "sfcanl" ]; then
    nth_max=$(($npe_node_max / $npe_node_sfcanl))

    export NTHREADS_CYCLE=${nth_sfcanl:-14}
    [[ $NTHREADS_CYCLE -gt $npe_node_max ]] && export NTHREADS_CYCLE=$npe_node_max
    npe_sfcanl=${ntiles:-6}
    export APRUN_CYCLE="$launcher $npe_sfcanl"

elif [ $step = "gldas" ]; then

    nth_max=$(($npe_node_max / $npe_node_gldas))

    export NTHREADS_GLDAS=${nth_gldas:-$nth_max}
    [[ $NTHREADS_GLDAS -gt $nth_max ]] && export NTHREADS_GLDAS=$nth_max
    export APRUN_GLDAS="$launcher $npe_gldas"

    export NTHREADS_GAUSSIAN=${nth_gaussian:-1}
    [[ $NTHREADS_GAUSSIAN -gt $nth_max ]] && export NTHREADS_GAUSSIAN=$nth_max
    export APRUN_GAUSSIAN="$launcher $npe_gaussian"

elif [ $step = "eobs" ]; then

    nth_max=$(($npe_node_max / $npe_node_eobs))

    export NTHREADS_GSI=${nth_gsi:-$nth_max}
    [[ $NTHREADS_GSI -gt $nth_max ]] && export NTHREADS_GSI=$nth_max
    export APRUN_GSI="$launcher ${npe_gsi:-${npe_eobs:-$PBS_NP}}"

elif [ $step = "eupd" ]; then

    nth_max=$(($npe_node_max / $npe_node_eupd))

    export NTHREADS_ENKF=${nth_enkf:-$nth_max}
    [[ $NTHREADS_ENKF -gt $nth_max ]] && export NTHREADS_ENKF=$nth_max
    export APRUN_ENKF="$launcher ${npe_enkf:-${npe_eupd:-$PBS_NP}}"

elif [ $step = "fcst" ]; then

    #PEs and PEs/node can differ for GFS and GDAS forecasts if threading differs
    if [[ $CDUMP == "gfs" ]]; then
        npe_fcst=$npe_fcst_gfs
        npe_node_fcst=$npe_node_fcst_gfs
        nth_fv3=$nth_fv3_gfs
    fi

    nth_max=$(($npe_node_max / $npe_node_fcst))

    export NTHREADS_FV3=${nth_fv3:-$nth_max}
    [[ $NTHREADS_FV3 -gt $nth_max ]] && export NTHREADS_FV3=$nth_max
    export cores_per_node=$npe_node_max
    export APRUN_FV3="$launcher"

    export NTHREADS_REGRID_NEMSIO=${nth_regrid_nemsio:-1}
    [[ $NTHREADS_REGRID_NEMSIO -gt $nth_max ]] && export NTHREADS_REGRID_NEMSIO=$nth_max
    export APRUN_REGRID_NEMSIO="$launcher $LEVS"

    export NTHREADS_REMAP=${nth_remap:-2}
    [[ $NTHREADS_REMAP -gt $nth_max ]] && export NTHREADS_REMAP=$nth_max
    export APRUN_REMAP="$launcher ${npe_remap:-${npe_fcst:-$PBS_NP}}"
    export I_MPI_DAPL_UD="enable"

elif [ $step = "efcs" ]; then

    nth_max=$(($npe_node_max / $npe_node_efcs))

    export NTHREADS_FV3=${nth_fv3:-$nth_max}
    [[ $NTHREADS_FV3 -gt $nth_max ]] && export NTHREADS_FV3=$nth_max
    export cores_per_node=$npe_node_max
    export APRUN_FV3="$launcher ${npe_fv3:-${npe_efcs:-$PBS_NP}}"

    export NTHREADS_REGRID_NEMSIO=${nth_regrid_nemsio:-1}
    [[ $NTHREADS_REGRID_NEMSIO -gt $nth_max ]] && export NTHREADS_REGRID_NEMSIO=$nth_max
    export APRUN_REGRID_NEMSIO="$launcher $LEVS"

elif [ $step = "post" ]; then

    nth_max=$(($npe_node_max / $npe_node_post))

    export NTHREADS_NP=${nth_np:-1}
    [[ $NTHREADS_NP -gt $nth_max ]] && export NTHREADS_NP=$nth_max
    export APRUN_NP="$launcher"

    export NTHREADS_DWN=${nth_dwn:-1}
    [[ $NTHREADS_DWN -gt $nth_max ]] && export NTHREADS_DWN=$nth_max
    export APRUN_NP="$launcher"

elif [ $step = "ecen" ]; then

    nth_max=$(($npe_node_max / $npe_node_ecen))

    export NTHREADS_ECEN=${nth_ecen:-$nth_max}
    [[ $NTHREADS_ECEN -gt $nth_max ]] && export NTHREADS_ECEN=$nth_max
    export APRUN_ECEN="$launcher ${npe_ecen:-$PBS_NP}"

    export NTHREADS_CHGRES=${nth_chgres:-12}
    [[ $NTHREADS_CHGRES -gt $npe_node_max ]] && export NTHREADS_CHGRES=$npe_node_max
    export APRUN_CHGRES="time"

    export NTHREADS_CALCINC=${nth_calcinc:-1}
    [[ $NTHREADS_CALCINC -gt $nth_max ]] && export NTHREADS_CALCINC=$nth_max
    export APRUN_CALCINC="$launcher ${npe_ecen:-$PBS_NP}"

elif [ $step = "esfc" ]; then

    nth_max=$(($npe_node_max / $npe_node_esfc))

    export NTHREADS_ESFC=${nth_esfc:-$nth_max}
    [[ $NTHREADS_ESFC -gt $nth_max ]] && export NTHREADS_ESFC=$nth_max
    export APRUN_ESFC="$launcher ${npe_esfc:-$PBS_NP}"

    export NTHREADS_CYCLE=${nth_cycle:-12}
    [[ $NTHREADS_CYCLE -gt $npe_node_max ]] && export NTHREADS_CYCLE=$npe_node_max
    export APRUN_CYCLE="$launcher $npe_esfc"


elif [ $step = "epos" ]; then

    nth_max=$(($npe_node_max / $npe_node_epos))

    export NTHREADS_EPOS=${nth_epos:-$nth_max}
    [[ $NTHREADS_EPOS -gt $nth_max ]] && export NTHREADS_EPOS=$nth_max
    export APRUN_EPOS="$launcher ${npe_epos:-$PBS_NP}"

elif [ $step = "init" ]; then

    export APRUN="$launcher"

elif [ $step = "postsnd" ]; then

    nth_max=$(($npe_node_max / $npe_node_postsnd))

    export NTHREADS_POSTSND=${nth_postsnd:-1}
    [[ $NTHREADS_POSTSND -gt $nth_max ]] && export NTHREADS_POSTSND=$nth_max
    export APRUN_POSTSND="$launcher $npe_postsnd"

    export NTHREADS_POSTSNDCFP=${nth_postsndcfp:-1}
    [[ $NTHREADS_POSTSNDCFP -gt $nth_max ]] && export NTHREADS_POSTSNDCFP=$nth_max
    export APRUN_POSTSNDCFP="$launcher $npe_postsndcfp"

elif [ $step = "awips" ]; then

     echo "WARNING: $step is not enabled on $machine!"

elif [ $step = "gempak" ]; then

     echo "WARNING: $step is not enabled on $machine!"

fi