plot_slices.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89725ff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import json\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import h5py\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.path as mplPath\n",
    "\n",
    "from fr3D.train.utils import setup_datasets\n",
    "from fr3D.data.utils import get_normalization\n",
    "from fr3D.models import ConvAutoencoder, ConvAutoencoderCGAN, ConvAutoencoderC\n",
    "\n",
    "tf.keras.backend.set_image_data_format('channels_last')\n",
    "\n",
    "def check_in_polygon(pts, vertices):\n",
    "    path = mplPath.Path(vertices)\n",
    "    return path.contains_points(pts)\n",
    "\n",
    "def mape_with_threshold(yp, yt, pcterror_threshold=np.inf, max_magnitude_threshold=0.0, eps=1e-7):\n",
    "    pct_errors = 100*tf.abs((yp-yt)/(eps + yt))\n",
    "    pcterror_mask = pct_errors < pcterror_threshold\n",
    "    max_magnitude_mask = tf.logical_not(tf.abs(yt) < (max_magnitude_threshold*tf.reduce_max(tf.abs(yt))))\n",
    "    filtering_indices = tf.where(tf.logical_and(pcterror_mask, max_magnitude_mask))\n",
    "    filtered_pcterrors = tf.gather_nd(pct_errors, filtering_indices)\n",
    "    return float(tf.reduce_mean(filtered_pcterrors))\n",
    "\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec780149",
   "metadata": {},
   "outputs": [],
   "source": [
    "Re = 500\n",
    "expt_variable = 'V'\n",
    "\n",
    "#dataset_path = f'/fr3D/postprocessed/annulus_64_{Re}.h5'\n",
    "dataset_path = f'/fr3D/postprocessed/annulus_64.h5'\n",
    "#dataset_path = f'/fr3D/postprocessed/annulus_64_plane.h5'\n",
    "\n",
    "#CGAN\n",
    "#experiment_config = f'/fr3D/configs/training/ConvAutoencoderCGAN_{expt_variable}.json'\n",
    "#weights_path = f'/storage/weights/ConvAutoencoderCGAN_{expt_variable}_Annulus64/ConvAutoencoderCGAN_{expt_variable}_Annulus64.h5'\n",
    "\n",
    "#ConvAutoencoderC\n",
    "experiment_config = f'/fr3D/configs/training/ConvAutoencoderC_{expt_variable}.json'\n",
    "weights_path = f'/storage/weights{Re}/ConvAutoencoderC_{expt_variable}_Annulus64/ConvAutoencoderC_{expt_variable}_Annulus64.h5'\n",
    "\n",
    "\n",
    "\n",
    "datasetf = h5py.File(dataset_path,'r')\n",
    "\n",
    "shuf_buf = 1\n",
    "config = json.load(open(experiment_config,'r'))\n",
    "train_dataset, test_dataset = setup_datasets(config, dataset_path, shuf_buf, case_names=True, evaluation=True)\n",
    "sensor_shape = train_dataset.element_spec[0][0].shape\n",
    "full_field_shape = train_dataset.element_spec[0][1].shape\n",
    "\n",
    "def get_normalization_type(node_configs):\n",
    "    normalization_spec = {'method': None}\n",
    "    for node in node_configs:\n",
    "        if node['nodetype'] == 'normalize':\n",
    "            normalization_spec = node['normalization_spec']\n",
    "            break\n",
    "    normalizer = get_normalization(**normalization_spec, batch_mode=True)\n",
    "    return normalizer\n",
    "normalizer = get_normalization_type(config['dataset']['node_configurations'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4d58b5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "generator = ConvAutoencoder(input_shape=full_field_shape[1:], **config['generator'])\n",
    "model = ConvAutoencoderCGAN(sensor_shape[1],generator,**config.get('discriminator',{}))\n",
    "model.compile(d_optimizer = tf.keras.optimizers.get(config['training']['d_optimizer']),\n",
    "            g_optimizer = tf.keras.optimizers.get(config['training']['g_optimizer']),\n",
    "            l_optimizer = tf.keras.optimizers.get(config['training']['l_optimizer']),\n",
    "            global_batch_size=config['dataset']['batch_size'])\n",
    "_ = model(next(iter(test_dataset))[0][0], autoencode=False)\n",
    "\n",
    "model.load_weights(weights_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75db0968",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ConvAutoencoderC(dense_input_units=sensor_shape[1],\n",
    "                         autoencoder_input_shape=full_field_shape[1:],\n",
    "                         **config['model'])\n",
    "loss_fn = \"mse\"#tf.keras.losses.get(config['training']['loss'])\n",
    "model.compile(l_optimizer= tf.keras.optimizers.get(config['training']['l_optimizer']),\n",
    "              loss=loss_fn,\n",
    "              optimizer = tf.keras.optimizers.get(config['training']['ae_optimizer']),\n",
    "              metrics = config['training'].get('metrics', None))\n",
    "model.load_weights(weights_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f29ef3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = iter(test_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9705851b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams['figure.dpi'] = 800\n",
    "get_new=True\n",
    "if get_new:\n",
    "    (inp, target, norm_params), case_names = next(ds)\n",
    "    pred = model(inp,autoencode=False)\n",
    "    ae_pred = model(target,autoencode=True)\n",
    "\n",
    "    targetn = normalizer.undo(target, norm_params[:,0,:])\n",
    "    predn = normalizer.undo(pred, norm_params[:,0,:])\n",
    "    ae_predn = normalizer.undo(ae_pred, norm_params[:,0,:])\n",
    "\n",
    "idx=23\n",
    "zsl=48\n",
    "\n",
    "print(f\"{case_names[idx].numpy().decode('utf-8')}\")\n",
    "\n",
    "vmax = 1.0\n",
    "vmin = 0.0\n",
    "\n",
    "stdf = 10.0\n",
    "meann = tf.reduce_mean(target[idx,:,:,zsl])\n",
    "stdn = tf.math.reduce_std(target[idx,:,:,zsl])\n",
    "vmaxn = meann+stdf*stdn\n",
    "vminn = meann-stdf*stdn\n",
    "\n",
    "Lm = 7*20/40\n",
    "coords = datasetf[case_names[idx].numpy()]['full_field_coords']\n",
    "x = np.reshape(coords[:,:,zsl,0],[-1])\n",
    "y = np.reshape(coords[:,:,zsl,1],[-1])\n",
    "\n",
    "p = np.reshape(pred[idx,:,:,zsl,0], [-1])\n",
    "pn = np.reshape(predn[idx,:,:,zsl,0], [-1])\n",
    "ap = np.reshape(ae_pred[idx,:,:,zsl,0], [-1])\n",
    "apn = np.reshape(ae_predn[idx,:,:,zsl,0], [-1])\n",
    "t = np.reshape(target[idx,:,:,zsl,0], [-1])\n",
    "tn = np.reshape(targetn[idx,:,:,zsl,0], [-1])\n",
    "\n",
    "outer_polygon = [(-3,3), (-3,-3), (10,-3), (10,3)]\n",
    "spatialmask = check_in_polygon(np.stack([x,y],-1), outer_polygon)\n",
    "spatialmasked_indices = np.where(spatialmask)[0]\n",
    "x = x[spatialmasked_indices]\n",
    "y = y[spatialmasked_indices]\n",
    "p = p[spatialmasked_indices]\n",
    "pn = pn[spatialmasked_indices]\n",
    "ap = ap[spatialmasked_indices]\n",
    "apn = apn[spatialmasked_indices]\n",
    "t = t[spatialmasked_indices]\n",
    "tn = tn[spatialmasked_indices]\n",
    "\n",
    "polygon_vertices = datasetf[case_names[idx].numpy()]['full_field_coords'][0,:,zsl]/Lm\n",
    "\n",
    "#Normalized plots\n",
    "plt.figure()\n",
    "plt.tripcolor(x/Lm,y/Lm,p, shading='gouraud', cmap='rainbow', vmax = vmax, vmin = vmin)\n",
    "plt.fill(polygon_vertices[...,0], polygon_vertices[...,1], color='purple')\n",
    "plt.colorbar()\n",
    "plt.axis('equal')\n",
    "plt.show()\n",
    "plt.close()\n",
    "\n",
    "plt.figure()\n",
    "plt.tripcolor(x/Lm,y/Lm,ap, shading='gouraud', cmap='rainbow', vmax = vmax, vmin = vmin)\n",
    "plt.fill(polygon_vertices[...,0], polygon_vertices[...,1], color='purple')\n",
    "plt.colorbar()\n",
    "plt.axis('equal')\n",
    "plt.show()\n",
    "plt.close()\n",
    "\n",
    "plt.figure()\n",
    "plt.tripcolor(x/Lm,y/Lm,t, shading='gouraud', cmap='rainbow', vmax = vmax, vmin = vmin)\n",
    "plt.fill(polygon_vertices[...,0], polygon_vertices[...,1], color='purple')\n",
    "plt.colorbar()\n",
    "plt.axis('equal')\n",
    "plt.show()\n",
    "plt.close()\n",
    "\n",
    "#Unnormalized plots\n",
    "plt.figure()\n",
    "plt.tripcolor(x/Lm,y/Lm,pn, shading='gouraud', cmap='rainbow', vmax = vmaxn, vmin = vminn)\n",
    "plt.fill(polygon_vertices[...,0], polygon_vertices[...,1], color='purple')\n",
    "plt.colorbar()\n",
    "plt.axis('equal')\n",
    "plt.show()\n",
    "plt.close()\n",
    "\n",
    "plt.figure()\n",
    "plt.tripcolor(x/Lm,y/Lm,apn, shading='gouraud', cmap='rainbow', vmax = vmaxn, vmin = vminn)\n",
    "plt.fill(polygon_vertices[...,0], polygon_vertices[...,1], color='purple')\n",
    "plt.colorbar()\n",
    "plt.axis('equal')\n",
    "plt.show()\n",
    "plt.close()\n",
    "\n",
    "plt.figure()\n",
    "plt.tripcolor(x/Lm,y/Lm,tn, shading='gouraud', cmap='rainbow', vmax = vmaxn, vmin = vminn)\n",
    "plt.fill(polygon_vertices[...,0], polygon_vertices[...,1], color='purple')\n",
    "plt.colorbar()\n",
    "plt.axis('equal')\n",
    "plt.show()\n",
    "plt.close()\n",
    "\n",
    "print(\"==Slice % err==\")\n",
    "print(mape_with_threshold(p,t,pcterror_threshold=200))\n",
    "print(mape_with_threshold(ap,t,pcterror_threshold=200))\n",
    "print(\"==Whole snapshot % err==\")\n",
    "print(mape_with_threshold(pred[idx],target[idx],pcterror_threshold=200))\n",
    "print(mape_with_threshold(ae_pred[idx],target[idx],pcterror_threshold=200))\n",
    "print(\"==Slice % err unnormalized==\")\n",
    "print(mape_with_threshold(pn,tn,pcterror_threshold=200))\n",
    "print(mape_with_threshold(apn,tn,pcterror_threshold=200))\n",
    "print(\"==Whole snapshot % err unnormalized==\")\n",
    "print(mape_with_threshold(predn[idx],targetn[idx],pcterror_threshold=200))\n",
    "print(mape_with_threshold(ae_predn[idx],targetn[idx],pcterror_threshold=200))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f801c55",
   "metadata": {},
   "outputs": [],
   "source": [
    "#spanwise averaged error plots\n",
    "plt.rcParams['figure.dpi'] = 800\n",
    "plt.rcParams['text.usetex'] = True\n",
    "plt.rcParams.update({'font.size': 18})\n",
    "get_new=False\n",
    "if get_new:\n",
    "    (inp, target, norm_params), case_names = next(ds)\n",
    "    pred = model(inp,autoencode=False)\n",
    "    ae_pred = model(target,autoencode=True)\n",
    "\n",
    "    targetn = normalizer.undo(target, norm_params[:,0,:])\n",
    "    predn = normalizer.undo(pred, norm_params[:,0,:])\n",
    "    ae_predn = normalizer.undo(ae_pred, norm_params[:,0,:])\n",
    "\n",
    "coords = datasetf[case_names[idx].numpy()]['full_field_coords']\n",
    "x = np.reshape(coords[:,:,zsl,0],[-1])\n",
    "y = np.reshape(coords[:,:,zsl,1],[-1])\n",
    "\n",
    "pn = np.reshape(predn[idx,:,:,zsl,0], [-1])\n",
    "apn = np.reshape(ae_predn[idx,:,:,zsl,0], [-1])\n",
    "tn = np.reshape(targetn[idx,:,:,zsl,0], [-1])\n",
    "\n",
    "spanwise_avg_normalized_error = lambda X: np.reshape(np.mean(np.abs(X[idx,...,0] - targetn[idx,...,0]), -1), [-1]) / np.std(targetn[idx,...,0])\n",
    "\n",
    "pnerr = spanwise_avg_normalized_error(predn)\n",
    "apnerr = spanwise_avg_normalized_error(ae_predn)\n",
    "\n",
    "outer_polygon = [(-3,3), (-3,-3), (10,-3), (10,3)]\n",
    "spatialmask = check_in_polygon(np.stack([x,y],-1), outer_polygon)\n",
    "spatialmasked_indices = np.where(spatialmask)[0]\n",
    "x = x[spatialmasked_indices]\n",
    "y = y[spatialmasked_indices]\n",
    "pnerr = pnerr[spatialmasked_indices]\n",
    "apnerr = apnerr[spatialmasked_indices]\n",
    "\n",
    "polygon_vertices = datasetf[case_names[idx].numpy()]['full_field_coords'][0,:,zsl]\n",
    "\n",
    "plt.figure()\n",
    "plt.tripcolor(x/Lm, y/Lm, pnerr, shading='gouraud', cmap='Blues', vmax = 3.0, vmin = 0.0)\n",
    "plt.fill(polygon_vertices[...,0]/Lm, polygon_vertices[...,1]/Lm, color='purple')\n",
    "#plt.colorbar()\n",
    "plt.axis('equal')\n",
    "plt.xlabel(r\"$x/L_m$\")\n",
    "plt.ylabel(r\"$y/L_m$\")\n",
    "plt.savefig(f'/storage/paper/errormap_zavg_{expt_variable}.pdf', bbox_inches='tight')\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6451e99c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from evtk import hl\n",
    "\n",
    "def export_vtk(c, **fields):\n",
    "    x = c[...,0].reshape(-1)\n",
    "    y = c[...,1].reshape(-1)\n",
    "    z = c[...,2].reshape(-1)\n",
    "    data = {}\n",
    "    for f in fields:\n",
    "        data[f] = fields[f].reshape(-1).copy()\n",
    "    hl.pointsToVTK(\"/storage/pressure3d\", x, y, z, data = data)\n",
    "    \n",
    "export_vtk(coords, pred=predn[0,...,0], ae_pred=ae_predn[0,...,0].numpy(), gt = targetn[0,...,0].numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dddb1175",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}