BunchTwissAnalysis::computeBunchMoments doesn't calculate central moments correctly

[woodtp]

Under suspicion that BunchTwissAnalysis::computeBunchMoments is not working correctly, I produced some tests to compare the results of BunchTwissAnalysis to bunch statistics computed via numpy, the results of which are outlined below:

tests/py/orbit/bunch_utils/test_bunch_twiss_analysis.py:322: AssertionError
=================================================== short test summary info ===================================================
FAILED tests/py/orbit/bunch_utils/test_bunch_twiss_analysis.py::test_computeBunchMoments_order2_M02 - assert 0.0 == 3.19521707947...e-08 ± 1.0e-15
FAILED tests/py/orbit/bunch_utils/test_bunch_twiss_analysis.py::test_computeBunchMoments_order2_emitnorm_M02 - assert 0.0 == 1.0007118962927186 ± 1.0e-15
FAILED tests/py/orbit/bunch_utils/test_bunch_twiss_analysis.py::test_computeBunchMoments_order2_dispersion_flag - assert 9.331838635982362e-08 == 9.33183590281...e-08 ± 1.0e-15
FAILED tests/py/orbit/bunch_utils/test_bunch_twiss_analysis.py::test_computeBunchMoments_order3_M03 - assert 0.0 == 4.13235583192...e-13 ± 1.0e-15
FAILED tests/py/orbit/bunch_utils/test_bunch_twiss_analysis.py::test_computeBunchMoments_macrosize_order2_M00_M10_M01 - assert -7.968925040526781e-05 == -0.0002851747...7959 ± 1.0e-15
FAILED tests/py/orbit/bunch_utils/test_bunch_twiss_analysis.py::test_computeBunchMoments_macrosize_order2_M20 - assert 9.368853867319752e-08 == 8.91593940978...e-08 ± 1.0e-12
================================================ 6 failed, 24 passed in 0.16s =================================================

Moments Formula

The principle issue is that, for a 2D distribution, the central moments of overall order $j+k$, weighted by the size of the $i^{th}$ macroparticle $w_i$, are given by:

$$\mu_{xy}^{jk} = \frac{1}{N}\sum_i w_i(x_i-\hat{x})^j(y_i-\hat{y})^k,$$

where $N = \sum_i w_i$.

Manually evaluating some of the low order terms yields the following:

j	k	$\mu_{xy}$
0	0	1
1	0	$\frac{1}{N} \sum_i w_i (x_i-\hat{x})$ = 0
0	1	$\frac{1}{N} \sum_i w_i (y_i-\hat{y})$ = 0
1	1	$\frac{1}{N} \sum_i w_i (x_i-\hat{x})(y_i-\hat{y})$ = Cov(x,y)
2	0	$\frac{1}{N} \sum_i w_i (x_i-\hat{x})^2$ = Var(x)
0	2	$\frac{1}{N} \sum_i w_i (y_i-\hat{y})^2$ = Var(y)

Instead, these lines:

PyORBIT3/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc

Lines 250 to 261 in 51f0568

	for(i = 0; i < _order; i++)
	momX[i+1] = momX[i]*m_size*((part_coord_arr[ip][0] - dispterm) - xAvg);
	for(i = 0; i< _order; i++)
	momY[i+1] = momY[i]*m_size*(part_coord_arr[ip][2] - yAvg);
	for(j = 0; j<_order; j++)
	for(i=0 ; i< _order+1-j; i++){
	momentXY[i][j] += momX[i]/pow(xbetaterm, double(i)) * momY[j]/pow(ybetaterm, double(j));
	momentXY[i][j] += momX[i] * momY[j];
	}
	}

are essentially calculating

$$ \mu_{xy}^{jk} = \frac{1}{N}\sum_i w_i^{j+k}(x_i-\hat{x})^j(y_i-\hat{y})^k\left(1+\frac{1}{\beta_x^j\beta_y^k}\right), $$

which, evaluates to:

j	k	$\mu_{xy}$
0	0	$\frac{1}{N} \sum_i 2 = \frac{2i}{N}$
1	0	$\frac{1}{N} \sum_i w_i (x_i-\hat{x})\left(1 + \frac{1}{\beta_x}\right)$
1	1	$\frac{1}{N} \sum_i w_i^2 (x_i-\hat{x})(y_i-\hat{y}) \left(1 + \frac{1}{\beta_x \beta_y}\right)$

and so forth.

This mixing of normalized and raw terms leads to tests on $i>0$, $j=0$ to fail:

tests/py/orbit/bunch_utils/test_bunch_twiss_analysis.py::test_computeBunchMoments_macrosize_order2_M00_M10_M01 - assert -7.968925040526781e-05 == -0.0002851747...7959 ± 1.0e-15
FAILED tests/py/orbit/bunch_utils/test_bunch_twiss_analysis.py::test_computeBunchMoments_macrosize_order2_M20 - assert 9.368853867319752e-08 == 8.91593940978...e-08 ± 1.0e-12

Tests for $i=0$, $j>0$, on the other hand, fail because the outer loop is missing a _order+1, so the element at the last index is always 0

FAILED tests/py/orbit/bunch_utils/test_bunch_twiss_analysis.py::test_computeBunchMoments_order3_M03 - assert 0.0 == 4.13235583192...e-13 ± 1.0e-15

PyORBIT3/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc

Line 256 in 51f0568

for(j = 0; j<_order; j++)

The branch where macrosize is not set has a bug due to missing braces, but I'm not sure this branch is ever used in-practice.

PyORBIT3/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc

Lines 286 to 289 in 51f0568

	for(j = 0; j<_order; j++)
	for(i=0 ; i< _order+1-j; i++)
	momentXY[i][j] += momX[i]/pow(xbetaterm, double(i)) * momY[j]/pow(ybetaterm, double(j));
	momentXY[i][j] += momX[i] * momY[j];

dispersionflag doesn't do anything

The dispersionflag branch is inaccessible because dispterm is explicitly initialized to 0 here:

PyORBIT3/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc

Line 153 in 51f0568

double dispterm = 0.;

and is never modified, so the argument to the if-statement on L167 is always false.

Performance Implications

Lastly, analyzeBunch correctly computes means/covariances. computeBunchMoments itself calls analyzeBunch, which essentially triggers a double loop over the entire bunch every time this method is called. For the time being, users should prefer to call analyzeBunch directly rather than computeBunchMoments for computing the covariance matrix, averages, etc.

[austin-hoover]

I agree that the macrosize needs to be moved outside the product sum when calculating the moments.

Maybe there should have been an option to compute normalized or unnormalized moments?

for(j = 0; j<_order; j++) 
    for(i=0 ; i< _order+1-j; i++){ 
        if (norm) {
            momentXY[i][j] += momX[i] / pow(xbetaterm, double(i)) * momY[j] / pow(ybetaterm, double(j)); 
        } 
        else {
            momentXY[i][j] += momX[i] * momY[j];
        }
    } 
}

Looks like xbetaterm is actually $\sqrt{\beta}$ or $\sqrt{\varepsilon \beta}$ (depending on emitnormflag), where $\varepsilon = \sqrt{\langle xx \rangle \langle x'x' \rangle - \langle xx' \rangle^2}$ and $\beta = \langle xx \rangle / \varepsilon$.

[woodtp]

Making it optional might have been the original intent, and I think that makes good sense. I could take the time to refactor a bit some point as I would also like to see the behavior of analyzeBunch and computeBunchMoments merged to avoid unnecessarily iterating the bunch twice.

Another idea that could potentially simplify things a bit is to use the methods introduced in #119 to reimplement BunchTwissAnalysis in pure Python and offload stats computation to numpy/scipy. Maybe even bring the functionality of BunchExtremaCalculator into a single BunchAnalyzer class or something.