mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-02-23 18:10:55 +00:00
1141 lines
2.6 MiB
Text
1141 lines
2.6 MiB
Text
|
{
|
||
|
|
"cells": [
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 1,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stdout",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"Populating the interactive namespace from numpy and matplotlib\n"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"%pylab inline\n",
|
||
|
|
"import camels_library as cl\n",
|
||
|
|
"import readgadget\n",
|
||
|
|
"from tqdm import tqdm"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 2,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"init_cond = '/data/CAMELS/Sims/IllustrisTNG_DM/CV_0/ICs/ics'\n",
|
||
|
|
"\n",
|
||
|
|
"header = readgadget.header(init_cond)\n",
|
||
|
|
"BoxSize = header.boxsize/1e3 #Mpc/h\n",
|
||
|
|
"Nall = header.nall #Total number of particles\n",
|
||
|
|
"Masses = header.massarr*1e10 #Masses of the particles in Msun/h\n",
|
||
|
|
"Omega_m = header.omega_m #value of Omega_m\n",
|
||
|
|
"Omega_l = header.omega_l #value of Omega_l\n",
|
||
|
|
"h = header.hubble #value of h\n",
|
||
|
|
"redshift = header.redshift #redshift of the snapshot\n",
|
||
|
|
"Hubble = 100.0*np.sqrt(Omega_m*(1.0+redshift)**3+Omega_l)#Value of H(z) in km/s/(Mpc/h)\n",
|
||
|
|
"\n",
|
||
|
|
"ptype = [1] #dark matter is particle type 1\n",
|
||
|
|
"ids_i = np.argsort(readgadget.read_block(init_cond, \"ID \", ptype)-1) #IDs starting from 0\n",
|
||
|
|
"pos_i = readgadget.read_block(init_cond, \"POS \", ptype)[ids_i]/1e3 #positions in Mpc/h\n",
|
||
|
|
"vel_i = readgadget.read_block(init_cond, \"VEL \", ptype)[ids_i] #peculiar velocities in km/s\n",
|
||
|
|
"\n",
|
||
|
|
"# Reordering data for simple reshaping\n",
|
||
|
|
"pos_i = pos_i.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
|
||
|
|
"vel_i = vel_i.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
|
||
|
|
"pos_i = (pos_i/BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
|
||
|
|
"vel_i = (vel_i / 100 * (1./(1+redshift)) / BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
|
||
|
|
"a_i = 1./(1+redshift)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 3,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stderr",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"100%|██████████| 34/34 [01:54<00:00, 3.36s/it]\n"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"scales = []\n",
|
||
|
|
"poss = []\n",
|
||
|
|
"vels = []\n",
|
||
|
|
"\n",
|
||
|
|
"# Loading all the intermediate snapshots\n",
|
||
|
|
"for i in tqdm(range(34)):\n",
|
||
|
|
" snapshot='/data/CAMELS/Sims/IllustrisTNG_DM/CV_0/snap_%03d.hdf5'%i\n",
|
||
|
|
" \n",
|
||
|
|
" header = readgadget.header(snapshot)\n",
|
||
|
|
" \n",
|
||
|
|
" redshift = header.redshift #redshift of the snapshot\n",
|
||
|
|
" h = header.hubble #value of h\n",
|
||
|
|
" \n",
|
||
|
|
" ptype = [1] #dark matter is particle type 1\n",
|
||
|
|
" ids = np.argsort(readgadget.read_block(snapshot, \"ID \", ptype)-1) #IDs starting from 0\n",
|
||
|
|
" pos = readgadget.read_block(snapshot, \"POS \", ptype)[ids]/1e3 #positions in Mpc/h\n",
|
||
|
|
" vel = readgadget.read_block(snapshot, \"VEL \", ptype)[ids] #peculiar velocities in km/s\n",
|
||
|
|
"\n",
|
||
|
|
" # Reordering data for simple reshaping\n",
|
||
|
|
" pos = pos.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
|
||
|
|
" vel = vel.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
|
||
|
|
" \n",
|
||
|
|
" pos = (pos / BoxSize * 64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
|
||
|
|
" vel = (vel / 100 * (1./(1+redshift)) / BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
|
||
|
|
" \n",
|
||
|
|
" scales.append((1./(1+redshift)))\n",
|
||
|
|
" poss.append(pos)\n",
|
||
|
|
" vels.append(vel)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 4,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"[<matplotlib.lines.Line2D at 0x7f82862318b0>]"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 4,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAXQAAAD4CAYAAAD8Zh1EAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAAgZklEQVR4nO3dd3xW9f3+8debsPcIO4RRZtgQiKtqRSogiltwolC01dbR2lqto2rrqNaqP6pSoSJoceBAoaAiKFgZYUPCCGEljATCCGTf9+f7R277i5RAAndy7vvO9Xw8eJj7nEPOxREuDp9zPueYcw4REQl/1bwOICIiwaFCFxGJECp0EZEIoUIXEYkQKnQRkQhR3asdR0dHuw4dOni1exGRsLRixYr9zrnmJ1rnWaF36NCBxMREr3YvIhKWzGxHaes05CIiEiFU6CIiEUKFLiISIVToIiIRQoUuIhIhTlnoZjbFzDLMbH0p683MXjazFDNba2YDgh9TREROpSxn6G8Cw06yfjjQJfBjAvDqmccSEZHyOmWhO+e+AbJOssko4C1XbAnQ2MxaByugiEik8Pkdf5qdRNrBnAr5/sEYQ28L7CrxOS2w7H+Y2QQzSzSzxMzMzCDsWkQkfPx5TjL/WLSNbzbvr5DvX6kXRZ1zk5xz8c65+ObNTzhzVUQkIk1bsoPJi7cx9pwO3JAQWyH7CEahpwPtSnyOCSwTERFg4aYMHp+1gSHdW/DIyLgK208wCn0WcEvgbpezgMPOuT1B+L4iImEvec8R7n5nFd1bNeDlMf2JqmYVtq9TPpzLzP4FXAhEm1ka8BhQA8A59xowBxgBpAA5wG0VFVZEJJxkHMlj3JvLqV+rOpNvHUS9WhX7PMRTfnfn3JhTrHfAXUFLJCISAXIKihg3NZFDuYW8f+fZtGpUu8L3qZmiIiJB5vM77pmxmg27D/PKmP70bNOoUvarQhcRCbKn5yTzRdI+Hh0Zx5AeLSttvyp0EZEgmr5kB28Ebk8ce27HSt23Cl1EJEi+3pzJY5Vwe2JpVOgiIkGQvOcId729km4tK/72xNKo0EVEzlDawRxunbKMBrWrM3lsfIXfnlgaFbqIyBk4eKyAW6YsI6/Qx9TbB9O6UR3Psnjz14iISATILfBx+9TlpB/MZfr4BLq2bOBpHp2hi4ichiKfn7vfWcmaXYd4eUx/BnVo6nUknaGLiJSXc46HPlrH/I0ZPHVFLy7p2crrSIDO0EVEyu2vX2zmvcQ0fjWkCzed1d7rOP+lQhcRKYdp323nla9SGD2oHfdd3MXrOD+gQhcRKaO56/fw6KwNXNyjBU9d0Quzyr/X/GRU6CIiZbA09QC/mrGa/u0a88qYAVSPCr36DL1EIiIhZuPeI4x/K5HYpnWZMnYQdWpGeR3phFToIiInsX3/MW6evIy6NaOYevtgGtet6XWkUqnQRURKsedwLje+sRSf3zF9XAJtG3s3C7QsVOgiIiew/2g+N76xlCO5hbx1+2C6eDwLtCw0sUhE5DiHcwu5ZfIydh/KZdq4BHq1rZw3Dp0pnaGLiJSQU1DE7W8uZ0tGNq/fHB8SU/rLSoUuIhKQX+TjjmkrWLXzIC+P7s8FXZt7HalcNOQiIkLxw7Z+9a9VLNqyn+ev7cvw3q29jlRuOkMXkSrP73f89oO1zNuwj8cvi+OagTFeRzotKnQRqdKcczz+6QY+XJXOb37atdJf7BxMKnQRqbKcczw3bxNvfbeDO87vxF0/6ex1pDOiQheRKuvFLzbz6sKt3JAQy4PDu4fcw7bKS4UuIlXS377czMuBx+A+NSr0npx4OlToIlLlvDx/C3/7cgvXDozhz1f2plq18C9zUKGLSBUzcUEKf/1iM1cPiOGZq/tETJmDCl1EqpBXF27lL/M2cVX/tjx3TR+iIqjMQYUuIlXEpG+28uzcjYzq14a/XNs34socVOgiUgW8sSiVP8/ZyGV92/BChJY5qNBFJMJNXryNp2Ync2nv1rx4Xd+QfHVcsETur0xEqrw3v93Gk58lMbxXK/42ul9Elzno4VwiEqHeWJTKU7OT+WlcS14e058aEV7moEIXkQj0yvwtvPDF5uJhluv7VYkyBxW6iEQQ5xzPf76JiQu2/vfWxEgfZimpTL9SMxtmZpvMLMXMHjzB+lgzW2Bmq8xsrZmNCH5UEZHSOed4anYyExdsZczgdjx/bWRfAD2RU/5qzSwKmAgMB+KAMWYWd9xmfwDec871B0YDfw92UBGR0vj9jkc+Wc/kxdsYe06HiJrOXx5l+etrMJDinEt1zhUAM4BRx23jgIaBrxsBu4MXUUSkdD6/47cz1zJ9yU7uuKATj10WFxEP2jodZRlDbwvsKvE5DUg4bpvHgc/N7JdAPeDiE30jM5sATACIjY0tb1YRkR8o9Pm5/701fLpmN/de3IV7hnSpsmUOwbsPfQzwpnMuBhgBTDOz//nezrlJzrl451x88+bh9fJVEQkt+UU+7n5nJZ+u2c2Dw7tz78Vdq3SZQ9nO0NOBdiU+xwSWlTQOGAbgnPvOzGoD0UBGMEKKiJSUV+jjzukrWLgpk8cviwvr18YFU1nO0JcDXcyso5nVpPii56zjttkJDAEwsx5AbSAzmEFFRAAO5xZyy+RlfL05k6ev6q0yL+GUZ+jOuSIzuxuYB0QBU5xzG8zsCSDROTcL+DXwDzO7j+ILpGOdc64ig4tI1ZORncetU5aTkpHNK2P6M7JPG68jhZQyTSxyzs0B5hy37NESXycB5wY3mojI/7fzQA43TV7K/qP5TL51EOd31XW442mmqIiEvOQ9R7hlyjIKfX7eHp9A/9gmXkcKSSp0EQlpy7dncfuby6lfqzrvjD+bLi0beB0pZKnQRSRkzU/exy/eXknbJnWYNi6Bto3reB0ppKnQRSQkzVyRxm9nrqVnm4b8c+wgmtWv5XWkkKdCF5GQ8/2zzM/5UTMm3RJP/VqqqrLQURKRkOH3O56dt5HXv07971uGalWP8jpW2FChi0hIyCv08Zv31/DZ2j3cmBDLE6N6RezLnCuKCl1EPHcop4AJb61g2fYsHhzenTvO71Tln8tyOlToIuKpXVk53PrPZaRl5fLS6H6M6tfW60hhS4UuIp5Zm3aI299MpKDIx7Rxg0no1MzrSGFNhS4inpifvI+731lF03o1mTEhgc4tNGHoTKnQRaTSTV+yg0c/WU9cm4ZMGTuIFg1qex0pIqjQRaTS+P2O5+Zt4rWvt3JR9xa8MqY/9XSPedDoSIpIpSh5W+INCbE8cXlPqkcF66VpAip0EakEGUfy+NlbiaxNP8zvhnXnzgt0W2JFUKGLSIVan36Y8VMTOZJXyGs3DeSSnq28jhSxVOgiUmHmrt/Dfe+uoUndGnxw5znEtWnodaSIpkIXkaBzzjFxQQrPf76Z/rGNef3mgbqTpRKo0EUkqPIKfTw4cy0fr97NFf3a8MzVfahdQw/YqgwqdBEJmozsPO6YtoJVOw/xwCXd+MWFP9LFz0qkQheRoEjafYTxU5dzMKeQ124awLBerb2OVOWo0EXkjH26Zje//WAtjerU4P07z6ZX20ZeR6qSVOgictqKfH7+Mm8Tr3+TSnz7Jvz9xgG0aKiLn15RoYvIaTl4rIBf/msVi1P2c/NZ7XlkZBw1q2vmp5dU6CJSbuvTD3Pn9BVkZOfz3DV9uC6+ndeRBBW6iJTTR6vSeHDmOprWq8n7d5xN33aNvY4kASp0ESmTQp+fp+dsZMq320jo2JSJNw4gun4tr2NJCSp0ETml/UfzuevtlSzdlsVt53bgoRE9qKEnJYYcFbqInNTKnQe5++2VHDhWwIvX9+XK/jFeR5JSqNBF5IScc0z5djtPz0mmdePazPz5Obq/PMSp0EXkfxzJK+S3769l7oa9DI1ryfPX9KVR3Rpex5JTUKGLyA+sTz/ML95eye5Dufzh0h6MO6+jnscSJlToIgIUD7G8s2wnf/w0iWb1avLuHWcxsH1Tr2NJOajQRYRj+UU8/NE6Pl69m/O7Nudv1/ejab2aXseSclKhi1RxW/Zl8/O3V5KaeZRfD+3KXT/pTLVqGmI
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 1 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"plot(scales)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 5,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"import jax\n",
|
||
|
|
"import jax.numpy as jnp\n",
|
||
|
|
"import jax_cosmo as jc\n",
|
||
|
|
"import haiku as hk\n",
|
||
|
|
"\n",
|
||
|
|
"from jax.experimental.ode import odeint\n",
|
||
|
|
"\n",
|
||
|
|
"from jaxpm.painting import cic_paint, cic_read\n",
|
||
|
|
"from jaxpm.pm import linear_field, lpt, make_ode_fn, pm_forces\n",
|
||
|
|
"from jaxpm.kernels import fftk, gradient_kernel, laplace_kernel, longrange_kernel"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 6,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"mesh_shape= [64,64, 64]\n",
|
||
|
|
"box_size = [25.,25.,25.]\n",
|
||
|
|
"cosmo = jc.Planck15(Omega_c= 0.3 - 0.049, Omega_b=0.049, n_s=0.9624, h=0.671 , sigma8=0.8)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 7,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAkwAAAJCCAYAAAAybpizAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9bZAcyZkmiD3uER4RGZlZWVUooFFgo9gEyG4M2RhOg2SDHz0fnNudXe3caj9Mc9pZkzRnq7O11YedrUwmu7G7k0xmkpn2/sh0Z6Zb2dqeTneSxvZu7XZHd5q5O+7scW63dzhoDtGcQZNENwlMs9CNRAOFQmVVZmREeLi7fri/Hh5RCXY3u7tI08HNyoDKyoyM8PBwf/x5n/d5mTEGT9qT9qQ9aU/ak/akPWlP2uMb/0mfwJP2pD1pT9qT9qQ9aU/aT3t7ApietCftSXvSnrQn7Ul70t6lPQFMT9qT9qQ9aU/ak/akPWnv0p4ApiftSXvSnrQn7Ul70p60d2lPANOT9qQ9aU/ak/akPWlP2ru0J4DpSXvSnrQn7Ul70p60J+1d2gcCTIyxP88Ye50x9gPG2G9+WCf1pP13qz0ZR0/aB21PxtCT9mG0J+PoSftRjf24PkyMsQjAGwD+LIC3AHwTwK8bY7774Z3ek/b/7+3JOHrSPmh7MoaetA+jPRlHT9q7tQ/CML0I4AfGmNvGmBrAPwDwlz6c03rS/jvUnoyjJ+2Dtidj6En7MNqTcfSk/cgWf4DPfgzAneD3twBc/VEfyPMtM5mch1ISxmgYoxFFApwLGKPcawaAAWMcjHFwHsMYA60b/xnAsmL0fsYYAIAxDoCBMQbGOOI4BsBgDPOfbT9DWNHAGANjlD+G/V27v2kA9vicRwCY+xfunExwDDq+csdm7tyYux4W9Ib9G10fXVN4XbZP6G/2vYDti/Y6I3Aedc4DAOI4A+cxhJgB4GiaCYzR7rOrG+cxGGv7SikZ9LVx/UPX0H7/4eHbWC4fstVHfdf2vsZRkgyNSidIjIIGQ+LuYwEgNQrCaGgwaMagwKAYQ87tMNdaoTYaI90g5FWZu0oDQDHu/yYZRxZnfqxo3WAe9B+HAad+ARDBgBkDwxi4MRBGIzIGMTSUO6eSRZCM288CiIwBc9/IjUHkjsmD40vGocH8+xp3PhwGDeNY8BiR+z4DwDCGyBhERiM1GhwG2n1auXPLXP9pxtCAQTKOhnEoxpAYDW4MSh6hZhHWxMA9g/a5rbXy56YYg2IcCgxD99wBdgwvjUYWvFczBuPuiQn6vWYRDAA9n+4ZY06/x3ETtvc9F42jxGzFmb2vzF5/ohVS/+za3SRz/a8Yw9KNo8RoCKMxcGPBAL4vaTwod60MBopxHHKBxvVxphuMtfTjM3HjRAf9wowBh7Hjyo8x9yy672pnpnb2YMHnAfixGMFAu/PkwTgy7viKMVQsQgSD2Gj/HZHRYACEGxOdcwSC8W+vU7vjKjCEE4JiDLMo8f0XG43UjY2oN44MgNrNsfYZ0kjceTB3Dsod3TDmr1UzhgdNiUMlT2QuAoBJJMzpOPPnA9h5QBgNZuz9o/6W7vpoHKVGIXbPCH1auWuje6/A/DU2jOMoEmjAUXOOgVYYaenvK/UrjT06Rv88+i18VbszsffX+GeUxk04TmkFle6Zj9zcW/KonY/cOOWwYyUO5iP0zonGk2IM2j0b9FzRmNNgOHTjKA6eQzpfmk1pLqvc3CKMRgTjxxxz72ncOKVrpT5/R1WYrRhHHwQwvafGGPubAP4mAEwmO/jX/rU/xOHhW35Rnkx2sLX1DZTlBdT1NspSQqkKSTJyC/4jKLWB2WwKKReYz+85EGMXMK0bRFECxjiSZIQoSjAYnEKSDDEacWidQUo4kKZQVUdomiWapoIxClIWHoxZkJUhihIIkaOu5yiKPXAeu3PJEUWJ/z4ADljUaJoSVXUYgA37Gv2eJCPEcYY83wLnMbRuwLlAmo5RVUcoigce0FjgwlGWB9C6AWMcUZQgy9ahdYPlct/3QRxnECL3n6F/h8MzSNM1nDp1HwBQ1+dQlgmOju5CiBxJMoJSNbRuHDAC8vwUIjcYlapRFA8RRcL3RVkeeCAbxwP3tyH+w//wF05sDGXZOn7pS/8rNE2FPD+FjY0L/v48enQbh4dv+T6gnyhKYIxGWR7gTj3HS/N7AOyDXrAIRSQgXH9ORQ7JOGZRgr04w5fz04giAcY46nqOqjrEK+5BvVw8xHazhASDdBP8jAvcyE9BaIXcKOxUR7hUHvhr2U3HfuEAgMJNnsJoTFSNi+XMnhuPkOsGuW4651mwCNdGT6HgMcaqtpMVj7FTz3Gl2EOuGwijIRmHBMNESwitME2G/ji5UbhYHUKCYRankGAoIoGpyDGLEj8BSjBMkyH+zFM/C8Y4pCzcM/EQ11WFrabEXpzhBR7jFaOhAPxivoU4zpAkI9T1HH+8/wN7nSxqrzkS+CKPcU03SIzGF6MUnEf4p//03/zhRzB8fAvH0Uac4W984qu4ma3j+uAUPrn1HL63/wP8b+59GwWPfd/lurG/M46pyCGMxsXqEKeaEherQ9tPjGMvznAYJRi4/j8M+rHgMW5m65hFCaYix049x+XiIYRbWM/Xc0xUjb04s/fNAavdZOT/TveV/qVzmkUJGsYRu/G75DEGusFWU/pzo3FE95Z+bqdrkIxjomo/Ziaqxvl67r+LroGeDzpHOmbuQCMds+Cx/3kYZ368FSzCb68/gyISyJXEjlzg8nIfE1VjTdV+EzAVOQoe+/O5nm9homp89ehup6/34gzCaAitsN0skesGe3GG/8Pbr3yUQ+jYONqMM/zGJ34Zf3/rEgoW4fk4xfDwLfzGwzd8301UjdhoLF2/TOMBcqNwoTrEVlNip553rm0WJfbagvtM9/ZmOsFenOFWNsGl8gBXF/f9vd+WhR9HNK9IxnFHDJEb5ccRjaH+OKLvoc/RGADg3+/nluDnTjICAKypGoAFImuqxrYs/HfRMel4U5H7vgm/B7DPS3g+u8kIh1Hix/DX1p7213exnOGlxTsY9K7rTjLy86wEs+MFBherQw+07oghbmUTbMsCW02JNVUj1w3uiRz/xvT6ynv/QQDT2wDOB78/7V7rNGPM3wPw9wDgYx/7nLELmJ0goyjBePwWJpPfRxQdgTEJYAdKDcA5wFgDQMAYDmJ7LAPUMiW0WFqmxf5LC5yUI/ceC5gsgLE7SKUqSLmE1tIfx7JR9thNU3lAAcB/L/3QQmyvUTsglbpjRWiaEpzHaJoSWjeI4wxxPPCgjo4ZRQkaN7m1rI5tdB5CDHp9qlf+znnswZxSNaQsUFU7YGwJpUbQeumuE51roPOxfQlE0RxxHIOx0x78aa08IIuiBHGcIooScB71mLP33d51HIVjaDI5bxjjEGLgziPz74vjzINKe09S148SWjd4ramwo2oI7VgEDsAt5H7CD4DJi4z78ZEkQzDG8VpTYeIWuomWmDQVikhAamMnehhs1wsAgIBBbixwIgBDk7zkkV9g6PsB+O+exSly3fgJRWjlX5eOEQpBCL0GwE9qACC0ggh2uJJHgFKQbocnjAaCyZFAFJ0XAEhZgDGOpimhVA3AtH0IO46vOqaTxtNweAZJMsZnmhJNU0LKAgDwitG46p7Zq7Dj0I6hD5R/8r7novPZxFAfSh75Z4gWGwkGuDmcQAKB2ombWMMm3O6eWrjQSTDftyLW/p6EEzz90LFCIBOCpfD7ALs49Rc6uIU5XNQ699O9Zs9DdcYL/V0E1xL+7q/JHVP05iI6t1w3ODQagr4/EsiNglD2unMl/Rhqet9N57qbjPDOuc9jqhtM77+GbVngUnmAvTjDVOTYakpMtGqfBx7j+Nm8r/a+x9Ez6Zrx/cwjvGI0XnLAJ+wvAnSScTsnKOnv749qBY99/xQ89vPJpKmQ6wZxb3z0x5NkHBMtjwEKOjf6PdzEUZPB2KJ73f98f6z0Pw/AA7NVfwvHUP+8wv/TxqXgsR23jCNXErlbz+n+Txxoo/l2L86wm47x/Z2XYIzG/Ue3cak8wFeP7mI3GeF
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 720x720 with 16 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"figure(figsize=[10,10])\n",
|
||
|
|
"for i in range(16):\n",
|
||
|
|
" subplot(4,4,i+1)\n",
|
||
|
|
" imshow(cic_paint(jnp.zeros(mesh_shape), poss[::2][i]).sum(axis=0), cmap='gist_stern', vmin=0)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "markdown",
|
||
|
|
"metadata": {},
|
||
|
|
"source": [
|
||
|
|
"## Defining the force correction kernel"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 8,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"def deBoorVectorized(x, t, c, p):\n",
|
||
|
|
" \"\"\"\n",
|
||
|
|
" Evaluates S(x).\n",
|
||
|
|
"\n",
|
||
|
|
" Args\n",
|
||
|
|
" ----\n",
|
||
|
|
" x: position\n",
|
||
|
|
" t: array of knot positions, needs to be padded as described above\n",
|
||
|
|
" c: array of control points\n",
|
||
|
|
" p: degree of B-spline\n",
|
||
|
|
" \"\"\"\n",
|
||
|
|
" k = jnp.digitize(x, t) -1\n",
|
||
|
|
" \n",
|
||
|
|
" d = [c[j + k - p] for j in range(0, p+1)]\n",
|
||
|
|
" for r in range(1, p+1):\n",
|
||
|
|
" for j in range(p, r-1, -1):\n",
|
||
|
|
" alpha = (x - t[j+k-p]) / (t[j+1+k-r] - t[j+k-p])\n",
|
||
|
|
" d[j] = (1.0 - alpha) * d[j-1] + alpha * d[j]\n",
|
||
|
|
" return d[p]"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 9,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"class FreeFormFilter(hk.Module):\n",
|
||
|
|
" \"\"\"This is a NN that predicts the score.\"\"\"\n",
|
||
|
|
" def __call__(self, x, a):\n",
|
||
|
|
" \n",
|
||
|
|
" net = jax.nn.leaky_relu(hk.Linear(128)(a))\n",
|
||
|
|
" net = jax.nn.leaky_relu(hk.Linear(128)(net))\n",
|
||
|
|
" \n",
|
||
|
|
" w = hk.Linear(32)(net) # hk.get_parameter(\"w\", [32], init=jnp.ones)\n",
|
||
|
|
" k = hk.Linear(31)(net) # hk.get_parameter(\"k\", [31], init=jnp.ones)\n",
|
||
|
|
" \n",
|
||
|
|
" # make sure the knots sum to 1 and are in the interval 0,1\n",
|
||
|
|
" k = jnp.concatenate([jnp.zeros((1,)),\n",
|
||
|
|
" jnp.cumsum(jax.nn.softmax(k))])\n",
|
||
|
|
"\n",
|
||
|
|
" # Augment with repeating points\n",
|
||
|
|
" ak = jnp.concatenate([jnp.zeros((3,)), k, jnp.ones((3,))])\n",
|
||
|
|
"\n",
|
||
|
|
" return deBoorVectorized(x, ak, w, 3)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 10,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"model = hk.without_apply_rng(hk.transform(lambda x,a : FreeFormFilter()(x,a)))\n",
|
||
|
|
"rng_seq = hk.PRNGSequence(12)\n",
|
||
|
|
"params = model.init(next(rng_seq), jnp.zeros([64]), jnp.ones([1]))\n",
|
||
|
|
"paramsi = model.init(next(rng_seq), jnp.zeros([64]), jnp.ones([1]))"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 11,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"[<matplotlib.lines.Line2D at 0x7f82484566a0>]"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 11,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAXwAAAD6CAYAAACiefy7AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAABCMElEQVR4nO29eXibZ5X3/zmyLdnyKnl3vKVx0qy0TdydFmgDXWAahmFp2QoDdAamw8yPgXk7wzu8M8wFb4EZfrNQlkIZSlkKFIZmIG2hhUK3tEkX0rhJmtVLvG+SN0mWdL9/SI+jOLKt5ZFkSffnunJFy2M99+Plq6Nzf885opRCo9FoNLmPJdML0Gg0Gk160IKv0Wg0eYIWfI1Go8kTtOBrNBpNnqAFX6PRaPIELfgajUaTJ5gi+CJyvYgcEZFjInLHMsf9iYgoEek047wajUajiZ3CZF9ARAqAu4A3An3APhHZrZR6ZdFx5cBfAc/G8ro1NTWqvb092eVpNBpNXvH888+PKqVqoz2XtOADlwDHlFInAETkfmAX8Mqi4/4Z+ALwqVhetL29nf3795uwPI1Go8kfRKR7qefMSOmsAXoj7veFH4tcwHagRSn1SxPOp9FoNJoESPmmrYhYgC8DfxPDsbeJyH4R2T8yMpLqpWk0Gk1eYYbgnwZaIu43hx8zKAe2Ao+LyCngMmB3tI1bpdTdSqlOpVRnbW3UFJRGo9FoEsQMwd8HrBeRtSJiBW4GdhtPKqVcSqkapVS7Uqod2AvcpJTSCXqNRqNJI0kLvlLKD9wOPAIcAn6slOoSkc+KyE3Jvr5Go9FozMEMlw5KqT3AnkWPfWaJY19vxjk1Go1GEx+60laj0WjyBFMifI0mGQJBxYBrjp7xWXrHZ5n2Brj18jYKC3Q8YibBoKKr38225spML0WTIbTgazLCwdMuvvDwYXrGZzk9MYc/ePbktY0N5VzZUZOh1eUmv3plkD//3gt87T3buWFbY6aXo8kAWvA1GeF//tDP08fHuGFrA2/e1kir006L046t0MLbv/4M3WOzXNmR6VXmFi/2TgJw58OHuXZTPdZC/Qkq39CCr8kIAy4PzY4SvvLu7Wc9HggqrAUWusdnMrSy3KXrtJsyWyHdY7Pct7ebD712baaXpEkz+i1ekxEG3R7qK4rPebzAIjQ7Sugdn83AqnIXpRRd/S7evK2Rq9bX8B+PHcU1O5/pZWnSjBZ8TUYYdHlorDxX8AFaq+10j2nBN5N+l4eJ2Xm2rqng72/chNszz3/+5miml6VJM1rwNWlHKcWg20NDlAgfoNVpp2dsFqVU1Oc18dN12gXA5qZKNjVW8I4dzdz7zCl69BtrXqEFX5N2Jmbn8fmDNCwV4TvtTHn9uOZ0ysEsuvrdWAQ2NZYD8DdvOp9Ci4UvPHI4wyvTpBMt+Jq0M+CaA1g6peO0A+i0jol09bs4r7YMuzXk06ivKOa2q8/jlwcGeL57IsOr06QLLfiatDPk9gBE3bSFUA4foEdv3JpGV7+brU0VZz1229XnUVtu43O/fEWnz/IELfiatDPgCgl+Y2VJ1OeNCF8LvjmMTXsZcHnY0nR2hW2prZBPvmkDL/RMsuflwQytTpNOtOBr0s6Qy4NFoKbMGvV5u7WQmjKb3lA0ia5+NwBbFkX4AG/f0cL59eV84eHDeP2BdC9Nk2a04GvSzoDLQ1158bK9ctqq7TrCN4mD/SGHzuIIH0J1D39340Z6xmf5xR8G0r00TZrRgq9JO4Nuz5IOHYNWpxZ8s+jqd9PsKKHSXhT1+avX11JqLeBA32R6F6ZJO1rwNWln0LW0B9+gxWmn3zWHzx9M06pyl67TLrZGie4NLBbh/IZyDg1MpXFVmkxgiuCLyPUickREjonIHVGe/3MReVlEXhKRJ0Vksxnn1WQng66VI/w2px2loG9CR/nJMOWZ59TYbNT8fSSbGis4NOjWbp0cJ2nBF5EC4C7gBmAzcEsUQf+BUmqbUupC4IvAl5M9ryY7mfb6mfL6V07paGumKRhR+9Y1y/fA39RYwZTHT9/EXDqWpckQZkT4lwDHlFInlFI+4H5gV+QBSil3xN1SQIcRecrggiVz5QgftOAny8HTxobtyhE+wKEB97LHabIbMwR/DdAbcb8v/NhZiMhfiMhxQhH+x6O9kIjcJiL7RWT/yMiICUvTrDYMwV8ph19bbsNWaNHWzCTp6ndTU2ajboXv98aGckTQefwcJ22btkqpu5RS64D/BfzvJY65WynVqZTqrK2tTdfSNGlkMFxlu1JKR0S0U8cEuvpdbF2zfHQPoSKsNqddR/g5jhmCfxpoibjfHH5sKe4H3mrCeTVZyGC4j85SbRUi0V785PDMBzg6PL1iOsdgU2MFhwe14OcyZgj+PmC9iKwVEStwM7A78gARWR9x982AbsSdpwy6PTjsRRQXFax4bEs4wtfOkcQ4MjhFIKiWtWRGsqmxgu7xWWa8/hSvTJMpkhZ8pZQfuB14BDgE/Fgp1SUinxWRm8KH3S4iXSLyEvAJ4NZkz6vJTkKWzOg9dBbT5rQz6wswOu1L8apykzMtFWIXfKXg8KDO4+cqpsy0VUrtAfYseuwzEbf/yozzaLKfgWUmXS0m0ppZW25L5bJykoP9LsqLC2lxxvYGa/TKPzTgZkebI5VL02QIXWmrSStDS8yyjUarsxRAz7dNkK5+N1uaKhCRmI5fU1VCRXGh3rjNYbTga9KG1x9Kz8Qa4Tc7QpGpHoQSP/5AkMMD7pjz9xByRm1srNCCn8NowdekjWG3F1jZg29QXFRAQ0WxduokwPGRGbz+IFtisGRGsrmxgsODUwSDeqM8F9GCr0kbsXrwI2mtttMzPpOqJeUsXcu0RF6OTY3lzPoC+k02R9GCr0kbAzG2VYhEF18lxsHTboqLLJxXUxrX1+kWC7mNFnxN2hgKC359HILf5rQz5PbimdfTmOKhq9/FxoaKZYfMRGNDfTkW0YKfq2jB16SNAZeHUmsB5bbY3cCGNVM7dWInGFS80u+OqaXCYoqLCjivtoxXdE+dnEQLviZtDLrnqK8sjtkmCKFqW9BdM+Ohd2KWKa8/7vy9wcaGch3h5yha8DVpYzCOoisDo02ytmbGjlFhG48lM5JNjRWcnpzDNTdv5rKyiu88dZIfPteT6WWYjhZ8TdoIjTaMrerTwFlqpdRaoCP8ODg04KbAIqyvL0vo6zeHN24P53GUf89TJ/na48czvQzT0YKvSQuBoGJ4yktDZXwtEkSE1upSncOPg/5JD7Vltpga1EXDcOrka08drz/A6Yk5esZnGQ5biXMFLfiatDA27cUfVDE3Touk1VlCtxb8mBme8lBXkXjvofoKGw57Ud7m8XvHZzHqzvZ3T2R2MSajBV+TFgZinHQVjVanPfRHqKs/Y2JkyktdefzfZwMRCQ01z1PBPzFyptBv/ykt+BpN3BhVtvFu2gK0Vpfi9QcZnvKavaycZHjKm1SED6G0zpGhUD/9fOPkaEjwNzVWsL97PMOrMRct+Jq0sDDLNhHB19bMmPH5g4zP+KhLsp30psYKPPPBBfHLJ06OzlBdauWajbV09buZ9eXOQBhTBF9ErheRIyJyTETuiPL8J0TkFRE5ICKPiUibGefVZA+Dbg9FBYLTbo37a89YM/NPfOJlZDr0KSiZlA6c3Rs/3zgxOsPamlI6250EgoqXeiczvSTTSFrwRaQAuAu4AdgM3CIimxcd9iLQqZR6DfAA8MVkz6vJLgZdoT74FkvsRVcGTVUlWERX28aC4SqpTzKl01FXRqFF8lLwT4UFf3urA5HcyuObEeFfAhxTSp1QSvkIDSnfFXmAUuq3Sinjr3UvoUHnmjxiwDWX0IYtgLXQQmNliU7pxICxz5FshG8rLKCjrizvBH/a62d4ysva2lIqS4rYUFeeU04dMwR/DdAbcb8v/NhSfAh4yITzarKIIbc3ofy9QVu1XVszY2BB8JOM8IGwUye/vPinwnsWRpfRznYHL3RP5MzmdVo3bUXkvUAn8KUlnr9NRPaLyP6RkZF0Lk2TQpRSDLjmEnLoGBjWTM3yDLs9iEB1afx
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 1 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"plot(jnp.linspace(0.01,0.99), model.apply(params, jnp.linspace(0.01,0.99), jnp.ones([1])))"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 12,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"def neural_nbody_ode(state, a, cosmo, params):\n",
|
||
|
|
" \"\"\"\n",
|
||
|
|
" state is a tuple (position, velocities)\n",
|
||
|
|
" \"\"\"\n",
|
||
|
|
" pos, vel = state\n",
|
||
|
|
" \n",
|
||
|
|
" kvec = fftk(mesh_shape)\n",
|
||
|
|
"\n",
|
||
|
|
" delta = cic_paint(jnp.zeros(mesh_shape), pos)\n",
|
||
|
|
" \n",
|
||
|
|
" delta_k = jnp.fft.rfftn(delta)\n",
|
||
|
|
" \n",
|
||
|
|
" # Computes gravitational potential\n",
|
||
|
|
" pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec, r_split=0)\n",
|
||
|
|
" \n",
|
||
|
|
" # Apply a correction filter\n",
|
||
|
|
" kk = sum((ki/pi)**2 for ki in kvec)/4.\n",
|
||
|
|
" pot_k = pot_k *(1. + model.apply(params, kk, jnp.atleast_1d(a)))\n",
|
||
|
|
" \n",
|
||
|
|
" # Computes gravitational forces\n",
|
||
|
|
" forces = jnp.stack([cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i)*pot_k), pos) \n",
|
||
|
|
" for i in range(3)],axis=-1)\n",
|
||
|
|
" \n",
|
||
|
|
" forces = forces * 1.5 * cosmo.Omega_m\n",
|
||
|
|
"\n",
|
||
|
|
" # Computes the update of position (drift)\n",
|
||
|
|
" dpos = 1. / (a**3 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * vel\n",
|
||
|
|
" \n",
|
||
|
|
" # Computes the update of velocity (kick)\n",
|
||
|
|
" dvel = 1. / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * forces\n",
|
||
|
|
" \n",
|
||
|
|
" return dpos, dvel"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 13,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stderr",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"<ipython-input-12-b14f0c49fe59>:17: DeprecationWarning: Calling np.sum(generator) is deprecated, and in the future will give a different result. Use np.sum(np.fromiter(generator)) or the python sum builtin instead.\n",
|
||
|
|
" kk = sum((ki/pi)**2 for ki in kvec)/4.\n"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"# Run the simulation at the same steps\n",
|
||
|
|
"resi = odeint(neural_nbody_ode, [pos_i, vel_i], jnp.array([a_i]+scales), cosmo, paramsi, rtol=1e-5, atol=1e-5)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 14,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAkwAAAJCCAYAAAAybpizAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9a4wlyZUmiH1m/rh+PR43MjMyK6Me0exKdrGGZJHNZLOySdbMdO/OUiPMAjsD7C60AwxGGgGjP4IeWAEaCIJ+z68F9peABnaBkaDBahuS9jE7mG3NgrNSNTlZ3Z3V7GL3ZLGZ2ZyoYkVWZmRmPG749evmbqYfZuf4cbseWQ9mRRHaNCAQEffhD3Nzt8++853vKOccnrVn7Vl71p61Z+1Ze9aetbOb/rwP4Fl71p61Z+1Ze9aetWftl709A0zP2rP2rD1rz9qz9qw9ax/RngGmZ+1Ze9aetWftWXvWnrWPaM8A07P2rD1rz9qz9qw9a8/aR7RngOlZe9aetWftWXvWnrVn7SPaM8D0rD1rz9qz9qw9a8/as/YR7RcCTEqpv66Uelcp9VOl1D94Wgf1rP0Pqz0bR8/aL9qejaFn7Wm0Z+PoWXtSU5/Wh0kplQD4CYB/C8D7AP4AwH/gnPuzp3d4z9r/v7dn4+hZ+0XbszH0rD2N9mwcPWsf1X4Rhul1AD91zt11zjUA/jMA/87TOaxn7X9A7dk4etZ+0fZsDD1rT6M9G0fP2hNb+gt89wUA74n/3wdw40lfKMtLbnPzRThnxasKSunwmme7nHNQymO5JMngnIO1Bs45ONfBs2Lys8pvSWnenlIaSZLxPpzrYG3H36NG2xoybS46RkTbVyPfdeH4rPiu488rlYTv0bEqPrbV/dOx+T7xxw3+vlKKPy+3Kf9PkgxKaaSpg3MKzqWwthP9h5V+0zqFUhbO+evRdcsz+qf/rtYpjo9/jsXikcKna59oHKmsdEmxBQ0HC4UyjJ0lHBJnkTgHpxQcgFZpdFBY1wkAoLItUmcxcRbKOWj0o8EB/D0HhU4pGKUx1Rn3qXMWlTVwob8TZ6HDFhQA7fq/VXhfAfwZB7/NTimo8FlasaiwVe0cEueQwvJ+aDQZpWFV380O/nht2F6C1TGUOgcdtkfbUA7IYP130Z9zq5TfVugPozRapVEmud9fGNvGtnyOjj+boByMc39/LLsGGv6c+BqG86Xjb8J5tfN7B865y2dd+ye0T/ws2khydykt4OhYAaTOIhu571ulYaGw1AmUc5i4DolzyMNneQyJa+PvXP9/pxQqncJCoVUaueswtS2ffxI+SdvRPDYc97Mfr37My2M+qzmsjsnhGO8bvWfhr30CN7gvACBxbjCW/Wf76yrHZRfuH7ldqxQWOuV+TpxDFu6ffhz571ko1DqBpfsMDqVtuQ+oX+meoTFslcKDtsZxZ87lWQQAm0nmLqcFjNLhPuufQ/4YfXPw46BTGkuloQAUtkMCNxhzNG6U6Ge6Zh18H3ZKoUUYR67j72rn+L6S/dMfSz+e5BiQY89h2HXxOJHbj8e9Cs9eC/9co7HnxGf8MxP8/HPh2SWfkdQ6pdDRvkIfWCgswvM8Cc+8zFno6P6gfqx1gi7cdwkc1jvDY46Oje4vuY8PuyWORsbRLwKYPlZTSv19AH8fADY3X8Lf+3tvomnmAdDkmEw2UZZ3sVzuwhigaeawtoXWKbROkWUlrG0xn99D08xxenofztnBZ7TOkCQZ0rSAUgnK8hLStECer/NxtG2Nuj6CtQbWtmjbGta2DG4kQLK2hbWtPAsAQJZNoXUa9qMZ6LXtEs516LomAI0GbVujbWv+Lh1TlpUCICJ8v0bTzPl/zwwDbbtA2y7Rtgs4Z5FlJW/PA8AWSmlonSJJciRJjjQtoHWK6fQi8nwdV67ch1It6vplLJcF5vN74tr007XWCYpiK2yjhrUF5vMDPheaLD2o0vwzmWziH/2jf/OpjJWzmhxDRbGFv/zt/y3o762tL6Bp5qjrQywWj/Dm8pgfQL+Z5Hy9rG1x9/Q+dppTXGtOkAUQWiUZjNI4SAtUOoVRmv//UlpgOr0IAqBNM8et6iFm1qC0LWbtEqV4YAEeZFQ6RWY7lK5D5ixK2yILk/Fevo6jAEDo8wD4ffpe2Rlk4SFShrFI3z1Kchil+fMAMLMGO82p3xYcKp2i0im22xqlbbHd1jBKYz8rkTmL7bYe7B8A98GRzlAlGSqVYD9fw79x6UvQOkHTzNE0c/zR6QM+RqMTGKVxlOT4RrhfaRx2XYMfnN5HaVuUtvXbTnI+NwMFoxN8N5kgSTL883/+D/710xs1q02Oowtpgf/g5b+GN9ev4iAtsDfZwBvze/ibhz8b9D0A7m869leXR5iGPqVG44bG0EKn/HqlEuxNNlCpBFWSYcdU2G3mKG2Laegb6h+6rrQ9AIPxQ6/Jz8ZtZUxFnzlKcizC92n71Ka2xaxrVs5r1jWDbR0lOY/teJ809h4GIEF99/2N5wEA222N7bbGteUx90Er+q/SKY/FvXwdpW1xvTrgsW50gkqnvl/Ea0Zp/MP3/+XHHQ6fuslxdDEt8Le++Ndx67mv4x1T4Ub1AK9VD3F98XCl74+SnO+v0nX46uIR35vU6Pyp3+inDX2zn5WDe/va8ngwhqj/qdH35VgobTsYX/S5uMXfBTA4n4NwfVulkYqxAICPh75D25JjWe43HkfUF/LZ3NI4Wt8B4MfRTrvAl+rDlfNaiO8e6Qzf33yBxxE9uwE/buSx0n7/jx/84ei1/0UA088BvCT+fzG8NmjOud8B8DsA8Pzz33Q0gdHvLHuMLHsA5zIodRlKbTIQAOSkDhALIoGC386EQYPcPoESpXQAR10AS0t0XTMARbQf+o48BmoEiJRKmMGR30+SPLzfv07bi4EGHV8nHk6SnaJjJhzfs0BJOL749XTw0x/zBpQycC4Nn0/4vSTJBsflQVeLorgL51Io9TKMARaLR4Kd60G31slgX5+yfeQ4kmNoNnvJyX6kPtA6Q5aVeAMI17bja0QAmx+uSsMk/kFAkxg9hOgmPUpyvgZ0js5ZDxToARUAEd3kPGnYDllY+WS2G4Ih2zJQkN/hvzVQWQD9ZfLbdNZ/L37QhRuezo0m+sx2KIHBewBgoICRhyPtg7ZJ/fKNcO7WdrzAMDrhY6xUAqMTXA+gJ00LpGmBixe/COcsfnuyCWMqVNVDHi9tW+NW1/B50Jj+Bdonfha9VMwc4Psrg0PZGQBhUlMJX1ugP8fMWWRwK0xUDJZo8jdKw0ChCky3BGEEFGIgRE1uX05WY/uV/1NrR16Xk0K8D8AzP/T5sQlybD9ntcxZpOG8KpVw3wy2E8ZzJvqLPlfpFHcmm3j8/LcAOPze0R6uLY/xxvweb2e7rTGzHd9LT6F94nH0K8XMZbbDTWswC9e3SjIPKMWiCRAgN4w3ee0lqJA/NPHLa53ZDtApMmdXwFLc4rET/y9B0dhYkq/Jz8jvT8W4GjsG+d7YGI7B01nHfhwAJ11vup/akWOl/490htvFFrLdN9Aqjf/u5AP8RXWAv/3op/xslM9HajHTRu0Xme3+AMCvKaV+FX5Q/U8A/O0nf0UJUGOR5z/HZLKHoriLNH2MrttA216AtQWcy+BciuXyBQ5JaZ0w49R1DdK0CNvL+L0kyZnFIYYnTSewtsVyebLCMPnJUHH4jibXtq0Z2NBE0XUNA4whWemBBB2bBHT0PwE6yQjRecj90neV0gzqCIRpnQyAQr+fhLdJYNF/36BpdqBUi7Yt0HULpOnEXwmlkWVl+HwDpVoo9RhZ9hibm29C6xrOpVgsXsGjR3+Nj4OuhQ+REuP0aRnwTzeO/D59XxhTwTkbQHOCPN+AMad8rQBguTyGDQAHAD/MMqVxlE4G7BIxJa8EFpFANV2v7bZG2RnMrOHjoYnR6GSUWSo7w69VXYPMdszgQKcwUMjQsxHQGT8MDBR2lMasXXo2DGo4mdpu5SGVOYvSWcAaZgbGGj8UBYDLbIcjneEoneAoyfFrSR7CuBbGVOi6Bt/QKd5SGsYluB7u5zxf5/E3nV7El77
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 720x720 with 16 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"figure(figsize=[10,10])\n",
|
||
|
|
"for i in range(16):\n",
|
||
|
|
" subplot(4,4,i+1)\n",
|
||
|
|
" imshow(cic_paint(jnp.zeros(mesh_shape), resi[0][1:][::2][i]).sum(axis=0), cmap='gist_stern',\n",
|
||
|
|
" vmax=cic_paint(jnp.zeros(mesh_shape), poss[::2][i]).sum(axis=0).max(),vmin=0)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 15,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.colorbar.Colorbar at 0x7f82486381f0>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 15,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAA2kAAAEiCAYAAABjiJogAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAADRHElEQVR4nOz9aZAl2XUeCJ7ry9tfxIs9I/esrCoU1gJBEKSaFAUSpASxpQa7W8JImuFAbIxVt40WtiRrAdKYDbtnJDPIpo0S2ySTpkykAI1RBNFsaYCWIJIQSAxFigRR2FFr1pJ7xr68fXH3Oz8iMs53Tr7IJSIjMl7G+czCwu9z9+vX73Lc/d7vfMd578lgMBgMBoPBYDAYDEcDwaMugMFgMBgMBoPBYDAYGPaRZjAYDAaDwWAwGAxHCPaRZjAYDAaDwWAwGAxHCPaRZjAYDAaDwWAwGAxHCPaRZjAYDAaDwWAwGAxHCPaRZjAYDAaDwWAwGAxHCNGjLoDBYDAYDIbRxp/6sbJfXUsf+Lyvf6f3m977Dx9AkQwGg4GI9mafjoJtso80g8FgMBgM+8LKWkpf/c3TD3xePP/G9AEUx2AwGHawF/t0FGyT0R0NBoPBYDDsE55Snz3w373gnPtl59ySc+57Q/b9Leecd85Nb6edc+5/cc697pz7jnPufQdwowaDYeTw4PbpKMA+0gwGg8FgMOwLnogy8g/8dx/4NBHdQTlyzp0hoj9JRFfh5z9NRE9t/z1HRP90v/dlMBhGH3uxT0cB9pFmMBgMBoPhSMJ7/7tEtDZk1z8kor9NJN6mPkJE/9Jv4Q+JqOacmz+EYhoMBsNDh/mkGQwGg8Fg2Dcy2hNFaNo59wKkn/feP3+3E5xzHyGiG977bzvncNcpIroG6evbv93aS8EMBsPjgz3ap0cK+0gzGAwGg8GwL3jylPo9UYRWvPfvv9+DnXMlIvq7tEV1NBgMhntiH/bpkcI+0gwGg8FgMOwbh+THcZGILhDR7VW000T0DefcB4joBhGdgWNPb/9mMBiOOY6Kn9mDwD7SDAaDwWAw7AueiNJDeAny3n+XiGZvp51zl4no/d77FefcF4jorzrnPktEP0hEm957ozoaDMcch2WfHjbsI81gMBgMBsO+cRAz1c65XyWiD9KW79p1Ivp57/0v7XL4F4nop4jodSJqE9HPPvQCGQyGkYStpBkMBoPBYDh28EQH4vPhvf+L99h/HrY9Ef2Vh14Ig8Ew0jgo+3TQ2JcEv3Puw865V7cDR37yYRXKYDAYDAbDaCHbw5/BYDAcBkbRNu15Jc05FxLRPyGin6QtmduvOee+4L1/6WEVzmAwGAwGw9GHJz+SPh8Gg+Hxx6jap/3QHT9ARK97798kItp21P0IEe36kRZWyz6anthJu4GIb0LBQB7vYZ0vUyV1fvdjh8GlKp2oA4L720dElIV3L0uWVz+k6j4xf30fKm+9/477dmr/3T7/dV73mCrIYpkOVL3o9gq7cEBf7vSFnEgnJXmjgWofjTR/9/14L7r9dJ0Fifwhi1Qlaqglch/y8boOdZ2lBXlutdQV6Ua7oK6lyhLy+Xq83NF3orv3u7Avd9/R/nD6HeNJXRrbI1lbo7TVukclGgyGxxqeKB29dyCDwXAcMKL2aT8facOCRv7gXS82PUHz/9Nf5fQt+eJeuinf85ISb3fm5Btl2NfHyv1OvaBGTZkurMh0Bi+duQ31Ep+Tx3YnRZLCnky3npZvw+GafHPPr3N+gXpxHlRlWu+/10eaLgu+XOu84vbde2xrXtXZqtxfuSm/rKov8QH+5qIs59vOifTKs/JG83XZfj6Q165fkF8N+gMj7PB2cUX1lYG8z8Kq/IDsTqovKwX9Udet8Qdm1JX72rOynJvvkF+MP/59cg7jt7/xDpHW/dZX+PxYjRc9Bron5LWiDfkhXL1Ccn9bpsM+38ugLPNO8zLdeILr+MYv/CMyGAzHG56ODkXIYDAYEKNqnw5cOMQ59xwRPUdEFE7VDvpyBoPBYDAYDh2OUr3kbjAYDEcCo2mf9vORdl9BI733zxPR80RE+XOnvW/DKkRbVlj7lFyVGMzCioeiemVqxcnFaiVGlaNfk+moLflzSN2LW/LsRGemVnnaZ+WKki6LXvVJypBhSe4LEpm3XjkrLmrqndyfFtT5sLjSOSHP7U+oDqvKqal6/TF5fL8hV43WfmB6Z7u0VBP7GqflalVSUnmNyxvR9zWoyrKMv0a7YlCReeevy/bJYlnu3rhMl2/JFan+mKZmcll642rFqSjLEm3Kc7+2cEak3UBe25cU7xNW1u6g0WpKYiQbMFWrYRvPyHRuU6bjOmfYm5bXSgtqfNVgwGiapcFgOHbwRJSZKTAYDEcQo2qf9vOR9jUieso5d4G2Ps7+AhH9pYdSKoPBYDAYDCOFUZypNhgMxwOjaJ/2/JHmvU+cc3+ViH6TiEIi+mXv/YsPrWQGg8FgMBhGAp5G8yXIYDA8/hhV+7QvnzTv/ReJ6IsPqSwGg8FgMBhGFJlWsjIYDIYjglG0TwcuHHIHoI56U9LPJVN+L6dOru1sb3akVHmlIGUMx/NS2vxcZU2kf+sb7xZp7fPkILv2nHT2yXLqWB0qoKRU9W5Kf7dkTPoZFW9xtWuFPd2HKjdlnQyUL1d+U/kOdUimoez5NaWIqfyncg2Z7k7J4wfj8lqdnirLOm+vvEtWWtRRfn7KF6/9Ttl+xbJs37SrJPyvy/4Qt+Bcpe5YvNUS6aSq/BET6S8XDOT5QSr7Q+sEdx7tM5gqocjxN2S6/aTyvcsp/8We8lEDfy8/J+skW5T34VRexQnZGTpNdfyqrEPsi10ddmJcdvowB31aV4LBYDh2GNWZaoPB8PhjVO3TPaKLGQwGg8FgMBgMBoPhMHH4K2kGg8FgMBgeK3hylNq8r8FgOIIYVftkH2kGg8FgMBj2jVH0+TAYDMcDo2ifDvkjzZGDuE/hCeWQlUh/nVaP/ZBOjW+KfZdXJkV6oymdnDZ70t8mKEufmt60/KIedCEelYoPpn3lzl9clGW5MiOvpXzWXLZ7xyjfUn55kTpWufvEbflDriH93YK+zK99gp2kdGyyuCnzTmSVUSjdxMir+HBpXu2H3tQ+Jcsx/qq6tvJ/G3tBXjwpqsKouF06FlquwdcLe8rPqy99Bp1yMpz81rpId0/KIHytWdkv2/NclqgryxFItzFqzcv93TXpCBhUZWfx67JsQYX3/8TTr8i8n5AN8LVrZ0W6kJN5p0XZ57sn5fAfjPH+dFz2q7go80q64Hw3gobPYDA8XIyqz4fBYHj8Mar2yVbSDAaDwWAw7BOOUj96dCKDwXAcMJr2yT7SDAaDwWAw7AueiLIR9PkwGAyPP0bVPh3uR5rz5ItMo4pjSalK+rI4rTbTuV7blHTGXEHSrzJFKay3JV0u68q8446Slz/R37XYE1OSF3h9eUKkQyXBn12UaV+XFLbuHN930JdUulDR5TRVMlMt1huX+4uriqoJ9Eknq1vQE4nupC/qsuh0IG+TcJIityEHQ9T1Ki0piR0l919UVMw0f3caaH6T8+uNyzrNjynqa19WhI/l8VFL3lhxTe5PC3xvemKmN6lCDVTkfVKk2qcn8y4uyXQ340b5TzcuiH1/8qykPwaBzLvdlQ2az8n76oUqfAOGikhVuIYVFa8BbysdPQqBwWB4+BhFOpHBYDgeGEX7ZCtpBoPBYDAY9gXvR5NOZDAYHn+Mqn2yjzSDwWAwGAz7RjaCM9UGg+F4YBTtk32kGQwGg8Fg2Be21NNGb6baYDA8/hhV+3SoH2ku9JSvsmNTe036mbmckv1WPjSIMJS+PqW89FFbXauIdFCQeaUF6fsTrLOkeFaTx26sl0Xad9W5LZl2A/W1ruXMN0HqXPmBFZRPWf0JuT8tyv0uURLvKrQAyujHDXWuChVATuYVKjc97YNWXpBtkOb4/OKS8neSbnnUOCPL2Z+QeZWvy/2Tr8j9UWf3a3vZHLT6Ltl++r5KS8qHUI1jl8p7QV8+HbYgGVPy/zV1sQ1ZEWOvKn+4jvJpK/N9NZfkfXy+/axIZyu
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 1080x360 with 5 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"figure(figsize=[15,5])\n",
|
||
|
|
"subplot(131)\n",
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), resi[0][0]).sum(axis=0))\n",
|
||
|
|
"subplot(132)\n",
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), resi[0][1]).sum(axis=0)); colorbar()\n",
|
||
|
|
"subplot(133)\n",
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), poss[0]).sum(axis=0)); colorbar()"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 16,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"ref_pos = jnp.stack(poss, axis=0)\n",
|
||
|
|
"ref_vel = jnp.stack(vels, axis=0)\n",
|
||
|
|
"\n",
|
||
|
|
"@jax.jit\n",
|
||
|
|
"def loss_fn(params):\n",
|
||
|
|
" res = odeint(neural_nbody_ode, [pos_i, vel_i], jnp.array([a_i]+scales), cosmo, params, rtol=1e-5, atol=1e-5) \n",
|
||
|
|
" distance = jnp.sum((res[0][1:] - ref_pos)**2, axis=-1)\n",
|
||
|
|
" vel = jnp.sum((res[1][1:] - ref_vel)**2, axis=-1)\n",
|
||
|
|
" w = jnp.where(distance < 100, distance, 0.)\n",
|
||
|
|
" wv = jnp.where(distance < 100, vel, 0.)\n",
|
||
|
|
" return jnp.mean(w) #+ 0.001*jnp.mean(wv)\n",
|
||
|
|
"\n",
|
||
|
|
"@jax.jit\n",
|
||
|
|
"def update(params, opt_state):\n",
|
||
|
|
" \"\"\"Single SGD update step.\"\"\"\n",
|
||
|
|
" loss, grads = jax.value_and_grad(loss_fn)(params)\n",
|
||
|
|
" updates, new_opt_state = optimizer.update(grads, opt_state)\n",
|
||
|
|
" new_params = optax.apply_updates(params, updates)\n",
|
||
|
|
" return loss, new_params, new_opt_state"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 17,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"losses = []"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 18,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"import optax\n",
|
||
|
|
"from tqdm import tqdm\n",
|
||
|
|
"learning_rate=0.01\n",
|
||
|
|
"optimizer = optax.adam(learning_rate)\n",
|
||
|
|
"opt_state = optimizer.init(params)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 45,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stderr",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
" 23%|██▎ | 23/100 [03:31<11:48, 9.20s/it]\n"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"ename": "KeyboardInterrupt",
|
||
|
|
"evalue": "",
|
||
|
|
"output_type": "error",
|
||
|
|
"traceback": [
|
||
|
|
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
|
||
|
|
"\u001b[0;31mKeyboardInterrupt\u001b[0m Traceback (most recent call last)",
|
||
|
|
"\u001b[0;32m<ipython-input-45-26234aa997fa>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mstep\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtqdm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m100\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0ml\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mopt_state\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mupdate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mopt_state\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 3\u001b[0m \u001b[0mlosses\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ml\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
||
|
|
"\u001b[0;32m~/.local/lib/python3.8/site-packages/haiku/_src/data_structures.py\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m(treedef, leaves)\u001b[0m\n\u001b[1;32m 260\u001b[0m \u001b[0mFlatMapping\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 261\u001b[0m \u001b[0;32mlambda\u001b[0m \u001b[0ms\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0ms\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_leaves\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ms\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_structure\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 262\u001b[0;31m lambda treedef, leaves: FlatMapping(FlatComponents(leaves, treedef)))\n\u001b[0m\u001b[1;32m 263\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 264\u001b[0m \u001b[0;31m# _ _ _\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
||
|
|
"\u001b[0;31mKeyboardInterrupt\u001b[0m: "
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"for step in tqdm(range(100)):\n",
|
||
|
|
" l, params, opt_state = update(params, opt_state)\n",
|
||
|
|
" losses.append(l)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 48,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"[<matplotlib.lines.Line2D at 0x7f7e846b6d00>]"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 48,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAXoAAAD4CAYAAADiry33AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAAj0ElEQVR4nO3deXSc9X3v8fd3RpstjfZdtiXZlnfjHTBbAimFQApNkyYmaRqaNuTkhl6anrSlN+emKe29t23Skp6WtoekCU1aYiiBhAaC4xISCBjwjjdsy7JsSbZWa7P25Xf/mEdGGNmW7Bk9M6PP65w5mmeZma+eY3/m0e/3e36POecQEZHEFfC7ABERiS4FvYhIglPQi4gkOAW9iEiCU9CLiCS4JL8LOF9+fr6rqKjwuwwRkbiyc+fOVudcwUTbYi7oKyoq2LFjh99liIjEFTM7caFtaroREUlwCnoRkQSnoBcRSXAKehGRBKegFxFJcAp6EZEEp6AXEUlwCnrAOceWA40cber2uxQRkYib8UHf3T/E/d/fzee+t5OHfnzQ73JERCIu5q6MnU77Gzr5wuO7qG/vo6owgx217QwOj5KSNOO//0QkgczIRHPO8W+v1fIb//Qag8OjbL7vWr5022L6hkbYW9/hd3kiIhE1487ou/qH+JOn3uIn+xu5ZUkhX//NVeSmp9DZO4QZvFbdxoaKXL/LFBGJmBkV9G/Vd/CFx3dxqqOf/3XHEn7vhvkEAgZA1uxklpVksq2mlQeo8rlSEZHImTFNN68da+Uj//waIyOOJz+3kftuWnAu5MdctyCPXSc66B8a8alKEZHImzFBv/VgE0mBAM8/cCPrynMm3GfjgjwGR0bZdaJ9mqsTEYmeGRP0ta09VOSnkz075YL7bKjIJRgwttW0TWNlIiLRNWOC/kRbL5X5sy+6TygtmZVlWbx2TEEvIoljRgT98MgoJ8/0Up6Xfsl9Ny7IY29dBz0Dw9NQmYhI9M2IoD/V0c/wqKNyEkF/3YI8hkcd22vPTENlIiLRNyOCvratB4DyvIs33QCsL88lOah2ehFJHDMq6CvzL31GPyslyJq5ObyudnoRSRAzI+hbe5mVHKQglDqp/a9dkMe+hk66+oeiXJmISPRNKujN7HYzO2xm1Wb24ATbHzazPd7jiJl1jNs2z8x+amaHzOygmVVErvzJqW3roTxvNmZ26Z2BjfPzGHXwZo3a6UUk/l0y6M0sCDwCfBBYBtxjZsvG7+Oc+6JzbrVzbjXwD8DT4zZ/F/iac24pcDXQHKHaJ622rWdSzTZj1szLJjUpoGGWIpIQJnNGfzVQ7Zyrcc4NApuBuy+y/z3A9wG8L4Qk59xWAOfcWedc7xXWPCXDI6PUTXJo5Zi05CDrynPUISsiCWEyQV8G1I1brvfWvYeZlQOVwM+8VYuADjN72sx2m9nXvL8Qzn/dfWa2w8x2tLS0TO03uITTnf0MjbhLXix1vo3z8zh0uov2nsGI1iMiMt0i3Rm7CXjKOTc2K1gScCPwJWADMB+49/wXOecedc6td86tLygoiGhB7wytnPwZPcB1C/MAeF1n9SIS5yYT9A3A3HHLc7x1E9mE12zjqQf2eM0+w8APgbWXUedlq22d/NDK8a6ak83slKCab0Qk7k0m6LcDVWZWaWYphMP82fN3MrMlQA6w7bzXZpvZ2Gn6LcC03pi1tq2XtOQAhZMcWjkmORhgQ0WuOmRFJO5dMui9M/H7gS3AIeBJ59wBM3vIzO4at+smYLNzzo177QjhZpsXzWwfYMA3I/kLXEptaw8VeemTHlo53sYFeVQ3n6W5uz8KlYmITI9J3WHKOfc88Px5675y3vJXL/DarcBVl1nfFatt66GqMHRZr904f6yd/gx3rSqNZFkiItMmoa+MHRl11J3po3yKI27GLC/NJJSWxLZjrRGuTERk+iR00J/q6GNwZHRSs1ZOJCkY4JrKXLapnV5E4lhCB/2JtvC1WVMdWjnexgX51Lb1cqqjL1JliYhMq4QO+uNTmLXyQsba6XVWLyLxKqGD/kRrz2UNrRxvSXGInNnJGk8vInEroYO+tq2H8tx0AoGpD60cEwgY187PY9uxNsaNHBURiRsJHvS9VFzmiJvxNi7Io6Gjj7ozaqcXkfiTsEE/Muo42dZLxRV0xI7ZUJELwO669it+LxGR6ZawQX+6Mzy0suIKOmLHLCzMICUpwIFTXRGoTERkeiVs0L8ztPLKm26SgwEWF4U4cKrzit9LRGS6JWzQH7/MWSsvZEVZJvsbutQhKyJxJ2GD/kRbD6lJAYpCaRF5v2WlWXT2DdGgC6dEJM4kbNAfb+2lPG/2FQ2tHG9FaSaA2ulFJO4kbNCfaOuJyIibMUuKMwkYHGhQO72IxJeEDPrRUceJM70RGXEzZlZKkAUFGTqjF5G4k5BBf7qrn8Hh0Yie0QOsKMtiv0beiEicScigP+GNuKmIwNDK8ZaXZtLUNUDr2YGIvq+ISDQlZNCPzVoZyaYbgOWlWYA6ZEUkviRk0J9o6yUlKUBxZmSGVo5Z5o282a8OWRGJIwkZ9MdbeyjPjdzQyjFZs5KZmzuLgzqjF5E4kpBBf6KtJ+LNNmNWlKpDVkTiS8IF/eio40Rbb8Q7YscsL83kRFsvXf1DUXl/EZFIS7igb+zqZ2A4MrNWTmSsQ/aQmm9EJE4kXNDXjo24ifAY+jHLy7wOWQW9iMSJxAv61vD0xNE6oy8MpVEQStWUxSISNxIu6E+09ZCSFKAkwkMrx1tRmqmRNyISNxIu6KM1tHK85aVZHG0+S//QSNQ+Q0QkUhIu6E+09VIepfb5MctLMxkZdRxu7I7q54iIREJCBf3oqKO2rSdqQyvHrCgLj7zReHoRiQcJFfRN3dEdWjlmTs4sMtOSNOeNiMSFhAr6cyNuotx0Y2YsL81S0ItIXEisoD83a2V0m24g3E7/9ukuhkdGo/5ZIiJXIuGCPiUYoCRrVtQ/a3lZJgPDoxxr6Yn6Z4mIXInECvrWHublzSYYxaGVY1Z4UyFoymIRiXWTCnozu93MDptZtZk9OMH2h81sj/c4YmYd523PNLN6M/vHCNU9oWhOZna++QUZpCUH1E4vIjEv6VI7mFkQeAS4FagHtpvZs865g2P7OOe+OG7/3wfWnPc2fwG8HJGKL2BsaOX1C/Oj+THnBAPG0pJMTYUgIjFvMmf0VwPVzrka59wgsBm4+yL73wN8f2zBzNYBRcBPr6TQS2nuHqB/KPpDK8db7k2FMDrqpu0zRUSmajJBXwbUjVuu99a9h5mVA5XAz7zlAPC3wJcu9gFmdp+Z7TCzHS0tLZOp+z0KQqn84o/ez50rSy7r9ZdjeWkW3QPD1LX3TttniohMVaQ7YzcBTznnxiaB+R/A8865+ou9yDn3qHNuvXNufUFBwWV9cDBglOelk5ueclmvvxzvdMiqnV5EYtdkgr4BmDtueY63biKbGNdsA2wE7jezWuDrwG+b2V9dRp0xaVFxBkkBUzu9iMS0S3bGAtuBKjOrJBzwm4BPnL+TmS0BcoBtY+ucc58ct/1eYL1z7j2jduJValKQhYUZGnkjIjHtkmf0zrlh4H5gC3AIeNI5d8DMHjKzu8btugnY7JybUT2TK8qyOHCqkxn2a4tIHJnMGT3OueeB589b95Xzlr96ifd4DHhsStXFgeWlmTy1s57m7gGKonizExGRy5VQV8b64dyUxbpCVkRilIL+Ci0tycQMtdOLSMxS0F+hjNQkKvLSNfJGRGKWgj4Clpdmsq9eQS8isUlBHwFr5+VwqrOfxs5+v0sREXkPBX0ErC3PAWDXyXafKxEReS8FfQQsK8kkNSnAzhMKehGJPQr6CEhJCrBqTraCXkRikoI+QtaUZ3PgVCf9QyOX3llEZBop6CNk3bwchkacLpwSkZijoI+QsQ5ZNd+ISKxR0EdIfkYq5XmzNfJGRGKOgj6C1s3LYeeJDs1kKSIxRUEfQWvLc2g9O0DdmT6/SxEROUdBH0Fr5+nCKRGJPQr6CFpcHCI9JagOWRGJKQr6CAoGjDXzchT0IhJTFPQRtnZeNm83dtEzMOx
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 1 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"plot(losses[10:])"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 71,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"import pickle\n",
|
||
|
|
"pickle.dump( params, open( \"node.params\", \"wb\" ) )"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 49,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"res = odeint(neural_nbody_ode, [pos_i, vel_i], jnp.array([a_i]+scales), cosmo, params, rtol=1e-5, atol=1e-5) "
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 50,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.colorbar.Colorbar at 0x7f7e845df820>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 50,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAA2kAAAEiCAYAAABjiJogAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAADWt0lEQVR4nOz9aZAl2XUeCJ7ry9vfixcv9tyzsqpQKBRQIASC4JCiQIKSQDYpsHskDMkZCmJjrFo2lKiWZC2A6raBekwcg6xllCijRpoakU1QRhFEs8UGWoJIQiDRJCUCBKqw1ppZlWtkxh5v33y58yMi43zn5ovMjIhc4mWczyws/D53v379Lsfd7/3Od4y1lhQKhUKhUCgUCoVCcTjgPewCKBQKhUKhUCgUCoWCoR9pCoVCoVAoFAqFQnGIoB9pCoVCoVAoFAqFQnGIoB9pCoVCoVAoFAqFQnGIoB9pCoVCoVAoFAqFQnGIoB9pCoVCoVAoFAqFQnGIEDzsAigUCoVCoRhv/MXvL9r1jWTP573wzcHvWms/cB+KpFAoFES0P/t0GGyTfqQpFAqFQqE4ENY2Evry757Y83nhwhvT96E4CoVCsYP92KfDYJv0I02hUCgUCsUBYSmx6cMuhEKhUIzAeNon/UhTKBQKhUJxIFgiSsk+7GIoFArFLRhX+6TCIQqFQqFQKBQKhUJxiKAraQqFQqFQKA6MlMaPTqRQKI4GxtE+6UeaQqFQKBSKA8GSpcSOH51IoVA8+hhX+6QfaQqFQqFQKA6McfT5UCgURwPjaJ/0I02hUCgUCsWBYIkoGcOXIIVC8ehjXO2TfqQpFAqFQqE4MMZxplqhUBwNjKN90o80hUKhUCgUB4IlGkufD4VC8ehjXO3TgST4jTEfMMa8Zoy5YIz52L0qlEKhUCgUivFCuo8/hUKheBAYR9u075U0Y4xPRP+ciP48EV0joq8YYz5rrX35XhVOoVAoFArF4YclO5Y+HwqF4tHHuNqng9Ad30NEF6y1bxIRGWM+RUQfJKJdP9L8ctEG05M7aRMZsd+L5PEW1vlSp6TG7n7sKJjEScfOAd7d7SMiSv3blyXNOj8kzn1i/u59OHm7+2+5b+Psv93nv5vXHaYK0lCmPade3Pby+3DAUO60uYxIxwV5o57TPi6S7O3347247efWmRfLH9LAqUQXzhK59fl4tw7dOkty8txyoS/SrW7OuZZTFp/Pd8fLLX0nuH2/84dy9y3tD6ffMp6cS2N7xBsblHQ6d6hEhULxSMMSJeP3DqRQKI4CxtQ+HeQj7TgRXYX0NSL6rttebHqSFv7Hv8HpG/LFvXBdvufFBd7uzck3Sn/oHiv3G+cFNWjLdG5NplN46czUnZf4jDy2XxNJ8gcy3XlSvg37G/LNPbvJ+XnOi3NUlml3/50+0tyy4Mu1m1fYvX2P7Sw4dbYu95euyy+r8st8gL2+LMv5ltMivfasvNFsU7af9eS1m2flV4P7geH3eDu/5vSVSN5nbl1+QPZrzpeVA/ejrl/lD8ygL/d1Z2U5G0/LL8Yf+A45h/H7Lz4t0m6/tSU+P3TGizsG+vPyWkFdfgiXL5Pc35Vpf8j3EhVl3klWpluPcR0v/sI/JYVCcbRh6fBQhBQKhQIxrvbpvguHGGOeI6LniIj8qer9vpxCoVAoFIoHDkOJu+SuUCgUhwLjaZ8O8pG2SEQnIX1i+zcBa+3zRPQ8EVH29Alru7AK0ZUV1j0uVyWiWVjxcKheqbPiZEJnJcYpx7Aq00FX8ueQuhd25Nmxm5mzytM9JVeU3LK4qz5xETIsyH1eLPN2V87yyy71Tu5Pcs75sLjSm5fnDiedDuuU06XqDSvy+GFLrhptfOf0znZhpSr2tU7I1aq44OQ1IW/Eva+oLMsy8Trtiqgk885ek+2ThrLcgwmZLt6QK1LDikvN5LIMJpwVp7wsS9CQ535l6aRIm0he2xYc3iesrN1Co3UpiYFswMRZDas/JdOZhkyHTc5wMC2vleSc8VWFAePSLBUKxZGDJaJUTYFCoTiEGFf7dJCPtK8Q0RPGmLO09XH240T0k/ekVAqFQqFQKMYK4zhTrVAojgbG0T7t+yPNWhsbY/4GEf0uEflE9CvW2pfuWckUCoVCoVCMBSyN50uQQqF49DGu9ulAPmnW2s8R0efuUVkUCoVCoVCMKVJXyUqhUCgOCcbRPt134ZBbAHU0mJJ+Lqnj93L82MbOdqMnpcpLOSljOJGV0uanSxsi/Xsvvl2kXZ8nA9l156SzT5pxjnVDBRQcVb3r0t8trkg/o/wNrnZXYc/tQ6Xrsk4ix5cr23B8h3ok01D27IajiOn4T2VaMt2fksdHE/JavYFTlk3eXntGVlrQc/z8HF+87ttk++WLsn2TviPhf032h7AD5zrqjvkbHZGOy44/Yiz95bxInu8lsj905rnzuD6DiSMUOfGGTHcfd3zvMo7/4sDxUQN/Lzsn6yRdlvdhnLzyk7Iz9NrO8euyDrEv9t2wExOy0/sZ6NNuJSgUiiOHcZ2pVigUjz7G1T7dIbqYQqFQKBQKhUKhUCgeJB78SppCoVAoFIpHCpYMJTrvq1AoDiHG1T7pR5pCoVAoFIoDYxx9PhQKxdHAONqnB/yRZshA3Cd/3nHIiqW/TmfAfkjHJxpi36W1mkjX29LJqTGQ/jZeUfrUDKblF3XUh3hUTnww11fuzLllWZbLM/Jajs+aSXfvGMUbjl9e4BzruPuEXflDpiX93byhzK87z05SbmyysC3zjmWVkS/dxMg68eGSrLMfelP3uCzHxGvOtR3/t8pX5cXjvFMYJ26XGwst0+Lr+QPHz2sofQaN42RY+/qmSPePySB8nVnZL7sLXJagL8vhSbcx6izI/f0N6QjolWVnsZuybF6J9//gk6/KvB+TDfCVq6dEOpeReSd52ef7x+Twjyq8P5mQ/SrMy7ziPjjfjaHhUygU9xbj6vOhUCgefYyrfdKVNIVCoVAoFAeEocSOH51IoVAcBYynfdKPNIVCoVAoFAeCJaJ0DH0+FArFo49xtU8P9iPNWLJ5plGFoaRUxUNZnE6X6VyvNySdMZOT9KvUoRQ2u5Iul/Zl3mHPkZefH+5a7MkpyQu8tjop0r4jwZ+ek2nblBS2/hzftzeUVDrfocu5VMnUabHBhNyfX3eomkCfNLK6BT2R6Fb6olsWN+3J2yScpMjU5WAI+tZJS0piz5H7zztUzCR7expotsH5DSZknWYrDvV1KCvChvL4oCNvLL8h9yc5vjd3YmZQc0INlOR9UuC0z0DmnV+R6X7KjfKfF8+KfX/hlKQ/ep7Mu9uXDZrNyPsa+E74BgwVkTjhGtaceA14W8n4UQgUCsW9xzjSiRQKxdHAONonXUlTKBQKhUJxIFg7nnQihULx6GNc7ZN+pCkUCoVCoTgw0jGcqVYoFEcD42if9CNNoVAoFArFgbClnjZ+M9UKheLRx7japwf6kWZ8S9kyOzZ1N6Sfmck4st+ODw3C96WvTyErfdTWN0oi7eVkXklO+v54mywpnlblsfXNokjbvnNuR6ZN5Hytu3LmDZA6d/zAco5PWfMxuT/Jy/0mdiTendACKKMftpxznVABZGRevuOm5/qgFZdkGyQZPj+/4vg7Sbc8ap2U5RxOyryK1+T+2qtyf9Db/dpWNgetPyPbz72vworjQ+iMY5PIe0FfPjdsQVxx5P+rzsXqsiIqrzn+cD3Hp63I99Vekffxme6zIp2uOZ3J8VGLarIspuv4wy3xjaeh2EWDWcePLwN5G8dBUKFQHEGMJ51IoVAcBYynfdKVNIVCoVAoFAfCuKqnKRSKRx/jap/0I02hUCgUCsWBkWhge4VCcUgxjvZJP9IUCoVCoVAcCJbMWPp8KBSKRx/jap8e6EeajT0arrIfmutPRaHjdwR+Z5mCDNJVK3ZFOkqkf02x3Bfp0Jc+Nf28dMjqtdifJ3dR+va4vli9+cTZL++juOjEmNqQ1TyYZB8ep1gUFZ14Yctyf+zsTx1fL9cfqw0x24qXZDmGE25ML+e+Isf3zvFh60/KshhovkzT8YcqyGODnswrqri+dY5
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 1080x360 with 5 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"figure(figsize=[15,5])\n",
|
||
|
|
"subplot(131)\n",
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), res[0][0]).sum(axis=0))\n",
|
||
|
|
"subplot(132)\n",
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), res[0][1]).sum(axis=0));colorbar()\n",
|
||
|
|
"subplot(133)\n",
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), ref_pos[0]).sum(axis=0));colorbar()"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 74,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.lines.Line2D at 0x7f8267c7c4c0>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 74,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAWoAAAD8CAYAAABekO4JAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOydd3hU1daH3+npvfdAQgoJLfTei1TpvSpFFAQFFRFFsaCoKCK9d+mINOkgvQSSkIQkhPTey/Q53x/DTYwBRcXvqnfe58mDzt5nzz4nyS9r1l5FJAgCJkyYMGHi74v4v70BEyZMmDDx65iE2oQJEyb+5piE2oQJEyb+5piE2oQJEyb+5piE2oQJEyb+5piE2oQJEyb+5kj/ikWdnJwEPz+/v2JpEyZMmHimPMirAKCOs+V/dR83b97MFwTB+XFjf4lQ+/n5cePGjb9iaRMmTJh4pgxbdRmAXVNa/Vf3IRKJUp40ZnJ9mDBhwsTfHJNQmzBhwsTfHJNQmzBhwsTfHJNQmzBhwsTfHJNQmzBhwsTfHJNQmzBhwsTfnL8kPM+ECRP/fxhUKjSpqegLCjBv0ACx5X83HtjEs8ck1CZM/APR5uRSvGsnJYe+R5ueXvW6SKHAqn17bPr1xbprV0Qi0X9xlyaeFU8t1CKRSALcADIEQejz123JhAkTT0J17x4Fa9dReuIE6PVYtm2L7cDnkfv6IrGxofzsOcpOnKDsxx+xbN0Kt/ffR+7l9ZvrVly+TNnJU1Reu4bY1ga3+fMxCw7+f7gjE0/D77GoZwKxgM1ftBcTJkw8AWVkJPkrVlJ+7hxiKyscRo3CftRI5D4+NeZZtWuH67y3KN69m9zPlvCgbz+cZ8zAYfQoRDJZrXV1+flkL/qQsmPHEJmbY9G4Mar4eJIHD8H55ek4TZ36/3WLJn6FpxJqkUjkBfQGPgRm/6U7MmHCBAC6vDxKjx6l5PvDqKKikNjZ4fzqTOxHjUJibf3E60QSCfbDh2PVoQNZ771H7uLFlOzbh+v8+Vg0b4ZIJEJfWkrR9h0UbNiAUFmJ86uv4jBxAmK5HF1RETkffEDe0q+Q+/pi06vX/+Ndm3gcT2tRLwXmAk/+6TBhwsTvQl9ejr6oCH1xCfqSEnT5eWjTM9Cmp6NJSUF55w4YDChCQnB9603sBg/+XQeFMnd3vFeupPzUKXI++pjUceMQW1gg8/ZGm5GBobwcyw7tcX3jDRR16lRdJ7W3x+PTT9GkZ5C98H0smjVD6uT0VzwCE0/Jbwq1SCTqA+QKgnBTJBJ1/JV5k4HJAD6/+DhmwsT/OoJeT+XNm6ju3kV5NwpldBS6zKzaE0UipC4uyLy8cJz8IrZ9+qAICPjD7ysSibDu2hXLNm0oOXgIdVIS2tRUFEH1cJwwAbOQkMdfJ5Xi8fFHJD8/kKx338Prm2Wmg8n/Ik9jUbcB+olEoucAM8BGJBJtFQRh9M8nCYKwGlgN0LRpU1NrcxMmAMFgoOz4cfKWL0eTmASAzMsL84YNMRs+AqmTExI7OyR2tkjs7ZF5eCBWKJ75PsTm5tgPH/a7rlHUrYvzzJnkfvYZFRd/wqpd22e+LxNPx28KtSAIbwFvATyyqF//pUibMGGiJoLBQNmPJ8n/5hvUCQnIA+ri8dlnWLZtg9Te/i95T6VOibnU/JmuaT9mNIXbtpK3bBmWbduYrOr/EqY4ahMmnjH6sjIyZr9GxYULyP398ViyBJtePRFJJM/0fUo1pRxLPsbNnJtE5kaSWZGJq4UroY6hRLhGMKTeECxkFn/qPcRyOU7TppH9zgLKz57FulOnZ7R7E7+H3yXUgiCcBc7+JTsxYeJfgCY9g/RpU1EnP8T1nfnYDx/+zAU6rSyNbbHb2JewD6VOibO5M41dGjMgcAAppSnE5MdwJu0MW+5tYXbEbHr59/pNS9ggGLiWfY3DSYfRGrQMqTeECNcIRCIRdgMGULB6DXnLlmHVsaPJqv4vYLKoTZh4RigjI0mb/jKCVovP2jVYtmz5TNe/m3eXjTEbOZV6CjFievn3YnToaEIcQmqJ562cW3xy7RPeuPAGu+J3Mb/lfALtAx+77qXMS3x09SNSSlOwklkhFok5knyERs6NWNltJZYyS5ymv0TWm29RfuYs1p1NVvX/NyJBePbnfk2bNhVMrbhM/C9ReuQImW++hdTNDe+VK2qEu/0Zsiuy+SnjJw4lHeJW7i2s5dYMrTeUEcEjcLV0/dVr9QY9BxIP8OWtL6nQVDCm/hiGBQ3D08oTQRCIKYhh873NHE0+ip+NHy81eolO3p0QEDiYeJBPrn1CZ5/OfN7hc9DrSerZC4mDA367dv6rrOq/USuum4IgNH3cmMmiNmHiT1K8dx9Zb7+NeUQEXt8se+rDQkEQKNeWU6AsoFBVSIGqgEJlIYWqQvKUedzOvU1icSIAnlaevNHsDQYGDnxqv7NELGFQvUF09unMFze/YEP0BjZEb8Db2hsRIlLLUpGL5UxrOI1J4ZNQSKqjTYYHD0epU/LFzS/YfG8z4+qPw/HFF8l+910qLl3Cqk2b3/+gTPxhTBa1CRN/gtJjx8mYPRvLli3xWvHtr4bWCYLA/aL7XM68zKXMS0TmRaLUKR8711ZhS4hDCG0929LGow117er+aSs2uSSZS5mXuJJ5Ba1BS3e/7nTx6YKtwvaJ+519djZn0s6wt99e/C28SerWHbm3N75bt/ypvfydMFnUJkz8iym/cIGMOXMwb9gQr2+WPVGkdQYd38V/x9qoteQp8wAIsAtgQMAAPK08cTBzwNHMEQdz4792ZnbIxLXrcvxZ/G398bf1Z1TIqKeaLxKJeLfVu1zNusrSm0tZ1mUZjpMmkfPRR1RcuYplyxbPfI8mHo9JqE2Y+ANU3rhB+iszUAQE4L1qJWKLx7sjInMj+fDqh8QVxtHMrRmvNH6FVh6tcLN0e7YbMhigOAXkVmBuD5Jn86ttZ2bHpPBJLL21lOvZ14kYNpSC9evJW7oUix3b/1W+6r8zJqE2YeJ3ooyJIW3qNGTu7visXYPEpnZByWJVMV/c/IL9iftxsXBhSYcldPft/myFTa+FhBMQf9T4b3lO9ZidD7R8CSLGg+zPJcGMChnFjrgdfHHjC7b13obTS9PIXvCuKa76/xGTUJsw8TvQPHxI2gsvIraxxmf9OqSOjrXmpJelM+XHKWSWZzKh/gSmNpz6pxNPaqAqgVub4cpKKE0HhS0EdAH/dqDXQWUBPLwAx96Ei19C+znQdCKIfyWe26A3Cv6VFVCWCZ4RUK8nhA/GTGrGy41f5p2f3uFM2hk6P/88BWvXkbf0K6zat3/mceImamMSahMmnhJdYSGpk6cA4Lt+PTJ391pz4gvjmXpyKhq9hg09N9DIpdGz20DOPbi5ASJ3gKYM/NpB7yUQ0BUkv/RpvwUPL8LZT+DI63B7C/T50ijAvyThRzg+D/Lvg60PuDcwXhu1G0ozoc0M+tTpw5q7a1h5ZyWdvTvjPHMGma+9Tsn+/dgNHvzs7tHEYzEJtQkTT4FBqSRt2jR0OTn4btqI3M+v1pwb2TeYcXoGFjILNvfaTF27un/uTdXlkBUJGbcg/gikXgaJHOo/Dy2ngUfjX7/ery2M+x5i9sOxt2BNF6jTERqOAOd6kHkbYg9D0ilwqAuD10NIf6N/22CAvZPgx3fA0glpo5FMbjCZ+T/N50zaGTo99xxF27aT++VSrHv0+NX62Cb+PCahNmHiNxD0ejLmzEF1NwrPr7/CvFGjWnNOp55mzrk5eFp7srrb6t8+LDToIScaih5CcRqUpEFZltGtoSwGZZHxNcFgnO8YAN3eh0ajwbK2u+WJiEQQNtBodV/5FiK3wf7J1eMWTtB9ETSfAlJ59etiMTy/EpSFcPBlcAund53erLq7ihV3VtDJuxOu8+bxcMgQ8pd/i+ubbzz9nkz8bkxCbcLEryAIAjkff0L5yVO4zpuHTbduteacSjnF7HOzCXMMY3mX5diZ2T1pMaOFHLUHovcahfk/yK3Bxt0YsWHlAk6B0GgkeDQxWs5Wzn/uRsxsoOOb0H4upF0xHjy6NwJ7P6O
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 1 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"for i, a in enumerate(scales): \n",
|
||
|
|
" semilogx(jnp.linspace(0.00,0.99,100), model.apply(params, jnp.linspace(0.0,0.99,100), jnp.atleast_1d(a)))\n",
|
||
|
|
"axvline(3./4.)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 52,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"import tensorflow as tf\n",
|
||
|
|
"import flowpm"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 53,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"k, pk_ref = flowpm.power_spectrum(\n",
|
||
|
|
" tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), poss[0]).reshape([1,64,64,64])),\n",
|
||
|
|
" boxsize=np.array([25.] * 3),\n",
|
||
|
|
" kmin=np.pi / 25.,\n",
|
||
|
|
" dk=2 * np.pi / 25.)\n",
|
||
|
|
"\n",
|
||
|
|
"_, pki = flowpm.power_spectrum(\n",
|
||
|
|
" tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), resi[0][1]).reshape([1,64,64,64])),\n",
|
||
|
|
" boxsize=np.array([25.] * 3),\n",
|
||
|
|
" kmin=np.pi / 25.,\n",
|
||
|
|
" dk=2 * np.pi / 25.)\n",
|
||
|
|
"\n",
|
||
|
|
"_, pk = flowpm.power_spectrum(\n",
|
||
|
|
" tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), res[0][1]).reshape([1,64,64,64])),\n",
|
||
|
|
" boxsize=np.array([25.] * 3),\n",
|
||
|
|
" kmin=np.pi / 25.,\n",
|
||
|
|
" dk=2 * np.pi / 25.)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 54,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"[<matplotlib.lines.Line2D at 0x7f7e8436df10>]"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 54,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAXwAAAD8CAYAAAB0IB+mAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAAtmklEQVR4nO3dd3RVVdrH8e/OTSMVUgiQhBJT6DX0JiBVmoBIEZAOguLYx/KqY0FHRRyxgSDYQERBOlKE0CUJSm9JICRAEiAkpLf9/nHBQUYgIcktuc9nLdaanHv2PU8c8stm332eo7TWCCGEqPjszF2AEEII05DAF0IIGyGBL4QQNkICXwghbIQEvhBC2AgJfCGEsBH25i7g7yil+gH93N3dJ4aGhpq7HCGEsBpRUVEXtda+f/easuR9+OHh4ToyMtLcZQghhNVQSkVprcP/7jVZ0hFCCBshgS+EEDZCAl8IIWyEBL4QQtgICXwhhLAREvhCCGEjKmTgbzuRQkxKhrnLEEIIi2KRN16VhtaanCWjmeV2FU17wuqPpGfT2oT5uaOUMnd5QghhNhYZ+NfvtA0ODi75WKBGXV92Zp2lUK3lWOIqDh73xL6wDcFhD9OzSTAN/T0k/IUQNqfC3mmblnmRiN/nsSFuPXsKLpGrFLPPp5Kd2YDtLq2oXP9+ejQNo1lgZezsJPyFEBXD7e60rbCBf6Os3HR2/j6ftokx2B9bw+euhSz2cCck0wXHnCbUChpNz6YNCK/thUHCXwhhxWw+8P+iqJCoA1+x4uj3bM1O4IpB4VSkaZDpTO3LPTHU60PXpmG0DvLCwVAhP9MWQlRgEvi3UFiQT/Th79h4bBlFl2J4KTmRfG3gMe86ZBWEUKP6SHo2a037YB8c7SX8hRCWTwK/OLSGxGiuHPyBEedWc9ZeobQmKMcet6x7qFZ1JN2btaNzqC/ODgbT1CSEECUkgV9CuqiImJh1/HLwKzanHuGEPbyeconq6YGssWtBap3W9GzanS51q+LqZJEbnYQQNkoCvzS05uzprXjGRGA4uIoV+jxve3tRLU/hlRGAl+cAejS9n271/fBwdjBvrUIImyeBX4YuJuxl0/55/JIUSbR9AYVK4ZMPPeLbkBZ4P82btaBHfT8quziau1QhhA2SwC8naRf+4NfouZxM+I1nzp0AYKpPTZKLfPFw7k6Xxg/Sp2EgPm5OZq5UCGErJPBN4Uo8RYd/5vmji9hmyCbLzg6XQk31TG+87HvRufEwejasTjVPZ3NXKoSowCTwTSw39TR7oj9nw+ktbFdXeSQtnXtT3fhZtyCyegDt641gQNMQAqq4mLtUIUQFYxGBr5QKAl4EPLXWQ4ozxloD/0YFaefIP7YSfXA1Oy5H85SfN/ZaUyPLAzda0TJ0JIOb1qeOj6u5SxVCVADlFvhKqQVAXyBZa93whuO9gA8BA/CF1vrtG15bZkuBf6OijBQO/j6f9afWsK3gImcd7LHTmjGnQ0lw60Fg4070bRJIiJ+7uUsVQlip8gz8TkAG8NX1wFdKGYATQHcgAdgHDNdaH7n2us0G/o101mVO/PE120+sZMyZAzgU5fGkTzUOOHjgnt+IxnVGMLRZS+pXl86eQojiu13gl+quIa11hFKq9k2HWwGntNax1y6+BBgAHCnOeyqlJgGTAGrWrFma8iyacvEirO0MwtrOgJx0OPkLzfbP40LOGQ667+PUxX1ErHbEO6cZLe+ZQc/GgTQJ8JTwF0LctVKv4V8L/NU3zPCHAL201hOufT0KaA28AryJceb/hdZ65p3euyLP8G8pL5MLh5fxy+El/JoZS4PcHMZfKmBDYXO+83OjTo0HGNm8K+G1vKWtsxDif5TbDL8ktNaXgCnFObc0D0Cxeo6uVGs2htHNxjA6Pxt9ajP5B1fQJO4X/u3uQVzGIX7b8iZVsoOo4zeQYU170jaoKvbS2VMIcQflMcNvC7yqte557et/AhRnRn8zm5zh30pBLhkn1rP1wNdsTj3ITmd7su3sGJvozbmCbtjXa0eXhqF0CKou/X2EsGGmnuHvA0KUUnWARGAYMKIcrmNb7J1wqz+AvvUH0Lcwn5xTm9hx4Gvaqt24Fr3Hu+e8eD7DDc+tlXGxa8A9VTvRK6QNHe4JxLOS9PgRQpR+l85i4F7AB0gCXtFaz1dK9QFmY9yWuUBr/WYJ3/f6ks7EkydP3nV9NqGwAOJ3Ebl/ERsv7OaAfQFHHR0pVAqXQjs6xXTnondrPOo40a5WHbqHheHlKn1+hKioLOLGq7shSzolpDWkxpF1ahO/x2wgJekAA64kA9DHP4CzjnY457ngVBRC7SrhdK3VjvvrNcbPQ9o9CFFRSODbKq0h5TjERXDo1DqiLh7kDwdNlLMTlw0GGqR7UOV8NxIrN8c58ATtApoxoH4r6nh7mrtyIcRdsrrAlyWdclJUCEmH0LHbOB27CX0umqDsqyTaG+gV6A+AKrLDIb8WNVwa0aNWb+6v25Q6Pq6y/18IK2F1gX+dzPDLWWE+nPsd4rZxMe5X9qccYL+jHdHOzhx1dKT7hVokpbcnxtMX9xqHaFG1OX3D2tOmZpDcAyCEhZLAF8VTkAsJ+yAugsy4rRgSonEuymedixsv+3qRe22rv8qvgo9DXQbWmkzXkGDqVffAIL8AhLAIEvji7uRlQvweiIsgP24bJy4dYb+TI1HOlYhyqkTPuHCiChsR43sed69z1PdqQrc6rekb1gpXR/kgWAhzsLrAlzV8C5WTBmd2QVyE8U/SIQDme3qxxN2TCw6FAOgiezypz9jg1xndtjaO9nIXsBCmYnWBf53M8C1c5kU4vePPXwCXUmP43dmJ3yq5c0F543S+A8f9+vH8gBq0rRVi7mqFsAkS+MI00s/D6e0Qt834S+BKPAsrhfC+Xx7tfYfyn57P4mgvN30JUZ5uF/jyb21RdjyqQ+OhMOBjmHEABn7KA4Xn6ZOZzc6L39Pp24FEnjts7iqFsFkWGfhKqX5KqblpaWnmLkXcLaWg6Qg8p2znbceafJiUAgWJjP1lJK9HfGbu6oSwSRYZ+FrrVVrrSZ6ecsen1fMKQo1bT9fwx1h7LoEOWfn8uvt3/r3+GPmFReauTgibYpGBLyoYgwN0fRGv0auZk6v4JXsp9tvfpde895gTuZAiLcEvhClI4AvTqdUWw9Sd2DcazJMOy6hd+B2fH36fgT+O4dzVc+auTogKTwJfmJazJwyeB4PmMTftIv+8eJX4qwfp89NAFh/5EUveNSaEtbPIwJcPbW1A46EYpu5keOUwVifGE5hdxFv7XmXJgZ3mrkyICssiA18+tLURVWqhHllDQMfnWZEUy9sXcln+/V7e23CcE5djzF2dEBWORQa+sCEGe+j8LIZxG+jtVIkljm+QsfdlBq8cxLSNT5KWK//KE6KsSOALyxDYErupO7BrOpwXizYw5IoiInEzvZb1JyIhwtzVCVEhSOALy+HkDgM/wWHIl7ycfZmvz1+EzDymbZ7Gyzv+JR/oClFKEvjC8jQchN2ju2hSrQnbzh2hc6orm/dfJvJMqrkrE8KqWWTzNGmPLADjIxl3foje8iYpeDIjdyrV2wRS1e8MT7SYgbO99NwX4mbSLVNYt8RoipaNh9Q4JniEs887iQC32rzbeSYNfRqauzohLIp0yxTWzb+58QPd5qNZkL6PF8/bcf7KRUauGclH+z8ivzDf3BUKYRUk8IV1cHSF/v+Bh77hQdLYmBhLjbRqzD0wl7Wxm81dnRBWwd7cBQhRIvX6YfAPx3v5FNbFbeXTjMbM/i6d1Pvi8PI7THU3P1pWa2nuKoWwSLKGL6xTURHs+YSiTa+RpR15O28wq4OOUWh/nk4BnXii+ROEVJHHKgrbIx/aioor+Rh63bOouG0cMdRiunNzLvkeQatcBgT3Z3rT6fi5+pm7SiFMRj60FRVX1bqo0T/Dg4uo51rAlszlvHzWG8OlFvx8ag37z8WZu0IhLIas4QvrpxQ0GIgK6QE7PmDwzg/pn2HgnfT7efJEMpFtj+Dgu55qbt4MrzscR4M8SF3YJotc0pEbr0SpXI6DDS/A8bW
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 1 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"loglog(k,pk_ref[0])\n",
|
||
|
|
"loglog(k,pki[0])\n",
|
||
|
|
"loglog(k,pk[0],'--')"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 55,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"resi = odeint(neural_nbody_ode, [pos_i, vel_i], jnp.array([a_i]+scales), cosmo, paramsi, rtol=1e-5, atol=1e-5) "
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 56,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.colorbar.Colorbar at 0x7f7ebc541bb0>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 56,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAA28AAAEeCAYAAAAdLi2xAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAACEHElEQVR4nO29e7AlV33f+/3tc848NKN5afQYPYwEDCYiiQVXF+Hg64t5CuKKcF2HK5JrhK1r+VakxK44NwgqVWAwjlxlg3EZKxmMjJTCCAWbyxSlWBEyxIVjhCQQoAdCgx5oRqMZjeb9nrP37/6x++z1W33612f13rvP6d3n+6k6dXr37l69+vXr1Xt91/cnqgpCCCGEEEIIIc2ms9QVIIQQQgghhBCyMHx5I4QQQgghhJAJgC9vhBBCCCGEEDIB8OWNEEIIIYQQQiYAvrwRQgghhBBCyATAlzdCCCGEEEIImQCml7oChJB03vELa/Sl/d3K6z30/VP3qOrVNVSJEEIYmwghjWWY+NTk2MSXN0ImiJf2d/Hte36q8npTW57cXEN1CCEEAGMTIaS5DBOfmhyb+PJGyAShAHroLXU1CCEkgrGJENJU2haf+PJGyESh6Gp7AhAhpC0wNhFCmkq74hNf3giZIPq/HulSV4MQQiIYmwghTaVt8Ykvb4RMGG3q+ieEtAfGJkJIU2lTfOLLGyEThELR1fb8ekQIaQeMTYSQptK2+MSXN0ImjDZ1/RNC2gNjEyGkqbQpPvHljZAJQgF0WxSACCHtgLGJENJU2haf+PJGyITRpl+PCCHtgbGJENJU2hSf+PJGyAShQKt024SQdsDYRAhpKm2LT3x5I2TCaI9fEiGkTTA2EUKaSpviE1/eCJkgFNoq3TYhpB0wNhFCmkrb4hNf3giZJBTotif+EELaAmMTIaSptCw+8eWNkAlC0a6uf0JIO2BsIoQ0lbbFJ768ETJRCLqQpa4EIYTkYGwihDSVdsUnvrwRMkEogF6Luv4JIe2AsYkQ0lTaFp/48kbIhNGmX48IIe2BsYkQ0lTaFJ86S10BQgghhBBCCJkURGSViHxbRL4nIo+KyO9k8z8nIk+LyMPZ3xXZfBGRPxaRHSLyfRF5nSnrOhF5Mvu7bqFts+eNkAlC0a5fjwgh7YCxiRDSVGqKT6cAvFlVj4rIDIBvish/y777f1X1S7nl3wlga/Z3FYBbAVwlIpsAfBjAlVlVHxKR7ap6wNswX94ImTB6ygYSIaR5MDYRQprKuOOTqiqAo9nHmeyvbGTdNQDuyNb7lohsEJEtAN4E4F5V3Q8AInIvgKsBfMEriLJJQiaIuV+Pqv4RQkidMDYRQprKMPEpBRGZEpGHAexF/wXs/uyrj2fSyE+KyMps3kUAnjOr78zmefNd2PNGyAShEHT5mwshpGEwNhFCmsqQ8WmziDxoPm9T1W1RuapdAFeIyAYAXxaRfwjggwBeALACwDYAHwDw0WHrXgRf3giZMChNIoQ0EcYmQkhTGSI+7VPVK1MWVNWDIvJ1AFer6h9ks0+JyJ8D+HfZ510ALjGrXZzN24W+dNLO/0bZ9vgzGSETBKVJhJAmwthECGkqdcgmReTcrMcNIrIawNsA/DAbxwYREQDvBvBItsp2AO/LXCffAOCQqu4GcA+At4vIRhHZCODt2TwX9rwRMlEIusrfXAghTYOxiRDSVGqJT1sA3C4iU+h3ht2lql8Vkb8RkXMBCICHAfw/2fJ3A3gXgB0AjgP4VQBQ1f0i8jEAD2TLfXTOvMSDL2+ETBAKoMcOc0JIw2BsIoQ0lTrik6p+H8BrC+a/2VleAdzofHcbgNtSt82XN0ImDEqNCCFNhLGJENJU2hSf+PJGyAShSmkSIaR5MDYRQppK2+ITX94ImTB6Lfr1iBDSHhibCCFNpU3xiS9vhEwQfcek9vx6RAhpB4xNhJCm0rb4xJc3QiaKdnX9E0LaAmMTIaSptCs+tWdPCFkGzDkmVf0jhJA6qSs2icgqEfm2iHxPRB4Vkd/J5l8mIveLyA4R+aKIrMjmr8w+78i+v9SU9cFs/hMi8o6aDgUhpGEME5+aTLNrRwiZR1el8h8hhNRNTbHpFIA3q+rPALgCwNVZgtvfB/BJVX0lgAMArs+Wvx7AgWz+J7PlICKXA7gWwGsAXA3gT7P8TISQZUCb2k18eSNkglAIuuhU/iOEkDqpKzZpn6PZx5nsTwG8GcCXsvm3A3h3Nn1N9hnZ928REcnm36mqp1T1afQT5b5+DLtOCGk4w8SnJjNS7UTk6kx+sENEbh5XpQghi4eIXCIiXxeRxzJZ0m9m8z8iIrtE5OHs711mnUL5UZNiQpPqQggZHhGZEpGHAewFcC+AHwM4qKqz2SI7AVyUTV8E4DkAyL4/BOAcO79gnUWFsYkQMgpDG5ZkcoNPA3gb+kHwARHZrqqPjatyhJD59MY/6HYWwG+r6ndE5GwAD4nIvdl3n1TVP7AL5+RHFwL4moi8Kvu6ETGB8YmQxWfI2LRZRB40n7ep6ja7gKp2AVwhIhsAfBnAq4eu5BLD2ETI0lBD22nJGMVt8vUAdqjqUwAgIneiL0twA9AKWamrsGaETSbgyVQ1ZV0xk2Fa1a6shZOldUjZ9iIindwFbPe11w3zG1bvZMSegITzlVJkyTGD/c5u+szsYDK+hgJHcGCfqp6bWo867G5VdTeA3dn0ERF5HOW/SA/kRwCeFhErP6oUE2qkUnxalNg0CuJ9qHh9Nzw25YnisFfZhu9DhI0bTkwYqUwAMmWGcc2YR7yNR7OzWIhFjE37VPXKpG2oHhSRrwP4WQAbRGQ66127GMCubLFdAC4BsFNEpgGsB/CSmT+HXWcxaWbbaRQkoeHVgvg0rx3g3MPe877xOI+W2rYBp83d6y1YTBPaTkvJKC9vRRKEq8pWWIU1uKrz1vJS7UWf+qCzL11TxeOPtWeDiLkwJJxMu66Yh552zQuNmY7KtNXpxBEoZdvR/DpufHOMOqtXx1+tXDmY7h05Yqrk1Ns7R+LcGMPsm/cw8NbvFJ+76HzZ8+hut3gfOqtXxYutWBGm15wVvpgK6/f2vBimT50qLPdrvf/67MKVCijqHUibObO9FsD9AN4I4CYReR+AB9HvnTuA/r3/LbOalR9Vigk1Uik+RbGprgdvyr1ifzzx4pqdrnh95+Nj0j2xmA2RTly/zoqZUA3nhSPah4T4EOEcbzfGzdt4wrGx53Ha7E9U74UbK3GZxc8tAOhsWB+K3XJOWG7P/sF098WXFtx2U2KTiJwL4Ez24rYa/R6r3wfwdQC/DOBOANcB+Eq2yvbs899n3/+NqqqIbAfwFyLyCfQVA1sBfHvsFV6Y4dpOU2/vf6irrVD1WW6vwRmnKZkSn8w+yHRcTlJ8csoaG7bttHZt/N2ZM4Wr9E4Xz4/ijbONCG9/Ro1JTlk2lrjtv6rkrqWofexcQ70TJ4rLMvv2Nf1SI+LTUlF7njcRuQHADQCwCmctsDQhZCGGtLBdUJokImsB/CWA31LVwyJyK4CPof+j1ccA/CGAXxuu1s2DsYmQ8VKTvfYWALdncsMOgLtU9asi8hiAO0XkdwF8F8Bns+U/C+C/ZIqA/ehLvKGqj4rIXej3cM0CuDGTYzYSxidCxkvT7f+rMMrLW5IEIWsgbgOAdbJp4Z8EhpGV2Dd7M53yy2b0K4CdHuHXbchM/NEUpbPOLxjRrxP19sLp6dPuZ683McI9R06v4jD7lvrL92B5U67Zh2h/Un5NNEQ9sVvOi76bPXfdYPrQ1vBgnTkeylr3tcNhBafnrSqqGDbRZKk0SURm0H9x+7yq/lV/W7rHfP8ZAF/NPpbd+02QJQEJ8WlebKr5V1u3pz1leXMt2l7g3tFjCXUoLqdfQIIqAE5d65D95Y6L12PtxhevLe7tm+2R887JqPupI/yC7cWsEmZfFRTPT/7fYf9eeVtQXEy9FHrhxvX6MkJsWqBc/T76aoD8/KdQ4BapqicB/HOnrI8D+Pi461iR4dpOo/R+eHg90im9e/axbNpFVskTtyecHnLb82MULQCgJ839n7L/dciSTTm9Y8dz342
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 1080x360 with 6 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"figure(figsize=[15,5])\n",
|
||
|
|
"subplot(131)\n",
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), resi[0][-1]).sum(axis=0));colorbar()\n",
|
||
|
|
"subplot(132)\n",
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), res[0][-1]).sum(axis=0));colorbar()\n",
|
||
|
|
"subplot(133)\n",
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), poss[-1]).sum(axis=0));colorbar()"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 57,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"k, pk_ref = flowpm.power_spectrum(\n",
|
||
|
|
" tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), poss[-1]).reshape([1,64,64,64])),\n",
|
||
|
|
" boxsize=np.array([25.] * 3),\n",
|
||
|
|
" kmin=np.pi / 25.,\n",
|
||
|
|
" dk=2 * np.pi / 25.)\n",
|
||
|
|
"\n",
|
||
|
|
"_, pki = flowpm.power_spectrum(\n",
|
||
|
|
" tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), resi[0][-1]).reshape([1,64,64,64])),\n",
|
||
|
|
" boxsize=np.array([25.] * 3),\n",
|
||
|
|
" kmin=np.pi / 25.,\n",
|
||
|
|
" dk=2 * np.pi / 25.)\n",
|
||
|
|
"\n",
|
||
|
|
"_, pk = flowpm.power_spectrum(\n",
|
||
|
|
" tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), res[0][-1]).reshape([1,64,64,64])),\n",
|
||
|
|
" boxsize=np.array([25.] * 3),\n",
|
||
|
|
" kmin=np.pi / 25.,\n",
|
||
|
|
" dk=2 * np.pi / 25.)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 58,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"[<matplotlib.lines.Line2D at 0x7f7ee8151910>]"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 58,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAXcAAAD8CAYAAACMwORRAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAAsFUlEQVR4nO3dd3RVVd7G8e++aUBCGgk9hEgoCSAdFRBQQQFBZWwodsowdh1HR8dpjnXG1zZiB1HHEQtYQBRQQaQ3QXovSSgJJQnp5e73j4NSBjC5ucm9SZ7PWlmQk3v2+R2Fh7322WdvY61FRERqFpevCxAREe9TuIuI1EAKdxGRGkjhLiJSAyncRURqIIW7iEgNFOjrAgBiYmJsy5YtfV2GiEi1smLFigPW2thT/cyn4W6MGQYMS0xMZPny5b4sRUSk2jHG7Drdz3w6LGOtnWatHRsREeHLMkREahyNuYuI1EAKdxGRGkjhLiJSAyncRURqIIW7iEgN5NNwN8YMM8a8kZWV5dH52zJymL/lAMWlbi9XJiJSvfl0nru1dhowrXv37mM8Of+DJbv5z7Z/EhRYQoeoXlyTPJBLkhOoF+wX72aJiPhMtU7B3/drwl7TlMUH5rG2eDVrVr3Oo4vOok29AVzX/jIGJDUiOjTY12WKiFS5ah3udVdNZPyPE3EnXsTquK5Mzj3CD/uXsDtrOw9+8hMuVzHxrZYwMP4Cbuzam7joUF+XLCJSJap1uJPQH7L34NownS4bp9PFFQRn9aP4nPZsHdSOdzau48sDX/JB2nTe3xlOuLsTvZv24+auF9GhSQOMMb6+AxGRSmH8YQ/V7t272wqtLeN2Q9py2DANNnwBh3cCBuJ7kdl6AJ8FhPJZ6nK25/6INYXk7rid5vXa0bddEBclNadvq3gCXAp6EalejDErrLXdT/mzGhHux7MW9q87GvTTIH2dc7xJZwrbDuHb8GZkFvdh1vp0lma/hStiCa7CBJLCz+PqpEEMbd+BOkEB3qlFRKQS1a5wP9nBbceCPu3oNWLaQtIwVjZM5u3UjSzN+IE8mwaAzUukX8SfuTi5ERe0a0h4naDKqUtEpIJqd7gfLysNNn7pDN3sWgDWDREtIGkY2+J68k56Khv357B9a3cO5BQQ2vJVYkMSuLBFf27rOpDmUVq9UkT8h9+G+3HruY/ZsmVL1V489yBs/srp0W/7DkqLILQhtBuCu+0wFpHAEyueJqXgR3AVYt3BhLnbc0nza7il2wUkxGjmjYj4lt+G+8+qrOd+OgXZsHW2E/SbZ0FxLoREQNtBFLYZxGelwUzZMpfNRxaTk3YZJTntSWicS8sWu7iuw2AGtm6vmTciUuUU7uVRnA/b5zpBv2kG5B+GwLrQegC23TBSG/bh2+1F/HfDZPYG/RcAV3FjEsPOYXjbi7m6Yy9CAqv3DFMRqR4U7p4qLXbG5jdMgw3TIWcfHJ1LT9Iw1jfoyMRNi1i0bx5ZdhNYF66UxxjQLo6urdxcmtSWqHoavhGRyqFw94YzzKWn3VDS4/vy0a4Utqc04dsN6RQ3egFXyD6iXR3p26w/o7oOISG6ka/vQkRqEIW7t51hLj1JwyhuO5R3UnYybetsduQtxQZkYq0hpnQgN7e7k4vbN6ZZZF2f3oKIVH8K98p2hrn07naXMj2rmI/Wz2TXvnBSUltjAo4Q1WoCyRHncnXyYC5t05PAAI3Ti0j5KNyr0s9z6TdOg50LwJb+MpeepGHsqNueyatX8mnKePJcmzHGjXHXJ6Fed8Z1Gsslbdrj0lIIIlIGCndfOcNcepKGsT08mQmrv+GH1HkcsqvJ3X4vDeo0okubDM5qUsQtXYbQJCzW13chIn5K4e4PzjCXnqRhHGrai/k7C5m5bh/fH3wZE74Maw0RrkR6Ne7LTZ2H0LFhG1/fhYj4EYW7vykugO1zTjmXnqTLKEi4iI83b+KzzTPZkrMUG5yKu7AJPYIe55L2jWndLI+2sc0ICwrTy1MitZjC3Z+VlhybS79xOhzZe8Jc+tI2Q/gm7RAzN25m+eb67D6cRVibxzCuYowNoa5pQHhQDF2iL+DiFsNoEhlCSsEqWkU1o1FoI+oH1dc/ACI1lMK9unC7IW2FM4/+FHPpbbtLWZsfyjurvmZn1m4O5qeTXXyAAnuIoqxOFB/ugwk4QlibJ35p0mVDqBfQgHOirubCuCFEh5WyKecHWkU3p2lYYxrVa0R4cLj+ARCphhTu1dEZ59IPhaZdICoBIlvgNoEcyC1kT2YBuw5msTp9Lbuy97Avdz+HCjLIKT1A3qEulOYk4aqTSmjCyydcykUwfSPvpHeTiwipe5g12d+QGNWc5uFNaFSvEY1DG+sfABE/5Lfh7tNVIaubU82lBzABENEcohOcsD/h15YQUh+AnMIS9mTmk3I4h00ZaWw/nEbqkb2k5+0nsyiD7IxOlBQ0JSB0E3XjJmHMiX8u+tb/E10bnkNR4DbWZs8iPqIp8RFNaBzamCZhTUiMTMRlXFX4H0RE/Dbcf6aeeznlZMDBLXBoBxzeceKv+YdO/GxorBPy/xP8CRDWEI72xotL3ezPLiDtcD4pmTlsPbCHnVl7SDuyjwP5+zmc0Z6CgjAC6/9ESKMvMYFHMMb9y2VuS3iO3513oXaxEqlCCvfapCDr1KF/eCdkpQLH/f8OCnWC/+de/vHhHxEHAcd2obLWcjivmD2Z+aQezif1cA7bDu1jd9YedmamkbanNTFhIbRJmktCbBC3dryBxKjEKr55kdpF4S6OkkLI3H368C8tPPZZEwCRcafu8Ue1hJCwXz5qrWXhtoO89cN2Fma+RVDkMoyrhA7RXRl99k30i+tHoEvLK4h4m8Jdfp3b7UzDPD7sj/8HIP/wiZ8PjT0W+NGtoMNvIKY1W9NzeO2H1czY+TkmYhGuoEwGNLmW5wb+SQ9kRbxM4S4Vl5956t7+oR2QnQZYaDsEet0NLc7lcF4x/1myg0k/zuBQVhStIhMY1K2QA2YuN7YfSfsG7X18QyLVn8JdKldOBix7E5a+6TzQbdYdet8N7YZS5DZM/2kPE+bvYFPud9Rt/AW4ikiK7sjN7UdycfzFBB03ti8iZadwl6pRlAer3odF453efVRLOO9O6DwSG1SXxdsP8fr8dSzc/xVBUYtxBR+geWgCM678XEM2Ih5QuEvVcpc6yx4vfAlSl0HdaOgxGnqOhbBYdhzIZeL8bUzZMIcisugZczG39o5nWfYkBiVcQpeGXRT2ImWgcBffsBZSlsCCl5wF0gKCofN1Tm8+pjVZecX8d+lu3lm4k/0FKdRPeAXryqdNZFtGJl/P4ITB1A3UjlUip6NwF987sMUZrln1X2dd+7ZDoNdd0OJcit2WGWv28ub8jWzKmUedBosgeC/1gyJ4d/AkzZcXOQ2Fu/iPkx++Nu/hhHy7oVjjYvmuw7w1bzvf7FxEcMQqBjb8HWPOb01a8SKK3cVEhERQP7g+oUGhRIVEEVtPm5lI7aVwF//zy8PXl50plVEJcN4d0HkkBNdj98E83l64g4+WpZBbVEJM2/EUulJPaKJzbFfeHTwJYwwjpo8gIz+D+kH1CQsOIywojO6NuzO642gA3lv/HgZDWHDYL59pF92OiJAIH9y8iHco3MV/uUuddewXvOQsiFY3GnqOgR5jICyW7IJiPlyawtsLt7IvLw1cBZiAQoyrAOuuQ0BhW2JCg3E1mEVAUBaBgUWYgAKsyadlWHtGtLqL6NBg7pg/hNySnBMufWHchbx44Ys+unGRilO4i/+zFnYvdmbYbJoBASFHH77eBTGJlLotezLzOZRbxMHcQg7mFHEo1/k6kFPEodzCoz9zjuUVlZ50gVIIKCA4qIiI0FLqhxbQM64VDw/oR777IGsPrOWiFhdplo5UKwp3qV4ObHGGa1Z9cOzha++7Ie6cX1ax/DX5RaUcPD7wc47+o3D09/uyC5i/9QCRdYPo0mkhyzKn0LVhVx7s8SDtY/T2rFQPCnepnnLSnQevy9501raJbnVsxcrIFs7Xz78PawSu8q0nv25PFk98uYGF29Jp3PwniJpJbkkmQ88ayj1
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 1 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"loglog(k,pk_ref[0])\n",
|
||
|
|
"loglog(k,pki[0])\n",
|
||
|
|
"loglog(k,pk[0],'--')"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 59,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAkwAAAJCCAYAAAAybpizAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9bZAcyZkmiD3uER4RGZlZWVUooFFgo9gEyG4M2RhOg2SDHz0fnNudXe3caj9Mc9pZkzRnq7O11YedrUwmu7G7k0xmkpn2/sh0Z6Zb2dqeTneSxvZu7XZHd5q5O+7scW63dzhoDtGcQZNENwlMs9CNRAOFQmVVZmREeLi7fri/Hh5RCXY3u7tI08HNyoDKyoyM8PBwf/x5n/d5mTEGT9qT9qQ9aU/ak/akPWlP2uMb/0mfwJP2pD1pT9qT9qQ9aU/aT3t7ApietCftSXvSnrQn7Ul70t6lPQFMT9qT9qQ9aU/ak/akPWnv0p4ApiftSXvSnrQn7Ul70p60d2lPANOT9qQ9aU/ak/akPWlP2ru0J4DpSXvSnrQn7Ul70p60J+1d2gcCTIyxP88Ye50x9gPG2G9+WCf1pP13qz0ZR0/aB21PxtCT9mG0J+PoSftRjf24PkyMsQjAGwD+LIC3AHwTwK8bY7774Z3ek/b/7+3JOHrSPmh7MoaetA+jPRlHT9q7tQ/CML0I4AfGmNvGmBrAPwDwlz6c03rS/jvUnoyjJ+2Dtidj6En7MNqTcfSk/cgWf4DPfgzAneD3twBc/VEfyPMtM5mch1ISxmgYoxFFApwLGKPcawaAAWMcjHFwHsMYA60b/xnAsmL0fsYYAIAxDoCBMQbGOOI4BsBgDPOfbT9DWNHAGANjlD+G/V27v2kA9vicRwCY+xfunExwDDq+csdm7tyYux4W9Ib9G10fXVN4XbZP6G/2vYDti/Y6I3Aedc4DAOI4A+cxhJgB4GiaCYzR7rOrG+cxGGv7SikZ9LVx/UPX0H7/4eHbWC4fstVHfdf2vsZRkgyNSidIjIIGQ+LuYwEgNQrCaGgwaMagwKAYQ87tMNdaoTYaI90g5FWZu0oDQDHu/yYZRxZnfqxo3WAe9B+HAad+ARDBgBkDwxi4MRBGIzIGMTSUO6eSRZCM288CiIwBc9/IjUHkjsmD40vGocH8+xp3PhwGDeNY8BiR+z4DwDCGyBhERiM1GhwG2n1auXPLXP9pxtCAQTKOhnEoxpAYDW4MSh6hZhHWxMA9g/a5rbXy56YYg2IcCgxD99wBdgwvjUYWvFczBuPuiQn6vWYRDAA9n+4ZY06/x3ETtvc9F42jxGzFmb2vzF5/ohVS/+za3SRz/a8Yw9KNo8RoCKMxcGPBAL4vaTwod60MBopxHHKBxvVxphuMtfTjM3HjRAf9wowBh7Hjyo8x9yy672pnpnb2YMHnAfixGMFAu/PkwTgy7viKMVQsQgSD2Gj/HZHRYACEGxOdcwSC8W+vU7vjKjCEE4JiDLMo8f0XG43UjY2oN44MgNrNsfYZ0kjceTB3Dsod3TDmr1UzhgdNiUMlT2QuAoBJJMzpOPPnA9h5QBgNZuz9o/6W7vpoHKVGIXbPCH1auWuje6/A/DU2jOMoEmjAUXOOgVYYaenvK/UrjT06Rv88+i18VbszsffX+GeUxk04TmkFle6Zj9zcW/KonY/cOOWwYyUO5iP0zonGk2IM2j0b9FzRmNNgOHTjKA6eQzpfmk1pLqvc3CKMRgTjxxxz72ncOKVrpT5/R1WYrRhHHwQwvafGGPubAP4mAEwmO/jX/rU/xOHhW35Rnkx2sLX1DZTlBdT1NspSQqkKSTJyC/4jKLWB2WwKKReYz+85EGMXMK0bRFECxjiSZIQoSjAYnEKSDDEacWidQUo4kKZQVUdomiWapoIxClIWHoxZkJUhihIIkaOu5yiKPXAeu3PJEUWJ/z4ADljUaJoSVXUYgA37Gv2eJCPEcYY83wLnMbRuwLlAmo5RVUcoigce0FjgwlGWB9C6AWMcUZQgy9ahdYPlct/3QRxnECL3n6F/h8MzSNM1nDp1HwBQ1+dQlgmOju5CiBxJMoJSNbRuHDAC8vwUIjcYlapRFA8RRcL3RVkeeCAbxwP3tyH+w//wF05sDGXZOn7pS/8rNE2FPD+FjY0L/v48enQbh4dv+T6gnyhKYIxGWR7gTj3HS/N7AOyDXrAIRSQgXH9ORQ7JOGZRgr04w5fz04giAcY46nqOqjrEK+5BvVw8xHazhASDdBP8jAvcyE9BaIXcKOxUR7hUHvhr2U3HfuEAgMJNnsJoTFSNi+XMnhuPkOsGuW4651mwCNdGT6HgMcaqtpMVj7FTz3Gl2EOuGwijIRmHBMNESwitME2G/ji5UbhYHUKCYRankGAoIoGpyDGLEj8BSjBMkyH+zFM/C8Y4pCzcM/EQ11WFrabEXpzhBR7jFaOhAPxivoU4zpAkI9T1HH+8/wN7nSxqrzkS+CKPcU03SIzGF6MUnEf4p//03/zhRzB8fAvH0Uac4W984qu4ma3j+uAUPrn1HL63/wP8b+59GwWPfd/lurG/M46pyCGMxsXqEKeaEherQ9tPjGMvznAYJRi4/j8M+rHgMW5m65hFCaYix049x+XiIYRbWM/Xc0xUjb04s/fNAavdZOT/TveV/qVzmkUJGsYRu/G75DEGusFWU/pzo3FE95Z+bqdrkIxjomo/Ziaqxvl67r+LroGeDzpHOmbuQCMds+Cx/3kYZ368FSzCb68/gyISyJXEjlzg8nIfE1VjTdV+EzAVOQoe+/O5nm9homp89ehup6/34gzCaAitsN0skesGe3GG/8Pbr3yUQ+jYONqMM/zGJ34Zf3/rEgoW4fk4xfDwLfzGwzd8301UjdhoLF2/TOMBcqNwoTrEVlNip553rm0WJfbagvtM9/ZmOsFenOFWNsGl8gBXF/f9vd+WhR9HNK9IxnFHDJEb5ccRjaH+OKLvoc/RGADg3+/nluDnTjICAKypGoAFImuqxrYs/HfRMel4U5H7vgm/B7DPS3g+u8kIh1Hix/DX1p7213exnOGlxTsY9K7rTjLy86wEs+MFBherQw+07oghbmUTbMsCW02JNVUj1w3uiRz/xvT6ynv/QQDT2wDOB78/7V7rNGPM3wPw9wDgYx/7nLELmJ0goyjBePwWJpPfRxQdgTEJYAdKDcA5wFgDQMAYDmJ7LAPUMiW0WFqmxf5LC5yUI/ceC5gsgLE7SKUqSLmE1tIfx7JR9thNU3lAAcB/L/3QQmyvUTsglbpjRWiaEpzHaJoSWjeI4wxxPPCgjo4ZRQkaN7m1rI5tdB5CDHp9qlf+znnswZxSNaQsUFU7YGwJpUbQeumuE51roPOxfQlE0RxxHIOx0x78aa08IIuiBHGcIooScB71mLP33d51HIVjaDI5bxjjEGLgziPz74vjzINKe09S148SWjd4ramwo2oI7VgEDsAt5H7CD4DJi4z78ZEkQzDG8VpTYeIWuomWmDQVikhAamMnehhs1wsAgIBBbixwIgBDk7zkkV9g6PsB+O+exSly3fgJRWjlX5eOEQpBCL0GwE9qACC0ggh2uJJHgFKQbocnjAaCyZFAFJ0XAEhZgDGOpimhVA3AtH0IO46vOqaTxtNweAZJMsZnmhJNU0LKAgDwitG46p7Zq7Dj0I6hD5R/8r7novPZxFAfSh75Z4gWGwkGuDmcQAKB2ombWMMm3O6eWrjQSTDftyLW/p6EEzz90LFCIBOCpfD7ALs49Rc6uIU5XNQ699O9Zs9DdcYL/V0E1xL+7q/JHVP05iI6t1w3ODQagr4/EsiNglD2unMl/Rhqet9N57qbjPDOuc9jqhtM77+GbVngUnmAvTjDVOTYakpMtGqfBx7j+Nm8r/a+x9Ez6Zrx/cwjvGI0XnLAJ+wvAnSScTsnKOnv749qBY99/xQ89vPJpKmQ6wZxb3z0x5NkHBMtjwEKOjf6PdzEUZPB2KJ73f98f6z0Pw/AA7NVfwvHUP+8wv/TxqXgsR23jCNXErlbz+n+Txxoo/l2L86wm47x/Z2XYIzG/Ue3cak8wFeP7mI3GeF
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 720x720 with 16 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"figure(figsize=[10,10])\n",
|
||
|
|
"for i in range(16):\n",
|
||
|
|
" subplot(4,4,i+1)\n",
|
||
|
|
" imshow(cic_paint(jnp.zeros(mesh_shape), poss[::2][i]).sum(axis=0), cmap='gist_stern',vmin=0)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 60,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAkwAAAJCCAYAAAAybpizAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9a4wlyZUmiH1m/rh+PR43MjMyK6Me0exKdrGGZJHNZLOySdbMdO/OUiPMAjsD7C60AwxGGgGjP4IeWAEaCIJ+z68F9peABnaBkaDBahuS9jE7mG3NgrNSNTlZ3Z3V7GL3ZLGZ2ZyoYkVWZmRmPG749evmbqYfZuf4cbseWQ9mRRHaNCAQEffhD3Nzt8++853vKOccnrVn7Vl71p61Z+1Ze9aetbOb/rwP4Fl71p61Z+1Ze9aetWftl709A0zP2rP2rD1rz9qz9qw9ax/RngGmZ+1Ze9aetWftWXvWnrWPaM8A07P2rD1rz9qz9qw9a8/aR7RngOlZe9aetWftWXvWnrVn7SPaM8D0rD1rz9qz9qw9a8/as/YR7RcCTEqpv66Uelcp9VOl1D94Wgf1rP0Pqz0bR8/aL9qejaFn7Wm0Z+PoWXtSU5/Wh0kplQD4CYB/C8D7AP4AwH/gnPuzp3d4z9r/v7dn4+hZ+0XbszH0rD2N9mwcPWsf1X4Rhul1AD91zt11zjUA/jMA/87TOaxn7X9A7dk4etZ+0fZsDD1rT6M9G0fP2hNb+gt89wUA74n/3wdw40lfKMtLbnPzRThnxasKSunwmme7nHNQymO5JMngnIO1Bs45ONfBs2Lys8pvSWnenlIaSZLxPpzrYG3H36NG2xoybS46RkTbVyPfdeH4rPiu488rlYTv0bEqPrbV/dOx+T7xxw3+vlKKPy+3Kf9PkgxKaaSpg3MKzqWwthP9h5V+0zqFUhbO+evRdcsz+qf/rtYpjo9/jsXikcKna59oHKmsdEmxBQ0HC4UyjJ0lHBJnkTgHpxQcgFZpdFBY1wkAoLItUmcxcRbKOWj0o8EB/D0HhU4pGKUx1Rn3qXMWlTVwob8TZ6HDFhQA7fq/VXhfAfwZB7/NTimo8FlasaiwVe0cEueQwvJ+aDQZpWFV380O/nht2F6C1TGUOgcdtkfbUA7IYP130Z9zq5TfVugPozRapVEmud9fGNvGtnyOjj+boByMc39/LLsGGv6c+BqG86Xjb8J5tfN7B865y2dd+ye0T/ws2khydykt4OhYAaTOIhu571ulYaGw1AmUc5i4DolzyMNneQyJa+PvXP9/pxQqncJCoVUaueswtS2ffxI+SdvRPDYc97Mfr37My2M+qzmsjsnhGO8bvWfhr30CN7gvACBxbjCW/Wf76yrHZRfuH7ldqxQWOuV+TpxDFu6ffhz571ko1DqBpfsMDqVtuQ+oX+meoTFslcKDtsZxZ87lWQQAm0nmLqcFjNLhPuufQ/4YfXPw46BTGkuloQAUtkMCNxhzNG6U6Ge6Zh18H3ZKoUUYR67j72rn+L6S/dMfSz+e5BiQY89h2HXxOJHbj8e9Cs9eC/9co7HnxGf8MxP8/HPh2SWfkdQ6pdDRvkIfWCgswvM8Cc+8zFno6P6gfqx1gi7cdwkc1jvDY46Oje4vuY8PuyWORsbRLwKYPlZTSv19AH8fADY3X8Lf+3tvomnmAdDkmEw2UZZ3sVzuwhigaeawtoXWKbROkWUlrG0xn99D08xxenofztnBZ7TOkCQZ0rSAUgnK8hLStECer/NxtG2Nuj6CtQbWtmjbGta2DG4kQLK2hbWtPAsAQJZNoXUa9qMZ6LXtEs516LomAI0GbVujbWv+Lh1TlpUCICJ8v0bTzPl/zwwDbbtA2y7Rtgs4Z5FlJW/PA8AWSmlonSJJciRJjjQtoHWK6fQi8nwdV67ch1It6vplLJcF5vN74tr007XWCYpiK2yjhrUF5vMDPheaLD2o0vwzmWziH/2jf/OpjJWzmhxDRbGFv/zt/y3o762tL6Bp5qjrQywWj/Dm8pgfQL+Z5Hy9rG1x9/Q+dppTXGtOkAUQWiUZjNI4SAtUOoVRmv//UlpgOr0IAqBNM8et6iFm1qC0LWbtEqV4YAEeZFQ6RWY7lK5D5ixK2yILk/Fevo6jAEDo8wD4ffpe2Rlk4SFShrFI3z1Kchil+fMAMLMGO82p3xYcKp2i0im22xqlbbHd1jBKYz8rkTmL7bYe7B8A98GRzlAlGSqVYD9fw79x6UvQOkHTzNE0c/zR6QM+RqMTGKVxlOT4RrhfaRx2XYMfnN5HaVuUtvXbTnI+NwMFoxN8N5kgSTL883/+D/710xs1q02Oowtpgf/g5b+GN9ev4iAtsDfZwBvze/ibhz8b9D0A7m869leXR5iGPqVG44bG0EKn/HqlEuxNNlCpBFWSYcdU2G3mKG2Laegb6h+6rrQ9AIPxQ6/Jz8ZtZUxFnzlKcizC92n71Ka2xaxrVs5r1jWDbR0lOY/teJ809h4GIEF99/2N5wEA222N7bbGteUx90Er+q/SKY/FvXwdpW1xvTrgsW50gkqnvl/Ea0Zp/MP3/+XHHQ6fuslxdDEt8Le++Ndx67mv4x1T4Ub1AK9VD3F98XCl74+SnO+v0nX46uIR35vU6Pyp3+inDX2zn5WDe/va8ngwhqj/qdH35VgobTsYX/S5uMXfBTA4n4NwfVulkYqxAICPh75D25JjWe43HkfUF/LZ3NI4Wt8B4MfRTrvAl+rDlfNaiO8e6Qzf33yBxxE9uwE/buSx0n7/jx/84ei1/0UA088BvCT+fzG8NmjOud8B8DsA8Pzz33Q0gdHvLHuMLHsA5zIodRlKbTIQAOSkDhALIoGC386EQYPcPoESpXQAR10AS0t0XTMARbQf+o48BmoEiJRKmMGR30+SPLzfv07bi4EGHV8nHk6SnaJjJhzfs0BJOL749XTw0x/zBpQycC4Nn0/4vSTJBsflQVeLorgL51Io9TKMARaLR4Kd60G31slgX5+yfeQ4kmNoNnvJyX6kPtA6Q5aVeAMI17bja0QAmx+uSsMk/kFAkxg9hOgmPUpyvgZ0js5ZDxToARUAEd3kPGnYDllY+WS2G4Ih2zJQkN/hvzVQWQD9ZfLbdNZ/L37QhRuezo0m+sx2KIHBewBgoICRhyPtg7ZJ/fKNcO7WdrzAMDrhY6xUAqMTXA+gJ00LpGmBixe/COcsfnuyCWMqVNVDHi9tW+NW1/B50Jj+Bdonfha9VMwc4Psrg0PZGQBhUlMJX1ugP8fMWWRwK0xUDJZo8jdKw0ChCky3BGEEFGIgRE1uX05WY/uV/1NrR16Xk0K8D8AzP/T5sQlybD9ntcxZpOG8KpVw3wy2E8ZzJvqLPlfpFHcmm3j8/LcAOPze0R6uLY/xxvweb2e7rTGzHd9LT6F94nH0K8XMZbbDTWswC9e3SjIPKMWiCRAgN4w3ee0lqJA/NPHLa53ZDtApMmdXwFLc4rET/y9B0dhYkq/Jz8jvT8W4GjsG+d7YGI7B01nHfhwAJ11vup/akWOl/490htvFFrLdN9Aqjf/u5AP8RXWAv/3op/xslM9HajHTRu0Xme3+AMCvKaV+FX5Q/U8A/O0nf0UJUGOR5z/HZLKHoriLNH2MrttA216AtQWcy+BciuXyBQ5JaZ0w49R1DdK0CNvL+L0kyZnFIYYnTSewtsVyebLCMPnJUHH4jibXtq0Z2NBE0XUNA4whWemBBB2bBHT0PwE6yQjRecj90neV0gzqCIRpnQyAQr+fhLdJYNF/36BpdqBUi7Yt0HULpOnEXwmlkWVl+HwDpVoo9RhZ9hibm29C6xrOpVgsXsGjR3+Nj4OuhQ+REuP0aRnwTzeO/D59XxhTwTkbQHOCPN+AMad8rQBguTyGDQAHAD/MMqVxlE4G7BIxJa8EFpFANV2v7bZG2RnMrOHjoYnR6GSUWSo7w69VXYPMdszgQKcwUMjQsxHQGT8MDBR2lMasXXo2DGo4mdpu5SGVOYvSWcAaZgbGGj8UBYDLbIcjneEoneAoyfFrSR7CuBbGVOi6Bt/QKd5SGsYluB7u5zxf5/E3nV7El77
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 720x720 with 16 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"figure(figsize=[10,10])\n",
|
||
|
|
"for i in range(16):\n",
|
||
|
|
" subplot(4,4,i+1)\n",
|
||
|
|
" imshow(cic_paint(jnp.zeros(mesh_shape), resi[0][1:][::2][i]).sum(axis=0), cmap='gist_stern',\n",
|
||
|
|
" vmax=cic_paint(jnp.zeros(mesh_shape), poss[::2][i]).sum(axis=0).max(),vmin=0)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 61,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAkwAAAJCCAYAAAAybpizAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9W5Bl2Zkehn1r7cvZZ+fJPJlVmdWVfclpVAGNGgAFTBfArgHQcyNBaKihxVHIVJjjoKiwJUYoLDsc5oNoMRyhCPtBfnGE9eCHEU2JUpiSTNsxvJMQaUyQzYGqMSgAUw2wugdVQGdfsroqqzJP5jn77LPX3mv5Ya3/3//eebK70ZcEwqoVkZGZ57Kva6/1re///u9Xzjk8bo/b4/a4PW6P2+P2uD1upzf9sz6Ax+1xe9wet8ftcXvcHref9/YYMD1uj9vj9rg9bo/b4/a4vUd7DJget8ftcXvcHrfH7XF73N6jPQZMj9vj9rg9bo/b4/a4PW7v0R4DpsftcXvcHrfH7XF73B6392iPAdPj9rg9bo/b4/a4PW6P23u0DwWYlFK/qZR6VSn1I6XUX/2oDupx+x9Xe9yPHrcP2x73ocfto2iP+9Hj9m5NfVAfJqVUBOA1AH8awJsAvg3gLzjnfvjRHd7j9v/v7XE/etw+bHvchx63j6I97keP23u1D8MwvQDgR865u865CsB/C+DPfTSH9bj9j6g97keP24dtj/vQ4/ZRtMf96HF71xZ/iO8+BeAN8f+bAK6/2xfy/LxbW3sG1tbhFQelImgdwdoGzlkAnvFSSgNQiKIUgEPTVHDOwtoGgINkxpRS4jv+t1IaUTSAUg5AA+ciNI3h7VPz22m3p5SCczYcS7fRdmk/9Bl5TP4Ya/EdBaU0tI4BqKXH3O5Pdd6j7dJv2u/J7bT/02taJ9A6QpJUABSaZgjnbLgG3c/SOWkNABbOJXDOoq7LcF4nrwUdRxQlODp6C/P5I3XKh96r/VT9KElyZ7N1JM6iUQpDFcE5C+MsIjjocB8dAKM0aqUx0r6bF9YgcRZD20CJfuCgYAE45V+l/2ulMYjScG0iWFtj1lSwUMgBVGGftCW6E8o5aACJs1BwUA5wym+3Uhq1UnzHtHNQADT1ewCJtUhdw8djw9EapWHFd+m9WilEziFxtnMOCMegnUMMCwXAQEPDf9YqhSZ8Tl6vCA7KOdRKo1Iawzjzn3ENrG1gXQMVjtsqf60qFWEl9HP5DDX1Av5JpvuioMP2nVKw4ZpYpWCm9/adc1vvo8/02089Fq1GiduMh7B8vfy18vcMfC8QrksDhbmOoeAwtDVi55C6hrfnwhk6enbFvhqlMdMxamgsdITM1hhZA9pTHPptQ2MCfB+Kwl2ka93uq+2r8jX6Lt1/Hbarep9plIKDOvEMNEpB9/oRH6Oz/r6JbSkAkXOd8waAOmxLHg9dA4TvxM4iDdvUzoV+5L9noVDoGI1SUA5IYDFqjH+Weufanq/fx/26xFFjzmQsAoBxlLjNOEOlND/zkbOIwnWX19gojUZpzFUEBfh+BIfUNnxd6fzpe3T9XXhWCx2jVhqljpDZBivW8Hej0HdpvKD75Y+lHV+WNce/Vec1uq9qyfk0SnM/oM/QcWr4e0zbo+fM9yN0xg/q6539ir7Qf64K0Y8iOAzC9dNw/Hn67jz0I6M0Yucwbha8f+rj9IzQfiyAe3WFiT3Zjz4MYHpfTSn1lwH8ZQAYj3fw7//7NzCbPfA7jwdYXc2xvv5PMZ1+EUWxhaqaomkqRFEKrWNk2QhNYzGZ7GKxOMLR0Vth4K4RRQMkyTCAKsV/Z9kGkmSI4fAclLJQqkZVWRTFQxAAKMtD1HXJYKAPhFpg0XaJNF2B1gmSZBgmBg92jClgbc0/db1A0yxgzJyvw+rqNpIkRxxnAvgA1tYwpsB8/khcMw8iabtVNQUARNEgHF8TjrGC1jG0jpEkefg7QRwPkGXryLJ17Oz8cwAG8/lnMZ8/h8PD+6fdKWTZOpJkiHTwAE29gcPD+2iaSpyH4/tC12s4PIf/4r/4tQ/UN95vk30oyzbwq1/5K2iaCmtrT+Ppp38Z0+k9HB7+BLPZfdwoD5EHwHolyRFFKZJkBdYavD55A9v1HFfFtS5UhEk8wH6codAxTLg3+3GGSZTiy/kmoihFFKVYLI7wndkD5K5B3hjsmBk26xKFivy2ogRGaUyiFImzyG3tfzcGSQApu8kKJvHg1HOlz49tC2zHTYW8Mbi5ssXHBQAr4TwTZ7FZl7i8OOLvTKIURmls1iW/DwB7SY7EWewsjvmYaRt7SY79OMN+nPF12E1H+Nefug6tI5TlIcryEDcmbyB3DcZNBQOFIkowiVJci1Ik4ZpHUYq6LvGTw9f53PfjDJeSHN+1NWLrwYbREb4SDaB1hH/2z/7j1z9QB3mfTfajjTjDv/XJfw3fWHsae/EQzw83MD/4Mf69/dt8vNT24wyFirAfZ8hdgxdm9zFuKoybij9T6Jj7DwFPI3520xEmOsEkHmDbFLi8OEJuawxtjc26RG5rvmeAn1yPQj8ahn4kj0n2VQCd94zSmOuYv0uN+uPdwZrvo+EeyL5GxyO3ZZRu+3LYXhG2T/8bpfkY6Fq8kY74c5Moxd8690kYKOyYGXYWx7haHiC3Ndaaiq/ZJEpR6Bh7SY5CRbidrWNsDV6c3kNiGyRwmEQpJlHKz0USJlsDhf/k7T/8QH3jp2myH52LM/zZT/3rePXpX0ZVzYDjt3Ct2PfH6yzicE1qpbEfDfhZyW2Nr07vIbc1X1sAfP6y79D1nYvrsh9n2K7nuLw4wlpT8XaW9aMT45HoE3IfQNuP6HX6ruxfsh8VOub/+33t3foR/RD4kcdEja7DXpLjKIx5+9EAf/vcZQDATjXFTjXF88U+nzu1SZRirmPspiPsxxm+sfoUctfgxek95I1BLhY79CzTvZrrGH/l3veW3vsPA5jeAvCM+P/p8FqnOed+F8DvAsCTT37J+cndA444zpCmu0iS+xgMdtE0GaLoPJqmOgFgCNjE8YAZnDge8ASudcwDNQEna2tmd6wtYK1hBqhpKjSNZ18AIIo076NpDKw1HWBDAAUArPWTjNaxYL08qyOP23+2hnNWsDhx+Kxn1ZSynX04Z6E1YK0Lx9kOZlpHUEqjCfda7oeO5ySY24JSNazNYG0qPu3ZIR06rGfkEmgNDNI9uOQBrP0cqmqGsjwEAU0Cc3KfSn3QBR2A99GPZB8aj3dcex2Bul6EaxYjTUf4ZaVR1yWsrZEkebgODZqmQu4a/5CqiAfZIkp4gC90DAMFoyMGANRf6F4TUMhtjVzcG6MjHnjo9dzWPMjT33ncwNiaP0tgiwd9pVFECRI4PpbENki05e/QRAegM/jI1/LG+O+Gv2nbBgpJOF56jb7L30tyBmXP6zgsUDzLa23tz9U6BkuFinAtGiCKkvD85Xjiic+jrhcYDNZQVVOU5SHOhefzWjXDDdXt13Q/P2D7qceiZ7KxS6zvDwkcn9d+YNPongH+HhVRwkBZXmtqcoKrlW4nvXAPaZt0rwgsScAhm5yA+ve4P6n1v9s/LsBP2J0+YpvOpNHf9rLX+/tdtq/+5+l5m+gExIX090HXDPCghybq29k6kp0XcWBr/B8fvobLiyNcn77Dz1piG+ToPoP2Yx6LgG4/enaw5pIwNzjnmQ4CfOOm6lw3WlAltukAh36T93W+BDwlcNyXhgIEJc527gttOz+lH/X3KX+f1vrbp368bAwy/T7X+wyd07LjkcdP5/kwzni8XtaH+PqEfVM/ujU8h6cu/2kAwL949CO8Mn+E3z78CXJbM+jub+s0ZfeHGaW+DeBTSqlPwHeq/xmA33m3LyhlEUU1smw9/L6L4fA1DAZvQKkaabqHut6AtRmaZhXWZijLSwwaaFIkQBXHQ6TpCoMRv7odYDBYg1IaVTWFUhHSdAVNU6GqpswCEZNFjI8MOznXoKpmPJD77zTM6DhnEUUDxHHL+CgVIUmGPMFGUcr7pO9JgOOBnkVVmbANy+xUFHngRRM/Aa4oagFP03RDhAQY03TEn2sag9n
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 720x720 with 16 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"figure(figsize=[10,10])\n",
|
||
|
|
"for i in range(16):\n",
|
||
|
|
" subplot(4,4,i+1)\n",
|
||
|
|
" imshow(cic_paint(jnp.zeros(mesh_shape), res[0][1:][::2][i]).sum(axis=0), cmap='gist_stern',\n",
|
||
|
|
" vmax=cic_paint(jnp.zeros(mesh_shape), poss[::2][i]).sum(axis=0).max(),vmin=0)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 62,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stdout",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"Validation parameter for omega_m 0.3\n"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"# So let's see what happens on a different simulation\n",
|
||
|
|
"val_snap = '/data/CAMELS/Sims/IllustrisTNG_DM/1P_20/'\n",
|
||
|
|
"init_cond_val = val_snap+'ICs/ics'\n",
|
||
|
|
"\n",
|
||
|
|
"header = readgadget.header(init_cond_val)\n",
|
||
|
|
"Omega_m_val = header.omega_m #value of Omega_m\n",
|
||
|
|
"Omega_l_val = header.omega_l #value of Omega_l\n",
|
||
|
|
"h_val = header.hubble #value of h \n",
|
||
|
|
"print('Validation parameter for omega_m', Omega_m_val)\n",
|
||
|
|
"\n",
|
||
|
|
"# For 1P_20, sigma8 = 0.96000\n",
|
||
|
|
"cosmo_val = jc.Planck15(Omega_c= Omega_m_val - 0.049, Omega_b=0.049, n_s=0.9624, h=0.671 , sigma8=0.96)\n",
|
||
|
|
"\n",
|
||
|
|
"redshift = header.redshift #redshift of the snapshot\n",
|
||
|
|
"\n",
|
||
|
|
"ptype = [1] #dark matter is particle type 1\n",
|
||
|
|
"ids_val_i = np.argsort(readgadget.read_block(init_cond_val, \"ID \", ptype)-1) #IDs starting from 0\n",
|
||
|
|
"pos_val_i = readgadget.read_block(init_cond_val, \"POS \", ptype)[ids_i]/1e3 #positions in Mpc/h\n",
|
||
|
|
"vel_val_i = readgadget.read_block(init_cond_val, \"VEL \", ptype)[ids_i] #peculiar velocities in km/s\n",
|
||
|
|
"\n",
|
||
|
|
"# Reordering data for simple reshaping\n",
|
||
|
|
"pos_val_i = pos_val_i.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
|
||
|
|
"vel_val_i = vel_val_i.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
|
||
|
|
"pos_val_i = (pos_val_i/BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
|
||
|
|
"vel_val_i = (vel_val_i / 100 * (1./(1+redshift)) / BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 63,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stderr",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"100%|██████████| 34/34 [02:29<00:00, 4.39s/it]\n"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"poss_val = []\n",
|
||
|
|
"vels_val = []\n",
|
||
|
|
"\n",
|
||
|
|
"# Loading all the intermediate snapshots\n",
|
||
|
|
"for i in tqdm(range(34)):\n",
|
||
|
|
" snapshot= val_snap+ 'snap_%03d.hdf5'%i\n",
|
||
|
|
" \n",
|
||
|
|
" header = readgadget.header(snapshot)\n",
|
||
|
|
" \n",
|
||
|
|
" redshift = header.redshift #redshift of the snapshot\n",
|
||
|
|
" \n",
|
||
|
|
" ptype = [1] #dark matter is particle type 1\n",
|
||
|
|
" ids = np.argsort(readgadget.read_block(snapshot, \"ID \", ptype)-1) #IDs starting from 0\n",
|
||
|
|
" pos = readgadget.read_block(snapshot, \"POS \", ptype)[ids]/1e3 #positions in Mpc/h\n",
|
||
|
|
" vel = readgadget.read_block(snapshot, \"VEL \", ptype)[ids] #peculiar velocities in km/s\n",
|
||
|
|
"\n",
|
||
|
|
" # Reordering data for simple reshaping\n",
|
||
|
|
" pos = pos.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
|
||
|
|
" vel = vel.reshape(4,4,4,64,64,64,3).transpose(0,3,1,4,2,5,6).reshape(-1,3)\n",
|
||
|
|
" \n",
|
||
|
|
" pos = (pos / BoxSize * 64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
|
||
|
|
" vel = (vel / 100 * (1./(1+redshift)) / BoxSize*64).reshape([256,256,256,3])[::4,::4,::4,:].reshape([-1,3])\n",
|
||
|
|
" \n",
|
||
|
|
" poss_val.append(pos)\n",
|
||
|
|
" vels_val.append(vel)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 64,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAkwAAAJCCAYAAAAybpizAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9a5Ak15UeCH7+9vCIjMhHVFVmAZUEKkGgmkSRRJFEESSa3Wh1sznqkcSV1DNS74yk6dG2mdbWpN2RbKUd087ujq3taHdm1ySz1fRYSxqtZqSelbQz0+pZUhLVaqqnQVKFbhYfBZCFRxbArEcmqiIf8fLw8Mf1/XHvOX7cMwoECCBJ26lrVpaVkREe7tev3/vd73znO1ZZlnjQHrQH7UF70B60B+1Be9Du3+wf9Qk8aA/ag/agPWgP2oP2oP24tweA6UF70B60B+1Be9AetAftB7QHgOlBe9AetAftQXvQHrQH7Qe0B4DpQXvQHrQH7UF70B60B+0HtAeA6UF70B60B+1Be9AetAftB7QHgOlBe9AetAftQXvQHrQH7Qe0dwWYLMv6vGVZL1uW9ZplWX/lvTqpB+1/Wu3BOHrQ3m17MIYetPeiPRhHD9pbNeuH9WGyLMsB8AqAnwNwC8DvAfiTZVl+9707vQft/9/bg3H0oL3b9mAMPWjvRXswjh60H9TeDcP0NIDXyrK8UZZlCuD/DeCPvDen9aD9T6g9GEcP2rttD8bQg/ZetAfj6EF7y+a+i88+BOCm+P0WgMtv9YEgWCqj6BTKUgEooVRh/lJiWpawUSKEBcuyYNseAECpjP/Wshw4jsfHK0uFsiz5eLbtArDM74BlWbW/x2WJlnndshxYlmXOoQTM99LrlmXzZ23bgWXZKIoMZVnw91qWBcCC6wawzLlZlgPbdqFUBqVy84++g867rP1uWTZs24Vl2QAs89mCz7t+bjb3QVmWcBwPjhOY4ygopfslz2f8vfR9+noKPkfLsht9UPWrOTPui/q5w5yDhbIsMJnsIUmG1lvd+7do72gcBW5YdvwO/LKAU5ZwSwULgI0SFgBLMKZTx0MBC16pYKOEVyoUsJDaDr/HKRVsAOrYN1XNAuCYMWib45eWhcRykFs2grKAUyooWFCWhZnlorSAEhYslLBKYG47SC0bbZXDLwt9TuZYyrIwtxzMbQfK8aH73cZMZegWKeyyhG3OlTq5BJCb+1JCH39ie3BQwikVukWGlspRWvoTCoCChcKykVk2prYLZVkoAUS2B8fxzXiwefzFxRyBKhCpHLYZH06p+3FuO5hbjj6uZfG1+krBQsmvQZwv3RnH9GNm2Sgsm/t+Gt8blGV56i1uxf3aO56LOo5fnnJDhGUBi+6t6ePqnHX/ZJYNZVk8vqgv6G5Y3AdVc0r9aXoPj0/zdzlOU9tBAQt+WfBYVGaclub/ljm/uWVjbjnwSgUXJUKV89gsYaGwLExtF7teG15ZoGXuX6Ry0B2xSzoX/Rl9H/QYzCwbia2P75cK7SJDUBZ8L+ms8wWfKR3fjCM9T2XZFGWR6XMtFZZUBqcs4aDkfiwsGzkslKb/CvN/XxWwzPcpy+LX6XOO6D/6bGbZ2M9mmBbpicxFALBsu+VDjg+/LLhvcwAF9E9AL7IOLNjmPHleANA80cLMIZa5R9RXyqrGmlXWj00tN+PUKxWsskRpWVDm9Woc6WOmZh6g47XMHGYDfB9i28WeFyFUBZZUiqBUfF+AakzLcy8sCzPbRWbGaahyhGWBQBVwUfL4oZ/KfGZuOVCWBQUL+24Iyw1Ba5aXJ2iZ8euUJdoq53mO+pCeGZrvaEy5Yp4tLQu5OGMbJdyy5GeX/pLDwp0ixZHKj42jdwOY3lazLOtXAPwKAIThCp555j+AbTtQqkBRzA34IIBj135XqsAVlWGpVPBUgdjx8Cm/A9t2YdsegxLXDWHbLjwvAgDM5yMolaMsFZTK8XvpBADQAfAJL4LrhgiCLhzHRxwPUBQpAMBxfLRaqzXgkqZTeF6L35umU9i2Bhx0nIceehq+30EYLsPzIvjBPUzGLYzHdzCfj5DnCabTu8jzuQFX+joty9bv95fQbp/CdHoPcXzPfM+Ez5H+RVGfgY7uA5f7KwyXEQRd9JYLeO4h7t1bxWx2gMHgOooiRVGkSNMJ4ngfvt+G73d4ciuKOQCg230Ytu0iSY4AwHyXA9cNkKYTzOcjFEVa6/M0neB3f/c/ObEx1HNb+JUPfBaZ7aBXpNicjxE7HoaODwDwVIF+niAqCwzcEJkBFV6psJHFiG0XO36HF5HMTBzbQRdDx+fXN9IpACB2PP3ZdAoPJSKVY2h7GLoBrofLGDo+Lk3voZ8n8FAig4Xr4TIy28HQrsB97HiILQcb+Uyfn8rhmQ3D0A1wpX0a11qruNy/gCQ5wqujW9hIp+jnCQZuiKEbYDOdoFek6OcJAGDghtj1IpSddbhuCM+L8K3xHXxudAtb8xH6eYIdr83fPXQDXAtXMHBD/GTYg+P4sG0PeT7DV9MJPmPGRBguw7JsTKd38fx8hJ89egO9IkWkcvTzBBvplI+Z2XrBHLghfwcAbGQxXzvdg9h2MXBDXIz3sZlNsetFiG1X90Wp8B+9/JvfP6lxdNrx8V+c/jC8UsErFXpFymNBvuaVCrFZAGLbhVcqRCrn62r+i20XuWWjaz5L107X2DNzjfz8jt/ByPHx5Oygdh4DN0Rsuxg5Ph9n6PgY2h6P8XPpBJ4AebHt4uVwGX+7fwEXkiN8bnQLkcrRUjkvZJHK4Ypr3vMiXA+X8esrWxi6AS65IU4dvYFfPLyB9SzWY96cAx1jO+hi14twpX0aAzfExItwcXaAS/GA/7aRxYiKDFFZoFekuBQP0M8TnJ+P+JzpWqnd9DsY2h4y24GnCpxPx/w+AHw/Ro7P96ll+tYrFX753vfer+HDTY6jh2wXX1r+AI8Jui46HwDYMfd8y/QdjTH5fjl+ZH/QuIltt/b7xHzfingvPU+b6aT296HjI7bdao4sFb9G57mZTmrjmubJL/U2sZUM8Znpm/ze5pimzwwdH9tBF//V2uOILQex4+EXhjv4I0dvHLsOOcb3vAjPd9b5tX6eoJ8n2HVbGLoBLtJzAQseSpyfj9ArUmzNR3ztzX7d9SIes3R99P3yOjLRPzwvlwo/M7q18N6/G8B0G8A58fvD5rVaK8vy1wD8GgD0eptlxSyVhonRwIZ2uBIw2bbCM0qDKzglHMeH64Z0ZH6vohttJmla1IkR4oEAi1mtokgZtNjmJsrvp/OybQdlqfj9kokhTJqmE9i2W12DcpHnCYOMPE+Q54kBZiWDHc+LGKRoUOQxEKM+oXNTKkeaTvgcHceH50WCpdPX77qPwPNOwbIOeBEFgDxPuN8JPBHoIWaB/mXZDEDJf89z3V90DdTftu3yNb2L9gPHkRxDZ1srpQYmusWOpx8KVejFW7BHUaHfRYs4AZ1dL+KHo/neqMg0MCpNvxZAVBaIDCtELYNV+3xmO4A5h6gskKmS/x7brj4mLP5e+VlAL4QXvMgAdgeP5glesWz9vSirz5kdO6AnvgwW2kHX9JPi76PJkfqH3r+RzxCVBV4pC8SWgyfdAC/nCS5kMV7OEwwdH5/MEziOj7JUuGx7GJhnrpfP4Sl9Pl6pEAGIFWqBffrezXTCmxzZb7JRH/Ty+cK/v4P2jueiJ/x2eb/zko0AEvWhXNRkv8qfrjguTc50HLnAyJ/N1+V3HzsfO0evSOGKY9NnR46P6+EyllfO48rwJp6d7PHfaAHp5wm6ZhH3GseIVI6b6QSP5XM+3qLz7OeJvr8GJA/cEBeSI1xIjvSGI4sZcA7cEBksDG0Pnl3UrmVRfwMw73UR57NavwGaMZGLo+yrH5ZaMu0dj6OPumHZPA9qdI6rlo2DUuGwVECp0LGrZVcCCHkM+bt8nf4dmOPK72p+/4p4TY6jJqiT918eb98NcTXqwysVPjN9s7bxaYImecxekSKCBeQJevmcr7H5HXKjciE5qgGmtTxB39Xz0WY6gVsq7Jt5iDZZi/q8efz7Patv1W9eqe47jt4NYPo
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 720x720 with 16 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"figure(figsize=[10,10])\n",
|
||
|
|
"for i in range(16):\n",
|
||
|
|
" subplot(4,4,i+1)\n",
|
||
|
|
" imshow(cic_paint(jnp.zeros(mesh_shape), poss_val[::2][i]).sum(axis=0), cmap='gist_stern',vmin=0)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 65,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stderr",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"<ipython-input-12-b14f0c49fe59>:17: DeprecationWarning: Calling np.sum(generator) is deprecated, and in the future will give a different result. Use np.sum(np.fromiter(generator)) or the python sum builtin instead.\n",
|
||
|
|
" kk = sum((ki/pi)**2 for ki in kvec)/4.\n"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"resi_val = odeint(neural_nbody_ode, [pos_val_i, vel_val_i], jnp.array([a_i]+scales), cosmo_val, paramsi, rtol=1e-5, atol=1e-5) "
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 66,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAkwAAAJCCAYAAAAybpizAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9e5BlyX0eiH2Z5+S59566Vbequrqna4ApANN4NAfTAKYJTHOAgUloRa4oalfUbpBhyd6VY9ehiH2E12E7vPLu/mHHyg5thHfttR0hi7t2hBhhhiSu15S0lCWRoiRq+Ogh0ADZANgDoQeDmkf1o7qrbtWtc+89eU6m/8j85fmdvOdWz7MG4e2MqKiq+zjPPJlffr/v9/2EtRaP2+P2uD1uj9vj9rg9bo/b8iY/7AN43B63x+1xe9wet8ftcftRb48B0+P2uD1uj9vj9rg9bo/bI9pjwPS4PW6P2+P2uD1uj9vj9oj2GDA9bo/b4/a4PW6P2+P2uD2iPQZMj9vj9rg9bo/b4/a4PW6PaI8B0+P2uD1uj9vj9rg9bo/bI9p7AkxCiD8lhHhFCPF9IcRffr8O6nH771d73I8et/faHvehx+39aI/70eN2WhPv1odJCJEA+B6AnwbwBoA/APDnrbXfff8O73H7//f2uB89bu+1Pe5Dj9v70R73o8ftUe29MEzPA/i+tfZVa20J4G8C+LPvz2E9bv89ao/70eP2XtvjPvS4vR/tcT963E5t6Xv47kcAvM7+fwPAtdO+0Out2ZWV8wAAa43/sbDWALA48WxXDkDKFEIIGFOH7wshIWUKwPptWAAWxJIJIeiT/rdd8hkRtun23bzutu/+p2Oj7dI26Njpc1IqSJkiTfthO8Zo1LWGtXU4z64mhAjnJYTDr/H36HUpk3CM9D0hEkiZ+M+I8L26Ltk1NuwYLAD6rozOzUafEex3+3rx909O7mI2GzcfemftHfUjlQ5s3htCWYPEWihrIGEhrXW/AQhrYYXATCSohQh3W1gLIwRquNckbOs9ALD+XAW7X/RZ2oeBgBEC7koJJLAQ1qIWAgYCpUhgBW0XELCwEK3XFEzYngVQigSFTJGoHICFMTUqo7Fa63Cc/BgtAAP3uxYScyEhE4UTa5HAYrXW6Juqde2sENCQqIXAXCaoINCTCaRMIaUK/aGq5jCmgrUGc7+txBoIAAksEmvCNShFgopdY2XcUQkgHBtAT6xrdC3dNRQohYQRAsXJvX1r7fll9/6U9o7Hojzt2XNpHz3b3IfEWqT+PHmj46v9O/x9QWOR/0v6/+ORKHze2tCfLPueFQLS2tBXjBDQYnE9a/x3at8H+qZG4q+n9cc4SRROsiGmpsJmNUffVOjZGolt7gtvtRCohUQpJGYiwSRRyKxBZmoMjUbPmvAd2k8V7r9EJSRKkUCnPSRJDzSOzKspclOH67Ja6/CsSFik1l0DIwQqIWHg+jTg+pmg/fnP0DHQCM6PCQBmMsFDPcVJXZ7JWAQAo0TZi0mGjPUjCaC2FnV0pZVwvaPy95WPPfQ3nSc973LhbqH9Xf+bRm16TbDXqo5+ROMHjVn90D/cViohUcgU99IBclNhvZ4jZeNtVz+iMWHu+8OJTF0/sjUGpoYKc0fz3NT+vldCoobr8+MkQ532YK1FAeB8NUPfVDD+GclNFfp8/CwZIcIzQv0u9CP/Hr+G/DrR/6WQ2KtKjI1e6EfvBTC9rSaE+EsA/hIADAab+OpX/yNImcKYCloXqKoZqmoGaw1etgbKGihT44rKIWUSAJOUCdK0j35/vQVY+N8cXLlmYUyFsjwJryg1QJJkoK5W13NYawJgybJh2K6bOHQAF8ZPQDSZJEkGpXJsbn4Seb6Fra3LAaycnNxDUexjPj8K58mPT/qJqt9fR5oO0O+PUBT7mE4foiwnqKpZOCalciRJhl5vDUnSQ56fgxDSgzn3MCRJBilTKJVDCImi2IfWBSaTO6iqGabTh6iqGbQuwvlk2RBK5WE/dH51rf02HRBMkh7qeg6tp6iqaevYAOCf//P/w/vSV5Y13oeGKscvPv3TULAYVXNsV9PwIFLfya17OMdSoUgUCpG492FRyBT7aR+5qZCbClpIFCKB8g8w3xYAaCGhrMGoLpHXGgoW+2kf4yQLE9pWNQufL0SCvWyldfzK1NAyCfsCgO1qitxf70KmuJFv4Ua+hS8/8TkUxT7eOHgV27rAzvwYWiZhO3Qe1PbSAV4aXsTzagClcuwevYmvTfawVc2Qmwp7KocWEhoC47SHW/117KUDvJifC9swpkJdl3BgPAmLBgCoqjny8Q+xVc2wXZ4gtzW2qlmzf5VjnGThetB1GNUlAGCcZK1rUMgUO+UEo7psXWMtJP7TW7/2w7ffK9554/1oK+nhv7j4Baz548hN1e5H/jV+fPye02v02UKmmMoUA7+dPAKrfLv0dyFT1yf89eb9SAuJ/bQf/qZGk9+eylHIFJfmR1iry9CXH6R93Bxs4q+dfwb/8sk9/MUH30PK9su3payBFhJTmWI3G+JX1z+BIlEYJxl+5ugN/Nx4FwP/nIyTDJU/50KmuN1bw37ax43BudC36fkbGd26nuMkQ24qvDi5g61qhqfKSXimqO2nfRQyxRG7xgBwzve1SkikfnsVu2bKmnD+hUzx79z79jvqE++m8X60LRV+bePpcL503+O+dGANNjr6EbUwfsgUhUzD88FBRtfnqVF/1EKGfkT74P0ofubousfP5J7KsZsN8U+H27g8O8RPH7/ZuV/aFh37bjbE31n/eBhzXpzcwVdO7obrw49TC4k9leMoyXCrv+6eB7bEGBkdxl0AKBIFZQ0+MzvEqC6x4/sRPya6fvRshW2xvhYfO32exkwA+JPj1zs//14A05sAnmL/f9S/1mrW2l8C8EsAsL7+MWutCYADgJ+Q3aD6ZfjJOgHStOeZF4ZKabXqt0GNgAixK/HnrK3D9210w6VMw3eESFqMTIyhOUhxq/IGNDT7q1HXJYypwkTE/2+YHYk07WMw2AzbaH4nSJIsfI7OQ+vCT2y2tW/6nJQp6rpkLJnbjmOs3LHRMQBAXZf+ewpJoth51P795twdY1WHe8ev/3usR/jIfsT70IXBpuWAQUNAM8ADCWjTvK+sgfb9C6YOD4gytZvEEwkFGx46ejhDk3DgyoMlDoxU1D+0cNvqeji1sW4wSKLPW+PAjHQTVZr2kab9AOrGaa890cb7lElYXNR1CQW3H5pQaOAC3Hls6wK5qXCDbSO3NfJaB2C2n/bxXHgubRhQ6HryASo3lb8HzSRGn6P3w30w3f3ktInhbbZ3PBZd6q1awE2+8b75/8smCO2ZFTp2ZQ0q9jf9xAClC5Ava/Fn+edHHiSk0Xs00X2qv469+XFzrNECYGG7EGFCujQ/wqiaQwuJ1N/XKjoPmpQviyQAprzWoS/R4oSOl/q5hui83/yadV2XtOOYu7bxbqkl395xP3pWDZYOfnR8B9bgof97A933XgsZQEUMtuj9ZeceX6t4u/H95sCFA7MwHvl7fpRk2MtWsF1NW+8ta7RNwD33o7rEyOjWMcXfp7Fyp5y0+iAB74GpAlB/u+20fsQ/s+zZBBaZYWrvBTD9AYBPCSE+Adep/ocA/sJpX7DWoqrmYXIm0EETNjFPvDm2w4bPu+0YDxwQ/rfWQOspABsAGE3sr1SzcDO+bHk4zYEkByzcMTlgU0f7FOHz7jUVVuJJkrWAUV2XKMsJynKC2WyMspwEoGNMjSRRgZnKsiFGox1U1Qyz2SGEkAEoceBG2y2K/XAcSZIhy4YBKNE5ZdkQSZJhMDgHKRMolTOAasP1bFi5CsqHgej6al2Ec6/reQRAawYo0Qks32F7x/2IWpEojD2IoYdMCemYGA9eNBBWLspP/Ptp3z0cftAP34VtsQqAAzq5qQJrBTSTCw02GgJgDx9NJsSqjKXrW1omyD0YbQ1qMsF+2sdXsiGybIi6LvHZ/jq+5ftdvPrjq9GxVBjl56B1gZvFPratwW5v1a2WGAhyoM+tyugcaKIL5+r3t5etYCwVvt1fh5YJdvz1UrANy8W
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 720x720 with 16 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"figure(figsize=[10,10])\n",
|
||
|
|
"for i in range(16):\n",
|
||
|
|
" subplot(4,4,i+1)\n",
|
||
|
|
" imshow(cic_paint(jnp.zeros(mesh_shape), resi_val[0][1:][::2][i]).sum(axis=0), cmap='gist_stern',\n",
|
||
|
|
" vmax=cic_paint(jnp.zeros(mesh_shape), poss_val[::2][i]).sum(axis=0).max(),vmin=0)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 67,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"res_val = odeint(neural_nbody_ode, [pos_val_i, vel_val_i], jnp.array([a_i]+scales), cosmo_val, params, rtol=1e-5, atol=1e-5) "
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 68,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAkwAAAJCCAYAAAAybpizAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9a5Ak15UeCH7+jAiPiIx81SPxSIJVIFAkUU2imkTxgW6SUjfV09JOszVqaSSzkWYlmdY0ttpds5HZ9u7M/h7N7ki2s2srzbbtS7MzvZJ6NdOSiS1NN6e7NUKTLJAskl0gWQBRRTABVBaqIh8RGeHhz3v3x73n+PGbkQBIAEnaTl2ztHx5uF93v+73u9/5znc8rTUetAftQXvQHrQH7UF70B60k5v/k+7Ag/agPWgP2oP2oD1oD9pPe3sAmB60B+1Be9AetAftQXvQ3qI9AEwP2oP2oD1oD9qD9qA9aG/RHgCmB+1Be9AetAftQXvQHrS3aA8A04P2oD1oD9qD9qA9aA/aW7QHgOlBe9AetAftQXvQHrQH7S3aOwJMnuf9kud5L3qe97Lneb/+bnXqQfsfV3swjh60d9oejKEH7d1oD8bRg/ZmzftxfZg8zwsAvATgFwG8BuBrAP6i1vq77173HrT/f28PxtGD9k7bgzH0oL0b7cE4etDeqr0ThukZAC9rrW9rrQsA/wjAr7w73XrQ/kfUHoyjB+2dtgdj6EF7N9qDcfSgvWkL38FnHwbwqvj9NQBX3+wDnc6K7vfPAAC01tBatb4WADxoaHgYBDEAQKkKgAbgwfN8BEEEw4pp5zvgeZ44Gv1Mx9FiPx5/V6oGoOF5PgAPQRDx/7RWUKqG5/n2d7MPparW/oIggu9HiKIePC+A7wNVVaKqMihVQ2v1phfS9wP4fsT9quvCHkO1+mb6Ya4B/c33Q/tVwPMq1HUCrWtUVQ6taz4/8x3QWvF+fD+A5/nivGp7TvYK8nGb60rXsrmGPubzN5BlE3nxf5T2I42jOOzpXmeAUGsEWiHSCgE0Aq3hQ8PXmu/83A9Rez58cU4aQAXztwDH2VWPxpLdVnsePG229e13BQ/K81DDgwYQaQUPQG3/tvCbx8pcpeYznu1fV9cItOIVSwUPsyCCHw/4fmRVjrU6N/uw5wcACh60PZaCh8rzUPgBAj+E1hoLAKO6QE9Vtod07h5Kz0fl+Vj4AUrPR88LEIYdBEEH9KyYsaOgdY1MK4yqgq9VoBVCaNTwUHk+Cs9H6fl8rh1Vcz810PoftcBeR3qCMj+EAjBP74+11mdOuvdv0n7kd1E/iPV61ENHK76/gdaI7LPq8Tl4KDwfyjPnC4DvIW3nAXx/zfjTJ65E6bM+n7/5nAYQ2nFE+8o9Hxpea/zS/yr7va8qhFohsOO29jxMgxjzeIiZqrBW50hUhURVCGzfqPf0c23PLfMCZH6AuR+ioxViVWOgSnTE+8vXpt/0mdwLeBzUUYIoSqBUCaVqZFWGrq6hYe75ap03zynMWOLz8Xwznr3mmaJnmY7H15C+xLMKAIXnY1xlmNXFqbyLAGDkR/p8GCO2946e6Vpr1NCo7HYhgMjz+Xl1z4M6rDwPCobJ8MUz776paHtfa2h7zWgb3+6TrmVpx5H7LlCeeR/U8NBXJUJ7PPpM6oe4GyVIVIXVukBH1Yh1zfvXXvsyKwC1HUc5jSNVo6Nr9FTderb4fO34y70Ate3PXthFFHahVI1UK2zYcVTbZ2FYl+b9LZ41Ov9ajCOI+0H9pZEsrz1dP9+Op9LzcacucKiqY+PonQCmt9U8z/sbAP4GAPR66/j5n/+P4fsRlCqR51PUdYGqynBNVejAvOgB4FKUwPdDC5hMi+MBkuQMtK5R1wUDLdkIGFDTukaeH4EmgyhKEIZdBEEMz/OR51MoVVkgEqPTWQEA1HXBfSNQQscjIBQEEaKoj3PnLmMwOI+zZx+BUl1UVYWjozuYze6iKGaoqgxlmaKuS+oVACAIYiTJJuJ4gF5vHdPpa5hOX0dVLVBVOQOjTmfFbnsGnc4QKyuPQKmqdQ18P0IYdhCGXQAesuwARTHD/v4tlOUceT5FVWUoihl8P4Tn+eh2VxHHA74WBNKqKgcAO5HGCIIYdV2gLFNk2SGKYsafAYA/+qP/w7s4Yo43OYb6UR9//sIvIoLGqMqxXc55u0grRKrGSJWItMIkiJH6IVIBYFI/xE48QKIqbFYZUj9EKabzCGbSpIc79UNEWmGrTBGpGomuMQ67GIddpF4AANgu50jqEqUfIPVD3LJjiPoTQaOEx/9P/RAX8ymSuuT/3Ug28NzgPK48/Axms7vY33sJ28UM28UMEz9CGkS8PfUt0go78QD3Vx9DGHYRRQlemuzgC4evYFQXSOoSu3Gfzz/1Q1xPNrETD/ALo20+56rKoVTZep7o3tZ1geHe97FZZdiqFhhVOTarDKVvzn03SlrXg/q3VaYoPR+TIOZrMA672A17uLLYw3Z+BAAo/QCRfWb/9sv/8ofvyoA5oclxtBl08PfOfxQrdYFIK37vyHufqAqRVgz6xmEXpQWIpedj4Yfo2W1SP8Q0iJGoCj0LUORnqdHf6TP0VXo+tosZElXxZLUbJa3Pyv29Gg94HG1UGRL7npwEMV7oreMfnPkQrs7ewF88uMWfXdboPG52RvgHZz+MwPPheQG+cPgKvjD5IV8f6mNlz30nHmAnHuC3Vx9D6fnwPZ+PcTGfYqtMW8eItMKzs7vYrDJcyKeItOI+07WdBDH2wi4/jxE0NqoMAFB5PkJ73RZ+iIldUAPAo/a6AcC/N37xbYyEd9bkONryI/z22gWM7HWSzyb9PFMV9rXCuudjzfMxDru8Lzk+ElVhEsSYBLF5fsVYcceR3L8ck6XnY6tMcaAV1uzvNI7cMUD/mwQxLuZTPiZg7sdOPMAfDh/CpewQn57dbZ3bsv7Tu++/3HiC36ufP3odvzzZOXYe1NdJEGMcdvGHgy2kQcTv1AgaW2WKzSrj917qh4hUjQvFEUZ1gYt2HMlG45Te/dTo/tBYpOds2TUAgM9NX1t6798JYHodwKPi90fs31pNa/0bAH4DAFZX36cJnBAQMkxHiE/ZCZiYEZrUff94F+Wk3v57zQwKtboucE2V/Pun7AWu65IZFjqWBA11XR4DY/S77JvvB62JRmufGSKlKlRVjqrK+G90zp7nIwy76PfPWVYLzPoQGKEv058CWXaIus5b50jHpv5EFmhSP8Ow05oMTb/MBBUEHd4/nRNdRwAtsNqcU906rr0yx+7Fj9DechzJMXS2t64jcTz3RQIfKG3/6cGQDw5NRomqzETu+YDz0ooE6I483zywdtKnh04et/R8lJaxAXDs5cnbqOXXqfQDAyaiBN3uKsoyxcSPMA67J050EtDF8QAAUFUZIlXzi6UM2/2MVG2An1aYzd7gv19TJfcZAFIvQBpE+KRd5U3CDhJdI5Lg1F6PpC6ReAESLwACtK4NXX+6FomqkOjm2hLoit7Z+AF+jHfRxc6wdVB6kfL52T67f6NtZZPbhksmzGXbuftadp9p8lo20Y3qApGqETqT0F7YxU7Ux9n+Wdyqstb/iCGjPkrAwsesS2xXRxipkp8Vd8KOtMKoyrEFD5fTPb6PxGRtVhlGdYHUC3hsR6rmsbHsusoxQhMeVPMcy+3KJdfDvX8/ZvuRx9FTUY/HkbyXy/qyL667fE5OGnv0M32X27l/l8cvPR8Q+3vLcbTsXSVAO4FfOn7qzMnHxrKqsakqXlzJfbrnmajKgJ/iiBcPchz1bN8r55lzr4H8mxxj7vO17L4su18nUZTvBDB9DcAHPM97P8yg+ncB/KU3+4DWGlWVHZvkaaIPw65gYmz4wV4wpSoLBLRlf3L+PG1TFHMAusWY1HXBq93S81H7EYKA9lkyuAjsikX2D0ALSBC4IBaH/l9VmQVFCbNJRTFDlk2QZQf2nAp7zgEzU53OCs6d+yDyPMXh4Q/heT5fhwaAKRTFDEpVmM/v8d/DsItud7V1DZSq+POrq4/B90N0OisMAJWqcI0mMFXjI37AIDMIYkRRAs/zUZY
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 720x720 with 16 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"figure(figsize=[10,10])\n",
|
||
|
|
"for i in range(16):\n",
|
||
|
|
" subplot(4,4,i+1)\n",
|
||
|
|
" imshow(cic_paint(jnp.zeros(mesh_shape), res_val[0][1:][::2][i]).sum(axis=0), cmap='gist_stern',\n",
|
||
|
|
" vmax=cic_paint(jnp.zeros(mesh_shape), poss_val[::2][i]).sum(axis=0).max(),vmin=0)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 69,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"k, pk_ref = flowpm.power_spectrum(\n",
|
||
|
|
" tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), poss_val[-1]).reshape([1,64,64,64])),\n",
|
||
|
|
" boxsize=np.array([25.] * 3),\n",
|
||
|
|
" kmin=np.pi / 25.,\n",
|
||
|
|
" dk=2 * np.pi / 25.)\n",
|
||
|
|
"\n",
|
||
|
|
"_, pki = flowpm.power_spectrum(\n",
|
||
|
|
" tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), resi_val[0][-1]).reshape([1,64,64,64])),\n",
|
||
|
|
" boxsize=np.array([25.] * 3),\n",
|
||
|
|
" kmin=np.pi / 25.,\n",
|
||
|
|
" dk=2 * np.pi / 25.)\n",
|
||
|
|
"\n",
|
||
|
|
"_, pk = flowpm.power_spectrum(\n",
|
||
|
|
" tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), res_val[0][-1]).reshape([1,64,64,64])),\n",
|
||
|
|
" boxsize=np.array([25.] * 3),\n",
|
||
|
|
" kmin=np.pi / 25.,\n",
|
||
|
|
" dk=2 * np.pi / 25.)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 70,
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"[<matplotlib.lines.Line2D at 0x7f7e841d4550>]"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 70,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAXcAAAD8CAYAAACMwORRAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/d3fzzAAAACXBIWXMAAAsTAAALEwEAmpwYAAAuR0lEQVR4nO3dd3gVVf7H8ffcm0J6AkkIpJKEktCL9N6RoojdFRugrrtr2fbboquurn3Xtq4ixV2xoLtiAQTpIEgJvbdAgBAgIZCQXu78/hgQUMDccJOb8nk9Tx7JZObMuQofjmfOfI9hmiYiIlK32NzdARERcT2Fu4hIHaRwFxGpgxTuIiJ1kMJdRKQOUriLiNRBHu7uAEBoaKgZFxfn7m6IiNQq69evzzJNM+xSP6sR4R4XF0dKSoq7uyEiUqsYhpF2uZ9pWkZEpA5ya7gbhjHaMIzJOTk57uyGiEid49ZwN03zK9M0JwUFBbmzGyIidY6mZURE6iCFu4hIHaQ5dxGROqhWz7nvO5HHir2ZFJWWu7hnIiK1W41Y515ZM9cd4t87p+Bhd5AUdA3Dm3dlQMsmJIT5YRiGu7snIuI2tTrcHx0Uz/ayfLae/o49LGL3Xm9e2ZxAYNk1DIoZSt/mYfRMDCXIx9PdXRURqVa1Otx9N03nw4OryY3vz9rQaBYUnmLlsbX4lKYze3MGH61No0HEHGJ82zI4rhdDWsXSLioYu02jehGp2wx3brNnGMZoYHRiYuLEvXv3Ot/AnvmQMh0OLIfSfLB7YcZ0pzRhIEbCUOZll/KX9ZMoNQsxTRuOwmg8SlpyTegQhrVsTd/mYUQENXD55xIRqQ6GYaw3TbPLJX9WE/ZQ7dKli3lVtWXKiuHQati30Po6scM6HhhJacIAtkS0YHF5MYvT13KkYA+emfeTnRWHzfsYTRsfo09UL65Nas01cQ1p4Gl3zYcSEalidT/cfygnHfYvgr0LIHUpFOeCYYfobuQ0641P86HstzXn7+veZvXpDwBwlDTCLGhB84DOjGo+UA9mRaTGq3/hfqHyUjiScnZUvwAyNlvH/cIxEwZyMLojy+ww7/B6dp3eiMNhkLv7ccBO4/A0OkZHMLplV3o3b6wHsyJSo9TvcP+hvBOwf/HZsF8EhdmAAZGdKU0YwMEmrWkQNpjl+7N5fdcDFHAEs9yH8oJEmnq1Z0BMb0YmJ+vBrIi4ncL9chzlcHTT+bn69BQwHeATAgkDyY7rySofX+Yc2cSGE2socGRTmtuOovTbCfb1JDk+g2ub92Bwqxg9mBWRaldjw/2qV8u4WkE2pC6xRvT7FkLecet4RFvMhMGkNm3NqaCWpGc1Yt6uHXxb/Bimaae8IIYQow09mvRgTHJXujUL1YNZEalyNTbcz3HbyP1KTBOObT0/fXN4NTjKwCsA4vtRmjCQ9cHhzD66g1VHV5FZkgpAYfqteBR2onMzb66JD2BUcpIezIpIlVC4u0JRLhxYZoX93oWQe8Q6HtYKEgdzMqYby40yPEvbknKghHmHPyHf/zPKi8PxKU2iQ1hXxrToTf+W0XowKyIuoXB3NdOErD3WUst9CyFtJZSXgKcvxPWBxMEcbpLMZ8d3sShtBWn5W3FQiunwpHDfE3SICqdrvDeDW8XRIbqhHsyKSKUo3KtaST4c/Pb8g9lsa4qGhvGQOJji+H6s9fZhRfp+vAv6sXxPJntsL2HzPo69uAUtArswIqEvI5Nb6cGsiFSYwr26ndx//qHsgeVQVgh2b4jtCYmDIXEwn57cw5d7F7Pj1DpKyAWg5HQX4hz30LdFKD0Tg+kR31gPZkXkshTu7lRaBIe+Oz+qz9xlHQ+KhsRBOBIGsickki8OriPzlC8ZR1uyLu0oXvHPYhbFEt2gA/1jejOuzTUkhvvrwayIfK/GhnuNWwpZHU4fPh/0qcug5AzYPCC62/ej+sO+4by4+m02nFhDrsN6cOso88c/92cMjO1Dn8RQ+rYMx9+7Vhf1FJGrVGPD/Zw6PXK/kvJSOLzWKouwb6G19BLAvzEkDILmgznepA1zjmxmfupyvM8MZ+N+Twq9U/ButJLYBt0Z22I4N3foqBU4IvWQwr22OHPMKo2wd4H1z6LTYNggsvPZUf0QShu3Y9qmeXywewqnyq0Ht47iJkR5duXu1vcxvHUkIX5e7v0cIlItFO61kaMc0jecL3iWvgEwwachJAyEpFEciWjD+7sWs+DgQk4WnSJn7yPYbTaSEw/QP6EV4zv1IixAq29E6iqFe12Qf/JsaYSF1si+IMtaV99iGLQeS1GzPuzLtjN7yxE+PDYB05aHoySEMFsXhsYN4Z4u/WgS5OvuTyEiLqRwr2sc5ZC2CrbPgh1fnA16P2g5HFqP5VRUFz7es5Sv9s7ncNFmMMooyRpAW99bGdY6nGFtGhPTMMDdn0JErpLCvS4rL4O0b88G/ZdWCWMvf2h5LbQeS15MVz7ds4xDx4JZs9uL3TmbaBA5gyCzA30jBzKh81ASw0Pc/SlEpBIU7vVFeRkcXG4F/c6voPAUeAd+H/QkDGBx2g7e3DiVfXlrMI0izHJv/MrbcWPcQ9zQoSUJYf7u/hQiUkEK9/qovNQqdHYu6ItywDsIWo2E1mMpie3J7NTvmLlzLntztpK9+1eAnejonbSPCuG+ziPpENlYL02J1GA1Ntzr5UtM7lBWYu0lu30W7JoDxTnQIBiSRkHrsZhxfTme72Detgze2vMohbb9mA473mWt6BLaj7s7jKR7XLSCXqSGqbHhfo5G7tWorBj2L7GCfvdca/NwnxBIGg2tx+KI7c3SIxuZsWU2m7JXUGqcpDSnA+FF9zKibQS9WjSgT3ycgl6kBlC4y6WVFlkvS50L+pI88G10NuhvwIzpyepjO0g5kMP6fQ1YlbYDr7i/41HajDbBvflZ25EMbZGETSWLRdxC4S4/rbTQqmS5fRbs/hpK88EvDJLGWA9jY3uy/9RxXl37PmtPLKOAwwAYJdEMaPhLOjdNIsC3jHD/BsSENCQswFsVLUWqmMJdnFNSYL0Vu30W7JkPpQXgFw7J11lBH9OdnScPMnXjV6zKWEZ26p0UFXvjFboQ77CFmGW+OEpDsDtC8LGF0cx2E40D/An0KyLM35fo4IY0DvQhLMCb8ABvgnw8Nc0jUgkKd6m8knzY+83ZoP/Gqk3vH2EFfZsbIKor5Rhk55ew8sh61mWsIz3vKCcKj3G65DiFjtPEF7xMVl4pWd4zsAWtxSz3xlEaglkagqMkFPPkaMICvAkOzCPMz5+mAaEXhX94oDdhAd6E+XvjYbe5+9+ISI2hcBfXKM6DvfOtoN+7AMqKIKApJI+BJu0hJA6CYyGgCdisEDZN8/tRecqxFNYf28KB00c4ciad4wUZOBw2Bgc/R+aZYlbnP88Z23ZweFFeGoxZGkx5YTQlWUMACPAtYWirZoxq35ReCaF4eSjopX5TuIvrFZ+B3fOsoN+3wNpD9hy7NwTHQEjs+cAPiTv7FQsNgi7Z5JqMNew7vY/0vHTSzxzlcG46QV7h3B73BMdzC3lzz/3kF5dTfLodDUq6MLxla0a2a0rPhEZ4akQv9ZDCXapWWQnkHIZTB62v02lnf332n0WnLz7fJ+SCwL8g+INjrR2qPH5csrjcUc6nez5lTupcNmVutA4WRVOQOYAAR3uGJUcwsl0TeijopR5RuIt7FZ7+ceCf+/70oYtH/YYNAiPPB39w3Plfh7WEBkFk5GUw7+A85qTOpWvIOI6mt2TB7t2UeG/Hr6wjI5ITGdm2Kd3jG2qOXuo0hbvUXA4HnMm49Ij/dJr1s3PsXlb5hI4/g/gBYLN/P6c/Y8dHvLDubxjYceQ3p+h0OwIcHRjROo5RbZvQtZmCXuoehbvUXqWF1uj+1EHrzdotM63Kl4GR0P426HA7NErANE12n9rN3ANz+Tp1HscKMrDTgJLUP1NQ7EGovxfD20Qwsm1TujZriF0vXkkdoHCXuqOs2HrJatMH1sYlpgNie0GHO6zlmd7+OEwHmzM3s/PkTsYm3MLS3Sd4YeP/kZljUHS6HcFGa0a2jWR4mwiuiWuoOXq
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 1 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"loglog(k,pk_ref[0])\n",
|
||
|
|
"loglog(k,pki[0])\n",
|
||
|
|
"loglog(k,pk[0],'--')"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"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.8"
|
||
|
|
}
|
||
|
|
},
|
||
|
|
"nbformat": 4,
|
||
|
|
"nbformat_minor": 4
|
||
|
|
}
|