Changeset 32


Ignore:
Timestamp:
May 20, 2014 3:23:10 PM (10 years ago)
Author:
JohnLightfoot
Message:

parallel processing

File:
1 edited

Legend:

Unmodified
Added
Removed

  • trunk/beamsgenerator.py

    r17 r32  
    33import collections
    44import numpy as np
    5 
    6 
    7 class Axis(object):
    8     """Class to hold axis information.
    9     """
    10 
    11     def __init__(self, data, title='', units=''):
    12         self.data = data
    13         self.title = title
    14         self.units = units
    15 
    16 
    17 class Image(object):
    18     """Class to hold 2d image, with axes, units, title, etc.
    19     """
    20     def __init__(self, data, axes=[], title='', units=''):
    21         self.data = data
    22         self.axes = axes
    23         self.title = title
    24         self.units = units
     5import pp
     6
     7import common.commonobjects as co
     8
     9def calculate_dirty_beam(npix, rpix, mult, bmax, bxbylist, wn_max, wn):
     10    lamb = 1.0 / (wn * 100.0)
     11    lambda_min = 1.0 / (wn_max * 100.0)
     12    maxbx = bmax / lambda_min
     13
     14    uv = numpy.zeros([mult*npix,mult*npix])
     15
     16    for bxby in bxbylist:
     17        ibx = int(round(((bxby[0] / lamb) / maxbx) * mult * rpix))
     18        if ibx < 0:
     19            ibx = mult*npix + ibx
     20        iby = int(round(((bxby[1] / lamb) / maxbx) * mult * rpix))
     21        if iby < 0:
     22            iby = mult*npix + iby
     23        uv[ibx,iby] = 1.0
     24 
     25    temp = numpy.fft.fftshift(uv)
     26    temp = numpy.fft.fft2(temp)
     27    temp = numpy.fft.fftshift(temp)
     28    temp = numpy.power(temp, 2)
     29    # convert to real, should be real anway as FT of symmetric function
     30    dirty_beam = numpy.abs(temp)
     31    # and normalise
     32    dirty_beam /= numpy.max(dirty_beam)
     33
     34    # truncate dirty beam pattern
     35    dirty_beam = dirty_beam[
     36      (mult-1)*rpix:(mult+1)*rpix, (mult-1)*rpix:(mult+1)*rpix]
     37
     38    return dirty_beam
     39
     40def calculate_primary_beam(npix, rpix, mult, mx, my, bmax, m1_diameter,
     41  wn_max, wn):
     42    # assume uniform illumination of primary. uv goes out to
     43    # +- bmax / lambda_min, fill uv plane appropriate to primary
     44    # size
     45    mirror = numpy.ones([mult*npix,mult*npix])
     46    lambda_min = 1.0 / (wn_max * 100.0)
     47    lamb = 1.0 / (wn * 100.0)
     48    mirror[(numpy.sqrt(mx*mx + my*my) / (mult*rpix)) *
     49      (bmax / lambda_min) > (0.5 * m1_diameter / lamb)] = 0.0
     50
     51    # fftshift mirror to be centred at [0,0],
     52    # do a 2d ftt of it,
     53    # fftshift transform so that 0 freq is in centre of array
     54    temp = numpy.fft.fftshift(mirror)
     55    temp = numpy.fft.fft2(temp)
     56    temp = numpy.fft.fftshift(temp)
     57    # convert to real, should be real anyway as FT of symmetric function
     58    # square to give intensity response
     59    primary_amplitude_beam = numpy.real(temp)
     60    primary_intensity_beam = numpy.power(numpy.abs(temp), 2)
     61    # and normalise
     62    primary_amplitude_beam /= numpy.max(primary_amplitude_beam)
     63    primary_intensity_beam /= numpy.max(primary_intensity_beam)
     64
     65    # truncate primary beam pattern
     66    primary_amplitude_beam = primary_amplitude_beam[
     67      (mult-1)*rpix:(mult+1)*rpix,
     68      (mult-1)*rpix:(mult+1)*rpix]
     69    primary_intensity_beam = primary_intensity_beam[
     70      (mult-1)*rpix:(mult+1)*rpix,
     71      (mult-1)*rpix:(mult+1)*rpix]
     72
     73    return primary_intensity_beam, primary_amplitude_beam
    2574
    2675
    … …  
    3281        self.previous_results = previous_results
    3382        self.result = collections.OrderedDict()
     83
     84        # start parallel python (pp), find the number of CPUS available
     85        ppservers = ()
     86        self.job_server = pp.Server(ppservers=ppservers)
     87        print 'Sampler starting pp with %s workers' % \
     88          self.job_server.get_ncpus()
    3489
    3590    def run(self):
    … …  
    63118        # go out to radius of first null
    64119        rpix = int(max_beam_radius / max_pixsize)
    65 #        rpix = 32
    66120        npix = 2 * rpix
    67121        self.result['npix'] = npix
    … …  
    72126        axis *= 206264.806247
    73127        self.result['spatial axis [arcsec]'] = axis
     128        axis1 = co.Axis(data=-axis, title='RA offset', units='arcsec')
     129        axis2 = co.Axis(data=axis, title='Dec offset', units='arcsec')
     130        axis3 = co.Axis(data=fts_wn, title='Frequency', units='cm-1')
    74131
    75132        # calculate beams for each point on fts_wn
    … …  
    82139        self.result['dirty beam'] = collections.OrderedDict()
    83140
     141        wn_max = np.max(fts_wn)
     142        m1_diameter = telescope['Primary mirror diameter']
     143
     144        # first calculate primary beam for each wn
     145        jobs = {}
    84146        for wn in fts_wn:
    85             # first calculate primary beam
    86             #
    87             # assume uniform illumination of primary. uv goes out to
    88             # +- bmax / lambda_min, fill uv plane appropriate to primary
    89             # size
    90             mirror = np.ones([mult*npix,mult*npix])
    91             lambda_min = 1.0 / (np.max(fts_wn) * 100.0)
    92             lamb = 1.0 / (wn * 100.0)
    93             mirror[(np.sqrt(mx*mx + my*my) / (mult*rpix)) *
    94               (bmax / lambda_min) > \
    95               (0.5 * telescope['Primary mirror diameter'] / lamb)] = 0.0
    96 
    97             # fftshift mirror to be centred at [0,0],
    98             # do a 2d ftt of it,
    99             # fftshift transform so that 0 freq is in centre of array
    100             temp = np.fft.fftshift(mirror)
    101             temp = np.fft.fft2(temp)
    102             temp = np.fft.fftshift(temp)
    103             # convert to real, should be real anyway as FT of symmetric function
    104             # square to give intensity response
    105             primary_amplitude_beam = np.real(temp)
    106             primary_intensity_beam = np.power(np.abs(temp), 2)
    107             # and normalise
    108             primary_amplitude_beam /= np.max(primary_amplitude_beam)
    109             primary_intensity_beam /= np.max(primary_intensity_beam)
    110 
    111             # truncate primary beam pattern
    112             primary_amplitude_beam = primary_amplitude_beam[
    113               (mult-1)*rpix:(mult+1)*rpix,
    114               (mult-1)*rpix:(mult+1)*rpix]
    115             primary_intensity_beam = primary_intensity_beam[
    116               (mult-1)*rpix:(mult+1)*rpix,
    117               (mult-1)*rpix:(mult+1)*rpix]
    118 
    119             # axes
    120             axis1 = Axis(data=-axis, title='RA offset', units='arcsec')
    121             axis2 = Axis(data=axis, title='Dec offset', units='arcsec')
    122             image = Image(data=primary_intensity_beam, axes=[axis1, axis2],
     147            # submit jobs
     148            indata = (npix, rpix, mult, mx, my, bmax, m1_diameter, wn_max,
     149              wn,)
     150            jobs[wn] = self.job_server.submit(calculate_primary_beam,
     151              indata, (), ('numpy',))
     152
     153        primary_amplitude_beam = np.zeros([npix,npix,len(fts_wn)], np.complex)
     154        for iwn,wn in enumerate(fts_wn):
     155            # collect and store results
     156            primary_intensity_beam, primary_amplitude_beam[:,:,iwn] = jobs[wn]()
     157            image = co.Image(data=primary_intensity_beam, axes=[axis1, axis2],
    123158              title='Primary Beam %06.4g cm-1' % wn)
    124159            self.result['primary beam'][wn] = image
    125160
    126             image = Image(data=primary_amplitude_beam, axes=[axis1, axis2],
    127               title='Amplitude Primary Beam %06.4g cm-1' % wn)
    128             self.result['primary amplitude beam'][wn] = image
    129 
    130             # second calculate dirty beam
     161        cube = co.Cube(data=primary_amplitude_beam, axes=[axis1, axis2, axis3],
     162          title='Amplitude Primary Beam %06.4g cm-1' % wn)
     163        self.result['primary amplitude beam'] = cube
     164
     165        # second calculate dirty beam, same use of pp
     166        for wn in fts_wn:
    131167            #
    132168            # 1. uv plane covered by discrete points so should use a
    … …  
    140176            # positions. Fastest but adequate?
    141177
    142             uv = np.zeros([mult*npix,mult*npix])
    143             maxbx = bmax / np.min(fts_lambda)
    144             for bxby in uvmapgen['bxby']:
    145                 ibx = int(round(((bxby[0] / lamb) / maxbx) * mult * rpix))
    146                 if ibx < 0:
    147                     ibx = mult*npix + ibx
    148                 iby = int(round(((bxby[1] / lamb) / maxbx) * mult * rpix))
    149                 if iby < 0:
    150                     iby = mult*npix + iby
    151                 uv[ibx,iby] = 1.0
    152  
    153             temp = np.fft.fftshift(uv)
    154             temp = np.fft.fft2(temp)
    155             temp = np.fft.fftshift(temp)
    156             temp = np.power(temp, 2)
    157             # convert to real, should be real anway as FT of symmetric function
    158             dirty_beam = np.abs(temp)
    159             # and normalise
    160             dirty_beam /= np.max(dirty_beam)
    161 
    162             # truncate dirty beam pattern
    163             dirty_beam = dirty_beam[
    164               (mult-1)*rpix:(mult+1)*rpix,
    165               (mult-1)*rpix:(mult+1)*rpix]
    166 
    167             # axes
    168             axis = np.arange(-rpix, rpix, dtype=np.float)
    169             axis *= (np.min(fts_lambda) / bmax)
    170             axis *= 206264.806247
    171             axis1 = Axis(data=-axis, title='RA offset', units='arcsec')
    172             axis2 = Axis(data=axis, title='Dec offset', units='arcsec')
    173            
    174             image = Image(data=dirty_beam, axes=[axis1, axis2],
     178            indata = (npix, rpix, mult, bmax, uvmapgen['bxby'], wn_max, wn,)
     179            jobs[wn] = self.job_server.submit(calculate_dirty_beam,
     180              indata, (), ('numpy',))
     181
     182
     183        # axes for dirty beam
     184        axis = np.arange(-rpix, rpix, dtype=np.float)
     185        axis *= (np.min(fts_lambda) / bmax)
     186        axis *= 206264.806247
     187        axis1 = co.Axis(data=-axis, title='RA offset', units='arcsec')
     188        axis2 = co.Axis(data=axis, title='Dec offset', units='arcsec')
     189
     190        for wn in fts_wn:
     191            dirty_beam = jobs[wn]()
     192            image = co.Image(data=dirty_beam, axes=[axis1, axis2],
    175193              title='Dirty Beam %06.4g cm-1' % wn)
    176            
    177194            self.result['dirty beam'][wn] = image
    178195
Note: See TracChangeset for help on using the changeset viewer.