To use gridded state vector and perturbation scaling factor in HEMCO for IMI

src/components/jacobian_component/jacobian.sh

@@ -21,9 +21,19 @@ setup_jacobian() {
 
     # make dir for jacobian ics/bcs
     mkdir -p jacobian_1ppb_ics_bcs/Restarts
-    mkdir -p jacobian_1ppb_ics_bcs/BCs
-    OrigBCFile=${fullBCpath}/GEOSChem.BoundaryConditions.${StartDate}_0000z.nc4
-    python ${InversionPath}/src/components/jacobian_component/make_jacobian_icbc.py $OrigBCFile ${RunDirs}/jacobian_1ppb_ics_bcs/BCs $StartDate
+    if "$isRegional"; then
+        mkdir -p jacobian_1ppb_ics_bcs/BCs
+        OrigBCFile="${fullBCpath}/GEOSChem.BoundaryConditions.${StartDate}_0000z.nc4"
+        python ${InversionPath}/src/components/jacobian_component/make_jacobian_icbc.py $OrigBCFile ${RunDirs}/jacobian_1ppb_ics_bcs/BCs $StartDate
+    fi
+    OrigRestartFile="${RestartFilePrefix}${StartDate}_0000z.nc4"
+    python ${InversionPath}/src/components/jacobian_component/make_jacobian_icbc.py $OrigRestartFile ${RunDirs}/jacobian_1ppb_ics_bcs/Restarts $StartDate
+    cd ${RunDirs}/jacobian_1ppb_ics_bcs/Restarts/
+    if [ -f GEOSChem.BoundaryConditions.1ppb.${StartDate}_0000z.nc4 ]; then
+        mv GEOSChem.BoundaryConditions.1ppb.${StartDate}_0000z.nc4 GEOSChem.Restart.1ppb.${StartDate}_0000z.nc4
+        ncrename -v SpeciesBC_CH4,SpeciesRst_CH4 GEOSChem.Restart.1ppb.${StartDate}_0000z.nc4
+    fi
+    cd ${RunDirs}
 
     # Create directory that will contain all Jacobian run directories
     mkdir -p -v jacobian_runs
@@ -57,12 +67,18 @@ setup_jacobian() {
     else
         sed -i -e "s:{JOBS}::g" jacobian_runs/submit_jacobian_simulations_array.sh
     fi
-    cp ${InversionPath}/src/geoschem_run_scripts/run_prior_simulation.sh jacobian_runs/
-    sed -i -e "s:{RunName}:${RunName}:g" \
-        -e "s:{InversionPath}:${InversionPath}:g" jacobian_runs/run_prior_simulation.sh
-    cp ${InversionPath}/src/geoschem_run_scripts/run_bkgd_simulation.sh jacobian_runs/
-    sed -i -e "s:{RunName}:${RunName}:g" \
-        -e "s:{InversionPath}:${InversionPath}:g" jacobian_runs/run_bkgd_simulation.sh
+
+    if "$KalmanMode"; then
+        jacobian_period=${period_i}
+    else
+        jacobian_period=1
+    fi
+
+    set -e
+    # generate gridded perturbation values for all state vector elements
+    printf "\n=== GENERATE GRIDDED PERTURBATION SFs ===\n"
+    python ${InversionPath}/src/components/jacobian_component/make_perturbation_sf.py $ConfigPath $jacobian_period $PerturbValue
+    printf "\n=== DONE GENERATE GRIDDED PERTURBATION SFs ===\n"
 
     # Initialize (x=0 is base run, i.e. no perturbation; x=1 is state vector element=1; etc.)
     x=0
@@ -114,7 +130,7 @@ create_simulation_dir() {
     cd $runDir
 
     # Link to GEOS-Chem executable instead of having a copy in each rundir
-    ln -s ../../GEOSChem_build/gcclassic .
+    ln -nsf ../../GEOSChem_build/gcclassic .
 
     # link to restart file
     RestartFileFromSpinup=${RunDirs}/spinup_run/Restarts/GEOSChem.Restart.${SpinupEnd}_0000z.nc4
@@ -126,7 +142,7 @@ create_simulation_dir() {
             sed -i -e "s|SpeciesRst|SpeciesBC|g" HEMCO_Config.rc
         fi
     fi
-    ln -s $RestartFile Restarts/GEOSChem.Restart.${StartDate}_0000z.nc4
+    ln -nsf $RestartFile Restarts/GEOSChem.Restart.${StartDate}_0000z.nc4
 
     # Modify HEMCO_Config.rc to turn off individual emission inventories
     # and use total emissions (without soil absorption) saved out from prior
@@ -235,24 +251,20 @@ create_simulation_dir() {
                 sed -i -e "s| OH_pert_factor 1.0| OH_pert_factor ${PerturbValueOH}|g" HEMCO_Config.rc
             else
                 # Perturb OH by hemisphere if this is a global simulation
-                # Add and edit perturbations txt file
-                cp Perturbations.txt PerturbationsOH.txt
-                sed -i -e "s|CH4_STATE_VECTOR|HEMIS_MASK|g" PerturbationsOH.txt
+                # Apply hemispheric OH perturbation values using mask file
+                Output_fpath="./gridded_perturbation_oh_scale.nc"
+                Hemis_mask_fpath="${DataPath}/HEMCO/MASKS/v2024-08/hemisphere_mask.01x01.nc"
+                OptimizeNorth='False'
+                OptimizeSouth='False'
                 if [ $start_element -eq $((ohThreshold + 1)) ]; then
-                    OHPertNewLine="N_HEMIS    1     ${PerturbValueOH}"
+                    OptimizeNorth='True'
                 else
-                    OHPertNewLine="S_HEMIS    2     ${PerturbValueOH}"
+                    OptimizeSouth='True'
                 fi
-                OHPertPrevLine='DEFAULT    0     1.0'
-                sed -i "/$OHPertPrevLine/a $OHPertNewLine" PerturbationsOH.txt
+                gridded_optimized_OH $PerturbValueOH $PerturbValueOH $Hemis_mask_fpath $Output_fpath $OptimizeNorth $OptimizeSouth
 
                 # Modify OH scale factor in HEMCO config
-                sed -i -e "s| OH_pert_factor  1.0 - - - xy 1 1| OH_pert_factor PerturbationsOH.txt - - - xy 1 1|g" HEMCO_Config.rc
-
-                HcoPrevLineMask='CH4_STATE_VECTOR'
-                HcoNextLineMask='* HEMIS_MASK $ROOT\/MASKS\/v2024-08\/hemisphere_mask.01x01.nc Hemisphere 2000\/1\/1\/0 C xy 1 * - 1 1 
-'
-                sed -i "/${HcoPrevLineMask}/a ${HcoNextLineMask}" HEMCO_Config.rc
+                sed -i -e "s| OH_pert_factor  1.0 - - - xy 1 1| OH_pert_factor ${Output_fpath} oh_scale 2000\/1\/1\/0 C xy 1 1|g" HEMCO_Config.rc
             fi
         fi
 
@@ -289,29 +301,32 @@ create_simulation_dir() {
     else
         # set 1ppb CH4 boundary conditions and restarts for all other perturbation simulations
         # Note that we use the timecycle flag C to avoid having to make additional files
-        RestartFile=${RunDirs}/jacobian_1ppb_ics_bcs/Restarts/GEOSChem.Restart.1ppb.${StartDate}_0000z.nc4
-        BCFile1ppb=${RunDirs}/jacobian_1ppb_ics_bcs/BCs/GEOSChem.BoundaryConditions.1ppb.${StartDate}_0000z.nc4
-        BCSettings1ppb="SpeciesBC_CH4  1980-2021/1-12/1-31/* C xyz 1 CH4 - 1 1"
-        sed -i -e "s|.*GEOSChem\.BoundaryConditions.*|\* BC_CH4 ${BCFile1ppb} ${BCSettings1ppb}|g" HEMCO_Config.rc
+        if "$isRegional"; then
+            BCFile1ppb=${RunDirs}/jacobian_1ppb_ics_bcs/BCs/GEOSChem.BoundaryConditions.1ppb.${StartDate}_0000z.nc4
+            BCSettings1ppb="SpeciesBC_CH4  1980-2021/1-12/1-31/* C xyz 1 CH4 - 1 1"
+            sed -i -e "s|.*GEOSChem\.BoundaryConditions.*|\* BC_CH4 ${BCFile1ppb} ${BCSettings1ppb}|g" HEMCO_Config.rc
+        fi
         # create symlink to 1ppb restart file
-        ln -sf $RestartFile Restarts/GEOSChem.Restart.${StartDate}_0000z.nc4
+        RestartFile1ppb=${RunDirs}/jacobian_1ppb_ics_bcs/Restarts/GEOSChem.Restart.1ppb.${StartDate}_0000z.nc4
+        RestartFile=$RestartFile1ppb
+        ln -nsf $RestartFile Restarts/GEOSChem.Restart.${StartDate}_0000z.nc4
         # Also, set emissions to zero for default CH4 tracer by applying ZERO scale factor (id 5)
         sed -i -e "s|CH4 - 1 500|CH4 5 1 500|g" HEMCO_Config.rc
     fi
 
     # Modify restart and BC entries in HEMCO_Config.rc to look for CH4 only
     # instead of all advected species
     sed -i -e "s/SPC_/SPC_CH4/g" -e "s/?ALL?/CH4/g" -e "s/EFYO xyz 1 \*/EFYO xyz 1 CH4/g" HEMCO_Config.rc
-    sed -i -e "s/BC_ /BC_CH4 /g" -e "s/?ADV?/CH4/g" -e "s/EFY xyz 1 \*/EFY xyz 1 CH4/g" HEMCO_Config.rc
-
+    if "$isRegional"; then
+        sed -i -e "s/BC_ /BC_CH4 /g" -e "s/?ADV?/CH4/g" -e "s/EFY xyz 1 \*/EFY xyz 1 CH4/g" HEMCO_Config.rc
+    fi
     # Initialize previous lines to search
     GcPrevLine='- CH4'
     HcoPrevLine1='EFYO xyz 1 CH4 - 1 '
     HcoPrevLine2='1 500'
-    HcoPrevLine3='Perturbations.txt - - - xy count 1'
+    HcoPrevLine3="#300N SCALE_ELEM_000N ${RunDirs}/StateVector.nc StateVector 2000/1/1/0 C xy 1 1 N"
     HcoPrevLine4='\* BC_CH4'
-    PertPrevLine='DEFAULT    0     0.0'
-
+
     # Loop over state vector element numbers for this run and add each element
     # as a CH4 tracer in the configuraton files
     if is_number "$x"; then
@@ -339,15 +354,9 @@ add_new_tracer() {
         istr="${i}"
     fi
 
-    # by default remove all emissions except for in the prior simulation
-    # and the OH perturbation simulation
-    if [ $x -gt 0 ]; then
-        sed -i -e "s/DEFAULT    0     1.0/$PertPrevLine/g" Perturbations.txt
-    fi
-
-    # Start HEMCO scale factor ID at 2000 to avoid conflicts with
+    # Start HEMCO scale factor ID at 3000 to avoid conflicts with
     # preexisting scale factors/masks
-    SFnum=$((2000 + i))
+    SFnum=$((3000 + i))
 
     # Add lines to geoschem_config.yml
     # Spacing in GcNewLine is intentional
@@ -361,33 +370,23 @@ add_new_tracer() {
     SpcNewLines='CH4_'$istr':\n  << : *CH4properties\n  Background_VV: 1.8e-6\n  FullName: Methane'
     sed -i -e "s|$SpcNextLine|$SpcNewLines\n$SpcNextLine|g" species_database.yml
 
-    # Add lines to HEMCO_Config.yml
-    HcoNewLine1='\
-* SPC_CH4_'$istr' - - - - - - CH4_'$istr' - 1 1'
+    # Add lines for new tracers to HEMCO_Config.rc
+    HcoNewLine2='0 CH4_Emis_Prior_'$istr' - - - - - - CH4_'$istr' '4/$SFnum' 1 500'
+    sed -i -e "\|$HcoPrevLine2|a $HcoNewLine2" HEMCO_Config.rc
+    HcoPrevLine2=$HcoNewLine2
+
+    HcoNewLine3="$SFnum SCALE_ELEM_$istr ${RunDirs}/StateVector.nc StateVector 2000/1/1/0 C xy 1 1 $i"
+    sed -i -e "\|$HcoPrevLine3|a $HcoNewLine3" HEMCO_Config.rc
+    HcoPrevLine3=$HcoNewLine3
+
+    # Add lines for restarts of new tracers to HEMCO_Config.rc
+    HcoNewLine1='* SPC_CH4_'$istr' - - - - - - CH4_'$istr' - 1 1'
     sed -i -e "/$HcoPrevLine1/a $HcoNewLine1" HEMCO_Config.rc
     HcoPrevLine1='SPC_CH4_'$istr
-
-    HcoNewLine2='\
-0 CH4_Emis_Prior_'$istr' - - - - - - CH4_'$istr' '$SFnum' 1 500'
-    sed -i "/$HcoPrevLine2/a $HcoNewLine2" HEMCO_Config.rc
-    HcoPrevLine2='CH4_'$istr' '$SFnum' 1 500'
-
-    HcoNewLine3='\
-'$SFnum' SCALE_ELEM_'$istr' Perturbations_'$istr'.txt - - - xy count 1'
-    sed -i "/$HcoPrevLine3/a $HcoNewLine3" HEMCO_Config.rc
-    HcoPrevLine3='SCALE_ELEM_'$istr' Perturbations_'$istr'.txt - - - xy count 1'
-
-    HcoNewLine4='\
-* BC_CH4_'$istr' - - - - - - CH4_'$istr' - 1 1'
-    sed -i -e "/$HcoPrevLine4/a $HcoNewLine4" HEMCO_Config.rc
-    HcoPrevLine4='BC_CH4_'$istr
-
-    # Add new Perturbations.txt and update for non prior runs
-    cp Perturbations.txt Perturbations_${istr}.txt
-    if [ $x -gt 0 ]; then
-        PertNewLine='\
-ELEM_'$istr'  '$i'     '0.0''
-        sed -i "/$PertPrevLine/a $PertNewLine" Perturbations_${istr}.txt
+    if "$isRegional"; then
+        HcoNewLine4='* BC_CH4_'$istr' - - - - - - CH4_'$istr' - 1 1'
+        sed -i -e "/$HcoPrevLine4/a $HcoNewLine4" HEMCO_Config.rc
+        HcoPrevLine4='BC_CH4_'$istr
     fi
 
 }
@@ -406,8 +405,9 @@ run_jacobian() {
     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
+        -e "s:{ReDoJacobian}:${ReDoJacobian}:g" \
+        -e "s:{StartDate}:${StartDate}:g" \
+        -e "s:{EndDate}:${EndDate}:g" jacobian_runs/run_jacobian_simulations.sh
 
     popd
 
@@ -417,37 +417,37 @@ run_jacobian() {
         jacobian_period=1
     fi
 
-    set -e
-    # update perturbation values before running jacobian simulations
-    printf "\n=== UPDATING PERTURBATION SFs ===\n"
-    python ${InversionPath}/src/components/jacobian_component/make_perturbation_sf.py $ConfigPath $jacobian_period $PerturbValue
-
     if ! "$PrecomputedJacobian"; then
 
         cd ${RunDirs}/jacobian_runs
         jacobian_start=$(date +%s)
 
-        # create 1ppb restart file
-        OrigRestartFile=$(readlink ${RunName}_0000/Restarts/GEOSChem.Restart.${StartDate}_0000z.nc4)
-        python ${InversionPath}/src/components/jacobian_component/make_jacobian_icbc.py $OrigRestartFile ${RunDirs}/jacobian_1ppb_ics_bcs/Restarts $StartDate
-        cd ${RunDirs}/jacobian_1ppb_ics_bcs/Restarts/
-        if [ -f GEOSChem.BoundaryConditions.1ppb.${StartDate}_0000z.nc4 ]; then
-            mv GEOSChem.BoundaryConditions.1ppb.${StartDate}_0000z.nc4 GEOSChem.Restart.1ppb.${StartDate}_0000z.nc4
-        fi
-        cd ${RunDirs}/jacobian_runs
         set +e
 
         printf "\n=== SUBMITTING JACOBIAN SIMULATIONS ===\n"
         # Submit job to job scheduler
         source submit_jacobian_simulations_array.sh
 
         if "$LognormalErrors"; then
+            # Submit background simulation to job scheduler
+            printf "\n=== SUBMITTING BACKGROUND SIMULATION ===\n"
+
+            cp ${InversionPath}/src/geoschem_run_scripts/run_bkgd_simulation.sh ./
+
+            sed -i -e "s:{RunName}:${RunName}:g" \
+                -e "s:{InversionPath}:${InversionPath}:g" \
+                -e "s:{StartDate}:${StartDate}:g" \
+                run_bkgd_simulation.sh
+
             sbatch --mem $RequestedMemory \
                 -c $RequestedCPUs \
+                -N 1 \
                 -t $RequestedTime \
                 -p $SchedulerPartition \
                 -W run_bkgd_simulation.sh
             wait
+
+            printf "\n=== DONE BACKGROUND SIMULATION ===\n"
         fi
 
         # check if any jacobians exited with non-zero exit code
@@ -468,16 +468,22 @@ run_jacobian() {
         fi
 
         precomputedJacobianCache=${precomputedJacobianCachePrefix}/data_converted
-        ln -s $precomputedJacobianCache data_converted_reference
+        ln -nsf $precomputedJacobianCache data_converted_reference
 
         # Run the prior simulation
         JacobianRunsDir=${RunDirs}/jacobian_runs
         cd ${JacobianRunsDir}
 
         # Submit prior simulation to job scheduler
         printf "\n=== SUBMITTING PRIOR SIMULATION ===\n"
+        cp ${InversionPath}/src/geoschem_run_scripts/run_prior_simulation.sh ./
+        sed -i -e "s:{RunName}:${RunName}:g" \
+            -e "s:{InversionPath}:${InversionPath}:g" \
+            -e "s:{StartDate}:${StartDate}:g" \
+            run_prior_simulation.sh
         sbatch --mem $RequestedMemory \
             -c $RequestedCPUs \
+            -N 1 \
             -t $RequestedTime \
             -o imi_output.tmp \
             -p $SchedulerPartition \
@@ -495,6 +501,7 @@ run_jacobian() {
             printf "\n=== SUBMITTING BACKGROUND SIMULATION ===\n"
             sbatch --mem $RequestedMemory \
                 -c $RequestedCPUs \
+                -N 1 \
                 -t $RequestedTime \
                 -p $SchedulerPartition \
                 -W run_bkgd_simulation.sh

src/components/jacobian_component/make_perturbation_sf.py

@@ -203,19 +203,55 @@ def make_perturbation_sf(config, period_number, perturb_value=1e-8):
         state_vector, hemco_emis, perturb_value, prior_sf
     )
 
-    # update jacobian perturbation files with new perturbation scale factors
-    # before we run the jacobian simulations
-    update_jacobian_perturbation_files(
-        jacobian_dir, state_vector["StateVector"], perturbation_dict["jacobian_pert_sf"]
-    )
-
     # archive npz file of perturbation scale factor dictionary for later calculation of sensitivity
     archive_dir = os.path.join(base_directory, "archive_perturbation_sfs")
     os.makedirs(archive_dir, exist_ok=True)
     np.savez(
         os.path.join(archive_dir, f"pert_sf_{period_number}.npz"), **perturbation_dict
     )
+    grid_pert_fpath = os.path.join(archive_dir, f"gridded_pert_scale_{period_number}.nc")
+    make_gridded_perturbation_sf_latlon(perturbation_dict["jacobian_pert_sf"], state_vector, grid_pert_fpath)
+
+def make_gridded_perturbation_sf_latlon(pert_vector, statevector, save_pth):
+    lat = statevector['lat'].values
+    lon = statevector['lon'].values
 
+    # Map the input vector (e.g., scale factors) to the state vector grid
+    if "time" in statevector.StateVector.dims:
+        sv_index = statevector.StateVector.values
+    else:
+        sv_index = statevector.StateVector.values[None,...]
+    valid_mask = ~np.isnan(sv_index)
+    # Create a placeholder array for integer indices, initialize with -1
+    sv_index_int = np.full_like(sv_index, fill_value=-1, dtype=int)
+    sv_index_int[valid_mask] = sv_index[valid_mask].astype(int)
+
+    gridded_pert = np.ones_like(sv_index, dtype=float)
+    gridded_pert[valid_mask] = pert_vector[sv_index_int[valid_mask] - 1]
+
+    refyear = 2000
+    gridded_pert = xr.DataArray(gridded_pert, dims=['time', 'lat', 'lon'], 
+                                coords=dict(time=(['time'], [0.]), lat=(['lat'], lat), lon=(['lon'], lon)),
+                                attrs=dict(units='1', long_name='scaling factor for state vector elements'))
+    gridded_pert_ds = xr.Dataset({'scale': gridded_pert})
+    # Add attribute metadata
+    gridded_pert_ds.lat.attrs["units"] = "degrees_north"
+    gridded_pert_ds.lat.attrs["long_name"] = "Latitude"
+    gridded_pert_ds.lon.attrs["units"] = "degrees_east"
+    gridded_pert_ds.lon.attrs["long_name"] = "Longitude"
+    gridded_pert_ds['time'].attrs = dict(units='days since {}-01-01 00:00:00'.format(refyear),
+                                        delta_t='0000-01-00 00:00:00', axis='T', standard_name='Time',
+                                        long_name='Time', calendar='standard')
+
+    # Save
+    if save_pth is not None:
+        print("Saving file {}".format(save_pth))
+        gridded_pert_ds.to_netcdf(
+            save_pth,
+            encoding={
+                v: {"zlib": True, "complevel": 1} for v in gridded_pert_ds.data_vars
+            },
+        )
 
 if __name__ == "__main__":
     config_path = sys.argv[1]

src/components/kalman_component/kalman.sh

@@ -118,7 +118,7 @@ run_period() {
     python ${InversionPath}/src/components/kalman_component/change_dates.py $StartDate_i $EndDate_i $PosteriorRunDir
     wait
     echo "Edited Start/End dates in geoschem_config.yml for prior/perturbed/posterior simulations: $StartDate_i to $EndDate_i"
-
+    
     # Prepare initial (prior) emission scale factors for the current period
     echo "python path = $PYTHONPATH"
     python ${InversionPath}/src/components/kalman_component/prepare_sf.py $ConfigPath $period_i ${RunDirs} $NudgeFactor
@@ -159,7 +159,8 @@ run_period() {
 
     # Make link to restart file from posterior run directory in prior, OH, and background simulation
     # and link to 1ppb restart file for perturbations
-    python ${InversionPath}/src/components/jacobian_component/make_jacobian_icbc.py ${PosteriorRunDir}/Restarts/GEOSChem.Restart.${EndDate_i}_0000z.nc4 ${RunDirs}/jacobian_1ppb_ics_bcs/Restarts $EndDate_i
+    org_restart=${PosteriorRunDir}/Restarts/GEOSChem.Restart.${EndDate_i}_0000z.nc4
+    python ${InversionPath}/src/components/jacobian_component/make_jacobian_icbc.py $org_restart ${RunDirs}/jacobian_1ppb_ics_bcs/Restarts $EndDate_i
     rundir_num=$(get_last_rundir_suffix $JacobianRunsDir)
     for ((idx = 0; idx <= rundir_num; idx++)); do
         # Add zeros to string name