Coverage for pySDC/helpers/vtkIO.py: 0%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3""" 

4Helper functions for VTK files IO (to be used with Paraview or PyVista) 

5""" 

6import os 

7import vtk 

8from vtkmodules.util import numpy_support 

9import numpy as np 

10 

11 

12def writeToVTR(fileName: str, data, coords, varNames): 

13 """ 

14 Write a data array containing variables from a 3D rectilinear grid into a VTR file. 

15 

16 Parameters 

17 ---------- 

18 fileName : str 

19 Name of the VTR file, can be with or without the .vtr extension. 

20 data : np.4darray 

21 Array containing all the variables with [nVar, nX, nY, nZ] shape. 

22 coords : list[np.1darray] 

23 Coordinates in each direction. 

24 varNames : list[str] 

25 Variable names. 

26 

27 Returns 

28 ------- 

29 fileName : str 

30 Name of the VTR file. 

31 """ 

32 data = np.asarray(data) 

33 nVar, *gridSizes = data.shape 

34 

35 assert len(gridSizes) == 3, "function can be used only for 3D grid data" 

36 assert nVar == len(varNames), "varNames must have as many variable as data" 

37 assert [len(np.ravel(coord)) for coord in coords] == gridSizes, "coordinate size incompatible with data shape" 

38 

39 nX, nY, nZ = gridSizes 

40 vtr = vtk.vtkRectilinearGrid() 

41 vtr.SetDimensions(nX, nY, nZ) 

42 

43 def vect(x): 

44 return numpy_support.numpy_to_vtk(num_array=x, deep=True, array_type=vtk.VTK_FLOAT) 

45 

46 x, y, z = coords 

47 vtr.SetXCoordinates(vect(x)) 

48 vtr.SetYCoordinates(vect(y)) 

49 vtr.SetZCoordinates(vect(z)) 

50 

51 def field(u): 

52 return numpy_support.numpy_to_vtk(num_array=u.ravel(order='F'), deep=True, array_type=vtk.VTK_FLOAT) 

53 

54 pointData = vtr.GetPointData() 

55 for name, u in zip(varNames, data): 

56 uVTK = field(u) 

57 uVTK.SetName(name) 

58 pointData.AddArray(uVTK) 

59 

60 writer = vtk.vtkXMLRectilinearGridWriter() 

61 if not fileName.endswith(".vtr"): 

62 fileName += ".vtr" 

63 writer.SetFileName(fileName) 

64 writer.SetInputData(vtr) 

65 writer.Write() 

66 

67 return fileName 

68 

69 

70def readFromVTR(fileName: str): 

71 """ 

72 Read a VTR file into a numpy 4darray 

73 

74 Parameters 

75 ---------- 

76 fileName : str 

77 Name of the VTR file, can be with or without the .vtr extension. 

78 

79 Returns 

80 ------- 

81 data : np.4darray 

82 Array containing all the variables with [nVar, nX, nY, nZ] shape. 

83 coords : list[np.1darray] 

84 Coordinates in each direction. 

85 varNames : list[str] 

86 Variable names. 

87 """ 

88 if not fileName.endswith(".vtr"): 

89 fileName += ".vtr" 

90 assert os.path.isfile(fileName), f"{fileName} does not exist" 

91 

92 reader = vtk.vtkXMLRectilinearGridReader() 

93 reader.SetFileName(fileName) 

94 reader.Update() 

95 

96 vtr = reader.GetOutput() 

97 gridSizes = vtr.GetDimensions() 

98 assert len(gridSizes) == 3, "can only read 3D data" 

99 

100 def vect(x): 

101 return numpy_support.vtk_to_numpy(x) 

102 

103 coords = [vect(vtr.GetXCoordinates()), vect(vtr.GetYCoordinates()), vect(vtr.GetZCoordinates())] 

104 pointData = vtr.GetPointData() 

105 varNames = [pointData.GetArrayName(i) for i in range(pointData.GetNumberOfArrays())] 

106 data = [numpy_support.vtk_to_numpy(pointData.GetArray(name)).reshape(gridSizes, order="F") for name in varNames] 

107 data = np.array(data) 

108 

109 return data, coords, varNames