PyAutoLens

repository·main·Indexed 41 days ago

https://github.com/pyautolabs/pyautolens

An open-source framework for modeling strong gravitational lenses. PyAutoLens-JAX leverages JAX for high-performance GPU acceleration and provides a unified API for imaging, interferometer, and point lensing tools, including galaxy modeling and configuration utilities.

Tokens
32.6K
Snippets
67
Records
163
Agent score
51%

What's inside PyAutoLens

  1. Overview of Lens Modeling and Simulation workflows

    main

    PyAutoLens supports two primary workflows for strong lensing research:

    1. Lens Modeling

    Goal: Fit physical models to strong-lensing data to infer properties of lens and source galaxies. Key Features:

    • Simplicity and Scalability: Designed to handle large datasets.
    • Speed: GPU acceleration is provided via JAX.
    • Iterative Fitting: Models are improved progressively to match observed images.

    2. Simulations

    Goal: Generate synthetic lensing data. Use Cases:

    • Practicing modeling before applying it to real data.
    • Generating large training sets for machine learning.
    • Testing lensing theory in controlled environments.

    Both workflows support various lensing regimes (galaxy-scale, multi-galaxy, group-scale, and cluster-scale) and data types (e.g., CCD imaging, interferometer data).

  2. Core capabilities of PyAutoLens-JAX

    main

    PyAutoLens-JAX is a differentiable, GPU-accelerated framework for strong and weak gravitational lensing modelling. It provides the following core capabilities:

    • Differentiable Modelling: Uses JAX for automatic differentiation, enabling gradient-based optimization and sampling that scales to a high number of free parameters.
    • GPU Acceleration: Supports Just-In-Time (JIT) compilation and GPU execution for high-speed likelihood evaluation.
    • Multi-Scale & Multi-Regime Analysis: Supports galaxy-, group-, and cluster-scale models.
    • Joint Dataset Integration: Allows combining multiple data types into a single probabilistic model, including:
      • CCD imaging (multi-band optical/infrared)
      • Interferometer visibilities (radio/submillimetre)
      • Point-source observables (lensed quasars or supernovae, including image positions, time delays, and flux)
      • Weak-lensing catalogues (shear catalogues)
    • API Preservation: Maintains the established object-oriented PyAutoLens API while providing the JAX-compatible numerical backend.
  3. Understand the PyAutoLens workspace structure

    main

    The autolens_workspace is organized into several key directories that facilitate lensing workflows:

    • Scripts / Notebooks: Contains Python scripts and Jupyter notebooks demonstrating lensing calculations, lens modeling, and other PyAutoLens features.
    • Config: Contains configuration files that customize default settings, visualization backends, priors, notation, and Cartesian grid behavior.
    • Dataset: Stores datasets used for lens modeling, including simulator-generated datasets.
    • Output: The destination folder where all modeling results are saved.
    • SLaM: Contains advanced lens modeling pipelines using the Source, Light and Mass (SLaM) approach.
    • HowToLens: While part of the ecosystem, this is a separate repository containing a lecture series of Jupyter notebooks for building model fitting projects.
  4. Compose Cluster-scale Lens Models

    main
    Cluster-scale models utilize the group mass framework (including a group halo and tidally truncated dPIE members on scaling relations) but employ a different source strategy. Instead of extended sources, they use many point-source multiple-image position datasets, each at its own redshift, fitted via AnalysisPoint.
  5. Prevent demagnified solutions using PositionsLH

    main

    When using pixelized source reconstruction, non-linear searches often converge on unphysical 'demagnified' solutions where the mass model is too low (or too high) and the source is reconstructed to look like the observed lensed image.

    To prevent this, PyAutoLens uses a PositionsLH likelihood term. This term penalizes mass models that do not trace specified multiple image positions within a given threshold in the source-plane. If the images trace closer together in the source-plane, the penalty decreases, guiding the search (e.g., Dynesty) toward the physical mass model.

    import autolens as al
    
    # 1. Define the locations of the multiple images in the image plane
    positions = al.Grid2DIrregular(
        grid=[(0.4, 1.6), (1.58, -0.35), (-0.43, -1.59), (-1.45, 0.2)]
    )
    
    # 2. Create the likelihood object with a threshold (e.g., 0.3")
    # This requires images to trace within 0.3" of each other in the source plane
    positions_likelihood = al.PositionsLH(positions=positions, threshold=0.3)
    
    # 3. Pass the likelihood to the Analysis object
    analysis = al.AnalysisImaging(
        dataset=dataset, 
        positions_likelihood_list=[positions_likelihood]
    )
  6. Select a dataset type for galaxy-scale lenses

    main

    If you are working at the Galaxy Scale, choose your analysis path based on your data type:

    • CCD Imaging: For image data from telescopes like Hubble or James Webb.
    • Interferometer: For radio/sub-mm interferometer data (e.g., ALMA).
    • Data Cubes: For spectral-line data cubes (e.g., ALMA CO cubes) where multiple frequency channels are fitted simultaneously with a shared lens model.
    • Point Sources: For strongly lensed point sources like lensed quasars or supernovae.
  7. Work with Interferometer and UV-plane data

    main

    For datasets from interferometers (e.g., ALMA), PyAutoLens provides structures to handle data in the uv-plane. This includes performing Fast Fourier Transforms (FFT) to map data between the image plane and the uv-plane.

    Key classes include:

    • Interferometer: The primary structure for interferometer datasets.
    • SimulatorInterferometer: Used for simulating interferometer data.
    • Visibilities: Represents the visibility data.
    • TransformerDFT: Performs Discrete Fourier Transforms.
    • TransformerNUFFT: Performs Non-Uniform Fast Fourier Transforms.
  8. Decompose galaxy light using Multi Gaussian Expansion (MGE)

    main

    MGE decomposes a galaxy's light into hundreds of 2D Gaussians. This is used for clean lens subtraction, modeling stellar mass tied to light, and capturing departures from elliptical symmetry.

    Implementation details:

    • The model is built from two bases of Gaussians (e.g., 30 Gaussians each).
    • The intensity parameters are solved for via inversion rather than sampling.

    For a full tutorial, see autolens_workspace/notebooks/features/multi_gaussian_expansion.ipynb.

  9. Model galaxies and gravitational lensing with Galaxy, Galaxies, and Tracer

    main

    PyAutoLens uses a hierarchy of objects to model galaxies and perform gravitational lensing ray-tracing:

    • Galaxy and Galaxies: Model individual galaxies or groups of galaxies with specific light and mass profiles at a defined redshift.
    • Tracer: Groups galaxies into redshift planes to perform multi-plane gravitational lensing ray-tracing. It computes lensed images, convergence, deflection angles, magnification, critical curves, and caustics.
    • MassField: Represents mass outside the modeled system (e.g., external shear, mass sheet, or external potential) at a specific redshift. It has no light profile.

    To include external mass in a lensing calculation, pass a list of MassField objects to a Tracer using the fields argument. Each field is placed in the plane corresponding to its own redshift.

    Note: Tracer(galaxies=..., fields=field) (passing a single field instead of a list) is also accepted, but tracer.fields will always be stored as a list.

  10. Advanced Model Composition Techniques

    main

    PyAutoLens supports several advanced modeling workflows:

    • Model Linking: During non-linear search chaining, the inferred model from one phase can be linked to the model of the next phase.
    • Across Datasets: When fitting multiple datasets, you can compose models where components are shared across datasets, but specific parameters are allowed to vary independently per dataset.
    • Relations: You can compose models where free parameters vary according to a user-specified function (e.g., defining intensity as a linear function of wavelength).
  11. Use the Redshift object for free-parameter modeling

    main

    When using PyAutoFit to perform model-fitting where the redshift of a galaxy is a free parameter, you must use the Redshift object. This is required because PyAutoFit requires all parameters to be instances of a Python class.

    For general PyAutoGalaxy usage where redshift is not a free parameter, the Redshift object is not required.

  12. Core concepts of PyAutoLens lens modeling

    main

    PyAutoLens uses an object-oriented approach to represent gravitational lens systems. The primary abstractions are:

    • Galaxy: Represents a single galaxy. It contains one or more LightProfiles (representing starlight distribution) and MassProfiles (representing mass distribution).
    • LightProfile & MassProfile: Analytic profiles used to define the 2D distribution of light and mass for a galaxy.
    • Plane: A grouping of Galaxy objects located at a specific redshift (distance) from the observer.
    • Tracer: An object that performs ray-tracing through multiple Planes using an astropy Cosmology.
    • Grid2D: An object used to compute strong lens quantities (like ray-tracing sightlines) by passing it through the system.
    • Pixelization: Used to reconstruct the background source's light using either a rectangular grid or a Voronoi mesh.
    • Point: Used for fitting lensed quasars and supernovae by utilizing observed positions, flux-ratios, and time-delays.

    These objects are extensible, allowing you to compose highly customized lensing systems.