vedo.ugrid

  1import numpy as np
  2
  3try:
  4    import vedo.vtkclasses as vtk
  5except ImportError:
  6    import vtkmodules.all as vtk
  7
  8import vedo
  9from vedo import settings
 10from vedo import colors
 11from vedo import utils
 12from vedo.base import BaseGrid
 13from vedo.file_io import download, loadUnStructuredGrid
 14
 15
 16__docformat__ = "google"
 17
 18__doc__ = """
 19Work with unstructured grid datasets
 20"""
 21
 22__all__ = ["UGrid"]
 23
 24#########################################################################
 25class UGrid(BaseGrid, vtk.vtkActor):
 26    """Support for UnstructuredGrid objects."""
 27
 28    def __init__(self, inputobj=None):
 29        """
 30        Support for UnstructuredGrid objects.
 31
 32        Arguments:
 33            inputobj : (list, vtkUnstructuredGrid, str)
 34                A list in the form `[points, cells, celltypes]`,
 35                or a vtkUnstructuredGrid object, or a filename
 36
 37        Celltypes are identified by the following convention:
 38            - VTK_TETRA = 10
 39            - VTK_VOXEL = 11
 40            - VTK_HEXAHEDRON = 12
 41            - VTK_WEDGE = 13
 42            - VTK_PYRAMID = 14
 43            - VTK_HEXAGONAL_PRISM = 15
 44            - VTK_PENTAGONAL_PRISM = 16
 45        """
 46
 47        vtk.vtkActor.__init__(self)
 48        BaseGrid.__init__(self)
 49
 50        inputtype = str(type(inputobj))
 51        self._data = None
 52        self._polydata = None
 53        self._bfprop = None
 54        self.name = "UGrid"
 55
 56        ###################
 57        if inputobj is None:
 58            self._data = vtk.vtkUnstructuredGrid()
 59
 60        elif utils.is_sequence(inputobj):
 61
 62            pts, cells, celltypes = inputobj
 63
 64            self._data = vtk.vtkUnstructuredGrid()
 65
 66            if not utils.is_sequence(cells[0]):
 67                tets = []
 68                nf = cells[0] + 1
 69                for i, cl in enumerate(cells):
 70                    if i in (nf, 0):
 71                        k = i + 1
 72                        nf = cl + k
 73                        cell = [cells[j + k] for j in range(cl)]
 74                        tets.append(cell)
 75                cells = tets
 76
 77            # This would fill the points and use those to define orientation
 78            vpts = utils.numpy2vtk(pts, dtype=np.float32)
 79            points = vtk.vtkPoints()
 80            points.SetData(vpts)
 81            self._data.SetPoints(points)
 82
 83            # This fill the points and use cells to define orientation
 84            # points = vtk.vtkPoints()
 85            # for c in cells:
 86            #       for pid in c:
 87            #           points.InsertNextPoint(pts[pid])
 88            # self._data.SetPoints(points)
 89
 90            # Fill cells
 91            # https://vtk.org/doc/nightly/html/vtkCellType_8h_source.html
 92            for i, ct in enumerate(celltypes):
 93                cell_conn = cells[i]
 94                if ct == vtk.VTK_HEXAHEDRON:
 95                    cell = vtk.vtkHexahedron()
 96                elif ct == vtk.VTK_TETRA:
 97                    cell = vtk.vtkTetra()
 98                elif ct == vtk.VTK_VOXEL:
 99                    cell = vtk.vtkVoxel()
100                elif ct == vtk.VTK_WEDGE:
101                    cell = vtk.vtkWedge()
102                elif ct == vtk.VTK_PYRAMID:
103                    cell = vtk.vtkPyramid()
104                elif ct == vtk.VTK_HEXAGONAL_PRISM:
105                    cell = vtk.vtkHexagonalPrism()
106                elif ct == vtk.VTK_PENTAGONAL_PRISM:
107                    cell = vtk.vtkPentagonalPrism()
108                else:
109                    print("UGrid: cell type", ct, "not implemented. Skip.")
110                    continue
111                cpids = cell.GetPointIds()
112                for j, pid in enumerate(cell_conn):
113                    cpids.SetId(j, pid)
114                self._data.InsertNextCell(ct, cpids)
115
116        elif "UnstructuredGrid" in inputtype:
117            self._data = inputobj
118
119        elif isinstance(inputobj, str):
120            if "https://" in inputobj:
121                inputobj = download(inputobj, verbose=False)
122            self._data = loadUnStructuredGrid(inputobj)
123            self.filename = inputobj
124
125        else:
126            vedo.logger.error(f"cannot understand input type {inputtype}")
127            return
128
129        # self._mapper = vtk.vtkDataSetMapper()
130        self._mapper = vtk.vtkPolyDataMapper()
131        self._mapper.SetInterpolateScalarsBeforeMapping(settings.interpolate_scalars_before_mapping)
132
133        if settings.use_polygon_offset:
134            self._mapper.SetResolveCoincidentTopologyToPolygonOffset()
135            pof, pou = settings.polygon_offset_factor, settings.polygon_offset_units
136            self._mapper.SetResolveCoincidentTopologyPolygonOffsetParameters(pof, pou)
137        self.GetProperty().SetInterpolationToFlat()
138
139        if not self._data:
140            return
141
142        # now fill the representation of the vtk unstr grid
143        sf = vtk.vtkShrinkFilter()
144        sf.SetInputData(self._data)
145        sf.SetShrinkFactor(1.0)
146        sf.Update()
147        gf = vtk.vtkGeometryFilter()
148        gf.SetInputData(sf.GetOutput())
149        gf.Update()
150        self._polydata = gf.GetOutput()
151
152        self._mapper.SetInputData(self._polydata)
153        sc = None
154        if self.useCells:
155            sc = self._polydata.GetCellData().GetScalars()
156        else:
157            sc = self._polydata.GetPointData().GetScalars()
158        if sc:
159            self._mapper.SetScalarRange(sc.GetRange())
160
161        self.SetMapper(self._mapper)
162        self.property = self.GetProperty()
163
164        self.pipeline = utils.OperationNode(
165            self, comment=f"#cells {self._data.GetNumberOfCells()}",
166            c="#4cc9f0",
167        )
168    # ------------------------------------------------------------------
169
170    def _repr_html_(self):
171        """
172        HTML representation of the UGrid object for Jupyter Notebooks.
173
174        Returns:
175            HTML text with the image and some properties.
176        """
177        import io
178        import base64
179        from PIL import Image
180
181        library_name = "vedo.ugrid.UGrid"
182        help_url = "https://vedo.embl.es/docs/vedo/ugrid.html"
183
184        arr = self.thumbnail()
185        im = Image.fromarray(arr)
186        buffered = io.BytesIO()
187        im.save(buffered, format="PNG", quality=100)
188        encoded = base64.b64encode(buffered.getvalue()).decode("utf-8")
189        url = "data:image/png;base64," + encoded
190        image = f"<img src='{url}'></img>"
191
192        bounds = "<br/>".join(
193            [
194                utils.precision(min_x,4) + " ... " + utils.precision(max_x,4)
195                for min_x, max_x in zip(self.bounds()[::2], self.bounds()[1::2])
196            ]
197        )
198
199        help_text = ""
200        if self.name:
201            help_text += f"<b> {self.name}: &nbsp&nbsp</b>"
202        help_text += '<b><a href="' + help_url + '" target="_blank">' + library_name + "</a></b>"
203        if self.filename:
204            dots = ""
205            if len(self.filename) > 30:
206                dots = "..."
207            help_text += f"<br/><code><i>({dots}{self.filename[-30:]})</i></code>"
208
209        pdata = ""
210        if self._data.GetPointData().GetScalars():
211            if self._data.GetPointData().GetScalars().GetName():
212                name = self._data.GetPointData().GetScalars().GetName()
213                pdata = "<tr><td><b> point data array </b></td><td>" + name + "</td></tr>"
214
215        cdata = ""
216        if self._data.GetCellData().GetScalars():
217            if self._data.GetCellData().GetScalars().GetName():
218                name = self._data.GetCellData().GetScalars().GetName()
219                cdata = "<tr><td><b> cell data array </b></td><td>" + name + "</td></tr>"
220
221        pts = self.points()
222        cm = np.mean(pts, axis=0)
223
224        all = [
225            "<table>",
226            "<tr>",
227            "<td>", image, "</td>",
228            "<td style='text-align: center; vertical-align: center;'><br/>", help_text,
229            "<table>",
230            "<tr><td><b> bounds </b> <br/> (x/y/z) </td><td>" + str(bounds) + "</td></tr>",
231            "<tr><td><b> center of mass </b></td><td>" + utils.precision(cm,3) + "</td></tr>",
232            # "<tr><td><b> average size </b></td><td>" + str(average_size) + "</td></tr>",
233            "<tr><td><b> nr. points&nbsp/&nbspcells </b></td><td>"
234            + str(self.npoints) + "&nbsp/&nbsp" + str(self.ncells) + "</td></tr>",
235            pdata,
236            cdata,
237            "</table>",
238            "</table>",
239        ]
240        return "\n".join(all)
241
242    def clone(self):
243        """Clone the UGrid object to yield an exact copy."""
244        ugCopy = vtk.vtkUnstructuredGrid()
245        ugCopy.DeepCopy(self._data)
246
247        cloned = UGrid(ugCopy)
248        pr = self.GetProperty()
249        if isinstance(pr, vtk.vtkVolumeProperty):
250            prv = vtk.vtkVolumeProperty()
251        else:
252            prv = vtk.vtkProperty()
253        prv.DeepCopy(pr)
254        cloned.SetProperty(prv)
255        cloned.property = prv
256
257        # assign the same transformation to the copy
258        cloned.SetOrigin(self.GetOrigin())
259        cloned.SetScale(self.GetScale())
260        cloned.SetOrientation(self.GetOrientation())
261        cloned.SetPosition(self.GetPosition())
262        cloned.name = self.name
263
264        cloned.pipeline = utils.OperationNode(
265            "clone", parents=[self], shape='diamond', c='#bbe1ed',
266        )
267        return cloned
268
269    def color(self, c=False, alpha=None):
270        """
271        Set/get UGrid color.
272        If None is passed as input, will use colors from active scalars.
273        Same as `ugrid.c()`.
274        """
275        if c is False:
276            return np.array(self.GetProperty().GetColor())
277        if c is None:
278            self._mapper.ScalarVisibilityOn()
279            return self
280        self._mapper.ScalarVisibilityOff()
281        cc = colors.get_color(c)
282        self.property.SetColor(cc)
283        if self.trail:
284            self.trail.GetProperty().SetColor(cc)
285        if alpha is not None:
286            self.alpha(alpha)
287        return self
288
289    def alpha(self, opacity=None):
290        """Set/get mesh's transparency. Same as `mesh.opacity()`."""
291        if opacity is None:
292            return self.property.GetOpacity()
293
294        self.property.SetOpacity(opacity)
295        bfp = self.GetBackfaceProperty()
296        if bfp:
297            if opacity < 1:
298                self._bfprop = bfp
299                self.SetBackfaceProperty(None)
300            else:
301                self.SetBackfaceProperty(self._bfprop)
302        return self
303
304    def opacity(self, alpha=None):
305        """Set/get mesh's transparency. Same as `mesh.alpha()`."""
306        return self.alpha(alpha)
307
308    def wireframe(self, value=True):
309        """Set mesh's representation as wireframe or solid surface.
310        Same as `mesh.wireframe()`."""
311        if value:
312            self.property.SetRepresentationToWireframe()
313        else:
314            self.property.SetRepresentationToSurface()
315        return self
316
317    def linewidth(self, lw=None):
318        """Set/get width of mesh edges. Same as `lw()`."""
319        if lw is not None:
320            if lw == 0:
321                self.property.EdgeVisibilityOff()
322                self.property.SetRepresentationToSurface()
323                return self
324            self.property.EdgeVisibilityOn()
325            self.property.SetLineWidth(lw)
326        else:
327            return self.property.GetLineWidth()
328        return self
329
330    def lw(self, linewidth=None):
331        """Set/get width of mesh edges. Same as `linewidth()`."""
332        return self.linewidth(linewidth)
333
334    def linecolor(self, lc=None):
335        """Set/get color of mesh edges. Same as `lc()`."""
336        if lc is not None:
337            if "ireframe" in self.property.GetRepresentationAsString():
338                self.property.EdgeVisibilityOff()
339                self.color(lc)
340                return self
341            self.property.EdgeVisibilityOn()
342            self.property.SetEdgeColor(colors.get_color(lc))
343        else:
344            return self.property.GetEdgeColor()
345        return self
346
347    def lc(self, linecolor=None):
348        """Set/get color of mesh edges. Same as `linecolor()`."""
349        return self.linecolor(linecolor)
350
351    def extract_cell_type(self, ctype):
352        """Extract a specific cell type and return a new `UGrid`."""
353        uarr = self._data.GetCellTypesArray()
354        ctarrtyp = np.where(utils.vtk2numpy(uarr) == ctype)[0]
355        uarrtyp = utils.numpy2vtk(ctarrtyp, deep=False, dtype="id")
356        selection_node = vtk.vtkSelectionNode()
357        selection_node.SetFieldType(vtk.vtkSelectionNode.CELL)
358        selection_node.SetContentType(vtk.vtkSelectionNode.INDICES)
359        selection_node.SetSelectionList(uarrtyp)
360        selection = vtk.vtkSelection()
361        selection.AddNode(selection_node)
362        es = vtk.vtkExtractSelection()
363        es.SetInputData(0, self._data)
364        es.SetInputData(1, selection)
365        es.Update()
366        ug = UGrid(es.GetOutput())
367
368        ug.pipeline = utils.OperationNode(
369            "extract_cell_type", comment=f"type {ctype}",
370            c="#edabab", parents=[self],
371        )
372        return ug