Skip to content

Commit dce3ca4

Browse files
tepickeringtepickering
authored
Merge pull request #86 from kecnry/background-subtract
Basic (boxcar) background subtraction
2 parents 38ff816 + 623248b commit dce3ca4

File tree

5 files changed

+685
-86
lines changed

5 files changed

+685
-86
lines changed

‎notebook_sandbox/jwst_boxcar/boxcar_extraction.ipynb‎

Lines changed: 355 additions & 28 deletions
Large diffs are not rendered by default.

‎specreduce/background.py‎

Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
1+
# Licensed under a 3-clause BSD style license - see LICENSE.rst
2+
3+
from dataclasses import dataclass, field
4+
5+
import numpy as np
6+
from astropy.nddata import NDData
7+
8+
from specreduce.extract import _ap_weight_image
9+
from specreduce.tracing import Trace, FlatTrace
10+
11+
__all__ = ['Background']
12+
13+
14+
@dataclass
15+
class Background:
16+
"""
17+
Determine the background from an image for subtraction
18+
19+
Parameters
20+
----------
21+
image : `~astropy.nddata.NDData` or array-like
22+
image with 2-D spectral image data
23+
traces : List
24+
list of trace objects (or integers to define FlatTraces) to
25+
extract the background
26+
width : float
27+
width of extraction aperture in pixels
28+
statistic: string
29+
statistic to use when computing the background. 'average' will
30+
account for partial pixel weights, 'median' will include all partial
31+
pixels.
32+
disp_axis : int
33+
dispersion axis
34+
crossdisp_axis : int
35+
cross-dispersion axis
36+
"""
37+
# required so numpy won't call __rsub__ on individual elements
38+
# https://stackoverflow.com/a/58409215
39+
__array_ufunc__ = None
40+
41+
image: NDData
42+
traces: list = field(default_factory=list)
43+
width: float = 5
44+
statistic: str = 'average'
45+
disp_axis: int = 1
46+
crossdisp_axis: int = 0
47+
48+
def __post_init__(self):
49+
"""
50+
Determine the background from an image for subtraction.
51+
52+
Parameters
53+
----------
54+
image : `~astropy.nddata.NDData` or array-like
55+
image with 2-D spectral image data
56+
traces : List
57+
list of trace objects (or integers to define FlatTraces) to
58+
extract the background
59+
width : float
60+
width of each background aperture in pixels
61+
statistic: string
62+
statistic to use when computing the background. 'average' will
63+
account for partial pixel weights, 'median' will include all partial
64+
pixels.
65+
disp_axis : int
66+
dispersion axis
67+
crossdisp_axis : int
68+
cross-dispersion axis
69+
"""
70+
def _to_trace(trace):
71+
if not isinstance(trace, Trace):
72+
trace = FlatTrace(self.image, trace)
73+
74+
# TODO: this check can be removed if/when implemented as a check in FlatTrace
75+
if isinstance(trace, FlatTrace):
76+
if trace.trace_pos < 1:
77+
raise ValueError('trace_object.trace_pos must be >= 1')
78+
return trace
79+
80+
bkg_wimage = np.zeros_like(self.image, dtype=np.float64)
81+
for trace in self.traces:
82+
trace = _to_trace(trace)
83+
bkg_wimage += _ap_weight_image(trace,
84+
self.width,
85+
self.disp_axis,
86+
self.crossdisp_axis,
87+
self.image.shape)
88+
89+
if np.any(bkg_wimage > 1):
90+
raise ValueError("background regions overlapped")
91+
92+
if self.statistic == 'median':
93+
# make it clear in the expose image that partial pixels are fully-weighted
94+
bkg_wimage[bkg_wimage > 0] = 1
95+
96+
self.bkg_wimage = bkg_wimage
97+
if self.statistic == 'average':
98+
self.bkg_array = np.average(self.image, weights=self.bkg_wimage, axis=0)
99+
elif self.statistic == 'median':
100+
med_image = self.image.copy()
101+
med_image[np.where(self.bkg_wimage) == 0] = np.nan
102+
self.bkg_array = np.nanmedian(med_image, axis=0)
103+
else:
104+
raise ValueError("statistic must be 'average' or 'median'")
105+
106+
@classmethod
107+
def two_sided(cls, image, trace_object, separation, **kwargs):
108+
"""
109+
Determine the background from an image for subtraction centered around
110+
an input trace.
111+
112+
Parameters
113+
----------
114+
image : nddata-compatible image
115+
image with 2-D spectral image data
116+
trace_object: Trace
117+
estimated trace of the spectrum to center the background traces
118+
separation: float
119+
separation from ``trace_object`` for the background regions
120+
width : float
121+
width of each background aperture in pixels
122+
statistic: string
123+
statistic to use when computing the background. 'average' will
124+
account for partial pixel weights, 'median' will include all partial
125+
pixels.
126+
disp_axis : int
127+
dispersion axis
128+
crossdisp_axis : int
129+
cross-dispersion axis
130+
"""
131+
kwargs['traces'] = [trace_object-separation, trace_object+separation]
132+
return cls(image=image, **kwargs)
133+
134+
@classmethod
135+
def one_sided(cls, image, trace_object, separation, **kwargs):
136+
"""
137+
Determine the background from an image for subtraction above
138+
or below an input trace.
139+
140+
Parameters
141+
----------
142+
image : nddata-compatible image
143+
image with 2-D spectral image data
144+
trace_object: Trace
145+
estimated trace of the spectrum to center the background traces
146+
separation: float
147+
separation from ``trace_object`` for the background, positive will be
148+
above the trace, negative below.
149+
width : float
150+
width of each background aperture in pixels
151+
statistic: string
152+
statistic to use when computing the background. 'average' will
153+
account for partial pixel weights, 'median' will include all partial
154+
pixels.
155+
disp_axis : int
156+
dispersion axis
157+
crossdisp_axis : int
158+
cross-dispersion axis
159+
"""
160+
kwargs['traces'] = [trace_object+separation]
161+
return cls(image=image, **kwargs)
162+
163+
def bkg_image(self, image=None):
164+
"""
165+
Expose the background tiled to the dimension of ``image``.
166+
167+
Parameters
168+
----------
169+
image : nddata-compatible image or None
170+
image with 2-D spectral image data. If None, will use ``image`` passed
171+
to extract the background.
172+
173+
Returns
174+
-------
175+
array with same shape as ``image``.
176+
"""
177+
if image is None:
178+
image = self.image
179+
180+
return np.tile(self.bkg_array, (image.shape[0], 1))
181+
182+
def sub_image(self, image=None):
183+
"""
184+
Subtract the computed background from ``image``.
185+
186+
Parameters
187+
----------
188+
image : nddata-compatible image or None
189+
image with 2-D spectral image data. If None, will use ``image`` passed
190+
to extract the background.
191+
192+
Returns
193+
-------
194+
array with same shape as ``image``
195+
"""
196+
if image is None:
197+
image = self.image
198+
199+
if isinstance(image, NDData):
200+
# https://docs.astropy.org/en/stable/nddata/mixins/ndarithmetic.html
201+
return image.subtract(self.bkg_image(image)*image.unit)
202+
else:
203+
return image - self.bkg_image(image)
204+
205+
def __rsub__(self, image):
206+
"""
207+
Subtract the background from an image.
208+
"""
209+
return self.sub_image(image)

‎specreduce/extract.py‎

Lines changed: 60 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,66 @@
1616
__all__ = ['BoxcarExtract', 'HorneExtract', 'OptimalExtract']
1717

1818

19+
def _get_boxcar_weights(center, hwidth, npix):
20+
"""
21+
Compute weights given an aperture center, half width,
22+
and number of pixels
23+
"""
24+
weights = np.zeros((npix))
25+
26+
# pixels with full weight
27+
fullpixels = [max(0, int(center - hwidth + 1)),
28+
min(int(center + hwidth), npix)]
29+
weights[fullpixels[0]:fullpixels[1]] = 1.0
30+
31+
# pixels at the edges of the boxcar with partial weight
32+
if fullpixels[0] > 0:
33+
w = hwidth - (center - fullpixels[0] + 0.5)
34+
if w >= 0:
35+
weights[fullpixels[0] - 1] = w
36+
else:
37+
weights[fullpixels[0]] = 1. + w
38+
if fullpixels[1] < npix:
39+
weights[fullpixels[1]] = hwidth - (fullpixels[1] - center - 0.5)
40+
41+
return weights
42+
43+
44+
def _ap_weight_image(trace, width, disp_axis, crossdisp_axis, image_shape):
45+
46+
"""
47+
Create a weight image that defines the desired extraction aperture.
48+
49+
Parameters
50+
----------
51+
trace : Trace
52+
trace object
53+
width : float
54+
width of extraction aperture in pixels
55+
disp_axis : int
56+
dispersion axis
57+
crossdisp_axis : int
58+
cross-dispersion axis
59+
image_shape : tuple with 2 elements
60+
size (shape) of image
61+
62+
Returns
63+
-------
64+
wimage : 2D image
65+
weight image defining the aperture
66+
"""
67+
wimage = np.zeros(image_shape)
68+
hwidth = 0.5 * width
69+
image_sizes = image_shape[crossdisp_axis]
70+
71+
# loop in dispersion direction and compute weights.
72+
for i in range(image_shape[disp_axis]):
73+
# TODO trace must handle transposed data (disp_axis == 0)
74+
wimage[:, i] = _get_boxcar_weights(trace[i], hwidth, image_sizes)
75+
76+
return wimage
77+
78+
1979
@dataclass
2080
class BoxcarExtract(SpecreduceOperation):
2181
"""
@@ -67,64 +127,6 @@ def __call__(self, image, trace_object, width=5,
67127
The extracted 1d spectrum with flux expressed in the same
68128
units as the input image, or u.DN, and pixel units
69129
"""
70-
def _get_boxcar_weights(center, hwidth, npix):
71-
"""
72-
Compute weights given an aperture center, half width,
73-
and number of pixels
74-
"""
75-
weights = np.zeros((npix))
76-
77-
# pixels with full weight
78-
fullpixels = [max(0, int(center - hwidth + 1)),
79-
min(int(center + hwidth), npix)]
80-
weights[fullpixels[0]:fullpixels[1]] = 1.0
81-
82-
# pixels at the edges of the boxcar with partial weight
83-
if fullpixels[0] > 0:
84-
w = hwidth - (center - fullpixels[0] + 0.5)
85-
if w >= 0:
86-
weights[fullpixels[0] - 1] = w
87-
else:
88-
weights[fullpixels[0]] = 1. + w
89-
if fullpixels[1] < npix:
90-
weights[fullpixels[1]] = hwidth - (fullpixels[1] - center - 0.5)
91-
92-
return weights
93-
94-
def _ap_weight_image(trace, width, disp_axis, crossdisp_axis, image_shape):
95-
96-
"""
97-
Create a weight image that defines the desired extraction aperture.
98-
99-
Parameters
100-
----------
101-
trace : Trace
102-
trace object
103-
width : float
104-
width of extraction aperture in pixels
105-
disp_axis : int
106-
dispersion axis
107-
crossdisp_axis : int
108-
cross-dispersion axis
109-
image_shape : tuple with 2 elements
110-
size (shape) of image
111-
112-
Returns
113-
-------
114-
wimage : 2D image
115-
weight image defining the aperture
116-
"""
117-
wimage = np.zeros(image_shape)
118-
hwidth = 0.5 * width
119-
image_sizes = image_shape[crossdisp_axis]
120-
121-
# loop in dispersion direction and compute weights.
122-
for i in range(image_shape[disp_axis]):
123-
# TODO trace must handle transposed data (disp_axis == 0)
124-
wimage[:, i] = _get_boxcar_weights(trace[i], hwidth, image_sizes)
125-
126-
return wimage
127-
128130
# TODO: this check can be removed if/when implemented as a check in FlatTrace
129131
if isinstance(trace_object, FlatTrace):
130132
if trace_object.trace_pos < 1:

‎specreduce/tests/test_background.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import numpy as np
2+
3+
import astropy.units as u
4+
from astropy.nddata import CCDData
5+
6+
from specreduce.background import Background
7+
from specreduce.tracing import FlatTrace
8+
9+
10+
# NOTE: same test image as in test_extract.py
11+
# Test image is comprised of 30 rows with 10 columns each. Row content
12+
# is row index itself. This makes it easy to predict what should be the
13+
# value extracted from a region centered at any arbitrary Y position.
14+
image = np.ones(shape=(30, 10))
15+
for j in range(image.shape[0]):
16+
image[j, ::] *= j
17+
image = CCDData(image, unit=u.Jy)
18+
19+
20+
def test_background():
21+
#
22+
# Try combinations of extraction center, and even/odd
23+
# extraction aperture sizes.
24+
#
25+
trace_pos = 15.0
26+
trace = FlatTrace(image, trace_pos)
27+
bkg_sep = 5
28+
bkg_width = 2
29+
# all the following should be equivalent:
30+
bg1 = Background(image, [trace-bkg_sep, trace+bkg_sep], width=bkg_width)
31+
bg2 = Background.two_sided(image, trace, bkg_sep, width=bkg_width)
32+
bg3 = Background.two_sided(image, trace_pos, bkg_sep, width=bkg_width)
33+
assert np.allclose(bg1.bkg_array, bg2.bkg_array)
34+
assert np.allclose(bg1.bkg_array, bg3.bkg_array)
35+
36+
# test that creating a one_sided background works
37+
Background.one_sided(image, trace, bkg_sep, width=bkg_width)
38+
39+
# test that image subtraction works
40+
sub1 = image - bg1
41+
sub2 = bg1.sub_image(image)
42+
sub3 = bg1.sub_image()
43+
assert np.allclose(sub1, sub2)
44+
assert np.allclose(sub1, sub3)

0 commit comments

Comments
 (0)