mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-02-23 18:10:55 +00:00
914 lines
295 KiB
Text
914 lines
295 KiB
Text
|
{
|
||
|
|
"cells": [
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 1,
|
||
|
|
"id": "2bb083b7",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stdout",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"%pylab is deprecated, use %matplotlib inline and import the required libraries.\n",
|
||
|
|
"Populating the interactive namespace from numpy and matplotlib\n"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"%pylab inline\n",
|
||
|
|
"%load_ext autoreload\n",
|
||
|
|
"%autoreload 2"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 2,
|
||
|
|
"id": "458efc98",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"import jax\n",
|
||
|
|
"import jax.numpy as jnp\n",
|
||
|
|
"import jax.lax as lax\n",
|
||
|
|
"import jax_cosmo as jc"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 3,
|
||
|
|
"id": "da00e618",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"import tensorflow as tf\n",
|
||
|
|
"import flowpm\n",
|
||
|
|
"from flowpm.tfpower import linear_matter_power\n",
|
||
|
|
"import flowpm.scipy.interpolate as interpolate"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 4,
|
||
|
|
"id": "b588c18c",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"from jaxpm.kernels import *\n",
|
||
|
|
"from jaxpm.painting import *\n",
|
||
|
|
"from jaxpm.growth import *"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 5,
|
||
|
|
"id": "f0cba7a4",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"# Below are a few parameters\n",
|
||
|
|
"box_size = [100.,100., 100.] # Transverse comoving size of the simulation volume\n",
|
||
|
|
"nc = [100, 100, 100] # Number of transverse voxels in the simulation volume\n",
|
||
|
|
"batch_size = 1 # Number of simulations to run in parallel"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 6,
|
||
|
|
"id": "8dad416f",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stderr",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"2022-02-14 00:21:13.509002: E tensorflow/stream_executor/cuda/cuda_blas.cc:232] failed to create cublas handle: CUBLAS_STATUS_NOT_INITIALIZED\n",
|
||
|
|
"2022-02-14 00:21:13.509025: E tensorflow/stream_executor/cuda/cuda_blas.cc:234] Failure to initialize cublas may be due to OOM (cublas needs some free memory when you initialize it, and your deep-learning framework may have preallocated more than its fair share), or may be because this binary was not built with support for the GPU in your machine.\n",
|
||
|
|
"2022-02-14 00:21:13.509034: W tensorflow/stream_executor/stream.cc:1260] attempting to perform BLAS operation using StreamExecutor without BLAS support\n"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"ename": "InternalError",
|
||
|
|
"evalue": "BlasScal failed : in.shape=[1,100,100,100] [Op:IFFT3D]",
|
||
|
|
"output_type": "error",
|
||
|
|
"traceback": [
|
||
|
|
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
|
||
|
|
"\u001b[0;31mInternalError\u001b[0m Traceback (most recent call last)",
|
||
|
|
"Input \u001b[0;32mIn [6]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m 6\u001b[0m pk \u001b[38;5;241m=\u001b[39m linear_matter_power(cosmology, k)\n\u001b[1;32m 7\u001b[0m pk_fun \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mlambda\u001b[39;00m x: tf\u001b[38;5;241m.\u001b[39mcast(tf\u001b[38;5;241m.\u001b[39mreshape(interpolate\u001b[38;5;241m.\u001b[39minterp_tf(tf\u001b[38;5;241m.\u001b[39mreshape(tf\u001b[38;5;241m.\u001b[39mcast(x, tf\u001b[38;5;241m.\u001b[39mfloat32), [\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]), k, pk), x\u001b[38;5;241m.\u001b[39mshape), tf\u001b[38;5;241m.\u001b[39mcomplex64)\n\u001b[0;32m----> 8\u001b[0m initial_conditions \u001b[38;5;241m=\u001b[39m \u001b[43mflowpm\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlinear_field\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnc\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 9\u001b[0m \u001b[43m \u001b[49m\u001b[43mbox_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m 10\u001b[0m \u001b[43m \u001b[49m\u001b[43mpk_fun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m 11\u001b[0m \u001b[43m \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbatch_size\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 13\u001b[0m initial_state \u001b[38;5;241m=\u001b[39m flowpm\u001b[38;5;241m.\u001b[39mlpt_init(cosmology, initial_conditions, \u001b[38;5;241m0.1\u001b[39m, order\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n",
|
||
|
|
"File \u001b[0;32m~/repo/flowpm/flowpm/tfpm.py:62\u001b[0m, in \u001b[0;36mlinear_field\u001b[0;34m(nc, boxsize, pk, kvec, batch_size, seed, dtype, name)\u001b[0m\n\u001b[1;32m 59\u001b[0m whitec \u001b[38;5;241m=\u001b[39m white_noise(nc, batch_size\u001b[38;5;241m=\u001b[39mbatch_size, seed\u001b[38;5;241m=\u001b[39mseed, \u001b[38;5;28mtype\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mcomplex\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m 60\u001b[0m lineark \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mmultiply(whitec, (pkmesh \u001b[38;5;241m/\u001b[39m\n\u001b[1;32m 61\u001b[0m (boxsize[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m*\u001b[39m boxsize[\u001b[38;5;241m1\u001b[39m] \u001b[38;5;241m*\u001b[39m boxsize[\u001b[38;5;241m2\u001b[39m]))\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m0.5\u001b[39m)\n\u001b[0;32m---> 62\u001b[0m linear \u001b[38;5;241m=\u001b[39m \u001b[43mc2r3d\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlineark\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnorm\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnc\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mnc\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mnc\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdtype\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdtype\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 63\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m linear\n",
|
||
|
|
"File \u001b[0;32m~/repo/flowpm/flowpm/utils.py:271\u001b[0m, in \u001b[0;36mc2r3d\u001b[0;34m(cfield, norm, dtype, name)\u001b[0m\n\u001b[1;32m 268\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 269\u001b[0m norm \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mcast(norm, dtype)\n\u001b[1;32m 270\u001b[0m rfield \u001b[38;5;241m=\u001b[39m tf\u001b[38;5;241m.\u001b[39mmultiply(\n\u001b[0;32m--> 271\u001b[0m tf\u001b[38;5;241m.\u001b[39mcast(\u001b[43mtf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msignal\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mifft3d\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcfield\u001b[49m\u001b[43m)\u001b[49m, dtype), norm, name\u001b[38;5;241m=\u001b[39mname)\n\u001b[1;32m 272\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m rfield\n",
|
||
|
|
"File \u001b[0;32m/usr/lib/python3.10/site-packages/tensorflow/python/ops/gen_spectral_ops.py:809\u001b[0m, in \u001b[0;36mifft3d\u001b[0;34m(input, name)\u001b[0m\n\u001b[1;32m 807\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _result\n\u001b[1;32m 808\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m _core\u001b[38;5;241m.\u001b[39m_NotOkStatusException \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[0;32m--> 809\u001b[0m \u001b[43m_ops\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mraise_from_not_ok_status\u001b[49m\u001b[43m(\u001b[49m\u001b[43me\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 810\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m _core\u001b[38;5;241m.\u001b[39m_FallbackException:\n\u001b[1;32m 811\u001b[0m \u001b[38;5;28;01mpass\u001b[39;00m\n",
|
||
|
|
"File \u001b[0;32m/usr/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:7186\u001b[0m, in \u001b[0;36mraise_from_not_ok_status\u001b[0;34m(e, name)\u001b[0m\n\u001b[1;32m 7184\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mraise_from_not_ok_status\u001b[39m(e, name):\n\u001b[1;32m 7185\u001b[0m e\u001b[38;5;241m.\u001b[39mmessage \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m (\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m name: \u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m+\u001b[39m name \u001b[38;5;28;01mif\u001b[39;00m name \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m-> 7186\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m core\u001b[38;5;241m.\u001b[39m_status_to_exception(e) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;28mNone\u001b[39m\n",
|
||
|
|
"\u001b[0;31mInternalError\u001b[0m: BlasScal failed : in.shape=[1,100,100,100] [Op:IFFT3D]"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"# Instantiates a cosmology with desired parameters\n",
|
||
|
|
"cosmology = flowpm.cosmology.Planck15()\n",
|
||
|
|
"\n",
|
||
|
|
"# Create some initial conditions\n",
|
||
|
|
"k = tf.constant(np.logspace(-4, 1, 128), dtype=tf.float32)\n",
|
||
|
|
"pk = linear_matter_power(cosmology, k)\n",
|
||
|
|
"pk_fun = lambda x: tf.cast(tf.reshape(interpolate.interp_tf(tf.reshape(tf.cast(x, tf.float32), [-1]), k, pk), x.shape), tf.complex64)\n",
|
||
|
|
"initial_conditions = flowpm.linear_field(nc,\n",
|
||
|
|
" box_size, \n",
|
||
|
|
" pk_fun, \n",
|
||
|
|
" batch_size=batch_size)\n",
|
||
|
|
"\n",
|
||
|
|
"initial_state = flowpm.lpt_init(cosmology, initial_conditions, 0.1, order=1)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 7,
|
||
|
|
"id": "5ad12a1e",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"@tf.function\n",
|
||
|
|
"def solve_tf(init_state):\n",
|
||
|
|
" final_state = flowpm.nbody(cosmology, initial_state, linspace(0.1,1.,40), nc)\n",
|
||
|
|
" return final_state"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 8,
|
||
|
|
"id": "7d4780a8",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"final_state = solve_tf(initial_state)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 29,
|
||
|
|
"id": "6a94dae9",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stdout",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"1.3 s ± 2.59 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"%timeit final_state = solve_tf(initial_state)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 9,
|
||
|
|
"id": "e5148fed",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.image.AxesImage at 0x7fd56040fdc0>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 9,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD7CAYAAACscuKmAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAA19UlEQVR4nO19a4wl13HeV33vzM7MzuzOvnf5iEhLtCz5Icoh9IACg5Ysh3EM80egwA4cMIEAAoYTyIgDi0qAAA6QgEEAw/lhBCBixwTs2BYs2yQUxzbBWFZsC5KoNymKWkkkV8td7nv2MTs7r678uHfmfFV9z5ne1x1aXR8wmNv3dJ+uPt3ndtWpqq9EVREIBL73UW23AIFAYDyIyR4IdAQx2QOBjiAmeyDQEcRkDwQ6gpjsgUBHcFOTXUQeEpGXRORbIvLYrRIqEAjcesiN+tlFpAfgmwA+COA4gM8D+DlV/fqtEy8QCNwq9G/i2HcB+JaqfgcAROT3ATwMIDvZJ6spna7mBhv+R4a2/c+P8EZFysjtCAjiPkVsG2/Whf2qtK2Tk7b7Xmqrrq3Y49bWqQ+6Tt9/XXupU//ZMSmMqj+GT3eDQywZ+dXJLj3ar9eznazTePB4N+R145MVyghY2NH3T/tW7lwsC19bXRi4qnDu4jPNz5zrYyjW0tplrNRLIwfkZib7nQC+S9vHAby7dMB0NYf37noYAKDr7oFdXd386Nv4gZCpHWm/1bX8yUqDJoUHnW+Ye/iEtnVlJb/f9FTq7t47TNvqXJr8U18/bk+9cDH1QdeJCfeDsbhIJ7P3VddoTMwEcePN1+LGUfjHqtRH4Yexmp5OTTvStdRXr9r9Znemjb3ztv/zC+lUS9fS9/7HjsffTyT6wZA+Pe50jwZ98qRdN01C16LTO2wb/WDr0lL6fHUJdsc0PjIzgyz4/vnxZhl3ODn6gzH42zN/kO36Zib7qF+PxgwTkUcBPAoAU9XOxgGBQGA8uJnJfhzA3bR9F4ATfidVfQLAEwCwu7dfN97a5s0I+9Zs/IpMTGx+rOkXXhrqc8v1Rvq192ov9ynuja3rGTV73b4J6ouXUx9fO2radtCbTFdWTRu/lXUx/arLlP+Fz6vxQmYDv2k8+G0uXjUlFVGqdC6t/fiSHAW12Ggbvu3a8ubn6vKibdu3J7VdupK+X3XjRn34e2HeqHOz6Rj/hl4mzXLKalK4SlrF5IRtY62opVlpNLMCqoP7zfbyvWn7/Fut/Ic/fW7Y4Ewh7q/VWUfj8wDuE5F7RWQSwM8CePom+gsEArcRN/xmV9U1EflXAP4cQA/Ab6nqC7dMskAgcEtxM2o8VPVPAfzpLZIlEAjcRtzUZL8ZeHvbrJQ6sM3H+zVsaloH8Cv6bG/LRLvLrpeXzbY5H8vv1wrYdnNtukxrFQ37kt2KZCv79Q0eq4I3oaJV38a1gNYtnJuI7XT7vVv5J/H9vTAgG7tx3/k49wzoDHskqO3E6fy5/DoOjw8/A8sFu3/Orpbr/FzaWHP3jO8nrw+41XKzHnPNP1e0RrJ71+bnpbccMPu99uO0HuOGe+WlwXqEHrs9NnsgEPg7hJjsgUBHMF41fqIPOTxUTS5ctG2kdqtXOVm9ywW2wLnRfEACg9RW9ao0q5WuD1Z3WfUqhRz7/qUQtGMi71bpXOL69zLnzk0mhDeTuEffv4lqM/LZ7/laqv37bP/s9qvyY2oiJ53La31n2u5dSc+E3mMDlXrnL6U2dpMBJkCGTaiG2UGq+vqsVcEZ1bJzI/aS2s0qvizZZxikujdcxvxMH0nutXM/aOXok8du+rS9Zxe+f7Dv2tfy0YTxZg8EOoKY7IFARxCTPRDoCMZqs2uvQr174NboOZtGr6RwSG/LmnDLUgIK2aUNK7oeHbLZCBU12/n+jf3q1w4KCTpKNqv0J7L7NeTiPnwSEYOOE7btS5lhjTEYbes3QovJvbR+aN609c7SmgaFuvr1Bl7TqBZteG/vSkpWWZtL57p6h01i2f08udH8WhCFtwqvCfh1j3661/Wkve8Tpy7Rhp0yK4dSCG7/SpJj7c7dZr/eUnom+q+dt+dm1/KFFGp95NP2vqztTmNwwYXL7lgY3JuqkBsWb/ZAoCOIyR4IdARjVeOlVlRXBypvI+uIo5vccfWFhdRGWV2y00U6LaZcaZ9RllWLC6ZAtWvO7stuolIkXMm1R1FyuuZkZLWe5fJECAUV35Jq5N18FeVzN8Yq4+r07jvZQa6x01Z91ivkJ+Isw1IGnCe2WExutF4/jdv85y/Y4+j5aLhBKa9chTImfU75maRaT5yxarZx2bkxmCR3m06l+7d2wJoaytF1h/eYtupCMnOE7h+7FAGgt5DG4OCCTRff6L9azrtl480eCHQEMdkDgY5g7KvxG9FJ/VOOWIFWTdf37jJNFat3FF2nPimBVkqrSaeCk6paSpgxqvqipVAyqnCBO81EybVVuQETXSY9Wjl2q8OsPuOgjVzD6XPpM12bOQaA7k7jUy07LjyOZuSkDUfXZMan4NVQijprRBQyR58fD1pZN2+lNb+iT+Ndii40Xgbn8blGMroITk4uqn0SC+0rO5NqPbPmoi8n0vmqc1Y9NyYh3SdP0sF0Z9Xr50xbfXjwHPhgS0a82QOBjiAmeyDQEcRkDwQ6grHa7PWEYOnIgJZ39tVTpo0zgXpnXRRUzk3k3Tjs4inYbpaU0WZJacs+2maeNTLbStTMLMdK3t2DKRqPSdtWccQbu948aQTZno21CV7fYFeWsyGNDezIGkzkGn+W/H7wBJyUOWeo7D3ZRoGX3u2Y9nO2t+Wod33wdbouDXHnarpnl95m3Ws7FtJ+U+cvIwdjp/toQx5H52LciPrTwhpRvNkDgY4gJnsg0BGMVY3vXV7G7KcHPOqNSCrmlpuwCSIm2o7Vbq/2sQru3CeGeGKCeewcRxy7cbyazdFv3F8pacUnxZRINXLc655P71Jy3chlqxLWbDbQeNTXUyqL1UeOoGuUw+LKOgUue3JJ6ayNXBPi2F931WLY5BEeR8+Fx+7ShiuV3I8VRy+6KjglTkGOfpt1hU44um63dRkzKnLFeZIOfh7X96T+q+NnbCcXiKTDPd/V2UHUX4M0g/fJtgQCge8pxGQPBDqCmOyBQEcw3nDZukY9zIZqcLezC8bZZGyHsK3i7WFjxzj7N8eFXrLxvIsHGO1ua/Kut7Nlm7XBMnztE46DvJRFRn1Kxn5vbLuxMusRzC/vxzQjEwBDsAiyc8Xz6GdqAvg2y/vvMiZZrkm3vmGq2uZdkXwnGiG9haq/1f69m5/X96UQ5N2fe832zwQnc3m7v3ea5J2y11mfTdl42TWHwtrMlm92EfktETktIs/Td3tF5BkROTr8v6fURyAQ2H60UeN/G8BD7rvHADyrqvcBeHa4HQgE3sDYUo1X1U+LyD3u64cBPDj8/CSATwH46Nan00211qvgMknqoo9uMhlr9LnA9dZwz6yOVk1LpYYb5Y5aRr8Z7jrfv+ZVyZwsDRW5YV4QeF92Z/ooPI4mc6YMX7flyvdyFdxtbKZxFt1Fl/GVK4PtZLY8hG6/kppNJBUyn1xjesm6LE0UoRtf5dO5iEse494qZ865TEK+7xevuCZ69un5LpVEa7hBW+BGF+gOqepJABj+P3iD/QQCgTHhti/QicijAB4FgCnMbLF3IBC4XbjRyX5KRI6o6kkROQIgW1ZTVZ8A8AQA7Kr26WZElg/yJ/Wl2jPvpMyI6RIW6hWiJe6XKKL5xIVklJK6TCp3I4LOqON5DrrScUa9bZt048/NKu31UEnnTlcqh+W8K0YFL3hQ+NzMiwfAEIQYNd4/O1w+yY8Vr4Kzx8ePB5sQ3qwxpCX5xCaTxOIJMOjc6iMFCcYD4av3sunlab1bqPU3qsY/DeCR4edHADx1g/0EAoExoY3r7fcAfAbAW0XkuIh8GMDjAD4oIkcBfHC4HQgE3sBosxr/c5mmD9xiWQKBwG3EeEs2Mwour0YUUH90JJjPQJLL5NIouXG4/wInuzjjlW3UypQVcpFOroyRkaNEQGkytGi/gvuuyUtP18PX4m1
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 1 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"imshow(flowpm.cic_paint(tf.zeros([1]+nc), final_state[0]).numpy()[0].sum(axis=0))"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 7,
|
||
|
|
"id": "cf61ff87",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"mesh_shape = nc\n",
|
||
|
|
"kvec = fftk(mesh_shape)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 26,
|
||
|
|
"id": "bc83bc0c",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"# We need the vel and forces from lpt\n",
|
||
|
|
"def lpt1(delta_k, pos): \n",
|
||
|
|
" pot_k = delta_k * laplace_kernel(kvec)\n",
|
||
|
|
" return jnp.stack([cic_read(jnp.fft.irfftn(gradient_kernel(kvec, i)*pot_k), pos) for i in range(3)],axis=-1)\n",
|
||
|
|
" \n",
|
||
|
|
"def lpt(cosmo, linear, pos, a):\n",
|
||
|
|
" delta_k = jnp.fft.rfftn(linear)\n",
|
||
|
|
" a = jnp.atleast_1d(a)\n",
|
||
|
|
" lpt_dx = lpt1(delta_k, pos)\n",
|
||
|
|
" dx = growth_factor(cosmo, a) * lpt_dx\n",
|
||
|
|
" p = a**2 * growth_rate(cosmo, a) * jnp.sqrt(jc.background.Esqr(cosmo, a)) * dx\n",
|
||
|
|
" f = a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a)) * dGfa(cosmo, a) / growth_factor(cosmo, a) * dx\n",
|
||
|
|
" return lpt_dx, dx, p, f"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 27,
|
||
|
|
"id": "ab85a109",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"pos0 = jnp.stack(jnp.meshgrid(jnp.arange(100), jnp.arange(100), jnp.arange(100)), axis=-1).reshape([-1,3])"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 30,
|
||
|
|
"id": "02141439",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"lpt_dx, x, p,f = lpt(jc.Planck15(), initial_conditions[0].numpy(), pos0, 1.)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 31,
|
||
|
|
"id": "4225880e",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.colorbar.Colorbar at 0x7fce887a5c30>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 31,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAATcAAAD7CAYAAAAPf9NJAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAADK1UlEQVR4nOz9fayu23YXhv3GmPN53netvfc5914bGweDTEXaiiKVVlW/qCq3pCohKBZVQaTiK1CcP0pCGqpg+AckK5Kr8hGkVkimSWsaokAJFS5ySyiVFSGlbjBFgeA6hYSAwdi+1/ees/da632fZ84x+sf4mPNde+/zcc/aH+fePaVz1l7vet7ne445xm/8xm+QquLdeDfejXfjG23wmz6Bd+PdeDfejVcx3hm3d+PdeDe+Icc74/ZuvBvvxjfkeGfc3o134934hhzvjNu78W68G9+Q451xezfejXfjG3K8MuNGRL+GiH6SiP4OEX3fqzrOu/FuvBvvxosGvQqeGxEVAP8JgP8RgJ8C8B8C+OdU9W8/+MHejXfj3Xg3XjDqK9rvfxPA31HV/xQAiOjfAfA9AF5o3FY66pEegWzj8QdVKF5kfOm53z6LiX7uuPNg/zwPoJcH0zjD6XPb4div6r0/0b3vTvseW734XF90ni/az/1dPbdren7biwNN18CTgx+LId3fMQBi+3y+Z74fkH1H59MnGqchAKmO/cexVC/vfZwc09jvvL3o5TnpfI0vuNYXLO73P3np+/GiZ/+mh58n3XvvAOCp/vyXVfUXfJbd/4//B4/0Kz/fP3a7H/+Pzn9JVX/NZznWZx2vyrj9IgD/YPr9pwD8t+YNiOh7AXwvABxxjf/O8dfaJCKyn70DItAu40vq/6ZpsjGBiC63e8GgwrZfEdz3VimOWYp9IL4vZtBhteOpjMnT7eGq2r+1d/tcBSoKYgJKyf1q788fK8b98xG1c70/5nsz7wuw/auO+6XTvfBzz/Oa7x+Tnb//TUX9zwQQ5z2jw8G2JQJi/3GvpklPywKUAj2u41zZ70Eh6FIgq79yBGjxzwuhnDp466Ddnjt1BbYdtO15X6Bq+y0MXartvxZoTOjYvvVLoyUCbR2QDuzN7re84H259514xlTK8++G6uVz7/2FhvJBxr0F/7nPpvlApdicKGWcn4+/vP3b//lnPZUv/3zHj/2l7/zY7Zbv+Lvf+lmP9VnHqzJuL3I9Lp68qv4ggB8EgPf5WzSMS0xsFbGJNX+nxSSdX0x/6XjaUu69ZGFs1hXaGkgkDdQnvqBafVJoGgWiYi+ZCJQBCF+cWxi/+ZyUYd8tJQ3Axc0qdiwwAaI5Eel4uPD48jyIQFpt220zLzYuLQy/EiDt+Xsi6kYstrWJmpOlFDvPWkHVfqI1qIgZVrL7itkYh0E/LtC1pkGSpbiBY0gxz0uZoAxoIfSVUbYCPpuRKzfb2Fct47qJgOr3jsww5n1ZqjmL0wKgcY/2BnSB3t2Buj9/N+I54jNgLGhhXIELw4Z4h4h9Hz299AcdL4soXmDkiMnOucMWqvuL2oMMRdePdiTelvGqjNtPAfjF0+/fCeAfvWzjDDaIbAKVAlpW/6OYtxBGIiZ23ODeoRcrFwPl3sMsBbRUm6S+T40J4EY0v3v/fFQzpLLzY+iG9OCewyyJc+WM/avqMCax3/BsRMZkCe8sznVZQL1DWzfjxpT3Qlsf94ALQDKMN/cLz+PC2Pv52femiUEEFDOE5EbLjlmGYavuFYRhDcMTEyx+Fh7GpzKUCLoypDD6ge37ALTA/manD6mEcvLHXBmgxY45TtL2t1a7PgGod6BP1zd5cmkMmdzoCagt0CKA1HEPwuvSydgVBpRAVYcRjAVn2y6fefz9VU76fMdffow0ZHmufM8R+OxDAchbE4N/9HhVxu0/BPBPEtEvBfAPAfwmAP+zj/1WeFi1Xq5YrUH3HVC1yS5tPGwPCcIdt/3wWEWJhmErDFV/aWPf83FKsbCQdJyPHwPw0CQwJSag4/nwhodxCINF90NP9nCPKF8TVR34lhs2vTqAutg1r4tdS+tAZ5CqeWgqfp7FvDC4cY7r90VhfulfiNu5MSbV9KAjzMzwrxZQJ6BYSAn28BC4zLsLoMzQCD2ZIIUhC0ErAQ0gNa8qDRwjnXDAvg+GubphqEsxz+9QQU3ATYC9D6wuFo1aLjwbJXveSrboUCwqAVMAl5583CsCsPi9nBeF1p7z0Cys58ltfqAR1zUO9PXt5+v93guGfISBfZvGKzFuqtqI6HcD+EuwV/bfVNX/+GXbXwC23d37Wm31LAWoxbyYZYe2BtzeTZ6JhVOJqUWoF0YowqowOHOoCPhnkwfE8FXbt9vbWPlVgcMKWlegVsj2NPdhRmFcRxwvjBp56BpeW55f87+XAloXYFlBj66hhwVyvU43yYwV3+2GRbVuxmf25iBAYRDquKZ9mhizt3r/GcTiUMrw1q6O0KVCj0tel4WUDKmGl8nq164An7sZ410ABqgL5KpCKkMWM3bUFWBAiNBXSgCjM4GE0B5V8M7gTSx8bAJqMIxrYTOqALgJaO8Dk3OjpuG1+2RWdk9+J1ATYF2GYdqnUJ0mzBXI+0fLcnH/VXUsSmEw5vfpVYSmQIaWOnvh97BVmhdj4fz8uW0/w1Ao9m/ysBSq+iMAfuRTfSnwCweNM+wRyhDRQsN9YFGqFtPcB9rZDSSxrdT8gpUrV+P7uF0ZoRzde1EzJHN8kOjCK8tjx++gxA+fw/lEh8dEZIbtsEKPK3StBr47NgVVkCioq13z3oan4vsh5ilg8FD7/rUSmecxT5LZqE0Lii4WiupaLcvJZEatMqRaiCgL2Tn5fyBY2Kl6EUHNWVJlglQyT47s8Zknp2Y8hRw/onHOcV8Jl8ZjCoXVDW+GpYxMZqArUPTSY3dDSarmOcZJ3vdyZs8pcNC4b/fx3dc1ZmwUk2f+IsPzgMZIAfRv8rD0048Ix6qFf7rUnGSWmXIQuzOIH9sEjVVbBRqrcGToarXwNkJJ4HJSxOfEeUwDlKeQVsowimGYWjOPpjDo6srwsNaGQauXt9RCwJOd//R5enTuoVKt0MfXkOsD+pMjZDED0o8F7YpQ7xS8C0AErgwmyqwgTdeXeF54HLOXFiHkeRv3tI6w8yLMBKCHCi0FshbIymjHgn5kSAXakS0ZwMBypyhnBTcFu93kvYNUQM3CfO5muLQQ2hWjL8D+mEAClHMckFA2BcDgZj/zPnoYC7j3p2pG7GAQhi4jeyt1NloELQwOIyTi2F1kFR268PA0Pfzw1OYQtHfzkv2emgenCd6/sqH6Qo/tIvudrIH733348/pmx9y+viHd04k+2F9aVWBxQycCtG7ZLyJ/+ehexrGMVW3CVgYP6gVDdRi2PP4AxnPy1zKMxPEAah16Gl9JL3FOENQKIrlMVEQYGEb4sEIfHdGvVwvhintGK6EfzAjkZHdgnfy8zbvVybOdjPJ8iZEMOKzpCWKp0FpG6JlYnVpovDDao4q+MtrVMEz94IZGFNzIjVFc4HRPxTxNw98I7UhoB4JUd5TY1pD4jnmC9rssDD2wX7s/ki6gTSzhwGohJ3l43MVCzzh+17weUocVyIydrgtob1D0TCRZlnrC1mQsHBfe2vx8uwy8DXi1Rg64PNa98TqEZxVA/5wI3L4dxs3pFOgdWmu+hDGJw+Do4hO4K7A1S2KeYAYusqvSHTujS0Plnkry1eaRL+4LMKnIti3VDMFhufzutl9SS9IjpMwaUq1QNMPCgMuwttp+5fqI/mhFv6ppCKQS+mLYFDfkJFfCxLuCTXC/Lxpg+pTRzGtsHUpm0JIguy7QpaA/WhMToy4gUfSDeW3tmtEOZtz2x4S+ArI4GrARytm9sosVBm5UNOe7VjOK/eDJSrLJoiUyp5YxpcXAf1kt7M1HIQo9A3UT8+IKm0fDti2f9YWEW1JkRlo97NZDsec2L0JC95JR7r0HX4wYgEUQGjCDJ5ZmnuArwd2mMHTG32gyxLmdbfSwx5/G5wNxe2uMmwPaopYZzc99VV7YwGQyzEWJUE4NhRm0VMsgBnGzdwONL8B
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 2 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), pos0+x).sum(axis=0)); colorbar()"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 13,
|
||
|
|
"id": "167d193f",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"# Define the ODE\n",
|
||
|
|
"def f(state, a, cosmo, lpt_dx):\n",
|
||
|
|
" # Extracts positions and velocity at a given point in\n",
|
||
|
|
" # the simulation\n",
|
||
|
|
" res_pos, res_vel = state\n",
|
||
|
|
" a = jnp.atleast_1d(a)\n",
|
||
|
|
" pos = pos0 + growth_factor(cosmo, a) * lpt_dx + res_pos\n",
|
||
|
|
" \n",
|
||
|
|
" # Computes the potential given the current positions\n",
|
||
|
|
" delta_k = jnp.fft.rfftn(cic_paint(jnp.zeros(mesh_shape), pos))\n",
|
||
|
|
" pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec,r_split=0)\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",
|
||
|
|
" 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))) * res_vel\n",
|
||
|
|
" \n",
|
||
|
|
" # Computes the update of velocity (kick)\n",
|
||
|
|
" dvel = 1. / (a**2 * jnp.sqrt(jc.background.Esqr(cosmo, a))) * forces - dGfa(cosmo, a) * lpt_dx\n",
|
||
|
|
" \n",
|
||
|
|
" return dpos, dvel"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 14,
|
||
|
|
"id": "ce83d3b8",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"import diffrax\n",
|
||
|
|
"from diffrax import diffeqsolve, ODETerm, Dopri5, LeapfrogMidpoint, PIDController"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 15,
|
||
|
|
"id": "d4e7b055",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"cosmo = jc.Planck15()\n",
|
||
|
|
"\n",
|
||
|
|
"def f_diffrax(t, state, args):\n",
|
||
|
|
" return jnp.stack(f(state, t, cosmo, lpt_dx),axis=0)\n",
|
||
|
|
"\n",
|
||
|
|
"\n",
|
||
|
|
"term = ODETerm(f_diffrax)\n",
|
||
|
|
"solver = Dopri5()\n",
|
||
|
|
"stepsize_controller = PIDController(rtol=1e-5, atol=1e-5)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 31,
|
||
|
|
"id": "37285eb4",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stderr",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"/home/francois/.local/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:6655: UserWarning: Explicitly requested dtype <class 'jax._src.numpy.lax_numpy.int64'> requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.\n",
|
||
|
|
" lax._check_user_dtype_supported(dtype, \"astype\")\n"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"name": "stdout",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"(1000000, 8)\n",
|
||
|
|
"(1000000, 8)\n",
|
||
|
|
"(1000000, 8)\n",
|
||
|
|
"(1000000, 8)\n",
|
||
|
|
"(1000000, 8)\n",
|
||
|
|
"(1000000, 8)\n",
|
||
|
|
"(1000000, 8)\n",
|
||
|
|
"(1000000, 8)\n"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"name": "stderr",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"2022-02-14 00:10:24.488749: E external/org_tensorflow/tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc:2124] Execution of replica 0 failed: INTERNAL: Failed to load in-memory CUBIN: CUDA_ERROR_OUT_OF_MEMORY: out of memory\n"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"ename": "RuntimeError",
|
||
|
|
"evalue": "INTERNAL: Failed to load in-memory CUBIN: CUDA_ERROR_OUT_OF_MEMORY: out of memory",
|
||
|
|
"output_type": "error",
|
||
|
|
"traceback": [
|
||
|
|
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
|
||
|
|
"\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)",
|
||
|
|
"Input \u001b[0;32mIn [31]\u001b[0m, in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m 1\u001b[0m init_state \u001b[38;5;241m=\u001b[39m jnp\u001b[38;5;241m.\u001b[39mstack([jnp\u001b[38;5;241m.\u001b[39mzeros((\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m3\u001b[39m))\u001b[38;5;241m.\u001b[39mreshape((\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m,\u001b[38;5;241m3\u001b[39m)),\n\u001b[1;32m 2\u001b[0m jnp\u001b[38;5;241m.\u001b[39mzeros((\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m100\u001b[39m,\u001b[38;5;241m3\u001b[39m))\u001b[38;5;241m.\u001b[39mreshape((\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m,\u001b[38;5;241m3\u001b[39m))], axis\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m)\n\u001b[0;32m----> 4\u001b[0m solution \u001b[38;5;241m=\u001b[39m \u001b[43mdiffeqsolve\u001b[49m\u001b[43m(\u001b[49m\u001b[43mterm\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msolver\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt0\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt1\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1.\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdt0\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.1\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 5\u001b[0m \u001b[43m \u001b[49m\u001b[43my0\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minit_state\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstepsize_controller\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstepsize_controller\u001b[49m\u001b[43m)\u001b[49m\n",
|
||
|
|
"File \u001b[0;32m~/.local/lib/python3.10/site-packages/equinox/jit.py:90\u001b[0m, in \u001b[0;36mfilter_jit.<locals>.fun_wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 88\u001b[0m static_leaves, static_treedef \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mtree_flatten((static_args, static_kwargs))\n\u001b[1;32m 89\u001b[0m static_leaves \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtuple\u001b[39m(static_leaves)\n\u001b[0;32m---> 90\u001b[0m dynamic_out, static_out \u001b[38;5;241m=\u001b[39m \u001b[43m_filter_jit_cache\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mjitkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 91\u001b[0m \u001b[43m \u001b[49m\u001b[43mstatic_leaves\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 92\u001b[0m \u001b[43m \u001b[49m\u001b[43mstatic_treedef\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 93\u001b[0m \u001b[43m \u001b[49m\u001b[43mdynamic_args\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 94\u001b[0m \u001b[43m \u001b[49m\u001b[43mdynamic_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 95\u001b[0m \u001b[43m \u001b[49m\u001b[43mfilter_spec_return\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 96\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 97\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m combine(dynamic_out, static_out\u001b[38;5;241m.\u001b[39mvalue)\n",
|
||
|
|
" \u001b[0;31m[... skipping hidden 6 frame]\u001b[0m\n",
|
||
|
|
"File \u001b[0;32m~/.local/lib/python3.10/site-packages/jax/_src/dispatch.py:444\u001b[0m, in \u001b[0;36m_execute_compiled\u001b[0;34m(name, compiled, output_buffer_counts, result_handlers, kept_var_idx, *args)\u001b[0m\n\u001b[1;32m 441\u001b[0m device, \u001b[38;5;241m=\u001b[39m compiled\u001b[38;5;241m.\u001b[39mlocal_devices()\n\u001b[1;32m 442\u001b[0m input_bufs \u001b[38;5;241m=\u001b[39m util\u001b[38;5;241m.\u001b[39mflatten(\n\u001b[1;32m 443\u001b[0m device_put(x, device) \u001b[38;5;28;01mfor\u001b[39;00m i, x \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(args) \u001b[38;5;28;01mif\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m kept_var_idx)\n\u001b[0;32m--> 444\u001b[0m out_bufs \u001b[38;5;241m=\u001b[39m \u001b[43mcompiled\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecute\u001b[49m\u001b[43m(\u001b[49m\u001b[43minput_bufs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 445\u001b[0m check_special(name, out_bufs)\n\u001b[1;32m 446\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m output_buffer_counts \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
|
||
|
|
"\u001b[0;31mRuntimeError\u001b[0m: INTERNAL: Failed to load in-memory CUBIN: CUDA_ERROR_OUT_OF_MEMORY: out of memory"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"init_state = jnp.stack([jnp.zeros((100,100,100,3)).reshape((-1,3)),\n",
|
||
|
|
" jnp.zeros((100,100,100,3)).reshape((-1,3))], axis=0)\n",
|
||
|
|
"\n",
|
||
|
|
"solution = diffeqsolve(term, solver, t0=0.1, t1=1., dt0=0.1,\n",
|
||
|
|
" y0=init_state, stepsize_controller=stepsize_controller)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 26,
|
||
|
|
"id": "9fc54299",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stdout",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"421 ms ± 2.28 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"%timeit solution = diffeqsolve(term, solver, t0=0.1, t1=1., dt0=0.1, y0=init_state+0, stepsize_controller=stepsize_controller)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 22,
|
||
|
|
"id": "325ee24c",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"Solution(t0=DeviceArray(0.1, dtype=float32), t1=DeviceArray(1., dtype=float32), ts=DeviceArray([1.], dtype=float32, weak_type=True), ys=DeviceArray([[[[-0.05698265, -0.4233138 , -0.23527317],\n",
|
||
|
|
" [-0.4048839 , -0.18278353, 0.10613398],\n",
|
||
|
|
" [-0.39811632, 0.36003193, 0.5207157 ],\n",
|
||
|
|
" ...,\n",
|
||
|
|
" [-0.12596521, 0.04650298, 0.12780993],\n",
|
||
|
|
" [ 0.10532238, 0.22021866, -0.00881826],\n",
|
||
|
|
" [ 0.21299808, 0.42117628, -0.05347536]],\n",
|
||
|
|
"\n",
|
||
|
|
" [[-0.07113229, -0.40208206, -0.18125573],\n",
|
||
|
|
" [-0.3594675 , -0.14042078, 0.1074788 ],\n",
|
||
|
|
" [-0.29909593, 0.34923548, 0.42418486],\n",
|
||
|
|
" ...,\n",
|
||
|
|
" [-0.2394397 , 0.07868364, 0.09896102],\n",
|
||
|
|
" [ 0.06142797, 0.22736594, 0.05615562],\n",
|
||
|
|
" [ 0.18805356, 0.38284227, -0.00324537]]]], dtype=float32, weak_type=True), interpolation=None, stats={'compiled_num_steps': None, 'max_steps': 4096, 'num_accepted_steps': DeviceArray(22, dtype=int32, weak_type=True), 'num_rejected_steps': DeviceArray(6, dtype=int32, weak_type=True), 'num_steps': DeviceArray(28, dtype=int32, weak_type=True)}, result=DeviceArray(0, dtype=int32, weak_type=True), solver_state=None, controller_state=None, made_jump=None)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 22,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"solution"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 23,
|
||
|
|
"id": "f131301b",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"(1, 2, 1000000, 3)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 23,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"solution.ys.shape"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 25,
|
||
|
|
"id": "a5f4d029",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.colorbar.Colorbar at 0x7f22ef156230>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 25,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAATcAAAD7CAYAAAAPf9NJAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAACtN0lEQVR4nO39fZQk13UfCP5eZGRWdnZ2VVdWEegCGyUSVQX0EMCqCbNRosnDhYSRphvUmpxzDI9EcsTxyCsdr2jKs/KMKO0ftmcO92iOvdoRpTk65tqypEPCEqmxJY4FwKS55HCopQpNQRiDoBrsKhAuNFENoCrro7OzsiIj4+0f8W7kjZvvRURWZ38AnfecOpUZGfG+4r37fvfzKa01xjSmMY3prUbezW7AmMY0pjFdDxoztzGNaUxvSRoztzGNaUxvSRoztzGNaUxvSRoztzGNaUxvSRoztzGNaUxvSbom5qaUOquUelEptaqU+tSoGjWmMY1pTNdK6rB+bkqpEoDvAfhxAJcAnAfw01rr746ueWMa05jGdDjyr+HZhwGsaq1fAgCl1B8A+BAAJ3M7opSeVB5KA41QAIAQGgfmWmiYrq9U6t5QMOMwo4G8c3RfDxolKOd9vlLwAUxAYYLV3dMaHWj0TFmh1knbGsrDFBRCAAfQ8AFUlAcNgFpb0jr5/cB8DqFxVHk4Zu6NoKAA9HQPHQBXdZSMB7VzAgqhaUcJQBUKJaUwBQWtFJQpmz8TqbjcQEfoANjT0cA9MOXBlDmpvOR6CKCpIxygP/Y0Pj6Ao1A4gMaWjtAD0DHvqMrGj+rj48brpjElqiqFo2ZMQ1MvnycA0DP3dkyfffTni5wnRYjaWzJ1HUAPzC9eboh4PtnIPr4q+e4rhQlTz4xXwoHqz8qK7mE9ChHqwfqrZszleNG83dMR9rVWuAY6e/as3tzcLHTvX/zFX/w7rfXZa6nvetG1MLe3A3iFfb8EYFnepJT6OQA/BwDHoPBflGtoqLQ0PGO+b+kIq1H8Ops6AoCBe+k60Zb47iqXP++qn+prKA9Lno9Fz089dzEK0dTRQNuWSxWc86vJbw3lYb5UAQB0zT21KMS26R+Vs6UjLJcq+ECljq7y0PZ8lHWE3aCFi1GIlV6QjAfVt+T52GL1LHk+GsrDWb+KrvJQ1hG2WX+nWbnrvQAXoxBPhZ3UOFI/6P+S5+OcX031/amwg4usLTQ+M8rDouen7qFxWDRt42Mmx47+y3Glsvn7k9/p/tUoxJaOMGPeHf9tGKJ331AeZpSX9Lcp5g/RFmuzjeRc4+MLIOnjf3mkgfVKPblvPmjhE63LyRzhtCTmpGz/U2GncH9dtLn5Bs6fXyl0r+eVZ6+5wutE18LcbLvDwDamtf4sgM8CwIxX0vTCOUOxkZz8RPyF55UhFwEvg5fLy6N7t3SERfO7ZH7n/GpqoTV1hM9328n3GeVh2cE8OINqmHs2wg4W2KSle1ajEE3Rx4tRmGJqy6VKUn5ZR1gzzJMv9An2PF3njMT2PlZ6QWrsiGnxcSOGREyA981G/D3QM7JeYoic5IKWjJKXI5mPpLw5wylvYy06h+X9kqFvdttY0BF2SxVslGvoGuZOG5vcnIls/TsMYrWR1lky0ZuDroW5XQJwN/t+EsCrwxbCF9qwlIXagOF3b2JMfPHQROQTecnzseD52DYMaqUXDOzixIRszJxP+qS+KMScWcS1AguQo4xpC8Ox1WUrg3+2MQ3aTFxlyuv0J8twLUr5DhcxuJnYSJYpKW9uUNlZZcw4+inv4XXZGB7NK0lU5kovwDnlocyY+Izy0LS0L6tfW0LdcHjS0LqXf9stTtfC3M4DWFJKvRPADwD8FICPZD0Qao2mEc2aDFkAbDfLmVAuxOWivDJsz0pRasmIXYR65ksVbBqxcwrAR8IOnjDITZbBEUpD9JkWBjHARi/AtPISsVLeT/2ZYcgNALZNGRzV8u9LAn1xpt1QHh7zqzhTqmDNiMK83bbxI1F0uVTBqnnGRlusXXxcOFPh77ChPPxG9TjWK3XUohBTvQAtJkZTeSSG8jHh5VN5vHz5bmzkmls03jR+QBpN2sRTvlFkMWsa60YvwFmjQqhFYaISWTUovAizHhVprW9v5Ka1DpVSnwDw7xDrYH9Ha/1CkWe5TkYyFw7Z5c7KF2RT5zM2eS8vowhdjMLUJKaF2fb85I/0XEtmgctJSH0lpk1MifeTxEjJADnio3KliMvbaltgfJFJHQ6Vf8aIQ/NhJyWOZm0MS56PiVIFi0BKLKVnbcQZW0PZmc9GuYa1iUlM9QJM9QLMddtodtuJPpbrPl11ZJENWUqy6QX5uyDmTrQIOEXIIpIJMey1KMRsL0CZJAhHW7k0cH2YXoQounbd3c2ma0Fu0Fo/CeDJoveHGNRvrfQCK0KxTVKbMcD2DN/1XHqSvB2d65oaYsEvddtY6AXoKg9d5WGDKdtp0nH0x9EU/7zK6t7Sg8pjW7upDo4IOXLgk14uLpdhZVtHKJs/2X/bdypzIQqxkcFQXcSZINdDNZSH2bCDhYO9BLmREcTGnF3lcspiALb55Jp3nJZLFdxrDABTvWDAUMSfkSKraw7T+5oz99McfsrRjsOqcorS7S6WDk096OTl8sVsE9k4YgKyX+aWWCRcZFtlui/6PU9Upfq4Di4lsnTbKWZDE5PXu+D5aJgyGkLMc9XPxfUZNaj050RIYVWUkaAKYYywKeJd7bDpHG3Execi5fF75e90/avdNpZ6Ab5vxuxiFFpRURZJq6lE9652ZqEhLimcKVWwUj2Oso4w121jqhdgK2gNGGFcbXONx5aOUItCdJWHupk/WcjNVpfLNWU4us3F0mshl8jJF3MRixsnlw6m6PN597iQzABCs1g4gb7uJqutJKbfbyxmjW57wKpI9XGmBgDn/GpiXOgqD4uwK7JtYmlTR5jvBVgzjISLwFJ1QEyA+mqzSHMEPaM8rBj0JUkyNq6nlOKny/J9GJL1cpLvxlbfWhRi1rhcTPUCHBiXnSxRmYuyEoFxPWGbjEpC52nbkIpuQsPSba9zu1aSE4G7L9gWSpa4I6/LBTzMgpCLyWbZk4udnlmNQsC4QtBiljo2AGgwZkjPz6jYrWOiVMHzRxpoez4e3G9iWRgqppU3gAQbysNUpY5mqYIrRryc6gVY0hEWjG6wzPpDDIzadjFK+95RfxY9Hx8zfolkkV1j95FKQRKVO2d87xrdNmCYgRxDyfClzs/lS+Z6ry4dLL0DG0Ln48ivScMFlf1U2MHPBy0AwHNGJcEZcUPMD2JsHyvXsKWjxECSzAFT96LnY9OvohaFCXOztZW39/owuDFzGxnJiZVlzeLiIoCEQRDxxWBjiMPoXySRyGLzi3OVR21w6YMkke6LmBG3nhIqmwYGDB1tz8duqZIsjN1SBbVSBWuen1gfFw72cH8U4qxYOMSw5OJc8nzUJibxWqmCS6buhYM9wCj4pZHD1hdyYuZM3EVy8V4PZTm9i8MwAj7v+GaRJXbz6xKh8u/UtqkoTByx8zbl64XcgLGf28jIZjnLs2plTfxVsXh5WXIi8glbBOHZdDJSDKM2NHWEZwQD4Lodea1pRL25bjuFtgg1tRkKO22u7ZYqaBsGtleqYLIXoBaFuGysjmsTk/jasbswG3bwwZ3/iKmoi5pZQCeMvmiu247RFRuDGeXh8Uod/+PsKayXj2LTr6LrlfD3Xv8OzuyuO8eLI6RWFMZiuoUJcnQIAFsZVlqXfilPnyeZJd9w6F1xaYGLzjbUKJkiifHSWmwbE65X5pZQIoqIOBt2EsYmN+4s3SZdGw1LihBFB/m33eJ0SzA3F2UpvqVC/7Dkcro8LLkYstTR0WfbPVN60HpJREiILLWbfhVtz8eW+e+b5zZLE1iv1HGhehyzR2YAaPy2sfDtemWUofH49kt4cL+J+zo7WFAeTosQq67ysFuqYNOvYtefQBcKYcYCc/WHW26JbPrAw+jT8sbbdl0iJlvbbcQZVNvzsTBEO7nPX9b8WnP4FxYxio2Kxjq3QxAFYNsWwcAu7GBsHBkkYqrDiRQYdLHgylv+Xe7
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 2 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"resx = solution.ys[0,0]+pos0+x\n",
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), resx).sum(axis=0),cmap='gist_stern'); colorbar()"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 57,
|
||
|
|
"id": "654d9c69",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.colorbar.Colorbar at 0x7ff54ab2b1c0>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 57,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAATcAAAD7CAYAAAAPf9NJAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAzrklEQVR4nO2df6xlV3Xfv9/75s1Pe/xrbDN4XHASJw1ExSQjQ0SbGojBJVGc/oFqqiSkQpoqdSqi0sa2KpWmkiVXbWmatEGaAsUoBGTlR7FoAnHcWAgVDOPEgH/E2MaOPczg8fjXjGfGM++9u/rHPe+dtde5e7197zvv3vvOWx/p6t17fuyzz4/Zc9baa30XRQRBEARdozftDgRBEKwHMbgFQdBJYnALgqCTxOAWBEEnicEtCIJOEoNbEASdZE2DG8kbSD5G8gmSt7bVqSAIgrXCcePcSM4B+C6A6wEcBvBNAB8QkUfa614QBMF4bFnDvtcCeEJEvgcAJD8P4EYA2cFtK7fJduxawyFnCKrvsxoHzewPkOnvsVD/MRb/J9k4rFrQRkC5aZ8sNE70sc21EemrH96xnYdivZ+R6tivySmck9fWdHPf+85d8sKLS0XbPvDts18WkRvWcrz1Yi2D2xUAnlW/DwN4m7fDduzC2+beM/jRNxePhQ+5t13pP9jSB95toj6W9O2D3EcRLfTDb15dj7m5dN0WdetLBxXTBhYWVr72z54t65NtQ10DWSr7B+W23zOD+LZt9Q9nAIM+di+9L3LuXP3d3mt9LH1u5hkoPrfSZ988O8vH/vrCl8qO43D8xSXc/+V9RdvO731yz5oPuE6sZXAbNpI07gzJAwAOAMB27FzD4YIgmAyCpdL/oGeYtQxuhwFcqX7vA3DEbiQiBwEcBIDdvLge/Oz/nKVvD+Nul7zxqRtn3568dW3jPUDq2PZtpLx5fbmd69Y3/ZifX/naO0+5EbbOJ5v1n3+h/uFdq9LrPeYbr/uGqn6Xms5i3kKTtzXbx9x5N5YXvrm14S5YIwKgP7O+lnLWMrh9E8DVJK8C8H0ANwH4p630KgiCqdLHJn5zE5FFkr8O4MsA5gB8SkQebq1nQRBMBYFgYZObpRCRPwXwpy31JQiCGUAALG1ys3Q8lv9HGGXaPzdD6vkn3HXKH+L8D2V9XYkPS882Li6m2y06fSztv+qXjDuJ6J2nDuNYsuvOqe9qEsjMFCbXAKmfKp0VNTOkGexMqjczme6o/JN2NlZv5lxvfaRxZ785X1+PNmZ+3WNZP+xKH9sZlDa7zy0Igg4iAJY6IGIbg1sQBA02vsdtmoOba5IVBjKOS2EIhosKl2j0V5up3rkUR/WPec6Oaatb5JwJWlVmav/UadUN0w8VGtIzYSLJ9VHBvjYcQxaUqW8DcHUYhzLzGtvpYzsBuNyxY2jfAUBeeBFFNEKM1Cptipaa1GPCrVuT3/o814pAwucWBEH3EAEWNv7YFoNbEAQWYmloAtLGIga3IAgSBOtuVU+E2RzcxvUxFStAlKVYNabzdZK35+MoTe5vI+WsZ0Ifxgm+tOET2meoU7PMdv0rX7fynQvmWn3vmbpLOtTE+suc1LLEF6i3s//ytJ/N3DPt45PTyn94piwMp9kpZz8vTcujMDQpCXnRggAAZGH5nrXzxhVvbkEQdI5BEG8MbkEQdAwBsCAbvwLBxh7cxlTt8KLYvchyZkwjHZkOIDHfxGQvpA0WKqN42zmZB6UZHI2wgow2Gy+5KPl9+g21Ysi5Xem9uOTUmZXv/aPPZbshzj3kdmV6navDSWzCgw5REXtv1f2sTTc0TfGkQUe5pLGpusZa42+U6cZEGaXs/jXcIi1mKAiIpQ6UV9nYg1sQBOtCX8IsDYKgY4TPba2MkoWQe1V3hAM90zNpepQEZw6fvUvMncY+IzwkpdvaGdKkjcxicz16u8+vf2xLzVKeqU1K7tie7d95j9RilWIj/tVMJy/YXS8/fSbZrrdVmce7jFKz7vOJk3UbxtTvn3mtbsNmSuRiGqwMuBf74MxmJmIKZY+cjzvTrlZpM133o5UQDmIpfG5BEHSNgRJvDG5BEHQMEeKctPEaOl1icAuCoEE/fG4tMo7wpBdGYJUutP/CC/fQYRzGD5P45xxXXRIeMEJYQRbre9ESH8510+fSuB6Lqo2FU/k2lG9OTMiLHKlDPHrGXyYXKJ/ehbXPDSdfTbfTPxbNRVW/tZ/NDdc5z9TFVccrL683puBPYdk/tx/FWSs2BKg9kaLBhEKYpUEQdI5uTChs/DMIgqBVlicUSj6rQfJKkn9J8lGSD5P8cLX8YpL3kHy8+nuR2uc2kk+QfIzke9XynyL5nWrd79DTjcdUaihUr90tJMc3BBL1YUxdgFKTxK1RWUgiDNlG7dMWBDrt9aA+N5tFMZ+5ruYaJqauDhkBINvUo6UT28256JAGOXEiWdfLhKGktRvSDBHu3JGs031m4lYw55K4EkxYi9523GeiBZkN3cc22vNYai+IdxHAR0Tkr0ieD+ABkvcA+FUA94rIHSRvBXArgFtIvgmDMqFvBvB6AH9B8kdFZAnAxzEo8P51DApT3QDgz3IHjje3IAgSBMSCbCn6rNqWyFER+avq+0kAjwK4AsCNAO6sNrsTwC9W328E8HkROSsiTwF4AsC1JPcC2C0iX5OBA/0zap+hhM8tCIKE9ZpQIPlGAG8FcD+Ay0XkKDAYAEleVm12BQZvZsscrpYtVN/t8iwxuAVBkCDgKGbpHpKH1O+DInLQbkTyPAB/BOA3ROSE4y4btkKc5VkmP7gtpw6NUKQl8YfolBxb51L7b9Z72r+UNtq34QFO+lVW/NGmDamUJdu+9mnJSRUm0jfXVF9/3R4Avqz8Z+pe9I1/z4aoJH18TamTOH47HYIhJtREq4kk/6CMEoo+F/sPr69ELqU/5hvNOM+Buc/JfTHpV634dhUjZCgcF5H93gYk5zEY2D4rIn9cLX6O5N7qrW0vgGPV8sMArlS77wNwpFq+b8jyLKueAclPkTxG8iG1LDvTEQTBxkYEWJJe0Wc1qhnNTwJ4VEQ+plbdDeCD1fcPAviCWn4TyW0krwJwNYBvVCbsSZJvr9r8FbXPUEqG509jMCuhuRWDmY6rAdxb/Q6CoAMMJhTmij4FvAPALwN4F8kHq8/7ANwB4HqSjwO4vvoNEXkYwF0AHgHwJQA3VzOlAPBrAD6BwSTDk3BmSoECs1REvlI5AjU3Ariu+n4ngPsA3LJaW6Aym8x0u6uRr9EmjlXjGKceaRsZBOtNQyUl30dtNnmhA3REEXU2B7VwpSME2X8tFbhMlDqcjA1RbhMvq6SnTTKrCqL6aO+ydk8k4TDWPNZmqgkxSkJPbNhMTn3GekVK6yQoU7S3Pa2TkITDzGXa6LdVQ6EdM1dEvop8YYd3Z/a5HcDtQ5YfAvATpcce1+eWm+kIgmCDI2CIVZZA8gAGgXfYjp2rbB0EwSywmXNLczMdDapp4YMAsLt3iSy/xtsZqcRkKBUSHMcMXQUvcb6YUtPWK/vnUXhuxf3vOe3p6P8dafS/vKqS0s+Y2btMP9jLC4zaLApk6hNYQUo5o01PZ5Zct2c2c+srqIyNhqmoZ2qTkoiO6IJFZ93o9u3zoLMt7CxzdQ+52EI2C4D+Js4tzc10BEGw4RlUnC/5zDKrvrmR/BwGkwd7SB4G8FEMZjbuIvkhAM8AeP96djIIgskxKO23CcQqReQDmVVDZzqCINjYiLATZunkMxQqX4TYzIlCP1Ua3uBkMhhK/XbjRKA3itH0dCS5qS+Z7Djma73usxfKMkbBHACpb2dHLTTZv2R3shlP1ZH77WRimDbUvdAR+b0L0n70ttV+qkaGAjMhRp7fsm8UVLQA5vmpGObi5RfU2y2o62Yety3P1m7ppZdeTttXYShUxXrkrHl2cmEtAFD9llKxy1Xogp5b5JYGQZAw0HObbX9aCTG4BUFg6IYS74QHN1kxPcYNs0jDCsz/Ll5NBRWC4B671LxyIs65RU3n25CAxXzIRDGFoSD6+nDeudUmTCGJ6l9Q5uCJ08l2fcfkzprBYwuAqj6
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 2 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"imshow((flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]).numpy()[0]).sum(axis=0)); colorbar()"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 58,
|
||
|
|
"id": "80da1b98",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.colorbar.Colorbar at 0x7ff546f8ceb0>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 58,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAASwAAAD7CAYAAADQMalWAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAke0lEQVR4nO2df6xlV3XfP+u9mWF+ePxjYmxP7VEHEkNjOWDIxIG6pcaG1CEI548YQQWagqNREaGmEIEdpKKoioTaioLUCHUUHBzhBix+xBai2O5QN40ExmNswMY2GOLagwcPGKeYsZmZd9/qH/e8efuc886++5577vnx5vuRjt49956z97rnnLfuXmuvvZa5O0IIMQQWuhZACCFSkcISQgwGKSwhxGCQwhJCDAYpLCHEYJDCEkIMhpkUlpldZWaPmNmjZnZ9U0IJIcRaWN04LDNbBL4HvB44BNwDvNXdv9uceEIIscqGGc69FHjU3X8IYGafAa4GKhXWhi3bfOMZO2boUggR48T/+xlLzx+1Wdr4l6/d5k//bJR07L3fPna7u181S3/TMIvCOh94Itg/BPx27ISNZ+zgV9/+vvGOAuzrk/o4hte4cI4tR06b6XGfE6FMdZ+d1DZi33/e5834f/GDT390tgaAn/5sxN23X5B07MadPzh75g6nYBaFtdbtKV1uM9sH7APYuP2sGboTQrSDM/LIL1qHzKKwDgG7gv0LgCeLB7n7fmA/wJbzds13XNXgL9XE9kO6HC0W+25idBS0kRttFUdpgdVgBTlqjdKK51Rd19TjisdGjvNg+qn4XRqh7uirAxxY7ptQGbMorHuAC83sRcCPgLcA/6oRqYQQnbLMOhthufuSmf0RcDuwCNzo7g82JpkQohMc58Q6NAlx9y8DX25IFiFED3BgtA5NwnrM8zo00XbM5xHsz93nUZcqWSI+ppL8wX7OhXWicNhi8LoYghy2UXHdJlJ1L4r+oMg9C2dDc30Xj/Pqz1JJ9fcNgfXowxJCrEMcGPU0sacUlhCiRD89WF0orJWh8rwV+DRT33WOiwz/5/7dasgYCxSNhUbkziuYc6GZE5qHUDADY2Z2nen+iDlXDKfImb6BvMuFJ38h/C4R87ZI2LctRdpIvQZNBMXOiOON+bDM7EbgjcARd784e28H8FlgN/AY8GZ3fyalPWVrEELkcIcTiVsCnwKKS3euBw64+4XAgWw/CSksIUQBY5S4TcLd/xb4WeHtq4Gbstc3Ab+fKpl8WEKIHA4sz9etca67HwZw98Nmdk7qiesrrKHFfhZCf0Xbi4VT/T4R307O71PwP1W2UThu+QXBZ4Wx+sLxoK9UD27iPau7DCiUcaEQZhD1s8XaDH1k0QNrftYRKaOnjLPN7GCwvz9bjjcXNMISQuQYB44mK6yfuvueKbt4ysx2ZqOrncCR1BPlwxJC5HDghC8kbTW5Ddibvd4L3Jp6okZYU1AVKV00eVo1EVNDEmLnRUyS4vR/yGjT6utNzxbEiIQQVMoRo+FMHNH/tSna79Uqh4ZwjFFDYxkz+2vgcsam4yHgw8BHgFvM7FrgceCa1PaksIQQJZYb+tV197dWfHRlnfaksIQQOab0YbWKFFaMyMxabAauqf7W6mviOVVyRSLAlzfmP1s8tvo6FwFeaGPDc0HzsVm3UIy6JlTqDGLRDK5KRrjeTLlGv48xqu+fmitSWEKIHOOMo1JYQogB4G4cLwXn9QMpLCFEiWX5sOZAzH9TzCyQuPI/+bN5EEtSN2vTxSj1YH+hkJgvd1zwhJSi2U9Uf5bL5BC2ERcz30boj4pcj/C4Dc/ne1jasnYoesyXNo1Pspcl0WZk7HSXSSiEGARyugshBoKc7i2RS6aWOM3eRF9zD2uYov0qucLkcgAbQjMwllBuVPE+hbz2hes92rJ2+6Wwg5hpFpit0ZzogVzHT88LmVvwHDF9U0ldeN1IncaoIIX9ht0Xo57auutKYQkhZscxTng/VUM/pRJCdIac7kKIweCYTMJ50GWWhLn0VeWHaKCv6JKYWA290E91nEpKcYbBvYkt78n1W/T7VLRfKjQRiXGskzyw5B6qWt5Du9ka2qyF2Ven+0SpzOxGMztiZg8E7+0wszvN7PvZ37PmK6YQoi3cYeQLSVvbpPT4KRqseiGE6Ddjp/ti0tY2E01Cd/9bM9tdePtqxkm5YFz14i7gg00KlkJPzezmmaKOXZ385sVzFqqS+01Try8sEZ8qU2Ku+uJhYf74YrLAKpOw1Fcs8WHkZ71S5gZMtpLZl5ijvwmadLqb2b8D/pDxN/gO8A53/2WdtupKlat6ASRXvRBC9BvHWPa0bRJmdj7wb4E9WSHVReAtdWWbu9PdzPYB+wA2bperS4gh0HBYwwZgi5mdALYCT87SUB2Sq15kJX/2A2w5b1f1wLXN5GoN9DWXSPdcB8HriGlQN+K5anFysc2wnBmFSPFi4r865K5jxOSsSp5Y3C9Gs1fNrMX6KuWgD2cJIyXoFyKzqFXnxCh9zdR8/TMyrkvYjMJy9x+Z2X9mnLv9eeAOd7+jbnt1papd9UII0Xemqvx8tpkdDLZ9uZbGEQRXAy8C/hGwzczeVleyiSOspqteCCH6zbjMV/IM4KS6hK8D/t7dfwJgZl8A/inw6TqypcwSNlr1QgjRb9ytMZOQ8YDmVWa2lbFJeCVwMH5KNe1Huq/Y4TG/zLxpoq5dxJ+Q+3FquehClT+tNA0evF4oZHIIn9WlzRUnrXFeHapqPUIk80Qhc0PoSyuWoK/sK0IxRGAxmIDf8Fz+w1CWY2euCllKmBjUcCz6usLPFp+vliMW7d94WENzPqy7zexzwDeBJeA+Mp92HQa9NEcI0TzjfFjNzSS5+4cZu5JmRgpLCFFAGUdXWRm6zjkBWasUvks4LR7W+IN0sySVVNOgtGB47VTnFW9k5xRkz7VRM0d6bHq+0ryN+IOj1yMxlKOY7HDz06tCbno2L/BvfPhbJ1//n5t/8+TrYnhFaD7HvnPOvC2a6pHr2GRozTisoZ/LSDTCEkLkWFlL2EeksIQQJfqaXkYKSwiRY5xeRiZhOgPzb5Xq9UWm++ss6SlNb8eyB+RODF7W8BUV+7a6RRwSl0LVvR6hz6loyYSHLob3JRIa8dzOfAejzasX/JcvzJ93+AcvPfl69z3PnXz905dtyR0X1keM3YuFSNGPKA3rF/mwhBCDYJytQSahEGIAjJfmSGHliEZoRyKec0nMihMZfTEdY1kBQhJN39LovEZJ+1h2guh54TUu3LPwvixHvmes5HzuqySai7Go/WIGiTBKffGXqx388oXVF27js/nPwkj0Igvf23by9XM7V7/oaFO1kLGMD9H7Eltd0agFpxGWEGJANBnp3iRSWEKIHJolXINprkdlRPUUpavyDab3XYvUoXsDyQNL5nM4Yxbc3eIIP2amVSXEK0aAL22tkClGzNyPEC4yLi1OPrb6xolt+S+zIThvFJupC9p8wbPVcmw4mr+Q4TU+et7imu8X2489m6l57eetT2QSCiEGwUpO9z4ihSWEyOHAkkZYQoihIJOwSM1o9tRy3VH3VmIxghjRiPWa7VeGb8SenZg/Z2O1zybWZi5cIXhCrOgHi8gY+pw2/mL1wOcL4QQbwoR1Eb9S7u1COEvotyquMjhxWtBXEOIQKxhRDI2IrQNeDj5bCoLbS8UwIkkdGylB36RfNrGEVypmdibwF8DFjCV9p7t/rU5bGmEJIXI0ncAP+DjwFXf/AzPbxLjUVy2ksIQQJZoaYZnZ6cBrgH8N4O7HgdSCaCWGobAqQgFSF+1OOjalryK5Uu+Rxc/FCPDFwFSIRjynJscrtBGaXOFC2mIiwZCi+RKac8svCDsr9B3IGJp2AFt+smrfLR5fPfHo+QX7KixY3kDU/mKhAHouSj0xDKaUqDDYL17HXKhHxf0rUlqgHewvRK6Hx9pv0CRsOIHfi4GfAH9pZi8H7gWuc/ejdRrrp2dNCNEZjrG0vJC0MaEuIeNB0SuBT7j7K4CjwPV1ZRvGCEsI0SpT+LAm1SU8BBxy97uz/c8xg8LSCEsIkcfHJmH
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 2 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"imshow((cic_paint(jnp.zeros(mesh_shape), resx) - \n",
|
||
|
|
" flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]).numpy()[0]).sum(axis=0)); colorbar()"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": null,
|
||
|
|
"id": "32357464",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": []
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": null,
|
||
|
|
"id": "547af8a9",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": []
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": null,
|
||
|
|
"id": "8f5cac4d",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": []
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 35,
|
||
|
|
"id": "aa14a004",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"from jax.experimental.ode import odeint"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 39,
|
||
|
|
"id": "973ca031",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"init_state = [jnp.zeros((64,64,64,3)).reshape((-1,3)),\n",
|
||
|
|
" jnp.zeros((64,64,64,3)).reshape((-1,3))]\n",
|
||
|
|
"\n",
|
||
|
|
"@jax.jit\n",
|
||
|
|
"def solve_ode(init_state):\n",
|
||
|
|
" cosmo = jc.Planck15()\n",
|
||
|
|
" return odeint(f, init_state, \n",
|
||
|
|
" jnp.linspace(0.1,1.0,10), \n",
|
||
|
|
" cosmo,\n",
|
||
|
|
" lpt_dx, rtol=1e-5, atol=1e-5)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 37,
|
||
|
|
"id": "e801c3eb",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"res = solve_ode(init_state)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 38,
|
||
|
|
"id": "de770afb",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"name": "stdout",
|
||
|
|
"output_type": "stream",
|
||
|
|
"text": [
|
||
|
|
"765 ms ± 7.45 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
|
||
|
|
]
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"%timeit res = solve_ode(init_state)"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 22,
|
||
|
|
"id": "4bdfe23f",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.colorbar.Colorbar at 0x7fd5603ba4d0>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 22,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAATcAAAD7CAYAAAAPf9NJAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAA/iElEQVR4nO2de4wld3Xnv6du90x3T89Mz8OeGXu82MBAMCSYZGQgrLIEkuDNohitROREZEmEZCnyJmQ3q4Cz0kZZyStWu0LJSiHSKCRxFALx8lgsNuGxThxEQmzGYMAPjA1+jWc8755HT09P962zf1R11/ec6vp1dc+d2913zke66ntvVf3qV7+qrlvn/M75HlFVBEEQDBrZancgCILgShA3tyAIBpK4uQVBMJDEzS0IgoEkbm5BEAwkcXMLgmAguaybm4jcJiJPicgzIvLhXnUqCILgcpGVxrmJSAfA9wH8LIBDAL4B4JdU9YnedS8IgmBlDF3GtrcCeEZVfwgAIvIpALcDaLy5bchGdDTbXHzwN1X67G+3wh8yeti8EgHI3KbIytbLqs+6YYPdrFMtyy5estvNdakNOk7ffp4nutU0JolR9dvw7lY4xNLQf3V9lw6t1+nYRro0Hjl1pNbfxHky65kOJlb07dO6mdsX94WPLU8MXJbYd+qa5mW+jbJb03PncCmfbjkgi/Oun96kJ091l14RwCPfmfmSqt52Ofu7UlzOze16AC/S50MA3pzaYDTbjLduuR0AoF33Dzo7u/DWL+N/ABnZWK03O9e8s9RFIol/bL5A/T8bX1DUX7+ejI5Uzd10nVk2u7m62Y08ccjuevJM1QYdJ4bdDXJqinZmr2OdozExNwQ33tRnP47CN+dUG4kbfDY6Wi3aWB1LfuGCXW98U/Vh+4Rt/9Rktavpi9X3/ubO4+//6ekGKUN0udM5Ktrkm5T9xxY6Fh3daJfRD5ROT1fvL0zDrliNj4yNoRE+f268+f9CRmz/ZagYg386/tfNbbfkxKkuHvrS3lbrDu/5wc7L3uEV4nJubov9OtTuKCJyJ4A7AWAk21TbIAiCtYai63/I1iGXc3M7BOAG+rwXwGG/kqoeAHAAALZ2dur8r49esiaZ0K9v7a45PLzwNqdfcKmZgy3nR+jX3Jtx3Ka4JzLtNpiNXftLn585V7Xx3afNso30pKKXZs0yfurSKf6VdhdawiwVMoP5ScLDT2viTS0ywySr9qW5H1/qR8LMM0+TftnFmYX32bkpu2zHtmrZ2fPV97Nu3KgNfy7ME9Pm8Wob/wQ2Q5bDiH1SxgV6atwwbJfxU29LN4l58k6QXWsfii7dVH0+9Vrb/91fPVkucNbGClAA+Ur9EWuIy7m5fQPAPhG5CcBLAO4A8Ms96VUQBKtKjqv4yU1V50Tk3wP4EoAOgD9V1cd71rMgCFYFhWL2KjdLoap/A+BvetSXIAjWAAqge5WbpZeF95eZmSwH+2x4vZpPjPx4fsaV/WUy3O6w85kZ89nsj/vvfX2JKXudIV9jzT/EYS7k6/L+SR6rxGxvRrNytWMB+R1d2AL72ez3bmaWuu/PhYF8ZLXzztu5a0DHeMaYlh0+1rwv74fl8eFrYCbht9tsZzN1YnP1Yc6dMz6f7N/baH1ixp960V9X5OPcumXh/fRrrjXrvfSvKn+fuuG+9FThT9QXLt/nBoTPLQiCAUQBdAdAxDZubkEQ1Fj/Hrd+39yGhyC7rynenz5jl5EZqd6EYnOFg0+duWbCOlIOUTLD1JuGbCb5IErajk2JVAqbb1+SQcIUPDtL+xLXvu9z077JJPZmP7fo2zdZA6Z/9ns+lmznDts+h6FkzWNqMlNcCEZ3U/W5c766JvRGGxjdOXW2WsZhG4AJyGWXQM2MJtOzO25NSiabcWEtncqMZJNVpu01DDJFayFMfE3vqcI9Tr7OjscQRZCMHrPn7PRrij7PffeykhOKPkDD5xYEweChCsyu/3tb3NyCIPAIuosmIK0v4uYWBIFBkc77Xy/09eamnQz51mKaveN8Enq+Sq/xviiTvsN+Np+wTn6l2rnJF08BqqUemc/N7Rv/k/f9JRL6lXxOMjTcuF6tX9yGFx1gaDth31xKOaM2Bov76mqpahTu0N01YZZ1TpBPklKnvL+QfZLZlE0X65yvksPnNlf7unCdTRrf+hiFdXhfLqVLCfv0vN9yqDrX+QZ73oePnqUP9l/m0q4qpWvofNWPueu3mvU609U1MfTSKbtvDnU6XaXu7fmqPS9zW6sxOO3SrzZOFucmS2hJLIdBeHILJd4gCAxFEK+0ei2FiPypiBwTkcfou+0i8hURebr8u42W3V2K3z4lIu+i739CRL5bLvtfUpuVqRM3tyAIDApgVrNWrxb8OQCv9/ZhAA+o6j4AD5SfISI3o8hRf325zcdKUVwA+GMU6kL7yteSGnJ9NUslV2QXChOupsrA0eNuu/z0ZLWMVC9kk4skZz0tp7jRaOYlTNtsy2a7LoctpDINUqEmlIWgc66PbKZyv7wDJGGymnWlOewkIz2z2lg1hN74cBLZSKEax6w5qOcpboFVWFIKIV7IcqoK6+gMVeM28Y3Tdju6PmphOaSrpkKKMl5T7XhlKg4ft2ajCSFxY7CBwj90pDp/c9dY01k5e2H3NrMsO12Z7ULnj0NcAKAzWY3BtZNWPmy+/WymXZhQCoWg26PnHlX9qojc6L6+HcDby/f3AngQwIfK7z+lqjMAnhWRZwDcKiLPAdiiql8HABH5CwDvAfC3qX3HhEIQBDVyvaI+t12qegQAVPWIiMznmV0P4J9pvUPld7Ple/99kri5BUFgmPe5tWSniBykzwdKDceV0CSA20oY19P32dL56O+ho05IkWa1utu3mEUZmyuUveDNGDahsq3OpCTTK5Vgb0zPKSuJbUy7hPa/yUJoa0ICJnpfOjSz52bv2BzEtTYzAMdOVu9ZlnqjkyrfUs3yZc4sNdkinOTt5LfN+CRmnZWi+msZG1xjwo8HzXwaI2nOz7jSeKeyN8wssJuRv0h9dBkyLEaQ+6R3Wlc2Vabi2JzLbhmma/OkNTeNi4POkxflZPn67OWTZlm+u7gOfDLLyhB02/nTAOCEqu5f5g6Oisie8qltD4B5JYQmAdxD5Xv/fZKYUAiCwFAo8WatXivkfgDvL9+/H8Dn6fs7RGRjKYK7D8DDpQl7TkTeUs6S/jvappEwS4MgMKgKLnlNpRUiIp9EMXmwU0QOAfg9AB8BcJ+IfADACwDeW+xXHxeR+1BU0JsDcJfqgrDWr6OYeR1FMZGQnEwA4uYWBMEi5D0K4lXVX2pY9M6G9e8BcM8i3x8E8Ibl7LuvN7d8WDC9pyiTNv78UbOMY/I6J1yUeVPYQs1XRNPyCd+LLaJiVSS0ZRttlTlqyh+pUnncj0vN4Qegsm66wS7LOKOAQ0G8SCT5eWq+RfZPcmiF8wEZH5YTZzSZASwwmljPn08O7TGlVL24ZqIuqluxWs/5zmyNVNcGH6dr0hTama3G9OzrbLjHxslqvZFT59CE8bP5bA4eUxfyMp9VoSkfb0uKCYX177GKJ7cgCBzLmlBYs8TNLQgCw/yEwnqnrze3zvkZjH/tGQD1qHjWz5dhm1BushnYjPRmDJuUbjrfCE1S8rMXZjRhBd5s5OwCbi+V5O6T6FMimk21P309iLNVKIGcsyZOzmYwjUe+HNloNoc4Q8Gn8/HnxHGZeqHjNjNAqMZr11WjZxNeeBx9LQcO36mF9lA4TMbZIfa8SKomBmcXjLvC4py9sNWGMDEZhYZ4UU6+HrvbqvazQ8dtI6dJlNNd39mJIquiJpK5QrpXNoi3L8STWxAEBoVgVtf/rWH9H0EQBD0lJhSCIBhIFBJm6XLRbo68FC6s1Q7lkADnU2E/AvsavD/L+CGc/6qpFmfKR+NDDoDFwz/qdT/b+aJqNUcbag7JsKuBmVLZoDalwf9W++zGyvgTub6pH9OGPgEwBVFAfirxdVwbatL6ZbburFOU4X5tcP5JHp+8OTSGz0QtRcz4Fu05y3ZuX3jf3VGl/G19+CXbPguabm7223WOUX9H7HH
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 2 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"resx = res[0][-1]+pos0+x\n",
|
||
|
|
"imshow(cic_paint(jnp.zeros(mesh_shape), resx).sum(axis=0)); colorbar()"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 23,
|
||
|
|
"id": "9abade92",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.colorbar.Colorbar at 0x7fd53e682e00>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 23,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAATcAAAD7CAYAAAAPf9NJAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAA88klEQVR4nO2de5xdVZXnf6sq7ydUQqDIw5CkgCYgjw6JNv1hUEQDqGFssXlE7RY/cWaAwIx2Q5ixBXuwse12JPhoM0Ib5CXQCBmFKKI0YkMeQBQChDwpilQSKo8iDwiVqjV/3FOnfmdV7V3n3txUpW6t7+dTn9r37nP22WefU7v2Wns9RFXhOI5TaVT1dgccx3EOBT65OY5Tkfjk5jhOReKTm+M4FYlPbo7jVCQ+uTmOU5Ec1OQmIrNFZI2IrBOR68vVKcdxnINFSrVzE5FqAK8BOA9AA4AVAC5V1ZfL1z3HcZzSGHAQ584EsE5VNwCAiNwHYA6A4OR2hFRpbVU1AGC/mVT3ouPzu6ZuiEharpGOxebeyMR8AOG6AZDgca1UHkzHAcBg6sfb2hY8bhz1cfOg4Zm6vVUdQz5l/9uZus1tB9Iy3+cwyS6wG+m4AebaPI4HaHwOIAuP6V66FwCopjZbEW6Dx6ra1PEYjOHxMNeaKB1nvlM9MFM3tLUlLW/Tjqu1IguPv32h36X+D6d+jDRjyqPYat6rtqqOPu435w3i96Cto2ebNdtLfk5HVdnR6uA9as+ON38eZfrRmrT/RlsLtre1Zl+KIpk9e7Y2NTXlOva55577parOPpjrHSoOZnIbD+AN+twAYFbshNqqaiweWgMAWNuWfXTLWt9Ly+tM3TSaEOYOHNblOZbt5o+I4T82e9wO+lxXlR0e7sdjB94NHnfVoBFp+WvHzsjULRs+Li3/dMOvM3V/t3dbWub7PG3AkMxx33h3V1oeY15yHhO+F3uf3Gc7jjyxchs7ImM12vRjHo0B38uNZkJfOOSItLxq1IRM3WlvN3Qc996eYD/4XmpMP/g9m1U9KC2fTf0DgIHU5h7z/u0ePCotr6cyAEyifk2le7vJ3Cc/py+ZNpj6wPMDss/wXBpTAGhO7u285jdwsDQ1vYUVK5blOraqauDYg77gIeJgJreu/jt0Wi6JyDwA8wDgGPH9C8fpC6jadWPf42AmtwYAE+nzBACb7UGqugjAIgCYUjVA2/+T8soHyK6KppmVEP/H5f/g9r+0/RyC/5vb/47chu3HuoDYaFehF+7rWNJPM6uzpbRSWWrGILR6vdy0H1tNnU+rPB4rC1/LjptdDeYhNvaxFfbdLfvS8mf3ZUWhZaMnpeVraVXbZsaN27DPgu/l1ERqAIDnzeqpltqoN6u6E2mlvMWsmMZSX3ZGpAWGV94xPnXE5Mzn742fmZYX112QqfuLZ78DANi9Z0uutuMoVK3w3/c4mMltBYA6ETkOwJsALgFwWVl65ThOr6Gq/XvlpqoHROQqAL9EQZ98h6quLlvPHMfpJdrQ1vZu94cd5hzMyg2q+iiAR8vUF8dxDhP6u1h6UFgdDevVLKyz4eOsToz1eHbHNbRrFoN1OfZ6Mb0U35u9z3ty6od4Z8zqJ7n/sd3eG0mvZO8lpncMYe8ltrPMxPR7PKbN5rmsIf3kdtIlzmx6NXgt+1x4fIbRPdea8eBdzxeGZTcAnxx5bFoebfSH3M6RdO25RjfHYxDTN/8DXXv2xLMyx33wQzel5WozVovXFtYY298qh5lpPxdLHcepTPq9zs1xnErFJ7eiGThgMGprpgEAbt2dtRpZTyLDXUZkYDGsLmBIC8SNVkPHWdGQ24+JfCxKxMQ6Kx7HRLmQeYlt3/Y51D63FxPFbftW3O+qPXvel40B7kgyYh1B7Q0012LzidXGBOPFIUem5TPe2Z6WHznm9MxxF7zdYbg61RjPskFuE71Xw8w9Pj3imLT8/NAxCDGpZW/m8z5qZxSJnrYfF0dMmPg9uH3siWl5ykl/kTluz56taXnbthczdcOnFZwEqjb+Jtj3YvDJzXGcCqQNbW37e7sTB41Pbo7jZHCdm+M4FYtPbkWyp2ogfp/oNqxOYsc7O9Ky1UWFttHtcXnNSZiYC1fMXIXrrO4v5m7EeqpYf2PuTFaPFzqPdXN5TVfsZ+6j1c1x3b2JLrWdObs2dZxHblVWX8j3colxS2I9Gwcc2Hz0+zPHTXrt52l5k3E/mk0mJK1knjHQ9ONtupdGE8nlc9tfS8vbTRCDxWOO7+gv3edt407OHMcuXP/FmGuwecl5FCyg5T/+OXMc6wUHHX9hpm5XMt6tkXcvP5XhfuWe7I7jGApiaZ6f7hCRO0Rkm4i8RN/ViMjjIrI2+X0k1S1Igt+uEZGP0fd/KiIvJnULRaTbsE4+uTmOk0FV0db2bq6fHPwYgI33dj2AJ1S1DsATyWeIyEko+KhPT875fhIUFwB+gEJ0obrkp9sYcj0qlg7R1nR5buNinUJiwnHmvGUkanDUixlkwQ4AK2npb0XFkJgXE20vMZbqbMbQFPE0iJmacD+s+BoyeYlFLrGEIoZYsXQ+mV08GvGA4H5MMmJ0I4l5n965IdsRUjPwGMRE9mFmrE6mNth74eoX78kc10pmRWvNWHFctTFkjmFjqrEYzWUg61ViVQls/vEGjek9RkxvofG/k0RZICuK8jvGJi4AcM6exrT8eyoDQEsSgexlo+4pjfKJpar6lIhMNl/PAXBOUl4M4EkA1yXf36eq+wFsFJF1AGaKyCYAo1T1GQAQkTsBXATgsdi1fUPBcRzDId8tPVpVGwFAVRtFpF2hOh7As3RcQ/JdS1K230fxyc1xHENRk9tYEVlJnxclMRxLIRQAN1dgXEuPTm77qgak1t+f27E2U9dE4uajoyZm6v6BxJXVJCJYJ2YWGWYbkZIDQ8Yc7NkrYaXZvWMRbSqVrWgb8hKwxMJIzwyIhgBwGYmDDx45JVPH4iF7fYw1jty/p523T5mgliweskO5Db/dQuPT6T4D4rcdb1YzWO+F6caLpR373JdRmzHvDe6jddJvpvu0HjL8nGwAgpm083kDqUlsH/n9vqC5PlPHzvcs6o816gJW5fAOLtAh6h5U8oQE1aJMQZpUdUb3h2XYKiK1yaqtFkB7JNJQANyGpGy/j+IbCo7jGMq3WxpgCYDPJ+XPA3iEvr9ERAYnQXDrACxPRNjdIvKBZJf0c3ROEBdLHccxlM/9SkTuRWHzYKyINAD4GoBbANwvIlcAqAdwMQCo6moRuR+FDHoHAFypHTsb/xWFndehKGwkRDcTAJ/cHMfpRPk2FFT10kDVuYHjbwZwcxffrwRwcuczwvTo5PbWgKH4wbjpAID/teWFTB1vlc9pfj1Tx7qGGtJD1Bu9ButbYlb8rOe5zeib1uVsI6bbYWL6uFjkEtYL7jD6oYvJ5KAxkOINACaxjsnoxDhn6olGt8j6yWXUx1gKQBucsZZ1TFS2ujk+rt7omDjBzY30fcw0JhahheuazXMvNWnQchqTVdT/H0w5L3vx3W+mRTb9sLCezZrG8LO2+snGAUMBZP+OSqVIndthi6/cHMcxVIb7lU9ujuMYPCpI0fzJvrfws+d/BAB4OGLGAeOcvIHE0vNInGo0Ysy6SMDLkMN6LDepFRvHBMSfYnJ2xsSmUDtWPP4COWjX1D+dqfsJ3U8peRLsebFs7qHxAIBaKk8gEwmbn2AseZ/YbPQ7AiKx9bbg42KmPTHvkFjABL7e14xnA3svjCTnfguLmzYvKpuJcL6Gq7e9lDnuQ2Qa84R5v49NRN2B+3cH+5Afn9wcx6lA2n1L+zo+uTmOY3Cdm+M4FYmLpUVT39aK+YnZwUxj3sDmGVsjZgvsfmX1JqxniwVnZGI6GqtHCumtymGaEMOO1bpIlI28SWxigSxDUU1iY9rJNIaexa2kc9tnzE5GB3LSAmG3rVnG7CQ23txGKGKK/dwpqQ+1v9Ocd+LoSWn5QSoveunezHHsYvW80Tuyixu7z1nd3K/JhCSkM9zfvctlLiphcut2Big22JzjOH2dQ+5+1SPkWd78GDmDzTmO0/dRLbhf5fk
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 2 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"imshow((flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]).numpy()[0]).sum(axis=0),cmap='gist_stern'); colorbar()"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 24,
|
||
|
|
"id": "a7f8b5d5",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.colorbar.Colorbar at 0x7fd5382000d0>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 24,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAATYAAAD7CAYAAADgvbh3AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAA5ZklEQVR4nO2df6yk51Xfv2dm7q+9u+vd9WZtx17sBNwkbgpOtBiHIOo4mBo3jaWWoARRAkrlgpIqtCDiFAlUpEquWlGKEoi2kCaIND8EIbZSk2BcrEAFwWvihNiO8ZI49mbXXu+u9/f9NTOnf8zc+37Peed57jv3zs7MnT0faXTfue/zPu95f9znPuc854eoKoIgCCaJ2qgFCIIgGDQxsAVBMHHEwBYEwcQRA1sQBBNHDGxBEEwcMbAFQTBxbGpgE5E7ReRpETksIvcOSqggCILNIBv1YxOROoC/B3AHgCMAHgXwLlV9cnDiBUEQ9E9jE8feAuCwqn4TAETkUwDuBpAc2Bpz8zq9Y0/vnTy+SuasVdttFO6zndmXQXPtaI4srUz/uevcaj7VuWuh7/6+Cd//QVwz9S/+2fKpNipj4pjOgdXkylKhj+Vzp9BcuLCpv4x/9pZ5PXnKv5y9eexrS19U1TuTYol8FMDbABxX1df32H8bgPsBfKv7q8+q6q/3K7NnMwPbtQCep+9HAPxA7oDpHXvwmn/17wEA6pTgWrPYbk2L21c8UR4MfDtG2u4tSA0aDm0UDeuLtmGb7pbW0+duZ+5qc1uxPXPa9m/6pD8Ydf0J3atay/XBYmX+kHODr2iiXeaP2venNbqPK0WH/t60pyS5b+oCPffMuYxM7tmm5K8vpftoTbs+6rzPnryx0PvdLD2zzCBtrjvzzLKDcffv6ek/+u/pRhU5caqFL3/xukptp675h73rNPkYgA8B+P1Mm79Q1bdVk64amxnYer1epSFDRO4BcA8ATG3fvYnTBUEwHBQtzYyi/fSk+iURuWEgnfXBZga2IwD20/frABz1jVT1IICDALBt335d/Q+am0GUZlsE/3fjmUC5D3fcFO1LzUgACM0Om3OZGSGfug81afpsWkY/iy1kst95ltZuWBnrS8U+7q/Ut1Gv/HSr2gXlZk48y2mb2ZY/V7FZX7S7+Ds/91IXdf5SbV9rxrbjZ+Fnc63ZYnvqonvn6L62aabnr7OG3s8FSGsryzttu+s/8e217fNvuNbsW9jTudDcrK4qCqA9XHvHm0Tkq+iMH7+kqk9stsPNDGyPArhRRF4F4DsA3gngJzcrUBAEo6ddMjAn2Ssih+j7we5kpip/C+B6VT0vIncB+ByAG/s4vicbHthUtSki7wPwRQB1AB8dxEgbBMFoUShWqquiJ1T1wIbPpXqWth8Ukd8Wkb2qemKjfQKbm7FBVR8E8OBm+giCYLxQAK0hqaIicjWAF1VVReQWdPwGTm62300NbBthdeVPnH2MbQO8UtY9qtgi20XbL4ER3vYkyn1kBOSVs+Xc8imdy8nLx7Vm7D5eFa0vpld/jY3GyWvO5/6xtmZ7r+qWVgvZ7QRpW6XB2yMzt6exmLApeftdu+jU2MMcU7T6uLzDPVv2TCitilJbtoc5Oxf3sfCKtBxzx+0Jast0rjpvp28Orxh3vhdtT3/vytr2t/7F/zTtbvr+n1rbvuHel8y+2T3bAQD1xcEY/QdlYxORTwK4DR2V9QiAXwMwBQCq+hEAPw7g50WkCWABwDt1AEkihz6wBUEw3iiA1oAS0Krqu9bZ/yF03EEGSgxsQRCUGMy8b3QMd2CTYroubkk9t2RvnDN5uu+W1NlNpJZxGTFuIRmHznY9p86yv0TyVMaBEwDay6RKT9m23CdfS8kZuOJbZ1xcnCO5UUUzERal6IgUGfWeXRi8ep9zSWG1ne+bb9dYoHZOnTWuQ8ZZ1wrcmiu257+TltG7ALEjLj9rLyObJGrOfYffn9f8buHj8kMP/VvT7PW/cHht+/F/84/Mvt3dZbvW4c3ntVDo0Gxsl4qYsQVBYFAFVrb2uBYDWxAEHkHrkgRiD48Y2IIgMCiAjCVnSzDUgU1awMzLnTu2sM+5OhSr3CYsCLA2GraJcfgT4GxFPuwm9Q8oY8+ruf6Nu0BmOT8H2+nYNtTZR/1zCFHLC5nuX9l9gtwKfB+50Btj02RbXMZlpLnNCjV1nuydnMTA2ZeEOm3t8P4kveW/eLVtN3OKzutDnhq9baGlBAdkx1zZbvufPpew08G+q/wOz5yx/Z96bXGgt+Gxze3kP9m+tr3voedNu9O/XASmz7zZCnL21Z2fPoB/o8SMLQiCiaLjoBsDWxAEE4QCWEllZNgiDHVgUwGa3WX16dM5Nc/ukoQKwVN/AFjaQ0vqy3YfqyHGlSLz/LSUl4v6p2WjUoaQTOYP4z7hvThIFr42n6cspSp29vVWkUsZJTJuHMYdJnFewKpvdadWs3sDq2vexcW4Qbhnxsct7yza7fh2Jk+ev06TCy49C2EV1udca5IriL+/HAWxjfpYnvd524ptjg4BgDb3SZsv3b7ftsv8tTYursqXblMVhaC1xcuhxIwtCIIS7Wzc4fgTA1sQBIawsfWJNgp1ceq828exym410qov7Llv2zUuZtQO4+GeDiLnJI5eqeOVOROV4GftvAqYSXzoPdA1oVJ5lTuXg5+vzaQNd3Jw2u2WC+L3SRhT7Uwkhgvsri33VvdL6a5JrsYFKySrqauqFlAOls+ZE8wqeiutbtpaFOnIA0+DkmGyyuqf2fzR4sL9uU2ECN0fTnDp+5w90dvkUIpq2BCCVtjYgiCYJDoZdGNgC4JgglAVLOdySG0BYmALgqBEO2xsfdAuCnR4O45Z6veeILSP7RgzZ1z3XEDEuRVwH8bdw9viuIRaLW1TMjP1UlGWdCYKax/LZCCp994G8hk3TBEV8tz39hq2Qfp7xRibTS7vppOxnnie2ZKJvmgNJWBk21Pp+tk26TO+tHrbAb3bBp8798xK5+Z3gk23zlWozdftc21OJd4Xd6tWiqCEkmuML+W4GTqLB6GKBkEwUcTiQRAEE0YsHvSLYG16Xaohye4BmbqR7DG/4gKveZ+fqktCZfDkvNOtGpJ2pTB1S/37wfu8fbaiNsFqTdWoB3+/Ge8iYGsvcOe545yrxlRvta/tTBB1Ss7o60OY3vlNdWmrTQV2rwKit8pda3kTRM4swEkiXQA7JUzl/pd3+boGxfb8UdvHIplXuFaEV3v5/njzwWpkxqAmWq1w0A2CYJJQCFZ0aw8NW1v6IAgGTiweBEEwcSgkVNF+kHZhUykVKOEkiz7khOwmbHeoZRKzl5f9aV+df199mdw8a2NH89dSsc9cs9x7xV4nXv5En7miNaXuK/6zzmWbMHY1ep4NlwiybmxU6awdnD3EuwpVTahpw6tcH+xe48PD6D3z7irsksJhdj5DDYcQlgq9kCuOyQKScY3xyVjXrmdAXh9bffFgXelF5KMiclxEvk6/2yMiD4nIM92fuy+tmEEQDAtVoKW1Sp9xpYpkHwNwp/vdvQAeVtUbATzc/R4EwQTQWTyoV/qMK+uqoqr6JRG5wf36bnTK1gPAxwE8AuADlc7YnV2X1EhWr9w+np5zPvpSUj1WD92VsRuHcf0opfCgXU6dNR7o1IfkVM/MrlLgQaqtFyOXKDOhvZQuk29BqX/ax/0P4FraLkMIR1+UMn+Qumg8+f018jPLmAXYPcVn3+CIFi977h6kXE3mXkqr1SU3pUTUg3ct0VSGGgBYrcExINPY5bp4cJWqHgMAVT0mIvsGKFMQBCNEIZFocj1E5B4A9wDA1PYwxQXBVuBynbG9KCLXdGdr1wA4nmqoqgcBHASAbfv26+pU3qskLIlf6eO88jxt957qnNywtFrIMhmV0u2jbSnpJMku0/h/fBkVMKlSZdqVywemkz8aqt4fUg/bfiWbv7r+eCXU1ofwHvkUReFMNg2KSsgF6qfqPADAyo7e/W87ln62pefO5/JmB843auor2Gb83rYbbmV4gc0kJFIpoWY6AmKQKID2GC8MVGGj0j8A4N3d7XcDuH8w4gRBMHo
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 2 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"imshow((cic_paint(jnp.zeros(mesh_shape), resx) - \n",
|
||
|
|
" flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]).numpy()[0]).sum(axis=0)); colorbar()"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 25,
|
||
|
|
"id": "68b72b68",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"from DifferentiableHOS.pk import power_spectrum"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 26,
|
||
|
|
"id": "785818e3",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": [
|
||
|
|
"k, pk = power_spectrum(flowpm.cic_paint(tf.zeros([1]+mesh_shape), final_state[0]),\n",
|
||
|
|
" boxsize=array([64.,64.,64.]), \n",
|
||
|
|
" kmin=0.1,dk=2*np.pi/64.)\n",
|
||
|
|
"\n",
|
||
|
|
"k, pk_jax = power_spectrum(tf.convert_to_tensor(cic_paint(jnp.zeros(mesh_shape), resx).reshape([1,64,64,64])),\n",
|
||
|
|
" boxsize=array([64.,64.,64.]), \n",
|
||
|
|
" kmin=0.1,dk=2*np.pi/64.)\n"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 27,
|
||
|
|
"id": "18e3d30b",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"<matplotlib.legend.Legend at 0x7fd53838db40>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 27,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAXYAAAD8CAYAAABjAo9vAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAjJklEQVR4nO3deVxVZeLH8c9z2cEtQU1wwV1BFBAQTafSFs1c2svKaizHMes37WrO1GRli5XZMmVlq2X9bFHLamrMn5qigOKKW66IKaKgIMj2/P6oaczRQlnOvZfv+/Xyj3u655wvll+fnvOcc4y1FhER8R4upwOIiEj1UrGLiHgZFbuIiJdRsYuIeBkVu4iIl1Gxi4h4GV+nAwCEhYXZyMhIp2OIiHiU9PT0A9baJidud4tij4yMJC0tzekYIiIexRiz82TbNRUjIuJlVOwiIl7G0WI3xgw2xkzPz893MoaIiFdxdI7dWjsPmJeQkHCbkzlExFmlpaVkZWVRXFzsdBS3FBgYSIsWLfDz86vU993i4qmI1G1ZWVnUr1+fyMhIjDFOx3Er1lpyc3PJysqiTZs2ldpHc+wi4rji4mJCQ0NV6idhjCE0NPS0/m/Go4t9+4ZUNi7/p9MxRKQaqNRP7XR/bzy62HO/eoLOX15F5mO9yPj2AyrKy52OJCIebNq0aXTp0oWIiAjGjh3rdJwz5tHFHjVqBimdHqBRaQ6xS0az67HupH72IiXHdAFGRE7fyy+/zPz583nsscecjlIlHl3swfUaknzdBMImrCct/kksPiRmPMihyVGkzHyEwiN5TkcUEQ8xevRotm3bxpAhQzh06NAv23fu3En//v3p1q0b/fv3Z9euXZSXl9O2bVusteTl5eFyuVi0aBEAffv2ZevWrTz88MPceOON9OvXjw4dOvDaa68BsHDhQs4991yuvvpqOnbsyLhx45g5cyZJSUnExMTwww8/VPln8YpVMX7+ASQMGY29dBSr/+9jfJc9T/KWZ8h/5hWWRVxNxyH3EtqshdMxRaQS/j5vPRuyD1frMaPCG/DQ4Ojf/M4rr7zCV199xXfffcfnn3/+y/axY8cyYsQIbrrpJmbMmMGdd97JZ599RseOHdmwYQPbt2+nR48eLF68mJ49e5KVlUX79u0BWLNmDSkpKRQWFhIXF8egQYMAWL16NZmZmTRu3Ji2bdty6623smLFCp5//nleeOEFpk6dWqWf16NH7CcyLhfdz7+K6AlL2HTpp2wLiaNn1luEvBzL8hduZs+2TKcjioiHWbZsGcOHDwfgxhtvZMmSJcBPI/NFixaxaNEixo8fz5IlS0hNTSUxMfGXfYcOHUpQUBBhYWGcf/75rFixAoDExESaN29OQEAA7dq146KLLgIgJiaGHTt2VDmzV4zYT6ZTQj9I6MeuzRn8+OVTxB2Yh8/bn5He4DwaXnAf7buf43REETmJ3xtZO+3fK1T69u3LK6+8QnZ2No888ghPP/00Cxcu5A9/+MN/fffEzwEBAb9sc7lcv3x2uVyUlZVVOaNXjdhPplXHWJL+533y/5TOivDr6Xg4hfafXsLayeezfvEcbEWF0xFFxI317t2bWbNmATBz5kz69OkDQM+ePVm6dCkul4vAwEBiY2N59dVX6du37y/7zpkzh+LiYnJzc1m4cOGvRvM1yeuL/d+ahEfS608vUXHXepa1vYPmx7YR/a8RbH08kZXzZ1BeDX9Lioj3mTZtGm+++SbdunXj3Xff5fnnnwd+GnW3bNmS5ORk4KcR/JEjR4iJifll36SkJAYNGkRycjJ//etfCQ8Pr5XMxlpbKyf6LQkJCba2n8deXFTI6i9eJXz9a7S02ewxZ7Mn6la6XTqGwKCQWs0iUtdlZmbSpUsXp2NUq4cffph69epx7733VsvxTvZ7ZIxJt9YmnPjdOjNiP1FgUAg9r7yb8AfXsjL5eQp9GpC0/lEKn+xCylsTyD90wOmIIiJnpM6O2E9kKypYv+wLyhc/T/fiVAptIGubX07bwffTNKJyD94RkTPjjSP26qYR+xkwLhddzxlM93HfsvWKr8hs0IeEvbNoNL0HqVOvY+fGVU5HFBGpFK9d7lgV7WN6QczHZG/fyK4vnqJ7zjyCZs0nI7g3Qf3uoVPCBU5HFBE5JY3Yf0N4m84kj51B0e0ZLGtxK5FH19Dp8yvY8Pg5rFnwoZZKiohbUrFXQuOmEfS69Rn87llPSsf7aFzyI90WjWLHo7GkzX2Z0pJjTkcUEflFtRe7MaaLMeYVY8xsY8yfq/v4Tgqp34jk4RMJnbCB1NjJACSsHE/u41Esf38SRwvynA0oIlXSu3dvpyNUi0oVuzFmhjFmvzFm3QnbBxhjNhljthpjxgFYazOttaOBq4H/ulrrDfz8A0gcNobWD2aQ0Xc6B/3OpufmKZRMiSbl9bs5lJPtdEQROQNLly51OkK1qOyI/S1gwPEbjDE+wEvAQCAKuM4YE/XzPxsCLAH+VW1J3ZDLx0Vs/2uIevB7Ng76hG3B3UnOeoPAF7uz4qU/kr19o9MRReQ01KtXj4KCAvr37098fDwxMTHMmTMHgNTUVLp160ZxcTGFhYVER0ezbt263zmiMyq9jt0YEwl8bq3t+vPnXsDD1tqLf/48HsBaO/m4fb6w1g46xfFGAaMAWrVq1WPnzp1V+DHcx85Nq9j35VPEHvoaF5bVDc/nrIvup23XZKejibitX63R/nIc/Li2ek9wdgwMfOJ3v1avXj3y8vI4evQoDRo04MCBAyQnJ7NlyxaMMUycOJHi4mKKiopo0aIF48ePr96cv+F01rFXZbljBLD7uM9ZQE9jzHnA5UAAMP9UO1trpwPT4acblKqQw6207hRH604fsG/PNrbNe5qYvZ9Sb/bFrPk8AZ++dxHV6xKMS9esRdyVtZYJEyawaNEiXC4Xe/bsYd++fZx99tn87W9/IzExkcDAQKZNm+Z01FOqSrGf7O2q1lq7EFhYheN6hWYRbWk2+h/kH5zEsrnP0mHHTMK+uZ7N33WkMGEM3S+4EZevbiMQ+S+VGFnXpJkzZ5KTk0N6ejp+fn5ERkZSXPzT6zYPHjxIQUEBpaWlFBcXExLins+VqsrQMQtoedznFoCuGp6gYeMwet38OPUe2EBK1F8JLj9CXMpf2PNYDKmzn+FYcaHTEUXkOPn5+TRt2hQ/Pz++++47jp8mHjVqFJMmTeL666/ngQcecDDlb6tKsacCHYwxbYwx/sC1wNzTOYAxZrAxZnp+fn4VYniGwKAQkq++l+YPriM96TmKXcEkrnuEI09EkfLORA7n5TodUaTOM8Zw/fXXk5aWRkJCAjNnzqRz584AvPPOO/j6+jJ8+HDGjRtHamoqCxYscDjxyVXq4qkx5gPgPCAM2Ac8ZK19wxhzCTAV8AFmWGvP6NXe7vAQsNpmKypY9/3n2CVT6XYsnQIbxLqWw+l+3d8JCqnvdDyRWuUODwHLzc0lPj4ed13IUe0XT621151i+3x+4wKpnJpxuYjpOwT6DmFLxvcc/vZJkrPe4Men55DZ80HiLr5ZF1lFakl2djbnnXdetT073Wm6eucGOsSeA7Fz2ZDyNQHfjCN++V2sX/UmIcOeITIqyel4Il4vPDyczZs3Ox2j2jg6JKxLc+yVEZV8Ma3HrSAl6kEiSrbR4sOLWf7SSPIP5TgdTUQ8iKPFbq2dZ60d1bBhQydjuBVfPz+Sr74fxq4kLWwoCfs/puL5eFI/fpYKvZdVvJg7vPTHXZ3u740mcd1Uo7BmJN/xFtuv+JK9fq1IXPt3fniiJ5vSvnU6mki1CwwMJDc3V+V+EtZacnNzCQwMrPQ+ejWeB7AVFaTPf51WaZNpykHSGl5M5LVPE9a8tdPRRKpFaWkpWVlZv9wIJL8WGBhIixYt8PPz+9X2U62KcbTYjTGDgcHt27e/bcuWLY7l8BQFR/JY88HfSNgzk1J8WddhNPFXj8fPv/J/k4uI93DLYv83jdhPz+6t68j9+B5ii1LY5Yog/9xJxJx7hdOxRKSW6WXWXqRl+67EPvA1GX2nY6wl5rs/kvHUQLK3ZzodTUTcgIrdg8X
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 1 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"loglog(k,pk[0], label='flowpm')\n",
|
||
|
|
"loglog(k,pk_jax[0], label='jax')\n",
|
||
|
|
"legend()"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": 28,
|
||
|
|
"id": "c2c1ecdb",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"text/plain": [
|
||
|
|
"[<matplotlib.lines.Line2D at 0x7fd530740190>]"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"execution_count": 28,
|
||
|
|
"metadata": {},
|
||
|
|
"output_type": "execute_result"
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"data": {
|
||
|
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAY8AAAD8CAYAAACPWyg8AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAr1UlEQVR4nO3deXiU5b3/8fc3GxCWQCAJIQl7WMIqhMUNcQGBarEuiNVKWy11O9a2p5aec3pau5za/rQed0Vt61YBbVWwVERwwwISViEsCUFICISEnQSy3r8/MnowZpvMJDOTfF7XlWtmnue+n/k+KPnwbPdtzjlERES8ERboAkREJPQoPERExGsKDxER8ZrCQ0REvKbwEBERryk8RETEaxGBLqCl9OjRw/Xt2zfQZYiIhJT169cXOefiai5vM+HRt29fMjIyAl2GiEhIMbO9tS3XaSsREfGawkNERLym8BAREa8pPERExGsKDxER8ZrCQ0REvKbwEGnDTp4pJ//Y6UCXISHIL+FhZtPMbKeZZZvZvFrWm5k94lm/xczGNNTXzGLNbLmZZXleu521bqSZrTazbWb2qZm198d+iLQ1d7+ykSl//IC9h4sDXYqEGJ/Dw8zCgceB6UAacIOZpdVoNh1I9fzMBZ5sRN95wArnXCqwwvMZM4sAXgJuc84NAyYD5b7uh0hb81lRMe/tLKS4rJIfLtxERWVVoEuSEOKPI4/xQLZzLsc5VwYsAGbWaDMTeMFVWwN0NbPEBvrOBJ73vH8euMrzfiqwxTm3GcA5d9g5V+mH/RBpU/76yT7Cw4z/+tpQNuw7xmPvZQe6JAkh/giPJCD3rM95nmWNaVNf3wTn3AEAz2u8Z/kgwJnZMjPbYGb31lWYmc01swwzyygsLPRyt0RarzPllSzKyGVqWgK3Xtifb5yTxKMrs9mw72igS5MQ4Y/wsFqW1ZwYva42jelbUwRwAXCj5/UbZnZpbQ2dc/Odc+nOufS4uK+M6yXSZv1jywGOlZTzrYl9ALhv5jB6dmnPPQs2caq0IsDVSSjwR3jkASlnfU4G8hvZpr6+BZ5TW3heD521rQ+cc0XOuRJgKTAGEWm0F9fspX9cR84d0B2ALu0j+d/Zo8k7WsIvF28LcHUSCvwRHuuAVDPrZ2ZRwGxgcY02i4GbPXddTQSOe05F1dd3MTDH834O8Kbn/TJgpJlFey6eXwRk+mE/RNqErfuPsyn3GDdN6IPZ/x38j+sby50XD+S19Xks/fRAACuUUOBzeDjnKoC7qP6lvh1Y5JzbZma3mdltnmZLgRwgG3gGuKO+vp4+9wNTzCwLmOL5jHPuKPBHqoNnE7DBOfcPX/dDpK14ac1e2keGcc3Y5K+su/vSVEaldOVnf/+UA8f1/IfUzZxr6BJD65Cenu40n4e0dcdPlzPhf95l5qgkfn/tyFrb7Ckq5muPfMTolK68dMsEwsJquzQpbYWZrXfOpddcrifMRdqQv63P40x5Fd86t0+dbfr16MgvrkzjX7sP8+yqnBasTkKJwkOkjXDO8dLavYxO6crwpJh6285KT+HyYQn8v2U72ZZ/vIUqlFCi8BAJYkeLy3jmwxyK/XD77Ordh8kpLOamiXUfdXzOzLj/6pF0i47iBws2caa84edwnXNUVbWN0+DShuYwFwk1zjl+tGgT7+0sJGPvEZ68caxP1x9eWruXrtGRXDEysVHtu3WM4sFZo/jWc5/wu6XbuW/m8C+tP3yqlC15x9mcd4zNucfYkneciHDjz98eT1qvLk2uU0KDwkMkSP354894b2chFwzswbJtBTy8IosfThnUpG0VnDjDsm0FfPf8vrSPDG90vwtT47jlgn48t2oP/Xp0pKyyis15x9mce4y8o9V3Y5lBanwnLhkSz8fZRXzrubUs/P65DIzv1KRaJTQoPESC0Nb9x7n/nzu4bGgCz9w8lntf28LDK7IYlNCZrzXyyOFsr3yyj8oqx40TGj5lVdNPLh/Mx9lF/HJJ9eNUSV07MDqlKzef24eRydXXTzq1q/5VklN4illPr+HGZ9fw6vfPo3f3aK+/T0KDbtUVCTIlZRVc8egqiksr+OcPJhHbMYrSikq++cxatuUf57XbzmvwgvfZyiuruOD3KxncswsvfHd8k2o6dOIMmQdOMDwphh6d2tXbdufBk1w/fzWd2kXw6m3nkhjToUnfKcFBt+qKhIj7Fmeyp6iYh64fTWzHKADaRYTz1E1j6RYdxdwXMig8Wdro7a3YXkDBiVJumtC7yTXFd2nP5MHxDQYHwOCenXnxuxM4XlLOjc+s9apWCR0KD5Eg8taWfBZm5HL7RQM4b0CPL62L69yOZ25O50hJGbe/tJ6yisbNv/HSmn30imnPJUPiG27sJyOSY/jzd8Zx4PgZvvXcWo4Wl7XYd0vLUHiIBIncIyX87O+fMjqla50XxocnxfDAdaPI2HuUn7+xlYZOO+8uPMWq7CK+OaE3EeEt+9c9vW8sz85JJ6eomDl//oQTZzRnW2ui8BAJAhWVVdyzcBPOwSOzzyGynl/0V4zsxb9dMpCFGbn85V+f1bvdl9fsIyLMmDUupd52zeX8gT146qYxZOaf4Ja/rKOkTMO9txYKD5Eg8MiKLNbvPcpvvzG8UXco/fCyQUxNS+DXb2XyUVbtE52dLqvktfW5TBvek/jO7f1dcqNdMiSBh2efw/q9R5n7wvpGPXAowU/hIRJga3MO89h72VwzJpmZo2tOwlm7sDDjoetHMyihM3f9dSN7ioq/0mbJ5nxOnKn4YsKnQPrayET+cO0oVmUXMXv+GgpOnAl0SeIjhYdIAB0rKeOehZvo070j980c5lXfju0ieObmdMIMbn1+3VeuKby0di+DEjoxvl+sP0tusmvHJvPUTWPZVXCSrz+2is25xwJdkvhA4SESIM45fvq3LRSdKuWR2ed88aCdN1Jio3nixrHsPVzCD17ZSKVnbKnPhwu5aeKXJ3wKtGnDe/K3288jMjyM655ezesb8wJdkjSRwkMkQP76yT6WbSvg3suHMCK58Q/91XTugO788uvDeG9nIX9YtgOonmY2Oiqcb5zTuNNgLWloYhcW33UB56R05YcLN/O7pdu/CD0JHRqeRCQAdhWc5FdLMrkwtQe3XNDP5+3dNLEPOw6e4OkPcujZpT1LNudzzdhkOreP9EO1/hfbMYqXbp3Ar5Zk8vSHOewqOMnDN5xDlyCtV75KRx4iLexMeSV3v7KRzu0jeHDWKL/N1PeLK4cxsX8s9y3JpLSiipuaMI5VS4oMD+PXVw3nN1cN56OsIq56/GNyCk8FuixpJIWHSAv7n6Xb2XHwJA9cN8qvt9BGhofxxI1j6dM9mvMGdA+ZYdFvmtiHl26dwLGScq558l96mDBEKDxEWtDyzAJeWL2XWy/ox+TB/h8uJLZjFMvumcRzc8b5fdvNaWL/7jz+zTEcLSnnw121P7ciwUXhIdJCDh4/w72vbWZYry78ZNrgZvue9pHhdIhq/JwdwWJ8v1i6RkeycvuhQJcijeCX8DCzaWa208yyzWxeLevNzB7xrN9iZmMa6mtmsWa23MyyPK/damyzt5mdMrN/98c+iDSnyirHDxduorSiikdvOId2EaH3y725hYcZFw+O572dh3T3VQjwOTzMLBx4HJgOpAE3mFlajWbTgVTPz1zgyUb0nQescM6lAis8n8/2EPBPX+sXaU6VVY6Psgq58+UNrM45zC+/Poz+cZphry6XDo3naEk5G/cdDXQp0gB/3Ko7Hsh2zuUAmNkCYCaQeVabmcALrnoI0DVm1tXMEoG+9fSdCUz29H8eeB/4qafdVUAO8NUxGUQCzDnHlrzjvLkpnyVb8ik8WUrndhHcefEArhubHOjygtqkQXFEhBnvbj9Eet/geDJeaueP8EgCcs/6nAdMaESbpAb6JjjnDgA45w6YWTyAmXWkOkSmAPWesjKzuVQf6dC7d9MnwhFpjM+Kinlj034Wb8onp6iYqPAwLh4Sx1Wjk7h4SLxXc4e3VV3aRzK+XywrdxQwb/qQQJcj9fBHeNR2k3rNE5Z1tWlM35ruAx5yzp1qaNg
|
||
|
|
"text/plain": [
|
||
|
|
"<Figure size 432x288 with 1 Axes>"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
"metadata": {
|
||
|
|
"needs_background": "light"
|
||
|
|
},
|
||
|
|
"output_type": "display_data"
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"source": [
|
||
|
|
"semilogx(k,(pk_jax[0] - pk[0])/pk[0])"
|
||
|
|
]
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": null,
|
||
|
|
"id": "703c8195",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": []
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": null,
|
||
|
|
"id": "d6921f17",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": []
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": null,
|
||
|
|
"id": "b3083d29",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": []
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": null,
|
||
|
|
"id": "65116b88",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": []
|
||
|
|
},
|
||
|
|
{
|
||
|
|
"cell_type": "code",
|
||
|
|
"execution_count": null,
|
||
|
|
"id": "c6bd8cfd",
|
||
|
|
"metadata": {},
|
||
|
|
"outputs": [],
|
||
|
|
"source": []
|
||
|
|
}
|
||
|
|
],
|
||
|
|
"metadata": {
|
||
|
|
"kernelspec": {
|
||
|
|
"display_name": "Python 3 (ipykernel)",
|
||
|
|
"language": "python",
|
||
|
|
"name": "python3"
|
||
|
|
},
|
||
|
|
"language_info": {
|
||
|
|
"codemirror_mode": {
|
||
|
|
"name": "ipython",
|
||
|
|
"version": 3
|
||
|
|
},
|
||
|
|
"file_extension": ".py",
|
||
|
|
"mimetype": "text/x-python",
|
||
|
|
"name": "python",
|
||
|
|
"nbconvert_exporter": "python",
|
||
|
|
"pygments_lexer": "ipython3",
|
||
|
|
"version": "3.10.2"
|
||
|
|
}
|
||
|
|
},
|
||
|
|
"nbformat": 4,
|
||
|
|
"nbformat_minor": 5
|
||
|
|
}
|