Alphafold Database Fetch And Analyze

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Analyzes Predicted Aligned Error (PAE) and detects domain boundaries."""

# /// script
# requires-python = ">=3.10"
# dependencies = []
# ///

import argparse
import itertools
import json
import os


def find_sub_domains(pae_matrix, distance_cutoff=7.0, min_domain_size=40):
  """Identifies structurally independent sub-domains based on the PAE matrix."""
  n_res = len(pae_matrix)
  domains = []
  current_domain = []

  for i in range(n_res):
    if not current_domain:
      current_domain.append(i)
      continue

    window_size = min(20, len(current_domain))
    recent_res = current_domain[-window_size:]

    pae_sum = sum(pae_matrix[r][i] + pae_matrix[i][r] for r in recent_res)
    avg_pae = pae_sum / (2.0 * window_size)

    if avg_pae < distance_cutoff:
      current_domain.append(i)
    else:
      if len(current_domain) >= min_domain_size:
        domains.append(current_domain)
      current_domain = [i]

  if len(current_domain) >= min_domain_size:
    domains.append(current_domain)

  domain_boundaries = []
  for comp in domains:
    start = comp[0] + 1
    end = comp[-1] + 1
    domain_boundaries.append([start, end])

  return domain_boundaries


def merge_global_domains(boundaries, pae_matrix, merge_cutoff=15.0):
  """Merges sub-domains if the average PAE between them is below cutoff."""
  if not boundaries:
    return []

  if len(boundaries) == 1:
    merged = boundaries
  else:
    merged = [boundaries[0]]

    for i in range(1, len(boundaries)):
      prev_end = merged[-1][1] - 1
      curr_start = boundaries[i][0] - 1

      lookback = max(merged[-1][0] - 1, prev_end - 30)
      lookfwd = min(boundaries[i][1] - 1, curr_start + 30)

      pae_sum = 0
      n_pairs = 0
      for r1 in range(lookback, prev_end + 1):
        for r2 in range(curr_start, lookfwd + 1):
          pae_sum += pae_matrix[r1][r2] + pae_matrix[r2][r1]
          n_pairs += 2

      if n_pairs > 0 and (pae_sum / n_pairs) < merge_cutoff:
        merged[-1][1] = boundaries[i][1]
      else:
        merged.append(boundaries[i])

  filtered_merged = [dom for dom in merged if (dom[1] - dom[0] + 1) > 50]

  return filtered_merged


def analyze_pae(pae_file):
  """Parses a PAE JSON file and calculates structural domain metrics."""
  print(
      "\n[*] Analyzing Predicted Aligned Error (PAE) from"
      f" {os.path.basename(pae_file)}..."
  )
  try:
    with open(pae_file, "r") as f:
      data = json.load(f)[0]

    if "predicted_aligned_error" in data:
      pae = data["predicted_aligned_error"]
    elif "distance" in data:
      pae = data["distance"]
    else:
      print(
          "     [!] Could not locate PAE matrix in JSON keys:"
          f" {list(data.keys())}"
      )
      return

    flat_pae = list(itertools.chain.from_iterable(pae))
    if not flat_pae:
      print("     [!] PAE matrix is empty.")
      return

    mean_pae = sum(flat_pae) / len(flat_pae)
    max_pae = max(flat_pae)
    min_pae = min(flat_pae)
    confident_pairs = sum(1 for p in flat_pae if p < 5.0) / len(flat_pae) * 100

    print(f"  -> PAE Matrix Shape: {len(pae)}x{len(pae[0])}")
    print(f"  -> Mean Error: {mean_pae:.2f} Å")
    print(
        f"  -> Max Error: {max_pae:.2f} Å (suggests max possible distance"
        " between domains)"
    )
    print(f"  -> Min Error: {min_pae:.2f} Å")
    print(
        "  -> Fraction of confident residue pairs (<5Å PAE):"
        f" {confident_pairs:.1f}%"
    )

    sub_domains = find_sub_domains(pae, distance_cutoff=7.0, min_domain_size=40)
    global_domains = merge_global_domains(sub_domains, pae, merge_cutoff=15.0)

    print("\n[*] Domain Boundary Analysis:")
    if not global_domains:
      print("  -> No distinct rigidly-folded domains detected (>50 AAs).")
    else:
      print(
          "  -> Number of distinct Global Domains detected:"
          f" {len(global_domains)}"
      )
      for i, (start, end) in enumerate(global_domains, 1):
        print(
            f"     Domain {i}: residues {start} - {end} (Length:"
            f" {end - start + 1} AAs)"
        )

    print("\n[*] PAE Structural Conclusion:")
    if len(global_domains) == 1:
      conclusion = (
          "The protein consists of a single well-folded, rigid composite"
          " domain."
      )
    elif len(global_domains) > 1:
      conclusion = (
          f"The protein has {len(global_domains)} independently positioned"
          " global domains separated by truly flexible joints."
      )
    else:
      conclusion = (
          "The protein is likely entirely disordered or lacks rigid"
          " tertiary structure."
      )
    print(f"  -> {conclusion}")

    return {
        "pae_file": os.path.basename(pae_file),
        "matrix_shape": f"{len(pae)}x{len(pae[0])}",
        "mean_pae": round(mean_pae, 2),
        "max_pae": round(max_pae, 2),
        "min_pae": round(min_pae, 2),
        "confident_pairs_pct": round(confident_pairs, 1),
        "domains": [
            {"start": s, "end": e, "length": e - s + 1}
            for s, e in global_domains
        ],
        "conclusion": conclusion,
    }

  except (IOError, json.JSONDecodeError) as e:
    print(f"     [!] Failed to analyze PAE file: {e}")
    return None


if __name__ == "__main__":
  parser = argparse.ArgumentParser(
      description=(
          "Analyze PAE matrix and detect domain boundaries from an AlphaFold"
          " PAE JSON file"
      )
  )
  parser.add_argument(
      "pae_file",
      help=(
          "Path to the PAE JSON file (e.g.,"
          " AF-P04637-F1-predicted_aligned_error_v6.json)"
      ),
  )

  args = parser.parse_args()

  analyze_pae(args.pae_file)

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Analyzes pLDDT confidence metrics from a saved AFDB metadata JSON file."""

# /// script
# requires-python = ">=3.10"
# dependencies = []
# ///

import argparse
import json
import sys

CONFIDENT_THRESHOLD = 0.7
MODERATE_THRESHOLD = 0.4
NOTABLE_DISORDER_THRESHOLD = 0.15
MIXED_DISORDER_THRESHOLD = 0.3
MOSTLY_DISORDERED_THRESHOLD = 0.5


def analyze_plddt(metadata_file):
  """Loads an AFDB metadata JSON and analyzes pLDDT metrics."""
  try:
    with open(metadata_file, "r") as f:
      entry = json.load(f)
  except (IOError, json.JSONDecodeError) as e:
    print(f"[!] Error reading metadata file: {e}")
    sys.exit(1)

  return _analyze_entry(entry)


def _analyze_entry(entry):
  """Parses and analyzes pLDDT confidence metrics from the API payload."""
  accession = entry.get("uniprotAccession", "Unknown")
  print(f"\n[*] AlphaFold pLDDT Metrics for Accession: {accession}")
  global_plddt = entry.get("globalMetricValue", 0.0)
  frac_vlow = entry.get("fractionPlddtVeryLow", 0.0)
  frac_low = entry.get("fractionPlddtLow", 0.0)
  frac_conf = entry.get("fractionPlddtConfident", 0.0)
  frac_vhigh = entry.get("fractionPlddtVeryHigh", 0.0)

  print("-" * 65)
  print(f"  -> Overall Global pLDDT   : {global_plddt:.2f}")
  print(f"  -> Fraction Very Low (<50): {frac_vlow:.3f} ({frac_vlow*100:.1f}%)")
  print(f"  -> Fraction Low (50-70)   : {frac_low:.3f} ({frac_low*100:.1f}%)")
  print(f"  -> Fraction Confident     : {frac_conf:.3f} ({frac_conf*100:.1f}%)")
  print(
      f"  -> Fraction Very High     : {frac_vhigh:.3f} ({frac_vhigh*100:.1f}%)"
  )
  print("-" * 65)

  conf_total = frac_conf + frac_vhigh

  print("[*] pLDDT Conclusion:")
  if conf_total >= CONFIDENT_THRESHOLD:
    if frac_vlow > NOTABLE_DISORDER_THRESHOLD:
      conclusion = (
          "Protein is mostly confidently predicted, but contains notable"
          " disordered regions."
      )
    else:
      conclusion = (
          "Protein is confidently predicted and likely fully"
          " ordered/structured."
      )
  elif conf_total >= MODERATE_THRESHOLD:
    if frac_vlow >= MIXED_DISORDER_THRESHOLD:
      conclusion = (
          "Protein has a mixture of confidently predicted structured"
          " domains and significant intrinsically disordered regions."
      )
    else:
      conclusion = (
          "Protein has moderate prediction confidence. Certain regions"
          " might be flexible or poorly predicted."
      )
  else:
    if frac_vlow >= MOSTLY_DISORDERED_THRESHOLD:
      conclusion = (
          "Protein is mostly poorly predicted, likely being highly"
          " intrinsically disordered."
      )
    else:
      conclusion = "Protein prediction is of low confidence overall."
  print(f"  -> {conclusion}")
  print()

  return {
      "uniprot_id": accession,
      "global_plddt": global_plddt,
      "fractions": {
          "very_low": frac_vlow,
          "low": frac_low,
          "confident": frac_conf,
          "very_high": frac_vhigh,
      },
      "conclusion": conclusion,
  }


if __name__ == "__main__":
  parser = argparse.ArgumentParser(
      description="Analyze pLDDT confidence metrics from an AFDB metadata file"
  )
  parser.add_argument(
      "metadata_file",
      help="Path to the metadata JSON file (e.g., AF-P04637-F1-metadata.json)",
  )

  args = parser.parse_args()

  analyze_plddt(args.metadata_file)

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Fetches AlphaFold structure files (mmCIF + PAE) for a UniProt ID."""

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "scienceskillscommon",
# ]
# [tool.uv.sources]
# scienceskillscommon = { path = "../../scienceskillscommon" }
# ///

import argparse
import json
import os
import sys
from science_skills.skills.scienceskillscommon import http_client

CLIENT = http_client.HttpClient("https://alphafold.ebi.ac.uk", qps=1.0)


def fetch_structure(uniprot_id, output_dir):
  """Downloads the mmCIF file and PAE JSON for a given UniProt ID from AFDB."""
  uniprot_id = uniprot_id.strip().upper()
  api_url = f"https://alphafold.ebi.ac.uk/api/prediction/{uniprot_id}"
  is_fragment = False

  os.makedirs(output_dir, exist_ok=True)
  print(f"[*] Requesting AlphaFold data for UniProt ID: {uniprot_id}")

  try:
    data = CLIENT.fetch_json(api_url)
  except http_client.HttpError as e:
    if e.status_code == 404:
      print(
          f"\n[!] Error: UniProt ID '{uniprot_id}' was not found in"
          " the AlphaFold Database."
      )
      print(
          "    Please double-check the ID for typos, or verify it"
          " has an AFDB entry."
      )
    else:
      print(f"[!] HTTP error fetching API data: {e}")
    sys.exit(1)

  if not data:
    print(f"[!] No AlphaFold data returned for {uniprot_id}")
    sys.exit(1)

  # The API may return multiple entries (e.g. isoforms) for a single
  # UniProt ID. Prefer the canonical entry whose accession matches exactly.
  entry = None
  for e in data:
    if e.get("uniprotAccession") == uniprot_id:
      entry = e
      break
  # If no canonical entry exists (common for very large proteins like
  # Dystrophin), fall back to the longest available isoform so the user
  # gets the most complete structure.
  if entry is None:
    entry = max(data, key=lambda e: e.get("sequenceEnd", 0))
    entry_acc = entry.get("uniprotAccession", "unknown")
    entry_len = entry.get("sequenceEnd", 0)
    print(
        "[!] WARNING: No canonical AFDB entry found for"
        f" '{uniprot_id}'. Using longest available isoform"
        f" '{entry_acc}' ({entry_len} amino acids) instead."
        " The full-length protein may not be available in AFDB."
    )

  max_seq_len = max((e.get("sequenceEnd", 0) for e in data), default=0)
  if max_seq_len > 2700:
    print(
        f"[!] WARNING: Protein {uniprot_id} is massive"
        f" ({max_seq_len} amino acids). Only the first entry has"
        " been downloaded. The full protein may span many more"
        " fragments in AFDB."
    )
    is_fragment = True

  entry_acc = entry.get("uniprotAccession", uniprot_id)
  metadata_filename = f"AF-{entry_acc}-F1-metadata.json"
  metadata_path = os.path.join(output_dir, metadata_filename)
  with open(metadata_path, "w") as f:
    json.dump(entry, f, indent=2)
  print(f"  -> Saved API metadata to: {metadata_path}")

  cif_url = entry.get("cifUrl")
  pae_url = entry.get("paeDocUrl")

  urls_to_fetch = []
  if cif_url:
    urls_to_fetch.append(cif_url)
  if pae_url:
    urls_to_fetch.append(pae_url)

  success_count = 0

  for url in urls_to_fetch:
    filename = url.split("/")[-1]
    file_path = os.path.join(output_dir, filename)

    print(f"  -> Fetching {filename}...")

    try:
      file_bytes = CLIENT.fetch_bytes(url)
      with open(file_path, "wb") as f:
        f.write(file_bytes)

      print(f"     [+] Saved to: {file_path}")
      success_count += 1

    except http_client.HttpError as e:
      if e.status_code == 404:
        print(f"     [!] Error 404: {filename} not found.")
      else:
        print(f"     [!] Download error: {e}")

  if success_count == 0:
    print(
        f"\n[!] Failed to download any data for {uniprot_id}. Please check"
        " the ID."
    )
    sys.exit(1)
  else:
    print(
        f"\n[*] Successfully downloaded {success_count}/{len(urls_to_fetch)}"
        " files."
    )

  return {
      "uniprot_id": uniprot_id,
      "output_dir": output_dir,
      "is_fragment": is_fragment,
      "files_downloaded": success_count,
      "metadata_file": metadata_path,
  }


if __name__ == "__main__":
  parser = argparse.ArgumentParser(
      description="Download AlphaFold structure files for a UniProt ID"
  )
  parser.add_argument(
      "uniprot_id", help="The UniProt ID (e.g., P04637 or A0A1B0GX81)"
  )
  parser.add_argument(
      "-o",
      "--output-dir",
      help="Output directory to save the files (required)",
      required=True,
  )

  args = parser.parse_args()

  fetch_structure(args.uniprot_id, args.output_dir)