diff --git a/calculations/2in3body_HO_B2R_EC.jl b/calculations/2in3body_HO_B2R_EC.jl
index 2e103a3..87a7a52 100644
--- a/calculations/2in3body_HO_B2R_EC.jl
+++ b/calculations/2in3body_HO_B2R_EC.jl
@@ -23,8 +23,8 @@ n_max = max(maximum(basis.n1s), maximum(basis.n2s))
 
 println("Basis size = ", basis.dim)
 
-@time "T1" T1 = get_sp_T_matrix(basis.n1s, basis.l1s; mask=mask1(basis), μω_gen=μ1ω1, μ=μ1)
-@time "T2" T2 = get_sp_T_matrix(basis.n2s, basis.l2s; mask=mask2(basis), μω_gen=μ2ω2, μ=μ2)
+@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)
 
 @time "V" Vb = get_jacobi_V_matrix(V_of_r, basis, μ1ω1, μω_global)
 
diff --git a/calculations/3body_HO_B2R_EC.jl b/calculations/3body_HO_B2R_EC.jl
index 246a221..e2a0c8a 100644
--- a/calculations/3body_HO_B2R_EC.jl
+++ b/calculations/3body_HO_B2R_EC.jl
@@ -24,8 +24,8 @@ n_max = max(maximum(basis.n1s), maximum(basis.n2s))
 
 println("Basis size = ", basis.dim)
 
-@time "T1" T1 = get_sp_T_matrix(basis.n1s, basis.l1s; mask=mask1(basis), μω_gen=μ1ω1, μ=μ1)
-@time "T2" T2 = get_sp_T_matrix(basis.n2s, basis.l2s; mask=mask2(basis), μω_gen=μ2ω2, μ=μ2)
+@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)
 
 @time "Va" Va = get_jacobi_V_matrix(Va_of_r, basis, μ1ω1, μω_global)
 @time "Vb" Vb = get_jacobi_V_matrix(Vb_of_r, basis, μ1ω1, μω_global)
diff --git a/calculations/3body_dis_HO_EC.jl b/calculations/3body_dis_HO_EC.jl
index 90d24c1..bbce32e 100644
--- a/calculations/3body_dis_HO_EC.jl
+++ b/calculations/3body_dis_HO_EC.jl
@@ -26,8 +26,8 @@ n_max = max(maximum(basis.n1s), maximum(basis.n2s))
 
 println("Basis size = ", basis.dim)
 
-@time "T1" T1 = get_sp_T_matrix(basis.n1s, basis.l1s; mask=mask1(basis), μω_gen=μ1ω1, μ=μ1)
-@time "T2" T2 = get_sp_T_matrix(basis.n2s, basis.l2s; mask=mask2(basis), μω_gen=μ2ω2, μ=μ2)
+@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)
 
 @time "Va" Va = get_jacobi_V1_matrix(Va_of_r, basis, μ1ω1)
 @time "Vb" Vb = get_jacobi_V2_matrix(Vb_of_r, basis, μω_global)
diff --git a/ho_basis.jl b/ho_basis.jl
index 9873407..9db92cc 100644
--- a/ho_basis.jl
+++ b/ho_basis.jl
@@ -68,9 +68,6 @@ struct ho_basis_2B
     end
 end
 
-mask1(basis::ho_basis_2B) = (basis.n2s .== basis.n2s') .&& (basis.l2s .== basis.l2s')
-mask2(basis::ho_basis_2B) = (basis.n1s .== basis.n1s') .&& (basis.l1s .== basis.l1s')
-
 function V_numerical(V_of_r, l, n1, n2; μω_gen=1.0, atol=0, maxevals=10^7)
     const_part = sqrt(μω_gen) * ho_basis_const(l, n1) * ho_basis_const(l, n2)
     integrand(r) = ho_basis_func(l, n1, sqrt(μω_gen) * r) * ho_basis_func(l, n2, sqrt(μω_gen) * r) * V_of_r(r)
@@ -78,11 +75,13 @@ function V_numerical(V_of_r, l, n1, n2; μω_gen=1.0, atol=0, maxevals=10^7)
     return const_part * integral
 end
 
-function get_sp_T_matrix(ns, ls; mask=trues(length(ns),length(ns)), μω_gen=1.0, μ=1.0)
+"KE matrix for a single DOF. Set kron_deltas=[other quantum numbers] for other DOFs which the operator does not act on.
+    E.g. get_sp_T_matrix(n1s, l1s, kron_deltas=[n2s l2s])"
+function get_sp_T_matrix(ns, ls, kron_deltas=[]; μω_gen=1.0, μ=1.0)
     mat = spzeros(length(ns), length(ns))
     for idx in CartesianIndices(mat)
-        if !mask[idx]; continue; end
         (i, j) = Tuple(idx)
+        all(arr -> arr[i]==arr[j], kron_deltas) || continue # check if all Kronecker deltas are non-zero
         if ls[i] == ls[j]
             if ns[i] == ns[j]
                 mat[idx] = ns[j] + ls[i]/2 + 3/4
@@ -95,11 +94,13 @@ function get_sp_T_matrix(ns, ls; mask=trues(length(ns),length(ns)), μω_gen=1.0
     return (μω_gen / μ) .* mat
 end
 
-function get_sp_V_matrix(V_l, ns, ls; mask=trues(length(ns),length(ns)), dtype=Float64, cache=LRU{Tuple{UInt8, UInt8, UInt8}, dtype}(maxsize=(1+maximum(ns))^2))
+"PE matrix for a single DOF. Set kron_deltas=[other quantum numbers] for other DOFs which the operator does not act on.
+    E.g. get_sp_V_matrix(n1s, l1s, kron_deltas=[n2s l2s])"
+function get_sp_V_matrix(V_l, ns, ls, kron_deltas=[]; dtype=Float64, cache=LRU{Tuple{UInt8, UInt8, UInt8}, dtype}(maxsize=(1+maximum(ns))^2))
     mat = zeros(dtype, length(ns), length(ns))
     Threads.@threads for idx in CartesianIndices(mat) 
-        if !mask[idx]; continue; end
         (i, j) = Tuple(idx)
+        all(arr -> arr[i]==arr[j], kron_deltas) || continue # check if all Kronecker deltas are non-zero
         if ls[i] == ls[j]
             l = UInt8(ls[i])
             n1, n2 = UInt8.(minmax(ns[i], ns[j])) # assuming transpose symmetry
@@ -160,7 +161,7 @@ function get_jacobi_V1_matrix(V_of_r, basis::ho_basis_2B, μ1ω1; atol=10^-6, ma
     
     V1_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μ1ω1, atol=atol, maxevals=maxevals)
     V1_cache = LRU{Tuple{UInt8, UInt8, UInt8}, ComplexF64}(maxsize=(1+l_max)*(1+n_max)^2)
-    V1 = get_sp_V_matrix(V1_elem, basis.n1s, basis.l1s; mask=mask1(basis), dtype=ComplexF64, cache=V1_cache)
+    V1 = get_sp_V_matrix(V1_elem, basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; dtype=ComplexF64, cache=V1_cache)
 
     return V1
 end
@@ -172,7 +173,7 @@ function get_jacobi_V2_matrix(V_of_r, basis::ho_basis_2B, μω_global; atol=10^-
     V_relative_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μω_global, atol=atol, maxevals=maxevals)
     V_relative_cache = LRU{Tuple{UInt8, UInt8, UInt8}, ComplexF64}(maxsize=(1+l_max)*(1+n_max)^2)
 
-    V_relative = get_sp_V_matrix(V_relative_elem, basis.n1s, basis.l1s; mask=mask1(basis), dtype=ComplexF64, cache=V_relative_cache) + get_sp_V_matrix(V_relative_elem, basis.n2s, basis.l2s; mask=mask2(basis), dtype=ComplexF64, cache=V_relative_cache)
+    V_relative = get_sp_V_matrix(V_relative_elem, basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; dtype=ComplexF64, cache=V_relative_cache) + get_sp_V_matrix(V_relative_elem, basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; dtype=ComplexF64, cache=V_relative_cache)
     U = Moshinsky_transform(basis)
     V2 =  U' * V_relative * U
 
@@ -200,8 +201,8 @@ function get_src_V_matrix(V_of_r, basis::ho_basis_2B, μω, μω_global; atol=10
     V_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μω, atol=atol, maxevals=maxevals)
     V_cache = LRU{Tuple{UInt8, UInt8, UInt8}, ComplexF64}(maxsize=(1+l_max)*(1+n_max)^2)
 
-    V1 = get_sp_V_matrix(V_elem, basis.n1s, basis.l1s; mask=mask1(basis), dtype=ComplexF64, cache=V_cache)
-    V2 = get_sp_V_matrix(V_elem, basis.n2s, basis.l2s; mask=mask2(basis), dtype=ComplexF64, cache=V_cache)
+    V1 = get_sp_V_matrix(V_elem, basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; dtype=ComplexF64, cache=V_cache)
+    V2 = get_sp_V_matrix(V_elem, basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; dtype=ComplexF64, cache=V_cache)
 
     V12 = get_src_V12_matrix(V_of_r, basis, μω_global; atol=atol, maxevals=maxevals)
 
@@ -215,7 +216,7 @@ function get_src_V12_matrix(V_of_r, basis::ho_basis_2B, μω_global; atol=10^-6,
     V_relative_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μω_global, atol=atol, maxevals=maxevals)
     V_relative_cache = LRU{Tuple{UInt8, UInt8, UInt8}, ComplexF64}(maxsize=(1+l_max)*(1+n_max)^2)
 
-    V_relative = get_sp_V_matrix(V_relative_elem, basis.n1s, basis.l1s; mask=mask1(basis), dtype=ComplexF64, cache=V_relative_cache)
+    V_relative = get_sp_V_matrix(V_relative_elem, basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; dtype=ComplexF64, cache=V_relative_cache)
     U = Moshinsky_transform(basis)
     V12 =  U' * V_relative * U
 
diff --git a/ho_basis_3body.jl b/ho_basis_3body.jl
index 9053573..6b13c6c 100644
--- a/ho_basis_3body.jl
+++ b/ho_basis_3body.jl
@@ -20,8 +20,8 @@ println("No of threads = ", Threads.nthreads())
 println("Basis size = ", basis.dim)
 
 println("Constructing KE matrices")
-@time "T1" T1 = get_sp_T_matrix(basis.n1s, basis.l1s; mask=mask1(basis), μω_gen=μω, μ=μ)
-@time "T2" T2 = get_sp_T_matrix(basis.n2s, basis.l2s; mask=mask2(basis), μω_gen=μω, μ=μ)
+@time "T1" T1 = get_sp_T_matrix(basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; μω_gen=μω, μ=μ)
+@time "T2" T2 = get_sp_T_matrix(basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; μω_gen=μω, μ=μ)
 @time "T_cross" T_cross = get_2p_p1p2_matrix(basis, μω, μω) ./ (2*μ)
 
 println("Constructing PE matrices")
diff --git a/ho_basis_3body_resonance.jl b/ho_basis_3body_resonance.jl
index 1c775c2..998621a 100644
--- a/ho_basis_3body_resonance.jl
+++ b/ho_basis_3body_resonance.jl
@@ -20,8 +20,8 @@ println("Basis size = ", basis.dim)
 
 println("Constructing KE matrices")
 
-@time "T1" T1 = get_sp_T_matrix(basis.n1s, basis.l1s; mask=mask1(basis), μω_gen=μω, μ=μ)
-@time "T2" T2 = get_sp_T_matrix(basis.n2s, basis.l2s; mask=mask2(basis), μω_gen=μω, μ=μ)
+@time "T1" T1 = get_sp_T_matrix(basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; μω_gen=μω, μ=μ)
+@time "T2" T2 = get_sp_T_matrix(basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; μω_gen=μω, μ=μ)
 @time "T_cross" T_cross = get_2p_p1p2_matrix(basis, μω, μω; dtype=ComplexF64) ./ (2*μ)
 
 println("Constructing PE matrices")