Line data    Source code

       1             : /* 
       2             : * Copyright (C) 2020-2024 MEmilio
       3             : *
       4             : * Authors: Daniel Abele, Jan Kleinert, Martin J. Kuehn
       5             : *
       6             : * Contact: Martin J. Kuehn <Martin.Kuehn@DLR.de>
       7             : *
       8             : * Licensed under the Apache License, Version 2.0 (the "License");
       9             : * you may not use this file except in compliance with the License.
      10             : * You may obtain a copy of the License at
      11             : *
      12             : *     http://www.apache.org/licenses/LICENSE-2.0
      13             : *
      14             : * Unless required by applicable law or agreed to in writing, software
      15             : * distributed under the License is distributed on an "AS IS" BASIS,
      16             : * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
      17             : * See the License for the specific language governing permissions and
      18             : * limitations under the License.
      19             : */
      20             : 
      21             : #ifndef ODESIR_MODEL_H
      22             : #define ODESIR_MODEL_H
      23             : 
      24             : #include "memilio/compartments/compartmentalmodel.h"
      25             : #include "memilio/epidemiology/age_group.h"
      26             : #include "memilio/epidemiology/populations.h"
      27             : #include "memilio/epidemiology/contact_matrix.h"
      28             : #include "ode_sir/infection_state.h"
      29             : #include "ode_sir/parameters.h"
      30             : 
      31             : namespace mio
      32             : {
      33             : namespace osir
      34             : {
      35             : 
      36             : /********************
      37             :  * define the model *
      38             :  ********************/
      39             : 
      40             : template <typename FP = ScalarType>
      41             : class Model
      42             :     : public mio::CompartmentalModel<FP, InfectionState, mio::Populations<FP, AgeGroup, InfectionState>, Parameters<FP>>
      43             : {
      44             :     using Base =
      45             :         mio::CompartmentalModel<FP, InfectionState, mio::Populations<FP, AgeGroup, InfectionState>, Parameters<FP>>;
      46             : 
      47             : public:
      48             :     using typename Base::ParameterSet;
      49             :     using typename Base::Populations;
      50             : 
      51             :     Model(const Populations& pop, const ParameterSet& params)
      52             :         : Base(pop, params)
      53             :     {
      54             :     }
      55             : 
      56          72 :     Model(int num_agegroups)
      57          72 :         : Base(Populations({AgeGroup(num_agegroups), InfectionState::Count}), ParameterSet(AgeGroup(num_agegroups)))
      58             :     {
      59          72 :     }
      60             : 
      61        4077 :     void get_derivatives(Eigen::Ref<const Vector<FP>> pop, Eigen::Ref<const Vector<FP>> y, FP t,
      62             :                          Eigen::Ref<Vector<FP>> dydt) const override
      63             :     {
      64        4077 :         auto params                              = this->parameters;
      65        4077 :         AgeGroup n_agegroups                     = params.get_num_groups();
      66        4077 :         ContactMatrixGroup const& contact_matrix = params.template get<ContactPatterns<FP>>();
      67             : 
      68        8154 :         for (auto i = AgeGroup(0); i < n_agegroups; i++) {
      69             : 
      70        4077 :             size_t Si = this->populations.get_flat_index({i, InfectionState::Susceptible});
      71        4077 :             size_t Ii = this->populations.get_flat_index({i, InfectionState::Infected});
      72        4077 :             size_t Ri = this->populations.get_flat_index({i, InfectionState::Recovered});
      73             : 
      74        8154 :             for (auto j = AgeGroup(0); j < n_agegroups; j++) {
      75             : 
      76        4077 :                 size_t Sj = this->populations.get_flat_index({j, InfectionState::Susceptible});
      77        4077 :                 size_t Ij = this->populations.get_flat_index({j, InfectionState::Infected});
      78        4077 :                 size_t Rj = this->populations.get_flat_index({j, InfectionState::Recovered});
      79             : 
      80        4077 :                 const ScalarType Nj    = pop[Sj] + pop[Ij] + pop[Rj];
      81        4077 :                 const ScalarType divNj = (Nj < Limits<ScalarType>::zero_tolerance()) ? 0.0 : 1.0 / Nj;
      82             : 
      83        8154 :                 ScalarType coeffStoI = contact_matrix.get_matrix_at(t)(static_cast<Eigen::Index>((size_t)i),
      84        8154 :                                                                        static_cast<Eigen::Index>((size_t)j)) *
      85        4077 :                                        params.template get<TransmissionProbabilityOnContact<FP>>()[i] * divNj;
      86             : 
      87        4077 :                 dydt[Si] += -coeffStoI * y[Si] * pop[Ij];
      88        4077 :                 dydt[Ii] += coeffStoI * y[Si] * pop[Ij];
      89             :             }
      90        4077 :             dydt[Ii] -= (1.0 / params.template get<TimeInfected<FP>>()[i]) * y[Ii];
      91        4077 :             dydt[Ri] = (1.0 / params.template get<TimeInfected<FP>>()[i]) * y[Ii];
      92             :         }
      93        8154 :     }
      94             : 
      95             :     /**
      96             :      * serialize this. 
      97             :      * @see mio::serialize
      98             :      */
      99             :     template <class IOContext>
     100             :     void serialize(IOContext& io) const
     101             :     {
     102             :         auto obj = io.create_object("Model");
     103             :         obj.add_element("Parameters", this->parameters);
     104             :         obj.add_element("Populations", this->populations);
     105             :     }
     106             : 
     107             :     /**
     108             :      * deserialize an object of this class.
     109             :      * @see mio::deserialize
     110             :      */
     111             :     template <class IOContext>
     112             :     static IOResult<Model> deserialize(IOContext& io)
     113             :     {
     114             :         auto obj = io.expect_object("Model");
     115             :         auto par = obj.expect_element("Parameters", Tag<ParameterSet>{});
     116             :         auto pop = obj.expect_element("Populations", Tag<Populations>{});
     117             :         return apply(
     118             :             io,
     119             :             [](auto&& par_, auto&& pop_) {
     120             :                 return Model{pop_, par_};
     121             :             },
     122             :             par, pop);
     123             :     }
     124             : };
     125             : 
     126             : } // namespace osir
     127             : } // namespace mio
     128             : 
     129             : #endif // ODESIR_MODEL_H