Use unit support in Astropy models

[pllim]

Now that astropy/astropy#4855 is merged, we can refactor the code to use that feature. But that feature is only for Astropy v2 or greater.

Related:

• Enable unit support on all functional models astropy/astropy#6183
• https://gist.github.com/pllim/babea1d1337da665740439887c5b6f5f
• http://astropy.readthedocs.io/en/latest/modeling/units.html

[pllim]

I thought about making this change backward-compatible but it is not feasible. Not only it reduces performance but also increases code complexity. Since this package is still in "beta" version, I think it is okay to require Astropy version to be 2.0 or greater without much consequence.

c/c @bhilbert4

[pllim]

This now needs Astropy v3 to work due to astropy/astropy#6529

[pllim]

Let's wait for Perry's astropy.modeling refactor. This package will still work as-is for now, by passing unitless values to underlying models and wrap unit handling around those.

[lpsinger]

Where does this stand? I am trying to use my own model classes with Synphot but I am getting really weird behavior. Here is a minimal reproducer:

from astropy.modeling import Model
from astropy import units as u
import numpy as np
from synphot.units import FLAM
from synphot import SourceSpectrum


class Foobar(Model):

    n_inputs = 1
    n_outputs = 1
    input_units = {"x": u.nm}
    return_units = {"y": FLAM}
    input_units_equivalencies = {"x": u.spectral()}
    _input_units_strict = True

    waverange = [380, 450] * u.nm
    waveset = np.linspace(*waverange)

    @classmethod
    def evaluate(cls, x):
        lo, hi = cls.waverange
        return np.where((x >= lo) & (x <= hi), 1, 0) * FLAM

Evaluating the model itself works fine:

foobar = Foobar()
print(foobar(400 * u.nm))  # prints `1.0 FLAM`

However, if I wrap it in a SourceSpectrum so that I can use it with synphot, it blows up in my face. The following code:

foobar_spectrum = SourceSpectrum(foobar)
print(foobar_spectrum(400 * u.nm))

raises this exception:

Traceback (most recent call last):
  File "/Users/lpsinger/src/m4opt/test.py", line 30, in <module>
    print(foobar_spectrum(400 * u.nm))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/lpsinger/src/m4opt/.vscode/env/lib/python3.11/site-packages/synphot/spectrum.py", line 943, in __call__
    y = self.model(w.value) * self._internal_flux_unit
        ^^^^^^^^^^^^^^^^^^^
  File "/Users/lpsinger/src/m4opt/.vscode/env/lib/python3.11/site-packages/astropy/modeling/core.py", line 417, in __call__
    new_call = make_function_with_signature(
   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/lpsinger/src/m4opt/.vscode/env/lib/python3.11/site-packages/astropy/modeling/core.py", line 394, in __call__
    return super(cls, self).__call__(*inputs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/lpsinger/src/m4opt/.vscode/env/lib/python3.11/site-packages/astropy/modeling/core.py", line 1123, in __call__
    evaluate, inputs, broadcasted_shapes, kwargs = self._pre_evaluate(
                                                   ^^^^^^^^^^^^^^^^^^^
  File "/Users/lpsinger/src/m4opt/.vscode/env/lib/python3.11/site-packages/astropy/modeling/core.py", line 970, in _pre_evaluate
    inputs, broadcasted_shapes = self.prepare_inputs(*args, **kwargs)
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/lpsinger/src/m4opt/.vscode/env/lib/python3.11/site-packages/astropy/modeling/core.py", line 2093, in prepare_inputs
    inputs = self._validate_input_units(inputs, equivalencies, inputs_map)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/lpsinger/src/m4opt/.vscode/env/lib/python3.11/site-packages/astropy/modeling/core.py", line 2189, in _validate_input_units
    raise UnitsError(
astropy.units.core.UnitsError: Foobar: Units of input 'x', (dimensionless), could not be converted to required input units of nm (length)

[pllim]

Many years ago, I started a branch but ended up abandoning it because the diff got too big and too convoluted. At that time, unit support in modeling was very new and somewhat incomplete. Additionally, this package pre-dated that feature and thus has its own way to deal with spectral axis and flux units.

Now, this package is used in JWST ETC (last I heard), so I am hesitant to go down that path again as any breaking change might impact operations.

To use a custom model with this package currently, you basically need to throw away astropy way of unit support and probably subclass models from https://github.com/spacetelescope/synphot_refactor/blob/master/synphot/models.py

[pllim]

Your example looks like a box? There is a Box1D but I don't think it is used as a source, but rather more for bandpass.

[lpsinger]

Yes, I know. This is a contrived example. My real example is quite a bit more complicated, but the synphot way of doing units is way more painful than astropy's.

[lpsinger]

Now, this package is used in JWST ETC (last I heard), so I am hesitant to go down that path again as any breaking change might impact operations.

Wouldn't breaking changes just end up in a major version bump?

[pllim]

Wouldn't breaking changes just end up in a major version bump?

Theoretically, yes. In reality, I need to ask my stakeholders and supervisor if this is what they want to spend money on. 😬

The internals inherited from PySynphot is that all the important calculations are done in Angstrom and PHOTLAM (don't ask me why). So, this does not play very well with native unit support in modeling that does not have such assumption.

synphot way of doing units is way more painful than astropy's

More painful, maybe, but typically a user does not have to worry about that. Maybe give me a DM on Astropy Slack so I can better understand your use case? I apologize for any inconvenience caused.

[pllim]

Also, I am perplexed by how a Box1D defined in FLAM turns into a trapezoid going through this code, but since it is not your real use case, I don't want to waste time going into that rabbit hole right now.

from astropy import units as u

from synphot import SourceSpectrum
from synphot.units import FLAM
from synphot.models import Box1D

sp = SourceSpectrum(Box1D, amplitude=1 * FLAM, x_0=415 * u.nm, width=70 * u.nm)

sp.plot(flux_unit=FLAM)

[pllim]

A more hacky way is to turn your model into empirical model (a lookup table) like this:

https://github.com/spacetelescope/synphot_refactor/blob/55810c69979120bb45acc836332953bcf56455f3/synphot/spectrum.py#L1360C20-L1361

It is not analytical but it is very flexible.

Alternately, if specutils objects better suit your needs, we have a translator: https://synphot.readthedocs.io/en/latest/synphot/spectrum.html#specutils