diff --git a/calculations/3body_HO_B2R_EC.jl b/calculations/3body_HO_B2R_EC.jl
index cd4acee..ee52f80 100644
--- a/calculations/3body_HO_B2R_EC.jl
+++ b/calculations/3body_HO_B2R_EC.jl
@@ -1,18 +1,25 @@
+include("../ho_basis.jl")
 include("../EC.jl")
 
+V_of_r(r) = 2 * exp(-(r-3)^2 / (1.5)^2)
+Λ = 0
+m = 1.0
+
+μω_global = 0.5 * exp(-2im * pi / 9)
+E_max = 30
+
+H0 = get_3b_H_matrix(src, V_of_r, μω_global, E_max, Λ, m, true, true)
+
+# Vp = perturbation to make the state artificially bound
+Vp_of_r(r) = -exp(-(r/3)^2)
+@time "Vp" Vp = get_3b_H_matrix(src, Vp_of_r, μω_global, E_max, Λ, m, false, true)
+
 training_ref = -0.72763
 extrapolating_ref = 4.0766890719636635 - 0.01275892774109674im
 
 training_c = [2.0, 1.9, 1.8]
 extrapolating_c = 0.0 : 0.2 : 1.2
 
-include("../ho_basis_3body_resonance.jl")
-H0 = H
-
-# Vp = perturbation to make the state artificially bound
-Vp_of_r(r) = -exp(-(r/3)^2)
-@time "Vp" Vp =  get_src_V_matrix(Vp_of_r, basis, μω, μω_global)
-
 EC = affine_EC(H0, Vp)
 train!(EC, training_c; ref_eval=training_ref, CAEC=true)
 extrapolate!(EC, extrapolating_c; ref_eval=extrapolating_ref)
diff --git a/calculations/3body_HO_R2R_EC.jl b/calculations/3body_HO_R2R_EC.jl
index 6c85598..aee2417 100644
--- a/calculations/3body_HO_R2R_EC.jl
+++ b/calculations/3body_HO_R2R_EC.jl
@@ -1,14 +1,22 @@
+include("../ho_basis.jl")
 include("../EC.jl")
 
+V_of_r(r) = 2 * exp(-(r-3)^2 / (1.5)^2)
+Λ = 0
+m = 1.0
+
+μω_global = 0.5 * exp(-2im * pi / 9)
+E_max = 30
+
+T = get_3b_H_matrix(src, V_of_r, μω_global, E_max, Λ, m, true, false)
+V = get_3b_H_matrix(src, V_of_r, μω_global, E_max, Λ, m, false, true)
+
 ref_E = 5.9673 - 0.0006im
 
 training_c = 2.0 : -0.2 : 1.2
 extrapolating_c = 1.05 .- [0.0 : 0.1 : 0.4; 0.45 : 0.05 : 0.60]
 
-include("../ho_basis_3body_resonance.jl")
-@time "H0" H0 = T1 + T2 + T_cross
-
-EC = affine_EC(H0, V)
+EC = affine_EC(T, V)
 train!(EC, training_c; ref_eval=ref_E, CAEC=false)
 extrapolate!(EC, extrapolating_c)
 
diff --git a/common.jl b/common.jl
index 69973c6..1201230 100644
--- a/common.jl
+++ b/common.jl
@@ -1,5 +1,7 @@
 using LinearAlgebra, DelimitedFiles, SparseArrays
 
+@enum coordinate_system jacobi src
+
 "Sum over array while minimizing catastrophic cancellation as much as possible"
 function better_sum(arr::Array{T}) where T<:Real
     pos_arr = arr[arr .> 0]
diff --git a/ho_basis.jl b/ho_basis.jl
index 7d7219c..9bbc913 100644
--- a/ho_basis.jl
+++ b/ho_basis.jl
@@ -237,10 +237,7 @@ function get_W_matrix(basis_p, basis::ho_basis_2B, μ1ω1, μ2ω2=μ1ω1; weight
     return sparse(W)
 end
 
-@enum coordinate_system jacobi src
-
-function get_3b_H_matrix(coord_system::coordinate_system, V_of_r, μω_global, E_max, Λ, m = 1.0; atol = 10^-5, maxevals = 10^5, verbose=true)
-
+function get_3b_H_matrix(coord_system::coordinate_system, V_of_r, μω_global, E_max, Λ, m=1.0, kinetic_part=true, potential_part=true; atol = 10^-5, maxevals = 10^5, verbose=true)
     if coord_system == jacobi
         μ1ω1 = μω_global * 1/2
         μ2ω2 = μω_global * 2
@@ -254,27 +251,27 @@ function get_3b_H_matrix(coord_system::coordinate_system, V_of_r, μω_global, E
     verbose || println("No of threads = ", Threads.nthreads())
 
     @time "Basis" basis = ho_basis_2B(E_max, Λ)
-
     verbose || println("Basis size = ", basis.dim)
 
-    verbose || println("Constructing KE matrices")
-    @time "T1" T1 = get_sp_T_matrix(basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; μω_gen=μ1ω1, μ=μ1)
-    @time "T2" T2 = get_sp_T_matrix(basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; μω_gen=μ2ω2, μ=μ2)
-    if coord_system == src
-        @time "T_cross" T_cross = get_2p_p1p2_matrix(basis, μ1ω1, μ2ω2; dtype=ComplexF64) ./ (2*μ)
-    else
-        T_cross = spzeros(basis.dim, basis.dim)
-    end
-    
-    verbose || println("Constructing PE matrices")
-    if coord_system == jacobi
-        @time "V" V = get_jacobi_V_matrix(V_of_r, basis, μ1ω1, μω_global; atol=atol, maxevals=maxevals)
-    elseif coord_system == src
-        @time "V" V = get_src_V_matrix(V_of_r, basis, μω, μω_global; atol=atol, maxevals=maxevals)
+    out = spzeros(basis.dim, basis.dim)
+
+    if kinetic_part
+        verbose || println("Constructing KE matrices")
+        @time "T1" out += get_sp_T_matrix(basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; μω_gen=μ1ω1, μ=μ1)
+        @time "T2" out += get_sp_T_matrix(basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; μω_gen=μ2ω2, μ=μ2)
+        if coord_system == src
+            @time "T_cross" out += get_2p_p1p2_matrix(basis, μ1ω1, μ2ω2; dtype=ComplexF64) ./ (2*μ)
+        end
     end
 
-    verbose || println("Adding up all matrices")
-    @time "H" H = T1 + T2 + T_cross + V
+    if potential_part
+        verbose || println("Constructing PE matrices")
+        if coord_system == jacobi
+            @time "V" out += get_jacobi_V_matrix(V_of_r, basis, μ1ω1, μω_global; atol=atol, maxevals=maxevals)
+        elseif coord_system == src
+            @time "V" out += get_src_V_matrix(V_of_r, basis, μω, μω_global; atol=atol, maxevals=maxevals)
+        end
+    end
 
-    return H
+    return out
 end
\ No newline at end of file