ztf_axs_with_lsdb_and_tape.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8a84516b-5bd7-44bb-8f9b-33a6479d687a",
   "metadata": {},
   "source": [
    "# Load and process ZTF AXS (DR14) with LSDB and Tape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5d485791",
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import monotonic\n",
    "\n",
    "import light_curve as licu\n",
    "import lsdb\n",
    "import numpy as np\n",
    "from dask.distributed import Client\n",
    "from tape import Ensemble, ColumnMapper"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "105ba99e-6c8a-46b2-a82e-668851d07189",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Use LSDB to load 'object' and 'source' catalogs\n",
    "\n",
    "We do not really read or process anything (but some metadata) until `.compute()` is called in the very end\n",
    "\n",
    "We load few columns only, but actual analysis bellow doesn't really use them all\n",
    "\n",
    "Paths are for PSC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "50c8d823",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 29.1 s, sys: 2.78 s, total: 31.9 s\n",
      "Wall time: 39.1 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "objects = lsdb.read_hipscat(\n",
    "    '/ocean/projects/phy210048p/shared/hipscat/catalogs/ztf_axs/ztf_dr14',\n",
    "    # Select few columns only\n",
    "    columns=['ps1_objid', 'ra', 'dec', 'nobs_g', 'nobs_r', 'nobs_i',\n",
    "             # HiPSCat-specific columns\n",
    "             'Norder', 'Dir', 'Npix'],\n",
    ")\n",
    "sources = lsdb.read_hipscat(\n",
    "    '/ocean/projects/phy210048p/shared/hipscat/catalogs/ztf_axs/ztf_source',\n",
    "    # Select few columns only\n",
    "    columns=['ps1_objid', 'mjd', 'mag', 'magerr', 'catflags', 'band',\n",
    "             # HiPSCat-specific columns\n",
    "             'Norder', 'Dir', 'Npix'],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8aeaaecb-56fa-46f9-8221-e6d8bc052921",
   "metadata": {},
   "source": [
    "### Use LSDB to join objects and sources\n",
    "\n",
    "This would assign sources object's `_hipscat_index`, which we are going to use as a primary key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c64382a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2min 25s, sys: 1.75 s, total: 2min 27s\n",
      "Wall time: 2min 27s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "joined_sources = objects.join(\n",
    "    sources,\n",
    "    left_on='ps1_objid',\n",
    "    right_on='ps1_objid',\n",
    "    output_catalog_name='ztf_axs_sources'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "27db1b22",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><strong>Dask DataFrame Structure:</strong></div>\n",
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>ps1_objid_ztf_dr14</th>\n",
       "      <th>ra_ztf_dr14</th>\n",
       "      <th>dec_ztf_dr14</th>\n",
       "      <th>nobs_g_ztf_dr14</th>\n",
       "      <th>nobs_r_ztf_dr14</th>\n",
       "      <th>nobs_i_ztf_dr14</th>\n",
       "      <th>Norder_ztf_dr14</th>\n",
       "      <th>Dir_ztf_dr14</th>\n",
       "      <th>Npix_ztf_dr14</th>\n",
       "      <th>ps1_objid_ztf_source</th>\n",
       "      <th>mjd_ztf_source</th>\n",
       "      <th>mag_ztf_source</th>\n",
       "      <th>magerr_ztf_source</th>\n",
       "      <th>catflags_ztf_source</th>\n",
       "      <th>band_ztf_source</th>\n",
       "      <th>Norder_ztf_source</th>\n",
       "      <th>Dir_ztf_source</th>\n",
       "      <th>Npix_ztf_source</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>npartitions=311037</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>int64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>int32</td>\n",
       "      <td>int32</td>\n",
       "      <td>int32</td>\n",
       "      <td>int32</td>\n",
       "      <td>int32</td>\n",
       "      <td>int32</td>\n",
       "      <td>int64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float32</td>\n",
       "      <td>float32</td>\n",
       "      <td>int16</td>\n",
       "      <td>string</td>\n",
       "      <td>int32</td>\n",
       "      <td>int32</td>\n",
       "      <td>int32</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>281474976710656</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13834987686537986048</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18446744073709551615</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "<div>Dask Name: from-delayed, 311040 graph layers</div>"
      ],
      "text/plain": [
       "Dask DataFrame Structure:\n",
       "                     ps1_objid_ztf_dr14 ra_ztf_dr14 dec_ztf_dr14 nobs_g_ztf_dr14 nobs_r_ztf_dr14 nobs_i_ztf_dr14 Norder_ztf_dr14 Dir_ztf_dr14 Npix_ztf_dr14 ps1_objid_ztf_source mjd_ztf_source mag_ztf_source magerr_ztf_source catflags_ztf_source band_ztf_source Norder_ztf_source Dir_ztf_source Npix_ztf_source\n",
       "npartitions=311037                                                                                                                                                                                                                                                                                                   \n",
       "0                                 int64     float64      float64           int32           int32           int32           int32        int32         int32                int64        float64        float32           float32               int16          string             int32          int32           int32\n",
       "281474976710656                     ...         ...          ...             ...             ...             ...             ...          ...           ...                  ...            ...            ...               ...                 ...             ...               ...            ...             ...\n",
       "...                                 ...         ...          ...             ...             ...             ...             ...          ...           ...                  ...            ...            ...               ...                 ...             ...               ...            ...             ...\n",
       "13834987686537986048                ...         ...          ...             ...             ...             ...             ...          ...           ...                  ...            ...            ...               ...                 ...             ...               ...            ...             ...\n",
       "18446744073709551615                ...         ...          ...             ...             ...             ...             ...          ...           ...                  ...            ...            ...               ...                 ...             ...               ...            ...             ...\n",
       "Dask Name: from-delayed, 311040 graph layers"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joined_sources"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e3d9d5b-866b-4681-bee0-9eae1e25f4b8",
   "metadata": {},
   "source": [
    "### Let's take only first few partitions of the object table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e1fa1b94-1b5f-48e8-aaf6-b7520037ff3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# n_object_partitions = 5\n",
    "# # off by one because we need last partition **end**\n",
    "# last_object_division = objects._ddf.divisions[n_object_partitions+1]\n",
    "\n",
    "# n_joined_source_partitions = np.searchsorted(joined_sources._ddf.divisions, last_object_division) - 1\n",
    "\n",
    "# print(f'{n_object_partitions = } / {objects._ddf.npartitions}')\n",
    "# print(f'{n_joined_source_partitions = } / {joined_sources._ddf.npartitions}')\n",
    "\n",
    "# object_frame = objects._ddf.partitions[:n_object_partitions]\n",
    "# source_frame = joined_sources._ddf.partitions[:n_joined_source_partitions]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5b87b0a-0240-4b38-9761-a9f348aa38a4",
   "metadata": {},
   "source": [
    "### Or just take all the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "abde14c8-4036-4b3a-b147-d644554152c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "object_frame = objects._ddf\n",
    "source_frame = joined_sources._ddf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b0537b8-7073-4399-8667-9660136115c7",
   "metadata": {},
   "source": [
    "### Create Dask client\n",
    "\n",
    "On a SLURM cluster we may use `dask_jobqueue.SLURMCluster` to scale our job.\n",
    "In this case the current node would be a manager node, and would run none of Dask workers itself.\n",
    "Instead it would run SLURM jobs, each with few workers, and assign Dask tasks for them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fbf58b27",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-e9b360d9-af49-11ee-9efa-fdcdae23d7d5</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> dask_jobqueue.SLURMCluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://10.8.9.182:8787/status\" target=\"_blank\">http://10.8.9.182:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">SLURMCluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">28f8f36f</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://10.8.9.182:8787/status\" target=\"_blank\">http://10.8.9.182:8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 0\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 0\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 0 B\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-6ff62076-d418-49d6-b34b-3ac77dbe3dfe</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://10.8.9.182:35581\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://10.8.9.182:8787/status\" target=\"_blank\">http://10.8.9.182:8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 0\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 0 B\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tcp://10.8.9.182:35581' processes=0 threads=0, memory=0 B>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### Set larger timeouts\n",
    "\n",
    "import dask\n",
    "\n",
    "dask.config.set({\n",
    "    'distributed.comm.timeouts.connect': '3600s',\n",
    "    'distributed.comm.timeouts.tcp': '3600s',\n",
    "})\n",
    "\n",
    "### Create a SLURM cluster\n",
    "\n",
    "from dask_jobqueue import SLURMCluster\n",
    "\n",
    "cluster = SLURMCluster(\n",
    "    # Number of Dask workers per node\n",
    "    processes=8,\n",
    "    # Regular memory node type on PSC bridges2\n",
    "    queue=\"RM\",\n",
    "    # Infiniband should be faster, but it doesn't work well =(\n",
    "    # interface='ib0',\n",
    "    # dask_jobqueue requires cores and memory to be specified\n",
    "    # We set them to match RM specs\n",
    "    cores=128,\n",
    "    memory=\"256GB\",\n",
    "    walltime=\"12:00:00\",\n",
    "    death_timeout=7200.0,\n",
    ")\n",
    "# Scale to up to 20 nodes\n",
    "cluster.scale(jobs=10)\n",
    "# cluster.adapt(maximum_jobs=20)\n",
    "client = Client(cluster)\n",
    "\n",
    "### Or create a local cluster\n",
    "# client = Client()\n",
    "\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aab02951-6acc-4d6e-8c3a-ab8c82be79a3",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Show how to access Dask dashboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bf1aa2f1-c997-4d5c-b8b2-cfd150da04a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Copy-paste and run in your terminal:\n",
      "\n",
      "ssh -N -L 127.0.0.1:8787:10.8.9.182:8787 malanche@bridges2.psc.edu\n",
      "\n",
      "And open this URL in your browser to see the dashboard:\n",
      "http://127.0.0.1:8787/\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Make a command for dashboard ssh-tunneling\n",
    "\n",
    "import socket\n",
    "from getpass import getuser\n",
    "from urllib.parse import urlparse\n",
    "\n",
    "local_addr = '127.0.0.1:8787'\n",
    "remote_host = 'bridges2.psc.edu'\n",
    "\n",
    "with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:\n",
    "    s.connect(('1.1.1.1', 53))\n",
    "    ip = s.getsockname()[0]\n",
    "username = getuser()\n",
    "dashboard_port = urlparse(client.dashboard_link).port\n",
    "\n",
    "print(f'''\n",
    "Copy-paste and run in your terminal:\n",
    "\n",
    "ssh -N -L {local_addr}:{ip}:{dashboard_port} {username}@{remote_host}\n",
    "\n",
    "And open this URL in your browser to see the dashboard:\n",
    "http://{local_addr}/\n",
    "''')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59aca60e-eb45-4e66-be33-bdb8edee8d33",
   "metadata": {},
   "source": [
    "### Create Tape Ensemble and plan the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2fbd5b5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/jet/home/malanche/.virtualenvs/tape/lib/python3.9/site-packages/dask/dataframe/core.py:5400: UserWarning: New index has same name as existing, this is a no-op.\n",
      "  warnings.warn(\n",
      "/jet/home/malanche/.virtualenvs/tape/lib/python3.9/site-packages/dask/dataframe/core.py:5400: UserWarning: New index has same name as existing, this is a no-op.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7056\n",
      "1557537\n",
      "2179612\n",
      "CPU times: user 1.25 s, sys: 91.9 ms, total: 1.34 s\n",
      "Wall time: 1.35 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "ens = Ensemble(client)\n",
    "column_mapper = ColumnMapper(\n",
    "    id_col='_hipscat_index',\n",
    "    time_col='mjd_ztf_source',\n",
    "    flux_col='mag_ztf_source',\n",
    "    err_col='magerr_ztf_source',\n",
    "    band_col='band_ztf_source',\n",
    ")\n",
    "ens.from_dask_dataframe(\n",
    "    object_frame=object_frame,\n",
    "    source_frame=source_frame,\n",
    "    sorted=True,\n",
    "    sort=False,\n",
    "    sync_tables=False,\n",
    "    column_mapper=column_mapper,\n",
    ")\n",
    "\n",
    "# ens.source.query('catflags_ztf_source == 0 & magerr_ztf_source > 0').update_ensemble()\n",
    "# ens = ens.calc_nobs(by_band=False, label=\"nobs\", temporary=False)\n",
    "# ens.object.query('nobs_total >= 2000').update_ensemble()\n",
    "# features = ens.batch(licu.ReducedChi2(), band_to_calc=None, label='features', compute=False)\n",
    "# ens.object.merge(features).update_ensemble()\n",
    "# ens.object.query('chi2 >= 3.0').update_ensemble()\n",
    "ens.object.query('nobs_r == 1000 and nobs_g == 0 and nobs_i == 0').update_ensemble()\n",
    "print(len(ens.object.dask))\n",
    "new_object = ens.object.repartition(npartitions=1)\n",
    "new_object.ensemble = ens\n",
    "new_object.update_ensemble()\n",
    "# ens.source.query('catflags_ztf_source == 0 & magerr_ztf_source > 0').update_ensemble()\n",
    "features = ens.batch(licu.Amplitude(), band_to_calc=None, label='features')\n",
    "print(len(features.dask))\n",
    "max_amplitude = features['amplitude'].max()\n",
    "print(len(max_amplitude.dask))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fb15e074-223d-4310-8882-6a670d0d98ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.179612"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(max_amplitude.dask) / 1e6\n",
    "# len(ens.source.dask) / 1e6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a893428b-b085-4c72-bf2b-dec8b05ead99",
   "metadata": {},
   "source": [
    "### Run the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6b5ec24",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/jet/home/malanche/.virtualenvs/tape/lib/python3.9/site-packages/distributed/client.py:3163: UserWarning: Sending large graph of size 121.21 MiB.\n",
      "This may cause some slowdown.\n",
      "Consider scattering data ahead of time and using futures.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "t = monotonic()\n",
    "result = max_amplitude.compute()\n",
    "dt = monotonic() - t\n",
    "dt, result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f02b0c8-f63e-46a6-8a47-95cc7c89a921",
   "metadata": {},
   "source": [
    "### Save results to disk, just in case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c86cb6c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('cluster_result.txt', 'w') as f:\n",
    "    f.write(f'{dt = }\\n{result = }\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adc931fb-8e1a-4180-9de1-05acd08f877b",
   "metadata": {},
   "source": [
    "### Shut down Dask cluster\n",
    "\n",
    "If we run SLURM cluster it would also cancel all associated slurm jobs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b89a106e",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tape-venv-python3.9",
   "language": "python",
   "name": "tape-venv-python3.9"
  },
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}