From 541f2659bcc0f57ab7576711513ca607ec3dad97 Mon Sep 17 00:00:00 2001 From: Cordero Core <127983572+uwcdc@users.noreply.github.com> Date: Fri, 8 Dec 2023 11:30:39 -0800 Subject: [PATCH] docs: Update docs with Jupyter Book and NumPy docstrings (#110) This is a PR from https://github.com/uw-ssec/caustics/pull/25. This improved the documentation to use Jupyter Book instead of vanilla sphinx. Google docstring style has been updated to Numpy docstring style. --- Ref: https://github.com/uw-ssec/caustics/issues/11 Original author: https://github.com/sophietao127 - docs: Add tutorials - Added index.md to tutorials - Updated _config.yml and _toc.yml - docs: Lint docs and update notebooks --------- Co-authored-by: Landung 'Don' Setiawan --- .devcontainer/postBuild.sh | 2 +- .github/CONTRIBUTING.md | 12 +- .gitignore | 4 + .pre-commit-config.yaml | 9 +- .readthedocs.yml | 8 +- LICENSE | 2 +- README.md | 8 +- docs/requirements.txt | 5 +- docs/source/_config.yml | 47 ++ docs/source/_toc.yml | 21 + docs/source/citation.rst | 2 +- docs/source/conf.py | 187 ------- docs/source/contributing.rst | 2 +- docs/source/getting_started.rst | 2 +- docs/source/index.rst | 54 -- docs/source/install.rst | 2 +- docs/source/intro.md | 29 + docs/source/modules.rst | 2 +- docs/source/pull_notebooks.py | 18 - docs/source/tutorials.rst | 16 - docs/source/tutorials/BasicIntroduction.ipynb | 190 +++++++ .../source/tutorials/InvertLensEquation.ipynb | 145 +++++ docs/source/tutorials/LensZoo.ipynb | 517 ++++++++++++++++++ docs/source/tutorials/MultiplaneDemo.ipynb | 248 +++++++++ docs/source/tutorials/Playground.ipynb | 235 ++++++++ docs/source/tutorials/Simulators.ipynb | 279 ++++++++++ docs/source/tutorials/VisualizeCaustics.ipynb | 164 ++++++ docs/source/tutorials/index.md | 5 + noxfile.py | 1 - src/caustics/cosmology.py | 264 +++++---- src/caustics/data/hdf5dataset.py | 16 +- src/caustics/data/probes.py | 5 +- src/caustics/lenses/base.py | 362 +++++++----- src/caustics/lenses/epl.py | 149 +++-- src/caustics/lenses/external_shear.py | 93 ++-- src/caustics/lenses/mass_sheet.py | 43 +- src/caustics/lenses/multiplane.py | 117 ++-- src/caustics/lenses/nfw.py | 307 +++++++---- src/caustics/lenses/pixelated_convergence.py | 234 +++++--- src/caustics/lenses/point.py | 117 ++-- src/caustics/lenses/pseudo_jaffe.py | 168 ++++-- src/caustics/lenses/sie.py | 119 ++-- src/caustics/lenses/singleplane.py | 81 ++- src/caustics/lenses/sis.py | 90 +-- src/caustics/lenses/tnfw.py | 309 +++++++---- src/caustics/lenses/utils.py | 64 ++- src/caustics/light/base.py | 23 +- src/caustics/light/sersic.py | 92 ++-- src/caustics/namespace_dict.py | 12 +- src/caustics/parameter.py | 18 +- src/caustics/parametrized.py | 110 ++-- src/caustics/sims/lens_source.py | 109 ++-- src/caustics/sims/simulator.py | 2 +- src/caustics/tests.py | 9 +- src/caustics/utils.py | 308 +++++++---- tests/test_parametrized.py | 4 +- tests/test_sersic.py | 2 +- tests/test_simulator.py | 6 +- tests/utils.py | 2 +- 59 files changed, 4017 insertions(+), 1434 deletions(-) create mode 100644 docs/source/_config.yml create mode 100644 docs/source/_toc.yml delete mode 100644 docs/source/conf.py delete mode 100644 docs/source/index.rst create mode 100644 docs/source/intro.md delete mode 100644 docs/source/pull_notebooks.py delete mode 100644 docs/source/tutorials.rst create mode 100644 docs/source/tutorials/BasicIntroduction.ipynb create mode 100644 docs/source/tutorials/InvertLensEquation.ipynb create mode 100644 docs/source/tutorials/LensZoo.ipynb create mode 100644 docs/source/tutorials/MultiplaneDemo.ipynb create mode 100644 docs/source/tutorials/Playground.ipynb create mode 100644 docs/source/tutorials/Simulators.ipynb create mode 100644 docs/source/tutorials/VisualizeCaustics.ipynb create mode 100644 docs/source/tutorials/index.md diff --git a/.devcontainer/postBuild.sh b/.devcontainer/postBuild.sh index 743f496f..92c47745 100644 --- a/.devcontainer/postBuild.sh +++ b/.devcontainer/postBuild.sh @@ -1,5 +1,5 @@ # For writing commands that will be executed after the container is created -# Installs `caustic` as local library without resolving dependencies (--no-deps) +# Installs `caustics` as local library without resolving dependencies (--no-deps) python3 -m pip install -e /workspaces/caustics --no-deps python3 -m pip install -e ".[dev]" diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md index 7453e0b6..4b247715 100644 --- a/.github/CONTRIBUTING.md +++ b/.github/CONTRIBUTING.md @@ -11,8 +11,9 @@ install `nox`: ## Codespaces -Nox is pre-installed in the Codespaces environment. So, after activating a Codespace, -you can just open the terminal and run `nox` to run all the checks and tests. +Nox is pre-installed in the Codespaces environment. So, after activating a +Codespace, you can just open the terminal and run `nox` to run all the checks +and tests. ## Local @@ -32,10 +33,9 @@ brew install nox ### What is it? -`nox` is a command-line tool that automates testing in multiple Python environments, -similar to tox. Unlike tox, -Nox uses a standard Python file for configuration, -you can find this configuration in [`noxfile.py`](../noxfile.py). +`nox` is a command-line tool that automates testing in multiple Python +environments, similar to tox. Unlike tox, Nox uses a standard Python file for +configuration, you can find this configuration in [`noxfile.py`](../noxfile.py). ### How do I use it? diff --git a/.gitignore b/.gitignore index 12601fd5..ab13c8bb 100644 --- a/.gitignore +++ b/.gitignore @@ -137,3 +137,7 @@ package.json # version _version.py + +# Jupyter book +conf.py +caustics*.rst diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c93e201b..7cc05c4a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,8 +1,7 @@ exclude: | (?x)^( tests/utils.py | - .devcontainer/ | - docs/source/ + .devcontainer/ ) ci: @@ -77,6 +76,12 @@ repos: hooks: - id: shellcheck + - repo: https://github.com/kynan/nbstripout + rev: 0.6.1 + hooks: + - id: nbstripout + args: [--extra-keys=metadata.kernelspec metadata.language_info.version] + - repo: local hooks: - id: disallow-caps diff --git a/.readthedocs.yml b/.readthedocs.yml index 05041f5a..e1ab16a6 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -25,10 +25,16 @@ build: python: "3.9" apt_packages: - pandoc # Specify pandoc to be installed via apt-get + - graphviz # Specify graphviz to be installed via apt-get + jobs: + pre_build: + # Generate the Sphinx configuration for this Jupyter Book so it builds. + - "jupyter-book config sphinx docs/source/" + - "sphinx-apidoc -f -o docs/source/ src/caustics/" python: install: - - requirements: requirements.txt # Path to your requirements.txt file - requirements: docs/requirements.txt # Path to your requirements.txt file + - requirements: requirements.txt # Path to your requirements.txt file - method: pip path: . # Install the package itself diff --git a/LICENSE b/LICENSE index 3f4c46db..d3a9c27c 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) [2023] [caustic authors] +Copyright (c) [2023] [caustics authors] Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index 4a8da577..132978a9 100644 --- a/README.md +++ b/README.md @@ -26,15 +26,15 @@ pip install caustics ## Documentation -Please see our [documentation page](https://caustics.readthedocs.io/en/latest/) for -more detailed information. +Please see our [documentation page](https://caustics.readthedocs.io/en/latest/) +for more detailed information. ## Contribution We welcome contributions from collaborators and researchers interested in our work. If you have improvements, suggestions, or new findings to share, please -submit an issue or pull request. Your contributions help advance our research and analysis -efforts. +submit an issue or pull request. Your contributions help advance our research +and analysis efforts. To get started with your development (or fork), click the "Open with GitHub Codespaces" button below to launch a fully configured development environment diff --git a/docs/requirements.txt b/docs/requirements.txt index 80e683f8..3aee37f7 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -1,4 +1,5 @@ -nbsphinx +jupyter-book sphinx sphinx_rtd_theme -wheel +matplotlib +ipywidgets diff --git a/docs/source/_config.yml b/docs/source/_config.yml new file mode 100644 index 00000000..3920dac0 --- /dev/null +++ b/docs/source/_config.yml @@ -0,0 +1,47 @@ +# Book settings +# Learn more at https://jupyterbook.org/customize/config.html + +title: Caustics +author: Ciela Institute +logo: ../../media/caustics_logo.png + +# Force re-execution of notebooks on each build. +# See https://jupyterbook.org/content/execute.html +execute: + execute_notebooks: force + allow_errors: false + # Per-cell notebook execution limit (seconds) + timeout: 10 + +# Define the name of the latex output file for PDF builds +latex: + latex_documents: + targetname: book.tex + +# Add a bibtex file so that we can create citations +# bibtex_bibfiles: +# - references.bib + +# Information about where the book exists on the web +repository: + url: https://github.com/Ciela-Institute/caustics # Online location of your book + path_to_book: docs/source # Optional path to your book, relative to the repository root + branch: main # Which branch of the repository should be used when creating links (optional) + +# Add GitHub buttons to your book +# See https://jupyterbook.org/customize/config.html#add-a-link-to-your-repository +html: + favicon: ../../media/caustics_favicon.ico + use_issues_button: true + use_repository_button: true + +sphinx: + extra_extensions: + - "sphinx.ext.autodoc" + - "sphinx.ext.autosummary" + - "sphinx.ext.napoleon" + - "sphinx.ext.doctest" + - "sphinx.ext.coverage" + - "sphinx.ext.mathjax" + - "sphinx.ext.ifconfig" + - "sphinx.ext.viewcode" diff --git a/docs/source/_toc.yml b/docs/source/_toc.yml new file mode 100644 index 00000000..a93ba2cd --- /dev/null +++ b/docs/source/_toc.yml @@ -0,0 +1,21 @@ +# Table of contents +# Learn more at https://jupyterbook.org/customize/toc.html + +format: jb-book +root: intro +chapters: + - file: getting_started + - file: install + - file: tutorials/index + sections: + - file: tutorials/BasicIntroduction + - file: tutorials/LensZoo + - file: tutorials/VisualizeCaustics + - file: tutorials/MultiplaneDemo + - file: tutorials/InvertLensEquation + - file: tutorials/Simulators + - file: tutorials/Playground + - file: contributing + - file: citation + - file: license + - file: modules diff --git a/docs/source/citation.rst b/docs/source/citation.rst index 0d200348..bb0c021e 100644 --- a/docs/source/citation.rst +++ b/docs/source/citation.rst @@ -3,4 +3,4 @@ Citation ======== -Paper comming soon! We will put all the citation information here when its ready. +Paper coming soon! We will put all the citation information here when its ready. diff --git a/docs/source/conf.py b/docs/source/conf.py deleted file mode 100644 index 2011b358..00000000 --- a/docs/source/conf.py +++ /dev/null @@ -1,187 +0,0 @@ -# -*- coding: utf-8 -*- -# -# Configuration file for the Sphinx documentation builder. -# -# This file does only contain a selection of the most common options. For a -# full list see the documentation: -# http://www.sphinx-doc.org/en/master/config - -# -- Path setup -------------------------------------------------------------- - -# If extensions (or modules to document with autodoc) are in another directory, -# add these directories to sys.path here. If the directory is relative to the -# documentation root, use os.path.abspath to make it absolute, like shown here. -# -from caustics import __version__ -# sys.path.insert(0, os.path.abspath('../src')) - - -# -- Project information ----------------------------------------------------- - -project = "caustics" -copyright = "2023, Ciela Institute" -author = "Ciela Institute" - -# The short X.Y version -version = __version__ -# The full version, including alpha/beta/rc tags -release = version - - -# -- General configuration --------------------------------------------------- - -# If your documentation needs a minimal Sphinx version, state it here. -# -# needs_sphinx = '1.0' - -# Add any Sphinx extension module names here, as strings. They can be -# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom -# ones. -extensions = [ - "nbsphinx", - "sphinx.ext.autodoc", - "sphinx.ext.autosummary", - "sphinx.ext.napoleon", - "sphinx.ext.doctest", - "sphinx.ext.coverage", - "sphinx.ext.mathjax", - "sphinx.ext.ifconfig", - "sphinx.ext.viewcode", -] - -# Add any paths that contain templates here, relative to this directory. -templates_path = ["_templates"] - -# The suffix(es) of source filenames. -# You can specify multiple suffix as a list of string: -# -# source_suffix = ['.rst', '.md'] -source_suffix = ".rst" - -# The master toctree document. -master_doc = "index" - -# The language for content autogenerated by Sphinx. Refer to documentation -# for a list of supported languages. -# -# This is also used if you do content translation via gettext catalogs. -# Usually you set "language" from the command line for these cases. -language = "en" - -# List of patterns, relative to source directory, that match files and -# directories to ignore when looking for source files. -# This pattern also affects html_static_path and html_extra_path. -exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] - -# The name of the Pygments (syntax highlighting) style to use. -pygments_style = "sphinx" - - -# -- Options for HTML output ------------------------------------------------- - -# The theme to use for HTML and HTML Help pages. See the documentation for -# a list of builtin themes. -# -html_theme = "sphinx_rtd_theme" -# html_theme = 'alabaster' - -# Theme options are theme-specific and customize the look and feel of a theme -# further. For a list of options available for each theme, see the -# documentation. -# -# html_theme_options = {} - -# Add any paths that contain custom static files (such as style sheets) here, -# relative to this directory. They are copied after the builtin static files, -# so a file named "default.css" will overwrite the builtin "default.css". -html_static_path = ["_static"] - -# Custom sidebar templates, must be a dictionary that maps document names -# to template names. -# -# The default sidebars (for documents that don't match any pattern) are -# defined by theme itself. Builtin themes are using these templates by -# default: ``['localtoc.html', 'relations.html', 'sourcelink.html', -# 'searchbox.html']``. -# -# html_sidebars = {} -html_favicon = "../media/caustics_favicon.ico" - - -# -- Options for HTMLHelp output --------------------------------------------- - -# Output file base name for HTML help builder. -htmlhelp_basename = "causticsdoc" - - -# -- Options for LaTeX output ------------------------------------------------ - -latex_elements = { - # The paper size ('letterpaper' or 'a4paper'). - # - # 'papersize': 'letterpaper', - # The font size ('10pt', '11pt' or '12pt'). - # - # 'pointsize': '10pt', - # Additional stuff for the LaTeX preamble. - # - # 'preamble': '', - # Latex figure (float) alignment - # - # 'figure_align': 'htbp', -} - -# Grouping the document tree into LaTeX files. List of tuples -# (source start file, target name, title, -# author, documentclass [howto, manual, or own class]). -latex_documents = [ - (master_doc, "caustics.tex", "caustics Documentation", "Ciela Institute", "manual"), -] - - -# -- Options for manual page output ------------------------------------------ - -# One entry per manual page. List of tuples -# (source start file, name, description, authors, manual section). -man_pages = [(master_doc, "caustics", "caustics Documentation", [author], 1)] - - -# -- Options for Texinfo output ---------------------------------------------- - -# Grouping the document tree into Texinfo files. List of tuples -# (source start file, target name, title, author, -# dir menu entry, description, category) -texinfo_documents = [ - ( - master_doc, - "caustics", - "caustics Documentation", - author, - "caustics", - "The gravitational lensing pipeline of the future.", - "Miscellaneous", - ), -] - - -# -- Options for Epub output ------------------------------------------------- - -# Bibliographic Dublin Core info. -epub_title = project - -# The unique identifier of the text. This can be a ISBN number -# or the project homepage. -# -# epub_identifier = '' - -# A unique identification for the text. -# -# epub_uid = '' - -# A list of files that should not be packed into the epub file. -epub_exclude_files = ["search.html"] - - -# -- Extension configuration ------------------------------------------------- -# -- Options for nbsphinx -------------------------------------------------- -nbsphinx_execute = "never" diff --git a/docs/source/contributing.rst b/docs/source/contributing.rst index fbed92c6..4929d38b 100644 --- a/docs/source/contributing.rst +++ b/docs/source/contributing.rst @@ -9,7 +9,7 @@ Thanks for helping make caustics better! Here you will learn the full process ne Create An Issue --------------- -Before actually writing any code, its best to create an issue on the GitHub. Describe the issue in detail and let us know the desired solution. Here it will be possible to address concerns (maybe its aready solved and just not yet documented) and plan out the best solution. We may also assign someone to work on it if that seems better. Note that submitting an issue is a contribution to caustics, we appreciate your ideas! Still, if after discussion it seems that the problem does need some work and you're the person to do it, then we can move on to the next steps. +Before actually writing any code, its best to create an issue on the GitHub. Describe the issue in detail and let us know the desired solution. Here it will be possible to address concerns (maybe its already solved and just not yet documented) and plan out the best solution. We may also assign someone to work on it if that seems better. Note that submitting an issue is a contribution to caustics, we appreciate your ideas! Still, if after discussion it seems that the problem does need some work and you're the person to do it, then we can move on to the next steps. 1. Navigate to the **Issues** tab of the GitHub repository. 2. Click on **New Issue**. diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index 02d7db8d..513fef20 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -17,4 +17,4 @@ We have created a repository of tutorial Jupyter notebooks that can help initiat Read The Docs ------------- -Docs for all the main functions in caustics are avaialble at :doc:`caustics` at varying degrees of completeness. Further development of the docs is always ongoing. +Docs for all the main functions in caustics are available at :doc:`caustics` at varying degrees of completeness. Further development of the docs is always ongoing. diff --git a/docs/source/index.rst b/docs/source/index.rst deleted file mode 100644 index 04db23c3..00000000 --- a/docs/source/index.rst +++ /dev/null @@ -1,54 +0,0 @@ -.. caustics documentation master file, created by - sphinx-quickstart on Thu Jun 1 09:33:32 2023. - You can adapt this file completely to your liking, but it should at least - contain the root `toctree` directive. - -.. image:: https://github.com/Ciela-Institute/caustics/blob/main/media/caustics_logo.png?raw=true - :width: 100 % - :target: https://github.com/Ciela-Institute/caustics - -|br| - -Welcome to Caustics' documentation! -=================================== - -The lensing pipeline of the future: GPU-accelerated, automatically-differentiable, highly modular and extensible. - -.. note:: - - Caustics is in its early development phase. This means the API will change with time. These changes are a good thing, but they can be annoying. Watch the version numbers, when we get to 1.0.0 that will be the first stable release! - -Intallation -=========== - -The easiest way to install is to make a new virtual environment then run:: - - pip install caustics - -this will install all the required libraries and then install caustics and you are ready to go! You can check out the tutorials afterwards to see some of caustics's capabilities. If you want to help out with building the caustics code base check out the developer installation instructions instead. - -Read The Docs -============= - -.. toctree:: - :maxdepth: 2 - :caption: Contents: - - getting_started.rst - install.rst - tutorials.rst - contributing.rst - citation.rst - license.rst - -Indices and tables -================== - -.. toctree:: - :maxdepth: 1 - - modules.rst - -* :ref:`genindex` -* :ref:`modindex` -* :ref:`search` diff --git a/docs/source/install.rst b/docs/source/install.rst index 4d67dc7a..874a8a3d 100644 --- a/docs/source/install.rst +++ b/docs/source/install.rst @@ -19,7 +19,7 @@ First clone the repo with:: git clone git@github.com:Ciela-Institute/caustics.git -this will create a directory `caustics` wherever you ran the command. Next go into the directory and install in developer mode:: +this will create a directory ``caustics`` wherever you ran the command. Next go into the directory and install in developer mode:: pip install -e ".[dev]" diff --git a/docs/source/intro.md b/docs/source/intro.md new file mode 100644 index 00000000..c36a4a2a --- /dev/null +++ b/docs/source/intro.md @@ -0,0 +1,29 @@ +# Welcome to Caustics’ documentation! + +!["caustics_logo"](https://github.com/Ciela-Institute/caustics/blob/main/media/caustics_logo.png?raw=true) + +The lensing pipeline of the future: GPU-accelerated, +automatically-differentiable, highly modular and extensible. + +```{note} +Caustics is in its early development phase. This means the API will change with time. These changes are a good thing, but they can be annoying. Watch the version numbers, when we get to 1.0.0 that will be the first stable release! +``` + +## Installation + +The easiest way to install is to make a new virtual environment then run: + +```console +pip install caustics +``` + +this will install all the required libraries and then install caustics and you +are ready to go! You can check out the tutorials afterwards to see some of +caustics' capabilities. If you want to help out with building the caustics code +base check out the developer installation instructions instead. + +### Contents + +```{tableofcontents} + +``` diff --git a/docs/source/modules.rst b/docs/source/modules.rst index 446b67c7..d6f97e96 100644 --- a/docs/source/modules.rst +++ b/docs/source/modules.rst @@ -1,5 +1,5 @@ caustics -======= +======== .. toctree:: :maxdepth: 4 diff --git a/docs/source/pull_notebooks.py b/docs/source/pull_notebooks.py deleted file mode 100644 index cc928977..00000000 --- a/docs/source/pull_notebooks.py +++ /dev/null @@ -1,18 +0,0 @@ -import requests - -tutorials = [ - "https://raw.github.com/Ciela-Institute/caustics-tutorials/main/tutorials/BasicIntroduction.ipynb", - "https://raw.github.com/Ciela-Institute/caustics-tutorials/main/tutorials/LensZoo.ipynb", - "https://raw.github.com/Ciela-Institute/caustics-tutorials/main/tutorials/VisualizeCaustics.ipynb", - "https://raw.github.com/Ciela-Institute/caustics-tutorials/main/tutorials/MultiplaneDemo.ipynb", - "https://raw.github.com/Ciela-Institute/caustics-tutorials/main/tutorials/InvertLensEquation.ipynb", -] -for url in tutorials: - try: - R = requests.get(url) - with open(url[url.rfind("/") + 1 :], "w") as f: - f.write(R.text) - except: - print( - f"WARNING: couldn't find tutorial: {url[url.rfind('/')+1:]} check internet conection" - ) diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst deleted file mode 100644 index e9dea14b..00000000 --- a/docs/source/tutorials.rst +++ /dev/null @@ -1,16 +0,0 @@ -========= -Tutorials -========= - -Here you will find the jupyter notebook tutorials. It is recommended -that you go through the tutorials yourself, but for quick reference a -version of each tutorial is available here. - -.. toctree:: - :maxdepth: 1 - - BasicIntroduction - LensZoo - VisualizeCaustics - MultiplaneDemo - InvertLensEquation diff --git a/docs/source/tutorials/BasicIntroduction.ipynb b/docs/source/tutorials/BasicIntroduction.ipynb new file mode 100644 index 00000000..67c288a3 --- /dev/null +++ b/docs/source/tutorials/BasicIntroduction.ipynb @@ -0,0 +1,190 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Welcome to Caustic!\n", + "\n", + "In need of a differentiable strong gravitational lensing simulation package? Look no further! We have all your lensing simulator needs. In this tutorial we will cover the basics of caustic and how to get going making your own lensing configurations. Caustic is easy to use and very powerful, you will get to see some of that power here, but there will be more notebooks which demo specific use cases." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2\n", + "\n", + "import torch\n", + "from torch.nn.functional import avg_pool2d\n", + "import matplotlib.pyplot as plt\n", + "from astropy.io import fits\n", + "import numpy as np\n", + "\n", + "import caustics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Specify the image/cosmology parameters\n", + "n_pix = 100\n", + "res = 0.05\n", + "upsample_factor = 2\n", + "fov = res * n_pix\n", + "thx, thy = caustics.get_meshgrid(\n", + " res / upsample_factor,\n", + " upsample_factor * n_pix,\n", + " upsample_factor * n_pix,\n", + " dtype=torch.float32,\n", + ")\n", + "z_l = torch.tensor(0.5, dtype=torch.float32)\n", + "z_s = torch.tensor(1.5, dtype=torch.float32)\n", + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology.to(dtype=torch.float32)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Simulating an SIE lens\n", + "\n", + "Here we will demo the very basics of lensing with a classic `SIE` lens model. We will see what it takes to make an `SIE` model, lens a backgorund `Sersic` source, and sample some examples in a simulator. Caustic simulators can generalize to very complex scenarios. In these cases there can be a lot of parameters moving through the simulator, and the order/number of parameters may change depending on what lens or source is being used. To streamline this process, caustic impliments a class called `Parametrized` which has some knowledge of the parameters moving through it, this way it can keep track of everything for you. For this to work, you must put the parameters into a `Packed` object which it can recognize, each sub function can then unpack the parameters it needs. Below we will show some examples of what this looks like." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# demo simulator with sersic source, SIE lens. then sample some examples. demo the model graph\n", + "\n", + "\n", + "class Simple_Sim(caustics.Simulator):\n", + " def __init__(\n", + " self,\n", + " lens,\n", + " src,\n", + " z_s=None,\n", + " name: str = \"sim\",\n", + " ):\n", + " super().__init__(name) # need this so `Parametrized` can do its magic\n", + "\n", + " # These are the lens and source objects to keep track of\n", + " self.lens = lens\n", + " self.src = src\n", + "\n", + " # Here we can add a parameter to the simulator, in this case it is `z_s` which we will need later\n", + " self.add_param(\"z_s\", z_s)\n", + "\n", + " def forward(self, params): # define the forward model\n", + " # Here the simulator unpacks the parameter it needs\n", + " z_s = self.unpack(params)\n", + "\n", + " # Note this is very similar to before, except the packed up `x` is all the raytrace function needs to work\n", + " bx, by = self.lens.raytrace(thx, thy, z_s, params)\n", + " mu_fine = self.src.brightness(bx, by, params)\n", + "\n", + " # We return the sampled brightness at each pixel location\n", + " return avg_pool2d(mu_fine.squeeze()[None, None], upsample_factor)[0, 0]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sie = caustics.lenses.SIE(cosmology, name=\"sie\")\n", + "src = caustics.light.Sersic(name=\"src\")\n", + "\n", + "sim = Simple_Sim(sie, src, torch.tensor(0.8))\n", + "\n", + "sim.get_graph(True, True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(sim)\n", + "print(sie)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Reading the x_keys above we can input the parameters that we would like the simulator to evaluate\n", + "x = torch.tensor(\n", + " [\n", + " z_l.item(), # sie z_l\n", + " 0.7, # sie x0\n", + " 0.13, # sie y0\n", + " 0.4, # sie q\n", + " np.pi / 5, # sie phi\n", + " 1.0, # sie b\n", + " 0.2, # src x0\n", + " 0.5, # src y0\n", + " 0.5, # src q\n", + " -np.pi / 4, # src phi\n", + " 1.5, # src n\n", + " 2.5, # src Re\n", + " 1.0, # src Ie\n", + " ]\n", + ")\n", + "plt.imshow(sim(x), origin=\"lower\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Where to go next?\n", + "\n", + "The caustic tutorials are generally short and to the point, that way you can idenfity what you want and jump right to some useful code that demo's the particular problem you face. Below is a list of caustic tutorials and a quick description of what you will learn in each one::\n", + "\n", + "- `LensZoo`: here you can see all the built-in lens mass distributions in `caustic` and how they distort the same background Seric source.\n", + "- `Playground`: here we demo the main visualizations of a lensing system (deflection angles, convergence, potential, time delay, magnification) in an interactive display so you can change the parameters by hand and see how the visuals change!\n", + "- `VisualizeCaustics`: here you can see how to find and display caustics, a must when using `caustic`!\n", + "- `Simulators`: here we describe the powerful simulator framework and how it can be used to quickly swap models, parameters, and other features and turn a complex forward model into a simple function.\n", + "- `InvertLensEquation`: here we demo forward ray tracing in `caustic` the process of mapping from the source plane to the image plane." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/source/tutorials/InvertLensEquation.ipynb b/docs/source/tutorials/InvertLensEquation.ipynb new file mode 100644 index 00000000..db75ab0c --- /dev/null +++ b/docs/source/tutorials/InvertLensEquation.ipynb @@ -0,0 +1,145 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "b7d08f1d", + "metadata": {}, + "source": [ + "# Inverting the Lens Equation\n", + "\n", + "The lens equation $\\vec{\\beta} = \\vec{\\theta} - \\vec{\\alpha}(\\vec{\\theta})$ allows us to find a point in the source plane given a point in the image plane. However, sometimes we know a point in the source plane and would like to see where it ends up in the image plane. This is not easy to do since a point in the source plane may map to multiple locations in the image plane. There is no closed form function to invert the lens equation, in large part because the deflection angle $\\vec{\\alpha}$ depends on the position in the image plane $\\vec{\\theta}$. To invert the lens equation, we will need to rely on optimization and a little luck to find all the images for a given source plane point. Below we will demonstrate how this is done in caustic!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4027aaf9", + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2\n", + "\n", + "from functools import partial\n", + "\n", + "import torch\n", + "from torch.nn.functional import avg_pool2d\n", + "import matplotlib.pyplot as plt\n", + "from ipywidgets import interact\n", + "from astropy.io import fits\n", + "import numpy as np\n", + "from time import process_time as time\n", + "\n", + "import caustics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2118e1c1", + "metadata": {}, + "outputs": [], + "source": [ + "# initialization stuff for an SIE lens\n", + "\n", + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology.to(dtype=torch.float32)\n", + "n_pix = 100\n", + "res = 0.05\n", + "upsample_factor = 2\n", + "fov = res * n_pix\n", + "thx, thy = caustics.get_meshgrid(\n", + " res / upsample_factor,\n", + " upsample_factor * n_pix,\n", + " upsample_factor * n_pix,\n", + " dtype=torch.float32,\n", + ")\n", + "z_l = torch.tensor(0.5, dtype=torch.float32)\n", + "z_s = torch.tensor(1.5, dtype=torch.float32)\n", + "lens = caustics.SIE(\n", + " cosmology=cosmology,\n", + " name=\"sie\",\n", + " z_l=z_l,\n", + " x0=torch.tensor(0.0),\n", + " y0=torch.tensor(0.0),\n", + " q=torch.tensor(0.4),\n", + " phi=torch.tensor(np.pi / 5),\n", + " b=torch.tensor(1.0),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "98e46aa1", + "metadata": {}, + "outputs": [], + "source": [ + "# Point in the source plane\n", + "sp_x = torch.tensor(0.2)\n", + "sp_y = torch.tensor(0.2)\n", + "\n", + "# Points in image plane\n", + "x, y = lens.forward_raytrace(sp_x, sp_y, z_s)\n", + "\n", + "# Raytrace to check\n", + "bx, by = lens.raytrace(x, y, z_s)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "eb73147c", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots()\n", + "\n", + "A = lens.jacobian_lens_equation(thx, thy, z_s)\n", + "detA = torch.linalg.det(A)\n", + "CS = ax.contour(thx, thy, detA, levels=[0.0], colors=\"b\", zorder=1)\n", + "# Get the path from the matplotlib contour plot of the critical line\n", + "paths = CS.collections[0].get_paths()\n", + "caustic_paths = []\n", + "for path in paths:\n", + " # Collect the path into a descrete set of points\n", + " vertices = path.interpolated(5).vertices\n", + " x1 = torch.tensor(list(float(vs[0]) for vs in vertices))\n", + " x2 = torch.tensor(list(float(vs[1]) for vs in vertices))\n", + " # raytrace the points to the source plane\n", + " y1, y2 = lens.raytrace(x1, x2, z_s)\n", + "\n", + " # Plot the caustic\n", + " ax.plot(y1, y2, color=\"r\", zorder=1)\n", + "ax.scatter(x, y, color=\"b\", label=\"forward raytrace\", zorder=10)\n", + "ax.scatter(bx, by, color=\"r\", marker=\"x\", label=\"source plane\", zorder=9)\n", + "ax.scatter([sp_x.item()], [sp_y.item()], color=\"g\", label=\"true pos\", zorder=8)\n", + "plt.legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a0aa515f", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/tutorials/LensZoo.ipynb b/docs/source/tutorials/LensZoo.ipynb new file mode 100644 index 00000000..d50e3f4a --- /dev/null +++ b/docs/source/tutorials/LensZoo.ipynb @@ -0,0 +1,517 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "e4054d16", + "metadata": {}, + "source": [ + "# A Menagerie of Lenses\n", + "\n", + "Here we have a quick visual demo of every type of lens in caustic. This is a good way to pick out what you need and quickly copy paste. For all of these lenses we have placed a Sersic source with the same parameters as the background, that way you can visualize the differences between them." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "beeb58fa", + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2\n", + "\n", + "import torch\n", + "from torch.nn.functional import avg_pool2d\n", + "import matplotlib.pyplot as plt\n", + "from ipywidgets import interact\n", + "from astropy.io import fits\n", + "import numpy as np\n", + "from time import process_time as time\n", + "\n", + "import caustics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "72cbde6d", + "metadata": {}, + "outputs": [], + "source": [ + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology.to(dtype=torch.float32)\n", + "z_s = torch.tensor(1.0)\n", + "base_sersic = caustics.light.Sersic(\n", + " name=\"sersic\",\n", + " x0=torch.tensor(0.1),\n", + " y0=torch.tensor(0.1),\n", + " q=torch.tensor(0.6),\n", + " phi=torch.tensor(np.pi / 3),\n", + " n=torch.tensor(2.0),\n", + " Re=torch.tensor(1.0),\n", + " Ie=torch.tensor(1.0),\n", + ")\n", + "n_pix = 100\n", + "res = 0.05\n", + "upsample_factor = 2\n", + "fov = res * n_pix\n", + "thx, thy = caustics.get_meshgrid(\n", + " res / upsample_factor,\n", + " upsample_factor * n_pix,\n", + " upsample_factor * n_pix,\n", + " dtype=torch.float32,\n", + ")\n", + "z_l = torch.tensor(0.5, dtype=torch.float32)\n", + "\n", + "plt.imshow(np.log10(base_sersic.brightness(thx, thy).numpy()), origin=\"lower\")\n", + "plt.gca().axis(\"off\")\n", + "plt.title(\"Base Sersic\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "aa183e3d", + "metadata": {}, + "source": [ + "## Point (Point)\n", + "\n", + "The simplest lens, an infinitely small point of mass (did someone say black holes?)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4b4b2faa", + "metadata": {}, + "outputs": [], + "source": [ + "lens = caustics.lenses.Point(\n", + " cosmology=cosmology,\n", + " name=\"point\",\n", + " x0=torch.tensor(0.0),\n", + " y0=torch.tensor(0.0),\n", + " th_ein=torch.tensor(1.0),\n", + ")\n", + "sim = caustics.sims.Lens_Source(\n", + " lens=lens,\n", + " source=base_sersic,\n", + " pixelscale=res,\n", + " pixels_x=n_pix,\n", + " z_s=z_s,\n", + " upsample_factor=2,\n", + ")\n", + "fig, axarr = plt.subplots(1, 2, figsize=(14, 7))\n", + "# axarr[0].imshow(np.log10(convergence.numpy()), origin = \"lower\")\n", + "axarr[0].axis(\"off\")\n", + "axarr[0].set_title(\"Point Convergence not defined\")\n", + "axarr[1].imshow(np.log10(sim([z_l]).numpy()), origin=\"lower\")\n", + "axarr[1].axis(\"off\")\n", + "axarr[1].set_title(\"Lensed Sersic\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "2ca96b29", + "metadata": {}, + "source": [ + "## Singular Isothermal Sphere (SIS)\n", + "\n", + "An SIS is a mass distribution represented by a constant temperature velocity dispersion of masses. Alternatively, a constant temperature gas in a spherical gravitational potential." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6d563d58", + "metadata": {}, + "outputs": [], + "source": [ + "lens = caustics.lenses.SIS(\n", + " cosmology=cosmology,\n", + " name=\"sis\",\n", + " x0=torch.tensor(0.0),\n", + " y0=torch.tensor(0.0),\n", + " th_ein=torch.tensor(1.0),\n", + ")\n", + "sim = caustics.sims.Lens_Source(\n", + " lens=lens,\n", + " source=base_sersic,\n", + " pixelscale=res,\n", + " pixels_x=n_pix,\n", + " z_s=z_s,\n", + " upsample_factor=2,\n", + ")\n", + "fig, axarr = plt.subplots(1, 2, figsize=(14, 7))\n", + "convergence = avg_pool2d(\n", + " lens.convergence(thx, thy, z_s, z_l).squeeze()[None, None], upsample_factor\n", + ").squeeze()\n", + "axarr[0].imshow(np.log10(convergence.numpy()), origin=\"lower\")\n", + "axarr[0].axis(\"off\")\n", + "axarr[0].set_title(\"SIS Convergence\")\n", + "axarr[1].imshow(np.log10(sim([z_l]).numpy()), origin=\"lower\")\n", + "axarr[1].axis(\"off\")\n", + "axarr[1].set_title(\"Lensed Sersic\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "e09f874d", + "metadata": {}, + "source": [ + "## Singular Isothermal Ellipsoid (SIE)\n", + "\n", + "The SIE is just like the SIS except it has been compressed along one axis." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "63c78c9a", + "metadata": {}, + "outputs": [], + "source": [ + "lens = caustics.lenses.SIE(\n", + " cosmology=cosmology,\n", + " name=\"sie\",\n", + " x0=torch.tensor(0.0),\n", + " y0=torch.tensor(0.0),\n", + " q=torch.tensor(0.6),\n", + " phi=torch.tensor(np.pi / 2),\n", + " b=torch.tensor(1.0),\n", + ")\n", + "sim = caustics.sims.Lens_Source(\n", + " lens=lens,\n", + " source=base_sersic,\n", + " pixelscale=res,\n", + " pixels_x=n_pix,\n", + " z_s=z_s,\n", + " upsample_factor=2,\n", + ")\n", + "fig, axarr = plt.subplots(1, 2, figsize=(14, 7))\n", + "convergence = avg_pool2d(\n", + " lens.convergence(thx, thy, z_s, z_l).squeeze()[None, None], upsample_factor\n", + ").squeeze()\n", + "axarr[0].imshow(np.log10(convergence.numpy()), origin=\"lower\")\n", + "axarr[0].axis(\"off\")\n", + "axarr[0].set_title(\"SIE Convergence\")\n", + "axarr[1].imshow(np.log10(sim([z_l]).numpy()), origin=\"lower\")\n", + "axarr[1].axis(\"off\")\n", + "axarr[1].set_title(\"Lensed Sersic\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "794060c3", + "metadata": {}, + "source": [ + "## Elliptical Power Law (EPL)\n", + "\n", + "This is a power law mass distribution with an elliptical isodensity contour." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "557bad9f", + "metadata": {}, + "outputs": [], + "source": [ + "lens = caustics.lenses.EPL(\n", + " cosmology=cosmology,\n", + " name=\"epl\",\n", + " x0=torch.tensor(0.0),\n", + " y0=torch.tensor(0.0),\n", + " q=torch.tensor(0.6),\n", + " phi=torch.tensor(np.pi / 2),\n", + " b=torch.tensor(1.0),\n", + " t=torch.tensor(0.5),\n", + ")\n", + "sim = caustics.sims.Lens_Source(\n", + " lens=lens,\n", + " source=base_sersic,\n", + " pixelscale=res,\n", + " pixels_x=n_pix,\n", + " z_s=z_s,\n", + " upsample_factor=2,\n", + ")\n", + "fig, axarr = plt.subplots(1, 2, figsize=(14, 7))\n", + "convergence = avg_pool2d(\n", + " lens.convergence(thx, thy, z_s, z_l).squeeze()[None, None], upsample_factor\n", + ").squeeze()\n", + "axarr[0].imshow(np.log10(convergence.numpy()), origin=\"lower\")\n", + "axarr[0].axis(\"off\")\n", + "axarr[0].set_title(\"EPL Convergence\")\n", + "axarr[1].imshow(np.log10(sim([z_l]).numpy()), origin=\"lower\")\n", + "axarr[1].axis(\"off\")\n", + "axarr[1].set_title(\"Lensed Sersic\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "1aa4a0b4", + "metadata": {}, + "source": [ + "## Navarro Frenk White profile (NFW)\n", + "\n", + "The NFW profile is a classic mass profile that approximates the mass distribution of halos in large dark matter simulations.\n", + "\n", + "$$\\rho(r) = \\frac{\\rho_0}{\\frac{r}{r_s}\\left(1 + \\frac{r}{r_s}\\right)^2}$$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bb209aba", + "metadata": {}, + "outputs": [], + "source": [ + "lens = caustics.lenses.NFW(\n", + " cosmology=cosmology,\n", + " name=\"nfw\",\n", + " x0=torch.tensor(0.0),\n", + " y0=torch.tensor(0.0),\n", + " m=torch.tensor(1e13),\n", + " c=torch.tensor(20.0),\n", + ")\n", + "sim = caustics.sims.Lens_Source(\n", + " lens=lens,\n", + " source=base_sersic,\n", + " pixelscale=res,\n", + " pixels_x=n_pix,\n", + " z_s=z_s,\n", + " upsample_factor=2,\n", + ")\n", + "fig, axarr = plt.subplots(1, 2, figsize=(14, 7))\n", + "convergence = avg_pool2d(\n", + " lens.convergence(thx, thy, z_s, z_l).squeeze()[None, None], upsample_factor\n", + ").squeeze()\n", + "axarr[0].imshow(np.log10(convergence.numpy()), origin=\"lower\")\n", + "axarr[0].axis(\"off\")\n", + "axarr[0].set_title(\"NFW Convergence\")\n", + "axarr[1].imshow(np.log10(sim([z_l]).numpy()), origin=\"lower\")\n", + "axarr[1].axis(\"off\")\n", + "axarr[1].set_title(\"Lensed Sersic\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "d96d3992", + "metadata": {}, + "source": [ + "## Truncated NFW profile (TNFW)\n", + "\n", + "The TNFW profile is a slight modification to the classic NFW mass profile that approximates the mass distribution of halos in large dark matter simulations. The new density profile is defined as:\n", + "\n", + "$$\\rho_{\\rm tnfw}(r) = \\rho_{\\rm nfw}(r)\\frac{\\tau^2}{\\tau^2 + \\frac{r^2}{r_s^2}}$$\n", + "\n", + "where $\\tau = \\frac{r_t}{r_s}$ is the ratio of the truncation radius to the scale radius. Note that the truncation happens smoothly so there are no sharp boundaries. In the TNFW profile, the mass quantity now actually corresponds the to the total mass since it is no longer divergent. This often means the mass values are larger." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dd5805a5", + "metadata": {}, + "outputs": [], + "source": [ + "lens = caustics.lenses.TNFW(\n", + " cosmology=cosmology,\n", + " name=\"tnfw\",\n", + " x0=torch.tensor(0.0),\n", + " y0=torch.tensor(0.0),\n", + " mass=torch.tensor(1e12),\n", + " scale_radius=torch.tensor(1.0),\n", + " tau=torch.tensor(3.0),\n", + ")\n", + "sim = caustics.sims.Lens_Source(\n", + " lens=lens,\n", + " source=base_sersic,\n", + " pixelscale=res,\n", + " pixels_x=n_pix,\n", + " z_s=z_s,\n", + " upsample_factor=2,\n", + ")\n", + "fig, axarr = plt.subplots(1, 2, figsize=(14, 7))\n", + "convergence = avg_pool2d(\n", + " lens.convergence(thx, thy, z_s, z_l).squeeze()[None, None], upsample_factor\n", + ").squeeze()\n", + "axarr[0].imshow(np.log10(convergence.numpy()), origin=\"lower\")\n", + "axarr[0].axis(\"off\")\n", + "axarr[0].set_title(\"Truncated NFW Convergence\")\n", + "axarr[1].imshow(np.log10(sim([z_l]).numpy()), origin=\"lower\")\n", + "axarr[1].axis(\"off\")\n", + "axarr[1].set_title(\"Lensed Sersic\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "f6f19942", + "metadata": {}, + "source": [ + "## Pseudo Jaffe (PseudoJaffe)\n", + "\n", + "The Pseudo Jaffe closely approximates an isothermal mass distribution except that it is easier to compute and has finite mass.\n", + "\n", + "$$ \\rho(r) = \\frac{\\rho_0}{\\left(1 + \\frac{r^2}{r_c^2}\\right)\\left(1 + \\frac{r^2}{r_s^2}\\right)} $$\n", + "\n", + "where $\\rho_0$ is the central density limit, $r_c$ is the core radius, $r_s$ is the scale radius." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0c8985d8", + "metadata": {}, + "outputs": [], + "source": [ + "lens = caustics.lenses.PseudoJaffe(\n", + " cosmology=cosmology,\n", + " name=\"pseudojaffe\",\n", + " x0=torch.tensor(0.0),\n", + " y0=torch.tensor(0.0),\n", + " mass=torch.tensor(1e13),\n", + " core_radius=torch.tensor(5e-1),\n", + " scale_radius=torch.tensor(15.0),\n", + ")\n", + "sim = caustics.sims.Lens_Source(\n", + " lens=lens,\n", + " source=base_sersic,\n", + " pixelscale=res,\n", + " pixels_x=n_pix,\n", + " z_s=z_s,\n", + " upsample_factor=2,\n", + ")\n", + "fig, axarr = plt.subplots(1, 2, figsize=(14, 7))\n", + "convergence = avg_pool2d(\n", + " lens.convergence(thx, thy, z_s, z_l).squeeze()[None, None], upsample_factor\n", + ").squeeze()\n", + "axarr[0].imshow(np.log10(convergence.numpy()), origin=\"lower\")\n", + "axarr[0].axis(\"off\")\n", + "axarr[0].set_title(\"Pseudo Jaffe Convergence\")\n", + "axarr[1].imshow(np.log10(sim([z_l]).numpy()), origin=\"lower\")\n", + "axarr[1].axis(\"off\")\n", + "axarr[1].set_title(\"Lensed Sersic\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "b232c8b0", + "metadata": {}, + "source": [ + "## External Shear (ExternalShear)\n", + "\n", + "It is often necessary to embed a lens in an external shear field to account for the fact that the lensing mass is not the only mass in the universe." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e1d7b927", + "metadata": {}, + "outputs": [], + "source": [ + "lens = caustics.lenses.ExternalShear(\n", + " cosmology=cosmology,\n", + " name=\"externalshear\",\n", + " x0=torch.tensor(0.0),\n", + " y0=torch.tensor(0.0),\n", + " gamma_1=torch.tensor(1.0),\n", + " gamma_2=torch.tensor(-1.0),\n", + ")\n", + "sim = caustics.sims.Lens_Source(\n", + " lens=lens,\n", + " source=base_sersic,\n", + " pixelscale=res,\n", + " pixels_x=n_pix,\n", + " z_s=z_s,\n", + " upsample_factor=2,\n", + ")\n", + "fig, axarr = plt.subplots(1, 2, figsize=(14, 7))\n", + "# convergence = avg_pool2d(lens.convergence(thx, thy, z_s, z_l).squeeze()[None, None], upsample_factor).squeeze()\n", + "# axarr[0].imshow(np.log10(convergence.numpy()), origin = \"lower\")\n", + "axarr[0].axis(\"off\")\n", + "axarr[0].set_title(\"External Shear Convergence not defined\")\n", + "axarr[1].imshow(np.log10(sim([z_l]).numpy()), origin=\"lower\")\n", + "axarr[1].axis(\"off\")\n", + "axarr[1].set_title(\"Lensed Sersic\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "688d0796", + "metadata": {}, + "source": [ + "## Mass Sheet (MassSheet)\n", + "\n", + "This is a simple case of an external shear field which represents an infinite constant surface density sheet." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cdfba784", + "metadata": {}, + "outputs": [], + "source": [ + "lens = caustics.lenses.MassSheet(\n", + " cosmology=cosmology,\n", + " name=\"masssheet\",\n", + " x0=torch.tensor(0.0),\n", + " y0=torch.tensor(0.0),\n", + " surface_density=torch.tensor(1.5),\n", + ")\n", + "sim = caustics.sims.Lens_Source(\n", + " lens=lens,\n", + " source=base_sersic,\n", + " pixelscale=res,\n", + " pixels_x=n_pix,\n", + " z_s=z_s,\n", + " upsample_factor=2,\n", + ")\n", + "fig, axarr = plt.subplots(1, 2, figsize=(14, 7))\n", + "convergence = avg_pool2d(\n", + " lens.convergence(thx, thy, z_s, z_l).squeeze()[None, None], upsample_factor\n", + ").squeeze()\n", + "axarr[0].imshow(np.log10(convergence.numpy()), origin=\"lower\")\n", + "axarr[0].axis(\"off\")\n", + "axarr[0].set_title(\"Mass Sheet Convergence\")\n", + "axarr[1].imshow(np.log10(sim([z_l]).numpy()), origin=\"lower\")\n", + "axarr[1].axis(\"off\")\n", + "axarr[1].set_title(\"Lensed Sersic\")\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d6c44b6c", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/tutorials/MultiplaneDemo.ipynb b/docs/source/tutorials/MultiplaneDemo.ipynb new file mode 100644 index 00000000..e6c12d53 --- /dev/null +++ b/docs/source/tutorials/MultiplaneDemo.ipynb @@ -0,0 +1,248 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "2a324070-a163-4b79-8e77-819da73083f3", + "metadata": {}, + "source": [ + "# Multiplane Lensing\n", + "\n", + "The universe is three dimensional and filled with stuff. A light ray traveling to our telescope may encounter more than a single massive object on its way to our telescopes. This is handled by a multiplane lensing framework. Multiplane lensing involves tracing the path of a ray backwards from our telescope through each individual plane (which is treated similarly to typical single plane lensing, though extra factors account for the ray physically moving in 3D space) getting perturbed at each step until it finally lands on the source we'd like to image. For more mathmatical details see [Petkova et al. 2014](https://doi.org/10.1093/mnras/stu1860) for the formalism we use internally.\n", + "\n", + "The main concept to keep in mind is that a lot of quantities we are used to working with, such as \"reduced deflection angles\" don't really exist in multiplane lensing since these are normalized by the redshift of the source and lens, however there is no single \"lens redshift\" for multiplane! Instead we define everything with respect to results from full raytracing, once the raytracing is done we can define effective quantities (like effective reduced deflection angle) which behave similarly in intuition but are not quite the same in detail." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "45b6a8b4", + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2\n", + "\n", + "import torch\n", + "from torch.nn.functional import avg_pool2d\n", + "import matplotlib.pyplot as plt\n", + "from ipywidgets import interact\n", + "from astropy.io import fits\n", + "import numpy as np\n", + "\n", + "import caustics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ab43e042", + "metadata": {}, + "outputs": [], + "source": [ + "# initialization stuff for lenses\n", + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology.to(dtype=torch.float32)\n", + "n_pix = 100\n", + "res = 0.05\n", + "upsample_factor = 2\n", + "fov = res * n_pix\n", + "thx, thy = caustics.get_meshgrid(\n", + " res / upsample_factor,\n", + " upsample_factor * n_pix,\n", + " upsample_factor * n_pix,\n", + " dtype=torch.float32,\n", + ")\n", + "z_s = torch.tensor(1.5, dtype=torch.float32)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ea49d25d", + "metadata": {}, + "outputs": [], + "source": [ + "N_planes = 10\n", + "N_lenses = 2 # per plane\n", + "\n", + "z_plane = np.linspace(0.1, 1.0, N_planes)\n", + "planes = []\n", + "\n", + "for p, z_p in enumerate(z_plane):\n", + " lenses = []\n", + "\n", + " for _ in range(N_lenses):\n", + " lenses.append(\n", + " caustics.PseudoJaffe(\n", + " cosmology=cosmology,\n", + " z_l=z_p,\n", + " x0=torch.tensor(np.random.uniform(-fov / 3.0, fov / 3.0)),\n", + " y0=torch.tensor(np.random.uniform(-fov / 3.0, fov / 3.0)),\n", + " mass=torch.tensor(3e10),\n", + " core_radius=0.01,\n", + " scale_radius=torch.tensor(10 ** np.random.uniform(-0.5, 0.5)),\n", + " s=torch.tensor(0.001),\n", + " )\n", + " )\n", + "\n", + " planes.append(\n", + " caustics.lenses.SinglePlane(\n", + " z_l=z_p, cosmology=cosmology, lenses=lenses, name=f\"plane_{p}\"\n", + " )\n", + " )\n", + "\n", + "lens = caustics.lenses.Multiplane(name=\"multiplane\", cosmology=cosmology, lenses=planes)" + ] + }, + { + "cell_type": "markdown", + "id": "4aa429c8", + "metadata": {}, + "source": [ + "## Effective Reduced Deflection Angles" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f2e0a341", + "metadata": {}, + "outputs": [], + "source": [ + "# Effective reduced deflection angles for the multiplane lens system\n", + "ax, ay = lens.effective_reduced_deflection_angle(thx, thy, z_s)\n", + "\n", + "# Plot\n", + "fig, axarr = plt.subplots(1, 2, figsize=(12, 5))\n", + "im = axarr[0].imshow(\n", + " ax, extent=(thx[0][0], thx[0][-1], thy[0][0], thy[-1][0]), origin=\"lower\"\n", + ")\n", + "axarr[0].set_title(\"Deflection angle X\")\n", + "plt.colorbar(im)\n", + "im = axarr[1].imshow(\n", + " ay, extent=(thx[0][0], thx[0][-1], thy[0][0], thy[-1][0]), origin=\"lower\"\n", + ")\n", + "axarr[1].set_title(\"Deflection angle Y\")\n", + "plt.colorbar(im)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "c7ad98c6", + "metadata": {}, + "source": [ + "## Critical Lines" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b2e23c3e", + "metadata": {}, + "outputs": [], + "source": [ + "# Compute critical curves using effective difflection angle\n", + "A = lens.jacobian_lens_equation(thx, thy, z_s)\n", + "\n", + "# Here we compute detA at every point\n", + "detA = torch.linalg.det(A)\n", + "\n", + "# Plot the critical line\n", + "im = plt.imshow(\n", + " detA, extent=(thx[0][0], thx[0][-1], thy[0][0], thy[-1][0]), origin=\"lower\"\n", + ")\n", + "plt.colorbar(im)\n", + "CS = plt.contour(thx, thy, detA, levels=[0.0], colors=\"b\")\n", + "plt.title(\"Critical curves\")\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7cd1f948", + "metadata": {}, + "outputs": [], + "source": [ + "# For completeness, here are the caustics!\n", + "paths = CS.collections[0].get_paths()\n", + "\n", + "for path in paths:\n", + " # Collect the path into a descrete set of points\n", + " vertices = path.interpolated(5).vertices\n", + " x1 = torch.tensor(list(float(vs[0]) for vs in vertices))\n", + " x2 = torch.tensor(list(float(vs[1]) for vs in vertices))\n", + " # raytrace the points to the source plane\n", + " y1, y2 = lens.raytrace(x1, x2, z_s)\n", + "\n", + " # Plot the caustic\n", + " plt.plot(y1, y2, color=\"k\", linewidth=0.5)\n", + "plt.gca().axis(\"off\")\n", + "plt.title(\"Caustics\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "94144e25", + "metadata": {}, + "source": [ + "## Effective Convergence" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d8a84fde", + "metadata": {}, + "outputs": [], + "source": [ + "C = lens.effective_convergence_div(thx, thy, z_s)\n", + "\n", + "plt.imshow(C.detach().cpu().numpy(), origin=\"lower\")\n", + "plt.colorbar()\n", + "plt.title(\"Effective Convergence\")\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "708338df", + "metadata": {}, + "outputs": [], + "source": [ + "Curl = lens.effective_convergence_curl(thx, thy, z_s)\n", + "\n", + "plt.imshow(Curl.detach().cpu().numpy(), origin=\"lower\")\n", + "plt.colorbar()\n", + "plt.title(\"Curl\")\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "00e9c94a", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/tutorials/Playground.ipynb b/docs/source/tutorials/Playground.ipynb new file mode 100644 index 00000000..28be0533 --- /dev/null +++ b/docs/source/tutorials/Playground.ipynb @@ -0,0 +1,235 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "1fafae89", + "metadata": {}, + "source": [ + "# Lensing playground\n", + "\n", + "This is just a fun notebook where you can interactively change lensing parameters. It is a great way to build some intuition around lensing and make some cool pictures!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c0d85608", + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2\n", + "\n", + "import torch\n", + "from torch.nn.functional import avg_pool2d\n", + "import matplotlib.pyplot as plt\n", + "from ipywidgets import interact\n", + "from astropy.io import fits\n", + "import numpy as np\n", + "\n", + "import caustics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b1ac7bf7", + "metadata": {}, + "outputs": [], + "source": [ + "n_pix = 100\n", + "res = 0.05\n", + "upsample_factor = 2\n", + "fov = res * n_pix\n", + "thx, thy = caustics.get_meshgrid(\n", + " res / upsample_factor,\n", + " upsample_factor * n_pix,\n", + " upsample_factor * n_pix,\n", + " dtype=torch.float32,\n", + ")\n", + "z_l = torch.tensor(0.5, dtype=torch.float32)\n", + "z_s = torch.tensor(1.5, dtype=torch.float32)\n", + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology.to(dtype=torch.float32)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "964a76a6", + "metadata": {}, + "outputs": [], + "source": [ + "# SIE lens model, kappa map, alpha map, magnification, time delay, caustics\n", + "\n", + "\n", + "def plot_lens_metrics(thx0, thy0, q, phi, b):\n", + " lens = caustics.lenses.SIE(\n", + " cosmology=cosmology,\n", + " z_l=z_l,\n", + " x0=torch.tensor(thx0),\n", + " y0=torch.tensor(thy0),\n", + " q=torch.tensor(q),\n", + " phi=torch.tensor(phi),\n", + " b=torch.tensor(b),\n", + " name=\"sie\",\n", + " )\n", + " fig, axarr = plt.subplots(2, 3, figsize=(9, 6))\n", + " kappa = avg_pool2d(\n", + " lens.convergence(thx, thy, z_s)[None, None, :, :], upsample_factor\n", + " )[0, 0]\n", + " axarr[0][0].imshow(torch.log10(kappa), origin=\"lower\")\n", + " axarr[0][0].set_title(\"log(convergence)\")\n", + " psi = avg_pool2d(lens.potential(thx, thy, z_s)[None, None, :, :], upsample_factor)[\n", + " 0, 0\n", + " ]\n", + " axarr[0][1].imshow(psi, origin=\"lower\")\n", + " axarr[0][1].set_title(\"potential\")\n", + " timedelay = avg_pool2d(\n", + " lens.time_delay(thx, thy, z_s)[None, None, :, :], upsample_factor\n", + " )[0, 0]\n", + " axarr[0][2].imshow(timedelay, origin=\"lower\")\n", + " axarr[0][2].set_title(\"time delay\")\n", + " magnification = avg_pool2d(\n", + " lens.magnification(thx, thy, z_s)[None, None, :, :], upsample_factor\n", + " )[0, 0]\n", + " axarr[1][0].imshow(torch.log10(magnification), origin=\"lower\")\n", + " axarr[1][0].set_title(\"log(magnification)\")\n", + " alpha = lens.reduced_deflection_angle(thx, thy, z_s)\n", + " alpha0 = avg_pool2d(alpha[0][None, None, :, :], upsample_factor)[0, 0]\n", + " alpha1 = avg_pool2d(alpha[1][None, None, :, :], upsample_factor)[0, 0]\n", + " axarr[1][1].imshow(alpha0, origin=\"lower\")\n", + " axarr[1][1].set_title(\"deflection angle x\")\n", + " axarr[1][2].imshow(alpha1, origin=\"lower\")\n", + " axarr[1][2].set_title(\"deflection angle y\")\n", + " plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3b04c973", + "metadata": {}, + "outputs": [], + "source": [ + "p = interact(\n", + " plot_lens_metrics,\n", + " thx0=(-2.5, 2.5, 0.1),\n", + " thy0=(-2.5, 2.5, 0.1),\n", + " q=(0.01, 0.99, 0.01),\n", + " phi=(0.0, np.pi, np.pi / 25),\n", + " b=(0.1, 2.0, 0.1),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dce1edef", + "metadata": {}, + "outputs": [], + "source": [ + "# Sersic source, demo lensed source\n", + "def plot_lens_distortion(\n", + " x0_lens,\n", + " y0_lens,\n", + " q_lens,\n", + " phi_lens,\n", + " b_lens,\n", + " x0_src,\n", + " y0_src,\n", + " q_src,\n", + " phi_src,\n", + " n_src,\n", + " Re_src,\n", + " Ie_src,\n", + "):\n", + " lens = caustics.lenses.SIE(\n", + " cosmology,\n", + " z_l,\n", + " torch.tensor(x0_lens),\n", + " torch.tensor(y0_lens),\n", + " torch.tensor(q_lens),\n", + " torch.tensor(phi_lens),\n", + " torch.tensor(b_lens),\n", + " name=\"sie\",\n", + " )\n", + " source = caustics.light.Sersic(\n", + " torch.tensor(x0_src),\n", + " torch.tensor(y0_src),\n", + " torch.tensor(q_src),\n", + " torch.tensor(phi_src),\n", + " torch.tensor(n_src),\n", + " torch.tensor(Re_src),\n", + " torch.tensor(Ie_src),\n", + " name=\"sersic\",\n", + " )\n", + " fig, axarr = plt.subplots(1, 3, figsize=(18, 6))\n", + " brightness = avg_pool2d(\n", + " source.brightness(thx, thy)[None, None, :, :], upsample_factor\n", + " )[0, 0]\n", + " axarr[0].imshow(brightness, origin=\"lower\")\n", + " axarr[0].set_title(\"Sersic source\")\n", + " kappa = avg_pool2d(\n", + " lens.convergence(thx, thy, z_s)[None, None, :, :], upsample_factor\n", + " )[0, 0]\n", + " axarr[1].imshow(torch.log10(kappa), origin=\"lower\")\n", + " axarr[1].set_title(\"lens log(convergence)\")\n", + " beta_x, beta_y = lens.raytrace(thx, thy, z_s)\n", + " mu = avg_pool2d(\n", + " source.brightness(beta_x, beta_y)[None, None, :, :], upsample_factor\n", + " )[0, 0]\n", + " axarr[2].imshow(mu, origin=\"lower\")\n", + " axarr[2].set_title(\"Sersic lensed\")\n", + " plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c03161b9", + "metadata": {}, + "outputs": [], + "source": [ + "p = interact(\n", + " plot_lens_distortion,\n", + " x0_lens=(-2.5, 2.5, 0.1),\n", + " y0_lens=(-2.5, 2.5, 0.1),\n", + " q_lens=(0.01, 0.99, 0.01),\n", + " phi_lens=(0.0, np.pi, np.pi / 25),\n", + " b_lens=(0.1, 2.0, 0.1),\n", + " x0_src=(-2.5, 2.5, 0.1),\n", + " y0_src=(-2.5, 2.5, 0.1),\n", + " q_src=(0.01, 0.99, 0.01),\n", + " phi_src=(0.0, np.pi, np.pi / 25),\n", + " n_src=(0.5, 4, 0.1),\n", + " Re_src=(0.1, 2, 0.1),\n", + " Ie_src=(0.1, 2.0, 0.1),\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "35336e43", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/tutorials/Simulators.ipynb b/docs/source/tutorials/Simulators.ipynb new file mode 100644 index 00000000..ae6101bc --- /dev/null +++ b/docs/source/tutorials/Simulators.ipynb @@ -0,0 +1,279 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "f3ed209f", + "metadata": {}, + "source": [ + "# Now you're thinking with Simulators\n", + "\n", + "At least, you will by the end of the tutorial." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "89275b65", + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2\n", + "\n", + "import torch\n", + "from torch.nn.functional import avg_pool2d\n", + "import matplotlib.pyplot as plt\n", + "from astropy.io import fits\n", + "import numpy as np\n", + "\n", + "import caustics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "88a9200a", + "metadata": {}, + "outputs": [], + "source": [ + "# Specify the image/cosmology parameters\n", + "n_pix = 100\n", + "res = 0.05\n", + "upsample_factor = 2\n", + "fov = res * n_pix\n", + "thx, thy = caustics.get_meshgrid(\n", + " res / upsample_factor,\n", + " upsample_factor * n_pix,\n", + " upsample_factor * n_pix,\n", + " dtype=torch.float32,\n", + ")\n", + "z_l = torch.tensor(0.5, dtype=torch.float32)\n", + "z_s = torch.tensor(1.5, dtype=torch.float32)\n", + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology.to(dtype=torch.float32)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1665abeb", + "metadata": {}, + "outputs": [], + "source": [ + "# demo simulator with sersic source, SIE lens. then sample some examples. demo the model graph\n", + "\n", + "\n", + "class Simple_Sim(caustics.Simulator):\n", + " def __init__(\n", + " self,\n", + " lens,\n", + " src,\n", + " z_s=None,\n", + " name: str = \"sim\",\n", + " ):\n", + " super().__init__(name) # need this so `Parametrized` can do its magic\n", + "\n", + " # These are the lens and source objects to keep track of\n", + " self.lens = lens\n", + " self.src = src\n", + "\n", + " # Here we can add a parameter to the simulator, in this case it is `z_s` which we will need later\n", + " self.add_param(\"z_s\", z_s)\n", + "\n", + " def forward(self, params): # define the forward model\n", + " # Here the simulator unpacks the parameter it needs\n", + " z_s = self.unpack(params)\n", + "\n", + " # Note this is very similar to before, except the packed up `x` is all the raytrace function needs to work\n", + " bx, by = self.lens.raytrace(thx, thy, z_s, params)\n", + " mu_fine = self.src.brightness(bx, by, params)\n", + "\n", + " # We return the sampled brightness at each pixel location\n", + " return avg_pool2d(mu_fine.squeeze()[None, None], upsample_factor)[0, 0]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0babaead", + "metadata": {}, + "outputs": [], + "source": [ + "sie = caustics.lenses.SIE(cosmology, name=\"sie\")\n", + "src = caustics.light.Sersic(name=\"src\")\n", + "\n", + "sim = Simple_Sim(sie, src, torch.tensor(0.8))\n", + "\n", + "sim.get_graph(True, True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e672be73", + "metadata": {}, + "outputs": [], + "source": [ + "# Reading the x_keys above we can input the parameters that we would like the simulator to evaluate\n", + "x = torch.tensor(\n", + " [\n", + " z_l.item(), # sie z_l\n", + " 0.7, # sie x0\n", + " 0.13, # sie y0\n", + " 0.4, # sie q\n", + " np.pi / 5, # sie phi\n", + " 1.0, # sie b\n", + " 0.2, # src x0\n", + " 0.5, # src y0\n", + " 0.5, # src q\n", + " -np.pi / 4, # src phi\n", + " 1.5, # src n\n", + " 2.5, # src Re\n", + " 1.0, # src Ie\n", + " ]\n", + ")\n", + "plt.imshow(sim(x), origin=\"lower\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fe16052c", + "metadata": {}, + "outputs": [], + "source": [ + "sie = caustics.lenses.SIE(cosmology, name=\"sie\")\n", + "hdu = fits.open(\n", + " \"https://www.legacysurvey.org/viewer/fits-cutout?ra=36.3684&dec=-25.6389&size=250&layer=ls-dr9&pixscale=0.262&bands=r\"\n", + ")\n", + "image_data = np.array(hdu[0].data, dtype=np.float64)\n", + "src = caustics.light.Pixelated(\n", + " name=\"ESO479_G1\", image=torch.tensor(image_data, dtype=torch.float32)\n", + ")\n", + "\n", + "sim2 = Simple_Sim(sie, src, torch.tensor(0.8))\n", + "\n", + "sim2.get_graph(True, True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b9921f68", + "metadata": {}, + "outputs": [], + "source": [ + "# x can also be a dict\n", + "x = torch.tensor(\n", + " [\n", + " z_l.item(), # sie z_l\n", + " 0.2, # sie x0\n", + " 0.1, # sie y0\n", + " 0.4, # sie q\n", + " np.pi / 5, # sie phi\n", + " 1.0, # sie b\n", + " 0.0, # src x0\n", + " 0.0, # src y0\n", + " 0.01, # src pixelscale\n", + " ]\n", + ")\n", + "fig, axarr = plt.subplots(1, 2, figsize=(12, 6))\n", + "axarr[0].imshow(image_data, origin=\"lower\")\n", + "axarr[0].set_title(\"Source image\")\n", + "axarr[1].imshow(sim2(x), origin=\"lower\")\n", + "axarr[1].set_title(\"Lensed image\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "e16ca698", + "metadata": {}, + "source": [ + "## Setting static/dynamic parameters\n", + "\n", + "So far we have assumed that all parameters for the source and lens are being modelled. However, it is often the case that a parameter can/must be fixed for a given science case. It is very easy to do this in Caustic, simply pass the fixed value when constructing the lens/source objects.\n", + "\n", + "Below we have fixed some parameters in our simulator by providing them when constructing the objects. In the graph they now appear as greyed boxes. In fact we can see now that the cosmology object `FlatLambdaCDM` had fixed parameters all along. This is because there are natural default parameters for such a cosmology. It is possible, of course, to make a different cosmology object which has alternate values, or which leaves some values as free parameters.\n", + "\n", + "In general, to set a parameter as dynamic (must be passed to the simulator) just set it to `None`, to fix a parameter give it a value." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "91d3cf26", + "metadata": {}, + "outputs": [], + "source": [ + "sief = caustics.lenses.SIE(\n", + " name=\"sie\",\n", + " cosmology=cosmology,\n", + " z_l=torch.tensor(0.5),\n", + " x0=torch.tensor(0.0),\n", + " y0=torch.tensor(0.0),\n", + ")\n", + "srcf = caustics.light.Sersic(name=\"src\", n=torch.tensor(2.0))\n", + "\n", + "simf = Simple_Sim(sief, srcf, z_s=torch.tensor(0.8))\n", + "\n", + "simf.get_graph(True, True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6b029159", + "metadata": {}, + "outputs": [], + "source": [ + "# Reading the x_keys above we can input the parameters that we would like the simulator to evaluate\n", + "x = torch.tensor(\n", + " [\n", + " 0.4, # sie q\n", + " np.pi / 5, # sie phi\n", + " 1.0, # sie b\n", + " -0.2, # src x0\n", + " -0.5, # src y0\n", + " 0.5, # src q\n", + " -np.pi / 4, # src phi\n", + " 2.5, # src Re\n", + " 1.0, # src Ie\n", + " ]\n", + ")\n", + "plt.imshow(simf(x), origin=\"lower\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6e244b74", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb6619ec-9057-4daa-916b-9384b74a5e29", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/tutorials/VisualizeCaustics.ipynb b/docs/source/tutorials/VisualizeCaustics.ipynb new file mode 100644 index 00000000..8e163d5c --- /dev/null +++ b/docs/source/tutorials/VisualizeCaustics.ipynb @@ -0,0 +1,164 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "3bb2cd70-e98d-4e8d-b195-64d7e5de6e13", + "metadata": {}, + "source": [ + "# Visualize Caustics\n", + "\n", + "Here we will demonstrate how to collect caustic lines using caustic! Since caustic (the code) uses autodiff and can get exact derivatives, it is actually very acurate at computing caustics. \n", + "\n", + "Conceptually a caustic occurs where the magnification of a lens diverges to infinity. A convenient way to measure the magnification in the image plane is by taking the determinant ($det$) of the jacobian of the lens equation ($A$), its reciprocal is the magnification. This means that anywhere that $det(A) = 0$ is a critical line in the image plane (magnification goes to infinity). If we take this line and raytrace it back to the source plane we can see the caustics which define boundaries for lensing phenomena." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f716feef", + "metadata": {}, + "outputs": [], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2\n", + "\n", + "import torch\n", + "from torch.nn.functional import avg_pool2d\n", + "import matplotlib.pyplot as plt\n", + "from ipywidgets import interact\n", + "from astropy.io import fits\n", + "import numpy as np\n", + "from time import process_time as time\n", + "\n", + "import caustics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bdede2df", + "metadata": {}, + "outputs": [], + "source": [ + "# initialization stuff for an SIE lens\n", + "\n", + "cosmology = caustics.FlatLambdaCDM(name=\"cosmo\")\n", + "cosmology.to(dtype=torch.float32)\n", + "sie = caustics.lenses.SIE(cosmology, name=\"sie\")\n", + "n_pix = 100\n", + "res = 0.05\n", + "upsample_factor = 2\n", + "fov = res * n_pix\n", + "thx, thy = caustics.get_meshgrid(\n", + " res / upsample_factor,\n", + " upsample_factor * n_pix,\n", + " upsample_factor * n_pix,\n", + " dtype=torch.float32,\n", + ")\n", + "z_l = torch.tensor(0.5, dtype=torch.float32)\n", + "z_s = torch.tensor(1.5, dtype=torch.float32)\n", + "x = torch.tensor(\n", + " [\n", + " z_l.item(), # sie z_l\n", + " 0.0, # sie x0\n", + " 0.0, # sie y0\n", + " 0.4, # sie q\n", + " np.pi / 5, # sie phi\n", + " 1.0, # sie b\n", + " ]\n", + ")\n", + "packparams = sie.pack(x)" + ] + }, + { + "cell_type": "markdown", + "id": "38dd09e3", + "metadata": {}, + "source": [ + "## Critical Lines\n", + "\n", + "Before we can see the caustics, we need to find the critical lines. The critical lines are the locus of points in the lens plane (the plane of the mass causing the lensing) at which the magnification of the source becomes theoretically infinite for a point source. In simpler terms, it is where the lensing effect becomes so strong that it can create highly magnified and distorted images of the source. The shape and size of the critical curve depend on the distribution of mass in the lensing object. These lines can be found using the Jacobian of the lensing deflection, specifically $A = \\mathbb{I} - J$. When ${\\rm det}(A) = 0$, that point is on the critical line. Interestingly, $\\frac{1}{{\\rm det}(A)}$ is the magnification, which is why ${\\rm det}(A) = 0$ defines the points of infinite magnification." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "487b3030", + "metadata": {}, + "outputs": [], + "source": [ + "# Conveniently caustic has a function to compute the jacobian of the lens equation\n", + "A = sie.jacobian_lens_equation(thx, thy, z_s, packparams)\n", + "# Note that if this is too slow you can set `method = \"finitediff\"` to run a faster version. You will also need to provide `pixelscale` then\n", + "\n", + "# Here we compute A's determinant at every point\n", + "detA = torch.linalg.det(A)\n", + "\n", + "# Plot the critical line\n", + "im = plt.imshow(\n", + " detA, extent=(thx[0][0], thx[0][-1], thy[0][0], thy[-1][0]), origin=\"lower\"\n", + ")\n", + "plt.colorbar(im)\n", + "CS = plt.contour(thx, thy, detA, levels=[0.0], colors=\"b\")\n", + "\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "3b099cfd", + "metadata": {}, + "source": [ + "## Caustics\n", + "\n", + "Critical lines show us where the magnification approaches infinity, they are important structures in understanding a lensing system. These lines are also very useful when mapped into the source plane. When the critical lines are raytraced back to the source plane they are called caustics (see what we did there?). In the source plane these lines deliniate when a source will be multiply imaged. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0c1f1177", + "metadata": {}, + "outputs": [], + "source": [ + "# Get the path from the matplotlib contour plot of the critical line\n", + "paths = CS.collections[0].get_paths()\n", + "\n", + "for path in paths:\n", + " # Collect the path into a descrete set of points\n", + " vertices = path.interpolated(5).vertices\n", + " x1 = torch.tensor(list(float(vs[0]) for vs in vertices))\n", + " x2 = torch.tensor(list(float(vs[1]) for vs in vertices))\n", + " # raytrace the points to the source plane\n", + " y1, y2 = sie.raytrace(x1, x2, z_s, packparams)\n", + "\n", + " # Plot the caustic\n", + " plt.plot(y1, y2)\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1e77da2e", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/source/tutorials/index.md b/docs/source/tutorials/index.md new file mode 100644 index 00000000..75d2e53c --- /dev/null +++ b/docs/source/tutorials/index.md @@ -0,0 +1,5 @@ +# Tutorials + +Here you will find the jupyter notebook tutorials. +It is recommended that you go through the tutorials yourself, +but for quick reference a version of each tutorial is available here. \ No newline at end of file diff --git a/noxfile.py b/noxfile.py index 5dc68404..95165cff 100644 --- a/noxfile.py +++ b/noxfile.py @@ -1,6 +1,5 @@ from __future__ import annotations -import argparse import shutil from pathlib import Path diff --git a/src/caustics/cosmology.py b/src/caustics/cosmology.py index 6b55112d..9f72de7b 100644 --- a/src/caustics/cosmology.py +++ b/src/caustics/cosmology.py @@ -42,20 +42,27 @@ class Cosmology(Parametrized): This class provides an interface for cosmological computations used in lensing such as comoving distance and critical surface density. - Units: - - Distance: Mpc - - Mass: solar mass - - Attributes: - name (str): Name of the cosmological model. + Units + ----- + Distance + Mpc + Mass + solar mass + + Attributes + ---------- + name: str + Name of the cosmological model. """ def __init__(self, name: str = None): """ Initialize the Cosmology. - Args: - name (str): Name of the cosmological model. + Parameters + ---------- + name: str + Name of the cosmological model. """ super().__init__(name) @@ -64,12 +71,17 @@ def critical_density(self, z: Tensor, params: Optional["Packed"] = None) -> Tens """ Compute the critical density at redshift z. - Args: - z (Tensor): The redshifts. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - Tensor: The critical density at each redshift. + Parameters + ---------- + z: Tensor + The redshifts. + params: Packed, optional + Dynamic parameter container for the computation. + + Returns + ------- + Tensor + The critical density at each redshift. """ ... @@ -81,12 +93,17 @@ def comoving_distance( """ Compute the comoving distance to redshift z. - Args: - z (Tensor): The redshifts. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - Tensor: The comoving distance to each redshift. + Parameters + ---------- + z: Tensor + The redshifts. + params: (Packed, optional0 + Dynamic parameter container for the computation. + + Returns + ------- + Tensor + The comoving distance to each redshift. """ ... @@ -98,12 +115,17 @@ def transverse_comoving_distance( """ Compute the transverse comoving distance to redshift z (Mpc). - Args: - z (Tensor): The redshifts. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - Tensor: The transverse comoving distance to each redshift in Mpc. + Parameters + ---------- + z: Tensor + The redshifts. + params: (Packed, optional) + Dynamic parameter container for the computation. + + Returns + ------- + Tensor + The transverse comoving distance to each redshift in Mpc. """ ... @@ -114,13 +136,19 @@ def comoving_distance_z1z2( """ Compute the comoving distance between two redshifts. - Args: - z1 (Tensor): The starting redshifts. - z2 (Tensor): The ending redshifts. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - Tensor: The comoving distance between each pair of redshifts. + Parameters + ---------- + z1: Tensor + The starting redshifts. + z2: Tensor + The ending redshifts. + params: (Packed, optional) + Dynamic parameter container for the computation. + + Returns + ------- + Tensor + The comoving distance between each pair of redshifts. """ return self.comoving_distance(z2, params) - self.comoving_distance(z1, params) @@ -131,13 +159,19 @@ def transverse_comoving_distance_z1z2( """ Compute the transverse comoving distance between two redshifts (Mpc). - Args: - z1 (Tensor): The starting redshifts. - z2 (Tensor): The ending redshifts. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - Tensor: The transverse comoving distance between each pair of redshifts in Mpc. + Parameters + ---------- + z1: Tensor + The starting redshifts. + z2: Tensor + The ending redshifts. + params: (Packed, optional) + Dynamic parameter container for the computation. + + Returns + ------- + Tensor + The transverse comoving distance between each pair of redshifts in Mpc. """ return self.transverse_comoving_distance( z2, params @@ -150,12 +184,17 @@ def angular_diameter_distance( """ Compute the angular diameter distance to redshift z. - Args: - z (Tensor): The redshifts. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - Tensor: The angular diameter distance to each redshift. + Parameters + ----------- + z: Tensor + The redshifts. + params: (Packed, optional) + Dynamic parameter container for the computation. + + Returns + ------- + Tensor + The angular diameter distance to each redshift. """ return self.comoving_distance(z, params) / (1 + z) @@ -166,13 +205,19 @@ def angular_diameter_distance_z1z2( """ Compute the angular diameter distance between two redshifts. - Args: - z1 (Tensor): The starting redshifts. - z2 (Tensor): The ending redshifts. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - Tensor: The angular diameter distance between each pair of redshifts. + Parameters + ---------- + z1: Tensor + The starting redshifts. + z2: Tensor + The ending redshifts. + params: (Packed, optional) + Dynamic parameter container for the computation. + + Returns + ------- + Tensor + The angular diameter distance between each pair of redshifts. """ return self.comoving_distance_z1z2(z1, z2, params) / (1 + z2) @@ -183,13 +228,19 @@ def time_delay_distance( """ Compute the time delay distance between lens and source planes. - Args: - z_l (Tensor): The lens redshifts. - z_s (Tensor): The source redshifts. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - Tensor: The time delay distance for each pair of lens and source redshifts. + Parameters + ---------- + z_l: Tensor + The lens redshifts. + z_s: Tensor + The source redshifts. + params: (Packed, optional) + Dynamic parameter container for the computation. + + Returns + ------- + Tensor + The time delay distance for each pair of lens and source redshifts. """ d_l = self.angular_diameter_distance(z_l, params) d_s = self.angular_diameter_distance(z_s, params) @@ -203,13 +254,19 @@ def critical_surface_density( """ Compute the critical surface density between lens and source planes. - Args: - z_l (Tensor): The lens redshifts. - z_s (Tensor): The source redshifts. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - Tensor: The critical surface density for each pair of lens and source redshifts. + Parameters + ---------- + z_l: Tensor + The lens redshifts. + z_s: Tensor + The source redshifts. + params: (Packed, optional) + Dynamic parameter container for the computation. + + Returns + ------- + Tensor + The critical surface density for each pair of lens and source redshifts. """ d_l = self.angular_diameter_distance(z_l, params) d_s = self.angular_diameter_distance(z_s, params) @@ -232,11 +289,16 @@ def __init__( """ Initialize a new instance of the FlatLambdaCDM class. - Args: - name (str): Name of the cosmology. - h0 (Optional[Tensor]): Hubble constant over 100. Default is h0_default. - critical_density_0 (Optional[Tensor]): Critical density at z=0. Default is critical_density_0_default. - Om0 (Optional[Tensor]): Matter density parameter at z=0. Default is Om0_default. + Parameters + ---------- + name: str + Name of the cosmology. + h0: Optional[Tensor] + Hubble constant over 100. Default is h0_default. + critical_density_0: (Optional[Tensor]) + Critical density at z=0. Default is critical_density_0_default. + Om0: Optional[Tensor] + Matter density parameter at z=0. Default is Om0_default. """ super().__init__(name) @@ -266,11 +328,15 @@ def hubble_distance(self, h0): """ Calculate the Hubble distance. - Args: - h0 (Tensor): Hubble constant. + Parameters + ---------- + h0: Tensor + Hubble constant. - Returns: - Tensor: Hubble distance. + Returns + ------- + Tensor + Hubble distance. """ return c_Mpc_s / (100 * km_to_Mpc) / h0 @@ -287,12 +353,17 @@ def critical_density( """ Calculate the critical density at redshift z. - Args: - z (Tensor): Redshift. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - torch.Tensor: Critical density at redshift z. + Parameters + ---------- + z: Tensor + Redshift. + params: (Packed, optional) + Dynamic parameter container for the computation. + + Returns + ------- + torch.Tensor + Critical density at redshift z. """ Ode0 = 1 - Om0 return central_critical_density * (Om0 * (1 + z) ** 3 + Ode0) @@ -304,11 +375,15 @@ def _comoving_distance_helper( """ Helper method for computing comoving distances. - Args: - x (Tensor): Input tensor. + Parameters + ---------- + x: Tensor + Input tensor. - Returns: - Tensor: Computed comoving distances. + Returns + ------- + Tensor + Computed comoving distances. """ return interp1d( self._comoving_distance_helper_x_grid, @@ -329,12 +404,17 @@ def comoving_distance( """ Calculate the comoving distance to redshift z. - Args: - z (Tensor): Redshift. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - Tensor: Comoving distance to redshift z. + Parameters + ---------- + z: Tensor + Redshift. + params: (Packed, optional) + Dynamic parameter container for the computation. + + Returns + ------- + Tensor + Comoving distance to redshift z. """ Ode0 = 1 - Om0 ratio = (Om0 / Ode0) ** (1 / 3) diff --git a/src/caustics/data/hdf5dataset.py b/src/caustics/data/hdf5dataset.py index 02236973..11c43311 100644 --- a/src/caustics/data/hdf5dataset.py +++ b/src/caustics/data/hdf5dataset.py @@ -22,11 +22,17 @@ def __init__( dtypes: Union[Dict[str, torch.dtype], torch.dtype] = torch.float32, ): """ - Args: - path: location of dataset. - keys: dataset keys to read. - dtypes: either a numpy datatype to which the items will be converted - or a dictionary specifying the datatype corresponding to each key. + Parameters + ---------- + path: string + location of dataset. + keys: List[str] + dataset keys to read. + device: torch.deviec + the device for torch + dtypes: torch.dtype + either a numpy datatype to which the items will be converted + or a dictionary specifying the datatype corresponding to each key. """ super().__init__() self.keys = keys diff --git a/src/caustics/data/probes.py b/src/caustics/data/probes.py index def03ef9..05df1304 100644 --- a/src/caustics/data/probes.py +++ b/src/caustics/data/probes.py @@ -24,6 +24,9 @@ def __len__(self): def __getitem__(self, i: Union[int, slice]) -> Tensor: """ - Returns image `i` with channel as first dimension. + Returns + ------- + Tensor + image `i` with channel as first dimension. """ return self.ds[i][self.key].movedim(-1, 0) diff --git a/src/caustics/lenses/base.py b/src/caustics/lenses/base.py index 768205f1..c6a26d2c 100644 --- a/src/caustics/lenses/base.py +++ b/src/caustics/lenses/base.py @@ -24,9 +24,12 @@ def __init__(self, cosmology: Cosmology, name: str = None): """ Initializes a new instance of the Lens class. - Args: - name (str): The name of the lens model. - cosmology (Cosmology): An instance of a Cosmology class that describes the cosmological parameters of the model. + Parameters + ---------- + name: string + The name of the lens model. + cosmology: Cosmology + An instance of a Cosmology class that describes the cosmological parameters of the model. """ super().__init__(name) self.cosmology = cosmology @@ -75,14 +78,21 @@ def magnification( """ Compute the gravitational magnification at the given coordinates. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - - Returns: - Tensor: Gravitational magnification at the given coordinates. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + + Returns + ------- + Tensor + Gravitational magnification at the given coordinates. """ return get_magnification(partial(self.raytrace, params=params), x, y, z_s) @@ -102,17 +112,27 @@ def forward_raytrace( """ Perform a forward ray-tracing operation which maps from the source plane to the image plane. - Args: - bx (Tensor): Tensor of x coordinate in the source plane (scalar). - by (Tensor): Tensor of y coordinate in the source plane (scalar). - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - epsilon (Tensor): maximum distance between two images (arcsec) before they are considered the same image. - n_init (int): number of random initialization points used to try and find image plane points. - fov (float): the field of view in which the initial random samples are taken. - - Returns: - tuple[Tensor, Tensor]: Ray-traced coordinates in the x and y directions. + Parameters + ---------- + bx: Tensor + Tensor of x coordinate in the source plane (scalar). + by: Tensor + Tensor of y coordinate in the source plane (scalar). + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + epsilon: Tensor + maximum distance between two images (arcsec) before they are considered the same image. + n_init: int + number of random initialization points used to try and find image plane points. + fov: float + the field of view in which the initial random samples are taken. + + Returns + ------- + tuple[Tensor, Tensor] + Ray-traced coordinates in the x and y directions. """ bxy = torch.stack((bx, by)).repeat(n_init, 1) # has shape (n_init, Dout:2) @@ -155,8 +175,10 @@ class ThickLens(Lens): Base class for modeling gravitational lenses that cannot be treated using the thin lens approximation. It is an abstract class and should be subclassed for different types of lens models. - Attributes: - cosmology (Cosmology): An instance of a Cosmology class that describes the cosmological parameters of the model. + Attributes + ---------- + cosmology: Cosmology + An instance of a Cosmology class that describes the cosmological parameters of the model. """ @unpack(3) @@ -172,14 +194,18 @@ def reduced_deflection_angle( """ ThickLens objects do not have a reduced deflection angle since the distance D_ls is undefined - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. - Raises: - NotImplementedError + Raises:: NotImplementedError """ warnings.warn( "ThickLens objects do not have a reduced deflection angle since they have no unique lens redshift. The distance D_{ls} is undefined in the equation $\alpha_{reduced} = \frac{D_{ls}}{D_s}\alpha_{physical}$. See `effective_reduced_deflection_angle`. Now using effective_reduced_deflection_angle, please switch functions to remove this warning" @@ -204,11 +230,16 @@ def effective_reduced_deflection_angle( angular coordinates, and $\beta$ are the angular coordinates to the source plane. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. """ bx, by = self.raytrace(x, y, z_s, params) @@ -228,14 +259,21 @@ def physical_deflection_angle( plane. ThickLens objects have no unique definition of a lens plane and so cannot compute a physical_deflection_angle - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. - Returns: - tuple[Tensor, Tensor]: Tuple of Tensors representing the x and y components of the deflection angle, respectively. + Returns + ------- + tuple[Tensor, Tensor] + Tuple of Tensors representing the x and y components of the deflection angle, respectively. """ raise NotImplementedError( @@ -257,13 +295,19 @@ def raytrace( source plance associated with a given input observed angular coordinate x,y. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - - Returns: + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + + Returns + ------- tuple[Tensor, Tensor]: Tuple of Tensors representing the x and y coordinates of the ray-traced light rays, respectively. """ @@ -283,14 +327,21 @@ def surface_density( """ Computes the projected mass density at given coordinates. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - - Returns: - Tensor: The projected mass density at the given coordinates in units of solar masses per square Megaparsec. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + + Returns + ------- + Tensor + The projected mass density at the given coordinates in units of solar masses per square Megaparsec. """ ... @@ -308,14 +359,21 @@ def time_delay( """ Computes the gravitational time delay at given coordinates. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor ofsource redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - - Returns: - Tensor: The gravitational time delay at the given coordinates. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor ofsource redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + + Returns + ------- + Tensor + The gravitational time delay at the given coordinates. """ ... @@ -490,10 +548,14 @@ class ThinLens(Lens): lensing quantities such as the deflection angle, convergence, potential, surface mass density, and gravitational time delay. - Args: - name (str): Name of the lens model. - cosmology (Cosmology): Cosmology object that encapsulates cosmological parameters and distances. - z_l (Optional[Tensor], optional): Redshift of the lens. Defaults to None. + Attributes + ---------- + name: string + Name of the lens model. + cosmology: Cosmology + Cosmology object that encapsulates cosmological parameters and distances. + z_l: (Optional[Tensor], optional) + Redshift of the lens. Defaults to None. """ @@ -520,14 +582,21 @@ def reduced_deflection_angle( """ Computes the reduced deflection angle of the lens at given coordinates [arcsec]. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - - Returns: - tuple[Tensor, Tensor]: Reduced deflection angle in x and y directions. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + + Returns + -------- + tuple[Tensor, Tensor] + Reduced deflection angle in x and y directions. """ d_s = self.cosmology.angular_diameter_distance(z_s, params) d_ls = self.cosmology.angular_diameter_distance_z1z2(z_l, z_s, params) @@ -550,14 +619,21 @@ def physical_deflection_angle( """ Computes the physical deflection angle immediately after passing through this lens's plane. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - - Returns: - tuple[Tensor, Tensor]: Physical deflection angle in x and y directions in arcseconds. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + + Returns + ------- + tuple[Tensor, Tensor] + Physical deflection angle in x and y directions in arcseconds. """ d_s = self.cosmology.angular_diameter_distance(z_s, params) d_ls = self.cosmology.angular_diameter_distance_z1z2(z_l, z_s, params) @@ -580,14 +656,21 @@ def convergence( """ Computes the convergence of the lens at given coordinates. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - - Returns: - Tensor: Convergence at the given coordinates. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + + Returns + ------- + Tensor + Convergence at the given coordinates. """ ... @@ -605,13 +688,21 @@ def potential( """ Computes the gravitational lensing potential at given coordinates. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - - Returns: Tensor: Gravitational lensing potential at the given coordinates in arcsec^2. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + + Returns + ------- + Tensor + Gravitational lensing potential at the given coordinates in arcsec^2. """ ... @@ -629,14 +720,21 @@ def surface_density( """ Computes the surface mass density of the lens at given coordinates. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - - Returns: - Tensor: Surface mass density at the given coordinates in solar masses per Mpc^2. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + + Returns + ------- + Tensor + Surface mass density at the given coordinates in solar masses per Mpc^2. """ critical_surface_density = self.cosmology.critical_surface_density( z_l, z_s, params @@ -656,14 +754,21 @@ def raytrace( """ Perform a ray-tracing operation by subtracting the deflection angles from the input coordinates. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - - Returns: - tuple[Tensor, Tensor]: Ray-traced coordinates in the x and y directions. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + + Returns + ------- + tuple[Tensor, Tensor] + Ray-traced coordinates in the x and y directions. """ ax, ay = self.reduced_deflection_angle(x, y, z_s, params) return x - ax, y - ay @@ -682,14 +787,21 @@ def time_delay( """ Compute the gravitational time delay for light passing through the lens at given coordinates. - Args: - x (Tensor): Tensor of x coordinates in the lens plane. - y (Tensor): Tensor of y coordinates in the lens plane. - z_s (Tensor): Tensor of source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. Defaults to None. - - Returns: - Tensor: Time delay at the given coordinates. + Parameters + ---------- + x: Tensor + Tensor of x coordinates in the lens plane. + y: Tensor + Tensor of y coordinates in the lens plane. + z_s: Tensor + Tensor of source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. Defaults to None. + + Returns + ------- + Tensor + Time delay at the given coordinates. """ d_l = self.cosmology.angular_diameter_distance(z_l, params) d_s = self.cosmology.angular_diameter_distance(z_s, params) diff --git a/src/caustics/lenses/epl.py b/src/caustics/lenses/epl.py index 4945527e..378bf97f 100644 --- a/src/caustics/lenses/epl.py +++ b/src/caustics/lenses/epl.py @@ -18,17 +18,27 @@ class EPL(ThinLens): This class represents a thin gravitational lens model with an elliptical power law profile. The lensing equations are solved iteratively using an approach based on Tessore et al. 2015. - Attributes: - n_iter (int): Number of iterations for the iterative solver. - s (float): Softening length for the elliptical power-law profile. - - Parameters: - z_l (Optional[Union[Tensor, float]]): This is the redshift of the lens. In the context of gravitational lensing, the lens is the galaxy or other mass distribution that is bending the light from a more distant source. - x0 and y0 (Optional[Union[Tensor, float]]): These are the coordinates of the lens center in the lens plane. The lens plane is the plane perpendicular to the line of sight in which the deflection of light by the lens is considered. - q (Optional[Union[Tensor, float]]): This is the axis ratio of the lens, i.e., the ratio of the minor axis to the major axis of the elliptical lens. - phi (Optional[Union[Tensor, float]]): This is the orientation of the lens on the sky, typically given as an angle measured counter-clockwise from some reference direction. - b (Optional[Union[Tensor, float]]): This is the scale length of the lens, which sets the overall scale of the lensing effect. In some contexts, this is referred to as the Einstein radius. - t (Optional[Union[Tensor, float]]): This is the power-law slope parameter of the lens model. In the context of the EPL model, t is equivalent to the gamma parameter minus one, where gamma is the power-law index of the radial mass distribution of the lens. + Attributes + ---------- + n_iter: int + Number of iterations for the iterative solver. + s: float + Softening length for the elliptical power-law profile. + + Parameters + ---------- + z_l: Optional[Union[Tensor, float]] + This is the redshift of the lens. In the context of gravitational lensing, the lens is the galaxy or other mass distribution that is bending the light from a more distant source. + x0 and y0: Optional[Union[Tensor, float]] + These are the coordinates of the lens center in the lens plane. The lens plane is the plane perpendicular to the line of sight in which the deflection of light by the lens is considered. + q: Optional[Union[Tensor, float]] + This is the axis ratio of the lens, i.e., the ratio of the minor axis to the major axis of the elliptical lens. + phi: Optional[Union[Tensor, float]] + This is the orientation of the lens on the sky, typically given as an angle measured counter-clockwise from some reference direction. + b: Optional[Union[Tensor, float]] + This is the scale length of the lens, which sets the overall scale of the lensing effect. In some contexts, this is referred to as the Einstein radius. + t: Optional[Union[Tensor, float]] + This is the power-law slope parameter of the lens model. In the context of the EPL model, t is equivalent to the gamma parameter minus one, where gamma is the power-law index of the radial mass distribution of the lens. """ @@ -49,18 +59,30 @@ def __init__( """ Initialize an EPL lens model. - Args: - name (str): Name of the lens model. - cosmology (Cosmology): Cosmology object that provides cosmological distance calculations. - z_l (Optional[Tensor]): Redshift of the lens. If not provided, it is considered as a free parameter. - x0 (Optional[Tensor]): X coordinate of the lens center. If not provided, it is considered as a free parameter. - y0 (Optional[Tensor]): Y coordinate of the lens center. If not provided, it is considered as a free parameter. - q (Optional[Tensor]): Axis ratio of the lens. If not provided, it is considered as a free parameter. - phi (Optional[Tensor]): Position angle of the lens. If not provided, it is considered as a free parameter. - b (Optional[Tensor]): Scale length of the lens. If not provided, it is considered as a free parameter. - t (Optional[Tensor]): Power law slope (`gamma-1`) of the lens. If not provided, it is considered as a free parameter. - s (float): Softening length for the elliptical power-law profile. - n_iter (int): Number of iterations for the iterative solver. + Parameters + ----------- + name: string + Name of the lens model. + cosmology: Cosmology + Cosmology object that provides cosmological distance calculations. + z_l: Optional[Tensor] + Redshift of the lens. If not provided, it is considered as a free parameter. + x0: Optional[Tensor] + X coordinate of the lens center. If not provided, it is considered as a free parameter. + y0: Optional[Tensor] + Y coordinate of the lens center. If not provided, it is considered as a free parameter. + q: Optional[Tensor] + Axis ratio of the lens. If not provided, it is considered as a free parameter. + phi: Optional[Tensor] + Position angle of the lens. If not provided, it is considered as a free parameter. + b: Optional[Tensor] + Scale length of the lens. If not provided, it is considered as a free parameter. + t: Optional[Tensor] + Power law slope (`gamma-1`) of the lens. If not provided, it is considered as a free parameter. + s: float + Softening length for the elliptical power-law profile. + n_iter: int + Number of iterations for the iterative solver. """ super().__init__(cosmology, z_l, name=name) @@ -94,14 +116,21 @@ def reduced_deflection_angle( """ Compute the reduced deflection angles of the lens. - Args: - x (Tensor): X coordinates in the lens plane. - y (Tensor): Y coordinates in the lens plane. - z_s (Tensor): Source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. + Parameters + ---------- + x: Tensor + X coordinates in the lens plane. + y: Tensor + Y coordinates in the lens plane. + z_s: Tensor + Source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. - Returns: - tuple[Tensor, Tensor]: Reduced deflection angles in the x and y directions. + Returns + -------- + tuple[Tensor, Tensor] + Reduced deflection angles in the x and y directions. """ x, y = translate_rotate(x, y, x0, y0, phi) @@ -121,13 +150,19 @@ def _r_omega(self, z, t, q): """ Iteratively computes `R * omega(phi)` (eq. 23 in Tessore et al 2015). - Args: - z (Tensor): `R * e^(i * phi)`, position vector in the lens plane. - t (Tensor): Power law slow (`gamma-1`). - q (Tensor): Axis ratio. + Parameters + ---------- + z: Tensor + `R * e^(i * phi)`, position vector in the lens plane. + t: Tensor + Power law slow (`gamma-1`). + q: Tensor + Axis ratio. - Returns: - Tensor: The value of `R * omega(phi)`. + Returns + -------- + Tensor + The value of `R * omega(phi)`. """ # constants f = (1.0 - q) / (1.0 + q) @@ -164,14 +199,21 @@ def potential( """ Compute the lensing potential of the lens. - Args: - x (Tensor): X coordinates in the lens plane. - y (Tensor): Y coordinates in the lens plane. - z_s (Tensor): Source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. + Parameters + ---------- + x: Tensor + X coordinates in the lens plane. + y: Tensor + Y coordinates in the lens plane. + z_s: Tensor + Source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. - Returns: - Tensor: The lensing potential. + Returns + ------- + Tensor + The lensing potential. """ ax, ay = self.reduced_deflection_angle(x, y, z_s, params) ax, ay = derotate(ax, ay, -phi) @@ -198,14 +240,21 @@ def convergence( """ Compute the convergence of the lens, which describes the local density of the lens. - Args: - x (Tensor): X coordinates in the lens plane. - y (Tensor): Y coordinates in the lens plane. - z_s (Tensor): Source redshifts. - params (Packed, optional): Dynamic parameter container for the lens model. + Parameters + ---------- + x: Tensor + X coordinates in the lens plane. + y: Tensor + Y coordinates in the lens plane. + z_s: Tensor + Source redshifts. + params: (Packed, optional) + Dynamic parameter container for the lens model. - Returns: - Tensor: The convergence of the lens. + Returns + ------- + Tensor + The convergence of the lens. """ x, y = translate_rotate(x, y, x0, y0, phi) psi = (q**2 * (x**2 + self.s**2) + y**2).sqrt() diff --git a/src/caustics/lenses/external_shear.py b/src/caustics/lenses/external_shear.py index b1f41450..a2276b14 100644 --- a/src/caustics/lenses/external_shear.py +++ b/src/caustics/lenses/external_shear.py @@ -14,14 +14,22 @@ class ExternalShear(ThinLens): """ Represents an external shear effect in a gravitational lensing system. - Attributes: - name (str): Identifier for the lens instance. - cosmology (Cosmology): The cosmological model used for lensing calculations. - z_l (Optional[Union[Tensor, float]]): The redshift of the lens. - x0, y0 (Optional[Union[Tensor, float]]): Coordinates of the shear center in the lens plane. - gamma_1, gamma_2 (Optional[Union[Tensor, float]]): Shear components. - - Note: The shear components gamma_1 and gamma_2 represent an external shear, a gravitational + Attributes + ---------- + name: str + Identifier for the lens instance. + cosmology: Cosmology + The cosmological model used for lensing calculations. + z_l: Optional[Union[Tensor, float]] + The redshift of the lens. + x0, y0: Optional[Union[Tensor, float]] + Coordinates of the shear center in the lens plane. + gamma_1, gamma_2: Optional[Union[Tensor, float]] + Shear components. + + Notes + ------ + The shear components gamma_1 and gamma_2 represent an external shear, a gravitational distortion that can be caused by nearby structures outside of the main lens galaxy. """ @@ -62,14 +70,21 @@ def reduced_deflection_angle( """ Calculates the reduced deflection angle. - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Returns: - tuple[Tensor, Tensor]: The reduced deflection angles in the x and y directions. + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + tuple[Tensor, Tensor] + The reduced deflection angles in the x and y directions. """ x, y = translate_rotate(x, y, x0, y0) # Meneghetti eq 3.83 @@ -100,14 +115,21 @@ def potential( """ Calculates the lensing potential. - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The lensing potential. + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The lensing potential. """ ax, ay = self.reduced_deflection_angle(x, y, z_s, params) x, y = translate_rotate(x, y, x0, y0) @@ -131,14 +153,21 @@ def convergence( """ The convergence is undefined for an external shear. - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Raises: - NotImplementedError: This method is not implemented as the convergence is not defined + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Raises + ------ + NotImplementedError + This method is not implemented as the convergence is not defined for an external shear. """ raise NotImplementedError("convergence undefined for external shear") diff --git a/src/caustics/lenses/mass_sheet.py b/src/caustics/lenses/mass_sheet.py index 4e8e5a58..1e7a79ea 100644 --- a/src/caustics/lenses/mass_sheet.py +++ b/src/caustics/lenses/mass_sheet.py @@ -15,14 +15,22 @@ class MassSheet(ThinLens): """ Represents an external shear effect in a gravitational lensing system. - Attributes: - name (str): Identifier for the lens instance. - cosmology (Cosmology): The cosmological model used for lensing calculations. - z_l (Optional[Union[Tensor, float]]): The redshift of the lens. - x0, y0 (Optional[Union[Tensor, float]]): Coordinates of the shear center in the lens plane. - gamma_1, gamma_2 (Optional[Union[Tensor, float]]): Shear components. + Attributes + ---------- + name: string + Identifier for the lens instance. + cosmology: Cosmology + The cosmological model used for lensing calculations. + z_l: Optional[Union[Tensor, float]] + The redshift of the lens. + x0, y0: Optional[Union[Tensor, float]] + Coordinates of the shear center in the lens plane. + gamma_1, gamma_2: Optional[Union[Tensor, float]] + Shear components. - Note: The shear components gamma_1 and gamma_2 represent an external shear, a gravitational + Notes + ------ + The shear components gamma_1 and gamma_2 represent an external shear, a gravitational distortion that can be caused by nearby structures outside of the main lens galaxy. """ @@ -58,14 +66,21 @@ def reduced_deflection_angle( """ Calculates the reduced deflection angle. - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. - Returns: - tuple[Tensor, Tensor]: The reduced deflection angles in the x and y directions. + Returns + ------- + tuple[Tensor, Tensor] + The reduced deflection angles in the x and y directions. """ x, y = translate_rotate(x, y, x0, y0) # Meneghetti eq 3.84 diff --git a/src/caustics/lenses/multiplane.py b/src/caustics/lenses/multiplane.py index e24ae2da..59192ed2 100644 --- a/src/caustics/lenses/multiplane.py +++ b/src/caustics/lenses/multiplane.py @@ -16,13 +16,19 @@ class Multiplane(ThickLens): """ Class for handling gravitational lensing with multiple lens planes. - Attributes: - lenses (list[ThinLens]): List of thin lenses. - - Args: - name (str): Name of the lens. - cosmology (Cosmology): Cosmological parameters used for calculations. - lenses (list[ThinLens]): List of thin lenses. + Attributes + ---------- + lenses (list[ThinLens]) + List of thin lenses. + + Parameters + ---------- + name: string + Name of the lens. + cosmology: Cosmology + Cosmological parameters used for calculations. + lenses: list[ThinLens] + List of thin lenses. """ def __init__(self, cosmology: Cosmology, lenses: list[ThinLens], name: str = None): @@ -38,11 +44,15 @@ def get_z_ls( """ Get the redshifts of each lens in the multiplane. - Args: - params (Packed, optional): Dynamic parameter container. + Parameters + ---------- + params: (Packed, optional) + Dynamic parameter container. - Returns: - List[Tensor]: Redshifts of the lenses. + Returns + -------- + List[Tensor] + Redshifts of the lenses. """ # Relies on z_l being the first element to be unpacked, which should always # be the case for a ThinLens @@ -78,14 +88,21 @@ def raytrace( Here we set as initialization :math:`\vec{\theta}^0 = theta` the observation angular coordinates and :math:`\vec{x}^0 = 0` the initial physical coordinates (i.e. the observation rays come from a point at the observer). The indexing of :math:`\vec{x}^i` and :math:`\vec{\theta}^i` indicates the properties at the plane :math:`i`, and 0 means the observer, 1 is the first lensing plane (infinitesimally after the plane since the deflection has been applied), and so on. Note that in the actual implementation we start at :math:`\vec{x}^1` and :math:`\vec{\theta}^0` and begin at the second step in the recursion formula. - Args: - x (Tensor): angular x-coordinates from the observer perspective. - y (Tensor): angular y-coordinates from the observer perspective. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Returns: - tuple[Tensor, Tensor]: The reduced deflection angle. + Parameters + ---------- + x: Tensor + angular x-coordinates from the observer perspective. + y: Tensor + angular y-coordinates from the observer perspective. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + tuple[Tensor, Tensor] + The reduced deflection angle. """ return self.raytrace_z1z2(x, y, torch.zeros_like(z_s), z_s, params) @@ -172,17 +189,26 @@ def surface_density( """ Calculate the projected mass density. - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: Projected mass density [solMass / Mpc^2]. - - Raises: - NotImplementedError: This method is not yet implemented. + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + Projected mass density [solMass / Mpc^2]. + + Raises + ------- + NotImplementedError + This method is not yet implemented. """ # TODO: rescale mass densities of each lens and sum raise NotImplementedError() @@ -200,17 +226,26 @@ def time_delay( """ Compute the time delay of light caused by the lensing. - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: Time delay caused by the lensing. - - Raises: - NotImplementedError: This method is not yet implemented. + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + Time delay caused by the lensing. + + Raises + ------ + NotImplementedError + This method is not yet implemented. """ # TODO: figure out how to compute this raise NotImplementedError() diff --git a/src/caustics/lenses/nfw.py b/src/caustics/lenses/nfw.py index 2a1d0297..4f70ed27 100644 --- a/src/caustics/lenses/nfw.py +++ b/src/caustics/lenses/nfw.py @@ -21,31 +21,47 @@ class NFW(ThinLens): The NFW profile is a spatial density profile of dark matter halo that arises in cosmological simulations. - Attributes: - z_l (Optional[Tensor]): Redshift of the lens. Default is None. - x0 (Optional[Tensor]): x-coordinate of the lens center in the lens plane. - Default is None. - y0 (Optional[Tensor]): y-coordinate of the lens center in the lens plane. - Default is None. - m (Optional[Tensor]): Mass of the lens. Default is None. - c (Optional[Tensor]): Concentration parameter of the lens. Default is None. - s (float): Softening parameter to avoid singularities at the center of the lens. - Default is 0.0. - use_case (str): Due to an idyosyncratic behaviour of PyTorch, the NFW/TNFW profile - specifically cant be both batchable and differentiable. You may select which version - you wish to use by setting this parameter to one of: batchable, differentiable. - - Methods: - get_scale_radius: Returns the scale radius of the lens. - get_scale_density: Returns the scale density of the lens. - get_convergence_s: Returns the dimensionless surface mass density of the lens. - _f: Helper method for computing deflection angles. - _g: Helper method for computing lensing potential. - _h: Helper method for computing reduced deflection angles. - deflection_angle_hat: Computes the reduced deflection angle. - deflection_angle: Computes the deflection angle. - convergence: Computes the convergence (dimensionless surface mass density). - potential: Computes the lensing potential. + Attributes + ----------- + z_l: Optional[Tensor] + Redshift of the lens. Default is None. + x0: Optional[Tensor] + x-coordinate of the lens center in the lens plane. Default is None. + y0: Optional[Tensor] + y-coordinate of the lens center in the lens plane. Default is None. + m: Optional[Tensor] + Mass of the lens. Default is None. + c: Optional[Tensor] + Concentration parameter of the lens. Default is None. + s: float + Softening parameter to avoid singularities at the center of the lens. Default is 0.0. + use_case: str + Due to an idyosyncratic behaviour of PyTorch, the NFW/TNFW profile + specifically cant be both batchable and differentiable. You may select which version + you wish to use by setting this parameter to one of: batchable, differentiable. + + Methods + ------- + get_scale_radius + Returns the scale radius of the lens. + get_scale_density + Returns the scale density of the lens. + get_convergence_s + Returns the dimensionless surface mass density of the lens. + _f + Helper method for computing deflection angles. + _g + Helper method for computing lensing potential. + _h + Helper method for computing reduced deflection angles. + deflection_angle_hat + Computes the reduced deflection angle. + deflection_angle + Computes the deflection angle. + convergence + Computes the convergence (dimensionless surface mass density). + potential + Computes the lensing potential. """ def __init__( @@ -63,19 +79,28 @@ def __init__( """ Initialize an instance of the NFW lens class. - Args: - name (str): Name of the lens instance. - cosmology (Cosmology): An instance of the Cosmology class which contains - information about the cosmological model and parameters. - z_l (Optional[Union[Tensor, float]]): Redshift of the lens. Default is None. - x0 (Optional[Union[Tensor, float]]): x-coordinate of the lens center in the lens plane. + Parameters + ---------- + name: str + Name of the lens instance. + cosmology: Cosmology + An instance of the Cosmology class which contains + information about the cosmological model and parameters. + z_l: Optional[Union[Tensor, float]] + Redshift of the lens. Default is None. + x0: Optional[Union[Tensor, float]] + x-coordinate of the lens center in the lens plane. Default is None. - y0 (Optional[Union[Tensor, float]]): y-coordinate of the lens center in the lens plane. + y0: Optional[Union[Tensor, float]] + y-coordinate of the lens center in the lens plane. Default is None. - m (Optional[Union[Tensor, float]]): Mass of the lens. Default is None. - c (Optional[Union[Tensor, float]]): Concentration parameter of the lens. Default is None. - s (float): Softening parameter to avoid singularities at the center of the lens. - Default is 0.0. + m: Optional[Union[Tensor, float]] + Mass of the lens. Default is None. + c: Optional[Union[Tensor, float]] + Concentration parameter of the lens. Default is None. + s: float + Softening parameter to avoid singularities at the center of the lens. + Default is 0.0. """ super().__init__(cosmology, z_l, name=name) @@ -102,14 +127,21 @@ def get_scale_radius( """ Calculate the scale radius of the lens. - Args: - z_l (Tensor): Redshift of the lens. - m (Tensor): Mass of the lens. - c (Tensor): Concentration parameter of the lens. - x (dict): Dynamic parameter container. - - Returns: - Tensor: The scale radius of the lens in Mpc. + Parameters + ---------- + z_l: Tensor + Redshift of the lens. + m: Tensor + Mass of the lens. + c: Tensor + Concentration parameter of the lens. + x: dict + Dynamic parameter container. + + Returns + ------- + Tensor + The scale radius of the lens in Mpc. """ critical_density = self.cosmology.critical_density(z_l, params) r_delta = (3 * m / (4 * pi * DELTA * critical_density)) ** (1 / 3) @@ -122,13 +154,19 @@ def get_scale_density( """ Calculate the scale density of the lens. - Args: - z_l (Tensor): Redshift of the lens. - c (Tensor): Concentration parameter of the lens. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The scale density of the lens in solar masses per Mpc cubed. + Parameters + ---------- + z_l: Tensor + Redshift of the lens. + c: Tensor + Concentration parameter of the lens. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The scale density of the lens in solar masses per Mpc cubed. """ return ( DELTA @@ -145,15 +183,23 @@ def get_convergence_s( """ Calculate the dimensionless surface mass density of the lens. - Args: - z_l (Tensor): Redshift of the lens. - z_s (Tensor): Redshift of the source. - m (Tensor): Mass of the lens. - c (Tensor): Concentration parameter of the lens. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The dimensionless surface mass density of the lens. + Parameters + ---------- + z_l: Tensor + Redshift of the lens. + z_s: Tensor + Redshift of the source. + m: Tensor + Mass of the lens. + c: Tensor + Concentration parameter of the lens. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The dimensionless surface mass density of the lens. """ critical_surface_density = self.cosmology.critical_surface_density( z_l, z_s, params @@ -169,11 +215,15 @@ def _f_differentiable(x: Tensor) -> Tensor: """ Helper method for computing deflection angles. - Args: - x (Tensor): The scaled radius (xi / xi_0). + Parameters + ---------- + x: Tensor + The scaled radius (xi / xi_0). - Returns: - Tensor: Result of the deflection angle computation. + Returns + ------- + Tensor + Result of the deflection angle computation. """ # TODO: generalize beyond torch, or patch Tensor f = torch.zeros_like(x) @@ -196,11 +246,15 @@ def _f_batchable(x: Tensor) -> Tensor: """ Helper method for computing deflection angles. - Args: - x (Tensor): The scaled radius (xi / xi_0). + Parameters + ---------- + x: Tensor + The scaled radius (xi / xi_0). - Returns: - Tensor: Result of the deflection angle computation. + Returns + ------- + Tensor + Result of the deflection angle computation. """ # TODO: generalize beyond torch, or patch Tensor return torch.where( @@ -218,11 +272,15 @@ def _g_differentiable(x: Tensor) -> Tensor: """ Helper method for computing lensing potential. - Args: - x (Tensor): The scaled radius (xi / xi_0). + Parameters + ---------- + x: Tensor + The scaled radius (xi / xi_0). - Returns: - Tensor: Result of the lensing potential computation. + Returns + ------- + Tensor + Result of the lensing potential computation. """ # TODO: generalize beyond torch, or patch Tensor term_1 = (x / 2).log() ** 2 @@ -236,11 +294,15 @@ def _g_batchable(x: Tensor) -> Tensor: """ Helper method for computing lensing potential. - Args: - x (Tensor): The scaled radius (xi / xi_0). + Parameters + ---------- + x: Tensor + The scaled radius (xi / xi_0). - Returns: - Tensor: Result of the lensing potential computation. + Returns + ------- + Tensor + Result of the lensing potential computation. """ # TODO: generalize beyond torch, or patch Tensor term_1 = (x / 2).log() ** 2 @@ -260,11 +322,15 @@ def _h_differentiable(x: Tensor) -> Tensor: """ Helper method for computing reduced deflection angles. - Args: - x (Tensor): The scaled radius (xi / xi_0). + Parameters + ---------- + x: Tensor + The scaled radius (xi / xi_0). - Returns: - Tensor: Result of the reduced deflection angle computation. + Returns + ------- + Tensor + Result of the reduced deflection angle computation. """ term_1 = (x / 2).log() term_2 = torch.ones_like(x) @@ -277,11 +343,15 @@ def _h_batchable(x: Tensor) -> Tensor: """ Helper method for computing reduced deflection angles. - Args: - x (Tensor): The scaled radius (xi / xi_0). + Parameters + ---------- + x: Tensor + The scaled radius (xi / xi_0). - Returns: - Tensor: Result of the reduced deflection angle computation. + Returns + ------- + Tensor + Result of the reduced deflection angle computation. """ term_1 = (x / 2).log() term_2 = torch.where( @@ -311,14 +381,21 @@ def reduced_deflection_angle( """ Compute the reduced deflection angle. - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Returns: - tuple[Tensor, Tensor]: The reduced deflection angles in the x and y directions. + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + tuple[Tensor, Tensor] + The reduced deflection angles in the x and y directions. """ x, y = translate_rotate(x, y, x0, y0) th = (x**2 + y**2).sqrt() + self.s @@ -362,14 +439,21 @@ def convergence( """ Compute the convergence (dimensionless surface mass density). - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The convergence (dimensionless surface mass density). + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The convergence (dimensionless surface mass density). """ x, y = translate_rotate(x, y, x0, y0) th = (x**2 + y**2).sqrt() + self.s @@ -398,14 +482,21 @@ def potential( """ Compute the lensing potential. - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The lensing potential. + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The lensing potential. """ x, y = translate_rotate(x, y, x0, y0) th = (x**2 + y**2).sqrt() + self.s diff --git a/src/caustics/lenses/pixelated_convergence.py b/src/caustics/lenses/pixelated_convergence.py index 84477eb1..b0cb2d7c 100644 --- a/src/caustics/lenses/pixelated_convergence.py +++ b/src/caustics/lenses/pixelated_convergence.py @@ -39,25 +39,38 @@ def __init__( grid using either Fast Fourier Transform (FFT) or a 2D convolution. - Attributes: - name (str): The name of the PixelatedConvergence object. - fov (float): The field of view in arcseconds. - n_pix (int): The number of pixels on each side of the grid. - cosmology (Cosmology): An instance of the cosmological parameters. - z_l (Optional[Tensor]): The redshift of the lens. - x0 (Optional[Tensor]): The x-coordinate of the center of the grid. - y0 (Optional[Tensor]): The y-coordinate of the center of the grid. - convergence_map (Optional[Tensor]): A 2D tensor representing the convergence map. - shape (Optional[tuple[int, ...]]): The shape of the convergence map. - convolution_mode (str, optional): The convolution mode for calculating deflection angles and lensing potential. - It can be either "fft" (Fast Fourier Transform) or "conv2d" (2D convolution). Default is "fft". - use_next_fast_len (bool, optional): If True, adds additional padding to speed up the FFT by calling - `scipy.fft.next_fast_len`. The speed boost can be substantial when `n_pix` is a multiple of a - small prime number. Default is True. - padding (str): Specifies the type of padding to use. "zero" will do zero padding, "circular" will do - cyclic boundaries. "reflect" will do reflection padding. "tile" will tile the image at 2x2 which - basically identical to circular padding, but is easier. Generally you should use either "zero" - or "tile". + Attributes + ---------- + name: string + The name of the PixelatedConvergence object. + fov: float + The field of view in arcseconds. + n_pix: int + The number of pixels on each side of the grid. + cosmology: Cosmology + An instance of the cosmological parameters. + z_l: Optional[Tensor] + The redshift of the lens. + x0: Optional[Tensor] + The x-coordinate of the center of the grid. + y0: Optional[Tensor] + The y-coordinate of the center of the grid. + convergence_map: Optional[Tensor] + A 2D tensor representing the convergence map. + shape: Optional[tuple[int, ...]] + The shape of the convergence map. + convolution_mode: (str, optional) + The convolution mode for calculating deflection angles and lensing potential. + It can be either "fft" (Fast Fourier Transform) or "conv2d" (2D convolution). Default is "fft". + use_next_fast_len: (bool, optional) + If True, adds additional padding to speed up the FFT by calling + `scipy.fft.next_fast_len`. The speed boost can be substantial when `n_pix` is a multiple of a + small prime number. Default is True. + padding: string + Specifies the type of padding to use. "zero" will do zero padding, "circular" will do + cyclic boundaries. "reflect" will do reflection padding. "tile" will tile the image at 2x2 which + basically identical to circular padding, but is easier. Generally you should use either "zero" + or "tile". """ @@ -108,9 +121,12 @@ def to( """ Move the ConvergenceGrid object and all its tensors to the specified device and dtype. - Args: - device (Optional[torch.device]): The target device to move the tensors to. - dtype (Optional[torch.dtype]): The target data type to cast the tensors to. + Parameters + ---------- + device: Optional[torch.device] + The target device to move the tensors to. + dtype: Optional[torch.dtype] + The target data type to cast the tensors to. """ super().to(device, dtype) self.potential_kernel = self.potential_kernel.to(device=device, dtype=dtype) @@ -127,11 +143,14 @@ def _fft2_padded(self, x: Tensor) -> Tensor: """ Compute the 2D Fast Fourier Transform (FFT) of a tensor with zero-padding. - Args: - x (Tensor): The input tensor to be transformed. + Parameters + x: Tensor + The input tensor to be transformed. - Returns: - Tensor: The 2D FFT of the input tensor with zero-padding. + Returns + ------- + Tensor + The 2D FFT of the input tensor with zero-padding. """ pad = 2 * self.n_pix if self.use_next_fast_len: @@ -153,11 +172,15 @@ def _unpad_fft(self, x: Tensor) -> Tensor: """ Remove padding from the result of a 2D FFT. - Args: - x (Tensor): The input tensor with padding. + Parameters + ---------- + x: Tensor + The input tensor with padding. - Returns: - Tensor: The input tensor without padding. + Returns + ------- + Tensor + The input tensor without padding. """ return torch.roll(x, (-self._s[0] // 2, -self._s[1] // 2), dims=(-2, -1))[ ..., : self.n_pix, : self.n_pix @@ -167,11 +190,15 @@ def _unpad_conv2d(self, x: Tensor) -> Tensor: """ Remove padding from the result of a 2D convolution. - Args: - x (Tensor): The input tensor with padding. + Parameters + ---------- + x: Tensor + The input tensor with padding. - Returns: - Tensor: The input tensor without padding. + Returns + ------- + Tensor + The input tensor without padding. """ return x # torch.roll(x, (-self.padding_range * self.ax_kernel.shape[0]//4,-self.padding_range * self.ax_kernel.shape[1]//4), dims = (-2,-1))[..., :self.n_pix, :self.n_pix] #[..., 1:, 1:] @@ -180,8 +207,10 @@ def convolution_mode(self): """ Get the convolution mode of the ConvergenceGrid object. - Returns: - str: The convolution mode, either "fft" or "conv2d". + Returns + ------- + string + The convolution mode, either "fft" or "conv2d". """ return self._convolution_mode @@ -190,8 +219,10 @@ def convolution_mode(self, convolution_mode: str): """ Set the convolution mode of the ConvergenceGrid object. - Args: - mode (str): The convolution mode to be set, either "fft" or "conv2d". + Parameters + ---------- + mode: string + The convolution mode to be set, either "fft" or "conv2d". """ if convolution_mode == "fft": # Create FFTs of kernels @@ -225,14 +256,21 @@ def reduced_deflection_angle( """ Compute the deflection angles at the specified positions using the given convergence map. - Args: - x (Tensor): The x-coordinates of the positions to compute the deflection angles for. - y (Tensor): The y-coordinates of the positions to compute the deflection angles for. - z_s (Tensor): The source redshift. - params (Packed, optional): A dictionary containing additional parameters. - - Returns: - tuple[Tensor, Tensor]: The x and y components of the deflection angles at the specified positions. + Parameters + ---------- + x: Tensor + The x-coordinates of the positions to compute the deflection angles for. + y: Tensor + The y-coordinates of the positions to compute the deflection angles for. + z_s: Tensor + The source redshift. + params: (Packed, optional) + A dictionary containing additional parameters. + + Returns + ------- + tuple[Tensor, Tensor] + The x and y components of the deflection angles at the specified positions. """ if self.convolution_mode == "fft": deflection_angle_x_map, deflection_angle_y_map = self._deflection_angle_fft( @@ -257,11 +295,15 @@ def _deflection_angle_fft(self, convergence_map: Tensor) -> tuple[Tensor, Tensor """ Compute the deflection angles using the Fast Fourier Transform (FFT) method. - Args: - convergence_map (Tensor): The 2D tensor representing the convergence map. + Parameters + ---------- + convergence_map: Tensor + The 2D tensor representing the convergence map. - Returns: - tuple[Tensor, Tensor]: The x and y components of the deflection angles. + Returns + ------- + tuple[Tensor, Tensor] + The x and y components of the deflection angles. """ convergence_tilde = self._fft2_padded(convergence_map) deflection_angle_x = torch.fft.irfft2( @@ -278,16 +320,20 @@ def _deflection_angle_conv2d( """ Compute the deflection angles using the 2D convolution method. - Args: - convergence_map (Tensor): The 2D tensor representing the convergence map. + Parameters + ---------- + convergence_map: Tensor + The 2D tensor representing the convergence map. - Returns: - tuple[Tensor, Tensor]: The x and y components of the deflection angles. + Returns + ------- + tuple[Tensor, Tensor] + The x and y components of the deflection angles. """ # Use convergence_map as kernel since the kernel is twice as large. Flip since # we actually want the cross-correlation. - pad = 2 * self.n_pix + 2 * self.n_pix convergence_map_flipped = convergence_map.flip((-1, -2))[ None, None ] # F.pad(, ((pad - self.n_pix)//2, (pad - self.n_pix)//2, (pad - self.n_pix)//2, (pad - self.n_pix)//2), mode = self.padding_mode) @@ -318,14 +364,21 @@ def potential( """ Compute the lensing potential at the specified positions using the given convergence map. - Args: - x (Tensor): The x-coordinates of the positions to compute the lensing potential for. - y (Tensor): The y-coordinates of the positions to compute the lensing potential for. - z_s (Tensor): The source redshift. - params (Packed, optional): A dictionary containing additional parameters. - - Returns: - Tensor: The lensing potential at the specified positions. + Parameters + ---------- + x: Tensor + The x-coordinates of the positions to compute the lensing potential for. + y: Tensor + The y-coordinates of the positions to compute the lensing potential for. + z_s: Tensor + The source redshift. + params: (Packed, optional) + A dictionary containing additional parameters. + + Returns + ------- + Tensor + The lensing potential at the specified positions. """ if self.convolution_mode == "fft": potential_map = self._potential_fft(convergence_map) @@ -342,11 +395,15 @@ def _potential_fft(self, convergence_map: Tensor) -> Tensor: """ Compute the lensing potential using the Fast Fourier Transform (FFT) method. - Args: - convergence_map (Tensor): The 2D tensor representing the convergence map. + Parameters + ---------- + convergence_map: Tensor + The 2D tensor representing the convergence map. - Returns: - Tensor: The lensing potential. + Returns + ------- + Tensor + The lensing potential. """ convergence_tilde = self._fft2_padded(convergence_map) potential = torch.fft.irfft2( @@ -358,11 +415,15 @@ def _potential_conv2d(self, convergence_map: Tensor) -> Tensor: """ Compute the lensing potential using the 2D convolution method. - Args: - convergence_map (Tensor): The 2D tensor representing the convergence map. + Parameters + ---------- + convergence_map: Tensor + The 2D tensor representing the convergence map. - Returns: - Tensor: The lensing potential. + Returns + ------- + Tensor + The lensing potential. """ # Use convergence_map as kernel since the kernel is twice as large. Flip since # we actually want the cross-correlation. @@ -389,17 +450,26 @@ def convergence( """ Compute the convergence at the specified positions. This method is not implemented. - Args: - x (Tensor): The x-coordinates of the positions to compute the convergence for. - y (Tensor): The y-coordinates of the positions to compute the convergence for. - z_s (Tensor): The source redshift. - params (Packed, optional): A dictionary containing additional parameters. - - Returns: - Tensor: The convergence at the specified positions. - - Raises: - NotImplementedError: This method is not implemented. + Parameters + ---------- + x: Tensor + The x-coordinates of the positions to compute the convergence for. + y: Tensor + The y-coordinates of the positions to compute the convergence for. + z_s: Tensor + The source redshift. + params: (Packed, optional) + A dictionary containing additional parameters. + + Returns + ------- + Tensor + The convergence at the specified positions. + + Raises + ------ + NotImplementedError + This method is not implemented. """ return interp2d( convergence_map, diff --git a/src/caustics/lenses/point.py b/src/caustics/lenses/point.py index a0efca85..0822c1ff 100644 --- a/src/caustics/lenses/point.py +++ b/src/caustics/lenses/point.py @@ -15,14 +15,22 @@ class Point(ThinLens): """ Class representing a point mass lens in strong gravitational lensing. - Attributes: - name (str): The name of the point lens. - cosmology (Cosmology): The cosmology used for calculations. - z_l (Optional[Union[Tensor, float]]): Redshift of the lens. - x0 (Optional[Union[Tensor, float]]): x-coordinate of the center of the lens. - y0 (Optional[Union[Tensor, float]]): y-coordinate of the center of the lens. - th_ein (Optional[Union[Tensor, float]]): Einstein radius of the lens. - s (float): Softening parameter to prevent numerical instabilities. + Attributes + ---------- + name: str + The name of the point lens. + cosmology: Cosmology + The cosmology used for calculations. + z_l: Optional[Union[Tensor, float]] + Redshift of the lens. + x0: Optional[Union[Tensor, float]] + x-coordinate of the center of the lens. + y0: Optional[Union[Tensor, float]] + y-coordinate of the center of the lens. + th_ein: Optional[Union[Tensor, float]] + Einstein radius of the lens. + s: float + Softening parameter to prevent numerical instabilities. """ def __init__( @@ -38,14 +46,22 @@ def __init__( """ Initialize the Point class. - Args: - name (str): The name of the point lens. - cosmology (Cosmology): The cosmology used for calculations. - z_l (Optional[Tensor]): Redshift of the lens. - x0 (Optional[Tensor]): x-coordinate of the center of the lens. - y0 (Optional[Tensor]): y-coordinate of the center of the lens. - th_ein (Optional[Tensor]): Einstein radius of the lens. - s (float): Softening parameter to prevent numerical instabilities. + Parameters + ---------- + name: string + The name of the point lens. + cosmology: Cosmology + The cosmology used for calculations. + z_l: Optional[Tensor] + Redshift of the lens. + x0: Optional[Tensor] + x-coordinate of the center of the lens. + y0: Optional[Tensor] + y-coordinate of the center of the lens. + th_ein: Optional[Tensor] + Einstein radius of the lens. + s: float + Softening parameter to prevent numerical instabilities. """ super().__init__(cosmology, z_l, name=name) @@ -71,14 +87,21 @@ def reduced_deflection_angle( """ Compute the deflection angles. - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Returns: - tuple[Tensor, Tensor]: The deflection angles in the x and y directions. + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + tuple[Tensor, Tensor] + The deflection angles in the x and y directions. """ x, y = translate_rotate(x, y, x0, y0) th = (x**2 + y**2).sqrt() + self.s @@ -103,14 +126,21 @@ def potential( """ Compute the lensing potential. - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The lensing potential. + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The lensing potential. """ x, y = translate_rotate(x, y, x0, y0) th = (x**2 + y**2).sqrt() + self.s @@ -133,14 +163,21 @@ def convergence( """ Compute the convergence (dimensionless surface mass density). - Args: - x (Tensor): x-coordinates in the lens plane. - y (Tensor): y-coordinates in the lens plane. - z_s (Tensor): Redshifts of the sources. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The convergence (dimensionless surface mass density). + Parameters + ---------- + x: Tensor + x-coordinates in the lens plane. + y: Tensor + y-coordinates in the lens plane. + z_s: Tensor + Redshifts of the sources. + params: (Packed, optional) + Dynamic parameter container. + + Returns + -------- + Tensor + The convergence (dimensionless surface mass density). """ x, y = translate_rotate(x, y, x0, y0) return torch.where((x == 0) & (y == 0), torch.inf, 0.0) diff --git a/src/caustics/lenses/pseudo_jaffe.py b/src/caustics/lenses/pseudo_jaffe.py index cca8ab09..f6990841 100644 --- a/src/caustics/lenses/pseudo_jaffe.py +++ b/src/caustics/lenses/pseudo_jaffe.py @@ -19,16 +19,26 @@ class PseudoJaffe(ThinLens): based on `Eliasdottir et al 2007 `_ and the `lenstronomy` source code. - Attributes: - name (str): The name of the Pseudo Jaffe lens. - cosmology (Cosmology): The cosmology used for calculations. - z_l (Optional[Union[Tensor, float]]): Redshift of the lens. - x0 (Optional[Union[Tensor, float]]): x-coordinate of the center of the lens (arcsec). - y0 (Optional[Union[Tensor, float]]): y-coordinate of the center of the lens (arcsec). - mass (Optional[Union[Tensor, float]]): Total mass of the lens (Msol). - core_radius (Optional[Union[Tensor, float]]): Core radius of the lens (arcsec). - scale_radius (Optional[Union[Tensor, float]]): Scaling radius of the lens (arcsec). - s (float): Softening parameter to prevent numerical instabilities. + Attributes + ---------- + name: string + The name of the Pseudo Jaffe lens. + cosmology: Cosmology + The cosmology used for calculations. + z_l: Optional[Union[Tensor, float]] + Redshift of the lens. + x0: Optional[Union[Tensor, float]] + x-coordinate of the center of the lens (arcsec). + y0: Optional[Union[Tensor, float]] + y-coordinate of the center of the lens (arcsec). + mass: Optional[Union[Tensor, float]] + Total mass of the lens (Msol). + core_radius: Optional[Union[Tensor, float]] + Core radius of the lens (arcsec). + scale_radius: Optional[Union[Tensor, float]] + Scaling radius of the lens (arcsec). + s: float + Softening parameter to prevent numerical instabilities. """ def __init__( @@ -46,16 +56,26 @@ def __init__( """ Initialize the PseudoJaffe class. - Args: - name (str): The name of the Pseudo Jaffe lens. - cosmology (Cosmology): The cosmology used for calculations. - z_l (Optional[Tensor]): Redshift of the lens. - x0 (Optional[Tensor]): x-coordinate of the center of the lens. - y0 (Optional[Tensor]): y-coordinate of the center of the lens. - mass (Optional[Tensor]): Total mass of the lens (Msol). - core_radius (Optional[Tensor]): Core radius of the lens. - scale_radius (Optional[Tensor]): Scaling radius of the lens. - s (float): Softening parameter to prevent numerical instabilities. + Parameters + ---------- + name: string + The name of the Pseudo Jaffe lens. + cosmology: Cosmology + The cosmology used for calculations. + z_l: Optional[Tensor] + Redshift of the lens. + x0: Optional[Tensor] + x-coordinate of the center of the lens. + y0: Optional[Tensor] + y-coordinate of the center of the lens. + mass: Optional[Tensor] + Total mass of the lens (Msol). + core_radius: Optional[Tensor] + Core radius of the lens. + scale_radius: Optional[Tensor] + Scaling radius of the lens. + s: float + Softening parameter to prevent numerical instabilities. """ super().__init__(cosmology, z_l, name=name) @@ -109,13 +129,19 @@ def mass_enclosed_2d( """ Calculate the mass enclosed within a two-dimensional radius. - Args: - theta (Tensor): Radius at which to calculate enclosed mass (arcsec). - z_s (Tensor): Source redshift. - params (Packed, optional): Dynamic parameter container. + Parameters + ---------- + theta: Tensor + Radius at which to calculate enclosed mass (arcsec). + z_s: Tensor + Source redshift. + params: (Packed, optional) + Dynamic parameter container. - Returns: - Tensor: The mass enclosed within the given radius. + Returns + ------- + Tensor + The mass enclosed within the given radius. """ theta = theta + self.s d_l = self.cosmology.angular_diameter_distance(z_l, params) @@ -150,16 +176,25 @@ def central_convergence( """ Compute the central convergence. - Args: - z_l (Tensor): Lens redshift. - z_s (Tensor): Source redshift. - rho_0 (Tensor): Central mass density. - core_radius (Tensor): Core radius of the lens (must be in Mpc). - scale_radius (Tensor): Scaling radius of the lens (must be in Mpc). - cosmology (Cosmology): The cosmology used for calculations. + Parameters + ----------- + z_l: Tensor + Lens redshift. + z_s: Tensor + Source redshift. + rho_0: Tensor + Central mass density. + core_radius: Tensor + Core radius of the lens (must be in Mpc). + scale_radius: Tensor + Scaling radius of the lens (must be in Mpc). + cosmology: Cosmology + The cosmology used for calculations. - Returns: - Tensor: The central convergence. + Returns + -------- + Tensor + The central convergence. """ return ( pi @@ -188,14 +223,21 @@ def reduced_deflection_angle( ) -> tuple[Tensor, Tensor]: """Calculate the deflection angle. - Args: - x (Tensor): x-coordinate of the lens. - y (Tensor): y-coordinate of the lens. - z_s (Tensor): Source redshift. - params (Packed, optional): Dynamic parameter container. + Parameters + ---------- + x: Tensor + x-coordinate of the lens. + y: Tensor + y-coordinate of the lens. + z_s: Tensor + Source redshift. + params: (Packed, optional) + Dynamic parameter container. - Returns: - Tuple[Tensor, Tensor]: The deflection angle in the x and y directions. + Returns + -------- + Tuple[Tensor, Tensor] + The deflection angle in the x and y directions. """ x, y = translate_rotate(x, y, x0, y0) R = (x**2 + y**2).sqrt() + self.s @@ -233,14 +275,21 @@ def potential( """ Compute the lensing potential. This calculation is based on equation A18. - Args: - x (Tensor): x-coordinate of the lens. - y (Tensor): y-coordinate of the lens. - z_s (Tensor): Source redshift. - params (Packed, optional): Dynamic parameter container. + Parameters + -------- + x: Tensor + x-coordinate of the lens. + y: Tensor + y-coordinate of the lens. + z_s: Tensor + Source redshift. + params: (Packed, optional) + Dynamic parameter container. - Returns: - Tensor: The lensing potential. + Returns + -------- + Tensor + The lensing potential. """ x, y = translate_rotate(x, y, x0, y0) R_squared = x**2 + y**2 + self.s @@ -278,14 +327,21 @@ def convergence( """ Calculate the projected mass density, based on equation A6. - Args: - x (Tensor): x-coordinate of the lens. - y (Tensor): y-coordinate of the lens. - z_s (Tensor): Source redshift. - params (Packed, optional): Dynamic parameter container. + Parameters + ----------- + x: Tensor + x-coordinate of the lens. + y: Tensor + y-coordinate of the lens. + z_s: Tensor + Source redshift. + params: (Packed, optional) + Dynamic parameter container. - Returns: - Tensor: The projected mass density. + Returns + ------- + Tensor + The projected mass density. """ x, y = translate_rotate(x, y, x0, y0) R_squared = x**2 + y**2 + self.s diff --git a/src/caustics/lenses/sie.py b/src/caustics/lenses/sie.py index f5bcd917..5e9e50b5 100644 --- a/src/caustics/lenses/sie.py +++ b/src/caustics/lenses/sie.py @@ -15,16 +15,26 @@ class SIE(ThinLens): A class representing a Singular Isothermal Ellipsoid (SIE) strong gravitational lens model. This model is based on Keeton 2001, which can be found at https://arxiv.org/abs/astro-ph/0102341. - Attributes: - name (str): The name of the lens. - cosmology (Cosmology): An instance of the Cosmology class. - z_l (Optional[Union[Tensor, float]]): The redshift of the lens. - x0 (Optional[Union[Tensor, float]]): The x-coordinate of the lens center. - y0 (Optional[Union[Tensor, float]]): The y-coordinate of the lens center. - q (Optional[Union[Tensor, float]]): The axis ratio of the lens. - phi (Optional[Union[Tensor, float]]): The orientation angle of the lens (position angle). - b (Optional[Union[Tensor, float]]): The Einstein radius of the lens. - s (float): The core radius of the lens (defaults to 0.0). + Attributes + ---------- + name: str + The name of the lens. + cosmology: Cosmology + An instance of the Cosmology class. + z_l: Optional[Union[Tensor, float]] + The redshift of the lens. + x0: Optional[Union[Tensor, float]] + The x-coordinate of the lens center. + y0: Optional[Union[Tensor, float]] + The y-coordinate of the lens center. + q: Optional[Union[Tensor, float]] + The axis ratio of the lens. + phi: Optional[Union[Tensor, float]] + The orientation angle of the lens (position angle). + b: Optional[Union[Tensor, float]] + The Einstein radius of the lens. + s: float + The core radius of the lens (defaults to 0.0). """ def __init__( @@ -55,13 +65,19 @@ def _get_potential(self, x, y, q): """ Compute the radial coordinate in the lens plane. - Args: - x (Tensor): The x-coordinate in the lens plane. - y (Tensor): The y-coordinate in the lens plane. - q (Tensor): The axis ratio of the lens. - - Returns: - Tensor: The radial coordinate in the lens plane. + Parameters + ---------- + x: Tensor + The x-coordinate in the lens plane. + y: Tensor + The y-coordinate in the lens plane. + q: Tensor + The axis ratio of the lens. + + Returns + -------- + Tensor + The radial coordinate in the lens plane. """ return (q**2 * (x**2 + self.s**2) + y**2).sqrt() @@ -84,14 +100,21 @@ def reduced_deflection_angle( """ Calculate the physical deflection angle. - Args: - x (Tensor): The x-coordinate of the lens. - y (Tensor): The y-coordinate of the lens. - z_s (Tensor): The source redshift. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tuple[Tensor, Tensor]: The deflection angle in the x and y directions. + Parameters + ---------- + x: Tensor + The x-coordinate of the lens. + y: Tensor + The y-coordinate of the lens. + z_s: Tensor + The source redshift. + params: (Packed, optional) + Dynamic parameter container. + + Returns + -------- + Tuple[Tensor, Tensor] + The deflection angle in the x and y directions. """ x, y = translate_rotate(x, y, x0, y0, phi) psi = self._get_potential(x, y, q) @@ -120,14 +143,21 @@ def potential( """ Compute the lensing potential. - Args: - x (Tensor): The x-coordinate of the lens. - y (Tensor): The y-coordinate of the lens. - z_s (Tensor): The source redshift. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The lensing potential. + Parameters + ---------- + x: Tensor + The x-coordinate of the lens. + y: Tensor + The y-coordinate of the lens. + z_s: Tensor + The source redshift. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The lensing potential. """ ax, ay = self.reduced_deflection_angle(x, y, z_s, params) ax, ay = derotate(ax, ay, -phi) @@ -153,14 +183,21 @@ def convergence( """ Calculate the projected mass density. - Args: - x (Tensor): The x-coordinate of the lens. - y (Tensor): The y-coordinate of the lens. - z_s (Tensor): The source redshift. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The projected mass. + Parameters + ---------- + x: Tensor + The x-coordinate of the lens. + y: Tensor + The y-coordinate of the lens. + z_s: Tensor + The source redshift. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The projected mass. """ x, y = translate_rotate(x, y, x0, y0, phi) psi = self._get_potential(x, y, q) diff --git a/src/caustics/lenses/singleplane.py b/src/caustics/lenses/singleplane.py index 7c6ff11a..577d336e 100644 --- a/src/caustics/lenses/singleplane.py +++ b/src/caustics/lenses/singleplane.py @@ -15,10 +15,14 @@ class SinglePlane(ThinLens): A class for combining multiple thin lenses into a single lensing plane. This model inherits from the base `ThinLens` class. - Attributes: - name (str): The name of the single plane lens. - cosmology (Cosmology): An instance of the Cosmology class. - lenses (List[ThinLens]): A list of ThinLens objects that are being combined into a single lensing plane. + Attributes + ---------- + name: str + The name of the single plane lens. + cosmology: Cosmology + An instance of the Cosmology class. + lenses: List[ThinLens] + A list of ThinLens objects that are being combined into a single lensing plane. """ def __init__( @@ -46,14 +50,21 @@ def reduced_deflection_angle( """ Calculate the total deflection angle by summing the deflection angles of all individual lenses. - Args: - x (Tensor): The x-coordinate of the lens. - y (Tensor): The y-coordinate of the lens. - z_s (Tensor): The source redshift. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tuple[Tensor, Tensor]: The total deflection angle in the x and y directions. + Parameters + ---------- + x: Tensor + The x-coordinate of the lens. + y: Tensor + The y-coordinate of the lens. + z_s: Tensor + The source redshift. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tuple[Tensor, Tensor] + The total deflection angle in the x and y directions. """ ax = torch.zeros_like(x) ay = torch.zeros_like(x) @@ -76,14 +87,21 @@ def convergence( """ Calculate the total projected mass density by summing the mass densities of all individual lenses. - Args: - x (Tensor): The x-coordinate of the lens. - y (Tensor): The y-coordinate of the lens. - z_s (Tensor): The source redshift. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The total projected mass density. + Parameters + ---------- + x: Tensor + The x-coordinate of the lens. + y: Tensor + The y-coordinate of the lens. + z_s: Tensor + The source redshift. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The total projected mass density. """ convergence = torch.zeros_like(x) for lens in self.lenses: @@ -104,14 +122,21 @@ def potential( """ Compute the total lensing potential by summing the lensing potentials of all individual lenses. - Args: - x (Tensor): The x-coordinate of the lens. - y (Tensor): The y-coordinate of the lens. - z_s (Tensor): The source redshift. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The total lensing potential. + Parameters + ----------- + x: Tensor + The x-coordinate of the lens. + y: Tensor + The y-coordinate of the lens. + z_s: Tensor + The source redshift. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The total lensing potential. """ potential = torch.zeros_like(x) for lens in self.lenses: diff --git a/src/caustics/lenses/sis.py b/src/caustics/lenses/sis.py index 244fe6ea..bb72118f 100644 --- a/src/caustics/lenses/sis.py +++ b/src/caustics/lenses/sis.py @@ -15,14 +15,21 @@ class SIS(ThinLens): A class representing the Singular Isothermal Sphere (SIS) model. This model inherits from the base `ThinLens` class. - Attributes: - name (str): The name of the SIS lens. - cosmology (Cosmology): An instance of the Cosmology class. - z_l (Optional[Union[Tensor, float]]): The lens redshift. - x0 (Optional[Union[Tensor, float]]): The x-coordinate of the lens center. - y0 (Optional[Union[Tensor, float]]): The y-coordinate of the lens center. + Attributes + ---------- + name: str + The name of the SIS lens. + cosmology: Cosmology + An instance of the Cosmology class. + z_l: Optional[Union[Tensor, float]] + The lens redshift. + x0: Optional[Union[Tensor, float]] + The x-coordinate of the lens center. + y0: Optional[Union[Tensor, float]] + The y-coordinate of the lens center. th_ein (Optional[Union[Tensor, float]]): The Einstein radius of the lens. - s (float): A smoothing factor, default is 0.0. + s: float + A smoothing factor, default is 0.0. """ def __init__( @@ -62,14 +69,21 @@ def reduced_deflection_angle( """ Calculate the deflection angle of the SIS lens. - Args: - x (Tensor): The x-coordinate of the lens. - y (Tensor): The y-coordinate of the lens. - z_s (Tensor): The source redshift. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tuple[Tensor, Tensor]: The deflection angle in the x and y directions. + Parameters + ---------- + x: Tensor + The x-coordinate of the lens. + y: Tensor + The y-coordinate of the lens. + z_s: Tensor + The source redshift. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tuple[Tensor, Tensor] + The deflection angle in the x and y directions. """ x, y = translate_rotate(x, y, x0, y0) R = (x**2 + y**2).sqrt() + self.s @@ -94,14 +108,21 @@ def potential( """ Compute the lensing potential of the SIS lens. - Args: - x (Tensor): The x-coordinate of the lens. - y (Tensor): The y-coordinate of the lens. - z_s (Tensor): The source redshift. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The lensing potential. + Parameters + ---------- + x: Tensor + The x-coordinate of the lens. + y: Tensor + The y-coordinate of the lens. + z_s: Tensor + The source redshift. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The lensing potential. """ x, y = translate_rotate(x, y, x0, y0) th = (x**2 + y**2).sqrt() + self.s @@ -124,14 +145,21 @@ def convergence( """ Calculate the projected mass density of the SIS lens. - Args: - x (Tensor): The x-coordinate of the lens. - y (Tensor): The y-coordinate of the lens. - z_s (Tensor): The source redshift. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The projected mass density. + Parameters + ---------- + x: Tensor + The x-coordinate of the lens. + y: Tensor + The y-coordinate of the lens. + z_s: Tensor + The source redshift. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The projected mass density. """ x, y = translate_rotate(x, y, x0, y0) th = (x**2 + y**2).sqrt() + self.s diff --git a/src/caustics/lenses/tnfw.py b/src/caustics/lenses/tnfw.py index 093ee53f..1132448d 100644 --- a/src/caustics/lenses/tnfw.py +++ b/src/caustics/lenses/tnfw.py @@ -29,7 +29,8 @@ class TNFW(ThinLens): https://ui.adsabs.harvard.edu/abs/2009JCAP...01..015B/abstract - Note: + Notes + ------ The mass `m` in the TNFW profile corresponds to the total mass of the lens. This is different from the NFW profile where the mass `m` parameter corresponds to the mass within R200. If you @@ -39,25 +40,37 @@ class TNFW(ThinLens): NFW profile, not a TNFW profile. This is in line with how lenstronomy inteprets the mass parameter. - Args: - name (str): Name of the lens instance. - cosmology (Cosmology): An instance of the Cosmology class which contains - information about the cosmological model and parameters. - z_l (Optional[Tensor]): Redshift of the lens. - x0 (Optional[Tensor]): Center of lens position on x-axis (arcsec). - y0 (Optional[Tensor]): Center of lens position on y-axis (arcsec). - mass (Optional[Tensor]): Mass of the lens (Msol). - scale_radius (Optional[Tensor]): Scale radius of the TNFW lens (arcsec). - tau (Optional[Tensor]): Truncation scale. Ratio of truncation radius to scale radius (rt/rs). - s (float): Softening parameter to avoid singularities at the center of the lens. - Default is 0.0. - interpret_m_total_mass (bool): Indicates how to interpret the mass variable "m". If true - the mass is intepreted as the total mass of the halo (good because it makes sense). If - false it is intepreted as what the mass would have been within R200 of a an NFW that - isn't truncated (good because it is easily compared with an NFW). - use_case (str): Due to an idyosyncratic behaviour of PyTorch, the NFW/TNFW profile - specifically cant be both batchable and differentiable. You may select which version - you wish to use by setting this parameter to one of: batchable, differentiable. + Parameters + ----- + name: string + Name of the lens instance. + cosmology: Cosmology + An instance of the Cosmology class which contains + information about the cosmological model and parameters. + z_l: Optional[Tensor] + Redshift of the lens. + x0: Optional[Tensor] + Center of lens position on x-axis (arcsec). + y0: Optional[Tensor] + Center of lens position on y-axis (arcsec). + mass: Optional[Tensor] + Mass of the lens (Msol). + scale_radius: Optional[Tensor] + Scale radius of the TNFW lens (arcsec). + tau: Optional[Tensor] + Truncation scale. Ratio of truncation radius to scale radius (rt/rs). + s: float + Softening parameter to avoid singularities at the center of the lens. + Default is 0.0. + interpret_m_total_mass: boolean + Indicates how to interpret the mass variable "m". If true + the mass is intepreted as the total mass of the halo (good because it makes sense). If + false it is intepreted as what the mass would have been within R200 of a an NFW that + isn't truncated (good because it is easily compared with an NFW). + use_case: str + Due to an idyosyncratic behaviour of PyTorch, the NFW/TNFW profile + specifically cant be both batchable and differentiable. You may select which version + you wish to use by setting this parameter to one of: batchable, differentiable. """ @@ -143,17 +156,27 @@ def get_concentration( """ Calculate the scale radius of the lens. This is the same formula used for the classic NFW profile. - Args: - z_l (Tensor): Redshift of the lens. - x0 (Tensor): Center of lens position on x-axis (arcsec). - y0 (Tensor): Center of lens position on y-axis (arcsec). - mass (Optional[Tensor]): Mass of the lens (Msol). - scale_radius (Optional[Tensor]): Scale radius of the TNFW lens (arcsec). - tau (Optional[Tensor]): Truncation scale. Ratio of truncation radius to scale radius (rt/rs). - params (dict): Dynamic parameter container. - - Returns: - Tensor: The scale radius of the lens in Mpc. + Parameters + ---------- + z_l: Tensor + Redshift of the lens. + x0: Tensor + Center of lens position on x-axis (arcsec). + y0: Tensor + Center of lens position on y-axis (arcsec). + mass: Optional[Tensor] + Mass of the lens (Msol). + scale_radius: Optional[Tensor] + Scale radius of the TNFW lens (arcsec). + tau: Optional[Tensor] + Truncation scale. Ratio of truncation radius to scale radius (rt/rs). + params: dict + Dynamic parameter container. + + Returns + ------- + Tensor + The scale radius of the lens in Mpc. """ critical_density = self.cosmology.critical_density(z_l, params) d_l = self.cosmology.angular_diameter_distance(z_l, params) @@ -176,17 +199,27 @@ def get_truncation_radius( """ Calculate the truncation radius of the TNFW lens. - Args: - z_l (Tensor): Redshift of the lens. - x0 (Tensor): Center of lens position on x-axis (arcsec). - y0 (Tensor): Center of lens position on y-axis (arcsec). - mass (Optional[Tensor]): Mass of the lens (Msol). - scale_radius (Optional[Tensor]): Scale radius of the TNFW lens (arcsec). - tau (Optional[Tensor]): Truncation scale. Ratio of truncation radius to scale radius (rt/rs). - params (dict): Dynamic parameter container. - - Returns: - Tensor: The truncation radius of the lens in arcsec. + Parameters + ---------- + z_l: Tensor + Redshift of the lens. + x0: Tensor + Center of lens position on x-axis (arcsec). + y0: Tensor + Center of lens position on y-axis (arcsec). + mass: Optional[Tensor] + Mass of the lens (Msol). + scale_radius: Optional[Tensor] + Scale radius of the TNFW lens (arcsec). + tau: Optional[Tensor] + Truncation scale. Ratio of truncation radius to scale radius (rt/rs). + params: dictionary + Dynamic parameter container. + + Returns + ------- + Tensor + The truncation radius of the lens in arcsec. """ return tau * scale_radius @@ -206,17 +239,27 @@ def get_M0( """ Calculate the reference mass. This is an abstract reference mass used internally in the equations from Baltz et al. 2009. - Args: - z_l (Tensor): Redshift of the lens. - x0 (Tensor): Center of lens position on x-axis (arcsec). - y0 (Tensor): Center of lens position on y-axis (arcsec). - mass (Optional[Tensor]): Mass of the lens (Msol). - scale_radius (Optional[Tensor]): Scale radius of the TNFW lens (arcsec). - tau (Optional[Tensor]): Truncation scale. Ratio of truncation radius to scale radius (rt/rs). - params (dict): Dynamic parameter container. - - Returns: - Tensor: The reference mass of the lens in Msol. + Parameters + ---------- + z_l: Tensor + Redshift of the lens. + x0: Tensor + Center of lens position on x-axis (arcsec). + y0: Tensor + Center of lens position on y-axis (arcsec). + mass: Optional[Tensor] + Mass of the lens (Msol). + scale_radius: Optional[Tensor] + Scale radius of the TNFW lens (arcsec). + tau: Optional[Tensor] + Truncation scale. Ratio of truncation radius to scale radius (rt/rs). + params: dictionary + Dynamic parameter container. + + Returns + ------- + Tensor + The reference mass of the lens in Msol. """ if self.interpret_m_total_mass: return ( @@ -252,17 +295,27 @@ def get_scale_density( """ Calculate the scale density of the lens. - Args: - z_l (Tensor): Redshift of the lens. - x0 (Tensor): Center of lens position on x-axis (arcsec). - y0 (Tensor): Center of lens position on y-axis (arcsec). - mass (Optional[Tensor]): Mass of the lens (Msol). - scale_radius (Optional[Tensor]): Scale radius of the TNFW lens (arcsec). - tau (Optional[Tensor]): Truncation scale. Ratio of truncation radius to scale radius (rt/rs). - params (dict): Dynamic parameter container. - - Returns: - Tensor: The scale density of the lens in solar masses per Mpc cubed. + Parameters + ---------- + z_l: Tensor + Redshift of the lens. + x0: Tensor + Center of lens position on x-axis (arcsec). + y0: Tensor + Center of lens position on y-axis (arcsec). + mass: Optional[Tensor] + Mass of the lens (Msol). + scale_radius: Optional[Tensor] + Scale radius of the TNFW lens (arcsec). + tau: Optional[Tensor] + Truncation scale. Ratio of truncation radius to scale radius (rt/rs). + params: dict + Dynamic parameter container. + + Returns + -------- + Tensor + The scale density of the lens in solar masses per Mpc cubed. """ c = self.get_concentration(params) return ( @@ -292,17 +345,27 @@ def convergence( """ TNFW convergence as given in Baltz et al. 2009. This is unitless since it is Sigma(x) / Sigma_crit. - Args: - z_l (Tensor): Redshift of the lens. - x0 (Tensor): Center of lens position on x-axis (arcsec). - y0 (Tensor): Center of lens position on y-axis (arcsec). - mass (Optional[Tensor]): Mass of the lens (Msol). - scale_radius (Optional[Tensor]): Scale radius of the TNFW lens (arcsec). - tau (Optional[Tensor]): Truncation scale. Ratio of truncation radius to scale radius (rt/rs). - params (dict): Dynamic parameter container. - - Returns: - Tensor: unitless convergence at requested position + Parameters + ---------- + z_l: Tensor + Redshift of the lens. + x0: Tensor + Center of lens position on x-axis (arcsec). + y0: Tensor + Center of lens position on y-axis (arcsec). + mass: Optional[Tensor] + Mass of the lens (Msol). + scale_radius: Optional[Tensor] + Scale radius of the TNFW lens (arcsec). + tau: Optional[Tensor] + Truncation scale. Ratio of truncation radius to scale radius (rt/rs). + params: dict + Dynamic parameter container. + + Returns + --------- + Tensor + unitless convergence at requested position """ x, y = translate_rotate(x, y, x0, y0) @@ -343,17 +406,27 @@ def mass_enclosed_2d( """ Total projected mass (Msol) within a radius r (arcsec). - Args: - z_l (Tensor): Redshift of the lens. - x0 (Tensor): Center of lens position on x-axis (arcsec). - y0 (Tensor): Center of lens position on y-axis (arcsec). - mass (Optional[Tensor]): Mass of the lens (Msol). - scale_radius (Optional[Tensor]): Scale radius of the TNFW lens (arcsec). - tau (Optional[Tensor]): Truncation scale. Ratio of truncation radius to scale radius (rt/rs). - params (dict): Dynamic parameter container. - - Returns: - Tensor: Integrated mass projected in infinite cylinder within radius r. + Parameters + ----------- + z_l: Tensor + Redshift of the lens. + x0: Tensor + Center of lens position on x-axis (arcsec). + y0: Tensor + Center of lens position on y-axis (arcsec). + mass: Optional[Tensor] + Mass of the lens (Msol). + scale_radius: Optional[Tensor] + Scale radius of the TNFW lens (arcsec). + tau: Optional[Tensor] + Truncation scale. Ratio of truncation radius to scale radius (rt/rs). + params: dict + Dynamic parameter container. + + Returns + ------- + Tensor + Integrated mass projected in infinite cylinder within radius r. """ g = r / scale_radius t2 = tau**2 @@ -388,17 +461,27 @@ def physical_deflection_angle( naturally represented as a physical deflection angle, this is easily internally converted to a reduced deflection angle. - Args: - z_l (Tensor): Redshift of the lens. - x0 (Tensor): Center of lens position on x-axis (arcsec). - y0 (Tensor): Center of lens position on y-axis (arcsec). - mass (Optional[Tensor]): Mass of the lens (Msol). - scale_radius (Optional[Tensor]): Scale radius of the TNFW lens (arcsec). - tau (Optional[Tensor]): Truncation scale. Ratio of truncation radius to scale radius (rt/rs). - params (dict): Dynamic parameter container. - - Returns: - tuple[Tensor, Tensor]: The physical deflection angles in the x and y directions (arcsec). + Parameters + ---------- + z_l: Tensor + Redshift of the lens. + x0: Tensor + Center of lens position on x-axis (arcsec). + y0: Tensor + Center of lens position on y-axis (arcsec). + mass: Optional[Tensor] + Mass of the lens (Msol). + scale_radius: Optional[Tensor] + Scale radius of the TNFW lens (arcsec). + tau: Optional[Tensor] + Truncation scale. Ratio of truncation radius to scale radius (rt/rs). + params: dict + Dynamic parameter container. + + Returns + -------- + tuple[Tensor, Tensor] + The physical deflection angles in the x and y directions (arcsec). """ d_l = self.cosmology.angular_diameter_distance(z_l, params) @@ -434,17 +517,27 @@ def potential( TODO: convert to dimensionless potential. - Args: - z_l (Tensor): Redshift of the lens. - x0 (Tensor): Center of lens position on x-axis (arcsec). - y0 (Tensor): Center of lens position on y-axis (arcsec). - mass (Optional[Tensor]): Mass of the lens (Msol). - scale_radius (Optional[Tensor]): Scale radius of the TNFW lens (arcsec). - tau (Optional[Tensor]): Truncation scale. Ratio of truncation radius to scale radius (rt/rs). - params (dict): Dynamic parameter container. - - Returns: - Tensor: The lensing potential. + Parameters + ----------- + z_l: Tensor + Redshift of the lens. + x0: Tensor + Center of lens position on x-axis (arcsec). + y0: Tensor + Center of lens position on y-axis (arcsec). + mass: Optional[Tensor] + Mass of the lens (Msol). + scale_radius: Optional[Tensor] + Scale radius of the TNFW lens (arcsec). + tau: Optional[Tensor] + Truncation scale. Ratio of truncation radius to scale radius (rt/rs). + params: dict + Dynamic parameter container. + + Returns + ------- + Tensor + The lensing potential. """ x, y = translate_rotate(x, y, x0, y0) r = (x**2 + y**2).sqrt() + self.s diff --git a/src/caustics/lenses/utils.py b/src/caustics/lenses/utils.py index 4ccb2b54..3d4d52ff 100644 --- a/src/caustics/lenses/utils.py +++ b/src/caustics/lenses/utils.py @@ -16,14 +16,20 @@ def get_pix_jacobian( (:math:`\\partial \beta / \\partial \theta`). This is done at a single point on the lensing plane. - Args: - raytrace: A function that maps the lensing plane coordinates to the source plane coordinates. - x (Tensor): The x-coordinate on the lensing plane. - y (Tensor): The y-coordinate on the lensing plane. - z_s (Tensor): The redshift of the source. - - Returns: - The Jacobian matrix of the image position with respect to the source position at the given point. + Parameters + ----------- + raytrace: function + A function that maps the lensing plane coordinates to the source plane coordinates. + x: Tensor + The x-coordinate on the lensing plane. + y: Tensor + The y-coordinate on the lensing plane. + z_s: Tensor + The redshift of the source. + + Returns + -------- + The Jacobian matrix of the image position with respect to the source position at the given point. """ jac = torch.func.jacfwd(raytrace, (0, 1))(x, y, z_s) # type: ignore @@ -35,13 +41,20 @@ def get_pix_magnification(raytrace, x, y, z_s) -> Tensor: Computes the magnification at a single point on the lensing plane. The magnification is derived from the determinant of the Jacobian matrix of the image position with respect to the source position. - Args: - raytrace: A function that maps the lensing plane coordinates to the source plane coordinates. - x (Tensor): The x-coordinate on the lensing plane. - y (Tensor): The y-coordinate on the lensing plane. - z_s (Tensor): The redshift of the source. - - Returns: + Parameters + ---------- + raytrace: function + A function that maps the lensing plane coordinates to the source plane coordinates. + x: Tensor + The x-coordinate on the lensing plane. + y: Tensor + The y-coordinate on the lensing plane. + z_s: Tensor + The redshift of the source. + + Returns + ------- + Tensor The magnification at the given point on the lensing plane. """ jac = get_pix_jacobian(raytrace, x, y, z_s) @@ -53,13 +66,20 @@ def get_magnification(raytrace, x, y, z_s) -> Tensor: Computes the magnification over a grid on the lensing plane. This is done by calling `get_pix_magnification` for each point on the grid. - Args: - raytrace: A function that maps the lensing plane coordinates to the source plane coordinates. - x (Tensor): The x-coordinates on the lensing plane. - y (Tensor): The y-coordinates on the lensing plane. - z_s (Tensor): The redshift of the source. - - Returns: + Parameters + ---------- + raytrace: function + A function that maps the lensing plane coordinates to the source plane coordinates. + x: Tensor + The x-coordinates on the lensing plane. + y: Tensor + The y-coordinates on the lensing plane. + z_s: Tensor + The redshift of the source. + + Returns + -------- + Tensor A tensor representing the magnification at each point on the grid. """ return vmap_n(get_pix_magnification, 2, (None, 0, 0, None))(raytrace, x, y, z_s) diff --git a/src/caustics/light/base.py b/src/caustics/light/base.py index f6a0e915..e457a58f 100644 --- a/src/caustics/light/base.py +++ b/src/caustics/light/base.py @@ -28,21 +28,28 @@ def brightness( Abstract method that calculates the brightness of the source at the given coordinates. This method is expected to be implemented in any class that derives from Source. - Args: - x (Tensor): The x-coordinate(s) at which to calculate the source brightness. + Parameters + ---------- + x: Tensor + The x-coordinate(s) at which to calculate the source brightness. This could be a single value or a tensor of values. - y (Tensor): The y-coordinate(s) at which to calculate the source brightness. + y: Tensor + The y-coordinate(s) at which to calculate the source brightness. This could be a single value or a tensor of values. - params (Packed, optional): Dynamic parameter container that might be required to calculate - the brightness. The exact contents will depend on the specific implementation in derived classes. + params: (Packed, optional) + Dynamic parameter container that might be required to calculate + the brightness. The exact contents will depend on the specific implementation in derived classes. - Returns: - Tensor: The brightness of the source at the given coordinate(s). The exact form of the output + Returns + ------- + Tensor + The brightness of the source at the given coordinate(s). The exact form of the output will depend on the specific implementation in the derived class. - Note: + Note + ----- This method must be overridden in any class that inherits from `Source`. """ ... diff --git a/src/caustics/light/sersic.py b/src/caustics/light/sersic.py index 4e37263a..c913534d 100644 --- a/src/caustics/light/sersic.py +++ b/src/caustics/light/sersic.py @@ -17,16 +17,26 @@ class Sersic(Source): The Sersic profile is often used to describe elliptical galaxies and spiral galaxies' bulges. - Attributes: - x0 (Optional[Tensor]): The x-coordinate of the Sersic source's center. - y0 (Optional[Tensor]): The y-coordinate of the Sersic source's center. - q (Optional[Tensor]): The axis ratio of the Sersic source. - phi (Optional[Tensor]): The orientation of the Sersic source (position angle). - n (Optional[Tensor]): The Sersic index, which describes the degree of concentration of the source. - Re (Optional[Tensor]): The scale length of the Sersic source. - Ie (Optional[Tensor]): The intensity at the effective radius. - s (float): A small constant for numerical stability. - lenstronomy_k_mode (bool): A flag indicating whether to use lenstronomy to compute the value of k. + Attributes + ----------- + x0: Optional[Tensor] + The x-coordinate of the Sersic source's center. + y0: Optional[Tensor] + The y-coordinate of the Sersic source's center. + q: Optional[Tensor] + The axis ratio of the Sersic source. + phi: Optional[Tensor] + The orientation of the Sersic source (position angle). + n: Optional[Tensor] + The Sersic index, which describes the degree of concentration of the source. + Re: Optional[Tensor] + The scale length of the Sersic source. + Ie: Optional[Tensor] + The intensity at the effective radius. + s: float + A small constant for numerical stability. + lenstronomy_k_mode: bool + A flag indicating whether to use lenstronomy to compute the value of k. """ def __init__( @@ -45,17 +55,28 @@ def __init__( """ Constructs the `Sersic` object with the given parameters. - Args: - name (str): The name of the source. - x0 (Optional[Tensor]): The x-coordinate of the Sersic source's center. - y0 (Optional[Tensor]): The y-coordinate of the Sersic source's center. - q (Optional[Tensor]): The axis ratio of the Sersic source. - phi (Optional[Tensor]): The orientation of the Sersic source. - n (Optional[Tensor]): The Sersic index, which describes the degree of concentration of the source. - Re (Optional[Tensor]): The scale length of the Sersic source. - Ie (Optional[Tensor]): The intensity at the effective radius. - s (float): A small constant for numerical stability. - use_lenstronomy_k (bool): A flag indicating whether to use lenstronomy to compute the value of k. + Parameters + ---------- + name: str + The name of the source. + x0: Optional[Tensor] + The x-coordinate of the Sersic source's center. + y0: Optional[Tensor] + The y-coordinate of the Sersic source's center. + q: Optional[Tensor]) + The axis ratio of the Sersic source. + phi: Optional[Tensor] + The orientation of the Sersic source. + n: Optional[Tensor] + The Sersic index, which describes the degree of concentration of the source. + Re: Optional[Tensor] + The scale length of the Sersic source. + Ie: Optional[Tensor] + The intensity at the effective radius. + s: float + A small constant for numerical stability. + use_lenstronomy_k: bool + A flag indicating whether to use lenstronomy to compute the value of k. """ super().__init__(name=name) self.add_param("x0", x0) @@ -89,17 +110,24 @@ def brightness( Implements the `brightness` method for `Sersic`. The brightness at a given point is determined by the Sersic profile formula. - Args: - x (Tensor): The x-coordinate(s) at which to calculate the source brightness. - This could be a single value or a tensor of values. - y (Tensor): The y-coordinate(s) at which to calculate the source brightness. - This could be a single value or a tensor of values. - params (Packed, optional): Dynamic parameter container. - - Returns: - Tensor: The brightness of the source at the given point(s). The output tensor has the same shape as `x` and `y`. - - Notes: + Parameters + ---------- + x: Tensor + The x-coordinate(s) at which to calculate the source brightness. + This could be a single value or a tensor of values. + y: Tensor + The y-coordinate(s) at which to calculate the source brightness. + This could be a single value or a tensor of values. + params: (Packed, optional) + Dynamic parameter container. + + Returns + ------- + Tensor + The brightness of the source at the given point(s). The output tensor has the same shape as `x` and `y`. + + Notes + ----- The Sersic profile is defined as: I(r) = Ie * exp(-k * ((r / r_e)^(1/n) - 1)), where Ie is the intensity at the effective radius r_e, n is the Sersic index that describes the concentration of the source, and k is a parameter that diff --git a/src/caustics/namespace_dict.py b/src/caustics/namespace_dict.py index 6e0ac77d..7106f5df 100644 --- a/src/caustics/namespace_dict.py +++ b/src/caustics/namespace_dict.py @@ -38,8 +38,10 @@ def flatten(self) -> NamespaceDict: """ Flatten the nested dictionary into a NamespaceDict - Returns: - NamespaceDict: Flattened dictionary as a NamespaceDict + Returns + ------- + NamespaceDict + Flattened dictionary as a NamespaceDict """ flattened_dict = NamespaceDict() @@ -59,8 +61,10 @@ def collapse(self) -> NamespaceDict: Flatten the nested dictionary and collapse keys into the first level of the NamespaceDict - Returns: - NamespaceDict: Flattened dictionary as a NamespaceDict + Returns + ------- + NamespaceDict + Flattened dictionary as a NamespaceDict """ flattened_dict = NamespaceDict() diff --git a/src/caustics/parameter.py b/src/caustics/parameter.py index d49a490d..1fcda9ff 100644 --- a/src/caustics/parameter.py +++ b/src/caustics/parameter.py @@ -12,9 +12,12 @@ class Parameter: A static parameter has a fixed value, while a dynamic parameter must be passed in each time it's required. - Attributes: - value (Optional[Tensor]): The value of the parameter. - shape (tuple[int, ...]): The shape of the parameter. + Attributes + ---------- + value: (Optional[Tensor]) + The value of the parameter. + shape: (tuple[int, ...]) + The shape of the parameter. """ def __init__( @@ -75,9 +78,12 @@ def to( """ Moves and/or casts the values of the parameter. - Args: - device (Optional[torch.device], optional): The device to move the values to. Defaults to None. - dtype (Optional[torch.dtype], optional): The desired data type. Defaults to None. + Parameters + ---------- + device: (Optional[torch.device], optional) + The device to move the values to. Defaults to None. + dtype: (Optional[torch.dtype], optional) + The desired data type. Defaults to None. """ if self.static: self.value = self._value.to(device=device, dtype=dtype) diff --git a/src/caustics/parametrized.py b/src/caustics/parametrized.py index 753207e5..ecac2de9 100644 --- a/src/caustics/parametrized.py +++ b/src/caustics/parametrized.py @@ -37,14 +37,22 @@ class Parametrized: - Attributes can be Params, Parametrized, tensor buffers or just normal attributes. - Need to make sure an attribute of one of those types isn't rebound to be of a different type. - Attributes: - name (str): The name of the Parametrized object. Default to class name. - parents (NestedNamespaceDict): Nested dictionary of parent Parametrized objects (higher level, more abstract modules). - params (OrderedDict[str, Parameter]): Dictionary of parameters. - childs NestedNamespaceDict: Nested dictionary of childs Parametrized objects (lower level, more specialized modules). - dynamic_size (int): Size of dynamic parameters. - n_dynamic (int): Number of dynamic parameters. - n_static (int): Number of static parameters. + Attributes + ---------- + name: str + The name of the Parametrized object. Default to class name. + parents: NestedNamespaceDict + Nested dictionary of parent Parametrized objects (higher level, more abstract modules). + params: OrderedDict[str, Parameter] + Dictionary of parameters. + childs: NestedNamespaceDict + Nested dictionary of childs Parametrized objects (lower level, more specialized modules). + dynamic_size: int + Size of dynamic parameters. + n_dynamic: int + Number of dynamic parameters. + n_static: int + Number of static parameters. """ def __init__(self, name: str = None): @@ -155,10 +163,14 @@ def add_param( """ Stores a parameter in the _params dictionary and records its size. - Args: - name (str): The name of the parameter. - value (Optional[Tensor], optional): The value of the parameter. Defaults to None. - shape (Optional[tuple[int, ...]], optional): The shape of the parameter. Defaults to an empty tuple. + Parameters + ---------- + name: str + The name of the parameter. + value: (Optional[Tensor], optional) + The value of the parameter. Defaults to None. + shape: (Optional[tuple[int, ...]], optional) + The shape of the parameter. Defaults to an empty tuple. """ self._params[name] = Parameter(value, shape) # __setattr__ inside add_param to catch all uses of this method @@ -189,18 +201,26 @@ def pack( into arguments to this component and its childs. Also, add a batch dimension to each Tensor without such a dimension. - Args: - x (Union[list[Tensor], dict[str, Union[list[Tensor], Tensor, dict[str, Tensor]]], Tensor): - The input to be packed. Can be a list of tensors, a dictionary of tensors, or a single tensor. - - Returns: - Packed: The packed input, and whether or not the input was batched. - - Raises: - ValueError: If the input is not a list, dictionary, or tensor. - ValueError: If the input is a dictionary and some keys are missing. - ValueError: If the number of dynamic arguments does not match the expected number. - ValueError: If the input is a tensor and the shape does not match the expected shape. + Parameters + ---------- + x: (Union[list[Tensor], dict[str, Union[list[Tensor], Tensor, dict[str, Tensor]]], Tensor) + The input to be packed. Can be a list of tensors, a dictionary of tensors, or a single tensor. + + Returns + ------- + Packed + The packed input, and whether or not the input was batched. + + Raises + ------ + ValueError + If the input is not a list, dictionary, or tensor. + ValueError + If the input is a dictionary and some keys are missing. + ValueError + If the number of dynamic arguments does not match the expected number. + ValueError + If the input is a tensor and the shape does not match the expected shape. """ if isinstance(x, (dict, Packed)): missing_names = [ @@ -211,7 +231,6 @@ def pack( # TODO: check structure! return Packed(x) - elif isinstance(x, (list, tuple)): n_passed = len(x) n_dynamic_params = len(self.params.dynamic.flatten()) @@ -261,18 +280,24 @@ def unpack( Unpacks a dict of kwargs, list of args or flattened vector of args to retrieve this object's static and dynamic parameters. - Args: - x (Optional[dict[str, Union[list[Tensor], dict[str, Tensor], Tensor]]]): - The packed object to be unpacked. + Parameters + ---------- + x: (Optional[dict[str, Union[list[Tensor], dict[str, Tensor], Tensor]]]) + The packed object to be unpacked. - Returns: - list[Tensor]: Unpacked static and dynamic parameters of the object. Note that + Returns + ------- + list[Tensor] + Unpacked static and dynamic parameters of the object. Note that parameters will have an added batch dimension from the pack method. - Raises: - ValueError: If the input is not a dict, list, tuple or tensor. - ValueError: If the argument type is invalid. It must be a dict containing key {self.name} - and value containing args as list or flattened tensor, or kwargs. + Raises + ------ + ValueError + If the input is not a dict, list, tuple or tensor. + ValueError + If the argument type is invalid. It must be a dict containing key {self.name} + and value containing args as list or flattened tensor, or kwargs. """ # Check if module has dynamic parameters if self.module_params.dynamic: @@ -399,12 +424,17 @@ def get_graph( """ Returns a graph representation of the object and its parameters. - Args: - show_dynamic_params (bool, optional): If true, the dynamic parameters are shown in the graph. Defaults to False. - show_static_params (bool, optional): If true, the static parameters are shown in the graph. Defaults to False. - - Returns: - graphviz.Digraph: The graph representation of the object. + Parameters + ---------- + show_dynamic_params: (bool, optional) + If true, the dynamic parameters are shown in the graph. Defaults to False. + show_static_params: (bool, optional) + If true, the static parameters are shown in the graph. Defaults to False. + + Returns + ------- + graphviz.Digraph + The graph representation of the object. """ import graphviz diff --git a/src/caustics/sims/lens_source.py b/src/caustics/sims/lens_source.py index 8e31cc57..80401d7a 100644 --- a/src/caustics/sims/lens_source.py +++ b/src/caustics/sims/lens_source.py @@ -6,7 +6,11 @@ import torch from .simulator import Simulator -from ..utils import get_meshgrid, gaussian_quadrature_grid, gaussian_quadrature_integrator +from ..utils import ( + get_meshgrid, + gaussian_quadrature_grid, + gaussian_quadrature_integrator, +) __all__ = ("Lens_Source",) @@ -36,18 +40,30 @@ class Lens_Source(Simulator): plt.imshow(img, origin = "lower") plt.show() - Attributes: - lens: caustics lens mass model object - source: caustics light object which defines the background source - pixelscale: pixelscale of the sampling grid. - pixels_x: number of pixels on the x-axis for the sampling grid - lens_light (optional): caustics light object which defines the lensing object's light - psf (optional): An image to convolve with the scene. Note that if ``upsample_factor > 1`` the psf must also be at the higher resolution. - pixels_y (optional): number of pixels on the y-axis for the sampling grid. If left as ``None`` then this will simply be equal to ``gridx`` - upsample_factor (default 1): Amount of upsampling to model the image. For example ``upsample_factor = 2`` indicates that the image will be sampled at double the resolution then summed back to the original resolution (given by pixelscale and gridx/y). Note that if you are using a PSF then the PSF must also be at the higher resolution. - psf_pad (default True): If convolving the PSF it is important to sample the model in a larger FOV equal to half the PSF size in order to account for light that scatters from outside the requested FOV inwards. Internally this padding will be added before sampling, then cropped off before returning the final image to the user. - z_s (optional): redshift of the source - name (default "sim"): a name for this simulator in the parameter DAG. + Attributes + ---------- + lens + caustics lens mass model object + source + caustics light object which defines the background source + pixelscale: float + pixelscale of the sampling grid. + pixels_x: int + number of pixels on the x-axis for the sampling grid + lens_light: (optional) + caustics light object which defines the lensing object's light + psf: (optional) + An image to convolve with the scene. Note that if ``upsample_factor > 1`` the psf must also be at the higher resolution. + pixels_y: Optional[int] + number of pixels on the y-axis for the sampling grid. If left as ``None`` then this will simply be equal to ``gridx`` + upsample_factor (default 1) + Amount of upsampling to model the image. For example ``upsample_factor = 2`` indicates that the image will be sampled at double the resolution then summed back to the original resolution (given by pixelscale and gridx/y). + psf_pad: Boolean(default True) + If convolving the PSF it is important to sample the model in a larger FOV equal to half the PSF size in order to account for light that scatters from outside the requested FOV inwards. Internally this padding will be added before sampling, then cropped off before returning the final image to the user. + z_s: optional + redshift of the source + name: string (default "sim") + a name for this simulator in the parameter DAG. Notes: ----- @@ -106,7 +122,9 @@ def __init__( self.gridding[1] + self.psf_pad[1] * 2, ) self.grid = get_meshgrid( - pixelscale/self.upsample_factor, self.n_pix[0]*self.upsample_factor, self.n_pix[1]*self.upsample_factor + pixelscale / self.upsample_factor, + self.n_pix[0] * self.upsample_factor, + self.n_pix[1] * self.upsample_factor, ) if self.psf is not None: @@ -132,11 +150,15 @@ def _unpad_fft(self, x): """ Remove padding from the result of a 2D FFT. - Args: - x (Tensor): The input tensor with padding. + Parameters + --------- + x: Tensor + The input tensor with padding. - Returns: - Tensor: The input tensor without padding. + Returns + ------- + Tensor + The input tensor without padding. """ return torch.roll(x, (-self.psf_pad[0], -self.psf_pad[1]), dims=(-2, -1))[ ..., : self.n_pix[0], : self.n_pix[1] @@ -152,11 +174,20 @@ def forward( quad_level=None, ): """ - params: Packed object - source_light: when true the source light will be sampled - lens_light: when true the lens light will be sampled - lens_source: when true, the source light model will be lensed by the lens mass distribution - psf_convolve: when true the image will be convolved with the psf + forward function + + Parameters + ---------- + params: + Packed object + source_light: boolean + when true the source light will be sampled + lens_light: boolean + when true the lens light will be sampled + lens_source: boolean + when true, the source light model will be lensed by the lens mass distribution + psf_convolve: boolean + when true the image will be convolved with the psf """ (z_s,) = self.unpack(params) @@ -166,37 +197,29 @@ def forward( if self.lens_light is None: lens_light = False if self.psf is None: - psf_convolve = False + psf_convolve = False - if quad_level is not None and quad_level > 1: + if quad_level is not None and quad_level > 1: finegrid_x, finegrid_y, weights = gaussian_quadrature_grid( - self.pixelscale/self.upsample_factor, *self.grid, quad_level - ) + self.pixelscale / self.upsample_factor, *self.grid, quad_level + ) # Sample the source light if source_light: if lens_source: # Source is lensed by the lens mass distribution if quad_level is not None and quad_level > 1: - bx, by = self.lens.raytrace( - finegrid_x, finegrid_y, z_s, params - ) + bx, by = self.lens.raytrace(finegrid_x, finegrid_y, z_s, params) mu_fine = self.source.brightness(bx, by, params) - mu = gaussian_quadrature_integrator( - mu_fine, weights - ) + mu = gaussian_quadrature_integrator(mu_fine, weights) else: bx, by = self.lens.raytrace(*self.grid, z_s, params) mu = self.source.brightness(bx, by, params) else: # Source is imaged without lensing if quad_level is not None and quad_level > 1: - mu_fine = self.source.brightness( - finegrid_x, finegrid_y, params - ) - mu = gaussian_quadrature_integrator( - mu_fine, weights - ) + mu_fine = self.source.brightness(finegrid_x, finegrid_y, params) + mu = gaussian_quadrature_integrator(mu_fine, weights) else: mu = self.source.brightness(*self.grid, params) else: @@ -206,12 +229,8 @@ def forward( # Sample the lens light if lens_light and self.lens_light is not None: if quad_level is not None and quad_level > 1: - mu_fine = self.lens_light.brightness( - finegrid_x, finegrid_y, params - ) - mu += gaussian_quadrature_integrator( - mu_fine, weights - ) + mu_fine = self.lens_light.brightness(finegrid_x, finegrid_y, params) + mu += gaussian_quadrature_integrator(mu_fine, weights) else: mu += self.lens_light.brightness(*self.grid, params) diff --git a/src/caustics/sims/simulator.py b/src/caustics/sims/simulator.py index d24f5260..c005f93e 100644 --- a/src/caustics/sims/simulator.py +++ b/src/caustics/sims/simulator.py @@ -4,7 +4,7 @@ class Simulator(Parametrized): - """A caustic simulator using Parametrized framework. + """A caustics simulator using Parametrized framework. Defines a simulator class which is a callable function that operates on the Parametrized framework. Users define the `forward` diff --git a/src/caustics/tests.py b/src/caustics/tests.py index 62574e86..e86114ad 100644 --- a/src/caustics/tests.py +++ b/src/caustics/tests.py @@ -10,6 +10,7 @@ __all__ = ["test"] + def _test_simulator_runs(): # Model cosmology = FlatLambdaCDM(name="cosmo") @@ -89,6 +90,7 @@ def _test_simulator_runs(): ) ) + def _test_jacobian_autograd_vs_finitediff(): # Models cosmology = FlatLambdaCDM(name="cosmo") @@ -205,20 +207,21 @@ def _test_multiplane_effective_convergence(): curl = lens.effective_convergence_curl(thx, thy, z_s, lens.pack(x)) assert curl.shape == (10, 10) + def test(): """ Run tests for caustics basic functionallity. Run this function to ensure that caustics is working properly. Simply call:: - + >>> import caustics >>> caustics.test() all tests passed! To run the checks. """ - + _test_simulator_runs() _test_jacobian_autograd_vs_finitediff() _test_multiplane_jacobian() _test_multiplane_jacobian_autograd_vs_finitediff() _test_multiplane_effective_convergence() - print("all tests passed!") \ No newline at end of file + print("all tests passed!") diff --git a/src/caustics/utils.py b/src/caustics/utils.py index f010c640..b784bd5e 100644 --- a/src/caustics/utils.py +++ b/src/caustics/utils.py @@ -13,12 +13,17 @@ def flip_axis_ratio(q, phi): """ Makes the value of 'q' positive, then swaps x and y axes if 'q' is larger than 1. - Args: - q (Tensor): Tensor containing values to be processed. - phi (Tensor): Tensor containing the phi values for the orientation of the axes. + Parameters + ---------- + q: Tensor + Tensor containing values to be processed. + phi: Tensor + Tensor containing the phi values for the orientation of the axes. - Returns: - Tuple[Tensor, Tensor]: Tuple containing the processed 'q' and 'phi' Tensors. + Returns + ------- + Tuple[Tensor, Tensor] + Tuple containing the processed 'q' and 'phi' Tensors. """ q = q.abs() return torch.where(q > 1, 1 / q, q), torch.where(q > 1, phi + pi / 2, phi) @@ -28,15 +33,23 @@ def translate_rotate(x, y, x0, y0, phi: Optional[Tensor] = None): """ Translates and rotates the points (x, y) by subtracting (x0, y0) and applying rotation angle phi. - Args: - x (Tensor): Tensor containing the x-coordinates. - y (Tensor): Tensor containing the y-coordinates. - x0 (Tensor): Tensor containing the x-coordinate translation values. - y0 (Tensor): Tensor containing the y-coordinate translation values. - phi (Optional[Tensor], optional): Tensor containing the rotation angles. If None, no rotation is applied. Defaults to None. + Parameters + ---------- + x: Tensor + Tensor containing the x-coordinates. + y: Tensor + Tensor containing the y-coordinates. + x0: Tensor + Tensor containing the x-coordinate translation values. + y0: Tensor + Tensor containing the y-coordinate translation values. + phi: Optional[Tensor], optional) + Tensor containing the rotation angles. If None, no rotation is applied. Defaults to None. - Returns: - Tuple[Tensor, Tensor]: Tuple containing the translated and rotated x and y coordinates. + Returns + ------- + Tuple: [Tensor, Tensor] + Tuple containing the translated and rotated x and y coordinates. """ xt = x - x0 yt = y - y0 @@ -55,13 +68,19 @@ def derotate(vx, vy, phi: Optional[Tensor] = None): """ Applies inverse rotation to the velocity components (vx, vy) using the rotation angle phi. - Args: - vx (Tensor): Tensor containing the x-component of velocity. - vy (Tensor): Tensor containing the y-component of velocity. - phi (Optional[Tensor], optional): Tensor containing the rotation angles. If None, no rotation is applied. Defaults to None. + Parameters + ---------- + vx: Tensor + Tensor containing the x-component of velocity. + vy: Tensor + Tensor containing the y-component of velocity. + phi: Optional[Tensor], optional) + Tensor containing the rotation angles. If None, no rotation is applied. Defaults to None. - Returns: - Tuple[Tensor, Tensor]: Tuple containing the derotated x and y components of velocity. + Returns + ------- + Tuple: [Tensor, Tensor] + Tuple containing the derotated x and y components of velocity. """ if phi is None: return vx, vy @@ -75,13 +94,19 @@ def to_elliptical(x, y, q: Tensor): """ Converts Cartesian coordinates to elliptical coordinates. - Args: - x (Tensor): Tensor containing the x-coordinates. - y (Tensor): Tensor containing the y-coordinates. - q (Tensor): Tensor containing the elliptical parameters. + Parameters + ---------- + x: Tensor + Tensor containing the x-coordinates. + y: Tensor + Tensor containing the y-coordinates. + q: Tensor + Tensor containing the elliptical parameters. - Returns: - Tuple[Tensor, Tensor]: Tuple containing the x and y coordinates in elliptical form. + Returns + ------- + Tuple: Tensor, Tensor + Tuple containing the x and y coordinates in elliptical form. """ return x, y / q @@ -92,15 +117,23 @@ def get_meshgrid( """ Generates a 2D meshgrid based on the provided pixelscale and dimensions. - Args: - pixelscale (float): The scale of the meshgrid in each dimension. - nx (int): The number of grid points along the x-axis. - ny (int): The number of grid points along the y-axis. - device (torch.device, optional): The device on which to create the tensor. Defaults to None. - dtype (torch.dtype, optional): The desired data type of the tensor. Defaults to torch.float32. + Parameters + ---------- + pixelscale: float + The scale of the meshgrid in each dimension. + nx: int + The number of grid points along the x-axis. + ny: int + The number of grid points along the y-axis. + device: torch.device, optional + The device on which to create the tensor. Defaults to None. + dtype: torch.dtype, optional + The desired data type of the tensor. Defaults to torch.float32. - Returns: - Tuple[Tensor, Tensor]: The generated meshgrid as a tuple of Tensors. + Returns + ------- + Tuple: [Tensor, Tensor] + The generated meshgrid as a tuple of Tensors. """ xs = ( torch.linspace(-1, 1, nx, device=device, dtype=dtype) @@ -116,6 +149,7 @@ def get_meshgrid( ) return torch.meshgrid([xs, ys], indexing="xy") + @lru_cache(maxsize=32) def _quad_table(n, p, dtype, device): """ @@ -145,15 +179,16 @@ def _quad_table(n, p, dtype, device): W = torch.outer(w, w) / 4.0 - X, Y = X.reshape(-1), Y.reshape(-1) # flatten + X, Y = X.reshape(-1), Y.reshape(-1) # flatten return X, Y, W.reshape(-1) + def gaussian_quadrature_grid( - pixelscale, - X, - Y, - quad_level = 3, - device=None, + pixelscale, + X, + Y, + quad_level=3, + device=None, dtype=torch.float32, ): """ @@ -190,22 +225,21 @@ def gaussian_quadrature_grid( """ # collect gaussian quadrature weights - abscissaX, abscissaY, weight = _quad_table( - quad_level, pixelscale, dtype, device - ) + abscissaX, abscissaY, weight = _quad_table(quad_level, pixelscale, dtype, device) # Gaussian quadrature evaluation points - Xs = torch.repeat_interleave(X[..., None], quad_level ** 2, -1) + abscissaX - Ys = torch.repeat_interleave(Y[..., None], quad_level ** 2, -1) + abscissaY + Xs = torch.repeat_interleave(X[..., None], quad_level**2, -1) + abscissaX + Ys = torch.repeat_interleave(Y[..., None], quad_level**2, -1) + abscissaY return Xs, Ys, weight + def gaussian_quadrature_integrator( F: Tensor, weight: Tensor, ): """ Performs a pixel-wise integration using Gaussian quadrature. - It takes the brightness function evaluated at the quadrature points `F` and the quadrature weights `weight` as input. + It takes the brightness function evaluated at the quadrature points `F` and the quadrature weights `weight` as input. The result is the integrated brightness function at each pixel. @@ -231,19 +265,20 @@ def gaussian_quadrature_integrator( F = your_brightness_function(Xs, Ys, other, parameters) res = gaussian_quadrature_integrator(F, weight) """ - + return (F * weight).sum(axis=-1) + def quad( - F: Callable, - pixelscale: float, - X: Tensor, + F: Callable, + pixelscale: float, + X: Tensor, Y: Tensor, args: Optional[Tuple] = None, - quad_level: int = 3, - device=None, + quad_level: int = 3, + device=None, dtype=torch.float32, - ): +): """ Performs a pixel-wise integration on a function using Gaussian quadrature. @@ -271,9 +306,7 @@ def quad( Tensor The integrated brightness function at each pixel. """ - X, Y, weight = gaussian_quadrature_grid( - pixelscale, X, Y, quad_level, device, dtype - ) + X, Y, weight = gaussian_quadrature_grid(pixelscale, X, Y, quad_level, device, dtype) F = F(X, Y, *args) return gaussian_quadrature_integrator(F, weight) @@ -282,12 +315,17 @@ def safe_divide(num, denom, places=7): """ Safely divides two tensors, returning zero where the denominator is zero. - Args: - num (Tensor): The numerator tensor. - denom (Tensor): The denominator tensor. + Parameters + ---------- + num: Tensor + The numerator tensor. + denom: Tensor + The denominator tensor. - Returns: - Tensor: The result of the division, with zero where the denominator was zero. + Returns + ------- + Tensor + The result of the division, with zero where the denominator was zero. """ out = torch.zeros_like(num) where = denom != 0 @@ -299,11 +337,15 @@ def safe_log(x): """ Safely applies the logarithm to a tensor, returning zero where the tensor is zero. - Args: - x (Tensor): The input tensor. + Parameters + ---------- + x: Tensor + The input tensor. - Returns: - Tensor: The result of applying the logarithm, with zero where the input was zero. + Returns + ------- + Tensor + The result of applying the logarithm, with zero where the input was zero. """ out = torch.zeros_like(x) where = x != 0 @@ -315,11 +357,15 @@ def _h_poly(t): """Helper function to compute the 'h' polynomial matrix used in the cubic spline. - Args: - t (Tensor): A 1D tensor representing the normalized x values. + Parameters + ---------- + t: Tensor + A 1D tensor representing the normalized x values. - Returns: - Tensor: A 2D tensor of size (4, len(t)) representing the 'h' polynomial matrix. + Returns + ------- + Tensor + A 2D tensor of size (4, len(t)) representing the 'h' polynomial matrix. """ @@ -336,19 +382,26 @@ def interp1d(x: Tensor, y: Tensor, xs: Tensor, extend: str = "extrapolate") -> T """Compute the 1D cubic spline interpolation for the given data points using PyTorch. - Args: - x (Tensor): A 1D tensor representing the x-coordinates of the known data points. - y (Tensor): A 1D tensor representing the y-coordinates of the known data points. - xs (Tensor): A 1D tensor representing the x-coordinates of the positions where - the cubic spline function should be evaluated. - extend (str, optional): The method for handling extrapolation, either "const", "extrapolate", or "linear". - Default is "extrapolate". - "const": Use the value of the last known data point for extrapolation. - "linear": Use linear extrapolation based on the last two known data points. - "extrapolate": Use cubic extrapolation of data. + Parameters + ---------- + x: Tensor + A 1D tensor representing the x-coordinates of the known data points. + y: Tensor + A 1D tensor representing the y-coordinates of the known data points. + xs: Tensor + A 1D tensor representing the x-coordinates of the positions where + the cubic spline function should be evaluated. + extend: (str, optional) + The method for handling extrapolation, either "const", "extrapolate", or "linear". + Default is "extrapolate". + "const": Use the value of the last known data point for extrapolation. + "linear": Use linear extrapolation based on the last two known data points. + "extrapolate": Use cubic extrapolation of data. - Returns: - Tensor: A 1D tensor representing the interpolated values at the specified positions (xs). + Returns + ------- + Tensor + A 1D tensor representing the interpolated values at the specified positions (xs). """ m = (y[1:] - y[:-1]) / (x[1:] - x[:-1]) @@ -381,23 +434,37 @@ def interp2d( Interpolates a 2D image at specified coordinates. Similar to `torch.nn.functional.grid_sample` with `align_corners=False`. - Args: - im (Tensor): A 2D tensor representing the image. - x (Tensor): A 0D or 1D tensor of x coordinates at which to interpolate. - y (Tensor): A 0D or 1D tensor of y coordinates at which to interpolate. - method (str, optional): Interpolation method. Either 'nearest' or 'linear'. Defaults to 'linear'. - padding_mode (str, optional): Defines the padding mode when out-of-bound indices are encountered. - Either 'zeros' or 'extrapolate'. Defaults to 'zeros'. - - Raises: - ValueError: If `im` is not a 2D tensor. - ValueError: If `x` is not a 0D or 1D tensor. - ValueError: If `y` is not a 0D or 1D tensor. - ValueError: If `padding_mode` is not 'extrapolate' or 'zeros'. - ValueError: If `method` is not 'nearest' or 'linear'. - - Returns: - Tensor: Tensor with the same shape as `x` and `y` containing the interpolated values. + Parameters + ---------- + im: Tensor + A 2D tensor representing the image. + x: Tensor + A 0D or 1D tensor of x coordinates at which to interpolate. + y: Tensor + A 0D or 1D tensor of y coordinates at which to interpolate. + method: (str, optional) + Interpolation method. Either 'nearest' or 'linear'. Defaults to 'linear'. + padding_mode: (str, optional) + Defines the padding mode when out-of-bound indices are encountered. + Either 'zeros' or 'extrapolate'. Defaults to 'zeros'. + + Raises + ------ + ValueError + If `im` is not a 2D tensor. + ValueError + If `x` is not a 0D or 1D tensor. + ValueError + If `y` is not a 0D or 1D tensor. + ValueError + If `padding_mode` is not 'extrapolate' or 'zeros'. + ValueError + If `method` is not 'nearest' or 'linear'. + + Returns + ------- + Tensor + Tensor with the same shape as `x` and `y` containing the interpolated values. """ if im.ndim != 2: raise ValueError(f"im must be 2D (received {im.ndim}D tensor)") @@ -458,18 +525,28 @@ def vmap_n( Returns `func` transformed `depth` times by `vmap`, with the same arguments passed to `vmap` each time. - Args: - func (Callable): The function to transform. - depth (int, optional): The number of times to apply `torch.vmap`. Defaults to 1. - in_dims (Union[int, Tuple], optional): The dimensions to vectorize over in the input. Defaults to 0. - out_dims (Union[int, Tuple[int, ...]], optional): The dimensions to vectorize over in the output. Defaults to 0. - randomness (str, optional): How to handle randomness. Defaults to 'error'. - - Raises: - ValueError: If `depth` is less than 1. + Parameters + ---------- + func: Callable + The function to transform. + depth: (int, optional) + The number of times to apply `torch.vmap`. Defaults to 1. + in_dims: (Union[int, Tuple], optional) + The dimensions to vectorize over in the input. Defaults to 0. + out_dims: (Union[int, Tuple[int, ...]], optional): + The dimensions to vectorize over in the output. Defaults to 0. + randomness: (str, optional) + How to handle randomness. Defaults to 'error'. + + Raises + ------ + ValueError + If `depth` is less than 1. - Returns: - Callable: The transformed function. + Returns + ------- + Callable + The transformed function. TODO: test. """ @@ -487,12 +564,17 @@ def get_cluster_means(xs: Tensor, k: int): """ Computes cluster means using the k-means++ initialization algorithm. - Args: - xs (Tensor): A tensor of data points. - k (int): The number of clusters. + Parameters + ---------- + xs: Tensor + A tensor of data points. + k: int + The number of clusters. - Returns: - Tensor: A tensor of cluster means. + Returns + ------- + Tensor + A tensor of cluster means. """ b = len(xs) mean_idxs = [int(torch.randint(high=b, size=(), device=xs.device).item())] diff --git a/tests/test_parametrized.py b/tests/test_parametrized.py index 9deed57e..660321c3 100644 --- a/tests/test_parametrized.py +++ b/tests/test_parametrized.py @@ -123,12 +123,12 @@ def test_parametrized_name_setter_bad_names(): # Make sure bad names are catched by our added method. Bad names are name which cannot be used as class attributes. good_names = ["variable", "_variable", "var_iable2"] for name in good_names: - module = Sersic(name=name) + Sersic(name=name) bad_names = ["for", "2variable", "variable!", "var-iable", "var iable", "def"] for name in bad_names: print(name) with pytest.raises(NameError): - module = Sersic(name=name) + Sersic(name=name) def test_parametrized_name_collision(): diff --git a/tests/test_sersic.py b/tests/test_sersic.py index e4c5e9a8..96e1405a 100644 --- a/tests/test_sersic.py +++ b/tests/test_sersic.py @@ -9,7 +9,7 @@ def test(): - # Caustic setup + # Caustics setup res = 0.05 nx = 200 ny = 200 diff --git a/tests/test_simulator.py b/tests/test_simulator.py index 71d06ed1..50968514 100644 --- a/tests/test_simulator.py +++ b/tests/test_simulator.py @@ -89,7 +89,5 @@ def test_simulator_runs(): ) # Check quadrature integration is accurate - assert torch.allclose(sim(), sim(quad_level=3), rtol = 1e-1) - assert torch.allclose(sim(quad_level=3), sim(quad_level=5), rtol = 1e-2) - - + assert torch.allclose(sim(), sim(quad_level=3), rtol=1e-1) + assert torch.allclose(sim(quad_level=3), sim(quad_level=5), rtol=1e-2) diff --git a/tests/utils.py b/tests/utils.py index 6ab48e0b..0dd96ca7 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -157,7 +157,7 @@ def get_default_cosmologies(): def setup_grids(res=0.05, n_pix=100): - # Caustic setup + # Caustics setup thx, thy = get_meshgrid(res, n_pix, n_pix) # Lenstronomy setup