test_aiida_sweep_cell_size.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_cell_sizes"
   ]
  },
  {
   "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",
    "# 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",
    "alpha = [1.0, 0.5]\n",
    "onsite_term = 0.0\n",
    "potential_type = \"null\"\n",
    "potential_params = {\n",
    "    \"amplitude\": 1.0,\n",
    "    # \"width\": 0.5,\n",
    "    # \"height\": 0.5,\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",
    "}\n",
    "\n",
    "sizes = [1, 2, (2, 3)]  # (nx, ny), or n (nx = ny = 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_cell_sizes(\n",
    "    structure_label=structure_label,\n",
    "    initial_structure=initial_structure,\n",
    "    distances=distances,\n",
    "    nearest_neighbor=nearest_neighbor,\n",
    "    hopping_parameters=hopping_parameters,\n",
    "    interlayer_coupling=interlayer_coupling,\n",
    "    potential_type=potential_type,\n",
    "    potential_params=potential_params,\n",
    "    onsite_term=onsite_term,\n",
    "    alpha=alpha,\n",
    "    band_params=band_params,\n",
    "    sizes=sizes,\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",
    "\n",
    "# wg_node = load_node(wg.pk)\n",
    "# for child in wg_node.called:\n",
    "#     band_calculation = child.called[-1]\n",
    "#     band_structure = band_calculation.outputs.result.value\n",
    "#     ax = band_structure.plot(title=child.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
}