experimente.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Not for the public Eye"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ants\n",
    "import numpy as np\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ANTsImage (LPI)\n",
      "\t Pixel Type : float (float32)\n",
      "\t Components : 6\n",
      "\t Dimensions : (240, 240, 155)\n",
      "\t Spacing    : (1.0, 1.0, 1.0)\n",
      "\t Origin     : (239.0, 0.0, 0.0)\n",
      "\t Direction  : [-1.  0.  0.  0. -1.  0.  0.  0.  1.]\n",
      "\n",
      "ANTsImage (LPI)\n",
      "\t Pixel Type : float (float32)\n",
      "\t Components : 1\n",
      "\t Dimensions : (240, 240, 155)\n",
      "\t Spacing    : (1.0, 1.0, 1.0)\n",
      "\t Origin     : (239.0, 0.0, 0.0)\n",
      "\t Direction  : [-1.  0.  0.  0. -1.  0.  0.  0.  1.]\n",
      "\n",
      "ANTsImage (RAI)\n",
      "\t Pixel Type : float (float32)\n",
      "\t Components : 1\n",
      "\t Dimensions : (240, 240, 155)\n",
      "\t Spacing    : (1.0, 1.0, 1.0)\n",
      "\t Origin     : (-0.0, -239.0, 0.0)\n",
      "\t Direction  : [1. 0. 0. 0. 1. 0. 0. 0. 1.]\n",
      "\n",
      "/tmp/tmp7vmg5dwb1Warp.nii.gz\n"
     ]
    }
   ],
   "source": [
    "_atlas_fiber_DTI_path = \"tissue_reconstruction/data/FSL_HCP1065_tensor_1mm_space-HPC-AntsIndexSpace_SRI.nii.gz\"\n",
    "_atlas_t1_path = \"tissue_reconstruction/data/sub-mni152_t1-inside-brain_space-sri.nii.gz\"\n",
    "_patient_path = \"tissue_reconstruction/data/sample_patient.nii.gz\"\n",
    "\n",
    "\n",
    "tensor_image = ants.image_read(_atlas_fiber_DTI_path)\n",
    "atlas_t1 = ants.image_read(_atlas_t1_path)\n",
    "patient = ants.image_read(_patient_path)\n",
    "\n",
    "reg = ants.registration(patient, atlas_t1, \"SyN\")\n",
    "\n",
    "transform = reg['fwdtransforms']\n",
    "\n",
    "print(tensor_image)\n",
    "print(atlas_t1)\n",
    "print(patient)\n",
    "print(transform[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optionally for the Fututre combine the affine and the warp ...\n",
    "\n",
    "def combine_warp(fwd_transform_paths):\n",
    "    return ants.image_read(fwd_transform_paths[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ANTsImage (RAI)\n",
      "\t Pixel Type : float (float32)\n",
      "\t Components : 6\n",
      "\t Dimensions : (240, 240, 155)\n",
      "\t Spacing    : (1.0, 1.0, 1.0)\n",
      "\t Origin     : (-0.0, -239.0, 0.0)\n",
      "\t Direction  : [1. 0. 0. 0. 1. 0. 0. 0. 1.]\n",
      "\n",
      "FIELD: /tmp/tmpxcf5772f1Warp.nii.gz\n",
      "moving_image_filename: out/transformed_dti.nii.gz\n",
      "output_image_filename: out/transformed_rotated_dti.nii.gz\n",
      "[0/1]: FIELD: /tmp/tmpxcf5772f1Warp.nii.gz\n",
      "Warp transform\n",
      "Iterating over image\n"
     ]
    }
   ],
   "source": [
    "def func1(fixed_image, dti_img, fwd_transform_paths):\n",
    "    # Using ANTs registration throught the system\n",
    "    warp_path = fwd_transform_paths[0]    \n",
    "    \n",
    "    # 1. Split into Six components\n",
    "    dti_components = ants.split_channels(dti_img)\n",
    "    #print(dti_components)\n",
    "\n",
    "    # 2. Transform every component\n",
    "    transformed_components = []\n",
    "    for component in dti_components:\n",
    "        transformed_component = ants.apply_transforms(fixed=fixed_image, \n",
    "                                                    moving=component, \n",
    "                                                    transformlist=fwd_transform_paths)\n",
    "        transformed_components.append(transformed_component)\n",
    "    transformed_dti_img = ants.merge_channels(transformed_components)\n",
    "\n",
    "    print(transformed_dti_img)\n",
    "    \n",
    "    # 3. Reorient Image ?!\n",
    "    transformed_dti_save_path = \"out/transformed_dti.nii.gz\"\n",
    "    transformed_rotated_dti_save_path = \"out/transformed_rotated_dti.nii.gz\"\n",
    "\n",
    "    ants.image_write(transformed_dti_img, transformed_dti_save_path)\n",
    "\n",
    "    reorientTensor = f\"ReorientTensorImage 3 {transformed_dti_save_path} {transformed_rotated_dti_save_path} {warp_path}\"\n",
    "    os.system(reorientTensor) \n",
    "\n",
    "func1(patient, tensor_image, transform)\n",
    "rotated_tensor1 = ants.image_read(\"out/transformed_rotated_dti.nii.gz\")\n",
    "transformed_tensor_img = ants.image_read(\"out/transformed_dti.nii.gz\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 8/240 [01:26<41:45, 10.80s/it]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[54], line 139\u001b[0m\n\u001b[1;32m    136\u001b[0m     transformed_dti_img \u001b[38;5;241m=\u001b[39m ants\u001b[38;5;241m.\u001b[39mfrom_numpy(transformed_dti, fixed_image\u001b[38;5;241m.\u001b[39morigin, fixed_image\u001b[38;5;241m.\u001b[39mspacing, fixed_image\u001b[38;5;241m.\u001b[39mdirection, has_components\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m    137\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m transformed_dti_img\n\u001b[0;32m--> 139\u001b[0m rotated_tensor2 \u001b[38;5;241m=\u001b[39m \u001b[43mfunc2\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpatient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtensor_image\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtransform\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[54], line 114\u001b[0m, in \u001b[0;36mfunc2\u001b[0;34m(fixed_image, dti_img, fwd_transform_paths)\u001b[0m\n\u001b[1;32m    111\u001b[0m point \u001b[38;5;241m=\u001b[39m (x, y, z)\n\u001b[1;32m    113\u001b[0m \u001b[38;5;66;03m# Compute local Jacobian and extract rotation\u001b[39;00m\n\u001b[0;32m--> 114\u001b[0m jacobian \u001b[38;5;241m=\u001b[39m \u001b[43mget_local_jacobian\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdef_field_np\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpoint\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfixed_image\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mspacing\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    115\u001b[0m rotation \u001b[38;5;241m=\u001b[39m get_rotation_from_jacobian(jacobian)\n\u001b[1;32m    117\u001b[0m \u001b[38;5;66;03m# Transform tensor\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[54], line 45\u001b[0m, in \u001b[0;36mget_local_jacobian\u001b[0;34m(deformation_field, point, spacing)\u001b[0m\n\u001b[1;32m     41\u001b[0m         jacobian[\u001b[38;5;241m2\u001b[39m, i] \u001b[38;5;241m=\u001b[39m (deformation_field[x, y, z\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m, i] \u001b[38;5;241m-\u001b[39m \n\u001b[1;32m     42\u001b[0m                         deformation_field[x, y, z\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, i]) \u001b[38;5;241m/\u001b[39m (\u001b[38;5;241m2\u001b[39m \u001b[38;5;241m*\u001b[39m spacing[\u001b[38;5;241m2\u001b[39m])\n\u001b[1;32m     44\u001b[0m \u001b[38;5;66;03m# Add identity matrix since the deformation field represents displacement\u001b[39;00m\n\u001b[0;32m---> 45\u001b[0m jacobian \u001b[38;5;241m=\u001b[39m \u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43meye\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m3\u001b[39;49m\u001b[43m)\u001b[49m \u001b[38;5;241m+\u001b[39m jacobian\n\u001b[1;32m     46\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m jacobian\n",
      "File \u001b[0;32m~/Arbeit/brain-tumor-tissue-reconstruction/.venv/lib/python3.8/site-packages/numpy/lib/twodim_base.py:215\u001b[0m, in \u001b[0;36meye\u001b[0;34m(N, M, k, dtype, order, like)\u001b[0m\n\u001b[1;32m    213\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m M \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    214\u001b[0m     M \u001b[38;5;241m=\u001b[39m N\n\u001b[0;32m--> 215\u001b[0m m \u001b[38;5;241m=\u001b[39m \u001b[43mzeros\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43mN\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mM\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdtype\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdtype\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43morder\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43morder\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    216\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m k \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m M:\n\u001b[1;32m    217\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import ants\n",
    "import numpy as np\n",
    "from scipy.ndimage import map_coordinates\n",
    "from numpy.linalg import svd\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "\n",
    "def get_local_jacobian(deformation_field, point, spacing=(1.0, 1.0, 1.0)):\n",
    "    \"\"\"\n",
    "    Compute the Jacobian matrix of the deformation field at a given point.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    deformation_field : numpy.ndarray\n",
    "        Array of shape [x, y, z, 3] containing the deformation vectors\n",
    "    point : tuple\n",
    "        (x, y, z) coordinates where to compute the Jacobian\n",
    "    spacing : tuple\n",
    "        Voxel spacing in mm\n",
    "        \n",
    "    Returns:\n",
    "    --------\n",
    "    numpy.ndarray\n",
    "        3x3 Jacobian matrix\n",
    "    \"\"\"\n",
    "    x, y, z = point\n",
    "    jacobian = np.zeros((3, 3))\n",
    "    \n",
    "    # Compute partial derivatives using central differences\n",
    "    for i in range(3):  # For each component of the deformation field\n",
    "        # x derivative\n",
    "        if 0 < x < deformation_field.shape[0]-1:\n",
    "            jacobian[0, i] = (deformation_field[x+1, y, z, i] - \n",
    "                            deformation_field[x-1, y, z, i]) / (2 * spacing[0])\n",
    "        # y derivative\n",
    "        if 0 < y < deformation_field.shape[1]-1:\n",
    "            jacobian[1, i] = (deformation_field[x, y+1, z, i] - \n",
    "                            deformation_field[x, y-1, z, i]) / (2 * spacing[1])\n",
    "        # z derivative\n",
    "        if 0 < z < deformation_field.shape[2]-1:\n",
    "            jacobian[2, i] = (deformation_field[x, y, z+1, i] - \n",
    "                            deformation_field[x, y, z-1, i]) / (2 * spacing[2])\n",
    "    \n",
    "    # Add identity matrix since the deformation field represents displacement\n",
    "    jacobian = np.eye(3) + jacobian\n",
    "    return jacobian\n",
    "\n",
    "def get_rotation_from_jacobian(jacobian):\n",
    "    \"\"\"\n",
    "    Extract the rotation matrix from a Jacobian matrix using SVD.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    jacobian : numpy.ndarray\n",
    "        3x3 Jacobian matrix\n",
    "        \n",
    "    Returns:\n",
    "    --------\n",
    "    numpy.ndarray\n",
    "        3x3 rotation matrix\n",
    "    \"\"\"\n",
    "    U, _, Vh = svd(jacobian)\n",
    "    R = np.dot(U, Vh)\n",
    "    \n",
    "    # Ensure proper rotation (determinant = 1)\n",
    "    if np.linalg.det(R) < 0:\n",
    "        U[:, -1] *= -1\n",
    "        R = np.dot(U, Vh)\n",
    "    return R\n",
    "\n",
    "def func2(fixed_image, dti_img, fwd_transform_paths):\n",
    "    \"\"\"\n",
    "    Transform DTI data using a deformation field with proper tensor reorientation.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    \"\"\"\n",
    "    \n",
    "    # Load deformation field using ANTs and convert to numpy\n",
    "    combined_warp = combine_warp(fwd_transform_paths)\n",
    "    def_field_np = combined_warp.numpy()\n",
    "    \n",
    "    # Initialize output array\n",
    "    dti_data = dti_img.numpy()\n",
    "    transformed_dti = np.zeros_like(dti_data)\n",
    "    \n",
    "    # Convert between 6D vector and 3x3 matrix representations\n",
    "    def vector_to_matrix(vector):\n",
    "        return np.array([[vector[0], vector[1], vector[2]],\n",
    "                        [vector[1], vector[3], vector[4]],\n",
    "                        [vector[2], vector[4], vector[5]]])\n",
    "    \n",
    "    def matrix_to_vector(matrix):\n",
    "        return np.array([matrix[0,0], matrix[0,1], matrix[0,2],\n",
    "                        matrix[1,1], matrix[1,2], matrix[2,2]])\n",
    "    \n",
    "    # Create coordinate grid\n",
    "    shape = dti_data.shape[:3]\n",
    "    grid = np.stack(np.meshgrid(\n",
    "        np.arange(shape[0]),\n",
    "        np.arange(shape[1]),\n",
    "        np.arange(shape[2]),\n",
    "        indexing='ij'\n",
    "    ), axis=-1)\n",
    "    \n",
    "    # Apply deformation and tensor reorientation\n",
    "    for x in tqdm(range(shape[0])):\n",
    "        for y in range(shape[1]):\n",
    "            for z in range(shape[2]):\n",
    "                # Get current point\n",
    "                point = (x, y, z)\n",
    "                \n",
    "                # Compute local Jacobian and extract rotation\n",
    "                jacobian = get_local_jacobian(def_field_np, point, fixed_image.spacing)\n",
    "                rotation = get_rotation_from_jacobian(jacobian)\n",
    "                \n",
    "                # Transform tensor\n",
    "                tensor_vec = dti_data[x, y, z, :]\n",
    "                tensor_mat = vector_to_matrix(tensor_vec)\n",
    "                transformed_tensor = np.matmul(\n",
    "                    np.matmul(rotation, tensor_mat),\n",
    "                    rotation.T\n",
    "                )\n",
    "                \n",
    "                # Get transformed coordinates\n",
    "                new_coords = np.array([x, y, z]) + def_field_np[x, y, z]\n",
    "                \n",
    "                # Store transformed tensor if coordinates are within bounds\n",
    "                if (0 <= new_coords[0] < shape[0] and \n",
    "                    0 <= new_coords[1] < shape[1] and \n",
    "                    0 <= new_coords[2] < shape[2]):\n",
    "                    transformed_dti[int(new_coords[0]), \n",
    "                                 int(new_coords[1]), \n",
    "                                 int(new_coords[2])] = matrix_to_vector(transformed_tensor) \n",
    "    \n",
    "    transformed_dti_img = ants.from_numpy(transformed_dti, fixed_image.origin, fixed_image.spacing, fixed_image.direction, has_components=True)\n",
    "    return transformed_dti_img\n",
    "\n",
    "rotated_tensor2 = func2(patient, tensor_image, transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MAE  0.0023838751778361986\n",
      "(240, 240, 155, 6)\n"
     ]
    }
   ],
   "source": [
    "from scipy.ndimage import sobel\n",
    "\n",
    "def compute_jacobian_field(deformation_field):\n",
    "    \"\"\"\n",
    "    Compute the Jacobian field using vectorized gradient operations on the displacement field.\n",
    "    \n",
    "    Parameters:\n",
    "    - displacement_field: A 4D numpy array of shape (X, Y, Z, 3).\n",
    "    \n",
    "    Returns:\n",
    "    - jacobian_field: A 5D numpy array of shape (X, Y, Z, 3, 3).\n",
    "    \"\"\"\n",
    "    jacobian_field = np.zeros((*deformation_field.shape[:3], 3, 3))\n",
    "\n",
    "    # Calculate partial derivatives using finite differences\n",
    "    for i in range(3):  # x, y, z components\n",
    "        jacobian_field[..., i, 0] = sobel(deformation_field[..., i], axis=0, mode='nearest')  # ∂u/∂x, ∂v/∂x, ∂w/∂x\n",
    "        jacobian_field[..., i, 1] = sobel(deformation_field[..., i], axis=1, mode='nearest')  # ∂u/∂y, ∂v/∂y, ∂w/∂y\n",
    "        jacobian_field[..., i, 2] = sobel(deformation_field[..., i], axis=2, mode='nearest')  # ∂u/∂z, ∂v/∂z, ∂w/∂z\n",
    "\n",
    "    # Add identity matrix to account for the fact we’re working with a displacement field\n",
    "    identity = np.eye(3)\n",
    "    jacobian_field += identity[None, None, None, :, :]\n",
    "\n",
    "    return jacobian_field\n",
    "\n",
    "\n",
    "def func3(fixed_image, dti_image, fwd_transforms):\n",
    "    \"\"\"\n",
    "    Vectorized relocation and reorientation of a 6-component DTI image based on a warp field.\n",
    "\n",
    "    Parameters:\n",
    "    - fixed_image: ANTsImage, reference image (e.g., FA or b0)\n",
    "    - dti_image: ANTsImage, the 6-component DTI image (4D image)\n",
    "    - transforms: str-pair, paths to the forward transformations\n",
    "\n",
    "    Returns:\n",
    "    - ANTsImage, the reoriented 6-component DTI image\n",
    "    \"\"\"\n",
    "    # Step 0: Combine Warps\n",
    "    combined_warp = combine_warp(fwd_transforms) \n",
    "\n",
    "    # Step 1: Split the 6-component DTI image\n",
    "    dti_components = ants.split_channels(dti_image)\n",
    "    \n",
    "    # Step 2: Relocate each component by applying the warp field\n",
    "    relocated_components = [\n",
    "        ants.apply_transforms(fixed=fixed_image, moving=component, transformlist=fwd_transforms)\n",
    "        for component in dti_components\n",
    "    ]\n",
    "    \n",
    "    # Convert relocated components to a 4D array: shape = (X, Y, Z, 6)\n",
    "    relocated_dti_array = np.stack([comp.numpy() for comp in relocated_components], axis=-1)\n",
    "\n",
    "    # Step 3: Compute the Jacobian from the warp field as a 4D array\n",
    "    jacobian_array = compute_jacobian_field(combined_warp.numpy())\n",
    "\n",
    "    # Step 4: Construct the 3x3 tensor for each voxel using broadcasting\n",
    "    T = np.zeros((*relocated_dti_array.shape[:3], 3, 3))\n",
    "    T[..., 0, 0] = relocated_dti_array[..., 0]  # Dxx\n",
    "    T[..., 0, 1] = relocated_dti_array[..., 1]  # Dxy\n",
    "    T[..., 1, 0] = relocated_dti_array[..., 1]  # Dxy\n",
    "    T[..., 0, 2] = relocated_dti_array[..., 2]  # Dxz\n",
    "    T[..., 2, 0] = relocated_dti_array[..., 2]  # Dxz\n",
    "    T[..., 1, 1] = relocated_dti_array[..., 3]  # Dyy\n",
    "    T[..., 1, 2] = relocated_dti_array[..., 4]  # Dyz\n",
    "    T[..., 2, 1] = relocated_dti_array[..., 4]  # Dyz\n",
    "    T[..., 2, 2] = relocated_dti_array[..., 5]  # Dzz\n",
    "\n",
    "    # Step 5: Reorient the tensor using batched matrix multiplication\n",
    "    # Reoriented tensor T' = J * T * J^T, performed across all voxels simultaneously\n",
    "    reoriented_T_2 = np.einsum('...ij,...jk,...lk->...il', jacobian_array, T, jacobian_array)\n",
    "\n",
    "    # First, perform the multiplication J * T for each voxel\n",
    "    JT = np.matmul(jacobian_array, T)\n",
    "\n",
    "    # Then, multiply by J^T (transpose of jacobian_field along the last two axes)\n",
    "    reoriented_T = np.matmul(JT, jacobian_array.transpose(0, 1, 2, 4, 3))\n",
    "\n",
    "    print(\"MAE \", np.mean(np.abs(reoriented_T, reoriented_T_2)))\n",
    "\n",
    "    # Step 6: Extract the reoriented components from the reoriented tensor\n",
    "    reoriented_dti = np.zeros((*relocated_dti_array.shape[:3], 6))\n",
    "    reoriented_dti[..., 0] = reoriented_T[..., 0, 0]  # Dxx\n",
    "    reoriented_dti[..., 1] = reoriented_T[..., 0, 1]  # Dxy\n",
    "    reoriented_dti[..., 2] = reoriented_T[..., 0, 2]  # Dxz\n",
    "    reoriented_dti[..., 3] = reoriented_T[..., 1, 1]  # Dyy\n",
    "    reoriented_dti[..., 4] = reoriented_T[..., 1, 2]  # Dyz\n",
    "    reoriented_dti[..., 5] = reoriented_T[..., 2, 2]  # Dzz\n",
    "\n",
    "    print(reoriented_dti.shape)\n",
    "\n",
    "    # Step 7: Convert the reoriented tensor back into an ANTsImage\n",
    "    reoriented_dti_image = ants.from_numpy(reoriented_dti, \n",
    "                                           origin=fixed_image.origin, \n",
    "                                           spacing=fixed_image.spacing, \n",
    "                                           direction=fixed_image.direction, has_components=True)\n",
    "\n",
    "    return reoriented_dti_image\n",
    "\n",
    "rotated_tensor3 = func3(patient, tensor_image, transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ANTsImage (LPI)\n",
      "\t Pixel Type : float (float32)\n",
      "\t Components : 6\n",
      "\t Dimensions : (240, 240, 155)\n",
      "\t Spacing    : (1.0, 1.0, 1.0)\n",
      "\t Origin     : (239.0, 0.0, 0.0)\n",
      "\t Direction  : [-1.  0.  0.  0. -1.  0.  0.  0.  1.]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def compute_jacobian_field(deformation_field, spacing):\n",
    "    \"\"\"\n",
    "    Efficiently compute Jacobian matrices for entire deformation field.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    deformation_field : numpy.ndarray (shape=[x, y, z, 3])\n",
    "        The deformation field\n",
    "    spacing : tuple\n",
    "        Voxel spacing in mm\n",
    "    \n",
    "    Returns:\n",
    "    --------\n",
    "    numpy.ndarray (shape=[x, y, z, 3, 3])\n",
    "        Jacobian matrices for each voxel\n",
    "    \"\"\"\n",
    "    # Initialize gradients array\n",
    "    gradients = np.zeros((*deformation_field.shape[:3], 3, 3))\n",
    "    \n",
    "    # Compute gradients using numpy's gradient function\n",
    "    for i in range(3):\n",
    "        grad = np.gradient(deformation_field[..., i], \n",
    "                         spacing[0], spacing[1], spacing[2])\n",
    "        for j in range(3):\n",
    "            gradients[..., i, j] = grad[j]\n",
    "    \n",
    "    # Add identity matrix to gradients\n",
    "    gradients += np.eye(3)\n",
    "    \n",
    "    return gradients\n",
    "\n",
    "def compute_rotation_field(jacobian_field):\n",
    "    \"\"\"\n",
    "    Extract rotation matrices from Jacobian field using vectorized SVD.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    jacobian_field : numpy.ndarray (shape=[x, y, z, 3, 3])\n",
    "        Field of Jacobian matrices\n",
    "    \n",
    "    Returns:\n",
    "    --------\n",
    "    numpy.ndarray (shape=[x, y, z, 3, 3])\n",
    "        Field of rotation matrices\n",
    "    \"\"\"\n",
    "    shape = jacobian_field.shape[:-2]\n",
    "    \n",
    "    # Reshape for vectorized SVD\n",
    "    J_flat = jacobian_field.reshape(-1, 3, 3)\n",
    "    \n",
    "    # Perform SVD\n",
    "    U, _, Vh = np.linalg.svd(J_flat)\n",
    "    \n",
    "    # Compute rotation matrices\n",
    "    R = np.einsum('...ij,...jk->...ik', U, Vh)\n",
    "    \n",
    "    # Fix improper rotations\n",
    "    det = np.linalg.det(R)\n",
    "    fix_idx = det < 0\n",
    "    U[fix_idx, :, -1] *= -1\n",
    "    R[fix_idx] = np.einsum('...ij,...jk->...ik', U[fix_idx], Vh[fix_idx])\n",
    "    \n",
    "    return R.reshape(*shape, 3, 3)\n",
    "\n",
    "\n",
    "def func4(fixed_image, dti_image, fwd_transforms):\n",
    "    \"\"\"\n",
    "    Transform DTI data using vectorized operations -> Claude implementation \n",
    "    \"\"\"\n",
    "    print(dti_image)\n",
    "\n",
    "    dti_arr = dti_image.numpy()\n",
    "    deformation_field = combine_warp(fwd_transforms).numpy()\n",
    "\n",
    "    shape = dti_arr.shape[:3]\n",
    "    \n",
    "    # Create coordinate grid\n",
    "    grid = np.stack(np.meshgrid(\n",
    "        np.arange(shape[0]),\n",
    "        np.arange(shape[1]),\n",
    "        np.arange(shape[2]),\n",
    "        indexing='ij'\n",
    "    ), axis=-1).astype(float)\n",
    "    \n",
    "    # Apply deformation to coordinates\n",
    "    new_coords = grid + deformation_field\n",
    "    \n",
    "    # Compute Jacobian field\n",
    "    jacobian_field = compute_jacobian_field(deformation_field, fixed_image.spacing)\n",
    "    \n",
    "    # Get rotation field\n",
    "    rotation_field = compute_rotation_field(jacobian_field)\n",
    "    \n",
    "    # Convert 6D vectors to 3x3 matrices\n",
    "    tensor_matrices = np.zeros((*shape, 3, 3))\n",
    "    tensor_matrices[..., 0, 0] = dti_arr[..., 0]  # Dxx\n",
    "    tensor_matrices[..., 0, 1] = tensor_matrices[..., 1, 0] = dti_arr[..., 1]  # Dxy\n",
    "    tensor_matrices[..., 0, 2] = tensor_matrices[..., 2, 0] = dti_arr[..., 2]  # Dxz\n",
    "    tensor_matrices[..., 1, 1] = dti_arr[..., 3]  # Dyy\n",
    "    tensor_matrices[..., 1, 2] = tensor_matrices[..., 2, 1] = dti_arr[..., 4]  # Dyz\n",
    "    tensor_matrices[..., 2, 2] = dti_arr[..., 5]  # Dzz\n",
    "    \n",
    "    # Apply rotation to tensors: R * T * R^T\n",
    "    # First multiply: T * R^T\n",
    "    temp = np.einsum('...ij,...jk->...ik', tensor_matrices, np.transpose(rotation_field, (0,1,2,4,3)))\n",
    "    # Then multiply: R * (T * R^T)\n",
    "    rotated_tensors = np.einsum('...ij,...jk->...ik', rotation_field, temp)\n",
    "    \n",
    "    # Convert back to 6D vectors\n",
    "    transformed_dti = np.zeros_like(dti_arr)\n",
    "    transformed_dti[..., 0] = rotated_tensors[..., 0, 0]  # Dxx\n",
    "    transformed_dti[..., 1] = rotated_tensors[..., 0, 1]  # Dxy\n",
    "    transformed_dti[..., 2] = rotated_tensors[..., 0, 2]  # Dxz\n",
    "    transformed_dti[..., 3] = rotated_tensors[..., 1, 1]  # Dyy\n",
    "    transformed_dti[..., 4] = rotated_tensors[..., 1, 2]  # Dyz\n",
    "    transformed_dti[..., 5] = rotated_tensors[..., 2, 2]  # Dzz\n",
    "    \n",
    "    # Apply spatial transformation using interpolation\n",
    "    coords = new_coords.reshape(-1, 3)\n",
    "    valid_mask = np.all((coords >= 0) & (coords < np.array(shape)), axis=1)\n",
    "    \n",
    "    result = np.zeros_like(transformed_dti)\n",
    "    for i in range(6):\n",
    "        coordinates = np.moveaxis(new_coords, -1, 0)\n",
    "        result[..., i] = map_coordinates(transformed_dti[..., i], coordinates, order=1, mode='nearest')\n",
    "    \n",
    "\n",
    "\n",
    "    return ants.from_numpy(result, (0.0, -239.0, 0.0), fixed_image.spacing, fixed_image.direction, has_components=True)\n",
    "\n",
    "rotated_tensor4 = func4(patient, tensor_image, transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ANTsImage (RAI)\n",
      "\t Pixel Type : float (float32)\n",
      "\t Components : 6\n",
      "\t Dimensions : (240, 240, 155)\n",
      "\t Spacing    : (1.0, 1.0, 1.0)\n",
      "\t Origin     : (0.0, -239.0, 0.0)\n",
      "\t Direction  : [1. 0. 0. 0. 1. 0. 0. 0. 1.]\n",
      "\n",
      "ANTsImage (RAI)\n",
      "\t Pixel Type : float (float32)\n",
      "\t Components : 6\n",
      "\t Dimensions : (240, 240, 155)\n",
      "\t Spacing    : (1.0, 1.0, 1.0)\n",
      "\t Origin     : (0.0, -239.0, 0.0)\n",
      "\t Direction  : [1. 0. 0. 0. 1. 0. 0. 0. 1.]\n",
      "\n",
      "ANTsImage (RAI)\n",
      "\t Pixel Type : float (float32)\n",
      "\t Components : 6\n",
      "\t Dimensions : (240, 240, 155)\n",
      "\t Spacing    : (1.0, 1.0, 1.0)\n",
      "\t Origin     : (-0.0, -239.0, 0.0)\n",
      "\t Direction  : [1. 0. 0. 0. 1. 0. 0. 0. 1.]\n",
      "\n",
      "ANTsImage (RAI)\n",
      "\t Pixel Type : float (float32)\n",
      "\t Components : 6\n",
      "\t Dimensions : (240, 240, 155)\n",
      "\t Spacing    : (1.0, 1.0, 1.0)\n",
      "\t Origin     : (0.0, -239.0, 0.0)\n",
      "\t Direction  : [1. 0. 0. 0. 1. 0. 0. 0. 1.]\n",
      "\n",
      "ANTsImage (RAI)\n",
      "\t Pixel Type : float (float32)\n",
      "\t Components : 1\n",
      "\t Dimensions : (240, 240, 155)\n",
      "\t Spacing    : (1.0, 1.0, 1.0)\n",
      "\t Origin     : (-0.0, -239.0, 0.0)\n",
      "\t Direction  : [1. 0. 0. 0. 1. 0. 0. 0. 1.]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "\n",
    "print(transformed_tensor_img)\n",
    "print(rotated_tensor1)\n",
    "#print(rotated_tensor2)\n",
    "print(rotated_tensor3)\n",
    "print(rotated_tensor4)\n",
    "\n",
    "print(patient)\n",
    "\n",
    "# REORIENT THE DAMN TENSOR BEFORE TRANFORMTAION AND ROTATION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "COMPARISON\n",
      "===========\n",
      "[ 7.8866823e-04  1.6968910e-05  7.9412370e-05  8.7626872e-04\n",
      " -3.3947985e-05  9.1161032e-04]\n",
      "[ 7.8185508e-04  1.1979826e-05  7.6760836e-05  8.7911834e-04\n",
      " -3.0642372e-05  9.1557385e-04]\n",
      "[ 7.8866823e-04  1.6968910e-05  7.9412370e-05  8.7626872e-04\n",
      " -3.3947985e-05  9.1161032e-04]\n",
      "Comparing my method with the Ants method\n",
      "MSE:  6.2243984e-08\n",
      "MAE:  7.001484e-05 \n",
      "\n",
      "\n",
      "Comparing my method with the original tensor image\n",
      "MSE:  6.2242556e-08\n",
      "MAE:  7.0020826e-05\n",
      "Comparing ANTS method with the original tensor image\n",
      "MSE:  1.0579147e-11\n",
      "MAE:  7.350872e-07\n"
     ]
    }
   ],
   "source": [
    "print(\"COMPARISON\\n===========\")\n",
    "\n",
    "x, y, z = 100, 100, 100\n",
    "\n",
    "transformed_tensor_arr = transformed_tensor_img.numpy()\n",
    "rotated_tensor_arr1 = rotated_tensor1.numpy()\n",
    "#rotated_tensor_arr2 = rotated_tensor2.numpy()\n",
    "rotated_tensor_arr3 = rotated_tensor4.numpy()\n",
    "rotated_tensor_arr4 = rotated_tensor4.numpy()\n",
    "\n",
    "print(transformed_tensor_arr[x, y, z])\n",
    "print(rotated_tensor_arr1[x, y, z])\n",
    "print(rotated_tensor_arr2[x, y, z])\n",
    "\n",
    "print(\"Comparing my method with the Ants method\")\n",
    "print(\"MSE: \", np.mean(np.square(rotated_tensor_arr1 - rotated_tensor_arr3)))\n",
    "print(\"MAE: \", np.mean(np.abs(rotated_tensor_arr1 - rotated_tensor_arr3)), \"\\n\\n\")\n",
    "\n",
    "print(\"Comparing my method with the original tensor image\")\n",
    "print(\"MSE: \", np.mean(np.square(transformed_tensor_arr - rotated_tensor_arr3)))\n",
    "print(\"MAE: \", np.mean(np.abs(transformed_tensor_arr - rotated_tensor_arr3)))\n",
    "\n",
    "print(\"Comparing ANTS method with the original tensor image\")\n",
    "print(\"MSE: \", np.mean(np.square(transformed_tensor_arr - rotated_tensor_arr1)))\n",
    "print(\"MAE: \", np.mean(np.abs(transformed_tensor_arr - rotated_tensor_arr1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "ants.image_write(rotated_tensor4, \"out/transformed_rotated_dti_test.nii.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Project TODOs\n",
    "\n",
    "- Lazy Loading of the Images that are transformed\n",
    "- Remove ANTs dependency (RotateTensorImage)\n",
    "- add other modalities\n",
    "- more or better print statements"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}