python_tools/beam_profile.py at master · jpyanez/python_tools · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
from mpl_toolkits import mplot3d
import jp_mpl as jplot
from scipy import interpolate, ndimage
from copy import deepcopy

def profile2D(image, pixel_size):

    image_dimensions = np.array(image.shape)*pixel_size # in mm
    #print 'Image dimensions', image_dimensions, ' (mm) '
    xaxis = np.linspace(0, 1, image.shape[1])*image_dimensions[1]
    yaxis = np.linspace(0, 1, image.shape[0])*image_dimensions[0]

    fig, ax = plt.subplots(2,2,
                           figsize = (12,12),
                           sharey='row',sharex='col',
                           gridspec_kw={'width_ratios':[1,4],
                                        'height_ratios':[4,1]})
    fig.delaxes(ax[1,0])

    xmax = image.max(axis=0)
    ax[1,1].plot(xaxis,xmax)

    ymax = image.max(axis=1)
    ax[0,0].plot(ymax, yaxis[::-1])

    ax[0,1].imshow(image, plt.cm.gray,
                   aspect='auto',
                   extent = [0, image_dimensions[1], 0, image_dimensions[0]])

    plt.subplots_adjust(wspace=0, hspace=0)

    ax[0,0].set_ylabel('y (mm)')
    ax[1,1].set_xlabel('x (mm)')
    return fig

def profileContour(image, pixel_size):
    image_dimensions = np.array(image.shape)*pixel_size # in mm
    xaxis = np.linspace(0, 1, image.shape[1])*image_dimensions[1]
    yaxis = np.linspace(0, 1, image.shape[0])*image_dimensions[0]

    fig = plt.figure( figsize = (12,9))
    ax = fig.add_subplot(111)

    levels = np.linspace(noise, image.max(), 10)

    ####
    cmap = plt.cm.viridis
    cmaplist = [cmap(i) for i in range(cmap.N)]
    cmap = mpl.colors.LinearSegmentedColormap.from_list(
        'Custom cmap', cmaplist, cmap.N)
    bounds = levels
    norm = mpl.colors.BoundaryNorm(bounds, cmap.N)
    #######

    c = ax.imshow(image, cmap=cmap,
                  norm=norm,
                  aspect='auto',
                  extent = [0, image_dimensions[1], 0, image_dimensions[0]])

    cf = ax.contour(xaxis, yaxis[::-1], image,
                    levels = levels,
                    colors = 'white',
                    alpha = 0.5)

    plt.colorbar(c)

    plt.ylim(3, 7)
    plt.xlim(1.5, 8.5)

    return fig

def profile3D(image, pixel_size):
    image_dimensions = np.array(image.shape)*pixel_size # in mm
    xaxis = np.linspace(0, 1, image.shape[1])*image_dimensions[1]
    yaxis = np.linspace(0, 1, image.shape[0])*image_dimensions[0]

    fig = plt.figure(figsize=(18,16))
    ax = plt.axes(projection='3d')
    y = np.linspace(0, image_dimensions[0], image.shape[0])
    x = np.linspace(0, image_dimensions[1], image.shape[1])
    X, Y = np.meshgrid(x, y)
    ax.view_init(30, 40)
    ax.plot_surface(X, Y, image,
                   cmap='viridis')
    ax.set_xlabel('x (mm)')
    ax.set_ylabel('y (mm)')
    ax.set_zlabel('Intensity')

    zoom = True
    if zoom:
        step= 6
        ax.set_xlim(3, 3+step)
        ax.set_ylim(2, 2+step)
    return fig

import re

def read_pgm(filename, byteorder='>'):
    """Return image data from a raw PGM file as numpy array.

    Format specification: http://netpbm.sourceforge.net/doc/pgm.html

    """
    #print filename
    with open(filename, 'rb') as f:
        buffer = f.read()
    try:
        header, width, height, maxval = re.search(
            b"(^P5\s(?:\s*#.*[\r\n])*"
            b"(\d+)\s(?:\s*#.*[\r\n])*"
            b"(\d+)\s(?:\s*#.*[\r\n])*"
            b"(\d+)\s(?:\s*#.*[\r\n]\s)*)", buffer).groups()
    except AttributeError:
        raise ValueError("Not a raw PGM file: '%s'" % filename)
    return np.frombuffer(buffer,
                            dtype='u1' if int(maxval) < 256 else byteorder+'u2',
                            count=int(width)*int(height),
                            offset=len(header)
                            ).reshape((int(height), int(width)))