Source code for PVGeo._helpers.arrays

__all__ = [
    'get_selected_array_name',
    'get_selected_array_field',
    'copy_arrays_to_point_data',
    'get_numpy_array',
    'get_vtk_array',
    'add_array',
    'get_selected_array',
    'search_for_array',
    'get_all_array_names',
]

import vtk
from vtk.numpy_interface import dataset_adapter as dsa

from . import errors as _helpers



[docs]
def get_selected_array_name(algorithm, idx):
    """Gets the name of the input array for a given index on a VTK algorithm

    Args:
        algorithm (vtkAlgorithm): A vtkAlgorithm class instantiation
        idx (int): the input array index

    Return:
        str : the name of the input array for the given index
    """
    info = algorithm.GetInputArrayInformation(idx)
    return info.Get(vtk.vtkDataObject.FIELD_NAME())




[docs]
def get_selected_array_field(algorithm, idx):
    """Gets the field of the input array for a given index on a VTK algorithm

    Args:
        algorithm (vtkAlgorithm) : A vtkAlgorithm class instantiation
        idx (int) : the input array index

    Return:
        int : the field type of the input array for the given index
    """
    info = algorithm.GetInputArrayInformation(idx)
    return info.Get(vtk.vtkDataObject.FIELD_ASSOCIATION())



def get_field_id_by_name(field):
    """Get the field ID by name."""
    fields = dict(
        point=0,
        pt=0,
        p=0,
        cell=1,
        c=1,
        field=2,
        f=2,
        row=6,
        r=6,
    )
    field = field.lower()
    try:
        return fields[field]
    except KeyError:
        raise _helpers.PVGeoError(
            'Field association not defined. Try inputing `point`, `cell`, `field`, or `row`.'
        )



[docs]
def copy_arrays_to_point_data(pdi, pdo, field):
    """Copies arrays from an input to an ouput's point data.

    Args:
        pdi (vtkDataObject) : The input data object to copy from
        pdo (vtkDataObject) : The output data object to copy over to
        field (int or str) : the field type id or name

    Return:
        vtkDataObject : returns the output data object parameter
    """
    if isinstance(field, str):
        field = get_field_id_by_name(field)
    # Point Data
    if field == 0:
        for i in range(pdi.GetPointData().GetNumberOfArrays()):
            arr = pdi.GetPointData().GetArray(i)
            pdo.GetPointData().AddArray(arr)
    # Cell Data: DO NOT USE
    elif field == 1:
        for i in range(pdi.GetCellData().GetNumberOfArrays()):
            arr = pdi.GetCellData().GetArray(i)
            pdo.GetPointData().AddArray(arr)
    # Field Data:
    elif field == 2:
        for i in range(pdi.GetFieldData().GetNumberOfArrays()):
            arr = pdi.GetFieldData().GetArray(i)
            pdo.GetPointData().AddArray(arr)
    # Row Data:
    elif field == 6:
        for i in range(pdi.GetRowData().GetNumberOfArrays()):
            arr = pdi.GetRowData().GetArray(i)
            pdo.GetPointData().AddArray(arr)
    else:
        raise _helpers.PVGeoError(
            'Field association ({}) not defined. Try inputing Point, Cell, Field, or Row data.'.format(
                field
            )
        )

    # Field Data
    return pdo




[docs]
def get_numpy_array(wpdi, field, name):
    """Grabs an array from vtkDataObject given its name and field association.

    Args:
        wpdi  (wrapped vtkDataObject) : the input data object wrapped using
            vtk dataset adapter
        field (int or str) : the field type id or name
        name (str) : the name of the input array for the given index

    Return:
        numpy.array : a wrapped ``vtkDataArray`` for NumPy
    """
    if isinstance(field, str):
        field = get_field_id_by_name(field)
    if not isinstance(wpdi, vtk.numpy_interface.dataset_adapter.DataObject):
        wpdi = dsa.WrapDataObject(wpdi)
    # Point Data
    if field == 0:
        arr = wpdi.PointData[name]
    # Cell Data:
    elif field == 1:
        arr = wpdi.CellData[name]
    # Field Data:
    elif field == 2:
        arr = wpdi.FieldData[name]
    # Row Data:
    elif field == 6:
        arr = wpdi.RowData[name]
    else:
        raise _helpers.PVGeoError(
            'Field association ({}) not defined. Try inputing Point, Cell, Field, or Row data.'.format(
                field
            )
        )
    return arr




[docs]
def get_vtk_array(pdi, field, name):
    """Grabs an array from vtkDataObject given its name and field association.

    Args:
        pdi  (vtkDataObject) : the input data object
        field (int or str) : the field type id or name
        name (str) : the name of the input array for the given index

    Return:
        vtkDataArray : the array from input field and name
    """
    if isinstance(field, str):
        field = get_field_id_by_name(field)
    # Point Data
    if field == 0:
        arr = pdi.GetPointData().GetArray(name)
    # Cell Data:
    elif field == 1:
        arr = pdi.GetCellData().GetArray(name)
    # Field Data:
    elif field == 2:
        arr = pdi.GetFieldData().GetArray(name)
    # Row Data:
    elif field == 6:
        arr = pdi.GetRowData().GetArray(name)
    else:
        raise _helpers.PVGeoError(
            'Field association ({}) not defined. Try inputing Point, Cell, Field, or Row data.'.format(
                field
            )
        )
    return arr




[docs]
def get_selected_array(algorithm, wpdi, idx):
    """Gets selectected array at index idx wrapped for NumPy

    Args:
        algorithm (vtkAlgorithm) : A vtkAlgorithm class instantiation
        wpdi (wrapped vtkDataObject) : the input data object wrapped using vtk
            dataset adapter
        idx (int) : the input array index

    Return:
        numpy.array : a wrapped ``vtkDataArray`` for NumPy
    """
    name = get_selected_array_name(algorithm, idx)
    field = get_selected_array_field(algorithm, idx)
    return get_array(wpdi, field, name)




[docs]
def add_array(pdo, field, vtkArray):
    """Adds an array to a vtkDataObject given its field association.

    Args:
        pdo (vtkDataObject) : the output data object
        field (int or str) : the field type id or name
        vtkArray (vtkDataArray) : the data array to add to the output

    Return:
        vtkDataObject : the output data object with the data array added
    """
    if isinstance(field, str):
        field = get_field_id_by_name(field)
    # Point Data
    if field == 0:
        pdo.GetPointData().AddArray(vtkArray)
    # Cell Data:
    elif field == 1:
        pdo.GetCellData().AddArray(vtkArray)
    # Field Data:
    elif field == 2:
        pdo.GetFieldData().AddArray(vtkArray)
    # Row Data:
    elif field == 6:
        pdo.GetRowData().AddArray(vtkArray)
    else:
        raise _helpers.PVGeoError(
            'Field association ({}) not defined. Try inputing Point, Cell, Field, or Row data.'.format(
                field
            )
        )
    return pdo



def _get_data(pdi, field):
    """Gets data field from input vtkDataObject"""
    data = None
    if isinstance(field, str):
        field = get_field_id_by_name(field)
    try:
        # Point Data
        if field == 0:
            data = pdi.GetPointData()
        # Cell Data:
        elif field == 1:
            data = pdi.GetCellData()
        # Field Data:
        elif field == 2:
            data = pdi.GetFieldData()
        # Row Data:
        elif field == 6:
            data = pdi.GetRowData()
        else:
            raise _helpers.PVGeoError(
                'Field association ({}) not defined. Try inputing Point, Cell, Field, or Row data.'.format(
                    field
                )
            )
    except AttributeError:
        raise _helpers.PVGeoError(
            'Input data does not have field type `{}`.'.format(field)
        )
    return data


def get_array(pdi, field, name):
    """Gets an array from a vtkDataObject given its field association and name.

    Notes:
        - Point Data: 0
        - Cell Data: 1
        - Field Data: 2
        - Row Data: 6

    Args:
        pdi (vtkDataObject) : the input data object
        field (int or str) : the field type id or name
        name (str) : the data array name

    Return:
        vtkDataObject: the output data object
    """
    if isinstance(field, str):
        field = get_field_id_by_name(field)
    data = _get_data(pdi, field)
    return data.GetArray(name)



[docs]
def search_for_array(pdi, name):
    def _search_field(field):
        data = _get_data(pdi, field)
        for i in range(data.GetNumberOfArrays()):
            if data.GetArrayName(i) == name:
                return data.GetArray(i)
        return None

    fields = [0, 1, 2, 6]
    for field in fields:
        try:
            arr = _search_field(field)
        except _helpers.PVGeoError:
            continue
        if arr is not None:
            # We found it!
            return arr, field

    raise _helpers.PVGeoError('Array `{}` not found in input data.'.format(name))
    return None




[docs]
def get_all_array_names(dataset, field):
    if isinstance(field, str):
        field = get_field_id_by_name(field)
    if not isinstance(dataset, vtk.numpy_interface.dataset_adapter.DataObject):
        wpdi = dsa.WrapDataObject(dataset)
    else:
        wpdi = dataset
    # Point Data
    if field == 0:
        return wpdi.PointData.keys()
    # Cell Data:
    elif field == 1:
        return wpdi.CellData.keys()
    # Field Data:
    elif field == 2:
        return wpdi.FieldData.keys()
    # Row Data:
    elif field == 6:
        return wpdi.RowData.keys()
    else:
        raise _helpers.PVGeoError(
            'Field association ({}) not defined. Try inputing Point, Cell, Field, or Row data.'.format(
                field
            )
        )
    return None