<filename>se-1st-phot-function.py
# 1st photometry
# estimate FWHM value and auto mag zp
import os
import glob
import astropy.io.fits as fits
import numpy as np
import subprocess
from sys import flags
from astropy.table import Table
from astropy.stats import sigma_clip
from astropy.stats import sigma_clipped_stats
from astropy.stats import sigma_clipping
import matplotlib.pyplot as plt
import seaborn as sns
import astropy.io.ascii as ascii
from scipy.interpolate import UnivariateSpline
from multiprocessing import Process,Pool

# growth curve fit to estimate optimal aperture size

def puthdr(inim, hdrkey, hdrval, hdrcomment=''):
	from astropy.io import fits
	hdr = fits.getheader(inim)
	fits.setval(inim, hdrkey, value=hdrval, comment=hdrcomment)
	comment = inim+'\t'+'('+hdrkey+'\t'+str(hdrval)+')'

def pixelscale(i):
	cd11 = fits.getheader(i)['CD1_1']
	cd12 = fits.getheader(i)['CD1_2']
	cd21 = fits.getheader(i)['CD2_1']
	cd22 = fits.getheader(i)['CD2_2']
	pixscale = round(np.sqrt(cd11**2 + cd21**2) * 3600, 4)
	puthdr(i, 'PSCALE', round(pixscale, 3))
	# print('Pixel scale =', pixscale,'\"')
	return pixscale

# input files, config and params
seconfigdir = '/data7/cschoi/code/cspy/sex.config/'
seconfig = 'se1.sex'
separam = 'se1.param'
growthparam = 'growth.param'
separam_noPSF = 'se1_noPSF.param'
seconv = 'default.conv'
sennw = 'default.nnw'
DETECT_MINAREA = str(5)
DETECT_THRESH = str(3)
DEBLEND_NTHRESH = str(32)
DEBLEND_MINCONT = str(0.005)

def segrowthcom(im, psf=False):
	PSCALE = pixelscale(im)
	fn = os.path.splitext(im)[0]
	aper_list=[s for s in range(1,51)]
	aper_input = ''
	for i in aper_list: aper_input += '{},'.format(i,1)
	aper_input = aper_input[:-1]
	opt1 = seconfigdir+seconfig+' -CATALOG_TYPE ASCII_HEAD -CATALOG_NAME ' + fn+'.growth'
	opt2=' -FILTER_NAME '+seconfigdir+seconv +' -STARNNW_NAME '+seconfigdir+sennw
	#opt2a = ' -PARAMETERS_NAME '+seconfigdir+separam
	#opt2b = ' -PARAMETERS_NAME '+seconfigdir+separam_noPSF
	opt2b = ' -PARAMETERS_NAME '+seconfigdir+growthparam
	opt2 = ' -FILTER_NAME '+seconfigdir+seconv + ' -STARNNW_NAME '+seconfigdir+sennw
	opt3 = ' -DETECT_MINAREA ' + DETECT_MINAREA + ' -DETECT_THRESH '+DETECT_THRESH
	opt4 = ' -DEBLEND_NTHRESH ' + DEBLEND_NTHRESH + \
	    ' -DEBLEND_MINCONT ' + DEBLEND_MINCONT
	opt5=' -CHECKIMAGE_TYPE NONE '
	opt6 = ' -PHOT_APERTURES '+aper_input+' '
	opt7 = ' -PSF_NAME '+fn+'.psf '
	opt8 = ' -PIXEL_SCALE '+str(PSCALE)+' '
	secommand = 'sex -c '+opt1+opt2+opt2b+opt3+opt4+opt5+opt6+opt8 + im
	print(secommand)
	sexout = subprocess.getoutput(secommand)
	line = [s for s in sexout.split('\n') if 'RMS' in s]
	skymed = float(line[0].split('Background:')[1].split('RMS:')[0])
	skysig= float(line[0].split('RMS:')[1].split('/')[0])
	os.system(secommand)
	return skymed, skysig

def fwhm_img(im,mtbl1):
	fwhm_img=sigma_clipped_stats(mtbl1['FWHM_IMAGE'], sigma=3, maxiters=10)
	filtered_data=sigma_clip(mtbl1['FWHM_IMAGE'],sigma=3,maxiters=10)
	selected, nonselected= ~filtered_data.mask, filtered_data.mask
	print('FWHM_IMAGE','{}'.format(round(fwhm_img[0],3)),
		len(mtbl1[selected]),'stars from',len(mtbl1))
	puthdr(im, 'FWHM_PIX', round(fwhm_img[0],3), hdrcomment='FWHM PIXEL VALUE')
	return round(fwhm_img[0],3)

def find_opt_aper(im):
	skymed, skysig=segrowthcom(im)
	fn=os.path.splitext(im)[0]
	growthcat=ascii.read(fn+'.growth')
	aper = np.linspace(1,50,50)
	# catalog cleaning
	idx= np.where( (growthcat['FLAGS']==0) &
            (growthcat['SNR_WIN']>100) &
            (growthcat['MAGERR_AUTO']<0.1)
        	)
	mcat=growthcat[idx]
	# signal to noise ratio
	fwhm=fwhm_img(im,mcat)
	plt.plot(np.ones(len(mcat))*aper[0],mcat['FLUX_APER']/mcat['FLUXERR_APER'],'k.')
	for q in aper[:-1]:
		p=int(q)
		plt.plot(np.ones(len(mcat))*aper[p],mcat['FLUX_APER_'+str(p)]/mcat['FLUXERR_APER_'+str(p)],'k.')
	# mean values
	plt.plot(aper[0],np.mean(mcat['FLUX_APER']/mcat['FLUXERR_APER']),'bo')
	for q in aper[:-1]:
		p=int(q)
		plt.plot(aper[p],np.mean(mcat['FLUX_APER_'+str(p)]/mcat['FLUXERR_APER_'+str(p)]),'bo')
	# plt.plot(np.ones(len(mcat))*0,mcat['FLUX_PSF']/mcat['FLUXERR_PSF'],'ro')
	# mcat['FLUX_AUTO']/mcat['FLUXERR_AUTO']
	# plt.vlines(fwhm*1.5, 0,10000)
	# curve fitting
	snrval=[0]
	snrval[0]=np.mean(mcat['FLUX_APER']/mcat['FLUXERR_APER'])
	for q in aper[:-1]:
		p=int(q)
		snrval.append(np.mean(mcat['FLUX_APER_'+str(p)]/mcat['FLUXERR_APER_'+str(p)]))
	sarr=sarr=np.asarray(snrval)
	sarr1=sarr[~np.isnan(sarr)]
	aper1=aper[~np.isnan(sarr)]
	s = UnivariateSpline(aper1, sarr1, s=1)
	x=np.linspace(1,50,500)
	ys = s(x)
	idx_opt=np.where(ys==np.max(ys))
	plt.plot(x,ys)
	plt.vlines(x[idx_opt],0,10000)
	plt.title(fn+' '+'optimal ap size')
	plt.xlabel('aperture size (pixel)')
	plt.ylabel('SNR (FLUX / FLUX err)')
	plt.savefig(fn+'_growth.png')
	plt.close()
	os.system('rm '+fn+'.growth')
	print('Optimal aperture',round(x[idx_opt][0],2))
	return skymed, skysig, fwhm, round(x[idx_opt][0],2)
# first maximum
def opt_ap_fwhm(im):
	skyval, skysig,fwhm,opt_ap=find_opt_aper(im)
	puthdr(im, 'SKYVAL', skyval,
		hdrcomment='sky median value form sextractor')
	puthdr(im, 'SKYSIG', skysig,
		hdrcomment='sky sigma value form sextractor')
	puthdr(im, 'FWHM_PIX', round(fwhm,3), hdrcomment='FWHM PIXEL VALUE')
	puthdr(im, 'OPT_AP', opt_ap, hdrcomment='Optimal aperture size')
	return skyval, skysig,fwhm,opt_ap

def fwhm_img(im,mtbl1):
	fwhm_img=sigma_clipped_stats(mtbl1['FWHM_IMAGE'], sigma=3, maxiters=10)
	filtered_data=sigma_clip(mtbl1['FWHM_IMAGE'],sigma=3,maxiters=10)
	selected, nonselected= ~filtered_data.mask, filtered_data.mask
	print('FWHM_IMAGE','{}'.format(round(fwhm_img[0],3)),
		len(mtbl1[selected]),'stars from',len(mtbl1))
	puthdr(im, 'FWHM_PIX', round(fwhm_img[0],3), hdrcomment='FWHM PIXEL VALUE')
	puthdr(im, 'NUM_FWHM', len(mtbl1[selected]), hdrcomment='Number of stars for FHWM estimate')
	return round(fwhm_img[0],3)

# input files, config and params
seconfigdir ='/data7/cschoi/code/cspy/sex.config/'
seconfig    ='se1.sex'
separam     ='se1.param'
separam_noPSF = 'se1_noPSF.param'
growthparam = 'growth.param'
seconv      ='default.conv'
sennw       ='default.nnw'
DETECT_MINAREA = str(5)
DETECT_THRESH  = str(3)
DEBLEND_NTHRESH = str(32)
DEBLEND_MINCONT = str(0.005)
lowmag=13
highmag=19
filname,filerr='R','Rerr'
magtypes=['MAG_AUTO', 'MAG_PSF',
		'MAG_APER','MAG_APER_1','MAG_APER_2',
		'MAG_APER_3','MAG_APER_4','MAG_APER_5','MAG_APER_6','MAG_APER_7',
		'MAG_APER_8']
magtype=magtypes[0]
#refcat='../../ps1-Tonry-NGC3367.cat'
# source extractor command
psf=True
def secom(im,psf=psf):
    #PSCALE=fits.getheader(i)['PSCALE']
	PSCALE=pixelscale(im)
	skyval, skysig,fwhm,opt_ap=opt_ap_fwhm(im)
	aper_list,aper_list2=[3,5,7],[1.0*fwhm, 1.5*fwhm, 2.0*fwhm, 2.5*fwhm, 3.0*fwhm, opt_ap]
	aper_input = ''
	for i in aper_list: aper_input += '{},'.format(round(i/PSCALE,1))
	for i in aper_list2: aper_input += '{},'.format(round(i,1))
	aper_input = aper_input[:-1]
	fn = os.path.splitext(im)[0]
	opt1= seconfigdir+seconfig+' -CATALOG_TYPE ASCII_HEAD -CATALOG_NAME '+ fn+'.se1'
	opt2a=' -PARAMETERS_NAME '+seconfigdir+separam
	opt2b= ' -PARAMETERS_NAME '+seconfigdir+separam_noPSF
	opt2=' -FILTER_NAME '+seconfigdir+seconv +' -STARNNW_NAME '+seconfigdir+sennw
	opt3=' -DETECT_MINAREA '+ DETECT_MINAREA + ' -DETECT_THRESH '+DETECT_THRESH
	opt4=' -DEBLEND_NTHRESH '+ DEBLEND_NTHRESH +' -DEBLEND_MINCONT '+ DEBLEND_MINCONT
	opt5=' -CHECKIMAGE_TYPE SEGMENTATION,APERTURES ' +\
	 		' -CHECKIMAGE_NAME '+fn+'_seg.fits'+','+fn+'_ap.fits'
	opt5=' -CHECKIMAGE_TYPE NONE '
	opt6=' -PHOT_APERTURES '+aper_input+' '
	opt7=' -PSF_NAME '+fn+'.psf '
	opt8=' -PIXEL_SCALE '+str(PSCALE)+' '
	opt9=' -SEEING_FWHM '+str(round(PSCALE*fwhm,3))+' '
	if psf==True:
		secommand= 'sex -c '+opt1+opt2+opt2a+opt3+opt4+opt5+opt6+opt7+opt8+opt9 +im
	else:
		secommand= 'sex -c '+opt1+opt2+opt2b+opt3+opt4+opt5+opt6+opt8+opt9 +im
	print(secommand)
	#sexout = subprocess.getoutput(secommand)
	#line = [s for s in sexout.split('\n') if 'RMS' in s]
	#skymed, skysig = float(line[0].split('Background:')[1].split('RMS:')[0]), float(line[0].split('RMS:')[1].split('/')[0])
	os.system(secommand)
	return skyval, skysig

#macthing
def matching(intbl, reftbl, inra, indec, refra, refdec, sep=2.0):
    """
    MATCHING TWO CATALOG WITH RA, Dec COORD. WITH python
    INPUT   :   SE catalog, SDSS catalog file name, sepertation [arcsec]
    OUTPUT  :   MATCED CATALOG FILE & TABLE
    """
    import numpy as np
    import astropy.units as u
    from astropy.table import Table, Column
    from astropy.coordinates import SkyCoord
    from astropy.io import ascii
    incoord     = SkyCoord(inra, indec, unit=(u.deg, u.deg))
    refcoord    = SkyCoord(refra, refdec, unit=(u.deg, u.deg))
    #   INDEX FOR REF.TABLE
    indx, d2d, d3d  = incoord.match_to_catalog_sky(refcoord)
    mreftbl         = reftbl[indx]
    mreftbl['sep']  = d2d
    mergetbl        = intbl
    for col in mreftbl.colnames:
        mergetbl[col]   = mreftbl[col]
    indx_sep        = np.where(mergetbl['sep']*3600.<sep)
    mtbl            = mergetbl[indx_sep]
    #mtbl.write(mergename, format='ascii', overwrite=True)
    return mtbl

def limitmag(N, zp, aper, skysigma):			# 3? 5?, zp, diameter [pixel], skysigma
	import numpy as np
	R           = float(aper)/2.				# to radius
	braket      = N*skysigma*np.sqrt(np.pi*(R**2))
	upperlimit  = float(zp)-2.5*np.log10(braket)
	return round(upperlimit, 3)

def starcut(mtbl,lowmag=lowmag,highmag=highmag,filname=filname,magtype=magtype):
	idx=np.where( (mtbl['SNR_WIN'] >10) &
				(mtbl['FLAGS'] == 0) &
				(mtbl[filname] < highmag) &
				(mtbl[filname] > lowmag) &
				(mtbl[magtype[:3]+'ERR'+magtype[3:]]<0.5)		)
	return mtbl[idx]

#zp calculation
magtypes=['MAG_AUTO', 'MAG_PSF',
		'MAG_APER','MAG_APER_1','MAG_APER_2',
		'MAG_APER_3','MAG_APER_4','MAG_APER_5','MAG_APER_6','MAG_APER_7',
		'MAG_APER_8']

def zpcal(mtbl1,filname, magtype):
    zp=mtbl1[filname]-mtbl1[magtype]
    #zp3=sigma_clipped_stats(zp, sigma=3, maxiters=10)
    zp2=sigma_clipped_stats(zp, sigma=2, maxiters=10)
    print ('zp ',zp2[0], 'zp err',zp2[2])
    filtered_data=sigma_clip(zp,sigma=2,maxiters=10)
    selected, nonselected= ~filtered_data.mask, filtered_data.mask
    zperrp=np.sqrt( np.sum(mtbl1[filerr][selected]**2 + \
						mtbl1[magtype[:3]+'ERR'+magtype[3:]][selected]**2)\
						/ len(mtbl1) )
    print(magtype, 'zp', '{},'.format(round(zp2[0],3)),
	 	'zperr', '{},'.format(round(zperrp,3)),
		len(mtbl1[selected]),'stars from',len(mtbl1))
    return zp2, selected,zperrp

def fwhm_img(im,mtbl1):
	fwhm_img=sigma_clipped_stats(mtbl1['FWHM_IMAGE'], sigma=3, maxiters=10)
	filtered_data=sigma_clip(mtbl1['FWHM_IMAGE'],sigma=3,maxiters=10)
	selected, nonselected= ~filtered_data.mask, filtered_data.mask
	print('FWHM_IMAGE','{}'.format(round(fwhm_img[0],3)),
		len(mtbl1[selected]),'stars from',len(mtbl1))
	puthdr(im, 'FWHM_PIX', round(fwhm_img[0],3), hdrcomment='FWHM PIXEL VALUE')
	return fwhm_img[0]

'''
def fwhm_wcs(mtbl1):
	fwhm_wcs=sigma_clipped_stats(mtbl1['FWHM_WORLD'], sigma=3, maxiters=10)
	filtered_data=sigma_clip(mtbl1['FWHM_WORLD'],sigma=3,maxiters=10)
	selected, nonselected= ~filtered_data.mask, filtered_data.mask
	print('FWHM_WORLD','{},'.format(round(fwhm_wcs[0]*3600,3)),
		len(mtbl1[selected]),'stars from',len(mtbl1))
	return fwhm_wcs[0]*3600
'''
# 5sigma detection limit estimate for MAG_AUTO, MAG_PSF
# error fitting polinomial
def UL_5sig_err(im,setbl,mtbl,mtbl1,magtype,zp2):
	fn=os.path.splitext(im)[0]
	from astropy.modeling import models, fitting
	import numpy as np
	magerrtype = magtype[:3]+'ERR'+magtype[3:]
	if len(mtbl1) < 3:
		return 0
	x,y = mtbl1[magtype],mtbl1[magerrtype]
	#x,y=setbl['MAG_AUTO'],setbl['MAGERR_AUTO']
	x,y=x[np.where(y<1)],y[np.where(y<1)]
	#fit_init=models.Polynomial1D(7)
	fit_init=models.Exponential1D()
	fit_t=fitting.LevMarLSQFitter()
	#fit_t=fitting.LinearLSQFitter()
	t=fit_t(fit_init,x,y)
	plt.plot(setbl[magtype]+zp2[0],setbl[magerrtype],'ro')
	plt.plot(mtbl[magtype]+zp2[0],mtbl[magerrtype],'ko')
	plt.plot(x+zp2[0],y,'bo')
	plt.xlabel('Mag')
	plt.ylabel('Error')
	plt.xlim(10,25)
	plt.ylim(-0.1,0.5)
	xp=np.linspace(-20,0,20001)
	plt.plot(xp+zp2[0],t(xp),'--')
	plt.hlines(0.2,10,25)
	idx_min=np.where(np.abs(t(xp)-0.198) == np.min(np.abs(t(xp)-0.198)))
	xp[idx_min]+zp2[0] #
	plt.vlines(xp[idx_min]+zp2[0],-0.1,0.5)
	# result print, file, header
	plt.text(12,0.4,'5 sigma Detection Limit error=0.198')
	plt.text(12,0.3,'5sig_UL = '+'{}'.format(round((xp[idx_min]+zp2[0])[0],3)))
	plt.title(fn+ ' '+ magtype+' '+'5 sig Detection Limit')
	plt.savefig(fn+'_'+magtype+'_'+'5sigUL.png')
	plt.close()
	return round((xp[idx_min]+zp2[0])[0],3)

def zp_plot(mtbl1, zp2, selected, magtype, im, filname=filname, filerr=filerr):
	fn=os.path.splitext(im)[0]
	zp=mtbl1[filname]-mtbl1[magtype]
	xr=np.linspace(np.min(mtbl1[filname]), np.max(mtbl1[filname]), len(zp))
	magerrtype=magtype[:3]+'ERR'+magtype[3:]
	zperrp=np.sqrt( np.sum(mtbl1[filerr][selected]**2 + \
						mtbl1[magtype[:3]+'ERR'+magtype[3:]][selected]**2)\
						/ len(mtbl1) )
#	plt.plot(mtbl1[filname],zp,'o',markersize=12)
	plt.errorbar(mtbl1[filname],zp,yerr=zperrp,fmt='bo',markersize=12)
	plt.ylim(zp2[0]-1,zp2[0]+1)
	plt.xlim(np.min(mtbl1[filname]),np.max(mtbl1[filname]))
	#plt.hlines(zp3[0],xmin=12,xmax=20,color='b')
	plt.hlines(zp2[0],xmin=12,xmax=20,color='r')
	sig2=np.ones(len(mtbl1))*zp2[2]
	zp2a=np.ones(len(mtbl1))*zp2[0]
	#plt.fill_between(xr,zp3a+sig3,zp3a-sig3,color='b',alpha=0.5)
	plt.fill_between(xr,zp2a+sig2,zp2a-sig2,color='yellow',alpha=0.5)
	plt.plot(mtbl1[filname][~selected],zp[~selected],'ro')
	plt.plot(mtbl1[filname][selected],zp[selected],'ko')
	plt.text(18,zp2[0]+0.8,magtype+' ZP '+str(round(zp2[0],3))+' err '+str(round(zp2[2],3)))
	plt.title(fn+', '+filname+' '+magtype)
	plt.ylabel('Zeropoint (AB Mag)')
	plt.xlabel(filname +' Reference Mag (AB)')
	plt.savefig(fn+'_'+filname+'_'+magtype+'_zp.png')
	plt.close()

# fits to png with regions
def fitplot(im, mtbl1, magtype, selected):
	import matplotlib.pyplot as plt
	from astropy.wcs import WCS
	from astropy.io import fits
	from astropy.visualization import MinMaxInterval,ZScaleInterval,PercentileInterval
	from astropy.visualization import SqrtStretch,LinearStretch
	from astropy.visualization import ImageNormalize
	imdata,imhdr=fits.getdata(im,header=True)
	norm = ImageNormalize(imdata,
			interval=ZScaleInterval(),
			stretch=LinearStretch() )
	wcs=WCS(imhdr)
	fig,ax=plt.subplots()
	ax=plt.subplot(projection=wcs)
	ax.set_xlabel('RA')
	ax.set_ylabel('DEC')
	#ax.invert_xaxis()
	ax.set_title(im+' '+magtype)
	ax.scatter(mtbl1[selected]['ALPHA_J2000'],mtbl1[selected]['DELTA_J2000'],
		transform=ax.get_transform('fk5'),s=20, edgecolor='green',facecolor='none')
	ax.scatter(mtbl1[~selected]['ALPHA_J2000'],mtbl1[~selected]['DELTA_J2000'],
		transform=ax.get_transform('fk5'),s=20, edgecolor='red',facecolor='none')
	img=ax.imshow(imdata,cmap='gray',norm=norm,origin='lower')
	ax.invert_yaxis()
	fig.colorbar(img)
	plt.savefig(os.path.splitext(im)[0]+'_'+magtype+'_FOV.png')
	plt.close()

'''
	param_dict={
	'CATALOG_NAME'     : fn+'.se1',
	'PARAMETERS_NAME'  : seconfigdir+'se1.param',
	'DETECT_MINAREA'   : DETECT_MINAREA,
	'DETECT_THRESH'    : DETECT_THRESH,
	'FILTER_NAME'      : seconfigdir+seconv,
	'DEBLEND_NTHRESH'  : DEBLEND_NTHRESH,
	'DEBLEND_MINCONT'  : DEBLEND_MINCONT,
	'BACK_TYPE'        : 'AUTO',
	'BACK_DEFAULT'     : '0.0',
	'BACK_SIZE'        : '64',
	'BACK_FILTERSIZE'  : '3',
	'BACKPHOTO_TYPE'   : 'LOCAL',
	'PHOT_APERTURES'   : aper_input,
	'SATUR_LEVEL'      : '60000',
	'GAIN'             : '1.0',
	'PIXEL_SCALE'      : str(PSCALE),
	'SEEING_FWHM'      : '1.2',
	'STARNNW_NAME'     :seconfigdir+sennw,
	'CHECKIMAGE_TYPE'  : 'SEGMENTATION,OBJECTS,BACKGROUND',
	'CHECKIMAGE_NAME'  : fn+'_seg.fits,'+fn+'_obj.fits,'+fn+'bg.fits',
	'PSF_NAME'         : fn+'.psf'
	}
	optstr=''
	for i in param_dict:
		#print(' -{} {}'.format(i,param_dict[i]))
		optstr += ' -{} {}'.format(i,param_dict[i])
	secom='sex -c '+seconfigdir+seconfig +' '+ im + optstr
'''
'''
# Default configuration file for SExtractor 2.19.5
# EB 2014-03-19
#

#-------------------------------- Catalog ------------------------------------

CATALOG_NAME     test.cat       # name of the output catalog
CATALOG_TYPE     ASCII_HEAD     # NONE,ASCII,ASCII_HEAD, ASCII_SKYCAT,
                                # ASCII_VOTABLE, FITS_1.0 or FITS_LDAC
PARAMETERS_NAME  default.param  # name of the file containing catalog contents

#------------------------------- Extraction ----------------------------------

DETECT_TYPE      CCD            # CCD (linear) or PHOTO (with gamma correction)
DETECT_MINAREA   5              # min. # of pixels above threshold
DETECT_MAXAREA   0              # max. # of pixels above threshold (0=unlimited)
THRESH_TYPE      RELATIVE       # threshold type: RELATIVE (in sigmas)
                                # or ABSOLUTE (in ADUs)
DETECT_THRESH    1.5            # <sigmas> or <threshold>,<ZP> in mag.arcsec-2
ANALYSIS_THRESH  1.5            # <sigmas> or <threshold>,<ZP> in mag.arcsec-2

FILTER           Y              # apply filter for detection (Y or N)?
FILTER_NAME      default.conv   # name of the file containing the filter
FILTER_THRESH                   # Threshold[s] for retina filtering

DEBLEND_NTHRESH  32             # Number of deblending sub-thresholds
DEBLEND_MINCONT  0.005          # Minimum contrast parameter for deblending

CLEAN            Y              # Clean spurious detections? (Y or N)?
CLEAN_PARAM      1.0            # Cleaning efficiency

MASK_TYPE        CORRECT        # type of detection MASKing: can be one of
                                # NONE, BLANK or CORRECT

#-------------------------------- WEIGHTing ----------------------------------

WEIGHT_TYPE      NONE           # type of WEIGHTing: NONE, BACKGROUND,
                                # MAP_RMS, MAP_VAR or MAP_WEIGHT
RESCALE_WEIGHTS  Y              # Rescale input weights/variances (Y/N)?
WEIGHT_IMAGE     weight.fits    # weight-map filename
WEIGHT_GAIN      Y              # modulate gain (E/ADU) with weights? (Y/N)
WEIGHT_THRESH                   # weight threshold[s] for bad pixels

#-------------------------------- FLAGging -----------------------------------

FLAG_IMAGE       flag.fits      # filename for an input FLAG-image
FLAG_TYPE        OR             # flag pixel combination: OR, AND, MIN, MAX
                                # or MOST

#------------------------------ Photometry -----------------------------------

PHOT_APERTURES   5              # MAG_APER aperture diameter(s) in pixels
PHOT_AUTOPARAMS  2.5, 3.5       # MAG_AUTO parameters: <Kron_fact>,<min_radius>
PHOT_PETROPARAMS 2.0, 3.5       # MAG_PETRO parameters: <Petrosian_fact>,
                                # <min_radius>
PHOT_AUTOAPERS   0.0,0.0        # <estimation>,<measurement> minimum apertures
                                # for MAG_AUTO and MAG_PETRO
PHOT_FLUXFRAC    0.5            # flux fraction[s] used for FLUX_RADIUS

SATUR_LEVEL      50000.0        # level (in ADUs) at which arises saturation
SATUR_KEY        SATURATE       # keyword for saturation level (in ADUs)

MAG_ZEROPOINT    0.0            # magnitude zero-point
MAG_GAMMA        4.0            # gamma of emulsion (for photographic scans)
GAIN             0.0            # detector gain in e-/ADU
GAIN_KEY         GAIN           # keyword for detector gain in e-/ADU
PIXEL_SCALE      1.0            # size of pixel in arcsec (0=use FITS WCS info)

#------------------------- Star/Galaxy Separation ----------------------------

SEEING_FWHM      1.2            # stellar FWHM in arcsec
STARNNW_NAME     default.nnw    # Neural-Network_Weight table filename

#------------------------------ Background -----------------------------------

BACK_TYPE        AUTO           # AUTO or MANUAL
BACK_VALUE       0.0            # Default background value in MANUAL mode
BACK_SIZE        64             # Background mesh: <size> or <width>,<height>
BACK_FILTERSIZE  3              # Background filter: <size> or <width>,<height>

BACKPHOTO_TYPE   GLOBAL         # can be GLOBAL or LOCAL
BACKPHOTO_THICK  24             # thickness of the background LOCAL annulus
BACK_FILTTHRESH  0.0            # Threshold above which the background-
                                # map filter operates

#------------------------------ Check Image ----------------------------------

CHECKIMAGE_TYPE  NONE           # can be NONE, BACKGROUND, BACKGROUND_RMS,
                                # MINIBACKGROUND, MINIBACK_RMS, -BACKGROUND,
                                # FILTERED, OBJECTS, -OBJECTS, SEGMENTATION,
                                # or APERTURES
CHECKIMAGE_NAME  check.fits     # Filename for the check-image

#--------------------- Memory (change with caution!) -------------------------

MEMORY_OBJSTACK  3000           # number of objects in stack
MEMORY_PIXSTACK  300000         # number of pixels in stack
MEMORY_BUFSIZE   1024           # number of lines in buffer

#------------------------------- ASSOCiation ---------------------------------

ASSOC_NAME       sky.list       # name of the ASCII file to ASSOCiate
ASSOC_DATA       2,3,4          # columns of the data to replicate (0=all)
ASSOC_PARAMS     2,3,4          # columns of xpos,ypos[,mag]
ASSOCCOORD_TYPE  PIXEL          # ASSOC coordinates: PIXEL or WORLD
ASSOC_RADIUS     2.0            # cross-matching radius (pixels)
ASSOC_TYPE       NEAREST        # ASSOCiation method: FIRST, NEAREST, MEAN,
                                # MAG_MEAN, SUM, MAG_SUM, MIN or MAX
ASSOCSELEC_TYPE  MATCHED        # ASSOC selection type: ALL, MATCHED or -MATCHED

#----------------------------- Miscellaneous ---------------------------------

VERBOSE_TYPE     NORMAL         # can be QUIET, NORMAL or FULL
HEADER_SUFFIX    .head          # Filename extension for additional headers
WRITE_XML        N              # Write XML file (Y/N)?
XML_NAME         sex.xml        # Filename for XML output
XSL_URL          file:///usr/share/sextractor/sextractor.xsl
                                # Filename for XSL style-sheet
NTHREADS         1              # 1 single thread

FITS_UNSIGNED    N              # Treat FITS integer values as unsigned (Y/N)?
INTERP_MAXXLAG   16             # Max. lag along X for 0-weight interpolation
INTERP_MAXYLAG   16             # Max. lag along Y for 0-weight interpolation
INTERP_TYPE      ALL            # Interpolation type: NONE, VAR_ONLY or ALL

#--------------------------- Experimental Stuff -----------------------------

PSF_NAME         default.psf    # File containing the PSF model
PSF_NMAX         1              # Max.number of PSFs fitted simultaneously
PATTERN_TYPE     RINGS-HARMONIC # can RINGS-QUADPOLE, RINGS-OCTOPOLE,
                                # RINGS-HARMONICS or GAUSS-LAGUERRE
SOM_NAME         default.som    # File containing Self-Organizing Map weights
'''


<reponame>adam-coogan/diffbank
from itertools import product

import click
from diffjeom import get_ricci_scalar
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
from scipy.interpolate import griddata

from diffbank.constants import MSUN
from diffbank.metric import get_g
from diffbank.noise import Sn_O3a as Sn
from diffbank.utils import ms_to_Mc_eta
from diffbank.waveforms.twoPN_chirptimes import (
    Amp,
    Psi,
    get_th_boundary_interps,
    phys_to_th,
)

plt.style.use("../plot_style.mplstyle")

"""
Plots the scalar curvature (Ricci scalar) for the metric for the waveform in
https://arxiv.org/abs/gr-qc/0604037.

To reproduce:

    >>> python plot_scalar_curvature.py

"""

##### Frequency settings
f_u = 512.0  # Hz
f_l = 10.0  # Hz
N_fbins = 1000
#####

@click.command()
@click.option("--n-m1s", type=int, default=200)
@click.option("--n-m2s", type=int, default=100)
@click.option("--n-th0s", type=int, default=200)
@click.option("--n-th3s", type=int, default=100)
@click.option("--fig-path", default="figures/scalar-curvature.pdf")
def run(n_m1s, n_m2s, n_th0s, n_th3s, fig_path):
    fs = jnp.linspace(f_l, f_u, N_fbins)

    g_fun = lambda theta: get_g(theta, Amp, Psi, fs, Sn)

    # Set parameter grid
    m_min = jnp.array(1.0) * MSUN
    m_max = jnp.array(20.0) * MSUN
    M_max = m_min + m_max
    m1s = jnp.geomspace(m_min, m_max, n_m1s)
    m2s = jnp.geomspace(m_min, m_max, n_m2s)
    m1s, m2s = jnp.array(list(product(m1s, m2s))).T
    m1s, m2s = m1s[m1s >= m2s], m2s[m1s >= m2s]  # remove redundant systems
    m1s, m2s = m1s[m1s + m2s <= M_max], m2s[m1s + m2s <= M_max]
    M_chirps, etas = ms_to_Mc_eta(jnp.stack([m1s, m2s]))

    # (th0, th3) boundaries
    (th0_min, th0_max), bound_itp_low, bound_itp_high = get_th_boundary_interps(
        m_min, m_max, f_l
    )

    # Plotting configuration
    vmin = -10.0
    vmax = -7.5
    levels = jnp.linspace(vmin, vmax, 60)
    cbar_ticks = jnp.arange(vmin, vmax + 0.05, 0.5)

    thetas = phys_to_th(jnp.stack([M_chirps, etas]), f_l).T  # type: ignore
    Rss = jax.vmap(lambda x: get_ricci_scalar(x, g_fun))(thetas)

    # Plot!
    th0s = jnp.linspace(thetas[:, 0].min(), thetas[:, 0].max(), n_th0s)
    th3s = jnp.linspace(thetas[:, 1].min(), thetas[:, 1].max(), n_th3s)

    cs = plt.contourf(
        th0s / 1e4,
        th3s / 1e2,
        jnp.clip(
            griddata(
                thetas[:, :2],
                jnp.log10(jnp.abs(Rss)),
                jnp.stack(jnp.meshgrid(th0s, th3s)).reshape([2, -1]).T,
            ).reshape([len(th3s), len(th0s)]),
            vmin,
            vmax,
        ),
        levels=levels,
        cmap="viridis",
    )
    plt.colorbar(cs, label=r"$\log_{10}(|R|)$", ticks=cbar_ticks)

    # Mask outside boundaries
    th0_grid = jnp.linspace(th0_min, th0_max, 200)
    plt.fill_between(
        th0_grid / 1e4,
        bound_itp_low(th0_grid) / 1e2,
        jnp.full_like(th0_grid, -1e3) / 1e2,
        where=jnp.full_like(th0_grid, True),
        color="w",
    )
    plt.fill_between(
        th0_grid / 1e4,
        bound_itp_high(th0_grid) / 1e2,
        jnp.full_like(th0_grid, 1e3) / 1e2,
        where=jnp.full_like(th0_grid, True),
        color="w",
    )

    plt.xlabel(r"$\theta_0 / 10^4$")
    plt.ylabel(r"$\theta_3 / 10^2$")
    plt.xlim(0.0, 10.0)
    plt.ylim(0.8, 8)
    # plt.tight_layout()
    plt.savefig(fig_path, bbox_inches="tight")

if __name__ == "__main__":
    run()


#!/usr/bin/env python3
import argparse
import array
import asyncio as aio
import base64
import itertools
import logging
import random
import statistics
from abc import ABCMeta, abstractmethod
from asyncio import DatagramTransport, Lock, StreamReader, StreamWriter, Task
from types import TracebackType
from typing import (ClassVar, Iterable, Iterator, List, Optional, Sequence,
                    Set, SupportsFloat, Tuple, Type, Union)

import httpx

DEFAULT_LISTEN_ADDRESSES = \
[
	'127.0.0.1',
	'::1',
]

DEFAULT_LISTEN_PORTS = \
[
	5053,
]

DEFAULT_UPSTREAMS = \
[
	'https://192.168.3.11:443/dns-query',
	'https://172.16.31.10:443/dns-query',
	'https://[fc00:db20:35b:7399::5]:443/dns-query',
	'https://[2606:fd00:a516:7c1b:17cd:6d81:2137:bd2a:2c5b]:443/dns-query',
]

async def main(args) -> None:
	# Setup event loop
	loop = aio.get_running_loop()

	# Setup DNS resolver to cache/forward queries and answers
	async with AsyncDnsResolver(args.upstreams, AsyncDohUpstreamContext) as resolver:
		transports = []

		# Setup listening transports
		for addr in args.listen_address:
			for port in args.listen_port:
				# Setup UDP server
				logging.info('Starting UDP server listening on %s#%d' % (addr, port))
				udp, _ = await loop.create_datagram_endpoint(lambda: UdpResolverProtocol(resolver), local_addr=(addr, port))
				transports.append(udp)

				# Setup TCP server
				if args.tcp:
					logging.info('Starting TCP server listening on %s#%d' % (addr, port))
					tcp = await aio.start_server(TcpResolverProtocol(resolver).ahandle_peer, addr, port)
					transports.append(tcp)

		# Serve forever
		try:
			while True:
				await aio.sleep(3600)
				logging.info(resolver.get_stats())

		except (KeyboardInterrupt, SystemExit):
			pass

		logging.info('Shutting down DNS over HTTPS forwarder')

		wait_closers = []
		for transport in transports:
			transport.close()
			if hasattr(transport, 'wait_closed'):
				wait_closers.append(aio.create_task(transport.wait_closed()))

		await aio.wait(wait_closers)

	await aio.sleep(0.3)

class AsyncDnsUpstreamContext(metaclass=ABCMeta):
	"""A base class used to manage upstream DNS server connections and metadata."""

	RTT_WINDOW_SIZE: ClassVar[int] = 10

	def __init__(self, host: str) -> None:
		self.host = host
		self.queries = 0
		self.answers = 0
		self._rtts = array.array('d', [0.0])
		self._rtts_index: Iterator[int] = itertools.cycle(range(self.RTT_WINDOW_SIZE))

	async def __aenter__(self) -> 'AsyncDnsUpstreamContext':
		return self

	async def __aexit__(self,
		exc_type: Optional[Type[BaseException]],
		exc_val: Optional[BaseException],
		exc_tb: Optional[TracebackType]) -> None:
		await self.aclose()

	@property
	def avg_rtt(self) -> float:
		"""The average rtt or latency (in seconds) for DNS requests to this upstream DNS server."""
		return statistics.fmean(self._rtts)

	def add_rtt_sample(self, rtt: SupportsFloat) -> None:
		"""Add a new rtt sample to help compute the average rtt for this upstream DNS server."""
		i = next(self._rtts_index)
		self._rtts[i:i+1] = array.array('d', [float(rtt)])

	def get_stats(self) -> str:
		"""Returns a formatted string of statistics for this upstream server."""
		return f'{self.host} (rtt: {self.avg_rtt:.3f} s, queries: {self.queries}, answers: {self.answers})'

	@abstractmethod
	async def aforward_query(self, query: bytes) -> bytes:
		"""Resolve a DNS query via forwarding to a upstream DNS server.

		Params:
			query - A wireformat DNS query packet.

		Returns:
			A wireformat DNS answer packet.

		Notes:
			This coroutine is be safely cancellable. That is, even if the
			coroutine is cancelled it still leaves any internal state
			it uses in a consistent and usable state.
		"""
		...

	@abstractmethod
	async def aclose(self) -> None:
		"""Close any open connections to the upstream DNS server."""
		...

class AsyncDohUpstreamContext(AsyncDnsUpstreamContext):
	"""A class used to manage upstream DoH server connections and metadata."""

	SESSION_LIMITS: ClassVar[httpx.Limits] = httpx.Limits(max_keepalive_connections=1, max_connections=3, keepalive_expiry=60.0)
	SESSION_TIMEOUTS: ClassVar[httpx.Timeout] = httpx.Timeout(None)

	def __init__(self, url: str) -> None:
		super().__init__(url)
		self.session = httpx.AsyncClient(
			limits=self.SESSION_LIMITS,
			timeout=self.SESSION_TIMEOUTS,
			headers={'accept': 'application/dns-message'},
			http2=True)

	async def aforward_post(self, query: bytes) -> Tuple[bytes, float]:
		"""Resolve a DNS query via forwarding to a upstream DoH server (POST).

		Params:
			query - A wireformat DNS query packet.

		Returns:
			A wireformat DNS answer packet and rtt sample.

		Notes:
			Using DNS over HTTPS POST format as described here:
			https://datatracker.ietf.org/doc/html/rfc8484
			https://developers.cloudflare.com/1.1.1.1/dns-over-https/wireformat/
		"""
		# Send HTTP request to upstream DoH server and wait for the response
		response = await aio.shield(
			self.session.post(
				self.host,
				headers={'content-type': 'application/dns-message'},
				content=query))

		# Parse HTTP response
		response.raise_for_status()
		answer = response.read()
		rtt = response.elapsed.total_seconds()

		# Return the DNS answer
		return (answer, rtt)

	async def aforward_get(self, query: bytes) -> Tuple[bytes, float]:
		"""Resolve a DNS query via forwarding to a upstream DoH server (GET).

		Params:
			query - A wireformat DNS query packet.

		Returns:
			A wireformat DNS answer packet and rtt sample.

		Notes:
			Using DNS over HTTPS GET format as described here:
			https://datatracker.ietf.org/doc/html/rfc8484
			https://developers.cloudflare.com/1.1.1.1/dns-over-https/wireformat/
		"""
		# Encode DNS query into url
		url = ''.join([self.host, '?dns=', base64.urlsafe_b64encode(query).rstrip(b'=').decode()])

		# Send HTTP request to upstream DoH server and wait for the response
		response = await aio.shield(self.session.get(url))

		# Parse HTTP response
		response.raise_for_status()
		answer = response.read()
		rtt = response.elapsed.total_seconds()

		# Return the DNS answer
		return (answer, rtt)

	async def aforward_query(self, query: bytes) -> bytes:
		self.queries += 1

		query = memoryview(query)
		qid = query[:2]

		# Forward the DNS query to the upstream DoH server
		try:
			logging.debug(f'Sending query {qid.hex()} to {self.host} --->')
			answer, rtt = await self.aforward_post(b''.join([b'\0' * 2, query[2:]]))

			self.add_rtt_sample(rtt)
			self.answers += 1

			logging.debug(f'Receiving answer {qid.hex()} from {self.host} ({rtt}) <---')
			return b''.join([qid, memoryview(answer)[2:]])

		# Raise connection error
		except (httpx.NetworkError, httpx.RemoteProtocolError):
			raise ConnectionError(f'DNS query to DoH server {self.host} failed due to network errors')

		# Raise abnormal HTTP status codes
		except httpx.HTTPStatusError as exc:
			raise ConnectionError(f'received HTTP error status from DoH server {self.host} ({exc.response.status_code})')

	async def aclose(self) -> None:
		await self.session.aclose()

class AsyncDnsResolver:
	"""A class that manages upstream DNS server contexts and resolves DNS queries."""

	DEFAULT_QUERY_TIMEOUT: ClassVar[float] = 3.0

	def __init__(self, upstreams: Iterable[str], context_class: Type[AsyncDnsUpstreamContext]) -> None:
		self._upstreams = tuple(context_class(upstream) for upstream in upstreams)

		if not self._upstreams:
			raise ValueError('iterable of upstreams must have at least one entry')

	async def __aenter__(self) -> 'AsyncDnsResolver':
		return self

	async def __aexit__(self,
		exc_type: Optional[Type[BaseException]],
		exc_val: Optional[BaseException],
		exc_tb: Optional[TracebackType]) -> None:
		await self.aclose()

	@property
	def queries(self) -> int:
		return sum(upstream.queries for upstream in self._upstreams)

	@property
	def answers(self) -> int:
		return sum(upstream.answers for upstream in self._upstreams)

	@property
	def avg_rtt(self) -> float:
		return statistics.fmean(upstream.avg_rtt for upstream in self._upstreams)

	def get_stats(self) -> str:
		"""Returns a formatted string of statistics for this resolver."""
		return f'Statistics for resolver at 0x{id(self)} (avg_rtt: {self.avg_rtt:.3f} s, total_queries: {self.queries}, total_answers: {self.answers})'

	async def aresolve(self, query: bytes, timeout: float = DEFAULT_QUERY_TIMEOUT) -> bytes:
		"""Resolve a DNS query via forwarding to upstream DNS servers.

		Params:
			query - A wireformat DNS query packet.
			timeout - The maximum amount of time (in seconds) to wait for the receipt of the DNS answer packet.

		Returns:
			A wireformat DNS answer packet.
		"""
		# Forward the DNS query and return the DNS answer
		# (perform a staggered race and accept the earliest response)
		async def astaggered_resolution(upstreams: Sequence[AsyncDnsUpstreamContext], period: float) -> bytes:
			assert len(upstreams) > 0

			winner: Task = None
			racers: Set[Task] = set()
			errors: List[BaseException] = []

			# Wait for the first racer to finish and cleanup exceptions
			async def await_first_racer(timeout: float = None) -> bool:
				nonlocal winner
				nonlocal racers
				nonlocal errors

				done, racers = await aio.wait(racers, timeout=timeout, return_when=aio.FIRST_COMPLETED)

				for racer in done:
						error = racer.exception()

						if error is None:
							winner = racer
							break
						else:
							errors.append(error)

				return winner is not None

			try:
				for upstream in upstreams:
					racers.add(aio.create_task(upstream.aforward_query(query)))

					if await await_first_racer(period):
						return winner.result()

				while racers:
					if await await_first_racer():
						return winner.result()

			finally:
				for loser in racers:
					loser.cancel()

			def raise_multi_error(errors: Iterable[BaseException]) -> None:
				class MultiError(*frozenset(type(error) for error in errors)):
					pass

				raise MultiError

			assert len(errors) == len(upstreams)
			raise_multi_error(errors)

		# Weighted random shuffle the upstream servers by average latency
		k = len(self._upstreams)
		rtts = tuple(upstream.avg_rtt for upstream in self._upstreams)
		max_rtt = max(rtts)
		weights = (max_rtt - rtt + 0.001 for rtt in rtts)
		upstreams = random.choices(self._upstreams, weights=weights, k=k)
		period = (timeout / 2) / k if timeout is not None else 0.1

		# Forward the DNS query and return the DNS answer
		try:
			return await aio.wait_for(astaggered_resolution(upstreams, period), timeout)

		# Raise timeout error
		except aio.TimeoutError:
			raise TimeoutError(f'DNS query expired and was cancelled')

	async def aclose(self) -> None:
		"""Close all upstream DoH server connections."""
		for upstream in self._upstreams:
			await upstream.aclose()

class UdpResolverProtocol(aio.DatagramProtocol):
	"""Protocol for serving UDP DNS requests via a DnsResolver instance."""

	def __init__(self, resolver: AsyncDnsResolver) -> None:
		self.resolver = resolver
		self.buffer = []
		self.worker = None

	def connection_made(self, transport: DatagramTransport) -> None:
		self.transport = transport

	def datagram_received(self, data: bytes, peer: Tuple[str, int]) -> None:
		logging.debug(f'Got UDP DNS query from {peer}')

		# Add query to buffer
		self.buffer.append((peer, data))

		# Schedule query processing task if necessary
		if self.worker is None:
			self.worker = aio.create_task(self.ahandle_queries())

	async def ahandle_queries(self) -> None:
		while self.buffer:
			tasks = set(aio.create_task(self.ahandle_query(peer, query)) for peer, query in self.buffer)

			del self.buffer[:]

			while tasks:
				done, tasks = await aio.wait(tasks, timeout=0.05)

				for task in done:
					error = task.exception()

					if error is None:
						peer, answer = task.result()
						self.transport.sendto(answer, peer)
					else:
						logging.warning(f'UDP DNS query resolution encountered an error - {error!r}')

		self.worker = None

	async def ahandle_query(self, peer: Tuple[str, int], query: bytes) -> Tuple[Tuple[str, int], bytes]:
		return (peer, await self.resolver.aresolve(query))

class TcpResolverProtocol:
	"""Protocol for serving TCP DNS requests via a DnsResolver instance."""

	def __init__(self, resolver: AsyncDnsResolver) -> None:
		self.resolver = resolver

	async def ahandle_peer(self, reader: StreamReader, writer: StreamWriter) -> None:
		"""Read all DNS queries from the peer stream and schedule their resolution via a DnsResolver instance."""
		tasks: Union[List[Task], Set[Task]] = []
		wlock = aio.Lock()

		logging.debug(f'Got TCP DNS query stream from {writer.transport.get_extra_info("peername")}')

		while True:
			# Parse a DNS query packet off of the wire
			try:
				query_size = int.from_bytes(await reader.readexactly(2), 'big')
				query = await reader.readexactly(query_size)

			# Check if our peer has finished writing to the stream
			except aio.IncompleteReadError:
				break

			# Schedule the processing of the query
			tasks.append(aio.create_task(self.ahandle_query(writer, wlock, query)))

		# Wait for all scheduled query processing to finish
		while tasks:
			done, tasks = await aio.wait(tasks, return_when=aio.FIRST_COMPLETED)

			for task in done:
				error = task.exception()

				if error is not None:
					logging.warning(f'TCP DNS query resolution encountered an error - {error!r}')

		if not writer.is_closing():
			# Indicate we are done writing to the stream
			if writer.can_write_eof():
				writer.write_eof()

			# Close the stream
			writer.close()
			await writer.wait_closed()

	async def ahandle_query(self, writer: StreamWriter, wlock: Lock, query: bytes) -> None:
		"""Resolve a DNS query and write the DNS answer to the peer stream."""
		if writer.is_closing():
			return

		# Resolve DNS query
		answer = await self.resolver.aresolve(query)

		if writer.is_closing():
			return

		# Create the DNS answer packet
		answer_size = len(answer).to_bytes(2, 'big')
		answer = b''.join([answer_size, answer])

		# Write the DNS answer to the peer stream
		async with wlock:
			if writer.is_closing():
				return

			await writer.drain()

			if writer.is_closing():
				return

			writer.write(answer)

if __name__ == '__main__':
	# Handle command line arguments
	parser = argparse.ArgumentParser()
	parser.add_argument('-l', '--listen-address', nargs='+', default=DEFAULT_LISTEN_ADDRESSES,
						help='addresses to listen on for DNS over HTTPS requests (default: %(default)s)')
	parser.add_argument('-p', '--listen-port', nargs='+', type=int, default=DEFAULT_LISTEN_PORTS,
						help='ports to listen on for DNS over HTTPS requests (default: %(default)s)')
	parser.add_argument('-u', '--upstreams', nargs='+', default=DEFAULT_UPSTREAMS,
						help='upstream servers to forward DNS queries and requests to (default: %(default)s)')
	parser.add_argument('-t', '--tcp', action='store_true', default=False,
						help='serve TCP based queries and requests along with UDP (default: %(default)s)')
	parser.add_argument('-f', '--file', default=None,
						help='file to store logging output to (default: %(default)s)')
	parser.add_argument('-d', '--debug', action='store_true', default=False,
						help='enable debugging on the internal asyncio event loop (default: %(default)s)')
	args = parser.parse_args()

	# Setup logging
	log_level = 'DEBUG' if args.debug else 'INFO'
	logging.basicConfig(level=log_level, filename=args.file, filemode='w', format='(%(asctime)s)[%(levelname)s] %(message)s')
	logging.info('Starting DNS over HTTPS forwarder')
	logging.info('Args: %r' % (vars(args)))

	aio.run(main(args), debug=args.debug)


"""
Question Source:Leetcode
Level: Hard
Topic: Binary search
Solver: Tayyrov
Date: 25.05.2022
"""
from bisect import bisect_left
from cmath import inf
from typing import List

def maxEnvelopes(self, envelopes: List[List[int]]) -> int:
    sor = sorted(envelopes, key=lambda x: (x[0], -x[-1]))

    ans = [-inf]

    for _, h in sor:
        idx = bisect_left(ans, h)
        print(idx)
        if idx == len(ans):
            ans.append(h)
        else:
            ans[idx] = h
        print(ans)
    return len(ans) - 1


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import pandas as pd
import numpy as np
from scipy import interpolate

# Download Task (10 MB)
# Egger, Sebastian, et al. "“Time is bandwidth”?
# Narrowing the gap between subjective time perception and Quality of Experience."
# Communications (ICC), 2012 IEEE International Conference on. IEEE, 2012.
class WebDLUtility(object):

    def __init__(self, scaled=False, min_mos=1, max_mos=5):

        self._scaled = scaled

        self._min_mos = min_mos
        self._max_mos = max_mos

        # Create inverse by interpolation
        x = np.linspace(0.1, 150)
        y = list(map(self.predict, x))

        self._finv = interpolate.interp1d(y, x,
                                 fill_value="extrapolate",
                                 bounds_error=False)

    def predict(self, dl):

        temp_mos = self._predict_mos(dl)

        if not (temp_mos >= self._min_mos and temp_mos <= self._max_mos):
            raise Exception("Invalid MOS %.5f (bounds: [%.5f, %.5f]!" % (temp_mos, self._min_mos, self._max_mos))

        if self._scaled:
            utility = np.interp(temp_mos, (self._min_mos, self._max_mos), (1, 5))
        else:
            utility = temp_mos

        return utility

    def _predict_mos(self, dl):
        """
        :param dl: Download Time (s)
        """
        # Restrict to range 1.0 - 5.0
        return max(1.0, min(-1.68 * np.log(dl) + 9.61, 5.0))

    def inverse(self, mos):

        assert(mos >= 1.0 and mos <= 5.0)

        r = self._finv(mos)

        return r

if __name__ == "__main__":

    model = WebDLUtility()

    print(model.predict(10))
    print(model.predict(150))

    import numpy as np
    import matplotlib.pylab as plt

    f, ax = plt.subplots()

    x = np.linspace(0.01, 180.0)
    y = list(map(model.predict, x))

    ax.step(x, y)
    ax.grid()

    ax.set_yticks(range(1, 6))

    ax.set_xlabel("Download Time 10MB dl (s)")
    ax.set_ylabel(r"$Utility_{WEBDL}(dl)$")

    # Inverse function
    f, ax = plt.subplots()

    x = np.linspace(1, 5)
    y = list(map(model.inverse, x))

    ax.step(x, y)
    ax.grid()

    ax.set_ylabel("Download Time 10MB dl (s)")
    ax.set_xlabel(r"$Utility_{WEBDL}(dl)$")

    f.savefig("utility_webdl.png", dpi=200)


<filename>cobra/evaluation/evaluator.py

import numpy as np
import pandas as pd

import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import seaborn as sns

from numpy import sqrt
from scipy.stats import norm

# classification
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score
from sklearn.metrics import f1_score
from sklearn.metrics import accuracy_score
from sklearn.metrics import roc_curve
from sklearn.metrics import confusion_matrix
from sklearn.metrics import roc_auc_score
from sklearn.metrics import matthews_corrcoef
from sklearn.exceptions import NotFittedError

# regression
from sklearn.metrics import mean_absolute_error
from sklearn.metrics import mean_squared_error
from sklearn.metrics import r2_score

class ClassificationEvaluator():
    """Evaluator class encapsulating classification model metrics
    and plotting functionality.

    Attributes
    ----------
    y_true : np.ndarray
        True binary target data labels.
    y_pred : np.ndarray
        Target scores of the model.
    confusion_matrix : np.ndarray
        Confusion matrix computed for a particular cut-off.
    cumulative_gains : tuple
        Data for plotting cumulative gains curve.
    evaluation_metrics : dict
        Map containing various scalar evaluation metrics (precision, recall, accuracy, AUC, F1, etc.).
    lift_at : float
        Parameter to determine at which top level percentage the lift of the
        model should be computed.
    lift_curve : tuple
        Data for plotting lift curve(s).
    probability_cutoff : float
        Probability cut off to convert probability scores to a binary score.
    roc_curve : dict
        Map containing true-positive-rate, false-positive-rate at various
        thresholds (also incl.).
    n_bins : int, optional
        Defines the number of bins used to calculate the lift curve for
        (by default 10, so deciles).
    """

    def __init__(self,
                 probability_cutoff: float=None,
                 lift_at: float=0.05,
                 n_bins: int = 10):

        self.y_true = None
        self.y_pred = None

        self.lift_at = lift_at
        self.probability_cutoff = probability_cutoff
        self.n_bins = n_bins

        # Placeholder to store fitted output
        self.scalar_metrics = None
        self.roc_curve = None
        self.confusion_matrix = None
        self.lift_curve = None
        self.cumulative_gains = None

    def fit(self, y_true: np.ndarray, y_pred: np.ndarray):
        """Fit the evaluator by computing the relevant evaluation metrics on
        the inputs.

        Parameters
        ----------
        y_true : np.ndarray
            True labels.
        y_pred : np.ndarray
            Model scores (as probability).
        """
        fpr, tpr, thresholds = roc_curve(y_true=y_true, y_score=y_pred)

        # if probability_cutoff is not set, take the optimal cut-off
        if not self.probability_cutoff:
            self.probability_cutoff = (ClassificationEvaluator.
                                       _compute_optimal_cutoff(fpr, tpr,
                                                               thresholds))

        # Transform probabilities to binary array using cut-off
        y_pred_b = np.array([0 if pred <= self.probability_cutoff else 1
                             for pred in y_pred])

        # Compute the various evaluation metrics
        self.scalar_metrics = ClassificationEvaluator._compute_scalar_metrics(
            y_true,
            y_pred,
            y_pred_b,
            self.lift_at
        )

        self.y_true = y_true
        self.y_pred = y_pred

        self.roc_curve = {"fpr": fpr, "tpr": tpr, "thresholds": thresholds}
        self.confusion_matrix = confusion_matrix(y_true, y_pred_b)
        self.lift_curve = ClassificationEvaluator._compute_lift_per_bin(y_true, y_pred, self.n_bins)
        self.cumulative_gains = ClassificationEvaluator._compute_cumulative_gains(y_true, y_pred)

    @staticmethod
    def _compute_scalar_metrics(y_true: np.ndarray,
                                y_pred: np.ndarray,
                                y_pred_b: np.ndarray,
                                lift_at: float) -> pd.Series:
        """Convenient function to compute various scalar performance measures
        and return them in a pd.Series.

        Parameters
        ----------
        y_true : np.ndarray
            True binary target data labels.
        y_pred : np.ndarray
            Target scores of the model.
        y_pred_b : np.ndarray
            Predicted target data labels (binary).
        lift_at : float
            At what top level percentage the lift should be computed.

        Returns
        -------
        pd.Series
            Contains various performance measures of the model, being:
                Accuracy
                AUC
                Precision
                Recall
                F1
                Matthews correlation coefficient
                Lift at given percentage
        """
        return pd.Series({
            "accuracy": accuracy_score(y_true, y_pred_b),
            "AUC": roc_auc_score(y_true, y_pred),
            "precision": precision_score(y_true, y_pred_b),
            "recall": recall_score(y_true, y_pred_b),
            "F1": f1_score(y_true, y_pred_b, average=None)[1],
            "matthews_corrcoef": matthews_corrcoef(y_true, y_pred_b),
            "lift at {}".format(lift_at): np.round(ClassificationEvaluator
                                                   ._compute_lift(y_true=y_true,
                                                                  y_pred=y_pred,
                                                                  lift_at=lift_at), 2)
        })

    def plot_roc_curve(self, path: str=None, dim: tuple=(12, 8)):
        """Plot ROC curve of the model.

        Parameters
        ----------
        path : str, optional
            Path to store the figure.
        dim : tuple, optional
            Tuple with width and length of the plot.
        """

        if self.roc_curve is None:
            msg = ("This {} instance is not fitted yet. Call 'fit' with "
                   "appropriate arguments before using this method.")

            raise NotFittedError(msg.format(self.__class__.__name__))

        auc = float(self.scalar_metrics.loc["AUC"])

        with plt.style.context("seaborn-whitegrid"):

            fig, ax = plt.subplots(figsize=dim)

            ax.plot(self.roc_curve["fpr"],
                    self.roc_curve["tpr"],
                    color="cornflowerblue", linewidth=3,
                    label="ROC curve (area = {s:.3})".format(s=auc))

            ax.plot([0, 1], [0, 1], color="darkorange", linewidth=3,
                    linestyle="--")
            ax.set_xlabel("False Positive Rate", fontsize=15)
            ax.set_ylabel("True Positive Rate", fontsize=15)
            ax.legend(loc="lower right")
            ax.set_title("ROC curve", fontsize=20)

            if path:
                plt.savefig(path, format="png", dpi=300, bbox_inches="tight")

        plt.show()

    def plot_confusion_matrix(self, path: str=None, dim: tuple=(12, 8),
                              labels: list=["0", "1"]):
        """Plot the confusion matrix.

        Parameters
        ----------
        path : str, optional
            Path to store the figure.
        dim : tuple, optional
            Tuple with width and length of the plot.
        labels : list, optional
            Optional list of labels, default "0" and "1".
        """

        if self.confusion_matrix is None:
            msg = ("This {} instance is not fitted yet. Call 'fit' with "
                   "appropriate arguments before using this method.")

            raise NotFittedError(msg.format(self.__class__.__name__))

        fig, ax = plt.subplots(figsize=dim)
        ax = sns.heatmap(self.confusion_matrix,
                         annot=self.confusion_matrix.astype(str),
                         fmt="s", cmap="Blues",
                         xticklabels=labels, yticklabels=labels)
        ax.set_title("Confusion matrix", fontsize=20)

        if path:
            plt.savefig(path, format="png", dpi=300, bbox_inches="tight")

        plt.show()

    def plot_cumulative_response_curve(self, path: str=None, dim: tuple=(12, 8)):
        """Plot cumulative response curve.

        Parameters
        ----------
        path : str, optional
            Path to store the figure.
        dim : tuple, optional
            Tuple with width and length of the plot.
        """

        if self.lift_curve is None:
            msg = ("This {} instance is not fitted yet. Call 'fit' with "
                   "appropriate arguments before using this method.")

            raise NotFittedError(msg.format(self.__class__.__name__))

        x_labels, lifts, inc_rate = self.lift_curve

        lifts = np.array(lifts)*inc_rate*100

        with plt.style.context("seaborn-ticks"):
            fig, ax = plt.subplots(figsize=dim)

            plt.bar(x_labels[::-1], lifts, align="center",
                    color="cornflowerblue")
            plt.ylabel("response (%)", fontsize=16)
            plt.xlabel("decile", fontsize=16)
            ax.set_xticks(x_labels)
            ax.set_xticklabels(x_labels)

            plt.axhline(y=inc_rate*100, color="darkorange", linestyle="--",
                        xmin=0.05, xmax=0.95, linewidth=3, label="Incidence")

            # Legend
            ax.legend(loc="upper right")

            # Set Axis - make them pretty
            sns.despine(ax=ax, right=True, left=True)

            # Remove white lines from the second axis
            ax.grid(False)

            # Description
            ax.set_title("Cumulative Response curve", fontsize=20)

            if path is not None:
                plt.savefig(path, format="png", dpi=300, bbox_inches="tight")

            plt.show()

    def plot_lift_curve(self, path: str=None, dim: tuple=(12, 8)):
        """Plot lift per decile.

        Parameters
        ----------
        path : str, optional
            Path to store the figure.
        dim : tuple, optional
            Tuple with width and length of the plot.
        """

        if self.lift_curve is None:
            msg = ("This {} instance is not fitted yet. Call 'fit' with "
                   "appropriate arguments before using this method.")

            raise NotFittedError(msg.format(self.__class__.__name__))

        x_labels, lifts, _ = self.lift_curve

        with plt.style.context("seaborn-ticks"):
            fig, ax = plt.subplots(figsize=dim)

            plt.bar(x_labels[::-1], lifts, align="center",
                    color="cornflowerblue")
            plt.ylabel("lift", fontsize=16)
            plt.xlabel("decile", fontsize=16)
            ax.set_xticks(x_labels)
            ax.set_xticklabels(x_labels)

            plt.axhline(y=1, color="darkorange", linestyle="--",
                        xmin=0.05, xmax=0.95, linewidth=3, label="Baseline")

            # Legend
            ax.legend(loc="upper right")

            # Set Axis - make them pretty
            sns.despine(ax=ax, right=True, left=True)

            # Remove white lines from the second axis
            ax.grid(False)

            # Description
            ax.set_title("Cumulative Lift curve", fontsize=20)

            if path is not None:
                plt.savefig(path, format="png", dpi=300, bbox_inches="tight")

            plt.show()

    def plot_cumulative_gains(self, path: str=None, dim: tuple=(12, 8)):
        """Plot cumulative gains per decile.

        Parameters
        ----------
        path : str, optional
            Path to store the figure.
        dim : tuple, optional
            Tuple with width and length of the plot.
        """

        with plt.style.context("seaborn-whitegrid"):
            fig, ax = plt.subplots(figsize=dim)

            ax.plot(self.cumulative_gains[0]*100, self.cumulative_gains[1]*100,
                    color="cornflowerblue", linewidth=3,
                    label="cumulative gains")
            ax.plot(ax.get_xlim(), ax.get_ylim(), linewidth=3,
                    ls="--", color="darkorange", label="random selection")

            ax.set_title("Cumulative Gains curve", fontsize=20)

            # Format axes
            ax.set_xlim([0, 100])
            ax.set_ylim([0, 105])

            # Format ticks
            ticks_loc_y = ax.get_yticks().tolist()
            ax.yaxis.set_major_locator(mticker.FixedLocator(ticks_loc_y))
            ax.set_yticklabels(["{:3.0f}%".format(x) for x in ticks_loc_y])

            ticks_loc_x = ax.get_xticks().tolist()
            ax.xaxis.set_major_locator(mticker.FixedLocator(ticks_loc_x))
            ax.set_xticklabels(["{:3.0f}%".format(x) for x in ticks_loc_x])

            # Legend
            ax.legend(loc="lower right")

            if path is not None:
                plt.savefig(path, format="png", dpi=300, bbox_inches="tight")
            plt.show()

    @staticmethod
    def _find_optimal_cutoff(y_true: np.ndarray,
                             y_pred: np.ndarray) -> float:
        """Find the optimal probability cut off point for a
        classification model. Wrapper around _compute_optimal_cutoff.

        Parameters
        ----------
        y_true : np.ndarray
            True binary target data labels.
        y_pred : np.ndarray
            Target scores of the model.

        Returns
        -------
        float
            Optimal cut-off probability for the model.
        """
        return ClassificationEvaluator._compute_optimal_cutoff(roc_curve(y_true=y_true,
                                                                         y_score=y_pred))

    @staticmethod
    def _compute_optimal_cutoff(fpr: np.ndarray, tpr: np.ndarray,
                                thresholds: np.ndarray) -> float:
        """Find the optimal probability cut-off point for a
        classification model.

        The optimal cut-off would be where TPR is high and FPR is low, hence
        TPR - (1-FPR) should be zero or close to zero for the optimal cut-off.

        Parameters
        ----------
        fpr : np.ndarray
            False positive rate for various thresholds.
        tpr : np.ndarray
            True positive rate for various thresholds.
        thresholds : np.ndarray
            List of thresholds for which fpr and tpr were computed.

        Returns
        -------
        float
            Optimal probability cut-off point.
        """
        temp = np.absolute(tpr - (1-fpr))

        # index for optimal value is the one for which temp is minimal
        optimal_index = np.where(temp == min(temp))[0]

        return thresholds[optimal_index][0]

    @staticmethod
    def _compute_cumulative_gains(y_true: np.ndarray,
                                  y_pred: np.ndarray) -> tuple:
        """Compute cumulative gains of the model, returns percentages and
        gains cumulative gains curves.

        Code from (https://github.com/reiinakano/scikit-plot/blob/
                   2dd3e6a76df77edcbd724c4db25575f70abb57cb/
                   scikitplot/helpers.py#L157)

        Parameters
        ----------
        y_true : np.ndarray
            True binary target data labels.
        y_pred : np.ndarray
            Target scores of the model.

        Returns
        -------
        tuple
            With x-labels, and gains.
        """

        # make y_true a boolean vector
        y_true = (y_true == 1)

        sorted_indices = np.argsort(y_pred)[::-1]
        y_true = y_true[sorted_indices]
        gains = np.cumsum(y_true)

        percentages = np.arange(start=1, stop=len(y_true) + 1)

        gains = gains / float(np.sum(y_true))
        percentages = percentages / float(len(y_true))

        gains = np.insert(gains, 0, [0])
        percentages = np.insert(percentages, 0, [0])

        return percentages, gains

    @staticmethod
    def _compute_lift_per_bin(y_true: np.ndarray,
                              y_pred: np.ndarray,
                              n_bins: int=10) -> tuple:
        """Compute lift of the model for a given number of bins, returns x-labels,
        lifts and the target incidence to create cumulative response curves.

        Parameters
        ----------
        y_true : np.ndarray
            True binary target data labels.
        y_pred : np.ndarray
            Target scores of the model.
        n_bins : int, optional
            Defines the number of bins used to calculate the lift curve for
            (by default 10, so deciles).

        Returns
        -------
        tuple
            Includes x-labels, lifts per decile, and target incidence.
        """

        lifts = [ClassificationEvaluator._compute_lift(y_true=y_true,
                                                       y_pred=y_pred,
                                                       lift_at=perc_lift)
                 for perc_lift in np.linspace(1/n_bins, 1, num=n_bins, endpoint=True)]

        x_labels = [len(lifts)-x for x in np.arange(0, len(lifts), 1)]

        return x_labels, lifts, y_true.mean()

    @staticmethod
    def _compute_lift(y_true: np.ndarray, y_pred: np.ndarray,
                      lift_at: float=0.05) -> float:
        """Calculates lift given two arrays on specified level.

        Parameters
        ----------
        y_true : np.ndarray
            True binary target data labels.
        y_pred : np.ndarray
            Target scores of the model.
        lift_at : float, optional
            At what top level percentage the lift should be computed.

        Returns
        -------
        float
            Lift of the model.
        """

        # Make sure it is numpy array
        y_true_ = np.array(y_true)
        y_pred_ = np.array(y_pred)

        # Make sure it has correct shape
        y_true_ = y_true_.reshape(len(y_true_), 1)
        y_pred_ = y_pred_.reshape(len(y_pred_), 1)

        # Merge data together
        y_data = np.hstack([y_true_, y_pred_])

        # Calculate necessary variables
        nrows = len(y_data)
        stop = int(np.floor(nrows*lift_at))
        avg_incidence = np.einsum("ij->j", y_true_)/float(len(y_true_))

        # Sort and filter data
        data_sorted = (y_data[y_data[:, 1].argsort()[::-1]][:stop, 0]
                       .reshape(stop, 1))

        # Calculate lift (einsum is a very fast way of summing, but needs specific shape)
        inc_in_top_n = np.einsum("ij->j", data_sorted)/float(len(data_sorted))

        lift = np.round(inc_in_top_n/avg_incidence, 2)[0]

        return lift

class RegressionEvaluator():
    """Evaluator class encapsulating regression model metrics
    and plotting functionality.

    Attributes
    ----------
    y_true : np.ndarray
        True binary target data labels.
    y_pred : np.ndarray
        Target scores of the model.
    scalar_metrics : dict
        Map containing various scalar evaluation metrics (R-squared, MAE, MSE, RMSE)
    qq : pd.Series
        Theoretical quantiles and associated actual residuals.
    """

    def __init__(self):

        self.y_true = None
        self.y_pred = None

        # Placeholder to store fitted output
        self.scalar_metrics = None
        self.qq = None

    def fit(self, y_true: np.ndarray, y_pred: np.ndarray):
        """Fit the evaluator by computing the relevant evaluation metrics on
        the inputs.

        Parameters
        ----------
        y_true : np.ndarray
            True labels.
        y_pred : np.ndarray
            Model scores.
        """
        # Compute the various evaluation metrics
        self.scalar_metrics = RegressionEvaluator._compute_scalar_metrics(y_true, y_pred)

        self.y_true = y_true
        self.y_pred = y_pred

        # Compute qq info
        self.qq = RegressionEvaluator._compute_qq_residuals(y_true, y_pred)

    @staticmethod
    def _compute_scalar_metrics(y_true: np.ndarray,
                                y_pred: np.ndarray) -> pd.Series:
        """Convenient function to compute various scalar performance measures
        and return them in a pd.Series.

        Parameters
        ----------
        y_true : np.ndarray
            True binary target data labels.
        y_pred : np.ndarray
            Target scores of the model.

        Returns
        -------
        pd.Series
            Contains various performance measures of the model, being:
                R-squared (coefficient of determination, usually denoted as R-squared)
                Mean absolute error (expected value of the absolute error loss)
                Mean squared error (expected value of the quadratic error)
                Root mean squared error (sqrt of expected value of the quadratic error)
        """
        return pd.Series({
            "R2": r2_score(y_true, y_pred),
            "MAE": mean_absolute_error(y_true, y_pred),
            "MSE": mean_squared_error(y_true, y_pred),
            "RMSE": sqrt(mean_squared_error(y_true, y_pred))
        })

    @staticmethod
    def _compute_qq_residuals(y_true: np.ndarray,
                              y_pred: np.ndarray) -> pd.Series:
        """Convenience function to compute various scalar performance measures
        and return them in a pd.Series.

        Parameters
        ----------
        y_true : np.ndarray
            True binary target data labels.
        y_pred : np.ndarray
            Target scores of the model.

        Returns
        -------
        pd.Series
            Theoretical quantiles and associated actual residuals.
        """
        ## also possible directly via statsmodels.api.qqplot()

        n = len(y_true)

        df = pd.DataFrame({"res": sorted((y_true - y_pred))})  # ascending order
        m, s = df["res"].mean(), df["res"].std()

        df["z_res"] = df["res"].apply(lambda x: (x-m)/s)
        df["rank"] = df.index+1
        df["percentile"] = df["rank"].apply(lambda x: x/(n+1))  # divide by n+1 to avoid inf
        df["q_theoretical"] = norm.ppf(df["percentile"])

        return pd.Series({
            "quantiles": df["q_theoretical"].values,
            "residuals": df["z_res"].values,
        })

    def plot_predictions(self, path: str=None, dim: tuple=(12, 8)):
        """Plot predictions from the model against actual values.

        Parameters
        ----------
        path : str, optional
            Path to store the figure.
        dim : tuple, optional
            Tuple with width and length of the plot.
        """
        if self.y_true is None and self.y_pred is None:
            msg = ("This {} instance is not fitted yet. Call 'fit' with "
                   "appropriate arguments before using this method.")

        y_true = self.y_true
        y_pred = self.y_pred

        with plt.style.context("seaborn-whitegrid"):

            fig, ax = plt.subplots(figsize=dim)

            x = np.arange(1, len(y_true)+1)

            ax.plot(x, y_true, ls="--", label="actuals", color="darkorange", linewidth=3)
            ax.plot(x, y_pred, label="predictions", color="cornflowerblue", linewidth=3)

            ax.set_xlabel("Index", fontsize=15)
            ax.set_ylabel("Value", fontsize=15)
            ax.legend(loc="best")
            ax.set_title("Predictions vs. Actuals", fontsize=20)

            if path:
                plt.savefig(path, format="png", dpi=300, bbox_inches="tight")

        plt.show()

    def plot_qq(self, path: str=None, dim: tuple=(12, 8)):
        """Display a Q-Q plot from the standardized prediction residuals.

        Parameters
        ----------
        path : str, optional
            Path to store the figure.
        dim : tuple, optional
            Tuple with width and length of the plot.
        """

        if self.qq is None:
            msg = ("This {} instance is not fitted yet. Call 'fit' with "
                   "appropriate arguments before using this method.")

            raise NotFittedError(msg.format(self.__class__.__name__))

        with plt.style.context("seaborn-whitegrid"):

            fig, ax = plt.subplots(figsize=dim)

            x = self.qq["quantiles"]
            y = self.qq["residuals"]

            ax.plot(x, x, ls="--", label="perfect model", color="darkorange", linewidth=3)
            ax.plot(x, y, label="current model", color="cornflowerblue", linewidth=3)

            ax.set_xlabel("Theoretical quantiles", fontsize=15)
            ax.set_xticks(range(int(np.floor(min(x))), int(np.ceil(max(x[x < float("inf")])))+1, 1))

            ax.set_ylabel("Standardized residuals", fontsize=15)
            ax.set_yticks(range(int(np.floor(min(y))), int(np.ceil(max(y[x < float("inf")])))+1, 1))

            ax.legend(loc="best")
            ax.set_title("Q-Q plot", fontsize=20)

            if path:
                plt.savefig(path, format="png", dpi=300, bbox_inches="tight")

        plt.show()


<filename>test/test_biplot.py<gh_stars>1-10
import unittest
import numpy as np
import pandas as pd
from biofes.biplot import *
from biofes import biplot
from scipy import stats

class test_functions(unittest.TestCase):
    def test_standardize(self):
        A = np.random.uniform(-300,300,size=(300,30))
        A_st = standardize(A, meth=1)
        A_ref = (A-A.mean(axis = 0))/A.std(axis = 0)
        self.assertAlmostEqual(np.mean(A_ref - A_st), 0, msg='standardization error')

    def test_Factor2Binary(self):
        target = list(np.random.randint(np.random.randint(2, 10), size = 100))
        Z = Factor2Binary(target,Name = None)

        Z_ref = pd.get_dummies(target)
        self.assertAlmostEqual(np.mean(Z_ref.values - Z.values), 0, msg='Factor2Binary error')

    def test_matrixsqrt(self):
        A = np.random.randint(low = 0, high = 200, size=(300, 30))
        d = np.random.randint(30)
        tol = np.finfo(float).eps

        Sinv = matrixsqrt(A, d, tol, inv=True)
        U, Sigma, VT = SVD(A, d, niter=5, state=0)
        nz = Sigma > tol
        Sinv_ref = U.dot(np.diag(1/np.sqrt(Sigma[nz]))).dot(VT[nz,:])
        self.assertAlmostEqual(np.mean(Sinv_ref - Sinv), 0, delta=1e-3, msg='matrixsqrt (inv=True) error')

        ###############################################################################

        A = np.random.randint(low = 0, high = 200, size=(300, 30))
        d = np.random.randint(30)
        tol = np.finfo(float).eps

        S = matrixsqrt(A, d, tol, inv=False)
        U, Sigma, VT = SVD(A, d, niter=5, state=0)
        nz = Sigma > tol
        S_ref = U.dot(np.diag(np.sqrt(Sigma[nz]))).dot(VT[nz,:])
        self.assertAlmostEqual(np.mean(S_ref - S), 0, delta=1e-3, msg='matrixsqrt (inv=False) error')

class test_biplot(unittest.TestCase):
    def test_Classic(self):
        n, p = np.random.randint(70,500), np.random.randint(30,50)

        A = np.random.uniform(-300,300,size=(n,p))
        d = np.random.randint(p)
        a = np.random.random(1)[0]
        methods = [None, 1]
        m = methods[np.random.randint(2)]

        data_st = standardize(A, m)
        U, Sigma, VT = SVD(data_st, d, niter = 35, state = 0)

        EV_ref = np.power(Sigma,2)
        Inert_ref = EV_ref/np.sum(EV_ref) * 100

        # Contributions

        R = U.dot(np.diag(Sigma[:d]))
        C = np.transpose(VT).dot(np.diag(Sigma[:d]))

        sf = np.sum(np.power(A,2),axis=1)
        cf = np.zeros((n,d))
        for k in range(0,d):
            cf[:,k] = np.power(R[:,k],2)*100/sf

        sc = np.sum(np.power(A, 2),axis=0)
        cc = np.zeros((p,d))
        for k in range(0,d):
            cc[:,k] = np.power(C[:,k],2)*100/sc

        # Coordinates

        R = R.dot(np.diag(np.power(Sigma,a)))
        C = C.dot(np.diag(np.power(Sigma,1-a)))

        sca = np.sum(np.power(R,2))/n
        scb = np.sum(np.power(C,2))/p
        scf = np.sqrt(np.sqrt(scb/sca))

        RowCoord_ref = R*scf
        ColCoord_ref = C/scf

        # biplot from biofes

        BCla = biplot.Classic(data = A ,dim = d, alpha = a, method = m, niter = 35, state = 0)

        # DIMENSION TEST
        self.assertEqual(BCla.RowCoord.shape, (n, d), msg='dimension output error (Classic Biplot)')
        self.assertEqual(BCla.ColCoord.shape, (p, d) , msg='dimension output error (Classic Biplot)')
        self.assertEqual(len(BCla.Inert), d, msg='dimension output error (Classic Biplot)')
        self.assertEqual(len(BCla.EV)   , d, msg='dimension output error (Classic Biplot)')

        # INERTIA / EV TEST
        try:
            if str((EV_ref - EV).mean()) == 'nan':
                pass
            else:
                self.assertAlmostEqual(np.mean(EV_ref - BCla.EV), 0, msg='EV error')
                self.assertAlmostEqual(np.mean(Inert_ref - BCla.Inert), 0, msg='Inertia error')
        except:
            pass

        # CONTRIBUTIONS TEST
        try:
            if str((cf - BCla.RowCont).mean()) == 'nan':
                pass
            else:
                els = cf.shape[0]*cf.shape[1]
                self.assertAlmostEqual(np.mean(cf - BCla.RowCont), 0, delta=els*(1e-2), msg='Row Contributions error')
                els = cc.shape[0]*cc.shape[1]
                self.assertAlmostEqual(np.mean(cc - BCla.ColCont), 0, delta=els*(1e-2), msg='Column Contributions error')
        except:
            pass

        # COORDINATES TEST
        self.assertAlmostEqual(np.mean(RowCoord_ref - BCla.RowCoord), 0, delta=1e-3, msg='Row Coordinates error')
        self.assertAlmostEqual(np.mean(ColCoord_ref - BCla.ColCoord), 0, delta=1e-3, msg='Col Coordinates error')

    def test_Canonical(self):
        n, m = np.random.randint(70,500), np.random.randint(10,50)
        A = np.random.uniform(-300,300,size=(n,m))
        target = list(np.random.randint(np.random.randint(2, 10), size = A.shape[0]))
        gn = list(set(target))
        g = len(gn)
        d = np.random.randint(len(gn)+1, m)
        methods = [None, 1]
        met = methods[np.random.randint(2)]

        data_std = standardize(A, met)
        r = np.array([len(gn) - 1, m]).min()
        #Groups to Binary
        Z = Factor2Binary(target)
        ng = Z.sum(axis=0)
        S11 = (Z.T).dot(Z).values

        Xb = np.linalg.inv(S11).dot(Z.T).dot(data_std)
        B = (Xb.T).dot(S11).dot(Xb)
        S = (data_std.T).dot(data_std) - B
        Y = np.power(S11,0.5).dot(Xb).dot(matrixsqrt(S,d,inv=True))

        U, Sigma, VT = SVD(Y, d, niter = 15, state = 0)

        #Variable_Coord
        H = matrixsqrt(S, d, inv=False).dot(np.transpose(VT[0:r,:]))
        #Canonical_Weights
        B = matrixsqrt(S, d, inv=True ).dot(np.transpose(VT[0:r,:]))
        #Group_Coord
        J = Xb.dot(B)
        #Individual_Coord
        V = data_std.dot(B)

        sct = np.diag((V.T).dot(V))
        sce = np.diag((J.T).dot(S11).dot(J))
        scr = sct -sce
        fs = (sce/(g - 1))/(scr/(n - g))

        #eigenvectors
        vprop = Sigma[:r]
        #Inertia
        iner = (np.power(vprop,2)/(np.power(vprop,2).sum()))*100

        lamb = np.power(vprop,2)
        pill = 1/(1 + lamb)
        pillai = np.linalg.det(np.diag(pill))
        glh = g - 1
        gle = n - g
        t = np.sqrt((np.power(glh,2) * np.power(m,2) - 4)/(np.power(m,2) + np.power(glh,2) - 5))
        w = gle + glh - 0.5 * (m + glh + 1)
        df1 = m * glh
        df2 = w * t - 0.5 * (m * glh - 2)

        # Wilks
        Wilksf = (1 - np.power(pillai,1/t))/(np.power(pillai,1/t)) * (df2/df1)
        Wilksp = stats.f.pdf(Wilksf, df1, df2)
        Wilks = {'f-val': Wilksf,'p-val': Wilksp}

        # Radius

        falfau = stats.t.ppf(1 - (0.025), (n - g))
        falfab = stats.t.ppf(1 - (0.025/(g * m)), (n - g))
        falfam = np.sqrt(stats.f.ppf(1 - 0.05, m, (n - g - m + 1)) * (((n - g) * m)/(n - g - m + 1)))
        falfac = 2.447747

        UnivRad = falfau * np.diag(np.linalg.inv(np.sqrt(S11)))/np.sqrt(n - g)
        BonfRad = falfab * np.diag(np.linalg.inv(np.sqrt(S11)))/np.sqrt(n - g)
        MultRad = falfam * np.diag(np.linalg.inv(np.sqrt(S11)))/np.sqrt(n - g)
        ChisRad = falfac * np.diag(np.linalg.inv(np.sqrt(S11)))/np.sqrt(n - g)

        Radius = {'Uni': UnivRad,'Bonf': BonfRad, 'Mult': MultRad, 'Chis': ChisRad}

        BCan = biplot.Canonical(data = A, dim = d, GroupNames = gn, y = target, method = met, niter = 35, state = 0)

        # DIMENSION TEST
        self.assertEqual(BCan.Ind_Coord.shape, (n, len(gn)-1), msg='dimension output error (Canonical Biplot) Ind_Coord')
        self.assertEqual(BCan.Var_Coord.shape, (m, len(gn)-1), msg='dimension output error (Canonical Biplot) Var_Coord')
        self.assertEqual(BCan.Group_Coord.shape, (len(gn), len(gn)-1), msg='dimension output error (Canonical Biplot) Group_Coord')
        self.assertEqual(len(BCan.inert), len(gn)-1, msg='dimension output error (Canonical Biplot)')

        # COORDINATES TEST
        els = H.shape[0]*H.shape[1]
        self.assertAlmostEqual(np.mean(H - BCan.Var_Coord), 0, delta=els*(1e-2), msg='Var Coordinates error')
        els = V.shape[0]*V.shape[1]
        self.assertAlmostEqual(np.mean(V - BCan.Ind_Coord), 0, delta=els*(1e-2), msg='Ind Coordinates error')
        els = J.shape[0]*J.shape[1]
        self.assertAlmostEqual(np.mean(J - BCan.Group_Coord), 0, delta=els*(1e-2), msg='Group Coordinates error')

        # CANONICAL WEIGHTS TEST
        els = B.shape[0]*B.shape[1]
        self.assertAlmostEqual(np.mean(B - BCan.Can_Weights), 0, delta=els*(1e-2), msg='Canonical Weights error')

        # EV / INERTIA TEST
        try:
            if str((vprop - BCan.vprop).mean()) == 'nan':
                pass
            else:
                self.assertAlmostEqual(np.mean(vprop - BCan.vprop), 0, msg='EV error')
                self.assertAlmostEqual(np.mean(iner -  BCan.inert), 0, msg='Inertia error')
        except:
            pass

        # WILKS TEST
        self.assertAlmostEqual(Wilks['f-val'] - BCan.Wilks['f-val'], 0, delta=(1e-2), msg='f-val Wilks error')
        self.assertAlmostEqual(Wilks['p-val'] - BCan.Wilks['p-val'], 0, delta=(1e-2), msg='p-val Wilks error')

        # RADIUS
        self.assertAlmostEqual(np.mean(Radius['Uni'] - BCan.Radius['Uni']), 0, delta=(1e-3), msg='Uni Radius error')
        self.assertAlmostEqual(np.mean(Radius['Bonf'] - BCan.Radius['Bonf']), 0, delta=(1e-3), msg='Bonferroni Radius error')
        self.assertAlmostEqual(np.mean(Radius['Mult'] - BCan.Radius['Mult']), 0, delta=(1e-3), msg='Mult Radius error')
        self.assertAlmostEqual(np.mean(Radius['Chis'] - BCan.Radius['Chis']), 0, delta=(1e-3), msg='Chi-sqr Radius error')

    def test_CA(self):
        n, p = np.random.randint(100, 500), np.random.randint(50, 90)
        A = np.random.randint(np.random.randint(1,10), np.random.randint(30,200), (n,p))
        dim = np.random.randint(p)
        a = np.random.random(1)[0]

        BCA = biplot.CA(data=A, dim=dim, alpha=a ,niter=15, state=0)

        A = A / A.sum()

        dr = np.matrix(A.sum(axis=1))
        dc = np.matrix(A.sum(axis=0))

        A = A - dr.T.dot(dc)

        Dr = np.diagflat(1/np.sqrt(dr))
        Dc = np.diagflat(1/np.sqrt(dc))
        A = Dr.dot(A).dot(Dc)

        U, Sigma, VT = SVD(A, dim, niter = 15, state = 0)

        d = Sigma[:np.min(A.shape)]
        r = np.min(A.shape)

        inertia = numpy.power(d,2)*100 / numpy.sum(numpy.power(d,2))

        U = Dr.dot(U[:,:r])
        V = Dc.dot(VT.T[:,:r])

        D = np.diagflat(d)
        A = U.dot(D)
        B = V.dot(D)

        sf = np.power(A,2).sum(axis = 1)
        cf = np.linalg.inv(np.diagflat(sf)).dot(np.power(A,2))

        sc = np.power(B,2).sum(axis = 1)
        cc = np.linalg.inv(np.diagflat(sc)).dot(np.power(B,2))

        A = U.dot(np.diagflat(np.power(d,a)))
        B = V.dot(np.diagflat(np.power(d,1-a)))

        # AB
        AB = A[:, :dim].dot(B[:, :dim].T)

        # Eigenvalues
        eigen_values = np.power(d,2)

        # Coordinates
        RowCoordinates = A[:, :dim]
        ColCoordinates = B[:, :dim]

        # Contributions
        RowContributions = cf[:, :dim] * 100
        ColContributions = cc[:, :dim] * 100

        # INERTIA TEST
        try:
            if str((inertia -  BCA.inertia).mean()) == 'nan':
                pass
            else:
                self.assertAlmostEqual(np.mean(inertia -  BCA.inertia), 0, msg='Inertia error')
                self.assertAlmostEqual(np.mean(eigen_values -  BCA.eigen_values), 0, msg='eigen values error')
        except:
            pass

        # AB TEST
        self.assertAlmostEqual(np.mean(AB - BCA.AB), 0, delta=(1e-2), msg='AB error')

        # COORDINATES TEST
        els = RowCoordinates.shape[0]*RowCoordinates.shape[1]
        self.assertAlmostEqual(np.mean(RowCoordinates - BCA.RowCoordinates), 0, delta=els*(1e-2), msg='Row Coordinates error')
        els = ColCoordinates.shape[0]*ColCoordinates.shape[1]
        self.assertAlmostEqual(np.mean(ColCoordinates - BCA.ColCoordinates), 0, delta=els*(1e-2), msg='Col Coordinates error')

        # CONTRIBUTIONS TEST
        try:
            if str((cf - BCla.RowCont).mean()) == 'nan':
                pass
            else:
                els = RowContributions.shape[0]*RowContributions.shape[1]
                self.assertAlmostEqual(np.mean(RowContributions - BCA.RowContributions), 0, delta=els*(1e-2), msg='Row Contributions error')
                els = ColContributions.shape[0]*ColContributions.shape[1]
                self.assertAlmostEqual(np.mean(ColContributions - BCA.ColContributions), 0, delta=els*(1e-2), msg='Column Contributions error')
        except:
            pass

if __name__ == '__main__':
    unittest.main()