Add threeml Gaussian and Powerlaw function with torch

[GallegoSav]

@scipascal Found that the nominal threeml function (like Gaussian or Band) which uses numba are much more slower than just using basic numpy method.

This is even faster with torch. So this pr add a Gaussian and Powerlaw function that uses torch.tensor to cosipy.

This is to use for unbinned analysis in order to improve the fitting time.

[read-the-docs-community]

Documentation build overview

📚 cositools-cosipy | 🛠️ Build #33480515 | 📁 Comparing 8138c9c against latest (f4f8741)

  🔍 Preview build  

14 files changed · ± 9 modified · - 5 deleted

± Modified

• genindex.html
• api/index.html
• api/polarization.html
• api/threeml.html
• tutorials/DataIO/DataIO_example.html
• tutorials/polarization/ASAD_method.html
• tutorials/polarization/Stokes_method.html
• tutorials/polarization/maximum_likelihood_method.html
• tutorials/source_injector/Point_source_injector.html

- Deleted

• py-modindex.html
• _modules/index.html
• _modules/cosipy/polarization/conventions.html
• _modules/cosipy/polarization/polarization_angle.html
• _modules/cosipy/polarization/polarization_axis.html

[codecov]

Codecov Report

❌ Patch coverage is 85.71429% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.73%. Comparing base (4f70ca2) to head (8138c9c).

Files with missing lines	Patch %	Lines
cosipy/threeml/ml/function_torch.py	85.71%	6 Missing ⚠️

Files with missing lines	Coverage Δ
cosipy/threeml/ml/function_torch.py	85.71% <85.71%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[israelmcmc]

Thanks @GallegoSav.

For the test, you'd have to add this guard:

    if not cosipy.with_ml:
        pytest.skip(reason="Optional [ml] dependencies not installed")

The _ml runner should be able to execute this and have full coverage.

I'm tagging @ndilalla and @omodei for awareness since it's 3ML related.

[jdbuhler]

A couple of things...

(1) This code needs to have test coverage. (This seems to be in progress).

(2) The power-law implementation explicitly sets the tensor device to the CPU. The Gaussian implementation does not specify a device. What model do you want to follow for where to run the computation if the user has both a CPU and a GPU? You could not specify a device at all and rely on the user to use torch.set_default_device() (though this incurs some extra overhead on every tensor construction), or you could pass in the device as an optional argument to the function object's constructor (defaulting to None to use whatever device is current).

Longer-term thoughts (that do not need to be addressed in this PR):

• If this is giving useful speedups on the CPU, I presume it is because it is implicitly distributed across cores. Ideally, Numba would provide the same benefit with parallel=True, but as we know, Numba unfortunately uses very slow transcendental functions (at least on x86).

• I'm not sufficiently familiar with PyTorch tensor arithmetic to know if there are implicit conversions back to scalar hiding in this code. Asking torch.compile() to JIT the function and looking at its graph might be a way to check this.

• This implementation uses tensor operations for individual arithmetic steps but still drives the computation from the CPU at the interpreter level. It might be preferable to use torch.compile on the whole function, or to use a Google JAX implementation that could avoid having to write explicit tensor ops. But this needs to be explored more thoroughly to see what is feasible and what gives the best performance.

[GallegoSav]

@jdbuhler thanks for your comments.

I added the unit tests and removed the explicit device. Then everything has to orient itself by default on the input tensor device set by the user.

Yes torch seems to be efficient for distributing tasks across cores. For the others long term comments I don't know

[jdbuhler]

Further minor points:

• The power law class moves the result back to a numpy array at the end of evaluation. But I think the Gaussian class returns a tensor? Is that your intent?

• Do you want to explicitly force float64 for your tensors? This is ok for the CPU, but if the code ever runs on a GPU, the speed of double-precision arithmetic is highly dependent on which GPU you use. A high-end NVidia datacenter GPU (eg, the A100 or H100) has dedicated FP64 units, but cheaper RTX-family devices run FP64 at 1/32 or 1/64 the speed of FP32.

• .ravel() might be a more readable alias for .view(-1).

Let me know your thoughts - if we resolve the first one, I can go ahead and approve.

[GallegoSav]

@jdbuhler for the first point, yes threeml expect that the source evaluation return an array , so we use tensor for the internal computation but return an array at the end. So I added this for the Gaussian as well. I also changed to cpu().numpy() which is apparently better for converting the tensor into an numpy array.

I explicitly set the devices to CPU per default but the user could change it to GPU. But then the float64 vs float32 question could be raise.

However the idea for cosipy is to use it with cpu as torch seems more efficient for multiprocessing

[jdbuhler]

Code changes for returning numpy array and passing in device look good. We'll have to revisit the float width issue more generally in the future -- we want to be consistent in not using FP64 on the GPU anywhere unless absolutely necessary. So AFAIK, the code is now fine...

[Update: it's failing some of the test cases, which is why the coverage is still 0%. Not sure what is happening there.]

[jdbuhler]

Take a look at the Band_Eflux class in threeml/custom_functions.py to see how to create custom class member variables in Astromodels-derived function classes. I believe init() is being overridden, and you have to use the _setup() hook instead.

But the better question is how you can pass additional arguments like devices to the class and have them be available to _setup(). I think you have to add a "properties" section to the docstring for the function, which is parsed to create the init procedure, to add a string parameter that does what you want. Do a search for "properties" in

https://astromodels.readthedocs.io/en/latest/notebooks/Functions_tutorial.html

Astromodels functions are deeply mysterious...

[GallegoSav]

@jdbuhler ok I fixed the unit test. Could you merge it ?

[jdbuhler]

Yay, fixed!

[israelmcmc]

Thanks @GallegoSav and @jdbuhler !