refactor(fssh): reduce repetition in rescaling (#196)

src/DynamicsMethods/SurfaceHoppingMethods/fssh.jl

@@ -62,14 +62,26 @@ Evaluates the probability of hopping from the current state to all other states
 - 'd' is skew-symmetric so here the indices are important.
 """
 function evaluate_hopping_probability!(sim::Simulation{<:FSSH}, u, dt)
-    v = DynamicsUtils.get_velocities(u)
+    v = get_hopping_velocity(sim, u)
     σ = DynamicsUtils.get_quantum_subsystem(u)
     s = sim.method.state
-    d = sim.calculator.nonadiabatic_coupling
+    d = get_hopping_nonadiabatic_coupling(sim)
 
     fewest_switches_probability!(sim.method.hopping_probability, v, σ, s, d, dt)
 end
 
+function get_hopping_nonadiabatic_coupling(sim::Simulation{<:SurfaceHopping})
+    sim.calculator.nonadiabatic_coupling
+end
+
+function get_hopping_velocity(::Simulation{<:SurfaceHopping}, u)
+    DynamicsUtils.get_velocities(u)
+end
+
+function get_hopping_eigenvalues(sim::Simulation{<:SurfaceHopping})
+    sim.calculator.eigenvalues
+end
+
 function fewest_switches_probability!(probability, v, σ, s, d, dt)
     probability .= 0 # Set all entries to 0
     for m in axes(σ, 1)
@@ -102,9 +114,9 @@ function rescale_velocity!(sim::AbstractSimulation{<:FSSH}, u)::Bool
 
     old_state = sim.method.state
     new_state = sim.method.new_state
-    velocity = DynamicsUtils.get_velocities(u)
+    velocity = get_hopping_velocity(sim, u)
 
-    c = calculate_potential_energy_change(sim.calculator, new_state, old_state)
+    c = calculate_potential_energy_change(sim, new_state, old_state)
     a, b = evaluate_a_and_b(sim, velocity, new_state, old_state)
     discriminant = b.^2 .- 2a.*c
 
@@ -114,32 +126,35 @@ function rescale_velocity!(sim::AbstractSimulation{<:FSSH}, u)::Bool
     plus = (b .+ root) ./ a
     minus = (b .- root) ./ a 
     velocity_rescale = sum(abs.(plus)) < sum(abs.(minus)) ? plus : minus
-    perform_rescaling!(sim, velocity, velocity_rescale, new_state, old_state)
+    perform_rescaling!(sim, DynamicsUtils.get_velocities(u), velocity_rescale, new_state, old_state)
 
     return true
 end
 
-function evaluate_a_and_b(sim::Simulation{<:SurfaceHopping}, velocity, new_state, old_state)
+function evaluate_a_and_b(sim::AbstractSimulation{<:SurfaceHopping}, velocity, new_state, old_state)
     a = zeros(length(sim.atoms))
     b = zero(a)
+    d = get_hopping_nonadiabatic_coupling(sim)
     @views for i in range(sim.atoms)
-        coupling = [sim.calculator.nonadiabatic_coupling[j,i][new_state, old_state] for j=1:ndofs(sim)]
+        coupling = [d[j,i][new_state, old_state] for j=1:ndofs(sim)]
         a[i] = dot(coupling, coupling) / sim.atoms.masses[i]
         b[i] = dot(velocity[:,i], coupling)
     end
     return (a, b)
 end
 
-function perform_rescaling!(sim::Simulation{<:SurfaceHopping}, velocity, velocity_rescale, new_state, old_state)
+function perform_rescaling!(sim::AbstractSimulation{<:SurfaceHopping}, velocity, velocity_rescale, new_state, old_state)
+    d = get_hopping_nonadiabatic_coupling(sim)
     for i in range(sim.atoms)
-        coupling = [sim.calculator.nonadiabatic_coupling[j,i][new_state, old_state] for j=1:ndofs(sim)]
+        coupling = [d[j,i][new_state, old_state] for j=1:ndofs(sim)]
         velocity[:,i] .-= velocity_rescale[i] .* coupling ./ sim.atoms.masses[i]
     end
     return nothing
 end
 
-function calculate_potential_energy_change(calc::Calculators.AbstractDiabaticCalculator, new_state::Integer, current_state::Integer)
-    return calc.eigenvalues[new_state] - calc.eigenvalues[current_state]
+function calculate_potential_energy_change(sim::AbstractSimulation{<:SurfaceHopping}, new_state::Integer, current_state::Integer)
+    eigs = get_hopping_eigenvalues(sim)
+    return eigs[new_state] - eigs[current_state]
 end
 
 function Estimators.diabatic_population(sim::AbstractSimulation{<:FSSH}, u)

src/DynamicsMethods/SurfaceHoppingMethods/iesh.jl

@@ -196,7 +196,7 @@ function rescale_velocity!(sim::Simulation{<:IESH}, u)::Bool
     new_state, old_state = symdiff(sim.method.new_state, sim.method.state)
     velocity = DynamicsUtils.get_velocities(u)
 
-    c = calculate_potential_energy_change(sim.calculator, new_state, old_state)
+    c = calculate_potential_energy_change(sim, new_state, old_state)
     a, b = evaluate_a_and_b(sim, velocity, new_state, old_state)
     discriminant = b.^2 .- 2a.*c
 

src/DynamicsMethods/SurfaceHoppingMethods/rpsh.jl

@@ -26,61 +26,27 @@ function DynamicsMethods.motion!(du, u, sim::RingPolymerSimulation{<:SurfaceHopp
     DynamicsUtils.set_quantum_derivative!(dσ, v, σ, sim)
 end
 
-function evaluate_hopping_probability!(sim::RingPolymerSimulation{<:FSSH}, u, dt)
-    v = RingPolymers.get_centroid(DynamicsUtils.get_velocities(u))
-    σ = DynamicsUtils.get_quantum_subsystem(u)
-    s = sim.method.state
-    d = sim.calculator.centroid_nonadiabatic_coupling
-
-    fewest_switches_probability!(sim.method.hopping_probability, v, σ, s, d, dt)
+function get_hopping_nonadiabatic_coupling(sim::RingPolymerSimulation{<:FSSH})
+    sim.calculator.centroid_nonadiabatic_coupling
 end
 
-function rescale_velocity!(sim::RingPolymerSimulation{<:FSSH}, u)::Bool
-    sim.method.rescaling === :off && return true
-
-    old_state = sim.method.state
-    new_state = sim.method.new_state
-    velocity = DynamicsUtils.get_velocities(u)
-    centroid_velocity = RingPolymers.get_centroid(DynamicsUtils.get_velocities(u))
-
-    c = calculate_potential_energy_change(sim.calculator, new_state, old_state)
-    a, b = evaluate_a_and_b(sim, centroid_velocity, new_state, old_state)
-    discriminant = b.^2 .- 2a.*c
-
-    any(discriminant .< 0) && return false
-
-    root = sqrt.(discriminant)
-    plus = (b .+ root) ./ a
-    minus = (b .- root) ./ a 
-    velocity_rescale = sum(abs.(plus)) < sum(abs.(minus)) ? plus : minus
-    perform_rescaling!(sim, velocity, velocity_rescale, new_state, old_state)
-
-    return true
+function get_hopping_velocity(::RingPolymerSimulation{<:FSSH}, u)
+    RingPolymers.get_centroid(DynamicsUtils.get_velocities(u))
 end
 
-function evaluate_a_and_b(sim::RingPolymerSimulation{<:SurfaceHopping}, velocity, new_state, old_state)
-    a = zeros(length(sim.atoms))
-    b = zero(a)
-    @views for i in range(sim.atoms)
-        coupling = [sim.calculator.centroid_nonadiabatic_coupling[j,i][new_state, old_state] for j=1:ndofs(sim)]
-        a[i] = dot(coupling, coupling) / sim.atoms.masses[i]
-        b[i] = dot(velocity[:,i], coupling)
-    end
-    return (a, b)
+function get_hopping_eigenvalues(sim::RingPolymerSimulation{<:FSSH})
+    sim.calculator.centroid_eigenvalues
 end
 
-function perform_rescaling!(sim::RingPolymerSimulation{<:SurfaceHopping}, velocity, velocity_rescale, new_state, old_state)
+function perform_rescaling!(sim::RingPolymerSimulation{<:FSSH}, velocity, velocity_rescale, new_state, old_state)
+    d = get_hopping_nonadiabatic_coupling(sim)
     for i in range(sim.atoms)
-        coupling = [sim.calculator.centroid_nonadiabatic_coupling[j,i][new_state, old_state] for j=1:ndofs(sim)]
+        coupling = [d[j,i][new_state, old_state] for j=1:ndofs(sim)]
         velocity[:,i,:] .-= velocity_rescale[i] .* coupling ./ sim.atoms.masses[i]
     end
     return nothing
 end
 
-function calculate_potential_energy_change(calc::RingPolymerDiabaticCalculator, new_state::Integer, current_state::Integer)
-    return calc.centroid_eigenvalues[new_state] - calc.centroid_eigenvalues[current_state]
-end
-
 function DynamicsUtils.classical_hamiltonian(sim::RingPolymerSimulation{<:FSSH}, u)
     kinetic = DynamicsUtils.classical_kinetic_energy(sim, DynamicsUtils.get_velocities(u))
     spring = RingPolymers.get_spring_energy(sim.beads, sim.atoms.masses, DynamicsUtils.get_positions(u))

test/Dynamics/fssh.jl

@@ -64,16 +64,16 @@ atoms = Atoms(1)
 
     @testset "calculate_potential_energy_change" begin
         integrator.p.calculator.eigenvalues = SVector{2}([0.9, -0.3])
-        @test 1.2 ≈ SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p.calculator, 1, 2)
-        @test -1.2 ≈ SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p.calculator, 2, 1)
+        @test 1.2 ≈ SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p, 1, 2)
+        @test -1.2 ≈ SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p, 2, 1)
     end
 
     @testset "execute_hop!" begin
         Calculators.update_electronics!(integrator.p.calculator, get_positions(integrator.u))
         DynamicsUtils.get_velocities(integrator.u) .= 2 # Set high momentum to ensure successful hop
         integrator.u.state = 1
         integrator.p.method.new_state = 2
-        ΔE = SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p.calculator, 2, 1)
+        ΔE = SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p, 2, 1)
 
         KE_initial = DynamicsUtils.classical_kinetic_energy(integrator.p, get_velocities(integrator.u))
         H_initial = DynamicsUtils.classical_hamiltonian(integrator.p, integrator.u)
@@ -126,16 +126,16 @@ end
 
     @testset "calculate_potential_energy_change" begin
         integrator.p.calculator.centroid_eigenvalues = SVector{2}([0.9, -0.3])
-        @test 1.2 ≈ SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p.calculator, 1, 2)
-        @test -1.2 ≈ SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p.calculator, 2, 1)
+        @test 1.2 ≈ SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p, 1, 2)
+        @test -1.2 ≈ SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p, 2, 1)
     end
 
     @testset "execute_hop!" begin
         Calculators.update_electronics!(integrator.p.calculator, get_positions(integrator.u))
         get_velocities(integrator.u) .= 5 # Set high momentum to ensure successful hop
         integrator.u.state = 1
         integrator.p.method.new_state = 2
-        ΔE = SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p.calculator, 2, 1)
+        ΔE = SurfaceHoppingMethods.calculate_potential_energy_change(integrator.p, 2, 1)
 
         KE_initial = DynamicsUtils.classical_kinetic_energy(integrator.p, get_velocities(integrator.u))
         H_initial = DynamicsUtils.classical_hamiltonian(integrator.p, integrator.u)