{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":118765,"databundleVersionId":15231210,"sourceType":"competition"}],"dockerImageVersionId":31259,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os, random, time, warnings\nwarnings.filterwarnings(\"ignore\")\n\nDATA = \"/kaggle/input/stanford-rna-3d-folding-2/\"\n\ntrain_seq = pd.read_csv(DATA+\"train_sequences.csv\")\ntest_seq  = pd.read_csv(DATA+\"test_sequences.csv\")\ntrain_lab = pd.read_csv(DATA+\"train_labels.csv\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ndef build_coord_dict(df):\n    d={}\n    prefixes = df.ID.str.rsplit(\"_\",n=1).str[0]\n    for tid,g in df.groupby(prefixes):\n        d[tid]=g.sort_values(\"resid\")[[\"x_1\",\"y_1\",\"z_1\"]].values\n    return d\n\ntrain_coords = build_coord_dict(train_lab)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ndef align(q,t,match=2,mismatch=-1,gap=-2):\n\n    L1,L2=len(q),len(t)\n    dp=np.zeros((L1+1,L2+1))\n    ptr=np.zeros((L1+1,L2+1),int)\n\n    for i in range(1,L1+1): dp[i,0]=i*gap\n    for j in range(1,L2+1): dp[0,j]=j*gap\n\n    for i in range(1,L1+1):\n        for j in range(1,L2+1):\n\n            s = match if q[i-1]==t[j-1] else mismatch\n\n            vals=[\n                dp[i-1,j-1]+s,\n                dp[i-1,j]+gap,\n                dp[i,j-1]+gap\n            ]\n            best=np.argmax(vals)\n            dp[i,j]=vals[best]\n            ptr[i,j]=best\n\n    i,j=L1,L2\n    pairs=[]\n    while i>0 and j>0:\n        p=ptr[i,j]\n        if p==0:\n            pairs.append((i-1,j-1))\n            i-=1;j-=1\n        elif p==1:\n            i-=1\n        else:\n            j-=1\n\n    return pairs[::-1], dp[L1,L2]\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_records=list(train_seq.itertuples())\n\n# -------------------------------------------------\n# FAST PREFILTER USING KMER SIMILARITY\n# -------------------------------------------------\n\ndef kmer_score(a,b,k=3):\n    A={a[i:i+k] for i in range(len(a)-k+1)}\n    B={b[i:i+k] for i in range(len(b)-k+1)}\n    return len(A&B)/(len(A)+1e-6)\n\n\ndef find_templates(seq, top=20):\n\n    # fast shortlist\n    rough=[]\n\n    for r in train_records:\n        tid=r.target_id\n        tseq=r.sequence\n\n        # length filter\n        if abs(len(seq)-len(tseq))/max(len(seq),len(tseq))>0.4:\n            continue\n\n        rough.append((kmer_score(seq,tseq),tid,tseq))\n\n    # keep only best 60\n    rough.sort(reverse=True)\n    rough=rough[:60]\n\n\n    # accurate alignment only on shortlist\n    scores=[]\n    for _,tid,tseq in rough:\n        _,score=align(seq,tseq)\n        scores.append((score,tid,tseq))\n\n    scores.sort(reverse=True)\n    return scores[:top]\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def kabsch(P,Q):\n    Pc=P-P.mean(0)\n    Qc=Q-Q.mean(0)\n    C=Pc.T@Qc\n    V,S,W=np.linalg.svd(C)\n    d=np.sign(np.linalg.det(V@W))\n    R=V@np.diag([1,1,d])@W\n    return (Pc@R)+Q.mean(0)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ndef map_coords(query,template,coords):\n\n    pairs,_=align(query,template)\n\n    out=np.full((len(query),3),np.nan)\n\n    for qi,ti in pairs:\n        if ti < len(coords):\n            out[qi]=coords[ti]\n\n    # interpolation\n    for i in range(len(out)):\n        if np.isnan(out[i,0]):\n            l=i-1\n            while l>=0 and np.isnan(out[l,0]): l-=1\n            r=i+1\n            while r<len(out) and np.isnan(out[r,0]): r+=1\n\n            if l>=0 and r<len(out):\n                w=(i-l)/(r-l)\n                out[i]=(1-w)*out[l]+w*out[r]\n            elif l>=0:\n                out[i]=out[l]+[6,0,0]\n            elif r<len(out):\n                out[i]=out[r]+[6,0,0]\n            else:\n                out[i]=[i*6,0,0]\n\n    return out\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def refine(X,strength=0.3,passes=2):\n    X=X.copy()\n    for _ in range(passes):\n\n        d=X[1:]-X[:-1]\n        dist=np.linalg.norm(d,axis=1)+1e-6\n        target=6\n\n        adj=(target-dist)/dist\n        v=d*adj[:,None]*0.5*strength\n        X[:-1]-=v\n        X[1:]+=v\n\n        lap=0.5*(X[:-2]+X[2:])-X[1:-1]\n        X[1:-1]+=0.2*strength*lap\n\n    return X","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def rotate(X,rng):\n    axis=rng.normal(size=3)\n    axis/=np.linalg.norm(axis)\n    ang=rng.uniform(-0.5,0.5)\n    x,y,z=axis\n    c,s=np.cos(ang),np.sin(ang)\n\n    R=np.array([\n        [c+x*x*(1-c),x*y*(1-c)-z*s,x*z*(1-c)+y*s],\n        [y*x*(1-c)+z*s,c+y*y*(1-c),y*z*(1-c)-x*s],\n        [z*x*(1-c)-y*s,z*y*(1-c)+x*s,c+z*z*(1-c)]\n    ])\n    return X@R.T\n\n\ndef jitter(X,rng,amp=0.5):\n    return X+rng.normal(0,amp,X.shape)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict(row):\n\n    seq=row.sequence\n    temps=find_templates(seq,25)\n\n    coords_list=[]\n    weights=[]\n\n    for score,tid,tseq in temps:\n        mapped=map_coords(seq,tseq,train_coords[tid])\n        coords_list.append(mapped)\n        weights.append(score)\n\n    weights=np.array(weights)\n    weights-=weights.min()\n    weights+=1e-6\n    weights/=weights.sum()\n\n    # align templates before averaging\n    base=coords_list[0].copy()\n    aligned=[base]\n\n    for c in coords_list[1:]:\n        aligned.append(kabsch(c,base))\n\n    base=sum(w*c for w,c in zip(weights,aligned))\n\n    base=refine(base,0.6,3)\n\n    preds=[]\n    for k in range(5):\n\n        rng=np.random.default_rng(abs(hash(row.target_id))+k)\n\n        X=base.copy()\n\n        if k==1: X=jitter(X,rng,0.4)\n        if k==2: X=rotate(X,rng)\n        if k==3: X=jitter(rotate(X,rng),rng,0.7)\n        if k==4: X=refine(X,0.5,4)\n\n        preds.append(X)\n\n    return preds","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nrows=[]\nt0=time.time()\n\nfor i,row in test_seq.iterrows():\n\n    if i%10==0:\n        print(i,\"time\",time.time()-t0)\n\n    preds=predict(row)\n\n    for j,res in enumerate(row.sequence):\n\n        d={\n            \"ID\":f\"{row.target_id}_{j+1}\",\n            \"resname\":res,\n            \"resid\":j+1\n        }\n\n        for p in range(5):\n            x,y,z=preds[p][j]\n            d[f\"x_{p+1}\"]=x\n            d[f\"y_{p+1}\"]=y\n            d[f\"z_{p+1}\"]=z\n\n        rows.append(d)\n\nsub=pd.DataFrame(rows)\n\ncols=[\"ID\",\"resname\",\"resid\"]+[f\"{c}_{i}\" for i in range(1,6) for c in \"xyz\"]\nsub[cols[3:]]=sub[cols[3:]].clip(-999.999,9999.999)\n\nsub[cols].to_csv(\"submission.csv\",index=False)\n\nprint(\"DONE\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}