Fast Hardy-Weinberg equilibrium simulation

Question

I was very bored over one of my breaks this year, so I built a Hardy-Weinberg equilibrium simulator for two unrelated alleles of the same gene. Hardy-Weinberg equilibrium is when there is no evolution, random mating, and a constant population size. The idea behind the simulation is to show the effects of genetic drift over time on small or large populations for a particular gene.

The goal of the project is to teach me how to use C (this is my first non-trivial program) and how to heavily optimize programs. The current best benchmark I have on my computer is 380µs per generation of 65536 ~= 170m matings per second on my 2.5 GHz Intel Core i5. In the core logic of the program, I don't believe there are any memory accesses, so it should be totally CPU bound. The project is built to run on computers with SSE4.2. It will probably run on computers without it, but I imagine it will be much slower. GCC is not required, but heavily recommended as the code becomes ~25% faster with GCC as opposed to clang. The project as a whole is on Github.

The idea behind my implementation is that for each mating, because there are only three attributes possible, (aa, ab or ba, and bb) we don't have to actually pick out each organism, we can just count how many there are of a certain type. Additionally, there are only two possibilities in a mating: both alleles from one parent or one allele from each parent. The number of times each occur is represented in the both_one_parent variable, which should be approximately normal. The current maximum number of organisms is 65536, so we can used an unsigned short to find one of the organisms and get its' allele. If there's only one parent, then we're done. If there are two parents, then we have to mix their genes, which is in the second for loop.

One of the keys to fast simulation is a fast random function, and the fastest I could find is the FastRand algorithm by Ivan Dimkovic. I benchmarked FastRand, the system random function and Intel's rdrand instruction at respectively 4300 bytes/µs, 790 bytes/µs, and 190 bytes/µs (on generating 1GB. If anyone knows of a better one I would appreciate it. FastRand works in C by initializing a FastRand struct using InitFastRand and then calling FastRand on it, which refills the FastRand.res buffer, which is theoretically an int[4] array but is actually a 128 bit register.

I am fairly new to C and to optimizing programs like this in general. I would appreciate any advice. Below is the core loop of the program logic, which probably consumes >99% of runtime. I slightly changed the FastRand from Ivan Dimkovic and my version is here.

#include <stdlib.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/time.h>
#include <time.h>
#include <errno.h>
#include <string.h>
#include <x86intrin.h>

#include "fast_rand.h"

void progress_generation(int thresh_aa, int thresh_ab, int thresh_bb, int next_members, int* result) {

// 0 | thresh_aa | thresh_ab | end = thresh_bb

int counts[5] = { 0,0,0,0,0 };
if (next_members != 65536) {
double offset = 65536 / next_members;
thresh_aa *= offset;
thresh_ab *= offset;
thresh_bb *= offset;
printf("next_members != 65536");
}

fastrand rand_index = InitFastRand();
fastrand rand_choice = InitFastRand();
/*
* This is a little bit complicated. In the old way, the choice over whether to have both bits
* come from one parent or not was made more or less on the fly, which costs a shr and most
* importantly, an evil evil branch on every iteration. This branch was so awful because if the
* random number generator is good, it is unpredictable, which means you have a branch miss half
* of the time. Instead, we can refactor the decision making so it happens up here. What that means
* is that we decide ahead of time how many have both bits come from one parent and how many
* have one bit from each parent. This means that we don't have a branch in a loop here - there's
* only the loop branch, which should have >99% prediction. This should make the code much faster.
*/
int both_one_parent = 0;
for (int i = 0; i < (next_members >> 7); i += 1) { // FastRand is 128 bits = 2^7 members per fastrand
FastRand(&rand_choice);
both_one_parent += __builtin_popcountll(*(rand_choice.res)) +
__builtin_popcountll(*(rand_choice.res + 2)); // hopefully the compiler gets it
}
unsigned short *short_res = (unsigned short *)rand_index.res;
// normally rand_index.res would be typed as an int *, though really it's a 128 bit register
// location. Unsigned shorts cover the 0-65536 range. I could handle ints also, but it would be
// something like half as fast.

// For ones where both bits come from one parent
for (int i = 0; i < (both_one_parent >> 3); i++) {
FastRand(&rand_index);
for (int k = 0; k < 8; k++) {
unsigned short firstIndex = short_res[k];
counts[(firstIndex < thresh_aa) + (firstIndex < thresh_ab)] += 1;
// 0 = bb, 1 = ab, 2 = aa
}
}
counts[3] = counts[0];
counts[0] = counts[2];
counts[2] = 0;
// Counts output: 0: aa 1: ab 2: ab 3: bb 4: bb , so we have to fiddle around a bit

for (int i = 0; i < ((next_members - both_one_parent) >> 2); i++) {
// >> 2 and not 3 because the cycle only gives 4 results, not 8, because it needs 32 bits
FastRand(&rand_index);
for (int k = 0; k < 8; k++) {
unsigned short firstIndex = short_res[k];
k += 1;
unsigned short secondIndex = short_res[k];
char allele = 0;

// Three tests for first bit: firstIndex > thresh_ab in which case it's 1. firstIndex < thresh_aa in which case it's 1. threst_aa < firstIndex < thresh_ab 50% chance
if (firstIndex > thresh_ab) { // if it's a bb, then both alleles are b so result is b
allele = 1;
}
else if (firstIndex > thresh_aa) {
allele = secondIndex & 1; // subtle small error. change to secondIndex?
}

if (secondIndex > thresh_ab) { // We know second one is BB
counts[allele + 2] += 1;
}

else if (secondIndex > thresh_aa) { // Second one is AB
counts[allele + (firstIndex & 2)] += 1; // This used to be firstIndex. This caused an insidious bug where aa and bb were relatively favored 10:12:10 when they should be 8:16:8
// This used to be allele+secondIndex&2 and the &2 operated after the + so it was always 0 or 2
}

else { // second one is AA
counts[allele] += 1;
}
}
}

// 0: aa 1: ab 2: ab 3: bb 4: bb
result[0] = counts[0]; // aa
result[1] = counts[1] + counts[2]; // ab
result[2] = counts[3] + counts[4]; // bb
}

void initialize_generation(int number, int* gen) { // 14 milliseconds faster on generating 1.6 million organisms lol. 224 organisms / microsecond is good
   // AA consistently has ~17.5k while BB consistently has ~15k.
   // This is equivalent to a normal distribution. There's probably some normal distribution function I could use that's faster.
int num_aa = 0;
int num_ab = 0;
int num_bb = 0;

for (int i = 0; i < number / 8; i++) {  // 8 instead of 16 because the upper bit is always 0.
// It would be more efficient to make this 15 or something but who cares
int n = random();

#pragma clang loop unroll(full)
for (int j = 0; j < 8; j++) { // This needs to be unrolled
char organism = n & 3;
if (organism == 0) {
num_aa += 1;
}
else if (organism == 3) {
num_bb += 1;
}
else {
num_ab += 1;
}
n >>= 2;
}
}

gen[0] = num_aa;
gen[1] = num_ab;
gen[2] = num_bb;
}

int main(int argc, char **argv) {
int num_organisms = 0;
int num_generations = 0;
if (argc < 3) {
num_organisms = 65536;
num_generations = 1000;
}
else {
num_organisms = atoi(argv[1]);
num_generations = atoi(argv[2]);
if (__builtin_popcount(num_organisms) != 1) {
printf("num_organisms must be a multiple of 2");
exit(1);
}
}
printf("Simulating %d organisms for %d generations
",num_organisms,num_generations);
srandomdev(); // Seeds random() using information from /dev/random
int initial_values[3];
initialize_generation(num_organisms, initial_values);
int thresh_aa = initial_values[0];
int thresh_ab = thresh_aa + initial_values[1];
int thresh_bb = thresh_ab + initial_values[2];
int **results = (int **)malloc(num_generations * 3 * sizeof(int));
int result[3] = { 0,0,0 };

int time_takens[20];
int average = 0;
for (int i = 0; i < 20; i++) {
struct timeval start;
struct timeval end;
gettimeofday(&start, NULL);
for (int j = 0; j < 1000; j++) {
progress_generation(thresh_aa, thresh_ab, thresh_bb, num_organisms, result);
thresh_aa = result[0];
thresh_ab = thresh_aa + result[1];
thresh_bb = thresh_ab + result[2];
}
gettimeofday(&end, NULL);
int time_taken = (1000000 * end.tv_sec + end.tv_usec) - (1000000 * start.tv_sec + start.tv_usec);
printf("%d microseconds     ", time_taken);
printf("aa: %d      ab:%d   bb:%d
",result[0],result[1],result[2]);
time_takens[i] = time_taken;
average += time_taken;

initialize_generation(num_organisms, initial_values);
thresh_aa = initial_values[0];
thresh_ab = thresh_aa + initial_values[1];
thresh_bb = thresh_ab + initial_values[2];
}
printf("Took on average %d microseconds per 1000 generationss or %d microseconds per generation
",average/20,average/20000);
}

Answer 1

Indent your code to show its nesting. Although all whitespace is identical as far as the compiler is concerned, it conveys a lot of information to the more important (human) reader.

---

It's not legal for a string literal to contain a newline character like this:

printf("Simulating %d organisms for %d generations
",num_organisms,num_generations);

We need to write that as

printf("Simulating %d organisms for %d generations\n",
       num_organisms, num_generations);

---

The code would be more reviewable if it supported use of standard random-number generation (particularly useful on platforms where your FastRand library is unavailable).

---

There are some dubious implicit conversions between signed and unsigned types - enable more compiler warnings to catch these. For example:

202364.c:76:49: warning: conversion to ‘long long unsigned int’ from ‘int’ may change the sign of the result [-Wsign-conversion]
   76 |         both_one_parent += __builtin_popcountll(*(rand_choice.res)) +
      |                                                 ^~~~~~~~~~~~~~~~~~
202364.c:77:34: warning: conversion to ‘long long unsigned int’ from ‘int’ may change the sign of the result [-Wsign-conversion]
   77 |             __builtin_popcountll(*(rand_choice.res + 2)); // hopefully the compiler gets it
      |                                  ^~~~~~~~~~~~~~~~~~~~~~

Those look obfuscated anyway - what's wrong with the usual [] operator for array access?

---

time_takens and results are never used.

---

Don't cast the pointer returned from malloc(). Its type is void*, which implicitly converts to any other object pointer.

This calculation looks incorrect:

int **results = (int **)malloc(num_generations * 3 * sizeof(int));

Since we're creating an array of pointers, then we need to be allocating a multiple of sizeof (int*), not sizeof (int). Prefer to use the pointed-to object as sizeof argument when allocating (i.e. sizeof *results here) - that makes its correctness more obvious, rather than requiring the user to check the type manually.

---

Be careful with comments like this:

// Unsigned shorts cover the 0-65536 range.

You probably have "at least" in your head there, but the next reader might misinterpret that and be misled by it.

---

Careful with this calculation:

int time_taken = (1000000 * end.tv_sec + end.tv_usec) - (1000000 * start.tv_sec + start.tv_usec);

That certainly overflows a minimal int (guaranteed range at least -32767 to 32767), and probably overflows larger int (after all, that's at least partly why struct timeval has separate fields for these). Signed integer overflow is UB. There's no need for this anyway, since timersub() is available:

 struct timeval duration;
 timersub(&end, &start, &duration);
 printf("%lu%06lu microseconds     aa: %d      ab:%d   bb:%d",
        (unsigned long)duration.tv_sec, (unsigned long)duration.tv_usec,
        result[0],result[1],result[2]);

Using wall-clock time for the calculation is problematic, as it's not necessarily a steady clock. Consider instead using an interface such as POSIX clock_gettime() which allows access to monotonic clocks such as (on Linux) CLOCK_MONOTONIC_RAW. As this seems to be used as a kind of primitive benchmarking, perhaps Linux CLOCK_PROCESS_CPUTIME_ID is more apt?