LoopStructural 1.0.3__zip → 1.0.71.dev0__zip

Miniconda/envs/loop/Lib/site-packages/LoopStructural/visualisation/model_visualisation.py

@@ -4,15 +4,19 @@ A wrapper for lavavu
 """
 
 import logging
-
-import lavavu
+from LoopStructural.utils import getLogger
+logger = getLogger(__name__)
+
+try:
+    import lavavu
+    from lavavu.vutils import is_notebook
+except ImportError:
+    logger.error("Please install lavavu: pip install lavavu")
 import numpy as np
-from lavavu.vutils import is_notebook
-from skimage.measure import marching_cubes_lewiner as marching_cubes
-
+from skimage.measure import marching_cubes
+from LoopStructural.modelling.features import GeologicalFeature
 from LoopStructural.utils.helper import create_surface, get_vectors, create_box
 
-logger = logging.getLogger(__name__)
 # adapted/copied from pyvista for sphinx scraper
 _OPEN_VIEWERS = {}
 
@@ -40,12 +44,13 @@ class LavaVuModelViewer:
         ----------
         **kwargs : lavavu viewer kwargs
 
-        Attributes
-        ----------
-        lv Lavavu.Viewer object
+
         objects : dictionary of objects that have been plotted
         """
         # copied from pyvista
+        if lavavu is None:
+            logger.error("Lavavu isn't installed: pip install lavavu")
+            return
         self._id_name = "{}-{}".format(str(hex(id(self))), len(_OPEN_VIEWERS))
         _OPEN_VIEWERS[self._id_name] = self
         #
@@ -63,11 +68,55 @@ class LavaVuModelViewer:
             logger.error("Plot area has not been defined.")
         self.bounding_box = np.array(self.bounding_box)
         self.nsteps = np.array(self.nsteps)
-        self.model = model
+        self._model = model
         # prerotate to a nice view
         # self.lv.rotate([-57.657936096191406, -13.939384460449219, -6.758780479431152])
     def close(self):
         pass
+    
+    @property
+    def model(self):
+        return self._model
+    
+    @model.setter
+    def model(self, model):
+        if model is not None:
+            self.bounding_box = np.array(model.bounding_box)
+            self.nsteps = np.array(model.nsteps)
+            self._model = model
+            self._nelements = self.nsteps[0]*self.nsteps[1]*self.nsteps[2]
+            logger.debug("Using bounding box from model")
+    @property
+    def nelements(self):
+        """The number of elements to use for evaluating the isosurface
+
+        Returns
+        -------
+        nelements : int
+            number of elements to use for isosurfacing
+        """
+        return self._nelements
+    
+    @nelements.setter
+    def nelements(self, nelements : int):
+        """Setter for nelements, automatically caculates the number of equally sized elements
+        to isosurface. Better than specifying step distance manually
+
+        Parameters
+        ----------
+        nelements : int
+            [description]
+        """        
+        box_vol = (self.bounding_box[1, 0]-self.bounding_box[0, 0]) * (self.bounding_box[1, 1]-self.bounding_box[0, 1]) * (self.bounding_box[1, 2]-self.bounding_box[0, 2])
+        ele_vol = box_vol / nelements
+        # calculate the step vector of a regular cube
+        step_vector = np.zeros(3)
+        step_vector[:] = ele_vol ** (1. / 3.)
+        # step_vector /= np.array([1,1,2])
+        # number of steps is the length of the box / step vector
+        nsteps = np.ceil((self.bounding_box[1, :] - self.bounding_box[0, :]) / step_vector).astype(int)
+        self.nsteps = nsteps
+        logger.info("Using grid with dimensions {} {} {}".format(nsteps[0],nsteps[1],nsteps[2]))
 
     def deep_clean(self):
         """[summary]
@@ -77,8 +126,8 @@ class LavaVuModelViewer:
         self.lv.clear()
         self.lv.cleardata()
         pass
-
-    def add_section(self, geological_feature=None, axis='x', value=None, **kwargs):
+    
+    def add_section(self, geological_feature=None, axis='x', value=None,  **kwargs):
         """
 
         Plot a section/map thru the model and paint with a geological feature
@@ -98,29 +147,26 @@ class LavaVuModelViewer:
         -------
 
         """
-
         if axis == 'x':
             tri, yy, zz = create_surface(self.bounding_box[:, [1, 2]], self.nsteps[[1, 2]])
             xx = np.zeros(zz.shape)
             if value is None:
-                xx[:] = np.nanmean(self.bounding_box[:, 0])
-            else:
-                xx[:] = value
+                value = np.nanmean(self.bounding_box[:, 0])
+            xx[:] = value
         if axis == 'y':
             tri, xx, zz = create_surface(self.bounding_box[:, [0, 2]], self.nsteps[[0, 2]])
             yy = np.zeros(xx.shape)
             if value is None:
-                yy[:] = np.nanmean(self.bounding_box[:, 1])
-            else:
-                yy[:] = value
+                value = np.nanmean(self.bounding_box[:, 1])
+            yy[:] = value
         if axis == 'z':
             tri, xx, yy = create_surface(self.bounding_box[:, 0:2], self.nsteps[0:2])
             zz = np.zeros(xx.shape)
             if value is None:
-                zz[:] = np.nanmean(self.bounding_box[:, 2])
-            else:
-                zz[:] = value
-        name = kwargs.get('name', axis + '_slice')
+                value = np.nanmean(self.bounding_box[:, 2])
+            zz[:] = value
+        name = kwargs.get('name', geological_feature.name)
+        name = '{}_section_at_{}_of_{}'.format(axis,value,name)
         colour = kwargs.get('colour', 'red')
 
         # create an array to evaluate the feature on for the section
@@ -130,12 +176,13 @@ class LavaVuModelViewer:
         points[:, 2] = zz
 
         surf = self.lv.triangles(name)
-        surf.vertices(points)
+        surf.vertices(self.model.rescale(points,inplace=False))
         surf.indices(tri)
         logger.info("Adding %s section at %f" % (axis, value))
         if geological_feature is None:
             surf.colours(colour)
-        if geological_feature is not None:
+
+        if geological_feature is not None and type(geological_feature) == GeologicalFeature:
             if 'norm' in kwargs:
                 surf.values(np.linalg.norm(
                     geological_feature.evaluate_gradient(points), axis=1),
@@ -150,10 +197,19 @@ class LavaVuModelViewer:
             logger.info("Colouring section with %s min: %f, max: %f" % (
                 geological_feature.name, geological_feature.min(), geological_feature.max()))
             surf.colourmap(cmap, range=[geological_feature.min(), geological_feature.max()])
+        if geological_feature == 'model' and self.model is not None:
+            name = kwargs.get('name','model_section')
+            surf.values(self.model.evaluate_model(points,scale=True),
+                            name)
+            surf["colourby"] = name
+            cmap = lavavu.cubehelix(100)
+            if 'cmap' in kwargs:
+                cmap = kwargs['cmap']
+
 
     def add_isosurface(self, geological_feature, value = None, isovalue=None,
                      paint_with=None, slices=None, colour='red', nslices=None, 
-                     cmap=None, filename=None, **kwargs):
+                     cmap=None, filename=None, names=None, colours=None,**kwargs):
         """ Plot the surface of a geological feature 
 
         [extended_summary]
@@ -177,14 +233,19 @@ class LavaVuModelViewer:
         cmap : [type], optional
             [description], by default None
         filename: string, optional
-            filename for exporting 
+            filename for exporting
+        names: list, optional
+            list of names same length as slices
+        colours: list, optional
+            list of colours same length as slices
 
         Returns
         -------
         [type]
             [description]
         """
-
+        if geological_feature is None:
+            logger.error("Cannot add isosurface GeologicalFeature does not exist")
         # update the feature to make sure its current
         if 'update' in kwargs:
             geological_feature.update()
@@ -230,7 +291,7 @@ class LavaVuModelViewer:
         if region is not None:
             val[~region(np.array([xx.flatten(), yy.flatten(), zz.flatten()]).T)] = np.nan
         step_vector = np.array([x[1] - x[0], y[1] - y[0], z[1] - z[0]])
-        for isovalue in slices_:
+        for i, isovalue in enumerate(slices_):
             logger.info("Creating isosurface of %s at %f" % (geological_feature.name, isovalue))
 
             if isovalue > np.nanmax(val) or isovalue < np.nanmin(val):
@@ -242,23 +303,41 @@ class LavaVuModelViewer:
                     isovalue,
                     spacing=step_vector)
                 verts += np.array([self.bounding_box[0, 0], self.bounding_box[0, 1], self.bounding_box[1, 2]])
-            except ValueError:
-                logger.warning("no surface to mesh, skipping")
+                self.model.rescale(verts)
+
+            except (ValueError, RuntimeError) as e:
+                print(e)
+                logger.warning("Cannot isosurface {} at {}, skipping".format(geological_feature.name,isovalue))
                 continue
-            
 
+            
             name = geological_feature.name
             name = kwargs.get('name', name)
             name += '_iso_%f' % isovalue
+            if names is not None and len(names) == len(slices_):
+                name = names[i]
+            if name in self.lv.objects:
+                ii = 0
+                newname = name+"_{}".format(ii)
+                while newname in self.lv.objects:
+                    ii+=1
+                    newname = name+"_{}".format(ii)
+                name  = newname
+                
+            if colours is not None and len(colours) == len(slices_):
+                colour=colours[i]
             if filename is not None:
+                svalues = None
+                # svalues[:] = np.nan
                 try:
                     import meshio
+                    meshio.write_points_cells(filename.format(name),
+                    verts,
+                    [("triangle", faces)]
+                    )
                 except ImportError:
                     logger.error("Could not save surfaces, meshio is not installed")
-                meshio.write_points_cells(filename.format(name),
-                self.model.rescale(verts),
-                [("triangle", faces)]
-                )
+
             surf = self.lv.triangles(name)
             surf.vertices(verts)
             surf.indices(faces)
@@ -268,7 +347,7 @@ class LavaVuModelViewer:
                 # add a property to the surface nodes for visualisation
                 # calculate the mode value, just to get the most common value
                 surfaceval = np.zeros(verts.shape[0])
-                surfaceval[:] = painter.evaluate_value(verts)
+                surfaceval[:] = painter.evaluate_value(self.model.scale(verts))
                 if painter.name is geological_feature.name:
                     logger.info("Setting surface value to %f"%isovalue)
                     surfaceval[:] = isovalue
@@ -278,43 +357,63 @@ class LavaVuModelViewer:
                 vmax = kwargs.get('vmax', max_property_val)
                 surf.colourmap(cmap, range=(vmin, vmax))  # nodes.shape[0]))
         
-    def add_scalar_field(self, geological_feature, **kwargs):
-        """
+    def add_scalar_field(self, geological_feature, name=None, cmap='rainbow', vmin=None, vmax = None, **kwargs):
+        """Add a block the size of the model area painted with the scalar field value
 
         Parameters
         ----------
         geological_feature : GeologicalFeature
             the geological feature to colour the scalar field by
-        kwargs
-            kwargs for lavavu
-
-        Returns
-        -------
-
+        name : string, optional
+            Name of the object for lavavu, needs to be unique for the viewer object, by default uses feature name
+        cmap : str, optional
+            mpl colourmap reference, by default 'rainbow'
+        vmin : double, optional
+            minimum value of the colourmap, by default None
+        vmax : double, optional
+            maximum value of the colourmap, by default None
         """
-        name = kwargs.get('name', geological_feature.name + '_scalar_field')
+        if name == None:
+            if geological_feature is None:
+                name = 'unnamed scalar field'
+            else:
+                name = geological_feature.name + '_scalar_field'
+
         points, tri = create_box(self.bounding_box,self.nsteps)
 
         surf = self.lv.triangles(name)
-        surf.vertices(points)
+        surf.vertices(self.model.rescale(points))
         surf.indices(tri)
-        val =geological_feature.evaluate_value(points)
+        val =geological_feature.evaluate_value(self.model.scale(points))
         surf.values(val, geological_feature.name)
         surf["colourby"] = geological_feature.name
-        cmap = kwargs.get('cmap',lavavu.cubehelix(100))
-
         logger.info("Adding scalar field of %s to viewer. Min: %f, max: %f" % (geological_feature.name,
                                                                                geological_feature.min(),
                                                                                geological_feature.max()))
-        vmin = kwargs.get('vmin', np.nanmin(val))
-        vmax = kwargs.get('vmax', np.nanmax(val))
+        if vmin == None:
+            vmin =np.nanmin(val)
+        if vmax == None:
+            vmax = np.nanmax(val)
         surf.colourmap(cmap, range=(vmin, vmax))
 
-    def add_model(self, **kwargs):
+    def add_box(self,bounding_box,name,colour='red'):
+        points, tri = create_box(bounding_box,self.nsteps)
+
+        surf = self.lv.triangles(name)
+        surf.vertices(self.model.rescale(points))
+        surf.indices(tri)
+        surf.colours(colour)
+
+    def add_model(self, cmap = None, **kwargs):
         """Add a block model painted by stratigraphic id to the viewer
 
         Calls self.model.evaluate_model() for a cube surrounding the model.
 
+        Parameters
+        ----------
+        cmap : matplotlib cmap, optional
+            colourmap name or object from mpl
+
         Notes
         ------
         It is sensible to increase the viewer step sizes before running this function to
@@ -324,17 +423,36 @@ class LavaVuModelViewer:
         >>> viewer.nsteps = np.array([100,100,100])
 
         """
+        import matplotlib.colors as colors
+        from matplotlib import cm
+
         name = kwargs.get('name', 'geological_model')
         points, tri = create_box(self.bounding_box, self.nsteps)
 
         surf = self.lv.triangles(name)
-        surf.vertices(points)
+        surf.vertices(self.model.rescale(points))
         surf.indices(tri)
-        val = self.model.evaluate_model(points,rescale=False)
+        val = self.model.evaluate_model(points,scale=True)
         surf.values(val, 'model')
         surf["colourby"] = 'model'
-        cmap = kwargs.get('cmap', lavavu.cubehelix(100))
 
+        if cmap is None:
+            import matplotlib.colors as colors
+            colours = []
+            boundaries = []
+            data = []
+            for g in self.model.stratigraphic_column.keys():
+                if g == 'faults':
+                    continue
+                for u, v  in self.model.stratigraphic_column[g].items():
+                    data.append((v['id'],v['colour']))
+                    colours.append(v['colour'])
+                    boundaries.append(v['id'])#print(u,v)
+            cmap = colors.ListedColormap(colours).colors
+        # else:
+            # cmap = cm.get_cmap(cmap,n_units)
+
+        
         # logger.info("Adding scalar field of %s to viewer. Min: %f, max: %f" % (geological_feature.name,
         #                                                                        geological_feature.min(),
         #                                                                        geological_feature.max()))
@@ -342,7 +460,55 @@ class LavaVuModelViewer:
         vmax = kwargs.get('vmax', np.nanmax(val))
         surf.colourmap(cmap, range=(vmin, vmax))
 
-    def add_model_surfaces(self, faults = True, cmap='tab20', **kwargs):
+    def add_fault_displacements(self, cmap = 'rainbow', **kwargs):
+        """Add a block model painted by the fault displacement magnitude
+
+        Calls fault.displacementfeature.evaluate_value(points) for all faults
+
+        Parameters
+        ----------
+        cmap : matplotlib cmap, optional
+            colourmap name or object from mpl
+
+        Notes
+        ------
+        It is sensible to increase the viewer step sizes before running this function to
+        increase the resolution of the model as its not possible to interpolate a discrete
+        colourmap and this causes the model to look like a lego block.
+        You can update the model resolution by changing the attribute nsteps
+        >>> viewer.nsteps = np.array([100,100,100])
+
+        """
+        
+        name = kwargs.get('name', 'fault_displacements')
+        points, tri = create_box(self.bounding_box, self.nsteps)
+
+        surf = self.lv.triangles(name)
+        surf.vertices(self.model.rescale(points))
+        surf.indices(tri)
+        vals = self.model.evaluate_fault_displacements(points)
+        surf.values(vals, 'displacement')
+        surf["colourby"] = 'displacement'
+
+        vmin = kwargs.get('vmin', np.nanmin(vals))
+        vmax = kwargs.get('vmax', np.nanmax(vals))
+        surf.colourmap(cmap, range=(vmin, vmax))
+        
+    def add_fault(self,fault,step=100):
+        self.add_isosurface(fault,value=0,name=fault.name)
+        self.add_vector_field(fault,locations=self.model.regular_grid()[::step])
+
+    def unfault_grid(self,feature,grid=None):
+        if grid is None:
+            grid = self.model.regular_grid()
+        # apply all faults associated with a feature to a regular grid
+        self.add_value_data(self.model.rescale(grid,inplace=False),grid[:,2],name='Regular grid before faults',pointsize=10,)
+   
+        for f in feature.faults:
+            grid = f.apply_to_points(grid)
+        self.add_value_data(self.model.rescale(grid,inplace=False),grid[:,2],name='Regular grid after faults',pointsize=10,)
+
+    def add_model_surfaces(self, strati=True, faults = True, cmap=None, fault_colour='black',**kwargs):
         """Add surfaces for all of the interfaces in the model
 
 
@@ -357,26 +523,79 @@ class LavaVuModelViewer:
         Other parameters are passed to self.add_isosurface() 
 
         """
+        import time
         from matplotlib import cm
+        from matplotlib import colors
+        from tqdm.auto import tqdm
+        start = time.time()
         n_units = 0 #count how many discrete colours
-        for g in self.model.stratigraphic_column.keys():
-            for u in self.model.stratigraphic_column[g].keys():
-                n_units+=1
-        tab = cm.get_cmap(cmap,n_units)
-        ci = 0
-
         for g in self.model.stratigraphic_column.keys():
             if g in self.model.feature_name_index:
-                feature = self.model.features[self.model.feature_name_index[g]]
-                for u, vals in self.model.stratigraphic_column[g].items():
-                    self.add_isosurface(feature, isovalue=vals['max'],name=u,colour=tab.colors[ci,:],**kwargs)
-                    ci+=1
+                for u in self.model.stratigraphic_column[g].keys():
+                    n_units+=1
+        n_faults = 0
+        for f in self.model.features:
+            if f.type=='fault':
+                n_faults+=1
+        
+        if cmap is None:            
+            colours = []
+            boundaries = []
+            data = []
+            for g in self.model.stratigraphic_column.keys():
+                if g == 'faults':
+                    # skip anything saved in faults here
+                    continue
+                for u, v  in self.model.stratigraphic_column[g].items():
+                    data.append((v['id'],v['colour']))
+                    colours.append(v['colour'])
+                    boundaries.append(v['id'])
+            cmap = colors.ListedColormap(colours)
+        else:
+            cmap = cm.get_cmap(cmap,n_units)
+        ci = 0
+        cmap_colours = colors.to_rgba_array(cmap.colors)
+        n_surfaces = 0
+        if strati:
+            n_surfaces+=n_units
         if faults:
-            for f in self.model.features:
-                if f.type == 'fault':
-                    self.add_isosurface(f,isovalue=0,**kwargs)
-
+            n_surfaces+=n_faults
+        with tqdm(total=n_surfaces) as pbar:
+
+            if strati:
+                for g in self.model.stratigraphic_column.keys():
+                    if g in self.model.feature_name_index:
+                        feature = self.model.features[self.model.feature_name_index[g]]
+                        names = []
+                        values = []
+                        colours = []
+                        for u, vals in self.model.stratigraphic_column[g].items():
+                            names.append(u)
+                            values.append(vals['min'])
+                            colours.append(cmap_colours[ci,:])
+                            ci+=1
+                        pbar.set_description('Isosurfacing {}'.format(feature.name))
+                        self.add_isosurface(feature, slices=values,names=names,colours=colours,**kwargs)
+                        pbar.update(len(values))
+            
 
+            if faults:
+                for f in self.model.features:
+                    if f.type == 'fault':
+                        def mask(x):
+                            val = f.displacementfeature.evaluate_value(x)
+                            val[np.isnan(val)] = 0
+                            maskv = np.zeros(val.shape).astype(bool)
+                            maskv[np.abs(val) > 0.001] = 1
+                            return maskv
+                        if f.name in self.model.stratigraphic_column['faults']:
+                            fault_colour = self.model.stratigraphic_column['faults'][f.name].get('colour',['red'])
+                        pbar.set_description('Isosurfacing {}'.format(f.name))
+
+                        region = kwargs.pop('region',None) 
+                        self.add_isosurface(f,isovalue=0,region=mask,colour=fault_colour[0],name=f.name,**kwargs)
+                        pbar.update(1)
+            print("Adding surfaces took {} seconds".format(time.time()-start))
     def add_vector_field(self, geological_feature, **kwargs):
         """
 
@@ -406,7 +625,7 @@ class LavaVuModelViewer:
         vector[mask, :] /= np.linalg.norm(vector[mask, :], axis=1)[:, None]
         vectorfield = self.lv.vectors(geological_feature.name + "_grad",
                                       **kwargs)
-        vectorfield.vertices(locations[mask, :])
+        vectorfield.vertices(self.model.rescale(locations[mask, :],inplace=False))
         vectorfield.vectors(vector[mask, :])
         return
 
@@ -429,6 +648,7 @@ class LavaVuModelViewer:
         add_grad = True
         add_value = True
         add_tang = True
+        add_interface = True
         if 'name' in kwargs:
             name = kwargs['name']
             del kwargs['name']
@@ -438,26 +658,49 @@ class LavaVuModelViewer:
             add_value = kwargs['value']
         if 'tang' in kwargs:
             add_tang = kwargs['tang']
+        if 'interface' in kwargs:
+            add_interface = kwargs['interface']
         grad = feature.builder.get_gradient_constraints()
         norm = feature.builder.get_norm_constraints()
         value = feature.builder.get_value_constraints()
         tang = feature.builder.get_tangent_constraints()
+        interface = feature.builder.get_interface_constraints()
+
         if grad.shape[0] > 0 and add_grad:
-            self.add_vector_data(grad[:, :3], grad[:, 3:6], name + "_grad_cp",
+            self.add_vector_data(self.model.rescale(grad[:, :3],inplace=False), grad[:, 3:6], name + "_grad_cp",
                                  **kwargs)
 
         if norm.shape[0] > 0 and add_grad:
-            self.add_vector_data(norm[:, :3], norm[:, 3:6], name + "_norm_cp",
+            self.add_vector_data(self.model.rescale(norm[:, :3],inplace=False), norm[:, 3:6], name + "_norm_cp",
                                  **kwargs)
         if value.shape[0] > 0 and add_value:
             kwargs['range'] = [feature.min(), feature.max()]
-            self.add_value_data(value[:, :3], value[:, 3], name + "_value_cp",
+            self.add_value_data(self.model.rescale(value[:, :3],inplace=False), value[:, 3], name + "_value_cp",
                                 **kwargs)
         if tang.shape[0] > 0 and add_tang:
-            self.add_vector_data(tang[:, :3], tang[:, 3:6], name + "_tang_cp",
+            self.add_vector_data(self.model.rescale(tang[:, :3],inplace=False), tang[:, 3:6], name + "_tang_cp",
                                  **kwargs)
+        if interface.shape[0] > 0 and add_interface:
+            self.add_points(self.model.rescale(interface[:,:3],inplace=False), name + "_interface_cp")
 
-
+    def add_intersection_lineation(self, feature, **kwargs):
+        name = feature.name
+        if 'name' in kwargs:
+            name = kwargs['name']
+            del kwargs['name']
+        intersection = feature.fold.foldframe.calculate_intersection_lineation(
+            feature.builder)
+        gpoints = feature.builder.interpolator.get_gradient_constraints()[:,:6]
+        npoints = feature.builder.interpolator.get_norm_constraints()[:,:6]
+        points = []
+        if gpoints.shape[0] > 0:
+            points.append(gpoints)
+        if npoints.shape[0] > 0:
+            points.append(npoints)
+        points = np.vstack(points)
+        if intersection.shape[0] > 0:
+            self.add_vector_data(self.model.rescale(points[:,:3],inplace=False), intersection, name + "_intersection")
+        
     def add_points(self, points, name, **kwargs):
         """
 
@@ -523,7 +766,7 @@ class LavaVuModelViewer:
         if "pointsize" not in kwargs:
             kwargs["pointsize"] = 4
         # set the colour map to diverge unless user decides otherwise
-        cmap = kwargs.get('cmap', "spot")
+        cmap = kwargs.get('cmap', "rainbow")
         p = self.lv.points(name, **kwargs)
         p.vertices(position)
         p.values(value, "v")
@@ -533,7 +776,7 @@ class LavaVuModelViewer:
             logger.info('vmin {} and vmax {}'.format(kwargs['vmin'],kwargs['vmax']))
             p.colourmap(cmap, range=(kwargs['vmin'],kwargs['vmax']))
         else:
-            p.colourmap(cmap)
+            p.colourmap(cmap, range=(np.nanmin(value),np.nanmax(value)))
 
     def add_fold(self, fold, **kwargs):
         """
@@ -556,6 +799,7 @@ class LavaVuModelViewer:
             xx, yy, zz = np.meshgrid(x, y, z, indexing='ij')
             locations = np.array([xx.flatten(), yy.flatten(), zz.flatten()]).T
         r2r, fold_axis, dgz = fold.get_deformed_orientation(locations)
+        locations = self.model.rescale(locations,inplace=False)
         self.add_vector_data(locations, r2r, fold.name + '_direction', colour='red')
         self.add_vector_data(locations, fold_axis, fold.name + '_axis', colour='black')
         self.add_vector_data(locations, dgz, fold.name + '_norm', colour='green')
@@ -578,7 +822,26 @@ class LavaVuModelViewer:
             self.lv.control.ObjectList()
             self.lv.interactive()
 
+    def add_support_box(self,geological_feature, paint=False, **kwargs):
+        name = kwargs.get('name', geological_feature.name + '_support')
+        box = np.vstack([geological_feature.interpolator.support.origin,geological_feature.interpolator.support.maximum])
+        points, tri = create_box(box,self.nsteps)
 
+        surf = self.lv.triangles(name)
+        surf.vertices(self.model.rescale(points))
+        surf.indices(tri)
+        if paint:
+            val =geological_feature.evaluate_value(self.model.scale(points))
+            surf.values(val, geological_feature.name)
+            surf["colourby"] = geological_feature.name
+            cmap = kwargs.get('cmap',lavavu.cubehelix(100))
+
+            logger.info("Adding scalar field of %s to viewer. Min: %f, max: %f" % (geological_feature.name,
+                                                                                geological_feature.min(),
+                                                                                geological_feature.max()))
+            vmin = kwargs.get('vmin', np.nanmin(val))
+            vmax = kwargs.get('vmax', np.nanmax(val))
+            surf.colourmap(cmap, range=(vmin, vmax))
     def set_zscale(self,zscale):
         """ Set the vertical scale for lavavu
 
@@ -648,6 +911,16 @@ class LavaVuModelViewer:
 
         """
         self.lv.image(name)
+    
+    def image_array(self, **kwargs):
+        """Return the current viewer image image data as a numpy array
+
+        Returns
+        -------
+        image : np.array
+            image as a numpy array
+        """
+        return self.lv.rawimage(**kwargs).data
 
     def rotatex(self, r):
         """
@@ -730,3 +1003,78 @@ class LavaVuModelViewer:
         """
         self.lv.rotation(xyz)
 
+    @property
+    def border(self):
+        """The width of the border around the model area
+
+        Returns
+        -------
+        border : double
+            [description]
+        """
+        return self.lv['border']
+    
+    @border.setter
+    def border(self, border):
+        """Setter for the border
+
+        Parameters
+        ----------
+        border : double
+            set the thickness of the border around objects
+        """
+        self.lv['border'] = border
+
+    def clear(self):
+        """Remove all objects from the viewer
+        """
+        self.lv.clear()
+    
+    @property
+    def xmin(self):
+        return self.lv['xmin']
+    
+    @xmin.setter
+    def xmin(self, xmin):
+        self.lv['xmin'] = xmin
+
+    @property
+    def xmax(self):
+        return self.lv['xmax']
+    
+    @xmax.setter
+    def xmax(self, xmax):
+        self.lv['xmax'] = xmax
+
+    @property
+    def ymin(self):
+        return self.lv['ymin']
+    
+    @ymin.setter
+    def ymin(self, ymin):
+        self.lv['ymin'] = ymin
+
+    @property
+    def ymax(self):
+        return self.lv['ymax']
+    
+    @ymax.setter
+    def ymax(self, ymax):
+        self.lv['ymax'] = ymax
+    
+    @property
+    def zmin(self):
+        return self.lv['zmax']
+
+    @zmin.setter
+    def zmin(self, zmin):
+        self.lv['zmin'] = zmin
+
+    @property
+    def zmax(self):
+        return self.lv['zmax']
+    
+    @zmax.setter
+    def zmax(self, zmax):
+        self.lv['zmax'] = zmax
+