Bugfixes

sefffal · Apr 17, 2023 · 93bf069 · 93bf069
1 parent 2fd5a5e
commit 93bf069
Show file tree

Hide file tree

Showing 9 changed files with 65 additions and 56 deletions.
diff --git a/OctofitterImages/Project.toml b/OctofitterImages/Project.toml
@@ -12,6 +12,7 @@ Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 Octofitter = "daf3887e-d01a-44a1-9d7e-98f15c5d69c9"
 PlanetOrbits = "fd6f9641-d78f-43ce-a379-ceb0bddb468a"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9"
 

diff --git a/OctofitterImages/src/OctofitterImages.jl b/OctofitterImages/src/OctofitterImages.jl
@@ -9,7 +9,7 @@ using ImageTransformations
 using CoordinateTransformations
 using Interpolations
 using AstroImages
-
+using Statistics
 
 const images_cols = (:band, :image, :epoch, :platescale,)
 
@@ -59,8 +59,8 @@ struct ImageLikelihood{TTable<:Table} <: Octofitter.AbstractLikelihood
         return new{typeof(table)}(table)
     end
 end
-Images(observations::NamedTuple...) = ImageLikelihood(observations)
-export Images
+ImageLikelihood(observations::NamedTuple...) = ImageLikelihood(observations)
+export ImageLikelihood
 
 
 """
@@ -115,7 +115,15 @@ function contrast(image::AstroImage; step=2)
 end
 
 
-
+function imgsep(image::AstroImage)
+    dx = dims(image,X)
+    dy = collect(dims(image,Y))
+    dr = sqrt.(
+        dx.^2 .+ (dy').^2
+    )
+    return dr
+end
+
 
 
 """
@@ -168,13 +176,15 @@ function Octofitter.ln_like(images::ImageLikelihood, θ_planet, orbit)
         # TODO: verify this is type stable
         f_band = getproperty(θ_planet, band)
 
+        # σ_add = θ_planet.H_σ
 
         # When we get a position that falls outside of our available
         # data (e.g. under the coronagraph) we cannot say anything
         # about the likelihood. This is equivalent to σₓ→∞ or log likelihood 
         # of zero.
         if !isfinite(σₓ) || !isfinite(f̃ₓ)
             continue
+            # return convert(T, -Inf)
         end
 
         # Direct imaging likelihood.
@@ -187,7 +197,7 @@ function Octofitter.ln_like(images::ImageLikelihood, θ_planet, orbit)
         # Ruffio et al 2017, eqn (31)
         # Mawet et al 2019, eqn (8)
 
-        σₓ² = σₓ^2
+        σₓ² = σₓ^2 #+ σ_add^2
         ll += -1 / (2*σₓ²) * (f_band^2 - 2f_band * f̃ₓ)
     end
 
@@ -197,39 +207,36 @@ end
 
 
 # Generate new images
-function Octofitter.generate_from_params(like::ImageLikelihood, θ_system,  elements::Vector{<:Visual{KepOrbit}})
+# function Octofitter.generate_from_params(like::ImageLikelihood, θ_system,  elements::Vector{<:Visual{KepOrbit}})
+function Octofitter.generate_from_params(like::ImageLikelihood, θ_planet,  orbit::PlanetOrbits.AbstractOrbit)
 
     newrows = map(like.table) do row
         (;band, image, platescale, epoch, psf) = row
 
         injected = copy(image)
-
-        for i in eachindex(elements)
-            θ_planet = θ_system.planets[i]
-            elem = elements[i]
-            # Generate new astrometry point
-            os = orbitsolve(elem, epoch)
-
-            # TODO: this does not consider the shift to the images due to the motion of the star
-            ra = raoff(os)
-            dec = decoff(os)
-
-            phot = θ_planet[band]
-
-            # TODO: verify signs
-            dx = ra/platescale
-            dy = -dec/platescale
-            translation_tform = Translation(
-                mean(axes(psf,1))-mean(axes(image,1))+mean(dims(image,1))+dx,
-                mean(axes(psf,2))-mean(axes(image,2))+mean(dims(image,2))+dy
-            )
-            # TBD if we want to support rotations for handling negative sidelobes.
-
-            psf_positioned = warp(psf, translation_tform, axes(image), fillvalue=0)
-            psf_positioned[.! isfinite.(psf_positioned)] .= 0
-            psf_scaled = psf_positioned .* phot ./ maximum(filter(isfinite, psf_positioned))
-            injected .+= psf_scaled
-        end
+
+        # Generate new astrometry point
+        os = orbitsolve(orbit, epoch)
+
+        # TODO: this does not consider the shift to the images due to the motion of the star
+        ra = raoff(os)
+        dec = decoff(os)
+
+        phot = θ_planet[band]
+
+        # TODO: verify signs
+        dx = ra/platescale
+        dy = -dec/platescale
+        translation_tform = Translation(
+            mean(axes(psf,1))-mean(axes(image,1))+mean(dims(image,1))+dx,
+            mean(axes(psf,2))-mean(axes(image,2))+mean(dims(image,2))+dy
+        )
+        # TBD if we want to support rotations for handling negative sidelobes.
+
+        psf_positioned = warp(psf, translation_tform, axes(image), fillvalue=0)
+        psf_positioned[.! isfinite.(psf_positioned)] .= 0
+        psf_scaled = psf_positioned .* phot ./ maximum(filter(isfinite, psf_positioned))
+        injected .+= psf_scaled
 
         return merge(row, (;image=injected))
     end

diff --git a/Project.toml b/Project.toml
@@ -57,14 +57,14 @@ DataDeps = "0.7"
 Distributions = "0.25"
 ForwardDiff = "0.10"
 Interpolations = "0.13, 0.14"
-MCMCChains = "5.0"
+MCMCChains = "6.0"
 NamedTupleTools = "0.13, 0.14"
 RecipesBase = "1.2"
 RuntimeGeneratedFunctions = "0.5"
 StaticArrays = "1.2"
 StatsBase = "0.33"
 Tables = "1.6"
-PlanetOrbits = "0.5"
+PlanetOrbits = "0.6"
 julia = "1.9"
 
 [extras]

diff --git a/docs/src/api.md b/docs/src/api.md
@@ -6,7 +6,7 @@ Priors
 Derived
 AstrometryLikelihood
 PhotometryLikelihood
-ProperMotionAnom
+ProperMotionAnomLikelihood
 HGCALikelihood
 ImageLikelihood
 Sine

diff --git a/ext/PlotsExt.jl b/ext/PlotsExt.jl
@@ -555,7 +555,8 @@ function Octofitter.timeplotgrid(
     Plots.annotate!(pz,((0.5,1.0),Plots.text("top down",:bottom,:middle,10)))
     Plots.scatter!(pz, [0],[0],marker=(:star, :white, :black, 5),label="")
     Plots.annotate!(pz,((-0.2,0.00),Plots.text("↓","Helvetica",:center,15)))
-    Plots.annotate!(pz,((-0.2,-0.15),Plots.text("🌍","Helvetica",:center,15)))
+    # Plots.annotate!(pz,((-0.2,-0.15),Plots.text("🌍","Helvetica",:center,15)))
+    Plots.annotate!(pz,((-0.2,-0.15),Plots.text("Earth",:center,5)))
 
     py = Plots.plot()
     for (i,planet_key) in enumerate(planet_keys)

diff --git a/src/io.jl b/src/io.jl
@@ -25,12 +25,12 @@ function savechain(fname, chain::MCMCChains.Chains)
         # Expand any vectors into multiple headers
         for k in ks
             i = findfirst(==(k),ks)
-            if vals[i] isa AbstractArray{<:Number}
+            if typeof(vals[i]) <: AbstractArray{<:Number}
                 for ki in eachindex(vals[i])
                     push!(ks, "$(k)_$ki")
                     push!(vals, vals[i][ki])
-                    vals[i] = "ARRAY"
                 end
+                vals[i] = "ARRAY"
             end
         end
 

diff --git a/src/likelihoods/hgca.jl b/src/likelihoods/hgca.jl
@@ -26,7 +26,7 @@ export gaia_plx
 
 
 const pma_cols = (:ra_epoch, :dec_epoch, :dt, :pm_ra, :pm_dec, :σ_pm_ra, :σ_pm_dec,)
-struct ProperMotionAnom{TTable<:Table} <: AbstractLikelihood
+struct ProperMotionAnomLikelihood{TTable<:Table} <: AbstractLikelihood
     table::TTable
     function ProperMotionAnom(observations...)
         table = Table(observations...)
@@ -90,7 +90,7 @@ export HGCALikelihood
 General proper motion likelihood at any number of epochs.
 Each epoch is averaged over 5 measurements at +-dt/2.
 """
-function ln_like(pma::ProperMotionAnom, θ_system, elements)
+function ln_like(pma::ProperMotionAnomLikelihood, θ_system, elements)
     ll = 0.0
 
     # How many points over Δt should we average the proper motion at each
@@ -182,8 +182,8 @@ function ln_like(pma::HGCALikelihood, θ_system, elements, _L=1 #=length of obse
             # meaning we calculate the orbit 2x as much as we need.
             o_ra = orbitsolve(orbit, years2mjd(hgca.epoch_ra_hip[1])+δt)
             o_dec = orbitsolve(orbit, years2mjd(hgca.epoch_dec_hip[1])+δt)
-            ra_hip_model += -raoff(o_ra) * θ_planet.mass*mjup2msol/orbit.M
-            dec_hip_model += -decoff(o_dec) * θ_planet.mass*mjup2msol/orbit.M
+            ra_hip_model += -raoff(o_ra) * θ_planet.mass*mjup2msol/totalmass(orbit)
+            dec_hip_model += -decoff(o_dec) * θ_planet.mass*mjup2msol/totalmass(orbit)
             pmra_hip_model += pmra(o_ra, θ_planet.mass*mjup2msol)
             pmdec_hip_model += pmdec(o_dec, θ_planet.mass*mjup2msol)
         end
@@ -213,8 +213,8 @@ function ln_like(pma::HGCALikelihood, θ_system, elements, _L=1 #=length of obse
             # make it the same unit as the stellar mass (element.M)
             o_ra = orbitsolve(orbit, years2mjd(hgca.epoch_ra_gaia[1])+δt)
             o_dec = orbitsolve(orbit, years2mjd(hgca.epoch_dec_gaia[1])+δt)
-            ra_gaia_model += -raoff(o_ra) * θ_planet.mass*mjup2msol/orbit.M
-            dec_gaia_model += -decoff(o_dec) * θ_planet.mass*mjup2msol/orbit.M
+            ra_gaia_model += -raoff(o_ra) * θ_planet.mass*mjup2msol/totalmass(orbit)
+            dec_gaia_model += -decoff(o_dec) * θ_planet.mass*mjup2msol/totalmass(orbit)
             pmra_gaia_model += pmra(o_ra, θ_planet.mass*mjup2msol)
             pmdec_gaia_model += pmdec(o_dec, θ_planet.mass*mjup2msol)
         end
@@ -330,8 +330,8 @@ function generate_from_params(like::HGCALikelihood, θ_system, orbit::PlanetOrbi
             # meaning we calculate the orbit 2x as much as we need.
             o_ra = orbitsolve(orbit, years2mjd(hgca.epoch_ra_hip[1])+δt)
             o_dec = orbitsolve(orbit, years2mjd(hgca.epoch_dec_hip[1])+δt)
-            ra_hip_model += -raoff(o_ra) * θ_planet.mass*mjup2msol/orbit.M
-            dec_hip_model += -decoff(o_dec) * θ_planet.mass*mjup2msol/orbit.M
+            ra_hip_model += -raoff(o_ra) * θ_planet.mass*mjup2msol/totalmass(orbit)
+            dec_hip_model += -decoff(o_dec) * θ_planet.mass*mjup2msol/totalmass(orbit)
             pmra_hip_model += pmra(o_ra, θ_planet.mass*mjup2msol)
             pmdec_hip_model += pmdec(o_dec, θ_planet.mass*mjup2msol)
         end
@@ -361,8 +361,8 @@ function generate_from_params(like::HGCALikelihood, θ_system, orbit::PlanetOrbi
             # make it the same unit as the stellar mass (element.M)
             o_ra = orbitsolve(orbit, years2mjd(hgca.epoch_ra_gaia[1])+δt)
             o_dec = orbitsolve(orbit, years2mjd(hgca.epoch_dec_gaia[1])+δt)
-            ra_gaia_model += -raoff(o_ra) * θ_planet.mass*mjup2msol/orbit.M
-            dec_gaia_model += -decoff(o_dec) * θ_planet.mass*mjup2msol/orbit.M
+            ra_gaia_model += -raoff(o_ra) * θ_planet.mass*mjup2msol/totalmass(orbit)
+            dec_gaia_model += -decoff(o_dec) * θ_planet.mass*mjup2msol/totalmass(orbit)
             pmra_gaia_model += pmra(o_ra, θ_planet.mass*mjup2msol)
             pmdec_gaia_model += pmdec(o_dec, θ_planet.mass*mjup2msol)
         end

diff --git a/src/sampling.jl b/src/sampling.jl
@@ -324,7 +324,7 @@ function construct_elements(chain::Chains, planet_key::Union{String,Symbol}, ii:
         error("Unrecognized chain format")
     end
 end
-construct_elements(chain::Chains, planet_key::Union{String,Symbol}, ii::Colon) = construct_elements(chain, planet_key, 1:size(chain,1))
+construct_elements(chain::Chains, planet_key::Union{String,Symbol}, ii::Colon) = construct_elements(chain, planet_key, 1:size(chain,1)*size(chain,3))
 function construct_elements(chain, planet_key::Union{String,Symbol}, ii::AbstractArray{<:Union{Integer,CartesianIndex}})
     pk = string(planet_key)
     Ms=chain[:,"M"]
@@ -1007,12 +1007,12 @@ function advancedhmc(
     for (i,mc_samples) in enumerate(mc_samples_all_chains)
         stat = map(s->s.stat, mc_samples)
         numerical_error = getproperty.(stat, :numerical_error)
-        tree_depth = getproperty.(stat, :tree_depth)
+        actual_tree_depth = getproperty.(stat, :tree_depth)
 
         mean_accept = mean(getproperty.(stat, :acceptance_rate))
         num_err_frac = mean(numerical_error)
-        mean_tree_depth = mean(tree_depth)
-        max_tree_depth_frac = mean(getproperty.(stat, :tree_depth) .== tree_depth)
+        mean_tree_depth = mean(actual_tree_depth)
+        max_tree_depth_frac = mean(==(tree_depth), actual_tree_depth)
 
         verbosity >= 1 && println("""
         Sampling report for chain $i:
@@ -1079,7 +1079,7 @@ function advancedhmc(
             stop_time,
             # model=model.system,
             # logpost=logposts_mat,
-            states=mc_samples_all_chains,
+            # states=mc_samples_all_chains,
             # pathfinder=pathfinder_chain_with_info,
         )
     )

diff --git a/src/variables.jl b/src/variables.jl
@@ -215,7 +215,7 @@ orbittype(::Planet{TElem}) where TElem = TElem
 Construct a model of a system.
 Must be constructed with a block of priors, and optionally
 additional derived parameters.
-You may provide `ProperMotionAnom()` and/or `Images()` of the system.
+You may provide `ProperMotionAnomLikelihood()` and/or `Images()` of the system.
 Finally, planet models are listed last.
 `name` must be a symbol e.g. `:HD56441`.
 """
@@ -273,7 +273,7 @@ export astrometry
 
 function propermotionanom(system::System)
     for like in system.observations
-        if like isa ProperMotionAnom || like isa HGCALikelihood
+        if like isa ProperMotionAnomLikelihood || like isa HGCALikelihood
             return like
         end
     end