Source code for PVGeo.interface

"""The ``interface`` module provides functions to convert/cast between common
VTK and NumPy/Pandas data types. These methods provide a simple to use interface
for VTK data types so that users can make changes to VTK data strucutres via
Python data structures that are a bit easier to perform numerical operations
upon.
"""


__all__ = [
    'get_vtk_type',
    'convert_string_array',
    'convert_array',
    'data_frame_to_table',
    'table_to_data_frame',
    'place_array_in_table',
    'get_dtypes',
    'points_to_poly_data',
    'add_arrays_from_data_frame',
    'convert_cell_conn',
    'get_array',
    'get_data_dict',
]


__displayname__ = 'Interface'


import numpy as np
import pandas as pd
import vtk
from vtk.util import numpy_support as nps

import pyvista as pv
from pyvista.utilities import convert_string_array, get_vtk_type

from . import _helpers


[docs]def convert_array(arr, name='Data', deep=0, array_type=None, pdf=False): """A helper to convert a NumPy array to a vtkDataArray or vice versa Args: arr (ndarray or vtkDataArry) : A numpy array or vtkDataArry to convert name (str): the name of the data array for VTK deep (bool, int): if input is numpy array then deep copy values pdf (bool): if input is vtkDataArry, make a pandas DataFrame of the array Return: vtkDataArray, ndarray, or DataFrame: the converted array (if input is a NumPy ndaray then returns ``vtkDataArray`` or is input is ``vtkDataArray`` then returns NumPy ``ndarray``). If pdf==True and the input is ``vtkDataArry``, return a pandas DataFrame. """ num_data = pv.convert_array(arr, name=name, deep=deep, array_type=array_type) if not isinstance(num_data, np.ndarray): return num_data if not pdf: return num_data return pd.DataFrame(data=num_data, columns=[arr.GetName()])
[docs]def data_frame_to_table(df, pdo=None): """Converts a pandas DataFrame to a vtkTable""" if not isinstance(df, pd.DataFrame): raise _helpers.PVGeoError('Input is not a pandas DataFrame') if pdo is None: pdo = pv.Table(df) else: pdo.DeepCopy(pv.Table(df)) return pdo
[docs]def table_to_data_frame(table): """Converts a vtkTable to a pandas DataFrame""" if not isinstance(table, vtk.vtkTable): raise _helpers.PVGeoError('Input is not a vtkTable') if not isinstance(table, pv.Table): table = pv.Table(table) return table.to_pandas()
[docs]def place_array_in_table(ndarr, titles, pdo): """Takes a 1D/2D numpy array and makes a vtkTable of it Args: ndarr (numpy.ndarray) : The 1D/2D array to be converted to a table titles (list or tuple): The titles for the arrays in the table. Must have same number of elements as columns in input ndarray pdo (vtkTable) : The output data object pointer Return: vtkTable : returns the same input pdo table """ # Put columns into table if len(np.shape(ndarr)) > 2: raise _helpers.PVGeoError('Input np.ndarray must be 1D or 2D to be converted to vtkTable.') if len(np.shape(ndarr)) == 1: # First check if it is an array full of tuples (varying type) if isinstance(ndarr[0], (tuple, np.void)): for i, title in enumerate(titles): place_array_in_table(ndarr['f%d' % i], title, pdo) return pv.wrap(pdo) # Otherwise it is just a 1D array which needs to be 2D else: ndarr = np.reshape(ndarr, (-1, 1)) cols = np.shape(ndarr)[1] for i in range(cols): VTK_data = convert_array(ndarr[:,i]) VTK_data.SetName(titles[i]) pdo.AddColumn(VTK_data) return pv.wrap(pdo)
[docs]def get_dtypes(dtype='', endian=None): """This converts char dtypes and an endian to a numpy and VTK data type. Return: tuple (numpy.dtype, int): the numpy data type and the integer type id specified in vtkType.h for VTK data types """ # If native `@` was chosen then do not pass an endian if endian == '@': #print('WARNING: Native endianness no longer supported for packed binary reader. Please chose `>` or `<`. This defaults to big `>`.') endian = '' # No endian specified: elif endian is None: endian = '' # Get numpy and VTK data types and return them both if dtype == 'd': vtktype = vtk.VTK_DOUBLE elif dtype == 'f': vtktype = vtk.VTK_FLOAT elif dtype == 'i': vtktype = vtk.VTK_INT else: raise _helpers.PVGeoError('dtype \'%s\' unknown:' % dtype) # Return data types dtype = np.dtype('%s%s' % (endian, dtype)) return dtype, vtktype
[docs]def points_to_poly_data(points, copy_z=False): """Create ``vtkPolyData`` from a numpy array of XYZ points. If the points have more than 3 dimensions, then all dimensions after the third will be added as attributes. Assume the first three dimensions are the XYZ coordinates. Args: points (np.ndarray or pandas.DataFrame): The points and pointdata copy_z (bool): A flag on whether to append the z values as a PointData array Return: vtkPolyData : points with point-vertex cells """ # This prevents an error that occurs when only one point is passed if points.ndim < 2: points = points.reshape((1,-1)) keys = ['Field %d' % i for i in range(points.shape[1] - 3)] # Check if input is anything other than a NumPy array and cast it # e.g. you could send a Pandas dataframe if not isinstance(points, np.ndarray): if isinstance(points, pd.DataFrame): # If a pandas data frame, lets grab the keys keys = points.keys()[3::] points = np.array(points) # If points are not 3D if points.shape[1] < 2: raise RuntimeError('Points must be 3D. Try adding a third dimension of zeros.') atts = points[:, 3::] points = points[:, 0:3].astype(np.float) # Create polydata pdata = pv.PolyData(points) # Add attributes if given scalSet = False for i, key in enumerate(keys): data = convert_array(atts[:, i], name=key) pdata.GetPointData().AddArray(data) if not scalSet: pdata.GetPointData().SetActiveScalars(key) scalSet = True if copy_z: z = convert_array(points[:, 2], name='Elevation') pdata.GetPointData().AddArray(z) return pv.wrap(pdata)
[docs]def add_arrays_from_data_frame(pdo, field, df): """Add all of the arrays from a given data frame to an output's data""" for key in df.keys(): VTK_data = convert_array(df[key].values, name=key) _helpers.add_array(pdo, field, VTK_data) return pv.wrap(pdo)
[docs]def convert_cell_conn(cell_connectivity): """Converts cell connectivity arrays to a cell matrix array that makes sense for VTK cell arrays. """ cellsMat = np.concatenate( ( np.ones((cell_connectivity.shape[0], 1), dtype=np.int64)*cell_connectivity.shape[1], cell_connectivity ), axis=1).ravel() return nps.numpy_to_vtk(cellsMat, deep=True, array_type=vtk.VTK_ID_TYPE)
[docs]def get_array(dataset, name, vtk_object=False): """Given an input dataset, this will return the named array as a NumPy array or a vtkDataArray if spceified """ arr, field = _helpers.search_for_array(dataset, name) if vtk_object: return arr return convert_array(arr)
[docs]def get_data_dict(dataset, field='cell'): """Given an input dataset, this will return all the arrays in that object's cell/point/field/row data as named NumPy arrays in a dictionary. """ data = {} for key in _helpers.get_all_array_names(dataset, field): data[key] = np.array(_helpers.get_numpy_array(dataset, field, key)) return data