{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3b91c190-c52c-4448-b237-4433115b6191",
   "metadata": {},
   "source": [
    "# PYMBSE Example\n",
    "\n",
    "* We construct a specific model\n",
    "  * Magnetic\n",
    "  * Mechanical\n",
    "  * Geometry\n",
    "  * 2D/3D\n",
    "* Each Model is defined by its dependencies and input files\n",
    "## Overview\n",
    "Each notebook consists of 3 stages\n",
    "* Dependency loading (inputs, dependent models)\n",
    "* Execution, processing\n",
    "* Storage of results\n",
    "\n",
    "Let's look at an example\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88c6c7e6-596f-4b0c-8012-02b92ea50bb4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pymbse\n",
    "from pymbse.query.api import PyMBSE\n",
    "from pymbse.microservice.api import init_upload_and_run, download_artefact, get_artefact_names\n",
    "\n",
    "from roxieapi.commons.roxie_constants import PlotLabels\n",
    "from roxieapi.commons.types import Plot2D, Plot3D\n",
    "from roxieapi.input.builder import RoxieInputBuilder\n",
    "from roxieapi.output.plots import RoxiePlotOutputs\n",
    "\n",
    "import scrapbook as sb\n",
    "\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3349482-4ddd-4588-af0d-52a4b4d8b7a2",
   "metadata": {},
   "source": [
    "## Loading of the model\n",
    "\n",
    "First, we load the model, define our inputs, outputs and temporary folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c61e260-a20e-4b0a-bf0b-eb4fd0784ffa",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# General settings\n",
    "config_file = './pymbse_example.yml'\n",
    "tmp_folder = Path('./tmp')\n",
    "input_folder = Path('./inputs')\n",
    "output_folder = Path('./outputs')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b29dfb53-c6a5-43dd-b1e8-97633824cc9b",
   "metadata": {},
   "source": [
    "* We have a model name defined in the yml.\n",
    "\n",
    "* We know, if we run roxie, we need a certain set of files, so we define a structure to hold them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d130ba72-82de-4a78-9d81-f90918833880",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Model specific settings\n",
    "\n",
    "model_name = 'pymbse_example'\n",
    "\n",
    "# Files required by a roxie run. This will collect inputs from artefacts and input files\n",
    "filelist = {\n",
    "    'data' : None,\n",
    "    'cadata' : None,\n",
    "    'bhdata' : None,\n",
    "    'hmo' : None,\n",
    "    'iron': None,\n",
    "}\n",
    "\n",
    "# A list of all input files (for later use)\n",
    "input_files = []"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62060ca2-38c5-49ca-806a-01a5cdc6566c",
   "metadata": {},
   "source": [
    "## Parameters\n",
    "\n",
    "A notebook cell marked with Parameters can store externally injected parameters as well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66af2078-1fbc-4e2e-824a-aa7bea32ec00",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "R_i = 28.0\n",
    "R_o = 43.9"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d34d245c-acb9-4051-be25-c53d57c80c2e",
   "metadata": {},
   "source": [
    "### Input file loading\n",
    "First we load all input files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58269931-7a52-4dc3-8191-8a07408dae4d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(f\"R_i is {R_i}\")\n",
    "print(f\"R_o = {R_o}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89e7815d-0694-42b9-9c33-7ee4a343e8de",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load values from config\n",
    "config = pymbse.query.config.pymbse_config.load_pymbse_config(config_file)\n",
    "model_config = config.models[model_name]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4409a86c-d484-4623-b40a-c151a868df81",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_config.input_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6020009e-63f2-487f-8500-e02e86cb6823",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for file in model_config.input_files:\n",
    "    suf = Path(file).suffix[1:]\n",
    "    if suf in filelist:\n",
    "        filelist[suf]=file\n",
    "    if suf != 'data':\n",
    "        input_files.append(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d797440-ec56-42e1-b838-bccef7604241",
   "metadata": {},
   "source": [
    "## Dependency execution\n",
    "\n",
    "Next step is to run dependend models. Iterates through all dependend models, executes them, and stores the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4f00fa3-bfec-4392-87be-ee5ab205dde4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_config.needs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "840b236a-7204-4667-986d-3076c15e5f0a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dependend_models = {}\n",
    "artefacts = {}\n",
    "for dm in model_config.needs:\n",
    "    dependend_models[dm] = PyMBSE(config_path=config_file,source_model=model_name,target_model=dm).build()\n",
    "    model_artefacts = dependend_models[dm].get_artefacts()\n",
    "    for ma in model_artefacts:\n",
    "        filename = ma\n",
    "        if ma in filelist:\n",
    "            filename = f\"{dm}.{ma}\"\n",
    "            filelist[ma] = filename\n",
    "        with open(tmp_folder / filename, 'w') as f:\n",
    "            f.write(model_artefacts[ma])\n",
    "        input_files.append(str(tmp_folder / filename))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fe477f0-3774-489e-9ee7-0497f687bb60",
   "metadata": {},
   "source": [
    "## Sanity check - Are all files present\n",
    "\n",
    "A quick check to see if all required input files are defined.\n",
    "\n",
    "The minimum requirements are a datafile + a cadata file.\n",
    "\n",
    "If there is a hmo or iron file linked, there also needs to be a bhdata file present"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da9b43b1-c6d6-476c-bded-38b0ded2f055",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert filelist['data'], \"Missing Datafile\"\n",
    "assert filelist['cadata'], \"Missing cadata file\"\n",
    "if filelist['iron'] or filelist['hmo']:\n",
    "    assert filelist['bhdata'], \"Iron mesh given, but no bhdata file\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03d8f6db-1ccf-47ff-a87d-e9a3a1d5fb74",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Prepare datafile for execution\n",
    "\n",
    "Generate a new datafile, and replace paths to dependend files\n",
    "\n",
    "Here, we could also add modifications on the datafile (adding blocks, removing flags, etc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ce7d14a-4427-403a-86b8-0f39e1288325",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "roxie_input = RoxieInputBuilder.from_datafile(filelist['data'])\n",
    "roxie_input.cadata_path = Path(filelist['cadata']).name\n",
    "if filelist['bhdata']:\n",
    "    roxie_input.bhdata_path = Path(filelist['bhdata']).name\n",
    "if filelist['hmo'] and roxie_input.flags['LHMO']:\n",
    "    roxie_input.iron_path = Path(filelist['hmo']).name\n",
    "if filelist['iron']:\n",
    "    roxie_input.iron_path = Path(filelist['iron']).name\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "148fd784-6f5c-44b1-aca9-5e7ab2d16625",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "roxie_input.block.loc[roxie_input.block['no'].isin([1,2]),'radius'] = R_o\n",
    "roxie_input.block.loc[roxie_input.block['no'].isin([3,4,5,6]),'radius'] = R_i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc830e2a-9a84-4274-945e-b1a2ac4f57c5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "    \n",
    "# Store in temp folder\n",
    "datafile_name = str(tmp_folder / (model_name + '.data'))\n",
    "tmp_folder.mkdir(exist_ok=True)\n",
    "roxie_input.build(datafile_name)\n",
    "input_files.append(datafile_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6b54b08-abea-4ded-b5e8-dbfdd7acea0a",
   "metadata": {},
   "source": [
    "## Roxie execution\n",
    "\n",
    "With dependencies resolved and input generated, Roxie can be called from the script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95439fce-ec51-4fe4-b8f9-ad72dc25a320",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "url_with_port = \"http://roxie:8080\"\n",
    "\n",
    "model_name, timestamp, output_lines, error_lines = init_upload_and_run(url_with_port, model_name, input_files)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7feccd7-60e3-4180-a402-6a71790d3215",
   "metadata": {},
   "source": [
    "# Retreive outputs\n",
    "\n",
    "From the results, we can download the artefacts into the tmp folder\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ba2eb79-bd8f-4aa3-981d-9c018b2a7091",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for ol in output_lines:\n",
    "    print(ol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6656e3b0-3039-42d0-8f72-16be2b8f455c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get a list of provided artefacts\n",
    "artefact_list = get_artefact_names(url_with_port,model_name,timestamp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "193b41d2-5581-458a-ad2f-ffa1aa15ee99",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(artefact_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0157c0e1-ba12-4f56-9567-c3f10c6a50dc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Download artefacts into temp folder (for now, only output and post.xml file\n",
    "for artefact in [f'{model_name}.output', f'{model_name}.post.xml']:\n",
    "    download_artefact(url_with_port, model_name, timestamp, tmp_folder, artefact)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38e402b5-24b0-4430-8dce-07639f0ff198",
   "metadata": {},
   "source": [
    "## Output plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "697a2efc-e03f-4afc-b924-5df069a9e786",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plots = RoxiePlotOutputs(\n",
    "    tmp_folder/f'{model_name}.post.xml',\n",
    "    tmp_folder/f'{model_name}.data',\n",
    ")\n",
    "\n",
    "plots.output_report()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98fe4015-95c3-4d7b-800a-4f85fb9257d9",
   "metadata": {},
   "source": [
    "## Storage of Results\n",
    "\n",
    "For this notebook to act as a model, we need to define its outputs and store them.\n",
    "The following lines of code create two dictionaries: FOM and artefacts.\n",
    "\n",
    "### fom\n",
    "\n",
    "The Figures of Merit (fom) are a list of Key-Value pairs for output values. This could be objectives, or results for specific purposes.\n",
    "\n",
    "More generally, the artefacts output is a list of files, which will be attached as a whole."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69bd7695-ec37-4425-81a9-7ced2e91dde0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "fom_dct = {}\n",
    "artefacts_dct = {'xml': str(tmp_folder/f'{model_name}.post.xml'), 'data': str(tmp_folder/f'{model_name}.data')}\n",
    "data = {'fom': fom_dct, 'artefacts': artefacts_dct}\n",
    "sb.glue('model_results', data=data, encoder='json')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
