Merge branch 'develop'

Author: Joshuaalbert

dsa2000_cal/README.md

@@ -61,10 +61,10 @@ cd DSA2000-Cal
 
 ```bash
 conda activate dsa2000_cal_py
-pip install dsa2000_call
+pip install dsa2000_cal
 ```
 
-8. Set up PyCharm for development
+8. Set up PyCharm for development (optional but recommended).
 
     1. Make sure you have created a `dsa2000_cal_py` conda env as above, and installed requirements files.
     2. Create a new project in PyCharm in the repo root directory `/home/username/git/DSA2000-Cal`. Use an empty project

dsa2000_cal/benchmarking/DFT_predict/main.py

@@ -8,7 +8,7 @@
 import astropy.units as au
 import jax
 import numpy as np
-from tomographic_kernel.frames import ENU
+from dsa2000_common.common.enu_frame import ENU
 
 from dsa2000_common.common.quantity_utils import quantity_to_jnp, time_to_jnp
 from dsa2000_common.delay_models.base_far_field_delay_engine import build_far_field_delay_engine

…000_cal/benchmarking/calibration/main.py …/benchmarking/actual_calibration/main.pydsa2000_cal/benchmarking/calibration/main.py renamed to dsa2000_cal/benchmarking/actual_calibration/main.py dsa2000_cal/benchmarking/calibration/main.py renamed to dsa2000_cal/benchmarking/actual_calibration/main.py

@@ -1,5 +1,6 @@
 import os
 
+from dsa2000_common.common.ray_utils import TimerLog
 from dsa2000_common.visibility_model.source_models.celestial.base_point_source_model import BasePointSourceModel
 
 os.environ['JAX_PLATFORMS'] = 'cuda'
@@ -12,9 +13,7 @@
 
 # os.environ["XLA_FLAGS"] = f"--xla_force_host_platform_device_count={os.cpu_count()}"
 
-import dataclasses
 import os
-import time
 from functools import partial
 from typing import Generator, Tuple, List
 
@@ -31,7 +30,7 @@
 import tensorflow_probability.substrates.jax as tfp
 from astropy import units as au, coordinates as ac
 from matplotlib import pyplot as plt
-from tomographic_kernel.frames import ENU
+from dsa2000_common.common.enu_frame import ENU
 
 from dsa2000_assets.content_registry import fill_registries
 from dsa2000_assets.registries import array_registry
@@ -65,21 +64,6 @@
 tfpd = tfp.distributions
 
 
-@dataclasses.dataclass
-class TimerLog:
-    msg: str
-
-    def __post_init__(self):
-        self.t0 = time.time()
-
-    def __enter__(self):
-        print(f"{self.msg}")
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        print(f"... took {time.time() - self.t0:.3f} seconds")
-        return False
-
-
 def build_mock_obs_setup(array_name: str, num_sol_ints_time: int, frac_aperture: float = 1.):
     fill_registries()
     array = array_registry.get_instance(array_registry.get_match(array_name))

dsa2000_cal/benchmarking/calibration/run_benchmark.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+conda create -n cal_benchmark python=3.11
+conda activate cal_benchmark
+
+pip install jax[cuda12] jaxlib 'numpy<2' nvtx
+
+python standalone_lm_multi_step.py
+
+conda deactivate
+conda remove -n cal_benchmark --all