bd_averager_2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# bd_averager_2.py notebook\n",
    "\n",
    "### Simon Perkins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import argparse\n",
    "import numpy as np\n",
    "from xarrayms import xds_from_ms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Can't create a parser since this is a notebook, its interactive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def create_parser():\n",
    "#     p = argparse.ArgumentParser()\n",
    "#     p.add_argument(\"ms\")\n",
    "#     return p\n",
    "# args = create_parser().parse_args()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _avg_time_bins(max_ew, max_uvw, longest_baseline_bins=2):\n",
    "    return np.ceil(max_ew / max_uvw).astype(np.int32) * longest_baseline_bins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def baseline_average_scan(time, ant1, ant2, uvw, flag_row, max_uvw):\n",
    "    # Find unique baselines\n",
    "    baselines = np.stack([ant1, ant2], axis=1)\n",
    "    ubl, inv = np.unique(baselines, return_inverse=True, axis=0)\n",
    "\n",
    "    # Generate a mask for each unflagged row\n",
    "    # containing the unique baseline\n",
    "    bl_mask = np.arange(ubl.shape[0])[:, None] == inv[None, :]\n",
    "    # Remove flagged data\n",
    "    unflagged = flag_row == False # noqa\n",
    "    bl_mask[:, :] &= unflagged[None, :]\n",
    "\n",
    "    # Compute the maximum EW distance and\n",
    "    # number of rows for each baseline\n",
    "    max_ew_dist = np.empty(bl_mask.shape[0], dtype=uvw.dtype)\n",
    "    nrows = np.empty(bl_mask.shape[0], dtype=np.int32)\n",
    "\n",
    "    for i, mask in enumerate(bl_mask):\n",
    "        max_ew_dist[i] = np.abs(uvw[mask, 0]).sum()\n",
    "        nrows[i] = np.count_nonzero(mask)\n",
    "\n",
    "    # Compute the average time bins for each baseline\n",
    "    avg_time_bins = _avg_time_bins(max_ew_dist, max_uvw)\n",
    "    # Clamp number of bins to number of rows\n",
    "    avg_time_bins = np.minimum(avg_time_bins, nrows)\n",
    "    # Number of output rows for each baseline\n",
    "    out_nrows = nrows // avg_time_bins\n",
    "    out_rem = nrows % avg_time_bins\n",
    "\n",
    "    for mask, in_rows, out_rows, rem, bins in zip(bl_mask, nrows, out_nrows, out_rem, avg_time_bins):\n",
    "        # No averaging required\n",
    "        if out_rows == in_rows:\n",
    "            continue\n",
    "\n",
    "        bl_uvw = uvw[mask, ...]\n",
    "        assert bl_uvw.shape[0] == in_rows\n",
    "\n",
    "        tot_rows = out_rows if rem == 0 else out_rows + 1\n",
    "\n",
    "        avg_uvw = np.empty((tot_rows, 3), dtype=uvw.dtype)\n",
    "        avg_uvw[:out_rows, :] = bl_uvw[:out_rows*bins, :].reshape(out_rows, bins, 3).mean(axis=1)\n",
    "\n",
    "        if rem > 0:\n",
    "            avg_uvw[out_rows:, :] = bl_uvw[out_rows*bins:, :].mean(axis=0)\n",
    "\n",
    "    return avg_uvw"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Read the MS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Successful readonly open of default-locked table 1491291289.1GC.ms: 26 columns, 129480 rows\n"
     ]
    }
   ],
   "source": [
    "xds = list(xds_from_ms(\"1491291289.1GC.ms\",\n",
    "                       columns=[\"TIME\", \"ANTENNA1\", \"ANTENNA2\",\n",
    "                                \"UVW\", \"FLAG_ROW\"],\n",
    "                       group_cols=[],\n",
    "                       index_cols=[],\n",
    "                       chunks={\"row\": 1e9}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xds[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "xarray.core.dataset.Dataset"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the max_uvw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_uvw = np.sqrt(np.max(np.sum(ds.UVW.data ** 2, axis=1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "ename": "NotImplementedError",
     "evalue": "The out parameter is not fully supported. Received type ndarray, expected Dask Array",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNotImplementedError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-9-d422a4701bfe>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m baseline_average_scan(ds.TIME.data, ds.ANTENNA1.data, ds.ANTENNA2.data,\n\u001b[0;32m----> 2\u001b[0;31m                       ds.UVW.data, ds.FLAG_ROW.data, max_uvw)\n\u001b[0m",
      "\u001b[0;32m<ipython-input-3-8d5c42eb7410>\u001b[0m in \u001b[0;36mbaseline_average_scan\u001b[0;34m(time, ant1, ant2, uvw, flag_row, max_uvw)\u001b[0m\n\u001b[1;32m      9\u001b[0m     \u001b[0;31m# Remove flagged data\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m     \u001b[0munflagged\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mflag_row\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mFalse\u001b[0m \u001b[0;31m# noqa\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 11\u001b[0;31m     \u001b[0mbl_mask\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m&=\u001b[0m \u001b[0munflagged\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     12\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m     \u001b[0;31m# Compute the maximum EW distance and\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/fundamentals/lib/python2.7/site-packages/dask/array/core.pyc\u001b[0m in \u001b[0;36m__array_ufunc__\u001b[0;34m(self, numpy_ufunc, method, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m   1214\u001b[0m                 \u001b[0;32mreturn\u001b[0m \u001b[0mda_ufunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1215\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1216\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0melemwise\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnumpy_ufunc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1217\u001b[0m         \u001b[0;32melif\u001b[0m \u001b[0mmethod\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'outer'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1218\u001b[0m             \u001b[0;32mfrom\u001b[0m \u001b[0;34m.\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mufunc\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/fundamentals/lib/python2.7/site-packages/dask/array/core.pyc\u001b[0m in \u001b[0;36melemwise\u001b[0;34m(op, *args, **kwargs)\u001b[0m\n\u001b[1;32m   3499\u001b[0m                   **atop_kwargs)\n\u001b[1;32m   3500\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3501\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mhandle_out\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3502\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3503\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/fundamentals/lib/python2.7/site-packages/dask/array/core.pyc\u001b[0m in \u001b[0;36mhandle_out\u001b[0;34m(out, result)\u001b[0m\n\u001b[1;32m   3527\u001b[0m         msg = (\"The out parameter is not fully supported.\"\n\u001b[1;32m   3528\u001b[0m                \" Received type %s, expected Dask Array\" % type(out).__name__)\n\u001b[0;32m-> 3529\u001b[0;31m         \u001b[0;32mraise\u001b[0m \u001b[0mNotImplementedError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmsg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3530\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3531\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNotImplementedError\u001b[0m: The out parameter is not fully supported. Received type ndarray, expected Dask Array"
     ]
    }
   ],
   "source": [
    "baseline_average_scan(ds.TIME.data, ds.ANTENNA1.data, ds.ANTENNA2.data,\n",
    "                      ds.UVW.data, ds.FLAG_ROW.data, max_uvw)"
   ]
  },
  {
   "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": []
  },
  {
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}