data_gen.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "688d086c-fae5-4b03-b6d3-18f67371c91e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from phi.flow import *\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30538592-1467-43f4-b909-f15d627572c3",
   "metadata": {},
   "source": [
    "## Translation: Two Buoyant forces flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0ea6c0c-e180-470d-b7e8-5640bde65f14",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    os.mkdir(\"Translation\")\n",
    "except:\n",
    "    print(\"Translation folder already exists\")\n",
    "\n",
    "\n",
    "inflow_lst = [(i, j) for i in np.arange(8,64,8) for j in np.arange(5,30,5)]\n",
    "\n",
    "res_x = 64\n",
    "res_y = 64\n",
    "pressure = None\n",
    "DOMAIN = dict(x=res_x, y=res_y)\n",
    "DOMAIN_LEFT = dict(x=res_x, y=res_y)\n",
    "DOMAIN_RIGHT = dict(x=res_x, y=res_y)\n",
    "\n",
    "LEFT = StaggeredGrid(HardGeometryMask(Box[:res_x//2, :]), extrapolation.ZERO, **DOMAIN_LEFT)\n",
    "RIGHT = StaggeredGrid(HardGeometryMask(Box[res_x//2:, :]), extrapolation.ZERO, **DOMAIN_RIGHT)\n",
    "\n",
    "smoke = CenteredGrid(0, extrapolation.BOUNDARY, x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])  # sampled at cell centers\n",
    "velocity = StaggeredGrid(0, extrapolation.ZERO, x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])# sampled in staggered form at face centers\n",
    "\n",
    "INFLOW_LOCATION = tensor(inflow_lst, batch('inflow_loc'), channel('vector'))\n",
    "INFLOW = CenteredGrid(Sphere(center=INFLOW_LOCATION, radius=5), extrapolation.BOUNDARY, x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])\n",
    "\n",
    "data = []\n",
    "for _ in view(smoke, velocity).range(310):\n",
    "    smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW\n",
    "    buoyancy_force1 = smoke * LEFT * (0, 0.1) @ velocity#(velocity * LEFT)\n",
    "    buoyancy_force2 = smoke * RIGHT * (0, 0.2) @ velocity\n",
    "    velocity = advect.semi_lagrangian(velocity, velocity, dt=0.5) + buoyancy_force1 + buoyancy_force2\n",
    "    velocity, _ = fluid.make_incompressible(velocity, (), Solve('auto', 1e-5, 0, x0=pressure))\n",
    "    data.append(torch.from_numpy(velocity.staggered_tensor().numpy(['inflow_loc', 'vector', 'y', 'x'])).float().unsqueeze(1))\n",
    "    \n",
    "data = torch.cat(data, dim = 1)[:,:,:,:-1,:-1]\n",
    "for k in range(len(inflow_lst)):\n",
    "    torch.save(data[k].double().float(), \"Translation/raw_data_\" + str(inflow_lst[k][0]) + \"_\" + str(inflow_lst[k][1]) + \".pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a612c623-4b20-4d96-830d-2fbf0e47d3b1",
   "metadata": {},
   "source": [
    "## Rotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d27654a-89c4-40c0-a0fa-1bbd3ad25c97",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.mkdir(\"Rotation\")\n",
    "\n",
    "inflow_pos = [[(32,5), (0,0.01)], [(64-5,32), (-0.01, 0)], [(32,64-5),(0, -0.011)], [(5,32), (0.01, 0)]]\n",
    "\n",
    "k = 0\n",
    "for i in np.arange(9):\n",
    "    for j in range(4):\n",
    "        k += 1\n",
    "        res_x = 64\n",
    "        res_y = 64\n",
    "        pressure = None\n",
    "        DOMAIN = dict(x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])\n",
    "        smoke = CenteredGrid(0, extrapolation.BOUNDARY, x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])  # sampled at cell centers\n",
    "        velocity = StaggeredGrid(0, extrapolation.ZERO, x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])  # sampled in staggered form at face centers\n",
    "        INFLOW_LOCATION = tensor([inflow_pos[j][0]], batch('inflow_loc'), channel('vector'))\n",
    "        INFLOW = CenteredGrid(Sphere(center=INFLOW_LOCATION, radius=5), extrapolation.BOUNDARY, x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])\n",
    "        \n",
    "        data = []\n",
    "        for _ in view(smoke, velocity).range(310):\n",
    "            smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW\n",
    "            buoyancy_force = smoke * tuple(e*k for e in inflow_pos[j][1]) @ velocity\n",
    "            velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force\n",
    "            velocity, _ = fluid.make_incompressible(velocity, (), Solve('auto', 1e-5, 0, x0=pressure))\n",
    "            vel_tensor = torch.from_numpy(velocity.staggered_tensor().numpy(['inflow_loc', 'vector', 'y', 'x'])).float()\n",
    "            data.append(vel_tensor)\n",
    "        data = torch.cat(data, dim = 0)\n",
    "        torch.save(data.float()[1:,:,:64,:64], \"Rotation/raw_data_\"+ str(k) + \"_\" + str(j) + \".pt\")\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a472d56-5a83-467a-9cbe-52ab01527541",
   "metadata": {},
   "source": [
    "## Scaling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02785f41-8b70-4938-8d7b-e506f37f0f62",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.mkdir(\"Scale\")\n",
    "\n",
    "factors = np.linspace(1, 5, 41)\n",
    "for i in range(0, 38):\n",
    "    factor = factors[i]\n",
    "    res_x = int(64*factor)\n",
    "    res_y = int(64*factor)\n",
    "    pressure = None\n",
    "    DOMAIN = dict(x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])\n",
    "    smoke = CenteredGrid(0, extrapolation.BOUNDARY, **DOMAIN)  # sampled at cell centers\n",
    "    velocity = StaggeredGrid(0, extrapolation.ZERO, **DOMAIN)  # sampled in staggered form at face centers\n",
    "    INFLOW_LOCATION = tensor([(32*factor, 5*factor)], batch('inflow_loc'), channel('vector'))\n",
    "    INFLOW = CenteredGrid(Sphere(center=INFLOW_LOCATION, radius=5*factor), extrapolation.BOUNDARY, **DOMAIN)\n",
    "    \n",
    "    data = []\n",
    "    for _ in view(smoke, velocity).range(310):\n",
    "        smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW\n",
    "        buoyancy_force = smoke * (0, 0.1*factor) @ velocity\n",
    "        velocity = advect.semi_lagrangian(velocity, velocity, dt = 0.5) + buoyancy_force\n",
    "        velocity, _ = fluid.make_incompressible(velocity, (), Solve('auto', 1e-5, 0, x0=pressure))\n",
    "        vel_tensor = torch.from_numpy(velocity.staggered_tensor().numpy(['inflow_loc', 'vector', 'y', 'x'])).float()\n",
    "        #print(vel_tensor.shape)\n",
    "        data.append(vel_tensor)\n",
    "        \n",
    "    data = torch.cat(data, dim  = 0).transpose(0,1).unsqueeze(0)\n",
    "    data = F.interpolate(data, scale_factor = (1, 1/factor, 1/factor), mode=\"trilinear\", align_corners = True).squeeze(0).transpose(0,1)/factor\n",
    "    #print(data.shape)\n",
    "    torch.save(data.float()[1:,:,:64,:64], \"Scale/raw_data_\"+ str(i) + \".pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95faccfe-ff62-485d-812c-52ff493f5615",
   "metadata": {},
   "source": [
    "## Create datasets with various rotation equivariance errors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d55270bc-ce76-461f-84b2-719f6437b669",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.mkdir(\"equivariance_test\")\n",
    "\n",
    "for i in np.arange(50):\n",
    "    res_x = 64\n",
    "    res_y = 64\n",
    "    pressure = None\n",
    "    DOMAIN = dict(x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])\n",
    "    smoke = CenteredGrid(0, extrapolation.BOUNDARY, x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])  # sampled at cell centers\n",
    "    velocity = StaggeredGrid(0, extrapolation.ZERO, x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])  # sampled in staggered form at face centers\n",
    "    INFLOW_LOCATION = tensor([(32, 5)], batch('inflow_loc'), channel('vector'))\n",
    "    INFLOW = CenteredGrid(Sphere(center=INFLOW_LOCATION, radius=5), extrapolation.BOUNDARY, x=res_x, y=res_y, bounds=Box[0:res_x, 0:res_y])\n",
    "    data = []\n",
    "    for _ in view(smoke, velocity).range(61):\n",
    "        smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW\n",
    "        buoyancy_force = smoke * (0, 0.005*(i+1)) @ velocity\n",
    "        velocity = advect.semi_lagrangian(velocity, velocity, dt=0.5) + buoyancy_force\n",
    "        velocity, _ = fluid.make_incompressible(velocity, (), Solve('auto', 1e-5, 0, x0=pressure))\n",
    "        vel_tensor = torch.from_numpy(velocity.staggered_tensor().numpy(['inflow_loc', 'vector', 'y', 'x'])).float()\n",
    "\n",
    "        data.append(vel_tensor)\n",
    "    data = torch.cat(data, dim = 0)\n",
    "    torch.save(data.float()[10:60,:,:64,:64], \"equivariance_test/raw_data_\"+ str(i) + \".pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "11a1daec-24e9-47d3-802b-b4c18a262b1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rot_vector(inp, theta):\n",
    "    #inp shape: c x 2 x 64 x 64\n",
    "    theta = torch.tensor(theta).float()\n",
    "    rot_matrix = torch.tensor([[torch.cos(theta), -torch.sin(theta)], [torch.sin(theta), torch.cos(theta)]]).float()\n",
    "    out = torch.einsum(\"ab, bc... -> ac...\",(rot_matrix, inp.transpose(0,1))).transpose(0,1)\n",
    "    return out\n",
    "\n",
    "def get_rot_mat(theta):\n",
    "    theta = torch.tensor(theta).float()\n",
    "    return torch.tensor([[torch.cos(theta), -torch.sin(theta), 0],\n",
    "                         [torch.sin(theta), torch.cos(theta), 0]]).float()\n",
    "\n",
    "def rot_img(x, theta):\n",
    "    rot_mat = get_rot_mat(theta)[None, ...].float().repeat(x.shape[0],1,1)\n",
    "    grid = F.affine_grid(rot_mat, x.size()).float()\n",
    "    x = F.grid_sample(x, grid)\n",
    "    return x.float()\n",
    "\n",
    "def rot_field(x, theta):\n",
    "    x_rot = torch.cat([rot_img(rot_vector(x, theta)[:,:1],  theta), \n",
    "                       rot_img(rot_vector(x, theta)[:,-1:], theta)], dim = 1)\n",
    "    return x_rot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4f78c050-fe31-432b-a73c-97ab2ac51171",
   "metadata": {},
   "outputs": [],
   "source": [
    "direc = \"equivariance_test/raw_data_\"\n",
    "lst = [(1,1,1,1), (1,2,3,4), (1,3,5,7), (1,4,7,10), (1,5,9,13), (1,6,11,16), \n",
    "       (1,7,13,19), (1,8,15,22), (1,9,17,25), (1,10,19,28)]\n",
    "\n",
    "EE = []\n",
    "for i, item in enumerate(lst):\n",
    "    d1 = torch.load(direc + str(item[0]) + \".pt\")\n",
    "    d2 = torch.load(direc + str(item[1]) + \".pt\")\n",
    "    d3 = torch.load(direc + str(item[2]) + \".pt\")\n",
    "    d4 = torch.load(direc + str(item[3]) + \".pt\")\n",
    "    torch.save(d1, \"equivariance_test/E_\" + str(i) + \"/data_0.pt\")\n",
    "    torch.save(rot_field(d2, np.pi/2), \"equivariance_test/E_\" + str(i) + \"/data_1.pt\")\n",
    "    torch.save(rot_field(d3, np.pi), \"equivariance_test/E_\" + str(i) + \"/data_2.pt\")\n",
    "    torch.save(rot_field(d4, np.pi/2*3), \"equivariance_test/E_\" + str(i) + \"/data_3.pt\")\n",
    "    # Equivariance error of each dataset\n",
    "    EE.append(np.round(torch.mean((torch.abs(d1[10:-1]-d2[10:-1]) + torch.abs(d1[10:-1]-d3[10:-1]) + torch.abs(d1[10:-1]-d4[10:-1]))/3).numpy().item(),3))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch-1.13.1",
   "language": "python",
   "name": "pytorch-1.13.1"
  },
  "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}