diff --git a/src/io.f90 b/src/io.f90
index f4928e4..4aab723 100644
--- a/src/io.f90
+++ b/src/io.f90
@@ -373,7 +373,8 @@ module io
         !NOTE: This code doesn't work for arbitrary number of basis atoms per node. It assumes that the 
         !each element has only 1 atom type at the node.
         character(len=100), intent(in) :: file
-        integer :: interp_max, i, j, inod, ibasis, ip, unique_index(10), unique_num
+        integer :: interp_max, i, j, inod, ibasis, ip, unique_index(50), unique_size(50), unique_num, &
+                   etype
         real(kind=dp) :: box_vec(3)
 
 1  format('time' / i16, f23.15)
@@ -405,6 +406,22 @@ module io
 
         !Below writes the header information for the restart file
 
+
+            !First figure out all of the unique element types
+            unique_num = 0
+            unique_index(:) = 0
+            eleloop:do i = 1, ele_num
+                do j =1 , unique_num
+                    if ( ( size_ele(i) == size_ele( unique_index(j) ) ).and. &
+                        ( lat_ele(i) == lat_ele(unique_index(j)) ) ) then
+                        cycle eleloop
+                    end if
+                end do
+                unique_num = unique_num + 1
+                unique_index(unique_num) = i
+                unique_size(unique_num) = size_ele(i)
+            end do eleloop
+
         !Calculate the max number of atoms per element
         select case(max_ng_node)
         case(8)
@@ -415,7 +432,7 @@ module io
         write(11,20) interp_max
         write(11,3) node_num
         write(11,19) max_ng_node
-        write(11,4) lattice_types
+        write(11,4) unique_num
         write(11,5) atom_num
         write(11,6) 1 !Grain_num is ignored 
         write(11,16) lattice_types, atom_types
@@ -454,26 +471,21 @@ module io
         !write the element information
         if(ele_num > 0) then
             write(11,12)
-            !First figure out all of the unique element types
-            unique_num = 0
-            unique_index(:) = 0
-            eleloop:do i = 1, ele_num
-                do j =1 , unique_num
-                    if ( ( size_ele(i) == size_ele( unique_index(j) ) ).and. &
-                        ( lat_ele(i) == lat_ele(unique_index(j)) ) ) then
-                        cycle eleloop
-                    end if
-                end do
-                unique_num = unique_num + 1
-                unique_index(unique_num) = i
-            end do eleloop
+           
             do i = 1, unique_num
-                write(11,'(3i16)') i, size_ele(i)-1, basis_type(1,i)
+                write(11,'(3i16)') i, size_ele(unique_index(i))-1, basis_type(1,lat_ele(unique_index(i)))
             end do
             ip = 0
             write(11,13)
             do i = 1, ele_num
-                write(11, '(4i16)') i, lat_ele(i), 1, basis_type(1,lat_ele(i))
+                !Figure out the ele type
+                do j = 1, unique_num
+                    if ( unique_size(j) == size_ele(i)) then
+                        etype = j
+                        exit
+                    endif
+                end do
+                write(11, '(4i16)') i, etype, 1, basis_type(1,lat_ele(i))
                 do inod = 1, ng_node(lat_ele(i))
                     do ibasis = 1, basisnum(lat_ele(i))
                         ip = ip + 1
@@ -633,15 +645,16 @@ module io
 
         print *, "displace", displace
         do i = 1, 3
-            if (displace(i) > lim_zero) then 
+            if (abs(displace(i)) > lim_zero) then 
                 newdisplace(i) = displace(i) - temp_box_bd(2*i-1)
             else
-                newdisplace=displace(i)
+                newdisplace(i)=displace(i)
             end if 
             temp_box_bd(2*i-1) = temp_box_bd(2*i-1) + newdisplace(i)
             temp_box_bd(2*i) = temp_box_bd(2*i) + newdisplace(i)
         end do 
-        
+       
+        print *, "newdisplace", newdisplace 
         call grow_box(temp_box_bd)
         !Read in the number of sub_boxes and allocate the variables
         read(11, *) n