From ba9829017df5885c95c89792944eaecd36f61138 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Mon, 6 Jun 2022 14:56:41 +0200 Subject: [PATCH 01/15] Update gitignore --- .gitignore | 87 ++------------------ surfvis/__pycache__/__init__.cpython-37.pyc | Bin 276 -> 0 bytes surfvis/__pycache__/surfvis.cpython-37.pyc | Bin 6813 -> 0 bytes 3 files changed, 6 insertions(+), 81 deletions(-) delete mode 100644 surfvis/__pycache__/__init__.cpython-37.pyc delete mode 100644 surfvis/__pycache__/surfvis.cpython-37.pyc diff --git a/.gitignore b/.gitignore index 84229f4..07b02c6 100644 --- a/.gitignore +++ b/.gitignore @@ -1,18 +1,9 @@ -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class - -# C extensions -*.so - # Distribution / packaging .Python env/ build/ develop-eggs/ dist/ -downloads/ eggs/ .eggs/ lib/ @@ -22,81 +13,15 @@ sdist/ var/ wheels/ *.egg-info/ -.installed.cfg *.egg -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec - -# Installer logs -pip-log.txt -pip-delete-this-directory.txt - -# Unit test / coverage reports -htmlcov/ -.tox/ -.coverage -.coverage.* -.cache -nosetests.xml -coverage.xml -*.cover -.hypothesis/ -.pytest_cache/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - # Sphinx documentation docs/_build/ -# PyBuilder -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# pyenv -.python-version - -# celery beat schedule file -celerybeat-schedule - -# SageMath parsed files -*.sage.py - -# dotenv -.env - -# virtualenv -.venv -venv/ -ENV/ - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site +# Unit test / coverage reports +.pytest_cache/ -# mypy -.mypy_cache/ +# Compiled bytecode +surfvis/*.pyc +surfvis/__pycache__/ +tests/__pycache__/ diff --git a/surfvis/__pycache__/__init__.cpython-37.pyc b/surfvis/__pycache__/__init__.cpython-37.pyc deleted file mode 100644 index a712b68c853665bc78a5b9771c784870532691ef..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 276 zcmYL^u};G<5QdW!6lCZNFnR4zJI%y^P{jlbs>IgiWa6u0aD2sf3hl_-@Jd;E1h!5% zNQhs$@4vf$xx1ZCS(bG9U8;}dU*qtnJT=Qy;UYn2>rChCTo;eoY;r)hLO*{}0j+V! zqv?=uJihX+r5G4feQe3|Ilotq-{Jm^NHgq|u{twcm{t+_7*@5c zu4G-+3wAiIL?spvvtlR#G&KOjIlvBmFysI}Ge1eU{o+{5mO7MHIo&BISDv1+i9v?m dZ+7P20gN*dKy3HJ*>y*{w|G;Ie-;TR`41TtOj`f| diff --git a/surfvis/__pycache__/surfvis.cpython-37.pyc b/surfvis/__pycache__/surfvis.cpython-37.pyc deleted file mode 100644 index f2dabc0baae0f2314fe14da5215647ce70d0003b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 6813 zcmbVRO;8+3cCNqvp=l5T#E(FzkN~L#4dO?CNE!+8Gt#bvv_dm@OHy)zz)4 zLV|Mk#zruER1627eAsBh7jMLd4?ge7;nT(*e3-B!DtxfR7c0WCo{KMYvA>tyAkb(o zZlW?@zI^%K%lDI2FL}_`mXh#0z0_v@YFLu~g9e8`6~KKw{#!ImVzMI{k|ot;rX0vJ zVN^#ctA;9I+0hJ5P{oY6=o3al^hqNr`WB-Fz2>CKtwt->#8~_r$w;#VOQLUMEi8q; zowc$w`cq6^lQL}|AX{2B)P-U(BS-0{)xfT}&rHv^Jd9g+n4c?nm2$4^=VcXn$%EA0u*!C74mV*@X5&6GoJ;N{gan8m7|?FIl6Wy|-?qV+q;i(~l@ z(^H$3vNh$HE-Rbv)HAPAvhsmH~)M{*}vmQ4yOt*OZryUM2740s)4X4xFO#;$+V11`R?8upA< zy-HDkSrKvLzt+pXYnCnY_BOdVRI*MO&t@Il5284$h(6^970=2A-nJE9evBFYiDkms zR@rg`ea#AVFs)i%U|YV9pwO)yJ6OqRk;W`yW9BJHBW=@is!@Vj1#{a8Ac=+PY?hgU znXfq8WjBg9`f%z~_6@uqRP?G-3G~ZN7W!R%Y2n2}#?VBkdTsEgcZ!-0G>}jT6WMIR zwj374L>CUN&`5u@^f_)TVs$@=5JJP5Pr^sH|iCL~r=+#Zrw2}#H9^V~kGLecF zeC$xoX6-0xx`E}o=5#o`(17&m`G#34uCuELn>P7C(0uq`{IIk{OG`&fXTqz8OK0Zw zZQs)0*}Qy71dI8bNDIY~IqVmbpcGh<%h1Cw3L{Ff@@M#CxAgM>GXi zt})w-b^cj2lAKa{vx(8u=V8ujs73~izo7uY1Q2?4bLN@6BeYYAG>r!h{UZ>!O$y_OD6ad}T>@|NoT zI_O{uy^w;_^-iYJOVw@XGSm8{dKZfw$h9^eXYs9s_qU)-oTLo8OK0jmyp1J$B~Z5W z_IfXC0dz`0DL@^7`VR0uJfdwCw5OrNY3S9-Ph(9QlwAU9I}ox-eyv+z+A((q&`B<* z1eYFg=>eBs-UBWjpzIS6N$v-9VhwO-gZ}zieipQ+4`hr7xU#Q!e+ULxr`Xvb>*8&A zNn=0FaX0llGr>8gLDs1~c~4oAUhB0%tU1F6`6)iwFYyj)3u=#8wJ*P^uS-ATiG4SD zLSyh;BA&cc@VhSf_42`Bh=XfCC;@e|zJUhRMYzVOpdCs1{xecwm!_0{A?1~j(krB_ z)&`;Rz*Zc4{-&vaKXHbvzJ}DpJ-9>gHtp{x=J%C1Uw@PiiwaBc$K8OZS?3m6^`8+J zXL}{C!kz=Y5~A*WNv{uc_!_g{L;fI7?SJ6N8I$@5 z@2hW0>yqp4ID$Weq*p(Z^xaQN`mQNS^9GxeG;i!9ckKO0DkOckE}{J+T5abCKD31X z_4l=5cypLC2$s1ZU|GNy`31sqSh8R0lbGHm0fRYU&^9pe*abEWoM;bm>_V4R8wrMj zOMHY~9FW-X9!?El2G|H>jRHHihi^3aVHD38@RvmV{fft6<+rs6J?x_iGD(4pNxziZD#aLEgSTgGJKqm@+@*$!T0Ipd!;rmuoHaj1a@LS z_QqS6_;~3G*{U{)sE@&NZHP;pnM^%AJ+XtS=I&r~u?ytH@!i&@IKCu-5qrELiSbz; z$I9tmiN&(d54%^;*qe{PnfK>wNgET3bG z-6ZuIpZKvh{MvCGTPmn`BtBiALs>S-XBlGcK&@Tp*X!5VvWTQvrU{+?8o29hMdbM^ zpXHNG?vupKANlptjrvV`y9wwP`%pH zq_>GkoWmPg+tO&nQ{wzedLul==jwM3q;9EpQ^0oxe9P^DJ-*;`*zG)f#OG}TvTd-5Iyv%7Rc>Ti$Mwz5^4=bvBIPIT-k z_&j@$(;NG{&whdEo2BTcy&?CX!=h`@9NISYr?X){{^qyq(#{1~eIa;I%9a)nWMP-v z{C0hby%2UuzWFcMWxBM?zQW(v_*-Y|MUDLhX%vjtSNN2WF<$?Y?2as-Uq|Gj()#7< z2T@2|edeH2vUI;{|PQUP+=$=1aM~_PE z*iXps5{S|4b$>KFJ=bW$d4X)KA#rO*tArjepF9j-AEER^eFZw~`;_6#)dtrhl#+h^ z&E!Hje2hIQ@sZWxULrYXB$}V(Mr@_B?J^^|NFUp%mBM&=3!nS?5AuZI3AXn?&_{c7 z^l$eX!2NIUhn>&hFtRc3QIzKZl?gp;O+9E9h4a{LHdOSNp{8$@Ux&$*sIBMqjPilJ{edcR zh)EReqHl**@ows0{-%2O&-YO1b3~j^Nwla#*wOR4(u;K=$gBUKlbgy8x3|5*8y6;DtxgAt8 zDy}0a8v{ejJ9atJsLC@Eu4Q^iCsbN9;A^zu;>+@?6$j3(xJ|4#i>Zu)5#_dMPQbEQ z8@4S{g$hTq?R=nzrzw(-qvpd#eCtb8E~Ap2K)DJ(Pfr(LeftmpbNSVO-z)xQ_uqZ_ z@8MtFD;C7>&-XH_5qpQ4jv28{(cv)22bqVWc#Be^*vpw-(f3+Ox6^1M6`Szr8qDjW z<_tAd4|@I7P*I3W5vO7X-T)y6sX0ej6|v`05g4idt{-W1ofE}9Kk)1-4x?QU7K;Z| z`ukB#^hUgB1@IScd~obUF}MSle(`PBrii19I&tL(r{nvJpq2PE;|!>X2Z+;vpg6w`$dLagpOu>Umed)M(9D+&rpX7Z2z+%__R;u48T( zaYP-is3I+}b`VV9=vW`)$Vi$SKEfKylziFHik`V^v{HTTZ*S0vYNUKC$nAiWhwZ&Sz^r(lrWE)j_d-)JQw+>Y8XTcpC}Xl&PjDVPzr z%l>BNozX%xfgLy&C5_7}8>J{!KDvM~v=_9OG*K0hEtFJ7I-od*_}$8OU?hC=t%aRO zvAk39t<0S$(Y)$I{GSL9l7O2qoBd>Td=eH$ym2jIB*jAb9s1%vByZV_wkJTk=kSd< zuFPm&a{(AhGzHIk;~On`r0YAI1)E0PB^ac*t7?z}5F?ln-f;uLDpG>LNaCu60pmxq zZL|pKG-kvzm1s#%?;X2wXs+ zlK{qP%p2)Lel*iqMo=Up1>786Mi~j5l|mmFrr7YFN-mItNZvA%6is4r3-P44w`C~i zj*$@2k+Y4MjjP2S(IFr(LP$h%<>IJN&<)oZe z+La{wRIEeHC6x{u;l3v&AHoEs<&Ib!Gc?;R>rJRbzJ3IyWL+5nuXHj6s&;Tr$?Ym0 zC8Z=)4Xe`NsWqfv6={+JRVV1^PiWD^^{DM?8ua(&PNf5DHSi@0O}!&1NJ3KR-7WXX zgK7%+4!J{Zhm;XOX{_#6yFojI=PW(SAa-_Hp?yIr^rIHL2X$Wq13e%C{NmU>;nBl7 zoyw5XEq5q=nD<`6o|z8s867Dqf-OoSX&M)3^bYraP7`Ty@J%92aPE1x3IBkao78+k z6me$yTV>y)SCcsJ#33q-T!m+-fRX%5{ Gj{gH)eHG;Z From aba96b7ede3720531975c9acb5ec8dfa921ad6c3 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Mon, 6 Jun 2022 15:00:12 +0200 Subject: [PATCH 02/15] Convert indentation to spaces and add option to specify weights from sigma_spectrum column --- surfvis/surfchi2.py | 557 ++++++++++++++++++++++---------------------- surfvis/surfvis.py | 496 +++++++++++++++++++-------------------- 2 files changed, 528 insertions(+), 525 deletions(-) diff --git a/surfvis/surfchi2.py b/surfvis/surfchi2.py index 144a63c..0e99920 100755 --- a/surfvis/surfchi2.py +++ b/surfvis/surfchi2.py @@ -21,284 +21,287 @@ # COMMAND LINE OPTIONS def create_parser(): - parser = OptionParser(usage='%prog [options] msname') - parser.add_option('--rcol', default='RESIDUAL', - help='Residual column (default = RESIDUAL)') - parser.add_option('--wcol', default='WEIGHT_SPECTRUM', - help='Weight column (default = WEIGHT_SPECTRUM)') - parser.add_option('--fcol', default='FLAG', - help='Flag column (default = FLAG)') - parser.add_option('--dataout', default='', - help='Output name of zarr dataset') - parser.add_option('--imagesout', default=None, - help='Output folder to place images. ' - 'If None (default) no plots are saved') - parser.add_option('--nthreads', default=4, type=int, - help='Number of dask threads to use') - parser.add_option('--ntimes', default=-1, type=int, - help='Number of unique times in each chunk.') - parser.add_option('--nfreqs', default=128, type=int, - help='Number of frequencies in a chunk.') - parser.add_option("--use-corrs", type=str, - help='Comma seprated list of correlations to use (do ' - 'not use spaces). Default = diagonal correlations') - return parser + parser = OptionParser(usage='%prog [options] msname') + parser.add_option('--rcol', default='RESIDUAL', + help='Residual column (default = RESIDUAL)') + parser.add_option('--wcol', default='WEIGHT_SPECTRUM', + help='Weight column (default = WEIGHT_SPECTRUM)') + parser.add_option('--fcol', default='FLAG', + help='Flag column (default = FLAG)') + parser.add_option('--dataout', default='', + help='Output name of zarr dataset') + parser.add_option('--imagesout', default=None, + help='Output folder to place images. ' + 'If None (default) no plots are saved') + parser.add_option('--nthreads', default=4, type=int, + help='Number of dask threads to use') + parser.add_option('--ntimes', default=-1, type=int, + help='Number of unique times in each chunk.') + parser.add_option('--nfreqs', default=128, type=int, + help='Number of frequencies in a chunk.') + parser.add_option("--use-corrs", type=str, + help='Comma seprated list of correlations to use (do ' + 'not use spaces). Default = diagonal correlations') + return parser def main(): - (options,args) = create_parser().parse_args() - - if options.dataout == '': - options.dataout = os.getcwd() + '/chi2' - - if os.path.isdir(options.dataout): - print(f"Removing existing {options.dataout} folder") - os.system(f"rm -r {options.dataout}") - - # Some error trapping - if len(args) != 1: - print('Please specify a single Measurement Set to plot.') - sys.exit(-1) - else: - msname = args[0].rstrip('/') - - # chunking info - schema = {} - schema[options.fcol] = {'dims': ('chan', 'corr')} - xds = xds_from_ms(msname, - chunks={'row': -1}, - columns=['TIME', options.fcol], - group_cols=['FIELD_ID', 'DATA_DESC_ID', 'SCAN_NUMBER'], - table_schema=schema) - - chunks = [] - rbin_idx = [] - rbin_counts = [] - tbin_idx = [] - tbin_counts = [] - fbin_idx = [] - fbin_counts = [] - t0s = [] - tfs = [] - for ds in xds: - time = ds.TIME.values - ut, counts = np.unique(time, return_counts=True) - if options.ntimes in [0, -1]: - utpc = ut.size - else: - utpc = options.ntimes - row_chunks = [np.sum(counts[i:i+utpc]) - for i in range(0, ut.size, utpc)] - - nchan = ds.chan.size - if options.nfreqs in [0, -1]: - options.nfreqs = nchan - - # list per ds - chunks.append({'row': tuple(row_chunks), 'chan': options.nfreqs}) - - ridx = np.zeros(len(row_chunks)) - ridx[1:] = np.cumsum(row_chunks)[0:-1] - rbin_idx.append(da.from_array(ridx.astype(int), chunks=1)) - rbin_counts.append(da.from_array(row_chunks, chunks=1)) - - ntime = ut.size - tidx = np.arange(0, ntime, utpc) - tbin_idx.append(da.from_array(tidx.astype(int), chunks=1)) - tidx2 = np.append(tidx, ntime) - tcounts = tidx2[1:] - tidx2[0:-1] - tbin_counts.append(da.from_array(tcounts, chunks=1)) - - t0 = ut[tidx] - t0s.append(da.from_array(t0, chunks=1)) - tf = ut[tidx + tcounts -1] - tfs.append(da.from_array(tf, chunks=1)) - - fidx = np.arange(0, nchan, options.nfreqs) - fbin_idx.append(da.from_array(fidx, chunks=1)) - fidx2 = np.append(fidx, nchan) - fcounts = fidx2[1:] - fidx2[0:-1] - fbin_counts.append(da.from_array(fcounts, chunks=1)) - - schema = {} - schema[options.rcol] = {'dims': ('chan', 'corr')} - schema[options.wcol] = {'dims': ('chan', 'corr')} - schema[options.fcol] = {'dims': ('chan', 'corr')} - - xds = xds_from_ms(msname, - columns=[options.rcol, options.wcol, options.fcol, - 'ANTENNA1', 'ANTENNA2', 'TIME'], - chunks=chunks, - group_cols=['FIELD_ID', 'DATA_DESC_ID', 'SCAN_NUMBER'], - table_schema=schema) - - if options.use_corrs is None: - print('Using only diagonal correlations') - if len(xds[0].corr) > 1: - use_corrs = [0, -1] - else: - use_corrs = [0] - else: - use_corrs = list(map(int, options.use_corrs.split(','))) - print(f"Using correlations {use_corrs}") - ncorr = len(use_corrs) - - out_ds = [] - idts = [] - for i, ds in enumerate(xds): - ds = ds.sel(corr=use_corrs) - - resid = ds.get(options.rcol).data - weight = ds.get(options.wcol).data - flag = ds.get(options.fcol).data - ant1 = ds.ANTENNA1.data - ant2 = ds.ANTENNA2.data - - # ncorr = resid.shape[0] - - # time = ds.TIME.values - # utime = np.unique(time) - - # spw = xds_from_table(msname + '::SPECTRAL_WINDOW') - # freq = spw[0].CHAN_FREQ.values - - field = ds.FIELD_ID - ddid = ds.DATA_DESC_ID - scan = ds.SCAN_NUMBER - - tmp = surfchisq(resid, weight, flag, ant1, ant2, - rbin_idx[i], rbin_counts[i], - fbin_idx[i], fbin_counts[i]) - - d = xr.Dataset( - data_vars={'data': (('time', 'freq', 'corr', 'p', 'q', '2'), tmp), - 'fbin_idx': (('freq'), fbin_idx[i]), - 'fbin_counts': (('freq'), fbin_counts[i]), - 'tbin_idx': (('time'), tbin_idx[i]), - 'tbin_counts': (('time'), tbin_counts[i])}, - attrs = {'FIELD_ID': ds.FIELD_ID, - 'DATA_DESC_ID': ds.DATA_DESC_ID, - 'SCAN_NUMBER': ds.SCAN_NUMBER}, - # coords={'time': (('time'), utime), - # 'freq': (('freq'), freq), - # 'corr': (('corr'), np.arange(ncorr))} - ) - - idt = f'::F{ds.FIELD_ID}_D{ds.DATA_DESC_ID}_S{ds.SCAN_NUMBER}' - out_ds.append(xds_to_zarr(d, options.dataout + idt)) - idts.append(idt) - - - with ProgressBar(): - dask.compute(out_ds) - - # primitive plotting - if options.imagesout is not None: - foldername = options.imagesout.rstrip('/') - if not os.path.isdir(foldername): - os.system('mkdir '+ foldername) - - for idt in idts: - xds = xds_from_zarr(options.dataout + idt) - for ds in xds: - field = ds.FIELD_ID - if not os.path.isdir(foldername + f'/field{field}'): - os.system('mkdir '+ foldername + f'/field{field}') - - spw = ds.DATA_DESC_ID - if not os.path.isdir(foldername + f'/field{field}' + f'/spw{spw}'): - os.system('mkdir '+ foldername + f'/field{field}' + f'/spw{spw}') - - scan = ds.SCAN_NUMBER - if not os.path.isdir(foldername + f'/field{field}' + f'/spw{spw}' + f'/scan{scan}'): - os.system('mkdir '+ foldername + f'/field{field}' + f'/spw{spw}'+ f'/scan{scan}') - - tmp = ds.data.values - tbin_idx = ds.tbin_idx.values - tbin_counts = ds.tbin_counts.values - fbin_idx = ds.fbin_idx.values - fbin_counts = ds.fbin_counts.values - - ntime, nfreq, ncorr, _, _, _ = tmp.shape - - basename = foldername + f'/field{field}' + f'/spw{spw}'+ f'/scan{scan}/' - if len(os.listdir(basename)): - print(f"Removing contents of {basename} folder") - os.system(f'rm {basename}*.png') - for t in range(ntime): - for f in range(nfreq): - for c in range(ncorr): - chi2 = tmp[t, f, c, :, :, 0] - N = tmp[t, f, c, :, :, 1] - chi2_dof = np.zeros_like(chi2) - chi2_dof[N>0] = chi2[N>0]/N[N>0] - chi2_dof[N==0] = np.nan - t0 = tbin_idx[t] - tf = tbin_idx[t] + tbin_counts[t] - chan0 = fbin_idx[f] - chanf = fbin_idx[f] + fbin_counts[f] - makeplot(chi2_dof, basename + f't{t}_f{f}_c{c}.png', - f't {t0}-{tf}, chan {chan0}-{chanf}, corr {c}') - - # reduce over corr - chi2 = np.nansum(tmp[t, f, (0, -1), :, :, 0], axis=0) - N = np.nansum(tmp[t, f, (0, -1), :, :, 1], axis=0) - chi2_dof = np.zeros_like(chi2) - chi2_dof[N>0] = chi2[N>0]/N[N>0] - chi2_dof[N==0] = np.nan - t0 = tbin_idx[t] - tf = tbin_idx[t] + tbin_counts[t] - chan0 = fbin_idx[f] - chanf = fbin_idx[f] + fbin_counts[f] - makeplot(chi2_dof, basename + f't{t}_f{f}.png', - f't {t0}-{tf}, chan {chan0}-{chanf}') - - # reduce over freq - chi2 = np.nansum(tmp[t, :, (0, -1), :, :, 0], axis=(0,1)) - N = np.nansum(tmp[t, :, (0, -1), :, :, 1], axis=(0,1)) - chi2_dof = np.zeros_like(chi2) - chi2_dof[N>0] = chi2[N>0]/N[N>0] - chi2_dof[N==0] = np.nan - t0 = tbin_idx[t] - tf = tbin_idx[t] + tbin_counts[t] - makeplot(chi2_dof, basename + f't{t}.png', - f't {t0}-{tf}') - - # now the entire scan - chi2 = np.nansum(tmp[:, :, (0, -1), :, :, 0], axis=(0,1,2)) - N = np.nansum(tmp[:, :, (0, -1), :, :, 1], axis=(0,1,2)) - chi2_dof = np.zeros_like(chi2) - chi2_dof[N>0] = chi2[N>0]/N[N>0] - chi2_dof[N==0] = np.nan - makeplot(chi2_dof, basename + f'scan.png', - f'scan {scan}.png') + (options,args) = create_parser().parse_args() + + if options.dataout == '': + options.dataout = os.getcwd() + '/chi2' + + if os.path.isdir(options.dataout): + print(f"Removing existing {options.dataout} folder") + os.system(f"rm -r {options.dataout}") + + # Some error trapping + if len(args) != 1: + print('Please specify a single Measurement Set to plot.') + sys.exit(-1) + else: + msname = args[0].rstrip('/') + + # chunking info + schema = {} + schema[options.fcol] = {'dims': ('chan', 'corr')} + xds = xds_from_ms(msname, + chunks={'row': -1}, + columns=['TIME', options.fcol], + group_cols=['FIELD_ID', 'DATA_DESC_ID', 'SCAN_NUMBER'], + table_schema=schema) + + chunks = [] + rbin_idx = [] + rbin_counts = [] + tbin_idx = [] + tbin_counts = [] + fbin_idx = [] + fbin_counts = [] + t0s = [] + tfs = [] + for ds in xds: + time = ds.TIME.values + ut, counts = np.unique(time, return_counts=True) + if options.ntimes in [0, -1]: + utpc = ut.size + else: + utpc = options.ntimes + row_chunks = [np.sum(counts[i:i+utpc]) + for i in range(0, ut.size, utpc)] + + nchan = ds.chan.size + if options.nfreqs in [0, -1]: + options.nfreqs = nchan + + # list per ds + chunks.append({'row': tuple(row_chunks), 'chan': options.nfreqs}) + + ridx = np.zeros(len(row_chunks)) + ridx[1:] = np.cumsum(row_chunks)[0:-1] + rbin_idx.append(da.from_array(ridx.astype(int), chunks=1)) + rbin_counts.append(da.from_array(row_chunks, chunks=1)) + + ntime = ut.size + tidx = np.arange(0, ntime, utpc) + tbin_idx.append(da.from_array(tidx.astype(int), chunks=1)) + tidx2 = np.append(tidx, ntime) + tcounts = tidx2[1:] - tidx2[0:-1] + tbin_counts.append(da.from_array(tcounts, chunks=1)) + + t0 = ut[tidx] + t0s.append(da.from_array(t0, chunks=1)) + tf = ut[tidx + tcounts -1] + tfs.append(da.from_array(tf, chunks=1)) + + fidx = np.arange(0, nchan, options.nfreqs) + fbin_idx.append(da.from_array(fidx, chunks=1)) + fidx2 = np.append(fidx, nchan) + fcounts = fidx2[1:] - fidx2[0:-1] + fbin_counts.append(da.from_array(fcounts, chunks=1)) + + schema = {} + schema[options.rcol] = {'dims': ('chan', 'corr')} + schema[options.wcol] = {'dims': ('chan', 'corr')} + schema[options.fcol] = {'dims': ('chan', 'corr')} + + xds = xds_from_ms(msname, + columns=[options.rcol, options.wcol, options.fcol, + 'ANTENNA1', 'ANTENNA2', 'TIME'], + chunks=chunks, + group_cols=['FIELD_ID', 'DATA_DESC_ID', 'SCAN_NUMBER'], + table_schema=schema) + + if options.use_corrs is None: + print('Using only diagonal correlations') + if len(xds[0].corr) > 1: + use_corrs = [0, -1] + else: + use_corrs = [0] + else: + use_corrs = list(map(int, options.use_corrs.split(','))) + print(f"Using correlations {use_corrs}") + ncorr = len(use_corrs) + + out_ds = [] + idts = [] + for i, ds in enumerate(xds): + ds = ds.sel(corr=use_corrs) + + resid = ds.get(options.rcol).data + if options.wcol == 'SIGMA_SPECTRUM': + weight = 1.0/ds.get(options.wcol).data**2 + else: + weight = ds.get(options.wcol).data + flag = ds.get(options.fcol).data + ant1 = ds.ANTENNA1.data + ant2 = ds.ANTENNA2.data + + # ncorr = resid.shape[0] + + # time = ds.TIME.values + # utime = np.unique(time) + + # spw = xds_from_table(msname + '::SPECTRAL_WINDOW') + # freq = spw[0].CHAN_FREQ.values + + field = ds.FIELD_ID + ddid = ds.DATA_DESC_ID + scan = ds.SCAN_NUMBER + + tmp = surfchisq(resid, weight, flag, ant1, ant2, + rbin_idx[i], rbin_counts[i], + fbin_idx[i], fbin_counts[i]) + + d = xr.Dataset( + data_vars={'data': (('time', 'freq', 'corr', 'p', 'q', '2'), tmp), + 'fbin_idx': (('freq'), fbin_idx[i]), + 'fbin_counts': (('freq'), fbin_counts[i]), + 'tbin_idx': (('time'), tbin_idx[i]), + 'tbin_counts': (('time'), tbin_counts[i])}, + attrs = {'FIELD_ID': ds.FIELD_ID, + 'DATA_DESC_ID': ds.DATA_DESC_ID, + 'SCAN_NUMBER': ds.SCAN_NUMBER}, + # coords={'time': (('time'), utime), + # 'freq': (('freq'), freq), + # 'corr': (('corr'), np.arange(ncorr))} + ) + + idt = f'::F{ds.FIELD_ID}_D{ds.DATA_DESC_ID}_S{ds.SCAN_NUMBER}' + out_ds.append(xds_to_zarr(d, options.dataout + idt)) + idts.append(idt) + + + with ProgressBar(): + dask.compute(out_ds) + + # primitive plotting + if options.imagesout is not None: + foldername = options.imagesout.rstrip('/') + if not os.path.isdir(foldername): + os.system('mkdir '+ foldername) + + for idt in idts: + xds = xds_from_zarr(options.dataout + idt) + for ds in xds: + field = ds.FIELD_ID + if not os.path.isdir(foldername + f'/field{field}'): + os.system('mkdir '+ foldername + f'/field{field}') + + spw = ds.DATA_DESC_ID + if not os.path.isdir(foldername + f'/field{field}' + f'/spw{spw}'): + os.system('mkdir '+ foldername + f'/field{field}' + f'/spw{spw}') + + scan = ds.SCAN_NUMBER + if not os.path.isdir(foldername + f'/field{field}' + f'/spw{spw}' + f'/scan{scan}'): + os.system('mkdir '+ foldername + f'/field{field}' + f'/spw{spw}'+ f'/scan{scan}') + + tmp = ds.data.values + tbin_idx = ds.tbin_idx.values + tbin_counts = ds.tbin_counts.values + fbin_idx = ds.fbin_idx.values + fbin_counts = ds.fbin_counts.values + + ntime, nfreq, ncorr, _, _, _ = tmp.shape + + basename = foldername + f'/field{field}' + f'/spw{spw}'+ f'/scan{scan}/' + if len(os.listdir(basename)): + print(f"Removing contents of {basename} folder") + os.system(f'rm {basename}*.png') + for t in range(ntime): + for f in range(nfreq): + for c in range(ncorr): + chi2 = tmp[t, f, c, :, :, 0] + N = tmp[t, f, c, :, :, 1] + chi2_dof = np.zeros_like(chi2) + chi2_dof[N>0] = chi2[N>0]/N[N>0] + chi2_dof[N==0] = np.nan + t0 = tbin_idx[t] + tf = tbin_idx[t] + tbin_counts[t] + chan0 = fbin_idx[f] + chanf = fbin_idx[f] + fbin_counts[f] + makeplot(chi2_dof, basename + f't{t}_f{f}_c{c}.png', + f't {t0}-{tf}, chan {chan0}-{chanf}, corr {c}') + + # reduce over corr + chi2 = np.nansum(tmp[t, f, (0, -1), :, :, 0], axis=0) + N = np.nansum(tmp[t, f, (0, -1), :, :, 1], axis=0) + chi2_dof = np.zeros_like(chi2) + chi2_dof[N>0] = chi2[N>0]/N[N>0] + chi2_dof[N==0] = np.nan + t0 = tbin_idx[t] + tf = tbin_idx[t] + tbin_counts[t] + chan0 = fbin_idx[f] + chanf = fbin_idx[f] + fbin_counts[f] + makeplot(chi2_dof, basename + f't{t}_f{f}.png', + f't {t0}-{tf}, chan {chan0}-{chanf}') + + # reduce over freq + chi2 = np.nansum(tmp[t, :, (0, -1), :, :, 0], axis=(0,1)) + N = np.nansum(tmp[t, :, (0, -1), :, :, 1], axis=(0,1)) + chi2_dof = np.zeros_like(chi2) + chi2_dof[N>0] = chi2[N>0]/N[N>0] + chi2_dof[N==0] = np.nan + t0 = tbin_idx[t] + tf = tbin_idx[t] + tbin_counts[t] + makeplot(chi2_dof, basename + f't{t}.png', + f't {t0}-{tf}') + + # now the entire scan + chi2 = np.nansum(tmp[:, :, (0, -1), :, :, 0], axis=(0,1,2)) + N = np.nansum(tmp[:, :, (0, -1), :, :, 1], axis=(0,1,2)) + chi2_dof = np.zeros_like(chi2) + chi2_dof[N>0] = chi2[N>0]/N[N>0] + chi2_dof[N==0] = np.nan + makeplot(chi2_dof, basename + f'scan.png', + f'scan {scan}.png') def makeplot(data, name, subt): - nant, _ = data.shape - fig = plt.figure() - ax = plt.gca() - im = ax.imshow(data, cmap='inferno') - ax.set_xticks(np.arange(0, nant, 2)) - ax.set_yticks(np.arange(nant)) - ax.tick_params(axis='both', which='major', - length=1, width=1, labelsize=4) - - divider = make_axes_locatable(ax) - cax = divider.append_axes("bottom", size="3%", pad=0.2) - cb = fig.colorbar(im, cax=cax, orientation="horizontal") - cb.outline.set_visible(False) - cb.ax.tick_params(length=1, width=1, labelsize=6, pad=0.1) - - rax = divider.append_axes("right", size="50%", pad=0.025) - x = data[~ np.isnan(data)] - hist(x, bins='scott', ax=rax, histtype='stepfilled', - alpha=0.5, density=False) - rax.set_yticks([]) - rax.tick_params(axis='y', which='both', - bottom=False, top=False, - labelbottom=False) - rax.tick_params(axis='x', which='both', - length=1, width=1, labelsize=8) - - fig.suptitle(subt, fontsize=20) - plt.savefig(name, dpi=250) - plt.close(fig) + nant, _ = data.shape + fig = plt.figure() + ax = plt.gca() + im = ax.imshow(data, cmap='inferno') + ax.set_xticks(np.arange(0, nant, 2)) + ax.set_yticks(np.arange(nant)) + ax.tick_params(axis='both', which='major', + length=1, width=1, labelsize=4) + + divider = make_axes_locatable(ax) + cax = divider.append_axes("bottom", size="3%", pad=0.2) + cb = fig.colorbar(im, cax=cax, orientation="horizontal") + cb.outline.set_visible(False) + cb.ax.tick_params(length=1, width=1, labelsize=6, pad=0.1) + + rax = divider.append_axes("right", size="50%", pad=0.025) + x = data[~ np.isnan(data)] + hist(x, bins='scott', ax=rax, histtype='stepfilled', + alpha=0.5, density=False) + rax.set_yticks([]) + rax.tick_params(axis='y', which='both', + bottom=False, top=False, + labelbottom=False) + rax.tick_params(axis='x', which='both', + length=1, width=1, labelsize=8) + + fig.suptitle(subt, fontsize=20) + plt.savefig(name, dpi=250) + plt.close(fig) diff --git a/surfvis/surfvis.py b/surfvis/surfvis.py index 4b07184..ef2d6d1 100755 --- a/surfvis/surfvis.py +++ b/surfvis/surfvis.py @@ -21,253 +21,253 @@ def ri(message): # COMMAND LINE OPTIONS def create_parser(): - parser = OptionParser(usage='%prog [options] msname') - #parser.add_option('-m','--msname',dest='msname',help='Measurement Set to open',metavar='DIRECTORY') - parser.add_option('-l','--list',dest='dolist',action='store_true',help='List Measurement Set properties and exit',default=False) - parser.add_option('-d','--datacolumn',dest='column',help='Measurement Set column to plot (default = DATA)',default='DATA') - parser.add_option('-f','--field',dest='field',help='Field ID to plot (default = 0)',default=0) - parser.add_option('-s','--spw',dest='myspw',help='Comma separated list of SPWs to plot (default = all)',default='') - parser.add_option('-p','--plot',dest='plot',help='Set to amp, phase, real or imag (default = amp)',default='amp') - parser.add_option('-i','--i',dest='antenna1',help='Antenna 1: plot only this antenna',default=-1) - parser.add_option('-j','--j',dest='antenna2',help='Antenna 2: use with -i to plot a single baseline',default=-1) - parser.add_option('--noflags',dest='noflags',help='Disable flagged data overlay',action='store_true',default=False) - parser.add_option('--scale',dest='scale',help='Scale the image peak to this multiple of the per-corr min/max (default = scale image max to 5 sigma, this parameter is ignored for phase plots)',default=-1) - parser.add_option('--cmap',dest='mycmap',help='Matplotlib colour map to use (default = jet)',default='jet') - parser.add_option('-o','--opdir',dest='foldername',help='Output folder to store plots (default = msname___plots)',default='') - return parser + parser = OptionParser(usage='%prog [options] msname') + #parser.add_option('-m','--msname',dest='msname',help='Measurement Set to open',metavar='DIRECTORY') + parser.add_option('-l','--list',dest='dolist',action='store_true',help='List Measurement Set properties and exit',default=False) + parser.add_option('-d','--datacolumn',dest='column',help='Measurement Set column to plot (default = DATA)',default='DATA') + parser.add_option('-f','--field',dest='field',help='Field ID to plot (default = 0)',default=0) + parser.add_option('-s','--spw',dest='myspw',help='Comma separated list of SPWs to plot (default = all)',default='') + parser.add_option('-p','--plot',dest='plot',help='Set to amp, phase, real or imag (default = amp)',default='amp') + parser.add_option('-i','--i',dest='antenna1',help='Antenna 1: plot only this antenna',default=-1) + parser.add_option('-j','--j',dest='antenna2',help='Antenna 2: use with -i to plot a single baseline',default=-1) + parser.add_option('--noflags',dest='noflags',help='Disable flagged data overlay',action='store_true',default=False) + parser.add_option('--scale',dest='scale',help='Scale the image peak to this multiple of the per-corr min/max (default = scale image max to 5 sigma, this parameter is ignored for phase plots)',default=-1) + parser.add_option('--cmap',dest='mycmap',help='Matplotlib colour map to use (default = jet)',default='jet') + parser.add_option('-o','--opdir',dest='foldername',help='Output folder to store plots (default = msname___plots)',default='') + return parser def main(): - (options,args) = create_parser().parse_args() - - dolist = options.dolist - column = options.column - fieldid = int(options.field) - myspw = options.myspw - plot = options.plot - antenna1 = int(options.antenna1) - antenna2 = int(options.antenna2) - noflags = options.noflags - scale = float(options.scale) - mycmap = options.mycmap - foldername = options.foldername - - # Some error trapping - - if len(args) != 1: - ri('Please specify a single Measurement Set to plot.') - sys.exit(-1) - else: - msname = args[0].rstrip('/') - - if plot not in ['amp','phase','real','imag']: - ri('Requested plot not valid, must be one of amp, phase, real or imag.') - sys.exit(-1) - - - # MEASUREMENT SET INFO - - fieldtab = pyrap.tables.table(msname+'/FIELD') - sourceids = fieldtab.getcol('SOURCE_ID') - sourcenames = fieldtab.getcol('NAME') - fieldtab.done() - - spwtab = pyrap.tables.table(msname+'/SPECTRAL_WINDOW') - nspw = len(spwtab) - spwfreqs = spwtab.getcol('REF_FREQUENCY') - chanwidth = spwtab.getcol('CHAN_WIDTH')[0][0] # probably needs changing if SPWs have different widths - nchans = spwtab.getcol('NUM_CHAN') - spwtab.done() - - anttab = pyrap.tables.table(msname+'/ANTENNA') - nant = len(anttab) - antpos = anttab.getcol('POSITION') - antnames = anttab.getcol('NAME') - anttab.done() - - tt = pyrap.tables.table(msname) - usedants = numpy.unique(tt.getcol('ANTENNA1')) - - - # PRINT SUMMARY - - if dolist: - print('') - gi(' '+msname+'/FIELD') - gi(' ROW ID NAME') - for i in range(0,len(sourceids)): - print(' %-6s%-14s%-14s' % (i,sourceids[i],sourcenames[i])) - print('') - gi(' '+msname+'/SPECTRAL_WINDOW') - gi(' ROW CHANS WIDTH[MHz] REF_FREQ[MHz]') - for i in range(0,nspw): - print(' %-6s%-14s%-20s%-14s' % (i,str(nchans[i]),str(chanwidth/1e6),str(spwfreqs[i]/1e6))) - print('') - gi(' '+msname+'/ANTENNA') - gi(' ROW NAME POSITION') - for i in range(0,nant): - if i in usedants: - print(' %-6s%-14s%-14s' % (i,(antnames[i]),str(antpos[i]))) - else: - ri(' %-6s%-14s%-14s' % (i,(antnames[i]),str(antpos[i]))) - - print('') - tt.done() - - - # OR ELSE DO THE PLOTS - - - else: - # Create output folder if it doesn't exist - if foldername == '': - foldername = msname+'_'+column+'__plots' - if os.path.isdir(foldername): - print('Found',foldername) - else: - print('Creating',foldername) - os.system('mkdir '+foldername) - - # Make a complete list of SPWs if one isn't provided - if myspw == '': - myspw = numpy.arange(0,nspw) - else: - myspw = myspw.split(',') - - fieldname = sourcenames[fieldid] - - # Make a list of baseline pairs based on the antenna selections - baselines = [] - if antenna1 != -1 and antenna2 != -1: - baselines = [(antenna1,antenna2)] - elif antenna1 != -1: - i = antenna1 - for j in usedants: - if i != j: - pair = [i,j] - pair = sorted(pair) - if pair not in baselines: - if antenna1 != -1: - if antenna1 in pair: - baselines.append(pair) - else: - baselines.append(pair) - else: - for i in usedants: - for j in usedants: - if i != j: - pair = [i,j] - pair = sorted(pair) - if pair not in baselines: - if antenna1 != -1: - if antenna1 in pair: - baselines.append(pair) - else: - baselines.append(pair) - - # Loop over baselines - for baseline in baselines: - - # Determine unprojected baseline length - ap1 = antpos[baseline[0]] - ap2 = antpos[baseline[1]] - blength = (((ap1[0]-ap2[0])**2.0)+((ap1[1]-ap2[1])**2.0)+((ap1[2]-ap1[2])**2.0))**0.5 - blength = str(round(blength/1000.0,2)) - - print('Plotting baseline:',baseline,' Deprojected length:',blength,'km') - - # Get the data - datacols = [] - flagcols = [] - # Loop over SPWs - print('SPW:') - for spw in myspw: - print(spw) - subtab = tt.query(query='ANTENNA1=='+str(baseline[0]) - +' && ANTENNA2=='+str(baseline[1]) - +' && DATA_DESC_ID=='+str(spw) - +' && FIELD_ID=='+str(fieldid)) - datacol = subtab.getcol(column) - flagcol = subtab.getcol('FLAG') - datacols.append(datacol) - flagcols.append(flagcol) - print('') - # Reshape the data - baselinedata = datacols[0] - flagdata = flagcols[0] - for p in range(1,len(datacols)): - baselinedata = numpy.concatenate((baselinedata,datacols[p]),axis=1) - flagdata = numpy.concatenate((flagdata,flagcols[p]),axis=1) - - # Get number of corr products from the data shape - n_corr = baselinedata.shape[2] - - # Generate png name - pngname = foldername+'/'+msname.split('/')[-1].rstrip('/')+'_baseline_'+str(baseline[0])+'_'+str(baseline[1]) - pngname+='_field'+str(fieldid) - pngname+='_'+plot+'.png' - - # Generate figure title - figtitle = 'MS: '+msname.rstrip('/') - figtitle += '\nColumn: '+column+', '+plot - figtitle += '\nBaseline: '+str(baseline[0])+'-'+str(baseline[1]) - figtitle += ' ['+blength+' km]' - figtitle+='\nField: '+fieldname - - # Create the figure - fig = pylab.figure(figsize=(20,15)) - t = fig.text(0.5, 0.945, figtitle,horizontalalignment='center',color='blue') - - # A panel for each corr product - for k in range(0,n_corr): - if plot == 'phase': - plotdata = numpy.angle(baselinedata[:,:,k]) # radians - elif plot == 'real': - plotdata = baselinedata[:,:,k].real - elif plot == 'imag': - plotdata = baselinedata[:,:,k].imag - else: - plotdata = numpy.absolute(baselinedata[:,:,k]) - - flagimage = pylab.cm.gray(plotdata*0.0) - flagimage[:,:,3] = (flagdata[:,:,k]) - - ax = fig.add_subplot(1,n_corr,k+1) - ax.set_xlabel('Channel number') - if k==0: - ax.set_ylabel('Time slot') - elif k==n_corr-1: - ax.yaxis.tick_right() - ax.yaxis.set_label_position('right') - ax.set_ylabel('Time slot') - else: - for ytick_i in ax.get_yticklabels(): - ytick_i.set_visible(False) - if plot != 'phase' and len(plotdata)>0: - if scale != -1: - immax = scale*plotdata.max() - immin = scale*plotdata.min() - else: - imstd = numpy.std(plotdata) - immax = 5.0*imstd - immin = 0.0 - ax.imshow(plotdata,aspect='auto',clim=(immin,immax),cmap=mycmap) - if not noflags: - ax.imshow(flagimage,aspect='auto',interpolation='nearest') - elif len(plotdata)>0: - ax.imshow(plotdata,aspect='auto',cmap=mycmap) - if not noflags: - ax.imshow(flagimage,aspect='auto',interpolation='nearest') - else: - ax.imshow(((0,0),(0,0)),aspect='auto') - - ax.set_title('Corr product '+str(k)) - - print(' Corr product:',k,' Data min,max:',plotdata.min(),plotdata.max()) - - for o in fig.findobj(matplotlib.text.Text): - o.set_fontsize('11') - - fig.tight_layout(w_pad=0.98,h_pad=0.98,rect=[0.02,0.02,0.95,0.95]) - - pylab.savefig(pngname) - pylab.close() + (options,args) = create_parser().parse_args() + + dolist = options.dolist + column = options.column + fieldid = int(options.field) + myspw = options.myspw + plot = options.plot + antenna1 = int(options.antenna1) + antenna2 = int(options.antenna2) + noflags = options.noflags + scale = float(options.scale) + mycmap = options.mycmap + foldername = options.foldername + + # Some error trapping + + if len(args) != 1: + ri('Please specify a single Measurement Set to plot.') + sys.exit(-1) + else: + msname = args[0].rstrip('/') + + if plot not in ['amp','phase','real','imag']: + ri('Requested plot not valid, must be one of amp, phase, real or imag.') + sys.exit(-1) + + + # MEASUREMENT SET INFO + + fieldtab = pyrap.tables.table(msname+'/FIELD') + sourceids = fieldtab.getcol('SOURCE_ID') + sourcenames = fieldtab.getcol('NAME') + fieldtab.done() + + spwtab = pyrap.tables.table(msname+'/SPECTRAL_WINDOW') + nspw = len(spwtab) + spwfreqs = spwtab.getcol('REF_FREQUENCY') + chanwidth = spwtab.getcol('CHAN_WIDTH')[0][0] # probably needs changing if SPWs have different widths + nchans = spwtab.getcol('NUM_CHAN') + spwtab.done() + + anttab = pyrap.tables.table(msname+'/ANTENNA') + nant = len(anttab) + antpos = anttab.getcol('POSITION') + antnames = anttab.getcol('NAME') + anttab.done() + + tt = pyrap.tables.table(msname) + usedants = numpy.unique(tt.getcol('ANTENNA1')) + + + # PRINT SUMMARY + + if dolist: + print('') + gi(' '+msname+'/FIELD') + gi(' ROW ID NAME') + for i in range(0,len(sourceids)): + print(' %-6s%-14s%-14s' % (i,sourceids[i],sourcenames[i])) + print('') + gi(' '+msname+'/SPECTRAL_WINDOW') + gi(' ROW CHANS WIDTH[MHz] REF_FREQ[MHz]') + for i in range(0,nspw): + print(' %-6s%-14s%-20s%-14s' % (i,str(nchans[i]),str(chanwidth/1e6),str(spwfreqs[i]/1e6))) + print('') + gi(' '+msname+'/ANTENNA') + gi(' ROW NAME POSITION') + for i in range(0,nant): + if i in usedants: + print(' %-6s%-14s%-14s' % (i,(antnames[i]),str(antpos[i]))) + else: + ri(' %-6s%-14s%-14s' % (i,(antnames[i]),str(antpos[i]))) + + print('') + tt.done() + + + # OR ELSE DO THE PLOTS + + + else: + # Create output folder if it doesn't exist + if foldername == '': + foldername = msname+'_'+column+'__plots' + if os.path.isdir(foldername): + print('Found',foldername) + else: + print('Creating',foldername) + os.system('mkdir '+foldername) + + # Make a complete list of SPWs if one isn't provided + if myspw == '': + myspw = numpy.arange(0,nspw) + else: + myspw = myspw.split(',') + + fieldname = sourcenames[fieldid] + + # Make a list of baseline pairs based on the antenna selections + baselines = [] + if antenna1 != -1 and antenna2 != -1: + baselines = [(antenna1,antenna2)] + elif antenna1 != -1: + i = antenna1 + for j in usedants: + if i != j: + pair = [i,j] + pair = sorted(pair) + if pair not in baselines: + if antenna1 != -1: + if antenna1 in pair: + baselines.append(pair) + else: + baselines.append(pair) + else: + for i in usedants: + for j in usedants: + if i != j: + pair = [i,j] + pair = sorted(pair) + if pair not in baselines: + if antenna1 != -1: + if antenna1 in pair: + baselines.append(pair) + else: + baselines.append(pair) + + # Loop over baselines + for baseline in baselines: + + # Determine unprojected baseline length + ap1 = antpos[baseline[0]] + ap2 = antpos[baseline[1]] + blength = (((ap1[0]-ap2[0])**2.0)+((ap1[1]-ap2[1])**2.0)+((ap1[2]-ap1[2])**2.0))**0.5 + blength = str(round(blength/1000.0,2)) + + print('Plotting baseline:',baseline,' Deprojected length:',blength,'km') + + # Get the data + datacols = [] + flagcols = [] + # Loop over SPWs + print('SPW:') + for spw in myspw: + print(spw) + subtab = tt.query(query='ANTENNA1=='+str(baseline[0]) + +' && ANTENNA2=='+str(baseline[1]) + +' && DATA_DESC_ID=='+str(spw) + +' && FIELD_ID=='+str(fieldid)) + datacol = subtab.getcol(column) + flagcol = subtab.getcol('FLAG') + datacols.append(datacol) + flagcols.append(flagcol) + print('') + # Reshape the data + baselinedata = datacols[0] + flagdata = flagcols[0] + for p in range(1,len(datacols)): + baselinedata = numpy.concatenate((baselinedata,datacols[p]),axis=1) + flagdata = numpy.concatenate((flagdata,flagcols[p]),axis=1) + + # Get number of corr products from the data shape + n_corr = baselinedata.shape[2] + + # Generate png name + pngname = foldername+'/'+msname.split('/')[-1].rstrip('/')+'_baseline_'+str(baseline[0])+'_'+str(baseline[1]) + pngname+='_field'+str(fieldid) + pngname+='_'+plot+'.png' + + # Generate figure title + figtitle = 'MS: '+msname.rstrip('/') + figtitle += '\nColumn: '+column+', '+plot + figtitle += '\nBaseline: '+str(baseline[0])+'-'+str(baseline[1]) + figtitle += ' ['+blength+' km]' + figtitle+='\nField: '+fieldname + + # Create the figure + fig = pylab.figure(figsize=(20,15)) + t = fig.text(0.5, 0.945, figtitle,horizontalalignment='center',color='blue') + + # A panel for each corr product + for k in range(0,n_corr): + if plot == 'phase': + plotdata = numpy.angle(baselinedata[:,:,k]) # radians + elif plot == 'real': + plotdata = baselinedata[:,:,k].real + elif plot == 'imag': + plotdata = baselinedata[:,:,k].imag + else: + plotdata = numpy.absolute(baselinedata[:,:,k]) + + flagimage = pylab.cm.gray(plotdata*0.0) + flagimage[:,:,3] = (flagdata[:,:,k]) + + ax = fig.add_subplot(1,n_corr,k+1) + ax.set_xlabel('Channel number') + if k==0: + ax.set_ylabel('Time slot') + elif k==n_corr-1: + ax.yaxis.tick_right() + ax.yaxis.set_label_position('right') + ax.set_ylabel('Time slot') + else: + for ytick_i in ax.get_yticklabels(): + ytick_i.set_visible(False) + if plot != 'phase' and len(plotdata)>0: + if scale != -1: + immax = scale*plotdata.max() + immin = scale*plotdata.min() + else: + imstd = numpy.std(plotdata) + immax = 5.0*imstd + immin = 0.0 + ax.imshow(plotdata,aspect='auto',clim=(immin,immax),cmap=mycmap) + if not noflags: + ax.imshow(flagimage,aspect='auto',interpolation='nearest') + elif len(plotdata)>0: + ax.imshow(plotdata,aspect='auto',cmap=mycmap) + if not noflags: + ax.imshow(flagimage,aspect='auto',interpolation='nearest') + else: + ax.imshow(((0,0),(0,0)),aspect='auto') + + ax.set_title('Corr product '+str(k)) + + print(' Corr product:',k,' Data min,max:',plotdata.min(),plotdata.max()) + + for o in fig.findobj(matplotlib.text.Text): + o.set_fontsize('11') + + fig.tight_layout(w_pad=0.98,h_pad=0.98,rect=[0.02,0.02,0.95,0.95]) + + pylab.savefig(pngname) + pylab.close() From 1b2a3bb4741a0778106c49daab303811da6b88a2 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Tue, 14 Jun 2022 10:47:49 +0200 Subject: [PATCH 03/15] flag above and unflag below --- surfvis/flagchi2.py | 20 +++++++++++++++----- surfvis/utils.py | 11 ++++++----- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/surfvis/flagchi2.py b/surfvis/flagchi2.py index 64cacfa..675b74b 100644 --- a/surfvis/flagchi2.py +++ b/surfvis/flagchi2.py @@ -26,8 +26,10 @@ def create_parser(): help='Weight column (default = WEIGHT_SPECTRUM)') parser.add_option('--fcol', default='FLAG', help='Flag column (default = FLAG)') - parser.add_option('--sigma', default=3, type=float, - help='chisq threshold (default = 3)') + parser.add_option('--flag-above', default=3, type=float, + help='flag data with chisq above this value (default = 3)') + parser.add_option('--unflag-below', default=1.15, type=float, + help='unflag data with chisq below this value (default = 1.15)') parser.add_option('--nthreads', default=4, type=int, help='Number of dask threads to use') parser.add_option('--nrows', default=10000, type=int, @@ -51,7 +53,7 @@ def main(): schema = {} schema[options.rcol] = {'dims': ('chan', 'corr')} schema[options.fcol] = {'dims': ('chan', 'corr')} - columns = [options.rcol, options.fcol] + columns = [options.rcol, options.fcol, 'FLAG_ROW'] if options.scol is not None: schema[options.scol] = {'dims': ('chan', 'corr')} columns.append(options.scol) @@ -86,12 +88,20 @@ def main(): weight = ds.get(options.wcol).data flag = ds.get(options.fcol).data - uflag = flagchisq(resid, weight, flag, tuple(use_corrs), sigma=options.sigma) + uflag = flagchisq(resid, weight, flag, tuple(use_corrs), + flag_above=options.flag_above, + unflag_below=options.unflag_below) out_ds = ds.assign(**{options.fcol: (("row", "chan", "corr"), uflag)}) + + # update FLAG_ROW + flag_row = da.all(uflag.rechunk({'1':-1, '2':-1}), axis=(1,2)) + + out_ds = ds.assign(**{'FLAG_ROW': (("row",), flag_row)}) + out_data.append(out_ds) - writes = xds_to_table(out_data, msname, columns=[options.fcol]) + writes = xds_to_table(out_data, msname, columns=[options.fcol, 'FLAG_ROW']) with ProgressBar(): dask.compute(writes) diff --git a/surfvis/utils.py b/surfvis/utils.py index 6a4fb46..5ccf99b 100644 --- a/surfvis/utils.py +++ b/surfvis/utils.py @@ -95,20 +95,21 @@ def _surfchisq(resid, weight, flag, ant1, ant2, return out -def flagchisq(resid, weight, flag, use_corrs, sigma=25): +def flagchisq(resid, weight, flag, use_corrs, flag_above, unflag_below): res = da.blockwise(_flagchisq, 'rfc', resid, 'rfc', weight, 'rfc', flag, 'rfc', use_corrs, None, - sigma, None, + flag_above, None, + unflag_below, None, dtype=bool) return res @njit(fastmath=True, nogil=True) -def _flagchisq(resid, weight, flag, use_corrs, sigma): +def _flagchisq(resid, weight, flag, use_corrs, flag_above, unflag_below): nrow, nchan, ncorr = resid.shape sigmasq = sigma**2 for r in range(nrow): @@ -117,8 +118,8 @@ def _flagchisq(resid, weight, flag, use_corrs, sigma): res = resid[r, f, c] w = weight[r, f, c] chi2 = (np.conj(res) * w * res).real - if chi2 > sigmasq or chi2 == 0: + if chi2 > flag_above or chi2 == 0: flag[r, f, c] = True - else: + elif chi2 <= unflag_below : flag[r, f, c] = False return flag From 58491f883e4374ec669382abe004dc582be63f30 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Tue, 14 Jun 2022 12:31:49 +0200 Subject: [PATCH 04/15] pass int not str to chunk --- surfvis/flagchi2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/surfvis/flagchi2.py b/surfvis/flagchi2.py index 675b74b..fd34b7d 100644 --- a/surfvis/flagchi2.py +++ b/surfvis/flagchi2.py @@ -95,7 +95,7 @@ def main(): out_ds = ds.assign(**{options.fcol: (("row", "chan", "corr"), uflag)}) # update FLAG_ROW - flag_row = da.all(uflag.rechunk({'1':-1, '2':-1}), axis=(1,2)) + flag_row = da.all(uflag.rechunk({1:-1, 2:-1}), axis=(1,2)) out_ds = ds.assign(**{'FLAG_ROW': (("row",), flag_row)}) From d32fe51cfd99f940da46165cb2306c70c9598a34 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Tue, 14 Jun 2022 12:35:33 +0200 Subject: [PATCH 05/15] remove unpassed sigma --- surfvis/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/surfvis/utils.py b/surfvis/utils.py index 5ccf99b..3074a17 100644 --- a/surfvis/utils.py +++ b/surfvis/utils.py @@ -111,7 +111,6 @@ def flagchisq(resid, weight, flag, use_corrs, flag_above, unflag_below): @njit(fastmath=True, nogil=True) def _flagchisq(resid, weight, flag, use_corrs, flag_above, unflag_below): nrow, nchan, ncorr = resid.shape - sigmasq = sigma**2 for r in range(nrow): for f in range(nchan): for c in use_corrs: From 4b7be78ea2484a8f7a3105f4b3f787cc6481da93 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Thu, 16 Jun 2022 11:37:47 +0200 Subject: [PATCH 06/15] Remove scol in favour of special wcol keyword SIGMA_SPECTRUM' --- surfvis/flagchi2.py | 21 +++++++-------------- surfvis/surfchi2.py | 15 ++++++++++++--- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/surfvis/flagchi2.py b/surfvis/flagchi2.py index fd34b7d..3e8337e 100644 --- a/surfvis/flagchi2.py +++ b/surfvis/flagchi2.py @@ -20,10 +20,10 @@ def create_parser(): parser = OptionParser(usage='%prog [options] msname') parser.add_option('--rcol', default='RESIDUAL', help='Residual column (default = RESIDUAL)') - parser.add_option('--scol', - help='sigma column to use. Overides wcol if provided.') parser.add_option('--wcol', default='WEIGHT_SPECTRUM', - help='Weight column (default = WEIGHT_SPECTRUM)') + help='Weight column (default = WEIGHT_SPECTRUM). ' + 'The special value SIGMA_SPECTRUM can be passed to ' + 'initialise the weights as 1/sigma**2') parser.add_option('--fcol', default='FLAG', help='Flag column (default = FLAG)') parser.add_option('--flag-above', default=3, type=float, @@ -53,14 +53,8 @@ def main(): schema = {} schema[options.rcol] = {'dims': ('chan', 'corr')} schema[options.fcol] = {'dims': ('chan', 'corr')} - columns = [options.rcol, options.fcol, 'FLAG_ROW'] - if options.scol is not None: - schema[options.scol] = {'dims': ('chan', 'corr')} - columns.append(options.scol) - else: - schema[options.wcol] = {'dims': ('chan', 'corr')} - columns.append(options.wcol) - + schema[options.wcol] = {'dims': ('chan', 'corr')} + columns = [options.wcol, options.rcol, options.fcol, 'FLAG_ROW'] xds = xds_from_ms(msname, columns=columns, @@ -81,9 +75,8 @@ def main(): out_data = [] for i, ds in enumerate(xds): resid = ds.get(options.rcol).data - if options.scol is not None: - sigma = ds.get(options.scol).data - weight = 1/sigma**2 + if options.wcol == 'SIGMA_SPECTRUM': + weight = 1.0/ds.get(options.wcol).data**2 else: weight = ds.get(options.wcol).data flag = ds.get(options.fcol).data diff --git a/surfvis/surfchi2.py b/surfvis/surfchi2.py index 0e99920..be7798b 100755 --- a/surfvis/surfchi2.py +++ b/surfvis/surfchi2.py @@ -25,14 +25,16 @@ def create_parser(): parser.add_option('--rcol', default='RESIDUAL', help='Residual column (default = RESIDUAL)') parser.add_option('--wcol', default='WEIGHT_SPECTRUM', - help='Weight column (default = WEIGHT_SPECTRUM)') + help='Weight column (default = WEIGHT_SPECTRUM). ' + 'The special value SIGMA_SPECTRUM can be passed to ' + 'initialise the weights as 1/sigma**2') parser.add_option('--fcol', default='FLAG', help='Flag column (default = FLAG)') parser.add_option('--dataout', default='', help='Output name of zarr dataset') - parser.add_option('--imagesout', default=None, + parser.add_option('--imagesout', default='', help='Output folder to place images. ' - 'If None (default) no plots are saved') + 'Saved in CWD/chi2 by default. ') parser.add_option('--nthreads', default=4, type=int, help='Number of dask threads to use') parser.add_option('--ntimes', default=-1, type=int, @@ -54,6 +56,13 @@ def main(): print(f"Removing existing {options.dataout} folder") os.system(f"rm -r {options.dataout}") + if options.imagesout == '': + options.imagesout = os.getcwd() + '/chi2' + + if os.path.isdir(options.imagesout): + print(f"Removing existing {options.imagesout} folder") + os.system(f"rm -r {options.imagesout}") + # Some error trapping if len(args) != 1: print('Please specify a single Measurement Set to plot.') From aecbd121e93576dff7608bc71f913684bfbffbb4 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Tue, 30 Aug 2022 16:42:17 +0200 Subject: [PATCH 07/15] Add option to leave certain antennas as is in flagchi2 (respect-ants) --- surfvis/flagchi2.py | 29 ++++++++++++++++++++++------- surfvis/surfchi2.py | 3 ++- surfvis/utils.py | 18 ++++++++++++++---- 3 files changed, 38 insertions(+), 12 deletions(-) diff --git a/surfvis/flagchi2.py b/surfvis/flagchi2.py index 3e8337e..f79f8e6 100644 --- a/surfvis/flagchi2.py +++ b/surfvis/flagchi2.py @@ -13,7 +13,9 @@ import dask.array as da from dask.diagnostics import ProgressBar from surfvis.utils import flagchisq -from daskms import xds_from_ms, xds_from_table, xds_to_table +from daskms import xds_from_storage_ms as xds_from_ms +from daskms import xds_from_storage_table as xds_from_table +from daskms import xds_to_storage_table as xds_to_table def create_parser(): @@ -38,6 +40,8 @@ def create_parser(): help='Number of frequencies in a chunk (default=128)') parser.add_option("--use-corrs", type=str, help='Comma seprated list of correlations to use (do not use spaces)') + parser.add_option("--respect-ants", type=str, + help='Comma seprated list of antennas to respect (do not use spaces)') return parser def main(): @@ -52,12 +56,12 @@ def main(): schema = {} schema[options.rcol] = {'dims': ('chan', 'corr')} - schema[options.fcol] = {'dims': ('chan', 'corr')} schema[options.wcol] = {'dims': ('chan', 'corr')} - columns = [options.wcol, options.rcol, options.fcol, 'FLAG_ROW'] + schema[options.fcol] = {'dims': ('chan', 'corr')} xds = xds_from_ms(msname, - columns=columns, + columns=[options.rcol, options.wcol, options.fcol, + 'ANTENNA1', 'ANTENNA2'], chunks={'row': options.nrows, 'chan': options.nfreqs}, group_cols=['FIELD_ID', 'DATA_DESC_ID', 'SCAN_NUMBER'], table_schema=schema) @@ -69,9 +73,14 @@ def main(): else: use_corrs = [0] else: - use_corrs = list(map(int, options.use_corrs.split(','))) + use_corrs = tuple(map(int, options.use_corrs.split(','))) print(f"Using correlations {use_corrs}") + if options.respect_ants is not None: + rants = list(map(int, options.respect_ants.split(','))) + else: + rants = [] + out_data = [] for i, ds in enumerate(xds): resid = ds.get(options.rcol).data @@ -80,10 +89,16 @@ def main(): else: weight = ds.get(options.wcol).data flag = ds.get(options.fcol).data + ant1 = ds.ANTENNA1.data + ant2 = ds.ANTENNA2.data + + # import pdb; pdb.set_trace() - uflag = flagchisq(resid, weight, flag, tuple(use_corrs), + uflag = flagchisq(resid, weight, flag, ant1, ant2, + use_corrs=tuple(use_corrs), flag_above=options.flag_above, - unflag_below=options.unflag_below) + unflag_below=options.unflag_below, + respect_ants=tuple(rants)) out_ds = ds.assign(**{options.fcol: (("row", "chan", "corr"), uflag)}) diff --git a/surfvis/surfchi2.py b/surfvis/surfchi2.py index be7798b..b1fdf72 100755 --- a/surfvis/surfchi2.py +++ b/surfvis/surfchi2.py @@ -13,7 +13,8 @@ import dask.array as da from dask.diagnostics import ProgressBar from surfvis.utils import surfchisq -from daskms import xds_from_ms, xds_from_table +from daskms import xds_from_storage_ms as xds_from_ms +from daskms import xds_from_storage_table as xds_from_table from daskms.experimental.zarr import xds_from_zarr, xds_to_zarr # might make for cooler histograms but doesn't work out of the box from astropy.visualization import hist diff --git a/surfvis/utils.py b/surfvis/utils.py index 3074a17..0313cb9 100644 --- a/surfvis/utils.py +++ b/surfvis/utils.py @@ -22,7 +22,7 @@ def surf(p, q, gp, gq, data, resid, weight, flag, basename): "is\n%s\n" % traceback_str) def surfchisq(resid, weight, flag, ant1, ant2, - rbin_idx, rbin_counts, fbin_idx, fbin_counts): + rbin_idx, rbin_counts, fbin_idx, fbin_counts): nant = da.maximum(ant1.max(), ant2.max()).compute() + 1 res = da.blockwise(_surfchisq, 'tfcpq2', @@ -45,7 +45,7 @@ def surfchisq(resid, weight, flag, ant1, ant2, @njit(fastmath=True, nogil=True) def _surfchisq(resid, weight, flag, ant1, ant2, - rbin_idx, rbin_counts, fbin_idx, fbin_counts): + rbin_idx, rbin_counts, fbin_idx, fbin_counts): nrow, nchan, ncorr = resid.shape nto = rbin_idx.size @@ -95,23 +95,33 @@ def _surfchisq(resid, weight, flag, ant1, ant2, return out -def flagchisq(resid, weight, flag, use_corrs, flag_above, unflag_below): +def flagchisq(resid, weight, flag, ant1, ant2, + use_corrs=(), flag_above=5, unflag_below=1, + respect_ants=()): + + # import pdb; pdb.set_trace() res = da.blockwise(_flagchisq, 'rfc', resid, 'rfc', weight, 'rfc', flag, 'rfc', + ant1, 'r', + ant2, 'r', use_corrs, None, flag_above, None, unflag_below, None, + respect_ants, None, dtype=bool) return res @njit(fastmath=True, nogil=True) -def _flagchisq(resid, weight, flag, use_corrs, flag_above, unflag_below): +def _flagchisq(resid, weight, flag, ant1, ant2, + use_corrs, flag_above, unflag_below, respect_ants): nrow, nchan, ncorr = resid.shape for r in range(nrow): + if ant1[r] in respect_ants or ant2[r] in respect_ants: + continue for f in range(nchan): for c in use_corrs: res = resid[r, f, c] From de753a67275e2d8b1b0de8e8f981a5e6e14288c8 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Tue, 6 Sep 2022 10:56:23 +0200 Subject: [PATCH 08/15] Add option to plot auto-correlations --- surfvis/surfvis.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/surfvis/surfvis.py b/surfvis/surfvis.py index ef2d6d1..0924132 100755 --- a/surfvis/surfvis.py +++ b/surfvis/surfvis.py @@ -31,6 +31,7 @@ def create_parser(): parser.add_option('-i','--i',dest='antenna1',help='Antenna 1: plot only this antenna',default=-1) parser.add_option('-j','--j',dest='antenna2',help='Antenna 2: use with -i to plot a single baseline',default=-1) parser.add_option('--noflags',dest='noflags',help='Disable flagged data overlay',action='store_true',default=False) + parser.add_option('--doacorr',dest='noflags',help='Disable flagged data overlay',action='store_true',default=False) parser.add_option('--scale',dest='scale',help='Scale the image peak to this multiple of the per-corr min/max (default = scale image max to 5 sigma, this parameter is ignored for phase plots)',default=-1) parser.add_option('--cmap',dest='mycmap',help='Matplotlib colour map to use (default = jet)',default='jet') parser.add_option('-o','--opdir',dest='foldername',help='Output folder to store plots (default = msname___plots)',default='') @@ -142,7 +143,7 @@ def main(): elif antenna1 != -1: i = antenna1 for j in usedants: - if i != j: + if i != j or options.doacorr: pair = [i,j] pair = sorted(pair) if pair not in baselines: @@ -154,7 +155,7 @@ def main(): else: for i in usedants: for j in usedants: - if i != j: + if i != j or options.doacorr: pair = [i,j] pair = sorted(pair) if pair not in baselines: From efc27c988c4a035f3925fdb7c9eb47e7506d37a8 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Tue, 6 Sep 2022 11:15:42 +0200 Subject: [PATCH 09/15] Document doacorr option and correct dest --- surfvis/surfvis.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/surfvis/surfvis.py b/surfvis/surfvis.py index 0924132..4242cca 100755 --- a/surfvis/surfvis.py +++ b/surfvis/surfvis.py @@ -31,7 +31,7 @@ def create_parser(): parser.add_option('-i','--i',dest='antenna1',help='Antenna 1: plot only this antenna',default=-1) parser.add_option('-j','--j',dest='antenna2',help='Antenna 2: use with -i to plot a single baseline',default=-1) parser.add_option('--noflags',dest='noflags',help='Disable flagged data overlay',action='store_true',default=False) - parser.add_option('--doacorr',dest='noflags',help='Disable flagged data overlay',action='store_true',default=False) + parser.add_option('--doacorr',dest='doacorr',help='Plots auto-correlations',action='store_true',default=False) parser.add_option('--scale',dest='scale',help='Scale the image peak to this multiple of the per-corr min/max (default = scale image max to 5 sigma, this parameter is ignored for phase plots)',default=-1) parser.add_option('--cmap',dest='mycmap',help='Matplotlib colour map to use (default = jet)',default='jet') parser.add_option('-o','--opdir',dest='foldername',help='Output folder to store plots (default = msname___plots)',default='') @@ -48,6 +48,7 @@ def main(): antenna1 = int(options.antenna1) antenna2 = int(options.antenna2) noflags = options.noflags + doacorr = options.doacorr scale = float(options.scale) mycmap = options.mycmap foldername = options.foldername @@ -143,7 +144,7 @@ def main(): elif antenna1 != -1: i = antenna1 for j in usedants: - if i != j or options.doacorr: + if i != j or doacorr: pair = [i,j] pair = sorted(pair) if pair not in baselines: @@ -155,7 +156,7 @@ def main(): else: for i in usedants: for j in usedants: - if i != j or options.doacorr: + if i != j or doacorr: pair = [i,j] pair = sorted(pair) if pair not in baselines: From 6f2fc9ad4285b049565ee0e0cfd3b9c8b54cf742 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Sun, 16 Oct 2022 15:56:55 +0200 Subject: [PATCH 10/15] rechunk in xds_to_table --- surfvis/flagchi2.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/surfvis/flagchi2.py b/surfvis/flagchi2.py index f79f8e6..1f6db13 100644 --- a/surfvis/flagchi2.py +++ b/surfvis/flagchi2.py @@ -109,7 +109,9 @@ def main(): out_data.append(out_ds) - writes = xds_to_table(out_data, msname, columns=[options.fcol, 'FLAG_ROW']) + writes = xds_to_table(out_data, msname, + columns=[options.fcol, 'FLAG_ROW'], + rechunk=True) with ProgressBar(): dask.compute(writes) From a129c2e000dac31648e19849bad2f9739783f3c4 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Sun, 16 Oct 2022 16:44:49 +0200 Subject: [PATCH 11/15] set nthreads --- surfvis/flagchi2.py | 7 +++++-- surfvis/surfchi2.py | 3 +++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/surfvis/flagchi2.py b/surfvis/flagchi2.py index 1f6db13..6d68df9 100644 --- a/surfvis/flagchi2.py +++ b/surfvis/flagchi2.py @@ -34,9 +34,9 @@ def create_parser(): help='unflag data with chisq below this value (default = 1.15)') parser.add_option('--nthreads', default=4, type=int, help='Number of dask threads to use') - parser.add_option('--nrows', default=10000, type=int, + parser.add_option('--nrows', default=100000, type=int, help='Number of rows in each chunk (default=10000)') - parser.add_option('--nfreqs', default=128, type=int, + parser.add_option('--nfreqs', default=256, type=int, help='Number of frequencies in a chunk (default=128)') parser.add_option("--use-corrs", type=str, help='Comma seprated list of correlations to use (do not use spaces)') @@ -54,6 +54,9 @@ def main(): else: msname = args[0].rstrip('/') + from multiprocessing.pool import ThreadPool + dask.config.set(pool=ThreadPool(options.nthreads)) + schema = {} schema[options.rcol] = {'dims': ('chan', 'corr')} schema[options.wcol] = {'dims': ('chan', 'corr')} diff --git a/surfvis/surfchi2.py b/surfvis/surfchi2.py index b1fdf72..1c98e32 100755 --- a/surfvis/surfchi2.py +++ b/surfvis/surfchi2.py @@ -71,6 +71,9 @@ def main(): else: msname = args[0].rstrip('/') + from multiprocessing.pool import ThreadPool + dask.config.set(pool=ThreadPool(options.nthreads)) + # chunking info schema = {} schema[options.fcol] = {'dims': ('chan', 'corr')} From d2e7f4220e8626535613c43394a51d93127eaf4b Mon Sep 17 00:00:00 2001 From: landmanbester Date: Mon, 17 Oct 2022 10:55:01 +0200 Subject: [PATCH 12/15] remove pdb import --- surfvis/flagchi2.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/surfvis/flagchi2.py b/surfvis/flagchi2.py index 6d68df9..b20c97a 100644 --- a/surfvis/flagchi2.py +++ b/surfvis/flagchi2.py @@ -34,9 +34,9 @@ def create_parser(): help='unflag data with chisq below this value (default = 1.15)') parser.add_option('--nthreads', default=4, type=int, help='Number of dask threads to use') - parser.add_option('--nrows', default=100000, type=int, + parser.add_option('--nrows', default=250000, type=int, help='Number of rows in each chunk (default=10000)') - parser.add_option('--nfreqs', default=256, type=int, + parser.add_option('--nfreqs', default=512, type=int, help='Number of frequencies in a chunk (default=128)') parser.add_option("--use-corrs", type=str, help='Comma seprated list of correlations to use (do not use spaces)') @@ -95,8 +95,6 @@ def main(): ant1 = ds.ANTENNA1.data ant2 = ds.ANTENNA2.data - # import pdb; pdb.set_trace() - uflag = flagchisq(resid, weight, flag, ant1, ant2, use_corrs=tuple(use_corrs), flag_above=options.flag_above, From 57afce6de7d6c382c14c91c7affa518e0a866127 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Mon, 17 Oct 2022 11:44:39 +0200 Subject: [PATCH 13/15] reuse out_ds in assign --- surfvis/flagchi2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/surfvis/flagchi2.py b/surfvis/flagchi2.py index b20c97a..1778d59 100644 --- a/surfvis/flagchi2.py +++ b/surfvis/flagchi2.py @@ -106,7 +106,7 @@ def main(): # update FLAG_ROW flag_row = da.all(uflag.rechunk({1:-1, 2:-1}), axis=(1,2)) - out_ds = ds.assign(**{'FLAG_ROW': (("row",), flag_row)}) + out_ds = out_ds.assign(**{'FLAG_ROW': (("row",), flag_row)}) out_data.append(out_ds) From 0bfb155af707318cbee0d472b235506b1a3ac2d3 Mon Sep 17 00:00:00 2001 From: landmanbester Date: Tue, 23 May 2023 15:04:05 +0200 Subject: [PATCH 14/15] remove progressbar --- surfvis/flagchi2.py | 3 +-- surfvis/surfchi2.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/surfvis/flagchi2.py b/surfvis/flagchi2.py index 1778d59..0970858 100644 --- a/surfvis/flagchi2.py +++ b/surfvis/flagchi2.py @@ -114,5 +114,4 @@ def main(): columns=[options.fcol, 'FLAG_ROW'], rechunk=True) - with ProgressBar(): - dask.compute(writes) + dask.compute(writes) diff --git a/surfvis/surfchi2.py b/surfvis/surfchi2.py index 1c98e32..6cadd1e 100755 --- a/surfvis/surfchi2.py +++ b/surfvis/surfchi2.py @@ -204,8 +204,7 @@ def main(): idts.append(idt) - with ProgressBar(): - dask.compute(out_ds) + dask.compute(out_ds) # primitive plotting if options.imagesout is not None: From c528f14cd2257404037a290b828879d3b81afa3d Mon Sep 17 00:00:00 2001 From: landmanbester Date: Thu, 14 Sep 2023 13:41:25 +0200 Subject: [PATCH 15/15] bump python versions in test matrix 3.8-3.10 --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ed341b6..4967c71 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -7,7 +7,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.7", "3.8", "3.9"] + python-version: ["3.8", "3.9", "3.10"] steps: - name: Set up Python ${{ matrix.python-version }}