# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright 2023-2024 TotalEnergies.
# SPDX-FileContributor: Martin Lemay
from typing import Union, cast
from vtkmodules.vtkCommonDataModel import (
vtkCompositeDataSet,
vtkDataObject,
vtkDataObjectTreeIterator,
vtkMultiBlockDataSet,
)
__doc__ = """Functions to explore and process multiblock inspector tree."""
[docs]
def getBlockName(input: Union[vtkMultiBlockDataSet, vtkCompositeDataSet]) -> str:
"""Get the name of input block.
If input is a vtkMultiBlockDataSet or vtkCompositeDataSet, returns the name
of the lowest level unique child block.
Args:
input (vtkMultiBlockDataSet | vtkCompositeDataSet): input multi block object.
Returns:
str: name of the block in the tree.
"""
iter: vtkDataObjectTreeIterator = vtkDataObjectTreeIterator()
iter.SetDataSet(input)
iter.VisitOnlyLeavesOff()
iter.GoToFirstItem()
blockName: str = "Block"
while iter.GetCurrentDataObject() is not None:
blockName = iter.GetCurrentMetaData().Get(vtkMultiBlockDataSet.NAME())
block: vtkDataObject = iter.GetCurrentDataObject()
nbBlocks: int = 99
if isinstance(block, vtkMultiBlockDataSet):
block1: vtkMultiBlockDataSet = cast(vtkMultiBlockDataSet, block)
nbBlocks = block1.GetNumberOfBlocks()
# stop if multiple children
if nbBlocks > 1:
break
iter.GoToNextItem()
return blockName
[docs]
def getBlockNameFromIndex(
input: Union[vtkMultiBlockDataSet, vtkCompositeDataSet], index: int
) -> str:
"""Get the name of a block from input multiblock and block flat index.
Args:
input (vtkMultiBlockDataSet | vtkCompositeDataSet): input multi block object.
index (int): flat index of the block to get the name
Returns:
str: name of the block in the tree.
"""
iter: vtkDataObjectTreeIterator = vtkDataObjectTreeIterator()
iter.SetDataSet(input)
iter.VisitOnlyLeavesOff()
iter.GoToFirstItem()
blockName: str = "Block"
while iter.GetCurrentDataObject() is not None:
blockName = iter.GetCurrentMetaData().Get(vtkMultiBlockDataSet.NAME())
if iter.GetCurrentFlatIndex() == index:
break
iter.GoToNextItem()
return blockName
[docs]
def getBlockIndexFromName(
input: Union[vtkMultiBlockDataSet, vtkCompositeDataSet], name: str
) -> int:
"""Get block flat index from name of node in the vtkMultiBlockDataSet tree.
Args:
input (vtkMultiBlockDataSet | vtkCompositeDataSet): input multi block object.
name (str): name of the block to get the index in the tree.
Returns:
int: index of the block if found, -1 otherwise.
"""
blockIndex: int = -1
# initialize data object tree iterator
iter: vtkDataObjectTreeIterator = vtkDataObjectTreeIterator()
iter.SetDataSet(input)
iter.VisitOnlyLeavesOff()
iter.GoToFirstItem()
found: bool = False
while iter.GetCurrentDataObject() is not None:
blockName: str = iter.GetCurrentMetaData().Get(vtkMultiBlockDataSet.NAME())
blockIndex = iter.GetCurrentFlatIndex()
if blockName.lower() == name.lower():
found = True
break
iter.GoToNextItem()
return blockIndex if found else -1
[docs]
def getElementaryCompositeBlockIndexes(
input: Union[vtkMultiBlockDataSet, vtkCompositeDataSet]
) -> dict[str, int]:
"""Get indexes of composite block that contains elementrary blocks.
Args:
input (vtkMultiBlockDataSet | vtkCompositeDataSet): input multi block object.
Returns:
dict[str, int]: dictionary that contains names as keys and flat indices
as values of the parent composite blocks of elementary blocks.
"""
# initialize data object tree iterator
iter: vtkDataObjectTreeIterator = vtkDataObjectTreeIterator()
iter.SetDataSet(input)
iter.VisitOnlyLeavesOff()
iter.GoToFirstItem()
elementaryBlockIndexes: dict[str, int] = {}
while iter.GetCurrentDataObject() is not None:
curIndex: int = iter.GetCurrentFlatIndex()
curName: str = iter.GetCurrentMetaData().Get(vtkMultiBlockDataSet.NAME())
curIsComposite = iter.GetCurrentDataObject().IsA("vtkMultiBlockDataSet")
iter.GoToNextItem()
nextIsNotNone = iter.GetCurrentDataObject() is not None
if (
curIsComposite
and nextIsNotNone
and (not iter.GetCurrentDataObject().IsA("vtkMultiBlockDataSet"))
):
elementaryBlockIndexes[curName] = curIndex
return elementaryBlockIndexes
[docs]
def getBlockElementIndexesFlatten(
input: Union[vtkMultiBlockDataSet, vtkCompositeDataSet]
) -> list[int]:
"""Get a flatten list that contains flat indexes of elementary blocks.
Args:
input (vtkMultiBlockDataSet | vtkCompositeDataSet): input multi block object.
Returns:
list[int]: list of flat indexes
"""
return [i for li in getBlockElementIndexes(input) for i in li]
[docs]
def getBlockElementIndexes(
input: Union[vtkMultiBlockDataSet, vtkCompositeDataSet]
) -> list[list[int]]:
"""Get a list of list that contains flat indexes of elementary blocks.
Each sublist contains the indexes of elementary blocks that belongs to a
same parent node.
Args:
input (vtkMultiBlockDataSet | vtkCompositeDataSet): input multi block object.
Returns:
list[list[int]]: list of list of flat indexes
"""
# initialize data object tree iterator
iter: vtkDataObjectTreeIterator = vtkDataObjectTreeIterator()
iter.SetDataSet(input)
iter.VisitOnlyLeavesOff()
iter.GoToFirstItem()
blockElementIndexes: list[list[int]] = []
indexes: list[int] = []
while iter.GetCurrentDataObject() is not None:
curIndex: int = iter.GetCurrentFlatIndex()
if iter.GetCurrentDataObject().IsA("vtkMultiBlockDataSet"):
# change of parent node, then add the indexes of the previous
# vtkMultiBlockDataSet if needed
if len(indexes) > 0:
blockElementIndexes += [indexes]
# reinitialize the list of indexes of included blocks
indexes = []
else:
indexes += [curIndex]
iter.GoToNextItem()
# add the indexes of the last vtkMultiBlockDataSet if needed
if len(indexes) > 0:
blockElementIndexes += [indexes]
return blockElementIndexes
[docs]
def getBlockFromFlatIndex(
multiBlock: Union[vtkMultiBlockDataSet, vtkCompositeDataSet], blockIndex: int
) -> vtkDataObject:
"""Get the block with blockIndex from input vtkMultiBlockDataSet.
Args:
multiBlock (vtkMultiBlockDataSet | vtkCompositeDataSet): input multi block
blockIndex (int): block index
Returns:
vtkMultiBlockDataSet: block if it exists, None otherwise
"""
block: vtkDataObject
# initialize data object tree iterator
iter: vtkDataObjectTreeIterator = vtkDataObjectTreeIterator()
iter.SetDataSet(multiBlock)
iter.VisitOnlyLeavesOff()
iter.GoToFirstItem()
while iter.GetCurrentDataObject() is not None:
if iter.GetCurrentFlatIndex() == blockIndex:
block = iter.GetCurrentDataObject()
break
iter.GoToNextItem()
return block
[docs]
def getBlockFromName(
multiBlock: Union[vtkMultiBlockDataSet, vtkCompositeDataSet], blockName: str
) -> vtkDataObject:
"""Get the block named blockName from input vtkMultiBlockDataSet.
Args:
multiBlock (vtkMultiBlockDataSet | vtkCompositeDataSet): input multi block
blockName (str): block name
Returns:
vtkDataObject: block if it exists, None otherwise
"""
block: vtkDataObject
# initialize data object tree iterator
iter: vtkDataObjectTreeIterator = vtkDataObjectTreeIterator()
iter.SetDataSet(multiBlock)
iter.VisitOnlyLeavesOff()
iter.GoToFirstItem()
while iter.GetCurrentDataObject() is not None:
if iter.GetCurrentMetaData().Get(vtkMultiBlockDataSet.NAME()) == blockName:
block = vtkMultiBlockDataSet.SafeDownCast(iter.GetCurrentDataObject())
break
iter.GoToNextItem()
return block