diff --git a/healvis/observatory.py b/healvis/observatory.py index 88ce253..bfc9c8e 100644 --- a/healvis/observatory.py +++ b/healvis/observatory.py @@ -26,7 +26,8 @@ class Baseline(object): - def __init__(self, ant1_enu=None, ant2_enu=None, enu_vec=None): + + def __init__(self, ant1_enu=None, ant2_enu=None, ant1=None, ant2=None, enu_vec=None): if enu_vec is not None: self.enu = enu_vec else: @@ -36,6 +37,12 @@ def __init__(self, ant1_enu=None, ant2_enu=None, enu_vec=None): assert self.enu.size == 3, f"Wronge enu vector shape {self.enu.shape}" + # Antenna indexes, for indexing beam list if necessary. + # Must be numbers from 0 ..., so the right beams are found. + self.ant1 = ant1 + self.ant2 = ant2 + + def get_uvw(self, freq_Hz): return np.outer(self.enu, 1 / (c_ms / freq_Hz)) # In wavelengths @@ -66,6 +73,8 @@ class Observatory(object): ---------- latitude, longitude: float Decimal degrees position of the observatory on Earth. + height: float + Decimal meters height of the observatory on Earth. fov: float Field of view in degrees (Defaults to 180 deg for horizon to horizon). baseline_array: array_like of Baseline instances @@ -82,6 +91,7 @@ def __init__( self, latitude, longitude, + height=0.0, fov=None, baseline_array=None, freqs=None, @@ -109,7 +119,7 @@ def __init__( self.pointing_centers = None # List of [ra, dec] positions. One for each time. `set_pointings` sets this to zenith. self.north_poles = None # [ra,dec] ICRS position of the Earth's north pole. Set by `set_pointings`. self.telescope_location = EarthLocation.from_geodetic( - longitude * units.degree, latitude * units.degree + longitude * units.degree, latitude * units.degree, height ) self.do_horizon_taper = False @@ -230,20 +240,24 @@ def set_beam(self, beam="uniform", freq_interp_kind="linear", **kwargs): Set the beam of the array. Args: - beam : str - Input to PowerBeam or AnalyticBeam. If beam is - a viable input to AnalyticBeam, then instantiates - an AnalyticBeam, otherwise assumes beam is a filepath - to a beamfits and instantiates a PowerBeam. + beam : str, or list of beam objects + str: If it is a viable input to AnalyticBeam, + then instantiates an AnalyticBeam, otherwise assumes beam is + a filepath to a beamfits and instantiates a PowerBeam. + list: List of beam objects. This allows for external beams to be + used, and different beams for each antenna. They should not be + power beams. Each beam must have an interp method: + interp(self, az_array, za_array, freq_array) freq_interp_kind : str For PowerBeam, frequency interpolation option. kwargs : keyword arguments kwargs to pass to AnalyticBeam instantiation. """ - if beam in ["uniform", "gaussian", "airy"] or callable(beam): - self.beam = AnalyticBeam(beam, **kwargs) + if isinstance(beam, list): self.beam = beam + elif beam in ['uniform', 'gaussian', 'airy'] or callable(beam): + self.beam = AnalyticBeam(beam, **kwargs) else: self.beam = PowerBeam(beam) self.beam.interp_freq(self.freqs, inplace=True, kind=freq_interp_kind) @@ -261,6 +275,9 @@ def beam_sq_int(self, freqs, Nside, pointing, beam_pol="pI"): pointing : len-2 list Pointing center [Dec, RA] in J2000 degrees """ + + if isinstance(self.beam, list): + raise RuntimeError("beam_sq_int not implemented for multiple antenna beams") za, az = self.calc_azza(pointing) beam_sq_int = np.sum( self.beam.beam_val(az, za, freqs, pol=beam_pol) ** 2, axis=0 @@ -270,6 +287,14 @@ def beam_sq_int(self, freqs, Nside, pointing, beam_pol="pI"): return beam_sq_int + def external_beam_val(self, beam, az_arr, za_arr, freqs, pol="XX"): + """ + Call interp() on a beam, and provide results in the right format. + """ + interp_data, interp_basis_vector = beam.interp(az_array=az_arr, za_array=za_arr, + freq_array=freqs) + return interp_data[0, 0, 1].T # just want Npix, Nfreq + def _horizon_taper(self, za_arr): """ For pixels near the edge of the FoV downweight flux @@ -307,16 +332,51 @@ def _vis_calc(self, pcents, tinds, shell, vis_array, Nfin, beam_pol="pI"): memory_usage_GB = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6 north = self.north_poles[tinds[count]] if haspoles else None za_arr, az_arr, pix = self.calc_azza(c_, north, return_inds=True) - beam_cube = self.beam.beam_val(az_arr, za_arr, self.freqs, pol=beam_pol) + if isinstance(self.beam, list): + # Adds another dimension to beam_cube: the baselines. + # Beams may be different for each antenna, and they are + # not power beams. + # Multiplies the beams for the 2 antennas in a baseline. + + # Accumulate the antenna numbers + antennas = set() + for bi, baseline in enumerate(self.array): + assert baseline.ant1 is not None and baseline.ant2 is not None, \ + "Antenna number not set for baseline "+str(bi) + antennas.add(baseline.ant1) + antennas.add(baseline.ant2) + assert len(antennas) == len(self.beam), "Number of beams does not match number of antennas" + + beam_cube = [ None for i in range(len(self.array)) ] + beam_val = [ None for i in range(len(antennas)) ] + for i in antennas: + bv = self.external_beam_val(self.beam[i], az_arr, za_arr, self.freqs, pol=beam_pol) + bv[np.argwhere(za_arr>np.pi/2)[:, 0], :] = 0 # Sources below horizon + beam_val[i] = bv + for bi, bl in enumerate(self.array): + # Multiply beam correction for the two antennas in each baseline + beam_cube[bi] = beam_val[bl.ant1]*beam_val[bl.ant2] + else: + beam_cube = self.beam.beam_val(az_arr, za_arr, self.freqs, pol=beam_pol) + beam_cube[np.argwhere(za_arr>np.pi/2)[:, 0], :] = 0 # Sources below horizon + if self.do_horizon_taper: horizon_taper = self._horizon_taper(za_arr).reshape(1, za_arr.size, 1) else: horizon_taper = 1.0 - sky = shell[..., pix, :] * horizon_taper * beam_cube - for bi, bl in enumerate(self.array): - fringe_cube = bl.get_fringe(az_arr, za_arr, self.freqs) - vis = np.sum(sky * fringe_cube, axis=-2) - vis_array.put((tinds[count], bi, vis.tolist())) + if isinstance(self.beam, list): + # Beams are possibly different for each baseline + sky = shell[..., pix, :] * horizon_taper + for bi, bl in enumerate(self.array): + fringe_cube = bl.get_fringe(az_arr, za_arr, self.freqs) + vis = np.sum(sky * fringe_cube * beam_cube[bi], axis=-2) + vis_array.put((tinds[count], bi, vis.tolist())) + else: + sky = shell[..., pix, :] * horizon_taper * beam_cube + for bi, bl in enumerate(self.array): + fringe_cube = bl.get_fringe(az_arr, za_arr, self.freqs) + vis = np.sum(sky * fringe_cube, axis=-2) + vis_array.put((tinds[count], bi, vis.tolist())) with Nfin.get_lock(): Nfin.value += 1 if mp.current_process().name == "0" and Nfin.value > 0: diff --git a/healvis/sky_model.py b/healvis/sky_model.py index ea69f55..4418bfc 100644 --- a/healvis/sky_model.py +++ b/healvis/sky_model.py @@ -406,7 +406,7 @@ def gsm_shell(Nside, freqs, use_2016=False): def construct_skymodel( - sky_type, freqs=None, Nside=None, ref_chan=0, Nskies=1, sigma=None, amplitude=None + sky_type, freqs=None, Nside=None, ref_chan=0, Nskies=1, sigma=None, amplitude=None, seed=None ): """ Construct a SkyModel object or read from disk @@ -426,10 +426,16 @@ def construct_skymodel( If sky_type == 'flat_spec', this is the power spectrum amplitude amplitude : float Monopole amplitude in K + seed : int + Seed to initialise random number generator. Only relevant for flat + spectrum noise shell. Returns: SkyModel object """ + + if seed is not None: np.random.seed(seed) + sky = SkyModel() sky.Nside = Nside sky.freqs = freqs diff --git a/healvis/tests/test_external_beam.py b/healvis/tests/test_external_beam.py new file mode 100644 index 0000000..c3009ec --- /dev/null +++ b/healvis/tests/test_external_beam.py @@ -0,0 +1,112 @@ +# -*- mode: python; coding: utf-8 -* +# Copyright (c) 2019 Radio Astronomy Software Group +# Licensed under the 3-clause BSD License + +import numpy as np +from healvis.observatory import Baseline, Observatory +from healvis.sky_model import construct_skymodel + +class ExternalBeam: + """ + A dummy beam object for testing. + + It returns an array of the right size from interp(). + Demonstrates that any beam with an interp() method + can be used. + """ + + def interp(self, az_array, za_array, freq_array): + """ + Evaluate the primary beam at given az, za locations (in radians). + Parameters + ---------- + az_array : array_like + Azimuth values in radians (same length as za_array). The azimuth + here has the UVBeam convention: North of East(East=0, North=pi/2) + + za_array : array_like + Zenith angle values in radians (same length as az_array). + + freq_array : array_like + Frequency values to evaluate at. + + Returns + ------- + interp_data : array_like + Array of beam values, shape (Naxes_vec, Nspws, Nfeeds or Npols, + Nfreqs or freq_array.size if freq_array is passed, + Npixels/(Naxis1, Naxis2) or az_array.size if az/za_arrays are passed) + + interp_basis_vector : array_like + Array of interpolated basis vectors (or self.basis_vector_array + if az/za_arrays are not passed), shape: (Naxes_vec, Ncomponents_vec, + Npixels/(Naxis1, Naxis2) or az_array.size if az/za_arrays are passed) + """ + + # Empty data array + interp_data = np.zeros((2, 1, 2, freq_array.size, za_array.size)) + + values = np.empty((freq_array.size, za_array.size)) + values[:] = np.cos(za_array) + + interp_data[1, 0, 0, :, :] = values + interp_data[0, 0, 1, :, :] = values + interp_basis_vector = None + return interp_data, interp_basis_vector + + +def test_external_beam(): + """ + None of this should fail. + """ + + obs_latitude = -30.7215277777 + obs_longitude = 21.4283055554 + obs_height = 1073.0 + + # Create baselines + + # Random antenna locations + number = 3 + x = np.random.random(number)*40 + y = np.random.random(number)*40 + z = np.random.random(number)*1 + ants = {} + for i in range(number): + ants[i] = ( x[i], y[i], z[i] ) + + baselines = [] + for i in range(len(ants)): + for j in range(i+1, len(ants)): + bl = Baseline(ants[i], ants[j], i, j) + baselines.append(bl) + + # Set frequency axis + freqs = np.linspace(100e6, 150e6, 2, endpoint=False) + Nfreqs = len(freqs) + + # Set times + times = np.linspace(2458000.1, 2458000.6, 2) + Ntimes = len(times) + + # Create the observatory + fov = 360 # Deg + obs = Observatory(obs_latitude, obs_longitude, obs_height, array=baselines, freqs=freqs) + obs.set_pointings(times) + obs.set_fov(fov) + + # Create external beam + one_beam = ExternalBeam() + beam = [ one_beam for i in range(len(ants)) ] + + # set beam + obs.set_beam(beam) + + # create a noise-like sky + np.random.seed(0) + Nside = 64 + eor = construct_skymodel("flat_spec", freqs=freqs, Nside=Nside, ref_chan=0, sigma=1e-3) + + # Compute Visibilities for eor + eor_vis, times, bls = obs.make_visibilities(eor) +