Merge pull request #551 from ibusko/redshift_enhancement

Template redshift enhancements

docs/analysis.rst

@@ -171,7 +171,7 @@ The ~`specutils.analysis.template_comparison.template_match` function takes an
 observed spectrum and ``n`` template spectra and returns the best template that
 matches the observed spectrum via chi-square minimization.
 
-If the redshift is known, the user can set that for the ``known_redshift`` parameter
+If the redshift is known, the user can set that for the ``redshift`` parameter
 and then run the
 ~`specutils.analysis.template_comparison.template_match` function.
 This function will:
@@ -180,20 +180,32 @@ This function will:
     2. Compute the chi-square between the observed spectrum and each template.
     3. Return the lowest chi-square and its corresponding template spectrum,
        normalized to the observed spectrum (and the index of the template
-       spectrum if the list of templates is iterable).
+       spectrum if the list of templates is iterable). It also
+
+If the redshift is unknown, the user specifies a grid of redshift values in the
+form of an iterable object such as a list, tuple, or numpy array with the redshift
+values to use. As an example, a simple linear grid can be built with:
+
+.. code-block:: python
+
+    >>> rs_values = np.arange(1., 3.25, 0.25)
 
-If the redshift is unknown, the user specifies the expected window in redshift
-(``min_redshift`` and ``max_redshift``) and the redshift spacing (``delta_redshift``).
 The ~`specutils.analysis.template_comparison.template_match` function will then:
 
-    1. Move each template to the first term in the redshift grid (``min_redshift``).
+    1. Move each template to the first term in the redshift grid.
     2. Run steps 1 and 2 of the case with known redshift.
-    3. Move to the next term in the redshift grid (``min_redshift + delta_redshift``).
+    3. Move to the next term in the redshift grid.
     4. Run steps 1 and 2 of the case with known redshift.
-    5. Repeat the steps until ``max_redshift`` is reached.
+    5. Repeat the steps until the end of the grid is reached.
     6. Return the best redshift, the lowest chi-square and its corresponding
-       template spectrum, redshifted and normalized to the observed spectrum
-       (and the index of the template spectrum if the list of templates is iterable).
+       template spectrum, and a list with all chi2 values, one per template.
+       The returned template spectrum corresponding to the lowest chi2 is redshifted
+       and normalized to the observed spectrum (and the index of the template spectrum if
+       the list of templates is iterable). When multiple templates are matched
+       with a redshift grid, a list-of-lists is returned with the trial chi-square
+       values computed for every combination redshift-template. The external list
+       spans the range of templates in the collection/list, while each internal list
+       contains all chi2 values for a given template.
 
 An example of how to do template matching with an unknown redshift is:
 
@@ -206,10 +218,11 @@ An example of how to do template matching with an unknown redshift is:
    >>> max_redshift = 3.0
    >>> delta_redshift = .25
    >>> resample_method = "flux_conserving"
+   >>> rs_values = np.arange(min_redshift, max_redshift+delta_redshift, delta_redshift)
 
    >>> observed_spectrum = Spectrum1D(spectral_axis=spec_axis*(1+observed_redshift), flux=np.random.randn(50) * u.Jy, uncertainty=StdDevUncertainty(np.random.sample(50), unit='Jy'))
    >>> spectral_template = Spectrum1D(spectral_axis=spec_axis, flux=np.random.randn(50) * u.Jy, uncertainty=StdDevUncertainty(np.random.sample(50), unit='Jy'))
-   >>> tm_result = template_comparison.template_match(observed_spectrum=observed_spectrum, spectral_templates=spectral_template, resample_method=resample_method, min_redshift=min_redshift, max_redshift=max_redshift, delta_redshift=delta_redshift) # doctest:+FLOAT_CMP
+   >>> tm_result = template_comparison.template_match(observed_spectrum=observed_spectrum, spectral_templates=spectral_template, resample_method=resample_method, redshift=rs_values) # doctest:+FLOAT_CMP
 
 
 Reference/API

specutils/analysis/template_comparison.py

@@ -113,8 +113,8 @@ def _chi_square_for_templates(observed_spectrum, template_spectrum, resample_met
 
 
 def template_match(observed_spectrum, spectral_templates,
-                   resample_method="flux_conserving", known_redshift=None,
-                   min_redshift=None, max_redshift=None, delta_redshift=None):
+                   resample_method="flux_conserving",
+                   redshift=None):
     """
     Find which spectral templates is the best fit to an observed spectrum by
     computing the chi-squared. If two template_spectra have the same chi2, the
@@ -135,12 +135,12 @@ def template_match(observed_spectrum, spectral_templates,
     known_redshift: `float`
         If the user knows the redshift they want to apply to the spectrum/spectra within spectral_templates,
         then this redshift can be applied to each template before attempting the match.
-    min_redshift : `float`
-        The minimum redshift allowed.
-    max_redshift : `float`
-        The maximum redshift allowed.
-    delta_redshift : `float`
-        The amount the redshift will change between loops.
+    redshift : 'float', `int`, `list`, `tuple`, 'numpy.array`
+        If the user knows the redshift they want to apply to the spectrum/spectra within spectral_templates,
+        then this float or int value redshift can be applied to each template before attempting the match.
+        Or, alternatively, an iterable with redshift values to be applied to each template, before computation
+        of the corresponding chi2 value, can be passed via this same parameter. For each template, the redshift
+        value that results in the smallest chi2 is used.
 
     Returns
     -------
@@ -151,43 +151,40 @@ def template_match(observed_spectrum, spectral_templates,
         normalized template spectrum.
     smallest_chi_index : `int`
         The index of the spectrum with the smallest chi2 in spectral templates.
+    chi2_list : `list`
+        A list with all chi2 values found for each template spectrum.
     """
     if hasattr(spectral_templates, 'flux') and len(spectral_templates.flux.shape) == 1:
 
         # Account for redshift if provided
-        if all(x is not None for x in (min_redshift, max_redshift, delta_redshift)):
-            redshift, redshifted_spectrum = template_redshift(observed_spectrum, spectral_templates,
-                                                               min_redshift, max_redshift, delta_redshift)
-            spectral_templates = redshifted_spectrum
-        elif known_redshift:
-            redshift, redshifted_spectrum = template_redshift(observed_spectrum, spectral_templates,
-                                                              known_redshift, known_redshift, 1.0)
+        chi2_list = []
+        if redshift is not None:
+            _, redshifted_spectrum, chi2_list = template_redshift(observed_spectrum, spectral_templates,
+                                                               redshift=redshift)
             spectral_templates = redshifted_spectrum
 
         normalized_spectral_template, chi2 = _chi_square_for_templates(
             observed_spectrum, spectral_templates, resample_method)
 
-        return normalized_spectral_template, chi2
+        return normalized_spectral_template, chi2, 0, chi2_list
 
     # At this point, the template spectrum is either a ``SpectrumCollection``
     # or a multi-dimensional``Spectrum1D``. Loop through the object and return
     # the template spectrum with the lowest chi square and its corresponding
     # chi square.
     chi2_min = None
     smallest_chi_spec = None
+    chi2_list = []
 
     for index, spectrum in enumerate(spectral_templates):
 
         # Account for redshift if provided
-        if all(x is not None for x in (min_redshift, max_redshift, delta_redshift)):
-            redshift, redshifted_spectrum = template_redshift(observed_spectrum, spectrum,
-                                                               min_redshift, max_redshift, delta_redshift)
+        if redshift is not None:
+            _, redshifted_spectrum, chi2_inner_list = template_redshift(
+                observed_spectrum, spectrum, redshift=redshift)
             spectrum = redshifted_spectrum
 
-        elif known_redshift:
-            redshift, redshifted_spectrum = template_redshift(observed_spectrum, spectrum,
-                                                              known_redshift, known_redshift, 1.0)
-            spectrum = redshifted_spectrum
+            chi2_list.append(chi2_inner_list)
 
         normalized_spectral_template, chi2 = _chi_square_for_templates(
             observed_spectrum, spectrum, resample_method)
@@ -197,10 +194,10 @@ def template_match(observed_spectrum, spectral_templates,
             smallest_chi_spec = normalized_spectral_template
             smallest_chi_index = index
 
-    return smallest_chi_spec, chi2_min, smallest_chi_index
+    return smallest_chi_spec, chi2_min, smallest_chi_index, chi2_list
 
 
-def template_redshift(observed_spectrum, template_spectrum, min_redshift, max_redshift, delta_redshift):
+def template_redshift(observed_spectrum, template_spectrum, redshift):
     """
     Find the best-fit redshift for template_spectrum to match observed_spectrum using chi2.
 
@@ -210,12 +207,8 @@ def template_redshift(observed_spectrum, template_spectrum, min_redshift, max_re
         The observed spectrum.
     template_spectrum : :class:`~specutils.Spectrum1D`
         The template spectrum, which will have it's redshift calculated.
-    min_redshift : `float`
-        The minimum redshift allowed.
-    max_redshift : `float`
-        The maximum redshift allowed.
-    delta_redshift : `float`
-        The amount the redshift will change between loops.
+    redshift : `float`, `int`, `list`, `tuple`, 'numpy.array`
+        A scalar or iterable with the redshift values to test.
 
     Returns
     -------
@@ -224,29 +217,30 @@ def template_redshift(observed_spectrum, template_spectrum, min_redshift, max_re
     redshifted_spectrum: :class:`~specutils.Spectrum1D`
         A new Spectrum1D object which incorporates the template_spectrum with a spectral_axis
         that has been redshifted using the final_redshift.
+    chi2_list : `list`
+        A list with the chi2 values corresponding to each input redshift value.
     """
-    if min_redshift > max_redshift:
-        raise ValueError("The `max_redshift` value must be greater than `min_redshift`.")
-        return
-
-    redshift = min_redshift
     chi2_min = None
     final_redshift = None
+    chi2_list = []
+
+    redshift = np.array(redshift).reshape((np.array(redshift).size,))
 
-    # Loop which goes through available redshift options and finds the smallest chi2
-    while redshift <= max_redshift:
+    # Loop which goes through available redshift values and finds the smallest chi2
+    for rs in redshift:
 
         # Create new redshifted spectrum and run it through the chi2 method
-        redshifted_spectrum = Spectrum1D(spectral_axis=template_spectrum.spectral_axis*(1+redshift),
+        redshifted_spectrum = Spectrum1D(spectral_axis=template_spectrum.spectral_axis*(1+rs),
                         flux=template_spectrum.flux, uncertainty=template_spectrum.uncertainty,
                                          meta=template_spectrum.meta)
         normalized_spectral_template, chi2 = _chi_square_for_templates(
                         observed_spectrum, redshifted_spectrum, "flux_conserving")
 
+        chi2_list.append(chi2)
+
         # Set new chi2_min if suitable replacement is found
         if not np.isnan(chi2) and (chi2_min is None or chi2 < chi2_min):
             chi2_min = chi2
-            final_redshift = redshift
-        redshift += delta_redshift
+            final_redshift = rs
 
-    return final_redshift, redshifted_spectrum
+    return final_redshift, redshifted_spectrum, chi2_list

specutils/tests/test_template_comparison.py

@@ -37,6 +37,7 @@ def test_template_match_no_overlap():
     # assert quantity_allclose(tm_result[0].flux, spec_result.flux, atol=0.01*u.Jy)
     assert np.isnan(tm_result[1])
 
+
 def test_template_match_minimal_overlap():
     """
     Test template_match when both observed and template spectra have minimal overlap on the wavelength axis.
@@ -70,6 +71,7 @@ def test_template_match_minimal_overlap():
     # TODO: investigate why the all elements in tm_result[1] are NaN even with overlap
     assert np.isnan(tm_result[1])
 
+
 def test_template_match_spectrum():
     """
     Test template_match when both observed and template spectra have the same wavelength axis.
@@ -154,6 +156,10 @@ def test_template_match_list():
 
     np.testing.assert_almost_equal(tm_result[1], 40093.28353756253)
 
+    # make sure that multiple template spectra will create a list of
+    # chi2 values, one per template.
+    assert len(tm_result) == 4
+
 
 def test_template_match_spectrum_collection():
     """
@@ -205,6 +211,7 @@ def test_template_match_multidim_spectrum():
 
     np.testing.assert_almost_equal(tm_result[1], 250.26870401777543)
 
+
 def test_template_unknown_redshift():
     """
     Test template redshift when redshift is unknown.
@@ -216,7 +223,7 @@ def test_template_unknown_redshift():
     spec_axis = np.linspace(0, 50, 50) * u.AA
     perm_flux = np.random.randn(50) * u.Jy
 
-    redshift = 3
+    redshift = 2.5
 
     # Observed spectrum
     spec = Spectrum1D(spectral_axis=spec_axis * (1+redshift),
@@ -232,12 +239,14 @@ def test_template_unknown_redshift():
     min_redshift = .5
     max_redshift = 5.5
     delta_redshift = .25
+    redshift_trial_values = np.arange(min_redshift, max_redshift, delta_redshift)
 
     tr_result = template_comparison.template_redshift(observed_spectrum=spec, template_spectrum=spec1,
-                                                      min_redshift=min_redshift, max_redshift=max_redshift,
-                                                      delta_redshift=delta_redshift)
+                                                      redshift=redshift_trial_values)
+
+    assert len(tr_result) == 3
+    assert tr_result[0] == 2.5
 
-    assert tr_result[0] == 3
 
 def test_template_redshift_with_one_template_spectrum_in_match():
     # Seed np.random so that results are consistent
@@ -264,14 +273,16 @@ def test_template_redshift_with_one_template_spectrum_in_match():
     min_redshift = .5
     max_redshift = 5.5
     delta_redshift = .25
+    redshift_trial_values = np.arange(min_redshift, max_redshift+delta_redshift, delta_redshift)
 
     tm_result = template_comparison.template_match(observed_spectrum=spec, spectral_templates=spec1,
-                                                      resample_method="flux_conserving", known_redshift=None,
-                                                      min_redshift=min_redshift, max_redshift=max_redshift,
-                                                      delta_redshift=delta_redshift)
+                                                      resample_method="flux_conserving",
+                                                      redshift=redshift_trial_values)
 
+    assert len(tm_result) == 4
     np.testing.assert_almost_equal(tm_result[1], 73484.0053895151)
 
+
 def test_template_redshift_with_multiple_template_spectra_in_match():
     # Seed np.random so that results are consistent
     np.random.seed(42)
@@ -303,14 +314,23 @@ def test_template_redshift_with_multiple_template_spectra_in_match():
     min_redshift = .5
     max_redshift = 5.5
     delta_redshift = .25
+    redshift_trial_values = np.arange(min_redshift, max_redshift+delta_redshift, delta_redshift)
 
     tm_result = template_comparison.template_match(observed_spectrum=spec, spectral_templates=spec_coll,
-                                                      resample_method="flux_conserving", known_redshift=None,
-                                                      min_redshift=min_redshift, max_redshift=max_redshift,
-                                                      delta_redshift=delta_redshift)
-
+                                                      resample_method="flux_conserving",
+                                                      redshift=redshift_trial_values)
+    assert len(tm_result) == 4
     np.testing.assert_almost_equal(tm_result[1], 6803.922741644725)
 
+    # When a spectrum collection is matched with a redshift
+    # grid, a list-of-lists is returned with the trial chi2
+    # values computed for every combination redshift-template.
+    # The external list spans the templates in the collection,
+    # while each internal list contains all chi2 values
+    # for a given template.
+    assert len(tm_result[3]) == 2
+
+
 def test_template_known_redshift():
     """
     Test template match when the redshift is known.
@@ -335,7 +355,7 @@ def test_template_known_redshift():
                        uncertainty=StdDevUncertainty(np.random.sample(50), unit='Jy'))
 
     tm_result = template_comparison.template_match(observed_spectrum=spec, spectral_templates=spec1,
-                                                      resample_method="flux_conserving", known_redshift=redshift,
-                                                      min_redshift=None, max_redshift=None,
-                                                      delta_redshift=None)
+                                                      resample_method="flux_conserving",
+                                                      redshift=redshift)
+    assert len(tm_result) == 4
     np.testing.assert_almost_equal(tm_result[1], 1.9062409482056814e-31)