test_aiida_onsite_sweep.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "from aiida import load_profile\n",
    "from ase.io import read\n",
    "from tb_hamiltonian.workflows import sweep_onsite_parameters"
   ]
  },
  {
   "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": [
    "# Control parameters\n",
    "debug = False\n",
    "\n",
    "# Initial structure\n",
    "structure_label = \"BLG\"\n",
    "input_path = Path(\"examples/BLG/AB/rectangular\")\n",
    "initial_structure = read(input_path / \"POSCAR\", format=\"vasp\")\n",
    "\n",
    "# Supercell parameters\n",
    "nx = 1  # number of repetitions in x direction\n",
    "ny = 1  # number of repetitions in y direction\n",
    "nz = 1  # number of repetitions in z direction\n",
    "\n",
    "# Distances\n",
    "distances = [0.0, 1.425, 2.468, 2.850]\n",
    "\n",
    "# Hamiltonian parameters\n",
    "hopping_parameters = [0.0, -2.7, 0.0, -0.27]\n",
    "interlayer_coupling = 0.33\n",
    "nearest_neighbor = 1  # degree of nearest neighbors | don't use 0!\n",
    "\n",
    "# Potential parameters\n",
    "sweep_params = {\n",
    "    \"onsite_term\": [],\n",
    "    \"potential_type\": [\"null\", \"kronig-penney\", \"sine\"],\n",
    "    \"alpha\": [[1.0, 0.3]],\n",
    "    \"potential_params\": {\n",
    "        \"amplitude\": [0.25, 0.5, 0.75],\n",
    "        \"width\": [],\n",
    "        \"height\": [],\n",
    "    },\n",
    "}\n",
    "\n",
    "# Band structure parameters\n",
    "band_params = {\n",
    "    \"high_sym_points\": {\n",
    "        \"Γ\": (0.00000, 0.00000, 0.00000),\n",
    "        \"P\": (0.00000, 0.33333, 0.00000),\n",
    "        \"X\": (0.00000, 0.50000, 0.00000),\n",
    "        \"W\": (0.50000, 0.50000, 0.00000),\n",
    "        \"Y\": (0.50000, 0.00000, 0.00000),\n",
    "    },\n",
    "    \"path\": \"Γ P X W Y Γ W\",\n",
    "    \"points_per_segment\": 100,\n",
    "    \"use_sparse_solver\": False,\n",
    "    \"sparse_solver_params\": {\"k\": 3, \"sigma\": 1e-8},\n",
    "    \"save_data\": True,\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",
    "NODES = 1\n",
    "TASKS_PER_NODE = 1\n",
    "CPUS_PER_TASK = 1\n",
    "\n",
    "CONDA_ENV = \"\"\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\": NODES,\n",
    "            \"num_mpiprocs_per_machine\": TASKS_PER_NODE,\n",
    "            \"num_cores_per_mpiproc\": CPUS_PER_TASK,\n",
    "        },\n",
    "        \"max_wallclock_seconds\": h * 3600 + m * 60 + s,\n",
    "    }\n",
    "}\n",
    "\n",
    "REMOTE_ENV_PATH = \"\"\n",
    "assert REMOTE_ENV_PATH, \"Please set the REMOTE_ENV_PATH variable\"\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\": False,\n",
    "        \"resources\": {\n",
    "            \"num_machines\": NODES,\n",
    "            \"num_mpiprocs_per_machine\": TASKS_PER_NODE,\n",
    "            \"num_cores_per_mpiproc\": CPUS_PER_TASK,\n",
    "        },\n",
    "        \"max_wallclock_seconds\": h * 3600 + m * 60 + s,\n",
    "    }\n",
    "}\n",
    "\n",
    "input_metadata = {\n",
    "    \"build_hamiltonian\": {\n",
    "        \"computer\": \"localhost\",\n",
    "        \"metadata\": {\n",
    "            \"options\": {\n",
    "                **local_metadata[\"options\"],\n",
    "                \"additional_retrieve_list\": [\"POSCAR\"],\n",
    "            }\n",
    "        },\n",
    "    },\n",
    "    \"apply_onsite_term\": {\n",
    "        \"computer\": \"localhost\",\n",
    "        \"metadata\": local_metadata,\n",
    "    },\n",
    "    \"get_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 = sweep_onsite_parameters(\n",
    "    structure_label=structure_label,\n",
    "    initial_structure=initial_structure,\n",
    "    repetitions=[nx, ny, nz],\n",
    "    distances=distances,\n",
    "    nearest_neighbor=nearest_neighbor,\n",
    "    hopping_parameters=hopping_parameters,\n",
    "    interlayer_coupling=interlayer_coupling,\n",
    "    sweep_params=sweep_params,\n",
    "    band_params=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 import orm\n",
    "# from tb_hamiltonian.utils import BandStructure\n",
    "\n",
    "# band_calculations = orm.QueryBuilder().append(\n",
    "#     orm.WorkChainNode,\n",
    "#     filters=orm.WorkChainNode.fields.pk == wg.pk,\n",
    "#     tag=\"wg\",\n",
    "# ).append(\n",
    "#     orm.CalcJobNode,\n",
    "#     filters=orm.CalcJobNode.fields.label.like(\"get_band_structure%\"),\n",
    "#     with_incoming=\"wg\",\n",
    "# ).all()\n",
    "\n",
    "# for calc, in band_calculations:\n",
    "#     band_structure: BandStructure = calc.outputs.result.value\n",
    "#     label = calc.label.split(\"get_band_structure_\")[1]\n",
    "#     band_structure.plot(title=label)"
   ]
  }
 ],
 "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
}