mu_sic.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e6589599",
   "metadata": {},
   "source": [
    "# Multi-scale Intensity Characterisation ($\\mu_s$IC)\n",
    "\n",
    "Denoising, feature identification, error estimation, and signal/background/noise separation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a70c4a8-dcfe-4853-ba56-42fc084a579c",
   "metadata": {},
   "source": [
    "# 1. Init"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "621e702e-c938-46ff-b867-f010626596b3",
   "metadata": {},
   "source": [
    "## System setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f82e57f1-d121-4973-a6ee-b3eb5c99145e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib ipympl\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib import colors\n",
    "from matplotlib.ticker import AutoMinorLocator\n",
    "\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display\n",
    "\n",
    "import numpy as np\n",
    "from time import time\n",
    "from scipy import ndimage, special\n",
    "\n",
    "from astropy import stats\n",
    "#import pywt\n",
    "#import pynkowski as mf   # For Minkowski Functionals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb81320d-6c10-48b3-a7f5-62e91befa390",
   "metadata": {},
   "outputs": [],
   "source": [
    "import importlib  # will not be needed in production\n",
    "import scripts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f99496b-23b9-4e19-8650-5fbdff57ba00",
   "metadata": {},
   "source": [
    "## Function definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76b6fae7-db30-4a2e-8506-9aba14bd9492",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "def uniform_filter(cumulant, h):\n",
    "    '''Uniform filter, based on cumulant'''\n",
    "    scale = 2*h + 1\n",
    "    if scale > cumulant.size:\n",
    "        print(f'ERROR: The requested scale, 2*{h} = {scale}, is larger than size = {cumulant.size}')\n",
    "        return cumulant*np.nan\n",
    "    \n",
    "    f = np.empty_like(cumulant)\n",
    "    f[:h+1] = cumulant[h:scale] / (np.arange(h+1) + 1)\n",
    "    f[h+1:f.size-h] = (cumulant[scale:] - cumulant[:-scale]) / scale\n",
    "    f[f.size-h:] = (cumulant[-1] - cumulant[-scale:-h-1]) / (h - np.arange(h))\n",
    "    return f\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a57c4552-1697-4f2c-8bdc-5bb7efa4d4d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def nanfilter(cumulant, cumulant_n, h):\n",
    "    '''Uniform filter, taking NaN into account (faster)'''\n",
    "    scale = 2*h + 1\n",
    "    if scale > data.size:\n",
    "        print(f'WARNING: The requested scale, 2*{h} + 1 = {scale}, is larger than size = {data.size}')\n",
    "        return np.full_like(data, np.nanmean(data))\n",
    "    \n",
    "    f = np.empty_like(cumulant)\n",
    "    f[:h+1] = np.where(cumulant_n[h:scale] > 0, cumulant[h:scale] / cumulant_n[h:scale], 0)\n",
    "    f[h+1:f.size-h] = np.where(cumulant_n[scale:] > cumulant_n[:-scale],\n",
    "                               (cumulant[scale:] - cumulant[:-scale]) / (cumulant_n[scale:] - cumulant_n[:-scale]), 0)\n",
    "    f[f.size-h:] = np.where(cumulant_n[-1] > cumulant_n[-scale:-h-1],\n",
    "                            (cumulant[-1] - cumulant[-scale:-h-1]) / (cumulant_n[-1] - cumulant_n[-scale:-h-1]), 0)\n",
    "    return f\n",
    "\n",
    "\n",
    "def nan_uniform_filter(x, h):\n",
    "    return nanfilter(np.nancumsum(x), np.nancumsum(np.isfinite(x)), h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cef89c0-bef5-4d6d-8cd8-25c2ef8dfaeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mexican_top_hat(x, h):\n",
    "    cumulant_x = np.nancumsum(x)\n",
    "    cumulant_n = np.nancumsum(np.isfinite(x))\n",
    "    #return np.sqrt(1.5) * (nanfilter(cumulant_x, cumulant_n, h) - nanfilter(cumulant_x, cumulant_n, 3*h + 1))\n",
    "    return nanfilter(cumulant_x, cumulant_n, h) - nanfilter(cumulant_x, cumulant_n, 3*h + 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26d0d920-50a8-48c8-850a-85d8e31391a6",
   "metadata": {},
   "source": [
    "## Read data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9152318f-7712-418e-9c99-47f6802039ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "importlib.reload(scripts.read_data)\n",
    "object_name, data, beam_FWHM_pix, true_signal = scripts.read_data.run(11, (0, 0, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7cf1529-c801-4398-a7a5-8302c826d697",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.zeros(666)\n",
    "sigma = 1\n",
    "#data[142:169] = 1\n",
    "data[42:84] = 1\n",
    "data[142:184] = -1\n",
    "data[242:284] = 1\n",
    "true_signal = data.copy()\n",
    "data += np.random.normal(0, sigma, data.size)\n",
    "'''\n",
    "'''\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e77f3ac7-a762-4517-9481-9b414c38a292",
   "metadata": {},
   "source": [
    "# 2. Feature extraction / variance estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67814b02-ccb0-4bb7-954c-997207b6abcd",
   "metadata": {},
   "source": [
    "Actually, I'm not sure whether this might be an overkill.\n",
    "Unless it is able to correctly account for correlated noise!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb9af4a7-a5ed-41eb-b1d6-29f33521f530",
   "metadata": {},
   "outputs": [],
   "source": [
    "cumulative_n = np.nancumsum(np.isfinite(data))\n",
    "\n",
    "s = [1]\n",
    "while s[-1] < data.size/3:\n",
    "    s.append(3 * s[-1])\n",
    "s = np.array(s)\n",
    "h_list = (s - 1) // 2\n",
    "n_h = h_list.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f969876-2500-4783-94fd-fb243434e5a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_features(x):\n",
    "    '''Multi-scale feature extraction'''\n",
    "    \n",
    "    cumulative_data = np.nancumsum(x)\n",
    "    cumulative_data2 = np.nancumsum(x**2)\n",
    "    cumulative_n = np.nancumsum(np.isfinite(x))\n",
    "\n",
    "    mth = np.empty(h_list.shape + x.shape)\n",
    "    mth_var = np.empty_like(mth)\n",
    "    \n",
    "    for i in range(h_list.size):\n",
    "        mth[i] = nanfilter(cumulative_data, cumulative_n, h_list[i])\n",
    "        mth_var[i] = nanfilter(cumulative_data2, cumulative_n, h_list[i])\n",
    "    mth_var -= mth**2\n",
    "    mth_var[1:, :] /= 2*h_list[1:, np.newaxis]\n",
    "\n",
    "    #for i in range(h_list.size-1):\n",
    "    #    mth[i] -= mth[i+1]\n",
    "    #    mth_var[i] += mth_var[i+1]\n",
    "    # WARNING: race condition?\n",
    "    mth[:-1] -= mth[1:]\n",
    "    mth_var[:-1] += mth_var[1:]\n",
    "\n",
    "    mth[-1] -= np.nanmean(x)\n",
    "    mth_var[0] += np.sum(mth_var, axis=0)\n",
    "    \n",
    "    return mth, mth_var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3335d398-7ebc-4acf-a27f-90b10653060f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only used to compute data variance\n",
    "mu_s_data, var_s_data = get_features(data)\n",
    "data_var = np.sum(var_s_data, axis=0)\n",
    "print(np.sqrt(np.nanpercentile(data_var, [16, 50, 84])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cab8d63c-a31a-4753-be7d-f2cac194699e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#if true_signal is not None:\n",
    "#    mu_s_true, var_s_true = get_features(true_signal)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c2236e4-1081-4e53-9ec8-0100a9de2fe3",
   "metadata": {},
   "source": [
    "# 3. Reconstruction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "766132eb-b883-4335-b2f0-5b7739e3839d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mu_s = np.zeros(h_list.shape + data.shape)\n",
    "mu_var = np.zeros_like(mu_s)\n",
    "#mu_estimate = np.zeros_like(mu_s)\n",
    "data_estimate = np.zeros_like(data)\n",
    "estimate_var = np.zeros_like(data)\n",
    "\n",
    "iteration = 0\n",
    "max_iter = 30\n",
    "rms = np.inf\n",
    "evidence = 0\n",
    "while rms > 1 and iteration < max_iter:\n",
    "    iteration += 1\n",
    "    old_rms = rms\n",
    "    old_evidence = evidence\n",
    "\n",
    "    residual = data - data_estimate\n",
    "    \n",
    "    delta_s = np.empty_like(mu_s)\n",
    "    delta = np.zeros_like(data)\n",
    "    for i, h in enumerate(h_list):\n",
    "        estimate = nan_uniform_filter(residual, h+1)\n",
    "        estimate *= np.nansum(residual * estimate) / np.sum(estimate**2)\n",
    "        #estimate *= np.sqrt(2*h+1) #* np.exp(-.5*np.nansum((residual-estimate)**2/data_var))\n",
    "        delta_s[i] = estimate\n",
    "        delta += delta_s[i]\n",
    "    mu_s += delta_s * np.nansum(residual * delta) / np.sum(delta**2)\n",
    "    \n",
    "    data_estimate = np.zeros_like(data)\n",
    "    for i, h in enumerate(h_list):\n",
    "        data_estimate += nan_uniform_filter(mu_s[i], h+1)\n",
    "\n",
    "    rms = np.sqrt(np.nanmean((data-data_estimate)**2/data_var))\n",
    "    evidence = np.exp(-.5 * np.nanmean((data - data_estimate)**2 / estimate_var)) / np.nanmean(estimate_var)\n",
    "    print(iteration, rms, np.nanmean(residual**2), np.nanmean(estimate_var), np.nanmean(data_var), evidence)\n",
    "    if evidence < old_evidence:        \n",
    "        for i, h in enumerate(h_list):\n",
    "            mu_s[i] -= delta[i]\n",
    "        break\n",
    "\n",
    "#w1 = nan_uniform_filter(np.exp(-.5*(mth - mth_smooth)**2/mth_var), h)\n",
    "#w2 = nan_uniform_filter(np.exp(-.5*(mth - estimate)**2/mth_var), h)\n",
    "#final = (w1*mth_smooth + w2*estimate) / (w1 + w2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0531c1d2-d83d-4bfc-85ba-a4314413256d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#mu_s_est, var_s_est = get_features(data_estimate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7930188e-e38c-41d3-85f7-13919bda08bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "signal_s = np.empty_like(mu_s)\n",
    "significance_s = np.empty_like(mu_s)\n",
    "for i, h in enumerate(h_list):\n",
    "    signal_s[i] = nan_uniform_filter(mu_s[i], h+1)\n",
    "    significance_s[i] = np.sqrt(nan_uniform_filter(mu_s[i]**2, h+1))    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1bba4a3-fdd9-492e-971c-2a67f65cbcf7",
   "metadata": {},
   "source": [
    "# 4. Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b77de12-61b8-4c23-a847-fe9f1d47a843",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = scripts.utils.new_figure(object_name, nrows=3, figsize=(12, 8))\n",
    "\n",
    "scripts.utils.colour_map(axes[0, 0], '$\\mu_s$', mu_s)\n",
    "scripts.utils.colour_map(axes[1, 0], '$S_s$', signal_s)\n",
    "scripts.utils.colour_map(axes[2, 0], '$\\Sigma_s$', significance_s)\n",
    "\n",
    "for ax in axes.ravel():\n",
    "    ax.contour(significance_s, levels=np.nanmax(significance_s) / 2**(15-np.arange(16)), colors='k')\n",
    "\n",
    "'''\n",
    "def plot_significance(ax, label, mu_s, **kwargs):\n",
    "    estimate = np.empty_like(mu_s)\n",
    "    significance = np.empty_like(mu_s)\n",
    "    for i, h in enumerate(h_list):\n",
    "        estimate[i] = nan_uniform_filter(mu_s[i], h+1)\n",
    "        significance[i] = np.sqrt(nan_uniform_filter(mu_s[i]**2, h+1))\n",
    "    #scripts.utils.colour_map(ax, label, significance, **kwargs)\n",
    "    scripts.utils.colour_map(ax, label, estimate, **kwargs)\n",
    "    ax.contour(significance, levels=np.nanmax(significance) / 2**(15-np.arange(16)), colors='k')\n",
    "\n",
    "p0, p50, p100 = np.nanpercentile(mu_s_est**2, [0, 50, 100])\n",
    "norm = colors.SymLogNorm(vmin=-p100, linthresh=p50, vmax=p100)\n",
    "norm = None\n",
    "cmap = 'rainbow_r'\n",
    "plot_significance(axes[0, 0], '$\\mu_s$ data', mu_s_data, norm=norm, cmap=cmap)\n",
    "plot_significance(axes[1, 0], '$\\mu_s$ estimate', mu_s, norm=norm, cmap=cmap)\n",
    "#plot_significance(axes[2, 0], '$\\mu_s$ true', mu_s_true, norm=norm, cmap=cmap)\n",
    "'''\n",
    "\n",
    "#axes[0, 0].set_xlim(950, 1250)\n",
    "#axes[0, 0].set_xlim(3350, 3850)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11ffbb05-6e99-4c0b-8458-5ee1c30ef3b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = scripts.utils.new_figure('filter', nrows=n_h, sharey=True)  #, figsize=(12, 8))\n",
    "\n",
    "for i, h in enumerate(h_list):\n",
    "    ax = axes[i, 0]\n",
    "    \n",
    "    if true_signal is not None:\n",
    "        ax.plot(true_signal, 'k:', label='true signal')\n",
    "        #x = 1.5 * np.sqrt(np.clip(mexican_top_hat(mu_s_true[i]**2, h+1), 0, np.inf))\n",
    "        #ax.plot(x, 'k--', label='true intensity')\n",
    "\n",
    "    x = data\n",
    "    noise = np.sqrt(data_var)\n",
    "    ax.plot(x, 'r-', alpha=.2, label='data')\n",
    "    ax.fill_between(np.arange(data.size), x-noise, x+noise, color='r', alpha=.1)\n",
    "\n",
    "    x = data_estimate\n",
    "    noise = np.sqrt(estimate_var)\n",
    "    ax.plot(x, 'k--', alpha=.5, label=f'estimate')\n",
    "    ax.fill_between(np.arange(data.size), x-noise, x+noise, color='k', alpha=.1)\n",
    "\n",
    "    #x = mu_s[i]\n",
    "    #x = nan_uniform_filter(mu_s[i], h+1)\n",
    "    x = signal_s[i]\n",
    "    noise = np.sqrt(data_var)\n",
    "    ax.plot(x, 'b-', alpha=.5, label=f'h={h}')\n",
    "    ax.fill_between(np.arange(data.size), x-noise, x+noise, color='b', alpha=.1)\n",
    "\n",
    "    x = np.sum(signal_s[:i+1, :], axis=0)\n",
    "    ax.plot(x, 'r-', alpha=.5, label=f'h<={h}')\n",
    "\n",
    "    if i < n_h-1:\n",
    "        x = np.sum(signal_s[i+1:, :], axis=0)\n",
    "        ax.plot(x, 'r:', alpha=.5, label=f'h>{h}')\n",
    "\n",
    "    '''\n",
    "    x = 1.5 * np.sqrt(np.clip(mexican_top_hat(mu_s_data[i]**2, h+1), 0, np.inf))\n",
    "    noise = np.sqrt(nan_uniform_filter(var_s_data[i], h))\n",
    "    ax.plot(x, 'k-', label='feature intensity')\n",
    "    ax.fill_between(np.arange(data.size), x-noise, x+noise, color='g', alpha=.1)\n",
    "    '''\n",
    "\n",
    "    #ax.set_xlim(950, 1250); ax.set_ylim(-10, 250)\n",
    "    #ax.set_xlim(3350, 3850); ax.set_ylim(-50, 450)\n",
    "    ax.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34d6543c-be16-4b48-9e64-3fba06bd5607",
   "metadata": {
    "tags": []
   },
   "source": [
    "# --- STOP ---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6b7a330-47b3-41e5-9f08-c72f013eb3fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise -1"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  },
  "toc-autonumbering": false,
  "toc-showcode": false,
  "toc-showmarkdowntxt": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}