From ea96c804e9f1ae646cb82e816b55095705f7d5df Mon Sep 17 00:00:00 2001 From: JulianGitW Date: Sun, 12 Jul 2026 23:56:02 +0200 Subject: [PATCH 1/2] Different Alanine CG Mapping-Notebook --- config/data/cg_mapping_script.ipynb | 382 ++++++++++++++++++++++++++++ 1 file changed, 382 insertions(+) create mode 100644 config/data/cg_mapping_script.ipynb diff --git a/config/data/cg_mapping_script.ipynb b/config/data/cg_mapping_script.ipynb new file mode 100644 index 0000000..76e9225 --- /dev/null +++ b/config/data/cg_mapping_script.ipynb @@ -0,0 +1,382 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "06b0c6bb", + "metadata": {}, + "source": [ + "***Create an adjustable mapping matrix from fine to coarse grain representations***" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "92afbdb0", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "import mdtraj as md\n", + "import numpy as np\n", + "import os\n", + "import json" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f21b07b8", + "metadata": {}, + "outputs": [], + "source": [ + "#main path to the project folder\n", + "main_path=\"../..\"" + ] + }, + { + "cell_type": "markdown", + "id": "cd93817c", + "metadata": {}, + "source": [ + "Create Groupings\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "6d29bc69", + "metadata": {}, + "outputs": [], + "source": [ + "#Create various center of mass groupings of Alanine Dipeptide by combining the atoms by their indices\n", + "ramachandran_COM=[[0,1,2,3,4,5],\n", + " [6,7],\n", + " [8,9,10,11,12,13],\n", + " [14,15],\n", + " [16,17,18,19,20,21]]\n", + "ramachandran_list_COM=[\"ramachandran_COM\",ramachandran_COM]\n", + "\n", + "ramachandran_beta_COM=[[0,1,2,3,4,5],\n", + " [6,7],\n", + " [8,9],\n", + " [10,11,12,13], #beta c-atom group for chirality\n", + " [14,15],\n", + " [16,17,18,19,20,21]]\n", + "ramachandran_beta_list_COM=[\"ramachandran_beta_COM\",ramachandran_beta_COM]\n", + "\n", + "backbone_COM=[[0,1,2,3],\n", + " [4,5],\n", + " [6,7],\n", + " [8,9,10,11,12,13],\n", + " [14,15],\n", + " [16,17],\n", + " [18,19,20,21]]\n", + "backbone_list_COM=[\"backbone_COM\",backbone_COM]\n", + "\n", + "backbone_beta_COM=[[0,1,2,3],\n", + " [4,5],\n", + " [6,7],\n", + " [8,9],\n", + " [10,11,12,13],\n", + " [14,15],\n", + " [16,17],\n", + " [18,19,20,21]]\n", + "backbone_beta_list_COM=[\"backbone_beta_COM\",backbone_beta_COM]\n", + "\n", + "short_COM=[[0,1,2,3,4,5,6,7],\n", + " [8,9,10,11,12,13],\n", + " [14,15,16,17,18,19,20,21]]\n", + "short_list_COM=[\"short_COM\",short_COM]\n", + "\n", + "short_beta_COM=[[0,1,2,3,4,5,6,7],\n", + " [8,9],\n", + " [10,11,12,13],\n", + " [14,15,16,17,18,19,20,21]]\n", + "short_beta_list_COM=[\"short_beta_COM\",short_beta_COM]\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "85d5be4d", + "metadata": {}, + "outputs": [], + "source": [ + "#Create coarse grained mappings of sliced atoms: central atom for latent grouping\n", + "ramachandran_slice=[[4,[0,1,2,3,5]],\n", + " [6,[7]],\n", + " [8,[9,10,11,12,13]],\n", + " [14,[15]],\n", + " [16,[17,18,19,20,21]]]\n", + "ramachandran_list_slice=[\"ramachandran_Slice\",ramachandran_slice]\n", + "\n", + "ramachandran_beta_slice=[[4,[0,1,2,3,5]],\n", + " [6,[7]],\n", + " [8,[9]],\n", + " [10,[11,12,13]], #beta c-atom group for chirality\n", + " [14,[15]],\n", + " [16,[17,18,19,20,21]]]\n", + "ramachandran_beta_list_slice=[\"ramachandran_beta_Slice\",ramachandran_beta_slice]\n", + "\n", + "backbone_slice=[[0,[1,2,3]],\n", + " [4,[5]], \n", + " [6,[7]],\n", + " [8,[9,10,11,12,13]], #beta c-atom group for chirality\n", + " [14,[15]],\n", + " [16,[17]],\n", + " [18,[19,20,21]]]\n", + "backbone_list_slice=[\"backbone_Slice\",backbone_slice]\n", + "\n", + "backbone_beta_slice=[[0,[1,2,3]],\n", + " [4,[5]], \n", + " [6,[7]],\n", + " [8,[9]],\n", + " [10,[11,12,13]], #beta c-atom group for chirality\n", + " [14,[15]],\n", + " [16,[17]],\n", + " [18,[19,20,21]]]\n", + "backbone_beta_list_slice=[\"backbone_beta_Slice\",backbone_beta_slice]\n", + "\n", + "short_slice=[[6,[0,1,2,3,4,5,7]],\n", + " [8,[9,10,11,12,13]],\n", + " [14,[15,16,17,18,19,20,21]]]\n", + "short_list_slice=[\"short_Slice\",short_slice]\n", + "\n", + "short_beta_slice=[[6,[0,1,2,3,4,5,7]],\n", + " [8,[9]],\n", + " [10,[11,12,13]],\n", + " [14,[15,16,17,18,19,20,21]]]\n", + "short_beta_list_slice=[\"short_beta_Slice\",short_beta_slice]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0cbe41b7", + "metadata": {}, + "outputs": [], + "source": [ + "def create_splitflow_COM_mapping(main_path, grouping_list: list[str, list[int]]):\n", + " \"\"\"center of mass coarse grain mapping: creates mapping matrix for SplitFlow and a CG PDB for visual verification.\"\"\"\n", + "\n", + " grouping_name = grouping_list[0]\n", + " grouping = grouping_list[1]\n", + " \n", + " # 1. Load Topology and Coordinates\n", + " aa_path=os.path.join(main_path,\"config/data/ala2/ala2.pdb\")\n", + " traj = md.load(aa_path)\n", + " top = traj.topology\n", + " n_atoms = top.n_atoms #get the number of atoms in the topology\n", + " n_beads = len(grouping) #get the number of beads of the predefined coarse grained grouping\n", + " \n", + " # 2. Initialize the Mapping Matrix (Shape: Atoms x Beads)\n", + " # SplitFlow calls: self.map_matrix.T @ x\n", + " # (Atoms x Beads).T @ (Atoms x 3)=(Beads x Atoms) @ (Atoms x 3) = (Beads x 3)\n", + " mapping_matrix = np.zeros((n_atoms, n_beads), dtype=np.float32)\n", + " \n", + " # 3. Fill matrix with Mass-Weighted values\n", + " for bead_idx, atom_indices in enumerate(grouping):\n", + " masses = np.array([top.atom(i).element.mass for i in atom_indices])\n", + " total_mass = np.sum(masses)\n", + " \n", + " for i, atom_idx in enumerate(atom_indices):\n", + " mapping_matrix[atom_idx, bead_idx] = masses[i] / total_mass \n", + "\n", + " # 4. Save the pure matrix for SplitFlow\n", + " matrix_tensor = torch.from_numpy(mapping_matrix)\n", + " #save as a pytorch tensor .pt\n", + " matrix_tensor_path=os.path.join(main_path,f\"config/data/mapping_matrices/ala2_{grouping_name}.pt\")\n", + " torch.save(matrix_tensor, matrix_tensor_path)\n", + " print(f\"Saved mapping matrix\")\n", + "\n", + "\n", + "\n", + " # 5. Create a Visualization PDB\n", + " # We calculate the COM positions for the first frame [0]\n", + " fg_coords = traj.xyz[0] # Shape (22 dialanine Atoms, 3 xyz)\n", + " # Manual COM calculation matching the matrix logic:\n", + " cg_coords = mapping_matrix.T @ fg_coords # Shape (Beads, 3) #this is also done in the SplitFlow\n", + " \n", + " # Create a simple CG topology for the PDB visualization (for simplicity made up of carbon, since masses are not taken into account in backmapping)\n", + " cg_top = md.Topology() \n", + " chain = cg_top.add_chain() #create new molecule topology with chain as bead container\n", + " for i in range(n_beads):\n", + " res = cg_top.add_residue(f\"BD{i+1}\", chain) #several beads called BD1,BD2 etc. inside the add_chain container\n", + " cg_top.add_atom(f\"BEAD\", md.element.carbon, res) #add a representative carbon atom at each bead position\n", + " \n", + " #optional: adding the bonds for a simple cg chain\n", + " for i in range(n_beads - 1):\n", + " atom1 = cg_top.atom(i)\n", + " atom2 = cg_top.atom(i + 1)\n", + " cg_top.add_bond(atom1, atom2)\n", + "\n", + " # Save the CG PDB\n", + " cg_traj = md.Trajectory(cg_coords[np.newaxis, :], cg_top) #create the new beads at the positions of the calculated COMs\n", + " #np.newaxis, because previously deleted axis by traj.xyz[0]\n", + " matrix_pdb_path=os.path.join(main_path,f\"config/data/ala2/ala2_{grouping_name}.pdb\")\n", + " cg_traj.save_pdb(matrix_pdb_path)\n", + " print(f\"Saved visualization PDB\")\n", + "\n", + "\n", + "\n", + "def create_splitflow_sliced_mapping(main_path, grouping_list: list[str, list[int, list]]): \n", + " \"\"\"sliced coarse grain mapping such that the atom group bead is mapped onto the position of the central atom. Creates .pdb file\"\"\"\n", + "\n", + " grouping_name = grouping_list[0]\n", + " grouping=grouping_list[1]\n", + " \n", + " #create new coarsened molecule topology\n", + " cg_top = md.Topology()\n", + " chain=cg_top.add_chain()\n", + " for i in range(len(grouping)):\n", + " res = cg_top.add_residue(f\"BD{i+1}\",chain)\n", + " cg_top.add_atom(f\"BEAD\",md.element.carbon,res)\n", + "\n", + " # 1. Load Topology and Coordinates\n", + " aa_path=os.path.join(main_path,\"config/data/ala2/ala2.pdb\")\n", + " traj=md.load(aa_path)\n", + " aa_coords=traj.xyz[0]\n", + " \n", + " #indices to keep\n", + " idx = []\n", + " for i in grouping:\n", + " idx.append(i[0])\n", + " cg_coords=aa_coords[idx]\n", + "\n", + " #create a .pdb file\n", + " cg_traj=md.Trajectory(cg_coords[np.newaxis, :], cg_top)\n", + " sliced_pdb_path=os.path.join(main_path,f\"config/data/ala2/ala2_{grouping_name}.pdb\")\n", + " cg_traj.save_pdb(sliced_pdb_path)\n", + " print(\"Saved visualization PDB\")\n", + "\n", + " #save the latent_groupings in a json file to be accessed by .yaml\n", + " latent_dir=os.path.join(main_path,\"config/data/latent_groupings\")\n", + " os.makedirs(latent_dir,exist_ok=True)\n", + " json_path = os.path.join(latent_dir, f\"ala2_{grouping_name}.json\")\n", + "\n", + " with open(json_path, 'w') as f:\n", + " # We save only the grouping list (the nested indices), not the name\n", + " json.dump(grouping, f)\n", + " print(\"saved latent_grouping as .json\")" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "6d11dc2c", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Saved mapping matrix\n", + "Saved visualization PDB\n" + ] + } + ], + "source": [ + "create_splitflow_COM_mapping(\n", + " main_path=main_path,\n", + " grouping_list=short_beta_list_COM\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "a71e09b9", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Saved visualization PDB\n", + "saved latent_grouping as .json\n" + ] + } + ], + "source": [ + "create_splitflow_sliced_mapping(\n", + " main_path=main_path,\n", + " grouping_list=short_beta_list_slice\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "9a65f71c", + "metadata": {}, + "source": [ + "Visualizing the coarse grained molecule using VMD and compare it to the all atom configuration:" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "9fc09af5", + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess, platform" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "51ad1931", + "metadata": {}, + "outputs": [], + "source": [ + "def open_vmd(main_path,aa_pdb, cg_pdb=None):\n", + "\n", + " if platform.system == \"Windows\":\n", + " # Standard path for VMD on Windows; adjust if your installation is elsewhere\n", + " vmd_path = r\"C:\\Program Files (x86)\\University of Illinois\\VMD\\vmd.exe\"\n", + " else: #for linux users\n", + " vmd_path = \"vmd\"\n", + " \n", + " aa_pdb_path=os.path.join(main_path, aa_pdb)\n", + " if not os.path.exists(aa_pdb_path):\n", + " print(f\"Error: {aa_pdb_path} not found.\")\n", + " return\n", + "\n", + " # Command: vmd -m aa_file.pdb cg_file.pdb\n", + " cmd = [vmd_path, aa_pdb_path]\n", + " if cg_pdb:\n", + " cg_pdb_path=os.path.join(main_path, cg_pdb)\n", + " if not os.path.exists(cg_pdb_path):\n", + " print(f\"Error: {cg_pdb_path} not found.\")\n", + " return\n", + " cmd.append(cg_pdb_path)\n", + " \n", + " subprocess.Popen(cmd)\n", + "\n", + "# Usage\n", + "open_vmd(main_path,\"config/data/ala2/ala2.pdb\", \"config/data/ala2/ala2_ramachandran_beta_Slice.pdb\") " + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "split-flows-winkler (3.12.12.final.0)", + "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.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 3394e17787b1c0b0b97178d86fa86d70ca5be54a Mon Sep 17 00:00:00 2001 From: JulianGitW Date: Sun, 12 Jul 2026 23:57:48 +0200 Subject: [PATCH 2/2] Included hutchinson_trace calculation --- split_flows/utils/utils.py | 47 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 45 insertions(+), 2 deletions(-) diff --git a/split_flows/utils/utils.py b/split_flows/utils/utils.py index b3fb862..7d26857 100644 --- a/split_flows/utils/utils.py +++ b/split_flows/utils/utils.py @@ -65,8 +65,8 @@ def to_one_hot( def gradient( - output: Tensor, - x: Tensor, + output: Tensor, #velocity field + x: Tensor, #the position of system at time t grad_outputs: Tensor | None = None, create_graph: bool = False, ) -> Tensor: @@ -84,3 +84,46 @@ def gradient( output, x, grad_outputs=grad_outputs, create_graph=create_graph )[0] return grad + + +def hutchinson_trace( + output: Tensor, + x: Tensor, + num_samples: int = 100, +) -> Tensor: + """Compute the trace of the Jacobian using Hutchinson trace estimator. + + Estimates tr(J) = E[z^T * J * z] where z ~ N(0, I). + This is much more efficient than computing the full Jacobian and avoids + graph retention issues with autograd.grad() in loops. + + :param output: Output tensor (typically velocity field), shape (batch, dim) + :param x: Input tensor with respect to which Jacobian is computed, shape (batch, dim) + :param num_samples: Number of random samples for trace estimation (default: 100) + :return: Estimated trace (divergence) for each sample in batch, shape (batch,) + """ + batch_size, dim = output.shape + trace_estimate = torch.zeros(batch_size, device=output.device, dtype=output.dtype) + + for i in range(num_samples): + # Sample random vector z ~ N(0, I) + z = torch.randn_like(output) + + #only retain graph if the output is needed for the next iteration: + is_last_sample = (i == num_samples - 1) + + # Compute z^T * J by computing gradients of (output * z).sum() + # This gives us one row of the Jacobian contracted with z + jvp = torch.autograd.grad( + outputs=output, + inputs=x, + grad_outputs=z, + create_graph=False, + retain_graph=not is_last_sample, + )[0] + + # Compute z^T * (J * z) = (z * jvp).sum(dim=-1) + # This estimates one sample of tr(J) + trace_estimate += (z * jvp).sum(dim=-1) + + return trace_estimate / num_samples