Add PINQI and MoDL examples

examples/notebooks/modl.ipynb

@@ -0,0 +1,266 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {
+    "lines_to_next_cell": 0
+   },
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/PTB-MR/mrpro/blob/main/examples/notebooks/modl.ipynb)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {
+    "tags": [
+     "remove-cell"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "import importlib\n",
+    "\n",
+    "if not importlib.util.find_spec('mrpro'):\n",
+    "    %pip install mrpro[notebooks]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline\n",
+    "from collections.abc import Sequence\n",
+    "from pathlib import Path\n",
+    "from typing import TypedDict\n",
+    "\n",
+    "import matplotlib.axes\n",
+    "import matplotlib.pyplot as plt\n",
+    "import mrpro\n",
+    "import torch\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "\n",
+    "class BatchType(TypedDict):\n",
+    "    \"\"\"A single Batch.\"\"\"\n",
+    "\n",
+    "    data: mrpro.data.KData\n",
+    "    target: mrpro.data.IData\n",
+    "    csm: mrpro.data.CsmData\n",
+    "\n",
+    "\n",
+    "class AcceleratedFastMRI(torch.utils.data.Dataset):\n",
+    "    \"\"\"An undersampled FastMRI Dataset.\"\"\"\n",
+    "\n",
+    "    def __init__(self, path: Path, acceleration: float = 12, noise_level: float = 0.1):\n",
+    "        \"\"\"Create an undersampled FastMRI Dataset.\n",
+    "\n",
+    "        Parameters\n",
+    "        ----------\n",
+    "        path\n",
+    "            Path to the FastMRI dataset.\n",
+    "        acceleration\n",
+    "            Undersampling factor; higher values mean more acceleration. Default is 12.\n",
+    "        noise_level\n",
+    "            Level of additive Gaussian noise applied to the FastMRI dataset. Default is 0.1.\n",
+    "        \"\"\"\n",
+    "        self.acceleration = acceleration\n",
+    "        files = list(path.glob('*AXT1*'))\n",
+    "        self.dataset = mrpro.phantoms.FastMRIKDataDataset(files)\n",
+    "        self.noise_level = noise_level\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        \"\"\"Get length of the dataset.\"\"\"\n",
+    "        return len(self.dataset)\n",
+    "\n",
+    "    def __getitem__(self, index: int) -> BatchType:\n",
+    "        \"\"\"Get a single batch of data.\n",
+    "\n",
+    "        Parameters\n",
+    "        ----------\n",
+    "        index\n",
+    "            Index of the batch.\n",
+    "\n",
+    "        Returns\n",
+    "        -------\n",
+    "        A single batch of data with keys 'data', 'target', and 'csm'.and\n",
+    "        \"\"\"\n",
+    "        data = self.dataset[index]\n",
+    "        data = data.remove_readout_os()\n",
+    "        data.data /= data.data.std()\n",
+    "        reconstruction = mrpro.algorithms.reconstruction.DirectReconstruction(\n",
+    "            data, csm=lambda data: mrpro.data.CsmData.from_idata_inati(data, downsampled_size=64)\n",
+    "        )\n",
+    "        csm = reconstruction.csm\n",
+    "        target = reconstruction(data)\n",
+    "\n",
+    "        n = max(data.data.shape[-2:])\n",
+    "        distance = (torch.linspace(-1, 1, n)[:, None] ** 2 + torch.linspace(-1, 1, n) ** 2).sqrt()\n",
+    "        random = 0.1 / (distance + 0.1) + torch.rand_like(distance)\n",
+    "        threshold = torch.kthvalue(random.ravel(), int(n**2 * (1 - 1 / self.acceleration))).values\n",
+    "        undersampling_mask = mrpro.utils.pad_or_crop(random > threshold, data.data.shape[-2:])\n",
+    "        data_undersampled = data[..., undersampling_mask].rearrange('k ... 1 -> ... k')\n",
+    "\n",
+    "        noise = mrpro.utils.RandomGenerator(seed=index).randn_like(data_undersampled.data)\n",
+    "        data_undersampled.data += self.noise_level * noise\n",
+    "\n",
+    "        assert csm is not None  # for mypy\n",
+    "        return {'data': data_undersampled, 'target': target, 'csm': csm}\n",
+    "\n",
+    "\n",
+    "class MODL(torch.nn.Module):\n",
+    "    \"\"\"MODL network.\"\"\"\n",
+    "\n",
+    "    def __init__(self, iterations: int = 8, n_features: Sequence[int] = (64, 64, 64, 64)):\n",
+    "        \"\"\"Initialize MODL network.\n",
+    "\n",
+    "        Parameters\n",
+    "        ----------\n",
+    "        iterations\n",
+    "            Number of iterations.\n",
+    "        n_features\n",
+    "            Number of features in the network.\n",
+    "        \"\"\"\n",
+    "        super().__init__()\n",
+    "        cnn = mrpro.nn.nets.BasicCNN(\n",
+    "            dim=2,\n",
+    "            channels_in=2,\n",
+    "            channels_out=2,\n",
+    "            n_features=n_features,\n",
+    "            batch_norm=True,\n",
+    "        )\n",
+    "        self.network = mrpro.nn.Residual(mrpro.nn.ComplexAsChannel(mrpro.nn.PermutedBlock((-1, -2), cnn)))\n",
+    "        self.network = torch.compile(self.network, dynamic=True, fullgraph=True)\n",
+    "        self.iterations = iterations\n",
+    "        self.regularization_weights = torch.nn.Parameter(0.2 * torch.ones(iterations))\n",
+    "\n",
+    "    def __call__(self, kdata: mrpro.data.KData, csm: mrpro.data.CsmData) -> mrpro.data.IData:\n",
+    "        \"\"\"Apply MODL network.\n",
+    "\n",
+    "        Parameters\n",
+    "        ----------\n",
+    "        kdata\n",
+    "            The k-space data.\n",
+    "        csm\n",
+    "            The coil sensitivity maps.\n",
+    "\n",
+    "        Returns\n",
+    "        -------\n",
+    "        The reconstructed image.\n",
+    "        \"\"\"\n",
+    "        return super().__call__(kdata, csm)\n",
+    "\n",
+    "    def forward(self, kdata: mrpro.data.KData, csm: mrpro.data.CsmData) -> mrpro.data.IData:\n",
+    "        \"\"\"Apply the MODL network.\"\"\"\n",
+    "        fourier_op = mrpro.operators.FourierOp.from_kdata(kdata)\n",
+    "        acquisition_op = fourier_op @ csm.as_operator()\n",
+    "        (zero_filled_image,) = acquisition_op.H(kdata.data)\n",
+    "        gram = acquisition_op.gram\n",
+    "        data_consistency_op = mrpro.operators.ConjugateGradientOp(\n",
+    "            operator_factory=lambda _image, weight: gram + weight,\n",
+    "            rhs_factory=lambda image, weight: zero_filled_image + weight * image,\n",
+    "        )\n",
+    "\n",
+    "        (image,) = mrpro.algorithms.optimizers.cg(gram, zero_filled_image, max_iterations=5)\n",
+    "        for iteration in range(self.iterations):\n",
+    "            regularization = self.network(image)\n",
+    "            (image,) = data_consistency_op(regularization, self.regularization_weights[iteration])\n",
+    "\n",
+    "        return mrpro.data.IData(image, header=mrpro.data.IHeader.from_kheader(kdata.header))\n",
+    "\n",
+    "\n",
+    "def plot(batch: BatchType, prediction: mrpro.data.IData, step: int) -> None:\n",
+    "    \"\"\"Plot the direct, sense, and modl reconstructions.\"\"\"\n",
+    "    target = batch['target'].rss().cpu().squeeze()\n",
+    "    direct = mrpro.algorithms.reconstruction.DirectReconstruction(batch['data'], csm=batch['csm'])(batch['data'])\n",
+    "    direct = direct.rss().cpu().squeeze()\n",
+    "    direct *= target.std() / direct.std()\n",
+    "    sense = mrpro.algorithms.reconstruction.IterativeSENSEReconstruction(batch['data'], csm=batch['csm'])(batch['data'])\n",
+    "    sense = sense.rss().cpu().squeeze()\n",
+    "    prediction_ = prediction.rss().cpu().squeeze().detach()\n",
+    "\n",
+    "    ssim = mrpro.operators.functionals.SSIM(mrpro.utils.pad_or_crop(target[None], (320, 320)))\n",
+    "\n",
+    "    def show(ax: matplotlib.axes.Axes, data: torch.Tensor, label: str):\n",
+    "        data = mrpro.utils.pad_or_crop(data, (320, 320))\n",
+    "        ax.imshow(data, vmin=0, vmax=target.max().item(), cmap='gray')\n",
+    "        if label != 'Ground Truth':\n",
+    "            (ssim_value,) = ssim(data[None])\n",
+    "            ax.text(\n",
+    "                0.98,\n",
+    "                0.1,\n",
+    "                f'SSIM: {ssim_value.item():.2f}',\n",
+    "                color='white',\n",
+    "                horizontalalignment='right',\n",
+    "                verticalalignment='top',\n",
+    "                transform=ax.transAxes,\n",
+    "            )\n",
+    "        ax.set_title(label)\n",
+    "        ax.set_axis_off()\n",
+    "\n",
+    "    fig, ax = plt.subplots(1, 4)\n",
+    "    show(ax[0], direct, 'Direct')\n",
+    "    show(ax[1], sense, 'CG-SENSE')\n",
+    "    show(ax[2], prediction_, 'MODL')\n",
+    "    show(ax[3], target, 'Ground Truth')\n",
+    "    fig.tight_layout()\n",
+    "    fig.savefig(f'modl_{step}.pdf', bbox_inches='tight', pad_inches=0)\n",
+    "\n",
+    "\n",
+    "# %%.\n",
+    "path = Path('/echo/allgemein/resources/publicTrainingData/fastmri/brain_multicoil_train/')\n",
+    "dataset = AcceleratedFastMRI(path)\n",
+    "dataloader = torch.utils.data.DataLoader(dataset, num_workers=16, shuffle=True, collate_fn=lambda batch: batch[0])\n",
+    "modl = MODL().cuda()\n",
+    "optimizer = torch.optim.Adam(modl.parameters(), lr=1e-3)\n",
+    "pbar = tqdm(dataloader)\n",
+    "for i, batch in enumerate(pbar):\n",
+    "    optimizer.zero_grad()\n",
+    "    kdata, csm, target = (batch['data'].cuda(), batch['csm'].cuda(), batch['target'].cuda())\n",
+    "    prediction = modl(kdata, csm)\n",
+    "    objective = 0.5 * mrpro.operators.functionals.MSE(target.data) - mrpro.operators.functionals.SSIM(target.data)\n",
+    "    (loss,) = objective(prediction.data)\n",
+    "    loss.backward()\n",
+    "    torch.nn.utils.clip_grad_norm_(modl.parameters(), 5.0)\n",
+    "    optimizer.step()\n",
+    "\n",
+    "    pbar.set_postfix(loss=loss.item())\n",
+    "    if i % 200 == 0:\n",
+    "        plot(batch, prediction, i)\n",
+    "        print(modl.regularization_weights)\n",
+    "        state = {'modl': modl.state_dict(), 'optimizer': optimizer.state_dict()}\n",
+    "        torch.save(state, f'modl_{i}.pt')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": []
+  },
+  "jupytext": {
+   "cell_metadata_filter": "mystnb,tags,-all"
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}