Compute log evidence per chain for diagnosis

harmonic/evidence.py

@@ -70,13 +70,15 @@ def __init__(self, nchains: int, model, shift=Shifting.MEAN_SHIFT):
         self.nchains = nchains
         self.ndim = model.ndim
         self.evidence_inv = 0.0
+
         self.evidence_inv_var = 0.0
         self.evidence_inv_var_var = 0.0
         self.kurtosis = 0.0
         self.n_eff = 0
 
         # For statistics computed purely in log-space.
         self.ln_evidence_inv = 0.0
+        self.ln_evidence_inv_per_chain = None
         self.ln_evidence_inv_var = 0.0
         self.ln_evidence_inv_var_var = 0.0
         self.ln_kurtosis = 0.0
@@ -140,6 +142,10 @@ def process_run(self):
 
         evidence_inv /= nsamples
 
+        self.ln_evidence_inv_per_chain = (
+            jnp.log(self.running_sum) - self.shift_value - jnp.log(nsamples_per_chain)
+        )
+
         """
         The following code computes the exponents of the variance and variance 
         of the variance where possible in log space to avoid overflow errors.