Merge pull request #51 from marius311/modular_sampling

make sample_joint code more modular / customizable

Author: marius311

src/CMBLensing.jl

@@ -11,12 +11,14 @@ using Base: @kwdef, @propagate_inbounds, Bottom, OneTo, showarg, show_datatype,
 using Combinatorics
 using DataStructures
 using DelimitedFiles
-using Distributed: pmap, nworkers, myid, workers, addprocs, @everywhere, remotecall_wait, @spawnat, pgenerate, procs, @fetchfrom
+using Distributed: pmap, nworkers, myid, workers, addprocs, @everywhere, remotecall_wait, 
+    @spawnat, pgenerate, procs, @fetchfrom, default_worker_pool
 using FileIO
 using FFTW
 using InteractiveUtils
 using IterTools: flagfirst
 using JLD2
+using JLD2: jldopen, JLDWriteSession
 using KahanSummation
 using Loess
 using LinearAlgebra
@@ -86,7 +88,14 @@ export
     simulate, SymmetricFuncOp, symplectic_integrate, Taylens, toCℓ, toDℓ,
     ud_grade, unbatch, unmix, white_noise, Ð, Ł,  
     ℓ², ℓ⁴, ∇, ∇², ∇ᵢ, ∇ⁱ
-    
+
+# bunch of sampling-related exports
+export gibbs_initialize_f!, gibbs_initialize_ϕ!, gibbs_initialize_θ!, 
+    gibbs_sample_f!, gibbs_sample_ϕ!, gibbs_sample_slice_θ!, 
+    gibbs_mix!, gibbs_unmix!, gibbs_postprocess!, 
+    once_every, start_after_burnin, mass_matrix_ϕ, hmc_step
+
+
 # generic stuff
 include("util.jl")
 include("util_fft.jl")

src/dataset.jl

@@ -6,20 +6,35 @@ getproperty(ds::DS, k::Symbol) where {DS<:DataSet{<:DataSet}} =
     hasfield(DS, k) ? getfield(ds, k) : getproperty(getfield(ds, :_super), k)
 setproperty!(ds::DS, k::Symbol, v) where {DS<:DataSet{<:DataSet}} = 
     hasfield(DS, k) ? setfield!(ds, k, v) : setproperty!(getfield(ds, :_super), k, v)
-propertynames(ds::DS) where {DS′<:DataSet, DS<:DataSet{DS′}} = 
-    union(fieldnames(DS), fieldnames(DS′))
+propertynames(ds::DS) where {DS<:DataSet} = propertynames(DS)
+propertynames(::Type{DS}) where {DS′<:DataSet, DS<:DataSet{DS′}} = 
+    union(propertynames(DS′), fieldnames(DS))
+propertynames(::Type{DS}) where {DS<:DataSet{Nothing}} = fieldnames(DS)
 
 function new_dataset(::Type{DS}; kwargs...) where {DS′<:DataSet, DS<:DataSet{DS′}}
-    kw  = filter(((k,_),)->  k in fieldnames(DS),  kwargs)
-    kw′ = filter(((k,_),)->!(k in fieldnames(DS)), kwargs)
-    DS(_super=DS′(;kw′...); kw...)
+    kw  = Dict(k => v for (k,v) in kwargs if   k in fieldnames(DS))
+    kw′ = Dict(k => v for (k,v) in kwargs if !(k in fieldnames(DS)))
+    DS(; kw..., _super=new_dataset(DS′; kw′...))
+end
+
+function new_dataset(::Type{DS}; kwargs...) where {DS<:DataSet{Nothing}}
+    DS(; kwargs...)
 end
 
 copy(ds::DS) where {DS<:DataSet} = 
     DS(((k==:_super ? copy(v) : v) for (k,v) in pairs(fields(ds)))...)
 
 hash(ds::DataSet, h::UInt64) = foldr(hash, (typeof(ds), fieldvalues(ds)...), init=h)
 
+function show(io::IO, ds::DataSet)
+    println(io, typeof(ds), ": ")
+    ds_dict = OrderedDict(k => getproperty(ds,k) for k in propertynames(ds) if k!=Symbol("_super"))
+    for line in split(sprint(show, MIME"text/plain"(), ds_dict, context = (:limit => true)), "\n")[2:end]
+        println(io, line)
+    end
+end
+
+
 # needed until fix to https://github.com/FluxML/Zygote.jl/issues/685
 Zygote.grad_mut(ds::DataSet) = Ref{Any}((;(propertynames(ds) .=> nothing)...))
 

src/healpix.jl

@@ -15,42 +15,43 @@ function xy_to_θϕ((x,y))
     θ, ϕ
 end
 
-function healpix_to_flat(healpix_map::Vector{T}, proj::ProjLambert{T}; rot=(0,0,0)) where {T}
+function healpix_to_flat(healpix_map::Vector{T}, proj::ProjLambert{T}; rots=[((0,90,0),)]) where {T}
 
     Nside_sphere = hp.npix2nside(length(healpix_map))
-    @unpack Δx = proj
+    @unpack Δx, Ny, Nx = proj
 
     # compute projection coordinate mapping
-    xs = ys = Δx*((-Nside÷2:Nside÷2-1) .+ 0.5)
+    ys = @. Δx * ((-Nx÷2:Nx÷2-1) + 0.5) # x/y switch here intentional
+    xs = @. Δx * ((-Ny÷2:Ny÷2-1) + 0.5)
     xys = tuple.(xs,ys')[:]
     θϕs = xy_to_θϕ.(xys)
     (θs, ϕs) = first.(θϕs), last.(θϕs)
 
-    # rotate the pole to the equator to match Healpy's azeqview convention, in
-    # addition to applying the user rotation
-    R = hp.Rotator((0,90,0), eulertype="ZYX") * hp.Rotator(rot, eulertype="ZYX")
+    # the default rots=[(0,90,0)] makes it so you get a view of the
+    # equator, to match Helapy's azeqview convention
+    R = prod([hp.Rotator(rot..., eulertype="ZYX") for rot in rots])
     (θs′, ϕs′) = eachrow(R.get_inverse()(θs, ϕs))
 
     # interpolate map
-    FlatMap(reshape(hp.get_interp_val(healpix_map, θs′, ϕs′), Nside, Nside), proj)
+    FlatMap(reshape(hp.get_interp_val(healpix_map, θs′, ϕs′), Ny, Nx), proj)
 
 end
 
-function healpix_pixel_centers_to_flat(f::FlatField, Nside_sphere; rots=[(0,90,0)], healpix_pixels=0:(12*Nside_sphere^2-1))
+function healpix_pixel_centers_to_flat(f::FlatField, Nside_sphere; rots=[((0,90,0),)], healpix_pixels=0:(12*Nside_sphere^2-1))
 
     @unpack Δx, Ny, Nx = f
 
     (θs, ϕs) = hp.pix2ang(Nside_sphere, healpix_pixels)
 
-    # rotate the pole to the equator to match Healpy's azeqview convention, in
-    # addition to applying the user rotation
-    R = prod([hp.Rotator(rot, eulertype="ZYX") for rot in rots])
+    # the default rots=[(0,90,0)] makes it so you get a view of the
+    # equator, to match Helapy's azeqview convention
+    R = prod([hp.Rotator(rot..., eulertype="ZYX") for rot in rots])
     (θs′, ϕs′) = eachrow(R(θs, ϕs))
 
     # compute projection coordinate mapping
     # (using Ny for xs and vice-versa intentional)
     xys = @. θϕ_to_xy(tuple(θs′, ϕs′))
-    xs = @. first(xys) / Δx + Ny÷2 + 0.5
+    xs = @. first(xys) / Δx + Ny÷2 + 0.5 # x/y switch here intentional
     ys = @.  last(xys) / Δx + Nx÷2 + 0.5
 
     (xs, ys)