Skip to content

Commit

Permalink
now randomise the osc parameters rather than energies. right now seem…
Browse files Browse the repository at this point in the history 
…s to be a lot slower....
  • Loading branch information
@ewanwm
ewanwm committed Jan 7, 2025
1 parent ab46c55 commit 54f93c6
Showing 1 changed file with 98 additions and 69 deletions.
167 changes: 98 additions & 69 deletions benchmarks/benchmarks.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,63 +8,81 @@
// want this to be fixed for reproducibility
const int randSeed = 123;

// set the PMNS parameters to use
// Will very likely change the benchmark so that energies are fixed
// and these get randomised but for now just set them here
const float m1 = 0.1;
const float m2 = 0.2;
const float m3 = 0.3;
const std::complex<float> i(0.0, 1.0);

const float th12 = 0.12;
const float th23 = 0.23;
const float th13 = 0.13;
class PMNSmatrix
{
public:
PMNSmatrix()
{
// set up the three matrices to build the PMNS matrix
_m1 = Tensor::zeros({1, 3, 3}, NTdtypes::kComplexFloat).requiresGrad(false);
_m2 = Tensor::zeros({1, 3, 3}, NTdtypes::kComplexFloat).requiresGrad(false);
_m3 = Tensor::zeros({1, 3, 3}, NTdtypes::kComplexFloat).requiresGrad(false);
}

void build(const Tensor &theta12, const Tensor &theta13, const Tensor &theta23, const Tensor &deltaCP)
{
_m1.requiresGrad(false);
_m2.requiresGrad(false);
_m3.requiresGrad(false);

_m1.setValue({0, 0, 0}, 1.0);
_m1.setValue({0, 1, 1}, Tensor::cos(theta23));
_m1.setValue({0, 1, 2}, Tensor::sin(theta23));
_m1.setValue({0, 2, 1}, -Tensor::sin(theta23));
_m1.setValue({0, 2, 2}, Tensor::cos(theta23));
_m1.requiresGrad(true);

_m2.setValue({0, 1, 1}, 1.0);
_m2.setValue({0, 0, 0}, Tensor::cos(theta13));
_m2.setValue({0, 0, 2}, Tensor::mul(Tensor::sin(theta13), Tensor::exp(Tensor::scale(deltaCP, -i))));
_m2.setValue({0, 2, 0}, -Tensor::mul(Tensor::sin(theta13), Tensor::exp(Tensor::scale(deltaCP, i))));
_m2.setValue({0, 2, 2}, Tensor::cos(theta13));
_m2.requiresGrad(true);

_m3.setValue({0, 2, 2}, 1.0);
_m3.setValue({0, 0, 0}, Tensor::cos(theta12));
_m3.setValue({0, 0, 1}, Tensor::sin(theta12));
_m3.setValue({0, 1, 0}, -Tensor::sin(theta12));
_m3.setValue({0, 1, 1}, Tensor::cos(theta12));
_m3.requiresGrad(true);

// Build PMNS
matrix = Tensor::matmul(_m1, Tensor::matmul(_m2, _m3));
matrix.requiresGrad(true);
}

const float dcp = 0.5;
Tensor matrix;

Tensor buildPMNS(const Tensor &theta12, const Tensor &theta13, const Tensor &theta23, const Tensor &deltaCP)
{
// set up the three matrices to build the PMNS matrix
Tensor M1 = Tensor::zeros({1, 3, 3}, NTdtypes::kComplexFloat).requiresGrad(false);
Tensor M2 = Tensor::zeros({1, 3, 3}, NTdtypes::kComplexFloat).requiresGrad(false);
Tensor M3 = Tensor::zeros({1, 3, 3}, NTdtypes::kComplexFloat).requiresGrad(false);

M1.setValue({0, 0, 0}, 1.0);
M1.setValue({0, 1, 1}, Tensor::cos(theta23));
M1.setValue({0, 1, 2}, Tensor::sin(theta23));
M1.setValue({0, 2, 1}, -Tensor::sin(theta23));
M1.setValue({0, 2, 2}, Tensor::cos(theta23));
M1.requiresGrad(true);

M2.setValue({0, 1, 1}, 1.0);
M2.setValue({0, 0, 0}, Tensor::cos(theta13));
std::complex<float> i(0.0, 1.0);
M2.setValue({0, 0, 2}, Tensor::mul(Tensor::sin(theta13), Tensor::exp(Tensor::scale(deltaCP, -i))));
M2.setValue({0, 2, 0}, -Tensor::mul(Tensor::sin(theta13), Tensor::exp(Tensor::scale(deltaCP, i))));
M2.setValue({0, 2, 2}, Tensor::cos(theta13));
M2.requiresGrad(true);

M3.setValue({0, 2, 2}, 1.0);
M3.setValue({0, 0, 0}, Tensor::cos(theta12));
M3.setValue({0, 0, 1}, Tensor::sin(theta12));
M3.setValue({0, 1, 0}, -Tensor::sin(theta12));
M3.setValue({0, 1, 1}, Tensor::cos(theta12));
M3.requiresGrad(true);

// Build PMNS
Tensor PMNS = Tensor::matmul(M1, Tensor::matmul(M2, M3));
PMNS.requiresGrad(true);

return PMNS;
}
private:
Tensor _m1;
Tensor _m2;
Tensor _m3;
};

static void batchedOscProbs(const Propagator &prop, long batchSize, long nBatches)
static void batchedOscProbs(Propagator &prop, PMNSmatrix &matrix, Tensor &theta23, Tensor &theta13, Tensor &theta12,
Tensor &deltaCP, Tensor &masses, const Tensor &energies, long batchSize, long nBatches)
{
for (int _ = 0; _ < nBatches; _++)
{

Tensor energies =
Tensor::scale(Tensor::rand({batchSize, 1}).dType(NTdtypes::kFloat).requiresGrad(false), 10000.0) +
Tensor({100.0});
// set random values of the oscillation parameters
masses.setValue({0, 0}, Tensor::rand({1}));
masses.setValue({0, 1}, Tensor::rand({1}));
masses.setValue({0, 2}, Tensor::rand({1}));

theta23.setValue({0}, Tensor::rand({1}));
theta13.setValue({0}, Tensor::rand({1}));
theta12.setValue({0}, Tensor::rand({1}));

deltaCP.setValue({0}, Tensor::scale(Tensor::rand({1}), 2.0 * 3.1415));

// calculate new values of the PMNS matrix
matrix.build(theta12, theta13, theta23, deltaCP);

prop.setPMNS(matrix.matrix);
prop.setMasses(masses);

// calculate the osc probabilities
// static_cast<void> to discard the return value that we're not supposed to discard :)
Expand All @@ -74,21 +92,23 @@ static void batchedOscProbs(const Propagator &prop, long batchSize, long nBatche

static void BM_vacuumOscillations(benchmark::State &state)
{
// make some random test energies
Tensor energies =
Tensor::scale(Tensor::rand({state.range(0), 1}).dType(NTdtypes::kFloat).requiresGrad(false), 10000.0) +
Tensor({100.0});

// set up the inputs
Tensor masses = Tensor({m1, m2, m3}, NTdtypes::kFloat).requiresGrad(false).addBatchDim();
Tensor masses = Tensor::zeros({1, 3}, NTdtypes::kFloat).requiresGrad(false);

Tensor theta23 = Tensor({th23}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor theta13 = Tensor({th13}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor theta12 = Tensor({th12}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor deltaCP = Tensor({dcp}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor theta23 = Tensor::zeros({1}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor theta13 = Tensor::zeros({1}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor theta12 = Tensor::zeros({1}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor deltaCP = Tensor::zeros({1}).dType(NTdtypes::kComplexFloat).requiresGrad(false);

Tensor PMNS = buildPMNS(theta12, theta13, theta23, deltaCP);
PMNSmatrix PMNS;

// set up the propagator
Propagator vacuumProp(3, 100.0);
vacuumProp.setPMNS(PMNS);
vacuumProp.setMasses(masses);
Propagator vacuumProp(3, 295000.0);

// seed the random number generator for the energies
std::srand(randSeed);
Expand All @@ -98,28 +118,36 @@ static void BM_vacuumOscillations(benchmark::State &state)
for (auto _ : state)
{
// This code gets timed
batchedOscProbs(vacuumProp, state.range(0), state.range(1));
batchedOscProbs(vacuumProp, PMNS, theta23, theta13, theta12, deltaCP, masses, energies, state.range(0),
state.range(1));
}
}

static void BM_constMatterOscillations(benchmark::State &state)
{
// make some random test energies
Tensor energies =
Tensor::scale(Tensor::rand({state.range(0), 1}).dType(NTdtypes::kFloat).requiresGrad(false), 10000.0) +
Tensor({100.0});

// set up the inputs
Tensor masses = Tensor({m1, m2, m3}, NTdtypes::kFloat).requiresGrad(false).addBatchDim();
Tensor masses = Tensor::zeros({1, 3}, NTdtypes::kFloat).requiresGrad(false);

Tensor theta23 = Tensor({th23}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor theta13 = Tensor({th13}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor theta12 = Tensor({th12}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor deltaCP = Tensor({dcp}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor theta23 = Tensor::zeros({1}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor theta13 = Tensor::zeros({1}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor theta12 = Tensor::zeros({1}).dType(NTdtypes::kComplexFloat).requiresGrad(false);
Tensor deltaCP = Tensor::zeros({1}).dType(NTdtypes::kComplexFloat).requiresGrad(false);

Tensor PMNS = buildPMNS(theta12, theta13, theta23, deltaCP);
PMNSmatrix PMNS;
PMNS.build(theta12, theta13, theta23, deltaCP);

// set up the propagator
Propagator matterProp(3, 100.0);
std::shared_ptr<BaseMatterSolver> matterSolver = std::make_shared<ConstDensityMatterSolver>(3, 2.6);
matterProp.setPMNS(PMNS);
Propagator matterProp(3, 295000.0);
matterProp.setPMNS(PMNS.matrix);
matterProp.setMasses(masses);

std::shared_ptr<BaseMatterSolver> matterSolver = std::make_shared<ConstDensityMatterSolver>(3, 2.6);

matterProp.setMatterSolver(matterSolver);

// seed the random number generator for the energies
Expand All @@ -130,7 +158,8 @@ static void BM_constMatterOscillations(benchmark::State &state)
for (auto _ : state)
{
// This code gets timed
batchedOscProbs(matterProp, state.range(0), state.range(1));
batchedOscProbs(matterProp, PMNS, theta23, theta13, theta12, deltaCP, masses, energies, state.range(0),
state.range(1));
}
}

Expand Down

0 comments on commit 54f93c6

Please sign in to comment.