Skip to content

Performance audit for genome_array module #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 20, 2024
Merged

Performance audit for genome_array module #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions src/common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,13 @@ pub struct ParentalGenome<'a> {
pub genome: usize,
}

#[derive(Debug)]
pub struct ParentalGenomeMut<'a> {
pub mutations: &'a mut [u32],
pub current_mutation_index: usize,
pub genome: usize,
}

// NOTE: we use usize::MAX to indicate a
// "no genome" state. Production
// code would do better.
Expand Down
258 changes: 201 additions & 57 deletions src/genome_array.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,11 @@ use forrustts::prelude::*;

use crate::common::generate_mutations;
use crate::common::generate_offspring_genome;
use crate::common::set_fixation_counts_to_zero;
use crate::common::DiploidGenome;
use crate::common::Mutation;
use crate::common::MutationRange;
use crate::common::ParentalGenome;
use crate::common::ParentalGenomeMut;
use crate::common::SimParams;

#[derive(Debug, Default)]
Expand Down Expand Up @@ -40,6 +40,24 @@ impl HaploidGenomes {
}
}

fn get_genome_mut(&mut self, genome: usize) -> ParentalGenomeMut {
if genome != usize::MAX {
let index_range = self.genome_spans[genome];
let mutations = &mut self.mutations[index_range.start..index_range.stop];
ParentalGenomeMut {
mutations,
current_mutation_index: 0,
genome,
}
} else {
ParentalGenomeMut {
mutations: &mut [],
current_mutation_index: 0,
genome,
}
}
}

fn add_range(&mut self, range: MutationRange) -> usize {
if range.start == range.stop {
usize::MAX
Expand Down Expand Up @@ -68,6 +86,50 @@ fn get_parental_genomes(
)
}

fn compact_mutations(pop: &mut DiploidPopulation) {
let new_indexes = {
let mut temp = vec![];
let mut i = 0_u32;
temp.resize_with(pop.mutations.len(), || {
let j = i;
i += 1;
j
});
let twon = 2 * (pop.individuals.len() as u32);
let (mut left, _): (Vec<u32>, Vec<u32>) = temp.into_iter().partition(|&index| {
let mcount = *unsafe { pop.mutation_counts.get_unchecked(index as usize) };
mcount > 0 && mcount < twon
});
left.sort_unstable_by_key(|&m| {
unsafe { pop.mutations.get_unchecked(m as usize) }.position()
});
left
};
let reindex = {
let mut temp = vec![u32::MAX; pop.mutations.len()];
for i in 0..new_indexes.len() {
temp[new_indexes[i] as usize] = i as u32;
}
temp
};

// remap the genomes
for m in &mut pop.genomes.mutations {
*m = reindex[*m as usize];
}

let mut mutations = vec![];
std::mem::swap(&mut mutations, &mut pop.mutations);
// convert all mutations into options
let mut mutations = mutations.into_iter().map(Some).collect::<Vec<_>>();
let mut mutation_counts = vec![];
for i in new_indexes {
mutation_counts.push(pop.mutation_counts[i as usize]);
pop.mutations.push(mutations[i as usize].take().unwrap());
}
std::mem::swap(&mut mutation_counts, &mut pop.mutation_counts);
}

/// When will the borrow checker hate this?
pub struct DiploidPopulation {
genomes: HaploidGenomes,
Expand Down Expand Up @@ -97,12 +159,18 @@ impl DiploidPopulation {
}

/// Generate indexes of extinct mutations
#[inline(never)]
fn mutation_recycling(&self) -> Vec<usize> {
// NOTE: the Fn generic is for API flexibility
// It is not always the case that we want to remove fixations:
// * fixations affect phenotype (unlike popgen models of multiplicative fitness)
// * multiple hits can make mutations "un-fixed" at a site
fn mutation_recycling<F>(&self, f: F) -> Vec<usize>
where
F: Fn((usize, &u32)) -> Option<usize>,
{
self.mutation_counts
.iter()
.enumerate()
.filter_map(|(index, value)| if value == &0 { Some(index) } else { None })
.filter_map(f)
.collect::<Vec<usize>>()
}

Expand Down Expand Up @@ -133,64 +201,129 @@ impl DiploidPopulation {
}
}

fn erase(mutation_counts: &[u32], twon: u32, mutations: &mut [u32]) -> usize {
let mut rv = 0;

// We only call this if we KNOW that the first mutation is a fixation,
// so we can start to index from 1.
for i in 1..mutations.len() {
if *unsafe { mutation_counts.get_unchecked(mutations[i] as usize) } != twon {
*unsafe { mutations.get_unchecked_mut(rv) } = mutations[i];
rv += 1;
}
}
rv
}

// This normally wouldn't be pub,
// but should be unit tested anyways.
#[inline(never)]
fn fixation_removal_check(mutation_counts: &[u32], twon: u32, output: &mut HaploidGenomes) -> bool {
fn fixation_removal_check(mutation_counts: &[u32], twon: u32, input: &mut HaploidGenomes) -> bool {
if mutation_counts.iter().any(|m| *m == twon) {
// NOTE: using this as an assert
// eliminates all performance
// benefits of the unsafe code below
debug_assert!(output
.mutations
.iter()
.all(|m| (*m as usize) < mutation_counts.len()));
let x = output.mutations.len();
// SAFETY: None in general. :(
// In the specific case of this function not being public,
// genomes only contain mutation keys generated
// as they are added to the pop's mutation array.
// As the same time, we make sure that mutation_counts'
// length is also correct.
//
// This is a sticky issue:
// 1. The perf gain of unsafe here is quite big.
// 2. But it is gonna be hard to encapsulate things
// to the point where we feel generally good about saftety
// 3. Therefore, it seems that the function itself should
// be marked unsafe in the absence of such encapsulation.
output
.mutations
.retain(|m| *unsafe { mutation_counts.get_unchecked(*m as usize) } < twon);
let delta = x - output.mutations.len();
assert_eq!(delta % output.genome_spans.len(), 0);

let delta_per_genome = delta / output.genome_spans.len();

// NOTE: could be SIMD later
output
.genome_spans
.iter_mut()
.enumerate()
.for_each(|(i, h)| {
let c = h.start;
h.start -= i * delta_per_genome;
if i > 0 {
assert!(c > h.start);
}
assert!(h.start < output.mutations.len());
h.stop -= (i + 1) * delta_per_genome;
assert!(h.stop >= h.start);
assert!(
h.stop <= output.mutations.len(),
"{h:?} {}",
output.mutations.len()
);
});
let first_fixation = {
let genome = input.get_genome(0);
let index = genome
.mutations
.iter()
.position(|&k| mutation_counts[k as usize] == twon)
.unwrap();
genome.mutations[index]
};
//let mut start = 0_usize;
for i in 0..input.genome_spans.len() {
let genome = input.get_genome_mut(i);
let first_fixation_position = genome
.mutations
.iter()
.position(|&m| m == first_fixation)
.unwrap();
let tpoint = erase(
mutation_counts,
twon,
&mut genome.mutations[first_fixation_position..],
);
input.genome_spans[i].stop -=
(genome.mutations.len() - tpoint - first_fixation_position);
//output
// .mutations
// .extend_from_slice(&genome.mutations[..first_fixation_position]);
//output.mutations.extend(
// genome
// .mutations
// .iter()
// .skip(first_fixation_position + 1)
// .filter(|&&m| {
// let count = *unsafe { mutation_counts.get_unchecked(m as usize) };
// count < twon
// }),
//);
// NOTE: this push can contradict what "add_range" is doing
// in the impl of HaploidGenomes.
//output.genome_spans.push(MutationRange {
// start,
// stop: output.mutations.len(),
//});
//start = output.mutations.len();
}
true
} else {
false
}

//if mutation_counts.iter().any(|m| *m == twon) {
// // NOTE: using this as an assert
// // eliminates all performance
// // benefits of the unsafe code below
// debug_assert!(output
// .mutations
// .iter()
// .all(|m| (*m as usize) < mutation_counts.len()));
// let x = output.mutations.len();
// // SAFETY: None in general. :(
// // In the specific case of this function not being public,
// // genomes only contain mutation keys generated
// // as they are added to the pop's mutation array.
// // As the same time, we make sure that mutation_counts'
// // length is also correct.
// //
// // This is a sticky issue:
// // 1. The perf gain of unsafe here is quite big.
// // 2. But it is gonna be hard to encapsulate things
// // to the point where we feel generally good about saftety
// // 3. Therefore, it seems that the function itself should
// // be marked unsafe in the absence of such encapsulation.
// output
// .mutations
// .retain(|m| *unsafe { mutation_counts.get_unchecked(*m as usize) } < twon);
// let delta = x - output.mutations.len();
// assert_eq!(delta % output.genome_spans.len(), 0);

// let delta_per_genome = delta / output.genome_spans.len();

// // NOTE: could be SIMD later
// output
// .genome_spans
// .iter_mut()
// .enumerate()
// .for_each(|(i, h)| {
// let c = h.start;
// h.start -= i * delta_per_genome;
// if i > 0 {
// assert!(c > h.start);
// }
// assert!(h.start < output.mutations.len());
// h.stop -= (i + 1) * delta_per_genome;
// assert!(h.stop >= h.start);
// assert!(
// h.stop <= output.mutations.len(),
// "{h:?} {}",
// output.mutations.len()
// );
// });
// true
//} else {
// false
//}
}

#[derive(Copy, Clone, Debug)]
Expand Down Expand Up @@ -344,9 +477,20 @@ pub fn evolve_pop(params: SimParams, genetic_map: GeneticMap) -> Option<DiploidP
2 * params.num_individuals,
&mut pop.genomes,
) {
set_fixation_counts_to_zero(2 * params.num_individuals, &mut pop.mutation_counts);
};
queue = pop.mutation_recycling();
//std::mem::swap(&mut pop.genomes, &mut offspring_genomes);
//offspring_genomes.clear();
}
queue = pop.mutation_recycling(|(index, &value)| {
if value == 0 || value == 2 * params.num_individuals {
Some(index)
} else {
None
}
});
}
if generation % 100 == 0 {
compact_mutations(&mut pop);
queue = vec![];
}
}
Some(pop)
Expand Down