diff --git a/berggren.jl b/berggren.jl
new file mode 100644
index 0000000..5bb4dfc
--- /dev/null
+++ b/berggren.jl
@@ -0,0 +1,18 @@
+using SparseArrays
+include("ho_basis.jl")
+
+"Basis transformation from HO to momentum space"
+function get_W_matrix(basis_p, basis_HO, μ1ω1, μ2ω2=μ1ω1; weights=true)
+    Es, n1s, l1s, n2s, l2s = basis_HO
+    W = zeros(ComplexF64, length(basis), length(Es))
+    Threads.@threads for idx in CartesianIndices(W)
+        (i1, i2) = Tuple(idx)
+        ((k1, w1), (k2, w2), (j1, j2)) = basis_p[i1]
+        if j1 == l1s[i2] && j2 == l2s[i2]
+            elem1 = 1/sqrt(sqrt(μ1ω1)) * (-1)^n1s[i2] * ho_basis(j1, n1s[i2], 1/sqrt(μ1ω1) * k1)
+            elem2 = 1/sqrt(sqrt(μ2ω2)) * (-1)^n2s[i2] * ho_basis(j2, n2s[i2], 1/sqrt(μ2ω2) * k2)
+            W[idx] = elem1 * elem2 * (weights ? w1 * w2 : 1)
+        end
+    end
+    return sparse(W)
+end
diff --git a/berggren_3body.jl b/berggren_3body.jl
index df1394b..5c74246 100644
--- a/berggren_3body.jl
+++ b/berggren_3body.jl
@@ -1,7 +1,7 @@
 using LinearAlgebra, SparseArrays, Arpack
 include("helper.jl")
 include("p_space.jl")
-include("ho_basis.jl")
+include("berggren.jl")
 
 println("No of threads = ", Threads.nthreads())
 
@@ -41,25 +41,6 @@ end
     V1 = blockdiag(sparse.(V1_blocks)...)
 end
 
-#############################################################################################
-
-function get_W_matrix(basis, basis_HO, μ1ω1, μ2ω2=μ1ω1; weights=true)
-    Es, n1s, l1s, n2s, l2s = basis_HO
-    W = zeros(ComplexF64, length(basis), length(Es))
-    Threads.@threads for idx in CartesianIndices(W)
-        (i1, i2) = Tuple(idx)
-        ((k1, w1), (k2, w2), (j1, j2)) = basis[i1]
-        if j1 == l1s[i2] && j2 == l2s[i2]
-            elem1 = 1/sqrt(sqrt(μ1ω1)) * (-1)^n1s[i2] * ho_basis(j1, n1s[i2], 1/sqrt(μ1ω1) * k1)
-            elem2 = 1/sqrt(sqrt(μ2ω2)) * (-1)^n2s[i2] * ho_basis(j2, n2s[i2], 1/sqrt(μ2ω2) * k2)
-            W[idx] = elem1 * elem2 * (weights ? w1 * w2 : 1)
-        end
-    end
-    return sparse(W)
-end
-
-############################################################################################
-
 E_max = 30
 μω_global = 0.5