test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-10-08 22:28:14.179047: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2024-10-08 22:28:14.179107: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import os\n",
    "import json\n",
    "import functools\n",
    "import csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _parse(proto, meta):\n",
    "  \"\"\"Parses a trajectory from tf.Example.\"\"\"\n",
    "  feature_lists = {k: tf.io.VarLenFeature(tf.string)\n",
    "                   for k in meta['field_names']}\n",
    "  features = tf.io.parse_single_example(proto, feature_lists)\n",
    "  out = {}\n",
    "  for key, field in meta['features'].items():\n",
    "    data = tf.io.decode_raw(features[key].values, getattr(tf, field['dtype']))\n",
    "    data = tf.reshape(data, field['shape'])\n",
    "    if field['type'] == 'static':\n",
    "      data = tf.tile(data, [meta['trajectory_length'], 1, 1])\n",
    "    elif field['type'] == 'dynamic_varlen':\n",
    "      length = tf.io.decode_raw(features['length_'+key].values, tf.int32)\n",
    "      length = tf.reshape(length, [-1])\n",
    "      data = tf.RaggedTensor.from_row_lengths(data, row_lengths=length)\n",
    "    elif field['type'] != 'dynamic':\n",
    "      raise ValueError('invalid data format')\n",
    "    out[key] = data\n",
    "  return out\n",
    "\n",
    "def load_dataset(path, split):\n",
    "  \"\"\"Load dataset.\"\"\"\n",
    "  with open(os.path.join(path, 'meta.json'), 'r') as fp:\n",
    "    meta = json.loads(fp.read())\n",
    "  ds = tf.data.TFRecordDataset(os.path.join(path, split+'.tfrecord'))\n",
    "  ds = ds.map(functools.partial(_parse, meta=meta), num_parallel_calls=8)\n",
    "  ds = ds.prefetch(1)\n",
    "  return ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = load_dataset(\"./DATA/cylinder_flow\", 'test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'tensorflow.python.data.ops.dataset_ops.PrefetchDataset'>\n",
      "{'cells': TensorSpec(shape=(600, None, 3), dtype=tf.int32, name=None), 'mesh_pos': TensorSpec(shape=(600, None, 2), dtype=tf.float32, name=None), 'node_type': TensorSpec(shape=(600, None, 1), dtype=tf.int32, name=None), 'velocity': TensorSpec(shape=(600, None, 2), dtype=tf.float32, name=None), 'pressure': TensorSpec(shape=(600, None, 1), dtype=tf.float32, name=None)}\n",
      "dict_keys(['cells', 'mesh_pos', 'node_type', 'velocity', 'pressure'])\n",
      "600 <class 'numpy.ndarray'>\n",
      "600 <class 'numpy.ndarray'>\n",
      "600 <class 'numpy.ndarray'>\n",
      "600 <class 'numpy.ndarray'>\n",
      "600 <class 'numpy.ndarray'>\n"
     ]
    }
   ],
   "source": [
    "print(type(ds))\n",
    "print(ds.element_spec)\n",
    "element = next(iter(ds.take(1).as_numpy_iterator()))\n",
    "print(element.keys())\n",
    "\n",
    "csv_file = 'test.csv'\n",
    "\n",
    "for i in element.keys():\n",
    "    print(len(element[i]), type(element[i]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'tensorflow.python.data.ops.dataset_ops.PrefetchDataset'>\n",
      "dict_keys(['cells', 'mesh_pos', 'node_type', 'velocity', 'pressure'])\n",
      "<class 'dict'>\n",
      "(600, 3612, 3)\n",
      "(600, 1923, 2)\n",
      "(600, 1923, 1)\n",
      "(600, 1923, 2)\n",
      "(600, 1923, 1)\n",
      "tf.Tensor(3.1386888, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "print(type(ds))\n",
    "element = next(iter(ds))\n",
    "\n",
    "print(element.keys())\n",
    "print(type(element))\n",
    "for name in element.keys():\n",
    "    \n",
    "    print(element[name].shape)\n",
    "\n",
    "print(element[\"pressure\"][0, 1, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "10\n",
      "20\n",
      "30\n",
      "40\n",
      "50\n",
      "60\n",
      "70\n",
      "80\n",
      "90\n",
      "100\n",
      "110\n",
      "120\n",
      "130\n",
      "140\n",
      "150\n",
      "160\n",
      "170\n",
      "180\n",
      "190\n",
      "200\n",
      "210\n",
      "220\n",
      "230\n",
      "240\n",
      "250\n",
      "260\n",
      "270\n",
      "280\n",
      "290\n",
      "300\n",
      "310\n",
      "320\n",
      "330\n",
      "340\n",
      "350\n",
      "360\n",
      "370\n",
      "380\n",
      "390\n",
      "400\n",
      "410\n",
      "420\n",
      "430\n",
      "440\n",
      "450\n",
      "460\n",
      "470\n",
      "480\n",
      "490\n",
      "500\n",
      "510\n",
      "520\n",
      "530\n",
      "540\n",
      "550\n",
      "560\n",
      "570\n",
      "580\n",
      "590\n"
     ]
    }
   ],
   "source": [
    "num_elements = 1876\n",
    "num_iter = 600\n",
    "\n",
    "csv_name = './DATA/cylinder_flow/csv/data_'\n",
    "\n",
    "header = [\"mesh_pos_x\", \"mesh_pos_y\", \"velocity_x\", \"velocity_y\", \"pressure\"]\n",
    "\n",
    "for i in range(num_iter):\n",
    "    if i % 10 != 0:\n",
    "        continue\n",
    "    file = csv_name + f'{i:03}' + '.csv'\n",
    "    rows = []\n",
    "    for j in range(num_elements):\n",
    "        row = []\n",
    "        row.append(float(element[\"mesh_pos\"][i, j, 0]))\n",
    "        row.append(float(element[\"mesh_pos\"][i, j, 1]))\n",
    "        row.append(float(element[\"velocity\"][i, j, 0]))\n",
    "        row.append(float(element[\"velocity\"][i, j, 1]))\n",
    "        row.append(float(element[\"pressure\"][i, j, 0]))\n",
    "        rows.append(row)\n",
    "    \n",
    "    with open(file, mode=\"w\", newline=\"\") as file:\n",
    "        writer = csv.writer(file)\n",
    "        writer.writerow(header)  # Write the header row\n",
    "        writer.writerows(rows)   # Write all the data rows\n",
    "    print(i)\n",
    "        \n",
    "        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}