mesh_grpc.py

# Copyright (C) 2024 ANSYS, Inc. and/or its affiliates.
# SPDX-License-Identifier: MIT
#
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Module to manage downloading and parsing the FEM from the MAPDL gRPC server."""
from functools import wraps
import os
import re
import time
from typing import Dict
import weakref

from ansys.api.mapdl.v0 import ansys_kernel_pb2 as anskernel
import numpy as np

from ansys.mapdl.core.common_grpc import DEFAULT_CHUNKSIZE, parse_chunks
from ansys.mapdl.core.mapdl_grpc import MapdlGrpc
from ansys.mapdl.core.misc import requires_package, supress_logging, threaded

TMP_NODE_CM = "__NODE__"


def requires_model(output=None):
    def decorator(method):
        """
        This function wrap some methods to check if the model contains elements or nodes.
        """

        @wraps(method)
        def wrapper(self, *args, **kwargs):
            if self._has_nodes and self._has_elements:
                return method(self, *args, **kwargs)
            else:
                if not output or output == "array":
                    return np.array([])
                elif output == "list":
                    return []
                elif output == "dict":
                    return {}
                else:  # pragma: no cover
                    raise ValueError("Output type not allowed.")

        return wrapper

    return decorator


class MeshGrpc:
    """Provides an interface to the gRPC mesh from MAPDL."""

    def __init__(self, mapdl):
        """Initialize grpc geometry data"""
        if not isinstance(mapdl, MapdlGrpc):  # pragma: no cover
            raise TypeError("Must be initialized using MapdlGrpc class")
        self._mapdl_weakref = weakref.ref(mapdl)
        mapdl._log.debug("Attached MAPDL object to MapdlMesh.")

        self.logger = mapdl._log
        self._log = mapdl._log

        self._ignore_cache_reset = False
        self._reset_cache()

    def __repr__(self):
        txt = "ANSYS Mesh\n"
        txt += f"  Number of Nodes:              {len(self.nnum)}\n"
        txt += f"  Number of Elements:           {len(self.enum)}\n"
        txt += f"  Number of Element Types:      {len(self.ekey)}\n"
        txt += f"  Number of Node Components:    {len(self.node_components)}\n"
        txt += f"  Number of Element Components: {len(self.element_components)}\n"
        return txt

    @property
    def _mapdl(self):
        """Return the weakly referenced instance of mapdl"""
        return self._mapdl_weakref()

    @property
    def _surf(self):
        """External surface"""
        if self._surf_cache is None:
            self._surf_cache = self._grid.extract_surface()
        return self._surf_cache

    @property
    def _has_nodes(self):
        """Returns True when has nodes"""
        # if isinstance(self._nodes, np.ndarray):
        # return bool(self._nodes.size)
        return self.nodes.size != 0

    @property
    def _has_elements(self):
        """Returns True when geometry has elements"""
        return self._elem.size != 0

    def _set_log_level(self, level):
        """Wraps set_log_level"""
        self._mapdl._set_log_level(level)

    def _reset_cache(self):
        """Reset entire mesh cache"""
        if not self._ignore_cache_reset:
            self.logger.debug("Resetting cache")

            self._cache_elem = None
            self._cache_elem_off = None
            self._cache_element_desc = None
            self._cache_nnum = None
            self._cached_elements = None  # cached list of elements
            self._chunk_size = None
            self._elem = None
            self._elem_off = None
            self._enum = None  # cached element numbering
            self._esys = None  # cached element coordinate system
            self._etype = None
            self._etype_cache = None  # cached ansys element type numbering
            self._etype_id = None  # cached element type id
            self._grid = None
            self._grid_cache = None
            self._keyopt = {}
            self._mtype = None  # cached ansys material type
            self._nnum = None
            self._node_angles = None  # cached node angles
            self._node_coord = None  # cached node coordinates
            self._rcon = None  # cached ansys element real constant
            self._rdat = None
            self._rnum = None
            self._secnum = None  # cached section number
            self._surf_cache = None
            self._tshape = None
            self._tshape_key = None

    def _update_cache(self):
        """Threaded local cache update.

        Used when needing all the geometry entries from MAPDL.
        """
        self.logger.debug("Updating cache")
        # elements must have their underlying nodes selected to avoid
        # VTK segfaul
        with self._mapdl.save_selection:
            self._mapdl.nsle("S", mute=True)

            # not thread safe
            self._update_cache_elem()

            threads = [
                self._update_cache_element_desc(),
                self._update_cache_nnum(),
                self._update_node_coord(),
            ]

            for thread in threads:
                thread.join()

            # must occur after read
            self._ignore_cache_reset = True

            # somehow requesting path seems to help windows avoid an
            # outright segfault prior to running CMSEL
            if os.name == "nt":
                _ = self._mapdl.path

            # TODO: flaky
            time.sleep(0.05)

        self._ignore_cache_reset = False

    @threaded
    def _update_cache_nnum(self):
        if self._cache_nnum is None:
            self.logger.debug("Updating nodes cache")
            nnum = self._mapdl.get_array("NODE", item1="NLIST")
            self._cache_nnum = nnum.astype(np.int32)
        if self._cache_nnum.size == 1:
            if self._cache_nnum[0] == 0:
                self._cache_nnum = np.empty(0, np.int32)

    @property
    def _nnum(self):
        """Return node number cache"""
        self._update_cache_nnum().join()
        return self._cache_nnum

    @_nnum.setter
    def _nnum(self, value):
        self._cache_nnum = value

    @threaded
    def _update_cache_element_desc(self):
        if self._cache_element_desc is None:
            self.logger.debug("Updating elements (desc) cache")
            self._cache_element_desc = self._load_element_types()

    @property
    def _ekey(self):
        """Element key description"""
        self._update_cache_element_desc().join()

        # convert to ekey format
        if self._cache_element_desc:
            ekey = []
            for einfo in self._cache_element_desc:
                ekey.append(einfo[:2])
            return np.vstack(ekey).astype(np.int32)
        return np.array([])

    @threaded
    def _update_node_coord(self):
        if self._node_coord is None:
            self._node_coord = self._load_nodes()

    @property
    def _ans_etype(self):
        """FIELD 1 : element type number"""
        if self._etype_cache is None:
            self._etype_cache = self._elem[self._elem_off[:-1] + 1]
        return self._etype_cache

    @property
    def local(self):
        return self._mapdl._local

    @property
    @requires_model()
    def et_id(self):
        """Element type id (ET) for each element."""
        if self._etype_id is None:
            etype_elem_id = self._elem_off[:-1] + 1
            self._etype_id = self._elem[etype_elem_id]
        return self._etype_id

    @property
    @requires_model()
    def tshape(self):
        """Tshape of contact elements."""
        if self._tshape is None:
            shape_elem_id = self._elem_off[:-1] + 7
            self._tshape = self._elem[shape_elem_id]
        return self._tshape

    @property
    @requires_model("dict")
    def tshape_key(self):
        """Dict with the mapping between element type and element shape.

        TShape is only applicable to contact elements.
        """
        if self._tshape_key is None:
            self._tshape_key = np.unique(np.vstack((self.et_id, self.tshape)), axis=1).T
        return {elem_id: tshape for elem_id, tshape in self._tshape_key}

    @property
    @requires_model()
    def material_type(self):
        """Material type index of each element in the archive."""
        # FIELD 0 : material reference number
        if self._mtype is None:
            self._mtype = self._elem[self._elem_off[:-1]]
        return self._mtype

    @property
    @requires_model()
    def etype(self):
        """Element type of each element.

        This is the ansys element type for each element.

        Notes
        -----
        Element types are listed below.  Please see the APDL Element
        Reference for more details:

        https://www.mm.bme.hu/~gyebro/files/vem/ansys_14_element_reference.pdf
        """
        if self._etype is None:
            arr = np.empty(self._ekey[:, 0].max() + 1, np.int32)
            arr[self._ekey[:, 0]] = self._ekey[:, 1]
            self._etype = arr[self._ans_etype]
        return self._etype

    @property
    @requires_model()
    def section(self):
        """Section number"""
        if self._secnum is None:
            self._secnum = self._elem[self._elem_off[:-1] + 3]  # FIELD 3
        return self._secnum

    @property
    @requires_model()
    def element_coord_system(self):
        """Element coordinate system number"""
        if self._esys is None:
            self._esys = self._elem[self._elem_off[:-1] + 4]  # FIELD 4
        return self._esys

    @property
    @requires_model("list")
    def elem(self):
        """List of elements containing raw ansys information.

        Each element contains 10 items plus the nodes belonging to the
        element.  The first 10 items are:

        - FIELD 0 : material reference number
        - FIELD 1 : element type number
        - FIELD 2 : real constant reference number
        - FIELD 3 : section number
        - FIELD 4 : element coordinate system
        - FIELD 5 : death flag (0 - alive, 1 - dead)
        - FIELD 6 : solid model reference
        - FIELD 7 : coded shape key
        - FIELD 8 : element number
        - FIELD 9 : base element number (applicable to reinforcing elements only)
        - FIELDS 10 - 30 : The nodes belonging to the element in ANSYS numbering.

        """
        if self._cached_elements is None:
            self._cached_elements = np.split(self._elem, self._elem_off[1:-1])
        return self._cached_elements

    @property
    def element_components(self):
        """Element components for the archive.

        Output is a dictionary of element components.  Each entry is an
        array of MAPDL element numbers corresponding to the element
        component.  The keys are element component names.
        """
        return self._mapdl.components._get_all_components_type("ELEM")

    @property
    def node_components(self):
        """Node components for the archive.

        Output is a dictionary of node components.  Each entry is an
        array of MAPDL node numbers corresponding to the node
        component.  The keys are node component names.
        """
        return self._mapdl.components._get_all_components_type("NODE")

    @property
    @requires_model()
    def elem_real_constant(self):
        """Real constant reference for each element.

        Use the data within ``rlblock`` and ``rlblock_num`` to get the
        real constant datat for each element.
        """
        # FIELD 2 : real constant reference number
        if self._rcon is None:
            self._rcon = self._elem[self._elem_off[:-1] + 2]
        return self._rcon

    @property
    def ekey(self):
        """Element type key

        Array containing element type numbers in the first column and
        the element types (like SURF154) in the second column.

        """
        return self._ekey

    @property
    def rlblock(self):
        """Real constant data from the RLBLOCK."""
        return self._parse_rlist()

    @property
    def rlblock_num(self):
        """Indices from the real constant data"""
        return list(self._parse_rlist().keys())

    @property
    def nnum(self) -> np.ndarray:
        """Array of currently selected node numbers.

        Examples
        --------
        >>> mapdl.mesh.nnum
        array([    1,     2,     3, ..., 19998, 19999, 20000])
        """
        return self._nnum

    @property
    def nnum_all(self) -> np.ndarray:
        """Array of all node numbers, even those not selected.

        Examples
        --------
        >>> mapdl.mesh.nnum_all
        array([    1,     2,     3, ..., 19998, 19999, 20000])
        """
        self._ignore_cache_reset = True
        with self._mapdl.save_selection:
            self._mapdl.nsel("all", mute=True)

            nnum = self._mapdl.get_array("NODE", item1="NLIST")
            nnum = nnum.astype(np.int32)
            if nnum.size == 1:
                if nnum[0] == 0:
                    nnum = np.empty(0, np.int32)

        self._ignore_cache_reset = False

        return nnum

    @property
    def enum_all(self) -> np.ndarray:
        """Array of all element numbers, even those not selected.

        Examples
        --------
        >>> mapdl.mesh.enum_all
        array([    1,     2,     3, ..., 19998, 19999, 20000])
        """
        self._ignore_cache_reset = True
        with self._mapdl.save_selection:
            self._mapdl.esel("all", mute=True)

            enum = self._mapdl.get_array("ELEM", item1="ELIST")
            enum = enum.astype(np.int32)
            if enum.size == 1:
                if enum[0] == 0:
                    enum = np.empty(0, np.int32)

        self._ignore_cache_reset = False
        return enum

    @property
    @supress_logging
    def n_node(self) -> int:
        """Number of currently selected nodes in MAPDL.

        Examples
        --------
        >>> mapdl.mesh.n_node
        7217
        """
        return int(self._mapdl.get(entity="NODE", item1="COUNT"))

    @property
    @supress_logging
    def n_elem(self) -> int:
        """Number of currently selected elements in MAPDL.

        Examples
        --------
        >>> mapdl.mesh.n_elem
        1520
        """
        return int(self._mapdl.get(entity="ELEM", item1="COUNT"))

    @property
    def node_angles(self):
        """Not yet implemented"""
        return

    @property
    def enum(self) -> np.ndarray:
        """Element numbers of currently selected elements

        Examples
        --------
        >>> mapdl.mesh.enum
        array([    1,     2,     3, ...,  9998,  9999, 10000])
        """
        if self._enum is None:
            if self._mapdl.mesh.n_elem == 0:
                return np.array([], dtype=np.int32)
            else:
                self._enum = self._mapdl.get_array("ELEM", item1="ELIST").astype(
                    np.int32
                )
        return self._enum

    @property
    def key_option(self):
        """Key options of selected element types."""
        self._update_cache_element_desc().join()

        key_opt = {}
        for einfo in self._cache_element_desc:
            keyopt = einfo[2:20]
            if keyopt.any():
                # convert to fortran/ANSYS format
                valid_keyopt = keyopt[keyopt.astype(np.bool_)]
                ans_keyopt = np.vstack((np.nonzero(keyopt)[0] + 1, valid_keyopt))
                key_opt[einfo[0]] = ans_keyopt.T.tolist()

        return key_opt

    @property
    def nodes(self) -> np.ndarray:
        """Array of nodes in Global Cartesian coordinate system.

        Examples
        --------
        >>> mapdl.mesh.nodes
        array([[0.   0.   0.  ]
               [1.   0.   0.  ]
               [0.25 0.   0.  ]
               ...,
               [0.75 0.5  3.5 ]
               [0.75 0.5  4.  ]
               [0.75 0.5  4.5 ]]
        """
        self._update_node_coord().join()
        if self._node_coord is None:
            return np.empty(0)
        return self._node_coord

    @property
    def nodes_in_current_CS(self) -> np.ndarray:
        """Returns the nodes array in the current coordinate system."""
        CS_id = self._mapdl.get_value("active", 0, "CSYS")
        return self.nodes_in_coordinate_system(CS_id=CS_id)

    def nodes_in_coordinate_system(self, CS_id):
        """Return nodes in the desired coordinate system."""
        if CS_id == 0:
            return self.nodes
        else:
            self._mapdl.parameters["__node_loc__"] = self.nodes
            self._mapdl.vfun("__node_loc_cs__", "local", "__node_loc__", CS_id)
            return self._mapdl.parameters["__node_loc_cs__"]

    @property
    def nodes_rotation(self):
        """Returns an array of node rotations"""
        return self._mapdl.nlist(kinternal="").to_array()[:, 4:]

    def _load_nodes(self, chunk_size=DEFAULT_CHUNKSIZE):
        """Loads nodes from server.

        Parameters
        ----------
        chunk_size : int
            Size of the chunks to request from the server.  Default
            256 kB

        Returns
        -------
        np.ndarray
            Numpy array of nodes
        """
        if self._chunk_size:
            chunk_size = self._chunk_size

        request = anskernel.StreamRequest(chunk_size=chunk_size)
        chunks = self._mapdl._stub.Nodes(request)
        nodes = parse_chunks(chunks, np.double).reshape(-1, 3)
        return nodes

    def _update_cache_elem(self):
        """Update the element and element offset cache"""
        if self._cache_elem is None:
            (
                self._cache_elem,
                self._cache_elem_off,
            ) = self._load_elements_offset()

    @property
    def _elem(self):
        """Contingious array of elements.

        Array of elements, with each element starting at the indices
        in offset.  Each element contains 10 items plus the nodes
        belonging to the element.
        """
        self._update_cache_elem()
        return self._cache_elem

    @_elem.setter
    def _elem(self, value):
        self._cache_elem = value

    @property
    def _elem_off(self):
        """Element offset array"""
        self._update_cache_elem()
        return self._cache_elem_off

    @_elem_off.setter
    def _elem_off(self, value):
        self._cache_elem_off = value

    def _load_elements_offset(self, chunk_size=DEFAULT_CHUNKSIZE):
        """Loads elements from server

        Parameters
        ----------
        chunk_size : int, optional
            Size of the chunks to request from the server.

        Returns
        -------
        elements : np.ndarray
            Array of elements, with each element starting at the
            indices in offset.  Each element contains 10 items plus
            the nodes belonging to the element.  The first 10 items
            are:

            mat    - material reference number
            type   - element type number
            real   - real constant reference number
            secnum - section number
            esys   - element coordinate system
            death  - death flag (0 - alive, 1 - dead)
            solidm - solid model reference
            shape  - coded shape key
            elnum  - element number
            baseeid- base element number (applicable to reinforcing
                     elements only)
            nodes  - The nodes belonging to the elementoffset

        offset : np.ndarray
            Array of indices indicating the start of each element.

        """
        if self._chunk_size:
            chunk_size = self._chunk_size

        request = anskernel.StreamRequest(chunk_size=chunk_size)
        chunks = self._mapdl._stub.LoadElements(request)
        elem_raw = parse_chunks(chunks, np.int32)

        if len(elem_raw) == 0:  # for empty mesh.
            return np.array([]), np.array([])
        n_elem = elem_raw[0]

        # ignore zeros
        elem_off_raw = elem_raw[:n_elem]
        elem_off_raw = elem_off_raw[elem_off_raw != 0]
        # TODO: arrays from gRPC interface should include size of the elem array
        lst_value = np.array(elem_raw.size - n_elem, np.int32)
        offset = np.hstack((elem_off_raw - n_elem, lst_value))

        # overwriting the last column to include element numbers
        elems_ = elem_raw.copy()  # elem_raw is only-read
        elems_ = elems_[n_elem:]
        indx_elem = offset[:-1] + 8
        elems_[indx_elem] = self.enum
        return elems_, offset

    def _load_element_types(self, chunk_size=DEFAULT_CHUNKSIZE):
        """Loads element types from the MAPDL server.

        Returns
        -------
        element_types : np.ndarray
            Array of numpy element types.

        Parameters
        ----------
        int
            Size of the chunks to request from the server.
        """
        request = anskernel.StreamRequest(chunk_size=chunk_size)
        chunks = self._mapdl._stub.LoadElementTypeDescription(request)
        data = parse_chunks(chunks, np.int32)
        n_items = data[0]
        split_ind = data[1 : 1 + n_items]
        # empty items sometimes...
        split_ind = split_ind[split_ind != 0]
        return np.split(data, split_ind)[1:]

    @property
    def grid(self):
        """VTK representation of the underlying finite element mesh.

        Examples
        --------
        Store the finite element mesh as a VTK UnstructuredGrid.

        >>> grid = mapdl.mesh.grid
        UnstructuredGrid (0x7f99b4135760)
          N Cells:	32198
          N Points:	50686
          X Bounds:	-1.181e+00, 1.181e+00
          Y Bounds:	-2.362e-01, 0.000e+00
          Z Bounds:	-2.394e+00, 2.509e+00
          N Arrays:	10

        Plot this grid.

        >>> grid.plot()

        Access the node numbers of grid.

        >>> grid.point_data
        Contains keys:
            ansys_node_num
            vtkOriginalPointIds
            origid
            VTKorigID

        >>> grid.point_data['ansys_node_num']
        pyvista_ndarray([    1,     2,     3, ..., 50684, 50685, 50686],
                        dtype=int32)

        Save this grid to disk

        >>> grid.save('grid.vtk')

        Load this grid externally with ParaView or with pyvista

        >>> import pyvista
        >>> pyvista.read('grid.vtk')

        """
        return self._grid

    @property
    @requires_package("pyvista")
    def _grid(self):
        if self._grid_cache is None:
            self._update_cache()
            self._grid_cache = self._parse_vtk(force_linear=True)
        return self._grid_cache

    @_grid.setter
    def _grid(self, value):
        self._grid_cache = value

    def save(
        self,
        filename,
        binary=True,
        force_linear=False,
        allowable_types=None,
        null_unallowed=False,
    ):
        """Save the geometry as a vtk file

        Parameters
        ----------
        filename : str, pathlib.Path
            Filename of output file. Writer type is inferred from
            the extension of the filename.

        binary : bool, optional
            If ``True``, write as binary, else ASCII.

        force_linear : bool, optional
            This parser creates quadratic elements if available.  Set
            this to True to always create linear elements.  Defaults
            to False.

        allowable_types : list, optional
            Allowable element types.  Defaults to all valid element
            types in ``ansys.mapdl.reader.elements.valid_types``

            See ``help(ansys.mapdl.reader.elements)`` for available element types.

        null_unallowed : bool, optional
            Elements types not matching element types will be stored
            as empty (null) elements.  Useful for debug or tracking
            element numbers.  Default False.

        Notes
        -----
        Binary files write much faster than ASCII and have a smaller
        file size.
        """
        grid = self._parse_vtk(
            allowable_types=allowable_types,
            force_linear=force_linear,
            null_unallowed=null_unallowed,
        )
        if grid:
            return grid.save(str(filename), binary=binary)
        else:
            raise ValueError("The mesh is empty, hence no file has been written.")

    def _parse_vtk(
        self,
        allowable_types=None,
        force_linear=False,
        null_unallowed=False,
        fix_midside=True,
        additional_checking=False,
    ):
        from ansys.mapdl.core.mesh.mesh import _parse_vtk

        return _parse_vtk(
            self,
            allowable_types,
            force_linear,
            null_unallowed,
            fix_midside,
            additional_checking,
        )

    def _parse_rlist(self) -> Dict[int, float]:
        # mapdl.rmore(*list)
        with self._mapdl.force_output:
            rlist = self._mapdl.rlist()

        # removing ueless part
        rlist = rlist.replace(
            """   *****MAPDL VERIFICATION RUN ONLY*****
     DO NOT USE RESULTS FOR PRODUCTION
""",
            "",
        )
        constants_ = re.findall(
            r"REAL CONSTANT SET.*?\n\n", rlist + "\n\n", flags=re.DOTALL
        )

        const_ = {}
        for each in constants_:
            values = [0 for i in range(18)]
            set_ = int(re.match(r"REAL CONSTANT SET\s+(\d+)\s+", each).groups()[0])
            limits = (
                int(re.match(r".*ITEMS\s+(\d+)\s+", each).groups()[0]),
                int(re.match(r".*TO\s+(\d+)\s*", each).groups()[0]),
            )
            values_ = [float(i) for i in each.strip().splitlines()[1].split()]

            if not set_ in const_:
                const_[set_] = values

            for i, jlimit in enumerate(range(limits[0] - 1, limits[1])):
                const_[set_][jlimit] = values_[i]

        return const_