finish hw4

1 files changed, 287 insertions, 0 deletions

diff --git a/hw4/4-17.jl b/hw4/4-17.jl
new file mode 100644
index 0000000..eb2b3b3
--- /dev/null
+++ b/hw4/4-17.jl
@@ -0,0 +1,287 @@
+#!/Applications/Julia-1.8.app/Contents/Resources/julia/bin/julia
+# FOR PROBLEM 4.17
+# author: sotech117
+
+# Simulate a solar system
+
+using Plots # for plotting trajectory
+using DifferentialEquations # for solving ODEs
+using LinearAlgebra # for dot and cross products
+
+println("Number of threads = ", Threads.nthreads())
+
+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
+   v::Vector{Float64} # velocity vector
+end 
+
+function star(name="Sun", m = 1.0, r = zeros(3), v = zeros(3))
+   return body(name, m, r, v)
+end
+
+function planet(name="Earth", m=3.0e-6, a=1.0, ϵ=0.017, i=0.0)
+   perihelion = (1.0 - ϵ) * a
+	aphelion = (1.0 + ϵ) * a
+	speed = sqrt(G * (1.0 + ϵ)^2 / (a * (1.0 - ϵ^2)))
+   phi = 0.0
+	r = [perihelion * cos(i*pi/180.0) * cos(phi), perihelion * cos(i*pi/180.0) * sin(phi), perihelion * sin(i*pi/180.0)]
+	v = [-speed * sin(phi), speed * cos(phi), 0.0]
+   return body(name, m, r, v)
+end
+
+function momentum(b::body)
+   return b.m * b.v
+end
+
+function angularMomentum(b::body)
+   return b.m * cross(b.r, b.v)
+end
+
+function kineticEnergy(b::body)
+   v = b.v
+   m = b.m
+   return 0.5 * m * dot(v, v)
+end
+
+function potentialEnergy(body1::body, body2::body)
+   r = body1.r - body2.r
+   rSquared = dot(r, r)
+   return -G * body1.m * body2.m / sqrt(rSquared)
+end
+
+function rv(b::body)
+   return [b.r; b.v]
+end
+
+function force(body1::body, body2::body)
+   r = body1.r - body2.r
+   rSquared = dot(r, r)
+   return -G * r * rSquared^(-1.5)
+   #return -G * r / (rSquared * sqrt(rSquared))
+end
+
+mutable struct SolarSystem
+   bodies::Vector{body}
+   numberOfBodies::Int64
+   phaseSpace::Matrix{Float64} # 6N dimensional phase space
+end
+
+function SolarSystem()
+
+   bodies = Vector{body}()
+   push!(bodies, star()) # default = Sun
+   # push!(bodies, planet("Venus", 2.44e-6, 0.72, 0.0068, 3.39))
+   # push!(bodies, planet("Earth", 3.0e-6, 1.0, 0.017, 0.0))
+   push!(bodies, planet("Jupiter", 0.00095, 5.2, 0.049, 1.3))
+   # push!(bodies, planet("Saturn", 0.000285, 9.58, 0.055, 2.48))
+   push!(bodies, body("Astroid 1", 1e-13, [3.0, 0.0, 0.0], [0.0, 3.628, 0.0]))
+   push!(bodies, body("Astroid 2 ", 1e-13, [3.276, 0.0, 0.0], [0.0, 3.471, 0.0]))
+   # push!(bodies, body("Astroid 3 (no gap)", 1e-13,[3.7, 0.0, 0.0], [0.0, 3.267, 0.0]))
+   numberOfBodies = size(bodies)[1]
+
+   phaseSpace = zeros(6, 0)
+   for b in bodies
+      phaseSpace = [phaseSpace rv(b)]
+   end
+   
+   return SolarSystem(bodies, numberOfBodies, phaseSpace)
+
+end
+
+function TotalMass(s::SolarSystem)
+   M = 0.0
+   for b in s.bodies
+      M += b.m
+   end
+   return M
+end
+
+function TotalLinearMomentum(s::SolarSystem)
+   P = zeros(3)
+   for b in s.bodies
+      P += momentum(b)
+   end
+   return P
+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 CenterOfMassFrame!(s::SolarSystem)
+
+   M = TotalMass(s)
+   P = TotalLinearMomentum(s)
+   V = P / M
+   
+   s.phaseSpace = zeros(6, 0)
+   for body in s.bodies
+      body.v -= V #boost to COM frame
+      s.phaseSpace = [s.phaseSpace rv(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.v = 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.v  #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] += b2.m * f
+            dps[4:6, j] -= b1.m * f # Newton's 3rd law
+      end
+   end
+
+   return nothing
+end
+
+function parallel_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.v = ps[4:6, i]
+      i += 1
+   end
+
+   # find velocities of bodies and forces on them. O(N^2) computational cost
+   N = s.numberOfBodies
+   Threads.@threads for i in 1:N
+      b1 = s.bodies[i]
+      dps[1:3, i] = b1.v  #dr/dt
+      dps[4:6, i] = zeros(3)
+      for j in 1:N
+            if (j != i)
+               b2 = s.bodies[j]
+               f = force(b1, b2)
+               dps[4:6, i] += b2.m * f
+            end
+      end
+   end
+
+   return nothing
+end
+
+s = SolarSystem()
+CenterOfMassFrame!(s)
+println(typeof(s))
+println("Initial total energy = ", TotalEnergy(s))
+println("Initial total linear momentum = ", TotalLinearMomentum(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 = 200.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, maxiters=1e8, reltol=1e-10, abstol=1e-10) # solve using Tsit5 algorithm to specified accuracy
+
+println("\n\t Results")
+println("Final time  = ", sol.t[end])
+println("Final total energy = ", TotalEnergy(s))
+println("Final total linear momentum = ", TotalLinearMomentum(s))
+println("Final total angular momentum  = ", TotalAngularMomentum(s))
+
+function eccentricity(body::Int64, solution)
+   n = size(solution[1, 1, :])[1]
+   samples = 1000
+   interval = floor(Int, n/samples)
+   ϵ = zeros(samples-1)
+   time = zeros(samples-1)
+   for i in 1:samples-1
+      r2 = sol[1, body, i*interval:(i+1)*interval].^2 + sol[2, body, i*interval:(i+1)*interval].^2 + sol[3, body, i*interval:(i+1)*interval].^2
+      aphelion = sqrt(maximum(r2))
+      perihelion = sqrt(minimum(r2))
+      ϵ[i] = (aphelion - perihelion) / (aphelion + perihelion)
+      time[i] = solution.t[i*interval]
+   end
+   return (time, ϵ)
+end
+
+function obliquity(body::Int64, solution)
+   n = size(solution[1, 1, :])[1]
+   samples = 1000
+   interval = floor(Int, n/samples)
+   ob = zeros(samples)
+   time = zeros(samples)
+   for i in 1:samples
+      r = solution[1:3, body, i*interval]
+      v = solution[4:6, body, i*interval]
+      ell = cross(r, v)
+      norm = sqrt(dot(ell, ell))
+      ob[i] = (180.0 / pi) * acos(ell[3]/norm)
+      time[i] = solution.t[i*interval]
+   end
+   return (time, ob)
+end
+
+
+body1 = 2 # Jupiter
+body2 = 3 # Astroid 1
+body3 = 4 # Astroid 2
+# body4 = 5 # Astroid 3
+
+# Plot of orbit
+xy = scatter([(sol[1, body1, :], sol[2, body1, :]), (sol[1, body2, :], sol[2, body2, :]),
+ (sol[1, body3, :], sol[2, body3, :]), 
+ # (sol[1, body4, :], sol[2, body4, :])
+ ], 
+ xlabel = "x (AU)", ylabel = "y (AU)", title = "Orbit",
+ label = ["Jupiter" "Asteroid not in gap (3.0AU)" "Asteroid 2/1 gap (3.27AU)" "Asteroid not in gap (3.7AU)"],
+ aspect_ratio=:equal, markersize=.5, markerstrokewidth = 0, legend = :bottomright,
+ colors=[:red :blue :green :black]
+ )
+
+# Plot of obliquity
+tilt = obliquity(body2, sol)
+obliquityPlot = plot(tilt, xlabel = "t", ylabel = "tilt", legend = false, title = "obliquity")
+
+# Plot of eccentricity
+eccentric = eccentricity(body2, sol)
+eccentricityPlot = plot(eccentric, xlabel = "t", ylabel = "ϵ", legend = false, title = "eccentricity")
+
+#plot(obliquityPlot, eccentricityPlot)
+plot(xy, aspect_ratio=:equal)
+
+savefig("4-17.png")