lsst_scan_strategy_video.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import PIL.Image\n",
    "import healpy\n",
    "import numpy as np\n",
    "import os\n",
    "import skvideo.io\n",
    "import astropy.table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# img1 is made from a 2mass image\n",
    "# img2 is made from a healpix image generated from Flask\n",
    "img1=PIL.Image.open(\"img1.png\")\n",
    "img2=PIL.Image.open(\"img2.png\")\n",
    "arr1=np.array(img1)\n",
    "arr2=np.array(img2)\n",
    "T = arr1[:,:,0]\n",
    "D = arr2[:,:,0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Crop the images\n",
    "ins = (T!=255) | (D!=255)\n",
    "\n",
    "ymin=np.where(ins)[0].min()\n",
    "ymax=np.where(ins)[0].max()\n",
    "xmin=np.where(ins)[1].min()\n",
    "xmax=np.where(ins)[1].max()\n",
    "\n",
    "# Need to be even for some video codecs\n",
    "if (ymax-ymin)%2==1:\n",
    "    ymax +=1 \n",
    "\n",
    "if (xmax-xmin)%2==1:\n",
    "    xmax += 1\n",
    "\n",
    "arr1=np.array(img1)[ymin:ymax,xmin:xmax]\n",
    "arr2=np.array(img2)[ymin:ymax,xmin:xmax]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make the combined image - white stars + purple dark matter\n",
    "T = arr1[:,:,0]\n",
    "D = arr2[:,:,0]\n",
    "\n",
    "out = (T==255) & (D==255) \n",
    "ins = (T!=255) | (D!=255)\n",
    "full_image=np.zeros(arr1.shape, dtype=float)\n",
    "full_image[:,:,0] = 255-T + 0.5*D\n",
    "full_image[:,:,1] = 255-T\n",
    "full_image[:,:,2] = 255-T + 0.5*D\n",
    "full_image[out,:] = 0\n",
    "full_image[:,:,3] = 255\n",
    "full_image=full_image.clip(0,255).astype('uint8')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get coordinates\n",
    "exps = astropy.table.Table.read('exps.fits')\n",
    "\n",
    "mjd = exps['mjd']\n",
    "order = np.argsort(mjd)\n",
    "mjd = mjd[order]\n",
    "ra = exps['ra'][order]\n",
    "dec = exps['dec'][order]\n",
    "\n",
    "exp_phi = ra\n",
    "exp_theta = np.pi/2 - dec\n",
    "\n",
    "# Rotate to galactic coordinates\n",
    "r = healpy.Rotator(coord=['C','G'])\n",
    "exp_theta, exp_phi = r(exp_theta, exp_phi)\n",
    "\n",
    "fov_deg = 3.5\n",
    "vecs = healpy.ang2vec(exp_theta, exp_phi)\n",
    "radius = np.radians(fov_deg/2.)\n",
    "n_exp = len(vecs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "saving\n",
      "Found pixel angles\n"
     ]
    }
   ],
   "source": [
    "# Get healpix angles of pixel coordinates\n",
    "ny, nx, _ = full_image.shape\n",
    "yy, xx = np.mgrid[0:ny,0:nx]\n",
    "proj=healpy.projector.MollweideProj(xsize=nx)\n",
    "xx=xx.astype(float)/nx * 4. - 2\n",
    "yy=1 - yy.astype(float)/ny * 2.\n",
    "xx=xx.flatten()\n",
    "yy=yy.flatten()\n",
    "\n",
    "\n",
    "# This takes a long time so cache\n",
    "angs_file = 'angs.npy'\n",
    "\n",
    "if os.path.exists(angs_file):\n",
    "    print(\"loading\")\n",
    "    theta,phi = np.load(angs_file)\n",
    "else:\n",
    "    angs = []\n",
    "    for xi, yi in zip(xx,yy):\n",
    "        angs.append(proj.xy2ang(xi, yi))\n",
    "\n",
    "    angs = np.array(angs).T\n",
    "    np.save(angs_file, angs)\n",
    "    print(\"saving\")\n",
    "    theta, phi = angs\n",
    "\n",
    "print(\"Found pixel angles\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Black out the outer region of the image, outside the sky region\n",
    "bad = np.isnan(theta) | np.isnan(phi) | (theta<0) | (theta>np.pi)\n",
    "theta[bad] = 0\n",
    "phi[bad] = 0\n",
    "\n",
    "bad2 = bad.reshape((ny,nx))\n",
    "full_image[bad2,0]=0\n",
    "full_image[bad2,1]=0\n",
    "full_image[bad2,2]=0\n",
    "full_image[bad2,3]=255"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "nside = 2048\n",
    "npix = healpy.nside2npix(nside)\n",
    "pix=healpy.ang2pix(nside,theta,phi)\n",
    "pix=pix.reshape((ny,nx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do this many exposures in the video.\n",
    "# It gets a bit long otherwise!\n",
    "exposure_max = np.where(mjd > mjd[0] + 365)[0][0]\n",
    "\n",
    "# This normalizes the colour brightness.\n",
    "# alpha/alpha_max = sqrt(count/count_max) with alpha_max=255\n",
    "count_max = 20.0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "/usr/local/bin/ffmpeg -y -f rawvideo -pix_fmt rgb24 -s 4106x2054 -i - -pix_fmt yuv420p -r 60 -c:v libx264 /Users/jaz/working/anims/exps.mov\n",
      "1000\n",
      "2000\n",
      "3000\n",
      "4000\n",
      "5000\n",
      "6000\n",
      "7000\n",
      "8000\n",
      "9000\n",
      "10000\n",
      "11000\n",
      "12000\n",
      "13000\n",
      "14000\n",
      "15000\n",
      "16000\n",
      "17000\n",
      "18000\n",
      "19000\n",
      "20000\n",
      "21000\n",
      "22000\n",
      "23000\n",
      "24000\n",
      "25000\n",
      "26000\n",
      "27000\n",
      "28000\n",
      "29000\n",
      "30000\n",
      "31000\n",
      "32000\n",
      "33000\n",
      "34000\n",
      "35000\n",
      "36000\n",
      "37000\n",
      "38000\n",
      "39000\n",
      "40000\n",
      "41000\n",
      "42000\n"
     ]
    }
   ],
   "source": [
    "outputdict={\"-pix_fmt\":\"yuv420p\", \"-r\":\"60\", \"-c:v\": \"libx264\"}\n",
    "writer = skvideo.io.FFmpegWriter(\"exps.mov\", outputdict=outputdict, verbosity=1)\n",
    "mask = np.zeros((ny,nx), dtype=bool)\n",
    "pixmask = np.zeros(npix, dtype=bool)\n",
    "rgb_image = full_image[:,:,:3].copy()\n",
    "\n",
    "\n",
    "# set alpha to zero\n",
    "\n",
    "output_cumulative = np.zeros((ny,nx,3), dtype=np.uint8)\n",
    "output_flash = np.zeros((ny,nx,3), dtype=np.uint8)\n",
    "\n",
    "count = np.zeros((ny,nx), dtype=float)\n",
    "alpha = np.zeros((ny,nx), dtype=float)\n",
    "frame_alpha = np.zeros((ny,nx), dtype=float)\n",
    "output_image = np.zeros_like(rgb_image)\n",
    "\n",
    "\n",
    "\n",
    "for i in range(exposure_max):\n",
    "    if (i%1000)==0:\n",
    "        print(i)\n",
    "    disc = healpy.query_disc(nside, vecs[i,:], radius)\n",
    "\n",
    "    # healpix index to x,y index\n",
    "    # am sure we could speed this up\n",
    "    pixmask[disc] = 1\n",
    "    mask[:,:] = pixmask[pix[:,:]]\n",
    "    w = np.where(mask)\n",
    "\n",
    "    count[w] += 1\n",
    "    alpha_w = 255 * np.sqrt(count[w]/count_max)\n",
    "    alpha[w] = np.clip(alpha_w, 0, 255)\n",
    "\n",
    "    output_cumulative[w] = np.einsum('ij,i->ij', rgb_image[w], (alpha[w]/255.0))\n",
    "    output_flash[w] = 255\n",
    "\n",
    "    writer.writeFrame(output_cumulative)\n",
    "    writer.writeFrame(output_flash)\n",
    "\n",
    "    output_flash[w] = output_cumulative[w]\n",
    "\n",
    "\n",
    "    # reset masks\n",
    "    pixmask[disc] = 0\n",
    "    mask[w] = False\n",
    "\n",
    "\n",
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}