<jupyter_start><jupyter_text>Computes the burst drift of an FRB using a 2d autocorrelationauthor: moh, based on victor's ectrture_matrice_intensite.py dat: march 28, 2020<jupyter_code>#!/usr/bin/python3

from __future__ import division
import math
import os
import sys
import time
import numpy as np
import scipy.stats
from scipy.optimize import curve_fit
from math import log10
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from scipy import pi as nombrepi
from scipy import signal, ndimage
from tqdm import tqdm
from matplotlib import colors as mcolors
import functools
print = functools.partial(print, flush=True) # print doesn't happen til script ends so force it to flush... windows thing?

burst = 16

junk, nchan, nbin, I, Q, U, V = np.loadtxt('data/{}_puppi_57772_C0531+33_0007_2695.dm559.72.calibP.RM.DD.ASCII'.format(burst), delimiter=' ', unpack=True)
n = len(junk)
print("Data loaded")
binmax = int(nbin[n-1])+1
frequencymax = (int(nchan[n-1])+1)
print("dimensions:", binmax, frequencymax)
intensity = np.zeros((frequencymax, binmax))

X = np.zeros(binmax)
Y = np.zeros(frequencymax)

tmin = 500

tmax = 1500

intensitynoise1 = np.zeros(tmin-1)
intensitynoise2 = np.zeros(binmax-tmax)
for i in tqdm(range(frequencymax-50,51,-1), desc='noise removal'):

	Y[i-1] = 4.15 + (i-1) * 1.5625 # ?

	for j in range(1,tmin) :

		intensitynoise1[j-1] = (I[j-1 + binmax*(frequencymax-i)])/(tmin-1)


	for j in range(tmax+1,binmax+1) :

		intensitynoise2[j-1-tmax] = (I[j-1 + binmax*(frequencymax-i)])/(binmax-tmax)

	a = sum(intensitynoise1)
	b = sum(intensitynoise2)

	for j in range(1,binmax+1) :
		X[j-1] = j-1
		intensity[i-1,j-1] = I[j-1 + binmax*(frequencymax-i)] - (a+b)/2

for i in tqdm(range(52,frequencymax-49), desc='variance'):
	for j in range(tmax+1,binmax+1) :
		intensitynoise2[j-1-tmax] = intensity[i-1,j-1]

	mean = np.mean(intensitynoise2)
	variance = np.var(intensitynoise2)
# del intensitynoise1
# del intensitynoise2
# del junk, nchan, nbin, I, Q, U, V
print(len(intensity), len(intensity[0]))
edge = 700
burstwindow = intensity[:,edge:edge+frequencymax]
cmap = plt.get_cmap('gray')
cmap.set_bad(color = 'b', alpha=1) # potentially hides important errors in the data!
plt.figure(figsize=(10,8))
plt.title("Burst #{}".format(burst))
plt.imshow(burstwindow, cmap=cmap, interpolation='bicubic',aspect='equal', origin="lower")
print(np.min(burstwindow), np.max(burstwindow))
print(burstwindow)
# plt.imshow(burstwindow, cmap=cmap, interpolation='bicubic',aspect='equal', origin="lower", norm=mcolors.LogNorm(vmin=0.05, vmax=np.max(burstwindow)))
plt.colorbar()
print("finding auto-correlation...")
corr = signal.correlate2d(burstwindow, burstwindow, mode='full')
print('omg finally finished')
cmap = "gray"
plt.figure(figsize=(15,6))
plt.subplot(121)
plt.title("Burst #{}".format(burst))
plt.imshow(burstwindow, cmap=cmap, interpolation='bicubic',aspect='equal', origin="lower")
plt.colorbar()
plt.subplot(122)
plt.imshow(corr, cmap=cmap, interpolation='bicubic',aspect='auto', origin="lower")
plt.title("Correlation #{}".format(burst))
plt.colorbar()
plt.tight_layout()
print(burstwindow.shape)<jupyter_output>(512, 512)<jupyter_text>Gaussian 2d Fit<jupyter_code># Source: https://gist.github.com/andrewgiessel/6122739
# Source: https://stackoverflow.com/questions/21566379/fitting-a-2d-gaussian-function-using-scipy-optimize-curve-fit-valueerror-and-m
def gaussian(height, center_x, center_y, width_x, width_y, rotation):
    """Returns a gaussian function with the given parameters"""
    width_x = float(width_x)
    width_y = float(width_y)

    rotation = np.deg2rad(rotation)
    center_x_old = center_x
    center_x = center_x * np.cos(rotation) - center_y * np.sin(rotation)
    center_y = center_x_old * np.sin(rotation) + center_y * np.cos(rotation)

    def rotgauss(x,y):
        xp = x * np.cos(rotation) - y * np.sin(rotation)
        yp = x * np.sin(rotation) + y * np.cos(rotation)
        g = height*np.exp( -(((center_x-xp)/width_x)**2  +((center_y-yp)/width_y)**2)/2.  )
        return g
    return rotgauss

def moments(data):
    """Returns (height, x, y, width_x, width_y)
    the gaussian parameters of a 2D distribution by calculating its
    moments """
    total = data.sum()
    X, Y = np.indices(data.shape)
    x = (X*data).sum()/total
    y = (Y*data).sum()/total
    col = data[:, int(y)]
    width_x = np.sqrt(abs((np.arange(col.size)-y)**2*col).sum()/col.sum())
    row = data[int(x), :]
    width_y = np.sqrt(abs((np.arange(row.size)-x)**2*row).sum()/row.sum())
    height = data.max()
    return height, x, y, width_x, width_y, 2.0

def twoD_Gaussian(point, amplitude, xo, yo, sigma_x, sigma_y, theta):
    x, y = point
    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 = amplitude*np.exp( - (a*((x-xo)**2) + 2*b*(x-xo)*(y-yo) + c*((y-yo)**2)))
    return g.ravel()

def fitgaussian(data):
    """Returns (height, x, y, width_x, width_y)
    the gaussian parameters of a 2D distribution found by a fit"""
    params = moments(data)
    errorfunction = lambda p: np.ravel(gaussian(*p)(*np.indices(data.shape)) - data)
    p, success = scipy.optimize.leastsq(errorfunction, params)
    return p, success

def fitgaussiannlsq(data):
    # use curve-fit (non-linear leastsq)
    x = range(0, 1023); y = range(0, 1023)
    x, y = np.meshgrid(x, y)
    params = moments(data)#+ (0.,)
    popt, pcov = scipy.optimize.curve_fit(twoD_Gaussian, (x, y), corr.ravel(), p0=params)
    return popt, pcov

p, success = fitgaussian(corr)
popt, pcov = fitgaussiannlsq(corr)
print('success? (nonzero):', success)
fit = gaussian(*p)
x = range(0, 1023); y = range(0, 1023)
x, y = np.meshgrid(x, y)
fitmap2 = twoD_Gaussian((x, y), *popt).reshape(1023, 1023)
# fit = gaussian(12309.30062645451, 511.0000000000002, 511.0000000000002, 126.10568482796336, 55.80036552617748, 2)
print('solution:', p)
print('solution nlsq:', popt)
print('initial guess:', moments(corr))
fitmap = np.zeros((len(corr), len(corr[0])))
for y in tqdm(range(0, len(corr))):
    for x in range(0, len(corr[0])):
        fitmap[x][y] = fit(x, y)
plt.figure(figsize=(16,7))
plt.subplot(121)
plt.title('Fit 1: Least Squares')
plt.imshow(fitmap, cmap=cmap, interpolation='bicubic',aspect='equal', origin="lower")
plt.colorbar()
plt.subplot(122)
plt.title('Fit 2: Non-linear Least Squares')
plt.imshow(fitmap2, cmap=cmap, interpolation='bicubic',aspect='equal', origin="lower")
plt.colorbar()
plt.tight_layout()
cmap = "gray"
plt.figure(figsize=(17,8))
plt.subplot(121)
plt.title("Burst #{}".format(burst))
plt.imshow(burstwindow, cmap=cmap, interpolation='bicubic',aspect='auto', origin="lower")
plt.colorbar()

plt.subplot(122)
plt.title("Correlation #{} with fit, tilt = {:.2f} rad".format(burst, popt[-1]))
plt.imshow(corr, cmap=cmap, interpolation='bicubic',aspect='auto', origin="lower")
plt.colorbar()
plt.contour(*np.meshgrid(range(0, 1023), range(0, 1023)), fitmap, [600, 2000], colors='w', alpha=0.5)
plt.contour(*np.meshgrid(range(0, 1023), range(0, 1023)), fitmap2, [600, 2000], colors='r', alpha=0.5)

plt.tight_layout()<jupyter_output><empty_output><jupyter_text>Correlate2d Tests<jupyter_code>dirac = np.zeros((frequencymax, frequencymax))
dirac[256][256] = 1

diraccorr = signal.correlate2d(dirac, dirac, mode='full')
print('done')
plt.figure(figsize=(18,8))

plt.subplot(121)
plt.title('Dirac pulse')
plt.imshow(dirac[128:384,128:384], cmap=cmap, interpolation='bicubic', aspect='equal', origin="lower")
plt.colorbar()

plt.subplot(122)
plt.title('Autocorrelation of Dirac input')
plt.imshow(diraccorr[384:640,384:640], cmap=cmap, interpolation='bicubic', aspect='equal', origin="lower")
plt.colorbar()
plt.tight_layout()
rect = np.zeros((frequencymax, frequencymax))
dims = (512/2-50, 512/2+50)
for y in range(0, len(rect)):
    for x in range(0, len(rect[0])):
        if y > dims[0] and y < dims[1] and x > dims[0] and x < dims[1]:
            rect[y][x] = 1
    
rectcorr = signal.correlate2d(rect, rect, mode='full')
print('done')
plt.figure(figsize=(16,7))
plt.subplot(121)
plt.title('rectangle')
plt.imshow(rect, cmap=cmap, origin='lower')
plt.colorbar()
plt.subplot(122)
plt.title('autocorrelation of rectangle')
plt.imshow(rectcorr, cmap=cmap, origin='lower')
plt.colorbar()
plt.tight_layout()
gaussinput = np.zeros((frequencymax, frequencymax))
for y in range(0, len(gaussinput)):
    for x in range(0, len(gaussinput[0])):
        gaussinput[y][x] = (1/(50*np.sqrt(np.pi))) * np.exp(-((y-256)/50)**2-((x-256)/50)**2)
    
guasscorr = signal.correlate2d(gaussinput, gaussinput, mode='full')
print('done')

def center(matrix ,w):
    lx, ly = matrix.shape
    print(lx, ly, w)
    return matrix[int(lx/2-w/2):int(lx/2+w/2) , int(ly/2-w/2):int(ly/2+w/2)]
    
plt.figure(figsize=(16,7))
plt.subplot(121)
plt.title('gaussian')
plt.imshow(gaussinput, cmap=cmap, origin='lower')
plt.colorbar()

plt.subplot(122)
plt.title('gaussian autocorrelation')
plt.imshow(center(guasscorr, 512), cmap=cmap, origin='lower')
plt.colorbar()

plt.tight_layout()
# check the peak of the autocorrelation
print(np.max(guasscorr))
print(np.max(gaussinput)**2)
print(1/(np.pi* (np.sqrt(2)*50)**2))

peak = 0 
for y in range(0, len(gaussinput)):
    for x in range(0, len(gaussinput[0])):
        peak += gaussinput[x][y]**2
print(peak)<jupyter_output>0.49999999999995715
0.00012732395447351627
6.366197723675813e-05
0.49999999999995715
{'fraction_non_alphanumeric': 0.1813341684935797, 'fraction_numerical': 0.05741726693809375, 'mean_word_length': 5.265533028122956, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 4, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '268721', 'license': 'MIT', 'n_tokens_mistral': 4268, 'n_tokens_neox': 3770, 'n_words': 743, 'path': 'BurstDrift.ipynb', 'repo_name': 'mef51/subdriftlaw', 'size': '9579'}