{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.11.2\n"
     ]
    }
   ],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import h5py\n",
    "\n",
    "import scipy\n",
    "print(scipy.__version__)\n",
    "scipy.interpolate.interpolate.RegularGridInterpolator = scipy.interpolate.RegularGridInterpolator\n",
    "\n",
    "sys.path.append('..')\n",
    "from darkhistory.config import load_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "****** Loading transfer functions... ******\n",
      "Using data at /n/holylabs/LABS/iaifi_lab/Users/yitians/darkhistory/DHdata_v1_1\n",
      "    for propagating photons...   Done!\n",
      "    for low-energy photons...  Done!\n",
      "    for low-energy electrons...  Done!\n",
      "    for high-energy deposition...  Done!\n",
      "    for total upscattered CMB energy rate...  Done!\n",
      "****** Loading complete! ******\n",
      "highengphot <class 'darkhistory.spec.transferfunclist.TransferFuncInterp'>\n",
      "lowengphot <class 'darkhistory.spec.transferfunclist.TransferFuncInterp'>\n",
      "lowengelec <class 'darkhistory.spec.transferfunclist.TransferFuncInterp'>\n",
      "highengdep <class 'darkhistory.history.histools.IonRSInterp'>\n",
      "CMB_engloss <class 'darkhistory.history.histools.IonRSInterp'>\n"
     ]
    }
   ],
   "source": [
    "tfs = {}\n",
    "tfs['dep_tf'] = load_data('dep_tf')\n",
    "for k, v in tfs['dep_tf'].items():\n",
    "    print(k, type(v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "****** Loading transfer functions... ******\n",
      "    for inverse Compton (Thomson)...  Done!\n",
      "    for inverse Compton (relativistic)...  Done!\n",
      "    for inverse Compton (energy loss)...  Done!\n",
      "****** Loading complete! ******\n",
      "dep_tf <class 'dict'>\n",
      "tf_helper <class 'dict'>\n",
      "ics_tf <class 'dict'>\n"
     ]
    }
   ],
   "source": [
    "tfs['ics_tf'] = load_data('ics_tf')\n",
    "for k, v in tfs.items():\n",
    "    print(k, type(v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf_E <class 'darkhistory.history.histools.IonRSInterp'>\n",
      "hep_lb <class 'darkhistory.history.histools.IonRSInterp'>\n",
      "lci <class 'darkhistory.history.histools.IonRSInterp'>\n",
      "hci <class 'darkhistory.history.histools.IonRSInterp'>\n"
     ]
    }
   ],
   "source": [
    "tfs['tf_helper'] = load_data('tf_helper')\n",
    "for k, v in tfs['tf_helper'].items():\n",
    "    print(k, type(v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "highengphot (500,)\n",
      "lowengphot (500,)\n",
      "lowengelec (500,)\n",
      "highengdep (4,)\n",
      "CMB_engloss ()\n",
      "thomson (5000,)\n",
      "rel (5000,)\n",
      "engloss (5000,)\n",
      "tf_E (4,)\n"
     ]
    }
   ],
   "source": [
    "with h5py.File('data/test_tf.h5', 'w') as hf:\n",
    "\n",
    "    # dep_tf\n",
    "    for k  in ['highengphot', 'lowengphot', 'lowengelec']:\n",
    "        tf = tfs['dep_tf'][k]\n",
    "        z = tf.get_tf(0.433, 0.302, 2244).sum_specs(np.sin(np.arange(500)))\n",
    "        z += tf.get_tf(0.760, 0.276, 384).sum_specs(np.sin(np.arange(500)))\n",
    "        z += tf.get_tf(0.930, 0.088, 18).sum_specs(np.sin(np.arange(500)))\n",
    "        z = z.N\n",
    "        print(k, z.shape)\n",
    "        hf.create_dataset(k, data=z)\n",
    "    for k in ['highengdep', 'CMB_engloss']:\n",
    "        tf = tfs['dep_tf'][k]\n",
    "        z = np.sin(np.arange(500)) @ tf.get_val(0.433, 0.302, 2244)\n",
    "        z += np.sin(np.arange(500)) @ tf.get_val(0.760, 0.276, 384)\n",
    "        z += np.sin(np.arange(500)) @ tf.get_val(0.930, 0.088, 18)\n",
    "        print(k, z.shape)\n",
    "        hf.create_dataset(k, data=z)\n",
    "\n",
    "    # ics_tf\n",
    "    for k  in ['thomson', 'rel', 'engloss']:\n",
    "        tf = tfs['ics_tf'][k]\n",
    "        z = tf.sum_specs(np.sin(np.arange(5000)))\n",
    "        z = z.N\n",
    "        print(k, z.shape)\n",
    "        hf.create_dataset(k, data=z)\n",
    "\n",
    "    # tf_helper\n",
    "    for k in ['tf_E']:\n",
    "        tf = tfs['tf_helper'][k]\n",
    "        z = tf.get_val(0.433, 0.302, 2244) @ np.sin(np.arange(500))\n",
    "        z += tf.get_val(0.760, 0.276, 384) @ np.sin(np.arange(500))\n",
    "        z += tf.get_val(0.930, 0.088, 18) @ np.sin(np.arange(500))\n",
    "        print(k, z.shape)\n",
    "        hf.create_dataset(k, data=z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RNG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.433 0.302 0.760 0.276 0.930 0.088 0.614 0.442 0.369 0.061'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "' '.join([f'{x:.3f}' for x in np.random.uniform(0.01, 0.99, 10)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2264\n",
      "384\n",
      "18\n"
     ]
    }
   ],
   "source": [
    "print(np.random.randint(1600, 3000))\n",
    "print(np.random.randint(40, 1600))\n",
    "print(np.random.randint(4, 40))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dm21cm",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
