Merge pull request #998 from ChrisRackauckas-Claude/runic-formatting

Switch from JuliaFormatter to Runic.jl for code formatting

Author: ChrisRackauckas

.JuliaFormatter.toml

@@ -1,9 +1,19 @@
-name: Format suggestions
+name: format-check
 
-on: [pull_request]
+on:
+  push:
+    branches:
+      - 'master'
+      - 'main'
+      - 'release-'
+    tags: '*'
+  pull_request:
 
 jobs:
-  code-style:
+  runic:
     runs-on: ubuntu-latest
     steps:
-      - uses: julia-actions/julia-format@v4
+      - uses: actions/checkout@v4
+      - uses: fredrikekre/runic-action@v1
+        with:
+          version: '1'

.github/workflows/FormatCheck.yml

@@ -9,15 +9,19 @@ ENV["DATADEPS_ALWAYS_ACCEPT"] = true
 
 include("pages.jl")
 
-makedocs(; sitename = "DiffEqFlux.jl",
+makedocs(;
+    sitename = "DiffEqFlux.jl",
     authors = "Chris Rackauckas et al.",
     clean = true,
     doctest = false,
     # linkcheck = true,
     warnonly = [:docs_block, :missing_docs, :linkcheck],
     modules = [DiffEqFlux],
-    format = Documenter.HTML(; assets = ["assets/favicon.ico"],
-        canonical = "https://docs.sciml.ai/DiffEqFlux/stable/"),
-    pages = pages)
+    format = Documenter.HTML(;
+        assets = ["assets/favicon.ico"],
+        canonical = "https://docs.sciml.ai/DiffEqFlux/stable/"
+    ),
+    pages = pages
+)
 
 deploydocs(; repo = "github.com/SciML/DiffEqFlux.jl.git", push_preview = true)

docs/make.jl

@@ -5,7 +5,8 @@ using DiffEqFlux: DiffEqFlux
 
 @views function DiffEqFlux.collocate_data(
         data::AbstractMatrix{T}, tpoints::AbstractVector{T},
-        tpoints_sample::AbstractVector{T}, interp, args...) where {T}
+        tpoints_sample::AbstractVector{T}, interp, args...
+    ) where {T}
     u = zeros(T, size(data, 1), length(tpoints_sample))
     du = zeros(T, size(data, 1), length(tpoints_sample))
     for d1 in axes(data, 1)

ext/DiffEqFluxDataInterpolationsExt.jl

@@ -11,14 +11,14 @@ using LuxLib: batched_matmul
 using Random: Random, AbstractRNG, randn!
 using Reexport: @reexport
 using SciMLBase: SciMLBase, DAEProblem, DDEFunction, DDEProblem, EnsembleProblem,
-                 ODEFunction, ODEProblem, ODESolution, SDEFunction, SDEProblem, remake,
-                 solve
+    ODEFunction, ODEProblem, ODESolution, SDEFunction, SDEProblem, remake,
+    solve
 using SciMLSensitivity: SciMLSensitivity, AdjointLSS, BacksolveAdjoint, EnzymeVJP,
-                        ForwardDiffOverAdjoint, ForwardDiffSensitivity, ForwardLSS,
-                        ForwardSensitivity, GaussAdjoint, InterpolatingAdjoint, NILSAS,
-                        NILSS, QuadratureAdjoint, ReverseDiffAdjoint, ReverseDiffVJP,
-                        SteadyStateAdjoint, TrackerAdjoint, TrackerVJP, ZygoteAdjoint,
-                        ZygoteVJP
+    ForwardDiffOverAdjoint, ForwardDiffSensitivity, ForwardLSS,
+    ForwardSensitivity, GaussAdjoint, InterpolatingAdjoint, NILSAS,
+    NILSS, QuadratureAdjoint, ReverseDiffAdjoint, ReverseDiffVJP,
+    SteadyStateAdjoint, TrackerAdjoint, TrackerVJP, ZygoteAdjoint,
+    ZygoteVJP
 using Setfield: @set!
 using Static: True, False
 
@@ -37,22 +37,22 @@ include("collocation.jl")
 include("multiple_shooting.jl")
 
 export NeuralODE, NeuralDSDE, NeuralSDE, NeuralCDDE, NeuralDAE, AugmentedNDELayer,
-       NeuralODEMM
+    NeuralODEMM
 export FFJORD, FFJORDDistribution
 export DimMover
 
 export EpanechnikovKernel, UniformKernel, TriangularKernel, QuarticKernel, TriweightKernel,
-       TricubeKernel, GaussianKernel, CosineKernel, LogisticKernel, SigmoidKernel,
-       SilvermanKernel
+    TricubeKernel, GaussianKernel, CosineKernel, LogisticKernel, SigmoidKernel,
+    SilvermanKernel
 export collocate_data
 
 export multiple_shoot
 
 # Reexporting only certain functions from SciMLSensitivity
 export BacksolveAdjoint, QuadratureAdjoint, GaussAdjoint, InterpolatingAdjoint,
-       TrackerAdjoint, ZygoteAdjoint, ReverseDiffAdjoint, ForwardSensitivity,
-       ForwardDiffSensitivity, ForwardDiffOverAdjoint, SteadyStateAdjoint, ForwardLSS,
-       AdjointLSS, NILSS, NILSAS
+    TrackerAdjoint, ZygoteAdjoint, ReverseDiffAdjoint, ForwardSensitivity,
+    ForwardDiffSensitivity, ForwardDiffOverAdjoint, SteadyStateAdjoint, ForwardLSS,
+    AdjointLSS, NILSS, NILSAS
 export TrackerVJP, ZygoteVJP, EnzymeVJP, ReverseDiffVJP
 
 # Precompilation workload - must be at the end

src/DiffEqFlux.jl

@@ -76,7 +76,7 @@ function collocate_data(data, tpoints, kernel = TriangularKernel(), bandwidth =
     e2 = [_zero; _one; _zero]
     n = length(tpoints)
     bandwidth = bandwidth === nothing ?
-                (n^(-1 / 5)) * (n^(-3 / 35)) * ((log(n))^(-1 / 16)) : bandwidth
+        (n^(-1 / 5)) * (n^(-3 / 35)) * ((log(n))^(-1 / 16)) : bandwidth
 
     Wd = similar(data, n, size(data, 1))
     WT1 = similar(data, n, 2)
@@ -101,8 +101,10 @@ function collocate_data(data, tpoints, kernel = TriangularKernel(), bandwidth =
     return estimated_derivative, estimated_solution
 end
 
-@views function collocate_data(data::AbstractVector, tpoints::AbstractVector,
-        tpoints_sample::AbstractVector, interp, args...)
+@views function collocate_data(
+        data::AbstractVector, tpoints::AbstractVector,
+        tpoints_sample::AbstractVector, interp, args...
+    )
     du, u = collocate_data(reshape(data, 1, :), tpoints, tpoints_sample, interp, args...)
     return du[1, :], u[1, :]
 end

src/collocation.jl

@@ -59,12 +59,15 @@ preprint arXiv:1810.01367 (2018).
 end
 
 function LuxCore.initialstates(rng::AbstractRNG, n::FFJORD)
-    return (; model = LuxCore.initialstates(rng, n.model),
-        regularize = false, monte_carlo = true)
+    return (;
+        model = LuxCore.initialstates(rng, n.model),
+        regularize = false, monte_carlo = true,
+    )
 end
 
 function FFJORD(
-        model, tspan, input_dims, args...; ad = nothing, basedist = nothing, kwargs...)
+        model, tspan, input_dims, args...; ad = nothing, basedist = nothing, kwargs...
+    )
     !(model isa AbstractLuxLayer) && (model = FromFluxAdaptor()(model))
     return FFJORD(model, basedist, ad, input_dims, tspan, args, kwargs)
 end
@@ -75,8 +78,10 @@ end
 
 @inline __norm_batched(x) = sqrt.(sum(abs2, x; dims = 1:(ndims(x) - 1)))
 
-function __ffjord(model::StatefulLuxLayer, u::AbstractArray{T, N}, p, ad = nothing,
-        regularize::Bool = false, monte_carlo::Bool = true) where {T, N}
+function __ffjord(
+        model::StatefulLuxLayer, u::AbstractArray{T, N}, p, ad = nothing,
+        regularize::Bool = false, monte_carlo::Bool = true
+    ) where {T, N}
     L = size(u, N - 1)
     z = selectdim(u, N - 1, 1:(L - ifelse(regularize, 3, 1)))
     @set! model.ps = p
@@ -91,17 +96,23 @@ function __ffjord(model::StatefulLuxLayer, u::AbstractArray{T, N}, p, ad = nothi
             trace_jac = dropdims(
                 sum(
                     batched_matmul(
-                        reshape(e, 1, :, size(e, N)), reshape(Je, :, 1, size(Je, N)));
-                    dims = (1, 2));
-                dims = (1, 2))
+                        reshape(e, 1, :, size(e, N)), reshape(Je, :, 1, size(Je, N))
+                    );
+                    dims = (1, 2)
+                );
+                dims = (1, 2)
+            )
         elseif ad isa AutoZygote
             eJ = Lux.vector_jacobian_product(model, AutoZygote(), z, e)
             trace_jac = dropdims(
                 sum(
                     batched_matmul(
-                        reshape(eJ, 1, :, size(eJ, N)), reshape(e, :, 1, size(e, N)));
-                    dims = (1, 2));
-                dims = (1, 2))
+                        reshape(eJ, 1, :, size(eJ, N)), reshape(e, :, 1, size(e, N))
+                    );
+                    dims = (1, 2)
+                );
+                dims = (1, 2)
+            )
         else
             error("`ad` must be `nothing` or `AutoForwardDiff` or `AutoZygote`.")
         end
@@ -136,11 +147,14 @@ function __forward_ffjord(n::FFJORD, x::AbstractArray{T, N}, ps, st) where {T, N
     ffjord(u, p, t) = __ffjord(model, u, p, n.ad, regularize, monte_carlo)
 
     _z = ChainRulesCore.@ignore_derivatives fill!(
-        similar(x, S[1:(N - 2)]..., ifelse(regularize, 3, 1), S[N]), zero(T))
+        similar(x, S[1:(N - 2)]..., ifelse(regularize, 3, 1), S[N]), zero(T)
+    )
 
     prob = ODEProblem{false}(ffjord, cat(x, _z; dims = Val(N - 1)), n.tspan, ps)
-    sol = solve(prob, n.args...; sensealg, n.kwargs...,
-        save_everystep = false, save_start = false, save_end = true)
+    sol = solve(
+        prob, n.args...; sensealg, n.kwargs...,
+        save_everystep = false, save_start = false, save_end = true
+    )
     pred = __get_pred(sol)
     L = size(pred, N - 1)
 
@@ -185,11 +199,14 @@ function __backward_ffjord(::Type{T1}, n::FFJORD, n_samples::Int, ps, st, rng) w
     ffjord(u, p, t) = __ffjord(model, u, p, n.ad, regularize, monte_carlo)
 
     _z = ChainRulesCore.@ignore_derivatives fill!(
-        similar(x, S[1:(N - 2)]..., ifelse(regularize, 3, 1), S[N]), zero(T))
+        similar(x, S[1:(N - 2)]..., ifelse(regularize, 3, 1), S[N]), zero(T)
+    )
 
     prob = ODEProblem{false}(ffjord, cat(x, _z; dims = Val(N - 1)), reverse(n.tspan), ps)
-    sol = solve(prob, n.args...; sensealg, n.kwargs...,
-        save_everystep = false, save_start = false, save_end = true)
+    sol = solve(
+        prob, n.args...; sensealg, n.kwargs...,
+        save_everystep = false, save_start = false, save_end = true
+    )
     pred = __get_pred(sol)
     L = size(pred, N - 1)
 
@@ -222,7 +239,8 @@ function Distributions._logpdf(d::FFJORDDistribution, x::AbstractArray)
     return first(first(__forward_ffjord(d.model, x, d.ps, d.st)))
 end
 function Distributions._rand!(
-        rng::AbstractRNG, d::FFJORDDistribution, x::AbstractArray{<:Real})
+        rng::AbstractRNG, d::FFJORDDistribution, x::AbstractArray{<:Real}
+    )
     copyto!(x, __backward_ffjord(eltype(d), d.model, size(x, ndims(x)), d.ps, d.st, rng))
     return x
 end

src/ffjord.jl

@@ -33,9 +33,11 @@ Arguments:
     The parameter 'continuity_term' should be a relatively big number to enforce a large penalty
     whenever the last point of any group doesn't coincide with the first point of next group.
 """
-function multiple_shoot(p, ode_data, tsteps, prob::ODEProblem, loss_function::F,
+function multiple_shoot(
+        p, ode_data, tsteps, prob::ODEProblem, loss_function::F,
         continuity_loss::C, solver::SciMLBase.AbstractODEAlgorithm,
-        group_size::Integer; continuity_term::Real = 100, kwargs...) where {F, C}
+        group_size::Integer; continuity_term::Real = 100, kwargs...
+    ) where {F, C}
     datasize = size(ode_data, ndims(ode_data))
     griddims = ntuple(_ -> Colon(), ndims(ode_data) - 1)
 
@@ -47,12 +49,17 @@ function multiple_shoot(p, ode_data, tsteps, prob::ODEProblem, loss_function::F,
     ranges = group_ranges(datasize, group_size)
 
     # Multiple shooting predictions
-    sols = [solve(
-                remake(prob; p, tspan = (tsteps[first(rg)], tsteps[last(rg)]),
-                    u0 = ode_data[griddims..., first(rg)]),
+    sols = [
+        solve(
+                remake(
+                    prob; p, tspan = (tsteps[first(rg)], tsteps[last(rg)]),
+                    u0 = ode_data[griddims..., first(rg)]
+                ),
                 solver;
                 saveat = tsteps[rg],
-                kwargs...) for rg in ranges]
+                kwargs...
+            ) for rg in ranges
+    ]
     group_predictions = Array.(sols)
 
     # Abort and return infinite loss if one of the integrations failed
@@ -76,10 +83,14 @@ function multiple_shoot(p, ode_data, tsteps, prob::ODEProblem, loss_function::F,
     return loss, group_predictions
 end
 
-function multiple_shoot(p, ode_data, tsteps, prob::ODEProblem, loss_function::F,
-        solver::SciMLBase.AbstractODEAlgorithm, group_size::Integer; kwargs...) where {F}
-    return multiple_shoot(p, ode_data, tsteps, prob, loss_function,
-        _default_continuity_loss, solver, group_size; kwargs...)
+function multiple_shoot(
+        p, ode_data, tsteps, prob::ODEProblem, loss_function::F,
+        solver::SciMLBase.AbstractODEAlgorithm, group_size::Integer; kwargs...
+    ) where {F}
+    return multiple_shoot(
+        p, ode_data, tsteps, prob, loss_function,
+        _default_continuity_loss, solver, group_size; kwargs...
+    )
 end
 
 """
@@ -117,20 +128,22 @@ Arguments:
     The parameter 'continuity_term' should be a relatively big number to enforce a large penalty
     whenever the last point of any group doesn't coincide with the first point of next group.
 """
-function multiple_shoot(p, ode_data, tsteps, ensembleprob::EnsembleProblem,
+function multiple_shoot(
+        p, ode_data, tsteps, ensembleprob::EnsembleProblem,
         ensemblealg::SciMLBase.BasicEnsembleAlgorithm, loss_function::F,
         continuity_loss::C, solver::SciMLBase.AbstractODEAlgorithm,
-        group_size::Integer; continuity_term::Real = 100, kwargs...) where {F, C}
+        group_size::Integer; continuity_term::Real = 100, kwargs...
+    ) where {F, C}
     ntraj = size(ode_data, ndims(ode_data))
-    datasize = size(ode_data, ndims(ode_data)-1)
+    datasize = size(ode_data, ndims(ode_data) - 1)
     griddims = ntuple(_ -> Colon(), ndims(ode_data) - 2)
     prob = ensembleprob.prob
 
     if group_size < 2 || group_size > datasize
         throw(DomainError(group_size, "group_size can't be < 2 or > number of data points"))
     end
 
-    @assert ndims(ode_data)>=3 "ode_data must have at least three dimension: `size(ode_data) = (problem_dimension,length(tsteps),trajectories)"
+    @assert ndims(ode_data) >= 3 "ode_data must have at least three dimension: `size(ode_data) = (problem_dimension,length(tsteps),trajectories)"
     @assert datasize == length(tsteps)
     @assert ntraj == kwargs[:trajectories]
 
@@ -142,14 +155,16 @@ function multiple_shoot(p, ode_data, tsteps, ensembleprob::EnsembleProblem,
         rg -> begin
             newprob = remake(prob; p = p, tspan = (tsteps[first(rg)], tsteps[last(rg)]))
             function prob_func(prob, i, repeat)
-                remake(prob; u0 = ode_data[griddims..., first(rg), i])
+                return remake(prob; u0 = ode_data[griddims..., first(rg), i])
             end
             newensembleprob = EnsembleProblem(
                 newprob, prob_func, ensembleprob.output_func, ensembleprob.reduction,
-                ensembleprob.u_init, ensembleprob.safetycopy)
+                ensembleprob.u_init, ensembleprob.safetycopy
+            )
             solve(newensembleprob, solver, ensemblealg; saveat = tsteps[rg], kwargs...)
         end,
-        ranges)
+        ranges
+    )
     group_predictions = Array.(sols)
 
     # Abort and return infinite loss if one of the integrations did not converge?
@@ -176,12 +191,16 @@ function multiple_shoot(p, ode_data, tsteps, ensembleprob::EnsembleProblem,
     return loss, group_predictions
 end
 
-function multiple_shoot(p, ode_data, tsteps, ensembleprob::EnsembleProblem,
+function multiple_shoot(
+        p, ode_data, tsteps, ensembleprob::EnsembleProblem,
         ensemblealg::SciMLBase.BasicEnsembleAlgorithm, loss_function::F,
         solver::SciMLBase.AbstractODEAlgorithm, group_size::Integer;
-        continuity_term::Real = 100, kwargs...) where {F}
-    return multiple_shoot(p, ode_data, tsteps, ensembleprob, ensemblealg, loss_function,
-        _default_continuity_loss, solver, group_size; continuity_term, kwargs...)
+        continuity_term::Real = 100, kwargs...
+    ) where {F}
+    return multiple_shoot(
+        p, ode_data, tsteps, ensembleprob, ensemblealg, loss_function,
+        _default_continuity_loss, solver, group_size; continuity_term, kwargs...
+    )
 end
 
 """
@@ -207,8 +226,12 @@ julia> group_ranges(10, 5)
 ```
 """
 function group_ranges(datasize::Integer, groupsize::Integer)
-    2 ≤ groupsize ≤ datasize || throw(DomainError(groupsize,
-        "datasize must be positive and groupsize must to be within [2, datasize]"))
+    2 ≤ groupsize ≤ datasize || throw(
+        DomainError(
+            groupsize,
+            "datasize must be positive and groupsize must to be within [2, datasize]"
+        )
+    )
     return [i:min(datasize, i + groupsize - 1) for i in 1:(groupsize - 1):(datasize - 1)]
 end