sapsan/examples/pimlturb1d_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Physics-Informed Convolutional model to predict\n",
    "the diagonal Reynolds stress tensor terms, which\n",
    "can later be used to calculate Turbulent Pressure.\n",
    "\n",
    "The example data is a 1D mapped version of the 3D CCSN\n",
    "12 solar mass simulation from Burrows et al, 2020.\n",
    "\n",
    "Note: only select checkpoints are provided for this example,\n",
    "      while the full mapped dataset is available through \n",
    "      documentation\n",
    "\n",
    "The algorithm is best suited to be trained on CPU/s,\n",
    "given the limited data and the shallow network employed.\n",
    "\n",
    "Note: due to a custom loss implementation, Catalyst is NOT used,\n",
    "hence logging is handled manually for this example!\n",
    "\n",
    "-pikarpov\n",
    "\"\"\"\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "from scipy import interpolate\n",
    "import scipy.signal as sig\n",
    "\n",
    "sys.path.append(\"../../\")\n",
    "\n",
    "from sapsan.lib.backends import MLflowBackend, FakeBackend\n",
    "from sapsan.lib.data import HDF5Dataset\n",
    "from sapsan.lib.estimator.pimlturb1d.pimlturb1d_estimator import PIMLTurb1D, PIMLTurb1DConfig\n",
    "from sapsan.lib import Train, Evaluate\n",
    "from sapsan.utils.plot import log_plot, line_plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# --- Experiment tracking backend ---\n",
    "# MLflow - the server will be launched automatically\n",
    "# in case it won't, type in cmd: mlflow ui --port=5000\n",
    "# to use it, set mlflow = True\n",
    "\n",
    "mlflow = False\n",
    "experiment_name = \"PIMLTurb 1D experiment\"\n",
    "\n",
    "if mlflow: \n",
    "    tracking_backend = MLflowBackend(experiment_name, host=\"localhost\", port=5000)\n",
    "else: tracking_backend = FakeBackend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false,
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# --- Data setup ---\n",
    "# In the intereset of loading and training multiple timesteps\n",
    "# one can specify which checkpoints to use and where\n",
    "# they appear in the path via syntax: {checkpoint:format}\n",
    "# \n",
    "# Next, you need to specify which features to load; let's assume \n",
    "#         path = \"{feature}.h5\"\n",
    "# \n",
    "#  1. If in different files, then specify features directly;\n",
    "#     The default HDF5 label will be the last label in the file\n",
    "#     Ex: features = ['velocity', 'denisty', 'pressure']\n",
    "#  2. If in the same, then duplicate the name in features\n",
    "#     and specify which labels to pull\n",
    "#     Ex: features = [\"data\", \"data\", \"data\"]\n",
    "#         feature_labels = ['velocity', 'density', 'pressure']\n",
    "\n",
    "smass = 's12.0'\n",
    "base  = 'data/1dccsn/%s.swbj15.horo.3d/'%smass\n",
    "path  = base+'dump_{checkpoint:05d}.h5'\n",
    "\n",
    "# rho, Pgas, Vsound, T, entropy\n",
    "features_label = ['rho','eos0','eos1','eos2','eos3']\n",
    "target_label   = ['Pturb']\n",
    "\n",
    "# Target will be Pgas/Pturb which will need to be calculated separately\n",
    "target         = ['dummy']\n",
    "\n",
    "# Checkpoints for training\n",
    "checkpoints = [30, 40, 50, 60, 70, 80, 90, 100, 110, 120, \n",
    "               130, 140, 150, 160, 170, 180, 190, 200, 210] \n",
    "               \n",
    "# Checkpoints for validation\n",
    "checkpoints_valid = [35, 65, 95, 125, 155, 185]\n",
    "\n",
    "# Checkpoints to load\n",
    "checkpoints_load = checkpoints + checkpoints_valid\n",
    "\n",
    "# Dimensionality of your data in format (D,H,W)\n",
    "# It is 1D in this case, with 678 grid points in radius\n",
    "INPUT_SIZE = [678]\n",
    "\n",
    "# Interpolation size to equite the region between the proto neutron star\n",
    "# and the stalled shock across all checkpoints\n",
    "mlin_grid_size = 200  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data pre-processing functions to:\n",
    "#   1. find proto neutron star position \n",
    "#   2. find shock position\n",
    "#   3. interpolate the region between them\n",
    "#   4. scale the data\n",
    "\n",
    "def shock_position(x, v, pturb_pgas):    \n",
    "    shock_r = 0    \n",
    "    shock_i = np.argmin(v)\n",
    "    \n",
    "    # find the local minima right before the shock to exclude \n",
    "    # the diffused shock (otherwise can contaminate Pturb/Pgas prediction)\n",
    "    minima  = np.array(sig.argrelmin(pturb_pgas[:shock_i]))[0]\n",
    "    \n",
    "    # a few extra grid points to be safe against the diffused shock\n",
    "    buffer = 2\n",
    "    \n",
    "    for j in range(len(minima)-1,-1,-1):        \n",
    "        shock_i = minima[j]-buffer\n",
    "        if pturb_pgas[shock_i]<1: break\n",
    "    shock_r = x[shock_i]\n",
    "    \n",
    "    return shock_i\n",
    "        \n",
    "def pns_radius(rho, rho_threshold=1e12):  \n",
    "    # find when density is reaching the threshold for the proto neutron star density  \n",
    "    for i in range(np.size(rho)-1,-1,-1):\n",
    "        if rho[i] >= rho_threshold:\n",
    "            pns_x = rho[i]\n",
    "            pns_ind = i\n",
    "            break\n",
    "    return pns_ind\n",
    "\n",
    "def process_var(var, v, new_r, pns_i, shock_i):\n",
    "    f = interpolate.interp1d(np.arange(np.size(v))[pns_i:shock_i+1], var[pns_i:shock_i+1])\n",
    "    return f(new_r)\n",
    "\n",
    "def interpolate_data(x, y, velocity, mlin_grid_size):\n",
    "    \n",
    "    mlin_x = np.zeros((x.shape[0], x.shape[1], mlin_grid_size))\n",
    "    mlin_y = np.zeros((y.shape[0], y.shape[1], mlin_grid_size))\n",
    "    \n",
    "    for i in range(x.shape[0]):\n",
    "        v          = velocity[i]\n",
    "        pturb_pgas = y[i,0]\n",
    "        \n",
    "        shock_i = shock_position(np.arange(np.size(v)), v, pturb_pgas)\n",
    "        pns_i   = pns_radius(x[i,0])\n",
    "        \n",
    "        new_r = np.linspace(pns_i,shock_i,mlin_grid_size)                           \n",
    "        \n",
    "        for j in range(x.shape[1]): \n",
    "            mlin_x[i,j] = process_var(x[i,j], v, new_r, pns_i,shock_i)\n",
    "            \n",
    "        mlin_y[i,0] = process_var(y[i,0], v, new_r, pns_i,shock_i)  \n",
    "            \n",
    "        print(f'Checkpoint ind {i} | Convective region from/to: {pns_i}, {shock_i} | Total size: {shock_i-pns_i}')\n",
    "\n",
    "    return mlin_x, mlin_y, pns_i,shock_i\n",
    "\n",
    "def scale_data(mlin_x, mlin_y, units, features_label):\n",
    "    # scale input fatures and target based on provided units\n",
    "    \n",
    "    print('----------')\n",
    "    for i in range(mlin_x.shape[1]):    \n",
    "        mlin_x[:,i] *= units[features_label[i]]\n",
    "        print(f'{features_label[i]:10s} min: {np.amin(mlin_x[:,i]):.3e}, max: {np.amax(mlin_x[:,i]):.3e}')   \n",
    "        \n",
    "    mlin_y[:,0] *= units[target_label[0]]\n",
    "    print(f'{target_label[0]}/Pgas min: {np.amin(mlin_y[:,0]):.3e}, max: {np.amax(mlin_y[:,0]):.3e}')   \n",
    "    print('----------')\n",
    "    \n",
    "    return mlin_x, mlin_y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Load the radial grid to find proto neutron star and shock positions\n",
    "grid =  HDF5Dataset(path           = base+'grid.h5',\n",
    "                    features_label = ['Z'],\n",
    "                    checkpoints    = [0],\n",
    "                    input_size     = INPUT_SIZE)\n",
    "r = grid.load_numpy().flatten()[:-1]\n",
    "\n",
    "# Load the Time metadata\n",
    "data_loader = HDF5Dataset(path           = path,                         \n",
    "                          features_label = ['Time'],                          \n",
    "                          checkpoints    = checkpoints_load,\n",
    "                          input_size     = INPUT_SIZE,                          \n",
    "                          shuffle        = False,\n",
    "                          train_fraction = len(checkpoints))\n",
    "\n",
    "# in miliseconds\n",
    "time = data_loader.load_numpy().flatten()\n",
    "\n",
    "# Load velocity profiles to find the shock position\n",
    "data_loader = HDF5Dataset(path           = path,                         \n",
    "                          features_label = ['u1'],                          \n",
    "                          checkpoints    = checkpoints_load,\n",
    "                          input_size     = INPUT_SIZE,                          \n",
    "                          shuffle        = False,\n",
    "                          train_fraction = len(checkpoints))\n",
    "\n",
    "velocity = data_loader.load_numpy()\n",
    "\n",
    "# Load the training and validation data\n",
    "data_loader = HDF5Dataset(path           = path,\n",
    "                          target         = target,\n",
    "                          features_label = features_label,\n",
    "                          target_label   = target_label,\n",
    "                          checkpoints    = checkpoints_load,\n",
    "                          input_size     = INPUT_SIZE,\n",
    "                          shuffle        = False,\n",
    "                          train_fraction = len(checkpoints))\n",
    "\n",
    "x, y = data_loader.load_numpy()\n",
    "\n",
    "# Calculate the target Pturb/Pgas\n",
    "for i in range(x.shape[1]):\n",
    "    if 'eos0' in features_label[i]: y[:,0] = np.divide(y[:,0],x[:,i])\n",
    "\n",
    "# Find the proto neutron star and the shock position, and interpolate\n",
    "# the convective in-between region to mlin_grid_size for each checkpoint\n",
    "mlin_x, mlin_y, pns_i, shock_i = interpolate_data(x, y, velocity, mlin_grid_size)\n",
    "\n",
    "# Scaling coefficients to equilibrate all maximums closer to 1e0\n",
    "# (leaving dependency on Pgas to be greater than others)\n",
    "units = {'u1'   : 1,\n",
    "         'mach2': 1,\n",
    "         'rho'  : 2e-12,\n",
    "         'eos0' : 2e-31,\n",
    "         'eos1' : 1e-8,\n",
    "         'eos2' : 1,\n",
    "         'eos3' : 1e-1,\n",
    "         target_label[0]:1/3}\n",
    "\n",
    "mlin_x, mlin_y = scale_data(mlin_x, mlin_y, units, features_label)\n",
    "\n",
    "loaders = data_loader.convert_to_torch([mlin_x, mlin_y])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Machine Learning model to use\n",
    "\n",
    "# Configuration of the model parameters:\n",
    "#    n_epochs  = number of epochs (need ~1000 for sensible results using PIMLTurb1D)\n",
    "#    patience  = number of epochs to run beyond convergence (not used for SmoothL1_KSLoss)\n",
    "#    min_delta = loss based convergence cut-off (not used for SmoothL1_KSLoss)\n",
    "#    lr        = learning rate\n",
    "#    min_lr    = minimum learning rate\n",
    "#    device    = device to train the model on: ['cpu', 'cuda', 'cuda:0', ...]\n",
    "#    activ     = activation function\n",
    "#    loss      = loss function\n",
    "#    ks_stop   = value of the KS (Kolmogorov-Smirnov stat) to stop training,\n",
    "#                checks both 'train' and 'valid' KS values \n",
    "#                Note: NOT the total loss value, but only the KS component!\n",
    "#    ks_frac   = contribution fraction of KS to the total loss. l1_frac = 1-ks_frac\n",
    "#    l1_scale  = amount to scale the initial L1 loss by, so that it is\n",
    "#                reduced first, with KS loss dominating later in the training\n",
    "#    ks_scale  = amount to scale the initial KS loss by (should always be = 1)\n",
    "#    l1_beta   = beta for the Smooth L1 (cutoff point for smooth portion)\n",
    "#    sigma     = sigma for the Gaussian kernel for the filtering/smoothing layer\n",
    "\n",
    "estimator = PIMLTurb1D(\n",
    "        config   = PIMLTurb1DConfig(n_epochs   = 3, \n",
    "                                    patience   = 10, \n",
    "                                    min_delta  = 1e15, \n",
    "                                    lr         = 1e-4, \n",
    "                                    min_lr     = 1e-4*1e-5, \n",
    "                                    device     = 'cpu'),\n",
    "        loaders  = loaders,\n",
    "        activ    = \"Tanhshrink\", \n",
    "        loss     = \"SmoothL1_KSLoss\",\n",
    "        ks_stop  = 0.13,\n",
    "        l1_scale = 1e3,\n",
    "        ks_scale = 1,\n",
    "        l1_beta  = 1,\n",
    "        sigma    = 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set the experiment\n",
    "training_experiment = Train(model           = estimator,\n",
    "                            backend         = tracking_backend,\n",
    "                            data_parameters = data_loader)\n",
    "\n",
    "#Train the model\n",
    "estimator = training_experiment.run()\n",
    "if mlflow: \n",
    "    tracking_backend.log_parameter(\"ks_stop\",  ks_stop)\n",
    "    tracking_backend.log_parameter(\"ks_scale\", ks_scale)\n",
    "    tracking_backend.log_parameter(\"l1_scale\", l1_scale)\n",
    "    tracking_backend.log_parameter(\"units\",    units) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plotting training results\n",
    "# Note: not done automatically since Catalyst is NOT used\n",
    "log_path       = 'train_log.txt'\n",
    "valid_log_path = 'valid_log.txt'\n",
    "delimiter      = '\\t'\n",
    "\n",
    "log = log_plot(show_log       = True,\n",
    "               log_path       = log_path, \n",
    "               valid_log_path = valid_log_path, \n",
    "               delimiter      = delimiter)\n",
    "\n",
    "log_name = 'runtime_log.html'\n",
    "log.write_html(log_name)\n",
    "tracking_backend.log_artifact(log_name)\n",
    "\n",
    "data       = np.genfromtxt(log_path, delimiter=delimiter, \n",
    "                            skip_header=1, dtype=np.float32)\n",
    "data_valid = np.genfromtxt(valid_log_path, delimiter=delimiter, \n",
    "                            skip_header=1, dtype=np.float32)\n",
    "    \n",
    "tracking_backend.log_metric('train - loss', data[-1,1])\n",
    "tracking_backend.log_metric('valid - loss', data_valid[-1,1])\n",
    "tracking_backend.log_metric('train - final epoch', data[-1,0])  \n",
    "\n",
    "#--------------------------------------------------\n",
    "# plotting individual evolutions of L1 and KS losses\n",
    "\n",
    "losses = [[f\"mean(L1_loss)\", f\"mean(L1_valid)\"],\n",
    "          [f\"mean(KS_loss)\", f\"mean(KS_valid)\"]]\n",
    "\n",
    "valid_log_path = 'valid_l1ks_log.txt'\n",
    "log_path       = 'train_l1ks_log.txt'\n",
    "\n",
    "for idx, (train_name, valid_name) in enumerate(losses):\n",
    "    if   'L1' in train_name: mark='L1'; idx = 2\n",
    "    elif 'KS' in train_name: mark='KS'; idx = 1\n",
    "    \n",
    "    log = log_plot(show_log       = True, \n",
    "                   valid_log_path = valid_log_path,\n",
    "                   log_path       = log_path,\n",
    "                   delimiter      = delimiter,\n",
    "                   train_name     = train_name, \n",
    "                   valid_name     = valid_name,\n",
    "                   train_column   = idx, \n",
    "                   valid_column   = idx, \n",
    "                   epoch_column   = None)\n",
    "    \n",
    "    log_name = f\"{train_name}.html\"\n",
    "    log.write_html(log_name)    \n",
    "    tracking_backend.log_artifact(log_name) \n",
    "    \n",
    "    data       = np.genfromtxt(log_path, delimiter=delimiter, \n",
    "                                skip_header=1, dtype=np.float32)\n",
    "    data_valid = np.genfromtxt(valid_log_path, delimiter=delimiter, \n",
    "                            skip_header=1, dtype=np.float32)\n",
    "        \n",
    "    tracking_backend.log_metric(f'train - {mark} loss', data[-1,idx])\n",
    "    tracking_backend.log_metric(f'valid - {mark} loss', data_valid[-1,idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "target_checkpoint = [205]\n",
    "\n",
    "# Load velocity profiles to find the shock position\n",
    "data_loader = HDF5Dataset(path           = path,                         \n",
    "                          features_label = ['u1'],                          \n",
    "                          checkpoints    = target_checkpoint,\n",
    "                          input_size     = INPUT_SIZE,                          \n",
    "                          shuffle        = False,\n",
    "                          train_fraction = len(checkpoints)-len(checkpoints_valid))\n",
    "\n",
    "velocity = data_loader.load_numpy()\n",
    "\n",
    "#Load the test data\n",
    "data_loader_eval = HDF5Dataset(path           = path,\n",
    "                               target         = target,\n",
    "                               features_label = features_label,\n",
    "                               target_label   = target_label,\n",
    "                               checkpoints    = target_checkpoint,\n",
    "                               input_size     = INPUT_SIZE)\n",
    "\n",
    "x, y = data_loader_eval.load_numpy()\n",
    "\n",
    "# Calculate the target Pturb/Pgas\n",
    "for i in range(x.shape[1]):\n",
    "    if 'eos0' in features_label[i]: y[:,0] = np.divide(y[:,0],x[:,i])\n",
    "        \n",
    "# Find the proto neutron star and the shock position, and interpolate\n",
    "# the convective in-between region to mlin_grid_size for each checkpoint\n",
    "mlin_x, mlin_y, pns_i, shock_i = interpolate_data(x, y, velocity, mlin_grid_size)\n",
    "\n",
    "mlin_x, mlin_y = scale_data(mlin_x, mlin_y, units, features_label)\n",
    "\n",
    "loaders = data_loader_eval.convert_to_torch([mlin_x, mlin_y])\n",
    "\n",
    "#Set the test experiment\n",
    "estimator.loaders = loaders\n",
    "data_loader_eval.input_size=[mlin_grid_size]\n",
    "data_loader_eval.batch_size=[mlin_grid_size]\n",
    "\n",
    "evaluation_experiment = Evaluate(model           = estimator,\n",
    "                                 backend         = tracking_backend,\n",
    "                                 data_parameters = data_loader_eval)\n",
    "\n",
    "# Let's make an extra plot from the evaluation\n",
    "# evaluation returns a dict, cubes = {'pred_cube':np.ndarray, 'target_cube':np.ndarray}\n",
    "cubes = evaluation_experiment.run() \n",
    "\n",
    "ax = line_plot([[np.arange(mlin_grid_size),cubes['target'][0,0]],\n",
    "                [np.arange(mlin_grid_size),cubes['predict'][0,0]]],\n",
    "               label     = ['target', 'predict'],\n",
    "               plot_type ='semilogy', \n",
    "               figsize   = (10,6))\n",
    "\n",
    "ax.set_xlabel('index')\n",
    "ax.set_ylabel(r'$P_{turb}/P_{gas}$')\n",
    "ax.set_title(f'checkpoint={target_checkpoint[0]}')\n",
    "ax.set_ylim(1e-6,1e1)\n",
    "ax.legend(loc=2)\n",
    "ax.get_figure().tight_layout()\n",
    "\n",
    "# If mlflow is active, then we want to open the last evaluation ID\n",
    "# and add this new plot to it, instead of creating a new experiment entry\n",
    "if mlflow:\n",
    "    logy_name = 'spatial_plot_logy.png'\n",
    "    ax.get_figure().savefig(logy_name)\n",
    "    \n",
    "    run_id = evaluation_experiment.run_id\n",
    "    tracking_backend.resume(run_id = run_id)\n",
    "    tracking_backend.log_artifact(logy_name)           \n",
    "    tracking_backend.end()       "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('py310')",
   "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.12"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  },
  "vscode": {
   "interpreter": {
    "hash": "55761d48b13336894241222244b506485a2c10caa3a7f51cb6df6e994d5120a6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}