WCOSS2.env

#! /usr/bin/env bash

if [[ $# -ne 1 ]]; then

    echo "Must specify an input argument to set runtime environment variables!"
    exit 1

fi

step=$1

# WCOSS2 information
export launcher="mpiexec -l"
export mpmd_opt="--cpu-bind verbose,core cfp"

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

    nth_max=$((npe_node_max / npe_node_prep))

    export POE=${POE:-"YES"}
    export BACK=${BACK:-"off"}
    export sys_tp="wcoss2"
    export launcher_PREP="mpiexec"

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

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

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

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

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

    export USE_CFP="YES"
    export wavempexec="${launcher} -np"
    export wave_mpmd=${mpmd_opt}

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

    nth_max=$((npe_node_max / npe_node_atmanlvar))

    export NTHREADS_ATMANLVAR=${nth_atmanlvar:-${nth_max}}
    [[ ${NTHREADS_ATMANLVAR} -gt ${nth_max} ]] && export NTHREADS_ATMANLVAR=${nth_max}
    export APRUN_ATMANLVAR="${launcher} -n ${npe_atmanlvar}"

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

    nth_max=$((npe_node_max / npe_node_atmensanlletkf))

    export NTHREADS_ATMENSANLLETKF=${nth_atmensanlletkf:-${nth_max}}
    [[ ${NTHREADS_ATMENSANLLETKF} -gt ${nth_max} ]] && export NTHREADS_ATMENSANLLETKF=${nth_max}
    export APRUN_ATMENSANLLETKF="${launcher} -n ${npe_atmensanlletkf}"

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

    nth_max=$((npe_node_max / npe_node_atmensanlfv3inc))

    export NTHREADS_ATMENSANLFV3INC=${nth_atmensanlfv3inc:-${nth_max}}
    [[ ${NTHREADS_ATMENSANLFV3INC} -gt ${nth_max} ]] && export NTHREADS_ATMENSANLFV3INC=${nth_max}
    export APRUN_ATMENSANLFV3INC="${launcher} -n ${npe_atmensanlfv3inc}"

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

    export APRUNCFP="${launcher} -np \$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}" = "prepobsaero" ]]; then

    nth_max=$((npe_node_max / npe_node_prepaeroobs))

    export NTHREADS_PREPOBSAERO=${nth_prepobsaero:-1}
    export APRUN_PREPOBSAERO="${launcher} -n ${npe_prepobsaero} --ppn ${npe_node_prepobsaero}--cpu-bind depth --depth=${NTHREADS_PREPOBSAERO}"

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

    nth_max=$((npe_node_max / npe_node_snowanl))

    export NTHREADS_SNOWANL=${nth_snowanl:-${nth_max}}
    [[ ${NTHREADS_SNOWANL} -gt ${nth_max} ]] && export NTHREADS_SNOWANL=${nth_max}
    export APRUN_SNOWANL="${launcher} -n ${npe_snowanl}"

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

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

    nth_max=$((npe_node_max / npe_node_atmanlfv3inc))

    export NTHREADS_ATMANLFV3INC=${nth_atmanlfv3inc:-${nth_max}}
    [[ ${NTHREADS_ATMANLFV3INC} -gt ${nth_max} ]] && export NTHREADS_ATMANLFV3INC=${nth_max}
    export APRUN_ATMANLFV3INC="${launcher} -n ${npe_atmanlfv3inc}"

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

    export OMP_PLACES=cores
    export OMP_STACKSIZE=1G
    export FI_OFI_RXM_SAR_LIMIT=3145728

    if [[ "${step}" = "analcalc" ]]; then
      export MPICH_MPIIO_HINTS="*:romio_cb_write=disable"
    fi

    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}} -ppn ${npe_node_anal} --cpu-bind depth --depth ${NTHREADS_GSI}"

    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:-14}
    [[ ${NTHREADS_CYCLE} -gt ${npe_node_max} ]] && export NTHREADS_CYCLE=${npe_node_max}
    npe_cycle=${ntiles:-6}
    export APRUN_CYCLE="${launcher} -n ${npe_cycle} -ppn ${npe_node_cycle} --cpu-bind depth --depth ${NTHREADS_CYCLE}"

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

    export NTHREADS_CHGRES=${nth_echgres:-14}
    [[ ${NTHREADS_CHGRES} -gt ${npe_node_max} ]] && export NTHREADS_CHGRES=${npe_node_max}
    export APRUN_CHGRES=""

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

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 OMP_PLACES=cores
    export OMP_STACKSIZE=1G
    export FI_OFI_RXM_SAR_LIMIT=3145728

    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}} -ppn ${npe_node_eobs} --cpu-bind depth --depth ${NTHREADS_GSI}"

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

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

    export OMP_PLACES=cores
    export OMP_STACKSIZE=2G
    export MPICH_COLL_OPT_OFF=1
    export FI_OFI_RXM_SAR_LIMIT=3145728

    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}} -ppn ${npe_node_eupd} --cpu-bind depth --depth ${NTHREADS_ENKF}"

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

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

    ppn="npe_node_${step}_${RUN}"
    [[ -z "${!ppn+0}" ]] && ppn="npe_node_${step}"
    nprocs="npe_${step}_${RUN}"
    [[ -z ${!nprocs+0} ]] && nprocs="npe_${step}"

    (( 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} -ppn ${!ppn} --cpu-bind depth --depth 1"
    unset nprocs ppn nnodes ntasks

    # TODO: Why are fcst and efcs so different on WCOSS2?
    # TODO: Compare these with the ufs-weather-model regression test job card at:
    # https://github.com/ufs-community/ufs-weather-model/blob/develop/tests/fv3_conf/fv3_qsub.IN_wcoss2
    export FI_OFI_RXM_RX_SIZE=40000
    export FI_OFI_RXM_TX_SIZE=40000
    if [[ "${step}" = "fcst" ]] || [[ "${step}" = "efcs" ]]; then
        export OMP_PLACES=cores
        export OMP_STACKSIZE=2048M
        export MPICH_MPIIO_HINTS="*:romio_cb_write=disable"
        export FI_OFI_RXM_SAR_LIMIT=3145728
    fi

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} -ppn ${npe_node_upp} --cpu-bind depth --depth ${NTHREADS_UPP}"

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

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

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

    nth_max=$((npe_node_max / npe_node_oceanice_products))

    export NTHREADS_OCNICEPOST=${nth_oceanice_products:-1}
    export APRUN_OCNICEPOST="${launcher} -n 1 -ppn ${npe_node_oceanice_products} --cpu-bind depth --depth ${NTHREADS_OCNICEPOST}"

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} -ppn ${npe_node_ecen} --cpu-bind depth --depth ${NTHREADS_ECEN}"

    export NTHREADS_CHGRES=${nth_chgres:-14}
    [[ ${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}"

    export NTHREADS_CYCLE=${nth_cycle:-14}
    [[ ${NTHREADS_CYCLE} -gt ${npe_node_max} ]] && export NTHREADS_CYCLE=${npe_node_max}
    export APRUN_CYCLE="${launcher} -n ${npe_ecen} -ppn ${npe_node_cycle} --cpu-bind depth --depth ${NTHREADS_CYCLE}"

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} -ppn ${npe_node_esfc} --cpu-bind depth --depth ${NTHREADS_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} -ppn ${npe_node_cycle} --cpu-bind depth --depth ${NTHREADS_CYCLE}"

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} -ppn ${npe_node_epos} --cpu-bind depth --depth ${NTHREADS_EPOS}"

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

    export MPICH_MPIIO_HINTS_DISPLAY=1
    export OMP_NUM_THREADS=1

    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} --depth=${NTHREADS_POSTSND} --cpu-bind depth"

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

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

    nth_max=$((npe_node_max / npe_node_awips))

    export NTHREADS_AWIPS=${nth_awips:-2}
    [[ ${NTHREADS_AWIPS} -gt ${nth_max} ]] && export NTHREADS_AWIPS=${nth_max}
    export APRUN_AWIPSCFP="${launcher} -np ${npe_awips} ${mpmd_opt}"

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

    nth_max=$((npe_node_max / npe_node_gempak))

    export NTHREADS_GEMPAK=${nth_gempak:-1}
    [[ ${NTHREADS_GEMPAK} -gt ${nth_max} ]] && export NTHREADS_GEMPAK=${nth_max}
    export APRUN_GEMPAKCFP="${launcher} -np ${npe_gempak} ${mpmd_opt}"

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} -np ${npe_fit2obs}"

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

    unset PERL5LIB

elif [[ "${step:0:3}" = "mos" ]]; then

    export FORT_BUFFERED=TRUE

    if [[ "${step}" = "mos_stn_prep" ]]; then
       export OMP_PROC_BIND=true
    fi

fi