Fix RETURN_SELECTION_JS

output.py

@@ -352,14 +352,12 @@ STARTUP_JS = """
     window.selected_elements = {}
     
     function handleMessage(event) {
-        if ("atom" in event.data.data) {
-            if ("add" in event.data.data) {
-                console.log("New message: ", event.data)
-                let atom = event.data.data["atom"];
-                let add = event.data.data["add"];
-                console.log("add: ", add, " atom: ", atom);
-                window.selected_elements[atom] = add;
-            }
+        if (event.data.name == "atom_selection") {   
+            console.log("New message: ", event.data)
+            let atom = event.data.data["atom"];
+            let add = event.data.data["add"];
+            console.log("add: ", add, " atom: ", atom);
+            window.selected_elements[atom] = add;
         }
     }
     
@@ -376,14 +374,16 @@ RETURN_SELECTION_JS = """
       if (add) {
         console.log("Adding atom ", atom);
         selected.push(String(atom));
-        window.parent.postMessage({
-          name: "atom_selection",
-          data: {"atom": parseInt(atom), "add": false} 
-        }, "*");
+        //window.parent.postMessage({
+        //  name: "atom_selection",
+        //  data: {"atom": parseInt(atom), "add": false} 
+        //}, "*");
       }
     }
     console.log("Finished parsing");
     console.log(window.selected_elements);
+    window.selected_elements = {}
+    console.log(window.selected_elements);
     return [input_file, input_protein_file, n_steps, n_atoms, n_samples, selected.join(",")];
 }
 """