Merge pull request #324 from LilithHafner/lh/format

Run JuliaFormatter.format()
SciML · Feb 15, 2024 · dbb96c9 · dbb96c9
2 parents 788dbf5 + 0709095
commit dbb96c9
Show file tree

Hide file tree

Showing 13 changed files with 173 additions and 161 deletions.
diff --git a/docs/make.jl b/docs/make.jl
@@ -8,21 +8,19 @@ ENV["GKSwstype"] = "100"
 include("pages.jl")
 
 makedocs(sitename = "DiffEqBayes.jl",
-         authors = "Chris Rackauckas, Vaibhav Kumar Dixit et al.",
-         clean = true,
-         doctest = false,
-         modules = [DiffEqBayes],
-         strict = [
-             :doctest,
-             :linkcheck,
-             :parse_error,
-             :example_block,
-             # Other available options are
-             # :autodocs_block, :cross_references, :docs_block, :eval_block, :example_block, :footnote, :meta_block, :missing_docs, :setup_block
-         ],
-         format = Documenter.HTML(assets = ["assets/favicon.ico"],
-                                  canonical = "https://docs.sciml.ai/DiffEqBayes/stable/"),
-         pages = pages)
+    authors = "Chris Rackauckas, Vaibhav Kumar Dixit et al.",
+    clean = true,
+    doctest = false,
+    modules = [DiffEqBayes],
+    strict = [
+        :doctest,
+        :linkcheck,
+        :parse_error,
+        :example_block        # Other available options are        # :autodocs_block, :cross_references, :docs_block, :eval_block, :example_block, :footnote, :meta_block, :missing_docs, :setup_block
+    ],
+    format = Documenter.HTML(assets = ["assets/favicon.ico"],
+        canonical = "https://docs.sciml.ai/DiffEqBayes/stable/"),
+    pages = pages)
 
 deploydocs(repo = "github.com/SciML/DiffEqBayes.jl.git";
-           push_preview = true)
+    push_preview = true)
diff --git a/docs/pages.jl b/docs/pages.jl
@@ -1,4 +1,4 @@
 pages = ["index.md",
     "Methods" => "methods.md",
-    "Examples" => ["examples.md", "examples/pendulum.md"],
+    "Examples" => ["examples.md", "examples/pendulum.md"]
 ]
diff --git a/docs/src/examples/pendulum.md b/docs/src/examples/pendulum.md
@@ -76,15 +76,15 @@ length of the pendulum L is probably around 3.0:
 ```@example pendulum
 priors = [
     truncated(Normal(0.1, 1.0), lower = 0.0),
-    truncated(Normal(3.0, 1.0), lower = 0.0),
+    truncated(Normal(3.0, 1.0), lower = 0.0)
 ]
 ```
 
 Finally, let's run the estimation routine from DiffEqBayes.jl with the Turing.jl backend to check if we indeed recover the parameters!
 
 ```@example pendulum
 bayesian_result = turing_inference(prob1, Tsit5(), t, data, priors; num_samples = 10_000,
-                                   syms = [:omega, :L])
+    syms = [:omega, :L])
 ```
 
 Notice that while our guesses had the wrong means, the learned parameters converged
@@ -120,15 +120,15 @@ to get better understanding of the performance.
 
 ```@example pendulum
 @btime bayesian_result = turing_inference(prob1, Tsit5(), t, data, priors;
-                                          syms = [:omega, :L], num_samples = 10_000)
+    syms = [:omega, :L], num_samples = 10_000)
 ```
 
 ```@example pendulum
 @btime bayesian_result = stan_inference(prob1, t, data, priors; num_samples = 10_000,
-                                        print_summary = false)
+    print_summary = false)
 ```
 
 ```@example pendulum
 @btime bayesian_result = dynamichmc_inference(prob1, Tsit5(), t, data, priors;
-                                              num_samples = 10_000)
+    num_samples = 10_000)
 ```
diff --git a/docs/src/methods.md b/docs/src/methods.md
@@ -12,9 +12,9 @@ using DiffEqBayes
 
 ```julia
 stan_inference(prob::ODEProblem, t, data, priors = nothing; alg = :rk45,
-               num_samples = 1000, num_warmups = 1000, reltol = 1e-3,
-               abstol = 1e-6, maxiter = Int(1e5), likelihood = Normal,
-               vars = (StanODEData(), InverseGamma(2, 3)))
+    num_samples = 1000, num_warmups = 1000, reltol = 1e-3,
+    abstol = 1e-6, maxiter = Int(1e5), likelihood = Normal,
+    vars = (StanODEData(), InverseGamma(2, 3)))
 ```
 
 `stan_inference` uses [Stan.jl](https://stanjulia.github.io/CmdStan.jl/latest/INTRO/)
@@ -37,8 +37,8 @@ parameter list.
 
 ```julia
 function turing_inference(prob::DiffEqBase.DEProblem, alg, t, data, priors;
-                          likelihood_dist_priors, likelihood, num_samples = 1000,
-                          sampler = Turing.NUTS(num_samples, 0.65), parallel_type = MCMCSerial(), n_chains = 1, syms, kwargs...)
+        likelihood_dist_priors, likelihood, num_samples = 1000,
+        sampler = Turing.NUTS(num_samples, 0.65), parallel_type = MCMCSerial(), n_chains = 1, syms, kwargs...)
 end
 ```
 
@@ -55,7 +55,7 @@ type. `num_samples` is the number of samples per MCMC chain. Sampling from multi
 
 ```julia
 dynamichmc_inference(prob::DEProblem, alg, t, data, priors, transformations;
-                     σ = 0.01, ϵ = 0.001, initial = Float64[])
+    σ = 0.01, ϵ = 0.001, initial = Float64[])
 ```
 
 `dynamichmc_inference` uses [DynamicHMC.jl](https://github.com/tpapp/DynamicHMC.jl) to
@@ -71,8 +71,8 @@ parameter values to specific domains. `initial` values for the parameters can be
 
 ```julia
 abc_inference(prob::DEProblem, alg, t, data, priors; ϵ = 0.001,
-              distancefunction = euclidean, ABCalgorithm = ABCSMC, progress = false,
-              num_samples = 500, maxiterations = 10^5, kwargs...)
+    distancefunction = euclidean, ABCalgorithm = ABCSMC, progress = false,
+    num_samples = 500, maxiterations = 10^5, kwargs...)
 ```
 
 `abc_inference` uses [ApproxBayes.jl](https://github.com/marcjwilliams1/ApproxBayes.jl), which uses Approximate Bayesian Computation (ABC) to

diff --git a/src/abc_inference.jl b/src/abc_inference.jl
@@ -1,5 +1,5 @@
 function createabcfunction(prob, t, distancefunction, alg; save_idxs = nothing,
-                           sample_u0 = false, kwargs...)
+        sample_u0 = false, kwargs...)
     function simfunc(params, constants, data)
         local u0
         if sample_u0
@@ -14,7 +14,7 @@ function createabcfunction(prob, t, distancefunction, alg; save_idxs = nothing,
             u0 = prob.u0
         end
         sol = solve(prob, alg, u0 = u0, p = params, saveat = t, save_idxs = save_idxs,
-                    kwargs...)
+            kwargs...)
         if size(sol, 2) < length(t)
             return Inf, nothing
         else
@@ -25,18 +25,19 @@ function createabcfunction(prob, t, distancefunction, alg; save_idxs = nothing,
 end
 
 function abc_inference(prob::DiffEqBase.DEProblem, alg, t, data, priors; ϵ = 0.001,
-                       distancefunction = euclidean, ABCalgorithm = ABCSMC,
-                       progress = false,
-                       num_samples = 500, maxiterations = 10^5, save_idxs = nothing,
-                       sample_u0 = false, parallel = false, kwargs...)
-    abcsetup = ABCalgorithm(createabcfunction(prob, t, distancefunction, alg;
-                                              save_idxs = save_idxs, sample_u0 = sample_u0,
-                                              kwargs...),
-                            length(priors),
-                            ϵ,
-                            ApproxBayes.Prior(priors);
-                            nparticles = num_samples,
-                            maxiterations = maxiterations)
+        distancefunction = euclidean, ABCalgorithm = ABCSMC,
+        progress = false,
+        num_samples = 500, maxiterations = 10^5, save_idxs = nothing,
+        sample_u0 = false, parallel = false, kwargs...)
+    abcsetup = ABCalgorithm(
+        createabcfunction(prob, t, distancefunction, alg;
+            save_idxs = save_idxs, sample_u0 = sample_u0,
+            kwargs...),
+        length(priors),
+        ϵ,
+        ApproxBayes.Prior(priors);
+        nparticles = num_samples,
+        maxiterations = maxiterations)
 
     abcresult = runabc(abcsetup, data, progress = progress, parallel = parallel)
     return abcresult

diff --git a/src/dynamichmc_inference.jl b/src/dynamichmc_inference.jl
@@ -44,12 +44,13 @@ function (P::DynamicHMCPosterior)(θ)
         end
     end
     _saveat = t === nothing ? Float64[] : t
-    sol = solve(problem, algorithm; u0 = u0, p = p, saveat = _saveat, save_idxs = save_idxs,
-                solve_kwargs...)
+    sol = solve(
+        problem, algorithm; u0 = u0, p = p, saveat = _saveat, save_idxs = save_idxs,
+        solve_kwargs...)
     failure = size(sol, 2) < length(_saveat)
     failure && return T(0) * sum(σ) + T(-Inf)
     log_likelihood = sum(sum(map(logpdf, Normal.(0.0, σ), sol[:, i] .- data[:, i]))
-                         for (i, t) in enumerate(t))
+    for (i, t) in enumerate(t))
     log_prior_parameters = sum(map(logpdf, parameter_priors, parameters))
     log_prior_σ = sum(map(logpdf, σ_priors, σ))
     log_likelihood + log_prior_parameters + log_prior_σ
@@ -91,24 +92,24 @@ posterior values (transformed from `ℝⁿ`).
 - `mcmc_kwargs` are passed on as keyword arguments to `DynamicHMC.mcmc_with_warmup`
 """
 function dynamichmc_inference(problem::DiffEqBase.DEProblem, algorithm, t, data,
-                              parameter_priors,
-                              parameter_transformations = as(Vector, asℝ₊,
-                                                             length(parameter_priors));
-                              σ_priors = fill(Normal(0, 5), size(data, 1)),
-                              sample_u0 = false, rng = Random.GLOBAL_RNG,
-                              num_samples = 1000, AD_gradient_kind = Val(:ForwardDiff),
-                              save_idxs = nothing, solve_kwargs = (),
-                              mcmc_kwargs = (initialization = (q = zeros(length(parameter_priors) +
-                                                                         (save_idxs ===
-                                                                          nothing ?
-                                                                          length(data[:, 1]) :
-                                                                          length(save_idxs))),),))
+        parameter_priors,
+        parameter_transformations = as(Vector, asℝ₊,
+            length(parameter_priors));
+        σ_priors = fill(Normal(0, 5), size(data, 1)),
+        sample_u0 = false, rng = Random.GLOBAL_RNG,
+        num_samples = 1000, AD_gradient_kind = Val(:ForwardDiff),
+        save_idxs = nothing, solve_kwargs = (),
+        mcmc_kwargs = (initialization = (q = zeros(length(parameter_priors) +
+                                                   (save_idxs ===
+                                                    nothing ?
+                                                    length(data[:, 1]) :
+                                                    length(save_idxs))),),))
     P = DynamicHMCPosterior(; algorithm = algorithm, problem = problem, t = t, data = data,
-                            parameter_priors = parameter_priors, σ_priors = σ_priors,
-                            solve_kwargs = solve_kwargs, sample_u0 = sample_u0,
-                            save_idxs = save_idxs)
+        parameter_priors = parameter_priors, σ_priors = σ_priors,
+        solve_kwargs = solve_kwargs, sample_u0 = sample_u0,
+        save_idxs = save_idxs)
     trans = as((parameters = parameter_transformations,
-                σ = as(Vector, asℝ₊, length(σ_priors))))
+        σ = as(Vector, asℝ₊, length(σ_priors))))
     ℓ = TransformedLogDensity(trans, P)
     ∇ℓ = LogDensityProblemsAD.ADgradient(AD_gradient_kind, ℓ)
     results = mcmc_with_warmup(rng, ∇ℓ, num_samples; mcmc_kwargs...)

diff --git a/src/stan_inference.jl b/src/stan_inference.jl
@@ -20,7 +20,7 @@ function generate_priors(n, priors)
     else
         for i in 1:n
             priors_string = string(priors_string, "theta_$i ~ ", stan_string(priors[i]),
-                                   ";")
+                ";")
         end
     end
     priors_string
@@ -39,7 +39,7 @@ function generate_theta(n, priors)
         end
         if lower_bound != "" && upper_bound != ""
             theta = string(theta, "real", "<$lower_bound", ",", "$upper_bound>",
-                           " theta_$i", ";")
+                " theta_$i", ";")
         elseif lower_bound != ""
             theta = string(theta, "real", "<$lower_bound", ">", " theta_$i", ";")
         elseif upper_bound != ""
@@ -52,23 +52,23 @@ function generate_theta(n, priors)
 end
 
 function stan_inference(prob::DiffEqBase.DEProblem,
-                        # Positional arguments
-                        t, data, priors = nothing, stanmodel = nothing;
-                        # DiffEqBayes keyword arguments
-                        likelihood = Normal, vars = (StanODEData(), InverseGamma(3, 3)),
-                        sample_u0 = false, save_idxs = nothing, diffeq_string = nothing,
-                        # Stan differential equation function keyword arguments
-                        alg = :rk45, reltol = 1e-3, abstol = 1e-6, maxiter = Int(1e5),
-                        # stan_sample keyword arguments
-                        num_samples = 1000, num_warmups = 1000,
-                        num_cpp_chains = 1, num_chains = 1, num_threads = 1,
-                        delta = 0.8,
-                        # read_samples arguments
-                        output_format = :mcmcchains,
-                        # read_summary arguments
-                        print_summary = true,
-                        # pass in existing tmpdir
-                        tmpdir = mktempdir())
+        # Positional arguments
+        t, data, priors = nothing, stanmodel = nothing;
+        # DiffEqBayes keyword arguments
+        likelihood = Normal, vars = (StanODEData(), InverseGamma(3, 3)),
+        sample_u0 = false, save_idxs = nothing, diffeq_string = nothing,
+        # Stan differential equation function keyword arguments
+        alg = :rk45, reltol = 1e-3, abstol = 1e-6, maxiter = Int(1e5),
+        # stan_sample keyword arguments
+        num_samples = 1000, num_warmups = 1000,
+        num_cpp_chains = 1, num_chains = 1, num_threads = 1,
+        delta = 0.8,
+        # read_samples arguments
+        output_format = :mcmcchains,
+        # read_summary arguments
+        print_summary = true,
+        # pass in existing tmpdir
+        tmpdir = mktempdir())
     save_idxs !== nothing && length(save_idxs) == 1 ? save_idxs = save_idxs[1] :
     save_idxs = save_idxs
     length_of_y = length(prob.u0)
@@ -113,7 +113,7 @@ function stan_inference(prob::DiffEqBase.DEProblem,
                 hyper_params = string(hyper_params, "sigma$(i-1) ~ $dist;")
                 tuple_hyper_params = string(tuple_hyper_params, "sigma$(i-1)", ",")
                 setup_params = string(setup_params,
-                                      "row_vector<lower=0>[$(length(save_idxs))] sigma$(i-1);")
+                    "row_vector<lower=0>[$(length(save_idxs))] sigma$(i-1);")
             end
         end
 
@@ -157,12 +157,12 @@ function stan_inference(prob::DiffEqBase.DEProblem,
 
         if isnothing(diffeq_string)
             diffeq_string = ModelingToolkit.build_function(ModelingToolkit.equations(sys),
-                                                           ModelingToolkit.states(sys),
-                                                           ModelingToolkit.parameters(sys),
-                                                           ModelingToolkit.get_iv(sys);
-                                                           expression = Val{true},
-                                                           fname = :sho,
-                                                           target = ModelingToolkit.StanTarget())
+                ModelingToolkit.states(sys),
+                ModelingToolkit.parameters(sys),
+                ModelingToolkit.get_iv(sys);
+                expression = Val{true},
+                fname = :sho,
+                target = ModelingToolkit.StanTarget())
         end
 
         parameter_estimation_model = "
@@ -192,16 +192,16 @@ function stan_inference(prob::DiffEqBase.DEProblem,
             "
 
         stanmodel = SampleModel("parameter_estimation_model",
-                                parameter_estimation_model,
-                                tmpdir)
+            parameter_estimation_model,
+            tmpdir)
     end
 
     data = Dict("u0" => prob.u0, "T" => length(t),
-                "internal_var___u" => view(data, :, 1:length(t)),
-                "t0" => prob.tspan[1], "ts" => t)
+        "internal_var___u" => view(data, :, 1:length(t)),
+        "t0" => prob.tspan[1], "ts" => t)
 
     @time rc = stan_sample(stanmodel; data, num_threads, num_cpp_chains,
-                           num_samples, num_warmups, num_chains, delta)
+        num_samples, num_warmups, num_chains, delta)
 
     if success(rc)
         return StanResult(stanmodel, rc, read_samples(stanmodel, output_format))

diff --git a/src/turing_inference.jl b/src/turing_inference.jl
@@ -1,21 +1,21 @@
 function turing_inference(prob::DiffEqBase.DEProblem,
-                          alg,
-                          t,
-                          data,
-                          priors;
-                          likelihood_dist_priors = [InverseGamma(2, 3)],
-                          likelihood = (u, p, t, σ) -> MvNormal(u,
-                                                                Diagonal((σ[1])^2 *
-                                                                         ones(length(u)))),
-                          num_samples = 1000,
-                          sampler = Turing.NUTS(0.65),
-                          parallel_type = MCMCSerial(),
-                          n_chains = 1,
-                          syms = [Turing.@varname(theta[i]) for i in 1:length(priors)],
-                          sample_u0 = false,
-                          save_idxs = nothing,
-                          progress = false,
-                          kwargs...)
+        alg,
+        t,
+        data,
+        priors;
+        likelihood_dist_priors = [InverseGamma(2, 3)],
+        likelihood = (u, p, t, σ) -> MvNormal(u,
+            Diagonal((σ[1])^2 *
+                     ones(length(u)))),
+        num_samples = 1000,
+        sampler = Turing.NUTS(0.65),
+        parallel_type = MCMCSerial(),
+        n_chains = 1,
+        syms = [Turing.@varname(theta[i]) for i in 1:length(priors)],
+        sample_u0 = false,
+        save_idxs = nothing,
+        progress = false,
+        kwargs...)
     N = length(priors)
     Turing.@model function infer(x, ::Type{T} = Float64) where {T <: Real}
         theta = Vector{T}(undef, length(priors))
@@ -37,7 +37,7 @@ function turing_inference(prob::DiffEqBase.DEProblem,
         end
         _saveat = t === nothing ? Float64[] : t
         sol = solve(prob, alg; u0 = u0, p = p, saveat = _saveat, progress = progress,
-                    save_idxs = save_idxs, kwargs...)
+            save_idxs = save_idxs, kwargs...)
         failure = size(sol, 2) < length(_saveat)
 
         if failure