Some Python function

Python code posted
created at 12 Jan 09:40

Edit | Back 
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
 
def plot_field(field, levels=16, fill=True, antialiased=False, comp=None, factor=1.0, **kwargs):
    '''Plot scalar real fields as two-dimensional color map'''
    #TODO documentation

    data = field.array

    if type(comp) == str:
        comp = plask.config.axes.index(comp)

    if type(field.mesh) in (plask.mesh.Regular2D, plask.mesh.Rectilinear2D):
        axis0 = field.mesh.axis0
        axis1 = field.mesh.axis1
        if len(data.shape) == 3:
            if comp is None:
                raise TypeError("specify vector component to plot")
            else:
                data = data[:,:,comp]
        data = data.transpose()
    elif type(field.mesh) in (plask.mesh.Regular3D, plask.mesh.Rectilinear3D):
        axes = [ axis for axis in (field.mesh.axis0, field.mesh.axis1, field.mesh.axis2) if len(axis) > 1 ]
        if len(axes) != 2:
            raise TypeError("'plot_field' only accepts 3D mesh with exactly one axis of size 1")
        axis0, axis1 = axes
        if len(data.shape) == 4:
            if comp is None:
                raise TypeError("specify vector component to plot")
            else:
                data = data[:,:,:,comp]
        data = data.reshape((len(axis0), len(axis1))).transpose()
    else:
        raise NotImplementedError("mesh type not supported")

    if 'cmap' in kwargs and type(kwargs['cmap']) == str: # contourf requires that cmap were cmap instance, not a string
        kwargs = kwargs.copy()
        kwargs['cmap'] = get_cmap(kwargs['cmap'])

    if fill:
        result = contourf(axis0, axis1, data*factor, levels, antialiased=antialiased, **kwargs)
    else:
        if 'colors' not in kwargs and 'cmap' not in kwargs:
            result = contour(axis0, axis1, data, levels, colors='k', antialiased=antialiased, **kwargs)
        else:
            result = contour(axis0, axis1, data, levels, antialiased=antialiased, **kwargs)
    return result


def plot_vectors(field, angles='xy', scale_units='xy', **kwargs):
    '''Plot vector field'''
    #TODO documentation

    m = field.mesh

    if type(m) in (plask.mesh.Regular2D, plask.mesh.Rectilinear2D):
        axis0, axis1 = m.axis0, m.axis1
        i0, i1 = -2, -1
        data = field.array.transpose()
    elif type(m) in (plask.mesh.Regular3D, plask.mesh.Rectilinear3D):
        iaxes = [ iaxis for iaxis in enumerate,(m.axis0, m.axis1, m.axis2) if len(axis) > 1 ]
        if len(axes) != 2:
            raise TypeError("'plot_field' only accepts 3D mesh with exactly one axis of size 1")
        (i0, axis0), (i1, axis1) = iaxes
        data = field.array.reshape((len(axis0), len(axis1), field.array.shape[-1]))[:,:,[i0,i1]].transpose((1,0,2))
    else:
        raise NotImplementedError("mesh type not supported")

    quiver(array(axis0), array(axis1), data[:,:,0], data[:,:,1], angles=angles, scale_units=scale_units, **kwargs)
2.86 KB in 7 ms with coderay