Merge pull request #128 from ojustino/poly-trace

ojustino · web-flow · commit 2a63e7d2a132 · 2022-11-17T11:29:35.000-05:00
Introducing polynomial tracing
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -4,9 +4,17 @@
 New Features
 ^^^^^^^^^^^^
 
+- The new FitTrace class (see "API Changes" below) introduces the ability to
+take a polynomial trace of an image [#128]
+
 API Changes
 ^^^^^^^^^^^
 
+- Renamed KosmosTrace as FitTrace, a conglomerate class for traces that are fit
+to images instead of predetermined [#128]
+- The default number of bins for FitTrace is now its associated image's number
+of dispersion pixels instead of 20. Its default peak_method is now 'max'. [#128]
+
 Bug Fixes
 ^^^^^^^^^
 
diff --git a/docs/extraction_quickstart.rst b/docs/extraction_quickstart.rst
@@ -16,7 +16,7 @@ the remaining steps of the extraction process.  Supported trace types include:
 
 * `~specreduce.tracing.ArrayTrace`
 * `~specreduce.tracing.FlatTrace`
-* `~specreduce.tracing.KosmosTrace`
+* `~specreduce.tracing.FitTrace`
 
 
 Each of these trace classes takes the 2D spectral image as input, as well as additional information
@@ -47,17 +47,17 @@ or, equivalently::
   bg = specreduce.tracing.Background.one_sided(image, 15, separation=5, width=2)
 
 
-The background image can be accessed via `~specreduce.background.Background.bkg_image` and the 
+The background image can be accessed via `~specreduce.background.Background.bkg_image` and the
 background-subtracted image via `~specreduce.background.Background.sub_image` (or ``image - bg``).
 
-The background and trace steps can be done iteratively, to refine an automated trace using the 
+The background and trace steps can be done iteratively, to refine an automated trace using the
 background-subtracted image as input.
 
 Extraction
 ----------
 
 The `specreduce.extract` module extracts a 1D spectrum from an input 2D spectrum (likely a
-background-extracted spectrum from the previous step) and a defined window, using one of the 
+background-extracted spectrum from the previous step) and a defined window, using one of the
 implemented methods:
 
 * `~specreduce.extract.BoxcarExtract`
diff --git a/licenses/KOSMOS_LICENSE b/licenses/KOSMOS_LICENSE
@@ -1,4 +1,4 @@
-# NOTE: This license applies only to code used in the KosmosTrace class.
+# NOTE: This license applies only to code used in part of the FitTrace class.
 
 MIT License
 
diff --git a/notebook_sandbox/compare_extractions.ipynb b/notebook_sandbox/compare_extractions.ipynb
@@ -22,7 +22,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "fc741345-720e-4189-9748-b3bb98d2d693",
    "metadata": {},
    "outputs": [],
@@ -94,8 +94,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "index_arr = np.tile(np.arange(nrows), (ncols, 1))\n",
-    "index_arr.T"
+    "index_arr = np.broadcast_to(np.arange(nrows)[:, np.newaxis], (nrows, ncols))\n",
+    "index_arr"
    ]
   },
   {
@@ -105,7 +105,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "img = col_model(index_arr.T) + noise"
+    "img = col_model(index_arr) + noise"
    ]
   },
   {
@@ -214,7 +214,7 @@
     "img_obj = CCDData(img, uncertainty=var_obj, mask=mask, unit=u.DN)\n",
     "\n",
     "hrn2 = HorneExtract(img_obj, trace)\n",
-    "hrn2_result1d_whole = hrn()"
+    "hrn2_result1d_whole = hrn2()"
    ]
   },
   {
@@ -244,7 +244,7 @@
    "metadata": {},
    "source": [
     "## Compare results\n",
-    "The whole-image extractions come out as expected, with the Horne-extracted 1D spectrum showing a noticeably better signal-to-noise ratio than its boxcar equivalent."
+    "The whole-image extractions come out as expected, with the Horne-extracted 1D spectrum showing a noticeably better signal-to-noise ratio than its windowless boxcar equivalent."
    ]
   },
   {
diff --git a/notebook_sandbox/horne_extract/optimal_extract_VLT.ipynb b/notebook_sandbox/horne_extract/optimal_extract_VLT.ipynb
@@ -8,6 +8,14 @@
     "# Optimal/Horne extraction exploration"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "5b9dc060-5262-463c-ab9b-5790dcf3e9e7",
+   "metadata": {},
+   "source": [
+    "<p style=\"font-size:125%; background-color:#ce1141;\"><b><i>Note:</i> This is an experimental notebook created for an older version of <tt>specreduce</tt>. To learn the package's current best practices, please visit our other notebooks.</b></p>"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "fc5ad51c-fc94-4dfc-8750-3e58b6224278",
diff --git a/notebook_sandbox/jwst_boxcar/boxcar_extraction.ipynb b/notebook_sandbox/jwst_boxcar/boxcar_extraction.ipynb
diff --git a/notebook_sandbox/jwst_boxcar/jwst_boxcar_algorithm.ipynb b/notebook_sandbox/jwst_boxcar/jwst_boxcar_algorithm.ipynb
@@ -7,6 +7,13 @@
     "# Spectral extraction tests (JDAT-1855)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "<p style=\"font-size:125%; background-color:#ce1141;\"><b><i>Note:</i> This is an experimental notebook created for an older version of <tt>specreduce</tt>. To learn the package's current best practices, please visit our other notebooks.</b></p>"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
diff --git a/notebook_sandbox/jwst_boxcar/tracing_options.ipynb b/notebook_sandbox/jwst_boxcar/tracing_options.ipynb
diff --git a/notebook_sandbox/tracing_options.ipynb b/notebook_sandbox/tracing_options.ipynb
diff --git a/specreduce/background.py b/specreduce/background.py
@@ -145,7 +145,7 @@ def two_sided(cls, image, trace_object, separation, **kwargs):
 
         Example: ::
 
-            trace = KosmosTrace(image, guess=trace_pos)
+            trace = FitTrace(image, guess=trace_pos)
             bg = Background.two_sided(image, trace, bkg_sep, width=bkg_width)
 
         Parameters
@@ -178,7 +178,7 @@ def one_sided(cls, image, trace_object, separation, **kwargs):
 
         Example: ::
 
-            trace = KosmosTrace(image, guess=trace_pos)
+            trace = FitTrace(image, guess=trace_pos)
             bg = Background.one_sided(image, trace, bkg_sep, width=bkg_width)
 
         Parameters
diff --git a/specreduce/tests/test_tracing.py b/specreduce/tests/test_tracing.py
@@ -3,7 +3,7 @@
 
 from astropy.modeling import models
 from specreduce.utils.synth_data import make_2dspec_image
-from specreduce.tracing import Trace, FlatTrace, ArrayTrace, KosmosTrace
+from specreduce.tracing import Trace, FlatTrace, ArrayTrace, FitTrace
 
 IM = make_2dspec_image()
 
@@ -66,8 +66,9 @@ def test_array_trace():
     assert t_short.shape[0] == IM.shape[1]
 
 
-# test KOSMOS trace algorithm
-def test_kosmos_trace():
+# test fitted traces
+@pytest.mark.filterwarnings("ignore:Model is linear in parameters")
+def test_fit_trace():
     # create image (process adapted from compare_extractions.ipynb)
     np.random.seed(7)
     nrows = 200
@@ -82,7 +83,7 @@ def test_kosmos_trace():
     img = col_model(index_arr.T) + noise
 
     # calculate trace on normal image
-    t = KosmosTrace(img)
+    t = FitTrace(img, bins=20)
 
     # test shifting
     shift_up = int(-img.shape[0]/4)
@@ -96,17 +97,20 @@ def test_kosmos_trace():
     t.shift(shift_out)
     assert t.trace.mask.all(), 'invalid values not masked'
 
-    # test peak_method options
-    tg = KosmosTrace(img, peak_method='gaussian')
-    tc = KosmosTrace(img, peak_method='centroid')
-    tm = KosmosTrace(img, peak_method='max')
+    # test peak_method and trace_model options
+    tg = FitTrace(img, bins=20,
+                  peak_method='gaussian', trace_model=models.Legendre1D(3))
+    tc = FitTrace(img, bins=20,
+                  peak_method='centroid', trace_model=models.Chebyshev1D(2))
+    tm = FitTrace(img, bins=20,
+                  peak_method='max', trace_model=models.Spline1D(degree=3))
     # traces should all be close to 100
     # (values may need to be updated on changes to seed, noise, etc.)
     assert np.max(abs(tg.trace-100)) < sigma_pix
     assert np.max(abs(tc.trace-100)) < 3 * sigma_pix
     assert np.max(abs(tm.trace-100)) < 6 * sigma_pix
     with pytest.raises(ValueError):
-        t = KosmosTrace(img, peak_method='invalid')
+        t = FitTrace(img, peak_method='invalid')
 
     # create same-shaped variations of image with invalid values
     img_all_nans = np.tile(np.nan, (nrows, ncols))
@@ -116,25 +120,37 @@ def test_kosmos_trace():
     img_win_nans = img.copy()
     img_win_nans[guess - window:guess + window] = np.nan
 
-    # ensure a low bin number is rejected
+    # ensure float bin values trigger a warning but no issues otherwise
+    with pytest.warns(UserWarning, match='TRACE: Converting bins to int'):
+        FitTrace(img, bins=20., trace_model=models.Polynomial1D(2))
+
+    # ensure non-equipped models are rejected
+    with pytest.raises(ValueError, match=r'trace_model must be one of*'):
+        FitTrace(img, trace_model=models.Hermite1D(3))
+
+    # ensure a bin number below 4 is rejected
     with pytest.raises(ValueError, match='bins must be >= 4'):
-        KosmosTrace(img, bins=3)
+        FitTrace(img, bins=3)
+
+    # ensure a bin number below degree of trace model is rejected
+    with pytest.raises(ValueError, match='bins must be > '):
+        FitTrace(img, bins=4, trace_model=models.Chebyshev1D(5))
 
     # ensure number of bins greater than number of dispersion pixels is rejected
     with pytest.raises(ValueError, match=r'bins must be <*'):
-        KosmosTrace(img, bins=ncols)
+        FitTrace(img, bins=ncols + 1)
 
     # error on trace of otherwise valid image with all-nan window around guess
     try:
-        KosmosTrace(img_win_nans, guess=guess, window=window)
+        FitTrace(img_win_nans, guess=guess, window=window)
     except ValueError as e:
         print(f"All-NaN window error message: {e}")
     else:
         raise RuntimeError('Trace was erroneously calculated on all-NaN window')
 
     # error on trace of all-nan image
     try:
-        KosmosTrace(img_all_nans)
+        FitTrace(img_all_nans)
     except ValueError as e:
         print(f"All-NaN image error message: {e}")
     else:
diff --git a/specreduce/tracing.py b/specreduce/tracing.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-# NOTE: This license applies only to code used in the KosmosTrace class.`
	`1`	`+# NOTE: This license applies only to code used in part of the FitTrace class.`
`2`	`2`
`3`	`3`	`MIT License`
`4`	`4`
Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,7 @@`
`22`	`22`	`},`
`23`	`23`	`{`
`24`	`24`	`"cell_type": "code",`
`25`		`- "execution_count": 1,`
	`25`	`+ "execution_count": null,`
`26`	`26`	`"id": "fc741345-720e-4189-9748-b3bb98d2d693",`
`27`	`27`	`"metadata": {},`
`28`	`28`	`"outputs": [],`
`@@ -94,8 +94,8 @@`
`94`	`94`	`"metadata": {},`
`95`	`95`	`"outputs": [],`
`96`	`96`	`"source": [`
`97`		`- "index_arr = np.tile(np.arange(nrows), (ncols, 1))\n",`
`98`		`- "index_arr.T"`
	`97`	`+ "index_arr = np.broadcast_to(np.arange(nrows)[:, np.newaxis], (nrows, ncols))\n",`
	`98`	`+ "index_arr"`
`99`	`99`	`]`
`100`	`100`	`},`
`101`	`101`	`{`
`@@ -105,7 +105,7 @@`
`105`	`105`	`"metadata": {},`
`106`	`106`	`"outputs": [],`
`107`	`107`	`"source": [`
`108`		`- "img = col_model(index_arr.T) + noise"`
	`108`	`+ "img = col_model(index_arr) + noise"`
`109`	`109`	`]`
`110`	`110`	`},`
`111`	`111`	`{`
`@@ -214,7 +214,7 @@`
`214`	`214`	`"img_obj = CCDData(img, uncertainty=var_obj, mask=mask, unit=u.DN)\n",`
`215`	`215`	`"\n",`
`216`	`216`	`"hrn2 = HorneExtract(img_obj, trace)\n",`
`217`		`- "hrn2_result1d_whole = hrn()"`
	`217`	`+ "hrn2_result1d_whole = hrn2()"`
`218`	`218`	`]`
`219`	`219`	`},`
`220`	`220`	`{`
`@@ -244,7 +244,7 @@`
`244`	`244`	`"metadata": {},`
`245`	`245`	`"source": [`
`246`	`246`	`"## Compare results\n",`
`247`		`- "The whole-image extractions come out as expected, with the Horne-extracted 1D spectrum showing a noticeably better signal-to-noise ratio than its boxcar equivalent."`
	`247`	`+ "The whole-image extractions come out as expected, with the Horne-extracted 1D spectrum showing a noticeably better signal-to-noise ratio than its windowless boxcar equivalent."`
`248`	`248`	`]`
`249`	`249`	`},`
`250`	`250`	`{`
Original file line number	Diff line number	Diff line change
`@@ -7,6 +7,13 @@`
`7`	`7`	`"# Spectral extraction tests (JDAT-1855)"`
`8`	`8`	`]`
`9`	`9`	`},`
	`10`	`+ {`
	`11`	`+ "cell_type": "markdown",`
	`12`	`+ "metadata": {},`
	`13`	`+ "source": [`
	`14`	`+ "<p style=\"font-size:125%; background-color:#ce1141;\"><b><i>Note:</i> This is an experimental notebook created for an older version of <tt>specreduce</tt>. To learn the package's current best practices, please visit our other notebooks.</b></p>"`
	`15`	`+ ]`
	`16`	`+ },`
`10`	`17`	`{`
`11`	`18`	`"cell_type": "markdown",`
`12`	`19`	`"metadata": {},`