timelags_from_diffrot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install ndcube ipympl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!conda install -y ffmpeg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!jupyter labextension install @jupyter-widgets/jupyterlab-manager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask.bytes\n",
    "import dask.array as da\n",
    "import dask.dataframe as df\n",
    "import dask\n",
    "from dask import delayed, compute\n",
    "from astropy.io import fits\n",
    "from astropy.time import Time\n",
    "import astropy.units as u\n",
    "from functools import partial\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as anim\n",
    "import sunpy.map\n",
    "from sunpy.util.metadata import MetaDict\n",
    "from astropy.coordinates import SkyCoord\n",
    "from sunpy.physics.differential_rotation import solar_rotate_coordinate\n",
    "from sunpy.physics.differential_rotation import diffrot_map\n",
    "from header_helpers import validate_dtype_shape\n",
    "from sunpy.instr.aia import aiaprep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gcsfs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "gcs = gcsfs.GCSFileSystem(token='cache')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "wavs = [94, 131, 171, 193, 211, 335]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{94: 692, 131: 738, 171: 743, 193: 744, 211: 737, 335: 753}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gcs.invalidate_cache()\n",
    "{wav: len(gcs.ls(f'pangeo-data/SDO_AIA_Images/diffrot/{wav:03d}/')) for wav in wavs}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "cluster.close()\n",
    "c.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask_kubernetes import KubeCluster\n",
    "cluster = KubeCluster(n_workers=200)\n",
    "cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client, progress\n",
    "c = Client(cluster)\n",
    "c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "read_template = \"gcs://pangeo-data/SDO_AIA_Images/diffrot/{wav:03d}/aia_diffrot_{wav:02d}a_*.fits\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "read_templates = {wav: read_template.format(wav=wav) for wav in wavs}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "filebytes = {wav: dask.bytes.open_files(temp, token='anon') for wav, temp in read_templates.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_header(fn, hdu=0):\n",
    "    with fn as fi:\n",
    "        return MetaDict(sunpy.io.fits.get_header(fi)[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "middle_n = len(filenames[171])//2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "futures = {wav: c.map(get_header, fnames) for wav, fnames in filenames.items()}\n",
    "headers = {wav: c.gather(fut) for wav, fut in futures.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "middle_t = Time(headers[171][middle_n]['DATE-OBS'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtypes = {wav: validate_dtype_shape(headers) for wav, headers in headers.items()}\n",
    "dtypes\n",
    "assert all([d == list(dtypes.values())[0] for d in dtypes.values()])\n",
    "dtype, shape = list(dtypes.values())[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DelayedFITS:\n",
    "    def __init__(self, file, shape, dtype, hdu=0):\n",
    "        self.shape = shape\n",
    "        self.dtype = dtype\n",
    "        self.file = file\n",
    "        self.hdu = hdu\n",
    "    \n",
    "    def __getitem__(self, item):\n",
    "        with self.file as fi:\n",
    "            with fits.open(fi) as hdul:\n",
    "                return hdul[self.hdu].section[item]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arrs = {wav: [da.from_array(DelayedFITS(fn, shape=shape, dtype=dtype, hdu=0), chunks=shape)\n",
    "         for fn in files] for wav, files in filenames.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "maps = {}\n",
    "for wav in arrs.keys():\n",
    "    maps[wav] = [delayed(sunpy.map.sources.AIAMap)(arr, header) for arr, header in zip(arrs[wav], headers[wav])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rarrs = {wav: [da.from_delayed(m.data, dtype=np.float64, shape=shape) for m in maps] for wav, maps in derot.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stacked = {wav: da.stack(arrs) for wav, arrs in rarrs.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stacked = {wav: d.rechunk((1800, 256, 256)) for wav, d in stacked.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times = {wav: (Time([head['DATE-OBS'] for head in header]) - Time(header[0]['DATE-OBS'])).to(u.s) for wav, header in headers.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timelags = {wav: u.Quantity(np.hstack((-1*time[::-1], time[1:])).value, time.unit) for wav, time in times.items()}\n",
    "timelags = timelags[171]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def correlation_2d(stacked, timelags, channel_a, channel_b):\n",
    "    \"\"\"\n",
    "    Create Dask task graph to compute cross-correlation using FFT for each pixel in an AIA map\n",
    "    \"\"\"\n",
    "    darray_a = stacked[channel_a][::-1]\n",
    "    darray_b = stacked[channel_b]\n",
    "    # Normalize\n",
    "    std_a = darray_a.std(axis=0)\n",
    "    std_a = dask.array.where(std_a == 0, 1, std_a)\n",
    "    v_a = (darray_a - darray_a.mean(axis=0)[np.newaxis]) / std_a[np.newaxis]\n",
    "    std_b = darray_b.std(axis=0)\n",
    "    std_b = dask.array.where(std_b == 0, 1, std_b)\n",
    "    v_b = (darray_b - darray_b.mean(axis=0)[np.newaxis]) / std_b[np.newaxis]\n",
    "    # Fast Fourier Transform of both channels\n",
    "    fft_a = dask.array.fft.rfft(v_a, axis=0, n=timelags.shape[0])\n",
    "    fft_b = dask.array.fft.rfft(v_b, axis=0, n=timelags.shape[0])\n",
    "    # Inverse of product of FFTS to get cross-correlation (by convolution theorem)\n",
    "    return dask.array.fft.irfft(fft_a * fft_b, axis=0, n=timelags.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xcorrelation = correlation_2d(stacked, timelags, 335, 171)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = timelags[xcorrelation.argmax(axis=0)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outmap = sunpy.map.Map(a, headers[171][90])\n",
    "outmap.plot_settings = {'cmap': 'RdBu',\n",
    "                        'norm': plt.Normalize(),\n",
    "                        'interpolation': 'nearest',\n",
    "                        'origin': 'lower'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8,8))\n",
    "outmap.plot(vmin=-1e4, vmax=1e4)\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda root]",
   "language": "python",
   "name": "conda-root-py"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}