Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- def tri2triang(tri_):
- import matplotlib.tri as tri
- xy = tri_['vertices']
- return tri.Triangulation(xy[:, 0], xy[:, 1], tri_['triangles'])
- def mesh2triang(mesh):
- import matplotlib.tri as tri
- xy = mesh.coordinates()
- return tri.Triangulation(xy[:, 0], xy[:, 1], mesh.cells())
- def plot_field(u, refine=1,
- ax=None, cax=None, plot_mesh=False, draw_contour=False, tricontourf=False, tripcolor=False,
- exps=None, exp_min=None, exp_max=None, exp_step=None, extend='neither', graph_mesh=None,
- zlim=(None, None), zlim_exp=(None, None), norm_type='lin', mantissa=1, base=10, cmap=None,
- debug_show_mesh_diff=False, elem=None):
- import matplotlib.colors as colors
- import matplotlib.ticker as ticker
- if ax is None:
- ax = plt.gca()
- tol = 10e-6
- exp_min_min = -5
- exp_step_def = 1
- f = u.cpp_object()
- mesh = f.function_space().mesh()
- if graph_mesh is None:
- graph_mesh = mesh
- for i in range(refine):
- graph_mesh = fc.refine(graph_mesh)
- if debug_show_mesh_diff:
- fig_tg, ax_tg = init_ax_tg()
- ax_tg.triplot(mesh2triang(graph_mesh), zorder=2)
- ax_tg.triplot(mesh2triang(mesh), zorder=2)
- fig_tg.tight_layout()
- #family, shape, r_, k = 'P', 'triangle', 1, 0
- if elem is None:
- family, shape, r_, k = 'P', 'triangle', 2, 0
- elem = fc.FiniteElement(family, shape, r_, k)
- Q = fc.FunctionSpace(graph_mesh, elem)
- u.set_allow_extrapolation(True)
- graph_u = fc.interpolate(u, Q)
- graph_f = graph_u.cpp_object()
- z = graph_f.compute_vertex_values(graph_mesh)
- graph_triang = mesh2triang(graph_mesh)
- triang = mesh2triang(mesh)
- #from IPython.core.debugger import set_trace; set_trace()
- if norm_type == 'lin':
- # n = 40
- # a = zlim[0] if zlim[0] is not None else min(z)
- # b = zlim[1] if zlim[1] is not None else max(z) - tol
- # levels = np.linspace(a, b, n+1)
- # ticks = np.linspace(a, b, n+1)
- # norm = format_ = None
- ticks = None
- levels = 40
- norm = format_ = None
- elif norm_type == 'log':
- if exps is None:
- if exp_min is None:
- if min(z) == 0:
- exp_min = exp_min_min
- else:
- exp_min = calc_m_e(min(z), base)[1] + 1
- if exp_max is None:
- exp_max = calc_m_e(max(z), base)[1] - 1
- if exp_step is None:
- exp_step = exp_step_def
- from tools import range_
- exps = range_(exp_min, exp_max, exp_step)
- levels = ticks = [mantissa * base**exp for exp in exps]
- norm = colors.SymLogNorm(mantissa * base**exps[0], base=base)
- import math as mt
- format_ = ticker.FuncFormatter(lambda z, pos: '${}^{{{:.4g}}}$'.format(base, mt.log(z, base)))
- if plot_mesh:
- ax.triplot(triang, color='#808080', alpha=0.5, zorder=2)
- if draw_contour:
- ax.tricontour(graph_triang, z, levels=levels, colors='k', norm=norm, zorder=3)
- if tricontourf:
- im = ax.tricontourf(graph_triang, z, levels, norm=norm, cmap=cmap, extend=extend, zorder=1)
- plt.colorbar(im, ax=ax, ticks=ticks, extend=extend, format=format_)
- if tripcolor:
- tpc = ax.tripcolor(graph_triang, z, norm=norm, cmap=cmap, zorder=1)
- plt.colorbar(tpc, ax=ax, ticks=ticks, extend=extend, format=format_)
- #ax.figure.colorbar(tpc)
- return graph_mesh
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement