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

if [[ "${step}" = "waveinit" ]]; then

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

elif [[ "${step}" = "fcst" ]]; 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}" = "atmos_products" ]]; then

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

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 --cpus-per-task=${NTHREADS_OCNICEPOST}"

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} --cpus-per-task=${NTHREADS_FIT2OBS}"

fi