add equirectangular projection

marius311 commented 3 years ago

@EthanAnderes here's the start of having what I'm calling "equirectangular projection". Best place to start is look at this notebook which gives you a feel for how they are defined / how to construct them. A bunch of stuff already works, including:

constructing them
broadcasting with them
having them be GPU backed
projecting between them and Healpix (for spin-0 only)

A rough summary of the remaining work is:

[x] defining basis transformations (ie the block diagonal matrices)
- [x] spin-0
- [x] spin-2
[x] defining dot products / logdets / traces
[x] making sure all the factors of pixel area are consistent with themselves and with the auto-diff code
[x] ~defining gradients (the code already had notion of co and contravariant derivatives ∇ⁱ vs ∇ᵢ and I believe my LenseFlow code handles these correctly, so once gradients work, LenseFlow should work without change)~ will do as a separate followup
[x] writing plotting functions
[x] oher loose ends which will surely pop up
[x] lazy operator sqrt and logdet

If you search the file proj_equirect.jl for TODO you'll see I've marked exactly where in the code some of these things would go. Perhaps you can review some of what is there and let me know if you have comments (there's almost definitely things I've done that need explaining / could be improved), and then I'd say start trying to implement the most basic basis transformation, Fourier(f::EquiRectMap), and lets go from there. I'm around to chat about this, which I'm sure we'll want / need to to make more progress.

For reference, spt3g names: https://github.com/SouthPoleTelescope/spt3g_software/blob/e344c30bd45de2ebbaa1b2ef72f718b64fe7e36a/maps/include/maps/FlatSkyProjection.h#L10-L43

marius311 commented 3 years ago

Fyi the cause of the test failures is since you bumped the upper limit on FFTW (the limit was there because I had noticed this error before). Unless you really need 1.4, I'd say revert that change and I'll figure out what's needed to upgrade to 1.4 in a separate PR later.

EthanAnderes commented 3 years ago

Is it just complaining about using the new FFTW.set_provider!()? Before the upgrade to v1.4 I was getting some issues when switching to/from FFTW/MKL so my preference would be to keep the upgrade version. I can try to fix it over the weekend

marius311 commented 3 years ago

OK, will work on getting the upgaded version then. It should be easy but it was a pain to figure out how to do it in a backwards compatible way when I looked earlier, so I just went with the easier option of capping it at 1.3.

marius311 commented 3 years ago

Basically there was a rename from FFTW.fftw_vendor to FFTW.fftw_provider (or vice versa, dont remember) but somehow it was really hard to track down which versions use which.

marius311 commented 3 years ago

Ok, I added a set of tests here which should give you a concrete target to hit and then this PR can basically be merged. As soon as all of these work, then pretty much posterior sampling for lense-free datasets will work with no changes to existing posterior / sampling code since all it does is depend on these things (granted there'll probably be some remaining little thing, which I can help iron out).

EthanAnderes commented 3 years ago

For testing equality of fields, like in @test AzFourier(f0) ≈ f0 for eg, I'm running into errors that look like the following. Not sure what unbatch is or how it works. Any tips?

julia> @test AzFourier(f0) ≈ f0
Error During Test at REPL[58]:1
  Test threw exception
  Expression: AzFourier(f0) ≈ f0
  MethodError: no method matching unbatch(::ComplexF64)
  Closest candidates are:
    unbatch(::CMBLensing.BatchedReal; dims) at /Users/ethananderes/Dropbox/CMBLensing.jl/src/batched_reals.jl:34
    unbatch(::Real; dims) at /Users/ethananderes/Dropbox/CMBLensing.jl/src/batched_reals.jl:35
    unbatch(::BaseField{B, M, T, A} where {M<:ProjLambert, T, A<:(AbstractArray{T, N} where N)}) where B at /Users/ethananderes/Dropbox/CMBLensing.jl/src/proj_lambert.jl:474
    ...
  Stacktrace:
    [1] (::CMBLensing.var"#_norm#112")(x::EquiRectAzFourier{Array{ComplexF64, 2}})
      @ CMBLensing ~/Dropbox/CMBLensing.jl/src/generic.jl:287
    [2] isapprox(x::EquiRectAzFourier{Array{ComplexF64, 2}}, y::EquiRectAzFourier{Array{ComplexF64, 2}}; atol::Int64, rtol::Float64)
      @ CMBLensing ~/Dropbox/CMBLensing.jl/src/generic.jl:288
    [3] isapprox(x::EquiRectAzFourier{Array{ComplexF64, 2}}, y::EquiRectAzFourier{Array{ComplexF64, 2}})
      @ CMBLensing ~/Dropbox/CMBLensing.jl/src/generic.jl:287
    [4] eval_test(evaluated::Expr, quoted::Expr, source::LineNumberNode, negate::Bool)
      @ Test ~/Software/julia1.6/usr/share/julia/stdlib/v1.6/Test/src/Test.jl:277
    [5] top-level scope
      @ REPL[58]:1
...

marius311 commented 3 years ago

The underlying problem looks to be that dot(f,f) for an AzFourier object returns a complex number. Thats how norm is implemented and norm is used in isapprox, which applies unbatch to the output of norm which it expected to be real. This will work after property defining that dot product for AzFourier fields.

I would recomend implementing dot by just going to pixel-space and doing the trivial thing there. One extra FFT really doesnt matter plus the code will be much easier to work on GPU. Here's how its done for Lambert:

https://github.com/marius311/CMBLensing.jl/blob/7c7173e4e1eabedb061a2bef802d5849977adb46/src/proj_lambert.jl#L303-L307

Something nearly identical (or perhaps exactly identical, modulo deciding where a factor of Ω goes) should be the right thing for EquiRect too.

If you're wondering what the batching is btw, its a mainly GPU oriented optimization where you pack multiple fields into one and operate on all of them at once. So if you pack eg 4 fields together, the dot product returns 4 numbers. That's what all the batch/unbatch stuff is. Eg:

julia> f = FlatMap(rand(10,10,4))
400-element 10×10(×4)-pixel 1.0′-resolution LambertMap{Array{Float64, 4}}:
 0.761341248724758
 0.7349820676347125
 0.5456431687843781
 ⋮
 0.0020472162107947334
 0.5248011451792878

julia> dot(f,f)
Batched[28.972324801946723, 28.760678890273393, 36.43664766782598, 38.16361592796549]

marius311 commented 3 years ago

Tests are now running on Github Actions, 9 passing / 17 failing (same as I see locally).

The immediate TODO as I see it is:

Fix basis transformations, QUAzFourier and QUMap don't work.
Need to implement dot. What you said about unitary transforms yesterday was a good point but then I remembered QUAzFourier isnt just an FFT, its also a rearranging / scaling, so I'm not sure if its not still just easier to do the dot product in map space. Up to you.

That should get us alot further of the way there.

EthanAnderes commented 3 years ago

Tests are now running on Github Actions, 9 passing / 17 failing (same as I see locally).

The immediate TODO as I see it is:

Fix basis transformations, QUAzFourier and QUMap don't work.

Need to implement dot. What you said about unitary transforms yesterday was a good point but then I remembered QUAzFourier isnt just an FFT, its also a rearranging / scaling, so I'm not sure if its not still just easier to do the dot product in map space. Up to you.

That should get us alot further of the way there.

Looking into problem with QUMap it appears the constructor. I've just been using a constructor based on the struct fields, assuming this was coming from Base Julia. The one you have in the test set comes from proj_cartesian.jl (I think??) and appears to load the arr field with a 3-dim array (rather than a 2-dim complex array as intended). Using the original constructor it works.

julia> let Nside = (128,128), θspan  = (π/2 .- deg2rad.((-40,-70))), φspanᵒ = deg2rad.((-50, 50))

           proj = ProjEquiRect(;Ny=Nside[1], Nx=Nside[2], θspan, φspan=φspanᵒ)
           f0 = EquiRectMap(rand(Nside...), proj)
           f2 = EquiRectQUMap(rand(ComplexF64, Nside...), proj)

           @test f0 isa EquiRectMap
           @test f2 isa EquiRectQUMap

           # transform
           @test Map(AzFourier(f0)) ≈ f0
           @test QUMap(QUAzFourier(f2)) ≈ f2

       end
Test Passed

EthanAnderes commented 3 years ago

Need to implement dot. What you said about unitary transforms yesterday was a good point but then I remembered QUAzFourier isnt just an FFT, its also a rearranging / scaling, so I'm not sure if its not still just easier to do the dot product in map space. Up to you.

This was already done in

commit 84ee35dc93f155ce21f4e1159f576d2e94ec5562
Author: Ethan Anderes <anderes@ucdavis.edu>
Date:   Fri Aug 13 16:25:41 2021 -0700

    fixed dot. added white_noise

Lets see if the tests pass after the most recent change with the f0 and f2 constructor in the test set

marius311 commented 3 years ago

No I think there's a misunderstanding. A QUMap is stored as a real Ny×Nx×2 array. The old tests used EquiRectQUMap(rand(Nside...), rand(Nside...); θspan, φspan)) and that form of the constructor, defined here https://github.com/marius311/CMBLensing.jl/blob/ecb98af31c6485153223649c338b1078044db087/src/proj_cartesian.jl#L25-L30 concatenated the two random matrices along the third dimension.

Your current form EquiRectQUMap(rand(ComplexF64, Nside...), proj) may not error (it probably should) but it needs to be EquiRectQUMap(rand(Float64, Nside..., 2), proj) to actually be a QUMap of the right format.

marius311 commented 3 years ago

This was already done

Ah missed that sorry.

EthanAnderes commented 3 years ago

No I think there's a misunderstanding. A QUMap is stored as a real Ny×Nx×2 array. The old tests used EquiRectQUMap(rand(Nside...), rand(Nside...); θspan, φspan)) and that form of the constructor, defined here

https://github.com/marius311/CMBLensing.jl/blob/ecb98af31c6485153223649c338b1078044db087/src/proj_cartesian.jl#L25-L30

concatenated the two random matrices along the third dimension. Your current form EquiRectQUMap(rand(ComplexF64, Nside...), proj) may not error (it probably should) but it needs to be EquiRectQUMap(rand(Float64, Nside..., 2), proj) to actually be a QUMap of the right format.

To me if feels closer to the physical field itself for the polarization field to be a complex Ny×Nx array. Also the QU operators are more naturally expressed in terms of complex matrices. So my preference would be to keep this format, but in the end I just want the least amount of change to get this done. Any idea what would be the least amount of code change to get this to where you want it: define an alternative to QUMap Basis type, or change complex Ny×Nx arrays to real Ny×Nx×2 arrays?

marius311 commented 3 years ago

The reason for the Ny×Nx×2 is mainly to ease handling of IQU, which is just a Ny×Nx×3, otherwise you have to have separate arrays for I and P. I think the fix in this case is simple, at the top of your implementation of QUAzFourier(f::EquiRectQUMap), you essentially wrote that function assuming f.arr is a complex Ny×Nx array, so just do arr_complex = f.arr[:,:,1] + im * f.arr[:,:,2] or something and then just go with your existing code from there. Its fine that the QUAzFourier array doesn't have the 3rd dimension, I was always thinking of QUAzFourier as this basis which was just specifically the thing your code needs, and its fine to be in a somewhat different format.

Edit: actually note that for a QUMap, f.Qx and f.Ux are aliases for view(f.arr, :, :, 1) and view(f.arr, :, :, 2) so the clearest / fastest this is probably f.Qx + im * f.Ux

marius311 commented 3 years ago

Is there any sense in which we can make E and B maps (for visualization) using what we have here?

marius311 commented 3 years ago

Basic AD should be working now and tests pass :tada:. I'll work on making it more robust after we merge this. Two last things we need:

Constructor for beam operator. Can you write this?
Just thought of this, but how do you scale the BlockDiagEquiRect for r, ie scale just the B component? Presumably this can be done without recomputing the entire matrix? If you show me the code, I can put it into a CMBLensing-like API, which will look like Cf(r=0.1) where Cf::BlockDiagEquiRect{QUAzFourier} (which matches how it currently looks for LambertFields via ParamDependentOp).

marius311 commented 3 years ago

Gonna merge when tests pass then continue future work on master / followup PRs. Thanks a ton for this!

marius311 / CMBLensing.jl

add equirectangular projection #64