{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Learning to Invert Heat Conduction with Scale-invariant Updates\n",
    "\n",
    "This example can be run with PyTorch and TensorFlow.\n",
    "Adjust the import below to select the library."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [],
   "source": [
    "from phi.torch.flow import *\n",
    "# from phi.tf.flow import *"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Generate batched 2D source `x` (non-observable) from 10 random rectangles"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [],
   "source": [
    "def generate_heat_example(*shape, bounds: Box = None):\n",
    "    shape = math.merge_shapes(*shape)\n",
    "    heat_t0 = CenteredGrid(0, extrapolation.PERIODIC, bounds, resolution=shape.spatial)\n",
    "    bounds = heat_t0.bounds\n",
    "    component_counts = math.to_int32(4 + 7 * math.random_uniform(shape.batch))\n",
    "    positions = (math.random_uniform(shape.batch, batch(components=10), channel(vector=shape.spatial.names)) - 0.5) * bounds.size * 0.8 + bounds.size * 0.5\n",
    "    for i in range(10):\n",
    "        position = positions.components[i]\n",
    "        half_size = math.random_uniform(shape.batch, channel(vector=shape.spatial.names)) * 10\n",
    "        strength = math.random_uniform(shape.batch) * math.to_float(i < component_counts)\n",
    "        position = math.clip(position, bounds.lower + half_size, bounds.upper - half_size)\n",
    "        component_box = Cuboid(position, half_size)\n",
    "        component_mask = SoftGeometryMask(component_box)\n",
    "        component_mask = component_mask.at(heat_t0)\n",
    "        heat_t0 += component_mask * strength\n",
    "    return heat_t0"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Look at two examples, generate data on-the-fly later during training."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 864x360 with 3 Axes>",
      "image/png": "[encoded data, 13464 characters]\n"
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": "<Figure size 432x288 with 0 Axes>"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "vis.plot(generate_heat_example(batch(view_examples=2), spatial(x=64, y=64)))"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Define scale-invariant updates.\n",
    "Take probabilistic viewpoint to avoid numerical explosions (see `probability_signal` in `apply_damping()`)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "outputs": [],
   "source": [
    "def apply_damping(kernel, inv_kernel, amp, f_uncertainty, log_kernel):\n",
    "    signal_prior = 0.5\n",
    "    expected_amp = 1. * kernel.shape.get_size('x') * inv_kernel  # This can be measured\n",
    "    signal_likelihood = math.exp(-0.5 * (abs(amp) / expected_amp) ** 2) * signal_prior  # this can be NaN\n",
    "    signal_likelihood = math.where(math.isfinite(signal_likelihood), signal_likelihood, math.zeros_like(signal_likelihood))\n",
    "    noise_likelihood = math.exp(-0.5 * (abs(amp) / f_uncertainty) ** 2) * (1 - signal_prior)\n",
    "    probability_signal = math.divide_no_nan(signal_likelihood, (signal_likelihood + noise_likelihood))\n",
    "    action = math.where((0.5 >= probability_signal) | (probability_signal >= 0.68), 2 * (probability_signal - 0.5), 0.)  # 1 sigma required to take action\n",
    "    prob_kernel = math.exp(log_kernel * action)\n",
    "    return prob_kernel, probability_signal\n",
    "\n",
    "\n",
    "def inv_diffuse(grid: Grid, amount: float, uncertainty: Grid):\n",
    "    f_uncertainty: math.Tensor = math.sqrt(math.sum(uncertainty.values ** 2, dim='x,y'))  # all frequencies have the same uncertainty, 1/N in iFFT\n",
    "    k_squared: math.Tensor = math.sum(math.fftfreq(grid.shape, grid.dx) ** 2, 'vector')\n",
    "    fft_laplace: math.Tensor = -(2 * np.pi) ** 2 * k_squared\n",
    "    # --- Compute sharpening kernel with damping ---\n",
    "    log_kernel = fft_laplace * -amount\n",
    "    log_kernel_clamped = math.minimum(log_kernel, math.to_float(math.floor(math.log(math.wrap(np.finfo(np.float32).max)))))  # avoid overflow\n",
    "    raw_kernel = math.exp(log_kernel_clamped)  # inverse diffusion FFT kernel, all values >= 1\n",
    "    inv_kernel = math.exp(-log_kernel)\n",
    "    amp = math.fft(grid.values)\n",
    "    kernel, sig_prob = apply_damping(raw_kernel, inv_kernel, amp, f_uncertainty, log_kernel)\n",
    "    # --- Apply and compute uncertainty ---\n",
    "    data = math.real(math.ifft(amp * math.to_complex(kernel)))\n",
    "    uncertainty = math.sqrt(math.sum(((f_uncertainty * kernel) ** 2))) / grid.shape.get_size('x')  # 1/N normalization in iFFT\n",
    "    uncertainty = grid * 0 + uncertainty\n",
    "    return grid.with_values(data), uncertainty, abs(amp), raw_kernel, kernel, sig_prob"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Set up neural network + Adam\n",
    "We will use 64 bit precision for physics but 32 for network.\n",
    "Batch size 128."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "math.set_global_precision(64)\n",
    "BATCH = batch(batch=128)\n",
    "\n",
    "math.seed(0)\n",
    "net = u_net(1, 1)\n",
    "optimizer = adam(net, 0.001)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Define loss function.\n",
    "Gradients will always computed as d(loss_function)/dθ (network weights).\n",
    "For SIP\n",
    "* Invert the physics\n",
    "* Define proxy loss L2(prediction - correction) where correction is taken as constant (based on `x = field.stop_gradient(prediction)`)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [],
   "source": [
    "# @math.jit_compile\n",
    "def loss_function(x_gt: CenteredGrid, sip: bool):\n",
    "    y_target = diffuse.fourier(x_gt, 8., 1)\n",
    "    with math.precision(32):\n",
    "        prediction = field.native_call(net, field.to_float(y_target)).vector[0]\n",
    "        prediction += field.mean(x_gt) - field.mean(prediction)\n",
    "    prediction = field.to_float(prediction)  # is this call necessary?\n",
    "    x = field.stop_gradient(prediction)\n",
    "    if sip:\n",
    "        y = diffuse.fourier(x, 8., 1)\n",
    "        dx, _, amp, raw_kernel, kernel, sig_prob = inv_diffuse(y_target - y, 8., uncertainty=abs(y_target - y) * 1e-6)\n",
    "        correction = x + dx\n",
    "        loss = field.l2_loss(prediction - correction)\n",
    "        y_l2 = field.l2_loss(y - y_target)\n",
    "    else:\n",
    "        y = diffuse.fourier(prediction, 8., 1)\n",
    "        y_l2 = loss = field.l2_loss(y - y_target)\n",
    "    return loss, x, y, field.l2_loss(x_gt - x), y_l2"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Training loop, generate data on-the-fly.\n",
    "(Switch between SIP and GD with the use_SIP checkbox or make two loops)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m10.493 ± 6.369\u001B[0m \u001B[37m(2e+00...3e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m10.187 ± 6.252\u001B[0m \u001B[37m(2e+00...3e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m11.047 ± 6.128\u001B[0m \u001B[37m(2e+00...3e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m10.072 ± 6.347\u001B[0m \u001B[37m(1e+00...4e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m8.682 ± 5.031\u001B[0m \u001B[37m(1e+00...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m8.241 ± 4.486\u001B[0m \u001B[37m(1e+00...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m8.466 ± 4.927\u001B[0m \u001B[37m(1e+00...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m9.183 ± 6.186\u001B[0m \u001B[37m(1e+00...5e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m8.269 ± 5.729\u001B[0m \u001B[37m(1e+00...4e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.814 ± 4.365\u001B[0m \u001B[37m(8e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m8.390 ± 5.162\u001B[0m \u001B[37m(1e+00...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.937 ± 4.595\u001B[0m \u001B[37m(2e+00...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.315 ± 4.118\u001B[0m \u001B[37m(1e+00...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.617 ± 4.776\u001B[0m \u001B[37m(1e+00...3e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.730 ± 4.983\u001B[0m \u001B[37m(7e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.126 ± 3.905\u001B[0m \u001B[37m(5e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.256 ± 6.348\u001B[0m \u001B[37m(8e-01...6e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.313 ± 3.860\u001B[0m \u001B[37m(8e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m8.344 ± 5.242\u001B[0m \u001B[37m(5e-01...3e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m6.463 ± 4.300\u001B[0m \u001B[37m(6e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.400 ± 4.349\u001B[0m \u001B[37m(9e-01...3e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.453 ± 4.910\u001B[0m \u001B[37m(5e-01...3e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.298 ± 5.256\u001B[0m \u001B[37m(1e+00...4e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.223 ± 4.199\u001B[0m \u001B[37m(7e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m6.610 ± 4.331\u001B[0m \u001B[37m(3e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m6.666 ± 3.827\u001B[0m \u001B[37m(3e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m6.979 ± 4.320\u001B[0m \u001B[37m(1e+00...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.171 ± 4.482\u001B[0m \u001B[37m(3e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m6.919 ± 3.900\u001B[0m \u001B[37m(1e+00...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m6.410 ± 4.150\u001B[0m \u001B[37m(6e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m6.273 ± 3.477\u001B[0m \u001B[37m(5e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.019 ± 4.924\u001B[0m \u001B[37m(9e-01...3e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.876 ± 4.400\u001B[0m \u001B[37m(2e+00...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m6.431 ± 4.248\u001B[0m \u001B[37m(5e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m6.557 ± 3.771\u001B[0m \u001B[37m(8e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m5.735 ± 3.723\u001B[0m \u001B[37m(5e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m7.207 ± 4.388\u001B[0m \u001B[37m(6e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m6.715 ± 3.744\u001B[0m \u001B[37m(1e+00...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m6.823 ± 4.558\u001B[0m \u001B[37m(8e-01...2e+01)\u001B[0m\n",
      "\u001B[92m(batchᵇ=128)\u001B[0m \u001B[94m5.899 ± 3.558\u001B[0m \u001B[37m(5e-01...2e+01)\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "for training_step in range(500):\n",
    "    data = generate_heat_example(spatial(x=64, y=64), BATCH)\n",
    "    loss_value, x, y, x_l2, y_l2 = update_weights(net, optimizer, loss_function, data, sip=True)\n",
    "    if training_step % 10 == 0: print(x_l2)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "outputs": [
    {
     "data": {
      "text/plain": "<Figure size 864x360 with 6 Axes>",
      "image/png": "[encoded data, 74484 characters]\n"
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": "<Figure size 432x288 with 0 Axes>"
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "test_data = generate_heat_example(batch(test_examples=2), spatial(x=64, y=64))\n",
    "_, test_pred, test_y, *_ = loss_function(test_data, True)\n",
    "vis.plot({\"Ground Truth\": test_data, \"Observation\": test_y, \"NN Prediction\": test_pred}, show_color_bar=False, row_dims='test_examples')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}