Implement time-dependent quaternions and save memory when computing pointings

[ziotom78]

This PR is a WIP to implement the ideas in discussion #312. It contains several changes to the code base to implement the following features:

• A new type, RotQuaternion, which encodes a rotation quaternion that can vary in time
• Make get_pointings work with generic time-varying quaternions
• Implement Observation.get_pointings so that it always requires the index of the detector for which pointings are requested
• Modify the class HWP so that it accepts a method apply_hwp_to_pointings to modify the pointings (useful to simulate some families of HWP systematics)
• Implement a class PointingProvider that computes pointings on the fly
• Create HDF5 I/O functions for RotMatrix objects
• Create a way to save the definition of a generic HWP into a HDF5 file (needed to encode the way HWP angles are calculated by Observation.get_pointings)
• Create HDF5 I/O functions for PointingProvider objects
• Update all the tests in test/test_io.py to check that quaternions and pointings are saved correctly
• Separate Simulation.prepare_pointings() into a high-level method in Simulation and a low-level function prepare_pointings()
• Add an option to store the full pointings once prepare_pointings() is called
• Add the possibility to specify an arbitrary list of detectors to Observation.get_pointings()
• Support both Ecliptic and Galactic coordinates
• Update the code in hwp_sys.py to make use of time-varying quaternions
• Update the code in the analysis modules so that the smarter API is used:
  • Binner
  • Destriper
  • scan_map.py
  • dipole.py
• Provide an handy way to create RotQuaternion objects that encode periodic fluctuations around the z axis
• Improve in the documentation related to the simulation of pointings:
  • Fixes in the terminology of the existing functions
  • Examples that show how to encode time-varying quaternions
• Update the documentation to explain how to compute pointings on the fly vs how to pre-compute them
• Update the example notebooks

The type RotQuaternion encodes a time-varying quaternion that is normalized to 1, i.e., it encodes a time-dependent rotation. The rotation is saved internally in a (N, 4) matrix, where N is the number of quaternions that are sampled evenly. The type RotQuaternion has a field start_time and a field sampling_rate_hz that encode the time t0 and the number of quaternions per second encoded in the matrix; thus, the sample cover the time range from t0 to t0 + N / sampling_rate_hz. You can create it by passing a (N, 4) matrix:

quat = RotQuaternion(
    quats = np.array([
        [a, b, c, d],  # Time t0
        [e, f, g, h],  # Time t0 + δt
        [i, j, k, l],  # Time t0 + 2δt
        …
    ]),
    start_time=astropy.time.Time("2024-04-29T00:00:00"),
    sampling_rate_hz=1.0 / 19.0,
)

There is one significant breaking change in the API. In the current API for get_pointings, the function requires a matrix of quaternions for the parameter detector_quats to specify the rotation to be applied to the boresight of each detector to convert its reference frame into the RF of the focal plane; this matrix has the shape (N_det, 4), with N_det the number of detectors. However, now that the code supports time-varying quaternions, the shape (N, 4) is used to encode the value of the rotation quaternion as time passes for one detector.

This means that when one spots this matrix in the parameter detector_quats

[
    [a, b, c, d],
    [e, f, g, h],
]

it could either be interpreted as

1. Two quaternions that do not vary with time and are associated to two detectors:

   [
       [a, b, c, d],  # Det #1
       [e, f, g, h],  # Det #2
   ]

2. One time-varying quaternion sampled at two times t0 and t1 that is associated to one detector:

   [
       [a, b, c, d],  # Det #1 at time t = t0
       [e, f, g, h],  # Det #1 (the same) at time t = t0 + δt
   ]

To prevent ambiguities, in this PR I have implemented the following breaking change to the API: now you are expected to always pass a list of RotQuaternion objects to the parameter detector_quats:

pointings = get_pointings(
    …,
    # Two quaternions
    detector_quats=[
        RotQuaternion(np.array([
            [0., 0.70710678, 0., 0.70710678],
        ])),
        RotQuaternion(np.array([
            [0.40824829, 0., 0.81649658, 0.40824829],
        ])),
    ],
)

I played a bit with some “smart” conversions to avoid this breaking change, but I disliked the result because it seemed to me that they looked confusing and made the code of RotQuaternion.__init__() quite complex.

Note that RotMatrix has a default constructor that creates the identity rotation (constant in time), so in the simplest cases the new API is quite readable:

q1 = RotQuaternion(np.array([0.0, 0.0, 0.0, 1.0]))  # Identity transformation
q2 = RotQuaternion()  # Same as q1

[ziotom78]

The tests are currently not passing because of the fact that hwp_sys.py obviously needs to access the detector quaternions directly. @nraffuzz , this should be fixed before merging this PR, but I believe it's wise to wait for the implementation of the new “smart” API for get_pointings (the ones that returns just the pointings for one detector) before fixing it. For the moment, I would leave the failing test as is.

[ziotom78]

@paganol , @giuspugl , @sgiardie , @ggalloni , @mreineck I'm implementing the PointingProvider class, as specified in #312 , but I have a few doubts about the implementation:

1. The old API kept pointings in two variables: one is a (N, 2) matrix that holds θ and φ, while another stores ψ. However, now that buffers need to be passed much more often, it's easier if we store θ, φ, and ψ in a (N, 3) matrix. This should not be too burdensome to fix in the existing code, what do you think?

2. Currently, when we save observations to HDF5 files we store the pointings as well. It's easy to do this with the new API, but what about reading HDF5 files, now that we always compute them on the fly?

   • Should the code just load the “slow” quaternions from the HDF5 file?
   • Or should it instead load the pointings and forget about the slow quaternions?
   • Or should it load the slow quaternions, compute the pointings, and then check that the result with the actual pointings read from the HDF5 and warn the user if it is not so?

Probably the biggest problem I have is that I do not know if we really need to load Observation objects from HDF5 files back into LBS. (The only use I'm aware of is to export them so that they can used by Commander, but in this case we just need to save HDF5 files, not to load them.) Has the HDF5 loading feature been used in the past by any of us?

Implementing any of these solutions but the first one has a number of drawbacks, and I would like to be sure that this is really worth the effort.

[mreineck]

I agree that the switch to the (N, 3) format should not be a big deal, especially since this can be emulated with views if the old array shapes are needed in any components. Doin it the other way round would be more problematic.

Concerning the reading: if we have indeed not really read the observations yet, I'd vote for:

• the slow quaternions are the only "real" pointings
• if pointings exist, they are only for software that cannot deal with slow quaternions (like current Commander, but I expect this will change at some point)

[paganol]

1. The old API kept pointings in two variables: one is a (N, 2) matrix that holds θ and φ, while another stores ψ. However, now that buffers need to be passed much more often, it's easier if we store θ, φ, and ψ in a (N, 3) matrix. This should not be too burdensome to fix in the existing code, what do you think?

I agree on this, and this was what we had at the begining. The reason why we moved to two different variables was to handle more efficiently the coordinate rotation. For example in dipole.py we do not rotate polarization. But we can easly find a workaorund here.

2. Currently, when we save observations to HDF5 files we store the pointings as well. It's easy to do this with the new API, but what about reading HDF5 files, now that we always compute them on the fly?

   • Should the code just load the “slow” quaternions from the HDF5 file?
   • Or should it instead load the pointings and forget about the slow quaternions?
   • Or should it load the slow quaternions, compute the pointings, and then check that the result with the actual pointings read from the HDF5 and warn the user if it is not so?

Probably the biggest problem I have is that I do not know if we really need to load Observation objects from HDF5 files back into LBS. (The only use I'm aware of is to export them so that they can used by Commander, but in this case we just need to save HDF5 files, not to load them.) Has the HDF5 loading feature been used in the past by any of us?

Implementing any of these solutions but the first one has a number of drawbacks, and I would like to be sure that this is really worth the effort.

For this I would go for the first option, only "slow" quaternions. Already with the new sims saving all the pointing is kind of unfeasable. Codes should be able to deal with quaternions by themselves, or import litebird_sim for computing pointings.
On my side I worked a bit with the HDF5 loading feature. The purpouse was to test map-makings, and this is tipically what we want to do with save timelines, potentialy process the timelines and do the map-making. So I guess that if we save the TODs and "slow" quaternions, when, for example, we run the map-making the code computes on the fly the pointing, and so could do any reporcessing that needs the pointing.

[paganol]

Hi @ziotom78, everything seems ok, but I'm still getting stupid syntax issues that are not detected locally. But, a part from that, I think we are ready for merging and the a release.

[anand-avinash]

I have fixed the syntax issue and changed the github action image for macos to macos-13 (macos-latest is arm based, macos-latest-large is x86 based but available only to enterprise account. Apologies!). Now the tests seem to be passing on my branch.