# -*- coding: utf-8 -*-
"""
Laue_Spot_Patterns.py

Plots reflection Laue image spot patterns for arbitrary orientionations of
    some simple cubic crystals.

This program is far from fully developed, but it is useful for simple Laue
    undergraduate experiments.  Since it is straight Python without a
    GUI (Graphical User Interface), and only needs common python libraries
    such as numpy, matplotlib, math, pprint, and time,
    it should run easily on most Python Installations.

To use this program:
    First make sure it works. Running it should produce Laue Spot pattern
        determined by the default parameters. The default parameters are
        chosen to match the 001 Tungsten example shown at
        http://www.physics.utoronto.ca/~phy326/laue
        (i.e. file Example_W001_45kV_18mA_30mm.jpg).
    To run it with your desired parameters, you should edit:
        - Crystal_axis :        orientation of crystal relative to the 
                                    beam (100) axis.
        - Z_axis :              this choice rotates the image around the
                                    beam axis
        - d_film :              distance (mm) between crystal face and Laue
                                    film plane
        - a_c :                 cubic lattic constant (Angstroms)
        - beam_yz :             horizontal (y) and vertical (z) location of
                                    beam spot
        - min_y, max_y,
          min_z, max_z :        horizontal (y) and vertical (z) edges of
                                    useful area of film
        - E_beam :              X-ray machine voltage (maximum X-ray energy)
        - maximum_plane_index : maximum index used to generate planes
        - spot_scale          : scale factor to adjust size of plotted spots
        - intensity_threshold : intensity threshold required for a spot to be
                                    labelled with its hkl indices
        - hkl_threshold :       spots with hkl values with h**2+k**2+l**2 less
                                    than this value will be labeled
                                    irrespective of their intensity        
        - saved_image_file    : name of saved image file
    Running this script will produce a Laue spot pattern onscreen,
            and save it in the image file.


The intensity of the spots is based on Preuss, but should not be trusted
    very much. For example, the X-ray intensity spectrum use in Preuss is
    likely not very accurate. Feel free to correct, modify and improve
    as desired. 

Cited references:
    Preuss: "Laue Atlas" by E. Preuss, B. Krahl-Urban, R. Butz, Wiley 1974
    Toronto : Laue Back-Reflection of X-Rays experiment manual for
        University of Toronto undergraduate Advanced Physics Lab, see
        http://www.physics.utoronto.ca/~phy326/laue
 
Please let me know if you create a modified version.
     For example, modifying for tetragonal or orthorhombic crystal structures
     requires care but should be straightforward. Changing the input to read
     from a file is also easy to do.
This code has not between particularly optimized, either for speed or length.
    It could be made more recursive, but possibly at the cost of speed.

Copyright (c) 2016 University of Toronto
Original Version:   May 2016 by David Bailey
Contact: David Bailey <dbailey@physics.utoronto.ca>
            (www.physics.utoronto.ca/~dbailey)
License: Released under the MIT License; the full terms are this license
            are appended to the end of this file, and are also available
            at http://www.opensource.org/licenses/mit-license.php.
"""

# Use python 3 consistent printing and unicode
from __future__ import print_function
from __future__ import unicode_literals

from time import time
from numpy import array, matrix, cross
from numpy.linalg import norm
from matplotlib.pyplot import show, scatter, text, figure, savefig
from math import sqrt, cos, sin, pi
from pprint import pprint

start_time = time()

# Default values below roughly match those needed to reproduce pattern of
#            Example_W001_45kV_18mA_30mm

# Define Crystal and Z axes. (The beam axis is [1,0,0].)
Crystal_axis = [3.147, 0.172, 0.202]
Z_axis       = [-0.262, 1.654, 2.678] 

# Distance between crystal face and Laue film plane. It is assumed that the
#    film is perpendicular to the beam.
d_film = 29.5     # mm

# Cubic crystal lattice constant
a_c = 3.1585     # Angstroms

# Cubic crystal space group 
#     (simple cubic 221, fcc 225, fcc diamond 227, bcc 229)
space_group = 229

# Beam spot location
beam_yz =[-10, 5]

# Define active size of filem
min_y, max_y, min_z, max_z = -100, 100, -100, 100  # mm

# Electron beam energy (i.e. maximum x-ray energy)
E_beam = 45000     # eV

# Maximum index of planes that will be generated
maximum_plane_index = 20

# Spot size scale factor
spot_scale = 2

# How intense should spots be in order to be labelled with hkl values on plot?
intensity_threshold = 12
# hkl values with h**2+k**2+l**2 less than this value will be labeled
hkl_threshold = 25

# Name of file with saved image
#  supported formats include png, jpg, pdf, ps, svg, tif
saved_image_file = 'LauePattern.png'

# Calculate the minimum X-ray wavelength from the electron energy
min_wave = 12398.419/E_beam    # Angstroms

# Calculate maximum h**2+k**2+l**2 value that can produce a spot given the 
#    beam energy (See Toronto experiment manual.)
hkl2_max = (2*a_c/min_wave)**2

print("\nInput values")
print("  Crystal axis :        ", Crystal_axis)
print("  Z axis :              ", Z_axis)
print("  Crystal Space Group : ", space_group)
print("  Lattice constant      ", a_c, "Angstroms")
print("  Distance to film is   ", d_film, "mm")
print("  Beam spot location is ", beam_yz, "mm")
print("  Film edges :          ", min_y, max_y, min_z, max_z, "mm")
print("  X-ray tube voltage :  ", E_beam, "eV")
print("  Maximum plane index : ", maximum_plane_index)
print("  Spot size scale :     ", spot_scale)
print("  Thresholds for labelling a spot with its hkl indices")
print("      Intensity :       ", intensity_threshold)
print("      h^2+k^2+l^2 :     ", hkl_threshold)
print("  Imaged saved in :     ", saved_image_file)

print("\nCalculated values")
print("  Minimum wavelength :  ", min_wave, "Angstroms")
print("  Maximum possible h**2+k**2+l**2 :",hkl2_max)

beam_yz=array(beam_yz)
# Figure out rotation matrix to orient spot pattern
X_axis    = array(Crystal_axis)/norm(Crystal_axis)
Z_axis    = array(Z_axis)
Y_axis    = cross(Z_axis,X_axis)
Y_axis    = array(Y_axis)/norm(Y_axis)
Z_axis    = cross(X_axis,Y_axis)
Z_axis    = array(Z_axis)/norm(Z_axis)
R = matrix([X_axis,Y_axis,Z_axis])
print("  Rotation matrix : ")
pprint(R, indent=5)

# Renormalize intensity threshold for labelling spots with their hkl values
intensity_threshold *= spot_scale

# Bunch of utility functions to switch between hkl, angles, spots

def plane_angles_from_hkl(hkl) :
    from math import atan2,sqrt
    h,k,l = hkl
    theta = atan2(sqrt(k**2+l**2),h)
    phi   = atan2(float(l),k)
    return (theta,phi)

def spot_yz_from_plane_angles(theta_phi) :
    from math import tan,cos,sin,pi
    theta,phi = theta_phi
    if theta > pi/2:
        return None, None
    y = d_film*tan(2*theta)*cos(phi)
    z = d_film*tan(2*theta)*sin(phi)
    return y,z

def spot_yz_from_hkl(hkl) :
    # Returns spot location for a plane
    return spot_yz_from_plane_angles(plane_angles_from_hkl(hkl))

# Calculate all possible independent hkl combinations
hkl_all = []       # Storage for hkl values
wavelengths = []   # Storage for wavelength picked out by that hkl
# Loop over hkl starting with 1 and working up, then 0 and down
hkl_range = ( list(range(1,maximum_plane_index+1))
              + list(range(0,-maximum_plane_index-1,-1)) )
for h in hkl_range :
    for k in hkl_range :
        for l in hkl_range :
            if h==0 and k==0 and l==0 : continue
            # Since this is reflection Laue, only hkl directions within
            #     45 degrees of the the beam direction can produce a spot,
            #     so the cosine of the angle between them must be more than
            #     1/sqrt(2)
            cosang  =  ( (h*X_axis[0]+k*X_axis[1]+l*X_axis[2]) /
                                sqrt( (h**2+k**2+l**2)
                                 *(X_axis[0]**2+X_axis[1]**2+X_axis[2]**2) ) )
            if cosang > 1/sqrt(2) :
                if h**2+k**2+l**2 < hkl2_max :
                    hkl_all.append([h,k,l])
                    wavelengths.append(a_c*cosang/sqrt(h**2+k**2+l**2))
print("\nNumber of planes calculated : ", len(hkl_all))

def F(h,k,l,space_group) :
    from sys import exit
    # Calculate reflection conditions the space group of the crystal structure
    #     For now, this only works for a few simple cubic groups with simple
    #     unit celss  
    if   space_group == 221 : # Simple cubic 
        # Any h,k,l are allowed.
        return 1
    elif space_group == 225 :  # Face centred cubic
        # h,k,l must be all even or all odd, not mixed
        if (h%2==0 and l%2==0 and k%2==0) or (h%2!=0 and l%2!=0 and k%2!=0) :
            return 1
        else:
            return 0
    elif space_group == 227 :  # Face centred cubic diamond
        # Either h,k,l all even and h+k+l = 4N (N an integer) or
        #    or  h,k,l all odd and h+k+l = 4N±1
        if (h%2==0 and l%2==0 and k%2==0) :
            if (h+k+l) % 4 == 0 :
                return 1
            else :
                return 0
        elif (h%2!=0 and l%2!=0 and k%2!=0) :
            if ((h+k+l) % 4 == 1) or ((h+k+l) % 4 == 3) :
                return 1
            else :
                return 0
        else:
            return 0
    elif space_group == 229 :  # Body centred cubic
        # Sum of h,k,l must be even
        if (h+k+l) % 2 == 0 :
            return 1
        else:
            return 0
    else :
        print("\n!!!! SORRY, space group ", space_group,
                      "is currently not installed !!!!")
        exit()
            
spots = []
for i in range(len(hkl_all)) :
    hkl = hkl_all[i]
    h,k,l = hkl
    rotated_hkl = (((R*matrix(hkl).T).T).A)[0]
    y, z = spot_yz_from_hkl(rotated_hkl)+beam_yz
    # Skip if no valid y, z or if spot would be outside the area of the film.
    if ( y==None or z == None or y > max_y or y < min_y 
                              or z > max_z or z < min_z ) :
        continue
    theta,phi = plane_angles_from_hkl(rotated_hkl)
    # Calculate intensity using I_refl from Equation 2.26 in Preuss, but
    #     have simplified by noticing that (1+cot(theta)**2) = 1/sin(theta)**2
    # The wavelength dependence is probably not a good match for real X-ray
    #     beam spectra, but it is a start.
    #
    # theta defined here is complement of theta in Preuss, so need to convert
    theta = pi/2-theta
    if theta != pi/2 :
        w = wavelengths[i]
        # Skip if required wavelength is below minimum wavelength available
        if w < min_wave :
            continue
        I = spot_scale *( (1+cos(2*theta)**2)* ((w/min_wave-1)/w**3) *
                F(h,k,l,space_group)*
                (sin(theta)**2-cos(theta)**2)**3 /sin(theta)**4 )
    else :
        # Intensity of spots on beam axis is set to 1, since swamped by beam
        I=1
    spots.append([y,z,I,h,k,l])
y,z,I,h,k,l = zip(*spots)
# Convert tuplea to list so it can be modified later
I=list(I)
h=list(h)
k=list(k)
l=list(l)

## Create plot showing spot positions (spot positions in mm)
fig = figure(figsize=(8.5,10.))
ax = fig.add_subplot(111, aspect='equal')
# Set axes to just contain the active area of the image
ax.axis([min_y, max_y, min_z, max_z])

fig.suptitle("Calculated Spot Positions", fontsize=14)

# Define how close spot positions need to be to be considered the same
epsilon = 0.01
# Loop over spots adding their intensities if they are the same location
for i in range(len(y)) :
    if I[i] > 0 :     # Spots with zero intensity are ignored.
        for j in range(i+1,len(y)) :
            if I[j] > 0 :
                delta = (y[i]-y[j])**2+(z[i]-z[j])**2
                if delta < epsilon :
                    # Each plane contributes to the spot intensity
                    I[i] += I[j]
                    # Zero plane intensity once it is combined with another
                    I[j] = 0
                    # Label the spot with the lowest index for that location
                    hkli2 = h[i]**2+k[i]**2+l[i]**2
                    hklj2 = h[j]**2+k[j]**2+l[j]**2
                    if hklj2 < hkli2 :
                        h[i],k[i],l[i] = h[j],k[j],l[j]
        # Annotate spots on image with their hkl index if they have 
        #    sufficiently high intensity or low index.
        if ( I[i] > intensity_threshold 
                or h[i]**2+k[i]**2+l[i]**2 < hkl_threshold ) :
            text( y[i], z[i], str(h[i])+str(k[i])+str(l[i]),
                    size='xx-small', color="green")
            #print(h[i], k[i], l[i], I[i], y[i], z[i])

# Plot calculated spots
#   s is the size of the spots
scatter(y, z, s=I, color="red", linewidths=0, alpha = 0.9)

# Mark beam spot
scatter([beam_yz[0]], [beam_yz[1]], color='k', marker="+", s=400.0)

# Invert axis so it matches QLaue orientation
#     Note: This is the x axis in pyplot coordinates, but the y axis for
#     this program and QLaue.
ax.invert_xaxis()

print("\nTotal time : ", time()-start_time, "seconds")

savefig(saved_image_file, dpi=400, bbox_inches='tight')

show()

##### Laue_Spot_Identify.py

"""
Full text of MIT License:

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""
