Skip to content

Allowing using NamedTuple as initial_params #632

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 12 commits into from
Jul 23, 2024
Merged

Allowing using NamedTuple as initial_params #632

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
52a9b76
Initial work on `NamedTuple` as `initial_params`
torfjelde Jul 15, 2024
94c2ef7
Apply suggestions from code review
torfjelde Jul 15, 2024
5b5f4dd
add some tests
sunxd3 Jul 15, 2024
4693fa2
Update src/sampler.jl
sunxd3 Jul 15, 2024
bab789e
fix type error for inits with `nothing`
sunxd3 Jul 15, 2024
c3ce3fb
Update src/sampler.jl
sunxd3 Jul 15, 2024
aab6606
Merge branch 'master' into torfjelde/namedtuple-initial-params
sunxd3 Jul 16, 2024
22942d4
use better variable names
sunxd3 Jul 22, 2024
a3969d8
remove init with scalar
sunxd3 Jul 22, 2024
1b2ce47
move `update_values!!` out of `TestUtils`
sunxd3 Jul 22, 2024
14bb2cf
fix error
sunxd3 Jul 22, 2024
8b79ac0
Merge branch 'master' into torfjelde/namedtuple-initial-params
sunxd3 Jul 22, 2024
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
60 changes: 38 additions & 22 deletions src/sampler.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -142,38 +142,54 @@ By default, it returns an instance of [`SampleFromPrior`](@ref).
"""
initialsampler(spl::Sampler) = SampleFromPrior()

function initialize_parameters!!(
vi::AbstractVarInfo, initial_params, spl::Sampler, model::Model
function set_values!!(
varinfo::AbstractVarInfo,
initial_params::AbstractVector{<:Union{Real,Missing}},
Copy link
Member

@sunxd3 sunxd3 Jul 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I changed the array element type from Real to Union{Real, Missing} because we allow initialization vector with missing in it.

Copy link
Member Author

@torfjelde torfjelde Jul 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Really? In what scenario would this be used?

Copy link
Member

@sunxd3 sunxd3 Jul 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DynamicPPL.jl/test/sampler.jl

Line 135 in dfdc155

chain = sample(model, sampler, 1; initial_params=[missing, -1], progress=false)

Copy link
Member Author

@torfjelde torfjelde Jul 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, not sure how I feel about that. I guess it means to sample that parameter from the prior? But that's so simple to do by hand anyways these days.

Copy link
Member

@sunxd3 sunxd3 Jul 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suppose so, although it made sense to me because the initialization vector need to be the same dimension as the model, so it makes sense for someone to say "I don't care about these" by setting them to missing

Copy link
Member Author

@torfjelde torfjelde Jul 18, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, but you can also just do

rand(Vector, model)

and alter those values 😕

Basically, it's a question of whether we want to maintain this or just leave it up to the user when it is a simple one-liner.

Copy link
Member

@sunxd3 sunxd3 Jul 22, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Got your argument, 👍
But in this case it's not really a one-liner, as user also need to set individual values.
For models with small dimensions, current syntax can still be useful. I am for keeping this option.

spl::AbstractSampler,
)
@debug "Using passed-in initial variable values" initial_params

# Flatten parameters.
init_theta = mapreduce(vcat, initial_params) do x
vec([x;])
end

# Get all values.
linked = islinked(vi, spl)
if linked
vi = invlink!!(vi, spl, model)
end
theta = vi[spl]
length(theta) == length(init_theta) || throw(
flattened_param_vals = varinfo[spl]
length(flattened_param_vals) == length(initial_params) || throw(
DimensionMismatch(
"Provided initial value size ($(length(init_theta))) doesn't match the model size ($(length(theta)))",
"Provided initial value size ($(length(initial_params))) doesn't match the model size ($(length(theta)))",
),
)

# Update values that are provided.
for i in eachindex(init_theta)
x = init_theta[i]
for i in eachindex(initial_params)
x = initial_params[i]
if x !== missing
theta[i] = x
flattened_param_vals[i] = x
end
end

# Update in `vi`.
vi = setindex!!(vi, theta, spl)
# Update in `varinfo`.
return setindex!!(varinfo, flattened_param_vals, spl)
end

function set_values!!(
varinfo::AbstractVarInfo, initial_params::NamedTuple, spl::AbstractSampler
Copy link
Member

@mhauru mhauru Jul 19, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are we happy with spl not doing anything in this function? I get that we need to leave it there to keep a uniform interface, but could this end up being called with varinfo that has not been sampled with spl, which would then produce unexpected results?

)
initial_params = NamedTuple(k => v for (k, v) in pairs(initial_params) if v !== missing)
Copy link
Member

@sunxd3 sunxd3 Jul 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure this is good, but changed to support a more uniform interface: allowing NamedTuple with missing values.

Copy link
Member Author

@torfjelde torfjelde Jul 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as in last comment: when are we initializing using missing?

Copy link
Member

@sunxd3 sunxd3 Jul 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

see reply above

return update_values!!(
varinfo, initial_params, map(k -> VarName{k}(), keys(initial_params))
)
end

function initialize_parameters!!(
vi::AbstractVarInfo, initial_params, spl::AbstractSampler, model::Model
)
@debug "Using passed-in initial variable values" initial_params

# `link` the varinfo if needed.
linked = islinked(vi, spl)
if linked
vi = invlink!!(vi, spl, model)
end

# Set the values in `vi`.
vi = set_values!!(vi, initial_params, spl)

# `invlink` if needed.
if linked
vi = link!!(vi, spl, model)
end
Expand Down
14 changes: 1 addition & 13 deletions src/test_utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,19 +11,7 @@ using Bijectors: Bijectors
using Accessors: Accessors

# For backwards compat.
using DynamicPPL: varname_leaves

"""
update_values!!(vi::AbstractVarInfo, vals::NamedTuple, vns)

Return instance similar to `vi` but with `vns` set to values from `vals`.
"""
function update_values!!(vi::AbstractVarInfo, vals::NamedTuple, vns)
for vn in vns
vi = DynamicPPL.setindex!!(vi, get(vals, vn), vn)
end
return vi
end
using DynamicPPL: varname_leaves, update_values!!

"""
test_values(vi::AbstractVarInfo, vals::NamedTuple, vns)
Expand Down
12 changes: 12 additions & 0 deletions src/utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -796,6 +796,18 @@ function nested_getindex(values::AbstractDict, vn::VarName)
return child(value)
end

"""
update_values!!(vi::AbstractVarInfo, vals::NamedTuple, vns)

Return instance similar to `vi` but with `vns` set to values from `vals`.
"""
function update_values!!(vi::AbstractVarInfo, vals::NamedTuple, vns)
for vn in vns
vi = DynamicPPL.setindex!!(vi, get(vals, vn), vn)
end
return vi
end

"""
float_type_with_fallback(x)

Expand Down
6 changes: 6 additions & 0 deletions src/varinfo.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -892,6 +892,12 @@ Base.keys(vi::TypedVarInfo{<:NamedTuple{()}}) = VarName[]
return expr
end

# FIXME(torfjelde): Don't use `_getvns`.
Copy link
Member

@mhauru mhauru Jul 19, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this still relevant? Don't know why _getvns should be shunned.

Copy link
Member Author

@torfjelde torfjelde Jul 19, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's just unnecessary; Base.keys should just be it + IIRC _getvns is overly complicated.

Base.keys(vi::UntypedVarInfo, spl::AbstractSampler) = _getvns(vi, spl)
function Base.keys(vi::TypedVarInfo, spl::AbstractSampler)
return mapreduce(values, vcat, _getvns(vi, spl))
end

"""
setgid!(vi::VarInfo, gid::Selector, vn::VarName)

Expand Down
112 changes: 59 additions & 53 deletions test/sampler.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,23 +84,25 @@
model = coinflip()
sampler = Sampler(alg)
lptrue = logpdf(Binomial(25, 0.2), 10)
chain = sample(model, sampler, 1; initial_params=0.2, progress=false)
@test chain[1].metadata.p.vals == [0.2]
@test getlogp(chain[1]) == lptrue

# parallel sampling
chains = sample(
model,
sampler,
MCMCThreads(),
1,
10;
initial_params=fill(0.2, 10),
progress=false,
)
for c in chains
@test c[1].metadata.p.vals == [0.2]
@test getlogp(c[1]) == lptrue
let inits = (; p=0.2)
chain = sample(model, sampler, 1; initial_params=inits, progress=false)
@test chain[1].metadata.p.vals == [0.2]
@test getlogp(chain[1]) == lptrue

# parallel sampling
chains = sample(
model,
sampler,
MCMCThreads(),
1,
10;
initial_params=fill(inits, 10),
progress=false,
)
for c in chains
@test c[1].metadata.p.vals == [0.2]
@test getlogp(c[1]) == lptrue
end
end

# model with two variables: initialization s = 4, m = -1
Expand All @@ -110,45 +112,49 @@
end
model = twovars()
lptrue = logpdf(InverseGamma(2, 3), 4) + logpdf(Normal(0, 2), -1)
chain = sample(model, sampler, 1; initial_params=[4, -1], progress=false)
@test chain[1].metadata.s.vals == [4]
@test chain[1].metadata.m.vals == [-1]
@test getlogp(chain[1]) == lptrue

# parallel sampling
chains = sample(
model,
sampler,
MCMCThreads(),
1,
10;
initial_params=fill([4, -1], 10),
progress=false,
)
for c in chains
@test c[1].metadata.s.vals == [4]
@test c[1].metadata.m.vals == [-1]
@test getlogp(c[1]) == lptrue
for inits in ([4, -1], (; s=4, m=-1))
chain = sample(model, sampler, 1; initial_params=inits, progress=false)
@test chain[1].metadata.s.vals == [4]
@test chain[1].metadata.m.vals == [-1]
@test getlogp(chain[1]) == lptrue

# parallel sampling
chains = sample(
model,
sampler,
MCMCThreads(),
1,
10;
initial_params=fill(inits, 10),
progress=false,
)
for c in chains
@test c[1].metadata.s.vals == [4]
@test c[1].metadata.m.vals == [-1]
@test getlogp(c[1]) == lptrue
end
end

# set only m = -1
chain = sample(model, sampler, 1; initial_params=[missing, -1], progress=false)
@test !ismissing(chain[1].metadata.s.vals[1])
@test chain[1].metadata.m.vals == [-1]

# parallel sampling
chains = sample(
model,
sampler,
MCMCThreads(),
1,
10;
initial_params=fill([missing, -1], 10),
progress=false,
)
for c in chains
@test !ismissing(c[1].metadata.s.vals[1])
@test c[1].metadata.m.vals == [-1]
for inits in ([missing, -1], (; s=missing, m=-1), (; m=-1))
chain = sample(model, sampler, 1; initial_params=inits, progress=false)
@test !ismissing(chain[1].metadata.s.vals[1])
@test chain[1].metadata.m.vals == [-1]

# parallel sampling
chains = sample(
model,
sampler,
MCMCThreads(),
1,
10;
initial_params=fill(inits, 10),
progress=false,
)
for c in chains
@test !ismissing(c[1].metadata.s.vals[1])
@test c[1].metadata.m.vals == [-1]
end
end

# specify `initial_params=nothing`
Expand Down