Implements a simple Nutpie style adaptation (using both positions and gradients, but not changing the schedule). #473

nsiccha · 2025-11-19T12:00:26Z

WIP that partially addresses #311 and supersedes #312.

There's a demo in tmp/demo.jl, which certainly does something that finishes quicker than the current default.

There are currently no additional tests, and I'm sure a few things are currently broken due to my changes.

Gonna tag @sethaxen, @aseyboldt, @svilupp, and maybe @yebai.

nsiccha · 2025-11-19T12:09:22Z

Gonna ask some questions before I move on:

Currently, the way I change the used mass matrix adaptor feels a bit hacky, reproduced below:

adaptor = AdvancedHMC.StanHMCAdaptor(
    AdvancedHMC.Adaptation.NutpieVar(size(metric); var=copy(metric.M⁻¹)), 
    AdvancedHMC.StepSizeAdaptor(spl.δ, integrator)
)
h, t = AdvancedHMC.sample_init(rng, hamiltonian, initial_params)
# Using the below uses Nutpie (as in position and gradients)
initial_state = AdvancedHMC.HMCState(0, t, metric, κ, adaptor)
# Using the below uses Stan (as in only positions)
# initial_state = nothing
@time samples = AdvancedHMC.sample(
    rng,
    model,
    spl,
    n_adapts + n_samples;
    n_adapts=n_adapts, initial_state,
    progress=true, 
)

Is there currently no easier way to specify what kind of adaptation to use, ideally just via some (keyword) argument to the sample function? Gonna also tag @penelopeysm and @mhauru who might know or have opinions on how to change the public API :)

nsiccha · 2025-11-20T09:33:21Z

After chatting with or at @penelopeysm I've opened #475 and think that this PR should only implement what it's currently doing.

I don't know whether we even want to export the defined struct currently - maybe.
In any case, I think I'll be adding new tests, docstrings, and make sure that the current tests don't fail anymore.

The main thing where I might need help is to figure out whether the needed changes to the adapt! function break things outside of this repo - which they might very well do!

nsiccha · 2025-11-21T09:27:56Z

Reproducing the code to demo the changes in this PR at the end of this comment.

There's maybe one thing I'm unhappy with in this PR: There's a bunch of code duplication for the adapt! method.

I needed to pass the position+gradient information through to the mass matrix adaptor, and the easiest way to do that was to allow a PhasePoint to be passed instead of a simple vector of the position of the last draw. However, I think due to backward compatibility I also had to still keep the old methods, accepting only the position instead of a PhasePoint.

using AdvancedHMC, PosteriorDB, StanLogDensityProblems, Random, MCMCDiagnosticTools

if !@isdefined pdb 
    const pdb = PosteriorDB.database()
end
stan_problem(path, data) = StanProblem(
    path, data;
    nan_on_error=true,
    make_args=["STAN_THREADS=TRUE"],
    warn=false
)
stan_problem(posterior_name::AbstractString) = stan_problem(
    PosteriorDB.path(PosteriorDB.implementation(PosteriorDB.model(PosteriorDB.posterior(pdb, (posterior_name))), "stan")), 
    PosteriorDB.load(PosteriorDB.dataset(PosteriorDB.posterior(pdb, (posterior_name))), String)
)

begin
lpdf = stan_problem("radon_mn-radon_county_intercept")

n_adapts = n_samples = 1000
rng = Xoshiro(2)
spl = NUTS(0.8)
initial_params = nothing
model = AdvancedHMC.AbstractMCMC._model(lpdf)
(;logdensity) = model
metric = AdvancedHMC.make_metric(spl, logdensity)
hamiltonian = AdvancedHMC.Hamiltonian(metric, model)
initial_params = AdvancedHMC.make_initial_params(rng, spl, logdensity, initial_params)
ϵ = AdvancedHMC.make_step_size(rng, spl, hamiltonian, initial_params)
integrator = AdvancedHMC.make_integrator(spl, ϵ)
κ = AdvancedHMC.make_kernel(spl, integrator)
adaptor = AdvancedHMC.StanHMCAdaptor(
    AdvancedHMC.Adaptation.NutpieVar(size(metric); var=copy(metric.M⁻¹)), 
    AdvancedHMC.StepSizeAdaptor(spl.δ, integrator)
)
h, t = AdvancedHMC.sample_init(rng, hamiltonian, initial_params)
performances = map((;nutpie=AdvancedHMC.HMCState(0, t, metric, κ, adaptor), stan=nothing)) do initial_state
    dt = @elapsed samples = AdvancedHMC.sample(
        rng,
        model,
        spl,
        n_adapts + n_samples;
        n_adapts=n_adapts, initial_state,
        progress=true, 
    )
    ess(reshape(mapreduce(sample->sample.z.θ , hcat, samples[n_adapts+1:end])', (n_samples, 1, :))) |> minimum  |> Base.Fix2(/, dt)
end
@info (;performances)
end

nsiccha · 2025-11-21T09:38:23Z

Hm - I'm pretty sure the failings tests are not due to my changes - what's up with that?

mhauru

I did a hasty review, just reading through the files and noting things that stuck out to me. I didn't actually try to understand the logic of the code, of how this all works. (I don't really know AHMC and it's existing structures.)

mhauru · 2025-11-21T18:10:25Z

src/adaptation/Adaptation.jl

    adapt!(nca.pc, θ, α)
    return nothing
 end
+adapt!(


Could you please run JuliaFormatter.jl on this PR? I believe it would format this function, and some other things, differently. Note that we use JuliaFormatter v1, not v2! v2 is still quite buggy, so be sure to install explicitly with ] add JuliaFormatter@v1. If you like using pre-commit, let me know, I've got a config for running the formatter in pre-commit.

If you like using pre-commit, let me know, I've got a config for running the formatter in pre-commit.

I don't know whether I like that - time to find out!
(So yes, would be great if you could share your config :) )

Shared on Slack

I think this remains unformatted. At least when I run the formatter locally I get this diff:

- adapt!( - nca::NaiveHMCAdaptor, - z::PhasePoint, - α::AbstractScalarOrVec{<:AbstractFloat}, - ) = adapt!(nca, z.θ, α) + function adapt!( + nca::NaiveHMCAdaptor, z::PhasePoint, α::AbstractScalarOrVec{<:AbstractFloat} + ) + return adapt!(nca, z.θ, α) + end

Instead of pre-commit, if you prefer you can also do julia -e 'using JuliaFormatter; format(("src", "test"))', although that will format all files, not just the ones you edited.

Huh, I wasn't able to comment on github for some reason, but I can via VSCode? I'll format this with the next commit! I hadn't changed this file in the commits since installing pre-commit.

Ah, I see - I had to scroll up to the first review comment thread.

mhauru · 2025-11-21T18:21:57Z

src/adaptation/massmatrix.jl


+## Nutpie-style diagonal mass matrix estimator (using positions and gradients)
+
+mutable struct NutpieVar{T<:AbstractFloat,E<:AbstractVecOrMat{T},V<:AbstractVecOrMat{T}} <: DiagMatrixEstimator{T}


Does this have to be mutable?

I thought that it doesn't have to be - but to adhere to the implicit internal interface, having it be mutable makes implementation easier. WelfordVar e.g. is also mutable.

See e.g. here, which implies among other things the presence of a (mutable) n field.

I think the function you linked should be fine. It does a comparison on ve.n and then conditionally mutates ve.var, but it does the mutation with ve.var .=, which doesn't require changing the value of the field ve.var by making it point to a new Array, but rather changes the elements within ve.var, for which ve itself doesn't have to be mutable. Note that an immutable type can have mutable objects as its field values.

If it is an interface demand that subtypes of DiagMatrixEstimator have to be mutable then I think that should be changed because

I think it's better to define interfaces based on functions and methods rather than particular fields of a type. So rather than say "there has to be field .var and we must be able to mutate it" we could say "there has to be function setvar!! with argument types blahblah".

Mutable types are generally harder to reason about and often slower, and thus good to avoid when possible.

However, that's probably out of scope for this PR.

Oh right, I mostly meant that ve.n has to be able to be mutated, because otherwise the conditional will always evaluate to the same value.

You can of course get around that by having some ...!! function acting on the adaptation, but I do find this a bit awkward and would prefer a solution where ve.n would just be a Ref (inside an otherwise non-mutable struct) and then the conditional would look something like nobs(ve) >= min_nobs(ve).

If it is an interface demand that subtypes of DiagMatrixEstimator have to be mutable then I think that should be changed.

I absolutely agree!

Edit: though I don't think it's a strict interface demand - for it was just so that the method using that interface was already defined (and would be used for my subtype), so I thought I'd just rely on that already present method.

for it was just so that the method using that interface was already defined (and would be used for my subtype), so I thought I'd just rely on that already present method.

I think that's fair. Would raise this as something to change later for the whole abstact type though.

You can of course get around that by having some ...!! function acting on the adaptation, but I do find this a bit awkward and would prefer a solution where ve.n would just be a Ref (inside an otherwise non-mutable struct) and then the conditional would look something like nobs(ve) >= min_nobs(ve).

Do you have a particular reason for preferring the Ref? I have the opposite preference, of avoiding mutable types and fields whenever possible. (I do find a single Ref preferable to having the whole type be mutable though.) Partially because I've seen it help performance, and partially because I find it easier to trust and understand code with as little mutable state as possible. The only downside I see is that you need to construct a lot of objects like

new_foo = ImmutableFooType(old_foo.a, old_foo.b, c_the_only_field_that_has_changed, old_foo.d, old_foo.e)

but there's no performance penalty for that and the code can be simplified with

new_foo = Accessors.@set old_foo.c = c_the_only_field_that_has_changed

Hmmm, I think it's mainly overkill and an overoptimization. Yes, it will probably be slightly more efficient, but it will also make interacting with the struct more awkward. For simple counters such as n in this case, a Ref is probably perfectly fine and not performance critical.

I do think the "better" use case for ...!! type functions is when you can clearly anticipate that the type will change from one call to the other.

I don't see it as overkill, because I see it as being simpler than using a Ref, and don't find it awkward at all. I just always think mutability makes things harder to reason about, and I feel safe and secure when surrounded my immutable things. It's like transcending our grimy, contingent reality and living in Plato's world of pure forms.

I do think the "better" use case for ...!! type functions is when you can clearly anticipate that the type will change from one call to the other.

I love me a !! function. I think most ! functions should actually be !! functions, to leave it up to the implementer of the type to choose whether the internal data structures are mutating or not.

Anyway, this is a bigger design conversation than this PR, and warrants an issue for a proper discussion.

Anyway, this is a bigger design conversation than this PR, and warrants an issue for a proper discussion.

I'm ready to die on this hill!

mhauru · 2025-11-21T18:28:37Z

src/adaptation/massmatrix.jl

+function Base.show(io::IO, ::NutpieVar{T}) where {T}
+    return print(io, "NutpieVar{", T, "} adaptor")
+end


The two-argument version of show should, according to the docs,

The representation used by show generally includes Julia-specific formatting and type information, and should be parseable Julia code when possible.

We break this rule all the time in TuringLang, so not too fussed about it, but I would still slightly prefer making a nice human readable version of show to be defined with the three-argument version.

Should we then simultaneously also fix e.g. WelfordVar (which was my template)?

Interestingly, the current state of the show methods is due to #466.

I'm happy to leave this as is, follow the current intra-package convention, and maybe open an issue about generally fixing our use of show. I didn't fully understand #466 at first glance, but it might be implying that something is calling print(x) when it should rather be calling something likeshow(io, MIME("text/plain"), x).

Yeah, I'm also unsure about #466. I'd also be in favour of opening another issue to potentially address any "weird" implementations.

mhauru · 2025-11-21T19:30:43Z

src/adaptation/massmatrix.jl

+
+function resize_adaptor!(nv::NutpieVar{T}, size_θ::Tuple{Int,Int}) where {T<:AbstractFloat}
+    if size_θ != size(nv.var)
+        @assert nv.n == 0 "Cannot resize a var estimator when it contains samples."


This looks like something that could plausibly be hit sometimes. Could it a throw error rather than an @assert? From the docstring of @assert:

│ Warning
│
│ An assert might be disabled at some optimization levels. Assert should therefore only be used as a debugging tool and
│ not used for authentication verification (e.g., verifying passwords or checking array bounds). The code must not rely
│ on the side effects of running cond for the correct behavior of a function.

I didn't know that! I'd assume as before, we might then also want to fix WelfordVar and friends?

Yep, I would propose changing those too. Can be a different PR though, don't mean to turn this into a general code style refactor.

mhauru · 2025-11-21T19:31:02Z

src/adaptation/massmatrix.jl

+function resize_adaptor!(nv::NutpieVar{T}, size_θ::Tuple{Int}) where {T<:AbstractFloat}
+    length_θ = first(size_θ)
+    if length_θ != size(nv.var, 1)
+        @assert nv.n == 0 "Cannot resize a var estimator when it contains samples."


mhauru · 2025-11-21T19:32:41Z

src/adaptation/massmatrix.jl

+end
+
+# Ref: TODO
+function get_estimation(nv::NutpieVar{T}) where {T<:AbstractFloat}


Suggested change

function get_estimation(nv::NutpieVar{T}) where {T<:AbstractFloat}

function get_estimation(nv::NutpieVar)

Same as above I think - I copied over the stuff from WelvordVar, but didn't get rid of superfluous code :(

mhauru · 2025-11-21T19:33:03Z

src/adaptation/stan_adaptor.jl

+    tp::StanHMCAdaptor,
+    z::PhasePoint,
+    α::AbstractScalarOrVec{<:AbstractFloat},
+)


Same question as above for the is_update argument.

Oops, sorry, I had somehow failed to click send on that comment. I meant this one: #473 (comment)

mhauru · 2025-11-21T19:34:12Z

src/sampler.jl

    adaptor::AbstractAdaptor,
    i::Int,
-    n_adapts::Int,
+    n_adapts::Int,    


Suggested change

n_adapts::Int,

n_adapts::Int,

Did your autoformatter catch that? I had seen at least one of these trailing whitespaces, but couldn't be bothered to fix them because I thought "where's one, there are many, and how would I find all of them...".

This one I just picked up reading the diff, since it's new. I do see that there's more trailing whitespace around the codebase that would warrant cleaning up.

Makes sense!

mhauru · 2025-11-21T19:37:34Z

src/sampler.jl

+    isadapted = false
+    if i <= n_adapts
+        i == 1 && Adaptation.initialize!(adaptor, n_adapts)
+        adapt!(adaptor, z, α)
+        i == n_adapts && finalize!(adaptor)
+        h = update(h, adaptor)
+        κ = update(κ, adaptor)
+        isadapted = true
+    end
+    return h, κ, isadapted


Suggested change

isadapted = false

if i <= n_adapts

i == 1 && Adaptation.initialize!(adaptor, n_adapts)

adapt!(adaptor, z, α)

i == n_adapts && finalize!(adaptor)

h = update(h, adaptor)

κ = update(κ, adaptor)

isadapted = true

end

return h, κ, isadapted

adapt = i <= n_adapts

if adapt

i == 1 && Adaptation.initialize!(adaptor, n_adapts)

adapt!(adaptor, z, α)

i == n_adapts && finalize!(adaptor)

h = update(h, adaptor)

κ = update(κ, adaptor)

end

return h, κ, adapt

Just a bit simpler, should be equivalent.

Maybe - do we then also want to change the function which I copied?

In my preference, yes but in a different PR. Not very fussed though, this is nitpicking about code style.

I'd say let's change the two functions in an eventual refactor then!

mhauru · 2025-11-21T19:47:20Z

docs/src/api.md

      + This is lowered to `UnitMassMatrix`, `WelfordVar` or `WelfordCov` based on the type of the mass matrix `metric`
+      + There is an experimental way to improve the *diagonal* mass matrix adaptation using gradient information (similar to [nutpie](https://github.com/pymc-devs/nutpie)), 
+      currently to be initialized for a `metric` of type `DiagEuclideanMetric`
+      via `mma = AdvancedHMC.Adaptation.NutpieVar(size(metric); var=copy(metric.M⁻¹))` 


Is the idea to not export any new functionality in this PR? Do that later together with a change of interface?

I thought so. But we can also export it - I've seen that AdvancedHMC does export e.g. WelfordVar as well.

I'm happy not to export and figure out the interface part later if you think that's better. I can see that it might be cleaner, having to do less changes to the interface if we collect them into a single bigger one. Just wanted to make sure I'm understanding correctly.

If we would export something new then we would need a version bump and a HISTORY.md entry explaining the change.

Do people ever just check out main and work with it to try out the newest experimental features? That's what I'd do!

I do anticipate the Nutpie interface and implementation changing soonish - but it would do so along with its relatives, so we'd need a breaking version bump anyways.

If it's bound to change soon, and this isn't urgently requested by any users, then I agree on not exporting anything now. Could we open an issue once this is merged to track finalising this and exporting it? Maybe leave a TODO note next to the struct definition linking to that issue too. I'm just thinking, if I find code that's not needed by any exported features, my inclination is to just go and delete it.

That would be this issue: #475
I've added it to the comment next to the new struct.

nsiccha · 2025-11-21T20:16:55Z

Thanks @mhauru! The main questions are AFAICT a) do we want to export NutpieVar?, b) do we want to fix the non-ideal code on which the code in this PR is based in the rest of the codebase as well?

We might want to do at least a slight refactor anyways when we change the interface - I guess we could then apply similar improvements both to the old code and the new code in this PR?

mhauru · 2025-11-24T11:09:21Z

src/adaptation/Adaptation.jl

+adapt!(
+    nca::NaiveHMCAdaptor,
+    z::PhasePoint,
+    α::AbstractScalarOrVec{<:AbstractFloat},


Is it intentional that this method of adapt! does not have the is_update argument at all? I'm not familiar with how the call stacks around any of this stuff works, but it made me wonder if somewhere adapt! might get called with that argument and NaiveHMCAdaptor, resulting in a MethodError.

Hmm, good question. To me, it looks like the is_update flag is exclusive to functions acting on MassMatrixAdaptors - so I guess this is intended (in the code which I copied to implement this function).

mhauru · 2025-11-24T11:14:06Z

do we want to fix the non-ideal code on which the code in this PR is based in the rest of the codebase as well?

I left the individual comments before reading this, but generally I'd be in favour of making the same changes elsewhere too, but in a different PR. I understand this may not be a priority though, and we may now just create an issue flagging it as something to do in the future. The show method I wouldn't change in this PR because I think it's something where being consistent within a package is valuable. The other things I'd personally change in the new code and flag to be changed in the old code later, but I'm not too bothered if you want to leave the new code consistent with old code and change all of it in one go later.

nsiccha · 2025-11-24T12:32:06Z

Great! I've implemented the minimal changes needed in the commits above.

mhauru

Thanks for addressing my earlier comments. I did a slower, more thoughtful read through and came up with more comments. Sorry to be annoying about this, I think I'm done after this round.

mhauru · 2025-11-25T10:21:41Z

src/adaptation/Adaptation.jl

    adapt!(nca.pc, θ, α)
    return nothing
 end
+adapt!(


I think this remains unformatted. At least when I run the formatter locally I get this diff:

- adapt!( - nca::NaiveHMCAdaptor, - z::PhasePoint, - α::AbstractScalarOrVec{<:AbstractFloat}, - ) = adapt!(nca, z.θ, α) + function adapt!( + nca::NaiveHMCAdaptor, z::PhasePoint, α::AbstractScalarOrVec{<:AbstractFloat} + ) + return adapt!(nca, z.θ, α) + end

Instead of pre-commit, if you prefer you can also do julia -e 'using JuliaFormatter; format(("src", "test"))', although that will format all files, not just the ones you edited.

src/adaptation/massmatrix.jl

src/adaptation/stan_adaptor.jl

src/sampler.jl

test/adaptation.jl

mhauru · 2025-11-25T11:26:29Z

test/adaptation.jl

                res = runnuts(ℓπ, DenseEuclideanMetric(D))
                @test res.adaptor.pc.cov ≈ Σ rtol = 0.25
            end
+            @test n_nutpie_superior > n_tests / 2


Could the right hand side here be larger? Currently, even if Nutpie was just normal NUTS this would pass with a 50% chance.

Maybe. I think Nupie was superior 90% of the time. With just 5 tests, NUTS should pass this test 18.75% of the time, while Nutpie (if it's true that it's superior roughly 90% of the time) should pass roughly 90% of the time - which actually sounds a bit low.

~~(compare ccdf(Binomial(5, .5), 3) (NUTS) vs ccdf(Binomial(5, .9), 3) (nutpie))~~

Edit: Whoops, tried to be careful about that CDF, but I think I messed up...

I do want to note that "nutpie" is NUTS. And that NUTS (IMO) is just the sampler, not the adaptation, i.e. there's nothing forcing anyone to use NUTS with dual averaging or with the mass matrix adaptation strategy as implemented in Stan, PyMC or Turing.jl.

In any case, IMO the "proper" way to test superiority of (in this case) nutpie would be to test it over more seeds, and for shorter runs, but I didn't want to mess with the previous tests too much. Also, I don't think we really know how often nutpie would be superior, or by how much - so what would be the frequency we test against?

Furthermore, I believe even a 50% test failure rate (if we were swapping in the standard adaptation instead of nutpie) would be pretty high, ~~as this test will be run many times, even for a single PR.~~ ah, I guess the seed will actually be equal across all runs, I did miss that!

Edit2: Huh, but shouldn't this way of (re)setting the global RNG in every iteration just result in exactly the same code/numbers being run/produced?

I've changed the threshold such that what I had written in the beginning but then struck through is actually true.

Edit2: Huh, but shouldn't this way of (re)setting the global RNG in every iteration just result in exactly the same code/numbers being run/produced?

Hadn't even realised that was here. I agree, it looks like it rendered the whole loop pointless. Nice spot.

In any case, IMO the "proper" way to test superiority of (in this case) nutpie would be to test it over more seeds, and for shorter runs, but I didn't want to mess with the previous tests too much. Also, I don't think we really know how often nutpie would be superior, or by how much - so what would be the frequency we test against?

Yeah, I agree that this could be made more robust. I was just concerned that if there was some bug that made Nutpie be equivalent to NUTS as it's usually used, this test would still pass typically. With the current limit I guess Nutpie needs to do better either 5/5 or 4/5 times, which feels like a non-trivial if not super robust test.

src/adaptation/massmatrix.jl

nsiccha · 2025-11-26T12:10:29Z

Once #479 is merged into main, I'll merge main into this branch and then we can also see what CI says.

mhauru

Looks good except for a question about exports.

mhauru · 2025-11-26T17:48:17Z

src/AdvancedHMC.jl

    MassMatrixAdaptor,
    UnitMassMatrix,
    WelfordVar,
+    NutpieVar,


Did you change your mind about whether to export anything? If yes, I think we should do a version bump and HISTORY.md entry.

Right! I forgot, I had wanted to highlight this also to you because I was unsure about exactly what you mention!

What to mention in HISTORY.md you mean? A description of what's been changed (or in this case, added), why/what it does, and how/when to use it. The HISTORY.md entry can have a basic explanation of the gist of the change, the way you would explain it to someone who asked about it in person, e.g. through an example if that feels helpful. You can refer to the docstring for all the details of what all the optional arguments are etc.

The most important HISTORY.md entries are the ones where something is being broken/removed, there we try to give clear instructions for how to cope with the change, like how to change your code that uses the feature that is being removed. That's obviously not relevant here though.

Note that the AHMC HISTORY.md doc isn't suuuper detailed yet, but we would like to slowly improve this across TuringLang. Currently we keep detailed notes for DynamicPPL and Turing, and the other packages are quite variable.

For the version bump, I think this can be a patch version bump.

What to mention in HISTORY.md you mean?

Yeah, exactly, and also whether that would be needed and whether we'd do a version bump. Would we then also directly trigger registration? And how'd Turing.jl be affected downstream? 🤔 Just wondering what the usual workflow is here 😅

Hm, I've added a line to the HISTORY.md, but I doubt that many people would actually use this newly exported feature. For one, because most people are probably interacting with AdvancedHMC via Turing, but also because it seems to be a bit convoluted currently to switch out the default adaptation for this one, see e.g. #473 (comment).

If we do want to export NutpieVar then I would do the version bump and immediate release, because otherwise I kinda don't see the point of exporting it in the first place, since presumably that is done to give users access to it. I think Turing should be unaffected since this doesn't break the existing interface (hence bumping just the patch version is fine) and we probably wouldn't want to use this new feature in Turing (before the interface rework).

I think it's up to you to decide if this is something AHMC should ship to users straight away, even if using it is a clunky, or wait for the interface changes. I'm happy to merge either way, though I think the NutpieVar docstring may still have a mention saying that it isn't exported, so that would need harmonising.

I'm happy to merge either way, though I think the NutpieVar docstring may still have a mention saying that it isn't exported, so that would need harmonising.

Oh, you're right!

test/adaptation.jl

mhauru · 2025-11-26T17:58:17Z

test/adaptation.jl

 preconditioned_cond(a::DiagMatrixEstimator, cov::AbstractMatrix) = cond(sqrt(Diagonal(a.var)) \ cov / sqrt(Diagonal(a.var)))

 @testset "Adaptation" begin
+    Random.seed!(1)


@testset actually resets the global RNG at the beginning of every @testset (source). Thus, I wonder if we should move this seed-setting to the very beginning of test/runtests.jl, and that should fix the seed to be same from run-to-run for the whole test suite.

oh wow ,that docstring is so long!

mhauru · 2025-11-26T18:08:35Z

test/adaptation.jl

                res = runnuts(ℓπ, DenseEuclideanMetric(D))
                @test res.adaptor.pc.cov ≈ Σ rtol = 0.25
            end
+            @test n_nutpie_superior > n_tests / 2


Edit2: Huh, but shouldn't this way of (re)setting the global RNG in every iteration just result in exactly the same code/numbers being run/produced?

Hadn't even realised that was here. I agree, it looks like it rendered the whole loop pointless. Nice spot.

In any case, IMO the "proper" way to test superiority of (in this case) nutpie would be to test it over more seeds, and for shorter runs, but I didn't want to mess with the previous tests too much. Also, I don't think we really know how often nutpie would be superior, or by how much - so what would be the frequency we test against?

Yeah, I agree that this could be made more robust. I was just concerned that if there was some bug that made Nutpie be equivalent to NUTS as it's usually used, this test would still pass typically. With the current limit I guess Nutpie needs to do better either 5/5 or 4/5 times, which feels like a non-trivial if not super robust test.

nsiccha · 2025-11-27T10:40:01Z

Do we maybe want to merge #479 first, @mhauru?

mhauru

Approving since I'm happy with the code, but would like to see CI pass before actually merging. I guess that's waiting on #479. EDIT: Just realised that's exactly what you said above...

mhauru · 2025-11-27T10:56:00Z

Oh wait, actually one of the new tests is failing, so that would need fixing first.

nsiccha · 2025-11-27T10:59:52Z

Ah, great! I thought I sometimes saw something like that, but it was a bit random whether any of the "actual" tests ran (the JET tests passed on the LTS), or whether the JET tests failed first on 1.12.

nsiccha · 2025-11-27T11:01:47Z

'Dense' MvNormal target: Test Failed at /Users/runner/work/AdvancedHMC.jl/AdvancedHMC.jl/test/adaptation.jl:216
  Expression: n_nutpie_superior > 1 + n_tests / 2
   Evaluated: 3 > 3.5

Oh no!

codecov · 2025-11-27T11:22:55Z

Codecov Report

❌ Patch coverage is 82.69231% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.71%. Comparing base (0dbc2ee) to head (e82cc19).

Files with missing lines	Patch %	Lines
src/adaptation/massmatrix.jl	76.92%	9 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #473      +/-   ##
==========================================
+ Coverage   77.58%   77.71%   +0.13%     
==========================================
  Files          21       21              
  Lines        1240     1270      +30     
==========================================
+ Hits          962      987      +25     
- Misses        278      283       +5

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

nsiccha · 2025-11-27T12:55:29Z

Could be merged now, @mhauru - but I do wonder how often nutpie is now better in the tests.

initial changes to get a working demo

caabd0f

github-actions bot assigned nsiccha Nov 19, 2025

nsiccha marked this pull request as draft November 19, 2025 12:05

This was referenced Nov 19, 2025

Draft of Nutpie/Nuts-rs mass matrix adaption #312

Closed

Make it easier to select the used adaptation #475

Open

nsiccha added 3 commits November 21, 2025 10:09

fix tests, add new ones, and add documentation

e030eec

fix type

7cdea36

fix some stray tests

989d284

nsiccha marked this pull request as ready for review November 21, 2025 09:18

delete tmp folder with demo

842fe06

mhauru reviewed Nov 21, 2025

View reviewed changes

mhauru reviewed Nov 24, 2025

View reviewed changes

nsiccha added 2 commits November 24, 2025 12:48

address review comments

0f27172

reference interface refactor issue

c395715

nsiccha mentioned this pull request Nov 24, 2025

Clean up internal interface and streamline a few implementations #477

Open

4 tasks

mhauru requested changes Nov 25, 2025

View reviewed changes

nsiccha added 5 commits November 26, 2025 10:18

refactor and fix test rng handling

f1d1c80

improve docstring for NutpieVar

76a1373

remove superfluous white space

45a1915

fix JET tests

3d47e8a

Merge branch 'fix-tests' into simple-nutpie

98c1ef7

mhauru reviewed Nov 26, 2025

View reviewed changes

nsiccha added 2 commits November 27, 2025 10:38

add entry to history, bump version

5888c0f

fix NutpieVar docstring

ca680bb

mhauru previously approved these changes Nov 27, 2025

View reviewed changes

Merge branch 'main' into simple-nutpie

4ee6cd4

increase number of tests for mass matrix adaptation

e82cc19

nsiccha dismissed mhauru’s stale review via e82cc19 November 27, 2025 11:04


		## Nutpie-style diagonal mass matrix estimator (using positions and gradients)

		mutable struct NutpieVar{T<:AbstractFloat,E<:AbstractVecOrMat{T},V<:AbstractVecOrMat{T}} <: DiagMatrixEstimator{T}

	function get_estimation(nv::NutpieVar{T}) where {T<:AbstractFloat}
	function get_estimation(nv::NutpieVar)

Implements a simple Nutpie style adaptation (using both positions and gradients, but not changing the schedule). #473

Are you sure you want to change the base?

Implements a simple Nutpie style adaptation (using both positions and gradients, but not changing the schedule). #473

Uh oh!

Conversation

nsiccha commented Nov 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nsiccha commented Nov 19, 2025

Uh oh!

nsiccha commented Nov 20, 2025

Uh oh!

nsiccha commented Nov 21, 2025

Uh oh!

nsiccha commented Nov 21, 2025

Uh oh!

mhauru left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nsiccha Nov 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

nsiccha commented Nov 19, 2025 •

edited

Loading

nsiccha Nov 24, 2025 •

edited

Loading

nsiccha Nov 24, 2025 •

edited

Loading