{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Tutorial 5: Ray Tracing\n",
        "=======================\n",
        "\n",
        "In the previous tutorial, we used planes to perform strong lens calculations. The use of planes was a bit cumbersome:\n",
        "we had to set up each galaxy, pass them to each plane and manually use each plane to perform every ray-tracing\n",
        "calculation ourselves. It was easy to make a mistake!\n",
        "\n",
        "However, remember how the `Galaxy` objects contains its `redshift` as an attribute? Given a list of galaxies, there\n",
        "should be no need for us to manually specify each plane setup and manually perform each ray-tracing calculations\n",
        "ourself. All of the information required to do this is contained in the galaxies!\n",
        "\n",
        "In this tutorial, we introduce potentially the most important object, the `Tracer`. This exploits the\n",
        "redshift information of galaxies to automatically perform ray-tracing calculations."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "%matplotlib inline\n",
        "from pyprojroot import here\n",
        "workspace_path = str(here())\n",
        "%cd $workspace_path\n",
        "print(f\"Working Directory has been set to `{workspace_path}`\")\n",
        "\n",
        "import autolens as al\n",
        "import autolens.plot as aplt"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "__Initial Setup__\n",
        "\n",
        "Let use the same grid we've grown love by now!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "image_plane_grid = al.Grid2D.uniform(shape_native=(100, 100), pixel_scales=0.05)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "For our lens galaxy, we'll use the same spherical isothermal mass profile again."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "sis_mass_profile = al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6)\n",
        "\n",
        "lens_galaxy = al.Galaxy(redshift=0.5, mass=sis_mass_profile)\n",
        "\n",
        "print(lens_galaxy)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "And for the source, the same spherical Sersic light profile."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "sersic_light_profile = al.lp.SersicSph(\n",
        "    centre=(0.0, 0.0), intensity=1.0, effective_radius=1.0, sersic_index=1.0\n",
        ")\n",
        "\n",
        "source_galaxy = al.Galaxy(redshift=1.0, light=sersic_light_profile)\n",
        "\n",
        "print(source_galaxy)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "__Tracers__\n",
        "\n",
        "Now, lets use the lens and source galaxies to ray-trace the grid, however this time using a `Tracer` object.  \n",
        "When we pass our galaxies into the `Tracer` below, the following happens:\n",
        "\n",
        "1) The galaxies are ordered in ascending redshift.\n",
        "2) Planes are created at every one of these redshifts, with each galaxy associated with the plane at its redshift."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "tracer = al.Tracer.from_galaxies(galaxies=[lens_galaxy, source_galaxy])"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The has the list of planes as an attribute, which in this example is two planes (an image and source plane)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "print(tracer.planes)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can access these using the `image_plane` and `source_plane` attributes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "print(\"Image Plane:\")\n",
        "print(tracer.planes[0])\n",
        "print(tracer.image_plane)\n",
        "print()\n",
        "print(\"Source Plane:\")\n",
        "print(tracer.planes[1])\n",
        "print(tracer.source_plane)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "__Ray Tracing__\n",
        "\n",
        "With a `Tracer` we create fully ray-traced images without manually setting up the planes and performing lensing \n",
        "calculations one-by-one to do this. The function below does the following:\n",
        "\n",
        " 1) Using the lens's total mass distribution, the deflection angle of every image-plane $(y,x)$ grid coordinate is \n",
        " computed.\n",
        " 2) These deflection angles are used to trace every image-plane coordinate to the source-plane.\n",
        " 3) The light of each traced source-plane coordinate is evaluated using the source plane galaxy's light profiles."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "traced_image_2d = tracer.image_2d_from(grid=image_plane_grid)\n",
        "print(\"traced image pixel 1\")\n",
        "print(traced_image_2d.native[0, 0])\n",
        "print(\"traced image pixel 2\")\n",
        "print(traced_image_2d.native[0, 1])\n",
        "print(\"traced image pixel 3\")\n",
        "print(traced_image_2d.native[0, 2])"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This image appears as the Einstein ring we saw in the previous tutorial."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "tracer_plotter = aplt.TracerPlotter(tracer=tracer, grid=image_plane_grid)\n",
        "tracer_plotter.figures_2d(image=True)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can also use the tracer to compute the traced grid of every plane:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "traced_grid_list = tracer.traced_grid_2d_list_from(grid=image_plane_grid)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The first traced grid corresponds to the image-plane grid (i.e. before lensing), whereas the second grid is that of\n",
        "the source-plane."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "print(\"grid image-plane (y,x) coordinate 1\")\n",
        "print(traced_grid_list[0].native[0, 0])\n",
        "print(\"grid image-plane (y,x) coordinate 2\")\n",
        "print(traced_grid_list[0].native[0, 1])\n",
        "print(\"grid source-plane (y,x) coordinate 1\")\n",
        "print(traced_grid_list[1].native[0, 0])\n",
        "print(\"grid source-plane (y,x) coordinate 2\")\n",
        "print(traced_grid_list[1].native[0, 1])"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can use the `TracerPlotter` to plot these planes and grids."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "include = aplt.Include2D(grid=True)\n",
        "\n",
        "tracer_plotter = aplt.TracerPlotter(\n",
        "    tracer=tracer, grid=image_plane_grid, include_2d=include\n",
        ")\n",
        "tracer_plotter.figures_2d_of_planes(plane_image=True, plane_grid=True, plane_index=0)\n",
        "tracer_plotter.figures_2d_of_planes(plane_image=True, plane_grid=True, plane_index=1)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "A ray-tracing subplot plots the following:\n",
        "\n",
        " 1) The image, computed by ray-tracing the source-galaxy's light from the source-plane to the image-plane.\n",
        " 2) The source-plane image, showing the source-galaxy's intrinsic appearance (i.e. if it were not lensed).\n",
        " 3) The image-plane convergence, computed using the lens galaxy's total mass distribution.\n",
        " 4) The image-plane gravitational potential, computed using the lens galaxy's total mass distribution.\n",
        " 5) The image-plane deflection angles, computed using the lens galaxy's total mass distribution."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "tracer_plotter.subplot_tracer()"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Just like for a profile, galaxy or plane, these quantities attributes can be computed via a `_from_grid` method."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "convergence_2d = tracer.convergence_2d_from(grid=image_plane_grid)\n",
        "\n",
        "print(\"Tracer convergence at coordinate 1:\")\n",
        "print(convergence_2d.native[0, 0])\n",
        "print(\"Tracer convergence at coordinate 2:\")\n",
        "print(convergence_2d.native[0, 1])\n",
        "print(\"Tracer convergence at coordinate 101:\")\n",
        "print(convergence_2d.native[1, 0])"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Of course, these convergences are identical to the image-plane convergences, as it`s only the lens galaxy that \n",
        "contributes to the overall mass of the ray-tracing system."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "image_plane_convergence_2d = tracer.image_plane.convergence_2d_from(\n",
        "    grid=image_plane_grid\n",
        ")\n",
        "\n",
        "print(\"Image-Plane convergence at coordinate 1:\")\n",
        "print(image_plane_convergence_2d.native[0, 0])\n",
        "print(\"Image-Plane convergence at coordinate 2:\")\n",
        "print(image_plane_convergence_2d.native[0, 1])\n",
        "print(\"Image-Plane convergene at coordinate 101:\")\n",
        "print(image_plane_convergence_2d.native[1, 0])"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "I've left the rest below commented to avoid too many print statements, but if you're feeling adventurous go ahead \n",
        "and uncomment the lines below!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "# print(\"Potential:\")\n",
        "# print(tracer.potential_2d_from(grid=image_plane_grid))\n",
        "# print(tracer.image_plane.potential_2d_from(grid=image_plane_grid))\n",
        "# print(\"Deflections:\")\n",
        "# print(tracer.deflections_yx_2d_from(grid=image_plane_grid))\n",
        "# print(tracer.image_plane.deflections_yx_2d_from(grid=image_plane_grid))"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The `TracerPlotter` can also plot the above attributes as individual figures:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "tracer_plotter = aplt.TracerPlotter(tracer=tracer, grid=image_plane_grid)\n",
        "tracer_plotter.figures_2d(\n",
        "    image=True,\n",
        "    convergence=True,\n",
        "    potential=False,\n",
        "    deflections_y=False,\n",
        "    deflections_x=False,\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "__Caustics__\n",
        "\n",
        "In the previous tutorial, we plotted the critical curves of the mass profile on the image-plane. We will now plot the\n",
        "'caustics', which correspond to each critical curve ray-traced to the source-plane. This is computed by using the \n",
        "lens galaxy mass profile's to calculate the deflection angles at the critical curves and ray-trace them to the \n",
        "source-plane.\n",
        "\n",
        "As discussed in the previous tutorial, critical curves mark regions of infinite magnification. Thus, if a source\n",
        "appears near a caustic in the source plane it will appear significantly brighter than its true luminosity. \n",
        "\n",
        "We again have to use a mass profile with a slope below 2.0 to ensure a radial critical curve and therefore radial\n",
        "caustic is formed.\n",
        "\n",
        "We can plot both the tangential and radial critical curves and caustics using an `Include2D` object. Note how the \n",
        "critical curves appear only for the image-plane grid, whereas the caustic only appears in the source plane."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "sis_mass_profile = al.mp.IsothermalSph(centre=(0.0, 0.0), einstein_radius=1.6)\n",
        "\n",
        "lens_galaxy = al.Galaxy(redshift=0.5, mass=sis_mass_profile)\n",
        "\n",
        "tracer = al.Tracer.from_galaxies(galaxies=[lens_galaxy, source_galaxy])\n",
        "\n",
        "include = aplt.Include2D(\n",
        "    tangential_critical_curves=True,\n",
        "    tangential_caustics=True,\n",
        "    radial_critical_curves=True,\n",
        "    radial_caustics=True,\n",
        ")\n",
        "tracer_plotter = aplt.TracerPlotter(tracer=tracer, grid=image_plane_grid)\n",
        "\n",
        "tracer_plotter.figures_2d_of_planes(plane_grid=True, plane_index=0)\n",
        "tracer_plotter.figures_2d_of_planes(plane_grid=True, plane_index=1)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can also plot the caustic on the source-plane image."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "tracer_plotter.figures_2d_of_planes(plane_index=1, plane_image=True)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Caustics also mark the regions in the source-plane where the multiplicity of the strong lens changes. That is,\n",
        "if a source crosses a caustic, it goes from 2 images to 1 image. Try and show this yourself by changing the (y,x) \n",
        "centre of the source-plane galaxy's light profile!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [
        "sersic_light_profile = al.lp.SersicSph(\n",
        "    centre=(0.0, 0.0), intensity=1.0, effective_radius=1.0, sersic_index=1.0\n",
        ")\n",
        "source_galaxy = al.Galaxy(redshift=1.0, light=sersic_light_profile)\n",
        "tracer = al.Tracer.from_galaxies(galaxies=[lens_galaxy, source_galaxy])\n",
        "\n",
        "tracer_plotter = aplt.TracerPlotter(tracer=tracer, grid=image_plane_grid)\n",
        "tracer_plotter.figures_2d_of_planes(plane_index=1, plane_image=True)"
      ],
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "__Wrap Up__\n",
        "\n",
        "You might be wondering why do both the tracer and its image-plane have the attributes convergence / potential / \n",
        "deflection angles, when the two are identical? Afterall, only mass profiles contribute to these quantities, and \n",
        "only image-plane galaxies have had mass profiles so far!\n",
        "\n",
        "The reason is due 'multi-plane' lensing, which is basically any lensing configuration where there are more than 2 \n",
        "galaxies at more than 2 redshifts. For example, we could have a lensing system with 3 galaxies, at redshifts 0.5, 1.0 \n",
        "and 2.0 (you can make such a system in the `Tracer` if you wish!).\n",
        " \n",
        "When there are more than 2 planes, the convergence, potential and deflections at each individual plane is different to \n",
        "that of the entire multi-plane strong lens system. This is beyond the scope of this chapter, but be reassured that \n",
        "what you're learning now will prepare you for the advanced chapters later on!\n",
        "\n",
        "And with that, we're done. You`ve performed your first ray-tracing! There are no exercises for this \n",
        "chapter, and we're going to take a deeper look at ray-tracing in the next chapter."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": [],
      "outputs": [],
      "execution_count": null
    }
  ],
  "metadata": {
    "anaconda-cloud": {},
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.1"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}