constrained_SDE.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "114b1f21-621b-474e-a8ae-c0c0ea00aa14",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import h5py\n",
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import ase\n",
    "from tqdm import tqdm, trange\n",
    "from functools import partial\n",
    "import subprocess\n",
    "\n",
    "from e3nn import o3\n",
    "from e3nn.o3 import Irreps\n",
    "from e3_layers.data import Batch, CondensedDataset\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a47bfff-4c3b-4d30-a1e6-9dd15131df5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def relax(x, loss_fn):\n",
    "    step_size = 1e-3\n",
    "    while True:\n",
    "        x.requires_grad = True\n",
    "        x.grad = None\n",
    "        loss = loss_fn(x)\n",
    "        if loss < 1e-3*20 or loss.isnan():\n",
    "            break\n",
    "        loss.backward()\n",
    "        grad = x.grad\n",
    "        x.requires_grad = False\n",
    "        x -= grad*step_size\n",
    "    x.requires_grad = False\n",
    "    x.grad = None\n",
    "    return x\n",
    "    \n",
    "def loss_fn(x):\n",
    "    tmp = x[1:]-x[:-1]\n",
    "    tmp = (torch.linalg.norm(tmp, dim=-1) - 1)**2\n",
    "    return tmp.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cf8943e-de4d-417b-a0ba-86cbd848e4da",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 20\n",
    "x = torch.randn((N, 3))\n",
    "fig = plt.figure()\n",
    "ax0 = fig.add_subplot(211, projection='3d')\n",
    "ax1 = fig.add_subplot(212, projection='3d')\n",
    "ax0.plot3D(x[:,0], x[:,1], x[:,2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81233875-23a9-4af2-a03c-41dd46a0ad51",
   "metadata": {},
   "source": [
    "# Relax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cb2d789-e8ab-4e4f-a9d2-54c5fa4e040a",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = relax(x, loss_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8caec615-ddcb-4b48-b6a0-5b3d4177463e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax1.plot3D(x[:,0], x[:,1], x[:,2])\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79661eb1-3033-417e-979f-01a314aba1bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp = x[1:]-x[:-1]\n",
    "torch.linalg.norm(tmp, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75938556-bee1-4fb2-8b9f-5f402d701bde",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab6d14ab-b328-44cf-955a-6c2759860ba1",
   "metadata": {},
   "source": [
    "# Diffusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "391aac8b-6a2d-41a4-bc96-11788e2fd914",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_STEP = 1000\n",
    "for i in range(N_STEP):\n",
    "    x = x + torch.randn(x.shape)/N_STEP**0.5\n",
    "    if i % 100 == 0:\n",
    "        x = relax(x, loss_fn)"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}