Display units: model fitting

[kecnry]

Description

This pull request updates the model fitting plugin to be aware of display units and to handle the compatibility/incompatibility between model components and the current display units.

Screen.Recording.2023-05-25.at.11.43.34.AM.mov

Changelog will point to #2127 so is ignored from CI here.

Closes #2026

Change log entry

• Is a change log needed? If yes, is it added to CHANGES.rst? If you want to avoid merge conflicts,
  list the proposed change log here for review and add to CHANGES.rst before merge. If no, maintainer
  should add a no-changelog-entry-needed label.

Checklist for package maintainer(s)

This checklist is meant to remind the package maintainer(s) who will review this pull request of some common things to look for. This list is not exhaustive.

• Are two approvals required? Branch protection rule does not check for the second approval. If a second approval is not necessary, please apply the trivial label.
• Do the proposed changes actually accomplish desired goals? Also manually run the affected example notebooks, if necessary.
• Do the proposed changes follow the STScI Style Guides?
• Are tests added/updated as required? If so, do they follow the STScI Style Guides?
• Are docs added/updated as required? If so, do they follow the STScI Style Guides?
• Did the CI pass? If not, are the failures related?
• Is a milestone set? Set this to bugfix milestone if this is a bug fix and needs to be released ASAP; otherwise, set this to the next major release milestone.
• After merge, any internal documentations need updating (e.g., JIRA, Innerspace)?

[codecov]

Codecov Report

Patch coverage: 98.82% and project coverage change: +0.05 🎉

Comparison is base (75d542e) 91.49% compared to head (1aaf310) 91.55%.

Additional details and impacted files

@@                  Coverage Diff                  @@
##           dev-display-units    #2216      +/-   ##
=====================================================
+ Coverage              91.49%   91.55%   +0.05%     
=====================================================
  Files                    148      148              
  Lines                  16576    16657      +81     
=====================================================
+ Hits                   15167    15250      +83     
+ Misses                  1409     1407       -2     

Impacted Files	Coverage Δ
jdaviz/core/helpers.py	87.67% <ø> (ø)
...igs/default/plugins/model_fitting/model_fitting.py	86.15% <97.91%> (+1.15%)	⬆️
...efault/plugins/model_fitting/tests/test_fitting.py	99.56% <100.00%> (+0.08%)	⬆️

... and 1 file with indirect coverage changes

☔ View full report in Codecov by Sentry.
📢 Do you have feedback about the report comment? Let us know in this issue.

[rosteen]

Working on testing this now. The only thing I've noticed so far is that, after changing units, I chose to re-estimate the model parameters on my linear and gaussian components before refitting. After that, running the fit updated the linear model parameters successfully, but I seemed to have gotten into a state where the gaussian parameters wouldn't update, and the model would never successfully fit the feature. After deleting and re-adding the gaussian component, I was able to fit the feature successfully. I'll try to replicate this in a screen recording soon.

Edit: I didn't have the components checked to be fixed, and G was still in the model equation. The second time I tried reproducing, the gaussian parameters updated once on the first button click (but the fit was bad), and then after I edited them manually they no longer updated on fitting.

[rosteen]

Alright, after further testing I've mostly convinced myself that the bug reported above was some combination of user error and bad initial guesses resulting in a poor fit. However, I've run into another bug that I think is real. If you convert the flux from Jy to something like erg / (Angstrom s cm2) and then try to use the "Re-estimate free parameters" button under the model component, you get the following error trace:

UnitConversionError                       Traceback (most recent call last)
File ~/opt/anaconda3/envs/viz_dev/lib/python3.11/site-packages/ipyvue/VueTemplateWidget.py:60, in Events._handle_event(self, _, content, buffers)
     58     getattr(self, "vue_" + event)(data, buffers)
     59 else:
---> 60     getattr(self, "vue_" + event)(data)

File ~/projects/jdaviz/jdaviz/configs/default/plugins/model_fitting/model_fitting.py:636, in ModelFitting.vue_reestimate_model_parameters(self, model_component_label, **kwargs)
    635 def vue_reestimate_model_parameters(self, model_component_label=None, **kwargs):
--> 636     self.reestimate_model_parameters(model_component_label=model_component_label)

File ~/projects/jdaviz/jdaviz/configs/default/plugins/model_fitting/model_fitting.py:662, in ModelFitting.reestimate_model_parameters(self, model_component_label)
    659 # store user-fixed parameters so we can revert after re-initializing
    660 fixed_params = {p['name']: p for p in model_comp['parameters'] if p['fixed']}
--> 662 new_model = self._initialize_model_component(model_comp['model_type'],
    663                                              model_component_label,
    664                                              poly_order=model_comp['model_kwargs'].get('degree', None))  # noqa
    666 # revert fixed parameters to user-value
    667 new_model['parameters'] = [fixed_params.get(p['name'], p) for p in new_model['parameters']]

File ~/projects/jdaviz/jdaviz/configs/default/plugins/model_fitting/model_fitting.py:491, in ModelFitting._initialize_model_component(self, model_comp, comp_label, poly_order)
    487         initial_val = default_param.quantity.to(default_units)
    489     initial_values[param_name] = initial_val
--> 491 masked_spectrum = self._apply_subset_masks(self._get_1d_spectrum(), self.spectral_subset)
    492 mask = masked_spectrum.mask
    493 initialized_model = initialize(
    494     MODELS[model_comp](name=comp_label,
    495                        **initial_values,
    496                        **new_model.get("model_kwargs", {})),
    497     masked_spectrum.spectral_axis[~mask] if mask is not None else masked_spectrum.spectral_axis,  # noqa
    498     masked_spectrum.flux[~mask] if mask is not None else masked_spectrum.flux)

File ~/projects/jdaviz/jdaviz/configs/default/plugins/model_fitting/model_fitting.py:315, in ModelFitting._get_1d_spectrum(self)
    313 def _get_1d_spectrum(self):
    314     # retrieves the 1d spectrum (accounting for spatial subset for cubeviz, if necessary)
--> 315     return self.dataset.selected_spectrum_for_spatial_subset(self.spatial_subset_selected)

File ~/projects/jdaviz/jdaviz/core/template_mixin.py:1543, in DatasetSelect.selected_spectrum_for_spatial_subset(self, spatial_subset, use_display_units)
   1541 if spatial_subset == SPATIAL_DEFAULT_TEXT:
   1542     spatial_subset = None
-> 1543 return self.plugin._specviz_helper.get_data(data_label=self.selected,
   1544                                             spatial_subset=spatial_subset,
   1545                                             use_display_units=use_display_units)

File ~/projects/jdaviz/jdaviz/configs/specviz/helper.py:333, in Specviz.get_data(self, data_label, spectral_subset, cls, use_display_units, **kwargs)
    330     spatial_subset = None
    331     function = None
--> 333 return self._get_data(data_label=data_label, spatial_subset=spatial_subset,
    334                       spectral_subset=spectral_subset, function=function,
    335                       cls=cls, use_display_units=use_display_units)

File ~/projects/jdaviz/jdaviz/core/helpers.py:497, in ConfigHelper._get_data(self, data_label, spatial_subset, spectral_subset, mask_subset, function, cls, use_display_units)
    494     else:
    495         data = data.get_object(cls=cls, **object_kwargs)
--> 497     return _handle_display_units(data, use_display_units)
    499 if not cls and spatial_subset:
    500     raise AttributeError(f"A valid cls must be provided to"
    501                          f" apply subset {spatial_subset} to data. "
    502                          f"Instead, {cls} was given.")

File ~/projects/jdaviz/jdaviz/core/helpers.py:434, in ConfigHelper._get_data.<locals>._handle_display_units(data, use_display_units)
    429     # TODO: any other attributes (meta, wcs, etc)?
    430     # TODO: implement uncertainty.to upstream
    431     new_uncert = data.uncertainty.__class__(data.uncertainty.quantity.to(flux_unit)) if data.uncertainty is not None else None  # noqa
    432     data = Spectrum1D(spectral_axis=data.spectral_axis.to(spectral_unit,
    433                                                           u.spectral()),
--> 434                       flux=data.flux.to(flux_unit),
    435                       uncertainty=new_uncert)
    436 else:  # pragma: nocover
    437     raise NotImplementedError(f"converting {data.__class__.__name__} to display units is not supported")  # noqa

File ~/opt/anaconda3/envs/viz_dev/lib/python3.11/site-packages/astropy/units/quantity.py:935, in Quantity.to(self, unit, equivalencies, copy)
    931 unit = Unit(unit)
    932 if copy:
    933     # Avoid using to_value to ensure that we make a copy. We also
    934     # don't want to slow down this method (esp. the scalar case).
--> 935     value = self._to_value(unit, equivalencies)
    936 else:
    937     # to_value only copies if necessary
    938     value = self.to_value(unit, equivalencies)

File ~/opt/anaconda3/envs/viz_dev/lib/python3.11/site-packages/astropy/units/quantity.py:888, in Quantity._to_value(self, unit, equivalencies)
    885     equivalencies = self._equivalencies
    886 if not self.dtype.names or isinstance(self.unit, StructuredUnit):
    887     # Standard path, let unit to do work.
--> 888     return self.unit.to(
    889         unit, self.view(np.ndarray), equivalencies=equivalencies
    890     )
    892 else:
    893     # The .to() method of a simple unit cannot convert a structured
    894     # dtype, so we work around it, by recursing.
    895     # TODO: deprecate this?
    896     # Convert simple to Structured on initialization?
    897     result = np.empty_like(self.view(np.ndarray))

File ~/opt/anaconda3/envs/viz_dev/lib/python3.11/site-packages/astropy/units/core.py:1195, in UnitBase.to(self, other, value, equivalencies)
   1193     return UNITY
   1194 else:
-> 1195     return self._get_converter(Unit(other), equivalencies)(value)

File ~/opt/anaconda3/envs/viz_dev/lib/python3.11/site-packages/astropy/units/core.py:1124, in UnitBase._get_converter(self, other, equivalencies)
   1121             else:
   1122                 return lambda v: b(converter(v))
-> 1124 raise exc

File ~/opt/anaconda3/envs/viz_dev/lib/python3.11/site-packages/astropy/units/core.py:1107, in UnitBase._get_converter(self, other, equivalencies)
   1105 # if that doesn't work, maybe we can do it with equivalencies?
   1106 try:
-> 1107     return self._apply_equivalencies(
   1108         self, other, self._normalize_equivalencies(equivalencies)
   1109     )
   1110 except UnitsError as exc:
   1111     # Last hope: maybe other knows how to do it?
   1112     # We assume the equivalencies have the unit itself as first item.
   1113     # TODO: maybe better for other to have a `_back_converter` method?
   1114     if hasattr(other, "equivalencies"):

File ~/opt/anaconda3/envs/viz_dev/lib/python3.11/site-packages/astropy/units/core.py:1085, in UnitBase._apply_equivalencies(self, unit, other, equivalencies)
   1082 unit_str = get_err_str(unit)
   1083 other_str = get_err_str(other)
-> 1085 raise UnitConversionError(f"{unit_str} and {other_str} are not convertible")

UnitConversionError: 'Jy' (spectral flux density) and 'erg / (Angstrom s cm2)' (power density/spectral flux density wav) are not convertible

Looks like the spectral density equivalency needs to be specified somewhere.

[cshanahan1]

Ok I took an initial look at this, following along with the screen recording you posted and also just playing around with it.

What is the reasoning behind needing to manually re-fit when you change units? If the units are compatible, could the model components not just be tied to the X axis and automatically update?

My only other comment so far is that I think maybe the equation editor could be more relaxed about whitespace? 'L + L_1' is invalid, only 'L+L1' is accepted.

[kecnry]

What is the reasoning behind needing to manually re-fit when you change units? If the units are compatible, could the model components not just be tied to the X axis and automatically update?

The thinking is that for the case of a linear model defined when x units are in wavelength space but then later changed to frequency, the user would still expect fitting that to result in a straight line on the plot currently (but would not be surprised that that line becomes not-straight if later changing display units).

The alternative would be that the names/shapes of the model components only apply to the units at the time in which they are created, and we'd have to modify the equation under-the-hood to map each individual component. (This would require a complete re-implementation, so we can reconsider it if we want, but its not just a simple switch).

My only other comment so far is that I think maybe the equation editor could be more relaxed about whitespace? 'L + L_1' is invalid, only 'L+L1' is accepted.

Added support for this - good catch!

[cshanahan1]

Looks good to me then! I can't figure out anything that breaks this. Auto-updating model units can be discussed and addressed elsewhere.

[rosteen]

Tested again and confirmed that the bug I found is fixed, approved.