import glob
import json
import os
import pickle
from functools import wraps
import astropy.units as u
import numpy as np
[docs]
__author__ = ["Sijie Yu"]
[docs]
__email__ = "sijie.yu@njit.edu"
[docs]
def img2html_movie(imgprefix, outname='movie', img_fmt='png'):
imgfiles = glob.glob(imgprefix + '*.' + img_fmt)
imgfiles = sorted(imgfiles)
# try:
import matplotlib.image as mpimg
img = mpimg.imread(imgfiles[0])
height, width, dummy = img.shape
# except:
# from PIL import Image
# img = Image.open(imgfiles[0])
# width, height = img.size
# img.close()
nfiles = len(imgfiles)
htmlfile = os.path.join(os.path.dirname(imgprefix), '{}.html'.format(outname))
fi = open(htmlfile, 'w')
fi.write('<HTML>\n')
fi.write('\n')
fi.write('<HEAD>\n')
fi.write('<TITLE> Javascript Movie Player </TITLE>\n')
fi.write('\n')
fi.write('</HEAD>\n')
fi.write('\n')
fi.write('<BODY FONT=3 BGCOLOR=FFFFFF LINK=#0000CC VLINK=#0000CC TEXT=#000000> \n')
fi.write('\n')
fi.write('<SCRIPT language="Javascript">\n')
fi.write('<!--\n')
fi.write('//\n')
fi.write('// DYNAMICALLY CREATED HTML - DO NOT EDIT\n')
fi.write('//\n')
fi.write('// Javascript program to produce fast animations by reading from cache\n')
fi.write('// Written: 7 July 1997, Zarro, NASA/GSFC\n')
fi.write('// Modified: 27 July 1997, Freeland, LSAL - added nice SWING option\n')
fi.write('// Modified: 9 Oct 1998, Zarro, SMA/GSFC - added context window button\n')
fi.write('// Modified: 14 Oct 1998, Zarro, SMA/GSFC - load images asynchronously\n')
fi.write('// Modified: 23 Aug 1999, William Thompson, GSFC - make button text black\n')
fi.write('\n')
fi.write('// Following variables control movie speed, labels, etc.\n')
fi.write('\n')
fi.write('var imax = {}, inc = 1.50, delay = 250;\n'.format(nfiles))
fi.write('var num_loaded_images = 0;\n')
fi.write('var frame=-1, speed=10;\n')
fi.write('var timeout_id=null;\n')
fi.write('var dir=1, playing=0, swingon=0, run=0;\n')
fi.write('var bname = "Reverse";\n')
fi.write('var url_path = ".";\n')
fi.write('var context = "frame00.gif";\n')
fi.write('var ctitle ="test";\n')
fi.write('var url_context=url_path+"/"+context;\n')
fi.write('var iwidth = {},iheight = {};\n'.format(width, height))
fi.write('var index=0;\n')
fi.write('\n')
fi.write('// -->\n')
fi.write('</SCRIPT>\n')
fi.write('\n')
fi.write('<CENTER>\n')
fi.write('\n')
fi.write('<H1> Javascript Movie Player </H1>\n')
fi.write('<P>\n')
fi.write('\n')
fi.write('<TABLE BORDER="10" CELLPADDING="8">\n')
fi.write('<TR>\n')
fi.write('<TD align="center"> \n')
fi.write('<img NAME=animation ALT="FRAME" width={} height={}>\n'.format(width, height))
fi.write('</TR>\n')
fi.write('</TABLE>\n')
fi.write('<P>\n')
fi.write('\n')
fi.write('<FORM NAME="form">\n')
fi.write(' <FONT COLOR="Black">\n')
fi.write(' <INPUT TYPE=button VALUE="Start" onClick="start_play();">\n')
fi.write(' <INPUT TYPE=button VALUE="Pause" onClick="stop_play();">\n')
fi.write(' <INPUT TYPE=button VALUE="Faster" onClick="speed*=inc; show_speed();">\n')
fi.write(' <INPUT TYPE=button VALUE="Slower" onClick="speed/=inc; show_speed();">\n')
fi.write(' <INPUT TYPE=button VALUE="Step" onClick="oneStep();">\n')
fi.write(' <INPUT TYPE=button NAME="direction" VALUE="Reverse" onClick="reverse();">\n')
fi.write(' <INPUT TYPE=button VALUE="Swing Mode:" onClick="swing_mode();">\n')
fi.write(' <INPUT TYPE=text VALUE="OFF" NAME="swing" SIZE=3>\n')
fi.write('\n')
fi.write(' </FONT>\n')
fi.write(' <BR>\n')
fi.write(' Frame: <INPUT TYPE=text VALUE="" NAME="frame" SIZE=22> \n')
fi.write(' Speed:<INPUT TYPE=text VALUE="" NAME="rate" SIZE=4> (frames/sec)\n')
fi.write('</FORM>\n')
fi.write(' \n')
fi.write('</FORM>\n')
fi.write('</CENTER>\n')
fi.write('\n')
fi.write('<P>\n')
fi.write('<HR>\n')
fi.write('<B>Document</B>: <I><SCRIPT>document.write(document.title);</SCRIPT></I><BR>\n')
fi.write('<B>URL</B>: <I><SCRIPT>document.write(document.URL);</SCRIPT></I><BR>\n')
fi.write('<B>Last Update</B>: <I><SCRIPT>document.write(document.lastModified);</SCRIPT></I><BR>\n')
fi.write('\n')
fi.write('\n')
fi.write('<SCRIPT LANGUAGE="JavaScript">\n')
fi.write('<!--\n')
fi.write('\n')
fi.write('function get_window_width(fac,def) { // Return window width\n')
fi.write('\n')
fi.write(' \n')
fi.write('var width;\n')
fi.write('if (!fac) {fac=.75;}\n')
fi.write('if (!def) {def=512;}\n')
fi.write('\n')
fi.write('if (window.screen) {\n')
fi.write(' width=parseInt(fac*parseFloat(window.screen.width));\n')
fi.write('} else {\n')
fi.write(' width=def;\n')
fi.write('}\n')
fi.write('\n')
fi.write('return width;\n')
fi.write('\n')
fi.write('}\n')
fi.write('\n')
fi.write('/////////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function get_window_height(fac,def) { // Return window height\n')
fi.write('\n')
fi.write(' \n')
fi.write('var height;\n')
fi.write('if (!fac) {fac=.75;}\n')
fi.write('if (!def) {def=512;}\n')
fi.write('\n')
fi.write('if (window.screen) {\n')
fi.write(' height=parseInt(fac*parseFloat(window.screen.height));\n')
fi.write('} else {\n')
fi.write(' height=def;\n')
fi.write('}\n')
fi.write('\n')
fi.write('return height;\n')
fi.write('\n')
fi.write('}\n')
fi.write('\n')
fi.write('/////////////////////////////////////////////////////////////////////////////\n')
fi.write('// Javascript Pop-up image Window\n')
fi.write('// Written: Zarro (SAC/GSFC), October 1998, (dzarro@solar.stanford.edu)\n')
fi.write('\n')
fi.write('var win;\n')
fi.write('\n')
fi.write('function pop_img(url,title,width,height) {\n')
fi.write('\n')
fi.write('if (!url) {\n')
fi.write(" alert('Image URL not entered');\n")
fi.write(' return false;\n')
fi.write('}\n')
fi.write('\n')
fi.write('// default to fit within 75% of browser window, or 512x512\n')
fi.write('\n')
fi.write('if (!width) {width=get_window_width(.75,512);}\n')
fi.write('if (!height) {height=get_window_height(.75,512);}\n')
fi.write('\n')
fi.write('if (!win || win.closed) {\n')
fi.write(
' win = open("","img","width="+width.toString()+",height="+height.toString()+",scrollbars=yes,resizable=yes"); \n')
fi.write('}\n')
fi.write('\n')
fi.write('// dynamically create HTML, adding a CLOSE button\n')
fi.write('\n')
fi.write('if (win) {\n')
fi.write(' d=win.document;\n')
fi.write(' d.write(\'<html><head><title>Image Window</title></head><body bgcolor="white"><center>\');\n')
fi.write(' if (title) {\n')
fi.write(" d.write('<h1>'+title+'</h1>');\n")
fi.write(' }\n')
fi.write(" d.write('<img src='+url+'>');\n")
fi.write(" d.write('<br><br><br>');\n")
fi.write(
' d.write(\'<form><b><input type="button" value="CLOSE" onClick="self.close();"></b></form></center>\'); \n')
fi.write(" d.write('</html></body>');\n")
fi.write(' d.close();\n')
fi.write(' win.focus();\n')
fi.write('}\n')
fi.write('\n')
fi.write('return true; \n')
fi.write('}\n')
fi.write('\n')
fi.write('///////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function load_img() { // asynchronously load all images into cache\n')
fi.write(' for (i=0; i < imax ; i++) {\n')
fi.write(' id[i]=setTimeout("load_src()",0);\n')
fi.write(' }\n')
fi.write(' return;\n')
fi.write('}\n')
fi.write('/////////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function load_src() { // load individual images into cache\n')
fi.write('\n')
fi.write(' if (index < imax) {\n')
fi.write(' if (iwidth && iheight) {\n')
fi.write(' images[index] = new Image(iwidth,iheight);\n')
fi.write(' } else {\n')
fi.write(' images[index] = new Image();\n')
fi.write(' }\n')
fi.write(' images[index].onload=count_images;\n')
fi.write(' images[index].src = urls[index];\n')
fi.write(' index++;\n')
fi.write(' }\n')
fi.write(' return;\n')
fi.write('}\n')
fi.write('\n')
fi.write('/////////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write("function clear_ids() { // clear asynchronous id's\n")
fi.write(' for (i=0; i < imax ; i++) {clearTimeout(id[i]);}\n')
fi.write(' return;\n')
fi.write('}\n')
fi.write('\n')
fi.write('\n')
fi.write('/////////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function count_images() // count images as they are loaded into cache\n')
fi.write('{ \n')
fi.write(' if (++num_loaded_images == imax) {\n')
fi.write(' show_speed();\n')
fi.write(' clear_ids();\n')
fi.write(' animate();\n')
fi.write(' } else {\n')
fi.write(' document.animation.src=images[num_loaded_images-1].src;\n')
fi.write(' document.form.frame.value="Loading "+num_loaded_images+" of "+imax; \n')
fi.write(' }\n')
fi.write('}\n')
fi.write('\n')
fi.write('///////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function image_abort() // abort loading images\n')
fi.write('{ \n')
fi.write(' imax=num_loaded_images;\n')
fi.write(' if (!images[num_loaded_images].complete) imax=imax-1;\n')
fi.write(' alert("Aborting");\n')
fi.write(' if (imax > -1) animate(); \n')
fi.write('}\n')
fi.write('\n')
fi.write('///////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function image_error(message) // abort loading images\n')
fi.write('{ \n')
fi.write(' alert(message);\n')
fi.write('}\n')
fi.write('\n')
fi.write('///////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function start_play() // start movie\n')
fi.write('{\n')
fi.write(' if (playing == 0) {\n')
fi.write(' if (timeout_id == null && num_loaded_images==imax) animate();\n')
fi.write(' }\n')
fi.write('} \n')
fi.write('\n')
fi.write('///////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function stop_play() // stop movie\n')
fi.write('{ \n')
fi.write(' if (timeout_id) clearTimeout(timeout_id); \n')
fi.write(' timeout_id=null;\n')
fi.write(' playing = 0;\n')
fi.write('} \n')
fi.write('\n')
fi.write('///////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function swing_mode() // set swing mode\n')
fi.write('{\n')
fi.write(' if (swingon) {\n')
fi.write(' swingon=0;\n')
fi.write(' document.form.swing.value="OFF";\n')
fi.write(' }\n')
fi.write(' else {\n')
fi.write(' swingon=1;\n')
fi.write(' document.form.swing.value="ON";\n')
fi.write(' }\n')
fi.write('}\n')
fi.write('\n')
fi.write('///////////////////////////////////////////////////////////////////////////\n')
fi.write(' \n')
fi.write('function animate() // control movie loop\n')
fi.write('{\n')
fi.write(' var j;\n')
fi.write(' frame=(frame+dir+imax)%imax;\n')
fi.write(' j=frame+1;\n')
fi.write(' if (images[frame].complete) {\n')
fi.write(' document.animation.src=images[frame].src;\n')
fi.write(' document.form.frame.value="Displaying "+j+" of "+imax;\n')
fi.write(' if (swingon && (j == imax || frame == 0)) reverse();\n')
fi.write(' timeout_id=setTimeout("animate()",delay);\n')
fi.write(' playing=1;\n')
fi.write(' }\n')
fi.write('}\n')
fi.write('\n')
fi.write('///////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function oneStep() // step frames\n')
fi.write('{\n')
fi.write(' var j;\n')
fi.write(' if (timeout_id) clearTimeout(timeout_id); timeout_id=null;\n')
fi.write(' frame=(frame+dir+imax)%imax;\n')
fi.write(' j=frame+1;\n')
fi.write(' if (images[frame].complete) {\n')
fi.write(' document.animation.src=images[frame].src;\n')
fi.write(' document.form.frame.value="Displaying "+j+" of "+imax;\n')
fi.write(' playing=0;\n')
fi.write(' }\n')
fi.write('}\n')
fi.write('\n')
fi.write('///////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function reverse() // reverse direction\n')
fi.write('{\n')
fi.write(' dir=-dir;\n')
fi.write(' if (dir > 0) document.form.direction.value="Reverse"; bname="Reverse";\n')
fi.write(' if (dir < 0) document.form.direction.value="Forward"; bname="Forward";\n')
fi.write('}\n')
fi.write('\n')
fi.write('///////////////////////////////////////////////////////////////////////////\n')
fi.write('\n')
fi.write('function show_speed() // show speed\n')
fi.write('{\n')
fi.write(' document.form.rate.value=Math.round(speed);\n')
fi.write(' delay = 1000.0/speed; \n')
fi.write('}\n')
fi.write('\n')
fi.write('///////////////////////////////////////////////////////////////////////////\n')
fi.write('// actual image loading is done here\n')
fi.write('\n')
fi.write('show_speed();\n')
fi.write('images = new Array(imax); \n')
fi.write('urls= new Array(imax);\n')
fi.write('id= new Array(imax);\n')
fi.write('\n')
fi.write('\n')
for i in range(nfiles):
fi.write('urls[{}]=url_path+"/{}";\n'.format(i, os.path.basename(imgfiles[i])))
fi.write('\n')
fi.write('load_img();\n')
fi.write('\n')
fi.write('// -->\n')
fi.write('</SCRIPT>\n')
fi.write('\n')
fi.write('</BODY>\n')
fi.write('</HTML>\n')
fi.write('\n')
fi.close() # pdb.set_trace()
print('\nWrite movie to {}'.format(htmlfile))
# print('To open -----> firefox {} &'.format(os.path.abspath(htmlfile)))
[docs]
def my_timer(orig_func):
import time
@wraps(orig_func)
def wrapper(*args, **kwargs):
t1 = time.time()
result = orig_func(*args, **kwargs)
t2 = time.time() - t1
print('{} ran in: {} sec'.format(orig_func.__name__, t2))
return result
return wrapper
[docs]
def spectrogram2wav(spec, threshld=None, gama=1, fs=1.0, t_length=None, w=1, wavfile='output.wav'):
'''
Convert spectrogram to audio in WAV format
:param spec: spec.shape (nfreq, ntime)
:param threshld: below which is set to be threshold
:param gama:
:param fs:
:param t_length: time duration of output WAV file
:param w: width of the smooth window, if apply
:param wavfile:
:return:
'''
from scipy import signal
from scipy.io.wavfile import write
import numpy as np
if np.sum(np.isnan(spec)) > 0:
for idx in range(spec.shape[1]):
p_slice = spec[:, idx]
mask_nan = np.isnan(p_slice)
if np.sum(~mask_nan) > 0:
p_slice[mask_nan] = np.interp(np.flatnonzero(mask_nan), np.flatnonzero(~mask_nan), p_slice[~mask_nan])
spec[:, idx] = p_slice
for idx in range(spec.shape[0]):
p_slice = spec[idx, :]
mask_nan = np.isnan(p_slice)
if np.sum(~mask_nan) > 0:
p_slice[mask_nan] = np.interp(np.flatnonzero(mask_nan), np.flatnonzero(~mask_nan), p_slice[~mask_nan])
spec[idx, :] = p_slice
# smooth the dynamic spectrum
if w > 1:
k = 1 - np.abs(np.linspace(-1, 1, w))
kernel = k.reshape(w, 1) * k.reshape(1, w)
kernel /= kernel.sum() # kernel should sum to 1! :)
Zxx = signal.convolve2d(spec, kernel, mode='same')
else:
Zxx = spec
if threshld:
Zxx = np.where(Zxx >= threshld, Zxx, threshld)
_, xrec = signal.istft(Zxx ** gama, fs)
scaled = np.int16(xrec / np.max(np.abs(xrec)) * 32767)
if t_length:
sample_rate = int(np.round(len(scaled) / t_length))
else:
sample_rate = 44100
write(wavfile, sample_rate, scaled)
return
[docs]
def smooth(x, window_len=11, window='hanning'):
"""smooth the data using a window with requested size.
This method is based on the convolution of a scaled window with the signal.
The signal is prepared by introducing reflected copies of the signal
(with the window size) in both ends so that transient parts are minimized
in the begining and end part of the output signal.
input:
x: the input signal
window_len: the dimension of the smoothing window; should be an odd integer
window: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'
flat window will produce a moving average smoothing.
output:
the smoothed signal
example:
t=linspace(-2,2,0.1)
x=sin(t)+randn(len(t))*0.1
y=smooth(x)
see also:
numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
scipy.signal.lfilter
NOTE: length(output) != length(input), to correct this: return y[(window_len/2-1):-(window_len/2)] instead of just y.
"""
if x.ndim != 1:
raise ValueError("smooth only accepts 1 dimension arrays.")
if x.size < window_len:
raise ValueError("Input vector needs to be bigger than window size.")
if window_len < 3:
return x
if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
raise ValueError("Window is on of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'")
s = np.r_[x[window_len - 1:0:-1], x, x[-2:-window_len - 1:-1]]
# print(len(s))
if window == 'flat': # moving average
w = np.ones(window_len, 'd')
else:
w = eval('np.' + window + '(window_len)')
y = np.convolve(w / w.sum(), s, mode='same')
if len(x) == len(y):
return y
else:
return y[window_len - 1:-(window_len - 1)]
# def smooth(x, window_len=11, window='hanning', mode='same'):
# """smooth the data using a window with requested size.
#
# This method is based on the convolution of a scaled window with the signal.
# The signal is prepared by introducing reflected copies of the signal
# (with the window size) in both ends so that transient parts are minimized
# in the begining and end part of the output signal.
#
# input:
# x: the input signal
# window_len: the dimension of the smoothing window; should be an odd integer
# window: the type of window from 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'
# flat window will produce a moving average smoothing.
#
# output:
# the smoothed signal
#
# example:
#
# t=linspace(-2,2,0.1)
# x=sin(t)+randn(len(t))*0.1
# y=smooth(x)
#
# see also:
#
# numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman, numpy.convolve
# scipy.signal.lfilter
#
# TODO: the window parameter could be the window itself if an array instead of a string
# NOTE: length(output) != length(input), to correct this: return y[(window_len/2-1):-(window_len/2)] instead of just y.
# """
# import numpy
# if x.ndim != 1:
# raise ValueError("smooth only accepts 1 dimension arrays.")
#
# if x.size < window_len:
# raise ValueError("Input vector needs to be bigger than window size.")
#
# if window_len < 3:
# return x
#
# if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
# raise ValueError("Window is on of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'")
#
# s = numpy.r_[x[window_len - 1:0:-1], x, x[-2:-window_len - 1:-1]]
# # print(len(s))
# if window == 'flat': # moving average
# w = numpy.ones(window_len, 'd')
# else:
# w = eval('numpy.' + window + '(window_len)')
#
# y = numpy.convolve(w / w.sum(), s, mode=mode)
# if mode == 'same':
# return y[np.int_(window_len) - 1:-np.int_(window_len) + 1]
# else:
# return y[np.int_(window_len / 2 - 1):-np.int_(window_len / 2)]
[docs]
def img2movie(imgprefix='', img_ext='png', outname='movie', size=None, start_num=0, crf=15, fps=10, overwrite=False,
crop=[], title=[], dpi=200, keeptmp=False, usetmp=False, autorotate=True):
'''
:param imgprefix:
:param img_ext:
:param outname:
:param size:
:param start_num:
:param crf:
:param fps:
:param overwrite:
:param title: the timestamp on each frame
:param crop: 4-tuple of integer specifies the cropping pixels [x0, x1, y0, y1]
:param dpi:
:param keeptmp:
:param usetmp: use the image in the default tmp folder
:return:
'''
import subprocess, os
from tqdm import tqdm
import matplotlib.pyplot as plt
if type(imgprefix) is list:
imgs = imgprefix
imgprefix = os.path.dirname(imgprefix[0])
else:
imgs = glob.glob(imgprefix + '*.' + img_ext)
if imgs:
imgs = sorted(imgs)
tmpdir = os.path.join(os.path.dirname(imgprefix), 'img_tmp') + '/'
if not os.path.exists(tmpdir):
os.makedirs(tmpdir)
if usetmp:
pass
else:
if crop != [] or title != []:
plt.ioff()
fig, ax = plt.subplots(figsize=[float(ll) / dpi for ll in size.split('x')])
for idx, ll in enumerate(tqdm(imgs)):
data = plt.imread(ll)
if crop != []:
x0, x1, y0, y1 = crop
data = data[y0:y1 + 1, x0:x1 + 1, :]
if idx == 0:
im = ax.imshow(data)
ax.set_axis_off()
else:
im.set_array(data)
if title != []:
ax.set_title(title[idx])
if idx == 0:
fig.tight_layout()
fig.savefig('{}/{:04d}.{}'.format(tmpdir, idx, img_ext), dpi=dpi, format=img_ext)
plt.close(fig)
else:
for idx, ll in enumerate(imgs):
os.system('cp {} {}/{:04d}.{}'.format(ll, tmpdir, idx, img_ext))
outd = {'r': fps, 's': size, 'start_number': start_num, 'crf': crf}
if size is None:
outd.pop('s')
if overwrite:
ow = '-y'
else:
ow = ''
if autorotate:
noautorotate = ''
else:
noautorotate = '-noautorotate'
try:
outdstr = ' '.join(['-{} {}'.format(k, v) for k, v in outd.iteritems()])
except:
outdstr = ' '.join(['-{} {}'.format(k, outd[k]) for k in outd])
try:
cmd = 'ffmpeg {4} -r {3} -f image2 -i {0}%04d.{1} -vcodec libx264 -pix_fmt yuv420p {2} '.format(tmpdir,
img_ext,
outdstr,
fps,
noautorotate) + '{0} {1}.mp4'.format(
ow, os.path.join(os.path.dirname(imgprefix), outname))
subprocess.check_output(['bash', '-c', cmd])
except:
cmd = 'ffmpeg {4} -r {3} -f image2 -i {0}%04d.{1} -vcodec libx264 -pix_fmt yuv420p {2} -vf "pad=ceil(iw/2)*2:ceil(ih/2)*2"'.format(
tmpdir, img_ext,
outdstr,
fps, noautorotate) + '{0} {1}.mp4'.format(
ow, os.path.join(os.path.dirname(imgprefix), outname))
subprocess.check_output(['bash', '-c', cmd])
print(cmd)
if not keeptmp:
os.system('rm -rf {}'.format(tmpdir))
else:
print('Images not found!')
[docs]
def image_fill_gap(image):
for idx in range(image.shape[0]):
image_slice = image[idx, :]
mask_nan = np.isnan(image_slice)
if np.sum(~mask_nan) > 0 and np.sum(mask_nan) > 0:
image_slice[mask_nan] = np.interp(np.flatnonzero(mask_nan), np.flatnonzero(~mask_nan),
image_slice[~mask_nan])
image[idx, :] = image_slice
for idx in range(image.shape[1]):
image_slice = image[:, idx]
mask_nan = np.isnan(image_slice)
if np.sum(~mask_nan) > 0 and np.sum(mask_nan) > 0:
image_slice[mask_nan] = np.interp(np.flatnonzero(mask_nan), np.flatnonzero(~mask_nan),
image_slice[~mask_nan])
image[:, idx] = image_slice
return image
[docs]
def getspwfromfreq(vis, freqrange):
from taskinit import ms
ms.open(vis)
axisInfo = ms.getdata(["axis_info"], ifraxis=True)
spwInfo = ms.getspectralwindowinfo()
freqInfo = axisInfo["axis_info"]["freq_axis"]["chan_freq"].swapaxes(0, 1) / 1e9
freqInfo_ravel = freqInfo.ravel()
timeInfo = axisInfo["axis_info"]["time_axis"]['MJDseconds']
mstimran = ms.range(["time"])
ms.close()
freq0, freq1 = freqrange.split(' ')[0].split('~')
freq0, freq1 = float(freq0), float(freq1)
for ll in [freq0, freq1]:
if not freqInfo_ravel[0] <= ll <= freqInfo_ravel[-1]:
raise ValueError('Selected frequency out of range!!!')
freqIdx0 = np.where(freqInfo == freq0)
freqIdx1 = np.where(freqInfo == freq1)
sz_freqInfo = freqInfo.shape
ms_spw = ['{}'.format(ll) for ll in range(freqIdx0[0], freqIdx1[0] + 1)]
if len(ms_spw) == 1:
ms_chan = ['{}~{}'.format(freqIdx0[1][0], freqIdx1[1][0])]
else:
ms_chan = ['{}~{}'.format(freqIdx0[1][0], sz_freqInfo[1] - 1)] + ['0~{}'.format(sz_freqInfo[1] - 1) for ll in
range(freqIdx0[0] + 1, freqIdx1[0])]
ms_chan.append('0~{}'.format(freqIdx1[1][0]))
spw = ','.join('{}:{}'.format(t[0], t[1]) for t in zip(ms_spw, ms_chan))
return spw
[docs]
def initconfig(suncasa_dir):
if not os.path.exists(suncasa_dir + 'DataBrowser/config.json'):
os.system(
'cp {} {}'.format(suncasa_dir + 'DataBrowser/config_init.json', suncasa_dir + 'DataBrowser/config.json'))
return True
else:
return False
# def mkunidir(dirlist, isdir=True):
# '''
# get the list of all unique directories from the filelist, make these directories
# :param dirlist:
# :param isdir: if dirlist is a list of directories.
# if not, dirlist is a list of full path of files,
# the base name will be removed from the paths.
# :return:
# '''
# DEll19432017 import os
# if not isdir:
# dirlist = [os.path.dirname(ff) for ff in dirlist]
# dirs = list(set(dirlist))
# for ll in dirs:
# if not os.path.exists(ll):
# os.makedirs(ll)
[docs]
def ProgressBar(iteration, total, prefix='', suffix='', decimals=1, length=100, empfill=' ', fill='#'):
"""
Call in a loop to create terminal progress bar
@params:
iteration - Required : current iteration (Int)
total - Required : total iterations (Int)
prefix - Optional : prefix string (Str)
suffix - Optional : suffix string (Str)
decimals - Optional : positive number of decimals in percent complete (Int)
length - Optional : character length of bar (Int)
fill - Optional : bar fill character (Str)
empfill - Optional : empty bar fill character (Str)
"""
percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
filledLength = int(length * iteration // total)
bar = fill * filledLength + empfill * (length - filledLength)
# return '%s |%s| %s%% %s' % (prefix, bar, percent, suffix)
return '{} |{}| {}% {}'.format(prefix, bar, percent, suffix)
[docs]
def getcurtimstr(prefix='CleanID_', suffix=''):
import time
return prefix + time.strftime("%Y%m%d_%H%M%S") + suffix
[docs]
def getlatestfile(directory='./', prefix='CleanID_', suffix=''):
filelist = glob.glob('{}/{}*{}'.format(directory, prefix, suffix))
if len(filelist) > 0:
latest_file = max(filelist, key=os.path.getctime)
print(latest_file)
return {'items': filelist, 'latest': latest_file, 'idx': filelist.index(latest_file)}
else:
print('No file found!')
return None
[docs]
def loadjsonfile(jsonfile, mustexist=True):
if os.path.exists(jsonfile):
with open(jsonfile, 'r') as fp:
data = json.load(fp)
return data
else:
if mustexist:
raise SystemExit('{} not found!'.format(jsonfile))
else:
return None
[docs]
def updatejsonfile(jsonfile, data):
with open(jsonfile, 'w') as fp:
json.dump(data, fp)
[docs]
def getSDOdir(config, database_dir, suncasa_dir):
try:
if config['datadir']['SDOdir']:
SDOdir = config['datadir']['SDOdir']
if SDOdir.startwith('$'):
SDOdir = os.path.expandvars(SDOdir)
if not os.path.exists(SDOdir):
os.makedirs(SDOdir)
else:
raise ValueError
except:
SDOdir = database_dir + 'Download/'
config['datadir']['SDOdir'] = SDOdir
fout = suncasa_dir + 'DataBrowser/config.json'
updatejsonfile(fout, config)
return SDOdir
[docs]
def getsdodir(filename, unique=True):
'''
return a list of the data path relative to the SDOdir
:param filename:
:return:
'''
if type(filename) == str:
filename = [filename]
dirlist = []
timstrs = []
if unique:
for ll in filename:
timstr = os.path.basename(ll).split('.')[2].split('T')[0]
ymd = timstr.replace('-', '/')
dirlist.append('/{}/_{}'.format(ymd, timstr))
dirlist = list(set(dirlist))
for ll, dd in enumerate(dirlist):
dirlist[ll], timstr = dd.split('_')
timstrs.append(timstr)
else:
for ll in filename:
timstr = os.path.basename(ll).split('.')[2].split('T')[0]
ymd = timstr.replace('-', '/')
dirlist.append('/{}/'.format(ymd))
timstrs.append(timstr)
return {'dir': dirlist, 'timstr': timstrs}
[docs]
def FileNotInList(file2chk, filelist):
'''
return the index of files not in the list
:param file2chk: files to be check
:param filelist: the list
:return:
'''
idxs = []
if len(file2chk) > 0:
filelist = [os.path.basename(ll) for ll in filelist]
for idx, items in enumerate(file2chk):
if items not in filelist:
idxs.append(idx)
return idxs
[docs]
def getfreeport():
import socket
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind(('localhost', 0))
port = s.getsockname()[1]
s.close()
return port
[docs]
def normalize_aiamap(aiamap):
'''
do expisure normalization of an aia map
:param aia map made from sunpy.map:
:return: normalised aia map
'''
import sunpy.map as smap
try:
if aiamap.observatory == 'SDO' and aiamap.instrument[0:3] == 'AIA':
data = aiamap.data.copy().astype(float)
idxpix = ~np.isnan(data)
data[idxpix] = data[idxpix] / aiamap.exposure_time.value
data[data < 0] = 0
aiamap.meta['exptime'] = 1.0
aiamap = smap.Map(data, aiamap.meta)
return aiamap
else:
raise ValueError('input sunpy map is not from aia.')
except:
raise ValueError('check your input map. There are some errors in it.')
[docs]
def tplt(mapcube):
from astropy.time import Time
t = []
for idx, mp in enumerate(mapcube):
if mp.meta.has_key('t_obs'):
tstr = mp.meta['t_obs']
else:
tstr = mp.meta['date-obs']
t.append(tstr)
return Time(t)
[docs]
def sdo_aia_scale_hdr(amap, sigma=None):
import sunpy.map as smap
from astropy.coordinates import SkyCoord
wavelnth = '{:0.0f}'.format(amap.wavelength.value)
pxscale_x, pxscal_y = amap.scale
pxscale = (pxscale_x + pxscal_y) / 2
r_sun = amap.rsun_obs / pxscale
x = np.arange(amap.dimensions.x.value)
y = np.arange(amap.dimensions.y.value)
xx, yy = np.meshgrid(x, y) * u.pix
crpix1, crpix2 = amap.world_to_pixel(SkyCoord(0 * u.arcsec, 0 * u.arcsec, frame=amap.coordinate_frame))
rdist = np.sqrt((xx - (crpix1 - 1.0 * u.pix)) ** 2 + (yy - (crpix2 - 1.0 * u.pix)) ** 2)
ind_disk = np.where(rdist <= r_sun)
# ind_limb = np.where(rdist < r_sun)
rdist[ind_disk] = r_sun
rfilter = rdist / r_sun - 1
rfilter = rfilter.value
if sigma:
mapdata = amap.data * np.exp(rfilter * sigma)
else:
if wavelnth == '94':
mapdata = amap.data * np.exp(rfilter * 4)
elif wavelnth == '131':
mapdata = amap.data * (np.sqrt(rfilter * 5) + 1)
elif wavelnth == '171':
mapdata = amap.data * np.exp(rfilter * 5)
elif wavelnth == '193':
mapdata = amap.data * np.exp(rfilter * 3)
elif wavelnth == '211':
mapdata = amap.data * np.exp(rfilter * 3)
elif wavelnth == '304':
mapdata = amap.data * np.exp(rfilter * 5)
elif wavelnth == '335':
mapdata = amap.data * np.exp(rfilter * 3)
elif wavelnth == '6173':
pass
elif wavelnth == '1':
pass
else:
sigma = 5.0
mapdata = amap.data * np.exp(rfilter * sigma)
return smap.Map(mapdata, amap.meta)
[docs]
def sdo_aia_scale_dict(wavelength=None, imagetype='image'):
'''
rescale the aia image
:param image: normalised aia image data
:param wavelength:
:return: byte scaled image data
'''
if type(wavelength) is not str:
wavelength = '{:0.0f}'.format(wavelength)
if wavelength == '94':
if imagetype == 'image':
return {'low': 0.1, 'high': 3000, 'log': True}
elif imagetype == 'RDimage':
return {'low': -30, 'high': 30, 'log': False}
elif imagetype == 'BDimage':
return {'low': -30, 'high': 30, 'log': False}
elif imagetype == 'RDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif imagetype == 'BDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif wavelength == '131':
if imagetype == 'image':
return {'low': 0.5, 'high': 10000, 'log': True}
elif imagetype == 'RDimage':
return {'low': -100, 'high': 100, 'log': False}
elif imagetype == 'BDimage':
return {'low': -100, 'high': 100, 'log': False}
elif imagetype == 'RDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif imagetype == 'BDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif wavelength == '171':
if imagetype == 'image':
return {'low': 20, 'high': 5000, 'log': True}
elif imagetype == 'RDimage':
return {'low': -400, 'high': 400, 'log': False}
elif imagetype == 'BDimage':
return {'low': -400, 'high': 400, 'log': False}
elif imagetype == 'RDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif imagetype == 'BDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif wavelength == '193':
if imagetype == 'image':
return {'low': 30, 'high': 5000, 'log': True}
elif imagetype == 'RDimage':
return {'low': -1500, 'high': 1500, 'log': False}
elif imagetype == 'BDimage':
return {'low': -1500, 'high': 1500, 'log': False}
elif imagetype == 'RDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif imagetype == 'BDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif wavelength == '211':
if imagetype == 'image':
return {'low': 10, 'high': 8000, 'log': True}
elif imagetype == 'RDimage':
return {'low': -300, 'high': 300, 'log': False}
elif imagetype == 'BDimage':
return {'low': -300, 'high': 300, 'log': False}
elif imagetype == 'RDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif imagetype == 'BDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif wavelength == '304':
if imagetype == 'image':
return {'low': 1, 'high': 10000, 'log': True} # return {'low': 1, 'high': 500, 'log': True}
elif imagetype == 'RDimage':
return {'low': -300, 'high': 300, 'log': False}
elif imagetype == 'BDimage':
return {'low': -300, 'high': 300, 'log': False}
elif imagetype == 'RDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif imagetype == 'BDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif wavelength == '335':
if imagetype == 'image':
return {'low': 0.1, 'high': 800, 'log': True}
elif imagetype == 'RDimage':
return {'low': -15, 'high': 15, 'log': False}
elif imagetype == 'BDimage':
return {'low': -15, 'high': 15, 'log': False}
elif imagetype == 'RDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif imagetype == 'BDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif wavelength == '1600':
if imagetype == 'image':
return {'low': 20, 'high': 2500, 'log': True}
elif imagetype == 'RDimage':
return {'low': -800, 'high': 800, 'log': False}
elif imagetype == 'BDimage':
return {'low': -800, 'high': 800, 'log': False}
elif imagetype == 'RDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif imagetype == 'BDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif wavelength == '1700':
if imagetype == 'image':
return {'low': 300, 'high': 10000, 'log': True}
elif imagetype == 'RDimage':
return {'low': -1500, 'high': 1500, 'log': False}
elif imagetype == 'BDimage':
return {'low': -1500, 'high': 1500, 'log': False}
elif imagetype == 'RDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
elif imagetype == 'BDRimage':
return {'low': -1.5, 'high': 1.5, 'log': False}
else:
return {'high': None, 'log': False, 'low': None}
[docs]
def sdo_aia_scale(image=None, wavelength=None):
'''
rescale the aia image
:param image: normalised aia image data
:param wavelength:
:return: byte scaled image data
'''
from scipy.misc import bytescale
clrange = sdo_aia_scale_dict(wavelength)
image[image > clrange['high']] = clrange['high']
image[image < clrange['low']] = clrange['low']
image = np.log10(image)
return bytescale(image)
[docs]
def insertchar(source_str, insert_str, pos):
return source_str[:pos] + insert_str + source_str[pos:]
# def get_trange_files(trange):
# '''
# Given a timerange, this routine will take all relevant SDO files from that time range,
# put them in a list, and return that list.
# :param trange: two elements list of timestamps in Julian days
# :return:
# '''
[docs]
def readsdofile(datadir=None, wavelength=None, trange=None, isexists=False, timtol=1):
'''
read sdo file from local database
:param datadir:
:param wavelength:
:param trange: the timestamp or timerange in Julian days. if is timerange, return a list of files in the timerange
:param isexists: check if file exist. if files exist, return file name
:param timtol: time difference tolerance in days for considering data as the same timestamp
:return:
'''
from astropy.time import Time
import sunpy.map
from datetime import date
from datetime import timedelta as td
trange = Time(trange)
wavelength = str(wavelength)
wavelength = wavelength.lower()
if timtol < 12. / 3600 / 24:
timtol = 12. / 3600 / 24
if len(trange.iso) == 2:
if trange[1] < trange[0]:
raise ValueError('start time must be occur earlier than end time!')
else:
sdofitspath = []
jdtimestr = [Time(ll, format='jd').iso for ll in trange]
ymd = [ll.split(' ')[0].split('-') for ll in jdtimestr]
d1 = date(int(ymd[0][0]), int(ymd[0][1]), int(ymd[0][2]))
d2 = date(int(ymd[1][0]), int(ymd[1][1]), int(ymd[1][2]))
delta = d2 - d1
for i in range(delta.days + 1):
ymd = d1 + td(days=i)
sdofitspathtmp = glob.glob(
datadir + '/{:04d}/{:02d}/{:02d}/aia.lev1_*Z.{}.image_lev1.fits'.format(ymd.year, ymd.month,
ymd.day, wavelength))
if len(sdofitspathtmp) > 0:
sdofitspath = sdofitspath + sdofitspathtmp
if len(sdofitspath) == 0:
if isexists:
return sdofitspath
else:
raise ValueError(
'No SDO file found under {} at the time range of {} to {}. Download the data with EvtBrowser first.'.format(
datadir,
jdtimestr[0],
jdtimestr[1]))
sdofits = [os.path.basename(ll) for ll in sdofitspath]
sdotimeline = Time(
[insertchar(insertchar(ll.split('.')[2].replace('T', ' ').replace('Z', ''), ':', -4), ':', -2) for ll in
sdofits],
format='iso', scale='utc')
sdofitspathnew = [x for (y, x) in sorted(zip(sdotimeline.jd, sdofitspath))]
sdofitsnew = [os.path.basename(ll) for ll in sdofitspathnew]
sdotimelinenew = Time(
[insertchar(insertchar(ll.split('.')[2].replace('T', ' ').replace('Z', ''), ':', -4), ':', -2) for ll in
sdofitsnew], format='iso',
scale='utc')
sdofile = list(
np.array(sdofitspathnew)[
np.where(np.logical_and(trange[0].jd <= sdotimelinenew.jd, sdotimelinenew.jd <= trange[1].jd))[0]])
return sdofile
else:
jdtimstr = trange.iso
ymd = jdtimstr.split(' ')[0].split('-')
sdofitspath = glob.glob(
datadir + '/{}/{}/{}/aia.lev1_*Z.{}.image_lev1.fits'.format(ymd[0], ymd[1], ymd[2], wavelength))
if len(sdofitspath) == 0:
return [] # raise ValueError('No SDO file found under {}.'.format(datadir))
sdofits = [os.path.basename(ll) for ll in sdofitspath]
sdotimeline = Time(
[insertchar(insertchar(ll.split('.')[2].replace('T', ' ').replace('Z', ''), ':', -4), ':', -2) for ll in
sdofits],
format='iso', scale='utc')
if timtol < np.min(np.abs(sdotimeline.jd - trange.jd)):
return [] # raise ValueError('No SDO file found at the select timestamp. Download the data with EvtBrowser first.')
idxaia = np.argmin(np.abs(sdotimeline.jd - trange.jd))
sdofile = sdofitspath[idxaia]
if isexists:
return sdofile
else:
try:
sdomap = sunpy.map.Map(sdofile)
return sdomap
except:
raise ValueError('File not found or invalid input')
[docs]
def readsdofileX(datadir=None, filelist=None, wavelength=None, trange=None, isexists=False, timtol=1):
'''
read sdo file from local database
:param datadir:
:param wavelength:
:param trange: time object. the timestamp or timerange. if is timerange, return a list of files in the timerange
:param isexists: check if file exist. if files exist, return file name
:param timtol: time difference tolerance in days for considering data as the same timestamp
:return:
'''
from astropy.time import Time
import sunpy.map
from datetime import date
from datetime import timedelta as td
trange = Time(trange)
wavelength = str(wavelength)
wavelength = wavelength.lower()
if timtol < 12. / 3600 / 24:
timtol = 12. / 3600 / 24
# import pdb
# pdb.set_trace()
if len(trange.iso) == 2:
if trange.jd[1] < trange.jd[0]:
raise ValueError('start time must be occur earlier than end time!')
else:
if filelist:
sdofitspath = filelist
else:
sdofitspath = []
jdtimestr = trange.iso
ymd = [ll.split(' ')[0].split('-') for ll in jdtimestr]
d1 = date(int(ymd[0][0]), int(ymd[0][1]), int(ymd[0][2]))
d2 = date(int(ymd[1][0]), int(ymd[1][1]), int(ymd[1][2]))
delta = d2 - d1
for i in range(delta.days + 1):
ymd = d1 + td(days=i)
sdofitspathtmp = glob.glob(
datadir + '/aia.lev1_*{0}*{1}*{2}*Z.{3}.image*.*'.format(ymd.year, ymd.month, ymd.day,
wavelength))
if len(sdofitspathtmp) > 0:
sdofitspath = sdofitspath + sdofitspathtmp
if len(sdofitspath) == 0:
if isexists:
return sdofitspath
else:
raise ValueError(
'No SDO file found under {} at the time range of {} to {}. Download the data with EvtBrowser first.'.format(
datadir,
jdtimestr[0],
jdtimestr[1]))
sdofits = [os.path.basename(ll) for ll in sdofitspath]
if 'hmi' in sdofits[0]:
sdotimeline = Time(
[insertchar(insertchar(
insertchar(insertchar(ll.split('.')[2].replace('_TAI', '').replace('_', ' '), ':', -4), ':',
-2), '-',
6), '-', 4) for ll in
sdofits],
format='iso', scale='utc')
else:
sdotimeline = Time(
[insertchar(insertchar(ll.split('.')[2].replace('T', ' ').replace('Z', ''), ':', -4), ':', -2) for
ll in
sdofits],
format='iso', scale='utc')
sdofitspathnew = [x for (y, x) in sorted(zip(sdotimeline.jd, sdofitspath))]
sdofitsnew = [os.path.basename(ll) for ll in sdofitspathnew]
if 'hmi' in sdofits[0]:
sdotimelinenew = Time(
[insertchar(insertchar(
insertchar(insertchar(ll.split('.')[2].replace('_TAI', '').replace('_', ' '), ':', -4), ':',
-2), '-',
6), '-', 4) for ll in
sdofitsnew],
format='iso', scale='utc')
else:
sdotimelinenew = Time(
[insertchar(insertchar(ll.split('.')[2].replace('T', ' ').replace('Z', ''), ':', -4), ':', -2) for
ll in
sdofitsnew],
format='iso', scale='utc')
sdofile = list(
np.array(sdofitspathnew)[
np.where(np.logical_and(trange.jd[0] <= sdotimelinenew.jd, sdotimelinenew.jd <= trange.jd[1]))[0]])
return sdofile
else:
jdtimstr = trange.iso
ymd = jdtimstr.split(' ')[0].split('-')
if filelist:
sdofitspath = filelist
else:
sdofitspath = glob.glob(
datadir + '/aia.lev1_*{0}*{1}*{2}*Z.{3}.image*.*'.format(ymd[0], ymd[1], ymd[2], wavelength))
if len(sdofitspath) == 0:
sdofitspath = glob.glob(
datadir + '/hmi*{0}*{1}*{2}*.*'.format(ymd[0], ymd[1], ymd[2]))
if len(sdofitspath) == 0:
return [] # raise ValueError('No SDO file found under {}.'.format(datadir))
sdofits = [os.path.basename(ll) for ll in sdofitspath]
if 'hmi' in sdofits[0]:
sdotimeline = Time(
[insertchar(insertchar(
insertchar(insertchar(ll.split('.')[2].replace('_TAI', '').replace('_', ' '), ':', -4), ':', -2),
'-',
6), '-', 4) for ll in
sdofits],
format='iso', scale='utc')
else:
sdotimeline = Time(
[insertchar(insertchar(ll.split('.')[2].replace('T', ' ').replace('Z', ''), ':', -4), ':', -2) for ll in
sdofits],
format='iso', scale='utc')
if timtol <= np.min(np.abs(sdotimeline.jd - trange.jd)):
return [] # raise ValueError('No SDO file found at the select timestamp. Download the data with EvtBrowser first.')
idxaia = np.argmin(np.abs(sdotimeline.jd - trange.jd))
sdofile = sdofitspath[idxaia]
if isexists:
return sdofile
else:
try:
sdomap = sunpy.map.Map(sdofile)
return sdomap
except:
raise ValueError('File not found or invalid input')
[docs]
def findDist(x, y):
dx = np.diff(x)
dy = np.diff(y)
dist = np.hypot(dx, dy)
return np.insert(dist, 0, 0.0)
[docs]
def paramspline(x, y, length, s=0):
from scipy.interpolate import splev, splprep
tck, u = splprep([x, y], s=s)
unew = np.linspace(0, u[-1], length)
out = splev(unew, tck)
xs, ys = out[0], out[1]
grads = get_curve_grad(xs, ys)
return {'xs': xs, 'ys': ys, 'grads': grads['grad'], 'posangs': grads['posang']}
[docs]
def polyfit(x, y, length, deg, keepxorder=False):
if keepxorder:
xs = np.linspace(x[0], x[-1], length)
else:
xs = np.linspace(np.nanmin(x), np.nanmax(x), length)
z = np.polyfit(x=x, y=y, deg=deg)
p = np.poly1d(z)
ys = p(xs)
if keepxorder:
return {'xs': xs, 'ys': ys}
else:
grads = get_curve_grad(xs, ys)
return {'xs': xs, 'ys': ys, 'grads': grads['grad'], 'posangs': grads['posang']}
[docs]
def htfit_warren2011(x, y, cutlength):
'''
reference: https://ui.adsabs.harvard.edu/abs/2011ApJ...742...92W/abstract
'''
from scipy.optimize import curve_fit
def fit_func(t, h0, vt, a0, tau):
return h0 + vt * t + a0 * tau ** 2 * (np.exp(-t / tau) - 1)
print(x,y)
x0 = x[0]
params = curve_fit(fit_func, x - x0, y)
[h0, vt, a0, tau] = params[0]
print([h0, vt, a0, tau])
xs = np.linspace(np.nanmin(x), np.nanmax(x), cutlength)
ys = fit_func(xs - x0, h0, vt, a0, tau)
grads = get_curve_grad(xs, ys)
return {'xs': xs, 'ys': ys, 'grads': grads['grad'], 'posangs': grads['posang']}
[docs]
def spline(x, y, length, s=0):
from scipy.interpolate import splev, splrep
xs = np.linspace(x.min(), x.max(), length)
tck = splrep(x, y, s=s)
ys = splev(xs, tck)
grads = get_curve_grad(xs, ys)
return {'xs': xs, 'ys': ys, 'grads': grads['grad'], 'posangs': grads['posang']}
[docs]
def get_curve_grad(x, y):
'''
get the grad of at data point
:param x:
:param y:
:return: grad,posang
'''
deltay = np.roll(y, -1) - np.roll(y, 1)
deltay[0] = y[1] - y[0]
deltay[-1] = y[-1] - y[-2]
deltax = np.roll(x, -1) - np.roll(x, 1)
deltax[0] = x[1] - x[0]
deltax[-1] = x[-1] - x[-2]
grad = deltay / deltax
posang = np.arctan2(deltay, deltax)
return {'grad': grad, 'posang': posang}
[docs]
def improfile(z, xi, yi, interp='cubic'):
'''
Pixel-value cross-section along line segment in an image
:param z: an image array
:param xi and yi: equal-length vectors specifying the pixel coordinates of the endpoints of the line segment
:param interp: interpolation type to sampling, 'nearest' or 'cubic'
:return: the intensity values of pixels along the line
'''
import scipy.ndimage
imgshape = z.shape
if len(xi) != len(yi):
raise ValueError('xi and yi must be equal-length!')
if len(xi) < 2:
raise ValueError('xi or yi must contain at least two elements!')
for idx, ll in enumerate(xi):
if not 0 < ll < imgshape[1]:
return np.nan
if not 0 < yi[idx] < imgshape[0]:
return np.nan
if len(xi) == 2:
length = np.hypot(np.diff(xi), np.diff(yi))[0]
x, y = np.linspace(xi[0], xi[1], length), np.linspace(yi[0], yi[1], length)
else:
x, y = xi, yi
if interp == 'cubic':
zi = scipy.ndimage.map_coordinates(z, np.vstack((x, y)))
else:
zi = z[np.floor(y).astype(int), np.floor(x).astype(int)]
return zi
[docs]
def canvaspix_to_data(smap, x, y):
import astropy.units as u
'''
Convert canvas pixel coordinates in MkPlot to data (world) coordinates by using
`~astropy.wcs.WCS.wcs_pix2world`.
:param smap: sunpy map
:param x: canvas Pixel coordinates of the CTYPE1 axis. (Normally solar-x)
:param y: canvas Pixel coordinates of the CTYPE2 axis. (Normally solar-y)
:return: world coordinates
'''
# xynew = smap.pixel_to_data(x * u.pix, y * u.pix)
mesh = smap.pixel_to_world(x * u.pix, y * u.pix)
mapx, mapy = mesh.Tx, mesh.Tx
xnew = mapx[0].value
ynew = mapy[1].value
return [xnew, ynew]
[docs]
def data_to_mappixel(smap, x, y):
import astropy.units as u
'''
Convert data (world) coordinates in MkPlot to pixel coordinates in smap by using
`~astropy.wcs.WCS.wcs_pix2world`.
:param smap: sunpy map
:param x: Data coordinates of the CTYPE1 axis. (Normally solar-x)
:param y: Data coordinates of the CTYPE2 axis. (Normally solar-y)
:return: pixel coordinates
'''
xynew = smap.data_to_pixel(x * u.arcsec, y * u.arcsec)
xnew = xynew[0].value
ynew = xynew[1].value
return [xnew, ynew]
[docs]
def transfitdict2DF(datain, gaussfit=True, getcentroid=False):
'''
convert the results from pimfit or pmaxfit tasks to pandas DataFrame structure.
:param datain: The component list from pimfit or pmaxfit tasks
:param gaussfit: True if the results is from pimfit, otherwise False.
:param getcentroid: If True returns the centroid
:return: the pandas DataFrame structure.
'''
import pandas as pd
ra2arcsec = 180. * 3600. / np.pi
dspecDF0 = pd.DataFrame()
for ll in datain['timestamps']:
tidx = datain['timestamps'].index(ll)
if datain['succeeded'][tidx]:
pols = datain['outputs'][tidx].keys()
dspecDFlist = []
dspecDF = pd.DataFrame()
for ppit in pols:
dspecDFtmp = pd.DataFrame()
shape_latitude = []
shape_longitude = []
shape_latitude_err = []
shape_longitude_err = []
shape_majoraxis = []
shape_minoraxis = []
shape_positionangle = []
peak = []
beam_major = []
beam_minor = []
beam_positionangle = []
freqstrs = []
fits_local = []
for comp in datain['outputs'][tidx][ppit]['results'].keys():
if comp.startswith('component'):
if gaussfit:
majoraxis = datain['outputs'][tidx][ppit]['results'][comp]['shape']['majoraxis']['value']
minoraxis = datain['outputs'][tidx][ppit]['results'][comp]['shape']['minoraxis']['value']
positionangle = datain['outputs'][tidx][ppit]['results'][comp]['shape']['positionangle'][
'value']
bmajor = datain['outputs'][tidx][ppit]['results'][comp]['beam']['beamarcsec']['major'][
'value']
bminor = datain['outputs'][tidx][ppit]['results'][comp]['beam']['beamarcsec']['minor'][
'value']
bpositionangle = \
datain['outputs'][tidx][ppit]['results'][comp]['beam']['beamarcsec']['positionangle'][
'value']
shape_majoraxis.append(majoraxis)
shape_minoraxis.append(minoraxis)
shape_positionangle.append(positionangle)
beam_major.append(bmajor)
beam_minor.append(bminor)
beam_positionangle.append(bpositionangle)
fluxpeak = datain['outputs'][tidx][ppit]['results'][comp]['peak']['value']
else:
fluxpeak = datain['outputs'][tidx][ppit]['results'][comp]['flux']['value'][0]
if getcentroid:
mkey = 'centroid'
else:
mkey = 'shape'
longitude = datain['outputs'][tidx][ppit]['results'][comp][mkey]['direction']['m0'][
'value'] * ra2arcsec
latitude = datain['outputs'][tidx][ppit]['results'][comp][mkey]['direction']['m1'][
'value'] * ra2arcsec
longitude_err = \
datain['outputs'][tidx][ppit]['results'][comp]['shape']['direction']['error']['longitude'][
'value']
latitude_err = \
datain['outputs'][tidx][ppit]['results'][comp]['shape']['direction']['error']['latitude'][
'value']
shape_longitude.append(longitude)
shape_latitude.append(latitude)
shape_longitude_err.append(longitude_err)
shape_latitude_err.append(latitude_err)
peak.append(fluxpeak)
freqstrs.append('{:.3f}'.format(
datain['outputs'][tidx][ppit]['results'][comp]['spectrum']['frequency']['m0']['value']))
fits_local.append(datain['imagenames'][tidx].split('/')[-1])
if gaussfit:
dspecDFtmp['shape_latitude{}'.format(ppit)] = pd.Series(shape_latitude)
dspecDFtmp['shape_longitude{}'.format(ppit)] = pd.Series(shape_longitude)
dspecDFtmp['shape_latitude_err{}'.format(ppit)] = pd.Series(shape_latitude_err)
dspecDFtmp['shape_longitude_err{}'.format(ppit)] = pd.Series(shape_longitude_err)
dspecDFtmp['peak{}'.format(ppit)] = pd.Series(peak)
dspecDFtmp['shape_majoraxis{}'.format(ppit)] = pd.Series(shape_majoraxis)
dspecDFtmp['shape_minoraxis{}'.format(ppit)] = pd.Series(shape_minoraxis)
dspecDFtmp['shape_positionangle{}'.format(ppit)] = pd.Series(shape_positionangle)
dspecDFtmp['beam_major{}'.format(ppit)] = pd.Series(beam_major)
dspecDFtmp['beam_minor{}'.format(ppit)] = pd.Series(beam_minor)
dspecDFtmp['beam_positionangle{}'.format(ppit)] = pd.Series(beam_positionangle)
else:
dspecDFtmp['shape_latitude{}'.format(ppit)] = pd.Series(shape_latitude)
dspecDFtmp['shape_longitude{}'.format(ppit)] = pd.Series(shape_longitude)
dspecDFtmp['shape_latitude_err{}'.format(ppit)] = pd.Series(shape_latitude_err)
dspecDFtmp['shape_longitude_err{}'.format(ppit)] = pd.Series(shape_longitude_err)
dspecDFtmp['peak{}'.format(ppit)] = pd.Series(peak)
dspecDFtmp['freqstr'.format(ppit)] = pd.Series(freqstrs)
dspecDFtmp['fits_local'.format(ppit)] = pd.Series(fits_local)
dspecDFlist.append(dspecDFtmp)
for DFidx, DFit in enumerate(dspecDFlist):
if DFidx == 0:
dspecDF = dspecDFlist[0]
else:
dspecDF = pd.merge(dspecDF.copy(), DFit, how='outer', on=['freqstr', 'fits_local'])
dspecDF0 = dspecDF0.append(dspecDF, ignore_index=True)
return dspecDF0
[docs]
def getcolctinDF(dspecDF, col):
'''
return the count of how many times of the element starts with col occurs in columns of dspecDF
:param dspecDF:
:param col: the start string
:return: the count and items
'''
itemset1 = col
itemset2 = dspecDF.columns.tolist()
items = []
for ll in itemset2:
if ll.startswith(itemset1):
items.append(ll)
itemct = [ll.startswith(itemset1) for ll in itemset2].count(True)
columns = items
return [itemct, columns]
[docs]
def dspecDFfilter(dspecDF, pol):
'''
filter the unselect polarisation from dspecDF
:param dspecDF: the original dspecDF
:param pol: selected polarisation, dtype = string
:return: the output dspecDF
'''
colnlistcom = ['shape_latitude', 'shape_longitude', 'peak', 'shape_latitude_err', 'shape_longitude_err']
colnlistgaus = ['shape_majoraxis', 'shape_minoraxis', 'shape_positionangle', 'beam_major', 'beam_minor',
'beam_positionangle']
## above are the columns to filter
colnlistess = dspecDF.columns.tolist()
if getcolctinDF(dspecDF, 'peak')[0] > 0:
for ll in colnlistcom + colnlistgaus:
colinfo = getcolctinDF(dspecDF, ll)
if colinfo[0] > 0:
for cc in colinfo[1]:
if cc in colnlistess:
colnlistess.remove(cc)
dspecDF1 = dspecDF.copy()[colnlistess]
for ll in colnlistcom:
dspecDF1[ll] = dspecDF.copy()[ll + pol]
if getcolctinDF(dspecDF, 'shape_majoraxis')[0] > 0:
for ll in colnlistgaus:
dspecDF1[ll] = dspecDF.copy()[ll + pol]
print('dspedDF is filtered')
return dspecDF1
else:
print('dspedDF no need filter')
return dspecDF
[docs]
def dspecDF2text(DFfile, outfile=None):
if DFfile:
if os.path.exists(DFfile):
if outfile:
with open(DFfile, 'rb') as f:
dspecDF0 = pickle.load(f)
dspecDF0.drop(['dspec', 'fits_global', 'fits_local'], axis=1, inplace=True)
dspecDF0.to_csv(outfile, sep='\t')
else:
raise ValueError('provide output file name!')
else:
raise ValueError('input file "{}" does not existed!'.format(DFfile))
else:
raise ValueError('provide input file name!')
[docs]
def smapmeshgrid2(smap, angle=None, rescale=1.0, origin=1):
import astropy.units as u
if angle is None:
mrot = smap.rotation_matrix
else:
sin = np.sin(angle)
cos = np.cos(angle)
mrot = np.array([[cos, -sin], [sin, cos]])
ref_pix = smap.reference_pixel
ref_coord = smap.reference_coordinate
scale = smap.scale
XX, YY = np.meshgrid(np.arange(smap.data.shape[1] * rescale) / rescale,
np.arange(smap.data.shape[0] * rescale) / rescale)
x, y = XX * u.pix, YY * u.pix
try:
x = (x - ref_pix[0] + origin * u.pix) * scale[0] + ref_coord.x
y = (y - ref_pix[1] + origin * u.pix) * scale[1] + ref_coord.y
except:
x = (x - ref_pix[0] + origin * u.pix) * scale[0] + ref_coord.Tx
y = (y - ref_pix[1] + origin * u.pix) * scale[1] + ref_coord.Ty
xnew = mrot[0, 0] * x + mrot[0, 1] * y
ynew = mrot[1, 0] * x + mrot[1, 1] * y
return xnew, ynew
[docs]
def map2wcsgrids(snpmap, cell=True, antialiased=True):
'''
:param snpmap:
:param cell: if True, return the coordinates of the pixel centers. if False, return the coordinates of the pixel boundaries
:return:
'''
# embed()
import astropy.units as u
ny, nx = snpmap.data.shape
x0, x1 = snpmap.xrange.to(u.arcsec).value
y0, y1 = snpmap.yrange.to(u.arcsec).value
dx = snpmap.scale.axis1.to(u.arcsec / u.pix).value
dy = snpmap.scale.axis2.to(u.arcsec / u.pix).value
if cell:
mapx, mapy = np.linspace(x0, x1, nx), np.linspace(y0, y1, ny)
mapx = np.tile(mapx, ny).reshape(ny, nx)
mapy = np.tile(mapy, nx).reshape(nx, ny).transpose()
else:
nx += 1
ny += 1
mapx, mapy = np.linspace(x0 - dx / 2.0, x1 + dx / 2.0, nx), np.linspace(y0 - dy / 2.0, y1 + dy / 2.0, ny)
mapx = np.tile(mapx, ny).reshape(ny, nx)
mapy = np.tile(mapy, nx).reshape(nx, ny).transpose()
return mapx, mapy
[docs]
def smapradialfilter(sunpymap, factor=5, grid=None):
from sunpy import map as smap
if grid:
x, y = grid
else:
x, y = smapmeshgrid2(sunpymap)
r = sunpymap.rsun_obs
rr = np.sqrt(x * x + y * y)
maskout = rr > r
try:
sunpymap.data[maskout] = sunpymap.data[maskout] * np.exp(factor * (rr[maskout] / r - 1))
except:
data = sunpymap.data.copy()
data[maskout] = data[maskout] * np.exp(factor * (rr[maskout] / r - 1))
sunpymap = smap.Map(data, sunpymap.meta)
return sunpymap
[docs]
def regridimage(values, x, y, grid=None, resize=[1.0, 1.0]):
'''
re-grid the data on a regular grid with uneven grid spacing to an uniform grid
:param values: The image data on the regular grid
:param x: the points defining the regular grid in x
:param y: the points defining the regular grid in y
:param grid: new uniform mesh grid [gridx,gridy]
:param resize: list of re-size ratio factors of x and y. if resize is not [1.0,1.0], grid is neglected.
:return: re-gridded image
'''
from scipy.interpolate import RegularGridInterpolator
ny, nx = values.shape
if grid and resize == [1.0, 1.0]:
gridx, gridy = grid
else:
gridx, gridy = np.meshgrid(np.linspace(x[0], x[-1], nx * resize[0]), np.linspace(y[0], y[-1], ny * resize[1]))
ny, nx = gridx.shape
rgi = RegularGridInterpolator(points=(y, x), values=values, bounds_error=False)
datanew = rgi(np.stack(np.stack((gridy.ravel(), gridx.ravel()), axis=-1))).reshape(ny, nx)
if grid:
return datanew
else:
return [datanew, gridx, gridy]
[docs]
def regridspec(spec, x, y, nxmax=None, nymax=None, interp=False):
'''
:param spec: ndarray of float or complex, shape (npol,nbl,nf,nt) Data values.
:param x: Data point x coordinates.
:param y: Data point y coordinates.
:param nxmax:
:param nymax:
:return:
'''
npol, nbl, nf, nt = spec.shape
if interp:
if nxmax:
if nt > nxmax:
nt = nxmax
if nymax:
if nf > nymax:
nf = nymax
specnew = np.zeros((npol, nbl, nf, nt))
tt = np.linspace(x[0], x[-1], nt)
ff = np.linspace(y[0], y[-1], nf)
grid_x, grid_y = np.meshgrid(tt, ff)
for p in range(npol):
for b in range(nbl):
specnew[p, b, :, :] = regridimage(spec[p, b, :, :], x, y, grid=[grid_x, grid_y])
else:
xstep, ystep = 1, 1
if nxmax:
if nt > nxmax:
import math
xstep = int(math.ceil(float(nt) / nxmax))
if nymax:
if nf > nymax:
ystep = int(float(nf) / nymax)
specnew = spec[:, :, ::ystep, ::xstep]
return [specnew, xstep, ystep]
[docs]
def get_contour_data(X, Y, Z, levels=[0.5, 0.7, 0.9]):
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from bokeh.models import (ColumnDataSource)
try:
cs = plt.contour(X, Y, Z, levels=(np.array(levels) * np.nanmax(Z)).tolist(), cmap=cm.Greys_r)
# dx = X[0,1]-X[0,0]
# dy = Y[1,0]-Y[0,0]
xs = []
ys = []
xt = []
yt = []
col = []
text = []
isolevelid = 0
for isolevel in cs.collections:
# isocol = isolevel.get_color()[0]
# thecol = 3 * [None]
theiso = '{:.0f}%'.format(cs.get_array()[isolevelid] / Z.max() * 100)
isolevelid += 1
# for i in range(3):
# thecol[i] = int(255 * isocol[i])
thecol = '#%02x%02x%02x' % (220, 220, 220)
# thecol = '#03FFF9'
for path in isolevel.get_paths():
v = path.vertices
# x = v[:, 0]+dx
# y = v[:, 1]+dy
x = v[:, 0]
y = v[:, 1]
xs.append(x.tolist())
ys.append(y.tolist())
xt.append(x[len(x) / 2])
yt.append(y[len(y) / 2])
text.append(theiso)
col.append(thecol)
for coll in cs.collections:
coll.remove()
source = ColumnDataSource(data={'xs': xs, 'ys': ys, 'line_color': col, 'xt': xt, 'yt': yt, 'text': text})
except:
source = ColumnDataSource(data={'xs': [], 'ys': [], 'line_color': [], 'xt': [], 'yt': [], 'text': []})
return source
# def twoD_Gaussian((x, y), amplitude, xo, yo, sigma_x, sigma_y, theta, offset):
# xo = float(xo)
# yo = float(yo)
# a = (np.cos(theta) ** 2) / (2 * sigma_x ** 2) + (np.sin(theta) ** 2) / (2 * sigma_y ** 2)
# b = -(np.sin(2 * theta)) / (4 * sigma_x ** 2) + (np.sin(2 * theta)) / (4 * sigma_y ** 2)
# c = (np.sin(theta) ** 2) / (2 * sigma_x ** 2) + (np.cos(theta) ** 2) / (2 * sigma_y ** 2)
# g = offset + amplitude * np.exp(- (a * ((x - xo) ** 2) + 2 * b * (x - xo) * (y - yo) + c * ((y - yo) ** 2)))
# return g.ravel()
[docs]
def c_correlate(a, v, returnx=False):
a = (a - np.mean(a)) / (np.std(a) * len(a))
v = (v - np.mean(v)) / np.std(v)
if returnx:
return np.arange(len(a)) - np.floor(len(a) / 2.0), np.correlate(a, v, mode='same')
else:
return np.correlate(a, v, mode='same')
[docs]
def c_correlateX(a, v, returnx=False, returnav=False, s=0):
'''
:param a:
:param v: a and v can be a dict in following format {'x':[],'y':[]}. The length of a and v can be different.
:param returnx:
:return:
'''
from scipy.interpolate import splev, splrep
import numpy.ma as ma
if isinstance(a, dict):
max_a = np.nanmax(a['x'])
min_a = np.nanmin(a['x'])
max_v = np.nanmax(v['x'])
min_v = np.nanmin(v['x'])
max_ = min(max_a, max_v)
min_ = max(min_a, min_v)
if not max_ > min_:
print('the x ranges of a and v have no overlap.')
return None
a_x = ma.masked_outside(a['x'].copy(), min_, max_)
if isinstance(a['y'], np.ma.core.MaskedArray):
a_y = ma.masked_array(a['y'].copy(), a['y'].mask | a_x.mask)
a_x = ma.masked_array(a_x, a_y.mask)
else:
a_y = ma.masked_array(a['y'].copy(), a_x.mask)
v_x = ma.masked_outside(v['x'].copy(), min_, max_)
if isinstance(v['y'], np.ma.core.MaskedArray):
v_y = ma.masked_array(v['y'].copy(), v['y'].mask | v_x.mask)
v_x = ma.masked_array(v_x, v_y.mask)
else:
v_y = ma.masked_array(v['y'].copy(), v_x.mask)
dx_a = np.abs(np.nanmean(np.diff(a_x)))
dx_v = np.abs(np.nanmean(np.diff(v_x)))
if dx_a >= dx_v:
v_ = v_y.compressed()
x_ = v_x.compressed()
tck = splrep(a_x.compressed(), a_y.compressed(), s=s)
ys = splev(x_, tck)
a_ = ys
elif dx_a < dx_v:
a_ = a_y.compressed()
x_ = a_x.compressed()
tck = splrep(v_x.compressed(), v_y.compressed(), s=s)
ys = splev(x_, tck)
v_ = ys
else:
a_ = a.copy()
v_ = v.copy()
x_ = None
a_ = (a_ - np.nanmean(a_)) / (np.nanstd(a_) * len(a_))
v_ = (v_ - np.nanmean(v_)) / np.nanstd(v_)
if returnx:
if x_ is None:
return np.arange(len(a_)) - np.floor(len(a_) / 2.0), np.correlate(a_, v_, mode='same')
else:
return (np.arange(len(a_)) - np.floor(len(a_) / 2.0)) * np.nanmean(np.diff(x_)), np.correlate(a_, v_,
mode='same'), x_, a_, v_
else:
return np.correlate(a_, v_, mode='same')
[docs]
def XCorrMap(z, x, y, doxscale=True):
'''
get the cross correlation map along y axis
:param z: data
:param x: x axis
:param y: y axis
:return:
'''
from tqdm import tqdm
from scipy.interpolate import splev, splrep
if doxscale:
xfit = np.linspace(x[0], x[-1], 10 * len(x) + 1)
zfit = np.zeros((len(y), len(xfit)))
for yidx1, yq in enumerate(y):
xx = x
yy = z[yidx1, :]
s = len(yy) # (len(yy) - np.sqrt(2 * len(yy)))*2
tck = splrep(xx, yy, s=s)
ys = splev(xfit, tck)
zfit[yidx1, :] = ys
else:
xfit = x
zfit = z
ny, nxfit = zfit.shape
ccpeak = np.empty((ny - 1, ny - 1))
ccpeak[:] = np.nan
ccmax = ccpeak.copy()
ya = ccpeak.copy()
yv = ccpeak.copy()
yidxa = ccpeak.copy()
yidxv = ccpeak.copy()
for idx1 in tqdm(range(1, ny)):
for idx2 in range(0, idx1):
lightcurve1 = zfit[idx1, :]
lightcurve2 = zfit[idx2, :]
ccval = c_correlate(lightcurve1, lightcurve2)
if sum(lightcurve1) != 0 and sum(lightcurve2) != 0:
cmax = np.amax(ccval)
cpeak = np.argmax(ccval) - (nxfit - 1) / 2
else:
cmax = 0
cpeak = 0
ccmax[idx2, idx1 - 1] = cmax
ccpeak[idx2, idx1 - 1] = cpeak
ya[idx2, idx1 - 1] = y[idx1 - 1]
yv[idx2, idx1 - 1] = y[idx2]
yidxa[idx2, idx1 - 1] = idx1 - 1
yidxv[idx2, idx1 - 1] = idx2
if idx1 - 1 != idx2:
ccmax[idx1 - 1, idx2] = cmax
ccpeak[idx1 - 1, idx2] = cpeak
ya[idx1 - 1, idx2] = y[idx2]
yv[idx1 - 1, idx2] = y[idx1 - 1]
yidxa[idx1 - 1, idx2] = idx2
yidxv[idx1 - 1, idx2] = idx1 - 1
return {'zfit': zfit, 'ccmax': ccmax, 'ccpeak': ccpeak, 'x': x, 'nx': len(x), 'xfit': xfit, 'nxfit': nxfit, 'y': y,
'ny': ny, 'yv': yv, 'ya': ya,
'yidxv': yidxv, 'yidxa': yidxa}
# class FileDialog(): # ''' # produce a file dialog button widget for bokeh plot # ''' # import Tkinter # import tkFileDialog # # def __init__(self, plot_width=30, labels={'dir':'...','open':'open','save':'save'}, *args, # **kwargs): # from bokeh.models import Button # buttons = {} # for k,v in labels.items(): # buttons[k] = Button(label=v, width=plot_width) # self.buttons = buttons # # def askdirectory(self): # tkRoot = Tkinter.Tk() # tkRoot.withdraw() # Close the root window # in_path = tkFileDialog.askdirectory() # tkRoot.destroy() # if in_path: # return in_path # # def askopenfilename(self): # tkRoot = Tkinter.Tk() # tkRoot.withdraw() # Close the root window # in_path = tkFileDialog.askopenfilename() # tkRoot.destroy() # if in_path: # return in_path # # def asksaveasfilename(self): # tkRoot = Tkinter.Tk() # tkRoot.withdraw() # Close the root window # in_path = tkFileDialog.asksaveasfilename() # tkRoot.destroy() # if in_path: # return in_path