{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.11"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":87793,"databundleVersionId":12276181,"sourceType":"competition"},{"sourceId":11769888,"sourceType":"datasetVersion","datasetId":7389227},{"sourceId":11769898,"sourceType":"datasetVersion","datasetId":7388010},{"sourceId":11858978,"sourceType":"datasetVersion","datasetId":7406538},{"sourceId":11871860,"sourceType":"datasetVersion","datasetId":7403946},{"sourceId":11908669,"sourceType":"datasetVersion","datasetId":7405022},{"sourceId":11922858,"sourceType":"datasetVersion","datasetId":7495841},{"sourceId":11938875,"sourceType":"datasetVersion","datasetId":7505933},{"sourceId":11988454,"sourceType":"datasetVersion","datasetId":7540404},{"sourceId":11995580,"sourceType":"datasetVersion","datasetId":7544186}],"dockerImageVersionId":31011,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"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)\nimport traceback;\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\n#for 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":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:44.780708Z","iopub.execute_input":"2025-05-29T13:01:44.780915Z","iopub.status.idle":"2025-05-29T13:01:46.600847Z","shell.execute_reply.started":"2025-05-29T13:01:44.780898Z","shell.execute_reply":"2025-05-29T13:01:46.600074Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ndef save_mmcif(mmcif_dict, output_file):\n    \"\"\"\n    Save the MMCIF dictionary using BioPython's MMCIFIO\n    \n    Args:\n        mmcif_dict: MMCIF dictionary with atom_site data\n        output_file: Output file path\n    \"\"\"\n    from Bio.PDB.mmcifio import MMCIFIO\n    \n    mmio = MMCIFIO()\n    mmio.set_dict(mmcif_dict)\n    with tempfile.TemporaryDirectory() as tempdir:\n        outname_ = tempdir+\"/bio-pdb-mmcifio-out.cif_\";\n        mmio.save(outname_);\n        fout = open(output_file,\"wt\");\n        with open(outname_,\"rt\") as fin:\n            for ll in fin:\n                if ll.startswith(\"ATOM\") or ll.startswith(\"HETATM\"):\n                    ll = re.sub(r\"([\\s])'([^\\s]+)'([\\s])\",r'\\1\"\\2\"\\3',ll);\n                fout.write(ll);\n        fout.close();\n# 今は一本の chain にしか対応してない\ndef calc_average_plddt(mdict,selected_index = None):\n    numatom = len(mdict[\"_atom_site.label_seq_id\"]);\n    dummy_list = [\"\" for _ in range(numatom)];\n    label_seq_id  = mdict.get(\"_atom_site.label_seq_id\",dummy_list);\n    label_atom_id  = mdict.get(\"_atom_site.label_atom_id\",dummy_list);\n    B_iso_or_equiv  = mdict.get(\"_atom_site.B_iso_or_equiv\",[\"0.0\" for _ in range(numatom)]);\n    asum = 0;\n    acou = 0;\n    for eii,(aid,sid,bfac) in enumerate(zip(\n        label_atom_id\n        ,label_seq_id\n        ,B_iso_or_equiv\n    )):\n        if selected_index is not None:\n            if int(sid) not in selected_index:\n                continue;\n        if aid == \"C1'\":\n            asum += float(bfac);\n            acou += 1;\n    if acou == 0:\n        return 0;\n    return asum/acou;\n        \ndef split_chain_mmcif(filename):\n    if filename.endswith(\".gz\"):\n        fin = gzip.open(filename,\"rt\");\n    else:\n        fin = open(filename,\"rt\");\n        \n    from Bio.PDB.MMCIF2Dict import MMCIF2Dict\n    mdict = MMCIF2Dict(fin);\n    fin.close();\n    \n    res_3_1_ = { \"ALA\":\"A\", \"ARG\":\"R\", \"ASN\":\"N\", \"ASP\":\"D\", \"CYS\":\"C\", \"GLN\":\"Q\", \"GLU\":\"E\", \"GLY\":\"G\", \"HIS\":\"H\", \"ILE\":\"I\", \"LEU\":\"L\", \"LYS\":\"K\", \"MET\":\"M\", \"PHE\":\"F\", \"PRO\":\"P\", \"SER\":\"S\", \"THR\":\"T\", \"TRP\":\"W\", \"TYR\":\"Y\", \"VAL\":\"V\", \"MSE\":\"M\"\n    , \"A\": \"A\", \"U\": \"U\", \"G\": \"G\", \"C\": \"C\"\n    , \"DA\": \"A\", \"DT\": \"T\", \"DG\": \"G\", \"DC\": \"C\"};\n    \n    atom_site_keys = [\n        \"_atom_site.label_comp_id\", \n        \"_atom_site.label_asym_id\", \n        \"_atom_site.label_entity_id\", \n        \"_atom_site.label_seq_id\", \n        \"_atom_site.pdbx_PDB_ins_code\",\n        \"_atom_site.auth_asym_id\", \n        \"_atom_site.pdbx_PDB_model_num\", \n    ];\n    numatom = len(mdict[\"_atom_site.label_seq_id\"]);\n    dummy_list = [\"\" for _ in range(numatom)];\n    \n    label_comp_id  = mdict.get(\"_atom_site.label_comp_id\",dummy_list);\n    label_asym_id  = mdict.get(\"_atom_site.label_asym_id\",dummy_list);\n    label_entity_id  = mdict.get(\"_atom_site.label_entity_id\",dummy_list);\n    label_seq_id  = mdict.get(\"_atom_site.label_seq_id\",dummy_list);\n    auth_seq_id  = mdict.get(\"_atom_site.auth_seq_id\",dummy_list);\n    pdbx_PDB_ins_code  = mdict.get(\"_atom_site.pdbx_PDB_ins_code\",dummy_list);\n    auth_asym_id  = mdict.get(\"_atom_site.auth_asym_id\",dummy_list);\n    pdbx_PDB_model_num  = mdict.get(\"_atom_site.pdbx_PDB_model_num\",dummy_list);\n    B_iso_or_equiv  = mdict.get(\"_atom_site.B_iso_or_equiv\",[\"0.0\" for _ in range(numatom)]);\n\n    \n    chain_to_index = {};\n    ret = {};\n    for eii,z in enumerate(zip(\n        label_asym_id\n        ,label_entity_id\n        ,auth_asym_id\n        ,pdbx_PDB_model_num\n    )):\n        chaincode = \"#\".join(z);\n        if chaincode not in chain_to_index:\n            chain_to_index[chaincode] = [];\n        chain_to_index[chaincode].append(eii);\n        \n    ret = [];\n\n    atom_site_keys = [];\n    for k in list(mdict.keys()):\n        if \"_atom_site.\" in k:\n            atom_site_keys.append(k);\n\n    for k in list(chain_to_index.keys()):\n        retdict = {};\n        retfas = [];\n        processed = set();\n        bfac = 0.0;\n        for i in list(chain_to_index[k]):\n            # chain は別れているのでこれだけ\n            code = label_comp_id[i]+\"#\"+label_seq_id[i]+\"#\"+pdbx_PDB_ins_code[i]+\"#\"\n            if code in processed:\n                continue;\n            retfas.append(res_3_1_.get(label_comp_id[i],\"X\"));\n            processed.add(code);\n            bfac += float(B_iso_or_equiv[i]);\n        bfac /= len(processed);\n        for ak in list(atom_site_keys):\n            retdict[ak] = [];\n            for i in list(chain_to_index[k]): \n                retdict[ak].append(mdict[ak][i]);\n        ret.append({\"code\":k,\"seq\":\"\".join(retfas),\"atom_site\":retdict,\"average_plddt\":bfac});\n    print(chain_to_index.keys());\n    return ret;\n    \n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.602156Z","iopub.execute_input":"2025-05-29T13:01:46.602468Z","iopub.status.idle":"2025-05-29T13:01:46.621757Z","shell.execute_reply.started":"2025-05-29T13:01:46.602448Z","shell.execute_reply":"2025-05-29T13:01:46.620922Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def loadFasta(filename):\n    fin = open(filename,\"r\");\n    ret = [];\n    cdict = dict();\n    cdict[\"seq\"] = \"\";\n    ret.append(cdict);\n    \n    for ll in fin:\n        mat = re.search(r\"[\\s]*>\",ll);\n        if(not mat == None):\n            cdict = dict();\n            ret.append(cdict);\n            nmat = re.search(r\"[\\s]*>[\\s]*([^\\s]+)\",ll);\n            if(not nmat == None):\n                cdict[\"name\"] = nmat.group(1);\n            dmat = re.search(r\"[\\s]*>[\\s]*([^\\s]+)[\\s]+([^\\s][^\\r\\n]*)\",ll);\n            cdict[\"desc\"] = \"\";\n            if(not dmat == None):\n                cdict[\"desc\"] = dmat.group(2);\n            cdict[\"seq\"] = \"\";\n        else:\n            cdict[\"seq\"] += re.sub(r\"[\\s]\",\"\",ll);\n            \n    if(len(ret[0][\"seq\"]) == 0):\n        ret.pop(0);\n    fin.close();\n    return ret;","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.622538Z","iopub.execute_input":"2025-05-29T13:01:46.622766Z","iopub.status.idle":"2025-05-29T13:01:46.646958Z","shell.execute_reply.started":"2025-05-29T13:01:46.622747Z","shell.execute_reply":"2025-05-29T13:01:46.646217Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import shutil;\nimport subprocess;\nimport sys,re,gzip,os;\nimport time;\nimport random;\nimport copy;\nimport tempfile;\nrandom.seed(123);\n\n\nnotebook_start = time.time();","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.648696Z","iopub.execute_input":"2025-05-29T13:01:46.648944Z","iopub.status.idle":"2025-05-29T13:01:46.662758Z","shell.execute_reply.started":"2025-05-29T13:01:46.648923Z","shell.execute_reply":"2025-05-29T13:01:46.662039Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_input_dir = \"data/\";\nkaggle_env = False;\nif os.path.exists(\"/kaggle/input/\"):\n    data_input_dir = \"/kaggle/input/\";\n    kaggle_env = True;\n\nsys.path.append(data_input_dir+\"/rna2025-scripts/\");\nimport smithwaterman","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.663552Z","iopub.execute_input":"2025-05-29T13:01:46.663769Z","iopub.status.idle":"2025-05-29T13:01:46.69607Z","shell.execute_reply.started":"2025-05-29T13:01:46.663749Z","shell.execute_reply":"2025-05-29T13:01:46.695232Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.system(\"cp \"+data_input_dir+\"/rna2025-scripts/sequence_alignment_rust/sequence_alignment/seqalign ./\");\nos.system(\"chmod 755 seqalign\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.69683Z","iopub.execute_input":"2025-05-29T13:01:46.697438Z","iopub.status.idle":"2025-05-29T13:01:46.744782Z","shell.execute_reply.started":"2025-05-29T13:01:46.697419Z","shell.execute_reply":"2025-05-29T13:01:46.744148Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def rust_seqalign(a,b,go=5.0,ge=2.0):\n    # penalty は negative で与えることになっている\n    with tempfile.TemporaryDirectory() as tmpdir:\n        tmpfas1 = os.path.join(tmpdir,\"infas1.fas\");\n        with open(tmpfas1,\"wt\") as fout:\n            fout.write(\">seq_a\\n\");\n            fout.write(a+\"\\n\");\n        tmpfas2 = os.path.join(tmpdir,\"infas2.fas\");\n        with open(tmpfas2,\"wt\") as fout:\n            fout.write(\">seq_b\\n\");\n            fout.write(b+\"\\n\");\n        res = subprocess.run([\"./seqalign\",\"--fas1\",tmpfas1,\"--fas2\",tmpfas2\n                              ,\"--scoring_matrix\",\"LETTER2-3\",\"--gap_open_penalty\",str(go*-1.0),\"--gap_extension_penalty\",str(ge*-1.0)\n                              ,],stdout=subprocess.PIPE,encoding=\"utf-8\");\n        lines = re.split(r\"[\\r\\n]\",res.stdout);\n        score = -1;\n        res_a = \"\";\n        res_b = \"\";\n        aflag = False;\n        bflag = False;\n        for ll in lines:\n            mat = re.search(\"#score:[\\s]*([^\\s]+)\",ll);\n            if mat:\n                score = float(mat.group(1));\n            if ll.startswith(\"#\"):\n                continue;\n            if \">seq_a\" in ll:\n                aflag = True;\n                bflag = False;\n                continue;\n            if \">seq_b\" in ll:\n                aflag = False;\n                bflag = True;\n                continue;\n            if \">\" in ll:\n                raise Exception(\"????\"+str(lines));\n            if aflag:\n                res_a += re.sub(r\"[\\s]\",\"\",ll)\n            if bflag:\n                res_b += re.sub(r\"[\\s]\",\"\",ll)\n        return [res_a,res_b,score];","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.745623Z","iopub.execute_input":"2025-05-29T13:01:46.746108Z","iopub.status.idle":"2025-05-29T13:01:46.755067Z","shell.execute_reply.started":"2025-05-29T13:01:46.746081Z","shell.execute_reply":"2025-05-29T13:01:46.754358Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df=pd.read_csv(data_input_dir+'/stanford-rna-3d-folding/test_sequences.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.755872Z","iopub.execute_input":"2025-05-29T13:01:46.756211Z","iopub.status.idle":"2025-05-29T13:01:46.788566Z","shell.execute_reply.started":"2025-05-29T13:01:46.756188Z","shell.execute_reply":"2025-05-29T13:01:46.788024Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"seq_noise=0.01;\nnum_trial=3;\nlength_max_boltz = 850;\nlength_max_chai1 = 450;\nretry_len_max = 300; # 短いのだけ何回かやる\n\nweighted_representation_score = False;\n\ndebug = False;\nskip_chai1 = False;\n\nboltz_res_dir = \"boltz_res\";\nchai1_res_dir = \"chai1_res\";\nblast_db_dir = \"blast_db\";\nmakeblastdb_path = \"ncbi-blast-2.16.0+/bin/makeblastdb\"\nblastn_path = \"ncbi-blast-2.16.0+/bin/blastn\"\nboltz_cache_dir = data_input_dir+\"/boltz-0511/boltz_data\";\n\npdb_template_dir = data_input_dir+\"/rnatbm-templates-20250521-re/fortbm_clustered_sampled_20250526_gt_re/\";\n\npdb_rep_template_dir = data_input_dir+\"/ribonanza-rep-pdb20250521-re/fortbm_clustered_20250526_re_rep_set/\";\n\nribonanzanet_dir = data_input_dir+\"/ribonanzanet-custom-set/\";\n\nmatrix_align_dir = data_input_dir+\"/matrix-align-20250525/\";\n\nchai_downloads_dir = data_input_dir+\"/chai1-set/chai1_weights\";\nboltz_runner = data_input_dir+\"/rna2025-runner/boltz_runner.py\";\nchai1_runner = data_input_dir+\"/rna2025-runner/chai1_runner.py\";\n\nsample_submission_path= data_input_dir+\"/stanford-rna-3d-folding/sample_submission.csv\"\n\nif not kaggle_env:\n    length_max_boltz = 500;\n    length_max_chai1 = 250;\n    seq_noise=0.05;\n    num_trial=2;\n    retry_len_max = 120; # 短いのだけ何回かやる\n    skip_chai1 = True;\n    \n    boltz_res_dir = \"nogit/boltz_res\";\n    chai1_res_dir = \"nogit/chai1_res\";\n    blast_db_dir = \"nogit/blast_db\";\n    if os.path.exists(boltz_res_dir):\n        shutil.rmtree(boltz_res_dir);\n    if os.path.exists(chai1_res_dir):\n        shutil.rmtree(chai1_res_dir);\n    if os.path.exists(blast_db_dir):\n        shutil.rmtree(blast_db_dir);\n\nif not os.path.exists(boltz_res_dir):\n    os.mkdir(boltz_res_dir);\nif not os.path.exists(chai1_res_dir):\n    os.mkdir(chai1_res_dir);\nif not os.path.exists(blast_db_dir):\n    os.mkdir(blast_db_dir);","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.78902Z","iopub.execute_input":"2025-05-29T13:01:46.789196Z","iopub.status.idle":"2025-05-29T13:01:46.80647Z","shell.execute_reply.started":"2025-05-29T13:01:46.789179Z","shell.execute_reply":"2025-05-29T13:01:46.80593Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# fileprefix が必要！\nfull_target_sequences = [\n    #{\"name\":\"3dix\",\"fileprefix\":\"seq1_\",\"desc\":\"\"\n    #,\"seq\":\"GGCCGACGGAGGCGCGCCCGAGAUGAGUAGGCUGUCCCAUCAGGGGAGGAAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCGCUCUGCAUGCCUGGGGGUAUGGGGAAUACCCAUACCACUGUCACGGAGGUCUCUCCGUGGAGAGCCGUCGGUX\"\n    #,\"all_sequences\":\">....\"}\n\n];\nname_to_length = {};\nfor ii in range(len(test_df)):\n    seq = {};\n    seq[\"name\"] = test_df.iloc[ii][\"target_id\"];\n    seq[\"seq\"] = test_df.iloc[ii][\"sequence\"];\n    name_to_length[seq[\"name\"]] = len(seq[\"seq\"])\n    seq[\"all_sequences\"] = [];\n    seq[\"fileprefix\"] = \"seq\"+str(ii)+\"_\";\n    seq[\"desc\"] = \"\";\n    try:\n        with tempfile.TemporaryDirectory() as tempdir:\n            tfilename = tempdir+\"/tfas.fas\";\n            fout = open(tfilename,\"wt\");\n            fout.write(test_df.iloc[ii][\"all_sequences\"]+\"\\n\");\n            fout.close();\n            seq[\"all_sequences\"] = loadFasta(tfilename);\n    except Exception as e:\n        sys.stderr.write(str(e));\n    # 複数配列に対応\n    seq[\"multiseq\"] = [];\n    for ss in list(seq[\"all_sequences\"]):\n        d = re.split(r\"\\|\",ss[\"name\"]+\" \"+ss[\"desc\"]);\n        cnum = 0;\n        for dd in list(d):\n            if dd.startswith(\"Chain\"):\n                dd = re.sub(r\"\\[[^\\]]+\\]\",\"\",dd);\n                cnum = len(re.split(r\",\",dd));\n        tag = \"rna\";\n        if ss[\"seq\"] == seq[\"seq\"]:\n            tag = \"rna\";\n        elif len(re.sub(r\"[AGCUX]\",\"\",ss[\"seq\"])) == 0:\n            tag = \"rna\";\n        elif len(re.sub(r\"[AGCTNX]\",\"\",ss[\"seq\"])) == 0:\n            tag = \"dna\";\n        else:\n            tag  =\"protein\";\n        seq[\"multiseq\"].append({\"type\":tag,\"seq\":ss[\"seq\"],\"num\":cnum});\n        \n    if not kaggle_env: #!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! チェック中\n        if seq[\"name\"] != \"R1117v2\":\n            continue;\n        if len(seq[\"seq\"]) > retry_len_max:\n            continue;\n    full_target_sequences.append(seq);\n    # print(len(seq[\"sequence\"]))\n    \nif debug:\n    length_max_boltz = 150;\n    length_max_chai1 = 150;\n    num_trial=4;\n    full_target_sequences = [\n    {\"name\":\"3dix\",\"fileprefix\":\"debug_seq1_\",\"desc\":\"\"\n    ,\"seq\":\"GGCCGACGGAGGCGCGCCCGAGAUGAGUAGGCUGUCCCAUCAGGGGAGGAAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCGCUCUGCAUGCCUGGGGGUAUGGGGAAUACCCAUACCACUGUCACGGAGGUCUCUCCGUGGAGAGCCGUCGGUX\"\n    ,\"all_sequences\":\">testseq\\nGGCCGACGGAGGCGCGCCCGAGAUGAGUAGGCUGUCCCAUCAGGGGAGGAAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCGCUCUGCAUGCCUGGGGGUAUGGGGAAUACCCAUACCACUGUCACGGAGGUCUCUCCGUGGAGAGCCGUCGGUX\\n\"}\n    ,{\"name\":\"3dix_dummy\",\"fileprefix\":\"debug_seq2_\",\"desc\":\"\"\n    ,\"seq\":\"GAGGCGGAUAC\"\n    ,\"all_sequences\":\">testseq\\nGGCCGACGGAGGCGCGCCCGAGAUGAGUAGGCUGUCCCAUCAGGGGAGGAAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCGCUCUGCAUGCCUGGGGGUAUGGGGAAUACCCAUACCACUGUCACGGAGGUCUCUCCGUGGAGAGCCGUCGGUX\\n\"}\n    # {\"name\":\"3dix_900\",\"fileprefix\":\"seq1_\",\"desc\":\"\"\n    # \",\"seq\":\"GGCCGACGGAGGCGCGCCCGAGAUGAGUAGGCUGUCCCAUCAGGGGAGGAAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCGGCCGACGGAGGCGCGCCCGAGAUGAGUAGGCUGUCCCAUCAGGGGAGGAAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCGGCCGACGGAGGCGCGCCCGAGAUGAGUAGGCUGUCCCAUCAGGGGAGGAAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCC\"\n    # ,\"all_sequences\":\">testseq\\nGGCCGACGGAGGCGCGCCCGAGAUGAGUAGGCUGUCCCAUCAGGGGAGGAAUCGGGGACGGCUGAAAGGCGAGGGCGCCGAAGGGUGCAGAGUUCCUCCCGCUCUGCAUGCCUGGGGGUAUGGGGAAUACCCAUACCACUGUCACGGAGGUCUCUCCGUGGAGAGCCGUCGGUX\\n\"}\n\n    ];","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.81235Z","iopub.execute_input":"2025-05-29T13:01:46.812556Z","iopub.status.idle":"2025-05-29T13:01:46.863276Z","shell.execute_reply.started":"2025-05-29T13:01:46.81254Z","shell.execute_reply":"2025-05-29T13:01:46.862658Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"spacer_seq = \"AAAAAAAAAAAAAAAAAAAA\";\ndef split_seq(seq,maxlen,shiftlen,terminallen):\n    if len(seq) <= maxlen:\n        return [copy.deepcopy(seq)];\n    st = 0;\n    ret = [];\n    while True:\n        en = st+maxlen;\n        if en > len(seq):\n            en = len(seq);\n            st = en-maxlen;\n        ret.append(\n            [copy.deepcopy(seq[st:en]),]\n        );\n        st += shiftlen;\n        if en == len(seq):\n            break;\n\n    # これ以外も行いたいが、時間制限があるので・・・\n    ret.append(\n        [seq[:terminallen],seq[-terminallen:]]\n    );\n    if len(seq) > maxlen*2:\n        ret.append([seq[terminallen:terminallen*2],seq[-terminallen*2:-terminallen]]);\n    return ret;\n    \n# 長いやつの処理\nall_target_fragments = [];\nfileprefix_to_target = {};\nfragmentcount = 0;\nfor ff in list(full_target_sequences):\n    pf = copy.deepcopy(ff);\n    all_target_fragments.append(pf);\n    fileprefix_to_target[ff[\"fileprefix\"]] = pf;\n    for tag,lengthmax in [(\"boltz\",length_max_boltz),(\"chai1\",length_max_chai1)]:\n        if len(ff[\"seq\"]) > lengthmax:\n            shift_fraction = 0.5;\n            if len(ff[\"seq\"]) > 2000:\n                shift_fraction = 0.8;\n                \n            pseq = split_seq(ff[\"seq\"],lengthmax,int(lengthmax*shift_fraction),lengthmax//2-10);\n            for p in list(pseq):\n                pp = {};\n                pp[\"name\"] = ff[\"name\"];\n                pp[\"desc\"] = \"\";\n                pp[\"fileprefix\"] = \"frag\"+str(fragmentcount)+\"_\";\n                pp[\"fragment\"] = tag;\n                \n                if len(p) == 2:\n                    pp[\"seq\"] = p[0]+spacer_seq+p[1];\n                    pp[\"fragment_length\"] = [len(p[0]),len(spacer_seq),len(p[1])];\n                else:\n                    assert len(p) == 1;\n                    pp[\"seq\"] = p[0];\n                    \n                pp[\"all_sequences\"] = \"\";\n                \n                fileprefix_to_target[pp[\"fileprefix\"]] = pp;\n                fragmentcount += 1;\n                all_target_fragments.append(pp);","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.866552Z","iopub.execute_input":"2025-05-29T13:01:46.866779Z","iopub.status.idle":"2025-05-29T13:01:46.887553Z","shell.execute_reply.started":"2025-05-29T13:01:46.866758Z","shell.execute_reply":"2025-05-29T13:01:46.88678Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"chk = {};\nngflag = False;\nfor aa in list(all_target_fragments):\n    if aa[\"fileprefix\"] in chk:\n        sys.stderr.write(\"Error in code\\n\");\n        sys.stderr.write(aa[\"fileprefix\"]+\"\\n\"+str(aa)+\"\\n\");\n        ngflag = True;\n    chk[aa[\"fileprefix\"]] = 100;\nif ngflag:\n    if debug:\n        sys.stderr.write(str(all_target_fragments));\n        raise Exception(\"????\");\n    all_target_fragments = list(c.values());","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.888408Z","iopub.execute_input":"2025-05-29T13:01:46.888617Z","iopub.status.idle":"2025-05-29T13:01:46.917182Z","shell.execute_reply.started":"2025-05-29T13:01:46.888601Z","shell.execute_reply":"2025-05-29T13:01:46.915665Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_res_map(a,b):\n    apos = 0;\n    bpos = 0;\n    ret = [];\n    for aa,bb in zip(a,b):\n        if aa != \"-\" and bb != \"-\":\n            ret.append((apos,bpos));\n        elif bb != \"-\":\n            ret.append((-1,bpos));\n        elif aa != \"-\":\n            ret.append((apos,-1));\n            \n        if aa != \"-\":\n            apos += 1;\n        if bb != \"-\":\n            bpos += 1;\n    return ret;\n    \ndef split_sw(aseq,bseq,go=5.0,ge=2.0): # aseq に full seq が入ることを想定している\n    max_map = [];\n    sys.stderr.write(\"Perform split align.\\n\")\n    for fraglen in [1000,2000,3000]:\n        stepper = len(bseq)//fraglen+1;\n        st = 0;\n        lastmatch = -1;\n        mapped = [];\n        for ii in range(stepper):\n            en = st+fraglen;\n            if ii == stepper-1 or len(bseq)-en < 200:\n                en = len(bseq);\n            cseq = bseq[st:en]\n            try:\n                res = rust_seqalign(aseq,cseq,go=go,ge=ge);\n            except Exception as e:\n                sys.stderr.write(str(e)+\"\\n\");\n                sw = smithwaterman.SmithWaterman(nucleotide=True);\n                res = sw.align(aseq,cseq);\n                del sw;\n            mapped_ = get_res_map(res[0],res[1]);\n            for jj in range(len(mapped_)):\n                if mapped_[jj][1] != -1:\n                    mapped_[jj] = (mapped_[jj][0],mapped_[jj][1]+st);\n            lastmatch_b = -1;\n            for jj in range(len(mapped_)):\n                if mapped_[jj][0] != -1 and mapped_[jj][1] != -1:                \n                    if lastmatch > -1 and mapped_[jj][0] <= lastmatch:\n                        pass;\n                    else:\n                        mapped.append(mapped_[jj]);\n                        lastmatch = mapped_[jj][0];\n                        lastmatch_b = mapped_[jj][1];\n            if en == len(bseq):\n                break;\n            if lastmatch_b > en - 100: #少しはみ出ている場合補正する\n                st = lastmatch_b+1;\n            else:\n                st = en;\n        if len(mapped) > len(max_map):\n            max_map = mapped;\n    a = [];\n    b = [];\n    alas = -1;\n    blas = -1;\n    for apos,bpos in list(max_map):\n        while alas < apos - 1:  # 修正\n            assert alas + 1 < len(aseq);\n            alas += 1\n            a.append(aseq[alas])\n            b.append(\"-\")\n        while blas < bpos - 1:  # 修正\n            assert blas + 1 < len(bseq);\n            blas += 1\n            a.append(\"-\")\n            b.append(bseq[blas])\n        assert blas == bpos -1 and alas == apos -1;\n        assert apos < len(aseq) and bpos < len(bseq);\n        a.append(aseq[apos]);\n        b.append(bseq[bpos]);\n        alas = apos;\n        blas = bpos;\n        \n    while alas < len(aseq)-1:\n        alas += 1;\n        a.append(aseq[alas]);\n        b.append(\"-\");\n    while blas < len(bseq)-1:\n        blas += 1;\n        a.append(\"-\");\n        b.append(bseq[blas]);\n    res = [];\n    res.append(\"\".join(a));\n    res.append(\"\".join(b));\n    res.append(len(max_map)*7.0-len(a)*2.0); # 正確ではないスコア\n    del max_map, a, b  # メモリ解放\n    return res;","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.918009Z","iopub.execute_input":"2025-05-29T13:01:46.918218Z","iopub.status.idle":"2025-05-29T13:01:46.93505Z","shell.execute_reply.started":"2025-05-29T13:01:46.9182Z","shell.execute_reply":"2025-05-29T13:01:46.934413Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def sequence_alignment(a,b,ge0=False):\n    if len(a) > 4000 and len(b) > 4000:\n        if ge0:\n            return split_sw(a,b,ge=0.0);\n        else:\n            return split_sw(a,b);\n    if len(a) < 200 and len(b) < 200:\n        if debug:\n            sys.stderr.write(\"Align with script.\\n\")\n        sw = smithwaterman.SmithWaterman(nucleotide=True);\n        if ge0:\n            sw.ge = 0.0;\n        return sw.align(a,b);\n    try:\n        if ge0:\n            return rust_seqalign(a,b,ge=0.0);\n        return rust_seqalign(a,b);\n    except Exception as e:\n        sys.stderr.write(str(e));\n        if ge0:\n            return split_sw(a,b,ge=0.0);\n        else:\n            return split_sw(a,b);\ntest = False;\nif test:\n    ngcount = 0;\n    for ii in range(10):\n        a = [];\n        b = [];\n        for jj in range(2000):\n            a.append(\"AGUC\"[random.randint(0,3)])\n        b = copy.deepcopy(a);\n        for jj in range(2):\n            p = random.randint(0,len(b)-50);\n            print(\"delete\",p,len(b))\n            for kk in range(80):\n                if p < len(b):\n                    b.pop(p);\n                else:\n                    break;\n        a,b = \"\".join(a),\"\".join(b);\n        print(\"split...\")\n        res = split_sw(a,b)\n        print(\"normal...\")\n        sw = smithwaterman.SmithWaterman(nucleotide=True);\n        res2 = sw.align(a,b);\n        if res[0] != res2[0] or res[1] != res2[1]:\n            ngcount += 1;\n            print(res[0])\n            print(res[1])\n            print(res2[0])\n            print(res2[1])\n            print(\"ng\")\n    print(ngcount)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.935929Z","iopub.execute_input":"2025-05-29T13:01:46.93619Z","iopub.status.idle":"2025-05-29T13:01:46.95396Z","shell.execute_reply.started":"2025-05-29T13:01:46.936168Z","shell.execute_reply":"2025-05-29T13:01:46.953231Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# print(all_target_fragments)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.954597Z","iopub.execute_input":"2025-05-29T13:01:46.954833Z","iopub.status.idle":"2025-05-29T13:01:46.967324Z","shell.execute_reply.started":"2025-05-29T13:01:46.954813Z","shell.execute_reply":"2025-05-29T13:01:46.966647Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.system(\"cp -rp \"+data_input_dir+\"/rnatbm2025-set/ncbi-blast-2.16.0+ ./\");\nos.system(\"chmod 755 \"+makeblastdb_path);\nos.system(\"chmod 755 \"+blastn_path);","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:46.96805Z","iopub.execute_input":"2025-05-29T13:01:46.968292Z","iopub.status.idle":"2025-05-29T13:01:57.834458Z","shell.execute_reply.started":"2025-05-29T13:01:46.968274Z","shell.execute_reply":"2025-05-29T13:01:57.833696Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.system(\"cp -rp \"+matrix_align_dir+\"/matrix_align ./\");\nos.system(\"chmod 755 matrix_align\");","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:57.835289Z","iopub.execute_input":"2025-05-29T13:01:57.835567Z","iopub.status.idle":"2025-05-29T13:01:57.911931Z","shell.execute_reply.started":"2025-05-29T13:01:57.835545Z","shell.execute_reply":"2025-05-29T13:01:57.911218Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"prevpath = \"\";\nif \"PYTHONPATH\" in os.environ:\n    prevpath = str(os.environ[\"PYTHONPATH\"]);\nrpath = ribonanzanet_dir+\"/RibonanzaNet_custom/RibonanzaNet\";\nrpath = re.sub(r\"/+\",\"/\",rpath);\nif rpath not in prevpath:\n    if len(prevpath) == 0:\n        os.environ[\"PYTHONPATH\"] = rpath;\n    else:\n        os.environ[\"PYTHONPATH\"] = rpath+\":\"+prevpath;","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:57.91267Z","iopub.execute_input":"2025-05-29T13:01:57.912902Z","iopub.status.idle":"2025-05-29T13:01:57.98439Z","shell.execute_reply.started":"2025-05-29T13:01:57.912875Z","shell.execute_reply":"2025-05-29T13:01:57.983676Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"qseqname_to_targetindex = {};\n# なんか外部プロセス？呼び出しが遅い\ntempfas = \"testfas.fas\";\nwith open(tempfas,\"wt\") as fout:\n    for eii,a in enumerate(full_target_sequences):\n        seqname = \"query\"+str(eii);\n        qseqname_to_targetindex[seqname] = eii;\n        fout.write(\">\"+seqname+\"\\n\")\n        fout.write(a[\"seq\"]+\"\\n\");\nfor subject_dir,tag in [\n    [pdb_template_dir,\"pdb\"],\n]:\n    # blast の tbm 用 DB を作成する\n    template_fas = os.path.join(blast_db_dir,\"search_templates.fas\");\n    seqname_to_file = {};\n    seqcount = 0;\n    with open(template_fas,\"wt\") as fout:\n        allfas_ = os.listdir(subject_dir);\n        for ff in list(sorted(allfas_)):\n            if ff.endswith(\"fas\") or ff.endswith(\"fasta\"):\n                pass;\n            else:\n                continue;\n            fasfile = os.path.join(subject_dir,ff);\n            fass = loadFasta(fasfile);\n            assert len(fass) == 1;\n            seqname = \"tbm_seq_\"+str(seqcount);\n            seqcount += 1;\n            seqname_to_file[seqname] = fasfile\n            fout.write(\">\"+seqname+\" file:\"+fasfile+\"\\n\");\n            fout.write(fass[0][\"seq\"]+\"\\n\");\n            \n    subprocess.run([makeblastdb_path,\"-in\",template_fas,\"-dbtype\",\"nucl\",\"-parse_seqids\"])\n    \n    proc = subprocess.run([blastn_path,\"-query\",tempfas,\"-word_size\",\"5\",\"-evalue\",\"0.1\",\"-db\",template_fas,\"-outfmt\",\"6\"],stdout=subprocess.PIPE,encoding=\"utf-8\");\n    lines = re.split(r\"[\\r\\n]+\",proc.stdout);\n    #print(proc.stdout);\n    #print(proc.stderr);\n    for ll in list(lines):\n        ptt = re.split(r\"[\\s]+\",re.sub(r\"[\\s]+$\",\"\",ll));\n        if len(ptt) > 9:\n            if int(ptt[8]) > int(ptt[9]): # reverse\n                sys.stderr.write(\"reverse complement?: \"+ll+\"\\n\")\n                sys.stderr.write(ptt[8]+\" \"+ptt[9]+\"\\n\")\n                continue;\n            idxx = qseqname_to_targetindex[ptt[0]];\n            fasfile = seqname_to_file[ptt[1]];\n            \n            btag = tag+\"_blasthit\";\n            if btag not in full_target_sequences[idxx]:\n                full_target_sequences[idxx][btag] = [];\n            full_target_sequences[idxx][btag].append((fasfile,float(ptt[-1])));\n                \n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:01:57.985295Z","iopub.execute_input":"2025-05-29T13:01:57.98559Z","iopub.status.idle":"2025-05-29T13:05:26.738767Z","shell.execute_reply.started":"2025-05-29T13:01:57.985573Z","shell.execute_reply":"2025-05-29T13:05:26.738059Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"del lines,proc,seqname_to_file;\nos.remove(tempfas)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:05:26.73955Z","iopub.execute_input":"2025-05-29T13:05:26.739848Z","iopub.status.idle":"2025-05-29T13:05:26.743842Z","shell.execute_reply.started":"2025-05-29T13:05:26.73981Z","shell.execute_reply":"2025-05-29T13:05:26.743139Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# representation alignment\nmatrix_align_binary=\"./matrix_align\";\n# matrix_align_template_list=\"templist_mat.dat\"\nribonanza_net_script = ribonanzanet_dir+\"/ribonanzanet_process.py\";\n\nribonanza_net_config = ribonanzanet_dir+\"/RibonanzaNet_custom/RibonanzaNet/configs/pairwise.yaml\";\nribonanza_net_model = ribonanzanet_dir+\"/archive/RibonanzaNet.pt\";\n\nuse_representation = False;\n\ndef get_alignment_score(lines,query_seq,weighted_representation_score):\n    score = {};\n    plen = {};\n    current_name = \"\";\n    namecount = 0;\n    for eii,ll in enumerate(list(lines)):\n        mat = re.search(r\">([^\\s]+)\",ll);\n        if mat:\n            current_name = mat.group(1);\n            namecount += 1;\n        mat = re.search(r\"score[\\s]*:[\\s]*([^\\s]+)\",ll);\n        if mat:\n            if current_name in score:\n                raise Exception(\"???\"+ll+\"\\n\"+\"\\n\".join(lines[max([0,eii-10]):min([len(lines),eii+10])]));\n            score[current_name] = float(mat.group(1));\n        mat = re.search(r\"profile_length[\\s]*:[\\s]*([^\\s]+)\",ll);\n        if mat:\n            plen[current_name] = float(mat.group(1));     \n\n    res = [];\n    assert len(score.keys()) == namecount;\n    assert len(plen.keys()) == namecount;\n\n    for kk in list(sorted(score.keys())):\n        if weighted_representation_score:\n            res.append((kk,score[kk]/max([plen[kk],len(query_seq)]))); # key が足りなかったらエラーになるはず\n        else:\n            res.append((kk,score[kk]));\n    return list(reversed(sorted(res,key=lambda x:x[1])));\n\ndef run_rep_align(query_seq,tmpdirname,matrix_align_template_list,seqname_to_file,mx_other_options):\n    \n    matrix_align_options = [\n        matrix_align_binary,\n        \"--in_list\",\n        matrix_align_template_list,\n        \"--a3m_pairwise\",\n        \"true\",\n        \"--alignment_type\",\n        \"local\",\n        \"--score_type\",\n        \"cos_sim\",\n        \"--num_threads\",\n        \"2\"\n    ];\n    matrix_align_options.extend(mx_other_options)\n    ribonanza_net_options=[\n        \"python\",\n        ribonanza_net_script,\n        \"--config_path\",\n        ribonanza_net_config,\n        \"--use_chunk\",\n        \"True\",\n        \"--split_len\",\n        \"512\",\n        \"--model_pt_path\",\n        ribonanza_net_model,\n        \"--max_len\",\n        \"10000\"\n    ];\n    tmpname = os.path.join(tmpdirname,\"tmp.\"+str(random.random())+\".fas\");\n    with open(tmpname,\"wt\") as fout:\n        fout.write(\">dummy_query\\n\");\n        fout.write(query_seq+\"\\n\");\n    tmpout = tmpname+\".mat.gz\";\n\n    comm = copy.deepcopy(ribonanza_net_options);\n    comm.append(\"--in_fasta\");\n    comm.append(tmpname);\n    comm.append(\"--out_dir\");\n    comm.append(tmpname+\"_res\");\n    tmpout = tmpname+\"_res/seq_0.mat.gz\";\n    \n    proc = subprocess.run(comm);\n    \n    alignres=tmpname+\".align.res\";\n    comm = copy.deepcopy(matrix_align_options);\n    comm.append(\"--in\");\n    comm.append(tmpout);\n    comm.append(\"--out\");\n    comm.append(alignres);\n    \n    proc = subprocess.run(comm,stdout=subprocess.PIPE,stderr=subprocess.PIPE,encoding=\"utf-8\");\n\n    # print(proc.stderr);\n    res = get_alignment_score(list(re.split(r\"[\\r\\n]+\",proc.stdout)),query_seq,weighted_representation_score);\n    name_to_score = {k:v for k,v in res};\n    hitseq = smithwaterman.SmithWaterman.loadFasta(alignres);\n    for hh in list(hitseq):\n        if hh[\"name\"] == \"dummy_query\":\n            continue;\n        hh[\"score\"] = name_to_score[hh[\"name\"]];\n        hh[\"name\"] = seqname_to_file[hh[\"name\"]];\n    return hitseq;\ndef get_first_line(file):\n    if file.endswith(\".gz\"):\n        fin = gzip.open(file,\"rt\");\n    else:\n        fin = open(file,\"rt\");\n    l = fin.readline();\n    fin.close();\n    return l;\nwith tempfile.TemporaryDirectory() as tmpdirname:\n    for ddir,hittag,alitag,mxop in [\n                (pdb_rep_template_dir,\"pdb_blasthit\",\"template_alignment\",[\"--gap_open_penalty\",\"-5.0\",\"--gap_penalty_auto_adjust\",\"false\"]),\n                (pdb_rep_template_dir,\"pdb_blasthit\",\"template_alignment_B\",[\"--gap_penalty_auto_adjust\",\"true\"]),\n                ]:\n        matlist = tmpdirname+\"/templist_mat.dat\";\n        seqname_to_file = {};\n        with open(matlist,\"wt\")as fout:\n            for root,dirs,files in list(sorted(list(os.walk(ddir)))):\n                for ff in files:\n                    if ff.endswith(\".mat.gz\") or ff.endswith(\".mat\"):\n                        matfile = os.path.join(root,ff);\n                        l = get_first_line(matfile);\n                        mat = re.search(r\">([^\\s]+)\",l);\n                        if mat:\n                            sname = mat.group(1);\n                            if sname not in seqname_to_file:\n                                seqname_to_file[sname] = matfile;                \n                            else:\n                                sys.stderr.write(\"{} was already used in {}. {} \\n\".format(sname,root+\"/\"+ff,seqname_to_file[sname]));\n                                continue;\n                            fout.write(matfile+\"\\n\");\n                        else:\n                            sys.stderr.write(matfile+\" does not have name section?\\n\");\n                            \n        for ff in list(full_target_sequences):\n            # 時間がかかるし、sw よりスコアが悪いことがあるのでヒットがないときのみ行う\n            #if hittag not in ff or len(ff[hittag]) == 0:\n            if True:\n                try:\n                    res = run_rep_align(ff[\"seq\"],tmpdirname,matlist,seqname_to_file,mxop);\n                    sorter = [];\n                    for aa in list(res):\n                        if \"score\" in aa:\n                            sorter.append((aa,aa[\"score\"]));\n                    sorter = list(sorted(sorter,key=lambda x:x[1],reverse=True));\n                    if len(sorter) > 10:\n                        ff[alitag] = [x[0] for x in sorter[0:10]];\n                    else:\n                        ff[alitag] = [x[0] for x in sorter];\n                    del sorter;\n                except Exception as e:\n                    tr = sys.exc_info()\n                    sys.stderr.write(\"\\n\".join(traceback.format_exception(*tr)));\n#for ff in list(full_target_sequences):\n#    if \"pdb_blast_hit\" not in ff:\n        \n#os.system(\"python \"+ribonanzanet_dir+\"/ribonanzanet_process.py\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:05:26.744692Z","iopub.execute_input":"2025-05-29T13:05:26.744993Z","iopub.status.idle":"2025-05-29T13:15:06.453262Z","shell.execute_reply.started":"2025-05-29T13:05:26.74497Z","shell.execute_reply":"2025-05-29T13:15:06.452651Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# print(full_target_sequences)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:15:06.454003Z","iopub.execute_input":"2025-05-29T13:15:06.454213Z","iopub.status.idle":"2025-05-29T13:15:06.457722Z","shell.execute_reply.started":"2025-05-29T13:15:06.454197Z","shell.execute_reply":"2025-05-29T13:15:06.457029Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if kaggle_env:\n    ccwd = os.getcwd();\n    os.chdir(data_input_dir+\"/boltz-0511/boltz_wheel/\");\n    os.system(\"pip install *\");\n    os.chdir(ccwd);","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:15:06.458343Z","iopub.execute_input":"2025-05-29T13:15:06.458579Z","iopub.status.idle":"2025-05-29T13:17:42.306651Z","shell.execute_reply.started":"2025-05-29T13:15:06.458561Z","shell.execute_reply":"2025-05-29T13:17:42.305828Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# スコアの低いものを優先的に予測する\nname_to_score = {};\nfor t in list(all_target_fragments):\n    name_to_score[t[\"name\"]] = 0.0;","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:17:42.307325Z","iopub.execute_input":"2025-05-29T13:17:42.3075Z","iopub.status.idle":"2025-05-29T13:17:42.311508Z","shell.execute_reply.started":"2025-05-29T13:17:42.307486Z","shell.execute_reply":"2025-05-29T13:17:42.310844Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc;\ngc.collect();","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:17:42.312227Z","iopub.execute_input":"2025-05-29T13:17:42.312484Z","iopub.status.idle":"2025-05-29T13:17:42.371135Z","shell.execute_reply.started":"2025-05-29T13:17:42.31246Z","shell.execute_reply":"2025-05-29T13:17:42.370551Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"precalc_score = {};","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:17:42.371953Z","iopub.execute_input":"2025-05-29T13:17:42.37223Z","iopub.status.idle":"2025-05-29T13:17:42.383569Z","shell.execute_reply.started":"2025-05-29T13:17:42.372215Z","shell.execute_reply":"2025-05-29T13:17:42.383048Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n### BOLTZ による予測 ##################\n\ncounter = 0; # 名前を一意にするためのカウンタ\nrna_linker = \"XXXXXXXXXXXXXXXXXXXXXXXXX\";\nfor jj in range(0,num_trial): # 時間がかかるので一回だけにしておく\n    print(\"try\",jj);\n    \n    target_sequences = [];\n    if jj == 0:\n        target_sequences = copy.deepcopy(all_target_fragments);\n    else:\n        for a in list(all_target_fragments):\n            if a[\"name\"] not in name_to_score:\n                name_to_score[a[\"name\"]] = 0;\n            if \"fragment\" in a[\"desc\"]:\n                target_sequences.append(a);\n            elif name_to_score[a[\"name\"]] < 90.0:\n                target_sequences.append(a);\n    if len(target_sequences) == 0:\n        break;\n    target_sequences = list(sorted(target_sequences,key=lambda x:name_to_score[x[\"name\"]]))\n    \n    for ii in range(len(target_sequences)):\n        if \"fragment\" in target_sequences[ii]:\n            if \"boltz\" not in target_sequences[ii][\"fragment\"]:\n                print(\"skipped\",target_sequences[ii][\"fragment\"])\n                continue;\n        out_prefix = \"test\"+str(ii)+\"_\"+str(jj);\n        outyaml = out_prefix+\".yaml\";\n        outdir = out_prefix+\"_out\";\n        noisedseq = \"\"\n        \n        if len(target_sequences[ii][\"seq\"]) > retry_len_max and jj != 0:\n            # print(\"retry skip: \"+target_sequences[ii][\"seq\"]);\n            continue;\n            \n        if \"fragment\" not in target_sequences[ii] and len(target_sequences[ii][\"seq\"]) > length_max_boltz:\n            # print(\"Too long sequence!\",target_sequences[ii][\"seq\"])\n            continue;\n        seq_orig = target_sequences[ii][\"seq\"];\n        if jj == 0:\n            sseq = target_sequences[ii][\"seq\"];\n        else:\n            sseq = [];\n            for si in range(len(target_sequences[ii][\"seq\"])):\n                if random.random() < seq_noise:\n                    #sseq.append(\"X\");\n                    sseq.append(\"AGUC\"[random.randint(0,3)]);\n                else:\n                    sseq.append(target_sequences[ii][\"seq\"][si]);\n            sseq = \"\".join(sseq);\n            \n        print(\"process\",target_sequences[ii][\"name\"],target_sequences[ii][\"fileprefix\"],len(sseq))\n        multiflag = False;\n        if jj > 0 and \"multiseq\" in target_sequences[ii]:\n            len_total = 0;\n            monomercount = 0;\n            for m in list(target_sequences[ii][\"multiseq\"]):\n                if m[\"type\"] == \"rna\" or m[\"type\"] == \"dna\":\n                    len_total += len(m[\"seq\"])*m[\"num\"]\n                if m[\"type\"] == \"rna\":\n                    monomercount += 1;\n            if jj == 1:\n                len_total += (monomercount-1)*len(rna_linker);\n                \n            if kaggle_env and len_total <= length_max_boltz:\n                multiflag = True;\n            \n            if not kaggle_env and len_total <= 400:\n                multiflag = True;\n            \n        if multiflag:\n            fout = open(outyaml,\"wt\");\n            fout.write(\"version: 1\\n\");\n            fout.write(\"sequences:\\n\");\n            chains = \"ABCDEFGHIJKLMNOPQRSTUVWXYZ\";\n            ctot = 0;\n            if jj == 2:\n                for m in list(target_sequences[ii][\"multiseq\"]):\n                    if m[\"type\"] == \"rna\" or m[\"type\"] == \"dna\":\n                        for ccou in range(m[\"num\"]):\n                            fout.write(\"  - \"+m[\"type\"]+\":\\n\");\n                            fout.write(\"      id: \"+\",\".join(chains[ctot])+\"\\n\");\n                            ctot += 1;\n                            fout.write(\"      sequence: \"+m[\"seq\"]+\"\\n\");\n            elif jj == 1:\n                rnas = []; # リンカーでつないだものを作る\n                for m in list(target_sequences[ii][\"multiseq\"]):\n                    if m[\"type\"] == \"rna\":                    \n                        for ccou in range(m[\"num\"]):\n                            rnas.append(m[\"seq\"]);\n                        continue;\n                    if m[\"type\"] == \"dna\":\n                        for ccou in range(m[\"num\"]):\n                            fout.write(\"  - \"+m[\"type\"]+\":\\n\");\n                            fout.write(\"      id: \"+chains[ctot]+\"\\n\");\n                            ctot += 1;\n                            fout.write(\"      sequence: \"+m[\"seq\"]+\"\\n\");\n                if len(rnas) > 0:\n                    concatseq = rna_linker.join(rnas);\n                    print(\"concatenated:\",rnas);\n                    fout.write(\"  - rna:\\n\");\n                    fout.write(\"      id: \"+chains[ctot]+\"\\n\");\n                    ctot += 1;\n                    fout.write(\"      sequence: \"+concatseq+\"\\n\");\n            else:\n                for m in list(target_sequences[ii][\"multiseq\"]):\n                    if m[\"type\"] == \"rna\" or m[\"type\"] == \"dna\":\n                        for ccou in range(m[\"num\"]):\n                            fout.write(\"  - \"+m[\"type\"]+\":\\n\");\n                            fout.write(\"      id: \"+\",\".join(chains[ctot])+\"\\n\");\n                            ctot += 1;\n                            if jj == 1 or target_sequences[ii][\"seq\"] != m[\"seq\"]:\n                                fout.write(\"      sequence: \"+m[\"seq\"]+\"\\n\");\n                            else:\n                                fout.write(\"      sequence: \"+sseq+\"\\n\");\n            fout.close();\n            print(\"process multi chains\",ctot)\n\n            \n        else:\n            fout = open(outyaml,\"wt\");\n            fout.write(\"version: 1\\n\");\n            fout.write(\"sequences:\\n\");\n            fout.write(\"  - rna:\\n\");\n            fout.write(\"      id: A\\n\");\n            fout.write(\"      sequence: {}\\n\".format(sseq));\n            fout.close();\n        \n        try:\n            #predict(outyaml,outdir,seed=123,accelerator=\"gpu\",cache=\"/kaggle/input/boltz-0511/boltz_data\")\n            subprocess.run(\n                [\"python\",boltz_runner,\"--in_yaml\",outyaml,\"--random_seed\",str(jj*10000+ii)\n                 ,\"--out_dir\",outdir,\"--cache_dir\",boltz_cache_dir]\n            );\n        except Exception as e:\n            sys.stderr.write(str(e));\n        os.remove(outyaml);\n        if os.path.exists(outdir):\n            for root, dirs, files in os.walk(outdir):\n                for ff in files:\n                    if ff.endswith(\"cif\"):\n                        fileprefix = target_sequences[ii][\"fileprefix\"]+\".\"+str(jj)+\".\"+str(counter)+\".boltz\";\n                        counter += 1;\n\n                        cif_chains = split_chain_mmcif(os.path.join(root,ff));\n\n                        maxseq = None;\n                        for c in range(len(cif_chains)):\n                            if cif_chains[c][\"seq\"] == target_sequences[ii][\"seq\"]:\n                                maxseq = cif_chains[c][\"seq\"];\n\n                        if maxseq is None:                        \n                            maxscore = 0.0;\n                            for c in range(len(cif_chains)):\n                                res = sequence_alignment(cif_chains[c][\"seq\"],target_sequences[ii][\"seq\"]);\n                                if res[2] > maxscore:\n                                    maxscore = res[2];\n                                    maxseq = cif_chains[c][\"seq\"];\n                                    \n                        if maxseq is None:\n                            continue;\n                            \n                        maxplddt = -1000.0;\n                        maxchain = None;\n                        for c in range(len(cif_chains)):\n                            if cif_chains[c][\"seq\"] == maxseq:\n                                if maxplddt < cif_chains[c][\"average_plddt\"]:\n                                    maxplddt = cif_chains[c][\"average_plddt\"];\n                                    maxchain = cif_chains[c];\n                        if len(target_sequences[ii][\"seq\"]) != len(maxchain[\"seq\"]):\n                            sres = sequence_alignment(maxchain[\"seq\"],target_sequences[ii][\"seq\"]);\n                            apos = 0;\n                            bpos = 0;\n                            chkpos = [];\n                            for aaa,bbb in zip(sres[0],sres[1]):\n                                if aaa != \"-\" and bbb != \"-\":\n                                    chkpos.append(apos+1);\n                                if aaa != \"-\":\n                                    apos += 1;\n                                if bbb != \"-\":\n                                    bpos += 1;\n                            pmax = maxplddt;\n                            maxplddt = calc_average_plddt(maxchain[\"atom_site\"],set(chkpos));\n                            print(pmax , \"->\", maxplddt );\n                        print(\"predicted\",len(cif_chains),\"chains\",target_sequences[ii][\"seq\"],maxchain[\"seq\"],maxplddt);\n                        if target_sequences[ii][\"name\"] not in name_to_score:\n                            name_to_score[target_sequences[ii][\"name\"]] = maxplddt;\n                        else:\n                            name_to_score[target_sequences[ii][\"name\"]] = max([name_to_score[target_sequences[ii][\"name\"]], maxplddt]);\n                        del cif_chains;\n                        rdict = copy.deepcopy(maxchain[\"atom_site\"]);\n                        rdict[\"data_\"] = \"X\";\n                        \n                        save_mmcif(rdict,os.path.join(boltz_res_dir,fileprefix)+\".cif\")\n                        \n                        # concat された配列についてはアライン領域のみのスコアを保存しておく\n                        if len(target_sequences[ii][\"seq\"]) != len(maxchain[\"seq\"]): \n                            precalc_score[os.path.join(boltz_res_dir,fileprefix)+\".cif\"] = maxplddt;\n                            \n                        with open(os.path.join(boltz_res_dir,fileprefix)+\".cif.fas\",\"wt\") as fout:\n                            if target_sequences[ii][\"seq\"] == maxchain[\"seq\"]:\n                                fout.write(\">\"+target_sequences[ii][\"name\"]+\" \"+target_sequences[ii][\"desc\"]+\"\\n\");\n                                fout.write(maxchain[\"seq\"]+\"\\n\");\n                            else:\n                                fout.write(\">\"+target_sequences[ii][\"name\"]+\" mutated:True \"+target_sequences[ii][\"desc\"]+\"\\n\");\n                                fout.write(maxchain[\"seq\"]+\"\\n\");\n            shutil.rmtree(outdir);\n        gc.collect();\n        en = time.time();\n        if en - notebook_start > 60*60*5.5:\n            sys.stderr.write(\"Some sequences were not processed but break considering the time limit...\");\n            break;\n    if en - notebook_start > 60*60*5.5: # 1 つも取れないと困るので boltz の方を長めに取る\n        break;","metadata":{"scrolled":true,"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:17:42.384477Z","iopub.execute_input":"2025-05-29T13:17:42.384672Z","iopub.status.idle":"2025-05-29T13:27:03.277946Z","shell.execute_reply.started":"2025-05-29T13:17:42.384658Z","shell.execute_reply":"2025-05-29T13:27:03.275237Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc;\ngc.collect();","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:27:06.296518Z","iopub.execute_input":"2025-05-29T13:27:06.296775Z","iopub.status.idle":"2025-05-29T13:27:06.350237Z","shell.execute_reply.started":"2025-05-29T13:27:06.296754Z","shell.execute_reply":"2025-05-29T13:27:06.34963Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if kaggle_env:\n    ccwd = os.getcwd();\n    os.chdir(data_input_dir+\"/chai1-set/chai1_wheel/\");\n    os.system(\"pip install *\");\n    os.chdir(ccwd);","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:27:06.466466Z","iopub.execute_input":"2025-05-29T13:27:06.466899Z","iopub.status.idle":"2025-05-29T13:28:13.079154Z","shell.execute_reply.started":"2025-05-29T13:27:06.466882Z","shell.execute_reply":"2025-05-29T13:28:13.078352Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc;\ngc.collect();","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:28:13.080289Z","iopub.execute_input":"2025-05-29T13:28:13.080568Z","iopub.status.idle":"2025-05-29T13:28:13.135075Z","shell.execute_reply.started":"2025-05-29T13:28:13.080551Z","shell.execute_reply":"2025-05-29T13:28:13.134444Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pathlib,os;\nos.environ[\"CHAI_DOWNLOADS_DIR\"] = chai_downloads_dir;","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:28:13.135979Z","iopub.execute_input":"2025-05-29T13:28:13.136192Z","iopub.status.idle":"2025-05-29T13:28:13.148234Z","shell.execute_reply.started":"2025-05-29T13:28:13.136176Z","shell.execute_reply":"2025-05-29T13:28:13.147472Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"counter = 0;\n# 時間がかかるので 2 回だけにしておく\nfor jj in range(2):\n    if skip_chai1:\n        continue;\n    print(\"try\",jj)\n    target_sequences = [];\n    if jj == 0:\n        target_sequences = copy.deepcopy(all_target_fragments);\n    else:\n        for a in list(all_target_fragments):\n            if a[\"name\"] not in name_to_score:\n                name_to_score[a[\"name\"]] = 0;\n            if \"fragment\" in a[\"desc\"]:\n                target_sequences.append(a);\n            elif name_to_score[a[\"name\"]] < 90.0:\n                target_sequences.append(a);\n    if len(target_sequences) == 0:\n        break;\n    target_sequences = list(sorted(target_sequences,key=lambda x:name_to_score[x[\"name\"]]))\n    rna_complement = {\n        \"A\":\"U\",\n        \"U\":\"A\",\n        \"G\":\"C\",\n        \"C\":\"G\",\n    };\n    for ii in range(len(target_sequences)):\n        if \"fragment\" in target_sequences[ii]:\n            if \"chai1\" not in target_sequences[ii][\"fragment\"]:\n                print(\"skipped\",target_sequences[ii][\"fragment\"])\n                continue;\n        out_prefix = \"test\"+str(ii)+\"_\"+str(jj);\n        outfas = out_prefix+\".fas\";\n        outdir = out_prefix+\"_out\";\n        \n        if len(target_sequences[ii][\"seq\"]) > retry_len_max and jj != 0:\n            # print(\"retry skip: \"+target_sequences[ii][\"seq\"])\n            continue;\n            \n        if \"fragment\" not in target_sequences[ii] and len(target_sequences[ii][\"seq\"]) > length_max_chai1:\n            # print(\"Too long sequence!\",target_sequences[ii][\"seq\"])\n            continue;\n        \n        if jj == 0:\n            # X が入っていると Ligand にされてしまうので変更する\n            sseq = re.sub(r\"[^AGUC]\",\"A\",target_sequences[ii][\"seq\"]);\n        else:\n            sseq = [];\n            for si in range(len(target_sequences[ii][\"seq\"])):\n                if random.random() < seq_noise:\n                    sseq.append(rna_complement.get(target_sequences[ii][\"seq\"][si],\"X\"));\n                else:\n                    sseq.append(target_sequences[ii][\"seq\"][si]);\n            sseq = \"\".join(sseq);\n            sseq = re.sub(r\"[^AGUC]\",\"A\",sseq);\n            \n        print(\"process\",target_sequences[ii][\"fileprefix\"],len(sseq))\n        \n        fout = open(outfas,\"wt\");\n        fout.write(\">rna|name=testseq1\\n\");\n        fout.write(sseq+\"\\n\");  \n        fout.close();\n        \n        try:\n            res = subprocess.run([\"python\",chai1_runner\n                ,\"--in_fas\",outfas,\"--random_seed\",str(jj*10000+ii)\n                ,\"--out_dir\",outdir,\"--num_trunk_recycles\",\"3\",\"--low_memory\",\"False\"])\n        except Exception as e:\n            sys.stderr.write(str(e));\n        os.remove(outfas);\n        if os.path.exists(outdir):\n            max_in = None;\n            max_plddt = -1.0;\n            for root, dirs, files in os.walk(outdir):\n                for ff in files:\n                    if ff.endswith(\"cif\") and \"pred\" in ff:\n                        if \"model_idx_0\" in ff: # AF2 は iter0 は良くない。chai1 の idx_0 が何を表しているのかよくわからないが self estimate で概ねベストではないようなのでスキップしておく\n                            continue;\n                        fileprefix = target_sequences[ii][\"fileprefix\"]+\".\"+str(jj)+\".\"+str(counter)+\".chai1.\";\n                        counter += 1;\n                        \n                        from Bio.PDB.MMCIF2Dict import MMCIF2Dict;\n                        cifdic = MMCIF2Dict(os.path.join(root,ff));\n                        plddt = 0.0;\n                        if \"_ma_qa_metric_local.metric_value\" in cifdic:\n                            for a in list(cifdic[\"_ma_qa_metric_local.metric_value\"]):\n                                plddt += float(a);\n                            plddt /= len(cifdic[\"_ma_qa_metric_local.metric_value\"]);\n                            if target_sequences[ii][\"name\"] not in name_to_score:\n                                name_to_score[target_sequences[ii][\"name\"]] = plddt;\n                            else:\n                                name_to_score[target_sequences[ii][\"name\"]] = max([name_to_score[target_sequences[ii][\"name\"]], plddt]);\n                        del cifdic;\n                        if plddt > max_plddt:\n                            max_in = os.path.join(root,ff);\n                            max_out = os.path.join(chai1_res_dir,fileprefix)+\".cif\";\n                            max_plddt = plddt;\n                        gc.collect();             \n            if max_in is not None:    \n                shutil.copy2(max_in,max_out);\n                with open(os.path.join(max_out+\".fas\"),\"wt\") as fout:\n                    if target_sequences[ii][\"seq\"] == sseq:\n                        fout.write(\">\"+target_sequences[ii][\"name\"]+\" \"+target_sequences[ii][\"desc\"]+\"\\n\");\n                        fout.write(sseq+\"\\n\");\n                    else:\n                        fout.write(\">\"+target_sequences[ii][\"name\"]+\" mutated:True \"+target_sequences[ii][\"desc\"]+\"\\n\");\n                        fout.write(sseq+\"\\n\");\n            shutil.rmtree(outdir);\n        en = time.time();\n        if en-notebook_start > 60*60*6.5:\n            sys.stderr.write(\"Some sequences were not processed but break considering the time limit...\");\n            break;\n    if en-notebook_start > 60*60*6.5:\n        break;","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:28:13.149755Z","iopub.execute_input":"2025-05-29T13:28:13.150375Z","iopub.status.idle":"2025-05-29T13:29:28.519954Z","shell.execute_reply.started":"2025-05-29T13:28:13.150349Z","shell.execute_reply":"2025-05-29T13:29:28.51861Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# cif を json に変換する\nimport re,os,sys;\nimport Bio.PDB;\nimport gzip,datetime;\nimport copy;\nimport numpy as np;\nimport json;\n\ndef zipper(mdict,keynames):\n    alen = len(mdict[keynames[0]]);\n    ret = [];\n    klen = len(keynames);\n    for kk in list(keynames):\n        ret.append([]);\n        if len(mdict[kk]) != alen:\n            sys.stderr.write(kk+\" length must be {} but {}.\\n\".format(alen,len(mdict[kk])));\n            raise Exception();\n    \n    for aa in range(alen):\n        for kk in range(klen):\n            ret[kk].append(mdict[keynames[kk]][aa]);\n    return ret;\n\n\ndef process_rna(filename):\n    if filename.endswith(\".gz\"):\n        fin = gzip.open(filename,\"rt\");\n    else:\n        fin = open(filename,\"rt\");        \n    mdict = Bio.PDB.MMCIF2Dict.MMCIF2Dict(fin);\n    fin.close();\n    \n    atom_site_keys = [\n        \"_atom_site.label_atom_id\", \n        \"_atom_site.label_alt_id\", \n        \"_atom_site.label_comp_id\", \n        \"_atom_site.label_asym_id\", \n        \"_atom_site.label_entity_id\", \n        \"_atom_site.label_seq_id\", \n        \"_atom_site.pdbx_PDB_ins_code\",\n        \"_atom_site.Cartn_x\", \n        \"_atom_site.Cartn_y\", \n        \"_atom_site.Cartn_z\", \n        \"_atom_site.auth_asym_id\", \n        \"_atom_site.pdbx_PDB_model_num\", \n    ];\n\n    idx_label_atom_id  = atom_site_keys.index(\"_atom_site.label_atom_id\");\n    idx_label_alt_id  = atom_site_keys.index(\"_atom_site.label_alt_id\");\n    idx_label_comp_id  = atom_site_keys.index(\"_atom_site.label_comp_id\");\n    idx_label_asym_id  = atom_site_keys.index(\"_atom_site.label_asym_id\");\n    idx_label_entity_id  = atom_site_keys.index(\"_atom_site.label_entity_id\");\n    idx_label_comp_id  = atom_site_keys.index(\"_atom_site.label_comp_id\");\n    idx_label_seq_id  = atom_site_keys.index(\"_atom_site.label_seq_id\");\n    idx_pdbx_PDB_ins_code  = atom_site_keys.index(\"_atom_site.pdbx_PDB_ins_code\");\n\n    idx_x  = atom_site_keys.index(\"_atom_site.Cartn_x\");\n    idx_y  = atom_site_keys.index(\"_atom_site.Cartn_y\");\n    idx_z  = atom_site_keys.index(\"_atom_site.Cartn_z\");\n\n    idx_auth_asym_id  = atom_site_keys.index(\"_atom_site.auth_asym_id\");\n    idx_pdbx_PDB_model_num  = atom_site_keys.index(\"_atom_site.pdbx_PDB_model_num\");\n\n    chains = {};\n    processed = {};\n    processed_atom = {};\n    for a in zip(*zipper(mdict,atom_site_keys)):\n        if a[idx_label_seq_id] == \".\":\n            continue;\n\n        chaincode = a[idx_pdbx_PDB_model_num]+\"#\"+a[idx_label_entity_id]+\"#\"+a[idx_label_asym_id]+\"#\"+a[idx_auth_asym_id];\n        \n        # \n        reslabel  = chaincode+\"#\"+a[idx_label_seq_id];\n\n        # if a[idx_label_alt_id] != \"A\" and a[idx_label_alt_id] != \".\" and a[idx_label_alt_id] != \"?\" and a[idx_label_alt_id] != \" \":\n        #     continue;\n        \n        if reslabel in processed and processed[reslabel] != a[idx_label_comp_id]+\"_\"+a[idx_pdbx_PDB_ins_code]:\n            sys.stderr.write(\"Duplicated label for \"+reslabel+\".\"+\" \".join(a)+\" was ignored.\\n\");\n            continue;\n        \n        # ALTPOS がある場合初出の ATOM を採択する\n        atomlabel = reslabel+\"#\"+a[idx_label_atom_id];\n        if atomlabel in processed_atom:\n            continue;\n        processed_atom[atomlabel] = 100;\n        \n        processed[reslabel] = a[idx_label_comp_id]+\"_\"+a[idx_pdbx_PDB_ins_code];\n\n        if chaincode not in chains:\n            chains[chaincode] = {};\n            chains[chaincode][\"residues\"] = [];\n            chains[chaincode][\"label_to_index\"] = {};\n        if reslabel not in chains[chaincode][\"label_to_index\"]:\n            chains[chaincode][\"label_to_index\"][reslabel] = len(chains[chaincode][\"residues\"]);\n            chains[chaincode][\"residues\"].append([]);\n        chains[chaincode][\"residues\"][chains[chaincode][\"label_to_index\"][reslabel]].append(a);\n    \n    ret = {};\n\n    # seq_id\n    # one_letter_code\n    # pos: [{\"X\"->[3]}]\n    \n    acc = {\n        \"A\":\"A\",\n        \"U\":\"U\",\n        \"G\":\"G\",\n        \"C\":\"C\"\n    };\n    for cc in list(chains.keys()):\n        chain_data = [];\n        for rr in list(chains[cc][\"residues\"]):\n            resname = rr[0][idx_label_comp_id];\n            one_letter_code = acc.get(resname,\"X\");\n\n            seq_id = int(rr[0][idx_label_seq_id]);\n            atoms = {};\n            processedatom = set();\n            for aa in list(rr):\n                if aa[idx_label_atom_id].startswith(\"H\"):\n                    continue;\n                p = [\n                    float(aa[idx_x]),\n                    float(aa[idx_y]),\n                    float(aa[idx_z])\n                ];\n\n                if aa[idx_label_atom_id] in atoms:\n                    sys.stderr.write(aa[idx_label_atom_id]+\" already found.\\n\");\n                    sys.stderr.write(filename+\"\\n\");\n                    sys.stderr.write(\" \".join(aa)+\"\\n\");\n                    raise Exception();\n                atoms[aa[idx_label_atom_id]] = p;\n\n            chain_data.append(\n                {\n                    \"seq_id\":seq_id,\n                    \"one_letter_code\":one_letter_code,\n                    \"atoms\":atoms,\n                }\n            );\n        ret[cc] = chain_data;\n    return ret;\n\ndef create_gt(fastafile,ciffile):\n    fass = loadFasta(fastafile);\n    if len(fass) != 1:\n        sys.stderr.write(fastafile+\" must have only one sequence.\"+\"\\n\");\n        raise Exception();\n    seq = fass[0][\"seq\"];\n    chains = process_rna(ciffile);\n    if len(chains) != 1:\n        print(chains)\n        sys.stderr.write(ciffile+\" must have only one chain.\"+\"\\n\");\n        raise Exception();\n\n    for cc in list(chains.items())[0][1]:\n        if len(seq) <=  cc[\"seq_id\"]-1:\n            print(fastafile,ciffile,seq,list(chains.items())[0][1])\n        if seq[cc[\"seq_id\"]-1] != cc[\"one_letter_code\"] and seq[cc[\"seq_id\"]-1] != \"X\" and cc[\"one_letter_code\"] != \"X\":\n            sys.stderr.write(\"sequence discrepancy {} {} {} {} {}\".format(fastafile,ciffile,cc[\"seq_id\"],cc[\"one_letter_code\"],seq[cc[\"seq_id\"]-1])+\"\\n\");\n            raise Exception();\n    return chains,fass;\n\ndef create_json_gt(fasta_file,cif_file,out_file):\n    chains,fass =create_gt(fasta_file,cif_file);\n    \n    if out_file.endswith(\".gz\") :\n        fout = gzip.open(out_file,\"wt\");\n    else:\n        fout = open(out_file,\"wt\");\n    json.dump(chains,fout,indent=2);\n    fout.close();\n    \n    with open(out_file+\".fas\",\"wt\") as fout:\n        for ff in fass:\n            fout.write(\">\"+ff[\"name\"]+\" \"+ff[\"desc\"]+\"\\n\");\n            fout.write(ff[\"seq\"]+\"\\n\");\n    \n    del chains,fass","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:29:39.911625Z","iopub.execute_input":"2025-05-29T13:29:39.912064Z","iopub.status.idle":"2025-05-29T13:29:39.930258Z","shell.execute_reply.started":"2025-05-29T13:29:39.91204Z","shell.execute_reply":"2025-05-29T13:29:39.929615Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# しきい値とかで原子をフィルタリングする\nfrom Bio.PDB.MMCIF2Dict import MMCIF2Dict;\nfrom Bio.PDB.mmcifio import MMCIFIO;\ncifio = MMCIFIO();\ndef filter_cif(cifname,outname,minlddt,insert_residues_1index=None):\n    if cifname.endswith(\".gz\"):\n        with gzip.open(cifname,\"rt\") as fin:\n            cifdic = MMCIF2Dict(fin);\n    else:\n        with open(cifname) as fin:\n            cifdic = MMCIF2Dict(fin);\n    c1_plddt = 0.0;\n    counter = 0;\n    retain_index = set();\n    if \"_atom_site.B_iso_or_equiv\" in cifdic:\n        for ii in range(len(cifdic[\"_atom_site.B_iso_or_equiv\"])):\n            if insert_residues_1index is not None:\n                if int(cifdic[\"_atom_site.label_seq_id\"][ii]) in insert_residues_1index:\n                    continue;\n            if float(cifdic[\"_atom_site.B_iso_or_equiv\"][ii]) >= minlddt:\n                retain_index.add(ii);\n                if cifdic[\"_atom_site.label_atom_id\"][ii] == \"C1'\":\n                    c1_plddt += float(cifdic[\"_atom_site.B_iso_or_equiv\"][ii]);\n                    counter += 1;\n    if counter == 0:\n        return 0.0;\n    c1_plddt /= counter;\n    outdic = {};\n    outdic[\"data_\"] = \"X\";\n    for kk in list(cifdic.keys()):\n        if kk.startswith(\"_atom_site.\"):\n            outdic[kk] = [];\n            for ii in range(len(cifdic[kk])):\n                if ii in retain_index:\n                    outdic[kk].append(cifdic[kk][ii]);\n                    \n    cifio.set_dict(outdic);\n    if outname.endswith(\".gz\"):\n        fout = gzip.open(outname,\"wt\");\n    else:\n        fout = open(outname,\"wt\");\n    cifio.save(fout);\n    fout.close();\n    del cifdic;\n    return c1_plddt;","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:29:39.986629Z","iopub.execute_input":"2025-05-29T13:29:39.986802Z","iopub.status.idle":"2025-05-29T13:29:39.99378Z","shell.execute_reply.started":"2025-05-29T13:29:39.986788Z","shell.execute_reply":"2025-05-29T13:29:39.993108Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"prefix_to_name = {};\nfor ff in list(all_target_fragments):\n    prefix_to_name[ff[\"fileprefix\"]] = ff[\"name\"];\n    \nname_to_cif = {};\ncif_to_target = {};\nfor ff in os.listdir(boltz_res_dir):\n    if ff.endswith(\"cif\") or ff.endswith(\"cif.gz\"):\n        ptt = re.split(r\"\\.\",ff);\n        prefix = ptt[0];\n        name = prefix_to_name[prefix];\n        \n        if name not in name_to_cif:\n            name_to_cif[name] = [];\n        name_to_cif[name].append(os.path.join(boltz_res_dir,ff));\n        cif_to_target[os.path.join(boltz_res_dir,ff)] = fileprefix_to_target[prefix];\nfor ff in os.listdir(chai1_res_dir):\n    if ff.endswith(\"cif\") or ff.endswith(\"cif.gz\"):\n        ptt = re.split(r\"\\.\",ff);\n        prefix = ptt[0];\n        name = prefix_to_name[prefix];\n        if name not in name_to_cif:\n            name_to_cif[name] = [];\n        name_to_cif[name].append(os.path.join(chai1_res_dir,ff));\n        cif_to_target[os.path.join(chai1_res_dir,ff)] = fileprefix_to_target[prefix];\n        \n        \n#create_gt();","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:29:40.146535Z","iopub.execute_input":"2025-05-29T13:29:40.146703Z","iopub.status.idle":"2025-05-29T13:29:40.152902Z","shell.execute_reply.started":"2025-05-29T13:29:40.14669Z","shell.execute_reply":"2025-05-29T13:29:40.152209Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# TBM\n# coordinates_dir=pdb_template_dir;\n\n# このへんまだ使っていない\nmatrix_align_binary=\"matrix_align\";\nmatrix_align_template_list=\"clustered_20250526_rep.list\"\n# os.environ[\"PYTHONPATH\"]=\"/home/ubuntu8/work/kaggle/rna/ribonanzanet/RibonanzaNet\";\nribonanza_net_script = \"ribonanzanet_process.py\";\nribonanza_net_config = \"dataprocess/ribonanzanet.config.yaml\";\nribonanza_net_model = \"ribonanzanet/archive/RibonanzaNet.pt\";\n\nuse_representation = False;\n\nmatrix_align_options = [\n    matrix_align_binary,\n    \"--in_list\",\n    matrix_align_template_list,\n    \"--a3m_pairwise\",\n    \"true\",\n    \"--alignment_type\",\n    \"local\",\n    \"--score_type\",\n    \"cos_sim\",\n    \"--num_threads\",\n    \"2\"\n];\n\nribonanza_net_options=[\n    \"python\",\n    ribonanza_net_script,\n    \"--config_path\",\n    ribonanza_net_config,\n    \"--use_chunk\",\n    \"True\",\n    \"--split_len\",\n    \"512\",\n    \"--model_pt_path\",\n    ribonanza_net_model,\n    \"--max_len\",\n    \"10000\"\n];\n\n# ===\n# RNA 用 TBM\n\nimport re,os,sys,gzip;\nimport numpy as np;\nimport smithwaterman;\nimport subprocess;\nimport random;\nimport copy;\nimport json;\nimport tempfile;\nfrom Bio.SVDSuperimposer import SVDSuperimposer;\nsw = smithwaterman.SmithWaterman(nucleotide=True);\n\ndef align_points(basepos,newpos,allpos):\n    sup = SVDSuperimposer();\n    sup.set(basepos, newpos);\n    sup.run();\n    rot, tran = sup.get_rotran();\n    return np.dot(allpos, rot) + tran;\n\ndef dist(a,b):\n    assert a.shape == b.shape;\n    assert a.shape[-1] == 3;\n\n    return np.sqrt(((a-b)*(a-b)).sum(axis=-1));\n\n# backbone_copy = ['P',\"OP1\",\"OP2\",\"OP3\",\"O5'\",\"O4'\", \"C5'\", \"C4'\", \"C3'\", \"C2'\", \"C1'\", \"O3'\", \"O2'\"]\n# OP3 があると openmm が失敗する\nbackbone_copy = ['P',\"OP1\",\"OP2\",\"O5'\",\"O4'\", \"C5'\", \"C4'\", \"C3'\", \"C2'\", \"C1'\", \"O3'\", \"O2'\"]\npurine_copy = [\"C4\",\"C5\",\"N3\",\"C2\",\"N1\",\"C6\",\"N7\",\"C8\",\"N9\"];\npyrim_copy = [\"C5\",\"C6\",\"N1\",\"C2\",\"N3\",\"C4\"];\n\ns5 = np.sin(54.0/360.0*np.pi*2.0)\ns6 = np.sin(60.0/360.0*np.pi*2.0)\ndef copy_atom_pos(\n    a,apos ,b\n    ):\n    ret = {};\n    for p in list(backbone_copy):\n        if p in apos:\n            ret[p] = apos[p];\n    if a in \"AG\" and b in \"AG\":\n        for p in list(purine_copy):\n            if p in apos:\n                ret[p] = apos[p];\n    elif a in \"CU\" and b in \"CU\":\n        for p in list(pyrim_copy):\n            if p in apos:\n                ret[p] = apos[p];\n    elif a in \"AG\" and b in \"CU\":\n        if \"N9\" in apos:\n            ret[\"N1\"] = apos[\"N9\"];\n        if \"C8\" in apos and \"C4\" in apos:\n            center = (apos[\"C8\"]+apos[\"C4\"])/2.0;\n            ret[\"C6\"] = (apos[\"C8\"]-center)/s5*s6+center;\n            ret[\"C2\"] = (apos[\"C4\"]-center)/s5*s6+center;\n    elif a in \"CU\" and b in \"AG\":\n        if \"N1\" in apos:\n            ret[\"N9\"] = apos[\"N1\"];\n        if \"C6\" in apos and \"C2\" in apos:\n            center = (apos[\"C6\"]+apos[\"C2\"])/2.0;\n            ret[\"C8\"] = (apos[\"C6\"]-center)/s6*s5+center;\n            ret[\"C4\"] = (apos[\"C2\"]-center)/s6*s5+center;\n    return ret;\n\n# 破壊的だし返り値でも返す\ndef move_to_center(pos):\n    allpos = np.zeros((3,));\n    counter = 0;\n    for rr in pos:\n        if rr is None:\n            continue;\n        for k,v in rr.items():\n            allpos += v;\n            counter += 1;\n    allpos /= counter;\n    for rr in pos:\n        if rr is None:\n            continue;\n        for k in list(rr.keys()):\n            rr[k] -= allpos;\n    return pos;\n\ndef relativemove_to(pos,direc):\n    for rr in pos:\n        if rr is None:\n            continue;\n        for k in list(rr.keys()):\n            rr[k] += direc;\n    return pos;\n    \ndef get_alignment_score(lines,query_seq,weighted_representation_score):\n    score = {};\n    plen = {};\n    current_name = \"\";\n    namecount = 0;\n    for eii,ll in enumerate(list(lines)):\n        mat = re.search(r\">([^\\s]+)\",ll);\n        if mat:\n            current_name = mat.group(1);\n            namecount += 1;\n        mat = re.search(r\"score[\\s]*:[\\s]*([^\\s]+)\",ll);\n        if mat:\n            if current_name in score:\n                raise Exception(\"???\"+ll+\"\\n\"+\"\\n\".join(lines[max([0,eii-10]):min([len(lines),eii+10])]));\n            score[current_name] = float(mat.group(1));\n        mat = re.search(r\"profile_length[\\s]*:[\\s]*([^\\s]+)\",ll);\n        if mat:\n            plen[current_name] = float(mat.group(1));     \n\n    res = [];\n    assert len(score.keys()) == namecount;\n    assert len(plen.keys()) == namecount;\n\n    for kk in list(sorted(score.keys())):\n        if weighted_representation_score:\n            res.append((kk,score[kk]/max([plen[kk],len(query_seq)]))); # key が足りなかったらエラーになるはず\n        else:\n            res.append((kk,score[kk]));\n    return list(reversed(sorted(res,key=lambda x:x[1])));\n\ndef get_rmsd(apos_,bpos_):\n    sup = SVDSuperimposer();\n    apos = [];\n    bpos = [];\n    \n    for a,b in zip(apos_,bpos_):\n        if a is not None and b is not None:\n            if \"C1'\" in a and \"C1'\" in b:\n                apos.append(a[\"C1'\"]);\n                bpos.append(b[\"C1'\"]);\n    if len(apos) < 7:\n        return 1000.0;\n    if type(apos[0]) == np.ndarray:\n        apos = np.stack(apos,axis=0);\n    else:\n        apos = np.array(apos);\n        \n    if type(bpos[0]) == np.ndarray:\n        bpos = np.stack(bpos,axis=0);\n    else:\n        bpos = np.array(bpos);\n        \n    sup.set(apos, bpos);\n    sup.run();\n    rms = sup.get_rms();\n    return rms;\n\ndef run_tbm(query_seq,fasta_list=None,similar_gt= None,prev_model=None,min_fragment_length=5,coordinates_dir=None,precomputed_alignment=None):\n    ret_similar = [];\n    with tempfile.TemporaryDirectory() as tmpdirname:\n        if use_representation:\n            tmpname = os.path.join(tmpdirname,\"tmp.\"+str(random.random())+\".fas\");\n            with open(tmpname,\"wt\") as fout:\n                fout.write(\">query\\n\");\n                fout.write(query_seq+\"\\n\");\n            tmpout = tmpname+\".mat.gz\";\n    \n            comm = copy.deepcopy(ribonanza_net_options);\n            comm.append(\"--in_fasta\");\n            comm.append(tmpname);\n            comm.append(\"--out_dir\");\n            comm.append(tmpname+\"_res\");\n            tmpout = tmpname+\"_res/seq_0.mat.gz\";\n            \n            proc = subprocess.run(comm);\n            \n            alignres=tmpname+\".align.res\";\n            comm = copy.deepcopy(matrix_align_options);\n            comm.append(\"--in\");\n            comm.append(tmpout);\n            comm.append(\"--out\");\n            comm.append(alignres);\n            \n            proc = subprocess.run(comm,stdout=subprocess.PIPE,encoding=\"utf-8\");\n    \n            res = get_alignment_score(list(re.split(r\"[\\r\\n]+\",proc.stdout)),query_seq,weighted_representation_score);\n            hitseq = smithwaterman.SmithWaterman.loadFasta(alignres);\n            for hh in list(hitseq):\n                mat = re.search(r\"file:([^\\s]+)\",hh[\"desc\"]);\n                if mat:\n                    hh[\"name\"] = mat.group(1);\n                else:\n                    fname = re.sub(coordinates_dir,\"\",hh[\"name\"]);\n                    fname = os.path.join(coordinates_dir,fname);\n                    hh[\"name\"] = fname;\n        else:\n            def fas_to_a3m(a,b):\n                ret = [];\n                for aa,bb in zip(a,b):\n                    if aa == \"-\":\n                        ret.append(bb.lower());\n                    else:\n                        ret.append(bb);\n                return \"\".join(ret);\n            if fasta_list is None:\n                fastafiles_ = os.listdir(coordinates_dir);\n                fastafiles = [(os.path.join(coordinates_dir,ff),-1.0) for ff in list(fastafiles_) ];\n            else:\n                fastafiles = [];\n                with open(fasta_list,\"rt\") as ffin:\n                    for ll in ffin:\n                        ll = re.sub(r\"[\\s]+$\",\"\",ll);\n                        if len(ll) > 0:\n                            ptt = re.split(r\"[\\s]+\",ll);\n                            if len(ptt) == 1:\n                                fastafiles.append((ptt[0],-1.0));\n                            else:\n                                fastafiles.append((ptt[0],float(ptt[1])));\n            \n            hitseq = [];\n            res = [];\n            for ff,sc in fastafiles:\n                if ff.endswith(\".mat\") or  ff.endswith(\".mat.gz\") or ff.endswith(\".fas\") or ff.endswith(\".fasta\"):\n                    mat = re.search(r\"/([^/]+\\.json)\",ff);\n                    if mat:\n                        tagname = mat.group(1);\n                        if precomputed_alignment is not None and tagname in precomputed_alignment:\n                            sss = copy.deepcopy(precomputed_alignment[tagname]);\n                            sss[\"name\"] = ff;\n                            hitseq.append(sss);\n                            if sc < 0.0:\n                                res.append((ff,sss[\"score\"]));\n                            else:\n                                res.append((ff,sc));\n                            continue;\n                    if ff.endswith(\".mat\") or  ff.endswith(\".mat.gz\"):\n                        sys.stderr.write(ff+\" has no precomputed data.\\n\");\n                        continue;\n                if ff.endswith(\".fas\") or ff.endswith(\".fasta\"):\n                    ffile = ff;\n                    sss = smithwaterman.SmithWaterman.loadFasta(ffile);\n                    assert len(sss) == 1;\n                    sss = sss[0];\n                    mapped_flag = False;\n                        \n                    if sss[\"seq\"] in query_seq:\n                        idx = query_seq.index(sss[\"seq\"]);\n                        b = [];\n                        for ii in range(len(query_seq)):\n                            if ii < idx:\n                                b.append(\"-\");\n                            else:\n                                if len(sss[\"seq\"]) > ii-idx:\n                                    if sss[\"seq\"][ii-idx] != query_seq[ii]:\n                                        break;\n                                    b.append(sss[\"seq\"][ii-idx]);\n                                else:\n                                    b.append(\"-\");\n                        if len(b) == len(query_seq):\n                            mapped_flag = True;\n                            sss[\"seq\"] = fas_to_a3m(query_seq,\"\".join(b));\n                            # print(\"mapped\",sss[\"seq\"]);\n                            \n                    if mapped_flag != True:\n                        ge0 = False;\n                        if \"frag\" in ff:\n                            ge0 = True;\n                        swres = sequence_alignment(query_seq,sss[\"seq\"],ge0=ge0);\n                        gc.collect();\n                        \n                        sss[\"seq\"] = fas_to_a3m(swres[0],swres[1]);\n                    sss[\"name\"] = ff;\n                    hitseq.append(sss);\n                    if sc < 0.0:\n                        res.append((ff,swres[2]));\n                    else:\n                        res.append((ff,sc));\n            res = list(reversed(sorted(res,key=lambda x:x[1])))\n    \n        name_to_aligned = {};\n    \n        for hh in list(hitseq):\n            name_to_aligned[hh[\"name\"]] = hh[\"seq\"];\n        del hitseq;\n        \n        aligned_pos = None;\n        aligned_mask = np.zeros((len(query_seq),),dtype=np.float32);\n        \n        unplaced = [];\n        unplaced_mask = np.zeros((len(query_seq),),dtype=np.float32);\n        for res_key,score in list(res):\n            if coordinates_dir is not None:\n                kk = re.sub(coordinates_dir,\"\",res_key);\n            else:\n                kk = res_key;\n                \n            pmat = re.search(r\"^(.+)/([^\\/]+\\.json(\\.gz)?)\\.mat/\",kk); # kaggle のファイル名変更への対応\n            if pmat:\n                fasfile = pmat.group(1)+\"/\"+pmat.group(2)+\".fas\";\n                name_to_aligned[fasfile] = name_to_aligned[res_key];\n                jsonfile = pmat.group(2);\n                jsonfile = re.sub(r\"\\.gz$\",\"\",jsonfile);\n                gtfile = pmat.group(1)+\"/\"+jsonfile+\"/\"+jsonfile;\n            else:\n                if os.path.exists(kk) and (kk.endswith(\".fas\") or kk.endswith(\".fasta\") or kk.endswith(\".mat\") or kk.endswith(\".mat.gz\")):\n                    fasfile = kk;\n                    gtfile = re.sub(r\"\\.fas(ta)?$\",\"\",kk);\n                    gtfile = re.sub(r\"\\.mat(\\.gz)?$\",\"\",gtfile);\n                elif kk.endswith(\".fas\") or kk.endswith(\".fasta\"):\n                    gtfile = os.path.join(coordinates_dir,re.sub(r\"\\.fas(ta)?$\",\"\",kk));\n                    fasfile = os.path.join(coordinates_dir,kk);\n                else:\n                    gtfile = os.path.join(coordinates_dir,kk+\".cif.gz.gt.json.gz\");\n                    fasfile = os.path.join(coordinates_dir,kk+\".cif.gz.gt.json.gz.fas\");\n                if not os.path.exists(gtfile):\n                    gtfile = re.sub(r\"\\.gz$\",\"\",gtfile); #kaggle が勝手に gz を展開してフォルダの下に入れる。。。\n                    gtfile = gtfile+\"/\"+os.path.basename(gtfile);\n            \n            if similar_gt is not None:\n                if gtfile in similar_gt:\n                    sys.stderr.write(\"{} is in similar gt.\\n\".format(gtfile));\n                    continue;\n            # sys.stderr.write(\"{} check..\\n\".format(gtfile));# !!!!!!!!!!!!!!!debug 中\n            if os.path.exists(gtfile):\n                if gtfile.endswith(\".gz\"):\n                    gtin = gzip.open(gtfile,\"rt\");\n                else:\n                    gtin = open(gtfile,\"rt\");\n                gt = json.load(gtin);\n                gtin.close();\n                pk = list(gt.keys());\n                assert len(pk) == 1;# 今の所一本にしか対応していない\n                gt = gt[pk[0]];\n                \n                scount = 0;\n                hitseq = \"\"\n                with open(fasfile,\"rt\") as ffin:\n                    for ll in ffin:\n                        if \">\" in ll:\n                            scount += 1;\n                            continue;\n                        hitseq += re.sub(r\"[\\s]\",\"\",ll);\n                        \n                pos = [None for _ in range(len(hitseq))];\n                mask = np.zeros((len(hitseq),),dtype=np.float32);\n                for pp in list(gt):\n                    idxx = pp[\"seq_id\"]-1;\n                    if pp[\"one_letter_code\"] not in \"NX\" and hitseq[idxx] not in \"NX\":\n                        assert pp[\"one_letter_code\"] == hitseq[idxx]\n                    pos[idxx] = {k:np.array(v) for k,v in pp[\"atoms\"].items()};\n                    mask[idxx] = 1.0;\n                del gt;\n                current_pos = [None for _ in range(len(query_seq))];\n                current_mask = np.zeros_like(aligned_mask);\n    \n                c1pos = 0;\n                \n                sseq = name_to_aligned[fasfile];\n                apos = 0;\n                bpos = 0;\n                skip_pos = set();\n                checked = set();\n                if min_fragment_length > 0:\n                    for ppp_ in range(len(sseq)):\n                        ppp = sseq[ppp_];\n                        if ppp_ in checked:\n                            continue;\n                        if ppp != \"-\" and ppp.isupper():\n                            ucount = 0;\n                            for qqq_ in range(ppp_,len(sseq)):\n                                if sseq[qqq_] == \"-\" or (not sseq[qqq_].isupper()):\n                                    break;\n                                ucount += 1;\n                                checked.add(qqq_);\n                            if ucount < min_fragment_length:\n                                for qqq_ in range(ppp_,len(sseq)):\n                                    if sseq[qqq_] == \"-\" or (not sseq[qqq_].isupper()):\n                                        break;\n                                    skip_pos.add(qqq_);\n                ngflag = False;\n                for ppp_ in range(len(sseq)):\n                    ppp = sseq[ppp_];\n                    if ppp_ not in skip_pos:\n                        if ppp != \"-\" and ppp.isupper():\n                            if mask[bpos] == 1.0:\n                                if hitseq[bpos] in \"XN\" or ppp in \"XN\" or hitseq[bpos] == ppp:\n                                    pass;\n                                else:\n                                    sys.stderr.write(\"Sequence discrepancy. {} \\n\".format(gtfile));\n                                    ngflag = True;\n                                chkp = copy_atom_pos(ppp,pos[bpos],query_seq[apos]);\n                                if \"C1'\" in chkp:\n                                    current_pos[apos] = chkp;\n                                    current_mask[apos] = 1.0;\n                    if ppp != \"-\":\n                        bpos += 1;\n                    if ppp == \"-\" or ppp.isupper():\n                        apos += 1;\n                        \n                if prev_model is not None:\n                    minrmsd = 100000.0;\n                    for prr in list(prev_model):\n                        r = get_rmsd(prr,current_pos);\n                        minrmsd = min(r,minrmsd);\n                        if minrmsd < 3.0:\n                            # sys.stderr.write(\"Skipped {} because it has rmsd {} with prev model.\\n\".format(gtfile,minrmsd));\n                            ret_similar.append(gtfile);\n                            ngflag = True;\n                \n                if current_mask.sum() == 0.0:\n                    # sys.stderr.write(\"Possibly aligned with missing region {}.\\n\".format(gtfile));\n                    continue;\n                if ngflag or current_mask.sum() == 0.0:\n                    continue;\n    \n                current_pos = move_to_center(current_pos);\n    \n                if aligned_mask.sum() == 0.0:\n                    aligned_mask = current_mask;\n                    unplaced_mask = np.minimum(current_mask+unplaced_mask,1.0);\n                    aligned_pos = current_pos;\n                    # sys.stderr.write(\"mapped \"+gtfile+\"\\n\");\n                    # sys.stderr.write(sseq+\"\\n\");\n                else:\n                    chk = unplaced_mask*current_mask\n                    if chk.sum() < current_mask.sum()*0.98: #ほぼカバーされている場合は保持しない\n                        unplaced_mask = np.minimum(current_mask+unplaced_mask,1.0);\n                        unplaced.append({\"pos\":current_pos,\"mask\":current_mask,\"score\":score});\n                        # sys.stderr.write(\"mapped \"+gtfile+\"\\n\");\n                        # sys.stderr.write(sseq+\"\\n\");\n                    else:\n                        pass;\n                        # sys.stderr.write(\"covered large region \"+gtfile+\"\\n\");\n                        # sys.stderr.write(sseq+\"\\n\");\n                        \n            else:\n                sys.stderr.write(gtfile+\" is not found.\\n\")\n        \n        fragment_move = 50; # intersection がなかったときに crash を避けるための移動量\n        while True:\n            plen = len(unplaced);\n            updated = False;\n            for uu in range(plen):\n                if unplaced[uu] is None:\n                    continue;\n                current_mask = unplaced[uu][\"mask\"];\n                current_pos = unplaced[uu][\"pos\"];\n    \n                intersection = current_mask*aligned_mask;\n                ssum = intersection.sum();\n                if ssum < 7:\n                    # 少ないと SVDSuperimposer が失敗する\n                    continue;\n                if ssum == current_mask.sum():\n                    unplaced[uu] = None;\n                    continue;\n                basepos = [];\n                tpos = [];\n                mapper = [];\n                for ii in range(len(query_seq)):\n                    if intersection[ii] != 0.0:\n                        if \"C1'\" in aligned_pos[ii] and \"C1'\" in current_pos[ii]:\n                            basepos.append(aligned_pos[ii][\"C1'\"]);\n                            tpos.append(current_pos[ii][\"C1'\"]);\n                            mapper.append(ii);\n                if len(basepos) == 0:\n                    unplaced[uu] = None;\n                    continue;\n                updated = True;\n                pointmap_key = [];\n                pointmap = [];\n                for rr in range(len(current_pos)):\n                    if current_pos[rr] is None:\n                        continue;\n                    for aa in list(current_pos[rr].keys()):\n                        pointmap_key.append((rr,aa));\n                        pointmap.append(current_pos[rr][aa]);\n    \n                moved_ = align_points(np.array(basepos),np.array(tpos),np.array(pointmap));\n                moved = copy.deepcopy(current_pos);\n                for ii in range(len(pointmap_key)):\n                    r,a = pointmap_key[ii];\n                    moved[r][a] = moved_[ii];\n                \n                for ii in range(len(query_seq)):\n                    if moved[ii] is None:\n                        continue;\n                    if aligned_mask[ii] != 0.0:\n                        for aa in list(moved[ii].keys()):\n                            if aa in aligned_pos[ii]:\n                                pass; # average しないほうがいいっぽい\n                                #if dist(aligned_pos[ii][aa],moved[ii][aa]) < 3.0: # 距離は適当\n                                #    aligned_pos[ii][aa] = aligned_pos[ii][aa]*0.5+moved[ii][aa]*0.5; # 割合は適当\n                            else:\n                                aligned_pos[ii][aa] = moved[ii][aa];\n                    else:\n                        aligned_pos[ii] = moved[ii];\n                        aligned_mask[ii] = 1.0;\n                unplaced[uu] = None;\n                break;\n            if not updated:\n                unplaced = list(reversed(unplaced));\n                for _ in range(plen):\n                    if len(unplaced) == 0: # ないはずだが一応\n                        break;\n                    uu = unplaced.pop();\n                    if uu is None:\n                        continue;\n                    current_mask = uu[\"mask\"];\n                    if fragment_move > 1000:\n                        uup = copy.deepcopy(uu[\"pos\"]);\n                        uup = relativemove_to(uup,np.array([fragment_move%1000,0,0],dtype=np.float32));\n                        current_pos = uup;\n                        fragment_move += 50;\n                    else:\n                        uup = copy.deepcopy(uu[\"pos\"]);\n                        uup = relativemove_to(uup,np.array([0,fragment_move,0],dtype=np.float32));\n                        current_pos = uup;\n                        fragment_move += 50;\n    \n                    for ii in range(len(query_seq)):\n                        if aligned_mask[ii] == 0.0 and current_mask[ii] != 0.0:\n                            aligned_pos[ii] = current_pos[ii];\n                            aligned_mask[ii] = 1.0;\n                    break;\n                unplaced = list(reversed(unplaced));\n                if len(unplaced) == 0:\n                    break;\n    return {\"pos\":aligned_pos,\"mask\":aligned_mask,\"seq\":query_seq},ret_similar\n    \n\"\"\"\n    import pandas as pd\n    ddict = {\"x_1\": [], \"y_1\": [],\"z_1\": [],\"resname\":[],\"resid\":[] }\n    for ii in range(len(query_seq)):\n        rid = ii+1;\n        if aligned_mask[ii] == 0.0:\n            continue;\n        ddict[\"x_1\"].append(aligned_pos[ii][\"C1'\"][0]);\n        ddict[\"y_1\"].append(aligned_pos[ii][\"C1'\"][1]);\n        ddict[\"z_1\"].append(aligned_pos[ii][\"C1'\"][2]);\n        ddict[\"resname\"].append(query_seq[ii]);\n        ddict[\"resid\"].append(str(rid));\n\n    dff = pd.DataFrame(ddict);\n    ribonanza_tm_score.write2pdb(dff,\"1\",\"nogit/testout.pdb\");\n\"\"\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:29:40.555258Z","iopub.execute_input":"2025-05-29T13:29:40.555523Z","iopub.status.idle":"2025-05-29T13:29:40.624622Z","shell.execute_reply.started":"2025-05-29T13:29:40.555505Z","shell.execute_reply":"2025-05-29T13:29:40.623976Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\nimport argparse\nimport gzip;\nfrom collections import OrderedDict\nimport tempfile;\nimport re;\n\ndef parse_json_to_mmcif_dict(json_file_path):\n    \"\"\"\n    Parse the JSON file and reconstruct the MMCIF atom_site section.\n    \n    Args:\n        json_file_path: Path to the JSON file containing parsed MMCIF data\n        \n    Returns:\n        dictionary compatible with BioPython's MMCIFIO.set_dict\n    \"\"\"\n    # Read the JSON file\n    if json_file_path.endswith(\"gz\"):\n        with gzip.open(json_file_path, 'r') as f:\n            data = json.load(f)\n    else:\n        with open(json_file_path, 'r') as f:\n            data = json.load(f)\n    \n    # Extract the key components\n    # Format: pdbx_PDB_model_num#label_entity_id#label_asym_id#auth_asym_id\n    ret = [];\n    for key, residues in data.items():\n        model_num, entity_id, asym_id, auth_asym_id = key.split('#')\n        \n        # Initialize the MMCIF dictionary\n        mmcif_dict = OrderedDict();\n        mmcif_dict[\"data_\"] = \"X\";\n        fastaseq = [];\n        # Create the atom_site category if it doesn't exist\n        if \"_atom_site.group_PDB\" not in mmcif_dict:\n            # Initialize the atom_site arrays\n            mmcif_dict[\"_atom_site.group_PDB\"] = []\n            mmcif_dict[\"_atom_site.id\"] = []\n            mmcif_dict[\"_atom_site.type_symbol\"] = []\n            mmcif_dict[\"_atom_site.label_atom_id\"] = []\n            mmcif_dict[\"_atom_site.label_alt_id\"] = []\n            mmcif_dict[\"_atom_site.label_comp_id\"] = []\n            mmcif_dict[\"_atom_site.label_asym_id\"] = []\n            mmcif_dict[\"_atom_site.label_entity_id\"] = []\n            mmcif_dict[\"_atom_site.label_seq_id\"] = []\n            mmcif_dict[\"_atom_site.pdbx_PDB_ins_code\"] = []\n            mmcif_dict[\"_atom_site.Cartn_x\"] = []\n            mmcif_dict[\"_atom_site.Cartn_y\"] = []\n            mmcif_dict[\"_atom_site.Cartn_z\"] = []\n            mmcif_dict[\"_atom_site.occupancy\"] = []\n            mmcif_dict[\"_atom_site.B_iso_or_equiv\"] = []\n            mmcif_dict[\"_atom_site.pdbx_formal_charge\"] = []\n            mmcif_dict[\"_atom_site.auth_seq_id\"] = []\n            mmcif_dict[\"_atom_site.auth_comp_id\"] = []\n            mmcif_dict[\"_atom_site.auth_asym_id\"] = []\n            mmcif_dict[\"_atom_site.auth_atom_id\"] = []\n            mmcif_dict[\"_atom_site.pdbx_PDB_model_num\"] = []\n        \n        # Process each residue in the list\n        atom_id_counter = 1\n        for residue in residues:\n            seq_id = residue[\"seq_id\"]\n            while len(fastaseq) < int(seq_id):\n                fastaseq.append(\"X\");\n            fastaseq[int(seq_id)-1] = residue[\"one_letter_code\"];\n            one_letter_code = residue[\"one_letter_code\"]\n            \n            # Process each atom in the residue\n            for atom_id, coords in residue[\"atoms\"].items():\n                # Determine element symbol from atom_id (first character or first two for some elements)\n                if atom_id.startswith(('C', 'N', 'O', 'P', 'S','H')):\n                    element = atom_id[0]\n                else:\n                    raise Exception(\"Not implemented \"+atom_id)\n                    # For atoms like 'Mg', 'Fe', etc.\n                    element = atom_id[:2]\n                \n                # Append values to the arrays\n                if one_letter_code != \"X\":\n                    mmcif_dict[\"_atom_site.group_PDB\"].append(\"ATOM\");\n                    assert one_letter_code in \"ATGCU\";\n                else:\n                    mmcif_dict[\"_atom_site.group_PDB\"].append(\"HETATM\");\n                mmcif_dict[\"_atom_site.id\"].append(str(atom_id_counter))\n                mmcif_dict[\"_atom_site.type_symbol\"].append(element)\n                mmcif_dict[\"_atom_site.label_atom_id\"].append(atom_id)\n                mmcif_dict[\"_atom_site.label_alt_id\"].append(\".\")  # Default value\n                mmcif_dict[\"_atom_site.label_comp_id\"].append(one_letter_code)\n                mmcif_dict[\"_atom_site.label_asym_id\"].append(asym_id)\n                mmcif_dict[\"_atom_site.label_entity_id\"].append(entity_id)\n                mmcif_dict[\"_atom_site.label_seq_id\"].append(str(seq_id))\n                mmcif_dict[\"_atom_site.pdbx_PDB_ins_code\"].append(\"?\")  # Default value\n                mmcif_dict[\"_atom_site.Cartn_x\"].append(f\"{coords[0]:.3f}\")\n                mmcif_dict[\"_atom_site.Cartn_y\"].append(f\"{coords[1]:.3f}\")\n                mmcif_dict[\"_atom_site.Cartn_z\"].append(f\"{coords[2]:.3f}\")\n                mmcif_dict[\"_atom_site.occupancy\"].append(\"1.00\")  # Default value\n                mmcif_dict[\"_atom_site.B_iso_or_equiv\"].append(\"50.00\")  # Default value\n                mmcif_dict[\"_atom_site.pdbx_formal_charge\"].append(\"?\")  # Default value\n                mmcif_dict[\"_atom_site.auth_seq_id\"].append(str(seq_id))  # Same as label_seq_id\n                mmcif_dict[\"_atom_site.auth_comp_id\"].append(one_letter_code)  # Same as label_comp_id\n                mmcif_dict[\"_atom_site.auth_asym_id\"].append(auth_asym_id)\n                mmcif_dict[\"_atom_site.auth_atom_id\"].append(atom_id)  # Same as label_atom_id\n                mmcif_dict[\"_atom_site.pdbx_PDB_model_num\"].append(model_num)\n                \n                atom_id_counter += 1\n        ret.append((mmcif_dict,\"\".join(fastaseq)));\n    del data;\n    return ret;\n\ndef convert_json_to_cif(json_file,output_prefix):\n    \n    # Parse JSON and reconstruct MMCIF dictionary\n    mmcifs = parse_json_to_mmcif_dict(json_file)\n    \n    # Save the MMCIF file\n    counter = 0;\n    ret = [];\n    for mmcif_dict,fasta_seq in mmcifs:\n        output_file = output_prefix+\".\"+str(counter)+\".cif\"\n        save_mmcif(mmcif_dict, output_file)\n        fout = open(output_file+\".fas\",\"wt\");\n        fout.write(\">seq.\"+str(counter)+\"\\n\");\n        fout.write(fasta_seq+\"\\n\");\n        fout.close();\n        ret.append((output_file,output_file+\".fas\"))\n    del mmcifs;\n    return ret;","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:29:40.625682Z","iopub.execute_input":"2025-05-29T13:29:40.62596Z","iopub.status.idle":"2025-05-29T13:29:40.638603Z","shell.execute_reply.started":"2025-05-29T13:29:40.625943Z","shell.execute_reply":"2025-05-29T13:29:40.637973Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from Bio.PDB.MMCIF2Dict import MMCIF2Dict;                      \nsw = smithwaterman.SmithWaterman(nucleotide=True);\n\n# OPENMM は末端を埋めてくれないので補正する\n# OPENMM を通らなかった場合については変な CIF になる\ndef get_mmcif_c1p(fullseq,ciffile):\n    cifdic = MMCIF2Dict(ciffile);\n    compid = cifdic[\"_atom_site.label_comp_id\"];\n    atomid = cifdic[\"_atom_site.label_atom_id\"];\n    car_x = cifdic[\"_atom_site.Cartn_x\"];\n    car_y = cifdic[\"_atom_site.Cartn_y\"];\n    car_z = cifdic[\"_atom_site.Cartn_z\"];\n    cifseq = [];\n    cifpos = [];\n    for c,a,x,y,z in zip(compid,atomid,car_x,car_y,car_z):\n        if a == \"C1'\":\n            if len(c) != 1:\n                sys.stderr.write(\"???\"+c+\" is not expected.\\n\");\n                continue;\n            else:\n                cifseq.append(c);\n                cifpos.append([float(x),float(y),float(z)]);\n                \n    cifseq = \"\".join(cifseq)\n    if fullseq == cifseq:\n        return cifpos;\n        \n    res = sequence_alignment(fullseq,cifseq);\n    \n    apos = 0;\n    bpos = 0;\n    ret = [];\n    for _ in range(len(fullseq)):\n        ret.append(\n            None\n        );\n    mismatchcount = 0;\n    placed = 0;\n    for a,b in zip(res[0],res[1]):\n        if a != \"-\" and b != \"-\":\n            if a != b:\n                mismatchcount += 1;\n            ret[apos] = cifpos[bpos];\n            placed += 1;\n        if a != \"-\":\n            apos += 1;\n        if b != \"-\":\n            bpos += 1;\n    if placed < 2:\n        return None;\n    placed = {};\n    for ii in range(len(fullseq)):\n        if ret[ii] is None:\n            a = None;\n            b = None;\n            # おそらく N 末\n            for jj in range(ii+1,len(fullseq)):\n                if ret[jj] is not None:\n                    for kk in range(jj+1,len(fullseq)):\n                        if ret[kk] is not None:\n                            a = jj;\n                            b = kk;\n                            break;\n                    break;\n            if a is None:\n                # おそらく C 末\n                for jj_ in range(len(fullseq)):\n                    jj = ii-1-jj_;\n                    if jj < 0:\n                        break;\n                    if ret[jj] is not None:\n                        for kk_ in range(jj_+1,len(fullseq)):\n                            kk = ii-1-kk_;\n                            if kk < 0:\n                                break;\n                            if ret[kk] is not None:\n                                a = jj;\n                                b = kk;\n                                break;\n                        break;\n            if a is None:\n                continue;\n            p = [];\n            for pp in range(3):\n                mov = (ret[a][pp]-ret[b][pp])/abs(a-b)*abs(ii-a);\n                p.append(mov+ret[a][pp]);\n            placed[ii] = p;\n            # print(a,ret[a],b,ret[b])\n            # print(ii,\"placed\",p)\n    for ii in list(placed.keys()):\n        ret[ii] = placed[ii];\n    del cifdic;\n    return ret;","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:29:40.891765Z","iopub.execute_input":"2025-05-29T13:29:40.892388Z","iopub.status.idle":"2025-05-29T13:29:40.902265Z","shell.execute_reply.started":"2025-05-29T13:29:40.892369Z","shell.execute_reply":"2025-05-29T13:29:40.901744Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"name_to_modelfile = {};\nt = 0;\ndel t;\nfor f_target in list(full_target_sequences):\n    model1_ = [];\n    model2_ = [];\n    model3_ = [];\n    model4_ = [];\n    model5_ = [];\n    if f_target[\"name\"] in name_to_cif:\n        for cc in list(name_to_cif[f_target[\"name\"]]):\n            try:\n                if \"filtered\" in cc:\n                    continue;\n                fragbias = 1.0;\n                fragment_index = None;\n                filter_min_plddt = 0.0;\n                if \"frag\" in cc:\n                    fragbias = 0.5;\n                    filter_min_plddt = 60.0;\n                    if cc in cif_to_target:\n                        tar = cif_to_target[cc];\n                        if \"fragment_length\" in tar:\n                            st = tar[\"fragment_length\"][0]+1; # 1 開始を与える\n                            en = tar[\"fragment_length\"][0]+tar[\"fragment_length\"][1]+1;\n                            fragment_index = set(range(st,en));\n                outname = re.sub(r\"\\.cif(\\.gz)$\",\"\",cc)+\".filtered.cif.gz\";\n                score = filter_cif(cc,outname,filter_min_plddt,fragment_index)*fragbias;\n                if score == 0.0:\n                    continue;\n                if cc in precalc_score:\n                    score = precalc_score[cc];\n                outjson = re.sub(r\"\\.cif(\\.gz)$\",\"\",cc)+\".filtered.json.gz\";\n                create_json_gt(cc+\".fas\",cc,outjson);\n                model1_.append([outjson+\".fas\",score]);\n                model2_.append([outjson+\".fas\",score]);\n                model3_.append([outjson+\".fas\",score]);\n                if \"boltz\" in cc:\n                    model4_.append([outjson+\".fas\",score]);\n                else:\n                    model5_.append([outjson+\".fas\",score]);\n            except Exception as e:\n                sys.stderr.write(str(e)+\"\\n\");\n    has_blasthit = False;\n    if \"pdb_blasthit\" in f_target and len(f_target[\"pdb_blasthit\"]) > 0:\n        for bii,b in enumerate(list(f_target[\"pdb_blasthit\"])):\n            if bii > 20:\n                break;\n            model1_.append([b[0],b[1]+100000.0]); # BLASTHIT は優先する\n            model2_.append([b[0],b[1]+100000.0]);\n            has_blasthit = True;\n            \n    if \"template_alignment\" in f_target:\n        for bii,b in enumerate(list(f_target[\"template_alignment\"])):\n            if bii > 10:\n                break;\n            if not has_blasthit:\n                model1_.append([b[\"name\"],b[\"score\"]/100.0+100.0]); # precomputed alignment が計算されてしまうので、hit がある場合は除外\n            model2_.append([b[\"name\"],b[\"score\"]/100.0+100.0]);\n    \n    if \"template_alignment_B\" in f_target:\n        for bii,b in enumerate(list(f_target[\"template_alignment_B\"])):\n            if bii > 10:\n                break;\n            model3_.append([b[\"name\"],b[\"score\"]/100.0+100.0]);\n    \n            \n    model1 = [p[0]+\"\\t\"+str(p[1]) for p in list(sorted(model1_,key=lambda x:x[1],reverse=True))]\n    model2 = [p[0]+\"\\t\"+str(p[1]) for p in list(sorted(model2_,key=lambda x:x[1],reverse=True))]\n    model3 = [p[0]+\"\\t\"+str(p[1]) for p in list(sorted(model3_,key=lambda x:x[1],reverse=True))]\n    model4 = [p[0]+\"\\t\"+str(p[1]) for p in list(sorted(model4_,key=lambda x:x[1],reverse=True))]\n    model5 = [p[0]+\"\\t\"+str(p[1]) for p in list(sorted(model5_,key=lambda x:x[1],reverse=True))]\n    \n    if len(model1) == 0:\n        continue;\n    if len(model2) == 0:\n        model2 = copy.deepcopy(model1);\n    if len(model3) == 0:\n        model3 = copy.deepcopy(model2);\n    if len(model4) == 0:\n        model4 = copy.deepcopy(model3);\n    if len(model5) == 0:\n        model5 = copy.deepcopy(model4);\n        \n    name_to_modelfile[f_target[\"name\"]] = [model1,model2,model3,model4,model5];\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:29:41.046731Z","iopub.execute_input":"2025-05-29T13:29:41.046916Z","iopub.status.idle":"2025-05-29T13:29:42.028071Z","shell.execute_reply.started":"2025-05-29T13:29:41.046901Z","shell.execute_reply":"2025-05-29T13:29:42.027536Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# name_to_modelfile[\"R1117v2\"][3]\n# full_target_sequences[0]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:29:42.029114Z","iopub.execute_input":"2025-05-29T13:29:42.029365Z","iopub.status.idle":"2025-05-29T13:29:42.032494Z","shell.execute_reply.started":"2025-05-29T13:29:42.029341Z","shell.execute_reply":"2025-05-29T13:29:42.031894Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nname_to_c1 = {};\nname_to_prevmodel = {};\nname_to_similar_gt = {};\nfor zii in range(5):\n    for f_target in list(full_target_sequences):\n        sys.stderr.write(\"Construct model {} {}\\n\".format(f_target[\"name\"],zii));\n        target_name = f_target[\"name\"]\n        if target_name not in name_to_c1:\n            name_to_c1[target_name] = [];\n        if f_target[\"name\"] not in name_to_prevmodel:\n            name_to_prevmodel[target_name] = [];\n            name_to_similar_gt[target_name] = set();\n        try:\n            if len(name_to_modelfile[target_name]) <= zii:\n                continue;\n            m = name_to_modelfile[target_name][zii];\n            has_mat_flag = False;\n            for mm in list(m):\n                if \".mat\" in mm:\n                    has_mat_flag = True;\n            tmpname = \"tmplist.list\";\n            fout = open(tmpname,\"wt\");\n            fout.write(\"\\n\".join(m));\n            fout.write(\"\\n\");\n            fout.close();\n            \n            def get_precomputed_alignment(pretag=\"template_alignment\"):\n                precomputed_alignment = {};\n                if pretag in f_target:\n                    for ali in f_target[pretag]:\n                        filename = ali[\"name\"];\n                        if \"assembly\" in filename:\n                            mat = re.search(r\"/([^/]+.json)\",filename);\n                            if mat:\n                                precomputed_alignment[mat.group(1)] = ali\n                return precomputed_alignment;\n            if has_mat_flag:\n                if zii != 2:\n                    precomputed_alignment = get_precomputed_alignment();\n                else:\n                    precomputed_alignment = get_precomputed_alignment(\"template_alignment_B\");\n            else:\n                precomputed_alignment = {};\n                \n            if len(precomputed_alignment) == 0:\n                sys.stderr.write(\"Precompued alignment was not passed.\\n\");\n            tbmres = run_tbm(f_target[\"seq\"],tmpname\n                            ,similar_gt=name_to_similar_gt[target_name],prev_model=name_to_prevmodel[target_name]\n                            ,min_fragment_length=0\n                             ,precomputed_alignment=precomputed_alignment);\n                            # ,min_fragment_length=0,precomputed_alignment=None);\n            if type(tbmres) == tuple:\n                tbmres,similar = tbmres;\n                for sss in list(similar):\n                    name_to_similar_gt[target_name].add(sss);\n            \n            pos = tbmres[\"pos\"];\n            if pos is None:\n                m = name_to_modelfile[target_name][0];\n                has_mat_flag = False;\n                for mm in list(m):\n                    if \".mat\" in mm:\n                        has_mat_flag = True;\n                if has_mat_flag:\n                    precomputed_alignment = get_precomputed_alignment();\n                else:\n                    precomputed_alignment = {};\n                tmpname = \"tmplist.list\";\n                fout = open(tmpname,\"wt\");\n                fout.write(\"\\n\".join(m));\n                fout.write(\"\\n\");\n                fout.close();\n                if len(precomputed_alignment) == 0:\n                    sys.stderr.write(\"Precompued alignment was not passed.\\n\");\n                tbmres,_ = run_tbm(f_target[\"seq\"],tmpname\n                                ,similar_gt=name_to_similar_gt[target_name]\n                                ,prev_model=name_to_prevmodel[target_name]\n                                ,precomputed_alignment=precomputed_alignment\n                                );            \n                pos = tbmres[\"pos\"];\n                if pos is None:\n                    tbmres,_ = run_tbm(f_target[\"seq\"],tmpname\n                        ,similar_gt=None,prev_model=None\n                        ,precomputed_alignment=precomputed_alignment\n                        );\n                    pos = tbmres[\"pos\"];\n            \n            name_to_prevmodel[target_name].append(pos);\n            \n            mask = tbmres[\"mask\"];\n            \n            assert len(pos) == len(f_target[\"seq\"]);\n            residuelist = [];\n            for i in range(len(pos)):\n                if mask[i] == 0.0:\n                    continue;\n                r = {};\n                r[\"seq_id\"] = i+1;\n                r[\"one_letter_code\"] = f_target[\"seq\"][i];\n                r[\"atoms\"] = {};\n                for k2 in list(pos[i].keys()):\n                    if i == 0 and \"P\" in k2:\n                        continue;\n                    if \"tolist\" in dir(pos[i][k2]):\n                        r[\"atoms\"][k2] = pos[i][k2].tolist();\n                residuelist.append(r);\n            with open('test.json', 'w') as f:\n               json.dump({\"1#1#A#A\":residuelist}, f)\n            convert_json_to_cif('test.json',\"testout\");\n            del residuelist,tbmres;\n            \n            tout = \"testout.0.cif\";\n            ores = \"openmmres.cif\";\n            \n            if os.path.exists(ores):\n                name_to_c1[f_target[\"name\"]].append(get_mmcif_c1p(f_target[\"seq\"],ores));\n            elif os.path.exists(tout):\n                c1p_ = get_mmcif_c1p(f_target[\"seq\"],tout);\n                if c1p_ is None:\n                    print(zii,f_target[\"name\"],\"???\");\n                    os.system(\"cat \"+tout);\n                name_to_c1[f_target[\"name\"]].append(c1p_);\n            if os.path.exists(ores):\n                os.remove(ores);    \n            if os.path.exists(tout):\n                os.remove(tout);\n            if 'pos' in locals():\n                del pos\n            gc.collect()\n        except Exception as e:\n            tr = sys.exc_info()\n            sys.stderr.write(\"\\n\".join(traceback.format_exception(*tr)))\n            \n        if zii > 0 and time.time() - notebook_start > 60*60*7.5:\n            break;\n    if zii > 0 and time.time() - notebook_start > 60*60*7.5:\n        break;\n","metadata":{"scrolled":true,"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:29:44.576542Z","iopub.execute_input":"2025-05-29T13:29:44.576771Z","iopub.status.idle":"2025-05-29T13:29:50.573705Z","shell.execute_reply.started":"2025-05-29T13:29:44.576755Z","shell.execute_reply":"2025-05-29T13:29:50.573013Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.path.exists(\"data//ribonanza-rep-pdb20250521-re/fortbm_clustered_20250526_re_rep_set/mod_8FZA-assembly1.1.1.A.A.cif.gz.gt.json\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:29:50.574822Z","iopub.execute_input":"2025-05-29T13:29:50.575055Z","iopub.status.idle":"2025-05-29T13:29:50.579898Z","shell.execute_reply.started":"2025-05-29T13:29:50.575033Z","shell.execute_reply":"2025-05-29T13:29:50.579159Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nfinal_result = open(\"submission.csv\",\"wt\");\nfinal_result.write(\",\".join([\"ID\",\"resname\",\"resid\",\"x_1\",\"y_1\",\"z_1\",\"x_2\",\"y_2\",\"z_2\",\"x_3\",\"y_3\",\"z_3\",\"x_4\",\"y_4\",\"z_4\",\"x_5\",\"y_5\",\"z_5\"])+\"\\n\");\nfor f_target in list(full_target_sequences):\n    if f_target[\"name\"] not in name_to_c1:\n        continue;\n    c1p_pos = name_to_c1[f_target[\"name\"]];\n    if len(c1p_pos) > 0:\n        try:\n            for ii in range(len(c1p_pos)):\n                if c1p_pos[ii] is None:\n                    c1p_pos[ii] = [];\n                if len(c1p_pos[ii]) != len(f_target[\"seq\"]):\n                    sys.stderr.write(\"???? \"+f_target[\"name\"]+\" \"+str(ii)+\"\\n\"); # 何か間違っているがないよりマシなので通す\n            for jj in range(len(f_target[\"seq\"])):\n                ptt = [];\n                ptt.append(\n                    f_target[\"name\"]+\"_\"+str(jj+1)\n                );\n                ptt.append(f_target[\"seq\"][jj]);\n                ptt.append(str(jj+1));\n                \n                for ii in range(5):\n                    if len(c1p_pos) <= ii:\n                        ptt.append(\"0.0\");\n                        ptt.append(\"0.0\");\n                        ptt.append(\"0.0\");\n                        continue;\n                    if c1p_pos[ii][jj] is None:\n                        ptt.append(\"0.0\");\n                        ptt.append(\"0.0\");\n                        ptt.append(\"0.0\");\n                    else:\n                        ptt.append(str(c1p_pos[ii][jj][0]));\n                        ptt.append(str(c1p_pos[ii][jj][1]));\n                        ptt.append(str(c1p_pos[ii][jj][2]));\n                final_result.write(\",\".join(ptt)+\"\\n\");\n        except Exception as e:\n            tr = sys.exc_info()\n            sys.stderr.write(\"\\n\".join(traceback.format_exception(*tr)));\nfinal_result.close();","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:29:50.580546Z","iopub.execute_input":"2025-05-29T13:29:50.580716Z","iopub.status.idle":"2025-05-29T13:29:50.595122Z","shell.execute_reply.started":"2025-05-29T13:29:50.580703Z","shell.execute_reply":"2025-05-29T13:29:50.594523Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"display(pd.read_csv('submission.csv'))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:30:06.517733Z","iopub.execute_input":"2025-05-29T13:30:06.517984Z","iopub.status.idle":"2025-05-29T13:30:06.550869Z","shell.execute_reply.started":"2025-05-29T13:30:06.517965Z","shell.execute_reply":"2025-05-29T13:30:06.55036Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 何がしかのエラーがあった場合に対応する\ntry:\n    fin  = open(\"submission.csv\",\"rt\");\n    alllines = fin.readlines();\n    filtered = [];\n    has_data = set();\n    head = alllines.pop(0);\n    has_data.add(\"ID\")\n    filtered.append(re.sub(r\"[\\s]+$\",\"\",head));\n    for ll in alllines:\n        ll = re.sub(r\"[\\s]+$\",\"\",ll);\n        ptt = re.split(r\",\",ll);\n        if len(ptt) < 3:\n            continue;\n        while len(ptt) < 18:\n            ptt.append(\"0.0\");\n        filtered.append(\",\".join(ptt));\n        has_data.add(ptt[0]);\n    fin.close();\n    updated = False;\n    fin = open(sample_submission_path,\"rt\");\n    samlines = fin.readlines();\n    samlines.pop(0);\n    for ll in samlines:\n        ll = re.sub(r\"[\\s]+$\",\"\",ll);\n        ptt = re.split(r\",\",ll);\n        if len(ptt) != 18:\n            sys.stderr.write(ll+\"???\\n\");\n            continue;\n        if ptt[0] not in has_data:\n            updated = True;\n            filtered.append(ll);\n    fin.close();\n    if updated:\n        print(\"update!\")\n        fout = open(\"submission.csv_\",\"wt\");\n        fout.write(\"\\n\".join(filtered));\n        fout.write(\"\\n\");\n        os.system(\"mv submission.csv_ submission.csv\");\nexcept Exception as e:\n    tr = sys.exc_info()\n    sys.stderr.write(\"\\n\".join(traceback.format_exception(*tr)));","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:30:07.647417Z","iopub.execute_input":"2025-05-29T13:30:07.648017Z","iopub.status.idle":"2025-05-29T13:30:07.680521Z","shell.execute_reply.started":"2025-05-29T13:30:07.647997Z","shell.execute_reply":"2025-05-29T13:30:07.679836Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.system(\"ls\");","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:30:08.496759Z","iopub.execute_input":"2025-05-29T13:30:08.497434Z","iopub.status.idle":"2025-05-29T13:30:08.50661Z","shell.execute_reply.started":"2025-05-29T13:30:08.497409Z","shell.execute_reply":"2025-05-29T13:30:08.505926Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not debug and kaggle_env:\n    import shutil;\n    for f in [\"ncbi-blast-2.16.0+/\",\"blast_db/\",\"chai1_res/\",\"boltz_res/\",\"testout.0.cif.fas\",\"test.json\",\"tmplist.list\"]:\n        try:\n            if os.path.exists(f):\n                if os.path.isdir(f):\n                    shutil.rmtree(f);\n                else:\n                    os.remove(f);\n        except Exception as e:\n            sys.stderr.write(str(e)+\"\\n\");","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:30:11.642005Z","iopub.execute_input":"2025-05-29T13:30:11.64228Z","iopub.status.idle":"2025-05-29T13:30:11.647162Z","shell.execute_reply.started":"2025-05-29T13:30:11.642261Z","shell.execute_reply":"2025-05-29T13:30:11.646558Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.system(\"ls\");","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:30:12.036486Z","iopub.execute_input":"2025-05-29T13:30:12.0367Z","iopub.status.idle":"2025-05-29T13:30:12.043119Z","shell.execute_reply.started":"2025-05-29T13:30:12.036683Z","shell.execute_reply":"2025-05-29T13:30:12.042475Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not kaggle_env:\n    import kaggle_snippet;\n    val_df = pd.read_csv(data_input_dir+\"/stanford-rna-3d-folding/validation_labels.csv\");\n    sub_df = pd.read_csv(\"submission.csv\");\n    def get_structure_from_df(target_df):\n        id_to_structure = {};\n        for ii in range(len(target_df)):\n            idd = re.split(\"_\",target_df.iloc[ii][\"ID\"])[0];\n            resname = target_df.iloc[ii][\"resname\"];\n            resid = target_df.iloc[ii][\"resid\"];\n            if idd not in id_to_structure:\n                id_to_structure[idd] = {};\n            for p in range(1,6):\n                x = target_df.iloc[ii][\"x_\"+str(p)];\n                y = target_df.iloc[ii][\"y_\"+str(p)];\n                z = target_df.iloc[ii][\"z_\"+str(p)];\n                if x == y and y == z:\n                    continue;\n                if p not in id_to_structure[idd]:\n                    id_to_structure[idd][p] = [];\n                att = kaggle_snippet.PDBAtom();\n                att.head = \"ATOM  \";\n                att.serial_number = len(id_to_structure[idd][p])+1;\n                att.atom_name = \"C1'\";\n                att.alt_loc = \" \";\n                att.residue_name = resname;\n                att.chain_id = \"A\";\n                att.residue_pos = int(resid);\n                att.insertion_code = \" \";\n                # print(x,y,z);\n                att.x = float(x);\n                att.y = float(y);\n                att.z = float(z);\n                att.occupancy = \"1.00\";\n                att.bfactor = \"0.00\";\n                att.element = \"C\";\n                att.charge = \"  \";\n                try:\n                    id_to_structure[idd][p].append(att.make_line());\n                except Exception as e:\n                    tr = sys.exc_info()\n                    sys.stderr.write(\"\\n\".join(traceback.format_exception(*tr)))\n        return id_to_structure;\n    val_structure = get_structure_from_df(val_df);\n    sub_structure = get_structure_from_df(sub_df);","metadata":{"scrolled":true,"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:30:12.557404Z","iopub.execute_input":"2025-05-29T13:30:12.558015Z","iopub.status.idle":"2025-05-29T13:30:12.565687Z","shell.execute_reply.started":"2025-05-29T13:30:12.557994Z","shell.execute_reply":"2025-05-29T13:30:12.564975Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not kaggle_env:\n    for kk in list(val_structure.keys()):\n        if kk not in sub_structure:\n            sys.stderr.write(kk+\" ????\\n\");\n            continue;\n        with tempfile.TemporaryDirectory() as tempdir:\n            vout = open(tempdir+\"/val.pdb\",\"wt\");\n            vout.write(\"\\n\".join(val_structure[kk][1])+\"\\n\");\n            vout.close();\n            results = [kk,str(name_to_length[kk])];\n            for ii in range(1,6):\n                if ii not in sub_structure[kk]:\n                    sys.stderr.write(kk+\" \"+str(ii)+\" was not found.\\n\");\n                    continue;\n                sout = open(tempdir+\"/sub.pdb\",\"wt\");\n                sout.write(\"\\n\".join(sub_structure[kk][ii])+\"\\n\");\n                sout.close();\n                proc = subprocess.run([\n                    \"/home/ubuntu8/apps/USalign/USalign\"\n                    ,tempdir+\"/val.pdb\",tempdir+\"/sub.pdb\"\n                    ,\"-outfmt\",\"2\",\"-TMscore\",\"1\",\"-atom\",\" C1'\"       \n                ],stdout=subprocess.PIPE,encoding=\"utf-8\");\n                lines = re.split(r\"[\\r\\n]+\",proc.stdout);\n                ptt = re.split(r\"[\\s]+\",lines[1]);\n                results.append(ptt[2]);\n            print(\"\\t\".join(results))\n            ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:27:03.310078Z","iopub.status.idle":"2025-05-29T13:27:03.310411Z","shell.execute_reply.started":"2025-05-29T13:27:03.310228Z","shell.execute_reply":"2025-05-29T13:27:03.310241Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not kaggle_env:\n    #print(full_target_sequences[0])\n    print(name_to_modelfile[\"R1117v2\"][2])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-29T13:27:03.312353Z","iopub.status.idle":"2025-05-29T13:27:03.312653Z","shell.execute_reply.started":"2025-05-29T13:27:03.312496Z","shell.execute_reply":"2025-05-29T13:27:03.312509Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}