diff --git a/src/psfmachine/machine.py b/src/psfmachine/machine.py index e70e633..d470bab 100644 --- a/src/psfmachine/machine.py +++ b/src/psfmachine/machine.py @@ -123,6 +123,8 @@ def __init__( phi: numpy.ndarray Angle between pixel and source coordinates (polar coordinates), in units of radians + rough_mask: scipy.sparce.csr_matrix + Sparce mask matrix with pixels that are close to sources, simple round mask source_mask: scipy.sparce.csr_matrix Sparce mask matrix with pixels that contains flux from sources uncontaminated_source_mask: scipy.sparce.csr_matrix @@ -165,12 +167,13 @@ def __init__( self.rmax = rmax self.cut_r = cut_r self.sparse_dist_lim = sparse_dist_lim * u.arcsecond - self.time_corrector = "polynomial" + self.time_corrector = "centroid" # if pos_corr, then we can smooth the vector with a Gaussian kernel of size # poscorr_filter_size, if this is < 0.5 -> no smoothing, default is 12 # beacause of 6hr-CDPP self.poscorr_filter_size = 12 - self.cartesian_knot_spacing = "sqrt" + self.cartesian_knot_spacing = "linear" + # disble tqdm prgress bar when running in HPC self.quiet = False @@ -183,13 +186,25 @@ def __init__( self.nt = len(self.time) self.npixels = self.flux.shape[1] + self.source_flux_estimates = np.copy(np.asarray(self.sources.phot_g_mean_flux)) # Hardcoded: sparse implementation is efficient when nsourxes * npixels < 1e7 # (JMP profile this) # https://github.com/SSDataLab/psfmachine/pull/17#issuecomment-866382898 + self.ra_centroid, self.dec_centroid = np.zeros((2)) * u.deg + if self.nsources * self.npixels < 1e7: + self._update_delta_arrays() + else: + self._update_delta_sparse_arrays() + self._get_rough_source_mask() + self._get_centroid() if self.nsources * self.npixels < 1e7: - self._create_delta_arrays() + self._update_delta_arrays() else: - self._create_delta_sparse_arrays() + self._update_delta_sparse_arrays() + # if self.nsources * self.npixels < 1e7: + # self._update_delta_arrays() + # else: + # self._update_delta_sparse_arrays() @property def shape(self): @@ -198,49 +213,92 @@ def shape(self): def __repr__(self): return f"Machine (N sources, N times, N pixels): {self.shape}" - def _create_delta_arrays(self, centroid_offset=[0, 0]): + @property + def dx(self): + """Delta RA, corrected for centroid shift""" + if not sparse.issparse(self.dra): + return self.dra - self.ra_centroid.value + else: + ra_offset = sparse.csr_matrix( + ( + np.repeat( + self.ra_centroid.value, + self.dra.data.shape, + ), + (self.dra.nonzero()), + ), + shape=self.dra.shape, + dtype=float, + ) + return self.dra - ra_offset + + @property + def dy(self): + """Delta Dec, corrected for centroid shift""" + if not sparse.issparse(self.ddec): + return self.ddec - self.dec_centroid.value + else: + dec_offset = sparse.csr_matrix( + ( + np.repeat( + self.dec_centroid.value, + self.ddec.data.shape, + ), + (self.ddec.nonzero()), + ), + shape=self.ddec.shape, + dtype=float, + ) + return self.ddec - dec_offset + + def _update_delta_arrays(self, frame_indices="mean"): """ Creates dra, ddec, r and phi numpy ndarrays . Parameters ---------- - centroid_offset : list - Centroid offset for [ra, dec] to be included in dra and ddec computation. - Default is [0, 0]. + frame_indices : list or str + "mean" takes the mean of all the centroids in "time_mask" + """ # The distance in ra & dec from each source to each pixel # when centroid offset is 0 (i.e. first time creating arrays) create delta # arrays from scratch - if centroid_offset[0] == centroid_offset[1] == 0: + + if not hasattr(self, "dra"): self.dra, self.ddec = np.asarray( [ [ - self.ra - self.sources["ra"][idx] - centroid_offset[0], - self.dec - self.sources["dec"][idx] - centroid_offset[1], + self.ra - self.sources["ra"][idx], + self.dec - self.sources["dec"][idx], ] for idx in range(len(self.sources)) ] ).transpose(1, 0, 2) self.dra = self.dra * (u.deg) self.ddec = self.ddec * (u.deg) - # when offsets are != 0 (i.e. updating dra and ddec arrays) we just substract - # the ofsets avoiding the for loop - else: - self.dra -= centroid_offset[0] * u.deg - self.ddec -= centroid_offset[1] * u.deg # convertion to polar coordinates - self.r = np.hypot(self.dra, self.ddec).to("arcsec") - self.phi = np.arctan2(self.ddec, self.dra) + self.r = ( + np.hypot( + self.dx, + self.dy, + ) + * 3600 + ) + self.phi = np.arctan2( + self.dy, + self.dx, + ) - def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]): + def _update_delta_sparse_arrays(self, frame_indices="mean", dist_lim=50): """ Creates dra, ddec, r and phi arrays as sparse arrays to be used for dense data, e.g. Kepler FFIs or cluster fields. Assuming that there is no flux information further than `dist_lim` for a given source, we only keep pixels within the `dist_lim`. dra, ddec, ra, and phi are unitless because they are `sparse.csr_matrix`. But - keep same scale as '_create_delta_arrays()'. + keep same scale as '_update_delta_arrays()'. dra and ddec in deg. r in arcseconds and phi in rads Parameters @@ -251,11 +309,11 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]): Centroid offset for [ra, dec] to be included in dra and ddec computation. Default is [0, 0]. """ + if frame_indices == "mean": + frame_indices = np.where(self.time_mask)[0] # If not centroid offsets or centroid correction are larget than a pixel, # then we need to compute the sparse delta arrays from scratch - if (centroid_offset[0] == centroid_offset[1] == 0) or ( - np.maximum(*np.abs(centroid_offset)) > 4 / 3600 - ): + if not hasattr(self, "dra"): # iterate over sources to only keep pixels within dist_lim # this is inefficient, could be done in a tiled manner? only for squared data dra, ddec, sparse_mask = [], [], [] @@ -264,8 +322,8 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]): desc="Creating delta arrays", disable=self.quiet, ): - dra_aux = self.ra - self.sources["ra"].iloc[i] - centroid_offset[0] - ddec_aux = self.dec - self.sources["dec"].iloc[i] - centroid_offset[1] + dra_aux = self.ra - self.sources["ra"].iloc[i] + ddec_aux = self.dec - self.sources["dec"].iloc[i] box_mask = sparse.csr_matrix( (np.abs(dra_aux) <= self.sparse_dist_lim.to("deg").value) & (np.abs(ddec_aux) <= self.sparse_dist_lim.to("deg").value) @@ -282,32 +340,16 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]): sparse_mask.eliminate_zeros() # if centroid correction is less than 1 pixel, then we just update dra and ddec # sparse arrays arrays and r and phi. - else: - self.dra = self.dra - sparse.csr_matrix( - ( - np.repeat(centroid_offset[0], self.dra.data.shape), - (self.dra.nonzero()), - ), - shape=self.dra.shape, - dtype=float, - ) - self.ddec = self.ddec - sparse.csr_matrix( - ( - np.repeat(centroid_offset[1], self.ddec.data.shape), - (self.ddec.nonzero()), - ), - shape=self.ddec.shape, - dtype=float, - ) - sparse_mask = self.dra.astype(bool) + + sparse_mask = self.dra.astype(bool) # convertion to polar coordinates. We can't apply np.hypot or np.arctan2 to # sparse arrays. We keep track of non-zero index, do math in numpy space, # then rebuild r, phi as sparse. nnz_inds = sparse_mask.nonzero() # convert radial dist to arcseconds - r_vals = np.hypot(self.dra.data, self.ddec.data) * 3600 - phi_vals = np.arctan2(self.ddec.data, self.dra.data) + r_vals = np.hypot(self.dx.data, self.dy.data) * 3600 + phi_vals = np.arctan2(self.dy.data, self.dx.data) self.r = sparse.csr_matrix( (r_vals, (nnz_inds[0], nnz_inds[1])), shape=sparse_mask.shape, @@ -321,13 +363,12 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]): del r_vals, phi_vals, nnz_inds, sparse_mask return - def _get_source_mask( + def _get_rough_source_mask( self, - upper_radius_limit=28.0, + upper_radius_limit=60.0, lower_radius_limit=4.5, upper_flux_limit=2e5, - lower_flux_limit=100, - correct_centroid_offset=True, + lower_flux_limit=3, plot=False, ): """Find the pixel mask that identifies pixels with contributions from ANY NUMBER of Sources @@ -351,12 +392,6 @@ def _get_source_mask( plot: bool Whether to show diagnostic plot. Default is False """ - - if not hasattr(self, "source_flux_estimates"): - # gaia estimate flux values per pixel to be used as flux priors - self.source_flux_estimates = np.copy( - np.asarray(self.sources.phot_g_mean_flux) - ) # We will use the radius a lot, this is for readibility # don't do if sparse array if isinstance(self.r, u.quantity.Quantity): @@ -484,51 +519,17 @@ def _get_source_mask( source_radius_limit < lower_radius_limit ] = lower_radius_limit - # Here we set the radius for each source. We add two pixels, to be generous - self.radius = source_radius_limit + 2 + # Here we set the radius for each source. We add six pixels, to be generous + self.radius = source_radius_limit + 6 # This sparse mask is one where where is ANY number of sources in a pixel if not isinstance(self.r, sparse.csr_matrix): - self.source_mask = sparse.csr_matrix(self.r.value < self.radius[:, None]) + self.rough_mask = sparse.csr_matrix(self.r.value < self.radius[:, None]) else: # for a sparse matrix doing < self.radius is not efficient, it # considers all zero values in the sparse matrix and set them to True. # this is a workaround to this problem. - self.source_mask = sparse_lessthan(r, self.radius) - - self._get_uncontaminated_pixel_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() - # 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 - # (JMP profile this) - # https://github.com/SSDataLab/psfmachine/pull/17#issuecomment-866382898 - if self.nsources * self.npixels < 1e7: - self._create_delta_arrays( - centroid_offset=[ - self.ra_centroid_avg.value, - self.dec_centroid_avg.value, - ] - ) - else: - self._create_delta_sparse_arrays( - centroid_offset=[ - self.ra_centroid_avg.value, - self.dec_centroid_avg.value, - ] - ) - # if centroid offset id larger than 1" then we need to recalculate the - # source mask to include/reject correct pixels. - # this 1 arcsec limit only works for Kepler/K2 - if ( - np.abs(self.ra_centroid_avg.to("arcsec").value) > 1 - or np.abs(self.dec_centroid_avg.to("arcsec").value) > 1 - ): - self._get_source_mask(correct_centroid_offset=False) + self.rough_mask = sparse_lessthan(r, self.radius) if plot: k = np.isfinite(f_temp_mask) @@ -567,54 +568,112 @@ def _get_source_mask( xlabel=("log$_{10}$ Gaia Source Flux"), ) return fig - return - def _get_uncontaminated_pixel_mask(self): - """ - creates a mask of shape nsources x npixels where targets are not contaminated. - This mask is used to select pixels to build the PSF model. - """ - - # This could be a property, but it is a pain to calculate on the fly, perhaps with lru_cache - self.uncontaminated_source_mask = self.source_mask.multiply( - np.asarray(self.source_mask.sum(axis=0) == 1)[0] - ).tocsr() - # have to remove leaked zeros - self.uncontaminated_source_mask.eliminate_zeros() + def _get_source_mask(self): - # # reduce to good pixels - # self.uncontaminated_pixel_mask = sparse.csr_matrix( - # self.uncontaminated_source_mask.sum(axis=0) > 0 - # ) + self._get_uncontaminated_pixel_mask() + # self._get_centroids() + # Now we can update the r and phi estimates, allowing for a slight centroid + # calculate image centroids and correct dra,ddec for offset. + # re-estimate dra, ddec with centroid shifts, check if sparse case applies. + # Hardcoded: sparse implementation is efficient when nsourxes * npixels < 1e7 + # (JMP profile this) + # https://github.com/SSDataLab/psfmachine/pull/17#issuecomment-866382898 + if self.nsources * self.npixels < 1e7: + self._update_delta_arrays(frame_indices="mean") + else: + self._update_delta_sparse_arrays(frame_indices="mean") + # if centroid offset id larger than 1" then we need to recalculate the + # source mask to include/reject correct pixels. + # this 1 arcsec limit only works for Kepler/K2 return - # CH: We're not currently using this, but it might prove useful later so I will leave for now - def _get_centroids(self): - """ - Find the ra and dec centroid of the image, at each time. - """ - # centroids are astropy quantities - self.ra_centroid = np.zeros(self.nt) - 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 - # mask out non finite values and background pixels - k = (np.isfinite(wgts)) & ( - self.uncontaminated_source_mask.multiply(self.flux[t]).data > 100 - ) - 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() + # def dra(self): + # if not hasattr(self, 'ra_centroid'): + # return self.uncontaminated_source_mask.multiply(self.dra).data + # else: + + # def _get_centroids(self, plot=False): + # """ + # Find the ra and dec centroid of the image, at each time. + # """ + # # centroids are astropy quantities + # self.ra_centroid = np.zeros(self.nt) + # 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(self.flux[t]) + # .multiply(1 / self.source_flux_estimates[:, None]) + # .data + # ** 2 + # ) + # # mask out non finite values and background pixels + # k = (np.isfinite(wgts)) & (wgts > 0.01) + # self.ra_centroid[t] = np.average(dra_m[k], weights=wgts[k]) + # self.dec_centroid[t] = np.average(ddec_m[k], weights=wgts[k]) + # if plot: + # plt.figure() + # plt.scatter( + # dra_m, ddec_m, c=wgts ** 0.5, s=1, vmin=0, vmax=0.2, cmap="Greys" + # ) + # plt.scatter(dra_m[k], ddec_m[k], c=wgts[k] ** 0.5, s=1, vmin=0, vmax=0.2) + # plt.scatter(self.ra_centroid[t], self.dec_centroid[t], c="r") + # plt.gca().set_aspect("equal") + # + # self.ra_centroid *= u.deg + # self.dec_centroid *= u.deg + # del dra_m, ddec_m + # return + + def _get_centroid(self, frame_indices="mean", binsize=0.0005): + if isinstance(frame_indices, str): + if frame_indices == "mean": + frame_indices = np.where(self.time_mask)[0] + else: + frame_indices = np.atleast_1d(frame_indices) + + x, y = ( + self.rough_mask.multiply(self.dra).data, + self.rough_mask.multiply(self.ddec).data, + ) + ar, X, Y = np.histogram2d( + x, + y, + (np.arange(-0.01, 0.01, binsize), np.arange(-0.01, 0.011, binsize)), + ) + X, Y = np.meshgrid(X, Y) + j = np.zeros(np.asarray(ar.T.shape) + 1, bool) + j[:-1, :-1] = ar.T != 0 + c = ( + self.rough_mask.multiply(self.flux[frame_indices].mean(axis=0)) + .multiply(1 / self.source_flux_estimates[:, None]) + .data + ) + k = c < 0.8 + x_k, y_k, c_k = x[k], y[k], c[k] + ar2 = np.asarray( + [ + np.nanmedian( + c_k[ + (x_k >= X1) + & (x_k < X1 + binsize) + & (y_k >= Y1) + & (y_k < Y1 + binsize) + ] + ) + for X1, Y1 in zip(X[j], Y[j]) + ] + ) + self.ra_centroid = ( + np.average(X[j] + binsize / 2, weights=np.nan_to_num(ar2 ** 3)) * u.deg + ) + self.dec_centroid = ( + np.average(Y[j] + binsize / 2, weights=np.nan_to_num(ar2 ** 3)) * u.deg + ) return def _time_bin(self, npoints=200, downsample=False): @@ -661,8 +720,8 @@ def _time_bin(self, npoints=200, downsample=False): mpc2 = np.nanmedian(self.pos_corr2, axis=0) else: # if usig centroids need to convert to pixels - mpc1 = self.ra_centroid.to("arcsec").value / 4 - mpc2 = self.dec_centroid.to("arcsec").value / 4 + mpc1 = self.ra_centroids.to("arcsec").value / 4 + mpc2 = self.dec_centroids.to("arcsec").value / 4 # find poscorr discontinuity in each axis grads1 = np.gradient(mpc1, self.time) @@ -812,6 +871,23 @@ def _time_bin(self, npoints=200, downsample=False): return ta, tm, fm_raw, fm, fem + def _time_matrix(self, mask): + dx, dy = ( + mask.multiply(self.dx), + mask.multiply(self.dy), + ) + dx = dx.data * u.deg.to(u.arcsecond) + dy = dy.data * u.deg.to(u.arcsecond) + + A_c = _make_A_cartesian( + dx, + dy, + n_knots=self.n_time_knots, + radius=self.time_radius, + spacing=self.cartesian_knot_spacing, + ) + return A_c + def build_time_model(self, plot=False, downsample=False): """ Builds a time model that moves the PRF model to account for the scene movement @@ -857,29 +933,20 @@ def build_time_model(self, plot=False, downsample=False): ) = self._time_bin(npoints=self.n_time_points, downsample=downsample) self._whitened_time = time_original - # not necessary to take value from Quantity to do .multiply() - dx, dy = ( - self.uncontaminated_source_mask.multiply(self.dra), - self.uncontaminated_source_mask.multiply(self.ddec), + A_c = self._time_matrix(self.uncontaminated_source_mask) + A2 = sparse.vstack( + [A_c] * time_binned.shape[0], + format="csr", ) - dx = dx.data * u.deg.to(u.arcsecond) - dy = dy.data * u.deg.to(u.arcsecond) - A_c = _make_A_cartesian( - dx, - dy, - n_knots=self.n_time_knots, - radius=self.time_radius, - spacing=self.cartesian_knot_spacing, - ) - A2 = sparse.vstack([A_c] * time_binned.shape[0], format="csr") - # Cartesian spline with time dependence if hasattr(self, "pos_corr1") and self.time_corrector in [ "pos_corr", "centroid", ]: # Cartesian spline with poscor dependence - A3 = _combine_A(A2, poscorr=[poscorr1_binned, poscorr2_binned]) + A3 = _combine_A( + A2, poscorr=[poscorr1_binned, poscorr2_binned], time=time_binned + ) else: # Cartesian spline with time dependence A3 = _combine_A(A2, time=time_binned) @@ -963,22 +1030,11 @@ def plot_time_model(self): flux_err_binned, ) = self._time_bin(npoints=self.n_time_points) - # not necessary to take value from Quantity to do .multiply() - dx, dy = ( - self.uncontaminated_source_mask.multiply(self.dra), - self.uncontaminated_source_mask.multiply(self.ddec), - ) - dx = dx.data * u.deg.to(u.arcsecond) - dy = dy.data * u.deg.to(u.arcsecond) - - A_c = _make_A_cartesian( - dx, - dy, - n_knots=self.n_time_knots, - radius=self.time_radius, - spacing=self.cartesian_knot_spacing, + A_c = self._time_matrix(self.uncontaminated_source_mask) + A2 = sparse.vstack( + [A_c] * time_binned.shape[0], + format="csr", ) - A2 = sparse.vstack([A_c] * time_binned.shape[0], format="csr") # Cartesian spline with time dependence # Cartesian spline with time dependence if hasattr(self, "pos_corr1") and self.time_corrector in [ @@ -986,11 +1042,17 @@ def plot_time_model(self): "centroid", ]: # Cartesian spline with poscor dependence - A3 = _combine_A(A2, poscorr=[poscorr1_binned, poscorr2_binned]) + A3 = _combine_A( + A2, poscorr=[poscorr1_binned, poscorr2_binned], time=time_binned + ) else: # Cartesian spline with time dependence A3 = _combine_A(A2, time=time_binned) + dx, dy = ( + self.uncontaminated_source_mask.multiply(self.dx).data * 3600, + self.uncontaminated_source_mask.multiply(self.dy).data * 3600, + ) model = A3.dot(self.time_model_w).reshape(flux_binned.shape) + 1 fig, ax = plt.subplots(2, 2, figsize=(7, 6), facecolor="w") k1 = self._time_masked.reshape(flux_binned.shape)[0] @@ -1048,7 +1110,7 @@ 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_cut_off=1, frame_index="max", **kwargs ): """ Builds a sparse model matrix of shape nsources x npixels to be used when @@ -1070,39 +1132,43 @@ def build_shape_model( # Mask of shape nsources x number of pixels, one where flux from a # source exists - if not hasattr(self, "source_mask"): - self._get_source_mask(**kwargs) + if not hasattr(self, "uncontaminated_source_mask"): + mask = self.rough_mask + else: + mask = self.uncontaminated_source_mask + # self._get_source_mask(**kwargs) # Mask of shape npixels (maybe by nt) where not saturated, not faint, # not contaminated etc - self._get_uncontaminated_pixel_mask() + # self._get_uncontaminated_pixel_mask() # for iter in range(niters): flux_estimates = self.source_flux_estimates[:, None] if frame_index == "mean": f = (self.flux[self.time_mask]).mean(axis=0) - elif isinstance(frame_index, int): + elif frame_index == "max": + f = (self.flux[self.time_mask]).max(axis=0) + elif isinstance(frame_index, (int, np.int64, np.int32, np.int16)): f = self.flux[frame_index] + else: + raise ValueError(f"frame_index {frame_index} not valid") # f, fe = (self.flux[self.time_mask]).mean(axis=0), ( # (self.flux_err[self.time_mask] ** 2).sum(axis=0) ** 0.5 # ) / (self.nt) mean_f = np.log10( - self.uncontaminated_source_mask.astype(float) - .multiply(f) - .multiply(1 / flux_estimates) - .data + mask.astype(float).multiply(f).multiply(1 / flux_estimates).data ) # Actual Kepler errors cause all sorts of instability # mean_f_err = ( - # self.uncontaminated_source_mask.astype(float) + # mask.astype(float) # .multiply(fe / (f * np.log(10))) # .multiply(1 / flux_estimates) # .data # ) # We only need these weights for the wings, so we'll use poisson noise mean_f_err = ( - self.uncontaminated_source_mask.astype(float) + mask.astype(float) .multiply((f ** 0.5) / (f * np.log(10))) .multiply(1 / flux_estimates) .data @@ -1110,8 +1176,8 @@ def build_shape_model( mean_f_err.data = np.abs(mean_f_err.data) # take value from Quantity is not necessary - phi_b = self.uncontaminated_source_mask.multiply(self.phi).data - r_b = self.uncontaminated_source_mask.multiply(self.r).data + phi_b = mask.multiply(self.phi).data + r_b = mask.multiply(self.r).data # build a design matrix A with b-splines basis in radius and angle axis. A = _make_A_polar( @@ -1124,7 +1190,7 @@ def build_shape_model( n_phi_knots=self.n_phi_knots, ) prior_sigma = np.ones(A.shape[1]) * 10 - prior_mu = np.zeros(A.shape[1]) - 10 + prior_mu = np.zeros(A.shape[1]) - 100 nan_mask = np.isfinite(mean_f.ravel()) @@ -1139,7 +1205,7 @@ def build_shape_model( errors=True, ) - bad = sigma_clip(mean_f.ravel() - A.dot(psf_w), sigma=5).mask + bad = sigma_clip((mean_f.ravel() - A.dot(psf_w)), sigma=5).mask psf_w, psf_w_err = solve_linear_model( A, @@ -1156,91 +1222,103 @@ def build_shape_model( # We then build the same design matrix for all pixels with flux self._get_mean_model() + # self._update_source_mask_remove_bkg_pixels() # remove background pixels and recreate mean model - self._update_source_mask_remove_bkg_pixels( - flux_cut_off=flux_cut_off, frame_index=frame_index - ) + # self._update_source_mask_remove_bkg_pixels( + # flux_cut_off=flux_cut_off, frame_index=frame_index + # ) if plot: - return self.plot_shape_model(frame_index=frame_index) + return self.plot_shape_model(mask, frame_index=frame_index) return - def _update_source_mask_remove_bkg_pixels(self, flux_cut_off=1, 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. - - Parameters - ---------- - flux_cut_off : float - Lower limit for the normalized flux predicted from the mean model. - frame_index : string or int - The frame index to be used, if "mean" then use the - mean value across time - """ - - # Re-estimate source flux - # ----- - prior_mu = self.source_flux_estimates - prior_sigma = ( - np.ones(self.mean_model.shape[0]) * 10 * self.source_flux_estimates - ) - - if frame_index == "mean": - f, fe = (self.flux).mean(axis=0), ( - (self.flux_err ** 2).sum(axis=0) ** 0.5 - ) / (self.nt) - elif isinstance(frame_index, int): - f, fe = self.flux[frame_index], self.flux_err[frame_index] - - X = self.mean_model.copy() - X = X.T - - sigma_w_inv = X.T.dot(X.multiply(1 / fe[:, None] ** 2)).toarray() - sigma_w_inv += np.diag(1 / (prior_sigma ** 2)) - B = X.T.dot((f / fe ** 2)) - B += prior_mu / (prior_sigma ** 2) - ws = np.linalg.solve(sigma_w_inv, B) - werrs = np.linalg.inv(sigma_w_inv).diagonal() ** 0.5 - - # ----- - - # Rebuild source mask - ok = np.abs(ws - self.source_flux_estimates) / werrs > 3 - ok &= ((ws / self.source_flux_estimates) < 10) & ( - (self.source_flux_estimates / ws) < 10 - ) - ok &= ws > 10 - ok &= werrs > 0 - - self.source_flux_estimates[ok] = ws[ok] - - self.source_mask = ( - self.mean_model.multiply( - self.mean_model.T.dot(self.source_flux_estimates) - ).tocsr() - > flux_cut_off - ) + # + # def _update_source_mask_remove_bkg_pixels(self, flux_cut_off=1, frame_index="max"): + # """ + # 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. + # + # Parameters + # ---------- + # flux_cut_off : float + # Lower limit for the normalized flux predicted from the mean model. + # frame_index : string or int + # The frame index to be used, if "mean" then use the + # mean value across time + # """ + # + # # Re-estimate source flux + # # ----- + # prior_mu = self.source_flux_estimates + # prior_sigma = ( + # np.ones(self.mean_model.shape[0]) * 10 * self.source_flux_estimates + # ) + # + # if frame_index == "mean": + # f, fe = (self.flux).mean(axis=0), ( + # (self.flux_err ** 2).sum(axis=0) ** 0.5 + # ) / (self.nt) + # if frame_index == "max": + # f, fe = (self.flux).max(axis=0), ( + # (self.flux_err ** 2).sum(axis=0) ** 0.5 + # ) / (self.nt) + # elif isinstance(frame_index, int): + # f, fe = self.flux[frame_index], self.flux_err[frame_index] + # + # X = self.mean_model.copy() + # X = X.T + # + # sigma_w_inv = X.T.dot(X.multiply(1 / fe[:, None] ** 2)).toarray() + # sigma_w_inv += np.diag(1 / (prior_sigma ** 2)) + # B = X.T.dot((f / fe ** 2)) + # B += prior_mu / (prior_sigma ** 2) + # ws = np.linalg.solve(sigma_w_inv, B) + # werrs = np.linalg.inv(sigma_w_inv).diagonal() ** 0.5 + # + # # ----- + # + # # Rebuild source mask + # ok = np.abs(ws - self.source_flux_estimates) / werrs > 3 + # ok &= ((ws / self.source_flux_estimates) < 10) & ( + # (self.source_flux_estimates / ws) < 10 + # ) + # ok &= ws > 10 + # ok &= werrs > 0 + # + # self.source_flux_estimates[ok] = ws[ok] + # + # self.source_mask = ( + # self.mean_model.multiply( + # self.mean_model.T.dot(self.source_flux_estimates) + # ).tocsr() + # > flux_cut_off + # ) + # + # # Recreate uncontaminated mask + # self._get_uncontaminated_pixel_mask() + # # self.uncontaminated_source_mask = self.uncontaminated_source_mask.multiply( + # # (self.mean_model.max(axis=1) < 1) + # # ) + # + # # Recreate mean model! + # self._get_mean_model() + + def _get_mean_model(self, relative_flux_limit=0.01, absolute_flux_limit=1): - # Recreate uncontaminated mask - self._get_uncontaminated_pixel_mask() - # self.uncontaminated_source_mask = self.uncontaminated_source_mask.multiply( - # (self.mean_model.max(axis=1) < 1) - # ) - - # Recreate mean model! - self._get_mean_model() + if not hasattr(self, "source_mask"): + mask = self.rough_mask + else: + mask = self.source_mask - def _get_mean_model(self): """Convenience function to make the scene model""" Ap = _make_A_polar( - self.source_mask.multiply(self.phi).data, - self.source_mask.multiply(self.r).data, + mask.multiply(self.phi).data, + mask.multiply(self.r).data, rmin=self.rmin, rmax=self.rmax, cut_r=self.cut_r, @@ -1252,11 +1330,17 @@ def _get_mean_model(self): mean_model = sparse.csr_matrix(self.r.shape) m = 10 ** Ap.dot(self.psf_w) m[~np.isfinite(m)] = 0 - mean_model[self.source_mask] = m + mean_model[mask] = m + mean_model = mean_model.multiply(mean_model > relative_flux_limit) + + mean_model = mean_model.multiply( + mean_model.multiply(self.source_flux_estimates[:, None]) + > absolute_flux_limit + ) mean_model.eliminate_zeros() self.mean_model = mean_model - def plot_shape_model(self, radius=20, frame_index="mean"): + def plot_shape_model(self, mask, radius=20, frame_index="max"): """ Diagnostic plot of shape model. @@ -1276,22 +1360,29 @@ def plot_shape_model(self, radius=20, frame_index="mean"): if frame_index == "mean": mean_f = np.log10( - self.uncontaminated_source_mask.astype(float) + mask.astype(float) .multiply(self.flux[self.time_mask].mean(axis=0)) .multiply(1 / self.source_flux_estimates[:, None]) .data ) + if frame_index == "max": + mean_f = np.log10( + mask.astype(float) + .multiply(self.flux[self.time_mask].max(axis=0)) + .multiply(1 / self.source_flux_estimates[:, None]) + .data + ) elif isinstance(frame_index, int): mean_f = np.log10( - self.uncontaminated_source_mask.astype(float) + mask.astype(float) .multiply(self.flux[frame_index]) .multiply(1 / self.source_flux_estimates[:, None]) .data ) dx, dy = ( - self.uncontaminated_source_mask.multiply(self.dra), - self.uncontaminated_source_mask.multiply(self.ddec), + mask.multiply(self.dx), + mask.multiply(self.dy), ) dx = dx.data * u.deg.to(u.arcsecond) dy = dy.data * u.deg.to(u.arcsecond) @@ -1392,6 +1483,17 @@ def fit_model(self, fit_va=False): Fitting model accounting for velocity aberration. If `True`, then a time model has to be built previously with `build_time_model`. """ + + self._n_time_components = [ + 6 + if self.time_corrector + in [ + "pos_corr", + "centroid", + ] + else 4 + ][0] + prior_mu = self.source_flux_estimates # np.zeros(A.shape[1]) prior_sigma = ( np.ones(self.mean_model.shape[0]) @@ -1411,21 +1513,8 @@ def fit_model(self, fit_va=False): ) # not necessary to take value from Quantity to do .multiply() - dx, dy = ( - self.source_mask.multiply(self.dra), - self.source_mask.multiply(self.ddec), - ) - dx = dx.data * u.deg.to(u.arcsecond) - dy = dy.data * u.deg.to(u.arcsecond) - - A_cp = _make_A_cartesian( - dx, - dy, - n_knots=self.n_time_knots, - radius=self.time_radius, - spacing=self.cartesian_knot_spacing, - ) - A_cp3 = sparse.hstack([A_cp, A_cp, A_cp, A_cp], format="csr") + A_cp = self._time_matrix(self.mean_model != 0) + A_cp3 = sparse.hstack([A_cp] * self._n_time_components, format="csr") self.ws_va = np.zeros((self.nt, self.mean_model.shape[0])) self.werrs_va = np.zeros((self.nt, self.mean_model.shape[0])) @@ -1458,18 +1547,23 @@ def fit_model(self, fit_va=False): np.array( [ 1, + self._whitened_time[tdx], + self._whitened_time[tdx] ** 2, self.pos_corr1_smooth[tdx], self.pos_corr2_smooth[tdx], self.pos_corr1_smooth[tdx] * self.pos_corr2_smooth[tdx], ] )[:, None] - * np.ones(A_cp3.shape[1] // 4) + * np.ones(A_cp3.shape[1] // self._n_time_components) ) else: # use time t_mult = np.hstack( - (self._whitened_time[tdx] ** np.arange(4))[:, None] - * np.ones(A_cp3.shape[1] // 4) + ( + self._whitened_time[tdx] + ** np.arange(self._n_time_components) + )[:, None] + * np.ones(A_cp3.shape[1] // self._n_time_components) ) X.data *= A_cp3.multiply(t_mult).dot(self.time_model_w) + 1 X = X.T @@ -1510,7 +1604,7 @@ def fit_model(self, fit_va=False): self.werrs[tdx] = np.linalg.inv(sigma_w_inv).diagonal() ** 0.5 self.model_flux[tdx] = X.dot(self.ws[tdx]) - nodata = np.asarray(self.source_mask.sum(axis=1))[:, 0] == 0 + nodata = np.asarray(mask.sum(axis=1))[:, 0] == 0 # These sources are poorly estimated nodata |= (self.mean_model.max(axis=1) > 1).toarray()[:, 0] self.ws[:, nodata] *= np.nan diff --git a/src/psfmachine/utils.py b/src/psfmachine/utils.py index 7374f72..c01524b 100644 --- a/src/psfmachine/utils.py +++ b/src/psfmachine/utils.py @@ -197,21 +197,26 @@ def _make_A_polar(phi, r, cut_r=6, rmin=1, rmax=18, n_r_knots=12, n_phi_knots=15 def _make_A_cartesian(x, y, n_knots=10, radius=3.0, spacing="sqrt"): + # Must be odd + n_time_knots = n_knots if n_knots % 2 == 1 else n_knots + 1 if spacing == "sqrt": - x_knots = np.linspace(-np.sqrt(radius), np.sqrt(radius), n_knots) + x_knots = np.linspace(-np.sqrt(radius), np.sqrt(radius), n_time_knots) x_knots = np.sign(x_knots) * x_knots ** 2 else: - x_knots = np.linspace(-radius, radius, n_knots) + x_knots = np.linspace(-radius, radius, n_time_knots) x_spline = sparse.csr_matrix( np.asarray( dmatrix( "bs(x, knots=knots, degree=3, include_intercept=True)", - {"x": list(x), "knots": x_knots}, + { + "x": list(np.hstack([x_knots.min(), x, x_knots.max()])), + "knots": x_knots, + }, ) - ) + )[1:-1] ) if spacing == "sqrt": - y_knots = np.linspace(-np.sqrt(radius), np.sqrt(radius), n_knots) + y_knots = np.linspace(-np.sqrt(radius), np.sqrt(radius), n_time_knots) y_knots = np.sign(y_knots) * y_knots ** 2 else: y_knots = np.linspace(-radius, radius, n_knots) @@ -219,9 +224,12 @@ 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": list(np.hstack([y_knots.min(), y, y_knots.max()])), + "knots": y_knots, + }, ) - ) + )[1:-1] ) X = sparse.hstack( [x_spline.multiply(y_spline[:, idx]) for idx in range(y_spline.shape[1])], @@ -424,9 +432,11 @@ def _combine_A(A, poscorr=None, time=None): A2 = sparse.hstack( [ A, + A.multiply(time.ravel()[:, None]), + A.multiply(time.ravel()[:, None] ** 2), A.multiply(poscorr[0].ravel()[:, None]), A.multiply(poscorr[1].ravel()[:, None]), - A.multiply((poscorr[0] * poscorr[1]).ravel()[:, None]), + A.multiply((poscorr[0].ravel() * poscorr[1].ravel())[:, None]), ], format="csr", ) @@ -443,3 +453,14 @@ def _combine_A(A, poscorr=None, time=None): format="csr", ) return A2 + + +def _find_uncontaminated_pixels(mask): + """ + creates a mask of shape nsources x npixels where targets are not contaminated. + This mask is used to select pixels to build the PSF model. + """ + + new_mask = mask.multiply(np.asarray(mask.sum(axis=0) == 1)[0]).tocsr() + new_mask.eliminate_zeros() + return new_mask