Advertisement
Guest User

Untitled

a guest
Apr 6th, 2020
200
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.56 KB | None | 0 0
  1. def tri2triang(tri_):
  2.     import matplotlib.tri as tri
  3.     xy = tri_['vertices']
  4.     return tri.Triangulation(xy[:, 0], xy[:, 1], tri_['triangles'])
  5.  
  6.  
  7. def mesh2triang(mesh):
  8.     import matplotlib.tri as tri
  9.     xy = mesh.coordinates()
  10.     return tri.Triangulation(xy[:, 0], xy[:, 1], mesh.cells())
  11.  
  12.  
  13. def plot_field(u, refine=1,
  14.     ax=None, cax=None, plot_mesh=False, draw_contour=False, tricontourf=False, tripcolor=False,
  15.     exps=None, exp_min=None, exp_max=None, exp_step=None, extend='neither', graph_mesh=None,
  16.     zlim=(None, None), zlim_exp=(None, None), norm_type='lin', mantissa=1, base=10, cmap=None,
  17.     debug_show_mesh_diff=False, elem=None):
  18.  
  19.     import matplotlib.colors as colors
  20.     import matplotlib.ticker as ticker
  21.  
  22.     if ax is None:
  23.         ax = plt.gca()
  24.  
  25.     tol = 10e-6
  26.     exp_min_min = -5
  27.     exp_step_def = 1
  28.  
  29.     f = u.cpp_object()
  30.     mesh = f.function_space().mesh()
  31.    
  32.     if graph_mesh is None:
  33.         graph_mesh = mesh
  34.         for i in range(refine):
  35.             graph_mesh = fc.refine(graph_mesh)
  36.    
  37.  
  38.     if debug_show_mesh_diff:
  39.         fig_tg, ax_tg = init_ax_tg()
  40.  
  41.         ax_tg.triplot(mesh2triang(graph_mesh), zorder=2)
  42.         ax_tg.triplot(mesh2triang(mesh), zorder=2)
  43.         fig_tg.tight_layout()
  44.  
  45.  
  46.    
  47.     #family, shape, r_, k = 'P', 'triangle', 1, 0
  48.     if elem is None:
  49.         family, shape, r_, k = 'P', 'triangle', 2, 0
  50.         elem = fc.FiniteElement(family, shape, r_, k)
  51.     Q = fc.FunctionSpace(graph_mesh, elem)
  52.  
  53.     u.set_allow_extrapolation(True)
  54.     graph_u = fc.interpolate(u, Q)
  55.  
  56.     graph_f = graph_u.cpp_object()
  57.     z = graph_f.compute_vertex_values(graph_mesh)    
  58.    
  59.     graph_triang = mesh2triang(graph_mesh)
  60.     triang = mesh2triang(mesh)
  61.     #from IPython.core.debugger import set_trace; set_trace()
  62.  
  63.     if norm_type == 'lin':
  64. #         n = 40
  65. #         a = zlim[0] if zlim[0] is not None else min(z)
  66. #         b = zlim[1] if zlim[1] is not None else max(z) - tol
  67. #         levels = np.linspace(a, b, n+1)
  68. #         ticks = np.linspace(a, b, n+1)
  69. #         norm = format_ = None
  70.         ticks = None
  71.         levels = 40
  72.         norm = format_ = None
  73.  
  74.     elif norm_type == 'log':
  75.         if exps is None:
  76.             if exp_min is None:
  77.                 if min(z) == 0:
  78.                     exp_min = exp_min_min
  79.                 else:
  80.                     exp_min = calc_m_e(min(z), base)[1] + 1
  81.            
  82.             if exp_max is None:
  83.                 exp_max = calc_m_e(max(z), base)[1] - 1
  84.            
  85.             if exp_step is None:
  86.                 exp_step = exp_step_def
  87.  
  88.             from tools import range_
  89.             exps = range_(exp_min, exp_max, exp_step)
  90.  
  91.         levels = ticks = [mantissa * base**exp for exp in exps]
  92.         norm = colors.SymLogNorm(mantissa * base**exps[0], base=base)
  93.         import math as mt
  94.         format_ = ticker.FuncFormatter(lambda z, pos: '${}^{{{:.4g}}}$'.format(base, mt.log(z, base)))
  95.  
  96.     if plot_mesh:
  97.         ax.triplot(triang, color='#808080', alpha=0.5, zorder=2)
  98.  
  99.     if draw_contour:
  100.         ax.tricontour(graph_triang, z, levels=levels, colors='k', norm=norm, zorder=3)
  101.  
  102.     if tricontourf:
  103.         im = ax.tricontourf(graph_triang, z, levels, norm=norm, cmap=cmap, extend=extend, zorder=1)
  104.  
  105.         plt.colorbar(im, ax=ax, ticks=ticks, extend=extend, format=format_)
  106.  
  107.     if tripcolor:
  108.         tpc = ax.tripcolor(graph_triang, z, norm=norm, cmap=cmap, zorder=1)
  109.  
  110.         plt.colorbar(tpc, ax=ax, ticks=ticks, extend=extend, format=format_)
  111.         #ax.figure.colorbar(tpc)
  112.  
  113.     return graph_mesh
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement