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:

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

[ziotom78]

@paganol , @sgiardie , @nraffuzz , @yusuke-takase , @giuspugl , @ggalloni I have implemented the new API and fixed a number of issues in the code. The new API works as follows:

1. You use Simulation.set_instrument, Simulation.set_scanning_strategy, Simulation.set_hwp as before: nothing changes for these methods
2. Once you have created observations with Simulation.create_observations(), you must call Simulation.prepare_pointings(). This was the old compute_pointings(), but I renamed it as in this new incarnation it just prepares the quaternions that will be used by Observation.get_pointings().
3. Every object Observation has a method named get_pointings(IDX), which takes the zero-based index of the detector and returns a tuple of two arrays:
   1. A N×3 matrix containing the pointing angles θ, φ, ψ (in radians, as usual);
   2. A vector of N elements containing the angle of the HWP.

Thus, to use the new API you must write something like this:

# This is the same as usual…
sim = Simulation(…)

sim.set_instrument(…)
sim.set_scanning_strategy(…)
sim.set_hwp(…)

sim.create_observations(…)

# …until here! The call to `sim.prepare_pointings()` substitutes
# `sim.compute_pointings()`, and no pointings are actually computed
sim.prepare_pointings()

for cur_obs in sim.observations:
    for det_idx in range(cur_obs.n_detectors):
        # Now we need to access the pointings, which are computed on the fly.
        # Note that `Observation.get_pointings` returns both the N×3 pointing
        # matrix AND the HWP angle, which are now kept separated.
        pointings, hwp_angle = cur_obs.get_pointings(det_idx)
        theta, phi, psi = pointings[:, 0], pointings[:, 1], pointings[:, 2]

        # Use the pointings

I have made numerous changes to the Python code and feel quite confident in what I have done so far. However, I am concerned that I have been working on it in isolation for too long without receiving feedback. Could you have a look at the code?

I am also seeking assistance in addressing the failing tests, particularly those tied to the binner and the destriper, the dipole module, the scan_map module, and the hwp_sys.py file. Updating the tests to align with the new API may be challenging in some instances, and I'm not sure I'm the most adept at this task. In particular, the scan_map and mapmaking codes now need to properly handle the pointings and the HWP separately.

[ziotom78]

(BTW @mreineck , I have still to study your implementation of the periodic slerp in ducc, hope to have a look at it in the next days…)

[paganol]

Hi @ziotom78. Fantastic!! Thanks!! I can help in the debugging and fixing test. I started looking into test_dipole.py, and I have immediately a couple of comments. I think having a function like lbs.prepare_pointings(obs,...) could be useful. Previously the code had two running modes, one using the high level methods of the class Simulations, and the other one using the low level functions lbs.xxx. Can we add a function that tries to mimic the old get_pointings running mode? Before we had:

pointings = lbs.get_pointings(
    obs,
    spin2ecliptic_quats=sim.spin2ecliptic_quats,
    bore2spin_quat=sim.instrument.bore2spin_quat,
)

Now it could be:

lbs.prepare_pointings(
    obs,
    spin2ecliptic_quats=sim.spin2ecliptic_quats,
    bore2spin_quat=sim.instrument.bore2spin_quat,
    hwp=self.hwp,
)
pointings, hwp_angle = obs.get_pointings(0)

Also obs.get_pointings should have two extra features in my opinion:

• if detector_idx is None it should return the pointing for all the detectors in the observation, returning pointings in buffer (obs.n_detectors, obs.n_samples, 3)
• it sould have an extra option for storing the pointing in obs (as we had in the past with store_pointings_in_obs)

What do you think?

[paganol]

Hi @ziotom78 , I'm playing a bit with the PR. I have a couple of quick comments:

obs.get_pointings() returns different types if I ask one detector or all the detectors
See this

VS

Do we want to specify the type of pointings in prepare_pointings?

I made a quick test for checking the speed, and the new implementation seems to be faster Old

New

[ziotom78]

Ok, I added a keyword to specify the data type of the pointings, both to get_pointings() (used when pointings are computed on the fly) and to prepare_pointings() (used when pointings are pre-computed).

[mreineck]

@ziotom78 , @paganol please let me know if you think that ducc should support single -precision pointings natively (at least on the output side)! This is not much work to implement, I was just a bit reluctant to provide this because single-precision pointings can be deceptively inaccurate in some situations if no used carefully.

[mreineck]

I can (and probably should) also add support for writing the results directly into a preallocated array (via an optional out argument). This could save further temporary allocations.

[ziotom78]

I can (and probably should) also add support for writing the results directly into a preallocated array (via an optional out argument). This could save further temporary allocations.

Hi Martin, this is not super-urgent at the time, as we only use ducc to compute the high-frequency quaternions, which are handled with 64-bit floats; the conversion to pointing angles is done by an internal Numba routine, and it is already able to work with both 32-bit and 64-bit floats.

[ziotom78]

I implemented @paganol 's idea to split the preparation of the quaternions and the (optional) pre-computation of the pointings in two distinct functions prepare_pointings() and precompute_pointings(). I have also split the high-level wrapper Simulation.prepare_pointings() and have created a similar Simulation.precompute_pointings().

With the last commits, I have updated the documentation as well. It's now time to work on the modules that use the pointings, and for this we are looking for volunteers. Please, drop a note here if you are interested and say on which module you would like to work.

[paganol]

Hi @ziotom78, I can work on dipole.py and scan_map.py, and the relative tests.

[nraffuzz]

Hi @ziotom78, @sgiardie and I will take care of the hwp part

[yusuke-takase]

Hi @ziotom78,

I tried to use the new API, then I met a weird behavior. When I simulate pointings for two different detectors, it returns same theta and phi. (Only psi is not same.) I put a notebook on my gist to reproduce the situation. What I did are:

1. Clone the smart_pointing branch.
2. Create new environment and install it. (lbsim v0.12.0)
3. Test the notebook.

Basically, the code in the notebook follows the notebook in the litebird_sim example. I changed the detector name to avoid to select top and bottom detector because they have same quaternion. I think it is not due to the new API but I hope here is one of the appropriate places to discuss it.

[yusuke-takase]

The problem of pointing matching with different detectors, but the cause was IMO. I ran the first test with imo-vPTEP, but when I ran the same code using imo-v1.3 I observed that the loaded quaternion changed as follows:

# imo-vPTEP
Name: 000_000_003_QA_040_T  |  Quats: [0.         0.         0.99625809 0.08642811]
Name: 000_000_005_UA_040_T  |  Quats: [0.         0.         0.92225982 0.3865706 ]

# imo-v1.3
Name: 000_000_003_QA_040_T  |  Quats: [0.0238787  0.07619702 0.         0.99680681]
Name: 000_000_005_UA_040_T  |  Quats: [0.04050428 0.03790323 0.38209417 0.92245693]

I don't know where the appropriate place to point out the content of the IMO in vPTEP is, but the code that calculates the pointing seems to be fine for the moment.

[yusuke-takase]

Hi @ziotom, I could play new API eventually, I can work for the task: Provide an handy way to create RotQuaternion objects that encode periodic fluctuations around the z axis.

[nraffuzz]

Hi, @sgiardie and I "completed" the part on hwp_sys, however we plan to further update the example notebook because of its dependece on lbs binned map-maker. At the moment the cells regarding lbs.make_binned_maps() are NOT working because the map-making part about pointings is still missing the update.

One small issue that we had is just about coordinate transformation from ecl to gal: at line 198 of common.py there's curr_pointings_det, polang_all[idet] = rotate_coordinates_e2g(...) but rotate_coordinates_e2g returns only 1 object while 2 were given.

[paganol]

So, scan_map.py, dipole.py, binner.py have been fixed, and also the relevant tests are ok now. We are still failing the tests for coordinates, destriper and madam. test_coordinates.py and test_madam.py should be easy to fix. For test_destriper.py we should wait for PR #309. Then we should start working on the documentation.

[paganol]

store_full_pointings and pointings_dtype should be removed from sim.prepare_pointings? https://github.com/litebird/litebird_sim/blob/a3583c7ab0c0116d38e8f5640cd2b2f63a955295/litebird_sim/simulations.py#L1272-L1278

[ggalloni]

Also destriper.py has been fixed now.

All tests run smoothly now (exception made for TOAST breaking on macos, of course).

[yusuke-takase]

Hi @ziotom78,

I'm making a notebook to explain pointing systematics function as a task of "Examples that show how to encode time-varying quaternions". I thought it is good to use multiple detectors on a focal plane which is given by the public IMO (vPTEP). But due to the reason of the problem, I could not make a focal plane.

Are these IMO quaternions intentionally changed? Or is this something that should be fixed? If the former, I think it would constitute a virtual focal plane to show an example of pointing systematics.

[paganol]

Hi @yusuke-takase, in the imo-vPTEP we put on purpose all the detectors at the boresight. This was done in order not to expose potentially sensitive information to those who are not in the LiteBIRD collaboration, but still have a reasonable object to use in tests and notebooks.

[yusuke-takase]

Thank you for telling me, @paganol. Okay, I understood the reason, then I will go the way to make a mock focal plane.

[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.