Skip to content

fixation removal #1

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
Apr 13, 2023
Merged

fixation removal #1

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
99 changes: 48 additions & 51 deletions src/common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,34 +119,49 @@ pub fn generate_mutations(
rv
}

fn get_partition_index(genome: &[usize], mutations: &[Mutation], position: Position) -> usize {
genome.partition_point(|k| mutations[*k].position() < position)
}

fn extend_from_slice(
genome: &[usize],
mutations: &[Mutation],
position: Position,
offspring_haplotypes: &mut Vec<usize>,
) -> usize {
let n = get_partition_index(genome, mutations, position);
offspring_haplotypes.extend_from_slice(&genome[..n]);
n
}

fn update_genome(
mutations: &[Mutation],
position: Position,
parent_genome: &mut ParentalGenome,
offspring_haplotypes: &mut Vec<usize>,
) {
parent_genome.current_mutation_index += extend_from_slice(
&parent_genome.mutations[parent_genome.current_mutation_index..],
mutations,
position,
offspring_haplotypes,
)
}

#[inline(never)]
fn merge_mutations(
mutations: &[Mutation],
mutation_counts: &[u32],
current_total_size: u32,
new_mutations: &[usize],
offspring_haplotypes: &mut Vec<usize>,
current_genome: &mut ParentalGenome,
) {
for m in new_mutations.iter() {
let mpos = mutations[*m].position();
let n = current_genome.mutations[current_genome.current_mutation_index..]
.iter()
.take_while(|mutation| mutations[**mutation].position() < mpos)
.inspect(|x| {
if mutation_counts[**x] < current_total_size {
offspring_haplotypes.push(**x);
}
})
.count();
update_genome(mutations, mpos, current_genome, offspring_haplotypes);
offspring_haplotypes.push(*m);
current_genome.current_mutation_index += n;
}
for m in &current_genome.mutations[current_genome.current_mutation_index..] {
if mutation_counts[*m] < current_total_size {
offspring_haplotypes.push(*m);
}
}
offspring_haplotypes
.extend_from_slice(&current_genome.mutations[current_genome.current_mutation_index..]);
}

// This has had a lot of refactoring and still
Expand All @@ -155,8 +170,6 @@ fn merge_mutations(
pub fn generate_offspring_genome(
genomes: (ParentalGenome, ParentalGenome),
mutations: &[Mutation],
mutation_counts: &[u32],
current_total_size: u32,
new_mutations: Vec<usize>,
breakpoints: &[Breakpoint],
offspring_mutations: &mut Vec<usize>,
Expand All @@ -176,45 +189,24 @@ pub fn generate_offspring_genome(
.iter()
.take_while(|k| mutations[**k].position() < bpos)
.inspect(|k| {
// TODO: this should be abstracted out and the k-th
// mutation position cached
current_genome.current_mutation_index += current_genome.mutations
[current_genome.current_mutation_index..]
.iter()
.take_while(|gk| mutations[**gk].position() < mutations[**k].position())
.inspect(|gk| {
if mutation_counts[**gk] < current_total_size {
offspring_mutations.push(**gk);
}
})
.count();
let mpos = mutations[**k].position();
update_genome(mutations, mpos, &mut current_genome, offspring_mutations);
offspring_mutations.push(**k);
})
.count();
current_genome.current_mutation_index += current_genome.mutations
[current_genome.current_mutation_index..]
.iter()
.take_while(|gk| mutations[**gk].position() < bpos)
.inspect(|gk| {
if mutation_counts[**gk] < current_total_size {
offspring_mutations.push(**gk);
}
})
.count();
update_genome(mutations, bpos, &mut current_genome, offspring_mutations);

// Advance other genome
other_genome.current_mutation_index += other_genome.mutations
[other_genome.current_mutation_index..]
.iter()
.take_while(|gk| mutations[**gk].position() < bpos)
.count();
other_genome.current_mutation_index += get_partition_index(
&other_genome.mutations[other_genome.current_mutation_index..],
mutations,
bpos,
);

std::mem::swap(&mut current_genome, &mut other_genome);
}
merge_mutations(
mutations,
mutation_counts,
current_total_size,
&new_mutations[mut_index..],
offspring_mutations,
&mut current_genome,
Expand All @@ -223,6 +215,14 @@ pub fn generate_offspring_genome(
MutationRange { start, stop }
}

#[inline(never)]
pub fn set_fixation_counts_to_zero(twon: u32, mutation_counts: &mut [u32]) {
mutation_counts
.iter_mut()
.filter(|m| **m == twon)
.for_each(|m| *m = 0);
}

#[cfg(test)]
mod test_create_offspring_genome {
use super::*;
Expand Down Expand Up @@ -378,12 +378,9 @@ mod test_create_offspring_genome {
.windows(2)
.all(|w| mutations[w[0]].position() <= mutations[w[1]].position()),);
let mut offspring_genomes = Vec::<usize>::new();
let mutation_counts = vec![0_u32; mutations.len() + new_mutations.len()];
let range = generate_offspring_genome(
(parent1_genome, parent2_genome),
&mutations,
&mutation_counts,
u32::MAX,
new_mutations.clone(),
&breakpoints,
&mut offspring_genomes,
Expand Down
51 changes: 47 additions & 4 deletions src/first_pass.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ 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;
Expand Down Expand Up @@ -190,6 +191,44 @@ impl SimParams {
}
}

// 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 Haplotypes) -> bool {
if mutation_counts.iter().any(|m| *m == twon) {
let x = output.mutations.len();
output.mutations.retain(|m| mutation_counts[*m] < twon);
let delta = x - output.mutations.len();
assert_eq!(delta % output.haplotypes.len(), 0);

let delta_per_genome = delta / output.haplotypes.len();
assert_eq!(
delta_per_genome,
mutation_counts.iter().filter(|m| **m == twon).count()
);

// NOTE: could be SIMD later
output.haplotypes.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
}
}

// A proper implementation
// would be generic over "generating mutations"
#[inline(never)]
Expand Down Expand Up @@ -246,8 +285,6 @@ pub fn evolve_pop_with_haplotypes(
let range = generate_offspring_genome(
genomes,
&pop.mutations,
&pop.mutation_counts,
2 * params.num_individuals,
mutations,
genetic_map.breakpoints(),
&mut offspring_haplotypes.mutations,
Expand Down Expand Up @@ -276,8 +313,6 @@ pub fn evolve_pop_with_haplotypes(
let range = generate_offspring_genome(
genomes,
&pop.mutations,
&pop.mutation_counts,
2 * params.num_individuals,
mutations,
genetic_map.breakpoints(),
&mut offspring_haplotypes.mutations,
Expand All @@ -292,6 +327,14 @@ pub fn evolve_pop_with_haplotypes(
std::mem::swap(&mut pop.individuals, &mut offspring);
offspring.clear();
pop.count_mutations();

if fixation_removal_check(
&pop.mutation_counts,
2 * params.num_individuals,
&mut pop.haplotypes,
) {
set_fixation_counts_to_zero(2 * params.num_individuals, &mut pop.mutation_counts);
};
// for h in &pop.haplotypes.haplotypes {
// assert!(pop.haplotypes.mutations[h.start..h.stop]
// .windows(2)
Expand Down
32 changes: 25 additions & 7 deletions src/fwdpp_copy.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use rand::prelude::SeedableRng;

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::ParentalGenome;
Expand Down Expand Up @@ -153,13 +154,29 @@ enum NewGenomeType {
New(HaploidGenome),
}

#[inline(never)]
fn remove_fixations_from_extant_genomes(
twon: u32,
mutation_counts: &[u32],
genomes: &mut [HaploidGenome],
) -> bool {
if mutation_counts.iter().any(|c| *c == twon) {
genomes
.iter_mut()
.filter(|g| g.count > 0)
.for_each(|g| g.mutations.retain(|k| mutation_counts[*k] < twon));
true
} else {
false
}
}

#[inline(never)]
fn update_genomes(
genomes: (usize, usize),
mutations: Vec<usize>,
genetic_map: &GeneticMap,
pop: &mut DiploidPopulation,
fixation_threshold: u32,
genome_queue: &mut Vec<usize>,
temporary_mutations: &mut Vec<usize>,
) -> usize {
Expand Down Expand Up @@ -200,8 +217,6 @@ fn update_genomes(
generate_offspring_genome(
genomes,
&pop.mutations,
&pop.mutation_counts,
fixation_threshold,
mutations,
genetic_map.breakpoints(),
temporary_mutations,
Expand All @@ -212,8 +227,6 @@ fn update_genomes(
generate_offspring_genome(
genomes,
&pop.mutations,
&pop.mutation_counts,
fixation_threshold,
mutations,
genetic_map.breakpoints(),
&mut genome.mutations,
Expand Down Expand Up @@ -284,7 +297,6 @@ pub fn evolve_pop_with_haplotypes(
mutations,
&genetic_map,
&mut pop,
2 * params.num_individuals,
&mut genome_queue,
&mut temporary_mutations,
);
Expand Down Expand Up @@ -319,7 +331,6 @@ pub fn evolve_pop_with_haplotypes(
mutations,
&genetic_map,
&mut pop,
2 * params.num_individuals,
&mut genome_queue,
&mut temporary_mutations,
);
Expand All @@ -332,6 +343,13 @@ pub fn evolve_pop_with_haplotypes(
// assert!(pop.haplotypes[i.second].count > 0);
//}
pop.count_mutations();
if remove_fixations_from_extant_genomes(
2 * params.num_individuals,
&pop.mutation_counts,
&mut pop.haplotypes,
) {
set_fixation_counts_to_zero(2 * params.num_individuals, &mut pop.mutation_counts);
}
offspring.clear();
}
Some(pop)
Expand Down