src/empyre/vis/plot2d.py

+# -*- coding: utf-8 -*-
+# Copyright 2019 by Forschungszentrum Juelich GmbH
+# Author: J. Caron
+#
+"""This module provides functions for 2D plots that often wrap functions from `maptlotlib.pyplot`."""
+
+
+import logging
+import warnings
+
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.colors import TwoSlopeNorm
+from PIL import Image
+
+from . import colors
+from .tools import use_style
+from ..fields.field import Field
+
+
+__all__ = ['imshow', 'contour', 'colorvec', 'cosine_contours', 'quiver']
+
+_log = logging.getLogger(__name__)
+
+
+DIVERGING_CMAPS = ['PiYG', 'PRGn', 'BrBG', 'PuOr', 'RdGy', 'RdBu', 'RdYlBu', 'RdYlGn',
+                   'Spectral', 'coolwarm', 'bwr', 'seismic',  # all divergent maps from matplotlib!
+                   'balance', 'delta', 'curl', 'diff', 'tarn']  # all divergent maps from cmocean!
+# TODO: add seaborn and more?
+
+
+def imshow(field, axis=None, cmap=None, **kwargs):
+    """Display an image on a 2D regular raster. Wrapper for `matplotlib.pyplot.imshow`.
+
+    Parameters
+    ----------
+    field : `Field` or ndarray
+        The image data as a `Field` or a numpy array (in the latter case, `vector=False` and `scale=1.0` are assumed).
+    axis : `matplotlib.axes.Axes` object, optional
+        The axis to which the image should be added, by default None, which will pick the last use axis via `gca`.
+    cmap : str or `matplotlib.colors.Colormap`, optional
+        The Colormap that should be used for the display, either as a string or object, by default None, which will pick
+        `cmocean.cm.balance` if available. `imshow` will automatically detect if a divergent colormap is used and will
+        make sure that zero is pinned to the symmetry point of the colormap (this is done by creating a new colormap
+        with custom range under the hood).
+
+    Returns
+    -------
+    axis : `matplotlib.axes.Axes`
+        The plotting axis.
+
+    Notes
+    -----
+    Additional kwargs are passed to :meth:`~matplotlib.pyplot.imshow`.
+    Note that the y-axis of the plot is flipped in comparison to :meth:`~matplotlib.pyplot.imshow`, i.e. that the
+    origin is `'lower'` in this case instead of `'upper'`.
+    Uses the `empyre-image` stylesheet settings for plotting (and axis creation if none exists, yet).
+    Fields are squeezed before plotting, so non-2D fields work as long as their superfluous dimensions have length 1.
+
+    """
+    _log.debug('Calling imshow')
+    if not isinstance(field, Field):  # Try to convert input to Field if it is not already one:
+        field = Field(data=np.asarray(field), scale=1.0, vector=False)
+    assert not field.vector, 'Can only plot scalar fields!'
+    # Get squeezed data and make sure it's 2D scalar:
+    squeezed_field = field.squeeze()
+    assert len(squeezed_field.dim) == 2, 'Cannot plot more than 2 dimensions!'
+    # Determine colormap and related important properties and flags:
+    if cmap is None:
+        try:
+            import cmocean
+            cmap = cmocean.cm.balance
+        except ImportError:
+            _log.info('cmocean.balance not found, fallback to rRdBu!')
+            cmap = plt.get_cmap('RdBu_r')  # '_r' for reverse!
+    elif isinstance(cmap, str):  # make sure we have a Colormap object (and not a string):
+        cmap = plt.get_cmap(cmap)
+    if cmap.name.replace('_r', '') in DIVERGING_CMAPS:  # 'replace' also matches reverted cmaps!
+        kwargs.setdefault('norm', TwoSlopeNorm(0))  # Diverging colormap should have zero at the symmetry point!
+    # Set extent in data coordinates (left, right, bottom, top) to kwargs (if not set explicitely):
+    dim_v, dim_u = squeezed_field.dim
+    s_v, s_u = squeezed_field.scale
+    kwargs.setdefault('extent', (0, dim_u * s_u, 0, dim_v * s_v))
+    # Plot with the empyre style context:
+    with use_style('empyre-image'):  # Only works on axes created WITHIN context!
+        if axis is None:  # If no axis is set, find the current or create a new one:
+            axis = plt.gca()
+        return axis.imshow(squeezed_field, cmap=cmap, **kwargs)
+
+
+def contour(field, axis=None, **kwargs):
+    """Plot contours. Wrapper for `matplotlib.pyplot.contour`.
+
+    Parameters
+    ----------
+    field : `Field` or ndarray
+        The contour data as a `Field` or a numpy array (in the latter case, `vector=False` and `scale=1.0` are assumed).
+    axis : `matplotlib.axes.Axes` object, optional
+        The axis to which the contour should be added, by default None, which will pick the last use axis via `gca`.
+
+    Returns
+    -------
+    axis : `matplotlib.axes.Axes`
+        The plotting axis.
+
+    Notes
+    -----
+    Additional kwargs are passed to `matplotlib.pyplot.contour`.
+    Note that the y-axis of the plot is flipped in comparison to :meth:`~matplotlib.pyplot.imshow`, i.e. that the
+    origin is `'lower'` in this case instead of `'upper'`.
+    Uses the `empyre-image` stylesheet settings for plotting (and axis creation if none exists, yet).
+    Fields are squeezed before plotting, so non-2D fields work as long as their superfluous dimensions have length 1.
+
+    """
+    _log.debug('Calling contour')
+    if not isinstance(field, Field):  # Try to convert input to Field if it is not already one:
+        field = Field(data=np.asarray(field), scale=1.0, vector=False)
+    assert not field.vector, 'Can only plot scalar fields!'
+    # Get squeezed data and make sure it's 2D scalar:
+    squeezed_field = field.squeeze()
+    assert len(squeezed_field.dim) == 2, 'Cannot plot more than 2 dimensions!'
+    # Create coordinates (respecting the field scale, +0.5: pixel center!):
+    vv, uu = (np.indices(squeezed_field.dim) + 0.5) * np.asarray(squeezed_field.scale)[:, None, None]
+    # Set kwargs defaults without overriding possible user input:
+    kwargs.setdefault('levels', [0.5])
+    kwargs.setdefault('colors', 'k')
+    kwargs.setdefault('linestyles', 'dotted')
+    kwargs.setdefault('linewidths', 2)
+    # Plot with the empyre style context:
+    with use_style('empyre-image'):  # Only works on axes created WITHIN context!
+        if axis is None:  # If no axis is set, find the current or create a new one:
+            axis = plt.gca()
+            axis.set_aspect('equal')
+        return axis.contour(uu, vv, squeezed_field.data, **kwargs)
+
+
+def colorvec(field, axis=None, **kwargs):
+    """Plot an image of a 2D vector field with up to 3 components by color encoding the vector direction.
+
+    In-plane directions are encoded via hue ("color wheel"), making sure that all in-plane directions are isoluminant
+    (i.e. a greyscale image would result a homogeneously medium grey image). Out-of-plane directions are encoded via
+    brightness with upwards pointing vectors being white and downward pointing vectors being black. The length of the
+    vectors are encoded via saturation, with full saturation being fully chromatic (in-plane) or fully white/black
+    (up/down). The center of the "color sphere" desaturated in a medium gray and encodes vectors with length zero.
+
+    Parameters
+    ----------
+    field : `Field` or ndarray
+        The image data as a `Field` or a numpy array (in the latter case, `vector=True` and `scale=1.0` are assumed).
+    axis : `matplotlib.axes.Axes` object, optional
+        The axis to which the image should be added, by default None, which will pick the last use axis via `gca`.
+
+    Returns
+    -------
+    axis : `matplotlib.axes.Axes`
+        The plotting axis.
+
+    Notes
+    -----
+    Additional kwargs are passed to `matplotlib.pyplot.imshow`.
+    Note that the y-axis of the plot is flipped in comparison to :meth:`~matplotlib.pyplot.imshow`, i.e. that the
+    origin is `'lower'` in this case instead of `'upper'`.
+    Uses the `empyre-image` stylesheet settings for plotting (and axis creation if none exists, yet).
+    Fields are squeezed before plotting, so non-2D fields work as long as their superfluous dimensions have length 1.
+    Even though squeezing takes place, `colorvec` "remembers" the original orientation of the slice! This is important
+    if you want to plot a slice that should not represent the xy-plane. The colors chosen will respect the original
+    orientation of your slice, e.g. a vortex in the xz-plane will include black and white colors (up/down) if the
+    `Field` object given as the `field` parameter has `dim=(128, 1, 128)`. If you want to plot a slice of a 3D vector
+    with 3 components and make use of this functionality, make sure to not use an integer as an index, as that will
+    drop the dimension BEFORE it is passed to `colorvec`, which will have no way of knowing which dimension was dropped.
+    Instead, make sure to use a slice of length one (example with `dim=(128, 128, 128)`):
+    >>> colorvec(field[:, 15, :])  # Wrong! Shape: (128, 128), interpreted as xy-plane!
+    >>> colorvec(field[:, 15:16, :])  # Right! Shape: (128, 1, 128), passed as 3D to `colorvec`, squeezed internally!
+
+    """
+    _log.debug('Calling colorvec')
+    if not isinstance(field, Field):  # Try to convert input to Field if it is not already one:
+        field = Field(data=np.asarray(field), scale=1.0, vector=True)
+    assert field.vector, 'Can only plot vector fields!'
+    assert len(field.dim) <= 3, 'Unusable for vector fields with dimension higher than 3!'
+    # Get squeezed data and make sure it's 2D scalar:
+    squeezed_field = field.squeeze()
+    assert len(squeezed_field.dim) == 2, 'Cannot plot more than 2 dimensions!'
+    # Extract vector components (fill 3rd component with zeros if field.comp is only 2):
+    comp = squeezed_field.comp
+    x_comp = comp[0]
+    y_comp = comp[1]
+    z_comp = comp[2] if (squeezed_field.ncomp == 3) else np.zeros(squeezed_field.dim)
+    # Calculate image with color encoded directions:
+    cmap = kwargs.pop('cmap', None)
+    if cmap is None:
+        cmap = colors.cmaps.cyclic_cubehelix
+    rgb = cmap.rgb_from_vector(np.stack((x_comp, y_comp, z_comp), axis=0))
+    # Set extent in data coordinates (left, right, bottom, top) to kwargs (if not set explicitely):
+    dim_v, dim_u = squeezed_field.dim
+    s_v, s_u = squeezed_field.scale
+    kwargs.setdefault('extent', (0, dim_u * s_u, 0, dim_v * s_v))
+    # Plot with the empyre style context:
+    with use_style('empyre-image'):  # Only works on axes created WITHIN context!
+        if axis is None:  # If no axis is set, find the current or create a new one:
+            axis = plt.gca()
+        return axis.imshow(Image.fromarray(rgb), **kwargs)
+
+
+def cosine_contours(field, axis=None, gain='auto', cmap=None, **kwargs):
+    """Plots the cosine of the (amplified) field. Wrapper for `matplotlib.pyplot.imshow`.
+
+    Parameters
+    ----------
+    field : `Field` or ndarray
+        The contour data as a `Field` or a numpy array (in the latter case, `vector=False` and `scale=1.0` are assumed).
+    axis : `matplotlib.axes.Axes` object, optional
+        The axis to which the contour should be added, by default None, which will pick the last use axis via `gca`.
+    gain : float or 'auto', optional
+        Gain factor with which the `Field` is amplified before taking the cosine, by default 'auto', which calculates a
+        gain factor that would produce roughly 4 cosine contours.
+    cmap : str or `matplotlib.colors.Colormap`, optional
+        The Colormap that should be used for the display, either as a string or object, by default None, which will pick
+        `colors.cmaps['transparent_black']` that will alternate between regions with alpha=0, showing layers below and
+        black contours.
+
+    Returns
+    -------
+    axis : `matplotlib.axes.Axes`
+        The plotting axis.
+
+    Notes
+    -----
+    Additional kwargs are passed to `matplotlib.pyplot.imshow`.
+    Uses the `empyre-image` stylesheet settings for plotting (and axis creation if none exists, yet).
+    Fields are squeezed before plotting, so non-2D fields work as long as their superfluous dimensions have length 1.
+
+    """
+    _log.debug('Calling cosine_contours')
+    if not isinstance(field, Field):  # Try to convert input to Field if it is not already one:
+        field = Field(data=np.asarray(field), scale=1.0, vector=False)
+    assert not field.vector, 'Can only plot scalar fields!'
+    # Get squeezed data and make sure it's 2D scalar:
+    squeezed_field = field.squeeze()
+    assert len(squeezed_field.dim) == 2, 'Cannot plot more than 2 dimensions (Squeezing did not help)!'
+    # Determine colormap and related important properties and flags:
+    if cmap is None:
+        cmap = colors.cmaps['transparent_black']
+    # Calculate gain if 'auto' is selected:
+    if gain == 'auto':
+        gain = 4 * 2*np.pi / (squeezed_field.amp.data.max() + 1E-30)  # 4: roughly 4 contours!
+        gain = round(gain, -int(np.floor(np.log10(abs(gain)))))  # Round to last significant digit!
+        _log.info(f'Automatically calculated a gain of: {gain}')
+    # Calculate the contours:
+    contours = np.cos(gain * squeezed_field)  # Range: [-1, 1]
+    contours += 1  # Shift to positive values
+    contours /= 2  # Rescale to [0, 1]
+    # Set extent in data coordinates (left, right, bottom, top) to kwargs (if not set explicitely):
+    dim_v, dim_u = squeezed_field.dim
+    s_v, s_u = squeezed_field.scale
+    kwargs.setdefault('extent', (0, dim_u * s_u, 0, dim_v * s_v))
+    # Plot with the empyre style context:
+    with use_style('empyre-image'):  # Only works on axes created WITHIN context!
+        if axis is None:  # If no axis is set, find the current or create a new one:
+            axis = plt.gca()
+        return axis.imshow(contours, cmap=cmap, **kwargs)
+
+
+def quiver(field, axis=None, color_angles=False, cmap=None, n_bin='auto', bin_with_mask=True, **kwargs):
+    """Plot a 2D field of arrows. Wrapper for `matplotlib.pyplot.imshow`.
+
+    Parameters
+    ----------
+    field : `Field` or ndarray
+        The vector data as a `Field` or a numpy array (in the latter case, `vector=True` and `scale=1.0` are assumed).
+    axis : `matplotlib.axes.Axes` object, optional
+        The axis to which the image should be added, by default None, which will pick the last use axis via `gca`.
+    color_angles : bool, optional
+        Switch that turns on color encoding of the arrows, by default False. Encoding works the same as for the
+        `colorvec` function (see for details). If False, arrows are uniformly colored white with black border. In both
+        cases, the amplitude is encoded via the transparency of the arrow.
+    cmap : str or `matplotlib.colors.Colormap`, optional
+        The Colormap that should be used for the arrows, either as a string or object, by default None. Will only be
+        used if `color_angles=True`.
+    n_bin : float or 'auto', optional
+        Number of entries along each axis over which the average is taken, by default 'auto', which automatically
+        determines a bin size resulting in roughly 16 arrows along the largest dimension. Usually sensible to leave
+        this on to not clutter the image with too many arrows (also due to performance). Can be turned off by setting
+        `n_bin=1`. Uses the `..fields.field.Field.bin` method.
+    bin_with_mask : bool, optional
+        If True (default) and if `n_bin>1`, entries of the constructed binned `Field` that averaged over regions that
+        were outside the `..fields.field.Field.mask` will not be assigned an arrow and stay empty instead. This prevents
+        errouneous "fade-out" effects of the arrows that would occur even for homogeneous objects.
+
+    Returns
+    -------
+    quiv : Quiver instance
+        The quiver instance that was created.
+
+    Notes
+    -----
+    Additional kwargs are passed to `matplotlib.pyplot.quiver`.
+    Uses the `empyre-image` stylesheet settings for plotting (and axis creation if none exists, yet).
+    Fields are squeezed before plotting, so non-2D fields work as long as their superfluous dimensions have length 1.
+    Even though squeezing takes place, `quiver` "remembers" the original orientation of the slice and which dimensions
+    were squeezed! See `colorvec` for  more information and an example (the same principles apply here, too).
+    The transparency of the arrows denotes the 3D(!) amplitude, if you see dots in the plot, that means the amplitude
+    is not zero, but simply out of the current plane!
+
+    """
+    _log.debug('Calling quiver')
+    if not isinstance(field, Field):  # Try to convert input to Field if it is not already one:
+        field = Field(data=np.asarray(field), scale=1.0, vector=True)
+    assert field.vector, 'Can only plot vector fields!'
+    assert len(field.dim) <= 3, 'Unusable for vector fields with dimension higher than 3!'
+    if len(field.dim) < field.ncomp:
+        warnings.warn('Assignment of vector components to dimensions is ambiguous!'
+                      f'`ncomp` ({field.ncomp}) should match `len(dim)` ({len(field.dim)})!'
+                      'If you want to plot a slice of a 3D volume, make sure to use `from:to` notation!')
+    # Get squeezed data and make sure it's 2D scalar:
+    squeezed_field = field.squeeze()
+    assert len(squeezed_field.dim) == 2, 'Cannot plot more than 2 dimensions (Squeezing did not help)!'
+    # Determine binning size if necessary:
+    if n_bin == 'auto':
+        n_bin = int(np.max((1, np.max(squeezed_field.dim) / 16)))
+    # Save old limits in case binning has to use padding:
+    u_lim = squeezed_field.dim[1] * squeezed_field.scale[1]
+    v_lim = squeezed_field.dim[0] * squeezed_field.scale[0]
+    # Bin if necessary:
+    if n_bin > 1:
+        field_mask = squeezed_field.mask  # Get mask BEFORE binning!
+        squeezed_field = squeezed_field.bin(n_bin)
+        if bin_with_mask:  # Excludes regions where in and outside are binned together!
+            mask = (field_mask.bin(n_bin) == 1)
+            squeezed_field *= mask
+    # Extract normalized vector components (fill 3rd component with zeros if field.comp is only 2):
+    normalised_comp = (squeezed_field / squeezed_field.amp.data.max()).comp
+    amplitude = squeezed_field.amp.data / squeezed_field.amp.data.max()
+    x_comp = normalised_comp[0].data
+    y_comp = normalised_comp[1].data
+    z_comp = normalised_comp[2].data if (field.ncomp == 3) else np.zeros(squeezed_field.dim)
+    # Create coordinates (respecting the field scale, +0.5: pixel center!):
+    vv, uu = (np.indices(squeezed_field.dim) + 0.5) * np.asarray(squeezed_field.scale)[:, None, None]
+    # Calculate the arrow colors:
+    if color_angles:  # Color angles according to calculated RGB values (only with circular colormaps):
+        _log.debug('Encoding angles')
+        if cmap is None:
+            cmap = colors.cmaps.cyclic_cubehelix
+        rgb = cmap.rgb_from_vector(np.asarray((x_comp, y_comp, z_comp))) / 255
+        rgba = np.concatenate((rgb, amplitude[..., None]), axis=-1)
+        kwargs.setdefault('color', rgba.reshape(-1, 4))
+    else:  # Color amplitude with numeric values, according to cmap, overrides 'color':
+        _log.debug('Encoding amplitudes')
+        if cmap is None:
+            cmap = colors.cmaps['transparent_white']
+        C = amplitude  # Numeric values, used with cmap!
+    # Check which (if any) indices were squeezed to find out which components are passed to quiver:  # TODO: TEST!!!
+    squeezed_indices = np.flatnonzero(np.asarray(field.dim) == 1)
+    if not squeezed_indices:  # Separate check, because in this case squeezed_indices == []:
+        u_comp = x_comp
+        v_comp = y_comp
+    elif squeezed_indices[0] == 0:  # Slice of the xy-plane with z squeezed:
+        u_comp = x_comp
+        v_comp = y_comp
+    elif squeezed_indices[0] == 1:  # Slice of the xz-plane with y squeezed:
+        u_comp = x_comp
+        v_comp = z_comp
+    elif squeezed_indices[0] == 2:  # Slice of the zy-plane with x squeezed:
+        u_comp = y_comp
+        v_comp = z_comp
+    # Set specific defaults for quiver kwargs:
+    kwargs.setdefault('edgecolor', colors.cmaps['transparent_black'](amplitude).reshape(-1, 4))
+    kwargs.setdefault('scale', 1/np.max(squeezed_field.scale))
+    kwargs.setdefault('width', np.max(squeezed_field.scale))
+    kwargs.setdefault('clim', (0, 1))
+    kwargs.setdefault('pivot', 'middle')
+    kwargs.setdefault('units', 'xy')
+    kwargs.setdefault('scale_units', 'xy')
+    kwargs.setdefault('minlength', 0.05)
+    kwargs.setdefault('headlength', 2)
+    kwargs.setdefault('headaxislength', 2)
+    kwargs.setdefault('headwidth', 2)
+    kwargs.setdefault('minshaft', 2)
+    kwargs.setdefault('linewidths', 1)
+    # Plot with the empyre style context:
+    with use_style('empyre-image'):  # Only works on axes created WITHIN context!
+        if axis is None:  # If no axis is set, find the current or create a new one:
+            axis = plt.gca()
+        axis.set_xlim(0, u_lim)
+        axis.set_ylim(0, v_lim)
+        axis.set_aspect('equal')
+        if color_angles:
+            return axis.quiver(uu, vv, np.asarray(u_comp), np.asarray(v_comp), cmap=cmap, **kwargs)
+        else:
+            return axis.quiver(uu, vv, np.asarray(u_comp), np.asarray(v_comp), C, cmap=cmap, **kwargs)

src/empyre/vis/plot3d.py

+# -*- coding: utf-8 -*-
+# Copyright 2020 by Forschungszentrum Juelich GmbH
+# Author: J. Caron
+#
+"""This module provides functions for 3D plots based on the `mayavi` library."""
+
+
+import logging
+
+import numpy as np
+
+from . import colors
+
+
+__all__ = ['contour3d', 'mask3d', 'quiver3d']
+
+_log = logging.getLogger(__name__)
+
+
+# TODO: Docstrings and signature!
+
+def contour3d(field, title='Field Distribution', contours=10, opacity=0.25, size=None, new_fig=True, **kwargs):
+    """Plot a field as a 3D-contour plot.
+
+    Parameters
+    ----------
+    title: string, optional
+        The title for the plot.
+    contours: int, optional
+        Number of contours which should be plotted.
+    opacity: float, optional
+        Defines the opacity of the contours. Default is 0.25.
+
+    Returns
+    -------
+    plot : :class:`mayavi.modules.vectors.Vectors`
+        The plot object.
+
+    """
+    _log.debug('Calling contour3d')
+    try:
+        from mayavi import mlab
+    except ImportError:
+        _log.error('This extension recquires the mayavi package!')
+        return
+    if size is None:
+        size = (750, 700)
+    if new_fig:
+        mlab.figure(size=size, bgcolor=(0.5, 0.5, 0.5), fgcolor=(0., 0., 0.))
+    zzz, yyy, xxx = np.indices(field.dim) + np.reshape(field.scale, (3, 1, 1, 1)) / 2  # shifted by half of scale!
+    zzz, yyy, xxx = zzz.T, yyy.T, xxx.T  # Transpose because of VTK order!
+    field_amp = field.amp.data.T  # Transpose because of VTK order!
+    if not isinstance(contours, (list, tuple, np.ndarray)):  # Calculate the contours:
+        contours = list(np.linspace(field_amp.min(), field_amp.max(), contours))
+    extent = np.ravel(list(zip((0, 0, 0), field_amp.shape)))
+    cont = mlab.contour3d(xxx, yyy, zzz, field_amp, contours=contours, opacity=opacity, **kwargs)
+    mlab.outline(cont, extent=extent)
+    mlab.axes(cont, extent=extent)
+    mlab.title(title, height=0.95, size=0.35)
+    mlab.orientation_axes()
+    cont.scene.isometric_view()
+    return cont
+
+
+def mask3d(field, title='Mask', threshold=0, grid=True, labels=True,
+           orientation=True, size=None, new_fig=True, **kwargs):
+    """Plot the mask as a 3D-contour plot.
+
+    Parameters
+    ----------
+    title: string, optional
+        The title for the plot.
+    threshold : float, optional
+        A pixel only gets masked, if it lies above this threshold . The default is 0.
+
+    Returns
+    -------
+    plot : :class:`mayavi.modules.vectors.Vectors`
+        The plot object.
+
+    """
+    _log.debug('Calling mask3d')
+    try:
+        from mayavi import mlab
+    except ImportError:
+        _log.error('This extension recquires the mayavi package!')
+        return
+    if size is None:
+        size = (750, 700)
+    if new_fig:
+        mlab.figure(size=size, bgcolor=(0.5, 0.5, 0.5), fgcolor=(0., 0., 0.))
+    zzz, yyy, xxx = np.indices(field.dim) + np.reshape(field.scale, (3, 1, 1, 1)) / 2  # shifted by half of scale!
+    zzz, yyy, xxx = zzz.T, yyy.T, xxx.T  # Transpose because of VTK order!
+    mask = field.mask.data.T.astype(int)  # Transpose because of VTK order!
+    extent = np.ravel(list(zip((0, 0, 0), mask.shape)))
+    cont = mlab.contour3d(xxx, yyy, zzz, mask, contours=[1], **kwargs)
+    if grid:
+        mlab.outline(cont, extent=extent)
+    if labels:
+        mlab.axes(cont, extent=extent)
+        mlab.title(title, height=0.95, size=0.35)
+    if orientation:
+        oa = mlab.orientation_axes()
+        oa.marker.set_viewport(0, 0, 0.4, 0.4)
+        mlab.draw()
+    engine = mlab.get_engine()
+    scene = engine.scenes[0]
+    scene.scene.isometric_view()
+    return cont
+
+
+def quiver3d(field, title='Vector Field', limit=None, cmap=None, mode='2darrow',
+             coloring='angle', ar_dens=1, opacity=1.0, grid=True, labels=True,
+             orientation=True, size=(700, 750), new_fig=True, view='isometric',
+             position=None, bgcolor=(0.5, 0.5, 0.5)):
+    """Plot the vector field as 3D-vectors in a quiverplot.
+
+    Parameters
+    ----------
+    title : string, optional
+        The title for the plot.
+    limit : float, optional
+        Plotlimit for the vector field arrow length used to scale the colormap.
+    cmap : string, optional
+        String describing the colormap which is used for color encoding (uses `~.colors.cmaps.cyclic_cubehelix` if
+        left on the `None` default) or amplitude encoding (uses 'jet' if left on the `None` default).
+    ar_dens: int, optional
+        Number defining the arrow density which is plotted. A higher ar_dens number skips more
+        arrows (a number of 2 plots every second arrow). Default is 1.
+    mode: string, optional
+        Mode, determining the glyphs used in the 3D plot. Default is '2darrow', which
+        corresponds to 2D arrows. For smaller amounts of arrows, 'arrow' (3D) is prettier.
+    coloring : {'angle', 'amplitude'}, optional
+        Color coding mode of the arrows. Use 'angle' (default) or 'amplitude'.
+    opacity: float, optional
+        Defines the opacity of the arrows. Default is 1.0 (completely opaque).
+
+    Returns
+    -------
+    plot : :class:`mayavi.modules.vectors.Vectors`
+        The plot object.
+
+    """
+    _log.debug('Calling quiver_plot3D')
+    try:
+        from mayavi import mlab
+    except ImportError:
+        _log.error('This extension recquires the mayavi package!')
+        return
+    if limit is None:
+        limit = np.max(np.nan_to_num(field.amp))
+    ad = ar_dens
+    # Create points and vector components as lists:
+    zzz, yyy, xxx = (np.indices(field.dim) + 1 / 2)
+    zzz = zzz[::ad, ::ad, ::ad].ravel()
+    yyy = yyy[::ad, ::ad, ::ad].ravel()
+    xxx = xxx[::ad, ::ad, ::ad].ravel()
+    x_mag = field.data[::ad, ::ad, ::ad, 0].ravel()
+    y_mag = field.data[::ad, ::ad, ::ad, 1].ravel()
+    z_mag = field.data[::ad, ::ad, ::ad, 2].ravel()
+    # Plot them as vectors:
+    if new_fig:
+        mlab.figure(size=size, bgcolor=(0.5, 0.5, 0.5), fgcolor=(0., 0., 0.))
+    if coloring == 'angle':  # Encodes the full angle via colorwheel and saturation:
+        _log.debug('Encoding full 3D angles')
+        vecs = mlab.quiver3d(xxx, yyy, zzz, x_mag, y_mag, z_mag, mode=mode, opacity=opacity,
+                             scalars=np.arange(len(xxx)), line_width=2)
+        vector = np.asarray((x_mag.ravel(), y_mag.ravel(), z_mag.ravel()))
+        if cmap is None:
+            cmap = colors.cmaps.cyclic_cubehelix
+        rgb = cmap.rgb_from_vector(vector)
+        rgba = np.hstack((rgb, 255 * np.ones((len(xxx), 1), dtype=np.uint8)))
+        vecs.glyph.color_mode = 'color_by_scalar'
+        vecs.module_manager.scalar_lut_manager.lut.table = rgba
+        mlab.draw()
+    elif coloring == 'amplitude':  # Encodes the amplitude of the arrows with the jet colormap:
+        _log.debug('Encoding amplitude')
+        if cmap is None:
+            cmap = 'jet'
+        vecs = mlab.quiver3d(xxx, yyy, zzz, x_mag, y_mag, z_mag,
+                             mode=mode, colormap=cmap, opacity=opacity, line_width=2)
+        mlab.colorbar(label_fmt='%.2f')
+        mlab.colorbar(orientation='vertical')
+    else:
+        raise AttributeError('Coloring mode not supported!')
+    vecs.glyph.glyph_source.glyph_position = 'center'
+    vecs.module_manager.vector_lut_manager.data_range = np.array([0, limit])
+    extent = np.ravel(list(zip((0, 0, 0), (field.dim[2], field.dim[1], field.dim[0]))))
+    if grid:
+        mlab.outline(vecs, extent=extent)
+    if labels:
+        mlab.axes(vecs, extent=extent)
+        mlab.title(title, height=0.95, size=0.35)
+    if orientation:
+        oa = mlab.orientation_axes()
+        oa.marker.set_viewport(0, 0, 0.4, 0.4)
+        mlab.draw()
+    engine = mlab.get_engine()
+    scene = engine.scenes[0]
+    if view == 'isometric':
+        scene.scene.isometric_view()
+    elif view == 'x_plus_view':
+        scene.scene.x_plus_view()
+    elif view == 'y_plus_view':
+        scene.scene.y_plus_view()
+    if position:
+        scene.scene.camera.position = position
+    return vecs

src/empyre/vis/tools.py

+# -*- coding: utf-8 -*-
+# Copyright 2020 by Forschungszentrum Juelich GmbH
+# Author: J. Caron
+#
+"""This module provides helper functions to the vis module."""
+
+
+import os
+import glob
+import shutil
+import logging
+from numbers import Number
+from contextlib import contextmanager
+
+import numpy as np
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+
+from ..fields.field import Field
+
+
+__all__ = ['new', 'savefig', 'calc_figsize', 'use_style', 'copy_mpl_stylesheets']
+_log = logging.getLogger(__name__)
+
+
+def new(nrows=1, ncols=1, mode='image', figsize=None, textwidth=None, width_scale=1.0, aspect=None, **kwargs):
+    R"""Convenience function for the creation of a new subplot grid (wraps `~matplotlib.pyplot.subplots`).
+
+    If you use the `textwidth` parameter, plot sizes are fitting into publications with LaTeX. Requires two stylesheets
+    `empyre-image` and `empyre-plot` corresponding to its two `mode` settings. Those stylesheets use
+    `constrained_layout=True` to achieve well behaving plots without much whitespace around. This function should work
+    fine for a small number of images (e.g. 1, 2x2, etc.), for more fine grained control, the contexts can be used
+    directly if they are installed corretly, or use width_scale to build the images separately (e.g. 2 adjacent with
+    width=0.5). For images, it is assumed that most images are square (and therefore `aspect=1`).
+
+    Parameters
+    ----------
+    nrows : int, optional
+        Number of rows of the subplot grid, by default 1
+    ncols : int, optional
+        Number of columns of the subplot grid, by default 1
+    mode : {'image', 'plot'}, optional
+        Mode of the new subplot grid, by default 'image'. Both modes have dedicated matplotlib styles which are used
+        and which are installed together with EMPyRe. The 'image' mode disables axis labels and ticks, mainly intended
+        to be used with `~matplotlib.pyplot.imshow` with `~empyre.vis.decorators.scalebar`, while the 'plot'
+        mode should be used for traditional plots like with `~matplotlib.pyplot.plot` or `~matplotlib.pyplot.scatter`.
+    figsize : (float, float), optional
+        Width and height of the figure in inches, defaults to rcParams["figure.figsize"], which depends on the chosen
+        stylesheet. If set, this will overwrite all other following parameters.
+    textwidth : float, optional
+        The textwidth of your LaTeX document in points, which you can get by using :math:`\the\textwidth`. If this is
+        not None (the default), this will be used to define the figure size if it is not set explicitely.
+    width_scale : float, optional
+        Only meaningful if `textwidth` is set. If it is, `width_scale` will be a scaling factor for the figure width.
+        Example: if you set this to 0.5, your figure will span half of the textwidth. Default is 1.
+    aspect : float, optional
+        Aspect ratio of the figure height relative to the figure width. If None (default), the aspect is set to be 1
+        for `mode=image` and to 'golden' for `mode=plot`, which adjusts the aspect to represent the golden ratio of
+        0.6180... If `ncols!=nrows`, it often makes sense to use `aspect=nrows/ncols` here.
+
+    Returns
+    -------
+    fig : :class:`~matplotlib.figure.Figure`
+        The constructed figure.
+    axes : axes.Axes object or array of Axes objects.
+        axes can be either a single Axes object or an array of Axes objects if more than one subplot was created.
+        The dimensions of the resulting array can be controlled with the squeeze keyword argument.
+
+    Notes
+    -----
+    additional kwargs are passed to `~matplotlib.pyplot.subplots`.
+
+    """
+    _log.debug('Calling new')
+    assert mode in ('image', 'plot'), "mode has to be 'image', or 'plot'!"
+    with use_style(f'empyre-{mode}'):
+        if figsize is None:
+            if aspect is None:
+                aspect = 'golden' if mode == 'plot' else 1  # Both image modes have 'same' as default'!
+            elif isinstance(aspect, Field):
+                dim_uv = [d for d in aspect.dim if d != 1]
+                assert len(dim_uv) == 2, f"Couldn't find field aspect ({len(dim_uv)} squeezed dimensions, has to be 2)!"
+                aspect = dim_uv[0]/dim_uv[1]  # height/width
+            else:
+                assert isinstance(aspect, Number), 'aspect has to be None, a number or field instance squeezable to 2D!'
+            figsize = calc_figsize(textwidth=textwidth, width_scale=width_scale, aspect=aspect)
+        return plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize, **kwargs)
+
+
+def savefig(fname, **kwargs):
+    """Utility wrapper around :func:`~matplotlib.pyplot.savefig` to save the current figure.
+
+    Parameters
+    ----------
+    fname : str or PathLike or file-like object
+        Path to the file wherein the figure should be saved.
+
+    Notes
+    -----
+    Uses the 'empyre-save' stylesheet (installed together with EMPyRe to control the saving behaviour. Any kwargs are
+    passed to :func:`~matplotlib.pyplot.savefig`.
+
+    """
+    _log.debug('Calling savefig')
+    with use_style('empyre-save'):
+        plt.savefig(fname, **kwargs)
+
+
+def calc_figsize(textwidth=None, width_scale=1.0, aspect=1):
+    R"""Helper function to calculate the figure size from various parameters. Useful for publications via LaTeX.
+
+    Parameters
+    ----------
+    textwidth : float, optional
+        The textwidth of your LaTeX document in points, which you can get by using :math:`\the\textwidth`. If this is
+        None (default), the standard width in inches from the current stylesheet is used.
+    width_scale : float, optional
+        Scaling factor for the figure width. Example: if you set this to 0.5, your figure will span half of the
+        textwidth. Default is 1.
+    aspect : float, optional
+        Aspect ratio of the figure height relative to the figure width. If None (default), the aspect is set to be 1
+        for `mode=image` and to 'golden' for `mode=plot`, which adjusts the aspect to represent the golden ratio of
+        0.6180...
+
+    Returns
+    -------
+    figsize: (float, float)
+        The determined figure size
+
+    Notes
+    -----
+    Based on snippet by Florian Winkler.
+
+    """
+    _log.debug('Calling calc_figsize')
+    GOLDEN_RATIO = (1 + np.sqrt(5)) / 2   # Aesthetic ratio!
+    INCHES_PER_POINT = 1.0 / 72.27  # Convert points to inch, LaTeX constant, apparently...
+    if textwidth is not None:
+        textwidth_in = textwidth * INCHES_PER_POINT  # Width of the text in inches
+    else:  # If textwidth is not given, use the default from rcParams:
+        textwidth_in = mpl.rcParams["figure.figsize"][0]
+    fig_width = textwidth_in * width_scale  # Width in inches
+    if aspect == 'golden':
+        fig_height = fig_width / GOLDEN_RATIO
+    elif isinstance(aspect, Number):
+        fig_height = textwidth_in * aspect
+    else:
+        raise ValueError(f"aspect has to be either a number, or 'golden'! Was {aspect}!")
+    fig_size = [fig_width, fig_height]  # Both in inches
+    return fig_size
+
+
+@contextmanager
+def use_style(stylename):
+    """Context that uses a matplotlib stylesheet. Can fall back to local mpl stylesheets if necessary!
+
+    Parameters
+    ----------
+    stylename : str
+        A style specification.
+
+    Yields
+    -------
+    context
+        Context manager for using style settings temporarily.
+
+    """
+    try:  # Try to load the style directly (works if it is installed somewhere mpl looks for it):
+        with plt.style.context(stylename) as context:
+            yield context
+    except OSError:  # Stylesheet not found, use local ones:
+        mplstyle_path = os.path.join(os.path.dirname(__file__), 'mplstyles', f'{stylename}.mplstyle')
+        with plt.style.context(mplstyle_path) as context:
+            yield context
+
+
+def copy_mpl_stylesheets():
+    """Copy matplotlib styles to the users matplotlib config directory. Useful if you want to utilize them elsewhere.
+
+    Notes
+    -----
+    You might need to restart your Python session for the stylesheets to be recognized/found!
+
+    """
+    # Find matplotlib styles:
+    user_stylelib_path = os.path.join(mpl.get_configdir(), 'stylelib')
+    vis_dir = os.path.dirname(__file__)
+    style_files = glob.glob(os.path.join(vis_dir, 'mplstyles', '*.mplstyle'))
+    # Copy them to the local matplotlib styles folder:
+    if not os.path.exists(user_stylelib_path):
+        os.makedirs(user_stylelib_path)
+    for style_path in style_files:
+        _, fname = os.path.split(style_path)
+        dest = os.path.join(user_stylelib_path, fname)
+        shutil.copy(style_path, dest)

tests/conftest.py

+import os
+
+import pytest
+import numpy as np
+
+from empyre.fields import Field
+
+
+@pytest.fixture
+def fielddata_path():
+    return os.path.join(os.path.dirname(os.path.realpath(__file__)), 'test_fielddata')
+
+
+@pytest.fixture
+def vector_data():
+    magnitude = np.zeros((4, 4, 4, 3))
+    magnitude[1:-1, 1:-1, 1:-1] = 1
+    return Field(magnitude, 10.0, vector=True)
+
+
+@pytest.fixture
+def vector_data_asymm():
+    shape = (5, 7, 11, 3)
+    data = np.linspace(0, 1, np.prod(shape))
+    return Field(data.reshape(shape), 10.0, vector=True)
+
+
+@pytest.fixture
+def vector_data_asymm_2d():
+    shape = (5, 7, 2)
+    data = np.linspace(0, 1, np.prod(shape))
+    return Field(data.reshape(shape), 10.0, vector=True)
+
+
+@pytest.fixture
+def vector_data_asymmcube():
+    shape = (3, 3, 3, 3)
+    data = np.linspace(0, 1, np.prod(shape))
+    return Field(data.reshape(shape), 10.0, vector=True)
+
+
+@pytest.fixture
+def scalar_data():
+    magnitude = np.zeros((4, 4, 4))
+    magnitude[1:-1, 1:-1, 1:-1] = 1
+    return Field(magnitude, 10.0, vector=False)
+
+
+@pytest.fixture
+def scalar_data_asymm():
+    shape = (5, 7, 2)
+    data = np.linspace(0, 1, np.prod(shape))
+    return Field(data.reshape(shape), 10.0, vector=False)

tests/fields/test_fielddata.py

+# -*- coding: utf-8 -*-
+"""Testcase for the magdata module."""
+
+import pytest
+
+from numbers import Number
+
+import numpy as np
+import numpy.testing
+
+from empyre.fields import Field
+
+from utils import assert_allclose
+
+
+def test_copy(vector_data):
+    vector_data = vector_data.copy()
+    # Make sure it is a new object
+    assert vector_data != vector_data, 'Unexpected behaviour in copy()!'
+    assert np.allclose(vector_data, vector_data)
+
+
+def test_bin(vector_data):
+    binned_data = vector_data.bin(2)
+    reference = 1 / 8. * np.ones((2, 2, 2, 3))
+    assert_allclose(binned_data, reference,
+                    err_msg='Unexpected behavior in scale_down()!')
+    assert_allclose(binned_data.scale, (20, 20, 20),
+                    err_msg='Unexpected behavior in scale_down()!')
+
+
+def test_zoom(vector_data):
+    zoomed_test = vector_data.zoom(2, order=0)
+    reference = np.zeros((8, 8, 8, 3))
+    reference[2:6, 2:6, 2:6] = 1
+    assert_allclose(zoomed_test, reference,
+                    err_msg='Unexpected behavior in zoom()!')
+    assert_allclose(zoomed_test.scale, (5, 5, 5),
+                    err_msg='Unexpected behavior in zoom()!')
+
+
+@pytest.mark.parametrize(
+    'mode', [
+        'constant',
+        'edge',
+        'wrap'
+    ]
+)
+@pytest.mark.parametrize(
+    'pad_width,np_pad', [
+        (1, ((1, 1), (1, 1), (1, 1), (0, 0))),
+        ((1, 2, 3), ((1, 1), (2, 2), (3, 3), (0, 0))),
+        (((1, 2), (3, 4), (5, 6)), ((1, 2), (3, 4), (5, 6), (0, 0)))
+    ]
+)
+def test_pad(vector_data, mode, pad_width, np_pad):
+    magdata_test = vector_data.pad(pad_width, mode=mode)
+    reference = np.pad(vector_data, np_pad, mode=mode)
+    assert_allclose(magdata_test, reference,
+                    err_msg='Unexpected behavior in pad()!')
+
+
+@pytest.mark.parametrize(
+    'axis', [-1, 3]
+)
+def test_component_reduction(vector_data, axis):
+    # axis=-1 is supposed to reduce over the component dimension, if it exists. axis=3 should do the same here!
+    res = np.sum(vector_data, axis=axis)
+    ref = np.zeros((4, 4, 4))
+    ref[1:-1, 1:-1, 1:-1] = 3
+    assert res.shape == ref.shape, 'Shape mismatch!'
+    assert_allclose(res, ref, err_msg="Unexpected behavior of axis keyword")
+    assert isinstance(res, Field), 'Result is not a Field object!'
+    assert not res.vector, 'Result is a vector field, but should be reduced to a scalar!'
+
+
+@pytest.mark.parametrize(
+    'axis', [(0, 1, 2), (2, 1, 0), None, (-4, -3, -2)]
+)
+def test_full_reduction(vector_data, axis):
+    res = np.sum(vector_data, axis=axis)
+    ref = np.zeros((3,))
+    ref[:] = 8
+    assert res.shape == ref.shape
+    assert_allclose(res, ref, err_msg="Unexpected behavior of full or default reduction")
+    assert isinstance(res, np.ndarray)
+
+
+@pytest.mark.parametrize(
+    'axis', [-1, 2]
+)
+def test_last_reduction_scalar(scalar_data, axis):
+    # axis=-1 is supposed to reduce over the component dimension if it exists.
+    # In this case it doesn't!
+    res = np.sum(scalar_data, axis=axis)
+    ref = np.zeros((4, 4))
+    ref[1:-1, 1:-1] = 2
+    assert res.shape == ref.shape
+    assert_allclose(res, ref, err_msg="Unexpected behavior of axis keyword")
+    assert isinstance(res, Field)
+    assert not res.vector
+
+
+@pytest.mark.parametrize(
+    'axis', [(0, 1, 2), (2, 1, 0), None, (-1, -2, -3)]
+)
+def test_full_reduction_scalar(scalar_data, axis):
+    res = np.sum(scalar_data, axis=axis)
+    ref = 8
+    assert res.shape == ()
+    assert_allclose(res, ref, err_msg="Unexpected behavior of full or default reduction")
+    assert isinstance(res, Number)
+
+
+def test_binary_operator_vector_number(vector_data):
+    res = vector_data + 1
+    ref = np.ones((4, 4, 4, 3))
+    ref[1:-1, 1:-1, 1:-1] = 2
+    assert res.shape == ref.shape
+    assert_allclose(res, ref, err_msg="Unexpected behavior of addition")
+    assert isinstance(res, Field)
+    assert res.vector
+
+
+def test_binary_operator_vector_scalar(vector_data, scalar_data):
+    res = vector_data + scalar_data
+    ref = np.zeros((4, 4, 4, 3))
+    ref[1:-1, 1:-1, 1:-1] = 2
+    assert res.shape == ref.shape
+    assert_allclose(res, ref, err_msg="Unexpected behavior of addition")
+    assert isinstance(res, Field)
+    assert res.vector
+
+
+def test_binary_operator_vector_vector(vector_data):
+    res = vector_data + vector_data
+    ref = np.zeros((4, 4, 4, 3))
+    ref[1:-1, 1:-1, 1:-1] = 2
+    assert res.shape == ref.shape
+    assert_allclose(res, ref, err_msg="Unexpected behavior of addition")
+    assert isinstance(res, Field)
+    assert res.vector
+
+
+@pytest.mark.xfail
+def test_binary_operator_vector_broadcast(vector_data):
+    # Broadcasting between vector fields is currently not implemented
+    second = np.zeros((4, 4, 3))
+    second[1:-1, 1:-1] = 1
+    second = Field(second, 10.0, vector=True)
+    res = vector_data + second
+    ref = np.zeros((4, 4, 4, 3))
+    ref[1:-1, 1:-1, 1:-1] = 1
+    ref[:, 1:-1, 1:-1] += 1
+    assert res.shape == ref.shape
+    assert_allclose(res, ref, err_msg="Unexpected behavior of addition")
+    assert isinstance(res, Field)
+    assert res.vector
+
+
+def test_mask(vector_data):
+    mask = vector_data.mask
+    reference = np.zeros((4, 4, 4))
+    reference[1:-1, 1:-1, 1:-1] = True
+    assert_allclose(mask, reference,
+                    err_msg='Unexpected behavior in mask attribute!')
+
+
+def test_get_vector(vector_data):
+    mask = vector_data.mask
+    vector = vector_data.get_vector(mask)
+    reference = np.ones(np.sum(mask) * 3)
+    assert_allclose(vector, reference,
+                    err_msg='Unexpected behavior in get_vector()!')
+
+
+def test_set_vector(vector_data):
+    mask = vector_data.mask
+    vector = 2 * np.ones(np.sum(mask) * 3)
+    vector_data.set_vector(vector, mask)
+    reference = np.zeros((4, 4, 4, 3))
+    reference[1:-1, 1:-1, 1:-1] = 2
+    assert_allclose(vector_data, reference,
+                    err_msg='Unexpected behavior in set_vector()!')
+
+
+def test_flip(vector_data_asymm):
+    field_flipz = vector_data_asymm.flip(0)
+    field_flipy = vector_data_asymm.flip(1)
+    field_flipx = vector_data_asymm.flip(2)
+    field_flipxy = vector_data_asymm.flip((1, 2))
+    field_flipdefault = vector_data_asymm.flip()
+    field_flipcomp = vector_data_asymm.flip(-1)
+    assert_allclose(np.flip(vector_data_asymm.data, axis=0) * [1, 1, -1], field_flipz.data,
+                    err_msg='Unexpected behavior in flip()! (z)')
+    assert_allclose(np.flip(vector_data_asymm.data, axis=1) * [1, -1, 1], field_flipy.data,
+                    err_msg='Unexpected behavior in flip()! (y)')
+    assert_allclose(np.flip(vector_data_asymm.data, axis=2) * [-1, 1, 1], field_flipx.data,
+                    err_msg='Unexpected behavior in flip()! (x)')
+    assert_allclose(np.flip(vector_data_asymm.data, axis=(1, 2)) * [-1, -1, 1], field_flipxy.data,
+                    err_msg='Unexpected behavior in flip()! (xy)')
+    assert_allclose(np.flip(vector_data_asymm.data, axis=(0, 1, 2)) * [-1, -1, -1], field_flipdefault.data,
+                    err_msg='Unexpected behavior in flip()! (default)')
+    assert_allclose(np.flip(vector_data_asymm.data, axis=-1) * [1, 1, 1], field_flipcomp.data,
+                    err_msg='Unexpected behavior in flip()! (components)')
+
+
+def test_unknown_num_of_components():
+    shape = (5, 7, 7)
+    data = np.linspace(0, 1, np.prod(shape))
+    with pytest.raises(AssertionError):
+        Field(data.reshape(shape), 10.0, vector=True)
+
+
+def test_repr(vector_data_asymm):
+    string_repr = repr(vector_data_asymm)
+    data_str = str(vector_data_asymm.data)
+    string_ref = f'Field(data={data_str}, scale=(10.0, 10.0, 10.0), vector=True)'
+    print(f'reference: {string_ref}')
+    print(f'repr output: {string_repr}')
+    assert string_repr == string_ref, 'Unexpected behavior in __repr__()!'
+
+
+def test_str(vector_data_asymm):
+    string_str = str(vector_data_asymm)
+    string_ref = 'Field(dim=(5, 7, 11), scale=(10.0, 10.0, 10.0), vector=True, ncomp=3)'
+    print(f'reference: {string_str}')
+    print(f'str output: {string_str}')
+    assert string_str == string_ref, 'Unexpected behavior in __str__()!'
+
+
+@pytest.mark.parametrize(
+    "index,t,scale", [
+        ((0, 1, 2), tuple, None),
+        ((0, ), Field, (2., 3.)),
+        (0, Field, (2., 3.)),
+        ((0, 1, 2, 0), float, None),
+        ((0, 1, 2, 0), float, None),
+        ((..., 0), Field, (1., 2., 3.)),
+        ((0, slice(1, 3), 2), Field, (2.,)),
+    ]
+)
+def test_getitem(vector_data, index, t, scale):
+    vector_data.scale = (1., 2., 3.)
+    data_index = index
+    res = vector_data[index]
+    assert_allclose(res, vector_data.data[data_index])
+    assert isinstance(res, t)
+    if t is Field:
+        assert res.scale == scale
+
+
+def test_from_scalar_field(scalar_data):
+    sca_x, sca_y, sca_z = [i * scalar_data for i in range(1, 4)]
+    field_comb = Field.from_scalar_fields([sca_x, sca_y, sca_z])
+    assert field_comb.vector
+    assert field_comb.scale == scalar_data.scale
+    assert_allclose(sca_x, field_comb.comp[0])
+    assert_allclose(sca_y, field_comb.comp[1])
+    assert_allclose(sca_z, field_comb.comp[2])
+
+
+def test_squeeze():
+    magnitude = np.zeros((4, 1, 4, 3))
+    field = Field(magnitude, (1., 2., 3.), vector=True)
+    sq = field.squeeze()
+    assert sq.shape == (4, 4, 3)
+    assert sq.dim == (4, 4)
+    assert sq.scale == (1., 3.)
+
+
+def test_gradient():
+    pass
+
+
+def test_gradient_1d():
+    pass
+
+
+def test_curl():
+    pass
+
+
+def test_curl_2d():
+    pass
+
+
+def test_clip_scalar_noop():
+    shape = (3, 3, 3)
+    data = np.linspace(-2, 1, np.prod(shape)).reshape(shape)
+    field = Field(data, (1., 2., 3.), vector=False)
+    assert_allclose(field, field.clip())
+
+
+def test_clip_scalar_minmax():
+    shape = (3, 3, 3)
+    data = np.linspace(-2, 1, np.prod(shape)).reshape(shape)
+    field = Field(data, (1., 2., 3.), vector=False)
+    assert_allclose(np.clip(data, -1, 0.1), field.clip(vmin=-1, vmax=0.1))
+
+
+def test_clip_scalar_sigma():
+    shape = (3, 3, 3)
+    data = np.linspace(-2, 1, np.prod(shape)).reshape(shape)
+    data[0, 0, 0] = 1e6
+    field = Field(data, (1., 2., 3.), vector=False)
+    # We clip off the one outlier
+    assert_allclose(np.clip(data, -2, 1), field.clip(sigma=5))
+    assert field.clip(sigma=5)[0, 0, 0] == 1
+
+
+def test_clip_scalar_mask():
+    shape = (3, 3, 3)
+    data = np.linspace(-2, 1, np.prod(shape)).reshape(shape)
+    mask = np.zeros(shape, dtype=bool)
+    mask[0, 0, 0] = True
+    mask[0, 0, 1] = True
+    field = Field(data, (1., 2., 3.), vector=False)
+    assert_allclose(np.clip(data, data[0, 0, 0], data[0, 0, 1]), field.clip(mask=mask))
+
+
+def test_clip_vector_noop():
+    shape = (3, 3, 3, 3)
+    data = np.linspace(-2, 1, np.prod(shape)).reshape(shape)
+    field = Field(data, (1., 2., 3.), vector=True)
+    assert_allclose(field, field.clip())
+
+
+def test_clip_vector_max():
+    shape = (3, 3, 3, 3)
+    data = np.linspace(-2, 1, np.prod(shape)).reshape(shape)
+    field = Field(data, (1., 2., 3.), vector=True)
+    res = field.clip(vmax=0.1)
+    assert_allclose(np.max(res.amp), 0.1)
+
+
+def test_clip_vector_sigma():
+    shape = (3, 3, 3, 3)
+    data = np.linspace(-2, 1, np.prod(shape)).reshape(shape)
+    data[0, 0, 0] = (1e6, 1e6, 1e6)
+    field = Field(data, (1., 2., 3.), vector=True)
+    # We clip off the one outlier
+    res = field.clip(sigma=5)
+    assert np.max(res.amp) < 1e3
+
+
+# TODO: HyperSpy would need to be installed for the following tests (slow...):
+# def test_from_signal()
+#     raise NotImplementedError()
+#
+# def test_to_signal()
+#     raise NotImplementedError()

tests/fields/test_rotate.py

+import pytest
+
+from utils import assert_allclose
+
+from empyre.fields import Field
+import numpy as np
+
+
+@pytest.mark.parametrize(
+    'axis', ['x', 'y', 'z']
+)
+def test_rot90_360(vector_data_asymm, axis):
+    assert_allclose(vector_data_asymm.rot90(axis=axis).rot90(axis=axis).rot90(axis=axis).rot90(axis=axis),
+                    vector_data_asymm,
+                    err_msg=f'Unexpected behavior in rot90()! {axis}')
+
+
+@pytest.mark.parametrize(
+    'rot_axis,flip_axes', [
+        ('x', (0, 1)),
+        ('y', (0, 2)),
+        ('z', (1, 2))
+    ]
+)
+def test_rot90_180(vector_data_asymm, rot_axis, flip_axes):
+    res = vector_data_asymm.rot90(axis=rot_axis).rot90(axis=rot_axis)
+    ref = vector_data_asymm.flip(axis=flip_axes)
+    assert_allclose(res, ref, err_msg=f'Unexpected behavior in rot90()! {rot_axis}')
+
+
+@pytest.mark.parametrize(
+    'rot_axis', [
+        'x',
+        'y',
+        'z',
+    ]
+)
+def test_rotate_compare_rot90_1(vector_data_asymmcube, rot_axis):
+    res = vector_data_asymmcube.rotate(angle=90, axis=rot_axis)
+    ref = vector_data_asymmcube.rot90(axis=rot_axis)
+    print("input", vector_data_asymmcube.data)
+    print("ref", res.data)
+    print("res", ref.data)
+    assert_allclose(res, ref, err_msg=f'Unexpected behavior in rotate()! {rot_axis}')
+
+
+def test_rot90_manual():
+    data = np.zeros((3, 3, 3, 3))
+    diag = np.array((1, 1, 1))
+    diag_unity = diag / np.sqrt(np.sum(diag**2))
+    data[0, 0, 0] = diag_unity
+    data = Field(data, 10, vector=True)
+    print("data", data.data)
+
+    rot90_x = np.zeros((3, 3, 3, 3))
+    # Axis order z, y, x; vector components x, y, z
+    rot90_x[0, 2, 0] = diag_unity * (1, -1, 1)
+    rot90_x = Field(rot90_x, 10, vector=True)
+    print("rot90_x", rot90_x.data)
+    print("data rot90 x", data.rot90(axis='x').data)
+
+    rot90_y = np.zeros((3, 3, 3, 3))
+    # Axis order z, y, x; vector components x, y, z
+    rot90_y[2, 0, 0] = diag_unity * (1, 1, -1)
+    rot90_y = Field(rot90_y, 10, vector=True)
+    print("rot90_y", rot90_y.data)
+    print("data rot90 y", data.rot90(axis='y').data)
+
+    rot90_z = np.zeros((3, 3, 3, 3))
+    # Axis order z, y, x; vector components x, y, z
+    rot90_z[0, 0, 2] = diag_unity * (-1, 1, 1)
+    rot90_z = Field(rot90_z, 10, vector=True)
+    print("rot90_z", rot90_z.data)
+    print("data rot90 z", data.rot90(axis='z').data)
+
+    assert_allclose(rot90_x, data.rot90(axis='x'), err_msg='Unexpected behavior in rot90("x")!')
+    assert_allclose(rot90_y, data.rot90(axis='y'), err_msg='Unexpected behavior in rot90("y")!')
+    assert_allclose(rot90_z, data.rot90(axis='z'), err_msg='Unexpected behavior in rot90("z")!')
+
+
+def test_rot45_manual():
+    data = np.zeros((3, 3, 3, 3))
+    data[0, 0, 0] = (1, 1, 1)
+    data = Field(data, 10, vector=True)
+    print("data", data.data)
+
+    rot45_x = np.zeros((3, 3, 3, 3))
+    # Axis order z, y, x; vector components x, y, z
+    rot45_x[0, 1, 0] = (1, 0, np.sqrt(2))
+    rot45_x = Field(rot45_x, 10, vector=True)
+    print("rot45_x", rot45_x.data)
+    # Disable spline interpolation, use nearest instead
+    res_rot45_x = data.rotate(45, axis='x', order=0)
+    print("data rot45 x", res_rot45_x.data)
+
+    rot45_y = np.zeros((3, 3, 3, 3))
+    # Axis order z, y, x; vector components x, y, z
+    rot45_y[1, 0, 0] = (np.sqrt(2), 1, 0)
+    rot45_y = Field(rot45_y, 10, vector=True)
+    print("rot45_y", rot45_y.data)
+    # Disable spline interpolation, use nearest instead
+    res_rot45_y = data.rotate(45, axis='y', order=0)
+    print("data rot45 y", res_rot45_y.data)
+
+    rot45_z = np.zeros((3, 3, 3, 3))
+    # Axis order z, y, x; vector components x, y, z
+    rot45_z[0, 0, 1] = (0, np.sqrt(2), 1)
+    rot45_z = Field(rot45_z, 10, vector=True)
+    print("rot45_z", rot45_z.data)
+    # Disable spline interpolation, use nearest instead
+    res_rot45_z = data.rotate(45, axis='z', order=0)
+    print("data rot45 z", res_rot45_z.data)
+
+    assert_allclose(rot45_x, res_rot45_x, err_msg='Unexpected behavior in rotate(45, "x")!')
+    assert_allclose(rot45_y, res_rot45_y, err_msg='Unexpected behavior in rotate(45, "y")!')
+    assert_allclose(rot45_z, res_rot45_z, err_msg='Unexpected behavior in rotate(45, "z")!')
+
+
+def test_rot90_2d_360(vector_data_asymm_2d):
+    assert_allclose(vector_data_asymm_2d.rot90().rot90().rot90().rot90(), vector_data_asymm_2d,
+                    err_msg='Unexpected behavior in 2D rot90()!')
+
+
+def test_rot90_2d_180(vector_data_asymm_2d):
+    res = vector_data_asymm_2d.rot90().rot90()
+    ref = vector_data_asymm_2d.flip()
+    assert_allclose(res, ref, err_msg='Unexpected behavior in 2D rot90()!')
+
+
+@pytest.mark.parametrize(
+    'k', [0, 1, 2, 3, 4]
+)
+def test_rot90_comp_2d_with_3d(vector_data_asymm_2d, k):
+    data_x, data_y = [comp.data[np.newaxis, :, :] for comp in vector_data_asymm_2d.comp]
+    data_z = np.zeros_like(data_x)
+    data_3d = np.stack([data_x, data_y, data_z], axis=-1)
+    vector_data_asymm_3d = Field(data_3d, scale=10, vector=True)
+    print(f'2D shape, scale: {vector_data_asymm_2d.shape, vector_data_asymm_2d.scale}')
+    print(f'3D shape, scale: {vector_data_asymm_3d.shape, vector_data_asymm_3d.scale}')
+    vector_data_rot_2d = vector_data_asymm_2d.rot90(k=k)
+    vector_data_rot_3d = vector_data_asymm_3d.rot90(k=k, axis='z')
+    print(f'2D shape after rot: {vector_data_rot_2d.shape}')
+    print(f'3D shape after rot: {vector_data_rot_3d.shape}')
+    assert_allclose(vector_data_rot_2d, vector_data_rot_3d[0, :, :, :2], err_msg='Unexpected behavior in 2D rot90()!')
+
+
+@pytest.mark.parametrize(
+    'angle', [90, 45, 23, 11.5]
+)
+def test_rotate_comp_2d_with_3d(vector_data_asymm_2d, angle):
+    data_x, data_y = [comp.data[np.newaxis, :, :] for comp in vector_data_asymm_2d.comp]
+    data_z = np.zeros_like(data_x)
+    data_3d = np.stack([data_x, data_y, data_z], axis=-1)
+    vector_data_asymm_3d = Field(data_3d, scale=10, vector=True)
+    print(f'2D shape, scale: {vector_data_asymm_2d.shape, vector_data_asymm_2d.scale}')
+    print(f'3D shape, scale: {vector_data_asymm_3d.shape, vector_data_asymm_3d.scale}')
+    r2d = vector_data_asymm_2d.rotate(angle)
+    r3d = vector_data_asymm_3d.rotate(angle, axis='z')
+    print(f'2D shape after rot: {r2d.shape}')
+    print(f'3D shape after rot: {r3d.shape}')
+    assert_allclose(r2d, r3d[0, :, :, :2], err_msg='Unexpected behavior in 2D rotate()!')
+
+
+@pytest.mark.parametrize(
+    'angle', [180, 360, 90, 45, 23, 11.5],
+)
+@pytest.mark.parametrize(
+    'axis', ['x', 'y', 'z'],
+)
+def test_rotate_scalar(vector_data_asymm, angle, axis):
+
+    data = np.zeros((1, 2, 2, 3))
+    data[0, 0, 0] = 1
+    field = Field(data, scale=10., vector=True)
+    print(field)
+    print(field.amp)
+
+    assert_allclose(
+        field.rotate(angle, axis=axis).amp,
+        field.amp.rotate(angle, axis=axis)
+    )
+
+
+@pytest.mark.parametrize(
+    'angle,order', [(180, 3), (360, 3), (90, 3), (45, 0), (23, 0), (11.5, 0)],
+)
+@pytest.mark.parametrize(
+    'axis', ['x', 'y', 'z'],
+)
+@pytest.mark.parametrize(
+    'reshape', [True, False],
+)
+def test_rotate_scalar_asymm(vector_data_asymm, angle, axis, order, reshape):
+    assert_allclose(
+        vector_data_asymm.rotate(angle, axis=axis, reshape=reshape, order=order).amp,
+        vector_data_asymm.amp.rotate(angle, axis=axis, reshape=reshape, order=order)
+    )
+
+
+@pytest.mark.parametrize(
+    'axis', ['x', 'y', 'z'],
+)
+@pytest.mark.parametrize(
+    'k', [0, 1, 2, 3, 4],
+)
+def test_rot90_scalar(vector_data_asymm, axis, k):
+    assert_allclose(
+        vector_data_asymm.amp.rot90(k=k, axis=axis),
+        vector_data_asymm.rot90(k=k, axis=axis).amp
+    )

tests/test_analytic.py

-# -*- coding: utf-8 -*-
-"""Testcase for the analytic module."""
-
-
-import os
-import unittest
-import numpy as np
-from numpy import pi
-
-import pyramid.analytic as an
-
-
-class TestCaseAnalytic(unittest.TestCase):
-
-    def setUp(self):
-        self.path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'test_analytic/')
-        self.dim = (4, 4, 4)
-        self.res = 10.0
-        self.phi = pi/4
-        self.center = (self.dim[0]/2-0.5, self.dim[1]/2-0.5, self.dim[2]/2-0.5)
-        self.radius = self.dim[2]/4
-
-    def tearDown(self):
-        self.path = None
-        self.dim = None
-        self.res = None
-        self.phi = None
-        self.center = None
-        self.radius = None
-
-    def test_phase_mag_slab(self):
-        width = (self.dim[0]/2, self.dim[1]/2, self.dim[2]/2)
-        phase = an.phase_mag_slab(self.dim, self.res, self.phi, self.center, width)
-        reference = np.load(os.path.join(self.path, 'ref_phase_slab.npy'))
-        np.testing.assert_equal(phase, reference, 'Unexpected behavior in phase_mag_slab()')
-
-    def test_phase_mag_disc(self):
-        radius = self.dim[2]/4
-        height = self.dim[2]/2
-        phase = an.phase_mag_disc(self.dim, self.res, self.phi, self.center, radius, height)
-        reference = np.load(os.path.join(self.path, 'ref_phase_disc.npy'))
-        np.testing.assert_equal(phase, reference, 'Unexpected behavior in phase_mag_disc()')
-
-    def test_phase_mag_sphere(self):
-        radius = self.dim[2]/4
-        phase = an.phase_mag_sphere(self.dim, self.res, self.phi, self.center, radius)
-        reference = np.load(os.path.join(self.path, 'ref_phase_sphere.npy'))
-        np.testing.assert_equal(phase, reference, 'Unexpected behavior in phase_mag_sphere()')
-
-    def test_phase_mag_vortex(self):
-        radius = self.dim[2]/4
-        height = self.dim[2]/2
-        phase = an.phase_mag_vortex(self.dim, self.res, self.center, radius, height)
-        reference = np.load(os.path.join(self.path, 'ref_phase_vort.npy'))
-        np.testing.assert_equal(phase, reference, 'Unexpected behavior in phase_mag_vortex()')
-
-if __name__ == '__main__':
-    suite = unittest.TestLoader().loadTestsFromTestCase(TestCaseAnalytic)
-    unittest.TextTestRunner(verbosity=2).run(suite)

tests/test_compliance.py

-# -*- coding: utf-8 -*-
-"""Testcase for the magcreator module."""
-
-
-import os
-import sys
-import datetime
-import unittest
-
-import pep8
-
-
-class TestCaseCompliance(unittest.TestCase):
-    """Class for checking compliance of pyramid."""  # TODO: Docstring
-
-    def setUp(self):
-        self.path = os.path.split(os.path.dirname(os.path.realpath(__file__)))[0]  # Pyramid dir
-
-    def tearDown(self):
-        self.path = None
-
-    def get_files_to_check(self, rootdir):
-        filepaths = []
-        for root, dirs, files in os.walk(rootdir):
-            for filename in files:
-                if ((filename.endswith('.py') or filename.endswith('.pyx'))
-                and root != os.path.join(self.path, 'scripts', 'gui')):
-                    filepaths.append(os.path.join(root, filename))
-        return filepaths
-
-    def test_pep8(self):
-        # TODO: Docstring
-        files = self.get_files_to_check(os.path.join(self.path, 'pyramid')) \
-            + self.get_files_to_check(os.path.join(self.path, 'scripts')) \
-            + self.get_files_to_check(os.path.join(self.path, 'tests'))
-        ignores = ('E226', 'E128')
-        pep8.MAX_LINE_LENGTH = 99
-        pep8style = pep8.StyleGuide(quiet=False)
-        pep8style.options.ignore = ignores
-
-        stdout_buffer = sys.stdout
-        with open(os.path.join(self.path, 'output', 'pep8_log.txt'), 'w') as sys.stdout:
-            print '<<< PEP8 LOGFILE >>>'
-            print 'RUN:', datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-            print 'IGNORED RULES:', ', '.join(ignores)
-            print 'MAX LINE LENGTH:', pep8.MAX_LINE_LENGTH
-            print '\nERRORS AND WARNINGS:'
-            result = pep8style.check_files(files)
-            if result.total_errors == 0:
-                print 'No Warnings or Errors detected!'
-            else:
-                print '\n{} Warnings and Errors detected!'.format(result.total_errors)
-        sys.stdout = stdout_buffer
-
-        error_message = 'Found %s Errors and Warnings!' % result.total_errors
-        self.assertEqual(result.total_errors, 0, error_message)
-
-
-if __name__ == '__main__':
-    suite = unittest.TestLoader().loadTestsFromTestCase(TestCaseCompliance)
-    unittest.TextTestRunner(verbosity=2).run(suite)

tests/test_holoimage.py

-# -*- coding: utf-8 -*-
-"""Testcase for the holoimage module."""
-
-
-import os
-import unittest
-import numpy as np
-from numpy import pi
-
-from pyramid.phasemap import PhaseMap
-import pyramid.holoimage as hi
-
-
-class TestCaseHoloImage(unittest.TestCase):
-
-    def setUp(self):
-        self.path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'test_holoimage')
-        phase = np.zeros((4, 4))
-        phase[1:-1, 1:-1] = pi/4
-        self.phase_map = PhaseMap(10.0, phase)
-
-    def tearDown(self):
-        self.path = None
-        self.phase_map = None
-
-    def test_holo_image(self):
-        img = hi.holo_image(self.phase_map)
-        arr = np.array(img.getdata(), np.uint8).reshape(img.size[1], img.size[0], 3)
-        holo_img_r, holo_img_g, holo_img_b = arr[..., 0], arr[..., 1], arr[..., 2]
-        ref_holo_img_r = np.loadtxt(os.path.join(self.path, 'ref_holo_img_r.txt'))
-        ref_holo_img_g = np.loadtxt(os.path.join(self.path, 'ref_holo_img_g.txt'))
-        ref_holo_img_b = np.loadtxt(os.path.join(self.path, 'ref_holo_img_b.txt'))
-        hi.display(img)
-        np.testing.assert_equal(holo_img_r, ref_holo_img_r,
-                                'Unexpected behavior in holo_image() (r-component)!')
-        np.testing.assert_equal(holo_img_g, ref_holo_img_g,
-                                'Unexpected behavior in holo_image() (g-component)!')
-        np.testing.assert_equal(holo_img_b, ref_holo_img_b,
-                                'Unexpected behavior in holo_image() (b-component)!')
-
-
-if __name__ == '__main__':
-    suite = unittest.TestLoader().loadTestsFromTestCase(TestCaseHoloImage)
-    unittest.TextTestRunner(verbosity=2).run(suite)

tests/test_holoimage/ref_holo_img_b.txt

-0.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00
-0.000000000000000000e+00 0.000000000000000000e+00 9.700000000000000000e+01 2.550000000000000000e+02
-0.000000000000000000e+00 0.000000000000000000e+00 1.000000000000000000e+02 2.550000000000000000e+02
-0.000000000000000000e+00 3.000000000000000000e+00 3.000000000000000000e+00 0.000000000000000000e+00

tests/test_holoimage/ref_holo_img_g.txt

-0.000000000000000000e+00 1.000000000000000000e+00 1.000000000000000000e+00 0.000000000000000000e+00
-2.550000000000000000e+02 9.900000000000000000e+01 0.000000000000000000e+00 0.000000000000000000e+00
-2.550000000000000000e+02 1.950000000000000000e+02 9.500000000000000000e+01 0.000000000000000000e+00
-0.000000000000000000e+00 2.520000000000000000e+02 2.520000000000000000e+02 0.000000000000000000e+00

tests/test_holoimage/ref_holo_img_r.txt

-0.000000000000000000e+00 2.550000000000000000e+02 2.550000000000000000e+02 0.000000000000000000e+00
-2.530000000000000000e+02 1.950000000000000000e+02 9.800000000000000000e+01 0.000000000000000000e+00
-2.530000000000000000e+02 9.500000000000000000e+01 0.000000000000000000e+00 0.000000000000000000e+00
-0.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00 0.000000000000000000e+00

tests/test_magcreator.py

-# -*- coding: utf-8 -*-
-"""Testcase for the magcreator module."""
-
-
-import os
-import unittest
-import numpy as np
-from numpy import pi
-
-import pyramid.magcreator as mc
-
-
-class TestCaseMagCreator(unittest.TestCase):
-
-    def setUp(self):
-        self.path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'test_magcreator')
-
-    def tearDown(self):
-        self.path = None
-
-    def test_shape_slab(self):
-        test_slab = mc.Shapes.slab((5, 6, 7), (2, 3, 4), (1, 3, 5))
-        np.testing.assert_equal(test_slab, np.load(os.path.join(self.path, 'ref_slab.npy')),
-                                'Created slab does not match expectation!')
-
-    def test_shape_disc(self):
-        test_disc_z = mc.Shapes.disc((5, 6, 7), (2, 3, 4), 2, 3, 'z')
-        test_disc_y = mc.Shapes.disc((5, 6, 7), (2, 3, 4), 2, 3, 'y')
-        test_disc_x = mc.Shapes.disc((5, 6, 7), (2, 3, 4), 2, 3, 'x')
-        np.testing.assert_equal(test_disc_z, np.load(os.path.join(self.path, 'ref_disc_z.npy')),
-                                'Created disc in z-direction does not match expectation!')
-        np.testing.assert_equal(test_disc_y, np.load(os.path.join(self.path, 'ref_disc_y.npy')),
-                                'Created disc in y-direction does not match expectation!')
-        np.testing.assert_equal(test_disc_x, np.load(os.path.join(self.path, 'ref_disc_x.npy')),
-                                'Created disc in x-direction does not match expectation!')
-
-    def test_shape_sphere(self):
-        test_sphere = mc.Shapes.sphere((5, 6, 7), (2, 3, 4), 2)
-        np.testing.assert_equal(test_sphere, np.load(os.path.join(self.path, 'ref_sphere.npy')),
-                                'Created sphere does not match expectation!')
-
-    def test_shape_filament(self):
-        test_filament_z = mc.Shapes.filament((5, 6, 7), (2, 3), 'z')
-        test_filament_y = mc.Shapes.filament((5, 6, 7), (2, 3), 'y')
-        test_filament_x = mc.Shapes.filament((5, 6, 7), (2, 3), 'x')
-        np.testing.assert_equal(test_filament_z, np.load(os.path.join(self.path, 'ref_fil_z.npy')),
-                                'Created filament in z-direction does not match expectation!')
-        np.testing.assert_equal(test_filament_y, np.load(os.path.join(self.path, 'ref_fil_y.npy')),
-                                'Created filament in y-direction does not match expectation!')
-        np.testing.assert_equal(test_filament_x, np.load(os.path.join(self.path, 'ref_fil_x.npy')),
-                                'Created filament in x-direction does not match expectation!')
-
-    def test_shape_pixel(self):
-        test_pixel = mc.Shapes.pixel((5, 6, 7), (2, 3, 4))
-        np.testing.assert_equal(test_pixel, np.load(os.path.join(self.path, 'ref_pixel.npy')),
-                                'Created pixel does not match expectation!')
-
-    def test_create_mag_dist_homog(self):
-        mag_shape = mc.Shapes.disc((1, 10, 10), (0, 4.5, 4.5), 3, 1)
-        magnitude = mc.create_mag_dist_homog(mag_shape, pi/4)
-        np.testing.assert_equal(magnitude, np.load(os.path.join(self.path, 'ref_mag_disc.npy')),
-                                'Created homog. magnetic distribution does not match expectation')
-
-    def test_create_mag_dist_vortex(self):
-        mag_shape = mc.Shapes.disc((1, 10, 10), (0, 4.5, 4.5), 3, 1)
-        magnitude = mc.create_mag_dist_vortex(mag_shape)
-        np.testing.assert_equal(magnitude, np.load(os.path.join(self.path, 'ref_mag_vort.npy')),
-                                'Created vortex magnetic distribution does not match expectation')
-
-
-if __name__ == '__main__':
-    suite = unittest.TestLoader().loadTestsFromTestCase(TestCaseMagCreator)
-    unittest.TextTestRunner(verbosity=2).run(suite)