{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "31a94bbd-34ef-404b-a9be-4f3a216a8c3c",
   "metadata": {},
   "source": [
    "# Element distribution interpolation\n",
    "\n",
    "In this example we show how to perform interpolations that keep the same element structure.\n",
    "\n",
    "Two examples of this are:\n",
    "\n",
    "1. Mapping the GLL points of each element into a different distribution.\n",
    "2. Perfomring P refinement, i.e., map to points of the same distribution but a different degree."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9ad6337-6ba4-47c7-b687-0bf119a5b637",
   "metadata": {},
   "source": [
    "#### Import general modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8b4ca2be-46e8-443b-ab44-6d1a60abfec4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required modules\n",
    "from mpi4py import MPI #equivalent to the use of MPI_init() in C\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Get mpi info\n",
    "comm = MPI.COMM_WORLD"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1702ba96-7854-4a2f-89d2-7f2b2aaa61b1",
   "metadata": {},
   "source": [
    "#### Import modules from pynek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "09a4efe8-3557-4ff7-850e-0a26c8937bd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pysemtools.io.ppymech.neksuite import pynekread, pynekwrite\n",
    "from pysemtools.datatypes.msh import Mesh\n",
    "from pysemtools.datatypes.field import FieldRegistry"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63764358",
   "metadata": {},
   "source": [
    "## Read the data and build objects\n",
    "\n",
    "In this instance, we create connectivity for the mesh object, given that we wish to use direct stiffness summation to reduce discontinuities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "116a2e64",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-09-11 21:48:29,145 - Mesh - INFO - Initializing empty Mesh object.\n",
      "2024-09-11 21:48:29,147 - Field - INFO - Initializing empty Field object\n",
      "2024-09-11 21:48:29,147 - pynekread - INFO - Reading file: ../data/tc_channel0.f00001\n",
      "2024-09-11 21:48:29,163 - Mesh - INFO - Initializing Mesh object from x,y,z ndarrays.\n",
      "2024-09-11 21:48:29,165 - Mesh - INFO - Initializing common attributes.\n",
      "2024-09-11 21:48:29,167 - Mesh - INFO - Getting vertices\n",
      "2024-09-11 21:48:29,168 - Mesh - INFO - Getting vertices\n",
      "2024-09-11 21:48:29,177 - Mesh - INFO - Getting facet centers\n",
      "2024-09-11 21:48:29,187 - Mesh - INFO - Creating connectivity\n",
      "2024-09-11 21:48:29,575 - Mesh - INFO - Mesh object initialized.\n",
      "2024-09-11 21:48:29,576 - Mesh - INFO - Mesh data is of type: float64\n",
      "2024-09-11 21:48:29,577 - Mesh - INFO - Elapsed time: 0.41460768800000003s\n",
      "2024-09-11 21:48:29,578 - pynekread - INFO - Reading field data\n",
      "2024-09-11 21:48:29,586 - pynekread - INFO - File read\n",
      "2024-09-11 21:48:29,587 - pynekread - INFO - Elapsed time: 0.439429355s\n"
     ]
    }
   ],
   "source": [
    "msh = Mesh(comm, create_connectivity=True)\n",
    "fld = FieldRegistry(comm)\n",
    "pynekread('../data/tc_channel0.f00001', comm, data_dtype=np.double, msh=msh, fld=fld)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45d0bf3a",
   "metadata": {},
   "source": [
    "## Mapping to a new distribution\n",
    "\n",
    "We can map to a new distribution of points in any of the dimensions of the element at the same time.\n",
    "\n",
    "Here we show an example where this is done to an equidistant mesh in the z direction.\n",
    "\n",
    "For this we use the PMapper class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cea2e84e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-09-11 21:48:30,619 - Mesh - INFO - Initializing Mesh object from x,y,z ndarrays.\n",
      "2024-09-11 21:48:30,619 - Mesh - INFO - Initializing common attributes.\n",
      "2024-09-11 21:48:30,620 - Mesh - INFO - Getting vertices\n",
      "2024-09-11 21:48:30,621 - Mesh - INFO - Getting vertices\n",
      "2024-09-11 21:48:30,630 - Mesh - INFO - Getting facet centers\n",
      "2024-09-11 21:48:30,637 - Mesh - INFO - Creating connectivity\n",
      "2024-09-11 21:48:30,986 - Mesh - INFO - Mesh object initialized.\n",
      "2024-09-11 21:48:30,987 - Mesh - INFO - Mesh data is of type: float64\n",
      "2024-09-11 21:48:30,987 - Mesh - INFO - Elapsed time: 0.36897678200000006s\n",
      "2024-09-11 21:48:31,746 - Field - INFO - Initializing empty Field object\n",
      "2024-09-11 21:48:31,747 - pynekwrite - INFO - Writing file: mappedfield0.f00001\n",
      "2024-09-11 21:48:31,756 - pynekwrite - INFO - Elapsed time: 0.008410899999999888s\n"
     ]
    }
   ],
   "source": [
    "# Import the module\n",
    "from pysemtools.interpolation.mesh_to_mesh import PMapper\n",
    "\n",
    "# Here specify that it should be equal in the 3 direction\n",
    "mapper = PMapper(n=msh.lx, distribution=['GLL', 'GLL', 'EQ'])\n",
    "\n",
    "# Create the mesh with the new distribution\n",
    "eq_msh = mapper.create_mapped_mesh(comm, msh=msh)\n",
    "\n",
    "# Interpolate the fields. They are passed as a list and returned as a list\n",
    "mapped_fields = mapper.interpolate_from_field_list(comm, field_list=[fld.registry['u'],fld.registry['v'],fld.registry['w']])\n",
    "\n",
    "# Create an empty field object for the equidistant fields\n",
    "eq_fld = FieldRegistry(comm)\n",
    "eq_fld.add_field(comm, field_name='u', field=mapped_fields[0], dtype = np.double)\n",
    "eq_fld.add_field(comm, field_name='v', field=mapped_fields[1], dtype = np.double)\n",
    "eq_fld.add_field(comm, field_name='w', field=mapped_fields[2], dtype = np.double)\n",
    "\n",
    "# Write the new mesh and fields\n",
    "fname = \"mappedfield0.f00001\"\n",
    "pynekwrite(fname, comm, msh=eq_msh, fld=eq_fld, write_mesh=True, wdsz=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4741f38f",
   "metadata": {},
   "source": [
    "### Mapping all fields\n",
    "\n",
    "In the previous block, we showed how to map specific fields. However if you wish to map all the contents of the field object, you can use the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "50dda4df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-09-11 21:48:31,766 - Field - INFO - Initializing empty Field object\n",
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "eq_fld2 = mapper.create_mapped_field(comm, fld=fld)\n",
    "\n",
    "for key in eq_fld.registry.keys():\n",
    "    print(np.allclose(eq_fld.registry[key].data, eq_fld2.registry[key].data))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c45a389c",
   "metadata": {},
   "source": [
    "## Performing P refinement\n",
    "\n",
    "In some instances one wishes to keep the GLL distribution to perform integration or derivation but doing so at higher or lower polynomiar order. For this cases one can employ p refinement/coarsening. We show now how that can be done.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5a3844e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-09-11 21:48:33,932 - Mesh - INFO - Initializing Mesh object from x,y,z ndarrays.\n",
      "2024-09-11 21:48:33,933 - Mesh - INFO - Initializing common attributes.\n",
      "2024-09-11 21:48:33,933 - Mesh - INFO - Getting vertices\n",
      "2024-09-11 21:48:33,935 - Mesh - INFO - Getting vertices\n",
      "2024-09-11 21:48:33,940 - Mesh - INFO - Getting facet centers\n",
      "2024-09-11 21:48:33,943 - Mesh - INFO - Creating connectivity\n",
      "2024-09-11 21:48:34,030 - Mesh - INFO - Mesh object initialized.\n",
      "2024-09-11 21:48:34,031 - Mesh - INFO - Mesh data is of type: float64\n",
      "2024-09-11 21:48:34,033 - Mesh - INFO - Elapsed time: 0.10139660900000003s\n",
      "2024-09-11 21:48:34,922 - Field - INFO - Initializing empty Field object\n",
      "2024-09-11 21:48:34,923 - pynekwrite - INFO - Writing file: refinedfield0.f00001\n",
      "2024-09-11 21:48:34,926 - pynekwrite - INFO - Elapsed time: 0.002947181000000576s\n"
     ]
    }
   ],
   "source": [
    "# Import the module\n",
    "from pysemtools.interpolation.mesh_to_mesh import PRefiner\n",
    "\n",
    "# Here specify that it should be equal in the 3 direction\n",
    "pr = PRefiner(n_old = msh.lx, n_new = 3)\n",
    "\n",
    "# Create the mesh with the new distribution\n",
    "r_msh = pr.create_refined_mesh(comm, msh=msh)\n",
    "\n",
    "# Interpolate the fields. They are passed as a list and returned as a list\n",
    "r_fields = pr.interpolate_from_field_list(comm, field_list=[fld.registry['u'],fld.registry['v'],fld.registry['w']])\n",
    "\n",
    "# Create an empty field object for the equidistant fields\n",
    "r_fld = FieldRegistry(comm)\n",
    "r_fld.add_field(comm, field_name='u', field=r_fields[0], dtype = np.double)\n",
    "r_fld.add_field(comm, field_name='v', field=r_fields[1], dtype = np.double)\n",
    "r_fld.add_field(comm, field_name='w', field=r_fields[2], dtype = np.double)\n",
    "\n",
    "# Write the new mesh and fields\n",
    "fname = \"refinedfield0.f00001\"\n",
    "pynekwrite(fname, comm, msh=r_msh, fld=r_fld, write_mesh=True, wdsz=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96f7fdc7",
   "metadata": {},
   "source": [
    "### Refining all fields\n",
    "\n",
    "In the previous block, we showed how to refine specific fields. However if you wish to refine all the contents of the field object, you can use the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c9248d90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-09-11 21:48:34,935 - Field - INFO - Initializing empty Field object\n",
      "True\n",
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "r_fld2 = pr.create_refined_field(comm, fld=fld)\n",
    "\n",
    "for key in r_fld.registry.keys():\n",
    "    print(np.allclose(r_fld.registry[key].data, r_fld2.registry[key].data))"
   ]
  }
 ],
 "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}