Source code for Muscat.Experimental.Containers.PartitionedMeshTools

# -*- coding: utf-8 -*-
#
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
#
from typing import Optional, List, Callable, Generator, Any

import numpy as np

from Muscat.Types import MuscatIndex
from Muscat.Containers.Mesh import Mesh
from Muscat.Experimental.Containers.PartitionedMesh import PartitionedMesh, GlobalIdsFieldName
from Muscat.Containers.Filters.FilterObjects import ElementFilter
from Muscat.Containers.Filters.FilterBase import ElementsSelection
from Muscat.Containers.MeshModificationTools import CleanLonelyNodes


def ElementFilterOnPartitionedMesh(elementFilter: ElementFilter) -> Callable[[PartitionedMesh], Generator[ElementsSelection, Any, None]]:
    """convert an element filter to work on partitioned meshes

        for selection in ElementFilterOnPartitionedMesh(ElementFilter(dimensionality=3))(myParMesh):
            print("--------------------------------------------------")
            print(selection.part)

    Parameters
    ----------
    elementFilter : ElementFilter

    Returns
    -------
    IterOverPartitionedMesh
        a object to iterate in the same way as a filter

    """

    class IterOverPartitionedMesh():
        def __init__(self, elementFilter):
            self.elementFilter: ElementFilter = elementFilter

        def __call__(self, partitionedMesh: PartitionedMesh):
            for part, mesh in enumerate(partitionedMesh.storage):
                for selection in self.elementFilter(mesh):
                    yield ElementsSelection(mesh=selection.mesh, indices=selection.indices, elements=selection.elements, meshOffset=selection.meshOffset, selectionOffset=selection.selectionOffset, part=part)

    return IterOverPartitionedMesh(elementFilter)


def PartitionMeshFromMask(mesh: Mesh, mask: np.ndarray, nbParts: int = 0) -> PartitionedMesh:
    """Function to construct a PartitionedMesh from a mesh and a element partition (coloring)

    Parameters
    ----------
    mesh : Mesh
        the mesh to apply the partition
    mask : np.ndarray
        a numpy array of ints of size Number of elements with values starting from zero
    nbParts : int, optional
        the number of part in the partition (this must be equal to len(np.unique(mask)) ), by default the value is internally computed by a np.unique

    Returns
    -------
    PartitionedMesh
        partitioned mesh (the )
    """

    from Muscat.Containers.MeshInspectionTools import ExtractElementsByElementFilter
    from Muscat.Containers.MeshFieldOperations import CopyFieldsFromOriginalMeshToTargetMesh

    myParMesh = PartitionedMesh()

    if nbParts == 0:
        nbParts = len(np.unique(mask))

    for n in range(nbParts):
        boolMask = mask == n
        f = ElementFilter(eMask=boolMask)
        newMesh = ExtractElementsByElementFilter(mesh, f)
        CopyFieldsFromOriginalMeshToTargetMesh(mesh, newMesh)

        CleanLonelyNodes(newMesh, inPlace=True)
        myParMesh.AddMesh(newMesh)

    return myParMesh


def CheckIntegrity(GUI: bool = False):
    from Muscat.Containers.MeshGraphTools import PartitionMesh
    from Muscat.Containers.Filters.FilterObjects import ElementFilter
    from Muscat.Containers.MeshInspectionTools import ExtractElementsByElementFilter
    from Muscat.Containers.MeshCreationTools import CreateCube
    myMesh = CreateCube(dimensions=[10, 10, 10], spacing=[1, 1, 1], origin=[0, 0, 0], ofTetras=False)
    myMesh = ExtractElementsByElementFilter(myMesh, ElementFilter(dimensionality=3))

    part = np.asarray(PartitionMesh(myMesh, 4))
    myMesh.nodeFields[GlobalIdsFieldName] = np.arange(myMesh.GetNumberOfNodes())
    myMesh.elemFields[GlobalIdsFieldName] = np.arange(myMesh.GetNumberOfElements())

    myParMesh = PartitionMeshFromMask(myMesh, part)

    print(myMesh)
    if GUI:  # pragma: no cover
        from Muscat.Actions.OpenInParaView import OpenInParaView
        OpenInParaView(myParMesh)

    for selection in ElementFilterOnPartitionedMesh(ElementFilter(dimensionality=3))(myParMesh):
        print("--------------------------------------------------")
        print(selection.Size())

    return "OK"


if __name__ == '__main__':
    print(CheckIntegrity(True))  # pragma: no cover