/*

* Copyright (C) 2020-2026 MEmilio

*

* Authors: Rene Schmieding, Henrik Zunker

*

* Contact: Martin J. Kuehn <[email protected]>

*

* Licensed under the Apache License, Version 2.0 (the "License");

* you may not use this file except in compliance with the License.

* You may obtain a copy of the License at

*

* http://www.apache.org/licenses/LICENSE-2.0

*

* Unless required by applicable law or agreed to in writing, software

* distributed under the License is distributed on an "AS IS" BASIS,

* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

* See the License for the specific language governing permissions and

* limitations under the License.

*/

#include "load_test_data.h"

#include "memilio/compartments/compartmental_model.h"

#include "memilio/compartments/flow_simulation.h"

#include "memilio/compartments/simulation.h"

#include "ode_seir/infection_state.h"

#include "ode_seir/model.h"

#include "ode_seir/parameters.h"

#include "gtest/gtest.h"

using I = mio::oseir::InfectionState;

using Flows = mio::TypeList<mio::Flow<I::Susceptible, I::Exposed>, mio::Flow<I::Exposed, I::Infected>,

mio::Flow<I::Infected, I::Recovered>>;

struct CatA : public mio::Index<CatA> {

CatA(size_t i)

: mio::Index<CatA>(i)

{

}

};

struct CatB : public mio::Index<CatB> {

CatB(size_t i)

: mio::Index<CatB>(i)

{

}

};

struct CatC : public mio::Index<CatC> {

CatC(size_t i)

: mio::Index<CatC>(i)

{

}

};

class TestModel : public mio::FlowModel<double, I, mio::Populations<double, I, CatA, CatB, CatC>,

mio::oseir::Parameters<double>, Flows>

{

using Base =

mio::FlowModel<double, I, mio::Populations<double, I, CatA, CatB, CatC>, mio::oseir::Parameters<double>, Flows>;

public:

TestModel(Populations::Index dimensions)

: Base(Populations(dimensions, 0.), mio::oseir::Parameters<double>(mio::AgeGroup(1)))

{

}

void get_flows(Eigen::Ref<const Eigen::VectorXd> /*pop*/, Eigen::Ref<const Eigen::VectorXd> /*y*/, double /*t*/,

Eigen::Ref<Eigen::VectorXd> /*flows*/) const override

{

}

};

TEST(TestFlows, FlowChart)

{

EXPECT_EQ(Flows().size(), 3);

// Testing get (by index) function, verifying with source members.

auto flow0 = mio::Flow<I::Susceptible, I::Exposed>().source;

auto test_source0 = mio::type_at_index_t<0, Flows>::source;

EXPECT_EQ(flow0, test_source0);

auto flow1 = mio::Flow<I::Exposed, I::Infected>().source;

auto test_source1 = mio::type_at_index_t<1, Flows>::source;

EXPECT_EQ(flow1, test_source1);

auto flow2 = mio::Flow<I::Infected, I::Recovered>().source;

auto test_source2 = mio::type_at_index_t<2, Flows>::source;

EXPECT_EQ(flow2, test_source2);

// Testing get (by Flow) function.

const size_t index0 = mio::index_of_type_v<mio::Flow<I::Susceptible, I::Exposed>, Flows>;

EXPECT_EQ(index0, 0);

const size_t index1 = mio::index_of_type_v<mio::Flow<I::Exposed, I::Infected>, Flows>;

EXPECT_EQ(index1, 1);

const size_t index2 = mio::index_of_type_v<mio::Flow<I::Infected, I::Recovered>, Flows>;

EXPECT_EQ(index2, 2);

}

TEST(TestFlows, FlowSimulation)

{

double t0 = 0;

double tmax = 1;

double dt = 0.001;

mio::oseir::Model<double> model(1);

double total_population = 10000;

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Exposed}] = 100;

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Infected}] = 100;

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Recovered}] = 100;

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Susceptible}] =

total_population - model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Exposed}] -

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Infected}] -

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Recovered}];

// suscetible now set with every other update

// params.nb_sus_t0 = params.nb_total_t0 - params.nb_exp_t0 - params.nb_inf_t0 - params.nb_rec_t0;

model.parameters.set<mio::oseir::TimeExposed<double>>(5.2);

model.parameters.set<mio::oseir::TimeInfected<double>>(6);

model.parameters.set<mio::oseir::TransmissionProbabilityOnContact<double>>(0.04);

model.parameters.get<mio::oseir::ContactPatterns<double>>().get_cont_freq_mat()[0].get_baseline().setConstant(10);

mio::set_log_level(mio::LogLevel::off); // Suppress log output of check_constraints and the Simulation.

model.check_constraints();

auto seir = mio::simulate_flows<double, mio::oseir::Model<double>>(

t0, tmax, dt, model, std::make_unique<mio::DefaultIntegratorCore<double>>());

mio::set_log_level(mio::LogLevel::warn);

// verify results (computed using flows)

auto results = seir[0].get_last_value();

EXPECT_NEAR(results[0], 9660.5835936179408, 1e-14);

EXPECT_NEAR(results[1], 118.38410512338650, 1e-14);

EXPECT_NEAR(results[2], 104.06636087558745, 1e-14);

EXPECT_NEAR(results[3], 116.96594038308581, 1e-14);

// test flow results

auto flows_results = seir[1].get_last_value();

EXPECT_NEAR(flows_results[0], 39.416406382059762, 1e-14);

EXPECT_NEAR(flows_results[1], 21.032301258673261, 1e-14);

EXPECT_NEAR(flows_results[2], 16.965940383085815, 1e-14);

}

TEST(TestFlows, CompareSimulations)

{

double t0 = 0;

double tmax = 1;

double dt = 0.001;

mio::oseir::Model<double> model(1);

double total_population = 10000;

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Exposed}] = 100;

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Infected}] = 100;

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Recovered}] = 100;

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Susceptible}] =

total_population - model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Exposed}] -

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Infected}] -

model.populations[{mio::AgeGroup(0), mio::oseir::InfectionState::Recovered}];

// suscetible now set with every other update

// params.nb_sus_t0 = params.nb_total_t0 - params.nb_exp_t0 - params.nb_inf_t0 - params.nb_rec_t0;

model.parameters.set<mio::oseir::TimeExposed<double>>(5.2);

model.parameters.set<mio::oseir::TimeInfected<double>>(6);

model.parameters.set<mio::oseir::TransmissionProbabilityOnContact<double>>(0.04);

mio::ContactMatrixGroup<double>& contact_matrix =

model.parameters.get<mio::oseir::ContactPatterns<double>>().get_cont_freq_mat();

contact_matrix[0].get_baseline().setConstant(10);

mio::set_log_level(mio::LogLevel::off); // Suppress log output of check_constraints and the Simulation.

model.check_constraints();

auto seir_sim_flows = simulate_flows(t0, tmax, dt, model);

auto seir_sim = simulate(t0, tmax, dt, model);

mio::set_log_level(mio::LogLevel::warn);

auto results_flows = seir_sim_flows[0].get_last_value();

auto results = seir_sim.get_last_value();

EXPECT_NEAR(results[0], results_flows[0], 1e-10);

EXPECT_NEAR(results[1], results_flows[1], 1e-10);

EXPECT_NEAR(results[2], results_flows[2], 1e-10);

EXPECT_NEAR(results[3], results_flows[3], 1e-10);

}

TEST(TestFlows, GetInitialFlows)

{

mio::oseir::Model<double> m(1);

EXPECT_EQ(m.get_initial_flows().size(), 3); // 3 == Flows().size()

EXPECT_EQ(m.get_initial_flows().norm(), 0);

}

TEST(TestFlows, GetFlowIndex)

{

// test get_flat_flow_index with some prime number products

TestModel m({I::Count, CatA(11), CatB(5), CatC(7)});

auto idx0 = m.get_flat_flow_index<I::Susceptible, I::Exposed>({CatA(0), CatB(0), CatC(1)});

EXPECT_EQ(idx0, 3);

auto idx1 = m.get_flat_flow_index<I::Susceptible, I::Exposed>({CatA(0), CatB(1), CatC(0)});

EXPECT_EQ(idx1, 7 * 3);

auto idx2 = m.get_flat_flow_index<I::Susceptible, I::Exposed>({CatA(1), CatB(0), CatC(0)});

EXPECT_EQ(idx2, 5 * 7 * 3);

auto idx3 = m.get_flat_flow_index<I::Susceptible, I::Exposed>({CatA(1), CatB(1), CatC(1)});

EXPECT_EQ(idx3, (5 * 7 * 3) + (7 * 3) + (3));

auto idx4 = m.get_flat_flow_index<I::Susceptible, I::Exposed>({CatA(10), CatB(4), CatC(6)});

EXPECT_EQ(idx4, 10 * (5 * 7 * 3) + 4 * (7 * 3) + 6 * (3));

}