Skip to content

Commit

Permalink
some performance improvements and bugfix
Browse files Browse the repository at this point in the history
  • Loading branch information
@benjione
benjione committed Nov 7, 2023
1 parent da0a9a5 commit 684f970
Show file tree
Hide file tree
Showing 8 changed files with 67 additions and 46 deletions.
2 changes: 1 addition & 1 deletion src/PSDModels/feature_map/PSDModelFM.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ end
Returns the feature map of the PSD model at x.
"""
@inline function Φ(a::AbstractPSDModelFM, x::PSDdata{T}) where {T<:Number}
@inline function Φ(a::AbsPSD, x::PSDdata{T}) where {T<:Number, AbsPSD<:AbstractPSDModelFM{T}}
return a.Φ(x)
end

34 changes: 22 additions & 12 deletions src/Samplers/PSDModelSampler.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,10 @@ struct PSDModelSampler{d, T<:Number, S, R, dC} <: ConditionalSampler{d, T, R, dC
model::PSDModelOrthonormal{d, T, S} # model to sample from
margins::Vector{<:PSDModelOrthonormal{<:Any, T, S}} # start with x_{≤1}, then x_{≤2}, ...
integrals::Vector{<:OrthonormalTraceModel{T, S}} # integrals of marginals
variable_ordering::AbstractVector{Int} # variable ordering for the model
variable_ordering::Vector{Int} # variable ordering for the model
R_map::R # reference map from reference distribution to uniform
function PSDModelSampler(model::PSDModelOrthonormal{d, T, S},
variable_ordering::AbstractVector{Int},
variable_ordering::Vector{Int},
amount_cond_variable::Int) where {d, T<:Number, S}
@assert amount_cond_variable < d
# check that the last {amount_cond_variable} variables are the last ones in the ordering
Expand All @@ -20,7 +20,7 @@ struct PSDModelSampler{d, T<:Number, S, R, dC} <: ConditionalSampler{d, T, R, dC
new{d, T, S, Nothing, amount_cond_variable}(model, margins, integrals, variable_ordering, nothing)
end
function PSDModelSampler(model::PSDModelOrthonormal{d, T, S},
variable_ordering::AbstractVector{Int}) where {d, T<:Number, S}
variable_ordering::Vector{Int}) where {d, T<:Number, S}
PSDModelSampler(model, variable_ordering, 0)
end
end
Expand All @@ -46,7 +46,7 @@ function _pushforward_first_n(sampler::PSDModelSampler{d, T, S},
x = zeros(T, n)
u = @view u[sampler.variable_ordering[1:n]]
## T^{-1}(x_1,...,x_k) functions, z=x_k
f(k) = begin
f(k::Int) = begin
if k==1
z->sampler.integrals[k](T[z]) - u[k] #u[sampler.variable_ordering[k]]
else
Expand All @@ -56,12 +56,12 @@ function _pushforward_first_n(sampler::PSDModelSampler{d, T, S},
end
if S<:OMF
for k=1:n
x[k] = find_zero(f(k), zero(T))
x[k] = find_zero(f(k), zero(T))::T
end
else
for k=1:n
left, right = domain_interval(sampler.model, sampler.variable_ordering[k])
x[k] = find_zero(f(k), (left, right))
left, right = domain_interval(sampler.model, sampler.variable_ordering[k])::Tuple{T, T}
x[k] = find_zero(f(k), (left, right))::T
end
end
return invpermute!(x, sampler.variable_ordering[1:n])
Expand All @@ -72,19 +72,29 @@ function _pullback_first_n(sampler::PSDModelSampler{d, T},
x::PSDdata{T},
n::Int) where {d, T<:Number}
x = @view x[sampler.variable_ordering[1:n]]
f(k) = begin
f(k::Int) = begin
if k==1
z->sampler.integrals[k](T[z])
else
z->(sampler.integrals[k]([x[1:k-1]; z])/sampler.margins[k-1](x[1:k-1]))
end
end
u = zeros(T, n)
u = Vector{T}(undef, n)
for k=1:n
u[sampler.variable_ordering[k]] = f(k)(x[k])
end
u = map(x->x ≤ 0 ? (x ≥ one(T) ? x : one(T)) : zero(T) : x, u)
return u
# typestable function
_rounding(u::Vector{T}) = begin
map(u) do x
if zero(T) ≤ x ≤ one(T)
x
else
zero(T) > x ? zero(T) : one(T)
end
end
end
u = _rounding(u)
return u::Vector{T}
end


Expand All @@ -94,7 +104,7 @@ function Distributions.pdf(
sar::PSDModelSampler{d, T},
x::PSDdata{T}
) where {d, T<:Number}
return sar.model(x)
return sar.model(x)::T
end

@inline pushforward(sampler::PSDModelSampler{d, T, S},
Expand Down
4 changes: 2 additions & 2 deletions src/Samplers/SelfReinforcedSampler.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -506,7 +506,7 @@ function pushforward(
u = pushforward(sra.samplers[j], u)
end
u = _ref_pullback(sra, u)
return u
return u::Vector{T}
end

function pullback(
Expand All @@ -520,7 +520,7 @@ function pullback(
x = pullback(sra.samplers[j], x)
end
x = _ref_pullback(sra, x)
return x
return x::Vector{T}
end

## Interface of ConditionalSampler
Expand Down
37 changes: 22 additions & 15 deletions src/Samplers/SubsetSampler.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,10 @@ struct SubsetSampler{d, d_sub, T<:Number,
R<:ReferenceMap{d, T},
R_sub<:ReferenceMap{d_sub, T},
dC,
d_sub_cond # reduced dimension of conditional variables
d_sub_cond, # reduced dimension of conditional variables
sampler_type<:ConditionalSampler{d_sub, T, <:Any, d_sub_cond}
} <: ConditionalSampler{d, T, Nothing, dC}
sampler::ConditionalSampler{d_sub, T, <:Any, d_sub_cond}
sampler::sampler_type
X::AbstractMatrix{T} # bases X of subspace, X ∈ R^{d x d2}
P::AbstractMatrix{T} # Projection so that P * y ∈ S and (I - P) * y = X' * y ∈ S^⟂
P_tilde::AbstractMatrix{T} # transform of y ∈ R^d to subspace α ∈ R^d2
Expand All @@ -24,7 +25,11 @@ struct SubsetSampler{d, d_sub, T<:Number,
ones(T, length(sampler_variables)))
P_tilde = X'
P = X * P_tilde
new{d, d2, T, R, R2, dC, d_sub_cond}(sampler, X, P, P_tilde, R_map, R_map_sub)
new{d, d2, T, R, R2, dC, d_sub_cond, typeof(sampler)}(
sampler, X, P,
P_tilde, R_map,
R_map_sub
)
end
function SubsetSampler{d}(sampler::Sampler{d2, T},
sampler_variables::AbstractVector{Int},
Expand Down Expand Up @@ -55,7 +60,7 @@ struct SubsetSampler{d, d_sub, T<:Number,
@assert size(P, 2) == d
@assert size(P_tilde, 1) == d2
@assert size(P_tilde, 2) == d
new{d, d2, T, R, R2, dC, d_sub_cond}(sampler, X, P, P_tilde, R_map, R_map_sub)
new{d, d2, T, R, R2, dC, d_sub_cond, typeof(sampler)}(sampler, X, P, P_tilde, R_map, R_map_sub)
end
# function SubsetSampler{d, dC}(sampler::Sampler{d2, T},
# X::AbstractMatrix{T},
Expand All @@ -80,9 +85,9 @@ Hence, det(∂v (P * T(v) + (1-P) v)) = det(∂v T(v)) * det(∂v v)
Total expression of forward is with T being the sampler:
R^{-1}(y) = X * R_sub^{-1}(T(R_sub(X' R^{-1}(x)))) + (I - P) * R^{-1}(x)
"""
function Distributions.pdf(sampler::SubsetSampler{<:Any, <:Any, T, R, R_sub},
function Distributions.pdf(sampler::SubsetSampler{<:Any, <:Any, T, R, R_sub, samplerType},
x::PSDdata{T}
) where {T<:Number, R, R_sub}
) where {T<:Number, R, R_sub, samplerType}
xs = pullback(sampler.R_map, x) # from [0,1]^d to R^d
sub_x = sampler.P_tilde * xs # from R^d to R^d2
sub_x_maped = pushforward(sampler.R_map_sub, sub_x) # from R^d2 to [0,1]^d2
Expand Down Expand Up @@ -149,9 +154,9 @@ function project_to_subset(P_tilde::AbstractMatrix{T},
return sub_x
end

function pushforward(sampler::SubsetSampler{<:Any, <:Any, T, R, R_sub},
function pushforward(sampler::SubsetSampler{<:Any, <:Any, T, R, R_sub,<:Any, <:Any,ST},
u::PSDdata{T}
) where {T<:Number, R, R_sub}
) where {T<:Number, R, R_sub, ST}
us = pullback(sampler.R_map, u)
sub_u = sampler.P_tilde * us
sub_x = pushforward(sampler.sampler, pushforward(sampler.R_map_sub, sub_u))
Expand All @@ -160,7 +165,9 @@ function pushforward(sampler::SubsetSampler{<:Any, <:Any, T, R, R_sub},
return x
end

function pullback(sampler::SubsetSampler{<:Any, <:Any, T, R}, x::PSDdata{T}) where {T<:Number, R}
function pullback(sampler::SubsetSampler{<:Any, <:Any, T, R, <:Any, <:Any, <:Any, ST},
x::PSDdata{T}
) where {T<:Number, R, ST}
xs = pullback(sampler.R_map, x)
sub_x = sampler.P_tilde * xs
sub_u = pullback(sampler.sampler, pushforward(sampler.R_map_sub, sub_x))
Expand All @@ -177,8 +184,8 @@ end
"""
Distribution p(x) = ∫ p(x, y) d y
"""
function marg_pdf(sampler::SubsetSampler{d, d_sub, T, R, R2, dC, d_sub_cond},
x::PSDdata{T}) where {d, d_sub, T<:Number, R, R2, dC, d_sub_cond}
function marg_pdf(sampler::SubsetSampler{d, d_sub, T, R, R2, dC, d_sub_cond, ST},
x::PSDdata{T}) where {d, d_sub, T<:Number, R, R2, dC, d_sub_cond, ST}
xs = pullback(sampler.R_map, x) # from [0,1]^dx to R^dx
sub_x = _marg_P_tilde(sampler) * xs # from R^dx to R^d_sub_marg
sub_x_maped = pushforward(sampler.R_map_sub, sub_x) # from R^d_sub_marg to [0,1]^d_sub_marg
Expand All @@ -192,9 +199,9 @@ function marg_pdf(sampler::SubsetSampler{d, d_sub, T, R, R2, dC, d_sub_cond},
return inverse_Jacobian(sampler.R_map, x) * T_inv_jac_det * Jacobian(sampler.R_map, us)
end

function marg_pushforward(sampler::SubsetSampler{d, d_sub, T, R, R2, dC, d_sub_cond},
function marg_pushforward(sampler::SubsetSampler{d, d_sub, T, R, R2, dC, d_sub_cond, ST},
u::PSDdata{T}
) where {d, d_sub, T<:Number, R, R2, dC, d_sub_cond}
) where {d, d_sub, T<:Number, R, R2, dC, d_sub_cond, ST}
us = pullback(sampler.R_map, u)
sub_u = _marg_P_tilde(sampler) * us
sub_x = marg_pushforward(sampler.sampler, pushforward(sampler.R_map_sub, sub_u))
Expand All @@ -203,9 +210,9 @@ function marg_pushforward(sampler::SubsetSampler{d, d_sub, T, R, R2, dC, d_sub_c
return x
end

function marg_pullback(sampler::SubsetSampler{d, d_sub, T, R, R2, dC, d_sub_cond},
function marg_pullback(sampler::SubsetSampler{d, d_sub, T, R, R2, dC, d_sub_cond, ST},
x::PSDdata{T}
) where {d, d_sub, T<:Number, R, R2, dC, d_sub_cond}
) where {d, d_sub, T<:Number, R, R2, dC, d_sub_cond, ST}
xs = pullback(sampler.R_map, x)
sub_x = _marg_P_tilde(sampler) * xs
sub_u = marg_pullback(sampler.sampler, pushforward(sampler.R_map_sub, sub_x))
Expand Down
16 changes: 8 additions & 8 deletions src/Samplers/reference_maps/algebraic.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,34 +10,34 @@ end


function PSDModels.pushforward(
m::AlgebraicReference{d, T},
m::AlgebraicReference{<:Any, T},
x::PSDdata{T}
) where {d, T<:Number}
) where {T<:Number}
return ((x./sqrt.(1 .+ x.^2)).+1.0)/2.0
end


function PSDModels.pullback(
m::AlgebraicReference{d, T},
m::AlgebraicReference{<:Any, T},
u::PSDdata{T}
) where {d, T<:Number}
) where {T<:Number}
ξ = 2*(u .- 0.5)
return ξ./sqrt.(1 .- ξ.^2)
end


function Jacobian(
m::AlgebraicReference{d, T},
m::AlgebraicReference{<:Any, T},
x::PSDdata{T}
) where {d, T<:Number}
) where {T<:Number}
return mapreduce(xi->2.0/(1+xi^2)^(3/2), *, x)
end


function inverse_Jacobian(
mapping::AlgebraicReference{d, T},
mapping::AlgebraicReference{<:Any, T},
u::PSDdata{T}
) where {d, T<:Number}
) where {T<:Number}
# inverse function theorem
return mapreduce(ui->2.0/(1-ui^2)^(3/2), *, u)
end
2 changes: 1 addition & 1 deletion src/TraceModels/polynomial.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

abstract type OrthonormalTraceModel{T, S} <: TraceModel{T} end

function ΦΦT(a::OrthonormalTraceModel{T, Nothing}, x::PSDdata{T}) where {T<:Number}
function ΦΦT(a::OT, x::PSDdata{T}) where {T<:Number, OT<:OrthonormalTraceModel{T, Nothing}}
return a.ΦΦT(x)
end

Expand Down
6 changes: 3 additions & 3 deletions test/Sampler/conditional_sampler_test.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,12 +38,12 @@ end
amount_cond_variable=1,
)
# test densities are close
for x in range(-5, 5, length=20)
for y in range(-5, 5, length=20)
for x in range(-5.0, 5.0, length=20)
for y in range(-5.0, 5.0, length=20)
@test abs(pdf(sra, [x, y]) - f([x, y])) < 0.1
end
end
for x in range(-4, 4, length=100)
for x in range(-4.0, 4.0, length=100)
@test abs(PSDModels.marg_pdf(sra, [x]) - f_marg(x)) < 0.2
end
end
Expand Down
12 changes: 8 additions & 4 deletions test/Sampler/subset_sampler_test.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,10 +51,14 @@ end
X2 = rand(distr2, N2)
X = hcat(X1, X2)
T = Float64
bridge = DiffusionBrigdingDensity{2}(f, T[1.5, 1.3, 1.0, 0.8, 0.75, 0.5, 0.3, 0.25, 0.18, 0.13, 0.1, 0.07, 0.02, 0.01, 0.005, 0.0], T(2.0))
ref_map = PSDModels.ReferenceMaps.GaussianReference{3, T}(T(2.0))
to_subspace_ref_map = PSDModels.ReferenceMaps.GaussianReference{3, T}(T(2.0))
subspace_ref_map = PSDModels.ReferenceMaps.GaussianReference{2, T}(T(2.0))
bridge = DiffusionBrigdingDensity{2}(f, T[1.5, 1.3, 1.0, 0.8, 0.75,
0.5, 0.3, 0.25, 0.18, 0.13,
0.1, 0.07, 0.02, 0.01, 0.005, 0.0], T(2.0))
ref_map = PSDModels.ReferenceMaps.GaussianReference{3, T}(T(2.5))
to_subspace_ref_map = PSDModels.ReferenceMaps.GaussianReference{3, T}(T(3.0))
subspace_ref_map = PSDModels.ReferenceMaps.GaussianReference{2, T}(T(3.0))
# to_subspace_ref_map = PSDModels.ReferenceMaps.AlgebraicReference{3, T}()
# subspace_ref_map = PSDModels.ReferenceMaps.AlgebraicReference{2, T}()

model = PSDModel{T}(Legendre(T(0)..T(1))^2, :downward_closed, 3)

Expand Down

0 comments on commit 684f970

Please sign in to comment.