pyvista unstructured mesh

pyvista-unstructured-mesh.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from pyvista import set_plot_theme\n",
    "set_plot_theme('document')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import netCDF4 as nc\n",
    "import numpy as np\n",
    "import pyvista as pv\n",
    "import vtk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Approximate radius of the Earth\n",
    "RADIUS = 6371.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_unstructured(fname, node_face, node_x, node_y, radius):\n",
    "    with nc.Dataset(fname) as ds:\n",
    "        node_face = ds.variables[node_face][:].data\n",
    "        node_x = ds.variables[node_x][:].data\n",
    "        node_y = ds.variables[node_y][:].data\n",
    "    \n",
    "    # account for start_index = 1\n",
    "    node_x = np.concatenate(([0], node_x))\n",
    "    node_y = np.concatenate(([0], node_y))\n",
    "    \n",
    "    # convert lat/lon to cartesian coordinates\n",
    "    node_face_x = node_x[node_face]\n",
    "    node_face_y = 90.0 - node_y[node_face]\n",
    "    \n",
    "    node_face_x_rad = np.radians(node_face_x)\n",
    "    node_face_y_rad = np.radians(node_face_y)\n",
    "    \n",
    "    x = radius * np.sin(node_face_y_rad) * np.cos(node_face_x_rad)\n",
    "    y = radius * np.sin(node_face_y_rad) * np.sin(node_face_x_rad)\n",
    "    z = radius * np.cos(node_face_y_rad)\n",
    "    \n",
    "    # create unstructured grid\n",
    "    points = np.vstack((np.ravel(x), np.ravel(y), np.ravel(z))).T\n",
    "    N = points.shape[0]\n",
    "    print(N)\n",
    "    cell_type = np.broadcast_to(np.array([vtk.VTK_QUAD], np.uint8), (N // 4,))\n",
    "    cells = np.ravel(np.hstack((np.broadcast_to(np.array([4], np.int8), (N // 4, 1)),\n",
    "                                np.arange(0, N).reshape((-1, 4)))))\n",
    "    offset = np.arange(0, cells.shape[0], 4 + 1)\n",
    "    \n",
    "    grid = pv.UnstructuredGrid(offset, cells, cell_type, points)\n",
    "    \n",
    "    # add some synthetic data to the cells\n",
    "    cell_array = np.ones((grid.n_cells))\n",
    "    NC = N // 4\n",
    "    NCP = NC // 6\n",
    "    for i in range(6):\n",
    "        cell_array[NCP*i:NCP*(i+1)] = i\n",
    "    grid.cell_arrays[\"face values\"] = cell_array\n",
    "    \n",
    "    return grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = \"/home/bill/pyvista/data/mesh_C4.nc\"\n",
    "node_face = \"example_C4_face_nodes\"\n",
    "node_x = \"example_C4_node_x\"\n",
    "node_y = \"example_C4_node_y\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = \"/home/bill/pyvista/data/mesh_C12.nc\"\n",
    "node_face = \"dynamics_face_nodes\"\n",
    "node_x = \"dynamics_node_x\"\n",
    "node_y = \"dynamics_node_y\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = \"/home/bill/pyvista/data/mesh_C24.nc\"\n",
    "node_face = \"dynamics_face_nodes\"\n",
    "node_x = \"dynamics_node_x\"\n",
    "node_y = \"dynamics_node_y\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = \"/home/bill/pyvista/data/mesh_C48.nc\"\n",
    "node_face = \"dynamics_face_nodes\"\n",
    "node_x = \"dynamics_node_x\"\n",
    "node_y = \"dynamics_node_y\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = \"/home/bill/pyvista/data/mesh_C96.nc\"\n",
    "node_face = \"dynamics_face_nodes\"\n",
    "node_x = \"dynamics_node_x\"\n",
    "node_y = \"dynamics_node_y\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = \"/home/bill/pyvista/data/mesh_C1048.nc\"\n",
    "node_face = \"dynamics_face_nodes\"\n",
    "node_x = \"dynamics_node_x\"\n",
    "node_y = \"dynamics_node_y\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = to_unstructured(fname, node_face, node_x, node_y, RADIUS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = pv.PlotterITK()\n",
    "p.add_mesh(grid)\n",
    "p.add_points(grid.points, color=\"red\")\n",
    "p.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh_fname = \"unstructured_mesh.vtk\"\n",
    "grid.save(mesh_fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l *.vtk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = pv.read(mesh_fname)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = pv.PlotterITK()\n",
    "p.add_mesh(mesh)\n",
    "p.add_points(mesh.points, color=\"red\")\n",
    "p.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}