Repository: tyler-a-cox/fftvis
Pull request #45: Add Type 1 transform for gridded arrays
Author: login: tyler-a-cox
Created: 2025-05-08T03:31:56Z | Merged: 2025-05-12T21:46:50Z

For cases when arrays are gridded, Type 1 transforms can make `fftvis` substantially faster, especially when the maximum extent of the array is large. This PR implements methods to detect whether an input array is gridded and automatically triggers the use of Type 1 if the array is gridded.


Diff:
diff --git a/docs/tutorials/fftvis_gridded_array.ipynb b/docs/tutorials/fftvis_gridded_array.ipynb
new file mode 100644
index 0000000..b672aaa
--- /dev/null
+++ b/docs/tutorials/fftvis_gridded_array.ipynb
@@ -0,0 +1,348 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "611edba3-a666-412a-9aa2-5dd2e03344d8",
+   "metadata": {},
+   "source": [
+    "# Running `fftvis` on Gridded Arrays"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d6493a16-2ac4-47fd-8f6f-4ae294af31ed",
+   "metadata": {},
+   "source": [
+    "<div class=\"alert alert-warning\">\n",
+    "\n",
+    "__Warning__\n",
+    "\n",
+    "Before running this tutorial, you should make sure you understand the basic concepts and algorithm that `fftvis` uses. You can read up on that here.\n",
+    "</div>"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "3a26ed82-971a-47c3-807f-9e6cd74a6be5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "basis_vector_type was not defined, defaulting to azimuth and zenith_angle.\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Standard imports\n",
+    "import numpy as np\n",
+    "import healpy as hp\n",
+    "from astropy.time import Time\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# HERA-stack imports\n",
+    "import fftvis\n",
+    "import matvis\n",
+    "from hera_sim.antpos import hex_array\n",
+    "from pyuvdata.telescopes import Telescope\n",
+    "from pyuvdata.analytic_beam import AiryBeam"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6db45573-29bd-4ca3-bd5f-1020f959e9a7",
+   "metadata": {},
+   "source": [
+    "## Mathematical Prelimaries"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "db71ba69-473f-4c74-9009-bb364374513e",
+   "metadata": {},
+   "source": [
+    "Simulating radio interferometer visibilities involves calculating the contribution of many sky sources to the signal correlation measured by each baseline pair. For point sources, this is mathematically expressed as a sum:\n",
+    "\n",
+    "$$\n",
+    "V(\\mathbf{b}) \\approx \\sum_{j=1}^{N_{\\rm sources}} c_j \\, e^{-2\\pi i \\nu \\mathbf{b} \\cdot \\mathbf{x}_j/c}\n",
+    "$$\n",
+    "\n",
+    "Evaluating this sum directly is computationally intensive, scaling as $\\mathcal{O}(N_{\\rm sources} N_{\\rm bls})$, which becomes prohibitive for large source catalogs ($N_{\\rm sources}$) or many baselines ($N_{\\rm bls}$). `fftvis` accelerates this computation significantly by using the non-uniform Fast Fourier Transform (NUFFT), specifically using the high-performance `finufft` library. The exact NUFFT approach and its computational scaling depend on the geometric layout of the array baselines.\n",
+    "\n",
+    "For the general case where array baseline vectors $\\mathbf{b}_k$ do not form a regular grid (or when considering effects that break perfect regularity, such as errors in antenna placement), `fftvis` treats both the source positions $\\mathbf{x}_j$ and the target baseline coordinates $\\mathbf{u}_k = (\\nu/c)\\mathbf{b}_k$ as non-uniform. This \"non-uniform to non-uniform\" calculation is handled using a **Type 3 NUFFT**. This algorithm internally uses an intermediate uniform grid to perform the bulk of the computation via a standard FFT. The computational cost scales approximately as $\\mathcal{O}(N_{\\rm sources} w^d + N_{\\rm grid} \\log N_{\\rm grid} + N_{\\rm bls} w^d)$. Here, $w$ relates to the desired precision, $d$ is the number of dimensions (2 or 3), and $N_{\\rm grid,int}$ is the size of the intermediate grid. Most notably, $N_{\\rm grid}$ depends on the product of the maximum sky source extent and the maximum baseline extent (in wavelengths), and can become quite large, especially for wide fields or long baselines, impacting the computational performance.\n",
+    "\n",
+    "However, a more efficient approach is possible if the array's baseline vectors $\\mathbf{b}_k$ inherently form a regular Cartesian grid (meaning $b_x = k \\Delta b_x$, $b_y = l \\Delta b_y$, etc., for fixed spacings $\\Delta b_x, \\Delta b_y$). Since the source positions $\\mathbf{x}_j$ remain non-uniform, this becomes a \"non-uniform to uniform\" problem, ideally suited for a **Type 1 NUFFT**. The Type 1 NUFFT computes visibilities directly onto a uniform output grid covering the extent of the baseline coordinates. Let $N_{\\rm unif}$ be the number of points in this target uniform $uv$-grid. The computational scaling for the Type 1 NUFFT is approximately $\\mathcal{O}(N_{\\rm sources} w^d + N_{\\rm unif} \\log N_{\\rm unif})$.\n",
+    "\n",
+    "Comparing the two, the Type 1 NUFFT can be significantly faster than Type 3 if its $N_{\\rm unif} \\log N_{\\rm unif}$ term is considerably smaller than the corresponding $N_{\\rm grid} \\log N_{\\rm grid} + N_{\\rm bls} w^d$ terms from Type 3. This advantage occurs when the target uniform grid defined by the baselines ($N_{\\rm unif}$) is much smaller than the intermediate grid ($N_{\\rm grid}$) that Type 3 requires to handle the general non-uniform case across the same sky and baseline extents. Therefore, for arrays whose baselines naturally form a compact, regular grid, utilizing a Type 1 NUFFT offers a potential performance boost by using this structure. \n",
+    "\n",
+    "In the following cells, we simulate an gridded array using `fftvis` and compare the performance when simulating with a Type 1 and Type 3 NUFFT."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b35eeb4e-77f8-4847-bc64-d768564be996",
+   "metadata": {},
+   "source": [
+    "## Setup Telescope / Observation Parameters"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ba4b1972-7724-4fc8-bfbb-326d518cc4e4",
+   "metadata": {},
+   "source": [
+    "First, create our antenna positions. We define this as a dictionary, which maps an antenna number to its 3D East-North-Up position relative to the array centre. To test the difference between the Type 1 and Type 3 NUFFTs as used by `fftvis`, we'll use a hexagonally gridded array layout similar to the layout used by the HERA array."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "c34dee48-59b4-4a14-a9c6-10b2f63b556e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# define antenna array positions\n",
+    "antpos = hex_array(11, split_core=True, outriggers=2)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c740bc7b-517e-4f62-b288-6c9a1f094c9d",
+   "metadata": {},
+   "source": [
+    "Next, we define the beam to be used by all antennas in the array. Unlike `matvis` and `pyuvsim`, `fftvis` currently restricts users to a single beam for all antennas. The specified beam must be a `UVBeam` or `AnalyticBeam` object from `pyuvdata`. Alternatively, you can create a custom `AnalyticBeam` class (see the pyuvdata tutorial on `UVBeam` objects for guidance). For this simulation, we will use a simple, frequency-dependent Airy beam corresponding to a dish size of 14 meters."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "f68dd61e-4f2a-48a1-af2f-cc79c6ea38e2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# define antenna beam using pyuvdata.analytic_beam.AiryBeam with a dish size of 14 meters\n",
+    "beam = AiryBeam(diameter=14.0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e9bcf2cc-0ab5-4a22-8b87-6dab6ede2433",
+   "metadata": {},
+   "source": [
+    "We also required to provide `fftvis` with the observational configuration including a frequency array, a time array, and a telescope location. The frequency array specifies the observation frequencies in units of Hz. The time array defines the observation times using an `astropy.time.Time` object, with times specified in Julian Dates and configured with the appropriate format and scale. The telescope location specifies the geographic position of the array and can be defined either using `astropy.coordinates.EarthLocation` with a known site name or through `pyuvdata.telescopes.Telescope` by selecting a predefined telescope location supported within `pyuvdata`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "8fb30e2c-d98d-464f-83bb-6c0eaf78313c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# define a list of frequencies in units of Hz\n",
+    "nfreqs = 2\n",
+    "freqs = np.linspace(100e6, 120e6, nfreqs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "fee91e09-ec8b-4566-b027-bd7ea42ac016",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# define a list of times with an astropy time.Time object\n",
+    "ntimes = 3\n",
+    "times = Time(np.linspace(2459845, 2459845.05, ntimes), format='jd', scale='utc')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "cd7be8ff-3cf0-4e89-a865-14965b5d00cb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from astropy.coordinates import EarthLocation\n",
+    "\n",
+    "# define using astropy.coordinates.EarthLocation\n",
+    "telescope_loc = EarthLocation.of_site('meerkat')\n",
+    "\n",
+    "# define the telescope location using the pyuvdata.telescopes.Telescope\n",
+    "telescope_loc = Telescope.from_known_telescopes('hera').location"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3f585f86-b388-4e72-aa9d-493c68d181dd",
+   "metadata": {},
+   "source": [
+    "## Setup Sky Model\n",
+    "\n",
+    "Like `matvis`, `fftvis` makes the point source approximation -- that is it makes breaks a continuous sky model into a discrete number of point sources that it sums over when computing the visibilities. In this notebook, we'll assume the point source approximation by discretizing the sky with a randomly generated HEALpix map."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "c285bc67-1d9d-4c8f-a414-c5426fd6f883",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# number of sources\n",
+    "nside = 64\n",
+    "nsource = hp.nside2npix(nside)\n",
+    "\n",
+    "# pixels can be defined as point sources randomly distributed over the full sky\n",
+    "ra = np.deg2rad(np.random.uniform(0, 360, nsource))        # ra of each source (in rad)\n",
+    "dec = np.deg2rad(np.random.uniform(-90, 90.0, nsource))    # dec of each source (in rad)\n",
+    "\n",
+    "# define sky model using healpix map\n",
+    "dec, ra = hp.pix2ang(nside, np.arange(nsource))\n",
+    "dec -= np.pi / 2\n",
+    "\n",
+    "# define the flux of the sources as a function of frequency. Here, we define smooth spectrum sources\n",
+    "flux = np.random.uniform(0, 1, nsource)                              # flux of each source at 100MHz (in Jy)\n",
+    "alpha = np.ones(nsource) * -0.8                      # sp. index of each source\n",
+    "\n",
+    "# Now get the (Nsource, Nfreq) array of the flux of each source at each frequency.\n",
+    "flux_allfreq = ((freqs[:, np.newaxis] / freqs[0]) ** alpha.T * flux.T).T"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2920b621-7ece-4ca5-b0e4-85d428a9b3d5",
+   "metadata": {},
+   "source": [
+    "## Run `fftvis`\n",
+    "\n",
+    "Once the simulation parameters are set, the computation is run using the `fftvis.simulate` function. By default, `fftvis` first checks if the provided antenna positions form a regular grid. It does this by using the shortest two non-collinear baseline vectors to define a potential grid basis and then verifying if all antenna positions fit onto that grid. This approach allows `fftvis` to recognize various grid layouts, even those that are skewed or rotated relative to the coordinate axes (like hexagonal arrays, which are essentially skewed regular grids). If a grid is detected, `fftvis` automatically uses the computationally faster Type 1 NUFFT. Since this automatic detection might be unexpected or hide potential configuration issues, the `force_use_type3` parameter is available. Setting this allows a user to enforce the use of the more general Type 3 NUFFT, even if the antenna positions technically form a grid. Additionally, `fftvis` logs a message to inform the user whenever it detects a gridded array configuration and switches to the Type 1 transform.\n",
+    "\n",
+    "Below, we perform a simulation of a HERA-like array using both the Type 1 and Type 3 transforms and compare their output."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "41ed5c6a-e6ad-42f7-8975-905d219a3ce2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define subset of baselines we're interested in for simulating\n",
+    "baselines = [(i, j) for i in range(len(antpos)) for j in range(len(antpos))]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "b8c75eaa-1b3f-47cf-8f97-3d02cce21e4c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 635 ms, sys: 150 ms, total: 785 ms\n",
+      "Wall time: 482 ms\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "# simulate visibilities with the new API\n",
+    "vis_vc_gridded = fftvis.simulate_vis(\n",
+    "    ants=antpos,\n",
+    "    fluxes=flux_allfreq,\n",
+    "    ra=ra,\n",
+    "    dec=dec,\n",
+    "    freqs=freqs,\n",
+    "    times=times.jd,\n",
+    "    telescope_loc=telescope_loc,\n",
+    "    beam=beam,\n",
+    "    polarized=False,\n",
+    "    nprocesses=1,\n",
+    "    baselines=baselines,\n",
+    "    force_use_type3=False, # This is the default value\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "0e713a70-f88d-4342-a677-2b0cae932433",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 10.8 s, sys: 1.9 s, total: 12.7 s\n",
+      "Wall time: 6.69 s\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "# simulate visibilities \n",
+    "vis_vc_type3 = fftvis.simulate_vis(\n",
+    "    ants=antpos,\n",
+    "    fluxes=flux_allfreq,\n",
+    "    ra=ra,\n",
+    "    dec=dec,\n",
+    "    freqs=freqs,\n",
+    "    times=times.jd,\n",
+    "    telescope_loc=telescope_loc,\n",
+    "    beam=beam,\n",
+    "    polarized=False,\n",
+    "    nprocesses=1,\n",
+    "    baselines=baselines,\n",
+    "    force_use_type3=True, # Even though the antenna layout is gridded we force use type 3\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "8507e9ef-7bd1-4913-b14a-57d5729e11f2",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "np.allclose(vis_vc_gridded, vis_vc_type3)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/src/fftvis/__init__.py b/src/fftvis/__init__.py
index b02cf5b..20dc22b 100644
--- a/src/fftvis/__init__.py
+++ b/src/fftvis/__init__.py
@@ -2,7 +2,7 @@
 
 # Import key components for beams
 from .core.beams import BeamEvaluator
-from .cpu.cpu_beams import CPUBeamEvaluator
+from .cpu.beams import CPUBeamEvaluator
 from .wrapper import create_beam_evaluator
 
 # Import simulation functionality
diff --git a/src/fftvis/core/antenna_gridding.py b/src/fftvis/core/antenna_gridding.py
new file mode 100644
index 0000000..8ad1c2c
--- /dev/null
+++ b/src/fftvis/core/antenna_gridding.py
@@ -0,0 +1,210 @@
+import numpy as np
+from math import lcm
+from fractions import Fraction
+from typing import Any, Dict, Tuple
+
+
+def find_integer_multiplier(
+    arr: np.ndarray,
+    max_denominator: int = 10**6
+) -> int:
+    """
+    Return the smallest positive integer f such that f * arr[i] is integral
+    for all entries (up to rational approximation). Zeros in arr are ignored.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    max_denominator : int
+        Maximum denominator for the fractions.
+
+    Returns
+    -------
+    f : int
+        The integer factor used for scaling.
+    """
+    fracs = [
+        Fraction(val).limit_denominator(max_denominator)
+        for val in np.ravel(arr)
+        if val != 0
+    ]
+    if not fracs:
+        return 1
+    dens = (frac.denominator for frac in fracs)
+    return lcm(*dens)
+
+
+def can_scale_to_int(
+    arr: np.ndarray,
+    tol: float = 1e-9,
+    max_denominator: int = 10**6,
+    max_factor: int = None
+) -> Tuple[bool, int]:
+    """
+    Check if there exists an integer factor f such that f * arr is
+    approximately integers. Returns (is_griddable, factor).
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    tol : float
+        Tolerance for checking if values are close to integers in units of meters.
+    max_denominator : int
+        Maximum denominator for the fractions.
+    max_factor : int
+        Maximum allowed factor for scaling.
+
+    Returns
+    -------
+    is_griddable : bool
+        True if the array can be scaled to integers, False otherwise.
+    factor : int
+        The integer factor used for scaling.
+    """
+    f = find_integer_multiplier(arr, max_denominator)
+    if max_factor is not None and f > max_factor:
+        return False, f
+    scaled = f * np.array(arr, dtype=float)
+    if np.allclose(scaled, np.round(scaled), atol=tol):
+        return True, f
+    return False, f
+
+def find_lattice_basis(
+    antpos: Dict[Any, np.ndarray],
+    tol: float = 1e-9,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Find the lattice basis for the antenna positions. This function assumes that 
+    the antenna positions are in a 2D plane (x, y) and that the z-coordinate is
+    not relevant for the basis calculation. The function computes the pairwise
+    differences between the antenna positions and selects two non-collinear
+    vectors to form the basis. If all vectors are collinear, it returns a default
+    basis.
+
+    Parameters
+    ----------
+    antpos : dict
+        Dictionary of antenna positions in the form {ant_index: np.array([x,y,z])}
+        where the antenna positinons are in units of meters.
+    tol : float
+        Tolerance for checking if values are close to integers in units of meters.
+    max_denominator : int
+        Maximum denominator for the fractions.
+
+    Returns
+    -------
+    basis : np.ndarray
+        The lattice basis.
+    transform_matrix : np.ndarray
+        The transformation matrix.
+    """
+    antvecs = np.array([antpos[ant][:2] for ant in antpos], dtype=float)
+
+    # pairwise differences
+    blvec = np.reshape(antvecs[:, None, :] - antvecs[None, :, :], (-1, 2))
+
+    # filter out zeros
+    norms = np.linalg.norm(blvec, axis=1)
+    mask = norms > tol
+    blvec = blvec[mask]
+    norms = norms[mask]
+
+    # sort by ascending length
+    order = np.argsort(norms)
+    blvec = blvec[order]
+
+    # Pick the shortest non-zero baseline as a basis vector
+    basis_vec_1 = blvec[0]
+
+    # find second that’s not collinear (cross‑product ≠ 0)
+    for v in blvec[1:]:
+        # Compute the cross product of the x and y components
+        cross_x = basis_vec_1[0]*v[1] - basis_vec_1[1]*v[0]
+        if np.abs(cross_x) > tol:
+            basis_vec_2 = v
+            break
+    else:
+        # If all are collinear, use the shortest baseline
+        return np.vstack([basis_vec_1, np.array([0, 1])])
+
+    return np.column_stack([basis_vec_1, basis_vec_2])
+
+def check_antpos_griddability(
+    antpos: Dict[Any, np.ndarray],
+    tol: float = 1e-9,
+    max_denominator: int = 10**6,
+    max_factor: int = 1000,
+) -> Tuple[bool, Dict[Any, np.ndarray], np.ndarray]:
+    """
+    Check if the antenna positions can be gridded. This function checks if the
+    antenna positions can be scaled to integers using a linear transformation.
+    It first infers the 2D lattice basis from the antenna positions, then checks
+    if the positions can be scaled to integers. If they can, it returns the
+    modified antenna positions and the transformation matrix. If not, it returns
+    the original antenna positions and an identity transformation matrix.
+
+    Parameters
+    ----------
+    antpos : dict
+        Dictionary of antenna positions in the form {ant_index: np.array([x,y,z])}
+        where the antenna positinons are in units of meters.
+    tol : float
+        Tolerance for checking if values are close to integers in units of meters.
+    max_denominator : int
+        Maximum denominator for the fractions.
+    max_factor : int
+        Maximum allowed factor for scaling. If the factor exceeds this value,
+        the function will return False.
+
+    Returns:
+    -------
+    is_griddable : bool
+        True if the antenna positions can be gridded, False otherwise.
+    modified_antpos : dict
+        The modified antenna positions after scaling.
+    transform_matrix : np.ndarray
+        The transformation matrix used to scale the antenna positions.
+        This is a 3x3 matrix that transforms the original coordinates
+        to the grid coordinates.
+    """
+    # Map antenna keys to index and stack positions
+    antkey_to_idx = dict(zip(antpos.keys(), range(len(antpos))))
+    antvecs = np.array([antpos[ant] for ant in antpos], dtype=float)
+
+    # Infer the 2D lattice basis from the antenna positions
+    basis_2D = find_lattice_basis(
+        antpos,
+        tol=tol,
+    )
+    basis = np.zeros((3, 3))
+    basis[:2, :2] = basis_2D
+    basis[2, 2] = 1.0
+
+    # Rotate the basis to align with the inferred grid
+    modified_antvecs = np.linalg.solve(
+        basis,
+        (antvecs - antvecs[0]).T
+    ).T
+
+    # Scale the inferred lattice basis to integers
+    is_griddable, factor = can_scale_to_int(
+        np.ravel(modified_antvecs),
+        tol=tol,
+        max_denominator=max_denominator,
+        max_factor=max_factor,
+    )
+
+    if is_griddable:
+        # Scale the antenna positions, rounding to integers
+        antpos = {
+            ant: np.round(
+                factor * modified_antvecs[antkey_to_idx[ant]]
+            ).astype(int)
+            for ant in antpos
+        }
+        return True, antpos, basis / factor
+    else:
+        # If not griddable, check 
+        return False, antpos, np.eye(antvecs.shape[-1])
\ No newline at end of file
diff --git a/src/fftvis/core/simulate.py b/src/fftvis/core/simulate.py
index 814120d..69c5450 100644
--- a/src/fftvis/core/simulate.py
+++ b/src/fftvis/core/simulate.py
@@ -37,6 +37,7 @@ def simulate(
         precision: int = 2,
         polarized: bool = False,
         eps: float = None,
+        upsample_factor: Literal[1.25, 2] = 2,
         beam_spline_opts: dict = None,
         flat_array_tol: float = 0.0,
         interpolation_function: str = "az_za_map_coordinates",
@@ -91,6 +92,11 @@ def simulate(
             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.
+        upsample_factor : default = 2
+            Upsampling factor for the non-uniform fast fourier transform. This is the factor by which the
+            intermediate grid is upsampled. Only values of 1.25 or 2 are allowed. Can be useful for decreasing
+            the computation time and memory requirement for large arrays at the expensive of some accuracy. 
+            The default value is 2.
         beam_spline_opts : dict, optional
             Options to pass to :meth:`pyuvdata.uvbeam.UVBeam.interp` as `spline_opts`.
         flat_array_tol : float, default = 0.0
@@ -152,10 +158,13 @@ def _evaluate_vis_chunk(
         nfeeds: int,
         polarized: bool = False,
         eps: float = None,
+        upsample_factor: Literal[1.25, 2] = 2,
         beam_spline_opts: dict = None,
         interpolation_function: str = "az_za_map_coordinates",
         n_threads: int = 1,
         is_coplanar: bool = False,
+        basis_matrix: np.ndarray = None,
+        type1_n_modes: int = None,
         trace_mem: bool = False,
     ) -> np.ndarray:
         """
@@ -185,6 +194,8 @@ def _evaluate_vis_chunk(
             Whether to simulate polarized visibilities.
         eps : float, default = None
             Desired accuracy of the non-uniform fast fourier transform.
+        upsample_factor : int
+            Upsampling factor for the non-uniform FFT.
         beam_spline_opts : dict, default = None
             Options for beam interpolation.
         interpolation_function : str
@@ -193,6 +204,12 @@ def _evaluate_vis_chunk(
             Number of threads to use.
         is_coplanar : bool
             Whether the array is coplanar.
+        basis_matrix : np.ndarray, default = None
+            Lattice basis matrix used to grid baselines and sources for a type 1
+            non-uniform FFT. If None, a set of basis functions will be inferred from
+            the antenna positions.
+        type1_n_modes : int
+            Number of modes for type 1 non-uniform FFT.
         trace_mem : bool
             Whether to trace memory usage.
 
diff --git a/src/fftvis/core/utils.py b/src/fftvis/core/utils.py
index 37eb2dd..26519a7 100644
--- a/src/fftvis/core/utils.py
+++ b/src/fftvis/core/utils.py
@@ -4,7 +4,6 @@
 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
diff --git a/src/fftvis/cpu/__init__.py b/src/fftvis/cpu/__init__.py
index 692be72..f1e30e3 100644
--- a/src/fftvis/cpu/__init__.py
+++ b/src/fftvis/cpu/__init__.py
@@ -1,5 +1,5 @@
 """CPU-specific implementations for fftvis."""
 
-from .cpu_beams import CPUBeamEvaluator
+from .beams import CPUBeamEvaluator
 from .cpu_simulate import CPUSimulationEngine
-from .cpu_nufft import cpu_nufft2d, cpu_nufft3d
+from .nufft import cpu_nufft2d, cpu_nufft3d, cpu_nufft2d_type1
diff --git a/src/fftvis/cpu/cpu_beams.py b/src/fftvis/cpu/beams.py
similarity index 100%
rename from src/fftvis/cpu/cpu_beams.py
rename to src/fftvis/cpu/beams.py
diff --git a/src/fftvis/cpu/cpu_simulate.py b/src/fftvis/cpu/cpu_simulate.py
index 93bc9ca..84f11b5 100644
--- a/src/fftvis/cpu/cpu_simulate.py
+++ b/src/fftvis/cpu/cpu_simulate.py
@@ -22,12 +22,13 @@
 from matvis.core.coords import CoordinateRotation
 
 from ..core.simulate import SimulationEngine, default_accuracy_dict
+from ..core.antenna_gridding import check_antpos_griddability
 from .. import utils
 
 # Import the CPU beam evaluator
-from .cpu_beams import CPUBeamEvaluator
-from .cpu_nufft import cpu_nufft2d, cpu_nufft3d
-
+from .beams import CPUBeamEvaluator
+from .nufft import cpu_nufft2d, cpu_nufft3d, cpu_nufft2d_type1
+from . import utils as cpu_utils
 logger = logging.getLogger(__name__)
 
 # Create a global instance of CPUBeamEvaluator to use for beam evaluation
@@ -48,10 +49,14 @@ def _evaluate_vis_chunk_remote(
     nfeeds: int,
     polarized: bool = False,
     eps: float = None,
+    upsample_factor: Literal[1.25, 2] = 2,
     beam_spline_opts: dict = None,
     interpolation_function: str = "az_za_map_coordinates",
     n_threads: int = 1,
     is_coplanar: bool = False,
+    use_type1: bool = False,
+    basis_matrix: np.ndarray = None,
+    type1_n_modes: int = None,
     trace_mem: bool = False,
 ):
     """Ray-compatible remote version of _evaluate_vis_chunk."""
@@ -70,10 +75,14 @@ def _evaluate_vis_chunk_remote(
         nfeeds=nfeeds,
         polarized=polarized,
         eps=eps,
+        upsample_factor=upsample_factor,
         beam_spline_opts=beam_spline_opts,
         interpolation_function=interpolation_function,
         n_threads=n_threads,
         is_coplanar=is_coplanar,
+        use_type1=use_type1,
+        basis_matrix=basis_matrix,
+        type1_n_modes=type1_n_modes,
         trace_mem=trace_mem,
     )
 
@@ -95,6 +104,7 @@ def simulate(
         precision: int = 2,
         polarized: bool = False,
         eps: float = None,
+        upsample_factor: Literal[1.25, 2] = 2,
         beam_spline_opts: dict = None,
         flat_array_tol: float = 0.0,
         interpolation_function: str = "az_za_map_coordinates",
@@ -105,6 +115,7 @@ def simulate(
         ] = "CoordinateRotationERFA",
         coord_method_params: dict | None = None,
         force_use_ray: bool = False,
+        force_use_type3: bool = False,
         trace_mem: bool = False,
         enable_memory_monitor: bool = False,
     ) -> np.ndarray:
@@ -144,11 +155,6 @@ def simulate(
         # Get number of baselines
         nbls = len(baselines)
 
-        if isinstance(beam, UVBeam):
-            # Only try to interpolate the beam if it has more than one frequency
-            if hasattr(beam, "Nfreqs") and beam.Nfreqs > 1:
-                beam = beam.interp(freq_array=freqs, new_object=True, run_check=False) # pragma: no cover
-
         # Factor of 0.5 accounts for splitting Stokes between polarization channels
         Isky = 0.5 * fluxes
 
@@ -156,21 +162,53 @@ def simulate(
         antkey_to_idx = dict(zip(ants.keys(), range(len(ants))))
         antvecs = np.array([ants[ant] for ant in ants], dtype=real_dtype)
 
-        # Rotate the array to the xy-plane
-        rotation_matrix = utils.get_plane_to_xy_rotation_matrix(antvecs)
-        rotation_matrix = np.ascontiguousarray(rotation_matrix.astype(real_dtype).T)
-        rotated_antvecs = np.dot(rotation_matrix, antvecs.T)
-        rotated_ants = {ant: rotated_antvecs[:, antkey_to_idx[ant]] for ant in ants}
-
-        # Compute baseline vectors
-        bls = np.array([rotated_ants[bl[1]] - rotated_ants[bl[0]] for bl in baselines])[
-            :, :
-        ].T.astype(real_dtype)
-
-        # Check if the array is flat within tolerance
-        is_coplanar = np.all(np.less_equal(np.abs(bls[2]), flat_array_tol))
-
-        bls /= utils.speed_of_light
+        # If the array is flat within tolerance, we can check for griddability
+        if np.abs(antvecs[:, -1]).max() > flat_array_tol or force_use_type3:
+            is_gridded = False
+        else:
+            is_gridded, gridded_antpos, basis_matrix = check_antpos_griddability(ants)
+                
+        # Rotate antenna positions to XY plane if not gridded
+        if not is_gridded:
+            # Get the rotation matrix to rotate the array to the XY plane
+            rotation_matrix = utils.get_plane_to_xy_rotation_matrix(antvecs)
+            rotation_matrix = np.ascontiguousarray(rotation_matrix.T)
+            rotated_antvecs = np.dot(rotation_matrix, antvecs.T)
+            rotated_ants = {
+                ant: rotated_antvecs[:, antkey_to_idx[ant]] for ant in ants
+            }
+            rotation_matrix = rotation_matrix.astype(real_dtype)
+        
+            # Compute baseline vectors and convert to speed of light units
+            bls = np.array([rotated_ants[bl[1]] - rotated_ants[bl[0]] for bl in baselines])[
+                :, :
+            ].T
+            bls /= utils.speed_of_light
+            bls = bls.astype(real_dtype)
+
+            # Check if the array is flat within tolerance
+            is_coplanar = np.all(np.less_equal(np.abs(bls[2]), flat_array_tol))
+        else:
+            logger.info(
+                "Using gridded coordinates for the array. Type 1 transform will be used."
+            )
+            # Compute the baseline vectors in the gridded coordinate system
+            bls = np.array([
+                gridded_antpos[bl[1]] - gridded_antpos[bl[0]] 
+                for bl in baselines]
+            ).T
+            bls = np.round(bls).astype(int)
+            
+            # Find the maximum extent of the array in gridded coordinates
+            n_modes = 2 * int(np.round(np.max(np.abs(bls)))) + 1
+
+            # Get the maximum baseline length for proper coordinate scaling
+            basis_matrix *= 1 / utils.speed_of_light
+            basis_matrix = basis_matrix.astype(real_dtype)
+
+            # Assume the array is coplanar for gridded coordinates
+            is_coplanar = True
+            rotation_matrix = np.eye(3, dtype=real_dtype)
 
         # Get number of processes for multiprocessing
         if nprocesses is None:
@@ -298,10 +336,14 @@ def simulate(
                     nfeeds=nfeeds,
                     polarized=polarized,
                     eps=eps,
+                    upsample_factor=upsample_factor,
                     beam_spline_opts=beam_spline_opts,
                     interpolation_function=interpolation_function,
                     n_threads=nthi,
                     is_coplanar=is_coplanar,
+                    use_type1=is_gridded,
+                    basis_matrix=basis_matrix if is_gridded else None,
+                    type1_n_modes=n_modes if is_gridded else None,
                     trace_mem=(nprocesses > 1 or force_use_ray) and trace_mem,
                 )
             )
@@ -344,10 +386,14 @@ def _evaluate_vis_chunk(
         nfeeds: int,
         polarized: bool = False,
         eps: float = None,
+        upsample_factor: Literal[1.25, 2] = 2,
         beam_spline_opts: dict = None,
         interpolation_function: str = "az_za_map_coordinates",
         n_threads: int = 1,
         is_coplanar: bool = False,
+        use_type1: bool = False,
+        basis_matrix: float = None,
+        type1_n_modes: int = None,
         trace_mem: bool = False,
     ) -> np.ndarray:
         """
@@ -373,6 +419,9 @@ def _evaluate_vis_chunk(
 
         coord_mgr.setup()
 
+        # Check to see if the rotation matrix is an identity matrix
+        is_rotation_identity = np.allclose(rotation_matrix, np.eye(3))
+
         with threadpool_limits(limits=n_threads, user_api="blas"):
             for time_index, ti in enumerate(range(ntimes)[time_idx]):
                 coord_mgr.rotate(ti)
@@ -391,12 +440,21 @@ def _evaluate_vis_chunk(
                 )
 
                 # Rotate source coordinates with rotation matrix.
-                utils.inplace_rot(rotation_matrix, topo)
+                if not is_rotation_identity:
+                    cpu_utils.inplace_rot(rotation_matrix, topo)
+                
+                # Rotate the basis matrix
+                if basis_matrix is not None:
+                    # Rotate antenna positions with basis matrix
+                    cpu_utils.inplace_rot(basis_matrix.T, topo)
+
                 topo *= 2 * np.pi
 
                 for freqidx in range(nfreqs)[freq_idx]:
                     freq = freqs[freqidx]
-                    uvw = bls * freq
+
+                    if not use_type1:
+                        uvw = bls * freq
 
                     # Update beam evaluator for matvis compatibility
                     _cpu_beam_evaluator.beam_list = [beam]
@@ -451,28 +509,42 @@ def _evaluate_vis_chunk(
                         i_sky = i_sky.astype(complex_dtype)
 
                     # Compute visibilities w/ non-uniform FFT
-                    if is_coplanar:
-                        _vis_here = cpu_nufft2d(
-                            topo[0],
-                            topo[1],
+                    if use_type1:
+                        _vis_here = cpu_nufft2d_type1(
+                            topo[0] * freq,
+                            topo[1] * freq,
                             i_sky,
-                            uvw[0],
-                            uvw[1],
+                            n_modes=type1_n_modes,
+                            index=bls,
                             eps=eps,
                             n_threads=n_threads,
+                            upsample_factor=upsample_factor,
                         )
                     else:
-                        _vis_here = cpu_nufft3d(
-                            topo[0],
-                            topo[1],
-                            topo[2],
-                            i_sky,
-                            uvw[0],
-                            uvw[1],
-                            uvw[2],
-                            eps=eps,
-                            n_threads=n_threads,
-                        )
+                        if is_coplanar:
+                            _vis_here = cpu_nufft2d(
+                                topo[0],
+                                topo[1],
+                                i_sky,
+                                uvw[0],
+                                uvw[1],
+                                eps=eps,
+                                n_threads=n_threads,
+                                upsample_factor=upsample_factor,
+                            )
+                        else:
+                            _vis_here = cpu_nufft3d(
+                                topo[0],
+                                topo[1],
+                                topo[2],
+                                i_sky,
+                                uvw[0],
+                                uvw[1],
+                                uvw[2],
+                                eps=eps,
+                                n_threads=n_threads,
+                                upsample_factor=upsample_factor,
+                            )
 
                     vis[time_index, ..., freqidx] = np.swapaxes(
                         _vis_here.reshape(nfeeds, nfeeds, nbls), 2, 0
diff --git a/src/fftvis/cpu/cpu_nufft.py b/src/fftvis/cpu/nufft.py
similarity index 54%
rename from src/fftvis/cpu/cpu_nufft.py
rename to src/fftvis/cpu/nufft.py
index 2296876..9ae2521 100644
--- a/src/fftvis/cpu/cpu_nufft.py
+++ b/src/fftvis/cpu/nufft.py
@@ -6,7 +6,7 @@
 
 import numpy as np
 import finufft
-
+from typing import Literal
 
 def cpu_nufft2d(
     x: np.ndarray,
@@ -16,6 +16,7 @@ def cpu_nufft2d(
     v: np.ndarray,
     eps: float,
     n_threads: int = 1,
+    upsample_factor: Literal[1.25, 2] = 2,
 ) -> np.ndarray:
     """
     Perform a 2D non-uniform FFT on the CPU.
@@ -34,6 +35,8 @@ def cpu_nufft2d(
         V coordinates for baselines.
     eps : float
         Desired accuracy of the transform.
+    upsample_factor : default = 2
+        Upsampling factor for the non-uniform FFT.
     n_threads : int
         Number of threads to use.
 
@@ -52,6 +55,7 @@ def cpu_nufft2d(
         eps=eps,
         nthreads=n_threads,
         showwarn=0,
+        upsampfac=upsample_factor,
     )
 
 
@@ -64,6 +68,7 @@ def cpu_nufft3d(
     v: np.ndarray,
     w: np.ndarray,
     eps: float,
+    upsample_factor: int = 2,
     n_threads: int = 1,
 ) -> np.ndarray:
     """
@@ -87,6 +92,8 @@ def cpu_nufft3d(
         W coordinates for baselines.
     eps : float
         Desired accuracy of the transform.
+    upsample_factor : default = 2
+        Upsampling factor for the non-uniform FFT.
     n_threads : int
         Number of threads to use.
 
@@ -107,4 +114,62 @@ def cpu_nufft3d(
         eps=eps,
         nthreads=n_threads,
         showwarn=0,
+        upsampfac=upsample_factor,
+    )
+
+def cpu_nufft2d_type1(
+    x: np.ndarray,
+    y: np.ndarray,
+    weights: np.ndarray,
+    n_modes: int,
+    index: np.ndarray,
+    eps: float,
+    upsample_factor: int = 2,
+    n_threads: int = 1,
+) -> np.ndarray:
+    """
+    Perform a 2D non-uniform FFT of type 1 on the CPU.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        X coordinates of source positions.
+    y : np.ndarray
+        Y coordinates of source positions.
+    weights : np.ndarray
+        Weights of sources (typically beam-weighted fluxes).
+    n_modes : int
+        Number of fourier modes returned in the Type 1 NUFFT. The resulting
+        transform will be of shape (..., n_modes, n_modes) prior to indexing.
+        This value should be at least as large as the maximum difference between the 
+        largest and smallest index values in the index array.
+    index : np.ndarray
+        2D array of indices to select specific modes from the 2D transform.
+        The shape of index should be (2, n_modes).
+    eps : float
+        Desired accuracy of the transform.
+    upsample_factor : default = 2
+        Upsampling factor for the non-uniform FFT.
+    n_threads : int
+        Number of threads to use.
+
+    Returns
+    -------
+    np.ndarray
+        Visibility data.
+    """
+    # Model is a 2D array of shape (n_modes, n_modes)
+    model = finufft.nufft2d1(
+        x,
+        y,
+        weights,
+        n_modes,
+        modeord=1,
+        eps=eps,
+        nthreads=n_threads,
+        showwarn=0,
+        upsampfac=upsample_factor,
     )
+
+    # Select specific indices from the model
+    return model[..., index[0], index[1]]
\ No newline at end of file
diff --git a/src/fftvis/cpu/cpu_utils.py b/src/fftvis/cpu/utils.py
similarity index 100%
rename from src/fftvis/cpu/cpu_utils.py
rename to src/fftvis/cpu/utils.py
diff --git a/src/fftvis/gpu/__init__.py b/src/fftvis/gpu/__init__.py
index b92c460..b179673 100644
--- a/src/fftvis/gpu/__init__.py
+++ b/src/fftvis/gpu/__init__.py
@@ -1,5 +1,5 @@
 """GPU-specific implementations for fftvis."""
 
-from .gpu_beams import GPUBeamEvaluator
+from .beams import GPUBeamEvaluator
 from .gpu_simulate import GPUSimulationEngine
-from .gpu_nufft import gpu_nufft2d, gpu_nufft3d
+from .nufft import gpu_nufft2d, gpu_nufft3d
diff --git a/src/fftvis/gpu/gpu_beams.py b/src/fftvis/gpu/beams.py
similarity index 100%
rename from src/fftvis/gpu/gpu_beams.py
rename to src/fftvis/gpu/beams.py
diff --git a/src/fftvis/gpu/gpu_nufft.py b/src/fftvis/gpu/nufft.py
similarity index 100%
rename from src/fftvis/gpu/gpu_nufft.py
rename to src/fftvis/gpu/nufft.py
diff --git a/src/fftvis/gpu/gpu_utils.py b/src/fftvis/gpu/utils.py
similarity index 100%
rename from src/fftvis/gpu/gpu_utils.py
rename to src/fftvis/gpu/utils.py
diff --git a/src/fftvis/utils.py b/src/fftvis/utils.py
index d49af98..01efee4 100644
--- a/src/fftvis/utils.py
+++ b/src/fftvis/utils.py
@@ -29,12 +29,6 @@ def _use_gpu():
         return False
 
 
-# Import implementation-specific functions
-if _use_gpu():
-    from .gpu.gpu_utils import inplace_rot
-else:
-    from .cpu.cpu_utils import inplace_rot
-
 __all__ = [
     "IDEALIZED_BL_TOL",
     "speed_of_light",
diff --git a/src/fftvis/wrapper.py b/src/fftvis/wrapper.py
index 4a31861..29b4e0a 100644
--- a/src/fftvis/wrapper.py
+++ b/src/fftvis/wrapper.py
@@ -3,10 +3,11 @@
 from astropy.coordinates import EarthLocation
 from astropy.time import Time
 from pyuvdata.beam_interface import BeamInterface
+from pyuvdata import UVBeam
 from matvis.core.beams import prepare_beam_unpolarized
 
 from .core.beams import BeamEvaluator
-from .cpu.cpu_beams import CPUBeamEvaluator
+from .cpu.beams import CPUBeamEvaluator
 from .core.simulate import SimulationEngine, default_accuracy_dict
 from .cpu.cpu_simulate import CPUSimulationEngine
 
@@ -93,6 +94,7 @@ def simulate_vis(
     precision: int = 2,
     polarized: bool = False,
     eps: float = None,
+    upsample_factor: Literal[1.25, 2] = 2,
     beam_spline_opts: dict = None,
     use_feed: str = "x",
     flat_array_tol: float = 0.0,
@@ -103,6 +105,7 @@ def simulate_vis(
         "CoordinateRotationAstropy", "CoordinateRotationERFA"
     ] = "CoordinateRotationERFA",
     coord_method_params: dict | None = None,
+    force_use_type3: bool = False,
     force_use_ray: bool = False,
     trace_mem: bool = False,
     backend: Literal["cpu", "gpu"] = "cpu",
@@ -145,6 +148,11 @@ def simulate_vis(
         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.
+    upsample_factor : int, default = 2
+        Upsampling factor for the non-uniform fast fourier transform. This is the factor by which the
+        intermediate grid is upsampled. Only values of 1.25 or 2 are allowed. Can be useful for decreasing
+        the computation time and memory requirement for large arrays at the expensive of some accuracy. 
+        The default value is 2.
     beam_spline_opts : dict, optional
         Options to pass to :meth:`pyuvdata.uvbeam.UVBeam.interp` as `spline_opts`.
     use_feed : str, default = "x"
@@ -174,6 +182,9 @@ def simulate_vis(
         for the CoordinateRotationERFA method, there is the parameter ``update_bcrs_every``,
         which should be a time in seconds, for which larger values speed up the computation.
         See the documentation for the CoordinateRotation classes in matvis for more information.
+    force_use_type3: bool, default = False
+        Whether to force the use of type 3 NUFFT. If False, type 3 will only be used
+        if the array cannot be defined in a regular grid.
     force_use_ray: bool, default = False
         Whether to force the use of Ray for parallelization. If False, Ray will only be used
         if nprocesses > 1.
@@ -198,6 +209,15 @@ def simulate_vis(
     # Make sure antpos has the right format
     ants = {k: np.array(v) for k, v in ants.items()}
 
+    # Interpolate the beam to the desired frequencies to avoid redundant
+    # interpolation in the simulation engine
+    if isinstance(beam, UVBeam):
+        if hasattr(beam, "Nfreqs") and beam.Nfreqs > 1:
+            beam = beam.interp(freq_array=freqs, new_object=True, run_check=False) # pragma: no cover
+    elif isinstance(beam, BeamInterface) and beam._isuvbeam:
+        if hasattr(beam.beam, "Nfreqs") and beam.beam.Nfreqs > 1:
+            beam.beam = beam.beam.interp(freq_array=freqs, new_object=True, run_check=False)
+
     beam = BeamInterface(beam)
 
     # Prepare the beam
@@ -221,6 +241,7 @@ def simulate_vis(
         precision=precision,
         polarized=polarized,
         eps=eps,
+        upsample_factor=upsample_factor,
         beam_spline_opts=beam_spline_opts,
         flat_array_tol=flat_array_tol,
         interpolation_function=interpolation_function,
@@ -228,6 +249,7 @@ def simulate_vis(
         nthreads=nthreads,
         coord_method=coord_method,
         coord_method_params=coord_method_params,
+        force_use_type3=force_use_type3,
         force_use_ray=force_use_ray,
         trace_mem=trace_mem,
     )
diff --git a/tests/test_antenna_gridding.py b/tests/test_antenna_gridding.py
new file mode 100644
index 0000000..76b34ff
--- /dev/null
+++ b/tests/test_antenna_gridding.py
@@ -0,0 +1,77 @@
+import numpy as np
+import pytest
+
+from fftvis.core.antenna_gridding import check_antpos_griddability
+
+
+# ------------------------------------------------------------------
+# helpers / fixtures
+# ------------------------------------------------------------------
+def _linear_array(n=10, spacing=1.5):
+    return {i: np.array([i * spacing, 0.0, 0.0]) for i in range(n)}
+
+def _square_grid(n_side=5):
+    """Full n×n Cartesian grid of antennas in the xy‑plane."""
+    antpos = {}
+    idx = 0
+    for i in range(n_side):
+        for j in range(n_side):
+            antpos[idx] = np.array([i, j, 0.0])
+            idx += 1
+    return antpos
+
+def _square_grid_with_holes():
+    antpos = _square_grid()
+    for k in (3, 16):        # remove a couple of elements
+        antpos.pop(k)
+    return antpos
+
+def _scattered_xy(n=15, seed=42):
+    """
+    Deterministic pseudo‑random scatter in the xy‑plane.
+    No single linear transform can map these points onto
+    an integer lattice within sane tolerance.
+    """
+    rng = np.random.default_rng(seed)
+    return {
+        i: np.array([rng.uniform(0, 10), rng.uniform(0, 10), 0.0])
+        for i in range(n)
+    }
+
+def _hex_grid():
+    return {
+        0: np.array([-0.5,  np.sqrt(3) / 2,  0.0]),
+        1: np.array([ 0.5,  np.sqrt(3) / 2,  0.0]),
+        2: np.array([-1.0,  0.0,             0.0]),
+        3: np.array([ 0.0,  0.0,             0.0]),
+        4: np.array([ 1.0,  0.0,             0.0]),
+        5: np.array([-0.5, -np.sqrt(3) / 2,  0.0]),
+        6: np.array([ 0.5, -np.sqrt(3) / 2,  0.0]),
+    }
+
+
+# ------------------------------------------------------------------
+# parametrised test
+# ------------------------------------------------------------------
+@pytest.mark.parametrize(
+    "antpos, expected_griddable",
+    [
+        pytest.param(_linear_array(), True, id="linear"),
+        pytest.param(_square_grid(), True, id="square-full"),
+        pytest.param(_square_grid_with_holes(), True, id="square-holey"),
+        pytest.param(_hex_grid(), True, id="hex-grid"),
+        pytest.param(_scattered_xy(), False, id="non‑griddable"),
+    ],
+)
+def test_check_antpos_griddability(antpos, expected_griddable):
+    """All griddability scenarios in one parametrised test."""
+    is_griddable, modified_antpos, _ = check_antpos_griddability(antpos)
+
+    # 1) overall flag matches expectation
+    assert is_griddable is expected_griddable
+
+    # 2) if griddable, every coordinate should round cleanly to ints
+    if expected_griddable:
+        for ant_id, pos in modified_antpos.items():
+            msg = f"Antenna {ant_id} not on integer grid: {pos}"
+            assert np.allclose(pos, np.round(pos).astype(int)), msg
\ No newline at end of file
diff --git a/tests/test_beam_evaluator.py b/tests/test_beam_evaluator.py
index 3595dcf..b950f25 100644
--- a/tests/test_beam_evaluator.py
+++ b/tests/test_beam_evaluator.py
@@ -1,6 +1,6 @@
 import pytest
 import numpy as np
-from fftvis.cpu.cpu_beams import CPUBeamEvaluator
+from fftvis.cpu.beams import CPUBeamEvaluator
 from fftvis.core.beams import BeamEvaluator
 
 from pathlib import Path
@@ -199,7 +199,7 @@ def test_beam_evaluator_interp_branches():
     2. The conversion from tx/ty to az/za works properly
     3. All code paths in the interp method are exercised
     """
-    from fftvis.cpu.cpu_beams import CPUBeamEvaluator
+    from fftvis.cpu.beams import CPUBeamEvaluator
     import numpy as np
     from pyuvdata import UVBeam
     from pyuvdata.beam_interface import BeamInterface
@@ -271,7 +271,7 @@ def test_beam_evaluator_interp_polarized_branches():
     1. The interp method correctly handles polarized beams with beam_values.ndim == 3
     2. The missing branch in the interp method is exercised
     """
-    from fftvis.cpu.cpu_beams import CPUBeamEvaluator
+    from fftvis.cpu.beams import CPUBeamEvaluator
     import numpy as np
     from unittest.mock import patch, MagicMock
     
diff --git a/tests/test_core_utils.py b/tests/test_core_utils.py
index 876c719..d6e0082 100644
--- a/tests/test_core_utils.py
+++ b/tests/test_core_utils.py
@@ -174,23 +174,8 @@ def test_utils_module_imports():
     """Test that the fftvis.utils module correctly imports functions."""
     # Check that the utils module has all the expected functions
     for func_name in ["IDEALIZED_BL_TOL", "speed_of_light", "get_pos_reds", 
-                     "get_plane_to_xy_rotation_matrix", "get_task_chunks", "inplace_rot"]:
+                     "get_plane_to_xy_rotation_matrix", "get_task_chunks"]:
         assert hasattr(fftvis.utils, func_name)
-    
-    # inplace_rot should be imported from CPU or GPU
-    assert hasattr(fftvis.utils, "inplace_rot")
-    
-    # Test calling the function to check it's properly imported
-    # Create a simple rotation test
-    rot = np.eye(3)
-    b = np.zeros((3, 2))
-    
-    # This should not raise an error if properly imported from CPU
-    try:
-        fftvis.utils.inplace_rot(rot, b)
-    except NotImplementedError:
-        # If GPU version, this will raise NotImplementedError, which is fine
-        pass
 
 
 @pytest.mark.parametrize("with_cupy", [False, True])
@@ -230,4 +215,4 @@ def mock_import_error(name, *args, **kwargs):
         assert fftvis.utils._use_gpu() is False
         
         # Restore original import
-        monkeypatch.undo()
+        monkeypatch.undo()
\ No newline at end of file
diff --git a/tests/test_cpu_beams.py b/tests/test_cpu_beams.py
index 41d4e93..c89c2ff 100644
--- a/tests/test_cpu_beams.py
+++ b/tests/test_cpu_beams.py
@@ -1,6 +1,6 @@
 import pytest
 import numpy as np
-from fftvis.cpu.cpu_beams import CPUBeamEvaluator
+from fftvis.cpu.beams import CPUBeamEvaluator
 from fftvis.core.beams import BeamEvaluator
 
 from pathlib import Path
diff --git a/tests/test_cpu_simulate.py b/tests/test_cpu_simulate.py
index 33ff6b6..9a11022 100644
--- a/tests/test_cpu_simulate.py
+++ b/tests/test_cpu_simulate.py
@@ -47,6 +47,26 @@ def get_telescope(telescope_name, **kwargs):
 
 from fftvis.core.simulate import SimulationEngine
 
+def _hex_grid():
+    return {
+        0: np.array([-0.5,  np.sqrt(3) / 2,  0.0]),
+        1: np.array([ 0.5,  np.sqrt(3) / 2,  0.0]),
+        2: np.array([-1.0,  0.0,             0.0]),
+        3: np.array([ 0.0,  0.0,             0.0]),
+        4: np.array([ 1.0,  0.0,             0.0]),
+        5: np.array([-0.5, -np.sqrt(3) / 2,  0.0]),
+        6: np.array([ 0.5, -np.sqrt(3) / 2,  0.0]),
+    }
+
+def _square_grid(n_side=3, spacing=10.0):
+    """Full n×n Cartesian grid of antennas in the xy‑plane."""
+    antpos = {}
+    idx = 0
+    for i in range(n_side):
+        for j in range(n_side):
+            antpos[idx] = spacing * np.array([i, j, 0.0])
+            idx += 1
+    return antpos
 
 @pytest.mark.parametrize("polarized", [False, True])
 @pytest.mark.parametrize("precision", [2, 1])
@@ -98,6 +118,7 @@ def test_simulate(
         use_analytic_beam=use_analytic_beam, polarized=polarized
     )
     ants = params.pop("ants")
+
     if tilt_array:
         # Tilt the array
         ants = {
@@ -123,7 +144,7 @@ def test_simulate(
 
     # Use fftvis to simulate visibilities
     fvis = simulate_vis(
-        ants,
+        ants=ants,
         eps=1e-10 if precision == 2 else 6e-8,
         baselines=sim_baselines,
         precision=precision,
@@ -138,7 +159,7 @@ def test_simulate(
     )
 
     fvis_all_bls = simulate_vis(
-        ants,
+        ants=ants,
         eps=1e-10 if precision == 2 else 6e-8,
         precision=precision,
         nprocesses=nprocesses,
@@ -171,6 +192,80 @@ def test_simulate(
         )
 
 
+@pytest.mark.parametrize("polarized", [False, True])
+@pytest.mark.parametrize("precision", [2, 1])
+@pytest.mark.parametrize("shear_array", [True, False])
+@pytest.mark.parametrize("rotate_array", [True, False])
+@pytest.mark.parametrize("remove_antennas", [True, False])
+@pytest.mark.parametrize("ants", [_hex_grid(), _square_grid()])
+def test_simulate_gridded_type1_vs_type3(polarized, precision, shear_array, rotate_array, remove_antennas, ants):
+    rng = np.random.default_rng(42)
+
+    params, *_ = get_standard_sim_params(
+        use_analytic_beam=True, polarized=polarized
+    )
+    params.pop("ants")
+
+    # Remove antennas randomly
+    if remove_antennas:
+        ants = {k: ants[k] for k in ants if rng.uniform(0, 1) > 0.25}
+        ants = {ki: ants[k] for ki, k in enumerate(ants)}
+
+    # Rigidly rotate antenna layout
+    if rotate_array:
+        theta = np.pi / 2
+        rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1]])
+        # Rotate the array
+        ants = {
+            ant: np.dot(rotation_matrix, ants[ant])
+            for ai, ant in enumerate(ants)
+        }
+    if shear_array:
+        shear_matrix = np.array([[1, 0.5, 0], [0, 1, 0], [0, 0, 1]])
+        # shear the array
+        ants = {
+            ant: np.dot(shear_matrix, ants[ant])
+            for ai, ant in enumerate(ants)
+        }
+
+    # List of baselines
+    baselines = [(i, j) for i in ants for j in ants if j >= i]
+
+    # Get the beam
+    beam = params.pop("beams")[0]
+    times = params.pop("times").jd
+
+    # Use fftvis type 1
+    fvis_type1 = simulate_vis(
+        ants=ants,
+        eps=1e-10 if precision == 2 else 6e-8,
+        precision=precision,
+        coord_method_params={"source_buffer": 0.75},
+        beam=beam,
+        times=times,
+        force_use_type3=False,
+        baselines=baselines,
+        **params,
+    )
+
+    # Use fftvis type 3
+    fvis_type3 = simulate_vis(
+        ants=ants,
+        eps=1e-10 if precision == 2 else 6e-8,
+        precision=precision,
+        coord_method_params={"source_buffer": 0.75},
+        beam=beam,
+        times=times,
+        force_use_type3=True,
+        baselines=baselines,
+        **params,
+    )
+
+    # Compare the results of type 1 and type 3
+    np.testing.assert_allclose(
+        fvis_type1, fvis_type3, atol=1e-5 if precision == 2 else 1e-4
+    )
+
 def test_cpu_simulation_engine_init():
     """Test that the CPUSimulationEngine initializes correctly."""
     engine = CPUSimulationEngine()
diff --git a/tests/test_gpu_beams.py b/tests/test_gpu_beams.py
index 5c62df5..afc28d6 100644
--- a/tests/test_gpu_beams.py
+++ b/tests/test_gpu_beams.py
@@ -1,7 +1,7 @@
 import pytest
 import numpy as np
-from fftvis.gpu.gpu_beams import GPUBeamEvaluator
-from fftvis.gpu.gpu_utils import inplace_rot
+from fftvis.gpu.beams import GPUBeamEvaluator
+from fftvis.gpu.utils import inplace_rot
 
 
 def test_gpu_beam_evaluator_init():
diff --git a/tests/test_gpu_nufft.py b/tests/test_gpu_nufft.py
index e02d77e..b2b3f19 100644
--- a/tests/test_gpu_nufft.py
+++ b/tests/test_gpu_nufft.py
@@ -1,6 +1,6 @@
 import pytest
 import numpy as np
-from fftvis.gpu.gpu_nufft import gpu_nufft2d, gpu_nufft3d
+from fftvis.gpu.nufft import gpu_nufft2d, gpu_nufft3d
 from fftvis.gpu.gpu_simulate import GPUSimulationEngine
 
 


Reviews:
user: login: steven-murray
body: Hey Tyler, this is super cool, thanks! I have a few broader questions/comments:

1. I would move the new `cpu_utils.py` stuff into its own module (like `antenna_gridding.py` or something nicer)
2. I don't know why we have `cpu/cpu_nufft.py` etc instead of just `cpu/nufft.py` which is much nicer.
3. How does the antenna rotation work? Rigid-Rotating a hex array should not be able to shear the antennas, right?
4. I would like to see a test comparing a gridded to non-gridded calculation for the same ant positions and sky model. This would mean we also need an option to tell the simulator to NOT try gridding even if it could grid.
5. I also think that we should let the user tell the simulate function that yes, indeed, I want to use gridding, and I know just how to grid it. Probably by passing a rotation matrix and argument like "griddable = True" (or maybe just the rotation matrix). That way we can skip the checks and just assume the user knows what they're doing. Helps with speed a little bit.
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T10:02:06Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-08T20:29:46Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-08T20:35:08Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-08T20:51:21Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-08T23:06:04Z
user: login: steven-murray
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T08:37:07Z
user: login: steven-murray
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T08:39:06Z
user: login: steven-murray
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T08:41:23Z
user: login: steven-murray
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T08:41:46Z
user: login: steven-murray
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T08:44:34Z
user: login: steven-murray
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T08:45:23Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-09T14:58:43Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-09T17:05:45Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-10T06:09:27Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-10T06:09:49Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-10T06:20:46Z
user: login: steven-murray
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T13:22:17Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-12T18:00:23Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-12T18:01:00Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-12T18:01:15Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-12T18:02:05Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-12T18:04:28Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-12T18:04:55Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-12T18:07:43Z
user: login: tyler-a-cox
state: COMMENTED
author_association: OWNER
submitted_at: 2025-05-12T18:15:41Z
user: login: steven-murray
body: Thanks @tyler-a-cox this looks very good now. Maybe the only thing would be to add an FAQ or tutorial to the docs describing this option.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-12T19:46:56Z

Review comments:
diff_hunk: @@ -206,3 +223,134 @@ def inplace_rot_base(rot, b):
         out[1] = rot[1, 0] * b[0, n] + rot[1, 1] * b[1, n] + rot[1, 2] * b[2, n]
         out[2] = rot[2, 0] * b[0, n] + rot[2, 1] * b[1, n] + rot[2, 2] * b[2, n]
         b[:, n] = out
+
+
+def find_integer_multiplier(
+    arr: np.ndarray,
+    max_denominator: int = 10**6
+) -> int:
+    """
+    Return the smallest positive integer f such that f * arr[i] is integral
+    for all entries (up to rational approximation). Zeros in arr are ignored.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    max_denominator : int
+        Maximum denominator for the fractions.
+
+    Returns
+    -------
+    f : int
+        The integer factor used for scaling.
+    """
+    fracs = [
+        Fraction(val).limit_denominator(max_denominator)
+        for val in np.ravel(arr)
+        if val != 0
+    ]
+    if not fracs:
+        return 1
+    dens = (frac.denominator for frac in fracs)
+    return lcm(*dens)
+
+
+def can_scale_to_int(
+    arr: np.ndarray,
+    tol: float = 1e-9,
+    max_denominator: int = 10**6,
+    max_factor: int = None
+) -> Tuple[bool, int]:
+    """
+    Check if there exists an integer factor f such that f * arr is
+    approximately integers. Returns (is_griddable, factor).
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    tol : float
+        Tolerance for checking if values are close to integers.
+    max_denominator : int
+        Maximum denominator for the fractions.
+    max_factor : int
+        Maximum allowed factor for scaling.
+
+    Returns
+    -------
+    is_griddable : bool
+        True if the array can be scaled to integers, False otherwise.
+    factor : int
+        The integer factor used for scaling.
+    """
+    f = find_integer_multiplier(arr, max_denominator)
+    if max_factor is not None and f > max_factor:
+        return False, f
+    scaled = f * np.array(arr, dtype=float)
+    if np.allclose(scaled, np.round(scaled), atol=tol):
+        return True, f
+    return False, f
+
+
+def check_antpos_griddability(
path: src/fftvis/core/utils.py
user: login: steven-murray
body: I suggest we add an argument like "rotation_matrix" to this, so the user can supply their own transform in the case that it's not a rectangular/hexagonal layout. 
created_at: 2025-05-08T09:48:43Z
author_association: COLLABORATOR
diff_hunk: @@ -108,3 +108,52 @@ def cpu_nufft3d(
         nthreads=n_threads,
         showwarn=0,
     )
+
+def cpu_nufft2d_type1(
+    x: np.ndarray,
+    y: np.ndarray,
+    weights: np.ndarray,
+    n_modes: int,
+    index: np.ndarray,
+    eps: float,
+    n_threads: int = 1,
+) -> np.ndarray:
+    """
+    Perform a 2D non-uniform FFT of type 1 on the CPU.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        X coordinates of source positions.
+    y : np.ndarray
+        Y coordinates of source positions.
+    weights : np.ndarray
+        Weights of sources (typically beam-weighted fluxes).
+    n_modes : int
+        Number of modes in the transform.
+    index : np.ndarray
path: src/fftvis/cpu/cpu_nufft.py
user: login: steven-murray
body: If this is to be a public-facing API, then I suggest documenting this argument better (e.g. say that it needs to be a 2D array,a nd what the shape corresponds to).
created_at: 2025-05-08T09:50:23Z
author_association: COLLABORATOR
diff_hunk: @@ -391,12 +437,14 @@ def _evaluate_vis_chunk(
                 )
 
                 # Rotate source coordinates with rotation matrix.
-                utils.inplace_rot(rotation_matrix, topo)
+                inplace_rot(rotation_matrix, topo)
path: src/fftvis/cpu/cpu_simulate.py
user: login: steven-murray
body: Does the math not work out to be able to do the rotation on the equatorial coordinates, so we only have to do it once?
created_at: 2025-05-08T09:56:15Z
author_association: COLLABORATOR
diff_hunk: @@ -231,3 +231,58 @@ def mock_import_error(name, *args, **kwargs):
         
         # Restore original import
         monkeypatch.undo()
+
+def test_check_array_griddability():
path: tests/test_core_utils.py
user: login: steven-murray
body: I suggest making this test a class: `TestCheckArrayGriddability` and then make each subtest a new test function (or maybe use parametrization). That will help with test failure identification.
created_at: 2025-05-08T09:57:53Z
author_association: COLLABORATOR
diff_hunk: @@ -108,3 +108,52 @@ def cpu_nufft3d(
         nthreads=n_threads,
         showwarn=0,
     )
+
+def cpu_nufft2d_type1(
+    x: np.ndarray,
+    y: np.ndarray,
+    weights: np.ndarray,
+    n_modes: int,
+    index: np.ndarray,
+    eps: float,
+    n_threads: int = 1,
+) -> np.ndarray:
+    """
+    Perform a 2D non-uniform FFT of type 1 on the CPU.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        X coordinates of source positions.
+    y : np.ndarray
+        Y coordinates of source positions.
+    weights : np.ndarray
+        Weights of sources (typically beam-weighted fluxes).
+    n_modes : int
+        Number of modes in the transform.
+    index : np.ndarray
path: src/fftvis/cpu/cpu_nufft.py
user: login: tyler-a-cox
body: Added slightly more description in the docstring.
created_at: 2025-05-08T20:29:46Z
author_association: OWNER
diff_hunk: @@ -391,12 +437,14 @@ def _evaluate_vis_chunk(
                 )
 
                 # Rotate source coordinates with rotation matrix.
-                utils.inplace_rot(rotation_matrix, topo)
+                inplace_rot(rotation_matrix, topo)
path: src/fftvis/cpu/cpu_simulate.py
user: login: tyler-a-cox
body: I think there are some cases when you can avoid doing this. Certainly, if the rotation matrix is the identity, we could just skip this, but I think if the rotation matrix requires shearing like with the hex array, I don't think we can pre-rotate, but I may be wrong. I'll have a think about this, but in the meantime, I'll add a check that passes on this if the rotation matrix is the identity.
created_at: 2025-05-08T20:35:08Z
author_association: OWNER
diff_hunk: @@ -206,3 +223,134 @@ def inplace_rot_base(rot, b):
         out[1] = rot[1, 0] * b[0, n] + rot[1, 1] * b[1, n] + rot[1, 2] * b[2, n]
         out[2] = rot[2, 0] * b[0, n] + rot[2, 1] * b[1, n] + rot[2, 2] * b[2, n]
         b[:, n] = out
+
+
+def find_integer_multiplier(
+    arr: np.ndarray,
+    max_denominator: int = 10**6
+) -> int:
+    """
+    Return the smallest positive integer f such that f * arr[i] is integral
+    for all entries (up to rational approximation). Zeros in arr are ignored.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    max_denominator : int
+        Maximum denominator for the fractions.
+
+    Returns
+    -------
+    f : int
+        The integer factor used for scaling.
+    """
+    fracs = [
+        Fraction(val).limit_denominator(max_denominator)
+        for val in np.ravel(arr)
+        if val != 0
+    ]
+    if not fracs:
+        return 1
+    dens = (frac.denominator for frac in fracs)
+    return lcm(*dens)
+
+
+def can_scale_to_int(
+    arr: np.ndarray,
+    tol: float = 1e-9,
+    max_denominator: int = 10**6,
+    max_factor: int = None
+) -> Tuple[bool, int]:
+    """
+    Check if there exists an integer factor f such that f * arr is
+    approximately integers. Returns (is_griddable, factor).
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    tol : float
+        Tolerance for checking if values are close to integers.
+    max_denominator : int
+        Maximum denominator for the fractions.
+    max_factor : int
+        Maximum allowed factor for scaling.
+
+    Returns
+    -------
+    is_griddable : bool
+        True if the array can be scaled to integers, False otherwise.
+    factor : int
+        The integer factor used for scaling.
+    """
+    f = find_integer_multiplier(arr, max_denominator)
+    if max_factor is not None and f > max_factor:
+        return False, f
+    scaled = f * np.array(arr, dtype=float)
+    if np.allclose(scaled, np.round(scaled), atol=tol):
+        return True, f
+    return False, f
+
+
+def check_antpos_griddability(
path: src/fftvis/core/utils.py
user: login: tyler-a-cox
body: Added an argument that allows the user to pass in their own rotation matrix.
created_at: 2025-05-08T20:51:21Z
author_association: OWNER
diff_hunk: @@ -231,3 +231,58 @@ def mock_import_error(name, *args, **kwargs):
         
         # Restore original import
         monkeypatch.undo()
+
+def test_check_array_griddability():
path: tests/test_core_utils.py
user: login: tyler-a-cox
body: Reorganized this using parametrization.
created_at: 2025-05-08T23:06:04Z
author_association: OWNER
diff_hunk: @@ -110,7 +111,9 @@ def simulate(
             "CoordinateRotationAstropy", "CoordinateRotationERFA"
         ] = "CoordinateRotationERFA",
         coord_method_params: dict | None = None,
+        skew_to_cartesian: np.ndarray | None = None,
path: src/fftvis/cpu/cpu_simulate.py
user: login: steven-murray
body: Did we add docstrings for these into the base class?
created_at: 2025-05-09T08:37:07Z
author_association: COLLABORATOR
diff_hunk: @@ -174,6 +176,15 @@ def simulate_vis(
         for the CoordinateRotationERFA method, there is the parameter ``update_bcrs_every``,
         which should be a time in seconds, for which larger values speed up the computation.
         See the documentation for the CoordinateRotation classes in matvis for more information.
+    skew_to_cartesian : np.ndarray, default = None
+        Linear change of basis matrix that maps antenna positions expressed in a skewed/oblique grid 
+        into the standard Cartesian frame. This is used to convert the antenna positions from the skewed grid
+        to a Cartesian grid that is compatible with a more efficient NUFFT algorithm. By default fftvis checks
+        if the antenna positions are on a regular or hexagonal grid, but other grid types are not automatically
+        supported. The matrix should be of shape (3, 3).
+    force_use_type3: bool, default = False
+        Whether to force the use of type 3 NUFFT. If False, type 3 will only be used
+        if the array can be defined in a regular grid.
path: src/fftvis/wrapper.py
user: login: steven-murray
body: Isn't this the wrong way around? Type 3 will be used if the array *cannot* be defined on a regular grid?
created_at: 2025-05-09T08:39:05Z
author_association: COLLABORATOR
diff_hunk: @@ -148,11 +158,12 @@ def test_simulate(
         times=times,
         backend=backend,
         force_use_ray=force_use_ray,
+        force_use_type3=force_use_type3,
         **params,
     )
 
     # Check shape of result when no baselines are specified
-    nbls = (len(ants) * (len(ants) - 1)) // 2 + 1
+    nbls = int((len(ants) * (len(ants) - 1)) / 2)
path: tests/test_cpu_simulate.py
user: login: steven-murray
body: Why did this change?
created_at: 2025-05-09T08:41:22Z
author_association: COLLABORATOR
diff_hunk: @@ -206,3 +223,134 @@ def inplace_rot_base(rot, b):
         out[1] = rot[1, 0] * b[0, n] + rot[1, 1] * b[1, n] + rot[1, 2] * b[2, n]
         out[2] = rot[2, 0] * b[0, n] + rot[2, 1] * b[1, n] + rot[2, 2] * b[2, n]
         b[:, n] = out
+
+
+def find_integer_multiplier(
+    arr: np.ndarray,
+    max_denominator: int = 10**6
+) -> int:
+    """
+    Return the smallest positive integer f such that f * arr[i] is integral
+    for all entries (up to rational approximation). Zeros in arr are ignored.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    max_denominator : int
+        Maximum denominator for the fractions.
+
+    Returns
+    -------
+    f : int
+        The integer factor used for scaling.
+    """
+    fracs = [
+        Fraction(val).limit_denominator(max_denominator)
+        for val in np.ravel(arr)
+        if val != 0
+    ]
+    if not fracs:
+        return 1
+    dens = (frac.denominator for frac in fracs)
+    return lcm(*dens)
+
+
+def can_scale_to_int(
+    arr: np.ndarray,
+    tol: float = 1e-9,
+    max_denominator: int = 10**6,
+    max_factor: int = None
+) -> Tuple[bool, int]:
+    """
+    Check if there exists an integer factor f such that f * arr is
+    approximately integers. Returns (is_griddable, factor).
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    tol : float
+        Tolerance for checking if values are close to integers.
+    max_denominator : int
+        Maximum denominator for the fractions.
+    max_factor : int
+        Maximum allowed factor for scaling.
+
+    Returns
+    -------
+    is_griddable : bool
+        True if the array can be scaled to integers, False otherwise.
+    factor : int
+        The integer factor used for scaling.
+    """
+    f = find_integer_multiplier(arr, max_denominator)
+    if max_factor is not None and f > max_factor:
+        return False, f
+    scaled = f * np.array(arr, dtype=float)
+    if np.allclose(scaled, np.round(scaled), atol=tol):
+        return True, f
+    return False, f
+
+
+def check_antpos_griddability(
path: src/fftvis/core/utils.py
user: login: steven-murray
body: Sweet, thanks!
created_at: 2025-05-09T08:41:45Z
author_association: COLLABORATOR
diff_hunk: @@ -391,12 +437,14 @@ def _evaluate_vis_chunk(
                 )
 
                 # Rotate source coordinates with rotation matrix.
-                utils.inplace_rot(rotation_matrix, topo)
+                inplace_rot(rotation_matrix, topo)
path: src/fftvis/cpu/cpu_simulate.py
user: login: steven-murray
body: I think we want to stick away from "rotation matrix" terminology because rotation matrices don't shear. It's not obvious to me that transforming baselines with a shear is able to reproduce the same visibilities after inversely transforming the source positions, but it seems like it does! I'd love to be convinced of this :-)
created_at: 2025-05-09T08:44:34Z
author_association: COLLABORATOR
diff_hunk: @@ -231,3 +231,58 @@ def mock_import_error(name, *args, **kwargs):
         
         # Restore original import
         monkeypatch.undo()
+
+def test_check_array_griddability():
path: tests/test_core_utils.py
user: login: steven-murray
body: I didn't see the updated parameterized test, so maybe you forgot to push that?
created_at: 2025-05-09T08:45:23Z
author_association: COLLABORATOR
diff_hunk: @@ -174,6 +176,15 @@ def simulate_vis(
         for the CoordinateRotationERFA method, there is the parameter ``update_bcrs_every``,
         which should be a time in seconds, for which larger values speed up the computation.
         See the documentation for the CoordinateRotation classes in matvis for more information.
+    skew_to_cartesian : np.ndarray, default = None
+        Linear change of basis matrix that maps antenna positions expressed in a skewed/oblique grid 
+        into the standard Cartesian frame. This is used to convert the antenna positions from the skewed grid
+        to a Cartesian grid that is compatible with a more efficient NUFFT algorithm. By default fftvis checks
+        if the antenna positions are on a regular or hexagonal grid, but other grid types are not automatically
+        supported. The matrix should be of shape (3, 3).
+    force_use_type3: bool, default = False
+        Whether to force the use of type 3 NUFFT. If False, type 3 will only be used
+        if the array can be defined in a regular grid.
path: src/fftvis/wrapper.py
user: login: tyler-a-cox
body: Yeah, that's right. I'll change that. Thanks for the catch.
created_at: 2025-05-09T14:58:43Z
author_association: OWNER
diff_hunk: @@ -148,11 +158,12 @@ def test_simulate(
         times=times,
         backend=backend,
         force_use_ray=force_use_ray,
+        force_use_type3=force_use_type3,
         **params,
     )
 
     # Check shape of result when no baselines are specified
-    nbls = (len(ants) * (len(ants) - 1)) // 2 + 1
+    nbls = int((len(ants) * (len(ants) - 1)) / 2)
path: tests/test_cpu_simulate.py
user: login: tyler-a-cox
body: It was just differences in number of baselines returned due to me moving from a random array to a regular array. I've changed this back and separated the gridded tests.
created_at: 2025-05-09T17:05:45Z
author_association: OWNER
diff_hunk: @@ -110,7 +111,9 @@ def simulate(
             "CoordinateRotationAstropy", "CoordinateRotationERFA"
         ] = "CoordinateRotationERFA",
         coord_method_params: dict | None = None,
+        skew_to_cartesian: np.ndarray | None = None,
path: src/fftvis/cpu/cpu_simulate.py
user: login: tyler-a-cox
body: Removed this in favor of a more general gridding function that finds a set of basis vectors for the basis transformation.
created_at: 2025-05-10T06:09:27Z
author_association: OWNER
diff_hunk: @@ -231,3 +231,58 @@ def mock_import_error(name, *args, **kwargs):
         
         # Restore original import
         monkeypatch.undo()
+
+def test_check_array_griddability():
path: tests/test_core_utils.py
user: login: tyler-a-cox
body: You're right, I'd thought I'd pushed that earlier. Should be updated now!
created_at: 2025-05-10T06:09:49Z
author_association: OWNER
diff_hunk: @@ -391,12 +437,14 @@ def _evaluate_vis_chunk(
                 )
 
                 # Rotate source coordinates with rotation matrix.
-                utils.inplace_rot(rotation_matrix, topo)
+                inplace_rot(rotation_matrix, topo)
path: src/fftvis/cpu/cpu_simulate.py
user: login: tyler-a-cox
body: Yeah, good point! I've added a new parameter to `_evaluate_vis_chunk` that carries the basis transformation matrix. I've added a series of new tests that compares visibilities calculated with the Type 1 transform against matvis for gridded matrices that have been skewed, rotated, etc.
created_at: 2025-05-10T06:20:46Z
author_association: OWNER
diff_hunk: @@ -0,0 +1,194 @@
+import numpy as np
+from math import lcm
+from fractions import Fraction
+from typing import Any, Dict, Tuple
+
+
+def find_integer_multiplier(
+    arr: np.ndarray,
+    max_denominator: int = 10**6
+) -> int:
+    """
+    Return the smallest positive integer f such that f * arr[i] is integral
+    for all entries (up to rational approximation). Zeros in arr are ignored.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    max_denominator : int
+        Maximum denominator for the fractions.
+
+    Returns
+    -------
+    f : int
+        The integer factor used for scaling.
+    """
+    fracs = [
+        Fraction(val).limit_denominator(max_denominator)
+        for val in np.ravel(arr)
+        if val != 0
+    ]
+    if not fracs:
+        return 1
+    dens = (frac.denominator for frac in fracs)
+    return lcm(*dens)
+
+
+def can_scale_to_int(
+    arr: np.ndarray,
+    tol: float = 1e-9,
+    max_denominator: int = 10**6,
+    max_factor: int = None
+) -> Tuple[bool, int]:
+    """
+    Check if there exists an integer factor f such that f * arr is
+    approximately integers. Returns (is_griddable, factor).
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    tol : float
+        Tolerance for checking if values are close to integers.
+    max_denominator : int
+        Maximum denominator for the fractions.
+    max_factor : int
+        Maximum allowed factor for scaling.
+
+    Returns
+    -------
+    is_griddable : bool
+        True if the array can be scaled to integers, False otherwise.
+    factor : int
+        The integer factor used for scaling.
+    """
+    f = find_integer_multiplier(arr, max_denominator)
+    if max_factor is not None and f > max_factor:
+        return False, f
+    scaled = f * np.array(arr, dtype=float)
+    if np.allclose(scaled, np.round(scaled), atol=tol):
+        return True, f
+    return False, f
+
+def find_lattice_basis(
+    antpos: Dict[Any, np.ndarray],
+    tol: float = 1e-9,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Find the lattice basis for the antenna positions.
+
+    Parameters
+    ----------
+    antpos : dict
+        Dictionary of antenna positions in the form {ant_index: np.array([x,y,z])}.
+    tol : float
+        Tolerance for checking if values are close to integers.
path: src/fftvis/core/antenna_gridding.py
user: login: steven-murray
body: This is ultimately in meters, right? Or, in the units of whatever antpos is in? We should say so.
created_at: 2025-05-12T13:04:56Z
author_association: COLLABORATOR
diff_hunk: @@ -37,6 +37,7 @@ def simulate(
         precision: int = 2,
         polarized: bool = False,
         eps: float = None,
+        upsampfac: int = 2,
path: src/fftvis/core/simulate.py
user: login: steven-murray
body: This should not be `int` since it can be 1.25. In fact, since it can only be either 1.25 or 2, I would give it the type `Literal[1.25, 2]`
created_at: 2025-05-12T13:06:51Z
author_association: COLLABORATOR
diff_hunk: @@ -37,6 +37,7 @@ def simulate(
         precision: int = 2,
         polarized: bool = False,
         eps: float = None,
+        upsampfac: int = 2,
path: src/fftvis/core/simulate.py
user: login: steven-murray
body: I also don't like `upsampfac` as a name, sorry! How about just `upsample_factor`?
created_at: 2025-05-12T13:07:34Z
author_association: COLLABORATOR
diff_hunk: @@ -193,6 +204,11 @@ def _evaluate_vis_chunk(
             Number of threads to use.
         is_coplanar : bool
             Whether the array is coplanar.
+        basis_matrix : np.ndarray
path: src/fftvis/core/simulate.py
user: login: steven-murray
body: Maybe we should add that if None is given, we will automatically attempt to find one that would speed up the calculation.
created_at: 2025-05-12T13:09:35Z
author_association: COLLABORATOR
diff_hunk: @@ -193,6 +204,11 @@ def _evaluate_vis_chunk(
             Number of threads to use.
         is_coplanar : bool
             Whether the array is coplanar.
+        basis_matrix : np.ndarray
+            Lattice basis matrix used to grid baselines and sources for a type 1
+            non-uniform FFT.
+        type1_n_modes : int
+            Number of modes for type 1 non-uniform FFT.
path: src/fftvis/core/simulate.py
user: login: steven-murray
body: To me as a user, I would not understand what this means -- modes of what? What is the relevant range of values?
created_at: 2025-05-12T13:10:10Z
author_association: COLLABORATOR
diff_hunk: @@ -201,16 +203,12 @@ def simulate(
             n_modes = 2 * int(np.round(np.max(np.abs(bls)))) + 1
 
             # Get the maximum baseline length for proper coordinate scaling
-            max_baseline_coord_length = np.max(np.abs(
-               [
-                   ants[bl[1]] - ants[bl[0]]
-                   for bl in baselines
-               ]
-            ))
-            type1_coord_scalar = max_baseline_coord_length / utils.speed_of_light / (n_modes // 2)
+            basis_matrix *= 1 / utils.speed_of_light
+            basis_matrix = basis_matrix.astype(real_dtype)
 
             # Assume the array is coplanar for gridded coordinates
             is_coplanar = True
+            rotation_matrix = np.eye(3, dtype=real_dtype)
path: src/fftvis/cpu/cpu_simulate.py
user: login: steven-murray
body: Is this a stray reference to `rotation_matrix`?
created_at: 2025-05-12T13:12:27Z
author_association: COLLABORATOR
diff_hunk: @@ -437,7 +437,14 @@ def _evaluate_vis_chunk(
                 )
 
                 # Rotate source coordinates with rotation matrix.
-                inplace_rot(rotation_matrix, topo)
+                if not np.allclose(rotation_matrix, np.eye(3)):
path: src/fftvis/cpu/cpu_simulate.py
user: login: steven-murray
body: This seems like a small time-waste to me -- we should just have a boolean telling us if `rotation_matrix` is identity
created_at: 2025-05-12T13:14:27Z
author_association: COLLABORATOR
diff_hunk: @@ -32,9 +31,9 @@ def _use_gpu():
 
 # Import implementation-specific functions
 if _use_gpu():
-    from .gpu.gpu_utils import inplace_rot
+    from .gpu.utils import inplace_rot
path: src/fftvis/utils.py
user: login: steven-murray
body: I don't think this is the right way to do this -- this will set `inplace_rot` to either CPU or GPU for an entire session. This might not be what the user wants. We should instead use the correct `inplace_rot` whenever we need it, within a function.

created_at: 2025-05-12T13:16:36Z
author_association: COLLABORATOR
diff_hunk: @@ -171,6 +192,79 @@ def test_simulate(
         )
 
 
+@pytest.mark.parametrize("polarized", [False, True])
+@pytest.mark.parametrize("precision", [2, 1])
+@pytest.mark.parametrize("shear_array", [True, False])
+@pytest.mark.parametrize("rotate_array", [True, False])
+@pytest.mark.parametrize("remove_antennas", [True, False])
+@pytest.mark.parametrize("ants", [_hex_grid(), _square_grid()])
+@pytest.mark.parametrize("force_use_type3", [True, False])
+def test_simulate_gridded_type1_vs_type3(polarized, precision, shear_array, rotate_array, remove_antennas, ants, force_use_type3):
+    rng = np.random.default_rng(42)
+
+    params, *_ = get_standard_sim_params(
+        use_analytic_beam=True, polarized=polarized
+    )
+    params.pop("ants")
+
+    # Remove antennas randomly
+    if remove_antennas:
+        ants = {k: ants[k] for k in ants if rng.uniform(0, 1) > 0.25}
+        ants = {ki: ants[k] for ki, k in enumerate(ants)}
+
+    # Remove antennas from the array
path: tests/test_cpu_simulate.py
user: login: steven-murray
body: I think this comment is wrong
```suggestion
```
created_at: 2025-05-12T13:19:49Z
author_association: COLLABORATOR
diff_hunk: @@ -171,6 +192,79 @@ def test_simulate(
         )
 
 
+@pytest.mark.parametrize("polarized", [False, True])
+@pytest.mark.parametrize("precision", [2, 1])
+@pytest.mark.parametrize("shear_array", [True, False])
+@pytest.mark.parametrize("rotate_array", [True, False])
+@pytest.mark.parametrize("remove_antennas", [True, False])
+@pytest.mark.parametrize("ants", [_hex_grid(), _square_grid()])
+@pytest.mark.parametrize("force_use_type3", [True, False])
+def test_simulate_gridded_type1_vs_type3(polarized, precision, shear_array, rotate_array, remove_antennas, ants, force_use_type3):
+    rng = np.random.default_rng(42)
+
+    params, *_ = get_standard_sim_params(
+        use_analytic_beam=True, polarized=polarized
+    )
+    params.pop("ants")
+
+    # Remove antennas randomly
+    if remove_antennas:
+        ants = {k: ants[k] for k in ants if rng.uniform(0, 1) > 0.25}
+        ants = {ki: ants[k] for ki, k in enumerate(ants)}
+
+    # Remove antennas from the array
+    if rotate_array:
+        theta = np.pi / 2
+        rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1]])
+        # Rotate the array
+        ants = {
+            ant: np.dot(rotation_matrix, ants[ant])
+            for ai, ant in enumerate(ants)
+        }
+    if shear_array:
+        shear_matrix = np.array([[1, 0.5, 0], [0, 1, 0], [0, 0, 1]])
+        # Sheer the array
path: tests/test_cpu_simulate.py
user: login: steven-murray
body: ```suggestion
        # Shear the array
```
created_at: 2025-05-12T13:20:19Z
author_association: COLLABORATOR
diff_hunk: @@ -171,6 +192,79 @@ def test_simulate(
         )
 
 
+@pytest.mark.parametrize("polarized", [False, True])
+@pytest.mark.parametrize("precision", [2, 1])
+@pytest.mark.parametrize("shear_array", [True, False])
+@pytest.mark.parametrize("rotate_array", [True, False])
+@pytest.mark.parametrize("remove_antennas", [True, False])
+@pytest.mark.parametrize("ants", [_hex_grid(), _square_grid()])
+@pytest.mark.parametrize("force_use_type3", [True, False])
+def test_simulate_gridded_type1_vs_type3(polarized, precision, shear_array, rotate_array, remove_antennas, ants, force_use_type3):
+    rng = np.random.default_rng(42)
+
+    params, *_ = get_standard_sim_params(
+        use_analytic_beam=True, polarized=polarized
+    )
+    params.pop("ants")
+
+    # Remove antennas randomly
+    if remove_antennas:
+        ants = {k: ants[k] for k in ants if rng.uniform(0, 1) > 0.25}
+        ants = {ki: ants[k] for ki, k in enumerate(ants)}
+
+    # Remove antennas from the array
+    if rotate_array:
+        theta = np.pi / 2
+        rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1]])
+        # Rotate the array
+        ants = {
+            ant: np.dot(rotation_matrix, ants[ant])
+            for ai, ant in enumerate(ants)
+        }
+    if shear_array:
+        shear_matrix = np.array([[1, 0.5, 0], [0, 1, 0], [0, 0, 1]])
+        # Sheer the array
+        ants = {
+            ant: np.dot(shear_matrix, ants[ant])
+            for ai, ant in enumerate(ants)
+        }
+
+    # List of baselines
+    baselines = [(i, j) for i in ants for j in ants if j >= i]
+
+    # Use matvis as a reference
+    mvis = matvis.simulate_vis(
+        ants=ants,
+        precision=precision,
+        antpairs=np.array(baselines),
+        coord_method="CoordinateRotationERFA",
+        source_buffer=0.75,
+        **params,
+    )
+
+    # Get the beam
+    beam = params.pop("beams")[0]
+    times = params.pop("times").jd
+
+    # Use fftvis to simulate visibilities
+    fvis = simulate_vis(
+        ants=ants,
+        eps=1e-10 if precision == 2 else 6e-8,
+        precision=precision,
+        coord_method_params={"source_buffer": 0.75},
+        beam=beam,
+        times=times,
+        force_use_type3=force_use_type3,
+        baselines=baselines,
+        **params,
+    )
+
+    # Check shape of result when no baselines are specified
path: tests/test_cpu_simulate.py
user: login: steven-murray
body: ```suggestion
    # Check that results are the same
```
created_at: 2025-05-12T13:21:26Z
author_association: COLLABORATOR
diff_hunk: @@ -171,6 +192,79 @@ def test_simulate(
         )
 
 
+@pytest.mark.parametrize("polarized", [False, True])
+@pytest.mark.parametrize("precision", [2, 1])
+@pytest.mark.parametrize("shear_array", [True, False])
+@pytest.mark.parametrize("rotate_array", [True, False])
+@pytest.mark.parametrize("remove_antennas", [True, False])
+@pytest.mark.parametrize("ants", [_hex_grid(), _square_grid()])
+@pytest.mark.parametrize("force_use_type3", [True, False])
+def test_simulate_gridded_type1_vs_type3(polarized, precision, shear_array, rotate_array, remove_antennas, ants, force_use_type3):
+    rng = np.random.default_rng(42)
+
+    params, *_ = get_standard_sim_params(
+        use_analytic_beam=True, polarized=polarized
+    )
+    params.pop("ants")
+
+    # Remove antennas randomly
+    if remove_antennas:
+        ants = {k: ants[k] for k in ants if rng.uniform(0, 1) > 0.25}
+        ants = {ki: ants[k] for ki, k in enumerate(ants)}
+
+    # Remove antennas from the array
+    if rotate_array:
+        theta = np.pi / 2
+        rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1]])
+        # Rotate the array
+        ants = {
+            ant: np.dot(rotation_matrix, ants[ant])
+            for ai, ant in enumerate(ants)
+        }
+    if shear_array:
+        shear_matrix = np.array([[1, 0.5, 0], [0, 1, 0], [0, 0, 1]])
+        # Sheer the array
+        ants = {
+            ant: np.dot(shear_matrix, ants[ant])
+            for ai, ant in enumerate(ants)
+        }
+
+    # List of baselines
+    baselines = [(i, j) for i in ants for j in ants if j >= i]
+
+    # Use matvis as a reference
+    mvis = matvis.simulate_vis(
+        ants=ants,
+        precision=precision,
+        antpairs=np.array(baselines),
+        coord_method="CoordinateRotationERFA",
+        source_buffer=0.75,
+        **params,
+    )
path: tests/test_cpu_simulate.py
user: login: steven-murray
body: Why don't we test type1 vs type3 instead of testing both against matvis?
created_at: 2025-05-12T13:22:04Z
author_association: COLLABORATOR
diff_hunk: @@ -193,6 +204,11 @@ def _evaluate_vis_chunk(
             Number of threads to use.
         is_coplanar : bool
             Whether the array is coplanar.
+        basis_matrix : np.ndarray
path: src/fftvis/core/simulate.py
user: login: tyler-a-cox
body: Added more description here.
created_at: 2025-05-12T18:00:23Z
author_association: OWNER
diff_hunk: @@ -37,6 +37,7 @@ def simulate(
         precision: int = 2,
         polarized: bool = False,
         eps: float = None,
+        upsampfac: int = 2,
path: src/fftvis/core/simulate.py
user: login: tyler-a-cox
body: No worries! I was just sticking to the naming convention used in `finufft` I've changed this to `upsample_factor` for more clarity.
created_at: 2025-05-12T18:01:00Z
author_association: OWNER
diff_hunk: @@ -0,0 +1,194 @@
+import numpy as np
+from math import lcm
+from fractions import Fraction
+from typing import Any, Dict, Tuple
+
+
+def find_integer_multiplier(
+    arr: np.ndarray,
+    max_denominator: int = 10**6
+) -> int:
+    """
+    Return the smallest positive integer f such that f * arr[i] is integral
+    for all entries (up to rational approximation). Zeros in arr are ignored.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    max_denominator : int
+        Maximum denominator for the fractions.
+
+    Returns
+    -------
+    f : int
+        The integer factor used for scaling.
+    """
+    fracs = [
+        Fraction(val).limit_denominator(max_denominator)
+        for val in np.ravel(arr)
+        if val != 0
+    ]
+    if not fracs:
+        return 1
+    dens = (frac.denominator for frac in fracs)
+    return lcm(*dens)
+
+
+def can_scale_to_int(
+    arr: np.ndarray,
+    tol: float = 1e-9,
+    max_denominator: int = 10**6,
+    max_factor: int = None
+) -> Tuple[bool, int]:
+    """
+    Check if there exists an integer factor f such that f * arr is
+    approximately integers. Returns (is_griddable, factor).
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of values to check.
+    tol : float
+        Tolerance for checking if values are close to integers.
+    max_denominator : int
+        Maximum denominator for the fractions.
+    max_factor : int
+        Maximum allowed factor for scaling.
+
+    Returns
+    -------
+    is_griddable : bool
+        True if the array can be scaled to integers, False otherwise.
+    factor : int
+        The integer factor used for scaling.
+    """
+    f = find_integer_multiplier(arr, max_denominator)
+    if max_factor is not None and f > max_factor:
+        return False, f
+    scaled = f * np.array(arr, dtype=float)
+    if np.allclose(scaled, np.round(scaled), atol=tol):
+        return True, f
+    return False, f
+
+def find_lattice_basis(
+    antpos: Dict[Any, np.ndarray],
+    tol: float = 1e-9,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Find the lattice basis for the antenna positions.
+
+    Parameters
+    ----------
+    antpos : dict
+        Dictionary of antenna positions in the form {ant_index: np.array([x,y,z])}.
+    tol : float
+        Tolerance for checking if values are close to integers.
path: src/fftvis/core/antenna_gridding.py
user: login: tyler-a-cox
body: Changed this.
created_at: 2025-05-12T18:01:15Z
author_association: OWNER
diff_hunk: @@ -193,6 +204,11 @@ def _evaluate_vis_chunk(
             Number of threads to use.
         is_coplanar : bool
             Whether the array is coplanar.
+        basis_matrix : np.ndarray
+            Lattice basis matrix used to grid baselines and sources for a type 1
+            non-uniform FFT.
+        type1_n_modes : int
+            Number of modes for type 1 non-uniform FFT.
path: src/fftvis/core/simulate.py
user: login: tyler-a-cox
body: Added more description here to let users know that this will be the size of the output transform before indexing.
created_at: 2025-05-12T18:02:05Z
author_association: OWNER
diff_hunk: @@ -201,16 +203,12 @@ def simulate(
             n_modes = 2 * int(np.round(np.max(np.abs(bls)))) + 1
 
             # Get the maximum baseline length for proper coordinate scaling
-            max_baseline_coord_length = np.max(np.abs(
-               [
-                   ants[bl[1]] - ants[bl[0]]
-                   for bl in baselines
-               ]
-            ))
-            type1_coord_scalar = max_baseline_coord_length / utils.speed_of_light / (n_modes // 2)
+            basis_matrix *= 1 / utils.speed_of_light
+            basis_matrix = basis_matrix.astype(real_dtype)
 
             # Assume the array is coplanar for gridded coordinates
             is_coplanar = True
+            rotation_matrix = np.eye(3, dtype=real_dtype)
path: src/fftvis/cpu/cpu_simulate.py
user: login: tyler-a-cox
body: No, this is just a placeholder for the rotation matrix that handles rotation of the array that has a slope along the xy and z-axes. I could also just change this so that the rotation matrix is an optional parameter, but we were already storing it as a shared array in `ray` so I just decided to keep it.
created_at: 2025-05-12T18:04:28Z
author_association: OWNER
diff_hunk: @@ -437,7 +437,14 @@ def _evaluate_vis_chunk(
                 )
 
                 # Rotate source coordinates with rotation matrix.
-                inplace_rot(rotation_matrix, topo)
+                if not np.allclose(rotation_matrix, np.eye(3)):
path: src/fftvis/cpu/cpu_simulate.py
user: login: tyler-a-cox
body: Added an extra parameter which performs this check once outside the loop.
created_at: 2025-05-12T18:04:55Z
author_association: OWNER
diff_hunk: @@ -171,6 +192,79 @@ def test_simulate(
         )
 
 
+@pytest.mark.parametrize("polarized", [False, True])
+@pytest.mark.parametrize("precision", [2, 1])
+@pytest.mark.parametrize("shear_array", [True, False])
+@pytest.mark.parametrize("rotate_array", [True, False])
+@pytest.mark.parametrize("remove_antennas", [True, False])
+@pytest.mark.parametrize("ants", [_hex_grid(), _square_grid()])
+@pytest.mark.parametrize("force_use_type3", [True, False])
+def test_simulate_gridded_type1_vs_type3(polarized, precision, shear_array, rotate_array, remove_antennas, ants, force_use_type3):
+    rng = np.random.default_rng(42)
+
+    params, *_ = get_standard_sim_params(
+        use_analytic_beam=True, polarized=polarized
+    )
+    params.pop("ants")
+
+    # Remove antennas randomly
+    if remove_antennas:
+        ants = {k: ants[k] for k in ants if rng.uniform(0, 1) > 0.25}
+        ants = {ki: ants[k] for ki, k in enumerate(ants)}
+
+    # Remove antennas from the array
+    if rotate_array:
+        theta = np.pi / 2
+        rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1]])
+        # Rotate the array
+        ants = {
+            ant: np.dot(rotation_matrix, ants[ant])
+            for ai, ant in enumerate(ants)
+        }
+    if shear_array:
+        shear_matrix = np.array([[1, 0.5, 0], [0, 1, 0], [0, 0, 1]])
+        # Sheer the array
+        ants = {
+            ant: np.dot(shear_matrix, ants[ant])
+            for ai, ant in enumerate(ants)
+        }
+
+    # List of baselines
+    baselines = [(i, j) for i in ants for j in ants if j >= i]
+
+    # Use matvis as a reference
+    mvis = matvis.simulate_vis(
+        ants=ants,
+        precision=precision,
+        antpairs=np.array(baselines),
+        coord_method="CoordinateRotationERFA",
+        source_buffer=0.75,
+        **params,
+    )
path: tests/test_cpu_simulate.py
user: login: tyler-a-cox
body: Now testing type 1/3 directly.
created_at: 2025-05-12T18:07:43Z
author_association: OWNER
diff_hunk: @@ -32,9 +31,9 @@ def _use_gpu():
 
 # Import implementation-specific functions
 if _use_gpu():
-    from .gpu.gpu_utils import inplace_rot
+    from .gpu.utils import inplace_rot
path: src/fftvis/utils.py
user: login: tyler-a-cox
body: I'm now just getting the CPU or GPU function directly rather than doing this.
created_at: 2025-05-12T18:15:41Z
author_association: OWNER

PR comments:
user: login: codecov[bot]
created_at: 2025-05-08T22:01:22Z
body: ## [Codecov](https://app.codecov.io/gh/tyler-a-cox/fftvis/pull/45?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Tyler+Cox) Report
Attention: Patch coverage is `96.03960%` with `4 lines` in your changes missing coverage. Please review.
> Project coverage is 97.30%. Comparing base [(`3269f91`)](https://app.codecov.io/gh/tyler-a-cox/fftvis/commit/3269f91b6ecde629f6a804d11220053a554846df?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Tyler+Cox) to head [(`253491b`)](https://app.codecov.io/gh/tyler-a-cox/fftvis/commit/253491bfed1caca50dab0b5a7c52fae59406a573?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Tyler+Cox).
> Report is 2 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/tyler-a-cox/fftvis/pull/45?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Tyler+Cox) | Patch % | Lines |
|---|---|---|
| [src/fftvis/core/antenna\_gridding.py](https://app.codecov.io/gh/tyler-a-cox/fftvis/pull/45?src=pr&el=tree&filepath=src%2Ffftvis%2Fcore%2Fantenna_gridding.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Tyler+Cox#diff-c3JjL2ZmdHZpcy9jb3JlL2FudGVubmFfZ3JpZGRpbmcucHk=) | 91.66% | [2 Missing and 2 partials :warning: ](https://app.codecov.io/gh/tyler-a-cox/fftvis/pull/45?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Tyler+Cox) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main      #45      +/-   ##
==========================================
+ Coverage   97.28%   97.30%   +0.01%     
==========================================
  Files          17       18       +1     
  Lines         369      445      +76     
  Branches       43       56      +13     
==========================================
+ Hits          359      433      +74     
- Misses          6        7       +1     
- Partials        4        5       +1     
```

| [Flag](https://app.codecov.io/gh/tyler-a-cox/fftvis/pull/45/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Tyler+Cox) | Coverage Δ | |
|---|---|---|
| [unittests](https://app.codecov.io/gh/tyler-a-cox/fftvis/pull/45/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Tyler+Cox) | `97.30% <96.03%> (+0.01%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Tyler+Cox#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/tyler-a-cox/fftvis/pull/45?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Tyler+Cox).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Tyler+Cox).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: tyler-a-cox
created_at: 2025-05-10T07:09:06Z
body: > Hey Tyler, this is super cool, thanks! I have a few broader questions/comments:
> 
> 1. I would move the new `cpu_utils.py` stuff into its own module (like `antenna_gridding.py` or something nicer)
> 2. I don't know why we have `cpu/cpu_nufft.py` etc instead of just `cpu/nufft.py` which is much nicer.
> 3. How does the antenna rotation work? Rigid-Rotating a hex array should not be able to shear the antennas, right?
> 4. I would like to see a test comparing a gridded to non-gridded calculation for the same ant positions and sky model. This would mean we also need an option to tell the simulator to NOT try gridding even if it could grid.
> 5. I also think that we should let the user tell the simulate function that yes, indeed, I want to use gridding, and I know just how to grid it. Probably by passing a rotation matrix and argument like "griddable = True" (or maybe just the rotation matrix). That way we can skip the checks and just assume the user knows what they're doing. Helps with speed a little bit.

Hey Steven, thanks for taking a look!

1. All the gridding code has now been moved to the `antenna_gridding.py` module.
2. I've renamed these modules to `nufft.py` rather than `cpu_nufft.py`
3. Yeah, you're right, I think the "rotation" name in the code is a bit misleading. In practice, the code is much more general. As long as you can define some basis that can map the physical antenna positions onto integer grid points, then you should be able to use the Type 1 in place of the Type 3. It only really requires that you can apply the inverse of that matrix to the source coordinates and use the antenna positions projected in their new coordinates as indexing arrays. It also requires being a little careful about getting the scaling of the coordinates right, but it does seem to work pretty well.
4. Yeah, I do think it's a good idea to let users have the option to force `fftvis` to use the type 3 ever when the type 1 can be used. I've added a new keyword argument `force_use_type3` which requires `fftvis` to use the type 3. I've also added a parameterized test function that compares the Type 1 and Type 3 to matvis for a hex array and a rectangularly gridded array. I've also applied shears, rotations, and removed antennas to test the gridding function and make sure that the gridding is being done sensibly.
5. I decided to go a slightly different route on this point and just make the `check_antpos_griddability` better at detecting a range of gridded arrays. The code now looks for the two shortest, non-collinear baselines in the array to use as basis functions that are used for gridding.
author_association: OWNER
user: login: tyler-a-cox
created_at: 2025-05-10T07:17:16Z
body: I've also caught a potential issue using `precision=1`. I noticed that the Type 1 actually gave results that more closely match the Type 3 when using single precision. After digging into the `finufft` code, it looks like the Type 3 uses an upsampling factor of 1.25 for single precision, while the Type 1 uses an upsampling factor of 2.0. I decided to add another parameter that sets the upsampling factor (default=2.0) to maintain accuracy at single precision.
author_association: OWNER