small adjustments to plot.py, added plot_test.py

nifty5/plotting/plot.py

@@ -292,10 +292,12 @@ def _plot(f, ax, **kwargs):
 _plots = []
 _kwargs = []
 
+
 def plot(f, **kwargs):
     _plots.append(f)
     _kwargs.append(kwargs)
 
+
 def plot_finish(**kwargs):
     global _plots, _kwargs
     import matplotlib.pyplot as plt
@@ -303,14 +305,17 @@ def plot_finish(**kwargs):
     fig = plt.figure()
     if "title" in kwargs:
         plt.suptitle(kwargs.pop("title"))
-    nx = kwargs.pop("nx", 1)
-    ny = kwargs.pop("ny", 1)
+    nx = kwargs.pop("nx", int(np.ceil(np.sqrt(nplot))))
+    ny = kwargs.pop("ny", int(np.ceil(np.sqrt(nplot))))
+    if nx*ny < nplot:
+        raise ValueError('Figure dimensions not sufficient for number of plots')
     xsize = kwargs.pop("xsize", 6)
     ysize = kwargs.pop("ysize", 6)
     fig.set_size_inches(xsize, ysize)
     for i in range(nplot):
-        ax = fig.add_subplot(ny,nx,i+1)
+        ax = fig.add_subplot(ny, nx, i+1)
         _plot(_plots[i], ax, **_kwargs[i])
+    fig.tight_layout()
     _makeplot(kwargs.pop("name", None))
     _plots = []
     _kwargs = []

nifty5/plotting/plot_test.py

+import nifty5 as ift
+import numpy as np
+
+
+def plot_test():
+    rg_space1 = ift.makeDomain(ift.RGSpace((100,)))
+    rg_space2 = ift.makeDomain(ift.RGSpace((80, 80)))
+    hp_space = ift.makeDomain(ift.HPSpace(64))
+    gl_space = ift.makeDomain(ift.GLSpace(128))
+
+    fft = ift.FFTOperator(rg_space2)
+
+    field_rg1_1 = ift.Field(domain=rg_space1, val=np.random.randn(100))
+    field_rg1_2 = ift.Field(domain=rg_space1, val=np.random.randn(100))
+    field_rg2 = ift.Field(domain=rg_space2, val=np.random.randn(80 ** 2).reshape((80, 80)))
+    field_hp = ift.Field(domain=hp_space, val=np.random.randn(12*64**2))
+    field_gl = ift.Field(domain=gl_space, val=np.random.randn(32640))
+    field_ps = ift.power_analyze(fft.times(field_rg2))
+
+    ## Start various plotting tests
+
+    ift.plot(field_rg1_1, title='Single plot')
+    ift.plot_finish()
+
+    ift.plot(field_rg2, title='2d rg')
+    ift.plot([field_rg1_1, field_rg1_2], title='list 1d rg', label=['1', '2'])
+    ift.plot(field_rg1_2, title='1d rg, xmin, ymin', xmin=0.5, ymin=0., xlabel='xmin=0.5', ylabel='ymin=0')
+    ift.plot_finish(title='Three plots')
+
+    ift.plot(field_hp, title='HP planck-color', colormap='Planck-like')
+    ift.plot(field_rg1_2, title='1d rg')
+    ift.plot(field_ps)
+    ift.plot(field_gl, title='GL')
+    ift.plot(field_rg2, title='2d rg')
+    ift.plot_finish(nx=2, ny=3, title='Five plots')
+
+if __name__ == '__main__':
+    plot_test()
+