{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Managing Projects in `vipdopt`\n",
    "\n",
    "`vipdopt` includes a `Project` class for organizing all of your simulations,\n",
    "output data, and plots in one project directory. The `Project` class also enables\n",
    "saving the progress of an optimization and restarting from that checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "from pathlib import Path\n",
    "import sys  \n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "np.set_printoptions(threshold=100)\n",
    "\n",
    "# Get vipdopt directory path from Notebook\n",
    "parent_dir = str(Path().resolve().parents[2])\n",
    "\n",
    "# Add to sys.path\n",
    "sys.path.insert(0, parent_dir)\n",
    "\n",
    "# Imports from vipdopt\n",
    "from vipdopt.optimization import Device, Sigmoid, Scale, LumericalOptimization, BayerFilterFoM, AdamOptimizer, SuperFoM\n",
    "from vipdopt.simulation import LumericalSimulation, GaussianSource, DipoleSource, Power, Profile\n",
    "from vipdopt.configuration import SonyBayerConfig\n",
    "\n",
    "from vipdopt.project import Project, create_internal_folder_structure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you initialize a new `Project` it will be mostly empty. You'll need to set all of the\n",
    "parts (base simulation, optimization, optimizer, etc.) manually.\n",
    "\n",
    "Throughout this notebook, we will be using the \"Sony Color Router\" design that's appeared\n",
    "in the other tutorial notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating all of the parts first\n",
    "cfg = SonyBayerConfig()\n",
    "cfg.read_file('config_example_3d.yml')\n",
    "base_sim = LumericalSimulation('simulation_example.json')\n",
    "optimizer = AdamOptimizer(cfg['fixed_step_size'], (0.9, 0.999))\n",
    "\n",
    "coords = {\n",
    "    'x': np.linspace(0, 2.04e-6, 60),\n",
    "    'y': np.linspace(0, 2.04e-6, 60),\n",
    "    'z': np.linspace(0, 2.04e-6, 60),\n",
    "}\n",
    "lambda_vector = np.linspace(\n",
    "    cfg['lambda_min_um'],\n",
    "    cfg['lambda_max_um'],\n",
    "    cfg['num_bands'] * cfg['num_points_per_band']\n",
    ")\n",
    "n_freq = len(lambda_vector)\n",
    "\n",
    "device = Device(\n",
    "    (60, 60, 60),\n",
    "    (0, 1),\n",
    "    coords,\n",
    "    'color_router',\n",
    "    randomize=True,\n",
    "    init_seed=23,\n",
    "    filters=[Sigmoid(0.05, 0.1)]\n",
    ")\n",
    "\n",
    "foms = [\n",
    "    BayerFilterFoM(\n",
    "        'TE',\n",
    "        [GaussianSource(f'forward_src_{axis}')],\n",
    "        [GaussianSource(f'forward_src_{axis}'), DipoleSource(f'adjoint_src_{n}{axis}')],\n",
    "        [\n",
    "            Power(f'focal_monitor_{n}'),\n",
    "            Power(f'transmission_monitor_{n}'),\n",
    "            Profile('design_efield_monitor')\n",
    "        ],\n",
    "        [Profile('design_efield_monitor')],\n",
    "        range(n_freq),\n",
    "        [],\n",
    "        all_freqs=lambda_vector\n",
    "    ) for axis in 'xy' for n in range(4)\n",
    "]\n",
    "weights = [1] * len(foms)\n",
    "combined_fom = SuperFoM(((f,) for f in foms), weights)\n",
    "\n",
    "opt = LumericalOptimization(\n",
    "    base_sim,\n",
    "    device,\n",
    "    optimizer,\n",
    "    combined_fom,\n",
    "    epoch_list=np.linspace(\n",
    "        cfg['iter_per_epoch'],\n",
    "        cfg['iter_per_epoch'] * cfg['max_epochs'],\n",
    "        cfg['max_epochs'], dtype=int\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the parts are made we can make our `Project` manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "project = Project()\n",
    "project.config = cfg\n",
    "project.optimization = opt\n",
    "project.optimizer = optimizer\n",
    "project.base_sim = base_sim\n",
    "project.foms = foms\n",
    "project.weights = weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `Project` has a specific folder structure. This structure can be created \n",
    "automatically using `create_internal_structure`. Below is what the internal structure looks like:\n",
    "\n",
    "```bash\n",
    "├───.tmp\n",
    "├───data\n",
    "│   ├───checkpoints\n",
    "│   ├───opt_info\n",
    "│   │   └───plots\n",
    "│   └───saved_scripts\n",
    "├───device\n",
    "└───eval\n",
    "    ├───configs\n",
    "    └───utils\n",
    "```\n",
    "\n",
    "* `.tmp` stores intermediate files and completed simulation data\n",
    "* `data/checkpoints` contains checkpoints for optimiztion progress\n",
    "* `data/opt_info` contains information about the optimization\n",
    "* `data/opt_info/plots` contains plots generated during the optimization\n",
    "* `data/saved_scripts` contains scripts that are included as part of the optimization\n",
    "* `device` contains saved device parameters over each iteration\n",
    "* `eval` contains configurations and utility functions for evaluating the optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's use the directory \"color_router_proj\" as our project folder\n",
    "proj_dir = Path('color_router_proj')\n",
    "project.subdirectories = create_internal_folder_structure(proj_dir)\n",
    "\n",
    "project.save_as(proj_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we have a previously saved project folder, we can load our project into code for faster setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project.load_project(proj_dir)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}