diff --git a/src/Analysis/diatomic.jl b/src/Analysis/diatomic.jl
index e8a74601a..5fdfe392f 100644
--- a/src/Analysis/diatomic.jl
+++ b/src/Analysis/diatomic.jl
@@ -29,7 +29,7 @@ function surface_distance_condition(
     ) # Height of the molecule in the surface normal direction
     
     # Get the height of all substrate atoms in surface normal direction. 
-    substrate_heights = [surface_normal_height(get_positions(x)[:, substrate_atom_id]) for substrate_atom_id in symdiff(1:length(simulation.atoms.masses), indices, surface_normal)]
+    substrate_heights = [surface_normal_height(get_positions(x)[:, substrate_atom_id], surface_normal) for substrate_atom_id in symdiff(1:length(simulation.atoms.masses), indices)]
     
     # Ignore substrate above molecule in case PBC wrapping puts one above the diatomic
     highest_z = max(substrate_heights[substrate_heights.≤molecule_position]...) 
@@ -60,8 +60,8 @@ end
 
     Evaluate true if the diatomic bond length is below `threshold`.
 """
-function close_approach_condition(x::AbstractArray, indices::Vector{Int}, simulation::AbstractSimulation; threshold = 1.5u"Å")
-    if Structure.pbc_distance(x, indices..., simulation) ≤ threshold
+function close_approach_condition(x::AbstractArray, indices::Vector{Int}, simulation::AbstractSimulation; threshold = austrip(1.5u"Å"))
+    if austrip(Structure.pbc_distance(x, indices..., simulation)) ≤ threshold
         return true
     else
         return false
@@ -82,7 +82,7 @@ This is evaluated using two conditions:
 If the second condition is never reached (can happen for particularly quick desorptions), the `fallback_distance_threshold` is used to find the last point where the
 diatomic bond length is above the given value and saves from that point onwards. 
 """
-function get_desorption_frame(trajectory::AbstractVector, diatomic_indices::Vector{Int}, simulation::AbstractSimulation; surface_normal::Vector=[0, 0, 1], surface_distance_threshold=5.0 * u"Å", fallback_distance_threshold = 1.5u"Å")
+function get_desorption_frame(trajectory::AbstractVector, diatomic_indices::Vector{Int}, simulation::AbstractSimulation; surface_normal::Vector=[0, 0, 1], surface_distance_threshold=austrip(5.0 * u"Å"), fallback_distance_threshold = austrip(1.5u"Å"))
     desorbed_frame = findfirst(surface_distance_condition.(trajectory, Ref(diatomic_indices), Ref(simulation); surface_distance_threshold=surface_distance_threshold))
 
     if isa(desorbed_frame, Nothing)
diff --git a/src/structure.jl b/src/structure.jl
index b635ca1ac..b0263dd37 100644
--- a/src/structure.jl
+++ b/src/structure.jl
@@ -152,32 +152,32 @@ function minimum_distance_translation(config::AbstractVector, ind1::Int, ind2::I
     return minimum_distance_translation(get_positions(config),ind1,ind2,simulation.cell;cutoff=cutoff)
 end
 
-function pbc_distance(config::Matrix, ind1, ind2, sim::AbstractSimulation; args...)
+function pbc_distance(config::Matrix, ind1::int_or_index, ind2::int_or_index, sim::AbstractSimulation; args...)
     return pbc_distance(config,ind1,ind2,sim.cell; args...)
 end
 
-function pbc_distance(config::AbstractVector, ind1, ind2, sim::AbstractSimulation; args...)
+function pbc_distance(config::AbstractVector, ind1::int_or_index, ind2::int_or_index, sim::AbstractSimulation; args...)
     return pbc_distance(get_positions(config),ind1,ind2,sim.cell; args...)
 end
 
-function pbc_center_of_mass(config::Matrix, ind1, ind2, sim::AbstractSimulation; args...)
+function pbc_center_of_mass(config::Matrix, ind1::int_or_index, ind2::int_or_index, sim::AbstractSimulation; args...)
     return pbc_center_of_mass(config,ind1,ind2,sim.cell, sim.atoms; args...)
 end
 
-function pbc_center_of_mass(config::AbstractVector, ind1, ind2, sim::AbstractSimulation; args...)
+function pbc_center_of_mass(config::AbstractVector, ind1::int_or_index, ind2::int_or_index, sim::AbstractSimulation; args...)
     return pbc_center_of_mass(get_positions(config),ind1,ind2,sim.cell, sim.atoms; args...)
 end
 
-function velocity_center_of_mass(config::Matrix, ind1, ind2, sim::AbstractSimulation)
+function velocity_center_of_mass(config::Matrix, ind1::int_or_index, ind2::int_or_index, sim::AbstractSimulation)
     return velocity_center_of_mass(config,ind1,ind2, sim.atoms)
 end
 
-function velocity_center_of_mass(config::AbstractVector, ind1, ind2, sim::AbstractSimulation)
+function velocity_center_of_mass(config::AbstractVector, ind1::int_or_index, ind2::int_or_index, sim::AbstractSimulation)
     return velocity_center_of_mass(get_velocities(config),ind1,ind2, sim.atoms)
 end
 
 
-export minimum_distance_translation, pbc_distance, distance, angle_between, pbc_center_of_mass, velocity_center_of_mass, center_of_mass, OutputSubsetKineticEnergy, reduced_mass, fractional_mass
+export minimum_distance_translation, pbc_distance, distance, angle_between, pbc_center_of_mass, velocity_center_of_mass, center_of_mass, reduced_mass, fractional_mass
 
 
-end
\ No newline at end of file
+end
diff --git a/test/Analysis/diatomic.jl b/test/Analysis/diatomic.jl
index b80d8a46a..5daabfa50 100644
--- a/test/Analysis/diatomic.jl
+++ b/test/Analysis/diatomic.jl
@@ -13,6 +13,6 @@ using JLD2
     atoms, initial_positions, cell=NQCDynamics.convert_from_ase_atoms(desorption_trajectory_ase[1])
     simulation=Simulation(atoms, AdiabaticASEModel(desorption_trajectory_ase[1]), cell=cell)
     diatomic_indices=[55,56]
-    @test Analysis.Diatomic.get_desorption_frame(desorption_dynamicsvariables, diatomic_indices, simulation; surface_distance_threshold=2.4u"Å") == 2675
-    @test Analysis.Diatomic.get_desorption_angle(desorption_dynamicsvariables, diatomic_indices, simulation; surface_distance_threshold=2.4u"Å") ≈ 28.05088202518
-end
\ No newline at end of file
+    @test Analysis.Diatomic.get_desorption_frame(desorption_dynamicsvariables, diatomic_indices, simulation; surface_distance_threshold=austrip(2.4u"Å")) == 2675
+    @test Analysis.Diatomic.get_desorption_angle(desorption_dynamicsvariables, diatomic_indices, simulation; surface_distance_threshold=austrip(2.4u"Å")) ≈ 28.05088202518
+end