diversity_pool_functions.hpp

Go to the documentation of this file.

#ifndef GENESIS_POPULATION_FUNCTION_DIVERSITY_POOL_FUNCTIONS_H_
#define GENESIS_POPULATION_FUNCTION_DIVERSITY_POOL_FUNCTIONS_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/variant_filter.hpp"
#include "genesis/population/filter/sample_counts_filter.hpp"
#include "genesis/population/function/functions.hpp"
#include "genesis/population/variant.hpp"
 
#include <cassert>
#include <cmath>
#include <cstdint>
#include <iterator>
#include <limits>
#include <string>
#include <type_traits>
#include <vector>
 
namespace genesis {
namespace population {
 
// =================================================================================================
//     Diversity Pool Settings
// =================================================================================================
 
enum class TajimaDenominatorPolicy
{
    kEmpiricalMinReadDepth,
 
    kProvidedMinReadDepth,
 
    kWithPopoolationBugs,
 
    kPoolsize,
 
    kUncorrected
};
 
struct DiversityPoolSettings
{
    size_t min_count = 0;
    size_t min_read_depth = 0;
    size_t max_read_depth = 0;
 
    TajimaDenominatorPolicy tajima_denominator_policy = TajimaDenominatorPolicy::kUncorrected;
};
 
// =================================================================================================
//     Theta Pi
// =================================================================================================
 
double heterozygosity( SampleCounts const& sample, bool with_bessel = false );
 
template<class ForwardIterator>
double theta_pi(
    ForwardIterator begin,
    ForwardIterator end,
    bool with_bessel = true,
    bool only_passing_samples = true
) {
    double pi_sum = 0.0;
    for( auto& it = begin; it != end; ++it ) {
        if( only_passing_samples && !it->status.passing() ) {
            continue;
        }
        pi_sum += heterozygosity( *it, with_bessel );
    }
    return pi_sum;
}
 
template<class ForwardIterator>
double theta_pi_within_pool(
    size_t poolsize,
    ForwardIterator begin,
    ForwardIterator end,
    bool only_passing_samples = true
) {
    auto const psb = static_cast<double>( poolsize ) / ( static_cast<double>( poolsize ) - 1.0 );
    double pi_sum = 0.0;
    for( auto& it = begin; it != end; ++it ) {
        if( only_passing_samples && !it->status.passing() ) {
            continue;
        }
 
        // Apply heterozygosity with Bessel's for nucleotide count, and Bessel's for pool size.
        pi_sum += heterozygosity( *it, true ) * psb;
    }
    return pi_sum;
}
 
double theta_pi_pool_denominator(
    DiversityPoolSettings const& settings,
    size_t poolsize,
    size_t nucleotide_count
);
 
template<class ForwardIterator>
double theta_pi_pool( // get_pi_calculator
    DiversityPoolSettings const& settings,
    size_t poolsize,
    ForwardIterator begin,
    ForwardIterator end,
    bool only_passing_samples = true
) {
    // PoPoolation variable names:
    // poolsize:  n
    // min_count: b
 
    double pi_sum = 0.0;
    for( auto& it = begin; it != end; ++it ) {
        if( only_passing_samples && !it->status.passing() ) {
            continue;
        }
 
        double const pi_snp = heterozygosity( *it, true );
        double const denom  = theta_pi_pool_denominator(
            settings, poolsize, nucleotide_sum( *it )
        );
        if( std::isfinite( pi_snp ) && std::isfinite( denom )) {
            pi_sum += ( pi_snp / denom );
        }
    }
    return pi_sum;
}
 
inline double theta_pi_pool(
    DiversityPoolSettings const& settings,
    size_t poolsize,
    SampleCounts const& sample
) {
    auto const h = heterozygosity( sample, true );
    auto const d = theta_pi_pool_denominator( settings, poolsize, nucleotide_sum( sample ));
    return h / d;
}
 
// =================================================================================================
//     Theta Watterson
// =================================================================================================
 
double theta_watterson_pool_denominator(
    DiversityPoolSettings const& settings,
    size_t poolsize,
    size_t nucleotide_count
);
 
template<class ForwardIterator>
double theta_watterson_pool( // get_theta_calculator
    DiversityPoolSettings const& settings,
    size_t poolsize,
    ForwardIterator begin,
    ForwardIterator end,
    bool only_passing_samples = true
) {
    // PoPoolation variable names:
    // poolsize:  n
    // min_count: b
 
    double sum = 0.0;
    for( auto& it = begin; it != end; ++it ) {
        if( only_passing_samples && !it->status.passing() ) {
            continue;
        }
 
        auto const denom = theta_watterson_pool_denominator(
            settings, poolsize, nucleotide_sum( *it )
        );
        if( std::isfinite( denom )) {
            sum += 1.0 / denom;
        }
    }
    return sum;
}
 
inline double theta_watterson_pool(
    DiversityPoolSettings const& settings,
    size_t poolsize,
    SampleCounts const& sample
) {
    return 1.0 / theta_watterson_pool_denominator( settings, poolsize, nucleotide_sum( sample ));
}
 
// =================================================================================================
//     Tajima's D Helper Functions
// =================================================================================================
 
double a_n( double n );
 
double b_n( double n );
 
double f_star( double a_n, double n );
 
double alpha_star( double n );
 
double beta_star( double n );
 
double n_base_matrix( size_t read_depth, size_t poolsize );
 
double n_base( size_t read_depth, size_t poolsize );
 
// =================================================================================================
//     Tajima's D
// =================================================================================================
 
double tajima_d_pool_denominator(
    DiversityPoolSettings const& settings,
    double theta,
    size_t poolsize,
    double window_avg_denom,
    size_t empirical_min_read_depth
);
 
inline double tajima_d_pool(
    DiversityPoolSettings const& settings,
    double theta_pi,
    double theta_watterson,
    size_t poolsize,
    double window_avg_denom,
    size_t empirical_min_read_depth
) {
    // Edge case, following what PoPoolation does in this situation.
    // Deactivated - we want a nan instead.
    // if( window_avg_denom == 0 ) {
    //     return 0.0;
    // }
 
    // We already have the two theta statistics given here, but need to compute the
    // denominator according to Kofler et al for pooled sequences.
    auto const denom = tajima_d_pool_denominator(
        settings, theta_watterson, poolsize, window_avg_denom, empirical_min_read_depth
    );
    return ( theta_pi - theta_watterson ) / denom;
}
 
template<class ForwardIterator>
double tajima_d_pool( // get_D_calculator
    DiversityPoolSettings const& settings,
    double theta_pi,
    double theta_watterson,
    size_t poolsize,
    ForwardIterator begin,
    ForwardIterator end,
    bool only_passing_samples = true
) {
    // If we need the empirical min read depth, compute it.
    // If not, we can skip this step.
    size_t empirical_min_read_depth = std::numeric_limits<size_t>::max();
    auto entry_count = std::distance( begin, end );
    if(
        only_passing_samples ||
        settings.tajima_denominator_policy == TajimaDenominatorPolicy::kEmpiricalMinReadDepth
    ) {
        entry_count = 0;
        for( auto& it = begin; it != end; ++it ) {
            if( only_passing_samples && !it->status.passing() ) {
                continue;
            }
 
            auto const cov = nucleotide_sum( *it );
            if( cov < empirical_min_read_depth ) {
                empirical_min_read_depth = cov;
            }
            ++entry_count;
        }
    }
 
    return tajima_d_pool(
        settings, theta_pi, theta_watterson, poolsize, entry_count, empirical_min_read_depth
    );
}
 
template<class ForwardIterator>
double tajima_d_pool( // get_D_calculator
    DiversityPoolSettings const& settings,
    size_t poolsize,
    ForwardIterator begin,
    ForwardIterator end,
    bool only_passing_samples = true
) {
    // First compute the two theta statistics, then call the other version of this function.
    auto const pi = theta_pi_pool( settings, poolsize, begin, end, only_passing_samples );
    auto const tw = theta_watterson_pool( settings, poolsize, begin, end, only_passing_samples );
    return tajima_d_pool( settings, pi, tw, poolsize, begin, end, only_passing_samples );
}
 
} // namespace population
} // namespace genesis
 
#endif // include guard