如何优化rna逆合成模型

1

git clone https://ghfast.top/github.com/chaitjo/geometric-rna-design

1

cd /root/autodl-tmp/geometric-rna-design

1

mamba create -p /root/autodl-tmp/rna python=3.10

1

mamba activate /root/autodl-tmp/rna

1
2
3
4
5
6

export PROJECT_PATH='/root/autodl-tmp/geometric-rna-design/'
export ETERNAFOLD='/root/autodl-tmp/geometric-rna-design/tools/EternaFold'
export X3DNA='/root/autodl-tmp/geometric-rna-design/tools/x3dna-v2.4'
export PATH="/root/autodl-tmp/geometric-rna-design/tools/x3dna-v2.4/bin:$PATH"
export PATH="/root/autodl-tmp/cdhit:$PATH"
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

1
2
3
4
5
6
7
8
9

export PROJECT_PATH='/home/remote1/geometric-rna-design/'
export DATA_PATH='/home/remote1/geometric-rna-design/data/'
export WANDB_PROJECT='rna'
export WANDB_ENTITY='wenxy59-sun-yat-sen-university'
export WANDB_DIR='/home/remote1/geometric-rna-design/'
export ETERNAFOLD='/home/remote1/geometric-rna-design/tools/EternaFold'
export X3DNA='/home/remote1/geometric-rna-design/tools/x3dna-v2.4'
export PATH="/home/remote1/geometric-rna-design/tools/x3dna-v2.4/bin:$PATH"
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

1
2
3

git clone https://github.com/pylelab/USalign.git
cd USalign
g++ -O3 -o qTMclust qTMclust.cpp -lm

1

python gRNAde.py     --pdb_filepath data/raw/6J6G_1_L-E.pdb     --output_filepath tutorial/lnc/po/114.fasta     --split das     --max_num_conformers 1  --n_samples 16     --temperature 0.5

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

conda create -n rna2 python=3.10 -y
conda activate rna2
pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
pip install absl-py==0.15.0
pip install biopython==1.80
pip install dm_tree==0.1.7
pip install fair-esm==2.0.0
pip install ml_collections==0.1.1
pip install numpy==1.24.3
pip install scipy>=1.10.0
pip install tqdm==4.64.1
pip install torch-cluster==1.6.1+pt113cu116 -f https://data.pyg.org/whl/torch-1.13.0+cu116.html
pip install torch-scatter==2.1.1+pt113cu116 -f https://data.pyg.org/whl/torch-1.13.0+cu116.html
pip install torch-geometric==2.3.1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

import torch
from tqdm import tqdm
import numpy as np
import random
from models import *
from utils import *
from diffusion import NoiseScheduleVP
from sampling import get_sampling_fn
from datasets import utils as du
import functools
import tree
from configs.inference_ribodiffusion import get_config
config = get_config()
# Choose heckpoint name
checkpoint_path = './ckpts/exp_inf.pth'
# checkpoint_path = './ckpts/exp_inf_large.pth'
config.eval.sampling_steps = 100
# config.eval.sampling_steps = 100
NUM_TO_LETTER = np.array(['A', 'G', 'C', 'U'])

def get_optimizer(config, params):
  """Return a flax optimizer object based on `config`."""
  if config.optim.optimizer == 'Adam':
      optimizer = optim.Adam(params, lr=config.optim.lr, betas=(config.optim.beta1, 0.999), eps=config.optim.eps, weight_decay=config.optim.weight_decay)
  elif config.optim.optimizer == 'AdamW':
      optimizer = torch.optim.AdamW(params, lr=config.optim.lr, amsgrad=True, weight_decay=1e-12)
  else:
      raise NotImplementedError(f'Optimizer {config.optim.optimizer} not supported yet!')
  return optimizer
# Initialize model
model = create_model(config)
ema = ExponentialMovingAverage(model.parameters(), decay=config.model.ema_decay)
params = model.parameters()
optimizer = get_optimizer(config, model.parameters())
state = dict(optimizer=optimizer, model=model, ema=ema, step=0)

model_size = sum(p.numel() for p in model.parameters()) * 4 / 2 ** 20
print('model size: {:.1f}MB'.format(model_size))

# Load checkpoint
state = restore_checkpoint(checkpoint_path, state, device=config.device)
ema.copy_to(model.parameters())

# Initialize noise scheduler
noise_scheduler = NoiseScheduleVP(config.sde.schedule, continuous_beta_0=config.sde.continuous_beta_0,
                                  continuous_beta_1=config.sde.continuous_beta_1)
# Obtain data scalar and inverse scalar
inverse_scaler = get_data_inverse_scaler(config)

# Setup sampling function
test_sampling_fn = get_sampling_fn(config, noise_scheduler, config.eval.sampling_steps, inverse_scaler)
pdb2data = functools.partial(du.PDBtoData, num_posenc=config.data.num_posenc, num_rbf=config.data.num_rbf, knn_num=config.data.knn_num)
# Run inference on a single p
pdb_file= '/home/remote1/geometric-rna-design/data/raw/1FIR_1_A.pdb'
pdb_id = pdb_file.replace('.pdb', '')
if '/' in pdb_id:
    pdb_id = pdb_id.split('/')[-1]

config.eval.dynamic_threshold=True
config.eval.cond_scale=0.4
config.eval.n_samples=16
test_sampling_fn = get_sampling_fn(config, noise_scheduler, config.eval.sampling_steps, inverse_scaler)
struct_data = pdb2data(pdb_file)
struct_data = tree.map_structure(lambda x:x.unsqueeze(0).repeat_interleave(config.eval.n_samples, dim=0).to(config.device), struct_data)
samples = test_sampling_fn(model, struct_data)
print(f'PDB ID: {pdb_id}')
native_seq = ''.join(list(NUM_TO_LETTER[struct_data['seq'][0].cpu().numpy()]))
print(f'Native sequence: {native_seq}')
for i in range(len(samples)):
    # native_seq = ''.join(list(NUM_TO_LETTER[struct_data['seq'].squeeze(0).cpu().numpy()]))
    # print(f'Native sequence: {native_seq}')
    designed_seq = ''.join(list(NUM_TO_LETTER[samples[i].cpu().numpy()]))
    print(f'Generated sequence {i+1}: {designed_seq}')
    recovery_ = samples[i].eq(struct_data['seq'][0]).float().mean().item()
    print(f'Recovery rate {i+1}: {recovery_:.4f}')

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275

import torch
from tqdm import tqdm
import numpy as np
import random
import os
import sys
from datetime import datetime

# Add the gRNAde path to import evaluation tools
sys.path.append('/home/remote1/geometric-rna-design/src')
sys.path.append('/home/remote1/geometric-rna-design/')

from models import *
from utils import *
from diffusion import NoiseScheduleVP
from sampling import get_sampling_fn
from datasets import utils as du
import functools
import tree
from configs.inference_ribodiffusion import get_config

# Import gRNAde evaluation functions
from src.evaluator import (
    self_consistency_score_eternafold,
    edit_distance
)
from src.data.data_utils import pdb_to_tensor, get_c4p_coords
from src.constants import NUM_TO_LETTER, PROJECT_PATH

# Import BioPython for sequence handling
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord

config = get_config()

# Choose checkpoint name
checkpoint_path = './ckpts/exp_inf.pth'
config.eval.sampling_steps = 100

def get_optimizer(config, params):
    """Return a flax optimizer object based on `config`."""
    if config.optim.optimizer == 'Adam':
        optimizer = optim.Adam(params, lr=config.optim.lr, betas=(config.optim.beta1, 0.999), 
                              eps=config.optim.eps, weight_decay=config.optim.weight_decay)
    elif config.optim.optimizer == 'AdamW':
        optimizer = torch.optim.AdamW(params, lr=config.optim.lr, amsgrad=True, weight_decay=1e-12)
    else:
        raise NotImplementedError(f'Optimizer {config.optim.optimizer} not supported yet!')
    return optimizer

def extract_sec_struct_from_pdb(pdb_file):
    """
    Extract secondary structure from PDB file using gRNAde's data utilities.
    """
    try:
        sequence, coords, sec_struct, sasa = pdb_to_tensor(
            pdb_file, 
            return_sec_struct=True, 
            return_sasa=True,
            keep_insertions=False
        )
        return [sec_struct] if sec_struct else None
    except Exception as e:
        print(f"Warning: Could not extract secondary structure from {pdb_file}: {e}")
        return None

def prepare_raw_data_for_grnade_eval(pdb_file, native_seq):
    """
    Prepare raw data structure compatible with gRNAde evaluator.
    """
    try:
        sequence, coords, sec_struct, sasa = pdb_to_tensor(
            pdb_file, 
            return_sec_struct=True, 
            return_sasa=True,
            keep_insertions=False
        )
        
        raw_data = {
            'sequence': native_seq,
            'coords_list': [coords] if coords is not None else [],
            'sec_struct_list': [sec_struct] if sec_struct else ["."] * len(native_seq),
            'sasa_list': [sasa] if sasa is not None else [np.ones(len(native_seq))]
        }
        
        return raw_data
    except Exception as e:
        print(f"Warning: Could not prepare raw data from {pdb_file}: {e}")
        # Return minimal raw data structure
        return {
            'sequence': native_seq,
            'coords_list': [],
            'sec_struct_list': ["."] * len(native_seq),
            'sasa_list': [np.ones(len(native_seq))]
        }

def evaluate_ribodiffusion_with_grnade_sc(samples, native_seq, pdb_file, pdb_id, 
                                         output_dir=None, save_results=True):
    """
    Evaluate RiboDiffusion samples using gRNAde's self-consistency evaluation.
    """
    # Prepare raw data for gRNAde evaluator
    raw_data = prepare_raw_data_for_grnade_eval(pdb_file, native_seq)
    
    # Convert RiboDiffusion samples to numpy arrays
    sample_arrays = []
    for sample in samples:
        if isinstance(sample, torch.Tensor):
            sample_arrays.append(sample.cpu().numpy())
        else:
            sample_arrays.append(np.array(sample))
    
    # Create mask for coordinates (assume all positions are valid for now)
    mask_coords = np.ones(len(native_seq), dtype=bool)
    
    # Calculate basic metrics
    results = {
        'pdb_id': pdb_id,
        'native_seq': native_seq,
        'samples': [],
        'recovery_rates': [],
        'edit_distances': [],
        'sc_scores_eternafold': []
    }
    
    # Convert native sequence to numerical for recovery calculation
    letter_to_num = {letter: idx for idx, letter in enumerate(NUM_TO_LETTER)}
    native_array = np.array([letter_to_num[char] for char in native_seq])
    
    print(f"\nEvaluating {len(samples)} samples for {pdb_id}:")
    
    for i, sample_array in enumerate(sample_arrays):
        designed_seq = ''.join([NUM_TO_LETTER[num] for num in sample_array])
        results['samples'].append(designed_seq)
        
        # Calculate recovery rate
        recovery = (sample_array == native_array).mean()
        results['recovery_rates'].append(recovery)
        
        # Calculate edit distance using gRNAde's function
        edit_dist = edit_distance(designed_seq, native_seq)
        results['edit_distances'].append(edit_dist)
        
        print(f'Sample {i+1}: Recovery={recovery:.4f}, Edit_dist={edit_dist}')
    
    # Calculate self-consistency scores using gRNAde's EternaFold evaluator
    print("\nCalculating self-consistency scores with EternaFold...")
    try:
        sc_scores = self_consistency_score_eternafold(
            sample_arrays,
            raw_data['sec_struct_list'],
            mask_coords
        )
        results['sc_scores_eternafold'] = sc_scores.tolist()
        
        for i, sc_score in enumerate(sc_scores):
            print(f'Sample {i+1}: SC_score={sc_score:.4f}')
            
    except Exception as e:
        print(f"Warning: Could not calculate SC scores: {e}")
        print("This might be due to EternaFold not being properly installed or configured.")
        results['sc_scores_eternafold'] = [0.0] * len(samples)
    
    # Save results
    if save_results and output_dir:
        os.makedirs(output_dir, exist_ok=True)
        
        # Save as FASTA file compatible with gRNAde format
        sequences = []
        
        # First record: input sequence with metadata
        sequences.append(SeqRecord(
            Seq(native_seq), 
            id="input_sequence,",
            description=f"pdb_id={pdb_id}, ribodiffusion_evaluation"
        ))
        
        # Remaining records: designed sequences with metrics
        for i, (seq, recovery, edit_dist, sc_score) in enumerate(zip(
            results['samples'], 
            results['recovery_rates'], 
            results['edit_distances'],
            results['sc_scores_eternafold']
        )):
            sequences.append(SeqRecord(
                Seq(seq), 
                id=f"sample={i},",
                description=f"recovery={recovery:.4f}, edit_dist={edit_dist}, sc_score={sc_score:.4f}"
            ))
        
        # Save FASTA
        fasta_path = os.path.join(output_dir, f"{pdb_id}_ribodiffusion_designs.fasta")
        SeqIO.write(sequences, fasta_path, "fasta")
        print(f"\nResults saved to: {fasta_path}")
    
    # Print summary statistics
    print(f"\n{'='*50}")
    print(f"Summary for {pdb_id}:")
    print(f"Native sequence length: {len(native_seq)}")
    print(f"Number of samples: {len(samples)}")
    print(f"Mean Recovery: {np.mean(results['recovery_rates']):.4f} ± {np.std(results['recovery_rates']):.4f}")
    print(f"Mean Edit Distance: {np.mean(results['edit_distances']):.2f} ± {np.std(results['edit_distances']):.2f}")
    
    if results['sc_scores_eternafold'][0] != 0.0:
        print(f"Mean SC Score (EternaFold): {np.mean(results['sc_scores_eternafold']):.4f} ± {np.std(results['sc_scores_eternafold']):.4f}")
    else:
        print("SC Scores not calculated (EternaFold unavailable)")
    print(f"{'='*50}")
    
    return results

# Initialize model
model = create_model(config)
ema = ExponentialMovingAverage(model.parameters(), decay=config.model.ema_decay)
params = model.parameters()
optimizer = get_optimizer(config, model.parameters())
state = dict(optimizer=optimizer, model=model, ema=ema, step=0)

model_size = sum(p.numel() for p in model.parameters()) * 4 / 2 ** 20
print('Model size: {:.1f}MB'.format(model_size))

# Load checkpoint
state = restore_checkpoint(checkpoint_path, state, device=config.device)
ema.copy_to(model.parameters())

# Initialize noise scheduler
noise_scheduler = NoiseScheduleVP(config.sde.schedule, 
                                 continuous_beta_0=config.sde.continuous_beta_0,
                                 continuous_beta_1=config.sde.continuous_beta_1)

# Obtain data scalar and inverse scalar
inverse_scaler = get_data_inverse_scaler(config)

# Setup sampling function
test_sampling_fn = get_sampling_fn(config, noise_scheduler, config.eval.sampling_steps, inverse_scaler)
pdb2data = functools.partial(du.PDBtoData, num_posenc=config.data.num_posenc, 
                            num_rbf=config.data.num_rbf, knn_num=config.data.knn_num)

# Run inference
pdb_file = '/home/remote1/geometric-rna-design/data/raw/7PIC_1_5.pdb'
pdb_id = os.path.basename(pdb_file).replace('.pdb', '')

# Configure sampling
config.eval.dynamic_threshold = True
config.eval.cond_scale = 0.4
config.eval.n_samples = 16

# Generate samples
print(f'Processing PDB: {pdb_file}')
test_sampling_fn = get_sampling_fn(config, noise_scheduler, config.eval.sampling_steps, inverse_scaler)
struct_data = pdb2data(pdb_file)
struct_data = tree.map_structure(
    lambda x: x.unsqueeze(0).repeat_interleave(config.eval.n_samples, dim=0).to(config.device), 
    struct_data
)
samples = test_sampling_fn(model, struct_data)

# Get native sequence
native_seq = ''.join(list(NUM_TO_LETTER[struct_data['seq'][0].cpu().numpy()]))
print(f'Native sequence: {native_seq}')

# Create output directory
current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
output_dir = f"ribodiffusion_grnade_eval_{current_time}"

# Evaluate with gRNAde's self-consistency framework
results = evaluate_ribodiffusion_with_grnade_sc(
    samples=samples,
    native_seq=native_seq,
    pdb_file=pdb_file,
    pdb_id=pdb_id,
    output_dir=output_dir,
    save_results=True
)

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736

import os
import copy
import shutil
from datetime import datetime

import numpy as np
import pandas as pd
from tqdm import tqdm
import wandb

import torch
import torch.nn.functional as F
from torchmetrics.functional.classification import binary_matthews_corrcoef

from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord

from MDAnalysis.analysis.align import rotation_matrix
from MDAnalysis.analysis.rms import rmsd as get_rmsd

from src.data.data_utils import pdb_to_tensor, get_c4p_coords
from src.data.sec_struct_utils import (
    predict_sec_struct,
    dotbracket_to_paired,
    dotbracket_to_adjacency
)
from src.constants import (
    NUM_TO_LETTER, 
    PROJECT_PATH,
    RMSD_THRESHOLD,
    TM_THRESHOLD,
    GDT_THRESHOLD
)


def evaluate(
        model, 
        dataset, 
        n_samples, 
        temperature, 
        device, 
        model_name="eval",
        metrics=[
            'recovery', 'perplexity', 'sc_score_eternafold', 
            'sc_score_ribonanzanet', 'sc_score_rhofold'
        ],
        save_designs=False
    ):
    """
    Run evaluation suite for trained RNA inverse folding model on a dataset.

    The following metrics can be computed along with metadata per sample per residue:
    1. (recovery) Sequence recovery per residue (taking mean gives per sample recovery)
    2. (perplexity) Perplexity per sample
    3. (sc_score_eternafold) Secondary structure self-consistency score per sample, 
        using EternaFold for secondary structure prediction and computing MCC between
        the predicted and groundtruth 2D structures as adjacency matrices.
    4. (sc_score_ribonanzanet) Chemical modification self-consistency score per sample,
        using RibonanzaNet for chemical modification prediction of the groundtruth and
        designed sequences, and measuring MAE between them.
    5. (sc_score_rhofold) Tertiary structure self-consistency scores per sample,
        using RhoFold for tertiary structure prediction and measuring RMSD, TM-score,
        and GDT_TS between the predicted and groundtruth C4' 3D coordinates.
    6. (rmsd_within_thresh) Percentage of samples with RMSD within threshold (<=2.0A)
    7. (tm_within_thresh) Percentage of samples with TM-score within threshold (>=0.45)
    8. (gddt_within_thresh) Percentage of samples with GDT_TS within threshold (>=0.50)

    Args:
        model: trained RNA inverse folding model
        dataset: dataset to evaluate on
        n_samples: number of predicted samples/sequences per data point 
        temperature: sampling temperature
        device: device to run evaluation on
        model_name: name of model/dataset for plotting (default: 'eval')
        metrics: list of metrics to compute
        save_designs: whether to save designs as fasta with metrics
    
    Returns: Dictionary with the following keys:
        df: DataFrame with metrics and metadata per residue per sample for analysis and plotting
        samples_list: list of tensors of shape (n_samples, seq_len) per data point 
        recovery_list: list of mean recovery per data point
        perplexity_list: list of mean perplexity per data point
        sc_score_eternafold_list: list of 2D self-consistency scores per data point
        sc_score_ribonanzanet_list: list of 1D self-consistency scores per data point
        sc_score_rmsd_list: list of 3D self-consistency RMSDs per data point
        sc_score_tm_list: list of 3D self-consistency TM-scores per data point
        sc_score_gddt_list: list of 3D self-consistency GDTs per data point
        rmsd_within_thresh_list: list of % scRMSDs within threshold per data point
        tm_within_thresh_list: list of % scTMs within threshold per data point
        gddt_within_thresh_list: list of % scGDDTs within threshold per data point
    """
    assert 'recovery' in metrics, 'Sequence recovery must be computed for evaluation'

    #######################################################################
    # Optionally initialise other models used for self-consistency scoring
    #######################################################################

    if 'sc_score_ribonanzanet' in metrics:
        from tools.ribonanzanet.network import RibonanzaNet
        
        # Initialise RibonanzaNet for self-consistency score
        ribonanza_net = RibonanzaNet(
            os.path.join(PROJECT_PATH, 'tools/ribonanzanet/config.yaml'),
            os.path.join(PROJECT_PATH, 'tools/ribonanzanet/ribonanzanet.pt'),
            device
        )
        # Transfer model to device in eval mode
        ribonanza_net = ribonanza_net.to(device)
        ribonanza_net.eval()
    
    if 'sc_score_rhofold' in metrics:
        from tools.rhofold.rf import RhoFold
        from tools.rhofold.config import rhofold_config
        
        # Initialise RhoFold for 3D self-consistency score
        rhofold = RhoFold(rhofold_config, device)
        rhofold_path = os.path.join(PROJECT_PATH, "tools/rhofold/model_20221010_params.pt")
        print(f"Loading RhoFold checkpoint: {rhofold_path}")
        rhofold.load_state_dict(torch.load(rhofold_path, map_location=torch.device('cpu'))['model'])
        # Transfer model to device in eval mode
        rhofold = rhofold.to(device)
        rhofold.eval()
        current_datetime = datetime.now().strftime("%Y%m%d_%H%M%S")

    ####################################################
    # Evaluation loop over each data point sequentially
    ####################################################

    # per sample metric lists for storing evaluation results
    samples_list = []               # list of tensors of shape (n_samples, seq_len) per data point 
    recovery_list = []              # list of mean recovery per data point
    perplexity_list = []            # list of mean perplexity per data point
    sc_score_ribonanzanet_list = [] # list of 1D self-consistency scores per data point
    sc_score_eternafold_list = []   # list of 2D self-consistency scores per data point
    sc_score_rmsd_list = []         # list of 3D self-consistency RMSDs per data point
    rmsd_within_thresh_list = []    # list of % scRMSDs within threshold per data point
    sc_score_tm_list = []           # list of 3D self-consistency TM-scores per data point
    tm_within_thresh_list = []      # list of % scTMs within threshold per data point
    sc_score_gddt_list = []         # list of 3D self-consistency GDTs per data point
    gddt_within_thresh_list = []    # list of % scGDDTs within threshold per data point

    # DataFrame to store metrics and metadata per residue per sample for analysis and plotting
    df = pd.DataFrame(columns=['idx', 'recovery', 'sasa', 'paired', 'rmsds', 'model_name'])

    model.eval()
    if device.type == 'xpu':
        import intel_extension_for_pytorch as ipex
        model = ipex.optimize(model)
        if 'sc_score_ribonanzanet' in metrics:
            ribonanza_net = ipex.optimize(ribonanza_net)
        if 'sc_score_rhofold' in metrics:
            rhofold = ipex.optimize(rhofold)
    
    with torch.no_grad():
        for idx, raw_data in tqdm(
            enumerate(dataset.data_list),
            total=len(dataset.data_list)
        ):
            # featurise raw data
            data = dataset.featurizer(raw_data).to(device)

            # sample n_samples from model for single data point: n_samples x seq_len
            samples, logits = model.sample(data, n_samples, temperature, return_logits=True)
            samples_list.append(samples.cpu().numpy())
            
            # perplexity per sample: n_samples x 1
            n_nodes = logits.shape[1]
            perplexity = torch.exp(F.cross_entropy(
                logits.view(n_samples * n_nodes, model.out_dim), 
                samples.view(n_samples * n_nodes).long(), 
                reduction="none"
            ).view(n_samples, n_nodes).mean(dim=1)).cpu().numpy()
            perplexity_list.append(perplexity.mean())

            ###########
            # Metadata
            ###########

            # per residue average SASA: seq_len x 1
            mask_coords = data.mask_coords.cpu().numpy()
            sasa = np.mean(raw_data['sasa_list'], axis=0)[mask_coords]

            # per residue indicator for paired/unpaired: seq_len x 1
            paired = np.mean(
                [dotbracket_to_paired(sec_struct) for sec_struct in raw_data['sec_struct_list']], axis=0
            )[mask_coords]

            # per residue average RMSD: seq_len x 1
            if len(raw_data["coords_list"]) == 1:
                rmsds = np.zeros_like(sasa)
            else:
                rmsds = []
                for i in range(len(raw_data["coords_list"])):
                    for j in range(i+1, len(raw_data["coords_list"])):
                        coords_i = get_c4p_coords(raw_data["coords_list"][i])
                        coords_j = get_c4p_coords(raw_data["coords_list"][j])
                        rmsds.append(torch.sqrt(torch.sum((coords_i - coords_j)**2, dim=1)).cpu().numpy())
                rmsds = np.stack(rmsds).mean(axis=0)[mask_coords]

            ##########
            # Metrics
            ##########

            # sequence recovery per residue across all samples: n_samples x seq_len 
            recovery = samples.eq(data.seq).float().cpu().numpy()
            recovery_list.append(recovery.mean())

            # update per residue per sample dataframe
            df = pd.concat([
                df, 
                pd.DataFrame({
                    'idx': [idx] * len(recovery.mean(axis=0)),
                    'recovery': recovery.mean(axis=0),
                    'sasa': sasa,
                    'paired': paired,
                    'rmsds': rmsds,
                    'model_name': [model_name] * len(recovery.mean(axis=0))
                })
            ], ignore_index=True)

            # global 2D self consistency score per sample: n_samples x 1
            if 'sc_score_eternafold' in metrics:
                sc_score_eternafold, pred_sec_structs = self_consistency_score_eternafold(
                    samples.cpu().numpy(), 
                    raw_data['sec_struct_list'], 
                    mask_coords,
                    return_sec_structs = True
                )
                sc_score_eternafold_list.append(sc_score_eternafold.mean())

            # global 1D self consistency score per sample: n_samples x 1
            if 'sc_score_ribonanzanet' in metrics:
                sc_score_ribonanzanet, pred_chem_mods = self_consistency_score_ribonanzanet(
                    samples.cpu().numpy(), 
                    raw_data['sequence'],
                    mask_coords, 
                    ribonanza_net,
                    return_chem_mods = True
                )
                sc_score_ribonanzanet_list.append(sc_score_ribonanzanet.mean())
            
            # global 3D self consistency scores per sample: n_samples x 1, each
            if 'sc_score_rhofold' in metrics:
                try:
                    output_dir = os.path.join(
                        wandb.run.dir, f"designs_{model_name}/{current_datetime}/sample{idx}/")
                except AttributeError:
                    output_dir = os.path.join(
                        PROJECT_PATH, f"designs_{model_name}/{current_datetime}/sample{idx}/")

                sc_score_rmsd, sc_score_tm, sc_score_gdt = self_consistency_score_rhofold(
                    samples.cpu().numpy(), 
                    raw_data,
                    mask_coords,
                    rhofold,
                    output_dir,
                    save_designs = save_designs
                )
                sc_score_rmsd_list.append(sc_score_rmsd.mean())
                sc_score_tm_list.append(sc_score_tm.mean())
                sc_score_gddt_list.append(sc_score_gdt.mean())

                rmsd_within_thresh_list.append((sc_score_rmsd <= RMSD_THRESHOLD).sum() / n_samples)
                tm_within_thresh_list.append((sc_score_tm >= TM_THRESHOLD).sum() / n_samples)
                gddt_within_thresh_list.append((sc_score_gdt >= GDT_THRESHOLD).sum() / n_samples)

                if save_designs:
                    # collate designed sequences in fasta format
                    sequences = [SeqRecord(
                        Seq(raw_data["sequence"]), id=f"input_sequence,", 
                        description=f"pdb_id={raw_data['id_list'][0]} rfam={raw_data['rfam_list'][0]} eq_class={raw_data['eq_class_list'][0]} cluster={raw_data['cluster_structsim0.45']}"
                    )]
                    for idx, zipped in enumerate(zip(
                        samples.cpu().numpy(),
                        perplexity,
                        recovery.mean(axis=1),
                        sc_score_eternafold,
                        pred_sec_structs,
                        sc_score_ribonanzanet,
                        pred_chem_mods,
                        sc_score_rmsd,
                        sc_score_tm,
                        sc_score_gdt
                    )):
                        seq, perp, rec, sc, pred_ss, sc_ribo, pred_cm, sc_rmsd, sc_tm, sc_gdt = zipped
                        seq = "".join([NUM_TO_LETTER[num] for num in seq])
                        edit_dist = edit_distance(seq, raw_data['sequence'])
                        sequences.append(SeqRecord(
                            Seq(seq), id=f"sample={idx},",
                            description=f"temperature={temperature} perplexity={perp:.4f} recovery={rec:.4f} edit_dist={edit_dist} sc_score={sc:.4f} sc_score_ribonanzanet={sc_ribo:.4f} sc_score_rmsd={sc_rmsd:.4f} sc_score_tm={sc_tm:.4f} sc_score_gdt={sc_gdt:.4f}"
                        ))
                    # write all designed sequences to output filepath
                    SeqIO.write(sequences, os.path.join(output_dir, "all_designs.fasta"), "fasta")

    out = {
        'df': df,
        'samples_list': samples_list,
        'recovery_list': recovery_list,
        'perplexity_list': perplexity_list
    }
    if 'sc_score_eternafold' in metrics:
        out['sc_score_eternafold'] = sc_score_eternafold_list
    if 'sc_score_ribonanzanet' in metrics:
        out['sc_score_ribonanzanet'] = sc_score_ribonanzanet_list
    if 'sc_score_rhofold' in metrics:
        out['sc_score_rmsd'] = sc_score_rmsd_list
        out['sc_score_tm'] = sc_score_tm_list
        out['sc_score_gddt'] = sc_score_gddt_list
        out['rmsd_within_thresh'] = rmsd_within_thresh_list
        out['tm_within_thresh'] = tm_within_thresh_list
        out['gddt_within_thresh'] = gddt_within_thresh_list
        # =========================================================
    
    import gc 

    
    if 'ribonanza_net' in locals():
        del ribonanza_net
    if 'rhofold' in locals():
        del rhofold

    
    gc.collect()

 
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    # =========================================================
    return out


def self_consistency_score_eternafold(
        samples, 
        true_sec_struct_list, 
        mask_coords,
        n_samples_ss = 1,
        num_to_letter = NUM_TO_LETTER,
        return_sec_structs = False
    ):
    """
    Compute self consistency score for an RNA, given its true secondary structure(s)
    and a list of designed sequences. 
    EternaFold is used to 'forward fold' the designs.
    
    Args:
        samples: designed sequences of shape (n_samples, seq_len)
        true_sec_struct_list: list of true secondary structures (n_true_ss, seq_len)
        mask_coords: mask for missing sequence coordinates to be ignored during evaluation
        n_samples_ss: number of predicted secondary structures per designed sample
        num_to_letter: lookup table mapping integers to nucleotides
        return_sec_structs: whether to return the predicted secondary structures
    
    Workflow:
        
        Input: For a given RNA molecule, we are given:
        - Designed sequences of shape (n_samples, seq_len)
        - True secondary structure(s) of shape (n_true_ss, seq_len)
        
        For each designed sequence:
        - Predict n_sample_ss secondary structures using EternaFold
        - For each pair of true and predicted secondary structures:
            - Compute MCC score between their adjacency matrix representations
        - Take the average MCC score across all n_sample_ss predicted structures
        
        Take the average MCC score across all n_samples designed sequences
    """
    
    n_true_ss = len(true_sec_struct_list)
    sequence_length = mask_coords.sum()
    # map all entries from dotbracket to numerical representation
    true_sec_struct_list = np.array([dotbracket_to_adjacency(ss) for ss in true_sec_struct_list])
    # mask out missing sequence coordinates
    true_sec_struct_list = true_sec_struct_list[:, mask_coords][:, :, mask_coords]
    # reshape to (n_true_ss * n_samples_ss, seq_len, seq_len)
    true_sec_struct_list = torch.tensor(
        true_sec_struct_list
    ).unsqueeze(1).repeat(1, n_samples_ss, 1, 1).reshape(-1, sequence_length, sequence_length)

    mcc_scores = []
    pred_sec_structs = []
    for _sample in samples:
        # convert sample to string
        pred_seq = ''.join([num_to_letter[num] for num in _sample])
        # predict secondary structure(s) for each sample
        pred_sec_struct_list = predict_sec_struct(pred_seq, n_samples=n_samples_ss)
        if return_sec_structs:
            pred_sec_structs.append(copy.copy(pred_sec_struct_list))
        # map all entries from dotbracket to numerical representation
        pred_sec_struct_list = np.array([dotbracket_to_adjacency(ss) for ss in pred_sec_struct_list])
        # reshape to (n_samples_ss * n_true_ss, seq_len, seq_len)
        pred_sec_struct_list = torch.tensor(
            pred_sec_struct_list
        ).unsqueeze(0).repeat(n_true_ss, 1, 1, 1).reshape(-1, sequence_length, sequence_length)

        # compute mean MCC score between pairs of true and predicted secondary structures
        mcc_scores.append(
            binary_matthews_corrcoef(
                pred_sec_struct_list,
                true_sec_struct_list,
            ).float().mean()
        )

    if return_sec_structs:
        return np.array(mcc_scores), pred_sec_structs
    else:
        return np.array(mcc_scores)


def self_consistency_score_ribonanzanet(
    samples,
    true_sequence,
    mask_seq,
    ribonanza_net,
    num_to_letter=NUM_TO_LETTER,
    return_chem_mods=False,
):
    """Compute self consistency score for an RNA, given the (predicted) chemical modifications for
    the original RNA and a list of designed sequences. RibonanzaNet is used to 'forward fold' the
    designs.

    Args:
        samples: designed sequences of shape (n_samples, seq_len)
        true_sequence: true RNA sequence used to predict chemical modifications
        mask_seq: mask for missing sequence coordinates to be ignored during evaluation
        ribonanza_net: RibonanzaNet model
        num_to_letter: lookup table mapping integers to nucleotides
        return_chem_mods: whether to return the predicted chemical modifications

    Workflow:

        Input: For a given RNA molecule, we are given:
        - Designed sequences of shape (n_samples, seq_len)
        - Predicted chemical modifications for original sequence,
          of shape (n_samples, seq_len, 2), predicted via RibonanzaNet, of which we take
          the index 0 from the last channal --> 2A3/SHAPE.

        For each designed sequence:
        - Predict chemical modifications using RibonanzaNet
        - Compute mean absolute error between prediction and chemical modifications for
          the original sequence

        Take the average mean absolute error across all n_samples designed sequences
    """
    # Compute original sequence's chemical modifications using RibonanzaNet
    true_sequence = np.array([char for char in true_sequence])
    true_sequence = "".join(true_sequence[mask_seq])
    true_chem_mod = ribonanza_net.predict(true_sequence).unsqueeze(0).cpu().numpy()[:,:,0]

     # 2. 【核心修改】分批处理模型生成的样本
    _samples_char = np.array([[num_to_letter[num] for num in seq] for seq in samples])
    
    batch_size = 1  # 定义一个小的批次大小，可以根据显存调整（比如4, 8, 16）
    all_preds = [] # 用于收集所有批次的预测结果

    for i in range(0, len(_samples_char), batch_size):
        # 取出一个小批次
        batch_samples = _samples_char[i:i+batch_size]
        # 对这个小批次进行预测
        batch_pred = ribonanza_net.predict(batch_samples).cpu().numpy()[:,:,0]
        all_preds.append(batch_pred)
    
    # 将所有批次的结果拼接起来
    pred_chem_mod = np.concatenate(all_preds, axis=0)
   # _samples_char = np.array([[num_to_letter[num] for num in seq] for seq in samples])
   # _samples = np.array([[num_to_letter[num] for num in seq] for seq in samples])
   #pred_chem_mod = ribonanza_net.predict(_samples_char).cpu().numpy()[:,:,0]
   ##pred_chem_mod = ribonanza_net.predict(_samples[:, mask_seq]).cpu().numpy()[:,:,0]
    if return_chem_mods:
        return (np.abs(pred_chem_mod - true_chem_mod).mean(1)), pred_chem_mod
    else:
        return np.abs(pred_chem_mod - true_chem_mod).mean(1)


def self_consistency_score_ribonanzanet_sec_struct(
        samples, 
        true_sec_struct, 
        mask_coords, 
        ribonanza_net_ss,
        num_to_letter = NUM_TO_LETTER,
        return_sec_structs = False
    ):
    # map from dotbracket to numerical representation
    true_sec_struct = np.array(dotbracket_to_adjacency(true_sec_struct, keep_pseudoknots=True))
    # mask out missing sequence coordinates
    true_sec_struct = true_sec_struct[mask_coords][:, mask_coords]
    # (n_samples, seq_len, seq_len)
    true_sec_struct = torch.tensor(true_sec_struct)

    _samples = np.array([[num_to_letter[num] for num in seq] for seq in samples])
    _, pred_sec_structs = ribonanza_net_ss.predict(_samples)  # (n_samples, seq_len, seq_len)
    
    mcc_scores = []
    for pred_sec_struct in pred_sec_structs:
        # map from dotbracket to numerical representation
        pred_sec_struct = torch.tensor(dotbracket_to_adjacency(pred_sec_struct, keep_pseudoknots=True))
        # compute mean MCC score between pairs of true and predicted secondary structures
        mcc_scores.append(
            binary_matthews_corrcoef(
                pred_sec_struct,
                true_sec_struct,
            ).float().mean()
        )

    if return_sec_structs:
        return np.array(mcc_scores), pred_sec_structs
    else:
        return np.array(mcc_scores)


def self_consistency_score_rhofold(
        samples,
        true_raw_data,
        mask_coords,
        rhofold,
        output_dir,
        num_to_letter = NUM_TO_LETTER,
        save_designs = False,
        save_pdbs = False,
        use_relax = False,
    ):
    """
    Compute self consistency score for an RNA, given its true 3D structure(s)
    for the original RNA and a list of designed sequences.
    RhoFold is used to 'forward fold' the designs.

    Credit: adapted from Rishabh Anand

    Args:
        samples: designed sequences of shape (n_samples, seq_len)
        true_raw_data: Original RNA raw data with 3D structure(s) in `coords_list`
        mask_coords: mask for missing sequence coordinates to be ignored during evaluation
        rhofold: RhoFold model
        output_dir: directory to save designed sequences and structures
        num_to_letter: lookup table mapping integers to nucleotides
        save_designs: whether to save designs as fasta to output directory
        save_pdbs: whether to save PDBs of forward-folded designs to output directory
        use_relax: whether to perform Amber relaxation on designed structures

    Workflow:
            
        Input: For a given RNA molecule, we are given:
        - Designed sequences of shape (n_samples, seq_len)
        - True 3D structure(s) of shape (n_true_structs, seq_len, 3)
        
        For each designed sequence:
        - Predict the tertiary structure using RhoFold
        - For each pair of true and predicted 3D structures:
            - Compute RMSD, TM-score & GDT between their C4' coordinates
        
        Take the average self-consistency scores across all n_samples designed sequences

    Returns:
        sc_rmsds: array of RMSD scores per sample
        sc_tms: array of TM-score scores per sample
        sc_gddts: array of GDT scores per sample
    """
    os.makedirs(output_dir, exist_ok=True)

    # Collate designed sequences in fasta format
    # first record: input sequence and model metadata
    input_seq = SeqRecord(
        Seq(true_raw_data["sequence"]),
        id=f"input_sequence,",
        description=f"input_sequence"
    )
    # SeqIO.write(input_seq, os.path.join(output_dir, "input_seq.fasta"), "fasta")
    sequences = [input_seq]
    
    # remaining records: designed sequences and metrics
    sc_rmsds = []
    sc_tms = []
    sc_gddts = []
    for idx, seq in enumerate(samples):
        # Save designed sequence to fasta file (temporary)
        seq = SeqRecord(
            Seq("".join([num_to_letter[num] for num in seq])), 
            id=f"sample={idx},",
            description=f"sample={idx}"
        )
        sequences.append(seq)
        design_fasta_path = os.path.join(output_dir, f"design{idx}.fasta")
        SeqIO.write(seq, design_fasta_path, "fasta")
        
        # Forward fold designed sequence using RhoFold
        design_pdb_path = os.path.join(output_dir, f"design{idx}.pdb")
        rhofold.predict(design_fasta_path, design_pdb_path, use_relax)
        
        # Load C4' coordinates of designed structure
        _, coords, _, _ = pdb_to_tensor(
            design_pdb_path,
            return_sec_struct=False,
            return_sasa=False,
            keep_insertions=False,
        )
        coords = get_c4p_coords(coords)
        # zero-center coordinates
        coords = coords - coords.mean(dim=0)

        # Compute self-consistency between designed and groundtruth structures
        _sc_rmsds = []
        _sc_tms = []
        _sc_gddts = []
        for other_coords in true_raw_data["coords_list"]:
            _other = get_c4p_coords(other_coords)[mask_coords, :]
            # zero-center other coordinates
            _other = _other - _other.mean(dim=0)
            # globally align coordinates
            R_hat = rotation_matrix(
                _other,  # mobile set
                coords # reference set
            )[0]
            _other = _other @ R_hat.T
            # compute metrics
            _sc_rmsds.append(get_rmsd(
                coords, _other, superposition=True, center=True))
            _sc_tms.append(get_tmscore(coords, _other))
            _sc_gddts.append(get_gddt(coords, _other))

        sc_rmsds.append(np.mean(_sc_rmsds))
        sc_tms.append(np.mean(_sc_tms))
        sc_gddts.append(np.mean(_sc_gddts))

        # remove temporary files
        os.unlink(design_fasta_path)
        if save_pdbs is False:
            os.unlink(design_pdb_path)
    
    if save_designs is False:
        # remove output directory        
        shutil.rmtree(output_dir)
    else:
        # write all designed sequences to output filepath
        SeqIO.write(sequences, os.path.join(output_dir, "all_designs.fasta"), "fasta")

    return np.array(sc_rmsds), np.array(sc_tms), np.array(sc_gddts)


def get_tmscore(y_hat: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Template Modelling score (TM-score). 
    
    Credit: Arian Jamasb, graphein (https://github.com/a-r-j/graphein)

    https://en.wikipedia.org/wiki/Template_modeling_score

    TM-score is a measure of similarity between two protein structures.
    The TM-score is intended as a more accurate measure of the global
    similarity of full-length protein structures than the often used RMSD
    measure. The TM-score indicates the similarity between two structures
    by a score between ``[0, 1]``, where 1 indicates a perfect match
    between two structures (thus the higher the better). Generally scores
    below 0.20 corresponds to randomly chosen unrelated proteins whereas
    structures with a score higher than 0.5 assume roughly the same fold.
    A quantitative study shows that proteins of TM-score = 0.5 have a
    posterior probability of 37% in the same CATH topology family and of
    13% in the same SCOP fold family. The probabilities increase rapidly
    when TM-score > 0.5. The TM-score is designed to be independent of
    protein lengths.
    
    We have adapted the implementation to RNA (TM-score threshold = 0.45).
    Requires aligned C4' coordinates as input.
    """
    l_target = y.shape[0]
    d0_l_target = 1.24 * np.power(l_target - 15, 1 / 3) - 1.8
    di = torch.pairwise_distance(y_hat, y)
    out = torch.sum(1 / (1 + (di / d0_l_target) ** 2)) / l_target
    if torch.isnan(out):
        return torch.tensor(0.0)
    return out


def get_gddt(y_hat: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Global Distance Deviation Test metric (GDDT).

    Credit: Arian Jamasb, graphein (https://github.com/a-r-j/graphein)

    https://en.wikipedia.org/wiki/Global_distance_test

    The GDT score is calculated as the largest set of amino acid residues'
    alpha carbon atoms in the model structure falling within a defined
    distance cutoff of their position in the experimental structure, after
    iteratively superimposing the two structures. By the original design the
    GDT algorithm calculates 20 GDT scores, i.e. for each of 20 consecutive distance
    cutoffs (``0.5 Å, 1.0 Å, 1.5 Å, ... 10.0 Å``). For structure similarity assessment
    it is intended to use the GDT scores from several cutoff distances, and scores
    generally increase with increasing cutoff. A plateau in this increase may
    indicate an extreme divergence between the experimental and predicted structures,
    such that no additional atoms are included in any cutoff of a reasonable distance.
    The conventional GDT_TS total score in CASP is the average result of cutoffs at
    ``1``, ``2``, ``4``, and ``8`` Å.

    Random predictions give around 20; getting the gross topology right gets one to ~50; 
    accurate topology is usually around 70; and when all the little bits and pieces, 
    including side-chain conformations, are correct, GDT_TS begins to climb above 90.

    We have adapted the implementation to RNA.
    Requires aligned C4' coordinates as input.
    """
    # Get distance between points
    dist = torch.norm(y - y_hat, dim=1)

    # Return mean fraction of distances below cutoff for each cutoff (1, 2, 4, 8)
    count_1 = (dist < 1).sum() / dist.numel()
    count_2 = (dist < 2).sum() / dist.numel()
    count_4 = (dist < 4).sum() / dist.numel()
    count_8 = (dist < 8).sum() / dist.numel()
    out = torch.mean(torch.tensor([count_1, count_2, count_4, count_8]))
    if torch.isnan(out):
        return torch.tensor(0.0)
    return out


def edit_distance(s: str, t: str) -> int:
    """
    A Space efficient Dynamic Programming based Python3 program 
    to find minimum number operations to convert str1 to str2

    Source: https://www.geeksforgeeks.org/edit-distance-dp-5/
    """
    n = len(s)
    m = len(t)

    prev = [j for j in range(m+1)]
    curr = [0] * (m+1)

    for i in range(1, n+1):
        curr[0] = i
        for j in range(1, m+1):
            if s[i-1] == t[j-1]:
                curr[j] = prev[j-1]
            else:
                mn = min(1 + prev[j], 1 + curr[j-1])
                curr[j] = min(mn, 1 + prev[j-1])
        prev = curr.copy()

    return prev[m]