fftvis 0.0.1__py3-none-any.whl

fftvis/__init__.py

@@ -0,0 +1 @@
+from . import beams, utils, simulate

fftvis/beams.py

@@ -0,0 +1,75 @@
+import numpy as np
+from pyuvdata import UVBeam
+
+
+def _evaluate_beam(
+    A_s: np.ndarray,
+    beam: UVBeam,
+    az: np.ndarray,
+    za: np.ndarray,
+    polarized: bool,
+    freq: float,
+    check: bool = False,
+    spline_opts: dict = None,
+):
+    """Evaluate the beam on the CPU. Simplified version of the `_evaluate_beam_cpu` function
+    in matvis.
+
+    This function will either interpolate the beam to the given coordinates tx, ty,
+    or evaluate the beam there if it is an analytic beam.
+
+    Parameters
+    ----------
+    A_s
+        Array of shape (nax, nfeed, nsrcs_up) that will be filled with beam
+        values.
+    beam
+        UVBeam object to evaluate.
+    tx, ty
+        Coordinates to evaluate the beam at, in sin-projection.
+    polarized
+        Whether to use beam polarization.
+    freq
+        Frequency to interpolate beam to.
+    check
+        Whether to check that the beam has no inf/nan values. Set to False if you are
+        sure that the beam is valid, as it will be faster.
+    spline_opts
+        Extra options to pass to the RectBivariateSpline class when interpolating.
+    """
+    # Primary beam pattern using direct interpolation of UVBeam object
+    kw = (
+        {
+            "reuse_spline": True,
+            "check_azza_domain": False,
+            "spline_opts": spline_opts,
+        }
+        if isinstance(beam, UVBeam)
+        else {}
+    )
+    if isinstance(beam, UVBeam) and not beam.future_array_shapes:
+        beam.use_future_array_shapes()
+
+    interp_beam = beam.interp(
+        az_array=az,
+        za_array=za,
+        freq_array=np.atleast_1d(freq),
+        **kw,
+    )[0]
+
+    if polarized:
+        interp_beam = interp_beam[:, :, 0, :]
+    else:
+        # Here we have already asserted that the beam is a power beam and
+        # has only one polarization, so we just evaluate that one.
+        interp_beam = np.sqrt(interp_beam[0, 0, 0, :])
+
+    A_s[:, :] = interp_beam
+
+    # Check for invalid beam values
+    if check:
+        sm = np.sum(A_s)
+        if np.isinf(sm) or np.isnan(sm):
+            raise ValueError("Beam interpolation resulted in an invalid value")
+
+    return A_s

fftvis/simulate.py

@@ -0,0 +1,330 @@
+from __future__ import annotations
+
+import finufft
+import numpy as np
+from matvis import conversions
+
+from . import utils, beams
+
+# Default accuracy for the non-uniform fast fourier transform based on precision
+default_accuracy_dict = {
+    1: 6e-8,
+    2: 1e-13,
+}
+
+
+def simulate_vis(
+    ants: dict,
+    fluxes: np.ndarray,
+    ra: np.ndarray,
+    dec: np.ndarray,
+    freqs: np.ndarray,
+    lsts: np.ndarray,
+    beam,
+    baselines: list[tuple] = None,
+    precision: int = 2,
+    polarized: bool = False,
+    latitude: float = -0.5361913261514378,
+    eps: float = None,
+    use_feed: str = "x",
+):
+    """
+    Parameters:
+    ----------
+    ants : dict
+        Dictionary of antenna positions
+    fluxes : np.ndarray
+        Intensity distribution of sources/pixels on the sky, assuming intensity
+        (Stokes I) only. The Stokes I intensity will be split equally between
+        the two linear polarization channels, resulting in a factor of 0.5 from
+        the value inputted here. This is done even if only one polarization
+        channel is simulated.
+    ra, dec : array_like
+        Arrays of source RA and Dec positions in radians. RA goes from [0, 2 pi]
+        and Dec from [-pi, +pi].
+    freqs : np.ndarray
+        Frequencies to evaluate visibilities in Hz.
+    lsts : np.ndarray
+        Local sidereal time in radians. Range is [0, 2 pi].
+    beam : UVBeam
+        Beam object to use for the array. Per-antenna beams are not yet supported.
+    baselines : list of tuples, default = None
+        If provided, only the baselines within the list will be simulated and array of shape
+        (nbls, nfreqs, ntimes) will be returned if polarized is False, and (nbls, nfreqs, ntimes, 2, 2) if polarized is True.
+    precision : int, optional
+       Which precision level to use for floats and complex numbers
+       Allowed values:
+       - 1: float32, complex64
+       - 2: float64, complex128
+    polarized : bool, optional
+        Whether to simulate polarized visibilities. If True, the output will have
+        shape (nfreqs, ntimes, 2, 2, nants, nants), and if False, the output will
+        have shape (nfreqs, ntimes, nants, nants).
+    latitude : float, optional
+        Latitude of the array in radians. The default is the
+        HERA latitude = -30.7215 * pi / 180.
+    eps : float, default = None
+        Desired accuracy of the non-uniform fast fourier transform. If None, the default accuracy
+        for the given precision will be used. For precision 1, the default accuracy is 6e-8, and for
+        precision 2, the default accuracy is 1e-12.
+
+    Returns:
+    -------
+    vis : np.ndarray
+        Array of shape (nfreqs, ntimes, nants, nants) if polarized is False, and
+        (nfreqs, ntimes, 2, 2, nants, nants) if polarized is True.
+    """
+    # Get the accuracy for the given precision if not provided
+    if eps is None:
+        eps = default_accuracy_dict[precision]
+
+    # Source coordinate transform, from equatorial to Cartesian
+    crd_eq = conversions.point_source_crd_eq(ra, dec)
+
+    # Make sure antpos has the right format
+    ants = {k: np.array(v) for k, v in ants.items()}
+
+    # Get coordinate transforms as a function of LST
+    eq2tops = np.array([conversions.eci_to_enu_matrix(lst, latitude) for lst in lsts])
+
+    return simulate(
+        ants=ants,
+        freqs=freqs,
+        fluxes=fluxes,
+        beam=beam,
+        crd_eq=crd_eq,
+        eq2tops=eq2tops,
+        baselines=baselines,
+        precision=precision,
+        polarized=polarized,
+        eps=eps,
+        use_feed=use_feed,
+    )
+
+
+def simulate(
+    ants: dict,
+    freqs: np.ndarray,
+    fluxes: np.ndarray,
+    beam,
+    crd_eq: np.ndarray,
+    eq2tops: np.ndarray,
+    baselines: list[tuple] = None,
+    precision: int = 2,
+    polarized: bool = False,
+    eps: float = 1e-13,
+    use_feed: str = "x",
+):
+    """
+    Parameters:
+    ----------
+    ants : dict
+        Dictionary of antenna positions in the form {ant_index: np.array([x,y,z])}.
+    freqs : np.ndarray
+        Frequencies to evaluate visibilities at in Hz.
+    fluxes : np.ndarray
+        Intensity distribution of sources/pixels on the sky, assuming intensity
+        (Stokes I) only. The Stokes I intensity will be split equally between
+        the two linear polarization channels, resulting in a factor of 0.5 from
+        the value inputted here. This is done even if only one polarization
+        channel is simulated.
+    beam : UVBeam
+        Beam object to use for the array. Per-antenna beams are not yet supported.
+    crd_eq : np.ndarray
+        Cartesian unit vectors of sources in an ECI (Earth Centered
+        Inertial) system, which has the Earth's center of mass at
+        the origin, and is fixed with respect to the distant stars.
+        The components of the ECI vector for each source are:
+        (cos(RA) cos(Dec), sin(RA) cos(Dec), sin(Dec)).
+        Shape=(3, NSRCS).
+    eq2tops : np.ndarray
+        Set of 3x3 transformation matrices to rotate the RA and Dec
+        cosines in an ECI coordinate system (see `crd_eq`) to
+        topocentric ENU (East-North-Up) unit vectors at each
+        time/LST/hour angle in the dataset. Shape=(NTIMES, 3, 3).
+    baselines : list of tuples, default = None
+        If provided, only the baselines within the list will be simulated and array of shape
+        (nbls, nfreqs, ntimes) will be returned
+    precision : int, optional
+        Which precision level to use for floats and complex numbers
+        Allowed values:
+        - 1: float32, complex64
+        - 2: float64, complex128
+    eps : float, default = 6e-8
+        Desired accuracy of the non-uniform fast fourier transform.
+
+
+    Returns:
+    -------
+    vis : np.ndarray
+        Array of shape (nfreqs, ntimes, nants, nants) if polarized is False, and
+        (nfreqs, ntimes, 2, 2, nants, nants) if polarized is True.
+    """
+    # Get sizes of inputs
+    nfreqs = np.size(freqs)
+    nants = len(ants)
+    ntimes = len(eq2tops)
+
+    if polarized:
+        nax = nfeeds = 2
+    else:
+        nax = nfeeds = 1
+
+    if precision == 1:
+        real_dtype = np.float32
+        complex_dtype = np.complex64
+    else:
+        real_dtype = np.float64
+        complex_dtype = np.complex128
+
+    # Get the redundant groups - TODO handle this better
+    if not baselines:
+        reds = utils.get_pos_reds(ants, include_autos=True)
+        baselines = [red[0] for red in reds]
+        nbls = len(baselines)
+        bl_to_red_map = {red[0]: np.array(red) for red in reds}
+        expand_vis = True
+    else:
+        nbls = len(baselines)
+        expand_vis = False
+
+    # Prepare the beam
+    beam = conversions.prepare_beam(beam, polarized=polarized, use_feed=use_feed)
+
+    # Check if the beam is complex
+    beam_values, _ = beam.interp(
+        az_array=np.array([0]),
+        za_array=np.array([0]),
+        freq_array=np.array([freqs[0]]),
+    )
+    is_beam_complex = np.issubdtype(beam_values.dtype, np.complexfloating)
+
+    # Convert to correct precision
+    crd_eq = crd_eq.astype(real_dtype)
+    eq2tops = eq2tops.astype(real_dtype)
+
+    # Factor of 0.5 accounts for splitting Stokes between polarization channels
+    Isky = (0.5 * fluxes).astype(complex_dtype)
+
+    # Compute baseline vectors
+    blx, bly = np.array([ants[bl[1]] - ants[bl[0]] for bl in baselines])[
+        :, :2
+    ].T.astype(real_dtype)
+
+    # Generate visibility array
+    if expand_vis:
+        vis = np.zeros(
+            (ntimes, nants, nants, nfeeds, nfeeds, nfreqs), dtype=complex_dtype
+        )
+    else:
+        vis = np.zeros((ntimes, nbls, nfeeds, nfeeds, nfreqs), dtype=complex_dtype)
+
+    # Loop over time samples
+    for ti, eq2top in enumerate(eq2tops):
+        # Convert to topocentric coordinates
+        tx, ty, tz = np.dot(eq2top, crd_eq)
+
+        # Only simulate above the horizon
+        above_horizon = tz > 0
+        tx = tx[above_horizon]
+        ty = ty[above_horizon]
+
+        # Number of above horizon points
+        nsim_sources = above_horizon.sum()
+
+        # Form the visibility array
+        _vis = np.zeros((nfeeds, nfeeds, nbls, nfreqs), dtype=complex_dtype)
+
+        if is_beam_complex:
+            _vis_negatives = np.zeros(
+                (nfeeds, nfeeds, nbls, nfreqs), dtype=complex_dtype
+            )
+
+        # Compute azimuth and zenith angles
+        az, za = conversions.enu_to_az_za(enu_e=tx, enu_n=ty, orientation="uvbeam")
+
+        for fi in range(nfreqs):
+            # Compute uv coordinates
+            u, v = (
+                blx * freqs[fi] / utils.speed_of_light,
+                bly * freqs[fi] / utils.speed_of_light,
+            )
+
+            # Compute beams - only single beam is supported
+            A_s = np.zeros((nax, nfeeds, nsim_sources), dtype=complex_dtype)
+            A_s = beams._evaluate_beam(A_s, beam, az, za, polarized, freqs[fi])
+            A_s = A_s.transpose((1, 0, 2))
+            beam_product = np.einsum("abs,cbs->acs", A_s.conj(), A_s)
+            beam_product = beam_product.reshape(nax * nfeeds, nsim_sources)
+
+            # Compute sky beam product
+            i_sky = beam_product * Isky[above_horizon, fi]
+
+            # Compute visibilities w/ non-uniform FFT
+            _vis_here = finufft.nufft2d3(
+                2 * np.pi * tx,
+                2 * np.pi * ty,
+                i_sky,
+                u,
+                v,
+                modeord=0,
+                eps=eps,
+            )
+
+            # Expand out the visibility array
+            _vis[..., fi] = _vis_here.reshape(nfeeds, nfeeds, nbls)
+
+            # If beam is complex, we need to compute the reverse negative frequencies
+            # TODO: no way to store this in the loop
+            if is_beam_complex:
+                # Compute
+                _vis_here_neg = finufft.nufft2d3(
+                    2 * np.pi * tx,
+                    2 * np.pi * ty,
+                    i_sky,
+                    -u,
+                    -v,
+                    modeord=0,
+                    eps=eps,
+                )
+                _vis_negatives[..., fi] = _vis_here_neg.reshape(nfeeds, nfeeds, nbls)
+
+        # Expand out the visibility array in antenna by antenna matrix
+        if expand_vis:
+            for bi, bls in enumerate(baselines):
+                np.add.at(
+                    vis,
+                    (ti, bl_to_red_map[bls][:, 0], bl_to_red_map[bls][:, 1]),
+                    _vis[..., bi, :],
+                )
+
+                # Add the conjugate, avoid auto baselines twice
+                if bls[0] != bls[1]:
+                    if is_beam_complex:
+                        np.add.at(
+                            vis,
+                            (ti, bl_to_red_map[bls][:, 1], bl_to_red_map[bls][:, 0]),
+                            _vis_negatives[..., bi, :],
+                        )
+                    else:
+                        np.add.at(
+                            vis,
+                            (ti, bl_to_red_map[bls][:, 1], bl_to_red_map[bls][:, 0]),
+                            _vis[..., bi, :].conj(),
+                        )
+
+        else:
+            vis[ti] = np.swapaxes(_vis, 2, 0)
+
+    if expand_vis:
+        return (
+            np.transpose(vis, (5, 0, 3, 4, 1, 2))
+            if polarized
+            else np.moveaxis(vis[..., 0, 0, :], 3, 0)
+        )
+    else:
+        return (
+            np.transpose(vis, (4, 0, 2, 3, 1))
+            if polarized
+            else np.moveaxis(vis[..., 0, 0, :], 2, 0)
+        )

fftvis/tests/__init__.py

@@ -0,0 +1,165 @@
+"""Tests."""
+
+import numpy as np
+from astropy import units as un
+from astropy.coordinates import EarthLocation, Latitude, Longitude
+from astropy.time import Time
+from astropy.units import Quantity
+from pathlib import Path
+from pyradiosky import SkyModel
+from pyuvdata import UVBeam
+from pyuvsim import AnalyticBeam, simsetup
+from pyuvsim.telescope import BeamList
+
+from matvis import DATA_PATH, conversions
+
+nfreq = 1
+ntime = 5  # 20
+nants = 3  # 4
+nsource = 15  # 250
+beam_file = DATA_PATH / "NF_HERA_Dipole_small.fits"
+
+
+def get_standard_sim_params(
+    use_analytic_beam: bool,
+    polarized: bool,
+    nants=nants,
+    nfreq=nfreq,
+    ntime=ntime,
+    nsource=nsource,
+    first_source_antizenith=False,
+):
+    """Create some standard random simulation parameters for use in tests."""
+    hera_lat = -30.7215
+    hera_lon = 21.4283
+    hera_alt = 1073.0
+    obstime = Time("2018-08-31T04:02:30.11", format="isot", scale="utc")
+
+    # HERA location
+    location = EarthLocation.from_geodetic(lat=hera_lat, lon=hera_lon, height=hera_alt)
+
+    np.random.seed(1)
+
+    # Beam model
+    if use_analytic_beam:
+        n_freq = nfreq
+        beam = AnalyticBeam("gaussian", diameter=14.0)
+    else:
+        n_freq = min(nfreq, 2)
+        # This is a peak-normalized e-field beam file at 100 and 101 MHz,
+        # downsampled to roughly 4 square-degree resolution.
+        beam = UVBeam()
+        beam.read_beamfits(beam_file)
+        if not polarized:
+            uvsim_beam = beam.copy()
+            beam.efield_to_power(calc_cross_pols=False, inplace=True)
+            beam.select(polarizations=["xx"], inplace=True)
+
+        # Now, the beam we have on file doesn't actually properly wrap around in azimuth.
+        # This is fine -- the uvbeam.interp() handles the interpolation well. However, when
+        # comparing to the GPU interpolation, which first has to interpolate to a regular
+        # grid that ends right at 2pi, it's better to compare like for like, so we
+        # interpolate to such a grid here.
+        beam = beam.interp(
+            az_array=np.linspace(0, 2 * np.pi, 181),
+            za_array=np.linspace(0, np.pi / 2, 46),
+            az_za_grid=True,
+            new_object=True,
+        )
+
+    if polarized or use_analytic_beam:
+        uvsim_beams = BeamList([beam] * nants)
+    else:
+        uvsim_beams = BeamList([uvsim_beam] * nants)
+
+    # beams = [AnalyticBeam('uniform') for i in range(len(ants.keys()))]
+    beam_dict = {str(i): i for i in range(nants)}
+
+    # Random antenna locations
+    x = np.random.random(nants) * 400.0  # Up to 400 metres
+    y = np.random.random(nants) * 400.0
+    z = np.random.random(nants) * 0.0
+    ants = {i: (x[i], y[i], z[i]) for i in range(nants)}
+
+    # Observing parameters in a UVData object
+    uvdata = simsetup.initialize_uvdata_from_keywords(
+        Nfreqs=n_freq,
+        start_freq=100e6,
+        channel_width=97.3e3,
+        start_time=obstime.jd,
+        integration_time=182.0,  # Just over 3 mins between time samples
+        Ntimes=ntime,
+        array_layout=ants,
+        polarization_array=np.array(["XX", "YY", "XY", "YX"]),
+        telescope_location=(hera_lat, hera_lon, hera_alt),
+        telescope_name="test_array",
+        phase_type="drift",
+        vis_units="Jy",
+        complete=True,
+        write_files=False,
+    )
+    lsts = np.unique(uvdata.lst_array)
+
+    # One fixed source plus random other sources
+    sources = [
+        [
+            300 if first_source_antizenith else 125.7,
+            -30.72,
+            2,
+            0,
+        ],  # Fix a single source near zenith
+    ]
+    if nsource > 1:  # Add random other sources
+        ra = np.random.uniform(low=0.0, high=360.0, size=nsource - 1)
+        dec = -30.72 + np.random.random(nsource - 1) * 10.0
+        flux = np.random.random(nsource - 1) * 4
+        sources.extend([ra[i], dec[i], flux[i], 0] for i in range(nsource - 1))
+    sources = np.array(sources)
+
+    # Source locations and frequencies
+    ra_dec = np.deg2rad(sources[:, :2])
+    freqs = np.unique(uvdata.freq_array)
+
+    # Correct source locations so that matvis uses the right frame
+    ra_new, dec_new = conversions.equatorial_to_eci_coords(
+        ra_dec[:, 0], ra_dec[:, 1], obstime, location, unit="rad", frame="icrs"
+    )
+
+    # Calculate source fluxes for matvis
+    flux = ((freqs[:, np.newaxis] / freqs[0]) ** sources[:, 3].T * sources[:, 2].T).T
+
+    # Stokes for the first frequency only. Stokes for other frequencies
+    # are calculated later.
+    stokes = np.zeros((4, 1, ra_dec.shape[0]))
+    stokes[0, 0] = sources[:, 2]
+    reference_frequency = np.full(len(ra_dec), freqs[0])
+
+    # Set up sky model
+    sky_model = SkyModel(
+        name=[str(i) for i in range(len(ra_dec))],
+        ra=Longitude(ra_dec[:, 0], "rad"),
+        dec=Latitude(ra_dec[:, 1], "rad"),
+        frame="icrs",
+        spectral_type="spectral_index",
+        spectral_index=sources[:, 3],
+        stokes=stokes * un.Jy,
+        reference_frequency=Quantity(reference_frequency, "Hz"),
+    )
+
+    # Calculate stokes at all the frequencies.
+    sky_model.at_frequencies(Quantity(freqs, "Hz"), inplace=True)
+
+    return (
+        sky_model,
+        ants,
+        flux,
+        ra_new,
+        dec_new,
+        freqs,
+        lsts,
+        [beam],
+        uvsim_beams,
+        beam_dict,
+        hera_lat,
+        uvdata,
+    )

fftvis/tests/test_compare_matvis.py

@@ -0,0 +1,122 @@
+import matvis
+import pytest
+import numpy as np
+from fftvis import simulate
+from pyuvsim.analyticbeam import AnalyticBeam
+
+
+def test_simulate():
+    """ """
+    # Simulation parameters
+    ntimes = 10
+    nfreqs = 5
+    nants = 3
+    nsrcs = 20
+
+    # Set random set
+    np.random.seed(42)
+
+    # Define frequency and time range
+    freqs = np.linspace(100e6, 200e6, nfreqs)
+    lsts = np.linspace(0, np.pi, ntimes)
+
+    # Set up the antenna positions
+    antpos = {k: np.array([k * 10, 0, 0]) for k in range(nants)}
+
+    # Define a Gaussian beam
+    beam = AnalyticBeam("gaussian", diameter=14.0)
+
+    # Set sky model
+    ra = np.linspace(0.0, 2.0 * np.pi, nsrcs)
+    dec = np.linspace(-0.5 * np.pi, 0.5 * np.pi, nsrcs)
+    sky_model = np.random.uniform(0, 1, size=(nsrcs, 1)) * (freqs[None] / 150e6) ** -2.5
+
+    # Use matvis as a reference
+    mvis = matvis.simulate_vis(
+        antpos, sky_model, ra, dec, freqs, lsts, beams=[beam], precision=2
+    )
+
+    # Use fftvis to simulate visibilities
+    fvis = simulate.simulate_vis(
+        antpos, sky_model, ra, dec, freqs, lsts, beam, precision=2, eps=1e-10
+    )
+
+    # Should have shape (nfreqs, ntimes, nants, nants)
+    assert fvis.shape == (nfreqs, ntimes, nants, nants)
+
+    # Check that the results are the same
+    assert np.allclose(mvis, fvis, atol=1e-5)
+
+    # Test polarized visibilities
+    # Use matvis as a reference
+    mvis = matvis.simulate_vis(
+        antpos,
+        sky_model,
+        ra,
+        dec,
+        freqs,
+        lsts,
+        beams=[beam],
+        precision=2,
+        polarized=True,
+    )
+
+    # Use fftvis to simulate visibilities
+    fvis = simulate.simulate_vis(
+        antpos,
+        sky_model,
+        ra,
+        dec,
+        freqs,
+        lsts,
+        beam,
+        precision=2,
+        eps=1e-10,
+        polarized=True,
+    )
+
+    # Should have shape (nfreqs, ntimes, nfeeds, nfeeds, nants, nants)
+    assert fvis.shape == (nfreqs, ntimes, 2, 2, nants, nants)
+
+    # Check that the polarized results are the same
+    assert np.allclose(mvis, fvis, atol=1e-5)
+
+    # Simulate with specified baselines
+    sim_baselines = [(0, 1), (0, 2), (1, 2)]
+    fvis = simulate.simulate_vis(
+        antpos,
+        sky_model,
+        ra,
+        dec,
+        freqs,
+        lsts,
+        beam,
+        baselines=sim_baselines,
+        precision=2,
+        eps=1e-10,
+        polarized=True,
+    )
+
+    # Should have shape (nfreqs, ntimes, 2, 2, len(sim_baselines))
+    assert fvis.shape == (nfreqs, ntimes, 2, 2, len(sim_baselines))
+
+    # Check that the polarized results are the same
+    for bi, bl in enumerate(sim_baselines):
+        assert np.allclose(fvis[:, :, :, :, bi], mvis[:, :, :, :, bl[0], bl[1]])
+
+    # Test with precision 1
+    # Use matvis as a reference
+    mvis = matvis.simulate_vis(
+        antpos, sky_model, ra, dec, freqs, lsts, beams=[beam], precision=1
+    )
+
+    # Use fftvis to simulate visibilities
+    fvis = simulate.simulate_vis(
+        antpos, sky_model, ra, dec, freqs, lsts, beam, precision=1, eps=6e-8
+    )
+
+    # Should have shape (nfreqs, ntimes, nants, nants)
+    assert fvis.shape == (nfreqs, ntimes, nants, nants)
+
+    # Check that the results are the same
+    assert np.allclose(mvis, fvis, atol=1e-5)

fftvis/tests/test_compare_pyuvsim.py

@@ -0,0 +1,129 @@
+"""
+
+Compare matvis with pyuvsim visibilities.
+
+Tests copied from matvis/tests/test_compare_pyuvsim.py
+
+"""
+
+import pytest
+
+import numpy as np
+from pyuvsim import simsetup, uvsim
+
+from fftvis.simulate import simulate_vis
+from matvis import simulate_vis as simulate_vis_matvis
+
+from . import get_standard_sim_params, nants
+
+
+@pytest.mark.parametrize("use_analytic_beam", (True, False))
+@pytest.mark.parametrize("polarized", (True, False))
+def test_compare_pyuvsim(polarized, use_analytic_beam):
+    """Compare matvis and pyuvsim simulated visibilities."""
+    print("Polarized=", polarized, "Analytic Beam =", use_analytic_beam)
+    (
+        sky_model,
+        ants,
+        flux,
+        ra,
+        dec,
+        freqs,
+        lsts,
+        cpu_beams,
+        uvsim_beams,
+        beam_dict,
+        hera_lat,
+        uvdata,
+    ) = get_standard_sim_params(use_analytic_beam, polarized)
+    # ---------------------------------------------------------------------------
+    # (1) Run matvis
+    # ---------------------------------------------------------------------------
+    vis_fftvis = simulate_vis(
+        ants=ants,
+        fluxes=flux,
+        ra=ra,
+        dec=dec,
+        freqs=freqs,
+        lsts=lsts,
+        beam=cpu_beams[0],
+        polarized=polarized,
+        precision=2,
+        latitude=hera_lat * np.pi / 180.0,
+    )
+
+    # ---------------------------------------------------------------------------
+    # (2) Run pyuvsim
+    # ---------------------------------------------------------------------------
+    # uvd_uvsim = uvsim.run_uvdata_uvsim(
+    #    uvdata,
+    #    uvsim_beams,
+    #    beam_dict=beam_dict,
+    #    catalog=simsetup.SkyModelData(sky_model),
+    # )
+    vis_matvis = simulate_vis_matvis(
+        ants=ants,
+        fluxes=flux,
+        ra=ra,
+        dec=dec,
+        freqs=freqs,
+        lsts=lsts,
+        beams=cpu_beams,
+        polarized=polarized,
+        precision=2,
+        latitude=hera_lat * np.pi / 180.0,
+    )
+
+    # ---------------------------------------------------------------------------
+    # Compare
+    # ---------------------------------------------------------------------------
+    # Loop over baselines and compare
+    diff_re = 0.0
+    diff_im = 0.0
+    rtol = 2e-4 if use_analytic_beam else 0.01
+    atol = 5e-4
+
+    # If it passes this test, but fails the following tests, then its probably an
+    # ordering issue.
+    for i in range(nants):
+        for j in range(i, nants):
+            for if1, feed1 in enumerate(("X", "Y") if polarized else ("X",)):
+                for if2, feed2 in enumerate(("X", "Y") if polarized else ("X",)):
+                    # d_uvsim = uvd_uvsim.get_data(
+                    #    (i, j, feed1 + feed2)
+                    # ).T  # pyuvsim visibility
+                    d_uvsim = (
+                        vis_matvis[:, :, if1, if2, i, j]
+                        if polarized
+                        else vis_fftvis[:, :, i, j]
+                    )
+                    d_fftvis = (
+                        vis_fftvis[:, :, if1, if2, i, j]
+                        if polarized
+                        else vis_fftvis[:, :, i, j]
+                    )
+
+                    # Keep track of maximum difference
+                    delta = d_uvsim - d_fftvis
+                    if np.max(np.abs(delta.real)) > diff_re:
+                        diff_re = np.max(np.abs(delta.real))
+                    if np.max(np.abs(delta.imag)) > diff_im:
+                        diff_im = np.abs(np.max(delta.imag))
+
+                    err = f"\nMax diff: {diff_re:10.10e} + 1j*{diff_im:10.10e}\n"
+                    err += f"Baseline: ({i},{j},{feed1}{feed2})\n"
+                    err += f"Avg. diff: {delta.mean():10.10e}\n"
+                    err += f"Max values: \n    uvsim={d_uvsim.max():10.10e}"
+                    err += f"\n    matvis={d_fftvis.max():10.10e}"
+                    assert np.allclose(
+                        d_uvsim.real,
+                        d_fftvis.real,
+                        rtol=rtol if feed1 == feed2 else rtol * 100,
+                        atol=atol,
+                    ), err
+                    assert np.allclose(
+                        d_uvsim.imag,
+                        d_fftvis.imag,
+                        rtol=rtol if feed1 == feed2 else rtol * 100,
+                        atol=atol,
+                    ), err

fftvis/tests/test_utils.py

@@ -0,0 +1,6 @@
+import pytest
+
+
+def test_get_pos_reds():
+    """ """
+    pass

fftvis/utils.py

@@ -0,0 +1,57 @@
+import numpy as np
+
+IDEALIZED_BL_TOL = 1e-8  # bl_error_tol for redcal.get_reds when using antenna positions calculated from reds
+speed_of_light = 299792458.0  # m/s
+
+
+def get_pos_reds(antpos, decimals=3, include_autos=True):
+    """
+    Figure out and return list of lists of redundant baseline pairs. This function is a modified version of the
+    get_pos_reds function in redcal. It is used to calculate the redundant baseline groups from antenna positions
+    rather than from a list of baselines. This is useful for simulating visibilities with fftvis.
+
+    Parameters:
+    ----------
+        antpos: dict
+            dictionary of antenna positions in the form {ant_index: np.array([x,y,z])}.
+        decimals: int, optional
+            Number of decimal places to round to when determining redundant baselines. default is 3.
+        include_autos: bool, optional
+            if True, include autos in the list of pos_reds. default is False
+    Returns:
+    -------
+        reds: list of lists of redundant tuples of antenna indices (no polarizations),
+        sorted by index with the first index of the first baseline the lowest in the group.
+    """
+    # Create a dictionary of redundant baseline groups
+    uv_to_red_key = {}
+    reds = {}
+
+    # Compute baseline lengths and round to specified precision
+    baselines = np.round(
+        [
+            antpos[aj] - antpos[ai]
+            for ai in antpos
+            for aj in antpos
+            if ai < aj or include_autos and ai == aj
+        ],
+        decimals,
+    )
+
+    ci = 0
+    for ai in antpos:
+        for aj in antpos:
+            if ai < aj or include_autos and ai == aj:
+                u, v, _ = baselines[ci]
+
+                if (u, v) not in uv_to_red_key and (-u, -v) not in uv_to_red_key:
+                    reds[(ai, aj)] = [(ai, aj)]
+                    uv_to_red_key[(u, v)] = (ai, aj)
+                elif (-u, -v) in uv_to_red_key:
+                    reds[uv_to_red_key[(-u, -v)]].append((aj, ai))
+                elif (u, v) in uv_to_red_key:
+                    reds[uv_to_red_key[(u, v)]].append((ai, aj))
+
+                ci += 1
+
+    return [reds[k] for k in reds]

fftvis-0.0.1.dist-info/METADATA

@@ -0,0 +1,76 @@
+Metadata-Version: 2.1
+Name: fftvis
+Version: 0.0.1
+Summary: A package for simulating visibilities using FFTs
+Home-page: https://github.com/tyler-a-cox/fftvis
+Author: Tyler Cox
+Author-email: tyler.a.cox@berkeley.edu
+License: MIT
+Platform: any
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.6
+Classifier: Programming Language :: Python :: 3.7
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Scientific/Engineering :: Visualization
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy
+Requires-Dist: matvis
+Requires-Dist: finufft
+Requires-Dist: pyuvdata
+Provides-Extra: dev
+Requires-Dist: mpi4py ; extra == 'dev'
+Requires-Dist: pyuvsim[sim] ; extra == 'dev'
+Requires-Dist: pyradiosky ; extra == 'dev'
+Requires-Dist: pytest ; extra == 'dev'
+Requires-Dist: pre-commit ; extra == 'dev'
+Requires-Dist: pytest-cov ; extra == 'dev'
+Requires-Dist: hera-sim ; extra == 'dev'
+Requires-Dist: pytest-xdist ; extra == 'dev'
+
+# fftvis: A Non-Uniform Fast Fourier Transform-based Visibility Simulator
+
+![Tests](https://github.com/tyler-a-cox/fftvis/actions/workflows/ci.yml/badge.svg)
+![codecov](https://codecov.io/gh/tyler-a-cox/fftvis/branch/main/graph/badge.svg)
+![Black Formatting](https://img.shields.io/badge/code%20style-black-000000.svg)
+
+`fftvis` is a fast Python package designed for simulating interferometric visibilities using the Non-Uniform Fast Fourier Transform (NUFFT). It provides a convenient and efficient way to generate simulated visibilities.
+
+## Features
+
+- Utilizes the Flatiron Institute NUFFT (finufft) [algorithm](https://arxiv.org/abs/1808.06736) for fast visibility simulations that agree with similar methods ([`matvis`](https://github.com/HERA-team/matvis)) to nearly machine precision.
+- Designed to be a near drop-in replacement to `matvis` with a ~10x improvement in runtime
+
+## Limitations
+- Currently no support for per-antenna beams
+- Currently no support for polarized sky emission 
+- Currently no GPU support
+- Diffuse sky models must be pixelized
+
+## Installation
+
+You can install `fftvis` via pip:
+
+```bash
+git clone https://github.com/tyler-a-cox/fftvis
+cd fftvis
+pip install .
+```
+
+## Contributing
+Contributions to `fftvis` are welcome! If you find any issues, have feature requests, or want to contribute improvements, please open an issue or submit a pull request on the GitHub repository: `fftvis` on GitHub
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.
+
+## Acknowledgments
+This package relies on the `finufft` implementation provided by [finufft](https://github.com/flatironinstitute/finufft) library. Special thanks to the contributors and maintainers of open-source libraries used in this project.

fftvis-0.0.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+fftvis/__init__.py,sha256=8tSRabRYy3YZ8Oiu0yP-KhV2Pti6CNumjqrs4Qll6tw,37
+fftvis/beams.py,sha256=JP6jFyh2ulXCZvFLoiLRQ8GfYMLx0XWNkijqDTyQJU4,2143
+fftvis/simulate.py,sha256=YtyhNCAwq2ysggf2rl2Haur7qRyaU61SqjkVugYM4Oo,11422
+fftvis/utils.py,sha256=4nDW2pm5NCJofsxMB1DlSjyTAhYCUSxfuOmNwD2Bh3g,2159
+fftvis/tests/__init__.py,sha256=Dy1vOmUe-Q4lKQsNGE2vB0YR9n5XJn030IG6V55LCo8,5353
+fftvis/tests/test_compare_matvis.py,sha256=fRc2pvCS_jHV4Mb-5zCia7MMUZ3outlTBkLglDe3joU,3236
+fftvis/tests/test_compare_pyuvsim.py,sha256=N49B_903gjjtoeddGmu5niWNWxNM_heCQQxBITdtW2s,4367
+fftvis/tests/test_utils.py,sha256=gBmBDW-FFqEBOsh_315xALdpLk26RlkEQnmc8pw4hE8,62
+fftvis-0.0.1.dist-info/METADATA,sha256=-dbnCfVKg3oZ2idpm3MYliEYxw_kAkGQ2A2r5nDuOqQ,3177
+fftvis-0.0.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+fftvis-0.0.1.dist-info/top_level.txt,sha256=p91SGcGpU_MXL32mier1gBGXt2_nxk5H8hn1MwNki2I,7
+fftvis-0.0.1.dist-info/RECORD,,

fftvis-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.43.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

fftvis-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+fftvis