diversity_pool_calculator.hpp

Go to the documentation of this file.

#ifndef GENESIS_POPULATION_FUNCTION_DIVERSITY_POOL_CALCULATOR_H_
#define GENESIS_POPULATION_FUNCTION_DIVERSITY_POOL_CALCULATOR_H_
 
/*
    Genesis - A toolkit for working with phylogenetic data.
    Copyright (C) 2014-2024 Lucas Czech
 
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.
 
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.
 
    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
    Contact:
    Lucas Czech <lucas.czech@sund.ku.dk>
    University of Copenhagen, Globe Institute, Section for GeoGenetics
    Oster Voldgade 5-7, 1350 Copenhagen K, Denmark
*/
 
#include "genesis/population/filter/sample_counts_filter.hpp"
#include "genesis/population/function/diversity_pool_functions.hpp"
#include "genesis/population/function/functions.hpp"
#include "genesis/population/variant.hpp"
#include "genesis/utils/math/compensated_sum.hpp"
 
#include <cassert>
#include <cmath>
#include <cstdint>
#include <iterator>
#include <limits>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace genesis {
namespace population {
 
// =================================================================================================
//     Diversity Pool Calculator
// =================================================================================================
 
class DiversityPoolCalculator
{
public:
 
    // -------------------------------------------------------------------------
    //     Typedefs and Enums
    // -------------------------------------------------------------------------
 
    struct Result
    {
        // Results of the diversity statistics
        double theta_pi        = std::numeric_limits<double>::quiet_NaN();
        double theta_watterson = std::numeric_limits<double>::quiet_NaN();
        double tajima_d        = std::numeric_limits<double>::quiet_NaN();
    };
 
    // -------------------------------------------------------------------------
    //     Constructors and Rule of Five
    // -------------------------------------------------------------------------
 
    DiversityPoolCalculator( DiversityPoolSettings const& settings, size_t pool_size )
        : settings_( settings )
        , pool_size_( pool_size )
    {
        if( settings.min_count == 0 ) {
            throw std::invalid_argument(
                "In DiversityPoolCalculator, settings.min_count == 0 is not allowed."
            );
        }
        if( pool_size == 0 ) {
            throw std::invalid_argument(
                "In DiversityPoolCalculator, pool_size == 0 is not allowed."
            );
        }
    }
    ~DiversityPoolCalculator() = default;
 
    DiversityPoolCalculator( DiversityPoolCalculator const& ) = default;
    DiversityPoolCalculator( DiversityPoolCalculator&& )      = default;
 
    DiversityPoolCalculator& operator= ( DiversityPoolCalculator const& ) = default;
    DiversityPoolCalculator& operator= ( DiversityPoolCalculator&& )      = default;
 
    using self_type = DiversityPoolCalculator;
 
    // -------------------------------------------------------------------------
    //     Settings
    // -------------------------------------------------------------------------
 
    self_type& only_passing_samples( bool value )
    {
        only_passing_samples_ = value;
        return *this;
    }
 
    bool only_passing_samples() const
    {
        return only_passing_samples_;
    }
 
    self_type& enable_theta_pi( bool value )
    {
        enable_theta_pi_ = value;
        return *this;
    }
 
    bool enable_theta_pi() const
    {
        return enable_theta_pi_;
    }
 
    self_type& enable_theta_watterson( bool value )
    {
        enable_theta_watterson_ = value;
        return *this;
    }
 
    bool enable_theta_watterson() const
    {
        return enable_theta_watterson_;
    }
 
    self_type& enable_tajima_d( bool value )
    {
        enable_tajima_d_ = value;
        return *this;
    }
 
    bool enable_tajima_d() const
    {
        return enable_tajima_d_;
    }
 
    // -------------------------------------------------------------------------
    //     Calculator Functions
    // -------------------------------------------------------------------------
 
    void reset()
    {
        theta_pi_sum_.reset();
        theta_watterson_sum_.reset();
        filter_stats_.clear();
        empirical_min_read_depth_ = std::numeric_limits<size_t>::max();
    }
 
    void process( SampleCounts const& sample )
    {
        // We first check if we need to process this sample at all.
        // That is either the case if we process _all_ samples (!only_passing_samples_)
        // or if we only consider the passing ones and the sample is passing.
        // We assume that the Variant::status has already been checked before calling this,
        // for instance by the DiversityPoolProcessor.
        ++filter_stats_[sample.status.get()];
        if( ! sample.status.passing() ) {
            return;
        }
 
        // Now run the calculations that we need.
        double tp = std::numeric_limits<double>::quiet_NaN();
        double tw = std::numeric_limits<double>::quiet_NaN();
        if( enable_theta_pi_ || enable_tajima_d_ ) {
            tp = theta_pi_pool( settings_, pool_size_, sample );
            if( std::isfinite( tp )) {
                theta_pi_sum_ += tp;
            }
            // assert( std::isfinite( tp ) );
        }
        if( enable_theta_watterson_ || enable_tajima_d_ ) {
            tw = theta_watterson_pool( settings_, pool_size_, sample );
            if( std::isfinite( tw )) {
                theta_watterson_sum_ += tw;
            }
            // assert( std::isfinite( tw ) );
        }
 
        // We want to keep track of the minimum read depth of the data that we are processing.
        // This is only needed when using TajimaDenominatorPolicy::kEmpiricalMinReadDepth,
        // but cheap enough to just always keep track of here.
        auto const cov = nucleotide_sum( sample );
        if( cov > 0 && cov < empirical_min_read_depth_ ) {
            empirical_min_read_depth_ = cov;
        }
    }
 
    Result get_result( double window_avg_denom ) const
    {
        Result result;
        if( enable_theta_pi_ ) {
            result.theta_pi = theta_pi_sum_.get() / window_avg_denom;
        }
        if( enable_theta_watterson_ ) {
            result.theta_watterson = theta_watterson_sum_.get() / window_avg_denom;
        }
        if( enable_tajima_d_ ) {
            // Yet another problem in PoPoolation: For the |W| window size in the denominator
            // of Tajima's D, they use the number of SNPs in that window, which might or might not
            // be correct - we will have to figure this out. There is a chance that this is correct,
            // but it could also be that we want to use the the number of _all_ valid positions
            // (the ones that passed all filters, including any invariant positions) here again.
            // For now, we follow their approach, but might leave this to fix later.
            double const tajimas_window_avg_denom = filter_stats_[SampleCountsFilterTag::kPassed];
 
            // Potential fix:
            // auto tajimas_window_avg_denom = window_avg_denom;
            // if( settings_.tajima_denominator_policy == TajimaDenominatorPolicy::kWithPopoolationBugs ) {
            //     tajimas_window_avg_denom = filter_stats_[SampleCountsFilterTag::kPassed];
            // }
 
            result.tajima_d = tajima_d_pool(
                settings_,
                theta_pi_sum_.get(), theta_watterson_sum_.get(),
                pool_size_, tajimas_window_avg_denom, empirical_min_read_depth_
            );
        }
        return result;
    }
 
    SampleCountsFilterStats get_filter_stats() const
    {
        return filter_stats_;
    }
 
    // -------------------------------------------------------------------------
    //     Private Members
    // -------------------------------------------------------------------------
 
private:
 
    // Settings
    DiversityPoolSettings settings_;
    size_t pool_size_ = 0;
 
    bool only_passing_samples_   = true;
    bool enable_theta_pi_        = true;
    bool enable_theta_watterson_ = true;
    bool enable_tajima_d_        = true;
 
    // Data Accumulation
    utils::NeumaierSum theta_pi_sum_;
    utils::NeumaierSum theta_watterson_sum_;
    SampleCountsFilterStats filter_stats_;
 
    // Find the minimum empirical read depth that we see in the processed data.
    size_t empirical_min_read_depth_ = std::numeric_limits<size_t>::max();
};
 
} // namespace population
} // namespace genesis
 
#endif // include guard