docs(lens): add module-level docstrings to lens package

autolens/lens/sensitivity.py

@@ -1,3 +1,19 @@
+"""
+Sensitivity mapping for dark-matter subhalo detection.
+
+Sensitivity mapping asks: *given a lens model, how detectable would a dark-matter subhalo
+of mass M at position (y, x) be?*  It works by:
+
+1. Fixing the smooth lens model to the best-fit result.
+2. Simulating many lens datasets, each with a subhalo at a different position and mass,
+   using the smooth model plus a perturbation.
+3. Fitting each simulated dataset with the smooth model (no subhalo) to measure the change
+   in log-evidence caused by the subhalo.
+
+``SubhaloSensitivityResult`` wraps the generic ``PyAutoFit`` ``SensitivityResult`` with
+convenience properties for the subhalo grid positions (``y``, ``x``), the detection
+significance map, and Matplotlib visualisation helpers.
+"""
 import numpy as np
 from typing import Optional, List, Tuple
 

autolens/lens/subhalo.py

@@ -1,3 +1,18 @@
+"""
+Dark-matter subhalo detection via grid searches of non-linear searches.
+
+This module provides result containers and visualisation helpers for a subhalo detection
+workflow in which a grid of ``PyAutoFit`` non-linear searches is run.  Each cell of the
+grid confines the subhalo's (y, x) centre to a small sub-region of the image plane using
+uniform priors and fits the lens model with a subhalo included.
+
+``SubhaloGridSearchResult`` wraps ``af.GridSearchResult`` with:
+
+- ``y`` / ``x`` — the physical centre coordinates of each grid cell.
+- ``log_evidence_differences`` — the Bayesian evidence improvement from adding a subhalo
+  relative to a smooth-model fit, useful for building a detection significance map.
+- Plotting helpers that overlay the detection map on the lens image.
+"""
 import numpy as np
 from typing import List, Optional, Tuple
 

autolens/lens/to_inversion.py

@@ -1,3 +1,21 @@
+"""
+Interface between a ``Tracer`` and the linear-algebra inversion module.
+
+``TracerToInversion`` extends the ``autogalaxy`` ``GalaxiesToInversion`` pattern to the
+multi-plane lensing setting.  It is responsible for:
+
+- Extracting every ``LightProfileLinear`` and ``Pixelization`` from the tracer's galaxies.
+- Ray-tracing each linear profile's grid to the correct source plane using the tracer's
+  multi-plane deflection calculations.
+- Assembling the ``mapping_matrix`` that maps source-plane parameters (intensities, mesh
+  coefficients) to image-plane pixels.
+- Passing the assembled objects to ``autoarray``'s ``inversion_from`` factory so that the
+  linear system  ``F x = d`` can be solved for the best-fit intensities / reconstructed
+  image.
+
+This class is not used directly by the user; it is instantiated inside
+``FitImaging`` / ``FitInterferometer`` whenever the tracer contains linear components.
+"""
 from typing import Dict, List, Optional, Tuple, Type, Union
 
 import numpy as np

autolens/lens/tracer.py

@@ -1,3 +1,23 @@
+"""
+The core gravitational-lensing ray-tracing module for **PyAutoLens**.
+
+The central class is ``Tracer``, which groups a list of ``Galaxy`` objects by redshift
+into a series of *planes* and performs multi-plane gravitational lensing calculations.
+Key capabilities:
+
+- **Multi-plane ray tracing** — images are traced from the source plane through
+  intermediate lens planes to the observer using the cosmology-dependent angular
+  diameter distances between each pair of planes.
+- **Lensed images** — the light of each galaxy is ray-traced to its correct position
+  and the contributions from all planes are summed.
+- **Lensing maps** — convergence, shear, magnification, deflection angles, and potential
+  maps can all be computed on arbitrary ``Grid2D`` grids.
+- **Critical curves & caustics** — the image-plane loci where magnification diverges and
+  their source-plane counterparts.
+- **Lens modeling** — the ``Tracer`` is the core object used by
+  ``FitImaging`` / ``FitInterferometer`` / ``FitPointDataset`` when performing a
+  Bayesian model fit via a ``PyAutoFit`` non-linear search.
+"""
 from abc import ABC
 import numpy as np
 from scipy.interpolate import griddata

autolens/lens/tracer_util.py

@@ -1,3 +1,16 @@
+"""
+Utility functions supporting the ``Tracer`` ray-tracing calculations.
+
+This module contains lower-level helpers that are called by ``Tracer`` but kept separate
+to avoid cluttering the main class.  Key functions:
+
+- ``plane_redshifts_from`` — derives the list of unique plane redshifts from a list of
+  galaxies, collapsing multiple galaxies at the same redshift into a single plane.
+- ``ordered_plane_redshifts_with_slicing_from`` — extends the above with optional
+  redshift slicing for multi-plane calculations that include intermediate planes.
+- ``positions_in_ordered_planes_from`` — distributes a set of image-plane positions
+  across the ordered plane list so that multi-plane tracing can propagate them.
+"""
 import numpy as np
 from typing import List, Optional
 
@@ -145,7 +158,7 @@ def traced_grid_2d_list_from(
 
     for plane_index, galaxies in enumerate(planes):
 
-        scaled_grid = grid.array
+        scaled_grid = xp.asarray(grid.array)
 
         if plane_index > 0:
 
@@ -164,7 +177,7 @@ def traced_grid_2d_list_from(
                 scaled_grid = scaled_grid - scaled_deflections
 
         scaled_grid = aa.Grid2DIrregular(
-            values=scaled_grid,
+            values=scaled_grid, xp=xp
         )
 
         traced_grid_list.append(scaled_grid)