Incorporate multiprocessing into fftvis

.coveragerc

@@ -1,7 +1,14 @@
 [run]
 include = fftvis/*
-omit = fftvis/logutils.py
+omit = 
+    fftvis/logutils.py
+    fftvis/cli.py
 branch = True
 
 [report]
-omit = */tests/*
+omit = */tests/*
+exclude_lines = 
+    if trace_mem:
+    pragma: no cover
+    if isinstance(beam, UVBeam):
+    if _nbig != nprocesses:

.gitignore

@@ -158,4 +158,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
-.vscode
+.vscode
+*.flame
+cprofile*

fftvis/beams.py

@@ -1,10 +1,9 @@
 import numpy as np
-from pyuvdata import UVBeam
-
+from pyuvdata.beam_interface import BeamInterface
+import numba as nb
 
 def _evaluate_beam(
-    A_s: np.ndarray,
-    beam: UVBeam,
+    beam: BeamInterface,
     az: np.ndarray,
     za: np.ndarray,
     polarized: bool,
@@ -44,39 +43,61 @@ def _evaluate_beam(
         for interpolation orders greater than linear. 
     """
     # Primary beam pattern using direct interpolation of UVBeam object
-    kw = (
-        {
-            "reuse_spline": True,
-            "check_azza_domain": False,
-            "spline_opts": spline_opts,
-            "interpolation_function": interpolation_function,
-        }
-        if isinstance(beam, UVBeam)
-        else {}
-    )
-    if isinstance(beam, UVBeam) and not beam.future_array_shapes:
-        beam.use_future_array_shapes()
+    kw = {
+        "reuse_spline": True,
+        "check_azza_domain": False,
+        "spline_opts": spline_opts,
+        "interpolation_function": interpolation_function,
+    }
 
-    interp_beam = beam.interp(
+    interp_beam = beam.compute_response(
         az_array=az,
         za_array=za,
         freq_array=np.atleast_1d(freq),
         **kw,
-    )[0]
+    )
 
     if polarized:
         interp_beam = interp_beam[:, :, 0, :]
     else:
         # Here we have already asserted that the beam is a power beam and
         # has only one polarization, so we just evaluate that one.
-        interp_beam = np.sqrt(interp_beam[0, 0, 0, :])
-
-    A_s[:, :] = interp_beam
+        interp_beam = interp_beam[0, 0, 0, :]
 
     # Check for invalid beam values
     if check:
-        sm = np.sum(A_s)
+        sm = np.sum(interp_beam)
         if np.isinf(sm) or np.isnan(sm):
             raise ValueError("Beam interpolation resulted in an invalid value")
 
-    return A_s
+    return interp_beam
+
+@nb.jit(
+    nopython=True,
+    parallel=False,
+    nogil=False
+)
+def get_apparent_flux_polarized(beam: np.ndarray, flux: np.ndarray):  # pragma: no cover
+    """Calculate apparent flux of the sources."""
+    nax, nfd, nsrc = beam.shape
+
+    for isrc in range(nsrc):
+        c = np.conj(beam[:, :, isrc])
+
+        i00 = (
+            c[0, 0] * beam[0, 0, isrc] + 
+            c[1, 0] * beam[1, 0, isrc] 
+        )
+        i01 = (
+            c[0, 0] * beam[0, 1, isrc] + 
+            c[1, 0] * beam[1, 1, isrc] 
+        )
+
+        i11 = (
+            c[0, 1] * beam[0, 1, isrc] + 
+            c[1, 1] * beam[1, 1, isrc] 
+        )
+        beam[0, 0, isrc] = i00 * flux[isrc]
+        beam[0, 1, isrc] = i01 * flux[isrc]
+        beam[1, 0, isrc] = np.conj(i01) * flux[isrc]
+        beam[1, 1, isrc] = i11 * flux[isrc]

fftvis/cli.py

@@ -0,0 +1,147 @@
+import typer
+from matvis.cli import get_standard_sim_params, get_line_based_stats, get_summary_stats, get_label, get_redundancies
+from .simulate import simulate_vis, _evaluate_vis_chunk
+from pathlib import Path
+from typing import Optional
+import logging
+from rich.console import Console
+from rich.rule import Rule
+from line_profiler import LineProfiler
+import time
+from pyuvdata.telescopes import get_telescope
+import cProfile
+import pstats
+import os
+import numpy as np
+from hera_sim.antpos import hex_array
+
+cns = Console()
+
+logger = logging.getLogger('fftvis')
+profiler = LineProfiler()
+
+app = typer.Typer()
+
+@app.command()
+def run_profile(
+    analytic_beam: bool = True,
+    nfreq: int = 1,
+    ntimes: int = 1,
+    nants: int = 1,
+    nsource: int = 1,
+    double_precision: bool = True,
+    outdir: Path = Path(".").absolute(),
+    verbose: bool = True,
+    log_level: str = "INFO",
+    coord_method: str = "CoordinateRotationERFA",
+    naz: int = 360,
+    nza: int = 180,
+    nprocesses: int = 1,
+    update_bcrs_every: float = np.inf,
+    hera: int = 0,
+    nside: int = 0,
+    force_use_ray: bool = False,
+    trace_mem: bool = False,
+    beam_spline_order: int = 3,
+    freq_min: float = 100,  # MHz
+):
+    """Run the script."""
+    logger.setLevel(log_level.upper())
+
+    if nside > 0:
+        nsource = 12*nside**2
+
+    (
+        ants,
+        flux,
+        ra,
+        dec,
+        freqs,
+        times,
+        cpu_beams,
+        beam_idx,
+    ) = get_standard_sim_params(
+        analytic_beam, nfreq, ntimes, nants, nsource, nbeams=1, naz=naz, nza=nza,
+        freq_min=freq_min*1e6,
+    )
+    if hera > 0:
+        ants = hex_array(hera)
+
+    ants = {k: list(v) for k, v in ants.items()}
+    nants = len(ants)
+
+    cns.print(Rule("Running fftvis profile"))
+    cns.print(f"  NANTS:            {nants:>7}")
+    cns.print(f"  NTIMES:           {ntimes:>7}")
+    cns.print(f"  NFREQ:            {nfreq:>7}")
+    cns.print(f"  NSOURCE:          {nsource:>7}")
+    cns.print(f"  DOUBLE-PRECISION: {double_precision:>7}")
+    cns.print(f"  ANALYTIC-BEAM:    {analytic_beam:>7}")
+    cns.print(f"  COORDROT METHOD:  {coord_method:>7}")
+    #cns.print(f"  NPAIRS:           {len(pairs) if pairs is not None else nants**2:>7}")
+    cns.print(f"  NAZ:              {naz:>7}")
+    cns.print(f"  NZA:              {nza:>7}")
+    cns.print(f"  NPROCESSES:       {nprocesses:>7}")
+    cns.print(f"  INTERP ORDER:     {beam_spline_order:>7}")
+
+    if coord_method == "CoordinateRotationERFA":
+        cns.print(f"  BCRS UPDATE:       {update_bcrs_every:>7}")
+        coord_params = {"update_bcrs_every": update_bcrs_every}
+    else:
+        coord_params = {}
+
+    cns.print(Rule())
+
+    profiler.add_function(simulate_vis)
+    profiler.add_function(_evaluate_vis_chunk)
+
+    init_time = time.time()
+
+    str_id = get_label(
+        analytic_beam=analytic_beam,
+        nfreq=nfreq,
+        ntimes=ntimes,
+        nants=nants,
+        nsource=nsource,
+        nbeams=1,
+        matprod_method='',
+        gpu=False,
+        double_precision=double_precision,
+        coord_method=coord_method,
+        naz=naz,
+        nza=nza,
+    )
+
+    cProfile.runctx(
+        """simulate_vis(
+    ants=ants,
+    fluxes=flux,
+    ra=ra,
+    dec=dec,
+    freqs=freqs,
+    times=times.jd,
+    beam=cpu_beams[0],
+    polarized=True,
+    precision=2 if double_precision else 1,
+    telescope_loc=get_telescope("hera").location,
+    nprocesses=nprocesses,
+    coord_method_params=coord_params,
+    force_use_ray=force_use_ray,
+    trace_mem=trace_mem,
+    beam_spline_opts={'order': beam_spline_order},
+        )""", 
+        globals(),
+        locals(), 
+        str_id
+    )
+
+    out_time = time.time()
+    cns.print("TOTAL TIME: ", out_time - init_time)
+
+    p = pstats.Stats(str_id)
+    p.sort_stats("cumulative").print_stats(50)
+    os.system(f"flameprof --format=log {str_id} > {str_id}.flame")
+
+
+if __name__=="__main__":
+    app()

fftvis/logutils.py

@@ -5,6 +5,7 @@
 import logging
 import time
 import tracemalloc as tm
+import psutil
 
 
 logger = logging.getLogger(__name__)
@@ -26,13 +27,31 @@ def human_readable_size(size, decimal_places=2, indicate_sign=False):
     else:
         return f"{size:.{decimal_places}f} {unit}"
 
-
+def printmem(pr: psutil.Process, msg: str=""):
+    """Print memory usage of the process."""
+    #if logger.isEnabledFor(logging.INFO):
+    info = pr.memory_info()
+
+    rss = info.rss
+    try:    
+        shm = info.shared
+    except AttributeError:
+        # On macos there is no shared memory
+        shm = 0
+    used = rss - shm
+
+    shm = human_readable_size(shm)
+    used = human_readable_size(used)
+
+    #logger.info(f"{msg} Memory Usage [{pr.pid}]: {used} internal, {shm} shared.")
+    print(f"{msg} Memory Usage [{pr.pid}]: {used} internal, {shm} shared.")
+
 def memtrace(highest_peak) -> int:
     if logger.isEnabledFor(logging.INFO):
         cm, pm = tm.get_traced_memory()
         logger.info(f"Starting Memory usage  : {cm / 1024**3:.3f} GB")
         logger.info(f"Starting Peak Mem usage: {pm / 1024**3:.3f} GB")
-        logger.info(f"Traemalloc Peak Memory (tot)(GB): {highest_peak / 1024**3:.2f}")
+        logger.info(f"Tracemalloc Peak Memory (tot)(GB): {highest_peak / 1024**3:.2f}")
         tm.reset_peak()
         return max(pm, highest_peak)
 
@@ -48,9 +67,10 @@ def log_progress(start_time, prev_time, iters, niters, pr, last_mem):
     per_iter = total / iters
     expected = per_iter * niters
 
-    rss = pr.memory_info().rss
-    mem = human_readable_size(rss)
-    memdiff = human_readable_size(rss - last_mem, indicate_sign=True)
+    info = pr.memory_info()
+    used = info.rss - info.shared
+    mem = human_readable_size(used)
+    memdiff = human_readable_size(used - last_mem, indicate_sign=True)
 
     logger.info(
         f"""