lift_drag_pcterror.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78ffe1bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "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_type\n",
    "from fr3D.models import ConvAutoencoder, ConvAutoencoderCGAN, ConvAutoencoderC\n",
    "from fr3D.utils import mape_with_threshold"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b72df71c",
   "metadata": {},
   "outputs": [],
   "source": [
    "Re = 500\n",
    "model_type = 'ConvAutoencoderCGAN'\n",
    "expt_variables = ['Pressure', 'U', 'V', 'W']\n",
    "\n",
    "dataset_path = f'/fr3D/postprocessed/annulus_64.h5'\n",
    "#dataset_path = f'/fr3D/postprocessed/annulus_64_plane.h5'\n",
    "\n",
    "#ConvAutoencoderC\n",
    "experiment_configs = {expt_variable:f'/fr3D/configs/training/{model_type}_{expt_variable}.json' for expt_variable in expt_variables}\n",
    "weights_paths = {expt_variable:f'/storage/weights{Re}/{model_type}_{expt_variable}_Annulus64/{model_type}_{expt_variable}_Annulus64.h5' for expt_variable in expt_variables}\n",
    "\n",
    "#experiment_configs = {expt_variable:f'/fr3D/configs/training_piv/ConvAutoencoderC_{expt_variable}.json' for expt_variable in expt_variables}\n",
    "#weights_paths = {expt_variable:f'/storage/weights{Re}piv/ConvAutoencoderC_{expt_variable}_Annulus64/ConvAutoencoderC_{expt_variable}_Annulus64.h5' for expt_variable in expt_variables}\n",
    "\n",
    "\n",
    "\n",
    "datasetf = h5py.File(dataset_path,'r')\n",
    "\n",
    "shuf_buf = 1\n",
    "\n",
    "train_datasets = {}\n",
    "test_datasets = {}\n",
    "sensor_shapes = {}\n",
    "full_field_shapes = {}\n",
    "normalizers = {}\n",
    "\n",
    "for expt_variable in experiment_configs:\n",
    "    config = json.load(open(experiment_configs[expt_variable],'r'))\n",
    "    train_datasets[expt_variable], test_datasets[expt_variable] = setup_datasets(config, dataset_path, shuf_buf, case_names=True, evaluation=True)\n",
    "    sensor_shapes[expt_variable] = train_datasets[expt_variable].element_spec[0][0].shape\n",
    "    full_field_shapes[expt_variable] = train_datasets[expt_variable].element_spec[0][1].shape\n",
    "    normalizers[expt_variable] = get_normalization_type(config['dataset']['node_configurations'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48ab5246",
   "metadata": {},
   "outputs": [],
   "source": [
    "def circular_pad(x, coords=True):\n",
    "    if coords:\n",
    "        r = np.zeros((x.shape[0]+2, x.shape[1]+2, x.shape[2]))\n",
    "    else:\n",
    "        r = np.zeros((x.shape[0]+2, x.shape[1]+2))\n",
    "    r[1:-1,1:-1] = x\n",
    "    \n",
    "    r[0,1:-1] = x[-1]\n",
    "    r[-1,1:-1] = x[0]\n",
    "    \n",
    "    if coords:\n",
    "        r[:,0,:2] = r[:,1,:2]\n",
    "        r[:,0,2] = x[0,0,2] - x[0,1,2]\n",
    "        \n",
    "        r[:,-1,:2] = r[:,-2,:2]\n",
    "        r[:,-1,2] = 2*x[0,-1,2] - x[0,-2,2]\n",
    "    return r\n",
    "\n",
    "def compute_surface_normals(ds, normalize=True):\n",
    "    r = {}\n",
    "    for shape in ds:\n",
    "        deltas = ds[shape]['full_field_coords'][1] - ds[shape]['full_field_coords'][0]\n",
    "        if normalize:\n",
    "            normals = deltas / np.linalg.norm(deltas,axis=-1,keepdims=True)\n",
    "        else:\n",
    "            normals = deltas\n",
    "        r[shape] = tf.convert_to_tensor(normals)\n",
    "    return r\n",
    "\n",
    "@tf.function\n",
    "def surface_integrate(vertex_values, vertex_coords):\n",
    "    \n",
    "    padded_coords = tf.concat([vertex_coords, vertex_coords[:1]],0)\n",
    "    padded_values = tf.concat([vertex_values, vertex_values[:1]],0)\n",
    "    \n",
    "    cell_mean_values = tf.nn.conv2d(\n",
    "        tf.convert_to_tensor([[padded_values]]),\n",
    "        tf.constant([[[[0.25]],[[0.25]]],[[[0.25]],[[0.25]]]]),\n",
    "        1, \"VALID\", data_format=\"NCHW\"\n",
    "    )[0,0]\n",
    "    \n",
    "    z_direction_vector = padded_coords[:-1, 1:] - padded_coords[:-1, :-1]\n",
    "    theta_direction_vector = -(padded_coords[1:,:-1] - padded_coords[:-1,:-1])\n",
    "    cell_normals = tf.linalg.cross(theta_direction_vector, z_direction_vector)\n",
    "    cell_areas = tf.linalg.norm(cell_normals, axis=-1)\n",
    "    \n",
    "    return tf.einsum('tzn,tz->n', cell_normals, cell_mean_values)#/(0.5*1.0*(1.0**2)*tf.reduce_sum(cell_areas))\n",
    "\n",
    "@tf.function\n",
    "def surface_area(vertex_coords):\n",
    "    \n",
    "    padded_coords = tf.concat([vertex_coords, vertex_coords[:1]],0)\n",
    "    z_direction_vector = padded_coords[:-1, 1:] - padded_coords[:-1, :-1]\n",
    "    theta_direction_vector = -(padded_coords[1:,:-1] - padded_coords[:-1,:-1])\n",
    "    cell_normals = tf.linalg.cross(theta_direction_vector, z_direction_vector)\n",
    "    cell_areas = tf.linalg.norm(cell_normals, axis=-1)\n",
    "    \n",
    "    return tf.reduce_sum(cell_areas)\n",
    "\n",
    "@tf.function\n",
    "def compute_shear_stresses(velocities, surface_normals, nu=1.0):\n",
    "    #get surface tangent component of velocity    \n",
    "    R = tf.constant([[0,-1,0],[1,0,0],[0,0,1]],dtype=velocities.dtype)\n",
    "    surface_normals_norm = tf.linalg.norm(surface_normals, axis=-1, keepdims=True)\n",
    "    surface_unit_tangents = tf.einsum('ij,tzj->tzi', R, surface_normals/surface_normals_norm)\n",
    "    velocity_tangent_component = tf.einsum('tzi,tzi->tz', surface_unit_tangents, velocities[1])\n",
    "    \n",
    "    #compute shear stresses from du/dn\n",
    "    tau = nu*velocity_tangent_component/(2*surface_normals_norm[...,0]**2.0)\n",
    "    return tau\n",
    "\n",
    "@tf.function\n",
    "def compute_friction_forces(shear_stresses, vertex_coords):\n",
    "    Fnormal = surface_integrate(shear_stresses, vertex_coords)\n",
    "    return tf.convert_to_tensor([Fnormal[1], -Fnormal[0], Fnormal[2]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bc058b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "models = {}\n",
    "\n",
    "for v in expt_variables:\n",
    "    config = json.load(open(experiment_configs[v],'r'))\n",
    "    loss_fn = \"mse\"#tf.keras.losses.get(config['training']['loss'])\n",
    "    \n",
    "    if model_type == \"ConvAutoencoderC\":\n",
    "        model = ConvAutoencoderC(dense_input_units=sensor_shapes[v][1],\n",
    "                             autoencoder_input_shape=full_field_shapes[v][1:],\n",
    "                             **config['model'])\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",
    "    elif model_type == \"ConvAutoencoderCGAN\":\n",
    "        model = ConvAutoencoderCGAN(\n",
    "            sensor_shapes[v][1], \n",
    "            ConvAutoencoder(input_shape=full_field_shapes[v][1:], **config['generator']),\n",
    "            **config.get('discriminator', {})\n",
    "        )\n",
    "        model.compile(\n",
    "            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",
    "        )\n",
    "        model.build(full_field_shapes[v])\n",
    "    model.load_weights(weights_paths[v])\n",
    "    models[v] = model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdd8c4ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def _surface_integrate_wrap(x):\n",
    "    return surface_integrate(x[0], x[1])\n",
    "\n",
    "@tf.function\n",
    "def _friction_forces(x):\n",
    "    velocities, surface_normals, coords = x\n",
    "    nu = 0.01\n",
    "    tau = compute_shear_stresses(velocities, surface_normals, nu)\n",
    "    friction_forces = compute_friction_forces(tau, coords)\n",
    "    return friction_forces\n",
    "\n",
    "@tf.function\n",
    "def get_predictions(model, inp, target):\n",
    "    return model(inp, autoencode=False), model(target, autoencode=True)\n",
    "\n",
    "def undo_normalization(normalizer, norm_param, pred, ae_pred, target):\n",
    "    ntarget = normalizer.undo(target, norm_param[:,0,:])\n",
    "    npred = normalizer.undo(pred, norm_param[:,0,:])\n",
    "    nae_pred = normalizer.undo(ae_pred, norm_param[:,0,:])\n",
    "    return npred, nae_pred, ntarget\n",
    "\n",
    "@tf.function\n",
    "def check_in_box(coords, bounds_lower, bounds_upper):\n",
    "    newshape = tf.concat([tf.ones((tf.rank(coords)-1,), dtype=tf.int32), tf.shape(bounds_lower)[:1]],0)\n",
    "    bounds_lower_r = tf.reshape(bounds_lower, newshape)\n",
    "    bounds_upper_r = tf.reshape(bounds_upper, newshape)\n",
    "    lower_cond = tf.reduce_all(coords > bounds_lower, axis=-1)\n",
    "    upper_cond = tf.reduce_all(coords < bounds_upper, axis=-1)\n",
    "    return tf.logical_and(lower_cond, upper_cond)\n",
    "\n",
    "def compute_metrics(coords, v, mapes, umapes, mses, umses, pred, target, npred, ntarget):\n",
    "    lower_box_bounds = tf.constant([-3.0,-3.0,-0.1])\n",
    "    upper_box_bounds = tf.constant([8.0,3.0,10.1])\n",
    "    in_box = check_in_box(coords, lower_box_bounds, upper_box_bounds)\n",
    "    in_box_r = tf.reshape(in_box, tf.concat([tf.shape(in_box), tf.ones((tf.rank(p_npred) - tf.rank(in_box),), tf.int32)], 0))\n",
    "    in_box_rf = tf.cast(in_box_r, tf.float32)\n",
    "    \n",
    "    _lf =tf.keras.losses.MeanSquaredError(reduction=tf.keras.losses.Reduction.NONE)\n",
    "    \n",
    "    mapes[v].append(\n",
    "        mape_with_threshold(npred, ntarget, pcterror_threshold=100.0, max_magnitude_threshold=0.03, sample_weights=in_box_rf, axis=tf.range(1, tf.rank(in_box_rf)))\n",
    "    )\n",
    "    umapes[v].append(\n",
    "        mape_with_threshold(pred, target, pcterror_threshold=100.0, max_magnitude_threshold=0.03, sample_weights=in_box_rf, axis=tf.range(1, tf.rank(in_box_rf)))\n",
    "    )\n",
    "    mses[v].append(tf.reduce_mean(_lf(npred, ntarget, in_box_rf), axis=tf.range(1, tf.rank(in_box_rf)-1)))\n",
    "    umses[v].append(tf.reduce_mean(_lf(pred, target, in_box_rf), axis=tf.range(1, tf.rank(in_box_rf)-1)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acce48c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_iterators = {v: iter(test_datasets[v]) for v in test_datasets}\n",
    "surface_normals_map = compute_surface_normals(datasetf, normalize=False)\n",
    "\n",
    "pred_forces = []\n",
    "ae_pred_forces = []\n",
    "target_forces = []\n",
    "\n",
    "pred_mapes = {k:[] for k in dataset_iterators.keys()}\n",
    "pred_unnormalized_mapes = {k:[] for k in dataset_iterators.keys()}\n",
    "pred_mses = {k:[] for k in dataset_iterators.keys()}\n",
    "pred_unnormalized_mses = {k:[] for k in dataset_iterators.keys()}\n",
    "\n",
    "ae_pred_mapes = {k:[] for k in dataset_iterators.keys()}\n",
    "ae_pred_unnormalized_mapes = {k:[] for k in dataset_iterators.keys()}\n",
    "ae_pred_mses = {k:[] for k in dataset_iterators.keys()}\n",
    "ae_pred_unnormalized_mses = {k:[] for k in dataset_iterators.keys()}\n",
    "\n",
    "for pdata, udata, vdata, wdata in zip(*dataset_iterators.values()):\n",
    "    (p_inp, p_target, p_norm_param), case_name = pdata\n",
    "    (u_inp, u_target, u_norm_param), case_name = udata\n",
    "    (v_inp, v_target, v_norm_param), case_name = vdata\n",
    "    (w_inp, w_target, w_norm_param), case_name = wdata\n",
    "    \n",
    "    coords = tf.convert_to_tensor(np.stack([datasetf[c.decode()]['full_field_coords'] for c in case_name.numpy()],0))\n",
    "    \n",
    "    p_pred, p_ae_pred = get_predictions(models['Pressure'], p_inp, p_target)\n",
    "    u_pred, u_ae_pred = get_predictions(models['U'], u_inp, u_target)\n",
    "    v_pred, v_ae_pred = get_predictions(models['V'], v_inp, v_target)\n",
    "    w_pred, w_ae_pred = get_predictions(models['W'], w_inp, w_target)\n",
    "    \n",
    "    p_npred, p_nae_pred, p_ntarget = undo_normalization(normalizers['Pressure'], p_norm_param, p_pred, p_ae_pred, p_target)\n",
    "    u_npred, u_nae_pred, u_ntarget = undo_normalization(normalizers['U'], u_norm_param, u_pred, u_ae_pred, u_target)\n",
    "    v_npred, v_nae_pred, v_ntarget = undo_normalization(normalizers['V'], v_norm_param, v_pred, v_ae_pred, v_target)\n",
    "    w_npred, w_nae_pred, w_ntarget = undo_normalization(normalizers['W'], w_norm_param, w_pred, w_ae_pred, w_target)\n",
    "    \n",
    "    pred_velocities = tf.concat([u_npred, v_npred, w_pred], -1)\n",
    "    pred_pressures = p_npred[...,0]\n",
    "    \n",
    "    ae_pred_velocities = tf.concat([u_nae_pred, v_nae_pred, w_nae_pred], -1)\n",
    "    ae_pred_pressures = p_nae_pred[...,0]\n",
    "    \n",
    "    target_velocities = tf.concat([u_ntarget, v_ntarget, w_ntarget], -1)\n",
    "    target_pressures = p_ntarget[...,0]\n",
    "    \n",
    "    surface_coords = tf.stack([datasetf[c.decode()]['full_field_coords'][0,...] for c in case_name.numpy()], 0)\n",
    "    surface_normals = tf.stack([surface_normals_map[c.decode()] for c in case_name.numpy()], 0)\n",
    "    \n",
    "    target_pressure_forces = tf.map_fn(_surface_integrate_wrap, (target_pressures[:,0], surface_coords), fn_output_signature=target_pressures.dtype)\n",
    "    pred_pressure_forces = tf.map_fn(_surface_integrate_wrap, (pred_pressures[:,0], surface_coords), fn_output_signature=target_pressures.dtype)\n",
    "    ae_pred_pressure_forces = tf.map_fn(_surface_integrate_wrap, (ae_pred_pressures[:,0], surface_coords), fn_output_signature=target_pressures.dtype)\n",
    "    \n",
    "    target_friction_forces = tf.map_fn(_friction_forces, (target_velocities, surface_normals, surface_coords), fn_output_signature=target_velocities.dtype)\n",
    "    pred_friction_forces = tf.map_fn(_friction_forces, (pred_velocities, surface_normals, surface_coords), fn_output_signature=target_velocities.dtype)\n",
    "    ae_pred_friction_forces = tf.map_fn(_friction_forces, (ae_pred_velocities, surface_normals, surface_coords), fn_output_signature=target_velocities.dtype)\n",
    "    \n",
    "    target_forces.append(target_pressure_forces + target_friction_forces)\n",
    "    pred_forces.append(pred_pressure_forces + pred_friction_forces)\n",
    "    ae_pred_forces.append(ae_pred_pressure_forces + ae_pred_friction_forces)\n",
    "    \n",
    "    compute_metrics(coords, 'Pressure', pred_mapes, pred_unnormalized_mapes, pred_mses, pred_unnormalized_mses,\n",
    "                       p_pred, p_target, p_npred, p_ntarget)\n",
    "    \n",
    "    compute_metrics(coords, 'Pressure', ae_pred_mapes, ae_pred_unnormalized_mapes, ae_pred_mses, ae_pred_unnormalized_mses,\n",
    "                       p_ae_pred, p_target, p_nae_pred, p_ntarget)\n",
    "    \n",
    "    compute_metrics(coords, 'U', pred_mapes, pred_unnormalized_mapes, pred_mses, pred_unnormalized_mses,\n",
    "                       u_pred, u_target, u_npred, u_ntarget)\n",
    "    \n",
    "    compute_metrics(coords, 'U', ae_pred_mapes, ae_pred_unnormalized_mapes, ae_pred_mses, ae_pred_unnormalized_mses,\n",
    "                       u_ae_pred, u_target, u_nae_pred, u_ntarget)\n",
    "    \n",
    "    compute_metrics(coords, 'V', pred_mapes, pred_unnormalized_mapes, pred_mses, pred_unnormalized_mses,\n",
    "                       v_pred, v_target, v_npred, v_ntarget)\n",
    "    \n",
    "    compute_metrics(coords, 'V', ae_pred_mapes, ae_pred_unnormalized_mapes, ae_pred_mses, ae_pred_unnormalized_mses,\n",
    "                       v_ae_pred, v_target, v_nae_pred, v_ntarget)\n",
    "    \n",
    "    compute_metrics(coords, 'W', pred_mapes, pred_unnormalized_mapes, pred_mses, pred_unnormalized_mses,\n",
    "                       w_pred, w_target, w_npred, w_ntarget)\n",
    "    \n",
    "    compute_metrics(coords, 'W', ae_pred_mapes, ae_pred_unnormalized_mapes, ae_pred_mses, ae_pred_unnormalized_mses,\n",
    "                       w_ae_pred, w_target, w_nae_pred, w_ntarget)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9c6eb58",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_mapes = {k:tf.concat(pred_mapes[k],0) for k in dataset_iterators.keys()}\n",
    "pred_unnormalized_mapes = {k:tf.concat(pred_unnormalized_mapes[k],0) for k in dataset_iterators.keys()}\n",
    "pred_mses = {k:tf.concat(pred_mses[k],0) for k in dataset_iterators.keys()}\n",
    "pred_unnormalized_mses = {k:tf.concat(pred_unnormalized_mses[k],0) for k in dataset_iterators.keys()}\n",
    "\n",
    "ae_pred_mapes = {k:tf.concat(ae_pred_mapes[k],0) for k in dataset_iterators.keys()}\n",
    "ae_pred_unnormalized_mapes = {k:tf.concat(ae_pred_unnormalized_mapes[k],0) for k in dataset_iterators.keys()}\n",
    "ae_pred_mses = {k:tf.concat(ae_pred_mses[k],0) for k in dataset_iterators.keys()}\n",
    "ae_pred_unnormalized_mses = {k:tf.concat(ae_pred_unnormalized_mses[k],0) for k in dataset_iterators.keys()}\n",
    "\n",
    "pred_forces = tf.concat(pred_forces,0)\n",
    "ae_pred_forces = tf.concat(ae_pred_forces,0)\n",
    "target_forces = tf.concat(target_forces,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8c0b30c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_metrics(v, s, mapes, unnormalized_mapes, mses, unnormalized_mses):\n",
    "    print(f\">> {v} {s}\")\n",
    "    print(f'MAPE: {tf.reduce_mean(mapes[v]).numpy()}')\n",
    "    print(f'Unnormalized MAPE: {tf.reduce_mean(unnormalized_mapes[v]).numpy()}')\n",
    "    print(f'MSE: {tf.reduce_mean(mses[v]).numpy()}')\n",
    "    print(f'Unnormalized MSE: {tf.reduce_mean(unnormalized_mses[v]).numpy()}')\n",
    "    print('')\n",
    "\n",
    "print('===Overall metrics===')\n",
    "\n",
    "print_metrics('Pressure', 'Pred', pred_mapes, pred_unnormalized_mapes, pred_mses, pred_unnormalized_mses)\n",
    "print_metrics('Pressure', 'AE Pred', ae_pred_mapes, ae_pred_unnormalized_mapes, ae_pred_mses, ae_pred_unnormalized_mses)\n",
    "\n",
    "print_metrics('U', 'Pred', pred_mapes, pred_unnormalized_mapes, pred_mses, pred_unnormalized_mses)\n",
    "print_metrics('U', 'AE Pred', ae_pred_mapes, ae_pred_unnormalized_mapes, ae_pred_mses, ae_pred_unnormalized_mses)\n",
    "\n",
    "print_metrics('V', 'Pred', pred_mapes, pred_unnormalized_mapes, pred_mses, pred_unnormalized_mses)\n",
    "print_metrics('V', 'AE Pred', ae_pred_mapes, ae_pred_unnormalized_mapes, ae_pred_mses, ae_pred_unnormalized_mses)\n",
    "\n",
    "print_metrics('W', 'Pred', pred_mapes, pred_unnormalized_mapes, pred_mses, pred_unnormalized_mses)\n",
    "print_metrics('W', 'AE Pred', ae_pred_mapes, ae_pred_unnormalized_mapes, ae_pred_mses, ae_pred_unnormalized_mses)\n",
    "\n",
    "print('>> Drag')\n",
    "print(f'MAPE: {mape_with_threshold(pred_forces[:,0], target_forces[:,0], 200)}')\n",
    "print(f'MSE: {tf.reduce_mean((pred_forces[:,0]-target_forces[:,0])**2)}')\n",
    "\n",
    "print('>> Lift')\n",
    "print(f'MAPE: {mape_with_threshold(pred_forces[:,1], target_forces[:,1], 200, max_magnitude_threshold=0.1)}')\n",
    "print(f'MSE: {tf.reduce_mean((pred_forces[:,1]-target_forces[:,1])**2)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d078c1de",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.pylab as pylab\n",
    "\n",
    "params = {'legend.fontsize': 'xx-large',\n",
    "          'figure.figsize': (7, 5),\n",
    "         'axes.labelsize': 'xx-large',\n",
    "         'axes.titlesize':'xx-large',\n",
    "         'xtick.labelsize':'xx-large',\n",
    "         'ytick.labelsize':'xx-large'}\n",
    "pylab.rcParams.update(params)\n",
    "\n",
    "#TODO: automatically grab test geometry names\n",
    "test_geometry_names = ('shape_34', 'shape_358', 'shape_74', 'shape_22', 'shape_183', 'shape_102', 'shape_326', 'shape_16', 'shape_107', 'shape_163', 'shape_76', 'shape_95', 'shape_52', 'shape_75', 'shape_67', 'shape_159', 'shape_343', 'shape_46', 'shape_147', 'shape_86', 'shape_325')\n",
    "for shape_name in test_geometry_names:\n",
    "    start_idx = 801*test_geometry_names.index(shape_name)\n",
    "    end_idx = 801*(test_geometry_names.index(shape_name)+1)\n",
    "\n",
    "    tstart = 20.0\n",
    "    tend = 60.0\n",
    "    Lm = 3.5\n",
    "    Uinf = 1.0\n",
    "    rho = 1.0\n",
    "    S = surface_area(datasetf[shape_name]['full_field_coords'][0])\n",
    "    norm_factor = 0.5*rho*(Uinf**2)*S\n",
    "    taun = Uinf*np.linspace(tstart, tend, end_idx - start_idx)/Lm\n",
    "\n",
    "    plt.figure()\n",
    "    pred_Cd = pred_forces[start_idx:end_idx,0]/norm_factor\n",
    "    target_Cd = target_forces[start_idx:end_idx,0]/norm_factor\n",
    "    plt.plot(taun,pred_Cd, label='Predicted')\n",
    "    plt.plot(taun,target_Cd, label='Target')\n",
    "    #plt.legend()\n",
    "    #plt.ylim(0.0, tf.reduce_max(target_Cd))\n",
    "    plt.xlabel(r'$\\tau^*$')\n",
    "    plt.ylabel(r'$C_D$')\n",
    "    plt.savefig(f'{shape_name}_Cd.pdf', bbox_inches='tight')\n",
    "    plt.show()\n",
    "    plt.close()\n",
    "\n",
    "    plt.figure()\n",
    "    pred_Cl = pred_forces[start_idx:end_idx,1]/norm_factor\n",
    "    target_Cl = target_forces[start_idx:end_idx,1]/norm_factor\n",
    "    plt.plot(taun, pred_Cl, label='Predicted')\n",
    "    plt.plot(taun, target_Cl, label='Target')\n",
    "    #plt.legend()\n",
    "    plt.xlabel(r'$\\tau^*$')\n",
    "    plt.ylabel(r'$C_L$')\n",
    "    plt.savefig(f'{shape_name}_Cl.pdf', bbox_inches='tight')\n",
    "    plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65e481dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#export legend separately\n",
    "#https://stackoverflow.com/questions/4534480/get-legend-as-a-separate-picture-in-matplotlib\n",
    "def export_legend(legend, filename=\"legend.pdf\"):\n",
    "    fig  = legend.figure\n",
    "    fig.canvas.draw()\n",
    "    bbox  = legend.get_window_extent().transformed(fig.dpi_scale_trans.inverted())\n",
    "    fig.savefig(filename, dpi=\"figure\", bbox_inches=bbox)\n",
    "    \n",
    "plt.figure()\n",
    "plt.plot(taun, pred_Cl, label='Predicted')\n",
    "plt.plot(taun, target_Cl, label='Target')\n",
    "pltlegend = plt.legend(ncol=2, bbox_to_anchor=(11.0, 11.0))\n",
    "export_legend(pltlegend)\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d8ab9d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(norm_factor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1baae687",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4335f996",
   "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
}