feat: updating prior class and example for dgf

dtu-qmcm · Sep 20, 2024 · 050788a · 050788a
1 parent 3250374
commit 050788a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/scripts/mcmc_demo.py b/scripts/mcmc_demo.py
@@ -38,7 +38,10 @@ def main():
         log_kcat=ind_prior_from_truth(true_parameters.log_kcat, 0.1),
         log_enzyme=ind_prior_from_truth(true_parameters.log_enzyme, 0.1),
         log_drain=ind_prior_from_truth(true_parameters.log_drain, 0.1),
-        dgf=ind_prior_from_truth(true_parameters.dgf, 0.1),
+        dgf=(
+            ind_prior_from_truth(true_parameters.dgf, 0.1)[0], 
+            jnp.diag(ind_prior_from_truth(true_parameters.dgf, 0.1)[1])
+        ),
         log_km=ind_prior_from_truth(true_parameters.log_km, 0.1),
         log_conc_unbalanced=ind_prior_from_truth(
             true_parameters.log_conc_unbalanced, 0.1

diff --git a/src/enzax/mcmc.py b/src/enzax/mcmc.py
@@ -40,7 +40,7 @@ class AllostericMichaelisMentenPriorSet:
     log_kcat: Float[Array, "2 n_enzyme"]
     log_enzyme: Float[Array, "2 n_enzyme"]
     log_drain: Float[Array, "2 n_drain"]
-    dgf: Float[Array, "2 n_metabolite"]
+    dgf: tuple[Float[Array, "n_metabolite"], Float[Array, "n_metabolite n_metabolite"]]
     log_km: Float[Array, "2 n_km"]
     log_ki: Float[Array, "2 n_ki"]
     log_conc_unbalanced: Float[Array, "2 n_unbalanced"]
@@ -100,7 +100,7 @@ def posterior_logdensity_amm(
         ind_normal_prior_logdensity(parameters.log_kcat, prior.log_kcat)
         + ind_normal_prior_logdensity(parameters.log_enzyme, prior.log_enzyme)
         + ind_normal_prior_logdensity(parameters.log_drain, prior.log_drain)
-        + ind_normal_prior_logdensity(parameters.dgf, prior.dgf)
+        + mv_normal_prior_logdensity(parameters.dgf, prior.dgf)
         + ind_normal_prior_logdensity(parameters.log_km, prior.log_km)
         + ind_normal_prior_logdensity(
             parameters.log_conc_unbalanced, prior.log_conc_unbalanced