Cacher is the code snippet organizer for pro developers

# coding: utf-8
# Created 2014, Zack Gainsforth

# Definitions:
# CLEN: CLinoENstatite, unit cell a = 9 Å
# OREN: ORthoENstatite, unit cell a = 18 Å
# HALF: Half unit cell, a = 4.5 Å
# PEN: ProtoENstatite, unit cell a = 18 Å
# An OREN -> CLEN transition always has to transform an entire 18 Å cell.  Hence always even numbers of CLEN cells.
# For PEN origin, you can get half-planes.  Don't know why.

import matplotlib
matplotlib.use('Qt4Agg')
import matplotlib.pyplot as plt
import numpy as np
import pyximport; pyximport.install()
from scipy import ndimage
from PIL import Image

### INPUTS SECTION

# Open the tif of the HRTEM image
FileName = '/Volumes/Stardust/Desktop/TEM Cecil unprocessed/20131219 - Libra - Cecil A6 S2/Cecil Grid A6 S2 - 0087-89 SUM.tif'

# If the scale on the image is off a little bit, then correct it.
ScaleCorr = 0.9651 / ((4.41 - 0.339)/5) # (22.399/22.3181)
ScaleCorr = 1.18
print 'ScaleCorr = ' + str(ScaleCorr)

# Stacking sequence.  To start with, everything is a CLEN.
# OREN has a double unit cell, and some half cells can occur for a PEN origin.
CLEN = 0.9651
HALF = CLEN/2 # 0.51 #0.4405
OREN = 1.9 #1.8233
Sequence = np.ones(30)*CLEN

Sequence[[6, 12, 15, 17, 20, 24]] = HALF
Sequence[[25]] = OREN

# We'll draw lines at regular intervals starting at x0.
x0 = 1.907 #0.339

### CODE SECTION

def QuickPlot(x, y, xlabel=None, ylabel=None, title=None, legendstrs=None, savefile=None, boldlevel=1, figax=None):
    """
    :param x: Can be a 1D numpy array, or a 2D numpy array.  If 1D, then there is one plot.  If 2D, every row is the x-axis for another line on the plot.
    :param y: Identical dimension to x.
    :param xlabel: String (can include TeX) for the label of the x-axis.
    :param ylabel: String (can include TeX) for the label of the y-axis.
    :param title: String (can include TeX) for the label of the title.
    :param legendstrs: List of strings with the same length as x and y, axis 1.
    :param savefile: A filename to save to if desired.
    :param boldlevel: Sets the linewidths fatter and texts bigger.  1 is normal, 4 is usually good for publication.
    :param figax: None makes a new figure and axis.  Or plots on an existing one if a tuple is given: (fig, ax)
    :return: (fig, ax) The figure and axis so the user can do stuff of his own with it if he wants to.
    """

    # Set the bold level.
    FontSizeBasis = (boldlevel+2)*4    # Fonts get bigger as boldlevel increases
    TickMajorBasis = boldlevel*4    # As fonts get bigger, they need a larger padding from the axis.
    # Increase the size of the tick label fonts.
    matplotlib.rc('xtick', labelsize=FontSizeBasis)
    matplotlib.rc('ytick', labelsize=FontSizeBasis)
    # Increase their padding.
    matplotlib.rc('xtick.major', pad=TickMajorBasis)
    matplotlib.rc('ytick.major', pad=TickMajorBasis)

    # Create/reuse a figure to plot on.
    if figax is None:
        fig, ax = plt.subplots()
    else:
        fig, ax = figax

    # Figure out how many lines are going on this plot.
    dim = np.shape(x)
    if len(dim) == 1:
        # Just one line on this plot.
        ax.plot(x,y, linewidth=boldlevel)
    elif len(dim) == 2:
        # Multiple lines on this plot.
        for i in range(dim[0]):
            ax.plot(x[i,:], y[i,:], linewidth=boldlevel)
    else:
        raise TypeError('x is not a 1D or 2D numpy array.')

    # Write the x and y labels and the title.
    if xlabel is not None:
        plt.xlabel(xlabel, fontsize=FontSizeBasis)
    if ylabel is not None:
        plt.ylabel(ylabel, fontsize=FontSizeBasis)
    if title is not None:
        plt.title(title, fontsize=FontSizeBasis)

    # Show a legend if input.
    if legendstrs is not None:
        ax.legend(legendstrs)

    # Resize the figure appropriately
    plt.tight_layout()

    # Save it to a file if a name is given.
    if savefile is not None:
        plt.savefig(savefile)

    return fig, ax

def GetLinePlotx(Image, Startx, y, width, height):
    """ Given a numpy 2D array representing an image, this creates an integrated (along the y-axis)
    linescan from Startx to Endx.  It returns the linescan and the section of the image that was integrated.
    """

    z = np.zeros(width)
    Endx = Startx + width

    # Height has to be an odd number (middle plus two sides)
    if not height % 2:
        height += 1

    for i in range(len(z)):
        z[i] = sum(Image[y-height/2:y+height/2, Startx+i])

    SubImage = Image[y-height/2:y+height/2, Startx:Endx]

    return z, SubImage

# Open the trace of the HRTEM image.
# FileName = '/Volumes/Stardust/Desktop/TEM Cecil unprocessed/20131219 - Libra - Cecil A6 S2/Cecil Grid A6 S2 - 0087-89 SUM Trace.txt'
# d = np.genfromtxt(FileName)

# FileName = '/Volumes/Stardust/Desktop/TEM Cecil unprocessed/20131219 - Libra - Cecil A6 S2/testimage.tif'
img = plt.imread(FileName)
imgtags = Image.open(FileName)

# TIFF tag 282 is the number of pixels/unit
dx = imgtags.tag.tags[282]
dx = float(dx[0][1])/float(dx[0][0])
# Now dx is units/pixel.  We will assume units are nm.
#print dx

# Rotate it, plot it and draw some vertical lines so we can adjust the rotation iteratively.
imgrot = ndimage.rotate(img, -127.15)
plt.gray()
plt.imshow(imgrot)
plt.plot([800, 800], [0,np.shape(imgrot)[0]], 'r')
plt.plot([1500, 1500], [0,np.shape(imgrot)[0]], 'r')
plt.gcf().set_size_inches(15, 15, forward=True)
plt.gcf().canvas.toolbar.pan()

# Get the linescan.
z, SubImage = GetLinePlotx(imgrot, 435, 1150, 1000, 100)

# Make an x-axis to go with the extracted line plot.
x = np.array(range(len(z)))*dx

# Apply the scale correction if our image scale is a little off.
x *= ScaleCorr

# Now make a figure with two subplots.  The top has our linescan.  The bottom
# has the subimage from which the linescan was derived.
fig = plt.figure()
ax1 = plt.subplot(2,1,1)
QuickPlot(x,z, xlabel='nm', boldlevel=2, figax=(fig,ax1))
ax2 = plt.subplot(2,1,2)
plt.imshow(SubImage)

# Lines will be drawn vertically from 0 up to the maximum intensity in the plot.
linetop = max(z)

# And loop through the sequence.
xnow = x0
for i in range(len(Sequence)):
    # Choose the color for the line.
    if Sequence[i] == CLEN:
        c = 'r'
    elif Sequence[i] == OREN:
        c = 'g'
    else:
        c = 'k'
    # And draw it.
    ax1.plot([xnow, xnow], [0, linetop], c)
    # Add an index so we can type in the new sequence
    ax1.text(xnow+0.05, 0, i)
    # Increment our x position depending on what width this cell was.
    xnow += Sequence[i]

fig.set_size_inches(24, 6, forward=True)
ax1.set_xlim(min(x),max(x))
plt.draw()
plt.show()