Skip to content

⚡ tweaking PSFmachine to play with the K2 implementation #47

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
christinahedges wants to merge 2 commits into
base: master
Choose a base branch
from
Open

⚡ tweaking PSFmachine to play with the K2 implementation #47

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
80d54a0
:zap: tweaking
christinahedges Feb 3, 2022
9cfa03e
:zap: tweaking
christinahedges Feb 3, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 70 additions & 29 deletions src/psfmachine/machine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -323,12 +323,13 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):

def _get_source_mask(
self,
upper_radius_limit=28.0,
upper_radius_limit=36.0,
lower_radius_limit=4.5,
upper_flux_limit=2e5,
lower_flux_limit=100,
lower_flux_limit=40,
Comment on lines +326 to +329
Copy link
Contributor

@jorgemarpa jorgemarpa Mar 22, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this makes the rough mask more generous, good when scene motion is considerable, and allows for fainter pixels

correct_centroid_offset=True,
plot=False,
frame_index="mean",
Copy link
Contributor

@jorgemarpa jorgemarpa Mar 22, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add docstring for this argument.
Using main keeps the original behavior, right?

):
"""Find the pixel mask that identifies pixels with contributions from ANY NUMBER of Sources

Expand Down Expand Up @@ -477,6 +478,7 @@ def _get_source_mask(
source_radius_limit = np.polyval(
np.polyfit(test_f[ok], l[ok], 1), np.log10(self.source_flux_estimates)
)

source_radius_limit[
source_radius_limit > upper_radius_limit
] = upper_radius_limit
Expand All @@ -501,7 +503,7 @@ def _get_source_mask(
# Now we can update the r and phi estimates, allowing for a slight centroid
# calculate image centroids and correct dra,ddec for offset.
if correct_centroid_offset:
self._get_centroids()
self._get_centroids(frame_index=frame_index)
# print(self.ra_centroid_avg.to("arcsec"), self.dec_centroid_avg.to("arcsec"))
# re-estimate dra, ddec with centroid shifts, check if sparse case applies.
# Hardcoded: sparse implementation is efficient when nsourxes * npixels < 1e7
Expand All @@ -510,15 +512,23 @@ def _get_source_mask(
if self.nsources * self.npixels < 1e7:
self._create_delta_arrays(
centroid_offset=[
self.ra_centroid_avg.value,
self.dec_centroid_avg.value,
self.ra_centroid_avg.value
if frame_index == "mean"
else self.ra_centroid[frame_index].value,
self.dec_centroid_avg.value
if frame_index == "mean"
else self.dec_centroid[frame_index].value,
]
)
else:
self._create_delta_sparse_arrays(
centroid_offset=[
self.ra_centroid_avg.value,
self.dec_centroid_avg.value,
self.ra_centroid_avg.value
if frame_index == "mean"
else self.ra_centroid[frame_index].value,
self.dec_centroid_avg.value
if frame_index == "mean"
else self.dec_centroid[frame_index].value,
]
)
# if centroid offset id larger than 1" then we need to recalculate the
Expand Down Expand Up @@ -590,7 +600,7 @@ def _get_uncontaminated_pixel_mask(self):
return

# CH: We're not currently using this, but it might prove useful later so I will leave for now
def _get_centroids(self):
def _get_centroids(self, frame_index="mean"):
"""
Find the ra and dec centroid of the image, at each time.
"""
Expand All @@ -599,22 +609,30 @@ def _get_centroids(self):
self.dec_centroid = np.zeros(self.nt)
dra_m = self.uncontaminated_source_mask.multiply(self.dra).data
ddec_m = self.uncontaminated_source_mask.multiply(self.ddec).data
for t in range(self.nt):
wgts = self.uncontaminated_source_mask.multiply(
np.sqrt(np.abs(self.flux[t]))
).data
if frame_index == "mean":
ts = np.where(self.time_mask)[0]
else:
ts = np.atleast_1d(frame_index)
for t in ts:
wgts = (
self.uncontaminated_source_mask.multiply(self.flux[t].copy())
.multiply(1 / self.source_flux_estimates[:, None])
.data
** 2
)
# mask out non finite values and background pixels
k = (np.isfinite(wgts)) & (
self.uncontaminated_source_mask.multiply(self.flux[t]).data > 100
self.uncontaminated_source_mask.multiply(self.flux[t]).data > 10
Copy link
Contributor

@jorgemarpa jorgemarpa Mar 22, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree this value should be more generous.
Let's add a comment that this value works well for Kepler/K2 background-subtracted pixels.

)
self.ra_centroid[t] = np.average(dra_m[k], weights=wgts[k])
self.dec_centroid[t] = np.average(ddec_m[k], weights=wgts[k])

del dra_m, ddec_m
self.ra_centroid *= u.deg
self.dec_centroid *= u.deg
self.ra_centroid_avg = self.ra_centroid.mean()
self.dec_centroid_avg = self.dec_centroid.mean()

if frame_index == "mean":
self.ra_centroid_avg = np.nanmean(self.ra_centroid)
self.dec_centroid_avg = np.nanmean(self.dec_centroid)
return

def _time_bin(self, npoints=200, downsample=False):
Expand Down Expand Up @@ -691,14 +709,19 @@ def _time_bin(self, npoints=200, downsample=False):
if not downsample:
# time averaged between breaks
tm = np.vstack(
[t1.mean(axis=0) for t1 in np.array_split(self.time, breaks)]
[
t1.mean(axis=0)
for t1 in np.array_split(self.time, breaks)
if len(t1) != 0
]
).ravel()
# whiten the time array
ta = (self.time - tm.mean()) / (tm.max() - tm.mean())
# find the time index for each segments between breaks to then average flux
ms = [
np.in1d(np.arange(self.nt), i)
for i in np.array_split(np.arange(self.nt), breaks)
if len(i) != 0
]
# Average Pixel flux values
fm = np.asarray(
Expand Down Expand Up @@ -726,7 +749,11 @@ def _time_bin(self, npoints=200, downsample=False):
else:
dwns_idx = (
np.vstack(
[np.median(t1) for t1 in np.array_split(np.arange(self.nt), breaks)]
[
np.median(t1)
for t1 in np.array_split(np.arange(self.nt), breaks)
if len(t1) != 0
]
)
.ravel()
.astype(int)
Expand Down Expand Up @@ -789,10 +816,18 @@ def _time_bin(self, npoints=200, downsample=False):
# do poscorr binning
if not downsample:
pc1_bin = np.vstack(
[np.median(t1, axis=0) for t1 in np.array_split(pc1_smooth, breaks)]
[
np.median(t1, axis=0)
for t1 in np.array_split(pc1_smooth, breaks)
if len(t1) != 0
]
).ravel()[:, None] * np.ones(fm.shape)
pc2_bin = np.vstack(
[np.median(t1, axis=0) for t1 in np.array_split(pc2_smooth, breaks)]
[
np.median(t1, axis=0)
for t1 in np.array_split(pc2_smooth, breaks)
if len(t1) != 0
]
).ravel()[:, None] * np.ones(fm.shape)
else:
pc1_bin = pc1_smooth[dwns_idx][:, None] * np.ones(fm.shape)
Expand Down Expand Up @@ -1048,16 +1083,18 @@ def plot_time_model(self):
return fig

def build_shape_model(
self, plot=False, flux_cut_off=1, frame_index="mean", **kwargs
self, plot=False, flux_atol=1, flux_rtol=0.001, frame_index="mean", **kwargs
Copy link
Contributor

@jorgemarpa jorgemarpa Mar 22, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think these new arguments are better than the ones we had before.

):
"""
Builds a sparse model matrix of shape nsources x npixels to be used when
fitting each source pixels to estimate its PSF photometry

Parameters
----------
flux_cut_off: float
flux_atol: float
the flux in COUNTS at which to stop evaluating the model!
flux_rtol: float
the relative flux of the PSF at which to stop evaluating the model
frame_index : string or int
The frame index used to build the shape model, if "mean" then use the
mean value across time
Expand Down Expand Up @@ -1158,27 +1195,31 @@ def build_shape_model(
self._get_mean_model()
# remove background pixels and recreate mean model
self._update_source_mask_remove_bkg_pixels(
flux_cut_off=flux_cut_off, frame_index=frame_index
flux_atol=flux_atol, flux_rtol=flux_rtol, frame_index=frame_index
)

if plot:
return self.plot_shape_model(frame_index=frame_index)
return

def _update_source_mask_remove_bkg_pixels(self, flux_cut_off=1, frame_index="mean"):
def _update_source_mask_remove_bkg_pixels(
self, flux_atol=1, flux_rtol=0.001, frame_index="mean"
):
"""
Update the `source_mask` to remove pixels that do not contribuite to the PRF
shape.
First, re-estimate the source flux usign the precomputed `mean_model`.
This re-estimation is used to remove sources with bad prediction and update
the `source_mask` by removing background pixels that do not contribuite to
the PRF shape.
Pixels with normalized flux > `flux_cut_off` are kept.
Pixels with normalized flux > `flux_atol` are kept.

Parameters
----------
flux_cut_off : float
Lower limit for the normalized flux predicted from the mean model.
flux_atol : float
Lower limit for the normalized flux predicted from the mean model in ABSOLUTE flux.
flux_rtol : float
Lower limit for the normalized flux predicted from the mean model in RELATIVE flux.
frame_index : string or int
The frame index to be used, if "mean" then use the
mean value across time
Expand Down Expand Up @@ -1224,8 +1265,8 @@ def _update_source_mask_remove_bkg_pixels(self, flux_cut_off=1, frame_index="mea
self.mean_model.multiply(
self.mean_model.T.dot(self.source_flux_estimates)
).tocsr()
> flux_cut_off
)
> flux_atol
).multiply(self.mean_model > flux_rtol)

# Recreate uncontaminated mask
self._get_uncontaminated_pixel_mask()
Expand Down
4 changes: 2 additions & 2 deletions src/psfmachine/superstamp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,7 @@ def build_frame_shape_model(self, plot=False, **kwargs):
org_sm = self.source_mask
org_usm = self.uncontaminated_source_mask
for tdx in tqdm(range(self.nt), desc="Building shape model per frame"):
self._get_source_mask(frame_index=tdx)
fig = self.build_shape_model(frame_index=tdx, plot=plot, **kwargs)
self.mean_model_frame.append(self.mean_model)
if plot:
Expand Down Expand Up @@ -352,12 +353,11 @@ def fit_lightcurves(
# fit shape model at each cadence
self.build_frame_shape_model()
self.fit_frame_model()

self.lcs = []
for idx, s in self.sources.iterrows():
meta = {
"ORIGIN": "PSFMACHINE",
"APERTURE": "PSF + SAP" if sap else "PSF",
"APERTURE": "PSF + SAP" if hasattr(self, "sap_flux") else "PSF",
"LABEL": s.designation,
"MISSION": self.meta["MISSION"],
"RA": s.ra,
Expand Down
9 changes: 5 additions & 4 deletions src/psfmachine/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,10 +206,11 @@ def _make_A_cartesian(x, y, n_knots=10, radius=3.0, spacing="sqrt"):
np.asarray(
dmatrix(
"bs(x, knots=knots, degree=3, include_intercept=True)",
{"x": list(x), "knots": x_knots},
{"x": np.hstack([x_knots[0], x, x_knots[-1]]), "knots": x_knots},
)
)
)
)[1:-1]

if spacing == "sqrt":
y_knots = np.linspace(-np.sqrt(radius), np.sqrt(radius), n_knots)
y_knots = np.sign(y_knots) * y_knots ** 2
Expand All @@ -219,10 +220,10 @@ def _make_A_cartesian(x, y, n_knots=10, radius=3.0, spacing="sqrt"):
np.asarray(
dmatrix(
"bs(x, knots=knots, degree=3, include_intercept=True)",
{"x": list(y), "knots": y_knots},
{"x": np.hstack([y_knots[0], y, y_knots[-1]]), "knots": y_knots},
)
)
)
)[1:-1]
X = sparse.hstack(
[x_spline.multiply(y_spline[:, idx]) for idx in range(y_spline.shape[1])],
format="csr",
Expand Down