From 99948d2fbeaead68cf157c436e3156f936158a91 Mon Sep 17 00:00:00 2001
From: marius <mariusmillea@gmail.com>
Date: Thu, 30 Nov 2023 18:19:07 -0600
Subject: [PATCH] allow step-dependent step-size

---
 src/muse.jl | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/muse.jl b/src/muse.jl
index aebd81c..8604fd2 100644
--- a/src/muse.jl
+++ b/src/muse.jl
@@ -142,6 +142,12 @@ function muse!(
         save_MAPs = x -> nothing
     end
     
+    if !(α isa Base.Callable)
+        α = let α = α
+            i -> α
+        end
+    end
+
     xs_ẑs_sims = pmap(pool, split_rng(rng, nsims)) do rng
         (x, z) = sample_x_z(prob, rng, θ)
         (x, @something(z₀, z))
@@ -225,7 +231,7 @@ function muse!(
             )
 
             # Newton-Rhapson step
-            θunreg′ = θ′ .- T.(α) .* (H⁻¹_post′ * g_post′)
+            θunreg′ = θ′ .- T.(α(i)) .* (H⁻¹_post′ * g_post′)
             θunreg  = inv_transform_θ(prob, θunreg′)
             θ′ = regularize(θunreg′)
             θ  = inv_transform_θ(prob, θ′)