{"id":23221534,"url":"https://github.com/ickc/log-pyautolens","last_synced_at":"2026-03-20T01:01:50.149Z","repository":{"id":267847080,"uuid":"902457775","full_name":"ickc/log-PyAutoLens","owner":"ickc","description":"logs of my work on the PyAutoLens JAX project","archived":false,"fork":false,"pushed_at":"2025-03-21T17:59:39.000Z","size":464,"stargazers_count":0,"open_issues_count":0,"forks_count":1,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-03-28T09:40:40.476Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":null,"has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/ickc.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-12-12T15:51:12.000Z","updated_at":"2025-03-21T17:59:42.000Z","dependencies_parsed_at":"2024-12-12T20:29:30.041Z","dependency_job_id":"f4fc2d67-ee71-4148-889a-b32777f6b21a","html_url":"https://github.com/ickc/log-PyAutoLens","commit_stats":null,"previous_names":["ickc/log-pyautolens"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ickc%2Flog-PyAutoLens","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ickc%2Flog-PyAutoLens/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ickc%2Flog-PyAutoLens/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/ickc%2Flog-PyAutoLens/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/ickc","download_url":"https://codeload.github.com/ickc/log-PyAutoLens/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247361598,"owners_count":20926642,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-12-18T22:17:28.254Z","updated_at":"2026-01-19T15:06:40.134Z","avatar_url":"https://github.com/ickc.png","language":null,"funding_links":[],"categories":[],"sub_categories":[],"readme":"A logbook for the project \"DiRAC: revealing the nature of dark matter with the James Webb space telescope and JAX\".\n\nThis should contains all the notes about the project, notebooks and scripts I experimented with, environment(s) I setup for the project, which includes packages related to this project as submodules\n\nContribution worthy contents should not resides here, but as PRs to respective repos.\n\n# Introduction for others\n\nThis is a logbook written primarily for me, which may contains notes that may only make sense to me personally. But it also serves as a reproducible recipe to repeat what I experimented on or how I run other repos required for this project.\n\nAs such, the 3 main prerequisites are `git`, `pixi`, and `task`.\n\nGit submodules are used to record all the repos needed to run some of the scripts/examples/notebooks in this project. I.e. just git pull this repo and initiate submodules, conveniently available as `task init`.\n\n`pixi` is used to reproduce the software stack, including the use of the softwares pinned via git submodules. You can either follow [Getting Started - Pixi by prefix.dev](https://pixi.sh/latest/) to install pixi, or use `environment.yml` with conda if you're more comfortable with that.\n\nTaskfile is used to run reusable short one liners. You can either install it via [Installation | Task](https://taskfile.dev/installation/), or open `Taskfile.yml` and copy the command there and run it manually.\n\npixi creates a conda environment under the `.pixi/envs/default` subdirectory of current directory. You can use that as a Jupyter kernel to run notebooks for example.\n\nCaveats: there are some notebooks from the git submodules that when it is loading some datasets, you need to modify the path either because it is broken or because you're running it from a location different from it was written initially. The fix is usually obvious.\n\nLastly, there's some scripts from autojax that can only be run from the environment from this project, as it requires some dependencies from the git submodules here. Run it like this:\n\n```sh\npixi run packages/autojax/external/get_dataset.py\n# or if you use conda\nconda activate .pixi/envs/default\npython packages/autojax/external/get_dataset.py\n```\n\n# Project overview\n\n- DiRAC: revealing the nature of dark matter with the James Webb space telescope and JAX\n- aim: add JAX support in PyAutoLens\n- Steven is working on this project\n- end in March 7th 2025\n- join\n    - [x] Slack\n    - [x] bi-weekly meeting on the science of lensing, occasionally jax\n        - Gokmen will forward the Zoom link to me\n    - [x] DIRAC cluster?\n- people:\n    - James Nightingale (Durham)\n    - Gokmen Kilic\n    - Coleman Krawczyk is the one to ask about jax\n    - Rich developed PyAutoLens from scratch for 6-7 years\n- [x] schedule a weekly meeting in the week of 12/9\n- [x] learn jax\n- Gokmen will ask James about tasks split between us\n\n# Questions\n\n- [x] I have a question on rectangular grid and pixel. When reading through the code, it seems it may support a grid (and its mask) that is not square, and from the pixel_scales, it seems that even pixel itself can be a rectangle too. Do I need to make sure the code supports the non-square case?\n    - non-square pixel is used elsewhere but can be assumed in this project\n- [x] `regularization_matrix` is guaranteed to be symmetric? (H = B.T @ B?)\n    - This relates to the inherent property of `neighbors`\n        - [x] gen `neighbors` to ensure symmetry\n    - [x] take advantage of it using `jnp.linalg.cholesky`\n- [x] check and ask `w = C.T @ C + S.T @ S`\n- [x] `w_tilde` used in the likelihood calculation only through the eventual calculation of `curvature_matrix`?\n- [x] flag the JAX CPU multithreading limitation\n    - https://github.com/jax-ml/jax/issues/5022#issuecomment-1222336766\n    - Numba implementation might still be useful for CPU?\n- [x] scaling test to understand the computational requirement from the science case\n    - FLOP (speed)\n    - memory need\n    - perhaps keep multiple implementations I have and see how all of them scales in large M, K case\n- [x] check `mapping_matrix` sparseness\n    - is it the only matrix that is sparse?\n    - is `curvature_matrix_via_w_tilde_curvature_preload_interferometer_from` essentially calculating the `mapping_matrix` on the fly without writing it down explicitly?\n        - [x] understand the FLOPS and memory requirement of this function better\n- [x] mention N and `N_PRIME` and its implication in calculating w_preload/w_compact\n    - I use N = N_PRIME in my mock data for scaling test, but it probably requires some architectural change in the PyAuto* code base\n- mentioned 4 notebooks, last 2 repeated:\n    * [autogalaxy\\_workspace/notebooks/advanced/log\\_likelihood\\_function/imaging/light\\_profile/log\\_likelihood\\_function.ipynb at release · Jammy2211/autogalaxy\\_workspace](https://github.com/Jammy2211/autogalaxy\\_workspace/blob/release/notebooks/advanced/log\\_likelihood\\_function/imaging/light\\_profile/log\\_likelihood\\_function.ipynb)\n    * [autogalaxy\\_workspace/notebooks/advanced/log\\_likelihood\\_function/imaging/linear\\_light\\_profile at release · Jammy2211/autogalaxy\\_workspace](https://github.com/Jammy2211/autogalaxy\\_workspace/tree/release/notebooks/advanced/log\\_likelihood\\_function/imaging/linear\\_light\\_profile)\n    * [autogalaxy\\_workspace/notebooks/advanced/log\\_likelihood\\_function/interferometer/light\\_profile/log\\_likelihood\\_function.ipynb at release · Jammy2211/autogalaxy\\_workspace](https://github.com/Jammy2211/autogalaxy\\_workspace/blob/release/notebooks/advanced/log\\_likelihood\\_function/interferometer/light\\_profile/log\\_likelihood\\_function.ipynb)\n- difference between \u003chttps://github.com/Jammy2211/autogalaxy_workspace/blob/main/notebooks/advanced/log_likelihood_function/imaging/light_profile/log_likelihood_function.ipynb\u003e and \u003chttps://github.com/Jammy2211/autogalaxy_workspace/tree/main/notebooks/advanced/log_likelihood_function/imaging/linear_light_profile\u003e: is the only difference between the 2 approaches, practically speaking, the former does not optimized the intensity per profile, hence it eliminates 2 free parameters when sampling with the log-likelihood?\n\n# TODO\n\n- [ ] final release\n    - [ ] make a copy in UoE RSE repo\n    - [ ] make a GitHub Releases\n- [x] fix bugs on GPU:\n    - `test_constant_regularization_matrix_from_jax`\n    - `test_reconstruction_positive_negative_from_jax`: `XlaRuntimeError: INTERNAL: cuSolver internal error`\n- [ ] scaling test:\n    - include w_tilde and curvature matrix including running on GPU\n    - for checking if there's unexpected performance issue in memory/FLOP utilization in larger \n    - 3 variants on CPU: constraint to 1 thread, 256 thread, and no constraint\n    - run on A100 with `cosma8-shm`\n    - bonus: Grace Hopper\n        - solve dotfiles dispatch first\n        - support aarch64 in pixidataset size\n- [ ] documentation\n    - function docstrings: especially for new functions\n        - math\n        - algorithm: point to vmap, scan, loop ordering, etc.\n    - add \"advanced benchmarking\"\n    - general doc, such as README of autojax\n- [ ] presentation\n    - for general audience, albeit slides can be cherry-picked depending on audiences\n    - introduction of the project, expectation\n        - [x] ask about autodiff expectation: which variant of MCMC exactly?\n            - https://github.com/blackjax-devs/blackjax is a possible choice\n        - [x] what is the expected no. of iterations run on the log-likelihood function?\n            - ~100,000\n    - benchmark \u0026 profiling: JAX vs Numba, CPU vs GPU, etc.\n    - showcase examples\n    - slide deck\n        - https://docs.google.com/presentation/d/1FOsbmv58V4PwRNt_E0AvF54MFUNZ2e_vItHMXmVlTW8/edit\n        - https://docs.google.com/presentation/d/1XzNGCC_bs1xDpmCLaXOdTnG7cdpfvlpIkEhwuypvYwA/edit\n            - slide 6 of the second deck is good\n- bonus features\n    - generates BCOO from sparse mapping matrix\n    - shorten turn around of initial calculation of w_tilde/w_compact.\n    - circulant w_tilde approximation\n- [x] check number of non-zero element per row/column of mapping_matrix\n    - For the sample dataset, each image plane pixel maps to 3 source pixels (probably only for Delaunay triangulation). Curiously this is not P=10 here.\n- [x] explore how mapping_matrix is calculated?\n- [x] check NumbaPerformanceWarning (low priority as we aren't porting to Numba)\n- [x] request projects dp004 and do018 for benchmarking via https://safe.epcc.ed.ac.uk/dirac @wait(for approval)\n    - [ ] setup /snap8/scratch/dp004/dc-kili1/RAC16/PyAutoLens\n        \u003e You can create files under the snap8/do009 . So something like: /snap8/scratch/do009/dc-cheu2/RAC16/PyAutoLens\n- [x] refactor ported functions per implementation\n    - [x] setup unit test\n    - [x] setup doc\n    - [x] set up benchmark to compare implementations\n        - [x] repeat previous rudimentary benchmark with float64 and double check for consistency\n        - [x] compare results from pytest-benchmark to manually running it. The results seems wildly different\n            - perhaps write a manual benchmark with cli and ensure all available threads are used? Perhaps add a matmul in the beginning as a control.\n    - [x] Update duplicated code from \u003chttps://github.com/Jammy2211/dirac_rse_interferometer/\u003e to this repo\n- [x] explore calculation of symmetric matrix under JAX framework\n- Go through\n    - [x] \u003chttps://github.com/Jammy2211/autogalaxy_workspace_test/blob/master/jax_examples/task_2_simple_conversions/func_grad_manual.py\u003e\n    * [x] [autogalaxy_workspace/notebooks/advanced/log_likelihood_function/imaging/light_profile/log_likelihood_function.ipynb at release · Jammy2211/autogalaxy_workspace](https://github.com/Jammy2211/autogalaxy_workspace/blob/release/notebooks/advanced/log_likelihood_function/imaging/light_profile/log_likelihood_function.ipynb)\n    * [x] [autogalaxy_workspace/notebooks/advanced/log_likelihood_function/imaging/linear_light_profile at release · Jammy2211/autogalaxy_workspace](https://github.com/Jammy2211/autogalaxy_workspace/tree/release/notebooks/advanced/log_likelihood_function/imaging/linear_light_profile)\n    * [x] [autogalaxy_workspace/notebooks/advanced/log_likelihood_function/interferometer/light_profile/log_likelihood_function.ipynb at release · Jammy2211/autogalaxy_workspace](https://github.com/Jammy2211/autogalaxy_workspace/blob/release/notebooks/advanced/log_likelihood_function/interferometer/light_profile/log_likelihood_function.ipynb)\n    * [x] [autogalaxy_workspace/notebooks/advanced/log_likelihood_function/imaging/pixelization/log_likelihood_function.ipynb at release · Jammy2211/autogalaxy_workspace](https://github.com/Jammy2211/autogalaxy_workspace/blob/release/notebooks/advanced/log_likelihood_function/imaging/pixelization/log_likelihood_function.ipynb)?\n- go through \u003chttps://github.com/Jammy2211/dirac_rse_interferometer\u003e\n    \u003e At the moment, you only care about this script, which shows how a simple dataset we are going to fit is simulated:\n    \u003e\n    \u003e https://github.com/Jammy2211/dirac_rse_interferometer/blob/main/simulators/sma.py\n    \u003e\n    \u003e And this script, which illustrates the function we discussed which can currently take 2+ weeks to run:\n    \u003e\n    \u003e https://github.com/Jammy2211/dirac_rse_interferometer/blob/main/w_tilde.py\n\n    - [x] https://github.com/Jammy2211/dirac_rse_interferometer/blob/main/w_tilde.py\n\n        $$\\begin{aligned}\n        \\mathbf{V}_r \u0026: \\text{Real visibilities vector} \\\\\n        \\mathbf{N}_r \u0026: \\text{Real noise map vector} \\\\\n        \\mathbf{u} \u0026: \\text{UV wavelengths matrix} \\\\\n        \\mathbf{g} \u0026: \\text{Grid radians coordinates matrix} \\\\\n        M \u0026: \\text{Number of image pixels} \\\\\n        N \u0026: \\text{Number of visibility points}\n        \\end{aligned}$$\n\n        $$\\tilde{w}_{\\text{data},i} = \\sum_{j=1}^N \\left(\\frac{N_{r,j}^2}{V_{r,j}}\\right)^2 \\cos\\left(2\\pi(g_{i,1}u_{j,0} + g_{i,0}u_{j,1})\\right)$$\n\n        where $$i \\in [0,M-1]$$ and output is vector $$\\tilde{\\mathbf{w}}_\\text{data} \\in \\mathbb{R}^M$$\n\n    - [x] https://github.com/Jammy2211/dirac_rse_interferometer/blob/main/w_tilde/w_tilde_curvature.py\n\n        $$W̃_{ij} = \\sum_{k=1}^N \\frac{1}{n_k^2} \\cos(2\\pi[(g_{i1} - g_{j1})u_{k0} + (g_{i0} - g_{j0})u_{k1}])$$\n\n    - [x] https://github.com/Jammy2211/dirac_rse_interferometer/blob/main/simulators/sma.py\n        - seems to be from https://github.com/Jammy2211/autolens_workspace/blob/main/notebooks/simulators/interferometer/instruments/sma.ipynb\n\n\u003e Another interesting question which could prove a problem with the full JAX implementation  is how to get a 2d delaunay mesh to work, which currently uses scipy https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.Delaunay.html\n\u003e We also use a voronoi mesh with 2d natural neighbour interpolation.\n\u003e My understanding is getting these to run in jax could be very hard, but it would be good to scope out if it looks at all feasible.\n\u003e https://github.com/Jammy2211/PyAutoArray/blob/main/autoarray/structures/mesh/triangulation_2d.py\n\u003e triangulation_2d.py\n\u003e \n\u003e But the first round of JAX conversions I'll sort the examples for don't need these meshes, that kind of the problem we face after getting the easier to convert code sorted.\n\u003e \n\u003e It may also make sense to implement them via this tool https://arxiv.org/abs/2403.08847 first, which will have limitations, and worry about a full jax implementation later\n\n---\n\nMore notes from James on Jan 15th 2025:\n\n\u003e In preparation for your return, I've put together a lot of materials to hopefully help you understand the problem better, rather than it just being a bunch of abstract numpy arrays which get mushed together.\n\u003e\n\u003e I would approach the work via the following steps (you may of done some):\n\u003e\n\u003e 1. Read this example, which provides a very basic explanation of how a galaxy light profile is used to fit an image of a galaxy. This example uses imaging data, as opposed to interferometer data, because it is conceptually simpler and therefore a better starting point \u003chttps://github.com/Jammy2211/autogalaxy_workspace/blob/main/notebooks/advanced/log_likelihood_function/imaging/light_profile/log_likelihood_function.ipynb\u003e\n\u003e\n\u003e 2. Read this example, which again using imaging data (for simplicity) explains how the problem can use linear algebra to solve for the light of the galaxy: \u003chttps://github.com/Jammy2211/autogalaxy_workspace/tree/main/notebooks/advanced/log_likelihood_function/imaging/linear_light_profile\u003e\n\u003e\n\u003e 3. Read this example, which using imaging data explains how a pixelized reconstruction of a galaxy is performed: \u003chttps://github.com/Jammy2211/autogalaxy_workspace/blob/main/notebooks/advanced/log_likelihood_function/imaging/pixelization/log_likelihood_function.ipynb\u003e\n\u003e\n\u003e 4. Read this example, which now shows how the simple galaxy light profile calculation is performed for interferometer data, which includes a non-uniform Fast Fourier transform: \u003chttps://github.com/Jammy2211/autogalaxy_workspace/blob/main/notebooks/advanced/log_likelihood_function/interferometer/light_profile/log_likelihood_function.ipynb\u003e\n\u003e\n\u003e 5. Read this example, which explains how a pixelized source is performed on interferometer data: \u003chttps://github.com/Jammy2211/autogalaxy_workspace/blob/main/notebooks/advanced/log_likelihood_function/interferometer/pixelization/log_likelihood_function.ipynb\u003e\n\u003e\n\u003e 6. Read this example, which explains how the \"w-tilde\" linear algebra formalism (which you converted some functions to JAX already for) changes the calculation in step 5: \u003chttps://github.com/Jammy2211/autogalaxy_workspace/blob/main/notebooks/advanced/log_likelihood_function/interferometer/pixelization/w_tilde.ipynb\u003e\n\u003e\n\u003e 7. Begin to convert all these steps to JAX in the GitHub repo we are working on: \u003chttps://github.com/Jammy2211/dirac_rse_interferometer/blob/main/inversion/likelihood_function.py\u003e\n\u003e\n\u003e By building up a deeper understanding of the problem end-to-end, I think we'll find the overall JAX implementation goes smoother, especially as we will have to slowly extend the code to encompass every step of the likelihood function by the end. However, for now, I think the task is broken down in a way that should be manageable.\n\u003e\n\u003e I will look at your PR's in the next few days and as I said above, I am working on making the JAX feature branch source code a lot more useable (e.g. refactor, fix unit tests) so we can easily incorporate your JAX'd code outside the source code into the source code without too many headaches.\n\n# PyAutoLens Intro\n\nFrom @Jam:\n\n\u003e Busy day so will send expanded instructions later but this is the Google Doc which explains the project layout and how to get started:[https://docs.google.com/document/d/1jGlxdnpjX5t4rXKBc0hrRx9vWcgXXnoR_5RY2nWTeUs/edit?usp=sharing](https://docs.google.com/document/d/1jGlxdnpjX5t4rXKBc0hrRx9vWcgXXnoR_5RY2nWTeUs/edit?usp=sharing)The project is spread out over multiple repos so takes a bit of getting used to.This is the simple task which is probably a good starting point as an example of trying to convert code from numba to JAX:[https://github.com/Jammy2211/autogalaxy_workspace_test/tree/master/jax_examples/task_2_simple_conversions](https://github.com/Jammy2211/autogalaxy_workspace_test/tree/master/jax_examples/task_2_simple_conversions)\n\n- branch: feature/jax_wrapper for JAX\n- E.g. PR to fix JAX wrapper: https://github.com/Jammy2211/PyAutoArray/pull/156\n- [x] Work out of source code, porting, then worry about putting it back later (get initial out of source code repo from @jam)\n- [x] get notebooks, papers from @jam\n\n## Environment\n\nI created a reproducible-ish environment using pixi.\n\nStatus:\n\n- all `PyAuto*`'s dependencies added.\n\n    \u003e WARN These conda-packages will be overridden by pypi:\n    \u003e        networkx, anesthetic, jsonpickle\n\n- [x] setup `PyAuto*` as submodules and use editable dependencies documented in \u003chttps://pixi.sh/latest/reference/pixi_manifest/#pypi-dependencies\u003e.\n\n    ```bash\n    git submodule add -b main https://github.com/Jammy2211/autogalaxy_workspace.git packages/autogalaxy_workspace\n    git submodule add -b main https://github.com/Jammy2211/PyAutoArray.git packages/PyAutoArray\n    git submodule add -b main https://github.com/Jammy2211/PyAutoGalaxy.git packages/PyAutoGalaxy\n    git submodule add -b main https://github.com/Jammy2211/PyAutoLens.git packages/PyAutoLens\n    git submodule add -b main https://github.com/rhayes777/PyAutoConf.git packages/PyAutoConf\n    git submodule add -b main https://github.com/rhayes777/PyAutoFit.git packages/PyAutoFit\n    git submodule add -b main https://github.com/Jammy2211/dirac_rse_interferometer.git packages/dirac_rse_interferometer\n    git submodule add -b master https://github.com/Jammy2211/autogalaxy_workspace_test.git packages/autogalaxy_workspace_test\n    git submodule add -b main https://github.com/Jammy2211/autolens_workspace.git packages/autolens_workspace\n    git submodule add -b main git@github.com:ickc/python-autojax.git packages/autojax\n    ```\n\n- [x] run tests\n\n## Notes\n\n### Sizes\n\nM = number of image pixels in `real_space_mask` = ~70000  \nK = number of visibilitiies = ~1e7 (but this is only used to compute `w_tilde` otuside the likelihood function)  \nP = number of neighbors = 3 (for Delaunay) but can go up to 300 for Voronoi (but we can just focus on delaunay for now)  \nS = number of **source pixels** (e.g. `reconstruction.shape`) = 1000\n\n### `neighbors` spec\n\n`neighbors_sizes` is a 1D array of the number of neighbors in each `source_plane_mesh_grid` pixel.So if there are 3 entries which are non-zero before you get to minus ones, then `neighbor_sizes` for that entry will be 3.\n\nI guess this information could be expressed as a matrix of size `[source_pixels, source_pixels]` where non-zero entries mean they share a vertex, and then it could be done as some sort of matrix multplication. The for loop way of doing things made sense for `numba` but may not make as much sense for JAX?\n\nIn case you havent seen it, all these functgions have unitt ests which may help:\n\nhttps://github.com/Jammy2211/PyAutoArray/blob/main/test_autoarray/inversion/regularizations/test_constant.py\n\nThe unit test here might be even more useful:[https://github.com/Jammy2211/PyAutoArray/blob/main/test_autoarray/inversion/regularizations/test_regularization_util.py](https://github.com/Jammy2211/PyAutoArray/blob/main/test_autoarray/inversion/regularizations/test_regularization_util.py)\n\nThe `b_matrix` formalism is how you do this as matrix multiplication I think.\n\nSee the docstring here as well (glad I wrote all this down somewhere lol): [https://github.com/Jammy2211/PyAutoArray/blob/main/autoarray/inversion/regularization/abstract.py](https://github.com/Jammy2211/PyAutoArray/blob/main/autoarray/inversion/regularization/abstract.py)\n\nThis part of the docstring seems especially relevent:\n\n        # SKIPPING THE B MATRIX CALCULATION #\n\nThe `neighbors` array can come from one of 3 sources, depending on the pixel grid used to reconstruct the source:\n\n- `Rectangular`: pixels are rectangles and therefore they all have 4 neighbors (except those on the edge of the grid).\n- `Delaunay`: Pixels are Delaunay triangles and all have 3 neighbors (exception those on the edge).\n- `Voronoi`: Pixels are Voronoi cells and have variable numbers of neighbors.\n\nEach of the above 3 pixelizations has their own unique function to compute the `neighbors` array.The `rectangular` neighbor calcualtion is here:[https://github.com/Jammy2211/PyAutoArray/blob/main/autoarray/inversion/pixelization/mesh/mesh_util.py](https://github.com/Jammy2211/PyAutoArray/blob/main/autoarray/inversion/pixelization/mesh/mesh_util.py)For `Deluanay` and `Voroni` they use `scipy.sptial` libraries, and they could pose a significant challenge to the conversion to JAX (I think this is the biggest unknown in how we are going to get everything running in JAX).For your current work, this is why I set it up to assume we have `neighbors` and we can work our way upstream later to JAX-ify this part of the calculation.\n\n## Issues\n\nFrom @Jam:\n\n\u003e The way I approach this is the project has a dependency hierarchy of PyAutoFit -\u003e PyAutoArray -\u003e PyAutoGalaxy -\u003e PyAutoLens, so post it in the highest project the issue is relevant.\n\n### Chi-square bug\n\nThere seems something wrong with the following definition of chi-square:\n\n```py\nchi_squared_term_1 = np.linalg.multi_dot(\n    [\n        mapping_matrix.T,  # NOTE: shape = (N, )\n        w_tilde,  # NOTE: shape = (N, N)\n        mapping_matrix,\n    ]\n)\n\nchi_squared_term_2 = -np.multiply(\n    2.0, np.dot(mapping_matrix.T, dataset.w_tilde.dirty_image)\n)  # Need to double check dirty_image is the right input.\n\nchi_squared = chi_squared_term_1 + chi_squared_term_2\n```\n\nThe dimensions are wrong (these are not scalar) and also the 1st term is defined as curvature_matrix earlier.\n\n### Potential issues\n\n- \u003chttps://github.com/Jammy2211/PyAutoArray/blob/main/autoarray/numba_util.py\u003e: try and except requires better handlings\n\n# DiRAC COSMA\n\n## Links\n\n* [DiRAC High Performance Computing Facility – Supporting the STFC theory community](https://dirac.ac.uk/)\n* [DIRAC Project - hpcicc: Durham](https://safe.epcc.ed.ac.uk/dirac/Project/hpcicc/project\\_member.jsp)\n* [Memory Intensive Service: Durham – DiRAC High Performance Computing Facility](https://dirac.ac.uk/memory-intensive-durham/)\n* [COSMA Facilities — cosma 0.1 documentation](https://cosma.readthedocs.io/en/latest/facilities.html)\n\n## Notes\n\ndata (scratch?): `/cosma5/data/durham/dc-cheu2`\napps: `/cosma/apps/durham/dc-cheu2`\ncompute: only have access to the COSMA5 system:\n\n```bash\nsalloc --partition=cosma --time=02:00:00 --nodes=1 -A durham\nsalloc --partition=cosma5 --time=02:00:00 --nodes=1 -A durham\n```\n\n\u003e If you are a DiRAC user, please use login7 or [login8.cosma.dur.ac.uk](login8.cosma.dur.ac.uk),  \nand submit to the cosma7 or cosma8 queues as appropriate for your  \nproject.\n\n\u003e If you are not a DiRAC user, please use [login5.cosma.dur.ac.uk](login5.cosma.dur.ac.uk) and  \nsubmit to the cosma queue.\n\n\u003e The quota command will show you your disk quotas and current usage on  \nthe various different file systems.\n\n\u003e All file systems except for /cosma/home and /cosma/apps are Lustre  \nfile systems.\n\nTODO: use XDG env var to enforce the following instead of relying on root's symlink:\n\n```\nlrwxrwxrwx 1 root root 38 Dec  3 09:44 ./.apptainer -\u003e /cosma/apps/durham/dc-cheu2/.apptainer\nlrwxrwxrwx 1 root root 34 Dec  3 09:44 ./.cache -\u003e /cosma/apps/durham/dc-cheu2/.cache\nlrwxrwxrwx 1 root root 34 Dec  3 09:44 ./.conda -\u003e /cosma/apps/durham/dc-cheu2/.conda\nlrwxrwxrwx 1 root root 35 Dec  3 09:44 ./.config -\u003e /cosma/apps/durham/dc-cheu2/.config\nlrwxrwxrwx 1 root root 34 Dec  3 09:44 ./.local -\u003e /cosma/apps/durham/dc-cheu2/.local\nlrwxrwxrwx 1 root root 34 Dec  3 09:44 ./.spack -\u003e /cosma/apps/durham/dc-cheu2/.spack\n```\n\nWhich can be reproduce by:\n\n```bash\nTARGET_DIR=\"/cosma/apps/durham/dc-cheu2\"\nLINKS=(.apptainer .cache .conda .local .spack)\nfor LINK in \"${LINKS[@]}\"; do\n    ln -s \"$TARGET_DIR/$LINK\" \"$LINK\"\ndone  \n```\n\n### Submitting jobs\n\n[GPUs — cosma 0.1 documentation](https://cosma.readthedocs.io/en/latest/gpu.html#mi300x),\n\n```bash\n# MI-300X\nsrun --partition=mi300x --account=do018 --time=02:00:00 --pty /bin/bash\nsrun --partition=mi300x --account=do018 --time=02:00:00 --exclusive --pty /bin/bash\nsalloc --partition=mi300x --account=do018 --time=02:00:00 --exclusive\n# COSMA5\nsrun --partition=cosma5 --account=durham --time=72:00:00 --nodes=1 --ntasks=1 --cpus-per-task=256 --pty /bin/bash\n# A100\nsrun --partition=cosma8-shm --account=dr004 --time=00:30:00 --exclusive --pty /bin/bash\n\n```\n\n# References\n\n\u003e I have not read all that much so can't recommend a book but have seen this github recommended: [ratt-ru/foi-course: Fundamentals of Radio Interferometry and Aperture Synthesis Book](https://github.com/ratt-ru/foi-course)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fickc%2Flog-pyautolens","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fickc%2Flog-pyautolens","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fickc%2Flog-pyautolens/lists"}