diff --git a/src/DiffusionGarnet.jl b/src/DiffusionGarnet.jl
index 1fccb07..1787d71 100644
--- a/src/DiffusionGarnet.jl
+++ b/src/DiffusionGarnet.jl
@@ -23,7 +23,7 @@ using Preferences
 function set_backend(new_backend::String)
     if !(
         new_backend ∈
-        ("Threads_Float64_2D", "Threads_Float32_2D", "CUDA_Float64_2D", "CUDA_Float32_2D",        "Threads_Float64_3D", "Threads_Float32_3D", "CUDA_Float64_3D", "CUDA_Float32_3D")
+        ("Threads_Float64_2D", "Threads_Float32_2D", "CUDA_Float64_2D", "CUDA_Float32_2D", "Threads_Float64_3D", "Threads_Float32_3D", "CUDA_Float64_3D", "CUDA_Float32_3D")
     )
         throw(ArgumentError("Invalid backend: \"$(new_backend)\""))
     end
diff --git a/src/Discretisation/2D/semi_discretisation_2D.jl b/src/Discretisation/2D/semi_discretisation_2D.jl
index 2c51748..5885c5c 100644
--- a/src/Discretisation/2D/semi_discretisation_2D.jl
+++ b/src/Discretisation/2D/semi_discretisation_2D.jl
@@ -149,8 +149,6 @@ function semi_discretisation_diffusion_2D(du,u,p,t)
     dtCFe = @view du[:,:,2]
     dtCMn = @view du[:,:,3]
 
-    println(typeof(du))
-
     # update diffusive parameters
     @parallel Diffusion_coef_2D!(DMgMg, DMgFe, DMgMn, DFeMg, DFeFe, DFeMn, DMnMg, DMnFe, DMnMn,
                                  CMg, CFe ,CMn, D0, D_charact, grt_position)
diff --git a/src/input/initialconditions.jl b/src/input/initialconditions.jl
index 61a1748..921d766 100644
--- a/src/input/initialconditions.jl
+++ b/src/input/initialconditions.jl
@@ -4,14 +4,14 @@ using Parameters
 using Statistics
 
 
-@with_kw struct InitialConditions1D{T1<:AbstractArray{<:Real, 1} , T2 <: Float64}
+@with_kw struct InitialConditions1D{T1, T2, T3, T4}
     CMg0::T1
     CFe0::T1
     CMn0::T1
-    nx::Int
     Lx::T2
+    nx::T3
     Δx::T2
-    x::StepRangeLen
+    x::T4
     tfinal::T2
     function InitialConditions1D(CMg0::T1, CFe0::T1, CMn0::T1, Lx::T2, tfinal::T2) where {T1 <: AbstractArray{<:Real, 1}, T2 <: Float64}
         if Lx <= 0
@@ -24,19 +24,22 @@ using Statistics
             nx = size(CMg0, 1)
             Δx = Lx / (nx-1)
             x = range(0, length=nx, stop= Lx)
-            new{T1, T2}(CMg0, CFe0, CMn0, nx, Lx, Δx, x, tfinal)
+
+            T3 = typeof(nx)
+            T4 = typeof(x)
+            new{T1, T2, T3, T4}(CMg0, CFe0, CMn0, Lx, nx, Δx, x, tfinal)
         end
     end
 end
 
-@with_kw struct InitialConditionsSpherical{T1<:AbstractArray{<:Real, 1} , T2 <: Float64}
+@with_kw struct InitialConditionsSpherical{T1, T2, T3, T4}
     CMg0::T1
     CFe0::T1
     CMn0::T1
-    nr::Int
     Lr::T2
+    nr::T3
     Δr::T2
-    r::StepRangeLen
+    r::T4
     tfinal::T2
     function InitialConditionsSpherical(CMg0::T1, CFe0::T1, CMn0::T1, Lr::T2, tfinal::T2) where {T1 <: AbstractArray{<:Real, 1}, T2 <: Float64}
         if Lr <= 0
@@ -49,23 +52,25 @@ end
             nr = size(CMg0, 1)
             Δr = Lr / (nr-1)
             radius = range(0, length=nr, stop= Lr)
-            new{T1, T2}(CMg0, CFe0, CMn0, nr, Lr, Δr, radius, tfinal)
+            T3 = typeof(nr)
+            T4 = typeof(radius)
+            new{T1, T2, T3, T4}(CMg0, CFe0, CMn0, Lr, nr, Δr, radius, tfinal)
         end
     end
 end
 
-@with_kw struct InitialConditions2D{T1 <: AbstractArray{<:Real, 2}, T2 <: Float64}
+@with_kw struct InitialConditions2D{T1, T2, T3, T4}
     CMg0::T1
     CFe0::T1
     CMn0::T1
-    nx::Int
-    ny::Int
     Lx::T2
     Ly::T2
+    nx::T3
+    ny::T3
     Δx::T2
     Δy::T2
-    x::StepRangeLen
-    y::StepRangeLen
+    x::T4
+    y::T4
     grt_position::T1
     grt_boundary::T1
     # grid::NamedTuple{(:x, :y), Tuple{AbstractArray{T2, 2}, AbstractArray{T2, 2}}}
@@ -88,11 +93,14 @@ end
             # grid = (x=x' .* @ones(ny), y= (@ones(nx))' .* y)
 
             # define when grt is present
-            grt_position = similar(CMg0) .* 0.0
+            grt_position = similar(CMg0) .* 0
             # create a binary matrix with 1 where grt is present. 0 otherwise. This depends on if CMg0, CFe0 and CMn0 are equal to 0 or not
             grt_position .= (CMg0 .≠ 0) .& (CFe0 .≠ 0) .& (CMn0 .≠ 0)
 
-            new{T1, T2}(CMg0, CFe0, CMn0, nx, ny, Lx, Ly, Δx, Δy, x, y, grt_position, grt_boundary, tfinal)
+            T3 = typeof(nx)
+            T4 = typeof(x)
+
+            new{T1, T2, T3, T4}(CMg0, CFe0, CMn0, Lx, Ly, nx, ny, Δx, Δy, x, y, grt_position, grt_boundary, tfinal)
         end
     end
 end
@@ -228,12 +236,12 @@ end
 
         @unpack nx, Δx, tfinal, Lx, CMg0, CFe0, CMn0 = IC
 
-        D0::Vector{Float64} = zeros(Float64, 4)
+        D0 = zeros(Float64, 4)
         D_ini!(D0, T[1], P[1])  # compute initial diffusion coefficients
 
         D = (DMgMg = zeros(nx), DMgFe = zeros(nx), DMgMn = zeros(nx), DFeMg = zeros(nx), DFeFe = zeros(nx), DFeMn = zeros(nx), DMnMg = zeros(nx), DMnFe = zeros(nx), DMnMn = zeros(nx))  # tensor of interdiffusion coefficients
 
-        u0::Matrix{typeof(CMg0[1])} = similar(CMg0, (nx, 3))
+        u0 = similar(CMg0, (nx, 3))
         u0[:,1] .= CMg0
         u0[:,2] .= CFe0
         u0[:,3] .= CMn0
@@ -298,30 +306,29 @@ end
     end
 end
 
-@with_kw struct Domain2D{T1 <: Union{AbstractArray{Float64, 2}, Float64}, T2 <: Union{Float32, Float64}}
+@with_kw struct Domain2D{T1, T2, T3, T4, T5, T6}
     IC::InitialConditions2D
     T::T1
     P::T1
     time_update::T1
-    D0::AbstractArray{T2, 1}
-    D::NamedTuple{(:DMgMg, :DMgFe, :DMgMn, :DFeMg, :DFeFe, :DFeMn, :DMnMg, :DMnFe, :DMnMn),
-                  NTuple{9, Matrix{T2}}}  # tensor of interdiffusion coefficients
+    D0::T4
+    D::T5  # tensor of interdiffusion coefficients
     L_charact::T2
     D_charact::T2
     t_charact::T2
     Δxad_::T2
     Δyad_::T2
-    u0::Array{Real, 3}
+    u0::T6
     tfinal_ad::T2
     function Domain2D(IC::InitialConditions2D, T::T1, P::T1, time_update::T1) where {T1 <: Union{Float64, AbstractArray{Float64, 2}}}
         @unpack nx, ny, Δx, Δy, tfinal, Lx, CMg0, CFe0, CMn0 = IC
-
-        D0 = @zeros(4)
+        similar(CMg0, (nx,ny))
+        D0 = similar(CMg0, (4)) .* 0
         D_ini!(D0, T, P)  # compute initial diffusion coefficients
 
-        D = (DMgMg = @zeros(nx, ny), DMgFe = @zeros(nx, ny), DMgMn = @zeros(nx, ny), DFeMg = @zeros(nx, ny), DFeFe = @zeros(nx, ny), DFeMn = @zeros(nx, ny), DMnMg = @zeros(nx, ny), DMnFe = @zeros(nx, ny), DMnMn = @zeros(nx, ny))  # tensor of interdiffusion coefficients
+        D = (DMgMg = similar(CMg0, (nx,ny)) .* 0, DMgFe = similar(CMg0, (nx,ny)) .* 0, DMgMn = similar(CMg0, (nx,ny)) .* 0, DFeMg = similar(CMg0, (nx,ny)) .* 0, DFeFe = similar(CMg0, (nx,ny)) .* 0, DFeMn = similar(CMg0, (nx,ny)) .* 0, DMnMg = similar(CMg0, (nx,ny)) .* 0, DMnFe = similar(CMg0, (nx,ny)) .* 0, DMnMn = similar(CMg0, (nx,ny)) .* 0)  # tensor of interdiffusion coefficients
 
-        u0 = @zeros((nx, ny, 3))
+        u0 = similar(CMg0, (nx, ny, 3)) .* 0
         u0[:, :, 1] .= CMg0
         u0[:, :, 2] .= CFe0
         u0[:, :, 3] .= CMn0
@@ -333,7 +340,14 @@ end
         Δyad_ = 1 / (Δy / L_charact)  # inverse of nondimensionalised Δy
         tfinal_ad = tfinal / t_charact  # nondimensionalised total time
         time_update = time_update / t_charact  # nondimensionalised time update
-        new{T1, Union{Float32, Float64}}(IC, T, P, time_update, D0, D, L_charact, D_charact, t_charact, Δxad_, Δyad_, u0, tfinal_ad)
+
+        T2 = typeof(t_charact)
+        T3 = typeof(CMg0)
+        T4 = typeof(D0)
+        T5 = typeof(D)
+        T6 = typeof(u0)
+
+        new{T1, T2, T3, T4, T5, T6}(IC, T, P, time_update, D0, D, L_charact, D_charact, t_charact, Δxad_, Δyad_, u0, tfinal_ad)
     end
 end
 
diff --git a/src/simulate/simulate.jl b/src/simulate/simulate.jl
index 40e1913..4c83b0a 100644
--- a/src/simulate/simulate.jl
+++ b/src/simulate/simulate.jl
@@ -1,6 +1,6 @@
 
 """
-    simulate(domain; callback=nothing)
+    simulate(domain; callback=nothing, progressbar=true)
 
 Solve the coupled diffusion equation using finite differences for a given domain and return a solution type variable.
 
@@ -9,16 +9,18 @@ The time discretisation is based on the ROCK2 method, a stabilized explicit meth
 The solution type variable is following the format of OrdinaryDiffEq.jl (see https://docs.sciml.ai/DiffEqDocs/stable/basics/solution/), and can be used to plot the solution, and to extract the solution at a given time. As the system is nondimensionalised, the time of the solution is in nondimensional time.
 
 callbacks is an optional argument, which can be used to pass a callback function to the solver. It follows the format of DiffEqCallbacks.jl (see https://docs.sciml.ai/DiffEqCallbacks/stable/).
+
+progressbar is an optional argument, which can be used to display a progressbar during the simulation. Default is to true.
 """
 function simulate end
 
 """
-    simulate(domain::Domain1D; callback=nothing)
+    simulate(domain::Domain1D; callback=nothing, progressbar=true)
 
 Solve the coupled diffusion equation in 1D. Save all timesteps in the output solution type variable.
 
 """
-function simulate(domain::Domain1D; callback=nothing)
+function simulate(domain::Domain1D; callback=nothing, progressbar=true)
 
     @unpack tfinal_ad, u0 = domain
 
@@ -26,22 +28,18 @@ function simulate(domain::Domain1D; callback=nothing)
 
     prob = ODEProblem(semi_discretisation_diffusion_1D, u0, t, domain)
 
-    if callback === nothing
-        @time sol = solve(prob, ROCK2(), progress=true, progress_steps=1, save_start=true, abstol=1e-6,reltol=1e-6)
-    else
-        @time sol = solve(prob, ROCK2(), progress=true, progress_steps=1, save_start=true, abstol=1e-6,reltol=1e-6, callback=callback)
-    end
+    @time sol = solve(prob, ROCK2(), progress=progressbar, progress_steps=1, save_start=true, abstol=1e-6,reltol=1e-6, callback=callback)
 
     return sol
 end
 
 """
-    simulate(domain::DomainSpherical; callback=nothing)
+    simulate(domain::DomainSpherical; callback=nothing, progressbar=true)
 
 Solve the coupled diffusion equation in spherical coordinates. Save all timesteps in the output solution type variable.
 
 """
-function simulate(domain::DomainSpherical; callback=nothing)
+function simulate(domain::DomainSpherical; callback=nothing, progressbar=true)
 
     @unpack tfinal_ad, u0 = domain
 
@@ -49,22 +47,18 @@ function simulate(domain::DomainSpherical; callback=nothing)
 
     prob = ODEProblem(semi_discretisation_diffusion_spherical, u0, t, domain)
 
-    if callback === nothing
-        @time sol = solve(prob, ROCK2(), progress=true, progress_steps=1, save_start=true, abstol=1e-6,reltol=1e-6)
-    else
-        @time sol = solve(prob, ROCK2(), progress=true, progress_steps=1, save_start=true, abstol=1e-6,reltol=1e-6, callback=callback)
-    end
+    @time sol = solve(prob, ROCK2(), progress=progressbar, progress_steps=1, save_start=true, abstol=1e-6,reltol=1e-6, callback=callback)
 
     return sol
 end
 
 """
-    simulate(domain::Domain2D; callback=nothing)
+    simulate(domain::Domain2D; callback=nothing, progressbar=true)
 
 Solve the coupled diffusion equation in 2D. Save only the first and last timestep in the output solution type variable.
 
 """
-function simulate(domain::Domain2D; callback=nothing)
+function simulate(domain::Domain2D; callback=nothing, progressbar=true)
 
     @unpack tfinal_ad, u0 = domain
 
@@ -72,11 +66,7 @@ function simulate(domain::Domain2D; callback=nothing)
 
     prob = ODEProblem(semi_discretisation_diffusion_2D, u0, t, domain)
 
-    if callback === nothing
-        @time sol = solve(prob, ROCK2(), progress=true, progress_steps=1, save_start=true, abstol=1e-6,reltol=1e-6, save_everystep = false)
-    else
-        @time sol = solve(prob, ROCK2(), progress=true, progress_steps=1, save_start=true, abstol=1e-6,reltol=1e-6, save_everystep = false, callback=callback)
-    end
+    @time sol = solve(prob, ROCK2(), progress=progressbar, progress_steps=1, save_start=true, abstol=1e-6,reltol=1e-6, save_everystep = false, callback=callback)
 
     return sol
 end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 98fef41..daf3651 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -101,7 +101,7 @@ end
 
     domain1D = Domain(IC1D, T, P)
 
-    sol = simulate(domain1D)
+    sol = simulate(domain1D; progressbar=false)
 
     @test norm(sum.(sol[end][:,1] .+ sol[end][:,2] .+ sol[end][:,3])) ≈ 28.64886878627501
 end
@@ -128,11 +128,35 @@ end
 
     domainSph = Domain(ICSph, T, P)
 
-    sol = simulate(domainSph)
+    sol = simulate(domainSph; progressbar=false)
 
     @test norm(sum.(sol[end][:,1] .+ sol[end][:,2] .+ sol[end][:,3])) ≈ 20.268803083443927
 end
 
+@testset "2D Diffusion" begin
+
+    using LinearAlgebra: norm
+
+    CMg = DelimitedFiles.readdlm("./Data/2D/Xprp_LR.txt", '\t', '\n', header=false)
+    CFe = DelimitedFiles.readdlm("./Data/2D/Xalm_LR.txt", '\t', '\n', header=false)
+    CMn = DelimitedFiles.readdlm("./Data/2D/Xsps_LR.txt", '\t', '\n', header=false)
+    grt_boundary = DelimitedFiles.readdlm("./Data/2D/Contour_LR.txt", '\t', '\n', header=false)
+
+    Lx = 900.0u"µm"
+    Ly = 900.0u"µm"
+    tfinal = 1.0u"Myr"
+    T = 900u"°C"
+    P = 0.6u"GPa"
+    IC2D = InitialConditions2D(CMg, CFe, CMn, Lx, Ly, tfinal; grt_boundary = grt_boundary)
+    domain2D = Domain(IC2D, T, P)
+
+    sol = simulate(domain2D; progressbar=false)
+
+    @test norm(sol[end][:,:,1]) ≈ 12.783357041653609
+
+end
+
+
 @testset "Callback update D0" begin
 
     data = DelimitedFiles.readdlm("./Data/1D/Data_Grt_1D.txt", '\t', '\n', header=true)[1]
@@ -162,8 +186,8 @@ end
 
     update_diffusion_coef_call = PresetTimeCallback(time_update_ad, update_diffusion_coef)
 
-    sol_sph = simulate(domainSph; callback=update_diffusion_coef_call)
-    sol_1D = simulate(domain1D; callback=update_diffusion_coef_call)
+    sol_sph = simulate(domainSph; callback=update_diffusion_coef_call, progressbar=false)
+    sol_1D = simulate(domain1D; callback=update_diffusion_coef_call, progressbar=false)
 
     T=600  # in °C
     P=3  # in kbar