generate_configs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fd2ab28d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os.path\n",
    "import numpy as np\n",
    "\n",
    "from pynextsim.nextsim_config import NextsimConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4884615d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_grid_params(params, exporter_path_root):\n",
    "    vectors = list(params.values())\n",
    "    grids = np.meshgrid(*vectors, indexing='ij')\n",
    "    all_vals = np.vstack([a.flatten() for a in grids]).T\n",
    "\n",
    "    grid_params = {}\n",
    "    for i, vals in enumerate(all_vals):\n",
    "        grid_params[i] = {}\n",
    "        for j, key in enumerate(params.keys()):\n",
    "            grid_params[i][key] = all_vals[i,j]\n",
    "\n",
    "    for i in grid_params:\n",
    "        grid_params[i]['output.exporter_path'] = f'{exporter_path_root}/exp_{i:05d}'\n",
    "    return grid_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4c97d6fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'mesh': {'filename': 'small_arctic_5km.msh'}, 'setup': {'ice-type': 'topaz_osisaf_icesat', 'ocean-type': 'topaz', 'atmosphere-type': 'era5'}, 'simul': {'spinup_duration': '1', 'timestep': '450', 'time_init': '2006-11-15', 'duration': '100'}, 'thermo': {'h_young_max': '0.5'}, 'dynamics': {'time_relaxation_damage': '15', 'compression_factor': '10e3', 'C_lab': '2e6', 'substeps': '90'}, 'output': {'output_per_day': '4', 'exporter_path': 'None', 'export_after_regrid': 'true', 'export_before_regrid': 'true'}}\n",
      "#!/bin/bash -x\n",
      "\n",
      "## Project:\n",
      "#SBATCH --account=nn9878k\n",
      "\n",
      "## Job name:\n",
      "#SBATCH --job-name=EXP_NAME\n",
      "\n",
      "## Output file\n",
      "#SBATCH --output=/cluster/home/akorosov/harmony.%j.log         # Stdout & stderr\n",
      "\n",
      "## Wall time limit:\n",
      "#SBATCH --time=1-12:00:00\n",
      "\n",
      "## Number of nodes:\n",
      "#SBATCH --nodes=1\n",
      "\n",
      "## Number of tasks (total)\n",
      "#SBATCH --ntasks=32\n",
      "\n",
      "## Set OMP_NUM_THREADS\n",
      "#SBATCH --cpus-per-task=1\n",
      "\n",
      "## Queue Option (preproc for 1 node; normal for >4 nodes; also can be bigmem)\n",
      "#\n",
      "\n",
      "## Ensure we have a clean environment\n",
      "#SBATCH --export=NONE\n",
      "\n",
      "## Email info\n",
      "#SBATCH --mail-type=ALL # Mail events (NONE, BEGIN, END, FAIL, ALL)\n",
      "#SBATCH --mail-user=anton.korosov@nersc.no # email to the user\n",
      "\n",
      "source $HOME/pynextsimf.src\n",
      "mpirun -np 32 nextsim.exec --config-files=/cluster/home/akorosov/n/nextsim-env/config_files/harmony/EXP_NAME.cfg > /cluster/home/akorosov/EXP_NAME.log 2>&1 || exit\n",
      "\n"
     ]
    }
   ],
   "source": [
    "conf_dir = '/Home/antonk/n/nextsim-env/config_files'\n",
    "conf_template_name = f'{conf_dir}/hrmny_sa05free_template.cfg'\n",
    "sbatch_template_name = f'{conf_dir}/sbatch_nextsim_template.sh'\n",
    "exporter_path_root = '/cluster/work/users/akorosov/harmony'\n",
    "params = {\n",
    "    'dynamics.compression_factor': [5e6, 10e6, 15e6],\n",
    "    'dynamics.C_lab': [0.2e6, 0.4e6, 0.6e6, 1e6, 1.6e6, 2.5e6, 4e6],\n",
    "}\n",
    "\n",
    "grid_params = create_grid_params(params, exporter_path_root)\n",
    "\n",
    "template = NextsimConfig(conf_template_name)\n",
    "\n",
    "with open(sbatch_template_name) as fid:\n",
    "    sbatch_template = fid.read()\n",
    "    \n",
    "print(template)\n",
    "print(sbatch_template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7752a1c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in grid_params:\n",
    "    for param in grid_params[i]:\n",
    "        p0, p1 = param.split('.')\n",
    "        template[p0][p1] = grid_params[i][param]\n",
    "        conf_name = conf_template_name.replace('template.cfg', f'{i:05}.cfg')\n",
    "        with open(conf_name, 'wt') as fid:\n",
    "            template.write(fid)\n",
    "            \n",
    "        sbatch_name = sbatch_template_name.replace('template.sh', f'{i:05}.sh')\n",
    "        with open(sbatch_name, 'wt') as fid:\n",
    "            fid.write(sbatch_template.replace('EXP_NAME', os.path.splitext(os.path.basename(conf_name))[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c2f547f",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}