diff --git a/analysis/mean_squared_displacement.jl b/analysis/mean_squared_displacement.jl
index 586a547..7718e38 100644
--- a/analysis/mean_squared_displacement.jl
+++ b/analysis/mean_squared_displacement.jl
@@ -4,19 +4,21 @@ using Dates: Dates
 using Random: Random
 using StaticArrays: SVector
 using JLD2: JLD2
+using CellListMap: CellListMap
 
 using ReCo: ReCo
 
-# Disable the periodic boundary conditions
-ReCo.update_verlet_lists!(args...) = nothing
-ReCo.gen_cell_list(args...) = nothing
-ReCo.gen_cell_list_box(args...) = nothing
-ReCo.push_to_verlet_list!(args...) = nothing
-ReCo.restrict_coordinate(value::Float64, args...) = value
-ReCo.restrict_coordinates(v::SVector{2,Float64}, args...) = v
-ReCo.restrict_coordinates!(args...) = nothing
-ReCo.minimum_image_coordinate(value::Float64, args...) = value
-ReCo.minimum_image(v::SVector{2,Float64}, args...) = v
+# IMPORTANT: Disable the periodic boundary conditions
+# The arguments types have to match for the function to be overwritten!
+ReCo.update_verlet_lists!(::Any, ::CellListMap.CellList) = nothing
+ReCo.gen_cell_list(::Vector{SVector{2,Float64}}, ::CellListMap.Box) = nothing
+ReCo.gen_cell_list_box(::Float64, ::Float64) = nothing
+ReCo.push_to_verlet_list!(::Any, ::Any, ::Any) = nothing
+ReCo.restrict_coordinate(value::Float64, ::Float64) = value
+ReCo.restrict_coordinates(v::SVector{2,Float64}, ::Float64) = v
+ReCo.restrict_coordinates!(::Particle, ::Float64) = nothing
+ReCo.minimum_image_coordinate(value::Float64, ::Float64) = value
+ReCo.minimum_image(v::SVector{2,Float64}, ::Float64) = v
 
 function mean_squared_displacement(;
     n_simulations::Int64, v₀s::AbstractVector{Float64}, T::Float64
@@ -49,7 +51,13 @@ function mean_squared_displacement(;
 
             sim_dirs[sim_ind, v₀_ind] = dir
 
-            ReCo.run_sim(dir; duration=T, seed=rand(1:typemax(Int64)), snapshot_at=0.01)
+            ReCo.run_sim(
+                dir;
+                duration=T,
+                seed=rand(1:typemax(Int64)),
+                snapshot_at=0.01,
+                n_bundle_snapshots=200,
+            )
         end
     end
 
diff --git a/src/simulation.jl b/src/simulation.jl
index 2b4bb7a..4a9d170 100644
--- a/src/simulation.jl
+++ b/src/simulation.jl
@@ -12,7 +12,7 @@ function push_to_verlet_list!(verlet_lists, i, j)
     return nothing
 end
 
-function update_verlet_lists!(args, cl)
+function update_verlet_lists!(args, cl::CellListMap.CellList)
     @simd for pre_vec in args.verlet_lists
         reset!(pre_vec)
     end