disko 1.3.0__tar.gz → 1.3.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (53) hide show
  1. {disko-1.3.0 → disko-1.3.3}/CHANGES.md +4 -1
  2. {disko-1.3.0 → disko-1.3.3}/Makefile +1 -1
  3. {disko-1.3.0 → disko-1.3.3}/PKG-INFO +1 -1
  4. disko-1.3.3/disko/bayes_cli.py +434 -0
  5. {disko-1.3.0 → disko-1.3.3}/pyproject.toml +1 -1
  6. {disko-1.3.0 → disko-1.3.3}/uv.lock +1 -1
  7. disko-1.3.0/disko/bayes_cli.py +0 -341
  8. {disko-1.3.0 → disko-1.3.3}/.github/workflows/package.yaml +0 -0
  9. {disko-1.3.0 → disko-1.3.3}/.github/workflows/testing.yaml +0 -0
  10. {disko-1.3.0 → disko-1.3.3}/.gitignore +0 -0
  11. {disko-1.3.0 → disko-1.3.3}/.travis.yml +0 -0
  12. {disko-1.3.0 → disko-1.3.3}/CITATION.cff +0 -0
  13. {disko-1.3.0 → disko-1.3.3}/LICENSE.txt +0 -0
  14. {disko-1.3.0 → disko-1.3.3}/MANIFEST.in +0 -0
  15. {disko-1.3.0 → disko-1.3.3}/README.md +0 -0
  16. {disko-1.3.0 → disko-1.3.3}/disko/__init__.py +0 -0
  17. {disko-1.3.0 → disko-1.3.3}/disko/cli.py +0 -0
  18. {disko-1.3.0 → disko-1.3.3}/disko/disko.py +0 -0
  19. {disko-1.3.0 → disko-1.3.3}/disko/draw_cli.py +0 -0
  20. {disko-1.3.0 → disko-1.3.3}/disko/draw_sky.py +0 -0
  21. {disko-1.3.0 → disko-1.3.3}/disko/fov/__init__.py +0 -0
  22. {disko-1.3.0 → disko-1.3.3}/disko/fov/factory.py +0 -0
  23. {disko-1.3.0 → disko-1.3.3}/disko/healpix_sphere.py +0 -0
  24. {disko-1.3.0 → disko-1.3.3}/disko/ms_helper.py +0 -0
  25. {disko-1.3.0 → disko-1.3.3}/disko/multivariate_gaussian.py +0 -0
  26. {disko-1.3.0 → disko-1.3.3}/disko/parser_support.py +0 -0
  27. {disko-1.3.0 → disko-1.3.3}/disko/projection_lsqr.py +0 -0
  28. {disko-1.3.0 → disko-1.3.3}/disko/resolution.py +0 -0
  29. {disko-1.3.0 → disko-1.3.3}/disko/rime.py +0 -0
  30. {disko-1.3.0 → disko-1.3.3}/disko/sphere.py +0 -0
  31. {disko-1.3.0 → disko-1.3.3}/disko/sphere_mesh.py +0 -0
  32. {disko-1.3.0 → disko-1.3.3}/disko/svd_cli.py +0 -0
  33. {disko-1.3.0 → disko-1.3.3}/disko/telescope_operator.py +0 -0
  34. {disko-1.3.0 → disko-1.3.3}/disko/tests/__init__.py +0 -0
  35. {disko-1.3.0 → disko-1.3.3}/disko/tests/test_disko.py +0 -0
  36. {disko-1.3.0 → disko-1.3.3}/disko/tests/test_meshsphere.py +0 -0
  37. {disko-1.3.0 → disko-1.3.3}/disko/tests/test_multivariate.py +0 -0
  38. {disko-1.3.0 → disko-1.3.3}/disko/tests/test_projection_lsqr.py +0 -0
  39. {disko-1.3.0 → disko-1.3.3}/disko/tests/test_pylops_operator.py +0 -0
  40. {disko-1.3.0 → disko-1.3.3}/disko/tests/test_sphere.py +0 -0
  41. {disko-1.3.0 → disko-1.3.3}/disko/tests/test_subsphere.py +0 -0
  42. {disko-1.3.0 → disko-1.3.3}/disko/tests/test_telescope_operator.py +0 -0
  43. {disko-1.3.0 → disko-1.3.3}/disko/util.py +0 -0
  44. {disko-1.3.0 → disko-1.3.3}/docker/Dockerfile +0 -0
  45. {disko-1.3.0 → disko-1.3.3}/docker/Makefile +0 -0
  46. {disko-1.3.0 → disko-1.3.3}/docker/docker-compose.yml +0 -0
  47. {disko-1.3.0 → disko-1.3.3}/install.sh +0 -0
  48. {disko-1.3.0 → disko-1.3.3}/movie/.gitignore +0 -0
  49. {disko-1.3.0 → disko-1.3.3}/movie/Makefile +0 -0
  50. {disko-1.3.0 → disko-1.3.3}/prof.py +0 -0
  51. {disko-1.3.0 → disko-1.3.3}/requirements.txt +0 -0
  52. {disko-1.3.0 → disko-1.3.3}/sequential_test.sh +0 -0
  53. {disko-1.3.0 → disko-1.3.3}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  # Changes
2
2
 
3
- ## 1.3.0 (2026-06-19)
3
+ ## 1.3.3 (2026-06-19)
4
4
 
5
5
  - Vectorize `get_harmonics`, `make_gamma`, and matrix-free operators (`DiSkOOperator`, `DirectImagingOperator`) using broadcasting and blocked BLAS-level operations (10-50x speedup)
6
6
  - Fix FISTA solver: remove unconditional `eps = 1e-9` override, switch from broken analysis formulation (`SOp=Apre`) to synthesis formulation, use proper initial guess (`abs(Apre @ d)`) instead of zeros
@@ -12,6 +12,9 @@
12
12
  - CI: target `disko/` directory in flake8 to avoid scanning dependencies
13
13
  - Configure flake8 exclusions for `.venv`, `.git`, `__pycache__` in `setup.cfg`
14
14
  - Fix various flake8 lint issues (unused imports, variables, long lines)
15
+ - Fix `disko_bayes` MS reading: replace broken `DiSkO.from_ms()` with `disko_from_ms()`, add visibility conjugation, geo-location, and file-exists check
16
+ - Fix `disko_bayes` output: remove invalid `fov` argument from `to_fits()` and `to_svg()` calls
17
+ - Fix publish workflow: use concrete Python version (`3.12` instead of `3.x`), pin `actions/checkout@v4`, fix `--outdir` → `--out-dir`
15
18
 
16
19
  ## 1.2.0 (2026-06-18)
17
20
 
@@ -31,7 +31,7 @@ svd:
31
31
 
32
32
  bayes:
33
33
  #rm -f *.npz
34
- uv run disko_bayes --fov 155 --ms test_data/test.ms --mu --PNG --SVG --arcmin=90 --dir test_out --title 'bayes_tart' --sigma-v=0.15
34
+ uv run disko_bayes --healpix --fov 155 --res 30arcmin --ms test_data/test.ms --mu --var --PNG --SVG --arcmin=90 --dir test_out --title 'bayes_tart' --sigma-v=0.15
35
35
  # ${TIME} uv run disko --fov 155 --ms ../tart2ms/test.ms --SVG --arcmin=120 --title 'tart' --tikhonov --alpha=0.01
36
36
 
37
37
  ngc1194:
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: disko
3
- Version: 1.3.0
3
+ Version: 1.3.3
4
4
  Summary: Discrete Sky Operator (DiSkO) Aperture Synthesis Radio Imaging
5
5
  Author-email: Tim Molteno <tim@elec.ac.nz>
6
6
  License: GPLv3
@@ -0,0 +1,434 @@
1
+ #!/usr/bin/env python
2
+ import argparse
3
+ import datetime
4
+ import json
5
+ import logging
6
+ import os
7
+ import time
8
+ from copy import deepcopy
9
+
10
+ import matplotlib
11
+ import matplotlib.pyplot as plt
12
+ import numpy as np
13
+ import scipy.special
14
+ from tart.imaging import calibration, elaz, visibility
15
+ from tart.operation import settings
16
+ from tart_tools import api_imaging
17
+
18
+ from .cli import disko_from_ms
19
+ from .disko import DiSkO, vis_to_real
20
+ from .ms_helper import get_array_location
21
+ from .multivariate_gaussian import MultivariateGaussian
22
+ from .parser_support import sphere_args_parser, sphere_from_args
23
+ from .telescope_operator import TelescopeOperator
24
+
25
+ logger = logging.getLogger(__name__)
26
+ logger.addHandler(
27
+ logging.NullHandler()
28
+ ) # Add other handlers if you're using this as a library
29
+ logger.setLevel(logging.INFO)
30
+
31
+
32
+ def create_prior(vis_arr, sphere, hdf_prior):
33
+ """Based on the size of the visibilities, try and calculate
34
+ what range the image should have.
35
+ """
36
+ if hdf_prior is not None:
37
+ return MultivariateGaussian.from_hdf5(hdf_prior)
38
+
39
+ vabs = np.abs(vis_arr)
40
+
41
+ p05, p50, p95, p100 = np.percentile(vabs, [5, 50, 95, 100])
42
+
43
+ var = p95 * p95
44
+ logger.info("Extimated Sky Prior variance={}".format(var))
45
+ prior = MultivariateGaussian(
46
+ np.zeros(sphere.npix) + p50, sigma=p95 * np.identity(sphere.npix)
47
+ )
48
+
49
+ return prior
50
+
51
+
52
+ def do_inference(disko, sphere, prior, sigma_v=None):
53
+ real_vis = vis_to_real(disko.vis_arr)
54
+
55
+ to = TelescopeOperator(disko, sphere)
56
+
57
+ # Transform to the natural basis.
58
+ n_prior = prior.linear_transform(to.Vh)
59
+ n_v = real_vis.shape[0]
60
+
61
+ # TODO create a proper covariance that ensures the real and imaginary components are linked.
62
+ if sigma_v is None:
63
+ diag = np.diag(disko.rms**2)
64
+ else:
65
+ diag = np.diag(np.ones(n_v // 2) * (sigma_v) ** 2)
66
+
67
+ logger.info(f"do_inference(sigma_v={diag[0, 0]})")
68
+
69
+ sigma_vis = np.block([[diag, 0.5 * diag], [0.5 * diag, diag]]) # .rechunk('auto')
70
+
71
+ # now invert sigma_vis
72
+ sigma_precision = MultivariateGaussian.sp_inv(sigma_vis)
73
+ del sigma_vis
74
+
75
+ if True:
76
+ prior_r = n_prior.block(0, to.rank)
77
+ prior_n = n_prior.block(to.rank, to.n_s)
78
+
79
+ A_r = to.A_r
80
+ V = to.V
81
+
82
+ del to
83
+ posterior_r = prior_r.bayes_update(sigma_precision, real_vis, A_r)
84
+ posterior_n = prior_n
85
+
86
+ del A_r
87
+ del sigma_precision
88
+ del prior_r
89
+ del prior_n
90
+ del n_prior
91
+
92
+ posterior = MultivariateGaussian.outer(posterior_r, posterior_n)
93
+
94
+ del posterior_r
95
+ del posterior_n
96
+
97
+ logger.info("Transforming posterior")
98
+
99
+ posterior = posterior.linear_transform(V)
100
+
101
+ del V
102
+ else:
103
+ posterior = to.sequential_inference(n_prior, real_vis, sigma_precision)
104
+ del to
105
+ del sigma_precision
106
+ del n_prior
107
+
108
+ return posterior
109
+
110
+
111
+ def handle_bayes(ARGS):
112
+
113
+ sphere = sphere_from_args(ARGS)
114
+
115
+ # Create a prior.
116
+ if ARGS.file:
117
+ logger.info("Getting Data from file: {}".format(ARGS.file))
118
+ # Load data from a JSON file
119
+ with open(ARGS.file, "r") as json_file:
120
+ calib_info = json.load(json_file)
121
+
122
+ info = calib_info["info"]
123
+ ant_pos = calib_info["ant_pos"]
124
+ config = settings.from_api_json(info["info"], ant_pos)
125
+
126
+ flag_list = [] # [4, 5, 14, 22]
127
+
128
+ original_positions = deepcopy(config.get_antenna_positions())
129
+
130
+ gains_json = calib_info["gains"]
131
+ gains = np.asarray(gains_json["gain"])
132
+ phase_offsets = np.asarray(gains_json["phase_offset"])
133
+ config = settings.from_api_json(info["info"], ant_pos)
134
+
135
+ measurements = []
136
+ for d in calib_info["data"]:
137
+ vis_json, source_json = d
138
+ cv, timestamp = api_imaging.vis_calibrated(
139
+ vis_json, config, gains, phase_offsets, flag_list
140
+ )
141
+ src_list = elaz.from_json(source_json, 0.0)
142
+
143
+ if ARGS.sigma_v is None:
144
+ raise RuntimeError(
145
+ "The --sigma-v option must be supplied when --file JSON input is used"
146
+ )
147
+
148
+ prior = create_prior(cv.v, sphere, ARGS.prior)
149
+ timestamp = cv.get_timestamp()
150
+ disko = DiSkO.from_cal_vis(cv)
151
+
152
+ posterior = do_inference(disko, sphere, prior, sigma_v=ARGS.sigma_v)
153
+ handle_output(ARGS, timestamp, posterior, sphere, disko)
154
+
155
+ elif ARGS.hdf:
156
+ logger.info(f"Getting data from file {ARGS.hdf}")
157
+ if ARGS.sigma_v is None:
158
+ raise RuntimeError(
159
+ "The --sigma-v option must be supplied when HDF5 input is used"
160
+ )
161
+
162
+ data = visibility.from_hdf5(ARGS.hdf)
163
+
164
+ prior = create_prior(data["vis_list"][0].v, sphere, ARGS.prior)
165
+ posterior = None
166
+
167
+ for v in data["vis_list"]:
168
+ if posterior is not None:
169
+ prior = posterior
170
+ cv = calibration.CalibratedVisibility(v)
171
+ cv.set_config(v.config)
172
+ cv.set_phase_offset(
173
+ list(range(cv.get_config().get_num_antenna())),
174
+ np.array(data["phase_offset"]),
175
+ )
176
+ cv.set_gain(
177
+ list(range(cv.get_config().get_num_antenna())), np.array(data["gain"])
178
+ )
179
+ timestamp = cv.get_timestamp()
180
+ disko = DiSkO.from_cal_vis(cv)
181
+
182
+ # TODO Calibrate the vis with gains and phases?
183
+ posterior = do_inference(disko, sphere, prior, sigma_v=ARGS.sigma_v)
184
+ handle_output(ARGS, timestamp, posterior, sphere, disko)
185
+ else:
186
+ logger.info("Getting Data from MS file: {}".format(ARGS.ms))
187
+
188
+ if not os.path.exists(ARGS.ms):
189
+ raise RuntimeError("Measurement set {} not found".format(ARGS.ms))
190
+
191
+ min_res = sphere.min_res()
192
+ logger.info(f"Min Res {min_res}")
193
+ disko = disko_from_ms(
194
+ ARGS.ms,
195
+ "DATA",
196
+ ARGS.nvis,
197
+ res=min_res,
198
+ channel=ARGS.channel,
199
+ field_id=ARGS.field,
200
+ )
201
+ # CASAcore UVW is conjugated; conjugate visibilities for consistency.
202
+ disko.vis_arr = disko.vis_arr.conjugate()
203
+
204
+ # Convert from reduced Julian Date to timestamp.
205
+ timestamp = disko.timestamp
206
+
207
+ json_info = get_array_location(ARGS.ms)
208
+ lat = json_info["lat"]
209
+ lon = json_info["lon"]
210
+ height = json_info["height"]
211
+ sphere.set_info(timestamp=timestamp, lon=lon, lat=lat, height=height)
212
+
213
+ src_list = None
214
+
215
+ prior = create_prior(disko.vis_arr, sphere, ARGS.prior)
216
+
217
+ posterior = do_inference(disko, sphere, prior, sigma_v=ARGS.sigma_v)
218
+ handle_output(ARGS, timestamp, posterior, sphere, disko)
219
+
220
+
221
+ def handle_output(ARGS, timestamp, posterior, sphere, disko=None):
222
+
223
+ if not ARGS.show_sources:
224
+ src_list = None
225
+
226
+ time_repr = "{:%Y_%m_%d_%H_%M_%S_%Z}".format(timestamp)
227
+
228
+ # Now save the files.
229
+ if ARGS.posterior is not None:
230
+ posterior.to_hdf5(ARGS.posterior)
231
+
232
+ def path(ending, image_title):
233
+ os.makedirs(ARGS.dir, exist_ok=True)
234
+ fname = "{}.{}".format(image_title, ending)
235
+ return os.path.join(ARGS.dir, fname)
236
+
237
+ def save_images(image_title, source_list):
238
+ # Save as a FITS file
239
+
240
+ if ARGS.FITS:
241
+ sphere.to_fits(fname=path("fits", image_title), info=disko.info)
242
+
243
+ if ARGS.SVG:
244
+ fname = path("svg", image_title)
245
+ sphere.to_svg(
246
+ fname=fname,
247
+ show_grid=True,
248
+ src_list=source_list,
249
+ title=image_title,
250
+ show_cbar=True,
251
+ )
252
+ logger.info("Generating {}".format(fname))
253
+ if ARGS.PNG:
254
+ fname = path("png", image_title)
255
+ sphere.plot(plt, source_list)
256
+ plt.title(image_title)
257
+ plt.tight_layout()
258
+ plt.savefig(fname, dpi=300)
259
+ plt.close()
260
+ logger.info("Generating {}".format(fname))
261
+ if ARGS.PDF:
262
+ fname = path("pdf", image_title)
263
+ sphere.plot(plt, source_list)
264
+ plt.title(image_title)
265
+ plt.savefig(fname, dpi=600)
266
+ plt.close()
267
+ logger.info("Generating {}".format(fname))
268
+
269
+ if ARGS.PDF or ARGS.PNG or ARGS.SVG or ARGS.FITS:
270
+ if ARGS.mu:
271
+ logger.info("Computing pixels")
272
+ tic = time.perf_counter()
273
+ # mu_positive = np.array(da.clip(posterior.mu, 0, None))
274
+ logger.info(f" Took {time.perf_counter() - tic:0.4f} seconds")
275
+ stat = sphere.set_visible_pixels(np.array(posterior.mu), scale=False)
276
+ stat["sigma-v"] = ARGS.sigma_v
277
+ logger.info(json.dumps(stat, sort_keys=True))
278
+ save_images("{}_{}_mu".format(ARGS.title, time_repr), source_list=src_list)
279
+
280
+ if ARGS.var:
281
+ tic = time.perf_counter()
282
+ logger.info("Computing variance...")
283
+ variance = np.array(posterior.variance())
284
+ logger.info(f" Took {time.perf_counter() - tic:0.4f} seconds")
285
+ sphere.set_visible_pixels(variance, scale=False)
286
+ save_images("{}_{}_var".format(ARGS.title, time_repr), source_list=None)
287
+
288
+ if ARGS.pcf:
289
+ tic = time.perf_counter()
290
+ logger.info("Computing point covariance...")
291
+
292
+ brightest_pixel = np.argmax(posterior.mu)
293
+ pix_cov = np.array(posterior.sigma()[brightest_pixel, :])
294
+ logger.info(f" Took {time.perf_counter() - tic:0.4f} seconds")
295
+
296
+ sphere.set_visible_pixels(pix_cov, scale=False)
297
+ save_images("{}_{}_pcf".format(ARGS.title, time_repr), source_list=None)
298
+
299
+ for i in range(ARGS.nsamples):
300
+ sphere.set_visible_pixels(posterior.sample(), scale=False)
301
+ save_images(
302
+ image_title="{}_{}_s{:0>5}".format(ARGS.title, time_repr, i),
303
+ source_list=None,
304
+ )
305
+
306
+
307
+ def main():
308
+ np.random.seed(42)
309
+ sphere_parsers = sphere_args_parser()
310
+
311
+ parser = argparse.ArgumentParser(
312
+ description="DiSkO: Bayesian inference of a posterior sky",
313
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter,
314
+ parents=sphere_parsers,
315
+ )
316
+
317
+ parser.add_argument(
318
+ "--hdf", required=False, default=None, help="Exported Multi-visibility file"
319
+ )
320
+ parser.add_argument("--ms", required=False, default=None, help="visibility file")
321
+ parser.add_argument(
322
+ "--file",
323
+ required=False,
324
+ default=None,
325
+ help="Snapshot observation saved JSON file (visiblities, positions and more).",
326
+ )
327
+
328
+ parser.add_argument(
329
+ "--channel", type=int, default=0, help="Use this frequency channel."
330
+ )
331
+ parser.add_argument(
332
+ "--field",
333
+ type=int,
334
+ default=0,
335
+ help="Use this FIELD_ID from the measurement set.",
336
+ )
337
+
338
+ parser.add_argument("--dir", required=False, default=".", help="Output directory.")
339
+ parser.add_argument(
340
+ "--nvis", type=int, default=1000, help="Number of visibilities to use."
341
+ )
342
+ parser.add_argument(
343
+ "--arcmin",
344
+ type=float,
345
+ default=None,
346
+ help="Highest allowed res of the sky in arc minutes.",
347
+ )
348
+
349
+ parser.add_argument(
350
+ "--sigma-v",
351
+ type=float,
352
+ default=None,
353
+ help="Diagonal components of the visibility covariance. If not supplied use measurement set values",
354
+ )
355
+
356
+ parser.add_argument(
357
+ "--PNG", action="store_true", help="Generate a PNG format image."
358
+ )
359
+ parser.add_argument(
360
+ "--PDF", action="store_true", help="Generate a PDF format image."
361
+ )
362
+ parser.add_argument(
363
+ "--SVG", action="store_true", help="Generate a SVG format image."
364
+ )
365
+ parser.add_argument(
366
+ "--FITS", action="store_true", help="Generate a FITS format image."
367
+ )
368
+ parser.add_argument(
369
+ "--show-sources",
370
+ action="store_true",
371
+ help="Show known sources on images (only works on PNG & SVG).",
372
+ )
373
+
374
+ parser.add_argument(
375
+ "--prior", type=str, default=None, help="Load the from an HDF5 file."
376
+ )
377
+ parser.add_argument(
378
+ "--posterior",
379
+ type=str,
380
+ default=None,
381
+ help="Store the posterior in HDF5 format file.",
382
+ )
383
+
384
+ parser.add_argument("--uv", action="store_true", help="Plot the UV coverage.")
385
+ parser.add_argument("--mu", action="store_true", help="Save the mean image.")
386
+ parser.add_argument(
387
+ "--pcf", action="store_true", help="Save the point covariance function image."
388
+ )
389
+ parser.add_argument(
390
+ "--var", action="store_true", help="Save the pixel variance image."
391
+ )
392
+ parser.add_argument(
393
+ "--nsamples",
394
+ type=int,
395
+ default=0,
396
+ help="Number of samples to save from the posterior.",
397
+ )
398
+
399
+ parser.add_argument(
400
+ "--title", required=False, default="disko", help="Prefix the output files."
401
+ )
402
+
403
+ source_json = None
404
+
405
+ log_fmt = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
406
+ logging.basicConfig(format=log_fmt, level=logging.INFO)
407
+
408
+ root = logging.getLogger()
409
+
410
+ fh = logging.FileHandler("disko.log")
411
+ # fh.setLevel(logging.INFO)
412
+
413
+ # create console handler and set level to debug
414
+ ch = logging.StreamHandler()
415
+ # ch.setLevel(logging.INFO)
416
+
417
+ # create formatter
418
+ formatter = logging.Formatter(
419
+ "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
420
+ )
421
+
422
+ # add formatter to ch
423
+ fh.setFormatter(formatter)
424
+
425
+ # add ch to logger
426
+ # root.addHandler(ch)
427
+ root.addHandler(fh)
428
+
429
+ # client = Client()
430
+
431
+ handle_bayes(parser.parse_args())
432
+
433
+ # client.close()
434
+ # local_cluster.close()
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "disko"
3
- version = "1.3.0"
3
+ version = "1.3.3"
4
4
  description = "Discrete Sky Operator (DiSkO) Aperture Synthesis Radio Imaging"
5
5
  authors = [
6
6
  {name = "Tim Molteno",email = "tim@elec.ac.nz"}
@@ -370,7 +370,7 @@ wheels = [
370
370
 
371
371
  [[package]]
372
372
  name = "disko"
373
- version = "1.2.0"
373
+ version = "1.3.0"
374
374
  source = { editable = "." }
375
375
  dependencies = [
376
376
  { name = "dask-ms" },
@@ -1,341 +0,0 @@
1
- #!/usr/bin/env python
2
- import matplotlib
3
- import os
4
-
5
- import matplotlib.pyplot as plt
6
-
7
- import argparse
8
- import datetime
9
- import json
10
- import logging
11
- import time
12
-
13
- from copy import deepcopy
14
-
15
- import numpy as np
16
- import scipy.special
17
-
18
- from tart.operation import settings
19
-
20
- from tart_tools import api_imaging
21
- from tart.imaging import elaz
22
- from tart.imaging import visibility
23
- from tart.imaging import calibration
24
-
25
- from .disko import DiSkO, vis_to_real
26
- from .telescope_operator import TelescopeOperator
27
-
28
- from .multivariate_gaussian import MultivariateGaussian
29
-
30
- from .parser_support import sphere_from_args, sphere_args_parser
31
-
32
- logger = logging.getLogger(__name__)
33
- logger.addHandler(logging.NullHandler()) # Add other handlers if you're using this as a library
34
- logger.setLevel(logging.INFO)
35
-
36
-
37
- def create_prior(vis_arr, sphere, hdf_prior):
38
- ''' Based on the size of the visibilities, try and calculate
39
- what range the image should have.
40
- '''
41
- if hdf_prior is not None:
42
- return MultivariateGaussian.from_hdf5(hdf_prior)
43
-
44
- vabs = np.abs(vis_arr)
45
-
46
- p05, p50, p95, p100 = np.percentile(vabs, [5, 50, 95, 100])
47
-
48
- var = p95*p95
49
- logger.info("Extimated Sky Prior variance={}".format(var))
50
- prior = MultivariateGaussian(np.zeros(sphere.npix) + p50, sigma=p95*np.identity(sphere.npix))
51
-
52
- return prior
53
-
54
-
55
- def do_inference(disko, sphere, prior, sigma_v=None):
56
- real_vis = vis_to_real(disko.vis_arr)
57
-
58
- to = TelescopeOperator(disko, sphere)
59
-
60
- # Transform to the natural basis.
61
- n_prior = prior.linear_transform(to.Vh)
62
- n_v = real_vis.shape[0]
63
-
64
- # TODO create a proper covariance that ensures the real and imaginary components are linked.
65
- if sigma_v is None:
66
- diag = np.diag(disko.rms**2)
67
- else:
68
- diag = np.diag(np.ones(n_v // 2)*(sigma_v)**2)
69
-
70
- logger.info(f"do_inference(sigma_v={diag[0,0]})")
71
-
72
- sigma_vis = np.block([[diag, 0.5*diag],[0.5*diag, diag]]) # .rechunk('auto')
73
-
74
- # now invert sigma_vis
75
- sigma_precision = MultivariateGaussian.sp_inv(sigma_vis)
76
- del sigma_vis
77
-
78
- if True:
79
- prior_r = n_prior.block(0,to.rank)
80
- prior_n = n_prior.block(to.rank,to.n_s)
81
-
82
- A_r = to.A_r
83
- V = to.V
84
-
85
- del to
86
- posterior_r = prior_r.bayes_update(sigma_precision, real_vis, A_r)
87
- posterior_n = prior_n
88
-
89
- del A_r
90
- del sigma_precision
91
- del prior_r
92
- del prior_n
93
- del n_prior
94
-
95
- posterior = MultivariateGaussian.outer(posterior_r, posterior_n)
96
-
97
- del posterior_r
98
- del posterior_n
99
-
100
- logger.info("Transforming posterior")
101
-
102
- posterior = posterior.linear_transform(V)
103
-
104
- del V
105
- else:
106
- posterior = to.sequential_inference(n_prior, real_vis, sigma_precision)
107
- del to
108
- del sigma_precision
109
- del n_prior
110
-
111
- return posterior
112
-
113
-
114
- def handle_bayes(ARGS):
115
-
116
- sphere = sphere_from_args(ARGS)
117
-
118
- # Create a prior.
119
- if ARGS.file:
120
- logger.info("Getting Data from file: {}".format(ARGS.file))
121
- # Load data from a JSON file
122
- with open(ARGS.file, 'r') as json_file:
123
- calib_info = json.load(json_file)
124
-
125
- info = calib_info['info']
126
- ant_pos = calib_info['ant_pos']
127
- config = settings.from_api_json(info['info'], ant_pos)
128
-
129
- flag_list = [] # [4, 5, 14, 22]
130
-
131
- original_positions = deepcopy(config.get_antenna_positions())
132
-
133
- gains_json = calib_info['gains']
134
- gains = np.asarray(gains_json['gain'])
135
- phase_offsets = np.asarray(gains_json['phase_offset'])
136
- config = settings.from_api_json(info['info'], ant_pos)
137
-
138
- measurements = []
139
- for d in calib_info['data']:
140
- vis_json, source_json = d
141
- cv, timestamp = api_imaging.vis_calibrated(vis_json, config, gains, phase_offsets, flag_list)
142
- src_list = elaz.from_json(source_json, 0.0)
143
-
144
- if ARGS.sigma_v is None:
145
- raise RuntimeError("The --sigma-v option must be supplied when --file JSON input is used")
146
-
147
- prior = create_prior(cv.v, sphere, ARGS.prior)
148
- timestamp = cv.get_timestamp()
149
- disko = DiSkO.from_cal_vis(cv)
150
-
151
- posterior = do_inference(disko, sphere, prior, sigma_v=ARGS.sigma_v)
152
- handle_output(ARGS, timestamp, posterior, sphere, disko)
153
-
154
- elif ARGS.hdf:
155
- logger.info(f"Getting data from file {ARGS.hdf}")
156
- if ARGS.sigma_v is None:
157
- raise RuntimeError("The --sigma-v option must be supplied when HDF5 input is used")
158
-
159
- data = visibility.from_hdf5(ARGS.hdf)
160
-
161
- prior = create_prior(data['vis_list'][0].v, sphere, ARGS.prior)
162
- posterior = None
163
-
164
- for v in data['vis_list']:
165
- if posterior is not None:
166
- prior = posterior
167
- cv = calibration.CalibratedVisibility(v)
168
- cv.set_config(v.config)
169
- cv.set_phase_offset(list(range(cv.get_config().get_num_antenna())),np.array(data['phase_offset']))
170
- cv.set_gain(list(range(cv.get_config().get_num_antenna())),np.array(data['gain']))
171
- timestamp = cv.get_timestamp()
172
- disko = DiSkO.from_cal_vis(cv)
173
-
174
- # TODO Calibrate the vis with gains and phases?
175
- posterior = do_inference(disko, sphere, prior, sigma_v=ARGS.sigma_v)
176
- handle_output(ARGS, timestamp, posterior, sphere, disko)
177
- else:
178
- logger.info("Getting Data from MS file: {}".format(ARGS.ms))
179
-
180
- disko = DiSkO.from_ms(ARGS.ms, ARGS.nvis, res_arcmin=sphere.res_arcmin, channel=ARGS.channel, field_id=ARGS.field)
181
- # Convert from reduced Julian Date to timestamp.
182
- timestamp = disko.timestamp
183
- src_list = None
184
-
185
- prior = create_prior(disko.vis_arr, sphere, ARGS.prior)
186
-
187
- posterior = do_inference(disko, sphere, prior, sigma_v=ARGS.sigma_v)
188
- handle_output(ARGS, timestamp, posterior, sphere, disko)
189
-
190
-
191
- def handle_output(ARGS, timestamp, posterior, sphere, disko=None):
192
-
193
- if not ARGS.show_sources:
194
- src_list = None
195
-
196
- time_repr = "{:%Y_%m_%d_%H_%M_%S_%Z}".format(timestamp)
197
-
198
- # Now save the files.
199
- if ARGS.posterior is not None:
200
- posterior.to_hdf5(ARGS.posterior)
201
-
202
- def path(ending, image_title):
203
- os.makedirs(ARGS.dir, exist_ok=True)
204
- fname = '{}.{}'.format(image_title, ending)
205
- return os.path.join(ARGS.dir, fname)
206
-
207
- def save_images(image_title, source_list):
208
- # Save as a FITS file
209
-
210
- if ARGS.FITS:
211
- sphere.to_fits(fname=path('fits', image_title), fov=ARGS.fov, info=disko.info)
212
-
213
- if ARGS.SVG:
214
- fname = path('svg', image_title)
215
- sphere.to_svg(fname=fname, show_grid=True, src_list=source_list, fov=ARGS.fov, title=image_title, show_cbar=True)
216
- logger.info("Generating {}".format(fname))
217
- if ARGS.PNG:
218
- fname = path('png', image_title)
219
- sphere.plot(plt, source_list)
220
- plt.title(image_title)
221
- plt.tight_layout()
222
- plt.savefig(fname, dpi=300)
223
- plt.close()
224
- logger.info("Generating {}".format(fname))
225
- if ARGS.PDF:
226
- fname = path('pdf', image_title)
227
- sphere.plot(plt, source_list)
228
- plt.title(image_title)
229
- plt.savefig(fname, dpi=600)
230
- plt.close()
231
- logger.info("Generating {}".format(fname))
232
-
233
- if ARGS.PDF or ARGS.PNG or ARGS.SVG or ARGS.FITS:
234
-
235
-
236
- if ARGS.mu:
237
- logger.info("Computing pixels")
238
- tic = time.perf_counter()
239
- #mu_positive = np.array(da.clip(posterior.mu, 0, None))
240
- logger.info(f" Took {time.perf_counter() - tic:0.4f} seconds")
241
- stat = sphere.set_visible_pixels(np.array(posterior.mu), scale=False)
242
- stat['sigma-v'] = ARGS.sigma_v
243
- logger.info(json.dumps(stat, sort_keys=True))
244
- save_images('{}_{}_mu'.format(ARGS.title, time_repr), source_list=src_list)
245
-
246
- if ARGS.var:
247
- tic = time.perf_counter()
248
- logger.info("Computing variance...")
249
- variance = np.array(posterior.variance())
250
- logger.info(f" Took {time.perf_counter() - tic:0.4f} seconds")
251
- sphere.set_visible_pixels(variance, scale=False)
252
- save_images('{}_{}_var'.format(ARGS.title, time_repr), source_list=None)
253
-
254
- if ARGS.pcf:
255
- tic = time.perf_counter()
256
- logger.info("Computing point covariance...")
257
-
258
- brightest_pixel = np.argmax(posterior.mu)
259
- pix_cov=np.array(posterior.sigma()[brightest_pixel,:])
260
- logger.info(f" Took {time.perf_counter() - tic:0.4f} seconds")
261
-
262
- sphere.set_visible_pixels(pix_cov, scale=False)
263
- save_images('{}_{}_pcf'.format(ARGS.title, time_repr), source_list=None)
264
-
265
- for i in range(ARGS.nsamples):
266
- sphere.set_visible_pixels(posterior.sample(), scale=False)
267
- save_images(image_title = '{}_{}_s{:0>5}'.format(ARGS.title, time_repr, i), source_list=None)
268
-
269
-
270
- def main():
271
- np.random.seed(42)
272
- sphere_parsers = sphere_args_parser()
273
-
274
- parser = argparse.ArgumentParser(
275
- description='DiSkO: Bayesian inference of a posterior sky',
276
- formatter_class=argparse.ArgumentDefaultsHelpFormatter,
277
- parents=sphere_parsers)
278
-
279
-
280
- parser.add_argument('--hdf', required=False, default=None, help="Exported Multi-visibility file")
281
- parser.add_argument('--ms', required=False, default=None, help="visibility file")
282
- parser.add_argument('--file', required=False, default=None, help="Snapshot observation saved JSON file (visiblities, positions and more).")
283
-
284
-
285
- parser.add_argument('--channel', type=int, default=0, help="Use this frequency channel.")
286
- parser.add_argument('--field', type=int, default=0, help="Use this FIELD_ID from the measurement set.")
287
-
288
- parser.add_argument('--dir', required=False, default='.', help="Output directory.")
289
- parser.add_argument('--nvis', type=int, default=1000, help="Number of visibilities to use.")
290
- parser.add_argument('--arcmin', type=float, default=None, help="Highest allowed res of the sky in arc minutes.")
291
-
292
- parser.add_argument('--sigma-v', type=float, default=None, help="Diagonal components of the visibility covariance. If not supplied use measurement set values")
293
-
294
- parser.add_argument('--PNG', action="store_true", help="Generate a PNG format image.")
295
- parser.add_argument('--PDF', action="store_true", help="Generate a PDF format image.")
296
- parser.add_argument('--SVG', action="store_true", help="Generate a SVG format image.")
297
- parser.add_argument('--FITS', action="store_true", help="Generate a FITS format image.")
298
- parser.add_argument('--show-sources', action="store_true", help="Show known sources on images (only works on PNG & SVG).")
299
-
300
- parser.add_argument('--prior', type=str, default=None, help="Load the from an HDF5 file.")
301
- parser.add_argument('--posterior', type=str, default=None, help="Store the posterior in HDF5 format file.")
302
-
303
- parser.add_argument('--uv', action="store_true", help="Plot the UV coverage.")
304
- parser.add_argument('--mu', action="store_true", help="Save the mean image.")
305
- parser.add_argument('--pcf', action="store_true", help="Save the point covariance function image.")
306
- parser.add_argument('--var', action="store_true", help="Save the pixel variance image.")
307
- parser.add_argument('--nsamples', type=int, default=0, help="Number of samples to save from the posterior.")
308
-
309
- parser.add_argument('--title', required=False, default="disko", help="Prefix the output files.")
310
-
311
- source_json = None
312
-
313
-
314
- log_fmt = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
315
- logging.basicConfig(format=log_fmt, level=logging.INFO)
316
-
317
- root = logging.getLogger()
318
-
319
- fh = logging.FileHandler('disko.log')
320
- #fh.setLevel(logging.INFO)
321
-
322
- # create console handler and set level to debug
323
- ch = logging.StreamHandler()
324
- #ch.setLevel(logging.INFO)
325
-
326
- # create formatter
327
- formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
328
-
329
- # add formatter to ch
330
- fh.setFormatter(formatter)
331
-
332
- # add ch to logger
333
- #root.addHandler(ch)
334
- root.addHandler(fh)
335
-
336
- #client = Client()
337
-
338
- handle_bayes(parser.parse_args())
339
-
340
- #client.close()
341
- #local_cluster.close()
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes