FB12-Satellite rate model.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "import sys\n",
    "sys.path.insert(0, \"/home/yymao/Dropbox/Academia/Collaborations/SAGA-local/saga\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import SAGA\n",
    "print(SAGA.__version__)\n",
    "\n",
    "from SAGA import ObjectCuts as C\n",
    "from easyquery import *\n",
    "from SAGA.utils import *\n",
    "from SAGA.utils.distance import d2m\n",
    "from SAGA.targets import add_cut_scores\n",
    "from SAGA.objects.object_catalog import calc_fiducial_p_sat\n",
    "from astropy.table import vstack, Table, join\n",
    "from astropy.coordinates import Distance\n",
    "\n",
    "from collections import defaultdict\n",
    "from itertools import chain\n",
    "\n",
    "import numexpr as ne\n",
    "from scipy.stats import binned_statistic\n",
    "from scipy.stats import iqr\n",
    "from scipy.stats import norm\n",
    "from scipy.optimize import minimize_scalar\n",
    "from functools import reduce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from statsmodels.discrete.discrete_model import Logit\n",
    "\n",
    "class MyModel(Logit):\n",
    "    def __init__(self, endog, exog, **kwargs):\n",
    "        self._max_value = float(kwargs.get(\"max_value\", 1))\n",
    "        if not (self._max_value > 0 and self._max_value <= 1):\n",
    "            raise ValueError\n",
    "        super(MyModel, self).__init__(endog, exog, **kwargs)\n",
    "        \n",
    "    def _Lambda(self, X):\n",
    "        return \n",
    "        \n",
    "    def cdf(self, X):\n",
    "        return self._max_value * super(MyModel, self).cdf(X)\n",
    "    \n",
    "    def pdf(self, X):\n",
    "        return self._max_value * super(MyModel, self).pdf(X)\n",
    "    \n",
    "    def loglike(self, params):\n",
    "        return np.sum(self.loglikeobs(params))\n",
    "        \n",
    "    def loglikeobs(self, params):\n",
    "        y = self.endog\n",
    "        X = self.exog\n",
    "        L = self.cdf(np.dot(X, params))\n",
    "        return np.log(L) * y + np.log1p(-L) * (1-y)\n",
    "    \n",
    "    def hessian(self, params):\n",
    "        X = self.exog\n",
    "        L = super(MyModel, self).cdf(np.dot(X,params))\n",
    "        return -np.dot(self._max_value*L*(1-L)*X.T,X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_label_arr(data, valid_query, *value_setters):\n",
    "    \n",
    "    label_arr = np.empty(len(data))\n",
    "    label_arr.fill(-1.0)\n",
    "    label_arr[Query(valid_query).mask(data)] = 0\n",
    "\n",
    "    for q, setter in value_setters:\n",
    "        q_this = Query(valid_query, q)\n",
    "        if callable(setter):\n",
    "            value = setter(q_this.filter(data))\n",
    "        else:\n",
    "            value = setter\n",
    "        label_arr[q_this.mask(data)] = value\n",
    "    \n",
    "    return label_arr\n",
    "\n",
    "class LogitFit:\n",
    "    def __init__(self, data, label_arr, feature_cols=(\"r_mag\", \"sb_r\", \"gr\"), max_prob=1, add_prob_column=None, silent=False):\n",
    "        valid_label_mask = (label_arr >= 0) & (label_arr <= 1)       \n",
    "        self._feature_cols = list(feature_cols)\n",
    "        feature_arr = get_feature_array(data, self._feature_cols, True)\n",
    "        logit_mod = MyModel(label_arr[valid_label_mask], feature_arr[valid_label_mask], max_value=max_prob)\n",
    "        self.logit_res = logit_mod.fit(disp=False)\n",
    "        if add_prob_column:\n",
    "            if add_prob_column is True:\n",
    "                add_prob_column = \"p_sat\"\n",
    "            data[add_prob_column] = self.calc_prob(feature_arr, False)\n",
    "        if not silent:\n",
    "            display(self.logit_res.summary())\n",
    "\n",
    "    def calc_prob(self, data, convert_arr_format=True):\n",
    "        feature_arr = get_feature_array(data, self._feature_cols, True) if convert_arr_format else data\n",
    "        return self.logit_res.predict(feature_arr)\n",
    "    \n",
    "def calc_prob_with_params(base, params, max_prob=1, feature_cols=(\"r_mag\", \"sb_r\", \"gr\")):\n",
    "        mu = np.zeros(len(base))\n",
    "        for p, col in zip(params, feature_cols):\n",
    "            mu += base[col] * p\n",
    "        mu += params[-1]\n",
    "        mu = np.where(np.isnan(mu), np.inf, mu)\n",
    "        p = max_prob / (1 + np.exp(-mu))\n",
    "        return p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "saga = SAGA.QuickStart()\n",
    "saga.set_default_base_version(\"paper2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_feature_array(table, cols, add_constant=False):\n",
    "    data = [table[col].astype(np.float64) for col in cols]\n",
    "    if add_constant:\n",
    "        data.append(np.ones(len(table), dtype=np.float64))\n",
    "    return np.vstack(data).T "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def annotate_base(base):\n",
    "    base[\"Mr_\"] = base[\"r_mag\"] - d2m(base[\"HOST_DIST\"])\n",
    "    base[\"Mr\"] = np.where(np.isfinite(base[\"Mr\"]) & C.has_spec.mask(base), base[\"Mr\"], base[\"Mr_\"])\n",
    "    base[\"log_sm\"] = np.where(np.isfinite(base[\"log_sm\"]), base[\"log_sm\"], 1.254 + 1.0976 * base[\"gr\"] - 0.4 * base[\"Mr\"])\n",
    "    \n",
    "    base[\"human_selected\"] = 0\n",
    "    hs = saga.database[\"human_selected\"].read()\n",
    "    for i in range(1, 4):\n",
    "        base[\"human_selected\"] += np.in1d(base[\"OBJID\"], Query(f\"score == {i}\").filter(hs, \"OBJID\")).astype(np.int) * i\n",
    "\n",
    "    base[\"gr_c\"] = np.where(\n",
    "        C.valid_g_mag.mask(base), \n",
    "        base[\"gr\"],\n",
    "        np.where(\n",
    "            C.valid_i_mag.mask(base), \n",
    "            base[\"ri\"]+0.1, \n",
    "            np.where(C.valid_z_mag.mask(base), base[\"rz\"]+0.2, 0.92-0.03*base[\"r_mag\"]),\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    base[\"sb_c\"] = np.where(\n",
    "        C.valid_sb.mask(base), \n",
    "        base[\"sb_r\"],\n",
    "        20.75+0.6*(base[\"r_mag\"]-14),\n",
    "    )\n",
    "\n",
    "    return base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base = SAGA.database.FitsTable(\"/home/yymao/Documents/Research/SAGA/PaperII/data-archive/saga_base_all.fits\").read()\n",
    "base = annotate_base(base)\n",
    "hosts = np.unique(base[\"HOSTID\"])\n",
    "print(len(hosts))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_other = SAGA.database.FitsTable(\"/home/yymao/Documents/Research/SAGA/PaperII/data-archive/saga_base_other.fits\").read()\n",
    "base_other = annotate_base(base_other)\n",
    "hosts_other = np.unique(base_other[\"HOSTID\"])\n",
    "print(len(hosts_other))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HostQuery = saga.host_catalog.construct_host_query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sup = np.testing.suppress_warnings()\n",
    "sup.filter(message=\"invalid value encountered in true_divide\")  # module must match exactly\n",
    "\n",
    "@sup\n",
    "def plot_prob_hist(ax, data, col_x, bins, additonal_label=None, label_query=C.is_sat):\n",
    "    ys = [label_query.mask(data), np.ones(len(data), dtype=np.int), (data['p_sat'] if 'p_sat' in data.colnames else np.zeros(len(data)))]\n",
    "    s = binned_statistic(data[col_x], ys, bins=bins, statistic='sum')\n",
    "    \n",
    "    dn = np.where(s.statistic[1], s.statistic[1], 1)\n",
    "    p = s.statistic[0] / dn\n",
    "    e = np.maximum(np.sqrt(s.statistic[0] * (s.statistic[1] - s.statistic[0])), 1) * 0.84 / (dn**1.5)\n",
    "    eh = np.where(e+p <= 1, e, 1-p)\n",
    "    el = np.where(p-e >= 0, e, p)\n",
    "    x = (s.bin_edges[1:] + s.bin_edges[:-1]) * 0.5\n",
    "    \n",
    "    ax.bar(x, p, s.bin_edges[1:] - s.bin_edges[:-1], alpha=0.6, color='C0');\n",
    "    ax.errorbar(x, p, [el, eh], ls='', c='C0');\n",
    "    \n",
    "    if 'p_sat' in data.colnames:\n",
    "        ax.plot(x, s.statistic[2] / s.statistic[1], c='C1')\n",
    "    ax.set_xlabel(col_x + (\", {}\".format(additonal_label) if additonal_label else \"\"))\n",
    "    \n",
    "    ax.set_ylim(0, max(np.nanmax(s.statistic[2] / s.statistic[1]), np.nanmax(p))*1.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sat_prob_setter = (\n",
    "    #(Query(\"human_selected > 0\", C.paper2_targeting_cut, ~C.has_spec), 0.4),\n",
    "    (C.is_sat, 1),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid = Query(\n",
    "    HostQuery(\"paper2_need_spec / paper2_total < 0.2\"),\n",
    "    \"r_mag > 13\", \n",
    "    \"r_mag < 20.75\", \n",
    "    reduce(lambda a, b: (a | b), [q[0] for q in sat_prob_setter], Query(C.has_spec, \"SPEC_Z >= 0.003\")),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_cols = (\"r_mag\", \"sb_c\", \"gr_c\")\n",
    "valid_base = valid.filter(base, list(feature_cols) + [\"Mr_\", \"SATS\", \"OBJID\", \"SPEC_Z\", \"ZQUALITY\", \"gr\", \"gr_err\", \"human_selected\", \"g_mag\", \"HOST_VHOST\", \"radius\", \"HOST_DIST\", \"sb_r\", \"sb_r_err\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def obj_func(max_prob):\n",
    "    model_sat = LogitFit(\n",
    "        valid_base, \n",
    "        generate_label_arr(\n",
    "            valid_base, \n",
    "            Query(),\n",
    "            *sat_prob_setter\n",
    "        ), \n",
    "        max_prob=max_prob,\n",
    "        feature_cols=feature_cols,\n",
    "        add_prob_column=False,\n",
    "        silent=True,\n",
    "    )\n",
    "    return -model_sat.logit_res.llf\n",
    "\n",
    "res = minimize_scalar(obj_func, bounds=(0.2, 1), method=\"bounded\")\n",
    "max_prob = np.round(res.x, 3)\n",
    "print(res)\n",
    "print(max_prob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_sat = LogitFit(\n",
    "    valid_base, \n",
    "    generate_label_arr(\n",
    "        valid_base, \n",
    "        Query(),\n",
    "        *sat_prob_setter\n",
    "    ), \n",
    "    max_prob=max_prob,\n",
    "    feature_cols=feature_cols,\n",
    "    add_prob_column=False,\n",
    ")\n",
    "params_best = model_sat.logit_res.params.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = []\n",
    "\n",
    "pick = reduce(lambda a, b: (a | b), [q[0] for q in sat_prob_setter], C.is_sat)\n",
    "\n",
    "sats_idx = np.flatnonzero(pick.mask(valid_base))\n",
    "nonsats_idx = np.flatnonzero((~pick).mask(valid_base))\n",
    "n_sats = len(sats_idx)\n",
    "\n",
    "for i in range(100):\n",
    "    sats_idx_this = sats_idx[np.random.randint(n_sats, size=n_sats)]\n",
    "    base_this = valid_base[np.concatenate([sats_idx_this, nonsats_idx])]\n",
    "    model_sat = LogitFit(\n",
    "        base_this, \n",
    "        generate_label_arr(\n",
    "            base_this, \n",
    "            Query(),\n",
    "            *sat_prob_setter\n",
    "        ), \n",
    "        max_prob=max_prob,\n",
    "        feature_cols=feature_cols,\n",
    "        add_prob_column=False,\n",
    "        silent=True,\n",
    "    )\n",
    "    params.append(model_sat.logit_res.params.copy())\n",
    "\n",
    "params = np.array(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_probs(base, params):\n",
    "    prob_sat_collect = np.array([calc_prob_with_params(base, params_this, max_prob, feature_cols) for params_this in params])\n",
    "\n",
    "    base[\"p_sat_84_o\"], base[\"p_sat_50_o\"], base[\"p_sat_16_o\"] = np.percentile(prob_sat_collect, norm.sf([-1, 0, 1])*100.0, axis=0)\n",
    "    base[\"p_sat_o\"] = base[\"p_sat_50_o\"]\n",
    "    \n",
    "    for q, p in sat_prob_setter:\n",
    "        if p < 1:\n",
    "            prob_sat_collect[:,q.mask(base)] = p + np.random.randn(q.count(base))*1e-5\n",
    "        else:\n",
    "            prob_sat_collect[:,q.mask(base)] = p\n",
    "    \n",
    "    base[\"p_sat_84\"], base[\"p_sat_50\"], base[\"p_sat_16\"] = np.percentile(prob_sat_collect, norm.sf([-1, 0, 1])*100.0, axis=0)\n",
    "    base[\"p_sat\"] = base[\"p_sat_50\"]\n",
    "    return base, prob_sat_collect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base, prob_sat = add_probs(base, params)\n",
    "base_other, prob_sat_other = add_probs(base_other, params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "base[\"p_sat_best\"] = calc_prob_with_params(base, params_best, max_prob, feature_cols)\n",
    "base_other[\"p_sat_best\"] = calc_prob_with_params(base_other, params_best, max_prob, feature_cols)\n",
    "\n",
    "params_approx = np.round(params_best, 3).tolist()\n",
    "base[\"p_sat_approx_new\"] = calc_prob_with_params(base, params_approx, max_prob, feature_cols)\n",
    "base_other[\"p_sat_approx_new\"] = calc_prob_with_params(base_other, params_approx, max_prob, feature_cols)\n",
    "print(params_approx + [max_prob])\n",
    "\n",
    "for q, p in sat_prob_setter:\n",
    "    if p < 1:\n",
    "        base[\"p_sat_approx_new\"][q.mask(base)] = p + np.random.randn(q.count(base))*1e-5\n",
    "    else:\n",
    "        base[\"p_sat_approx_new\"][q.mask(base)] = p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt = norm.sf([-1, 0, 1])*100.0\n",
    "X = []\n",
    "\n",
    "high_priority_sb_new = Query(\"sb_r + abs(sb_r_err) - 0.7 * (r_mag - 14) > 18.5\") | (~C.valid_sb)\n",
    "main_targeting_cuts = high_priority_sb_new & C.high_priority_gr\n",
    "\n",
    "base[\"idx\"] = np.arange(len(base))\n",
    "base_other[\"idx\"] = np.arange(len(base_other))\n",
    "\n",
    "for base_this in chain(base.group_by(\"HOSTID\").groups, base_other.group_by(\"HOSTID\").groups):\n",
    "    hostid = base_this[\"HOSTID\"][0]\n",
    "    prob_this = (prob_sat if hostid in hosts else prob_sat_other)[:,base_this[\"idx\"]]\n",
    "    \n",
    "    mask = Query(C.faint_end_limit, ~C.has_spec).mask(base_this)\n",
    "    row = [\n",
    "        hostid, \n",
    "        Query(C.basic_cut2, main_targeting_cuts).count(base_this),\n",
    "        Query(C.basic_cut2, main_targeting_cuts, ~C.has_spec).count(base_this),\n",
    "        Query(C.is_sat, C.faint_end_limit).count(base_this),\n",
    "        base_this[\"p_sat_approx\"][mask].sum(),\n",
    "        base_this[\"p_sat_approx_new\"][mask].sum(),\n",
    "        base_this[\"p_sat_best\"][mask].sum(),\n",
    "    ]\n",
    "    row.extend(np.percentile(prob_this[:,mask].sum(axis=1), pt))\n",
    "    X.append(tuple(row))\n",
    "\n",
    "X = Table(np.array(X, np.dtype([\n",
    "    ('HOSTID', '<U10'), \n",
    "    ('really_need_spec_total', np.int), \n",
    "    ('really_need_spec', np.int), \n",
    "    ('n_sats', np.int), \n",
    "    ('sats_missed_approx', np.float),\n",
    "    ('sats_missed_approx_new', np.float),\n",
    "    ('sats_missed_best', np.float),\n",
    "    ('sats_missed_84', np.float),\n",
    "    ('sats_missed_pred', np.float),\n",
    "    ('sats_missed_16', np.float),\n",
    "])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(4.8, 4))\n",
    "\n",
    "X = X[np.lexsort([X[\"sats_missed_pred\"], X[\"really_need_spec\"]])]\n",
    "\n",
    "dx = np.ones(len(X))\n",
    "dx -= 0.01\n",
    "dx[::2] += 0.02\n",
    "X[\"really_need_spec_dx\"] = np.maximum(X[\"really_need_spec\"], 1) * dx\n",
    "\n",
    "p_cut = 0.8\n",
    "\n",
    "#complete_def_approx = Query(saga.paper2_complete_definition, QueryMaker.in1d(\"HOSTID\", hosts))\n",
    "#complete_def_approx = Query(f\"sats_missed_approx < 0.33\", f\"really_need_spec < {n_cut}\", QueryMaker.in1d(\"HOSTID\", hosts))\n",
    "complete_def = Query(f\"really_need_spec / really_need_spec_total < (1-{p_cut})\")\n",
    "#complete_def = saga.host_catalog.construct_host_query(\"paper2_complete\")\n",
    "p1_complete= HostQuery(\"paper1_complete\")\n",
    "\n",
    "#assert set(complete_def_approx.filter(X, \"HOSTID\")) == set(complete_def.filter(X, \"HOSTID\"))\n",
    "#assert Query(p1_complete, ~complete_def).count(X) == 0\n",
    "\n",
    "for q, color, label, marker in (\n",
    "    (Query(~complete_def), 'C0', \"Incomplete hosts ({})\", 'o'),\n",
    "    (Query(complete_def, ~p1_complete), 'C3', 'Newly complete hosts ({})', 's'),\n",
    "    (Query(p1_complete), 'C1', 'Paper I complete hosts ({})', '^'),\n",
    "):\n",
    "    X1 = q.filter(X)\n",
    "    n = 1#np.maximum(1,X1[\"n_sats\"])\n",
    "    plt.errorbar(1-X1[\"really_need_spec\"]/X1[\"really_need_spec_total\"], X1[\"sats_missed_pred\"] / n, \n",
    "                 yerr=[(X1[\"sats_missed_pred\"]-X1[\"sats_missed_16\"]) / n, (X1[\"sats_missed_84\"]-X1[\"sats_missed_pred\"]) / n], marker=marker, ls='', \n",
    "                 label=label.format(len(X1)), ms=3, color=color)\n",
    "\n",
    "    \n",
    "#plt.xscale(\"log\")\n",
    "plt.yscale(\"log\")\n",
    "#plt.xlim(1, None)\n",
    "\n",
    "\n",
    "xlim = plt.gca().get_xlim()\n",
    "ylim = plt.gca().get_ylim()\n",
    "\n",
    "#plt.axvspan(n_cut, xlim[-1], color=\"C7\", alpha=0.2)\n",
    "#y_cut = np.log(p_cut / ylim[0]) / np.log(ylim[1] / ylim[0])\n",
    "\n",
    "#plt.axvspan(xlim[0], n_cut, ymax=y_cut, color=\"C2\", alpha=0.25)\n",
    "#plt.axvspan(xlim[0], n_cut, ymin=y_cut, color=\"C6\", alpha=0.2)\n",
    "\n",
    "#plt.text(xlim[0], p_cut*1.2, \"    needs single-slit\\n$\\\\uparrow$ follow-up\", fontsize=12)\n",
    "#plt.text(xlim[0], p_cut/1.1, r\"$\\downarrow$ complete\", fontsize=12, va=\"top\")\n",
    "#plt.text(n_cut, ylim[1]*0.95, \"$\\\\rightarrow$ needs more\\n      pointings\", fontsize=12, va=\"top\")\n",
    "\n",
    "#plt.axhline(p_cut, color=\"C7\", ls=\"--\")\n",
    "plt.axvline(p_cut, color=\"C7\", ls=\"--\")\n",
    "\n",
    "\n",
    "plt.legend(loc=\"lower left\", markerfirst=True, handletextpad=-0.25, fontsize=13, borderpad=0)\n",
    "plt.ylim(0.05, 10)\n",
    "plt.xlabel(\"Spec coverage in primary targeting region\")\n",
    "plt.ylabel(\"Model incompleteness ($\\sum\\, \\mathcal{R}_{\\mathrm{sat},i}$)\")\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = valid.filter(base)\n",
    "\n",
    "fig, ax = plt.subplots(ncols=3, figsize=(10, 3.5))\n",
    "\n",
    "for i, ax_this in enumerate(ax):\n",
    "    if i == 0:\n",
    "        key = \"r_mag\"\n",
    "        ax_this.set_xlabel(r'$r_o$')\n",
    "        bins = np.linspace(13, 21, 10)\n",
    "        ax_this.set_xlim(bins[0], bins[-1])\n",
    "        bins[-1] = 20.75\n",
    "        #ax_this.set_xticks([13, 15, 17, 19, 21])\n",
    "        #ax_this.axvline(20.75, c=\"grey\", lw=1)\n",
    "        #ax_this.axvspan(-12.3, -11, color=\"grey\", alpha=0.2)\n",
    "        \n",
    "    elif i == 1:\n",
    "        key = \"sb_r\"\n",
    "        ax_this.set_xlabel(r'$\\mu_{r_o,{\\rm eff}}$')\n",
    "        bins = np.linspace(18.75, 26.25, 11)\n",
    "        ax_this.set_xlim(bins[0], bins[-1])\n",
    "        ax_this.set_xticks([19, 20, 21, 22, 23, 24, 25, 26]) \n",
    "        #ax_this.axvline(18.75, c=\"grey\", lw=1)\n",
    "        #ax_this.axvspan(18, 18.75, color=\"grey\", alpha=0.2)\n",
    "        \n",
    "    elif i == 2:\n",
    "        key = \"gr\"\n",
    "        ax_this.set_xlabel(r'$(g-r)_o$')\n",
    "        bins = np.linspace(-0.05, 0.85, 10)\n",
    "        ax_this.set_xlim(bins[0], bins[-1])\n",
    "        ax_this.set_xticks([0, 0.2, 0.4, 0.6, 0.8])\n",
    "        #ax_this.axvspan(0.9, 1, color=\"grey\", alpha=0.2)\n",
    "        \n",
    "    x = bins[:-1]\n",
    "    width = np.ediff1d(bins)\n",
    "    bin_index = np.digitize(data[key], bins)\n",
    "    \n",
    "    \n",
    "    y = np.bincount(bin_index, minlength=len(bins)+1)[1:-1]\n",
    "    y1 = np.zeros_like(y, dtype=np.float64)\n",
    "    for q, p in sat_prob_setter:\n",
    "        y1 += np.bincount(bin_index[q.mask(data)], minlength=len(bins)+1)[1:-1] * p    \n",
    "    ax_this.bar(x, y1/y, width=width, align='edge', alpha=0.3, label=\"Data\", color=\"C1\")\n",
    "    \n",
    "    p = (y1+1)/(y+2)\n",
    "    e = np.sqrt(p*(1-p) / y)\n",
    "    eh = np.where(e+p <= 1, e, 1-p)\n",
    "    el = np.where(p-e >= 0, e, p)\n",
    "    x = (bins[1:] + bins[:-1]) * 0.5\n",
    "    ax_this.errorbar(x, p, [el, eh], ls='', c='C1')\n",
    "\n",
    "    s1 = binned_statistic(data[key], data[\"p_sat_16_o\"], bins=bins, statistic='sum')\n",
    "    s2 = binned_statistic(data[key], data[\"p_sat_84_o\"], bins=bins, statistic='sum')\n",
    "\n",
    "    ax_this.fill_between(x, s1.statistic/y, s2.statistic/y, label=\"Model\", color=\"C9\", alpha=0.8)\n",
    "    ax_this.set_yscale(\"log\")\n",
    "    ax_this.set_ylim(1e-4, 1)\n",
    "\n",
    "ax[0].legend()\n",
    "ax[0].set_ylabel(\"Satellite rate\")\n",
    "plt.tight_layout(w_pad=0.1)\n",
    "plt.savefig('/home/yymao/Downloads/model_demo.pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(ncols=2, figsize=(10,4))\n",
    "\n",
    "for ax_this in ax:\n",
    "    ax_this.set_xlim(12.45, 21)\n",
    "    ax_this.set_xlabel(r\"$r_o$\")\n",
    "ax[0].set_ylabel(r\"$\\mu_{r_o, {\\rm eff}}$\")\n",
    "ax[1].set_ylabel(r\"$(g-r)_o$\")\n",
    "\n",
    "cut = Query(\"p_sat >= 1e-3\")\n",
    "\n",
    "t = Query(~cut, \"SATS != 3\").filter(base)\n",
    "ax[0].scatter(t[\"r_mag\"], t[\"sb_c\"], c='grey', s=0.1, lw=0, rasterized=True)\n",
    "ax[1].scatter(t[\"r_mag\"], t[\"gr_c\"], c='grey', s=0.1, lw=0, rasterized=True, label=\"Targets ($\\mathcal{R}_\\mathrm{sat}<10^{-3}$)\")\n",
    "\n",
    "t = Query(cut, \"SATS != 3\").filter(base)\n",
    "t.sort(\"p_sat\")\n",
    "cs = ax[0].scatter(t[\"r_mag\"], t[\"sb_c\"], c=np.log10(t[\"p_sat\"]), s=5, lw=0, cmap=\"winter\", vmin=-3, vmax=np.log10(0.5), rasterized=True, alpha=0.85,)\n",
    "ax[1].scatter(t[\"r_mag\"], t[\"gr_c\"], c=np.log10(t[\"p_sat\"]), s=5, lw=0, cmap=\"winter\", vmin=-3, vmax=np.log10(0.5), rasterized=True, alpha=0.8, label=r\"Targets ($\\mathcal{R}_\\mathrm{sat}\\geq10^{-3}$)\")\n",
    "\n",
    "r = np.linspace(11.8, 21, 10)\n",
    "#ax[0].plot(r, 0.6*(r-14) + 18.55, color='C3',  ls=\"--\",alpha=0.9)\n",
    "#ax[1].plot(r, -0.06*(r-14) + 0.9, color='C3', alpha=0.9, ls=\"--\", label='Targeting cuts')\n",
    "#ax[1].plot(r, -0.05*(r-14) + 0.85, color='C3', alpha=0.4, ls=\"--\", label='Targeting cuts')\n",
    "ax[0].plot(r, 0.7*(r-14) + 18.5, color='C4', ls=\"--\", label='Primary targeting region')\n",
    "ax[1].plot(r, -0.06*(r-14) + 0.9, color='C4', ls=\"--\",)\n",
    "#ax[1].plot(r, -0.06*(r-14) + 1.1, color='C3', ls=\"--\", alpha=0.4)\n",
    "\n",
    "#ax[0].plot(r, 0.7*(r-14) + 18.7, color='C3',  ls=\"--\",alpha=0.9)\n",
    "#ax[1].plot(r, -0.05*(r-14) + 1, color='C3', alpha=0.7, ls=\"--\")\n",
    "\n",
    "ax[0].legend(loc=\"lower right\", scatterpoints=3, frameon=True, fontsize=12)\n",
    "\n",
    "lgnd = ax[1].legend(loc=\"upper left\", scatterpoints=3, frameon=True, fontsize=12)\n",
    "#lgnd.legendHandles[0]._sizes = [10]\n",
    "lgnd.legendHandles[0]._sizes = [20]\n",
    "lgnd.legendHandles[1]._sizes = [20]\n",
    "#lgnd.legendHandles[1].set_color(plt.cm.winter([0.8, 0.5, 0.2]))\n",
    "lgnd.legendHandles[1].set_color(plt.cm.winter([0.8, 0.5, 0.2]))\n",
    "\n",
    "\n",
    "ax[0].set_ylim(17, 27)\n",
    "ax[1].set_ylim(-0.3, 2)\n",
    "fig.tight_layout()\n",
    "\n",
    "cbar = fig.colorbar(cs, ax=ax, pad=0.01, ticks=np.log10([0.001, 0.01, 0.1, 0.5]), label=\"Satellite Rate ($\\\\mathcal{R}_{\\\\rm sat}$)\")\n",
    "cbar.ax.set_yticklabels(['$10^{-3}$', '$0.01$', '$0.1$', '$0.5$']) \n",
    "cbar.ax.set_rasterized(True)\n",
    "plt.savefig('/home/yymao/Downloads/sat-prob.pdf', dpi=200)"
   ]
  },
  {
   "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}