Add unit check for TOD operations

[mj-gomes]

This is an idea from @silviamich and me after we faced issues with units in tod operations.

We could force all tod computations (scan_map, fill_tod_with_hwp_harmonics, add_2f, add_non_linearity, etc.) to have compatible units. I am thinking of something similar (or exactly the same) as what is done in the map operations currently. This would avoid e.g. adding differential emissions in Kelvins to a TOD in spectral radiance units. Every operation changing the tod of an observation would start by passing a unit check.

We can discuss this better in the next meeting, meanwhile I will try to come up with a draft PR to do this.

[mj-gomes]

Would it be useful to have something like pint (possibly our own tailored version) that would associate every variable with a unit and automatically keep the dimensional coherence?

As we can see in some of the pint examples, if we multiply some variable which is, let's say in W.sr-1.m-2 with one in m2 it will return the result in W.sr-1, maintaining coherence. And if then we try to sum something with different units, it returns an error. Likewise, if we sum meters with centimeters, it does the automatic conversion and returns the result in meters.

Something like pint would be amazing to ensure correctness everywhere. For example, it took me some time to debug something which was basically due to a normalization factor that I had forgotten, which was making the units incompatible between two variables being summed, and the result was very incorrect. If we had something like pint, I would have gotten immediately an error message saying that I am not summing two things with the same units, because the normalization operation would automatically include the correspondent change in dimensions (I think this is really the most important part, as in the input sky module, for example, we have unit compatibility check, but we must define the units ourselves and I believe they do not update automatically with the operations). I believe this dimension coherence across all lbs would be very useful.

@ziotom78 @paganol, and others, what do you think? If you agree that this could be a good option, would you say it is better to develop our own pint-like module, or could we use pint (I believe pint allows us to define new units that may not yet be present in there, and it seems to be a pretty much alive and up-to-date package).

[ziotom78]

Thanks Miguel for having had a look at this!

I have often used unit system libraries in Julia, but in that case the type system delegates all the checks at compilation time so that there is no performance penalty. On the other hand, I have no experience at all with Python libraries.

I did a few tests with Pint and Unyt and posted my scripts and findings here: https://gist.github.com/ziotom78/82add22a92eec2e134bb2474e91ef5f4

We should discuss this at the next ST telecon.

[mj-gomes]

Hi @ziotom78, thanks for having tested this. The fact that it does not work with numba without unwrapping, makes it a kind of useless as we lose the dimensional analysis part.

At the same time, I am thinking that we may not need this for the functions that use numba. Usually, we use numba for the methods that scan the map, filling the tod, but these are basically rotation matrices and hwp matrices applied to the input maps, so the output dimensions are the same as the input ones.

The cases in which this would be very useful are e.g. adding the 2f differential effect, adding the non-linearity, etc, where we really need to be sure we are using coherent units. These are in general simpler computations for which we don't need to optimize that much, so maybe it is ok to lose a bit of time. Of course we would lose the capability to compute non linearities on-the-fly, for example, but it consists of only one multiplication per sample, it should not be super slow.

[mj-gomes]

Hi, I started a personal project that may fix this. It is a package, which I called for now pumnum. I thought about doing it independently from litebird_sim, because I found other people online facing the same issue, so I thought that I could build something that could be use anywhere.

It is yet very simple and may not work in complex cases, but the idea is to perform unit checks and dimensional operations in numba no-python jitted functions, by using decorators that basically first pre-compute the final units of a series of operations in a given function, then call the function with numba's njit, and in the end simply recreate the pint/unyt object by giving it the pre-computed units. This is done by running the function once in python mode (as opposed to numba's nopython mode) with unitary objects.

The idea is for the compatibility between pint, unyt, astropy.units etc with numba to be as easy as changing the name of the decorator from @ njit (without the space) to @ pumnum, so that we don't have to change anything else in the package.

You can see the basic usage in the tutorial notebook that I put on the repository.

@ziotom78, What do you think about this solution for lbs case?

[ziotom78]

Hi Miguel, thanks a lot for this effort! I really like the “tracing” idea behind pumnum, and I bet that many people will find this kind of decorator useful for smaller projects.

For LBS, we must keep in mind that we want to support it for the next decade. I have a few questions about the architecture, and I would like to hear your opinion:

• In general, I would prefer an explicit conversion whenever needer, rather than an automatic conversion. But perhaps I did not understand your idea correctly, and with “automatic” you were referring to something else.

• Many of our @njit functions (like _get_invnpp in mapmaker/destriper.py) operate on pre-allocated buffers created via np.empty(). Since these buffers starts as “unitless” raw memory, how would pumnum know what unit to assign the result? We might end up having to manually annotate every single low-level function call, which could become a significant maintenance burden given the dozens of functions of this kind we already have.

• If a Numba function performs a physical conversion in-place, like an in-place calibration converting Volts into K_CMB, the tracing logic might simply re-apply the input unit to the output data. In this cases, “automatic” unit labeling might actually hide a physical error rather than reveal it.

• As we have seen recently, relying on external dependencies that wrap “moving targets” like Numba is risky. If Numba changes its internal dispatcher, pumnum would need an update to keep LBS running. You are extremely talented, and we might risk a few years from now to lose you because after your PhD you find your dream job elsewhere. In that case, who would update the pumnum repository? (This problem happened recently, as you might remember…)

What if we implement a slightly simpler architecture? We might use unyt or pint to hold measurement units in Observation objects and check them at the entry points of major functions like make_destriped_map. After the validation and any necessary unit conversion, we strip the units (avoiding new allocations!) and use raw views of the data for the whole computation. Then, at the end we put back the measurement units to the return values.

This approach has the advantage that it does not require to manually fix each Numba function declaration and all their calls, and it would not cause any significant overhead.

What do you think? Perhaps we could do some tests on a dedicated branch for just one simple simulation module, to find the most feasible, readable, and performant way to implement this.

[mj-gomes]

Hi Maurizio, thanks for the answer, I will try to reply to everything below.

• In general, I would prefer an explicit conversion whenever needer, rather than an automatic conversion. But perhaps I did not understand your idea correctly, and with “automatic” you were referring to something else.

Right now, pumnum allows the user to select explicitly to which units he wants the conversion to be. I implemented this via an argument that we can pass to the decorator, where we define the unit system (SI, CGS, etc...). The default is no conversion at all so things like km*m may appear, for example.

• Many of our @njit functions (like _get_invnpp in mapmaker/destriper.py) operate on pre-allocated buffers created via np.empty(). Since these buffers starts as “unitless” raw memory, how would pumnum know what unit to assign the result? We might end up having to manually annotate every single low-level function call, which could become a significant maintenance burden given the dozens of functions of this kind we already have.

I f I understood correctly the point, it shouldn't change anything in the current functionality of pumnum: the empty buffer will get filled with some data that will have units associated, so when pumnum does the first "fast-pass" to check the units in the end, it will get those.

• If a Numba function performs a physical conversion in-place, like an in-place calibration converting Volts into K_CMB, the tracing logic might simply re-apply the input unit to the output data. In this cases, “automatic” unit labeling might actually hide a physical error rather than reveal it.

Again, maybe I misunderstood the question, but that in-place calibration step would also be present in the first pass where pumnum gets the final units, so the units would be correct in the need. The only case where this could be a problem, is if some of the steps, the calibration for example, would also take place by some condition which is input-dependent. In that case, the output units of a function would depend on the values we give to it as input, not only their units, and thus the first pass done in pumnum would not be enough to get the correct output units of the function. At the same time, I believe that this is a very rare case, as usually (I suppose) a function in scientific applications has well-defined output dimensions, that should not change depending on the magnitude of the inputs.

• As we have seen recently, relying on external dependencies that wrap “moving targets” like Numba is risky. If Numba changes its internal dispatcher, pumnum would need an update to keep LBS running. You are extremely talented, and we might risk a few years from now to lose you because after your PhD you find your dream job elsewhere. In that case, who would update the pumnum repository? (This problem happened recently, as you might remember…)

Of course. I will continue updating pumnum for now, but the future is not certain. However, for lbs, the version that I have now (with some updates that I did not yet merge that should allow to use it in functions that call other njit(pumnum) functions, as is often the case in lbs), is already almost everything that we need for lbs (I believe), so if it does work and does not make lbs much slower, we could perhaps integrate in lbs as a new lbs module, independent from the repository pumnum itself. And, in that case, it would no longer be an external dependency and would be treated and maintained just like any other module in lbs.

What if we implement a slightly simpler architecture? We might use unyt or pint to hold measurement units in Observation objects and check them at the entry points of major functions like make_destriped_map. After the validation and any necessary unit conversion, we strip the units (avoiding new allocations!) and use raw views of the data for the whole computation. Then, at the end we put back the measurement units to the return values.

I had thought about this option before starting the pumnum option, but I believe this does not take into account changes in the units inside the njitted function, right? This method assumes the output units are the same as the input units. If that becomes the case for all njitted functions in lbs, then of course pumnum is not needed at all. For now it is not the case at least for the hwp systematics case, since in this function even if we do the conversion from temperature to spectral radiance before (on the input maps) the integration in band itself changes the output units.

[ziotom78]

Thanks for your reply, Miguel. To make sure I really understood how pumnum works, I decided to install it and play with it. I did a quick test with _get_invnpp to see how we could adapt our codes to pumnum, but the code crashed with a TypeError because pumnum currently only accepts unyt_quantity (scalars), not unyt_array.

Here is the code I wrote to test pumnum:

def test_pumnum():
    npix = 100_000
    nobs = unyt.unyt_array(np.random.rand(npix, 6), 1 / unyt.K)

    valid_pixel = unyt.unyt_array(
        np.array([True] * npix, dtype=bool), unyt.dimensionless
    )
    result = unyt.unyt_array(np.empty((npix, 3, 3), dtype=np.float64), unyt.K**2)

    final = _get_invnpp_pumnum(nobs, valid_pixel, result)
    print(final.units)

(I just copied-and-pasted _get_invnpp in my script, modified it to return result since pumnum needs a return value to attach units to, and added @pumnum(parallel=True) at the beginning.) Here is the error I got:

Traceback (most recent call last):
  File "/home/tomasi/tmp2/miguel_test/main.py", line 31, in <module>
    main_pumnum()
    ~~~~~~~~~~~^^
  File "/home/tomasi/tmp2/miguel_test/main.py", line 25, in main_pumnum
    final = _get_invnpp_pumnum(nobs, valid_pixel, result)
  File "/home/tomasi/tmp2/miguel_test/.venv/lib/python3.14/site-packages/pumnum.py", line 75, in __init__
    raise TypeError(
    ...<2 lines>...
    )
TypeError: object of type unyt_array is not compatible with pumnum; it should be either a pint.Quantity, unyt.array.unyt_quantity, or astropy.units.Quantity

Judging from lines 68-78, it seems that currently pumnum only accepts scalars (this is the only case tested in test_parallel.py).

I had a look at the function you mentioned, integrate_inband_signal_for_one_sample. It seems to me that this function wouldn’t work either with the current version of pumnum, as it requires arrays.

I had thought about this option before starting the pumnum option, but I believe this does not take into account changes in the units inside the njitted function, right? This method assumes the output units are the same as the input units. If that becomes the case for all njitted functions in lbs, then of course pumnum is not needed at all. For now it is not the case at least for the hwp systematics case, since in this function even if we do the conversion from temperature to spectral radiance before (on the input maps) the integration in band itself changes the output units.

No, this is not what I was thinking. In fact, I agree that the output units are often different from the input units. (I mentioned this case when I talked aboue “in-place calibration converting Volts into K_CMB”.) My “high-level” approach just requires unit changes to be explicit.

Consider integrate_inband_signal_for_one_sample again: in line #119 there is a * 1e9 factor to convert GHz to Hz. If pumnum tries to trace units automatically, it won’t be able to understand the meaning of this scaling factor and could easily return the wrong measurement unit metadata. But with a high-level approach, the Python function calling integrate_inband_signal_for_one_sample would strip the units, pass the raw arrays to Numbe, and then explicitly wrap the result into unyt_array(raw_result, input_units * unyt.Hz). This has the advantage to keep the Numba function untouched and at peak speed, while explicitly documenting the output measurement unit right in the code. I believe this improves both safety and readability.

Finally, I realized there is another problem with units we might have to face. Polynomial fitting/evaluation formulae like numpy.polyfit model polynomials by storing their coefficients in a NumPy array.

Suppose we want to implement a high-pass filter of the form $x_j \rightarrow x_j - \sum_{k=0}^N a_k x_j^k$, where we subtract a polynomial of degree $N$ from the data $x_j$ in the TOD. If $x_j$ is expressed in K_CMB, then each $a_k$ will have its own measurement unit ($a_0$ is in K_CMB, $a_1$ is dimensionless, $a_2$ is in K_CMB⁻¹, etc.). Because unit libraries like Pint and Unyt require a single unit for a whole array, pumnum would really struggle with this unless we fall back to slow object arrays. On the other hand, enforcing measurement units only at the boundary of high-level functions is totally transparent: the inner logic gets a raw array of coefficients and evaluates the polynomial seamlessly.

[mj-gomes]

Judging from lines 68-78, it seems that currently pumnum only accepts scalars (this is the only case tested in test_parallel.py).

Yes, that is indeed the case for now. I had not yet thought about arrays, but of course it's of major importance. That would be one of the next steps in pumnum.

Consider integrate_inband_signal_for_one_sample again: in line #119 there is a * 1e9 factor to convert GHz to Hz. If pumnum tries to trace units automatically, it won’t be able to understand the meaning of this scaling factor and could easily return the wrong measurement unit metadata. But with a high-level approach, the Python function calling integrate_inband_signal_for_one_sample would strip the units, pass the raw arrays to Numbe, and then explicitly wrap the result into unyt_array(raw_result, input_units * unyt.Hz). This has the advantage to keep the Numba function untouched and at peak speed, while explicitly documenting the output measurement unit right in the code. I believe this improves both safety and readability.

Ok, I am convinced with this one 😆 It could technically be solvable by simply changing the explicit conversions from *1e9 to a conversion method from the unit library being used, but that would, once again, not work with njit.

Because unit libraries like Pint and Unyt require a single unit for a whole array, pumnum would really struggle with this unless we fall back to slow object arrays. On the other hand, enforcing measurement units only at the boundary of high-level functions is totally transparent: the inner logic gets a raw array of coefficients and evaluates the polynomial seamlessly.

I was not aware that pint and unyt required single unit for the whole array.

I would say let's go with your suggestion then, I am out of arguments to prefer pumnum. I will keep working a bit on it, since it can be useful for some cases, but for now I will start testing an implementation of your suggestion in lbs.

Thanks for the discussion as always!