import sys
import importlib.metadata
def test_package(package_name):
    """Test if package exists and returns version or -1"""
    try:
        version = importlib.metadata.version(package_name)
    except importlib.metadata.PackageNotFoundError:
        version = '-1'
    return version

if test_package('pyTEMlib') < '0.2026.1.0':
    print('installing pyTEMlib')
    !{sys.executable} -m pip install  --upgrade pyTEMlib -q
print('done')

# import matplotlib and numpy
#                       use "inline" instead of "notebook" for non-interactive
#                       use widget for jupyterlab needs ipympl to be installed
# import matplotlib and numpy
%matplotlib widget  
import matplotlib.pylab as plt
import numpy as np

import sys
if 'google.colab' in sys.modules:
    from google.colab import output
    output.enable_custom_widget_manager()
    from google.colab import drive


%load_ext autoreload
%autoreload 2    

sys.path.insert(0, '../../')
import pyTEMlib

# For archiving reasons it is a good idea to print the version numbers out at this point
print('pyTEM version: ', pyTEMlib.__version__)
if 'google.colab' in sys.modules:
    drive.mount("/content/drive")
## Do all of registration

notebook_tags= {'notebook': 'Structure Analysis',
                'notebook_version': '2025_03_16',
                'pyTEM version': pyTEMlib.__version__}

# --------------- INPUT ------------------------
zone_hkl = np.array([1, 2, -2])
hkl_max = 34 #  maximum allowed Miller index
Sg_max = 0.03   # 1/Ang  maximum allowed excitation error

acceleration_voltage_V = 200.0 * 1000.0 #V
# -------------------------------------------
atoms = pyTEMlib.crystal_tools.structure_by_name('silicon')

hkl_all = pyTEMlib.diffraction_tools.get_all_miller_indices(hkl_max)
# all evaluated reciprocal_unit_cell points
g_non_rot = np.dot(hkl_all, atoms.cell.reciprocal())

ewald = pyTEMlib.diffraction_tools.ewald_sphere_center(acceleration_voltage_V, atoms, zone_hkl)

k0_magnitude = np.linalg.norm(ewald)

s = (k0_magnitude**2-np.linalg.norm(g_non_rot - ewald, axis=1)**2)/(2*k0_magnitude)

reflections = np.abs(s)<Sg_max

g = g_non_rot[reflections]
hkl = np.array(hkl_all[reflections])
s_g = s[reflections]
laue_zone = np.sum(hkl *zone_hkl, axis=1)

structure_factors = pyTEMlib.diffraction_tools.get_structure_factors(atoms, g)
    
allowed = structure_factors> 0.0001

def get_cylinder_coordinates (zone_hkl, g, k0_magnitude):
    theta, phi = pyTEMlib.diffraction_tools.find_angles(zone_hkl)
    rotation_matrix = pyTEMlib.diffraction_tools.basic.get_rotation_matrix([-phi, theta, 0], in_radians=True)
    center_rotated = [0, 0, k0_magnitude]
    
    g_rotated = np.dot(g, rotation_matrix)
    return  np.stack([np.arccos((g_rotated[:, 2]+k0_magnitude)/np.linalg.norm(g_rotated+center_rotated, axis=1)), 
                      np.arctan2(g_rotated[:, 1], g_rotated[:, 0]), 
                      g_rotated[:, 2]], 
                     axis=-1)


g_angles = get_cylinder_coordinates (zone_hkl, g, k0_magnitude)

np.sum(hkl*zone_hkl, axis=1), hkl,zone_hkl
g_spherical1 = (g_angles[allowed])
x = np.arctan(g_spherical1[:, 0])*k0_magnitude * np.cos(g_spherical1[:, 1])*10
y = np.arctan(g_spherical1[:, 0])*k0_magnitude  * np.sin(g_spherical1[:, 1])*10

plt.figure()
plt.scatter(x, y, cmap='tab10', c=laue_zone[allowed])

plt.gca().set_aspect('equal')
g_spherical
#np.abs(laue_zone).min(), sum(np.abs(laue_zone)==0), hkl[allowed][np.abs(laue_zone)==0]

tags = {'acceleration_voltage': acceleration_voltage_V,
        'zone_hkl': [1, 1, -0],
        'hkl_max': 15,
        'mistilt_alpha': np.radians(1),
        'mistilt_beta': np.radians(0),
        'Sg_max': 0.03}
diff_dict = {}
diff_dict = pyTEMlib.diffraction_tools.get_bragg_reflections(atoms, tags, verbose=True) 

diff_dict.setdefault('output',{})['plot_forbidden']=False
diff_dict['output']['plot_dynamically_allowed']=True
diff_dict['output']['plot_Kikuchi']=True
diff_dict['output']['plot_HOLZ']=True
diff_dict['convergence_angle'] = 3
pyTEMlib.diffraction_tools.plot_diffraction_pattern(diff_dict, unit='1/nm')

diff_dict['HOLZ'].keys()

spots = pyTEMlib.diffraction_tools.plotting_coordinates(diff_dict['allowed']['g'])

kikuchi = pyTEMlib.diffraction_tools.plotting_coordinates(diff_dict['Kikuchi']['g'], feature='line')
holz = pyTEMlib.diffraction_tools.plotting_coordinates(diff_dict['HOLZ']['g_deficient'], feature='line')

plt.figure()
plt.scatter(spots[:,0], spots[:, 1], cmap='tab10', c=diff_dict['allowed']['Laue_Zone'])
alpha =diff_dict['Kikuchi']['intensities']/ diff_dict['Kikuchi']['intensities'].max()*.5
for i, line in enumerate(kikuchi):
    plt.axline( (line[0], line[1]), slope=line[2], color='red',
                       alpha=alpha[i], linewidth=2)
alpha =diff_dict['HOLZ']['intensities']/ diff_dict['HOLZ']['intensities'].max()*.5
for i, line in enumerate(holz):
    plt.axline( (line[0], line[1]), slope=line[2], color='blue',
                       alpha=alpha[i], linewidth=2)
plt.gca().set_aspect('equal')
plt.xlabel('angle (mrad)');

g_spherical1 = diff_dict['allowed']['g']
inte = diff_dict['allowed']['intensities']
x = g_spherical1[:, 0] * np.cos(g_spherical1[:, 1])*1000
y = g_spherical1[:, 0] * np.sin(g_spherical1[:, 1])*1000

k = diff_dict['Kikuchi']
k['g'] = diff_dict['HOLZ']['g_deficient']
k['laue_circle'] = [0,0]
kx = ((k['g'][:, 0] * np.cos(k['g'][:, 1]+np.pi))/2)*1000
ky = ((k['g'][:, 0] * np.sin(k['g'][:, 1]+np.pi))/2)*1000
m = np.tan((k['g'][:, 1]-np.pi/2))      

plt.figure()
plt.scatter(x, y, cmap='tab10', c=diff_dict['allowed']['Laue_Zone'])
plt.xlabel('angle (mrad)')

for i, x in enumerate(kx):
    plt.axline( (x, ky[i]), slope=m[i], color='red',
                       alpha=.5, linewidth=1)
diff_dict.keys(), diff_dict['Kikuchi']['laue_circle'], tags['mistilt_alpha']
plt.gca().set_aspect('equal')
diff_dict['HOLZ'].keys()

((g_spherical1[:10,0]) - (np.arcsin(g_spherical1[:10,2]/g_spherical1[:10,3])))*1000 

s = pyTEMlib.diffraction_tools.get_structure_factors(atoms, g_non_rot)
s.shape, atoms.positions, list(zip(np.round(s[200:210],3), hkl_all[200:210]))

plt.close('all')

g_non_rot[:100].shape

F =  get_structure_factors(atoms, g_non_rot[:100])
F.shape, (F>0.0001).sum(), zip 

def get_structure_factors(atoms, g_hkl):
    form_factor = np.zeros((len(atoms.positions), g_hkl.shape[0]))
    for symbol in np.unique(atoms.symbols):
        atom_positions = atoms.symbols==symbol
        form_factor[(atom_positions)] = get_form_factor(symbol, np.linalg.norm(g_hkl, axis=1))    
    structure_factors = calculate_structure_factors(np.array(g_hkl), form_factor, np.array(atoms.positions))
    return structure_factor, form_factor.sum(axis=1) 

def calculate_structure_factors(all_g, form_factor, atom_positions):
    """ Calculate structure factors for given reciprocal lattice points g_hkl"""
    structure_factors = np.zeros(len(all_g))
    form_factor_sum = np.zeros(len(all_g))
    for i, g_i in enumerate(all_g):
        struct_factor = 0.0+0.0*1j
        for j, f_j in enumerate(form_factor):
            f_q_j = f_j[i]
            r_j = atom_positions[j]
            struct_factor += f_q_j * (np.exp(-2*np.pi*1j*np.dot(g_i, r_j))).sum() 
        structure_factors[i] = (struct_factor*struct_factor.conj()).real
        form_factor_sum[i] = (form_factor[:,i]).sum()
    return structure_factors

pyTEMlib.diffraction_tools.form_factor('Si', g_norm), form_factor('Si', g_norm[0]), g_norm[0]

def get_form_factor(element, q):
    
    if isinstance(q, float):
        q = np.array([q])
    if not isinstance(q, np.ndarray):
        print('not float')
        return
    parameter_dict = pyTEMlib.crystal_tools.electronFF[element]
    f_parameter = np.array([parameter_dict[key] for key in ['fa', 'fb', 'fc', 'fd']])
    q = (np.array([q, q, q])).T
    #f =  f_lorentzian + f_gauss
    f = ((f_parameter[0]/(q**2 + f_parameter[1])).sum(axis=1) 
         + (f_parameter[2]*np.exp(-q**2 * f_parameter[3])).sum(axis=1))
    
    return f

hkl2, g2=  pyTEMlib.diffraction_tools.get_all_reflections(atoms, 1)
g, hkl = pyTEMlib.diffraction_tools.get_all_g_vectors(1, atoms)

g[:5], g2[:5], hkl2[:5], hkl[:5]

fileWidget = pyTEMlib.file_tools.FileWidget()

image = fileWidget.selected_dataset
view = image.plot()

rigid_registered = pyTEMlib.image_tools.rigid_registration(image)
v= rigid_registered.plot()

image2 = rigid_registered[:, 300:-212, 512:].sum(axis=0)
image2 -= image2.max()/2
image2[image2<0] = 0.
image2.data_type='image'
v = image2.plot()

fft_image = pyTEMlib.image_tools.power_spectrum(image2)

v = fft_image.plot()
#, diffractogram_spots

spots = pyTEMlib.image_tools.diffractogram_spots(fft_image, 0.05)[0]
plt.scatter(spots[:, 0], spots[:,1], c='r')

angles = np.degrees(np.atan2(spots[1:, 1],spots[1:, 0]))

angles = angles-angles.min()- 180
angles, np.linalg.norm(spots[1:, :2],axis=1)

6.90524432/ 6.81851881

plt.figure()
plt.scatter(angles, r = [6.92]*len(angles)) #np.linalg.norm(spots[1:, :2],axis=1))

# --------------- INPUT ------------------------
structure = 'gold'
zone_hkl = np.array([1, 1, 1])
hkl_max = 8 #  maximum allowed Miller index
sg_max = 0.03   # 1/Ang  maximum allowed excitation error

acceleration_voltage = 200.0 * 1000.0 #V

rotation = np.radians(0)  # rotation of diffraction pattern
# -------------------------------------------
atoms = pyTEMlib.crystal_tools.structure_by_name(structure)
tags = {'zone_hkl': zone_hkl,
        'hkl_max': hkl_max,
        'Sg_max': sg_max,
        'mistilt_alpha': np.radians(3),
        'acceleration_voltage':  acceleration_voltage}

diff_dict ={}
diff_dict = pyTEMlib.diffraction_tools.get_bragg_reflections(atoms, tags, verbose=True) 

# Simple Plot
ZOLZ = diff_dict['allowed']['ZOLZ']
HOLZ = diff_dict['allowed']['HOLZ']
r = diff_dict['allowed']['g'][:, 0]
phi = diff_dict['allowed']['g'][:, 1]

x = r *np.cos(phi+rotation)*10
y = r * np.sin(phi+rotation)*10


plt.figure()
plt.scatter(x[ZOLZ], y[ZOLZ], label='ZOLZ allowed', c='r')
plt.scatter(x[HOLZ], y[HOLZ], label="HOLZ allowed", c ='orange')
plt.axis('equal')
plt.xlabel('reciprocal distance (1/nm)');

plt.figure()
plt.scatter(phi, r*10)

# -----Input for ring pattern calculation ----
structure = 'gold'
hkl_max = 15
verbose = True
# --------------------------------------------

atoms = pyTEMlib.crystal_tools.structure_by_name(structure)
image2.structures['Structure_000'] = atoms
image2.metadata['experiment']['hkl_max']  = hkl_max

pyTEMlib.diffraction_tools.ring_pattern_calculation(image2, verbose=verbose)

image3 = np.array(image2)-(image2.max()/2)
image3 [image3<0.] =0
im2 = np.fft.fftshift(np.fft.fft2(image3))
plt.close('all')
plt.figure()
plt.imshow(np.log(1+np.abs(im2)))
#plt.imshow(image3)
image3[346,263], image2.max(), image2.min(), image3.max(), image3.min()

chooser.dataset

dataset = chooSer.dataset
dataset

dataset = np.array(chooser.dataset).squeeze()
pixels = np.where(dataset==65535)[0]
frames = np.where(pixels%(215*178)==0)[0]
pixel_frames = np.searchsorted(pixels,frames)

spectra = np.zeros([215*178,2048]) 

def get_spectra(dataset, start, end):
    spectra = np.zeros([end-start, 2048]) 
    for frame in pixel_frames:
        last_pixel = pixels[frame+start]
        for index,  next_pixel in enumerate(pixels[frame+start:frame+end]):
            for pixel in range(last_pixel+1, next_pixel):
                spectra[index, int(dataset[pixel]/2)] += 1
            last_pixel = next_pixel

    return spectra
spectra = get_spectra(dataset, int(215*178/8), int(215*178/4))
plt.figure()
plt.plot(spectra.sum(axis=0))

pixels = np.where(np.array(dataset)==65535)[0]
print(pixels.shape[0]-dataset.shape[0])
frames = pixels[(215*178) *np.arange(63)]
print(frames)

spectra = np.zeros([215*178,2048]) 
for frame in frames[60:-1]:
    for i in range(40):
        for j in range(pixels[i+frame],pixels[i+1+frame]):
            #print(i, j, data[j])
            if data[j] < 65535:
                spectra[i, int(data[j]/2)] += 1
plt.figure()
plt.plot(spectra[:10000].sum(axis=0))

38282, 215*178

fileWidget.file_name
datasets = pyTEMlib.file_tools.open_file(fileWidget.file_name, eds_stream=True)
selector = sidpy.ChooseDataset(datasets)

dataset = selector.datasetsqueeze()
pixels = np.where(dataset==65535)[0]
frames = np.where(pixels%(215*178)==0)

frames
selector.dataset.squeeze().shape

for 

dataset = selector.dataset
dataset.data_type = 'IMAGE_STACK'
dataset.x.dimension_type = 'SPATIAL'
dataset.y.dimension_type = 'SPATIAL'

view = dataset.plot()
dataset.x.slope

rigid_registered= pyTEMlib.image_tools.rigid_registration(dataset)
datasets['rigid_registered'] = rigid_registeredrigid_registered
view = rigid_registered.plot()

demon_registered = pyTEMlib.image_tools.demon_registration(rigid_registered)
datasets['demon_registered'] = demon_registered
view2 = demon_registered.plot()

import importlib
importlib.reload(pyTEMlib.image_tools)

demon_registered, rigid_registered = pyTEMlib.image_tools.complete_registration(dataset)
demon_registered.data_type = 'image_stack'
view2 = demon_registered.plot()

Save these datasets

pyTEMlib.file_tools.save(FileWidget.filename{:-3] +'hdf5')

drift = rigid_registered.metadata['drift']
polynom_degree = 2 # 1 is linear fit, 2 is parabolic fit, ...

x = np.linspace(0,drift.shape[0]-1,drift.shape[0])

line_fit_x = np.polyfit(x, drift[:,0], polynom_degree)
poly_x = np.poly1d(line_fit_x)
line_fit_y = np.polyfit(x, drift[:,1], polynom_degree)
poly_y = np.poly1d(line_fit_y)

plt.figure()
plt.axhline(color = 'gray')
plt.plot(x, drift[:,0], label = 'drift x')
plt.plot(x, drift[:,1], label = 'drift y')
plt.plot(x, poly_x(x),  label = 'fit_drift_x')
plt.plot(x, poly_y(x),  label = 'fit_drift_y')

plt.legend();
ax_pixels = plt.gca()
ax_pixels.step(1, 1)

scaleX = 1000 #(rigid_registered.x[1]-rigid_registered.x[0])*1000.  #in pm

ax_pm = ax_pixels.twinx()
x_1, x_2 = ax_pixels.get_ylim()

ax_pm.set_ylim(x_1*scaleX, x_2*scaleX)

ax_pixels.set_ylabel('drift [pixels]')
ax_pm.set_ylabel('drift [pm]')
ax_pixels.set_xlabel('image number');
plt.tight_layout()

# ------- Input ------
atoms_size = .06 # in nm
# --------------------

# image = demon_registered.sum(axis=0)
# image = dataset.sum(axis=0)

out_tags = {}
image.metadata['experiment']= {'convergence_angle': 30, 'acceleration_voltage': 200000.}

scale_x =image.x.slope
gauss_diameter = atoms_size/scale_x
print(gauss_diameter)
if gauss_diameter < 3:
    print('smal')
    gauss_diameter = 3
print(gauss_diameter, gauss_diameter*scale_x)
gauss_probe = pyTEMlib.probe_tools.make_gauss(image.shape[0], image.shape[1], gauss_diameter)

print('Deconvolution of ', dataset.title)
LR_dataset = pyTEMlib.image_tools.decon_lr(image, gauss_probe, verbose=False)

extent = LR_dataset.get_extent([0,1])
fig, ax = plt.subplots(1, 2, figsize=(8, 4), sharex=True, sharey=True)
ax[0].imshow(image.T, extent = extent,vmax=np.median(np.array(image))+3*np.std(np.array(image)))
ax[1].imshow(LR_dataset.T, extent = extent, vmax=np.median(np.array(LR_dataset))+3*np.std(np.array(LR_dataset)));

datasets['lr_decon'] = LR_dataset

LR_dataset.metadata.update({'analysis': {'Lucy_Richardson': {
                        'notebook': 'Image_Registration' ,
                        # 'notebook_version': __notebook_version__,
                        'input': dataset.title,
                        'probe_diameter': gauss_diameter,
                        'kind_of_probe': 'Gauss',
                        'probe_width': atoms_size
                        }}})
LR_dataset.metadata

import skimage
# ------- Input ------
threshold = 20.7 #usally between 0.01 and 0.9  the smaller the more atoms
atom_size = .05 #in nm  
# ----------------------


image = LR_dataset
#image = image_choice.dataset
# scale_x = pyTEMlib.file_tools.get_slope(image.dim_1)
blobs =  skimage.feature.blob_log(image, max_sigma=atom_size/scale_x, threshold=threshold)

fig1, ax = plt.subplots(1, 1,figsize=(8,7), sharex=True, sharey=True)
plt.title("blob detection ")

plt.imshow(image.T, interpolation='nearest',cmap='gray', vmax=np.median(np.array(image))+3*np.std(np.array(image)))
plt.scatter(blobs[:, 0], blobs[:, 1], c='r', s=20, alpha = .5);

out_tags  =  {}
out_tags['analysis']= 'Atom Positions'
# out_tags['notebook']= __notebook__ 
# out_tags['notebook_version']= __notebook_version__

out_tags['atoms'] = blobs
out_tags['atom_size'] = atom_size #in nm gives the size of the atoms or resolution
out_tags['threshold'] =  threshold  #between 0.01 and 0.1 
out_tags['pixel_size'] = scale_x


out_tags['name'] = 'Atom finding'
out_tags['title'] = out_tags['name']
tags = {'atom_pixel':  out_tags}
if isinstance(image.metadata['analysis'], str):
    image.metadata['analysis']={image.metadata['analysis']:{}}

image.metadata['analysis'].update(tags)

image.metadata['analysis']

image_choice = sidpy.ChooseDataset(datasets)  

# ------- Input ------
atom_radius = 3  # in pixel
# ----------------------

if image_choice.dataset.data_type.name == 'IMAGE_STACK':
    image = image_choice.dataset.sum(axis=0)
else:
    image = image_choice.dataset

#atoms = atom_group['atoms'][()]
atoms = blobs
image = image-image.min()
print(image)

#atom_radius = 2
MaxInt = 0
MinInt = 0 
maxDist = 2
sym = pyTEMlib.atom_tools.atom_refine(np.array(image), atoms, atom_radius, max_int = 0, min_int = 0, max_dist = 2)
refined_atoms = np.array(sym['atoms'])

fig, ax = plt.subplots(1, 2, figsize=(8, 4), sharex=True, sharey=True)
ax[0].imshow(image.T)
ax[0].scatter(refined_atoms[:,0],refined_atoms[:,1],  s=10, alpha = 0.3, color = 'red')
ax[0].set_title('refined atom postion')
ax[1].imshow(image.T)
ax[1].scatter(atoms[:, 0], atoms[:, 1], c='r', s=10, alpha = .3);
ax[1].set_title('blobs on image');

out_tags  =  {}
out_tags['analysis']= 'Atom Positions'
# out_tags['notebook']= __notebook__ 
# out_tags['notebook_version']= __notebook_version__

out_tags['atoms'] = blobs
out_tags['atom_size'] = atom_size #in nm gives the size of the atoms or resolution
out_tags['threshold'] =  threshold  #between 0.01 and 0.1 
out_tags['pixel_size'] = scale_x


out_tags['name'] = 'Atom finding'
out_tags['title'] = out_tags['name']
tags = {'atom_pixel':  out_tags}
if isinstance(image.metadata['analysis'], str):
    image.metadata['analysis']={image.metadata['analysis']:{}}

image.metadata['analysis'].update(tags)

image.metadata['analysis']

crystal_name = 'Au'
crystal = pyTEMlib.crystal_tools.structure_by_name(crystal_name)
print(crystal_name)
crystal

crystal.cell.lengths()[0]/image.x.slope/10
crystal.info = {'experimental': {'zone_axis': [1, 1, 1], 'angle': 0}}
layer = pyTEMlib.crystal_tools.get_projection(crystal)
gamma =  layer.cell.angles()[2]
layer.cell.lengths() , layer.cell.lengths()/image.x.slope/10 , layer.cell.angles(), np.linalg.norm(layer.positions[:,:2], axis=1), layer.get_scaled_positions()

import scipy
import pyTEMlib.graph_tools
def breath_first(graph, initial, lattice_parameter, tolerance=1):
    
    neighbour_tree = scipy.spatial.KDTree(graph)
    distances, indices = neighbour_tree.query(graph,  # let's get all neighbours
                                              k=50)  # projection_tags['number_of_nearest_neighbours']*2 + 1)
    visited = []  # the atoms we visited
    angles = []  # atoms at ideal lattice
    sub_lattice = []
    sub_lattice = []  # atoms in base and disregarded
    queue = [initial]
    queue_angles=[0]
    
    while queue:
        node = queue.pop(0)
        angle = queue_angles.pop(0)
        
        if node not in visited and node not in sub_lattice:
            visited.append(node)
            angles.append(angle)
            neighbors = indices[node]
            for i, neighbour in enumerate(neighbors):
                if neighbour not in visited:
                    hopp = graph[node] - graph[neighbour]
                    distance_to_ideal = np.linalg.norm(hopp)
                    if distance_to_ideal < lattice_parameter - tolerance*5:
                        sub_lattice.append(neighbour) 
                    elif np.min(np.abs(distance_to_ideal- lattice_parameter)) < tolerance:
                        queue.append(neighbour) 
                        queue_angles.append(np.arctan2(hopp[1], hopp[0]))
    angles[0] = angles[1]
    out_atoms = np.stack([graph[visited][:, 0], graph[visited][:, 1], angles])
    return out_atoms.T, visited

def delete_rim_atoms(atoms, extent, rim_distance):
    rim = np.where(atoms[:, :2] - extent > -rim_distance)[0]
    middle_atoms = np.delete(atoms, rim, axis=0)
    rim = np.where(middle_atoms[:, :2].min(axis=1)<rim_distance)[0]
    middle_atoms = np.delete(middle_atoms, rim, axis=0)
    return middle_atoms
init = 633
#init = 6328
#init= 6589

%time hopped_atoms, indices = breath_first(blobs, init, 22, 2)
middle_atoms = pyTEMlib.graph_tools.delete_rim_atoms(hopped_atoms, image.shape, 10)

plt.close('all')
fig1, ax = plt.subplots(1, 1,figsize=(8,7), sharex=True, sharey=True)
plt.title("blob detection ")

plt.imshow(demon_registered.sum(axis=0).T, interpolation='nearest',cmap='gray', vmax=np.median(np.array(image))+3*np.std(np.array(image)))
plt.scatter(middle_atoms[:, 0], middle_atoms[:, 1], c=np.degrees(np.degrees(middle_atoms[:, 2])% 60), cmap = 'Reds', s=20, alpha = .5);
plt.scatter(middle_atoms[:, 0], middle_atoms[:, 1], c='red', s=20, alpha = .5);
#plt.scatter(blobs[init][0], blobs[init][1], c='orange')
angles = np.degrees(middle_atoms[:, 2])% 60

crystal_name = 'SrTiO3'
crystal = pyTEMlib.crystal_tools.structure_by_name(crystal_name)
print(crystal_name)
crystal

crystal.cell.lengths()[0]/image.x.slope/10
crystal.info = {'experimental': {'zone_axis': [0, 0, 1], 'angle': 0}}
layer = pyTEMlib.crystal_tools.get_projection(crystal)
print(layer.cell.angles())
gamma =  layer.cell.angles()[2]
layer.cell.lengths() , layer.cell.lengths()/image.x.slope/10 , layer.cell.angles(), np.linalg.norm(layer.positions[:,:2], axis=1), layer.get_scaled_positions()

print(len(blobs))
fig1, ax = plt.subplots(1, 1,figsize=(8,7), sharex=True, sharey=True)
plt.title("blob detection ")

plt.imshow(image.T, interpolation='nearest',cmap='gray', vmax=np.median(np.array(image))+3*np.std(np.array(image)))
plt.scatter(middle_atoms[:, 0], middle_atoms[:, 1], c=np.degrees(np.degrees(middle_atoms[:, 2])% gamma), cmap = 'viridis', s=20, alpha = .5);
angles = np.degrees(middle_atoms[:, 2])% gamma

print(f' Average unit cell angle {np.average(angles):.1f} with standard deviation {np.std(angles):.2f}; from {np.min(angles):.1f} to {np.max(angles):.1f}')

plt.figure()
plot_angles = np.append(angles, angles+gamma)
plot_angles[plot_angles<0] +=gamma 
counts, bins = np.histogram(plot_angles, bins = 30)
plt.stairs(counts, bins)

from sklearn.neighbors import KernelDensity
a = angles.reshape(-1,1)

kde = KernelDensity(kernel='gaussian', bandwidth=3).fit(a)

kde.score_samples(a)

X_plot = np.linspace(0, 90, 100)[:, np.newaxis] 

fig, ax = plt.subplots() 
  
  
# Calculating the density using the gaussian kernel with bandwidth 0.5 
kde = KernelDensity(kernel='gaussian', bandwidth=1).fit(a) 
  
# Calculating the log of the probability density function 
log_dens = kde.score_samples(X_plot) 
  
# Plotting the density curve 
ax.plot( 
    X_plot[:, 0], 
    np.exp(log_dens)*1000, 
    color="cornflowerblue", 
    linestyle="-", 
    label="Gaussian kernel density"
) 
counts, bins = np.histogram(angles, bins = 40)
ax.stairs(counts, bins, color='blue')
  

one_grain_indices = np.where(angles > 44)
one_grain = middle_atoms[one_grain_indices]
one_grain_angle = np.median(one_grain[:, 2]%np.radians(60))

two_grain_indices = np.where(angles < 44)
two_grain = middle_atoms[two_grain_indices]
two_grain_angle = np.median(two_grain[:, 2]%np.radians(60))

gamma = np.radians(layer.cell.angles()[2])

one_grain[100]
plt.figure()
plt.imshow(image.T, interpolation='nearest',cmap='gray', vmax=np.median(np.array(image))+3*np.std(np.array(image)))
#plt.scatter(one_grain[:, 0], one_grain[:, 1], c='blue', alpha = 0.2)

init = np.argmin(np.linalg.norm(blobs[:,:2]- [674, 594], axis=1))
#init = np.argmin(np.linalg.norm(blobs[:,:2]- [610, 552], axis=1))

plt.scatter(blobs[init][0], blobs[init][1], c='orange')
projection_tags = {'lattice_vector': {'a': np.array([np.cos(one_grain_angle)*21, np.sin(one_grain_angle)*21]),
                                      'b': np.array([np.cos(one_grain_angle+gamma)*21, np.sin(one_grain_angle+gamma)*21]) },
                   'allowed_variation': 1.5,
                   'distance_unit_cell':  21*1.04}
layer.info['projection'] = projection_tags
hop1, ideal = pyTEMlib.graph_tools.breadth_first_search2(blobs[:,:2], init, layer)
projection_tags = {'lattice_vector': {'a': np.array([np.cos(two_grain_angle)*21, np.sin(two_grain_angle)*21]),
                                      'b': np.array([np.cos(two_grain_angle+gamma)*21, np.sin(two_grain_angle+gamma)*21]) },
                   'allowed_variation': 1.5,
                   'distance_unit_cell':  21*1.04}
layer.info['projection'] = projection_tags
#init = np.argmin(np.linalg.norm(blobs[:,:2]- [574, 520], axis=1))

print(init)
hop2, ideal = pyTEMlib.graph_tools.breadth_first_search2(blobs[:,:2], init, layer)

plt.scatter(hop1[:,0], hop1[:,1], c='red', alpha = 0.3)
plt.scatter(hop2[:,0], hop2[:,1], c='blue', alpha = 0.3)


one_grain_angle

one_grain_indices = np.argmax(angles)
one_grain = middle_atoms # [one_grain_indices]
print(one_grain.shape)
one_grain_angle = np.max(one_grain[:, 2]%np.radians(60))


gamma = np.radians(layer.cell.angles()[2])
length = 22
plt.close('all')
plt.figure()
plt.imshow(image.T, interpolation='nearest',cmap='gray', vmax=np.median(np.array(image))+3*np.std(np.array(image)))
#plt.scatter(one_grain[:, 0], one_grain[:, 1], c='blue', alpha = 0.2)

init = np.argmin(np.linalg.norm(blobs[:,:2]- [812, 199], axis=1))

#init = np.argmin(np.linalg.norm(blobs[:,:2]- [610, 552], axis=1))
# init = 102
#plt.scatter(blobs[init][0], blobs[init][1], c='orange')
projection_tags = {'lattice_vector': {'a': np.array([np.cos(one_grain_angle)*length, np.sin(one_grain_angle)*length]),
                                      'b': np.array([np.cos(one_grain_angle+gamma)*length, np.sin(one_grain_angle+gamma)*length]) },
                   'allowed_variation': 5,
                   'distance_unit_cell':  length*1.04}
layer.info['projection'] = projection_tags
hop1, ideal = pyTEMlib.graph_tools.breadth_first_search2(blobs[:,:2], init, layer)
init = np.argmin(np.linalg.norm(blobs[:,:2]- [89, 989], axis=1))

hop2, ideal = pyTEMlib.graph_tools.breadth_first_search2(blobs[:,:2], init, layer)
plt.scatter(hop1[:,0], hop1[:,1], c='red', alpha = 0.3)
plt.scatter(hop2[:,0], hop2[:,1], c='blue', alpha = 0.3)


one_grain_angle

import ase
positions = hopped_atoms.copy()
positions[:,2] = 0
atoms = ase.Atoms('Sr'*len(positions),
             positions=positions,
             cell=[image.shape[0],image.shape[1],0],
             pbc=[0, 0, 0])
atoms.info['bond_radii'] = [3] * len(positions)
graph_dictionary = pyTEMlib.graph_tools.find_polyhedra(atoms)

len(positions), len (graph_dictionary['cyclicity'])

from matplotlib.collections import PatchCollection
from matplotlib import cm
import matplotlib
centers =  graph_dictionary['centers']
cyclicity = graph_dictionary['cyclicity']

unit_cells = PatchCollection(graph_dictionary['polygons'], alpha=0.4, cmap=matplotlib.cm.tab10)
plt.figure()
plt.imshow(image.T,  cmap = 'afmhot')

unit_cells.set_array(cyclicity)
unit_cells.set_edgecolor('black')
plt.gca().add_collection(unit_cells)
#plt.scatter(centers[:,0],centers[:,1],color='blue',alpha=0.5, s = 3)
#plt.scatter(hopped_atoms[:,0], hopped_atoms[:,1], c='red', alpha = 0.3)
cbar = plt.colorbar(unit_cells, label='cyclicity')

import simpleITK as sitk

def DemonReg(cube, verbose = False):
    """
    Diffeomorphic Demon Non-Rigid Registration 
    Usage:
    ------
    DemReg = DemonReg(cube, verbose = False)

    Input:
        cube: stack of image after rigid registration and cropping
    Output:
        DemReg: stack of images with non-rigid registration

    Dempends on:
        simpleITK and numpy
    
    Please Cite: http://www.simpleitk.org/SimpleITK/project/parti.html
    and T. Vercauteren, X. Pennec, A. Perchant and N. Ayache
    Diffeomorphic Demons Using ITK\'s Finite Difference Solver Hierarchy
    The Insight Journal, http://hdl.handle.net/1926/510 2007
    """
    
    DemReg =  np.zeros_like(cube)
    nimages = cube.shape[0]
    print(nimages)
    # create fixed image by summing over rigid registration

    fixed_np = np.average(current_dataset, axis=0)

    fixed = sitk.GetImageFromArray(fixed_np)
    fixed = sitk.DiscreteGaussian(fixed, 2.0)

    #demons = sitk.SymmetricForcesDemonsRegistrationFilter()
    demons = sitk.DiffeomorphicDemonsRegistrationFilter()

    demons.SetNumberOfIterations(200)
    demons.SetStandardDeviations(1.0)

    resampler = sitk.ResampleImageFilter()
    resampler.SetReferenceImage(fixed);
    resampler.SetInterpolator(sitk.sitkBspline)
    resampler.SetDefaultPixelValue(0)

    done = 0
        
    for i in range(nimages):
        if done < int((i+1)/nimages*50):
            done = int((i+1)/nimages*50)
            sys.stdout.write('\r')
            # progress output :
            sys.stdout.write("[%-50s] %d%%" % ('*'*done, 2*done))
            sys.stdout.flush()
        
        moving = sitk.GetImageFromArray(cube[i])
        movingf = sitk.DiscreteGaussian(moving, 2.0)
        displacementField = demons.Execute(fixed,movingf)
        outTx = sitk.DisplacementFieldTransform( displacementField )
        resampler.SetTransform(outTx)
        out = resampler.Execute(moving)
        DemReg[i,:,:] = sitk.GetArrayFromImage(out)
        #print('image ', i)
        
    
    print(':-)')
    print('You have succesfully completed Diffeomorphic Demons Registration')
    
    return DemReg