feat: color by pLDDT

backend/src/api/protein.py

@@ -105,11 +105,30 @@ def str_as_file_stream(input_str: str, filename_as: str) -> StreamingResponse:
  )
 
 
+def get_residue_bfactors(pdb: PDB):
+ chains = {}
+ for chain in pdb.structure.get_chains():
+ chains[chain.get_id()] = []
+ for r in chain.get_residues():
+ for a in r.get_atoms():
+ chains[chain.get_id()].append(a.bfactor)
+ break
+ return chains
+
+
 """
  ENDPOINTS TODO: add the other protein types here instead of in api_types.py
 """
 
 
+@router.get("/protein/pLDDT/{protein_name:str}", response_model=dict[str, list[float]])
+def get_pLDDT_given_protein(protein_name: str):
+ if protein_name_found(protein_name):
+ pdb = parse_protein_pdb(protein_name, encoding="file")
+ return get_residue_bfactors(pdb)
+ return {}
+
+
 class UploadPNGBody(CamelModel):
  protein_name: str
  base64_encoding: str

frontend/src/lib/Molstar.svelte

@@ -1,7 +1,12 @@
 <script lang="ts">
  import { onDestroy } from "svelte";
  import { PDBeMolstarPlugin } from "../../venome-molstar/lib";
- import { loseWebGLContext } from "./venomeMolstarUtils";
+ import {
+ loseWebGLContext,
+ colorResidues,
+ type ChainColors,
+ } from "./venomeMolstarUtils";
+ import type { QueryParam } from "../../venome-molstar/lib/helpers";
 
  export let url = "";
  export let format = "pdb";
@@ -10,7 +15,9 @@
  export let width = 500;
  export let height = 500;
  export let hideControls = true;
+ export let chainColors: ChainColors = {};
  let m: PDBeMolstarPlugin;
+ let subscribe: ReturnType<typeof colorByChain>;
 
  let divEl: HTMLDivElement;
  async function render() {
@@ -28,22 +35,42 @@
  bgColor,
  subscribeEvents: false,
  selectInteraction: true,
- alphafoldView: true,
+ alphafoldView: false,
  reactive: true,
  sequencePanel: true,
  hideControls,
  hideCanvasControls: ["animation"],
  });
  }
 
+ function colorByChain(chainColors: ChainColors) {
+ let allColors: QueryParam[] = [];
+ for (const [chainId, rgbPerResidue] of Object.entries(chainColors)) {
+ const colors = colorResidues({
+ struct_asym_id: chainId,
+ colors: rgbPerResidue,
+ });
+ // add to all colors
+ allColors = [...allColors, ...colors];
+ }
+ return m.events.loadComplete.subscribe(() => {
+ m.visual.select({ data: allColors });
+ console.log("color");
+ });
+ }
+
  onDestroy(() => {
  loseWebGLContext(divEl.querySelector("canvas")!);
  m.plugin.dispose();
+ subscribe.unsubscribe();
  });
 
  $: {
  if (url && divEl) {
  render();
+ if (chainColors) {
+ subscribe = colorByChain(chainColors);
+ }
  }
  }
 </script>

frontend/src/lib/venomeMolstarUtils.ts

@@ -1,6 +1,13 @@
 import { BACKEND_URL } from "./backend";
 import { PDBeMolstarPlugin } from "../../venome-molstar/lib";
 import type { InitParams } from "../../venome-molstar/lib/spec";
+import type { QueryParam } from "../../venome-molstar/lib/helpers";
+import * as d3 from "d3";
+
+export type ResidueColor = { r: number; g: number; b: number };
+export type ChainId = string;
+export type ChainColors = { [chainId: ChainId]: ResidueColor[] };
+export type ChainpLDDT = { [chainId: ChainId]: number[] };
 
 export async function screenshotMolstar(initParams: Partial<InitParams>) {
  const { div, molstar } = await renderHeadless(initParams);
@@ -74,3 +81,70 @@ export const defaultInitParams = (name: string): Partial<InitParams> => ({
  "controlInfo",
  ],
 });
+
+export function colorResidues({
+ colors = [],
+ entity_id = undefined,
+ struct_asym_id = "A",
+}: {
+ colors?: { r: number; g: number; b: number }[];
+ entity_id?: string;
+ struct_asym_id?: string;
+} = {}): QueryParam[] {
+ let selections: QueryParam[] = [];
+ for (let i = 0; i < colors.length; i++) {
+ const color = colors[i];
+ const residueIndex = i + 1;
+ const residueColoring = {
+ entity_id,
+ struct_asym_id,
+ color,
+ start_residue_number: residueIndex,
+ end_residue_number: residueIndex,
+ };
+ selections.push(residueColoring);
+ }
+ return selections;
+}
+
+const alphafoldColorscheme = [
+ "rgb(1,83,214)", // > 90
+ "rgb(100,203,243)", // > 70
+ "rgb(255,219,18)", // > 50
+ "rgb(255,125,69)", // < 50
+];
+export function pLDDTToAlphaFoldResidueColors(pLDDT: number[]): ResidueColor[] {
+ const colors = pLDDT.map((d) => {
+ if (d > 90) {
+ return alphafoldColorscheme[0];
+ } else if (d > 70) {
+ return alphafoldColorscheme[1];
+ } else if (d > 50) {
+ return alphafoldColorscheme[2];
+ } else {
+ return alphafoldColorscheme[3];
+ }
+ });
+ return colors.map((c) => {
+ const rgb = d3.color(c)!.rgb()!;
+ return { r: rgb.r, g: rgb.g, b: rgb.b };
+ });
+}
+
+export function pLDDTToResidueColors(pLDDT: number[]): ResidueColor[] {
+ const colors = pLDDT.map((d) => {
+ if (d > 90) {
+ return alphafoldColorscheme[0];
+ } else if (d > 70) {
+ return alphafoldColorscheme[1];
+ } else if (d > 50) {
+ return alphafoldColorscheme[2];
+ } else {
+ return alphafoldColorscheme[3];
+ }
+ });
+ return colors.map((c) => {
+ const rgb = d3.color(c)!.rgb()!;
+ return { r: rgb.r, g: rgb.g, b: rgb.b };
+ });
+}

frontend/src/routes/Protein.svelte

@@ -11,18 +11,31 @@
  } from "../lib/format";
  import { navigate } from "svelte-routing";
  import References from "../lib/References.svelte";
- import { ChevronDownSolid, PenOutline } from "flowbite-svelte-icons";
+ import {
+ ChevronDownSolid,
+ PaletteOutline,
+ PenOutline,
+ RefreshOutline,
+ UndoOutline,
+ UndoSolid,
+ } from "flowbite-svelte-icons";
  import EntryCard from "../lib/EntryCard.svelte";
  import SimilarProteins from "../lib/SimilarProteins.svelte";
  import DelayedSpinner from "../lib/DelayedSpinner.svelte";
  import { user } from "../lib/stores/user";
  import { AccordionItem, Accordion } from "flowbite-svelte";
+ import {
+ pLDDTToAlphaFoldResidueColors,
+ pLDDTToResidueColors,
+ } from "../lib/venomeMolstarUtils";
+ import type { ChainColors, ChainpLDDT } from "../lib/venomeMolstarUtils";
 
  const fileDownloadDropdown = ["pdb", "fasta"];
 
  export let urlId: string;
  let entry: ProteinEntry | null = null;
  let error = false;
+ let chainColors: ChainColors = {};
 
  // when this component mounts, request protein wikipedia entry from backend
  onMount(async () => {
@@ -123,13 +136,11 @@
 
  <div style="position: sticky; top: 55px; right: 0; z-index:999;">
  <EntryCard title="Provided Information">
- <Molstar
- format="pdb"
- url="http://localhost:8000/protein/pdb/{entry.name}"
- width={400}
- height={350}
- />
- <div id="info-grid" class="grid grid-cols-2 mt-5">
+ <div
+ id="info-grid"
+ class="grid grid-cols-2 mb-2"
+ style="width: 400px;"
+ >
  <b>Organism</b>
  <div>
  {entry.speciesName}
@@ -145,6 +156,44 @@
  >
  </div>
  </div>
+ <Molstar
+ format="pdb"
+ url="http://localhost:8000/protein/pdb/{entry.name}"
+ width={400}
+ height={350}
+ {chainColors}
+ />
+ <div class="mt-2 flex gap-2 items-center">
+ {#if Object.keys(chainColors).length > 0}
+ <Button
+ color="light"
+ size="xs"
+ on:click={() => {
+ chainColors = {};
+ }}><UndoOutline size="xs" /></Button
+ >
+ {/if}
+ <Button
+ color="light"
+ size="xs"
+ on:click={async () => {
+ if (!entry) return;
+ const pLDDTPerChain =
+ await Backend.getPLddtGivenProtein(
+ entry.name
+ );
+ for (const [
+ chainId,
+ pLDDTPerResidue,
+ ] of Object.entries(pLDDTPerChain)) {
+ chainColors[chainId] =
+ pLDDTToResidueColors(pLDDTPerResidue);
+ }
+ }}
+ >
+ Color by pLDDT</Button
+ >
+ </div>
  </EntryCard>
  </div>
  </div>