plot.py 3.35 KB
Newer Older
1 2 3
#!/usr/bin/env python
from pylab import *
from mpl_toolkits.mplot3d import axes3d
4
import numpy as np
5
from staticbs.common import mag
6

7

8
def plot_linearity(B_fn, obs_points_fn="./data/obs_points_linearity.npy"):
9 10
    obs_points = np.load(obs_points_fn)
    B = np.load(B_fn)
11

12 13 14
    x = obs_points[:, 0]
    y = obs_points[:, 1]
    z = obs_points[:, 2]
15

16 17 18 19
    plt.subplot(2, 1, 1)
    plt.title('B field as a function of the X position')
    plt.xlabel('x position when y = 0')
    plt.plot(x[0:x.shape[0]/2], B[0:x.shape[0]/2, 0], '.')
20

21 22 23 24
    plt.subplot(2,1,2)
    plt.title('B field as a function of the y position')
    plt.xlabel('y position when x = 0')
    plt.plot(y[y.shape[0]/2:-1], B[y.shape[0]/2:-1, 1], '.')
25

26
    # ax4.scatter(x, y)
27

28 29 30
    # ax4.quiver(x, y,
    #        B[:, 0],
    #        B[:, 1], color='blue')
31

32 33 34 35 36 37 38 39 40 41 42 43 44

def scatter_3D(B_fn, obs_points_fn="data/obs_points.npy"):
    """Plot a scatter graph in 3D of the B field"""
    obs_points = np.load(obs_points_fn)
    B = np.load(B_fn)

    xx = obs_points[:, :, 0]
    yy = obs_points[:, :, 1]
    zz = obs_points[:, :, 2]

    fig = plt.figure()
    ax = fig.gca(projection='3d')

Otthorn's avatar
Otthorn committed
45 46
    sc = ax.scatter(xx, yy, zz, c=mag(B))
    plt.colorbar(cs)
47 48 49 50


def plot_3D(B_fn, obs_points_fn="./data/obs_points.npy"):
    """Plot a 3D vector graph of the B field"""
51 52
    obs_points = np.load(obs_points_fn)
    B = np.load(B_fn)
53

54 55 56
    xx = obs_points[:, :, 0]
    yy = obs_points[:, :, 1]
    zz = obs_points[:, :, 2]
57

58 59
    fig3 = plt.figure()
    ax3 = fig3.gca(projection='3d')
60

61 62 63
    ax3.quiver(xx[::4, ::4], yy[::4, ::4], zz[::4, ::4],
           B[::4, ::4, 0]/float(np.amax(B[::4, ::4, 0])),
           B[::4, ::4, 1]/float(np.amax(B[::4, ::4, 1])),
Otthorn's avatar
Otthorn committed
64
           B[::4, ::4, 2]/float(np.amax(B[::4, ::4, 2])), color='blue', length=0.001)
65

66

67 68
def plot_B(B_fn, obs_points_fn="./data/obs_points.npy"):
    """Plot of 2 vector graph of the B field"""
69 70 71 72 73 74 75 76 77 78 79 80
    obs_points = np.load(obs_points_fn)
    B = np.load(B_fn)

    xx = obs_points[:, :, 0]
    yy = obs_points[:, :, 1]
    zz = obs_points[:, :, 2]

    x = xx[0, :]
    y = yy[:, 0].T

    fig1 = plt.figure()
    ax1 = fig1.gca()
Otthorn's avatar
Otthorn committed
81
    title_txt = 'XY View of the triangle B Field'
82 83 84 85 86 87
    ax1.set_title(title_txt)
    ax1.CS = contourf(x, y, mag(B))
    ax1.quiver(xx[::4, ::4], yy[::4, ::4],
           B[::4, ::4, 0],
           B[::4, ::4, 1], color='white')
    plt.colorbar(ax1.CS)
88

Otthorn's avatar
Otthorn committed
89

90 91
def plot_E(E_fn, obs_points_fn="./data/obs_points.npy"):
    """Plot of 2 vector graph of the E field"""
92 93 94 95 96 97 98 99 100 101 102 103
    obs_points = np.load(obs_points_fn)
    E = np.load(E_fn)

    xx = obs_points[:, :, 0]
    yy = obs_points[:, :, 1]
    zz = obs_points[:, :, 2]

    x = xx[0, :]
    y = yy[:, 0].T

    fig2 = plt.figure()
    ax2 = fig2.gca()
Otthorn's avatar
Otthorn committed
104
    title_txt2 = 'XY View of the triangle E Field'
105 106 107 108 109 110
    ax2.set_title(title_txt2)
    ax2.CS = contourf(x, y, mag(E))
    ax2.quiver(xx[::4, ::4], yy[::4, ::4],
           E[::4, ::4, 0],
           E[::4, ::4, 1], color='white')
    plt.colorbar(ax2.CS)
111

Otthorn's avatar
Otthorn committed
112 113

def plot_coil(coil):
114
    """Plot the vertices of a coil in a 3D graph"""
Otthorn's avatar
Otthorn committed
115 116 117 118 119 120 121 122 123 124 125
    
    vertices = coil.vertices
    xx = vertices[:, 0]
    yy = vertices[:, 1]
    zz = vertices[:, 2]

    fig = plt.figure()
    ax = fig.gca(projection='3d')
    ax.plot(xx, yy, zz, ":o")


126 127
if __name__ == '__main__':

128 129 130
    # plot_B(sys.argv[1])
    # plot_E(sys.argv[2])
    # plot_3D(sys.argv[1])
Otthorn's avatar
Otthorn committed
131
    scatter_3D(sys.argv[1])
132
    # plot_linearity(sys.argv[1])
133
    
134
    plt.show()