ribosome tilt-series simulation

figure-it-out-rln40.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 448,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import gemmi\n",
    "import mrcfile\n",
    "import starfile\n",
    "import einops\n",
    "import eulerangles\n",
    "from scipy.stats import special_ortho_group\n",
    "from scipy.spatial.transform import Rotation as R\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 449,
   "metadata": {},
   "outputs": [],
   "source": [
    "rootpath = Path(\"..\").resolve()\n",
    "tomoroot = rootpath.joinpath(\"simtomo40\")\n",
    "pdbdir = tomoroot.joinpath(\"Pdb\", \"job000\")\n",
    "pdbdir.mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 450,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdbfn = '4v6x-ribo.cif'\n",
    "pdbpath = pdbdir.joinpath(pdbfn)\n",
    "if not Path(pdbpath).exists():\n",
    "    import urllib\n",
    "    urllib.request.urlretrieve(\"https://files.rcsb.org/download/4V6X.cif\", pdbpath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 451,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<gemmi.Model 1 with 89 chain(s)>"
      ]
     },
     "execution_count": 451,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "structure = gemmi.read_structure(str(pdbpath))\n",
    "model = structure[0]\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 452,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(13338, 3)"
      ]
     },
     "execution_count": 452,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ca_coords = np.array([\n",
    "    [cra.atom.pos.x, cra.atom.pos.y, cra.atom.pos.z]\n",
    "    for cra in model.all()\n",
    "    if cra.atom.name == 'CA'\n",
    "])\n",
    "ca_coords.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 453,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(13338, 3)"
      ]
     },
     "execution_count": 453,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ca_coords_10apx = ca_coords / 10\n",
    "ca_coords_10apx_centered = ca_coords_10apx - np.mean(ca_coords_10apx, axis=0)\n",
    "assert np.allclose(np.mean(ca_coords_10apx_centered, axis=0), [0, 0, 0])\n",
    "ca_coords_10apx_centered.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 454,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model_coords = ca_coords_10apx_centered"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 455,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([-13.74219601, -12.59160924, -17.41221623]),\n",
       " array([12.15490399, 14.24029076, 12.01938377]))"
      ]
     },
     "execution_count": 455,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.min(model_coords, axis=0), np.max(model_coords, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### generate 64 ribos in a 500^3 box centered on 0,0,0 ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 456,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "n = 64\n",
    "model_positions = np.random.uniform(low=-200, high=200, size=(n, 3))\n",
    "model_orientations = special_ortho_group.rvs(dim=3, size=n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 457,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(64, 3, 3)"
      ]
     },
     "execution_count": 457,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_orientations.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 458,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model_orientations = model_orientations.reshape((-1, 1, 3, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 459,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(13338, 64, 3)"
      ]
     },
     "execution_count": 459,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "oriented_model_coordinates = model_orientations @ model_coords.reshape((-1, 3, 1))\n",
    "oriented_model_coordinates = np.squeeze(oriented_model_coordinates).transpose(1, 0, 2)\n",
    "oriented_model_coordinates.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 460,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((64, 3), (13338, 64, 3))"
      ]
     },
     "execution_count": 460,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "positioned_oriented_model_coordinates = oriented_model_coordinates + model_positions\n",
    "model_positions.shape, positioned_oriented_model_coordinates.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 461,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "final_coordinates = positioned_oriented_model_coordinates.reshape((-1, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 462,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def render_on_grid_3d(coords: np.ndarray) -> np.ndarray:\n",
    "    b = np.linspace(-250, 250, num=501, endpoint=True)\n",
    "    image, _ = np.histogramdd(coords, bins=[b, b, b])\n",
    "    return image\n",
    "\n",
    "\n",
    "def render_on_grid_2d(coords: np.ndarray) -> np.ndarray:\n",
    "    b = np.linspace(-250.5, 250.5, num=501, endpoint=True)\n",
    "    image, _ = np.histogramdd(coords[..., :2], bins=[b, b])\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 463,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(500, 500)"
      ]
     },
     "execution_count": 463,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test = render_on_grid_2d(final_coordinates)\n",
    "test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 464,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skip Napari\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    import napari\n",
    "\n",
    "    viewer = napari.Viewer()\n",
    "    viewer.add_image(test)\n",
    "except ImportError:\n",
    "    print(\"Skip Napari\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 465,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def Rx(angles_degrees: np.ndarray) -> np.ndarray:\n",
    "    \"\"\"Affine matrix for a rotation around the X-axis.\"\"\"\n",
    "    angles_degrees = np.asarray(angles_degrees).reshape(-1)\n",
    "    c = np.cos(np.deg2rad(angles_degrees))\n",
    "    s = np.sin(np.deg2rad(angles_degrees))\n",
    "    matrices = einops.repeat(\n",
    "        np.eye(4), 'i j -> n i j', n=len(angles_degrees)\n",
    "    )\n",
    "    matrices[:, 1, 1] = c\n",
    "    matrices[:, 1, 2] = -s\n",
    "    matrices[:, 2, 1] = s\n",
    "    matrices[:, 2, 2] = c\n",
    "    return np.squeeze(matrices)\n",
    "\n",
    "\n",
    "def Ry(angles_degrees: np.ndarray) -> np.ndarray:\n",
    "    \"\"\"Affine matrix for a rotation around the Y-axis.\"\"\"\n",
    "    angles_degrees = np.asarray(angles_degrees).reshape(-1)\n",
    "    c = np.cos(np.deg2rad(angles_degrees))\n",
    "    s = np.sin(np.deg2rad(angles_degrees))\n",
    "    matrices = einops.repeat(\n",
    "        np.eye(4), 'i j -> n i j', n=len(angles_degrees)\n",
    "    )\n",
    "    matrices[:, 0, 0] = c\n",
    "    matrices[:, 0, 2] = s\n",
    "    matrices[:, 2, 0] = -s\n",
    "    matrices[:, 2, 2] = c\n",
    "    return np.squeeze(matrices)\n",
    "\n",
    "\n",
    "def Rz(angles_degrees: float) -> np.ndarray:\n",
    "    \"\"\"Affine matrix for a rotation around the Z-axis.\"\"\"\n",
    "    angle_degrees = np.asarray(angles_degrees).reshape(-1)\n",
    "    c = np.cos(np.deg2rad(angle_degrees))\n",
    "    s = np.sin(np.deg2rad(angle_degrees))\n",
    "    matrices = einops.repeat(\n",
    "        np.eye(4), 'i j -> n i j', n=len(angle_degrees)\n",
    "    )\n",
    "    matrices[:, 0, 0] = c\n",
    "    matrices[:, 0, 1] = -s\n",
    "    matrices[:, 1, 0] = s\n",
    "    matrices[:, 1, 1] = c\n",
    "    return np.squeeze(matrices)\n",
    "\n",
    "\n",
    "def S(shifts: np.ndarray) -> np.ndarray:\n",
    "    \"\"\"Affine matrices for shifts.\n",
    "    Shifts supplied can be 2D or 3D.\n",
    "    \"\"\"\n",
    "    shifts = np.asarray(shifts, dtype=float)\n",
    "    if shifts.shape[-1] == 2:\n",
    "        shifts = _promote_2d_to_3d(shifts)\n",
    "    shifts = np.array(shifts).reshape((-1, 3))\n",
    "    matrices = einops.repeat(np.eye(4), 'i j -> n i j', n=shifts.shape[0])\n",
    "    matrices[:, 0:3, 3] = shifts\n",
    "    return np.squeeze(matrices)\n",
    "\n",
    "\n",
    "def _promote_2d_to_3d(shifts: np.ndarray) -> np.ndarray:\n",
    "    \"\"\"Promote 2D vectors to 3D with zeros in the last dimension.\"\"\"\n",
    "    shifts = np.asarray(shifts).reshape(-1, 2)\n",
    "    shifts = np.c_[shifts, np.zeros(shifts.shape[0])]\n",
    "    return np.squeeze(shifts)\n",
    "\n",
    "def homogenise_coordinates(coords: np.ndarray) -> np.ndarray:\n",
    "    return np.c_[coords, np.ones(len(coords))]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 466,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# simulate tilt-series\n",
    "def simulate_tilt_image(coords, rx=0, ry=0, rz=0, dx=0, dy=0):\n",
    "    coords = homogenise_coordinates(coords)\n",
    "    transformation = S([dx, dy, 0]) @ Rz(rz) @ Ry(ry) @ Rx(rx)\n",
    "    transformed_coords = np.squeeze(transformation @ coords.reshape((-1, 4, 1)))\n",
    "    return render_on_grid_2d(transformed_coords)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 467,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "test_0 = simulate_tilt_image(final_coordinates, ry=0)\n",
    "test_60 = simulate_tilt_image(final_coordinates, ry=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 468,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skip Napari\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    viewer = napari.Viewer()\n",
    "    viewer.add_image(test_0)\n",
    "    viewer.add_image(test_60)\n",
    "except NameError:\n",
    "    print(\"Skip Napari\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 469,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "raw_tilt_series = np.stack(\n",
    "    [\n",
    "        simulate_tilt_image(final_coordinates, ry=r)\n",
    "        for r\n",
    "        in np.arange(-90, 90, 3)\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 470,
   "metadata": {},
   "outputs": [],
   "source": [
    "tilt_series = raw_tilt_series * -1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 471,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(64, 60, 2)"
      ]
     },
     "execution_count": 471,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tilt_series = tilt_series.transpose(0, 2, 1)\n",
    "particle_positions = model_positions + np.array([250, 250, 250])\n",
    "\n",
    "transformations = S([250, 250]) @ Ry(np.arange(-90, 90, 3)) @ S([-250, -250, -250])\n",
    "projected_positions = np.squeeze(transformations @ homogenise_coordinates(particle_positions).reshape((-1, 1, 4, 1)))[..., :2]\n",
    "projected_positions.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 472,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3840, 3)"
      ]
     },
     "execution_count": 472,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "visualised_positions = np.zeros((n, 60, 3))\n",
    "visualised_positions[..., 0] = np.arange(60)\n",
    "visualised_positions[..., 1:] = projected_positions[..., ::-1]\n",
    "visualised_positions = visualised_positions.reshape((-1, 3))\n",
    "visualised_positions.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 473,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Skip Napari\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    viewer = napari.Viewer()\n",
    "    viewer.add_image(tilt_series)\n",
    "    viewer.add_points(visualised_positions)\n",
    "except NameError:\n",
    "    print(\"Skip Napari\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 474,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "importroot = tomoroot.joinpath(\"ImportTomo\", \"job000\")\n",
    "importroot.mkdir(parents=True, exist_ok=True)\n",
    "tomogram_dir = importroot.joinpath(\"tomograms\", \"TS_01\")\n",
    "tomogram_dir.mkdir(parents=True, exist_ok=True)\n",
    "tomofn = tomogram_dir / \"01.mrc\"\n",
    "mrcfile.write(str(tomofn), tilt_series.astype(np.float32), voxel_size=10, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 475,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "global_star_file = importroot.joinpath(\"aligned_tilt_series.star\")\n",
    "tsdir = importroot.joinpath(\"tilt_series\")\n",
    "tsdir.mkdir(parents=True, exist_ok=True)\n",
    "individual_star_file = tsdir.joinpath(\"TS_01.star\")\n",
    "particle_star_file = importroot.joinpath(\"particles.star\")\n",
    "optimisation_set_star_file = importroot.joinpath(\"optimisation_set.star\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 476,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "alignment_data = pd.DataFrame({\n",
    "    'rlnMicrographPreExposure': [0 for _ in range(len(np.arange(-90, 90, 3)))],\n",
    "    'rlnDefocusU': [0 for _ in range(len(np.arange(-90, 90, 3)))],\n",
    "    'rlnDefocusV': [0 for _ in range(len(np.arange(-90, 90, 3)))],\n",
    "    'rlnDefocusAngle': [0 for _ in range(len(np.arange(-90, 90, 3)))],\n",
    "    \"rlnPhaseShift\": [90 for _ in range(len(np.arange(-90, 90, 3)))],\n",
    "    \"rlnCtfScalefactor\": [1 for _ in range(len(np.arange(-90, 90, 3)))],\n",
    "})\n",
    "\n",
    "projection_matrix_labels = [f'rlnTomoProj{ax}' for ax in 'XYZW']\n",
    "for idx, label in enumerate(projection_matrix_labels):\n",
    "    rows = transformations[:, idx, :]\n",
    "    alignment_data[label] = [\n",
    "        f'[{r[0]:.13g},{r[1]:.13g},{r[2]:.13g},{r[3]:.13g}]'\n",
    "        for r in rows\n",
    "    ]\n",
    "\n",
    "starfile.write({'TS_01': alignment_data}, individual_star_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 477,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "global_data = pd.DataFrame({\n",
    "    'rlnTomoName': ['TS_01'],\n",
    "    'rlnTomoTiltSeriesName': [tomofn.relative_to(tomoroot)],\n",
    "    \"rlnTomoFrameCount\": (tilt_series.shape[0],),\n",
    "    \"rlnTomoSizeX\": (tilt_series.shape[1],),\n",
    "    \"rlnTomoSizeY\": (tilt_series.shape[2],),\n",
    "    \"rlnTomoSizeZ\": (tilt_series.shape[2],),\n",
    "    'rlnVoltage': [300],\n",
    "    'rlnSphericalAberration': [2.7],\n",
    "    'rlnAmplitudeContrast': [0.1],\n",
    "    'rlnMicrographOriginalPixelSize': [10],\n",
    "    'rlnTomoHand': [1],\n",
    "    'rlnTomoTiltSeriesPixelSize': [10],\n",
    "})\n",
    "\n",
    "starfile.write({'global': global_data, 'TS_01': alignment_data}, global_star_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 478,
   "metadata": {},
   "outputs": [],
   "source": [
    "inverse_model_orientations = model_orientations.copy()\n",
    "inverse_model_orientations[:, :, :, [2, 0]] = model_orientations[:, :, :, [0, 2]]\n",
    "inverse_model_orientations[:, :, [2, 0], :] = inverse_model_orientations[:, :, [0, 2], :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 479,
   "metadata": {},
   "outputs": [],
   "source": [
    "particle_eulers = eulerangles.matrix2euler(\n",
    "        rotation_matrices=np.linalg.pinv(model_orientations),\n",
    "        axes='zyz',\n",
    "        intrinsic=True,\n",
    "        right_handed_rotation=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 480,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "particle_optics = pd.DataFrame({\n",
    "    'rlnOpticsGroup': [1],\n",
    "    'rlnOpticsGroupName': ['optics'],\n",
    "    'rlnSphericalAberration': [2.7],\n",
    "    'rlnVoltage': [300],\n",
    "    'rlnTomoTiltSeriesPixelSize': [10],\n",
    "})\n",
    "\n",
    "# particle_eulers = eulerangles.matrix2euler(\n",
    "#         rotation_matrices=np.linalg.pinv(model_orientations),\n",
    "#         axes='zyz',\n",
    "#         intrinsic=True,\n",
    "#         right_handed_rotation=True,\n",
    "#     )\n",
    "particle_data = pd.DataFrame({\n",
    "    'rlnTomoName': ['TS_01' for _ in range(len(particle_positions))],\n",
    "    \"rlnTomoParticleName\": [\"TS_01/%06d\"%(no,) for no in range(len(particle_positions))],\n",
    "    'rlnTomoParticleId': [i for i in range(len(particle_positions))],\n",
    "    \"rlnTomoManifoldIndex\": [1 for _ in range(len(particle_positions))],\n",
    "    'rlnCoordinateX': particle_positions[:, 0],\n",
    "    'rlnCoordinateY': particle_positions[:, 1],\n",
    "    'rlnCoordinateZ': particle_positions[:, 2],\n",
    "    'rlnAngleRot': particle_eulers[:, 0],\n",
    "    'rlnAngleTilt': particle_eulers[:, 1],\n",
    "    'rlnAnglePsi': particle_eulers[:, 2],\n",
    "    \"rlnOriginXAngst\": [0. for _ in range(len(particle_positions))],\n",
    "    \"rlnOriginYAngst\": [0. for _ in range(len(particle_positions))],\n",
    "    \"rlnOriginZAngst\": [0. for _ in range(len(particle_positions))],\n",
    "    'rlnDefocusU': [0 for _ in range(len(particle_positions))],\n",
    "    'rlnDefocusV': [0 for _ in range(len(particle_positions))],\n",
    "    'rlnDefocusAngle': [0 for _ in range(len(particle_positions))],\n",
    "    \"rlnPhaseShift\": [90 for _ in range(len(particle_positions))],\n",
    "    \"rlnOpticsGroup\": [1 for _ in range(len(particle_positions))],\n",
    "    \"rlnClassNumber\": [-1 for _ in range(len(particle_positions))],\n",
    "    \"rlnRandomSubset\": [random.randint(1, 2) for _ in range(len(particle_positions))],\n",
    "})\n",
    "\n",
    "starfile.write({'optics': particle_optics, 'particles': particle_data}, particle_star_file, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 481,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimisation_set_data = pd.DataFrame({\n",
    "    \"rlnTomoParticlesFile\": (particle_star_file.relative_to(tomoroot), ),\n",
    "    \"rlnTomoTomogramsFile\": (global_star_file.relative_to(tomoroot), ),\n",
    "})\n",
    "\n",
    "starfile.write(optimisation_set_data, optimisation_set_star_file, overwrite=True, force_loop=False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.12 64-bit",
   "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.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "df0893f56f349688326838aaeea0de204df53a132722cbd565e54b24a8fec5f6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}