Add different alignment modes to sequences display (#24)

etowahadams · Nov 10, 2024 · 835c01d · 835c01d
1 parent 38c89fa
commit 835c01d
Show file tree

Hide file tree

Showing 10 changed files with 409 additions and 132 deletions.
diff --git a/viz/index.html b/viz/index.html
@@ -27,6 +27,7 @@
     gtag("js", new Date());
     gtag("config", "G-5S9BXEL0L4");
   </script>
+  <script src="https://cdn.jsdelivr.net/npm/biomsa/dist/biomsa.js"></script>
   <body>
     <div id="root"></div>
     <script type="module" src="/src/main.tsx"></script>

diff --git a/viz/package.json b/viz/package.json
@@ -17,6 +17,7 @@
     "@radix-ui/react-slider": "^1.2.1",
     "@radix-ui/react-slot": "^1.1.0",
     "@radix-ui/react-toggle": "^1.1.0",
+    "@radix-ui/react-toggle-group": "^1.1.0",
     "@radix-ui/react-tooltip": "^1.1.3",
     "class-variance-authority": "^0.7.0",
     "clsx": "^2.1.1",

diff --git a/viz/pnpm-lock.yaml b/viz/pnpm-lock.yaml
diff --git a/viz/src/SAEVisualizerPage.tsx b/viz/src/SAEVisualizerPage.tsx
@@ -1,25 +1,41 @@
 import "./App.css";
 import { useEffect, useState, useContext } from "react";
 import MolstarMulti from "./components/MolstarMulti";
-import SeqViewer, { SingleSeq } from "./components/SeqViewer";
 import CustomSeqPlayground from "./components/CustomSeqPlayground";
 import { Navigate } from "react-router-dom";
 
 import { SAEContext } from "./SAEContext";
 import { NUM_SEQS_TO_DISPLAY } from "./config";
 import { CONTRIBUTORS } from "./SAEConfigs";
+import SeqsViewer, { SeqWithSAEActs } from "./components/SeqsViewer";
+import { tokensToSequence } from "./utils";
 
 const SAEVisualizerPage: React.FC = () => {
   const { selectedFeature, selectedModel, SAEConfig } = useContext(SAEContext);
   const dimToCuratedMap = new Map(SAEConfig?.curated?.map((i) => [i.dim, i]) || []);
 
-  const [featureData, setFeatureData] = useState<SingleSeq[]>([]);
+  const [featureData, setFeatureData] = useState<SeqWithSAEActs[]>([]);
   useEffect(() => {
     const fileURL = `${SAEConfig.baseUrl}${selectedFeature}.json`;
     fetch(fileURL)
       .then((response) => response.json())
       .then((data) => {
-        setFeatureData(data.slice(0, NUM_SEQS_TO_DISPLAY));
+        // NOTE(liam): important data transformation
+        setFeatureData(
+          data
+            .slice(0, NUM_SEQS_TO_DISPLAY)
+            .map(
+              (seq: {
+                tokens_acts_list: number[];
+                tokens_list: number[];
+                alphafold_id: string;
+              }) => ({
+                sae_acts: seq.tokens_acts_list,
+                sequence: tokensToSequence(seq.tokens_list),
+                alphafold_id: seq.alphafold_id,
+              })
+            )
+        );
       });
   }, [SAEConfig, selectedFeature]);
 
@@ -52,16 +68,9 @@ const SAEVisualizerPage: React.FC = () => {
     <>
       <main className="text-left max-w-full overflow-x-auto">
         <h1 className="text-3xl font-semibold md:mt-0 mt-16">Feature {selectedFeature}</h1>
-        {dimToCuratedMap.has(selectedFeature) && <p className="mt-3">{desc}</p>}
+        {dimToCuratedMap.has(selectedFeature) && <div className="mt-3">{desc}</div>}
         {SAEConfig?.supportsCustomSequence && <CustomSeqPlayground feature={selectedFeature} />}
-        <h2 className="text-2xl font-semibold mt-8">Top activating sequences</h2>
-        <div className="p-4 mt-5 border-2 border-gray-200 border-dashed rounded-lg">
-          <div className="overflow-x-auto">
-            {featureData.map((seq) => (
-              <SeqViewer seq={seq} key={`seq-${seq.alphafold_id}`} />
-            ))}
-          </div>
-        </div>
+        <SeqsViewer seqs={featureData} />
         <MolstarMulti proteins={featureData} />
       </main>
     </>

diff --git a/viz/src/components/CustomSeqPlayground.tsx b/viz/src/components/CustomSeqPlayground.tsx
@@ -1,8 +1,7 @@
 import { useState, useEffect, useRef, useCallback } from "react";
-import SeqViewer from "./SeqViewer";
 import { Button } from "@/components/ui/button";
 import { Slider } from "@/components/ui/slider";
-import { isPDBID, getPDBSequence, sequenceToTokens } from "@/utils.ts";
+import { isPDBID, getPDBSequence } from "@/utils.ts";
 import CustomStructureViewer from "./CustomStructureViewer";
 import { getSAEDimActivations, getSteeredSequence } from "@/runpod.ts";
 import SeqInput from "./SeqInput";
@@ -225,12 +224,7 @@ const CustomSeqPlayground = ({ feature }: CustomSeqPlaygroundProps) => {
               {steeredActivations.length > 0 && (
                 <>
                   <div className="overflow-x-auto my-4">
-                    <SeqViewer
-                      seq={{
-                        tokens_acts_list: steeredActivations,
-                        tokens_list: sequenceToTokens(steeredSeq),
-                      }}
-                    />
+                    <FullSeqViewer sequence={steeredSeq} activations={steeredActivations} />
                   </div>
                   <CustomStructureViewer
                     viewerId="steered-viewer"

diff --git a/viz/src/components/MolstarMulti.tsx b/viz/src/components/MolstarMulti.tsx
@@ -7,14 +7,10 @@ import { Color } from "molstar/lib/mol-util/color";
 import proteinEmoji from "../protein.png";
 import { redColorMapRGB } from "@/utils.ts";
 import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
-
-interface ProteinData {
-  alphafold_id: string;
-  tokens_acts_list: Array<number>;
-}
+import { SeqWithSAEActs } from "./SeqsViewer";
 
 interface MolstarViewerProps {
-  proteins: ProteinData[];
+  proteins: SeqWithSAEActs[];
 }
 
 const PROTEIN_CANVAS_SIZE = 400;
@@ -74,15 +70,15 @@ const MolstarMulti: React.FC<MolstarViewerProps> = ({ proteins }) => {
 
   const loadStructure = async (
     plugin: PluginContext,
-    protein: ProteinData,
+    protein: SeqWithSAEActs,
     index: number,
     isInteractive: boolean = false
   ) => {
     try {
       const fileName = `https://alphafold.ebi.ac.uk/files/AF-${protein.alphafold_id}-F1-model_v4.cif`;
 
       const themeName = Math.random().toString(36).substring(7);
-      const ResidueColorTheme = createResidueColorTheme(protein.tokens_acts_list, themeName);
+      const ResidueColorTheme = createResidueColorTheme(protein.sae_acts, themeName);
       plugin.representation.structure.themes.colorThemeRegistry.add(
         ResidueColorTheme.colorThemeProvider!
       );

diff --git a/viz/src/components/SeqInput.tsx b/viz/src/components/SeqInput.tsx
@@ -34,7 +34,6 @@ export default function SeqInput({
         await getPDBSequence(input);
         return true;
       } catch (e) {
-        console.log("here", e);
         if (e instanceof Error) {
           setError(e.message);
         } else {

diff --git a/viz/src/components/SeqViewer.tsx b/viz/src/components/SeqViewer.tsx