From a020231a26e2b6aed05d114dc30fbfa1342649d6 Mon Sep 17 00:00:00 2001
From: Daniel Brown <ddb@star.sr.bham.ac.uk>
Date: Sat, 14 Dec 2013 02:40:08 +0000
Subject: [PATCH] 2d plotting examples, probably want to wrap these plots up in
 a function to make it easier in future. Fixed ordering of output rows/cols in
 the stored matrix for 2d outputs.

---
 bin/test_2d_plot.py | 20 +++++++++++++++++---
 pykat/finesse.py    |  5 ++++-
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/bin/test_2d_plot.py b/bin/test_2d_plot.py
index debb606..b178319 100644
--- a/bin/test_2d_plot.py
+++ b/bin/test_2d_plot.py
@@ -4,6 +4,7 @@ from pykat.detectors import *
 from pykat.components import *
 from pykat.commands import *
 from pykat.structs import *
+import matplotlib.image as mpimg
 from mpl_toolkits.mplot3d.axes3d import Axes3D
 
 import numpy as np
@@ -13,11 +14,12 @@ code = """
 l l1 1 0 0 n1
 s s1 10 1 n1 n2
 
+ad ad1 0 n2
 beam b1 0 n2
 maxtem 0 
 
-xaxis b1 x lin -7.0 7.0 50
-x2axis b1 y lin -7.0 7.0 50
+xaxis b1 x lin -10.0 10 100
+x2axis b1 y lin -6 6 100
 
 """
 
@@ -31,5 +33,17 @@ out = kat.run(printout=0,printerr=0)
 fig = pl.figure()
 ax = fig.add_subplot(1, 1, 1, projection='3d')
 x, y = np.meshgrid(out.x, out.y)
-p = ax.plot_surface(x, y, out.z[0,:])   
+p = ax.plot_wireframe(x, y, out.z[1])   
+pl.xlabel(out.xlabel)
+pl.ylabel(out.ylabel)
+pl.show()
+
+pl.figure()
+extent = [np.min(out.x),np.max(out.x),np.min(out.y),np.max(out.y)]
+imgplot = pl.imshow(out.z[1], extent=extent)
+#imgplot.set_interpolation('bicubic')
+imgplot.set_interpolation('nearest')
+pl.colorbar()
+pl.xlabel(out.xlabel)
+pl.ylabel(out.ylabel)
 pl.show()
diff --git a/pykat/finesse.py b/pykat/finesse.py
index 4469651..6e5b336 100644
--- a/pykat/finesse.py
+++ b/pykat/finesse.py
@@ -524,8 +524,11 @@ class kat(object):
             # written in linear form
             x = data[0::(1+self.x2axis.steps),0]
             y = data[0:(1+self.x2axis.steps),1]
+            # get rows and columns lined up so that we can reshape a single column of all x/y data
+            # into a matrix
             z = data[:,2:].transpose().reshape(data.shape[1]-2, 1+self.xaxis.steps, 1+self.x2axis.steps)
-            
+            # once you do this the data for y and x axes need swapping
+            z = z.swapaxes(1,2)
             return [x, y, z, hdr]
         else:
             shape_len = len(data.shape)
-- 
GitLab