diff --git a/demos/cornell_box.py b/demos/cornell_box.py
index f4d89c8..9b08fae 100644
--- a/demos/cornell_box.py
+++ b/demos/cornell_box.py
@@ -110,12 +110,12 @@
 # objects in enclosure
 box = Box(Point3D(-0.4, 0, -0.4), Point3D(0.3, 1.4, 0.3),
           parent=world,
-          transform=translate(0.4, -1 + 1e-6, 0.4)*rotate(30, 0, 0),
+          transform=translate(0.4, -1 + 1e-6, 0.4) * rotate(30, 0, 0),
           material=schott("N-BK7"))
 
 sphere = Sphere(0.4,
                 parent=world,
-                transform=translate(-0.4, -0.6 + 1e-6, -0.4)*rotate(0, 0, 0),
+                transform=translate(-0.4, -0.6 + 1e-6, -0.4) * rotate(0, 0, 0),
                 material=schott("N-BK7"))
 
 
@@ -154,6 +154,16 @@
 camera.ray_extinction_prob = 0.01
 
 
-from raysect_mayavi import visualise_scenegraph
+# Visualising with mayavi
+from raysect_mayavi.mayavi import visualise_scenegraph
+from mayavi import mlab
 
-visualise_scenegraph(world)
+fig = visualise_scenegraph(world)
+mlab.show()
+
+
+# Visualising with pyvista
+from raysect_mayavi.pyvista import visualise_scenegraph
+
+plotter = visualise_scenegraph(world)
+plotter.show()
diff --git a/demos/logging_trajectories.py b/demos/logging_trajectories.py
index 3ebe946..fe9b505 100644
--- a/demos/logging_trajectories.py
+++ b/demos/logging_trajectories.py
@@ -6,10 +6,10 @@
 import numpy as np
 
 # Raysect imports
-from raysect.optical import World, translate, rotate, Point3D, d65_white, Ray, Vector3D
+from raysect.optical import World, translate, Point3D, Vector3D
 from raysect.optical.material.absorber import AbsorbingSurface
 from raysect.optical.library import schott
-from raysect.primitive import Sphere, Box
+from raysect.primitive import Box
 from raysect.optical.loggingray import LoggingRay
 from raysect.primitive.lens.spherical import *
 
@@ -20,7 +20,7 @@
 lens_glass = schott("N-BK7")
 lens_glass.transmission_only = True
 lens = BiConvex(0.0254, 0.0052, 0.0506, 0.0506, parent=world, material=lens_glass)
-lens.meta['viz-color'] = (66/255, 188/255, 244/255)
+lens.meta['viz-color'] = (66 / 255, 188 / 255, 244 / 255)
 lens.meta['viz-opacity'] = 0.5
 
 # lens.meta['viz']['color'] = (66/255, 188/255, 244/255)
@@ -29,7 +29,7 @@
 # Create a target plane behind the lens.
 target = Box(lower=Point3D(-0.05, -0.05, -0), upper=Point3D(0.05, 0.05, 0), material=AbsorbingSurface(),
              transform=translate(0, 0, 0.1), parent=world)
-target.meta['viz-color'] = (224/255, 100/255, 17/255)
+target.meta['viz-color'] = (224 / 255, 100 / 255, 17 / 255)
 
 
 # for each sample direction trace a logging ray and plot the ray trajectory
@@ -51,21 +51,41 @@
 
     plt.plot(p[:, 0], p[:, 1], 'k-')
     plt.plot(p[:, 0], p[:, 1], 'r.')
+plt.draw()
+plt.pause(0.01)
 
 
-from raysect_mayavi import visualise_scenegraph
+# Visualising with mayavi
+from raysect_mayavi.mayavi import visualise_scenegraph, assign_visualiser
 from raysect_mayavi import to_mesh
 from mayavi import mlab
 
-
 fig = visualise_scenegraph(world)
 
 for v in np.linspace(-0.012, 0.012, 11):
     start = Point3D(v, 0, -0.05)
     log_ray = LoggingRay(start, Vector3D(0, 0, 1))
     log_ray.trace(world)
-    log_ray_mayavi = to_mesh(log_ray)
-    log_ray_mayavi.mayavi_plot(fig)
+    log_ray_mayavi = assign_visualiser(to_mesh(log_ray))
+    log_ray_mayavi.plot(fig)
+
+mlab.show()
+
+
+# Visualising with pyvista
+from raysect_mayavi.pyvista import visualise_scenegraph, assign_visualiser
+from raysect_mayavi import to_mesh
+
+plotter = visualise_scenegraph(world)
+
+for v in np.linspace(-0.012, 0.012, 11):
+    start = Point3D(v, 0, -0.05)
+    log_ray = LoggingRay(start, Vector3D(0, 0, 1))
+    log_ray.trace(world)
+    log_ray_mayavi = assign_visualiser(to_mesh(log_ray))
+    log_ray_mayavi.plot(plotter)
+
+plotter.show()
 
 plt.ioff()
 plt.show()
diff --git a/demos/mesh_csg/mesh_csg_examples.py b/demos/mesh_csg/mesh_csg_examples_mayavi.py
similarity index 87%
rename from demos/mesh_csg/mesh_csg_examples.py
rename to demos/mesh_csg/mesh_csg_examples_mayavi.py
index d553308..576a601 100644
--- a/demos/mesh_csg/mesh_csg_examples.py
+++ b/demos/mesh_csg/mesh_csg_examples_mayavi.py
@@ -4,9 +4,9 @@
 from raysect.core import translate, Point3D, rotate_basis, Vector3D, rotate
 from raysect.optical import World
 from raysect.primitive import Box, Sphere, Cylinder, Cone
-from raysect.primitive import Sphere, Mesh, Intersect, Subtract, Union
+from raysect.primitive import Intersect, Subtract, Union
 
-from raysect_mayavi import visualise_scenegraph
+from raysect_mayavi.mayavi import visualise_scenegraph
 
 
 ########################################################################################################################
@@ -105,7 +105,7 @@
 # Cone and Cylinder
 
 c1 = Cone(0.15, 1)
-c2 = Cylinder(0.15, 0.5, transform=translate(-0.25, 0, 0.5)*rotate_basis(Vector3D(1, 0, 0), Vector3D(0, 0, 1)))
+c2 = Cylinder(0.15, 0.5, transform=translate(-0.25, 0, 0.5) * rotate_basis(Vector3D(1, 0, 0), Vector3D(0, 0, 1)))
 
 world = World()
 Union(c1, c2, parent=world)
@@ -127,9 +127,9 @@
 # CSG Hell
 
 world = World()
-cyl_x = Cylinder(1, 4.2, transform=rotate(90, 0, 0)*translate(0, 0, -2.1))
-cyl_y = Cylinder(1, 4.2, transform=rotate(0, 90, 0)*translate(0, 0, -2.1))
-cyl_z = Cylinder(1, 4.2, transform=rotate(0, 0, 0)*translate(0, 0, -2.1))
+cyl_x = Cylinder(1, 4.2, transform=rotate(90, 0, 0) * translate(0, 0, -2.1))
+cyl_y = Cylinder(1, 4.2, transform=rotate(0, 90, 0) * translate(0, 0, -2.1))
+cyl_z = Cylinder(1, 4.2, transform=rotate(0, 0, 0) * translate(0, 0, -2.1))
 cube = Box(Point3D(-1.5, -1.5, -1.5), Point3D(1.5, 1.5, 1.5))
 sphere = Sphere(2.0)
 
diff --git a/demos/mesh_csg/mesh_csg_examples_pyvista.py b/demos/mesh_csg/mesh_csg_examples_pyvista.py
new file mode 100644
index 0000000..e56e83e
--- /dev/null
+++ b/demos/mesh_csg/mesh_csg_examples_pyvista.py
@@ -0,0 +1,118 @@
+
+from raysect.core import translate, Point3D, rotate_basis, Vector3D, rotate
+from raysect.optical import World
+from raysect.primitive import Box, Sphere, Cylinder, Cone
+from raysect.primitive import Intersect, Subtract, Union
+
+from raysect_mayavi.pyvista import visualise_scenegraph
+
+
+########################################################################################################################
+# Spheres
+
+
+s1 = Sphere(0.5, transform=translate(-0.25, 0, 0), name='s1')
+s2 = Sphere(0.5, transform=translate(0.25, 0, 0), name='s2')
+
+
+world = World()
+Union(s1, s2, parent=world)
+visualise_scenegraph(world).show()
+
+world = World()
+Intersect(s1, s2, parent=world)
+visualise_scenegraph(world).show()
+
+world = World()
+Subtract(s1, s2, parent=world)
+visualise_scenegraph(world).show()
+
+
+########################################################################################################################
+# Cubes
+
+b1 = Box(Point3D(0, 0, 0), Point3D(1, 1, 1))
+b2 = Box(Point3D(0, 0, 0), Point3D(1, 1, 1), transform=translate(0.6, 0.6, 0.6))
+
+world = World()
+Union(b1, b2, parent=world)
+visualise_scenegraph(world).show()
+
+world = World()
+Intersect(b1, b2, parent=world)
+visualise_scenegraph(world).show()
+
+world = World()
+Subtract(b1, b2, parent=world)
+visualise_scenegraph(world).show()
+
+########################################################################################################################
+# Cylinders
+
+c1 = Cylinder(0.5, 2)
+c2 = Cylinder(0.5, 2, transform=translate(0.15, 0, 0.4))
+
+world = World()
+Union(c1, c2, parent=world)
+visualise_scenegraph(world).show()
+
+world = World()
+Intersect(c1, c2, parent=world)
+visualise_scenegraph(world).show()
+
+world = World()
+Subtract(c1, c2, parent=world)
+visualise_scenegraph(world).show()
+
+
+########################################################################################################################
+# Box and Sphere
+
+s1 = Sphere(0.5)
+b2 = Box(Point3D(-0.5, -0.5, -0.5), Point3D(0.5, 0.5, 0.5), transform=translate(-0.2, 0, -0.2))
+
+world = World()
+Union(s1, b2, parent=world)
+visualise_scenegraph(world).show()
+
+world = World()
+Intersect(s1, b2, parent=world)
+visualise_scenegraph(world).show()
+
+world = World()
+Subtract(s1, b2, parent=world)
+visualise_scenegraph(world)
+
+
+########################################################################################################################
+# Cone and Cylinder
+
+c1 = Cone(0.15, 1)
+c2 = Cylinder(0.15, 0.5, transform=translate(-0.25, 0, 0.5) * rotate_basis(Vector3D(1, 0, 0), Vector3D(0, 0, 1)))
+
+world = World()
+Union(c1, c2, parent=world)
+visualise_scenegraph(world).show()
+
+world = World()
+Intersect(c1, c2, parent=world)
+visualise_scenegraph(world).show()
+
+world = World()
+Subtract(c1, c2, parent=world)
+visualise_scenegraph(world).show()
+
+
+########################################################################################################################
+# CSG Hell
+
+world = World()
+cyl_x = Cylinder(1, 4.2, transform=rotate(90, 0, 0) * translate(0, 0, -2.1))
+cyl_y = Cylinder(1, 4.2, transform=rotate(0, 90, 0) * translate(0, 0, -2.1))
+cyl_z = Cylinder(1, 4.2, transform=rotate(0, 0, 0) * translate(0, 0, -2.1))
+cube = Box(Point3D(-1.5, -1.5, -1.5), Point3D(1.5, 1.5, 1.5))
+sphere = Sphere(2.0)
+
+csg = Intersect(sphere, Subtract(cube, Union(Union(cyl_x, cyl_y), cyl_z)), world)
+
+visualise_scenegraph(world).show()
diff --git a/demos/optical_elements.py b/demos/optical_elements.py
index 15b167c..f50db27 100644
--- a/demos/optical_elements.py
+++ b/demos/optical_elements.py
@@ -1,11 +1,8 @@
-from raysect_mayavi import visualise_scenegraph
-
 from raysect.core import translate
 from raysect.optical import World
 from raysect.primitive import Cylinder
 from raysect.primitive.lens.spherical import BiConvex, BiConcave, PlanoConvex, PlanoConcave, Meniscus
 
-from mayavi import mlab
 # Display lens and cylinder primitives
 
 world = World()
@@ -22,5 +19,15 @@
 cylinder_primitive = Cylinder(radius=0.5 * diameter, height=center_thickness, parent=world, transform=translate(0, 0, 5))
 
 
+# Visualising with mayavi
+from raysect_mayavi.mayavi import visualise_scenegraph
+from mayavi import mlab
+
 visualise_scenegraph(world)
-mlab.show()
\ No newline at end of file
+mlab.show()
+
+# Visualising with pyvista
+from raysect_mayavi.pyvista import visualise_scenegraph
+
+plotter = visualise_scenegraph(world)
+plotter.show()
diff --git a/raysect_mayavi/mayavi/__init__.py b/raysect_mayavi/mayavi/__init__.py
index e69de29..13468db 100644
--- a/raysect_mayavi/mayavi/__init__.py
+++ b/raysect_mayavi/mayavi/__init__.py
@@ -0,0 +1,2 @@
+from .scenegraph_viewer import visualise_scenegraph
+from .parse import assign_visualiser
diff --git a/raysect_mayavi/mayavi/mesh.py b/raysect_mayavi/mayavi/mesh.py
index 5316f7f..6a64335 100644
--- a/raysect_mayavi/mayavi/mesh.py
+++ b/raysect_mayavi/mayavi/mesh.py
@@ -8,20 +8,19 @@ class TriangularMeshVisualiser(MayaviVisualiser):
     """
     This class serves as base class for Raysect objetcs visualised with triangular meshses.
     """
+
     def __init__(self, source):
 
-        
         self._init_plot_kwargs()
 
-        
         super().__init__(source)
         self.plot_method = mlab.triangular_mesh
 
     def set_source(self, source):
 
         if not isinstance(source, TriangularMeshSource):
-            raise TypeError("source has to be instance of SourceBase.")
-        
+            raise TypeError("source has to be instance of TriangularMeshSource.")
+
         self._source = source
 
     def _init_plot_kwargs(self):
@@ -33,5 +32,5 @@ def _init_plot_kwargs(self):
     def _mayavi_plot(self, figure):
         vertices = self._source.vertices
 
-        mlab.triangular_mesh(vertices[:,0], vertices[:,1], vertices[:,2], self._source.triangles,
-                             figure=figure, **self.plot_kwargs)
\ No newline at end of file
+        mlab.triangular_mesh(vertices[:, 0], vertices[:, 1], vertices[:, 2], self._source.triangles,
+                             figure=figure, **self.plot_kwargs)
diff --git a/raysect_mayavi/mayavi/observer.py b/raysect_mayavi/mayavi/observer.py
index 3789c16..68bd1fa 100644
--- a/raysect_mayavi/mayavi/observer.py
+++ b/raysect_mayavi/mayavi/observer.py
@@ -9,6 +9,7 @@
 class ObserverVisualiser(MayaviVisualiser):
     """
     Class for graphical representation of the Raysect Box primitive.
+
     :param raysect_object: Raysect Observer instance
     """
 
@@ -25,8 +26,8 @@ def _init_plot_kwargs(self):
     def set_source(self, source):
 
         if not isinstance(source, ObserverSource):
-            raise TypeError("source has to be instance of SourceBase.")
-        
+            raise TypeError("source has to be instance of ObserverSource.")
+
         self._source = source
 
     def _mayavi_plot(self, figure):
@@ -42,4 +43,4 @@ def _mayavi_plot(self, figure):
         w = np.array([direction_raysect.z], ndmin=1)
 
         mlab.quiver3d(x, y, z, u, v, w,
-                    figure=figure, **self.plot_kwargs)
+                      figure=figure, **self.plot_kwargs)
diff --git a/raysect_mayavi/mayavi/parse.py b/raysect_mayavi/mayavi/parse.py
index 7f41095..46ed5ca 100644
--- a/raysect_mayavi/mayavi/parse.py
+++ b/raysect_mayavi/mayavi/parse.py
@@ -9,17 +9,17 @@
 
 
 def parse_sources(sources):
-    
+
     visualisers = {}
-    
+
     for key, source in sources.items():
         visualisers[key] = assign_visualiser(source)
 
-    
     return visualisers
 
+
 def assign_visualiser(source):
-    
+
     if not isinstance(source, SourceBase):
         raise TypeError("sources items has to be of type SourceBase.")
     if isinstance(source, TriangularMeshSource):
@@ -30,4 +30,3 @@ def assign_visualiser(source):
         return ObserverVisualiser(source)
     else:
         raise TypeError("Source type not recognised")
-    
\ No newline at end of file
diff --git a/raysect_mayavi/mayavi/ray.py b/raysect_mayavi/mayavi/ray.py
index ac7d4e0..4d19678 100644
--- a/raysect_mayavi/mayavi/ray.py
+++ b/raysect_mayavi/mayavi/ray.py
@@ -1,5 +1,3 @@
-from raysect_mayavi.mayavi.visualiser import MayaviVisualiser
-
 from raysect_mayavi.primitives.ray import LoggingRaySource
 
 from raysect_mayavi.mayavi.visualiser import MayaviVisualiser
@@ -9,8 +7,10 @@
 
 class LoggingRayVisualiser(MayaviVisualiser):
     """
-    Class for graphical representation of the Raysect LoggingRay primitive. To display correctly with other scenegraph 
-    components observation has to be usualy done in the root node (an instance of World) of the scenegraph.
+    Class for graphical representation of the Raysect LoggingRay primitive.
+    To display correctly with other scenegraph  components observation has to be usualy
+    done in the root node (an instance of World) of the scenegraph.
+
     :param raysect_object: Raysect Loggingray instance containing path_vertices
     """
 
@@ -23,16 +23,16 @@ def __init__(self, source):
     def _init_plot_kwargs(self):
         self.plot_kwargs = {}
         self.plot_kwargs["tube_radius"] = None
-        
+
     def set_source(self, source):
 
         if not isinstance(source, LoggingRaySource):
-            raise TypeError("source has to be instance of SourceBase.")
-        
+            raise TypeError("source has to be instance of LoggingRaySource.")
+
         self._source = source
 
     def _mayavi_plot(self, figure):
         vertices = self._source._vertices
 
-        mlab.plot3d(vertices[:,0], vertices[:,1], vertices[:,2],
+        mlab.plot3d(vertices[:, 0], vertices[:, 1], vertices[:, 2],
                     figure=figure, **self.plot_kwargs)
diff --git a/raysect_mayavi/mayavi/scenegraph_viewer.py b/raysect_mayavi/mayavi/scenegraph_viewer.py
index 0f0730c..7135c4a 100644
--- a/raysect_mayavi/mayavi/scenegraph_viewer.py
+++ b/raysect_mayavi/mayavi/scenegraph_viewer.py
@@ -1,18 +1,24 @@
 
-from mayavi import mlab
-
 from raysect_mayavi.primitives.parse_nodes import parse_nodes
 from raysect_mayavi.mayavi.parse import parse_sources
 
+from mayavi import mlab
+from mayavi.core.scene import Scene
+
 
 def visualise_scenegraph(world, figure=None, show_axes=False, axes_length=1):
     """
     Automatic visualisation of the raysect scenegraph.
-    :param world: Instance of the Raysect World 
+
+    :param world: Instance of the Raysect World
+    :param figure: Optional, specifies the mlab.figure to plot in.
+
+    :return: mlab.figure
     """
 
-    if figure is None:
-        figure = mlab.figure(size=(1024, 768), bgcolor=(1, 1, 1), fgcolor=(0.5, 0.5, 0.5))
+    figure = figure or mlab.figure(size=(1024, 768), bgcolor=(1, 1, 1), fgcolor=(0.5, 0.5, 0.5))
+    if not isinstance(figure, Scene):
+        raise ValueError("figure has to be instance of mlab.figure")
 
     sources = parse_nodes(world)
     visualisers = parse_sources(sources)
@@ -22,8 +28,8 @@ def visualise_scenegraph(world, figure=None, show_axes=False, axes_length=1):
         visualiser.plot(figure)
 
     if show_axes:
-        mlab.plot3d([0, axes_length], [0, 0], [0, 0], tube_radius=axes_length/100, color=(1, 0, 0))
-        mlab.plot3d([0, 0], [0, axes_length], [0, 0], tube_radius=axes_length/100, color=(0, 1, 0))
-        mlab.plot3d([0, 0], [0, 0], [0, axes_length], tube_radius=axes_length/100, color=(0, 0, 1))
+        mlab.plot3d([0, axes_length], [0, 0], [0, 0], tube_radius=axes_length / 100, color=(1, 0, 0))
+        mlab.plot3d([0, 0], [0, axes_length], [0, 0], tube_radius=axes_length / 100, color=(0, 1, 0))
+        mlab.plot3d([0, 0], [0, 0], [0, axes_length], tube_radius=axes_length / 100, color=(0, 0, 1))
 
-    return figure
\ No newline at end of file
+    return figure
diff --git a/raysect_mayavi/mayavi/visualiser.py b/raysect_mayavi/mayavi/visualiser.py
index 5970050..7e6e21b 100644
--- a/raysect_mayavi/mayavi/visualiser.py
+++ b/raysect_mayavi/mayavi/visualiser.py
@@ -1,22 +1,22 @@
-from raysect.core import Point3D
 
-from raysect_mayavi.primitives.source import SourceBase, VisualiserBase, TriangularMeshSource
+from raysect_mayavi.primitives.source import VisualiserBase
 
 from mayavi import mlab
 from mayavi.core.scene import Scene
 
-class MayaviVisualiser(VisualiserBase):
 
+class MayaviVisualiser(VisualiserBase):
     """
     This is the base class for the raysect_mayavi representation of raysect objects.
     """
+
     def __init__(self, source):
 
         self.set_source(source)
         self._init_figure_kwargs()
 
     def _init_figure_kwargs(self):
-        self.figure_kwargs= {}
+        self.figure_kwargs = {}
         self.figure_kwargs["size"] = (1024, 768)
         self.figure_kwargs["bgcolor"] = (1, 1, 1)
         self.figure_kwargs["fgcolor"] = (0.5, 0.5, 0.5)
@@ -24,20 +24,21 @@ def _init_figure_kwargs(self):
     @property
     def source(self):
         return self._source
-    
+
     def set_source(self, source):
         raise NotImplementedError("Virtual method _set_source() has not been implemented.")
-    
+
     def plot(self, figure=None):
         """
         Plot the Mayavi representation of the Raysect object into the figure. The representation is done always
         in the root node.
-        :param figure:Optional, specifies the figure to plot in.
-        return mayavi figure
+
+        :param figure: Optional, specifies the mlab.figure to plot in.
+
+        :return: mlab.figure
         """
-        if figure is None:
-            figure = mlab.figure(**self.figure_kwargs)
-        elif not isinstance(figure, Scene):
+        figure = figure or mlab.figure(**self.figure_kwargs)
+        if not isinstance(figure, Scene):
             raise ValueError("figure has to be instance of mlab.figure")
 
         self._mayavi_plot(figure)
diff --git a/raysect_mayavi/primitives/geometric/cylindrical.py b/raysect_mayavi/primitives/geometric/cylindrical.py
index bbbd472..ed9d580 100644
--- a/raysect_mayavi/primitives/geometric/cylindrical.py
+++ b/raysect_mayavi/primitives/geometric/cylindrical.py
@@ -421,7 +421,7 @@ def __init__(self, raysect_object, vertical_divisions=10, cylindrical_divisions=
 
         super().__init__(raysect_object, vertical_divisions, cylindrical_divisions, radial_divisions)
 
-    def _generate_barrel_mesh(self):
+    def _generate_barrel_surface_mesh(self):
 
         radius = 0.5 * self.raysect_object.diameter
         height = self.raysect_object.edge_thickness
diff --git a/raysect_mayavi/pyvista/__init__.py b/raysect_mayavi/pyvista/__init__.py
index e69de29..13468db 100644
--- a/raysect_mayavi/pyvista/__init__.py
+++ b/raysect_mayavi/pyvista/__init__.py
@@ -0,0 +1,2 @@
+from .scenegraph_viewer import visualise_scenegraph
+from .parse import assign_visualiser
diff --git a/raysect_mayavi/pyvista/mesh.py b/raysect_mayavi/pyvista/mesh.py
index 51ea41c..13f892a 100644
--- a/raysect_mayavi/pyvista/mesh.py
+++ b/raysect_mayavi/pyvista/mesh.py
@@ -1,5 +1,3 @@
-import numpy as np
-
 from raysect_mayavi.primitives.mesh import TriangularMeshSource
 from raysect_mayavi.pyvista.visualiser import PyvistaVisualiser
 
@@ -11,21 +9,20 @@ class TriangularMeshVisualiser(PyvistaVisualiser):
     """
     This class serves as base class for Raysect objetcs visualised with triangular meshses.
     """
+
     def __init__(self, source):
 
-        
         self._init_plot_kwargs()
 
-        
         super().__init__(source)
 
     def set_source(self, source):
 
         if not isinstance(source, TriangularMeshSource):
-            raise TypeError("source has to be instance of SourceBase.")
-        
+            raise TypeError("source has to be instance of TriangularMeshSource.")
+
         self._source = source
-    
+
     def get_data_object(self):
 
         vertices = self._source.vertices
@@ -33,8 +30,7 @@ def get_data_object(self):
 
         return pv.make_tri_mesh(vertices, triangles)
 
-        
     def _add_object(self, plotter):
-        
+
         mesh = self.get_data_object()
-        plotter.add_mesh(mesh)
\ No newline at end of file
+        plotter.add_mesh(mesh)
diff --git a/raysect_mayavi/pyvista/observer.py b/raysect_mayavi/pyvista/observer.py
index 6971557..ba7ad8c 100644
--- a/raysect_mayavi/pyvista/observer.py
+++ b/raysect_mayavi/pyvista/observer.py
@@ -3,58 +3,58 @@
 from raysect_mayavi.primitives.observer import ObserverSource
 from raysect_mayavi.pyvista.visualiser import PyvistaVisualiser
 
-import pyvista as pv 
+import pyvista as pv
 pv.rcParams['use_ipyvtk'] = True
 
+
 class ObserverVisualiser(PyvistaVisualiser):
     """
     Class for graphical representation of the Raysect Box primitive.
+
     :param raysect_object: Raysect Observer instance
     """
 
     def __init__(self, source):
-        
+
         self.length = 0.1
         super().__init__(source)
 
     def _init_plot_kwargs(self):
-       self.plot_kwargs = {}
-       self.plot_kwargs["tip_length"] = 2 * self.length 
-       self.plot_kwargs["tip_radius"] = 2 * self.length
-       self.plot_kwargs["tip_resolution"] = 20
-       self.plot_kwargs["shaft_radius"] = 0.2 * self.length
-       self.plot_kwargs["shaft_resolution"] = 20
-    
+        self.plot_kwargs = {}
+        self.plot_kwargs["tip_length"] = 2 * self.length
+        self.plot_kwargs["tip_radius"] = 2 * self.length
+        self.plot_kwargs["tip_resolution"] = 20
+        self.plot_kwargs["shaft_radius"] = 0.2 * self.length
+        self.plot_kwargs["shaft_resolution"] = 20
+
     @property
     def length(self):
         return self._length
-    
+
     @length.setter
     def length(self, value):
 
         if not value > 0:
             raise ValueError("value has to be larger than 0")
-        
+
         self._length = value
-        
 
     def set_source(self, source):
 
         if not isinstance(source, ObserverSource):
-            raise TypeError("source has to be instance of SourceBase.")
-        
+            raise TypeError("source has to be instance of ObserverSource.")
+
         self._source = source
-    
+
     def get_data_object(self):
         origin_vector = self._source.origin_root
         origin = np.array([origin_vector.x, origin_vector.y, origin_vector.z])
         direction_vector = self.length * self._source.direction_root.normalise()
         direction = np.array([direction_vector.x, direction_vector.y, direction_vector.z])
-        
+
         return pv.Arrow(start=origin, direction=direction, **self.plot_kwargs)
 
     def _add_object(self, plotter):
 
         mesh = self.get_data_object()
         plotter.add_mesh(mesh, "r")
-
diff --git a/raysect_mayavi/pyvista/parse.py b/raysect_mayavi/pyvista/parse.py
index 1574795..a691e5c 100644
--- a/raysect_mayavi/pyvista/parse.py
+++ b/raysect_mayavi/pyvista/parse.py
@@ -9,17 +9,17 @@
 
 
 def parse_sources(sources):
-    
+
     visualisers = {}
-    
+
     for key, source in sources.items():
         visualisers[key] = assign_visualiser(source)
 
-    
     return visualisers
 
+
 def assign_visualiser(source):
-    
+
     if not isinstance(source, SourceBase):
         raise TypeError("sources items has to be of type SourceBase.")
     if isinstance(source, TriangularMeshSource):
@@ -30,4 +30,3 @@ def assign_visualiser(source):
         return ObserverVisualiser(source)
     else:
         raise TypeError("Source type not recognised")
-    
\ No newline at end of file
diff --git a/raysect_mayavi/pyvista/ray.py b/raysect_mayavi/pyvista/ray.py
index 7102036..830db1a 100644
--- a/raysect_mayavi/pyvista/ray.py
+++ b/raysect_mayavi/pyvista/ray.py
@@ -1,28 +1,29 @@
-import numpy as np
-
 from raysect_mayavi.primitives.ray import LoggingRaySource
 from raysect_mayavi.pyvista.visualiser import PyvistaVisualiser
 
-import pyvista as pv 
+import pyvista as pv
 pv.rcParams['use_ipyvtk'] = True
 
+
 class LoggingRayVisualiser(PyvistaVisualiser):
     """
-    Class for graphical representation of the Raysect LoggingRay primitive. To display correctly with other scenegraph 
-    components observation has to be usualy done in the root node (an instance of World) of the scenegraph.
+    Class for graphical representation of the Raysect LoggingRay primitive. To display correctly
+    with other scenegraph components observation has to be usualy done in the root node
+    (an instance of World) of the scenegraph.
+
     :param raysect_object: Raysect Loggingray instance containing path_vertices
     """
 
     def __init__(self, source):
-        
+
         self._init_plot_kwargs()
         super().__init__(source)
 
     def set_source(self, source):
 
         if not isinstance(source, LoggingRaySource):
-            raise TypeError("source has to be instance of SourceBase.")
-        
+            raise TypeError("source has to be instance of LoggingRaySource.")
+
         self._source = source
 
     def get_data_object(self):
@@ -36,4 +37,4 @@ def _add_object(self, plotter):
 
         mesh = self.get_data_object()
         if mesh is not None:
-            plotter.add_mesh(mesh, color="g")
\ No newline at end of file
+            plotter.add_mesh(mesh, color="g")
diff --git a/raysect_mayavi/pyvista/scenegraph_viewer.py b/raysect_mayavi/pyvista/scenegraph_viewer.py
index edcf467..d0fac80 100644
--- a/raysect_mayavi/pyvista/scenegraph_viewer.py
+++ b/raysect_mayavi/pyvista/scenegraph_viewer.py
@@ -1,25 +1,29 @@
+from raysect_mayavi.primitives.parse_nodes import parse_nodes
+from raysect_mayavi.pyvista.parse import parse_sources
+
 import pyvista as pv
 pv.rcParams['use_ipyvtk'] = True
 
-from raysect_mayavi.primitives.parse_nodes import parse_nodes
-from raysect_mayavi.mayavi.parse import parse_sources
-
 
 def visualise_scenegraph(world, plotter=None, show_axes=False, axes_length=1):
     """
     Automatic visualisation of the raysect scenegraph.
-    :param world: Instance of the Raysect World 
+
+    :param world: Instance of the Raysect World
+    :param plotter: Optional, specifies the pyvista.plotter to plot in.
+
+    :return: pyvista.plotter
     """
 
-    if plotter is None:
-        figure = pv.Plotter(window_size=(1024, 768))
+    plotter = plotter or pv.Plotter(window_size=(1024, 768))
+    if not isinstance(plotter, pv.Plotter):
+        raise ValueError("plotter has to be instance of pyvista.Plotter.")
 
     sources = parse_nodes(world)
     visualisers = parse_sources(sources)
 
     for _, visualiser in visualisers.items():
 
-        visualiser.plot(figure)
-
+        visualiser.plot(plotter)
 
-    return plotter
\ No newline at end of file
+    return plotter
diff --git a/raysect_mayavi/pyvista/visualiser.py b/raysect_mayavi/pyvista/visualiser.py
index 697b274..1366eae 100644
--- a/raysect_mayavi/pyvista/visualiser.py
+++ b/raysect_mayavi/pyvista/visualiser.py
@@ -1,14 +1,14 @@
-from raysect.core import Point3D
+from raysect_mayavi.primitives.source import VisualiserBase
 
-from raysect_mayavi.primitives.source import SourceBase, VisualiserBase, TriangularMeshSource
 import pyvista as pv
 pv.rcParams['use_ipyvtk'] = True
 
-class PyvistaVisualiser(VisualiserBase):
 
+class PyvistaVisualiser(VisualiserBase):
     """
     This is the base class for the raysect_mayavi representation of raysect objects.
     """
+
     def __init__(self, source):
 
         self.set_source(source)
@@ -16,7 +16,7 @@ def __init__(self, source):
         self._init_plot_kwargs()
 
     def _init_figure_kwargs(self):
-        self.figure_kwargs= {}
+        self.figure_kwargs = {}
         self.figure_kwargs["window_size"] = (512, 256)
 
     def _init_plot_kwargs(self):
@@ -25,7 +25,7 @@ def _init_plot_kwargs(self):
     @property
     def source(self):
         return self._source
-    
+
     def set_source(self, source):
         raise NotImplementedError("Virtual method _set_source() has not been implemented.")
 
@@ -34,18 +34,19 @@ def get_data_object(self):
         Constructs the pyvista.DataObject representation of the Raysect primitive
         """
         raise NotImplementedError("Virtual method get_vista_object() has not been implemented.")
-    
+
     def plot(self, plotter=None):
         """
-        Plot the Mayavi representation of the Raysect object into the figure. The representation is done always
-        in the root node.
-        :param figure:Optional, specifies the figure to plot in.
-        return mayavi figure
+        Plot the representation of the Raysect object into the pyvista.plotter.
+        The representation is done always in the root node.
+
+        :param plotter: Optional, specifies the pyvista.plotter to plot in.
+
+        :return: pyvista.plotter
         """
-        if plotter is None:
-            plotter = pv.Plotter(**self.figure_kwargs)
-        #elif not isinstance(plotter, pv.Plotter):
-        #    raise ValueError("figure has to be instance of mlab.figure")
+        plotter = plotter or pv.Plotter(**self.figure_kwargs)
+        if not isinstance(plotter, pv.Plotter):
+            raise ValueError("plotter has to be instance of pyvista.Plotter.")
 
         self._add_object(plotter)