teju85/covid-pooled-testing.cpp

## covid-pooled-testing.cpp
// Compilation and running:
//  g++ -std=c++11 covid-pooled-testing.cpp && ./a.out

#include <cstdio>
#include <stdint.h>
#include <cmath>
#include <cstdlib>
#include <algorithm>

/**
 * @brief Computes number of ones in the input number
 * @param num the input number
 * @return the number of ones
 */
int __popcnt64(uint64_t num) {
  int cnt = 0;
  while (num > 0) {
    cnt += num & 0x1;
    num >>= 1;
  }
  return cnt;
}

/**
 * @brief Computes the probability of the occurence of a particular combination of
 *        bits in the input number
 * @param num      input number
 * @param prob     probability of a bit having a value 1
 * @param poolSize max number of bits to consider in the input number
 * @return the probability
 *
 * @note this assumes IID across the bits!
 */
double probability(uint64_t num, double prob, int poolSize) {
  auto numPositive = __popcnt64(num);
  return std::pow(prob, numPositive) *
    std::pow(1.0 - prob, poolSize - numPositive);
}

/**
 * @brief total number of tests needed to ascertain the results of each
 *        individual for a given infection pattern. ie. i-th bit in the input
 *        number having a value of 1 means that i-th individual is positive.
 * @param num      input number containing the infection pattern
 * @param poolSize number of individuals in the pool to be tested for
 * @return total number of tests needed
 */
int numTestsNeeded(uint64_t num, int poolSize) {
  if (num == 0 || (num == 1 && poolSize == 1)) return 1;
  int ps = poolSize >> 1;
  uint64_t psMask = (1ULL << ps) - 1ULL;
  auto left = (num >> ps) & psMask;
  auto right = num & psMask;
  return numTestsNeeded(left, ps) + numTestsNeeded(right, ps) + 1;
}

/**
 * @brief average number of tests needed to ascertain the results of a pool
 *        where averaging is done across the whole population and also across
 *        whole possibility of infection patterns in a pool. For a pool of size
 *        4, there will be 2^4 patterns.
 * @param prob     probability of a person getting infected
 * @param poolSize number of individuals in the pool to be tested for
 * @return average number of tests needed per pool
 */
double averageTestsNeeded(double prob, int poolSize) {
  auto totalTests = 0.0;
  auto cnt = 1ULL << poolSize;
  for (uint64_t i = 0LL; i < cnt; ++i) {
    auto p = probability(i, prob, poolSize);
    auto nTests = numTestsNeeded(i, poolSize);
    totalTests += nTests * p;
  }
  return totalTests;
}

/**
 * @brief average number of tests needed to ascertain the results of a pool
 *        when it is given that atleast one of the member in the pool is
 *        positive.
 * @param prob     probability of a person getting infected
 * @param poolSize number of individuals in the pool to be tested for
 * @return average number of tests needed per pool
 */
double averageTestsNeededIfAtleastOnePositive(double prob, int poolSize) {
  auto totalTests = 0.0, totalProb = 0.0;
  auto cnt = 1ULL << poolSize;
  for (uint64_t i = 1LL; i < cnt; ++i) {
    auto p = probability(i, prob, poolSize);
    auto nTests = numTestsNeeded(i, poolSize);
    totalProb += p;
    totalTests += nTests * p;
  }
  return totalTests / totalProb;
}

/**
 * @brief total number of *sequence* of tests needed to ascertain the results
 *        of each individual for a given infection pattern, when it is given
 *        that atleast one of them is positive.
 * @param num      input number containing the infection pattern
 * @param poolSize number of individuals in the pool to be tested for
 * @return total number of steps needed
 */
int numStepsNeeded(uint64_t num, int poolSize) {
  if (num == 0 || (num == 1 && poolSize == 1)) return 1;
  int ps = poolSize >> 1;
  uint64_t psMask = (1ULL << ps) - 1ULL;
  auto left = (num >> ps) & psMask;
  auto right = num & psMask;
  return std::max(numStepsNeeded(left, ps), numStepsNeeded(right, ps)) + 1;
}

/**
 * @brief average number of *sequence* of tests needed to ascertain the
 *        results of a pool when it is given that atleast one of the member
 *        in the pool is positive.
 * @param prob     probability of a person getting infected
 * @param poolSize number of individuals in the pool to be tested for
 * @return average number of steps needed per pool
 */
double averageStepsNeededIfAtleastOnePositive(double prob, int poolSize) {
  auto totalSteps = 0.0, totalProb = 0.0;
  auto cnt = 1ULL << poolSize;
  for (uint64_t i = 1LL; i < cnt; ++i) {
    auto p = probability(i, prob, poolSize);
    auto nTests = numStepsNeeded(i, poolSize);
    totalProb += p;
    totalSteps += nTests * p;
  }
  return totalSteps / totalProb;
}

constexpr auto TotalPopulation = 8000000000ULL;

void checkForOneCase(uint64_t totalInfected) {
  const auto infectionProbability = totalInfected * 1.0 / TotalPopulation;
  std::printf("Infection probability = %lf\n", infectionProbability);
  for (int ps = 2; ps <= 16; ps <<= 1) {
    auto nPosTests = averageTestsNeededIfAtleastOnePositive(
      infectionProbability, ps);
    auto nPosSteps = averageStepsNeededIfAtleastOnePositive(
      infectionProbability, ps);
    std::printf("  PoolSize=%2d  avgTestsNeeded=%le avgPositiveTests=%le"
                " avgStepsForPositives=%le\n", ps,
                averageTestsNeeded(infectionProbability, ps), nPosTests,
                nPosSteps);
  }
}

int main(int argc, char** argv) {
  for (uint64_t infected = 1000000ULL; infected <= 100000000ULL; infected += 10000000ULL)
    checkForOneCase(infected);
  return 0;
}
	// Compilation and running:
	// g++ -std=c++11 covid-pooled-testing.cpp && ./a.out

	#include <cstdio>
	#include <stdint.h>
	#include <cmath>
	#include <cstdlib>
	#include <algorithm>

	/**
	* @brief Computes number of ones in the input number
	* @param num the input number
	* @return the number of ones
	*/
	int __popcnt64(uint64_t num) {
	int cnt = 0;
	while (num > 0) {
	cnt += num & 0x1;
	num >>= 1;
	}
	return cnt;
	}

	/**
	* @brief Computes the probability of the occurence of a particular combination of
	* bits in the input number
	* @param num input number
	* @param prob probability of a bit having a value 1
	* @param poolSize max number of bits to consider in the input number
	* @return the probability
	*
	* @note this assumes IID across the bits!
	*/
	double probability(uint64_t num, double prob, int poolSize) {
	auto numPositive = __popcnt64(num);
	return std::pow(prob, numPositive) *
	std::pow(1.0 - prob, poolSize - numPositive);
	}

	/**
	* @brief total number of tests needed to ascertain the results of each
	* individual for a given infection pattern. ie. i-th bit in the input
	* number having a value of 1 means that i-th individual is positive.
	* @param num input number containing the infection pattern
	* @param poolSize number of individuals in the pool to be tested for
	* @return total number of tests needed
	*/
	int numTestsNeeded(uint64_t num, int poolSize) {
	if (num == 0 \|\| (num == 1 && poolSize == 1)) return 1;
	int ps = poolSize >> 1;
	uint64_t psMask = (1ULL << ps) - 1ULL;
	auto left = (num >> ps) & psMask;
	auto right = num & psMask;
	return numTestsNeeded(left, ps) + numTestsNeeded(right, ps) + 1;
	}

	/**
	* @brief average number of tests needed to ascertain the results of a pool
	* where averaging is done across the whole population and also across
	* whole possibility of infection patterns in a pool. For a pool of size
	* 4, there will be 2^4 patterns.
	* @param prob probability of a person getting infected
	* @param poolSize number of individuals in the pool to be tested for
	* @return average number of tests needed per pool
	*/
	double averageTestsNeeded(double prob, int poolSize) {
	auto totalTests = 0.0;
	auto cnt = 1ULL << poolSize;
	for (uint64_t i = 0LL; i < cnt; ++i) {
	auto p = probability(i, prob, poolSize);
	auto nTests = numTestsNeeded(i, poolSize);
	totalTests += nTests * p;
	}
	return totalTests;
	}

	/**
	* @brief average number of tests needed to ascertain the results of a pool
	* when it is given that atleast one of the member in the pool is
	* positive.
	* @param prob probability of a person getting infected
	* @param poolSize number of individuals in the pool to be tested for
	* @return average number of tests needed per pool
	*/
	double averageTestsNeededIfAtleastOnePositive(double prob, int poolSize) {
	auto totalTests = 0.0, totalProb = 0.0;
	auto cnt = 1ULL << poolSize;
	for (uint64_t i = 1LL; i < cnt; ++i) {
	auto p = probability(i, prob, poolSize);
	auto nTests = numTestsNeeded(i, poolSize);
	totalProb += p;
	totalTests += nTests * p;
	}
	return totalTests / totalProb;
	}

	/**
	* @brief total number of sequence of tests needed to ascertain the results
	* of each individual for a given infection pattern, when it is given
	* that atleast one of them is positive.
	* @param num input number containing the infection pattern
	* @param poolSize number of individuals in the pool to be tested for
	* @return total number of steps needed
	*/
	int numStepsNeeded(uint64_t num, int poolSize) {
	if (num == 0 \|\| (num == 1 && poolSize == 1)) return 1;
	int ps = poolSize >> 1;
	uint64_t psMask = (1ULL << ps) - 1ULL;
	auto left = (num >> ps) & psMask;
	auto right = num & psMask;
	return std::max(numStepsNeeded(left, ps), numStepsNeeded(right, ps)) + 1;
	}

	/**
	* @brief average number of sequence of tests needed to ascertain the
	* results of a pool when it is given that atleast one of the member
	* in the pool is positive.
	* @param prob probability of a person getting infected
	* @param poolSize number of individuals in the pool to be tested for
	* @return average number of steps needed per pool
	*/
	double averageStepsNeededIfAtleastOnePositive(double prob, int poolSize) {
	auto totalSteps = 0.0, totalProb = 0.0;
	auto cnt = 1ULL << poolSize;
	for (uint64_t i = 1LL; i < cnt; ++i) {
	auto p = probability(i, prob, poolSize);
	auto nTests = numStepsNeeded(i, poolSize);
	totalProb += p;
	totalSteps += nTests * p;
	}
	return totalSteps / totalProb;
	}

	constexpr auto TotalPopulation = 8000000000ULL;

	void checkForOneCase(uint64_t totalInfected) {
	const auto infectionProbability = totalInfected * 1.0 / TotalPopulation;
	std::printf("Infection probability = %lf\n", infectionProbability);
	for (int ps = 2; ps <= 16; ps <<= 1) {
	auto nPosTests = averageTestsNeededIfAtleastOnePositive(
	infectionProbability, ps);
	auto nPosSteps = averageStepsNeededIfAtleastOnePositive(
	infectionProbability, ps);
	std::printf(" PoolSize=%2d avgTestsNeeded=%le avgPositiveTests=%le"
	" avgStepsForPositives=%le\n", ps,
	averageTestsNeeded(infectionProbability, ps), nPosTests,
	nPosSteps);
	}
	}

	int main(int argc, char** argv) {
	for (uint64_t infected = 1000000ULL; infected <= 100000000ULL; infected += 10000000ULL)
	checkForOneCase(infected);
	return 0;
	}