├── .gitmodules ├── LICENSE ├── README.md ├── build ├── cmakeclean.sh ├── cmakescript.sh └── machines │ ├── aws │ ├── aws_a100_gpu.env │ ├── job_1.sh │ ├── job_16.sh │ ├── job_2.sh │ ├── job_32.sh │ ├── job_4.sh │ └── job_8.sh │ ├── azure │ ├── azure_a100_gpu.env │ └── job_8.sh │ ├── crusher │ ├── crusher_gpu.env │ ├── crusher_gpu_debug.env │ └── job_16.sh │ ├── fhqwhgads │ ├── fhqwhgads_cpu.env │ ├── fhqwhgads_debug.env │ ├── fhqwhgads_gpu.env │ └── fhqwhgads_gpu_debug.env │ ├── summit │ ├── job_4096.sh │ ├── job_example.sh │ ├── summit_debug.env │ ├── summit_gpu.env │ └── summit_gpu_debug.env │ └── thatchroof │ ├── thatchroof_cpu.env │ ├── thatchroof_debug.env │ ├── thatchroof_gpu.env │ ├── thatchroof_gpu_debug.env │ └── thatchroof_openmp.env ├── experiments ├── community_benchmark │ ├── CMakeLists.txt │ ├── driver.cpp │ └── inputs │ │ ├── input_euler3d.yaml │ │ ├── input_euler3d_1024x1024x100.yaml │ │ ├── input_euler3d_2048x2048x100.yaml │ │ └── input_euler3d_4096x4096x100.yaml ├── simple_city │ ├── CMakeLists.txt │ ├── custom_modules │ │ ├── horizontal_sponge.h │ │ └── time_averager.h │ ├── driver.cpp │ └── inputs │ │ ├── input_building.yaml │ │ └── input_city.yaml ├── supercell_example │ ├── CMakeLists.txt │ ├── driver.cpp │ └── inputs │ │ └── input_euler3d.yaml └── supercell_kessler_surrogate │ ├── CMakeLists.txt │ ├── README.md │ ├── custom_modules │ ├── CMakeLists.txt │ ├── gather_micro_statistics.h │ ├── generate_micro_surrogate_data.h │ └── microphysics_kessler_ponni.h │ ├── gather_statistics.cpp │ ├── generate_micro_data.cpp │ ├── inference_ponni.cpp │ ├── inputs │ ├── examples │ │ ├── supercell_kessler_singlecell_model_weights.h5 │ │ ├── supercell_kessler_stencil_input_scaling.txt │ │ └── supercell_kessler_stencil_output_scaling.txt │ └── input_euler3d.yaml │ └── jupyter_notebooks │ ├── README.md │ ├── images │ ├── NN_framework.png │ ├── microphysics_Kessler.png │ └── microphysics_ML_singlecell.png │ ├── kessler_netcdf_to_numpy.ipynb │ └── kessler_singlecell_train_example.ipynb └── model ├── CMakeLists.txt ├── core ├── CMakeLists.txt ├── DataManager.h ├── MultipleFields.h ├── Options.h └── coupler.h ├── main_header.h └── modules ├── CMakeLists.txt ├── column_nudging.h ├── dynamics_euler_stratified_wenofv.h ├── helpers ├── TransformMatrices.h ├── WenoLimiter.h ├── WenoLimiter_recon.h └── microphysics_p3 │ ├── CMakeLists.txt │ ├── bfb_math.inc │ ├── micro_p3.F90 │ ├── micro_p3_utils.F90 │ ├── p3_lookup_table_1.dat-v4.1.1 │ ├── scream_abortutils.F90 │ ├── scream_types.F90 │ └── wv_sat_scream.F90 ├── microphysics_kessler.h ├── microphysics_p3.h ├── perturb_temperature.h └── sponge_layer.h /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "external/YAKL"] 2 | path = external/YAKL 3 | url = git@github.com:mrnorman/YAKL.git 4 | [submodule "external/ponni"] 5 | path = external/ponni 6 | url = git@github.com:mrnorman/ponni.git 7 | [submodule "external/yaml-cpp"] 8 | path = external/yaml-cpp 9 | url = git@github.com:jbeder/yaml-cpp.git 10 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 2-Clause License 2 | 3 | Copyright (c) 2022, Matt Norman 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | 1. Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | 2. Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 17 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 18 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 19 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 20 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 21 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 22 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 23 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 24 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # miniWeatherML 2 | 3 | Welcome to miniWeatherML: A playground for learning and developing Machine Learning (ML) surrogate models and workflows. It is based on a simplified weather model simulating flows such as [supercells](https://en.wikipedia.org/wiki/Supercell) that are realistic enough to be challenging and simple enough for rapid prototyping in: 4 | * Data generation and curation 5 | * Machine Learning model training 6 | * ML model deployment and analysis 7 | * End-to-end workflows 8 | 9 | 10 | 11 | ## Documentation: https://github.com/mrnorman/miniWeatherML/wiki 12 | 13 | Author: Matt Norman (Oak Ridge National Laboratory), https://mrnorman.github.io 14 | 15 | Contributors so far: 16 | * Matt Norman (Oak Ridge National Laboratory) 17 | * Murali Gopalakrishnan Meena (Oak Ridge National Laboratory) 18 | 19 | Written in portable C++, miniWeatherML runs out of the box on CPUs as well as Nvidia, AMD, and Intel GPUs. 20 | 21 | The core infrastructure of miniWeatherML is less than 1K lines of code, and the minimal meaningful module set is comprised of less than 3K lines of code, very little of which needs to be understood in full detail in order to effectively use miniWeatherML for its intended purposes. 22 | 23 | ### Check out the **[Kessler Microphysics Example Workflow](https://github.com/mrnorman/miniWeatherML/tree/main/experiments/supercell_kessler_surrogate)** to get started 24 | 25 | 33 | -------------------------------------------------------------------------------- /build/cmakeclean.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | rm -rf CMakeCache.txt CMakeFiles cmake_install.cmake CTestTestfile.cmake Makefile yakl \ 4 | core modules custom_modules inputs model 5 | 6 | -------------------------------------------------------------------------------- /build/cmakescript.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ./cmakeclean.sh 4 | 5 | if [ $# -ne 1 ]; then 6 | echo "Error: must pass exactly one parameter giving the directory to build" 7 | exit -1 8 | fi 9 | 10 | if [ ! -d $1 ]; then 11 | echo "Error: Passed directory does not exist" 12 | exit -1 13 | fi 14 | 15 | if [ ! -f $1/CMakeLists.txt ]; then 16 | echo "Error: Passed directory does not contain a CMakeLists.txt file" 17 | exit -1 18 | fi 19 | 20 | cmake \ 21 | -DYAKL_CUDA_FLAGS="${YAKL_CUDA_FLAGS}" \ 22 | -DYAKL_CXX_FLAGS="${YAKL_CXX_FLAGS}" \ 23 | -DYAKL_SYCL_FLAGS="${YAKL_SYCL_FLAGS}" \ 24 | -DYAKL_OPENMP_FLAGS="${YAKL_OPENMP_FLAGS}" \ 25 | -DYAKL_HIP_FLAGS="${YAKL_HIP_FLAGS}" \ 26 | -DYAKL_F90_FLAGS="${YAKL_F90_FLAGS}" \ 27 | -DMW_LINK_FLAGS="${MW_LINK_FLAGS}" \ 28 | -DYAKL_ARCH="${YAKL_ARCH}" \ 29 | -DYAKL_HAVE_MPI=ON \ 30 | -DYAKL_DEBUG="${YAKL_DEBUG}" \ 31 | -DYAKL_PROFILE="${YAKL_PROFILE}" \ 32 | -DYAKL_AUTO_PROFILE="${YAKL_AUTO_PROFILE}" \ 33 | -DYAKL_VERBOSE="${YAKL_VERBOSE}" \ 34 | -DYAKL_VERBOSE_FILE="${YAKL_VERBOSE_FILE}" \ 35 | -DYAKL_AUTO_FENCE="${YAKL_AUTO_FENCE}" \ 36 | -DCMAKE_CUDA_HOST_COMPILER="${CXX}" \ 37 | -DMINIWEATHER_ML_HOME="`pwd`/.." \ 38 | -Wno-dev \ 39 | $1 40 | 41 | ln -sf $1/inputs . 42 | 43 | -------------------------------------------------------------------------------- /build/machines/aws/aws_a100_gpu.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | source /usr/share/Modules/init/bash 4 | module use /home/$USER/spack/share/spack/modules/linux-amzn2-cascadelake 5 | module use /home/proj-share/spack/share/spack/modules/linux-amzn2-skylake_avx512 6 | module use /home/$USER/spack/share/spack/modules/linux-amzn2-skylake_avx512 7 | module purge 8 | module load cuda-11.4.2-gcc-9.3.0-2uqqcue cmake-3.23.1-gcc-9.3.0-s74zwan netcdf-c-4.8.1-gcc-9.3.0-jv4lbmr parallel-netcdf-1.12.2-gcc-9.3.0-cgdofig openmpi-4.1.4-gcc-7.3.1-v5liykd gcc-9.3.0-gcc-7.3.1-pzkp5qd 9 | module list 10 | 11 | export PARALLEL_NETCDF_PATH=/home/olcf5/spack/opt/spack/linux-amzn2-cascadelake/gcc-9.3.0/parallel-netcdf-1.12.2-cgdofigdd2zxpuv7muubix5k7qcyxc7w 12 | 13 | export OMPI_CXX=g++ 14 | 15 | export CC=mpicc 16 | export FC=gfortran 17 | export CXX=mpic++ 18 | 19 | export YAKL_ARCH=CUDA 20 | 21 | export YAKL_CUDA_FLAGS="-DMW_ORD=3 -DHAVE_MPI -O3 -arch sm_80 --use_fast_math -DYAKL_PROFILE -I`nc-config --includedir` -I${PARALLEL_NETCDF_PATH}/include -ccbin mpic++" 22 | export YAKL_F90_FLAGS="-O2 -ffree-line-length-none -DSCREAM_DOUBLE_PRECISION" 23 | export MW_LINK_FLAGS="`nc-config --libs` -L${PARALLEL_NETCDF_PATH}/lib -lpnetcdf" 24 | export GATOR_INITIAL_MB=39000 25 | 26 | unset CXXFLAGS 27 | unset FFLAGS 28 | unset F77FLAGS 29 | unset F90FLAGS 30 | -------------------------------------------------------------------------------- /build/machines/aws/job_1.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH -J miniWeatherML 3 | #SBATCH -N 1 4 | #SBATCH -t 00:15:00 5 | #SBATCH --partition eval-gpu 6 | #SBATCH --exclusive 7 | 8 | nodes=1 9 | mydir=job_$nodes 10 | ntasks=`echo "8*$nodes" | bc` 11 | 12 | cd /home/$USER/miniWeatherML/build 13 | source machines/aws/aws_a100_gpu.env 14 | 15 | echo "Nodes: $nodes" 16 | echo "Tasks: $ntasks" 17 | echo "Dir: `pwd`/$mydir" 18 | 19 | mkdir -p $mydir 20 | cp ./driver $mydir 21 | cp ./inputs/* $mydir 22 | cd $mydir 23 | 24 | export OMP_NUM_THREADS=1 25 | nvidia-smi 26 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 /home/$USER/hello_jobstep.exe 2>&1 | tee hello_jsrun_output.txt 27 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_1024x1024x100.yaml 2>&1 | tee job_output_1024x1024x100.txt 28 | 29 | 30 | -------------------------------------------------------------------------------- /build/machines/aws/job_16.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH -J miniWeatherML 3 | #SBATCH -N 16 4 | #SBATCH -t 00:40:00 5 | #SBATCH --partition eval-gpu 6 | #SBATCH --exclusive 7 | 8 | nodes=16 9 | mydir=job_$nodes 10 | ntasks=`echo "8*$nodes" | bc` 11 | 12 | cd /home/$USER/miniWeatherML/build 13 | source machines/aws/aws_a100_gpu.env 14 | 15 | echo "Nodes: $nodes" 16 | echo "Tasks: $ntasks" 17 | echo "Dir: `pwd`/$mydir" 18 | 19 | mkdir -p $mydir 20 | cp ./driver $mydir 21 | cp ./inputs/* $mydir 22 | cd $mydir 23 | 24 | export OMP_NUM_THREADS=1 25 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_1024x1024x100.yaml 2>&1 | tee job_output_1024x1024x100.txt 26 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_2048x2048x100.yaml 2>&1 | tee job_output_2048x2048x100.txt 27 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_4096x4096x100.yaml 2>&1 | tee job_output_4096x4096x100.txt 28 | 29 | 30 | -------------------------------------------------------------------------------- /build/machines/aws/job_2.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH -J miniWeatherML 3 | #SBATCH -N 2 4 | #SBATCH -t 00:30:00 5 | #SBATCH --partition eval-gpu 6 | #SBATCH --exclusive 7 | 8 | nodes=2 9 | mydir=job_$nodes 10 | ntasks=`echo "8*$nodes" | bc` 11 | 12 | cd /home/$USER/miniWeatherML/build 13 | source machines/aws/aws_a100_gpu.env 14 | 15 | echo "Nodes: $nodes" 16 | echo "Tasks: $ntasks" 17 | echo "Dir: `pwd`/$mydir" 18 | 19 | mkdir -p $mydir 20 | cp ./driver $mydir 21 | cp ./inputs/* $mydir 22 | cd $mydir 23 | 24 | export OMP_NUM_THREADS=1 25 | nvidia-smi 26 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 /home/$USER/hello_jobstep.exe 2>&1 | tee hello_jsrun_output.txt 27 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_1024x1024x100.yaml 2>&1 | tee job_output_1024x1024x100.txt 28 | 29 | 30 | -------------------------------------------------------------------------------- /build/machines/aws/job_32.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH -J miniWeatherML 3 | #SBATCH -N 32 4 | #SBATCH -t 00:40:00 5 | #SBATCH --partition eval-gpu 6 | #SBATCH --exclusive 7 | 8 | nodes=32 9 | mydir=job_$nodes 10 | ntasks=`echo "8*$nodes" | bc` 11 | 12 | cd /home/$USER/miniWeatherML/build 13 | source machines/aws/aws_a100_gpu.env 14 | 15 | echo "Nodes: $nodes" 16 | echo "Tasks: $ntasks" 17 | echo "Dir: `pwd`/$mydir" 18 | 19 | mkdir -p $mydir 20 | cp ./driver $mydir 21 | cp ./inputs/* $mydir 22 | cd $mydir 23 | 24 | export OMP_NUM_THREADS=1 25 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_1024x1024x100.yaml 2>&1 | tee job_output_1024x1024x100.txt 26 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_2048x2048x100.yaml 2>&1 | tee job_output_2048x2048x100.txt 27 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_4096x4096x100.yaml 2>&1 | tee job_output_4096x4096x100.txt 28 | 29 | 30 | -------------------------------------------------------------------------------- /build/machines/aws/job_4.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH -J miniWeatherML 3 | #SBATCH -N 4 4 | #SBATCH -t 00:30:00 5 | #SBATCH --partition eval-gpu 6 | #SBATCH --exclusive 7 | 8 | nodes=4 9 | mydir=job_$nodes 10 | ntasks=`echo "8*$nodes" | bc` 11 | 12 | cd /home/$USER/miniWeatherML/build 13 | source machines/aws/aws_a100_gpu.env 14 | 15 | echo "Nodes: $nodes" 16 | echo "Tasks: $ntasks" 17 | echo "Dir: `pwd`/$mydir" 18 | 19 | mkdir -p $mydir 20 | cp ./driver $mydir 21 | cp ./inputs/* $mydir 22 | cd $mydir 23 | 24 | export OMP_NUM_THREADS=1 25 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_1024x1024x100.yaml 2>&1 | tee job_output_1024x1024x100.txt 26 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_2048x2048x100.yaml 2>&1 | tee job_output_2048x2048x100.txt 27 | 28 | 29 | -------------------------------------------------------------------------------- /build/machines/aws/job_8.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH -J miniWeatherML 3 | #SBATCH -N 8 4 | #SBATCH -t 00:40:00 5 | #SBATCH --partition eval-gpu 6 | #SBATCH --exclusive 7 | 8 | nodes=8 9 | mydir=job_$nodes 10 | ntasks=`echo "8*$nodes" | bc` 11 | 12 | cd /home/$USER/miniWeatherML/build 13 | source machines/aws/aws_a100_gpu.env 14 | 15 | echo "Nodes: $nodes" 16 | echo "Tasks: $ntasks" 17 | echo "Dir: `pwd`/$mydir" 18 | 19 | mkdir -p $mydir 20 | cp ./driver $mydir 21 | cp ./inputs/* $mydir 22 | cd $mydir 23 | 24 | export OMP_NUM_THREADS=1 25 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_1024x1024x100.yaml 2>&1 | tee job_output_1024x1024x100.txt 26 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_2048x2048x100.yaml 2>&1 | tee job_output_2048x2048x100.txt 27 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_4096x4096x100.yaml 2>&1 | tee job_output_4096x4096x100.txt 28 | 29 | 30 | -------------------------------------------------------------------------------- /build/machines/azure/azure_a100_gpu.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | source /usr/share/modules/init/bash 4 | module use /shared/home/$USER/spack/share/spack/modules/linux-ubuntu18.04-zen 5 | module use /shared/home/$USER/spack/share/spack/modules/linux-ubuntu18.04-zen2 6 | module purge 7 | module load cmake-3.23.1-gcc-9.3.0-wroitcx cuda-11.4.2-gcc-9.3.0-eelhg76 netcdf-c-4.8.1-gcc-9.3.0-vxigvnu parallel-netcdf-1.12.2-gcc-9.3.0-7lrswsj mpi/openmpi gcc-9.3.0-gcc-7.5.0-6hxg3hz 8 | module list 9 | 10 | export PARALLEL_NETCDF_PATH=/shared/home/normanmr/spack/opt/spack/linux-ubuntu18.04-zen2/gcc-9.3.0/parallel-netcdf-1.12.2-7lrswsj5jmicqsefovhvyojk6ef7b6xl 11 | 12 | export OMPI_CXX=g++ 13 | 14 | export CC=mpicc 15 | export FC=gfortran 16 | export CXX=mpic++ 17 | 18 | export YAKL_ARCH=CUDA 19 | 20 | export YAKL_CUDA_FLAGS="-DMW_ORD=3 -DHAVE_MPI -O3 -arch sm_80 --use_fast_math -DYAKL_PROFILE -I`nc-config --includedir` -I${PARALLEL_NETCDF_PATH}/include -ccbin mpic++" 21 | export YAKL_F90_FLAGS="-O2 -ffree-line-length-none -DSCREAM_DOUBLE_PRECISION" 22 | export MW_LINK_FLAGS="`nc-config --libs` -L${PARALLEL_NETCDF_PATH}/lib -lpnetcdf" 23 | export GATOR_INITIAL_MB=39000 24 | 25 | unset CXXFLAGS 26 | unset FFLAGS 27 | unset F77FLAGS 28 | unset F90FLAGS 29 | -------------------------------------------------------------------------------- /build/machines/azure/job_8.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH -J miniWeatherML 3 | #SBATCH -N 8 4 | #SBATCH -t 00:25:00 5 | #SBATCH --partition hpc 6 | #SBATCH --exclusive 7 | 8 | nodes=8 9 | mydir=job_$nodes 10 | ntasks=`echo "8*$nodes" | bc` 11 | 12 | cd /shared/home/$USER/miniWeatherML/build 13 | source machines/azure/azure_a100_gpu.env 14 | 15 | echo "Nodes: $nodes" 16 | echo "Tasks: $ntasks" 17 | echo "Dir: `pwd`/$mydir" 18 | 19 | mkdir -p $mydir 20 | cp ./driver $mydir 21 | cp ./inputs/* $mydir 22 | cd $mydir 23 | 24 | export OMP_NUM_THREADS=1 25 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_1024x1024x100.yaml 2>&1 | tee job_output_1024x1024x100.txt 26 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_2048x2048x100.yaml 2>&1 | tee job_output_2048x2048x100.txt 27 | srun -N $nodes -n $ntasks --gpus-per-task=1 -c 1 ./driver ./input_euler3d_4096x4096x100.yaml 2>&1 | tee job_output_4096x4096x100.txt 28 | 29 | 30 | -------------------------------------------------------------------------------- /build/machines/crusher/crusher_gpu.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | source ${MODULESHOME}/init/bash 4 | module load PrgEnv-amd cray-parallel-netcdf cmake craype-accel-amd-gfx90a cray-hdf5 cray-netcdf 5 | 6 | unset CXXFLAGS 7 | unset FFLAGS 8 | unset F77FLAGS 9 | unset F90FLAGS 10 | 11 | export CC=cc 12 | export FC=ftn 13 | export CXX=CC 14 | 15 | export YAKL_ARCH=HIP 16 | 17 | export MPICH_GPU_SUPPORT_ENABLED=1 18 | 19 | export YAKL_HIP_FLAGS="-DMW_GPU_AWARE_MPI -munsafe-fp-atomics -O3 -I${ROCM_PATH}/include -D__HIP_ROCclr__ -D__HIP_ARCH_GFX90A__=1 --rocm-path=${ROCM_PATH} --offload-arch=gfx90a -x hip -Wno-unused-result -Wno-macro-redefined" 20 | export YAKL_F90_FLAGS="-O2 -DSCREAM_DOUBLE_PRECISION" 21 | export MW_LINK_FLAGS="--rocm-path=${ROCM_PATH} -L${ROCM_PATH}/lib -lamdhip64" 22 | export YAKL_DEBUG=OFF 23 | export YAKL_PROFILE=ON 24 | 25 | unset CXXFLAGS 26 | unset FFLAGS 27 | unset F77FLAGS 28 | unset F90FLAGS 29 | 30 | # -I`nc-config --includedir` -I$PNETCDF_DIR/include 31 | # `nc-config --libs` -L$PNETCDF_DIR/lib -lpnetcdf 32 | -------------------------------------------------------------------------------- /build/machines/crusher/crusher_gpu_debug.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | source ${MODULESHOME}/init/bash 4 | module load PrgEnv-amd cray-parallel-netcdf cmake craype-accel-amd-gfx90a cray-hdf5 cray-netcdf 5 | 6 | unset CXXFLAGS 7 | unset FFLAGS 8 | unset F77FLAGS 9 | unset F90FLAGS 10 | 11 | export CC=cc 12 | export FC=ftn 13 | export CXX=CC 14 | 15 | export YAKL_ARCH=HIP 16 | 17 | export MPICH_GPU_SUPPORT_ENABLED=1 18 | 19 | export YAKL_HIP_FLAGS="-DMW_GPU_AWARE_MPI -munsafe-fp-atomics -O0 -g -I${ROCM_PATH}/include -D__HIP_ROCclr__ -D__HIP_ARCH_GFX90A__=1 --rocm-path=${ROCM_PATH} --offload-arch=gfx90a -x hip -Wno-unused-result -Wno-macro-redefined" 20 | export YAKL_F90_FLAGS="-O0 -g -DSCREAM_DOUBLE_PRECISION" 21 | export MW_LINK_FLAGS="--rocm-path=${ROCM_PATH} -L${ROCM_PATH}/lib -lamdhip64" 22 | export YAKL_DEBUG=ON 23 | export YAKL_PROFILE=ON 24 | 25 | unset CXXFLAGS 26 | unset FFLAGS 27 | unset F77FLAGS 28 | unset F90FLAGS 29 | 30 | # -I`nc-config --includedir` -I$PNETCDF_DIR/include 31 | # `nc-config --libs` -L$PNETCDF_DIR/lib -lpnetcdf 32 | -------------------------------------------------------------------------------- /build/machines/crusher/job_16.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #SBATCH -A stf006 3 | #SBATCH -J miniWeatherML 4 | #SBATCH -o %x-%j.out 5 | #SBATCH -t 02:00:00 6 | #SBATCH -p batch 7 | #SBATCH -N 16 8 | 9 | NODES=16 10 | MW_HOME=/ccs/home/imn/miniWeatherML 11 | # JOBDIR=/lustre/orion/stf006/scratch/imn/miniWeatherML_jobs/nodes_$NODES 12 | JOBDIR=/gpfs/alpine/stf006/scratch/imn/miniWeatherML_jobs/nodes_$NODES 13 | 14 | TASKS=`echo "8*$NODES" | bc` 15 | echo "*** USING $TASKS TASKS ***" 16 | 17 | mkdir -p $JOBDIR 18 | cd $JOBDIR 19 | source $MW_HOME/build/machines/crusher/crusher_gpu.env 20 | cp -f $MW_HOME/build/driver* . 21 | cp -Lrf $MW_HOME/build/input* . 22 | 23 | srun -N $NODES -n $TASKS -c 1 --gpus-per-node=8 --gpus-per-task=1 --gpu-bind=closest ./driver ./inputs/input_euler3d.yaml | tee output.txt 24 | 25 | -------------------------------------------------------------------------------- /build/machines/fhqwhgads/fhqwhgads_cpu.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | unset YAKL_ARCH 4 | 5 | export CC=mpicc 6 | export CXX=mpic++ 7 | export FC=mpif90 8 | 9 | export YAKL_CXX_FLAGS="-O2 -DHAVE_MPI -I/usr/include/hdf5/serial" 10 | export YAKL_F90_FLAGS="-O2 -ffree-line-length-none" 11 | export MW_LINK_FLAGS="-lpnetcdf -lnetcdf -L/usr/lib/x86_64-linux-gnu/hdf5/serial -lhdf5" 12 | 13 | unset CXXFLAGS 14 | unset FFLAGS 15 | unset F77FLAGS 16 | unset F90FLAGS 17 | -------------------------------------------------------------------------------- /build/machines/fhqwhgads/fhqwhgads_debug.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | unset YAKL_ARCH 4 | 5 | export CC=mpicc 6 | export CXX=mpic++ 7 | export FC=mpif90 8 | 9 | export YAKL_CXX_FLAGS="-DHAVE_MPI -g -O0 -DYAKL_DEBUG -I/usr/include/hdf5/serial" 10 | export YAKL_F90_FLAGS="-O0 -g -ffree-line-length-none" 11 | export MW_LINK_FLAGS="-lpnetcdf -lnetcdf -L/usr/lib/x86_64-linux-gnu/hdf5/serial -lhdf5" 12 | 13 | unset CXXFLAGS 14 | unset FFLAGS 15 | unset F77FLAGS 16 | unset F90FLAGS 17 | -------------------------------------------------------------------------------- /build/machines/fhqwhgads/fhqwhgads_gpu.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | export YAKL_ARCH=CUDA 4 | 5 | export CC=mpicc 6 | export CXX=mpic++ 7 | export FC=mpif90 8 | 9 | export YAKL_CUDA_FLAGS="-DHAVE_MPI -O3 -arch sm_50 -DYAKL_PROFILE --use_fast_math -ccbin mpic++ -I/usr/include/hdf5/serial" 10 | export YAKL_F90_FLAGS="-O3 -ffree-line-length-none" 11 | export MW_LINK_FLAGS="-lpnetcdf -lnetcdf -L/usr/lib/x86_64-linux-gnu/hdf5/serial -lhdf5" 12 | 13 | unset CXXFLAGS 14 | unset FFLAGS 15 | unset F77FLAGS 16 | unset F90FLAGS 17 | -------------------------------------------------------------------------------- /build/machines/fhqwhgads/fhqwhgads_gpu_debug.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | export YAKL_ARCH=CUDA 4 | 5 | export CC=mpicc 6 | export CXX=mpic++ 7 | export FC=mpif90 8 | 9 | export YAKL_CUDA_FLAGS="-DHAVE_MPI -g -O0 -DYAKL_DEBUG -arch sm_50 -DYAKL_PROFILE --use_fast_math -ccbin mpic++ -I/usr/include/hdf5/serial" 10 | export YAKL_F90_FLAGS="-O0 -g -ffree-line-length-none" 11 | export MW_LINK_FLAGS="-lnetcdf -lpnetcdf -L/usr/lib/x86_64-linux-gnu/hdf5/serial -lhdf5" 12 | 13 | unset CXXFLAGS 14 | unset FFLAGS 15 | unset F77FLAGS 16 | unset F90FLAGS 17 | -------------------------------------------------------------------------------- /build/machines/summit/job_4096.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #BSUB -P stf006 3 | #BSUB -W 0:15 4 | #BSUB -nnodes 4096 5 | #BSUB -q debug 6 | #BSUB -J miniWeatherML 7 | #BSUB -o miniWeatherML.%J 8 | #BSUB -e miniWeatherML.%J 9 | 10 | nodes=4096 11 | mydir=job_$nodes 12 | ntasks=`echo "6*$nodes" | bc` 13 | 14 | cd /gpfs/alpine/proj-shared/stf006/imn/miniWeatherML/build 15 | source machines/summit/summit_gpu.env 16 | 17 | echo "Nodes: $nodes" 18 | echo "Tasks: $ntasks" 19 | echo "Dir: `pwd`/$mydir" 20 | 21 | mkdir -p $mydir 22 | cp ./driver $mydir 23 | cp ./inputs/* $mydir 24 | cd $mydir 25 | 26 | jsrun -r 6 -n $ntasks -a 1 -c 1 -g 1 ./driver ./input_euler3d_1024x1024x100.yaml 2>&1 | tee job_output_1024x1024x100.txt 27 | jsrun -r 6 -n $ntasks -a 1 -c 1 -g 1 ./driver ./input_euler3d_2048x2048x100.yaml 2>&1 | tee job_output_2048x2048x100.txt 28 | jsrun -r 6 -n $ntasks -a 1 -c 1 -g 1 ./driver ./input_euler3d_4096x4096x100.yaml 2>&1 | tee job_output_4096x4096x100.txt 29 | 30 | 31 | -------------------------------------------------------------------------------- /build/machines/summit/job_example.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | #BSUB -P stf006 3 | #BSUB -W 0:10 4 | #BSUB -nnodes 32 5 | #BSUB -q debug 6 | #BSUB -J miniWeatherML 7 | #BSUB -o miniWeatherML.%J 8 | #BSUB -e miniWeatherML.%J 9 | 10 | cd /gpfs/alpine/proj-shared/stf006/imn/miniWeatherML/build 11 | source machines/summit/summit_gpu.env 12 | 13 | jsrun -r 6 -n 192 -a 1 -c 1 -g 1 ./driver ./inputs/input_euler3d.yaml 14 | 15 | -------------------------------------------------------------------------------- /build/machines/summit/summit_debug.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | source $MODULESHOME/init/bash 4 | module purge 5 | # module load DefApps gcc/9.3.0 cuda/11.4.2 cmake hdf5 curl sz zlib netcdf-c parallel-netcdf 6 | module load DefApps gcc/9.3.0 cuda/11.0.3 cmake hdf5 curl sz zlib netcdf-c parallel-netcdf 7 | 8 | export CC=mpicc 9 | export FC=mpif90 10 | export CXX=mpic++ 11 | 12 | unset YAKL_ARCH 13 | 14 | export YAKL_CXX_FLAGS="-DMW_GPU_AWARE_MPI -DHAVE_MPI -O0 -g -DYAKL_DEBUG -I${OLCF_NETCDF_C_ROOT}/include -I${OLCF_PARALLEL_NETCDF_ROOT}/include" 15 | export YAKL_F90_FLAGS="-O0 -g -ffree-line-length-none -DSCREAM_DOUBLE_PRECISION" 16 | export MW_LINK_FLAGS="`${OLCF_NETCDF_C_ROOT}/bin/nc-config --libs` -L${OLCF_PARALLEL_NETCDF_ROOT}/lib -lpnetcdf" 17 | 18 | unset CXXFLAGS 19 | unset FFLAGS 20 | unset F77FLAGS 21 | unset F90FLAGS 22 | -------------------------------------------------------------------------------- /build/machines/summit/summit_gpu.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | source $MODULESHOME/init/bash 4 | module purge 5 | # module load DefApps gcc/9.3.0 cuda/11.4.2 cmake hdf5 curl sz zlib netcdf-c parallel-netcdf 6 | module load DefApps gcc/9.3.0 cuda/11.0.3 cmake hdf5 curl sz zlib netcdf-c parallel-netcdf 7 | 8 | export CC=mpicc 9 | export FC=mpif90 10 | export CXX=mpic++ 11 | 12 | export YAKL_ARCH=CUDA 13 | 14 | export YAKL_CUDA_FLAGS="-DMW_GPU_AWARE_MPI -DHAVE_MPI -O3 -arch sm_70 --use_fast_math -DYAKL_PROFILE -I${OLCF_NETCDF_C_ROOT}/include -I${OLCF_PARALLEL_NETCDF_ROOT}/include -ccbin mpic++" 15 | export YAKL_F90_FLAGS="-O2 -ffree-line-length-none -DSCREAM_DOUBLE_PRECISION" 16 | export MW_LINK_FLAGS="`${OLCF_NETCDF_C_ROOT}/bin/nc-config --libs` -L${OLCF_PARALLEL_NETCDF_ROOT}/lib -lpnetcdf" 17 | 18 | unset CXXFLAGS 19 | unset FFLAGS 20 | unset F77FLAGS 21 | unset F90FLAGS 22 | -------------------------------------------------------------------------------- /build/machines/summit/summit_gpu_debug.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | source $MODULESHOME/init/bash 4 | module purge 5 | # module load DefApps gcc/9.3.0 cuda/11.4.2 cmake hdf5 curl sz zlib netcdf-c parallel-netcdf 6 | module load DefApps gcc/9.3.0 cuda/11.0.3 cmake hdf5 curl sz zlib netcdf-c parallel-netcdf 7 | 8 | export CC=mpicc 9 | export FC=mpif90 10 | export CXX=mpic++ 11 | 12 | export YAKL_ARCH=CUDA 13 | 14 | export YAKL_CUDA_FLAGS="-DMW_GPU_AWARE_MPI -DHAVE_MPI -O0 -g -DYAKL_DEBUG -arch sm_70 -I${OLCF_NETCDF_C_ROOT}/include -I${OLCF_PARALLEL_NETCDF_ROOT}/include -ccbin mpic++" 15 | export YAKL_F90_FLAGS="-O0 -g -ffree-line-length-none -DSCREAM_DOUBLE_PRECISION" 16 | export MW_LINK_FLAGS="`${OLCF_NETCDF_C_ROOT}/bin/nc-config --libs` -L${OLCF_PARALLEL_NETCDF_ROOT}/lib -lpnetcdf" 17 | 18 | unset CXXFLAGS 19 | unset FFLAGS 20 | unset F77FLAGS 21 | unset F90FLAGS 22 | -------------------------------------------------------------------------------- /build/machines/thatchroof/thatchroof_cpu.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | spack unload --all 4 | 5 | source $MODULESHOME/init/bash 6 | module purge 7 | module load netcdf-c-4.9.2-gcc-11.1.0-mvu6i6y \ 8 | parallel-netcdf-1.12.2-gcc-11.1.0-33fpcbm 9 | 10 | export HDF5_PATH=`which h5cc | xargs dirname`/.. 11 | export NETCDF_C_PATH=`nc-config --includedir`/.. 12 | export PNETCDF_PATH=`pnetcdf-config --includedir`/.. 13 | 14 | unset YAKL_ARCH 15 | 16 | export CC=mpicc 17 | export CXX=mpic++ 18 | export FC=mpif90 19 | 20 | export YAKL_CUDA_FLAGS="-O3 -I$PNETCDF_PATH/include -I`nc-config --includedir` -I$HDF5_PATH/include" 21 | export YAKL_F90_FLAGS="-O3 -DSCREAM_DOUBLE_PRECISION -ffree-line-length-none" 22 | export MW_LINK_FLAGS="`nc-config --libs` -L$PNETCDF_PATH/lib -lpnetcdf -L$HDF5_PATH/lib -lhdf5" 23 | export YAKL_DEBUG=OFF 24 | 25 | unset CXXFLAGS 26 | unset FFLAGS 27 | unset F77FLAGS 28 | unset F90FLAGS 29 | 30 | -------------------------------------------------------------------------------- /build/machines/thatchroof/thatchroof_debug.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | spack unload --all 4 | 5 | source $MODULESHOME/init/bash 6 | module purge 7 | module load netcdf-c-4.9.2-gcc-11.1.0-mvu6i6y \ 8 | parallel-netcdf-1.12.2-gcc-11.1.0-33fpcbm 9 | 10 | export HDF5_PATH=`which h5cc | xargs dirname`/.. 11 | export NETCDF_C_PATH=`nc-config --includedir`/.. 12 | export PNETCDF_PATH=`pnetcdf-config --includedir`/.. 13 | 14 | unset YAKL_ARCH 15 | 16 | export CC=mpicc 17 | export CXX=mpic++ 18 | export FC=mpif90 19 | 20 | export YAKL_CXX_FLAGS="-O0 -g -I$PNETCDF_PATH/include -I`nc-config --includedir` -I$HDF5_PATH/include" 21 | export YAKL_F90_FLAGS="-O0 -g -DSCREAM_DOUBLE_PRECISION -ffree-line-length-none" 22 | export MW_LINK_FLAGS="`nc-config --libs` -L$PNETCDF_PATH/lib -lpnetcdf -L$HDF5_PATH/lib -lhdf5" 23 | export YAKL_DEBUG=ON 24 | 25 | unset CXXFLAGS 26 | unset FFLAGS 27 | unset F77FLAGS 28 | unset F90FLAGS 29 | 30 | -------------------------------------------------------------------------------- /build/machines/thatchroof/thatchroof_gpu.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | spack unload --all 4 | 5 | source $MODULESHOME/init/bash 6 | module purge 7 | module load netcdf-c-4.9.2-gcc-11.1.0-mvu6i6y \ 8 | parallel-netcdf-1.12.2-gcc-11.1.0-33fpcbm 9 | 10 | export HDF5_PATH=`which h5cc | xargs dirname`/.. 11 | export NETCDF_C_PATH=`nc-config --includedir`/.. 12 | export PNETCDF_PATH=`pnetcdf-config --includedir`/.. 13 | 14 | export YAKL_ARCH=CUDA 15 | 16 | export CC=mpicc 17 | export CXX=mpic++ 18 | export FC=mpif90 19 | 20 | export YAKL_CUDA_FLAGS="-O3 -arch sm_86 --use_fast_math -ccbin mpic++ -I$PNETCDF_PATH/include -I`nc-config --includedir` -I$HDF5_PATH/include" 21 | export YAKL_F90_FLAGS="-O2 -DSCREAM_DOUBLE_PRECISION -ffree-line-length-none" 22 | export MW_LINK_FLAGS="`nc-config --libs` -L$PNETCDF_PATH/lib -lpnetcdf -L$HDF5_PATH/lib -lhdf5" 23 | export YAKL_DEBUG=OFF 24 | 25 | unset CXXFLAGS 26 | unset FFLAGS 27 | unset F77FLAGS 28 | unset F90FLAGS 29 | 30 | -------------------------------------------------------------------------------- /build/machines/thatchroof/thatchroof_gpu_debug.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | spack unload --all 4 | 5 | source $MODULESHOME/init/bash 6 | module purge 7 | module load netcdf-c-4.9.2-gcc-11.1.0-mvu6i6y \ 8 | parallel-netcdf-1.12.2-gcc-11.1.0-33fpcbm 9 | 10 | export HDF5_PATH=`which h5cc | xargs dirname`/.. 11 | export NETCDF_C_PATH=`nc-config --includedir`/.. 12 | export PNETCDF_PATH=`pnetcdf-config --includedir`/.. 13 | 14 | export YAKL_ARCH=CUDA 15 | 16 | export CC=mpicc 17 | export CXX=mpic++ 18 | export FC=mpif90 19 | 20 | export YAKL_CUDA_FLAGS="-O0 -g -arch sm_86 -ccbin mpic++ -I$PNETCDF_PATH/include -I`nc-config --includedir` -I$HDF5_PATH/include" 21 | export YAKL_F90_FLAGS="-O0 -g -DSCREAM_DOUBLE_PRECISION -ffree-line-length-none" 22 | export MW_LINK_FLAGS="`nc-config --libs` -L$PNETCDF_PATH/lib -lpnetcdf -L$HDF5_PATH/lib -lhdf5" 23 | export YAKL_DEBUG=ON 24 | 25 | unset CXXFLAGS 26 | unset FFLAGS 27 | unset F77FLAGS 28 | unset F90FLAGS 29 | 30 | -------------------------------------------------------------------------------- /build/machines/thatchroof/thatchroof_openmp.env: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | export YAKL_ARCH=OPENMP 4 | 5 | export CC=mpicc 6 | export CXX=mpic++ 7 | export FC=mpif90 8 | 9 | export YAKL_OPENMP_FLAGS="-fopenmp -O3 -DYAKL_DEBUG -I/usr/include/hdf5/serial" 10 | export YAKL_F90_FLAGS="-O2 -DSCREAM_DOUBLE_PRECISION -ffree-line-length-none" 11 | export MW_LINK_FLAGS="`nc-config --libs` -L/usr/lib/x85_64-linux-gnu -lpnetcdf -L/usr/lib/x86_64-linux-gnu/hdf5/serial -lhdf5" 12 | 13 | unset CXXFLAGS 14 | unset FFLAGS 15 | unset F77FLAGS 16 | unset F90FLAGS 17 | 18 | -------------------------------------------------------------------------------- /experiments/community_benchmark/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.0) 2 | project(miniWeatherML) 3 | 4 | enable_language(Fortran) 5 | enable_language(CXX) 6 | enable_language(C) 7 | if ("${YAKL_ARCH}" STREQUAL "CUDA") 8 | enable_language(CUDA) 9 | endif() 10 | 11 | add_subdirectory(${MINIWEATHER_ML_HOME}/model model) 12 | 13 | add_executable(driver driver.cpp) 14 | include(${YAKL_HOME}/yakl_utils.cmake) 15 | yakl_process_target(driver) 16 | target_link_libraries(driver model "${MW_LINK_FLAGS}") 17 | 18 | -------------------------------------------------------------------------------- /experiments/community_benchmark/driver.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "coupler.h" 3 | #include "dynamics_euler_stratified_wenofv.h" 4 | #include "microphysics_kessler.h" 5 | #include "sponge_layer.h" 6 | #include "perturb_temperature.h" 7 | #include "column_nudging.h" 8 | 9 | int main(int argc, char** argv) { 10 | MPI_Init( &argc , &argv ); 11 | yakl::init(); 12 | { 13 | using yakl::intrinsics::abs; 14 | using yakl::intrinsics::maxval; 15 | yakl::timer_start("main"); 16 | 17 | // This holds all of the model's variables, dimension sizes, and options 18 | core::Coupler coupler; 19 | 20 | // Read the YAML input file for variables pertinent to running the driver 21 | if (argc <= 1) { endrun("ERROR: Must pass the input YAML filename as a parameter"); } 22 | std::string inFile(argv[1]); 23 | YAML::Node config = YAML::LoadFile(inFile); 24 | if ( !config ) { endrun("ERROR: Invalid YAML input file"); } 25 | auto sim_time = config["sim_time"].as(); 26 | auto nens = config["nens" ].as(); 27 | auto nx_glob = config["nx_glob" ].as(); 28 | auto ny_glob = config["ny_glob" ].as(); 29 | auto nz = config["nz" ].as(); 30 | auto xlen = config["xlen" ].as(); 31 | auto ylen = config["ylen" ].as(); 32 | auto zlen = config["zlen" ].as(); 33 | auto dtphys_in = config["dt_phys" ].as(); 34 | 35 | coupler.set_option( "out_prefix" , config["out_prefix"].as() ); 36 | coupler.set_option( "init_data" , config["init_data"].as() ); 37 | coupler.set_option( "out_freq" , config["out_freq" ].as() ); 38 | 39 | // Coupler state is: (1) dry density; (2) u-velocity; (3) v-velocity; (4) w-velocity; (5) temperature 40 | // (6+) tracer masses (*not* mixing ratios!) 41 | coupler.distribute_mpi_and_allocate_coupled_state(nz, ny_glob, nx_glob, nens); 42 | 43 | // Just tells the coupler how big the domain is in each dimensions 44 | coupler.set_grid( xlen , ylen , zlen ); 45 | 46 | // This is for the dycore to pull out to determine how to do idealized test cases 47 | coupler.set_option( "standalone_input_file" , inFile ); 48 | 49 | // The column nudger nudges the column-average of the model state toward the initial column-averaged state 50 | // This is primarily for the supercell test case to keep the the instability persistently strong 51 | modules::ColumnNudger column_nudger; 52 | // Microphysics performs water phase changess + hydrometeor production, transport, collision, and aggregation 53 | modules::Microphysics_Kessler micro; 54 | // They dynamical core "dycore" integrates the Euler equations and performans transport of tracers 55 | modules::Dynamics_Euler_Stratified_WenoFV dycore; 56 | 57 | // Run the initialization modules 58 | micro .init ( coupler ); // Allocate micro state and register its tracers in the coupler 59 | dycore.init ( coupler ); // Dycore should initialize its own state here 60 | column_nudger.set_column ( coupler ); // Set the column before perturbing 61 | modules::perturb_temperature( coupler ); // Randomly perturb bottom layers of temperature to initiate convection 62 | 63 | real etime = 0; // Elapsed time 64 | 65 | real dtphys = dtphys_in; 66 | yakl::timer_start("simulation_loop"); 67 | while (etime < sim_time) { 68 | // If dtphys <= 0, then set it to the dynamical core's max stable time step 69 | if (dtphys_in <= 0.) { dtphys = dycore.compute_time_step(coupler); } 70 | // If we're about to go past the final time, then limit to time step to exactly hit the final time 71 | if (etime + dtphys > sim_time) { dtphys = sim_time - etime; } 72 | 73 | // Run the runtime modules 74 | dycore.time_step ( coupler , dtphys ); // Move the flow forward according to the Euler equations 75 | micro .time_step ( coupler , dtphys ); // Perform phase changes for water + precipitation / falling 76 | modules::sponge_layer ( coupler , dtphys ); // Damp spurious waves to the horiz. mean at model top 77 | column_nudger.nudge_to_column( coupler , dtphys ); // Nudge slightly back toward unstable profile 78 | // so that supercell persists for all time 79 | 80 | etime += dtphys; // Advance elapsed time 81 | } 82 | yakl::timer_stop("simulation_loop"); 83 | 84 | // TODO: Add finalize( coupler ) modules here 85 | 86 | yakl::timer_stop("main"); 87 | } 88 | yakl::finalize(); 89 | MPI_Finalize(); 90 | } 91 | 92 | 93 | -------------------------------------------------------------------------------- /experiments/community_benchmark/inputs/input_euler3d.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | # Simulation time in seconds 3 | sim_time: 100 # 86400 4 | 5 | # Number of cells to use in the CRMs 6 | # nx_glob: 1024 7 | # ny_glob: 1024 8 | # nz : 100 9 | 10 | nens: 1 11 | nx_glob: 4096 12 | ny_glob: 4096 13 | nz : 100 14 | 15 | # Domain size of the CRMs 16 | xlen: 819200 17 | ylen: 819200 18 | zlen: 20000 19 | 20 | init_data: supercell 21 | 22 | # Output filename 23 | out_prefix: test.nc 24 | 25 | # GCM time step 26 | dt_gcm: 900 27 | 28 | # CRM physics time step 29 | dt_phys: 0. 30 | 31 | # Output frequency in seconds 32 | out_freq: 200. 33 | 34 | 35 | -------------------------------------------------------------------------------- /experiments/community_benchmark/inputs/input_euler3d_1024x1024x100.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | # Simulation time in seconds 3 | sim_time: 900 # 86400 4 | 5 | # Number of cells to use in the CRMs 6 | # nx_glob: 1024 7 | # ny_glob: 1024 8 | # nz : 100 9 | 10 | nens : 1 11 | nx_glob: 1024 12 | ny_glob: 1024 13 | nz : 100 14 | 15 | # Domain size of the CRMs 16 | xlen: 819200 17 | ylen: 819200 18 | zlen: 20000 19 | 20 | init_data: supercell 21 | 22 | # Output filename 23 | out_prefix: test.nc 24 | 25 | # GCM time step 26 | dt_gcm: 900 27 | 28 | # CRM physics time step 29 | dt_phys: 0. 30 | 31 | # Output frequency in seconds 32 | out_freq: -1. 33 | 34 | 35 | -------------------------------------------------------------------------------- /experiments/community_benchmark/inputs/input_euler3d_2048x2048x100.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | # Simulation time in seconds 3 | sim_time: 250 # 86400 4 | 5 | # Number of cells to use in the CRMs 6 | # nx_glob: 1024 7 | # ny_glob: 1024 8 | # nz : 100 9 | 10 | nens : 1 11 | nx_glob: 2048 12 | ny_glob: 2048 13 | nz : 100 14 | 15 | # Domain size of the CRMs 16 | xlen: 819200 17 | ylen: 819200 18 | zlen: 20000 19 | 20 | init_data: supercell 21 | 22 | # Output filename 23 | out_prefix: test.nc 24 | 25 | # GCM time step 26 | dt_gcm: 900 27 | 28 | # CRM physics time step 29 | dt_phys: 0. 30 | 31 | # Output frequency in seconds 32 | out_freq: -1. 33 | 34 | 35 | -------------------------------------------------------------------------------- /experiments/community_benchmark/inputs/input_euler3d_4096x4096x100.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | # Simulation time in seconds 3 | sim_time: 100 # 86400 4 | 5 | # Number of cells to use in the CRMs 6 | # nx_glob: 1024 7 | # ny_glob: 1024 8 | # nz : 100 9 | 10 | nens : 1 11 | nx_glob: 4096 12 | ny_glob: 4096 13 | nz : 100 14 | 15 | # Domain size of the CRMs 16 | xlen: 819200 17 | ylen: 819200 18 | zlen: 20000 19 | 20 | init_data: supercell 21 | 22 | # Output filename 23 | out_prefix: test.nc 24 | 25 | # GCM time step 26 | dt_gcm: 900 27 | 28 | # CRM physics time step 29 | dt_phys: 0. 30 | 31 | # Output frequency in seconds 32 | out_freq: -1. 33 | 34 | 35 | -------------------------------------------------------------------------------- /experiments/simple_city/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.0) 2 | project(miniWeatherML) 3 | 4 | enable_language(Fortran) 5 | enable_language(CXX) 6 | enable_language(C) 7 | if ("${YAKL_ARCH}" STREQUAL "CUDA") 8 | enable_language(CUDA) 9 | endif() 10 | 11 | add_subdirectory(${MINIWEATHER_ML_HOME}/model model) 12 | 13 | add_executable(driver driver.cpp) 14 | include(${YAKL_HOME}/yakl_utils.cmake) 15 | yakl_process_target(driver) 16 | target_link_libraries(driver model "${MW_LINK_FLAGS}") 17 | target_include_directories(driver PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/custom_modules) 18 | 19 | -------------------------------------------------------------------------------- /experiments/simple_city/custom_modules/horizontal_sponge.h: -------------------------------------------------------------------------------- 1 | 2 | #pragma once 3 | 4 | #include "coupler.h" 5 | 6 | namespace custom_modules { 7 | 8 | struct Horizontal_Sponge { 9 | real2d col_rho_d; 10 | real2d col_uvel; 11 | real2d col_vvel; 12 | real2d col_wvel; 13 | real2d col_temp; 14 | real2d col_rho_v; 15 | int sponge_cells; 16 | real time_scale; 17 | 18 | inline void init( core::Coupler &coupler , int sponge_cells = 10 , real time_scale = 1 ) { 19 | using yakl::c::parallel_for; 20 | using yakl::c::Bounds; 21 | 22 | YAKL_SCOPE( col_rho_d , this->col_rho_d ); 23 | YAKL_SCOPE( col_uvel , this->col_uvel ); 24 | YAKL_SCOPE( col_vvel , this->col_vvel ); 25 | YAKL_SCOPE( col_wvel , this->col_wvel ); 26 | YAKL_SCOPE( col_temp , this->col_temp ); 27 | YAKL_SCOPE( col_rho_v , this->col_rho_v ); 28 | 29 | auto nens = coupler.get_nens(); 30 | auto nz = coupler.get_nz(); 31 | 32 | auto &dm = coupler.get_data_manager_readonly(); 33 | auto rho_d = dm.get("density_dry"); 34 | auto uvel = dm.get("uvel"); 35 | auto vvel = dm.get("vvel"); 36 | auto wvel = dm.get("wvel"); 37 | auto temp = dm.get("temp"); 38 | auto rho_v = dm.get("water_vapor"); 39 | 40 | col_rho_d = real2d("col_rho_d",nz,nens); 41 | col_uvel = real2d("col_uvel ",nz,nens); 42 | col_vvel = real2d("col_vvel ",nz,nens); 43 | col_wvel = real2d("col_wvel ",nz,nens); 44 | col_temp = real2d("col_temp ",nz,nens); 45 | col_rho_v = real2d("col_rho_v",nz,nens); 46 | 47 | auto col_rho_d_host = col_rho_d.createHostObject(); 48 | auto col_uvel_host = col_uvel .createHostObject(); 49 | auto col_vvel_host = col_vvel .createHostObject(); 50 | auto col_wvel_host = col_wvel .createHostObject(); 51 | auto col_temp_host = col_temp .createHostObject(); 52 | auto col_rho_v_host = col_rho_v.createHostObject(); 53 | 54 | if (coupler.is_mainproc()) { 55 | parallel_for( YAKL_AUTO_LABEL() , Bounds<2>(nz,nens) , YAKL_LAMBDA (int k, int iens) { 56 | col_rho_d(k,iens) = rho_d(k,0,0,iens); 57 | col_uvel (k,iens) = uvel (k,0,0,iens); 58 | col_vvel (k,iens) = vvel (k,0,0,iens); 59 | col_wvel (k,iens) = wvel (k,0,0,iens); 60 | col_temp (k,iens) = temp (k,0,0,iens); 61 | col_rho_v(k,iens) = rho_v(k,0,0,iens); 62 | }); 63 | col_rho_d.deep_copy_to(col_rho_d_host); 64 | col_uvel .deep_copy_to(col_uvel_host ); 65 | col_vvel .deep_copy_to(col_vvel_host ); 66 | col_wvel .deep_copy_to(col_wvel_host ); 67 | col_temp .deep_copy_to(col_temp_host ); 68 | col_rho_v.deep_copy_to(col_rho_v_host); 69 | yakl::fence(); 70 | } 71 | 72 | MPI_Bcast( col_rho_d_host.data(), col_rho_d_host.size() , coupler.get_mpi_data_type() , 0 , MPI_COMM_WORLD ); 73 | MPI_Bcast( col_uvel_host .data(), col_uvel_host .size() , coupler.get_mpi_data_type() , 0 , MPI_COMM_WORLD ); 74 | MPI_Bcast( col_vvel_host .data(), col_vvel_host .size() , coupler.get_mpi_data_type() , 0 , MPI_COMM_WORLD ); 75 | MPI_Bcast( col_wvel_host .data(), col_wvel_host .size() , coupler.get_mpi_data_type() , 0 , MPI_COMM_WORLD ); 76 | MPI_Bcast( col_temp_host .data(), col_temp_host .size() , coupler.get_mpi_data_type() , 0 , MPI_COMM_WORLD ); 77 | MPI_Bcast( col_rho_v_host.data(), col_rho_v_host.size() , coupler.get_mpi_data_type() , 0 , MPI_COMM_WORLD ); 78 | 79 | if (! coupler.is_mainproc()) { 80 | col_rho_d_host.deep_copy_to(col_rho_d); 81 | col_uvel_host .deep_copy_to(col_uvel ); 82 | col_vvel_host .deep_copy_to(col_vvel ); 83 | col_wvel_host .deep_copy_to(col_wvel ); 84 | col_temp_host .deep_copy_to(col_temp ); 85 | col_rho_v_host.deep_copy_to(col_rho_v); 86 | } 87 | this->sponge_cells = sponge_cells; 88 | this->time_scale = time_scale; 89 | } 90 | 91 | 92 | void override_rho_d(real val) { col_rho_d = val; } 93 | void override_uvel (real val) { col_uvel = val; } 94 | void override_vvel (real val) { col_vvel = val; } 95 | void override_wvel (real val) { col_wvel = val; } 96 | void override_temp (real val) { col_temp = val; } 97 | void override_rho_v(real val) { col_rho_v = val; } 98 | 99 | 100 | inline void apply( core::Coupler &coupler , real dt , 101 | bool x1=true , bool x2=true , bool y1=true , bool y2=true ) { 102 | using yakl::c::parallel_for; 103 | using yakl::c::Bounds; 104 | 105 | auto nens = coupler.get_nens(); 106 | auto nx = coupler.get_nx(); 107 | auto ny = coupler.get_ny(); 108 | auto nz = coupler.get_nz(); 109 | 110 | real R_d = coupler.get_option("R_d" ); 111 | real cp_d = coupler.get_option("cp_d"); 112 | real p0 = coupler.get_option("p0" ); 113 | real kappa = R_d / cp_d; 114 | real gamma = cp_d / (cp_d - R_d); 115 | real C0 = pow( R_d * pow( p0 , -kappa ) , gamma ); 116 | 117 | 118 | auto &dm = coupler.get_data_manager_readwrite(); 119 | auto rho_d = dm.get("density_dry"); 120 | auto uvel = dm.get("uvel"); 121 | auto vvel = dm.get("vvel"); 122 | auto wvel = dm.get("wvel"); 123 | auto temp = dm.get("temp"); 124 | auto rho_v = dm.get("water_vapor"); 125 | 126 | YAKL_SCOPE( sponge_cells , this->sponge_cells ); 127 | YAKL_SCOPE( time_scale , this->time_scale ); 128 | YAKL_SCOPE( col_rho_d , this->col_rho_d ); 129 | YAKL_SCOPE( col_uvel , this->col_uvel ); 130 | YAKL_SCOPE( col_vvel , this->col_vvel ); 131 | YAKL_SCOPE( col_wvel , this->col_wvel ); 132 | YAKL_SCOPE( col_temp , this->col_temp ); 133 | YAKL_SCOPE( col_rho_v , this->col_rho_v ); 134 | 135 | real time_factor = dt / time_scale; 136 | 137 | if (coupler.get_px() == 0 && x1) { 138 | parallel_for( YAKL_AUTO_LABEL() , Bounds<4>(nz,ny,nx,nens) , 139 | YAKL_LAMBDA (int k, int j, int i, int iens) { 140 | real xloc = i / (sponge_cells-1._fp); 141 | real weight = i < sponge_cells ? (cos(M_PI*xloc)+1)/2 : 0; 142 | weight *= time_factor; 143 | rho_d(k,j,i,iens) = weight*col_rho_d(k,iens) + (1-weight)*rho_d(k,j,i,iens); 144 | uvel (k,j,i,iens) = weight*col_uvel (k,iens) + (1-weight)*uvel (k,j,i,iens); 145 | vvel (k,j,i,iens) = weight*col_vvel (k,iens) + (1-weight)*vvel (k,j,i,iens); 146 | wvel (k,j,i,iens) = weight*col_wvel (k,iens) + (1-weight)*wvel (k,j,i,iens); 147 | temp (k,j,i,iens) = weight*col_temp (k,iens) + (1-weight)*temp (k,j,i,iens); 148 | rho_v(k,j,i,iens) = weight*col_rho_v(k,iens) + (1-weight)*rho_v(k,j,i,iens); 149 | }); 150 | } 151 | if (coupler.get_px() == coupler.get_nproc_x()-1 && x2) { 152 | parallel_for( YAKL_AUTO_LABEL() , Bounds<4>(nz,ny,nx,nens) , 153 | YAKL_LAMBDA (int k, int j, int i, int iens) { 154 | real xloc = (nx-1-i) / (sponge_cells-1._fp); 155 | real weight = nx-1-i < sponge_cells ? (cos(M_PI*xloc)+1)/2 : 0; 156 | weight *= time_factor; 157 | rho_d(k,j,i,iens) = weight*col_rho_d(k,iens) + (1-weight)*rho_d(k,j,i,iens); 158 | uvel (k,j,i,iens) = weight*col_uvel (k,iens) + (1-weight)*uvel (k,j,i,iens); 159 | vvel (k,j,i,iens) = weight*col_vvel (k,iens) + (1-weight)*vvel (k,j,i,iens); 160 | wvel (k,j,i,iens) = weight*col_wvel (k,iens) + (1-weight)*wvel (k,j,i,iens); 161 | temp (k,j,i,iens) = weight*col_temp (k,iens) + (1-weight)*temp (k,j,i,iens); 162 | rho_v(k,j,i,iens) = weight*col_rho_v(k,iens) + (1-weight)*rho_v(k,j,i,iens); 163 | }); 164 | } 165 | if (coupler.get_py() == 0 && y1) { 166 | parallel_for( YAKL_AUTO_LABEL() , Bounds<4>(nz,ny,nx,nens) , 167 | YAKL_LAMBDA (int k, int j, int i, int iens) { 168 | real yloc = j / (sponge_cells-1._fp); 169 | real weight = j < sponge_cells ? (cos(M_PI*yloc)+1)/2 : 0; 170 | weight *= time_factor; 171 | rho_d(k,j,i,iens) = weight*col_rho_d(k,iens) + (1-weight)*rho_d(k,j,i,iens); 172 | uvel (k,j,i,iens) = weight*col_uvel (k,iens) + (1-weight)*uvel (k,j,i,iens); 173 | vvel (k,j,i,iens) = weight*col_vvel (k,iens) + (1-weight)*vvel (k,j,i,iens); 174 | wvel (k,j,i,iens) = weight*col_wvel (k,iens) + (1-weight)*wvel (k,j,i,iens); 175 | temp (k,j,i,iens) = weight*col_temp (k,iens) + (1-weight)*temp (k,j,i,iens); 176 | rho_v(k,j,i,iens) = weight*col_rho_v(k,iens) + (1-weight)*rho_v(k,j,i,iens); 177 | }); 178 | } 179 | if (coupler.get_py() == coupler.get_nproc_y()-1 && y2) { 180 | parallel_for( YAKL_AUTO_LABEL() , Bounds<4>(nz,ny,nx,nens) , 181 | YAKL_LAMBDA (int k, int j, int i, int iens) { 182 | real yloc = (ny-1-j) / (sponge_cells-1._fp); 183 | real weight = ny-1-j < sponge_cells ? (cos(M_PI*yloc)+1)/2 : 0; 184 | weight *= time_factor; 185 | rho_d(k,j,i,iens) = weight*col_rho_d(k,iens) + (1-weight)*rho_d(k,j,i,iens); 186 | uvel (k,j,i,iens) = weight*col_uvel (k,iens) + (1-weight)*uvel (k,j,i,iens); 187 | vvel (k,j,i,iens) = weight*col_vvel (k,iens) + (1-weight)*vvel (k,j,i,iens); 188 | wvel (k,j,i,iens) = weight*col_wvel (k,iens) + (1-weight)*wvel (k,j,i,iens); 189 | temp (k,j,i,iens) = weight*col_temp (k,iens) + (1-weight)*temp (k,j,i,iens); 190 | rho_v(k,j,i,iens) = weight*col_rho_v(k,iens) + (1-weight)*rho_v(k,j,i,iens); 191 | }); 192 | } 193 | } 194 | }; 195 | } 196 | 197 | 198 | -------------------------------------------------------------------------------- /experiments/simple_city/custom_modules/time_averager.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "coupler.h" 4 | 5 | namespace custom_modules { 6 | 7 | struct Time_Averager { 8 | real etime; 9 | 10 | void init(core::Coupler &coupler) { 11 | auto nens = coupler.get_nens(); 12 | auto nx = coupler.get_nx(); 13 | auto ny = coupler.get_ny(); 14 | auto nz = coupler.get_nz(); 15 | 16 | auto &dm = coupler.get_data_manager_readwrite(); 17 | 18 | dm.register_and_allocate("time_avg_density_dry","",{nz,ny,nx,nens}); 19 | dm.register_and_allocate("time_avg_uvel" ,"",{nz,ny,nx,nens}); 20 | dm.register_and_allocate("time_avg_vvel" ,"",{nz,ny,nx,nens}); 21 | dm.register_and_allocate("time_avg_wvel" ,"",{nz,ny,nx,nens}); 22 | dm.register_and_allocate("time_avg_temp" ,"",{nz,ny,nx,nens}); 23 | dm.register_and_allocate("time_avg_water_vapor","",{nz,ny,nx,nens}); 24 | 25 | dm.get("time_avg_density_dry") = 0; 26 | dm.get("time_avg_uvel" ) = 0; 27 | dm.get("time_avg_vvel" ) = 0; 28 | dm.get("time_avg_wvel" ) = 0; 29 | dm.get("time_avg_temp" ) = 0; 30 | dm.get("time_avg_water_vapor") = 0; 31 | 32 | etime = 0.; 33 | } 34 | 35 | void accumulate(core::Coupler &coupler , real dt) { 36 | using yakl::c::parallel_for; 37 | using yakl::c::SimpleBounds; 38 | 39 | auto nens = coupler.get_nens(); 40 | auto nx = coupler.get_nx(); 41 | auto ny = coupler.get_ny(); 42 | auto nz = coupler.get_nz(); 43 | 44 | auto &dm = coupler.get_data_manager_readwrite(); 45 | 46 | auto rho_d = dm.get("density_dry"); 47 | auto uvel = dm.get("uvel" ); 48 | auto vvel = dm.get("vvel" ); 49 | auto wvel = dm.get("wvel" ); 50 | auto temp = dm.get("temp" ); 51 | auto rho_v = dm.get("water_vapor"); 52 | 53 | auto tavg_rho_d = dm.get("time_avg_density_dry"); 54 | auto tavg_uvel = dm.get("time_avg_uvel" ); 55 | auto tavg_vvel = dm.get("time_avg_vvel" ); 56 | auto tavg_wvel = dm.get("time_avg_wvel" ); 57 | auto tavg_temp = dm.get("time_avg_temp" ); 58 | auto tavg_rho_v = dm.get("time_avg_water_vapor"); 59 | 60 | double inertia = etime / (etime + dt); 61 | 62 | parallel_for( YAKL_AUTO_LABEL() , SimpleBounds<4>(nz,ny,nx,nens) , 63 | YAKL_LAMBDA (int k, int j, int i, int iens) { 64 | tavg_rho_d(k,j,i,iens) = inertia * tavg_rho_d(k,j,i,iens) + (1-inertia) * rho_d(k,j,i,iens); 65 | tavg_uvel (k,j,i,iens) = inertia * tavg_uvel (k,j,i,iens) + (1-inertia) * uvel (k,j,i,iens); 66 | tavg_vvel (k,j,i,iens) = inertia * tavg_vvel (k,j,i,iens) + (1-inertia) * vvel (k,j,i,iens); 67 | tavg_wvel (k,j,i,iens) = inertia * tavg_wvel (k,j,i,iens) + (1-inertia) * wvel (k,j,i,iens); 68 | tavg_temp (k,j,i,iens) = inertia * tavg_temp (k,j,i,iens) + (1-inertia) * temp (k,j,i,iens); 69 | tavg_rho_v(k,j,i,iens) = inertia * tavg_rho_v(k,j,i,iens) + (1-inertia) * rho_v(k,j,i,iens); 70 | }); 71 | 72 | etime += dt; 73 | } 74 | 75 | void finalize(core::Coupler &coupler ) { 76 | using yakl::c::parallel_for; 77 | using yakl::c::SimpleBounds; 78 | 79 | int nx_glob = coupler.get_nx_glob(); 80 | int ny_glob = coupler.get_ny_glob(); 81 | auto nens = coupler.get_nens(); 82 | auto nx = coupler.get_nx(); 83 | auto ny = coupler.get_ny(); 84 | auto nz = coupler.get_nz(); 85 | auto dx = coupler.get_dx(); 86 | auto dy = coupler.get_dy(); 87 | auto dz = coupler.get_dz(); 88 | int i_beg = coupler.get_i_beg(); 89 | int j_beg = coupler.get_j_beg(); 90 | 91 | auto &dm = coupler.get_data_manager_readonly(); 92 | 93 | yakl::SimplePNetCDF nc; 94 | nc.create("time_averaged_fields.nc" , NC_CLOBBER | NC_64BIT_DATA); 95 | 96 | nc.create_dim( "x" , nx_glob ); 97 | nc.create_dim( "y" , ny_glob ); 98 | nc.create_dim( "z" , nz ); 99 | 100 | nc.create_var( "x" , {"x"} ); 101 | nc.create_var( "y" , {"y"} ); 102 | nc.create_var( "z" , {"z"} ); 103 | nc.create_var( "density_dry" , {"z","y","x"} ); 104 | nc.create_var( "uvel" , {"z","y","x"} ); 105 | nc.create_var( "vvel" , {"z","y","x"} ); 106 | nc.create_var( "wvel" , {"z","y","x"} ); 107 | nc.create_var( "temp" , {"z","y","x"} ); 108 | nc.create_var( "water_vapor" , {"z","y","x"} ); 109 | 110 | nc.enddef(); 111 | 112 | // x-coordinate 113 | real1d xloc("xloc",nx); 114 | parallel_for( YAKL_AUTO_LABEL() , nx , YAKL_LAMBDA (int i) { xloc(i) = (i+i_beg+0.5)*dx; }); 115 | nc.write_all( xloc.createHostCopy() , "x" , {i_beg} ); 116 | 117 | // y-coordinate 118 | real1d yloc("yloc",ny); 119 | parallel_for( YAKL_AUTO_LABEL() , ny , YAKL_LAMBDA (int j) { yloc(j) = (j+j_beg+0.5)*dy; }); 120 | nc.write_all( yloc.createHostCopy() , "y" , {j_beg} ); 121 | 122 | // z-coordinate 123 | real1d zloc("zloc",nz); 124 | parallel_for( YAKL_AUTO_LABEL() , nz , YAKL_LAMBDA (int k) { zloc(k) = (k +0.5)*dz; }); 125 | nc.begin_indep_data(); 126 | if (coupler.is_mainproc()) { 127 | nc.write( zloc.createHostCopy() , "z" ); 128 | } 129 | nc.end_indep_data(); 130 | 131 | std::vector varnames = {"density_dry","uvel","vvel","wvel","temp","water_vapor"}; 132 | 133 | for (int i=0; i < varnames.size(); i++) { 134 | real3d data("data",nz,ny,nx); 135 | auto var = dm.get(std::string("time_avg_")+varnames[i]); 136 | parallel_for( SimpleBounds<3>(nz,ny,nx) , YAKL_LAMBDA (int k, int j, int i) { data(k,j,i) = var(k,j,i,0); }); 137 | nc.write_all(data.createHostCopy(),varnames[i],{0,j_beg,i_beg}); 138 | } 139 | 140 | nc.close(); 141 | } 142 | }; 143 | 144 | } 145 | 146 | 147 | -------------------------------------------------------------------------------- /experiments/simple_city/driver.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "coupler.h" 3 | #include "dynamics_euler_stratified_wenofv.h" 4 | #include "horizontal_sponge.h" 5 | #include "time_averager.h" 6 | #include "sponge_layer.h" 7 | 8 | int main(int argc, char** argv) { 9 | MPI_Init( &argc , &argv ); 10 | yakl::init(); 11 | { 12 | using yakl::intrinsics::abs; 13 | using yakl::intrinsics::maxval; 14 | yakl::timer_start("main"); 15 | 16 | // This holds all of the model's variables, dimension sizes, and options 17 | core::Coupler coupler; 18 | 19 | // Read the YAML input file for variables pertinent to running the driver 20 | if (argc <= 1) { endrun("ERROR: Must pass the input YAML filename as a parameter"); } 21 | std::string inFile(argv[1]); 22 | YAML::Node config = YAML::LoadFile(inFile); 23 | if ( !config ) { endrun("ERROR: Invalid YAML input file"); } 24 | auto sim_time = config["sim_time"].as(); 25 | auto nens = config["nens" ].as(); 26 | auto nx_glob = config["nx_glob" ].as(); 27 | auto ny_glob = config["ny_glob" ].as(); 28 | auto nz = config["nz" ].as(); 29 | auto xlen = config["xlen" ].as(); 30 | auto ylen = config["ylen" ].as(); 31 | auto zlen = config["zlen" ].as(); 32 | auto dtphys_in = config["dt_phys" ].as(); 33 | 34 | coupler.set_option( "out_prefix" , config["out_prefix" ].as() ); 35 | coupler.set_option( "init_data" , config["init_data" ].as() ); 36 | coupler.set_option( "out_freq" , config["out_freq" ].as() ); 37 | coupler.set_option( "enable_gravity" , config["enable_gravity" ].as(true)); 38 | coupler.set_option( "file_per_process", config["file_per_process"].as(false)); 39 | 40 | // Coupler state is: (1) dry density; (2) u-velocity; (3) v-velocity; (4) w-velocity; (5) temperature 41 | // (6+) tracer masses (*not* mixing ratios!) 42 | coupler.distribute_mpi_and_allocate_coupled_state(nz, ny_glob, nx_glob, nens); 43 | 44 | // Just tells the coupler how big the domain is in each dimensions 45 | coupler.set_grid( xlen , ylen , zlen ); 46 | 47 | // This is for the dycore to pull out to determine how to do idealized test cases 48 | coupler.set_option( "standalone_input_file" , inFile ); 49 | 50 | // They dynamical core "dycore" integrates the Euler equations and performans transport of tracers 51 | modules::Dynamics_Euler_Stratified_WenoFV dycore; 52 | custom_modules::Horizontal_Sponge horiz_sponge; 53 | custom_modules::Time_Averager time_averager; 54 | 55 | coupler.add_tracer("water_vapor","water_vapor",true,true); 56 | coupler.get_data_manager_readwrite().get("water_vapor") = 0; 57 | 58 | // Run the initialization modules 59 | dycore .init( coupler ); // Dycore should initialize its own state here 60 | horiz_sponge .init( coupler , 10 , 1. ); 61 | time_averager.init( coupler ); 62 | 63 | real etime = 0; // Elapsed time 64 | 65 | real dtphys = dtphys_in; 66 | while (etime < sim_time) { 67 | // If dtphys <= 0, then set it to the dynamical core's max stable time step 68 | if (dtphys_in <= 0.) { dtphys = dycore.compute_time_step(coupler); } 69 | // If we're about to go past the final time, then limit to time step to exactly hit the final time 70 | if (etime + dtphys > sim_time) { dtphys = sim_time - etime; } 71 | 72 | horiz_sponge.apply ( coupler , dtphys , true , true , false , false ); 73 | dycore.time_step ( coupler , dtphys ); // Move the flow forward according to the Euler equations 74 | modules::sponge_layer ( coupler , dtphys , 1 ); 75 | time_averager.accumulate( coupler , dtphys ); 76 | 77 | etime += dtphys; // Advance elapsed time 78 | } 79 | 80 | time_averager.finalize( coupler ); 81 | 82 | yakl::timer_stop("main"); 83 | } 84 | yakl::finalize(); 85 | MPI_Finalize(); 86 | } 87 | 88 | 89 | -------------------------------------------------------------------------------- /experiments/simple_city/inputs/input_building.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | # Simulation time in seconds 3 | sim_time: 200. # 86400 4 | 5 | # Number of cells to use in the CRMs 6 | nens : 1 7 | nx_glob: 100 8 | ny_glob: 100 9 | nz : 25 10 | 11 | # Domain size of the CRMs 12 | xlen: 2000 13 | ylen: 2000 14 | zlen: 500 15 | 16 | init_data: building 17 | 18 | # Output filename 19 | out_prefix: test 20 | 21 | # GCM time step 22 | dt_gcm: 900 23 | 24 | # CRM physics time step 25 | dt_phys: 0. 26 | 27 | # Output frequency in seconds 28 | out_freq: 10. 29 | 30 | enable_gravity: false 31 | 32 | file_per_process: false 33 | 34 | -------------------------------------------------------------------------------- /experiments/simple_city/inputs/input_city.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | # Simulation time in seconds 3 | sim_time: 200. # 86400 4 | 5 | # Number of cells to use in the CRMs 6 | nens : 2 7 | nx_glob: 400 8 | ny_glob: 400 9 | nz : 60 10 | 11 | # Domain size of the CRMs 12 | xlen: 2000 13 | ylen: 2000 14 | zlen: 300 15 | 16 | init_data: city 17 | 18 | # Output filename 19 | out_prefix: test.nc 20 | 21 | # GCM time step 22 | dt_gcm: 900 23 | 24 | # CRM physics time step 25 | dt_phys: 0. 26 | 27 | # Output frequency in seconds 28 | out_freq: 10. 29 | 30 | enable_gravity: false 31 | 32 | 33 | -------------------------------------------------------------------------------- /experiments/supercell_example/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.0) 2 | project(miniWeatherML) 3 | 4 | enable_language(Fortran) 5 | enable_language(CXX) 6 | enable_language(C) 7 | if ("${YAKL_ARCH}" STREQUAL "CUDA") 8 | enable_language(CUDA) 9 | endif() 10 | 11 | add_subdirectory(${MINIWEATHER_ML_HOME}/model model) 12 | 13 | add_executable(driver driver.cpp) 14 | include(${YAKL_HOME}/yakl_utils.cmake) 15 | yakl_process_target(driver) 16 | target_link_libraries(driver model "${MW_LINK_FLAGS}") 17 | 18 | -------------------------------------------------------------------------------- /experiments/supercell_example/driver.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "coupler.h" 3 | #include "dynamics_euler_stratified_wenofv.h" 4 | #include "microphysics_kessler.h" 5 | #include "sponge_layer.h" 6 | #include "perturb_temperature.h" 7 | #include "column_nudging.h" 8 | 9 | int main(int argc, char** argv) { 10 | MPI_Init( &argc , &argv ); 11 | yakl::init(); 12 | { 13 | using yakl::intrinsics::abs; 14 | using yakl::intrinsics::maxval; 15 | yakl::timer_start("main"); 16 | 17 | // This holds all of the model's variables, dimension sizes, and options 18 | core::Coupler coupler; 19 | 20 | // Read the YAML input file for variables pertinent to running the driver 21 | if (argc <= 1) { endrun("ERROR: Must pass the input YAML filename as a parameter"); } 22 | std::string inFile(argv[1]); 23 | YAML::Node config = YAML::LoadFile(inFile); 24 | if ( !config ) { endrun("ERROR: Invalid YAML input file"); } 25 | auto sim_time = config["sim_time"].as(); 26 | auto nens = config["nens" ].as(); 27 | auto nx_glob = config["nx_glob" ].as(); 28 | auto ny_glob = config["ny_glob" ].as(); 29 | auto nz = config["nz" ].as(); 30 | auto xlen = config["xlen" ].as(); 31 | auto ylen = config["ylen" ].as(); 32 | auto zlen = config["zlen" ].as(); 33 | auto dtphys_in = config["dt_phys" ].as(); 34 | 35 | coupler.set_option( "out_prefix" , config["out_prefix"].as() ); 36 | coupler.set_option( "init_data" , config["init_data"].as() ); 37 | coupler.set_option( "out_freq" , config["out_freq" ].as() ); 38 | 39 | // Coupler state is: (1) dry density; (2) u-velocity; (3) v-velocity; (4) w-velocity; (5) temperature 40 | // (6+) tracer masses (*not* mixing ratios!) 41 | coupler.distribute_mpi_and_allocate_coupled_state(nz, ny_glob, nx_glob, nens); 42 | 43 | // Just tells the coupler how big the domain is in each dimensions 44 | coupler.set_grid( xlen , ylen , zlen ); 45 | 46 | // This is for the dycore to pull out to determine how to do idealized test cases 47 | coupler.set_option( "standalone_input_file" , inFile ); 48 | 49 | // The column nudger nudges the column-average of the model state toward the initial column-averaged state 50 | // This is primarily for the supercell test case to keep the the instability persistently strong 51 | modules::ColumnNudger column_nudger; 52 | // Microphysics performs water phase changess + hydrometeor production, transport, collision, and aggregation 53 | modules::Microphysics_Kessler micro; 54 | // They dynamical core "dycore" integrates the Euler equations and performans transport of tracers 55 | modules::Dynamics_Euler_Stratified_WenoFV dycore; 56 | 57 | // Run the initialization modules 58 | micro .init ( coupler ); // Allocate micro state and register its tracers in the coupler 59 | dycore.init ( coupler ); // Dycore should initialize its own state here 60 | column_nudger.set_column ( coupler ); // Set the column before perturbing 61 | modules::perturb_temperature( coupler ); // Randomly perturb bottom layers of temperature to initiate convection 62 | 63 | real etime = 0; // Elapsed time 64 | 65 | real dtphys = dtphys_in; 66 | while (etime < sim_time) { 67 | // If dtphys <= 0, then set it to the dynamical core's max stable time step 68 | if (dtphys_in <= 0.) { dtphys = dycore.compute_time_step(coupler); } 69 | // If we're about to go past the final time, then limit to time step to exactly hit the final time 70 | if (etime + dtphys > sim_time) { dtphys = sim_time - etime; } 71 | 72 | // Run the runtime modules 73 | dycore.time_step ( coupler , dtphys ); // Move the flow forward according to the Euler equations 74 | micro .time_step ( coupler , dtphys ); // Perform phase changes for water + precipitation / falling 75 | modules::sponge_layer ( coupler , dtphys ); // Damp spurious waves to the horiz. mean at model top 76 | column_nudger.nudge_to_column( coupler , dtphys ); // Nudge slightly back toward unstable profile 77 | // so that supercell persists for all time 78 | 79 | etime += dtphys; // Advance elapsed time 80 | } 81 | 82 | // TODO: Add finalize( coupler ) modules here 83 | 84 | yakl::timer_stop("main"); 85 | } 86 | yakl::finalize(); 87 | MPI_Finalize(); 88 | } 89 | 90 | 91 | -------------------------------------------------------------------------------- /experiments/supercell_example/inputs/input_euler3d.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | # Simulation time in seconds 3 | sim_time: 7200. # 86400 4 | 5 | # Number of cells to use in the CRMs 6 | nens : 1 7 | nx_glob: 100 8 | ny_glob: 1 9 | nz : 40 10 | 11 | # Domain size of the CRMs 12 | xlen: 100000 13 | ylen: 100000 14 | zlen: 20000 15 | 16 | init_data: supercell 17 | 18 | # Output filename 19 | out_prefix: test 20 | 21 | # GCM time step 22 | dt_gcm: 900 23 | 24 | # CRM physics time step 25 | dt_phys: 0. 26 | 27 | # Output frequency in seconds 28 | out_freq: 100. 29 | 30 | 31 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.0) 2 | project(miniWeatherML) 3 | 4 | enable_language(Fortran) 5 | enable_language(CXX) 6 | enable_language(C) 7 | if ("${YAKL_ARCH}" STREQUAL "CUDA") 8 | enable_language(CUDA) 9 | endif() 10 | 11 | add_subdirectory(${MINIWEATHER_ML_HOME}/model model) 12 | 13 | add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/custom_modules custom_modules) 14 | 15 | include(${YAKL_HOME}/yakl_utils.cmake) 16 | 17 | add_executable(generate_micro_data generate_micro_data.cpp) 18 | yakl_process_target(generate_micro_data) 19 | target_link_libraries(generate_micro_data model custom_modules "${MW_LINK_FLAGS}") 20 | 21 | add_executable(gather_statistics gather_statistics.cpp) 22 | yakl_process_target(gather_statistics) 23 | target_link_libraries(gather_statistics model custom_modules "${MW_LINK_FLAGS}") 24 | 25 | add_executable(inference_ponni inference_ponni.cpp) 26 | yakl_process_target(inference_ponni) 27 | target_link_libraries(inference_ponni model custom_modules "${MW_LINK_FLAGS}") 28 | 29 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/custom_modules/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | 2 | add_library(custom_modules INTERFACE) 3 | target_link_libraries(custom_modules INTERFACE core) 4 | target_include_directories(custom_modules INTERFACE ${CMAKE_CURRENT_SOURCE_DIR}) 5 | 6 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/custom_modules/gather_micro_statistics.h: -------------------------------------------------------------------------------- 1 | 2 | #pragma once 3 | 4 | #include "coupler.h" 5 | #include "YAKL_netcdf.h" 6 | #include 7 | 8 | namespace custom_modules { 9 | 10 | class StatisticsGatherer { 11 | public: 12 | double numer; 13 | double denom; 14 | int num_out; 15 | 16 | StatisticsGatherer() { numer = 0; denom = 0; num_out = 0; } 17 | 18 | 19 | void gather_micro_statistics( core::Coupler &input , core::Coupler &output , real dt , real etime ) { 20 | using yakl::c::parallel_for; 21 | using yakl::c::Bounds; 22 | using std::min; 23 | using std::abs; 24 | 25 | auto &dm_in = input .get_data_manager_readonly(); 26 | auto &dm_out = output.get_data_manager_readonly(); 27 | 28 | auto temp_in = dm_in .get("temp" ); 29 | auto temp_out = dm_out.get("temp" ); 30 | auto rho_v_in = dm_in .get("water_vapor" ); 31 | auto rho_v_out = dm_out.get("water_vapor" ); 32 | auto rho_c_in = dm_in .get("cloud_liquid" ); 33 | auto rho_c_out = dm_out.get("cloud_liquid" ); 34 | auto rho_p_in = dm_in .get("precip_liquid"); 35 | auto rho_p_out = dm_out.get("precip_liquid"); 36 | 37 | int nx = input.get_nx(); 38 | int ny = input.get_ny(); 39 | int nz = input.get_nz(); 40 | 41 | int3d active("active",nz,ny,nx); 42 | 43 | parallel_for( Bounds<3>(nz,ny,nx) , YAKL_LAMBDA (int k, int j, int i) { 44 | if ( is_active( temp_in (k,j,i,0) , temp_out (k,j,i,0) , 45 | rho_v_in(k,j,i,0) , rho_v_out(k,j,i,0) , 46 | rho_c_in(k,j,i,0) , rho_c_out(k,j,i,0) , 47 | rho_p_in(k,j,i,0) , rho_p_out(k,j,i,0) ) ) { 48 | active(k,j,i) = 1; 49 | } else { 50 | active(k,j,i) = 0; 51 | } 52 | }); 53 | 54 | if (etime > (num_out+1)*200) { print(input.is_mainproc(), input.get_mpi_data_type()); num_out++; } 55 | 56 | numer += yakl::intrinsics::sum( active ); 57 | denom += active.size(); 58 | } 59 | 60 | 61 | YAKL_INLINE static bool is_active( real temp_in , real temp_out , 62 | real rho_v_in , real rho_v_out , 63 | real rho_c_in , real rho_c_out , 64 | real rho_p_in , real rho_p_out ) { 65 | real tol = 1.e-10; 66 | real temp_diff = std::abs( temp_out - temp_in ); 67 | real rho_v_diff = std::abs( rho_v_out - rho_v_in ); 68 | real rho_c_diff = std::abs( rho_c_out - rho_c_in ); 69 | real rho_p_diff = std::abs( rho_p_out - rho_p_in ); 70 | 71 | if (temp_diff > tol || rho_v_diff > tol || rho_c_diff > tol || rho_p_diff > tol) { return true; } 72 | return false; 73 | } 74 | 75 | 76 | void print(bool mainproc, MPI_Datatype mpi_data_type) { 77 | double numer_all; 78 | double denom_all; 79 | MPI_Reduce( &numer , &numer_all , 1 , mpi_data_type , MPI_SUM , 0 , MPI_COMM_WORLD ); 80 | MPI_Reduce( &denom , &denom_all , 1 , mpi_data_type , MPI_SUM , 0 , MPI_COMM_WORLD ); 81 | if (mainproc) { 82 | std::cout << "*** Ratio Active ***: " << std::scientific << std::setw(10) << numer_all/denom_all << std::endl; 83 | } 84 | } 85 | 86 | 87 | void finalize( core::Coupler &coupler ) { print(coupler.is_mainproc(),coupler.get_mpi_data_type()); } 88 | 89 | }; 90 | 91 | 92 | } 93 | 94 | 95 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/custom_modules/generate_micro_surrogate_data.h: -------------------------------------------------------------------------------- 1 | 2 | #pragma once 3 | 4 | #include "coupler.h" 5 | #include "YAKL_netcdf.h" 6 | #include 7 | #include "gather_micro_statistics.h" 8 | #include 9 | 10 | namespace custom_modules { 11 | 12 | class DataGenerator { 13 | public: 14 | 15 | std::string fname; 16 | 17 | void init( core::Coupler &coupler ) { 18 | yakl::SimpleNetCDF nc; 19 | fname = std::string("supercell_kessler_data_task_") + std::to_string(coupler.get_myrank()) + 20 | std::string(".nc"); 21 | nc.create(fname); 22 | nc.createDim("nsamples"); 23 | nc.close(); 24 | 25 | std::ofstream myfile; 26 | myfile.open ("supercell_kessler_metadata.txt"); 27 | myfile << "This dataset contains data for training a surrogate model to emulate Kessler microphysics.\n\n"; 28 | myfile << "vars_in : temperature, dry air density, water vapor density, cloud liquid density, precipitation density\n"; 29 | myfile << "vars_out: temperature, water vapor density, cloud liquid density, precipitation density\n"; 30 | myfile.close(); 31 | 32 | } 33 | 34 | 35 | void generate_samples_stencil( core::Coupler &input , core::Coupler &output , real dt , real etime ) { 36 | using std::min; 37 | using std::max; 38 | using yakl::c::parallel_for; 39 | using yakl::c::Bounds; 40 | 41 | int nx = input.get_nx (); 42 | int ny = input.get_ny (); 43 | int nz = input.get_nz (); 44 | int nranks = input.get_nranks(); 45 | int myrank = input.get_myrank(); 46 | 47 | // This was gathered from the gather_statistics.cpp driver 48 | // On average, about 40% of the cells experience active microphysics at any given time 49 | double ratio_active = 0.4; 50 | double expected_num_active = ratio_active * nx*ny*nz; 51 | double expected_num_inactive = (1-ratio_active) * nx*ny*nz; 52 | 53 | double desired_samples_per_time_step = 50; 54 | 55 | double desired_ratio_active = 0.5; 56 | 57 | 58 | double desired_samples_active = desired_ratio_active * desired_samples_per_time_step / nranks; 59 | double desired_samples_inactive = (1-desired_ratio_active) * desired_samples_per_time_step / nranks; 60 | 61 | double active_threshold = desired_samples_active / expected_num_active; 62 | double inactive_threshold = desired_samples_inactive / expected_num_inactive; 63 | 64 | auto &dm_in = input .get_data_manager_readonly(); 65 | auto &dm_out = output.get_data_manager_readonly(); 66 | 67 | auto rho_d = dm_in .get("density_dry" ); 68 | 69 | auto temp_in = dm_in .get("temp" ); 70 | auto temp_out = dm_out.get("temp" ); 71 | auto rho_v_in = dm_in .get("water_vapor" ); 72 | auto rho_v_out = dm_out.get("water_vapor" ); 73 | auto rho_c_in = dm_in .get("cloud_liquid" ); 74 | auto rho_c_out = dm_out.get("cloud_liquid" ); 75 | auto rho_p_in = dm_in .get("precip_liquid"); 76 | auto rho_p_out = dm_out.get("precip_liquid"); 77 | 78 | bool3d do_sample("do_sample",nz,ny,nx); 79 | 80 | // Seed with time, but make sure the seed magnitude will not lead to an overflow of size_t's max value 81 | size_t max_mag = std::numeric_limits::max() / ( (size_t)nranks + (size_t)nz * (size_t)ny * (size_t)nx ); 82 | size_t seed = ( (size_t) time(nullptr) ) % max_mag; 83 | 84 | parallel_for( Bounds<3>(nz,ny,nx) , YAKL_LAMBDA (int k, int j, int i) { 85 | double thresh; 86 | if ( StatisticsGatherer::is_active( temp_in (k,j,i,0) , temp_out (k,j,i,0) , 87 | rho_v_in(k,j,i,0) , rho_v_out(k,j,i,0) , 88 | rho_c_in(k,j,i,0) , rho_c_out(k,j,i,0) , 89 | rho_p_in(k,j,i,0) , rho_p_out(k,j,i,0) ) ) { 90 | thresh = active_threshold; 91 | } else { 92 | thresh = inactive_threshold; 93 | } 94 | 95 | double rand_num = yakl::Random((seed+myrank)*nz*ny*nx + k*ny*nx + j*nx + i).genFP(); 96 | if (rand_num < thresh) { do_sample(k,j,i) = true; } 97 | else { do_sample(k,j,i) = false; } 98 | }); 99 | 100 | auto host_rho_d = rho_d .createHostCopy(); 101 | auto host_temp_in = temp_in .createHostCopy(); 102 | auto host_temp_out = temp_out .createHostCopy(); 103 | auto host_rho_v_in = rho_v_in .createHostCopy(); 104 | auto host_rho_v_out = rho_v_out.createHostCopy(); 105 | auto host_rho_c_in = rho_c_in .createHostCopy(); 106 | auto host_rho_c_out = rho_c_out.createHostCopy(); 107 | auto host_rho_p_in = rho_p_in .createHostCopy(); 108 | auto host_rho_p_out = rho_p_out.createHostCopy(); 109 | auto host_do_sample = do_sample.createHostCopy(); 110 | 111 | yakl::SimpleNetCDF nc; 112 | nc.open(fname,yakl::NETCDF_MODE_WRITE); 113 | int ulindex = nc.getDimSize("nsamples"); 114 | 115 | if ( ! nc.varExists("time_step_size" ) ) nc.write(dt ,"time_step_size" ); 116 | if ( ! nc.varExists("only_two_dimensions") ) nc.write(input.get_ny_glob()==1 ? 0 : 1,"only_two_dimensions"); 117 | if ( ! nc.varExists("dx" ) ) nc.write(input.get_dx () ,"dx" ); 118 | if ( ! nc.varExists("dy" ) ) nc.write(input.get_dy () ,"dy" ); 119 | if ( ! nc.varExists("dz" ) ) nc.write(input.get_dz () ,"dz" ); 120 | if ( ! nc.varExists("xlen" ) ) nc.write(input.get_xlen() ,"xlen" ); 121 | if ( ! nc.varExists("ylen" ) ) nc.write(input.get_ylen() ,"ylen" ); 122 | if ( ! nc.varExists("zlen" ) ) nc.write(input.get_zlen() ,"zlen" ); 123 | 124 | // Save in 32-bit to reduce file / memory size when training 125 | floatHost2d gen_input ("gen_input" ,5,2); 126 | floatHost1d gen_output("gen_output",4); 127 | 128 | for (int k=0; k < nz; k++) { 129 | for (int j=0; j < ny; j++) { 130 | for (int i=0; i < nx; i++) { 131 | if (host_do_sample(k,j,i)) { 132 | gen_input (0,0) = host_temp_in (k,j,i,0); 133 | gen_input (1,0) = host_rho_d (k,j,i,0); 134 | gen_input (2,0) = host_rho_v_in(k,j,i,0); 135 | gen_input (3,0) = host_rho_c_in(k,j,i,0); 136 | gen_input (4,0) = host_rho_p_in(k,j,i,0); 137 | gen_input (0,1) = host_temp_in (std::min(nz-1,k+1),j,i,0); 138 | gen_input (1,1) = host_rho_v_in(std::min(nz-1,k+1),j,i,0); 139 | gen_input (2,1) = host_rho_c_in(std::min(nz-1,k+1),j,i,0); 140 | gen_input (3,1) = host_rho_p_in(std::min(nz-1,k+1),j,i,0); 141 | gen_output(0) = host_temp_out (k,j,i,0); 142 | gen_output(1) = host_rho_v_out(k,j,i,0); 143 | gen_output(2) = host_rho_c_out(k,j,i,0); 144 | gen_output(3) = host_rho_p_out(k,j,i,0); 145 | nc.write1( gen_input , "inputs" , {"num_vars_in" ,"sten_size"} , ulindex , "nsamples" ); 146 | nc.write1( gen_output , "outputs" , {"num_vars_out" } , ulindex , "nsamples" ); 147 | ulindex++; 148 | } 149 | } 150 | } 151 | } 152 | nc.close(); 153 | } 154 | 155 | }; 156 | 157 | } 158 | 159 | 160 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/custom_modules/microphysics_kessler_ponni.h: -------------------------------------------------------------------------------- 1 | 2 | #pragma once 3 | 4 | #include "main_header.h" 5 | #include "coupler.h" 6 | #include "ponni.h" 7 | #include "ponni_load_h5_weights.h" 8 | 9 | namespace custom_modules { 10 | 11 | using ponni::Matvec; 12 | using ponni::Bias; 13 | using ponni::Relu; 14 | using ponni::Inference; 15 | 16 | class Microphysics_Kessler { 17 | public: 18 | int static constexpr num_tracers = 3; 19 | 20 | real R_d ; 21 | real cp_d ; 22 | real cv_d ; 23 | real gamma_d; 24 | real kappa_d; 25 | real R_v ; 26 | real cp_v ; 27 | real cv_v ; 28 | real p0 ; 29 | real grav ; 30 | 31 | int static constexpr ID_V = 0; // Local index for water vapor 32 | int static constexpr ID_C = 1; // Local index for cloud liquid 33 | int static constexpr ID_R = 2; // Local index for precipitated liquid (rain) 34 | 35 | int static constexpr MAX_LAYERS=10; 36 | 37 | real2d scl_out; 38 | real2d scl_in ; 39 | 40 | typedef decltype(ponni::create_inference_model(Matvec(), 41 | Bias (), 42 | Relu (), 43 | Matvec(), 44 | Bias ())) MODEL; 45 | MODEL model; 46 | 47 | 48 | Microphysics_Kessler() { 49 | R_d = 287.; 50 | cp_d = 1003.; 51 | cv_d = cp_d - R_d; 52 | gamma_d = cp_d / cv_d; 53 | kappa_d = R_d / cp_d; 54 | R_v = 461.; 55 | cp_v = 1859; 56 | cv_v = R_v - cp_v; 57 | p0 = 1.e5; 58 | grav = 9.81; 59 | } 60 | 61 | 62 | 63 | YAKL_INLINE static int get_num_tracers() { 64 | return num_tracers; 65 | } 66 | 67 | 68 | 69 | void init(core::Coupler &coupler) { 70 | using yakl::c::parallel_for; 71 | using yakl::c::Bounds; 72 | 73 | int nens = coupler.get_nens(); 74 | int nx = coupler.get_nx(); 75 | int ny = coupler.get_ny(); 76 | int nz = coupler.get_nz(); 77 | 78 | // Register tracers in the coupler 79 | // name description positive adds mass 80 | coupler.add_tracer("water_vapor" , "Water Vapor" , true , true); 81 | coupler.add_tracer("cloud_liquid" , "Cloud liquid" , true , true); 82 | coupler.add_tracer("precip_liquid" , "precip_liquid" , true , true); 83 | 84 | auto &dm = coupler.get_data_manager_readwrite(); 85 | 86 | // Register and allocation non-tracer quantities used by the microphysics 87 | dm.register_and_allocate( "precl" , "precipitation rate" , {ny,nx,nens} , {"y","x","nens"} ); 88 | 89 | // Initialize all micro data to zero 90 | dm.get("water_vapor" ) = 0; 91 | dm.get("cloud_liquid" ) = 0; 92 | dm.get("precip_liquid") = 0; 93 | dm.get("precl" ) = 0; 94 | 95 | coupler.set_option("micro","kessler"); 96 | 97 | auto inFile = coupler.get_option("standalone_input_file"); 98 | YAML::Node config = YAML::LoadFile(inFile); 99 | auto keras_weights_h5 = config["keras_weights_h5" ].as(); 100 | auto nn_input_scaling = config["nn_input_scaling" ].as(); 101 | auto nn_output_scaling = config["nn_output_scaling"].as(); 102 | 103 | ponni::Matvec matvec_1( ponni::load_h5_weights<2>( keras_weights_h5 , "/dense_6/dense_6" , "kernel:0" ) ); 104 | ponni::Bias bias_1 ( ponni::load_h5_weights<1>( keras_weights_h5 , "/dense_6/dense_6" , "bias:0" ) ); 105 | ponni::Relu relu_1 ( bias_1.get_num_outputs() , 0.1 ); // LeakyReLU with negative slope of 0.1 106 | ponni::Matvec matvec_2( ponni::load_h5_weights<2>( keras_weights_h5 , "/dense_7/dense_7" , "kernel:0" ) ); 107 | ponni::Bias bias_2 ( ponni::load_h5_weights<1>( keras_weights_h5 , "/dense_7/dense_7" , "bias:0" ) ); 108 | 109 | this->model = ponni::create_inference_model(matvec_1, bias_1, relu_1, matvec_2, bias_2); 110 | model.validate(); 111 | model.print(); 112 | 113 | // Load the data scaling arrays 114 | scl_out = real2d("scl_out",4,2); 115 | scl_in = real2d("scl_in" ,5,2); 116 | 117 | auto scl_out_host = scl_out.createHostObject(); 118 | auto scl_in_host = scl_in .createHostObject(); 119 | 120 | std::ifstream file1; 121 | // input scaler 122 | file1.open( nn_input_scaling ); 123 | for (int j = 0; j < 5; j++) { 124 | for (int i = 0; i < 2; i++) { 125 | file1 >> scl_in_host(j,i); 126 | } 127 | } 128 | file1.close(); 129 | file1.open( nn_output_scaling ); 130 | for (int j = 0; j < 4; j++) { 131 | for (int i = 0; i < 2; i++) { 132 | file1 >> scl_out_host(j,i); 133 | } 134 | } 135 | file1.close(); 136 | 137 | scl_in_host .deep_copy_to(scl_in ); 138 | scl_out_host.deep_copy_to(scl_out); 139 | yakl::fence(); 140 | 141 | std::cout << scl_in ; 142 | std::cout << "\n"; 143 | std::cout << scl_out; 144 | std::cout << "\n"; 145 | } 146 | 147 | 148 | 149 | void time_step( core::Coupler &coupler , real dt ) { 150 | using yakl::c::parallel_for; 151 | using yakl::c::Bounds; 152 | 153 | auto &dm = coupler.get_data_manager_readwrite(); 154 | 155 | // Grab the data 156 | auto rho_v = dm.get_lev_col("water_vapor" ); 157 | auto rho_c = dm.get_lev_col("cloud_liquid" ); 158 | auto rho_r = dm.get_lev_col("precip_liquid"); 159 | auto rho_d = dm.get_lev_col("density_dry" ); 160 | auto temp = dm.get_lev_col("temp" ); 161 | 162 | // Grab the dimension sizes 163 | real dz = coupler.get_dz(); 164 | int nz = coupler.get_nz(); 165 | int ny = coupler.get_ny(); 166 | int nx = coupler.get_nx(); 167 | int nens = coupler.get_nens(); 168 | int ncol = ny*nx*nens; 169 | 170 | YAKL_SCOPE( scl_out , this->scl_out ); 171 | YAKL_SCOPE( scl_in , this->scl_in ); 172 | 173 | ///////////////////////////////////////////////////////////////////////// 174 | // NEURAL NETWORK PONNI 175 | ///////////////////////////////////////////////////////////////////////// 176 | // Build inputs 177 | int constexpr num_in = 5; 178 | float2d ponni_in("ponni_in",num_in,nz*ncol); 179 | 180 | parallel_for( Bounds<2>(nz,ncol) , YAKL_LAMBDA (int k, int i) { 181 | int iglob = k*ncol+i; 182 | ponni_in( 0 , iglob ) = ( temp (k,i) - scl_in(0,0) ) / ( scl_in(0,1) - scl_in(0,0) ); 183 | ponni_in( 1 , iglob ) = ( rho_d(k,i) - scl_in(1,0) ) / ( scl_in(1,1) - scl_in(1,0) ); 184 | ponni_in( 2 , iglob ) = ( rho_v(k,i) - scl_in(2,0) ) / ( scl_in(2,1) - scl_in(2,0) ); 185 | ponni_in( 3 , iglob ) = ( rho_c(k,i) - scl_in(3,0) ) / ( scl_in(3,1) - scl_in(3,0) ); 186 | ponni_in( 4 , iglob ) = ( rho_r(k,i) - scl_in(4,0) ) / ( scl_in(4,1) - scl_in(4,0) ); 187 | }); 188 | 189 | auto ponni_out = model.forward_batch_parallel( ponni_in ); 190 | 191 | real2d temp_tmp = temp .createDeviceCopy(); 192 | real2d rho_v_tmp = rho_v.createDeviceCopy(); 193 | real2d rho_c_tmp = rho_c.createDeviceCopy(); 194 | real2d rho_r_tmp = rho_r.createDeviceCopy(); 195 | 196 | parallel_for( Bounds<2>(nz,ncol) , YAKL_LAMBDA (int k, int i) { 197 | int iglob = k*ncol+i; 198 | temp_tmp (k,i) = ponni_out( 0 , iglob ) * (scl_out(0,1) - scl_out(0,0)) + scl_out(0,0) ; 199 | rho_v_tmp(k,i) = std::max( 0._fp , ponni_out( 1 , iglob ) * (scl_out(1,1) - scl_out(1,0)) + scl_out(1,0) ); 200 | rho_c_tmp(k,i) = std::max( 0._fp , ponni_out( 2 , iglob ) * (scl_out(2,1) - scl_out(2,0)) + scl_out(2,0) ); 201 | rho_r_tmp(k,i) = std::max( 0._fp , ponni_out( 3 , iglob ) * (scl_out(3,1) - scl_out(3,0)) + scl_out(3,0) ); 202 | }); 203 | 204 | // std::cout << yakl::intrinsics::maxval( temp_tmp ) << "\n:"; 205 | // std::cout << yakl::intrinsics::maxval( rho_v_tmp ) << "\n:"; 206 | // std::cout << yakl::intrinsics::maxval( rho_c_tmp ) << "\n:"; 207 | // std::cout << yakl::intrinsics::maxval( rho_r_tmp ) << "\n:"; 208 | 209 | ///////////////////////////////////////////////////////////////////////// 210 | // NORMAL 211 | ///////////////////////////////////////////////////////////////////////// 212 | 213 | // These are inputs to kessler(...) 214 | real2d qv ("qv" ,nz,ncol); 215 | real2d qc ("qc" ,nz,ncol); 216 | real2d qr ("qr" ,nz,ncol); 217 | real2d pressure("pressure",nz,ncol); 218 | real2d theta ("theta" ,nz,ncol); 219 | real2d exner ("exner" ,nz,ncol); 220 | real2d zmid ("zmid" ,nz,ncol); 221 | 222 | // Force constants into local scope 223 | real R_d = this->R_d; 224 | real R_v = this->R_v; 225 | real cp_d = this->cp_d; 226 | real p0 = this->p0; 227 | 228 | // Save initial state, and compute inputs for kessler(...) 229 | parallel_for( "kessler timeStep 2" , Bounds<2>(nz,ncol) , YAKL_LAMBDA (int k, int i) { 230 | zmid (k,i) = (k+0.5_fp) * dz; 231 | qv (k,i) = rho_v(k,i) / rho_d(k,i); 232 | qc (k,i) = rho_c(k,i) / rho_d(k,i); 233 | qr (k,i) = rho_r(k,i) / rho_d(k,i); 234 | pressure(k,i) = R_d * rho_d(k,i) * temp(k,i) + R_v * rho_v(k,i) * temp(k,i); 235 | exner (k,i) = pow( pressure(k,i) / p0 , R_d / cp_d ); 236 | theta (k,i) = temp(k,i) / exner(k,i); 237 | }); 238 | 239 | auto precl = dm.get_collapsed("precl"); 240 | 241 | //////////////////////////////////////////// 242 | // Call Kessler code 243 | //////////////////////////////////////////// 244 | kessler(theta, qv, qc, qr, rho_d, precl, zmid, exner, dt, R_d, cp_d, p0); 245 | 246 | auto rho_v_diff = rho_v.createDeviceCopy(); 247 | auto rho_c_diff = rho_c.createDeviceCopy(); 248 | auto rho_r_diff = rho_r.createDeviceCopy(); 249 | auto temp_diff = temp .createDeviceCopy(); 250 | 251 | // Post-process microphysics changes back to the coupler state 252 | parallel_for( "kessler timeStep 3" , Bounds<2>(nz,ncol) , YAKL_LAMBDA (int k, int i) { 253 | rho_v (k,i) = qv(k,i)*rho_d(k,i); 254 | rho_c (k,i) = qc(k,i)*rho_d(k,i); 255 | rho_r (k,i) = qr(k,i)*rho_d(k,i); 256 | // While micro changes total pressure, thus changing exner, the definition 257 | // of theta depends on the old exner pressure, so we'll use old exner here 258 | temp (k,i) = theta(k,i) * exner(k,i); 259 | 260 | rho_v_diff(k,i) = rho_v_tmp(k,i) - rho_v(k,i); 261 | rho_c_diff(k,i) = rho_c_tmp(k,i) - rho_c(k,i); 262 | rho_r_diff(k,i) = rho_r_tmp(k,i) - rho_r(k,i); 263 | temp_diff (k,i) = temp_tmp (k,i) - temp (k,i); 264 | }); 265 | 266 | std::cout << "Relative diff rho_v: " << yakl::intrinsics::sum( rho_v_diff ) / rho_v_diff.size() << "\n"; 267 | std::cout << "Relative diff rho_c: " << yakl::intrinsics::sum( rho_c_diff ) / rho_c_diff.size() << "\n"; 268 | std::cout << "Relative diff rho_r: " << yakl::intrinsics::sum( rho_r_diff ) / rho_r_diff.size() << "\n"; 269 | std::cout << "Relative diff temp : " << yakl::intrinsics::sum( temp_diff ) / temp_diff .size() << "\n"; 270 | 271 | // Uncomment these to overwrite the original Kessler with the ML inference 272 | // I.e., uncommenting these lines runs the ML inferencing online 273 | // temp_tmp .deep_copy_to( temp ); 274 | // rho_v_tmp.deep_copy_to( rho_v ); 275 | // rho_c_tmp.deep_copy_to( rho_c ); 276 | // rho_r_tmp.deep_copy_to( rho_r ); 277 | yakl::fence(); 278 | } 279 | 280 | 281 | 282 | /////////////////////////////////////////////////////////////////////////////// 283 | // 284 | // Version: 2.0 285 | // 286 | // Date: January 22nd, 2015 287 | // 288 | // Change log: 289 | // v2 - Added sub-cycling of rain sedimentation so as not to violate 290 | // CFL condition. 291 | // 292 | // The KESSLER subroutine implements the Kessler (1969) microphysics 293 | // parameterization as described by Soong and Ogura (1973) and Klemp 294 | // and Wilhelmson (1978, KW). KESSLER is called at the end of each 295 | // time step and makes the final adjustments to the potential 296 | // temperature and moisture variables due to microphysical processes 297 | // occurring during that time step. KESSLER is called once for each 298 | // vertical column of grid cells. Increments are computed and added 299 | // into the respective variables. The Kessler scheme contains three 300 | // moisture categories: water vapor, cloud water (liquid water that 301 | // moves with the flow), and rain water (liquid water that falls 302 | // relative to the surrounding air). There are no ice categories. 303 | // 304 | // Variables in the column are ordered from the surface to the top. 305 | // 306 | // Parameters: 307 | // theta (inout) - dry potential temperature (K) 308 | // qv (inout) - water vapor mixing ratio (gm/gm) (dry mixing ratio) 309 | // qc (inout) - cloud water mixing ratio (gm/gm) (dry mixing ratio) 310 | // qr (inout) - rain water mixing ratio (gm/gm) (dry mixing ratio) 311 | // rho (in ) - dry air density (not mean state as in KW) (kg/m^3) 312 | // pk (in ) - Exner function (not mean state as in KW) (p/p0)**(R/cp) 313 | // dt (in ) - time step (s) 314 | // z (in ) - heights of thermodynamic levels in the grid column (m) 315 | // precl ( out) - Precipitation rate (m_water/s) 316 | // Rd (in ) - Dry air ideal gas constant 317 | // cp (in ) - Specific heat of dry air at constant pressure 318 | // p0 (in ) - Reference pressure (Pa) 319 | // 320 | // Output variables: 321 | // Increments are added into t, qv, qc, qr, and precl which are 322 | // returned to the routine from which KESSLER was called. To obtain 323 | // the total precip qt, after calling the KESSLER routine, compute: 324 | // 325 | // qt = sum over surface grid cells of (precl * cell area) (kg) 326 | // [here, the conversion to kg uses (10^3 kg/m^3)*(10^-3 m/mm) = 1] 327 | // 328 | // 329 | // Written in Fortran by: Paul Ullrich 330 | // University of California, Davis 331 | // Email: paullrich@ucdavis.edu 332 | // 333 | // Ported to C++ / YAKL by: Matt Norman 334 | // Oak Ridge National Laboratory 335 | // normanmr@ornl.gov 336 | // https://mrnorman.github.io 337 | // 338 | // Based on a code by Joseph Klemp 339 | // (National Center for Atmospheric Research) 340 | // 341 | // Reference: 342 | // 343 | // Klemp, J. B., W. C. Skamarock, W. C., and S.-H. Park, 2015: 344 | // Idealized Global Nonhydrostatic Atmospheric Test Cases on a Reduced 345 | // Radius Sphere. Journal of Advances in Modeling Earth Systems. 346 | // doi:10.1002/2015MS000435 347 | // 348 | /////////////////////////////////////////////////////////////////////////////// 349 | 350 | void kessler(real2d const &theta, real2d const &qv, real2d const &qc, real2d const &qr, realConst2d rho, 351 | real1d const &precl, realConst2d z, realConst2d pk, real dt, real Rd, real cp, real p0) const { 352 | using yakl::c::parallel_for; 353 | using yakl::c::Bounds; 354 | 355 | int nz = theta.dimension[0]; 356 | int ncol = theta.dimension[1]; 357 | 358 | // Maximum time step size in accordance with CFL condition 359 | if (dt <= 0) { endrun("kessler.f90 called with nonpositive dt"); } 360 | 361 | real psl = p0 / 100; // pressure at sea level (mb) 362 | real rhoqr = 1000._fp; // density of liquid water (kg/m^3) 363 | real lv = 2.5e6_fp; // latent heat of vaporization (J/kg) 364 | 365 | real2d r ("r" ,nz ,ncol); 366 | real2d rhalf("rhalf",nz ,ncol); 367 | real2d pc ("pc" ,nz ,ncol); 368 | real2d velqr("velqr",nz ,ncol); 369 | real2d dt2d ("dt2d" ,nz-1,ncol); 370 | 371 | parallel_for( "kessler main 1" , Bounds<2>(nz,ncol) , YAKL_LAMBDA (int k, int i) { 372 | r (k,i) = 0.001_fp * rho(k,i); 373 | rhalf(k,i) = sqrt( rho(0,i) / rho(k,i) ); 374 | pc (k,i) = 3.8_fp / ( pow( pk(k,i) , cp/Rd ) * psl ); 375 | // Liquid water terminal velocity (m/s) following KW eq. 2.15 376 | velqr(k,i) = 36.34_fp * pow( qr(k,i)*r(k,i) , 0.1364_fp ) * rhalf(k,i); 377 | // Compute maximum stable time step for each cell 378 | if (k < nz-1) { 379 | if (velqr(k,i) > 1.e-10_fp) { 380 | dt2d(k,i) = 0.8_fp * (z(k+1,i)-z(k,i))/velqr(k,i); 381 | } else { 382 | dt2d(k,i) = dt; 383 | } 384 | } 385 | // Initialize precip rate to zero 386 | if (k == 0) { 387 | precl(i) = 0; 388 | } 389 | }); 390 | 391 | // Reduce down the minimum time step among the cells 392 | real dt_max = yakl::intrinsics::minval(dt2d); 393 | 394 | // Number of subcycles 395 | int rainsplit = ceil(dt / dt_max); 396 | real dt0 = dt / static_cast(rainsplit); 397 | 398 | real2d sed("sed",nz,ncol); 399 | 400 | // Subcycle through rain process 401 | for (int nt=0; nt < rainsplit; nt++) { 402 | 403 | // Sedimentation term using upstream differencing 404 | parallel_for( "kessler main 2" , Bounds<2>(nz,ncol) , YAKL_LAMBDA (int k, int i) { 405 | if (k == 0) { 406 | // Precipitation rate (m/s) 407 | precl(i) = precl(i) + rho(0,i) * qr(0,i) * velqr(0,i) / rhoqr; 408 | } 409 | if (k == nz-1) { 410 | sed(nz-1,i) = -dt0*qr(nz-1,i)*velqr(nz-1,i)/(0.5_fp * (z(nz-1,i)-z(nz-2,i))); 411 | } else { 412 | sed(k,i) = dt0 * ( r(k+1,i)*qr(k+1,i)*velqr(k+1,i) - 413 | r(k ,i)*qr(k ,i)*velqr(k ,i) ) / ( r(k,i)*(z(k+1,i)-z(k,i)) ); 414 | } 415 | }); 416 | 417 | // Adjustment terms 418 | parallel_for( "kessler main 3" , Bounds<2>(nz,ncol) , YAKL_LAMBDA (int k, int i) { 419 | // Autoconversion and accretion rates following KW eq. 2.13a,b 420 | real qrprod = qc(k,i) - ( qc(k,i)-dt0*std::max( 0.001_fp * (qc(k,i)-0.001_fp) , 0._fp ) ) / 421 | ( 1 + dt0 * 2.2_fp * pow( qr(k,i) , 0.875_fp ) ); 422 | qc(k,i) = std::max( qc(k,i)-qrprod , 0._fp ); 423 | qr(k,i) = std::max( qr(k,i)+qrprod+sed(k,i) , 0._fp ); 424 | 425 | // Saturation vapor mixing ratio (gm/gm) following KW eq. 2.11 426 | real tmp = pk(k,i)*theta(k,i)-36._fp; 427 | real qvs = pc(k,i)*exp( 17.27_fp * (pk(k,i)*theta(k,i)-273._fp) / tmp ); 428 | real prod = (qv(k,i)-qvs) / (1._fp + qvs*(4093._fp * lv/cp)/(tmp*tmp)); 429 | 430 | // Evaporation rate following KW eq. 2.14a,b 431 | real tmp1 = dt0*( ( ( 1.6_fp + 124.9_fp * pow( r(k,i)*qr(k,i) , 0.2046_fp ) ) * 432 | pow( r(k,i)*qr(k,i) , 0.525_fp ) ) / 433 | ( 2550000._fp * pc(k,i) / (3.8_fp * qvs)+540000._fp) ) * 434 | ( std::max(qvs-qv(k,i),0._fp) / (r(k,i)*qvs) ); 435 | real tmp2 = std::max( -prod-qc(k,i) , 0._fp ); 436 | real tmp3 = qr(k,i); 437 | real ern = std::min( tmp1 , std::min( tmp2 , tmp3 ) ); 438 | 439 | // Saturation adjustment following KW eq. 3.10 440 | theta(k,i)= theta(k,i) + lv / (cp*pk(k,i)) * 441 | ( std::max( prod , -qc(k,i) ) - ern ); 442 | qv(k,i) = std::max( qv(k,i) - std::max( prod , -qc(k,i) ) + ern , 0._fp ); 443 | qc(k,i) = qc(k,i) + std::max( prod , -qc(k,i) ); 444 | qr(k,i) = qr(k,i) - ern; 445 | 446 | // Recalculate liquid water terminal velocity 447 | velqr(k,i) = 36.34_fp * pow( qr(k,i)*r(k,i) , 0.1364_fp ) * rhalf(k,i); 448 | if (k == 0 && nt == rainsplit-1) { 449 | precl(i) = precl(i) / static_cast(rainsplit); 450 | } 451 | }); 452 | 453 | } 454 | 455 | } 456 | 457 | 458 | std::string micro_name() const { return "kessler"; } 459 | 460 | }; 461 | 462 | } 463 | 464 | 465 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/gather_statistics.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "coupler.h" 3 | #include "dynamics_euler_stratified_wenofv.h" 4 | #include "microphysics_kessler.h" 5 | #include "sponge_layer.h" 6 | #include "perturb_temperature.h" 7 | #include "column_nudging.h" 8 | #include "gather_micro_statistics.h" 9 | 10 | int main(int argc, char** argv) { 11 | MPI_Init( &argc , &argv ); 12 | yakl::init(); 13 | { 14 | using yakl::intrinsics::abs; 15 | using yakl::intrinsics::maxval; 16 | yakl::timer_start("main"); 17 | 18 | // This holds all of the model's variables, dimension sizes, and options 19 | core::Coupler coupler; 20 | 21 | // Read the YAML input file for variables pertinent to running the driver 22 | if (argc <= 1) { endrun("ERROR: Must pass the input YAML filename as a parameter"); } 23 | std::string inFile(argv[1]); 24 | YAML::Node config = YAML::LoadFile(inFile); 25 | if ( !config ) { endrun("ERROR: Invalid YAML input file"); } 26 | real sim_time = config["sim_time"].as(); 27 | size_t nx_glob = config["nx_glob" ].as(); 28 | size_t ny_glob = config["ny_glob" ].as(); 29 | int nz = config["nz" ].as(); 30 | real xlen = config["xlen" ].as(); 31 | real ylen = config["ylen" ].as(); 32 | real zlen = config["zlen" ].as(); 33 | real dtphys_in = config["dt_phys" ].as(); 34 | 35 | int nens = 1; 36 | 37 | coupler.set_option( "out_prefix" , config["out_prefix"].as() ); 38 | coupler.set_option( "init_data" , config["init_data"].as() ); 39 | coupler.set_option( "out_freq" , config["out_freq" ].as() ); 40 | 41 | // Coupler state is: (1) dry density; (2) u-velocity; (3) v-velocity; (4) w-velocity; (5) temperature 42 | // (6+) tracer masses (*not* mixing ratios!) 43 | coupler.distribute_mpi_and_allocate_coupled_state(nz, ny_glob, nx_glob, nens); 44 | 45 | // Just tells the coupler how big the domain is in each dimensions 46 | coupler.set_grid( xlen , ylen , zlen ); 47 | 48 | // This is for the dycore to pull out to determine how to do idealized test cases 49 | coupler.set_option( "standalone_input_file" , inFile ); 50 | 51 | // The column nudger nudges the column-average of the model state toward the initial column-averaged state 52 | // This is primarily for the supercell test case to keep the the instability persistently strong 53 | modules::ColumnNudger column_nudger; 54 | // Microphysics performs water phase changess + hydrometeor production, transport, collision, and aggregation 55 | modules::Microphysics_Kessler micro; 56 | // They dynamical core "dycore" integrates the Euler equations and performans transport of tracers 57 | modules::Dynamics_Euler_Stratified_WenoFV dycore; 58 | // To gather statistics on how frequently microphysics is active 59 | custom_modules::StatisticsGatherer statistics_gatherer; 60 | 61 | // Run the initialization modules 62 | micro .init ( coupler ); // Allocate micro state and register its tracers in the coupler 63 | dycore.init ( coupler ); // Dycore should initialize its own state here 64 | column_nudger.set_column ( coupler ); // Set the column before perturbing 65 | modules::perturb_temperature( coupler ); // Randomly perturb bottom layers of temperature to initiate convection 66 | 67 | real etime = 0; // Elapsed time 68 | 69 | real dtphys = dtphys_in; 70 | while (etime < sim_time) { 71 | // If dtphys <= 0, then set it to the dynamical core's max stable time step 72 | if (dtphys_in <= 0.) { dtphys = dycore.compute_time_step(coupler); } 73 | // If we're about to go past the final time, then limit to time step to exactly hit the final time 74 | if (etime + dtphys > sim_time) { dtphys = sim_time - etime; } 75 | 76 | // Run the runtime modules 77 | dycore.time_step ( coupler , dtphys ); 78 | 79 | // Create a coupler and clone the current coupler's state into it before microphysics changes it 80 | core::Coupler input; 81 | coupler.clone_into(input); 82 | micro .time_step ( coupler , dtphys ); 83 | // Now that micro has run, the current coupler state is the "output" of micro 84 | // Let's use the custom module function below to determine what proportion of the cells have active microphysics 85 | statistics_gatherer.gather_micro_statistics(input,coupler,dtphys,etime); 86 | 87 | modules::sponge_layer ( coupler , dtphys ); // Damp spurious waves to the horiz. mean at model top 88 | column_nudger.nudge_to_column( coupler , dtphys ); 89 | 90 | etime += dtphys; // Advance elapsed time 91 | } 92 | 93 | // TODO: Add finalize( coupler ) modules here 94 | statistics_gatherer.finalize( coupler ); 95 | 96 | yakl::timer_stop("main"); 97 | } 98 | yakl::finalize(); 99 | MPI_Finalize(); 100 | } 101 | 102 | 103 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/generate_micro_data.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "coupler.h" 3 | #include "dynamics_euler_stratified_wenofv.h" 4 | #include "microphysics_kessler.h" 5 | #include "sponge_layer.h" 6 | #include "perturb_temperature.h" 7 | #include "column_nudging.h" 8 | #include "generate_micro_surrogate_data.h" 9 | 10 | int main(int argc, char** argv) { 11 | MPI_Init( &argc , &argv ); 12 | yakl::init(); 13 | { 14 | using yakl::intrinsics::abs; 15 | using yakl::intrinsics::maxval; 16 | yakl::timer_start("main"); 17 | 18 | // This holds all of the model's variables, dimension sizes, and options 19 | core::Coupler coupler; 20 | 21 | // Read the YAML input file for variables pertinent to running the driver 22 | if (argc <= 1) { endrun("ERROR: Must pass the input YAML filename as a parameter"); } 23 | std::string inFile(argv[1]); 24 | YAML::Node config = YAML::LoadFile(inFile); 25 | if ( !config ) { endrun("ERROR: Invalid YAML input file"); } 26 | real sim_time = config["sim_time"].as(); 27 | size_t nx_glob = config["nx_glob" ].as(); 28 | size_t ny_glob = config["ny_glob" ].as(); 29 | int nz = config["nz" ].as(); 30 | real xlen = config["xlen" ].as(); 31 | real ylen = config["ylen" ].as(); 32 | real zlen = config["zlen" ].as(); 33 | real dtphys_in = config["dt_phys" ].as(); 34 | 35 | int nens = 1; 36 | 37 | coupler.set_option( "out_prefix" , config["out_prefix"].as() ); 38 | coupler.set_option( "init_data" , config["init_data"].as() ); 39 | coupler.set_option( "out_freq" , config["out_freq" ].as() ); 40 | 41 | // Coupler state is: (1) dry density; (2) u-velocity; (3) v-velocity; (4) w-velocity; (5) temperature 42 | // (6+) tracer masses (*not* mixing ratios!) 43 | coupler.distribute_mpi_and_allocate_coupled_state(nz, ny_glob, nx_glob, nens); 44 | 45 | // Just tells the coupler how big the domain is in each dimensions 46 | coupler.set_grid( xlen , ylen , zlen ); 47 | 48 | // This is for the dycore to pull out to determine how to do idealized test cases 49 | coupler.set_option( "standalone_input_file" , inFile ); 50 | 51 | // The column nudger nudges the column-average of the model state toward the initial column-averaged state 52 | // This is primarily for the supercell test case to keep the the instability persistently strong 53 | modules::ColumnNudger column_nudger; 54 | // Microphysics performs water phase changess + hydrometeor production, transport, collision, and aggregation 55 | modules::Microphysics_Kessler micro; 56 | // They dynamical core "dycore" integrates the Euler equations and performans transport of tracers 57 | modules::Dynamics_Euler_Stratified_WenoFV dycore; 58 | // This is the class whose methods will generate samples for micro surrogate data 59 | custom_modules::DataGenerator data_generator; 60 | 61 | // Run the initialization modules 62 | micro .init ( coupler ); // Allocate micro state and register its tracers in the coupler 63 | dycore.init ( coupler ); // Dycore should initialize its own state here 64 | column_nudger.set_column ( coupler ); // Set the column before perturbing 65 | modules::perturb_temperature( coupler ); // Randomly perturb bottom layers of temperature to initiate convection 66 | data_generator.init ( coupler ); // Create the netcdf file that will hold micro surrogate data 67 | 68 | real etime = 0; // Elapsed time 69 | 70 | real dtphys = dtphys_in; 71 | while (etime < sim_time) { 72 | // If dtphys <= 0, then set it to the dynamical core's max stable time step 73 | if (dtphys_in <= 0.) { dtphys = dycore.compute_time_step(coupler); } 74 | // If we're about to go past the final time, then limit to time step to exactly hit the final time 75 | if (etime + dtphys > sim_time) { dtphys = sim_time - etime; } 76 | 77 | // Run the runtime modules 78 | dycore.time_step ( coupler , dtphys ); 79 | 80 | // Create a coupler snapshot before the micro is run as inputs to the micro routine 81 | core::Coupler input; 82 | coupler.clone_into(input); 83 | // Run microphysics 84 | micro .time_step ( coupler , dtphys ); 85 | // Generate samples for micro's effects in the coupler. Current coupler state is the output 86 | data_generator.generate_samples_stencil( input , coupler , dtphys , etime ); 87 | 88 | modules::sponge_layer ( coupler , dtphys ); // Damp spurious waves to the horiz. mean at model top 89 | column_nudger.nudge_to_column( coupler , dtphys ); 90 | 91 | etime += dtphys; // Advance elapsed time 92 | } 93 | 94 | // TODO: Add finalize( coupler ) modules here 95 | 96 | yakl::timer_stop("main"); 97 | } 98 | yakl::finalize(); 99 | MPI_Finalize(); 100 | } 101 | 102 | 103 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/inference_ponni.cpp: -------------------------------------------------------------------------------- 1 | 2 | #include "coupler.h" 3 | #include "dynamics_euler_stratified_wenofv.h" 4 | #include "microphysics_kessler_ponni.h" 5 | #include "sponge_layer.h" 6 | #include "perturb_temperature.h" 7 | #include "column_nudging.h" 8 | #include "gather_micro_statistics.h" 9 | 10 | int main(int argc, char** argv) { 11 | MPI_Init( &argc , &argv ); 12 | yakl::init(); 13 | { 14 | using yakl::intrinsics::abs; 15 | using yakl::intrinsics::maxval; 16 | yakl::timer_start("main"); 17 | 18 | // This holds all of the model's variables, dimension sizes, and options 19 | core::Coupler coupler; 20 | 21 | // Read the YAML input file for variables pertinent to running the driver 22 | if (argc <= 1) { endrun("ERROR: Must pass the input YAML filename as a parameter"); } 23 | std::string inFile(argv[1]); 24 | YAML::Node config = YAML::LoadFile(inFile); 25 | if ( !config ) { endrun("ERROR: Invalid YAML input file"); } 26 | real sim_time = config["sim_time"].as(); 27 | size_t nx_glob = config["nx_glob" ].as(); 28 | size_t ny_glob = config["ny_glob" ].as(); 29 | int nz = config["nz" ].as(); 30 | real xlen = config["xlen" ].as(); 31 | real ylen = config["ylen" ].as(); 32 | real zlen = config["zlen" ].as(); 33 | real dtphys_in = config["dt_phys" ].as(); 34 | 35 | int nens = 1; 36 | 37 | coupler.set_option( "out_prefix" , config["out_prefix"].as() ); 38 | coupler.set_option( "init_data" , config["init_data"].as() ); 39 | coupler.set_option( "out_freq" , config["out_freq" ].as() ); 40 | 41 | // Coupler state is: (1) dry density; (2) u-velocity; (3) v-velocity; (4) w-velocity; (5) temperature 42 | // (6+) tracer masses (*not* mixing ratios!) 43 | coupler.distribute_mpi_and_allocate_coupled_state(nz, ny_glob, nx_glob, nens); 44 | 45 | // Just tells the coupler how big the domain is in each dimensions 46 | coupler.set_grid( xlen , ylen , zlen ); 47 | 48 | // This is for the dycore to pull out to determine how to do idealized test cases 49 | coupler.set_option( "standalone_input_file" , inFile ); 50 | 51 | // The column nudger nudges the column-average of the model state toward the initial column-averaged state 52 | // This is primarily for the supercell test case to keep the the instability persistently strong 53 | modules::ColumnNudger column_nudger; 54 | // Microphysics performs water phase changess + hydrometeor production, transport, collision, and aggregation 55 | custom_modules::Microphysics_Kessler micro; 56 | // They dynamical core "dycore" integrates the Euler equations and performans transport of tracers 57 | modules::Dynamics_Euler_Stratified_WenoFV dycore; 58 | 59 | // Run the initialization modules 60 | micro .init ( coupler ); // Allocate micro state and register its tracers in the coupler 61 | dycore.init ( coupler ); // Dycore should initialize its own state here 62 | column_nudger.set_column ( coupler ); // Set the column before perturbing 63 | modules::perturb_temperature( coupler ); // Randomly perturb bottom layers of temperature to initiate convection 64 | 65 | real etime = 0; // Elapsed time 66 | 67 | real dtphys = dtphys_in; 68 | while (etime < sim_time) { 69 | // If dtphys <= 0, then set it to the dynamical core's max stable time step 70 | if (dtphys_in <= 0.) { dtphys = dycore.compute_time_step(coupler); } 71 | // If we're about to go past the final time, then limit to time step to exactly hit the final time 72 | if (etime + dtphys > sim_time) { dtphys = sim_time - etime; } 73 | 74 | // Run the runtime modules 75 | dycore.time_step ( coupler , dtphys ); 76 | micro .time_step ( coupler , dtphys ); 77 | modules::sponge_layer ( coupler , dtphys ); // Damp spurious waves to the horiz. mean at model top 78 | column_nudger.nudge_to_column( coupler , dtphys ); 79 | 80 | etime += dtphys; // Advance elapsed time 81 | } 82 | 83 | yakl::timer_stop("main"); 84 | } 85 | yakl::finalize(); 86 | MPI_Finalize(); 87 | } 88 | 89 | 90 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/inputs/examples/supercell_kessler_singlecell_model_weights.h5: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mrnorman/miniWeatherML/ebd64a60d5c37932a0126e4ba3f0da846050bdff/experiments/supercell_kessler_surrogate/inputs/examples/supercell_kessler_singlecell_model_weights.h5 -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/inputs/examples/supercell_kessler_stencil_input_scaling.txt: -------------------------------------------------------------------------------- 1 | 201.8189 302.2934 2 | 0.092945546 1.1441816 3 | 0.0 0.019461675 4 | 0.0 0.004399828 5 | 0.0 0.015578972 6 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/inputs/examples/supercell_kessler_stencil_output_scaling.txt: -------------------------------------------------------------------------------- 1 | 201.81894 302.29324 2 | 0.0 0.01946172 3 | 0.0 0.0044183163 4 | 0.0 0.0155778695 5 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/inputs/input_euler3d.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | # Simulation time in seconds 3 | sim_time: 86400 4 | 5 | # Number of cells to use in the CRMs 6 | nx_glob: 250 7 | ny_glob: 1 8 | nz : 50 9 | nens : 1 10 | 11 | # Domain size of the CRMs 12 | xlen: 100000 13 | ylen: 100000 14 | zlen: 20000 15 | 16 | init_data: supercell 17 | 18 | # Output filename 19 | out_prefix: test 20 | 21 | # GCM time step 22 | dt_gcm: 900 23 | 24 | # CRM physics time step 25 | dt_phys: 0. 26 | 27 | # Output frequency in seconds 28 | out_freq: 60. 29 | 30 | keras_weights_h5: "./inputs/examples/supercell_kessler_singlecell_model_weights.h5" 31 | 32 | nn_input_scaling: "./inputs/examples/supercell_kessler_stencil_input_scaling.txt" 33 | 34 | nn_output_scaling: "./inputs/examples/supercell_kessler_stencil_output_scaling.txt" 35 | 36 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/jupyter_notebooks/README.md: -------------------------------------------------------------------------------- 1 | # Introduction 2 | 3 | This folder contains Python scripts in IPython notebooks (Jupyter notebooks) framework to demonstrate the model training in machine learning (ML) for the Kessler microphysics surrogate. The ML tutorial is in IPython notebook formats, written using [Jupyter](https://jupyter.org/). JupyterLab or Google Colab is the recommended environment for running the notebooks and further development. The notebooks are meant to be used as a platform for further development, as a stepping stone. For some basics and lessons on ML, please go to this [page](https://github.com/muralikrishnangm/tutorial-ai4science-fluidflow/wiki/ML-lessons-courses-for-beginners) on this [repo](https://github.com/muralikrishnangm/tutorial-ai4science-fluidflow) which is a more general introduction to ML applications in fluid flows and climate science. 4 | 5 | 6 | Author: Matt Norman (Oak Ridge National Laboratory), https://mrnorman.github.io/ 7 | 8 | Contributors so far: 9 | * Matt Norman 10 | * Muralikrishnan Gopalakrishnan Meena (Oak Ridge National Laboratory), https://sites.google.com/view/muraligm/ 11 | 12 | # How to run the notebooks? 13 | 14 | Here are some basics of running the IPython notebooks. Mainly, two environments are recommended: 15 | 16 | 1. [Google Colab](https://colab.research.google.com/) 17 | 2. [JupyterLab](https://github.com/jupyterlab/jupyterlab) 18 | 19 | * To run the cells, use the different options in the `Run` option in the toolbar (at the top) or press `shift+return` after selecting a cell. 20 | 21 | ## Tips 22 | * Look for the `Table of Contents` option in the notebooks (usually on the left toolbar) to get a big picture view of the ML training procedure. Both Google Colab and JupyterLab automatically have a Table of Content for all the cells in the notebook. Check it out on the toolbar on the left side. Clicking on the sections allows you to easily navigate through different parts of the notebook. 23 | 24 | * You can convert the notebooks to executable scripts (Python) using [jupyter nbconvert](https://nbconvert.readthedocs.io/en/latest/usage.html#executable-script) 25 | ``` 26 | jupyter nbconvert --to script my_notebook.ipynb 27 | ``` 28 | This will be useful when you want to eventually train the model on HPC environment. 29 | 30 | [Back to Top](#introduction) 31 | 32 | ## Opening notebooks in Google Colab 33 | 34 | This is the easiest and recommended way to use the notebooks during this tutorial as Google Collab should have all the necessary libraries. Follow the procedure below to open the notebooks in Google Colab - all you need is a Google account: 35 | 36 | * Lookout for the badge [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/muralikrishnangm/tutorial-ai4science-fluidflow/blob/main/HelloWorld.ipynb) in this README file for each notebook. Click this badge to open the notebook in Google Colab. 37 | * You need to have a Google account to **run** the notebook, and Colab will ask you to sign into your Google account. 38 | * Google Colab will ask you to "Authorize with GitHub". This allows you to save a copy of your local notebook edits to GitHub. 39 | - If you choose to authorize, you will be given a pop-up window to sign into GitHub. Be sure to disable your pop-up blocker for this site. 40 | * When you see "Warning: This notebook was not authored by Google.", please click "Run anyway". We promise this notebook is safe. 41 | * To save changes made, you need to save a copy in one of the following 42 | 1. in your Google Drive: `File -> Save a copy in Drive` 43 | 2. in your GitHub: `File -> Save a copy in GitHub` 44 | 45 | **Are you running into strange errors?** If you're seeing errors to the tune of, "Could not load the JavaScript files needed to display output", it could be that you're running on a network that has a firewall that's interfering with Colab. If you're using a VPN, please try turning it off. 46 | 47 | ## Opening notebooks in JupyterLab 48 | 49 | The notes below are curated for a tutorial at Oak Ridge Leadership Computing Facility (OLCF). Nonetheless, a local or server version of JupyterLab would follow the same procedure. 50 | 51 | * Start Jupyterlab: 52 | * Normal Jupyter installation: 53 | * Run the following on terminal to open the user interface in local internet browser 54 | ``` 55 | jupyter-lab 56 | ``` 57 | * For OLCF: 58 | * The JupyetLab environment at OLCF can be accessed through [OLCF JupyterLab](https://jupyter.olcf.ornl.gov/). All members of currently enabled OLCF projects have access to the OLCF JupyterLab (through JupyterHub). Detailed documentation can be found [here](https://docs.olcf.ornl.gov/services_and_applications/jupyter/overview.html#jupyter-at-olcf). 59 | * **IMPORTANT**: The OLCF JupyterLab is free to use for all members but it is a VERY limited resource. Please try your best to use Google Colab and opt for this option only if necessary. 60 | * Sign in using your OLCF credentials (the one used for logging in to the OLCF machines). 61 | * When selecting the options for the Labs, select the one with the GPU: 62 | ``` 63 | Slate - GPU Lab 64 | JupyterLab 3 | 16 CPU | 16GB MEM | V100 GPU 65 | ``` 66 | * Clone this repo 67 | ``` 68 | git clone git@github.com:mrnorman/miniWeatherML.git 69 | ``` 70 | * Open this [README.md](README.md) file using `Markdown Preview`: 71 | ``` 72 | [Right-click-file] -> Open With -> Markdown Preview 73 | ``` 74 | * Follow the instructions in this README file. Clicking the highlighted text for each notebook should open the notebook in a new tab in the JupyterLab environment. If it does not work, please open the notebooks directly from the file browser. 75 | 76 | [Back to Top](#introduction) 77 | 78 | # Run: Training NNs 79 | 80 | We have created a Jupyter notebook describing the training phase of the ML model. We use the Keras modeling framework for creating the ML model and training it. 81 | 82 | * Open the notebook [kessler_singlecell_train_example.ipynb](kessler_singlecell_train_example.ipynb) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mrnorman/miniWeatherML/blob/main/experiments/supercell_kessler_surrogate/jupyter_notebooks/kessler_singlecell_train_example.ipynb) 83 | * ML model for emulating microphysics in supercell test case **using Keras**. 84 | * A simple introductory notebook to ML. 85 | * Advantages of Keras: 86 | - It is a high-level API capable of running on top of TensorFlow and other frameworks. 87 | - Very concise and readable. Easy to understand and implement the logical flow of ML. 88 | * Disadvantages of Keras: 89 | - Difficult to implement complex, custom architectures. 90 | - Not easy to debug these complex architectures. 91 | 92 | We are also providing a supplementary notebook describing the post-processing step involoved in converting the data generated from the solver to the ML readable format. 93 | 94 | * Open the notebook [kessler_netcdf_to_numpy.ipynb](kessler_netcdf_to_numpy.ipynb) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mrnorman/miniWeatherML/blob/main/experiments/supercell_kessler_surrogate/jupyter_notebooks/kessler_netcdf_to_numpy.ipynb) 95 | 96 | * Data conversion from netCDF file to Python numpy array. 97 | * The data required to train the model in the previous notebook is downloaded from a server. This notebook is only provided as a supplementray resource for clarity in data curation. 98 | 99 | [Back to Top](#introduction) 100 | 101 | -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/jupyter_notebooks/images/NN_framework.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mrnorman/miniWeatherML/ebd64a60d5c37932a0126e4ba3f0da846050bdff/experiments/supercell_kessler_surrogate/jupyter_notebooks/images/NN_framework.png -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/jupyter_notebooks/images/microphysics_Kessler.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mrnorman/miniWeatherML/ebd64a60d5c37932a0126e4ba3f0da846050bdff/experiments/supercell_kessler_surrogate/jupyter_notebooks/images/microphysics_Kessler.png -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/jupyter_notebooks/images/microphysics_ML_singlecell.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mrnorman/miniWeatherML/ebd64a60d5c37932a0126e4ba3f0da846050bdff/experiments/supercell_kessler_surrogate/jupyter_notebooks/images/microphysics_ML_singlecell.png -------------------------------------------------------------------------------- /experiments/supercell_kessler_surrogate/jupyter_notebooks/kessler_netcdf_to_numpy.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "# Introduction\n", 8 | "\n", 9 | "Python script for data acquisition for emulating microphysics (Kessler) in supercell (climate) test case - load and extract data from `netCDF` file to `numpy` array for ML model\n", 10 | "\n", 11 | "Microphysics consits of 4 flow variables - temperature, water vapor, cloud water \\[liquid\\] & precipitation/rain \\[liquid\\].\n", 12 | "\n", 13 | "* **Input data**: Microphysics of a single grid cell with dry air density\n", 14 | " - Size of a single input to NN model: $N_{\\text{micro}} + 1 = 5$ for 2D/3D simulation\n", 15 | "* **Output data**: Microphysics of the given cell after emmulation (at next time step)\n", 16 | " - Size of corresponding output from NN model: $[N_{\\text{micro}}] = [4]$ for 2D/3D simulation\n", 17 | "* **Training data size**:\n", 18 | " - Input: $5 \\times N_{\\text{train}}$\n", 19 | " - Output: $4 \\times N_{\\text{train}}$\n", 20 | "\n", 21 | "\n", 22 | "**By Matt Norman and Murali Gopalakrishnan Meena, ORNL**" 23 | ] 24 | }, 25 | { 26 | "cell_type": "code", 27 | "execution_count": 1, 28 | "metadata": { 29 | "colab": { 30 | "base_uri": "https://localhost:8080/" 31 | }, 32 | "id": "NMq3yYEeoIgv", 33 | "outputId": "61c9b09c-5c25-4d25-e7a6-fef80d2ab571" 34 | }, 35 | "outputs": [ 36 | { 37 | "name": "stdout", 38 | "output_type": "stream", 39 | "text": [ 40 | "Collecting netCDF4\n", 41 | " Downloading netCDF4-1.6.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (5.2 MB)\n", 42 | " |████████████████████████████████| 5.2 MB 3.8 MB/s \n", 43 | "\u001b[?25hCollecting cftime\n", 44 | " Downloading cftime-1.6.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (241 kB)\n", 45 | " |████████████████████████████████| 241 kB 115.3 MB/s \n", 46 | "\u001b[?25hRequirement already satisfied: numpy>=1.9 in /opt/conda/lib/python3.8/site-packages (from netCDF4) (1.19.5)\n", 47 | "Installing collected packages: cftime, netCDF4\n", 48 | "Successfully installed cftime-1.6.1 netCDF4-1.6.0\n", 49 | "Downloading data from:\n", 50 | " https://www.dropbox.com/s/nonpheml3309q7d/supercell_kessler_data.nc?dl=0...\n", 51 | "--2022-07-24 15:30:30-- https://www.dropbox.com/s/nonpheml3309q7d/supercell_kessler_data.nc?dl=0\n", 52 | "Resolving www.dropbox.com (www.dropbox.com)... 162.125.9.18, 2620:100:601f:18::a27d:912\n", 53 | "Connecting to www.dropbox.com (www.dropbox.com)|162.125.9.18|:443... connected.\n", 54 | "HTTP request sent, awaiting response... 301 Moved Permanently\n", 55 | "Location: /s/raw/nonpheml3309q7d/supercell_kessler_data.nc [following]\n", 56 | "--2022-07-24 15:30:31-- https://www.dropbox.com/s/raw/nonpheml3309q7d/supercell_kessler_data.nc\n", 57 | "Reusing existing connection to www.dropbox.com:443.\n", 58 | "HTTP request sent, awaiting response... 302 Found\n", 59 | "Location: https://uc65ee8f83cd11cb46f653f35835.dl.dropboxusercontent.com/cd/0/inline/Bps2o4mWePhwRRuY2Y3zqXeKSWf7CiM4BmqIjiOH5rBoNHbHbP4VvGXjrXD3Lw8FcyzFnAyxJk7PJQoIMWzwsapboZYGFOdIm6TscDeNoxrts3-FtGcJwkBMbalEpUCFG2C57sc03-sYtOD_1jnFqx0DGQ9i8Ilb9o1cm-Jr2xPfbw/file# [following]\n", 60 | "--2022-07-24 15:30:31-- https://uc65ee8f83cd11cb46f653f35835.dl.dropboxusercontent.com/cd/0/inline/Bps2o4mWePhwRRuY2Y3zqXeKSWf7CiM4BmqIjiOH5rBoNHbHbP4VvGXjrXD3Lw8FcyzFnAyxJk7PJQoIMWzwsapboZYGFOdIm6TscDeNoxrts3-FtGcJwkBMbalEpUCFG2C57sc03-sYtOD_1jnFqx0DGQ9i8Ilb9o1cm-Jr2xPfbw/file\n", 61 | "Resolving uc65ee8f83cd11cb46f653f35835.dl.dropboxusercontent.com (uc65ee8f83cd11cb46f653f35835.dl.dropboxusercontent.com)... 162.125.9.15, 2620:100:601f:15::a27d:90f\n", 62 | "Connecting to uc65ee8f83cd11cb46f653f35835.dl.dropboxusercontent.com (uc65ee8f83cd11cb46f653f35835.dl.dropboxusercontent.com)|162.125.9.15|:443... connected.\n", 63 | "HTTP request sent, awaiting response... 302 Found\n", 64 | "Location: /cd/0/inline2/BpvrlV0h5J0uT0DuW36Oh5sHtQvapVuU74HSaUzBSYaxtzptUns_vOgl_R6CDmIU9JdQwYYnSRs2YnqWwezR4gvID3MwPs_cmcPEkrP1-yDsQMUQHL_Obe_9DVSqGC0L_517ZAm_vmLRKKbW4osxdKfLmNiJC4T9r6pWRRf5Y1gPRdhpa-xrLkzh-skEHroK5yDE2uytzbPxKbi2o3yK8k3ikEnQjINnbhbCWCXb4-RuQ-o3ZKAjWlyvGYjVdhaCJAMlSacfXlhMsL-GeQPKV_DpIEP5wu8-iqa8ldJq3XSKXTk-XSn9Cvox7-4iuRwftu_YHp-V2qlyP1SrzI8p52p9X2k2igbl7D6NiTyk92TTu0u1Ks2RlNnJVSt78_rlyuOcfnibDWxXoKLmtVZI3OX4pb98gArMc4g3H8ws2weh6Q/file [following]\n", 65 | "--2022-07-24 15:30:31-- https://uc65ee8f83cd11cb46f653f35835.dl.dropboxusercontent.com/cd/0/inline2/BpvrlV0h5J0uT0DuW36Oh5sHtQvapVuU74HSaUzBSYaxtzptUns_vOgl_R6CDmIU9JdQwYYnSRs2YnqWwezR4gvID3MwPs_cmcPEkrP1-yDsQMUQHL_Obe_9DVSqGC0L_517ZAm_vmLRKKbW4osxdKfLmNiJC4T9r6pWRRf5Y1gPRdhpa-xrLkzh-skEHroK5yDE2uytzbPxKbi2o3yK8k3ikEnQjINnbhbCWCXb4-RuQ-o3ZKAjWlyvGYjVdhaCJAMlSacfXlhMsL-GeQPKV_DpIEP5wu8-iqa8ldJq3XSKXTk-XSn9Cvox7-4iuRwftu_YHp-V2qlyP1SrzI8p52p9X2k2igbl7D6NiTyk92TTu0u1Ks2RlNnJVSt78_rlyuOcfnibDWxXoKLmtVZI3OX4pb98gArMc4g3H8ws2weh6Q/file\n", 66 | "Reusing existing connection to uc65ee8f83cd11cb46f653f35835.dl.dropboxusercontent.com:443.\n", 67 | "HTTP request sent, awaiting response... 200 OK\n", 68 | "Length: 1599928996 (1.5G) [application/x-netcdf]\n", 69 | "Saving to: ‘supercell_kessler_data.nc’\n", 70 | "\n", 71 | "supercell_kessler_d 100%[===================>] 1.49G 44.7MB/s in 29s \n", 72 | "\n", 73 | "2022-07-24 15:31:01 (53.2 MB/s) - ‘supercell_kessler_data.nc’ saved [1599928996/1599928996]\n", 74 | "\n", 75 | "Reading dataset...\n", 76 | " * Finished reading chunk 1 of 20\n", 77 | " * Finished reading chunk 2 of 20\n", 78 | " * Finished reading chunk 3 of 20\n", 79 | " * Finished reading chunk 4 of 20\n", 80 | " * Finished reading chunk 5 of 20\n", 81 | " * Finished reading chunk 6 of 20\n", 82 | " * Finished reading chunk 7 of 20\n", 83 | " * Finished reading chunk 8 of 20\n", 84 | " * Finished reading chunk 9 of 20\n", 85 | " * Finished reading chunk 10 of 20\n", 86 | " * Finished reading chunk 11 of 20\n", 87 | " * Finished reading chunk 12 of 20\n", 88 | " * Finished reading chunk 13 of 20\n", 89 | " * Finished reading chunk 14 of 20\n", 90 | " * Finished reading chunk 15 of 20\n", 91 | " * Finished reading chunk 16 of 20\n", 92 | " * Finished reading chunk 17 of 20\n", 93 | " * Finished reading chunk 18 of 20\n", 94 | " * Finished reading chunk 19 of 20\n", 95 | " * Finished reading chunk 20 of 20\n", 96 | "Shuffling dataset...\n", 97 | "Saving data to file...\n" 98 | ] 99 | } 100 | ], 101 | "source": [ 102 | "!pip install netCDF4\n", 103 | "from netCDF4 import Dataset\n", 104 | "import numpy as np\n", 105 | "import os\n", 106 | "\n", 107 | "path = f'supercell_kessler_data.nc'\n", 108 | "data_link = \"https://www.dropbox.com/s/nonpheml3309q7d/supercell_kessler_data.nc?dl=0\"\n", 109 | "\n", 110 | "# Download the data if necessary\n", 111 | "if ( not os.path.isfile(path) ):\n", 112 | " print(f\"Downloading data from:\\n {data_link}...\")\n", 113 | " !wget {data_link} -O {path}\n", 114 | "\n", 115 | "print('Reading dataset...')\n", 116 | "\n", 117 | "# Open NetCDF4 file, allocate input and output data arrays\n", 118 | "nc = Dataset(path,'r')\n", 119 | "[num_samples, num_vars_in, stencil_size] = nc.variables['inputs'].shape\n", 120 | "input_from_file = np.ndarray(shape=nc.variables['inputs' ].shape,dtype=np.single)\n", 121 | "output_from_file = np.ndarray(shape=nc.variables['outputs'].shape,dtype=np.single)\n", 122 | "\n", 123 | "# We need to chunk the reading to avoid overflowing available memory\n", 124 | "num_chunks = 20\n", 125 | "chunk_size = int(np.ceil(num_samples / num_chunks))\n", 126 | "# Loop over chunks and load data\n", 127 | "for ichunk in range(num_chunks) :\n", 128 | " ibeg = int( ichunk *chunk_size)\n", 129 | " iend = int((ichunk+1)*chunk_size)\n", 130 | " if (ichunk == num_chunks-1) : # Ensure we don't go past the last index\n", 131 | " iend = num_samples\n", 132 | " input_from_file [ibeg:iend,:,:] = nc.variables['inputs' ][ibeg:iend,:,:]\n", 133 | " output_from_file[ibeg:iend,:] = nc.variables['outputs'][ibeg:iend,:]\n", 134 | " print(f' * Finished reading chunk {ichunk+1} of {num_chunks}')\n", 135 | "\n", 136 | "nc.close()\n", 137 | "\n", 138 | "print('Shuffling dataset...')\n", 139 | "\n", 140 | "# Randomly shuffle the samples before saving to file\n", 141 | "permuted_indices = np.random.permutation(np.arange(0, num_samples))\n", 142 | "input_from_file = input_from_file [permuted_indices[:],:,:]\n", 143 | "output_from_file = output_from_file[permuted_indices[:],:]\n", 144 | "\n", 145 | "print('Saving data to file...')\n", 146 | "\n", 147 | "np.savez('supercell_kessler_data.npz',\n", 148 | " input_from_file=input_from_file, output_from_file=output_from_file)\n" 149 | ] 150 | } 151 | ], 152 | "metadata": { 153 | "colab": { 154 | "name": "kessler_netcdf_to_numpy.ipynb", 155 | "provenance": [] 156 | }, 157 | "kernelspec": { 158 | "display_name": "OLCF-CUDA11 (ipykernel)", 159 | "language": "python", 160 | "name": "python3" 161 | }, 162 | "language_info": { 163 | "codemirror_mode": { 164 | "name": "ipython", 165 | "version": 3 166 | }, 167 | "file_extension": ".py", 168 | "mimetype": "text/x-python", 169 | "name": "python", 170 | "nbconvert_exporter": "python", 171 | "pygments_lexer": "ipython3", 172 | "version": "3.8.10" 173 | } 174 | }, 175 | "nbformat": 4, 176 | "nbformat_minor": 4 177 | } 178 | -------------------------------------------------------------------------------- /model/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | 2 | set(YAKL_HOME ${CMAKE_CURRENT_SOURCE_DIR}/../external/YAKL) 3 | set(YAKL_HOME ${CMAKE_CURRENT_SOURCE_DIR}/../external/YAKL PARENT_SCOPE) 4 | set(YAKL_BIN ${CMAKE_CURRENT_BINARY_DIR}/yakl) 5 | set(YAKL_BIN ${CMAKE_CURRENT_BINARY_DIR}/yakl PARENT_SCOPE) 6 | add_subdirectory(${YAKL_HOME} ${YAKL_BIN}) 7 | include_directories(${YAKL_BIN}) 8 | 9 | set(YAML_CPP_BUILD_TOOLS OFF CACHE BOOL "" FORCE) 10 | set(YAML_BUILD_SHARED_LIBS OFF CACHE BOOL "" FORCE) 11 | set(YAML_CPP_BUILD_CONTRIB OFF CACHE BOOL "" FORCE) 12 | add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../external/yaml-cpp ${CMAKE_CURRENT_BINARY_DIR}/yaml-cpp) 13 | 14 | add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../external/ponni ponni) 15 | 16 | include_directories(${CMAKE_CURRENT_SOURCE_DIR}) 17 | 18 | add_subdirectory(core core ) 19 | add_subdirectory(modules modules) 20 | 21 | add_library(model INTERFACE) 22 | target_link_libraries(model INTERFACE modules core ponni yaml-cpp) 23 | target_include_directories(model INTERFACE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/../external/yaml-cpp/include) 24 | 25 | -------------------------------------------------------------------------------- /model/core/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | 2 | add_library(core INTERFACE) 3 | target_include_directories(core INTERFACE ${CMAKE_CURRENT_SOURCE_DIR}) 4 | target_include_directories(core INTERFACE ${CMAKE_CURRENT_BINARY_DIR}) 5 | 6 | -------------------------------------------------------------------------------- /model/core/MultipleFields.h: -------------------------------------------------------------------------------- 1 | 2 | #pragma once 3 | 4 | #include "main_header.h" 5 | 6 | namespace core { 7 | 8 | // Aggregate multiple fields into a single field that makes it easier to operate 9 | // on them together inside the same kernel. used mostly for tracers 10 | template 11 | class MultipleFields { 12 | public: 13 | yakl::SArray fields; 14 | int num_fields; 15 | 16 | YAKL_INLINE MultipleFields() { num_fields = 0; } 17 | 18 | YAKL_INLINE MultipleFields(MultipleFields const &rhs) { 19 | this->num_fields = rhs.num_fields; 20 | for (int i=0; i < num_fields; i++) { 21 | this->fields(i) = rhs.fields(i); 22 | } 23 | } 24 | 25 | YAKL_INLINE MultipleFields & operator=(MultipleFields const &rhs) { 26 | this->num_fields = rhs.num_fields; 27 | for (int i=0; i < num_fields; i++) { 28 | this->fields(i) = rhs.fields(i); 29 | } 30 | return *this; 31 | } 32 | 33 | YAKL_INLINE MultipleFields(MultipleFields &&rhs) { 34 | this->num_fields = rhs.num_fields; 35 | for (int i=0; i < num_fields; i++) { 36 | this->fields(i) = rhs.fields(i); 37 | } 38 | } 39 | 40 | YAKL_INLINE MultipleFields& operator=(MultipleFields &&rhs) { 41 | this->num_fields = rhs.num_fields; 42 | for (int i=0; i < num_fields; i++) { 43 | this->fields(i) = rhs.fields(i); 44 | } 45 | return *this; 46 | } 47 | 48 | YAKL_INLINE void add_field( T field ) { 49 | this->fields(num_fields) = field; 50 | num_fields++; 51 | } 52 | 53 | YAKL_INLINE T &get_field( int tr ) const { 54 | return this->fields(tr); 55 | } 56 | 57 | YAKL_INLINE int get_num_fields() const { return num_fields; } 58 | 59 | YAKL_INLINE auto operator() (int tr, int i1) const -> 60 | decltype(fields(tr)(i1)) { 61 | return this->fields(tr)(i1); 62 | } 63 | YAKL_INLINE auto operator() (int tr, int i1, int i2) const -> 64 | decltype(fields(tr)(i1,i2)) { 65 | return this->fields(tr)(i1,i2); 66 | } 67 | YAKL_INLINE auto operator() (int tr, int i1, int i2, int i3) const -> 68 | decltype(fields(tr)(i1,i2,i3)) { 69 | return this->fields(tr)(i1,i2,i3); 70 | } 71 | YAKL_INLINE auto operator() (int tr, int i1, int i2, int i3, int i4) const -> 72 | decltype(fields(tr)(i1,i2,i3,i4)) { 73 | return this->fields(tr)(i1,i2,i3,i4); 74 | } 75 | YAKL_INLINE auto operator() (int tr, int i1, int i2, int i3, int i4, int i5) const -> 76 | decltype(fields(tr)(i1,i2,i3,i4,i5)) { 77 | return this->fields(tr)(i1,i2,i3,i4,i5); 78 | } 79 | YAKL_INLINE auto operator() (int tr, int i1, int i2, int i3, int i4, int i5, int i6) const -> 80 | decltype(fields(tr)(i1,i2,i3,i4,i5,i6)) { 81 | return this->fields(tr)(i1,i2,i3,i4,i5,i6); 82 | } 83 | YAKL_INLINE auto operator() (int tr, int i1, int i2, int i3, int i4, int i5, int i6, int i7) const -> 84 | decltype(fields(tr)(i1,i2,i3,i4,i5,i6,i7)) { 85 | return this->fields(tr)(i1,i2,i3,i4,i5,i6,i7); 86 | } 87 | YAKL_INLINE auto operator() (int tr, int i1, int i2, int i3, int i4, int i5, int i6, int i7, int i8) const -> 88 | decltype(fields(tr)(i1,i2,i3,i4,i5,i6,i7,i8)) { 89 | return this->fields(tr)(i1,i2,i3,i4,i5,i6,i7,i8); 90 | } 91 | }; 92 | 93 | 94 | 95 | template 96 | using MultiField = MultipleFields< max_fields , Array >; 97 | } 98 | 99 | 100 | -------------------------------------------------------------------------------- /model/core/Options.h: -------------------------------------------------------------------------------- 1 | 2 | #pragma once 3 | 4 | #include "main_header.h" 5 | 6 | // The Options class creates a way for users to set key,value pairs in the coupler 7 | // where the key is always a string, and the value can be different types 8 | 9 | namespace core { 10 | 11 | class Options { 12 | public: 13 | 14 | struct Option { 15 | std::string key; 16 | void * data; 17 | size_t type_hash; 18 | }; 19 | 20 | std::vector