plaid.utils.cgns_vtk¶

def CGNSBaseExtractGlobals(baseNode: list) -> dict:
    """Extract global fields from a CGNSBase_t node as a dictionary of name -> numpy array.

    Arguments:
        baseNode (list): CGNS ``CGNSBase_t`` node.

    Returns:
        dict: A dictionary mapping field names to numpy arrays.

    """
    globals = {}
    for xNode in baseNode[2]:
        if xNode[1] is not None:
            globals[xNode[0]] = np.asarray(xNode[1])
    return globals

def CGNSTreeToVtk(treeNode: list) -> Any:
    """Convert a full CGNS tree to VTK objects, one per base, and return either the single base or a multi-block of bases.

    Arguments:
        treeNode (list): CGNS tree as read by cgns_tree_from_json_payload.

    Returns:
        vtkStructuredGrid, vtkUnstructuredGrid, or vtkMultiBlockDataSet: the VTK
        representation of the tree. A single-zone base returns the zone object;
        a multi-zone base returns one block per zone; a multi-base tree returns
        one block per base.
    """
    _, _, _, _, vtkMultiBlockDataSet, numpy_support = _import_vtk_for_direct_cgns()

    bases = _cgns_children_by_label(treeNode, "CGNSBase_t")
    globals: dict[str, Any] = {}

    for baseNode in bases:
        if baseNode[0] == "Global":
            globals.update(CGNSBaseExtractGlobals(baseNode))

    baseVtkObjects = []
    basenames = []
    for baseNode in bases:
        if baseNode[0] == "Global":
            continue
        baseVtkObjects.append(CGNSBaseToVtk(baseNode))
        basenames.append(baseNode[0])

    new_output = baseVtkObjects[0]
    # Add field Data
    field_data = new_output.GetFieldData()

    for key, value in globals.items():
        if value.dtype == "|S1":
            from vtkmodules.vtkCommonCore import vtkStringArray

            labels = vtkStringArray()
            labels.SetName(key)
            labels.SetNumberOfValues(len(value))
            for i, v in enumerate(value):
                labels.SetValue(i, v)
            field_data.AddArray(labels)
            continue

        array = numpy_support.numpy_to_vtk(value)
        array.SetName(key)
        field_data.AddArray(array)

    if len(baseVtkObjects) == 1:
        return baseVtkObjects[0]

    multiBlock = vtkMultiBlockDataSet()
    multiBlock.SetNumberOfBlocks(len(baseVtkObjects))
    for i, (name, zoneVtkObject) in enumerate(zip(basenames, baseVtkObjects)):
        multiBlock.SetBlock(i, zoneVtkObject)
        multiBlock.GetMetaData(i).Set(multiBlock.NAME(), name)
    return multiBlock

def CGNSBaseToVtk(baseNode: list) -> Any:
    """Convert a CGNSBase_t node directly to a VTK object.

    This function intentionally bypasses Muscat mesh/conversion functions.  It
    reads the CGNS/Python tree node lists directly and creates native VTK data
    objects using only VTK and NumPy.

    Args:
        baseNode (list): CGNS ``CGNSBase_t`` node.

    Returns:
        vtkStructuredGrid, vtkUnstructuredGrid, or vtkMultiBlockDataSet: the VTK
        representation of the base. A single-zone base returns the zone object;
        a multi-zone base returns one block per zone.
    """
    if (
        not isinstance(baseNode, list)
        or len(baseNode) < 4
        or baseNode[3] != "CGNSBase_t"
    ):
        raise ValueError("CGNSBaseToVtk expects a CGNSBase_t node")
    if baseNode[1] is None:
        raise ValueError(f"CGNS base '{baseNode[0]}' has no base dimensionality value")

    baseDims = np.asarray(baseNode[1], dtype=int).ravel()
    physicalDim = int(baseDims[1]) if baseDims.size > 1 else 3
    zones = _cgns_children_by_label(baseNode, "Zone_t")
    if not zones:
        raise ValueError(f"CGNS base '{baseNode[0]}' has no Zone_t children")

    zoneVtkObjects = []
    for zoneNode in zones:
        zoneType = (
            _cgns_value_as_string(_cgns_child_by_name(zoneNode, "ZoneType"))
            or "Unstructured"
        )
        if zoneType == "Structured":
            zoneVtkObjects.append(_cgns_structured_zone_to_vtk(zoneNode, physicalDim))
        elif zoneType == "Unstructured":
            zoneVtkObjects.append(_cgns_unstructured_zone_to_vtk(zoneNode, physicalDim))
        else:
            raise NotImplementedError(
                f"CGNS ZoneType '{zoneType}' is not supported by direct VTK conversion"
            )

    if len(zoneVtkObjects) == 1:
        return zoneVtkObjects[0]

    _, _, _, _, vtkMultiBlockDataSet, _ = _import_vtk_for_direct_cgns()
    multiBlock = vtkMultiBlockDataSet()
    multiBlock.SetNumberOfBlocks(len(zoneVtkObjects))
    for i, (zoneNode, zoneVtkObject) in enumerate(zip(zones, zoneVtkObjects)):
        multiBlock.SetBlock(i, zoneVtkObject)
        multiBlock.GetMetaData(i).Set(multiBlock.NAME(), zoneNode[0])
    return multiBlock