Merge branch 'mf_plot_improvement' into 'NIFTy_6'

MF Plotting: plot MF fields with arbitray domain order

See merge request !441

nifty6/plot.py

@@ -331,14 +331,26 @@ def _plot2D(f, ax, **kwargs):
 
     # check for multifrequency plotting
     have_rgb = False
+    x_space = 0
     if len(dom) == 2:
-        if (not isinstance(dom[1], RGSpace)) or len(dom[1].shape) != 1:
-            raise TypeError("need 1D RGSpace as second domain")
-        if dom[1].shape[0] == 1:
+        f_space = kwargs.pop("freq_space_idx", 1)
+        if not f_space in [0, 1]:
+            raise ValueError("Invalid frequency space index")
+        if (not isinstance(dom[f_space], RGSpace)) \
+            or len(dom[f_space].shape) != 1:
+            raise TypeError("Need 1D RGSpace as frequency space domain")
+        x_space = 1 - f_space
+
+        # Only one frequency?
+        if dom[f_space].shape[0] == 1:
             from .sugar import makeField
-            f = makeField(f.domain[0], f.val[..., 0])
+            f = makeField(f.domain[x_space],
+                          f.val.squeeze(axis=dom.axes[f_space]))
         else:
-            rgb = _rgb_data(f.val)
+            val = f.val
+            if f_space == 0:
+                val = np.moveaxis(val, 0, -1)
+            rgb = _rgb_data(val)
             have_rgb = True
 
     foo = kwargs.pop("norm", None)
@@ -350,7 +362,7 @@ def _plot2D(f, ax, **kwargs):
     ax.set_title(kwargs.pop("title", ""))
     ax.set_xlabel(kwargs.pop("xlabel", ""))
     ax.set_ylabel(kwargs.pop("ylabel", ""))
-    dom = dom[0]
+    dom = dom[x_space]
     if not have_rgb:
         cmap = kwargs.pop("colormap", plt.rcParams['image.cmap'])
 
@@ -451,6 +463,9 @@ class Plot(object):
             `PowerSpace`, `HPSpace`, `GLSpace`.
             If it is a list, all list members must be Fields defined over the
             same one-dimensional `RGSpace` or `PowerSpace`.
+
+        Optional Parameters
+        -------------------
         title: string
             Title of the plot.
         xlabel: string
@@ -467,6 +482,8 @@ class Plot(object):
             Annotation string.
         alpha: float or list of floats
             Transparency value.
+        freq_space_idx: int
+            for multi-frequency plotting: index of frequency space in domain
         """
         from .multi_field import MultiField
         if isinstance(f, MultiField):
@@ -527,4 +544,6 @@ class Plot(object):
             ax = fig.add_subplot(ny, nx, i+1)
             _plot(self._plots[i], ax, **self._kwargs[i])
         fig.tight_layout()
-        _makeplot(kwargs.pop("name", None), block=kwargs.pop("block", True), dpi=kwargs.pop("dpi", None))
+        _makeplot(kwargs.pop("name", None),
+                  block=kwargs.pop("block", True),
+                  dpi=kwargs.pop("dpi", None))

test/test_plot.py

@@ -20,6 +20,11 @@ import numpy as np
 import nifty6 as ift
 from .common import setup_function, teardown_function
 
+nr = 0
+def name():
+    global nr
+    nr += 1
+    return 'plot{}.png'.format(nr)
 
 def test_plots():
     # FIXME Write to temporary folder?
@@ -40,14 +45,14 @@ def test_plots():
 
     plot = ift.Plot()
     plot.add(field_rg1_1, title='Single plot')
-    plot.output(name='plot1.png')
+    plot.output(name=name())
 
     plot = ift.Plot()
     plot.add(field_rg2, title='2d rg')
     plot.add([field_rg1_1, field_rg1_2], title='list 1d rg', label=['1', '2'])
     plot.add(field_rg1_2, title='1d rg, xmin, ymin', xmin=0.5, ymin=0.,
              xlabel='xmin=0.5', ylabel='ymin=0')
-    plot.output(title='Three plots', name='plot2.png')
+    plot.output(title='Three plots', name=name())
 
     plot = ift.Plot()
     plot.add(field_hp, title='HP planck-color', colormap='Planck-like')
@@ -55,4 +60,20 @@ def test_plots():
     plot.add(field_ps)
     plot.add(field_gl, title='GL')
     plot.add(field_rg2, title='2d rg')
-    plot.output(nx=2, ny=3, title='Five plots', name='plot3.png')
+    plot.output(nx=2, ny=3, title='Five plots', name=name())
+
+
+def test_mf_plot():
+    x_space = ift.RGSpace((16, 16))
+    f_space = ift.RGSpace(4)
+
+    d1 = ift.DomainTuple.make([x_space, f_space])
+    d2 = ift.DomainTuple.make([f_space, x_space])
+
+    f1 = ift.from_random('normal', d1)
+    f2 = ift.makeField(d2, np.moveaxis(f1.val, -1, 0))
+
+    plot = ift.Plot()
+    plot.add(f1, block=False, title='f_space_idx = 1')
+    plot.add(f2, freq_space_idx=0, title='f_space_idx = 0')
+    plot.output(nx=2, ny=1, title='MF-Plots, should look identical', name=name())