Coverage for frank/utilities.py: 94%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Frankenstein: 1D disc brightness profile reconstruction from Fourier data
2# using non-parametric Gaussian Processes
3#
4# Copyright (C) 2019-2020 R. Booth, J. Jennings, M. Tazzari
5#
6# This program is free software: you can redistribute it and/or modify
7# it under the terms of the GNU General Public License as published by
8# the Free Software Foundation, either version 3 of the License, or
9# (at your option) any later version.
10#
11# This program is distributed in the hope that it will be useful,
12# but WITHOUT ANY WARRANTY; without even the implied warranty of
13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14# GNU General Public License for more details.
15#
16# You should have received a copy of the GNU General Public License
17# along with this program. If not, see <https://www.gnu.org/licenses/>
18#
19"""This module has various functions useful for plotting and analyzing fit
20results.
21"""
22import logging
23import numpy as np
24from scipy.interpolate import interp1d
26from frank.constants import deg_to_rad, sterad_to_arcsec, rad_to_arcsec
27from frank.geometry import FixedGeometry
28from frank.hankel import DiscreteHankelTransform
29from frank.statistical_models import VisibilityMapping
31def arcsec_baseline(x):
32 """
33 Provide x as a radial scale [arcsec] to return the corresponding baseline
34 [lambda], or vice-versa
36 Parameters
37 ----------
38 x : float
39 Radial scale [arcsec] or baseline [lambda]
41 Returns
42 -------
43 converted : float
44 Baseline [lambda] or radial scale [arcsec]
46 """
48 converted = 1 / (x / 60 / 60 * np.pi / 180)
50 return converted
53def radius_convert(x, dist, conversion='arcsec_au'):
54 """
55 Provide x as a radius/radii in [arcsec] to convert to [au] (or vice-versa),
56 assuming a distance in [pc]
58 Parameters
59 ----------
60 x : float or array
61 Radius or radii ([arcsec] or [au])
62 dist : float
63 Distance to source [pc]
64 conversion : {'arcsec_au', 'au_arcsec'}, default = 'arcsec_au'
65 The unit conversion to perform, e.g. 'arcsec_au' converts from [arcsec]
66 to [au]
68 Returns
69 -------
70 converted : float or array
71 Radius or radii ([au] or [arcsec])
73 """
75 if conversion == 'arcsec_au':
76 converted = x * dist
77 elif conversion == 'au_arcsec':
78 converted = x / dist
79 else:
80 raise AttributeError("conversion must be one of {}"
81 "".format(['arcsec_au', 'au_arcsec']))
83 return converted
86def jy_convert(x, conversion, bmaj=None, bmin=None):
87 """
88 Provide x as a brightness in one of the units [Jy / beam], [Jy / arcsec^2],
89 [Jy / sterad] to convert x to another of these units
91 Parameters
92 ----------
93 x : array
94 Brightness in one of: [Jy / beam], [Jy / arcsec^2], [Jy / sterad]
95 conversion : { 'beam_sterad', 'beam_arcsec2', 'arcsec2_beam',
96 'arcsec2_sterad', 'sterad_beam', 'sterad_arcsec2'}
97 The unit conversion to perform, e.g., 'beam_sterad' converts
98 [Jy / beam] to [Jy / sterad]
99 bmaj : float, optional
100 Beam FWHM along the major axis [arcsec]
101 bmin : float, optional
102 Beam FWHM along the minor axis [arcsec]
104 Returns
105 -------
106 converted : float
107 Brightness in converted units
109 """
110 if (bmaj is None or bmin is None) and conversion in ['beam_sterad',
111 'beam_arcsec2',
112 'arcsec2_beam',
113 'sterad_beam']:
114 raise ValueError('bmaj and bmin must be specified to perform the'
115 ' conversion {}'.format(conversion))
117 if bmaj is not None and bmin is not None:
118 beam_solid_angle = np.pi * bmaj * bmin / (4 * np.log(2))
120 if conversion == 'beam_arcsec2':
121 converted = x / beam_solid_angle
122 elif conversion == 'arcsec2_beam':
123 converted = x * beam_solid_angle
124 elif conversion == 'arcsec2_sterad':
125 converted = x * sterad_to_arcsec
126 elif conversion == 'sterad_arcsec2':
127 converted = x / sterad_to_arcsec
128 elif conversion == 'beam_sterad':
129 converted = x / beam_solid_angle * sterad_to_arcsec
130 elif conversion == 'sterad_beam':
131 converted = x * beam_solid_angle / sterad_to_arcsec
132 else:
133 raise AttributeError("conversion must be one of {}"
134 "".format(['beam_sterad', 'beam_arcsec2',
135 'arcsec2_beam', 'arcsec2_sterad',
136 'sterad_beam', 'sterad_arcsec2']))
138 return converted
141def get_fit_stat_uncer(fit, return_linear=True):
142 r"""
143 Retrieve a model brightness profile's statistical uncertainty (that due to
144 the uncertainty on the observed baselines). This is only a lower bound on
145 the total uncertainty (which also has a systematic contribution due to
146 sparse sampling of the (u, v) plane).
148 Parameters
149 ----------
150 fit : FrankFitter result object
151 Fitted profile to obtain statistical uncertainty of
153 Returns
154 -------
155 sigma: array
156 Statistical uncertainty on the fitted brightness, at the fitted radial
157 points.
158 return_linear : bool, default=True
159 Whether to return linear uncertainty for a LogNormal fit
160 (if False, the logarithmic uncertainty is instead returned).
161 Has no effect for a Normal fit.
162 """
164 if 'method' not in fit._info.keys():
165 raise AttributeError("'fit' object lacks '_info.method' key. Should be one "
166 "of ['linear', 'log']")
168 variance = np.diag(fit.covariance)
170 if fit._info["method"] == "LogNormal" and return_linear == True:
171 # convert stored variance from var(log(brightness)) to var(brightness).
172 # see Appendix A.2 in https://ui.adsabs.harvard.edu/abs/2016A%26A...586A..76J/abstract;
173 # https://en.wikipedia.org/wiki/Log-normal_distribution
174 variance = (np.exp(variance) - 1) * np.exp(2 * np.log(fit.I))
176 sigma = np.sqrt(variance)
177 return sigma
180class UVDataBinner(object):
181 r"""
182 Average uv-data into bins of equal size.
184 Compute the weighted mean of the visibilities in each bin
186 Parameters
187 ----------
188 uv : array, unit = :math:`\lambda`
189 Baselines of the data to bin
190 V : array, unit = Jy
191 Observed visibility. If complex, both the real and imaginary
192 components will be binned. Else only the real part will be binned.
193 weights : array, unit = Jy^-2
194 Weights on the visibility points
195 bin_width : float, unit = :math:`\lambda`
196 Width of the uv-bins
198 Notes
199 -----
200 Uses numpy masked arrays to mask bins with no uv points.
202 """
204 def __init__(self, uv, V, weights, bin_width):
205 nbins = np.ceil(uv.max() / bin_width).astype('int')
206 #Protect against rounding
207 if nbins * bin_width < uv.max():
208 nbins += 1
209 bins = np.arange(nbins+1, dtype='float64') * bin_width
211 self._bins = bins
212 self._nbins = nbins
213 self._norm = 1 / bin_width
215 bin_uv, bin_wgt, bin_vis, bin_n = \
216 self.bin_quantities(uv, weights,
217 uv, np.ones_like(uv), V,
218 bin_counts=True)
220 # Normalize
221 idx = bin_n > 0
222 w = bin_wgt[idx]
223 bin_uv[idx] /= w
224 bin_vis[idx] /= w
226 # Store the binned data, masking empty bins:
227 mask = (bin_n == 0)
228 self._uv = np.ma.masked_where(mask, bin_uv)
229 self._V = np.ma.masked_where(mask, bin_vis)
230 self._w = np.ma.masked_where(mask, bin_wgt)
231 self._count = np.ma.masked_where(mask, bin_n)
233 self._uv_left = np.ma.masked_where(mask, bins[:-1])
234 self._uv_right = np.ma.masked_where(mask, bins[1:])
236 # Compute the uncertainty on the means:
237 bin_vis_err = np.full(nbins, np.nan, dtype=V.dtype)
239 # 1) Get the binned mean for each vis:
240 mu = self._V[self.determine_uv_bin(uv)]
242 # 2) Compute weighted error for bins with n > 1
243 quantities = (V-mu).real**2
244 if np.iscomplexobj(V):
245 quantities = quantities + 1j * (V-mu).imag**2
246 err = self.bin_quantities(uv, weights**2, quantities)
248 idx2 = bin_n > 1
249 err[idx2] /= bin_wgt[idx2]**2 * (1 - 1 / bin_n[idx2])
251 temp_error = np.sqrt(err.real[idx2])
252 if np.iscomplexobj(V):
253 temp_error = temp_error + 1.j * np.sqrt(err.imag[idx2])
254 bin_vis_err[idx2] = temp_error
256 # 3) Use a sensible error for bins with one baseline
257 idx1 = bin_n == 1
258 bin_vis_err[idx1].real = 1 / np.sqrt(bin_wgt[idx1])
259 if np.iscomplexobj(V):
260 bin_vis_err[idx1].imag = bin_vis_err[idx1].real
261 bin_vis_err[mask] = np.nan
263 # 4) Store the error
264 self._Verr = np.ma.masked_where(mask, bin_vis_err)
267 def determine_uv_bin(self, uv):
268 r"""Determine the bin that the given uv points belong too.
270 Parameters
271 ----------
272 uv : array, unit = :math:`\lambda`
273 Baselines to determine the bins of
275 Returns
276 -------
277 idx : array,
278 Bins that the uv point belongs to. Will be -1 if the bin does not
279 exist.
281 """
283 bins = self._bins
284 nbins = self._nbins
286 idx = np.floor(uv * self._norm).astype('int32')
287 idx[uv < bins[idx]] -= 1 # Fix rounding error
288 idx[uv == bins[nbins]] -= 1 # Handle point exaclty on outer boundary
290 too_high = idx >= nbins
291 idx[too_high] = -1
293 # Increase points on the RH boundary, unless it is the outer one
294 idx_tmp = idx[~too_high]
295 increment = (uv[~too_high] >= bins[idx_tmp+1]) & (idx_tmp+1 < nbins)
296 idx_tmp[increment] += 1
298 return idx
300 def bin_quantities(self, uv, w, *quantities, bin_counts=False):
301 r"""Bin the given quantities according the to uv points and weights.
303 Parameters
304 ----------
305 uv : array, unit = :math:`\lambda`
306 Baselines of the data to bin
307 weights : array, unit = Jy^-2
308 Weights on the visibility points
309 quantities : arrays,
310 Quantities evaluated at the uv points to bin.
311 bin_counts : bool, default = False
312 Determines whether to count the number of uv points per bin.
314 Returns
315 -------
316 results : arrays, same type as quantities
317 Binned data
318 bin_counts : array, int64, optional
319 If bin_counts=True, then this array contains the number of uv
320 points in each bin. Otherwise, it is not returned.
322 """
324 bins = self._bins
325 nbins = self._nbins
326 norm = self._norm
328 results = [np.zeros(nbins, dtype=x.dtype) for x in quantities]
330 if bin_counts:
331 counts = np.zeros(nbins, dtype='int64')
333 BLOCK = 65536
334 for i in range(0, len(uv), BLOCK):
335 tmp_uv = uv[i:i+BLOCK]
336 tmp_wgt = w[i:i+BLOCK]
338 idx = np.floor(tmp_uv * norm).astype('int32')
340 # Fix rounding:
341 idx[tmp_uv < bins[idx]] -= 1
342 # Move points exactly on the outer boundary inwards:
343 idx[idx == nbins] -= 1
345 # Only the last bin includes its edge
346 increment = (tmp_uv >= bins[idx+1]) & (idx+1 != nbins)
347 idx[increment] += 1
349 if bin_counts:
350 counts += np.bincount(idx, minlength=nbins)
352 for qty, res in zip(quantities, results):
353 tmp_qty = tmp_wgt*qty[i:i+BLOCK]
355 if np.iscomplexobj(qty):
356 res.real += np.bincount(idx, weights=tmp_qty.real,
357 minlength=nbins)
358 res.imag += np.bincount(idx, weights=tmp_qty.imag,
359 minlength=nbins)
360 else:
361 res += np.bincount(idx, weights=tmp_qty, minlength=nbins)
363 if bin_counts:
364 return results + [counts]
365 if len(results) == 1:
366 return results[0]
367 return results
369 def __len__(self):
370 return len(self._uv)
372 @property
373 def uv(self):
374 r"""Binned uv points, unit = :math:`\lambda`"""
375 return self._uv
377 @property
378 def V(self):
379 """Binned visibility, unit = Jy"""
380 return self._V
382 @property
383 def weights(self):
384 """Binned weights, unit = Jy^-2"""
385 return self._w
387 @property
388 def error(self):
389 """Uncertainty on the binned visibilities, unit = Jy"""
390 return self._Verr
392 @property
393 def bin_counts(self):
394 """Number of points in each bin"""
395 return self._count
397 @property
398 def bin_edges(self):
399 """Edges of the histogram bins"""
400 return [self._uv_left, self._uv_right]
403def check_uv(u, v, min_q=1e3, max_q=1e8):
404 r"""
405 Check if u,v distances are sensible for expected code unit of
406 [lambda], or if instead they're being supplied in [m].
408 Parameters
409 ----------
410 u, v : array, unit = :math:`\lambda`
411 u and v coordinates of observations
412 min_q : float, unit = :math:`\lambda`, default=1e3
413 Minimum baseline in code units expected for a dataset. The default
414 value of 1e3 is a conservative value for ALMA, assuming a minimum
415 antenna separation of ~12 m and maximum observing wavelength of 3.6 mm.
416 max_q : float, unit = :math:`\lambda`, default=1e5
417 Maximum baseline in code units expected for a dataset. The default
418 value of 1e8 is a conservative value for ALMA, assuming a maximum
419 antenna separation of ~16 km and minimum observing wavelength of 0.3 mm.
420 """
421 q = np.hypot(u, v)
423 if min(q) < min_q:
424 logging.warning("WARNING: "
425 f"Minimum baseline {min(q):.1e} < expected minimum {min_q:.1e} [lambda]. "
426 "'u' and 'v' distances must be in units of [lambda], "
427 "but it looks like they're in [m]."
428 )
431def normalize_uv(u, v, wle):
432 r"""
433 Normalize data u and v coordinates by the observing wavelength
435 Parameters
436 ----------
437 u, v : array, unit = [m]
438 u and v coordinates of observations
439 wle : float or array, unit = [m]
440 Observing wavelength of observations. If an array, it should be the
441 pointwise wavelength for each (u,v) point
443 Returns
444 -------
445 u_normed, v_normed : array, unit = :math:`\lambda`
446 u and v coordinates normalized by observing wavelength
448 """
450 logging.info(' Normalizing u and v coordinates by provided'
451 ' observing wavelength of {} m'.format(wle))
453 wle = np.atleast_1d(wle).astype('f8')
454 if len(wle) != 1 and len(wle) != len(u):
455 raise ValueError("len(wle) = {}. It should be equal to "
456 "len(u) = {} (or 1 if all wavelengths are the same)".format(len(wle), len(u)))
457 u_normed = u / wle
458 v_normed = v / wle
460 return u_normed, v_normed
463def cut_data_by_baseline(u, v, vis, weights, cut_range, geometry=None):
464 r"""
465 Truncate the data to be within a chosen baseline range.
467 The cut will be done in deprojected baseline space if the geometry is
468 provided.
470 Parameters
471 ----------
472 u, v : array, unit = :math:`\lambda`
473 u and v coordinates of observations
474 vis : array, unit = Jy
475 Observed visibilities (complex: real + imag * 1j)
476 weights : array, unit = Jy^-2
477 Weights assigned to observed visibilities, of the form
478 :math:`1 / \sigma^2`
479 cut_range : list of float, length = 2, unit = [\lambda]
480 Lower and upper baseline bounds outside of which visibilities are
481 truncated
482 geometry : SourceGeometry object, optional
483 Fitted geometry (see frank.geometry.SourceGeometry).
486 Returns
487 -------
488 u_cut, v_cut : array, unit = :math:`\lambda`
489 u and v coordinates in the chosen baseline range
490 vis_cut : array, unit = Jy
491 Visibilities in the chosen baseline range
492 weights_cut : array, unit = Jy^-2
493 Weights in the chosen baseline range
495 """
497 logging.info(' Cutting data outside of the minimum and maximum baselines'
498 ' of {} and {}'
499 ' klambda'.format(cut_range[0] / 1e3,
500 cut_range[1] / 1e3))
501 if geometry is not None:
502 up, vp = geometry.deproject(u, v)
503 else:
504 up, vp = u, v
506 baselines = np.hypot(up, vp)
507 above_lo = baselines >= cut_range[0]
508 below_hi = baselines <= cut_range[1]
509 in_range = above_lo & below_hi
510 u_cut, v_cut, vis_cut, weights_cut = [x[in_range] for x in [u, v, vis, weights]]
512 return u_cut, v_cut, vis_cut, weights_cut
515def estimate_weights(u, v=None, V=None, nbins=300, log=True, use_median=False,
516 verbose=True):
517 r"""
518 Estimate the weights using the variance of the binned visibilities.
520 The estimation is done assuming that the variation in each bin is dominated
521 by the noise. This will be true if:
522 1) The source is axi-symmetric,
523 2) The uv-points have been deprojected,
524 3) The bins are not too wide,
525 Otherwise the variance may be dominated by intrinsic variations in the
526 visibilities.
528 Parameters
529 ----------
530 u, v : array, unit = :math:`\lambda`
531 u and v coordinates of observations (deprojected). Data will be binned
532 by baseline. If v is not None, np.hpot(u,v) will be used instead. Note
533 that if V is None the argument v will be intepreted as V instead
534 V : array, unit = Jy, default = None
535 Observed visibility. If complex, the weights will be computed from the
536 average of the variance of the real and imaginary components, as in
537 CASA's statwt. Otherwise the variance of the real part is used.
538 nbins : int, default = 300
539 Number of bins used.
540 log : bool, default = True
541 If True, the uv bins will be constructed in log space, otherwise linear
542 spaced bins will be used.
543 use_median : bool, default = False
544 If True all of the weights will be set to the median of the variance
545 estimated across the bins. Otherwise, the baseline dependent variance
546 will be used.
547 verbose : bool, default = True
548 If true, the logger will record calls to this function, along with
549 whether the median estimate was used.
551 Returns
552 -------
553 weights : array
554 Estimate of the weight for each uv point.
556 Notes
557 -----
558 - This function does not use the original weights in the estimation.
559 - Bins with only one uv point do not have a variance estimate. Thus
560 the mean of the variance in the two adjacent bins is used instead.
562 Examples
563 --------
564 All of the following calls will work as expected:
565 `estimate_weights(u, v, V) `
566 `estimate_weights(u, V)`
567 `estimate_weights(u, V=V)`
568 In each case the variance of V in the uv-bins is used to estimate the
569 weights. The first call will use q = np.hypot(u, v) in the uv-bins. The
570 second and third calls are equivalent to the first with u=0.
572 """
574 if verbose:
575 logging.info(' Estimating visibility weights')
577 if V is None:
578 if v is not None:
579 V = v
580 q = np.abs(u)
581 else:
582 raise ValueError("The visibilities, V, must be supplied")
583 elif v is not None:
584 q = np.hypot(u,v)
585 else:
586 q = np.abs(u)
588 if log:
589 q = np.log(q)
590 q -= q.min()
592 bin_width = (q.max()-q.min()) / nbins
594 uvBin = UVDataBinner(q, V, np.ones_like(q), bin_width)
596 if uvBin.bin_counts.max() == 1:
597 raise ValueError("No bin contains more than one uv point, can't"
598 " estimate the variance. Use fewer bins.")
600 if np.iscomplex(V.dtype):
601 var = 0.5*(uvBin.error.real**2 + uvBin.error.imag**2) * uvBin.bin_counts
602 else:
603 var = uvBin.error.real**2 * uvBin.bin_counts
605 if use_median:
606 if verbose:
607 logging.info(' Setting all weights as median binned visibility '
608 'variance')
609 return np.full(len(u), 1/np.ma.median(var[uvBin.bin_counts > 1]))
610 else:
611 if verbose:
612 logging.info(' Setting weights according to baseline-dependent '
613 'binned visibility variance')
614 # For bins with 1 uv point, use the average of the adjacent bins
615 no_var = np.argwhere(uvBin.bin_counts == 1).reshape(-1)
616 if len(no_var) > 0:
617 # Find the location `loc` of the bad points in the array of good points
618 good_var = np.argwhere(uvBin.bin_counts > 1).reshape(-1)
619 loc = np.searchsorted(good_var, no_var, side='right')
621 # Set the variance to the average of the two adjacent bins
622 im = good_var[np.maximum(loc-1, 0)]
623 ip = good_var[np.minimum(loc, len(good_var)-1)]
624 var[no_var] = 0.5*(var[im] + var[ip])
626 bin_id = uvBin.determine_uv_bin(q)
627 assert np.all(bin_id != -1), "Error in binning" # Should never occur
629 weights = 1/var[bin_id]
631 return weights
634def draw_bootstrap_sample(u, v, vis, weights):
635 r"""
636 Obtain the sample for a bootstrap, drawing, with replacement, N samples from
637 a length N dataset
639 Parameters
640 ----------
641 u, v : array, unit = :math:`\lambda`
642 u and v coordinates of observations
643 vis : array, unit = Jy
644 Observed visibilities (complex: real + imag * 1j)
645 weights : array, unit = Jy^-2
646 Weights on the visibilities, of the form
647 :math:`1 / \sigma^2`
649 Returns
650 -------
651 u_boot, v_boot : array, unit = :math:`\lambda`
652 Bootstrap sampled u and v coordinates
653 vis_boot : array, unit = Jy
654 Bootstrap sampled visibilities
655 weights_boot : array, unit = Jy^-2
656 Boostrap sampled weights on the visibilities
658 """
659 idxs = np.random.randint(low=0, high=len(u), size=len(u))
661 u_boot = u[idxs]
662 v_boot = v[idxs]
663 vis_boot = vis[idxs]
664 weights_boot = weights[idxs]
666 return u_boot, v_boot, vis_boot, weights_boot
669def sweep_profile(r, I, project=False, phase_shift=False, geom=None, axis=0,
670 xmax=None, ymax=None, dr=None):
671 r"""
672 Sweep a 1D radial brightness profile over :math:`2 \pi` to yield a 2D
673 brightness distribution. Optionally project this sweep by a supplied
674 geometry.
676 Parameters
677 ----------
678 r : array
679 Radial coordinates at which the 1D brightness profile is defined
680 I : array
681 Brightness values at r
682 project : bool, default = False
683 Whether to project the swept profile by the supplied geom
684 phase_shift : bool, default = False
685 Whether to phase shift the projected profile by the supplied geom.
686 If False, the source will be centered in the image
687 geom : SourceGeometry object, default=None
688 Fitted geometry (see frank.geometry.SourceGeometry). Here we use
689 geom.inc [deg], geom.PA [deg], geom.dRA [arcsec], geom.dDec [arcsec] if
690 project=True
691 axis : int, default = 0
692 Axis over which to interpolate the 1D profile
693 xmax, ymax : float, optional, default = None
694 Value setting the x- and y-bounds of the image (same units as r). The
695 positive and negative bounds are both set to this value (modulo sign).
696 If not provided, these will be set to r.max()
697 dr : float, optional, default = None
698 Pixel size (same units as r). If not provided, it will be set at the
699 same spatial scale as r
701 Returns
702 -------
703 I2D : array, shape = (len(r), len(r))
704 2D brightness distribution (projected if project=True)
705 xmax : float
706 Maximum x-value of the 2D grid
707 ymax : float
708 Maximum y-value of the 2D grid
710 Notes
711 -----
712 Sign convention: a negative geom.dRA shifts the source to the right
713 in the image
715 """
716 if project:
717 inc, pa = geom.inc, geom.PA
718 inc *= deg_to_rad
719 pa *= deg_to_rad
721 cos_i = np.cos(inc)
722 cos_pa, sin_pa = np.cos(pa), np.sin(pa)
724 if xmax is None:
725 xmax = r.max()
726 if ymax is None:
727 ymax = r.max()
729 if dr is None:
730 dr = np.mean(np.diff(r))
732 x = np.linspace(-xmax, xmax, int(xmax/dr) + 1)
733 y = np.linspace(-ymax, ymax, int(ymax/dr) + 1)
735 if phase_shift:
736 xi, yi = np.meshgrid(x + geom.dRA, y - geom.dDec)
737 else:
738 xi, yi = np.meshgrid(x, y)
740 if project:
741 xp = xi * cos_pa + yi * sin_pa
742 yp = -xi * sin_pa + yi * cos_pa
743 xp /= cos_i
744 r1D = np.hypot(xp, yp)
745 else:
746 r1D = np.hypot(xi, yi)
748 im_shape = r1D.shape + I.shape[1:]
750 interp = interp1d(r, I, bounds_error=False, fill_value=0., axis=axis)
751 I2D = interp(r1D.ravel()).reshape(*im_shape)
753 return I2D, xmax, ymax
756def make_image(fit, Npix, xmax=None, ymax=None, project=True):
757 """Make an image of a model fit.
759 Parameters
760 ----------
761 fit : FrankFitter result object
762 Fitted profile to make an image of
763 Npix : int or list
764 Number of pixels in the x-direction, or [x-,y-] direction
765 xmax: float or None, unit=arcsec
766 Size of the image is [-xmax, xmax]. By default this is twice
767 fit.Rmax to avoid aliasing.
768 ymax: float or None, unit=arcsec
769 Size of the image is [-ymax,ymax]. Defaults to xmax if ymax=None
770 project: bool, default=True
771 Whether to produce a projected image.
773 Returns
774 -------
775 x : array, 1D; unit=arcsec
776 Locations of the x-points in the image.
777 y : array, 1D; unit=arcsec
778 Locations of the y-points in the image.
779 I : array, 2D; unit=Jy/Sr
780 Image of the surface brightness.
781 """
782 if xmax is None:
783 xmax = 2*fit.Rmax
784 if ymax is None:
785 ymax = xmax
787 try:
788 Nx, Ny = Npix
789 except TypeError:
790 Nx = Npix
791 Ny = int(Nx*(ymax/xmax))
793 dx = 2*xmax/(Nx*rad_to_arcsec)
794 dy = 2*ymax/(Ny*rad_to_arcsec)
797 # The easiest way to produce an image is to predict the visibilities
798 # on a regular grid in the Fourier plane and then transform it back.
799 # All frank models must be able to compute the visibilities so this
800 # method is completely general.
801 u = np.fft.fftfreq(Nx)/dx
802 v = np.fft.fftfreq(Ny)/dy
803 u, v = np.meshgrid(u,v, indexing='ij')
805 # Get the visibilities:
806 if project:
807 Vis = fit.predict(u.reshape(-1), v.reshape(-1)).reshape(*u.shape)
808 else:
809 q = np.hypot(u,v)
810 Vis = fit.predict_deprojected(q.reshape(-1)).reshape(*u.shape)
812 # Convert to the image plane
813 I = np.fft.ifft(np.fft.ifft(Vis, axis=0), axis=1).real
814 I /= dx*dy
816 # numpy's fft has zero in the corner. We want it in the middle so we need
817 # to wrap:
818 tmp = I.copy()
819 tmp[:Nx//2,], tmp[Nx//2:] = I[Nx//2:], I[:Nx//2]
820 I[:,:Ny//2], I[:,Ny//2:] = tmp[:,Ny//2:], tmp[:,:Ny//2]
822 xe = np.linspace(-xmax, xmax, Nx+1)
823 x = 0.5*(xe[1:] + xe[:-1])
824 ye = np.linspace(-ymax, ymax, Ny+1)
825 y = 0.5*(ye[1:] + ye[:-1])
827 return x, y, I
830def convolve_profile(r, I, disc_i, disc_pa, clean_beam,
831 n_per_sigma=5, axis=0):
832 r"""
833 Convolve a 1D radial brightness profile with a 2D Gaussian beam, degrading
834 the profile's resolution
836 Parameters
837 ----------
838 r : array
839 Radial coordinates at which the 1D brightness profile is defined
840 I : array
841 Brightness values at r
842 disc_i : float, unit = deg
843 Disc inclination
844 disc_pa : float, unit = deg
845 Disc position angle
846 clean_beam : dict
847 Dictionary with beam `bmaj` (FWHM of beam along its major axis) [arcsec],
848 `bmin` (FWHM of beam along its minor axis) [arcsec],
849 `pa` (beam position angle) [deg]
850 n_per_sigma : int, default = 50
851 Number of points per standard deviation of the Gaussian kernel (used
852 for gridding)
853 axis : int, default = 0
854 Axis over which to interpolate the 1D profile
856 Returns
857 -------
858 I_smooth : array, shape = (len(r), len(r))
859 Convolved brightness profile I at coordinates r
861 """
863 from scipy.constants import c
864 from scipy.interpolate import interp1d
865 from scipy.ndimage import gaussian_filter
867 # Set up the geometry for the smoothing mesh.
868 # We align the beam with the grid (major axis aligned) and rotate the
869 # image accordingly
871 # Convert beam FWHM to sigma
872 bmaj = clean_beam['bmaj'] / np.sqrt(8 * np.log(2))
873 bmin = clean_beam['bmin'] / np.sqrt(8 * np.log(2))
875 PA = (disc_pa - clean_beam['beam_pa']) * np.pi / 180.
877 cos_i = np.cos(disc_i * np.pi/180.)
878 cos_PA = np.cos(PA)
879 sin_PA = np.sin(PA)
881 # Pixel size in terms of bmin
882 rmax = r.max()
883 dx = bmin / n_per_sigma
885 xmax = rmax * (np.abs(cos_i * cos_PA) + np.abs(sin_PA))
886 ymax = rmax * (np.abs(cos_i * sin_PA) + np.abs(cos_PA))
888 xmax = int(xmax / dx + 1) * dx
889 ymax = int(ymax / dx + 1) * dx
891 x = np.arange(-xmax, xmax + dx / 2, dx)
892 y = np.arange(-ymax, ymax + dx / 2, dx)
894 xi, yi = np.meshgrid(x, y)
896 xp = xi * cos_PA + yi * sin_PA
897 yp = -xi * sin_PA + yi * cos_PA
898 xp /= cos_i
900 r1D = np.hypot(xi, yi)
902 im_shape = r1D.shape + I.shape[1:]
904 # Interpolate to grid and apply smoothing
905 interp = interp1d(r, I, bounds_error=False, fill_value=0., axis=axis)
907 I2D = interp(r1D.ravel()).reshape(*im_shape)
908 sigma = [float(n_per_sigma), (bmaj / bmin) * n_per_sigma]
909 I2D = gaussian_filter(I2D, sigma, mode='nearest', cval=0.)
911 # Now convert back to a 1D profile
912 edges = np.concatenate(([r[0] * r[0] / r[1]], r, [r[-1] * r[-1] / r[-2]]))
913 edges = 0.5 * (edges[:-1] + edges[1:])
915 I_smooth = np.empty_like(I)
916 I_smooth = np.histogram(r1D.ravel(), weights=I2D.ravel(), bins=edges)[0]
917 counts = np.maximum(np.histogram(r1D.ravel(), bins=edges)[0], 1)
918 I_smooth /= counts
920 return I_smooth
923def add_vis_noise(vis, weights, seed=None):
924 r"""
925 Add Gaussian noise to visibilities
927 Parameters
928 ----------
929 vis : array, unit = [Jy]
930 Visibilities to add noise to.
931 Can be complex (real + imag * 1j) or purely real.
932 weights : array, unit = [Jy^-2]
933 Weights on the visibilities, of the form :math:`1 / \sigma^2`.
934 Injected noise will be scaled proportional to `\sigma`.
935 seed : int, default = None
936 Number to initialize a pseudorandom number generator for the noise draws
938 Returns
939 -------
940 vis_noisy : array, shape = vis.shape
941 Visibilities with added noise
943 """
944 if seed is not None:
945 np.random.seed(seed)
947 dim0 = 1
948 if np.iscomplexobj(vis):
949 dim0 = 2
951 vis = np.array(vis)
952 noise = np.random.standard_normal((dim0,) + vis.shape)
953 noise *= weights ** -0.5
955 vis_noisy = vis + noise[0]
956 if np.iscomplexobj(vis):
957 vis_noisy += 1j * noise[1]
959 return vis_noisy
962def make_mock_data(r, I, Rmax, u, v, projection=None, geometry=None, N=500,
963 add_noise=False, weights=None, seed=None):
964 r"""
965 Generate mock visibilities from a provided brightness profile and (u,v)
966 distribution.
968 Parameters
969 ----------
970 r : array, unit = [arcsec]
971 Radial coordinates of I(r)
972 I : array, unit = [Jy / sr]
973 Brightness values at r
974 Rmax : float, unit = [arcsec], default=2.0
975 Maximum radius beyond which I(r) is zero. This should be larger than the
976 disk size
977 u, v : array, unit = :math:`\lambda`
978 u and v coordinates of observations
979 projection : str, default = None
980 One of [None, 'deproject', 'reproject']
981 If None, the visibilities will be neither deprojected nor reprojected.
982 If 'deproject' or 'reproject', the visibilites will be accordingly
983 deprojected or reprojected by the supplied `geometry` and their total
984 flux scaled by the inclination.
985 geometry : SourceGeometry object, default=None
986 Source geometry (see frank.geometry.SourceGeometry). Must be supplied
987 if `projection` is 'deproject' or 'reproject'.
988 N : integer, default=500
989 Number of terms to use in the Fourier-Bessel series
990 add_noise : bool, default = False
991 Whether to add noise to the mock visibilities
992 weights : array, unit = Jy^-2
993 Visibility weights, of the form :math:`1 / \sigma^2`.
994 If provided, injected noise will be scaled proportional to `\sigma`.
995 seed : int, default = None
996 Number to initialize a pseudorandom number generator for the noise draws
998 Returns
999 -------
1000 baselines : array, unit = :math:`\lambda`
1001 Baseline coordinates of the mock visibilities. These will be equal to
1002 np.hypot(u, v) if 'geometry' is None (or if its keys are all equal to 0)
1003 vis : array, unit = Jy
1004 Mock visibility amplitudes, including noise if 'add_noise' is True
1006 Notes
1007 -----
1008 'r' and 'I' should be sufficiently sampled to ensure an interpolation will
1009 be accurate, otherwise 'vis' may be a poor estimate of their transform.
1010 """
1011 proj_valid = [None, 'deproject', 'reproject']
1012 if projection not in proj_valid:
1013 raise AttributeError(f"projection is '{projection}'; must be one of {proj_valid}.")
1015 if projection in ['deproject', 'reproject']:
1016 if geometry is None:
1017 raise AttributeError(f"geometry must be supplied to perform {projection}.")
1019 if projection == 'deproject':
1020 u, v = geometry.deproject(u, v)
1021 else:
1022 u, v = geometry.reproject(u, v)
1024 else:
1025 if geometry is not None:
1026 raise AttributeError("projection is None; must be one of "
1027 "['deproject', 'reproject'] "
1028 "to perform projection.")
1029 geometry = FixedGeometry(0, 0, 0, 0)
1031 baselines = np.hypot(u, v)
1033 _, vis = generic_dht(r, I, Rmax, N, grid=baselines, inc=geometry.inc)
1035 if add_noise:
1036 vis = add_vis_noise(vis, weights, seed)
1038 return baselines, vis
1041def get_collocation_points(Rmax=2.0, N=500, direction='forward'):
1042 """
1043 Obtain the collocation points of a discrete Hankel transform for a given
1044 'Rmax' and 'N' (see frank.hankel.DiscreteHankelTransform)
1046 Parameters
1047 ----------
1048 Rmax : float, unit = [arcsec], default=2.0
1049 Maximum radius beyond which the real space function is zero
1050 N : integer, default=500
1051 Number of terms to use in the Fourier-Bessel series
1052 direction : { 'forward', 'backward' }, default='forward'
1053 Direction of the transform. 'forward' is real space -> Fourier space,
1054 returning real space radial collocation points needed for the transform.
1056 Returns
1057 -------
1058 coll_pts : array, unit = [lambda] or [arcsec]
1059 The DHT collocation points in either real or Fourier space.
1061 """
1062 if direction not in ['forward', 'backward']:
1063 raise AttributeError("direction must be one of ['forward', 'backward']")
1065 r_pts, q_pts = DiscreteHankelTransform.get_collocation_points(
1066 Rmax=Rmax / rad_to_arcsec, N=N, nu=0
1067 )
1069 if direction == 'forward':
1070 coll_pts = r_pts * rad_to_arcsec
1071 else:
1072 coll_pts = q_pts
1074 return coll_pts
1077def generic_dht(x, f, Rmax=2.0, N=500, direction='forward', grid=None,
1078 inc=0.0):
1079 """
1080 Compute the visibilities or brightness of a model by directly applying the
1081 Discrete Hankel Transform.
1083 The correction for inclination will also be applied, assuming an optically
1084 thick disc. For an optically thin disc, setting inc=0 (the default) will
1085 achieve the correct scaling.
1087 Parameters
1088 ----------
1089 x : array, unit = [arcsec] or [lambda]
1090 Radial or spatial frequency coordinates of f(x)
1091 f : array, unit = [Jy / sr] or [Jy]
1092 Amplitude values of f(x)
1093 Rmax : float, unit = [arcsec], default=2.0
1094 Maximum radius beyond which the real space function is zero
1095 N : integer, default=500
1096 Number of terms to use in the Fourier-Bessel series
1097 direction : { 'forward', 'backward' }, default='forward'
1098 Direction of the transform. 'forward' is real space -> Fourier space.
1099 grid : array, unit = [arcsec] or [lambda], default=None
1100 The radial or spatial frequency points at which to sample the DHT.
1101 If None, the DHT collocation points will be used.
1102 inc : float, unit = [deg], default = 0.0
1103 Source inclination. The total flux of the transform of f(x)
1104 will be scaled by cos(inc); this has no effect if inc=0.
1106 Returns
1107 -------
1108 grid : array, size=N, unit = [arcsec] or [lambda]
1109 Spatial frequency or radial coordinates of the Hankel transform of f(x)
1110 f_transform : array, size=N, unit = [Jy / sr] or [Jy]
1111 Hankel transform of f(x)
1113 Notes
1114 -----
1115 'x' and 'f' should be sufficiently sampled to ensure an interpolation will
1116 be accurate, otherwise 'f_transform' may be a poor estimate of their
1117 transform.
1118 """
1120 if direction not in ['forward', 'backward']:
1121 raise AttributeError("direction must be one of ['forward', 'backward']")
1123 DHT = DiscreteHankelTransform(Rmax=Rmax / rad_to_arcsec, N=N, nu=0)
1124 geom = FixedGeometry(inc, 0, 0, 0)
1125 VM = VisibilityMapping(DHT, geom)
1127 if direction == 'forward':
1128 # map the profile f(x) onto the DHT collocation points.
1129 # this is necessary for an accurate transform.
1130 y = np.interp(VM.r, x, f)
1132 if grid is None:
1133 grid = VM.q
1135 # perform the DHT
1136 f_transform = VM.predict_visibilities(y, grid, geometry=geom)
1138 else:
1139 y = np.interp(VM.q, x, f)
1141 if grid is None:
1142 grid = VM.r
1144 f_transform = VM.invert_visibilities(y, grid, geometry=geom)
1146 return grid, f_transform