doc/source/visualizing/FITSImageData.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Writing FITS Images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "yt has capabilities for writing 2D and 3D uniformly gridded data generated from datasets to FITS files. This is via the `FITSImageData` class. We'll test these capabilities out on an Athena dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "units_override = {\n",
    "    \"length_unit\": (1.0, \"Mpc\"),\n",
    "    \"mass_unit\": (1.0e14, \"Msun\"),\n",
    "    \"time_unit\": (1.0, \"Myr\"),\n",
    "}\n",
    "ds = yt.load(\"MHDSloshing/virgo_low_res.0054.vtk\", units_override=units_override)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating FITS images from Slices and Projections"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are several ways to make a `FITSImageData` instance. The most intuitive ways are to use the `FITSSlice`, `FITSProjection`, `FITSOffAxisSlice`, and `FITSOffAxisProjection` classes to write slices and projections directly to FITS. To demonstrate a useful example of creating a FITS file, let's first make a `ProjectionPlot`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj = yt.ProjectionPlot(\n",
    "    ds,\n",
    "    \"z\",\n",
    "    (\"gas\", \"temperature\"),\n",
    "    weight_field=(\"gas\", \"density\"),\n",
    "    width=(500.0, \"kpc\"),\n",
    ")\n",
    "prj.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose that we wanted to write this projection to a FITS file for analysis and visualization in other programs, such as ds9. We can do that using `FITSProjection`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits = yt.FITSProjection(\n",
    "    ds, \"z\", (\"gas\", \"temperature\"), weight_field=(\"gas\", \"density\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which took the same parameters as `ProjectionPlot` except the width, because `FITSProjection` and `FITSSlice` always make slices and projections of the width of the domain size, at the finest resolution available in the simulation, in a unit determined to be appropriate for the physical size of the dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also set width manually in `FITSProjection`. For example, set the width to 500 kiloparsec to get FITS file of the same projection plot as discussed above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits = yt.FITSProjection(\n",
    "    ds,\n",
    "    \"z\",\n",
    "    (\"gas\", \"temperature\"),\n",
    "    weight_field=(\"gas\", \"density\"),\n",
    "    width=(500.0, \"kpc\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want the center coordinates of the image in either a slice or a projection to be (0,0) instead of the domain coordinates, set `origin=\"image\"`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits_img = yt.FITSProjection(\n",
    "    ds,\n",
    "    \"z\",\n",
    "    (\"gas\", \"temperature\"),\n",
    "    weight_field=(\"gas\", \"density\"),\n",
    "    width=(500.0, \"kpc\"),\n",
    "    origin=\"image\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Making FITS images from Particle Projections"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "To create a FITS image from a particle field which is smeared onto the image, we can use\n",
    "`FITSParticleProjection`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "dsp = yt.load(\"gizmo_64/output/snap_N64L16_135.hdf5\")\n",
    "prjp_fits = yt.FITSParticleProjection(\n",
    "    dsp, \"x\", (\"PartType1\", \"particle_mass\"), deposition=\"cic\"\n",
    ")\n",
    "prjp_fits.writeto(\"prjp.fits\", overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Note that we used the \"Cloud-In-Cell\" interpolation method (`\"cic\"`) instead of the default\n",
    "\"Nearest-Grid-Point\" (`\"ngp\"`) method. \n",
    "\n",
    "If you want the projection to be divided by the pixel area (to make a projection of mass density, \n",
    "for example), supply the ``density`` keyword argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "prjpd_fits = yt.FITSParticleProjection(\n",
    "    dsp, \"x\", (\"PartType1\", \"particle_mass\"), density=True, deposition=\"cic\"\n",
    ")\n",
    "prjpd_fits.writeto(\"prjpd.fits\", overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "`FITSParticleOffAxisProjection` can be used to make a projection along any arbitrary sight line:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "L = [1, -1, 1]  # normal or \"line of sight\" vector\n",
    "N = [0, 0, 1]  # north or \"up\" vector\n",
    "poff_fits = yt.FITSParticleOffAxisProjection(\n",
    "    dsp, L, (\"PartType1\", \"particle_mass\"), deposition=\"cic\", north_vector=N\n",
    ")\n",
    "poff_fits.writeto(\"poff.fits\", overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Using `HDUList` Methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can call a number of the [AstroPy `HDUList`](https://astropy.readthedocs.io/en/latest/io/fits/api/hdulists.html) class's methods from a `FITSImageData` object. For example, `info` shows us the contents of the virtual FITS file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also look at the header for a particular field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits[\"temperature\"].header"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where we can see that the units of the temperature field are Kelvin and the cell widths are in kiloparsecs. Note that the length, time, mass, velocity, and magnetic field units of the dataset have been copied into the header "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " If we want the raw image data with units, we can use the `data` attribute of this field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits[\"temperature\"].data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Changing Aspects of the Images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the `set_unit` method to change the units of a particular field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits.set_unit(\"temperature\", \"R\")\n",
    "prj_fits[\"temperature\"].data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The length units of the image (and its coordinate system), as well as the resolution of the image, can be adjusted when creating it using the `length_unit` and `image_res` keyword arguments, respectively:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# length_unit defaults to that from the dataset\n",
    "# image_res defaults to 512\n",
    "slc_fits = yt.FITSSlice(\n",
    "    ds, \"z\", (\"gas\", \"density\"), width=(500, \"kpc\"), length_unit=\"ly\", image_res=256\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now check that this worked by looking at the header, notice in particular the `NAXIS[12]` and `CUNIT[12]` keywords (the `CDELT[12]` and `CRPIX[12]` values also change):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slc_fits[\"density\"].header"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Saving and Loading Images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The image can be written to disk using the `writeto` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits.writeto(\"sloshing.fits\", overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since yt can read FITS image files, it can be loaded up just like any other dataset. Since we created this FITS file with `FITSImageData`, the image will contain information about the units and the current time of the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds2 = yt.load(\"sloshing.fits\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and we can make a `SlicePlot` of the 2D image, which shows the same data as the previous image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slc2 = yt.SlicePlot(ds2, \"z\", (\"gas\", \"temperature\"), width=(500.0, \"kpc\"))\n",
    "slc2.set_log(\"temperature\", True)\n",
    "slc2.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating `FITSImageData` Instances Directly from FRBs, PlotWindow instances, and 3D Grids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want more fine-grained control over what goes into the FITS file, you can call `FITSImageData` directly, with various kinds of inputs. For example, you could use a `FixedResolutionBuffer`, and specify you want the units in parsecs instead:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slc3 = ds.slice(0, 0.0)\n",
    "frb = slc3.to_frb((500.0, \"kpc\"), 800)\n",
    "fid_frb = frb.to_fits_data(\n",
    "    fields=[(\"gas\", \"density\"), (\"gas\", \"temperature\")], length_unit=\"pc\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If one creates a `PlotWindow` instance, e.g. `SlicePlot`, `ProjectionPlot`, etc., you can also call this same method there:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fid_pw = prj.to_fits_data(\n",
    "    fields=[(\"gas\", \"density\"), (\"gas\", \"temperature\")], length_unit=\"pc\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A 3D FITS cube can also be created from regularly gridded 3D data. In yt, there are covering grids and \"arbitrary grids\". The easiest way to make an arbitrary grid object is using `ds.r`, where we can index the dataset like a NumPy array, creating a grid of 1.0 Mpc on a side, centered on the origin, with 64 cells on a side:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = ds.r[\n",
    "    (-0.5, \"Mpc\"):(0.5, \"Mpc\"):64j,\n",
    "    (-0.5, \"Mpc\"):(0.5, \"Mpc\"):64j,\n",
    "    (-0.5, \"Mpc\"):(0.5, \"Mpc\"):64j,\n",
    "]\n",
    "fid_grid = grid.to_fits_data(\n",
    "    fields=[(\"gas\", \"density\"), (\"gas\", \"temperature\")], length_unit=\"Mpc\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other `FITSImageData` Methods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating Images from Others"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `FITSImageData` instance can be generated from one previously written to disk using the `from_file` classmethod:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fid = yt.FITSImageData.from_file(\"sloshing.fits\")\n",
    "fid.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Multiple `FITSImageData` can be combined to create a new one, provided that the coordinate information is the same:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits2 = yt.FITSProjection(ds, \"z\", (\"gas\", \"density\"), width=(500.0, \"kpc\"))\n",
    "prj_fits3 = yt.FITSImageData.from_images([prj_fits, prj_fits2])\n",
    "prj_fits3.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, individual fields can be popped as well to produce new instances of `FITSImageData`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dens_fits = prj_fits3.pop(\"density\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So this new instance would only have the `\"density\"` field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dens_fits.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and the old one has the `\"density\"` field removed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits3.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding Sky Coordinates to Images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far, the FITS images we have shown have linear spatial coordinates. We can see this by looking at the header for one of the fields, and examining the `CTYPE1` and `CTYPE2` keywords:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits[\"temperature\"].header"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `WCSNAME` keyword is set to `\"yt\"` by default. \n",
    "\n",
    "However, one may want to take a projection of an object and make a crude mock observation out of it, with celestial coordinates. For this, we can use the `create_sky_wcs` method. Specify a center (RA, Dec) coordinate in degrees, as well as a linear scale in terms of angle per distance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sky_center = [30.0, 45.0]  # in degrees\n",
    "sky_scale = (2.5, \"arcsec/kpc\")  # could also use a YTQuantity\n",
    "prj_fits.create_sky_wcs(sky_center, sky_scale, ctype=[\"RA---TAN\", \"DEC--TAN\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, a tangent RA/Dec projection is used, but one could also use another projection using the `ctype` keyword. We can now look at the header and see it has the appropriate WCS now. The old `\"yt\"` WCS has been added to a second WCS in the header, where the parameters have an `\"A\"` appended to them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits[\"temperature\"].header"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and now the `WCSNAME` has been set to `\"celestial\"`. If you want the original WCS to remain in the original place, then you can make the call to `create_sky_wcs` and set `replace_old_wcs=False`, which will put the new, celestial WCS in the second one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "prj_fits3.create_sky_wcs(\n",
    "    sky_center, sky_scale, ctype=[\"RA---TAN\", \"DEC--TAN\"], replace_old_wcs=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prj_fits3[\"temperature\"].header"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Updating Header Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also add header keywords to a single field or for all fields in the FITS image using `update_header`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fid_frb.update_header(\"all\", \"time\", 0.1)  # Update all the fields\n",
    "fid_frb.update_header(\"temperature\", \"scale\", \"Rankine\")  # Update just one field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(fid_frb[\"density\"].header[\"time\"])\n",
    "print(fid_frb[\"temperature\"].header[\"scale\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Changing Image Names"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use the `change_image_name` method to change the name of an image in a `FITSImageData` instance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fid_frb.change_image_name(\"density\", \"mass_per_volume\")\n",
    "fid_frb.info()  # now \"density\" should be gone and \"mass_per_volume\" should be in its place"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convolving FITS Images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, you can convolve an image inside a `FITSImageData` instance with a kernel, either a Gaussian with a specific standard deviation, or any kernel provided by AstroPy. See AstroPy's [Convolution and filtering](http://docs.astropy.org/en/stable/convolution/index.html) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dens_fits.writeto(\"not_convolved.fits\", overwrite=True)\n",
    "# Gaussian kernel with standard deviation of 3.0 kpc\n",
    "dens_fits.convolve(\"density\", 3.0)\n",
    "dens_fits.writeto(\"convolved.fits\", overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's load these up as datasets and see the difference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds0 = yt.load(\"not_convolved.fits\")\n",
    "dsc = yt.load(\"convolved.fits\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slc3 = yt.SlicePlot(ds0, \"z\", (\"gas\", \"density\"), width=(500.0, \"kpc\"))\n",
    "slc3.set_log(\"density\", True)\n",
    "slc3.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slc4 = yt.SlicePlot(dsc, \"z\", (\"gas\", \"density\"), width=(500.0, \"kpc\"))\n",
    "slc4.set_log(\"density\", True)\n",
    "slc4.show()"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}