Improve search ux (#35)

* Improve search UX * URL state * Notes * Fix
etowahadams · Dec 16, 2024 · 19228b3 · 19228b3
1 parent e78b575
commit 19228b3
Show file tree

Hide file tree

Showing 5 changed files with 221 additions and 184 deletions.
diff --git a/viz/src/components/CustomSeqPlayground.tsx b/viz/src/components/CustomSeqPlayground.tsx
@@ -12,7 +12,7 @@ import {
 import CustomStructureViewer from "./CustomStructureViewer";
 import { getSAEDimActivations, getSteeredSequence } from "@/runpod.ts";
 import SeqInput, { ValidSeqInput } from "./SeqInput";
-import { useUrlState } from "@/hooks/useUrlState";
+import useUrlState from "@/hooks/useUrlState";
 import FullSeqsViewer from "./FullSeqsViewer";
 import PDBStructureViewer from "./PDBStructureViewer";
 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
@@ -50,7 +50,7 @@ const CustomSeqPlayground = ({ feature, saeName }: CustomSeqPlaygroundProps) =>
     initialState.steeredActivations
   );
 
-  const { urlInput, setUrlInput, clearUrlInput } = useUrlState();
+  const [urlState, setUrlState] = useUrlState<{ pdb?: string; seq?: string }>();
 
   const handleSubmit = useCallback(
     async (submittedInput: ValidSeqInput) => {
@@ -62,18 +62,16 @@ const CustomSeqPlayground = ({ feature, saeName }: CustomSeqPlaygroundProps) =>
       setSteeredActivations(initialState.steeredActivations);
 
       if (isPDBID(submittedInput)) {
-        setUrlInput("pdb", submittedInput);
         setInputProteinActivations(
           await constructProteinActivationsDataFromPDBID(submittedInput, feature, saeName)
         );
       } else {
-        setUrlInput("seq", submittedInput);
         setInputProteinActivations(
           await constructProteinActivationsDataFromSequence(submittedInput, feature, saeName)
         );
       }
     },
-    [feature, setUrlInput, saeName]
+    [feature, saeName]
   );
 
   // Reset some states when the user navigates to a new feature
@@ -87,13 +85,14 @@ const CustomSeqPlayground = ({ feature, saeName }: CustomSeqPlaygroundProps) =>
 
   // If an input is set in the URL, submit it
   useEffect(() => {
-    if (urlInput) {
-      if (isProteinSequence(urlInput) || isPDBID(urlInput)) {
-        setCustomSeqInput(urlInput);
-        handleSubmit(urlInput);
-      }
+    if (urlState.pdb && isPDBID(urlState.pdb)) {
+      setCustomSeqInput(urlState.pdb);
+      handleSubmit(urlState.pdb);
+    } else if (urlState.seq && isProteinSequence(urlState.seq)) {
+      setCustomSeqInput(urlState.seq);
+      handleSubmit(urlState.seq);
     }
-  }, [urlInput, setCustomSeqInput, handleSubmit]);
+  }, [urlState.pdb, urlState.seq, setCustomSeqInput, handleSubmit]);
 
   const handleSteer = async () => {
     setViewerState(PlaygroundState.LOADING_STEERED_SEQUENCE);
@@ -121,13 +120,19 @@ const CustomSeqPlayground = ({ feature, saeName }: CustomSeqPlaygroundProps) =>
         <SeqInput
           input={proteinInput}
           setInput={setCustomSeqInput}
-          onSubmit={handleSubmit}
+          onSubmit={(input) => {
+            if (isPDBID(input)) {
+              setUrlState({ pdb: input, seq: undefined });
+            } else {
+              setUrlState({ pdb: undefined, seq: input });
+            }
+          }}
           loading={playgroundState === PlaygroundState.LOADING_SAE_ACTIVATIONS}
           buttonText="Submit"
           onClear={() => {
             setCustomSeqInput("");
             setInputProteinActivations(initialState.inputProteinActivations);
-            clearUrlInput();
+            setUrlState({ pdb: undefined, seq: undefined });
           }}
         />
       </div>
@@ -152,13 +157,14 @@ const CustomSeqPlayground = ({ feature, saeName }: CustomSeqPlaygroundProps) =>
             </p>
           )}
 
-          {isPDBID(urlInput) ? (
+          {urlState.pdb && (
             <PDBStructureViewer
               viewerId="custom-viewer"
               proteinActivationsData={inputProteinActivations}
               onLoad={onStructureLoad}
             />
-          ) : (
+          )}
+          {urlState.seq && (
             <CustomStructureViewer
               viewerId="custom-viewer"
               proteinActivationsData={inputProteinActivations}