Extract the combined k-effective estimator and fix its two-estimate branch

[GuySten]

Description

openmc_get_keff() computes the minimum variance estimate of k-effective from
the collision, absorption and tracklength estimators. The calculation is
inlined in a function that also reads global tallies, builds the covariance and
handles the low-realization case, which makes it hard to test on its own.

This PR moves it into combine_estimates() in src/math_functions.cpp: a pure
function combining three correlated estimates of any quantity, with no
reference to k-effective or to the individual estimators. openmc_get_keff()
supplies the k-specific meaning.

The covariance is taken as a StaticTensor2D<double, 3, 3> rather than a
dynamically shaped Tensor, so its shape is checked at compile time and the
call cannot be made with a mismatched matrix. It is also stack allocated, which
matters for callers that combine more often than once per run.

The n <= 3 condition moves in too, since it is a property of the expressions
(an n-3 term in one denominator, n-2 in the other). The function reports
whether a combination was formed; the caller supplies its own estimate when it
was not. What stays in openmc_get_keff() is that random ray
has only one estimate, that the fallback is the average over generations, and
that a single realization reports an infinite standard deviation.

The fix

The two-estimate branch, used when two estimators coincide, has an error in its
standard deviation. Urbatsch's Eq. 40 in LA-12658-MS is written in terms of the
matrix S, while the code substitutes the sample covariance
Sigma = S / (n - 1). The mean is unaffected; the standard deviation is
understated by up to sqrt(n - 1).

-                    (g + n * f * f) / (n * (n - 2) * g * g);
+                    ((n - 1) * g + n * f * f) / (n * (n - 2) * g * g);

Found and diagnosed by @nuclearkevin in #4016, which was closed as "next to
impossible to test in continuous-energy transport without delta tracking" and
folded into #3971. Extracting the calculation removes that obstacle, which is
why the two changes are together here.

The branch is reached only in multi-group mode with survival biasing, where the
collision and absorption estimators are identical, so
tests/regression_tests/mg_survival_biasing is regolded. No other results
change.

Validation

Three correlated estimators of a known value (pi) are simulated. Each trial
draws n realizations, forms the sample means and covariance as
accumulate_tallies() and openmc_get_keff() do, and combines them. Over
40,000 trials the mean reported standard deviation is compared with the actual
spread of the combined estimate:

three-estimate branch                two-estimate branch
  n     rho    reported/empirical      n     corrected   develop
  10    0.90   0.9525                  10    0.9643      0.4286
  20    0.90   0.9871                  20    0.9869      0.3039
  100   0.90   1.0032                  100   0.9936      0.1352
  1000  0.90   1.0033                  1000  1.0018      0.0430
  100   0.00   0.9938
  100   0.50   0.9879
  100   0.99   1.0019

develop's two-estimate standard deviation is too small by a factor growing with
n -- 23x at n=1000 -- matching the two to three orders of magnitude reported in
#4016. The corrected expression tracks the truth to within a percent. Both
branches are unbiased in the mean.

Both branches read a few percent low at n=10. That is inherent to the
derivation: the variance is a finite-sample approximation for weights estimated
from the same realizations, leaving a residual of order 1/n (-4.5%, -1.5%,
-0.7%, -0.4% at n = 10, 20, 50, 100). About half the n=10 gap is also a
measurement artifact -- the estimate is mildly heavy-tailed there, and against
an interquartile spread the ratio is 0.978 rather than 0.954. Nothing is worth
doing about it: a standard deviation from n realizations carries an inherent
uncertainty of 1/sqrt(2(n-1)), 23.6% at n=10, so the systematic error sits far
inside the noise.

Validation script

"""Monte Carlo validation of OpenMC's combined k-effective estimator.
 
Three correlated estimators of a known value are simulated. Each trial draws n
realizations, forms the sample means and Bessel-corrected sample covariance the
way accumulate_tallies() and openmc_get_keff() do, and runs the combination.
Across many trials we check that the reported standard deviation matches the
actual spread of the combined estimate.
 
Run with:  python validate_combine.py
"""
import numpy as np
 
TRUE_VALUE = np.pi
FP_REL_PRECISION = 1e-14
 
 
def combine_estimates(kv, cov, n, develop=False):
    """Transcription of combine_estimates() from src/math_functions.cpp.
 
    With develop=True the two-estimate standard deviation is evaluated as on
    develop, using the sample covariance Sigma where Urbatsch's Eq. 40 calls
    for S = (n - 1) * Sigma.
    """
    pairs = ((0, 1, 2), (0, 2, 1), (1, 2, 1))
    for a, b, other in pairs:
        if (abs(kv[a] - kv[b]) / kv[a] < FP_REL_PRECISION and
                abs(cov[a, a] - cov[b, b]) / cov[a, a] < FP_REL_PRECISION):
            i, j = (0, 2) if (a, b) == (0, 1) else (0, 1)
            break
    else:
        i = j = None
 
    if i is None:
        # Three estimates, Urbatsch Eq. 33 and Eq. 35
        g, combined, S = 0.0, 0.0, np.zeros(3)
        for l, (i, j, k) in enumerate(((0, 1, 2), (1, 2, 0), (2, 0, 1))):
            f = (cov[j, j] * (cov[k, k] - cov[i, k]) - cov[k, k] * cov[i, j] +
                 cov[j, k] * (cov[i, j] + cov[i, k] - cov[j, k]))
            S[0] += f * cov[0, l]
            S[1] += (cov[j, j] + cov[k, k] - 2.0 * cov[j, k]) * kv[l] * kv[l]
            S[2] += ((cov[k, k] + cov[i, j] - cov[j, k] - cov[i, k]) *
                     kv[l] * kv[j])
            combined += f * kv[l]
            g += f
        S *= (n - 1)
        S[0] *= (n - 1) ** 2
        combined /= g
        g *= (n - 1) ** 2
        var = S[0] / (g * n * (n - 3)) * (1 + n * ((S[1] - 2 * S[2]) / g))
        return combined, np.sqrt(var), 3
 
    # Two estimates, Urbatsch Eq. 36 and Eq. 40
    f = kv[i] - kv[j]
    g = cov[i, i] + cov[j, j] - 2.0 * cov[i, j]
    combined = kv[i] - (cov[i, i] - cov[i, j]) / g * f
    term = (g + n * f * f) if develop else ((n - 1) * g + n * f * f)
    var = ((cov[i, i] * cov[j, j] - cov[i, j] * cov[i, j]) * term /
           (n * (n - 2) * g * g))
    return combined, np.sqrt(var), 2
 
 
def simulate(n, trials, rho=0.90, sigma=(0.30, 0.42, 0.21),
             coincident=False, develop=False, seed=0):
    """Draw `trials` independent runs of `n` realizations and combine each."""
    rng = np.random.default_rng(seed)
    sigma = np.asarray(sigma, float)
    cov_true = np.outer(sigma, sigma) * (rho + (1.0 - rho) * np.eye(3))
 
    if coincident:
        # Estimators 0 and 1 are the same random variable, as collision and
        # absorption are in multi-group mode with survival biasing. The 3x3
        # covariance is singular, so two variables are drawn and duplicated.
        chol = np.linalg.cholesky(cov_true[np.ix_([0, 2], [0, 2])])
    else:
        chol = np.linalg.cholesky(cov_true)
 
    means = np.empty(trials)
    reported = np.empty(trials)
    for t in range(trials):
        if coincident:
            y = TRUE_VALUE + rng.standard_normal((n, 2)) @ chol.T
            x = np.column_stack([y[:, 0], y[:, 0], y[:, 1]])
        else:
            x = TRUE_VALUE + rng.standard_normal((n, 3)) @ chol.T
        kv = x.mean(axis=0)
        cov = np.cov(x, rowvar=False, ddof=1)
        if coincident:
            cov[1] = cov[0]
            cov[:, 1] = cov[:, 0]
        means[t], reported[t], _ = combine_estimates(kv, cov, n, develop)
 
    empirical = means.std(ddof=1)
    return dict(
        bias=means.mean() - TRUE_VALUE,
        bias_err=empirical / np.sqrt(trials),
        empirical=empirical,
        reported=reported.mean(),
        ratio=reported.mean() / empirical,
    )
 
 
if __name__ == '__main__':
    TRIALS = 40000
 
    print("Three-estimate branch")
    print("  n      rho    bias (sigma)   reported/empirical")
    for n in (10, 20, 100, 1000):
        r = simulate(n, TRIALS, seed=n)
        print("  %-6d %-6.2f %5.1f          %.4f"
              % (n, 0.90, abs(r['bias']) / r['bias_err'], r['ratio']))
    for rho in (0.0, 0.5, 0.99):
        r = simulate(100, TRIALS, rho=rho, seed=int(100 * rho) + 7)
        print("  %-6d %-6.2f %5.1f          %.4f"
              % (100, rho, abs(r['bias']) / r['bias_err'], r['ratio']))
 
    print()
    print("Two-estimate branch (estimators 0 and 1 identical)")
    print("  n      corrected      develop")
    for n in (10, 20, 100, 1000):
        fixed = simulate(n, TRIALS, coincident=True, seed=n + 1)
        old = simulate(n, TRIALS, coincident=True, develop=True, seed=n + 1)
        print("  %-6d %-14.4f %.4f" % (n, fixed['ratio'], old['ratio']))

Testing

tests/cpp_unit_tests/test_combine_estimates.cpp tests the combination
directly rather than through a transport calculation:

• Estimates that all agree combine to exactly that value, and shifting all
  three shifts the result identically -- both follow from the weights summing
  to one.
• The two-estimate result is compared against Eq. 36 and Eq. 40 transcribed
  independently in terms of S. This pins down the standard deviation:
  agrees to 2e-15 after the fix, off by 90% before it.
• The standard deviation falls as 1/sqrt(n) with the covariance held fixed:
  2.008 between n = 100 and n = 400 after the fix, 3.85 before.
• A combination is refused at n = 0, 1, 2, 3 and formed at n = 4.
  The middle two fail on develop.

Checklist

• I have performed a self-review of my own code
• I have run clang-format (version 18) on any C++ source files (if applicable)
• I have followed the style guidelines for Python source files (if applicable)
• I have made corresponding changes to the documentation (if applicable)
• I have added tests that prove my fix is effective or that my feature works (if applicable)

[nuclearkevin]

Thanks for this fix @GuySten! I'd like to test it in the delta tracking branch before approving - I'll try to get to that soon (currently rather busy with M&C papers).