diff --git a/.github/workflows/Runic.yml b/.github/workflows/Runic.yml
index 2c8dd85..806ee67 100644
--- a/.github/workflows/Runic.yml
+++ b/.github/workflows/Runic.yml
@@ -12,7 +12,7 @@ jobs:
     name: Runic
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: fredrikekre/runic-action@v1
         with:
           version: '1'
\ No newline at end of file
diff --git a/README.md b/README.md
index 87d946b..253a95b 100644
--- a/README.md
+++ b/README.md
@@ -99,4 +99,4 @@ end
 
 Which produces the following result:
 
-![](/docs/assets/1D_linear_advection.png)
\ No newline at end of file
+![](/docs/src/assets/1D_linear_advection.png)
\ No newline at end of file
diff --git a/docs/make.jl b/docs/make.jl
index ada8552..a8cd5fa 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -1,5 +1,15 @@
-push!(LOAD_PATH,"../src/")
+push!(LOAD_PATH, "../src/")
 
 using Documenter, FiniteDiffWENO5
 
-makedocs(sitename="My Documentation", remotes = nothing)
\ No newline at end of file
+makedocs(
+    sitename = "FiniteDiffWENO5",
+    remotes = nothing,
+    authors = "Hugo Dominguez",
+    pages = [
+        "Home" => "index.md",
+        "Getting Started" => "GettingStarted.md",
+        "Background" => "background.md",
+        "API" => "API.md",
+    ],
+)
diff --git a/docs/src/API.md b/docs/src/API.md
new file mode 100644
index 0000000..eb7819d
--- /dev/null
+++ b/docs/src/API.md
@@ -0,0 +1,6 @@
+# [API](@id API)
+
+```@autodocs
+Modules  = [FiniteDiffWENO5]
+
+```
\ No newline at end of file
diff --git a/docs/src/GettingStarted.md b/docs/src/GettingStarted.md
new file mode 100644
index 0000000..03ca2fe
--- /dev/null
+++ b/docs/src/GettingStarted.md
@@ -0,0 +1,99 @@
+
+## Getting Started
+
+There are currently two exported functions in FiniteDiffWENO5.jl: `WENOScheme()` and `WENO_step!()`. The user must define the grid and the initial conditions themselves.
+
+`WENOScheme()` is used to create a WENO scheme structure containing all the necessary information for the WENO method, while `WENO_step!()` performs one step of the time integration using the WENO-Z method and a 3rd-order Runge-Kutta method. Refer to the docstrings to see the available options for each function.
+
+Here is a simple example of how to use FiniteDiffWENO5.jl to solve the 1D advection equation using the conservative form on a staggered grid. For more examples, please refer to the folder examples in the repository or the tests in the test folder.
+
+```julia
+using FiniteDiffWENO5
+
+nx = 400
+
+x_min = -1.0
+x_max = 1.0
+Lx = x_max - x_min
+
+x = range(x_min, stop = x_max, length = nx)
+
+# Courant number
+CFL = 0.4
+period = 4
+
+# Parameters for Shu test
+z = -0.7
+δ = 0.005
+β = log(2) / (36 * δ^2)
+a = 0.5
+α = 10
+
+# Functions
+G(x, β, z) = exp.(-β .* (x .- z) .^ 2)
+F(x, α, a) = sqrt.(max.(1 .- α^2 .* (x .- a) .^ 2, 0.0))
+
+# Grid x assumed defined
+u0_vec = zeros(length(x))
+
+# Gaussian-like smooth bump at x in [-0.8, -0.6]
+idx = (x .>= -0.8) .& (x .<= -0.6)
+u0_vec[idx] .= (1 / 6) .* (G(x[idx], β, z - δ) .+ 4 .* G(x[idx], β, z) .+ G(x[idx], β, z + δ))
+
+# Heaviside step at x in [-0.4, -0.2]
+idx = (x .>= -0.4) .& (x .<= -0.2)
+u0_vec[idx] .= 1.0
+
+# Piecewise linear ramp at x in [0, 0.2]
+# Triangular spike at x=0.1, base width 0.2
+idx = abs.(x .- 0.1) .<= 0.1
+u0_vec[idx] .= 1 .- 10 .* abs.(x[idx] .- 0.1)
+
+# Elliptic/smooth bell at x in [0.4, 0.6]
+idx = (x .>= 0.4) .& (x .<= 0.6)
+u0_vec[idx] .= (1 / 6) .* (F(x[idx], α, a - δ) .+ 4 .* F(x[idx], α, a) .+ F(x[idx], α, a + δ))
+
+
+u = copy(u0_vec)
+
+# fix here boundary to periodic when equal to 2
+# staggering = true means that the advection velocity is defined on the sides of the cells and should be of size nx+1 compared to the scalar field u.
+# Multithreading to false means that the computations will be done on a single thread.
+# Set to true to enable multithreading if the Julia session was started with multiple threads (e.g. `julia -t 4`).
+weno = WENOScheme(u; boundary = (2, 2), stag = true, multithreading = false)
+
+# advection velocity with size nx+1 for staggered grid (here constant)
+a = (; x = ones(nx + 1))
+
+# grid size
+Δx = x[2] - x[1]
+Δt = CFL * Δx^(5 / 3)
+
+tmax = period * (Lx + Δx) / maximum(abs.(a.x))
+
+t = 0
+
+while t < tmax
+    # take as input the scalar field u, the advection velocity a as a NamedTuple,
+    # the WENO scheme struct weno, the time step Δt and the grid size Δx
+    WENO_step!(u, a, weno, Δt, Δx)
+
+    t += Δt
+
+    if t + Δt > tmax
+        Δt = tmax - t
+    end
+end
+
+f = Figure(size = (800, 600), dpi = 400)
+ax = Axis(f[1, 1], title = "1D linear advection after $period periods", xlabel = "x", ylabel = "u")
+lines!(ax, x, u0_vec, label = "Exact", linestyle = :dash, color = :red)
+scatter!(ax, x, u, label = "WENO5")
+xlims!(ax, x_min, x_max)
+axislegend(ax)
+display(f)
+```
+
+which outputs:
+
+![1D advection](./assets/1D_linear_advection.png)
\ No newline at end of file
diff --git a/docs/assets/1D_linear_advection.png b/docs/src/assets/1D_linear_advection.png
similarity index 100%
rename from docs/assets/1D_linear_advection.png
rename to docs/src/assets/1D_linear_advection.png
diff --git a/docs/src/background.md b/docs/src/background.md
new file mode 100644
index 0000000..c40d0dd
--- /dev/null
+++ b/docs/src/background.md
@@ -0,0 +1,20 @@
+# [Background](@id background)
+
+The weighted essentially non-oscillatory (WENO) schemes form a class of high-order accurate numerical methods for solving hyperbolic partial differential equations (PDEs). They are particularly effective at resolving sharp gradients and discontinuities without introducing spurious oscillations. The WENO schemes achieve this by constructing nonlinear convex combinations of lower-order candidate polynomials, where the weights are determined by the local smoothness of the solution. This approach allows the method to retain high-order accuracy in smooth regions while automatically reducing to lower-order, more stable behavior near discontinuities. The WENO framework was first introduced by [Liu et al. 1994](https://doi.org/10.1006/jcph.1994.1187]), building upon the earlier essentially non-oscillatory (ENO) schemes developed by [Harten et al. 1987](https://doi.org/10.1016/0021-9991(87)90031-3).
+
+In practice, WENO schemes can be formulated either in a finite-volume or finite-difference framework. In this package, we adopt the finite-difference formulation, which is particularly well-suited for problems defined on structured grids due to its simplicity and computational efficiency. The time integration is based on a third order strong stability preserving Runge-Kutta (SSP-RK3) method.
+
+The implementation of a finite-difference WENO scheme involves the following main steps:
+
+2. **Smoothness Indicators**: For each candidate stencil, a smoothness indicator is computed. This indicator quantifies how smooth the approximation is within that stencil, with lower values indicating smoother regions.
+3. **Weight Calculation**: Nonlinear weights are computed based on the smoothness indicators. Stencils with lower smoothness indicators receive higher weights, allowing the scheme to adaptively favor smoother regions.
+4. **Reconstruction**: The final high-order approximation is obtained by combining the candidate polynomials using the computed weights.
+5. **Flux Evaluation**: The reconstructed values are used to approximate the variable of interest at the cell interfaces, which are then used in the numerical flux calculations.
+
+In FiniteDiffWENO5.jl, two forms of advection operators are currently supported:
+1. **Non-conservative Form** $\mathbf{v} \cdot \nabla u$, where the velocity field $\mathbf{v}$ and scalar field $u$ are both defined at the same grid locations (collocated grid), and
+2. **Conservative Form** $\nabla \cdot (\mathbf{v}u)$, where $\mathbf{v}$ is defined on cell faces and $u$ at cell centers (staggered grid).
+
+In both formulations, only the scalar field $u$ is reconstructed at the cell interfaces using the WENO scheme.
+
+The package currently implements the WENO-Z reconstruction developed by [Borges et al. (2008)](https://doi.org/10.1016/j.jcp.2007.11.038). This variant introduces a modified computation of the nonlinear weights that improves accuracy near critical points—where the first derivative of the solution vanishes—while preserving the robust, non-oscillatory behavior of the classical WENO methods. Additional reconstruction variants may be included in future versions.
diff --git a/docs/src/index.md b/docs/src/index.md
index 89bf36e..4a6b292 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -1,8 +1,17 @@
+# [FiniteDiffWENO5.jl](@id home)
 
-```@index
-```
+FiniteDiffWENO5.jl is a Julia package that implements fifth-order finite-difference weighted essentially non-oscillatory (WENO) schemes for solving hyperbolic partial differential equations (PDEs) in 1D, 2D and 3D on regular grids.
+
+The package currently focuses on the non-conservative form of the advection terms ($\mathbf{v} \cdot \nabla u$) on a collocated grid, and the conservative form ($\nabla \cdot$ ($\mathbf{v} u$)) where the velocity field $\mathbf{v}$ and scalar field $u$ are on a staggered grid with the advection velocity located on the sides of the cells.
+
+The core of the package is written in pure Julia, focusing on performance using CPUs, but GPU support is available using KernelAbstractions.jl and Chmy.jl via an extension.
 
-```@autodocs
-Modules  = [FiniteDiffWENO5]
+## Installation
 
-```
\ No newline at end of file
+FiniteDiffWENO5.jl is a registered package and may be installed directly with the following command in the Julia REPL
+
+```julia-repl
+julia>]
+  pkg> add FiniteDiffWENO5
+  pkg> test FiniteDiffWENO5
+```
diff --git a/examples/1D/1D_linear.jl b/examples/1D/1D_linear.jl
index 228979f..541ab10 100644
--- a/examples/1D/1D_linear.jl
+++ b/examples/1D/1D_linear.jl
@@ -1,13 +1,13 @@
 using FiniteDiffWENO5
 using GLMakie
 
-function main(nx=400)
+function main(nx = 400)
 
     x_min = -1.0
     x_max = 1.0
     Lx = x_max - x_min
 
-    x = range(x_min, stop=x_max, length=nx)
+    x = range(x_min, stop = x_max, length = nx)
 
     # Courant number
     CFL = 0.4
@@ -16,20 +16,20 @@ function main(nx=400)
     # Parameters for Shu test
     z = -0.7
     δ = 0.005
-    β = log(2)/(36*δ^2)
+    β = log(2) / (36 * δ^2)
     a = 0.5
     α = 10
 
     # Functions
-    G(x, β, z) = exp.(-β .* (x .- z).^2)
-    F(x, α, a) = sqrt.(max.(1 .- α^2 .* (x .- a).^2, 0.0))
+    G(x, β, z) = exp.(-β .* (x .- z) .^ 2)
+    F(x, α, a) = sqrt.(max.(1 .- α^2 .* (x .- a) .^ 2, 0.0))
 
     # Grid x assumed defined
     u0_vec = zeros(length(x))
 
     # Gaussian-like smooth bump at x in [-0.8, -0.6]
     idx = (x .>= -0.8) .& (x .<= -0.6)
-    u0_vec[idx] .= (1/6) .* (G(x[idx], β, z - δ) .+ 4 .* G(x[idx], β, z) .+ G(x[idx], β, z + δ))
+    u0_vec[idx] .= (1 / 6) .* (G(x[idx], β, z - δ) .+ 4 .* G(x[idx], β, z) .+ G(x[idx], β, z + δ))
 
     # Heaviside step at x in [-0.4, -0.2]
     idx = (x .>= -0.4) .& (x .<= -0.2)
@@ -42,18 +42,18 @@ function main(nx=400)
 
     # Elliptic/smooth bell at x in [0.4, 0.6]
     idx = (x .>= 0.4) .& (x .<= 0.6)
-    u0_vec[idx] .= (1/6) .* (F(x[idx], α, a - δ) .+ 4 .* F(x[idx], α, a) .+ F(x[idx], α, a + δ))
+    u0_vec[idx] .= (1 / 6) .* (F(x[idx], α, a - δ) .+ 4 .* F(x[idx], α, a) .+ F(x[idx], α, a + δ))
 
 
     u = copy(u0_vec)
-    weno = WENOScheme(u; boundary=(2, 2), stag=true, multithreading=false)
+    weno = WENOScheme(u; boundary = (2, 2), stag = true, multithreading = false)
 
     # advection velocity
-    a = (;x=ones(nx+1))
+    a = (; x = ones(nx + 1))
 
     # grid size
     Δx = x[2] - x[1]
-    Δt = CFL*Δx^(5/3)
+    Δt = CFL * Δx^(5 / 3)
 
     tmax = period * (Lx + Δx) / maximum(abs.(a.x))
 
@@ -69,14 +69,14 @@ function main(nx=400)
         end
     end
 
-    f = Figure(size = (800, 600), dpi=400)
+    f = Figure(size = (800, 600), dpi = 400)
     ax = Axis(f[1, 1], title = "1D linear advection after $period periods", xlabel = "x", ylabel = "u")
     lines!(ax, x, u0_vec, label = "Exact", linestyle = :dash, color = :red)
     scatter!(ax, x, u, label = "WENO5")
     xlims!(ax, x_min, x_max)
     axislegend(ax)
     save("1D_linear_advection.png", f)
-    display(f)
+    return display(f)
 end
 
-main()
\ No newline at end of file
+main()
diff --git a/examples/1D/1D_linear_chmy.jl b/examples/1D/1D_linear_chmy.jl
index 0bca2e9..bd175c0 100644
--- a/examples/1D/1D_linear_chmy.jl
+++ b/examples/1D/1D_linear_chmy.jl
@@ -3,7 +3,7 @@ using CairoMakie
 using Chmy
 using KernelAbstractions
 
-function main(backend=CPU(), nx=400)
+function main(backend = CPU(), nx = 400)
 
     arch = Arch(backend)
 
@@ -11,9 +11,9 @@ function main(backend=CPU(), nx=400)
     x_max = 1.0
     Lx = x_max - x_min
 
-    x = range(x_min, stop=x_max, length=nx)
+    x = range(x_min, stop = x_max, length = nx)
 
-    grid   = UniformGrid(arch; origin=(x_min,), extent=(Lx,), dims=(nx,))
+    grid = UniformGrid(arch; origin = (x_min,), extent = (Lx,), dims = (nx,))
 
     # Courant number
     CFL = 0.7
@@ -22,20 +22,20 @@ function main(backend=CPU(), nx=400)
     # Parameters
     z = -0.7
     δ = 0.005
-    β = log(2)/(36*δ^2)
+    β = log(2) / (36 * δ^2)
     a = 0.5
     α = 10
 
     # Functions
-    G(x, β, z) = exp.(-β .* (x .- z).^2)
-    F(x, α, a) = sqrt.(max.(1 .- α^2 .* (x .- a).^2, 0.0))
+    G(x, β, z) = exp.(-β .* (x .- z) .^ 2)
+    F(x, α, a) = sqrt.(max.(1 .- α^2 .* (x .- a) .^ 2, 0.0))
 
     # Grid x assumed defined
     u0_vec = zeros(length(x))
 
     # Gaussian-like smooth bump at x in [-0.8, -0.6]
     idx = (x .>= -0.8) .& (x .<= -0.6)
-    u0_vec[idx] .= (1/6) .* (G(x[idx], β, z - δ) .+ 4 .* G(x[idx], β, z) .+ G(x[idx], β, z + δ))
+    u0_vec[idx] .= (1 / 6) .* (G(x[idx], β, z - δ) .+ 4 .* G(x[idx], β, z) .+ G(x[idx], β, z + δ))
 
     # Heaviside step at x in [-0.4, -0.2]
     idx = (x .>= -0.4) .& (x .<= -0.2)
@@ -48,23 +48,23 @@ function main(backend=CPU(), nx=400)
 
     # Elliptic/smooth bell at x in [0.4, 0.6]
     idx = (x .>= 0.4) .& (x .<= 0.6)
-    u0_vec[idx] .= (1/6) .* (F(x[idx], α, a - δ) .+ 4 .* F(x[idx], α, a) .+ F(x[idx], α, a + δ))
+    u0_vec[idx] .= (1 / 6) .* (F(x[idx], α, a - δ) .+ 4 .* F(x[idx], α, a) .+ F(x[idx], α, a + δ))
 
 
     u = Field(backend, grid, Center())
     set!(u, u0_vec)
-    weno = WENOScheme(u, grid; boundary=(2, 2), stag=true)
+    weno = WENOScheme(u, grid; boundary = (2, 2), stag = true)
 
     # advection velocity
-    a_vec = ones(nx+1) .* -1
+    a_vec = ones(nx + 1) .* -1
     a = VectorField(backend, grid)
     set!(a, a_vec)
 
     # grid size
     Δx = x[2] - x[1]
-    Δt = CFL * Δx^(5/3)
+    Δt = CFL * Δx^(5 / 3)
 
-    tmax = period * (Lx+Δx) / maximum(abs.(a.x))
+    tmax = period * (Lx + Δx) / maximum(abs.(a.x))
 
     t = 0
 
@@ -86,7 +86,7 @@ function main(backend=CPU(), nx=400)
     scatter!(ax, x, interior(u) |> Array, label = "WENO5")
     xlims!(ax, x_min, x_max)
     axislegend(ax)
-    display(f)
+    return display(f)
 end
 
-main()
\ No newline at end of file
+main()
diff --git a/examples/2D/2D_linear.jl b/examples/2D/2D_linear.jl
index 6802abc..6192d8b 100644
--- a/examples/2D/2D_linear.jl
+++ b/examples/2D/2D_linear.jl
@@ -1,29 +1,29 @@
 using FiniteDiffWENO5
 using GLMakie
 
-function main(;nx=400, ny=400)
+function main(; nx = 400, ny = 400)
 
     Lx = 1.0
     Δx = Lx / nx
     Δy = Lx / ny
 
-    x = range(0, stop=Lx, length=nx)
+    x = range(0, stop = Lx, length = nx)
 
     # Courant number
     CFL = 0.7
     period = 1
 
     # Grid x assumed defined
-    x = range(0, length=nx, stop= Lx)
-    y = range(0, length=ny, stop= Lx)
+    x = range(0, length = nx, stop = Lx)
+    y = range(0, length = ny, stop = Lx)
     grid = (x .* ones(ny)', ones(nx) .* y')
 
     vx0 = ones(nx, ny)
     vy0 = ones(nx, ny)
 
-    v = (;x=vy0, y=vx0)
+    v = (; x = vy0, y = vx0)
 
-    x0 = 1/4
+    x0 = 1 / 4
     c = 0.08
 
     u0 = zeros(ny, nx)
@@ -33,11 +33,11 @@ function main(;nx=400, ny=400)
     end
 
     u = copy(u0)
-    weno = WENOScheme(u; boundary=(2, 2, 2, 2), stag=false, multithreading=true)
+    weno = WENOScheme(u; boundary = (2, 2, 2, 2), stag = false, multithreading = true)
 
 
     # grid size
-    Δt = CFL*min(Δx, Δy)^(5/3)
+    Δt = CFL * min(Δx, Δy)^(5 / 3)
 
     tmax = period * Lx / max(maximum(abs.(vx0)), maximum(abs.(vy0)))
 
@@ -45,9 +45,9 @@ function main(;nx=400, ny=400)
     counter = 0
 
     f = Figure(size = (800, 600))
-    ax = Axis(f[1, 1], title = "t = $(round(t, digits=2))")
+    ax = Axis(f[1, 1], title = "t = $(round(t, digits = 2))")
     u_obser = Observable(u0)
-    hm = heatmap!(ax, x, y, u_obser; colormap = cgrad(:roma, rev = true), colorrange=(0, 1))
+    hm = heatmap!(ax, x, y, u_obser; colormap = cgrad(:roma, rev = true), colorrange = (0, 1))
     Colorbar(f[1, 2], label = "u", hm)
     display(f)
 
@@ -63,13 +63,14 @@ function main(;nx=400, ny=400)
 
         if counter % 100 == 0
             u_obser[] = u
-            ax.title = "t = $(round(t, digits=2))"
+            ax.title = "t = $(round(t, digits = 2))"
         end
 
         counter += 1
 
     end
 
+    return
 end
 
-main(nx=400, ny=400)
\ No newline at end of file
+main(nx = 400, ny = 400)
diff --git a/examples/2D/2D_linear_chmy.jl b/examples/2D/2D_linear_chmy.jl
index 4bd48ed..8972245 100644
--- a/examples/2D/2D_linear_chmy.jl
+++ b/examples/2D/2D_linear_chmy.jl
@@ -3,7 +3,7 @@ using Chmy
 using KernelAbstractions
 using GLMakie
 
-function main(;backend=CPU(), nx=400, ny=400, stag=true)
+function main(; backend = CPU(), nx = 400, ny = 400, stag = true)
 
     arch = Arch(backend)
 
@@ -11,28 +11,28 @@ function main(;backend=CPU(), nx=400, ny=400, stag=true)
     Δx = Lx / nx
     Δy = Lx / ny
 
-    grid = UniformGrid(arch; origin=(0.0, 0.0), extent=(Lx, Lx), dims=(nx, ny))
+    grid = UniformGrid(arch; origin = (0.0, 0.0), extent = (Lx, Lx), dims = (nx, ny))
 
     # Courant number
     CFL = 0.7
     period = 1
 
     # Grid x assumed defined
-    x = range(0, length=nx, stop= Lx)
-    y = range(0, length=ny, stop= Lx)
+    x = range(0, length = nx, stop = Lx)
+    y = range(0, length = ny, stop = Lx)
     grid_array = (x .* ones(ny)', ones(nx) .* y')
 
     if stag
-        vx0 = ones(nx+1, ny)
-        vy0 = ones(nx, ny+1)
+        vx0 = ones(nx + 1, ny)
+        vy0 = ones(nx, ny + 1)
     else
         vx0 = ones(nx, ny)
         vy0 = ones(nx, ny)
     end
 
-    v = (;x=vy0, y=vx0)
+    v = (; x = vy0, y = vx0)
 
-    x0 = 1/4
+    x0 = 1 / 4
     c = 0.08
 
     u0 = zeros(ny, nx)
@@ -44,20 +44,22 @@ function main(;backend=CPU(), nx=400, ny=400, stag=true)
     u = Field(backend, grid, Center())
     set!(u, u0)
 
-    weno = WENOScheme(u, grid; boundary=(1, 1, 1, 1), stag=stag, multithreading=true)
+    weno = WENOScheme(u, grid; boundary = (1, 1, 1, 1), stag = stag, multithreading = true)
 
     if stag
         v = VectorField(backend, grid)
     else
-        v = (x=Field(arch, grid, Center()),
-        y=Field(arch, grid, Center()))
+        v = (
+            x = Field(arch, grid, Center()),
+            y = Field(arch, grid, Center()),
+        )
     end
 
     set!(v.x, vy0)
     set!(v.y, vx0)
 
     # grid size
-    Δt = CFL*min(Δx, Δy)^(5/3)
+    Δt = CFL * min(Δx, Δy)^(5 / 3)
 
     tmax = period * Lx / max(maximum(abs.(vx0)), maximum(abs.(vy0)))
 
@@ -67,9 +69,9 @@ function main(;backend=CPU(), nx=400, ny=400, stag=true)
     mass_ini = sum(u0) * Δx * Δy
 
     f = Figure(size = (800, 600))
-    ax = Axis(f[1, 1], title = "t = $(round(t, digits=2))")
+    ax = Axis(f[1, 1], title = "t = $(round(t, digits = 2))")
     u_obser = Observable(u0)
-    hm = heatmap!(ax, x, y, u_obser; colormap = cgrad(:roma, rev = true), colorrange=(0, 1))
+    hm = heatmap!(ax, x, y, u_obser; colormap = cgrad(:roma, rev = true), colorrange = (0, 1))
     Colorbar(f[1, 2], label = "u", hm)
     display(f)
 
@@ -89,14 +91,15 @@ function main(;backend=CPU(), nx=400, ny=400, stag=true)
             mass_ratio = (sum(u) * Δx * Δy) / mass_ini
 
             u_obser[] = interior(u) |> Array
-            ax.title = "t = $(round(t, digits=2)), mass ratio = $(round(mass_ratio, digits=6))"
+            ax.title = "t = $(round(t, digits = 2)), mass ratio = $(round(mass_ratio, digits = 6))"
         end
 
         counter += 1
 
     end
 
+    return
 end
 
 
-main(backend=CPU(), nx=400, ny=400, stag=true)
\ No newline at end of file
+main(backend = CPU(), nx = 400, ny = 400, stag = true)
diff --git a/examples/2D/2D_linear_rotation.jl b/examples/2D/2D_linear_rotation.jl
index b3573c5..ab542f0 100644
--- a/examples/2D/2D_linear_rotation.jl
+++ b/examples/2D/2D_linear_rotation.jl
@@ -1,30 +1,30 @@
 using FiniteDiffWENO5
 using GLMakie
 
-function main(;nx=400, ny=400)
+function main(; nx = 400, ny = 400)
 
     Lx = 1.0
     Δx = Lx / nx
     Δy = Lx / ny
 
-    x = range(0, stop=Lx, length=nx)
+    x = range(0, stop = Lx, length = nx)
 
     # Courant number
     CFL = 0.7
     period = 1
 
     # Grid x assumed defined
-    x = range(0, length=nx, stop= Lx)
-    y = range(0, length=ny, stop= Lx)
+    x = range(0, length = nx, stop = Lx)
+    y = range(0, length = ny, stop = Lx)
     grid = (x .* ones(ny)', ones(nx) .* y')
 
     w = π
-    vx0 = w .* (grid[1] .- Lx/2)
-    vy0 = -w .* (grid[2] .- Lx/2)
+    vx0 = w .* (grid[1] .- Lx / 2)
+    vy0 = -w .* (grid[2] .- Lx / 2)
 
-    v = (;x=vy0, y=vx0)
+    v = (; x = vy0, y = vx0)
 
-    x0 = 1/4
+    x0 = 1 / 4
     c = 0.08
 
     u0 = zeros(ny, nx)
@@ -34,21 +34,21 @@ function main(;nx=400, ny=400)
     end
 
     u = copy(u0)
-    weno = WENOScheme(u; boundary=(2, 2, 2, 2), stag=false, multithreading=true)
+    weno = WENOScheme(u; boundary = (2, 2, 2, 2), stag = false, multithreading = true)
 
 
     # grid size
-    Δt = CFL*min(Δx, Δy)^(5/3)
+    Δt = CFL * min(Δx, Δy)^(5 / 3)
 
-    tmax = period / (w/(2*π))
+    tmax = period / (w / (2 * π))
 
     t = 0
     counter = 0
 
     f = Figure(size = (800, 600))
-    ax = Axis(f[1, 1], title = "t = $(round(t, digits=2))")
+    ax = Axis(f[1, 1], title = "t = $(round(t, digits = 2))")
     u_obser = Observable(u0)
-    hm = heatmap!(ax, x, y, u_obser; colormap = cgrad(:roma, rev = true), colorrange=(0, 1))
+    hm = heatmap!(ax, x, y, u_obser; colormap = cgrad(:roma, rev = true), colorrange = (0, 1))
     Colorbar(f[1, 2], label = "u", hm)
     display(f)
 
@@ -64,13 +64,14 @@ function main(;nx=400, ny=400)
 
         if counter % 100 == 0
             u_obser[] = u
-            ax.title = "t = $(round(t, digits=2))"
+            ax.title = "t = $(round(t, digits = 2))"
         end
 
         counter += 1
 
     end
 
+    return
 end
 
-main(nx=400, ny=400)
\ No newline at end of file
+main(nx = 400, ny = 400)
diff --git a/examples/2D/2D_linear_rotation_chmy.jl b/examples/2D/2D_linear_rotation_chmy.jl
index ffdb427..5a68271 100644
--- a/examples/2D/2D_linear_rotation_chmy.jl
+++ b/examples/2D/2D_linear_rotation_chmy.jl
@@ -3,7 +3,7 @@ using Chmy
 using KernelAbstractions
 using GLMakie
 
-function main(;backend=CPU(), nx=400, ny=400)
+function main(; backend = CPU(), nx = 400, ny = 400)
 
     arch = Arch(backend)
 
@@ -11,24 +11,24 @@ function main(;backend=CPU(), nx=400, ny=400)
     Δx = Lx / nx
     Δy = Lx / ny
 
-    grid = UniformGrid(arch; origin=(0.0, 0.0), extent=(Lx, Lx), dims=(nx, ny))
+    grid = UniformGrid(arch; origin = (0.0, 0.0), extent = (Lx, Lx), dims = (nx, ny))
 
     # Courant number
     CFL = 0.7
     period = 1
 
     # Grid x assumed defined
-    x = range(0, length=nx, stop= Lx)
-    y = range(0, length=ny, stop= Lx)
+    x = range(0, length = nx, stop = Lx)
+    y = range(0, length = ny, stop = Lx)
     grid_array = (x .* ones(ny)', ones(nx) .* y')
 
     w = π
-    vx0 = w .* (grid_array[1] .- Lx/2)
-    vy0 = -w .* (grid_array[2] .- Lx/2)
+    vx0 = w .* (grid_array[1] .- Lx / 2)
+    vy0 = -w .* (grid_array[2] .- Lx / 2)
 
-    v = (;x=vy0, y=vx0)
+    v = (; x = vy0, y = vx0)
 
-    x0 = 1/4
+    x0 = 1 / 4
     c = 0.08
 
     u0 = zeros(ny, nx)
@@ -40,26 +40,28 @@ function main(;backend=CPU(), nx=400, ny=400)
     u = Field(backend, grid, Center())
     set!(u, u0)
 
-    weno = WENOScheme(u, grid; boundary=(2, 2, 2, 2), stag=false, multithreading=true)
+    weno = WENOScheme(u, grid; boundary = (2, 2, 2, 2), stag = false, multithreading = true)
 
-    v = (x=Field(arch, grid, Center()),
-        y=Field(arch, grid, Center()))
+    v = (
+        x = Field(arch, grid, Center()),
+        y = Field(arch, grid, Center()),
+    )
 
     set!(v.x, vy0)
     set!(v.y, vx0)
 
     # grid size
-    Δt = CFL*min(Δx, Δy)^(5/3)
+    Δt = CFL * min(Δx, Δy)^(5 / 3)
 
-    tmax = period / (w/(2*π))
+    tmax = period / (w / (2 * π))
 
     t = 0
     counter = 0
 
     f = Figure(size = (800, 600))
-    ax = Axis(f[1, 1], title = "t = $(round(t, digits=2))")
+    ax = Axis(f[1, 1], title = "t = $(round(t, digits = 2))")
     u_obser = Observable(u0)
-    hm = heatmap!(ax, x, y, u_obser; colormap = cgrad(:roma, rev = true), colorrange=(0, 1))
+    hm = heatmap!(ax, x, y, u_obser; colormap = cgrad(:roma, rev = true), colorrange = (0, 1))
     Colorbar(f[1, 2], label = "u", hm)
     display(f)
 
@@ -76,13 +78,14 @@ function main(;backend=CPU(), nx=400, ny=400)
         if counter % 100 == 0
             KernelAbstractions.synchronize(backend)
             u_obser[] = interior(u) |> Array
-            ax.title = "t = $(round(t, digits=2))"
+            ax.title = "t = $(round(t, digits = 2))"
         end
 
         counter += 1
 
     end
 
+    return
 end
 
-main(backend=CPU(), nx=400, ny=400)
\ No newline at end of file
+main(backend = CPU(), nx = 400, ny = 400)
diff --git a/examples/3D/3D_linear.jl b/examples/3D/3D_linear.jl
index c8dcf61..acd2689 100644
--- a/examples/3D/3D_linear.jl
+++ b/examples/3D/3D_linear.jl
@@ -1,15 +1,15 @@
 using FiniteDiffWENO5
 using GLMakie
 
-function main(; nx=50, ny=50, nz=50)
+function main(; nx = 50, ny = 50, nz = 50)
     L = 1.0
     Δx = L / nx
     Δy = L / ny
     Δz = L / nz
 
-    x = range(0, stop=L, length=nx)
-    y = range(0, stop=L, length=ny)
-    z = range(0, stop=L, length=nz)
+    x = range(0, stop = L, length = nx)
+    y = range(0, stop = L, length = ny)
+    z = range(0, stop = L, length = nz)
 
     # Courant number
     CFL = 0.7
@@ -24,28 +24,28 @@ function main(; nx=50, ny=50, nz=50)
     vy0 = ones(size(Y))
     vz0 = zeros(size(X)) # Rotation in XY plane only
 
-    v = (; x=vx0, y=vy0, z=vz0)
+    v = (; x = vx0, y = vy0, z = vz0)
 
-    x0 = 1/4
+    x0 = 1 / 4
     c = 0.08
 
     u0 = zeros(ny, nx, nz)
     for I in CartesianIndices((ny, nx, nz))
-        u0[I] = exp(-((X[I]-x0)^2 + (Y[I]-x0)^2 + (Z[I] - 0.5)^2) / c^2)
+        u0[I] = exp(-((X[I] - x0)^2 + (Y[I] - x0)^2 + (Z[I] - 0.5)^2) / c^2)
     end
 
     u = copy(u0)
-    weno = WENOScheme(u; boundary=(2, 2, 2, 2, 2, 2), stag=false, multithreading=true)
+    weno = WENOScheme(u; boundary = (2, 2, 2, 2, 2, 2), stag = false, multithreading = true)
 
-    Δt = CFL * min(Δx, Δy, Δz)^(5/3)
+    Δt = CFL * min(Δx, Δy, Δz)^(5 / 3)
     tmax = period * L / max(maximum(abs.(vx0)), maximum(abs.(vy0)), maximum(abs.(vz0)))
     t = 0
     counter = 0
 
     f = Figure(size = (800, 600))
-    ax = Axis(f[1, 1], title = "t = $(round(t, digits=2))")
+    ax = Axis(f[1, 1], title = "t = $(round(t, digits = 2))")
     u_obser = Observable(u[:, :, div(nz, 2)])
-    heatmap!(ax, u_obser; colormap = cgrad(:roma, rev = true), colorrange=(0, 1.0))
+    heatmap!(ax, u_obser; colormap = cgrad(:roma, rev = true), colorrange = (0, 1.0))
     Colorbar(f[1, 2], label = "u")
     display(f)
 
@@ -59,11 +59,12 @@ function main(; nx=50, ny=50, nz=50)
 
         if counter % 50 == 0
             u_obser[] = u[:, :, div(nz, 2)]
-            ax.title = "t = $(round(t, digits=2))"
+            ax.title = "t = $(round(t, digits = 2))"
         end
 
         counter += 1
     end
+    return
 end
 
-main()
\ No newline at end of file
+main()
diff --git a/examples/3D/3D_linear_chmy.jl b/examples/3D/3D_linear_chmy.jl
index a112f0e..0acb1d2 100644
--- a/examples/3D/3D_linear_chmy.jl
+++ b/examples/3D/3D_linear_chmy.jl
@@ -3,7 +3,7 @@ using Chmy
 using KernelAbstractions
 using CairoMakie
 
-function main(;backend=CPU(), nx=50, ny=50, nz=50)
+function main(; backend = CPU(), nx = 50, ny = 50, nz = 50)
 
     arch = Arch(backend)
 
@@ -12,16 +12,16 @@ function main(;backend=CPU(), nx=50, ny=50, nz=50)
     Δy = Lx / ny
     Δz = Lx / nz
 
-    grid = UniformGrid(arch; origin=(0.0, 0.0, 0.0), extent=(Lx, Lx, Lx), dims=(nx, ny, nz))
+    grid = UniformGrid(arch; origin = (0.0, 0.0, 0.0), extent = (Lx, Lx, Lx), dims = (nx, ny, nz))
 
     # Courant number
     CFL = 0.7
     period = 1
 
     # 3D grid
-    x = range(0, length=nx, stop=Lx)
-    y = range(0, length=ny, stop=Lx)
-    z = range(0, length=nz, stop=Lx)
+    x = range(0, length = nx, stop = Lx)
+    y = range(0, length = ny, stop = Lx)
+    z = range(0, length = nz, stop = Lx)
 
     X = reshape(x, 1, nx, 1) .* ones(ny, 1, nz)
     Y = reshape(y, ny, 1, 1) .* ones(1, nx, nz)
@@ -35,34 +35,36 @@ function main(;backend=CPU(), nx=50, ny=50, nz=50)
     vy0 = ones(size(Y3D))
     vz0 = zeros(size(Z3D)) # Rotation in XY plane only
 
-    v = (; x=Field(arch, grid, Center()),
-        y=Field(arch, grid, Center()),
-        z=Field(arch, grid, Center()))
+    v = (;
+        x = Field(arch, grid, Center()),
+        y = Field(arch, grid, Center()),
+        z = Field(arch, grid, Center()),
+    )
     set!(v.x, vy0)
     set!(v.y, vx0)
     set!(v.z, vz0)
 
-    x0 = 1/4
+    x0 = 1 / 4
     c = 0.08
 
     u0 = zeros(ny, nx, nz)
     for I in CartesianIndices((ny, nx, nz))
-        u0[I] = exp(-((X3D[I]-x0)^2 + (Y3D[I]-x0)^2 + (Z3D[I]-0.5)^2) / c^2)
+        u0[I] = exp(-((X3D[I] - x0)^2 + (Y3D[I] - x0)^2 + (Z3D[I] - 0.5)^2) / c^2)
     end
 
     u = Field(backend, grid, Center())
     set!(u, u0)
-    weno = WENOScheme(u, grid; boundary=(2, 2, 2, 2, 2, 2), stag=false, multithreading=true)
+    weno = WENOScheme(u, grid; boundary = (2, 2, 2, 2, 2, 2), stag = false, multithreading = true)
 
-    Δt = CFL * min(Δx, Δy, Δz)^(5/3)
+    Δt = CFL * min(Δx, Δy, Δz)^(5 / 3)
     tmax = period * Lx / max(maximum(abs.(vx0)), maximum(abs.(vy0)), maximum(abs.(vz0)))
     t = 0
     counter = 0
 
     f = Figure(size = (800, 600))
-    ax = Axis(f[1, 1], title = "t = $(round(t, digits=2))")
+    ax = Axis(f[1, 1], title = "t = $(round(t, digits = 2))")
     u_obser = Observable(u0[:, :, div(nz, 2)])
-    hm = heatmap!(ax, u_obser; colormap = cgrad(:roma, rev = true), colorrange=(0, 1))
+    hm = heatmap!(ax, u_obser; colormap = cgrad(:roma, rev = true), colorrange = (0, 1))
     Colorbar(f[1, 2], label = "u", hm)
     display(f)
 
@@ -79,7 +81,7 @@ function main(;backend=CPU(), nx=50, ny=50, nz=50)
             if backend == CPU()
                 KernelAbstractions.synchronize(backend)
                 u_obser[] = (interior(u) |> Array)[:, :, div(nz, 2)]
-                ax.title = "t = $(round(t, digits=2))"
+                ax.title = "t = $(round(t, digits = 2))"
             end
         end
 
@@ -88,12 +90,11 @@ function main(;backend=CPU(), nx=50, ny=50, nz=50)
 
     KernelAbstractions.synchronize(backend)
     u_obser[] = (interior(u) |> Array)[:, :, div(nz, 2)]
-    ax.title = "t = $(round(t, digits=2))"
+    ax.title = "t = $(round(t, digits = 2))"
 
-    save("weno5_cuda.png", f);
+    return save("weno5_cuda.png", f)
 end
 
 # using CUDA
 # main(backend=CUDABackend())
 main()
-
diff --git a/ext/ChmyExt.jl b/ext/ChmyExt.jl
index a88595e..93038d4 100644
--- a/ext/ChmyExt.jl
+++ b/ext/ChmyExt.jl
@@ -20,7 +20,7 @@ Create a WENO scheme structure for the given field `u` on the specified `grid` u
 - `stag::Bool`: Whether the grid is staggered (velocities on cell faces) or not (velocities on cell centers).
 - `multithreading::Bool`: Whether to use multithreading (only for 2D and 3D). Default is false.
 """
-function WENOScheme(u::AbstractField{T, N}, grid; boundary = (2, 2), stag = true, multithreading=true) where {T, N}
+function WENOScheme(u::AbstractField{T, N}, grid; boundary = (2, 2), stag = true, multithreading = true) where {T, N}
 
     # check that boundary conditions are correctly defined
     @assert length(boundary) == 2N "Boundary conditions must be a tuple of length $(2N) for $(N)D data."
@@ -46,16 +46,18 @@ function WENOScheme(u::AbstractField{T, N}, grid; boundary = (2, 2), stag = true
 end
 
 function WENOScheme(u::AbstractField; kwargs...)
-    error("""
-    You called `WENOScheme(u)` with a `$(typeof(u))`, which is a subtype of `AbstractField`.
-
-    To construct a WENO scheme for Chmy.jl fields, you must also provide the computational grid:
-        WENOScheme(u, grid; kwargs...)
-
-    Example:
-        grid = UniformGrid(arch; origin=(0.0, 0.0), extent=(Lx, Lx), dims=(nx, ny))
-        weno = WENOScheme(u, grid; boundary=(2,2,2,2), stag=false)
-    """)
+    error(
+        """
+        You called `WENOScheme(u)` with a `$(typeof(u))`, which is a subtype of `AbstractField`.
+
+        To construct a WENO scheme for Chmy.jl fields, you must also provide the computational grid:
+            WENOScheme(u, grid; kwargs...)
+
+        Example:
+            grid = UniformGrid(arch; origin=(0.0, 0.0), extent=(Lx, Lx), dims=(nx, ny))
+            weno = WENOScheme(u, grid; boundary=(2,2,2,2), stag=false)
+        """
+    )
 end
 
 include("ChmyExt1D.jl")
diff --git a/ext/ChmyExt1D.jl b/ext/ChmyExt1D.jl
index d3f9988..84ead4a 100644
--- a/ext/ChmyExt1D.jl
+++ b/ext/ChmyExt1D.jl
@@ -1,4 +1,3 @@
-
 @kernel inbounds = true function WENO_flux_chmy_1D(fl, fr, u, boundary, nx, χ, γ, ζ, ϵ, g::StructuredGrid, O)
 
     I = @index(Global, NTuple)
diff --git a/src/1D/time_stepping.jl b/src/1D/time_stepping.jl
index 44d1247..9c7d510 100644
--- a/src/1D/time_stepping.jl
+++ b/src/1D/time_stepping.jl
@@ -37,8 +37,9 @@ function WENO_step!(u::T, v, weno::WENOScheme, Δt, Δx) where {T <: AbstractVec
     WENO_flux!(fl, fr, ut, weno, nx)
     semi_discretisation_weno5!(du, v, weno, Δx_)
 
-    return @inbounds @maybe_threads multithreading for i in axes(u, 1)
+    @inbounds @maybe_threads multithreading for i in axes(u, 1)
         u[i] = @muladd 1.0 / 3.0 * u[i] + 2.0 / 3.0 * ut[i] - (2.0 / 3.0) * Δt * du[i]
     end
 
+    return nothing
 end
diff --git a/test/test_2D_advection.jl b/test/test_2D_advection.jl
index 534d39b..e0ff9b0 100644
--- a/test/test_2D_advection.jl
+++ b/test/test_2D_advection.jl
@@ -70,7 +70,7 @@
         Δx = Lx / nx
         Δy = Lx / ny
 
-        grid = UniformGrid(arch; origin=(0.0, 0.0), extent=(Lx, Lx), dims=(nx, ny))
+        grid = UniformGrid(arch; origin = (0.0, 0.0), extent = (Lx, Lx), dims = (nx, ny))
 
         x = range(0, stop = Lx, length = nx)