-
Notifications
You must be signed in to change notification settings - Fork 4
/
selections.jl
69 lines (55 loc) · 2.08 KB
/
selections.jl
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
using NQCDistributions: DynamicalDistribution, ProductDistribution
abstract type AbstractSelection end
"""
Select the initial conditions from the distribution in order.
"""
struct OrderedSelection{D,I} <: AbstractSelection
"Distribution that is sampled."
distribution::D
indices::I
end
"""
Obtain initial conditions by randomly sampling the distribution.
"""
struct RandomSelection{D} <: AbstractSelection
"Distribution that is sampled."
distribution::D
end
function Selection(distribution, selection::AbstractVector, trajectories)
if trajectories > 1
@info "Sampling the provided distribution in the range $selection."
end
OrderedSelection(distribution, selection)
end
function Selection(distribution, ::Any, trajectories)
if trajectories > 1
@info "Sampling randomly from provided distribution."
end
RandomSelection(distribution)
end
function (select::OrderedSelection)(prob, i, repeat)
j = select.indices[i]
u0 = sample_distribution(prob.p, select.distribution, j)
DynamicsMethods.create_problem(u0, prob.tspan, prob.p)
end
function sample_distribution(sim::AbstractSimulation, distribution::DynamicalDistribution, i)
u = distribution[i]
DynamicsMethods.DynamicsVariables(sim, u.v, u.r)
end
function sample_distribution(sim::AbstractSimulation, distribution::ProductDistribution, i)
u = distribution.nuclear[i]
DynamicsMethods.DynamicsVariables(sim, u.v, u.r, distribution.electronic)
end
function (select::RandomSelection)(prob, i, repeat)
u0 = sample_distribution(prob.p, select.distribution)
DynamicsMethods.create_problem(u0, prob.tspan, prob.p)
end
function sample_distribution(sim::AbstractSimulation, distribution::DynamicalDistribution)
u = rand(distribution)
DynamicsMethods.DynamicsVariables(sim, u.v, u.r)
end
function sample_distribution(sim::AbstractSimulation, distribution::ProductDistribution)
u = rand(distribution.nuclear)
DynamicsMethods.DynamicsVariables(sim, u.v, u.r, distribution.electronic)
end
sample_distribution(::AbstractSimulation, distribution) = copy(distribution)