multimolecule/regulatory-variant-effect
0
1# MultiMolecule2# Copyright (C) 2024-Present MultiMolecule3 4# This file is part of MultiMolecule.5 6# MultiMolecule is free software: you can redistribute it and/or modify7# it under the terms of the GNU Affero General Public License as published by8# the Free Software Foundation, either version 3 of the License, or9# any later version.10 11# MultiMolecule is distributed in the hope that it will be useful,12# but WITHOUT ANY WARRANTY; without even the implied warranty of13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the14# GNU Affero General Public License for more details.15 16# You should have received a copy of the GNU Affero General Public License17# along with this program. If not, see <http://www.gnu.org/licenses/>.18 19# For additional terms and clarifications, please refer to our License FAQ at:20# <https://multimolecule.danling.org/about/license-faq>.21 22from __future__ import annotations23 24import csv25import json26import re27import tempfile28import time29from functools import lru_cache30from typing import Any, Mapping31from urllib.parse import parse_qs, urlparse32 33import gradio as gr34import matplotlib35import numpy as np36import torch37from transformers import pipeline38 39matplotlib.use("Agg")40 41import matplotlib.pyplot as plt # noqa: E40242import multimolecule # noqa: E402, F401 - registers MultiMolecule models and pipelines with Transformers43 44DEFAULT_REFERENCE_SEQUENCE = "ACGT" * 25045DEFAULT_ALTERNATIVE_SEQUENCE = "ACGT" * 125 + "TCGA" + "ACGT" * 12446DEFAULT_MODEL_LABEL = "DeepSEA"47 48MODEL_OPTIONS = {49 "A2Z Chromatin": "multimolecule/a2zchromatin",50 "Basset": "multimolecule/basset",51 "DeepMEL": "multimolecule/deepmel",52 "DeepSEA": "multimolecule/deepsea",53 "DeepSTARR": "multimolecule/deepstarr",54 "Malinois": "multimolecule/malinois",55 "MPRA-DragoNN": "multimolecule/mpradragonn",56 "scBasset": "multimolecule/scbasset",57 "Xpresso": "multimolecule/xpresso",58}59MODEL_LABELS = {model_id: label for label, model_id in MODEL_OPTIONS.items()}60 61TABLE_HEADERS = ["position", "nucleotide", "channel", "delta_score", "reference_score", "alternative_score"]62DNA_ALPHABET = set("ACGTN")63FLOAT_PATTERN = re.compile(r"[-+]?(?:(?:\d*\.\d+)|(?:\d+\.?))(?:[eE][-+]?\d+)?")64 65 66def _device() -> int:67 return 0 if torch.cuda.is_available() else -168 69 70@lru_cache(maxsize=2)71def load_predictor(model_id: str):72 return pipeline("regulatory-variant-effect", model=model_id, device=_device())73 74 75def clean_sequence(sequence: str, label: str) -> str:76 sequence = "".join(str(sequence or "").split()).upper().replace("U", "T")77 if not sequence:78 raise gr.Error(f"{label} sequence is empty.")79 invalid = sorted(set(sequence) - DNA_ALPHABET)80 if invalid:81 invalid_text = ", ".join(invalid)82 raise gr.Error(f"{label} sequence contains unsupported symbols: {invalid_text}. Use A, C, G, T, or N.")83 return sequence84 85 86def parse_features(features_text: str) -> Any | None:87 text = str(features_text or "").strip()88 if not text:89 return None90 91 try:92 parsed = json.loads(text)93 except json.JSONDecodeError:94 values = FLOAT_PATTERN.findall(text)95 if not values:96 raise gr.Error("Features must be JSON or comma/space-separated numbers.")97 return [float(value) for value in values]98 99 if isinstance(parsed, Mapping):100 for key in ("features", "values", "reference_features", "alternative_features"):101 if key in parsed:102 return parsed[key]103 if all(isinstance(value, int | float) for value in parsed.values()):104 return list(parsed.values())105 raise gr.Error("Feature JSON objects must contain a features/values list or only numeric values.")106 if isinstance(parsed, str):107 return parse_features(parsed)108 return parsed109 110 111def feature_summary(features: Any | None) -> dict[str, Any]:112 if features is None:113 return {"provided": False}114 try:115 array = np.asarray(features, dtype=float)116 except (TypeError, ValueError):117 return {"provided": True, "shape": None}118 return {"provided": True, "shape": list(array.shape)}119 120 121def unpack_prediction_result(result: Any) -> dict[str, Any]:122 if isinstance(result, list):123 if len(result) != 1:124 raise gr.Error(f"Expected one prediction result, got {len(result)}.")125 result = result[0]126 if not isinstance(result, dict):127 raise gr.Error(f"Expected a prediction dictionary, got {type(result).__name__}.")128 return result129 130 131def build_delta_rows(result: Mapping[str, Any]) -> list[dict[str, Any]]:132 if "delta_score" in result:133 return [134 {135 "position": "",136 "nucleotide": "",137 "channel": "score",138 "delta_score": result.get("delta_score"),139 "reference_score": result.get("reference_score", ""),140 "alternative_score": result.get("alternative_score", ""),141 }142 ]143 144 delta_scores = result.get("delta_scores")145 if isinstance(delta_scores, Mapping):146 reference_scores = result.get("reference_scores") if isinstance(result.get("reference_scores"), Mapping) else {}147 alternative_scores = (148 result.get("alternative_scores") if isinstance(result.get("alternative_scores"), Mapping) else {}149 )150 return [151 {152 "position": "",153 "nucleotide": "",154 "channel": str(channel),155 "delta_score": value,156 "reference_score": reference_scores.get(channel, ""),157 "alternative_score": alternative_scores.get(channel, ""),158 }159 for channel, value in delta_scores.items()160 ]161 162 if isinstance(delta_scores, list):163 return build_axis_delta_rows(result, delta_scores)164 165 raise gr.Error("The selected model did not return delta scores.")166 167 168def build_axis_delta_rows(result: Mapping[str, Any], delta_scores: list[Any]) -> list[dict[str, Any]]:169 channels = [str(channel) for channel in result.get("channels", [])]170 reference_scores = _index_axis_rows(result.get("reference_scores"))171 alternative_scores = _index_axis_rows(result.get("alternative_scores"))172 output_rows: list[dict[str, Any]] = []173 174 for row_index, row in enumerate(delta_scores):175 if not isinstance(row, Mapping):176 continue177 position = row.get("position", row.get("bin", row_index))178 channel_names = channels or [179 str(key) for key in row if key not in {"position", "bin", "nucleotide"} and _is_number(row[key])180 ]181 ref_row = reference_scores.get(position, {})182 alt_row = alternative_scores.get(position, {})183 for channel in channel_names:184 if channel not in row:185 continue186 output_rows.append(187 {188 "position": position,189 "nucleotide": row.get("nucleotide", ""),190 "channel": channel,191 "delta_score": row[channel],192 "reference_score": ref_row.get(channel, ""),193 "alternative_score": alt_row.get(channel, ""),194 }195 )196 return output_rows197 198 199def _index_axis_rows(rows: Any) -> dict[Any, Mapping[str, Any]]:200 if not isinstance(rows, list):201 return {}202 indexed = {}203 for row_index, row in enumerate(rows):204 if isinstance(row, Mapping):205 indexed[row.get("position", row.get("bin", row_index))] = row206 return indexed207 208 209def _is_number(value: Any) -> bool:210 return isinstance(value, int | float | np.number)211 212 213def table_values(rows: list[Mapping[str, Any]]) -> list[list[Any]]:214 return [[row.get(header, "") for header in TABLE_HEADERS] for row in rows]215 216 217def plot_delta_rows(rows: list[Mapping[str, Any]], max_bars: int = 24):218 numeric_rows = [row for row in rows if _is_number(row.get("delta_score"))]219 fig, ax = plt.subplots(figsize=(7.0, 2.4))220 if not numeric_rows:221 ax.text(0.5, 0.5, "No numeric delta scores", ha="center", va="center", transform=ax.transAxes)222 ax.set_axis_off()223 fig.tight_layout()224 return fig225 226 top_rows = sorted(numeric_rows, key=lambda row: abs(float(row["delta_score"])), reverse=True)[:max_bars]227 labels = [_row_label(row) for row in top_rows]228 values = [float(row["delta_score"]) for row in top_rows]229 colors = ["#1b9e77" if value >= 0 else "#d95f02" for value in values]230 231 height = min(7.0, max(2.4, 0.28 * len(top_rows) + 1.2))232 fig.set_size_inches(7.0, height, forward=True)233 ax.barh(range(len(top_rows)), values, color=colors)234 ax.axvline(0, color="#333333", linewidth=0.8)235 ax.set_yticks(range(len(top_rows)), labels)236 ax.invert_yaxis()237 ax.set_xlabel("Alternative - reference")238 ax.set_title("Largest absolute delta scores")239 ax.tick_params(axis="y", labelsize=8)240 fig.tight_layout()241 return fig242 243 244def _row_label(row: Mapping[str, Any]) -> str:245 channel = str(row.get("channel", "score"))246 position = row.get("position")247 if position not in ("", None):248 nucleotide = row.get("nucleotide")249 suffix = f" {nucleotide}" if nucleotide not in ("", None) else ""250 return f"{position}{suffix} {channel}"251 return channel252 253 254def write_result_files(255 model_id: str,256 result: Mapping[str, Any],257 rows: list[Mapping[str, Any]],258 metadata: Mapping[str, Any],259) -> tuple[str, str]:260 csv_file = tempfile.NamedTemporaryFile("w", suffix=".csv", newline="", delete=False)261 writer = csv.DictWriter(csv_file, fieldnames=TABLE_HEADERS)262 writer.writeheader()263 writer.writerows({header: row.get(header, "") for header in TABLE_HEADERS} for row in rows)264 csv_file.close()265 266 json_file = tempfile.NamedTemporaryFile("w", suffix=".json", delete=False)267 json.dump(268 {269 "metadata": dict(metadata),270 "model": model_id,271 "result": result,272 "delta_table": [{header: row.get(header, "") for header in TABLE_HEADERS} for row in rows],273 },274 json_file,275 indent=2,276 default=_json_default,277 )278 json_file.close()279 return csv_file.name, json_file.name280 281 282def _json_default(value: Any):283 if isinstance(value, np.generic):284 return value.item()285 if isinstance(value, np.ndarray):286 return value.tolist()287 raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")288 289 290def predict(291 model_label: str,292 reference_sequence: str,293 alternative_sequence: str,294 reference_features_text: str,295 alternative_features_text: str,296):297 model_id = MODEL_OPTIONS[model_label]298 reference_sequence = clean_sequence(reference_sequence, "Reference")299 alternative_sequence = clean_sequence(alternative_sequence, "Alternative")300 if len(reference_sequence) != len(alternative_sequence):301 raise gr.Error(302 f"Reference and alternative sequences must have the same length. "303 f"Got {len(reference_sequence)} and {len(alternative_sequence)}."304 )305 306 reference_features = parse_features(reference_features_text)307 alternative_features = parse_features(alternative_features_text)308 started = time.perf_counter()309 310 predictor = load_predictor(model_id)311 try:312 result = predictor(313 reference_sequence,314 alternative=alternative_sequence,315 features=reference_features,316 alternative_features=alternative_features,317 )318 except Exception as error:319 raise gr.Error(f"Prediction failed for {model_id}: {error}") from error320 321 result = unpack_prediction_result(result)322 rows = build_delta_rows(result)323 if not rows:324 raise gr.Error("The selected model returned no tabular delta scores.")325 326 metadata = {327 "task": "regulatory-variant-effect",328 "model": model_id,329 "device": "cuda" if torch.cuda.is_available() else "cpu",330 "reference_length": len(reference_sequence),331 "alternative_length": len(alternative_sequence),332 "reference_features": feature_summary(reference_features),333 "alternative_features": feature_summary(alternative_features),334 "alternative_features_inherit_reference": alternative_features is None and reference_features is not None,335 "score_definition": "alternative_minus_reference",336 "num_delta_rows": len(rows),337 "has_reference_scores": any(row.get("reference_score") not in ("", None) for row in rows),338 "has_alternative_scores": any(row.get("alternative_score") not in ("", None) for row in rows),339 "elapsed_seconds": round(time.perf_counter() - started, 3),340 }341 csv_path, json_path = write_result_files(model_id, result, rows, metadata)342 343 return (344 table_values(rows),345 metadata,346 plot_delta_rows(rows),347 csv_path,348 json_path,349 )350 351 352def initial_model(request: gr.Request):353 if request is None:354 return DEFAULT_MODEL_LABEL355 356 query_params = getattr(request, "query_params", None)357 model_id = None358 if query_params is not None:359 model_id = query_params.get("model")360 if not model_id and getattr(request, "url", None):361 parsed = parse_qs(urlparse(str(request.url)).query)362 model_values = parsed.get("model")363 model_id = model_values[0] if model_values else None364 365 return MODEL_LABELS.get(model_id, DEFAULT_MODEL_LABEL)366 367 368with gr.Blocks(title="Regulatory Variant Effect") as demo:369 gr.Markdown(370 "# Regulatory Variant Effect\n"371 "Score matched reference and alternative DNA windows with MultiMolecule regulatory variant-effect models."372 )373 374 model = gr.Dropdown(375 choices=list(MODEL_OPTIONS.keys()),376 value=DEFAULT_MODEL_LABEL,377 label="Checkpoint",378 )379 380 with gr.Row():381 reference_sequence = gr.Textbox(label="Reference DNA sequence", value=DEFAULT_REFERENCE_SEQUENCE, lines=5)382 alternative_sequence = gr.Textbox(label="Alternative DNA sequence", value=DEFAULT_ALTERNATIVE_SEQUENCE, lines=5)383 384 with gr.Accordion("Optional numeric features", open=False), gr.Row():385 reference_features = gr.Textbox(386 label="Reference features JSON/text",387 placeholder='[0.1, 0.2, 0.3] or {"features": [0.1, 0.2, 0.3]}',388 lines=3,389 )390 alternative_features = gr.Textbox(391 label="Alternative features JSON/text",392 placeholder="Leave blank to reuse reference features when provided.",393 lines=3,394 )395 396 run = gr.Button("Run prediction", variant="primary")397 398 delta_table = gr.Dataframe(headers=TABLE_HEADERS, label="Delta scores", interactive=False, wrap=True)399 with gr.Row():400 metadata = gr.JSON(label="Run metadata")401 delta_plot = gr.Plot(label="Delta plot")402 403 with gr.Row():404 csv_download = gr.File(label="Download CSV")405 json_download = gr.File(label="Download JSON")406 407 run.click(408 predict,409 inputs=[model, reference_sequence, alternative_sequence, reference_features, alternative_features],410 outputs=[delta_table, metadata, delta_plot, csv_download, json_download],411 )412 demo.load(initial_model, outputs=model)413 414 415if __name__ == "__main__":416 demo.launch()417