Skip to content

Commit

Permalink
fix for misapplication of perturbation delta in jacobian sensitivity …
Browse files Browse the repository at this point in the history
…calculation (#290)

When we calculate our jacobian matrix we do the following calculation:
sensitivity_xch4 = delta_xch4 / perturbation
where the perturbation is the fractional change from the original emissions eg. 50% perturbation translates to a perturbation = 0.5.

However, in HEMCO_Config.rc, when we apply the perturbation, a 50% perturbation is represented by a scale factor of 1.5.

Since the introduction of #226, we did not account for this difference between the representations of the perturbation in HEMCO vs in the jacobian calculation. This introduces an error in the jacobian calculation. Luckily, since the emission perturbations are relatively high (targeting 1e-08 kg/m2/s) the impact on results should be relatively small. For example, if an emission element was perturbed by 52x, the jacobian calculations goes from sensitivity_xch4 = delta_xch4 / 52 to sensitivity_xch4 = delta_xch4 / 51. This expectation is confirmed by my CONUS work for 2019 which goes from 41.075 --> 40.810 Tg/a. Absolute differences shown in the image below.

For OH, we are using a relatively small (10%) perturbation, so this bug matters a lot and would lead to very low sensitivity of OH to the observations in the jacobian.

This bugfix addresses this issue for both the emission elements and the OH elements.

Co-authored-by: eastjames <[email protected]>
  • Loading branch information
laestrada and eastjames authored Nov 21, 2024
1 parent 08ce441 commit 052c7b0
Show file tree
Hide file tree
Showing 3 changed files with 95 additions and 63 deletions.
98 changes: 62 additions & 36 deletions src/components/jacobian_component/make_perturbation_sf.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,8 @@

def update_jacobian_perturbation_files(jacobian_dir, state_vector_labels, flat_sf):
"""
Update the perturbation files in the jacobian run directories with the new scale factors.
Update the perturbation files in the jacobian run directories with the new
perturbation scale factors.
"""
contents = os.listdir(jacobian_dir)
jacobian_rundirs = [
Expand All @@ -27,9 +28,8 @@ def update_jacobian_perturbation_files(jacobian_dir, state_vector_labels, flat_s
]

# loop through each jacobian run directory and update the perturbation file
# with the new scale factors
# with the new perturbation SFs
for jacobian_rundir in jacobian_rundirs:
dir_suffix = jacobian_rundir.split("/")[-1].split("_")[-1]
perturbation_files = glob.glob(f"{jacobian_rundir}/Perturbations_*.txt")
for perturbation_file in perturbation_files:

Expand All @@ -43,7 +43,7 @@ def update_jacobian_perturbation_files(jacobian_dir, state_vector_labels, flat_s
sv_element = int(sv_label)
sv_idx = sv_element - 1

# make sure we only apply scale factors to emission elements
# make sure we only apply perturbation SFs to emission elements
if sv_element <= int(state_vector_labels.max().item()):
# add the right amount of padding
padding = "".ljust(4 - len(str(sv_element)))
Expand All @@ -65,44 +65,65 @@ def update_jacobian_perturbation_files(jacobian_dir, state_vector_labels, flat_s
file.writelines(lines)


def calculate_sfs(state_vector, emis_prior, target_emission=1e-8, prior_sf=None):
def calculate_perturbation_sfs(
state_vector, emis_prior, target_emission=1e-8, prior_sf=None
):
"""
Calculate the scale factors to perturb each state vector
element by based on the target_emission. Return a dictionary containing
two flat numpy arrays of perturbation scale factors indexed by state vector
element. The first array contains the perturbation scale factors to apply
to the state vector elements in jacobian simulations (based on the
original prior emissions). The second array contains the effective scale
factors (based on the nudged prior emissions for kalman mode) used in the
inversion to calculate the sensitivity of observations to the perturbation.
For a standalone inversion, the effective scale factors == jacobian scale factors.
Calculate the perturbation scale factors to perturb each state vector element based on the
target_emission.
Args:
state_vector [xr.Dataset] : the state vector dataset
emis_prior [xr.Dataset] : the prior emissions dataset
target_emission [float] : the target emission value to perturb each state vector element by
prior_sf [xr.Dataset] : the prior scale factor dataset (for kalman mode)
Returns:
dictionary containing two flat numpy arrays of perturbation scale factors
indexed by state vector element. The dictionary arrays are as follows:
jacobian_pert_sf [array]: contains the perturbation scale factors to apply
to the state vector elements in jacobian simulations
(based on the original prior emissions). To accomodate
how perturbations are applied in HEMCO these are relative
to 1, so a 50% perturbation is represented as 1.5.
effective_pert_sf [array]: contains the perturbation scalefactors (based on the
nudged prior emissions for kalman mode) used in the
inversion to calculate the sensitivity of observations
to the perturbation. These are relative to 0, so a 50%
perturbation is represented as 0.5. For a standalone
inversion, effective_pert_sf + 1 == jacobian_pert_sf.
target_emission [float]: the target emission value used to calculate the perturbation
"""
# create a sf dataset with the same structure as the state vector
sf = state_vector.copy()
sf = sf.rename({"StateVector": "ScaleFactor"})
# create a dataset with the same structure as the state vector for perturbation SFs
pert_sf = state_vector.copy()
pert_sf = pert_sf.rename({"StateVector": "ScaleFactor"})

# Calculate scale factors such that applying them to the original emissions
# will result in a target_emission kg/m2/s2 emission.
sf["ScaleFactor"] = target_emission / emis_prior["EmisCH4_Total_ExclSoilAbs"]
# Calculate perturbation SFs such that applying them to the original
# emissions will result in a target_emission kg/m2/s2 emission.
pert_sf["ScaleFactor"] = target_emission / emis_prior["EmisCH4_Total_ExclSoilAbs"]

# Extract state vector labels
state_vector_labels = state_vector["StateVector"]

# Add state vector labels to sf Dataset
sf = sf.assign(StateVector=state_vector_labels)
# Add state vector labels to pert sf Dataset
pert_sf = pert_sf.assign(StateVector=state_vector_labels)

# Use groupby to find the median scale factor for each state vector label
max_sf_da = sf["ScaleFactor"].groupby(sf["StateVector"]).median()
# Use groupby to find the median perturbation scale factor for each state vector label
max_pert_sf_da = pert_sf["ScaleFactor"].groupby(pert_sf["StateVector"]).median()

# get flat, numpy array by converting to dataframe and sorting based on
# state vector element
max_sf_df = max_sf_da.to_dataframe().reset_index()
max_sf_df = max_sf_df[max_sf_df["StateVector"] > 0].sort_values(by="StateVector")
jacobian_pert_sf = max_sf_df["ScaleFactor"].values
max_pert_sf_df = max_pert_sf_da.to_dataframe().reset_index()
max_pert_sf_df = max_pert_sf_df[max_pert_sf_df["StateVector"] > 0].sort_values(
by="StateVector"
)
jacobian_pert_sf = max_pert_sf_df["ScaleFactor"].values

# Replace any values greater than the threshold to avoid issues
# with reaching infinity
# Replace any NaN values to 1.0
# with reaching infinity. Replace any NaN values with 1.0
max_sf_threshold = 15000000.0
jacobian_pert_sf[jacobian_pert_sf > max_sf_threshold] = max_sf_threshold
jacobian_pert_sf = np.nan_to_num(jacobian_pert_sf, nan=1.0)
Expand All @@ -122,10 +143,13 @@ def calculate_sfs(state_vector, emis_prior, target_emission=1e-8, prior_sf=None)
else:
flat_prior_sf = np.ones(len(jacobian_pert_sf))

# calculate the effective scale factors
# calculate the effective perturbation SFs
effective_pert_sf = jacobian_pert_sf / flat_prior_sf

# return dictionary of scale factor arrays
# make the effective perturbations relative to 0, as expected in the inversion
effective_pert_sf = effective_pert_sf - 1.0

# return dictionary of perturbation scale factor arrays
perturbation_dict = {
"effective_pert_sf": effective_pert_sf,
"jacobian_pert_sf": jacobian_pert_sf,
Expand All @@ -138,7 +162,7 @@ def calculate_sfs(state_vector, emis_prior, target_emission=1e-8, prior_sf=None)
def make_perturbation_sf(config, period_number, perturb_value=1e-8):
"""
Calculate the perturbations for each state vector element and update the perturbation files.
Write out an archive of the flat scale factors for later use in sensitivity calculations.
Write out an archive of the flat perturbation scale factors for later use in sensitivity calculations.
Args:
config [dict] : dictionary of configuration settings
period_number [int] : the period number for which to calculate the perturbations
Expand Down Expand Up @@ -174,16 +198,18 @@ def make_perturbation_sf(config, period_number, perturb_value=1e-8):
# load the state vector dataset
state_vector = xr.load_dataset(os.path.join(base_directory, "StateVector.nc"))

# calculate the scale factors to perturb each state vector element by
perturbation_dict = calculate_sfs(state_vector, hemco_emis, perturb_value, prior_sf)
# calculate the perturbation scale factors we perturb each state vector element by
perturbation_dict = calculate_perturbation_sfs(
state_vector, hemco_emis, perturb_value, prior_sf
)

# update jacobian perturbation files with new scale factors
# 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 scale factor dictionary for later calculation of sensitivity
# 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(
Expand Down
26 changes: 13 additions & 13 deletions src/inversion_scripts/invert.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def do_inversion(
Returns
xhat [float] : Posterior scaling factors
ratio [float] : Change from prior [xhat = 1 + ratio]
delta_optimized [float] : Change from prior [xhat = 1 + delta_optimized]
KTinvSoK [float] : K^T*inv(S_o)*K [part of inversion equation]
KTinvSoyKxA [float] : K^T*inv(S_o)*(y-K*xA) [part of inversion equation]
S_post [float] : Posterior error covariance matrix
Expand Down Expand Up @@ -92,8 +92,8 @@ def do_inversion(
#
# In the code below this becomes
# xhat = xA + inv(gamma*KTinvSoK + inv(S_a)) * gamma*KTinvSoyKxA
# = xA + ratio
# = 1 + ratio [since xA=1 when optimizing scale factors]
# = xA + delta_optimized
# = 1 + delta_optimized [since xA=1 when optimizing scale factors]
#
# We build KTinvSoK and KTinvSoyKxA "piece by piece", loading one jacobian .pkl file at a
# time. This is so that we don't need to assemble or invert the full Jacobian matrix, which
Expand Down Expand Up @@ -253,15 +253,15 @@ def do_inversion(
inv_Sa = np.diag(1 / Sa_diag) # Inverse of prior error covariance matrix

# Solve for posterior scale factors xhat
ratio = np.linalg.inv(gamma * KTinvSoK + inv_Sa) @ (gamma * KTinvSoyKxA)
delta_optimized = np.linalg.inv(gamma * KTinvSoK + inv_Sa) @ (gamma * KTinvSoyKxA)

# Update scale factors by 1 to match what GEOS-Chem expects
# xhat = 1 + ratio
# xhat = 1 + delta_optimized
# Notes:
# - If optimizing BCs, the last 4 elements are in concentration space,
# so we do not need to add 1
# - If optimizing OH, the last element also needs to be updated by 1
xhat = ratio.copy()
xhat = delta_optimized.copy()
xhat[:scale_factor_idx] += 1
if oh_optimization:
if is_Regional:
Expand All @@ -277,10 +277,10 @@ def do_inversion(
# Averaging kernel matrix
A = np.identity(n_elements) - S_post @ inv_Sa

# Calculate J_A, where ratio = xhat - xA
# Calculate J_A, where delta_optimized = xhat - xA
# J_A = (xhat - xA)^T * inv_Sa * (xhat - xA)
ratioT = ratio.transpose()
J_A = ratioT @ inv_Sa @ ratio
delta_optimizedT = delta_optimized.transpose()
J_A = delta_optimizedT @ inv_Sa @ delta_optimized
J_A_normalized = J_A / n_elements

# Print some statistics
Expand All @@ -297,7 +297,7 @@ def do_inversion(
xhat[:scale_factor_idx].max(),
)

return xhat, ratio, KTinvSoK, KTinvSoyKxA, S_post, A
return xhat, delta_optimized, KTinvSoK, KTinvSoyKxA, S_post, A


if __name__ == "__main__":
Expand Down Expand Up @@ -341,7 +341,7 @@ def do_inversion(
is_Regional,
)
xhat = out[0]
ratio = out[1]
delta_optimized = out[1]
KTinvSoK = out[2]
KTinvSoyKxA = out[3]
S_post = out[4]
Expand All @@ -353,13 +353,13 @@ def do_inversion(
nvar2 = dataset.createDimension("nvar2", n_elements)
nc_KTinvSoK = dataset.createVariable("KTinvSoK", np.float32, ("nvar1", "nvar2"))
nc_KTinvSoyKxA = dataset.createVariable("KTinvSoyKxA", np.float32, ("nvar1"))
nc_ratio = dataset.createVariable("ratio", np.float32, ("nvar1"))
nc_delta_optimized = dataset.createVariable("ratio", np.float32, ("nvar1"))
nc_xhat = dataset.createVariable("xhat", np.float32, ("nvar1"))
nc_S_post = dataset.createVariable("S_post", np.float32, ("nvar1", "nvar2"))
nc_A = dataset.createVariable("A", np.float32, ("nvar1", "nvar2"))
nc_KTinvSoK[:, :] = KTinvSoK
nc_KTinvSoyKxA[:] = KTinvSoyKxA
nc_ratio[:] = ratio
nc_delta_optimized[:] = delta_optimized
nc_xhat[:] = xhat
nc_S_post[:, :] = S_post
nc_A[:, :] = A
Expand Down
34 changes: 20 additions & 14 deletions src/inversion_scripts/operators/TROPOMI_operator.py
Original file line number Diff line number Diff line change
Expand Up @@ -241,34 +241,40 @@ def apply_average_tropomi_operator(
is_bc[e] = is_BC_element
is_emis = ~np.equal(is_oh | is_bc, True)

# get perturbations and calculate sensitivities
perturbations = np.full(n_elements, 1.0, dtype=float)

# fill pert base array with values
# array contains 1 entry for each state vector element
# fill array with nans
base_xch4 = np.full(n_elements, np.nan)
# fill emission elements with the base value
base_xch4 = np.where(is_emis, emis_base_xch4, base_xch4)
# fill BC elements with the base value, which is same as emis value
base_xch4 = np.where(is_bc, emis_base_xch4, base_xch4)

# emissions perturbations
perturbations[0 : is_emis.sum()] = emis_perturbations

# OH perturbations
if config["OptimizeOH"]:
# fill OH elements with the OH base value
base_xch4 = np.where(is_oh, oh_base_xch4, base_xch4)

# compute BC perturbation for jacobian construction
oh_perturbation = float(config["PerturbValueOH"]) - 1.0

# get perturbations and calculate sensitivities
perturbations = np.full(n_elements, 1.0, dtype=float)
# update perturbations array to include OH perturbations
perturbations = np.where(is_oh, oh_perturbation, perturbations)

if config["OptimizeOH"]:
oh_perturbation = config["PerturbValueOH"]
else:
oh_perturbation = 1.0
# BC perturbations
if config["OptimizeBCs"]:
# fill BC elements with the base value, which is same as emis value
base_xch4 = np.where(is_bc, emis_base_xch4, base_xch4)

# compute BC perturbation for jacobian construction
bc_perturbation = config["PerturbValueBCs"]
else:
bc_perturbation = 1.0

# fill perturbation array with OH and BC perturbations
perturbations[0 : is_emis.sum()] = emis_perturbations
perturbations = np.where(is_oh, oh_perturbation, perturbations)
perturbations = np.where(is_bc, bc_perturbation, perturbations)
# update perturbations array to include OH perturbations
perturbations = np.where(is_bc, bc_perturbation, perturbations)

# calculate difference
delta_xch4 = pert_jacobian_xch4 - base_xch4
Expand Down

0 comments on commit 052c7b0

Please sign in to comment.