From eeac60223d44e4472820443e9ca357aab4cd12d0 Mon Sep 17 00:00:00 2001
From: James East <47705516+eastjames@users.noreply.github.com>
Date: Sat, 23 Mar 2024 17:10:32 -0400
Subject: [PATCH] Re-run capability for Jacobian simulations (#188)

This PR enables the capability to re-run Jacobian simulations that have failed without running all Jacobian simulations. It is enabled with a new config switch ReDoJacobian. If no Jacobian simulations have been run, and DoJacobian is true, then IMI runs the Jacobians like normal, submitting all Jacobian simulations, regardless of the value of ReDoJacobian. If one or more Jacobian simulations failed in a previous run, and both ReDoJacobian and DoJacobian are true, then only the failed Jacobians will re-run. If DoJacobian is false, then setting ReDoJacobian to true will have no effect.
---
 config.yml                                    |  1 +
 envs/Harvard-Cannon/config.harvard-cannon.yml |  1 +
 src/components/jacobian_component/jacobian.sh | 19 ++++++++--
 .../run_jacobian_simulations.sh               | 38 ++++++++++++++++---
 4 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/config.yml b/config.yml
index 2a26c493..91816276 100644
--- a/config.yml
+++ b/config.yml
@@ -102,6 +102,7 @@ SetupPosteriorRun: false
 RunSetup: true
 DoSpinup: false
 DoJacobian: false
+ReDoJacobian: false
 DoInversion: false
 DoPosterior: false
 
diff --git a/envs/Harvard-Cannon/config.harvard-cannon.yml b/envs/Harvard-Cannon/config.harvard-cannon.yml
index 6a4f2ec2..b4634811 100644
--- a/envs/Harvard-Cannon/config.harvard-cannon.yml
+++ b/envs/Harvard-Cannon/config.harvard-cannon.yml
@@ -102,6 +102,7 @@ SetupPosteriorRun: false
 RunSetup: true
 DoSpinup: false
 DoJacobian: false
+ReDoJacobian: false
 DoInversion: false
 DoPosterior: false
 
diff --git a/src/components/jacobian_component/jacobian.sh b/src/components/jacobian_component/jacobian.sh
index 905321a3..d7025623 100644
--- a/src/components/jacobian_component/jacobian.sh
+++ b/src/components/jacobian_component/jacobian.sh
@@ -22,10 +22,6 @@ setup_jacobian() {
     # Create directory that will contain all Jacobian run directories
     mkdir -p -v jacobian_runs
 
-    # Copy run scripts
-    cp ${InversionPath}/src/geoschem_run_scripts/run_jacobian_simulations.sh jacobian_runs/
-    sed -i -e "s:{RunName}:${RunName}:g" \
-           -e "s:{InversionPath}:${InversionPath}:g" jacobian_runs/run_jacobian_simulations.sh
     cp ${InversionPath}/src/geoschem_run_scripts/submit_jacobian_simulations_array.sh jacobian_runs/
     sed -i -e "s:{START}:0:g" \
            -e "s:{END}:${nElements}:g" \
@@ -180,6 +176,21 @@ setup_jacobian() {
 # Usage:
 #   run_jacobian
 run_jacobian() {
+
+    pushd ${RunDirs}
+
+    # Copy run scripts
+    # need to re-copy since config vars are
+    # hardcoded and redojacobian might have changed
+    cp ${InversionPath}/src/geoschem_run_scripts/run_jacobian_simulations.sh jacobian_runs/
+    sed -i -e "s:{RunName}:${RunName}:g" \
+           -e "s:{InversionPath}:${InversionPath}:g" \
+           -e "s:{KalmanMode}:${KalmanMode}:g" \
+           -e "s:{EndDate}:${EndDate}:g" \
+           -e "s:{ReDoJacobian}:${ReDoJacobian}:g" jacobian_runs/run_jacobian_simulations.sh
+
+    popd
+
     if ! "$PrecomputedJacobian"; then
         jacobian_start=$(date +%s)
         printf "\n=== SUBMITTING JACOBIAN SIMULATIONS ===\n"
diff --git a/src/geoschem_run_scripts/run_jacobian_simulations.sh b/src/geoschem_run_scripts/run_jacobian_simulations.sh
index 4f56ee12..949f6cff 100755
--- a/src/geoschem_run_scripts/run_jacobian_simulations.sh
+++ b/src/geoschem_run_scripts/run_jacobian_simulations.sh
@@ -30,12 +30,40 @@ if test -f "$FILE"; then
     exit 1
 fi
 
-### Run GEOS-Chem in the directory corresponding to the cluster Id
-cd  ${RUNDIR}/{RunName}_${xstr}
-./{RunName}_${xstr}.run
+if {ReDoJacobian}; then
 
-# save the exit code of the jacobian simulation cmd
-retVal=$?
+    cd ${RUNDIR}/{RunName}_${xstr}
+
+    # check for last conc file
+    # it has only 1 timestep, so we
+    # only need check its existence
+    # rerun if it is not there
+    # get last date from geoschem_config file
+    date_str=$(grep end_date geoschem_config.yml)
+    yyyymmdd=$(python -c "import re; date_str='${date_str}'; print(re.split('\[|,', date_str)[1])")
+    LastConcFile=$(date -d ${yyyymmdd} +GEOSChem.SpeciesConc.%Y%m%d_0000z.nc4)
+        
+    cd OutputDir
+
+    if test -f "$LastConcFile"; then
+        echo "Not re-running jacobian simulation: ${xstr}" >> $output_log_file
+        exit 0
+    else
+        ### Run GEOS-Chem in the directory corresponding to the cluster Id
+        echo "Re-running jacobian simulation: ${xstr}" >> $output_log_file
+        cd ..
+        ./{RunName}_${xstr}.run
+        # save the exit code of the jacobian simulation cmd
+        retVal=$?
+    fi
+else
+    ### Run GEOS-Chem in the directory corresponding to the cluster Id
+    cd  ${RUNDIR}/{RunName}_${xstr}
+    ./{RunName}_${xstr}.run
+    
+    # save the exit code of the jacobian simulation cmd
+    retVal=$?
+fi
 
 # Check whether the jacobian finished successfully. If not, write to a hidden file. 
 # The presence of the .error_status_file.txt indicates whether an error ocurred.