lofarscripts/ms2matlabcube.py at master · tammojan/lofarscripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
#!/usr/bin/env python

from pyrap.tables import table, taql
from math import sqrt
import numpy as np
import scipy.io
import argparse
import sys

def ms2matlabcube(datacolname='DATA',msname='test1.MS',visfilename='bbsvis.mat',modelcolname='MODEL_DATA',applyweights=True,numtimes=0):
  """converts a measurement set to a matlab file"""
  t0=table(msname,ack=False)


  t=taql('SELECT DISTINCT TIME from $t0');
  if numtimes>0:
    ntimes=min(t.nrows(),numtimes);
  else:
    ntimes=t.nrows();
  print "Num times:",ntimes

  t=taql('SELECT DATA from $t0 LIMIT 1');
  tmp=t.col('DATA');
  nch=tmp[0].shape[0];
  print "Num channels:",nch

  t=taql('select DISTINCT ANTENNA1 from $t0');
  nants=t.nrows()+1;
  t=taql('select DISTINCT ANTENNA2 from $t0');
  nants=max(nants,t.nrows());
  print "Num ants:",nants

  V  = np.zeros((nants*2,nants*2,ntimes,nch),dtype=np.complex)
  Vm = np.zeros((nants*2,nants*2,ntimes,nch),dtype=np.complex)
  W  = np.zeros((nants*2,nants*2,ntimes,nch),dtype=np.float)
  UVW= np.zeros((nants,nants,ntimes,3),dtype=np.float)

  for ch in range(nch):
    print "Channel",ch
    for timeslot in range(ntimes):
      print "Timeslot",timeslot
      t=taql('SELECT *, MAX(WEIGHT_SPECTRUM) as MW FROM $t0 WHERE TIME IN (SELECT DISTINCT TIME FROM $t0 OFFSET $timeslot LIMIT 1)')

      data=t.col(datacolname)
      flag=t.col('FLAG')
      model=t.col(modelcolname)
      weight=t.col('WEIGHT_SPECTRUM')
      uvw=t.col('UVW')

      antenna1=t.col('ANTENNA1')
      antenna2=t.col('ANTENNA2')

      for bl in range(data.nrows()):
        ant1=antenna1[bl]
        ant2=antenna2[bl]
        UVW[ant1][ant2][timeslot][:] = uvw[timeslot]
        UVW[ant2][ant1][timeslot][:] =-uvw[timeslot]
        if ant1==ant2:
          continue
        for cor in range(4):
          if not flag[bl][0][cor]:
            if applyweights:
              V[ 2*ant1+cor/2][2*ant2+cor%2][ timeslot][ch] = data[bl][ch][cor]*sqrt(weight[bl][ch][cor])
              Vm[2*ant1+cor/2][2*ant2+cor%2][ timeslot][ch] =model[bl][ch][cor]*sqrt(weight[bl][ch][cor])
              # symmmetry
              V[ 2*ant2+cor%2][2*ant1+cor/2][ timeslot][ch] =( data[bl][ch][cor]*sqrt(weight[bl][ch][cor])).conjugate()
              Vm[2*ant2+cor%2][2*ant1+cor/2][ timeslot][ch] =(model[bl][ch][cor]*sqrt(weight[bl][ch][cor])).conjugate()
            else:
              V[ 2*ant1+cor/2][2*ant2+cor%2][ timeslot][ch] =  data[bl][ch][cor]
              Vm[2*ant1+cor/2][2*ant2+cor%2][ timeslot][ch] = model[bl][ch][cor]
              W[ 2*ant1+cor/2][2*ant2+cor%2][ timeslot][ch] =weight[bl][ch][cor]
              # symmetry
              V[ 2*ant2+cor%2][2*ant1+cor/2][ timeslot][ch] =(  data[bl][ch][cor]).conjugate()
              Vm[2*ant2+cor%2][2*ant1+cor/2][ timeslot][ch] =( model[bl][ch][cor]).conjugate()
              W[ 2*ant1+cor%2][2*ant2+cor/2][ timeslot][ch] =weight[bl][ch][cor]


  if applyweights:
    scipy.io.savemat(visfilename, dict(Vcube=V, Vmcube=Vm, UVW=UVW), oned_as="row")
  else:
    scipy.io.savemat(visfilename, dict(Vcube=V, Vmcube=Vm, Wgtcube=W, UVW=UVW), oned_as="row")

  print "Stored timeslot", timeslot, "of column", datacolname, "of file", msname, "as", visfilename
  print "Stored timeslot", timeslot, "of column", modelcolname, "of file", msname, "in", visfilename
  #if not applyweights:
    #print "Stored timeslot", timeslot, "weights of file", msname, "in", visfilename

if __name__ == '__main__':
  parser = argparse.ArgumentParser(description = "Extract data and model data from a MS, save in matlab format")
  parser.add_argument("msname", help="Name of the MS that contains data, model data and instrument table")
  parser.add_argument("-d", "--datacolname", help="Data column name", default="DATA")
  parser.add_argument("-m", "--modelcolname", help="Model column name", default="MODEL_DATA")
  parser.add_argument("-o", "--output", help="Save output in this output file", default="bbsvis.mat")
  parser.add_argument("-nw", "--no-weights", help="Do not apply weights, store them separately in the outputfile", action='store_true')
  parser.add_argument("-t", "--numtimes", help="Number of time slots to include (0 for all)", default=0, type=int)

  args = parser.parse_args()

  ms2matlabcube(msname=args.msname,datacolname=args.datacolname,modelcolname=args.modelcolname,applyweights=(not args.no_weights),visfilename=args.output,numtimes=args.numtimes)