test_aiida_continuum.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "from aiida import load_profile\n",
    "\n",
    "from tb_hamiltonian.continuum import create_Qn\n",
    "from tb_hamiltonian.workflows import compute_continuum_bands"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.tracebacklimit = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = load_profile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define parameters\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_mpi = False\n",
    "\n",
    "L = 500\n",
    "\n",
    "Qn = create_Qn(L)\n",
    "b1, b2 = Qn[0], Qn[1]\n",
    "\n",
    "model_params = dict(\n",
    "    bond_length=1.425,\n",
    "    interlayer_hopping=0.22,\n",
    "    superlattice_potential_periodicity=500,\n",
    "    superlattice_potential_amplitude=0.020,\n",
    "    gate_bias=0.024,\n",
    "    layer_potential_ratio=0.3,\n",
    "    nearest_neighbor_order=3,\n",
    ")\n",
    "\n",
    "band_params = dict(\n",
    "    high_sym_points={\n",
    "        \"Γ\": np.array([0.0, 0.0]),\n",
    "        \"M\": b1 / 2,\n",
    "        \"K\": (b1 + b2) / 3,\n",
    "    },\n",
    "    path=\"Γ K M Γ\",\n",
    "    Kpoint=\"K\",\n",
    "    total_points=100,\n",
    "    use_mpi=use_mpi,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Metadata\n",
    "\n",
    "For more scheduler options, see\n",
    "\n",
    "https://aiida.readthedocs.io/projects/aiida-core/en/latest/topics/calculations/usage.html#options\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h, m, s = 1, 0, 0\n",
    "\n",
    "CONDA_ENV = \"tbh\"\n",
    "assert CONDA_ENV, \"Please set the CONDA_ENV variable\"\n",
    "\n",
    "\n",
    "local_metadata = {\n",
    "    \"options\": {\n",
    "        \"prepend_text\": f\"\"\"\n",
    "source $(conda info --base)/etc/profile.d/conda.sh && conda activate {CONDA_ENV}\n",
    "        \"\"\",\n",
    "        \"withmpi\": False,\n",
    "        \"resources\": {\n",
    "            \"num_machines\": 1,\n",
    "            \"num_mpiprocs_per_machine\": 1,\n",
    "            \"num_cores_per_mpiproc\": 1,\n",
    "        },\n",
    "        \"max_wallclock_seconds\": h * 3600 + m * 60 + s,\n",
    "    }\n",
    "}\n",
    "\n",
    "REMOTE_ENV_PATH = \"\"\n",
    "# assert REMOTE_ENV_PATH, \\\n",
    "#     \"Please set the REMOTE_ENV_PATH variable or comment out the assertion if running locally\"\n",
    "\n",
    "remote_metadata = {\n",
    "    \"options\": {\n",
    "        \"prepend_text\": f\"\"\"\n",
    "module load cray-python\n",
    "source {REMOTE_ENV_PATH}/bin/activate\n",
    "        \"\"\",\n",
    "        \"withmpi\": use_mpi,\n",
    "        \"resources\": {\n",
    "            \"num_machines\": 1,\n",
    "            \"num_mpiprocs_per_machine\": 1,\n",
    "            \"num_cores_per_mpiproc\": 1,\n",
    "        },\n",
    "        \"max_wallclock_seconds\": h * 3600 + m * 60 + s,\n",
    "    }\n",
    "}\n",
    "\n",
    "input_metadata = {\n",
    "    \"get_continuum_model\": {\n",
    "        \"computer\": \"localhost\",\n",
    "        \"metadata\": local_metadata,\n",
    "    },\n",
    "    \"get_continuum_band_structure\": {\n",
    "        \"computer\": \"localhost\",\n",
    "        \"metadata\": local_metadata,\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build and submit work graph\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wg = compute_continuum_bands(\n",
    "    **model_params,\n",
    "    **band_params,\n",
    "    metadata=input_metadata,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wg.submit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Monitoring\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! verdi process status $wg.pk"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from aiida.orm import load_node\n",
    "from tb_hamiltonian.utils import BandStructure\n",
    "\n",
    "wg_node = load_node(wg.pk)\n",
    "band_calculation = wg_node.called[-1]\n",
    "band_structure: BandStructure = band_calculation.outputs.result.value\n",
    "band_structure.plot(\n",
    "    fig_params={\"ylim\": (-0.04, 0.04)},\n",
    "    plot_params={\"color\": \"k\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tbh",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}