jbrowse-plugin-msaview 2.2.2 → 2.2.4

package/src/MsaViewPanel/model.ts

@@ -1,22 +1,86 @@
+import { lazy } from 'react'
+
 import { BaseViewModel } from '@jbrowse/core/pluggableElementTypes'
 import { getSession } from '@jbrowse/core/util'
+import { addDisposer, cast, types } from '@jbrowse/mobx-state-tree'
 import { genomeToTranscriptSeqMapping } from 'g2p_mapper'
 import { autorun } from 'mobx'
-import { addDisposer, cast, types } from 'mobx-state-tree'
 import { MSAModelF } from 'react-msaview'
 
 import { doLaunchBlast } from './doLaunchBlast'
 import { genomeToMSA } from './genomeToMSA'
 import { msaCoordToGenomeCoord } from './msaCoordToGenomeCoord'
+import {
+  cleanupOldData,
+  generateDataStoreId,
+  retrieveMsaData,
+  storeMsaData,
+} from './msaDataStore'
+import { buildAlignmentMaps, runPairwiseAlignment } from './pairwiseAlignment'
+import {
+  gappedToUngappedPosition,
+  mapToRecord,
+  ungappedToGappedPosition,
+} from './structureConnection'
+import { getUniprotIdFromAlphaFoldUrl } from './util'
 
+import type { StructureConnection } from './structureConnection'
+import type { MafRegion, MsaViewInitState } from './types'
 import type { Feature } from '@jbrowse/core/util'
+import type { Instance } from '@jbrowse/mobx-state-tree'
 import type { LinearGenomeViewModel } from '@jbrowse/plugin-linear-genome-view'
-import type { Instance } from 'mobx-state-tree'
+
+const ConnectStructureDialog = lazy(
+  () => import('./components/ConnectStructureDialog'),
+)
 
 type LGV = LinearGenomeViewModel
 
 type MaybeLGV = LGV | undefined
 
+/**
+ * Highlights residues in connected protein structures based on current MSA hover position
+ */
+function highlightConnectedStructures(self: JBrowsePluginMsaViewModel) {
+  const { mouseCol, connectedProteinViews } = self
+  if (connectedProteinViews.length === 0) {
+    return
+  }
+
+  for (const conn of connectedProteinViews) {
+    const structure = conn.proteinView?.structures?.[conn.structureIdx]
+    if (!structure) {
+      continue
+    }
+
+    // Clear highlight if mouse left MSA
+    if (mouseCol === undefined) {
+      structure.clearHighlightFromExternal?.()
+      continue
+    }
+
+    const seq = self.getSequenceByRowName(conn.msaRowName)
+    if (!seq) {
+      continue
+    }
+
+    // Convert gapped MSA column to ungapped position
+    const msaUngapped = gappedToUngappedPosition(seq, mouseCol)
+    if (msaUngapped === undefined) {
+      structure.clearHighlightFromExternal?.()
+      continue
+    }
+
+    // Map to structure position and highlight
+    const structurePos = conn.msaToStructure[msaUngapped]
+    if (structurePos === undefined) {
+      structure.clearHighlightFromExternal?.()
+    } else {
+      structure.highlightFromExternal?.(structurePos)
+    }
+  }
+}
+
 export interface IRegion {
   refName: string
   start: number
@@ -28,7 +92,7 @@ export interface BlastParams {
   blastDatabase: string
   msaAlgorithm: string
   blastProgram: string
-  selectedTranscript: Feature
+  selectedTranscript?: Feature
   proteinSequence: string
 }
 
@@ -70,10 +134,48 @@ export default function stateModelFactory() {
          */
         querySeqName: 'QUERY',
 
+        /**
+         * #property
+         * UniProt ID extracted from AlphaFold MSA URL
+         */
+        uniprotId: types.maybe(types.string),
+
         /**
          * #property
          */
         zoomToBaseLevel: false,
+
+        /**
+         * #property
+         * used for loading the MSA view via session snapshots, e.g.
+         * {
+         *   "type": "MsaView",
+         *   "init": {
+         *     "msaUrl": "https://example.com/alignment.fa",
+         *     "treeUrl": "https://example.com/tree.nh",
+         *     "querySeqName": "ENST00000123_hg38"
+         *   }
+         * }
+         */
+        init: types.frozen<MsaViewInitState | undefined>(),
+
+        /**
+         * #property
+         * connections to protein 3D structure views for synchronized highlighting
+         */
+        connectedStructures: types.array(types.frozen<StructureConnection>()),
+
+        /**
+         * #property
+         * Reference ID for MSA data stored in IndexedDB (for large datasets)
+         */
+        dataStoreId: types.maybe(types.string),
+
+        /**
+         * #property
+         * MAF region for coordinate mapping (used when launched from MAF viewer)
+         */
+        mafRegion: types.frozen<MafRegion | undefined>(),
       }),
     )
 
@@ -90,28 +192,31 @@ export default function stateModelFactory() {
        * #volatile
        */
       error: undefined as unknown,
+      /**
+       * #volatile
+       * True when loading MSA data from IndexedDB
+       */
+      loadingStoredData: false,
     }))
 
     .views(self => ({
       /**
        * #method
+       * Get a row by name, returns [name, sequence] or undefined
        */
-      ungappedCoordMap(rowName: string, position: number) {
-        const row = self.rows.find(f => f[0] === rowName)
-        const seq = row?.[1]
-        if (seq && position < seq.length) {
-          let i = 0
-          let j = 0
-          for (; j < position; j++, i++) {
-            while (seq[i] === '-') {
-              i++
-            }
-          }
-          return i
-        }
-        return undefined
+      getRowByName(rowName: string) {
+        return self.rows.find(r => r[0] === rowName)
+      },
+
+      /**
+       * #method
+       * Get the sequence for a row by name
+       */
+      getSequenceByRowName(rowName: string) {
+        return self.rows.find(r => r[0] === rowName)?.[1]
       },
     }))
+
     .views(self => ({
       /**
        * #getter
@@ -121,38 +226,91 @@ export default function stateModelFactory() {
           ? genomeToTranscriptSeqMapping(self.connectedFeature)
           : undefined
       },
-    }))
-    .views(self => ({
+
       /**
        * #getter
        */
-      get mouseCol2(): number | undefined {
-        return genomeToMSA({ model: self as JBrowsePluginMsaViewModel })
+      get processing() {
+        return !!self.progress
       },
+
       /**
        * #getter
        */
-      get clickCol2() {
-        return undefined
+      get connectedView() {
+        const { views } = getSession(self)
+        return views.find(f => f.id === self.connectedViewId) as MaybeLGV
+      },
+
+      /**
+       * #getter
+       * Get connected protein views with their full model references
+       */
+      get connectedProteinViews() {
+        const { views } = getSession(self)
+        return self.connectedStructures
+          .map(conn => {
+            const proteinView = views.find(
+              (v: any) => v.id === conn.proteinViewId,
+            ) as any
+            return proteinView ? { ...conn, proteinView } : undefined
+          })
+          .filter((c): c is StructureConnection & { proteinView: any } => !!c)
       },
     }))
 
     .views(self => ({
       /**
        * #getter
+       * Get the MSA column that corresponds to the currently hovered structure position
+       * Returns the first match from any connected structure
        */
-      get processing() {
-        return !!self.progress
+      get structureHoverCol(): number | undefined {
+        for (const conn of self.connectedProteinViews) {
+          const structure = conn.proteinView?.structures?.[conn.structureIdx]
+          const structurePos = structure?.hoverPosition?.structureSeqPos
+          if (structurePos !== undefined) {
+            const msaUngapped = conn.structureToMsa[structurePos]
+            if (msaUngapped !== undefined) {
+              const seq = self.getSequenceByRowName(conn.msaRowName)
+              if (seq) {
+                // Convert ungapped position to global column, then to visible column
+                const globalCol = ungappedToGappedPosition(seq, msaUngapped)
+                if (globalCol !== undefined) {
+                  return self.globalColToVisibleCol(globalCol)
+                }
+              }
+            }
+          }
+        }
+        return undefined
       },
+    }))
 
+    .views(self => ({
       /**
        * #getter
+       * Returns a secondary highlight column from either:
+       * 1. Structure hover (from connected protein 3D view)
+       * 2. Genome hover (from connected linear genome view)
        */
-      get connectedView() {
-        const { views } = getSession(self)
-        return views.find(f => f.id === self.connectedViewId) as MaybeLGV
+      get mouseCol2(): number | undefined {
+        // Check structure hover first
+        const structureCol = self.structureHoverCol
+        if (structureCol !== undefined) {
+          return structureCol
+        }
+        // Fall back to genome hover
+        return genomeToMSA({ model: self as JBrowsePluginMsaViewModel })
+      },
+      /**
+       * #getter
+       */
+      get clickCol2() {
+        return undefined
       },
     }))
+
     .actions(self => ({
       /**
        * #action
@@ -202,7 +360,158 @@ export default function stateModelFactory() {
       setBlastParams(args?: BlastParams) {
         self.blastParams = args
       },
+      /**
+       * #action
+       */
+      setInit(arg?: MsaViewInitState) {
+        self.init = arg
+      },
+      /**
+       * #action
+       */
+      setQuerySeqName(arg: string) {
+        self.querySeqName = arg
+      },
+      /**
+       * #action
+       */
+      setUniprotId(arg?: string) {
+        self.uniprotId = arg
+      },
+      /**
+       * #action
+       */
+      setDataStoreId(arg?: string) {
+        self.dataStoreId = arg
+      },
+      /**
+       * #action
+       */
+      setMafRegion(arg?: MafRegion) {
+        self.mafRegion = arg
+      },
+      /**
+       * #action
+       */
+      setLoadingStoredData(arg: boolean) {
+        self.loadingStoredData = arg
+      },
+      /**
+       * #action
+       */
+      handleMsaClick(coord: number) {
+        const { connectedView, zoomToBaseLevel } = self
+        const { assemblyManager } = getSession(self)
+        const r2 = msaCoordToGenomeCoord({ model: self, coord })
+
+        if (!r2 || !connectedView) {
+          return
+        }
+
+        if (zoomToBaseLevel) {
+          connectedView.navTo(r2)
+        } else {
+          const r =
+            assemblyManager
+              .get(connectedView.assemblyNames[0]!)
+              ?.getCanonicalRefName(r2.refName) ?? r2.refName
+          connectedView.centerAt(r2.start, r)
+        }
+      },
+
+      /**
+       * #action
+       * Connect to a protein structure for synchronized highlighting
+       */
+      connectToStructure(
+        proteinViewId: string,
+        structureIdx: number,
+        msaRowName?: string,
+      ) {
+        const rowName = msaRowName ?? self.querySeqName
+        const msaSequence = self.getSequenceByRowName(rowName)
+        if (!msaSequence) {
+          throw new Error(`MSA row "${rowName}" not found`)
+        }
+
+        const ungappedMsaSequence = msaSequence.replaceAll('-', '')
+
+        const { views } = getSession(self)
+
+        const proteinView = views.find(
+          (v: any) => v.id === proteinViewId,
+        ) as any
+        if (!proteinView) {
+          throw new Error(`ProteinView "${proteinViewId}" not found`)
+        }
+
+        const structure = proteinView.structures?.[structureIdx]
+        if (!structure) {
+          throw new Error(`Structure at index ${structureIdx} not found`)
+        }
+
+        const structureSequence = structure.structureSequences?.[0]
+        if (!structureSequence) {
+          throw new Error('Structure sequence not available')
+        }
+
+        const alignment = runPairwiseAlignment(
+          ungappedMsaSequence,
+          structureSequence,
+        )
+        const { seq1ToSeq2, seq2ToSeq1 } = buildAlignmentMaps(alignment)
+
+        const connection: StructureConnection = {
+          proteinViewId,
+          structureIdx,
+          msaRowName: rowName,
+          msaToStructure: mapToRecord(seq1ToSeq2),
+          structureToMsa: mapToRecord(seq2ToSeq1),
+        }
+
+        self.connectedStructures.push(connection)
+      },
+
+      /**
+       * #action
+       * Disconnect from a protein structure
+       */
+      disconnectFromStructure(proteinViewId: string, structureIdx: number) {
+        const idx = self.connectedStructures.findIndex(
+          c =>
+            c.proteinViewId === proteinViewId &&
+            c.structureIdx === structureIdx,
+        )
+        if (idx !== -1) {
+          self.connectedStructures.splice(idx, 1)
+        }
+      },
+
+      /**
+       * #action
+       * Disconnect from all protein structures
+       */
+      disconnectAllStructures() {
+        self.connectedStructures.clear()
+      },
     }))
+    .actions(self => {
+      // store reference to the original action from react-msaview
+      const superSetMouseClickPos = self.setMouseClickPos.bind(self)
+
+      return {
+        /**
+         * #action
+         * overrides base setMouseClickPos to trigger navigation
+         */
+        setMouseClickPos(col?: number, row?: number) {
+          superSetMouseClickPos(col, row)
+          if (col !== undefined) {
+            self.handleMsaClick(col)
+          }
+        },
+      }
+    })
 
     .views(self => ({
       /**
@@ -219,26 +528,106 @@ export default function stateModelFactory() {
               self.setZoomToBaseLevel(!self.zoomToBaseLevel)
             },
           },
+          {
+            label: 'Connect to protein structure...',
+            onClick: () => {
+              getSession(self).queueDialog(handleClose => [
+                ConnectStructureDialog,
+                {
+                  model: self,
+                  handleClose,
+                },
+              ])
+            },
+          },
+          ...(self.connectedStructures.length > 0
+            ? [
+                {
+                  label: 'Disconnect from protein structures',
+                  onClick: () => {
+                    self.disconnectAllStructures()
+                  },
+                },
+              ]
+            : []),
         ]
       },
-      /**
-       * #getter
-       */
-      get processing() {
-        return !!self.progress
-      },
-
-      /**
-       * #getter
-       */
-      get connectedView() {
-        const { views } = getSession(self)
-        return views.find(f => f.id === self.connectedViewId) as MaybeLGV
-      },
     }))
 
     .actions(self => ({
       afterCreate() {
+        // Clean up old IndexedDB entries on startup
+        cleanupOldData().catch((e: unknown) => {
+          console.error('Failed to cleanup old MSA data:', e)
+        })
+
+        // Load MSA data from IndexedDB if dataStoreId exists and no data loaded
+        addDisposer(
+          self,
+          autorun(async () => {
+            const { dataStoreId, rows } = self
+            if (dataStoreId && rows.length === 0) {
+              try {
+                self.setLoadingStoredData(true)
+                const storedData = await retrieveMsaData(dataStoreId)
+                if (storedData) {
+                  if (storedData.msa) {
+                    self.setMSA(storedData.msa)
+                  }
+                  if (storedData.tree) {
+                    self.setTree(storedData.tree)
+                  }
+                }
+              } catch (e) {
+                console.error('Failed to load MSA data from IndexedDB:', e)
+              } finally {
+                self.setLoadingStoredData(false)
+              }
+            }
+          }),
+        )
+
+        // Store MSA data to IndexedDB when loaded from inline data (no filehandle)
+        // This ensures data persists across page refreshes even when
+        // react-msaview's postProcessSnapshot would strip it
+        addDisposer(
+          self,
+          autorun(async () => {
+            const { rows, dataStoreId } = self
+            // Only store if we have data and don't already have a dataStoreId
+            if (rows.length > 0 && !dataStoreId) {
+              // Only store if there's no filehandle (filehandles can reload from source)
+              const hasFilehandle = !!(
+                self.msaFilehandle ?? self.treeFilehandle
+              )
+              if (hasFilehandle) {
+                return
+              }
+
+              const msaData = self.data.msa
+              const treeData = self.data.tree
+
+              // Only store if we have actual data
+              if (msaData || treeData) {
+                try {
+                  const newId = generateDataStoreId()
+                  const success = await storeMsaData(newId, {
+                    msa: msaData,
+                    tree: treeData,
+                    treeMetadata: self.data.treeMetadata,
+                  })
+                  // Only set dataStoreId if storage was successful
+                  if (success) {
+                    self.setDataStoreId(newId)
+                  }
+                } catch (e) {
+                  console.error('Failed to store MSA data to IndexedDB:', e)
+                }
+              }
+            }
+          }),
+        )
+
         addDisposer(
           self,
           autorun(async () => {
@@ -261,6 +650,63 @@ export default function stateModelFactory() {
           }),
         )
 
+        // process init parameter for loading MSA from session snapshots
+        addDisposer(
+          self,
+          autorun(async () => {
+            const { init } = self
+            if (init) {
+              try {
+                self.setError(undefined)
+                const { msaData, msaUrl, treeData, treeUrl, querySeqName } =
+                  init
+
+                // Extract uniprotId from AlphaFold MSA URL and set querySeqName
+                if (msaUrl) {
+                  const id = getUniprotIdFromAlphaFoldUrl(msaUrl)
+                  if (id) {
+                    self.setUniprotId(id)
+                    // AlphaFold MSA files use 'query' as the row name
+                    self.setQuerySeqName('query')
+                  }
+                }
+
+                // User-provided querySeqName takes precedence
+                if (querySeqName) {
+                  self.setQuerySeqName(querySeqName)
+                }
+
+                if (msaData) {
+                  self.setMSA(msaData)
+                } else if (msaUrl) {
+                  const response = await fetch(msaUrl)
+                  if (!response.ok) {
+                    throw new Error(`Failed to fetch MSA: ${response.status}`)
+                  }
+                  const data = await response.text()
+                  self.setMSA(data)
+                }
+
+                if (treeData) {
+                  self.setTree(treeData)
+                } else if (treeUrl) {
+                  const response = await fetch(treeUrl)
+                  if (!response.ok) {
+                    throw new Error(`Failed to fetch tree: ${response.status}`)
+                  }
+                  const data = await response.text()
+                  self.setTree(data)
+                }
+
+                self.setInit(undefined)
+              } catch (e) {
+                self.setError(e)
+                console.error(e)
+              }
+            }
+          }),
+        )
+
         // this adds highlights to the genome view when mouse-ing over the MSA
         addDisposer(
           self,
@@ -274,36 +720,147 @@ export default function stateModelFactory() {
               mouseClickCol === undefined
                 ? undefined
                 : msaCoordToGenomeCoord({ model: self, coord: mouseClickCol })
+
             self.setConnectedHighlights([r1, r2].filter(f => !!f))
           }),
         )
 
-        // nav to genome position after click
+        // this highlights residues in connected protein structures when mousing over the MSA
         addDisposer(
           self,
           autorun(() => {
-            const { connectedView, zoomToBaseLevel, mouseClickCol } = self
-            const { assemblyManager } = getSession(self)
-            const r2 =
-              mouseClickCol === undefined
-                ? undefined
-                : msaCoordToGenomeCoord({ model: self, coord: mouseClickCol })
+            highlightConnectedStructures(self)
+          }),
+        )
+
+        // auto-connect to compatible ProteinViews
+        addDisposer(
+          self,
+          autorun(() => {
+            const { views } = getSession(self)
+            const { connectedViewId, uniprotId, rows, connectedStructures } =
+              self
 
-            if (!r2 || !connectedView) {
+            // Need MSA loaded and a uniprotId to auto-connect
+            if (!uniprotId || rows.length === 0) {
               return
             }
 
-            if (zoomToBaseLevel) {
-              connectedView.navTo(r2)
-            } else {
-              const r =
-                assemblyManager
-                  .get(connectedView.assemblyNames[0]!)
-                  ?.getCanonicalRefName(r2.refName) ?? r2.refName
-              connectedView.centerAt(r2.start, r)
+            // Find ProteinViews that share the same connectedViewId
+            for (const view of views) {
+              const v = view as any
+              if (v.type !== 'ProteinView' || !v.structures) {
+                continue
+              }
+
+              for (
+                let structureIdx = 0;
+                structureIdx < v.structures.length;
+                structureIdx++
+              ) {
+                const structure = v.structures[structureIdx]
+
+                // Check if structure shares the same genome view connection
+                if (structure.connectedViewId !== connectedViewId) {
+                  continue
+                }
+
+                // Check if structure has matching uniprotId
+                if (structure.uniprotId !== uniprotId) {
+                  continue
+                }
+
+                // Check if already connected
+                const alreadyConnected = connectedStructures.some(
+                  c =>
+                    c.proteinViewId === v.id && c.structureIdx === structureIdx,
+                )
+                if (alreadyConnected) {
+                  continue
+                }
+
+                // Check if structure sequence is available
+                if (!structure.structureSequences?.[0]) {
+                  continue
+                }
+
+                // Auto-connect
+                try {
+                  self.connectToStructure(v.id, structureIdx)
+                } catch (e) {
+                  console.error('Failed to auto-connect to ProteinView:', e)
+                }
+              }
             }
           }),
         )
+
+        // Observe protein3d genome highlights and update MSA highlighted columns
+        // This enables communication via the linear genome view coordinates
+        addDisposer(
+          self,
+          autorun(() => {
+            const { views } = getSession(self)
+            const { connectedViewId, transcriptToMsaMap, querySeqName } = self
+
+            if (!connectedViewId || !transcriptToMsaMap) {
+              return
+            }
+
+            const columns: number[] = []
+
+            // Find ProteinViews that share the same connected genome view
+            for (const view of views) {
+              const v = view as any
+              if (v.type !== 'ProteinView' || !v.structures) {
+                continue
+              }
+
+              for (const structure of v.structures) {
+                // Check if structure is connected to same genome view
+                if (structure.connectedViewId !== connectedViewId) {
+                  continue
+                }
+
+                // Check if structure has hover genome highlights
+                const highlights = structure.hoverGenomeHighlights
+                if (!highlights || highlights.length === 0) {
+                  continue
+                }
+
+                // Map genome coordinates to MSA columns
+                const { g2p } = transcriptToMsaMap
+                for (const highlight of highlights) {
+                  for (
+                    let coord = highlight.start;
+                    coord < highlight.end;
+                    coord++
+                  ) {
+                    const proteinPos = g2p[coord]
+                    if (proteinPos !== undefined) {
+                      const col = self.seqPosToGlobalCol(
+                        querySeqName,
+                        proteinPos,
+                      )
+                      if (!columns.includes(col)) {
+                        columns.push(col)
+                      }
+                    }
+                  }
+                }
+              }
+            }
+
+            // Convert global column indices to visible column indices
+            const visibleColumns = columns
+              .map(col => self.globalColToVisibleCol(col))
+              .filter((col): col is number => col !== undefined)
+
+            self.setHighlightedColumns(
+              visibleColumns.length > 0 ? visibleColumns : undefined,
+            )
+          }),
+        )
       },
     }))
 }
@@ -312,3 +869,5 @@ export type JBrowsePluginMsaViewStateModel = ReturnType<
   typeof stateModelFactory
 >
 export type JBrowsePluginMsaViewModel = Instance<JBrowsePluginMsaViewStateModel>
+
+export { type MafRegion, type MsaViewInitState } from './types'