Index: trunk/beamsgenerator.py
===================================================================
--- trunk/beamsgenerator.py	(revision 31)
+++ trunk/beamsgenerator.py	(revision 32)
@@ -3,24 +3,73 @@
 import collections
 import numpy as np
-
-
-class Axis(object):
-    """Class to hold axis information.
-    """
-
-    def __init__(self, data, title='', units=''):
-        self.data = data
-        self.title = title
-        self.units = units
-
-
-class Image(object):
-    """Class to hold 2d image, with axes, units, title, etc.
-    """
-    def __init__(self, data, axes=[], title='', units=''):
-        self.data = data
-        self.axes = axes
-        self.title = title
-        self.units = units
+import pp
+
+import common.commonobjects as co
+
+def calculate_dirty_beam(npix, rpix, mult, bmax, bxbylist, wn_max, wn):
+    lamb = 1.0 / (wn * 100.0)
+    lambda_min = 1.0 / (wn_max * 100.0)
+    maxbx = bmax / lambda_min
+
+    uv = numpy.zeros([mult*npix,mult*npix])
+
+    for bxby in bxbylist:
+        ibx = int(round(((bxby[0] / lamb) / maxbx) * mult * rpix))
+        if ibx < 0:
+            ibx = mult*npix + ibx 
+        iby = int(round(((bxby[1] / lamb) / maxbx) * mult * rpix))
+        if iby < 0:
+            iby = mult*npix + iby 
+        uv[ibx,iby] = 1.0
+ 
+    temp = numpy.fft.fftshift(uv)
+    temp = numpy.fft.fft2(temp)
+    temp = numpy.fft.fftshift(temp)
+    temp = numpy.power(temp, 2)
+    # convert to real, should be real anway as FT of symmetric function 
+    dirty_beam = numpy.abs(temp)
+    # and normalise
+    dirty_beam /= numpy.max(dirty_beam)
+
+    # truncate dirty beam pattern
+    dirty_beam = dirty_beam[
+      (mult-1)*rpix:(mult+1)*rpix, (mult-1)*rpix:(mult+1)*rpix]
+
+    return dirty_beam
+
+def calculate_primary_beam(npix, rpix, mult, mx, my, bmax, m1_diameter,
+  wn_max, wn):
+    # assume uniform illumination of primary. uv goes out to
+    # +- bmax / lambda_min, fill uv plane appropriate to primary
+    # size
+    mirror = numpy.ones([mult*npix,mult*npix])
+    lambda_min = 1.0 / (wn_max * 100.0)
+    lamb = 1.0 / (wn * 100.0)
+    mirror[(numpy.sqrt(mx*mx + my*my) / (mult*rpix)) * 
+      (bmax / lambda_min) > (0.5 * m1_diameter / lamb)] = 0.0
+
+    # fftshift mirror to be centred at [0,0],
+    # do a 2d ftt of it,
+    # fftshift transform so that 0 freq is in centre of array
+    temp = numpy.fft.fftshift(mirror)
+    temp = numpy.fft.fft2(temp)
+    temp = numpy.fft.fftshift(temp)
+    # convert to real, should be real anyway as FT of symmetric function 
+    # square to give intensity response
+    primary_amplitude_beam = numpy.real(temp)
+    primary_intensity_beam = numpy.power(numpy.abs(temp), 2)
+    # and normalise
+    primary_amplitude_beam /= numpy.max(primary_amplitude_beam)
+    primary_intensity_beam /= numpy.max(primary_intensity_beam)
+
+    # truncate primary beam pattern
+    primary_amplitude_beam = primary_amplitude_beam[
+      (mult-1)*rpix:(mult+1)*rpix,
+      (mult-1)*rpix:(mult+1)*rpix]
+    primary_intensity_beam = primary_intensity_beam[
+      (mult-1)*rpix:(mult+1)*rpix,
+      (mult-1)*rpix:(mult+1)*rpix]
+
+    return primary_intensity_beam, primary_amplitude_beam
 
 
@@ -32,4 +81,10 @@
         self.previous_results = previous_results
         self.result = collections.OrderedDict()
+
+        # start parallel python (pp), find the number of CPUS available
+        ppservers = ()
+        self.job_server = pp.Server(ppservers=ppservers)
+        print 'Sampler starting pp with %s workers' % \
+          self.job_server.get_ncpus()
 
     def run(self):
@@ -63,5 +118,4 @@
         # go out to radius of first null
         rpix = int(max_beam_radius / max_pixsize)
-#        rpix = 32
         npix = 2 * rpix
         self.result['npix'] = npix
@@ -72,4 +126,7 @@
         axis *= 206264.806247
         self.result['spatial axis [arcsec]'] = axis
+        axis1 = co.Axis(data=-axis, title='RA offset', units='arcsec') 
+        axis2 = co.Axis(data=axis, title='Dec offset', units='arcsec') 
+        axis3 = co.Axis(data=fts_wn, title='Frequency', units='cm-1')
 
         # calculate beams for each point on fts_wn
@@ -82,51 +139,30 @@
         self.result['dirty beam'] = collections.OrderedDict()
 
+        wn_max = np.max(fts_wn)
+        m1_diameter = telescope['Primary mirror diameter']
+
+        # first calculate primary beam for each wn
+        jobs = {}
         for wn in fts_wn:
-            # first calculate primary beam
-            #
-            # assume uniform illumination of primary. uv goes out to
-            # +- bmax / lambda_min, fill uv plane appropriate to primary
-            # size
-            mirror = np.ones([mult*npix,mult*npix])
-            lambda_min = 1.0 / (np.max(fts_wn) * 100.0)
-            lamb = 1.0 / (wn * 100.0)
-            mirror[(np.sqrt(mx*mx + my*my) / (mult*rpix)) * 
-              (bmax / lambda_min) > \
-              (0.5 * telescope['Primary mirror diameter'] / lamb)] = 0.0
-
-            # fftshift mirror to be centred at [0,0],
-            # do a 2d ftt of it,
-            # fftshift transform so that 0 freq is in centre of array
-            temp = np.fft.fftshift(mirror)
-            temp = np.fft.fft2(temp)
-            temp = np.fft.fftshift(temp)
-            # convert to real, should be real anyway as FT of symmetric function 
-            # square to give intensity response
-            primary_amplitude_beam = np.real(temp)
-            primary_intensity_beam = np.power(np.abs(temp), 2)
-            # and normalise
-            primary_amplitude_beam /= np.max(primary_amplitude_beam)
-            primary_intensity_beam /= np.max(primary_intensity_beam)
-
-            # truncate primary beam pattern
-            primary_amplitude_beam = primary_amplitude_beam[
-              (mult-1)*rpix:(mult+1)*rpix,
-              (mult-1)*rpix:(mult+1)*rpix]
-            primary_intensity_beam = primary_intensity_beam[
-              (mult-1)*rpix:(mult+1)*rpix,
-              (mult-1)*rpix:(mult+1)*rpix]
-
-            # axes
-            axis1 = Axis(data=-axis, title='RA offset', units='arcsec') 
-            axis2 = Axis(data=axis, title='Dec offset', units='arcsec') 
-            image = Image(data=primary_intensity_beam, axes=[axis1, axis2],
+            # submit jobs
+            indata = (npix, rpix, mult, mx, my, bmax, m1_diameter, wn_max,
+              wn,)
+            jobs[wn] = self.job_server.submit(calculate_primary_beam,
+              indata, (), ('numpy',))
+
+        primary_amplitude_beam = np.zeros([npix,npix,len(fts_wn)], np.complex)
+        for iwn,wn in enumerate(fts_wn):
+            # collect and store results
+            primary_intensity_beam, primary_amplitude_beam[:,:,iwn] = jobs[wn]()
+            image = co.Image(data=primary_intensity_beam, axes=[axis1, axis2],
               title='Primary Beam %06.4g cm-1' % wn)
             self.result['primary beam'][wn] = image
 
-            image = Image(data=primary_amplitude_beam, axes=[axis1, axis2],
-              title='Amplitude Primary Beam %06.4g cm-1' % wn)
-            self.result['primary amplitude beam'][wn] = image
-
-            # second calculate dirty beam
+        cube = co.Cube(data=primary_amplitude_beam, axes=[axis1, axis2, axis3],
+          title='Amplitude Primary Beam %06.4g cm-1' % wn)
+        self.result['primary amplitude beam'] = cube
+
+        # second calculate dirty beam, same use of pp
+        for wn in fts_wn:
             #
             # 1. uv plane covered by discrete points so should use a
@@ -140,39 +176,20 @@
             # positions. Fastest but adequate?
 
-            uv = np.zeros([mult*npix,mult*npix])
-            maxbx = bmax / np.min(fts_lambda)
-            for bxby in uvmapgen['bxby']:
-                ibx = int(round(((bxby[0] / lamb) / maxbx) * mult * rpix))
-                if ibx < 0:
-                    ibx = mult*npix + ibx 
-                iby = int(round(((bxby[1] / lamb) / maxbx) * mult * rpix))
-                if iby < 0:
-                    iby = mult*npix + iby 
-                uv[ibx,iby] = 1.0
- 
-            temp = np.fft.fftshift(uv)
-            temp = np.fft.fft2(temp)
-            temp = np.fft.fftshift(temp)
-            temp = np.power(temp, 2)
-            # convert to real, should be real anway as FT of symmetric function 
-            dirty_beam = np.abs(temp)
-            # and normalise
-            dirty_beam /= np.max(dirty_beam)
-
-            # truncate dirty beam pattern
-            dirty_beam = dirty_beam[
-              (mult-1)*rpix:(mult+1)*rpix,
-              (mult-1)*rpix:(mult+1)*rpix]
-
-            # axes
-            axis = np.arange(-rpix, rpix, dtype=np.float)
-            axis *= (np.min(fts_lambda) / bmax)
-            axis *= 206264.806247
-            axis1 = Axis(data=-axis, title='RA offset', units='arcsec') 
-            axis2 = Axis(data=axis, title='Dec offset', units='arcsec') 
-            
-            image = Image(data=dirty_beam, axes=[axis1, axis2],
+            indata = (npix, rpix, mult, bmax, uvmapgen['bxby'], wn_max, wn,)
+            jobs[wn] = self.job_server.submit(calculate_dirty_beam,
+              indata, (), ('numpy',))
+
+
+        # axes for dirty beam
+        axis = np.arange(-rpix, rpix, dtype=np.float)
+        axis *= (np.min(fts_lambda) / bmax)
+        axis *= 206264.806247
+        axis1 = co.Axis(data=-axis, title='RA offset', units='arcsec') 
+        axis2 = co.Axis(data=axis, title='Dec offset', units='arcsec') 
+
+        for wn in fts_wn:
+            dirty_beam = jobs[wn]()
+            image = co.Image(data=dirty_beam, axes=[axis1, axis2],
               title='Dirty Beam %06.4g cm-1' % wn)
-            
             self.result['dirty beam'][wn] = image
 
