Update app

app.py

@@ -26,6 +26,7 @@ from rdkit.Chem import Draw
 from rdkit.Chem import AllChem
 from rdkit.Chem import DataStructs
 
+st.set_page_config(layout="wide")
 
 basepath = os.path.dirname(__file__)
 datapath = os.path.join(basepath, "data")
@@ -459,17 +460,24 @@ def images_from_molecule(top_n):
 
         st.write("Retrieved images from the Cell Painting database")
         with st.container():
-            columns = st.columns(len(top_probs))
+            n_columns = 5
+            columns = st.columns(n_columns)
+
             for i, col in enumerate(columns):
-                id = top_ids[i]
-                id = f"{id}.npz"
-                image = images_dict[id]
+                l = list(range(len(top_ids)))
 
-                ## TODO: generalize and functionalize
-                im = reshape_image(image)
+                col_ids = [m for m in l if m % 5 == i]
 
-                index = i+1
-                col.image(im, caption=index)
+                for n in col_ids:
+                    id = top_ids[n]
+                    id = f"{id}.npz"
+                    image = images_dict[id]
+
+                    ## TODO: generalize and functionalize
+                    im = reshape_image(image)
+
+                    index = n+1
+                    col.image(im, caption=index)
 
 
 page_names_to_funcs = {