diff --git a/tests/SN_tests.ipynb b/tests/SN_tests.ipynb new file mode 100644 index 0000000..9789d1a --- /dev/null +++ b/tests/SN_tests.ipynb @@ -0,0 +1,159 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 22, + "id": "de9cdd42-f9bb-4dd6-80ba-d63d0dbc3ca2", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "import matplotlib.pyplot as plt\n", + "import linecache\n", + "\n", + "from astropy.io import fits" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "id": "72780ead-06f0-4fbb-b12e-602b203286a0", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "0.04197595 0.17233 0.0350472\n", + "-0.0372089 1.2312531999999994\n", + "-0.0186623 0.07995692199999997\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "<>:5: SyntaxWarning: invalid escape sequence '\\s'\n", + "<>:5: SyntaxWarning: invalid escape sequence '\\s'\n", + "/tmp/ipykernel_4414/3947829372.py:5: SyntaxWarning: invalid escape sequence '\\s'\n", + " df = pd.read_csv(fname, sep=\"\\s+\", skipinitialspace=True, skiprows=7, names=columns)\n" + ] + }, + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 30, + "metadata": {}, + "output_type": "execute_result" + }, + { + "data": { + "image/png": "iVBORw0KGgoAAAANSUhEUgAAAh8AAAGdCAYAAACyzRGfAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAp+0lEQVR4nO3df3RU9Z3/8deQH5MEQoL8yC8DDZCAiAiFlfLDDbUQF9R2l921NoroWlYFlUjdAFIk0pLwo6VxvyAtrMeNx0U826LrVkRiPcQfKQuy5IgBRZYgEchG2JCEJGSAfL5/5GSOYwIyycxnZpLn45x76tz7mc993485vS/v3Hs/DmOMEQAAgCW9Al0AAADoWQgfAADAKsIHAACwivABAACsInwAAACrCB8AAMAqwgcAALCK8AEAAKwKD3QB39TS0qJTp04pNjZWDocj0OUAAIBrYIxRfX29kpOT1avX1a9tBF34OHXqlFJTUwNdBgAA6ITKykpdf/31V20TdOEjNjZWUmvxffv2DXA1AADgWtTV1Sk1NdV9Hr+aoAsfbT+19O3bl/ABAECIuZZbJrjhFAAAWOV1+Dh58qTuu+8+9e/fXzExMRo7dqz279/v3m6MUV5enpKTkxUdHa1p06apvLzcp0UDAIDQ5VX4qKmp0ZQpUxQREaG33npLhw4d0q9//WvFx8e726xdu1br16/Xhg0btG/fPiUmJmrGjBmqr6/3de0AACAEOYwx5lobL1myRB9++KHef//9DrcbY5ScnKycnBwtXrxYktTc3KyEhAStWbNGDz/88Lfuo66uTnFxcaqtreWeDwBApxljdOnSJV2+fDnQpXQbYWFhCg8P7/C+Dm/O317dcPrGG2/o9ttv19///d+rpKREKSkpmj9/vubNmydJqqioUFVVlbKystzfcTqdyszMVGlpaYfho7m5Wc3NzR7FAwDQFS6XS6dPn1ZjY2OgS+l2YmJilJSUpMjIyE734VX4OHbsmDZt2qRFixbp6aef1t69e/XEE0/I6XTq/vvvV1VVlSQpISHB43sJCQn64osvOuyzoKBAzz77bCfLBwDAU0tLiyoqKhQWFqbk5GRFRkby0kofMMbI5XLpq6++UkVFhdLT07/1ZWJX4lX4aGlp0YQJE5Sfny9JGjdunMrLy7Vp0ybdf//97nbf/JdsjLniv/ilS5dq0aJF7s9tzwkDANAZLpdLLS0tSk1NVUxMTKDL6Vaio6MVERGhL774Qi6XS1FRUZ3qx6vIkpSUpFGjRnmsu+GGG3TixAlJUmJioiS5r4C0qa6ubnc1pI3T6XS/04N3ewAAfKWz/1WOq/PFuHrVw5QpU/TZZ595rDty5IiGDBkiSUpLS1NiYqKKi4vd210ul0pKSjR58uQuFwsAAEKfVz+7PPnkk5o8ebLy8/N19913a+/evdq8ebM2b94sqfXnlpycHOXn5ys9PV3p6enKz89XTEyMsrOz/XIAAABci5PnmlTT4LK2v369I5USH21tf6HEq/DxF3/xF3rttde0dOlSrVy5UmlpaSosLNS9997rbpObm6umpibNnz9fNTU1mjhxonbt2nVN73oHAMAfTp5r0vRfl6jpor3HbqMjwvTOzzKDNoAcP35caWlpOnDggMaOHWt1317P7XLnnXfqzjvvvOJ2h8OhvLw85eXldaUuAAB8pqbBpaaLl1X447EaPqiP3/d3tPq8cl4tU02DK2jDRyAF3cRyAAD4y/BBfTQ6JS7QZfR43AqMkHDyXJM+OVnrk+XkuaZAHw4AtDNt2jQ9/vjjysnJUb9+/ZSQkKDNmzeroaFBDz74oGJjYzVs2DC99dZb7u+UlJTolltukdPpVFJSkpYsWaJLly65t7e0tGjNmjUaPny4nE6nBg8erFWrVnW4/5aWFs2bN08ZGRlXfDeXr3DlA0HP17/VBvvvsAB6rqKiIuXm5mrv3r169dVX9eijj+r111/X3/zN3+jpp5/Wb37zG82ZM0cnTpxQTU2NZs2apQceeEAvvfSSPv30U82bN09RUVHuWx+WLl2qLVu26De/+Y2mTp2q06dP69NPP223X5fLpezsbP3P//yPPvjgAw0aNMivx0n4QNDz5W+1/A4LIJjdfPPN+vnPfy6pNTisXr1aAwYMcE9j8swzz2jTpk36+OOP9Z//+Z9KTU3Vhg0b5HA4NHLkSJ06dUqLFy/WM888o4aGBj333HPasGGD5s6dK0kaNmyYpk6d6rHP8+fP64477lBTU5N2796tuDj//yxF+EDI4LdaAN3dmDFj3P8cFham/v3766abbnKva3thZ3V1tQ4fPqxJkyZ5vEF8ypQpOn/+vL788ktVVVWpublZP/jBD666z5/85Ce6/vrr9ac//cnaG2G55wMAgCARERHh8dnhcHisawsaLS0tHU5d0jZRvcPhUHT0tV3dnTVrlj7++GPt2bOnK6V7hfABAEAIGjVqlEpLS92BQ5JKS0sVGxurlJQUpaenKzo6Wn/605+u2s+jjz6q1atX64c//KFKSkr8XbYkfnYBAPQgR6vPd5v9zJ8/X4WFhXr88cf12GOP6bPPPtOKFSu0aNEi9erVS1FRUVq8eLFyc3MVGRmpKVOm6KuvvlJ5ebkeeughj74ef/xxXb58WXfeeafeeuutdveF+BrhAwDQ7fXrHanoiDDlvFpmbZ/REWHq1zvSb/2npKRox44d+qd/+ifdfPPNuu666/TQQw+5b1iVpOXLlys8PFzPPPOMTp06paSkJD3yyCMd9peTk6OWlhbNmjVLO3fu9OucbA7z9es1QaCurk5xcXGqra1lhltIkj45Was7/98H+uPjU7t8w6kv+wIQnC5cuKCKigqlpaV5TPnO3C6+caXx9eb8zZUPAECPkBIf3S3DQCjihlMAAGAV4QMAAFhF+AAAAFYRPgAAgFXccIoeyRfP4HfXO9kBwN8IH+hRfPmsP7PjAkDnED7Qo6TER+udn2V2+Vl/ZscFgM4jfKDH4Vl/oIc6Vyk1nrW3v5j+Unyqvf2FEMIHAKD7O1cpbbxFuthob58RMdKCvV4FkGnTpmns2LEqLCz0X11BgPABAOj+Gs+2Bo/ZW6QBGf7f35kj0vZ5rfv14dUPY4wuX76s8PDQPn2HdvUAAHhjQIaUPDbQVXTogQceUElJiUpKSvTcc89Jkl588UU9+OCD2rlzp5YtW6aPP/5Yb7/9toqKinTu3Dm9/vrr7u/n5OSorKxMu3fvltQaVNatW6ff/va3On36tDIyMrR8+XL93d/9XQCOzhPhAwCAIPDcc8/pyJEjGj16tFauXClJKi8vlyTl5ubqV7/6lYYOHar4+Phr6u/nP/+5tm/frk2bNik9PV3vvfee7rvvPg0cOFCZmZn+OoxrQvgAACAIxMXFKTIyUjExMUpMTJQkffrpp5KklStXasaMGdfcV0NDg9avX693331XkyZNkiQNHTpUH3zwgX73u98RPgAAwNVNmDDBq/aHDh3ShQsX2gUWl8ulcePG+bK0TiF8AAAQ5Hr37u3xuVevXjLGeKy7ePGi+59bWlokSW+++aZSUlI82jmdTj9Vee0IHwAABInIyEhdvnz5W9sNHDhQn3zyice6srIyRURESJJGjRolp9OpEydOBPwnlo4QPgAACBLf+c539F//9V86fvy4+vTp476C8U233Xab1q1bp5deekmTJk3Syy+/rE8++cT9k0psbKyeeuopPfnkk2ppadHUqVNVV1en0tJS9enTR3PnzrV5WO0QPgAAPceZI0G9n6eeekpz587VqFGj1NTUpBdffLHDdrfffruWL1+u3NxcXbhwQf/wD/+g+++/XwcPHnS3+cUvfqFBgwapoKBAx44dU3x8vL773e/q6aef7lRtvkT4ALrAF7PjSsyQC/hdTP/WN45un2dvnxExrfv1QkZGhv785z97rHvggQc6bPvss8/q2WefvWJfDodDTzzxhJ544gmvarCB8AF0gi9nx5WYIRfwu/jU1ledM7dLUCB8AJ3gq9lxJWbIBayJTyUMBAnCB9BJzI4LAJ3TK9AFAACAnoXwAQAArCJ8AAC6pW++ARS+4YtxJXwAALqVtrd8NjY2BriS7qltXNvGuTO44RQA0K2EhYUpPj5e1dXVkqSYmBg5HI4AVxX6jDFqbGxUdXW14uPjFRYW1um+CB8AgG6nbUr6tgAC34mPj3ePb2cRPgAA3Y7D4VBSUpIGDRrkMdsruiYiIqJLVzzaED4AAN1WWFiYT06W8C1uOAUAAFYRPgAAgFWEDwAAYBXhAwAAWEX4AAAAVhE+AACAVYQPAABgFeEDAABY5VX4yMvLk8Ph8Fi+/opVY4zy8vKUnJys6OhoTZs2TeXl5T4vGgAAhC6vr3zceOONOn36tHs5ePCge9vatWu1fv16bdiwQfv27VNiYqJmzJih+vp6nxYNAABCl9fhIzw8XImJie5l4MCBklqvehQWFmrZsmWaPXu2Ro8eraKiIjU2Nmrr1q0+LxwAAIQmr8PH559/ruTkZKWlpemee+7RsWPHJEkVFRWqqqpSVlaWu63T6VRmZqZKS0uv2F9zc7Pq6uo8FgAA0H15FT4mTpyol156SW+//ba2bNmiqqoqTZ48WWfPnlVVVZUkKSEhweM7CQkJ7m0dKSgoUFxcnHtJTU3txGEAAIBQ4VX4mDlzpv72b/9WN910k6ZPn64333xTklRUVORu43A4PL5jjGm37uuWLl2q2tpa91JZWelNSQAAIMR06VHb3r1766abbtLnn3/ufurlm1c5qqur210N+Tqn06m+fft6LAAAoPvqUvhobm7W4cOHlZSUpLS0NCUmJqq4uNi93eVyqaSkRJMnT+5yoQAAoHsI96bxU089pbvuukuDBw9WdXW1fvnLX6qurk5z586Vw+FQTk6O8vPzlZ6ervT0dOXn5ysmJkbZ2dn+qh8AAIQYr8LHl19+qZ/85Cc6c+aMBg4cqO9973vas2ePhgwZIknKzc1VU1OT5s+fr5qaGk2cOFG7du1SbGysX4oHAAChx6vwsW3btqtudzgcysvLU15eXldqAgAA3RhzuwAAAKsIHwAAwCrCBwAAsIrwAQAArCJ8AAAAqwgfAADAKsIHAACwivABAACsInwAAACrCB8AAMAqwgcAALCK8AEAAKwifAAAAKsIHwAAwCrCBwAAsIrwAQAArCJ8AAAAqwgfAADAKsIHAACwivABAACsInwAAACrCB8AAMAqwgcAALCK8AEAAKwifAAAAKsIHwAAwCrCBwAAsIrwAQAArCJ8AAAAqwgfAADAKsIHAACwivABAACsInwAAACrCB8AAMAqwgcAALCK8AEAAKwifAAAAKsIHwAAwCrCBwAAsIrwAQAArCJ8AAAAqwgfAADAKsIHAACwivABAACsInwAAACrCB8AAMAqwgcAALCK8AEAAKwifAAAAKu6FD4KCgrkcDiUk5PjXmeMUV5enpKTkxUdHa1p06apvLy8q3UCAIBuotPhY9++fdq8ebPGjBnjsX7t2rVav369NmzYoH379ikxMVEzZsxQfX19l4sFAAChr1Ph4/z587r33nu1ZcsW9evXz73eGKPCwkItW7ZMs2fP1ujRo1VUVKTGxkZt3brVZ0UDAIDQ1anwsWDBAt1xxx2aPn26x/qKigpVVVUpKyvLvc7pdCozM1OlpaVdqxQAAHQL4d5+Ydu2bfrv//5v7du3r922qqoqSVJCQoLH+oSEBH3xxRcd9tfc3Kzm5mb357q6Om9LAgAAIcSrKx+VlZVauHChXn75ZUVFRV2xncPh8PhsjGm3rk1BQYHi4uLcS2pqqjclAQCAEONV+Ni/f7+qq6s1fvx4hYeHKzw8XCUlJfrnf/5nhYeHu694tF0BaVNdXd3uakibpUuXqra21r1UVlZ28lAAAEAo8Opnlx/84Ac6ePCgx7oHH3xQI0eO1OLFizV06FAlJiaquLhY48aNkyS5XC6VlJRozZo1HfbpdDrldDo7WT4AAAg1XoWP2NhYjR492mNd79691b9/f/f6nJwc5efnKz09Xenp6crPz1dMTIyys7N9VzUAAAhZXt9w+m1yc3PV1NSk+fPnq6amRhMnTtSuXbsUGxvr610BAIAQ1OXwsXv3bo/PDodDeXl5ysvL62rXAACgG2JuFwAAYBXhAwAAWEX4AAAAVhE+AACAVYQPAABgFeEDAABYRfgAAABWET4AAIBVhA8AAGAV4QMAAFhF+AAAAFYRPgAAgFWEDwAAYBXhAwAAWEX4AAAAVhE+AACAVYQPAABgFeEDAABYRfgAAABWET4AAIBVhA8AAGAV4QMAAFhF+AAAAFYRPgAAgFWEDwAAYBXhAwAAWEX4AAAAVhE+AACAVYQPAABgFeEDAABYRfgAAABWET4AAIBVhA8AAGAV4QMAAFhF+AAAAFYRPgAAgFWEDwAAYBXhAwAAWEX4AAAAVhE+AACAVYQPAABgFeEDAABYRfgAAABWET4AAIBVhA8AAGAV4QMAAFhF+AAAAFYRPgAAgFWEDwAAYBXhAwAAWOVV+Ni0aZPGjBmjvn37qm/fvpo0aZLeeust93ZjjPLy8pScnKzo6GhNmzZN5eXlPi8aAACELq/Cx/XXX6/Vq1fro48+0kcffaTbbrtNP/rRj9wBY+3atVq/fr02bNigffv2KTExUTNmzFB9fb1figcAAKHHq/Bx1113adasWcrIyFBGRoZWrVqlPn36aM+ePTLGqLCwUMuWLdPs2bM1evRoFRUVqbGxUVu3bvVX/QAAIMR0+p6Py5cva9u2bWpoaNCkSZNUUVGhqqoqZWVluds4nU5lZmaqtLT0iv00Nzerrq7OYwEAAN2X1+Hj4MGD6tOnj5xOpx555BG99tprGjVqlKqqqiRJCQkJHu0TEhLc2zpSUFCguLg495KamuptSQAAIIR4HT5GjBihsrIy7dmzR48++qjmzp2rQ4cOubc7HA6P9saYduu+bunSpaqtrXUvlZWV3pYEAABCSLi3X4iMjNTw4cMlSRMmTNC+ffv03HPPafHixZKkqqoqJSUludtXV1e3uxrydU6nU06n09syAABAiOryez6MMWpublZaWpoSExNVXFzs3uZyuVRSUqLJkyd3dTcAAKCb8OrKx9NPP62ZM2cqNTVV9fX12rZtm3bv3q2dO3fK4XAoJydH+fn5Sk9PV3p6uvLz8xUTE6Ps7Gx/1Q8AAEKMV+Hjf//3fzVnzhydPn1acXFxGjNmjHbu3KkZM2ZIknJzc9XU1KT58+erpqZGEydO1K5duxQbG+uX4gEAQOjxKny88MILV93ucDiUl5envLy8rtQEAAC6MeZ2AQAAVhE+AACAVYQPAABgFeEDAABYRfgAAABWET4AAIBVhA8AAGAV4QMAAFhF+AAAAFYRPgAAgFWEDwAAYBXhAwAAWEX4AAAAVnk1qy1wzc5VSo1nfdJV1JnzutFRoagzcZKjj0/67FBMfyk+1X/9AwAkET7gD+cqpY23SBcbfdLdcElvOiW95pPuriwiRlqwlwACAH5G+IDvNZ5tDR6zt0gDMrrc3dGvzmvhtjI9d89YDR/opysfZ45I2+e11k74AAC/InzAfwZkSMlju9zNBVOrclOrCwNukpLjul4XACCguOEUAABYRfgAAABWET4AAIBVhA8AAGAV4QMAAFhF+AAAAFYRPgAAgFWEDwAAYBXhAwAAWEX4AAAAVhE+AACAVYQPAABgFeEDAABYRfgAAABWET4AAIBV4YEuANfgXKXUeDbQVVy7M0cCXQEAIIgRPoLduUpp4y3SxcZAV+KdiBgppn+gqwAABCHCR7BrPNsaPGZvkQZkBLqaaxfTX4pPDXQVAIAgRPgIFQMypOSxga4CAIAu44ZTAABgFeEDAABYRfgAAABWET4AAIBVhA8AAGAV4QMAAFhF+AAAAFYRPgAAgFWEDwAAYBXhAwAAWMXr1eFXJ881qabB1aU+jlaf91E1AIBgQPiA35w816Tpvy5R08XLXe4rOiJM/XpH+qAqAECgET7gNzUNLjVdvKzCH4/V8EF9utRXv96RSomP9lFlAIBA8ip8FBQUaPv27fr0008VHR2tyZMna82aNRoxYoS7jTFGzz77rDZv3qyamhpNnDhRGzdu1I033ujz4hEahg/qo9EpcYEuAwAQJLy64bSkpEQLFizQnj17VFxcrEuXLikrK0sNDQ3uNmvXrtX69eu1YcMG7du3T4mJiZoxY4bq6+t9XjwAAAg9Xl352Llzp8fnF198UYMGDdL+/fv1l3/5lzLGqLCwUMuWLdPs2bMlSUVFRUpISNDWrVv18MMP+65yAAAQkrr0qG1tba0k6brrrpMkVVRUqKqqSllZWe42TqdTmZmZKi0t7bCP5uZm1dXVeSwAAKD76nT4MMZo0aJFmjp1qkaPHi1JqqqqkiQlJCR4tE1ISHBv+6aCggLFxcW5l9TU1M6WBAAAQkCnw8djjz2mjz/+WK+88kq7bQ6Hw+OzMabdujZLly5VbW2te6msrOxsSQAAIAR06lHbxx9/XG+88Ybee+89XX/99e71iYmJklqvgCQlJbnXV1dXt7sa0sbpdMrpdHamDAAAEIK8uvJhjNFjjz2m7du3691331VaWprH9rS0NCUmJqq4uNi9zuVyqaSkRJMnT/ZNxQAAIKR5deVjwYIF2rp1q/7jP/5DsbGx7vs44uLiFB0dLYfDoZycHOXn5ys9PV3p6enKz89XTEyMsrOz/XIAAAAgtHgVPjZt2iRJmjZtmsf6F198UQ888IAkKTc3V01NTZo/f777JWO7du1SbGysTwoGAAChzavwYYz51jYOh0N5eXnKy8vrbE0AAKAb69J7PgAAALxF+AAAAFYRPgAAgFWEDwAAYBXhAwAAWEX4AAAAVhE+AACAVYQPAABgFeEDAABYRfgAAABWET4AAIBVhA8AAGAV4QMAAFhF+AAAAFYRPgAAgFWEDwAAYBXhAwAAWEX4AAAAVhE+AACAVYQPAABgFeEDAABYRfgAAABWET4AAIBVhA8AAGAV4QMAAFhF+AAAAFYRPgAAgFWEDwAAYFV4oAsA0Opo9fku99Gvd6RS4qN9UA0A+A/hAwiwfr0jFR0RppxXy7rcV3REmN75WSYBBEBQI3wAAZYSH613fpapmgZXl/o5Wn1eOa+WqabBRfgAENQIH0AQSImPJjAA6DEIH0AoO1cpNZ6VJEWdOa8bHRWKOhMnOfoEuLCriOkvxacGugoAAUT4AELVuUpp4y3SxUZJ0nBJbzolvRbQqr5dRIy0YC8BBOjBCB9AqGo82xo8Zm+RBmTo6FfntXBbmZ67Z6yGDwzSKx9njkjb57XWTvgAeizCBxDqBmRIyWN1wdSq3NTqwoCbpOS4QFcFAFfES8YAAIBVhA8AAGAV4QMAAFhF+AAAAFYRPgAAgFWEDwAAYBWP2gJfd+ZIoCu4dqFUKwB8DeEDkFpf+R0R0/oCrFASEdNaOwCEEMIHILW+bXPBXvc8KSGDeVIAhCDCB9AmPpUTOQBYwA2nAADAKsIHAACwivABAACs8jp8vPfee7rrrruUnJwsh8Oh119/3WO7MUZ5eXlKTk5WdHS0pk2bpvLycl/VCwAAQpzX4aOhoUE333yzNmzY0OH2tWvXav369dqwYYP27dunxMREzZgxQ/X19V0uFgAAhD6vn3aZOXOmZs6c2eE2Y4wKCwu1bNkyzZ49W5JUVFSkhIQEbd26VQ8//HDXqgUAACHPp/d8VFRUqKqqSllZWe51TqdTmZmZKi0t7fA7zc3Nqqur81gAAED35dPwUVVVJUlKSEjwWJ+QkODe9k0FBQWKi4tzL6mpvGcBAIDuzC9PuzgcDo/Pxph269osXbpUtbW17qWystIfJQEAgCDh0zecJiYmSmq9ApKUlOReX11d3e5qSBun0ymn0+nLMgAAQBDzafhIS0tTYmKiiouLNW7cOEmSy+VSSUmJ1qxZ48tdoQMnzzWppsEV6DLcjlafD3QJAIAg5HX4OH/+vI4ePer+XFFRobKyMl133XUaPHiwcnJylJ+fr/T0dKWnpys/P18xMTHKzs72aeHwdPJck6b/ukRNFy8HuhQP0RFh6tc7MtBlAACCiNfh46OPPtL3v/999+dFixZJkubOnat//dd/VW5urpqamjR//nzV1NRo4sSJ2rVrl2JjY31XNdqpaXCp6eJlFf54rIYP6hPoctz69Y5USnx0oMsAAAQRr8PHtGnTZIy54naHw6G8vDzl5eV1pS500vBBfTQ6JS7QZQAAcEXM7QIAAKwifAAAAKsIHwAAwCrCBwAAsIrwAQAArCJ8AAAAqwgfAADAKsIHAACwivABAACsInwAAACrfDqrLYDA88VswszJA8CfCB9AN9Gvd6SiI8KU82pZl/uKjgjTOz/LJIAA8AvCB9BNpMRH652fZaqmwdWlfo5Wn1fOq2WqaXARPgD4BeED6EZS4qMJDACCHuEDAK7FuUqp8Wygq/BOTH8pPjXQVQDtED4A4Nucq5Q23iJdbAx0Jd6JiJEW7CWAIOgQPgDg2zSebQ0es7dIAzICXc21OXNE2j6vtXbCB4JMzwsfoXbp9MyRQFcAoM2ADCl5bKCrAEJezwofoXzpNKZ/oKsAAMAnelb4CMVLpxI3jQEAupWeFT7acOkUAICAYW4XAABgFeEDAABYRfgAAABWET4AAIBVhA8AAGAV4QMAAFjVMx+1BRBYofbm3lCrFwhyhA8A9sT0b31j7/Z5ga7Ee7xpGPAZwgcAe+JTW2dZDaX5ldrwpmHAZwgfAOyKT+UkDvRw3HAKAACsInwAAACrCB8AAMAqwgcAALCK8AEAAKziaRcA6M5C7QVpPNLcIxA+AKA7CtUXukXEtL4LhgDSrRE+OunkuSbVNLgCXYbb0erzgS4BQDAJxRe6nTnSGpYazxI+ujnCRyecPNek6b8uUdPFy4EuxUN0RJj69Y4MdBkAggUvdEOQInx0Qk2DS00XL6vwx2M1fFCfQJfj1q93pFLiowNdBgAAV0X46ILhg/podEpcoMsAACCk8KgtAACwivABAACsInwAAACrCB8AAMAqbjgFAKArzlWG1vtUpIC/SZbwAQBAZ52rlDbeIl1sDHQl3gnwm2QJHwAAdFbj2dbgMXuLNCAj0NVcmyB4k6zfwsfzzz+vdevW6fTp07rxxhtVWFioW2+91V+7AwB0F6E0GV5brQMypOSxAS0llPglfLz66qvKycnR888/rylTpuh3v/udZs6cqUOHDmnw4MH+2CUAINSF8mR4Mf0DXUVI8Uv4WL9+vR566CH99Kc/lSQVFhbq7bff1qZNm1RQUOCPXQIAQl0oToYnBfzmzVDk8/Dhcrm0f/9+LVmyxGN9VlaWSktL27Vvbm5Wc3Oz+3Ntba0kqa6uztelSfXnpWbT+r9d6P98fZ1amht1vr5OdXUOHxYIBB5/3wioXnFSnxCctsIf5yx/8dG58JvaztvGmG9vbHzs5MmTRpL58MMPPdavWrXKZGRktGu/YsUKI4mFhYWFhYWlGyyVlZXfmhX8dsOpw+H5X0zGmHbrJGnp0qVatGiR+3NLS4v+7//+TxERERo8eLAqKyvVt29ff5XZ49TV1Sk1NZVx9THG1T8YV/9hbP2jJ4+rMUb19fVKTk7+1rY+Dx8DBgxQWFiYqqqqPNZXV1crISGhXXun0ymn0+mxLj4+3n35pm/fvj3uX6ANjKt/MK7+wbj6D2PrHz11XOPi4q6pnc9frx4ZGanx48eruLjYY31xcbEmT57s690BAIAQ45efXRYtWqQ5c+ZowoQJmjRpkjZv3qwTJ07okUce8cfuAABACPFL+Pjxj3+ss2fPauXKlTp9+rRGjx6tHTt2aMiQIdfch9Pp1IoVK9r9JIOuYVz9g3H1D8bVfxhb/2Bcr43DmGt5JgYAAMA3fH7PBwAAwNUQPgAAgFWEDwAAYBXhAwAAWGUtfDz//PNKS0tTVFSUxo8fr/fff/+q7UtKSjR+/HhFRUVp6NCh+u1vf9uuzR/+8AeNGjVKTqdTo0aN0muvveav8oOWr8d1y5YtuvXWW9WvXz/169dP06dP1969e/15CEHLH3+zbbZt2yaHw6G//uu/9nHVwc8f43ru3DktWLBASUlJioqK0g033KAdO3b46xCCkj/GtbCwUCNGjFB0dLRSU1P15JNP6sKFC/46hKDkzbiePn1a2dnZGjFihHr16qWcnJwO23Huknw+t0tHtm3bZiIiIsyWLVvMoUOHzMKFC03v3r3NF1980WH7Y8eOmZiYGLNw4UJz6NAhs2XLFhMREWF+//vfu9uUlpaasLAwk5+fbw4fPmzy8/NNeHi42bNnj41DCgr+GNfs7GyzceNGc+DAAXP48GHz4IMPmri4OPPll1/aOqyg4I+xbXP8+HGTkpJibr31VvOjH/3Iz0cSXPwxrs3NzWbChAlm1qxZ5oMPPjDHjx8377//vikrK7N1WAHnj3F9+eWXjdPpNP/2b/9mKioqzNtvv22SkpJMTk6OrcMKOG/HtaKiwjzxxBOmqKjIjB071ixcuLBdG85drayEj1tuucU88sgjHutGjhxplixZ0mH73NxcM3LkSI91Dz/8sPne977n/nz33Xebv/qrv/Joc/vtt5t77rnHR1UHP3+M6zddunTJxMbGmqKioq4XHEL8NbaXLl0yU6ZMMf/yL/9i5s6d2+PChz/GddOmTWbo0KHG5XL5vuAQ4Y9xXbBggbnttts82ixatMhMnTrVR1UHP2/H9esyMzM7DB+cu1r5/WcXl8ul/fv3Kysry2N9VlaWSktLO/zOn//853btb7/9dn300Ue6ePHiVdtcqc/uxl/j+k2NjY26ePGirrvuOt8UHgL8ObYrV67UwIED9dBDD/m+8CDnr3F94403NGnSJC1YsEAJCQkaPXq08vPzdfnyZf8cSJDx17hOnTpV+/fvd//seuzYMe3YsUN33HGHH44i+HRmXK9FTz93tfHbrLZtzpw5o8uXL7ebVC4hIaHd5HNtqqqqOmx/6dIlnTlzRklJSVdsc6U+uxt/jes3LVmyRCkpKZo+fbrvig9y/hrbDz/8UC+88ILKysr8VXpQ89e4Hjt2TO+++67uvfde7dixQ59//rkWLFigS5cu6ZlnnvHb8QQLf43rPffco6+++kpTp06VMUaXLl3So48+qiVLlvjtWIJJZ8b1WvT0c1cbv4ePNg6Hw+OzMabdum9r/8313vbZHfljXNusXbtWr7zyinbv3q2oqCgfVBtafDm29fX1uu+++7RlyxYNGDDA98WGEF//zba0tGjQoEHavHmzwsLCNH78eJ06dUrr1q3rEeGjja/Hdffu3Vq1apWef/55TZw4UUePHtXChQuVlJSk5cuX+7j64OWP8wznLgvhY8CAAQoLC2uX6qqrq9ulvzaJiYkdtg8PD1f//v2v2uZKfXY3/hrXNr/61a+Un5+vd955R2PGjPFt8UHOH2NbXl6u48eP66677nJvb2lpkSSFh4frs88+07Bhw3x8JMHFX3+zSUlJioiIUFhYmLvNDTfcoKqqKrlcLkVGRvr4SIKLv8Z1+fLlmjNnjn76059Kkm666SY1NDToH//xH7Vs2TL16tW939TQmXG9Fj393NXG7389kZGRGj9+vIqLiz3WFxcXa/LkyR1+Z9KkSe3a79q1SxMmTFBERMRV21ypz+7GX+MqSevWrdMvfvEL7dy5UxMmTPB98UHOH2M7cuRIHTx4UGVlZe7lhz/8ob7//e+rrKxMqampfjueYOGvv9kpU6bo6NGj7jAnSUeOHFFSUlK3Dx6S/8a1sbGxXcAICwuTaX1QwYdHEJw6M67Xoqefu9xs3NXa9rjSCy+8YA4dOmRycnJM7969zfHjx40xxixZssTMmTPH3b7tMbAnn3zSHDp0yLzwwgvtHgP78MMPTVhYmFm9erU5fPiwWb16dY97XMkf47pmzRoTGRlpfv/735vTp0+7l/r6euvHF0j+GNtv6olPu/hjXE+cOGH69OljHnvsMfPZZ5+ZP/7xj2bQoEHml7/8pfXjCxR/jOuKFStMbGyseeWVV8yxY8fMrl27zLBhw8zdd99t/fgCxdtxNcaYAwcOmAMHDpjx48eb7Oxsc+DAAVNeXu7ezrmrlZXwYYwxGzduNEOGDDGRkZHmu9/9rikpKXFvmzt3rsnMzPRov3v3bjNu3DgTGRlpvvOd75hNmza16/Pf//3fzYgRI0xERIQZOXKk+cMf/uDvwwg6vh7XIUOGGEntlhUrVlg4muDij7/Zr+uJ4cMY/4xraWmpmThxonE6nWbo0KFm1apV5tKlS/4+lKDi63G9ePGiycvLM8OGDTNRUVEmNTXVzJ8/39TU1Fg4muDh7bh29P+fQ4YM8WjDucsYhzE94PoZAAAIGt37jiEAABB0CB8AAMAqwgcAALCK8AEAAKwifAAAAKsIHwAAwCrCBwAAsIrwAQAArCJ8AAAAqwgfAADAKsIHAACwivABAACs+v8gp5YZ+yOXhAAAAABJRU5ErkJggg==", + "text/plain": [ + "
" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "fname = '/data101/bartlett/fsigma8/PV_data/Foundation_DR1/Foundation_DR1.FITRES.TEXT'\n", + "\n", + "# Get header\n", + "columns = ['SN'] + linecache.getline(fname, 6).strip().split()[1:]\n", + "df = pd.read_csv(fname, sep=\"\\s+\", skipinitialspace=True, skiprows=7, names=columns)\n", + "\n", + "zCMB = df['zCMB']\n", + "m = df['mB']\n", + "m_err = df['mBERR']\n", + "\n", + "x1 = df['x1']\n", + "hyper_stretch_mu = np.median(x1)\n", + "hyper_stretch_sigma = (np.percentile(x1, 84) - np.percentile(x1, 16)) / 2\n", + "\n", + "c = df['c']\n", + "hyper_c_mu = np.median(c)\n", + "hyper_c_sigma = (np.percentile(c, 84) - np.percentile(c, 16)) / 2\n", + "\n", + "sigma_m = np.median(df['mBERR'])\n", + "sigma_stretch = np.median(df['x1ERR'])\n", + "sigma_c = np.median(df['cERR'])\n", + "print(sigma_m, sigma_stretch, sigma_c)\n", + "\n", + "print(hyper_stretch_mu, hyper_stretch_sigma)\n", + "print(hyper_c_mu, hyper_c_sigma)\n", + " \n", + "plt.figure()\n", + "speed_of_light = 299792 # km/s\n", + "mock_z = np.load('sn_z.npy')\n", + "mock_z /= speed_of_light\n", + "plt.hist(mock_z, bins=10, label='mock', histtype='step', density=True)\n", + "plt.hist(zCMB, bins=10, label='true', histtype='step', density=True)\n", + "plt.legend()" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "8e538aec-c47e-43ef-95ad-84afc388fc8c", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Filename: /data101/bartlett/fsigma8/PV_data/Foundation_DR1/kcor_PS1_none.fits\n", + "No. Name Ver Type Cards Dimensions Format\n", + " 0 PRIMARY 1 PrimaryHDU 41 (0,) \n", + " 1 ZPoff 1 BinTableHDU 19 11R x 5C [20A, 20A, 1E, 1E, 1E] \n", + " 2 SN SED 1 BinTableHDU 19 97626R x 1C [1E] \n", + " 3 KCOR 1 BinTableHDU 15 0R x 3C [1E, 1E, 1E] \n", + " 4 MAG+MWXTCOR 1 BinTableHDU 59 0R x 25C [1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E] \n", + " 5 FilterTrans 1 BinTableHDU 33 921R x 12C [1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E, 1E] \n", + " 6 PrimarySED 1 BinTableHDU 15 921R x 3C [1E, 1E, 1E] \n" + ] + } + ], + "source": [ + "# Open the FITS file\n", + "fname = '/data101/bartlett/fsigma8/PV_data/Foundation_DR1/kcor_PS1_none.fits'\n", + "with fits.open(fname) as hdul:\n", + " hdul.info() # Show summary of FITS file contents" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d012b67e-1600-4f5c-8724-6b926e4c32a9", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "borg_new", + "language": "python", + "name": "borg_new" + }, + "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.12.7" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/corner.png b/tests/corner.png index 54b13c4..287d7b5 100644 Binary files a/tests/corner.png and b/tests/corner.png differ diff --git a/tests/likelihood_scan_a_TFR.png b/tests/likelihood_scan_a_TFR.png index 7b02db2..6795c39 100644 Binary files a/tests/likelihood_scan_a_TFR.png and b/tests/likelihood_scan_a_TFR.png differ diff --git a/tests/likelihood_scan_alpha.png b/tests/likelihood_scan_alpha.png index d0df50c..4fe7b64 100644 Binary files a/tests/likelihood_scan_alpha.png and b/tests/likelihood_scan_alpha.png differ diff --git a/tests/likelihood_scan_b_TFR.png b/tests/likelihood_scan_b_TFR.png index 681a611..9fca64d 100644 Binary files a/tests/likelihood_scan_b_TFR.png and b/tests/likelihood_scan_b_TFR.png differ diff --git a/tests/likelihood_scan_mthresh.png b/tests/likelihood_scan_mthresh.png new file mode 100644 index 0000000..990fa0c Binary files /dev/null and b/tests/likelihood_scan_mthresh.png differ diff --git a/tests/likelihood_scan_sigma_TFR.png b/tests/likelihood_scan_sigma_TFR.png index cdd0074..7099d9a 100644 Binary files a/tests/likelihood_scan_sigma_TFR.png and b/tests/likelihood_scan_sigma_TFR.png differ diff --git a/tests/likelihood_scan_sigma_v.png b/tests/likelihood_scan_sigma_v.png index ef60705..52caf91 100644 Binary files a/tests/likelihood_scan_sigma_v.png and b/tests/likelihood_scan_sigma_v.png differ diff --git a/tests/sn_inference.py b/tests/sn_inference.py new file mode 100644 index 0000000..faa5e60 --- /dev/null +++ b/tests/sn_inference.py @@ -0,0 +1,533 @@ +import aquila_borg as borg +import pandas as pd +import linecache +import numpy as np +from astropy.coordinates import SkyCoord +import astropy.units as apu +import jax.numpy as jnp +import jax + +import matplotlib.pyplot as plt + +import borg_velocity.poisson_process as poisson_process +import borg_velocity.projection as projection +import borg_velocity.utils as utils + +from tfr_inference import get_fields, generateMBData + +# Output stream management +cons = borg.console() +myprint = lambda x: cons.print_std(x) if type(x) == str else cons.print_std(repr(x)) + + +def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha, + a_tripp, b_tripp, M_SN, sigma_SN, sigma_m, sigma_stretch, sigma_c, + hyper_stretch_mu, hyper_stretch_sigma, hyper_c_mu, hyper_c_sigma, + sigma_v, interp_order=1, bias_epsilon=1e-7): + """ + Create mock TFR catalogue from a density and velocity field + + Args: + - Nt (int): Number of tracers to produce + - L (float): Box length (Mpc/h) + - xmin (float): Coordinate of corner of the box (Mpc/h) + - cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters to use + - dens (np.ndarray): Over-density field (shape = (N, N, N)) + - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N)) + - Rmax (float): Maximum allowed comoving radius of a tracer (Mpc/h) + - alpha (float): Exponent for bias model + + - sigma_m (float): Uncertainty on the apparent magnitude measurements + + - sigma_v (float): Uncertainty on the velocity field (km/s) + - interp_order (int, default=1): Order of interpolation from grid points to the line of sight + - bias_epsilon (float, default=1e-7): Small number to add to 1 + delta to prevent 0^# + + Returns: + - all_RA (np.ndarrary): Right Ascension (degrees) of the tracers (shape = (Nt,)) + - all_Dec (np.ndarrary): Dec (np.ndarray): Delination (degrees) of the tracers (shape = (Nt,)) + - czCMB (np.ndarrary): Observed redshifts (km/s) of the tracers (shape = (Nt,)) + - all_mtrue (np.ndarrary): True apparent magnitudes of the tracers (shape = (Nt,)) + + - all_mobs (np.ndarrary): Observed apparent magnitudes of the tracers (shape = (Nt,)) + + - all_xtrue (np.ndarrary): True comoving coordinates of the tracers (Mpc/h) (shape = (3, Nt)) + - vbulk (np.ndarray): The bulk velocity of the box (km/s) + + """ + + # Initialize lists to store valid positions and corresponding sig_mu values + all_xtrue = np.empty((3, Nt)) + all_mtrue = np.empty(Nt) + all_stretchtrue = np.empty(Nt) + all_ctrue = np.empty(Nt) + all_mobs = np.empty(Nt) + all_stretchobs = np.empty(Nt) + all_cobs = np.empty(Nt) + all_RA = np.empty(Nt) + all_Dec = np.empty(Nt) + + # Counter for accepted positions + accepted_count = 0 + + # Bias model + phi = (1. + dens + bias_epsilon) ** alpha + + # Only use centre of box + x = np.linspace(xmin, xmin + L, dens.shape[0]+1) + i0 = np.argmin(np.abs(x + Rmax)) + i1 = np.argmin(np.abs(x - Rmax)) + L_small = x[i1] - x[i0] + xmin_small = x[i0] + phi_small = phi[i0:i1, i0:i1, i0:i1] + + # Loop until we have Nt valid positions + while accepted_count < Nt: + + # Generate positions (comoving) + xtrue = poisson_process.sample_3d(phi_small, Nt, L_small, (xmin_small, xmin_small, xmin_small)) + + # Convert to RA, Dec, Distance (comoving) + rtrue = np.sqrt(np.sum(xtrue** 2, axis=0)) # Mpc/h + c = SkyCoord(x=xtrue[0], y=xtrue[1], z=xtrue[2], representation_type='cartesian') + RA = c.spherical.lon.degree + Dec = c.spherical.lat.degree + r_hat = np.array(SkyCoord(ra=RA*apu.deg, dec=Dec*apu.deg).cartesian.xyz) + + # Compute cosmological redshift + zcosmo = utils.z_cos(rtrue, cpar.omega_m) + + # Compute luminosity distance + # DO I NEED TO DO /h??? + dL = (1 + zcosmo) * rtrue / cpar.h # Mpc + + # Compute true distance modulus + mutrue = 5 * np.log10(dL) + 25 + + # Sample true stretch and colour (c) from its prior + stretchtrue = hyper_stretch_mu + hyper_stretch_sigma * np.random.randn(Nt) + ctrue = hyper_c_mu + hyper_c_sigma * np.random.randn(Nt) + + # Obtain muSN from mutrue using the intrinsic scatter + muSN = mutrue + sigma_SN * np.random.randn(Nt) + + # Obtain apparent magnitude from the TFR + mtrue = muSN - (a_tripp * stretchtrue - b_tripp * ctrue) + M_SN + + # Scatter true observed apparent magnitudes and linewidths + mobs = mtrue + sigma_m * np.random.randn(Nt) + stretchobs = stretchtrue + sigma_stretch * np.random.randn(Nt) + cobs = ctrue + sigma_c * np.random.randn(Nt) + + # Apply apparement magnitude cut + m = np.ones(mobs.shape, dtype=bool) + mtrue = mtrue[m] + stretchtrue = stretchtrue[m] + ctrue = ctrue[m] + mobs = mobs[m] + stretchobs = stretchobs[m] + cobs = cobs[m] + xtrue = xtrue[:,m] + RA = RA[m] + Dec = Dec[m] + + # Calculate how many valid positions we need to reach Nt + remaining_needed = Nt - accepted_count + selected_count = min(xtrue.shape[1], remaining_needed) + + # Append only the needed number of valid positions + imin = accepted_count + imax = accepted_count + selected_count + all_xtrue[:,imin:imax] = xtrue[:,:selected_count] + all_mtrue[imin:imax] = mtrue[:selected_count] + all_stretchtrue[imin:imax] = stretchtrue[:selected_count] + all_ctrue[imin:imax] = ctrue[:selected_count] + all_mobs[imin:imax] = mobs[:selected_count] + all_stretchobs[imin:imax] = stretchobs[:selected_count] + all_cobs[imin:imax] = cobs[:selected_count] + all_RA[imin:imax] = RA[:selected_count] + all_Dec[imin:imax] = Dec[:selected_count] + + # Update the accepted count + accepted_count += selected_count + + myprint(f'\tMade {accepted_count} of {Nt}') + + # Obtain a bulk velocity + vhat = np.random.randn(3) + vhat = vhat / np.linalg.norm(vhat) + vbulk = np.random.randn() * utils.get_sigma_bulk(L, cpar) + vbulk = vhat * vbulk + + # Get the radial component of the peculiar velocity at the positions of the objects + myprint('Obtaining peculiar velocities') + tracer_vel = projection.interp_field( + vel, + np.expand_dims(all_xtrue, axis=2), + L, + np.array([xmin, xmin, xmin]), + interp_order + ) # km/s + myprint('Adding bulk velocity') + tracer_vel = tracer_vel + vbulk[:,None,None] + myprint('Radial projection') + vr_true = np.squeeze(projection.project_radial( + tracer_vel, + np.expand_dims(all_xtrue, axis=2), + np.zeros(3,) + )) # km/s + + # Recompute cosmological redshift + rtrue = jnp.sqrt(jnp.sum(all_xtrue ** 2, axis=0)) + zcosmo = utils.z_cos(rtrue, cpar.omega_m) + + # Obtain total redshift + vr_noised = vr_true + sigma_v * np.random.randn(Nt) + czCMB = ((1 + zcosmo) * (1 + vr_noised / utils.speed_of_light) - 1) * utils.speed_of_light + + return all_RA, all_Dec, czCMB, all_mtrue, all_stretchtrue, all_ctrue, all_mobs, all_stretchobs, all_cobs, all_xtrue, vbulk + + +def estimate_data_parameters(): + """ + Using Foundation DR1, estimate some parameters to use in mock generation. + + """ + + fname = '/data101/bartlett/fsigma8/PV_data/Foundation_DR1/Foundation_DR1.FITRES.TEXT' + + # Get header + columns = ['SN'] + linecache.getline(fname, 6).strip().split()[1:] + df = pd.read_csv(fname, sep="\s+", skipinitialspace=True, skiprows=7, names=columns) + + zCMB = df['zCMB'] + m = df['mB'] + m_err = df['mBERR'] + + x1 = df['x1'] + hyper_stretch_mu = np.median(x1) + hyper_stretch_sigma = (np.percentile(x1, 84) - np.percentile(x1, 16)) / 2 + + c = df['c'] + hyper_c_mu = np.median(c) + hyper_c_sigma = (np.percentile(c, 84) - np.percentile(c, 16)) / 2 + + sigma_m = np.median(df['mBERR']) + sigma_stretch = np.median(df['x1ERR']) + sigma_c = np.median(df['cERR']) + + return sigma_m, sigma_stretch, sigma_c, hyper_stretch_mu, hyper_stretch_sigma, hyper_c_mu, hyper_c_sigma + + + +def likelihood_vel(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk, + dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, + cz_obs, MB_pos): + """ + Evaluate the terms in the likelihood from the velocity and malmquist bias + + Args: + - alpha (float): Exponent for bias model + + - sigma_v (float): Uncertainty on the velocity field (km/s) + - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,)) + + - vbulk (np.ndarray): Bulk velocity of the box (km/s) (shape=(3,)) + - dens (np.ndarray): Over-density field (shape = (N, N, N)) + - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N)) + - omega_m (float): Matter density parameter Om + - h (float): Hubble constant H0 = 100 h km/s/Mpc + - L (float): Comoving box size (Mpc/h) + - xmin (float): Coordinate of corner of the box (Mpc/h) + - interp_order (int): Order of interpolation from grid points to the line of sight + - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^# + - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,)) + - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h). + The shape is (3, Nt, Nsig) + + Returns: + - loglike (float): The log-likelihood of the data + """ + + # Comoving radii of integration points (Mpc/h) + r = jnp.sqrt(jnp.sum(MB_pos ** 2, axis=0)) + + # p_r = r^2 n(r) N(mutrue; muTFR, sigmaTFR) + # Multiply by arbitrary number for numerical stability (cancels in p_r / p_r_norm) + number_density = projection.interp_field( + dens, + MB_pos, + L, + jnp.array([xmin, xmin, xmin]), + interp_order, + use_jitted=True, + ) + number_density = jax.nn.relu(1. + number_density) + number_density = jnp.power(number_density + bias_epsilon, alpha) + zcosmo = utils.z_cos(r, omega_m) + mutrue = 5 * jnp.log10((1 + zcosmo) * r / h) + 25 + mutripp = m_true + a_tripp * stretch_true - b_tripp * c_true - M_SN + d2 = ((mutrue - mutripp[:,None]) / sigma_SN) ** 2 + best = jnp.amin(jnp.abs(d2), axis=1) + d2 = d2 - jnp.expand_dims(jnp.nanmin(d2, axis=1), axis=1) + p_r = r ** 2 * jnp.exp(-0.5 * d2) * number_density + p_r_norm = jnp.expand_dims(jnp.trapezoid(p_r, r, axis=1), axis=1) + + # Peculiar velocity term + tracer_vel = projection.interp_field( + vel, + MB_pos, + L, + jnp.array([xmin, xmin, xmin]), + interp_order, + use_jitted=True, + ) + tracer_vel = tracer_vel + jnp.squeeze(vbulk)[...,None,None] + tracer_vr = projection.project_radial( + tracer_vel, + MB_pos, + jnp.zeros(3,) + ) + cz_pred = ((1 + zcosmo) * (1 + tracer_vr / utils.speed_of_light) - 1) * utils.speed_of_light + d2 = ((cz_pred - jnp.expand_dims(cz_obs, axis=1)) / sigma_v)**2 + scale = jnp.nanmin(d2, axis=1) + d2 = d2 - jnp.expand_dims(scale, axis=1) + + # Integrate to get likelihood + p_cz = jnp.trapezoid(jnp.exp(-0.5 * d2) * p_r / p_r_norm, r, axis=1) + lkl_ind = jnp.log(p_cz) - scale / 2 - 0.5 * jnp.log(2 * np.pi * sigma_v**2) + loglike = lkl_ind.sum() + + return loglike + + +def likelihood_stretch(stretch_true, stretch_obs, sigma_stretch): + """ + Evaluate the terms in the likelihood from stretch + + Args: + - stretch_true (np.ndarray): True stretch of the tracers (shape = (Nt,)) + - stretch_obs (np.ndarray): Observed stretch of the tracers (shape = (Nt,)) + - sigma_stretch (float): Uncertainty on the stretch measurements + + Returns: + - loglike (float): The log-likelihood of the data + """ + + Nt = stretch_obs.shape[0] + loglike = - ( + 0.5 * jnp.sum((stretch_obs - stretch_true) ** 2 / sigma_stretch ** 2) + + Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_stretch ** 2) + ) + + return loglike + + +def likelihood_c(c_true, c_obs, sigma_c): + """ + Evaluate the terms in the likelihood from colour + + Args: + - c_true (np.ndarray): True colours of the tracers (shape = (Nt,)) + - c_obs (np.ndarray): Observed colours of the tracers (shape = (Nt,)) + - sigma_c (float): Uncertainty on the colours measurements + + Returns: + - loglike (float): The log-likelihood of the data + """ + + Nt = c_obs.shape[0] + loglike = - ( + 0.5 * jnp.sum((c_obs - c_true) ** 2 / sigma_c ** 2) + + Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_c ** 2) + ) + + return loglike + + +def likelihood(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk, + dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, + cz_obs, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos): + """ + Evaluate the likelihood for SN sample + + Args: + - alpha (float): Exponent for bias model + + - sigma_v (float): Uncertainty on the velocity field (km/s) + - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,)) + + - vbulk (np.ndarray): Bulk velocity of the box (km/s) (shape=(3,)) + - dens (np.ndarray): Over-density field (shape = (N, N, N)) + - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N)) + - omega_m (float): Matter density parameter Om + - h (float): Hubble constant H0 = 100 h km/s/Mpc + - L (float): Comoving box size (Mpc/h) + - xmin (float): Coordinate of corner of the box (Mpc/h) + - interp_order (int): Order of interpolation from grid points to the line of sight + - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^# + - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,)) + - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,)) + + - sigma_m (float): Uncertainty on the apparent magnitude measurements + - sigma_eta (float): Uncertainty on the linewidth measurements + - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h). + The shape is (3, Nt, Nsig) + + Returns: + - loglike (float): The log-likelihood of the data + """ + + + loglike_vel = likelihood_vel(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk, + dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, + cz_obs, MB_pos) + loglike_stretch = likelihood_stretch(stretch_true, stretch_obs, sigma_stretch) + loglike_c = likelihood_c(c_true, c_obs, sigma_c) + + loglike = (loglike_vel + loglike_stretch + loglike_c) + + return loglike + + +def test_likelihood_scan(prior, alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk, + dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, + czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos): + """ + Plot likelihood as we scan through the paramaters [alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v] + to verify that the likelihood shape looks reasonable + + Args: + - prior (dict): Upper and lower bounds for a uniform prior for the parameters + - alpha (float): Exponent for bias model + + - sigma_v (float): Uncertainty on the velocity field (km/s) + - m_true (np.ndarray): True apparent magnitudes of the tracers (shape = (Nt,)) + + - dens (np.ndarray): Over-density field (shape = (N, N, N)) + - vel (np.ndarray): Velocity field (km/s) (shape = (3, N, N, N)) + - omega_m (float): Matter density parameter Om + - h (float): Hubble constant H0 = 100 h km/s/Mpc + - L (float): Comoving box size (Mpc/h) + - xmin (float): Coordinate of corner of the box (Mpc/h) + - interp_order (int): Order of interpolation from grid points to the line of sight + - bias_epsilon (float): Small number to add to 1 + delta to prevent 0^# + - cz_obs (np.ndarray): Observed redshifts (km/s) of the tracers (shape = (Nt,)) + - m_obs (np.ndarray): Observed apparent magnitudes of the tracers (shape = (Nt,)) + + - sigma_m (float): Uncertainty on the apparent magnitude measurements + - sigma_eta (float): Uncertainty on the apparent linewidth measurements + - MB_pos (np.ndarray): Comoving coordinates of integration points to use in likelihood (Mpc/h). + The shape is (3, Nt, Nsig) + + """ + + + pars = [alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v] + par_names = ['alpha', 'a_tripp', 'b_tripp', 'M_SN', 'sigma_SN', 'sigma_v'] + + orig_ll = - likelihood(*pars, m_true, stretch_true, c_true, vbulk, + dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, + czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos) + + for i, name in enumerate(par_names): + + myprint(f'Scanning {name}') + + if name in prior: + x = np.linspace(*prior[name], 20) + else: + pmin = pars[i] * 0.2 + pmax = pars[i] * 2.0 + x = np.linspace(pmin, pmax, 20) + + all_ll = np.empty(x.shape) + orig_x = pars[i] + for j, xx in enumerate(x): + pars[i] = xx + all_ll[j] = - likelihood(*pars, m_true, stretch_true, c_true, vbulk, + dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, + czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos) + pars[i] = orig_x + + plt.figure() + plt.plot(x, all_ll, '.') + plt.axvline(orig_x, ls='--', color='k') + plt.axhline(orig_ll, ls='--', color='k') + plt.xlabel(name) + plt.ylabel('Negative log-likelihood') + plt.savefig(f'sn_likelihood_scan_{name}.png') + fig = plt.gcf() + plt.clf() + plt.close(fig) + + return + + + + +def main(): + + myprint('Beginning') + + sigma_m, sigma_stretch, sigma_c, hyper_stretch_mu, hyper_stretch_sigma, hyper_c_mu, hyper_c_sigma = estimate_data_parameters() + + # Other parameters to use + L = 500.0 + N = 64 + xmin = -L/2 + R_lim = L / 2 + Rmax = 100 + Nt = 100 + alpha = 1.4 + sigma_v = 150 + interp_order = 1 + bias_epsilon = 1.e-7 + Nint_points = 201 + Nsig = 10 + frac_sigma_r = 0.07 # WANT A BETTER WAY OF DOING THIS - ESTIMATE THROUGH SIGMAS FROM Tripp formula + + # These values are from Table 6 of Boruah et al. 2020 + a_tripp = 0.140 + b_tripp = 2.78 + M_SN = - 18.558 + sigma_SN = 0.082 + + prior = { + 'alpha': [0.5, 4.5], + 'a_tripp': [0.01, 0.2], + 'b_tripp': [2.5, 4.5], + 'M_SN': [-19.5, -17.5], + 'hyper_mean_stretch': [hyper_stretch_mu - hyper_stretch_sigma, hyper_stretch_mu + hyper_stretch_sigma], + 'hyper_mean_c':[hyper_c_mu - hyper_c_sigma, hyper_c_mu + hyper_c_sigma], + 'sigma_v': [10, 3000], + } + + # Make mock + np.random.seed(123) + cpar, dens, vel = get_fields(L, N, xmin) + RA, Dec, czCMB, m_true, stretch_true, c_true, m_obs, stretch_obs, c_obs, xtrue, vbulk = create_mock( + Nt, L, xmin, cpar, dens, vel, Rmax, alpha, + a_tripp, b_tripp, M_SN, sigma_SN, sigma_m, sigma_stretch, sigma_c, + hyper_stretch_mu, hyper_stretch_sigma, hyper_c_mu, hyper_c_sigma, + sigma_v, interp_order=interp_order, bias_epsilon=bias_epsilon) + MB_pos = generateMBData(RA, Dec, czCMB, L, N, R_lim, Nsig, Nint_points, sigma_v, frac_sigma_r) + + # Test likelihood + loglike = likelihood(alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk, + dens, vel, cpar.omega_m, cpar.h, L, xmin, interp_order, bias_epsilon, + czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos) + myprint(f'loglike {loglike}') + + # Scan over parameters to make plots verifying behaviour + test_likelihood_scan(prior, alpha, a_tripp, b_tripp, M_SN, sigma_SN, sigma_v, m_true, stretch_true, c_true, vbulk, + dens, vel, cpar.omega_m, cpar.h, L, xmin, interp_order, bias_epsilon, + czCMB, m_obs, stretch_obs, c_obs, sigma_m, sigma_stretch, sigma_c, MB_pos) + + + return + + +if __name__ == "__main__": + main() + \ No newline at end of file diff --git a/tests/sn_likelihood_scan_M_SN.png b/tests/sn_likelihood_scan_M_SN.png new file mode 100644 index 0000000..f3ef1a0 Binary files /dev/null and b/tests/sn_likelihood_scan_M_SN.png differ diff --git a/tests/sn_likelihood_scan_a_tripp.png b/tests/sn_likelihood_scan_a_tripp.png new file mode 100644 index 0000000..b6ea99e Binary files /dev/null and b/tests/sn_likelihood_scan_a_tripp.png differ diff --git a/tests/sn_likelihood_scan_alpha.png b/tests/sn_likelihood_scan_alpha.png new file mode 100644 index 0000000..6fb0622 Binary files /dev/null and b/tests/sn_likelihood_scan_alpha.png differ diff --git a/tests/sn_likelihood_scan_b_tripp.png b/tests/sn_likelihood_scan_b_tripp.png new file mode 100644 index 0000000..7ef3c11 Binary files /dev/null and b/tests/sn_likelihood_scan_b_tripp.png differ diff --git a/tests/sn_likelihood_scan_sigma_SN.png b/tests/sn_likelihood_scan_sigma_SN.png new file mode 100644 index 0000000..225b44f Binary files /dev/null and b/tests/sn_likelihood_scan_sigma_SN.png differ diff --git a/tests/sn_likelihood_scan_sigma_v.png b/tests/sn_likelihood_scan_sigma_v.png new file mode 100644 index 0000000..9ba3d58 Binary files /dev/null and b/tests/sn_likelihood_scan_sigma_v.png differ diff --git a/tests/tfr_inference.py b/tests/tfr_inference.py index 3c570e3..1a76f96 100644 --- a/tests/tfr_inference.py +++ b/tests/tfr_inference.py @@ -2,7 +2,6 @@ import aquila_borg as borg import numpy as np from astropy.coordinates import SkyCoord import astropy.units as apu -import astropy.constants import pandas as pd import jax.numpy as jnp import jax.scipy.special @@ -181,7 +180,6 @@ def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh, - a_TFR (float): TFR relation intercept - b_TFR (float): TFR relation slope - sigma_TFR (float): Intrinsic scatter in the TFR - - sigma_v (float): Uncertainty on the velocity field (km/s) - sigma_m (float): Uncertainty on the apparent magnitude measurements - sigma_eta (float): Uncertainty on the linewidth measurements - hyper_eta_mu (float): Mean of the Gaussian hyper-prior for the true eta values @@ -263,8 +261,8 @@ def create_mock(Nt, L, xmin, cpar, dens, vel, Rmax, alpha, mthresh, etaobs = etatrue + sigma_eta * np.random.randn(Nt) # Apply apparement magnitude cut - # m = mobs <= mthresh - m = np.ones(mobs.shape, dtype=bool) + m = mobs <= mthresh + # m = np.ones(mobs.shape, dtype=bool) mtrue = mtrue[m] etatrue = etatrue[m] mobs = mobs[m] @@ -519,11 +517,10 @@ def likelihood_m(m_true, m_obs, sigma_m, mthresh): """ Nt = m_obs.shape[0] - # norm = 2 / (1 + jax.scipy.special.erf((mthresh - m_true) / (jnp.sqrt(2) * sigma_m))) / jnp.sqrt(2 * jnp.pi * sigma_m ** 2) - norm = jnp.sqrt(2 * jnp.pi * sigma_m ** 2) * jnp.ones(Nt) + norm = jnp.log(2) - jnp.log(jax.scipy.special.erfc(- (mthresh - m_true) / (jnp.sqrt(2) * sigma_m))) - 0.5 * jnp.log(2 * jnp.pi * sigma_m ** 2) loglike = - ( 0.5 * jnp.sum((m_obs - m_true) ** 2 / sigma_m ** 2) - + jnp.sum(jnp.log(norm)) + - jnp.sum(norm) + Nt * 0.5 * jnp.log(2 * jnp.pi * sigma_m ** 2) ) @@ -604,7 +601,7 @@ def test_likelihood_scan(prior, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, mthresh): """ - Plot likelihood as we scan through the paramaters [alpha, a_TFR, b_TFR, sigma_TFR, sigma_v] + Plot likelihood as we scan through the paramaters [alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, mthresh] to verify that the likelihood shape looks reasonable Args: @@ -673,6 +670,25 @@ def test_likelihood_scan(prior, alpha, a_TFR, b_TFR, sigma_TFR, sigma_v, m_true, fig = plt.gcf() plt.clf() plt.close(fig) + + # Now check the effect of varying mthresh on the likelihood + myprint(f'Scanning mthresh') + x = np.linspace(mthresh - 0.5, mthresh + 0.5, 20) + all_ll = np.empty(x.shape) + for j, xx in enumerate(x): + all_ll[j] = - likelihood(*pars, m_true, eta_true, vbulk, + dens, vel, omega_m, h, L, xmin, interp_order, bias_epsilon, + czCMB, m_obs, eta_obs, sigma_m, sigma_eta, MB_pos, xx) + plt.figure() + plt.plot(x, all_ll, '.') + plt.axvline(mthresh, ls='--', color='k') + plt.axhline(orig_ll, ls='--', color='k') + plt.xlabel('mthresh') + plt.ylabel('Negative log-likelihood') + plt.savefig(f'likelihood_scan_mthresh.png') + fig = plt.gcf() + plt.clf() + plt.close(fig) return @@ -966,7 +982,6 @@ if __name__ == "__main__": """ TO DO -- Reinsert magnitude cut - Deal with case where sigma_eta and sigma_m could be floats vs arrays """ \ No newline at end of file diff --git a/tests/trace.png b/tests/trace.png index 39659c2..17a126c 100644 Binary files a/tests/trace.png and b/tests/trace.png differ diff --git a/tests/true_predicted_eta.png b/tests/true_predicted_eta.png index 28b6580..7797e2a 100644 Binary files a/tests/true_predicted_eta.png and b/tests/true_predicted_eta.png differ diff --git a/tests/true_predicted_m.png b/tests/true_predicted_m.png index 78edaf1..217325e 100644 Binary files a/tests/true_predicted_m.png and b/tests/true_predicted_m.png differ