| import torch |
|
|
| def set_seed(seed=42): |
| """Set random seed for reproducibility.""" |
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| |
| def get_chain_organism(pdb_id, chain_id): |
| """ |
| Use RCSB PDB API to get the organism of the chain. |
| """ |
| import requests |
| entry_url = f"https://data.rcsb.org/rest/v1/core/entry/{pdb_id.lower()}" |
| res = requests.get(entry_url) |
| if res.status_code != 200: |
| return "Unknown" |
| entry_data = res.json() |
| |
| |
| chain_to_entity = {} |
| for entity_id in entry_data.get("rcsb_entry_container_identifiers", {}).get("polymer_entity_ids", []): |
| entity_url = f"https://data.rcsb.org/rest/v1/core/polymer_entity/{pdb_id.lower()}/{entity_id}" |
| entity_res = requests.get(entity_url) |
| if entity_res.status_code != 200: |
| continue |
| entity_data = entity_res.json() |
| chains = entity_data.get("rcsb_polymer_entity_container_identifiers", {}).get("auth_asym_ids", []) |
| for c in chains: |
| chain_to_entity[c] = entity_id |
| if chain_id in chains: |
| organism = entity_data.get("rcsb_entity_source_organism", [{}])[0].get("scientific_name", "Unknown") |
| return organism |
| return "Unknown" |
|
|
| def classify_antigen(organisms): |
| for org in organisms: |
| if "virus" in org.lower() or "coronavirus" in org.lower(): |
| return "viral" |
| elif "homo sapiens" in org.lower(): |
| return "human" |
| elif "bacteria" in org.lower() or "bacillus" in org.lower(): |
| return "bacterial" |
| elif "tumor" in org.lower() or "cancer" in org.lower(): |
| return "tumor" |
| return "other" |
