Add jax_zero_contour-based critical curve and caustic tracing

CLAUDE.md

@@ -9,24 +9,34 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 pip install -e ".[dev]"
 ```
 
-### Run Tests
-```bash
-# All tests
-python -m pytest test_autogalaxy/
+### Run Tests
+```bash
+# All tests
+python -m pytest test_autogalaxy/
 
 # Single test file
 python -m pytest test_autogalaxy/galaxy/test_galaxy.py
 
 # Single test
 python -m pytest test_autogalaxy/galaxy/test_galaxy.py::TestGalaxy::test_name
 
-# With output
-python -m pytest test_autogalaxy/imaging/test_fit_imaging.py -s
-```
-
-### Formatting
-```bash
-black autogalaxy/
+# With output
+python -m pytest test_autogalaxy/imaging/test_fit_imaging.py -s
+```
+
+### Codex / sandboxed runs
+
+When running Python from Codex or any restricted environment, set writable cache directories so `numba` and `matplotlib` do not fail on unwritable home or source-tree paths:
+
+```bash
+NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib python -m pytest test_autogalaxy/
+```
+
+This workspace is often imported from `/mnt/c/...` and Codex may not be able to write to module `__pycache__` directories or `/home/jammy/.cache`, which can cause import-time `numba` caching failures without this override.
+
+### Formatting
+```bash
+black autogalaxy/
 ```
 
 ## Architecture
@@ -213,4 +223,4 @@ find . -type f -name "*.py" | xargs dos2unix
 ```
 
 Prefer simple shell commands.
-Avoid chaining with && or pipes.
+Avoid chaining with && or pipes.

autogalaxy/config/visualize/general.yaml

@@ -2,6 +2,7 @@ general:
   backend: default
   imshow_origin: upper
   zoom_around_mask: true
+  critical_curves_method: zero_contour
 inversion:
   reconstruction_vmax_factor: 0.5
 zoom:

autogalaxy/galaxy/plot/galaxy_plots.py

@@ -143,4 +143,4 @@ def _deflections_x(mp):
             )
 
         plt.tight_layout()
-        _save_subplot(fig, output_path, f"subplot_{name}", output_format)
+        _save_subplot(fig, output_path, name, output_format)