Allow passing time_step to StatsTools functions

RimuQMC · Jan 6, 2025 · 9b4dd0f · 9b4dd0f
1 parent 5f48b35
commit 9b4dd0f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/src/StatsTools/growth_witness.jl b/src/StatsTools/growth_witness.jl
@@ -51,7 +51,7 @@ function growth_witness(
     shift=:shift, norm=:norm, time_step=nothing, kwargs...
 )
     df = DataFrame(sim)
-    time_step = determine_constant_time_step(df)
+    time_step = isnothing(time_step) ? determine_constant_time_step(df) : time_step
 
     shift_vec = getproperty(df, Symbol(shift))
     norm_vec = getproperty(df, Symbol(norm))

diff --git a/src/StatsTools/reweighting.jl b/src/StatsTools/reweighting.jl
@@ -6,7 +6,7 @@
 Given a `DataFrame` `df`, determine the time step that was used to compute it. Throw an
 error if time step is not constant.
 """
-function determine_constant_time_step(df)
+function determine_constant_time_step(df, default)
     # Using get for backwards compatibility with old data frames
     if get(metadata(df), "time_step_strategy", "ConstantTimeStep()") == "ConstantTimeStep()"
         if haskey(metadata(df), "time_step")
@@ -20,7 +20,7 @@ function determine_constant_time_step(df)
         elseif hasproperty(df, "dτ_1")
             return df.dτ_1[end]
         else
-            throw(ArgumentError("Time step not found in `df`"))
+            throw(ArgumentError("key `\"time_step\"` not found in `df` metadata"))
         end
     else
         throw(ArgumentError("Time step not constant"))
@@ -243,13 +243,14 @@ function growth_estimator_analysis(
     shift_name=:shift,
     norm_name=:norm,
     warn=true,
+    time_step=nothing,
     kwargs...
 )
     df = DataFrame(sim)
     shift_v = Vector(getproperty(df, Symbol(shift_name))) # casting to `Vector` to make SIMD loops efficient
     norm_v = Vector(getproperty(df, Symbol(norm_name)))
     num_reps = length(filter(startswith("norm"), names(df)))
-    time_step = determine_constant_time_step(df)
+    time_step = isnothing(time_step) ? determine_constant_time_step(df) : time_step
     se = blocking_analysis(shift_v; skip)
     E_r = se.mean
     correlation_estimate = 2^(se.k - 1)
@@ -412,14 +413,15 @@ function mixed_estimator_analysis(
     hproj_name=:hproj,
     vproj_name=:vproj,
     warn=true,
+    time_step=nothing,
     kwargs...
 )
     shift_v = Vector(getproperty(df, Symbol(shift_name))) # casting to `Vector` to make SIMD loops efficient
     hproj_v = Vector(getproperty(df, Symbol(hproj_name)))
     vproj_v = Vector(getproperty(df, Symbol(vproj_name)))
     num_reps = length(filter(startswith("norm"), names(df)))
 
-    time_step = determine_constant_time_step(df)
+    time_step = isnothing(time_step) ? determine_constant_time_step(df) : time_step
     se = blocking_analysis(shift_v; skip)
     E_r = se.mean
     correlation_estimate = 2^(se.k - 1)
@@ -555,11 +557,12 @@ function rayleigh_replica_estimator(
     h=0,
     skip=0,
     Anorm=1,
+    time_step=nothing,
     kwargs...
 )
     df = DataFrame(sim)
     num_reps = length(filter(startswith("norm"), names(df)))
-    time_step = determine_constant_time_step(df)
+    time_step = isnothing(time_step) ? determine_constant_time_step(df) : time_step
     T = eltype(df[!, Symbol(shift_name, "_1")])
     shift_v = Vector{T}[]
     for a in 1:num_reps
@@ -625,11 +628,12 @@ function rayleigh_replica_estimator_analysis(
     vec_name="dot",
     Anorm=1,
     warn=true,
+    time_step=nothing,
     kwargs...
 )
     df = DataFrame(sim)
     num_reps = length(filter(startswith("norm"), names(df)))
-    time_step = determine_constant_time_step(df)
+    time_step = isnothing(time_step) ? determine_constant_time_step(df) : time_step
     # estimate the correlation time by blocking the shift data
     T = eltype(df[!, Symbol(shift_name, "_1")])
     shift_v = Vector{T}[]