1from pyevtk.hl import _addDataToFile, _appendDataToFile 

2from pyevtk.vtk import VtkFile, VtkImageData 

3 

4 

5def image_to_vtk(path, cell_data, origin=(0.0, 0.0, 0.0), spacing=(1.0, 1.0, 1.0)): 

6 """ 

7 Writes numpy data to VTK 

8 

9 Numpy arrays have to be contiguous in memory - if this is a problem call :func:`numpy.ascontiguousarray` first 

10 

11 Patched version of same pyevtk function that also supports vector-valued data 

12 

13 Args: 

14 path: path with file name, without file ending (.vtk) where data should be stored 

15 cell_data: dictionary, mapping name of the data to a 3D numpy array, or to a 3-tuple of 3D numpy arrays 

16 in case of vector-valued data 

17 origin: 3-tuple describing the origin of the field in 3D 

18 spacing: 3-tuple describing the grid spacing in x,y, z direction 

19 

20 Returns: 

21 path to file that was written 

22 

23 Examples: 

24 >>> from tempfile import TemporaryDirectory 

25 >>> import os 

26 >>> import numpy as np 

27 >>> with TemporaryDirectory() as tmp_dir: 

28 ... path = os.path.join(tmp_dir, 'out') 

29 ... size = (20, 20, 20) 

30 ... res_file = image_to_vtk(path, cell_data={'vector': (np.ones(size), np.ones(size), np.ones(size)), 

31 ... 'scalar': np.zeros(size) 

32 ... }) 

33 """ 

34 

35 # Extract dimensions 

36 start = (0, 0, 0) 

37 end = None 

38 

39 keys = list(cell_data.keys()) 

40 data = cell_data[keys[0]] 

41 if hasattr(data, 'shape'): 

42 end = data.shape 

43 elif isinstance(data, tuple): 

44 shapes = set(d.shape for d in data) 

45 if len(shapes) > 1: 

46 raise ValueError("All components have to have the same shape") 

47 end = shapes.pop() 

48 

49 # Write data to file 

50 w = VtkFile(path, VtkImageData) 

51 w.openGrid(start=start, end=end, origin=origin, spacing=spacing) 

52 w.openPiece(start=start, end=end) 

53 _addDataToFile(w, cell_data, pointData=None) 

54 w.closePiece() 

55 w.closeGrid() 

56 _appendDataToFile(w, cell_data, pointData=None) 

57 w.save() 

58 return w.getFileName()