-- Classical RK4
-- July 8, 2009
-- Steve Kifowit

-- User input
dim = 2
x = 0
y = {1, 1}
steps = 40
Xend = 4
function func( x, y, z )
   z[1] = y[2]
   z[2] = -y[2] - 4.0*y[1]
   return
end
-- No user input required beyond this point, --
-- unless you want to edit the format strings --

-- Initialize variables, arrays, and format strings
h = ( Xend - x ) / steps
w, z = {}, {}
k1, k2, k3, k4 = {}, {}, {}, {}
form1 = "x = %.4f  "
form2 = "y%d = %.6f  "

-- Main program --

-- Print starting values
print( "Starting values:" )
io.write( string.format( form1, x ) )
for i = 1, dim do
   io.write( string.format( form2, i, y[i] ) )
end
print()

-- Also print to file
file = io.open("rk4sys.dat", "w+")
file:write( x, "  " )
for i = 1, dim do
   file:write( y[i], "  " )
end
file:write( "\n" )

-- Main loop - Classical RK4
for k = 1, steps do
   func( x, y, z )
   for i = 1, dim do
      k1[i] = h * z[i]
      w[i] = y[i] + k1[i]/2 
   end
   func( x+h/2, w, z )
   for i = 1, dim do
      k2[i] = h * z[i]
      w[i] = y[i] + k2[i]/2 
   end
   func( x+h/2, w, z )
   for i = 1, dim do
      k3[i] = h * z[i]
      w[i] = y[i] + k3[i] 
   end
   func( x+h, w, z )
   for i = 1, dim do
      k4[i] = h * z[i]
      y[i] = y[i] + ( k1[i]+2*k2[i]+2*k3[i]+k4[i] ) / 6 
   end
   x = x + h
-- Print to file
   file:write( x, "  " )
   for i = 1, dim do
      file:write( y[i], "  " )
   end
   file:write( "\n" )
end

--Close data file
file:close()

-- Print values at Xend
print("\nEnding values:")
io.write( string.format( form1, x ) )
for i = 1, dim do
   io.write( string.format( form2, i, y[i] ) )
end
print("\n\nData from individual steps written to rk4sys.dat")