experiment/graphene emt

other/experiments/relax_graphene_with_emt.ipynb

@@ -0,0 +1,363 @@
+{
+ "cells": [
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "# Relax Graphene with EMT"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "import base64\n",
+    "from ase.io import write\n",
+    "from ase.build import make_supercell\n",
+    "from IPython.display import HTML\n",
+    "import io\n",
+    "\n",
+    "\n",
+    "def visualize_material_base64(material, title: str, rotation: str = '0x', number_of_repetitions: int = 1):\n",
+    "    \"\"\"\n",
+    "    Returns an HTML string with a Base64-encoded image for visualization,\n",
+    "    including the name of the file, positioned horizontally.\n",
+    "    \"\"\"\n",
+    "    # Set the number of unit cell repetition for the structure\n",
+    "    n = number_of_repetitions\n",
+    "    material_repeat = make_supercell(material, [[n, 0, 0], [0, n, 0], [0, 0, 1]])\n",
+    "    text = f\"{material.symbols} - {title}\"\n",
+    "\n",
+    "    # Write image to a buffer to display in HTML\n",
+    "    buf = io.BytesIO()\n",
+    "    write(buf, material_repeat, format='png', rotation=rotation)\n",
+    "    buf.seek(0)\n",
+    "    img_str = base64.b64encode(buf.read()).decode('utf-8')\n",
+    "    html_str = f'''\n",
+    "    <div style=\"display: inline-block; margin: 10px; vertical-align: top;\">\n",
+    "        <p>{text}</p>\n",
+    "        <img src=\"data:image/png;base64,{img_str}\" alt=\"{title}\" />\n",
+    "    </div>\n",
+    "    '''\n",
+    "    return html_str"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "## 1. Get original structure"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "\n",
+    "MATERIALS_FILES = [\n",
+    "    \"C_Si_32_copy1.json\"\n",
+    "]\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "import json\n",
+    "from src.utils import to_pymatgen\n",
+    "from src.utils import ase_to_poscar, pymatgen_to_ase\n",
+    "\n",
+    "materials_in = []\n",
+    "FOLDER = \"uploads\"\n",
+    "for file in MATERIALS_FILES:\n",
+    "    with open(f\"{FOLDER}/{file}\", \"r\") as f:\n",
+    "        data = f.read()\n",
+    "        materials_in.append(json.loads(data))\n",
+    "\n",
+    "if \"materials_in\" in globals():\n",
+    "    pymatgen_materials = [to_pymatgen(item) for item in materials_in]\n",
+    "\n",
+    "\n",
+    "input_material = pymatgen_materials[0]\n",
+    "ase_interface = pymatgen_to_ase(input_material)\n",
+    "\n",
+    "\n",
+    "html_input = visualize_material_base64(ase_interface, \"input\", \"0x\")\n",
+    "html_input_z = visualize_material_base64(ase_interface, \"input_z\", \"-90x\")\n",
+    "# Display the interfaces before and after relaxation\n",
+    "html_content = f'<div style=\"display: flex;\">{html_input}{html_input_z}</div>'\n",
+    "display(HTML(html_content))\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Add small random z displacements.\n",
+    "Add small random z displacements to allow for relaxation in z-direction, otherwise no force in z-direction can be calculated."
+   ]
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "import plotly.graph_objs as go\n",
+    "from IPython.display import display\n",
+    "from plotly.subplots import make_subplots\n",
+    "\n",
+    "\n",
+    "import numpy as np\n",
+    "\n",
+    "# Set amplitude for z perturbation in Ångströms\n",
+    "amplitude = 0.1  # adjust as needed, e.g., 0.01 to 0.1 Å\n",
+    "\n",
+    "positions = ase_interface.get_positions()\n",
+    "perturbed_positions = positions.copy()\n",
+    "\n",
+    "# Apply random perturbation only to z\n",
+    "perturbed_positions[:, 2] += np.random.normal(-amplitude, amplitude, size=positions.shape[0]) + 5\n",
+    "\n",
+    "ase_interface.set_positions(perturbed_positions)\n",
+    "\n",
+    "ase_original_interface = ase_interface.copy()\n",
+    "html_original = visualize_material_base64(ase_original_interface, \"Original material with z displacement\", \"-90x\")\n",
+    "# Display the interfaces before and after relaxation\n",
+    "html_content = f'<div style=\"display: flex;\">{html_original}</div>'\n",
+    "display(HTML(html_content))"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "from ase.optimize import BFGS, FIRE\n",
+    "from ase.calculators.emt import EMT\n",
+    "from ase.constraints import FixCartesian\n",
+    "\n",
+    "RELAXATION_PARAMETERS = {\n",
+    "    \"FMAX\": 0.005,\n",
+    "}\n",
+    "\n",
+    "\n",
+    "# Apply constraint to all atoms to fix x and y\n",
+    "constraints = [FixCartesian(i, [True, True, False]) for i in range(len(ase_interface))]\n",
+    "ase_interface.set_constraint(constraints)\n",
+    "\n",
+    "\n",
+    "# Set up the calculator\n",
+    "calculator = EMT()\n",
+    "ase_interface.set_calculator(calculator)\n",
+    "dyn = FIRE(ase_interface)\n",
+    "\n",
+    "# Initialize empty lists to store steps and energies\n",
+    "steps = []\n",
+    "energies = []\n",
+    "\n",
+    "# Create a plotly figure widget\n",
+    "fig = make_subplots(rows=1, cols=1, specs=[[{\"type\": \"scatter\"}]])\n",
+    "scatter = go.Scatter(x=[], y=[], mode='lines+markers', name='Energy')\n",
+    "fig.add_trace(scatter)\n",
+    "fig.update_layout(title_text='Real-time Optimization Progress', xaxis_title='Step', yaxis_title='Energy (eV)')\n",
+    "\n",
+    "# Display figure widget\n",
+    "f = go.FigureWidget(fig)\n",
+    "display(f)\n",
+    "\n",
+    "\n",
+    "# Define a callback function to update the plot at each step\n",
+    "def plotly_callback():\n",
+    "    step = dyn.nsteps\n",
+    "    energy = ase_interface.get_total_energy()\n",
+    "\n",
+    "    # Add the new step and energy to the lists\n",
+    "    steps.append(step)\n",
+    "    energies.append(energy)\n",
+    "\n",
+    "    print(f\"Step: {step}, Energy: {energy:.4f} eV\")\n",
+    "\n",
+    "    # Update the figure with the new data\n",
+    "    with f.batch_update():\n",
+    "        f.data[0].x = steps\n",
+    "        f.data[0].y = energies\n",
+    "\n",
+    "\n",
+    "# Run the relaxation\n",
+    "dyn.attach(plotly_callback, interval=1)\n",
+    "dyn.run(fmax=RELAXATION_PARAMETERS[\"FMAX\"], steps =50)\n",
+    "\n",
+    "# Extract results\n",
+    "# ase_original_interface = pymatgen_to_ase(interface)\n",
+    "# ase_original_interface.set_calculator(calculator)\n",
+    "ase_final_interface = ase_interface\n",
+    "\n",
+    "# original_energy = ase_original_interface.get_total_energy()\n",
+    "relaxed_energy = ase_interface.get_total_energy()\n",
+    "\n",
+    "# Print out the final relaxed structure and energy\n",
+    "# print('Original structure:\\n', ase_to_poscar(ase_original_interface))\n",
+    "# print('\\nRelaxed structure:\\n', ase_to_poscar(ase_final_interface))\n",
+    "# print(f\"The final energy is {float(relaxed_energy):.3f} eV.\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "### 3. View structure before and after relaxation"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "\n",
+    "html_original = visualize_material_base64(ase_original_interface, \"original\", \"0x\")\n",
+    "html_original_z = visualize_material_base64(ase_original_interface, \"original\", \"-90x\")\n",
+    "html_relaxed = visualize_material_base64(ase_final_interface, \"relaxed\", \"0x\")\n",
+    "html_relaxed_z = visualize_material_base64(ase_final_interface, \"relaxed\", \"-90x\")\n",
+    "\n",
+    "# Display the interfaces before and after relaxation\n",
+    "html_content = f'<div style=\"display: flex;\">{html_original}{html_relaxed}</div>'\n",
+    "display(HTML(html_content))"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "# Display the interfaces before and after relaxation\n",
+    "html_content = f'<div style=\"display: flex;\">{html_original_z}{html_relaxed_z}</div>'\n",
+    "display(HTML(html_content))"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "# Get atomic positions from both interfaces\n",
+    "original_positions = ase_original_interface.get_positions()\n",
+    "final_positions = ase_final_interface.get_positions()\n",
+    "\n",
+    "# Extract z-coordinates\n",
+    "original_z = original_positions[:, 2]\n",
+    "final_z = final_positions[:, 2]\n",
+    "\n",
+    "# Compute delta_z for each atom\n",
+    "delta_z = final_z - original_z  # shape: (N,)\n",
+    "print(delta_z)"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "import numpy as np\n",
+    "import plotly.graph_objs as go\n",
+    "from scipy.stats import binned_statistic_2d\n",
+    "\n",
+    "# Step 1: Get original and final positions\n",
+    "original_pos = ase_original_interface.get_positions()\n",
+    "final_pos = ase_final_interface.get_positions()\n",
+    "\n",
+    "# Step 2: Calculate Δz\n",
+    "delta_z = final_pos[:, 2] - original_pos[:, 2]\n",
+    "delta_z_abs = np.abs(delta_z)\n",
+    "\n",
+    "# Step 3: Define 2D grid binning over x and y\n",
+    "x = original_pos[:, 0]\n",
+    "y = original_pos[:, 1]\n",
+    "z = delta_z_abs\n",
+    "\n",
+    "# Adjust number of bins to suit resolution\n",
+    "n_bins_x = 150\n",
+    "n_bins_y = 150\n",
+    "\n",
+    "# Step 4: Bin Δz over 2D x/y using mean (you can also use 'max', 'sum', etc.)\n",
+    "stat, x_edges, y_edges, _ = binned_statistic_2d(\n",
+    "    x, y, z,\n",
+    "    statistic='mean',\n",
+    "    bins=[n_bins_x, n_bins_y]\n",
+    ")\n",
+    "\n",
+    "# Step 5: Create heatmap\n",
+    "heatmap = go.Heatmap(\n",
+    "    z=stat.T,  # Transpose to align x and y correctly\n",
+    "    x=(x_edges[:-1] + x_edges[1:]) / 2,\n",
+    "    y=(y_edges[:-1] + y_edges[1:]) / 2,\n",
+    "    colorscale='Thermal',\n",
+    "    colorbar=dict(title='Δz (Å)'),\n",
+    ")\n",
+    "\n",
+    "layout = go.Layout(\n",
+    "    title='Δz Displacement Heatmap over x–y',\n",
+    "    xaxis=dict(title='X (Å)'),\n",
+    "    yaxis=dict(title='Y (Å)'),\n",
+    "    height=800,\n",
+    "    width=800,\n",
+    ")\n",
+    "\n",
+    "fig = go.Figure(data=[heatmap], layout=layout)\n",
+    "fig.show()\n"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "from mat3ra.made.material import Material\n",
+    "from mat3ra.made.tools.convert import from_ase\n",
+    "mat3ra_material = Material.create(from_ase(ase_final_interface))\n",
+    "# visualize_materials([mat3ra_material], viewer=\"wave\")\n",
+    "\n",
+    "from utils.jupyterlite import set_materials\n",
+    "n_atoms = len(mat3ra_material.basis.elements.values)\n",
+    "mat3ra_material.name = f\"C{n_atoms}, Relaxed with EMT\"\n",
+    "set_materials(mat3ra_material)"
+   ],
+   "outputs": [],
+   "execution_count": null
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "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.12"
+  },
+  "widgets": {
+   "application/vnd.jupyter.widget-state+json": {
+    "state": {},
+    "version_major": 2,
+    "version_minor": 0
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}