diff --git a/Hamiltonian.jl b/Hamiltonian.jl
index 83d232c..19aabfa 100644
--- a/Hamiltonian.jl
+++ b/Hamiltonian.jl
@@ -9,9 +9,7 @@ struct Hamiltonian{T}
     K_partial::Matrix{Complex{T}}
     K_diag::Union{CuTensor{Complex{T}},Nothing}
     K_mixed::Union{CuTensor{Complex{T}},Nothing}
-    K_partial_1x::Union{Matrix{Complex{T}},Nothing}
-    K_partial_1y::Union{Matrix{Complex{T}},Nothing}
-    K_partial_1z::Union{Matrix{Complex{T}},Nothing}
+    K_partial_1::Union{Tuple,Nothing}
     Vs::Union{Array{Complex{T}},CuArray{Complex{T}}}
     hermitian::Bool
     mode::Hamiltonian_backend
@@ -21,7 +19,7 @@ struct Hamiltonian{T}
         k = -s.N÷2:s.N÷2-1
         hermitian = ϕ == 0.0
         K_partial = (exp(-im * convert(T, ϕ)) * im / sqrt(2 * s.μ)) .* ∂_1DOF.(Ref(s), k, k')
-        K_partial_1x, K_partial_1y, K_partial_1z = sym_reduce(s, K_partial)
+        K_partial_1 = sym_reduce(s, K_partial)
         Vs = calculate_Vs(s, V_twobody, convert(T, ϕ), n_image)
         if mode == gpu_cutensor
             K_diag = CuTensor(CuArray(K_partial * K_partial), ['a', 'A'])
@@ -31,7 +29,7 @@ struct Hamiltonian{T}
             K_diag = nothing
             K_mixed = nothing    
         end
-        return new{T}(s, K_partial, K_diag, K_mixed, K_partial_1x, K_partial_1y, K_partial_1z, Vs, hermitian, mode)
+        return new{T}(s, K_partial, K_diag, K_mixed, K_partial_1, Vs, hermitian, mode)
     end
 end
 
@@ -39,7 +37,7 @@ Base.size(H::Hamiltonian, i::Int)::Int = (i == 1 || i == 2) ? H.s.N^(H.s.d * (H.
 Base.size(H::Hamiltonian)::Dims{2} = (size(H, 1), size(H, 2))
 
 "Dimensions of a vector to which 'H' can be applied"
-vectorDims(H::Hamiltonian)::Dims = tuple(H.dim1, fill(H.s.N, H.s.d * (H.s.n - 2))...)
+vectorDims(H::Hamiltonian)::Dims = tuple(length(H.s.unique_i), fill(H.s.N, H.s.d * (H.s.n - 2))...)
 
 "Apply 'H' on 'v' and store the result in 'out' using the 'cpu_tensor' backend"
 function LinearAlgebra.mul!(out::Array{Complex{T}}, H::Hamiltonian{T}, v::Array{Complex{T}})::Array{Complex{T}} where {T<:Float}
@@ -47,28 +45,26 @@ function LinearAlgebra.mul!(out::Array{Complex{T}}, H::Hamiltonian{T}, v::Array{
     # apply V operator
     @. out = H.Vs * v
     # apply K opereator
-    coords = H.s.n - 1
-    nconList_v_template = -collect(1:H.s.d*(coords))
+    nconList_v_template = -collect(1:(H.s.d * (H.s.n - 2) + 1))
     for dim = 1:H.s.d
-        for coord1 = 1:coords
+        for coord1 = 1:(H.s.n - 1)
             for coord2 = 1:coord1
-                    
-                if coord1 == 1 && coord2 == 1
-                
+                i1 = which_index(H.s, dim, coord1)
+                i2 = which_index(H.s, dim, coord2)
+                nconList_1 = [-i1, 1]
+                nconList_2 = [-i2, 2]
+                nconList_v = copy(nconList_v_template)
+                if i1 == i2
+                    nconList_2[1] = 1
                 else
-                    i1 = which_index(H.s, dim, coord1)
-                    i2 = which_index(H.s, dim, coord2)
-                    nconList_1 = [-i1, 1]
-                    nconList_2 = [-i2, 2]
-                    nconList_v = copy(nconList_v_template)
-                    if i1 == i2
-                        nconList_2[1] = 1
-                    else
-                        nconList_v[i1] = 1
-                    end
-                    nconList_v[i2] = 2
-                    v_new = @ncon((H.K_partial, H.K_partial, v), (nconList_1, nconList_2, nconList_v))
+                    nconList_v[i1] = 1
                 end
+                nconList_v[i2] = 2
+
+                tensor1 = coord1 == 1 ? H.K_partial_1[dim] : H.K_partial
+                tensor2 = coord2 == 1 ? H.K_partial_1[dim] : H.K_partial
+                v_new = @ncon((tensor1, tensor2, v), (nconList_1, nconList_2, nconList_v))
+                
                 out = axpy!(1, v_new, out)
             end
         end
diff --git a/common.jl b/common.jl
index 9380853..672551f 100644
--- a/common.jl
+++ b/common.jl
@@ -41,7 +41,7 @@ function ∂_1DOF(s::system{T}, k::Int, l::Int)::Complex{T} where {T<:Float}
 end
 
 "Which index (dimension of the multidimensional array) corresponds to spatial dimension 'dim' and particle 'p'?"
-which_index(s::system, dim::Int, p::Int)::Int = (dim - 1) * (s.n - 2) + p + 1
+which_index(s::system, dim::Int, p::Int)::Int = p == 1 ? 1 : (dim - 1) * (s.n - 2) + p + 1
 
 "Δk (distance in terms of lattice paramter) between two particles along the given dimension"
 function get_k(s::system, i::CartesianIndex, dim::Int, p::Int)::Int
@@ -90,7 +90,7 @@ function calculate_Vs(s::system{T}, V_twobody::Function, ϕ::T, n_image::Int)::A
                 end
             end
         end
-        Vs[i] *= s.multiplicity[i[1]]^2
+        Vs[i] *= s.multiplicity[i[1]]
     end
     return Vs
 end
diff --git a/testing-irrep.ipynb b/testing-irrep.ipynb
index 5fcb587..90112e0 100644
--- a/testing-irrep.ipynb
+++ b/testing-irrep.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -21,11 +21,13 @@
     "println(\"\\n$n-body system with N=$N\")\n",
     "\n",
     "for L::T in 5.0:9.0\n",
-    "    print(\"L=$L\")\n",
+    "    println(\"L=$L\")\n",
+    "    println(\"Constructing Hamiltonian\")\n",
     "    s=system{T}(3,n,N,L,0.5,all)\n",
     "    @time H=Hamiltonian{T}(s,V_test,0,0,cpu_tensor)\n",
-    "    #evals,_,_ = eig(H,5)\n",
-    "    #println(real.(evals))\n",
+    "    println(\"Solving eigenvalues\")\n",
+    "    @time evals,_,_ = eig(H,5)\n",
+    "    println(real.(evals))\n",
     "end"
    ]
   }