diff --git a/other/materials_designer/Introduction.ipynb b/other/materials_designer/Introduction.ipynb index dcd076a5..a2dd38f9 100644 --- a/other/materials_designer/Introduction.ipynb +++ b/other/materials_designer/Introduction.ipynb @@ -97,6 +97,8 @@ "\n", "This notebook demonstrates a workflow for converting materials data from the [JARVIS](https://jarvis.nist.gov/) database into ESSE format for use with the Mat3ra.com platform.\n", "\n", + "#### [6.1.4. Optimize film position on interface](optimize_film_position.ipynb).\n", + "\n", "## 7. Read more\n", "\n", "### 7.1. Under the hood.\n", diff --git a/other/materials_designer/optimize_film_position.ipynb b/other/materials_designer/optimize_film_position.ipynb new file mode 100644 index 00000000..0136b9f1 --- /dev/null +++ b/other/materials_designer/optimize_film_position.ipynb @@ -0,0 +1,288 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "source": [ + "# Optimize Interface Film Position\n", + "\n", + "Find most optimal position of the film on the substrate interface.\n", + "\n", + "

Usage

\n", + "\n", + "1. Make sure to select Input Material from the list of available materials. Must be an interface.\n", + "1. Set notebook parameters in cell 1.2. below (or use the default values).\n", + "1. Click \"Run\" > \"Run All\" to run all cells.\n", + "1. Wait for the run to complete.\n", + "1. Scroll down to view results.\n", + "\n", + "## Notes\n", + "\n", + "- The optimization is performed on a 2D grid of x,y translations.\n", + "- Interface material must have atoms labeled \"0\" for the substrate and \"1\" for the film.\n", + "\n", + "\n", + "## 1. Prepare the Environment\n", + "### 1.1. Install Packages\n" + ], + "metadata": { + "collapsed": false + }, + "id": "4dc7b2ed495d66e0" + }, + { + "cell_type": "code", + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if sys.platform == \"emscripten\":\n", + " import micropip\n", + "\n", + " await micropip.install('mat3ra-api-examples', deps=False)\n", + " from utils.jupyterlite import install_packages\n", + "\n", + " await install_packages(\"\")\n" + ], + "metadata": { + "collapsed": false + }, + "id": "dd86bee2985f1b50", + "execution_count": null + }, + { + "cell_type": "markdown", + "source": [ + "### 1.2. Set optimization parameters\n" + ], + "metadata": { + "collapsed": false + }, + "id": "cca70ab27ef1d01d" + }, + { + "cell_type": "code", + "outputs": [], + "source": [ + "MATERIAL_INDEX = 0 # Index of the material to optimize\n", + "# Grid parameters\n", + "GRID_SIZE = (20, 20) # Resolution of the x-y grid\n", + "GRID_RANGE_X = (-0.5, 0.5) # Range to search in x direction\n", + "GRID_RANGE_Y = (-0.5, 0.5) # Range to search in y direction\n", + "USE_CARTESIAN = False # Whether to use Cartesian coordinates\n", + "\n", + "# Visualization parameters\n", + "SHOW_3D_LANDSCAPE = False # Whether to show 3D energy landscape\n", + "STRUCTURE_REPETITIONS = [3, 3, 1] # Repetitions for structure visualization\n" + ], + "metadata": { + "collapsed": false + }, + "id": "12878fd61f5a6b13", + "execution_count": null + }, + { + "cell_type": "markdown", + "source": [ + "## 3. Load Material\n", + "### 3.1. Make sure that loaded material is an interface material (atoms must have labels \"0\" for the substrate and \"1\" for the film)" + ], + "metadata": { + "collapsed": false + }, + "id": "463af646361cd982" + }, + { + "cell_type": "code", + "outputs": [], + "source": [ + "from utils.jupyterlite import get_materials\n", + "\n", + "materials = get_materials(globals())\n", + "interface_material = materials[MATERIAL_INDEX]\n" + ], + "metadata": { + "collapsed": false + }, + "id": "3d982a1ca641f0d8" + }, + { + "cell_type": "markdown", + "source": [ + "### 3.2. Visualize the Material\n" + ], + "metadata": { + "collapsed": false + }, + "id": "e920a6dd4906d8e8" + }, + { + "cell_type": "code", + "outputs": [], + "source": [ + "from utils.visualize import visualize_materials\n", + "\n", + "visualize_materials([interface_material], repetitions=STRUCTURE_REPETITIONS)\n", + "visualize_materials([interface_material], repetitions=STRUCTURE_REPETITIONS, rotation='-90x')" + ], + "metadata": { + "collapsed": false + }, + "id": "5f4afdb7ac0c865b" + }, + { + "cell_type": "markdown", + "source": [ + "### 3.3. Optimize Film Position" + ], + "metadata": { + "collapsed": false + }, + "id": "90255d774f62d1da" + }, + { + "cell_type": "code", + "outputs": [], + "source": [ + "from mat3ra.made.tools.build.interface import get_optimal_film_displacement\n", + "from mat3ra.made.tools.modify import interface_displace_part\n", + "from mat3ra.made.tools.calculate.calculators import InterfaceMaterialCalculator\n", + "from mat3ra.made.tools.optimize import evaluate_calculator_on_xy_grid\n", + "calculator = InterfaceMaterialCalculator()\n", + "\n", + "# Calculate energy landscape\n", + "xy_matrix, energy_matrix = evaluate_calculator_on_xy_grid(\n", + " material=interface_material,\n", + " calculator_function=calculator.get_energy,\n", + " modifier=interface_displace_part,\n", + " grid_size_xy=GRID_SIZE,\n", + " grid_range_x=GRID_RANGE_X,\n", + " grid_range_y=GRID_RANGE_Y,\n", + " use_cartesian_coordinates=USE_CARTESIAN\n", + ")\n", + "\n", + "# Find optimal position\n", + "optimal_displacement = get_optimal_film_displacement(\n", + " material=interface_material,\n", + " calculator=calculator,\n", + " grid_size_xy=GRID_SIZE,\n", + " grid_range_x=GRID_RANGE_X,\n", + " grid_range_y=GRID_RANGE_Y,\n", + " use_cartesian_coordinates=USE_CARTESIAN\n", + ")\n", + "\n", + "print(f\"\\nOptimal displacement vector: {optimal_displacement}\")\n" + ], + "metadata": { + "collapsed": false + }, + "id": "eb0b6e59c24dda4" + }, + { + "cell_type": "markdown", + "source": [ + "## 4. Visualize Results\n", + "### 4.1. Plot Energy Landscape" + ], + "metadata": { + "collapsed": false + }, + "id": "2945179d3729935d" + }, + { + "cell_type": "code", + "outputs": [], + "source": [ + "from utils.plot import plot_energy_heatmap, plot_energy_landscape\n", + "# Plot energy landscape\n", + "plot_energy_heatmap(xy_matrix, energy_matrix, optimal_position=optimal_displacement[:2])\n", + "\n", + "if SHOW_3D_LANDSCAPE:\n", + " plot_energy_landscape(xy_matrix, energy_matrix, optimal_position=optimal_displacement[:2])\n", + "\n", + "# Create optimized material\n", + "optimized_material = interface_displace_part(\n", + " interface_material,\n", + " displacement=optimal_displacement,\n", + " use_cartesian_coordinates=USE_CARTESIAN\n", + ")\n" + ], + "metadata": { + "collapsed": false + }, + "id": "41ac6b383001db6b", + "execution_count": null + }, + { + "cell_type": "markdown", + "source": [ + "### 4.1. Visualize Original and Optimized Materials" + ], + "metadata": { + "collapsed": false + }, + "id": "82a1af573c6ca0e9" + }, + { + "cell_type": "code", + "outputs": [], + "source": [ + "print(\"\\nVisualization of original and optimized materials:\")\n", + "visualize_materials([interface_material, optimized_material],\n", + " repetitions=STRUCTURE_REPETITIONS)\n", + "visualize_materials([interface_material, optimized_material],\n", + " repetitions=STRUCTURE_REPETITIONS,\n", + " rotation='-90x')\n" + ], + "metadata": { + "collapsed": false + }, + "id": "e7972543ae747b68", + "execution_count": null + }, + { + "cell_type": "markdown", + "source": [ + "## 5. Save Results\n" + ], + "metadata": { + "collapsed": false + }, + "id": "b4f6308e795e4f3c" + }, + { + "cell_type": "code", + "outputs": [], + "source": [ + "from utils.jupyterlite import set_materials\n", + "\n", + "set_materials(optimized_material)" + ], + "metadata": { + "collapsed": false + }, + "id": "c81ec652fbb64316", + "execution_count": null + } + ], + "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": 5 +} diff --git a/other/materials_designer/specific_examples/optimize_film_position_graphene_nickel_interface.ipynb b/other/materials_designer/specific_examples/optimize_film_position_graphene_nickel_interface.ipynb new file mode 100644 index 00000000..75a5766a --- /dev/null +++ b/other/materials_designer/specific_examples/optimize_film_position_graphene_nickel_interface.ipynb @@ -0,0 +1,412 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "567b1b812acb1718", + "metadata": { + "collapsed": false + }, + "source": [ + "# Optimize Interface Film Position in Gr/Ni(111) Interface\n", + "\n", + "## 0. Introduction\n", + "\n", + "This notebook demonstrates how to create Gr/Ni(111) interface and optimize the film position.\n", + "\n", + "Following the manuscript:\n", + "\n", + "> **Arjun Dahal, Matthias Batzill**\n", + "> \"Graphene–nickel interfaces: a review\"\n", + "> Nanoscale, 6(5), 2548. (2014)\n", + "> [DOI: 10.1039/c3nr05279f](https://doi.org/10.1039/c3nr05279f)\n", + " \n", + "Recreating interface and shifting the film to the most favorable energy position showed in the image below. Fig 1. b.\n", + "\n", + "\"Optimal\n", + "\n", + "## 1. Prepare the Environment\n", + "### 1.1. Install Packages\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6cc535057854a4d1", + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "import sys\n", + "\n", + "if sys.platform == \"emscripten\":\n", + " import micropip\n", + "\n", + " await micropip.install('mat3ra-api-examples', deps=False)\n", + " from utils.jupyterlite import install_packages\n", + "\n", + " await install_packages(\"specific_examples|create_interface_with_min_strain_zsl.ipynb\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "63d36e1a062ae33a", + "metadata": { + "collapsed": false + }, + "source": [ + "### 1.2. Set Parameters\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "feed5532d2d5bf8", + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "# Material selection\n", + "SUBSTRATE_NAME = \"Nickel\"\n", + "FILM_NAME = \"Graphene\"\n", + "\n", + "# Slab parameters\n", + "FILM_MILLER_INDICES = (0, 0, 1)\n", + "FILM_THICKNESS = 1 # in atomic layers\n", + "FILM_VACUUM = 0.0 # in angstroms\n", + "FILM_XY_SUPERCELL_MATRIX = [[1, 0], [0, 1]]\n", + "FILM_USE_ORTHOGONAL_Z = True\n", + "\n", + "SUBSTRATE_MILLER_INDICES = (1, 1, 1)\n", + "SUBSTRATE_THICKNESS = 4 # in atomic layers\n", + "SUBSTRATE_VACUUM = 0.0 # in angstroms\n", + "SUBSTRATE_XY_SUPERCELL_MATRIX = [[1, 0], [0, 1]]\n", + "SUBSTRATE_USE_ORTHOGONAL_Z = True\n", + "\n", + "# Interface parameters\n", + "MAX_AREA = 50 # in Angstrom^2\n", + "INTERFACE_DISTANCE = 2.58 # in Angstrom\n", + "INTERFACE_VACUUM = 20.0 # in Angstrom\n", + "\n", + "# Optimization parameters\n", + "GRID_SIZE = (20, 20) # Resolution of the x-y grid\n", + "GRID_RANGE_X = (-0.5, 0.5) # Range to search in x direction\n", + "GRID_RANGE_Y = (-0.5, 0.5) # Range to search in y direction\n", + "USE_CARTESIAN = False # Whether to use Cartesian coordinates\n", + "\n", + "# Visualization parameters\n", + "SHOW_3D_LANDSCAPE = False # Whether to show 3D energy landscape\n", + "STRUCTURE_REPETITIONS = [3, 3, 1] # Repetitions for structure visualization\n" + ] + }, + { + "cell_type": "markdown", + "id": "d4ee2a1c17a07fa7", + "metadata": { + "collapsed": false + }, + "source": [ + "## 2. Create Interface Material\n", + "### 2.1. Load Materials\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f1047ca2cc30b87a", + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "from utils.visualize import visualize_materials\n", + "from mat3ra.standata.materials import Materials\n", + "from mat3ra.made.material import Material\n", + "\n", + "substrate = Material(Materials.get_by_name_first_match(SUBSTRATE_NAME))\n", + "film = Material(Materials.get_by_name_first_match(FILM_NAME))\n", + "\n", + "# Preview materials\n", + "visualize_materials([substrate, film], repetitions=STRUCTURE_REPETITIONS, rotation=\"0x\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "14c4e4f9aa6721a", + "metadata": { + "collapsed": false + }, + "source": [ + "### 2.2. Create Slabs" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9ca8a2ad2b10c1f7", + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "from mat3ra.made.tools.build.slab import SlabConfiguration, PymatgenSlabGeneratorParameters, get_terminations, \\\n", + " create_slab\n", + "\n", + "# Configure slabs\n", + "film_slab_configuration = SlabConfiguration(\n", + " bulk=film,\n", + " miller_indices=FILM_MILLER_INDICES,\n", + " thickness=FILM_THICKNESS,\n", + " vacuum=FILM_VACUUM,\n", + " xy_supercell_matrix=FILM_XY_SUPERCELL_MATRIX,\n", + " use_orthogonal_z=FILM_USE_ORTHOGONAL_Z\n", + ")\n", + "\n", + "substrate_slab_configuration = SlabConfiguration(\n", + " bulk=substrate,\n", + " miller_indices=SUBSTRATE_MILLER_INDICES,\n", + " thickness=SUBSTRATE_THICKNESS,\n", + " vacuum=SUBSTRATE_VACUUM,\n", + " xy_supercell_matrix=SUBSTRATE_XY_SUPERCELL_MATRIX,\n", + " use_orthogonal_z=SUBSTRATE_USE_ORTHOGONAL_Z,\n", + ")\n", + "\n", + "# Get terminations\n", + "params = PymatgenSlabGeneratorParameters(symmetrize=False)\n", + "film_slab_terminations = get_terminations(film_slab_configuration, params)\n", + "substrate_slab_terminations = get_terminations(substrate_slab_configuration, params)\n", + "\n", + "# Create slabs\n", + "film_slabs = [create_slab(film_slab_configuration, termination) for termination in film_slab_terminations]\n", + "substrate_slabs = [create_slab(substrate_slab_configuration, termination, params) for termination in\n", + " substrate_slab_terminations]\n", + "\n", + "# Visualize slabs\n", + "visualize_materials([{\"material\": slab, \"title\": slab.metadata[\"build\"][\"termination\"]} for slab in film_slabs],\n", + " repetitions=STRUCTURE_REPETITIONS, rotation=\"-90x\")\n", + "visualize_materials([{\"material\": slab, \"title\": slab.metadata[\"build\"][\"termination\"]} for slab in substrate_slabs],\n", + " repetitions=STRUCTURE_REPETITIONS, rotation=\"-90x\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "61ba88b1f053c4fb", + "metadata": { + "collapsed": false + }, + "source": [ + "### 2.3. Create Interface\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1af1471d3c1ee856", + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "from mat3ra.made.tools.build.interface import InterfaceConfiguration, ZSLStrainMatchingParameters, \\\n", + " ZSLStrainMatchingInterfaceBuilderParameters, ZSLStrainMatchingInterfaceBuilder\n", + "from itertools import product\n", + "\n", + "# Get termination pairs\n", + "termination_pairs = list(product(film_slab_terminations, substrate_slab_terminations))\n", + "print(\"Termination Pairs (Film, Substrate):\")\n", + "for idx, termination_pair in enumerate(termination_pairs):\n", + " print(f\" {idx}: {termination_pair}\")\n", + "\n", + "# Create interface with first termination pair\n", + "termination_pair = termination_pairs[0]\n", + "film_termination, substrate_termination = termination_pair\n", + "\n", + "interface_configuration = InterfaceConfiguration(\n", + " film_configuration=film_slab_configuration,\n", + " substrate_configuration=substrate_slab_configuration,\n", + " film_termination=film_termination,\n", + " substrate_termination=substrate_termination,\n", + " distance_z=INTERFACE_DISTANCE,\n", + " vacuum=INTERFACE_VACUUM\n", + ")\n", + "\n", + "# Build interface using ZSL matching\n", + "zsl_params = ZSLStrainMatchingParameters(max_area=MAX_AREA)\n", + "builder = ZSLStrainMatchingInterfaceBuilder(\n", + " build_parameters=ZSLStrainMatchingInterfaceBuilderParameters(\n", + " strain_matching_parameters=zsl_params\n", + " )\n", + ")\n", + "\n", + "interfaces = builder.get_materials(configuration=interface_configuration)\n", + "interface_material = interfaces[0] # Select first interface\n", + "interface_material.name = f\"{FILM_NAME}_{SUBSTRATE_NAME}_interface\"\n", + "\n", + "# Visualize interface\n", + "visualize_materials([interface_material], repetitions=STRUCTURE_REPETITIONS)\n", + "visualize_materials([interface_material], repetitions=STRUCTURE_REPETITIONS, rotation=\"-90x\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "8b20dc13f57cfce7", + "metadata": { + "collapsed": false + }, + "source": [ + "## 3. Optimize Interface Position\n", + "### 3.1. Calculate Energy Landscape\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3dd3a069a69b4be1", + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "from mat3ra.made.tools.build.interface import get_optimal_film_displacement\n", + "from mat3ra.made.tools.modify import interface_displace_part\n", + "from mat3ra.made.tools.optimize import evaluate_calculator_on_xy_grid\n", + "from mat3ra.made.tools.calculate.calculators import InterfaceMaterialCalculator\n", + "\n", + "calculator = InterfaceMaterialCalculator()\n", + "\n", + "# Calculate energy landscape\n", + "xy_matrix, energy_matrix = evaluate_calculator_on_xy_grid(\n", + " material=interface_material,\n", + " calculator_function=calculator.get_energy,\n", + " modifier=interface_displace_part,\n", + " grid_size_xy=GRID_SIZE,\n", + " grid_range_x=GRID_RANGE_X,\n", + " grid_range_y=GRID_RANGE_Y,\n", + " use_cartesian_coordinates=USE_CARTESIAN\n", + ")\n", + "\n", + "# Find optimal position\n", + "optimal_displacement = get_optimal_film_displacement(\n", + " material=interface_material,\n", + " calculator=calculator,\n", + " grid_size_xy=GRID_SIZE,\n", + " grid_range_x=GRID_RANGE_X,\n", + " grid_range_y=GRID_RANGE_Y,\n", + " use_cartesian_coordinates=USE_CARTESIAN\n", + ")\n", + "print(f\"\\nOptimal displacement vector: {optimal_displacement}\")\n" + ] + }, + { + "cell_type": "markdown", + "id": "4028cb9ac8d438f0", + "metadata": { + "collapsed": false + }, + "source": [ + "### 3.2. Visualize Energy Landscape\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "27d655253dd79d2d", + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "from utils.plot import plot_energy_heatmap, plot_energy_landscape\n", + "\n", + "# Plot energy landscape\n", + "plot_energy_heatmap(xy_matrix, energy_matrix, optimal_position=optimal_displacement[:2])\n", + "if SHOW_3D_LANDSCAPE:\n", + " plot_energy_landscape(xy_matrix, energy_matrix, optimal_position=optimal_displacement[:2])\n", + "\n", + "# Create and visualize optimized material\n", + "optimized_material = interface_displace_part(\n", + " interface_material,\n", + " displacement=optimal_displacement,\n", + " use_cartesian_coordinates=USE_CARTESIAN\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "262edee47baa8114", + "metadata": { + "collapsed": false + }, + "source": [ + "### 3.3. Visualize Material\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "489288e70deb7d73", + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "print(\"\\nVisualization of original and optimized materials:\")\n", + "visualize_materials([{\"material\": interface_material, \"title\": \"Original Interface\"},\n", + " {\"material\": optimized_material, \"title\": \"Optimized Interface\"}],\n", + " repetitions=STRUCTURE_REPETITIONS,\n", + " )\n", + "visualize_materials([{\"material\": interface_material, \"title\": \"Original Interface\"},\n", + " {\"material\": optimized_material, \"title\": \"Optimized Interface\"}],\n", + " repetitions=STRUCTURE_REPETITIONS,\n", + " rotation=\"-90x\")" + ] + }, + { + "cell_type": "markdown", + "id": "78304054eca1564e", + "metadata": { + "collapsed": false + }, + "source": [ + "# 4. Save optimized material" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb3a1781c32f144f", + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "from utils.jupyterlite import download_content_to_file\n", + "\n", + "# Save optimized material\n", + "download_content_to_file(optimized_material, f\"{interface_material.name}_optimized_xy.json\")" + ] + } + ], + "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": 5 +} diff --git a/utils/plot.py b/utils/plot.py index 331dcded..77d82005 100644 --- a/utils/plot.py +++ b/utils/plot.py @@ -1,6 +1,7 @@ from typing import Dict, List, Union import matplotlib.pyplot as plt +import numpy as np import plotly.graph_objs as go from ase.atoms import Atoms as ASEAtoms from ase.optimize import BFGS, FIRE @@ -151,3 +152,77 @@ def plot_rdf(material: Material, cutoff: float = 10.0, bin_size: float = 0.1): plt.legend() plt.grid() plt.show() + + +def plot_energy_landscape(xy_matrix, energy_matrix, optimal_position=None): + """ + Create a 3D surface plot of the energy landscape. + + Args: + xy_matrix (List[np.ndarray]): X and Y coordinate matrices + energy_matrix (np.ndarray): Matrix of energy values + optimal_position (tuple, optional): The optimal (x,y) position to highlight + """ + x_vals, y_vals = xy_matrix + X, Y = np.meshgrid(x_vals, y_vals) + + # Create the 3D surface plot + fig = go.Figure(data=[go.Surface(x=X, y=Y, z=energy_matrix, colorscale="Viridis")]) + + # Add optimal position marker if provided + if optimal_position is not None: + x_opt, y_opt = optimal_position[0], optimal_position[1] + z_opt = np.min(energy_matrix) + fig.add_trace( + go.Scatter3d( + x=[x_opt], + y=[y_opt], + z=[z_opt], + mode="markers", + marker=dict(size=8, color="red"), + name="Optimal Position", + ) + ) + + fig.update_layout( + title="Interface Energy Landscape", + scene=dict(xaxis_title="X Position", yaxis_title="Y Position", zaxis_title="Energy"), + width=800, + height=800, + ) + + fig.show() + + +def plot_energy_heatmap(xy_matrix, energy_matrix, optimal_position=None): + """ + Create a 2D heatmap of the energy landscape. + + Args: + xy_matrix (List[np.ndarray]): X and Y coordinate matrices + energy_matrix (np.ndarray): Matrix of energy values + optimal_position (tuple, optional): The optimal (x,y) position to highlight + """ + x_vals, y_vals = xy_matrix + + fig = go.Figure( + data=go.Heatmap(x=x_vals, y=y_vals, z=energy_matrix, colorscale="Viridis", colorbar=dict(title="Energy")) + ) + + if optimal_position is not None: + x_opt, y_opt = optimal_position[0], optimal_position[1] + fig.add_trace( + go.Scatter( + x=[x_opt], + y=[y_opt], + mode="markers", + marker=dict(size=12, color="red", symbol="x"), + name="Optimal Position", + ) + ) + + fig.update_layout( + title="Interface Energy Heatmap", xaxis_title="X Position", yaxis_title="Y Position", width=800, height=600 + ) + + fig.show()