From 571682c1e2843e09a4773e1ebb04db6d40647fdf Mon Sep 17 00:00:00 2001
From: Mark Nestor Costantini <mnc33@cam.ac.uk>
Date: Tue, 20 Jun 2023 17:27:57 +0100
Subject: [PATCH] .

---
 .../validphys/closuretest/multiclosure_pdf.py | 27 +++++++++++++++
 .../closuretest/multiclosure_pdf_output.py    | 33 +++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/validphys2/src/validphys/closuretest/multiclosure_pdf.py b/validphys2/src/validphys/closuretest/multiclosure_pdf.py
index fda9da3dd6..aa86047597 100644
--- a/validphys2/src/validphys/closuretest/multiclosure_pdf.py
+++ b/validphys2/src/validphys/closuretest/multiclosure_pdf.py
@@ -354,6 +354,33 @@ def fits_pdf_total_ratio(
     variance = np.mean(calc_chi2(sqrtcov, rep_diff.transpose(2, 1, 0)), axis=0)
     return np.mean(bias) / np.mean(variance)
 
+def fits_pdf_total_bias_variance(
+    fits_central_difference,
+    fits_replica_difference,
+    fits_covariance_matrix_totalpdf,
+    multiclosure_nx=4,
+):
+    """Calculate the total bias and variance for all flavours and x allowing for
+    correlations across flavour.
+
+    Returns
+    -------
+
+    """
+    central_diff = np.asarray(fits_central_difference).reshape(
+        -1, multiclosure_nx * len(XI_FLAVOURS)
+    )
+    rep_diff = np.asarray(fits_replica_difference).reshape(
+        len(fits_replica_difference), -1, multiclosure_nx * len(XI_FLAVOURS)
+    )
+
+    sqrtcov = la.cholesky(fits_covariance_matrix_totalpdf, lower=True)
+
+    bias = calc_chi2(sqrtcov, central_diff.T)
+    # need flav x on first axis
+    variance = calc_chi2(sqrtcov, rep_diff.transpose(2, 1, 0))
+    return bias, variance
+
 
 fits_xi_grid_values = collect("xi_grid_values", ("fits", "fitpdf"))
 
diff --git a/validphys2/src/validphys/closuretest/multiclosure_pdf_output.py b/validphys2/src/validphys/closuretest/multiclosure_pdf_output.py
index 60df7ba295..ae5144dd30 100644
--- a/validphys2/src/validphys/closuretest/multiclosure_pdf_output.py
+++ b/validphys2/src/validphys/closuretest/multiclosure_pdf_output.py
@@ -61,6 +61,39 @@ def plot_bias_distribution_flavour(fits_pdf_flavour_bias_variance):
         yield fig
 
 
+@figure
+def plot_variance_distribution(fits_pdf_total_bias_variance):
+    """
+    plot the distribution of variance (in PDF space).
+    The PDF covariance matrix includes flavours correlations.
+    """
+    _, var = fits_pdf_total_bias_variance
+    
+    
+    fig, ax = plotutils.subplots()
+    # plot distribution of all fits
+    var = np.concatenate(var, axis=0)
+    ax.hist(var, density=True, label = f"Variance Distribution, mean = {np.mean(var):.3f}")
+    
+    ax.legend()
+    return fig
+
+
+@figure
+def plot_bias_distribution(fits_pdf_total_bias_variance):
+    """
+    plot the distribution of bias (in PDF space).
+    The PDF covariance matrix includes flavours correlations.
+    """
+    bias, _ = fits_pdf_total_bias_variance
+    
+    fig, ax = plotutils.subplots()
+
+    ax.hist(bias, density=True, label = f"Bias Distribution, mean = {np.mean(bias):.3f}")
+    ax.legend()
+    return fig
+
+
 @table
 def xi_flavour_table(xi_flavour_x, xi_totalpdf):
     """For each flavour take the mean of xi_flavour_x across x to get a single