{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.modules[\"mpi4py\"] = None\n",
    "\n",
    "import numpy as np\n",
    "import astropy.units as u\n",
    "from astropy.cosmology import z_at_value\n",
    "\n",
    "import lenstools\n",
    "from lenstools.pipeline.simulation import SimulationBatch,LensToolsCosmology\n",
    "from lenstools.pipeline.settings import EnvironmentSettings,NGenICSettings,PlaneSettings,MapSettings\n",
    "from lenstools.simulations.camb import CAMBSettings\n",
    "from lenstools.simulations.gadget2 import Gadget2Settings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Interchangeable settings for simulation batch</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#Geometrical parameters\n",
    "zmax = 3.1\n",
    "box_size_Mpc_over_h = 240.0\n",
    "nside = 512\n",
    "lens_thickness_Mpc = 120.0\n",
    "\n",
    "#Code settings: interchangeable\n",
    "camb = CAMBSettings()\n",
    "ngenic = NGenICSettings()\n",
    "ngenic.GlassFile = lenstools.data(\"dummy_glass_little_endian.dat\")\n",
    "gadget2 = Gadget2Settings()\n",
    "\n",
    "#Code settings: lenstools multi-lens-plane algorithm\n",
    "planes = PlaneSettings.read(\"maps.ini\")\n",
    "maps = MapSettings.read(\"maps.ini\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Cosmological models ($\\Omega_m,w,\\sigma_8$)</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "models = np.array([[0.26,-1,0.8],[0.29,-1,0.8]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Lay down directory tree</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.260_Ol0.740_w1.000_si0.800 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.260_Ol0.740_w1.000_si0.800 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.260_Ol0.740_w1.000_si0.800/512b240 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.260_Ol0.740_w1.000_si0.800/512b240 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.260_Ol0.740_w1.000_si0.800/512b240/ic1 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.260_Ol0.740_w1.000_si0.800/512b240/ic1 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.260_Ol0.740_w1.000_si0.800/512b240/ic1/ics created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.260_Ol0.740_w1.000_si0.800/512b240/ic1/snapshots created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.260_Ol0.740_w1.000_si0.800/512b240/ic1/Planes created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.260_Ol0.740_w1.000_si0.800/512b240/ic1/Planes created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.260_Ol0.740_w1.000_si0.800/512b240/Maps created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.260_Ol0.740_w1.000_si0.800/512b240/Maps created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.290_Ol0.710_w1.000_si0.800 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.290_Ol0.710_w1.000_si0.800 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.290_Ol0.710_w1.000_si0.800/512b240 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.290_Ol0.710_w1.000_si0.800/512b240 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.290_Ol0.710_w1.000_si0.800/512b240/ic1 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.290_Ol0.710_w1.000_si0.800/512b240/ic1 created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.290_Ol0.710_w1.000_si0.800/512b240/ic1/ics created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.290_Ol0.710_w1.000_si0.800/512b240/ic1/snapshots created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.290_Ol0.710_w1.000_si0.800/512b240/ic1/Planes created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.290_Ol0.710_w1.000_si0.800/512b240/ic1/Planes created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.290_Ol0.710_w1.000_si0.800/512b240/Maps created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Storage/Om0.290_Ol0.710_w1.000_si0.800/512b240/Maps created on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.260_Ol0.740_w1.000_si0.800/512b240/camb.param written on localhost\n",
      "[+] /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.290_Ol0.710_w1.000_si0.800/512b240/camb.param written on localhost\n"
     ]
    }
   ],
   "source": [
    "batch = SimulationBatch.current(\"environment.ini\")\n",
    "\n",
    "#Cycle over parameters and create one model per parameter\n",
    "for Om,w,si8 in models:\n",
    "    \n",
    "    #Lay down directory tree\n",
    "\tcosmo = LensToolsCosmology(Om0=Om,Ode0=1-Om,w0=1,sigma8=si8)\n",
    "\tmodel = batch.newModel(cosmo,parameters=[\"Om\",\"Ol\",\"w\",\"si\"])\n",
    "\tcollection = model.newCollection(box_size=box_size_Mpc_over_h*model.Mpc_over_h,nside=nside)\n",
    "\tr = collection.newRealization(seed=111)\n",
    "\n",
    "\t#Plane and catalog directories\n",
    "\tpln = r.newPlaneSet(planes)\n",
    "\tmp = collection.newMapSet(maps)\n",
    "\n",
    "#CAMB settings\n",
    "for model in batch.models:\n",
    "\tcollection = model.collections[0]\n",
    "\tcollection.writeCAMB(z=np.array([0.0]),settings=camb)\n",
    "\n",
    "#Compute comoving distance to maximum redshift for each model\n",
    "d = list()\n",
    "for model in batch.available:\n",
    "\td.append(model.cosmology.comoving_distance(zmax))\n",
    "\n",
    "#Compute lens spacings\n",
    "d = np.array([dv.value for dv in d]) * d[0].unit\n",
    "\n",
    "#We want to make sure there are lenses up to the maximum of these distances\n",
    "lens_distances = np.arange(lens_thickness_Mpc,d.max().to(u.Mpc).value + lens_thickness_Mpc,lens_thickness_Mpc) * u.Mpc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3>Write the NGen-IC and Gadget2 parameter files</h3>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[+] Solving the linear growth ODE for Om0.260_Ol0.740_w1.000_si0.800...\n",
      "[+] Computing prefactors...\n",
      "[+] NGenIC parameter file /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.260_Ol0.740_w1.000_si0.800/512b240/ic1/ngenic.param written on localhost\n",
      "[+] Gadget2 parameter file /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.260_Ol0.740_w1.000_si0.800/512b240/ic1/gadget2.param written on localhost\n",
      "[+] Solving the linear growth ODE for Om0.290_Ol0.710_w1.000_si0.800...\n",
      "[+] Computing prefactors...\n",
      "[+] NGenIC parameter file /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.290_Ol0.710_w1.000_si0.800/512b240/ic1/ngenic.param written on localhost\n",
      "[+] Gadget2 parameter file /Users/andreapetri/Documents/Cosmology_software/LensTools/notebooks/SimTest/Home/Om0.290_Ol0.710_w1.000_si0.800/512b240/ic1/gadget2.param written on localhost\n"
     ]
    }
   ],
   "source": [
    "#Compute the lens redshifts in each models and write the N-body simulations parameter files\n",
    "for model in batch.models:\n",
    "\n",
    "\t#Compute the redshifts of the Gadget snapshots\n",
    "\tz = np.zeros_like(lens_distances.value)\n",
    "\tfor n,dlens in enumerate(lens_distances):\n",
    "\t\tz[n] = z_at_value(model.cosmology.comoving_distance,dlens)\n",
    "\n",
    "\t#Assgn values to gadget settings\n",
    "\tgadget2.OutputScaleFactor = np.sort(1/(1+z))\n",
    "\n",
    "\t#Write parameter files\t\t\n",
    "\tcollection = model.collections[0]\n",
    "\n",
    "\t#Write NGenIC and Gadget2 parameter files\n",
    "\tr = collection.realizations[0]\n",
    "\tr.writeNGenIC(ngenic)\n",
    "\tr.writeGadget2(gadget2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}