diff --git a/tests/test_vtkgrid.py b/tests/test_vtkgrid.py index 5f8d340..5804b52 100644 --- a/tests/test_vtkgrid.py +++ b/tests/test_vtkgrid.py @@ -13,8 +13,14 @@ def myVTKGrid(myVtkRectGrid) -> VTKGrid: def test_getDimensions(myVTKGrid, myGrid): assert myVTKGrid.getDimensions(0) == myGrid.nx() + assert myVTKGrid.getNX() == myGrid.nx() + assert myVTKGrid.getDimensions(1) == myGrid.ny() + assert myVTKGrid.getNY() == myGrid.ny() + assert myVTKGrid.getDimensions(2) == myGrid.nz() + assert myVTKGrid.getNZ() == myGrid.nz() + assert myVTKGrid.getDimensions() == [myGrid.nx(),myGrid.ny(),myGrid.nz()] # For getting coordinates @@ -67,6 +73,10 @@ def test_getExtentsZ(myVTKGrid, myGrid): def test_getArrayNone(myVTKGrid): assert myVTKGrid.getArray(None) is None +# test getting list of arrays +def test_getArrayList(myVtkRectGrid_V): + assert VTKGrid(myVtkRectGrid_V).getArrayList() == ["V"] + # test reading scalar data def test_getArrayScalar(myVtkRectGrid_S,myScalarData): s=VTKGrid(myVtkRectGrid_S).getArray("S") diff --git a/vtktonumpy/vtkgrid.py b/vtktonumpy/vtkgrid.py index 53599e1..a9eafa4 100644 --- a/vtktonumpy/vtkgrid.py +++ b/vtktonumpy/vtkgrid.py @@ -73,6 +73,24 @@ def getDimensions(self, d=None): else: return self.dims[d] + def getNX(self) -> int: + """ + $N_x$, the number of cells in $x$ direction + """ + return self.dims[0] + + def getNY(self) -> int: + """ + $N_y$, the number of cells in $y$ direction + """ + return self.dims[1] + + def getNZ(self) -> int: + """ + $N_z$, the number of cells in $z$ direction + """ + return self.dims[2] + # For getting coordinates def getXCoordinates(self) -> np.ndarray: """ @@ -173,6 +191,18 @@ def getExtentsZ(self) -> tuple: vtk_to_numpy(self.vtk_grid.GetZCoordinates()) ) + # For data arrays + def getArrayList(self): + """ + list[str]: A list of cell arrays present in the VTK data + """ + out = [] + cell_data = self.vtk_grid.GetCellData() + for i in range(cell_data.GetNumberOfArrays()): + out.append(cell_data.GetArrayName(i)) + + return out + def getArray(self, name: str) -> np.ndarray: r""" Get VTK array as a numpy array. If no such array exists, return :obj:`None` diff --git a/vtktonumpy/vtkreader.py b/vtktonumpy/vtkreader.py index eadb270..4b23e14 100644 --- a/vtktonumpy/vtkreader.py +++ b/vtktonumpy/vtkreader.py @@ -119,13 +119,7 @@ def removeArray(self, array) -> None: def getArrayList(self): """ - Returns a list of arrays present in the VTK data - - Parameters: - none - - Returns: - list[str]: a list of available cell arrays + list[str]: A list of cell arrays present in the VTK data """ out = [] for i in range(self.vtk_reader.GetNumberOfCellArrays()):