Line data    Source code

       1             : /* 
       2             : * Copyright (C) 2020-2024 MEmilio
       3             : *
       4             : * Authors: Lena Ploetzke
       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             : #ifndef MIO_EPI_LCT_POPULATIONS_H
      21             : #define MIO_EPI_LCT_POPULATIONS_H
      22             : 
      23             : #include "boost/type_traits/make_void.hpp"
      24             : #include "memilio/config.h"
      25             : #include "memilio/utils/uncertain_value.h"
      26             : #include "memilio/math/eigen.h"
      27             : #include "memilio/epidemiology/lct_infection_state.h"
      28             : #include "memilio/utils/type_list.h"
      29             : #include "memilio/utils/metaprogramming.h"
      30             : 
      31             : namespace mio
      32             : {
      33             : 
      34             : /**
      35             :  * @brief A class template for compartment populations of LCT models.
      36             :  *
      37             :  * Populations can be split up into different categories, e.g. by age group, yearly income group, gender etc. 
      38             :  * In LCT models, we want to use different numbers of subcompartments, i.e. different LctStates, 
      39             :  * for each group of a category.
      40             :  * (Therefore, we can't use the normal Populations class because it expects the same InfectionStates for each group.)
      41             :  * 
      42             :  * This template must be instantiated with one LctState for each group of a category. 
      43             :  * The purpose of the LctStates is to define the number of subcompartments for each InfectionState.
      44             :  * The number of LctStates also determines the number of groups. 
      45             :  * If you want to use more than one category, e.g. age and gender, you have to define num_age_groups * num_genders 
      46             :  * LctStates, because the number of subcompartments can be different 
      47             :  * even for (female, A05-A14) and (female, A80+) or (male, A05-A14).
      48             :  *
      49             :  * The class created from this template contains a "flat array" of compartment populations and some functions for 
      50             :  * retrieving or setting the populations. The order in the flat array is: First, all (sub-)compartments of the 
      51             :  * first group, afterwards all (sub-)compartments of the second group and so on.
      52             :  *
      53             :  */
      54             : 
      55             : template <typename FP = ScalarType, class... LctStates>
      56             : class LctPopulations
      57             : {
      58             : public:
      59             :     using Type                         = UncertainValue<FP>;
      60             :     using InternalArrayType            = Eigen::Array<Type, Eigen::Dynamic, 1>;
      61             :     using LctStatesGroups              = TypeList<LctStates...>;
      62             :     static size_t constexpr num_groups = sizeof...(LctStates); ///< Number of groups.
      63             :     static_assert(num_groups >= 1, "The number of LctStates provided should be at least one.");
      64             : 
      65             :     /// @brief Default constructor.
      66          19 :     LctPopulations()
      67          19 :     {
      68          19 :         set_count();
      69          19 :         m_y = InternalArrayType::Constant(m_count, 1, 0.);
      70          19 :     }
      71             : 
      72             :     /**
      73             :      * @brief get_num_compartments Returns the number of compartments.
      74             :      * @return Number of compartments which equals the flat array size.
      75             :      */
      76         602 :     size_t get_num_compartments() const
      77             :     {
      78         602 :         return m_count;
      79             :     }
      80             : 
      81             :     /**
      82             :      * @brief Returns a reference to the internally stored flat array.
      83             :      * @return Const reference to the InternalArrayType instance.
      84             :      */
      85             :     auto const& array() const
      86             :     {
      87             :         return m_y;
      88             :     }
      89             :     auto& array()
      90             :     {
      91             :         return m_y;
      92             :     }
      93             : 
      94             :     /**
      95             :      * @brief Returns the entry of the array given a flat index.
      96             :      * @param index A flat index.
      97             :      * @return The value of the internal array at the index.
      98             :      */
      99        1629 :     Type& operator[](size_t index)
     100             :     {
     101        1629 :         assert(index < m_count);
     102        1629 :         return m_y[index];
     103             :     }
     104             : 
     105             :     /**
     106             :     * @brief Gets the first index of a group in the flat array.
     107             :     * @tparam group The group for which the index should be returned.
     108             :     * @return The index of the first entry of group in the flat array.
     109             :     */
     110             :     template <size_t Group>
     111        2038 :     size_t get_first_index_of_group() const
     112             :     {
     113             :         static_assert((Group < num_groups) && (Group >= 0), "The template parameter Group should be valid.");
     114             :         if constexpr (Group == 0) {
     115        1318 :             return 0;
     116             :         }
     117             :         else {
     118         720 :             return get_first_index_of_group<Group - 1>() + type_at_index_t<Group - 1, LctStatesGroups>::Count;
     119             :         }
     120             :     }
     121             :     /**
     122             :      * @brief Returns an Eigen copy of the vector of populations. 
     123             :      * This can be used as initial conditions for the ODE solver.
     124             :      * @return Eigen::VectorXd of populations.
     125             :      */
     126         135 :     inline Vector<FP> get_compartments() const
     127             :     {
     128         135 :         return m_y.array().template cast<FP>();
     129             :     }
     130             : 
     131             :     /**
     132             :      * @brief Returns the total population of a group.
     133             :      * @tparam group The group for which the total population should be calculated.
     134             :      * @return Total population of the group.
     135             :      */
     136             :     template <size_t Group>
     137             :     FP get_group_total() const
     138             :     {
     139             :         return m_y.array()
     140             :             .template cast<FP>()
     141             :             .segment(get_first_index_of_group<Group>(), type_at_index_t<Group, LctStatesGroups>::Count)
     142             :             .sum();
     143             :     }
     144             : 
     145             :     /**
     146             :      * @brief Returns the total population of all compartments and groups.
     147             :      * @return Total population.
     148             :      */
     149             :     FP get_total() const
     150             :     {
     151             :         return m_y.array().template cast<FP>().sum();
     152             :     }
     153             : 
     154             :     /**
     155             :      * @brief Checks whether all compartments have non-negative values. 
     156             :      * This function can be used to prevent slightly negative function values in compartment sizes that are produced 
     157             :      * due to rounding errors if, e.g., population sizes were computed in a complex way.
     158             :      *
     159             :      * Attention: This function should be used with care. It can not and will not set model parameters and 
     160             :      *            compartments to meaningful values. In most cases it is preferable to use check_constraints,
     161             :      *            and correct values manually before proceeding with the simulation.
     162             :      *            The main usage for apply_constraints is in automated tests using random values for initialization.
     163             :      *
     164             :      * @return Returns true if one (or more) constraint(s) were corrected, otherwise false.
     165             :     */
     166             :     bool apply_constraints()
     167             :     {
     168             :         bool corrected = false;
     169             :         for (int i = 0; i < m_y.array().size(); i++) {
     170             :             if (m_y.array()[i] < 0) {
     171             :                 log_warning("Constraint check: Compartment size {:d} changed from {:.4f} to {:d}", i, m_y.array()[i],
     172             :                             0);
     173             :                 m_y.array()[i] = 0.;
     174             :                 corrected      = true;
     175             :             }
     176             :         }
     177             :         return corrected;
     178             :     }
     179             : 
     180             :     /**
     181             :      * @brief Checks whether all compartments have non-negative values and logs an error if constraint is not satisfied.
     182             :      * @return Returns true if one or more constraints are not satisfied, false otherwise.
     183             :      */
     184           9 :     bool check_constraints() const
     185             :     {
     186           9 :         if ((m_y.array() < 0.).any()) {
     187           2 :             log_error("Constraint check: At least one compartment size is smaller {}.", 0);
     188           1 :             return true;
     189             :         }
     190           8 :         return false;
     191             :     }
     192             : 
     193             : private:
     194             :     /**
     195             :      * @brief Sets recursively the total number of (sub-)compartments over all groups.
     196             :      * The number also corresponds to the size of the internal vector.
     197             :      */
     198             :     template <size_t Group = 0>
     199          60 :     void set_count()
     200             :     {
     201             :         if constexpr (Group == 0) {
     202          19 :             m_count = 0;
     203             :         }
     204             :         if constexpr (Group < num_groups) {
     205          41 :             m_count += type_at_index_t<Group, LctStatesGroups>::Count;
     206          41 :             set_count<Group + 1>();
     207             :         }
     208          60 :     }
     209             : 
     210             :     size_t m_count; //< Number of groups stored.
     211             :     InternalArrayType m_y{}; //< An array containing the number of people in the groups.
     212             : };
     213             : 
     214             : } // namespace mio
     215             : 
     216             : #endif // MIO_EPI_LCT_POPULATIONS_H