diff --git a/README.md b/README.md
index 5cf3c57..f73ba65 100644
--- a/README.md
+++ b/README.md
@@ -12,15 +12,15 @@ pip install openmc_source_plotter
 
 # Features
 
-The package simply extends the default ```openmc.IndependentSourceBase``` to provides additional functions that:
+The package simply extends the default ```openmc.IndependentSourceBase``` and ```openmc.Model``` to provides additional functions that:
 
 - extract the positions, directions and energy of particles
-- visualise an ```openmc.IndependentSourceBase``` with respect to:
+- visualise a source with respect to:
   - direction
   - energy
   - position
 
-Or just sample the with ```openmc.Model.sample_initial_particles```
+Or just provide the initial particles with ```sample_initial_particles```
 
 # Example plots
 
diff --git a/src/openmc_source_plotter/core.py b/src/openmc_source_plotter/core.py
index 3219fb3..6c2554e 100644
--- a/src/openmc_source_plotter/core.py
+++ b/src/openmc_source_plotter/core.py
@@ -18,14 +18,6 @@ def sample_initial_particles(self, n_samples: int = 1000, prn_seed: int = None):
         self.materials.export_to_xml()
         self.geometry.export_to_xml()
 
-        openmc.lib.init(output=False)
-        particles = openmc.lib.sample_external_source(
-            n_samples=n_samples, prn_seed=prn_seed
-        )
-        openmc.lib.finalize()
-
-        return particles
-
     else:  # source object
 
         settings = openmc.Settings()
@@ -43,13 +35,13 @@ def sample_initial_particles(self, n_samples: int = 1000, prn_seed: int = None):
 
         geometry.export_to_xml()
 
-        openmc.lib.init(output=False)
-        particles = openmc.lib.sample_external_source(
-            n_samples=n_samples, prn_seed=prn_seed
-        )
-        openmc.lib.finalize()
+    openmc.lib.init(output=False)
+    particles = openmc.lib.sample_external_source(
+        n_samples=n_samples, prn_seed=prn_seed
+    )
+    openmc.lib.finalize()
 
-        return particles
+    return particles
 
 
 def plot_source_energy(
@@ -94,7 +86,8 @@ def plot_source_energy(
     probability, bin_edges = np.histogram(e_values, bins=energy_bins, density=True)
 
     # scaling by strength
-    probability = probability * self.strength
+    if isinstance(self, openmc.SourceBase):
+        probability = probability * self.strength
 
     # Plot source energy histogram
     figure.add_trace(
@@ -239,5 +232,13 @@ def plot_source_direction(
 openmc.Model.sample_initial_particles = sample_initial_particles
 
 openmc.SourceBase.plot_source_energy = plot_source_energy
+openmc.model.Model.plot_source_energy = plot_source_energy
+openmc.Model.plot_source_energy = plot_source_energy
+
 openmc.SourceBase.plot_source_position = plot_source_position
+openmc.model.Model.plot_source_position = plot_source_position
+openmc.Model.plot_source_position = plot_source_position
+
 openmc.SourceBase.plot_source_direction = plot_source_direction
+openmc.model.Model.plot_source_direction = plot_source_direction
+openmc.Model.plot_source_direction = plot_source_direction
diff --git a/tests/test_core_with_model.py b/tests/test_core_with_model.py
new file mode 100644
index 0000000..5a5a5e0
--- /dev/null
+++ b/tests/test_core_with_model.py
@@ -0,0 +1,89 @@
+import openmc
+import openmc_source_plotter
+import numpy as np
+import plotly.graph_objects as go
+import pytest
+
+
+@pytest.fixture
+def test_model():
+    # initialises a new source object
+    my_source = openmc.Source()
+
+    # sets the location of the source to x=0 y=0 z=0
+    my_source.space = openmc.stats.Point((1.0, 2.0, 3.0))
+
+    # sets the direction to isotropic
+    my_source.angle = openmc.stats.Isotropic()
+
+    # sets the energy distribution to 100% 14MeV neutrons
+    my_source.energy = openmc.stats.Discrete([15e6], [1])
+
+    my_source.particle = "photon"
+
+    settings = openmc.Settings()
+    settings.particles = 1
+    settings.batches = 1
+    settings.source = my_source
+
+    materials = openmc.Materials()
+
+    sph = openmc.Sphere(r=9999999999, boundary_type="vacuum")
+    cell = openmc.Cell(region=-sph)
+    geometry = openmc.Geometry([cell])
+
+    model = openmc.Model(geometry, materials, settings)
+
+    return model
+
+
+def test_sample_initial_particles(test_model):
+    particles = test_model.sample_initial_particles(n_samples=43)
+    for particle in particles:
+        assert particle.E == 15e6
+        assert str(particle.particle) == "photon"
+        assert particle.r == (1.0, 2.0, 3.0)
+    assert len(particles) == 43
+
+
+def test_energy_plot_with_bins(test_model):
+    plot = test_model.plot_source_energy(
+        n_samples=10,
+        energy_bins=np.linspace(0, 20e6, 100),
+    )
+    assert isinstance(plot, go.Figure)
+
+
+def test_energy_plot(test_model):
+    plot = test_model.plot_source_energy(n_samples=10)
+    assert isinstance(plot, go.Figure)
+    assert len(plot.data[0]["x"]) == 1
+
+
+def test_position_plot(test_model):
+    plot = test_model.plot_source_position(n_samples=10)
+    assert isinstance(plot, go.Figure)
+
+
+def test_direction_plot(test_model):
+    plot = test_model.plot_source_direction(n_samples=10)
+    assert isinstance(plot, go.Figure)
+
+
+def test_energy_plot_with_figure(test_model):
+    base_figure = go.Figure()
+    plot = test_model.plot_source_energy(figure=base_figure, n_samples=10)
+    assert isinstance(plot, go.Figure)
+    assert len(plot.data[0]["x"]) == 1
+
+
+def test_position_plot_with_figure(test_model):
+    base_figure = go.Figure()
+    plot = test_model.plot_source_position(figure=base_figure, n_samples=10)
+    assert isinstance(plot, go.Figure)
+
+
+def test_direction_plot_with_figure(test_model):
+    base_figure = go.Figure()
+    plot = test_model.plot_source_direction(figure=base_figure, n_samples=10)
+    assert isinstance(plot, go.Figure)
diff --git a/tests/test_core.py b/tests/test_core_with_source.py
similarity index 64%
rename from tests/test_core.py
rename to tests/test_core_with_source.py
index e598e1f..56f75e6 100644
--- a/tests/test_core.py
+++ b/tests/test_core_with_source.py
@@ -25,39 +25,7 @@ def test_source():
     return my_source
 
 
-@pytest.fixture
-def test_model():
-    # initialises a new source object
-    my_source = openmc.Source()
-
-    # sets the location of the source to x=0 y=0 z=0
-    my_source.space = openmc.stats.Point((1.0, 2.0, 3.0))
-
-    # sets the direction to isotropic
-    my_source.angle = openmc.stats.Isotropic()
-
-    # sets the energy distribution to 100% 14MeV neutrons
-    my_source.energy = openmc.stats.Discrete([15e6], [1])
-
-    my_source.particle = "photon"
-
-    settings = openmc.Settings()
-    settings.particles = 1
-    settings.batches = 1
-    settings.source = my_source
-
-    materials = openmc.Materials()
-
-    sph = openmc.Sphere(r=9999999999, boundary_type="vacuum")
-    cell = openmc.Cell(region=-sph)
-    geometry = openmc.Geometry([cell])
-
-    model = openmc.Model(geometry, materials, settings)
-
-    return model
-
-
-def test_source_sample_initial_particles(test_source):
+def test_sample_initial_particles(test_source):
     particles = test_source.sample_initial_particles(n_samples=42)
     for particle in particles:
         assert particle.E == 14e6
@@ -66,15 +34,6 @@ def test_source_sample_initial_particles(test_source):
     assert len(particles) == 42
 
 
-def test_model_sample_initial_particles(test_model):
-    particles = test_model.sample_initial_particles(n_samples=43)
-    for particle in particles:
-        assert particle.E == 15e6
-        assert str(particle.particle) == "photon"
-        assert particle.r == (1.0, 2.0, 3.0)
-    assert len(particles) == 43
-
-
 def test_energy_plot_with_bins(test_source):
     plot = test_source.plot_source_energy(
         n_samples=10,