Add script for testing FP precision + numerics with ocean climate simulation

.github/workflows/Compile.yml

@@ -37,6 +37,9 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
+        simulation_type:
+          - 'data_free'
+          - 'ocean_climate'
         julia_version:
           - '1.10'
           - '1.11'
@@ -54,3 +57,4 @@ jobs:
       compile_or_run: 'compile'
       earlyoom_threshold: 4
       julia_optlevel: 2
+      simulation_type: ${{ matrix.simulation_type }}

.github/workflows/CompileOrRun.yml

@@ -33,10 +33,15 @@ on:
         required: true
         default: 2
         type: number
+      simulation_type:
+        description: 'data_free or ocean_climate'
+        required: true
+        default: ''
+        type: string
 
 jobs:
-  ocean_climate_simulation:
-    name: Julia ${{ inputs.julia_version }} - ${{ inputs.os }} - ${{ inputs.xla_runtime }} - ${{ inputs.compile_or_run }}
+  simulation:
+    name: Julia ${{ inputs.julia_version }} - ${{ inputs.simulation_type }} - ${{ inputs.os }} - ${{ inputs.xla_runtime }} - ${{ inputs.compile_or_run }}
     runs-on: ${{ inputs.os }}
     strategy:
       fail-fast: false
@@ -53,7 +58,7 @@ jobs:
       - uses: julia-actions/cache@v2
         if: ${{ startsWith(inputs.os, 'ubuntu') }}
         with:
-          cache-name: julia-cache;workflow=${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.xla_runtime }}-${{ inputs.compile_or_run }};job=${{ github.job }}
+          cache-name: julia-cache;workflow=${{ inputs.julia_version }}-${{ inputs.simulation_type }}-${{ inputs.os }}-${{ inputs.xla_runtime }}-${{ inputs.compile_or_run }};job=${{ github.job }}
       - name: Collect Workflow Telemetry
         uses: catchpoint/workflow-telemetry-action@v2
         with:
@@ -83,24 +88,26 @@ jobs:
         timeout-minutes: 10
         if: ${{ always() }}
         with:
-          name: 'environment-${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.xla_runtime }}-${{ inputs.compile_or_run }}'
+          name: 'environment-${{ inputs.julia_version }}-${{ inputs.simulation_type }}-${{ inputs.os }}-${{ inputs.xla_runtime }}-${{ inputs.compile_or_run }}'
           path: |
             Manifest.toml
             Project.toml
           retention-days: 90
           overwrite: false
-      - name: Ocean climate simulation
+      - name: ${{ inputs.simulation_type }} simulation
         timeout-minutes: 180
         run: |
           earlyoom -m ${{ inputs.earlyoom_threshold }} -s 100 -r 120 --prefer 'julia' &
-          julia --color=yes --project -O${{ inputs.julia_optlevel }} ocean-climate-simulation/ocean_climate_simulation_${{ inputs.compile_or_run }}.jl
+          julia --color=yes --project -O${{ inputs.julia_optlevel }} simulations/${{ inputs.simulation_type }}_simulation_${{ inputs.compile_or_run }}.jl
         env:
           XLA_FLAGS: "--xla_dump_to=xla_dump"
+          ECCO_PASSWORD: ${{ secrets.ECCO_PASSWORD }}
+          ECCO_USERNAME: ${{ secrets.ECCO_USERNAME }}
       - name: Show remaining jit calls
         if: ${{ inputs.compile_or_run == 'compile' }}
         timeout-minutes: 10
         run: |
-          JIT_CALLS=$(find . -name 'opt_ocean_climate_simulation.mlir' -exec grep 'enzymexla\.jit_call' '{}' \;)
+          JIT_CALLS=$(find . -name 'opt_${{ inputs.simulation_type }}.mlir' -exec grep 'enzymexla\.jit_call' '{}' \;)
           NUM_JIT_CALLS=$(echo "${JIT_CALLS}" | wc -l | xargs)
           if [[ -n "${JIT_CALLS}" ]]; then
               MSG="There are ${NUM_JIT_CALLS} remaining jit calls"
@@ -116,7 +123,7 @@ jobs:
         timeout-minutes: 10
         if: ${{ always() && inputs.compile_or_run == 'compile' }}
         with:
-          name: 'ocean-simulation-mlir-${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.xla_runtime }}'
+          name: '${{ inputs.simulation_type }}-mlir-${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.xla_runtime }}'
           path: '**/*.mlir'
           retention-days: 90
           overwrite: false
@@ -125,7 +132,7 @@ jobs:
         timeout-minutes: 10
         if: ${{ always() && inputs.compile_or_run == 'compile' }}
         with:
-          name: 'ocean-simulation-julia-profile-${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.xla_runtime }}'
+          name: '${{ inputs.simulation_type }}-julia-profile-${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.xla_runtime }}'
           path: |
             **/profile_*.txt
             **/profile_*.dat
@@ -136,7 +143,7 @@ jobs:
         timeout-minutes: 10
         if: ${{ always() && inputs.compile_or_run == 'run' }}
         with:
-          name: 'ocean-simulation-xla-dump-${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.xla_runtime }}'
+          name: '${{ inputs.simulation_type }}-xla-dump-${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.xla_runtime }}'
           path: '**/xla_dump'
           retention-days: 90
           overwrite: false
@@ -145,7 +152,8 @@ jobs:
         timeout-minutes: 10
         if: ${{ always() && inputs.compile_or_run == 'run' }}
         with:
-          name: 'ocean-simulation-xla-profile-${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.xla_runtime }}'
+          name: '${{ inputs.simulation_type }}-xla-profile-${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.xla_runtime }}'
           path: '**/plugins'
           retention-days: 90
           overwrite: false
+

.github/workflows/Run.yml

@@ -37,6 +37,9 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
+        simulation_type:
+          - 'data_free'
+          - 'ocean_climate'
         julia_version:
           - '1.11'
         os:
@@ -59,3 +62,4 @@ jobs:
       compile_or_run: 'run'
       earlyoom_threshold: ${{ matrix.earlyoom_threshold }}
       julia_optlevel: 0
+      simulation_type: ${{ matrix.simulation_type }}

.gitignore

@@ -4,3 +4,5 @@
 *.pdf
 *.mp4
 *.swp
+
+Manifest.toml

Project.toml

@@ -21,10 +21,10 @@ Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 [compat]
 BenchmarkTools = "1.6.0"
 CFTime = "0.1.4"
-ClimaOcean = "0.5.4"
+ClimaOcean = "0.5"
 Dates = "1"
 HDF5_jll = "=1.14.3"
-Oceananigans = "0.95.29"
+Oceananigans = "0.96"
 PrecompileTools = "1.2.1"
 Printf = "1"
 Profile = "1"

numerics_analysis/run_ocean_climate_simulation_numerics.ipynb

@@ -0,0 +1,214 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "6a7d5cc5-bb93-4849-b8c7-977080b9352b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m\u001b[1m  Activating\u001b[22m\u001b[39m project at `/global/cfs/cdirs/m4746/Users/nora/GB-25`\n"
+     ]
+    }
+   ],
+   "source": [
+    "using Pkg\n",
+    "Pkg.activate(\"../.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "1214a398-6f7e-48aa-97ef-ce8757522c13",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[33m\u001b[1m┌ \u001b[22m\u001b[39m\u001b[33m\u001b[1mWarning: \u001b[22m\u001b[39mYou are using Julia v1.11 or later!\"\n",
+      "\u001b[33m\u001b[1m│ \u001b[22m\u001b[39mOceananigans is currently tested on Julia v1.10.\"\n",
+      "\u001b[33m\u001b[1m│ \u001b[22m\u001b[39mIf you find issues with Julia v1.11 or later,\"\n",
+      "\u001b[33m\u001b[1m│ \u001b[22m\u001b[39mplease report at https://github.com/CliMA/Oceananigans.jl/issues/new\n",
+      "\u001b[33m\u001b[1m└ \u001b[22m\u001b[39m\u001b[90m@ Oceananigans /global/cfs/projectdirs/m4746/Users/nora/.julia/packages/Oceananigans/kw9Pz/src/Oceananigans.jl:118\u001b[39m\n"
+     ]
+    }
+   ],
+   "source": [
+    "using Oceananigans"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0d8a97dd-29ff-4391-b9b7-674f7af291d1",
+   "metadata": {},
+   "source": [
+    "### Tunable parameters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "bd4b9c0d-a5fb-48d2-89e1-8f527c8ebf6d",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Float64"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Architecture and precision\n",
+    "arch = GPU() # GPU() or ReactantState()\n",
+    "FT = Float64 # Float32 or Float64"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "aaacaf32-206f-425c-875e-952f77fb113f",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "300"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Resolution-dependent parameters\n",
+    "resolution = 4 # degree\n",
+    "Nz = 20 # vertical levels\n",
+    "Δt₁ = 120 # seconds, small time step for \"spinup\"\n",
+    "Δt₂ = 300 # seconds, bigger time step for after \"spinup\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "b91df47d-5ecf-466c-b81e-d522224b91cb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Another set of parameters for a finer resolution run, which should still fit on a single GPU:\n",
+    "# resolution = 1/4\n",
+    "# Nz = 50\n",
+    "# Δt₁ = 120 # tweak this\n",
+    "# Δt₂ = 300 # tweak this"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "b172cba1-6879-49c9-a96b-21e63296e907",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Numerical scheme orders (we might have to reduce these if simulations blow up)\n",
+    "momentum_advection_order = 5  # Order of the numerical scheme for momentum advection\n",
+    "tracer_advection_order = 5    # Order of the numerical scheme for tracer advection\n",
+    "\n",
+    "# Advection scheme order for vorticity flux (which is part of the momentum equation)\n",
+    "vorticity_order = nothing     # Set to `nothing` to use the same order as momentum advection\n",
+    "                              # If order needs to be reduced due to blow-ups, try reducing vorticity_order first\n",
+    "                              # before reducing the momentum_advection_order for all terms"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "1a54c3a9-3362-4997-8afc-2d8a8973c24b",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "\"/pscratch/sd/n/nloose/GB-25/simulations/\""
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# where to save output\n",
+    "target_dir = \"/pscratch/sd/n/nloose/GB-25/simulations/\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "00907ab5-a905-42e0-9218-e326f1a9ca1d",
+   "metadata": {},
+   "source": [
+    "### Run the numerics script"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0c695f86-9cec-4222-9f03-11e581072486",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[36m\u001b[1m┌ \u001b[22m\u001b[39m\u001b[36m\u001b[1mInfo: \u001b[22m\u001b[39mSkipping passes for interpolating bathymetry of size (21600, 10800, 1) \n",
+      "\u001b[36m\u001b[1m│ \u001b[22m\u001b[39mto target grid of size (90, 45, 20). Interpolation passes may only \n",
+      "\u001b[36m\u001b[1m│ \u001b[22m\u001b[39mbe used to coarsen bathymetry and require that the bathymetry \n",
+      "\u001b[36m\u001b[1m└ \u001b[22m\u001b[39mis finer than the target grid in both horizontal directions.\n"
+     ]
+    }
+   ],
+   "source": [
+    "include(\"../simulations/ocean_climate_simulation_numerics.jl\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2f8efa27-c91e-45ca-a350-7010efc298a1",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cb754ecf-5eb6-4c8a-ba1c-1c4431d9b9d5",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "NERSC Julia 1.11.4 1 Thread(s) with CUDNN Beta",
+   "language": "julia",
+   "name": "nersc-julia-1-1.11.4-cudnn-beta"
+  },
+  "language_info": {
+   "file_extension": ".jl",
+   "mimetype": "application/julia",
+   "name": "julia",
+   "version": "1.11.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}