diff --git a/benchmarks/benchmarks.cpp b/benchmarks/benchmarks.cpp index 913f3c4..79990df 100644 --- a/benchmarks/benchmarks.cpp +++ b/benchmarks/benchmarks.cpp @@ -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 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 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 to discard the return value that we're not supposed to discard :) @@ -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); @@ -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 matterSolver = std::make_shared(3, 2.6); - matterProp.setPMNS(PMNS); + Propagator matterProp(3, 295000.0); + matterProp.setPMNS(PMNS.matrix); matterProp.setMasses(masses); + + std::shared_ptr matterSolver = std::make_shared(3, 2.6); + matterProp.setMatterSolver(matterSolver); // seed the random number generator for the energies @@ -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)); } }