diff --git a/calculations/PMM.jl b/calculations/PMM.jl
index eee0769..d1b3cf2 100644
--- a/calculations/PMM.jl
+++ b/calculations/PMM.jl
@@ -20,21 +20,38 @@ H0 = H0 + transpose(H0) # symmetric
 H1 = randn(ComplexF64, N, N)
 H1 = H1 + transpose(H1) # symmetric
 
+# define the inner product for vectors
+inner(ψ1, ψ2) = only(transpose(ψ1) * ψ2)
+normalized(ψ) = ψ ./ sqrt(inner(ψ, ψ))
+overlap(ψ1_hat, ψ2) = inner(ψ1_hat, ψ2) / sqrt(inner(ψ2, ψ2))
+
 # training
 Es = ComplexF64[]
 ψs = Vector{ComplexF64}[]
 
 lr = 0.05
 epochs = 100000
+
 for epoch in 1:epochs
+    last_ψ = isempty(ψs) ? nothing : ψs[1]
     empty!(Es)
     empty!(ψs)
     for (c, E) in zip(data_c, data_E)
         H = H0 + c * H1
         evals, evecs = eigen(H)
-        i = nearestIndex(evals, E) # TODO: more robust way to identify the eigenvector
+        # identification of the eigenstate
+        if isnothing(last_ψ)
+            i = nearestIndex(evals, E)
+        else
+            overlaps = [abs(overlap(last_ψ, evecs[:, i])) for i in 1:N]
+            i = argmax(overlaps)
+            if epoch > 0.01epochs # check agreement with eigenvalues after 1% of epochs
+                i ≠ nearestIndex(evals, E) && @warn("Identification via overlap contradicts eigenvalues")
+            end
+        end
+        last_ψ = evecs[:, i] |> normalized
+        push!(ψs, last_ψ)
         push!(Es, evals[i])
-        push!(ψs, evecs[:, i])
     end
 
     if epoch % 1000 == 0