#! /usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

def colon(start,step,stop):
    # equivalent to matlab's start:step:stop
    from numpy import append
    from numpy import arange
    return append(arange(start,stop,step),stop)

sansfile='Sans_0mT_29p3K.dat'
#sansfile='Sans_104mT_29p3K.dat'
#sansfile='Sans_240mT_29p3K.dat'
#sansfile='Sans_350mT_29p3K.dat'
#sansfile='Sans_384mT_29p3K.dat'

sansmatrix=np.loadtxt(sansfile)

#%SANS-1
#%\begin{itemize}
# %\item $\lambda=\SI{5.5}{\angstrom}$ $\Delta\lambda/\lambda=0.1$
# %\item sample to detector distance \SI{10.2}{\meter}
# %\item collimation length \SI{12}{\meter} with \SI{20}{\milli\meter} incoming pin-hole and \SI{4}{\milli\meter} in front of the sample \pink{SD 10\,m, $B_{sx}$ = 496\,mm, $B_{sy}$ = 498\,mm}
#% \item closed cycle cryostate with superconducting \SI{5}{\tesla} magnet, field applied horizontally perpendicular to neutron beam
#% \item sample OFZ-59: spherical shaped \SI{6}{\milli\meter} diameter, $\langle 110\rangle$ axis vertical and $\langle 111\rangle$ horizontally perpendicular to neutron beam
#%\end{itemize}

nlambda = 5.5 # A wavelength of neutron beam
kin = 2*np.pi/nlambda # A-1 (incoming k is about 1.14 A-1)

SD=10.2 # distance between sample and detector in metres (Jonas PhD thesis mentions 1.38m, 2.6m, 2.525m, not obvious which one is relevant for SANS; a comment in the sharelatex says 10.2 m)
SD_offset=0.0
Beam_offset_x=0.0
Beam_offset_y=0.0

# calculate sd in mm 
SD=SD*1000+SD_offset

# produce a grid corresponding to the detector size, position in mm, middle of each pixel 
Pos_x=colon(-508,2*512/128,508)
Pos_y=colon(-508,2*512/128,508) # this is the other way around compared to the matlab
Pos_x_grid, Pos_y_grid=np.meshgrid(Pos_x,Pos_y)

# calculate real beam position for each point on the detector with respect to the direct beam 
Beam_pos_x=Pos_x_grid-Beam_offset_x
Beam_pos_y=Pos_y_grid-Beam_offset_y

# calculate cartesian components of scattering vector Q, z is along the beam, 
# x is horizontal and y vertical 
# Here we substract a normalized vector Kf from normalized vector Ki and multiply by kin = 2*pi/wav 
Q_x_grid= Beam_pos_x/(np.sqrt(Beam_pos_x**2 + Beam_pos_y**2 +SD**2)) * kin
Q_y_grid= Beam_pos_y/(np.sqrt(Beam_pos_x**2 + Beam_pos_y**2 +SD**2)) * kin
Q_z_grid= (SD/(np.sqrt(Beam_pos_x**2 + Beam_pos_y**2 +SD**2))-1) * kin 
Q_mod_grid=np.sqrt(Q_x_grid**2 + Q_y_grid**2 + Q_z_grid**2)

#plt.figure(sansfile)
#plt.imshow(sansmatrix, origin='lower')
#plt.tight_layout()
#plt.show()

fig, ax = plt.subplots()
# cs = ax.pcolormesh(1e3*Q_x_grid,1e3*Q_y_grid,sansmatrix*1e5, shading='auto', cmap='bwr') # linear intensity scale
cs = ax.pcolormesh(Q_y_grid,Q_x_grid,sansmatrix*1e5, shading='auto', cmap='bwr') # linear intensity scale
# ax.pcolormesh(1e3*Q_x_grid,1e3*Q_y_grid,np.log(sansmatrix), shading='auto', cmap='bwr') # log intensity scale
# cmap='jet' for a matlabby look, 'seismic', 'bwr', or 'coolwarm' for blue-white-red
cbautomin,cbautomax=(cs.get_clim())
cs.set_clim(0.0,cbautomax*1.15)
ax.set_aspect('equal')
#plt.xlabel('$q_x$ [ $10^{-3}$ Å$^{-1}$ ]')
#plt.ylabel('$q_y$ [ $10^{-3}$ Å$^{-1}$ ]')
plt.xlabel('$q_x$ [ Å$^{-1}$ ]')
plt.ylabel('$q_y$ [ Å$^{-1}$ ]')
divider = make_axes_locatable(ax)
cax = divider.new_vertical(size='5%', pad=0.05)
fig.add_axes(cax)
cb = fig.colorbar(cs, cax=cax, orientation='horizontal')
cb.ax.xaxis.set_ticks_position('top')
cb.ax.xaxis.set_label_position('top')
plt.tight_layout()
plt.show()

# In[98]:


#!/usr/bin/env python3
import numpy as np
from scipy.ndimage import map_coordinates
import matplotlib.pyplot as plt

xc, yc = 65, 63

x, y = np.mgrid[-0.06:0.06:0.0009375, -0.06:0.06:0.0009375]
z = sansmatrix

#plt.imshow(sansmatrix)

#zj = np.zeros(360)
#x = np.zeros(360)
#y = np.zeros(360)

ri=32.0 # inner radius
ro=60.0 # outer radius
rpoints=80 # radius points

plt.figure()
plt.xlim(0,127)
plt.ylim(0,127)

for i in range(360):
	phi = i*2*np.pi/360.0
	xi = xc+ri*np.cos(phi)
	yi = yc+ri*np.sin(phi)
	xo = xc+ro*np.cos(phi)
	yo = yc+ro*np.sin(phi)
	xlin,ylin = np.linspace(xi, xo, rpoints), np.linspace(yi, yo, rpoints)
	plt.imshow(z)
	plt.plot(xlin,ylin)

plt.show()

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

def Xinv(kpara,kperp,Q,J,acub,Xint,k,h):
    # k is the scattering vector "q"
    # k1 is [100]
    # k2 is [010]
    # k3 is [001]
	klen=np.linalg.norm(k)
	LC3=np.array([[[0,0,0],[0,0,1],[0,-1,0]],
			     [[0,0,-1],[0,0,0],[1,0,0]],
			     [[0,1,0],[-1,0,0],[0,0,0]]])
	LC=np.inner(LC3,k)
	Dperp=J*kperp**2
	Dpara=J*kpara**2
	Xinv=np.zeros([3,3], dtype=complex)
	Xinv=Xinv+(J*(Q**2+klen**2))*np.eye(3)+Xint*J*(Q/klen)**2*np.outer(k,k)+J*acub*np.diag(k**2)-2*1.j*J*Q*LC + Dperp*np.eye(3)+(Dpara-Dperp)*np.outer(h,h)
	return Xinv

def gauss(q,intens,qxc,qyc,qxw,qyw,angle):
	# q is the scattering vector at this pixel
	sigmax=qxw/2.355
	sigmay=qyw/2.355
	return 1e-8*intens*1/(2*np.pi*sigmax*sigmay)*np.exp(-0.5*((q[0]-qxc)/sigmax)**2)*np.exp(-0.5*((q[1]-qyc)/sigmay)**2)

# [10/5, 08:49] Basti Mühlbauer: Sollte ne 11-1 sein
# [10/5, 08:50] Basti Mühlbauer: Und ne 1-10

#%SANS-1
#%\begin{itemize}
# %\item $\lambda=\SI{5.5}{\angstrom}$ $\Delta\lambda/\lambda=0.1$
# %\item sample to detector distance \SI{10.2}{\meter}
# %\item collimation length \SI{12}{\meter} with \SI{20}{\milli\meter} incoming pin-hole and \SI{4}{\milli\meter} in front of the sample \pink{SD 10\,m, $B_{sx}$ = 496\,mm, $B_{sy}$ = 498\,mm}
#% \item closed cycle cryostate with superconducting \SI{5}{\tesla} magnet, field applied horizontally perpendicular to neutron beam
#% \item sample OFZ-59: spherical shaped \SI{6}{\milli\meter} diameter, $\langle 110\rangle$ axis vertical and $\langle 111\rangle$ horizontally perpendicular to neutron beam
#%\end{itemize}

hdir = np.array([1,1,-1]) # direction of magnetic field vector
h = hdir/np.linalg.norm(hdir) # unit vector

Labframe=True
# experimental geometry gives the qx and qy directions
if Labframe==True:
	# Lab frame, peaks will be left-right
	xdir = -hdir
	ydir = np.array([1,-1,0])

else:
	# MnFeSi_Tricritical_point.pdf, peaks will be up-down
	xdir = np.array([1,-1,0])
	ydir = hdir

zdir = np.cross(xdir,ydir)
xhat = xdir/np.linalg.norm(xdir) # unit vector
yhat = ydir/np.linalg.norm(ydir) # unit vector
zhat = np.cross(xhat,yhat) # unit vector

# components of q1, q2, q3 are in the crystal coordinates
q1dir= np.array([1,0,0])
q2dir= np.array([0,1,0])
q3dir= np.array([0,0,1])

Q = 0.039 # A-1
J = 2.8 # meV/A

acub = -0.041 # A-1; page 97 PhD Thesis Jonas
#acub = -0.023 # A-1; page 97 PhD Thesis Jonas
#acub = 0.0 # A-1; for testing - gives a ring of radius Q
Xint = 0.1

I0 = 3.2e-9 # intensity factor

# kperp = kpara = kappa = 0.00634 A-1 at zero field, 29.3K
kperp=0.00634; kpara=kperp

# 104mT, 29.3K
#kperp = 0.00807; kpara = 0.00956

# 240mT, 29.3K
#kperp=0.0076; kpara=0.02031

# 350mT, 29.3K
#kperp = 0.01375; kpara = 0.04368

# 384mT, 29.3K
#kperp = 0.0152; kpara = 0.06103


nlambda = 5.5 # A wavelength of neutron beam
kin = 2*np.pi/nlambda # A-1 (incoming k is about 1.14 A-1)

# detector is 128x128 pixels not 512x512 pixels
# qx and qy range is +/- 50e-3 A-1
xpix=128
ypix=xpix

qxrange=57e-3 # A-1
qyrange=qxrange

qxvals=np.linspace(-qxrange,qxrange,xpix)
qyvals=np.linspace(-qyrange,qyrange,ypix)

detector=np.zeros((xpix,ypix)) # to store detector "image"
qcoord=np.zeros((2,xpix,ypix)) # so we know how each pixel relates to qx and qy
qcoord[0]=qxvals
qcoord[1]=qyvals[:, np.newaxis]

# nested loop over all the pixels
for i in range(xpix):
    for j in range(ypix):
        qxpos=qcoord[0,i,j]
        qypos=qcoord[1,i,j]
        qzpos=np.sqrt(kin**2-qxpos**2-qypos**2)-kin
        qpos=np.array([qxpos,qypos,qzpos])
        # this pixel is detector[i,j]
        # xhat, yhat, zhat are unit vectors in the detector coordinates
        # q = [q1,q2,3] is the scattering vector in the crystal coordinates
        q = qxpos*xhat + qypos*yhat + qzpos*zhat
        q1 = np.dot(q[0],q1dir)
        q2 = np.dot(q[1],q2dir)
        q3 = np.dot(q[2],q3dir)

        Xinvk = Xinv(kpara,kperp,Q,J,acub,Xint,q,h)
        invXinv = np.linalg.inv(Xinvk)
        Sigma = I0*0.5*(np.trace(invXinv)-1/(np.linalg.norm(q)**2)*(np.inner(q,(np.inner(invXinv,q)))))
        #print("Xinv at q = {} A-1 gives Sigma = {}".format(q,Sigma))

        # Gaussian peaks:
        peakintens=0.021
        qxabs=0.04 # A-1
        qyabs=0.00 # A-1
        qxwidth=0.0035 # A-1
        qywidth=0.005 # A-1

        if Labframe==True: # peaks will be left-right
            peak1=gauss(qpos,peakintens,qxabs,qyabs,qxwidth,qywidth,0.0)
            peak2=gauss(qpos,peakintens,-qxabs,-qyabs,qxwidth,qywidth,0.0)

        else: # peaks will be up-down
            peak1=gauss(qpos,peakintens,qyabs,qxabs,qywidth,qxwidth,0.0)
            peak2=gauss(qpos,peakintens,-qyabs,-qxabs,qywidth,qxwidth,0.0)

        detector[i,j]=np.real(Sigma)+peak1+peak2

#plt.matshow(np.transpose(detector))

#fig, ax = plt.subplots()
#im = ax.imshow(np.transpose(detector), origin='lower')
#plt.show()

mA = False# set to True to plot in 10-3 Å

fig, ax = plt.subplots()
if mA:
	cs = ax.pcolormesh(1e3*qxvals,1e3*qyvals,(detector)*1e4, shading='auto', cmap='bwr') # linear intensity scale
	# cs = ax.pcolormesh(1e3*qxvals,1e3*qyvals,np.transpose(detector)*1e5, shading='auto', cmap='bwr') # linear intensity scale

else:
	cs = ax.pcolormesh(qxvals,qyvals,(detector)*1e5, shading='auto', cmap='bwr') # linear intensity scale
	# cs = ax.pcolormesh(qxvals,qyvals,np.transpose(detector)*1e5, shading='auto', cmap='bwr') # linear intensity scale

# cmap='jet' for a matlabby look, 'seismic', 'bwr', or 'coolwarm' for blue-white-red
cbautomin,cbautomax=(cs.get_clim())
cs.set_clim(0.0,cbautomax*1.15)
ax.set_aspect('equal')
if mA:
	plt.xlabel('$q_x$ [ $10^{-3}$ Å$^{-1}$ ]')
	plt.ylabel('$q_y$ [ $10^{-3}$ Å$^{-1}$ ]')

else:
	plt.xlabel('$q_x$ [ Å$^{-1}$ ]')
	plt.ylabel('$q_y$ [ Å$^{-1}$ ]')

divider = make_axes_locatable(ax)
cax = divider.new_vertical(size='5%', pad=0.05)
fig.add_axes(cax)
cb = fig.colorbar(cs, cax=cax, orientation='horizontal')
cb.ax.xaxis.set_ticks_position('top')
cb.ax.xaxis.set_label_position('top')
plt.tight_layout()
plt.show()

import scipy.ndimage

xcalc =1e3*qxvals
ycalc =1e3*qyvals
zcalc =(detector)*1e4


#plt.imshow(sansmatrix)

#zj = np.zeros(360)
#x = np.zeros(360)
#y = np.zeros(360)

ri=32.0 # inner radius
ro=60.0 # outer radius
rpoints=80 # radius points

plt.figure()
#plt.xlim(0,127)
#plt.ylim(0.0007,0.0008)
sumup2 = []
angle2 = np.linspace(0, 360, 360)

for i in range(360):
	phi = i*2*np.pi/360.0
	xi = xc+ri*np.cos(phi)
	yi = yc+ri*np.sin(phi)
	xo = xc+ro*np.cos(phi)
	yo = yc+ro*np.sin(phi)
	xlin,ylin = np.linspace(xi, xo, rpoints), np.linspace(yi, yo, rpoints)
	zi = scipy.ndimage.map_coordinates(zcalc, np.vstack((xlin,ylin)))
	sumup2.append(sum(zi))
plt.plot(angle2, sumup2)    
plt.show()

xcalc =1e3*qxvals
ycalc =1e3*qyvals
zcalc =(detector)


#plt.imshow(sansmatrix)

#zj = np.zeros(360)
#x = np.zeros(360)
#y = np.zeros(360)

ri=32.0 # inner radius
ro=60.0 # outer radius
rpoints=80 # radius points



xc, yc = 65, 63

x, y = np.mgrid[-0.06:0.06:0.0009375, -0.06:0.06:0.0009375]
z = sansmatrix


plt.figure()
#plt.xlim(0,127)
#plt.ylim(0.0007,0.0008)
sumup = []
angle2 = np.linspace(0, 360, 360)

plt.figure()
#plt.xlim(0,127)
#plt.ylim(0.0007,0.0008)
sumup2 = []
angle = np.linspace(90, 450, 360)

for i in range(360):
	phi = i*2*np.pi/360.0
	xi = xc+ri*np.cos(phi)
	yi = yc+ri*np.sin(phi)
	xo = xc+ro*np.cos(phi)
	yo = yc+ro*np.sin(phi)
	xlin,ylin = np.linspace(xi, xo, rpoints), np.linspace(yi, yo, rpoints)
	zi2 = scipy.ndimage.map_coordinates(zcalc, np.vstack((xlin,ylin)))
	zi = scipy.ndimage.map_coordinates(z, np.vstack((xlin,ylin)))
	sumup.append(sum(zi))
	sumup2.append(sum(zi2))
plt.plot(angle, sumup)
plt.plot(angle2, sumup2)    
plt.show()