{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":87793,"databundleVersionId":11553390,"sourceType":"competition"}],"dockerImageVersionId":31012,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# RiboDiffusion\n\nTrying out https://academic.oup.com/bioinformatics/article/40/Supplement_1/i347/7700903\n\nCopying from https://colab.research.google.com/drive/199D6B0FsIYf-gW-hfMEBCcKaai_hM_cU","metadata":{}},{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nassert torch.cuda.is_available(), \"WARNING! You are running on a non-GPU instance. A GPU is highly recommended.\"\nREQUIRED_VERSION = \"2.2.1+cu121\"\nTORCH_VERSION = torch.__version__\nif TORCH_VERSION != REQUIRED_VERSION:\n    print(f\"Detected torch version {TORCH_VERSION}, but notebook was created for {REQUIRED_VERSION}.\")\n    print(f\"Attempting installation of {REQUIRED_VERSION}\")\n    !pip install -q torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118\nprint(\"Correct version of torch detected. You are running on a machine with GPU.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install -q torch_geometric==2.3.1\n!pip install -q torch_scatter==2.1.1\n!pip install -q torch_cluster==1.6.1\n!pip install -q biopython==1.80 fair_esm==2.0.0 ml_collections==0.1.1","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!gdown 1TNab2MVPT0MXIxqizYfpd6YSAa5i8d4T\n!gdown 1-IfWkLa5asu4SeeZAQ09oWm4KlpBMPmq","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nif not os.path.exists(\"/content/RiboDiffusion\"):\n    !git clone --depth 1 -b main https://github.com/GRAPH-0/RiboDiffusion.git\n%cd /content/RiboDiffusion/","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!mkdir ckpts\n!mv ../exp_inf.pth ./ckpts\n!mv ../exp_inf_large.pth ./ckpts","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom tqdm import tqdm\nimport numpy as np\nimport random\nfrom models import *\nfrom utils import *\nfrom diffusion import NoiseScheduleVP\nfrom sampling import get_sampling_fn\nfrom datasets import utils as du\nimport functools\nimport tree\nfrom configs.inference_ribodiffusion import get_config","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\nconfig = get_config()\n# Choose heckpoint name\ncheckpoint_path = './ckpts/exp_inf.pth'\n# checkpoint_path = './ckpts/exp_inf_large.pth'","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"config.eval.sampling_steps = 50\n# config.eval.sampling_steps = 100","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\nNUM_TO_LETTER = np.array(['A', 'G', 'C', 'U'])\n\ndef get_optimizer(config, params):\n  \"\"\"Return a flax optimizer object based on `config`.\"\"\"\n  if config.optim.optimizer == 'Adam':\n      optimizer = optim.Adam(params, lr=config.optim.lr, betas=(config.optim.beta1, 0.999), eps=config.optim.eps, weight_decay=config.optim.weight_decay)\n  elif config.optim.optimizer == 'AdamW':\n      optimizer = torch.optim.AdamW(params, lr=config.optim.lr, amsgrad=True, weight_decay=1e-12)\n  else:\n      raise NotImplementedError(f'Optimizer {config.optim.optimizer} not supported yet!')\n  return optimizer\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Initialize model\nmodel = create_model(config)\nema = ExponentialMovingAverage(model.parameters(), decay=config.model.ema_decay)\nparams = model.parameters()\noptimizer = get_optimizer(config, model.parameters())\nstate = dict(optimizer=optimizer, model=model, ema=ema, step=0)\n\nmodel_size = sum(p.numel() for p in model.parameters()) * 4 / 2 ** 20\nprint('model size: {:.1f}MB'.format(model_size))\n\n# Load checkpoint\nstate = restore_checkpoint(checkpoint_path, state, device=config.device)\nema.copy_to(model.parameters())\n\n# Initialize noise scheduler\nnoise_scheduler = NoiseScheduleVP(config.sde.schedule, continuous_beta_0=config.sde.continuous_beta_0,\n                                  continuous_beta_1=config.sde.continuous_beta_1)\n# Obtain data scalar and inverse scalar\ninverse_scaler = get_data_inverse_scaler(config)\n\n# Setup sampling function\ntest_sampling_fn = get_sampling_fn(config, noise_scheduler, config.eval.sampling_steps, inverse_scaler)\npdb2data = functools.partial(du.PDBtoData, num_posenc=config.data.num_posenc, num_rbf=config.data.num_rbf, knn_num=config.data.knn_num)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}