1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

use core::marker::PhantomData;

use necsim_core::{
    cogs::{coalescence_sampler::CoalescenceRngSample, CoalescenceSampler, Habitat, MathsCore},
    landscape::{IndexedLocation, Location},
    lineage::LineageInteraction,
};

use crate::cogs::lineage_store::{
    coherent::globally::singleton_demes::SingletonDemesHabitat,
    independent::IndependentLineageStore,
};

#[allow(clippy::module_name_repetitions)]
#[derive(Debug)]
#[cfg_attr(feature = "cuda", derive(rust_cuda::lend::LendRustToCuda))]
#[cfg_attr(feature = "cuda", cuda(free = "M", free = "H"))]
pub struct IndependentCoalescenceSampler<M: MathsCore, H: Habitat<M>>(PhantomData<(M, H)>);

impl<M: MathsCore, H: Habitat<M>> Default for IndependentCoalescenceSampler<M, H> {
    fn default() -> Self {
        Self(PhantomData::<(M, H)>)
    }
}

impl<M: MathsCore, H: Habitat<M>> Clone for IndependentCoalescenceSampler<M, H> {
    fn clone(&self) -> Self {
        Self(PhantomData::<(M, H)>)
    }
}

#[contract_trait]
impl<M: MathsCore, H: Habitat<M>> CoalescenceSampler<M, H, IndependentLineageStore<M, H>>
    for IndependentCoalescenceSampler<M, H>
{
    #[must_use]
    // #[debug_ensures(ret.1 == LineageInteraction::Maybe, "always reports maybe")]
    default fn sample_interaction_at_location(
        &self,
        location: Location,
        habitat: &H,
        _lineage_store: &IndependentLineageStore<M, H>,
        coalescence_rng_sample: CoalescenceRngSample,
    ) -> (IndexedLocation, LineageInteraction) {
        let population = habitat.get_habitat_at_location(&location);

        let chosen_coalescence_index =
            coalescence_rng_sample.sample_coalescence_index::<M>(population);

        let indexed_location = IndexedLocation::new(location, chosen_coalescence_index);

        (indexed_location, LineageInteraction::Maybe)
    }
}

// Specialise for SingletonDemesHabitat as the compiler cannot yet optimise out
//  the call to `habitat.get_habitat_at_location(&location)`.
#[contract_trait]
impl<M: MathsCore, H: SingletonDemesHabitat<M>>
    CoalescenceSampler<M, H, IndependentLineageStore<M, H>>
    for IndependentCoalescenceSampler<M, H>
{
    #[must_use]
    #[debug_ensures(ret.1 == LineageInteraction::Maybe, "always reports maybe")]
    fn sample_interaction_at_location(
        &self,
        location: Location,
        _habitat: &H,
        _lineage_store: &IndependentLineageStore<M, H>,
        _coalescence_rng_sample: CoalescenceRngSample,
    ) -> (IndexedLocation, LineageInteraction) {
        (
            IndexedLocation::new(location, 0_u32),
            LineageInteraction::Maybe,
        )
    }
}