vort2d test case ready

added control scripts. To run the test case:
- go to config/vort2d_testcase and make changes, choose
  config/env/localhost and change inside the paths accordingly
- set -a; . config/vort2d_testcase
- run truth simulation by "cd models/vort2d; ./run_forecast.sh truth"
- run cycling by "cd scripts; ./run_cycle.sh"
- monitor job status by "cd scripts; ./jstat"
- go to $SCRATCH/vort2d_testcase/cycle... to see results

config/__init__.py

@@ -29,7 +29,7 @@
 ###check type if valid
 regress_type = os.environ.get('regress_type', 'linear')
 
-run_filter = os.environ.get('run_filter').lower()=='true'
+run_assim = os.environ.get('run_assim').lower()=='true'
 run_align_space = os.environ.get('run_align_space').lower()=='true'
 run_align_time = os.environ.get('run_align_time').lower()=='true'
 
@@ -221,3 +221,5 @@
 
 pid_show = 0  ##which pid is showing progress messages, default to root=0
 
+debug = os.environ.get('debug').lower()=='true'
+

config/defaults

@@ -23,7 +23,7 @@ assim_mode=serial ##batch or serial
 filter_type=EAKF
 regress_type=linear
 
-run_filter=true  ##if true, run filter/smoother update algorithms
+run_assim=true  ##if true, run filter/smoother update algorithms
 run_align_space=false  ##if true, run alignment after filter update in space/time
 run_align_time=false
 
@@ -138,3 +138,5 @@ perturb_variables="seaice_drift ocean_surf_height"
 
 ##physical parameters, ###set default values in namelist_<model>.sh first, if experiments frequently change a parameter, set it here
 
+debug=false
+

config/env/betzy/base.src

@@ -1,6 +1,9 @@
 #!/bin/bash
 host=betzy
-hostppn=128
+
+nnodes=$SLURM_NNODES
+ntasks=$SLURM_NTASKS
+tasks_per_node=128
 
 SCRATCH=/cluster/work/users/yingyue
 

config/env/betzy/vort2d.src

@@ -0,0 +1,15 @@
+ulimit -s 2000000
+
+# Load the right modules
+module --force purge
+ml load StdEnv
+ml load GCCcore/12.3.0
+ml load Python/3.11.3-GCCcore-12.3.0
+ml load FFTW/3.3.10-GCC-12.3.0
+ml load impi/2021.9.0-intel-compilers-2023.1.0
+
+export PYTHONPATH=$PYTHONPATH:$HOME/code/NEDAS
+
+
+##enter the project env
+source $HOME/yp/bin/activate

config/env/localhost/base.src

@@ -1,6 +1,9 @@
 #!/bin/bash
 
 host=localhost
-hostppn=4
+
+nnodes=1
+ntasks=4
+tasks_per_node=4
 
 SCRATCH=/Users/yueng/scratch

config/topaz5_testcase

@@ -23,7 +23,7 @@ assim_mode=serial ##batch or serial
 filter_type=EAKF
 regress_type=linear
 
-run_filter=true  ##if true, run filter/smoother update algorithms
+run_assim=true  ##if true, run filter/smoother update algorithms
 run_align_space=false  ##if true, run alignment after filter update in space/time
 run_align_time=false
 
@@ -140,3 +140,5 @@ perturb_variables="seaice_drift ocean_surf_height"
 
 ##physical parameters, ###set default values in namelist_<model>.sh first, if experiments frequently change a parameter, set it here
 
+debug=false
+

config/vort2d_testcase

@@ -23,7 +23,7 @@ assim_mode=batch ##batch or serial
 filter_type=ETKF
 regress_type=linear
 
-run_filter=true  ##if true, run filter/smoother update algorithms
+run_assim=true  ##if true, run filter/smoother update algorithms
 run_align_time=false
 run_align_space=false
 
@@ -111,4 +111,5 @@ use_synthetic_obs=true
 ##ensemble perturbation settings
 perturb_variables="velocity"
 
+debug=false
 

models/vort2d/generate_ic.py

@@ -0,0 +1,16 @@
+import numpy as np
+import models.vort2d as model
+import config as c
+import sys
+from conversion import s2t
+
+time = s2t(sys.argv[1])
+
+##input random seed is provided, set it
+if len(sys.argv) > 2:
+    np.random.seed(int(sys.argv[2]))
+
+state = model.initialize(c.grid, model.Vmax, model.Rmw, model.Vbg, model.Vslope)
+
+model.write_var('./', c.grid, state, name='velocity', is_vector=True, time=time)
+

models/vort2d/module_ens_forecast.sh

@@ -0,0 +1,66 @@
+#!/bin/bash
+. $config_file
+rundir=$work_dir/cycle/$time/vort2d
+if [[ ! -d $rundir ]]; then mkdir -p $rundir; echo waiting > $rundir/stat; fi
+
+cd $rundir
+if [[ `cat stat` == "complete" ]]; then exit; fi
+
+##check dependency
+if [[ $time == $time_start ]]; then wait_for_module ../perturb; fi
+if [[ $time -gt $time_start ]]; then
+    if $run_assim; then
+        wait_for_module ../analysis
+    else
+        wait_for_module ../../$prev_time/vort2d
+    fi
+fi
+
+echo running > stat
+
+echo "  Running ensemble forecast for vort2d model..."
+
+##load env if necessary
+src_file=$script_dir/../config/env/$host/vort2d.src
+if [[ -f $src_file ]]; then source $src_file; fi
+
+tid=0
+nt=$ntasks
+for m in `seq 1 $nens`; do
+    m_id=`padzero $m 3`
+    if [[ ! -d $m_id ]]; then mkdir -p $m_id; fi
+    touch $m_id/run.log
+
+    ##run the model for member m
+    cd $m_id
+
+    ##link files for model run
+    mv ../${time:0:8}_${time:8:2}_mem$m_id.nc .
+
+    $script_dir/job_submit.sh 1 1 0 python $script_dir/../models/vort2d/run.py $time $m >& run.log &
+
+    cd ..
+
+    ##wait if ntasks processors are all in use
+    tid=$((tid+1))
+    if [[ $tid == $nt ]]; then tid=0; wait; fi
+
+done
+wait
+
+nextdir=$work_dir/cycle/$next_time/vort2d
+if [[ ! -d $nextdir ]]; then mkdir -p $nextdir; fi
+
+##collect output files, make a copy of forecast files to next cycle
+for m in `seq 1 $nens`; do
+    m_id=`padzero $m 3`
+
+    watch_log $m_id/run.log successfully 5 $rundir
+
+    mv $m_id/${next_time:0:8}_${next_time:8:2}_mem$m_id.nc .
+    cp -L ${next_time:0:8}_${next_time:8:2}_mem$m_id.nc $nextdir/.
+done
+wait
+
+echo complete > stat
+

models/vort2d/module_forecast.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+. $config_file
+rundir=$work_dir/cycle/$time/vort2d
+if [[ ! -d $rundir ]]; then mkdir -p $rundir; echo waiting > $rundir/stat; fi
+
+cd $rundir
+if [[ `cat stat` == "complete" ]]; then exit; fi
+
+##check dependency
+if [[ $time == $time_start ]]; then wait_for_module ../icbc; fi
+
+echo running > stat
+
+echo "  Running forecast for vort2d model..."
+
+##load env if necessary
+src_file=$script_dir/../config/env/$host/vort2d.src
+if [[ -f $src_file ]]; then source $src_file; fi
+
+touch run.log
+
+$script_dir/job_submit.sh 1 1 0 python $script_dir/../models/vort2d/run.py $time >& run.log
+
+nextdir=$work_dir/cycle/$next_time/vort2d
+if [[ ! -d $nextdir ]]; then mkdir -p $nextdir; fi
+
+watch_log run.log successfully 5 $rundir
+
+cp -L ${next_time:0:8}_${next_time:8:2}.nc $nextdir/.
+
+echo complete > stat
+

models/vort2d/module_icbc.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+. $config_file
+
+if [[ $time -gt $time_start ]]; then exit; fi
+
+rundir=$work_dir/cycle/$time/icbc
+if [[ ! -d $rundir ]]; then mkdir -p $rundir; echo waiting > $rundir/stat; fi
+
+cd $rundir
+if [[ `cat stat` == "complete" ]]; then exit; fi
+
+echo running > stat
+
+echo "  Generating initial and boundary conditions..."
+
+##load env if necessary
+src_file=$script_dir/../config/env/$host/vort2d.src
+if [[ -f $src_file ]]; then source $src_file; fi
+
+$script_dir/job_submit.sh 1 1 0 python $script_dir/../models/vort2d/generate_ic.py $time >& icbc.log
+
+##check output
+icfile=${time:0:8}_${time:8:2}.nc
+watch_file $icfile 5 $rundir
+
+if [[ ! -d ../vort2d ]]; then mkdir -p ../vort2d; fi
+mv $icfile ../vort2d/.
+
+echo complete > stat
+

models/vort2d/module_perturb.sh

@@ -0,0 +1,35 @@
+#!/bin/bash
+. $config_file
+
+if [[ $time -gt $time_start ]]; then exit; fi
+
+rundir=$work_dir/cycle/$time/perturb
+if [[ ! -d $rundir ]]; then mkdir -p $rundir; echo waiting > $rundir/stat; fi
+
+cd $rundir
+if [[ `cat stat` == "complete" ]]; then exit; fi
+
+echo running > stat
+
+echo "  Generating perturbed ensemble members..."
+
+##load env if necessary
+src_file=$script_dir/../config/env/$host/vort2d.src
+if [[ -f $src_file ]]; then source $src_file; fi
+
+$script_dir/job_submit.sh 1 1 0 python $script_dir/../models/vort2d/perturb_ic.py $time $nens >& perturb.log
+
+if [[ ! -d ../vort2d ]]; then mkdir -p ../vort2d; fi
+
+##check output
+for m in `seq 1 $nens`; do
+    m_id=`padzero $m 3`
+
+    icfile=${time:0:8}_${time:8:2}_mem$m_id.nc
+    watch_file $icfile 1 $rundir
+
+    mv $icfile ../vort2d/.
+done
+
+echo complete > stat
+

models/vort2d/param.py

@@ -15,7 +15,7 @@
 restart_dt = 1
 
 ##model phys param
-gen = 5e-5    ##vorticity generation rate
+gen = 2e-5    ##vorticity generation rate
 diss = 3e3    ##dissipation rate
 
 ### vortex parameters

models/vort2d/perturb_ic.py

@@ -0,0 +1,14 @@
+import numpy as np
+import models.vort2d as model
+import config as c
+import sys
+from conversion import s2t
+
+time = s2t(sys.argv[1])
+nens = int(sys.argv[2])
+
+for m in range(nens):
+    state = model.initialize(c.grid, model.Vmax, model.Rmw, model.Vbg, model.Vslope, loc_sprd=model.loc_sprd)
+
+    model.write_var('./', c.grid, state, name='velocity', is_vector=True, time=time, member=m)
+

models/vort2d/run.py

@@ -0,0 +1,44 @@
+##run forecast using the vort2d model
+##this program is mimicking the run script of a real model, it is called by
+##top level control script run_cycle.sh to spawm several runs simultaneously
+##as if running ensemble forecasts
+
+import numpy as np
+import sys
+from datetime import datetime, timedelta
+import config as c
+import models.vort2d as model
+from conversion import s2t, t2s
+
+##current cycle time
+time = s2t(sys.argv[1])
+
+##member index, start from 0
+if len(sys.argv) > 2:
+    mem_id = int(sys.argv[2]) - 1
+else:
+    mem_id = None
+
+##time at next cycle, where current forecast will end
+next_time = time + timedelta(hours=1) * c.cycle_period
+
+path = './'
+
+##read the initial condition
+state = model.read_var(path, c.grid, name='velocity', is_vector=True, time=time, member=mem_id)
+
+t = time
+print('vort2d model forecast start at', t)
+while t < next_time:
+    t += timedelta(hours=1) * model.restart_dt
+
+    ##run the model
+    state = model.advance_time(state, model.dx, model.restart_dt, model.dt, model.gen, model.diss)
+
+    ##save restart file
+    model.write_var(path, c.grid, state, name='velocity', is_vector=True, time=t, member=mem_id)
+
+    print(t)
+
+print('vort2d model forecast finished successfully')
+