GAEAC5.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

export launcher="srun -l --export=ALL"
export mpmd_opt="--multi-prog --output=mpmd.%j.%t.out"

export OMP_STACKSIZE=2048000
export NTHSTACK=1024000000

ulimit -s unlimited
ulimit -a

# Calculate common variables
# Check first if the dependent variables are set
if [[ -n "${ntasks:-}" && -n "${max_tasks_per_node:-}" && -n "${tasks_per_node:-}" ]]; then
    max_threads_per_task=$((max_tasks_per_node / tasks_per_node))
    NTHREADSmax=${threads_per_task:-${max_threads_per_task}}
    NTHREADS1=${threads_per_task:-1}
    if [[ ${NTHREADSmax} -gt ${max_threads_per_task} ]]; then
        NTHREADSmax=${max_threads_per_task}
    fi
    if [[ ${NTHREADS1} -gt ${max_threads_per_task} ]]; then
        NTHREADS1=${max_threads_per_task}
    fi
    # This may be useful when GaeaC5 is fully ported, so ignore SC warning
    # shellcheck disable=SC2034
    APRUN_default="${launcher} -n ${ntasks}"
else
    echo "ERROR config.resources must be sourced before sourcing GAEAC5.env"
    exit 2
fi

case ${step} in
 "prep" | "prepbufr")

    export POE="NO"
    export BACK=${BACK:-"YES"}
    export sys_tp="GAEAC5"
    export launcher_PREP="srun"
 ;;
 "prep_emissions")

    export APRUN="${APRUN_default}"
 ;;
 "waveinit" | "waveprep" | "wavepostsbs" | "wavepostbndpnt" | "wavepostpnt" | "wavepostbndpntbll")

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

 ;;
 "atmanlvar")

    export NTHREADS_ATMANLVAR=${NTHREADSmax}
    export APRUN_ATMANLVAR="${APRUN_default} --cpus-per-task=${NTHREADS_ATMANLVAR}"
 ;;
 "atmanlfv3inc")

    export NTHREADS_ATMANLFV3INC=${NTHREADSmax}
    export APRUN_ATMANLFV3INC="${APRUN_default} --cpus-per-task=${NTHREADS_ATMANLFV3INC}"
 ;;
 "atmensanlobs")

    export NTHREADS_ATMENSANLOBS=${NTHREADSmax}
    export APRUN_ATMENSANLOBS="${APRUN_default} --cpus-per-task=${NTHREADS_ATMENSANLOBS}"
 ;;
 "atmensanlsol")

    export NTHREADS_ATMENSANLSOL=${NTHREADSmax}
    export APRUN_ATMENSANLSOL="${APRUN_default} --cpus-per-task=${NTHREADS_ATMENSANLSOL}"
 ;;
 "atmensanlletkf")

    export NTHREADS_ATMENSANLLETKF=${NTHREADSmax}
    export APRUN_ATMENSANLLETKF="${APRUN_default} --cpus-per-task=${NTHREADS_ATMENSANLLETKF}"
 ;;
 "atmensanlfv3inc")

    export NTHREADS_ATMENSANLFV3INC=${NTHREADSmax}
    export APRUN_ATMENSANLFV3INC="${APRUN_default} --cpus-per-task=${NTHREADS_ATMENSANLFV3INC}"
 ;;
 "aeroanlvar")

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

    export NTHREADS_AEROANL=${NTHREADSmax}
    export APRUN_AEROANL="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANL}"
 ;;
  "aeroanlgenb")

    export NTHREADS_AEROANLGENB=${NTHREADSmax}
    export APRUN_AEROANLGENB="${APRUN_default} --cpus-per-task=${NTHREADS_AEROANLGENB}"
 ;;
 "prepobsaero")

    export NTHREADS_PREPOBSAERO=${NTHREADS1}
    export APRUN_PREPOBSAERO="${APRUN_default} --cpus-per-task=${NTHREADS_PREPOBSAERO}"
;;
 "snowanl")

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

    export NTHREADS_SNOWANL=${NTHREADSmax}
    export APRUN_SNOWANL="${APRUN_default} --cpus-per-task=${NTHREADS_SNOWANL}"

    export APRUN_APPLY_INCR="${launcher} -n 6"
 ;;
 "esnowrecen")

    export NTHREADS_ESNOWRECEN=${NTHREADSmax}
    export APRUN_ESNOWRECEN="${APRUN_default} --cpus-per-task=${NTHREADS_ESNOWRECEN}"

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

 "marinebmat")

    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
    export APRUN_MARINEBMAT="${APRUN_default}"
 ;;
 "marineanlvar")

    export APRUNCFP="${launcher} -n \$ncmd ${mpmd_opt}"
    export APRUN_MARINEANLVAR="${APRUN_default}"
 ;;
"ocnanalecen")

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

    max_threads_per_task=$((max_tasks_per_node / tasks_per_node_ocnanalecen))

    export NTHREADS_OCNANALECEN=${threads_per_task_ocnanalecen:-${max_threads_per_task}}
    if [[ ${NTHREADS_OCNANALECEN} -gt ${max_threads_per_task} ]]; then
        export NTHREADS_OCNANALECEN=${max_threads_per_task}
    fi
    export APRUN_OCNANALECEN="${launcher} -n ${ntasks_ocnanalecen} --cpus-per-task=${NTHREADS_OCNANALECEN}"
;;
 "marineanlchkpt")

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

    export NTHREADS_OCNANAL=${NTHREADSmax}
    export APRUN_MARINEANLCHKPT="${APRUN_default} --cpus-per-task=${NTHREADS_OCNANAL}"
 ;;
 "anal" | "analcalc")

    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}"

    export NTHREADS_GSI=${threads_per_task_anal:-${max_threads_per_task}}
    export APRUN_GSI="${APRUN_default} --cpus-per-task=${NTHREADS_GSI}"

    export NTHREADS_CALCINC=${threads_per_task_calcinc:-1}
    if [[ ${NTHREADS_CALCINC} -gt ${max_threads_per_task} ]]; then
        export NTHREADS_CALCINC=${max_threads_per_task}
    fi
    export APRUN_CALCINC="${launcher} \$ncmd --cpus-per-task=${NTHREADS_CALCINC}"

    export NTHREADS_CYCLE=${threads_per_task_cycle:-12}
    if [[ ${NTHREADS_CYCLE} -gt ${max_tasks_per_node} ]]; then
        export NTHREADS_CYCLE=${max_tasks_per_node}
    fi
    ntasks_cycle=${ntiles:-6}
    export APRUN_CYCLE="${launcher} -n ${ntasks_cycle} --cpus-per-task=${NTHREADS_CYCLE}"

    export NTHREADS_GAUSFCANL=1
    ntasks_gausfcanl=${ntasks_gausfcanl:-1}
    export APRUN_GAUSFCANL="${launcher} -n ${ntasks_gausfcanl} --cpus-per-task=${NTHREADS_GAUSFCANL}"
 ;;
 "sfcanl")

    export NTHREADS_CYCLE=${threads_per_task:-14}
    if [[ ${NTHREADS_CYCLE} -gt ${max_tasks_per_node} ]]; then
        export NTHREADS_CYCLE=${max_tasks_per_node}
    fi
    export APRUN_CYCLE="${APRUN_default} --cpus-per-task=${NTHREADS_CYCLE}"

    # REGRID requires multiple of 6 tasks
    ntdiv6=$((ntasks/6))
    ntasks_regrid=$((ntdiv6*6))
    if [[ ${ntasks_regrid} -gt 36 ]]; then
        ntasks_regrid=36
    fi
    export APRUN_REGRID="${launcher} -n ${ntasks_regrid} "

 ;;
 "eobs")

    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}"


    export NTHREADS_GSI=${NTHREADSmax}
    if [[ ${NTHREADS_GSI} -gt ${max_threads_per_task} ]]; then
        export NTHREADS_GSI=${max_threads_per_task}
    fi
    export APRUN_GSI="${APRUN_default} --cpus-per-task=${NTHREADS_GSI}"
 ;;
 "eupd")

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


    export NTHREADS_ENKF=${NTHREADSmax}
    export APRUN_ENKF="${launcher} -n ${ntasks_enkf:-${ntasks}} --cpus-per-task=${NTHREADS_ENKF}"
 ;;
 "fcst" | "efcs")

    export OMP_STACKSIZE=512M

    (( nnodes = (ntasks+tasks_per_node-1)/tasks_per_node ))
    (( ufs_ntasks = nnodes*tasks_per_node ))
    # With ESMF threading, the model wants to use the full node
    export APRUN_UFS="${launcher} -n ${ufs_ntasks}"
    unset nnodes ufs_ntasks
 ;;

 "upp")

    export NTHREADS_UPP=${NTHREADS1}
    export APRUN_UPP="${APRUN_default} --cpus-per-task=${NTHREADS_UPP}"
 ;;

 "atmos_products")

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

"oceanice_products")

    export NTHREADS_OCNICEPOST=${NTHREADS1}
    export APRUN_OCNICEPOST="${launcher} -n 1 --cpus-per-task=${NTHREADS_OCNICEPOST}"
;;

 "ecen")

    export NTHREADS_ECEN=${NTHREADSmax}
    export APRUN_ECEN="${APRUN_default} --cpus-per-task=${NTHREADS_ECEN}"

    export NTHREADS_CHGRES=${threads_per_task_chgres:-12}
    if [[ ${NTHREADS_CHGRES} -gt ${max_tasks_per_node} ]]; then
        export NTHREADS_CHGRES=${max_tasks_per_node}
    fi
    export APRUN_CHGRES="time"

    export NTHREADS_CALCINC=${threads_per_task_calcinc:-1}
    if [[ ${NTHREADS_CALCINC} -gt ${max_threads_per_task} ]]; then
        export NTHREADS_CALCINC=${max_threads_per_task}
    fi
    export APRUN_CALCINC="${APRUN_default} --cpus-per-task=${NTHREADS_CALCINC}"

 ;;
 "esfc")

    export NTHREADS_ESFC=${NTHREADSmax}
    export APRUN_ESFC="${APRUN_default} --cpus-per-task=${NTHREADS_ESFC}"

    export NTHREADS_CYCLE=${threads_per_task_cycle:-14}
    if [[ ${NTHREADS_CYCLE} -gt ${max_tasks_per_node} ]]; then
        export NTHREADS_CYCLE=${max_tasks_per_node}
    fi
    export APRUN_CYCLE="${APRUN_default} --cpus-per-task=${NTHREADS_CYCLE}"

 ;;
 "epos")

    export NTHREADS_EPOS=${NTHREADSmax}
    export APRUN_EPOS="${APRUN_default} --cpus-per-task=${NTHREADS_EPOS}"

 ;;
 "postsnd")

    export CFP_MP="YES"

    export NTHREADS_POSTSND=${NTHREADS1}
    export APRUN_POSTSND="${APRUN_default} --cpus-per-task=${NTHREADS_POSTSND}"

    export NTHREADS_POSTSNDCFP=${threads_per_task_postsndcfp:-1}
    if [[ ${NTHREADS_POSTSNDCFP} -gt ${max_threads_per_task} ]]; then
        export NTHREADS_POSTSNDCFP=${max_threads_per_task}
    fi
    export APRUN_POSTSNDCFP="${launcher} -n ${ntasks_postsndcfp} ${mpmd_opt}"

 ;;
 "awips")

    export NTHREADS_AWIPS=${NTHREADS1}
    export APRUN_AWIPSCFP="${APRUN_default} ${mpmd_opt}"

 ;;
 "gempak")

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

 ;;
 "fit2obs")

    export NTHREADS_FIT2OBS=${NTHREADS1}
    export MPIRUN="${APRUN_default} --cpus-per-task=${NTHREADS_FIT2OBS}"

 ;;
 *)
    # Some other job not yet defined here
    echo "WARNING: The job step ${step} does not specify GAEAC5-specific resources"
 ;;
esac