new version of the background clipping. Appears to work for general c…

…utout sizes

tessreduce/helpers.py

@@ -18,6 +18,9 @@
 from scipy.interpolate import griddata
 from scipy.optimize import minimize
 from sklearn.cluster import OPTICS
+from sklearn.preprocessing import PolynomialFeatures
+from sklearn.linear_model import LinearRegression
+from sklearn.pipeline import make_pipeline
 from skimage.restoration import inpaint
 
 from astropy.table import Table
@@ -38,6 +41,8 @@
 
 from .psf_photom import create_psf
 
+
+
 fig_width_pt = 240.0  # Get this from LaTeX using \showthe\columnwidth
 inches_per_pt = 1.0/72.27			   # Convert pt to inches
 golden_mean = (np.sqrt(5)-1.0)/2.0		 # Aesthetic ratio
@@ -412,8 +417,8 @@ def sn_lookup(name,time='disc',buffer=0,print_table=True, df = False):
 	------
 	name : str
 		catalog name
-	time : str
-		reference time to use, can be either disc, or max
+	time : str or float
+		reference time to use, accepted string values are disc, or max. Float times are assumed to be MJD.
 	buffer : float
 		overlap buffer time in days 
 	
@@ -479,12 +484,16 @@ def sn_lookup(name,time='disc',buffer=0,print_table=True, df = False):
 		new_ind = [i for i in ind if i < len(sec_times)]
 
 		secs = sec_times.iloc[new_ind]
-		if (time.lower() == 'disc') | (time.lower() == 'discovery'):
-			disc_start = secs['mjd_start'].values - disc_t.mjd
-			disc_end = secs['mjd_end'].values - disc_t.mjd
-		elif (time.lower() == 'max') | (time.lower() == 'peak'):
-			disc_start = secs['mjd_start'].values - max_t.mjd
-			disc_end = secs['mjd_end'].values - max_t.mjd
+		if type(time) == str:
+			if (time.lower() == 'disc') | (time.lower() == 'discovery'):
+				disc_start = secs['mjd_start'].values - disc_t.mjd
+				disc_end = secs['mjd_end'].values - disc_t.mjd
+			elif (time.lower() == 'max') | (time.lower() == 'peak'):
+				disc_start = secs['mjd_start'].values - max_t.mjd
+				disc_end = secs['mjd_end'].values - max_t.mjd
+		else:
+			disc_start = secs['mjd_start'].values - time
+			disc_end = secs['mjd_end'].values - time
 
 		covers = []
 		differences = []
@@ -984,4 +993,50 @@ def regional_stats_mask(image,size=90,sigma=3,iters=10):
         m,me, s = sigma_clipped_stats(image[ry,rx],maxiters=iters)
         cut_ind = np.where((image[rx,ry] >= me+sigma*s) | (image[rx,ry] <= me-sigma*s))
         clip[rx[cut_ind],ry[cut_ind]] = 1
-    return clip
+    return clip
+
+
+def Surface_names2model(names):
+	# C[i] * X^n * Y^m
+	return ' + '.join([
+				f"C[{i}]*{n.replace(' ','*')}"
+				for i,n in enumerate(names)])
+
+def clip_background(bkg,mask,sigma=3):
+	regions, max_reg = subdivide_region(bkg)
+	b2 = deepcopy(bkg)
+	for j in range(2):
+		for region in range(int(max_reg)):
+			rx,ry = np.where(regions == region)
+			y = rx.reshape(ystep,xstep)
+			x = ry.reshape(ystep,xstep)
+			sm = abs((mask & 1)-1)[y,x] * 1.0
+			sm[sm==0] = np.nan
+			cut = b2[i,y,x]
+			if j > 0:
+				masked = cut * sm
+			else:
+				masked = cut
+			xx = x.reshape(-1,1)
+			yy = y.reshape(-1,1)
+			zz = masked.reshape(-1,1)
+			ind = np.where(np.isfinite(zz))
+			order = 6
+			model = make_pipeline(PolynomialFeatures(degree=order),
+										 LinearRegression(fit_intercept=False))
+			model.fit(np.c_[xx[ind], yy[ind]], zz[ind])
+			m = Surface_names2model(model[0].get_feature_names_out(['X', 'Y']))
+			C = model[1].coef_.T  # coefficients
+			r2 = model.score(np.c_[xx[ind], yy[ind]], zz[ind])  # R-squared
+			ZZ = model.predict(np.c_[x.flatten(), y.flatten()]).reshape(x.shape)
+			diff = cut - ZZ
+			m,me, s = sigma_clipped_stats(diff,maxiters=10)
+			ind_arr = (diff >= (me+sigma*s)) | (diff <= (me-sigma*s))
+			ind_arr = fftconvolve(ind_arr,np.ones((kern_size,kern_size)),mode='same')
+			ind_arr = ind_arr > 0.8
+			cut_ind = np.where(ind_arr)
+			bkg2 = deepcopy(cut)
+			bkg2[cut_ind] = ZZ[cut_ind]
+			#bkg2 = ZZ
+			b2[y,x] = bkg2
+	return b2

tessreduce/tessreduce.py

@@ -363,6 +363,7 @@ def harsh_mask(self):
 			ind = self.ref > self._harshmask_counts
 			self.ref[ind]
 
+
 	def make_mask(self,catalogue_path=None,maglim=19,scale=1,strapsize=6,useref=False):
 		"""
 		DESCRIPTION
@@ -608,6 +609,30 @@ def _bkg_round_3(self,iters=5):
 					bkg_3[i] = parallel_bkg3(self.bkg[i],dist_mask[i])
 			self.bkg = np.array(bkg_3)
 
+	def _clip_background(self,sigma=5):
+		"""
+	    DESCRIPTION
+
+	    Parameters
+	    ----------
+	    iters : TYPE, optional
+	        DESCRIPTION. The default is 5.
+
+	    Returns
+	    -------
+	    None
+
+	    """
+
+		if self.parallel:
+			bkg_clip = Parallel(n_jobs=self.num_cores)(delayed(clip_background)(self.bkg[i],self.mask,sigma) 
+													   for i in np.arange(len(self.bkg)))
+		else:
+			bkg_clip = []
+			for i in range(len(dist_mask)):
+				bkg_clip[i] = clip_background(self.bkg[i],self.mask)
+		self.bkg = np.array(bkg_clip)
+
 	def get_ref(self,start = None, stop = None):
 		"""
 		Get reference image to use for subtraction and mask creation.
@@ -1058,6 +1083,11 @@ def diff_lc(self,time=None,x=None,y=None,ra=None,dec=None,tar_ap=3,
 		if phot_method == 'psf':
 			tar = self.psf_photometry(x,y,diff=diff)
 			tar_err = sky_std # still need to work this out
+		nan_ind = np.where(np.nansum(self.flux,axis=(1,2))==0,True,False)
+		nan_ind[self.ref_ind] = False
+		tar[nan_ind] = np.nan
+		tar_err[nan_ind] = np.nan
+
 		#tar[tar_err > 100] = np.nan
 		#sky_med[tar_err > 100] = np.nan
 		if self.tpf is not None:
@@ -1105,7 +1135,7 @@ def dif_diag_plot(self,ap_tar,ap_sky,lc=None,sky=None,data=None):
 		plt.figure(figsize=(3*fig_width,1*fig_width))
 		plt.subplot(121)
 		plt.fill_between(lc[0],sky[1]-sky[2],sky[1]+sky[2],alpha=.5,color='C1')
-		plt.plot(sky[0],sky[1],'C1.',label='Sky')
+		plt.plot(sky[0],sky[1],'C3.',label='Sky')
 		plt.fill_between(lc[0],lc[1]-lc[2],lc[1]+lc[2],alpha=.5,color='C0')
 		plt.plot(lc[0],lc[1],'C0.',label='Target')
 		binned = self.bin_data(lc=lc)
@@ -1759,7 +1789,7 @@ def reduce(self, aper = None, align = None, parallel = None, calibrate=None,
 					print('background')
 				self.bkg_orig = deepcopy(self.bkg)
 				self.background(calc_qe = False,strap_iso = False,source_hunt=self._sourcehunt,gauss_smooth=1,interpolate=False)
-				self._bkg_round_3()
+				self._clip_background()
 				self.flux -= self.bkg
 				if self.corr_correction:
 					if self.verbose > 0: