{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "separated-percentage",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats.qmc import PoissonDisk\n",
    "import numpy as np\n",
    "from scipy.spatial.distance import cdist\n",
    "import napari"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "flush-howard",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "legitimate-defense",
   "metadata": {},
   "outputs": [],
   "source": [
    "scale = 128\n",
    "radius = 10 / scale\n",
    "border = 1.5\n",
    "fraction = 0.5\n",
    "\n",
    "k0=0.85\n",
    "k1 = 1.5\n",
    "delta = 0.01\n",
    "mean_scale = 0.95\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "adapted-tourism",
   "metadata": {},
   "outputs": [],
   "source": [
    "obj = PoissonDisk(d=3, radius=radius, hypersphere='volume', ncandidates=30, optimization=None, seed=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "comparable-coast",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "tmp = obj.fill_space()*scale\n",
    "selz = (tmp[:,0] > border*scale*radius) & (tmp[:,0]<(scale-border*scale*radius))\n",
    "sely = (tmp[:,1] > border*scale*radius) & (tmp[:,1]<(scale-border*scale*radius))\n",
    "selx = (tmp[:,2] > border*scale*radius) & (tmp[:,2]<(scale-border*scale*radius))\n",
    "sel = selz & sely & selx \n",
    "tmp = tmp[sel]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "existing-clerk",
   "metadata": {},
   "outputs": [],
   "source": [
    "volume = np.zeros( (int(scale), int(scale), int(scale)) )\n",
    "labels = np.zeros_like(volume)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "intense-soviet",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "sphere_radius = radius*scale*k0/2.0  # Set your desired radius for the sphere\n",
    "k2 = 1.0 / np.sqrt(k1)\n",
    "major_axis = sphere_radius*k1\n",
    "minor_axis = sphere_radius*k2\n",
    "\n",
    "\n",
    "\n",
    "def fill_sphere(center, radius, volume, labels, label =1 ):\n",
    "    # Create an array with the coordinates of each point in the volume\n",
    "    x, y, z = np.indices(volume.shape)\n",
    "\n",
    "    # Calculate the squared distance from each point to the center\n",
    "    dist_sq = (x - center[0])**2 + (y - center[1])**2 + (z - center[2])**2\n",
    "\n",
    "    # A point is inside the sphere if its distance to the center is less than the radius\n",
    "    inside_sphere = dist_sq < (radius**2)\n",
    "\n",
    "    # Set the value of points inside the sphere to 1\n",
    "    volume[inside_sphere] = 1   \n",
    "    labels[inside_sphere] = label   \n",
    "\n",
    "def fill_ellipsoid_simple(center, major_axis, minor_axis, volume, labels, label = 2):\n",
    "    # Create an array with the coordinates of each point in the volume\n",
    "    x, y, z = np.indices(volume.shape)\n",
    "\n",
    "    # Calculate the normalized squared distances along each axis\n",
    "    dist_sq_x = ((x - center[0]) / major_axis)**2\n",
    "    dist_sq_y = ((y - center[1]) / minor_axis)**2\n",
    "    dist_sq_z = ((z - center[2]) / minor_axis)**2\n",
    "\n",
    "    # A point is inside the ellipsoid if the sum of the squared distances is less than 1\n",
    "    inside_ellipsoid = dist_sq_x + dist_sq_y + dist_sq_z < 1\n",
    "\n",
    "    # Set the value of points inside the ellipsoid to 1\n",
    "    volume[inside_ellipsoid] = 1\n",
    "    labels[inside_ellipsoid] = label\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "def random_rotation_matrix():\n",
    "    theta, phi, z = np.random.uniform(0, 2*np.pi, 3)\n",
    "\n",
    "    # Rotation about the z-axis\n",
    "    rz = np.array([\n",
    "        [np.cos(z), -np.sin(z), 0],\n",
    "        [np.sin(z), np.cos(z), 0],\n",
    "        [0, 0, 1]\n",
    "    ])\n",
    "\n",
    "    # Rotation about the y-axis\n",
    "    ry = np.array([\n",
    "        [np.cos(phi), 0, np.sin(phi)],\n",
    "        [0, 1, 0],\n",
    "        [-np.sin(phi), 0, np.cos(phi)]\n",
    "    ])\n",
    "\n",
    "    # Rotation about the x-axis\n",
    "    rx = np.array([\n",
    "        [1, 0, 0],\n",
    "        [0, np.cos(theta), -np.sin(theta)],\n",
    "        [0, np.sin(theta), np.cos(theta)]\n",
    "    ])\n",
    "\n",
    "    return np.dot(rz, np.dot(ry, rx))\n",
    "\n",
    "def fill_ellipsoid(center, major_axis, minor_axis, volume, labels, label=2):\n",
    "    # Create an array with the coordinates of each point in the volume\n",
    "    x, y, z = np.indices(volume.shape).astype(float)\n",
    "\n",
    "    # Center the points\n",
    "    x -= center[0]\n",
    "    y -= center[1]\n",
    "    z -= center[2]\n",
    "\n",
    "    # Apply rotation\n",
    "    rotation_matrix = random_rotation_matrix()\n",
    "    rotated_coords = np.dot(rotation_matrix, np.array([x.ravel(), y.ravel(), z.ravel()]))\n",
    "    \n",
    "    x_rotated, y_rotated, z_rotated = rotated_coords.reshape(3, *volume.shape)\n",
    "\n",
    "    # Calculate the normalized squared distances\n",
    "    dist_sq_x = (x_rotated / major_axis)**2\n",
    "    dist_sq_y = (y_rotated / minor_axis)**2\n",
    "    dist_sq_z = (z_rotated / minor_axis)**2\n",
    "\n",
    "    # Check if points are inside the ellipsoid\n",
    "    inside_ellipsoid = dist_sq_x + dist_sq_y + dist_sq_z < 1\n",
    "\n",
    "    # Set the value of points inside the ellipsoid to 1\n",
    "    volume[inside_ellipsoid] = 1\n",
    "    labels[inside_ellipsoid] = label\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "integral-dynamics",
   "metadata": {},
   "outputs": [],
   "source": [
    "for point in tmp:\n",
    "    shape = 'sphere' if np.random.rand() > fraction else 'ellipsoid'\n",
    "    multi = np.random.rand()\n",
    "    multi = multi*delta - delta / 2.0 + mean_scale\n",
    "    \n",
    "    if shape == 'sphere':\n",
    "        fill_sphere(center=point, radius=sphere_radius*multi, volume=volume, labels=labels)\n",
    "    else:\n",
    "        # For ellipsoids, we can randomize the orientation and axis lengths\n",
    "        major_axis = sphere_radius * k1 * multi\n",
    "        minor_axis = sphere_radius * k2 * multi\n",
    "        fill_ellipsoid(center=point, \n",
    "                       major_axis=major_axis, \n",
    "                       minor_axis=minor_axis, \n",
    "                       volume=volume,\n",
    "                       labels = labels\n",
    "                      )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "younger-dialogue",
   "metadata": {},
   "outputs": [],
   "source": [
    "v = napari.view_image(volume)\n",
    "_ = v.add_labels(labels.astype(np.int8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "advance-latest",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(\"benchmark_volume\", volume)\n",
    "np.save(\"benchmark_labels\", labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "molecular-visiting",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dlsia-new",
   "language": "python",
   "name": "dlsia-new"
  },
  "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}