feat: Allow time step configuration

Allow configuration of the time step through a time_step_options field in `SimulationConfig`. One can use multiple dispatch to define an object which overloads the definition of `max_time_step` and/or `actual_time_step`. BREAKING CHANGE: SimulationConfig has two fields replaced by `time_step_options` and two constructor methods removed.
musoke · Sep 26, 2021 · 8abcb35 · 8abcb35 · musoke · Sep 26, 2021
1 parent c2437af
commit 8abcb35
Show file tree

Hide file tree

Showing 11 changed files with 82 additions and 51 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "UltraDark"
 uuid = "1c8d022d-dfc0-4b41-80ab-3fc7e88cdfea"
 authors = ["Nathan Musoke <[email protected]>"]
-version = "0.1.4"
+version = "0.2.0"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

diff --git a/examples/2d.ipynb b/examples/2d.ipynb
@@ -223,7 +223,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "options = Config.SimulationConfig(10);"
+    "options = Config.SimulationConfig();"
    ]
   },
   {

diff --git a/examples/growth.ipynb b/examples/growth.ipynb
@@ -699,7 +699,7 @@
     "output_times = t_init:(t_final-t_init)/num_saves:t_final\n",
     "output_config = OutputConfig(output_dir, output_times, summary_statistics=UltraDark.Output.SummaryStatisticsMeanMaxRms)\n",
     "\n",
-    "options = Config.SimulationConfig(10, a)"
+    "options = Config.SimulationConfig(a=a)"
    ]
   },
   {

diff --git a/examples/soliton_velocity.jl b/examples/soliton_velocity.jl
@@ -27,7 +27,7 @@ const output_dir = "output/vel"
 function run_sim()
 
     output_config = OutputConfig(output_dir, output_times; box=true, slice=false)
-    options = Config.SimulationConfig(10)
+    options = Config.SimulationConfig()
 
     grids = PencilGrids(10.0, resol)
     @info "Grids created" MPI.Comm_rank(comm) size(grids.ψk)

diff --git a/src/UltraDark.jl b/src/UltraDark.jl
@@ -9,7 +9,7 @@ using FFTW
 
 export simulate
 export Grids, PencilGrids
-export Config, SimulationConfig, constant_scale_factor
+export Config, SimulationConfig, constant_scale_factor, TimeStepOptions
 export OutputConfig
 
 include("grids.jl")
@@ -21,7 +21,7 @@ include("config.jl")
 
 import .Output: OutputConfig
 import .Output: output_summary_row, output_summary_header, output_grids
-import .Config: SimulationConfig, constant_scale_factor
+import .Config: SimulationConfig, constant_scale_factor, TimeStepOptions
 
 function outer_step!(Δt, grids, constants; a=1.0)
     psi_whole_step!(Δt, grids, constants)
@@ -155,10 +155,9 @@ function evolve_to!(t_start, t_end, grids, output_config, sim_config; constants=
         end
 
         Δt, n_steps = actual_time_step(
-            max_time_step(grids, sim_config.a(t)),
+            max_time_step(grids, sim_config.a(t), sim_config.time_step_options),
             t_end - t,
-            sim_config.time_step_update_period,
-            sim_config.time_step_multiplier,
+            sim_config.time_step_options,
         )
 
         t = take_steps!(grids, t, Δt, n_steps, output_config, sim_config.a, constants)

diff --git a/src/config.jl b/src/config.jl
@@ -1,37 +1,62 @@
 module Config
 
 export SimulationConfig, constant_scale_factor
+export TimeStepOptions
 
 struct SimulationConfig
-    time_step_update_period::Int64
-    time_step_multiplier::Float64
     "function defining a(t)"
     a
     phase_grad_limit::Float64
     density_threshold::Float64
+    time_step_options
 end
 
 function constant_scale_factor(t)
     1.0
 end
 
-function SimulationConfig(time_step_update_period)
-    SimulationConfig(time_step_update_period=time_step_update_period)
-end
-
-function SimulationConfig(time_step_update_period, a)
-    SimulationConfig(time_step_update_period=time_step_update_period, a=a)
-end
-
 function SimulationConfig(
                           ;
-                          time_step_update_period=10,
-                          time_step_multiplier=1,
                           a=constant_scale_factor,
                           phase_grad_limit=π/4,
-                          density_threshold=1e-6
+                          density_threshold=1e-6,
+                          time_step_options=TimeStepOptions(),
                          )
-    SimulationConfig(time_step_update_period, time_step_multiplier, constant_scale_factor, phase_grad_limit, density_threshold)
+    SimulationConfig(constant_scale_factor, phase_grad_limit, density_threshold, time_step_options)
+end
+
+"""
+    TimeStepOptions(update_period=10, multiplier=1.0)
+
+struct containing options controlling the size and calculation of time steps.
+
+update_period::Int64 controls how many steps are taken before the timestep is
+updated.
+
+multiplier::Float64 multiplies the calculated maximum time step by a constant
+
+See also: [`SimulationConfig`](@
+
+# Examples
+```jldoctest
+julia> using UltraDark
+
+julia> TimeStepOptions()
+TimeStepOptions(10, 1.0)
+
+```
+"""
+struct TimeStepOptions
+    update_period::Int64
+    multiplier::Float64
+end
+
+function TimeStepOptions(
+                         ;
+                         update_period=10,
+                         multiplier=1.0,
+                        )
+    TimeStepOptions(update_period, multiplier)
 end
 
 end # module
diff --git a/src/time_step.jl b/src/time_step.jl
@@ -1,26 +1,26 @@
 """
-    actual_time_step(max_time_step, time_interval, n::Integer)
+    actual_time_step(max_time_step, time_interval, time_step_options)
 
-Actual size and number of time steps that should be taken if the maximum 
-is `max_time_step`, no more than `n` steps should be taken, and they should
-fit in `time_interval`.
+Actual size and number of time steps that should be taken if the maximum
+is `max_time_step`. No more than `time_step_options.update_period` steps should
+be taken, and they should fit in `time_interval`.
 
 # Examples
 
 ```jldoctest
-julia> using UltraDark: actual_time_step
+julia> using UltraDark: actual_time_step, TimeStepOptions
 
-julia> actual_time_step(0.11, 1, 20, 1.0)
+julia> actual_time_step(0.11, 1, TimeStepOptions())
 (0.1, 10)
 ```
 """
-function actual_time_step(max_time_step, time_interval, n, multiplier)::Tuple{Float64, Integer}
-    max_time_step *= multiplier
-    if max_time_step * n > time_interval
+function actual_time_step(max_time_step, time_interval, time_step_options)::Tuple{Float64, Integer}
+    max_time_step *= time_step_options.multiplier
+    if max_time_step * time_step_options.update_period > time_interval
         num_steps = ceil(time_interval / max_time_step)
         time_interval / num_steps, num_steps
     else
-        max_time_step, n
+        max_time_step, time_step_options.update_period
     end
 end
 
@@ -29,7 +29,7 @@ end
 
 Calculate an upper bound on the time step
 """
-function max_time_step(grids, a)
+function max_time_step(grids, a, time_step_options)
     max_time_step_gravity = 2π / maximum(abs.(grids.Φx))
     max_time_step_pressure = 2π * 2 / maximum(grids.k)^2 * a^2  # TODO: cache k_max
 
@@ -46,7 +46,7 @@ end
 
 Calculate an upper bound on the time step for `PencilGrids`
 """
-function max_time_step(grids::PencilGrids, a)
+function max_time_step(grids::PencilGrids, a, time_step_options)
     # Find maximum on local grid
     local_max_time_step_gravity = 2π / maximum(abs.(grids.Φx))
     local_max_time_step_pressure = 2π * 2 / maximum(grids.k)^2 * a^2  # TODO: cache k_max

diff --git a/test/evolution.jl b/test/evolution.jl
@@ -1,25 +1,32 @@
 using UltraDark
+using UltraDark: TimeStepOptions
 using Test
 using MPI
 
 @testset "Actual time step" begin
-    @test UltraDark.actual_time_step(0.5, 1, 2, 1.0)[1] ≈ 0.5
-    @test UltraDark.actual_time_step(0.5, 1, 2, 1.0)[2] == 2
+    time_step_options = TimeStepOptions(update_period=2)
+    @test UltraDark.actual_time_step(0.5, 1, time_step_options)[1] ≈ 0.5
+    @test UltraDark.actual_time_step(0.5, 1, time_step_options)[2] == 2
 
-    @test UltraDark.actual_time_step(0.5, 1, 3, 1.0)[1] ≈ 0.5
-    @test UltraDark.actual_time_step(0.5, 1, 3, 1.0)[2] == 2
+    time_step_options = TimeStepOptions(update_period=3)
+    @test UltraDark.actual_time_step(0.5, 1, time_step_options)[1] ≈ 0.5
+    @test UltraDark.actual_time_step(0.5, 1, time_step_options)[2] == 2
 
-    @test UltraDark.actual_time_step(0.01, 1, 10, 1.0)[1] ≈ 0.01
-    @test UltraDark.actual_time_step(0.01, 1, 10, 1.0)[2] == 10
+    time_step_options = TimeStepOptions()
+    @test UltraDark.actual_time_step(0.01, 1, time_step_options)[1] ≈ 0.01
+    @test UltraDark.actual_time_step(0.01, 1, time_step_options)[2] == 10
 
-    @test UltraDark.actual_time_step(0.01, 1, 10, 10.)[1] ≈ 0.1
-    @test UltraDark.actual_time_step(0.01, 1, 10, 10.)[2] == 10
+    time_step_options = TimeStepOptions(multiplier=10.0)
+    @test UltraDark.actual_time_step(0.01, 1, time_step_options)[1] ≈ 0.1
+    @test UltraDark.actual_time_step(0.01, 1, time_step_options)[2] == 10
 
-    @test UltraDark.actual_time_step(0.7, 1, 2, 1.0)[1] ≈ 0.5
-    @test UltraDark.actual_time_step(0.7, 1, 2, 1.0)[2] == 2
+    time_step_options = TimeStepOptions(update_period=2)
+    @test UltraDark.actual_time_step(0.7, 1, time_step_options)[1] ≈ 0.5
+    @test UltraDark.actual_time_step(0.7, 1, time_step_options)[2] == 2
 
-    @test UltraDark.actual_time_step(0.7, 2, 3, 1.0)[1] ≈ 2/3
-    @test UltraDark.actual_time_step(0.7, 2, 3, 1.0)[2] == 3
+    time_step_options = TimeStepOptions(update_period=3)
+    @test UltraDark.actual_time_step(0.7, 2, time_step_options)[1] ≈ 2/3
+    @test UltraDark.actual_time_step(0.7, 2, time_step_options)[2] == 3
 end
 
 for grid_type in [Grids, PencilGrids]
@@ -44,7 +51,7 @@ for grid_type in [Grids, PencilGrids]
         grids.Φx .= grids.rfft_plan \ grids.Φk
 
         output_config = OutputConfig(mktempdir(), [])
-        options = UltraDark.Config.SimulationConfig(10)
+        options = UltraDark.Config.SimulationConfig()
         @test UltraDark.evolve_to!(t_begin, t_end, grids, output_config, options) ≈ t_end
     end
 end
diff --git a/test/phase_gradient.jl b/test/phase_gradient.jl
@@ -82,7 +82,7 @@ end
 
     grids_orig = deepcopy(grids)
 
-    options = UltraDark.Config.SimulationConfig(10)
+    options = UltraDark.Config.SimulationConfig()
     output_config = OutputConfig(mktempdir(), [1, 2])
 
     @test_throws "Phase gradient is too large to start" simulate(grids, options, output_config)

diff --git a/test/sims/full.jl b/test/sims/full.jl
@@ -9,7 +9,7 @@ output_dir = "output"
 output_times = 0.0:0.002:0.01
 
 output_config = OutputConfig(output_dir, output_times; box=false)
-options = Config.SimulationConfig(10)
+options = Config.SimulationConfig()
 
 for grid_type in [Grids, PencilGrids]
     grids = grid_type(1.0, 16)

diff --git a/test/sims/soliton_static.jl b/test/sims/soliton_static.jl
@@ -16,7 +16,7 @@ for grid_type in [Grids, PencilGrids]
     output_times = [0, 1]
 
     output_config = OutputConfig(this_output_dir, output_times; box=true, slice=false)
-    options = Config.SimulationConfig(10)
+    options = Config.SimulationConfig()
 
     grids = grid_type(10.0, resol)