Raw File
# -*- mode: python; coding: utf-8 -*-
# Copyright (c) 2019 Radio Astronomy Software Group
# Licensed under the 2-clause BSD License
"""Read in the Sujinto et al. full embedded element MWA Beam."""
import warnings

import numpy as np
import h5py
from scipy.special import factorial, lpmv  # associated Legendre function

from . import UVBeam
from .. import utils as uvutils

__all__ = ["P1sin", "P1sin_array", "MWABeam"]


def P1sin(nmax, theta):
    """
    Create the Legendre function flavors for FF expansion using spherical waves.

    Note this is not vectorized so is a bit slow, but it handles the special
    case of theta = 0 and pi. We primarily use the vectorized version
    (`P1sin_array`), but call this to handle the special cases.

    See:
    Calculating Far-Field Radiation Based on FEKO Spherical Wave Coefficients,
    draft 10 June 2015.
    Available at pyuvdata/docs/references/Far_field_spherical_FEKO_draft2.pdf
    This memo gives a full description of the equations implemented here,
    including descriptions of the approximations and numerical approaches used.
    In line comments below are helpful reminders, but see the memo for the full
    detail.
    Also see Sokolowski, M. et al, "Calibration and Stokes Imaging with Full
    Embedded Element Primary Beam Model for the Murchison Widefield Array",
    PASA, 2017 (10.1017/pasa.2017.54) for details specific to the MWA.

    Parameters
    ----------
    nmax : int
        Maximum n from FEKO Q1mn and Q2mn, n must be >=1
    theta : float
        The argument of the cosine or sine function used in the associated
        Legendre functions, in radians.

    Returns
    -------
    P_sin : array of float
        P_{n}^{abs(m)}(cos(theta))/sin(theta) with FEKO order M,N.
        Shape (nmax ** 2 + 2 * nmax).
    P1 : array of float
        P_{n}^{abs(m)+1}(cos(theta)) with FEKO order M,N.
        Shape (nmax ** 2 + 2 * nmax).

    """
    # initialize for nmax, we have 2(1+...+nmax)+nmax=nmax^2+2*nmax long array
    P_sin = np.zeros((nmax ** 2 + 2 * nmax))
    P1 = np.zeros((nmax ** 2 + 2 * nmax))

    # theta arguments
    cos_th = np.cos(theta)
    sin_th = np.sin(theta)
    delta_cos = 1e-6  # for slope estimation

    # step from 1 to nmax
    for n in range(1, nmax + 1):
        # legendre P_{n}^{abs(m)=0...n} (cos_th)
        orders = np.arange(0, n + 1)
        orders = orders.reshape(n + 1, 1)
        P = lpmv(orders, n, cos_th)

        # THESE ARE THE SAME:
        # legendre(2,0:0.1:0.2) (matlab)
        # scipy:
        # a=np.arange(0,3)
        # a=a.reshape(3,1)
        # lpmv(b,2,np.arange(0,0.3,0.1))

        # P_{n}^{abs(m)+1} (cos_th)
        Pm1 = np.append(P[1::], 0)
        Pm1 = Pm1.reshape(len(Pm1), 1)

        # P_{n}^{abs(m)}(cos_th)/sin_th
        Pm_sin = np.zeros((n + 1, 1))  # initialize

        if cos_th == 1:
            # special treatment depending on m;
            # for m=0, Pm_sin=inf so, the product m*Pm_sin is zero;
            # for m=1, we need a substitution
            # m>=2, value is 0, so initial values are OK
            # The first approach, to just use the analytical derivative
            # is not stable for n>~45
            # Instead use slope estimate with a small delta_cos
            # Pn(cos x)/sin x = -dPn(cos_th)/dcos_th
            Pm_cos_delta_cos = lpmv(orders, n, cos_th - delta_cos)
            # backward difference
            Pm_sin[1, 0] = -(P[0] - Pm_cos_delta_cos[0]) / delta_cos

        elif cos_th == -1:
            # The first approach, to just use the analytical derivative
            # is not stable for n>~45
            # Instead use slope estimate with a small delta_cos
            # Pn(cos x)/sin x = -dPn(cos_th)/dcos_th
            Pm_cos_delta_cos = lpmv(orders, n, cos_th - delta_cos)
            # forward difference
            Pm_sin[1, 0] = -(Pm_cos_delta_cos[0] - P[0]) / delta_cos
        else:
            Pm_sin = P / sin_th

        # accumulate Psin and P1 for the m values
        ind_start = (n - 1) ** 2 + 2 * (n - 1)  # start index to populate
        ind_stop = n ** 2 + 2 * n  # stop index to populate
        # assign
        P_sin[np.arange(ind_start, ind_stop)] = np.append(
            np.flipud(Pm_sin[1::, 0]), Pm_sin
        )
        P1[np.arange(ind_start, ind_stop)] = np.append(np.flipud(Pm1[1::, 0]), Pm1)

    return P_sin, P1


def P1sin_array(nmax, theta):
    """
    Calculate P^abs(m)_n(cos(theta))/sin(theta) and P^(abs(m)+1)_n(cos(theta)).

    Similar to the "P1sin" function, but calculates for all theta in one go.
    At the end of the function, patches are made using the original P1sin function
    to solve the 0/0 issue.


    Parameters
    ----------
    nmax : int
        Maximum n from FEKO Q1mn and Q2mn, n must be >=1
    theta : array of float
        The argument of the cosine or sine functions used in the associated
        Legendre functions, in radians.

    Returns
    -------
    P_sin : array of float
        P_{n}^{abs(m)}(cos(theta))/sin(theta) with FEKO order M,N.
        Shape (nmax ** 2 + 2 * nmax, theta.size).
    P1 : array of float
        P_{n}^{abs(m)+1}(cos(theta)) with FEKO order M,N.
        Shape (nmax ** 2 + 2 * nmax, theta.size).

    """
    cos_th = np.cos(theta)
    sin_theta = np.sin(theta)

    # Make sure that we don't divide by 0 (sin(0) = sin(pi) = 0 ) proper results
    # are inserted at the end of this function. Set to NaN for now
    sin_theta[(theta == 0) | (theta == np.pi)] = np.NaN

    # create at forehand
    P_sin = np.zeros((nmax ** 2 + 2 * nmax, np.size(theta)))
    P1 = np.zeros((nmax ** 2 + 2 * nmax, np.size(theta)))
    for n in range(1, nmax + 1):
        # legendre P_{n}^{abs(m)=0...n} (cos_th)
        orders = np.arange(0, n + 1)
        orders = orders.reshape(n + 1, 1)

        # fetch entire matrix in one go (for a particular n)
        # in theory, fetching for all n in one go should also be possible
        P = lpmv(orders, n, cos_th)

        # P_{n}^{abs(m)+1} (cos_th)
        Pm1 = np.vstack([P[1::, :], np.zeros((1, np.size(theta)))])

        # P_{n}^{abs(m)}(u)/sin_th
        Pm_sin = P / sin_theta

        # accumulate Psin and P1 for the m values
        # start index to populate
        ind_start = (n - 1) ** 2 + 2 * (n - 1)
        # stop index to populate
        ind_stop = n ** 2 + 2 * n
        # assign
        P_sin[np.arange(ind_start, ind_stop), :] = np.vstack(
            [np.flipud(Pm_sin[1::, :]), Pm_sin]
        )
        P1[np.arange(ind_start, ind_stop), :] = np.vstack([np.flipud(Pm1[1::, :]), Pm1])

    # fix for theta = 0 and theta = pi
    # (properly handled in P1sin, so use that function)
    P_sin[:, theta == 0] = np.array([P1sin(nmax, 0)[0]]).transpose()
    P_sin[:, theta == np.pi] = np.array([P1sin(nmax, np.pi)[0]]).transpose()

    return P_sin.transpose(), P1.transpose()


class MWABeam(UVBeam):
    """
    Defines an MWA-specific subclass of UVBeam for representing MWA beams.

    This class should not be interacted with directly, instead use the
    read_mwa_beam method on the UVBeam class.

    This is based on https://github.com/MWATelescope/mwa_pb/ but we don’t import
    that module because it’s not python 3 compatible
    """

    def _read_metadata(self, h5filepath):
        """
        Get metadata (frequencies, polarizations, dipole numbers) from input file.

        Parameters
        ----------
        h5filepath : str
            path to input h5 file containing the MWA full embedded element spherical
            harmonic modes.

        Returns
        -------
        freqs_hz : array of int
            Frequencies in Hz present in the file.
        pol_names : list of str
            Polarizations present in the file.
        dipole_names :
            Dipoles names present in the file.
        max_length : dict
            Dictionary keyed on pol and freq, giving max number of modes in the
            file for each pol and freq.
        """
        pol_names = set()
        dipole_names = set()
        freqs_hz = set()
        other_names = []
        max_length = {}
        with h5py.File(h5filepath, "r") as h5f:
            for name in h5f.keys():
                if name.startswith("X") or name.startswith("Y"):
                    pol = name[0]
                    dipole, freq = name[1:].split("_")
                    pol_names.add(pol)
                    dipole_names.add(dipole)
                    freq = np.int(freq)
                    freqs_hz.add(freq)

                    if pol not in max_length:
                        max_length[pol] = {}

                    this_length = h5f[name].shape[1] // 2
                    if freq not in max_length[pol]:
                        max_length[pol][freq] = this_length
                    elif this_length > max_length[pol][freq]:
                        max_length[pol][freq] = this_length

                else:
                    other_names.append(name)

        pol_names = sorted(pol_names)

        dipole_names = np.array(sorted(dipole_names))

        freqs_hz = np.array(sorted(freqs_hz))

        return freqs_hz, pol_names, dipole_names, max_length

    def _get_beam_modes(
        self,
        h5filepath,
        freqs_hz,
        pol_names,
        dipole_names,
        max_length,
        delays,
        amplitudes,
    ):
        """
        Get beam modes from input file and save as a dict to the object.

        Parameters
        ----------
        h5filepath : str
            path to input h5 file containing the MWA full embedded element spherical
            harmonic modes.
        freqs_hz : array of int
            Frequencies in Hz to get modes for. Must be present in the file.
        pol_names : list of str
            Polarizations  to get modes for. Must be present in the file.
        dipole_names : array of str
            Dipoles names present in the file.
        max_length : dict
            Dictionary keyed on pol and freq, giving max number of modes in the
            file for each pol and freq.
        delays : array of ints
            Array of MWA beamformer delay steps. Should be shape (n_pols, n_dipoles).
        amplitudes : array of floats
            Array of dipole amplitudes, these are absolute values
            (i.e. relatable to physical units). Should be shape (n_pols, n_dipoles).

        Returns
        -------
        beam_modes : dict
            A multi-level dict keyed on (in order) pol, freq, mode name (Q1, Q2, M, N).
        """
        beam_modes = {}
        for pol_i, pol in enumerate(pol_names):
            beam_modes[pol] = {}
            for freq in freqs_hz:
                # Calculate complex excitation voltages
                # convert delay to phase
                # 435e-12 is the delay step size in seconds (435 picosec)
                phases = 2 * np.pi * freq * (-delays[pol_i, :]) * 435e-12
                # complex excitation col voltage
                Vcplx = amplitudes[pol_i, :] * np.exp(1.0j * phases)

                Q1_accum = np.zeros(max_length[pol][freq], dtype=np.complex128)
                Q2_accum = np.zeros(max_length[pol][freq], dtype=np.complex128)

                # Read in modes
                with h5py.File(h5filepath, "r") as h5f:
                    Q_modes_all = h5f["modes"][()].T
                    Nmax = 0
                    M_accum = None
                    N_accum = None
                    for dp_i, dp in enumerate(dipole_names):
                        # re-initialise Q1 and Q2 for every dipole
                        Q1 = np.zeros(max_length[pol][freq], dtype=np.complex128)
                        Q2 = np.zeros(max_length[pol][freq], dtype=np.complex128)

                        # select spherical wave table
                        name = pol + dp + "_" + str(freq)
                        Q_all = h5f[name][()].T

                        # current length
                        my_len = np.max(Q_all.shape)
                        my_len_half = my_len // 2

                        # Get modes for this dipole
                        Q_modes = Q_modes_all[0:my_len, :]

                        # convert Qall to M, N, Q1, Q2 vectors for processing

                        # find s=1 and s=2 indices
                        # only find s1 and s2 for this dipole
                        # s = 1 and s = 2 refer to TE and TM modes, respectively
                        # see the Far_field_spherical_FEKO_draft2 memo under
                        # pyuvdata/docs/references/
                        s1 = Q_modes[0:my_len, 0] <= 1
                        s2 = Q_modes[0:my_len, 0] > 1

                        # grab m,n vectors
                        M = Q_modes[s1, 1]
                        N = Q_modes[s1, 2]

                        # update to the larger M and N
                        if np.max(N) > Nmax:
                            M_accum = M
                            N_accum = N
                            Nmax = np.max(N_accum)

                        # grab Q1mn and Q2mn and make them complex
                        Q1[0:my_len_half] = Q_all[s1, 0] * np.exp(
                            1.0j * np.deg2rad(Q_all[s1, 1])
                        )
                        Q2[0:my_len_half] = Q_all[s2, 0] * np.exp(
                            1.0j * np.deg2rad(Q_all[s2, 1])
                        )

                        # accumulate Q1 and Q2, scaled by excitation voltage
                        Q1_accum = Q1_accum + Q1 * Vcplx[dp_i]
                        Q2_accum = Q2_accum + Q2 * Vcplx[dp_i]

                    beam_modes[pol][freq] = {
                        "Q1": Q1_accum,
                        "Q2": Q2_accum,
                        "M": M_accum,
                        "N": N_accum,
                    }
        return beam_modes

    def _get_response(self, freqs_hz, pol_names, beam_modes, phi_arr, theta_arr):
        """
        Calculate full Jones matrix response (E-field) of beam on a regular az/za grid.

        Parameters
        ----------
        freqs_hz : array of int
            Frequencies in Hz to get modes for. Must be present in the file.
        pol_names : list of str
            Polarizations  to get modes for. Must be present in the file.
        beam_modes : dict
            A multi-level dict keyed on (in order) pol, freq, mode name (Q1, Q2, M, N).
        phi_arr : float or array of float
            azimuth angles (radians), north through east.
        theta_arr : float or array of float
            zenith angles (radian)

        Returns
        -------
        jones : array of float
            jones vectors, shape (Npol, 2, Nfreq, Nphi, Ntheta), e.g.
                [J_11=Xtheta J_12=Xphi]
                [J_21=Ytheta J_21=Yphi]

        """
        jones = np.zeros(
            (len(pol_names), 2, freqs_hz.size, phi_arr.size, theta_arr.size),
            dtype=np.complex128,
        )

        phi_arr = np.pi / 2 - phi_arr  # Convert to East through North (FEKO coords)
        phi_arr[phi_arr < 0] += 2 * np.pi  # 360 wrap

        for pol_i, pol in enumerate(pol_names):
            for freq_i, freq in enumerate(freqs_hz):
                M = beam_modes[pol][freq]["M"]
                N = beam_modes[pol][freq]["N"]
                Q1 = beam_modes[pol][freq]["Q1"]
                Q2 = beam_modes[pol][freq]["Q2"]

                # form P(cos(theta))/(sin\theta) and P^{m+1}(cos(theta))with
                # FEKO M,N order
                nmax = int(np.max(N))
                assert (
                    np.max(N) - nmax == 0
                ), "The maximum of N should be an integer value!"

                # calculate equation C_mn from equation 4 of
                # pyuvdata/docs/references/Far_field_spherical_FEKO_draft2.pdf
                # These are the normalization factors for the associated
                # Legendre function of order n and rank abs(m)
                C_MN = (
                    0.5 * (2 * N + 1) * factorial(N - abs(M)) / factorial(N + abs(M))
                ) ** 0.5

                # 1 for M<=0, -1 for odd M>0
                MabsM = np.ones(M.shape)
                MabsM[(M > 0) & (M % 2 != 0)] = -1

                # nomenclature:
                # T and P are the sky polarisations theta and phi
                # theta and phi are direction coordinates

                phi_comp = np.ascontiguousarray(
                    np.exp(1.0j * np.outer(phi_arr, range(-nmax, nmax + 1)))
                )

                (P_sin, P1) = P1sin_array(nmax, theta_arr)
                M_u = np.outer(np.cos(theta_arr), np.abs(M))
                phi_const = C_MN * MabsM / (N * (N + 1)) ** 0.5

                emn_T = (
                    (1.0j) ** N * (P_sin * (M_u * Q2 - M * Q1) + Q2 * P1) * phi_const
                )
                emn_P = (
                    (1.0j) ** (N + 1)
                    * (P_sin * (M * Q2 - Q1 * M_u) - Q1 * P1)
                    * phi_const
                )

                # Use a matrix multiplication to calculate Emn_P and Emn_T.
                # Sum results of Emn_P and emn_T for each unique M
                emn_P_sum = np.zeros(
                    (len(theta_arr), 2 * nmax + 1), dtype=np.complex128
                )
                emn_T_sum = np.zeros(
                    (len(theta_arr), 2 * nmax + 1), dtype=np.complex128
                )
                for m in range(-nmax, nmax + 1):
                    emn_P_sum[:, m + nmax] = np.sum(emn_P[:, M == m], axis=1)
                    emn_T_sum[:, m + nmax] = np.sum(emn_T[:, M == m], axis=1)

                Sigma_P = np.inner(phi_comp, emn_P_sum)
                Sigma_T = np.inner(phi_comp, emn_T_sum)

                jones[pol_i, 0, freq_i] = Sigma_T
                jones[pol_i, 1, freq_i] = -Sigma_P

        return jones

    def read_mwa_beam(
        self,
        h5filepath,
        delays=None,
        amplitudes=None,
        pixels_per_deg=5,
        freq_range=None,
        run_check=True,
        check_extra=True,
        run_check_acceptability=True,
    ):
        """
        Read in the full embedded element MWA beam.

        Parameters
        ----------
        h5filepath : str
            path to input h5 file containing the MWA full embedded element spherical
            harmonic modes. Download via
            `wget http://cerberus.mwa128t.org/mwa_full_embedded_element_pattern.h5`
        delays : array of ints
            Array of MWA beamformer delay steps. Should be shape (n_pols, n_dipoles).
        amplitudes : array of floats
            Array of dipole amplitudes, these are absolute values
            (i.e. relatable to physical units).
            Should be shape (n_pols, n_dipoles).
        pixels_per_deg : float
            Number of theta/phi pixels per degree. Sets the resolution of the beam.
        freq_range : array_like of float
            Range of frequencies to include in Hz, defaults to all available
            frequencies. Must be length 2.
        run_check : bool
            Option to check for the existence and proper shapes of
            required parameters after reading in the file.
        check_extra : bool
            Option to check optional parameters as well as required ones.
        run_check_acceptability : bool
            Option to check acceptable range of the values of
            required parameters after reading in the file.

        Returns
        -------
        None

        Raises
        ------
        ValueError
            If the amplitudes or delays are the wrong shape or there are delays
            greater than 32 or delays are not integer types.
            If the frequency range doesn't include any
            available frequencies.

        """
        freqs_hz, pol_names, dipole_names, max_length = self._read_metadata(h5filepath)

        n_dp = dipole_names.size
        n_pol = len(pol_names)

        if delays is None:
            delays = np.zeros([n_pol, n_dp], dtype="int")
        else:
            if not np.issubdtype(delays.dtype, np.integer):
                raise ValueError("Delays must be integers.")

        if amplitudes is None:
            amplitudes = np.ones([n_pol, n_dp])

        if amplitudes.shape != (n_pol, n_dp):
            raise ValueError(
                "amplitudes must be shape ({npol}, {nd})".format(npol=n_pol, nd=n_dp)
            )

        if delays.shape != (n_pol, n_dp):
            raise ValueError(
                "delays must be shape ({npol}, {nd})".format(npol=n_pol, nd=n_dp)
            )

        if (delays > 32).any():
            raise ValueError(
                "There are delays greater than 32: {delays}".format(delays=delays)
            )

        # check for terminated dipoles and reset delays and amplitudes
        terminated = delays == 32
        if (terminated).any():
            warnings.warn(
                "There are some terminated dipoles "
                "(delay setting 32). Setting the amplitudes and "
                "delays of terminated dipoles to zero."
            )
            delays[terminated] = 0
            amplitudes[terminated] = 0

        if freq_range is not None:
            if np.array(freq_range).size != 2:
                raise ValueError("freq_range must have 2 elements.")
            freqs_use = freqs_hz[
                np.nonzero((freqs_hz >= freq_range[0]) & (freqs_hz <= freq_range[1]))
            ]
            if freqs_use.size < 1:
                raise ValueError(
                    "No frequencies available in freq_range. "
                    "Available frequencies are between {fmin} Hz "
                    "and {fmax} Hz".format(fmin=np.min(freqs_hz), fmax=np.max(freqs_hz))
                )
            if freqs_use.size < 2:
                warnings.warn("Only one available frequency in freq_range.")
        else:
            freqs_use = freqs_hz

        beam_modes = self._get_beam_modes(
            h5filepath,
            freqs_hz,
            pol_names,
            dipole_names,
            max_length,
            delays,
            amplitudes,
        )

        n_phi = np.floor(360 * pixels_per_deg)
        n_theta = np.floor(90 * pixels_per_deg) + 1
        theta_arr = np.deg2rad(np.arange(0, n_theta) / pixels_per_deg)
        phi_arr = np.deg2rad(np.arange(0, n_phi) / pixels_per_deg)

        jones = self._get_response(freqs_use, pol_names, beam_modes, phi_arr, theta_arr)

        # work out zenith normalization
        # (MWA beams are peak normalized to 1 when pointed at zenith)

        # start filling in the object
        self.telescope_name = "MWA"
        self.feed_name = "MWA"
        self.feed_version = "1.0"
        self.model_name = "full embedded element"
        self.model_version = "1.0"
        self.history = (
            "Sujito et al. full embedded element beam, derived from "
            "https://github.com/MWATelescope/mwa_pb/"
        )

        delay_str_list = []
        gain_str_list = []
        for pol in range(n_pol):
            delay_str_list.append(
                "[" + ", ".join([str(x) for x in delays[pol, :]]) + "]"
            )
            gain_str_list.append(
                "[" + ", ".join([str(x) for x in amplitudes[pol, :]]) + "]"
            )
        delay_str = "[" + ", ".join(delay_str_list) + "]"
        gain_str = "[" + ", ".join(gain_str_list) + "]"

        self.history += "  delays set to " + delay_str + "  gains set to " + gain_str
        if not uvutils._check_history_version(self.history, self.pyuvdata_version_str):
            self.history += self.pyuvdata_version_str

        self.x_orientation = "east"

        self._set_efield()
        self.Naxes_vec = 2
        self.Ncomponents_vec = 2
        self.feed_array = np.array([str(pol.lower()) for pol in pol_names])
        self.Nfeeds = self.feed_array.size

        self.data_normalization = "physical"

        # for now this returns a simple beam because it requires amps & delays
        # to make the beam
        self.antenna_type = "simple"

        self.Nspws = 1
        self.spw_array = np.array([0])
        self.Nfreqs = freqs_use.size
        self.freq_array = freqs_use.astype(np.float)
        self.freq_array = self.freq_array[np.newaxis, :]
        self.bandpass_array = np.ones((self.Nspws, self.Nfreqs))

        self.pixel_coordinate_system = "az_za"
        self._set_cs_params()

        self.axis1_array = phi_arr
        self.Naxes1 = self.axis1_array.size
        self.axis2_array = theta_arr
        self.Naxes2 = self.axis2_array.size

        # The array that come from `_get_response` has shape shape
        # (Npol, 2, Nfreq, Nphi, Ntheta)
        # UVBeam wants shape
        # ('Naxes_vec', 'Nspws', 'Nfeeds', 'Nfreqs', 'Naxes2', 'Naxes1')
        # where the Naxes_vec dimension lines up with the 2 from `_get_response`,
        # Nfeeds is UVBeam's Npol for E-field beams,
        # and axes (2, 1) correspond to (theta, phi)
        # Then add an empty dimension for Nspws.
        self.data_array = np.transpose(jones, axes=[1, 0, 2, 4, 3])
        self.data_array = self.data_array[:, np.newaxis, :, :, :, :]

        self.basis_vector_array = np.zeros(
            (self.Naxes_vec, self.Ncomponents_vec, self.Naxes2, self.Naxes1)
        )
        self.basis_vector_array[0, 0, :, :] = 1.0
        self.basis_vector_array[1, 1, :, :] = 1.0

        if run_check:
            self.check(
                check_extra=check_extra, run_check_acceptability=run_check_acceptability
            )
back to top