diff --git a/examples/image.py b/examples/image.py index 65f23f8..82425c3 100644 --- a/examples/image.py +++ b/examples/image.py @@ -43,3 +43,14 @@ imageToVTK("./image", cellData = {"pressure" : pressure}, pointData = {"temp" : temp} ) +fluxx = np.random.rand(ncells).reshape( (nx, ny, nz), order='F') +fluxy = np.random.rand(ncells).reshape( (nx, ny, nz), order='F') +fluxz = np.random.rand(ncells).reshape( (nx, ny, nz), order='F') +flux = (fluxx, fluxy, fluxz) + +Efieldx = np.random.rand(npoints).reshape( (nx + 1, ny + 1, nz + 1), order='F') +Efieldy = np.random.rand(npoints).reshape( (nx + 1, ny + 1, nz + 1), order='F') +Efieldz = np.random.rand(npoints).reshape( (nx + 1, ny + 1, nz + 1), order='F') +Efield = (Efieldx,Efieldy,Efieldz) + +imageToVTK("./image", cellData={"flux" : flux}, pointData = {"Efield" : Efieldx} ) \ No newline at end of file diff --git a/pyevtk/hl.py b/pyevtk/hl.py index 3199370..e548060 100644 --- a/pyevtk/hl.py +++ b/pyevtk/hl.py @@ -79,13 +79,14 @@ def imageToVTK(path, origin = (0.0,0.0,0.0), spacing = (1.0,1.0,1.0), cellData = spacing: grid spacing (default = (1,1,1)) cellData: dictionary containing arrays with cell centered data. Keys should be the names of the data arrays. - Arrays must have the same dimensions in all directions and must contain - only scalar data. + Arrays must have the same dimensions in all directions and can contain + scalar data ([n,n,n]) or vector data ([n,n,n],[n,n,n],[n,n,n]). nodeData: dictionary containing arrays with node centered data. Keys should be the names of the data arrays. Arrays must have same dimension in each direction and they should be equal to the dimensions of the cell data plus one and - must contain only scalar data. + can contain scalar data ([n+1,n+1,n+1]) or + vector data ([n+1,n+1,n+1],[n+1,n+1,n+1],[n+1,n+1,n+1]). RETURNS: Full path to saved file. @@ -100,11 +101,17 @@ def imageToVTK(path, origin = (0.0,0.0,0.0), spacing = (1.0,1.0,1.0), cellData = if cellData != None: keys = list(cellData.keys()) data = cellData[keys[0]] - end = data.shape + if hasattr(data,'shape'): + end = data.shape + elif(data[0].ndim==3 and data[1].ndim==3 and data[2].ndim==3): + end = data[0].shape elif pointData != None: keys = list(pointData.keys()) data = pointData[keys[0]] - end = data.shape + if hasattr(data,'shape'): + end = data.shape + elif(data[0].ndim==3 and data[1].ndim==3 and data[2].ndim==3): + end = data[0].shape end = (end[0] - 1, end[1] - 1, end[2] - 1) # Write data to file