JET.env

#! /usr/bin/env bash

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 "atmanlrun atmensanlrun aeroanlrun landanl"
    echo "anal sfcanl fcst post 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
elif [[ "${PARTITION_BATCH}" = "kjet" ]]; then
  export npe_node_max=40
fi
export launcher="srun -l --epilog=/apps/local/bin/report-mem --export=ALL"
export mpmd_opt="--multi-prog --output=mpmd.%j.%t.out"

# Configure MPI environment
export OMP_STACKSIZE=2048000
export NTHSTACK=1024000000

ulimit -s unlimited
ulimit -a

if [[ "${step}" = "prep" ]] || [[ "${step}" = "prepbufr" ]]; then

    nth_max=$((npe_node_max / npe_node_prep))

    export POE="NO"
    export BACK="NO"
    export sys_tp="JET"
    export launcher_PREP="srun"

elif [[ "${step}" = "preplandobs" ]]; then

    export APRUN_CALCFIMS="${launcher} -n 1"

elif [[ "${step}" = "waveinit" ]] || [[ "${step}" = "waveprep" ]] || [[ "${step}" = "wavepostsbs" ]] || [[ "${step}" = "wavepostbndpnt" ]] || [[ "${step}" = "wavepostbndpntbll" ]] || [[ "${step}" = "wavepostpnt" ]]; then

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

elif [[ "${step}" = "atmanlrun" ]]; then

    nth_max=$((npe_node_max / npe_node_atmanlrun))

    export NTHREADS_ATMANL=${nth_atmanlrun:-${nth_max}}
    [[ ${NTHREADS_ATMANL} -gt ${nth_max} ]] && export NTHREADS_ATMANL=${nth_max}
    export APRUN_ATMANL="${launcher} -n ${npe_atmanlrun}"

elif [[ "${step}" = "atmensanlrun" ]]; then

    nth_max=$((npe_node_max / npe_node_atmensanlrun))

    export NTHREADS_ATMENSANL=${nth_atmensanlrun:-${nth_max}}
    [[ ${NTHREADS_ATMENSANL} -gt ${nth_max} ]] && export NTHREADS_ATMENSANL=${nth_max}
    export APRUN_ATMENSANL="${launcher} ${npe_atmensanlrun}"

elif [[ "${step}" = "aeroanlrun" ]]; then

    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"

    nth_max=$((npe_node_max / npe_node_aeroanlrun))

    export NTHREADS_AEROANL=${nth_aeroanlrun:-${nth_max}}
    [[ ${NTHREADS_AEROANL} -gt ${nth_max} ]] && export NTHREADS_AEROANL=${nth_max}
    export APRUN_AEROANL="${launcher} -n ${npe_aeroanlrun}"

elif [[ "${step}" = "landanl" ]]; then

    nth_max=$((npe_node_max / npe_node_landanl))

    export NTHREADS_LANDANL=${nth_landanl:-${nth_max}}
    [[ ${NTHREADS_LANDANL} -gt ${nth_max} ]] && export NTHREADS_LANDANL=${nth_max}
    export APRUN_LANDANL="${launcher} -n ${npe_landanl}"

    export APRUN_APPLY_INCR="${launcher} -n 6"

elif [[ "${step}" = "ocnanalbmat" ]]; then

    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"

    nth_max=$((npe_node_max / npe_node_ocnanalbmat))

    export NTHREADS_OCNANAL=${nth_ocnanalbmat:-${nth_max}}
    [[ ${NTHREADS_OCNANAL} -gt ${nth_max} ]] && export NTHREADS_OCNANAL=${nth_max}
    export APRUN_OCNANAL="${launcher} -n ${npe_ocnanalbmat}"

elif [[ "${step}" = "ocnanalrun" ]]; then

    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"

    nth_max=$((npe_node_max / npe_node_ocnanalrun))

    export NTHREADS_OCNANAL=${nth_ocnanalrun:-${nth_max}}
    [[ ${NTHREADS_OCNANAL} -gt ${nth_max} ]] && export NTHREADS_OCNANAL=${nth_max}
    export APRUN_OCNANAL="${launcher} -n ${npe_ocnanalrun}"

elif [[ "${step}" = "anal" ]] || [[ "${step}" = "analcalc" ]]; then

    export MKL_NUM_THREADS=4
    export MKL_CBWR=AUTO

    export CFP_MP=${CFP_MP:-"YES"}
    export USE_CFP=${USE_CFP:-"YES"}
    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"

    nth_max=$((npe_node_max / npe_node_anal))

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

    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} -n ${npe_cycle}"

    export NTHREADS_GAUSFCANL=1
    npe_gausfcanl=${npe_gausfcanl:-1}
    export APRUN_GAUSFCANL="${launcher} -n ${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} -n ${npe_sfcanl}"

elif [[ "${step}" = "eobs" ]]; then

    export MKL_NUM_THREADS=4
    export MKL_CBWR=AUTO

    nth_max=$((npe_node_max / npe_node_eobs))

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

    export CFP_MP=${CFP_MP:-"YES"}
    export USE_CFP=${USE_CFP:-"YES"}
    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"

elif [[ "${step}" = "eupd" ]]; then

    nth_max=$((npe_node_max / npe_node_eupd))

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

    export CFP_MP=${CFP_MP:-"YES"}
    export USE_CFP=${USE_CFP:-"YES"}
    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"

elif [[ "${step}" = "fcst" ]] || [[ "${step}" = "efcs" ]]; then

    if [[ "${CDUMP}" =~ "gfs" ]]; then
        nprocs="npe_${step}_gfs"
        ppn="npe_node_${step}_gfs" || ppn="npe_node_${step}"
    else
        nprocs="npe_${step}"
        ppn="npe_node_${step}"
    fi
    (( nnodes = (${!nprocs}+${!ppn}-1)/${!ppn} ))
    (( ntasks = nnodes*${!ppn} ))
    # With ESMF threading, the model wants to use the full node
    export APRUN_UFS="${launcher} -n ${ntasks}"
    unset nprocs ppn nnodes ntasks

elif [[ "${step}" = "upp" ]]; then

    nth_max=$((npe_node_max / npe_node_upp))

    export NTHREADS_UPP=${nth_upp:-1}
    [[ ${NTHREADS_UPP} -gt ${nth_max} ]] && export NTHREADS_UPP=${nth_max}
    export APRUN_UPP="${launcher} -n ${npe_upp}"

elif [[ "${step}" = "atmos_products" ]]; then

    export USE_CFP="YES"  # Use MPMD for downstream product generation

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} -n ${npe_ecen}"

    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} -n ${npe_ecen}"

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} -n ${npe_esfc}"

    export NTHREADS_CYCLE=${nth_cycle:-14}
    [[ ${NTHREADS_CYCLE} -gt ${npe_node_max} ]] && export NTHREADS_CYCLE=${npe_node_max}
    export APRUN_CYCLE="${launcher} -n ${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} -n ${npe_epos}"

elif [[ "${step}" = "postsnd" ]]; then

    export CFP_MP="YES"

    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} -n ${npe_postsnd}"

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

elif [[ "${step}" = "awips" ]]; then

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

elif [[ "${step}" = "gempak" ]]; then

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

elif [[ "${step}" = "fit2obs" ]]; then

    nth_max=$((npe_node_max / npe_node_fit2obs))

    export NTHREADS_FIT2OBS=${nth_fit2obs:-1}
    [[ ${NTHREADS_FIT2OBS} -gt ${nth_max} ]] && export NTHREADS_FIT2OBS=${nth_max}
    export MPIRUN="${launcher} -n ${npe_fit2obs}"

fi