Skip to content

Commit 2e46fa1

Browse files
ewanwmewanwm
authored
Add test for three flavour oscillations (#147)
* add three flavour osc test with comparisson to 'classic' barger propagator * make ThreeFlavourBarger interpret any density <= 0.0 as vacuum * add three flavour oscillation test for python too * change names of the python tests for more consistency * fix to make the PMNS matrix work... I have no idea how it was working before... I think it needs some reworking
1 parent 7c08b37 commit 2e46fa1

File tree

7 files changed

+433
-27
lines changed

7 files changed

+433
-27
lines changed

‎nuTens/propagator/pmns-matrix.hpp‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -58,10 +58,20 @@ class PMNSmatrix: public BaseMixingMatrix
5858
{
5959
NT_PROFILE();
6060

61+
_theta12.requiresGrad(false);
62+
_theta13.requiresGrad(false);
63+
_theta23.requiresGrad(false);
64+
_deltaCP.requiresGrad(false);
65+
6166
_theta12.setValue(theta12, 0);
6267
_theta13.setValue(theta13, 0);
6368
_theta23.setValue(theta23, 0);
64-
_deltaCP.setValue(deltaCP, 0);
69+
_deltaCP.setValue({0}, deltaCP);
70+
71+
_theta12.requiresGrad(true);
72+
_theta13.requiresGrad(true);
73+
_theta23.requiresGrad(true);
74+
_deltaCP.requiresGrad(true);
6575
}
6676

6777
/// @{Setters
@@ -74,10 +84,10 @@ class PMNSmatrix: public BaseMixingMatrix
7484
private:
7585

7686
// the mixing parameters
77-
AccessedTensor<float, 1, dtypes::kCPU> _theta12 = AccessedTensor<float, 1, dtypes::kCPU>::zeros({1});
78-
AccessedTensor<float, 1, dtypes::kCPU> _theta13 = AccessedTensor<float, 1, dtypes::kCPU>::zeros({1});
79-
AccessedTensor<float, 1, dtypes::kCPU> _theta23 = AccessedTensor<float, 1, dtypes::kCPU>::zeros({1});
80-
AccessedTensor<float, 1, dtypes::kCPU> _deltaCP = AccessedTensor<float, 1, dtypes::kCPU>::zeros({1});
87+
AccessedTensor<float, 1, dtypes::kCPU> _theta12 = AccessedTensor<float, 1, dtypes::kCPU>::zeros({1}, true);
88+
AccessedTensor<float, 1, dtypes::kCPU> _theta13 = AccessedTensor<float, 1, dtypes::kCPU>::zeros({1}, true);
89+
AccessedTensor<float, 1, dtypes::kCPU> _theta23 = AccessedTensor<float, 1, dtypes::kCPU>::zeros({1}, true);
90+
Tensor _deltaCP = Tensor::zeros({1}, dtypes::kComplexFloat, dtypes::kCPU, true);
8191

8292
// the sub-matrices
8393
Tensor _mat1;

‎tests/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ foreach(TESTNAME
1414
test-barger
1515
test-tensor
1616
test-two-flavour-osc
17+
test-three-flavour-osc
1718
test-pmns-matrix
1819
test-DP-propagator
1920
)

‎tests/barger-propagator.hpp‎

Lines changed: 35 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -352,41 +352,54 @@ class ThreeFlavourBarger
352352
/// @param a Row
353353
/// @param b Column
354354
/// @return Matrix element
355-
[[nodiscard]] inline std::complex<double> getTransitionMatrixElement(float energy, int a, int b) const {
355+
[[nodiscard]] inline std::complex<double> getTransitionMatrixElement(double energy, int a, int b) const {
356356

357357
std::complex<double> ret = 0.0;
358358

359-
for (int k = 0; k < 3; k++) {
359+
// interpret density <= 0.0 as vacuum, then transition matrix
360+
// is just matrix with exponential terms along diagonal
361+
if (_density <= 0.0 ) {
362+
363+
if ( a == b )
364+
ret = std::exp(- 0.5 * std::complex<double>(0.0, 1.0) * masses[a] * masses[a] * _baseline * 2.0 * M_PI / energy);
365+
366+
else
367+
ret = 0.0;
368+
}
360369

361-
std::complex<double> numerator = 4.0 * energy * energy * (
362-
getHamiltonianElement(energy, a, 0) * getHamiltonianElement(energy, 0, b) +
363-
getHamiltonianElement(energy, a, 1) * getHamiltonianElement(energy, 1, b) +
364-
getHamiltonianElement(energy, a, 2) * getHamiltonianElement(energy, 2, b)
365-
);
370+
else {
371+
for (int k = 0; k < 3; k++) {
366372

367-
std::complex<double> constant = 1.0;
368-
std::complex<double> denominator = 1.0;
373+
std::complex<double> numerator = 4.0 * energy * energy * (
374+
getHamiltonianElement(energy, a, 0) * getHamiltonianElement(energy, 0, b) +
375+
getHamiltonianElement(energy, a, 1) * getHamiltonianElement(energy, 1, b) +
376+
getHamiltonianElement(energy, a, 2) * getHamiltonianElement(energy, 2, b)
377+
);
369378

370-
for (int j = 0; j < 3; j++) {
379+
std::complex<double> constant = 1.0;
380+
std::complex<double> denominator = 1.0;
371381

372-
if (j == k) continue;
382+
for (int j = 0; j < 3; j++) {
373383

374-
numerator -= 2.0 * energy * getHamiltonianElement(energy, a, b) * calculateEffectiveM2(energy, j);
375-
denominator *= calculateEffectiveM2(energy, k) - calculateEffectiveM2(energy, j);
376-
constant *= calculateEffectiveM2(energy, j);
384+
if (j == k) continue;
377385

378-
}
386+
numerator -= 2.0 * energy * getHamiltonianElement(energy, a, b) * calculateEffectiveM2(energy, j);
387+
denominator *= calculateEffectiveM2(energy, k) - calculateEffectiveM2(energy, j);
388+
constant *= calculateEffectiveM2(energy, j);
379389

380-
std::complex<double> prod = numerator;
390+
}
381391

382-
if (a == b)
383-
prod += constant;
384-
385-
prod /= denominator;
392+
std::complex<double> prod = numerator;
386393

387-
std::complex<double> exponential = std::exp(-std::complex<double>(0.0, 1.0) * calculateEffectiveM2(energy, k) * _baseline * 2.0 * M_PI / (2.0 * energy));
394+
if (a == b)
395+
prod += constant;
396+
397+
prod /= denominator;
398+
399+
std::complex<double> exponential = std::exp(-std::complex<double>(0.0, 1.0) * calculateEffectiveM2(energy, k) * _baseline * 2.0 * M_PI / (2.0 * energy));
388400

389-
ret += prod * exponential;
401+
ret += prod * exponential;
402+
}
390403
}
391404

392405
return ret;

‎tests/python/test_three_flavour_oscillations.py‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
import unittest
2+
import math as m
3+
import nuTens as nt
4+
from nuTens.tensor import Tensor, matmul, scale
5+
from nuTens.testing import ThreeFlavourBarger
6+
from nuTens.propagator import ConstDensitySolver, PMNSmatrix
7+
8+
import numpy as np
9+
import typing
10+
11+
import pytest
12+
13+
@pytest.mark.parametrize("theta12", np.linspace(0.0, 2.0 * m.pi, 5, True))
14+
@pytest.mark.parametrize("theta13", np.linspace(0.0, 2.0 * m.pi, 5, True))
15+
@pytest.mark.parametrize("theta23", np.linspace(0.0, 2.0 * m.pi, 5, True))
16+
class TestTwoFlavourConstMatter:
17+
18+
m1 = 0.0
19+
m2 = 0.008 * nt.units.eV
20+
m3 = 0.01 * nt.units.eV
21+
22+
deltaCP = 0.25 * m.pi
23+
24+
baseline=295.0 * nt.units.km
25+
density=2.5
26+
27+
energy = 1.0 * nt.units.GeV
28+
29+
energy_tensor = Tensor.ones([1, 1], nt.dtype.scalar_type.complex_float, nt.dtype.device_type.cpu, False)
30+
energy_tensor.set_value([0, 0], energy)
31+
32+
def setup_tensor_inputs(self, theta12:float, theta13:float, theta23:float) -> typing.Tuple[Tensor]:
33+
34+
pmns = PMNSmatrix()
35+
pmns.set_parameter_values(theta12, theta13, theta23, self.deltaCP)
36+
37+
masses = Tensor([self.m1, self.m2, self.m3], nt.dtype.scalar_type.complex_float, nt.dtype.device_type.cpu, False).add_batch_dim()
38+
39+
return pmns.build(), masses
40+
41+
def test_compare_osc_probs(self, theta12:float, theta13:float, theta23:float):
42+
43+
print(f"------ const density oscillation probabilities ------\n")
44+
print(f" theta12 = {theta12}")
45+
print(f" theta13 = {theta13}")
46+
print(f" theta23 = {theta23}\n")
47+
48+
# set up barger
49+
barger = ThreeFlavourBarger()
50+
barger.set_params(self.m1, self.m2, self.m3, theta12, theta13, theta23, self.deltaCP, self.baseline, self.density)
51+
52+
pmns, masses = self.setup_tensor_inputs(theta12, theta13, theta23)
53+
54+
# set up tensor solver
55+
propagator = nt.propagator.Propagator(3, self.baseline)
56+
matter_solver = ConstDensitySolver(3, self.density)
57+
58+
propagator.set_matter_solver(matter_solver)
59+
propagator.set_mixing_matrix(pmns)
60+
propagator.set_masses(masses)
61+
propagator.set_energies(self.energy_tensor)
62+
63+
# calculate the evals + evecs + effective PMNS to print
64+
# for help when debugging
65+
evecs = nt.tensor.Tensor()
66+
evals = nt.tensor.Tensor()
67+
matter_solver.calculate_eigenvalues(evecs, evals)
68+
69+
tensor_osc_probs = propagator.calculate_probabilities()
70+
71+
print(f"Tensor solver evals: {evals.to_string()}\n")
72+
73+
print(f"Tensor solver effective Mi^2: {scale(evals, 2.0 * self.energy).to_string()}\n")
74+
75+
print(f"Barger M1^2: {barger.calculate_effective_m2(self.energy, 0)}")
76+
print(f"Barger M2^2: {barger.calculate_effective_m2(self.energy, 1)}")
77+
print(f"Barger M3^2: {barger.calculate_effective_m2(self.energy, 2)}\n")
78+
79+
print(f"Oscillation probabilities:")
80+
print(f"[0,0] tensor solver {tensor_osc_probs.get_value([0, 0, 0]):.5f} :: barger propagator {barger.calculate_prob(self.energy, i=0, j=0):.5f}")
81+
print(f"[0,1] tensor solver {tensor_osc_probs.get_value([0, 0, 1]):.5f} :: barger propagator {barger.calculate_prob(self.energy, i=0, j=1):.5f}")
82+
print(f"[0,2] tensor solver {tensor_osc_probs.get_value([0, 0, 2]):.5f} :: barger propagator {barger.calculate_prob(self.energy, i=0, j=2):.5f}")
83+
84+
print(f"[1,0] tensor solver {tensor_osc_probs.get_value([0, 1, 0]):.5f} :: barger propagator {barger.calculate_prob(self.energy, i=1, j=0):.5f}")
85+
print(f"[1,1] tensor solver {tensor_osc_probs.get_value([0, 1, 1]):.5f} :: barger propagator {barger.calculate_prob(self.energy, i=1, j=1):.5f}")
86+
print(f"[1,2] tensor solver {tensor_osc_probs.get_value([0, 1, 2]):.5f} :: barger propagator {barger.calculate_prob(self.energy, i=1, j=2):.5f}")
87+
88+
print(f"[2,0] tensor solver {tensor_osc_probs.get_value([0, 2, 0]):.5f} :: barger propagator {barger.calculate_prob(self.energy, i=2, j=0):.5f}")
89+
print(f"[2,1] tensor solver {tensor_osc_probs.get_value([0, 2, 1]):.5f} :: barger propagator {barger.calculate_prob(self.energy, i=2, j=1):.5f}")
90+
print(f"[2,2] tensor solver {tensor_osc_probs.get_value([0, 2, 2]):.5f} :: barger propagator {barger.calculate_prob(self.energy, i=2, j=2):.5f}")
91+
92+
assert (
93+
pytest.approx(tensor_osc_probs.get_value([0, 0, 0]), abs = 1e-5) == barger.calculate_prob(self.energy, i=0, j=0)
94+
), f"Const matter osc prob[0,0] != barger osc prob"
95+
assert (
96+
pytest.approx(tensor_osc_probs.get_value([0, 0, 1]), abs = 1e-5) == barger.calculate_prob(self.energy, i=0, j=1)
97+
), f"Const matter osc prob[0,1] != barger osc prob"
98+
assert (
99+
pytest.approx(tensor_osc_probs.get_value([0, 0, 2]), abs = 1e-5) == barger.calculate_prob(self.energy, i=0, j=2)
100+
), f"Const matter osc prob[0,2] != barger osc prob"
101+
102+
assert (
103+
pytest.approx(tensor_osc_probs.get_value([0, 1, 0]), abs = 1e-5) == barger.calculate_prob(self.energy, i=1, j=0)
104+
), f"Const matter osc prob[1,0] != barger osc prob"
105+
assert (
106+
pytest.approx(tensor_osc_probs.get_value([0, 1, 1]), abs = 1e-5) == barger.calculate_prob(self.energy, i=1, j=1)
107+
), f"Const matter osc prob[1,1] != barger osc prob"
108+
assert (
109+
pytest.approx(tensor_osc_probs.get_value([0, 1, 2]), abs = 1e-5) == barger.calculate_prob(self.energy, i=1, j=2)
110+
), f"Const matter osc prob[1,2] != barger osc prob"
111+
112+
assert (
113+
pytest.approx(tensor_osc_probs.get_value([0, 2, 0]), abs = 1e-5) == barger.calculate_prob(self.energy, i=2, j=0)
114+
), f"Const matter osc prob[2,0] != barger osc prob"
115+
assert (
116+
pytest.approx(tensor_osc_probs.get_value([0, 2, 1]), abs = 1e-5) == barger.calculate_prob(self.energy, i=2, j=1)
117+
), f"Const matter osc prob[2,1] != barger osc prob"
118+
assert (
119+
pytest.approx(tensor_osc_probs.get_value([0, 2, 2]), abs = 1e-5) == barger.calculate_prob(self.energy, i=2, j=2)
120+
), f"Const matter osc prob[2,2] != barger osc prob"
121+
122+
123+
if __name__ == '__main__':
124+
unittest.main()

tests/python/test_two-flavour-barger.py renamed to tests/python/test_two_flavour_barger.py

File renamed without changes.

tests/python/test_two-flavour-const-matter.py renamed to tests/python/test_two_flavour_oscillations.py

File renamed without changes.

0 commit comments

Comments
 (0)