Skip to content

CG mapping matrices and hutchinson trace calculation #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JulianGitW wants to merge 3 commits into
base: public
Choose a base branch
from
Open

CG mapping matrices and hutchinson trace calculation #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ea96c80
Different Alanine CG Mapping-Notebook
JulianGitW Jul 12, 2026
3394e17
Included hutchinson_trace calculation
JulianGitW Jul 12, 2026
ad246d6
Merge branch 'public' of https://github.com/JulianGitW/split-flows_Wi…
JulianGitW Jul 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
382 changes: 382 additions & 0 deletions config/data/cg_mapping_script.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
}
Loading
Loading