Merge remote-tracking branch 'origin/main' into feature/joss-paper

Author: tristan-deep

.github/workflows/tests.yaml

@@ -158,7 +158,7 @@ jobs:
       - name: Set up Git
         run: git config --global safe.directory $GITHUB_WORKSPACE
 
-      - name: Build documentation with Sphinx (fail on warnings)
+      - name: Build documentation with Sphinx (fails on warnings)
         run: |
           docker run --rm \
             -v "${{ github.workspace }}:/zea" \

docs/source/environment.rst

@@ -31,6 +31,10 @@ Here are the environment variables that ``zea`` uses at runtime. Arguably the mo
      - Timeout in seconds for calling ``nvidia-smi`` to get GPU information during :func:`zea.init_device`.
      - ``30``
      - Any positive integer, or ``<= 0`` to disable timeout.
+   * - ``ZEA_DOWNLOAD_TIMEOUT``
+     - Timeout in seconds for downloading files, e.g. during dataset conversion.
+     - ``60``
+     - Any positive integer, or ``<= 0`` to disable timeout.
    * - ``ZEA_FIND_H5_SHAPES_PARALLEL``
      - If set to ``1``, will use parallel processing when searching for HDF5 file shapes.
      - ``1``

docs/source/getting-started.rst

@@ -16,9 +16,7 @@ Let's take a quick look at how to use ``zea`` to load and process ultrasound dat
    zea.init_device()
 
    # loading a config file from Hugging Face, but can also load a local config file
-   config = zea.Config.from_hf(
-      "zeahub/configs", "config_picmus_rf.yaml", repo_type="dataset",
-   )
+   config = zea.Config.from_path("hf://zeahub/configs/config_picmus_rf.yaml")
 
    path = config.data.dataset_folder + "/" + config.data.file_path
    with zea.File(path) as file:
@@ -50,7 +48,7 @@ Similarly, we can easily load one of the pretrained models from the :mod:`zea.mo
    model = EchoNetDynamic.from_preset("echonet-dynamic")
 
    # we'll load a single file from the dataset
-   with zea.Dataset("hf://zeahub/camus-sample/", "image_sc") as dataset:
+   with zea.Dataset("hf://zeahub/camus-sample/") as dataset:
       file = dataset[0]
       image = file.load_data("image_sc", indices=0)
 

docs/source/notebooks/agent/agent_example.ipynb

@@ -266,23 +266,25 @@
     "positions = phantoms.fish()\n",
     "magnitudes = np.ones(len(positions), dtype=np.float32)\n",
     "\n",
-    "rf_data = simulate_rf(\n",
-    "    scatterer_positions=positions,\n",
-    "    scatterer_magnitudes=magnitudes,\n",
-    "    probe_geometry=probe.probe_geometry,\n",
-    "    apply_lens_correction=scan.apply_lens_correction,\n",
-    "    lens_thickness=scan.lens_thickness,\n",
-    "    lens_sound_speed=scan.lens_sound_speed,\n",
-    "    sound_speed=scan.sound_speed,\n",
-    "    n_ax=scan.n_ax,\n",
-    "    center_frequency=probe.center_frequency,\n",
-    "    sampling_frequency=probe.sampling_frequency,\n",
-    "    t0_delays=scan.t0_delays,\n",
-    "    initial_times=scan.initial_times,\n",
-    "    element_width=scan.element_width,\n",
-    "    attenuation_coef=scan.attenuation_coef,\n",
-    "    tx_apodizations=scan.tx_apodizations,\n",
-    ")\n",
+    "simulation_args = {\n",
+    "    \"scatterer_positions\": positions,\n",
+    "    \"scatterer_magnitudes\": magnitudes,\n",
+    "    \"probe_geometry\": probe.probe_geometry,\n",
+    "    \"apply_lens_correction\": scan.apply_lens_correction,\n",
+    "    \"lens_thickness\": scan.lens_thickness,\n",
+    "    \"lens_sound_speed\": scan.lens_sound_speed,\n",
+    "    \"sound_speed\": scan.sound_speed,\n",
+    "    \"n_ax\": scan.n_ax,\n",
+    "    \"center_frequency\": probe.center_frequency,\n",
+    "    \"sampling_frequency\": probe.sampling_frequency,\n",
+    "    \"t0_delays\": scan.t0_delays,\n",
+    "    \"initial_times\": scan.initial_times,\n",
+    "    \"element_width\": scan.element_width,\n",
+    "    \"attenuation_coef\": scan.attenuation_coef,\n",
+    "    \"tx_apodizations\": scan.tx_apodizations,\n",
+    "}\n",
+    "\n",
+    "rf_data = simulate_rf(**simulation_args)\n",
     "print(\"Simulated RF data shape:\", rf_data.shape)"
    ]
   },
@@ -334,9 +336,9 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m No transmit origins provided, using zeros\n",
       "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m No azimuth angles provided, using zeros\n",
-      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[33mDEBUG\u001b[0m [zea.Pipeline] The following input keys are not used by the pipeline: {'center_frequency', 'n_el', 'zlims', 'xlims'}. Make sure this is intended. This warning will only be shown once.\n"
+      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[38;5;214mWARNING\u001b[0m No transmit origins provided, using zeros\n",
+      "\u001b[1m\u001b[38;5;36mzea\u001b[0m\u001b[0m: \u001b[33mDEBUG\u001b[0m [zea.Pipeline] The following input keys are not used by the pipeline: {'center_frequency', 'xlims', 'n_el', 'zlims'}. Make sure this is intended. This warning will only be shown once.\n"
      ]
     }
    ],
@@ -384,6 +386,99 @@
    "source": [
     "That's it! You have now simulated ultrasound RF data and reconstructed a B-mode image using `zea`."
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c6dc06de",
+   "metadata": {},
+   "source": [
+    "## Speedup with Just-In-Time compilation (JIT)\n",
+    "\n",
+    "The `simulate_rf` function took quite some time to compute in this example. Larger experiments with more point scatterers or acquisitions can execute very slowly. In this case, it is advised to [JIT-compile](https://docs.jax.dev/en/latest/jit-compilation.html) the `simulate_rf` function. The way you do this depends on which machine learning backend (e.g., JAX, PyTorch, TensorFlow) you are using (see [documentation](../../installation.rst#Backend) for details). Starting with JAX, you can simply wrap the function with `jax.jit` as follows:\n",
+    "\n",
+    "**JAX**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "8aa8873d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from jax import jit\n",
+    "\n",
+    "simulate_rf_jit = jit(simulate_rf, static_argnames=[\"apply_lens_correction\", \"n_ax\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "89b0adf8",
+   "metadata": {},
+   "source": [
+    "Let's execute and time the JIT versus non-JIT versions of the `simulate_rf` function to see the speedup."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "83453631",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[1mFunction Timing Statistics\u001b[0m\n",
+      "=====================================================================================================\n",
+      "\u001b[36mFunction\u001b[0m              \u001b[32mMean\u001b[0m          \u001b[32mMedian\u001b[0m        \u001b[32mStd Dev\u001b[0m       \u001b[33mMin\u001b[0m           \u001b[33mMax\u001b[0m           \u001b[35mCount\u001b[0m    \n",
+      "-----------------------------------------------------------------------------------------------------\n",
+      "\u001b[36msimulate_rf\u001b[0m           \u001b[32m0.220066\u001b[0m      \u001b[32m0.218923\u001b[0m      \u001b[32m0.021850\u001b[0m      \u001b[33m0.189399\u001b[0m      \u001b[33m0.255911\u001b[0m      \u001b[35m30\u001b[0m       \n",
+      "\u001b[36msimulate_rf (JIT)\u001b[0m     \u001b[32m0.004081\u001b[0m      \u001b[32m0.003444\u001b[0m      \u001b[32m0.003207\u001b[0m      \u001b[33m0.003159\u001b[0m      \u001b[33m0.020947\u001b[0m      \u001b[35m30\u001b[0m       \n"
+     ]
+    }
+   ],
+   "source": [
+    "from zea.utils import FunctionTimer\n",
+    "\n",
+    "# Warm-up JIT compilation before benchmarking\n",
+    "simulate_rf_jit(**simulation_args)\n",
+    "\n",
+    "timer = FunctionTimer()\n",
+    "timed_rf = timer(simulate_rf, name=\"simulate_rf\")\n",
+    "timed_rf_jit = timer(simulate_rf_jit, name=\"simulate_rf (JIT)\")\n",
+    "\n",
+    "for _ in range(30):\n",
+    "    timed_rf_jit(**simulation_args)\n",
+    "    timed_rf(**simulation_args)\n",
+    "\n",
+    "timer.print()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "72b4387a",
+   "metadata": {},
+   "source": [
+    "If you are using another backend, a similar approach can be taken:"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4eba2540",
+   "metadata": {},
+   "source": [
+    "**PyTorch**\n",
+    "```python\n",
+    "import torch\n",
+    "simulate_rf_jit = torch.jit.script(simulate_rf)\n",
+    "```\n",
+    "**TensorFlow**\n",
+    "```python\n",
+    "import tensorflow as tf\n",
+    "simulate_rf_jit = tf.function(simulate_rf)\n",
+    "```"
+   ]
   }
  ],
  "metadata": {