examples/at_circular_piston_AS/at_circular_piston_AS.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install git+https://github.com/waltsims/k-wave-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from copy import deepcopy\n",
    "\n",
    "from kwave.data import Vector\n",
    "\n",
    "from kwave.kgrid import kWaveGrid\n",
    "from kwave.kmedium import kWaveMedium\n",
    "from kwave.ksensor import kSensor\n",
    "from kwave.ksource import kSource\n",
    "from kwave.utils.filters import extract_amp_phase\n",
    "from kwave.utils.mapgen import focused_bowl_oneil\n",
    "from kwave.utils.math import round_even\n",
    "from kwave.utils.kwave_array import kWaveArray\n",
    "from kwave.utils.signals import create_cw_signals\n",
    "\n",
    "from kwave.kspaceFirstOrderAS import kspaceFirstOrderASC\n",
    "\n",
    "from kwave.options.simulation_options import SimulationOptions, SimulationType\n",
    "from kwave.options.simulation_execution_options import SimulationExecutionOptions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# medium parameters\n",
    "c0: float            = 1500.0  # sound speed [m/s]\n",
    "rho0: float          = 1000.0  # density [kg/m^3]\n",
    "\n",
    "# source parameters\n",
    "source_f0            = 1.0e6              # source frequency [Hz]\n",
    "source_roc           = 30e-3              # bowl radius of curvature [m]\n",
    "source_diameter      = 30e-3              # bowl aperture diameter [m]\n",
    "source_amp           = np.array([1.0e6])  # source pressure [Pa]\n",
    "source_phase         = np.array([0.0])    # source phase [radians]\n",
    "\n",
    "# grid parameters\n",
    "axial_size: float    = 50.0e-3  # total grid size in the axial dimension [m]\n",
    "lateral_size: float  = 45.0e-3  # total grid size in the lateral dimension [m]\n",
    "\n",
    "# computational parameters\n",
    "ppw: int             = 3      # number of points per wavelength\n",
    "t_end: float         = 40e-6  # total compute time [s] (this must be long enough to reach steady state)\n",
    "record_periods: int  = 1      # number of periods to record\n",
    "cfl: float           = 0.05   # CFL number\n",
    "source_x_offset: int = 20     # grid points to offset the source\n",
    "bli_tolerance: float = 0.01   # tolerance for truncation of the off-grid source points\n",
    "upsampling_rate: int = 10     # density of integration points relative to grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the grid spacing based on the PPW and F0\n",
    "dx: float = c0 / (ppw * source_f0)   # [m]\n",
    "\n",
    "# compute the size of the grid\n",
    "Nx: int= round_even(axial_size / dx) + source_x_offset\n",
    "Ny: int = round_even(lateral_size / dx)\n",
    "\n",
    "grid_size_points = Vector([Nx, Ny])\n",
    "grid_spacing_meters = Vector([dx, dx])\n",
    "\n",
    "# create the k-space grid\n",
    "kgrid = kWaveGrid(grid_size_points, grid_spacing_meters)\n",
    "\n",
    "# compute points per temporal period\n",
    "ppp: int = round(ppw / cfl)\n",
    "\n",
    "# compute corresponding time spacing\n",
    "dt: float = 1.0 / (ppp * source_f0)\n",
    "\n",
    "# create the time array using an integer number of points per period\n",
    "Nt: int = round(t_end / dt)\n",
    "kgrid.setTime(Nt, dt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define source"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source = kSource()\n",
    "\n",
    "# create time varying source\n",
    "source_sig = create_cw_signals(np.squeeze(kgrid.t_array), source_f0, source_amp, source_phase)\n",
    "\n",
    "# set arc position and orientation\n",
    "arc_pos = [kgrid.x_vec[0].item() + source_x_offset * kgrid.dx, 0]\n",
    "focus_pos = [kgrid.x_vec[-1].item(), 0]\n",
    "\n",
    "# create empty kWaveArray\n",
    "karray = kWaveArray(axisymmetric=True,\n",
    "                    bli_tolerance=bli_tolerance,\n",
    "                    upsampling_rate=upsampling_rate,\n",
    "                    single_precision=True)\n",
    "\n",
    "# add bowl shaped element\n",
    "karray.add_arc_element(arc_pos, source_roc, source_diameter, focus_pos)\n",
    "\n",
    "# assign binary mask\n",
    "source.p_mask = karray.get_array_binary_mask(kgrid)\n",
    "\n",
    "# assign source signals\n",
    "source.p = karray.get_distributed_source_signal(kgrid, source_sig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up medium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "medium = kWaveMedium(sound_speed=c0, density=rho0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up sensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sensor = kSensor()\n",
    "\n",
    "# set sensor mask to record central plane, not including the source point\n",
    "sensor.mask = np.zeros((Nx, Ny), dtype=bool)\n",
    "sensor.mask[(source_x_offset + 1):, :] = True\n",
    "\n",
    "# record the pressure\n",
    "sensor.record = ['p']\n",
    "\n",
    "# record only the final few periods when the field is in steady state\n",
    "sensor.record_start_index = kgrid.Nt - (record_periods * ppp) + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define simulation options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulation_options = SimulationOptions(\n",
    "    simulation_type=SimulationType.AXISYMMETRIC,\n",
    "    data_cast='single',\n",
    "    data_recast=False,\n",
    "    save_to_disk=True,\n",
    "    save_to_disk_exit=False,\n",
    "    pml_inside=False)\n",
    "\n",
    "execution_options = SimulationExecutionOptions(\n",
    "    is_gpu_simulation=False,\n",
    "    delete_data=False,\n",
    "    verbose_level=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sensor_data = kspaceFirstOrderASC(medium=deepcopy(medium),\n",
    "                                  kgrid=deepcopy(kgrid),\n",
    "                                  source=deepcopy(source),\n",
    "                                  sensor=deepcopy(sensor),\n",
    "                                  simulation_options=simulation_options,\n",
    "                                  execution_options=execution_options)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Post-processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract amplitude from the sensor data\n",
    "amp, _, _  = extract_amp_phase(sensor_data['p'].T, 1.0 / kgrid.dt, source_f0,\n",
    "                               dim=1, fft_padding=1, window='Rectangular')\n",
    "\n",
    "# reshape data\n",
    "amp = np.reshape(amp, (Nx - (source_x_offset+1), Ny), order='F')\n",
    "\n",
    "# extract pressure on axis\n",
    "amp_on_axis = amp[:, 0]\n",
    "\n",
    "# define axis vectors for plotting\n",
    "x_vec = np.squeeze(kgrid.x_vec[(source_x_offset + 1):, :] - kgrid.x_vec[source_x_offset])\n",
    "y_vec = kgrid.y_vec\n",
    "\n",
    "# =========================================================================\n",
    "# ANALYTICAL SOLUTION\n",
    "# =========================================================================\n",
    "\n",
    "# calculate the wavenumber\n",
    "knumber = 2.0 * np.pi * source_f0 / c0\n",
    "\n",
    "# define radius and axis\n",
    "x_max = (Nx-1) * dx\n",
    "delta_x = x_max / 10000.0\n",
    "x_ref = np.arange(0.0, x_max + delta_x, delta_x)\n",
    "\n",
    "# calculate analytical solution\n",
    "p_ref_axial, _, _ = focused_bowl_oneil(source_roc,\n",
    "                                 source_diameter,\n",
    "                                 source_amp[0] / (c0 * rho0),\n",
    "                                 source_f0,\n",
    "                                 c0,\n",
    "                                 rho0,\n",
    "                                 axial_positions=x_ref)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the pressure along the focal axis of the piston\n",
    "fig1, ax1 = plt.subplots(1, 1)\n",
    "ax1.plot(1e3 * x_ref, 1e-6 * p_ref_axial, 'k-', label='Exact')\n",
    "ax1.plot(1e3 * x_vec, 1e-6 * amp_on_axis, 'b.', label='k-Wave')\n",
    "ax1.legend()\n",
    "ax1.set(xlabel='Axial Position [mm]',\n",
    "        ylabel='Pressure [MPa]',\n",
    "        title='Axial Pressure')\n",
    "ax1.set_xlim(0.0, 1e3 * axial_size)\n",
    "ax1.set_ylim(0.0, 20)\n",
    "ax1.grid()\n",
    "\n",
    "# get grid weights\n",
    "grid_weights = karray.get_array_grid_weights(kgrid)\n",
    "\n",
    "# define axis vectors for plotting\n",
    "yvec = np.squeeze(kgrid.y_vec) - kgrid.y_vec[0].item()\n",
    "y_vec = np.hstack((-np.flip(yvec)[:-1], yvec))\n",
    "x_vec = np.squeeze(kgrid.x_vec[(source_x_offset + 1):, :] - kgrid.x_vec[source_x_offset])\n",
    "\n",
    "data = np.hstack((np.fliplr(amp[:, :-1]), amp)) / 1e6\n",
    "sp =   np.hstack((np.fliplr(source.p_mask[:, :-1]), source.p_mask))\n",
    "gw =   np.hstack((np.fliplr(grid_weights[:, :-1]), grid_weights))\n",
    "\n",
    "fig3, (ax3a, ax3b) = plt.subplots(1, 2)\n",
    "ax3a.pcolormesh(1e3 * np.squeeze(y_vec),\n",
    "                1e3 * np.squeeze(kgrid.x_vec[:, :] - kgrid.x_vec[source_x_offset]),\n",
    "                sp,\n",
    "                shading='gouraud')\n",
    "ax3a.set(xlabel='y [mm]',\n",
    "         ylabel='x [mm]',\n",
    "         title='Source Mask')\n",
    "ax3b.pcolormesh(1e3 * np.squeeze(y_vec),\n",
    "                1e3 * np.squeeze(kgrid.x_vec[:, :] - kgrid.x_vec[source_x_offset]),\n",
    "                gw,\n",
    "                shading='gouraud')\n",
    "ax3b.set(xlabel='y [mm]',\n",
    "         ylabel='x [mm]',\n",
    "         title='Off-Grid Source Weights')\n",
    "fig3.tight_layout(pad=1.2)\n",
    "\n",
    "# plot the pressure field\n",
    "fig4, ax4 = plt.subplots(1, 1)\n",
    "ax4.pcolormesh(1e3 * np.squeeze(y_vec),\n",
    "               1e3 * np.squeeze(x_vec),\n",
    "               data,\n",
    "               shading='gouraud')\n",
    "ax4.set(xlabel='Lateral Position [mm]',\n",
    "        ylabel='Axial Position [mm]',\n",
    "        title='Pressure Field')\n",
    "ax4.invert_yaxis()\n",
    "\n",
    "# show figures\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyOTl1BLKuzg6aBmvj8vKKSi",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}