chore: fix regression tests for rf3 (#725)

models/rf3/tests/data/8vkf_from_file.cif

@@ -1,6 +1,8 @@
 data_8VKF
 # 
 _entry.id   8VKF 
+#
+_msa_paths_by_chain_id.A   models/rf3/tests/data/msas/8vkf_A.a3m
 # 
 _audit_conform.dict_name       mmcif_pdbx.dic 
 _audit_conform.dict_version    5.402 

models/rf3/tests/data/inference_regression_tests/5vht_from_file/5vht_from_file_metrics.csv

@@ -1,2 +0,0 @@
-example_id,ptm.ptm_0,ptm.ptm_1,ptm.ptm_2,ptm.ptm_3,ptm.ptm_4,iptm.iptm_0,iptm.iptm_1,iptm.iptm_2,iptm.iptm_3,iptm.iptm_4,iptm.iptm_protein_protein_0,iptm.iptm_protein_protein_1,iptm.iptm_protein_protein_2,iptm.iptm_protein_protein_3,iptm.iptm_protein_protein_4,iptm.iptm_protein_ligand_0,iptm.iptm_protein_ligand_1,iptm.iptm_protein_ligand_2,iptm.iptm_protein_ligand_3,iptm.iptm_protein_ligand_4,iptm.iptm_ligand_ligand_0,iptm.iptm_ligand_ligand_1,iptm.iptm_ligand_ligand_2,iptm.iptm_ligand_ligand_3,iptm.iptm_ligand_ligand_4,count_clashing_chains.has_clash_0,count_clashing_chains.has_clash_1,count_clashing_chains.has_clash_2,count_clashing_chains.has_clash_3,count_clashing_chains.has_clash_4
-5vht_from_file,0.9095256,0.90964013,0.90937024,0.9094581,0.9096535,0.909933,0.9100767,0.9099717,0.91014886,0.91009957,0.909933,0.9100767,0.9099717,0.91014886,0.91009957,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0,0,0,0,0

models/rf3/tests/data/inference_regression_tests/5vht_from_file/5vht_from_file_score.csv

@@ -1,16 +0,0 @@
-example_id,chain_chainwise,chainwise_plddt,chainwise_pde,chainwise_pae,overall_plddt,overall_pde,overall_pae,batch_idx,chain_i_interface,chain_j_interface,pae_interface,pde_interface,min_pae_interface,min_pde_interface
-5vht_from_file,A_1,0.8418933153152466,0.7738104462623596,4.3772172927856445,0.8415207266807556,0.8318539261817932,4.677595138549805,0,,,,,,
-5vht_from_file,B_1,0.8411479592323303,0.7790139317512512,4.396655082702637,0.8415207266807556,0.8318539261817932,4.677595138549805,0,,,,,,
-5vht_from_file,A_1,0.8418263792991638,0.7724140882492065,4.372034072875977,0.8414488434791565,0.8303730487823486,4.672811985015869,1,,,,,,
-5vht_from_file,B_1,0.8410713076591492,0.7779868841171265,4.392614364624023,0.8414488434791565,0.8303730487823486,4.672811985015869,1,,,,,,
-5vht_from_file,A_1,0.8417466878890991,0.7801038026809692,4.385378837585449,0.8413788676261902,0.8371773958206177,4.68752908706665,2,,,,,,
-5vht_from_file,B_1,0.8410109877586365,0.7858399748802185,4.405368804931641,0.8413788676261902,0.8371773958206177,4.68752908706665,2,,,,,,
-5vht_from_file,A_1,0.8417938947677612,0.77611243724823,4.379835605621338,0.8414207100868225,0.8332493901252747,4.680868625640869,3,,,,,,
-5vht_from_file,B_1,0.841047465801239,0.781791627407074,4.400188446044922,0.8414207100868225,0.8332493901252747,4.680868625640869,3,,,,,,
-5vht_from_file,A_1,0.841849148273468,0.7727293372154236,4.370999813079834,0.8414701819419861,0.8308433294296265,4.672425746917725,4,,,,,,
-5vht_from_file,B_1,0.8410910964012146,0.7783981561660767,4.391624450683594,0.8414701819419861,0.8308433294296265,4.672425746917725,4,,,,,,
-5vht_from_file,,,,,0.8415207266807556,0.8318539261817932,4.677595138549805,0,A_1,B_1,4.968254089355469,0.8872956037521362,0.7079706192016602,0.28655317425727844
-5vht_from_file,,,,,0.8414488434791565,0.8303730487823486,4.672811985015869,1,A_1,B_1,4.963299751281738,0.8855456709861755,0.6997964382171631,0.2865622043609619
-5vht_from_file,,,,,0.8413788676261902,0.8371773958206177,4.68752908706665,2,A_1,B_1,4.979684352874756,0.8913829326629639,0.7079532742500305,0.2865457534790039
-5vht_from_file,,,,,0.8414207100868225,0.8332493901252747,4.680868625640869,3,A_1,B_1,4.971724987030029,0.8875466585159302,0.6996998190879822,0.2865447700023651
-5vht_from_file,,,,,0.8414701819419861,0.8308433294296265,4.672425746917725,4,A_1,B_1,4.9635396003723145,0.8861227631568909,0.7079214453697205,0.286527156829834

models/rf3/tests/data/inference_regression_tests/5vht_from_file/5vht_from_file_summary_confidences.json

@@ -0,0 +1,26 @@
+{
+  "chain_ptm": [0.85,0.84],
+  "chain_pair_pae_min": [
+    [null,0.72],
+    [null,null]
+  ],
+  "chain_pair_pde_min": [
+    [null,0.29],
+    [null,null]
+  ],
+  "chain_pair_pae": [
+    [null,4.86],
+    [null,null]
+  ],
+  "chain_pair_pde": [
+    [null,0.86],
+    [null,null]
+  ],
+  "overall_plddt": 0.8449,
+  "overall_pde": 0.809,
+  "overall_pae": 4.5743,
+  "ptm": 0.9105080962181091,
+  "iptm": 0.9110991358757019,
+  "has_clash": false,
+  "ranking_score": 0.911
+}

models/rf3/tests/data/inference_regression_tests/8vkf_from_file/8vkf_from_file_metrics.csv

@@ -1,2 +0,0 @@
-example_id,ptm.ptm_0,ptm.ptm_1,ptm.ptm_2,ptm.ptm_3,ptm.ptm_4,iptm.iptm_0,iptm.iptm_1,iptm.iptm_2,iptm.iptm_3,iptm.iptm_4,iptm.iptm_protein_protein_0,iptm.iptm_protein_protein_1,iptm.iptm_protein_protein_2,iptm.iptm_protein_protein_3,iptm.iptm_protein_protein_4,iptm.iptm_protein_ligand_0,iptm.iptm_protein_ligand_1,iptm.iptm_protein_ligand_2,iptm.iptm_protein_ligand_3,iptm.iptm_protein_ligand_4,iptm.iptm_ligand_ligand_0,iptm.iptm_ligand_ligand_1,iptm.iptm_ligand_ligand_2,iptm.iptm_ligand_ligand_3,iptm.iptm_ligand_ligand_4,count_clashing_chains.has_clash_0,count_clashing_chains.has_clash_1,count_clashing_chains.has_clash_2,count_clashing_chains.has_clash_3,count_clashing_chains.has_clash_4
-8vkf_from_file,0.9397696,0.94022816,0.93984985,0.9406191,0.94042176,0.9700439,0.96997976,0.96992123,0.9700142,0.9700208,0.0,0.0,0.0,0.0,0.0,0.9700439,0.96997976,0.96992123,0.9700142,0.9700208,0.93240464,0.9323982,0.93233144,0.93221885,0.93262637,0,0,0,0,0

models/rf3/tests/data/inference_regression_tests/8vkf_from_file/8vkf_from_file_score.csv

@@ -1,76 +0,0 @@
-example_id,chain_chainwise,chainwise_plddt,chainwise_pde,chainwise_pae,overall_plddt,overall_pde,overall_pae,batch_idx,chain_i_interface,chain_j_interface,pae_interface,pde_interface,min_pae_interface,min_pde_interface
-8vkf_from_file,A_1,0.9320401549339294,0.9453674554824829,3.9356305599212646,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,,,,,,
-8vkf_from_file,B_1,0.9495342969894409,0.3442486524581909,1.6652324199676514,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,,,,,,
-8vkf_from_file,C_1,0.9272092580795288,0.3786144256591797,1.5746772289276123,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,,,,,,
-8vkf_from_file,D_1,0.9011284112930298,4.034590721130371,7.017383098602295,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,,,,,,
-8vkf_from_file,E_1,0.8282343745231628,13.62786865234375,8.54509449005127,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,,,,,,
-8vkf_from_file,A_1,0.932106077671051,0.93052077293396,3.9099152088165283,0.9322792887687683,0.8681702613830566,3.882534980773926,1,,,,,,
-8vkf_from_file,B_1,0.9495639204978943,0.34412628412246704,1.6654282808303833,0.9322792887687683,0.8681702613830566,3.882534980773926,1,,,,,,
-8vkf_from_file,C_1,0.9272972345352173,0.3788919746875763,1.574660062789917,0.9322792887687683,0.8681702613830566,3.882534980773926,1,,,,,,
-8vkf_from_file,D_1,0.9011081457138062,4.034159183502197,7.01707124710083,0.9322792887687683,0.8681702613830566,3.882534980773926,1,,,,,,
-8vkf_from_file,E_1,0.8281877636909485,13.573341369628906,8.574715614318848,0.9322792887687683,0.8681702613830566,3.882534980773926,1,,,,,,
-8vkf_from_file,A_1,0.932047963142395,0.9424294233322144,3.931513547897339,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,,,,,,
-8vkf_from_file,B_1,0.94952392578125,0.3441573977470398,1.6650241613388062,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,,,,,,
-8vkf_from_file,C_1,0.9272950291633606,0.37870317697525024,1.5746617317199707,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,,,,,,
-8vkf_from_file,D_1,0.9011602401733398,4.035135746002197,7.019920349121094,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,,,,,,
-8vkf_from_file,E_1,0.8280327320098877,13.63625717163086,8.625662803649902,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,,,,,,
-8vkf_from_file,A_1,0.9320767521858215,0.9201368093490601,3.8943328857421875,0.9322513937950134,0.859204888343811,3.8692407608032227,3,,,,,,
-8vkf_from_file,B_1,0.9496347904205322,0.344124972820282,1.6652103662490845,0.9322513937950134,0.859204888343811,3.8692407608032227,3,,,,,,
-8vkf_from_file,C_1,0.92729252576828,0.37906670570373535,1.5739272832870483,0.9322513937950134,0.859204888343811,3.8692407608032227,3,,,,,,
-8vkf_from_file,D_1,0.9012048244476318,4.029945373535156,7.012320518493652,0.9322513937950134,0.859204888343811,3.8692407608032227,3,,,,,,
-8vkf_from_file,E_1,0.8280571103096008,13.597679138183594,8.608673095703125,0.9322513937950134,0.859204888343811,3.8692407608032227,3,,,,,,
-8vkf_from_file,A_1,0.9320934414863586,0.9275834560394287,3.903271198272705,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,,,,,,
-8vkf_from_file,B_1,0.9496190547943115,0.344104528427124,1.6647707223892212,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,,,,,,
-8vkf_from_file,C_1,0.9273494482040405,0.37881240248680115,1.5747339725494385,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,,,,,,
-8vkf_from_file,D_1,0.9010905623435974,4.033513069152832,7.0145745277404785,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,,,,,,
-8vkf_from_file,E_1,0.8283198475837708,13.572347640991211,8.580305099487305,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,,,,,,
-8vkf_from_file,,,,,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,A_1,B_1,3.4717438220977783,0.605563223361969,0.6303903460502625,0.2651037275791168
-8vkf_from_file,,,,,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,A_1,C_1,4.8907575607299805,0.6995029449462891,0.7152637839317322,0.28492048382759094
-8vkf_from_file,,,,,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,A_1,D_1,4.7494916915893555,0.8386649489402771,2.0623931884765625,0.46002161502838135
-8vkf_from_file,,,,,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,A_1,E_1,9.049422264099121,3.8051230907440186,5.0449137687683105,2.8304648399353027
-8vkf_from_file,,,,,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,B_1,C_1,2.736272096633911,0.36354562640190125,1.1683330535888672,0.2732650637626648
-8vkf_from_file,,,,,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,B_1,D_1,3.264293909072876,0.48301342129707336,2.083875894546509,0.36788517236709595
-8vkf_from_file,,,,,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,B_1,E_1,7.824339389801025,2.5479378700256348,5.2585883140563965,1.8988345861434937
-8vkf_from_file,,,,,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,C_1,D_1,4.22511625289917,0.5288299918174744,2.2046451568603516,0.4303090572357178
-8vkf_from_file,,,,,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,C_1,E_1,9.384553909301758,2.672896146774292,7.508654594421387,2.3068904876708984
-8vkf_from_file,,,,,0.9322137236595154,0.8812259435653687,3.9045498371124268,0,D_1,E_1,9.492976188659668,4.036022186279297,8.316246032714844,4.03602409362793
-8vkf_from_file,,,,,0.9322792887687683,0.8681702613830566,3.882534980773926,1,A_1,B_1,3.459951162338257,0.5984316468238831,0.630462110042572,0.26509952545166016
-8vkf_from_file,,,,,0.9322792887687683,0.8681702613830566,3.882534980773926,1,A_1,C_1,4.877954006195068,0.6918987035751343,0.7154307961463928,0.2848895192146301
-8vkf_from_file,,,,,0.9322792887687683,0.8681702613830566,3.882534980773926,1,A_1,D_1,4.736978530883789,0.8306165933609009,2.0626161098480225,0.4596916437149048
-8vkf_from_file,,,,,0.9322792887687683,0.8681702613830566,3.882534980773926,1,A_1,E_1,9.124225616455078,3.7755820751190186,5.10176944732666,2.8448750972747803
-8vkf_from_file,,,,,0.9322792887687683,0.8681702613830566,3.882534980773926,1,B_1,C_1,2.73687744140625,0.3636227250099182,1.169067144393921,0.27333831787109375
-8vkf_from_file,,,,,0.9322792887687683,0.8681702613830566,3.882534980773926,1,B_1,D_1,3.262706756591797,0.4828460216522217,2.0824265480041504,0.36787867546081543
-8vkf_from_file,,,,,0.9322792887687683,0.8681702613830566,3.882534980773926,1,B_1,E_1,7.864833354949951,2.5528337955474854,5.208006858825684,1.8787106275558472
-8vkf_from_file,,,,,0.9322792887687683,0.8681702613830566,3.882534980773926,1,C_1,D_1,4.22160530090332,0.5286768078804016,2.1941330432891846,0.4300752282142639
-8vkf_from_file,,,,,0.9322792887687683,0.8681702613830566,3.882534980773926,1,C_1,E_1,9.301826477050781,2.655578374862671,7.413575172424316,2.298459053039551
-8vkf_from_file,,,,,0.9322792887687683,0.8681702613830566,3.882534980773926,1,D_1,E_1,9.516148567199707,3.908454418182373,8.317085266113281,3.908456325531006
-8vkf_from_file,,,,,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,A_1,B_1,3.4694740772247314,0.6041339039802551,0.6303707361221313,0.2650924623012543
-8vkf_from_file,,,,,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,A_1,C_1,4.888820171356201,0.698183000087738,0.7153699994087219,0.28490278124809265
-8vkf_from_file,,,,,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,A_1,D_1,4.746122360229492,0.8371659517288208,2.0624828338623047,0.4598386585712433
-8vkf_from_file,,,,,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,A_1,E_1,9.238248825073242,3.801729202270508,5.277400970458984,2.8974578380584717
-8vkf_from_file,,,,,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,B_1,C_1,2.7369372844696045,0.36360710859298706,1.1681721210479736,0.2720290422439575
-8vkf_from_file,,,,,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,B_1,D_1,3.2642736434936523,0.48311784863471985,2.0822458267211914,0.36787092685699463
-8vkf_from_file,,,,,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,B_1,E_1,7.957021236419678,2.5393381118774414,5.358416557312012,1.888516902923584
-8vkf_from_file,,,,,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,C_1,D_1,4.223407745361328,0.5279184579849243,2.2048254013061523,0.4302324652671814
-8vkf_from_file,,,,,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,C_1,E_1,9.441669464111328,2.6686744689941406,7.508003234863281,2.2565596103668213
-8vkf_from_file,,,,,0.9322216510772705,0.8786524534225464,3.9016854763031006,2,D_1,E_1,9.5327730178833,3.8588662147521973,8.329977989196777,3.85886812210083
-8vkf_from_file,,,,,0.9322513937950134,0.859204888343811,3.8692407608032227,3,A_1,B_1,3.4515013694763184,0.5935373902320862,0.6304340958595276,0.2650916874408722
-8vkf_from_file,,,,,0.9322513937950134,0.859204888343811,3.8692407608032227,3,A_1,C_1,4.87033224105835,0.6863856911659241,0.715640127658844,0.2848808169364929
-8vkf_from_file,,,,,0.9322513937950134,0.859204888343811,3.8692407608032227,3,A_1,D_1,4.7307820320129395,0.8252673745155334,2.062479019165039,0.45959168672561646
-8vkf_from_file,,,,,0.9322513937950134,0.859204888343811,3.8692407608032227,3,A_1,E_1,9.2240629196167,3.7983005046844482,5.149802207946777,2.77072811126709
-8vkf_from_file,,,,,0.9322513937950134,0.859204888343811,3.8692407608032227,3,B_1,C_1,2.734774589538574,0.36362209916114807,1.167860746383667,0.27200546860694885
-8vkf_from_file,,,,,0.9322513937950134,0.859204888343811,3.8692407608032227,3,B_1,D_1,3.2621781826019287,0.4826110005378723,2.07240629196167,0.3676457107067108
-8vkf_from_file,,,,,0.9322513937950134,0.859204888343811,3.8692407608032227,3,B_1,E_1,7.96647834777832,2.5434417724609375,5.3074846267700195,1.9209169149398804
-8vkf_from_file,,,,,0.9322513937950134,0.859204888343811,3.8692407608032227,3,C_1,D_1,4.221768379211426,0.5291066765785217,2.2043018341064453,0.4300277829170227
-8vkf_from_file,,,,,0.9322513937950134,0.859204888343811,3.8692407608032227,3,C_1,E_1,9.462154388427734,2.6579153537750244,7.632106304168701,2.2554261684417725
-8vkf_from_file,,,,,0.9322513937950134,0.859204888343811,3.8692407608032227,3,D_1,E_1,9.495095252990723,3.771697521209717,8.34188461303711,3.7716994285583496
-8vkf_from_file,,,,,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,A_1,B_1,3.453770637512207,0.5964496731758118,0.6305050253868103,0.26513877511024475
-8vkf_from_file,,,,,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,A_1,C_1,4.868809700012207,0.6897773146629333,0.7071991562843323,0.2849225699901581
-8vkf_from_file,,,,,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,A_1,D_1,4.734051704406738,0.828239917755127,2.0547072887420654,0.45981019735336304
-8vkf_from_file,,,,,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,A_1,E_1,9.135178565979004,3.786433219909668,5.102996826171875,2.852632761001587
-8vkf_from_file,,,,,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,B_1,C_1,2.734165668487549,0.36359599232673645,1.1680809259414673,0.273250013589859
-8vkf_from_file,,,,,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,B_1,D_1,3.2627716064453125,0.48312124609947205,2.0826122760772705,0.36937710642814636
-8vkf_from_file,,,,,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,B_1,E_1,7.866828441619873,2.5545129776000977,5.217214584350586,1.8813700675964355
-8vkf_from_file,,,,,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,C_1,D_1,4.222236156463623,0.528620183467865,2.2064924240112305,0.43011271953582764
-8vkf_from_file,,,,,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,C_1,E_1,9.318958282470703,2.6771984100341797,7.419789791107178,2.3011441230773926
-8vkf_from_file,,,,,0.9322678446769714,0.8655266165733337,3.8760082721710205,4,D_1,E_1,9.404740333557129,3.7400143146514893,8.203256607055664,3.740015983581543

models/rf3/tests/data/inference_regression_tests/8vkf_from_file/8vkf_from_file_summary_confidences.json

@@ -0,0 +1,38 @@
+{
+  "chain_ptm": [0.93,0.95,0.93,0.9,0.83],
+  "chain_pair_pae_min": [
+    [null,0.63,0.72,2.05,5.03],
+    [null,null,1.16,2.08,5.22],
+    [null,null,null,2.21,7.53],
+    [null,null,null,null,8.19],
+    [null,null,null,null,null]
+  ],
+  "chain_pair_pde_min": [
+    [null,0.27,0.28,0.46,2.88],
+    [null,null,0.27,0.37,1.92],
+    [null,null,null,0.43,2.33],
+    [null,null,null,null,3.92],
+    [null,null,null,null,null]
+  ],
+  "chain_pair_pae": [
+    [null,3.44,4.85,4.72,9.02],
+    [null,null,2.73,3.26,7.75],
+    [null,null,null,4.19,9.27],
+    [null,null,null,null,9.38],
+    [null,null,null,null,null]
+  ],
+  "chain_pair_pde": [
+    [null,0.59,0.68,0.81,3.79],
+    [null,null,0.36,0.48,2.55],
+    [null,null,null,0.53,2.66],
+    [null,null,null,null,3.92],
+    [null,null,null,null,null]
+  ],
+  "overall_plddt": 0.9324,
+  "overall_pde": 0.8463,
+  "overall_pae": 3.8409,
+  "ptm": 0.9411807060241699,
+  "iptm": 0.9701294898986816,
+  "has_clash": false,
+  "ranking_score": 0.9643
+}

models/rf3/tests/test_inference_regression.py

@@ -1,11 +1,11 @@
 #!/usr/bin/env -S /bin/sh -c '"$(dirname "$0")/../../../.ipd/shebang/rf3_exec.sh" "$0" "$@"'
 
+import json
 import tempfile
 from pathlib import Path
 
 import biotite.structure as struc
 import numpy as np
-import pandas as pd
 import pytest
 from atomworks.io.parser import STANDARD_PARSER_ARGS, parse
 from atomworks.io.utils.io_utils import load_any
@@ -18,7 +18,6 @@ from conftest import TEST_DATA_DIR
 from hydra import compose, initialize
 from hydra.utils import instantiate
 from omegaconf import OmegaConf
-from rf3.utils.inference import InferenceInput
 
 RUN_PARAM_KEYS = {
     "inputs",
@@ -36,136 +35,41 @@ RUN_PARAM_KEYS = {
 """Run parameters that should be passed to engine.run(), not __init__."""
 
 
-def assert_similar_csv_files(
-    predicted_file: Path, baseline_file: Path, tolerance: float = 1e-3
-):
-    """Compare CSV files with numerical tolerance for floating-point values."""
-    predicted_df = pd.read_csv(predicted_file)
-    baseline_df = pd.read_csv(baseline_file)
+def compute_rmsd(
+    predicted: struc.AtomArray | struc.AtomArrayStack,
+    baseline: struc.AtomArray | struc.AtomArrayStack,
+) -> float:
+    """Compute RMSD between two structures after superposition."""
+    # Handle AtomArrayStack - extract first model
+    if isinstance(predicted, struc.AtomArrayStack):
+        predicted = predicted[0]
+    if isinstance(baseline, struc.AtomArrayStack):
+        baseline = baseline[0]
 
-    # Check shape
-    assert (
-        predicted_df.shape == baseline_df.shape
-    ), f"Shape mismatch in {predicted_file.name}: {predicted_df.shape} vs {baseline_df.shape}"
+    # Mask: only consider atoms resolved in both
+    baseline_mask = np.ones(len(baseline), dtype=bool)
+    if hasattr(baseline, "occupancy") and baseline.occupancy is not None:
+        baseline_mask = baseline.occupancy == 1
 
-    # Check column names (order-independent)
-    predicted_cols = set(predicted_df.columns)
-    baseline_cols = set(baseline_df.columns)
-    assert (
-        predicted_cols == baseline_cols
-    ), f"Column mismatch in {predicted_file.name}: {predicted_cols} vs {baseline_cols}"
+    predicted_mask = np.ones(len(predicted), dtype=bool)
+    if hasattr(predicted, "occupancy") and predicted.occupancy is not None:
+        predicted_mask = predicted.occupancy == 1
 
-    # Compare values with tolerance for numeric columns
-    for col in predicted_df.columns:
-        if predicted_df[col].dtype in ["float64", "float32", "int64", "int32"]:
-            # Numeric comparison with tolerance
-            mean_pred = predicted_df[col].mean()
-            mean_base = baseline_df[col].mean()
-            assert abs(mean_pred - mean_base) <= tolerance, (
-                f"Numerical difference {abs(mean_pred - mean_base)} exceeds tolerance {tolerance} in column {col} of {predicted_file.name} "
-                f"(predicted_mean={mean_pred}, baseline_mean={mean_base})"
-            )
-        else:
-            # Exact comparison for non-numeric
-            assert predicted_df[col].equals(
-                baseline_df[col]
-            ), f"Non-numeric content mismatch in column {col} of {predicted_file.name}"
+    resolved_mask = baseline_mask & predicted_mask
+    baseline = baseline[resolved_mask]
+    predicted = predicted[resolved_mask]
 
-
-def compare_structures(
-    predicted: struc.AtomArray | struc.AtomArrayStack | list[struc.AtomArray],
-    baseline: struc.AtomArray | struc.AtomArrayStack | list[struc.AtomArray],
-) -> list[float]:
-    """Compare structures with RMSD after alignment.
-
-    Args:
-        predicted: Predicted structure(s).
-        baseline: Baseline structure(s).
-
-    Returns:
-        List of RMSD values (one per model).
-    """
-
-    def _compare_single(
-        baseline_array: struc.AtomArray,
-        predicted_array: struc.AtomArray,
-    ) -> float:
-        """Compare two AtomArray objects and return RMSD."""
-        assert isinstance(baseline_array, struc.AtomArray)
-        assert isinstance(predicted_array, struc.AtomArray)
-        assert (
-            len(baseline_array) == len(predicted_array)
-        ), f"Atom count mismatch: baseline {len(baseline_array)} vs predicted {len(predicted_array)}"
-
-        # Mask: only consider atoms resolved in both baseline and predicted arrays
-        # Check if occupancy exists, otherwise assume all atoms are resolved
-        baseline_mask = np.ones(len(baseline_array), dtype=bool)
-        if (
-            hasattr(baseline_array, "occupancy")
-            and baseline_array.occupancy is not None
-        ):
-            baseline_mask = baseline_array.occupancy == 1
-
-        predicted_mask = np.ones(len(predicted_array), dtype=bool)
-        if (
-            hasattr(predicted_array, "occupancy")
-            and predicted_array.occupancy is not None
-        ):
-            predicted_mask = predicted_array.occupancy == 1
-
-        resolved_mask = baseline_mask & predicted_mask
-        baseline_array = baseline_array[resolved_mask]
-        predicted_array = predicted_array[resolved_mask]
-
-        # Superimpose and calculate RMSD
-        superimposed, _ = struc.superimpose(baseline_array, predicted_array)
-        rmsd_value = struc.rmsd(baseline_array, superimposed)
-        return rmsd_value
-
-    # Convert to common format: list of AtomArrays
-    def _to_atom_array_list(structures):
-        """Convert various structure formats to a list of AtomArray objects."""
-        if isinstance(structures, struc.AtomArrayStack):
-            return [structures[i] for i in range(len(structures))]
-        elif isinstance(structures, list):
-            # If it's a list of AtomArrayStacks, ensure they all are length 1, and downcast to AtomArrays
-            if all(isinstance(s, struc.AtomArrayStack) for s in structures):
-                assert all(
-                    len(s) == 1 for s in structures
-                ), "All AtomArrayStacks in the list must contain exactly one structure."
-                return [s[0] for s in structures]
-            # Assert list of AtomArrays
-            assert all(
-                isinstance(s, struc.AtomArray) for s in structures
-            ), "All elements in the list must be AtomArray objects."
-            return structures
-        else:
-            # Single AtomArray - wrap in list
-            return [structures]
-
-    predicted_list = _to_atom_array_list(predicted)
-    baseline_list = _to_atom_array_list(baseline)
-
-    # If baseline is a single structure, replicate to match predicted length
-    if len(baseline_list) == 1 and len(predicted_list) > 1:
-        baseline_list = baseline_list * len(predicted_list)
-
-    # Compare each pair
-    return [
-        _compare_single(base, pred) for base, pred in zip(baseline_list, predicted_list)
-    ]
+    superimposed, _ = struc.superimpose(baseline, predicted)
+    return struc.rmsd(baseline, superimposed)
 
 
 @pytest.mark.gpu
 @pytest.mark.parametrize(
-    "example_id,rmsd_tolerance,csv_tolerance",
-    [
-        ("5vht_from_file", 0.1, 0.01),
-        ("8vkf_from_file", 0.1, 0.01),
-    ],
+    "example_id",
+    ["5vht_from_file", "8vkf_from_file"],
 )
-def test_inference_regression(example_id, rmsd_tolerance, csv_tolerance):
-    """Test RF3 inference by comparing predictions to ground truth structure and baseline metrics."""
+def test_inference_regression(example_id):
+    """Test RF3 inference comparing predictions to baseline structure and confidence metrics."""
 
     input_file = TEST_DATA_DIR / f"{example_id}.cif"
     baseline_dir = TEST_DATA_DIR / "inference_regression_tests" / example_id
@@ -179,153 +83,58 @@ def test_inference_regression(example_id, rmsd_tolerance, csv_tolerance):
         ground_truth_full = remove_protein_terminal_oxygen(
             parse(input_file, **STANDARD_PARSER_ARGS)["assemblies"]["1"][0]
         )
-        # Filter to heavy atoms only (no hydrogen)
         ground_truth = ground_truth_full[ground_truth_full.element != "H"]
 
-        # Predict and save the results to the temp_dir
+        # Run inference
         cfg = compose(
             config_name="inference",
             overrides=[
                 "inference_engine=rf3",
                 f"inputs={input_file}",
                 "annotate_b_factor_with_plddt=true",
-                "one_model_per_file=false",
                 f"out_dir={temp_dir}",
             ],
         )
 
-        # Separate config into init params and run params
         cfg_dict = OmegaConf.to_container(cfg, resolve=True)
         run_params = {k: v for k, v in cfg_dict.items() if k in RUN_PARAM_KEYS}
         init_cfg_dict = {k: v for k, v in cfg_dict.items() if k not in RUN_PARAM_KEYS}
         init_cfg = OmegaConf.create(init_cfg_dict)
 
-        # Instantiate engine (only __init__ params)
         inference_engine = instantiate(init_cfg, _convert_="partial", _recursive_=False)
-
-        # Run inference with the new API
         inference_engine.run(**run_params)
 
-        # Outputs are now nested in a subdirectory named after the example_id
         predicted_dir = Path(temp_dir) / example_id
 
-        # Compare CSV metrics to baseline
-        baseline_metrics_csv = baseline_dir / f"{example_id}_metrics.csv"
-        predicted_metrics_csv = predicted_dir / f"{example_id}_metrics.csv"
-        assert_similar_csv_files(
-            predicted_metrics_csv, baseline_metrics_csv, tolerance=csv_tolerance
-        )
+        # Compare predicted structure to ground truth
+        predicted_model = load_any(predicted_dir / f"{example_id}_model.cif")
+        baseline_model = load_any(baseline_dir / f"{example_id}_model.cif")
 
-        # Compare predicted structures to ground truth
-        # Load all predicted structures
-        predicted_files = sorted(predicted_dir.glob("*.cif.gz"))
-        predicted_structures = [load_any(f) for f in predicted_files]
-        mean_predicted_rmsd = np.mean(
-            compare_structures(predicted_structures, ground_truth)
-        )
+        predicted_rmsd = compute_rmsd(predicted_model, ground_truth)
+        baseline_rmsd = compute_rmsd(baseline_model, ground_truth)
 
-        # Load all baseline structures
-        baseline_files = sorted(baseline_dir.glob(f"{example_id}_model_*.cif.gz"))
-        baseline_structures = [load_any(f) for f in baseline_files]
-        mean_baseline_rmsd = np.mean(
-            compare_structures(baseline_structures, ground_truth)
-        )
-
-        # Assert mean RMSD difference between baseline and predicted < threshold
-        rmsd_difference = abs(mean_predicted_rmsd - mean_baseline_rmsd)
+        rmsd_tolerance = 0.3
+        rmsd_difference = abs(predicted_rmsd - baseline_rmsd)
         assert (
             rmsd_difference < rmsd_tolerance
-        ), f"Mean RMSD difference {rmsd_difference:.4f}Å exceeds {rmsd_tolerance}Å tolerance for {example_id}"
+        ), f"RMSD difference {rmsd_difference:.4f}Å exceeds {rmsd_tolerance}Å tolerance"
 
+        # Compare confidence metrics to baseline
+        with open(predicted_dir / f"{example_id}_summary_confidences.json") as f:
+            predicted_conf = json.load(f)
+        with open(baseline_dir / f"{example_id}_summary_confidences.json") as f:
+            baseline_conf = json.load(f)
 
-@pytest.mark.gpu
-@pytest.mark.parametrize(
-    "example_id,rmsd_tolerance,csv_tolerance",
-    [
-        ("5vht_from_file", 0.1, 0.01),
-        ("8vkf_from_file", 0.1, 0.01),
-    ],
-)
-def test_inference_regression_in_memory(example_id, rmsd_tolerance, csv_tolerance):
-    """Test in-memory inference against baseline predictions."""
-
-    inputs_path = TEST_DATA_DIR / f"{example_id}.cif"
-    baseline_dir = TEST_DATA_DIR / "inference_regression_tests" / example_id
-
-    with (
-        initialize(config_path="../configs"),
-        tempfile.TemporaryDirectory() as temp_dir,
-        rng_state(create_rng_state_from_seeds(1, 1, 1)),
-    ):
-        # Load ground truth structure from input CIF
-        ground_truth_full = remove_protein_terminal_oxygen(
-            parse(inputs_path, **STANDARD_PARSER_ARGS)["assemblies"]["1"][0]
-        )
-
-        # Filter to heavy atoms only (no hydrogen)
-        ground_truth = ground_truth_full[ground_truth_full.element != "H"]
-
-        # Create InferenceInput from CIF file
-        inference_input = InferenceInput.from_cif_path(
-            inputs_path, example_id=example_id
-        )
-
-        # Load config
-        cfg = compose(
-            config_name="inference",
-            overrides=[
-                "inference_engine=rf3",
-            ],
-        )
-
-        # Separate config into init params and run params
-        cfg_dict = OmegaConf.to_container(cfg, resolve=True)
-        init_cfg_dict = {k: v for k, v in cfg_dict.items() if k not in RUN_PARAM_KEYS}
-        init_cfg = OmegaConf.create(init_cfg_dict)
-
-        # Initialize engine
-        engine = instantiate(init_cfg, _convert_="partial", _recursive_=False)
-
-        # Run inference in-memory
-        results = engine.run(
-            inputs=inference_input,
-            out_dir=None,  # Return in-memory
-            annotate_b_factor_with_plddt=True,
-        )
-
-        # Extract results for this example
-        result = results[example_id]
-
-        # Compare CSV metrics to baseline
-        # (Write predicted metrics to temp CSV for comparison)
-        predicted_metrics_df = pd.DataFrame([result["metrics"]])
-        predicted_metrics_csv = Path(temp_dir) / f"{example_id}_metrics.csv"
-        predicted_metrics_df.to_csv(predicted_metrics_csv, index=False)
-        baseline_metrics_csv = baseline_dir / f"{example_id}_metrics.csv"
-        assert_similar_csv_files(
-            predicted_metrics_csv, baseline_metrics_csv, tolerance=csv_tolerance
-        )
-
-        # Compare predicted structures to ground truth
-        predicted_structures = result["predicted_structures"]
-        mean_predicted_rmsd = np.mean(
-            compare_structures(predicted_structures, ground_truth)
-        )
-
-        # Load baseline structures and compare to ground truth
-        baseline_structures = []
-        for i in range(len(predicted_structures)):
-            baseline_file = baseline_dir / f"{example_id}_model_{i}.cif.gz"
-            baseline_structures.append(load_any(baseline_file))
-        mean_baseline_rmsd = np.mean(
-            compare_structures(baseline_structures, ground_truth)
-        )
-
-        # Assert mean RMSD difference < threshold
-        rmsd_difference = abs(mean_predicted_rmsd - mean_baseline_rmsd)
-        assert (
-            rmsd_difference < rmsd_tolerance
-        ), f"Mean RMSD difference {rmsd_difference:.4f}Å exceeds {rmsd_tolerance}Å tolerance for {example_id}"
+        confidence_tolerance = 0.02
+        for key in ["overall_plddt", "ptm", "iptm"]:
+            pred_val = predicted_conf.get(key)
+            base_val = baseline_conf.get(key)
+            if pred_val is not None and base_val is not None:
+                diff = abs(pred_val - base_val)
+                assert diff <= confidence_tolerance, (
+                    f"{key} difference {diff:.4f} exceeds tolerance {confidence_tolerance} "
+                    f"(predicted={pred_val:.4f}, baseline={base_val:.4f})"
+                )
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -44,7 +44,7 @@ dependencies = [
     "zstandard",
     "pandas",
     # "biotite",
-    "atomworks==2.1.1",
+    "atomworks[ml]>=2.1.1",
     "ipykernel>=6.31.0",
     "assertpy",
     "toolz",
@@ -79,7 +79,7 @@ dev = [
     "pytest-dotenv>=0.5.2,<1", # load environment variables from .env file
     "pytest-cov>=4.1.0,<5", # generate coverage report
     "pytest-benchmark>=5.0.0,<6", # benchmark tests for speed
-    "atomworks==2.1.1",
+    "atomworks[ml,dev,openbabel]>=2.1.1",
     "pre-commit"
 ]
 [project.scripts]