Repository: kyleaoman/martini (language: python)

Issue:
Document source inspection, add preview tool
Some example code showing how to inspect and manipulate source particles before making the data cube would be a good addition to the docs. This could be a good starting point (the `plot_yz` function could be a good addition to the source module code as a "preview" function):

```
import numpy as np
import unyt as u
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from astropy import units as U
from swiftgalaxy import SWIFTGalaxy, Velociraptor
from velociraptor import load as load_catalogue
from os import path
from martini.sources import SWIFTGalaxySource
from martini.sources._L_align import L_align

np.seterr(divide="ignore", invalid="ignore")

snapnum = 23
base_dir = "106e3_104b2_norm_0p3_new_cooling_L006N188"
velociraptor_filebase = path.join(base_dir, f"halo_{snapnum:04d}")
snapshot_filename = path.join(base_dir, f"colibre_{snapnum:04d}.hdf5")

catalogue = load_catalogue(f"{velociraptor_filebase}.properties")
# select centrals with vmax > 80km/s
target_mask = np.logical_and.reduce(
    (
        catalogue.ids.hosthaloid == -1,
        catalogue.velocities.vmax > 80 * u.km / u.s,
    )
)
target_halo_indices = np.argwhere(target_mask).flatten()
# pick the first one from the list, why not...
target_halo_index = target_halo_indices[0]

sg = SWIFTGalaxy(
    snapshot_filename=snapshot_filename,
    halo_finder=Velociraptor(
        velociraptor_filebase=velociraptor_filebase,
        halo_index=target_halo_index,
    ),
)

source = SWIFTGalaxySource(
    sg,
    distance=3 * U.Mpc,
    vpeculiar=0 * U.km / U.s,
    rotation=dict(rotmat=np.eye(3)),  # identity matrix, we'll rotate "by hand"
    ra=0 * U.deg,
    dec=0 * U.deg,
)
# source.coordinates_g is an astropy CartesianRepresentation object.
# It has an attached CartesianDifferential object for the velocities.
# See https://docs.astropy.org/en/stable/coordinates/representations.html


def plot_yz(source, pdffile, prune=np.s_[::25], title=""):
    # prune: by default plot every 25th particle.
    fig = plt.figure(1)
    plt.clf()
    sp = fig.add_subplot(1, 1, 1, aspect="equal")
    scatter = sp.scatter(
        source.coordinates_g.y[prune].to_value(U.kpc),
        source.coordinates_g.z[prune].to_value(U.kpc),
        c=source.coordinates_g.differentials["s"].d_x[prune].to_value(U.km / U.s),
        marker="o",
        s=1,
        vmin=-150,
        vmax=150,
        cmap="RdBu",
    )
    sp.set_xlabel(r"$y\,[\mathrm{kpc}]$")
    sp.set_ylabel(r"$z\,[\mathrm{kpc}]$")
    sp.set_xlim((-30, 30))
    sp.set_ylim((-30, 30))
    sp.set_title(title)
    cb = fig.colorbar(mappable=scatter)
    cb.set_label(r"$v_x\,[\mathrm{km}\,\mathrm{s}^{-1}]$")
    plt.savefig(pdffile, format="pdf")


with PdfPages("forKatharina.pdf") as pdffile:
    plot_yz(source, pdffile, title="unrotated")
    # This is the same as you would get using
    # source.rotate(L_align=(0 * U.deg, 0 * U.deg, 270 * U.deg))
    # but you can adapt as needed:
    rotmat = L_align(
        source.coordinates_g.xyz,
        source.coordinates_g.differentials["s"].d_xyz,
        source.mHI_g,
        frac=0.3,
        Laxis="x",
    )
    source.rotate(rotmat=rotmat)
    plot_yz(source, pdffile, title="face-on")
    # rotation about the galactic pole:
    source.rotate(axis_angle=("x", 45 * U.deg))
    plot_yz(source, pdffile, title="face-on, polar rotation 45deg")
    # inclination:
    source.rotate(axis_angle=("y", 60 * U.deg))
    plot_yz(source, pdffile, title="inclined 60deg")

    # SWIFTGalaxy has already re-centred (x,y,z)=(0,0,0) on the velociraptor
    # centre of potential, and the (vx,vy,vz)=(0,0,0) on the velociraptor
    # subhalo bulk velocity. The centre of potential is already what we want,
    # but the bulk velocity probably does not result in the central gas
    # particles having 0 net velocity, so we can adjust like this:
    r = np.sqrt(np.sum(np.power(source.coordinates_g.xyz.to_value(U.kpc), 2), axis=0))
    rsort = np.argsort(r)
    # we'll reject particles with small HI fractions with this:
    HI_mask = (source.mHI_g / source.mHI_g.max()).to_value(U.one) > 0.3
    # LoS velocities:
    vxyz = source.coordinates_g.differentials["s"].d_xyz
    # mean velocity of 100 innermost HI-rich gas particles:
    vxyz0 = np.mean(vxyz[:, rsort][:, HI_mask[rsort]][:, :100], axis=1)
    # In principle we only need to adjust the LoS velocity, but if we adjust
    # the 3D velocity the correction will be ok for any orientation.
    source.boost(-vxyz0)
    plot_yz(source, pdffile, title="inclined 60deg, with velocity recentred")

```


Discussion / hints:

Fix patch:
diff --git a/docs/source/index.rst b/docs/source/index.rst
index c5bf364..bcbb1c1 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -14,7 +14,9 @@ Overview
    :caption: Documentation contents
 
    getting_started/index
+   sources/index
    sph_kernels/index
+   spectral_models/index
    martini/index
    modules/index
 
diff --git a/docs/source/sources/index.rst b/docs/source/sources/index.rst
new file mode 100644
index 0000000..1e34fd3
--- /dev/null
+++ b/docs/source/sources/index.rst
@@ -0,0 +1,231 @@
+Sources
+=======
+
+In the context of mock observations of a simulation, the source is the collection of particles that produce the emission to be "observed".
+
+Sources in MARTINI
+------------------
+
+MARTINI has a collection of classes used to store and manipulate the properties of the source particles. All of these inherit from the :class:`martini.sources.sph_source.SPHSource` class. This base class can be created by providing arrays containing the particle coordinates, velocities, masses, etc. The other source classes are tailored to specific simulations or data formats and will typically expect one or more filenames as input, along with an identifier specifying a group of particles (as defined in a group catalogue, for example). These classes will then take care of reading data from the input files, annotating them with units, doing any required calculations or conversions, and so on. There are source classes for the EAGLE_, IllustrisTNG_, Simba_, Colibre and Magneticum_ simulations. The SWIFTGalaxy_ interface to `SWIFT`_-based simulations with a variety of halo catalogue formats is also supported.
+
+.. _EAGLE: https://icc.dur.ac.uk/Eagle
+.. _IllustrisTNG: https://www.tng-project.org
+.. _Simba: http://simba.roe.ac.uk
+.. _Magneticum: http://www.magneticum.org/
+.. _SWIFTGalaxy: https://github.com/SWIFTSIM/swiftgalaxy
+.. _SWIFT: https://github.com/SWIFTSIM/SWIFT
+
+Units in MARTINI
+----------------
+
+MARTINI adopts the astropy_ system of :mod:`astropy.units`. A very brief introduction is given here for users unfamiliar with this module.
+
+.. _astropy: https://www.astropy.org
+
+The most common unit operation in MARTINI is to attach units to scalars or arrays. This is intuitively achieved by multiplying (or dividing) the scalar/array by the unit. Units can also be raised to powers.
+
+.. code-block :: python
+
+    import numpy as np
+    import astropy.units as U
+    
+    mass = 1 * U.Msun
+    speeds = np.array([1000, 1001, 1002]) * U.km / U.s
+    density = 1 * U.Msun / U.kpc ** 3
+
+These :class:`astropy.units.Quantity` objects can be used in most of the same ways as ordinary scalars or arrays, but will check unit consistency and/or propagate units during calculations. For example, attempting to add two quantities with incompatible dimensions will raise an exception, and dividing a quantity with dimensions of length by a quantity with dimensions of time will return a quantity with dimensions of speed. Conversion to different units with the same dimensions can be achieved with the :meth:`~astropy.units.Quantity.to` method, while the :meth:`~astropy.units.Quantity.to_value` method will return the (array or scalar) value in the specified units, without units attached.
+
+.. code-block :: python
+
+    from astropy.units import UnitConversionError
+
+    speeds.to(U.Mpc / U.Gyr)  # returns a Quantity object
+    (mass * speeds).to_value(U.kg * U.m / U.s)  # returns a numpy array of momenta in SI units
+    try:
+        mass + density  # incompatible units raise exception!
+    except UnitConversionError:
+        pass
+
+MARTINI functions that expect :class:`~astropy.units.Quantity` inputs accept any units with the correct dimensions, so a mass could be specified in kg, Msun, or other mass units. Any required unit conversions will then happen internally.
+
+The module :mod:`astropy.constants` provides a wide range of pre-defined physical constants compatible with :mod:`astropy.units`.
+
+Using MARTINI's source classes
+------------------------------
+
+Quick overview here (including preview), details below
+
+Particle arrays
++++++++++++++++
+
+The core information required by the :class:`~martini.sources.sph_source.SPHSource` module is the set of arrays containing the properties of the particles making up the source. These are particle coordinates ``xyz_g``, velocities ``vxyz_g``, HI masses ``mHI_g``, smoothing lengths ``hsm_g`` and temperatures ``T_g``. All of these require units attached (see above).
+
+The position of the source is defined by the `(0, 0, 0)` coordinate location. A particle at `y=0` will be placed at the nominal right ascension of the source (see below), and a particle at `z=0` at its nominal declination. Particles offset from the centre will correspondingly be offset in angle from the centre RA and Dec according to their offsets in `y` and `z` at the distance of the particles. The line of sight toward the centre of the source is along the `x`-axis, so loosely speaking the x-component of the particle velocities determines the channels in which they contribute flux. A particle at the centre of the source with zero velocity in the `x` direction will contribute flux in the channel corresponding to the systemic velocity of the source. More accurately, MARTINI implements a full perspective projection (in contrast to a parallel projection) and accounts for the 3D structure of the source. This means that the "observer" could even be placed *inside* a galaxy to simulate Galactic observations, yielding accurate results within the limitations of the resolution of the simulation. Users can rotate/translate particle coordinates to obtain the desired viewing angle to their source before passing the arrays to the source module, however the source module also offers features to manipulate the source orientation and visualise previews of the mock observation. These are explained below.
+
+The HI masses and temperatures of the particles are straightforward to understand. Both are often tabulated directly in simulation snapshots. If temperature is not present, usually a related quantity such as internal energy is, from which the temperatures can be calculated. If HI masses are not available, calculating neutral hydrogen fractions and partitioning the neutral hydrogen into its atomic and molecular phases may be necessary. How best to do this depends on the details of individual simulations. More information is best sought from relevant publications, the documentation of the simulations in question, or the teams that developed the simulations. For simulations with a corresponding source module in MARTINI, any calculations needed to obtain temperatures and HI masses are implemented in the code.
+
+The smoothing lengths are defined as the full width at half-maximum (FWHM) of the SPH kernel function. Because there is no standard convention for defining smoothing lengths, this is usually not equal to the smoothing lengths tabulated in snapshots. Users need to convert the smoothing lengths from their simulation to the equivalent FWHM before passing them to MARTINI. Further details are in the :doc:`SPH kernels </sph_kernels/index>` section of the documentation. Again, for simulations supported with MARTINI source modules, any conversion of smoothing lengths is handled internally.
+
+Simulation-specific source modules
+++++++++++++++++++++++++++++++++++
+
+MARTINI provides source modules to simplify working with publicly-available simulation data sets. These currently include the :class:`~martini.sources.eagle_source.EAGLESource`, :class:`~martini.sources.tng_source.TNGSource` and :class:`~martini.sources.simba_source.SimbaSource`. Example usage including how to obtain publicly available data is provided in a set of `Jupyter notebooks`_.
+
+The Magneticum_ simulations are supported with the :class:`~martini.sources.magneticum_source.MagneticumSource`, however detailed examples are not provided due to the lack of publicly available snapshot data. Interested users can refer to the API documentation of the class.
+
+There is also a :class:`~martini.sources.swiftgalaxy_source.SWIFTGalaxySource` source module for SWIFT_ simulation data in conjunction with Velociraptor_ or Caesar_ halo catalogues; support for `HBT+`_ and SOAP_ is also planned. A :class:`~martini.sources.colibre_source.ColibreSource` module exists, but since this galaxy formation model is still under development this is currently experimental functionality.
+
+Finally, there is an :class:`~martini.sources.so_source.SOSource` module that interfaces with the SimObj_ package. That package is no longer maintained, but may facilitate working with some simulations including APOSTLE, C-EAGLE/Hydrangea and Auriga.
+
+.. _Jupyter notebooks: https://github.com/kyleaoman/martini/tree/main/examples
+.. _Velociraptor: https://github.com/pelahi/VELOCIraptor-STF
+.. _Caesar: https://github.com/dnarayanan/caesar
+.. _HBT+: https://github.com/SWIFTSIM/HBTplus
+.. _SOAP: https://github.com/SWIFTSIM/SOAP
+.. _SimObj: https://github.com/kyleaoman/simobj
+
+Distance, peculiar velocity, right ascension and declination
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+
+Any source module can be configured with the ``distance``, ``vpeculiar``, ``ra`` and ``dec`` keyword arguments. The distance, RA and Dec define the offset of the particles relative to the observer. The source is given a recession velocity that places it in the Hubble flow according to its distance and the Hubble constant :math:`H_0=h(100\,\mathrm{km}\,\mathrm{s}^{-1}\,\mathrm{Mpc}^{-1})`. By default `h=0.7`, but this can be adjusted with the ``h`` keyword argument.
+
+The systemic velocity of the source is defined as the sum of the Hubble and peculiar velocities: :math:`v_\mathrm{sys}=v_\mathrm{Hubble}+v_\mathrm{peculiar}`. A positive peculiar velocity therefore makes the source recede faster than :math:`v_\mathrm{Hubble}=H_0D`, and vice-versa.
+
+Manipulating a source before making a mock
+------------------------------------------
+
+While particle arrays can be manipulated before passing them to one of MARTINI's source modules, sometimes it may be more convenient to manipulate them using tools provided by the source modules.
+
+Rotation and translation
+++++++++++++++++++++++++
+
+MARTINI allows source particles to be rotated on initialization, or later by calling the :meth:`~martini.sources.sph_source.SPHSource.rotate` method. Particles can be translated after initialization with the :meth:`~martini.sources.sph_source.SPHSource.translate` method (for "translations" in velocity use :meth:`~martini.sources.sph_source.SPHSource.boost`).
+
+MARTINI offers three ways of specifying rotations:
+
+ - ``axis_angle`` defines a rotation with a 2-tuple where the first element is one of the strings ``"x"``, ``"y"`` or ``"z"`` and the second element is an angle (with units specified with :mod:`astropy.units`). The position and velocity vectors of each particle are rotated about the specified axis (in the "source frame", not the "observer frame").
+ - ``rotmat`` defines a rotation with a 3x3 rotation matrix.
+ - ``L_coords`` is offered as a convenience when an (approximate) inclination angle and position angle are desired. This method first identifies the plane perpendicular to the angular momentum of the source (specifically, the angular momentum of the 1/3 of particles closest to the source centre). The source is then rotated to place this plane in the `y-z` plane. There is a degree of freedom (a rotation about the angular momentum vector) here; this is fixed with an arbitrary (but *not* random) choice. The ``L_coords`` specification is then a 2-tuple or a 3-tuple. If a 2-tuple, the first element is the desired inclination angle and the second element controls a rotation arond the pole (applied before the inclination). If a third element is provided, this sets the position angle on the sky (otherwise this defaults to 270 degrees). All of these angles should be specified using :mod:`astropy.units`.
+
+To specify a rotation when initializing the source, the rotation is given as a dictionary passed to the ``rotation`` keyword argument, for example:
+
+.. code-block:: python
+
+    SPHSource(rotation={"axis_angle": ("x", 30 * U.deg)}, ...)  # or "rotmat" or "L_coords"
+
+To specify a rotation for a source already initialized, use the ``axis_angle``, ``rotmat`` or ``L_coords`` keyword arguments as:
+
+.. code-block:: python
+
+    source = SPHSource(...)
+    source.rotate(axis_angle=("x", 30 * U.deg))  # or rotmat or L_coords
+
+The current rotation state of a source (relative to the state in which the particles were passed in to MARTINI) can be written out to a file (always as a rotation matrix) using :meth:`~martini.sources.sph_source.SPHSource.save_current_rotation`.
+
+Masking
++++++++
+
+A source can be masked to remove particles, if desired. The method :meth:`~martini.sources.sph_source.SPHSource.apply_mask` enables this. It accepts a boolean array (or other objects that can be used to index numpy arrays).
+
+Inspecting a source before making a mock
+----------------------------------------
+
+Creating a high-resolution mock observation can be computationally expensive, so it is useful to be able to view a quick representation of the source to check that it appears as desired before investing the effort of making a full mock observation. The source modules provide the :meth:`~martini.sources.sph_source.SPHSource.preview` to enable this. To illustrate its usage, let's set up a randomly-oriented source using MARTINI's simple "demo" toy model of a disc, and call :meth:`~martini.sources.sph_source.SPHSource.preview`:
+
+.. code-block:: python
+
+    import numpy as np
+    import astropy.units as U
+    from martini import demo_source
+
+    source = demo_source(N=20000)  # create simple disc with 20000 particles
+    # a random rotation matrix:
+    rotmat = np.array(
+        [
+            [-0.20808178, -0.97804544, -0.01136216],
+            [0.02991471, -0.01797457, 0.99939083],
+            [0.97765387, -0.20761513, -0.03299812],
+        ]
+    )
+    # apply it so that the source has no particular orientation:
+    source.rotate(rotmat=rotmat)
+
+    source.preview(fig=1)  # uses matplotlib `plt.figure(1)`
+
+The preview function returns the ``Figure`` object, so this can be captured manipulated, if desired. The resulting preview looks like:
+
+.. image:: preview1.png
+    :width: 800
+    :alt: Approximate moment 1 map and major & minor axis PV diagrams of randomly-oriented galaxy.
+
+The number of particles plotted is limited to 5000 to avoid excessively large ``Figure`` objects, but this can be controlled with the ``max_points`` keyword argument. The first panel shows the particle positions in the `y-z` plane (which approximately maps to RA and Dec), coloured by the `x`-component of the velocity (which approximately maps to the line-of-sight velocity). The second panel shows the PV diagram of particle along the `y` coordinate direction, and the third panel the same along the `z` coordinate direction. Let's rotate the disc to be edge on and preview it again:
+
+.. code-block:: python
+
+    source.rotate(L_coords=(90 * U.deg, 0 * U.deg))  # (inclination, polar rotation); implicit: PA=270deg
+    source.preview(fig=2)
+
+.. image:: preview2.png
+    :width: 800
+    :alt: Approximate moment 1 map and major & minor axis PV diagrams of edge-on galaxy.
+
+This does look like an edge on disc, and the PV diagrams now correspond to the major (middle panel) and minor (right panel) axis PV diagrams. Let's switch orientations again to 60 degrees inclined and a position angle offset 45 degrees from its previous value of 270 degrees, and this time we'll show only 100 points.
+
+source.rotate(L_coords=(90 * U.deg, 0 * U.deg, 225 * U.deg))
+source.preview(max_points=100, fig=3)
+
+.. image:: preview3.png
+    :width: 800
+    :alt: Approximate moment 1 map and major & minor axis PV diagrams of inclined galaxy, not very clear.
+
+This doesn't look great. The preview function tries to adjust the transparency and sizes of the points to qualitatively reflect their HI masses and smoothing lengths, but this may sometimes lead to results that aren't very useful. The points can be forced to a constant size and transparency with ``point_scaling="fixed"``.
+
+.. code-block:: python
+
+    source.preview(max_points=100, fig=4, point_scaling="fixed")
+
+.. image:: preview4.png
+    :width: 800
+    :alt: Approximate moment 1 map and major & minor axis PV diagrams of inclined galaxy.
+
+Let's look at a more realistic example from the TNG simulations.
+
+.. code-block:: python
+
+    from martini.sources import TNGSource
+
+    tng_source = TNGSource(
+        "TNG50-1",
+        99,
+        572840,
+        api_key="your-tng-api-key-goes-here",
+    )
+    tng_source.preview(fig=5)
+
+.. image:: preview5.png
+    :width: 800
+    :alt: Approximate moment 1 map and major & minor axis PV diagrams of TNG galaxy, zoomed out.
+
+The preview function scales the axes to enclose all particles in the source. In this case we're mostly seeing the hot circumgalactic gas extending out beyond 200 kpc. The ``lim`` keyword argument sets the (absolute value) offset from the centre shown in the preview, and the ``vlim`` keyword argument similarly limits the velocity offset from zero. Note that these keyword arguments have no influence on the particles contained in the source, just on what's visible in the preview. Let's take a guess at the likely size of a disc and see what's there.
+
+.. code-block:: python
+
+    tng_source.preview(fig=6, lim=20 * U.kpc, vlim=300 * U.km / U.s)
+
+.. image:: preview6.png
+    :width: 800
+    :alt: Approximate moment 1 map and major & minor axis PV diagrams of TNG galaxy, zoomed in.
+
+Aha, an inclined disc within the hot halo. From here we could continue to adjust the orientation, if desired.
+
+Because the :mod:`~martini.sources.tng_source.TNGSource` module has automatically retrieved and loaded the particle arrays, we did not need to pass them in, and therefore never had a chance to do anything with them, if we desired to. The particle arrays are accessible as attributes, for example:
+
+.. code-block:: python
+
+    tng_source.xyz_g
+    tng_source.vxyz_g
+    tng_source.mHI_g
+    tng_source.T_g
+    tng_source.hsm_g
+
+The same attributes exist for all source modules available in MARTINI.
diff --git a/docs/source/sources/preview1.png b/docs/source/sources/preview1.png
new file mode 100644
index 0000000..205e4d5
Binary files /dev/null and b/docs/source/sources/preview1.png differ
diff --git a/docs/source/sources/preview2.png b/docs/source/sources/preview2.png
new file mode 100644
index 0000000..ce38ff7
Binary files /dev/null and b/docs/source/sources/preview2.png differ
diff --git a/docs/source/sources/preview3.png b/docs/source/sources/preview3.png
new file mode 100644
index 0000000..00010dd
Binary files /dev/null and b/docs/source/sources/preview3.png differ
diff --git a/docs/source/sources/preview4.png b/docs/source/sources/preview4.png
new file mode 100644
index 0000000..80f472b
Binary files /dev/null and b/docs/source/sources/preview4.png differ
diff --git a/docs/source/sources/preview5.png b/docs/source/sources/preview5.png
new file mode 100644
index 0000000..bdb57f9
Binary files /dev/null and b/docs/source/sources/preview5.png differ
diff --git a/docs/source/sources/preview6.png b/docs/source/sources/preview6.png
new file mode 100644
index 0000000..9a77d67
Binary files /dev/null and b/docs/source/sources/preview6.png differ
diff --git a/docs/source/spectral_models/index.rst b/docs/source/spectral_models/index.rst
index da100ac..ae11b1e 100644
--- a/docs/source/spectral_models/index.rst
+++ b/docs/source/spectral_models/index.rst
@@ -24,8 +24,8 @@ The default is a constant 7 km/s line width, corresponding to a FWHM of about 16
 
 The :class:`~martini.spectral_models.DiracDeltaSpectrum` instead assigns all of the emission of a particle to the channel containing the particle's line-of-sight velocity. This is mostly intended as a debugging tool but can sometimes be useful to speed up a calculation in cases where a delta function is a good approximation (for instance if the channels are very wide).
 
-Using MARTINI's SPH kernel classes
-----------------------------------
+Using MARTINI's spectral model classes
+--------------------------------------
 
 Simply choose the class corresponding to your preferred line shape and initialise it, then pass it to the main :class:`~martini.martini.Martini` class, for example:
 
diff --git a/examples/martini_TNG.ipynb b/examples/martini_TNG.ipynb
index 18b9e8e..e1605ea 100644
--- a/examples/martini_TNG.ipynb
+++ b/examples/martini_TNG.ipynb
@@ -27,7 +27,7 @@
    "source": [
     "MARTINI is a modular package for the creation of synthetic resolved HI line observations (data cubes) of smoothed-particle hydrodynamics simulations of galaxies. The various aspects of the mock-observing process are divided logically into sub-modules handling the data cube, source, beam, noise, spectral model and SPH kernel. MARTINI is object-oriented: each sub-module provides a class (or classes) which can be configured as desired. For most sub-modules, base classes are provided to allow for straightforward customization. Instances of each sub-module class are given as parameters to the Martini class; a mock observation is then constructed by calling a handful of functions to execute the desired steps in the mock-observing process.\n",
     "\n",
-    "The MARTINI documentation is [hosted on ReadTheDocs](https://martini.readthedocs.io/en/latest/)."
+    "This tutorial focuses on particulars related to working with the IllustrisTNG simulations. More general information is available in the MARTINI documentation, [hosted on ReadTheDocs](https://martini.readthedocs.io/en/latest/)."
    ]
   },
   {
@@ -76,6 +76,22 @@
     "!{sys.executable} -m pip install matplotlib"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This cell may be needed in some cases to display figures below:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -215,26 +231,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The rotation configuration takes an inclination (here 60deg) and\n",
-    "rotation about the pole (here 90deg, relative to an arbitrary reference direction). The code attempts to\n",
-    "automatically align the galactic disk in the y-z plane by aligning\n",
-    "the angular momentum along the x-axis. The polar rotation is then\n",
-    "applied, and finally the disc inclined by a rotation around the\n",
-    "y-axis (the line of sight is along the x-axis). The automatic\n",
-    "alignment will work for typical reasonably isolated discs, but will\n",
-    "struggle when companions are present, when the angular momentum axis\n",
-    "is a poor tracer of the disc plane, and especially for satellites. If\n",
-    "finer control of the orientation is needed, derive the transformation\n",
-    "from the simulation box coordinates to the desired coordinates for\n",
-    "the 'observation', keeping in mind that the line of sight is along\n",
-    "the x-axis. This rotation matrix can then be passed to as\n",
-    "`rotation={'rotmat': np.eye(3)}` (here the identity rotation matrix used as a place holder). A common problem in this case is deriving the inverse\n",
-    "transform instead of the forward transform, if unexpected results are\n",
-    "obtained, first try passing the transpose of the rotation matrix.\n",
-    "\n",
     "Notice that your API key must be provided unless you are using `TNGSource` on the TNG JupyterLab (if you provide it anyways in that case, it will simply be ignored).\n",
     "\n",
-    "Any downloaded data can be cached using the `cutout_dir` parameter to specify a directory, if a cache for the requested galaxy is found it will be used instead of re-downloading the data. If running on the TNG JupyterLab environment, this parameter is ignored.\n"
+    "Any downloaded data can be cached using the `cutout_dir` parameter to specify a directory, if a cache for the requested galaxy is found it will be used instead of re-downloading the data. If running on the TNG JupyterLab environment, this parameter is ignored."
    ]
   },
   {
@@ -250,12 +249,100 @@
     "    api_key=api_key,\n",
     "    cutout_dir=\".\",\n",
     "    distance=4 * U.Mpc,\n",
-    "    rotation={\"L_coords\": (60 * U.deg, 0.0 * U.deg)},\n",
+    "    rotation={\"rotmat\": np.eye(3)},\n",
     "    ra=0.0 * U.deg,\n",
     "    dec=0.0 * U.deg,\n",
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The rotation argument above has been set to the identity matrix, so the source has the (random) orientation that it has within the simulation volume. The source class includes a function to make a quick plot to get an idea of the source's appearance:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preview_fig_unrotated = source.preview(title=\"unrotated\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The preview function defaults to apertures in position and velocity that enclose all particles in the source, so this preview emphasizes the diffuse circumgalactic gas. The apertures can be set manually using the `lim` and `vlim` keywords to set the maximum absolute offsets in position and velocity relative to the source centre to be plotted. For example, restricting the aperture to 50kpc and 300km/s makes the disc more clearly visible."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preview_fig_unrotated_zoom = source.preview(\n",
+    "    title=\"unrotated, zoomed-in\",\n",
+    "    lim=50 * U.kpc,\n",
+    "    vlim=300 * U.km / U.s,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This randomly-oriented viewing angle seems to be close to face-on. The source can be rotated to a different orientation. MARTINI's tool for quick/approximate manipulation of the orientation of the source aligns the source based on its angular momentum vector (\"L\"), for example:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "source.rotate(L_coords=(60 * U.deg, 90 * U.deg))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The rotation configuration takes an inclination (here 60deg) and rotation about the pole (here 90deg, relative to an arbitrary reference direction). The code attempts to\n",
+    "automatically align the galactic disk in the y-z plane by aligning\n",
+    "the angular momentum along the x-axis. The polar rotation is then\n",
+    "applied, and finally the disc inclined by a rotation around the\n",
+    "y-axis (the line of sight is along the x-axis). The automatic\n",
+    "alignment will work for typical reasonably isolated discs, but will\n",
+    "struggle when companions are present, when the angular momentum axis\n",
+    "is a poor tracer of the disc plane, and especially for satellites."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preview_fig_rotated_zoomed_in = source.preview(\n",
+    "    title=\"rotated, zoomed-in\",\n",
+    "    lim=50 * U.kpc,\n",
+    "    vlim=300 * U.km / U.s,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If finer control of the orientation is needed, derive the transformation from the simulation box coordinates (see [the documentation](https://martini.readthedocs.io/en/latest/) for examples) to the desired coordinates for the 'observation', keeping in mind that the line of sight is along the x-axis. This rotation matrix can then be passed to the rotate function as `rotmat=np.eye(3)` (here the identity rotation matrix used as a place holder). The rotation can also be provided when the source is initialized by using the `rotation` keyword argument.\n",
+    "\n",
+    "A common problem is deriving the inverse transform instead of the forward transform, if unexpected results are obtained, first try passing the transpose of the rotation matrix."
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -384,7 +471,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The calculation of the spectra can be done in parallel by providing a keyword argument `ncpu=N`, where `N` is the number of CPUs to use. However, the details of the implementation mean that for small numbers of particles running in parallel tends to slow down the calculation, so turning this on should be done with care. Significant speedups can be expected when the particl"
+    "The calculation of the spectra (that will happen when the `Martini` module is initialized below) can be done in parallel by providing a keyword argument `ncpu=N`, where `N` is the number of CPUs to use. However, the details of the implementation mean that for small numbers of particles running in parallel tends to slow down the calculation, so turning this on should be done with care. Significant speedups can be expected when the particle count is very large."
    ]
   },
   {
@@ -588,7 +675,7 @@
     ")\n",
     "im1 = sp1.imshow(mom0.T, cmap=\"Greys\", aspect=1.0, origin=\"lower\")\n",
     "plt.colorbar(im1, ax=sp1, label=\"mom0 [Jy/beam]\")\n",
-    "im2 = sp2.imshow((mom1 * mask).T, cmap=\"RdBu\", aspect=1.0, origin=\"lower\")\n",
+    "im2 = sp2.imshow((mom1 * mask).T, cmap=\"RdBu_r\", aspect=1.0, origin=\"lower\")\n",
     "plt.colorbar(im2, ax=sp2, label=\"mom1 [km/s]\")\n",
     "im3 = sp3.imshow(\n",
     "    (mom2 * mask).T, cmap=\"magma\", aspect=1.0, origin=\"lower\", vmin=0, vmax=300\n",
diff --git a/examples/martini_eagle.ipynb b/examples/martini_eagle.ipynb
index e754f64..613e083 100644
--- a/examples/martini_eagle.ipynb
+++ b/examples/martini_eagle.ipynb
@@ -20,7 +20,7 @@
    "source": [
     "MARTINI is a modular package for the creation of synthetic resolved HI line observations (data cubes) of smoothed-particle hydrodynamics simulations of galaxies. The various aspects of the mock-observing process are divided logically into sub-modules handling the data cube, source, beam, noise, spectral model and SPH kernel. MARTINI is object-oriented: each sub-module provides a class (or classes) which can be configured as desired. For most sub-modules, base classes are provided to allow for straightforward customization. Instances of each sub-module class are given as parameters to the Martini class; a mock observation is then constructed by calling a handful of functions to execute the desired steps in the mock-observing process.\n",
     "\n",
-    "The MARTINI documentation is [hosted on ReadTheDocs](https://martini.readthedocs.io/en/latest/)."
+    "This tutorial focuses on particulars related to working with the EAGLE simulations. More general information is available in the MARTINI documentation, [hosted on ReadTheDocs](https://martini.readthedocs.io/en/latest/)."
    ]
   },
   {
@@ -69,6 +69,22 @@
     "!{sys.executable} -m pip install requests matplotlib"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This cell may be needed in some cases to display figures below:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -314,30 +330,15 @@
     "disc plane. Normally it is advisable to set this to approximately\n",
     "the virial radius of the source object, or just large enough to\n",
     "capture the region of interest around it (e.g. enough to encompass\n",
-    "the host of a satellite galaxy).\n",
-    "\n",
-    "The rotation configuration takes an inclination (here 60deg) and\n",
-    "rotation about the pole (here 90deg, relative to an arbitrary reference direction). The code attempts to\n",
-    "automatically align the galactic disk in the y-z plane by aligning\n",
-    "the angular momentum along the x-axis. The polar rotation is then\n",
-    "applied, and finally the disc inclined by a rotation around the\n",
-    "y-axis (the line of sight is along the x-axis). The automatic\n",
-    "alignment will work for typical reasonably isolated discs, but will\n",
-    "struggle when companions are present, when the angular momentum axis\n",
-    "is a poor tracer of the disc plane, and especially for satellites. If\n",
-    "finer control of the orientation is needed, derive the transformation\n",
-    "from the simulation box coordinates to the desired coordinates for\n",
-    "the 'observation', keeping in mind that the line of sight is along\n",
-    "the x-axis. This rotation matrix can then be passed to as\n",
-    "`rotation={'rotmat': np.eye(3)}` (here the identity rotation matrix used as a place holder). A common problem in this case is deriving the inverse\n",
-    "transform instead of the forward transform, if unexpected results are\n",
-    "obtained, first try passing the transpose of the rotation matrix."
+    "the host of a satellite galaxy)."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "scrolled": true
+   },
    "outputs": [],
    "source": [
     "source = EAGLESource(\n",
@@ -350,7 +351,7 @@
     "    subBoxSize=150 * U.kpc,\n",
     "    distance=4 * U.Mpc,\n",
     "    vpeculiar=0 * U.km / U.s,\n",
-    "    rotation={\"L_coords\": (60 * U.deg, 90.0 * U.deg)},\n",
+    "    rotation={\"rotmat\": np.eye(3)},\n",
     "    ra=0.0 * U.deg,\n",
     "    dec=0.0 * U.deg,\n",
     ")"
@@ -360,7 +361,90 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "If you have not entered your username and password in an earlier cell, you can edit the `db_user` and `db_key` arguments to provide them. The `db_key` (password) argument can also be omitted, in which case you'll be prompted to enter it interactively."
+    "If you have not entered your username and password in an earlier cell, you can edit the `db_user` and `db_key` arguments to provide them. The `db_key` (password) argument can also be omitted, in which case you'll be prompted to enter it interactively.\n",
+    "\n",
+    "The rotation argument above has been set to the identity matrix, so the source has the (random) orientation that it has within the simulation volume. The source class includes a function to make a quick plot to get an idea of the source's appearance:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preview_fig_unrotated = source.preview(title=\"unrotated\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The preview function defaults to apertures in position and velocity that enclose all particles in the source, so this preview emphasizes the diffuse circumgalactic gas. The apertures can be set manually using the `lim` and `vlim` keywords to set the maximum absolute offsets in position and velocity relative to the source centre to be plotted. For example, restricting the aperture to 50kpc and 300km/s makes the disc more clearly visible."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preview_fig_unrotated_zoom = source.preview(\n",
+    "    title=\"unrotated, zoomed-in\",\n",
+    "    lim=50 * U.kpc,\n",
+    "    vlim=300 * U.km / U.s,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This randomly-oriented viewing angle seems to be close to edge-on. The source can be rotated to a different orientation. MARTINI's tool for quick/approximate manipulation of the orientation of the source aligns the source based on its angular momentum vector (\"L\"), for example:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "source.rotate(L_coords=(60 * U.deg, 90 * U.deg))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The rotation configuration takes an inclination (here 60deg) and rotation about the pole (here 90deg, relative to an arbitrary reference direction). The code attempts to\n",
+    "automatically align the galactic disk in the y-z plane by aligning\n",
+    "the angular momentum along the x-axis. The polar rotation is then\n",
+    "applied, and finally the disc inclined by a rotation around the\n",
+    "y-axis (the line of sight is along the x-axis). The automatic\n",
+    "alignment will work for typical reasonably isolated discs, but will\n",
+    "struggle when companions are present, when the angular momentum axis\n",
+    "is a poor tracer of the disc plane, and especially for satellites."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preview_fig_rotated_zoomed_in = source.preview(\n",
+    "    title=\"rotated, zoomed-in\",\n",
+    "    lim=50 * U.kpc,\n",
+    "    vlim=300 * U.km / U.s,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If finer control of the orientation is needed, derive the transformation from the simulation box coordinates (see [the documentation](https://martini.readthedocs.io/en/latest/) for examples) to the desired coordinates for the 'observation', keeping in mind that the line of sight is along the x-axis. This rotation matrix can then be passed to the rotate function as `rotmat=np.eye(3)` (here the identity rotation matrix used as a place holder). The rotation can also be provided when the source is initialized by using the `rotation` keyword argument.\n",
+    "\n",
+    "A common problem is deriving the inverse transform instead of the forward transform, if unexpected results are obtained, first try passing the transpose of the rotation matrix."
    ]
   },
   {
@@ -491,7 +575,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The calculation of the spectra can be done in parallel by providing a keyword argument `ncpu=N`, where `N` is the number of CPUs to use. However, the details of the implementation mean that for small numbers of particles running in parallel tends to slow down the calculation, so turning this on should be done with care. Significant speedups can be expected when the particl"
+    "The calculation of the spectra (that will happen when the `Martini` module is initialized below) can be done in parallel by providing a keyword argument `ncpu=N`, where `N` is the number of CPUs to use. However, the details of the implementation mean that for small numbers of particles running in parallel tends to slow down the calculation, so turning this on should be done with care. Significant speedups can be expected when the particle count is very large."
    ]
   },
   {
diff --git a/examples/martini_simba.ipynb b/examples/martini_simba.ipynb
index 8aa49ac..993bff7 100644
--- a/examples/martini_simba.ipynb
+++ b/examples/martini_simba.ipynb
@@ -20,7 +20,7 @@
    "source": [
     "MARTINI is a modular package for the creation of synthetic resolved HI line observations (data cubes) of smoothed-particle hydrodynamics simulations of galaxies. The various aspects of the mock-observing process are divided logically into sub-modules handling the data cube, source, beam, noise, spectral model and SPH kernel. MARTINI is object-oriented: each sub-module provides a class (or classes) which can be configured as desired. For most sub-modules, base classes are provided to allow for straightforward customization. Instances of each sub-module class are given as parameters to the Martini class; a mock observation is then constructed by calling a handful of functions to execute the desired steps in the mock-observing process.\n",
     "\n",
-    "The MARTINI documentation is [hosted on ReadTheDocs](https://martini.readthedocs.io/en/latest/)."
+    "This tutorial focuses on particulars related to working with the Simba simulations. More general information is available in the MARTINI documentation, [hosted on ReadTheDocs](https://martini.readthedocs.io/en/latest/)."
    ]
   },
   {
@@ -69,6 +69,22 @@
     "!{sys.executable} -m pip install requests matplotlib"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This cell may be needed in some cases to display figures below:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -255,22 +271,7 @@
     "capture the region of interest around it (e.g. enough to encompass\n",
     "the host of a satellite galaxy).\n",
     "\n",
-    "The rotation configuration takes an inclination (here 60deg) and\n",
-    "rotation about the pole (here 90deg, relative to an arbitrary reference direction). The code attempts to\n",
-    "automatically align the galactic disk in the y-z plane by aligning\n",
-    "the angular momentum along the x-axis. The polar rotation is then\n",
-    "applied, and finally the disc inclined by a rotation around the\n",
-    "y-axis (the line of sight is along the x-axis). The automatic\n",
-    "alignment will work for typical reasonably isolated discs, but will\n",
-    "struggle when companions are present, when the angular momentum axis\n",
-    "is a poor tracer of the disc plane, and especially for satellites. If\n",
-    "finer control of the orientation is needed, derive the transformation\n",
-    "from the simulation box coordinates to the desired coordinates for\n",
-    "the 'observation', keeping in mind that the line of sight is along\n",
-    "the x-axis. This rotation matrix can then be passed to as\n",
-    "`rotation={'rotmat': np.eye(3)}` (here the identity rotation matrix used as a place holder). A common problem in this case is deriving the inverse\n",
-    "transform instead of the forward transform, if unexpected results are\n",
-    "obtained, first try passing the transpose of the rotation matrix."
+    "Running the following cell temporarily uses about 20 GB of system memory as entire particle arrays are read and then pruned to retain only the particles belonging to the galaxy of interest."
    ]
   },
   {
@@ -288,12 +289,100 @@
     "    aperture=100.0 * U.kpc,\n",
     "    distance=4.0 * U.Mpc,\n",
     "    vpeculiar=0 * U.km / U.s,\n",
-    "    rotation={\"L_coords\": (60 * U.deg, 90.0 * U.deg)},\n",
+    "    rotation={\"rotmat\": np.eye(3)},\n",
     "    ra=0.0 * U.deg,\n",
     "    dec=0.0 * U.deg,\n",
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The rotation argument above has been set to the identity matrix, so the source has the (random) orientation that it has within the simulation volume. The source class includes a function to make a quick plot to get an idea of the source's appearance:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preview_fig_unrotated = source.preview(title=\"unrotated\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The preview function defaults to apertures in position and velocity that enclose all particles in the source, so this preview emphasizes the diffuse circumgalactic gas. The apertures can be set manually using the `lim` and `vlim` keywords to set the maximum absolute offsets in position and velocity relative to the source centre to be plotted. For example, restricting the aperture to 50kpc and 300km/s makes the disc more clearly visible."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preview_fig_unrotated_zoom = source.preview(\n",
+    "    title=\"unrotated, zoomed-in\",\n",
+    "    lim=50 * U.kpc,\n",
+    "    vlim=300 * U.km / U.s,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This randomly-oriented viewing angle seems to be moderately inclined with respect to the disc. The source can be rotated to a different orientation. MARTINI's tool for quick/approximate manipulation of the orientation of the source aligns the source based on its angular momentum vector (\"L\"), for example:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "source.rotate(L_coords=(60 * U.deg, 90 * U.deg))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The rotation configuration takes an inclination (here 60deg) and rotation about the pole (here 90deg, relative to an arbitrary reference direction). The code attempts to\n",
+    "automatically align the galactic disk in the y-z plane by aligning\n",
+    "the angular momentum along the x-axis. The polar rotation is then\n",
+    "applied, and finally the disc inclined by a rotation around the\n",
+    "y-axis (the line of sight is along the x-axis). The automatic\n",
+    "alignment will work for typical reasonably isolated discs, but will\n",
+    "struggle when companions are present, when the angular momentum axis\n",
+    "is a poor tracer of the disc plane, and especially for satellites."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preview_fig_rotated_zoomed_in = source.preview(\n",
+    "    title=\"rotated, zoomed-in\",\n",
+    "    lim=50 * U.kpc,\n",
+    "    vlim=300 * U.km / U.s,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "If finer control of the orientation is needed, derive the transformation from the simulation box coordinates (see [the documentation](https://martini.readthedocs.io/en/latest/) for examples) to the desired coordinates for the 'observation', keeping in mind that the line of sight is along the x-axis. This rotation matrix can then be passed to the rotate function as `rotmat=np.eye(3)` (here the identity rotation matrix used as a place holder). The rotation can also be provided when the source is initialized by using the `rotation` keyword argument.\n",
+    "\n",
+    "A common problem is deriving the inverse transform instead of the forward transform, if unexpected results are obtained, first try passing the transpose of the rotation matrix."
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -427,7 +516,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The calculation of the spectra can be done in parallel by providing a keyword argument `ncpu=N`, where `N` is the number of CPUs to use. However, the details of the implementation mean that for small numbers of particles running in parallel tends to slow down the calculation, so turning this on should be done with care. Significant speedups can be expected when the particl"
+    "The calculation of the spectra (that will happen when the `Martini` module is initialized below) can be done in parallel by providing a keyword argument `ncpu=N`, where `N` is the number of CPUs to use. However, the details of the implementation mean that for small numbers of particles running in parallel tends to slow down the calculation, so turning this on should be done with care. Significant speedups can be expected when the particle count is very large."
    ]
   },
   {
@@ -644,7 +733,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "This galaxy clearly has a very nice spiral morphology in HI, a central hole, and a nice rotation dominated velocity field. The alignment of the disc looks as expected - the inclination looks to be about 60 degrees, and the position angle is horizontal in the figure - in this case the automated orientation function has performed well, though it should never be assumed that this will always be the case!\n",
+    "This galaxy clearly has a very nice spiral morphology in HI, a central hole, and a nice rotation-dominated velocity field. The alignment of the disc looks as expected - the inclination looks to be about 60 degrees, and the position angle is horizontal in the figure - in this case the automated orientation function has performed well, though it should never be assumed that this will always be the case!\n",
     "\n",
     "For complete documentation, more usage examples, and further information, please take a look at the [MARTINI webpage](https://martini.readthedocs.io/en/latest/)."
    ]
diff --git a/martini/__init__.py b/martini/__init__.py
index 494cea0..4a2e3c1 100644
--- a/martini/__init__.py
+++ b/martini/__init__.py
@@ -1,6 +1,6 @@
 from .martini import Martini
 from .datacube import DataCube
-from ._demo import demo
+from ._demo import demo, demo_source
 from .__version__ import __version__
 
 __doc__ = """
diff --git a/martini/_demo.py b/martini/_demo.py
index 682b96b..e571767 100644
--- a/martini/_demo.py
+++ b/martini/_demo.py
@@ -9,26 +9,15 @@
 from scipy.optimize import fsolve
 
 
-def demo(cubefile="testcube.fits", beamfile="testbeam.fits", hdf5file="testcube.hdf5"):
+def demo_source(N=500):
     """
-    Demonstrates basic usage of MARTINI.
-
-    Creates a (very!) crude toy model of a galaxy with a linearly rising
-    rotation curve, exponential disk profile, exponential vertical structure. A
-    basic configuration of MARTINI is initialized and used to create and output
-    a datacube and an image of the beam.
-
-    Parameters
-    ----------
-    cubefile : string
-        File to write demonstration datacube.
+    Create a simple toy model of a galaxy.
 
-    beamfile : string
-        File to write demonstration beam.
+    Returns
+    -------
+    out : martini.source.SPHSource
+        An initialized MARTINI source module containing a toy model of a galaxy.
     """
-
-    # ------make a toy galaxy----------
-    N = 500
     phi = np.random.rand(N) * 2 * np.pi
     r = []
     for L in np.random.rand(N):
@@ -46,17 +35,19 @@ def f(r):
     x = r * np.cos(phi)
     y = r * np.sin(phi)
     xyz_g = np.vstack((x, y, z)) * U.kpc
-    # linear rotation curve
-    vphi = 100 * r / 6.0
+    # arctan rotation curve
+    vphi = 50 * np.arctan(r)
     vx = -vphi * np.sin(phi)
     vy = vphi * np.cos(phi)
-    # small pure random z velocities
+    # small random z velocities
     vz = (np.random.rand(N) * 2.0 - 1.0) * 5
     vxyz_g = np.vstack((vx, vy, vz)) * U.km * U.s**-1
     T_g = np.ones(N) * 8e3 * U.K
-    mHI_g = np.ones(N) / N * 5.0e9 * U.Msun
-    # ~mean interparticle spacing smoothing
-    hsm_g = np.ones(N) * 4 / np.sqrt(N) * U.kpc
+    # HI masses with some scatter
+    mHI_g = np.ones(N) + 0.01 * (np.random.rand(N) - 0.5)
+    mHI_g = mHI_g / mHI_g.sum() * 5.0e9 * U.Msun
+    # Smoothing lengths with some scatter
+    hsm_g = (40 / np.sqrt(N) * (np.ones(N) + 1.8 * (np.random.rand(N) - 0.5))) * U.kpc
     # ---------------------------------
 
     source = SPHSource(
@@ -72,6 +63,28 @@ def f(r):
         hsm_g=hsm_g,
     )
 
+    return source
+
+
+def demo(cubefile="testcube.fits", beamfile="testbeam.fits", hdf5file="testcube.hdf5"):
+    """
+    Demonstrates basic usage of MARTINI.
+
+    Creates a (very!) crude toy model of a galaxy with a linearly rising
+    rotation curve, exponential disk profile, exponential vertical structure. A
+    basic configuration of MARTINI is initialized and used to create and output
+    a datacube and an image of the beam.
+
+    Parameters
+    ----------
+    cubefile : string
+        File to write demonstration datacube.
+
+    beamfile : string
+        File to write demonstration beam.
+    """
+    source = demo_source()
+
     datacube = DataCube(
         n_px_x=128,
         n_px_y=128,
@@ -112,5 +125,3 @@ def f(r):
         print("h5py package not present, skipping hdf5 output demo.")
     else:
         print(f"Wrote demo hdf5 output to {hdf5file}.")
-
-    return
diff --git a/martini/_demo.pyi b/martini/_demo.pyi
new file mode 100644
index 0000000..4da859c
--- /dev/null
+++ b/martini/_demo.pyi
@@ -0,0 +1,4 @@
+from .sources import SPHSource
+
+def demo_source(N: int = ...) -> SPHSource: ...
+def demo(cubefile: str = ..., beamfile: str = ..., hdf5file: str = ...) -> None: ...
diff --git a/martini/sources/colibre_source.py b/martini/sources/colibre_source.py
index fc9b8ec..0b41f68 100644
--- a/martini/sources/colibre_source.py
+++ b/martini/sources/colibre_source.py
@@ -38,7 +38,7 @@ class ColibreSource(SWIFTGalaxySource):
         "sky" ('y-z'), rotated by the azimuthal angle about its angular \
         momentum pole (rotation about 'x'), and inclined (rotation about \
         'y'). A 3-tuple may be provided instead, in which case the third \
-        value specifies the position angle on the sky (rotation about 'x'). \
+        value specifies the position angle on the sky (second rotation about 'x'). \
         The default position angle is 270 degrees.
 
         (Default: identity rotation matrix.)
diff --git a/martini/sources/eagle_source.py b/martini/sources/eagle_source.py
index 39c608e..026d142 100644
--- a/martini/sources/eagle_source.py
+++ b/martini/sources/eagle_source.py
@@ -73,7 +73,7 @@ class EAGLESource(SPHSource):
         "sky" ('y-z'), rotated by the azimuthal angle about its angular \
         momentum pole (rotation about 'x'), and inclined (rotation about \
         'y'). A 3-tuple may be provided instead, in which case the third \
-        value specifies the position angle on the sky (rotation about 'x'). \
+        value specifies the position angle on the sky (second rotation about 'x'). \
         The default position angle is 270 degrees.
 
         (Default: identity rotation matrix.)
diff --git a/martini/sources/magneticum_source.py b/martini/sources/magneticum_source.py
index 6dbb5e3..42fc99c 100644
--- a/martini/sources/magneticum_source.py
+++ b/martini/sources/magneticum_source.py
@@ -90,7 +90,7 @@ class MagneticumSource(SPHSource):
         "sky" ('y-z'), rotated by the azimuthal angle about its angular \
         momentum pole (rotation about 'x'), and inclined (rotation about \
         'y'). A 3-tuple may be provided instead, in which case the third \
-        value specifies the position angle on the sky (rotation about 'x'). \
+        value specifies the position angle on the sky (second rotation about 'x'). \
         The default position angle is 270 degrees.
 
         (Default: identity rotation matrix.)
diff --git a/martini/sources/simba_source.py b/martini/sources/simba_source.py
index 6968497..1b98ec8 100644
--- a/martini/sources/simba_source.py
+++ b/martini/sources/simba_source.py
@@ -61,7 +61,7 @@ class SimbaSource(SPHSource):
         "sky" ('y-z'), rotated by the azimuthal angle about its angular \
         momentum pole (rotation about 'x'), and inclined (rotation about \
         'y'). A 3-tuple may be provided instead, in which case the third \
-        value specifies the position angle on the sky (rotation about 'x'). \
+        value specifies the position angle on the sky (second rotation about 'x'). \
         The default position angle is 270 degrees.
 
         (Default: identity rotation matrix.)
diff --git a/martini/sources/so_source.py b/martini/sources/so_source.py
index e38e3fe..89d8588 100644
--- a/martini/sources/so_source.py
+++ b/martini/sources/so_source.py
@@ -39,7 +39,7 @@ class SOSource(SPHSource):
         "sky" ('y-z'), rotated by the azimuthal angle about its angular \
         momentum pole (rotation about 'x'), and inclined (rotation about \
         'y'). A 3-tuple may be provided instead, in which case the third \
-        value specifies the position angle on the sky (rotation about 'x'). \
+        value specifies the position angle on the sky (second rotation about 'x'). \
         The default position angle is 270 degrees.
 
         (Default: identity rotation matrix.)
diff --git a/martini/sources/sph_source.py b/martini/sources/sph_source.py
index 7edc8c7..6a92e34 100644
--- a/martini/sources/sph_source.py
+++ b/martini/sources/sph_source.py
@@ -48,7 +48,7 @@ class SPHSource(object):
         "sky" ('y-z'), rotated by the azimuthal angle about its angular \
         momentum pole (rotation about 'x'), and inclined (rotation about \
         'y'). A 3-tuple may be provided instead, in which case the third \
-        value specifies the position angle on the sky (rotation about 'x'). \
+        value specifies the position angle on the sky (second rotation about 'x'). \
         The default position angle is 270 degrees.
 
         (Default: identity rotation matrix.)
@@ -255,7 +255,7 @@ def rotate(self, axis_angle=None, rotmat=None, L_coords=None):
             about its angular momentum pole (rotation about 'x'), and then
             inclined (rotation about 'y'). By default the position angle on the
             sky is 270 degrees, but if a third element is provided it sets the
-            position angle (rotation about 'z').
+            position angle (second rotation about 'x').
         """
 
         args_given = (axis_angle is not None, rotmat is not None, L_coords is not None)
@@ -350,3 +350,180 @@ def save_current_rotation(self, fname):
 
         np.savetxt(fname, self.current_rotation)
         return
+
+    def preview(
+        self,
+        max_points=5000,
+        fig=1,
+        lim=None,
+        vlim=None,
+        point_scaling="auto",
+        title="",
+        save=None,
+    ):
+        """
+        Produce a figure showing the source particle coordinates and velocities.
+
+        Makes a 3-panel figure showing the projection of the source as it will appear in
+        the mock observation. The first panel shows the particles in the y-z plane,
+        coloured by the x-component of velocity (MARTINI projects the source along the
+        x-axis). The second and third panels are position-velocity diagrams showing the
+        x-component of velocity against the y and z coordinates, respectively.
+
+        Parameters
+        ----------
+        max_points : int, optional
+            Maximum number of points to draw per panel, the particles will be randomly
+            subsampled if the source has more. (Default: 1000)
+
+        fig : int, optional
+            Number of the figure in matplotlib, it will be created as `plt.figure(fig)`.
+            (Default: 1)
+
+        lim : Quantity with dimensions of length, optional
+            The coordinate axes extend from -lim to lim. If unspecified, the maximum
+            absolute coordinate of particles in the source is used. (Default: None)
+
+        vlim : Quantity with dimensions of speed, optional
+            The velocity axes and colour bar extend from -vlim to vlim. If unspecified,
+            the maximum absolute velocity of particles in the source is used.
+            (Default: None)
+
+        point_scaling : str, optional
+            By default points are scaled in size and transparency according to their HI
+            mass and the smoothing length (loosely proportional to their surface
+            densities, but with different scaling to achieve a visually useful plot). For
+            some sources the automatic scaling may not give useful results, using
+            point_scaling="fixed" will plot points of constant size without opacity.
+            (Default: "auto")
+
+        title : str, optional
+            A title for the figure can be provided. (Default: "")
+
+        save : str, optional
+            If provided, the figure is saved using `plt.savefig(save)`. A `.png` or `.pdf`
+            suffix is recommended. (Default: None)
+
+        Returns
+        -------
+        out : matplotlib.figure instance
+            The preview figure.
+        """
+        import matplotlib.pyplot as plt
+
+        # every Nth particle to plot at most max_points, or all particles
+        lim = (
+            max(
+                np.max(np.abs(self.coordinates_g.y.to_value(U.kpc))),
+                np.max(np.abs(self.coordinates_g.z.to_value(U.kpc))),
+            )
+            if lim is None
+            else lim.to_value(U.kpc)
+        )
+        vlim = (
+            np.max(
+                np.abs(self.coordinates_g.differentials["s"].d_x.to_value(U.km / U.s))
+            )
+            if vlim is None
+            else vlim.to_value(U.km / U.s)
+        )
+        cmask = np.logical_and.reduce(
+            (
+                np.abs(self.coordinates_g.y.to_value(U.kpc)) < lim,
+                np.abs(self.coordinates_g.z.to_value(U.kpc)) < lim,
+                np.abs(self.coordinates_g.differentials["s"].d_x.to_value(U.km / U.s))
+                < vlim,
+            )
+        )
+        nparts = cmask.sum()
+        mask = np.arange(self.mHI_g.size)[cmask][:: max(nparts // max_points, 1)]
+        mass_factor = (
+            1
+            if (self.mHI_g.isscalar or mask.size <= 1)
+            else (self.mHI_g[mask] / self.mHI_g[mask].max()).to_value(
+                U.dimensionless_unscaled
+            )
+        )
+        hsm_factor = (
+            1
+            if (self.hsm_g.isscalar or mask.size <= 1)
+            else (1 - (self.hsm_g[mask] / self.hsm_g[mask].max())).to_value(
+                U.dimensionless_unscaled
+            )
+        )
+        alpha = (mass_factor**0 * hsm_factor**0.1) if point_scaling == "auto" else 1.0
+        size_scale = (
+            self.hsm_g.to_value(U.kpc) / lim
+            if (self.hsm_g.isscalar or mask.size <= 1)
+            else (self.hsm_g[mask].to_value(U.kpc) / lim)
+        )
+        size = 300 * size_scale if point_scaling == "auto" else 10
+        fig = plt.figure(fig, figsize=(12, 4))
+        fig.suptitle(title)
+
+        # ----- MOMENT 1 -----
+        sp1 = fig.add_subplot(1, 3, 1, aspect="equal")
+        sp1.set_facecolor("#222222")
+        scatter = sp1.scatter(
+            self.coordinates_g.y[mask].to_value(U.kpc),
+            self.coordinates_g.z[mask].to_value(U.kpc),
+            c=self.coordinates_g.differentials["s"].d_x[mask].to_value(U.km / U.s),
+            marker="o",
+            cmap="coolwarm",
+            edgecolor="None",
+            s=size,
+            alpha=alpha,
+            vmin=-vlim,
+            vmax=vlim,
+            zorder=0,
+        )
+        sp1.plot([0], [0], marker="+", ls="None", mec="grey", ms=6, zorder=1)
+        sp1.set_xlabel(r"$y\,[\mathrm{kpc}]$")
+        sp1.set_ylabel(r"$z\,[\mathrm{kpc}]$")
+        sp1.set_xlim((lim, -lim))
+        sp1.set_ylim((-lim, lim))
+        cb = fig.colorbar(mappable=scatter, ax=sp1, orientation="horizontal")
+        cb.set_label(r"$v_x\,[\mathrm{km}\,\mathrm{s}^{-1}]$")
+
+        # ----- PV Y -----
+        sp2 = fig.add_subplot(1, 3, 2)
+        sp2.set_facecolor("#222222")
+        sp2.scatter(
+            self.coordinates_g.y[mask].to_value(U.kpc),
+            self.coordinates_g.differentials["s"].d_x[mask].to_value(U.km / U.s),
+            c="white",
+            edgecolors="None",
+            marker="o",
+            s=size,
+            alpha=alpha,
+            zorder=0,
+        )
+        sp2.plot([0], [0], marker="+", ls="None", mec="red", ms=6, zorder=1)
+        sp2.set_xlim((lim, -lim))
+        sp2.set_ylim((-vlim, vlim))
+        sp2.set_xlabel(r"$y\,[\mathrm{kpc}]$")
+        sp2.set_ylabel(r"$v_x\,[\mathrm{km}\,\mathrm{s}^{-1}]$")
+
+        # ----- PV Z -----
+        sp3 = fig.add_subplot(1, 3, 3)
+        sp3.set_facecolor("#222222")
+        sp3.scatter(
+            self.coordinates_g.z[mask].to_value(U.kpc),
+            self.coordinates_g.differentials["s"].d_x[mask].to_value(U.km / U.s),
+            c="white",
+            edgecolors="None",
+            marker="o",
+            s=size,
+            alpha=alpha,
+            zorder=0,
+        )
+        sp3.plot([0], [0], marker="+", ls="None", mec="red", ms=6, zorder=1)
+        sp3.set_xlim((-lim, lim))
+        sp3.set_ylim((-vlim, vlim))
+        sp3.set_xlabel(r"$z\,[\mathrm{kpc}]$")
+        sp3.set_ylabel(r"$v_x\,[\mathrm{km}\,\mathrm{s}^{-1}]$")
+
+        fig.subplots_adjust(wspace=0.3)
+        if save is not None:
+            plt.savefig(save)
+        return fig
diff --git a/martini/sources/sph_source.pyi b/martini/sources/sph_source.pyi
index 5a095d7..b5cce8d 100644
--- a/martini/sources/sph_source.pyi
+++ b/martini/sources/sph_source.pyi
@@ -3,6 +3,7 @@ from astropy.coordinates import ICRS
 from martini.datacube import DataCube
 from numpy import ndarray
 import typing as T
+from matplotlib.figure import Figure
 
 class SPHSource:
     h: float
@@ -56,3 +57,13 @@ class SPHSource:
     def translate(self, translation_vector: U.Quantity[U.kpc]) -> None: ...
     def boost(self, boost_vector: U.Quantity[U.km / U.s]) -> None: ...
     def save_current_rotation(self, fname: str) -> None: ...
+    def preview(
+        self,
+        max_points: int = ...,
+        fig: int = ...,
+        lim: T.Optional[U.Quantity[U.kpc]] = ...,
+        vlim: T.Optional[U.Quantity[U.km / U.s]] = ...,
+        point_scaling: str = ...,
+        title: str = ...,
+        save: T.Optional[str] = ...,
+    ) -> Figure: ...
diff --git a/martini/sources/swiftgalaxy_source.py b/martini/sources/swiftgalaxy_source.py
index da6bdff..ffb0a92 100644
--- a/martini/sources/swiftgalaxy_source.py
+++ b/martini/sources/swiftgalaxy_source.py
@@ -38,7 +38,7 @@ class SWIFTGalaxySource(SPHSource):
         "sky" ('y-z'), rotated by the azimuthal angle about its angular \
         momentum pole (rotation about 'x'), and inclined (rotation about \
         'y'). A 3-tuple may be provided instead, in which case the third \
-        value specifies the position angle on the sky (rotation about 'x'). \
+        value specifies the position angle on the sky (second rotation about 'x'). \
         The default position angle is 270 degrees.
 
         (Default: identity rotation matrix.)
diff --git a/martini/sources/tng_source.py b/martini/sources/tng_source.py
index 1673d11..7d7822f 100644
--- a/martini/sources/tng_source.py
+++ b/martini/sources/tng_source.py
@@ -100,7 +100,7 @@ class TNGSource(SPHSource):
         "sky" ('y-z'), rotated by the azimuthal angle about its angular \
         momentum pole (rotation about 'x'), and inclined (rotation about \
         'y'). A 3-tuple may be provided instead, in which case the third \
-        value specifies the position angle on the sky (rotation about 'x'). \
+        value specifies the position angle on the sky (second rotation about 'x'). \
         The default position angle is 270 degrees.
 
         (Default: identity rotation matrix.)
diff --git a/martini/spectral_models.pyi b/martini/spectral_models.pyi
index 4ed5823..4ce6c91 100644
--- a/martini/spectral_models.pyi
+++ b/martini/spectral_models.pyi
@@ -15,18 +15,16 @@ class _BaseSpectrum(metaclass=abc.ABCMeta):
     spec_dtype: type
     spectral_function_extra_data: T.Optional[T.Dict[str, T.Any]]
 
-    def __init__(
-        self, ncpu: T.Optional[int] = None, spec_dtype: type = ...
-    ) -> None: ...
+    def __init__(self, ncpu: T.Optional[int] = ..., spec_dtype: type = ...) -> None: ...
     def init_spectra(self, source: SPHSource, datacube: DataCube) -> None: ...
     def evaluate_spectra(
         self,
         source: SPHSource,
         datacube: DataCube,
-        mask: T.Union[slice, EllipsisType] = np.s_[...],
+        mask: T.Union[slice, EllipsisType] = ...,
     ) -> None: ...
     @abstractmethod
-    def half_width(self, source: SPHSource): ...
+    def half_width(self, source: SPHSource) -> U.Quantity[U.km / U.s]: ...
     @abstractmethod
     def spectral_function(
         self,
@@ -38,7 +36,7 @@ class _BaseSpectrum(metaclass=abc.ABCMeta):
         self,
         source: SPHSource,
         datacube: DataCube,
-        mask: T.Union[slice, EllipsisType] = np.s_[...],
+        mask: T.Union[slice, EllipsisType] = ...,
     ) -> None: ...
 
 class GaussianSpectrum(_BaseSpectrum):
@@ -60,18 +58,16 @@ class GaussianSpectrum(_BaseSpectrum):
         self,
         source: SPHSource,
         datacube: DataCube,
-        mask: T.Union[slice, EllipsisType] = np.s_[...],
+        mask: T.Union[slice, EllipsisType] = ...,
     ) -> None: ...
-    def half_width(self, source: SPHSource) -> None: ...
+    def half_width(self, source: SPHSource) -> U.Quantity[U.km / U.s]: ...
 
 class DiracDeltaSpectrum(_BaseSpectrum):
-    def __init__(
-        self, ncpu: T.Optional[int] = None, spec_dtype: type = ...
-    ) -> None: ...
+    def __init__(self, ncpu: T.Optional[int] = ..., spec_dtype: type = ...) -> None: ...
     def spectral_function(
         self,
         a: U.Quantity[U.km / U.s],
         b: U.Quantity[U.km / U.s],
         vmids: U.Quantity[U.km / U.s],
     ) -> U.Quantity[U.dimensionless_unscaled]: ...
-    def half_width(self, source: SPHSource) -> None: ...
+    def half_width(self, source: SPHSource) -> U.Quantity[U.km / U.s]: ...
diff --git a/mypy.ini b/mypy.ini
index e9f20fd..0d06c7a 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -1,3 +1,3 @@
 [mypy]
 ignore_missing_imports = True
-modules = martini, martini.beams, martini.datacube, martini.martini, martini.noise, martini.spectral_models, martini.sph_kernels, martini.sources.sph_source, martini.sources._cartesian_translation, martini.sources._L_align
\ No newline at end of file
+modules = martini, martini.beams, martini.datacube, martini.martini, martini.noise, martini.spectral_models, martini.sph_kernels, martini.sources.sph_source, martini.sources._cartesian_translation, martini.sources._L_align, martini._demo
\ No newline at end of file
diff --git a/optional_requirements.txt b/optional_requirements.txt
index 25ab41c..75b3805 100644
--- a/optional_requirements.txt
+++ b/optional_requirements.txt
@@ -1,3 +1,4 @@
+matplotlib
 six
 requests
 multiprocess


Test patch:
diff --git a/run_stubtest b/run_stubtest
index 28f5298..714f242 100755
--- a/run_stubtest
+++ b/run_stubtest
@@ -1,1 +1,1 @@
-stubtest --mypy-config-file mypy.ini --allowlist stubtest_allowlist martini.martini martini.beams martini.datacube martini.noise martini.spectral_models martini.sph_kernels martini.sources.sph_source martini.sources._L_align martini.sources._cartesian_translation martini.sources._illustris_tools
\ No newline at end of file
+stubtest --mypy-config-file mypy.ini --allowlist stubtest_allowlist martini.martini martini.beams martini.datacube martini.noise martini.spectral_models martini.sph_kernels martini.sources.sph_source martini.sources._L_align martini.sources._cartesian_translation martini.sources._illustris_tools martini._demo
\ No newline at end of file
diff --git a/tests/test_demo.py b/tests/test_demo.py
index b2e5cb2..73c7440 100644
--- a/tests/test_demo.py
+++ b/tests/test_demo.py
@@ -1,9 +1,12 @@
-from martini import demo
+from martini._demo import demo, demo_source
 import os
 
 
 class TestDemo:
     def test_demo(self):
+        """
+        Check that demo completes and produces output files.
+        """
         filenames = dict(
             cubefile="pytest_testcube.fits",
             beamfile="pytest_testbeam.fits",
@@ -15,3 +18,11 @@ def test_demo(self):
             for filename in filenames.values():
                 if os.path.exists(filename):
                     os.remove(filename)
+
+    def test_demo_source(self):
+        """
+        Check that demo source created and contains particles.
+        """
+        N = 100
+        source = demo_source(N=N)
+        assert source.mHI_g.size == N
diff --git a/tests/test_sources.py b/tests/test_sources.py
index 6570d54..d0348fd 100644
--- a/tests/test_sources.py
+++ b/tests/test_sources.py
@@ -449,6 +449,34 @@ def test_save_current_rotation(self, s):
                 os.remove(testfile)
         assert np.allclose(saved_rotmat, rotmat)
 
+    def test_preview(self, s):
+        """
+        Simply check that the preview visualisation runs without error.
+        """
+        # with default arguments
+        s.preview()
+        # with non-default arguments
+        s.preview(
+            max_points=1000,
+            fig=2,
+            lim=10 * U.kpc,
+            vlim=100 * U.km / U.s,
+            point_scaling="fixed",
+            title="test",
+        )
+
+    @pytest.mark.parametrize("ext", ("pdf", "png"))
+    def test_preview_save(self, s, ext):
+        """
+        Check that we can output pdf and png preview images.
+        """
+        testfile = f"preview.{ext}"
+        try:
+            s.preview(save=testfile)
+        finally:
+            if os.path.exists(testfile):
+                os.remove(testfile)
+
 
 class TestSOSource:
     @pytest.mark.xfail