Merge pull request #84 from ericagol/langfzac/ics

Fix up initial conditions

docs/src/basic.md

@@ -24,21 +24,21 @@ b = Elements(
     m = 3e-6,     # Mass [Solar masses]
     t0 = 0.0,     # Initial time of transit [days]
     P = 365.256,  # Period [days]
-    ecosϖ = 0.01, # Eccentricity * cos(Argument of Periastron)
-    esinϖ = 0.0,  # Eccentricity * sin(Argument of Periastron)
+    ecosω = 0.01, # Eccentricity * cos(Argument of Periastron)
+    esinω = 0.0,  # Eccentricity * sin(Argument of Periastron)
     I = π/2,      # Inclination [Radians]
     Ω = 0.0       # Longitude of Ascending Node [Radians]
 );
 nothing # hide
 ```
-(`ecosϖ` can be typed as `ecos\varpi` and then hitting tab)
+(`ecosω` can be typed as `ecos\omega` and then hitting tab)
 
 Next we'll create a Jupiter analogue for the final body. Here the orbital elements are specified for the Keplerian ((a,b),c), or c orbiting the center of mass of a and b. (While this might not need to be stated explicitly here, this convention is useful for more complicated hierarchical systems).
 ```@example 1
 c = Elements(
     m = 9.54e-4,
     P = 4332.59,
-    ecosϖ = 0.05,
+    ecosω = 0.05,
     I = π/2
 );
 nothing # hide

docs/src/gradients.md

@@ -3,13 +3,16 @@ The main purpose of developing NbodyGradient.jl is to provide a differentiable N
 
 We assume you've taken a look at the [Basic Usage](@ref) tutorial, and are familiar with the orbital elements and units used in NbodyGradient.jl.
 
+!!! note "No gradient support for exactly circular orbits"
+    This package supports specifying initial conditions for circular orbits (ie. eccentricity=0). However, the derivative computations contain 1/e terms -- requiring a non-zero eccentricity to be computed correctly. We expect the need for derivatives for exactly circular initial orbits to be minimal, and intend to implement them in the future.
+
 Here, we will specify the orbital elements using the 'elements matrix' option (See [`ElementsIC`](@ref)).
 ```@example 2
 using NbodyGradient
 
 # Initial Conditions
 elements = [
-  # m    P       t0  ecosϖ esinϖ I   Ω
+  # m    P       t0  ecosω esinω I   Ω
     1.0  0.0     0.0 0.0   0.0   0.0 0.0; # Star
     3e-6 365.256 0.0 0.01  0.0   π/2 0.0; # Earth analogue
 ]

src/ics/InitialConditions.jl

@@ -9,53 +9,119 @@ Orbital elements of a binary, and mass of a 'outer' body. See [Tutorials](@ref)
 - `m::T` : Mass of outer body.
 - `P::T` : Period [Days].
 - `t0::T` : Initial time of transit [Days].
-- `ecosϖ` : Eccentricity vector x-component (eccentricity times cosine of the longitude of periastron)
-- `esinϖ` : Eccentricity vector y-component (eccentricity times sine of the longitude of periastron)
+- `ecosω::T` : Eccentricity vector x-component (eccentricity times cosine of the argument of periastron)
+- `esinω::T` : Eccentricity vector y-component (eccentricity times sine of the argument of periastron)
 - `I::T` : Inclination, as measured from sky-plane [Radians].
 - `Ω::T` : Longitude of ascending node, as measured from +x-axis [Radians].
 - `a::T` : Orbital semi-major axis [AU].
 - `e::T` : Eccentricity.
-- `ϖ::T` : Longitude of periastron [Radians].
+- `ω::T` : Argument of periastron [Radians].
+- `tp::T` : Time of periastron passage [Days].
 """
 struct Elements{T<:AbstractFloat} <: AbstractInitialConditions
     m::T
     P::T
     t0::T
-    ecosϖ::T
-    esinϖ::T
+    ecosω::T
+    esinω::T
     I::T
     Ω::T
     a::T
     e::T
-    ϖ::T
+    ω::T
+    tp::T
 end
 
 """
-    Elements(m,P,t0,ecosϖ,esinϖ,I,Ω)
+    Elements(m,P,t0,ecosω,esinω,I,Ω)
 
 Main [`Elements`](@ref) constructor. May use keyword arguments, see [Tutorials](@ref).
 """
-function Elements(m::T,P::T,t0::T,ecosϖ::T,esinϖ::T,I::T,Ω::T) where T<:Real
-    e = sqrt(ecosϖ^2 + esinϖ^2)
-    ϖ = atan(esinϖ,ecosϖ)
-    Elements(m,P,t0,ecosϖ,esinϖ,I,Ω,0.0,e,ϖ)
+function Elements(m::T,P::T,t0::T,ecosω::T,esinω::T,I::T,Ω::T) where T<:Real
+    e = sqrt(ecosω^2 + esinω^2)
+    ω = atan(esinω,ecosω)
+    Elements(m,P,t0,ecosω,esinω,I,Ω,0.0,e,ω,0.0)
 end
 
-function Base.show(io::IO, ::MIME"text/plain" ,elems::Elements{T}) where T <: Real
+function Base.show(io::IO, ::MIME"text/plain", elems::Elements{T}) where T <: Real
     fields = fieldnames(typeof(elems))
-    vals = Dict([fn => getfield(elems,fn) for fn in fields])
+    vals = [fn => getfield(elems,fn) for fn in fields]
     println(io, "Elements{$T}")
-    for key in keys(vals)
-        println(io,key,": ",vals[key])
+    for pair in vals
+        println(io,first(pair),": ",last(pair))
     end
     if elems.a == 0.0
-        println(io, "Semi-major axis: undefined")
+        println(io, "Orbital semi-major axis: undefined")
     end
     return
 end
 
+iszeroall(x...) = all(iszero.(x))
+isnanall(x...) = all(isnan.(x))
+replacenan(x) = isnan(x) ? zero(x) : x
+
 """Allows keyword arguments"""
-Elements(;m::T=0.0,P::T=0.0,t0::T=0.0,ecosϖ::T=0.0,esinϖ::T=0.0,I::T=0.0,Ω::T=0.0) where T<:Real = Elements(m,P,t0,ecosϖ,esinϖ,I,Ω)
+function Elements(;m::Real,P::Real=NaN,t0::Real=NaN,ecosω::Real=NaN,esinω::Real=NaN,I::Real=NaN,Ω::Real=NaN,a::Real=NaN,ω::Real=NaN,e::Real=NaN,tp::Real=NaN)
+    # Promote everything
+    m,P,t0,ecosω,esinω,I,Ω,a,ω,e,tp = promote(m,P,t0,ecosω,esinω,I,Ω,a,ω,e,tp)
+    T = typeof(P)
+
+    # Assume if only mass is set, we want the elements for the star (or central body in binary)
+    if isnanall(P,t0,ecosω,esinω,I,Ω,a,ω,e,tp); return Elements(m, zeros(T, length(fieldnames(Elements))-1)...); end
+    if isnanall(P,a); throw(ArgumentError("Must specify either P or a.")); end
+    if P > zero(T) && a > zero(T); throw(ArgumentError("Only one of P (period) or a (semi-major axis) can be defined")); end
+    if P < zero(T); throw(ArgumentError("Period must be positive")); end
+    if a < zero(T); throw(ArgumentError("Orbital semi-major axis must be positive")); end
+    if !(zero(T) <= e < one(T)) && !isnan(e); throw(ArgumentError("Eccentricity must be in [0,1), e=$(e)")); end
+    if !(isnan(ecosω) && isnan(esinω)) && !isnan(e); error("Must specify either e and ecosω/esinω, but not both."); end
+    if !(isnan(tp) || isnan(t0)); error("Must specify either initial transit time or time of periastron passage, but not both."); end
+
+    ω = replacenan(ω)
+    if isnanall(e,ecosω,esinω); e = ecosω = esinω = zero(T); end
+    if isnan(e)
+        ecosω = replacenan(ecosω)
+        esinω = replacenan(esinω)
+        e = sqrt(ecosω*ecosω + esinω*esinω)
+        @assert zero(T) <= e < one(T) "Eccentricity must be in [0,1)"
+        ω = atan(esinω, ecosω)
+    else
+        esinω, ecosω = e.*sincos(ω)
+        esinω = abs(esinω) < eps(T) ? zero(T) : esinω
+        ecosω = abs(ecosω) < eps(T) ? zero(T) : ecosω
+    end
+
+    t0 = replacenan(t0)
+    n = 2π/P
+    if isnan(tp) && !iszero(e)
+        tp = (t0 - sqrt(1.0-e*e)*ecosω/(n*(1.0-esinω)) - (2.0/n)*atan(sqrt(1.0-e)*(esinω+ecosω+e), sqrt(1.0+e)*(esinω-ecosω-e))) % P
+    elseif isnan(tp)
+        θ = π/2 + ω
+        tp = θ / n
+    end
+
+    a = replacenan(a)
+    I = replacenan(I)
+    Ω = replacenan(Ω)
+
+    return Elements(m,P,t0,ecosω,esinω,I,Ω,a,e,ω,tp)
+end
+
+# Handle the deprecated fields
+function Base.getproperty(obj::Elements, sym::Symbol)
+    if (sym === :ecosϖ)
+        Base.depwarn("ecosϖ (\\varpi) will be removed, use ecosω (\\omega).", :getproperty, force=true)
+        return obj.ecosω
+    end
+    if (sym === :esinϖ)
+        Base.depwarn("esinϖ (\\varpi) will be removed, use esinω (\\omega).", :getproperty, force=true)
+        return obj.esinω
+    end
+    if (sym === :ϖ)
+        Base.depwarn("ϖ (\\varpi) will be removed, use ω (\\omega).", :getproperty, force=true)
+        return obj.ω
+    end
+    return getfield(obj, sym)
+end
 
 """Abstract type for initial conditions specifications."""
 abstract type InitialConditions{T} end
@@ -146,7 +212,7 @@ Each method is simply populating the `ElementsIC.elements` field, which is a `Ma
 """
 function ElementsIC(t0::T, H::Matrix{<:Real}, elems::Elements{T}...) where T<:AbstractFloat
     elements = zeros(T,size(H)[1],7)
-    fields = [:m, :P, :t0, :ecosϖ, :esinϖ, :I, :Ω]
+    fields = [:m, :P, :t0, :ecosω, :esinω, :I, :Ω]
     for i in eachindex(elems)
         elements[i,:] .= [getfield(elems[i],f) for f in fields]
     end

src/ics/kepler_init.jl

@@ -81,8 +81,8 @@ function kepler_init(time::T,mass::T,elements::Array{T,1},jac_init::Array{T,2})
     ecosomega = elements[3]
     esinomega = elements[4]
     ecc=sqrt(esinomega^2+ecosomega^2)
-    deccdecos = ecosomega/ecc
-    deccdesin = esinomega/ecc
+    deccdecos = ecc != 0.0 ? ecosomega/ecc : zero(T)
+    deccdesin = ecc != 0.0 ? esinomega/ecc : zero(T)
     #omega = atan2(esinomega,ecosomega)
     # The true anomaly at the time of transit:
     #f1 = 1.5*pi-omega
@@ -91,9 +91,12 @@ function kepler_init(time::T,mass::T,elements::Array{T,1},jac_init::Array{T,2})
     #tp=(t0+period*sqrt1mecc2/2.0/pi*(ecc*sin(f1)/(1.0+ecc*cos(f1))
     #    -2.0/sqrt1mecc2*atan2(sqrt1mecc2*tan(0.5*f1),1.0+ecc)))
     den1 = esinomega-ecosomega-ecc
-    tp = (t0 - sqrt1mecc2/n*ecosomega/(1.0-esinomega)-2/n*
-          #     atan2(sqrt(1.0-ecc)*(esinomega+ecosomega+ecc),sqrt(1.0+ecc)*den1))
-    atan(sqrt(1.0-ecc)*(esinomega+ecosomega+ecc),sqrt(1.0+ecc)*den1))
+    tp = if ecc == 0.0
+        # If circular orbit, take periastron to be along +x-axis
+        t0 - 3*period/4
+    else
+        (t0 - sqrt1mecc2/n*ecosomega/(1.0-esinomega)-2/n*atan(sqrt(1.0-ecc)*(esinomega+ecosomega+ecc),sqrt(1.0+ecc)*den1))
+    end
     dtpdp = (tp-t0)/period
     fac = sqrt((1.0-ecc)/(1.0+ecc))
     den2 = 1.0/den1^2
@@ -125,7 +128,8 @@ function kepler_init(time::T,mass::T,elements::Array{T,1},jac_init::Array{T,2})
     capomega = elements[6]
     # Now, compute the positions
     coscapomega = cos(capomega) ; sincapomega = sin(capomega)
-    cosomega = ecosomega/ecc; sinomega = esinomega/ecc
+    cosomega = ecc != 0.0 ? ecosomega/ecc : one(T)
+    sinomega = ecc != 0.0 ? esinomega/ecc : zero(T)
     cosinc = cos(inc) ; sininc = sin(inc)
     # Define rotation matrices (M&D 2.119-2.120):
     P1 = [cosomega -sinomega 0.0; sinomega cosomega 0.0; 0.0 0.0 1.0]
@@ -172,11 +176,11 @@ function kepler_init(time::T,mass::T,elements::Array{T,1},jac_init::Array{T,2})
     #  println("position vs. period: term 1 ",dxda*dsemidp," term 2: ",dxdekep*dekepdm*dmdp," term 3: ",dxdekep*dekepdm*dmdtp*dtpdp," sum: ",jac_init[1:3,1])
     #end
     jac_init[1:3,2] = dxdekep*dekepdm*dmdtp*dtpdt0
-    jac_init[1:3,3] = dxdecos + dxdekep*(dekepdm*dmdtp*dtpdecos + dekepdecos)
+    jac_init[1:3,3] .= ecc != 0.0 ? dxdecos + dxdekep*(dekepdm*dmdtp*dtpdecos + dekepdecos) : zero(T)
     #if  T != BigFloat
     #  println("position vs. ecosom : term 1 ",dxdecos," term 2: ",dxdekep*dekepdm*dmdtp*dtpdecos," term 3: ",dxdekep*dekepdm*dekepdecos," sum: ",jac_init[1:3,3])
     #end
-    jac_init[1:3,4] = dxdesin + dxdekep*(dekepdm*dmdtp*dtpdesin + dekepdesin)
+    jac_init[1:3,4] .= ecc != 0.0 ? dxdesin + dxdekep*(dekepdm*dmdtp*dtpdesin + dekepdesin) : zero(T)
     #if  T != BigFloat
     #  println("position vs. esinom : term 1 ",dxdesin," term 2: ",dxdekep*dekepdm*dmdtp*dtpdesin," term 3: ",dxdekep*dekepdm*dekepdesin," sum: ",jac_init[1:3,4])
     #end
@@ -188,11 +192,11 @@ function kepler_init(time::T,mass::T,elements::Array{T,1},jac_init::Array{T,2})
     #  println("velocity vs. period: term 1 ",dvdp," term 2: ",dvda*dsemidp," term 3: ",dvdekep*dekepdm*dmdp," term 4: ",dvdekep*dekepdm*dmdtp*dtpdp," sum: ",jac_init[4:6,1])
     #end
     jac_init[4:6,2] = dvdekep*dekepdm*dmdtp*dtpdt0
-    jac_init[4:6,3] = dvdecos + dvdekep*(dekepdm*dmdtp*dtpdecos + dekepdecos)
+    jac_init[4:6,3] .= ecc != 0.0 ? dvdecos + dvdekep*(dekepdm*dmdtp*dtpdecos + dekepdecos) : zero(T)
     #if  T != BigFloat
     #  println("velocity vs. ecosom : term 1 ",dvdecos," term 2: ",dvdekep*dekepdm*dmdtp*dtpdecos," term 3: ",dvdekep*dekepdm*dekepdecos," sum: ",jac_init[4:6,3])
     #end
-    jac_init[4:6,4] = dvdesin + dvdekep*(dekepdm*dmdtp*dtpdesin + dekepdesin)
+    jac_init[4:6,4] .= ecc != 0.0 ? dvdesin + dvdekep*(dekepdm*dmdtp*dtpdesin + dekepdesin) : zero(T)
     #if  T != BigFloat
     #  println("velocity vs. esinom : term 1 ",dvdesin," term 2: ",dvdekep*dekepdm*dmdtp*dtpdesin," term 3: ",dvdekep*dekepdm*dekepdesin," sum: ",jac_init[4:6,4])
     #end

src/outputs/elements.jl

@@ -2,7 +2,7 @@
 
 """
 
-A 3d point in cartesian space 
+A 3d point in cartesian space
 """
 struct Point{T<:AbstractFloat}
     x::T
@@ -26,7 +26,7 @@ function get_relative_positions(x,v,ic::InitialConditions)
     n = ic.nbody
     X = zeros(3,n)
     V = zeros(3,n)
-    
+
     X .= permutedims(ic.amat*x')
     V .= permutedims(ic.amat*v')
 
@@ -51,7 +51,7 @@ mag(x) = sqrt(dot(x))
 mag(x,v) = sqrt(dot(x,v))
 Rdotmag(R,V,h) = sqrt(V^2 - (h/R)^2)
 
-function hvec(r,rdot) 
+function hvec(r,rdot)
     hx,hy,hz = unpack(cross(r,rdot))
     hz >= 0.0 ? hy *= -1 : hx *= -1
     return Point(hx,hy,hz)
@@ -60,10 +60,10 @@ end
 function calc_Ω(hx,hy,h,I)
     sinΩ = hx/(h*sin(I))
     cosΩ = hy/(h*sin(I))
-    return atan(sinΩ,cosΩ) 
+    return atan(sinΩ,cosΩ)
 end
 
-function calc_ϖ(x,R,Rdot,I,Ω,a,e,h)
+function calc_ω(x,R,Rdot,I,Ω,a,e,h)
     # Find ω + f
     wpf = 0.0
     if I != 0.0
@@ -76,8 +76,7 @@ function calc_ϖ(x,R,Rdot,I,Ω,a,e,h)
     cosf = (a*(1.0-e^2)/R - 1.0)/e
     f = atan(sinf,cosf)
 
-    # ϖ = Ω + ω
-    return Ω + wpf - f
+    return wpf - f
 end
 
 function convert_to_elements(x,v,M)
@@ -86,21 +85,25 @@ function convert_to_elements(x,v,M)
     h = mag(hvec(x,v))
     hx,hy,hz = unpack(hvec(x,v))
     Rdot = sign(dot(x,v))*Rdotmag(R,V,h)
-    
-    Gmm = M 
-    
+
+    Gmm = M
+
     a = 1.0/((2.0/R) - (V*V)/(Gmm))
     e = sqrt(1.0 - (h*h/(Gmm*a)))
     I = acos(hz/h)
-    
+
     # Make sure Ω is defined
     I != 0.0 ? Ω = calc_Ω(hx,hy,h,I) : Ω = 0.0
-    
-    ϖ = calc_ϖ(x,R,Rdot,I,Ω,a,e,h)
+
+    ω = calc_ω(x,R,Rdot,I,Ω,a,e,h)
     P = (2.0*π)*sqrt(a*a*a/Gmm)
 
-    return [P,0.0,e*cos(ϖ),e*sin(ϖ),I,Ω,a,e,ϖ]
-end    
+    n = 2π/P
+    ecosω, esinω = e.*sincos(ω)
+    tp = (-sqrt(1.0-e*e)*ecosω/(n*(1.0-esinω)) - (2.0/n)*atan(sqrt(1.0-e)*(esinω+ecosω+e), sqrt(1.0+e)*(esinω-ecosω-e))) % P
+
+    return [P,0.0,ecosω,esinω,I,Ω,a,e,ω,tp]
+end
 
 function get_orbital_elements(s::State{T},ic::InitialConditions{T}) where T<:AbstractFloat
     elems = Elements{T}[]
@@ -120,14 +123,14 @@ function get_orbital_elements(s::State{T},ic::InitialConditions{T}) where T<:Abs
 
         new_elems = convert_to_elements(X[i+b],V[i+b],μ[i+b])
         push!(elems,Elements(ic.m[i+1],new_elems...))
-        
+
         # Compensate for new binary in step
         if b > 0
             b -= 2
         elseif b < 0
             i += 1
         end
-        
+
     i+=1
     end
     return elems

test/test_initial_conditions.jl

@@ -28,7 +28,7 @@ function get_elements_ic_elems(t0, H, N)
     elements = readdlm(fname, ',', comments=true)
     elems = Elements{Float64}[]
     for i in 1:N
-        push!(elems, Elements(elements[i,:]...))
+        push!(elems, Elements(elements[i,:]..., zeros(4)...))
     end
     ic = ElementsIC(t0, H, elems)
     return ic
@@ -154,4 +154,11 @@ end
     @testset "Defaults" begin
         test_default_trappist()
     end
+
+    @testset "Deprecated" begin
+        el = Elements(m=1.0, P=1.0)
+        @test_logs (:warn, "ecosϖ (\\varpi) will be removed, use ecosω (\\omega).") Base.getproperty(el, :ecosϖ)
+        @test_logs (:warn, "esinϖ (\\varpi) will be removed, use esinω (\\omega).") Base.getproperty(el, :esinϖ)
+        @test_logs (:warn, "ϖ (\\varpi) will be removed, use ω (\\omega).") Base.getproperty(el, :ϖ)
+    end
 end