Preserve custom seq across features

etowahadams · Nov 9, 2024 · 0aef5ec · 0aef5ec
1 parent ce4c326
commit 0aef5ec
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git a/viz/src/SAEContext.tsx b/viz/src/SAEContext.tsx
@@ -59,7 +59,10 @@ export const SAEProvider = ({ children }: { children: React.ReactNode }) => {
         setSelectedFeature: (feature: number | undefined) => {
           setSelectedFeature(feature);
           if (feature !== undefined) {
-            navigate(`/sae-viz/${selectedModel}/${feature}`);
+            const seqMatch = path.match(/\?seq=([^&]+)/);
+            const seq = seqMatch ? seqMatch[1] : "";
+            const seqParam = seq ? `?seq=${seq}` : "";
+            navigate(`/sae-viz/${selectedModel}/${feature}${seqParam}`);
           }
         },
         SAEConfig: SAE_CONFIGS[selectedModel],

diff --git a/viz/src/components/CustomStructureViewer.tsx b/viz/src/components/CustomStructureViewer.tsx
@@ -7,6 +7,7 @@ import { Color } from "molstar/lib/mol-util/color";
 import { redColorMapRGB } from "@/utils.ts";
 import proteinEmoji from "../protein.png";
 import { useIsMobile } from "@/hooks/use-mobile";
+import { ESMFoldCache } from "@/utils";
 
 interface CustomStructureViewerProps {
   viewerId: string;
@@ -27,7 +28,6 @@ const CustomStructureViewer = ({
   const [warning, setWarning] = useState("");
   const isMobile = useIsMobile();
 
-  const ESMFoldCache = useRef<Record<string, string>>({});
   const pluginRef = useRef<PluginContext | null>(null);
 
   const createResidueColorTheme = (activationList: number[], name = "residue-colors") => {
@@ -156,8 +156,8 @@ const CustomStructureViewer = ({
     const renderStructure = async () => {
       setIsLoading(true);
       try {
-        const pdbData = ESMFoldCache.current[seq] || (await foldStructure(seq));
-        ESMFoldCache.current[seq] = pdbData;
+        const pdbData = ESMFoldCache[seq] || (await foldStructure(seq));
+        ESMFoldCache[seq] = pdbData;
         renderViewer(pdbData);
       } catch (error) {
         console.error("Error folding sequence:", error);

diff --git a/viz/src/utils.ts b/viz/src/utils.ts
@@ -67,3 +67,5 @@ export function tokensToSequence(tokens: Array<number>): string {
 export function sequenceToTokens(sequence: string): Array<number> {
   return sequence.split("").map((residue) => residueToToken(residue));
 }
+
+export const ESMFoldCache: Record<string, string> = {};