generate_data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Data generation\n",
    "\n",
    "Given covariance function $cov(k(x, y))$, need generate random field with given covariance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import scipy.sparse as scsp\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.linalg as sclin\n",
    "import scipy.sparse.linalg as scspl\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cov_func(points, alpha=1, beta=1):\n",
    "    '''\n",
    "    Function generated covariance matrix for given points \n",
    "    and parameters alpha and beta\n",
    "    '''\n",
    "    n = points.shape[0]\n",
    "    A = np.zeros((n, n))\n",
    "    for i in range(n):\n",
    "        for j in range(n):\n",
    "            norm = np.linalg.norm(points[i] - points[j])\n",
    "            A[i, j] = np.exp(-norm**alpha / beta) \n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def col_covmatrix(points, alpha=1, beta=1):\n",
    "    '''\n",
    "    Function generates first column of the \n",
    "    covariance matrix to store it efficiently\n",
    "    '''\n",
    "    n = points.shape[0]\n",
    "    A = np.zeros((n,))\n",
    "    for i in range(n):\n",
    "        norm = np.sqrt((points[i, 0] - points[0, 0])**2 + (points[i, 1] - points[0, 1])**2)\n",
    "        A[i, ] = np.exp(-norm**alpha / beta)\n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def symtoeplitz2circulant(A):\n",
    "    '''\n",
    "    Function converts the first column of \n",
    "    symmetrix toeplitz matrix to the first\n",
    "    column of circulant matrix\n",
    "    '''\n",
    "    n = A.shape[0]\n",
    "    col = np.zeros((2*(n-1), ))\n",
    "    col[:n] = A[:n]\n",
    "    col[n:] = A[-2:-n:-1]\n",
    "    return col"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def coltoeplitz2colcircuant(A):\n",
    "    '''\n",
    "    Function converts first column of \n",
    "    BTTB to first column of BCCB\n",
    "    '''\n",
    "    n = A.shape[0]\n",
    "    m = int(round(np.sqrt(n)))\n",
    "    num_blocks = n // m\n",
    "    B = np.zeros(((2*(m-1))**2, ))\n",
    "    for i in range(num_blocks):\n",
    "        B[i*(2*(m-1)):(i+1)*(2*(m-1))] = symtoeplitz2circulant(A[i*num_blocks:(i+1)*num_blocks])\n",
    "    for j in range(num_blocks-2):\n",
    "        B[2*(m-1)*(num_blocks+j):2*(m-1)*(num_blocks+j+1)] = B[2*(m-1)*(num_blocks-2-j):2*(m-1)*(num_blocks-1-j)]\n",
    "    return B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create matrix and solve PDE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_matrix(covar, rhs):\n",
    "    '''\n",
    "    Function create matrix and rhs to solve discretized \n",
    "    general Poisson equation in the form div(k grad(u)) = rhs\n",
    "    '''\n",
    "    n = covar.shape[0]\n",
    "    A = scsp.lil_matrix((n**2, n**2))\n",
    "    b = np.zeros((n*n, ))\n",
    "    inner_points2ind = {}\n",
    "    bound_points2ind = {}\n",
    "    k = 0\n",
    "    for i in range(n):\n",
    "        for j in range(n):\n",
    "            if i == 0 or j == 0 or i == n-1 or j == n-1:\n",
    "                bound_points2ind[(i, j)] = k\n",
    "                k += 1\n",
    "            else:\n",
    "                inner_points2ind[(i, j)] = k\n",
    "                k += 1\n",
    "    ind2inner_points = {inner_points2ind[key]: key for key in inner_points2ind}\n",
    "    ind2bound_points = {bound_points2ind[key]: key for key in bound_points2ind}\n",
    "    for k in ind2inner_points:\n",
    "        i, j = ind2inner_points[k]\n",
    "        b[k] = rhs[i, j]\n",
    "        A[k, k] = -(covar[i, j+1] + covar[i-1, j] + covar[i+1, j] + covar[i, j-1] + 4*covar[i, j]) / 2\n",
    "        if j + 1 != n -1:\n",
    "            l1 = inner_points2ind[(i, j+1)]\n",
    "        else:\n",
    "            l1 = bound_points2ind[(i, j+1)]\n",
    "        A[k, l1] = (covar[i, j+1] + covar[i, j]) / 2\n",
    "        \n",
    "        if j - 1 != 0:\n",
    "            l2 = inner_points2ind[(i, j-1)]\n",
    "        else:\n",
    "            l2 = bound_points2ind[(i, j-1)]\n",
    "        A[k, l2] = (covar[i, j-1] + covar[i, j]) / 2\n",
    "        \n",
    "        if i + 1 != n-1:\n",
    "            k1 = inner_points2ind[(i+1, j)]\n",
    "        else:\n",
    "            k1 = bound_points2ind[(i+1, j)]\n",
    "        A[k, k1] = (covar[i+1, j] + covar[i, j]) / 2\n",
    "        \n",
    "        if i-1 != 0:\n",
    "            k2 = inner_points2ind[(i-1, j)]\n",
    "        else:\n",
    "            k2 = bound_points2ind[(i-1, j)]\n",
    "        A[k, k2] = (covar[i-1, j] + covar[i, j]) / 2\n",
    "    \n",
    "    for k in ind2bound_points:\n",
    "        A[k ,k] = 1\n",
    "    A = A * (n-1)**2\n",
    "    return A.tocsr(), b, ind2inner_points, ind2bound_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve(random_field, rhs, disp=False):\n",
    "    '''\n",
    "    Function solves the PDE div(k grad(u)) = rhs for \n",
    "    given k and rhs\n",
    "    '''\n",
    "    n = rhs.shape[0]\n",
    "    A, b, ind2inner, ind2bound = create_matrix(random_field, rhs)\n",
    "    x = scspl.spsolve(A, b)\n",
    "    u = np.zeros((n, n))\n",
    "    for k in ind2inner:\n",
    "        i, j = ind2inner[k]\n",
    "        u[i, j] = x[k]\n",
    "    for k in ind2bound:\n",
    "        i, j = ind2bound[k]\n",
    "        u[i, j] = x[k]\n",
    "    if disp:\n",
    "        plt.figure()\n",
    "        plt.imshow(u, cmap=\"jet\", interpolation=\"bicubic\")\n",
    "        plt.colorbar()\n",
    "        plt.figure()\n",
    "        plt.imshow(random_field, cmap=\"jet\", interpolation=\"bicubic\")\n",
    "        plt.colorbar()\n",
    "    return u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_random_field(col_A, n):\n",
    "    '''\n",
    "    Function creates random field for given first\n",
    "    column of the covariance matrix\n",
    "    '''\n",
    "    col_circ = coltoeplitz2colcircuant(col_A)\n",
    "    key_col = col_circ\n",
    "    key_mat = np.reshape(key_col, (2*(n-1), 2*(n-1)))\n",
    "    fft_key_mat = np.fft.fft2(key_mat)\n",
    "    Lam = np.sqrt(fft_key_mat)\n",
    "    x = np.random.randn(2*(n-1), 2*(n-1)) + 1j*np.random.randn(2*(n-1), 2*(n-1))\n",
    "    rand_field = np.fft.ifft2(Lam * (np.fft.fft2(x)))[:n, :n]\n",
    "    random_field = rand_field.imag\n",
    "    return random_field "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def integrate(u):\n",
    "    n = u.shape[0]\n",
    "    h = 1.0 / n\n",
    "    val = h**2 * np.sum(u[40:80, 40:80])\n",
    "    return val"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demo script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 32\n",
    "sigma = 1e-3\n",
    "size = 5000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  0%|          | 0/5000 [00:00<?, ?it/s]\u001b[A\n",
      "  0%|          | 1/5000 [00:00<09:58,  8.35it/s]\u001b[A\n",
      "100%|██████████| 5000/5000 [08:20<00:00, 10.00it/s]\n"
     ]
    }
   ],
   "source": [
    "x = np.linspace(0, 1, n)\n",
    "y = np.linspace(0, 1, n)\n",
    "xv, yv = np.meshgrid(x, y)\n",
    "p = np.hstack((xv.reshape((n*n, 1)), yv.reshape((n*n, 1))))\n",
    "disp = False\n",
    "X = np.zeros((size, n, n))\n",
    "y = np.zeros(size)\n",
    "for t in tqdm(range(size)):\n",
    "    col_A = col_covmatrix(p, 2, sigma)\n",
    "    rhs = np.ones((n, n))\n",
    "    random_field = create_random_field(col_A, n)\n",
    "    u = solve(random_field, rhs, disp)\n",
    "    val = integrate(u)\n",
    "    X[t, :, :] = u\n",
    "    y[t] = val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.savez(\"./data_n{}_sigma{}_size{}\".format(n, sigma, size), X=X, y=y)"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python (pytorch)",
   "language": "python",
   "name": "pytorch"
  },
  "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.5.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}