├── README.md ├── license.txt ├── plot_vtk_matplotlib ├── __init__.py └── plot_vtk_matplotlib.py ├── setup.py ├── tutorial └── plot_vtk_matplotlib_tutorial.ipynb └── vector_field.png /README.md: -------------------------------------------------------------------------------- 1 | # Plot VTK files using Matplotlib 2 | 3 | This library allows to easily plot VTK files, for example, from FEniCS or any 4 | other software, in two dimensions using Matplotlib. The VTK file can be based 5 | on an Unstructured or Structured grid, binary or XML. 6 | 7 | By default, the class in this library extracts the data from a slice of the 8 | system, in the x-y plane, thus a range of z values must be specified. 9 | 10 | The data extracted from the VTK file is interpolated using SciPy but an option 11 | to use [natgrid](https://github.com/matplotlib/natgrid) is also available. 12 | 13 | Currently, there are two functions to plot the data: one to plot a colormap of 14 | the sliced data with an optional quiver plot on top and one to make a quiver 15 | plot of the vector field. 16 | 17 | The interpolation is performed in a square grid whose dimensions and spacing 18 | can be tuned. A complete tutorial with most of the capabilities of this library 19 | can be found in the `plot_vtk_matplotlib_tutorial.ipynb` IPython notebook. 20 | This library is mainly oriented to be used with IPython notebooks since we do 21 | not state `matplotlib.pyplot.show()` when generating the graphs, but it can be 22 | specified after the library function calls. 23 | 24 | 25 | ![Quiver plot](vector_field.png) 26 | 27 | # Installation 28 | 29 | Clone the repository, cd into the new directory and as a root user, do 30 | 31 | pip install . 32 | 33 | This will add `plot_vtk_matplotlib` to the Python libraries. 34 | 35 | Alternatively, use 36 | 37 | sudo pip install git+https://github.com/fangohr/plot_vtk_matplotlib.git 38 | 39 | # Software prerequisites 40 | 41 | The library needs `matplotlib`, `scipy` and `vtk` to work. It is recommended 42 | to install the latest versions using `pip`. 43 | 44 | # License 45 | 46 | A BSD license statement can be found in the `license.txt` file. 47 | 48 | # Authors 49 | 50 | David Cortes, Hans Fangohr, University of Southampton (2015) 51 | -------------------------------------------------------------------------------- /license.txt: -------------------------------------------------------------------------------- 1 | Copyright (c) 2015, David Cortes, Hans Fangohr, 2 | University of Southampton 3 | 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | 1. Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | 2. Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 17 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 18 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 19 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR 20 | ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 21 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 22 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND 23 | ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 24 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 25 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | -------------------------------------------------------------------------------- /plot_vtk_matplotlib/__init__.py: -------------------------------------------------------------------------------- 1 | from .plot_vtk_matplotlib import * 2 | -------------------------------------------------------------------------------- /plot_vtk_matplotlib/plot_vtk_matplotlib.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | plot_vtk_matplotlib library for plotting VTK files in two dimensions 4 | using Matplotlib 5 | 6 | Authors: David Cortes, Hans Fangohr 7 | 8 | License: BSD, stated in the Github repository: 9 | https://github.com/fangohr/plot_vtk_matplotlib 10 | 11 | 12 | FUTURE IDEAS:: 13 | 14 | A 3D plot is also possible to do, but the plt.quiver() function 15 | is still under development, and handling thousands of vectors 16 | makes matpltolib very slow. We could look into this in the future. 17 | 18 | """ 19 | 20 | import matplotlib 21 | import matplotlib.pyplot as plt 22 | from mpl_toolkits.axes_grid1 import make_axes_locatable 23 | import vtk 24 | from vtk.util.numpy_support import vtk_to_numpy 25 | try: 26 | import scipy.interpolate 27 | except: 28 | print('Could not load python-scipy. Install it.') 29 | import numpy as np 30 | # import matplotlib.mlab as ml 31 | 32 | import colorsys 33 | 34 | # We can change Matplotlib parameters as 35 | matplotlib.rcParams['font.size'] = 22 36 | # Or as 37 | matplotlib.rcParams.update({'xtick.labelsize': 16}) 38 | matplotlib.rcParams.update({'ytick.labelsize': 16}) 39 | 40 | 41 | class plot_vtk_matplotlib(): 42 | 43 | """ 44 | 45 | Class to load and plot a VTK file (from Fenics or any other simulation 46 | software) in a 2D matplotlib figure. The advantage is that we can save high 47 | quality PDF or SVG files and plot the systems quickly. 48 | 49 | The figure can be a colormap with the z (by default) component of the 50 | vector field (we assume the vectors are normalised) and a quiver plot with 51 | the vector field, which uses the colormap palette, but inverted. Or a 52 | quiver plot. 53 | 54 | To plot this figure, we must specify the path to the VTK file and a range 55 | of z values where the data is going to be extracted from the plot (i.e. we 56 | only plot a slice of the system), thus it is necessary to properly know the 57 | dimensions of the mesh. 58 | 59 | A quiver plot function is also provided, showing only the vector field 60 | 61 | """ 62 | 63 | def __init__(self, vtk_file, z_max, z_min=0, 64 | vtkfiletype='XMLUnstructuredGrid'): 65 | """ 66 | 67 | After initiating this class, call the .extract_data() function 68 | to process the data from the vtk file. After this, it is 69 | possible to plot the system. 70 | 71 | 72 | Attributes: 73 | 74 | vtk_file :: The path to the VTK file (.vtu in Fenics) 75 | 76 | z_max :: Right or maximum limit in the interval of z values 77 | to extract the data 78 | 79 | z_min :: Left or minimum limit in the interval of z_values 80 | to extract the data (zero by default 81 | 82 | vtkfiletype :: The type of the VTK file loaded with this class, By 83 | default we use the Fenics kind of VTK output, which 84 | is an Unstructured Grid XML file. Binary files are 85 | usually StructuredGrid or UnstructuredGrid 86 | Options: 87 | XMLUnstructuredGrid 88 | XMLStructuredGrid 89 | StructuredGrid 90 | UnstructuredGrid 91 | 92 | 93 | It is recommended to give a small interval of z values 94 | in order to avoid overlap of vectors from slices in the z 95 | direction outside the one that is going to be plotted, e.g. 96 | 97 | z_min = 0.38 and z_max = 0.41 98 | 99 | if the system goes from -0.4 to 0.4 100 | 101 | """ 102 | self.vtk_file = vtk_file 103 | self.z_max = z_max 104 | self.z_min = z_min 105 | self.vtkfiletype = vtkfiletype 106 | 107 | def extract_data(self, rotate=None): 108 | """ 109 | 110 | Extract the data from the VTK file using the python-vtk library 111 | The data is passed to numpy arrays for better manipulation 112 | 113 | """ 114 | 115 | # Check the type of file to load it according to its grid structure 116 | # which was specified when calling the class These dictionary entries 117 | # return the corresponding VTK Reader functions, so they can be easily 118 | # called according to the class argument 119 | # 120 | # Rotate argument is a rotation in the x-y plane. To use, set rotate 121 | # equal to an angle in radians. 122 | reader_type = { 123 | 'XMLUnstructuredGrid': lambda: vtk.vtkXMLUnstructuredGridReader(), 124 | 'XMLStructuredGrid': lambda: vtk.vtkXMLStructuredGridReader(), 125 | 'UnstructuredGrid': lambda: vtk.vtkUnstructuredGridReader(), 126 | 'StructuredGrid': lambda: vtk.vtkStructuredGridReader(), 127 | } 128 | 129 | if self.vtkfiletype.startswith('XML'): 130 | # Load the vtk file from the input file 131 | self.reader = reader_type[self.vtkfiletype]() 132 | self.reader.SetFileName(self.vtk_file) 133 | else: 134 | # Load the vtk file from the input file 135 | self.reader = reader_type[self.vtkfiletype]() 136 | self.reader.SetFileName(self.vtk_file) 137 | # For Non XML vtk files: 138 | self.reader.ReadAllVectorsOn() 139 | self.reader.ReadAllScalarsOn() 140 | 141 | self.reader.Update() 142 | 143 | # Get the coordinates of the nodes in the mesh 144 | nodes_vtk_array = self.reader.GetOutput().GetPoints().GetData() 145 | 146 | # Get The vector field (data of every node) 147 | vf_vtk_array = self.reader.GetOutput().GetPointData().GetArray(0) 148 | 149 | # Transform the coordinates of the nodes to a Numpy array and 150 | # save them to the corresponding class objects 151 | 152 | nodes = vtk_to_numpy(nodes_vtk_array) 153 | if rotate: 154 | self.x = nodes[:, 0]*np.cos(rotate) - nodes[:, 1]*np.sin(rotate) 155 | self.y = nodes[:, 0]*np.sin(rotate) + nodes[:, 1]*np.cos(rotate) 156 | self.z = nodes[:, 2] 157 | else: 158 | self.x, self.y, self.z = (nodes[:, 0], 159 | nodes[:, 1], 160 | nodes[:, 2] 161 | ) 162 | 163 | # Transform the magnetisation data to a Numpy array and save 164 | if rotate: 165 | vf = vtk_to_numpy(vf_vtk_array) 166 | vfx = vf[:, 0]*np.cos(rotate) - vf[:, 1]*np.sin(rotate) 167 | vfy = vf[:, 0]*np.sin(rotate) + vf[:, 1]*np.cos(rotate) 168 | vfz = vf[:, 2] 169 | self.vf = np.zeros_like(vf) 170 | self.vf[:, 0] = vfx 171 | self.vf[:, 1] = vfy 172 | self.vf[:, 2] = vfz 173 | else: 174 | self.vf = vtk_to_numpy(vf_vtk_array) 175 | 176 | def interpolate_field(self, 177 | x, y, 178 | x_min, x_max, y_min, y_max, 179 | nx_q=20, ny_q=20, 180 | interpolator='scipy', 181 | interpolator_method='linear' 182 | ): 183 | """ 184 | 185 | Function to interpolate the vector field 186 | 187 | x, y :: Spatial coordinates 188 | x_min, y_min, etc :: Range of interpolation for the vector field 189 | 190 | It returns quivx, quivy and quivz 191 | 192 | """ 193 | 194 | # Interpolate the arrows for the quiver plot if arrow_resolution was 195 | # passed as True 196 | # This option ONLY works with mlab natgrid (CHECK!) --> Try scipy 197 | try: 198 | # Quiver space 199 | xi_q = np.linspace(x_min, x_max, nx_q) 200 | yi_q = np.linspace(y_min, y_max, ny_q) 201 | 202 | # Interpolating with scipy, we can mask the values with 203 | # numpy nans 204 | if interpolator == 'scipy': 205 | xi_q, yi_q = np.meshgrid(xi_q, yi_q) 206 | 207 | quivx = scipy.interpolate.griddata( 208 | (x, y), 209 | self.vf[:, 0][self.data_filter], 210 | (xi_q, yi_q), 211 | method='linear', 212 | fill_value=np.nan 213 | ) 214 | 215 | quivy = scipy.interpolate.griddata( 216 | (x, y), 217 | self.vf[:, 1][self.data_filter], 218 | (xi_q, yi_q), 219 | method='linear', 220 | fill_value=np.nan 221 | ) 222 | 223 | quivz = scipy.interpolate.griddata( 224 | (x, y), 225 | self.vf[:, 2][self.data_filter], 226 | (xi_q, yi_q), 227 | method='linear', 228 | fill_value=np.nan 229 | ) 230 | 231 | # With natgrid we don't need to mask the values 232 | elif interpolator == 'natgrid': 233 | quivx = matplotlib.mlab.griddata(x, y, 234 | self.vf[:, 0][ 235 | self.data_filter], 236 | xi_q, yi_q, 237 | interp=interpolator_method 238 | ) 239 | 240 | quivy = matplotlib.mlab.griddata(x, y, 241 | self.vf[:, 1][ 242 | self.data_filter], 243 | xi_q, yi_q, 244 | interp=interpolator_method 245 | ) 246 | 247 | quivz = matplotlib.mlab.griddata(x, y, 248 | self.vf[:, 2][ 249 | self.data_filter], 250 | xi_q, yi_q, 251 | interp=interpolator_method 252 | ) 253 | 254 | return quivx, quivy, quivz, xi_q, yi_q 255 | 256 | except: 257 | print('Cannot interpolate vector field') 258 | return 259 | 260 | def plot_vtk(self, 261 | x_min, x_max, 262 | y_min, y_max, 263 | v_component='vz', 264 | normalize_data=True, 265 | nx=100, ny=100, 266 | interpolator='scipy', 267 | xlim=None, 268 | ylim=None, 269 | figsize=(8., 8.), 270 | cmap='gist_earth', 271 | hsv_map=False, 272 | cmap_alpha=1., 273 | quiver_map=None, 274 | colorbar=False, 275 | colorbar_label='', 276 | quiver_type='raw_cmap', 277 | quiver_color='k', 278 | pivot='middle', 279 | nx_q=20, 280 | ny_q=20, 281 | frame=True, 282 | predefined_axis=None, 283 | x_label=r'$x$', 284 | y_label=r'$y$', 285 | savefig=None, 286 | interpolator_method=None, 287 | interpolator_hsv_method=None, 288 | interpolator_quiver_method=None, 289 | **quiver_args 290 | ): 291 | """ 292 | 293 | Make a 2D plot from the data extracted of the VTK file, using 294 | a colormap with interpolated values. 295 | 296 | IT IS NECESSARY to run the extract_data() function before. 297 | 298 | If a new range of z_values is required, simply reassign the self.z_min 299 | and self.z_max attributes 300 | 301 | When setting quiver_type as interpolated, the numbers of arrows can be 302 | controled specifying the nx_q and ny_q parameeters, which are the 303 | number of entities along x and y respectively. 304 | 305 | OPTIONS: 306 | 307 | x_min, x_max :: Range of spatial x values to be used in the 2D 308 | plot to interpolate the data for the colormap 309 | 310 | y_min, y_max :: Range of spatial y values to be used in the 2D 311 | plot to interpolate the data for the colormap 312 | 313 | v_component :: Component of the vector field that is going 314 | to be shown as the magnitude of every entity 315 | in a colormap. By default, it is plotted 316 | the z component magnitude of the vectors. 317 | 318 | Options: 319 | 320 | 'vx', 'vy', 'vz' 321 | 322 | normalize_data :: Set False if the colorbar ticks values are in the 323 | range of the real data. By default, the colormap is 324 | normalised from -1 to 1 325 | 326 | nx, ny :: Resolution in the x and y directions 327 | for the interpolations using the data points, 328 | i.e. the number of divisions 329 | between x_min and x_max; y_min and y_max 330 | 331 | interpolator :: The interpolation from the irregular mesh 332 | of the VTK file is done by default using 333 | 'scipy'. It is also possible 334 | to use matplotlib.mlab.griddata passing 335 | the option 'natgrid' 336 | 337 | If an error about not having griddata from 338 | matplotlib, is raised, it can be installed 339 | from the instructions in the print statement 340 | 341 | xlim, ylim :: Plot ranges in the x and y directions, given 342 | as a list with the [min, max] values 343 | 344 | figsize :: Dimensions of the plot as a tuple, 345 | (8, 8) by default 346 | 347 | cmap :: Palette of the colourmap (not considered if 348 | using the hsv_map option) 349 | 350 | hsv_map :: With this option the colormap is going to use 351 | the HSV palette, where the x and y 352 | components of the 353 | vectors are mapped into the Hue values and the z 354 | component is done in the S and V, so that the 355 | maximum z values are shown in white and the 356 | minimum in black. For 2 dimensional vector 357 | fields, this makes the colormap to be only black 358 | (since all the z components are set to zero), 359 | thus this option can be passed as: 360 | 361 | '2d' or '3d'. 362 | 363 | The 2d option set S and V as 1, so the plot 364 | shows the full color. The 3d option makes the 365 | mapping black or white according to the z 366 | component of the field. 367 | 368 | This mapping is useful for showing a vector 369 | field without a quiver plot. 370 | 371 | cmap_alpha :: Transparency value of the colourmap 372 | 373 | quiver_map :: Colour palette of the arrows of the vector 374 | field. By default it is the inverted 375 | palette of cmap 376 | 377 | colorbar :: Set True to plot a color bar with the palette 378 | 379 | colorbar_label :: String with the colorbbar label 380 | (shown rotated in 270 degrees) 381 | 382 | quiver_type :: By default the quiver plot is not interpolated, 383 | it shows all the data points in the specified 384 | spatial ranges (raw data), and it is shown with 385 | a colormap. This option lets the user choose to 386 | interpolate the vector field and if a colormap 387 | or a single color is used. The options are: 388 | 389 | 'interpolated_cmap', 'interpolated_color', 390 | 'raw_cmap', 'raw_color' 391 | 392 | quiver_color :: Arrow color if one of the 'color' options was 393 | specified in the quiver_type argument 394 | 395 | pivot :: By default we make the arrows to be drawn at the 396 | center of the grid nodes. This option is from 397 | the matplotlib quiver function 398 | 399 | nx_q, ny_q :: Resolution in the x and y directions for the 400 | arrows in the quiver plot if one of the 401 | interpolated quiver_type options are passed 402 | (number of divisions between x_min and x_max; 403 | y_min and y_max). By default: 20 x 20 arrows are 404 | drawn 405 | 406 | frame :: Frame of the plot 407 | 408 | predefined_axis :: Can be a predefined matplotlib axis object to 409 | show the plot on it. This is useful to make 410 | a grid of plots 411 | 412 | x_label, y_label :: Axes labels 413 | 414 | savefig :: String with the route and/or name of the 415 | file if it is going to 416 | be saved. The format is obtained from the name, 417 | e.g. 'my_plot.pdf' 418 | 419 | interpolator_method :: Method for scipy or natgrid, default: 'cubic' 420 | or 'nn' 421 | 422 | interpolator_hsv_method :: Method for scipy, for the HSV mapping. 423 | Default: 'linear' 424 | 425 | interpolator_quiver_method :: Method for scipy or natgrid when 426 | interpolating the quiver plot, default: 427 | 'linear' or 'nn' 428 | 429 | **quiver_args :: Any extra keyword arguments for the quiver plot 430 | 431 | TODO: 432 | 433 | Add polar components 434 | 435 | Add titles 436 | 437 | """ 438 | 439 | # Set the interpolator methods according to the arguments 440 | if not interpolator_method: 441 | if interpolator == 'scipy': 442 | interpolator_method = 'cubic' 443 | elif interpolator == 'natgrid': 444 | interpolator_method = 'nn' 445 | else: 446 | print('Specify a valid interpolation method') 447 | return 448 | 449 | # The HSV interpolation is better with linear (cubic messes up things) 450 | if hsv_map: 451 | if not interpolator_hsv_method: 452 | interpolator_hsv_method = 'linear' 453 | 454 | if (quiver_type == 'interpolated_cmap' 455 | or quiver_type == 'interpolated_color'): 456 | if not interpolator_quiver_method: 457 | if interpolator == 'scipy': 458 | interpolator_quiver_method = 'linear' 459 | elif interpolator == 'natgrid': 460 | interpolator_quiver_method = 'nn' 461 | 462 | # Save the array with the filtered data indexes 463 | # (we put it here, since if the z ranges are updated, this will update 464 | # the values) 465 | self.data_filter = (self.z_max > self.z) & (self.z > self.z_min) 466 | if len(np.where(self.data_filter == True)[0]) == 0: 467 | print('No data in specified range!') 468 | return 469 | 470 | # Dictionary to use a specific vector component 471 | comp = {'vx': 0, 'vy': 1, 'vz': 2} 472 | 473 | # Leave only the components between the specified range of z values 474 | x = self.x[self.data_filter] 475 | y = self.y[self.data_filter] 476 | z = self.z[self.data_filter] 477 | 478 | # Define the side of the grid as the radius plus an extra 479 | # space to properly center the system in the plot 480 | # rgrid = x_max + 10 481 | 482 | xi = np.linspace(x_min, x_max, nx) 483 | yi = np.linspace(y_min, y_max, ny) 484 | 485 | # If the HSV map option was passed, make the mapping according 486 | # to the vector field components 487 | if hsv_map: 488 | # Angles of every spin (which defines the colour 489 | # when varying the H value in HSV) 490 | angles = np.arctan2(self.vf[:, comp['vy']][self.data_filter], 491 | self.vf[:, comp['vx']][self.data_filter]) 492 | 493 | # Redefine angles < 0 to got from pi to 2 pi 494 | angles[angles < 0] = angles[angles < 0] + 2 * np.pi 495 | 496 | # The m_z values will be white for m_z = +1 and 497 | # black for m_z = -1 498 | alphas = np.copy(self.vf[:, comp['vz']][self.data_filter]) 499 | alphas_inv = np.copy(self.vf[:, comp['vz']][self.data_filter]) 500 | 501 | if hsv_map == '3d': 502 | alphas[alphas > 0] = 1 503 | alphas[alphas < 0] = alphas[alphas < 0] + 1 504 | 505 | alphas_inv[alphas_inv > 0] = 1 - alphas_inv[alphas_inv > 0] 506 | alphas_inv[alphas_inv < 0] = 1 507 | 508 | # hsv_array = np.array((angles, np.ones_like(angles), alphas)).T 509 | hsv_array = np.array((angles, alphas_inv, alphas)).T 510 | 511 | elif hsv_map == '2d': 512 | hsv_array = np.array((angles, 513 | np.ones_like(angles), 514 | np.ones_like(angles) 515 | ) 516 | ).T 517 | else: 518 | print('Specify a dimension for the HSV mapping') 519 | return 520 | 521 | def convert_to_rgb(a): 522 | return np.array(colorsys.hsv_to_rgb(a[0] / (2 * np.pi), 523 | a[1], 524 | a[2] 525 | ) 526 | ) 527 | 528 | hsv_array = np.array(list(map(convert_to_rgb, hsv_array))) 529 | 530 | # Extract the z_component of the vector field 531 | # (or m_component if specified) to do the colormap 532 | if not hsv_map: 533 | try: 534 | if interpolator == 'natgrid': 535 | # ml.griddata may need the natgrid complement if matplotlib was 536 | # installed with pip. You can get it doing 537 | # git clone https://github.com/matplotlib/natgrid.git 538 | # and then: sudo pip install . 539 | # from the git folder 540 | 541 | zi = matplotlib.mlab.griddata(x, y, 542 | self.vf[:, comp[v_component]][ 543 | self.data_filter], 544 | xi, yi, 545 | interp=interpolator_method) 546 | 547 | elif interpolator == 'scipy': 548 | # Use scipy interpolation 549 | # We need to generate tuples 550 | # Ideas from: 551 | # http://stackoverflow.com/questions/9656489/ 552 | # griddata-scipy-interpolation-not-working-giving-nan 553 | 554 | xi, yi = np.meshgrid(xi, yi) 555 | 556 | # The method can be changed, but the 'nearest' is not 557 | # working well with the values outside the simulation mesh 558 | zi = scipy.interpolate.griddata((x, y), 559 | self.vf[:, comp[v_component]][ 560 | self.data_filter], 561 | (xi, yi), 562 | method=interpolator_method 563 | ) 564 | 565 | # Mask the NaN values (generally they are outside the 566 | # mesh defined in the simulation) so they are not plotted 567 | zi = np.ma.masked_where(np.isnan(zi), zi) 568 | 569 | except Exception('Interpolation Error'): 570 | print('An error ocurred while interpolating the data. ' 571 | 'One of the possible reasosn is that ' 572 | 'matplotlib.mlab.griddata may need the natgrid ' 573 | 'complement in case matplotlib was ' 574 | 'installed with pip. You can get natgrid by doing \n' 575 | 'git clone https://github.com/matplotlib/natgrid.git \n' 576 | 'and then: \n sudo pip install . \n' 577 | 'from the git folder' 578 | ) 579 | return 580 | 581 | # Otherwise, use the HSV colour map for the vector field 582 | else: 583 | xi, yi = np.meshgrid(xi, yi) 584 | zi = scipy.interpolate.griddata((x, y), 585 | hsv_array, 586 | (xi, yi), 587 | method=interpolator_hsv_method, 588 | fill_value=1 589 | ) 590 | 591 | # Quiver data in a dictionary 592 | quiv = {} 593 | 594 | # Interpolate the arrows for the quiver plot if arrow_resolution was 595 | # passed as True 596 | # This option ONLY works with mlab natgrid (CHECK!) --> Try scipy 597 | if (quiver_type == 'interpolated_cmap' 598 | or quiver_type == 'interpolated_colour'): 599 | 600 | (quiv['vx'], 601 | quiv['vy'], 602 | quiv['vz'], 603 | xi_q, yi_q) = self.interpolate_field(x, y, 604 | x_min, x_max, 605 | y_min, y_max, 606 | nx_q=nx_q, ny_q=ny_q, 607 | interpolator=interpolator, 608 | interpolator_method=interpolator_quiver_method 609 | ) 610 | 611 | # --------------------------------------------------------------------- 612 | # Now plot in matplotlib ---------------------------------------------- 613 | # --------------------------------------------------------------------- 614 | # Use a predefined axis if possible 615 | if predefined_axis: 616 | ax = predefined_axis 617 | else: 618 | fig = plt.figure(figsize=figsize, frameon=frame) 619 | ax = fig.add_subplot(111) 620 | 621 | if not hsv_map: 622 | # Plot the colour map with the interpolated values of v_i 623 | ax.pcolormesh(xi, yi, zi, cmap=plt.get_cmap(cmap), vmin=-1, vmax=1, 624 | alpha=cmap_alpha) 625 | else: 626 | # Plot the colour map with the HSV colours 627 | ax.imshow(zi, interpolation='None', 628 | extent=[np.min(xi), np.max(xi), 629 | np.min(yi), np.max(yi)], 630 | vmin=-1, vmax=1, 631 | origin='lower' 632 | ) 633 | if colorbar: 634 | if hsv_map: 635 | cmap_cb = matplotlib.cm.get_cmap(name='hsv') 636 | else: 637 | cmap_cb = matplotlib.cm.get_cmap(name=cmap) 638 | 639 | if normalize_data or hsv_map: 640 | norm = matplotlib.colors.Normalize(-1, 1) 641 | else: 642 | norm = matplotlib.colors.Normalize(vmin=np.min(zi), 643 | vmax=np.max(zi)) 644 | 645 | # Add axes for the colorbar with respect to the top image 646 | divider = make_axes_locatable(ax) 647 | cax = divider.append_axes("right", size="3%", pad=0.05) 648 | 649 | # Colorbar 650 | cbar = matplotlib.colorbar.ColorbarBase(cax, 651 | cmap=cmap_cb, 652 | norm=norm, 653 | # ticks=[-1, 0, 1], 654 | orientation='vertical', 655 | ) 656 | 657 | cbar.set_label(colorbar_label, rotation=270) 658 | 659 | # Label HSV colorbar accordingly 660 | if hsv_map: 661 | cbar.set_ticks([1, 0, -1]) 662 | cbar.set_ticklabels([r'$2\pi$', r'$\pi$', r'$0$']) 663 | # cbar.update_ticks() 664 | 665 | if not quiver_map: 666 | quiver_map = cmap + '_r' 667 | 668 | # Use whole data if the vector field is not inerpolated 669 | if (quiver_type == 'raw_cmap' 670 | or quiver_type == 'raw_colour'): 671 | quiv['vx'] = self.vf[:, 0][self.data_filter], 672 | quiv['vy'] = self.vf[:, 1][self.data_filter] 673 | quiv['vz'] = self.vf[:, 2][self.data_filter] 674 | 675 | xi_q, yi_q = x, y 676 | 677 | if (quiver_type == 'interpolated_cmap' 678 | or quiver_type == 'raw_cmap'): 679 | ax.quiver(xi_q, 680 | yi_q, 681 | quiv['vx'], 682 | quiv['vy'], 683 | # paint the vectors according to the 684 | # component of m_component 685 | quiv[v_component], 686 | cmap=quiver_map, 687 | pivot=pivot, 688 | **quiver_args 689 | ) 690 | elif (quiver_type == 'interpolated_colour' 691 | or quiver_type == 'raw_colour'): 692 | ax.quiver(xi_q, 693 | yi_q, 694 | quiv['vx'], 695 | quiv['vy'], 696 | color=quiver_color, 697 | pivot=pivot, 698 | **quiver_args 699 | ) 700 | elif not quiver_type: 701 | pass 702 | else: 703 | print('Specify an option for the quiver plot') 704 | return 705 | 706 | if not frame: 707 | ax.axis('off') 708 | 709 | if xlim: 710 | ax.set_xlim(xlim) 711 | if ylim: 712 | ax.set_ylim(ylim) 713 | 714 | # Axes labels 715 | ax.set_xlabel(x_label) 716 | ax.set_ylabel(y_label) 717 | 718 | if savefig: 719 | plt.savefig(savefig, bbox_inches='tight') 720 | 721 | # plt.show() 722 | 723 | def plot_quiver(self, 724 | v_component='vz', 725 | xlim=None, ylim=None, 726 | figsize=(8., 8.), 727 | cmap='gist_earth', 728 | alpha=1., 729 | colorbar=None, 730 | colorbar_label='', 731 | interpolator='scipy', 732 | quiver_type='raw_cmap', 733 | quiver_color='k', 734 | pivot='middle', 735 | linewidth=0.7, 736 | x_min=-10, y_min=-10, 737 | x_max=10, y_max=10, 738 | nx_q=20, ny_q=20, 739 | x_label=r'$x$', 740 | y_label=r'$y$', 741 | normalize_data=True, 742 | interpolator_method=None, 743 | predefined_axis=None, 744 | savefig=None, 745 | frame=True, 746 | **kwargs 747 | ): 748 | """ 749 | 750 | Make a 2D quiver plot from the data extracted from the VTK file, using 751 | matplotlib. 752 | 753 | Arrows are marked with a black face 754 | 755 | IT IS NECESSARY to run the extract_data() function before (only once). 756 | r'$v_{' 757 | + '%s' % v_component[1:] 758 | + r'}$' 759 | If a new range of z_values is required, simply reassign 760 | the self.z_min and self.z_max attributes 761 | 762 | 763 | OPTIONS: 764 | 765 | v_component :: 766 | Component of the vector field that is going to 767 | be shown as the magnitude of every entity in a 768 | colormap for the arrows. By default, it is 769 | plotted the z component magnitude of the vectors. 770 | Options: 771 | 772 | 'vx', 'vy', 'vz' 773 | 774 | xlim, ylim :: Plot ranges in the x and y directions, given 775 | as a list with the [min, max] values 776 | 777 | figsize :: Dimensions of the plot as a tuple, 778 | (8, 8) by default 779 | 780 | cmap :: Colour palette of the arrows of the vector 781 | field. 782 | 783 | alpha :: Transparency for the coloured arrows 784 | 785 | colorbar :: Set True to plot a color bar with the palette 786 | 787 | colorbar_label :: String with the colorbbar label 788 | (shown rotated in 270 degrees) 789 | 790 | quiver_type :: By default the quiver plot is not interpolated, 791 | it shows all the data points in the specified 792 | spatial ranges (raw data), and it is shown with 793 | a colormap. This option lets the user choose to 794 | interpolate the vector field and if a colormap 795 | or a single color is used. The options are: 796 | 797 | 'interpolated_cmap', 'interpolated_color', 798 | 'raw_cmap', 'raw_color' 799 | 800 | quiver_color :: Arrow color if one of the 'color' options was 801 | specified in the quiver_type argument 802 | 803 | pivot :: By default we make the arrows to be drawn at the 804 | center of the grid nodes. This option is from 805 | the matplotlib quiver function 806 | 807 | linewidth :: Arrows line width 808 | 809 | x_min, x_max :: Range of spatial x values to be used in the 810 | quiver plot if one of the interpolated options 811 | was passed to the quiver_type argument 812 | 813 | y_min, y_max :: Range of spatial y values to be used in the 814 | quiver plot if one of the interpolated options 815 | was passed to the quiver_type argument 816 | 817 | nx_q, ny_q :: Resolution in the x and y directions for the 818 | quiver plot if it was interpolated 819 | 820 | normalize_data :: Set False if the colorbar ticks values are in the 821 | range of the real data. By default, the colormap is 822 | normalised from -1 to 1r'$v_{' 823 | + '%s' % v_component[1:] 824 | + r'}$' 825 | 826 | interpolator_method :: Method for scipy or natgrid for interpolating the, 827 | quiver data, default: 'linear' or 'nn' 828 | 829 | savefig :: String with the route and/or name of the 830 | file if it is going to 831 | be saved. The format is obtained from the name, 832 | e.g. 'my_plot.pdf' 833 | 834 | **kwargs :: Extra keyword arguments for the quiver plot 835 | (see Matplotlib doc) 836 | 837 | TODO: 838 | 839 | Add polar components 840 | Add titles 841 | 842 | """ 843 | 844 | # Set the interpolator methods according to the arguments 845 | if (quiver_type == 'interpolated_cmap' 846 | or quiver_type == 'interpolated_colour'): 847 | if not interpolator_method: 848 | if interpolator == 'scipy': 849 | interpolator_method = 'linear' 850 | elif interpolator == 'natgrid': 851 | interpolator_method = 'nn' 852 | 853 | # Save the array with the filtered data indexes 854 | # (we put it here, since if the z ranges are updated, this will update 855 | # the values) 856 | self.data_filter = (self.z_max > self.z) & (self.z > self.z_min) 857 | if len(np.where(self.data_filter == True)[0]) == 0: 858 | print('No data in specified range!') 859 | return 860 | 861 | # Leave only the components between the specified z_range 862 | x = self.x[self.data_filter] 863 | y = self.y[self.data_filter] 864 | z = self.z[self.data_filter] 865 | 866 | # Interpolate the data into a square grid (nx x ny) 867 | # to plot with matplotlib 868 | # This is the Grid Resolution 869 | comp = {'vx': 0, 'vy': 1, 'vz': 2} 870 | 871 | # Quiver data in a dictionary 872 | quiv = {} 873 | 874 | # Interpolate if necessary. Notice that x and y variables CHANGE here 875 | # according to the interpolator (see the function) 876 | if (quiver_type == 'interpolated_cmap' 877 | or quiver_type == 'interpolated_colour'): 878 | 879 | (quiv['vx'], 880 | quiv['vy'], 881 | quiv['vz'], 882 | x, y) = self.interpolate_field(x, y, 883 | x_min, x_max, 884 | y_min, y_max, 885 | nx_q=nx_q, ny_q=ny_q, 886 | interpolator=interpolator, 887 | interpolator_method=interpolator_method 888 | ) 889 | 890 | # Use whole data if the vector field is not inerpolated 891 | if (quiver_type == 'raw_cmap' 892 | or quiver_type == 'raw_colour'): 893 | quiv['vx'] = self.vf[:, 0][self.data_filter], 894 | quiv['vy'] = self.vf[:, 1][self.data_filter] 895 | quiv['vz'] = self.vf[:, 2][self.data_filter] 896 | 897 | # Now plot in matplotlib --------------------------------------------- 898 | if not predefined_axis: 899 | f = plt.figure(figsize=figsize, frameon=frame) 900 | ax = f.add_subplot(111) 901 | 902 | else: 903 | ax = predefined_axis 904 | 905 | if (quiver_type == 'interpolated_cmap' 906 | or quiver_type == 'raw_cmap'): 907 | ax.quiver(x, 908 | y, 909 | quiv['vx'], 910 | quiv['vy'], 911 | # paint the vectors according to the 912 | # component of m_component 913 | quiv[v_component], 914 | cmap=cmap, 915 | alpha=alpha, 916 | linewidth=linewidth, 917 | edgecolor='k', 918 | facecolor=None, 919 | pivot=pivot, 920 | **kwargs 921 | ) 922 | 923 | if (quiver_type == 'interpolated_colour' 924 | or quiver_type == 'raw_colour'): 925 | 926 | ax.quiver(x, 927 | y, 928 | quiv['vx'], 929 | quiv['vy'], 930 | alpha=alpha, 931 | linewidth=linewidth, 932 | color=quiver_color, 933 | pivot=pivot, 934 | **kwargs 935 | ) 936 | 937 | # plt.tight_layout() 938 | 939 | if xlim: 940 | plt.xlim(xlim) 941 | if ylim: 942 | plt.ylim(ylim) 943 | 944 | if x_label: 945 | plt.xlabel(x_label) 946 | if y_label: 947 | plt.ylabel(y_label) 948 | 949 | if not frame: 950 | ax.axis('off') 951 | # This can help to remove all whitespace: 952 | # Get the figure extent and then use it in bbox_inches 953 | # extent = plt.gca().get_window_extent().transformed(plt.gcf().dpi_scale_trans.inverted()) 954 | 955 | if normalize_data: 956 | norm = matplotlib.colors.Normalize(vmin=-1, vmax=1) 957 | else: 958 | norm = matplotlib.colors.Normalize(vmin=np.min(quiv[v_component]), 959 | vmax=np.max(quiv[v_component])) 960 | if colorbar: 961 | divider = make_axes_locatable(ax) 962 | cax = divider.append_axes("right", size="3%", pad=0.05) 963 | 964 | # Colorbar. 965 | cbar = matplotlib.colorbar.ColorbarBase(cax, 966 | cmap=cmap, 967 | norm=norm, 968 | # ticks=[-1, 0, 1], 969 | orientation='vertical', 970 | ) 971 | 972 | # cb = plt.colorbar() 973 | cbar.set_label(colorbar_label, 974 | rotation=270, 975 | labelpad=10 976 | ) 977 | # plt.axes().set_aspect('equal', 'datalim') 978 | 979 | if savefig: 980 | plt.savefig(savefig, bbox_inches='tight') 981 | # plt.savefig(savefig, bbox_inches=extent) 982 | 983 | # plt.show() 984 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | version='0.1', 5 | name="plot_vtk_matplotlib", 6 | packages=['plot_vtk_matplotlib'], 7 | license='BSD', 8 | description='Library to plot VTK files in 2D with Matplotlib', 9 | # We comment this to avoid upgrading matplotlib when 10 | # instaling the library. Could be used in the future. 11 | # install_requires=['scipy', 'matplotlib'], 12 | author_email='d.i.cortes@soton.ac.uk' 13 | ) 14 | -------------------------------------------------------------------------------- /vector_field.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/fangohr/plot_vtk_matplotlib/501ffa6df02515308660282a6c9005009edc605f/vector_field.png --------------------------------------------------------------------------------