main-experimental.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# EXAMPLE CODES FOR THE EXPERIMENTAL RECONSTRUCTIONS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.forward_models.bornapprox import MVProd, MVProd_iterK\n",
    "import src.optimization as opt\n",
    "from src.nn_models.unet3D import denoiser\n",
    "\n",
    "from src.utils.stats import l2\n",
    "from src.utils.misc import Scene, super_clear_cuda\n",
    "from src.utils.data import load_dataset, load_json\n",
    "\n",
    "import torch \n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the experimental data (Please refer to \"format-exp-data.ipynb\" for the instructions on getting the experimental data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = load_dataset('Yarovoy','measurements')\n",
    "tx_positions = load_dataset('Yarovoy','tx_positions')\n",
    "rx_positions = load_dataset('Yarovoy','rx_positions')\n",
    "freqs = load_dataset('Yarovoy','freqs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the radar system's parameters as well as the number of voxels and the volume boundaries for each dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "RADAR_CONFIG={\n",
    "    'F_MIN':4e9, # Minimum frequency to use in the frequency band (Hz)\n",
    "    'F_MAX':16e9, # Maximum frequencyto use in the frequency band (Hz)\n",
    "    'N_F':11, # Total number of frequency samples  \n",
    "\n",
    "    'X_MIN':-0.15, # image cube's minimum x coordinate \n",
    "    'X_MAX':0.15, # image cube's maximum x coordinate \n",
    "    'N_X':61, # Number of voxels on the x direction\n",
    "\n",
    "    'Y_MIN':-0.15, # image cube's minimum y coordinate\n",
    "    'Y_MAX':0.15, # image cube's maximum y coordinate \n",
    "    'N_Y':61, # Number of voxels on the y direction\n",
    "\n",
    "    'Z_MIN':0.35, # image cube's minimum z coordinate\n",
    "    'Z_MAX':0.65, # image cube's maximum z coordinate \n",
    "    'N_Z':61 # Number of voxels on the z direction\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Utility codes for experimental measurements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def slice_freqs(y,freqs,f_min,f_max,N_f,partition='min-max-idx'):\n",
    "    assert torch.any(freqs == f_min)\n",
    "    assert torch.any(freqs == f_max)\n",
    "\n",
    "    if partition=='min-max-idx':\n",
    "        max_idx = torch.argwhere(freqs==f_max).item()\n",
    "        min_idx = torch.argwhere(freqs==f_min).item()\n",
    "\n",
    "        slice_idxs=torch.linspace(min_idx,max_idx,N_f).to(dtype=int)\n",
    "        \n",
    "        return freqs[slice_idxs], y[:,:,slice_idxs,...]\n",
    "\n",
    "    elif partition=='min-maxBound':\n",
    "\n",
    "        bw=f_max-f_min\n",
    "        df=freqs[1]-freqs[0]\n",
    "        bw_floor=  df*np.floor((bw/df)/(N_f-1))*(N_f-1)\n",
    "        f_max_used=f_min+bw_floor\n",
    "        \n",
    "        max_idx = torch.argwhere(freqs==f_max_used).item()\n",
    "        min_idx = torch.argwhere(freqs==f_min).item()\n",
    "\n",
    "        slice_idxs=torch.linspace(min_idx,max_idx,N_f).to(dtype=int)\n",
    "\n",
    "        assert slice_idxs.size()[0]==N_f\n",
    "        \n",
    "        print(\"fmin: {:.2e}, fmax: {:.2e}, fmax-error: {:.2e}, Nf: {}\".format(freqs[slice_idxs][0],freqs[slice_idxs][-1],f_max-freqs[slice_idxs][-1],freqs[slice_idxs].shape[0]))\n",
    "        \n",
    "        return freqs[slice_idxs], y[:,:,slice_idxs,...]\n",
    "\n",
    "def get_scene_coords(x_min,x_max,nx,y_min,y_max,ny,z_min,z_max,nz):\n",
    "    x_coords=torch.linspace(x_min,x_max,nx)\n",
    "    y_coords=torch.linspace(y_min,y_max,ny)\n",
    "    z_coords=torch.linspace(z_min,z_max,nz)\n",
    "    return x_coords,y_coords,z_coords\n",
    "\n",
    "class RadarSystemParams():\n",
    "    def __init__(self,freqs,tx_positions,rx_positions,x_coords,y_coords,z_coords) -> None:\n",
    "        self.freqs =freqs\n",
    "        self.tx_positions = tx_positions\n",
    "        self.rx_positions = rx_positions\n",
    "        self.x_coords = x_coords\n",
    "        self.y_coords = y_coords\n",
    "        self.z_coords = z_coords\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Slice the measurements to be used during the reconstruction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_used,y_sampled=slice_freqs(y,freqs,f_min=RADAR_CONFIG['F_MIN'],f_max=RADAR_CONFIG['F_MAX'],N_f=RADAR_CONFIG['N_F'],partition='min-maxBound')\n",
    "assert torch.all(torch.linspace(f_used.min(),f_used.max(),f_used.size()[0],dtype=torch.float64) == f_used)\n",
    "print(\"Utilized Frequnecies:\",f_used)\n",
    "\n",
    "x_coords,y_coords,z_coords = get_scene_coords(  x_min=RADAR_CONFIG['X_MIN'],    x_max=RADAR_CONFIG['X_MAX'], nx=RADAR_CONFIG['N_X'],\n",
    "                                                y_min=RADAR_CONFIG['Y_MIN'],    y_max=RADAR_CONFIG['Y_MAX'], ny=RADAR_CONFIG['N_Y'],\n",
    "                                                z_min=RADAR_CONFIG['Z_MIN'],    z_max=RADAR_CONFIG['Z_MAX'], nz=RADAR_CONFIG['N_Z'])\n",
    "\n",
    "radar_system_params = RadarSystemParams(freqs=f_used,tx_positions=tx_positions,rx_positions=rx_positions,x_coords=x_coords,y_coords=y_coords,z_coords=z_coords)\n",
    "\n",
    "MoverN=(tx_positions.shape[0]*rx_positions.shape[0]*RADAR_CONFIG['N_F'])/(RADAR_CONFIG['N_X']*RADAR_CONFIG['N_Y']*RADAR_CONFIG['N_Z'])\n",
    "\n",
    "del y\n",
    "del freqs\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate the forward model matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = MVProd( freqs = radar_system_params.freqs, # frequency list\n",
    "            tx_positions = radar_system_params.tx_positions, # transmitter positions\n",
    "            rx_positions = radar_system_params.rx_positions, # receiver positions\n",
    "            x_coords = radar_system_params.x_coords, # discretized x-coordinates\n",
    "            y_coords = radar_system_params.y_coords, # discretized y-coordinates\n",
    "            z_coords = radar_system_params.z_coords, # discretized z-coordinates\n",
    "            path = '', save = False # you can set a path to save the forward model matrix \n",
    "            )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For higher resolution images you may prefer to use the following module as it allows batchwise computation of measurements by iteratively computing the response for different frequencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A = MVProd_iterK( freqs = radar_system_params.freqs,\n",
    "#             tx_positions = radar_system_params.tx_positions,\n",
    "#             rx_positions = radar_system_params.rx_positions,\n",
    "#             x_coords = radar_system_params.x_coords,\n",
    "#             y_coords = radar_system_params.y_coords,\n",
    "#             z_coords = radar_system_params.z_coords,\n",
    "#             path = 'None', save = False,\n",
    "#             f_batch_size=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Forward model matrix shape: ',A.A.shape)\n",
    "print('100*(M/N) = ~{:.2f}'.format(100*A.A.shape[0]*A.A.shape[1]*A.A.shape[2]/(A.A.shape[3]*A.A.shape[4]*A.A.shape[5])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As explained in the paper, $\\epsilon$ is emprically set to match ~10dB measurement SNR:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SNR=10 # dB\n",
    "y_sampled=y_sampled\n",
    "eps = (l2(y_sampled)/(10**(SNR/20))).item()\n",
    "print('epsillon',eps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cpu\")\n",
    "torch.backends.cudnn.benchmark = True\n",
    "if device != \"cpu\":\n",
    "     super_clear_cuda(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Experimental Measurements' scale may be different than the simulated data. We therefore need to scale the input images before and after entering the DNN:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class scale_invariance_wrapper(torch.nn.Module):\n",
    "    def __init__(self,net) -> None:\n",
    "        super().__init__()\n",
    "        self.net=net\n",
    "      \n",
    "    def forward(self,xn):\n",
    "        x=xn[:,0:1,:,:,:]\n",
    "        n=xn[:,1:2,:,:,:]\n",
    "\n",
    "        sc = opt.maxmag3D(x)\n",
    "        x = x/sc\n",
    "        z  = self.net(torch.concat([x,n],dim=1))\n",
    "        return z*sc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reconstruct the scene:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "\n",
    "    model_state_dict= torch.load('trained_model/base_model.pt')\n",
    "\n",
    "    denoiser_config=load_json('trained_model/info.json')['arch']\n",
    "    denoiser_net = denoiser(**denoiser_config)\n",
    "    denoiser_net.load_state_dict(model_state_dict)\n",
    "\n",
    "    denoiser_net = scale_invariance_wrapper(denoiser_net) # make the dnn scale invariant\n",
    "\n",
    "    solver = opt.NNCG_CSALSA(forward_model=A, denoiser_net=denoiser_net, blind=False, noise_level_map='std', cg_max_iter=5, cg_th=1e-5, zero_phase=False, record=False)\n",
    "    # you can use other denoising networks as well,\n",
    "    # if you plan to use a blind denoiser set \"blind\" to True, \n",
    "    # if your noise level map is the variance of the noise, set \"noise_level_map\" to 'var'\n",
    "    # cg_max_iter is the maximum number of conjugate gradient iterations\n",
    "    # cg_th is the minimum relative improvement on the l2 norm before terminating the conjugate gradient iterations (cg_max_iter has priority over cg_th)\n",
    "    # you should always use zero_phase=False. If true, it enforces real-positive values on reflectivity distribution, which is an incorrect assumption\n",
    "    # by setting record to True, you can store the values of the data fidelity and regularization terms throughout the iterations (may be useful for debugging)\n",
    "\n",
    "\n",
    "    # Other solvers:\n",
    "    # solver = opt.TVCG_CSALSA(forward_model=A, chambolle_iter=5, cg_max_iter=5, cg_th=1e-5, zero_phase=False, record=False)\n",
    "    # solver = opt.l1CG_CSALSA(forward_model=A, cg_max_iter=5, cg_th=1e-5, zero_phase=False, record=False)\n",
    "    # solver = opt.BackProjection(radar_system_params=radar_system_params)\n",
    "    # solver = opt.KirchhoffMigration(radar_system_params=radar_system_params,compansate_limited_aparture=False)\n",
    "    \n",
    "    solver.to(device)\n",
    "    solver.clear_history()\n",
    "\n",
    "    kwargs = {}\n",
    "\n",
    "    # set the algorithm parameters\n",
    "    kwargs['y'] = y_sampled.to(device)\n",
    "    kwargs['eps'] = eps # epsillon parameter in the paper\n",
    "    kwargs['tqdm_active'] = True # Set to True if you want to see an iteration counter\n",
    "    kwargs['K'] = 500 # Kappa parameter in the paper\n",
    "    kwargs['noise_var'] = 7e-2 # alpha parameter in the paper\n",
    "    kwargs['max_iter']=30\n",
    "\n",
    "    # run the solver\n",
    "    solver.forward(**kwargs)\n",
    "    \n",
    "    # get estimates (estimates are returned as normalized magnitues, you can access complex valued reconstructions by \"solver.x\")\n",
    "    est = solver.get_estimate().cpu()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Show the reconstructed reflectivity magnitudes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scene=Scene()\n",
    "scene.show_experimental(x=est,radar_system_params=radar_system_params)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "TCI",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}