Merge branch 'develop'

Author: Joshuaalbert

dsa2000_cal/benchmarking/actual_calibration/main.py

@@ -59,7 +59,7 @@
     build_gaussian_source_model, \
     BaseGaussianSourceModel
 
-from dsa2000_fm.forward_models.streaming.average_utils import average_rule
+from dsa2000_fm.actors.average_utils import average_rule
 
 tfpd = tfp.distributions
 

dsa2000_cal/dashboards/dsa/pages/view_calibration_solutions.py

@@ -3,7 +3,7 @@
 import ray
 import streamlit as st
 
-from dsa2000_fm.forward_models.streaming.calibration_solution_cache import CalibrationSolutionCache, \
+from dsa2000_fm.actors.calibration_solution_cache import CalibrationSolutionCache, \
     CalibrationSolution
 
 

dsa2000_cal/notebooks/explore_calibration.ipynb

@@ -21,7 +21,7 @@
     "\n",
     "from dsa2000_common.common.noise import calc_baseline_noise\n",
     "from dsa2000_common.common.quantity_utils import time_to_jnp, quantity_to_jnp\n",
-    "from dsa2000_fm.forward_models.streaming.calibrator import Calibration\n",
+    "from dsa2000_fm.actors.calibrator import Calibration\n",
     "from dsa2000_fm.forward_models.utils import ObservationSetup\n",
     "from dsa2000_common.visibility_model.source_models.celestial.base_point_source_model import build_point_source_model\n",
     "\n",

dsa2000_cal/notebooks/explore_sky_loss.ipynb

@@ -42,8 +42,8 @@
    "source": [
     "\n",
     "\n",
-    "import glob\n",
     "import os\n",
+    "import queue\n",
     "import time\n",
     "from typing import NamedTuple, Dict\n",
     "\n",
@@ -318,6 +318,9 @@
     "\n",
     "\n",
     "def main(\n",
+    "        cpu,\n",
+    "        gpu,\n",
+    "        result_num: int,\n",
     "        seed: int,\n",
     "        save_folder: str,\n",
     "        array_name: str,\n",
@@ -338,9 +341,6 @@
     "):\n",
     "    plt.close('all')\n",
     "    t0 = time.time()\n",
-    "    result_num = len(glob.glob(os.path.join(save_folder, 'result_*.json')))\n",
-    "    cpu = jax.devices(\"cpu\")[0]\n",
-    "    gpu = jax.devices(\"cuda\")[0]  # or \"gpu\" depending on platform\n",
     "\n",
     "    key = jax.random.PRNGKey(seed)\n",
     "    fill_registries()\n",
@@ -702,36 +702,57 @@
     "    result_values: Dict[str, float]\n",
     "\n",
     "\n",
+    "gpu = jax.devices(\"cuda\")[0]  # or \"gpu\" depending on platform\n",
+    "\n",
+    "\n",
+    "def run_varying_systematics(result_idx, cpu, gpu, pointing_offset_stddev, axial_focus_error_stddev,\n",
+    "                            horizon_peak_astigmatism_stddev, with_smearing):\n",
+    "    main(\n",
+    "        cpu=cpu,\n",
+    "        gpu=gpu,\n",
+    "        result_num=result_idx,\n",
+    "        seed=0,\n",
+    "        save_folder='sky_loss_11Mar2025_varying_systematics_more_stats',\n",
+    "        array_name='dsa2000_optimal_v1',\n",
+    "        pointing=ac.ICRS(0 * au.deg, 0 * au.deg),\n",
+    "        num_measure_points=256,\n",
+    "        angular_radius=1.75 * au.deg,\n",
+    "        prior_psf_sidelobe_peak=1e-3,\n",
+    "        bright_source_id='nvss_calibrators',\n",
+    "        pointing_offset_stddev=pointing_offset_stddev,\n",
+    "        axial_focus_error_stddev=axial_focus_error_stddev,\n",
+    "        horizon_peak_astigmatism_stddev=horizon_peak_astigmatism_stddev,\n",
+    "        turbulent=True,\n",
+    "        dawn=True,\n",
+    "        high_sun_spot=True,\n",
+    "        with_ionosphere=True,\n",
+    "        with_dish_effects=True,\n",
+    "        with_smearing=with_smearing\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "cpus = jax.devices(\"cpu\")\n",
+    "gpus = jax.devices(\"cuda\")\n",
+    "queues = [queue.Queue() for _ in gpus]\n",
+    "\n",
+    "# fill queues with input args\n",
+    "result_idx = 0\n",
     "for pointing_offset_stddev in [0, 1, 2, 4] * au.arcmin:\n",
     "    for axial_focus_error_stddev in [0, 3, 5] * au.mm:\n",
     "        for horizon_peak_astigmatism_stddev in [0, 1, 2, 4] * au.mm:\n",
     "            for with_smearing in [True, False]:\n",
-    "                main(\n",
-    "                    seed=0,\n",
-    "                    save_folder='sky_loss_11Mar2025_varying_systematics_more_stats',\n",
-    "                    array_name='dsa2000_optimal_v1',\n",
-    "                    pointing=ac.ICRS(0 * au.deg, 0 * au.deg),\n",
-    "                    num_measure_points=256,\n",
-    "                    angular_radius=1.75 * au.deg,\n",
-    "                    prior_psf_sidelobe_peak=1e-3,\n",
-    "                    bright_source_id='nvss_calibrators',\n",
-    "                    pointing_offset_stddev=pointing_offset_stddev,\n",
-    "                    axial_focus_error_stddev=axial_focus_error_stddev,\n",
-    "                    horizon_peak_astigmatism_stddev=horizon_peak_astigmatism_stddev,\n",
-    "                    turbulent=True,\n",
-    "                    dawn=True,\n",
-    "                    high_sun_spot=True,\n",
-    "                    with_ionosphere=True,\n",
-    "                    with_dish_effects=True,\n",
-    "                    with_smearing=with_smearing\n",
-    "                )\n",
+    "                queue = queues[result_idx % len(gpus)]\n",
+    "                gpu = gpus[result_idx % len(gpus)]\n",
+    "                cpu = cpus[result_idx % len(cpus)]\n",
+    "                queue.put((run_varying_systematics, result_idx, cpu, gpu, pointing_offset_stddev,\n",
+    "                           axial_focus_error_stddev, horizon_peak_astigmatism_stddev, with_smearing))\n",
     "\n",
-    "fill_registries()\n",
-    "survey_pointings = misc_registry.get_instance(misc_registry.get_match('survey_pointings'))\n",
-    "pointings = survey_pointings.survey_pointings_v1()\n",
-    "for pointing in tqdm(pointings):\n",
-    "    print(pointing)\n",
+    "\n",
+    "def run_survey(result_idx, cpu, gpu):\n",
     "    main(\n",
+    "        cpu=cpu,\n",
+    "        gpu=gpu,\n",
+    "        result_num=result_idx,\n",
     "        seed=0,\n",
     "        save_folder='sky_loss_11Mar2025_full_survey_more_stats',\n",
     "        array_name='dsa2000_optimal_v1',\n",
@@ -752,6 +773,39 @@
     "    )\n",
     "\n",
     "\n",
+    "fill_registries()\n",
+    "survey_pointings = misc_registry.get_instance(misc_registry.get_match('survey_pointings'))\n",
+    "pointings = survey_pointings.survey_pointings_v1()\n",
+    "\n",
+    "# fill queues with input args\n",
+    "result_idx = 0\n",
+    "for pointing in tqdm(pointings):\n",
+    "    queue = queues[result_idx % len(gpus)]\n",
+    "    gpu = gpus[result_idx % len(gpus)]\n",
+    "    cpu = cpus[result_idx % len(cpus)]\n",
+    "    queue.put((run_survey, result_idx, cpu, gpu))\n",
+    "    result_idx += 1\n",
+    "\n",
+    "\n",
+    "# now run the jobs in thread pool\n",
+    "def worker(queue):\n",
+    "    while True:\n",
+    "        args = queue.get()\n",
+    "        if args is None:\n",
+    "            break\n",
+    "        f = args[0]\n",
+    "        args = args[1:]\n",
+    "        f(*args)\n",
+    "\n",
+    "\n",
+    "# now run the jobs in thread pool, each job processes a queue\n",
+    "import concurrent.futures\n",
+    "\n",
+    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
+    "    for queue in queues:\n",
+    "        executor.submit(worker, queue)\n",
+    "\n",
+    "\n",
     "\n",
     "\n"
    ],

dsa2000_cal/notebooks/performance_test_calibration.ipynb

@@ -72,7 +72,7 @@
     "from dsa2000_common.visibility_model.source_models.celestial.base_point_source_model import build_point_source_model, \\\n",
     "    BasePointSourceModel\n",
     "\n",
-    "from dsa2000_fm.forward_models.streaming.average_utils import average_rule\n",
+    "from dsa2000_fm.actors.average_utils import average_rule\n",
     "\n",
     "tfpd = tfp.distributions\n",
     "\n",

dsa2000_cal/scripts/dd_calibrate_from_casa_ms/main.py

@@ -1,6 +1,6 @@
 import os
 
-from dsa2000_fm.forward_models.streaming.calibrator import Calibration
+from dsa2000_fm.actors.calibrator import Calibration
 
 os.environ["XLA_FLAGS"] = f"--xla_force_host_platform_device_count={os.cpu_count()}"