diff --git a/Project.toml b/Project.toml
index 299dce0..67589ad 100644
--- a/Project.toml
+++ b/Project.toml
@@ -10,7 +10,9 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NonadiabaticDynamicsBase = "78c76ebc-5665-4934-b512-82d81b5cbfb7"
 NonadiabaticModels = "c814dc9f-a51f-4eaf-877f-82eda4edad48"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 julia = "1"
@@ -19,6 +21,7 @@ julia = "1"
 ASE = "51974c44-a7ed-5088-b8be-3e78c8ba416c"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+UnitfulAtomic = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
 
 [targets]
-test = ["Test", "PyCall", "ASE"]
+test = ["Test", "PyCall", "ASE", "UnitfulAtomic"]
diff --git a/plot/fig2.png b/plot/fig2.png
index eaabad2..b5d5c11 100644
Binary files a/plot/fig2.png and b/plot/fig2.png differ
diff --git a/plot/plot.jl b/plot/plot.jl
index 1b5ce97..09caa4d 100644
--- a/plot/plot.jl
+++ b/plot/plot.jl
@@ -14,42 +14,34 @@ no = build.molecule("NO")
 build.add_adsorbate(slab, no, 5, "hcp")
 build.add_vacuum(slab, 30)
 
-cell, atoms, R = extract_parameters_and_positions(slab)
+c, atoms, R = extract_parameters_and_positions(slab)
 
-model = NOAu(atoms.types, cell)
+model = NOAu(atoms.types, c, R)
+slab_plane = model.H11.image.slab_plane
 
 bond_length_ref = austrip(1.15077u"Å")
 R[3,end-1] = 40
 R[3,end] = 40 + bond_length_ref
-reference = ustrip.(auconvert.(u"kJ", eigvals(potential(model, R))[1]) .* Unitful.Na)
+const reference = ustrip.(auconvert.(u"kJ", eigvals(potential(model, R))[1]) .* Unitful.Na)
 
-function plot_figure_a!(ax, reference, bond, lims)
+function plot_figure_a!(ax, bond, lims)
     bond_length = austrip(bond)
-    zN = austrip.((collect(lims[1]:0.01:lims[2]) .+ 17.0668) .*u"Å" ) 
+    zN = austrip.((collect(lims[1]:0.01:lims[2]) .+ slab_plane) .*u"Å" ) 
     traj = Matrix[]
     for z in zN
-        positions = copy(R )
+        positions = copy(R)
         positions[3,end-1] = z
         positions[3,end] = z + bond_length
         push!(traj, positions)
     end
     potentials = potential.(model, traj)
-    potentials = [ustrip.(auconvert.(u"kJ", V) .* Unitful.Na) for V in potentials]
-    ground_state = [eigvals(V)[1] for V in potentials] .- reference
-    v11 = [V[1,1] for V in potentials] .- reference
-    v12 = [V[1,2] for V in potentials]
-    v22 = [V[2,2] for V in potentials] .- reference
-    x = ustrip.(auconvert.(u"Å", zN)) .- 17.0668
-    lines!(ax, x, v11, label="V11", color=:green, linestyle=:dash)
-    lines!(ax, x, v22, label="V22", color=:red, linestyle=:dashdotdot)
-    lines!(ax, x, -v12, label="V12", color=:purple, linestyle=:dashdot)
-    lines!(ax, x, ground_state, label="E0", color=:blue)
-    xlims!(ax, lims[1], lims[2])
+    x = ustrip.(auconvert.(u"Å", zN)) .- slab_plane
+    transform_and_plot_potentials!(ax, potentials, lims, x)
     return ax
 end
 
-function plot_figure_e!(ax, reference, lims)
-    z = austrip(1.6u"Å")
+function plot_figure_e!(ax, lims)
+    z = austrip((1.6 + slab_plane)u"Å")
     rs = austrip.(collect(lims[1]:0.01:lims[2]).*u"Å")
     traj = Matrix[]
     positions = copy(R)
@@ -60,18 +52,22 @@ function plot_figure_e!(ax, reference, lims)
         push!(traj, p)
     end
     potentials = potential.(model, traj)
+    x = ustrip.(auconvert.(u"Å", rs))
+    transform_and_plot_potentials!(ax, potentials, lims, x)
+    return ax
+end
+
+function transform_and_plot_potentials!(ax, potentials, lims, x)
     potentials = [ustrip.(auconvert.(u"kJ", V) .* Unitful.Na) for V in potentials]
     ground_state = [eigvals(V)[1] for V in potentials] .- reference
     v11 = [V[1,1] for V in potentials] .- reference
     v12 = [V[1,2] for V in potentials]
     v22 = [V[2,2] for V in potentials] .- reference
-    x = ustrip.(auconvert.(u"Å", rs))
     lines!(ax, x, v11, label="V11", color=:green, linestyle=:dash)
     lines!(ax, x, v22, label="V22", color=:red, linestyle=:dashdotdot)
     lines!(ax, x, -v12, label="V12", color=:purple, linestyle=:dashdot)
     lines!(ax, x, ground_state, label="E0", color=:blue)
     xlims!(ax, lims[1], lims[2])
-    return ax
 end
 
 theme = Theme(linewidth=3)
@@ -90,13 +86,13 @@ ax3 = f[3,1] = Axis(f)
 label_c = f[3,1,TopRight()] = Label(f, "E", textsize=24)
 label_c.padding = (-50, 0, -50, 0)
 
-plot_figure_a!(ax1, reference, 1.191u"Å", (1.0, 2.5))
+plot_figure_a!(ax1, 1.191u"Å", (1.0, 2.5))
 ylims!(ax1, -50, 400)
 
-plot_figure_a!(ax2, reference, 1.6u"Å", (0.5, 3.0))
+plot_figure_a!(ax2, 1.6u"Å", (0.5, 3.0))
 ylims!(ax2, -20, 800)
 
-plot_figure_e!(ax3, reference, (0.8, 2.0))
+plot_figure_e!(ax3, (0.8, 2.0))
 ylims!(ax3, -100, 2000)
 
 ax1.xlabel = "N atom height / Å"
diff --git a/src/NOAu111.jl b/src/NOAu111.jl
index 8e01e9e..a627377 100644
--- a/src/NOAu111.jl
+++ b/src/NOAu111.jl
@@ -13,6 +13,7 @@ export NOAu
 include("parameters.jl")
 include("au_au.jl")
 include("diabatic_elements.jl")
+include("h11.jl")
 
 struct NOAu{V1,V2,V3,V4,V5,A} <: DiabaticModel
     n_states::Int
@@ -22,11 +23,13 @@ struct NOAu{V1,V2,V3,V4,V5,A} <: DiabaticModel
     AuAu::V4
     image_potential::V5
     atoms::A
+    Nindex::Int
+    Oindex::Int
 end
 
-NOAu(jatoms) = NOAu(2, H00(), H11(), H01(), AuAu(), image(D, C, zimage), jatoms)
+NOAu(jatoms, Nindex, Oindex) = NOAu(2, H00(), H11(jatoms), H01(), AuAu(), image(D, C, zimage), jatoms, Nindex, Oindex)
 
-function NOAu(symbols::AbstractVector{Symbol}, cell)
+function NOAu(symbols::AbstractVector{Symbol}, cell, R)
     jatoms = JuLIP.Atoms(
         X=zeros(3,length(symbols)),
         M=JuLIP.atomic_mass.(symbols),
@@ -34,7 +37,11 @@ function NOAu(symbols::AbstractVector{Symbol}, cell)
         cell=au_to_ang.(cell.vectors'),
         pbc=cell.periodicity
         )
-    NOAu(jatoms)
+
+    JuLIP.set_positions!(jatoms, au_to_ang.(R))
+    Nindex = findfirst(isequal(7), jatoms.Z)
+    Oindex = findfirst(isequal(8), jatoms.Z)
+    NOAu(jatoms, Nindex, Oindex)
 end
 
 function NonadiabaticModels.potential(model::NOAu, R::AbstractMatrix)
@@ -45,8 +52,6 @@ function NonadiabaticModels.potential(model::NOAu, R::AbstractMatrix)
     V22 = eV_to_au(JuLIP.energy(model.H11, model.atoms) + ϕ - Eₐ) + Au
     V12 = eV_to_au(JuLIP.energy(model.H01, model.atoms))
 
-    V22 += eV_to_au(evaluate_image_potential(model))
-
     return Hermitian(SMatrix{2,2}(V11, V12, V12, V22))
 end
 
@@ -58,12 +63,6 @@ function NonadiabaticModels.derivative!(model::NOAu, D::AbstractMatrix{<:Hermiti
     D22 = -JuLIP.forces(model.H11, model.atoms) + Au
     D12 = -JuLIP.forces(model.H01, model.atoms)
 
-    Oderiv, Nderiv = evaluate_image_derivative(model)
-    Nindex = findfirst(isequal(7), model.atoms.Z)
-    Oindex = findfirst(isequal(8), model.atoms.Z)
-    D22[Nindex] += SVector{3}(0, 0, Nderiv)
-    D22[Oindex] += SVector{3}(0, 0, Oderiv)
-
     for i=1:length(model.atoms)
         for j=1:3
             d11 = eV_per_ang_to_au(D11[i][j])
@@ -76,26 +75,4 @@ function NonadiabaticModels.derivative!(model::NOAu, D::AbstractMatrix{<:Hermiti
     return D
 end
 
-function evaluate_image_potential(model::NOAu)
-    at = model.atoms
-    Nindex = findfirst(isequal(7), at.Z)
-    Oindex = findfirst(isequal(8), at.Z)
-    total_mass = at.M[Oindex] + at.M[Nindex]
-    zcom = (at[Oindex][3]*at.M[Oindex]+at[Nindex][3]*at.M[Nindex]) / total_mass
-    image = JuLIP.evaluate(model.image_potential, zcom)
-    return image
-end
-
-function evaluate_image_derivative(model::NOAu)
-    at = model.atoms
-    Nindex = findfirst(isequal(7), at.Z)
-    Oindex = findfirst(isequal(8), at.Z)
-    total_mass = at.M[Oindex] + at.M[Nindex]
-    zcom = (at[Oindex][3]*at.M[Oindex]+at[Nindex][3]*at.M[Nindex]) / total_mass
-    image = JuLIP.evaluate_d(model.image_potential, zcom)
-    Oderiv = image * at.M[Oindex] / (total_mass)
-    Nderiv = image * at.M[Nindex] / (total_mass)
-    return Oderiv, Nderiv
-end
-
 end
diff --git a/src/diabatic_elements.jl b/src/diabatic_elements.jl
index f727e45..b8588fb 100644
--- a/src/diabatic_elements.jl
+++ b/src/diabatic_elements.jl
@@ -1,11 +1,18 @@
 
-repulsive(A, α, r_cut) = JuLIP.@analytic r -> A * (exp(-α*r) - exp(-α*r_cut))
+function repulsive(A, α, r_cut)
+    cutoff = exp(-α*r_cut)
+    JuLIP.@analytic r -> A * (exp(-α*r) - cutoff)
+end
+
 morse(F, γ, r₀) = JuLIP.@analytic r -> F*(1 - exp(-γ * (r - r₀)))^2
-coupling(A2, A3, γ, r_cut) = JuLIP.@analytic r -> -A2*(1/(1+A3*exp(γ*r)) - 1/(1+A3*exp(γ*r_cut)))
 const zero_potential = JuLIP.@analytic r -> 0
-
 image(D, C, zimage) = JuLIP.@analytic z -> -D / sqrt(C^2 + (z-zimage)^2)
 
+function coupling(A2, A3, γ, r_cut)
+    cutoff = 1/(1+A3*exp(γ*r_cut))
+    JuLIP.@analytic r -> -A2*(1/(1+A3*exp(γ*r)) - cutoff)
+end
+
 struct DiabaticElement{V,Z} <: JuLIP.SitePotential
     potentials::V
     zlist::Z
@@ -44,41 +51,60 @@ function H01()
     DiabaticElement(potentials)
 end
 
-function H11()
-
-    V11AuO = repulsive(A₁, α₁, rcutoff)
-    V11AuN = repulsive(B₁, β₁, rcutoff)
-    V11NO = morse(F₁, γ₁, r₁NO)
-
-    potentials = DefaultDict{Tuple, JuLIP.AnalyticFunction}(zero_potential)
-    potentials[1,3] = V11AuO
-    potentials[1,2] = V11AuN
-    potentials[2,3] = V11NO
-
-    DiabaticElement(potentials)
-end
-
 function JuLIP.evaluate!(tmp, V::DiabaticElement, Rs, Zs, z0)
     Es = 0.0
     i0 = JuLIP.z2i(V, z0)
-    for (R, Z) in zip(Rs, Zs)
-        i = JuLIP.z2i(V, Z)
-        f = select_potential(V, i0, i)
-        r = norm(R)
+
+    # If X == N or O site, calculate X-Au potential
+    if (i0 == 2) || (i0 == 3)
+        for (R, Z) in zip(Rs, Zs)
+            i = JuLIP.z2i(V, Z)
+            if i == 1
+                f = select_potential(V, i0, i)
+                r = norm(R)
+                Es += JuLIP.evaluate!(tmp, f, r)
+            end
+        end
+    end
+
+    # If N site, calculate NO potential
+    if i0 == 2
+        closest = find_closest(Rs, Zs, 8)
+        f = select_potential(V, i0, 3)
+        r = norm(Rs[closest])
         Es += JuLIP.evaluate!(tmp, f, r)
     end
-    return Es / 2
+
+    return Es
 end
 
 function JuLIP.evaluate_d!(dEs, tmp, V::DiabaticElement, Rs, Zs, z0)
     i0 = JuLIP.z2i(V, z0)
-    for (j, (R, Z)) in enumerate(zip(Rs, Zs))
-        i = JuLIP.z2i(V, Z)
-        f = select_potential(V, i0, i)
-        r = norm(R)
-        R̂ = R / r
-        dEs[j] = JuLIP.evaluate_d(f, r) * R̂ / 2
+
+    for i=1:length(dEs)
+        dEs[i] = zero(dEs[i])
     end
+
+    if (i0 == 2) || (i0 == 3)
+        for (j, (R, Z)) in enumerate(zip(Rs, Zs))
+            i = JuLIP.z2i(V, Z)
+            if i == 1
+                f = select_potential(V, i0, i)
+                r = norm(R)
+                R̂ = R / r
+                dEs[j] = JuLIP.evaluate_d(f, r) * R̂
+            end
+        end
+    end
+
+    if i0 == 2
+        closest = find_closest(Rs, Zs, 8)
+        f = select_potential(V, i0, 3)
+        r = norm(Rs[closest])
+        R̂ = Rs[closest] / r
+        dEs[closest] = JuLIP.evaluate_d(f, r) * R̂
+    end
+
     return dEs
 end
 
diff --git a/src/h11.jl b/src/h11.jl
new file mode 100644
index 0000000..dd17e46
--- /dev/null
+++ b/src/h11.jl
@@ -0,0 +1,163 @@
+
+using UnPack
+using Zygote
+
+struct H11{I,A,V} <: JuLIP.AbstractCalculator
+    image::I
+    AuN::A
+    potential::V
+end
+
+function H11(jatoms)
+    image = ImagePotential(jatoms)
+    AuN = AuNCoupling()
+    potential = H11DiabaticElement()
+    H11(image, AuN, potential)
+end
+
+function H11DiabaticElement()
+    V11AuO = repulsive(A₁, α₁, rcutoff)
+    V11NO = morse(F₁, γ₁, r₁NO)
+
+    potentials = DefaultDict{Tuple, JuLIP.AnalyticFunction}(zero_potential)
+    potentials[1,3] = V11AuO
+    potentials[2,3] = V11NO
+
+    DiabaticElement(potentials)
+end
+
+function JuLIP.energy(calc::H11, at::JuLIP.AbstractAtoms; domain=1:length(at))
+    e1 = JuLIP.energy(calc.image, at)
+    e2 = JuLIP.energy(calc.AuN, at; domain=domain)
+    e3 = JuLIP.energy(calc.potential, at; domain=domain)
+    return e1 + e2 + e3
+end
+
+function JuLIP.forces(calc::H11, at::JuLIP.AbstractAtoms; kwargs...)
+    f1 = JuLIP.forces(calc.image, at)
+    f2 = JuLIP.forces(calc.AuN, at; kwargs...)
+    f3 = JuLIP.forces(calc.potential, at; kwargs...)
+    return f1 + f2 + f3
+end
+
+struct ImagePotential{V,T} <: JuLIP.AbstractCalculator
+    V::V
+    mO::T
+    mN::T
+    mtotal::T
+    iN::Int
+    iO::Int
+    slab_plane::T
+end
+
+function ImagePotential(jatoms)
+    iN = findfirst(isequal(7), jatoms.Z)
+    iO = findfirst(isequal(8), jatoms.Z)
+    mN = jatoms.M[iN]
+    mO = jatoms.M[iO]
+    mtotal = mO + mN
+
+    gold_positions = jatoms.X[jatoms.Z .== 79]
+    slab_plane = maximum(x->x[3], gold_positions)
+
+    V = image(D, C, zimage)
+    ImagePotential(V, mO, mN, mtotal, iN, iO, slab_plane)
+end
+
+JuLIP.energy(calc::ImagePotential, at::JuLIP.AbstractAtoms) = JuLIP.evaluate(calc, at.X)
+JuLIP.forces(calc::ImagePotential, at::JuLIP.AbstractAtoms) = -JuLIP.evaluate_d(calc, at.X)
+
+function JuLIP.evaluate(calc::ImagePotential, R)
+    @unpack mO, mN, mtotal, iN, iO = calc
+    zcom = (R[iO][3]*mO + R[iN][3]*mN) / mtotal - calc.slab_plane
+    return JuLIP.evaluate(calc.V, zcom)
+end
+
+function JuLIP.evaluate_d(calc::ImagePotential, R)
+    @unpack mO, mN, mtotal, iN, iO = calc
+    zcom = (R[iO][3]*mO + R[iN][3]*mN) / mtotal - calc.slab_plane
+    image = JuLIP.evaluate_d(calc.V, zcom)
+    out = zero(R)
+    out[iN] = SVector{3}(0, 0, image * mN / mtotal)
+    out[iO] = SVector{3}(0, 0, image * mO / mtotal)
+    return out
+end
+
+struct AuNCoupling{V,Z} <: JuLIP.SitePotential
+    V::V
+    zlist::Z
+end
+
+function AuNCoupling()
+    V(cosθ, r) = B₁*(exp(-2β₁*(r-r₁AuN)) - exp(-2β₁*(rcutoff-r₁AuN))
+                - 2cosθ^2*(exp(-β₁*(r-r₁AuN)) - exp(-β₁*(rcutoff-r₁AuN)))
+            )
+    zlist = JuLIP.Potentials.ZList([:Au, :N, :O])
+    AuNCoupling(V, zlist)
+end
+
+function JuLIP.evaluate!(tmp, V::AuNCoupling, Rs, Zs, z0)
+    Es = 0.0
+    i0 = JuLIP.z2i(V, z0)
+    if i0 == 2 # Evaluated only at N site
+        cosθ = calculate_cosθ(Rs, Zs)
+        for (R, Z) in zip(Rs, Zs)
+            i = JuLIP.z2i(V, Z)
+            if i == 1
+                r = norm(R)
+                Es += V.V(cosθ, r)
+            end
+        end
+    end
+    return Es
+end
+
+function calculate_cosθ(Rs, Zs)
+    closest = find_closest(Rs, Zs, 8)
+    NO_distance = Rs[closest]
+    cosθ = NO_distance[3] / norm(NO_distance)
+    return cosθ
+end
+
+"""
+Find the nearest atom with `z == ztarget`.
+This is used to make sure the NO molecule interacts only with itself,
+not the periodic images.
+"""
+function find_closest(Rs, Zs, ztarget)
+    target_indices = findall(isequal(ztarget), Zs)
+    max_dist = 100
+    closest = target_indices[1]
+    for i in target_indices
+        dist = norm(Rs[i])
+        if dist < max_dist
+            closest = i
+            max_dist = dist
+        end
+    end
+    return closest
+end
+
+"""
+This is the slowest part of the code, a hand-coded derivative would be faster,
+or possibly a different autodiff package.
+"""
+function JuLIP.evaluate_d!(dEs, tmp, V::AuNCoupling, Rs, Zs, z0)
+    grad = Zygote.gradient(x -> JuLIP.evaluate!(tmp, V, x, Zs, z0), Rs)[1]
+    if grad === nothing
+        for i=1:length(dEs)
+            dEs[i] = zero(dEs[i])
+        end
+    else
+        for i=1:length(grad)
+            if grad[i] === nothing
+                dEs[i] = zero(dEs[i])
+            else
+                dEs[i] = grad[i]
+            end
+        end
+    end
+    return dEs
+end
+
+JuLIP.cutoff(::AuNCoupling) = rcutoff
diff --git a/test/runtests.jl b/test/runtests.jl
index 96cec24..e9336f1 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -5,29 +5,35 @@ using PyCall
 using ASE
 using NonadiabaticDynamicsBase
 using NonadiabaticModels
-using LinearAlgebra: norm
+using Unitful, UnitfulAtomic
 
 build = pyimport("ase.build")
 
-slab = build.fcc111("Au", size=(2,2,3), vacuum=10.0)
+slab = build.fcc111("Au", size=(2,4,4), vacuum=10.0)
 no = build.molecule("NO")
 build.add_adsorbate(slab, no, 5, "ontop")
 
 at = JuLIP.Atoms(ASEAtoms(slab))
 rattle!(at, 0.5)
 
+@testset "H11" begin
+    @test JuLIP.Testing.fdtest(NOAu111.AuNCoupling(), at, verbose=false)
+    @test JuLIP.Testing.fdtest(NOAu111.H11DiabaticElement(), at, verbose=false)
+    @test JuLIP.Testing.fdtest(NOAu111.ImagePotential(at), at, verbose=false)
+end
+
 @testset "Components" begin
     @test JuLIP.Testing.fdtest(NOAu111.AuAu(), at, verbose=false)
     @test JuLIP.Testing.fdtest(NOAu111.H00(), at, verbose=false)
     @test JuLIP.Testing.fdtest(NOAu111.H01(), at, verbose=false)
-    @test JuLIP.Testing.fdtest(NOAu111.H11(), at, verbose=false)
+    @test JuLIP.Testing.fdtest(NOAu111.H11(at), at, verbose=false)
 end
 
 @testset "Final model" begin
-    R = ang_to_au.(Matrix(hcat(at.X...)))
-    c = PeriodicCell(ang_to_au.(cell(at)'))
+    R = austrip.(Matrix(hcat(at.X...))u"Å")
+    c = NonadiabaticDynamicsBase.PeriodicCell(austrip.(cell(at)' .* u"Å"))
 
-    model = NOAu(chemical_symbols(at), c)
+    model = NOAu(chemical_symbols(at), c, R)
 
     @test JuLIP.Testing.fdtest(x->potential(model, x), x->derivative(model, x), R, verbose=false)
 end