finish hw3 code and start hw4. add stencil code via upload

1 files changed, 179 insertions, 0 deletions

diff --git a/hw4/SolarSystem2.jl b/hw4/SolarSystem2.jl
new file mode 100644
index 0000000..b2e02b7
--- /dev/null
+++ b/hw4/SolarSystem2.jl
@@ -0,0 +1,179 @@
+#!/Applications/Julia-1.8.app/Contents/Resources/julia/bin/julia
+
+# Simulate a solar system
+
+using Plots # for plotting trajectory
+using DifferentialEquations # for solving ODEs
+using LinearAlgebra # for dot and cross products
+
+G = 4.0*pi^2 # time scale = year and length scale = AU
+
+mutable struct body
+   name::String # name of star or planet
+   m::Float64 # mass
+   r::Vector{Float64} # position vector
+   p::Vector{Float64} # momentum vector
+end 
+
+function angularMomentum(b::body)
+   r = b.r
+   p = b.p
+   return cross(r, p)
+end
+
+function kineticEnergy(b::body)
+   p = b.p
+   m = b.m
+   return dot(p, p) / (2.0 * m)
+end
+
+function rp(b::body)
+   return [b.r; b.p]
+end
+
+function force(body1::body, body2::body)
+   r = body1.r - body2.r
+   rSquared = dot(r, r)
+   return -G * body1.m * body2.m * r * rSquared^(-1.5)
+end
+
+function potentialEnergy(body1::body, body2::body)
+   r = body1.r - body2.r
+   rSquared = dot(r, r)
+   return -G * body1.m * body2.m * rSquared^(-0.5)
+end
+
+mutable struct SolarSystem
+   bodies::Vector{body}
+   numberOfBodies::Int64
+   phaseSpace::Matrix{Float64} # 6N dimensional phase space
+end
+
+function SolarSystem()
+
+   bodies = Vector{body}()
+   push!(bodies, body("Sun", 1.0, zeros(3), zeros(3)))
+   #push!(bodies, body("Earth", 1.0, [-1.0, 0.0, 0.0], [0.0, 1.0 * pi, 0.0]))
+   push!(bodies, body("Star", 1.0, [1.0, 0.0, 0.0], [0.0, sqrt(1.95) * 2.0 * pi, 0.0]))
+   #push!(bodies, body("Jupiter", 1.0, [3.0, 0.0, 0.0], [0.0, 0.25 * pi, 0.0]))
+   numberOfBodies = size(bodies)[1]
+
+   phaseSpace = zeros(6, 0)
+   for b in bodies
+      phaseSpace = [phaseSpace rp(b)]
+   end
+   
+   return SolarSystem(bodies, numberOfBodies, phaseSpace)
+
+end
+
+function TotalAngularMomentum(s::SolarSystem)
+   L = zeros(3)
+   for b in s.bodies
+      L += angularMomentum(b)
+   end
+
+   return L
+end
+
+function TotalEnergy(s::SolarSystem)
+   ke = 0.0
+   pe = 0.0
+
+   for body1 in s.bodies
+      ke += kineticEnergy(body1)
+      for body2 in s.bodies
+         if (body1 != body2)
+            pe += 0.5 * potentialEnergy(body1, body2)
+         end
+      end
+   end
+
+   return pe + ke
+end
+
+function ZeroOutLinearMomentum!(s::SolarSystem)
+
+   totalLinearMomentum = zeros(3)
+   totalMass = 0.0
+   for body in s.bodies
+      totalLinearMomentum += body.p
+      totalMass += body.m
+   end
+   
+   s.phaseSpace = zeros(6, 0)
+   for body in s.bodies
+      body.p -= body.m * totalLinearMomentum / totalMass
+      s.phaseSpace = [s.phaseSpace rp(body)]
+   end
+
+   return nothing
+end
+
+function tendency!(dps, ps, s::SolarSystem, t)
+
+   i = 1 # update phase space of individual bodies
+   for b in s.bodies
+      b.r = ps[1:3, i]
+      b.p = ps[4:6, i]
+      i += 1
+   end
+
+   # find velocities of bodies and forces on them. O(N^2) computational cost
+   N = s.numberOfBodies
+   for i in 1:N
+      b1 = s.bodies[i]
+      dps[1:3, i] = b1.p / b1.m  #dr/dt
+      dps[4:6, i] = zeros(3)
+      for j in 1:i-1
+            b2 = s.bodies[j]
+            f = force(b1, b2) # call only once
+            dps[4:6, i] += f
+            dps[4:6, j] -= f # Newton's 3rd law
+      end
+   end
+
+   return nothing
+end
+
+s = SolarSystem()
+ZeroOutLinearMomentum!(s)
+println(typeof(s))
+println("Initial total energy = ", TotalEnergy(s))
+println("Initial total angular momentum = ", TotalAngularMomentum(s))
+println("Number of bodies = ", s.numberOfBodies)
+for b in s.bodies
+   println("body name = ", b.name)
+end
+
+t_final = 10.0 # final time of simulation
+tspan = (0.0, t_final) # span of time to simulate
+prob = ODEProblem(tendency!, s.phaseSpace, tspan, s) # specify ODE
+sol = solve(prob, Tsit5(), reltol=1e-12, abstol=1e-12) # solve using Tsit5 algorithm to specified accuracy
+
+sample_times = sol.t
+println("\n\t Results")
+println("Final time  = ", sample_times[end])
+println("Final total energy = ", TotalEnergy(s))
+println("Final total angular momentum  = ", TotalAngularMomentum(s))
+
+body1 = 1
+body2 = 2
+# Plot of position vs. time
+xt = plot(sample_times, sol[1, body1, :], xlabel = "t", ylabel = "x(t)", legend = false, title = "x vs. t")
+
+# Plot of orbit
+xy = plot([(sol[1, body1, :], sol[2, body1, :]), (sol[1, body2, :], sol[2, body2, :])], colors = ("yellow", "green"), xlabel = "x", ylabel = "y", legend = false, title = "Orbit")
+#=
+plot(xy, aspect_ratio=:equal)
+=#
+
+samples = 1000
+interval = floor(Int,size(sol.t)[1] / samples)
+animation = @animate for i=1:samples-1
+   plot([(sol[1, body1, i*interval], sol[2, body1, i*interval]), (sol[1, body1, (i+1)*interval], sol[2, body1, (i+1)*interval])],
+   aspect_ratio=:equal, colors = ("yellow", "green"), xlabel = "x", ylabel = "y", legend = false, title = "Orbit", xlims=(-1, 1), ylims = (-1, 1))
+end
+
+gif(animation, fps=15)
+