create_bash.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "aeaf34a0-d107-486d-bf67-5d2dc7ffee14",
   "metadata": {},
   "source": [
    "Create bash scripts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2f145c4a-8458-4ca9-a8e0-e7684ba65f59",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/wi73yus/.conda/envs/lbl/lib/python3.9/site-packages/pyfftw/interfaces/__init__.py:268: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  has_scipy_fft = _LooseVersion(scipy.__version__) >= _LooseVersion('1.4.0')\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from copy import deepcopy\n",
    "\n",
    "from light_by_light.vacem_ini import W_to_E0, create_ini_file\n",
    "from light_by_light.utils import write_yaml, read_yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "07abb96d-480a-44d2-a21e-ba6578c5616a",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = '''#!/bin/bash\n",
    "#SBATCH --job-name={jobname}\n",
    "#SBATCH --partition=gpu\n",
    "#SBATCH --cpus-per-task={n_cpus}\n",
    "#SBATCH --time=10:00:00\n",
    "#SBATCH --mem=150GB\n",
    "#SBATCH --output=/work/wi73yus/bash_output/{jobname}.%j.log\n",
    "#SBATCH --error=/work/wi73yus/bash_output/{jobname}.%j.err\n",
    "\n",
    "source ~/.bashrc  #  not automatically read by slurm\n",
    "conda activate lbl  # activate here the anaconda environment you want to use\n",
    "\n",
    "python {script_path} {script_kwargs}\n",
    "'''\n",
    "\n",
    "home_path = '/home/wi73yus/light-by-light'\n",
    "work_path = '/work/wi73yus/light-by-light-experiments'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc115665-294b-46ea-9403-6102bee6c76c",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Bash single simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06ef61a2-ef4c-4215-a20f-b82a7005238a",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### ECS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "21087704-e969-4066-b874-7d063d5e0536",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 160\n",
    "beta = 90\n",
    "\n",
    "# 1st laser\n",
    "laser_params = []\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0x': 5.5*800e-9,\n",
    "    'w0y': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': 0,\n",
    "    'phi': 0,\n",
    "    'beta': 0,\n",
    "    'solution': 'gaussian_elliptical_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 7,\n",
    "    'trans': 15,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 2,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'simulation'\n",
    "geometry = 'z'\n",
    "low_memory_mode = False\n",
    "n_threads = 24\n",
    "pol_idx = 0\n",
    "\n",
    "discernible_spectral = False\n",
    "sphmap_params = {\n",
    "    'order': 1,\n",
    "    'angular_resolution': 0.5/180*np.pi\n",
    "}\n",
    "\n",
    "# Save yaml file\n",
    "data = {\n",
    "    '_mode': mode,\n",
    "    'lasers': laser_params,\n",
    "    'simbox_params': simbox_params,\n",
    "    'geometry': geometry,\n",
    "    'low_memory_mode': low_memory_mode,\n",
    "    'n_threads': n_threads,\n",
    "    'pol_idx': pol_idx,\n",
    "    'discernible_spectral': discernible_spectral,\n",
    "    'sphmap_params': sphmap_params,\n",
    "}\n",
    "\n",
    "folder = f'ECS_W'\n",
    "jobname = f'theta_{theta}_w0x_5.5_w0y_1'\n",
    "save_path = f'{work_path}/runs/{mode}/{folder}/{jobname}/'\n",
    "Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "default_yaml = f'{save_path}default_params.yml'\n",
    "write_yaml(default_yaml, data)\n",
    "\n",
    "# Create bash script\n",
    "script_path = f'{home_path}/cluster/simulation_cluster.py'\n",
    "script_kwargs = f'--save_path={save_path}'\n",
    "\n",
    "bash = template.format(jobname=jobname,\n",
    "                       n_cpus=n_threads,\n",
    "                       script_path=script_path,\n",
    "                       script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2f93ab85-9ee7-4093-93f9-8f023ae674bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_name = 'default_params.yml'\n",
    "\n",
    "default_yaml = f'{save_path}{default_name}'\n",
    "yaml_data = read_yaml(default_yaml)\n",
    "params = {}\n",
    "lasers = yaml_data['lasers']\n",
    "params['laser_params'] = [lasers[f'laser_{i}'] for i in range(len(lasers))]\n",
    "params['save_path'] = save_path\n",
    "keys = ['simbox_params', 'geometry', 'low_memory_mode']\n",
    "for key in keys:\n",
    "    params[key] = yaml_data[key]\n",
    "\n",
    "create_ini_file(**params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "576e1085-9d69-4b63-b7e0-ca330c5e60ae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 924401\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba21fda1-e291-477d-86c1-e460ea8a51e0",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Harmonics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "9300268e-6949-450f-9c0f-8e2263b48d0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 180\n",
    "beta = 90\n",
    "\n",
    "Wtotal = 25\n",
    "H = 1 + 1/2**2.5\n",
    "W = [Wtotal/(n**2.5*H) for n in [1,2]]\n",
    "# print(W)\n",
    "# print(np.sum(W))\n",
    "\n",
    "W = [25, 25]\n",
    "\n",
    "# 1st laser\n",
    "lam0 = 800e-9\n",
    "k0 = 2*np.pi/lam0\n",
    "order = 0\n",
    "gauss_type = 'gaussian_spectral'\n",
    "\n",
    "w0_harm = 1\n",
    "w0_usual = 1\n",
    "phi0_harm = [0, 0]\n",
    "\n",
    "laser_params = []\n",
    "for i in range(1,len(W)+1):\n",
    "    laser = {\n",
    "        'tau': 20e-15,\n",
    "        'lam': lam0/i,\n",
    "        'w0': w0_harm*800e-9/i,\n",
    "        'W': W[i-1],\n",
    "        'theta': 0,\n",
    "        'phi': 0,\n",
    "        'beta': 0,\n",
    "        'solution': gauss_type,\n",
    "        'order': order,\n",
    "        'phi0': phi0_harm[i-1],\n",
    "    }\n",
    "    laser['E0'] = float(W_to_E0(laser))\n",
    "    laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 20e-15,\n",
    "    'lam': lam0,\n",
    "    'w0': w0_usual*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': gauss_type,\n",
    "    'order': order,\n",
    "    'phi0': 0\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 8,\n",
    "    'trans': 18,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 3,\n",
    "    't': 2\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'simulation'\n",
    "geometry = 'z'\n",
    "low_memory_mode = False\n",
    "n_threads = 24\n",
    "pol_idx = 0\n",
    "\n",
    "discernible_spectral = True\n",
    "sphmap_params = {\n",
    "    'order': 1,\n",
    "    'angular_resolution': 0.5/180*np.pi\n",
    "}\n",
    "\n",
    "sph_limits = {\n",
    "    'k': [1.5*k0, 2.5*k0],\n",
    "    'theta': [np.pi/2, np.pi]\n",
    "}\n",
    "\n",
    "# Save yaml file\n",
    "data = {\n",
    "    '_mode': mode,\n",
    "    'lasers': laser_params,\n",
    "    'simbox_params': simbox_params,\n",
    "    'geometry': geometry,\n",
    "    'low_memory_mode': low_memory_mode,\n",
    "    'n_threads': n_threads,\n",
    "    'pol_idx': pol_idx,\n",
    "    'discernible_spectral': discernible_spectral,\n",
    "    'sphmap_params': sphmap_params,\n",
    "    'sph_limits': sph_limits\n",
    "}\n",
    "\n",
    "folder = f'harmonics_W'\n",
    "if type(w0_harm) == int:\n",
    "    w0_str = str(w0_harm)\n",
    "else:\n",
    "    w0_str = '_'.join(str(w0) for w0 in w0_harm)\n",
    "jobname = f'21-1omega_theta_{theta}_w0harm_{w0_str}_figure'\n",
    "save_path = f'{work_path}/runs/{mode}/{folder}/{jobname}/'\n",
    "Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "default_yaml = f'{save_path}default_params.yml'\n",
    "write_yaml(default_yaml, data)\n",
    "\n",
    "# Create bash script\n",
    "script_path = f'{home_path}/cluster/simulation_cluster.py'\n",
    "script_kwargs = f'--save_path={save_path}'\n",
    "\n",
    "bash = template.format(jobname=jobname,\n",
    "                       n_cpus=n_threads,\n",
    "                       script_path=script_path,\n",
    "                       script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "5e1aeb95-a9c5-4a4f-99d0-535d23b7f4c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_name = 'default_params.yml'\n",
    "\n",
    "default_yaml = f'{save_path}{default_name}'\n",
    "yaml_data = read_yaml(default_yaml)\n",
    "params = {}\n",
    "lasers = yaml_data['lasers']\n",
    "params['laser_params'] = [lasers[f'laser_{i}'] for i in range(len(lasers))]\n",
    "params['save_path'] = save_path\n",
    "keys = ['simbox_params', 'geometry', 'low_memory_mode']\n",
    "for key in keys:\n",
    "    params[key] = yaml_data[key]\n",
    "\n",
    "create_ini_file(**params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "04125fec-c2be-4007-858d-7816117b5dbb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 923820\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cfe4f84-00dd-49e5-a980-eed0f250c236",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "#### 3 pulses collision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "29371640-fc8c-4936-912d-9ef5c28e42f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 180\n",
    "beta = 90\n",
    "\n",
    "# 1st laser\n",
    "laser_params = []\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': 3*800e-9,\n",
    "    'W': 50,\n",
    "    'theta': 0,\n",
    "    'phi': 0,\n",
    "    'beta': 0,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "# # 2nd laser\n",
    "# laser = {\n",
    "#     'tau': 25e-15,\n",
    "#     'lam': 800e-9,\n",
    "#     'w0': 3*800e-9,\n",
    "#     'W': 25,\n",
    "#     'theta': 0,\n",
    "#     'phi': 0,\n",
    "#     'beta': 0,\n",
    "#     'solution': 'gaussian_paraxial'\n",
    "# }\n",
    "# laser['E0'] = float(W_to_E0(laser))\n",
    "# laser_params.append(laser)\n",
    "\n",
    "# 3rd laser\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 7,\n",
    "    'trans': 15,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 2,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions,\n",
    "                 'pick_largest_size': True}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'simulation'\n",
    "geometry = 'xz'\n",
    "low_memory_mode = False\n",
    "n_threads = 32\n",
    "pol_idx = 0\n",
    "\n",
    "discernible_spectral = False\n",
    "sphmap_params = {\n",
    "    'order': 1,\n",
    "    'angular_resolution': 0.5/180*np.pi\n",
    "}\n",
    "\n",
    "# Save yaml file\n",
    "data = {\n",
    "    '_mode': mode,\n",
    "    'lasers': laser_params,\n",
    "    'simbox_params': simbox_params,\n",
    "    'geometry': geometry,\n",
    "    'low_memory_mode': low_memory_mode,\n",
    "    'n_threads': n_threads,\n",
    "    'pol_idx': pol_idx,\n",
    "    'discernible_spectral': discernible_spectral,\n",
    "    'sphmap_params': sphmap_params,\n",
    "}\n",
    "\n",
    "folder = f'3_pulses'\n",
    "jobname = f'w0_3_counter_W_50'\n",
    "save_path = f'{work_path}/runs/{mode}/{folder}/{jobname}/'\n",
    "Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "default_yaml = f'{save_path}default_params.yml'\n",
    "write_yaml(default_yaml, data)\n",
    "\n",
    "# Create bash script\n",
    "script_path = f'{home_path}/cluster/simulation_cluster.py'\n",
    "script_kwargs = f'--save_path={save_path}'\n",
    "\n",
    "bash = template.format(jobname=jobname,\n",
    "                       n_cpus=n_threads,\n",
    "                       script_path=script_path,\n",
    "                       script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bf0cd82c-d9c2-498c-aeb9-c1698952566b",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_name = 'default_params.yml'\n",
    "\n",
    "default_yaml = f'{save_path}{default_name}'\n",
    "yaml_data = read_yaml(default_yaml)\n",
    "params = {}\n",
    "lasers = yaml_data['lasers']\n",
    "params['laser_params'] = [lasers[f'laser_{i}'] for i in range(len(lasers))]\n",
    "params['save_path'] = save_path\n",
    "keys = ['simbox_params', 'geometry', 'low_memory_mode']\n",
    "for key in keys:\n",
    "    params[key] = yaml_data[key]\n",
    "\n",
    "create_ini_file(**params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8590577a-970f-434e-8907-eb66e330614c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 1055857\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e024d457-c51c-41bc-883d-aca1bc41fe91",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Bash script for 1d grid scan"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d924c56e-4b35-45b2-9102-1a100fe6ff2a",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### CCS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5dc2fa95-9e17-4958-afab-d73c9f74f811",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 180\n",
    "beta = 90\n",
    "\n",
    "# tau = float(25e-15 * 2 * np.sqrt(np.log(2)))\n",
    "# 1st laser\n",
    "laser_params = []\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': 3*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': 0,\n",
    "    'phi': 0,\n",
    "    'beta': 0,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 7,\n",
    "    'trans': 15,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 2,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'gridscan'\n",
    "geometry = 'z'\n",
    "low_memory_mode = False\n",
    "n_threads = 24\n",
    "pol_idx = 0\n",
    "\n",
    "# Variable parameters \n",
    "laser_key = 'laser_0'\n",
    "param_key = 'W'\n",
    "grid = [1,50,10,1]\n",
    "vary_params = {\n",
    "    'lasers': {laser_key: {param_key: grid}}\n",
    "}\n",
    "\n",
    "# Save yaml file\n",
    "data = {\n",
    "    '_mode': mode,\n",
    "    'lasers': laser_params,\n",
    "    'simbox_params': simbox_params,\n",
    "    'geometry': geometry,\n",
    "    'low_memory_mode': low_memory_mode,\n",
    "    'n_threads': n_threads,\n",
    "    'pol_idx': pol_idx,\n",
    "    'do_energy_budget': False,\n",
    "}\n",
    "\n",
    "folder = f'CCS_W'\n",
    "jobname = f'{param_key}_theta_{theta}_{geometry}'\n",
    "save_path = f'{work_path}/runs/{mode}/{folder}/{jobname}/'\n",
    "Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "default_yaml = f'{save_path}default_params.yml'\n",
    "write_yaml(default_yaml, data)\n",
    "\n",
    "vary_yaml = f'{save_path}vary_params.yml'\n",
    "write_yaml(vary_yaml, vary_params)\n",
    "\n",
    "# Create bash script\n",
    "script_path = f'{home_path}/cluster/gridscan_cluster.py'\n",
    "script_kwargs = f'--save_path={save_path}'\n",
    "\n",
    "bash = template.format(jobname=jobname,\n",
    "                       n_cpus=n_threads,\n",
    "                       script_path=script_path,\n",
    "                       script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "8ea13073-02ac-4b96-b268-a138c8995eab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 1055918\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9163cbdd-5639-4a38-b026-3adf44b6a534",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Harmonics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "bcb31d5c-fac4-488a-9ac4-86c82ee2d291",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 180\n",
    "beta = 90\n",
    "\n",
    "Wtotal = 25\n",
    "H = 1 + 1/2**2.5\n",
    "W = [Wtotal/(n**2.5*H) for n in [1,2]]\n",
    "# print(W)\n",
    "# print(np.sum(W))\n",
    "\n",
    "n = 2\n",
    "W = [50/n for i in range(n)]\n",
    "\n",
    "# 1st laser\n",
    "lam0 = 800e-9\n",
    "k0 = 2*np.pi/lam0\n",
    "order = 0\n",
    "gauss_type = 'gaussian_spectral'\n",
    "\n",
    "w0_harm = 1\n",
    "w0_usual = 1\n",
    "phi0_harm = [0 for i in range(n)]\n",
    "\n",
    "laser_params = []\n",
    "for i in range(1,n+1):\n",
    "    laser = {\n",
    "        'tau': 20e-15,\n",
    "        'lam': lam0/i,\n",
    "        'w0': w0_harm*800e-9/i,\n",
    "        'W': W[i-1],\n",
    "        'theta': 0,\n",
    "        'phi': 0,\n",
    "        'beta': 0,\n",
    "        'solution': gauss_type,\n",
    "        'order': order,\n",
    "        'phi0': phi0_harm[i-1],\n",
    "    }\n",
    "    laser['E0'] = float(W_to_E0(laser))\n",
    "    laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 20e-15,\n",
    "    'lam': lam0,\n",
    "    'w0': w0_usual*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': gauss_type,\n",
    "    'order': order,\n",
    "    'phi0': 0\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 7,\n",
    "    'trans': 15,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 3,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'gridscan'\n",
    "geometry = 'z'\n",
    "low_memory_mode = False\n",
    "n_threads = 24\n",
    "pol_idx = 0\n",
    "\n",
    "# Variable parameters \n",
    "laser_key = 'laser_1'\n",
    "param_key = 'phi0'\n",
    "grid = [0,330,12]\n",
    "vary_params = {\n",
    "    'lasers': {laser_key: {param_key: grid}}\n",
    "}\n",
    "\n",
    "discernible_spectral = True\n",
    "sphmap_params = {\n",
    "    'order': 1,\n",
    "    'angular_resolution': 0.5/180*np.pi\n",
    "}\n",
    "\n",
    "sph_limits = {\n",
    "    'k': [1.5*k0, 2.5*k0],\n",
    "    'theta': [np.pi/2, np.pi]\n",
    "}\n",
    "\n",
    "# Save yaml file\n",
    "data = {\n",
    "    '_mode': mode,\n",
    "    'lasers': laser_params,\n",
    "    'simbox_params': simbox_params,\n",
    "    'geometry': geometry,\n",
    "    'low_memory_mode': low_memory_mode,\n",
    "    'n_threads': n_threads,\n",
    "    'pol_idx': pol_idx,\n",
    "    'discernible_spectral': discernible_spectral,\n",
    "    'sphmap_params': sphmap_params,\n",
    "    'sph_limits': sph_limits\n",
    "}\n",
    "\n",
    "folder = f'harmonics_W'\n",
    "if type(w0_harm) == int:\n",
    "    w0_str = str(w0_harm)\n",
    "else:\n",
    "    w0_str = '_'.join(str(w0) for w0 in w0_harm)\n",
    "jobname = f'21-1omega_theta_{theta}_w0harm_{w0_str}_phase'\n",
    "save_path = f'{work_path}/runs/{mode}/{folder}/{jobname}/'\n",
    "Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "default_yaml = f'{save_path}default_params.yml'\n",
    "write_yaml(default_yaml, data)\n",
    "\n",
    "vary_yaml = f'{save_path}vary_params.yml'\n",
    "write_yaml(vary_yaml, vary_params)\n",
    "\n",
    "# Create bash script\n",
    "script_path = f'{home_path}/cluster/gridscan_cluster.py'\n",
    "script_kwargs = f'--save_path={save_path}'\n",
    "\n",
    "bash = template.format(jobname=jobname,\n",
    "                       n_cpus=n_threads,\n",
    "                       script_path=script_path,\n",
    "                       script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "2bf69854-5b74-4e56-8cd1-0cace27d1b39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 923887\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3267f787-08a9-4ebf-a362-e889f936e845",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Bash script for optuna optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "86e74174-a270-4d7d-98e6-67dfe894af0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 180\n",
    "beta = 90\n",
    "\n",
    "# 1st laser\n",
    "laser_params = []\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': 4*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': 0,\n",
    "    'phi': 0,\n",
    "    'beta': 0,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 8,\n",
    "    'trans': 20,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 2,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'optuna'\n",
    "geometry = 'xz'\n",
    "low_memory_mode = False\n",
    "n_threads = 24\n",
    "pol_idx = 0\n",
    "\n",
    "# Optuna params\n",
    "obj_param = 'N_disc_num'\n",
    "n_trials = 5\n",
    "\n",
    "data = {\n",
    "    '_mode': mode,\n",
    "    'lasers': laser_params,\n",
    "    'simbox_params': simbox_params,\n",
    "    'geometry': geometry,\n",
    "    'low_memory_mode': low_memory_mode,\n",
    "    'n_threads': n_threads,\n",
    "    'pol_idx': pol_idx,\n",
    "    'obj_param': obj_param,\n",
    "    'n_trials': n_trials\n",
    "}\n",
    "\n",
    "# Variable parameters \n",
    "laser_key = 'laser_1'\n",
    "param_key = 'beta'\n",
    "grid = [0,180,1,'uniform']\n",
    "vary_params = {laser_key:\n",
    "    {param_key: grid}\n",
    "}\n",
    "\n",
    "# Save yaml files\n",
    "folder = f'CCS'\n",
    "jobname = f'{param_key}_theta_{theta}'\n",
    "save_path = f'{work_path}/runs/{mode}/{folder}/{jobname}/'\n",
    "Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "default_yaml = f'{save_path}default_params.yml'\n",
    "write_yaml(default_yaml, data)\n",
    "\n",
    "vary_yaml = f'{save_path}optuna_params.yml'\n",
    "write_yaml(vary_yaml, vary_params)\n",
    "\n",
    "script_path = f'{home_path}/cluster/optuna_cluster.py'\n",
    "script_kwargs = f'--save_path={save_path}'\n",
    "\n",
    "bash = template.format(jobname=f'opt_{jobname}',\n",
    "                       n_cpus=n_threads,\n",
    "                       script_path=script_path,\n",
    "                       script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fabfaea9-1c37-4edb-9d00-0ee18bd124bc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 600622\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05a705a1-ec0d-47cb-9814-8ffe3bf79b07",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Bash for numerical consistency check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "6a3a65bf-3bf9-49eb-bab0-c569788bab77",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_res = '''#!/bin/bash\n",
    "#SBATCH --job-name={jobname}\n",
    "#SBATCH --partition=hij-gpu\n",
    "#SBATCH --cpus-per-task={n_cpus}\n",
    "#SBATCH --time=8:00:00\n",
    "#SBATCH --mem=200GB\n",
    "#SBATCH --output=/work/wi73yus/bash_output/{jobname}.%A_%a.log\n",
    "#SBATCH --error=/work/wi73yus/bash_output/{jobname}.%A_%a.err\n",
    "#SBATCH --array=0-4\n",
    "\n",
    "arr=$SLURM_ARRAY_TASK_ID\n",
    "echo $arr\n",
    "\n",
    "source ~/.bashrc  #  not automatically read by slurm\n",
    "conda activate lbl  # activate here the anaconda environment you want to use\n",
    "\n",
    "python {script_path} {script_kwargs}\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd8227fd-ceb1-441c-abde-267ed2f8ef44",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### ECS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "909dd42d-5d29-4a48-b07d-19b2f676e2c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need to test the same scattering configuration with a couple of \n",
    "# better numerical resolutions\n",
    "# 1) increase time resolution,\n",
    "# 2) increase spatial resolution\n",
    "# 3) increase box size in time\n",
    "# 4) increase box size in space\n",
    "\n",
    "# Laser pulse parameters\n",
    "theta = 180\n",
    "beta = 90\n",
    "\n",
    "# 1st laser\n",
    "laser_params = []\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0x': 4*800e-9,\n",
    "    'w0y': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': 0,\n",
    "    'phi': 0,\n",
    "    'beta': 0,\n",
    "    'solution': 'gaussian_elliptical_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 8,\n",
    "    'trans': 20,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 2,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'gridscan'\n",
    "geometry = 'xz'\n",
    "low_memory_mode = False\n",
    "n_threads = 12\n",
    "pol_idx = 0\n",
    "\n",
    "data = {\n",
    "    'lasers': laser_params,\n",
    "    'simbox_params': simbox_params,\n",
    "    'geometry': geometry,\n",
    "    'low_memory_mode': low_memory_mode,\n",
    "    'n_threads': n_threads,\n",
    "    'pol_idx': pol_idx,\n",
    "}\n",
    "\n",
    "simbox_cases = {f'case_{i}': deepcopy(simbox_params) for i in range(0,5)}\n",
    "\n",
    "simbox_cases['case_1']['resolutions']['t'] = 2\n",
    "simbox_cases['case_2']['resolutions']['spatial'] = 3\n",
    "simbox_cases['case_3']['box_factors']['t'] = 4\n",
    "simbox_cases['case_4']['box_factors']['long'] = 10\n",
    "simbox_cases['case_4']['box_factors']['trans'] = 25\n",
    "\n",
    "folder = f'CCS'\n",
    "jobname = f'theta_{theta}'\n",
    "save_path = f'{work_path}/runs/resolution_check/{folder}/{jobname}/'\n",
    "\n",
    "for key in simbox_cases.keys():\n",
    "    data['simbox_params'] = simbox_cases[key]\n",
    "    default_yaml = f'{save_path}{key}/default_params.yml'\n",
    "    Path(os.path.dirname(default_yaml)).mkdir(parents=True, exist_ok=True)\n",
    "    write_yaml(default_yaml, data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "5ba0d1a6-58a2-470f-b228-8525e3443f86",
   "metadata": {},
   "outputs": [],
   "source": [
    "script_path = f'{home_path}/cluster/simulation_cluster.py'\n",
    "save_path_arr = save_path + 'case_$arr/'\n",
    "script_kwargs = f'--save_path={save_path_arr}'\n",
    "\n",
    "bash = template_res.format(jobname=f'numcheck_{jobname}',\n",
    "                           n_cpus=n_threads,\n",
    "                           script_path=script_path,\n",
    "                           script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "6af8a790-3c73-4d99-8ffb-1525a8c6003c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 600314\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3683e4e2-f2b8-402e-b685-e5f3f130185c",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Harmonics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "96400471-f79a-4357-9b61-dd943772a422",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[21.24447237944163, 3.7555276205583703]\n",
      "25.0\n"
     ]
    }
   ],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 180\n",
    "beta = 90\n",
    "\n",
    "Wtotal = 25\n",
    "H = 1 + 1/2**2.5\n",
    "W = [Wtotal/(n**2.5*H) for n in [1,2]]\n",
    "print(W)\n",
    "print(np.sum(W))\n",
    "\n",
    "# 1st laser\n",
    "w0_harm = 1\n",
    "w0_usual = 1\n",
    "# 1st laser\n",
    "laser_params = []\n",
    "for i in range(1,3):\n",
    "    laser = {\n",
    "        'tau': 20e-15,\n",
    "        'lam': 800e-9/i,\n",
    "        'w0': w0_harm*800e-9/i,\n",
    "        'W': W[i-1],\n",
    "        'theta': 0,\n",
    "        'phi': 0,\n",
    "        'beta': 0,\n",
    "        'solution': 'gaussian_paraxial'\n",
    "    }\n",
    "    laser['E0'] = float(W_to_E0(laser))\n",
    "    laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 20e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': w0_usual*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 7,\n",
    "    'trans': 15,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 2,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'resolution_check'\n",
    "geometry = 'z'\n",
    "low_memory_mode = False\n",
    "n_threads = 24\n",
    "pol_idx = 0\n",
    "\n",
    "discernible_spectral = True\n",
    "sphmap_params = {\n",
    "    'order': 1,\n",
    "    'angular_resolution': 0.5/180*np.pi\n",
    "}\n",
    "\n",
    "# Save yaml file\n",
    "data = {\n",
    "    '_mode': mode,\n",
    "    'lasers': laser_params,\n",
    "    'simbox_params': simbox_params,\n",
    "    'geometry': geometry,\n",
    "    'low_memory_mode': low_memory_mode,\n",
    "    'n_threads': n_threads,\n",
    "    'pol_idx': pol_idx,\n",
    "    'discernible_spectral': discernible_spectral,\n",
    "    'sphmap_params': sphmap_params\n",
    "}\n",
    "\n",
    "simbox_cases = {f'case_{i}': deepcopy(simbox_params) for i in range(0,5)}\n",
    "\n",
    "simbox_cases['case_1']['resolutions']['t'] = 2\n",
    "simbox_cases['case_2']['resolutions']['spatial'] = 4\n",
    "simbox_cases['case_3']['box_factors']['t'] = 4\n",
    "simbox_cases['case_4']['box_factors']['long'] = 10\n",
    "simbox_cases['case_4']['box_factors']['trans'] = 25\n",
    "\n",
    "folder = f'harmonics'\n",
    "jobname = f'21-1omega_theta_{theta}_beta_{beta}_w0harm_{w0_harm}_w0usual_{w0_usual}'\n",
    "save_path = f'{work_path}/runs/{mode}/{folder}/{jobname}/'\n",
    "\n",
    "for key in simbox_cases.keys():\n",
    "    data['simbox_params'] = simbox_cases[key]\n",
    "    default_yaml = f'{save_path}{key}/default_params.yml'\n",
    "    Path(os.path.dirname(default_yaml)).mkdir(parents=True, exist_ok=True)\n",
    "    write_yaml(default_yaml, data)\n",
    "\n",
    "default_yaml = f'{save_path}default_params.yml'\n",
    "write_yaml(default_yaml, data)\n",
    "\n",
    "# Create bash script\n",
    "script_path = f'{home_path}/cluster/simulation_cluster.py'\n",
    "save_path_arr = save_path + 'case_$arr/'\n",
    "script_kwargs = f'--save_path={save_path_arr}'\n",
    "\n",
    "bash = template_res.format(jobname=f'numcheck_{jobname}',\n",
    "                           n_cpus=n_threads,\n",
    "                           script_path=script_path,\n",
    "                           script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "43e3ba01-6b1e-4bf6-a500-fab010a9d0f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 721425\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d8b02ce-627a-4984-8463-060d57363446",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Bash for 2d grid scan where one dimension is performed in parallel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fb0dafe2-14f4-415f-b63a-eb0847b0376b",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_2d_gridscan = '''#!/bin/bash\n",
    "#SBATCH --job-name={jobname}\n",
    "#SBATCH --partition=hij-gpu\n",
    "#SBATCH --cpus-per-task={n_cpus}\n",
    "#SBATCH --time=24:00:00\n",
    "#SBATCH --mem=250GB\n",
    "#SBATCH --output=/work/wi73yus/bash_output/{jobname}.%A_%a.log\n",
    "#SBATCH --error=/work/wi73yus/bash_output/{jobname}.%A_%a.err\n",
    "#SBATCH --array={array_size}\n",
    "\n",
    "source ~/.bashrc  #  not automatically read by slurm\n",
    "conda activate lbl  # activate here the anaconda environment you want to use\n",
    "\n",
    "echo $SLURM_ARRAY_TASK_ID\n",
    "\n",
    "python {script_path} {script_kwargs}\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd716c71-9dd5-4c48-bbcc-09e08e3f77a7",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### ECS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1ff7f1d7-87c0-45f3-bdd1-0b9f207f826a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 180\n",
    "beta = 90\n",
    "\n",
    "# 1st laser\n",
    "laser_params = []\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0x': 1*800e-9,\n",
    "    'w0y': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': 0,\n",
    "    'phi': 0,\n",
    "    'beta': 0,\n",
    "    'solution': 'gaussian_elliptical_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 7,\n",
    "    'trans': 15,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 2,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'gridscan'\n",
    "geometry = 'z'\n",
    "low_memory_mode = False\n",
    "n_threads = 18\n",
    "pol_idx = 0\n",
    "\n",
    "# Variable parameters \n",
    "laser_key = 'laser_0'\n",
    "param_key = 'w0x'\n",
    "grid = [1,9,9,800e-9]\n",
    "vary_params = {\n",
    "    'lasers': {laser_key: {param_key: grid}}\n",
    "}\n",
    "\n",
    "# Additional variable parameter\n",
    "laser_key_add = 'laser_0'\n",
    "param_key_add = 'w0y'\n",
    "grid = [1,9,9,800e-9]\n",
    "param_grid = np.linspace(grid[0], grid[1], grid[2], endpoint=True)\n",
    "\n",
    "for idx, param_value in enumerate(param_grid):\n",
    "    laser_params[laser_key_add][param_key_add] = float(param_value*grid[-1])\n",
    "    laser_params[laser_key_add]['E0'] = float(W_to_E0(laser_params[laser_key_add]))\n",
    "\n",
    "    # Save yaml file\n",
    "    data = {\n",
    "        '_mode': mode,\n",
    "        'lasers': laser_params,\n",
    "        'simbox_params': simbox_params,\n",
    "        'geometry': geometry,\n",
    "        'low_memory_mode': low_memory_mode,\n",
    "        'n_threads': n_threads,\n",
    "        'pol_idx': pol_idx\n",
    "    }\n",
    "\n",
    "    folder = f'ECS_W'\n",
    "    jobname = f'{param_key}_theta_{theta}_{geometry}'\n",
    "    save_path = f'{work_path}/runs/{mode}/{folder}/{jobname}/'\n",
    "    save_path_local = f'{save_path}/{idx}/'\n",
    "    Path(os.path.dirname(save_path_local)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    default_yaml = f'{save_path_local}default_params.yml'\n",
    "    write_yaml(default_yaml, data)\n",
    "\n",
    "    vary_yaml = f'{save_path_local}vary_params.yml'\n",
    "    write_yaml(vary_yaml, vary_params)\n",
    "\n",
    "# Create bash script\n",
    "script_path = f'{home_path}/cluster/gridscan_cluster.py'\n",
    "script_kwargs = f'--save_path={save_path}$SLURM_ARRAY_TASK_ID/'\n",
    "array_size = f'0-{len(param_grid)-1}'\n",
    "\n",
    "bash = template_2d_gridscan.format(jobname=jobname,\n",
    "                       n_cpus=n_threads,\n",
    "                       array_size=array_size,\n",
    "                       script_path=script_path,\n",
    "                       script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "73f4809d-c6b8-4884-b325-a14da9e815e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 924422\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb93ebaf-6283-4b9e-96f0-86f2ca744860",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Harmonics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b1b77c9d-0315-41c3-b33a-b88e255802dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 180\n",
    "beta = 90\n",
    "\n",
    "Wtotal = 25\n",
    "H = 1 + 1/2**2.5\n",
    "W = [Wtotal/(n**2.5*H) for n in [1,2]]\n",
    "# print(W)\n",
    "# print(np.sum(W))\n",
    "\n",
    "n = 3\n",
    "W = [50/n for i in range(n)]\n",
    "\n",
    "# 1st laser\n",
    "lam0 = 800e-9\n",
    "k0 = 2*np.pi/lam0\n",
    "order = 0\n",
    "gauss_type = 'gaussian_spectral'\n",
    "\n",
    "w0_harm = 1\n",
    "w0_usual = 1\n",
    "phi0_harm = [0 for i in range(n)]\n",
    "\n",
    "laser_params = []\n",
    "for i in range(1,4):\n",
    "    laser = {\n",
    "        'tau': 20e-15,\n",
    "        'lam': lam0/i,\n",
    "        'w0': w0_harm*800e-9/i,\n",
    "        'W': W[i-1],\n",
    "        'theta': 0,\n",
    "        'phi': 0,\n",
    "        'beta': 0,\n",
    "        'solution': gauss_type,\n",
    "        'order': order,\n",
    "        'phi0': phi0_harm[i-1],\n",
    "    }\n",
    "    laser['E0'] = float(W_to_E0(laser))\n",
    "    laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 20e-15,\n",
    "    'lam': lam0,\n",
    "    'w0': w0_usual*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': gauss_type,\n",
    "    'order': order,\n",
    "    'phi0': 0\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 7,\n",
    "    'trans': 15,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 2,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'gridscan'\n",
    "geometry = 'z'\n",
    "low_memory_mode = False\n",
    "n_threads = 24\n",
    "pol_idx = 0\n",
    "\n",
    "# Variable parameters \n",
    "laser_key = 'laser_1'\n",
    "param_key = 'phi0'\n",
    "grid = [0,330,12]\n",
    "vary_params = {\n",
    "    'lasers': {laser_key: {param_key: grid}}\n",
    "}\n",
    "\n",
    "# Additional variable parameter\n",
    "laser_key_add = 'laser_2'\n",
    "param_key_add = 'phi0'\n",
    "grid = [0,330,12]\n",
    "param_grid = np.linspace(grid[0], grid[1], grid[2], endpoint=True)\n",
    "\n",
    "for idx, param_value in enumerate(param_grid):\n",
    "    laser_params[laser_key_add][param_key_add] = float(param_value)\n",
    "    laser_params[laser_key_add]['E0'] = float(W_to_E0(laser_params[laser_key_add]))\n",
    "\n",
    "    # Save yaml file\n",
    "    data = {\n",
    "        '_mode': mode,\n",
    "        'lasers': laser_params,\n",
    "        'simbox_params': simbox_params,\n",
    "        'geometry': geometry,\n",
    "        'low_memory_mode': low_memory_mode,\n",
    "        'n_threads': n_threads,\n",
    "        'pol_idx': pol_idx\n",
    "    }\n",
    "\n",
    "    folder = f'harmonics_W'\n",
    "    jobname = f'{param_key}_{laser_key}_theta_{theta}_{geometry}'\n",
    "    save_path = f'{work_path}/runs/{mode}/{folder}/{jobname}/'\n",
    "    save_path_local = f'{save_path}/{idx}/'\n",
    "    Path(os.path.dirname(save_path_local)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    default_yaml = f'{save_path_local}default_params.yml'\n",
    "    write_yaml(default_yaml, data)\n",
    "\n",
    "    vary_yaml = f'{save_path_local}vary_params.yml'\n",
    "    write_yaml(vary_yaml, vary_params)\n",
    "\n",
    "# Create bash script\n",
    "script_path = f'{home_path}/cluster/gridscan_cluster.py'\n",
    "script_kwargs = f'--save_path={save_path}$SLURM_ARRAY_TASK_ID/'\n",
    "array_size = f'0-{len(param_grid)-1}'\n",
    "\n",
    "bash = template_2d_gridscan.format(jobname=jobname,\n",
    "                       n_cpus=n_threads,\n",
    "                       array_size=array_size,\n",
    "                       script_path=script_path,\n",
    "                       script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e43279cd-ce0e-473a-851f-1cce700dce50",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 924025\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd453d00-7b07-464c-9f07-aee73f1298a5",
   "metadata": {},
   "source": [
    "### Bash for optuna optimization from several jobs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1814f425-4cd9-4d2b-b149-ae5a192c46a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_optuna = '''#!/bin/bash\n",
    "#SBATCH --job-name={jobname}\n",
    "#SBATCH --partition=gpu\n",
    "#SBATCH --cpus-per-task={n_cpus}\n",
    "#SBATCH --time=24:00:00\n",
    "#SBATCH --mem=200GB\n",
    "#SBATCH --output=/work/wi73yus/bash_output/{jobname}.%A_%a.log\n",
    "#SBATCH --error=/work/wi73yus/bash_output/{jobname}.%A_%a.err\n",
    "#SBATCH --array=0-1\n",
    "\n",
    "source ~/.bashrc  #  not automatically read by slurm\n",
    "conda activate lbl  # activate here the anaconda environment you want to use\n",
    "\n",
    "echo $SLURM_ARRAY_TASK_ID\n",
    "\n",
    "if [ $SLURM_ARRAY_TASK_ID -gt 0 ]\n",
    "then\n",
    "    sleep 300\n",
    "fi\n",
    "\n",
    "python {script_path} {script_kwargs}\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6daf8c0b-9526-4fce-b9bc-66b00b03e849",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### CCS optimizations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "5cbcbd5c-d359-4aa5-b6a7-594be7499231",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 160\n",
    "beta = 90\n",
    "\n",
    "# 1st laser\n",
    "laser_params = []\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': 0,\n",
    "    'phi': 0,\n",
    "    'beta': 0,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': 'gaussian_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 8,\n",
    "    'trans': 18,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 2,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'optuna'\n",
    "geometry = 'z'\n",
    "low_memory_mode = False\n",
    "n_threads = 18\n",
    "pol_idx = 0\n",
    "\n",
    "# Optuna params\n",
    "obj_param = 'N_disc_num'\n",
    "n_trials = 6\n",
    "consider_endpoints = True\n",
    "\n",
    "data = {\n",
    "    '_mode': mode,\n",
    "    'lasers': laser_params,\n",
    "    'simbox_params': simbox_params,\n",
    "    'geometry': geometry,\n",
    "    'low_memory_mode': low_memory_mode,\n",
    "    'n_threads': n_threads,\n",
    "    'pol_idx': pol_idx,\n",
    "    'obj_param': obj_param,\n",
    "    'n_trials': n_trials,\n",
    "    'consider_endpoints': consider_endpoints\n",
    "}\n",
    "\n",
    "# Variable parameters \n",
    "param_key = 'w0'\n",
    "grid = [1,9,800e-9,'uniform']\n",
    "vary_params = {\n",
    "    'laser_0':\n",
    "        {param_key: grid}\n",
    "}\n",
    "\n",
    "# Save yaml files\n",
    "folder = f'CCS_W'\n",
    "subfolder = f'{param_key}_theta_{theta}_better_res'\n",
    "save_path = f'{work_path}/runs/{mode}/{folder}/{subfolder}/'\n",
    "Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "default_yaml = f'{save_path}default_params.yml'\n",
    "write_yaml(default_yaml, data)\n",
    "\n",
    "vary_yaml = f'{save_path}optuna_params.yml'\n",
    "write_yaml(vary_yaml, vary_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "38ca01cd-ca62-4b3f-8d01-f9f5a22229ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "jobname = f'm_opt_{param_key}'\n",
    "\n",
    "script_path = f'{home_path}/cluster/optuna_cluster.py'\n",
    "script_kwargs = f'--save_path={save_path}'\n",
    "partition = 'long'\n",
    "\n",
    "bash = template_optuna.format(jobname=f'{jobname}',\n",
    "                              n_cpus=n_threads,\n",
    "                              script_path=script_path,\n",
    "                              script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "d8871165-73b0-40d2-8da1-3612e695d579",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 922310\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78548530-34ee-4047-b793-7e533cd1510b",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### ECS optimizations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "acfa75c8-60ba-4439-96de-64e852752e54",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 160\n",
    "beta = 90\n",
    "\n",
    "# 1st laser\n",
    "laser_params = []\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    'w0x': 1*800e-9,\n",
    "    'w0y': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': 0,\n",
    "    'phi': 0,\n",
    "    'beta': 0,\n",
    "    'solution': 'gaussian_elliptical_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 25e-15,\n",
    "    'lam': 800e-9,\n",
    "    # 'w0': 1*800e-9,\n",
    "    'w0x': 1*800e-9,\n",
    "    'w0y': 1*800e-9,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': 'gaussian_elliptical_paraxial'\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 7,\n",
    "    'trans': 15,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 1.5,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "# Other parameters\n",
    "mode = 'optuna'\n",
    "geometry = 'z'\n",
    "low_memory_mode = False\n",
    "n_threads = 24\n",
    "pol_idx = 0\n",
    "\n",
    "# Optuna params\n",
    "obj_param = 'N_disc_num'\n",
    "n_trials = 12\n",
    "consider_endpoints = True\n",
    "consider_prior = False\n",
    "n_startup_trials = 6\n",
    "\n",
    "data = {\n",
    "    '_mode': mode,\n",
    "    'lasers': laser_params,\n",
    "    'simbox_params': simbox_params,\n",
    "    'geometry': geometry,\n",
    "    'low_memory_mode': low_memory_mode,\n",
    "    'n_threads': n_threads,\n",
    "    'pol_idx': pol_idx,\n",
    "    'obj_param': obj_param,\n",
    "    'n_trials': n_trials,\n",
    "    'consider_endpoints': consider_endpoints,\n",
    "    'consider_prior': consider_prior,\n",
    "    'n_startup_trials': n_startup_trials,\n",
    "}\n",
    "\n",
    "# Variable parameters \n",
    "param_key = 'w0x_w0y'\n",
    "# grid = [1,15,800e-9,'uniform']\n",
    "# vary_params = {\n",
    "#     'laser_0':\n",
    "#         {param_key: grid}\n",
    "# }\n",
    "vary_params = {\n",
    "    'laser_0':\n",
    "        {'w0x': [1,9,800e-9,'uniform'],\n",
    "         'w0y': [1,9,800e-9,'uniform'],},\n",
    "}\n",
    "# vary_params = {\n",
    "#     'laser_0':\n",
    "#         {'w0x': [1,9,800e-9,'uniform'],\n",
    "#          'w0y': [1,9,800e-9,'uniform'],\n",
    "#          'W': [0,50,1,'uniform']},\n",
    "#     'laser_1':\n",
    "#         {'w0x': [1,9,800e-9,'uniform'],\n",
    "#          'w0y': [1,9,800e-9,'uniform']},\n",
    "# }\n",
    "\n",
    "# Save yaml files\n",
    "folder = f'ECS_W'\n",
    "subfolder = f'{param_key}_theta_{theta}_2_threads'\n",
    "save_path = f'{work_path}/runs/{mode}/{folder}/{subfolder}/'\n",
    "Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "default_yaml = f'{save_path}default_params.yml'\n",
    "write_yaml(default_yaml, data)\n",
    "\n",
    "vary_yaml = f'{save_path}optuna_params.yml'\n",
    "write_yaml(vary_yaml, vary_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7d92e092-96dd-423d-84c6-5d96f90246b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "jobname = f'm_opt_{param_key}'\n",
    "\n",
    "script_path = f'{home_path}/cluster/optuna_cluster.py'\n",
    "script_kwargs = f'--save_path={save_path}'\n",
    "\n",
    "bash = template_optuna.format(jobname=f'{jobname}',\n",
    "                              n_cpus=n_threads,\n",
    "                              script_path=script_path,\n",
    "                              script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c722e8c7-65c9-4690-89b4-aad6e3a3901e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Submitted batch job 1108854\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec34d80c-5c53-4583-9db5-229f6b43c001",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Harmonics optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "262ae492-3538-4f23-9975-a92649de434b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Laser pulse parameters\n",
    "theta = 180\n",
    "beta = 90\n",
    "\n",
    "lam0 = 800e-9\n",
    "k0 = 2*np.pi/lam0\n",
    "\n",
    "order = 0\n",
    "gauss_type = 'gaussian_spectral'\n",
    "\n",
    "w0_harm = 1\n",
    "w0_usual = 1\n",
    "# 1st laser\n",
    "laser_params = []\n",
    "for i in range(1,3):\n",
    "    laser = {\n",
    "        'tau': 20e-15,\n",
    "        'lam': lam0/i,\n",
    "        'w0': w0_harm*lam0/i,\n",
    "        'W': 12.5,\n",
    "        'theta': 0,\n",
    "        'phi': 0,\n",
    "        'beta': 0,\n",
    "        'solution': gauss_type,\n",
    "        'order': order,\n",
    "    }\n",
    "    laser['E0'] = float(W_to_E0(laser))\n",
    "    laser_params.append(laser)\n",
    "\n",
    "# 2nd laser\n",
    "laser = {\n",
    "    'tau': 20e-15,\n",
    "    'lam': lam0,\n",
    "    'w0': w0_usual*lam0,\n",
    "    'W': 25,\n",
    "    'theta': theta,\n",
    "    'phi': 0,\n",
    "    'beta': beta,\n",
    "    'solution': gauss_type,\n",
    "    'order': order,\n",
    "}\n",
    "laser['E0'] = float(W_to_E0(laser))\n",
    "laser_params.append(laser)\n",
    "\n",
    "laser_params = {f'laser_{i}': value for i,value in enumerate(laser_params)}\n",
    "\n",
    "# Simulation box parameters\n",
    "box_factors = {\n",
    "    'long': 8,\n",
    "    'trans': 20,\n",
    "    't': 2\n",
    "}\n",
    "\n",
    "resolutions = {\n",
    "    'spatial': 2,\n",
    "    't': 1\n",
    "}\n",
    "simbox_params = {'box_factors': box_factors,\n",
    "                 'resolutions': resolutions}\n",
    "\n",
    "mode = 'optuna'\n",
    "geometry = 'z'\n",
    "low_memory_mode = False\n",
    "n_threads = 20\n",
    "pol_idx = 0\n",
    "\n",
    "discernible_spectral = True\n",
    "sphmap_params = {\n",
    "    'order': 1,\n",
    "    'angular_resolution': 0.5/180*np.pi\n",
    "}\n",
    "\n",
    "sph_limits = {\n",
    "    'k': [1.5*k0, 2.5*k0],\n",
    "    'theta': [0, np.pi]\n",
    "}\n",
    "\n",
    "# Optuna params\n",
    "obj_param = 'N_disc_num'\n",
    "n_trials = 10\n",
    "consider_endpoints = True\n",
    "consider_prior = False\n",
    "\n",
    "data = {\n",
    "    '_mode': mode,\n",
    "    'lasers': laser_params,\n",
    "    'simbox_params': simbox_params,\n",
    "    'geometry': geometry,\n",
    "    'low_memory_mode': low_memory_mode,\n",
    "    'n_threads': n_threads,\n",
    "    'pol_idx': pol_idx,\n",
    "    'obj_param': obj_param,\n",
    "    'n_trials': n_trials,\n",
    "    'consider_endpoints': consider_endpoints,\n",
    "    'discernible_spectral': discernible_spectral,\n",
    "    'sphmap_params': sphmap_params,\n",
    "    'sph_limits': sph_limits,\n",
    "    'consider_prior': consider_prior,\n",
    "}\n",
    "\n",
    "# Variable parameters \n",
    "param_key = 'W'\n",
    "grid = [0,50,1,'uniform']\n",
    "vary_params = {\n",
    "    'laser_0':\n",
    "        {'W': grid}\n",
    "}\n",
    "\n",
    "# Save yaml files\n",
    "folder = f'harmonics_W'\n",
    "subfolder = f'{param_key}_2omega_channel'\n",
    "save_path = f'{work_path}/runs/{mode}/{folder}/{subfolder}/'\n",
    "Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "default_yaml = f'{save_path}default_params.yml'\n",
    "write_yaml(default_yaml, data)\n",
    "\n",
    "vary_yaml = f'{save_path}optuna_params.yml'\n",
    "write_yaml(vary_yaml, vary_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "864085a2-c1bb-4691-a1f5-f73bf8f66e69",
   "metadata": {},
   "outputs": [],
   "source": [
    "jobname = f'm_opt_{param_key}'\n",
    "\n",
    "script_path = f'{home_path}/cluster/optuna_cluster.py'\n",
    "script_kwargs = f'--save_path={save_path}'\n",
    "\n",
    "bash = template_optuna.format(jobname=f'{jobname}',\n",
    "                              n_cpus=n_threads,\n",
    "                              script_path=script_path,\n",
    "                              script_kwargs=script_kwargs)\n",
    "\n",
    "bash_path = f'{save_path}job.slurm'\n",
    "with open(bash_path, 'w') as f:\n",
    "    f.write(bash)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "8ecdabdb-6a8f-440e-b341-c85e248ad41f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sbatch: error: invalid partition specified: hij-jena\n",
      "sbatch: error: Batch job submission failed: Invalid partition name specified\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "256"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'sbatch {bash_path}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec17f10d-fc64-40fb-add6-38ee8cfe51af",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}