Source code for use_everybeam

import numpy as np
from astropy.coordinates import ITRS, SkyCoord, AltAz, EarthLocation
from astropy.time import Time, TimeDelta
import astropy.units as u
import argparse
from wodenpy.use_libwoden.create_woden_struct_classes import Woden_Struct_Classes
import erfa
from typing import Union, Tuple
import concurrent.futures
from line_profiler import profile
import os
import astropy
##Are we just making online documentation? If so, don't import everybeam
##Installing everybeam is non-trivial, so trying to get readthedocs to install
##it is a waste of time
read_the_docs_build = os.environ.get('READTHEDOCS', None) == 'True'


from wodenpy.wodenpy_setup.run_setup import check_for_library
have_everybeam = check_for_library('everybeam')

[docs] class EB_fake: """ A fake `everybeam` class so we can build the documentation online in ReadTheDocs without installing `everybeam`, which is non-trivial """ def __init__(self): """ Just set everything that is ever used to None """ self.OSKAR = None self.LOFAR = None self.MWA = None self.MWALocal = None self.load_telescope = None self.Telescope = None
if have_everybeam: import everybeam as eb else: eb = EB_fake() # if read_the_docs_build: # eb = EB_fake() # else: # import everybeam as eb ##This call is so we can use it as a type annotation woden_struct_classes = Woden_Struct_Classes() Source_Catalogue = woden_struct_classes.Source_Catalogue Woden_Settings = woden_struct_classes.Woden_Settings
[docs] def radec_to_xyz(ra : float, dec : float, time : Time): """ Convert RA and Dec ICRS coordinates to ITRS cartesian coordinates. Taken from the everybeam documentation https://everybeam.readthedocs.io/en/latest/tree/demos/lofar-array-factor.html Args: ra (astropy.coordinates.Angle): Right ascension dec (astropy.coordinates.Angle): Declination time float astropy time instance Returns: pointing_xyz (ndarray): NumPy array containing the ITRS X, Y and Z coordinates """ coord = SkyCoord(ra*u.rad, dec*u.rad, frame='icrs') coord_itrs = coord.transform_to(ITRS(obstime=time)) return np.asarray(coord_itrs.cartesian.xyz.transpose())
[docs] def load_OSKAR_telescope(ms_path : str, response_model = "skala40_wave", use_differential_beam : bool = True) -> eb.OSKAR: # type: ignore """Load an OSKAR telescope from a measurement set. Parameters ---------- ms_path : str Path to the measurement set response_model : str, optional Response model to use, by default "skala40_wave" use_differential_beam : bool, optional Use the differential beam a.k.a return a "normalised" beam, by default True Returns ------- eb.OSKAR Telescope object """ print("OSKAR response model", response_model) # Load the telescope telescope = eb.load_telescope( ms_path, use_differential_beam=use_differential_beam, element_response_model=response_model, ) # assert type(telescope) == eb.OSKAR if type(telescope) != eb.OSKAR: print(f'WARNING: Telescope specified in {ms_path} is not an OSKAR telescope. Proceeding, but you might get nonsense results.') return telescope
[docs] def load_LOFAR_telescope(ms_path : str, response_model : str = "hamaker", use_differential_beam : bool = False) -> eb.LOFAR: # type: ignore """Load an LOFAR telescope from a measurement set. Settings lifted directly from https://everybeam.readthedocs.io/en/latest/tree/demos/lofar-lobes.html Parameters ---------- ms_path : str Path to the measurement set response_model : str, optional Response model to use, by default "hamaker" use_differential_beam : bool, optional Use the differential beam a.k.a return a "normalised" beam, by default False Returns ------- eb.LOFAR Telescope object """ telescope = eb.load_telescope(ms_path, use_differential_beam=use_differential_beam, element_response_model=response_model) if type(telescope) != eb.LOFAR: print(f'WARNING: Telescope specified in {ms_path} is not an OSKAR telescope. Proceeding, but you might get nonsense results.') return telescope
[docs] def load_MWA_telescope(ms_path : str, coeff_path : str, use_local_mwa : bool = True) -> Union[eb.MWA, eb.MWALocal]: # type: ignore """Load an MWA telescope from a measurement set. Parameters ---------- ms_path : str Path to the measurement set response_model : str, optional Response model to use, by default "lobes" use_local_mwa : bool, optional Use the local MWA model, which takes za/az instead of RA/Dec. Defaults to True Returns ------- Union[eb.MWA, eb.MWALocal] Telescope object, either MWA (if use_local_mwa is False) or MWALocal (if use_local_mwa is True) """ ## Load the telescope. Adding use_differential_beam seems to do nothing, ## so always leave it as False telescope = eb.load_telescope(ms_path, use_differential_beam=False, coeff_path=coeff_path, use_local_mwa=use_local_mwa) # assert type(telescope) == eb.MWA if type(telescope) != eb.MWA and type(telescope) != eb.MWALocal: print(f'WARNING: Telescope specified in {ms_path} is not an MWA telescope. Proceeding, but you might get nonsense results.') return telescope
[docs] def eb_local_xyz_from_radec(ra : float, dec : float, altaz_frame : astropy.coordinates.AltAz, delta_az : float = (1/2)*np.pi, negative_azimuth=True): """ Get the local cartesian coordinates used by EveryBeam from a given RA, Dec, and AltAz frame. This function transforms the given right ascension (RA) and declination (Dec) into local cartesian coordinates based on the provided AltAz frame. The azimuth is adjusted by reversing it and adding 90 degrees (or a specified delta) to match the expected coordinates. This function is used to emulate the way the EveryBeam does it's parallactic rotation. `delta_az` is set to π/2 by experiment to match outputs from EveryBeam. Not really needed, as when can use EveryBeam to do the parallactic rotation for everything aside the MWA beam. Parameters ----------- ra : float Right ascension in radians. dec : float Declination in radians. altaz_frame : `astropy.coordinates.AltAz` The AltAz frame to transform the coordinates into. delta_az : float, optional The azimuth adjustment in radians. Default is π/2. negative_azimuth : bool, optional If True, the azimuth is reversed before adding the delta. Default is True. Returns -------- numpy.ndarray A 3xN array of the local cartesian coordinates. """ coord = SkyCoord(ra=ra*u.rad, dec=dec*u.rad, frame='icrs') coord = coord.transform_to(altaz_frame) delta_az = delta_az * u.rad if negative_azimuth: updated_coord = SkyCoord(az=-coord.az + delta_az, alt=coord.alt, distance=coord.distance, frame=coord.frame) else: updated_coord = SkyCoord(az=coord.az + delta_az, alt=coord.alt, distance=coord.distance, frame=coord.frame) return np.array(updated_coord.cartesian.xyz.transpose())
[docs] def eb_north_east(direction : np.ndarray, ncp_t : np.ndarray) -> Tuple[np.ndarray, np.ndarray]: """ Compute the north and east vectors given a direction ITRF vector, and the ITRF vector towards the north celestial pole. This function calculates the east vector by taking the cross product of the normal vector (ncp_t) and the direction vector, and then normalizes it. The north vector is then calculated as the cross product of the direction vector and the east vector. Translated from EveryBeam station.cc Station::ComputeElementResponse const vector3r_t east = normalize(cross(ncp_t, direction)); const vector3r_t north = cross(direction, east); options.east = east; options.north = north; Parameters ------------ direction : np.ndarray A 3-element array representing the direction vector. ncp_t : np.ndarray A 3-element array representing the normal vector. Returns -------- Tuple[np.ndarray, np.ndarray] A tuple containing the north and east vectors as 3-element arrays. """ east = np.cross(ncp_t, direction) east = east/np.linalg.norm(east) north = np.cross(direction, east) return north, east
[docs] def calc_everybeam_rotation(direction : np.ndarray, north : np.ndarray, east : np.ndarray) -> np.ndarray: """Given an ITRF 3-vector in the direction of interest `direction`, and the associated north and east vectors, calculate the 2x2 rotation matrix to rotate by parallactic angle. Translated from EveryBeam beamformer.cc BeamFormer::LocalResponse const vector3r_t e_phi = normalize(cross(direction)); const vector3r_t e_theta = cross(e_phi, direction); result *= {dot(e_theta, options.north), dot(e_theta, options.east), dot(e_phi, options.north), dot(e_phi, options.east)}; Parameters ------------ direction : np.ndarray A 3-element array representing the direction ITRF vector. north : np.ndarray A 3-element array representing the north vector. east : np.ndarray A 3-element array representing the east vector. Returns -------- np.ndarray A 2x2 rotation matrix. """ e_phi = np.cross([0.0, 0.0, 1.0], direction) e_phi = e_phi/np.linalg.norm(e_phi) e_theta = np.cross(e_phi, direction) e_theta = e_theta/np.linalg.norm(e_theta) rot_matrix = np.array([[np.dot(e_theta, north), np.dot(e_theta, east)], [np.dot(e_phi, north), np.dot(e_phi, east)]]) return rot_matrix
# @profile
[docs] def run_everybeam(ras: np.ndarray, decs: np.ndarray, beam_ra0: float, beam_dec0: float, j2000_latitudes: np.ndarray, j2000_lsts: np.ndarray, current_latitude: float, current_longitude: float, times: np.ndarray, freqs: np.ndarray, telescope: eb.Telescope, # type: ignore station_ids: np.ndarray, apply_beam_norms: bool = True, reorder_jones: bool = False, element_only: bool = False, eb_rotate: bool = False, parallactic_rotate: bool = False, para_angle_offset: float = 0) -> np.ndarray: """ Calculate the Jones matrices for a given set of coordinates, times, frequencies, and station ids using the EveryBeam library. `j2000_latitudes` should be the array latitude as precessed back to J2000, with `j2000_lsts` being the matching LST in J2000. `current_latitude` and `current_longitude` should be latitude and longitude of the array at the time of the observation. `telescope` should be an EveryBeam telescope object. Parameters ------------ ras : np.ndarray Right ascensions of the coordinates in radians. decs : np.ndarray Declinations of the coordinates in radians. beam_ra0 : float Right ascension of the beam center in radians. beam_dec0 : float Declination of the beam center in radians. j2000_latitudes : np.ndarray Latitudes in J2000 coordinates. j2000_lsts : np.ndarray Local sidereal times in J2000 coordinates. current_latitude : float Current latitude in radians. current_longitude : float Current longitude in radians. times : np.ndarray Array of observation times. freqs : np.ndarray Array of frequencies. telescope : (eb.Telescope): Telescope object from the EveryBeam library. station_ids : np.ndarray Array of station IDs. apply_beam_norms : bool, optional Whether to apply beam normalisation. Defaults to True. Achieved by calculating the beam response at beam centre, and multiplying all Jones by the inverse of this central beam response. reorder_jones : bool, optional Whether to reorder the Jones matrices. Defaults to False. Just rearranges the Jones matrix from [[0,0, 0,1,], [1,0, 1,1]] to [[1,1, 1,0,], [0,1, 0,0]]. element_only : bool, optional Whether to use only the element response. Defaults to False. Use this to look at the dipole response only, not the beam formed response. eb_rotate : bool, optional Whether to apply parallactic rotation using EveryBeam. Defaults to False. Should probably be used for everything apart from MWA beams. parallactic_rotate : bool, optional Whether to apply parallactic angle rotation using `wodenpy`. Defaults to False. Should be True for MWA beams if you want rotation. If True for a non-MWA beam, `wodenpy` should match the output as if `eb_rotate` was True. para_angle_offset : float, optional Offset to add to the parallactic angle. Defaults to 0. Returns -------- np.ndarray The calculated Jones matrices with shape (num_stations, num_times, num_freqs, num_coords, 2, 2). """ num_stations = len(station_ids) num_times = len(times) num_freqs = len(freqs) num_coords = len(ras) all_output_jones = np.zeros((num_stations, num_times, num_freqs, num_coords, 2, 2), dtype=np.complex128)*np.nan non_itrf_beams = [eb.MWA, eb.MWALocal] if parallactic_rotate: if type(telescope) not in non_itrf_beams: coords = SkyCoord(ras*u.rad, decs*u.rad, frame='icrs') location = EarthLocation(lat=current_latitude*u.rad, lon=current_longitude*u.rad) for time_ind, time in enumerate(times): if type(telescope) in non_itrf_beams: comp_has = j2000_lsts[time_ind] - ras azs, els = erfa.hd2ae(comp_has, decs, j2000_latitudes[time_ind]) zas = np.pi/2 - els if parallactic_rotate: beam_ha0 = j2000_lsts[time_ind] - beam_ra0 beam_az0, beam_el0 = erfa.hd2ae(beam_ha0, beam_dec0, j2000_latitudes[time_ind]) beam_za0 = np.pi/2 - beam_el0 else: phase_itrf = radec_to_xyz(beam_ra0, beam_dec0, time) dir_itrfs = radec_to_xyz(ras, decs, time) if parallactic_rotate: altaz_frame = AltAz(obstime=time, location=location) ncp_t = eb_local_xyz_from_radec(0, np.radians(90), altaz_frame) dir_local = eb_local_xyz_from_radec(ras, decs, altaz_frame) time_mjd_secs = time.mjd*3600*24 if parallactic_rotate: has = j2000_lsts[time_ind] - ras para_angles = erfa.hd2pa(has, decs, j2000_latitudes[time_ind]) rot_matrix = np.empty((num_coords, 2,2)) if type(telescope) in non_itrf_beams: rot_matrix[:,0,0] = np.sin(-para_angles) rot_matrix[:,0,1] = -np.cos(-para_angles) rot_matrix[:,1,0] = -np.cos(-para_angles) rot_matrix[:,1,1] = -np.sin(-para_angles) else: for dir_ind, dir_itrf in enumerate(dir_itrfs): dir_az = dir_local[dir_ind] north, east = eb_north_east(dir_az, ncp_t) rot = calc_everybeam_rotation(dir_az, north, east) rot_matrix[dir_ind] = rot for station_ind, station_id in enumerate(station_ids): for freq_ind, freq in enumerate(freqs): if apply_beam_norms: if type(telescope) == eb.MWA: ##Get the response norm_jones = telescope.station_response(time_mjd_secs, station_id, freq, beam_ra0, beam_dec0) elif type(telescope) == eb.MWALocal: norm_jones = telescope.station_response(time_mjd_secs, station_id, freq, beam_az0, beam_za0) if type(telescope) in non_itrf_beams: if parallactic_rotate: ha0 = j2000_lsts[time_ind] - beam_ra0 para_angles = erfa.hd2pa(ha0, beam_dec0, j2000_latitudes[time_ind]) rot = np.empty((2,2)) rot[0,0] = np.sin(-para_angles) rot[0,1] = -np.cos(-para_angles) rot[1,0] = -np.cos(-para_angles) rot[1,1] = -np.sin(-para_angles) else: element_id = 0 if element_only: norm_jones = telescope.element_response(time_mjd_secs, station_id, element_id, freq, phase_itrf, rotate=eb_rotate) else: norm_jones = telescope.station_response(time_mjd_secs, station_id, freq, phase_itrf, phase_itrf, rotate=eb_rotate) if parallactic_rotate: dir_phase_local = eb_local_xyz_from_radec(beam_ra0, beam_dec0, altaz_frame) north, east = eb_north_east(dir_phase_local, ncp_t) rot = calc_everybeam_rotation(dir_phase_local, north, east) if parallactic_rotate: norm_jones = np.matmul(norm_jones, rot) for coord_ind, (ra, dec) in enumerate(zip(ras, decs)): ##Only MWA uses ra,dec as a direct input if type(telescope) == eb.MWA: response = telescope.station_response(time_mjd_secs, station_id, freq, ra, dec) ##Only MWALocal uses az,za as a direct input elif type(telescope) == eb.MWALocal: response = telescope.station_response(time_mjd_secs, station_id, freq, zas[coord_ind], azs[coord_ind]) ##Everything else uses ITRF coordinates else: if element_only: response = telescope.element_response(time_mjd_secs, station_id, 0, freq, dir_itrfs[coord_ind], rotate=eb_rotate) else: response = telescope.station_response(time_mjd_secs, station_id, freq, dir_itrfs[coord_ind], phase_itrf, rotate=eb_rotate) all_output_jones[station_ind, time_ind, freq_ind, coord_ind] = response if parallactic_rotate: ##Parallactic angle doesn't change per station or freq, but ##if we are normalising the beam, we want to rotate before we normalise ##So do the rotation now rot_jones = np.einsum('klm,kmn->kln', all_output_jones[station_ind, time_ind, freq_ind, :, :, :], rot_matrix) all_output_jones[station_ind, time_ind, freq_ind, :, :, :] = rot_jones if apply_beam_norms: ##Each station, time, and freq gets it's own normalisation ##Same 2x2 normalisation for all directions inv_beam_norms = np.linalg.inv(norm_jones) output_jones = np.einsum('lm,kmn->kln', inv_beam_norms, all_output_jones[station_ind, time_ind, freq_ind, :, :, :]) all_output_jones[station_ind, time_ind, freq_ind, :, :, :] = output_jones if reorder_jones: ##swap all_output_jones[:,:,:,:,0,0] with all_output_jones[:,:,:,:,1,1] all_output_jones[:, :, :, :, [0, 1], [0, 1]] = all_output_jones[:, :, :, :, [1, 0], [1, 0]] ##swap all_output_jones[:,:,:,:,0,1] with all_output_jones[:,:,:,:,1,0] all_output_jones[:, :, :, :, [0, 1], [1, 0]] = all_output_jones[:, :, :, :, [1, 0], [0, 1]] return all_output_jones
[docs] def run_everybeam_thread(num_threads : int, thread_id : int, ms_path : str, coeff_path : str, ras : np.ndarray, decs : np.ndarray, ra0 : float, dec0 : float, j2000_latitudes : np.ndarray, j2000_lsts : np.ndarray, current_latitude : float, current_longitude : float, times : np.ndarray, freqs : np.ndarray, station_ids : np.ndarray, use_differential_beam : bool = True, apply_beam_norms : bool = True, reorder_jones : bool = False, element_only : bool = False, eb_rotate : bool = False, parallactic_rotate : bool = True, para_angle_offset : float = 0, element_response_model='hamaker', use_local_mwa : bool = True) -> Tuple[np.ndarray, int]: """ Thread function called by `run_everybeam_over_threads` to calculate the EveryBeam response in parrallel. Calls `run_everybeam` with a subset of the coordinates; see `run_everybeam` for more details of the parameters. Creates a new EveryBeam telescope object from `ms_path` for each thread. This has to be done because `concurrent.futures.ProcessPoolExecutor` has to pickle the function and all it's arguments, and EveryBeam objects can't be pickled. This is somewhat wasteful but I can't work out a better way to make things parallel. Parameters ------------ num_threads : int Number of threads being in call by `run_everybeam_over_threads`. thread_id : int ID of the current thread. Useds to work out what chunk of `ras` and `decs` to process. ms_path : str Path to the measurement set to load the EveryBeam telescope from. ras : np.ndarray Right ascensions of the coordinates in radians. decs : np.ndarray Declinations of the coordinates in radians. beam_ra0 : float Right ascension of the beam center in radians. beam_dec0 : float Declination of the beam center in radians. j2000_latitudes : np.ndarray Latitudes in J2000 coordinates. j2000_lsts : np.ndarray Local sidereal times in J2000 coordinates. current_latitude : float Current latitude in radians. current_longitude : float Current longitude in radians. times : np.ndarray Array of observation times. freqs : np.ndarray Array of frequencies. station_ids : np.ndarray Array of station IDs. apply_beam_norms : bool, optional Whether to apply beam normalisation. Defaults to True. Achieved by calculating the beam response at beam centre, and multiplying all Jones by the inverse of this central beam response. reorder_jones : bool, optional Whether to reorder the Jones matrices. Defaults to False. Just rearranges the Jones matrix from [[0,0, 0,1,], [1,0, 1,1]] to [[1,1, 1,0,], [0,1, 0,0]]. element_only : bool, optional Whether to use only the element response. Defaults to False. Use this to look at the dipole response only, not the beam formed response. eb_rotate : bool, optional Whether to apply parallactic rotation using EveryBeam. Defaults to False. Should probably be used for everything apart from MWA beams. parallactic_rotate : bool, optional Whether to apply parallactic angle rotation using `wodenpy`. Defaults to False. Should be True for MWA beams if you want rotation. If True for a non-MWA beam, `wodenpy` should match the output as if `eb_rotate` was True. para_angle_offset : float, optional Offset to add to the parallactic angle. Defaults to 0. element_response_model : str, optional The Everybeam element response model to use. Defaults to 'hamaker'. Avaible options are 'hamaker' (LOFAR), 'skala40_wave' (OSKAR), and 'MWA' (MWA). use_local_mwa : bool, optional Whether to use the local MWA model. Defaults to True. The local MWA model takes za/az instead of RA/Dec. Returns -------- Tuple[np.ndarray, int] The calculated Jones matrices with shape (num_stations, num_times, num_freqs, num_coords_in_thread, 2, 2), as well as the thread ID. Use the thread ID to insert this thread output into the correct place in the final Jones matrix. """ telescope = eb.load_telescope(ms_path, use_differential_beam=use_differential_beam, coeff_path=coeff_path, element_response_model=element_response_model, use_local_mwa=use_local_mwa) num_coords = len(ras) coords_per_thread = int(np.ceil(num_coords / num_threads)) low_coord = thread_id * coords_per_thread high_coord = (thread_id + 1) * coords_per_thread print(f"Thread {thread_id} processing coords {low_coord} to {high_coord}") jones = run_everybeam(ras[low_coord:high_coord], decs[low_coord:high_coord], ra0, dec0, j2000_latitudes, j2000_lsts, current_latitude, current_longitude, times, freqs, telescope, station_ids, apply_beam_norms=apply_beam_norms, reorder_jones=reorder_jones, element_only=element_only, eb_rotate=eb_rotate, parallactic_rotate=parallactic_rotate) print(f"Thread {thread_id} finished") return jones, thread_id
[docs] def run_everybeam_over_threads(num_threads : int, ms_path : str, coeff_path : str, ras : np.ndarray, decs : np.ndarray, ra0 : float, dec0 : float, j2000_latitudes : np.ndarray, j2000_lsts : np.ndarray, current_latitude : float, current_longitude : float, times : np.ndarray, freqs : np.ndarray, station_ids : np.ndarray, use_differential_beam : bool = True, apply_beam_norms : bool = True, reorder_jones : bool = False, element_only : bool = False, eb_rotate : bool = False, parallactic_rotate : bool = True, use_local_mwa : bool = True, para_angle_offset : float = 0, element_response_model='hamaker'): """ Runs `run_everybeam` in parallel over `num_threads` threads, using `concurrent.futures.ProcessPoolExecutor`. See `run_everybeam` for more details of what each parameter does. Creates a new EveryBeam telescope object from `ms_path` for each thread. This has to be done because `concurrent.futures.ProcessPoolExecutor` has to pickle the function and all it's arguments, and EveryBeam objects can't be pickled. This is somewhat wasteful but I can't work out a better way to make things parallel. Parameters ------------ num_threads : int Number of threads being in call by `run_everybeam_over_threads`. ms_path : str Path to the measurement set to load the EveryBeam telescope from. ras : np.ndarray Right ascensions of the coordinates in radians. decs : np.ndarray Declinations of the coordinates in radians. beam_ra0 : float Right ascension of the beam center in radians. beam_dec0 : float Declination of the beam center in radians. j2000_latitudes : np.ndarray Latitudes in J2000 coordinates. j2000_lsts : np.ndarray Local sidereal times in J2000 coordinates. current_latitude : float Current latitude in radians. current_longitude : float Current longitude in radians. times : np.ndarray Array of observation times. freqs : np.ndarray Array of frequencies. station_ids : np.ndarray Array of station IDs. apply_beam_norms : bool, optional Whether to apply beam normalisation. Defaults to True. Achieved by calculating the beam response at beam centre, and multiplying all Jones by the inverse of this central beam response. reorder_jones : bool, optional Whether to reorder the Jones matrices. Defaults to False. Just rearranges the Jones matrix from [[0,0, 0,1,], [1,0, 1,1]] to [[1,1, 1,0,], [0,1, 0,0]]. element_only : bool, optional Whether to use only the element response. Defaults to False. Use this to look at the dipole response only, not the beam formed response. eb_rotate : bool, optional Whether to apply parallactic rotation using EveryBeam. Defaults to False. Should probably be used for everything apart from MWA beams. parallactic_rotate : bool, optional Whether to apply parallactic angle rotation using `wodenpy`. Defaults to False. Should be True for MWA beams if you want rotation. If True for a non-MWA beam, `wodenpy` should match the output as if `eb_rotate` was True. para_angle_offset : float, optional Offset to add to the parallactic angle. Defaults to 0. element_response_model : str, optional The Everybeam element response model to use. Defaults to 'hamaker'. Avaible options are 'hamaker' (LOFAR), 'skala40_wave' (OSKAR), and 'MWA' (MWA). use_local_mwa : bool, optional Whether to use the local MWA model. Defaults to True. The local MWA model takes za/az instead of RA/Dec. Returns -------- np.ndarray The calculated Jones matrices with shape (num_stations, num_times, num_freqs, num_coord, 2, 2) """ with concurrent.futures.ProcessPoolExecutor(max_workers=num_threads) as executor: future_data = [executor.submit(run_everybeam_thread, num_threads, thread_id, ms_path, coeff_path, ras, decs, ra0, dec0, j2000_latitudes, j2000_lsts, current_latitude, current_longitude, times, freqs, station_ids, use_differential_beam=use_differential_beam, apply_beam_norms=apply_beam_norms, reorder_jones=reorder_jones, element_only=element_only, eb_rotate=eb_rotate, parallactic_rotate=parallactic_rotate, para_angle_offset=para_angle_offset, element_response_model=element_response_model, use_local_mwa=use_local_mwa) for thread_id in range(num_threads)] all_jones_chunks = [] all_thread_ids = [] for future in concurrent.futures.as_completed(future_data): jones_chunk, thread_id = future.result() all_jones_chunks.append(jones_chunk) all_thread_ids.append(thread_id) num_stations = len(station_ids) num_times = len(times) num_freqs = len(freqs) num_coords = len(ras) all_jones = np.zeros((num_stations, num_times, num_freqs, num_coords, 2, 2), dtype=np.complex128)*np.nan coords_per_thread = int(np.ceil(num_coords / num_threads)) for jones_chunk, thread_id in zip(all_jones_chunks, all_thread_ids): low_coord = thread_id * coords_per_thread high_coord = (thread_id + 1) * coords_per_thread all_jones[:, :, :, low_coord:high_coord, :, :] = jones_chunk return all_jones