initial commit

Dockerfile

@@ -0,0 +1,28 @@
+FROM python:3.11
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+
+# Switch to the "user" user
+USER user
+
+# Set home to the user's home directory
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Try and run pip command after setting the user with `USER user` to avoid permission issues with Python
+RUN pip install --no-cache-dir --upgrade pip
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+COPY --chown=user requirements.txt .
+
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user app.py .
+
+ENTRYPOINT ["solara", "run", "app.py", "--host=0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,281 @@
+import json
+from dataclasses import dataclass
+from io import StringIO
+from pathlib import Path
+from typing import Literal, Optional, TypedDict, cast
+from zipfile import ZipFile
+
+import polars as pl
+import solara
+from Bio.PDB import MMCIFParser, Structure
+from ipymolstar import PDBeMolstar
+from ipymolstar.widget import QueryParam
+from polarify import polarify
+from solara.components.file_drop import FileInfo
+
+
+class ColorData(TypedDict):
+    data: list[QueryParam]
+    NonSelectedColor: None
+
+
+class TooltipData(TypedDict):
+    data: list[QueryParam]
+
+
+class CustomData(TypedDict):
+    data: str
+    format: Literal["cif"]
+    binary: Literal[False]
+
+
+@dataclass
+class AlphaFoldData:
+    name: str
+    structure: Structure
+    atom_data: pl.DataFrame
+    residue_data: pl.DataFrame
+
+    custom_data: CustomData
+    color_data: ColorData
+    tooltip_data: TooltipData
+
+    def write_atoms(self):
+        return self.atom_data.write_csv()
+
+
+COLOR_LUT = {
+    "very-high": {"r": 16, "g": 109, "b": 255},
+    "confident": {"r": 16, "g": 207, "b": 241},
+    "low": {"r": 246, "g": 237, "b": 18},
+    "very-low": {"r": 239, "g": 130, "b": 30},
+}
+NO_COLOR_DATA = {"data": [], "nonSelectedColor": None}
+NO_TOOLTIP_DATA = {"data": []}
+
+
+PARSER = MMCIFParser()
+
+result_index = solara.reactive(0)
+file_info = solara.reactive(cast(Optional[FileInfo], None))
+
+
+@polarify
+def assign_confidence(x: pl.Expr) -> pl.Expr:
+    s = pl.lit("very-high")
+    if x < 50:
+        s = pl.lit("very-low")
+    elif x < 70:
+        s = pl.lit("low")
+    elif x < 90:
+        s = pl.lit("confident")
+
+    return s
+
+
+@solara.lab.task
+def load_result() -> Optional[AlphaFoldData]:
+    if file_info.value is None:
+        print("none done")
+        return "None"
+
+    f_idx = result_index.value
+    with ZipFile(file_info.value["file_obj"]) as zf:
+        files = zf.namelist()
+        structure_file = [f for f in files if f.endswith(".cif")][f_idx]
+        json_data_file = [f for f in files if "full_data" in f][f_idx]
+
+        with zf.open(json_data_file) as json_f:
+            json_load = json.load(json_f)
+
+        cif_str = zf.read(structure_file).decode("utf-8")
+
+    alphafold_name = structure_file.rstrip(f"_model_{f_idx}.cif")
+
+    sio = StringIO(cif_str)
+    sio.seek(0)
+    structure = PARSER.get_structure(structure_file.removesuffix(".cif"), sio)
+
+    names = pl.Series(
+        (atom.get_parent().resname for atom in structure.get_atoms()),
+        dtype=pl.Categorical,
+    )
+    resn = pl.Series(atom.get_parent().id[1] for atom in structure.get_atoms())
+    chain = pl.Series(json_load["atom_chain_ids"], dtype=pl.Categorical)
+
+    atoms_df = pl.DataFrame(
+        {
+            "name": names,
+            "resn": resn,
+            "chain": chain,
+            "plddt": json_load["atom_plddts"],
+        }
+    )
+
+    residue_df = (
+        atoms_df.group_by(["chain", "resn", "name"])
+        .agg(pl.col("plddt").mean().alias("mean_plddt"))
+        .sort(["chain", "resn"])
+        .with_columns(
+            assign_confidence(pl.col("mean_plddt"))
+            .alias("confidence")
+            .cast(pl.Categorical)
+        )
+    )
+
+    custom_data = {
+        "data": cif_str,
+        "format": "cif",
+        "binary": False,
+    }
+
+    color_query = []
+    tooltip_query = []
+    for chain, resn, alphafold_name, mean_plddt, confidence in residue_df.iter_rows():
+        res_color = {
+            "struct_asym_id": chain,
+            "residue_number": resn,
+            "color": COLOR_LUT[confidence],
+        }
+        res_tt = {
+            "struct_asym_id": chain,
+            "residue_number": resn,
+            "tooltip": f"Confidence: {confidence}; plddt: {mean_plddt:.2f}",
+        }
+        color_query.append(res_color)
+        tooltip_query.append(res_tt)
+
+    plddt_color_data = {"data": color_query, "nonSelectedColor": None}
+    plddt_tooltip_data = {"data": tooltip_query}
+
+    data = AlphaFoldData(
+        name=alphafold_name,
+        structure=structure,
+        atom_data=atoms_df,
+        residue_data=residue_df,
+        custom_data=custom_data,
+        color_data=plddt_color_data,
+        tooltip_data=plddt_tooltip_data,
+    )
+    return data
+
+
+assets_dir = Path(__file__).parent.parent / "assets"
+cif_file = assets_dir / "seca_secb" / "fold_seca_secb_model_0.cif"
+json_file = assets_dir / "seca_secb" / "fold_seca_secb_full_data_0.json"
+
+custom_data_initial = {
+    "data": cif_file.read_bytes(),
+    "format": "cif",
+    "binary": False,
+}
+
+
+@solara.component
+def Page():
+    color_data = solara.use_reactive(NO_COLOR_DATA)
+    tooltip_data = solara.use_reactive(NO_TOOLTIP_DATA)
+
+    color_mode = solara.use_reactive("chain")
+    dark_effective = solara.lab.use_dark_effective()
+
+    def on_color_mode(value: str):
+        color_mode.set(value)
+        if value == "chain":
+            color_data.set(NO_COLOR_DATA)
+            tooltip_data.set(NO_TOOLTIP_DATA)
+        else:
+            color_data.set(load_result.value.color_data)
+            tooltip_data.set(load_result.value.tooltip_data)
+
+    def set_result_index(value: int):
+        result_index.set(value)
+        load_result()
+
+    solara.Title("Solarafold result viewer")
+    with solara.AppBar():
+        solara.lab.ThemeToggle()
+
+    with solara.Sidebar():
+        solara.FileDrop(label="Upload zip file", on_file=file_info.set, lazy=True)
+        solara.Button(
+            label="Load result",
+            on_click=load_result,
+            block=True,
+            disabled=file_info.value is None,
+        )
+
+        if not load_result.not_called:
+            disabled = load_result.pending
+            solara.Select(
+                label="Result index",
+                value=result_index.value,
+                on_value=set_result_index,
+                values=list(range(5)),
+                disabled=disabled,
+            )
+            solara.Select(
+                label="Color mode",
+                values=["chain", "plddt"],
+                value=color_mode.value,
+                on_value=on_color_mode,
+                disabled=disabled,
+            )
+
+            def write_atoms():
+                return load_result.value.atom_data.write_csv()
+
+            solara.FileDownload(
+                write_atoms,
+                filename="NA" if disabled else f"{load_result.value.name}_atoms.csv",
+                children=[
+                    solara.Button(
+                        "Download atom plddt",
+                        block=True,
+                        disabled=disabled,
+                    )
+                ],
+            )
+
+            solara.Div(style={"height": "20px"})
+
+            def write_residues():
+                return load_result.value.residue_data.write_csv()
+
+            solara.FileDownload(
+                write_residues,
+                filename="NA" if disabled else f"{load_result.value.name}_residues.csv",
+                children=[
+                    solara.Button(
+                        "Download residue plddt",
+                        block=True,
+                        disabled=disabled,
+                    )
+                ],
+            )
+
+    if load_result.not_called:
+        solara.Text("Drop and drag an alphafold3 result .zip file to get started")
+    elif load_result.pending:
+        solara.ProgressLinear(load_result.pending)
+    elif load_result.finished:
+        fold_data: AlphaFoldData = load_result.value
+
+        with solara.Card():
+            theme = "dark" if dark_effective else "light"
+            PDBeMolstar.element(
+                height="calc(100vh - 150px)",
+                custom_data=fold_data.custom_data,
+                color_data=color_data.value,
+                tooltips=tooltip_data.value,
+                show_water=False,
+                theme=theme,
+            ).key(f"pdbemolstar-{dark_effective}")
+
+
+@solara.component
+def Layout(children):
+    dark_effective = solara.lab.use_dark_effective()
+    return solara.AppLayout(
+        children=children, toolbar_dark=dark_effective, color=None
+    )  # if dark_effective else "primary")

requirements.txt

@@ -0,0 +1,5 @@
+polars
+solara
+biopython
+ipymolstar
+polarify