blender_with_mda.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ce6d2ca-5f5f-4234-baee-5a4703e334e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import bpy\n",
    "import numpy as np\n",
    "import MolecularNodes as mn\n",
    "import MDAnalysis as mda\n",
    "\n",
    "# !pip install MDAnalysisTests\n",
    "from MDAnalysis.tests.datafiles import PSF, DCD\n",
    "from MDAnalysis.analysis import distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "404a2610-f445-4efe-a173-12d2e004c234",
   "metadata": {},
   "outputs": [],
   "source": [
    "bpy.context.scene.render.engine = 'CYCLES'\n",
    "bpy.context.scene.cycles.device = \"GPU\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c7cac45-b64c-4967-aeb4-202144709421",
   "metadata": {},
   "outputs": [],
   "source": [
    "u = mda.Universe('blender/md.tpr',\n",
    "                 'blender/md.xtc')\n",
    "\n",
    "custom_selections = {'name_ca': 'name CA',\n",
    "                     'ion': 'resname SOD CLA',\n",
    "                    }\n",
    "mda_session = mn.mda.MDAnalysisSession()\n",
    "mda_session.show(u,\n",
    "                selection='protein',\n",
    "                name='protein',\n",
    "                custom_selections=custom_selections\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a43f953",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show new atomgroups\n",
    "res_100 = u.select_atoms('segid seg_0_PROA and resid 96')\n",
    "acho = u.select_atoms('resname ACHO and resid 1045')\n",
    "\n",
    "mda_session.show(res_100, name='res_100')\n",
    "mda_session.show(acho, name='acho_1045')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d02ffedc",
   "metadata": {},
   "outputs": [],
   "source": [
    "u_2 = mda.Universe(PSF, DCD)\n",
    "\n",
    "mda_session.show(u_2, frame_offset=0)\n",
    "mda_session.show(u_2, frame_offset=40)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51496469-d7d7-4072-908f-63bee7f3a264",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add transformation\n",
    "from MDAnalysis.transformations import PositionAverager\n",
    "\n",
    "transformation = PositionAverager(10, check_reset=True)\n",
    "u.trajectory.add_transformations(transformation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d3fe7ef-acf9-428a-bad4-df111afffc7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "file_path = \"blender_and_mda.blend\"\n",
    "\n",
    "bpy.ops.wm.save_as_mainfile(filepath=file_path)\n",
    "\n",
    "# the .blend file can be opened by\n",
    "# file_path = \"blender_and_mda.blend\"\n",
    "# bpy.ops.wm.open_mainfile(filepath=file_path)\n",
    "# mda_session = bpy.context.scene.mda_session"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f87efde-2b49-4fdc-b1f9-fbe54ea918d7",
   "metadata": {},
   "source": [
    "## Visualize analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db2050ec-d78a-42bb-837f-413a965f5f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Distance_in_Blender:\n",
    "    def __init__(self, coord1_arr, coord2_arr, dist_arr, world_scale=0.01):\n",
    "        self.coord1_arr = coord1_arr\n",
    "        self.coord2_arr = coord2_arr\n",
    "        self.dist_arr = dist_arr\n",
    "        self.n_frames = len(dist_arr)\n",
    "        self.world_scale = world_scale\n",
    "        self.draw()\n",
    "        bpy.app.handlers.frame_change_post.append(self.update_position_handler_wrapper())\n",
    "\n",
    "    def draw(self, name='distance'):\n",
    "        line_data = bpy.data.curves.new(name=name, type='CURVE')\n",
    "        line_data.dimensions = '3D'\n",
    "        line_object = bpy.data.objects.new(f\"{name}_line\", line_data)\n",
    "        bpy.context.collection.objects.link(line_object)\n",
    "        bpy.context.view_layer.objects.active = line_object\n",
    "\n",
    "        text_data = bpy.data.curves.new(name=name, type='FONT')\n",
    "        text_object = bpy.data.objects.new(f\"{name}_text\", text_data)\n",
    "        bpy.context.collection.objects.link(text_object)\n",
    "        \n",
    "        text_data.size = 10 * self.world_scale\n",
    "        text_data.align_x = 'CENTER'\n",
    "        text_data.align_y = 'CENTER'\n",
    "        \n",
    "        line = line_data.splines.new('POLY')\n",
    "        line.points.add(1)\n",
    "\n",
    "        self.line_data = line_data\n",
    "        self.line_object = line_object\n",
    "        self.text_data = text_data\n",
    "        self.text_object = text_object\n",
    "        self.line = line\n",
    "        self.update_position(0)\n",
    "\n",
    "        line.resolution_u = 4\n",
    "        line.use_cyclic_u = False\n",
    "        line.use_endpoint_u = True\n",
    "        line.use_endpoint_v = True\n",
    "        line.use_smooth = False\n",
    "    \n",
    "    def update_position(self, frame):\n",
    "        coord1 = self.coord1_arr[frame]\n",
    "        coord2 = self.coord2_arr[frame]\n",
    "        dist = np.mean(self.dist_arr[frame])\n",
    "        \n",
    "        for point, coord in zip(self.line.points, [coord1, coord2]):\n",
    "            coord = list(coord * self.world_scale)\n",
    "            point.co = (coord[0], coord[1], coord[2], 1)\n",
    "        self.text_data.body = f'{dist:.2f} Å'\n",
    "        \n",
    "        middle_coord = np.mean([coord1, coord2], axis=0)\n",
    "        self.text_object.location = middle_coord * self.world_scale\n",
    "        # rotate the text object to face the camera\n",
    "        self.text_object.rotation_euler = bpy.context.scene.camera.rotation_euler\n",
    "\n",
    "    def update_position_handler_wrapper(self):\n",
    "        def update_position_handler(scene):\n",
    "            frame = scene.frame_current\n",
    "            if frame >= self.n_frames:\n",
    "                return None\n",
    "            self.update_position(frame)\n",
    "        return update_position_handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3338e4b6-3a4c-4d33-9c78-4d5b5a0af37c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run analysis\n",
    "coord1_arr = []\n",
    "coord2_arr = []\n",
    "dist_arr = []\n",
    "for ts in u.trajectory:\n",
    "    com_res_100 = res_100.center_of_mass()\n",
    "    com_acho = acho.center_of_mass()\n",
    "    dist_arr.append(distances.distance_array(com_res_100, com_acho, box=u.dimensions))\n",
    "\n",
    "    coord1_arr.append(com_res_100)\n",
    "    coord2_arr.append(com_acho)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a877e467-dd1f-4b9c-b533-6d066b20c1f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "dist_render = Distance_in_Blender(coord1_arr,\n",
    "                                  coord2_arr,\n",
    "                                  dist_arr,\n",
    "                                  world_scale=mda_session.world_scale)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "blender_3.6.1",
   "language": "python",
   "name": "blender_3.6.1"
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}