{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":11403143,"sourceType":"competition"},{"sourceId":10855324,"sourceType":"datasetVersion","datasetId":6742586},{"sourceId":10880419,"sourceType":"datasetVersion","datasetId":6760509},{"sourceId":11096987,"sourceType":"datasetVersion","datasetId":6917280}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"from datetime import datetime\nimport pytz\nprint('LOGGING TIME OF START:',  datetime.strftime(datetime.now(pytz.timezone('Asia/Singapore')), \"%Y-%m-%d %H:%M:%S\"))\n\n\ntry:\n    import Bio\nexcept:\n    #for drfold2 --------\n    #!pip install biopython\n    !pip install /kaggle/input/biopython/biopython-1.85-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n\nprint('PIP INSTALL OK !!!!')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-03-20T04:37:32.455875Z","iopub.execute_input":"2025-03-20T04:37:32.456331Z","iopub.status.idle":"2025-03-20T04:37:32.463206Z","shell.execute_reply.started":"2025-03-20T04:37:32.456291Z","shell.execute_reply":"2025-03-20T04:37:32.462300Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os,sys\n\nimport pandas as pd\npd.set_option('display.max_columns', 20)\npd.set_option('display.expand_frame_repr', False)\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom timeit import default_timer as timer\nimport re\n\nimport matplotlib \nimport matplotlib.pyplot as plt\n\n\n# helper--\nclass dotdict(dict):\n\t__setattr__ = dict.__setitem__\n\t__delattr__ = dict.__delitem__\n\n\tdef __getattr__(self, name):\n\t\ttry:\n\t\t\treturn self[name]\n\t\texcept KeyError:\n\t\t\traise AttributeError(name)\n\ndef time_to_str(t, mode='min'):\n\tif mode=='min':\n\t\tt  = int(t)/60\n\t\thr = t//60\n\t\tmin = t%60\n\t\treturn '%2d hr %02d min'%(hr,min) \n\telif mode=='sec':\n\t\tt   = int(t)\n\t\tmin = t//60\n\t\tsec = t%60\n\t\treturn '%2d min %02d sec'%(min,sec)\n\n\telse:\n\t\traise NotImplementedError\n\ndef gpu_memory_use():\n    if torch.cuda.is_available():\n        device = torch.device(0)\n        free, total = torch.cuda.mem_get_info(device)\n        used= (total - free) / 1024 ** 3\n        return int(round(used))\n    else:\n        return 0\n\ndef set_aspect_equal(ax):\n\tx_limits = ax.get_xlim()\n\ty_limits = ax.get_ylim()\n\tz_limits = ax.get_zlim()\n\n\t# Compute the mean of each axis\n\tx_middle = np.mean(x_limits)\n\ty_middle = np.mean(y_limits)\n\tz_middle = np.mean(z_limits)\n\n\t# Compute the max range across all axes\n\tmax_range = max(x_limits[1] - x_limits[0],\n\t\t\t\t\ty_limits[1] - y_limits[0],\n\t\t\t\t\tz_limits[1] - z_limits[0]) / 2.0\n\n\t# Set the new limits to ensure equal scaling\n\tax.set_xlim(x_middle - max_range, x_middle + max_range)\n\tax.set_ylim(y_middle - max_range, y_middle + max_range)\n\tax.set_zlim(z_middle - max_range, z_middle + max_range)\n\n\nprint('torch',torch.__version__)\nprint('torch.cuda',torch.version.cuda)\n\nprint('IMPORT OK!!!')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-20T04:37:32.464515Z","iopub.execute_input":"2025-03-20T04:37:32.464783Z","iopub.status.idle":"2025-03-20T04:37:32.486618Z","shell.execute_reply.started":"2025-03-20T04:37:32.464762Z","shell.execute_reply":"2025-03-20T04:37:32.485821Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"MODE = 'submit' #'local' # submit\n\nDATA_KAGGLE_DIR = '/kaggle/input/stanford-rna-3d-folding'\nif MODE == 'local':\n    valid_df = pd.read_csv(f'{DATA_KAGGLE_DIR}/validation_sequences.csv')\n    label_df = pd.read_csv(f'{DATA_KAGGLE_DIR}/validation_labels.csv')\n    label_df['target_id'] = label_df['ID'].apply(lambda x: '_'.join(x.split('_')[:-1]))\n\nif MODE == 'submit':\n\tvalid_df = pd.read_csv(f'{DATA_KAGGLE_DIR}/test_sequences.csv')\n\nprint('len(valid_df)',len(valid_df))\nprint(valid_df.iloc[0])\nprint('')\n\n\n# cfg = dotdict(\n#     num_conf = 5,\n#     max_length=480,\n# )\nNUM_CONF=5\nMAX_LENGTH=480\nDEVICE='cuda'#'cuda' #'cpu'#\n\nprint('MODE:', MODE)\nprint('SETTING OK!!!')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-20T04:37:32.488099Z","iopub.execute_input":"2025-03-20T04:37:32.488423Z","iopub.status.idle":"2025-03-20T04:37:32.516691Z","shell.execute_reply.started":"2025-03-20T04:37:32.488395Z","shell.execute_reply":"2025-03-20T04:37:32.515809Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sys.path.append('/kaggle/input/hengck23-rna-02/drfold2/cfg_97')\nfrom EvoMSA2XYZ.Model import MSA2XYZ\nfrom RNALM2.Model import RNA2nd\nfrom data import parse_seq, Get_base, BASE_COOR\nfrom data import (\n    write_frame_coor_to_pdb,\n    parse_pdb_to_xyz,\n    frame_coor_to_C1\n)\nfrom energy_scorer import score_energy_one\n\n\n###########################################################3\nKAGGLE_TRUTH_PDB_DIR ='/kaggle/input/hengck23-drfold2-dummy-00/kaggle-casp15-truth'\nUSALIGN = '/kaggle/working/USalign' \nos.system('cp /kaggle/input/usalign/USalign /kaggle/working/')\nos.system('sudo chmod u+x /kaggle/working/USalign')\n\n# evaluate helper\ndef get_truth_df(target_id, label_df):\n    truth_df = label_df[label_df['target_id'] == target_id]\n    truth_df = truth_df.reset_index(drop=True)\n    return truth_df\n\ndef parse_usalign_for_tm_score(output):\n    # Extract TM-score based on length of reference structure (second)\n    tm_score_match = re.findall(r'TM-score=\\s+([\\d.]+)', output)[1]\n    if not tm_score_match:\n        raise ValueError('No TM score found')\n    return float(tm_score_match)\n\ndef parse_usalign_for_transform(output):\n    # Locate the rotation matrix section\n    matrix_lines = []\n    found_matrix = False\n\n    for line in output.splitlines():\n        if \"The rotation matrix to rotate Structure_1 to Structure_2\" in line:\n            found_matrix = True\n        elif found_matrix and re.match(r'^\\d+\\s+[-\\d.]+\\s+[-\\d.]+\\s+[-\\d.]+\\s+[-\\d.]+$', line):\n            matrix_lines.append(line)\n        elif found_matrix and not line.strip():\n            break  # Stop parsing if an empty line is encountered after the matrix\n\n    # Parse the rotation matrix values\n    rotation_matrix = []\n    for line in matrix_lines:\n        parts = line.split()\n        row_values = list(map(float, parts[1:]))  # Skip the first column (index)\n        rotation_matrix.append(row_values)\n    return np.array(rotation_matrix)\n\n\n\n# data helper\ndef make_data(seq, device):\n    aa_type = parse_seq(seq)\n    base = Get_base(seq, BASE_COOR)\n    seq_idx = np.arange(len(seq)) + 1\n\n    msa = aa_type[None, :]\n    msa = torch.from_numpy(msa)\n    msa = torch.cat([msa, msa], 0)  # ???\n    msa = F.one_hot(msa.long(), 6).float()\n\n    base_x = torch.from_numpy(base).float()\n    seq_idx = torch.from_numpy(seq_idx).long()\n\n    msa, base_x, seq_idx = msa.to(device), base_x.to(device), seq_idx.to(device)\n    return msa, base_x, seq_idx\n\n    \ndef make_dummy_solution():\n    solution=dotdict()\n    for i, row in valid_df.iterrows():\n        target_id = row.target_id\n        sequence = row.sequence\n        solution[target_id]=dotdict(\n            target_id=target_id,\n            sequence=sequence,\n            coord=[],\n        )\n    return solution\n\ndef solution_to_submit_df(solution):\n    submit_df = []\n    for k,s in solution.items():\n        df = coord_to_df(s.sequence, s.coord, s.target_id)\n        submit_df.append(df)\n    \n    submit_df = pd.concat(submit_df)\n    return submit_df\n \n\ndef coord_to_df(sequence, coord, target_id):\n    L = len(sequence)\n    df = pd.DataFrame()\n    df['ID'] = [f'{target_id}_{i + 1}' for i in range(L)]\n    df['resname'] = [s for s in sequence]\n    df['resid'] = [i + 1 for i in range(L)]\n\n    num_coord = len(coord)\n    for j in range(num_coord):\n        df[f'x_{j+1}'] = coord[j][:, 0]\n        df[f'y_{j+1}'] = coord[j][:, 1]\n        df[f'z_{j+1}'] = coord[j][:, 2]\n    return df\n\n################### start here !!! #######################################################3\n\n\nout_dir = '/kaggle/working/model-output'\nos.makedirs(out_dir, exist_ok=True)\n\n\ndef run_submit(valid_df):\n    \n    #load model (these are moified versions, not the same from their github repo)\n    rnalm = RNA2nd(dict(\n        s_in_dim=5,\n        z_in_dim=2,\n        s_dim= 512,\n        z_dim= 128,\n        N_elayers=18,\n    ))\n    rnalm_file = '/kaggle/input/hengck23-rna-02/weight/RCLM/epoch_67000'\n    print(rnalm_file)\n    print(\n        rnalm.load_state_dict(torch.load(rnalm_file, map_location='cpu', weights_only=True), strict=False)\n        #Unexpected key(s) in state_dict: \"ss_head.linear.weight\", \"ss_head.linear.bias\".\n    )\n    rnalm = rnalm.to(DEVICE)\n    rnalm = rnalm.eval()\n    \n    #---\n    msa2xyz = MSA2XYZ(dict(\n        seq_dim=6,\n        msa_dim=7,\n        N_ensemble=1,\n        N_cycle=8,  # 8\n        m_dim=64,\n        s_dim=64,\n        z_dim=64,\n    ))\n    msa2xyz_file = [\n        f'/kaggle/input/hengck23-rna-02/weight/cfg_97/model_{i}'\n        for i in range(20)\n    ]\n    num_msa2xyz = len(msa2xyz_file) \n    msa2xyz_state_dict = []\n    for c in range(num_msa2xyz):\n        if c==0: print(msa2xyz_file[c])\n        m = torch.load(msa2xyz_file[c], map_location='cpu', weights_only=True)\n        msa2xyz_state_dict.append(m)\n        \n    print(msa2xyz.load_state_dict(msa2xyz_state_dict[0], strict=True))\n    msa2xyz = msa2xyz.to(DEVICE)\n    msa2xyz = msa2xyz.eval()\n    msa2xyz.msaxyzone.premsa.rnalm = rnalm\n\n    #---\n    # start here !!!!!!!!!!!!!!!!!!!!!!\n    #valid_df = valid_df.iloc[[0,1]].reset_index(drop=True)\n\n\n    submit_df = [] \n    total_time_taken = 0\n    max_gpu_mem_used = 0\n\n    for i, row in valid_df.iterrows():\n        start_timer = timer()\n        target_id = row.target_id  # 'R1116' #casp15 R1116: len(157)\n        sequence = row.sequence\n        seq = row.sequence  \n        L = len(seq)\n        if L > MAX_LENGTH:\n            seq = seq[:MAX_LENGTH]\n        print(i, target_id, L, len(seq), seq[:75] + '...')\n\n        msa, base_x, seq_idx = make_data(seq, DEVICE)\n        secondary = None  # secondary structure\n\n        if len(seq)>200:\n            model_to_try = [0,1,2,8,9]\n        elif  len(seq)>100:\n            model_to_try = [0,2,4,6,8,10,12,14,16,18]#list(range(min(num_msa2xyz,10)))\n        else:\n            model_to_try = list(range(num_msa2xyz))\n\n        energy = []\n        netout = []\n        coordinate=[]\n        for c in model_to_try:\n            msa2xyz.load_state_dict(msa2xyz_state_dict[c], strict=False)\n            #print(c, msa2xyz.training, msa2xyz.msa_predor.linear.weight.device, msa2xyz.msa_predor.linear.weight.reshape(-1)[:5])\n\n            #with torch.amp.autocast('cuda',dtype=torch.float16):\n            with torch.no_grad():\n                out = msa2xyz.pred(msa, seq_idx, secondary, base_x, np.array(list(seq)))\n\n            if len(model_to_try)>5:\n                e = score_energy_one(seq, target_id, out)\n            else:\n                e = 0\n            energy.append(e) #tranucated sequence\n\n \n            if L != len(seq):\n                out['coor'] = np.pad(out['coor'] ,((0, L - len(seq)), (0, 0), (0, 0)), 'constant', constant_values=0)\n            netout.append(out)\n            \n            xyz = frame_coor_to_C1(out['coor'], sequence)\n            coordinate.append(xyz)\n            \n            #print('out:', out['coor'].shape)\n\n            time_taken = timer() - start_timer\n            total_time_taken += time_taken\n            #print('time_taken:', time_to_str(time_taken, mode='sec'))\n\n            gpu_mem_used = gpu_memory_use()\n            max_gpu_mem_used = max(max_gpu_mem_used,gpu_mem_used)\n            #print('gpu_mem_used:', gpu_mem_used, 'GB')\n\n            print(f'{c:02d}   energy:{e:10.0f}   out{str(out[\"coor\"].shape)}  time:{time_to_str(time_taken, mode=\"sec\")}   gpu={gpu_mem_used} gb')\n\n\n        #------- \n        torch.cuda.empty_cache()\n        \n        #select top5\n        argsort = np.array(energy).argsort()\n        argsort = argsort[:5]\n        df = coord_to_df(row.sequence, [coordinate[k] for k in argsort], row.target_id)\n        submit_df.append(df)\n    \n    print('----------------------------------------')\n    print('MAX_LENGTH', MAX_LENGTH)\n    print('### total_time_taken:', time_to_str(total_time_taken, mode='min'))\n    print('### max_gpu_mem_used:', max_gpu_mem_used, 'GB')\n    print('')\n\n    submit_df = pd.concat(submit_df)\n    submit_df.to_csv(f'submission.csv', index=False)\n    print(submit_df) \n    return submit_df\n  \nrun_submit(valid_df)\n\nprint('SUBMIT OK!!!')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-20T04:37:32.647512Z","iopub.execute_input":"2025-03-20T04:37:32.647841Z","execution_failed":"2025-03-20T04:38:01.847Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if MODE == 'local':\n    from data import write2pdb_for_tm\n\n    # local validation\n    submit_df = pd.read_csv('submission.csv')\n    submit_df['target_id'] = submit_df['ID'].apply(lambda x: '_'.join(x.split('_')[:-1]))\n\n    tm_score = []\n    for i, row in valid_df.iterrows():\n        target_id = row.target_id  # 'R1116' #casp15 R1116: len(157)\n        seq = row.sequence\n        # -----------------------------------------------\n        print(i, target_id, len(seq), seq[:75] + '...')\n\n        truth_pdb = f'{KAGGLE_TRUTH_PDB_DIR}/kaggle_truth_{target_id}_C1.pdb'\n        # print(os.path.isfile(predict_pdb))\n\n        tm = []\n        for c in range(5):\n            # predict_pdb = f'{out_dir}/{target_id}-coor.{c:02d}.pdb'\n            # print(os.path.isfile(truth_pdb))\n            df = submit_df[submit_df['target_id'] == target_id]\n            xyz = df.loc[:, [f'x_{c+1}',f'y_{c+1}',f'z_{c+1}']].values\n            predict_pdb =  '/kaggle/working/predict.pdb'\n            write2pdb_for_tm(seq, xyz, target_id, predict_pdb)\n\n            command = f'{USALIGN} {predict_pdb} {truth_pdb} -atom \" C1\\'\" -m -'\n            output = os.popen(command).read()\n            # print(output)\n            try:\n                tm_c = parse_usalign_for_tm_score(output)\n            except:\n                tm_c = 0\n            tm.append(tm_c)\n        print('### tm:', tm)\n        tm_score.append(max(tm))\n\n    print('ALL\\n', tm_score)\n    print('MEAN', np.array(tm_score).mean())\n\n","metadata":{"trusted":true,"execution":{"execution_failed":"2025-03-20T04:38:01.848Z"}},"outputs":[],"execution_count":null}]}