Upload folder using huggingface_hub (part 2)

model/boltzgen/resources/config/train/boltzgen.no_distillation.yaml

@@ -0,0 +1,484 @@
+_target_: boltzgen.task.train.train.Training
+
+trainer:
+  accelerator: gpu
+  devices: 8
+  precision: bf16-mixed
+  gradient_clip_val: 10.0
+  accumulate_grad_batches: 1
+  max_epochs: -1
+  num_sanity_val_steps: 3
+  log_every_n_steps: 1
+
+wandb:
+  group: boltzgen
+  project: boltzgen
+  entity: yourwandb
+
+name: a_big_run_resume3
+slurm: true
+output: workdir
+strict_loading: false
+resume: null
+pretrained: null
+debug: false
+save_every_n_train_steps: 2500
+disable_checkpoint: false
+matmul_precision: null
+save_top_k: -1
+
+data:
+  datasets:
+    - _target_: boltzgen.task.train.data.DatasetConfig
+      target_dir: ./training_data/targets
+      msa_dir: ./training_data/msa
+      prob: 1.0
+      filters:
+        - _target_: boltzgen.data.filter.dynamic.size.SizeFilter
+          min_chains: 1
+          max_chains: 300
+        - _target_: boltzgen.data.filter.dynamic.date.DateFilter
+          date: "2023-06-01"
+          ref: released
+        - _target_: boltzgen.data.filter.dynamic.resolution.ResolutionFilter
+          resolution: 9.0
+      sampler:
+        _target_: boltzgen.data.sample.cluster.ClusterSampler
+      cropper:
+        _target_: boltzgen.data.crop.multimer.MultimerCropper
+        neighborhood_sizes: [ 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40 ]
+      split: ./model/boltzgen/resources/splits/validation_ids_boltz2_all.txt
+      symmetry_correction: false
+      val_group: "RCSB"
+
+  tokenizer:
+    _target_: boltzgen.data.tokenize.tokenizer.Tokenizer
+    atomize_modified_residues: false
+  featurizer:
+    _target_: boltzgen.data.feature.featurizer.Featurizer
+  moldir: ./training_data/mols
+  max_tokens: 512
+  max_atoms: 5120
+  max_seqs: 4096
+  pad_to_max_tokens: true
+  pad_to_max_atoms: true
+  pad_to_max_seqs: true
+  samples_per_epoch: 100000
+  batch_size: 1
+  num_workers: 2
+  random_seed: 42
+  pin_memory: false
+  overfit: null
+  return_train_symmetries: false
+  return_val_symmetries: false
+
+  
+  atoms_per_window_queries: 32
+  min_dist: 2.0
+  max_dist: 22.0
+  num_bins: 64
+  single_sequence_prop_training: 0.1
+  msa_sampling_training: true
+
+
+  # Design
+  design: true
+  backbone_only: false
+  atom14: true
+  atom37: false
+  selector:
+    _target_: boltzgen.data.select.protein.ProteinSelector
+    design_neighborhood_sizes: [2, 4, 6,8,10,12,14,16,18]
+    substructure_neighborhood_sizes: [2,4,6,8,10,12,24]
+    structure_condition_prob: 0.4
+    distance_noise_std:  1
+    run_selection: true
+    specify_binding_sites: true
+    ss_condition_prob: 0.1
+    select_all: false
+
+  # Design datasets
+  monomer_split: ./model/boltzgen/resources/splits/val_monomers_boltzgen_min50_max220.txt
+  monomer_target_dir: ./training_data/targets
+  monomer_target_structure_condition: true
+  monomer_seq_len: 100
+
+  ligand_split: ./model/boltzgen/resources/splits/val_ccd_pdb_pairs_boltzgen.txt
+  ligand_target_dir: ./training_data/targets
+  ligand_seq_len: 100
+
+
+model:
+  _target_: boltzgen.model.models.boltz.Boltz
+  atom_s: 128
+  atom_z: 16
+  token_s: 384
+  token_z: 128
+  num_bins: 64
+  atom_feature_dim: 388
+  atoms_per_window_queries: 32
+  atoms_per_window_keys: 128
+  use_miniformer: false
+  ema: true
+  ema_decay: 0.999
+  exclude_ions_from_lddt: true
+  num_val_datasets: 1               # New
+  ignore_ckpt_shape_mismatch: false # New
+  aggregate_distogram: true         # New
+  bond_type_feature: true
+  predict_bfactor: true
+  checkpoint_diffusion_conditioning: true
+  use_kernels: true
+
+
+  validators:
+    - _target_: boltzgen.model.validation.design.DesignValidator
+      val_names: ["RCSB"]
+      confidence_prediction: ${model.confidence_prediction}
+      atom14: ${data.atom14}
+      atom37: ${data.atom37}
+
+  masker_args:
+    mask: true
+    mask_backbone: false
+    mask_disto: true
+
+  embedder_args:
+    atom_encoder_depth: 3
+    atom_encoder_heads: 4
+    add_mol_type_feat: true
+    add_method_conditioning: true
+    add_modified_flag: true
+    add_cyclic_flag: true
+    add_design_mask_flag: true
+    add_binding_specification: true
+    add_ss_specification: true
+
+  freeze_template_weights: true
+  use_templates: true
+  template_args:
+    template_dim: 64
+    template_blocks: 2
+    activation_checkpointing: false
+
+
+  use_token_distances: true
+  token_distance_args:
+    token_distance_dim: 64
+    token_distance_blocks: 2
+    use_token_distance_feats: true
+    distance_gaussian_dim: 32
+    activation_checkpointing: true
+
+
+  msa_args:
+    msa_s: 64
+    msa_blocks: 4
+    msa_dropout: 0.15
+    z_dropout: 0.25
+    miniformer_blocks: false
+    pairwise_head_width: 32
+    pairwise_num_heads: 4
+    use_paired_feature: true
+    activation_checkpointing: true
+
+
+  pairformer_args:
+    num_blocks: 64
+    num_heads: 16
+    dropout: 0.25
+    post_layer_norm: false
+    activation_checkpointing: true
+
+
+  score_model_args:
+    sigma_data: 16
+    dim_fourier: 256
+    atom_encoder_depth: 3
+    atom_encoder_heads: 4
+
+    # token level args
+    token_layers: 1
+    token_transformer_depth: 24
+    token_transformer_heads: 16
+    diffusion_pairformer_args:
+      num_blocks: 0
+      num_heads: 2
+      dropout: 0
+      use_s_to_z: false
+
+
+
+    atom_decoder_depth: 3
+    atom_decoder_heads: 4
+    conditioning_transition_layers: 2
+    transformer_post_ln: false
+    activation_checkpointing: true
+
+  confidence_prediction: false
+  structure_prediction_training: true
+  
+  training_args:
+    recycling_steps: 3
+    sampling_steps: 20
+    diffusion_multiplicity: 32
+    diffusion_samples: 1
+    confidence_loss_weight: 1e-4
+    diffusion_loss_weight: 4.0
+    distogram_loss_weight: 3e-2
+    bfactor_loss_weight: 1e-3
+    adam_beta_1: 0.9
+    adam_beta_2: 0.95
+    adam_eps: 0.00000001
+    lr_scheduler: af3
+    base_lr: 0.0
+    max_lr: 0.0005
+    lr_warmup_no_steps: 1000
+    lr_start_decay_after_n_steps: 50000
+    lr_decay_every_n_steps: 50000
+    lr_decay_factor: 0.95
+    weight_decay: 0.003
+    weight_decay_exclude: true
+
+  validation_args:
+    recycling_steps: 3
+    sampling_steps: 200
+    diffusion_samples: 1
+    symmetry_correction: false
+
+  diffusion_process_args:
+    sigma_min: 0.0004  # min noise level
+    sigma_max: 160.0  # max noise level
+    sigma_data: 16.0  # standard deviation of data distribution
+    rho: 7  # controls the sampling schedule
+    P_mean: -1.2  # mean of log-normal distribution from which noise is drawn for training
+    P_std: 1.5  # standard deviation of log-normal distribution from which noise is drawn for training
+    gamma_0: 0.8
+    gamma_min: 1.0
+    noise_scale: 1.0
+    step_scale: 1.0
+    mse_rotational_alignment: true
+    coordinate_augmentation: true
+    alignment_reverse_diff: true
+    synchronize_sigmas: false
+
+  diffusion_loss_args:
+    add_smooth_lddt_loss: true
+    add_bond_loss: false
+    nucleotide_loss_weight: 5.0
+    ligand_loss_weight: 10.0
+
+  refolding_validator:
+    _target_: boltzgen.model.validation.refolding.RefoldingValidator
+    val_names: ["RCSB"]
+    step_scale: 1.5
+    noise_scale: 0.75
+    atom14: ${data.atom14}
+    atom37: ${data.atom37}
+    val_monomer: ${data.monomer_split}
+    val_ligand: ${data.ligand_split}
+    analyze_task:
+      _target_: boltzgen.task.analyze.analyze.Analyze
+      name: ${name}
+      debug: ${debug}
+      design_dir: null
+      num_processes: 1
+
+      # Common metrics to compute
+      affinity_metrics: false
+      allatom_fold_metrics: true
+      backbone_fold_metrics: true
+      noncovalents_original: false
+      noncovalents_refolded: false
+      delta_sasa_original: false
+      delta_sasa_refolded: false
+      largest_hydrophobic: false
+      largest_hydrophobic_refolded: false
+      run_clustering: false
+
+      # Liability analysis
+      liability_analysis: false
+      liability_modality: peptide
+      liability_peptide_type: linear
+
+      # Uncommon metrics
+      diversity_original: true
+      diversity_refolded: true
+      diversity_per_target_original: false
+      diversity_per_target_refolded: false
+      novelty_original: false
+      novelty_refolded: false
+      novelty_per_target_original: false
+      novelty_per_target_refolded: false
+
+      wandb: null
+
+      data:
+        _target_: boltzgen.task.predict.data_from_generated.FromGeneratedDataModule
+        cfg:
+          _target_: boltzgen.task.predict.data_from_generated.DataConfig
+          tokenizer:
+            _target_: boltzgen.data.tokenize.tokenizer.Tokenizer
+            atomize_modified_residues: false
+          featurizer:
+            _target_: boltzgen.data.feature.featurizer.Featurizer
+
+          suffix: .cif
+          suffix_metadata: .npz
+          suffix_native: _native.cif
+          samples_per_target: 1
+          num_targets: 100000000
+          moldir: ./training_data/mols
+
+          batch_size: 1
+          num_workers: 4
+          pin_memory: true
+        return_native: true
+
+    folding_checkpoint: ./training_data/boltz2_fold.ckpt
+
+    folding_args:
+      recycling_steps: 3
+      sampling_steps: 200
+      diffusion_samples: 1
+
+    folding_model_args:
+      atom_s: 128
+      atom_z: 16
+      token_s: 384
+      token_z: 128
+      num_bins: 64
+      atom_feature_dim: 388
+      atoms_per_window_queries: 32
+      atoms_per_window_keys: 128
+      compile_pairformer: false
+      compile_templates: false
+      compile_msa: false
+      use_miniformer: false
+      ema: true
+      ema_decay: 0.999
+      exclude_ions_from_lddt: true
+      num_val_datasets: 4
+      ignore_ckpt_shape_mismatch: false
+      aggregate_distogram: true
+      bond_type_feature: true
+      conditioning_cutoff_min: 4.0
+      conditioning_cutoff_max: 20.0
+      use_templates: true
+      predict_bfactor: true
+      checkpoint_diffusion_conditioning: false
+      use_kernels: true
+
+      validators: null
+
+      embedder_args:
+        atom_encoder_depth: 3
+        atom_encoder_heads: 4
+        add_mol_type_feat: true
+        add_method_conditioning: true
+        add_modified_flag: true
+        add_cyclic_flag: true
+
+      msa_args:
+        msa_s: 64
+        msa_blocks: 4
+        msa_dropout: 0.15
+        z_dropout: 0.25
+        miniformer_blocks: false
+        pairwise_head_width: 32
+        pairwise_num_heads: 4
+        use_paired_feature: true
+        activation_checkpointing: false
+
+
+      template_args:
+        template_dim: 64
+        template_blocks: 2
+        activation_checkpointing: false
+
+
+      pairformer_args:
+        num_blocks: 64
+        num_heads: 16
+        dropout: 0.25
+        post_layer_norm: false
+        activation_checkpointing: false
+
+
+      score_model_args:
+        sigma_data: 16
+        dim_fourier: 256
+        atom_encoder_depth: 3
+        atom_encoder_heads: 4
+        token_transformer_depth: 24
+        token_transformer_heads: 16
+        atom_decoder_depth: 3
+        atom_decoder_heads: 4
+        conditioning_transition_layers: 2
+        transformer_post_ln: false
+        activation_checkpointing: false
+
+      confidence_prediction: true
+      affinity_prediction: false
+      structure_prediction_training: true
+      affinity_model_args:
+        num_dist_bins: 64
+        max_dist: 22
+        no_trunk_feats: false
+        add_s_to_z_prod: false
+        add_s_input_to_s: false
+        confidence_args:
+          num_plddt_bins: 50
+          num_pde_bins: 64
+          num_pae_bins: 64
+
+      training_args:
+        recycling_steps: 3
+        sampling_steps: 20
+        diffusion_multiplicity: 48
+        diffusion_samples: 1
+        affinity_loss_weight: 3e-3
+        confidence_loss_weight: 1e-4
+        diffusion_loss_weight: 4.0
+        distogram_loss_weight: 3e-2
+        bfactor_loss_weight: 1e-3
+        adam_beta_1: 0.9
+        adam_beta_2: 0.95
+        adam_eps: 0.00000001
+        lr_scheduler: af3
+        base_lr: 0.0
+        max_lr: 0.001
+        lr_warmup_no_steps: 1000
+        lr_start_decay_after_n_steps: 50000
+        lr_decay_every_n_steps: 50000
+        lr_decay_factor: 0.95
+        weight_decay: 0.003
+        weight_decay_exclude: true
+
+      validation_args:
+        recycling_steps: 3
+        sampling_steps: 200
+        diffusion_samples: 5
+        symmetry_correction: false
+
+      diffusion_process_args:
+        sigma_min: 0.0004  # min noise level
+        sigma_max: 160.0  # max noise level
+        sigma_data: 16.0  # standard deviation of data distribution
+        rho: 7  # controls the sampling schedule
+        P_mean: -1.2  # mean of log-normal distribution from which noise is drawn for training
+        P_std: 1.5  # standard deviation of log-normal distribution from which noise is drawn for training
+        gamma_0: 0.8
+        gamma_min: 1.0
+        noise_scale: 1.0
+        step_scale: 1.0
+        mse_rotational_alignment: true
+        coordinate_augmentation: true
+        alignment_reverse_diff: true
+        synchronize_sigmas: false
+
+      diffusion_loss_args:
+        add_smooth_lddt_loss: true
+        add_bond_loss: false
+        nucleotide_loss_weight: 5.0
+        ligand_loss_weight: 10.0

model/boltzgen/resources/config/train/boltzgen.yaml

@@ -0,0 +1,578 @@
+_target_: boltzgen.task.train.train.Training
+
+trainer:
+  accelerator: gpu
+  devices: 8
+  precision: bf16-mixed
+  gradient_clip_val: 10.0
+  accumulate_grad_batches: 1
+  max_epochs: -1
+  num_sanity_val_steps: 3
+  log_every_n_steps: 1
+
+wandb:
+  group: boltzgen
+  project: boltzgen
+  entity: yourwandb
+
+name: a_big_run_resume3
+slurm: true
+output: workdir
+strict_loading: false
+resume: null
+pretrained: null
+debug: false
+save_every_n_train_steps: 2500
+disable_checkpoint: false
+matmul_precision: null
+save_top_k: -1
+
+data:
+  datasets:
+    - _target_: boltzgen.task.train.data.DatasetConfig
+      target_dir: ./training_data/targets
+      msa_dir: ./training_data/msa
+      prob: 0.6
+      filters:
+        - _target_: boltzgen.data.filter.dynamic.size.SizeFilter
+          min_chains: 1
+          max_chains: 300
+        - _target_: boltzgen.data.filter.dynamic.date.DateFilter
+          date: "2023-06-01"
+          ref: released
+        - _target_: boltzgen.data.filter.dynamic.resolution.ResolutionFilter
+          resolution: 9.0
+      sampler:
+        _target_: boltzgen.data.sample.cluster.ClusterSampler
+      cropper:
+        _target_: boltzgen.data.crop.multimer.MultimerCropper
+        neighborhood_sizes: [ 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40 ]
+      split: ./model/boltzgen/resources/splits/validation_ids_boltz2_all.txt
+      symmetry_correction: false
+      val_group: "RCSB"
+
+    # AFDB Distillation Data
+    - _target_: boltzgen.task.train.data.DatasetConfig
+      manifest_path: ./training_data/afdb/afdb_manifest_foldseek_c75_confidence.json
+      target_dir: ./training_data/afdb/targets
+      msa_dir: ./training_data/afdb/msa
+      prob: 0.3
+      filters:
+        - _target_: boltzgen.data.filter.dynamic.size.SizeFilter
+          min_chains: 1
+          max_chains: 300
+        - _target_: boltzgen.data.filter.dynamic.confidence.ConfidenceFilter
+          composition_op: "AND"
+          metrics: ["confidence_score"]
+          compare_ops: ["greater"]
+          thresholds: [70]
+      sampler:
+        _target_: boltzgen.data.sample.cluster.ClusterSampler
+      cropper:
+        _target_: boltzgen.data.crop.multimer.MultimerCropper
+        neighborhood_sizes: [ 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40 ]
+      symmetry_correction: true
+      override_method: "AFDB"
+      override_bfactor: true
+
+      # Protein-Ligand Distillation Data
+    - _target_: boltzgen.task.train.data.DatasetConfig
+      target_dir: ./training_data/protein_ligand/targets
+      msa_dir: ./training_data/protein_ligand/msa
+      moldir: ./training_data/protein_ligand/mols
+      prob: 0.03
+      filters:
+        - _target_: boltzgen.data.filter.dynamic.size.SizeFilter
+          min_chains: 1
+          max_chains: 300
+        - _target_: boltzgen.data.filter.dynamic.confidence.ConfidenceFilter
+          composition_op: "AND"
+          metrics: ["complex_ipde", "complex_pde", "iptm"]
+          compare_ops: ["lesser", "lesser", "greater"]
+          thresholds: [1.5, 1.5, 0.9]
+      sampler:
+        _target_: boltzgen.data.sample.cluster.ClusterSampler
+        beta_chain: 0.05
+      cropper:
+        _target_: boltzgen.data.crop.multimer.MultimerCropper
+        neighborhood_sizes: [ 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40 ]
+      symmetry_correction: true
+      override_method: "BOLTZ-1"
+
+    # RNA Distillation Data
+    - _target_: boltzgen.task.train.data.DatasetConfig
+      target_dir: ./training_data/rna/targets
+      msa_dir: ./training_data/rna/msa
+      prob: 0.04
+      filters:
+        - _target_: boltzgen.data.filter.dynamic.size.SizeFilter
+          min_chains: 1
+          max_chains: 300
+        - _target_: boltzgen.data.filter.dynamic.confidence.ConfidenceFilter
+          composition_op: "OR"
+          metrics: ["complex_pde"]
+          compare_ops: ["lesser"]
+          thresholds: [2.0]
+      sampler:
+        _target_: boltzgen.data.sample.cluster.ClusterSampler
+      cropper:
+        _target_: boltzgen.data.crop.multimer.MultimerCropper
+        neighborhood_sizes: [ 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40 ]
+      symmetry_correction: true
+      override_method: "BOLTZ-1"
+
+    # Protein-DNA Distillation Data
+    - _target_: boltzgen.task.train.data.DatasetConfig
+      target_dir: ./training_data/protein_dna/targets
+      msa_dir: ./training_data/protein_dna/msa
+      prob: 0.03
+      filters:
+        - _target_: boltzgen.data.filter.dynamic.size.SizeFilter
+          min_chains: 1
+          max_chains: 300
+        - _target_: boltzgen.data.filter.dynamic.confidence.ConfidenceFilter
+          composition_op: "AND"
+          metrics: ["complex_ipde", "complex_pde", "iptm"]
+          compare_ops: ["lesser", "lesser", "greater"]
+          thresholds: [1.0, 2.0, 0.7]
+      sampler:
+        _target_: boltzgen.data.sample.cluster.ClusterSampler
+        beta_chain: 0.05
+      cropper:
+        _target_: boltzgen.data.crop.multimer.MultimerCropper
+        neighborhood_sizes: [ 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40 ]
+      symmetry_correction: true
+      override_method: "BOLTZ-1"
+
+
+  tokenizer:
+    _target_: boltzgen.data.tokenize.tokenizer.Tokenizer
+    atomize_modified_residues: false
+  featurizer:
+    _target_: boltzgen.data.feature.featurizer.Featurizer
+  moldir: ./training_data/mols
+  max_tokens: 512
+  max_atoms: 5120
+  max_seqs: 4096
+  pad_to_max_tokens: true
+  pad_to_max_atoms: true
+  pad_to_max_seqs: true
+  samples_per_epoch: 100000
+  batch_size: 1
+  num_workers: 2
+  random_seed: 42
+  pin_memory: false
+  overfit: null
+  return_train_symmetries: false
+  return_val_symmetries: false
+
+  
+  atoms_per_window_queries: 32
+  min_dist: 2.0
+  max_dist: 22.0
+  num_bins: 64
+  single_sequence_prop_training: 0.1
+  msa_sampling_training: true
+
+
+  # Design
+  design: true
+  backbone_only: false
+  atom14: true
+  atom37: false
+  selector:
+    _target_: boltzgen.data.select.protein.ProteinSelector
+    design_neighborhood_sizes: [2, 4, 6,8,10,12,14,16,18]
+    substructure_neighborhood_sizes: [2,4,6,8,10,12,24]
+    structure_condition_prob: 0.4
+    distance_noise_std:  1
+    run_selection: true
+    specify_binding_sites: true
+    ss_condition_prob: 0.1
+    select_all: false
+
+  # Design datasets
+  monomer_split: ./model/boltzgen/resources/splits/val_monomers_boltzgen_min50_max220.txt
+  monomer_target_dir: ./training_data/targets
+  monomer_target_structure_condition: true
+  monomer_seq_len: 100
+
+  ligand_split: ./model/boltzgen/resources/splits/val_ccd_pdb_pairs_boltzgen.txt
+  ligand_target_dir: ./training_data/targets
+  ligand_seq_len: 100
+
+
+model:
+  _target_: boltzgen.model.models.boltz.Boltz
+  atom_s: 128
+  atom_z: 16
+  token_s: 384
+  token_z: 128
+  num_bins: 64
+  atom_feature_dim: 388
+  atoms_per_window_queries: 32
+  atoms_per_window_keys: 128
+  use_miniformer: false
+  ema: true
+  ema_decay: 0.999
+  exclude_ions_from_lddt: true
+  num_val_datasets: 1               # New
+  ignore_ckpt_shape_mismatch: false # New
+  aggregate_distogram: true         # New
+  bond_type_feature: true
+  predict_bfactor: true
+  checkpoint_diffusion_conditioning: true
+  use_kernels: true
+
+
+  validators:
+    - _target_: boltzgen.model.validation.design.DesignValidator
+      val_names: ["RCSB"]
+      confidence_prediction: ${model.confidence_prediction}
+      atom14: ${data.atom14}
+      atom37: ${data.atom37}
+
+  masker_args:
+    mask: true
+    mask_backbone: false
+    mask_disto: true
+
+  embedder_args:
+    atom_encoder_depth: 3
+    atom_encoder_heads: 4
+    add_mol_type_feat: true
+    add_method_conditioning: true
+    add_modified_flag: true
+    add_cyclic_flag: true
+    add_design_mask_flag: true
+    add_binding_specification: true
+    add_ss_specification: true
+
+  freeze_template_weights: true
+  use_templates: true
+  template_args:
+    template_dim: 64
+    template_blocks: 2
+    activation_checkpointing: false
+
+
+  use_token_distances: true
+  token_distance_args:
+    token_distance_dim: 64
+    token_distance_blocks: 2
+    use_token_distance_feats: true
+    distance_gaussian_dim: 32
+    activation_checkpointing: true
+
+
+  msa_args:
+    msa_s: 64
+    msa_blocks: 4
+    msa_dropout: 0.15
+    z_dropout: 0.25
+    miniformer_blocks: false
+    pairwise_head_width: 32
+    pairwise_num_heads: 4
+    use_paired_feature: true
+    activation_checkpointing: true
+
+
+  pairformer_args:
+    num_blocks: 64
+    num_heads: 16
+    dropout: 0.25
+    post_layer_norm: false
+    activation_checkpointing: true
+
+
+  score_model_args:
+    sigma_data: 16
+    dim_fourier: 256
+    atom_encoder_depth: 3
+    atom_encoder_heads: 4
+
+    # token level args
+    token_layers: 1
+    token_transformer_depth: 24
+    token_transformer_heads: 16
+    diffusion_pairformer_args:
+      num_blocks: 0
+      num_heads: 2
+      dropout: 0
+      use_s_to_z: false
+
+
+
+    atom_decoder_depth: 3
+    atom_decoder_heads: 4
+    conditioning_transition_layers: 2
+    transformer_post_ln: false
+    activation_checkpointing: true
+
+  confidence_prediction: false
+  structure_prediction_training: true
+  
+  training_args:
+    recycling_steps: 3
+    sampling_steps: 20
+    diffusion_multiplicity: 32
+    diffusion_samples: 1
+    confidence_loss_weight: 1e-4
+    diffusion_loss_weight: 4.0
+    distogram_loss_weight: 3e-2
+    bfactor_loss_weight: 1e-3
+    adam_beta_1: 0.9
+    adam_beta_2: 0.95
+    adam_eps: 0.00000001
+    lr_scheduler: af3
+    base_lr: 0.0
+    max_lr: 0.0005
+    lr_warmup_no_steps: 1000
+    lr_start_decay_after_n_steps: 50000
+    lr_decay_every_n_steps: 50000
+    lr_decay_factor: 0.95
+    weight_decay: 0.003
+    weight_decay_exclude: true
+
+  validation_args:
+    recycling_steps: 3
+    sampling_steps: 200
+    diffusion_samples: 1
+    symmetry_correction: false
+
+  diffusion_process_args:
+    sigma_min: 0.0004  # min noise level
+    sigma_max: 160.0  # max noise level
+    sigma_data: 16.0  # standard deviation of data distribution
+    rho: 7  # controls the sampling schedule
+    P_mean: -1.2  # mean of log-normal distribution from which noise is drawn for training
+    P_std: 1.5  # standard deviation of log-normal distribution from which noise is drawn for training
+    gamma_0: 0.8
+    gamma_min: 1.0
+    noise_scale: 1.0
+    step_scale: 1.0
+    mse_rotational_alignment: true
+    coordinate_augmentation: true
+    alignment_reverse_diff: true
+    synchronize_sigmas: false
+
+  diffusion_loss_args:
+    add_smooth_lddt_loss: true
+    add_bond_loss: false
+    nucleotide_loss_weight: 5.0
+    ligand_loss_weight: 10.0
+
+  refolding_validator:
+    _target_: boltzgen.model.validation.refolding.RefoldingValidator
+    val_names: ["RCSB"]
+    step_scale: 1.5
+    noise_scale: 0.75
+    atom14: ${data.atom14}
+    atom37: ${data.atom37}
+    val_monomer: ${data.monomer_split}
+    val_ligand: ${data.ligand_split}
+    analyze_task:
+      _target_: boltzgen.task.analyze.analyze.Analyze
+      name: ${name}
+      debug: ${debug}
+      design_dir: null
+      num_processes: 1
+
+      # Common metrics to compute
+      affinity_metrics: false
+      allatom_fold_metrics: true
+      backbone_fold_metrics: true
+      noncovalents_original: false
+      noncovalents_refolded: false
+      delta_sasa_original: false
+      delta_sasa_refolded: false
+      largest_hydrophobic: false
+      largest_hydrophobic_refolded: false
+      run_clustering: false
+
+      # Liability analysis
+      liability_analysis: false
+      liability_modality: peptide
+      liability_peptide_type: linear
+
+      # Uncommon metrics
+      diversity_original: true
+      diversity_refolded: true
+      diversity_per_target_original: false
+      diversity_per_target_refolded: false
+      novelty_original: false
+      novelty_refolded: false
+      novelty_per_target_original: false
+      novelty_per_target_refolded: false
+
+      wandb: null
+
+      data:
+        _target_: boltzgen.task.predict.data_from_generated.FromGeneratedDataModule
+        cfg:
+          _target_: boltzgen.task.predict.data_from_generated.DataConfig
+          tokenizer:
+            _target_: boltzgen.data.tokenize.tokenizer.Tokenizer
+            atomize_modified_residues: false
+          featurizer:
+            _target_: boltzgen.data.feature.featurizer.Featurizer
+
+          suffix: .cif
+          suffix_metadata: .npz
+          suffix_native: _native.cif
+          samples_per_target: 1
+          num_targets: 100000000
+          moldir: ./training_data/mols
+
+          batch_size: 1
+          num_workers: 4
+          pin_memory: true
+        return_native: true
+
+    folding_checkpoint: ./training_data/boltz2_fold.ckpt
+
+    folding_args:
+      recycling_steps: 3
+      sampling_steps: 200
+      diffusion_samples: 1
+
+    folding_model_args:
+      atom_s: 128
+      atom_z: 16
+      token_s: 384
+      token_z: 128
+      num_bins: 64
+      atom_feature_dim: 388
+      atoms_per_window_queries: 32
+      atoms_per_window_keys: 128
+      compile_pairformer: false
+      compile_templates: false
+      compile_msa: false
+      use_miniformer: false
+      ema: true
+      ema_decay: 0.999
+      exclude_ions_from_lddt: true
+      num_val_datasets: 4
+      ignore_ckpt_shape_mismatch: false
+      aggregate_distogram: true
+      bond_type_feature: true
+      conditioning_cutoff_min: 4.0
+      conditioning_cutoff_max: 20.0
+      use_templates: true
+      predict_bfactor: true
+      checkpoint_diffusion_conditioning: false
+      use_kernels: true
+
+      validators: null
+
+      embedder_args:
+        atom_encoder_depth: 3
+        atom_encoder_heads: 4
+        add_mol_type_feat: true
+        add_method_conditioning: true
+        add_modified_flag: true
+        add_cyclic_flag: true
+
+      msa_args:
+        msa_s: 64
+        msa_blocks: 4
+        msa_dropout: 0.15
+        z_dropout: 0.25
+        miniformer_blocks: false
+        pairwise_head_width: 32
+        pairwise_num_heads: 4
+        use_paired_feature: true
+        activation_checkpointing: false
+
+
+      template_args:
+        template_dim: 64
+        template_blocks: 2
+        activation_checkpointing: false
+
+
+      pairformer_args:
+        num_blocks: 64
+        num_heads: 16
+        dropout: 0.25
+        post_layer_norm: false
+        activation_checkpointing: false
+
+
+      score_model_args:
+        sigma_data: 16
+        dim_fourier: 256
+        atom_encoder_depth: 3
+        atom_encoder_heads: 4
+        token_transformer_depth: 24
+        token_transformer_heads: 16
+        atom_decoder_depth: 3
+        atom_decoder_heads: 4
+        conditioning_transition_layers: 2
+        transformer_post_ln: false
+        activation_checkpointing: false
+
+      confidence_prediction: false
+      affinity_prediction: false
+      structure_prediction_training: true
+      affinity_model_args:
+        num_dist_bins: 64
+        max_dist: 22
+        no_trunk_feats: false
+        add_s_to_z_prod: false
+        add_s_input_to_s: false
+        confidence_args:
+          num_plddt_bins: 50
+          num_pde_bins: 64
+          num_pae_bins: 64
+
+      training_args:
+        recycling_steps: 3
+        sampling_steps: 20
+        diffusion_multiplicity: 48
+        diffusion_samples: 1
+        affinity_loss_weight: 3e-3
+        confidence_loss_weight: 1e-4
+        diffusion_loss_weight: 4.0
+        distogram_loss_weight: 3e-2
+        bfactor_loss_weight: 1e-3
+        adam_beta_1: 0.9
+        adam_beta_2: 0.95
+        adam_eps: 0.00000001
+        lr_scheduler: af3
+        base_lr: 0.0
+        max_lr: 0.001
+        lr_warmup_no_steps: 1000
+        lr_start_decay_after_n_steps: 50000
+        lr_decay_every_n_steps: 50000
+        lr_decay_factor: 0.95
+        weight_decay: 0.003
+        weight_decay_exclude: true
+
+      validation_args:
+        recycling_steps: 3
+        sampling_steps: 200
+        diffusion_samples: 5
+        symmetry_correction: false
+
+      diffusion_process_args:
+        sigma_min: 0.0004  # min noise level
+        sigma_max: 160.0  # max noise level
+        sigma_data: 16.0  # standard deviation of data distribution
+        rho: 7  # controls the sampling schedule
+        P_mean: -1.2  # mean of log-normal distribution from which noise is drawn for training
+        P_std: 1.5  # standard deviation of log-normal distribution from which noise is drawn for training
+        gamma_0: 0.8
+        gamma_min: 1.0
+        noise_scale: 1.0
+        step_scale: 1.0
+        mse_rotational_alignment: true
+        coordinate_augmentation: true
+        alignment_reverse_diff: true
+        synchronize_sigmas: false
+
+      diffusion_loss_args:
+        add_smooth_lddt_loss: true
+        add_bond_loss: false
+        nucleotide_loss_weight: 5.0
+        ligand_loss_weight: 10.0

model/boltzgen/resources/config/train/boltzgen_small.yaml

@@ -0,0 +1,346 @@
+_target_: boltzgen.task.train.train.Training
+
+trainer:
+  accelerator: gpu
+  devices: 8
+  precision: bf16-mixed
+  gradient_clip_val: 10.0
+  accumulate_grad_batches: 16
+  max_epochs: -1
+  num_sanity_val_steps: 3
+  log_every_n_steps: 1
+
+wandb:
+  group: boltzgen
+  project: boltzgen
+  entity: yourwandb
+
+name: small
+output: workdir
+strict_loading: false
+resume: null
+pretrained: ./training_data/boltzgen1_structuretrained_small.ckpt
+debug: false
+save_every_n_train_steps: 2500
+disable_checkpoint: false
+matmul_precision: null
+save_top_k: -1
+# ddp_timeout_seconds: 1000
+
+data:
+  datasets:
+    - _target_: boltzgen.task.train.data.DatasetConfig
+      target_dir: ./training_data/targets
+      msa_dir: ./training_data/msa
+      prob: 1
+      filters:
+        - _target_: boltzgen.data.filter.dynamic.size.SizeFilter
+          min_chains: 1
+          max_chains: 300
+        - _target_: boltzgen.data.filter.dynamic.date.DateFilter
+          date: "2023-06-01"
+          ref: released
+        - _target_: boltzgen.data.filter.dynamic.resolution.ResolutionFilter
+          resolution: 9.0
+      sampler:
+        _target_: boltzgen.data.sample.cluster.ClusterSampler
+      cropper:
+        _target_: boltzgen.data.crop.multimer.MultimerCropper
+        neighborhood_sizes: [ 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40 ]
+      split: ./model/boltzgen/resources/splits/validation_ids_boltz2_all.txt
+      symmetry_correction: false
+      val_group: "RCSB"
+
+
+  tokenizer:
+    _target_: boltzgen.data.tokenize.tokenizer.Tokenizer
+    atomize_modified_residues: false
+  featurizer:
+    _target_: boltzgen.data.feature.featurizer.Featurizer
+  moldir: ./training_data/mols
+  max_tokens: 256
+  max_atoms: 2048
+  max_seqs: 1024
+  pad_to_max_tokens: true
+  pad_to_max_atoms: true
+  pad_to_max_seqs: true
+  samples_per_epoch: 100000
+  batch_size: 1
+  num_workers: 4
+  random_seed: 42
+  pin_memory: true
+  overfit: null
+  return_train_symmetries: false
+  return_val_symmetries: false
+
+  
+  atoms_per_window_queries: 32
+  min_dist: 2.0
+  max_dist: 22.0
+  num_bins: 64
+  single_sequence_prop_training: 0.1
+  msa_sampling_training: true
+
+
+  # Design
+  design: true
+  backbone_only: false
+  atom14: true
+  atom37: false
+  selector:
+    _target_: boltzgen.data.select.protein.ProteinSelector
+    design_neighborhood_sizes: [2, 4, 6,8,10,12,14,16,18]
+    substructure_neighborhood_sizes: [2,4,6,8,10,12,24]
+    structure_condition_prob: 0.4
+    distance_noise_std:  1
+    run_selection: true
+    specify_binding_sites: true
+    ss_condition_prob: 0.1
+    select_all: false
+    chain_reindexing: false
+
+  # Design datasets
+  monomer_split: ./model/boltzgen/resources/splits/val_monomers_boltzgen_min50_max220.txt
+  monomer_target_dir: ./training_data/targets
+  monomer_target_structure_condition: true
+  monomer_seq_len: 100
+
+  ligand_split: ./model/boltzgen/resources/splits/val_ccd_pdb_pairs_boltzgen.txt
+  ligand_target_dir: ./training_data/targets
+  ligand_seq_len: 100
+
+
+model:
+  _target_: boltzgen.model.models.boltz.Boltz
+  atom_s: 128
+  atom_z: 16
+  token_s: 384
+  token_z: 128
+  num_bins: 64
+  atom_feature_dim: 388
+  atoms_per_window_queries: 32
+  atoms_per_window_keys: 128
+  use_miniformer: true
+  ema: true
+  ema_decay: 0.999
+  exclude_ions_from_lddt: true
+  num_val_datasets: 1               # New
+  ignore_ckpt_shape_mismatch: false # New
+  aggregate_distogram: true         # New
+  bond_type_feature: true
+  predict_bfactor: true
+  predict_res_type: false
+  checkpoint_diffusion_conditioning: false
+  use_kernels: true
+
+
+  validators:
+    - _target_: boltzgen.model.validation.design.DesignValidator
+      val_names: ["RCSB"]
+      confidence_prediction: ${model.confidence_prediction}
+      backbone_only: ${data.backbone_only}
+      atom14: ${data.atom14}
+      atom37: ${data.atom37}
+
+  masker_args:
+    mask: true
+    mask_backbone: false
+    mask_disto: true
+
+  embedder_args:
+    atom_encoder_depth: 3
+    atom_encoder_heads: 4
+    add_mol_type_feat: true
+    add_method_conditioning: true
+    add_modified_flag: true
+    add_cyclic_flag: true
+    add_design_mask_flag: true
+    add_binding_specification: true
+    add_ss_specification: true
+
+  freeze_template_weights: true
+  use_templates: true
+  template_args:
+    template_dim: 64
+    template_blocks: 2
+    miniformer_blocks: true
+    activation_checkpointing: false
+
+  use_token_distances: true
+  token_distance_args:
+    token_distance_dim: 64
+    token_distance_blocks: 2
+    use_token_distance_feats: true
+    distance_gaussian_dim: 32
+
+  msa_args:
+    msa_s: 64
+    msa_blocks: 3
+    msa_dropout: 0.15
+    z_dropout: 0.25
+    miniformer_blocks: true
+    pairwise_head_width: 32
+    pairwise_num_heads: 4
+    use_paired_feature: true
+    activation_checkpointing: false
+
+  pairformer_args:
+    num_blocks: 12
+    num_heads: 16
+    dropout: 0.25
+    post_layer_norm: false
+    activation_checkpointing: false
+
+  score_model_args:
+    sigma_data: 16
+    dim_fourier: 256
+    atom_encoder_depth: 3
+    atom_encoder_heads: 4
+
+    # token level args
+    token_layers: 1
+    token_transformer_depth: 8
+    token_transformer_heads: 16
+    diffusion_pairformer_args:
+      num_blocks: 0
+      num_heads: 2
+      dropout: 0
+      use_s_to_z: false
+
+    atom_decoder_depth: 3
+    atom_decoder_heads: 4
+    conditioning_transition_layers: 2
+    transformer_post_ln: false
+    activation_checkpointing: false
+
+  confidence_prediction: false
+  structure_prediction_training: true
+  training_args:
+    recycling_steps: 3
+    sampling_steps: 20
+    diffusion_multiplicity: 12
+    diffusion_samples: 1
+    confidence_loss_weight: 1e-4
+    diffusion_loss_weight: 4.0
+    distogram_loss_weight: 3e-2
+    bfactor_loss_weight: 1e-3
+    res_type_loss_weight: 3e-2
+    adam_beta_1: 0.9
+    adam_beta_2: 0.95
+    adam_eps: 0.00000001
+    lr_scheduler: af3
+    base_lr: 0.0
+    max_lr: 0.0018
+    lr_warmup_no_steps: 1000
+    lr_start_decay_after_n_steps: 50000
+    lr_decay_every_n_steps: 50000
+    lr_decay_factor: 0.95
+    weight_decay: 0.003
+    weight_decay_exclude: true
+
+  validation_args:
+    recycling_steps: 3
+    sampling_steps: 200
+    diffusion_samples: 1
+    symmetry_correction: false
+
+  diffusion_process_args:
+    sigma_min: 0.0004  # min noise level
+    sigma_max: 160.0  # max noise level
+    sigma_data: 16.0  # standard deviation of data distribution
+    rho: 7  # controls the sampling schedule
+    P_mean: -1.2  # mean of log-normal distribution from which noise is drawn for training
+    P_std: 1.5  # standard deviation of log-normal distribution from which noise is drawn for training
+    gamma_0: 0.8
+    gamma_min: 1.0
+    noise_scale: 1.0
+    step_scale: 1.0
+    mse_rotational_alignment: true
+    coordinate_augmentation: true
+    alignment_reverse_diff: true
+    synchronize_sigmas: false
+
+  diffusion_loss_args:
+    add_smooth_lddt_loss: true
+    add_bond_loss: false
+    nucleotide_loss_weight: 5.0
+    ligand_loss_weight: 10.0
+
+  refolding_validator:
+    _target_: boltzgen.model.validation.refolding.RefoldingValidator
+    val_names: ["RCSB"]
+    step_scale: 1.5
+    noise_scale: 0.75
+    atom14: ${data.atom14}
+    atom37: ${data.atom37}
+    backbone_only: ${data.backbone_only}
+    val_monomer: ${data.monomer_split}
+    val_ligand: ${data.ligand_split}
+    analyze_task:
+      _target_: boltzgen.task.analyze.analyze.Analyze
+      name: ${name}
+      debug: ${debug}
+      design_dir: null
+      num_processes: 1
+
+      # Common metrics to compute
+      affinity_metrics: false
+      allatom_fold_metrics: true
+      backbone_fold_metrics: true
+      noncovalents_original: false
+      noncovalents_refolded: false
+      delta_sasa_original: false
+      delta_sasa_refolded: false
+      largest_hydrophobic: false
+      largest_hydrophobic_refolded: false
+      run_clustering: false
+
+      # Liability analysis
+      liability_analysis: false
+      liability_modality: peptide
+      liability_peptide_type: linear
+
+      # Uncommon metrics
+      diversity_original: true
+      diversity_refolded: true
+      diversity_per_target_original: false
+      diversity_per_target_refolded: false
+      novelty_original: false
+      novelty_refolded: false
+      novelty_per_target_original: false
+      novelty_per_target_refolded: false
+
+      wandb: null
+
+      data:
+        _target_: boltzgen.task.predict.data_from_generated.FromGeneratedDataModule
+        cfg:
+          _target_: boltzgen.task.predict.data_from_generated.DataConfig
+          tokenizer:
+            _target_: boltzgen.data.tokenize.tokenizer.Tokenizer
+            atomize_modified_residues: false
+          featurizer:
+            _target_: boltzgen.data.feature.featurizer.Featurizer
+
+          suffix: .cif
+          suffix_metadata: .npz
+          suffix_native: _native.cif
+          samples_per_target: 1
+          num_targets: 100000000
+          moldir: ./training_data/mols
+
+          batch_size: 1
+          num_workers: 4
+          pin_memory: true
+        target_templates: true
+        return_native: true
+
+    folding_checkpoint: ./training_data/boltz2_fold.ckpt
+
+    folding_args:
+      recycling_steps: 3
+      sampling_steps: 200
+      diffusion_samples: 1
+
+    folding_model_args:
+      validators: null

model/boltzgen/resources/config/train/inverse_folding.yaml

@@ -0,0 +1,376 @@
+_target_: boltzgen.task.train.train.Training
+
+trainer:
+  accelerator: cuda
+  devices: 4
+  precision: 32
+  gradient_clip_val: 10.0
+  accumulate_grad_batches: 1
+  max_epochs: 5
+  num_sanity_val_steps: 1
+  log_every_n_steps: 1
+
+wandb:
+  group: boltzgen
+  project: boltzgen
+  entity: yourwandb
+
+name: if_lr_scheduler
+output: workdir
+strict_loading: false
+resume: null
+debug: false
+save_every_n_train_steps: 2500
+disable_checkpoint: false
+matmul_precision: null
+save_top_k: -1
+
+data:
+  datasets:
+    - _target_: boltzgen.task.train.data.DatasetConfig
+      target_dir: ./training_data/targets
+      msa_dir: ./training_data/msa
+      prob: 1
+      filters:
+        - _target_: boltzgen.data.filter.dynamic.size.SizeFilter
+          min_chains: 1
+          max_chains: 300
+        - _target_: boltzgen.data.filter.dynamic.date.DateFilter
+          date: "2023-06-01"
+          ref: released
+        - _target_: boltzgen.data.filter.dynamic.resolution.ResolutionFilter
+          resolution: 9.0
+        - _target_: boltzgen.data.filter.dynamic.min_protein_residues.MinProteinResiduesFilter
+          min_residues: 5
+        - _target_: boltzgen.data.filter.dynamic.pdb_id_txtfile.FilterIDFromTXT
+          paths:
+            - data/exclude_ids/fibril.txt
+            - data/exclude_ids/transmembrane.txt
+      sampler:
+        _target_: boltzgen.data.sample.cluster.ClusterSampler
+      cropper:
+        _target_: boltzgen.data.crop.multimer.MultimerCropper
+        neighborhood_sizes: [ 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40 ]
+      split: ./model/boltzgen/resources/splits/validation_ids_boltz2_all.txt
+      symmetry_correction: false
+      val_group: "RCSB"
+
+
+  tokenizer:
+    _target_: boltzgen.data.tokenize.tokenizer.Tokenizer
+    atomize_modified_residues: false
+  featurizer:
+    _target_: boltzgen.data.feature.featurizer.Featurizer
+  moldir: ./training_data/mols
+  max_tokens: 1024 
+  max_atoms: 8192
+  max_seqs: 1
+  pad_to_max_tokens: true
+  pad_to_max_atoms: true
+  pad_to_max_seqs: true
+  samples_per_epoch: 600000
+  batch_size: 2
+  num_workers: 32
+  random_seed: 42
+  pin_memory: true
+  overfit: null
+  return_train_symmetries: false
+  return_val_symmetries: false
+  compute_frames: false
+
+  
+  atoms_per_window_queries: 32
+  min_dist: 2.0
+  max_dist: 22.0
+  num_bins: 64
+  single_sequence_prop_training: 0.05
+  msa_sampling_training: true
+
+  # Design
+  design: true
+  backbone_only: true
+  atom14: false
+  atom37: false
+  inverse_fold: ${model.inverse_fold}
+  use_msa: false
+  selector:
+    _target_: boltzgen.data.select.protein.ProteinSelector
+    design_neighborhood_sizes: [2, 4, 6,8,10,12,14,16,18]
+    substructure_neighborhood_sizes: [2,4,6,8,10,12,24]
+    structure_condition_prob: 0.5
+    distance_noise_std:  1
+    run_selection: true
+    specify_binding_sites: false
+    ss_condition_prob: 0
+    select_all: true
+    complete_structure_mask: true
+
+  # Design datasets
+  monomer_split: ./model/boltzgen/resources/splits/val_monomers_boltzgen_min50_max220.txt
+  monomer_target_dir: ./training_data/targets
+  monomer_target_structure_condition: true
+  monomer_seq_len: 100
+
+  ligand_split: null
+  ligand_target_dir: ./training_data/targets
+  ligand_seq_len: 100
+
+
+model:
+  _target_: boltzgen.model.models.boltz.Boltz
+  atom_s: 128
+  atom_z: 16
+  token_s: 384
+  token_z: ${model.inverse_fold_args.pair_dim}
+  num_bins: 64
+  atom_feature_dim: 388
+  atoms_per_window_queries: 32
+  atoms_per_window_keys: 128
+  use_miniformer: true
+  ema: true
+  ema_decay: 0.999
+  exclude_ions_from_lddt: true
+  num_val_datasets: 1               # New
+  ignore_ckpt_shape_mismatch: false # New
+  aggregate_distogram: true         # New
+  bond_type_feature: true
+  predict_bfactor: true
+  predict_res_type: true
+  checkpoint_diffusion_conditioning: false
+  inverse_fold: true
+  inverse_fold_args:
+    atom_s: ${model.atom_s}
+    atom_z: ${model.atom_z}
+    token_s: ${model.token_s}
+    token_z: ${model.token_z}
+    node_dim: 128
+    pair_dim: 128
+    hidden_dim: 128
+    dropout: 0.1
+    softmax_dropout: 0.2
+    num_encoder_layers: 6
+    num_decoder_layers: 3
+    autoregressive: true
+    transformation_scale_factor: 1.0
+    inverse_fold_noise: 0.2
+    topk: 30
+    num_heads: 4
+    enable_input_embedder: True
+    sampling_temperature: -1.0
+
+  validators:
+    - _target_: boltzgen.model.validation.design.DesignValidator
+      val_names: ["RCSB"]
+      confidence_prediction: ${model.confidence_prediction}
+      atom14: ${data.atom14}
+      atom37: ${data.atom37}
+      backbone_only: ${data.backbone_only}
+      inverse_fold: ${model.inverse_fold}
+
+  masker_args:
+    mask: true
+    mask_backbone: false
+    mask_disto: false
+
+  embedder_args:
+    atom_encoder_depth: 1
+    atom_encoder_heads: 4
+    add_mol_type_feat: true
+    add_method_conditioning: true
+    add_modified_flag: true
+    add_cyclic_flag: true
+    add_design_mask_flag: false
+    add_binding_specification: false
+    add_ss_specification: false
+
+  use_token_distances: false
+  token_distance_args:
+    token_distance_dim: ${model.inverse_fold_args.pair_dim}
+    token_distance_blocks: 0
+    use_token_distance_feats: true
+    distance_gaussian_dim: 32
+    disable_token_distance_transition: true
+    use_relative_position_encoding: true 
+
+  # MSA module is not used in inverse folding
+  msa_args:
+    msa_s: 2
+    msa_blocks: 0
+    msa_dropout: 0
+    z_dropout: 0
+    miniformer_blocks: true
+    pairwise_head_width: 2
+    pairwise_num_heads: 1
+    use_paired_feature: true
+    activation_checkpointing: false
+
+  pairformer_args:
+    num_blocks: 2
+    num_heads: 16
+    dropout: 0.25
+    post_layer_norm: false
+    activation_checkpointing: false
+
+  score_model_args:
+    sigma_data: 16
+    dim_fourier: 256
+    atom_encoder_depth: 3
+    atom_encoder_heads: 4
+
+    # token level args
+    token_layers: 1
+    token_transformer_depth: 3
+    token_transformer_heads: 16
+    diffusion_pairformer_args:
+      num_blocks: 0
+      num_heads: 2
+      dropout: 0
+      use_s_to_z: false
+
+    atom_decoder_depth: 3
+    atom_decoder_heads: 4
+    conditioning_transition_layers: 2
+    transformer_post_ln: false
+    activation_checkpointing: false
+
+  confidence_prediction: false
+  affinity_prediction: false
+  structure_prediction_training: true
+  affinity_model_args:
+    num_dist_bins: 64
+    max_dist: 22
+    no_trunk_feats: false
+    add_s_to_z_prod: false
+    add_s_input_to_s: false
+
+    confidence_args:
+      num_plddt_bins: 50
+      num_pde_bins: 64
+      num_pae_bins: 64
+
+  training_args:
+    recycling_steps: 0
+    sampling_steps: 20
+    diffusion_multiplicity: 2
+    diffusion_samples: 1
+    affinity_loss_weight: 3e-3
+    confidence_loss_weight: 1e-4
+    diffusion_loss_weight: 4.0
+    distogram_loss_weight: 3e-2
+    bfactor_loss_weight: 1e-3
+    res_type_loss_weight: 1
+    adam_beta_1: 0.9
+    adam_beta_2: 0.95 
+    adam_eps: 0.00000001
+    lr_scheduler: onecycle
+    base_lr: 0.0
+    max_lr: 0.001
+    weight_decay: 0.003
+    weight_decay_exclude: true 
+
+  validation_args:
+    recycling_steps: 0
+    sampling_steps: 200
+    diffusion_samples: 1
+    symmetry_correction: false
+
+  diffusion_process_args:
+    sigma_min: 0.0004  # min noise level
+    sigma_max: 160.0  # max noise level
+    sigma_data: 16.0  # standard deviation of data distribution
+    rho: 7  # controls the sampling schedule
+    P_mean: -1.2  # mean of log-normal distribution from which noise is drawn for training
+    P_std: 1.5  # standard deviation of log-normal distribution from which noise is drawn for training
+    gamma_0: 0.8
+    gamma_min: 1.0
+    noise_scale: 1.0
+    step_scale: 1.0
+    mse_rotational_alignment: true
+    coordinate_augmentation: true
+    alignment_reverse_diff: true
+    synchronize_sigmas: false
+
+  diffusion_loss_args:
+    add_smooth_lddt_loss: true
+    add_bond_loss: false
+    nucleotide_loss_weight: 5.0
+    ligand_loss_weight: 10.0
+
+  refolding_validator:
+    _target_: boltzgen.model.validation.refolding.RefoldingValidator
+    val_names: ["RCSB"]
+    step_scale: 1.5
+    noise_scale: 0.75
+    atom14: ${data.atom14}
+    atom37: ${data.atom37}
+    val_monomer: ${data.monomer_split}
+    val_ligand: ${data.ligand_split}
+    inverse_fold: ${model.inverse_fold}
+    analyze_task:
+      _target_: boltzgen.task.analyze.analyze.Analyze
+      name: ${name}
+      debug: ${debug}
+      design_dir: null
+      num_processes: 1
+
+      # Common metrics to compute
+      affinity_metrics: false
+      allatom_fold_metrics: true
+      backbone_fold_metrics: true
+      noncovalents_original: false
+      noncovalents_refolded: false
+      delta_sasa_original: false
+      delta_sasa_refolded: false
+      largest_hydrophobic: false
+      largest_hydrophobic_refolded: false
+      run_clustering: false
+
+      # Liability analysis
+      liability_analysis: false
+      liability_modality: peptide
+      liability_peptide_type: linear
+
+      # Uncommon metrics
+      diversity_original: true
+      diversity_refolded: true
+      diversity_per_target_original: false
+      diversity_per_target_refolded: false
+      novelty_original: false
+      novelty_refolded: false
+      novelty_per_target_original: false
+      novelty_per_target_refolded: false
+
+      wandb: null
+
+      data:
+        _target_: boltzgen.task.predict.data_from_generated.FromGeneratedDataModule
+        cfg:
+          _target_: boltzgen.task.predict.data_from_generated.DataConfig
+          tokenizer:
+            _target_: boltzgen.data.tokenize.tokenizer.Tokenizer
+            atomize_modified_residues: false
+          featurizer:
+            _target_: boltzgen.data.feature.featurizer.Featurizer
+
+          suffix: .cif
+          suffix_metadata: .npz
+          suffix_native: _native.cif
+          samples_per_target: 1
+          num_targets: 100000000
+          moldir: ./training_data/mols
+
+          batch_size: 1
+          num_workers: 1
+          pin_memory: false
+        target_templates: true
+        return_native: true
+
+    folding_checkpoint: ./training_data/boltz2_fold.ckpt
+
+    folding_args:
+      recycling_steps: 3
+      sampling_steps: 200
+      diffusion_samples: 1
+
+    folding_model_args:
+      validators: null

model/boltzgen/resources/main.py

@@ -0,0 +1,53 @@
+import sys
+from typing import List
+import hydra
+import omegaconf
+
+
+from boltzgen.task.task import Task
+
+
+def main(config: str, args: List) -> None:
+    """
+    This is just a wrapper for running the .run() function of our `Task` class.
+    If you run the pipeline (for example via `boltzgen run design_spec.yaml ...`) then this function reads the yaml files of the individual pipeline steps and executes the pipeline steps.
+
+    The possible tasks are:
+        - Train (GPU: BoltzGen diffusion model or inverse folding model training)
+        - Predict (GPU: Running BoltzGen diffusion, inverse folding, refolding, designfolding, or affinity prediction)
+        - Analyze (CPU: Compute CPU Metrics and aggregate metrics from GPU steps)
+        - Filter (CPU: Very fast (20s) computes ranking and writes final output files)
+
+    The files for these are:
+        - src/boltzgen/task/train/train.py
+        - src/boltzgen/task/predict/predict.py
+        - src/boltzgen/task/analyze/analyze.py
+        - src/boltzgen/task/filter/filter.py
+
+    Parameters
+    ----------
+    config : str
+        Path to the configuration yaml file. The yaml file contains something like `_target_: boltzgen.task.predict.predict.Predict` at the beginning which tells it which Task class to run
+    args : List
+        List of arguments to override the configuration.
+    """
+    # Load the configuration
+    args = omegaconf.OmegaConf.from_dotlist(args)
+    config = omegaconf.OmegaConf.load(config)
+    config = omegaconf.OmegaConf.merge(config, args)
+
+    # Instantiate the task
+    task = hydra.utils.instantiate(config)
+
+    if not isinstance(task, Task):
+        msg = "Config must be an instance of Task."
+        raise TypeError(msg)
+
+    # Run the task
+    task.run(config)
+
+
+if __name__ == "__main__":
+    config = sys.argv[1]
+    args = sys.argv[2:]
+    main(config, args)

model/boltzgen/resources/metrics_normalization.json

@@ -0,0 +1,38 @@
+{
+    "design_iiptm": {
+        "mean": 0.3828427453913554,
+        "std": 0.10621994686776445
+    },
+    "design_ptm": {
+        "mean": 0.6548983461494845,
+        "std": 0.041925380367595154
+    },
+    "min_design_to_target_pae": {
+        "mean": 8.665374909044512,
+        "std": 3.8643855151416453
+    },
+    "design_hydrophobicity": {
+        "mean": 43.95844996189343,
+        "std": 7.565721364606706
+    },
+    "design_largest_hydrophobic_patch_refolded": {
+        "mean": 608.7238163271427,
+        "std": 440.3199461402019
+    },
+    "delta_sasa_refolded": {
+        "mean": 1368.936707104108,
+        "std": 514.8573173243741
+    },
+    "plip_saltbridge_refolded": {
+        "mean": 1.746922289263117,
+        "std": 1.9177749177139145
+    },
+    "plip_hbonds_refolded": {
+        "mean": 7.099706654392513,
+        "std": 4.456755351208993
+    },
+    "affinity_probability_binary1": {
+        "mean": 0.4726671722992441,
+        "std": 0.19092615494037402
+    }
+}

model/boltzgen/resources/splits/val_ccd_pdb_pairs_boltzgen.txt

@@ -0,0 +1,101 @@
+WF6_7FZC
+WHI_7FZP
+LM9_7FWB
+WON_7G1L
+ZI8_8SHJ
+SWO_8J19
+SVI_7GO2
+IQH_7FVW
+QKO_8BOE
+WQC_7G1Y
+TQO_7FXR
+IUX_7ZG3
+VMR_8OG7
+LKX_7GCF
+W9K_7FYS
+LS9_8ADR
+5IY_8HXO
+ZV8_7G9Y
+JGX_7YMJ
+VI0_7FY6
+WPZ_7G1T
+TE0_7GP7
+WQH_7G1Z
+QLL_7FX7
+WEN_7FZB
+Q0O_7FWU
+OT9_7FWP
+ZDF_8GEM
+5YT_8HXP
+UTO_8CJ7
+S8X_7FXK
+Q63_5SO3
+U0O_7GQN
+Z4I_8GD4
+YXW_8G89
+ZCF_8GEU
+60J_8HYD
+5Z2_8HXU
+NOI_7GEK
+IWI_7FVY
+XE3_8WSM
+5ZX_8JAO
+ZQ2_8PH7
+WYW_8H7N
+WYO_7G15
+TJA_7GPN
+P6X_7FWR
+ZUG_7G9Q
+X7C_8P8K
+5ZU_8HY1
+P9X_8IHK
+SE1_8DJ3
+TZ9_7GO5
+TQC_7GOJ
+QXF_8BJK
+VJ3_7FYC
+OUF_7GGD
+W7U_7FYI
+SY0_8C11
+O6Q_7FWL
+O9U_7FWM
+M1U_8AGP
+R80_7FXC
+IKQ_7G21
+XNO_8FX5
+ZDK_8GEV
+WMQ_7G17
+NGX_7GED
+OD3_8I67
+LO0_7GCO
+KOX_8A0S
+WGW_7FZM
+YJN_8G3C
+WKN_7G0L
+Z5H_8GD2
+WHN_7FZR
+TJL_7GPE
+UOF_7G0K
+NQ3_7GEL
+NRX_7GEO
+WEA_7FZ6
+UOR_7FY2
+5ZD_8HXV
+ST0_7GNX
+VVL_7FYG
+U7F_8CC4
+NT6_7UMO
+MWR_7FWG
+W9W_7FYX
+WMF_7G10
+ITF_8HOC
+TZJ_7FXT
+Y2S_8E1O
+W8C_7FYK
+QG0_7FX0
+YW6_8PIF
+WYN_8H7I
+QBX_7FWY
+U2I_8DII
+R9O_7FXF
+Q7E_5SOA

model/boltzgen/resources/splits/val_monomers_boltzgen_min50_max220.txt

@@ -0,0 +1,110 @@
+8GM6
+8P35
+7Z3G
+8T9N
+8F8N
+8B6E
+8TGB
+8J9Q
+8P5N
+8BSA
+7Z3F
+8B4C
+8H0L
+8GM7
+8CGM
+8BAU
+7YSF
+8J0H
+8OXH
+8R8C
+8CM1
+8BL3
+8OQH
+8EY3
+8JNG
+8JN0
+8EV5
+7YD4
+8B2E
+8BAS
+8BBZ
+8BAQ
+8JNM
+8OYW
+8SYJ
+7Y5M
+7QRL
+7Y19
+7Z3B
+8GSY
+8AN4
+8UK7
+8GN3
+8V1K
+8OXL
+8T8K
+8OYV
+8T8L
+8B2G
+7WLH
+8HNA
+8TFS
+8FJG
+8JU8
+8OK3
+8P2A
+7WU8
+7XP9
+8J98
+8T4C
+8EOV
+8GJ9
+8OVU
+7TAL
+7VS2
+8OXJ
+8AH9
+8BL6
+8P5P
+8P4Y
+8BSB
+8EB9
+8SXS
+8GYR
+8P2B
+7YKM
+8OYS
+8JNL
+8T5J
+8VEH
+8OXK
+8TNM
+8BGB
+8EV4
+8BAR
+8EHN
+8J1X
+8FJF
+8HJJ
+8KCA
+8GN4
+7Z3I
+8F5G
+8G64
+8R8A
+8SOT
+8G8K
+8INI
+8J1W
+8JWS
+8U0X
+8GMX
+8J9R
+8AXJ
+8EC3
+8JWD
+8FR5
+8GTH
+8BHU
+8FZY

model/boltzgen/resources/splits/validation_ids_boltz2_all.txt

@@ -0,0 +1,398 @@
+8Q41
+8BH8
+8BH9
+8HIG
+8IA3
+8SVD
+7ZLQ
+8TPK
+8Q3Z
+8Q40
+8K3D
+8SVA
+8D9I
+8SSU
+8PE3
+8GN3
+8GN4
+8K4L
+8EB5
+8B4E
+8Q43
+8H0L
+8TP8
+8SSQ
+8F5G
+8SQU
+8DWJ
+8Q42
+8D9E
+8Q44
+8ISZ
+7XF1
+8HML
+8SSR
+8B4D
+8B4C
+7YSF
+8J0K
+8J0R
+8GBA
+8GTY
+8GBM
+7YL4
+8ANU
+8K3F
+8J7J
+8CQM
+8CLZ
+8HI7
+7YH3
+8OGG
+8CR3
+8GBH
+8BCK
+7YPR
+8ILL
+8OK7
+8FR5
+8BU0
+8OYY
+7XVI
+8JWS
+8R8A
+8Q1K
+8IN6
+8JI2
+8GD6
+8JDG
+8HHV
+8JT9
+8QFN
+8GKX
+8ACG
+8AVZ
+8BHU
+8GL4
+8CHX
+7Z3I
+8UPI
+8DQ2
+8HVC
+8OQH
+7Y9G
+7YH5
+8B2H
+8IQC
+7YN2
+8I3X
+7YR9
+8CH4
+8GOJ
+8FID
+7T4W
+8D40
+8HFP
+8K5K
+8GBI
+8OW8
+8Q70
+8SDY
+8JMR
+8IIB
+8SW0
+8DQ6
+8B6Q
+8H2N
+8SUT
+8J9Q
+8IW5
+8BL6
+8BBR
+8JWU
+8XBI
+8B2G
+8ASA
+8B2E
+8EY3
+8GJ9
+7ZAO
+8B2S
+8GD8
+8A9N
+8TN8
+8I34
+7Y1S
+8F9X
+8PXC
+8B5W
+8EC3
+8HFC
+8GYG
+8IR2
+8CJG
+8BUX
+7UWU
+8GYR
+8EHN
+7ZCB
+7YH1
+8IUB
+8BVK
+8OYE
+8GJ7
+8C4D
+8FTV
+8GIV
+8T8K
+8JZU
+8A26
+8HC0
+8EO2
+8H72
+8HMM
+8ORU
+8C05
+8G95
+8BSA
+7ZC9
+8IF7
+8QC6
+7QP5
+7YLS
+8TN3
+8ANI
+7ZET
+8AVH
+8D2Z
+8THR
+8K3B
+8J1C
+7Y4R
+8U2O
+8BTJ
+8HND
+8IN1
+8JX6
+7YLR
+8GK6
+8B7D
+8BGJ
+8P0S
+8B4L
+8JJ5
+8EVX
+8H24
+8OF7
+8IK2
+8P5N
+8GJA
+7YTL
+8H4P
+8KI5
+8T5J
+8GHH
+7ZN6
+8JNG
+8GR7
+8H7J
+8CI3
+8AH9
+8BS3
+8AGA
+8U01
+8HN2
+8H5A
+7YRO
+8AN5
+8HLG
+8OKS
+8K4X
+8D2J
+8PKD
+7VXT
+8H3Z
+8TJI
+8K76
+8CIL
+8HF2
+8P6K
+7YW0
+8C26
+8BV8
+8BS0
+8SME
+8I8A
+7YJI
+8EOZ
+8OML
+8BUY
+8OWI
+8DTQ
+8P26
+8EZT
+8FW7
+7Z8Z
+8DQA
+8K5L
+8PR6
+8K1F
+8AXR
+8JVM
+8ANT
+8ER5
+8OXU
+8H2R
+7YPE
+8WOJ
+8SS1
+8G5S
+8G2V
+8SAO
+8A50
+8BDK
+7YIL
+8ARD
+8KGZ
+8HO2
+8ECX
+8P32
+8AGY
+8SZB
+8ONB
+8DSP
+8BJW
+8U1J
+8FIH
+8KE8
+8SJJ
+8E2B
+8U1K
+8POC
+8HNO
+8K4W
+8UC6
+8AMR
+8ON4
+8PHB
+8IW0
+8OKW
+7Y16
+8PX4
+8BRF
+8BZ2
+8HGU
+8ERW
+8GAQ
+8HM5
+7YKH
+8SIU
+8G0K
+8ONQ
+7XPC
+8BIY
+8BNB
+8U12
+8FJE
+8A39
+8FIT
+8SRZ
+8HUC
+7W91
+7XRB
+8R79
+8AU0
+8OFJ
+8H1K
+8BCS
+8TTO
+8B3Z
+8EBB
+8QK8
+8OIJ
+8B45
+8EOX
+8EK4
+8E1E
+8ONF
+8HHJ
+8ENQ
+7YEQ
+8TNO
+8AN0
+8OK3
+8T4C
+8G8K
+8I3J
+8HDL
+8AXJ
+8GY4
+8H5S
+8CAR
+8FEH
+8UGC
+8GJY
+8OYX
+8DOT
+8D00
+8B6N
+8ANK
+8ARL
+8POF
+8HCY
+7Y3Z
+8KCA
+8CGM
+7QRL
+8GSY
+8T0J
+8HNA
+8OXL
+8P7A
+8EB9
+8JNM
+8I0A
+8F8N
+8P5P
+8OYS
+8JIY
+7XXE
+7YD4
+8FYG
+8HZZ
+8ES6
+7WU8
+8GJW
+8SOT
+8B4U
+8AY2
+8J1W
+8HM4
+8VEH
+8TMS
+8HX3
+8FJF
+8TVL
+8BKE
+8J69
+8EOV
+8OYW
+8BJV
+7TBO
+8OXH
+8B74
+8ANJ
+8HJJ
+8U0X
+8OYV
+7YJL
+8DHJ
+8EO7
+8JU8
+8JN0
+7VS2
+8F5D
+8ETQ
+8J0H
+8JNN
+8CQ3

model/boltzgen/task/analyze/analyze.py

@@ -0,0 +1,1483 @@
+from boltzgen.utils.quiet import quiet_startup
+
+
+quiet_startup()
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from concurrent.futures.process import BrokenProcessPool
+import copy
+import multiprocessing
+import numbers
+from pathlib import Path
+import traceback
+from typing import Optional, Dict, Any, List
+import subprocess
+import re
+import json
+
+from boltzgen.task.analyze.analyze_utils import (
+    TARGET_ID_RE,
+    calc_hydrophobicity,
+    compute_liability_metrics,
+    compute_novelty_foldseek,
+    compute_rmsd,
+    compute_ss_metrics,
+    get_best_folding_sample,
+    get_delta_sasa,
+    get_fold_metrics,
+    get_motif_set,
+    count_noncovalents,
+    largest_hydrophobic_patch_area,
+    make_histogram,
+    save_design_only_structure_to_cif,
+    save_design_only_structure_to_pdb,
+    vendi_scores,
+    vendi_sequences,
+)
+from matplotlib import pyplot as plt
+import matplotlib as mpl
+
+mpl.rcParams["figure.max_open_warning"] = 100
+import pydssp
+import rdkit
+from boltzgen.task.predict.data_from_generated import FromGeneratedDataModule
+
+import torch
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from collections import defaultdict
+
+from boltzgen.data import const
+from boltzgen.task.task import Task
+from boltzgen.data.data import Structure
+from boltzgen.data.write.mmcif import to_mmcif
+
+
+class Analyze(Task):
+    """
+    The Analyze step of the BoltzGen pipeline.
+    It computes various metrics on CPU.
+    It aggregates metrics from the Folding, Designfolding, and Affinity predictions step.
+    It writes metrics into a single csv file which is then used in the filter task.
+    It also creates a pickle file of sequences and structures for similarity assessments during filtering.
+    (needed for quality-diversity optimization algorithm)
+    """
+
+    def __init__(
+        self,
+        name: str,
+        data: FromGeneratedDataModule,
+        design_dir: str = None,
+        backbone_fold_metrics: bool = False,
+        allatom_fold_metrics: bool = True,
+        affinity_metrics: bool = False,
+        noncovalents_original: bool = False,
+        noncovalents_refolded: bool = False,
+        diversity_original: bool = False,
+        diversity_refolded: bool = False,
+        diversity_per_target_original: bool = False,
+        diversity_per_target_refolded: bool = False,
+        novelty_original: bool = False,
+        novelty_refolded: bool = False,
+        novelty_per_target_original: bool = False,
+        novelty_per_target_refolded: bool = False,
+        delta_sasa_original: bool = False,
+        delta_sasa_refolded: bool = False,
+        largest_hydrophobic: bool = False,
+        largest_hydrophobic_refolded: bool = False,
+        disulfide_quality: bool = False,
+        free_cys: bool = False,
+        compute_lddts: bool = True,  # computing LDDTs takes ~5-15 sec so it is optional
+        run_clustering: bool = False,
+        native: bool = False,
+        sequence_recovery: bool = False,
+        ss_conditioning_metrics: bool = False,
+        liability_analysis: bool = False,
+        liability_modality: str = "antibody",
+        liability_peptide_type: str = "linear",
+        debug: bool = False,
+        wandb: Optional[Dict[str, Any]] = None,
+        slurm: bool = False,
+        diversity_subset: int = None,
+        num_processes: int = 1,
+        foldseek_db: str = "/data/rbg/users/hstark/proteinblobs/data/foldseek_pdb/pdb",
+        foldseek_binary: str = "/data/rbg/users/hstark/foldseek/bin/foldseek",
+        skip_specific_ids: List[str] = None,
+        designfolding_metrics: bool = False,
+        use_design_mask_for_target: bool = False,
+    ) -> None:
+        """Initialize the task.
+
+        Parameters
+        ----------
+        fold_metrics : bool,
+            Compute folding metrics and assume that the folding directory exists.
+        """
+        super().__init__()
+        self.name = name
+        self.num_processes = num_processes
+        self.foldseek_db = Path(foldseek_db)
+        self.foldseek_binary = foldseek_binary
+        self.skip_specific_ids = set(skip_specific_ids or [])
+        self.data = data
+        self.noncovalents_original = noncovalents_original
+        self.noncovalents_refolded = noncovalents_refolded
+        self.diversity_original = diversity_original
+        self.diversity_refolded = diversity_refolded
+        self.diversity_per_target_original = diversity_per_target_original
+        self.diversity_per_target_refolded = diversity_per_target_refolded
+        self.novelty_original = novelty_original
+        self.novelty_refolded = novelty_refolded
+        self.novelty_per_target_original = novelty_per_target_original
+        self.novelty_per_target_refolded = novelty_per_target_refolded
+        self.delta_sasa_original = delta_sasa_original
+        self.delta_sasa_refolded = delta_sasa_refolded
+        self.largest_hydrophobic = largest_hydrophobic
+        self.largest_hydrophobic_refolded = largest_hydrophobic_refolded
+        self.compute_lddts = compute_lddts
+        self.run_clustering = run_clustering
+        self.affinity_metrics = affinity_metrics
+        self.fold_metrics = backbone_fold_metrics or allatom_fold_metrics
+        self.backbone_fold_metrics = backbone_fold_metrics
+        self.allatom_fold_metrics = allatom_fold_metrics
+        self.designfolding_metrics = designfolding_metrics
+        self.disulfide_quality = disulfide_quality
+        self.free_cys = free_cys
+        self.native = native
+        self.sequence_recovery = sequence_recovery
+        self.ss_conditioning_metrics = ss_conditioning_metrics
+        self.liability_analysis = liability_analysis
+        self.liability_modality = liability_modality
+        self.liability_peptide_type = liability_peptide_type
+        self.debug = debug
+        self.wandb = wandb
+        self.slurm = slurm
+        self.diversity_subset = diversity_subset
+        self.use_design_mask_for_target = use_design_mask_for_target
+
+        # Prevent each worker process from spawning its own multithreaded pools
+        torch.set_num_threads(1)
+        torch.set_num_interop_threads(1)
+
+        if design_dir is not None:
+            self.init_datasets(design_dir, load_dataset=False)
+
+        # Check that native structure is available if native metrics are desired
+        if self.native and not self.data.return_native:
+            msg = "native=True requires return_native=True in data config."
+            raise ValueError(msg)
+        if self.sequence_recovery and not self.native:
+            msg = "sequence_recovery=True requires native structure (native=True)."
+            raise ValueError(msg)
+
+        self.bindsite_adherence_thresholds = [3, 4, 5, 6, 7, 8, 9]
+
+    def init_datasets(self, design_dir: str, load_dataset: bool = False):
+        self.design_dir = Path(design_dir)
+
+        if load_dataset:
+            self.data.init_dataset(design_dir, skip_specific_ids=self.skip_specific_ids)
+
+        self.des_pdb_dir = self.design_dir / "des_pdbs"
+        self.des_pdb_dir.mkdir(parents=True, exist_ok=True)
+        self.des_refold_pdb_dir = self.design_dir / "des_refold_pdbs"
+        self.des_refold_pdb_dir.mkdir(parents=True, exist_ok=True)
+        self.refold_cif_dir = self.design_dir / const.refold_cif_dirname
+        self.refold_cif_dir.mkdir(parents=True, exist_ok=True)
+        self.metrics_dir = self.design_dir / const.metrics_dirname
+        self.metrics_dir.mkdir(exist_ok=True, parents=True)
+
+    def run_parallel(self, num, num_processes):
+        """
+        Run tasks  in parallel. If a worker crashes and the pool breaks,
+        restart a fresh pool and only rerun tasks that truly didn't finish.
+        """
+        ctx = multiprocessing.get_context("spawn")
+
+        all_task_ids = set(range(num))
+        completed_task_ids = set()
+        sample_ids = []
+
+        pbar = tqdm(total=num, desc="Processing samples")
+
+        while completed_task_ids != all_task_ids:
+            remaining = sorted(all_task_ids - completed_task_ids)
+
+            try:
+                with ProcessPoolExecutor(
+                    max_workers=num_processes, mp_context=ctx
+                ) as ex:
+                    fut2idx = {ex.submit(self.compute_metrics, i): i for i in remaining}
+
+                    # Iterate over futures that actually *completed* (finished or raised)
+                    for f in as_completed(fut2idx):
+                        idx = fut2idx[f]
+                        try:
+                            sid = f.result()
+                            if sid is not None:
+                                sample_ids.append(sid)
+                            # Count successful completion
+                            completed_task_ids.add(idx)
+                            pbar.update(1)
+
+                        except BrokenProcessPool:
+                            # Pool is dead, mark this idx completed.
+                            # Let the outer except restart a fresh pool for all unfinished.
+                            completed_task_ids.add(idx)
+                            pbar.update(1)
+                            raise
+
+            except BrokenProcessPool:
+                print("\nPOOL BROKEN: A worker died. Restarting with remaining tasks…")
+                # Nothing else to do: the while-loop will retry only the unfinished tasks.
+                continue
+
+        pbar.close()
+        return sample_ids
+
+    def run(self, config=None, run_prediction=False):
+        self.distribute_tasks()
+        self.aggregate_metrics()
+
+    def distribute_tasks(self):
+        # The rdkit thing is necessary to make multiprocessing with the rdkit molecules work.
+        rdkit.Chem.SetDefaultPickleProperties(rdkit.Chem.PropertyPickleOptions.AllProps)
+
+        # Compute metrics and write them to disk
+        sample_ids = []
+        num = len(self.data.predict_set)
+        if num == 0:
+            msg = "There were 0 samples to compute metrics for. Skipping the distribute_tasks step that calls compute_metrics"
+            print(msg)
+            return
+        num_processes = min(self.num_processes, multiprocessing.cpu_count())
+        if num_processes == 1:
+            for idx in tqdm(range(num)):
+                sample_id = self.compute_metrics(idx)
+                if sample_id is not None:
+                    sample_ids.append(sample_id)
+        else:
+            sample_ids = self.run_parallel(num, num_processes)
+        print(f"Computed metrics successfully for {len(sample_ids)} out of {num}.")
+
+    def aggregate_metrics(self):
+        # Load and aggregate saved metrics from disk
+
+        # Collect sample IDs for data_*.npz and metrics_*.npz
+        data_ids = {
+            f.stem.replace("data_", "") for f in self.metrics_dir.glob("data_*.npz")
+        }
+        metrics_ids = {
+            f.stem.replace("metrics_", "")
+            for f in self.metrics_dir.glob("metrics_*.npz")
+        }
+        sample_ids = sorted(data_ids & metrics_ids)
+        assert len(sample_ids) > 0
+
+        all_metrics, all_data = [], []
+        for sample_id in tqdm(
+            sample_ids, desc=f"Loading saved metrics from disk. 1% of total"
+        ):
+            data = np.load(
+                self.metrics_dir / f"data_{sample_id}.npz", allow_pickle=True
+            )
+            metrics = np.load(
+                self.metrics_dir / f"metrics_{sample_id}.npz", allow_pickle=True
+            )
+            data = {
+                k: v.item() if v.shape == () else torch.tensor(v)
+                for k, v in data.items()
+            }
+            metrics = {
+                k: v.item() if v.shape == () else torch.tensor(v)
+                for k, v in metrics.items()
+            }
+            all_metrics.append(metrics)
+            all_data.append(data)
+        df = pd.DataFrame(all_metrics)
+
+        # Cast per-motif integer fields and reconstruct a consolidated details column
+        try:
+            motif_keys = list(
+                get_motif_set(
+                    modality=self.liability_modality,
+                    peptide_type=self.liability_peptide_type,
+                ).keys()
+            )
+        except Exception:
+            motif_keys = []
+
+        # Build a consolidated details string if components exist
+        details_cols = []
+        for motif in motif_keys:
+            pos_col = f"liability_{motif}_position"
+            len_col = f"liability_{motif}_length"
+            sev_col = f"liability_{motif}_severity"
+            det_col = f"liability_{motif}_details"
+            cnt_col = f"liability_{motif}_count"
+            numpos_col = f"liability_{motif}_num_positions"
+            # Standardize dtypes: fill NA for ints, then cast
+            for col in [pos_col, len_col, sev_col, cnt_col, numpos_col]:
+                if col in df.columns:
+                    if df[col].dtype.kind in ("f", "O"):
+                        df[col] = df[col].fillna(-1).astype(int)
+            # Keep details as string
+            if det_col in df.columns:
+                df[det_col] = df[det_col].fillna("").astype(str)
+            # Track columns for a consolidated details view
+            if all(
+                c in df.columns for c in [pos_col, len_col, sev_col, det_col, cnt_col]
+            ):
+                details_cols.append(
+                    (motif, pos_col, len_col, sev_col, det_col, cnt_col)
+                )
+
+        # Optional: single consolidated details column combining motifs
+        if details_cols:
+
+            def _compose_details(row):
+                items = []
+                for motif, pos_col, len_col, sev_col, det_col, cnt_col in details_cols:
+                    cnt = row[cnt_col]
+                    if cnt and cnt > 0:
+                        pos = row[pos_col]
+                        length = row[len_col]
+                        sev = row[sev_col]
+                        det = row[det_col]
+                        base = f"{motif}x{cnt}"
+                        if pos >= 0:
+                            base += f"(pos{pos},len{length},sev{sev})"
+                        else:
+                            base += f"(sev{sev})"
+                        if det:
+                            base += f"[{det}]"
+                        items.append(base)
+                return "; ".join(items)
+
+            df["liability_details"] = df.apply(_compose_details, axis=1)
+
+        # Run clustering
+        if self.run_clustering:
+            df = self.run_foldseek_clustering(df)
+        # Write individual metrics to disk
+        csv_path = Path(self.design_dir) / f"aggregate_metrics_{self.name}.csv"
+        df.to_csv(csv_path, float_format="%.5f", index=False)
+
+        # Store ca coords and seq in a pickle file for later usage in e.g. diversity aware filtering
+        data_rows = []
+        for data in all_data:
+            data_rows.append(
+                {
+                    "id": data["sample_id"],
+                    "target_id": data["target_id"],
+                    "sequence": "".join(
+                        [
+                            const.prot_token_to_letter[const.tokens[t]]
+                            for t in data["design_seq"]
+                        ]
+                    ),
+                    "ca_coords": json.dumps(data["ca_coords"].numpy().tolist()),
+                }
+            )
+        ca_seq_df = pd.DataFrame(data_rows)
+        ca_seq_df.to_pickle(
+            Path(self.design_dir) / "ca_coords_sequences.pkl.gz", compression="gzip"
+        )
+
+        # Compute per target metrics
+        df["target_id"] = df["id"].apply(lambda s: TARGET_ID_RE.match(s).group(1))
+        per_target_df = df.groupby("target_id").mean(numeric_only=True).reset_index()
+        csv_path = Path(self.design_dir) / f"per_target_metrics_{self.name}.csv"
+        per_target_df.to_csv(csv_path, float_format="%.5f", index=False)
+
+        avg_metrics = df.mean(numeric_only=True).round(5).to_dict()
+        avg_metrics["num_targets"] = len(all_metrics)
+        if self.run_clustering:
+            if "cluster_07_seqidentity" in df.columns:
+                avg_metrics["num_cluster_07_seqidentity"] = len(
+                    np.unique(df["cluster_07_seqidentity"].to_numpy())
+                )
+            if "num_clusters_05_tmscore" in df.columns:
+                avg_metrics["num_clusters_05_tmscore"] = len(
+                    np.unique(df["clusters_05_tmscore"].to_numpy())
+                )
+
+        diversity_metrics, diversity_data = self.compute_diversity(
+            all_data, all_metrics
+        )
+        for k in diversity_metrics:
+            avg_metrics[k] = diversity_metrics[k]
+
+        novelty_metrics, novelty_data = self.compute_novelty()
+        for k in novelty_metrics:
+            avg_metrics[k] = novelty_metrics[k]
+
+        _, histograms = self.make_histograms(all_metrics)
+
+        # Log to Wandb
+        if self.wandb is not None and not self.debug:
+            import wandb
+
+            print("\nOverall average metrics:", avg_metrics)
+
+            # Make residue distribution plot
+            native_stats = np.load("data/native_statistics.npz")
+            design_freqs = np.array(
+                [
+                    avg_metrics[f"{k}_fraction"]
+                    for k in const.fake_atom_placements.keys()
+                ]
+            )
+            x = np.arange(len(const.fake_atom_placements.keys()))
+            width = 0.15
+            fig_res, ax = plt.subplots(figsize=(12, 6))
+            ax.bar(x - width / 2, design_freqs, width, label="Design frequency")
+            ax.bar(
+                x + width / 2, native_stats["res_dist"], width, label="Data frequency"
+            )
+            ax.set_xlabel("Res Type")
+            ax.set_ylabel("Probability")
+            ax.set_title("Res Type distributions")
+            ax.set_xticks(x)
+            ax.set_xticklabels(const.fake_atom_placements.keys())
+            ax.legend()
+            ax.grid(True, which="both", axis="y", linestyle="--", linewidth=0.5)
+            plt.tight_layout()
+            fig_res.savefig(Path(self.design_dir) / "res_type_distribution.png")
+
+            # Make secondary structure distribution plot
+            ss_dist = np.array(
+                [avg_metrics["loop"], avg_metrics["helix"], avg_metrics["sheet"]]
+            )
+            x = np.arange(3)
+            width = 0.15
+            fig_ss, ax = plt.subplots(figsize=(12, 6))
+            ax.bar(x - width / 2, ss_dist, width, label="Designed")
+            ax.bar(x + width / 2, native_stats["ss_dist"], width, label="Native data")
+            ax.set_xlabel("Secondary Structure type")
+            ax.set_ylabel("Frequency")
+            ax.set_title("Secondary Structure distributions")
+            ax.set_xticks(x)
+            ax.set_xticklabels(["loop", "helix", "sheet"])
+            ax.legend()
+            ax.grid(True, which="both", axis="y", linestyle="--", linewidth=0.5)
+            plt.tight_layout()
+            fig_ss.savefig(
+                Path(self.design_dir) / "secondary_structure_distribution.png"
+            )
+
+            wandb.init(name=self.name, **self.wandb)
+            wandb.log(avg_metrics)
+            wandb.log({"res_dist": wandb.Image(fig_res)})
+            wandb.log({"ss_dist": wandb.Image(fig_ss)})
+
+            # Log histograms
+            for name, fig in histograms.items():
+                wandb.log({f"{name}_hist": wandb.Image(fig)})
+
+            # plot per target vendiscore histogram
+            if self.diversity_per_target_original:
+                fig, ax = plt.subplots(figsize=(6, 4))
+                ax.hist(
+                    diversity_data["vendi_tm_fixed"], bins=50, color="blue", alpha=0.7
+                )
+                ax.set_title("Vendi Score Per Target Distribution")
+                ax.set_xlabel("Vendi Score")
+                ax.set_ylabel("Count")
+                plt.tight_layout()
+                wandb.log({"vendi_per_target": wandb.Image(fig)})
+                plt.close(fig)
+
+            if self.novelty_per_target_original:
+                fig, ax = plt.subplots(figsize=(6, 4))
+                ax.hist(
+                    novelty_data["nov_df"]["novelty"],
+                    bins=50,
+                    color="blue",
+                    alpha=0.7,
+                )
+                ax.set_title("Novelty Per Target Original Distribution")
+                ax.set_xlabel("Novelty")
+                ax.set_ylabel("Count")
+                plt.tight_layout()
+                wandb.log({"novelty_per_target_original_hist": wandb.Image(fig)})
+                plt.close(fig)
+
+            if self.novelty_per_target_refolded and (self.fold_metrics):
+                fig, ax = plt.subplots(figsize=(6, 4))
+                ax.hist(
+                    novelty_data["nov_df_refold"]["novelty"],
+                    bins=50,
+                    color="blue",
+                    alpha=0.7,
+                )
+                ax.set_title("Novelty Per Target Refolded Distribution")
+                ax.set_xlabel("Novelty")
+                ax.set_ylabel("Count")
+                plt.tight_layout()
+                wandb.log({"novelty_per_target_refolded_hist": wandb.Image(fig)})
+                plt.close(fig)
+
+    def compute_metrics(self, idx=None, sample_id=None, suffix=None, design_dir=None):
+        if sample_id is None:
+            sample_id = self.data.predict_set.generated_paths[idx].stem
+            feat = self.data.predict_set[idx]
+        else:
+            if design_dir is None:
+                design_dir = self.design_dir
+
+            feat = self.data.predict_set.get_sample(
+                design_dir=design_dir, sample_id=sample_id
+            )
+        path = feat["path"]
+        if feat["exception"]:
+            msg = f"Failed obtaining valid features for {path} due to feat['exception']. Skipping."
+            print(msg)
+            return None
+
+        # Get designed sequence
+        res_type_argmax = torch.argmax(feat["res_type"], dim=-1)
+        design_seq_tensor = res_type_argmax[
+            feat["design_mask"].bool() & feat["token_pad_mask"].bool()
+        ]
+        design_chain_id = feat["asym_id"][
+            torch.where(feat["design_mask"].bool() & feat["token_pad_mask"].bool())[0][
+                0
+            ]
+        ].item()
+        design_chain_seq = res_type_argmax[design_chain_id == feat["asym_id"]]
+        design_seq = "".join(
+            [
+                const.prot_token_to_letter.get(const.tokens[t], "X")
+                for t in design_seq_tensor
+            ]
+        )
+        design_chain_seq = "".join(
+            [
+                const.prot_token_to_letter.get(const.tokens[t], "X")
+                for t in design_chain_seq
+            ]
+        )
+
+        # initialize metrics
+        metrics = {
+            "id": sample_id,
+            "file_name": path.name,
+            "designed_sequence": design_seq,
+            "designed_chain_sequence": design_chain_seq,
+        }
+
+        # Add per-chain sequences to csv when designing multiple chains
+        design_token_indices = torch.where(feat["design_mask"].bool() & feat["token_pad_mask"].bool())[0]
+        designed_chain_ids = feat["asym_id"][design_token_indices].unique().tolist()
+        if len(designed_chain_ids) > 1:
+            for chain_id in designed_chain_ids:
+                chain_mask = feat["asym_id"] == chain_id
+
+                # Full chain sequence
+                chain_res_types = res_type_argmax[chain_mask]
+                full_chain_seq = "".join(
+                    [
+                        const.prot_token_to_letter.get(const.tokens[t], "X")
+                        for t in chain_res_types
+                    ]
+                )
+
+                # Designed residues only from this chain
+                design_chain_mask = feat["design_mask"].bool() & feat["token_pad_mask"].bool() & chain_mask
+                design_res_types = res_type_argmax[design_chain_mask]
+                design_seq = "".join(
+                    [
+                        const.prot_token_to_letter.get(const.tokens[t], "X")
+                        for t in design_res_types
+                    ]
+                )
+
+                metrics[f"designed_sequence_{chain_id}"] = design_seq
+                metrics[f"full_sequence_{chain_id}"] = full_chain_seq
+
+
+        target_id = re.search(rf"{self.data.cfg.target_id_regex}", sample_id).group(1)
+
+        # Get masks
+        design_mask = feat["design_mask"].bool()
+        chain_design_mask   = feat["chain_design_mask"].bool() 
+
+        design_resolved_mask = design_mask & feat["token_resolved_mask"].bool()
+
+        # For symmetric designs where all chains have designed residues, use design_mask
+        # instead of chain_design_mask so "target" = non-designed residues (not empty)
+        if self.use_design_mask_for_target:
+            target_resolved_mask = (~design_mask) & feat["token_resolved_mask"].bool()
+        else:
+            target_resolved_mask = (~chain_design_mask) & feat["token_resolved_mask"].bool()
+        atom_design_resolved_mask = (
+            (feat["atom_to_token"].float() @ design_resolved_mask.unsqueeze(-1).float())
+            .bool()
+            .squeeze()
+        )
+        atom_target_resolved_mask = (
+            (feat["atom_to_token"].float() @ target_resolved_mask.unsqueeze(-1).float())
+            .bool()
+            .squeeze()
+        )
+        atom_resolved_mask = feat["atom_resolved_mask"]
+        resolved_atoms_design_mask = atom_design_resolved_mask[atom_resolved_mask]
+        resolved_atoms_target_mask = atom_target_resolved_mask[atom_resolved_mask]
+        atom_chain_mask = (
+            (
+                feat["atom_to_token"].float()
+                @ chain_design_mask.unsqueeze(-1).float()
+            )
+            .bool()
+            .squeeze()
+        )
+
+        # Get masks for native structure
+        if self.native:
+            native_design_mask = feat["native_design_mask"].bool()
+            native_target_resolved_mask = (
+                ~native_design_mask & feat["native_token_resolved_mask"].bool()
+            )
+            native_atom_target_resolved_mask = (
+                (
+                    feat["native_atom_to_token"].float()
+                    @ native_target_resolved_mask.unsqueeze(-1).float()
+                )
+                .bool()
+                .squeeze()
+            )
+
+        # add to design_only directory for novelty computation
+        des_cif_path = None
+        if not suffix is None:
+            des_pdb_dir = self.des_pdb_dir / suffix
+            des_pdb_dir.mkdir(exist_ok=True, parents=True)
+        else:
+            des_pdb_dir = self.des_pdb_dir
+        des_pdb_path = des_pdb_dir / f"{feat['id']}_des.pdb"
+        des_cif_path = des_pdb_path.with_suffix(".cif")
+        if (
+            self.novelty_original
+            or self.novelty_refolded
+            or self.novelty_per_target_original
+            or self.novelty_per_target_refolded
+            or self.run_clustering
+        ):
+            try:
+                save_design_only_structure_to_pdb(
+                    atom_design_mask=atom_chain_mask,
+                    structure=feat["str_gen"],
+                    output_path=des_pdb_path,
+                )
+                des_cif_path = des_pdb_path.with_suffix(".cif")
+            except Exception as e:
+                print(
+                    f"[Warning] Could not save design-only structure for {feat['id']}: {e}. Skipping this file."
+                )
+                traceback.print_exc()
+                return None
+
+        # largest hydrophobic patch area original
+        if self.largest_hydrophobic:
+            if not des_cif_path.exists():
+                save_design_only_structure_to_cif(
+                    atom_design_mask=atom_chain_mask,
+                    structure=feat["str_gen"],
+                    output_path=des_cif_path,
+                )
+            area = largest_hydrophobic_patch_area(des_cif_path)
+            metrics["design_largest_hydrophobic_patch"] = area
+        if des_cif_path is not None:
+            des_cif_path.unlink(missing_ok=True)
+
+        # Count logging
+        metrics["num_prot_tokens"] = (
+            (feat["mol_type"] == const.chain_type_ids["PROTEIN"]).sum().item()
+        )
+        metrics["num_lig_atoms"] = (
+            (feat["mol_type"] == const.chain_type_ids["NONPOLYMER"]).sum().item()
+        )
+        metrics["num_resolved_tokens"] = feat["token_resolved_mask"].sum().item()
+        metrics["num_tokens"] = feat["token_pad_mask"].sum().item()
+        metrics["num_design"] = feat["design_mask"].sum().item()
+
+        # delta sasa for original
+        if self.delta_sasa_original:
+            (
+                delta_sasa_orig,
+                design_sasa_unbound,
+                design_sasa_bound,
+            ) = get_delta_sasa(
+                path,
+                atom_target_mask=resolved_atoms_target_mask,
+                atom_design_mask=resolved_atoms_design_mask,
+            )
+            metrics["delta_sasa_original"] = delta_sasa_orig
+            metrics["design_sasa_unbound_original"] = design_sasa_unbound
+            metrics["design_sasa_bound_original"] = design_sasa_bound
+
+        # Noncovalents metrics for original
+        try:
+            if self.noncovalents_original:
+                metrics.update(count_noncovalents(feat))
+        except Exception as e:
+            print(
+                f"[Error] computing noncovalents for {path}: {e}. Skipping this file."
+            )
+            traceback.print_exc()
+            return None
+
+        # Sequence metrics
+        if self.sequence_recovery:
+            native_seq = torch.argmax(feat["native_res_type"], dim=-1)[
+                native_design_mask
+            ]
+            metrics["seq_recovery"] = (
+                (design_seq_tensor == native_seq).float().mean().item()
+            )
+        for t in const.fake_atom_placements.keys():
+            metrics[f"{t}_fraction"] = (
+                (design_seq_tensor == const.token_ids[t]).float().mean().item()
+            )
+
+        # Secondary structure metrics
+        # Compute secondary structure distribution. First get backbone then use pydssp to compute.
+        bb_design_mask = (
+            feat["atom_pad_mask"].bool()
+            & atom_design_resolved_mask
+            & feat["backbone_mask"].bool()
+        )
+        bb_coords = feat["coords"][0][bb_design_mask]
+        num_atoms = bb_coords.shape[0]
+        if num_atoms % 4 != 0:
+            msg = f"BB atoms {num_atoms} is not divisible by 4 for {path}"
+            print(msg)
+            traceback.print_exc()
+            return None
+        bb = bb_coords.reshape(-1, 4, 3)
+        ca_coords = bb[:, 1, :]
+        if len(bb) > 5:
+            try:
+                dssp = (
+                    torch.zeros(bb.shape[0], dtype=torch.long)
+                    if torch.sum(bb_design_mask).item() == 0
+                    else pydssp.assign(bb, out_type="index")
+                )
+                # Secondary structure conditioning metric
+                if self.ss_conditioning_metrics:
+                    ss_conditioning_metricsed = feat["ss_type"][design_mask]
+                    dssp_adjusted = dssp + 1
+                    ss_metrics = compute_ss_metrics(
+                        dssp_adjusted, ss_conditioning_metricsed
+                    )
+                    metrics.update(ss_metrics)
+                metrics["loop"] = (dssp == 0).float().mean().item()
+                metrics["helix"] = (dssp == 1).float().mean().item()
+                metrics["sheet"] = (dssp == 2).float().mean().item()
+            except:
+                traceback.print_exc()
+                print(f"DSSP failed for {path}.")
+                return None
+        else:
+            metrics["loop"] = float("nan")
+            metrics["helix"] = float("nan")
+            metrics["sheet"] = float("nan")
+
+        # Liability analysis
+        if self.liability_analysis:
+            try:
+                liability_metrics = compute_liability_metrics(
+                    design_chain_seq,
+                    self.liability_modality,
+                    self.liability_peptide_type,
+                )
+                metrics.update(liability_metrics)
+            except Exception as e:
+                traceback.print_exc()
+                print(f"Liability analysis failed for {sample_id}: {e}")
+                return None
+
+        # Compute RMSD between native (input) and generated conditioning structures.
+        # conditioning structure does not have fake atoms, just parse the coordinates and compute rmsd.
+        metrics["native_rmsd"] = 0.0
+        metrics["native_rmsd_bb"] = 0.0
+        if self.native:
+            target_coords = feat["coords"][:, atom_target_resolved_mask]
+            native_target_coords = feat["native_coords"][
+                :, native_atom_target_resolved_mask
+            ]
+            target_rmsd = compute_rmsd(native_target_coords, target_coords)
+
+            bb_target_coords = feat["coords"][
+                :, atom_target_resolved_mask & feat["backbone_mask"].bool()
+            ]
+            bb_native_target_coords = feat["native_coords"][
+                :,
+                native_atom_target_resolved_mask & feat["native_backbone_mask"].bool(),
+            ]
+            bb_target_rmsd = compute_rmsd(bb_native_target_coords, bb_target_coords)
+
+            bb_coords = feat["coords"][:, feat["backbone_mask"].bool()]
+            bb_native_coords = feat["native_coords"][
+                :, feat["native_backbone_mask"].bool()
+            ]
+            bb_rmsd = compute_rmsd(bb_native_coords, bb_coords)
+            metrics["native_rmsd_all_bb"] = bb_rmsd.item()
+
+            metrics["native_rmsd"] = target_rmsd.item()
+            metrics["native_rmsd_bb"] = bb_target_rmsd.item()
+
+        # Check binding site adherence. For each binding site token, find closest design token
+        binding_site_mask = feat["binding_type"] == 1
+        if binding_site_mask.sum() > 1:
+            token_distances = torch.cdist(feat["center_coords"], feat["center_coords"])
+            bindsite_design_distances = token_distances[binding_site_mask][
+                :, feat["design_mask"]
+            ]
+            min_bindsite_design_distances = bindsite_design_distances.min(axis=1).values
+            for threshold in self.bindsite_adherence_thresholds:
+                metrics[f"bindsite_under_{threshold}rmsd"] = (
+                    (min_bindsite_design_distances < threshold).float().mean().item()
+                )
+
+        # Count free Cysteines
+        if self.free_cys:
+            cysteine_mask = (
+                torch.argmax(feat["res_type"], dim=-1)
+                == const.token_ids["CYS"] & feat["design_mask"]
+            )
+            cysteine_sulfur_indices = (
+                torch.argmax(feat["token_to_rep_atom"].int(), dim=-1)[cysteine_mask]
+                + const.ref_atoms["CYS"].index("SG")
+                - const.ref_atoms["CYS"].index("CA")
+            )
+            cysteine_coords = feat["coords"][0][cysteine_sulfur_indices]
+            free_cysteines = 0
+            dist = torch.cdist(cysteine_coords, cysteine_coords)
+            for i in range(dist.shape[0]):
+                if dist[i, torch.argsort(dist[i])[1]] > 4.0:
+                    free_cysteines += 1
+            metrics["free_cysteines"] = free_cysteines
+
+        # Quality of Cysteine-Cysteine bonds
+        if self.disulfide_quality:
+            bonds = feat["structure_bonds"]
+            sulfur_mask = (
+                torch.argmax(feat["ref_element"], dim=-1)
+                == const.element_to_atomic_num["S"]
+            )
+            dist = []
+            disulfide_bonds = []
+            for bond in bonds:
+                if (
+                    bond[6] == const.bond_type_ids["COVALENT"]
+                    and sulfur_mask[bond[4]]
+                    and sulfur_mask[bond[5]]
+                ):
+                    disulfide_bonds.append(bond)
+            for ds_bond in disulfide_bonds:
+                dist.append(
+                    torch.cdist(
+                        feat["coords"][0][ds_bond[4]].unsqueeze(0),
+                        feat["coords"][0][ds_bond[5]].unsqueeze(0),
+                    )
+                )
+            if len(dist) > 0:
+                min_dist = torch.cat(dist).min()
+                max_dist = torch.cat(dist).max()
+                metrics["disulfide_bond_len_qual"] = [max_dist.item(), min_dist.item()]
+
+        # Folding metrics
+        ca_coords_refolded = None
+        metrics["native_rmsd_refolded"] = 0.0
+        metrics["native_rmsd_bb_refolded"] = 0.0
+        if self.fold_metrics:
+            # Compute refolding metrics when refolding the design only in absence of the target (the whole design chain and anything covalently attached is refolded).
+            if self.designfolding_metrics:
+                folded_path = (
+                    self.design_dir / const.folding_design_dirname / f"{feat['id']}.npz"
+                )
+                if not folded_path.exists():
+                    print(f"Folded path does not exist. Skipping: {folded_path}")
+                    return None
+
+                folded = np.load(
+                    self.design_dir / const.folding_design_dirname / f"{feat['id']}.npz"
+                )
+                feat_design = {
+                    k: torch.from_numpy(folded[k]).squeeze(0)
+                    for k in [
+                        "input_coords",
+                        "res_type",
+                        "token_index",
+                        "atom_resolved_mask",
+                        "atom_to_token",
+                        "mol_type",
+                        "backbone_mask",
+                    ]
+                }
+                feat_design["design_mask"] = torch.ones_like(
+                    feat_design["token_index"]
+                ).bool()
+                feat_design["chain_design_mask"] = torch.ones_like(
+                    feat_design["token_index"]
+                ).bool()
+                # Use the same features as refolded, just need to change the coordinates back to the original designed coordinates
+                feat_design["coords"] = feat_design["input_coords"]
+                if (
+                    not len(folded["res_type"].squeeze())
+                    == len(feat_design["res_type"])
+                    or not (
+                        folded["res_type"].squeeze() == feat_design["res_type"]
+                    ).all()
+                ):
+                    msg = f"Skipping {path}. The sequences for which the refolding was run are not the same as the sequences in the design_dir. Maybe the designs in the design_dir were overwritten. Or maybe two processes are operating on the same design_dir."
+                    print(msg)
+                    return None
+
+                # Compute all-atom folding metrics when refolding only the designed part (this does not make sense when using inverse folded structures)
+                if self.allatom_fold_metrics:
+                    fold_metrics = get_fold_metrics(
+                        feat_design,
+                        folded,
+                        compute_lddts=self.compute_lddts,
+                    )
+                    fold_metrics = {
+                        f"designfolding-{k}": v for k, v in fold_metrics.items()
+                    }
+                    metrics.update(fold_metrics)
+
+                # Compute backbone folding metrics when refolding only the designed part
+                if self.backbone_fold_metrics:
+                    feat_bb = copy.deepcopy(feat_design)
+                    feat_bb["atom_resolved_mask"] = feat_bb["atom_resolved_mask"].to(
+                        bool
+                    ) & feat_bb["backbone_mask"].to(bool)
+                    fold_metrics_bb = get_fold_metrics(
+                        feat_bb,
+                        folded,
+                        compute_lddts=self.compute_lddts,
+                        prefix="bb_",
+                    )
+                    fold_metrics_bb = {
+                        f"designfolding-{k}": v for k, v in fold_metrics_bb.items()
+                    }
+                    metrics.update(fold_metrics_bb)
+
+            folded_path = self.design_dir / const.folding_dirname / f"{feat['id']}.npz"
+            if not folded_path.exists():
+                print(f"Folded path does not exist. Skipping: {folded_path}")
+                return None
+            folded = np.load(
+                self.design_dir / const.folding_dirname / f"{feat['id']}.npz"
+            )
+            if (
+                not len(folded["res_type"].squeeze()) == len(feat["res_type"])
+                or not (folded["res_type"].squeeze() == feat["res_type"].numpy()).all()
+            ):
+                msg = f"Skipping {path}. The sequences for which the refolding was run are not the same as the sequences in the design_dir. Maybe the designs in the design_dir were overwritten. Or maybe two processes are operating on the same design_dir."
+                print(msg)
+                return None
+
+            # Compute allatom folding metrics when refolding the whole complex (this does not make sense when the design_dir contains inverse folded structures).
+            if self.allatom_fold_metrics:
+                fold_metrics = get_fold_metrics(
+                    feat,
+                    folded,
+                    compute_lddts=self.compute_lddts,
+                )
+                metrics.update(fold_metrics)
+
+            # Compute backbone folding metrics when refolding the whole complex
+            if self.backbone_fold_metrics:
+                feat_bb = copy.deepcopy(feat)
+                feat_bb["atom_resolved_mask"] = feat_bb["atom_resolved_mask"].to(
+                    bool
+                ) & feat_bb["backbone_mask"].to(bool)
+                fold_metrics_bb = get_fold_metrics(
+                    feat_bb, folded, compute_lddts=self.compute_lddts, prefix="bb_"
+                )
+                fold_metrics_bb = {f"{k}": v for k, v in fold_metrics_bb.items()}
+                metrics.update(fold_metrics_bb)
+
+            # Construct features for refolded complex
+            feat_out = {}
+            for k in feat.keys():
+                if k == "coords":
+                    best_sample = get_best_folding_sample(folded)
+                    feat_out[k] = torch.from_numpy(best_sample["coords"])
+                else:
+                    feat_out[k] = feat[k]
+            refold_atom_target_resolved_mask = (
+                (
+                    feat_out["atom_to_token"].float()
+                    @ target_resolved_mask.unsqueeze(-1).float()
+                )
+                .bool()
+                .squeeze()
+            )
+            refold_target_coords = feat_out["coords"][
+                refold_atom_target_resolved_mask, :
+            ][None, ...]
+
+            # Compute reconstruction RMDS compared to a native binder structure (if a native binder exists).
+            if self.native:
+                refold_target_rmsd = compute_rmsd(
+                    native_target_coords,
+                    refold_target_coords,
+                )
+                bb_refold_target_coords = feat_out["coords"][
+                    refold_atom_target_resolved_mask & feat_out["backbone_mask"].bool()
+                ][None, ...]
+                bb_refold_target_rmsd = compute_rmsd(
+                    bb_native_target_coords,
+                    bb_refold_target_coords,
+                )
+                metrics["native_rmsd_refolded"] = refold_target_rmsd.item()
+                metrics["native_rmsd_bb_refolded"] = bb_refold_target_rmsd.item()
+
+            # Save the refolded structure of the design to a pdb file if novelty computation needs to be run on it later.
+            des_refold_pdb_path = (
+                self.des_refold_pdb_dir / f"{feat['id']}_des_refold.pdb"
+            )
+            des_refold_cif_path = des_refold_pdb_path.with_suffix(".cif")
+            if self.novelty_refolded or self.novelty_per_target_refolded:
+                structure, _, _ = Structure.from_feat(feat_out)
+                try:
+                    save_design_only_structure_to_pdb(
+                        atom_design_mask=atom_chain_mask,
+                        structure=structure,
+                        output_path=des_refold_pdb_path,
+                    )
+                except Exception as e:
+                    print(
+                        f"[Warning] Could not save design-only structure for {feat['id']}: {e}. Skipping this file."
+                    )
+                    traceback.print_exc()
+                    return None
+
+            # largest hydrophobic patch area refolded
+            if self.largest_hydrophobic_refolded:
+                if not des_refold_cif_path.exists():
+                    structure, _, _ = Structure.from_feat(feat_out)
+
+                    save_design_only_structure_to_cif(
+                        atom_design_mask=atom_chain_mask,
+                        structure=structure,
+                        output_path=des_refold_cif_path,
+                    )
+                area_refold = largest_hydrophobic_patch_area(des_refold_cif_path)
+                metrics["design_largest_hydrophobic_patch_refolded"] = area_refold
+            if des_refold_cif_path is not None:
+                des_refold_cif_path.unlink(missing_ok=True)
+
+            # Compute sequence based hydrophobicity
+            metrics["design_chain_hydrophobicity"] = calc_hydrophobicity(
+                design_chain_seq
+            )
+            metrics["design_hydrophobicity"] = calc_hydrophobicity(design_seq)
+
+            # delta sasa for refolded
+            if self.delta_sasa_refolded:
+                cif_path_refolded = self.refold_cif_dir / f"{feat['id']}.cif"
+
+                if not cif_path_refolded.exists():
+                    msg = f"Refolded cif path does not exist. This can happen if a process was interrupted between writing the refold .npz file and the refold .cif file. Missing path: {cif_path_refolded}"
+                    print(msg)
+                    return None
+
+                # Compute delta sasa
+                (
+                    delta_sasa_refolded,
+                    design_sasa_unbound,
+                    design_sasa_bound,
+                ) = get_delta_sasa(
+                    cif_path_refolded,
+                    atom_target_mask=resolved_atoms_target_mask,
+                    atom_design_mask=resolved_atoms_design_mask,
+                )
+
+                metrics["delta_sasa_refolded"] = delta_sasa_refolded
+                metrics["design_sasa_unbound_refolded"] = design_sasa_unbound
+                metrics["design_sasa_bound_refolded"] = design_sasa_bound
+
+            # noncovalents metrics for refolded structure
+            if self.noncovalents_refolded:
+                try:
+                    _metrics = count_noncovalents(feat_out)
+                    _metrics = {f"{k}_refolded": v for k, v in _metrics.items()}
+                    metrics.update(_metrics)
+                except Exception as e:
+                    print(
+                        f"[Error] computing noncovalents refolded for {path}: {e}. Skipping this file."
+                    )
+                    traceback.print_exc()
+                    return None
+
+            bb_out = feat_out["coords"][bb_design_mask].reshape(-1, 4, 3)
+            ca_coords_refolded = bb_out[:, 1, :].cpu()
+
+        # Affinity metrics
+        if self.affinity_metrics:
+            affinity_path = (
+                self.design_dir / const.affinity_dirname / f"{feat['id']}.npz"
+            )
+            if not affinity_path.exists():
+                print(f"Affinity path does not exist. Skipping: {affinity_path}")
+                return None
+
+            affinity = np.load(
+                self.design_dir / const.affinity_dirname / f"{feat['id']}.npz"
+            )
+
+            for key in const.eval_keys_affinity:
+                if key in affinity:
+                    metrics[key] = affinity[key].item()
+
+            if "affinity_probability_binary1" in metrics:
+                metrics["affinity_probability_binary1>50"] = (
+                    metrics["affinity_probability_binary1"] > 0.5
+                )
+                metrics["affinity_probability_binary1>75"] = (
+                    metrics["affinity_probability_binary1"] > 0.75
+                )
+        
+        for key in const.eval_keys_confidence:
+            if key in feat:
+                if isinstance(feat[key], torch.Tensor) and feat[key].numel() == 1:
+                    metrics[key] = feat[key].item()
+                elif isinstance(feat[key], (float, int)):
+                    metrics[key] = feat[key]
+
+        # Write outputs to files and return sample_id for conformation of successful processing
+        data = {
+            "target_id": target_id,
+            "sample_id": sample_id,
+            "design_seq": design_seq_tensor.cpu(),
+            "ca_coords": ca_coords.cpu(),
+            "ca_coords_refolded": ca_coords_refolded,
+        }
+        data_path = self.metrics_dir / f"data_{sample_id}.npz"
+        metrics_path = self.metrics_dir / f"metrics_{sample_id}.npz"
+        np.savez_compressed(metrics_path, **metrics)
+        np.savez_compressed(data_path, **data)
+        return sample_id
+
+    def compute_diversity(self, all_data, all_metrics):
+        avg_metrics = {}
+        metrics_data = {}
+        fold_metrics = self.fold_metrics
+
+        # Aggregate alpha carbon positions for diversity eval
+        ca_gen = defaultdict(list)
+        input_metrics = defaultdict(list)
+        ca_refold = defaultdict(list)
+        sequences = defaultdict(list)
+        for i, data in enumerate(all_data):
+            ca_gen[data["target_id"]].append(data["ca_coords"])
+            ca_refold[data["target_id"]].append(data["ca_coords_refolded"])
+            input_metrics[data["target_id"]].append(all_metrics[i])
+
+            seq = data["design_seq"]
+            try:
+                seq = "".join(
+                    [const.prot_token_to_letter[const.tokens[t]] for t in seq]
+                )
+                sequences[data["target_id"]].append(seq)
+            except KeyError as e:
+                print(
+                    f"[Error] KeyError '{e.args[0]}' for target_id: {data['target_id']}, sample_id: {data['sample_id']}"
+                )
+        print(
+            f"Number of targets: {len(ca_gen)}. Number of designs: {len(all_metrics)}."
+        )
+
+        if self.diversity_original:
+            print("Computing diveristy original.")
+            ca_filtered = [ca[0] for ca in ca_gen.values() if len(ca[0]) >= 3]
+            metrics_filtered = [
+                m[0]
+                for ca, m in zip(ca_gen.values(), input_metrics.values())
+                if len(ca[0]) >= 3
+            ]
+            scores = vendi_scores(
+                all_ca_coords=ca_filtered,
+                all_metrics=metrics_filtered,
+                fold_metrics=fold_metrics,
+                diversity_subset=self.diversity_subset,
+                compute_lddts=self.compute_lddts,
+                backbone_fold_metrics=self.backbone_fold_metrics,
+                allatom_fold_metrics=self.allatom_fold_metrics,
+            )
+            for k, v in scores.items():
+                avg_metrics[k + "_original"] = round(float(v), 5)
+
+            # Sequence diversity:
+            seqs_filtered = [seq[0] for seq in sequences.values()]
+            seq_scores = vendi_sequences(seqs_filtered, self.diversity_subset)
+            for k, v in seq_scores.items():
+                avg_metrics[k] = round(float(v), 5)
+
+        if self.diversity_per_target_original:
+            print("Computing diveristy original per target.")
+            vendi_per_target = []
+            for target_id, ca_list in ca_gen.items():
+                seq_list = sequences[target_id]
+                ca_filtered = [e for e in ca_list if len(e) >= 3]
+                metrics_filtered = [
+                    m
+                    for ca, m in zip(ca_list, input_metrics[target_id])
+                    if len(ca[0]) >= 3
+                ]
+                count = len(ca_filtered)
+                scores = vendi_scores(
+                    ca_filtered,
+                    all_metrics=metrics_filtered,
+                    fold_metrics=fold_metrics,
+                    diversity_subset=self.diversity_subset,
+                    compute_lddts=self.compute_lddts,
+                    backbone_fold_metrics=self.backbone_fold_metrics,
+                    allatom_fold_metrics=self.allatom_fold_metrics,
+                )
+                seq_scores = vendi_sequences(seq_list, self.diversity_subset)
+                scores.update(seq_scores)
+                scores.update(
+                    {
+                        "target_id": target_id,
+                        "num_filtered_ca": count,
+                    }
+                )
+                vendi_per_target.append(scores)
+            df_vendi = pd.DataFrame(vendi_per_target)
+            vendi_csv_path = Path(self.design_dir) / f"vendi_per_target_{self.name}.csv"
+            df_vendi.to_csv(vendi_csv_path, index=False, float_format="%.5f")
+
+            for k in vendi_per_target[0].keys():
+                if isinstance(vendi_per_target[0][k], numbers.Number):
+                    vendis = [e[k] for e in vendi_per_target if not np.isnan(e[k])]
+                    avg_metrics[k + "_mean_per_target"] = float(np.mean(vendis))
+                    avg_metrics[k + "_median_per_target"] = float(np.median(vendis))
+
+                    metrics_data[k] = vendis
+
+        if self.diversity_refolded and (self.fold_metrics):
+            print("Computing diveristy refolded.")
+            sample0 = [ca[0] for ca in ca_refold.values() if ca[0].shape[0] >= 3]
+            metrics_filtered = [
+                m[0]
+                for ca, m in zip(ca_refold.values(), input_metrics.values())
+                if len(ca[0]) >= 3
+            ]
+            scores = vendi_scores(
+                sample0,
+                metrics_filtered,
+                fold_metrics,
+                self.diversity_subset,
+                self.compute_lddts,
+                backbone_fold_metrics=self.backbone_fold_metrics,
+                allatom_fold_metrics=self.allatom_fold_metrics,
+            )
+
+            for k, v in scores.items():
+                avg_metrics[k + "_refolded"] = round(float(v), 5)
+
+        return avg_metrics, metrics_data
+
+    def compute_novelty(self, suffix=None):
+        """
+        Novelty computation using foldseek.
+        This function can be optionally run for computing novelty compared to a reference database which needs to be provided
+        The compute_metrics function writes pdb files which are then used by this function.
+        """
+        avg_metrics = {}
+        metrics_data = {}
+
+        des_pdb_dir = Path(self.des_pdb_dir)
+        design_dir = Path(self.design_dir)
+        des_refold_pdb_dir = Path(self.des_refold_pdb_dir)
+        if not suffix is None:
+            des_pdb_dir = des_pdb_dir / suffix
+            design_dir = design_dir / suffix
+            des_refold_pdb_dir = des_refold_pdb_dir / suffix
+            design_dir.mkdir(exist_ok=True, parents=True)
+
+        # novelty original
+        if self.novelty_original or self.novelty_per_target_original:
+            print("Computing novelty original.")
+            novelty_original_df = compute_novelty_foldseek(
+                indir=des_pdb_dir,
+                outdir=design_dir,
+                reference_db=self.foldseek_db,
+                files=[str(p) for p in des_pdb_dir.glob("*.pdb")],
+                foldseek_binary=self.foldseek_binary,
+            )
+
+        if self.novelty_original:
+            avg_metrics["novelty_original"] = round(
+                float(novelty_original_df["novelty"].mean()), 5
+            )
+
+        if self.novelty_per_target_original:
+            novelty_original_df["target_id"] = novelty_original_df["query"].apply(
+                lambda s: TARGET_ID_RE.match(s).group(1)
+            )
+            nov_df = (
+                novelty_original_df.groupby("target_id")["novelty"].mean().reset_index()
+            )
+            nov_csv = Path(design_dir) / f"novelty_per_target_original_{self.name}.csv"
+            nov_df.to_csv(nov_csv, index=False, float_format="%.5f")
+            avg_metrics["mean_novelty_per_target_original"] = (
+                nov_df["novelty"].mean().round(5)
+            )
+            avg_metrics["median_novelty_per_target_original"] = (
+                nov_df["novelty"].median().round(5)
+            )
+            metrics_data["nov_df"] = nov_df
+
+        # Novelty refolded
+        if (self.novelty_refolded or self.novelty_per_target_refolded) and (
+            self.fold_metrics
+        ):
+            print("Computing novelty refolded.")
+            novelty_refolded_df = compute_novelty_foldseek(
+                indir=des_refold_pdb_dir,
+                outdir=Path(design_dir),
+                reference_db=self.foldseek_db,
+                files=[str(p) for p in des_refold_pdb_dir.glob("*.pdb")],
+                foldseek_binary=self.foldseek_binary,
+            )
+
+        if self.novelty_refolded and (self.fold_metrics):
+            avg_metrics["novelty_refolded"] = round(
+                float(novelty_refolded_df["novelty"].mean()), 5
+            )
+
+        if self.novelty_per_target_refolded and (self.fold_metrics):
+            novelty_refolded_df["target_id"] = novelty_refolded_df["query"].apply(
+                lambda s: TARGET_ID_RE.match(s).group(1)
+            )
+            nov_df_refold = (
+                novelty_refolded_df.groupby("target_id")["novelty"].mean().reset_index()
+            )
+            nov_csv = Path(design_dir) / f"novelty_per_target_refolded_{self.name}.csv"
+            nov_df_refold.to_csv(nov_csv, index=False, float_format="%.5f")
+            avg_metrics["mean_novelty_per_target_refolded"] = (
+                nov_df_refold["novelty"].mean().round(5)
+            )
+            avg_metrics["median_novelty_per_target_refolded"] = round(
+                nov_df_refold["novelty"].median(), 5
+            )
+            metrics_data["nov_df_refold"] = nov_df_refold
+        return avg_metrics, metrics_data
+
+    def run_foldseek_clustering(self, df: pd.DataFrame, suffix=None) -> pd.DataFrame:
+        """
+        Annotates each design with a cluster based on foldseek clustering.
+        This function can be optionally run.
+        The compute_metrics function writes pdb files which are then used by this function.
+        """
+        des_pdb_dir = self.des_pdb_dir / suffix if suffix else self.des_pdb_dir
+        design_dir = Path(self.design_dir) / suffix if suffix else Path(self.design_dir)
+
+        cluster_output_dir = design_dir / "foldseek_cluster"
+        cluster_output_dir.mkdir(parents=True, exist_ok=True)
+
+        cluster_prefix = cluster_output_dir / "cluster"
+        tmp_dir = cluster_output_dir / "tmp"
+
+        min_num_design = int(df["num_design"].min())
+        cmd = [
+            self.foldseek_binary,
+            "easy-cluster",
+            str(des_pdb_dir),
+            str(cluster_prefix),
+            str(tmp_dir),
+            "--alignment-type",
+            "1",
+            "--cov-mode",
+            "0",
+            "--min-seq-id",
+            "0",
+            "--tmscore-threshold",
+            "0.5",
+        ]
+        if min_num_design < 20:
+            msg = f"[FoldSeek] Using --kmer-per-seq {2} due to short designs."
+            print(msg)
+            cmd += ["--kmer-per-seq", str(2)]
+
+        try:
+            subprocess.run(cmd, check=True)
+
+            df_cluster = pd.read_csv(
+                str(cluster_prefix) + "_cluster.tsv",
+                sep="\t",
+                header=None,
+                names=["file", "clusters_05_tmscore"],
+            )
+            df_cluster["file"] = df_cluster["file"].apply(lambda x: Path(x).stem)
+
+            df = df.merge(df_cluster, left_on="id", right_on="file", how="left").drop(
+                columns=["file"]
+            )
+        except Exception as e:
+            msg = f"Structure clustering was unsuccessful. No cluster labels are added to the dataframe / csv file output."
+            print(msg)
+        return df
+
+    def make_histograms(self, all_metrics):
+        df = pd.DataFrame(all_metrics)
+
+        # Make aggregate histograms
+        histograms = {}
+        cols = [
+            "delta_sasa_refolded",
+            "rmsd",
+            "iptm",
+            "ptm",
+            "rmsd",
+            "design_ptm",
+            "min_design_to_target_pae",
+            "helix",
+            "sheet",
+            "loop",
+            "plip_saltbridge",
+            "plip_hbonds",
+            "design_sasa_bound_original",
+            "design_sasa_unbound_original",
+            "delta_sasa_original",
+            "num_design",
+            "precision_loop",
+            "recall_loop",
+            "precision_helix",
+            "recall_helix",
+            "precision_sheet",
+            "recall_sheet",
+            "accuracy_overall",
+            "liability_score",
+            "liability_num_violations",
+        ]
+        for col in cols:
+            if col in df.columns:
+                histograms["hist" + col] = make_histogram(df, col)
+
+        # make per target histograms
+        df["target_id"] = df["id"].apply(lambda s: TARGET_ID_RE.match(s).group(1))
+        per_target_df = df.groupby("target_id").mean(numeric_only=True).reset_index()
+        cols += ["rmsd<2.5", "designability_rmsd_2"]
+        if self.compute_lddts:
+            cols += [
+                "designability_lddt_90",
+                "designability_lddt_85",
+                "designability_lddt_80",
+                "designability_lddt_75",
+                "designability_lddt_70",
+                "designability_lddt_65",
+                "designability_lddt_60",
+            ]
+        for col in cols:
+            if col in df.columns:
+                # Per target histograms
+                histograms["per_target" + col] = make_histogram(per_target_df, col)
+        return df, histograms

model/boltzgen/task/analyze/analyze_utils.py

@@ -0,0 +1,1260 @@
+import io
+import warnings
+from pathlib import Path
+import random
+from typing import List
+import subprocess
+import re
+import biotite
+import hydride
+from sklearn.cluster import DBSCAN
+from Bio import PDB
+from biotite import structure
+from Bio.Seq import Seq
+
+
+from matplotlib import pyplot as plt
+import matplotlib as mpl
+
+mpl.rcParams["figure.max_open_warning"] = 100
+
+from Bio import Align
+
+from boltzgen.data.rmsd_computation import get_true_coordinates
+from boltzgen.model.loss.diffusion import weighted_rigid_align
+from boltzgen.task.predict.data_from_generated import collate
+
+import torch
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+from boltzgen.data import const
+from boltzgen.data.data import Structure, biotite_array_from_feat
+from boltzgen.data.write.mmcif import to_mmcif
+from boltzgen.model.loss.validation import factored_lddt_loss, compute_subset_rmsd
+
+from biotite.structure.sasa import sasa
+from biotite.structure.info import vdw_radius_single, vdw_radius_protor
+import biotite.structure.io.pdbx as pdbx
+import biotite.structure.io.pdb as pdbio
+
+TARGET_ID_RE = re.compile(
+    r"^(?:(?:sample\d+_|batch\d+_|rank\d+_)+)?([^_]+)(?:_[^_]+)*?(?:_(?:gen))*$"
+)
+
+
+def _load_stack(path):
+    path = Path(path)
+    suffix = path.suffix.lower()
+    if suffix in {".cif", ".mmcif"}:
+        cif_file = pdbx.CIFFile.read(str(path))
+        stack = pdbx.get_structure(cif_file, use_author_fields=False)
+    elif suffix in {".pdb", ".ent"}:
+        pdb_file = pdbio.PDBFile.read(str(path))
+        stack = pdbio.get_structure(pdb_file, model=None)
+    else:
+        raise ValueError(f"Unsupported structure file extension: {suffix}")
+
+    return stack
+
+
+def compute_rmsd(atom_coords: torch.Tensor, pred_atom_coords: torch.Tensor):
+    rmsd, _ = compute_subset_rmsd(
+        atom_coords,
+        pred_atom_coords,
+        atom_mask=torch.ones_like(atom_coords[..., 0]),
+        align_weights=torch.ones_like(atom_coords[..., 0]),
+        subset_mask=torch.ones_like(atom_coords[..., 0]),
+        multiplicity=1,
+    )
+    return rmsd
+
+
+def make_histogram(
+    df,
+    column_name: str,
+):
+    data = df[column_name].dropna()
+    fig, ax = plt.subplots(figsize=(6, 4))
+
+    ax.hist(data, bins=50, alpha=0.7, edgecolor="black", linewidth=0.5)
+    ax.axvline(data.mean(), color="red", linestyle="dashed", linewidth=1)
+
+    ax.set_title(
+        f"{column_name.replace('_', ' ').capitalize()} Distribution", fontsize=12
+    )
+    ax.set_xlabel(column_name.replace("_", " "), fontsize=10)
+    ax.set_ylabel("Count", fontsize=10)
+
+    ax.grid(True, axis="y", linestyle="--", linewidth=0.5, alpha=0.6)
+    ax.tick_params(axis="both", which="major", labelsize=8)
+
+    plt.tight_layout()
+    return fig
+
+
+def get_best_folding_sample(folded):
+    confidence = 0.8 * folded["design_to_target_iptm"] + 0.2 * folded["design_ptm"]
+    best_idx = np.argmax(confidence)
+
+    # TODO: remove the "if k in folded"
+    best_sample = {
+        k: folded[k][best_idx] for k in const.eval_keys_confidence if k in folded
+    }
+    best_sample["coords"] = folded["coords"][best_idx]
+    return best_sample
+
+
+def get_fold_metrics(
+    feat,
+    folded,
+    compute_lddts=True,
+    prefix="",
+):
+    batch = collate([feat])
+    diffusion_samples = batch["coords"].shape[0]
+    best_sample = get_best_folding_sample(folded)
+
+    # Compute RMSDs
+    rmsd_out = get_true_coordinates(
+        batch=batch,
+        out={"sample_atom_coords": torch.from_numpy(best_sample["coords"])},
+        diffusion_samples=1,
+        symmetry_correction=False,
+        protein_lig_rmsd=True,
+    )
+    true_coords_resolved_mask = rmsd_out["true_coords_resolved_mask"]
+
+    # Add to metrics dictionary
+    metrics = {}
+    metrics["rmsd"] = rmsd_out.get("rmsd").item()
+    metrics["rmsd_design"] = rmsd_out.get("rmsd_design").item()
+    metrics["rmsd_target"] = rmsd_out.get("rmsd_target").item()
+    metrics["rmsd_design_target"] = rmsd_out.get("rmsd_design_target").item()
+    metrics["target_aligned_rmsd_design"] = rmsd_out.get(
+        "target_aligned_rmsd_design"
+    ).item()
+    metrics["rmsd<2.5"] = bool(metrics["rmsd"] <= 2.5)
+    metrics["target_aligned<2.5"] = bool(metrics["target_aligned_rmsd_design"] <= 2.5)
+    metrics["designability_rmsd_2"] = bool(metrics["rmsd_design"] <= 2.0)
+    metrics["designability_rmsd_4"] = bool(metrics["rmsd_design"] <= 4.0)
+
+    # Comput LDDTs
+    if compute_lddts:
+        all_lddt_dict, _ = factored_lddt_loss(
+            feats=batch,
+            atom_mask=true_coords_resolved_mask,
+            true_atom_coords=batch["coords"],
+            pred_atom_coords=torch.from_numpy(best_sample["coords"]),
+            multiplicity=diffusion_samples,
+            exclude_ions=False,
+        )
+        metrics.update({f"lddt_{k}": v.max().item() for k, v in all_lddt_dict.items()})
+        metrics["designability_lddt_60"] = bool(metrics["lddt_intra_design"] >= 0.6)
+        metrics["designability_lddt_65"] = bool(metrics["lddt_intra_design"] >= 0.65)
+        metrics["designability_lddt_70"] = bool(metrics["lddt_intra_design"] >= 0.7)
+        metrics["designability_lddt_75"] = bool(metrics["lddt_intra_design"] >= 0.75)
+        metrics["designability_lddt_80"] = bool(metrics["lddt_intra_design"] >= 0.8)
+        metrics["designability_lddt_85"] = bool(metrics["lddt_intra_design"] >= 0.85)
+        metrics["designability_lddt_90"] = bool(metrics["lddt_intra_design"] >= 0.9)
+
+    # metrics without prefix (backbone only is the same as all atom)
+    # TODO: remove the "if k in best_sample"
+    confs = {k: best_sample[k] for k in const.eval_keys_confidence if k in best_sample}
+    confs["min_interaction_pae<1.5"] = bool(confs["min_interaction_pae"] <= 1.5)
+    confs["min_interaction_pae<2"] = bool(confs["min_interaction_pae"] <= 2.0)
+    confs["min_interaction_pae<2.5"] = bool(confs["min_interaction_pae"] <= 2.5)
+    confs["min_interaction_pae<3"] = bool(confs["min_interaction_pae"] <= 3)
+    confs["min_interaction_pae<4"] = bool(confs["min_interaction_pae"] <= 4)
+    confs["min_interaction_pae<5"] = bool(confs["min_interaction_pae"] <= 5)
+    confs["design_ptm>80"] = bool(confs["design_ptm"] >= 0.8)
+    confs["design_ptm>75"] = bool(confs["design_ptm"] >= 0.75)
+    confs["design_iptm>80"] = bool(confs["design_iptm"] >= 0.8)
+    confs["design_iptm>70"] = bool(confs["design_iptm"] >= 0.7)
+    confs["design_iptm>60"] = bool(confs["design_iptm"] >= 0.6)
+    confs["design_iptm>50"] = bool(confs["design_iptm"] >= 0.5)
+
+    prefixed_metrics = {f"{prefix}{k}": v for k, v in metrics.items()}
+    prefixed_metrics.update(confs)
+    return prefixed_metrics
+
+
+def count_noncovalents(feat):
+    metrics = {}
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        biotite_array = biotite_array_from_feat(feat)
+        biotite_array, _ = hydride.add_hydrogen(biotite_array)
+        hbond = biotite.structure.hbond(biotite_array)
+    donor_idxs, acceptor_idxs = hbond[:, 0], hbond[:, 2]
+    donor_design_hbonds = int(
+        (
+            biotite_array.is_design[donor_idxs]
+            & ~biotite_array.is_chain_design[acceptor_idxs]
+        ).sum()
+    )
+    acceptor_design_hbonds = int(
+        (
+            ~biotite_array.is_chain_design[donor_idxs]
+            & biotite_array.is_design[acceptor_idxs]
+        ).sum()
+    )
+    metrics["plip_hbonds"] = donor_design_hbonds + acceptor_design_hbonds
+
+    # saltbridges
+    pos_atoms = biotite_array[biotite_array.charge > 0]
+    neg_atoms = biotite_array[biotite_array.charge < 0]
+    if len(neg_atoms) > 0 and len(pos_atoms) > 0:
+        pos_neg_distances = torch.cdist(
+            torch.as_tensor(pos_atoms.coord), torch.as_tensor(neg_atoms.coord)
+        )
+        pos_idxs, neg_idxs = torch.where(
+            (pos_neg_distances > 0.5) & (pos_neg_distances < 5.5)
+        )
+        # only keep the ones between design and non design
+        pos_design_sb = int(
+            (pos_atoms.is_design[pos_idxs] & ~neg_atoms.is_chain_design[neg_idxs]).sum()
+        )
+        neg_design_sb = int(
+            (~pos_atoms.is_chain_design[pos_idxs] & neg_atoms.is_design[neg_idxs]).sum()
+        )
+        metrics["plip_saltbridge"] = pos_design_sb + neg_design_sb
+    else:
+        metrics["plip_saltbridge"] = 0
+    return metrics
+
+
+def tm_score(coords1, coords2):
+    num_atoms1 = coords1.shape[0]
+    num_atoms2 = coords2.shape[0]
+
+    atom_array1 = structure.AtomArray(num_atoms1)
+    atom_array1.coord = coords1.numpy()
+    atom_array1.element = np.array(["C"] * num_atoms1)
+    atom_array1.atom_name = np.array(["CA"] * num_atoms1)
+    atom_array1.res_name = np.array(["ALA"] * num_atoms1)
+    atom_array1.chain_id = np.array(["A"] * num_atoms1)
+    atom_array1.res_id = np.arange(1, num_atoms1 + 1)
+
+    atom_array2 = structure.AtomArray(num_atoms2)
+    atom_array2.coord = coords2.numpy()
+    atom_array2.element = np.array(["C"] * num_atoms2)
+    atom_array2.atom_name = np.array(["CA"] * num_atoms2)
+    atom_array2.res_name = np.array(["ALA"] * num_atoms2)
+    atom_array2.chain_id = np.array(["A"] * num_atoms2)
+    atom_array2.res_id = np.arange(1, num_atoms2 + 1)
+
+    try:
+        # This fails with a value error if the structures are too dissimilar. In that event, we return 0 as the TM-Score
+        aligned, transform, fixed_indices, mobile_indices = (
+            structure.superimpose_structural_homologs(
+                atom_array1, atom_array2, max_iterations=25
+            )
+        )
+        tm_align_fixed = structure.tm_score(
+            atom_array1,
+            aligned,
+            fixed_indices,
+            mobile_indices,
+        )
+    except:
+        tm_align_fixed = 0
+
+    tm_score_rmsd_aligned = 0
+    if num_atoms1 == num_atoms2:
+        coords1 = weighted_rigid_align(
+            coords1.float()[None],
+            coords2.float()[None],
+            weights=torch.ones(len(coords1)).float()[None],
+            mask=torch.ones(len(coords2))[None],
+        ).squeeze()
+
+        atom_array1 = structure.AtomArray(num_atoms1)
+        atom_array1.coord = coords1.numpy()
+        atom_array1.element = np.array(["C"] * num_atoms1)
+        atom_array1.atom_name = np.array(["CA"] * num_atoms1)
+        atom_array1.res_name = np.array(["ALA"] * num_atoms1)
+        atom_array1.chain_id = np.array(["A"] * num_atoms1)
+        atom_array1.res_id = np.arange(1, num_atoms1 + 1)
+
+        atom_array2 = structure.AtomArray(num_atoms2)
+        atom_array2.coord = coords2.numpy()
+        atom_array2.element = np.array(["C"] * num_atoms2)
+        atom_array2.atom_name = np.array(["CA"] * num_atoms2)
+        atom_array2.res_name = np.array(["ALA"] * num_atoms2)
+        atom_array2.chain_id = np.array(["A"] * num_atoms2)
+        atom_array2.res_id = np.arange(1, num_atoms2 + 1)
+        try:
+            _, _, fixed_indices, mobile_indices = (
+                structure.superimpose_structural_homologs(
+                    atom_array1, atom_array2, max_iterations=25
+                )
+            )
+            tm_score_rmsd_aligned = structure.tm_score(
+                atom_array1,
+                atom_array2,
+                fixed_indices,
+                mobile_indices,
+            )
+        except:
+            pass
+
+    return tm_score_rmsd_aligned, tm_align_fixed
+
+
+def vendi_from_sim(mat):
+    mat = mat + mat.T
+    np.fill_diagonal(mat, 1.0)
+    eigvals, _ = np.linalg.eigh(mat / len(mat))
+    eigvals = np.clip(eigvals, 0.0, None)
+    return np.exp(np.nansum(-(eigvals * np.log(eigvals))))
+
+
+def vendi_scores(
+    all_ca_coords: List[np.ndarray],
+    all_metrics: list = None,
+    fold_metrics: bool = False,
+    diversity_subset: int = None,
+    compute_lddts: bool = True,
+    compute_iptms: bool = True,
+    compute_min_int_paes: bool = True,
+    backbone_fold_metrics: bool = False,
+    allatom_fold_metrics: bool = True,
+) -> float:
+    if fold_metrics or diversity_subset is not None:
+        assert all_metrics is not None
+    if all_metrics is not None:
+        assert len(all_ca_coords) == len(all_metrics)
+    if diversity_subset is not None and diversity_subset < len(all_ca_coords):
+        indices = random.sample(range(len(all_ca_coords)), diversity_subset)
+        all_metrics = [all_metrics[i] for i in indices]
+        all_ca_coords = [all_ca_coords[i] for i in indices]
+    N = len(all_ca_coords)
+    tm = np.zeros((N, N), dtype=np.float32)
+    tm_fixed = np.zeros((N, N), dtype=np.float32)
+
+    for i in tqdm(range(N), desc="Computing structure diversity."):
+        for j in range(i + 1, N):
+            tm_score_rmsd_aligned, tm_fixeds = tm_score(
+                all_ca_coords[i], all_ca_coords[j]
+            )
+            tm[i, j] = tm_score_rmsd_aligned
+            tm_fixed[i, j] = tm_fixeds
+
+    scores = {
+        "vendi_tm_fixed": vendi_from_sim(tm_fixed),
+        "vendi_tm_align": vendi_from_sim(tm),
+    }
+    prefixes = []
+    if allatom_fold_metrics:
+        prefixes.append("")
+    if backbone_fold_metrics:
+        prefixes.append("bb_")
+    for prefix in prefixes:
+        mask_2 = np.array([m[f"{prefix}designability_rmsd_2"] for m in all_metrics])
+        mask_4 = np.array([m[f"{prefix}designability_rmsd_4"] for m in all_metrics])
+        mask_25 = np.array([m[f"{prefix}rmsd<2.5"] for m in all_metrics])
+        mask_target_25 = np.array(
+            [m[f"{prefix}target_aligned<2.5"] for m in all_metrics]
+        )
+        scores.update(
+            {
+                f"vendi_tm_{prefix}rmsd<2.5": vendi_from_sim(
+                    tm_fixed[mask_25][:, mask_25]
+                )
+                if np.sum(mask_25) > 0
+                else 0.0,
+                f"vendi_tm_{prefix}rmsd_2": vendi_from_sim(tm_fixed[mask_2][:, mask_2])
+                if np.sum(mask_2) > 0
+                else 0.0,
+                f"vendi_tm_{prefix}rmsd_4": vendi_from_sim(tm_fixed[mask_4][:, mask_4])
+                if np.sum(mask_4) > 0
+                else 0.0,
+                f"vendi_tm_{prefix}target_aligned_rmsd<2.5": vendi_from_sim(
+                    tm_fixed[mask_target_25][:, mask_target_25]
+                )
+                if np.sum(mask_target_25) > 0
+                else 0.0,
+            }
+        )
+
+        if compute_lddts:
+            mask_60 = np.array(
+                [m[f"{prefix}designability_lddt_60"] for m in all_metrics]
+            )
+            mask_65 = np.array(
+                [m[f"{prefix}designability_lddt_65"] for m in all_metrics]
+            )
+            mask_70 = np.array(
+                [m[f"{prefix}designability_lddt_70"] for m in all_metrics]
+            )
+            mask_75 = np.array(
+                [m[f"{prefix}designability_lddt_75"] for m in all_metrics]
+            )
+            mask_80 = np.array(
+                [m[f"{prefix}designability_lddt_80"] for m in all_metrics]
+            )
+            mask_85 = np.array(
+                [m[f"{prefix}designability_lddt_85"] for m in all_metrics]
+            )
+            mask_90 = np.array(
+                [m[f"{prefix}designability_lddt_90"] for m in all_metrics]
+            )
+            scores.update(
+                {
+                    f"vendi_tm_{prefix}lddt_60": vendi_from_sim(
+                        tm_fixed[mask_60][:, mask_60]
+                    )
+                    if np.sum(mask_60) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}lddt_65": vendi_from_sim(
+                        tm_fixed[mask_65][:, mask_65]
+                    )
+                    if np.sum(mask_65) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}lddt_70": vendi_from_sim(
+                        tm_fixed[mask_70][:, mask_70]
+                    )
+                    if np.sum(mask_70) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}lddt_75": vendi_from_sim(
+                        tm_fixed[mask_75][:, mask_75]
+                    )
+                    if np.sum(mask_75) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}lddt_80": vendi_from_sim(
+                        tm_fixed[mask_80][:, mask_80]
+                    )
+                    if np.sum(mask_80) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}lddt_85": vendi_from_sim(
+                        tm_fixed[mask_85][:, mask_85]
+                    )
+                    if np.sum(mask_85) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}lddt_90": vendi_from_sim(
+                        tm_fixed[mask_90][:, mask_90]
+                    )
+                    if np.sum(mask_90) > 0
+                    else 0.0,
+                }
+            )
+        if compute_iptms:
+            mask_80 = np.array([m[f"design_iptm>80"] for m in all_metrics])
+            mask_70 = np.array([m[f"design_iptm>70"] for m in all_metrics])
+            mask_60 = np.array([m[f"design_iptm>60"] for m in all_metrics])
+            mask_50 = np.array([m[f"design_iptm>50"] for m in all_metrics])
+            scores.update(
+                {
+                    f"vendi_tm_{prefix}design_iptm_80": vendi_from_sim(
+                        tm_fixed[mask_80][:, mask_80]
+                    )
+                    if np.sum(mask_80) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}design_iptm_70": vendi_from_sim(
+                        tm_fixed[mask_70][:, mask_70]
+                    )
+                    if np.sum(mask_70) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}design_iptm_60": vendi_from_sim(
+                        tm_fixed[mask_60][:, mask_60]
+                    )
+                    if np.sum(mask_60) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}design_iptm_50": vendi_from_sim(
+                        tm_fixed[mask_50][:, mask_50]
+                    )
+                    if np.sum(mask_50) > 0
+                    else 0.0,
+                }
+            )
+        if compute_min_int_paes:
+            mask_15 = np.array([m[f"min_interaction_pae<1.5"] for m in all_metrics])
+            mask_2 = np.array([m[f"min_interaction_pae<2"] for m in all_metrics])
+            mask_25 = np.array([m[f"min_interaction_pae<2.5"] for m in all_metrics])
+            mask_3 = np.array([m[f"min_interaction_pae<3"] for m in all_metrics])
+            mask_4 = np.array([m[f"min_interaction_pae<4"] for m in all_metrics])
+            mask_5 = np.array([m[f"min_interaction_pae<5"] for m in all_metrics])
+            scores.update(
+                {
+                    f"vendi_tm_{prefix}min_interaction_pae_1.5": vendi_from_sim(
+                        tm_fixed[mask_15][:, mask_15]
+                    )
+                    if np.sum(mask_15) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}min_interaction_pae_2": vendi_from_sim(
+                        tm_fixed[mask_2][:, mask_2]
+                    )
+                    if np.sum(mask_2) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}min_interaction_pae_2.5": vendi_from_sim(
+                        tm_fixed[mask_25][:, mask_25]
+                    )
+                    if np.sum(mask_25) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}min_interaction_pae_3": vendi_from_sim(
+                        tm_fixed[mask_3][:, mask_3]
+                    )
+                    if np.sum(mask_3) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}min_interaction_pae_4": vendi_from_sim(
+                        tm_fixed[mask_4][:, mask_4]
+                    )
+                    if np.sum(mask_4) > 0
+                    else 0.0,
+                    f"vendi_tm_{prefix}min_interaction_pae_5": vendi_from_sim(
+                        tm_fixed[mask_5][:, mask_5]
+                    )
+                    if np.sum(mask_5) > 0
+                    else 0.0,
+                }
+            )
+
+    return scores
+
+
+def vendi_sequences(all_seqs: List[np.ndarray], diversity_subset: int = None) -> float:
+    if diversity_subset is not None and diversity_subset < len(all_seqs):
+        all_seqs = random.sample(all_seqs, diversity_subset)
+
+    N = len(all_seqs)
+    sims = np.zeros((N, N), dtype=np.float32)
+    aligner = Align.PairwiseAligner()
+    for i in tqdm(range(N), desc="Computing sequence diversity."):
+        for j in range(i + 1, N):
+            seq1 = Seq(all_seqs[i])
+            seq2 = Seq(all_seqs[j])
+            alignments = aligner.align(seq1, seq2)
+
+            similarity = alignments[0].score / max(len(seq1), len(seq2))
+            sims[i, j] = similarity
+
+    return {
+        "vendi_seq_sim": vendi_from_sim(sims),
+    }
+
+
+def compute_novelty_foldseek(
+    indir: Path,
+    outdir: Path,
+    reference_db: Path,
+    files: List[str],
+    foldseek_binary: str = "/data/rbg/users/hstark/foldseek/bin/foldseek",
+) -> pd.DataFrame:
+    if len(files) == 0:
+        return np.nan
+
+    aln_tsv = outdir / "aln.tsv"
+    tmp_dir = outdir / "tmp"
+
+    cmd = [
+        foldseek_binary,
+        "easy-search",
+        str(indir),
+        str(reference_db),
+        str(aln_tsv),
+        str(tmp_dir),
+        "--format-output",
+        "query,target,alntmscore,qtmscore,ttmscore",
+        "--alignment-type",
+        "1",
+        "--exhaustive-search",
+        "1",
+    ]
+
+    subprocess.run(cmd, check=True)
+
+    df = pd.read_csv(
+        aln_tsv,
+        sep="\t",
+        names=["query", "target", "alntmscore", "qtmscore", "ttmscore"],
+    )
+    df["tmscore"] = (df["qtmscore"] + df["ttmscore"]) / 2
+    df = df.groupby("query").max().reset_index()
+    queries = [Path(f).stem for f in files]
+    df = df.set_index("query").reindex(queries, fill_value=0.0).reset_index()
+    df_novelty = df[["query", "tmscore"]].rename(columns={"tmscore": "novelty"})
+    return df_novelty
+
+
+def _radius(res_name: str, atom_name: str, element: str) -> float:
+    """
+    ProtOr radius with element fallback.
+    """
+    try:
+        r = vdw_radius_protor(res_name, atom_name)
+        if r is not None:
+            return r
+    except KeyError:
+        pass
+    r = vdw_radius_single(element)
+    return r if r is not None else 1.8
+
+
+def compute_sasa(structure_path):
+    HYDROPHOBIC_RESIDUES = {"ALA", "VAL", "LEU", "ILE", "MET", "PHE", "PRO", "TRP"}
+    atoms = _load_stack(structure_path)[0]
+    res_names = [
+        bytes(r).decode() if isinstance(r, bytes) else str(r) for r in atoms.res_name
+    ]
+    atom_names = [
+        bytes(a).decode() if isinstance(a, bytes) else str(a) for a in atoms.atom_name
+    ]
+    elements = [
+        bytes(e).decode() if isinstance(e, bytes) else str(e) for e in atoms.element
+    ]
+
+    radii = np.array(
+        [
+            _radius(rn.strip(), an.strip(), el)
+            for rn, an, el in zip(res_names, atom_names, elements)
+        ],
+        dtype=float,
+    )
+    atom_sasa = sasa(atoms, probe_radius=1.4, point_number=960, vdw_radii=radii)
+    mask = (
+        np.array([rn.strip() in HYDROPHOBIC_RESIDUES for rn in res_names])
+        & np.char.startswith(atom_names, "C")
+        & (atom_sasa > 0)
+    )
+
+    return atoms.coord[mask], atom_sasa[mask]
+
+
+def largest_hydrophobic_patch_area(cif_path, distance_cutoff=6.0):
+    result = compute_sasa(cif_path)
+    if result is None or result[0].size == 0:
+        return np.nan
+
+    coords, sasa_vals = result
+    if len(coords) == 0:
+        return 0.0
+    clustering = DBSCAN(eps=distance_cutoff, min_samples=1).fit(coords)
+    labels = clustering.labels_
+
+    max_patch_area = 0.0
+    for label in np.unique(labels):
+        area = sasa_vals[labels == label].sum()
+        max_patch_area = max(max_patch_area, area)
+
+    return max_patch_area
+
+
+def get_delta_sasa(
+    path,
+    atom_target_mask,                
+    atom_design_mask,           
+):
+    stack = _load_stack(path)
+    atoms = stack[0]
+
+    res = [
+        r.decode().strip() if isinstance(r, bytes) else str(r).strip()
+        for r in atoms.res_name
+    ]
+    atm = [
+        a.decode().strip() if isinstance(a, bytes) else str(a).strip()
+        for a in atoms.atom_name
+    ]
+    elem = [e.decode() if isinstance(e, bytes) else str(e) for e in atoms.element]
+
+    radii = np.array(
+        [_radius(rn, an, el) for rn, an, el in zip(res, atm, elem)], dtype=float
+    )
+
+    
+    bound_mask = atom_design_mask | atom_target_mask
+    atoms_bound = atoms[bound_mask]
+    radii_bound = radii[bound_mask]
+
+    area_bound = sasa(
+        atoms_bound,
+        probe_radius=1.4,
+        point_number=960,
+        vdw_radii=radii_bound,
+    )
+    
+    target_in_bound = atom_target_mask[bound_mask]
+    target_bound    = area_bound[target_in_bound].sum()
+    
+    
+
+    target_atoms = atoms[atom_target_mask]
+    target_res = [r for r, m in zip(res, atom_target_mask) if m]
+    target_atm = [a for a, m in zip(atm, atom_target_mask) if m]
+    target_elem = [e for e, m in zip(elem, atom_target_mask) if m]
+
+    radii_lig = np.array(
+        [_radius(rn, an, el) for rn, an, el in zip(target_res, target_atm, target_elem)],
+        dtype=float,
+    )
+    target_area = sasa(
+        target_atoms,
+        probe_radius=1.4,
+        point_number=960,
+        vdw_radii=radii_lig,
+    )
+    delta = target_area.sum() - target_bound
+    return delta, target_area.sum(), target_bound
+
+
+def compute_ss_metrics(dssp_pred, ss_conditioning_metricsed):
+    ss_metrics = {}
+    conditioned_mask = ss_conditioning_metricsed != 0
+    if conditioned_mask.sum() == 0:
+        return {
+            "precision_loop": float("nan"),
+            "recall_loop": float("nan"),
+            "accuracy_loop": float("nan"),
+            "precision_helix": float("nan"),
+            "recall_helix": float("nan"),
+            "accuracy_helix": float("nan"),
+            "precision_sheet": float("nan"),
+            "recall_sheet": float("nan"),
+            "accuracy_sheet": float("nan"),
+            "accuracy_overall": float("nan"),
+        }
+    types = {1: "loop", 2: "helix", 3: "sheet"}
+    TP_total, total_conditioned = 0, conditioned_mask.sum().item()
+    for i, name in types.items():
+        TP = ((dssp_pred == i) & (ss_conditioning_metricsed == i)).sum().item()
+        FP = (
+            ((dssp_pred == i) & (ss_conditioning_metricsed != i) & conditioned_mask)
+            .sum()
+            .item()
+        )
+        FN = ((dssp_pred != i) & (ss_conditioning_metricsed == i)).sum().item()
+        precision = TP / (TP + FP) if (TP + FP) > 0 else float("nan")
+        recall = TP / (TP + FN) if (TP + FN) > 0 else float("nan")
+        accuracy = TP / (TP + FP + FN) if (TP + FP + FN) > 0 else float("nan")
+        ss_metrics[f"precision_{name}"] = precision
+        ss_metrics[f"recall_{name}"] = recall
+        ss_metrics[f"accuracy_{name}"] = accuracy
+        TP_total += TP
+
+    accuracy_overall = (
+        TP_total / total_conditioned if total_conditioned > 0 else float("nan")
+    )
+    ss_metrics["accuracy_overall"] = accuracy_overall
+    return ss_metrics
+
+
+def filter_resolved_atoms(structure: Structure) -> Structure:
+    resolved_atom_indices = np.where(structure.atoms["is_present"])[0]
+    return Structure.extract_atoms(structure, resolved_atom_indices)
+
+
+def save_design_only_structure_to_cif(atom_design_mask, structure, output_path: Path):
+    design_atom_indices = torch.where(atom_design_mask)[0].cpu().numpy()
+    design_only_str = Structure.extract_atoms(
+        structure, design_atom_indices, res_reindex=True
+    )
+    cif_text = to_mmcif(design_only_str)
+    output_path.write_text(cif_text)
+    return cif_text
+
+
+def save_design_only_structure_to_pdb(atom_design_mask, structure, output_path: Path):
+    cif_path = output_path.with_suffix(".cif")
+    cif_text = save_design_only_structure_to_cif(atom_design_mask, structure, cif_path)
+
+    cif_io = io.StringIO(cif_text)
+    mmcif_parser = PDB.MMCIFParser()
+    pdb_writer = PDB.PDBIO()
+    parsed_structure = mmcif_parser.get_structure("des_only", cif_io)
+    pdb_writer.set_structure(parsed_structure)
+    pdb_writer.save(str(output_path))
+
+
+########################################################################################################
+# Hydrophobycity computation functions. From Jeremie Alexander.
+########################################################################################################
+
+
+def calc_base_h(seq: str) -> float:
+    s = seq.upper()
+    n = len(s)
+    if n == 0:
+        return 0.0
+    H = 0.0
+
+    # position-specific coefficients
+    for i, aa in enumerate(s):
+        if aa not in const.hydrophobicity_info:
+            raise ValueError(f"Unknown residue '{aa}' in '{seq}'")
+        if i == 0:
+            key = "Rc1"
+        elif i == 1:
+            key = "Rc2"
+        elif i == n - 1:
+            key = "Rn"
+        elif i == n - 2:
+            key = "Rn1"
+        else:
+            key = "Rc"
+        H += const.hydrophobicity_info[aa][key]
+
+    # nearest-neighbor penalties around H/R/K
+    for i, aa in enumerate(s):
+        if aa in ("H", "R", "K"):
+            for j in (i - 1, i + 1):
+                if 0 <= j < n and s[j] in const.nn_penalty:
+                    H -= const.nn_penalty[s[j]]
+
+    # proline run penalties
+    i = 0
+    while i < n:
+        if s[i] == "P":
+            j = i
+            while j < n and s[j] == "P":
+                j += 1
+            run = j - i
+            if run >= 4:
+                H -= 5.0
+            elif run == 3:
+                H -= 3.5
+            elif run == 2:
+                H -= 1.2
+            i = j
+        else:
+            i += 1
+    return H
+
+
+def apply_length_weight(H: float, n: int) -> float:
+    if n < 8:
+        KL = 1.0 - 0.055 * (8 - n)
+    elif n > 20:
+        KL = 1.0 / (1.0 + 0.027 * (n - 20))
+    else:
+        KL = 1.0
+    return H * KL
+
+
+def overall_penalty(H: float) -> float:
+    if H <= 20:
+        return H
+    if H <= 30:
+        return H - 0.27 * (H - 18.0)
+    if H <= 40:
+        return H - 0.33 * (H - 18.0)
+    if H <= 50:
+        return H - 0.38 * (H - 18.0)
+    return H - 0.447 * (H - 18.0)
+
+
+def calc_hydrophobicity(seq: str) -> float:
+    s = (seq or "").strip().upper()
+    if not s or "X" in s:
+        return float("nan")
+    base = calc_base_h(s)
+    base = apply_length_weight(base, len(s))
+    return round(overall_penalty(base), 4)
+
+
+########################################################################################################
+# Liability computation functions
+########################################################################################################
+
+
+def get_motif_set(modality="antibody", peptide_type="linear"):
+    """
+    Returns a dict of motif_name -> compiled regex for the given modality.
+    modality: 'antibody' or 'peptide'
+    peptide_type: 'linear' or 'cyclic' (only for peptide modality)
+
+
+    """
+    if modality == "antibody":
+        return {
+            "DeAmdH": re.compile(r"N[GS]"),  # High-risk deamidation
+            "DeAmdM": re.compile(r"N[AHNT]"),  # Medium-risk deamidation
+            "DeAmdL": re.compile(r"[STK]N"),  # Low-risk deamidation
+            "Ngly": re.compile(r"N[^P][ST]"),  # N-glycosylation sites
+            "Isom": re.compile(r"D[DG HST]".replace(" ", "")),  # Isomerization
+            "Isomer": re.compile(r"DG|DS|DD"),  # Isomerization variants
+            "FragH": re.compile(r"DP"),  # High fragmentation risk
+            "FragM": re.compile(r"TS"),  # Medium fragmentation risk
+            "TrpOx": re.compile(r"W"),  # Tryptophan oxidation
+            "MetOx": re.compile(r"M"),  # Methionine oxidation
+            "Hydro": re.compile(r"NP"),  # Hydrolysis prone
+            "IntBind": re.compile(r"GPR|RGD|RYD|LDV|DGE|KGD|NGR"),  # Integrin binding
+            "Polyreactive": re.compile(
+                r"GGG|GG|RR|VG|VVV|WWW|YY|WxW"
+            ),  # Polyreactivity
+            "AggPatch": re.compile(r"FHW"),  # Aggregation patches
+            "ViscPatch": re.compile(r"HYF|HWH"),  # Viscosity patches
+            "DeAmdH": re.compile(r"N[GS]"),
+            "DeAmdM": re.compile(r"N[AHNT]"),
+            "DeAmdL": re.compile(r"[STK]N"),
+            "Ngly": re.compile(r"N[^P][ST]"),
+            "Isom": re.compile(r"D[DG HST]".replace(" ", "")),
+            "Isomer": re.compile(r"DG|DS|DD"),
+            "FragH": re.compile(r"DP"),
+            "FragM": re.compile(r"TS"),
+            "TrpOx": re.compile(r"W"),
+            "MetOx": re.compile(r"M"),
+            "Hydro": re.compile(r"NP"),
+            "IntBind": re.compile(r"GPR|RGD|RYD|LDV|DGE|KGD|NGR"),
+            "Polyreactive": re.compile(r"GGG|GG|RR|VG|VVV|WWW|YY|WxW"),
+            "AggPatch": re.compile(r"FHW"),
+            "ViscPatch": re.compile(r"HYF|HWH"),
+        }
+    elif modality == "peptide":
+        motifs = {
+            "AspBridge": re.compile(r"N[GSQA]"),  # Deamidation hotspots
+            "AspCleave": re.compile(r"D[PGS]"),  # Acidic cleavage sites
+            "NTCycl": re.compile(r"^[QN]"),  # N-terminal cyclization
+            "ProtTryp": re.compile(r"[KR](?=.)"),  # Trypsin cleavage sites
+            "DPP4": re.compile(r"^[PX]?[AP]"),  # DPP4 cleavage sites
+            "MetOx": re.compile(r"M"),  # Methionine oxidation
+            "TrpOx": re.compile(r"W"),  # Tryptophan oxidation
+            "HydroPatch": re.compile(r"[FILVWY]{3,}"),  # Hydrophobic patches
+        }
+        if peptide_type == "cyclic":
+            # remove N-term liabilities for cyclic peptides
+            motifs.pop("NTCycl", None)
+            motifs.pop("DPP4", None)
+        elif peptide_type == "linear":
+            # For linear peptides, we handle cysteine pairing separately
+            # so we don't include CysOx in the motif set
+            pass
+        return motifs
+    else:
+        raise ValueError(f"Unknown modality: {modality}")
+
+
+def severity_score(name):
+    return const.liability_severity.get(name, const.default_severity)
+
+
+def compute_liability_scores(sequences, modality="antibody", peptide_type="linear"):
+    """
+    Compute liability scores for given sequences.
+    modality: 'antibody' or 'peptide'; peptide_type: 'linear' or 'cyclic'.
+    For cyclic peptides, terminal CysOx flags are skipped.
+
+    Returns:
+        dict: sequence -> {'score': int, 'violations': list of dicts}
+    """
+    motifs = get_motif_set(modality, peptide_type)
+    results = {}
+    for seq in sequences:
+        violations = []
+        total_score = 0
+        length = len(seq)
+        # motif scanning
+        for name, pat in motifs.items():
+            for m in pat.finditer(seq):
+                pos = m.start() + 1
+                # skip terminal cysteines for cyclic peptides
+                if (
+                    modality == "peptide"
+                    and peptide_type == "cyclic"
+                    and name == "CysOx"
+                    and pos in (1, length)
+                ):
+                    continue
+                sev = severity_score(name)
+                violations.append(
+                    {"motif": name, "pos": pos, "len": len(m.group()), "severity": sev}
+                )
+                total_score += sev
+        # antibody-specific extras
+        if modality == "antibody":
+            # unpaired cysteines
+            cpos = [i for i, aa in enumerate(seq) if aa == "C"]
+            paired = set()
+            for i in range(len(cpos) - 1):
+                if abs(cpos[i + 1] - cpos[i]) in (1, 2):
+                    paired.update({cpos[i], cpos[i + 1]})
+            for i in cpos:
+                if i not in paired:
+                    sev = severity_score("UnpairedCys")
+                    violations.append(
+                        {
+                            "motif": "UnpairedCys",
+                            "pos": i + 1,
+                            "len": 1,
+                            "severity": sev,
+                        }
+                    )
+                    total_score += sev
+            # net charge
+            charge = seq.count("K") + seq.count("R") - seq.count("D") - seq.count("E")
+            if charge > 1:
+                sev = const.default_severity
+                violations.append(
+                    {"motif": "HighNetCharge", "pos": None, "len": 0, "severity": sev}
+                )
+                total_score += sev
+        # peptide-specific extras
+        elif modality == "peptide":
+            # For linear peptides, only flag unpaired cysteines (odd number of cysteines)
+            # For cyclic peptides, terminal cysteines are expected, so only flag internal unpaired cysteines
+            cpos = [i for i, aa in enumerate(seq) if aa == "C"]
+            if peptide_type == "linear":
+                # For linear peptides, if there's an odd number of cysteines, flag all cysteines as potential liabilities
+                # since we don't know which one is unpaired
+                if len(cpos) % 2 == 1:
+                    sev = severity_score("UnpairedCys")
+                    for cys_pos in cpos:
+                        violations.append(
+                            {
+                                "motif": "UnpairedCys",
+                                "pos": cys_pos + 1,
+                                "len": 1,
+                                "severity": sev,
+                            }
+                        )
+                        total_score += sev
+            elif peptide_type == "cyclic":
+                # For cyclic peptides, terminal cysteines are expected for cyclization
+                # Only flag internal unpaired cysteines
+                internal_cpos = [
+                    i for i in cpos if i != 0 and i != len(seq) - 1
+                ]  # exclude terminal positions
+                if len(internal_cpos) % 2 == 1:
+                    # Flag all internal cysteines as potential liabilities
+                    sev = severity_score("UnpairedCys")
+                    for cys_pos in internal_cpos:
+                        violations.append(
+                            {
+                                "motif": "UnpairedCys",
+                                "pos": cys_pos + 1,
+                                "len": 1,
+                                "severity": sev,
+                            }
+                        )
+                        total_score += sev
+
+                # Additional liability checks for cyclic peptides
+                # 1. Check for low hydrophilic content (< 40%)
+                hydrophilic_residues = (
+                    seq.count("D")
+                    + seq.count("E")
+                    + seq.count("K")
+                    + seq.count("R")
+                    + seq.count("H")
+                    + seq.count("N")
+                    + seq.count("Q")
+                    + seq.count("S")
+                    + seq.count("T")
+                )
+                hydrophilic_percentage = (hydrophilic_residues / len(seq)) * 100
+                if hydrophilic_percentage < 40:
+                    sev = severity_score("LowHydrophilic")
+                    violations.append(
+                        {
+                            "motif": "LowHydrophilic",
+                            "pos": None,
+                            "len": 0,
+                            "severity": sev,
+                            "details": f"{hydrophilic_percentage:.1f}% hydrophilic",
+                        }
+                    )
+                    total_score += sev
+
+                # 2. Check for consecutive identical residues
+                max_consec_identical = 1
+                current_consec = 1
+                for i in range(1, len(seq)):
+                    if seq[i] == seq[i - 1]:
+                        current_consec += 1
+                        max_consec_identical = max(max_consec_identical, current_consec)
+                    else:
+                        current_consec = 1
+
+                if max_consec_identical > 1:
+                    sev = severity_score("ConsecIdentical")
+                    violations.append(
+                        {
+                            "motif": "ConsecIdentical",
+                            "pos": None,
+                            "len": 0,
+                            "severity": sev,
+                            "details": f"{max_consec_identical} consecutive identical",
+                        }
+                    )
+                    total_score += sev
+
+                # 3. Check for more than 4 consecutive hydrophobic residues
+                max_consec_hydrophobic = 0
+                current_consec = 0
+                for aa in seq:
+                    if aa in "FILVWY":
+                        current_consec += 1
+                        max_consec_hydrophobic = max(
+                            max_consec_hydrophobic, current_consec
+                        )
+                    else:
+                        current_consec = 0
+
+                if max_consec_hydrophobic > 4:
+                    sev = severity_score("LongHydrophobic")
+                    violations.append(
+                        {
+                            "motif": "LongHydrophobic",
+                            "pos": None,
+                            "len": 0,
+                            "severity": sev,
+                            "details": f"{max_consec_hydrophobic} consecutive hydrophobic",
+                        }
+                    )
+                    total_score += sev
+        results[seq] = {"score": total_score, "violations": violations}
+    return results
+
+
+def compute_liability_metrics(sequence, liability_modality, liability_peptide_type):
+    metrics = {}
+    # check if sequence is valid
+    if not sequence or len(sequence) == 0:
+        raise ValueError(f"Sequence is empty: '{sequence}'")
+
+    liability_results = compute_liability_scores(
+        [sequence],
+        modality=liability_modality,
+        peptide_type=liability_peptide_type,
+    )
+    liability_data = liability_results[sequence]
+
+    # Store liability metrics
+    metrics["liability_score"] = liability_data["score"]
+    metrics["liability_num_violations"] = len(liability_data["violations"])
+
+    # Count violations by severity
+    high_severity_violations = [
+        v for v in liability_data["violations"] if v["severity"] >= 10
+    ]
+    medium_severity_violations = [
+        v for v in liability_data["violations"] if 5 <= v["severity"] < 10
+    ]
+    low_severity_violations = [
+        v for v in liability_data["violations"] if v["severity"] < 5
+    ]
+
+    metrics["liability_high_severity_violations"] = len(high_severity_violations)
+    metrics["liability_medium_severity_violations"] = len(medium_severity_violations)
+    metrics["liability_low_severity_violations"] = len(low_severity_violations)
+
+    # Count violations by type
+    violation_counts = {}
+    for v in liability_data["violations"]:
+        motif = v["motif"]
+        violation_counts[motif] = violation_counts.get(motif, 0) + 1
+
+    # Store individual violation type counts as metrics
+    for motif, count in violation_counts.items():
+        metrics[f"liability_{motif}_count"] = count
+
+    # Add detailed violation information
+    # Group violations by type for intelligent reporting
+    violations_by_type = {}
+    for v in liability_data["violations"]:
+        motif = v["motif"]
+        if motif not in violations_by_type:
+            violations_by_type[motif] = []
+        violations_by_type[motif].append(v)
+
+    # Initialize default values for all motifs to ensure consistent dataframe columns
+    # Use the full motif set for the configured modality/peptide_type so columns are consistent
+    all_motifs = set(
+        get_motif_set(
+            modality=liability_modality,
+            peptide_type=liability_peptide_type,
+        ).keys()
+    )
+    for motif in all_motifs:
+        # Initialize all possible fields with default values
+        metrics[f"liability_{motif}_count"] = 0
+        metrics[
+            f"liability_{motif}_position"
+        ] = -1  # use -1 for no position (keeps int dtype)
+        metrics[f"liability_{motif}_length"] = 0
+        metrics[f"liability_{motif}_severity"] = 0
+        metrics[f"liability_{motif}_details"] = ""
+        metrics[f"liability_{motif}_positions"] = ""
+        metrics[f"liability_{motif}_num_positions"] = 0
+        metrics[f"liability_{motif}_global_details"] = ""
+        metrics[f"liability_{motif}_avg_severity"] = 0.0
+
+    # Store detailed violation information
+    for motif, motif_violations in violations_by_type.items():
+        if len(motif_violations) == 1:
+            # Single violation - store all details
+            v = motif_violations[0]
+            # Ensure position is an integer; use -1 for non-positional violations
+            metrics[f"liability_{motif}_position"] = (
+                int(v["pos"]) if v["pos"] is not None else -1
+            )
+            metrics[f"liability_{motif}_length"] = v["len"]
+            metrics[f"liability_{motif}_severity"] = v["severity"]
+            if "details" in v:
+                metrics[f"liability_{motif}_details"] = v["details"]
+            else:
+                metrics[f"liability_{motif}_details"] = ""
+        else:
+            # Multiple violations - store summary information
+            positions = [v["pos"] for v in motif_violations if v["pos"] is not None]
+            global_violations = [v for v in motif_violations if v["pos"] is None]
+
+            if positions:
+                # Store position range for positional violations
+                metrics[f"liability_{motif}_positions"] = (
+                    f"{min(positions)}-{max(positions)}"
+                )
+                metrics[f"liability_{motif}_num_positions"] = len(positions)
+            else:
+                metrics[f"liability_{motif}_positions"] = ""
+                metrics[f"liability_{motif}_num_positions"] = 0
+
+            if global_violations:
+                # Store details for global violations
+                details = [
+                    v.get("details", "") for v in global_violations if v.get("details")
+                ]
+                if details:
+                    metrics[f"liability_{motif}_global_details"] = "; ".join(details)
+                else:
+                    metrics[f"liability_{motif}_global_details"] = ""
+            else:
+                metrics[f"liability_{motif}_global_details"] = ""
+
+            # Store average severity
+            avg_severity = sum(v["severity"] for v in motif_violations) / len(
+                motif_violations
+            )
+            metrics[f"liability_{motif}_avg_severity"] = round(avg_severity, 1)
+
+    # Add a comprehensive violation summary for easy interpretation
+    violation_summary = []
+    for motif, motif_violations in violations_by_type.items():
+        count = len(motif_violations)
+        if count == 1:
+            v = motif_violations[0]
+            if v["pos"] is not None:
+                violation_summary.append(f"{motif}(pos{v['pos']},sev{v['severity']})")
+            else:
+                details = v.get("details", "")
+                violation_summary.append(
+                    f"{motif}({details},sev{v['severity']})"
+                    if details
+                    else f"{motif}(sev{v['severity']})"
+                )
+        else:
+            positions = [v["pos"] for v in motif_violations if v["pos"] is not None]
+            if positions:
+                violation_summary.append(
+                    f"{motif}x{count}(pos{min(positions)}-{max(positions)},sev{motif_violations[0]['severity']})"
+                )
+            else:
+                violation_summary.append(
+                    f"{motif}x{count}(sev{motif_violations[0]['severity']})"
+                )
+
+    metrics["liability_violations_summary"] = (
+        "; ".join(violation_summary) if violation_summary else ""
+    )
+    return metrics

model/boltzgen/task/filter/filter.py

@@ -0,0 +1,1413 @@
+import json
+from boltzgen.utils.quiet import quiet_startup
+
+
+quiet_startup()
+from typing import Dict
+import pandas as pd
+from pathlib import Path
+import matplotlib.pyplot as plt
+import shutil
+from tqdm import tqdm
+from matplotlib.backends.backend_pdf import PdfPages
+import pandas as pd
+from pathlib import Path
+from Bio import Align
+import numpy as np
+import random
+from tqdm.auto import tqdm
+import heapq
+from textwrap import fill, dedent
+import re
+
+from boltzgen.task.analyze.analyze_utils import compute_liability_scores
+from boltzgen.task.filter.seqplot_utils import (
+    aa_composition_pie,
+    cdr_logo,
+    create_alignment_logo,
+    plot_seq_liabilities,
+)
+from boltzgen.task.task import Task
+
+
+class Filter(Task):
+    """
+    This step consumes the aggregated metrics produced by **Analyze** (i.e.
+    ``aggregate_metrics_*.csv`` in ``design_dir``), applies hard
+    filtering thresholds, ranks candidates by a composite quality key, and then
+    performs a lazy-greedy **diversity selection** over sequence identity.
+    It writes:
+        • the top-`budget` quality+diversity designs (+ refolded CIFs),
+        • a CSV with all metrics and ranks,
+        • a summary PDF with tables, histograms, scatter plots, and (optionally)
+            sequence logos and liability heat-maps.
+
+    The step is designed to be **fast** (≈30 s per 100k designs on typical CPUs);
+    the most expensive part is PDF plotting.
+
+    Parameters
+    ----------
+    design_dir : str | Path
+        Directory containing design artifacts and metrics, including:
+        - ``aggregate_metrics_*.csv`` (required)
+        - ``refold_cif/`` (for copying refolded structures)
+        - ``ca_coords_sequences.pkl.gz`` (id→sequence mapping for diversity)
+    outdir : str | Path, default=None
+        Defaults to design_dir. Parent directory where results will be written. The step creates
+        ``{outdir}/final_ranked_designs`` and subfolders.
+    budget : int, default=30
+        Number of designs to select by quality+diversity (Diverse set).
+    use_affinity : bool, default=False
+        Switch to affinity-oriented metrics (for small-molecule binders).
+    from_inverse_folded : bool, default=True
+        If True, use backbone RMSDs (``bb_rmsd``) for filtering and delta-SASA from
+        refolded structures. If False, use all-atom RMSDs and original ΔSASA.
+    filter_designfolding : bool, default=True
+        Also require the isolated design (no target) to refold to the same shape,
+        via ``designfolding-*_rmsd``.
+    filter_bindingsite : bool, default=False
+        Keep only designs with at least one residue within the specified distance of
+        a binding-site residue (e.g., ``bindsite_under_8rmsd > 0``).
+    filter_target_aligned : bool, default=False
+        Require target backbone alignment (``bb_target_aligned<2.5`` flag).
+    filter_biased : bool, default=True
+        Remove amino-acid composition outliers (default caps on ALA/GLY/GLU/LEU/VAL).
+    refolding_rmsd_threshold : float, default=2.5
+        Threshold used for the RMSD-based filters (lower is better).
+    modality : {"peptide","antibody"}, default="peptide"
+        Affects liability scoring and optional sequence visualizations.
+    alpha : float in [0,1], default=0.1
+        Trade-off for sequence diversity selection: 0=quality-only, 1=diversity-only.
+    metrics_override : Dict[str, float | None], default=None
+        Per-metric *inverse-importance* weights for ranking.
+        - A larger value **down-weights** that metric’s rank (rank / weight).
+        - ``None`` removes the metric entirely from the ranking key.
+        Example: ``{"plip_hbonds_refolded": 4, "delta_sasa_refolded": 2, "neg_min_design_to_target_pae": 1}``.
+    num_liability_plots : int, default=0
+        If >0, produce per-residue developability heat-maps for the first N top designs.
+    plot_seq_logos : bool, default=False
+        If True, include alignment logos and AA composition pies for All/Top/Diverse sets.
+    additional_filters : list[Dict], default=[]
+        Extra hard filters of the form:
+        ``{"feature": "<column>", "lower_is_better": bool, "threshold": float}``.
+    size_buckets : list[Dict], default=[]
+        Optional constraint for the maximum number of designs returned in a certain size range:
+        ``{"num_designs": int, "min": int, "max": int}``.
+
+    Ranking & Diversity
+    -------------------
+    **Filtering.** Each design must pass all hard thresholds (``pass_<feature>_filter``).
+    The step also adds convenience columns (e.g., ``filter_rmsd``,
+    ``designfolding-filter_rmsd``, signed variants like ``neg_min_design_to_target_pae``).
+
+    **Ranking (quality).** For each metric in ``self.metrics``, compute the row's
+    rank on the tuple ``(num_filters_passed, metric)`` (so designs that fail filters
+    are pushed down), then divide by the metric’s *inverse-importance* weight.
+    The **worst** (max) scaled rank across metrics becomes the design’s
+    quality key. The Top set is the best `budget` designs by this key (tie-broken
+    by iPTM).
+
+    **Diversity.** A lazy-greedy selection chooses `div_budget` designs maximizing:
+    ``(1 - alpha) * quality + alpha * (1 - seq_identity)``.
+    Sequence identity is computed via pairwise alignment on the full chain or the
+    designed segment (auto-chosen based on typical length ratio). Optional
+    ``size_buckets`` limit the number of selections per length range.
+
+    Inputs
+    ------
+    • ``design_dir/aggregate_metrics_*.csv`` (required)
+    • ``design_dir/ca_coords_sequences.pkl.gz`` (required for diversity)
+    • ``design_dir/refold_cif/`` (optional but recommended, copied to outputs)
+
+    Outputs
+    -------
+    On disk (under ``{outdir}/final_ranked_designs``):
+        • ``final_{budget}_designs`` folder with .mmcif files
+        • ``metrics_*.csv`` – full table with ranks/flags
+        • ``diverse_selected_{div_budget}.csv`` – rows/IDs for Diverse set
+        • ``results_overview_*.pdf`` – summary report with tables/plots
+
+    Notes
+    -----
+    - ``metrics_override`` lets you:
+        * increase a weight (de-emphasize that metric),
+        * set a metric to ``None`` (remove from ranking),
+        * include a new metric that is already in the csv for ranking.
+
+    Examples
+    --------
+    >>> from boltzgen.task.filter.filter import Filter
+    >>> filter = Filter(
+    ...     design_dir="workbench/run123",
+    ...     budget=50, div_budget=30, alpha=0.25,
+    ...     metrics_override={"design_ptm": 2, "neg_min_design_to_target_pae": 1},
+    ...     additional_filters=[{"feature":"design_ptm","lower_is_better":False,"threshold":0.7}],
+    ... )
+    >>> filter.run(jupyter_nb=False)  # writes CSVs, PDF, and copies selected structures
+    """
+
+    def __init__(
+        self,
+        design_dir: str,
+        budget: int = 30,
+        top_budget: int = 10,
+        outdir: str = None,
+        use_affinity: bool = False,  # This changes the filtering metrics to metrics more amenable to small molecule binder design
+        filter_cysteine: bool = True,  # This filters out all designs that have designed cysteins in them (prespecified cysteins in the design are not counted)
+        from_inverse_folded: bool = True,  # This makes it so that we use the backbone refolding rmsd instead of the all-atom RMSD
+        filter_designfolding: bool = True,  # Additionally filter based on the RMSD from refolding the design in isolation. This makes sure the design has the same shape with and without the target being present.
+        filter_bindingsite: bool = False,  # This filters out everything that does not have a residue within 4A of a binding site residue
+        filter_target_aligned: bool = False,
+        filter_biased: bool = True,  # This filters out sequences that are alanine rich, 30% alanine is threshold
+        refolding_rmsd_threshold: float = 2.5,
+        modality: str = "peptide",  # peptide, antibody
+        peptide_type: str = "linear",  # linear, cyclic
+        alpha: float = 0.1,  # 0 = quality-only, 1 = diversity-only
+        random_state: int = 0,
+        metrics_override: Dict = None,  # overrides metrics, None values delete keys
+        num_liability_plots: int = 0,
+        plot_seq_logos: bool = False,  # make sequence logo diagrams of designed sequence
+        additional_filters: list[
+            Dict
+        ] = [],  # For example: [{"feature": "design_ALA", "lower_is_better": True, "threshold": 0.3}],
+        size_buckets: list[Dict] = [],
+    ):
+        super().__init__()
+        assert modality in ["peptide", "antibody"]
+        assert peptide_type in ["linear", "cyclic"]
+        self.design_dir = Path(design_dir)
+        self.top_budget = top_budget
+        self.use_affinity = use_affinity
+        self.filter_cysteine = filter_cysteine
+        self.from_inverse_folded = from_inverse_folded
+        self.filter_bindingsite = filter_bindingsite
+        self.budget = budget
+        self.alpha = alpha
+        self.random_state = random_state
+        self.num_liability_plots = num_liability_plots
+        self.plot_seq_logos = plot_seq_logos
+        self.modality = modality
+        self.peptide_type = peptide_type
+        self.size_buckets = size_buckets
+
+        if outdir is None:
+            outdir = design_dir
+        self.outdir = Path(f"{outdir}") / "final_ranked_designs"
+        self.top_dir = self.outdir / f"intermediate_ranked_{top_budget}_designs"
+        self.div_dir = self.outdir / f"final_{budget}_designs"
+        self.outdir.mkdir(parents=True, exist_ok=True)
+        self.top_dir.mkdir(parents=True, exist_ok=True)
+        self.div_dir.mkdir(parents=True, exist_ok=True)
+
+        # we want to maximize all these metrics
+        self.metrics: dict = {
+            "design_to_target_iptm": 1,
+            "design_ptm": 1,
+            "neg_min_design_to_target_pae": 1,
+            "plip_hbonds" + ("_refolded" if from_inverse_folded else ""): 2,
+            "plip_saltbridge" + ("_refolded" if from_inverse_folded else ""): 2,
+            "delta_sasa_refolded" if from_inverse_folded else "delta_sasa_original": 2,
+        }
+        if use_affinity:
+            self.metrics: dict = {
+                "design_to_target_iptm": 1.1,
+                "design_ptm": 1.1,
+                "neg_min_design_to_target_pae": 1.1,
+                "affinity_probability_binary1": 1,
+                "plip_hbonds" + ("_refolded" if from_inverse_folded else ""): 2,
+                "plip_saltbridge" + ("_refolded" if from_inverse_folded else ""): 2,
+                "delta_sasa_refolded"
+                if from_inverse_folded
+                else "delta_sasa_original": 2,
+            }
+
+        # override metrics
+        if not metrics_override is None:
+            for k in metrics_override:
+                if metrics_override[k] is None:
+                    self.metrics.pop(k, None)
+                else:
+                    self.metrics[k] = metrics_override[k]
+
+        # Define how to Filter
+        self.filters = [
+            {"feature": "has_x", "lower_is_better": True, "threshold": 0},
+            {
+                "feature": "filter_rmsd",
+                "lower_is_better": True,
+                "threshold": refolding_rmsd_threshold,
+            },
+            {
+                "feature": "filter_rmsd_design",
+                "lower_is_better": True,
+                "threshold": refolding_rmsd_threshold,
+            },
+        ]
+        if filter_designfolding:
+            self.filters.append(
+                {
+                    "feature": "designfolding-filter_rmsd",
+                    "lower_is_better": True,
+                    "threshold": refolding_rmsd_threshold,
+                }
+            )
+        if filter_bindingsite:
+            self.filters.append(
+                {
+                    "feature": "bindsite_under_8rmsd",  # center_coord RMSD
+                    "lower_is_better": False,
+                    "threshold": 0.0001,  # at least one binding site residue
+                },
+            )
+        if filter_target_aligned:
+            self.filters.append(
+                {
+                    "feature": "bb_target_aligned<2.5",
+                    "lower_is_better": False,
+                }
+            )
+        if filter_cysteine:
+            self.filters.append(
+                {
+                    "feature": "CYS_fraction",
+                    "lower_is_better": True,
+                    "threshold": 0,
+                },
+            )
+        if filter_biased:
+            self.filters.extend(
+                [
+                    {
+                        "feature": "ALA_fraction",
+                        "lower_is_better": True,
+                        "threshold": 0.3,
+                    },
+                    {
+                        "feature": "GLY_fraction",
+                        "lower_is_better": True,
+                        "threshold": 0.3,
+                    },
+                    {
+                        "feature": "GLU_fraction",
+                        "lower_is_better": True,
+                        "threshold": 0.3,
+                    },
+                    {
+                        "feature": "LEU_fraction",
+                        "lower_is_better": True,
+                        "threshold": 0.3,
+                    },
+                    {
+                        "feature": "VAL_fraction",
+                        "lower_is_better": True,
+                        "threshold": 0.3,
+                    },
+                ]
+            )
+        self.filters.extend(additional_filters)
+
+        random.seed(self.random_state)
+        np.random.seed(self.random_state)
+
+    def run(self, config=None, jupyter_nb=False):
+        self.load_dataframe()
+        self.reset_outdir()
+        self.filter_df()
+        self.absolute_metrics()
+        self.sort_df()
+        self.optimize_diversity()
+        self.write_outdir()
+
+        # Visualizations
+        print(
+            "\nWriting design files is done. Now making plots for a final summary .pdf file with statistics."
+        )
+        (
+            hist_metrics,
+            extra_pairs,
+            row_headers,
+            rows,
+            metric_rows,
+            intro_text,
+            csv_expl_rows,
+        ) = self.prepare_visualization()
+        self.make_visualization(
+            hist_metrics,
+            extra_pairs,
+            row_headers,
+            rows,
+            metric_rows,
+            intro_text,
+            csv_expl_rows,
+            jupyter_nb=jupyter_nb,
+        )
+
+    def reset_outdir(self):
+        if self.outdir.exists():
+            shutil.rmtree(self.outdir)
+        self.outdir.mkdir(parents=True, exist_ok=True)
+        self.top_dir.mkdir(parents=True, exist_ok=True)
+        self.div_dir.mkdir(parents=True, exist_ok=True)
+
+    def load_dataframe(self):
+        csv_files = sorted(self.design_dir.glob("aggregate_metrics_*.csv"))
+        if not csv_files:
+            msg = (
+                f"No file starting with 'aggregate_metrics_' found in {self.design_dir}"
+            )
+            raise FileNotFoundError(msg)
+        df_in = pd.read_csv(csv_files[0])
+
+        self.df_in = df_in.copy()
+        df = df_in.copy()
+
+        if self.from_inverse_folded:
+            df["filter_rmsd"] = df["bb_rmsd"]
+            df["filter_rmsd_design"] = df["bb_rmsd_design"]
+        else:
+            df["filter_rmsd"] = df["rmsd"]
+            df["filter_rmsd_design"] = df["rmsd_design"]
+        if "designfolding-rmsd" in df:
+            df["designfolding-filter_rmsd"] = df["designfolding-rmsd"]
+        if "designfolding-bb_rmsd" in df and self.from_inverse_folded:
+            df["designfolding-filter_rmsd"] = df["designfolding-bb_rmsd"]
+        if "min_design_to_target_pae" in df:
+            df["neg_min_design_to_target_pae"] = -df["min_design_to_target_pae"]
+
+        if "design_hydrophobicity" in df:
+            df["neg_design_hydrophobicity"] = -df["design_hydrophobicity"]
+        if "design_largest_hydrophobic_patch_refolded" in df:
+            df["neg_design_largest_hydrophobic_patch_refolded"] = -df[
+                "design_largest_hydrophobic_patch_refolded"
+            ]
+        df["neg_min_interaction_pae"] = -df["min_interaction_pae"]
+        df["neg_filter_rmsd"] = -df["filter_rmsd"]
+        df["neg_filter_rmsd_design"] = -df["filter_rmsd_design"]
+        df["has_x"] = df["designed_sequence"].str.contains("X")
+        self.df = df
+
+        print(f"Total number of designs: {len(self.df):>5}")
+        before = len(self.df)
+        self.df = self.df.drop_duplicates(subset="designed_sequence", keep="first")
+        msg = f"Duplicates found: {before - len(self.df)}. Removing duplicates. {len(self.df)} designs remain.\n"
+        print(msg)
+
+    def filter_df(self):
+        filter_cols = []
+        self.df["num_filters_passed"] = 0
+
+        for filter in self.filters:
+            feat = filter["feature"]
+            low = filter["lower_is_better"]
+            threshold = filter["threshold"]
+
+            filter_col = f"pass_{feat}_filter"
+            filter_cols.append(filter_col)
+            if low:
+                self.df[filter_col] = self.df[feat] <= threshold
+            else:
+                self.df[filter_col] = self.df[feat] >= threshold
+
+            self.df["num_filters_passed"] += self.df[filter_cols].all(axis=1)
+            self.df["pass_filters"] = self.df[filter_cols].all(axis=1)
+
+            msg = f"Num designs that pass the {feat} filter with threshold {threshold} where {'lower' if low else 'higher'} is better: {self.df[filter_col].sum()}"
+            print(msg)
+            print(f"Remaining designs: {self.df['pass_filters'].sum()}")
+
+        num_pass = self.df["pass_filters"].sum()
+        if num_pass < self.top_budget or num_pass < self.budget:
+            print(
+                f"Only {num_pass} designs pass filters. We highly recommend relaxing the thresholds."
+            )
+        print("\n")
+
+    def absolute_metrics(self):
+        norm_path = Path(__file__).resolve().parents[2] / "resources" / "metrics_normalization.json"
+        if not norm_path.exists():
+            return
+
+        with norm_path.open("r") as f:
+            norm_stats = json.load(f)
+
+        for col, stats in norm_stats.items():
+            mean = stats["mean"]
+            std = stats["std"]
+            if col in self.df.columns:
+                self.df[col + "_z"] = (self.df[col] - mean) / std
+
+        importances = {
+            "affinity_probability_binary1": 1.5,
+            "design_iiptm": 1.0,
+            "design_ptm": 0.5,
+            "min_design_to_target_pae": -1.0,  # lower is better
+            "design_hydrophobicity": -0.125,  # lower is better
+            "design_largest_hydrophobic_patch_refolded": -0.15,  # lower is better
+            "delta_sasa_refolded": 0.25,
+            "plip_saltbridge_refolded": 0.25,
+            "plip_hbonds_refolded": 0.25,
+        }
+
+        self.df["absolute_score"] = 0.0
+        for base_col, weight in importances.items():
+            if base_col in self.df.columns:
+                norm_col = base_col + "_z"
+                self.df["absolute_score"] += weight * self.df[norm_col]
+        total_importance = sum(abs(w) for w in importances.values())
+        self.df["absolute_score"] /= total_importance
+
+        self.df["structure_confidence"] = 0.0
+        weight_sum = 0
+        for col in ["design_iiptm", "design_ptm", "min_design_to_target_pae"]:
+            weight = importances[col]
+            norm_col = col + "_z"
+            weight_sum += abs(weight)
+            self.df["structure_confidence"] += weight * self.df[norm_col]
+        self.df["structure_confidence"] /= weight_sum
+
+        for flt in self.filters:
+            feat = flt["feature"]
+            filter_col = f"pass_{feat}_filter"
+            if "fraction" in feat:
+                # If this is a "fraction" feature, meaning a res_type fraction filter, only apply the penalty if num_design > 8
+                mask_fail = (self.df["num_design"] > 8) & (self.df[filter_col] == False)
+            else:
+                mask_fail = self.df[filter_col] == False
+
+            self.df.loc[mask_fail, "absolute_score"] *= 0.1
+
+
+    def sort_df(self):
+        rank_df = pd.DataFrame(index=self.df.index)
+
+        # 1. For each row, compute its rank according to each metric
+        # Scale the ranks by importance (divide by inverse_importance)
+        # Use the feature that we are ranking by AND the number of filters that a design passes. Thus the things that do not pass ass many filters (e.g. they dont pass all the filters) are at the end
+        for col, inverse_importance in self.metrics.items():
+            if self.metrics[col] is None:
+                continue
+
+            rank_df[f"rank_{col}"] = (
+                self.df[["num_filters_passed", col]]
+                .apply(tuple, axis=1)
+                .rank(method="min", ascending=False)
+                .astype(int)
+                / inverse_importance
+            )
+        self.df = pd.concat([self.df, rank_df], axis=1)
+
+        # 2. For each row, find the max (worst) rank across the metrics.
+        # This single value determines its final rank group.
+        self.df["max_rank"] = rank_df.max(axis=1)
+
+        # 3. Sort by this new max_rank and then create the final dense rank,
+        # which is equivalent to the original 'rank_counter'.
+        self.df = self.df.sort_values("max_rank")
+        self.df["secondary_rank"] = self.df["max_rank"].rank(method="dense").astype(int)
+
+        # sort by ranking and resolve ties via design_iptm
+        self.df = self.df.sort_values(
+            by=[
+                "secondary_rank",
+                "design_to_target_iptm"
+                if "design_to_target_iptm" in self.df
+                else "design_iptm",
+            ],
+            ascending=[True, False],
+        )
+
+        self.df["final_rank"] = np.arange(1, len(self.df) + 1)
+        self.df["quality_score"] = 1 - (self.df["final_rank"] - 1) / (len(self.df) - 1)
+
+        # Reorder columns
+        priority_col_candidates = [
+            "id",
+            "final_rank",
+            "designed_sequence",
+            "designed_chain_sequence",
+            "num_design",
+            "affinity_probability_binary1",
+            "design_to_target_iptm"
+            if "design_to_target_iptm" in self.df
+            else "design_iptm",
+            "min_design_to_target_pae"
+            if "min_design_to_target_pae" in self.df
+            else "min_interaction_pae",
+            "design_ptm",
+            "filter_rmsd",
+            "designfolding-filter_rmsd"
+            if "designfolding-filter_rmsd" in self.df
+            else "filter_rmsd"
+            "plip_saltbridge" + ("_refolded" if self.from_inverse_folded else ""),
+            "plip_hbonds" + ("_refolded" if self.from_inverse_folded else ""),
+            "delta_sasa_refolded",
+            "design_largest_hydrophobic_patch_refolded",
+            "design_chain_hydrophobicity",
+            "design_hydrophobicity",
+            "loop",
+            "helix",
+            "sheet",
+        ]
+        priority_cols = [c for c in priority_col_candidates if c in self.df.columns]
+
+        other_cols = [col for col in self.df.columns if col not in priority_cols]
+        new_column_order = priority_cols + other_cols
+        self.df = self.df[new_column_order]
+
+    def write_outdir(self):
+        num_digits = len(str(len(self.df)))
+
+        top_dir2 = self.top_dir / "before_refolding"
+        top_dir2.mkdir(parents=True, exist_ok=True)
+        for i, (idx, row) in tqdm(
+            enumerate(self.df[: self.top_budget].iterrows()),
+            desc="copy top design files",
+        ):
+            filename = row["file_name"]
+            new_filename = f"rank{i:0{num_digits}d}_{filename}"
+            src = self.design_dir / filename
+            dst = top_dir2 / new_filename
+            shutil.copy2(src, dst)
+
+            src = self.design_dir / "refold_cif" / filename
+            dst = self.top_dir / new_filename
+            shutil.copy2(src, dst)
+
+        # save to output/diverse_* directory
+        self.div_dir.mkdir(parents=True, exist_ok=True)
+        div_dir2 = self.div_dir / "before_refolding"
+        div_dir2.mkdir(parents=True, exist_ok=True)
+        for i in tqdm(self.diverse_selection, desc="copy diversity files"):
+            src = self.design_dir / self.df_m.loc[i, "file_name"]
+            qualityrank = self.df_m.loc[i, "final_rank"]
+            filename = src.name
+            new_filename = f"rank{qualityrank:0{num_digits}d}_{filename}"
+            shutil.copy2(src, div_dir2 / new_filename)
+
+            src = self.design_dir / "refold_cif" / self.df_m.loc[i, "file_name"]
+            shutil.copy2(src, self.div_dir / new_filename)
+        self.df_div.to_csv(
+            self.outdir / f"final_designs_metrics_{self.budget}.csv", index=False
+        )
+        print("Files + CSV saved to", self.outdir)
+
+        self.df.to_csv(self.outdir / f"all_designs_metrics.csv", index=False)
+
+    def optimize_diversity(self):
+        # Load structures and sequences to compute similarities
+        seq_path = self.design_dir / "ca_coords_sequences.pkl.gz"
+        if not seq_path.exists():
+            raise FileNotFoundError(f"Expected {seq_path} to exist")
+        df_seq = pd.read_pickle(seq_path)[["id", "sequence"]]
+
+        self.df_m = pd.merge(self.df, df_seq, on="id", how="inner").reset_index(
+            drop=True
+        )
+        seqs = self.df_m["sequence"].tolist()
+        quality = self.df_m["quality_score"].to_numpy()
+
+        # sequence-only similarity
+        aligner = Align.PairwiseAligner()
+        pid_cache = {}
+
+        def sim_fn(i, j):
+            if i == j:
+                return 1.0
+            key = tuple(sorted((i, j)))
+            if key not in pid_cache:
+                seq1 = seqs[i]
+                seq2 = seqs[j]
+                aln = aligner.align(seq1, seq2)[0]
+                pid_cache[key] = aln.score / max(len(seqs[i]), len(seqs[j]))
+            return pid_cache[key]
+
+        random.seed(self.random_state)
+        np.random.seed(self.random_state)
+        diverse_selection = self.select_lazy_greedy(self.budget, quality, sim_fn)
+
+        self.diverse_selection = diverse_selection
+        self.df_div = self.df_m.iloc[diverse_selection].reset_index(drop=True)
+
+    def select_lazy_greedy(self, k, quality, sim_fn):
+        # Handle edge case where we have fewer items than requested
+        if len(quality) <= k:
+            return list(range(len(quality)))
+
+        selected = [int(np.argmax(quality))]
+        remaining = set(range(len(quality))) - set(selected)
+
+        heap = []
+        for i in remaining:
+            div = 1 - sim_fn(i, selected[0])
+            gain = (1 - self.alpha) * quality[i] + self.alpha * div
+            heapq.heappush(heap, (-gain, i))
+
+        buckets = np.zeros(len(self.size_buckets) + 1)
+        first = selected[0]
+        first_len = len(self.df_m["sequence"][first])
+        for idx, bucket_size in enumerate(self.size_buckets):
+            if first_len >= bucket_size["min"] and first_len < bucket_size["max"]:
+                buckets[idx] += 1
+                break
+        for _ in tqdm(
+            range(k - 1), desc="Performing lazy greedy diversity optimization."
+        ):
+            while True:
+                neg_gain, cand = heapq.heappop(heap)
+                num_design = len(self.df_m["sequence"][cand])
+                bucket_idx = None
+                for idx, bucket_size in enumerate(self.size_buckets):
+                    if (
+                        num_design < bucket_size["max"]
+                        and num_design >= bucket_size["min"]
+                    ):
+                        bucket_idx = idx
+
+                if bucket_idx is None:
+                    bucket_idx = len(self.size_buckets)
+
+                bucket_full = (
+                    not bucket_idx == len(self.size_buckets)
+                    and buckets[bucket_idx]
+                    == self.size_buckets[bucket_idx]["num_designs"]
+                )
+                if bucket_full:
+                    continue
+
+                true_div = 1 - max(sim_fn(cand, j) for j in selected)
+                true_g = (1 - self.alpha) * quality[cand] + self.alpha * true_div
+                heapq.heappush(heap, (-true_g, cand))
+
+                if heap[0][1] == cand:
+                    heapq.heappop(heap)
+                    selected.append(cand)
+                    remaining.remove(cand)
+                    buckets[bucket_idx] += 1
+                    break
+        return sorted(selected)
+
+    def prepare_visualization(self):
+        summary_metrics = [
+            "num_design",
+            "filter_rmsd",
+            "designfolding-filter_rmsd"
+            if "designfolding-filter_rmsd" in self.df
+            else "filter_rmsddesign_ptm",
+            "design_iptm",
+            "design_to_target_iptm"
+            if "design_to_target_iptm" in self.df
+            else "design_iptm",
+            "min_design_to_target_pae"
+            if "min_design_to_target_pae" in self.df
+            else "min_interaction_pae",
+            "design_ipsae_min"
+            if "design_ipsae_min" in self.df
+            else "design_iptm",
+            "design_to_target_ipsae"
+            if "design_to_target_ipsae" in self.df
+            else "design_iptm",
+            "delta_sasa_refolded"
+            if self.from_inverse_folded
+            else "delta_sasa_original",
+            "plip_saltbridge" + ("_refolded" if self.from_inverse_folded else ""),
+            "plip_hbonds" + ("_refolded" if self.from_inverse_folded else ""),
+        ]
+
+        # Scatter pairs (each will be one page)
+        extra_pairs = [
+            ("num_design", "rank"),
+            (
+                "num_design",
+                "plip_saltbridge" + ("_refolded" if self.from_inverse_folded else ""),
+            ),
+            (
+                "num_design",
+                "plip_hbonds" + ("_refolded" if self.from_inverse_folded else ""),
+            ),
+            (
+                "num_design",
+                "delta_sasa" + ("_refolded" if self.from_inverse_folded else ""),
+            ),
+            (
+                "num_design",
+                "design_ptm",
+            ),
+            (
+                "num_design",
+                "min_design_to_target_pae"
+                if "min_design_to_target_pae" in self.df
+                else "min_interaction_pae",
+            ),
+            (
+                "num_design",
+                "design_to_target_iptm"
+                if "design_to_target_iptm" in self.df
+                else "design_iptm",
+            ),
+            (
+                "num_design",
+                "design_iptm",
+            ),
+            (
+                "num_design",
+                "design_ipsae_min" if "design_ipsae_min" in self.df else "design_iptm",
+            ),
+            (
+                "num_design",
+                "design_to_target_ipsae" if "design_to_target_ipsae" in self.df else "design_iptm",
+            ),
+            (
+                "num_design",
+                "design_iiptm" if "design_iiptm" in self.df else "design_iptm",
+            ),
+            (
+                "num_design",
+                "design_largest_hydrophobic_patch_refolded"
+                if "design_largest_hydrophobic_patch_refolded" in self.df
+                else "design_iptm",
+            ),
+            (
+                "num_design",
+                "design_hydrophobicity"
+                if "design_hydrophobicity" in self.df
+                else "design_iptm",
+            ),
+        ]
+        if not self.from_inverse_folded:
+            extra_pairs.append(("delta_sasa_refolded", "delta_sasa_original"))
+            extra_pairs.append(("plip_saltbridge", "delta_sasa_original"))
+
+        # Histograms with selected overlay
+        hist_metrics = [
+            "num_design",
+            "filter_rmsd",
+            "designfolding-filter_rmsd"
+            if "designfolding-filter_rmsd" in self.df
+            else "filter_rmsddesign_ptm",
+            "design_to_target_iptm"
+            if "design_to_target_iptm" in self.df
+            else "design_iptm",
+            "design_iptm",
+            "design_ipsae_min"
+            if "design_ipsae_min" in self.df
+            else "design_iptm",
+            "design_to_target_ipsae"
+            if "design_to_target_ipsae" in self.df
+            else "design_iptm",
+            "min_design_to_target_pae"
+            if "min_design_to_target_pae" in self.df
+            else "min_interaction_pae",
+            "plip_saltbridge" + ("_refolded" if self.from_inverse_folded else ""),
+            "plip_hbonds" + ("_refolded" if self.from_inverse_folded else ""),
+            "plip_hydrophobic" + ("_refolded" if self.from_inverse_folded else ""),
+            "delta_sasa_refolded",
+            "design_largest_hydrophobic_patch_refolded"
+            if "design_largest_hydrophobic_patch_refolded" in self.df
+            else "delta_sasa_refolded",
+            "design_hydrophobicity"
+            if "design_hydrophobicity" in self.df
+            else "delta_sasa_refolded",
+        ]
+        if not self.from_inverse_folded:
+            hist_metrics.append("delta_sasa_original")
+
+        hist_metrics = list(dict.fromkeys(hist_metrics))
+        extra_pairs = list(dict.fromkeys(extra_pairs))
+
+        # Prepend any active ranking metrics not already in the lists
+        extra_ranking = [m for m in self.metrics if m in self.df.columns and m not in hist_metrics]
+        hist_metrics = extra_ranking + hist_metrics
+        summary_metrics = extra_ranking + summary_metrics
+        extra_pairs = [("num_design", m) for m in extra_ranking] + extra_pairs
+
+        if self.use_affinity:
+            summary_metrics.insert(2, "affinity_probability_binary1")
+            hist_metrics.insert(2, "affinity_probability_binary1")
+
+        avail = [m for m in summary_metrics if m in self.df.columns]
+        base_rows = [
+            ["Num designs", len(self.df), "-"],
+        ]
+
+        extra_mean = (
+            [
+                m,
+                f"{self.df[m].mean():.3f}",  # mean of ALL
+                f"{self.df[: self.top_budget][m].mean():.3f}",  # mean of red set
+                f"{self.df_div[m].mean():.3f}",  # mean of BLUE set
+            ]
+            for m in avail
+        )
+
+        for row in base_rows:
+            row.append("-")
+
+        rows = base_rows + list(extra_mean)
+
+        row_headers = [
+            "Metric",
+            f"Mean",
+            f"Mean top {self.top_budget}",
+            f"Mean top {self.budget} diverse",
+        ]
+
+        metric_rows = [[k, v] for k, v in self.metrics.items()]
+
+        n_total = len(self.df)
+        n_pass_filters = (
+            int(self.df["pass_filters"].sum()) if "pass_filters" in self.df else n_total
+        )
+
+        text = f"""
+        • Designs generated: {n_total}
+        • Designs passing all filters: {n_pass_filters}
+
+        • MMCIF files of final designs are in:  
+          {self.div_dir}  
+
+        • Metrics and sequences of {self.budget} final designs are in:  
+          {self.outdir}/final_designs_metrics_{self.budget}.csv
+          
+        • Metrics of all designs are in:  
+          {self.outdir}/all_designs_metrics.csv
+
+        You can rerun filtering (very quick), using this command with changed parameters:
+            -- boltzgen run input_spec.yaml --steps filtering --config filtering budget=60 alpha=0.05
+        You can also rerun filtering in a jupyter notebook if you want using `filter.ipynb`
+
+         What was run to produce this in the Filter task:
+         1. Filtering: each design is evaluated against mandatory thresholds. 
+         2. Ranking: for every metric we compute its rank, then scale it by the metric’s inverse-importance weight. Designs with fewer passed filters are automatically penalised because the ranking key is the pair (num_filters_passed, metric). The overall quality score is the worst (maximum) of these scaled ranks. The {self.top_budget} best designs form the Top set.
+         3. Diversity: a lazy-greedy algorithm selects {self.budget} designs that jointly maximise quality and
+            minimise sequence similarity (sequence-identity distance). The trade-off is controlled by α = {self.alpha}:
+              • α = 0   → 100 % quality focus (same as Top set)
+              • α = 1   → 100 % diversity focus (ignores quality)
+              • Quality – composite of metrics such as iPTM, salt-bridges, ΔSASA, etc. Each metric has an "inverse importance" weight (see "Sorting Criteria" table). A larger weight divides the rank by a bigger number and therefore down-weights that metric.
+              • Diversity – 1 − sequence identity between designs.
+            We use α = {self.alpha}, meaning {round((1 - self.alpha) * 100)} % emphasis on quality, {round(self.alpha * 100)} % on diversity.
+ 
+          """
+
+        csv_expl_rows = [
+            ["id", "filename to retrieve the file"],
+            ["design_sequence", "designed amino acids (may be subset of chain)"],
+            [
+                "designed_chain_sequence",
+                "full sequence of the chain containing designed residues (recommended for synthesis)",
+            ],
+            ["num_design", "number of designed residues"],
+            ["secondary_rank", "intermediate rank from the sorting procedure"],
+            [
+                "design_ptm",
+                "predicted TM score for intra-design contacts (higher = better)",
+            ],
+            [
+                "design_iptm",
+                "predicted TM score for design–target contacts (higher = better)",
+            ],
+            [
+                "design_to_target_iptm",
+                "same as design_iptm but for multi-chain designs",
+            ],
+            [
+                "design_ipsae_min",
+                "min interaction pSAE: PAE-based confidence score for interaction (higher = better)",
+            ],
+            [
+                "design_to_target_ipsae",
+                "interaction pSAE between design and target (higher = better)",
+            ],
+            [
+                "min_design_to_target_pae",
+                "minimum PAE between design & target (lower = better)",
+            ],
+            ["plip_saltbridge", "number of salt-bridge interactions"],
+            ["plip_hbonds", "number of hydrogen-bond interactions"],
+            ["plip_hydrophobic", "number of hydrophobic interactions"],
+            ["delta_sasa_original", "ΔSASA when binder present vs absent"],
+            ["delta_sasa_refolded", "same as above but on the refolded structure"],
+        ]
+
+        intro_text = text
+
+        return (
+            hist_metrics,
+            extra_pairs,
+            row_headers,
+            rows,
+            metric_rows,
+            intro_text,
+            csv_expl_rows,
+        )
+
+    def make_visualization(
+        self,
+        hist_metrics,
+        extra_pairs,
+        row_headers,
+        rows,
+        metric_rows,
+        intro_text,
+        csv_expl_rows,
+        jupyter_nb=False,
+    ):
+        pdf_path = self.outdir / f"results_overview.pdf"
+        pdf = PdfPages(pdf_path)
+
+        def _ensure_width(fig, target_w=8.5):
+            w, h = fig.get_size_inches()
+            if abs(w - target_w) > 0.01:
+                scale = target_w / w
+                fig.set_size_inches(target_w, h * scale, forward=True)
+
+        def show(fig):
+            _ensure_width(fig)
+            plt.tight_layout()
+            pdf.savefig(fig)
+            if jupyter_nb:
+                plt.show(fig)
+            plt.close(fig)
+
+        def _wrap(txt: str, width: int = 60):
+            return "\n".join(fill(line, width) for line in txt.splitlines())
+
+        def _format_body(txt: str, width: int | None = 95):
+            txt = dedent(txt).strip("\n")
+            if width is None:
+                return txt
+            out_lines = []
+            for line in txt.splitlines():
+                if not line.strip():
+                    out_lines.append("")
+                    continue
+                m = re.match(r"^(\s*)([•\d]+[.)]?)(\s+)(.*)", line)
+                if m:
+                    lead_ws, bullet, spacer, rest = m.groups()
+                    bullet_str = f"{lead_ws}{bullet}{spacer}"
+                    wrapped = fill(
+                        rest,
+                        width=max(20, width - len(bullet_str)),
+                        subsequent_indent=" " * len(bullet_str),
+                    )
+                    wrapped_lines = wrapped.split("\n")
+                    wrapped_lines[0] = bullet_str + wrapped_lines[0]
+                    out_lines.extend(wrapped_lines)
+                else:
+                    out_lines.extend(fill(line, width).split("\n"))
+            return "\n".join(out_lines)
+
+        def section_page(title: str, body: str = ""):
+            is_overview = title.startswith("Results Overview")
+            sec_height = 11 if is_overview else 4.5
+            fig, ax = plt.subplots(figsize=(8.5, sec_height))
+            ax.axis("off")
+            main, sub = (
+                (title.split("–", 1) + [None])[:2] if "–" in title else (title, None)
+            )
+            base_fs = 20
+            fs_main = max(14, base_fs - int(len(main) / 40) * 2)
+            fs_sub = fs_main - 2
+            ax.text(
+                0.6 if is_overview else 0.5,
+                0.9 if is_overview else 0.75,
+                _wrap(main.strip(), 40),
+                fontsize=fs_main,
+                ha="center",
+                va="top",
+                weight="bold",
+            )
+            if sub:
+                ax.text(
+                    0.5,
+                    0.66,
+                    _wrap(sub.strip(), 50),
+                    fontsize=fs_sub,
+                    ha="center",
+                    va="top",
+                    style="italic",
+                )
+
+            def _render_line(text, y, is_path=False):
+                ax.text(
+                    0.12,
+                    y,
+                    text,
+                    fontsize=10,
+                    ha="left",
+                    va="top",
+                    wrap=False,
+                    fontfamily="monospace" if is_path else None,
+                    color="blue" if is_path else "black",
+                )
+
+            if body:
+                if is_overview:
+                    y_pos = 0.78
+                    for para in re.split(r"\n\s*\n", body.strip()):
+                        lines = para.strip().split("\n")
+                        header = lines[0].strip()
+                        rest = (
+                            "\n".join(l.strip() for l in lines[1:])
+                            if len(lines) > 1
+                            else ""
+                        )
+                        ax.text(
+                            0.12,
+                            y_pos,
+                            header,
+                            fontsize=12,
+                            fontweight="bold",
+                            ha="left",
+                            va="top",
+                            wrap=True,
+                        )
+                        y_pos -= 0.05
+
+                        if rest:
+                            for line in _format_body(rest, width=100).split("\n"):
+                                _render_line(
+                                    line, y_pos, is_path=line.lstrip().startswith("--")
+                                )
+                                y_pos -= 0.03
+                            y_pos -= 0.05
+                else:
+                    ax.text(
+                        0.12,
+                        0.46,
+                        _format_body(body, width=90),
+                        fontsize=10,
+                        ha="left",
+                        va="top",
+                        wrap=True,
+                    )
+            show(fig)
+
+        section_page("Results Overview", body=_format_body(intro_text, width=110))
+
+        csv_expl_rows_wrapped = [[col, _wrap(desc, 100)] for col, desc in csv_expl_rows]
+        section_page("CSV Column Reference")
+        fig_height = 0.35 * len(csv_expl_rows_wrapped) + 1.5
+        fig, ax = plt.subplots(figsize=(8.5, fig_height))
+        ax.axis("off")
+        table = ax.table(
+            cellText=csv_expl_rows_wrapped,
+            colLabels=["Column", "Description"],
+            colLoc="left",
+            colWidths=[0.25, 0.70],
+            loc="center",
+        )
+        table.auto_set_font_size(False)
+        table.set_fontsize(7)
+        table.scale(1, 1.2)
+        show(fig)
+
+        section_page(
+            "Filtering & Sorting Criteria",
+            body=(
+                "This section first lists threshold filters that every design must pass (Filtering Criteria table)"
+                "then explains how designs are ranked (Sorting Criteria table) weighted by their inverse importance."
+            ),
+        )
+
+        filters_df = pd.DataFrame(self.filters)
+        filters_df["Pass"] = 0
+        for i, filter in enumerate(self.filters):
+            filters_df.at[i, "Pass"] = self.df[f"pass_{filter['feature']}_filter"].sum()
+
+        fig_height = 0.4 * len(filters_df) + 2
+        fig, ax = plt.subplots(figsize=(8.5, fig_height))
+        ax.axis("off")
+        ax.text(
+            0.5,
+            1.0,
+            "Filtering Criteria",
+            fontsize=14,
+            ha="center",
+            va="bottom",
+            transform=ax.transAxes,
+        )
+        table = ax.table(
+            cellText=filters_df.values,
+            colLabels=filters_df.columns.tolist(),
+            loc="center",
+        )
+        table.auto_set_font_size(False)
+        table.set_fontsize(8)
+        table.scale(1, 1.2)
+        show(fig)
+        fig_height = 0.4 * len(metric_rows) + 1
+        fig, ax = plt.subplots(figsize=(8.5, fig_height))
+        ax.axis("off")
+        ax.text(
+            0.5,
+            1.0,
+            "Sorting Criteria",
+            fontsize=14,
+            ha="center",
+            va="bottom",
+            transform=ax.transAxes,
+        )
+        table = ax.table(
+            cellText=metric_rows,
+            colLabels=["Metric", "Inverse Importance"],
+            loc="center",
+        )
+        table.auto_set_font_size(False)
+        table.set_fontsize(8)
+        table.scale(1, 1.2)
+        show(fig)
+
+        section_page(
+            "Results Summary – Aggregate statistics",
+            body=(
+                f"Quick numeric overview comparing (i) ALL incoming designs, (ii) the top-{self.top_budget} highest-quality "
+                f"designs, and (iii) the {self.budget} quality+diversity designs produced by the lazy-greedy "
+                "selection."
+            ),
+        )
+        if rows:
+            fig, ax = plt.subplots(figsize=(8.5, 0.4 * len(rows) + 1))
+            ax.axis("off")
+            table = ax.table(
+                cellText=rows,
+                colLabels=row_headers,
+                loc="center",
+            )
+            table.auto_set_font_size(False)
+            table.set_fontsize(8)
+            table.scale(1, 1.2)
+            show(fig)
+
+        if self.plot_seq_logos:
+            section_page(
+                "Sequence Logos & AA Composition – Motifs & biases",
+                body=(
+                    "Sequence logos display the per-position amino-acid preferences after multiple-sequence alignment, "
+                    "highlighting motifs that emerge in the design sets.  The accompanying pies summarise overall "
+                    "hydrophobicity and charge composition comparing ALL, Top-quality (red), and Diversity-optimised subsets. "
+                ),
+            )
+        vis = (
+            "designed_sequence"
+            if (
+                self.df["designed_chain_sequence"].str.len().mean()
+                > 1.5 * self.df["designed_sequence"].str.len().mean()
+            )
+            else "designed_chain_sequence"
+        )
+        seq_sets = [
+            (
+                f"All {len(self.df)} {vis}",
+                self.df[vis].tolist(),
+            ),
+            (
+                f"Top {self.top_budget} {vis}",
+                self.df[vis].tolist()[: self.top_budget],
+            ),
+            (
+                f"Diverse {self.budget} {vis}",
+                self.df_div[vis].tolist(),
+            ),
+        ]
+        if self.plot_seq_logos:
+            for name, sequences in seq_sets:
+                show(create_alignment_logo(sequences, name))
+            for name, sequences in seq_sets:
+                show(aa_composition_pie(sequences, name))
+
+            if self.modality == "antibody":
+                for name, sequences in seq_sets:
+                    show(cdr_logo(sequences, name))
+
+        section_page(
+            "Scatter Plots – Metric relationships",
+            body="""
+                Each scatter page contains two panels:
+                • Left – all designs (grey) with overlays of Top (red) and Diverse (blue).
+                • Right – same but limited to designs passing the RMSD filter.
+                """,
+        )
+        for x, y in extra_pairs:
+            if x in self.df.columns and y in self.df.columns:
+                fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8.5, 4.5))
+                self._scatter_plus(ax1, self.df, x, y, "All samples")
+                if self.df["pass_filter_rmsd_filter"].sum() > 0:
+                    self._scatter_plus(
+                        ax2,
+                        self.df[self.df["pass_filter_rmsd_filter"]],
+                        x,
+                        y,
+                        "(Designs passing RMSD threshold)",
+                    )
+                show(fig)
+
+        section_page(
+            "Metric Distributions – Histograms",
+            body=(
+                "Distribution of each metric across all designs (grey) with overlays for Top-quality (red outline) "
+                "and Diversity-optimised (blue dashed) subsets. The right panel repeats the histogram but only for designs "
+                "that pass the RMSD threshold."
+            ),
+        )
+        for m in hist_metrics:
+            if m not in self.df.columns:
+                continue
+            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8.5, 4.5))
+            self._hist_plus(
+                ax1, self.df[m], self.df[: self.top_budget][m], self.df_div[m], m, ""
+            )
+            if self.df["pass_filter_rmsd_filter"].sum() > 0:
+                self._hist_plus(
+                    ax2,
+                    self.df[self.df["pass_filter_rmsd_filter"]][m],
+                    self.df[: self.top_budget][m],
+                    self.df_div[m],
+                    m,
+                    " (Designs passing RMSD threshold)",
+                )
+            show(fig)
+
+        if self.num_liability_plots:
+            section_page(
+                "Liability Analysis – Developability flags",
+                body=(
+                    "For the first few top designs we compute biochemical liability scores (deamidation, oxidation, etc.) "
+                    "and visualise them along the sequence with a heat-map. Dark red indicates positions needing "
+                    "attention during optimisation."
+                ),
+            )
+        # Plot Liability Heatmaps for Top-budget subset
+        for idx, row in tqdm(
+            enumerate(self.df[: self.num_liability_plots].itertuples(index=False)),
+            desc=f"Making liability plots for top {self.num_liability_plots} sequences.",
+        ):
+            seq = row.designed_sequence
+            try:
+                res = compute_liability_scores(
+                    [seq], modality=self.modality, peptide_type=self.peptide_type
+                )
+                liab = res.get(seq, {"score": None, "violations": []})
+                fig = plot_seq_liabilities(
+                    seq,
+                    f"Qualityrank {idx} {row.id}",
+                    liab["violations"],
+                    total_score=liab["score"],
+                )
+                show(fig)
+            except Exception as e:
+                print(f"  Error processing  {seq[:20]}: {e}")
+                plt.close("all")
+                continue
+
+        pdf.close()
+
+        print(
+            "A description of metrics and summarizing plots was written to:", pdf_path
+        )
+
+    def _scatter_plus(self, ax, df, x, y, title=""):
+        ax.scatter(
+            df[x],
+            df[y],
+            color="lightgray",
+            alpha=0.4,
+            s=14,
+            zorder=1,
+        )
+
+        ax.scatter(
+            self.df[: self.top_budget][x],
+            self.df[: self.top_budget][y],
+            facecolors="none",
+            edgecolors="red",
+            linewidth=1.5,
+            s=30,
+            zorder=2,
+            alpha=0.5,
+            label="top-quality",
+        )
+
+        if not self.df_div.empty:
+            ax.scatter(
+                self.df_div[x],
+                self.df_div[y],
+                facecolors="none",
+                edgecolors="blue",
+                linewidth=1.5,
+                s=50,  #  ← slightly larger
+                zorder=3,
+                alpha=0.5,
+                label="quality+diversity",
+            )
+
+        ax.set_xlabel(x)
+        ax.set_ylabel(y)
+        ax.set_title(title)
+        ax.legend()
+
+    def _hist_plus(self, ax, data_all, data_red, data_blue, metric, suffix):
+        # Check if data contains valid values for histogram
+        if data_all is None or len(data_all) == 0 or data_all.isna().all():
+            ax.text(
+                0.5,
+                0.5,
+                f"No valid data for {metric}",
+                ha="center",
+                va="center",
+                transform=ax.transAxes,
+            )
+            ax.set_xlabel(metric)
+            ax.set_ylabel("count")
+            ax.set_title(f"{metric}{suffix} (No data)")
+            return
+
+        # Filter out NaN values for histogram plotting
+        valid_data_all = data_all.dropna()
+        valid_data_red = (
+            data_red.dropna() if data_red is not None else pd.Series(dtype=float)
+        )
+        valid_data_blue = (
+            data_blue.dropna() if data_blue is not None else pd.Series(dtype=float)
+        )
+
+        if len(valid_data_all) == 0:
+            ax.text(
+                0.5,
+                0.5,
+                f"No valid data for {metric}",
+                ha="center",
+                va="center",
+                transform=ax.transAxes,
+            )
+            ax.set_xlabel(metric)
+            ax.set_ylabel("count")
+            ax.set_title(f"{metric}{suffix} (No data)")
+            return
+
+        # Create histogram with valid data
+        ax.hist(valid_data_all, bins=30, color="lightgray", alpha=0.6, label="all")
+
+        if len(valid_data_red) > 0:
+            ax.hist(
+                valid_data_red,
+                bins=30,
+                histtype="step",
+                linewidth=1.7,
+                color="red",
+                label="top-quality",
+            )
+        if len(data_blue) and len(valid_data_blue) > 0:
+            ax.hist(
+                valid_data_blue,
+                bins=30,
+                histtype="step",
+                linewidth=1.7,
+                color="blue",
+                linestyle="--",
+                label="quality+diversity",
+            )
+
+        ax.set_xlabel(metric)
+        ax.set_ylabel("count")
+        ax.set_title(f"{metric}{suffix}")
+        ax.legend()

model/boltzgen/task/filter/seqplot_utils.py

@@ -0,0 +1,505 @@
+import os
+import tempfile
+from matplotlib import pyplot as plt
+import matplotlib.colors as mcolors
+from matplotlib.patches import Patch
+import numpy as np
+import pandas as pd
+import logomaker
+
+from boltzgen.data import const
+
+# Define amino acid properties for peptide visualization
+HYDROPHOBIC = set("ACILMFWV")
+NEUTRAL = set("GHPTY")
+HYDROPHILIC = set("RNDQEK")
+POSITIVE = set("RHK")
+NEGATIVE = set("DE")
+AA20 = list("ACDEFGHIKLMNPQRSTVWY")
+
+COLOR_TABLE = {
+    ("hydrophobic", "positive"): "#d62728",
+    ("hydrophobic", "negative"): "#1f77b4",
+    ("hydrophobic", "uncharged"): "#2ca02c",
+    ("neutral", "positive"): "#ff7f0e",
+    ("neutral", "negative"): "#17becf",
+    ("neutral", "uncharged"): "#bcbd22",
+    ("hydrophilic", "positive"): "#e377c2",
+    ("hydrophilic", "negative"): "#1f99d4",
+    ("hydrophilic", "uncharged"): "#9467bd",
+}
+
+
+def _hydropathy_class(res):
+    if res in HYDROPHOBIC:
+        return "hydrophobic"
+    elif res in HYDROPHILIC:
+        return "hydrophilic"
+    else:
+        return "neutral"
+
+
+def _charge_class(res):
+    if res in POSITIVE:
+        return "positive"
+    elif res in NEGATIVE:
+        return "negative"
+    else:
+        return "uncharged"
+
+
+hydrophobicity_colors = {
+    aa: COLOR_TABLE[(_hydropathy_class(aa), _charge_class(aa))] for aa in AA20
+}
+
+
+def create_alignment_logo(sequences, name, width=10):
+    """Create sequence logo from aligned sequences."""
+
+    aligned_sequences, _ = align_peptide_sequences(sequences)
+    if not aligned_sequences or len(aligned_sequences) < 2:
+        msg = f"Warning: Not enough sequences for logo: {len(aligned_sequences)}"
+        print(msg)
+        return
+
+    # Count amino acids at each position
+    max_len = max(len(seq) for seq in aligned_sequences)
+    counts = {aa: np.zeros(max_len, dtype=int) for aa in AA20}
+
+    for i, seq in enumerate(aligned_sequences):
+        for pos, aa in enumerate(seq):
+            if pos < max_len and aa in AA20:
+                counts[aa][pos] += 1
+
+    # Convert to DataFrame
+    counts_df = pd.DataFrame(counts)
+    counts_df.index.name = "position"
+
+    return draw_logo(counts_df, name, width=width)
+
+
+def aa_composition_pie(sequences, name):
+    # Generate amino acid composition analysis
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
+
+    # Hydrophobicity distribution
+    hydrophobicity_counts = {
+        "hydrophobic": 0,
+        "neutral": 0,
+        "hydrophilic": 0,
+    }
+    for seq in sequences:
+        for aa in seq:
+            if aa in AA20:
+                hydrophobicity_counts[_hydropathy_class(aa)] += 1
+
+    ax1.pie(
+        hydrophobicity_counts.values(),
+        labels=hydrophobicity_counts.keys(),
+        autopct="%1.1f%%",
+    )
+    ax1.set_title(f"{name}\n Hydrophobicity Distribution", y=1.02, wrap=True)
+
+    # Charge distribution
+    charge_counts = {"positive": 0, "negative": 0, "uncharged": 0}
+    for seq in sequences:
+        for aa in seq:
+            if aa in AA20:
+                charge_counts[_charge_class(aa)] += 1
+
+    ax2.pie(
+        charge_counts.values(),
+        labels=charge_counts.keys(),
+        autopct="%1.1f%%",
+    )
+    ax2.set_title(f"{name}\n  Charge Distribution", y=1.02, wrap=True)
+    fig.tight_layout(rect=[0, 0, 1, 0.92])
+    return fig
+
+
+def align_peptide_sequences(sequences):
+    """Align peptide sequences using BioPython's multiple sequence alignment."""
+
+    if len(sequences) < 2:
+        print("Warning: Only one sequence, no alignment needed")
+        return sequences, []
+
+    # Fallback to pairwise alignment (less optimal but functional)
+    from Bio import Align
+    from Bio.Align import substitution_matrices
+
+    # Use BLOSUM62 matrix for protein alignment
+    matrix = substitution_matrices.load("BLOSUM62")
+
+    # Create pairwise aligner
+    aligner = Align.PairwiseAligner()
+    aligner.substitution_matrix = matrix
+    aligner.open_gap_score = -11
+    aligner.extend_gap_score = -1
+
+    # Start with first sequence as reference
+    aligned_seqs = [sequences[0]]
+    alignment_scores = []
+
+    # Align each sequence to the reference
+    for i, seq in enumerate(sequences[1:], 1):
+        try:
+            alignment = aligner.align(sequences[0], seq)
+            if alignment:
+                # Get the QUERY sequence (the different one) with gaps
+                aligned_seq = str(alignment[0].query)
+
+                aligned_seqs.append(aligned_seq)
+                alignment_scores.append(alignment[0].score)
+            else:
+                aligned_seqs.append(seq)
+                alignment_scores.append(0)
+        except Exception as e:
+            print(f"      DEBUG: Alignment failed: {e}")
+            # Fallback: just add the sequence as-is
+            aligned_seqs.append(seq)
+            alignment_scores.append(0)
+    return aligned_seqs, alignment_scores
+
+
+def create_temp_fasta(sequences, names):
+    """Create a temporary FASTA file from sequences."""
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".fasta", delete=False) as f:
+        for i, (seq, name) in enumerate(zip(sequences, names)):
+            f.write(f">{name}\n{seq}\n")
+        return f.name
+
+
+def build_region_alignment(chains, region_name):
+    """Build alignment for a specific CDR region."""
+    if not chains:
+        return [], []
+
+    all_positions = set()
+    per_chain_pos_to_aa = []
+    for ch in chains:
+        reg = ch.regions.get(region_name, {})
+        pos_to_aa = {str(pos): aa for pos, aa in reg.items()}
+        per_chain_pos_to_aa.append(pos_to_aa)
+        all_positions.update(pos_to_aa.keys())
+
+    def pos_key(p):
+        p = p.replace("H", "")
+        num = ""
+        ins = ""
+        for c in p:
+            if c.isdigit():
+                num += c
+            else:
+                ins += c
+        return (int(num), ins or "")
+
+    ordered_positions = sorted(all_positions, key=pos_key)
+    msa = []
+    ids = []
+    for i, ch in enumerate(chains):
+        pos_to_aa = per_chain_pos_to_aa[i]
+        row = "".join(pos_to_aa.get(p, "-") for p in ordered_positions)
+        name = ch.name if getattr(ch, "name", None) else f"seq_{i + 1}"
+        ids.append(name)
+        msa.append(row)
+    return ids, msa
+
+
+def counts_matrix_from_msa(msa, alphabet):
+    """Compute counts matrix from MSA."""
+    if not msa:
+        return pd.DataFrame()
+
+    msa = [row.upper() for row in msa]
+    L = len(msa[0])
+    counts = {aa: np.zeros(L, dtype=int) for aa in alphabet}
+
+    for row in msa:
+        assert len(row) == L, "All rows must have same length"
+        for j, aa in enumerate(row):
+            if aa in ["-", "X", "J", "U", "O"]:
+                continue
+            if aa in counts:
+                counts[aa][j] += 1
+
+    df = pd.DataFrame(counts)
+    df.index.name = "position"
+    return df
+
+
+def draw_logo(counts, title, width=10):
+    """Draw sequence logo using logomaker."""
+    probs = counts.div(counts.sum(axis=1).replace(0, np.nan), axis=0).fillna(0)
+    entropy = (
+        -(probs.replace(0, np.nan) * np.log2(probs.replace(0, np.nan)))
+        .sum(axis=1)
+        .fillna(0)
+    )
+    info = np.log2(20) - entropy
+    heights = probs.mul(info, axis=0)
+
+    fig, ax = plt.subplots(figsize=(width, 4))
+    logo = logomaker.Logo(
+        heights,
+        color_scheme=hydrophobicity_colors,
+        shade_below=0.5,
+        fade_below=0.5,
+        vpad=0.05,
+        width=0.8,
+        ax=ax,
+    )
+
+    legend_elements = []
+    for (hydro, charge), color in COLOR_TABLE.items():
+        label = f"{hydro.capitalize()} + {charge}"
+        legend_elements.append(Patch(facecolor=color, edgecolor="black", label=label))
+
+    ax.legend(handles=legend_elements, bbox_to_anchor=(1.05, 1), loc="upper left")
+    ax.set_title(title)
+    ax.set_xlabel("Aligned position")
+    ax.set_ylabel("Information (bits)")
+
+    plt.tight_layout()
+    return fig
+
+
+def cdr_logo(sequences, name):
+    from abnumber import Chain
+
+    # Create temporary FASTA
+    names = [f"seq_{i + 1}" for i in range(len(sequences))]
+    temp_fasta = create_temp_fasta(sequences, names)
+
+    # Load chains using AbNumber
+    chains = []
+    for ch in Chain.from_fasta(
+        temp_fasta,
+        scheme="chothia",
+        cdr_definition="chothia",
+        as_generator=True,
+        allowed_species=["alpaca"],
+    ):
+        if ch.chain_type == "H":
+            chains.append(ch)
+
+    # Generate composite figure
+    fig, axes = plt.subplots(3, 1, figsize=(12, 8), constrained_layout=True)
+    for ax, cdr in zip(axes, ["CDR1", "CDR2", "CDR3"]):
+        ids, msa = build_region_alignment(chains, cdr)
+        if not msa:
+            ax.set_axis_off()
+            continue
+
+        counts = counts_matrix_from_msa(msa, AA20)
+        probs = counts.div(counts.sum(axis=1).replace(0, np.nan), axis=0).fillna(0)
+        entropy = (
+            -(probs.replace(0, np.nan) * np.log2(probs.replace(0, np.nan)))
+            .sum(axis=1)
+            .fillna(0)
+        )
+        info = np.log2(20) - entropy
+        heights = probs.mul(info, axis=0)
+
+        logo = logomaker.Logo(
+            heights,
+            ax=ax,
+            color_scheme=hydrophobicity_colors,
+            shade_below=0.5,
+            fade_below=0.5,
+            vpad=0.05,
+            width=0.8,
+        )
+        ax.set_title(f"{name} {cdr}")
+        ax.set_xlabel("")
+        ax.set_ylabel("bits")
+
+    # Add legend
+    legend_elements = [
+        Patch(
+            facecolor=color,
+            edgecolor="black",
+            label=f"{hydro.capitalize()} + {charge}",
+        )
+        for (hydro, charge), color in COLOR_TABLE.items()
+    ]
+    fig.legend(
+        handles=legend_elements,
+        loc="center left",
+        bbox_to_anchor=(1.02, 0.5),
+        borderaxespad=0.0,
+    )
+    os.unlink(temp_fasta)
+    return fig
+
+
+def plot_seq_liabilities(
+    sequence,
+    name,
+    violations,
+    total_score=None,
+):
+    """
+    Plot sequence with colored backgrounds for each motif region.
+
+    Args:
+        sequence: The amino acid sequence
+        violations: List of violation dictionaries
+        total_score: Total liability score
+    """
+    n = len(sequence)
+    sev_arr = [0] * n
+    violation_types = {}  # Track unique violation types and their max severity
+
+    # Process violations and compute positions for those that don't have explicit pos/len
+    for v in violations:
+        motif_name = v["motif"]
+
+        # Handle violations that need position computation
+        if motif_name == "ConsecIdentical" and v.get("pos") is None:
+            # Find consecutive identical residues
+            for i in range(n - 1):
+                if sequence[i] == sequence[i + 1]:
+                    # Color both positions
+                    sev_arr[i] = max(sev_arr[i], v["severity"])
+                    sev_arr[i + 1] = max(sev_arr[i + 1], v["severity"])
+        elif motif_name == "LongHydrophobic" and v.get("pos") is None:
+            # Find stretches of >4 consecutive hydrophobic residues
+            hydrophobic_residues = set("FILVWY")
+            for i in range(n - 4):
+                stretch = sequence[i : i + 5]
+                if all(aa in hydrophobic_residues for aa in stretch):
+                    # Color the entire stretch
+                    for j in range(i, i + 5):
+                        if j < n:
+                            sev_arr[j] = max(sev_arr[j], v["severity"])
+        elif v.get("pos") and v.get("len"):
+            # Standard violations with explicit positions
+            for i in range(v["pos"] - 1, v["pos"] - 1 + v["len"]):
+                if 0 <= i < n:  # Ensure index is within bounds
+                    sev_arr[i] = max(sev_arr[i], v["severity"])
+
+        # Track violation types for legend
+        if motif_name not in violation_types:
+            violation_types[motif_name] = v["severity"]
+        else:
+            violation_types[motif_name] = max(
+                violation_types[motif_name], v["severity"]
+            )
+
+    cmap = mcolors.LinearSegmentedColormap.from_list(
+        "HeatMap", ["white", "yellow", "red"]
+    )
+    norm = mcolors.Normalize(vmin=0, vmax=max(const.liability_severity.values()))
+
+    # Better figure size calculation to prevent letter cropping
+    # Ensure minimum width for short sequences and proper height
+    min_width = max(8, n * 0.3)  # At least 8 inches wide, or 0.3 inches per residue
+    fig_height = 4  # Fixed height for consistency
+
+    # Create figure with space for legend
+    fig, (ax, ax_legend) = plt.subplots(
+        2, 1, figsize=(min_width, fig_height), gridspec_kw={"height_ratios": [2.5, 1]}
+    )
+
+    # Main sequence plot
+    for idx, aa in enumerate(sequence):
+        color = cmap(norm(sev_arr[idx]))
+        ax.text(
+            idx,
+            0.5,
+            aa,
+            ha="center",
+            va="center",
+            fontsize=12,
+            bbox=dict(facecolor=color, edgecolor="none", boxstyle="square,pad=0.1"),
+        )
+    ax.set_xlim(-0.5, n - 0.5)
+    ax.set_ylim(0, 1)
+    ax.axis("off")
+
+    if total_score is not None:
+        name += f" (Score: {total_score})"
+    ax.set_title(name)
+
+    # Colorbar
+    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
+    sm.set_array([])
+    plt.colorbar(sm, ax=ax, orientation="horizontal", pad=0.2, label="Severity")
+
+    # Legend showing violation types with motif patterns
+    if violation_types:
+        ax_legend.axis("off")
+        legend_items = []
+        legend_labels = []
+
+        # Sort by severity (descending)
+        sorted_violations = sorted(
+            violation_types.items(), key=lambda x: x[1], reverse=True
+        )
+
+        # Motif pattern mapping for legend
+        motif_patterns = {
+            "DeAmdH": "N[GS]",
+            "DeAmdM": "N[AHNT]",
+            "DeAmdL": "[STK]N",
+            "Ngly": "N[^P][ST]",
+            "Isom": "D[DG HST]",
+            "Isomer": "DG|DS|DD",
+            "FragH": "DP",
+            "FragM": "TS",
+            "TrpOx": "W",
+            "MetOx": "M",
+            "Hydro": "NP",
+            "IntBind": "GPR|RGD|RYD|LDV|DGE|KGD|NGR",
+            "Polyreactive": "GGG|WWW|GG|RR|VG|VVV|YY|WxW",
+            "AggPatch": "FHW",
+            "ViscPatch": "HYF|HWH",
+            "UnpairedCys": "C",
+            "HighNetCharge": "net charge > +1",
+            "AspBridge": "N[GSQA]",
+            "AspCleave": "D[PGS]",
+            "NTCycl": "^[QN]",
+            "ProtTryp": "[KR](?=.)",
+            "DPP4": "^[PX]?[AP]",
+            "CysOx": "C",
+            "HydroPatch": "[FILVWY]{3,}",
+            "LowHydrophilic": "< 40% hydrophilic",
+            "ConsecIdentical": "consecutive identical",
+            "LongHydrophobic": "> 4 consecutive hydrophobic",
+        }
+
+        for motif_name, severity in sorted_violations:
+            color = cmap(norm(severity))
+            legend_items.append(
+                plt.Rectangle((0, 0), 1, 1, facecolor=color, edgecolor="black")
+            )
+            pattern = motif_patterns.get(motif_name, "unknown")
+            legend_labels.append(f"{motif_name}: {pattern}")
+
+        ax_legend.legend(
+            legend_items,
+            legend_labels,
+            loc="center",
+            title="Violations",
+            ncol=1,
+            fontsize=8,
+        )
+    else:
+        # No violations detected
+        ax_legend.axis("off")
+        ax_legend.text(
+            0.5,
+            0.5,
+            "No violations detected",
+            ha="center",
+            va="center",
+            fontsize=12,
+            bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgreen", alpha=0.7),
+        )
+
+        ax_legend.set_title("Violations", fontsize=10)
+
+    plt.tight_layout()
+
+    return fig

model/boltzgen/task/predict/data_from_generated.py

@@ -0,0 +1,882 @@
+from dataclasses import dataclass
+from pathlib import Path
+import random
+import re
+import warnings
+from typing import Dict, List, Optional
+from collections import defaultdict
+from rdkit.Chem import Mol
+import pickle
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from torch import Tensor
+from torch.utils.data import DataLoader
+
+from boltzgen.data import const
+from boltzgen.data.data import Input, Structure, Tokenized
+from boltzgen.data.feature.featurizer import Featurizer
+from boltzgen.data.mol import load_canonicals, load_molecules
+from boltzgen.data.pad import pad_to_max
+from boltzgen.data.parse import mmcif
+from boltzgen.data.parse.pdb_parser import parse_pdb
+from boltzgen.data.template.features import (
+    load_dummy_templates,
+)
+from boltzgen.data.parse.schema import parse_redesign_yaml
+from boltzgen.data.tokenize.tokenizer import Tokenizer
+
+
+class DataFetchException(Exception):
+    pass
+
+
+@dataclass
+class DataConfig:
+    """Data configuration."""
+
+    num_targets: int
+    samples_per_target: int
+    moldir: str
+    tokenizer: Tokenizer
+    featurizer: Featurizer
+    batch_size: int
+    num_workers: int
+    pin_memory: bool
+    suffix: str = ".cif"
+    suffix_native: str = "_native.cif"
+    suffix_metadata: str = ".npz"
+    target_id_regex: str = (
+        r"^(?:(?:sample\d+_|batch\d+_|rank\d+_)+)?([^_]+)(?:_[^_]+)*?(?:_(?:gen))*$"
+    )
+    design: bool = False
+    # Featurizer args (if design is True these should match with training config):
+    backbone_only: bool = False
+    atom14: bool = True
+    max_seqs: int = 1
+    inverse_fold: bool = False
+    extra_mol_dir: Optional[str] = None
+    disulfide_prob: float = 1.0
+    disulfide_on: bool = False
+    design_mask_override: Optional[str] = None
+    multiplicity: int = 1
+    return_designfolding: bool = False
+
+
+def collate(data: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
+    """Collate the data.
+
+    Parameters
+    ----------
+    data : List[Dict[str, Tensor]]
+        The data to collate.
+
+    Returns
+    -------
+    Dict[str, Tensor]
+        The collated data.
+
+    """
+    # Get the keys
+    keys = data[0].keys()
+
+    # Collate the data
+    collated = {}
+    for key in keys:
+        values = [d[key] for d in data]
+
+        if key not in [
+            "metadata",
+            "str_gen",
+            "id",
+            "path",
+            "native_metadata",
+            "native_str_gen",
+            "native_id",
+            "native_path",
+            "exception",
+            "native_exception",
+            "skip",
+            "native_skip",
+            "structure_bonds",
+            "native_structure_bonds",
+            "extra_mols",
+            "native_extra_mols",
+            "structure",
+            "tokenized",
+            "data_sample_idx",
+        ]:
+            # Check if all have the same shape
+            shape = values[0].shape
+            if not all(v.shape == shape for v in values):
+                values = pad_to_max(values, 0)
+            else:
+                values = torch.stack(values, dim=0)
+
+        # Stack the values
+        collated[key] = values
+
+    return collated
+
+
+@dataclass(frozen=True)
+class TemplateInfo:
+    """TemplateInfo datatype."""
+
+    name: str
+    query_chain: str
+    query_st: int
+    query_en: int
+    template_chain: str
+    template_st: int
+    template_en: int
+
+
+def template_from_tokens(
+    tokenized: Tokenized,
+    token_mask: np.ndarray[bool],
+    tdim: int = 1,
+) -> dict[str, torch.Tensor]:
+    """Get template features where the tokens specified in token_mask have their structure specified."""
+    # Get num token
+    num_tokens = len(tokenized.tokens)
+
+    # Allocate features
+    res_type = np.zeros((tdim, num_tokens), dtype=np.int64)
+    frame_rot = np.zeros((tdim, num_tokens, 3, 3), dtype=np.float32)
+    frame_t = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
+    cb_coords = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
+    ca_coords = np.zeros((tdim, num_tokens, 3), dtype=np.float32)
+    frame_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
+    cb_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
+    template_mask = np.zeros((tdim, num_tokens), dtype=np.float32)
+    query_to_template = np.zeros((tdim, num_tokens), dtype=np.int64)
+    visibility_ids = np.zeros((tdim, num_tokens), dtype=np.float32)
+
+    # Now create features per token
+    template_indices = np.where(token_mask)[0]
+    for token_idx in template_indices:
+        token = tokenized.tokens[token_idx]
+        res_type[:, token_idx] = token["res_type"]
+        frame_rot[:, token_idx] = token["frame_rot"].reshape(3, 3)
+        frame_t[:, token_idx] = token["frame_t"]
+        cb_coords[:, token_idx] = token["disto_coords"]
+        ca_coords[:, token_idx] = token["center_coords"]
+        cb_mask[:, token_idx] = token["disto_mask"]
+        frame_mask[:, token_idx] = token["frame_mask"]
+        template_mask[:, token_idx] = 1.0
+        visibility_ids[:, token_idx] = 1
+
+    # Convert to one-hot
+    res_type = torch.from_numpy(res_type)
+    res_type = torch.nn.functional.one_hot(res_type, num_classes=const.num_tokens)
+
+    return {
+        "template_restype": res_type,
+        "template_frame_rot": torch.from_numpy(frame_rot),
+        "template_frame_t": torch.from_numpy(frame_t),
+        "template_cb": torch.from_numpy(cb_coords),
+        "template_ca": torch.from_numpy(ca_coords),
+        "template_mask_cb": torch.from_numpy(cb_mask),
+        "template_mask_frame": torch.from_numpy(frame_mask),
+        "template_mask": torch.from_numpy(template_mask),
+        "query_to_template": torch.from_numpy(query_to_template),
+        "visibility_ids": torch.from_numpy(visibility_ids),
+    }
+
+
+class FromGeneratedDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        generated_paths: List[Path],
+        metadata_paths: List[Path],
+        native_paths: List[Path],
+        moldir: Path,
+        canonicals: dict[str, Mol],
+        tokenizer: Tokenizer,
+        featurizer: Featurizer,
+        return_native: bool = False,
+        reference_metadata_dir: Optional[Path] = None,
+        target_templates: bool = False,
+        design_mask_templates: bool = False,
+        compute_affinity: bool = False,
+        design: bool = False,
+        backbone_only: bool = False,
+        atom14: bool = True,
+        max_seqs: int = 1,
+        inverse_fold: bool = False,
+        extra_mol_dir: Optional[Path] = None,
+        extra_features: Optional[List[str]] = None,
+        disulfide_prob: float = 1.0,
+        disulfide_on: bool = False,
+        design_mask_override: Optional[str] = None,
+        use_new_design_mask: bool = False,
+        multiplicity: int = 1,
+        return_designfolding=False,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        design : bool
+            Set to True if this dataset is used to make predictions over (i.e. design some parts
+            of the structure). Set to False if this dataset is used to only evaluate the predictions
+            under the paths (i.e. no design is done).
+        """
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.moldir = moldir
+        self.canonicals = canonicals
+        self.featurizer = featurizer
+        self.metadata_paths = metadata_paths
+        self.generated_paths = generated_paths
+        self.native_paths = native_paths
+        self.return_native = return_native
+        self.reference_metadata_dir = reference_metadata_dir
+        self.target_templates = target_templates
+        self.design_mask_templates = design_mask_templates
+        self.compute_affinity = compute_affinity
+        self.design = design
+        self.backbone_only = backbone_only
+        self.atom14 = atom14
+        self.max_seqs = max_seqs
+        self.inverse_fold = inverse_fold
+        self.extra_mol_dir = extra_mol_dir
+        self.extra_features = (
+            set(extra_features) if extra_features is not None else set()
+        )
+        self.disulfide_prob = disulfide_prob
+        self.disulfide_on = disulfide_on
+        self.design_mask_override = design_mask_override
+        self.use_new_design_mask = use_new_design_mask
+        self.multiplicity = multiplicity
+        self.return_designfolding = return_designfolding
+
+    def __getitem__(self, idx: int) -> Dict:
+        """Get an item from the dataset.
+
+        Returns
+        -------
+        Dict[str, Tensor]
+
+        """
+        data_sample_idx = idx // len(self.generated_paths)
+        idx = idx % len(self.generated_paths)
+
+        try:
+            feat = self.getitem_from_paths(
+                self.metadata_paths[idx],
+                self.generated_paths[idx],
+                self.native_paths[idx],
+            )
+            if self.multiplicity > 1:
+                feat["data_sample_idx"] = data_sample_idx
+            return feat
+        except DataFetchException:
+            idx = random.randint(0, len(self) - 1)
+            feat = self.getitem_from_paths(
+                self.metadata_paths[idx],
+                self.generated_paths[idx],
+                self.native_paths[idx],
+            )
+            if self.multiplicity > 1:
+                feat["data_sample_idx"] = data_sample_idx
+            return feat
+
+    def get_sample(self, design_dir: Path, sample_id: Optional[str] = None) -> Dict:
+        metadata_path = design_dir / f"{sample_id}.npz"
+        generated_path = design_dir / f"{sample_id}.cif"
+        native_path = design_dir / f"{sample_id}_native.cif"
+        return self.getitem_from_paths(metadata_path, generated_path, native_path)
+
+    def getitem_from_paths(self, metadata_path, generated_path, native_path) -> Dict:
+        """Get an item from the dataset.
+
+        Returns
+        -------
+        Dict[str, Tensor]
+
+        """
+        # Get metadata
+
+        if self.reference_metadata_dir:
+            reference_metadata_path = self.reference_metadata_dir / metadata_path.name
+            metadata = np.load(reference_metadata_path)
+        else:
+            metadata = np.load(metadata_path)
+
+        # get conditioning information from metadata
+        metadata_design_mask = metadata["design_mask"]
+        if self.use_new_design_mask:
+            design_mask = metadata["inverse_fold_design_mask"].astype(np.float32)
+        else:
+            design_mask = metadata_design_mask
+
+        ss_type = None
+        if "ss_type" in metadata:
+            ss_type = metadata["ss_type"]
+
+        binding_type = None
+        if "binding_type" in metadata:
+            binding_type = metadata["binding_type"]
+
+        # Per-residue amino acid constraints for inverse folding
+        aa_constraint_mask = None
+        if "aa_constraint_mask" in metadata:
+            loaded_mask = metadata["aa_constraint_mask"]
+            # Validate the loaded mask is a proper array with expected shape
+            if (
+                isinstance(loaded_mask, np.ndarray)
+                and loaded_mask.ndim == 2
+                and loaded_mask.shape[1] == 20  # 20 canonical amino acids
+            ):
+                aa_constraint_mask = loaded_mask
+            else:
+                warnings.warn(
+                    f"Invalid aa_constraint_mask in NPZ: "
+                    f"type={type(loaded_mask)}, shape={getattr(loaded_mask, 'shape', 'N/A')}. "
+                    f"Expected ndarray with shape (N, 20). Ignoring constraints.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+
+        # Get features
+        feat = self.get_feat(generated_path, design_mask, ss_type, binding_type, aa_constraint_mask)
+
+        # Get native features
+        if self.return_native:
+            if "native_design_mask" in metadata.keys():
+                feat_native = self.get_feat(native_path, metadata["native_design_mask"])
+            else:
+                feat_native = self.get_feat(native_path, metadata_design_mask)
+
+            for k, v in feat_native.items():
+                feat[f"native_{k}"] = v
+
+        return feat
+
+    def get_feat(self, path, design_mask, ss_type=None, binding_type=None, aa_constraint_mask=None):
+        # Load design
+        if self.extra_mol_dir is not None:
+            mols = {
+                path.stem: pickle.load(path.open("rb"))
+                for path in self.extra_mol_dir.glob("*.pkl")
+            }
+            for mol_name, mol in mols.items():
+                element_counts = defaultdict(int)
+                for i, atom in enumerate(mol.GetAtoms()):
+                    symbol = atom.GetSymbol()
+                    element_counts[symbol] += 1
+                    atom_name = f"{symbol}{element_counts[symbol]}"
+                    atom.SetProp("name", atom_name)
+
+        try:
+            if path.suffix == ".cif":
+                structure = mmcif.parse_mmcif(
+                    path, mols, moldir=self.moldir, use_original_res_idx=False
+                ).data
+            elif path.suffix == ".pdb":
+                structure = parse_pdb(
+                    path, moldir=self.moldir, use_original_res_idx=False
+                ).data
+            else:
+                raise ValueError(f"Invalid path:{path}")  # noqa: T201
+        except Exception as e:  # noqa: BLE001
+            print(f"Failed to parse {path} with error {e}. Skipping.")  # noqa: T201
+            raise DataFetchException() from e
+
+        # Tokenize structure
+
+        try:
+            tokenized = self.tokenizer.tokenize(
+                structure, inverse_fold=self.inverse_fold
+            )
+        except Exception as e:  # noqa: BLE001
+            print(f"Tokenizer failed on {path} with error {e}. Skipping.")  # noqa: T201
+            raise DataFetchException() from e
+
+        # Propagate design mask to obtain chain_design_mask (True whenever something is covalently bound to any residue that is in a chain that contains a design residue).
+        chain_design_mask = design_mask.astype(bool)
+        asym_id = tokenized.tokens["asym_id"]
+        while True:
+            design_chains = np.unique(asym_id[chain_design_mask])
+            chain_propagated = np.isin(asym_id, design_chains)
+            for i, j, _ in tokenized.bonds:
+                if any([chain_propagated[i], chain_propagated[j]]):
+                    chain_propagated[i] = True
+                    chain_propagated[j] = True
+            if np.equal(chain_propagated, chain_design_mask).all():
+                break
+            chain_design_mask = chain_propagated.astype(bool)
+
+        # Extract design for refolding the design only
+        if self.return_designfolding:
+            residue_design_mask = np.zeros(tokenized.token_to_res.max() + 1, dtype=bool)
+            np.put_along_axis(
+                residue_design_mask, tokenized.token_to_res, chain_design_mask, axis=0
+            )
+            structure = Structure.extract_residues(structure, residue_design_mask)
+            tokenized = self.tokenizer.tokenize(structure)
+            design_mask = design_mask[chain_design_mask]
+            chain_design_mask = chain_design_mask[chain_design_mask]
+
+        # For inverse folding, condition even on structure selected for design
+        if self.inverse_fold:
+            tokenized.tokens["structure_group"] = 1
+
+        try:
+            # Try to find molecules in the dataset moldir if provided
+            # Find missing ones in global moldir and check if all found
+            molecules = {}
+            molecules.update(self.canonicals)
+            mol_names = set(tokenized.tokens["res_name"].tolist())
+            mol_names = mol_names - set(self.canonicals.keys())
+            if mols is not None:
+                molecules.update(mols)
+            mol_names = mol_names - set(molecules.keys())
+            if self.moldir is not None:
+                molecules.update(load_molecules(self.moldir, mol_names))
+            molecules.update(load_molecules(self.moldir, mol_names))
+        except Exception as e:  # noqa: BLE001
+            print(f"Molecule loading failed for {path} with error {e}. Skipping.")
+            raise DataFetchException() from e
+
+        # Set design mask for tokens. This will impact the featurization and add the atom14 features
+        if self.design:
+            tokenized.tokens["design_mask"] = torch.from_numpy(design_mask).bool()
+
+        # Finalize input data
+        input_data = Input(
+            tokens=tokenized.tokens,
+            bonds=tokenized.bonds,
+            token_to_res=tokenized.token_to_res,
+            structure=structure,
+            msa={},
+            templates=None,
+        )
+
+        # Compute features
+        try:
+            features = self.featurizer.process(
+                input_data,
+                molecules=molecules,
+                random=np.random.default_rng(None),
+                training=False,
+                max_seqs=self.max_seqs,
+                backbone_only=self.backbone_only,
+                atom14=self.atom14,
+                design=True,
+                compute_affinity=self.compute_affinity,
+                override_method="X-RAY DIFFRACTION",
+                disulfide_prob=self.disulfide_prob,
+                disulfide_on=self.disulfide_on,
+            )
+        except Exception as e:  # noqa: BLE001
+            print(f"Featurizer failed on {path} with error {e}. Skipping.")  # noqa: T201
+            raise DataFetchException() from e
+
+        # Set chain design mask
+        features["chain_design_mask"] = torch.from_numpy(chain_design_mask)
+
+        # Set conditioning variables that were set during design
+        if ss_type is not None:
+            features["ss_type"] = torch.from_numpy(ss_type).long()
+        if binding_type is not None:
+            features["binding_type"] = torch.from_numpy(binding_type).long()
+        # Per-residue amino acid constraints for inverse folding
+        if aa_constraint_mask is not None:
+            features["aa_constraint_mask"] = torch.from_numpy(aa_constraint_mask).float()
+
+        # If we do not want the design mask to impact the featurizer (e.g. represent atoms as atom14), we set the design mask only here.
+        if not self.design:
+            features["design_mask"] = torch.from_numpy(design_mask).bool()
+
+        # set chain_design_mask
+        # Override design mask for inverse folding if the part that should be inverse folded differs from the previously designed part.
+        if self.design and self.design_mask_override is not None:
+            msg = f"design mask being overridden with user input: {self.design_mask_override}"
+            print(msg)
+            new_design_mask = parse_redesign_yaml(
+                Path(self.design_mask_override), tokenized
+            )
+            features["inverse_fold_design_mask"] = torch.from_numpy(
+                new_design_mask
+            ).bool()
+
+        # Perform assertions
+        if len(tokenized.tokens) != len(design_mask):
+            print(
+                f"WARNING: len(tokenized.tokens) [{len(tokenized.tokens)}] != len(design_mask) "
+                f"[{len(design_mask)}] for {path}"
+            )
+            features["exception"] = True
+            return features
+        else:
+            features["exception"] = False
+
+        # Set templates
+        if self.target_templates:
+            if self.design_mask_templates:
+                template_mask = ~features["design_mask"].numpy()
+            else:
+                template_mask = ~features["chain_design_mask"].numpy()
+            templates_features = template_from_tokens(tokenized, template_mask)
+        else:
+            # Compute template features
+            templates_features = load_dummy_templates(
+                tdim=1, num_tokens=len(features["res_type"])
+            )
+        features.update(templates_features)
+
+        features["affinity_token_mask"] = (
+            features["mol_type"] == const.chain_type_ids["NONPOLYMER"]
+        )
+
+        # Set additional features
+        features["str_gen"] = structure
+        features["path"] = path
+        features["id"] = path.stem
+        if "structure" in self.extra_features:
+            features["structure"] = structure
+        if "tokenized" in self.extra_features:
+            features["tokenized"] = tokenized
+
+        return features
+
+    def __len__(self) -> int:
+        return len(self.generated_paths) * self.multiplicity
+
+
+class FromGeneratedDataModule(pl.LightningDataModule):
+    def __init__(
+        self,
+        cfg: DataConfig,
+        return_native: bool = False,
+        compute_affinity: bool = False,
+        target_templates: bool = False,
+        design_mask_templates: bool = False,
+        skip_existing: bool = False,
+        skip_existing_kind: str = None,
+        legacy_gen_suffix: str = "_gen.cif",
+        legacy_metadata_suffix: str = "_metadata.npz",
+        reference_metadata_dir: Optional[Path] = None,
+        design_dir: Optional[str] = None,
+        extra_features: Optional[List[str]] = None,
+        design_mask_override: Optional[str] = None,
+        subset_target_ids: Optional[str] = None,
+        skip_specific_ids: Optional[List[str]] = None,
+        use_new_design_mask: bool = False,
+        fail_if_no_designs: bool = False,
+        output_dir: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.return_native = return_native
+        self.skip_existing = skip_existing
+        self.skip_existing_kind = skip_existing_kind
+        self.reference_metadata_dir = (
+            Path(reference_metadata_dir) if reference_metadata_dir else None
+        )
+
+        self.legacy_gen_suffix = legacy_gen_suffix
+        self.legacy_metadata_suffix = legacy_metadata_suffix
+        self.compute_affinity = compute_affinity
+        self.target_templates = target_templates
+        self.design_mask_templates = design_mask_templates
+        self.extra_features = extra_features
+        self.disulfide_prob = cfg.disulfide_prob
+        self.disulfide_on = cfg.disulfide_on
+        self.design_mask_override = cfg.design_mask_override
+        self.collate = collate
+        self.fail_if_no_designs = fail_if_no_designs
+        self.subset_target_ids = subset_target_ids
+        self.output_dir = Path(output_dir) if output_dir else None
+
+        if design_dir is not None:
+            self.init_dataset(
+                design_dir,
+                skip_specific_ids=skip_specific_ids,
+                extra_features=extra_features,
+                use_new_design_mask=use_new_design_mask,
+            )
+        else:
+            # Load canonical molecules
+            canonicals = load_canonicals(self.cfg.moldir)
+
+            self.predict_set = FromGeneratedDataset(
+                generated_paths=[],
+                metadata_paths=[],
+                native_paths=[],
+                canonicals=canonicals,
+                moldir=Path(self.cfg.moldir),
+                tokenizer=self.cfg.tokenizer,
+                featurizer=self.cfg.featurizer,
+                return_native=self.return_native,
+                reference_metadata_dir=self.reference_metadata_dir,
+                target_templates=self.target_templates,
+                design_mask_templates=self.design_mask_templates,
+                compute_affinity=self.compute_affinity,
+                design=self.cfg.design,
+                backbone_only=self.cfg.backbone_only,
+                atom14=self.cfg.atom14,
+                max_seqs=self.cfg.max_seqs,
+                inverse_fold=self.cfg.inverse_fold,
+                extra_features=self.extra_features,
+                disulfide_prob=self.disulfide_prob,
+                disulfide_on=self.disulfide_on,
+                design_mask_override=self.design_mask_override,
+                use_new_design_mask=use_new_design_mask,
+                multiplicity=self.cfg.multiplicity,
+                return_designfolding=self.cfg.return_designfolding,
+            )
+
+    def init_dataset(
+        self,
+        design_dir,
+        skip_specific_ids: Optional[List[str]] = None,
+        extra_features: Optional[List[str]] = None,
+        use_new_design_mask: bool = False,
+    ):
+        print(f"Initializing FromGeneratedDataModule datasets for {design_dir}")
+        design_dir = Path(design_dir)
+        assert design_dir.exists(), f"Path does not exist design_dir: {design_dir}"
+
+        # Aggregate generated structure files (.cif or .pdb) while skipping companion native/metadata files.
+        generated_paths = sorted(
+            p
+            for p in design_dir.iterdir()
+            if p.suffix in {".cif", ".pdb"}
+            and "_native.cif" not in p.name
+            and "_metadata.npz" not in p.name
+        )
+        if self.fail_if_no_designs and len(generated_paths) == 0:
+            raise ValueError(f"No designs found in {design_dir}")
+
+        # skip certain ids
+        num_files_before = len(generated_paths)
+        print(
+            f"[Info] Number of files to process (including already processed ones): {num_files_before}"
+        )
+
+        if skip_specific_ids:
+            filtered_generated_paths = [
+                p
+                for p in generated_paths
+                if not any(prob_id in p.name for prob_id in skip_specific_ids)
+            ]
+            num_files_after = len(filtered_generated_paths)
+            print(f"[Info] Skipped specific IDs: {skip_specific_ids}")
+            print(f"[Info] Number of files after filtering: {num_files_after}")
+            generated_paths = filtered_generated_paths
+
+        if self.skip_existing:
+            # Functions to map an input path to a list of output paths.
+            # If all output paths exist, the input path is skipped.
+            def output_path_inverse_fold(input_path):
+                assert self.output_dir is not None
+                return [
+                    self.output_dir / f"{input_path.stem}.cif",
+                    self.output_dir / f"{input_path.stem}.npz",
+                ]
+
+            def output_path_folded(input_path):
+                output_dir = (
+                    design_dir / const.folding_dirname
+                    if self.output_dir is None
+                    else self.output_dir
+                )
+                return [
+                    output_dir / f"{input_path.stem}.npz",
+                    output_dir / f"{input_path.stem}.npz",
+                ]
+
+            def output_path_design_folded(input_path):
+                output_dir = (
+                    design_dir / const.refold_design_cif_dirname
+                    if self.output_dir is None
+                    else self.output_dir
+                )
+                return [
+                    output_dir / f"{input_path.stem}.cif",
+                ]
+
+            def output_path_affinity(input_path):
+                output_dir = (
+                    design_dir / const.affinity_dirname
+                    if self.output_dir is None
+                    else self.output_dir
+                )
+                return [
+                    output_dir / f"{input_path.stem}.npz",
+                ]
+
+            def output_path_analyzed(input_path):
+                output_dir = (
+                    design_dir / const.metrics_dirname
+                    if self.output_dir is None
+                    else self.output_dir
+                )
+                return [
+                    output_dir / f"data_{input_path.stem}.npz",
+                    output_dir / f"metrics_{input_path.stem}.npz",
+                ]
+
+            mappings = {
+                "inverse_fold": output_path_inverse_fold,
+                "folded": output_path_folded,
+                "design_folded": output_path_design_folded,
+                "affinity": output_path_affinity,
+                "analyzed": output_path_analyzed,
+            }
+            if self.skip_existing_kind not in mappings:
+                msg = f"Invalid skip_existing_kind: {self.skip_existing_kind}. Available kinds: {list(mappings.keys())}"
+                raise ValueError(msg)
+            selected_mapping = mappings[self.skip_existing_kind]
+
+            generated_paths = [
+                p
+                for p in generated_paths
+                if not all(output_path.exists() for output_path in selected_mapping(p))
+            ]
+            msg = f"[Info] Skipped already {self.skip_existing_kind} IDs. Number of files after filtering: {len(generated_paths)}"
+            print(msg)
+
+        target_ids = [
+            re.search(rf"{self.cfg.target_id_regex}", p.stem).group(1)
+            for p in generated_paths
+        ]
+        target_ids = list(set(target_ids))
+
+        if self.cfg.num_targets is not None:
+            target_ids = target_ids[: self.cfg.num_targets]
+            generated_paths = [
+                p
+                for p in generated_paths
+                if re.search(rf"{self.cfg.target_id_regex}", p.stem).group(1)
+                in target_ids
+            ]
+
+        filtered_paths = []
+        for target_id in target_ids:
+            paths_of_target = [
+                p
+                for p in generated_paths
+                if re.search(rf"{self.cfg.target_id_regex}", p.stem).group(1)
+                == target_id
+            ]
+
+            filtered_paths.extend(paths_of_target[: self.cfg.samples_per_target])
+
+        filtered_paths2 = []
+        if self.subset_target_ids is not None:
+            subset_ids = [
+                l.strip() for l in open(self.subset_target_ids, "r").readlines()
+            ]
+            for path in filtered_paths:
+                if any([sid in str(path) for sid in subset_ids]):
+                    filtered_paths2.append(path)
+            filtered_paths = filtered_paths2
+
+        metadata_paths = []
+        native_paths = []
+        # Sort the paths to make sure each subprocess (when using multiple GPUs) has the same order and the index distribution when fetching from the dataset fetches the correct paths instead of fetching the same paths multiple times.
+        filtered_paths = sorted(filtered_paths)
+        for path in filtered_paths:
+            ext = path.suffix
+
+            # Legacy files contain "_gen" before the extension.
+            if path.stem.endswith("_gen"):
+                metadata_path = path.with_name(
+                    path.name.replace(f"_gen{ext}", "_metadata.npz")
+                )
+                native_path = path.with_name(
+                    path.name.replace(f"_gen{ext}", "_native.cif")
+                )
+            else:
+                metadata_path = path.with_suffix(".npz")
+                native_path = path.with_name(f"{path.stem}_native.cif")
+
+            if not metadata_path.exists():
+                print(f"[WARNING] Path does not exist: {metadata_path}")
+            metadata_paths.append(metadata_path)
+            if self.return_native:
+                if not native_path.exists():
+                    print(f"[WARNING] Path does not exist: {native_path}")
+                native_paths.append(native_path)
+            else:
+                native_paths.append(None)
+        msg = f"Found {len(target_ids)} targets and {len(filtered_paths)} remaining designs that still need to be processed in this step."
+        print(msg)
+
+        # Load canonical molecules
+        canonicals = load_canonicals(self.cfg.moldir)
+
+        self.predict_set = FromGeneratedDataset(
+            generated_paths=filtered_paths,
+            metadata_paths=metadata_paths,
+            native_paths=native_paths,
+            canonicals=canonicals,
+            moldir=Path(self.cfg.moldir),
+            tokenizer=self.cfg.tokenizer,
+            featurizer=self.cfg.featurizer,
+            return_native=self.return_native,
+            reference_metadata_dir=self.reference_metadata_dir,
+            target_templates=self.target_templates,
+            design_mask_templates=self.design_mask_templates,
+            compute_affinity=self.compute_affinity,
+            design=self.cfg.design,
+            backbone_only=self.cfg.backbone_only,
+            atom14=self.cfg.atom14,
+            max_seqs=self.cfg.max_seqs,
+            inverse_fold=self.cfg.inverse_fold,
+            extra_mol_dir=design_dir / const.molecules_dirname,
+            extra_features=self.extra_features,
+            disulfide_prob=self.disulfide_prob,
+            disulfide_on=self.disulfide_on,
+            design_mask_override=self.design_mask_override,
+            use_new_design_mask=use_new_design_mask,
+            multiplicity=self.cfg.multiplicity,
+            return_designfolding=self.cfg.return_designfolding,
+        )
+
+    def predict_dataloader(self) -> DataLoader:
+        return DataLoader(
+            self.predict_set,
+            batch_size=self.cfg.batch_size,
+            num_workers=self.cfg.num_workers,
+            pin_memory=self.cfg.pin_memory,
+            shuffle=False,
+            collate_fn=collate,
+        )
+
+    def transfer_batch_to_device(
+        self,
+        batch: Dict,
+        device: torch.device,
+        dataloader_idx: int = 0,
+    ) -> Dict:
+        for key in batch:
+            if key not in [
+                "metadata",
+                "str_gen",
+                "id",
+                "path",
+                "native_metadata",
+                "native_str_gen",
+                "native_id",
+                "native_path",
+                "exception",
+                "native_exception",
+                "skip",
+                "native_skip",
+                "structure_bonds",
+                "native_structure_bonds",
+                "extra_mols",
+                "native_extra_mols",
+                "structure",
+                "tokenized",
+                "data_sample_idx",
+            ]:
+                batch[key] = batch[key].to(device)
+
+        return batch

model/boltzgen/task/predict/data_from_yaml.py

@@ -0,0 +1,449 @@
+from dataclasses import dataclass
+from pathlib import Path
+import re
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from torch import Tensor
+from torch.utils.data import DataLoader
+from rdkit.Chem import Mol
+from boltzgen.data import const
+from boltzgen.data.data import Input
+from boltzgen.data.feature.featurizer import Featurizer
+from boltzgen.data.pad import pad_to_max
+from boltzgen.data.mol import load_canonicals, load_molecules
+from boltzgen.data.parse.schema import YamlDesignParser
+from boltzgen.data.template.features import load_dummy_templates
+from boltzgen.data.tokenize.tokenizer import Tokenizer
+from boltzgen.data.data import Input
+from boltzgen.data.select.protein import ProteinSelector
+
+
+@dataclass
+class DataConfig:
+    """Data configuration."""
+
+    moldir: str
+    multiplicity: int
+    yaml_path: Union[List[str], str]
+    tokenizer: Tokenizer
+    featurizer: Featurizer
+    backbone_only: bool = False
+    atom14: bool = False
+    atom37: bool = False
+    design: bool = True
+    compute_affinity: bool = False
+    disulfide_prob: float = 1.0
+    disulfide_on: bool = False
+    skip_existing: bool = False
+    skip_offset: int = 0
+    diffusion_samples: int = 1
+    output_dir: Optional[str] = None
+  
+   
+
+
+@dataclass
+class Dataset:
+    yaml_path: Union[List[str], str]
+    tokenizer: Tokenizer
+    featurizer: Featurizer
+    multiplicity: int = 1
+
+
+def collate(data: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
+    """Collate the data.
+
+    Parameters
+    ----------
+    data : List[Dict[str, Tensor]]
+        The data to collate.
+
+    Returns
+    -------
+    Dict[str, Tensor]
+        The collated data.
+
+    """
+    # Get the keys
+    keys = data[0].keys()
+
+    # Collate the data
+    collated = {}
+    for key in keys:
+        values = [d[key] for d in data]
+
+        if key not in [
+            "all_coords",
+            "all_resolved_mask",
+            "crop_to_all_atom_map",
+            "chain_symmetries",
+            "amino_acids_symmetries",
+            "ligand_symmetries",
+            "activity_name",
+            "activity_qualifier",
+            "sid",
+            "cid",
+            "aid",
+            "normalized_protein_accession",
+            "pair_id",
+            "record",
+            "id",
+            "structure",
+            "tokenized",
+            "structure_bonds",
+            "extra_mols",
+            "data_sample_idx",
+        ]:
+            # Check if all have the same shape
+            shape = values[0].shape
+            if not all(v.shape == shape for v in values):
+                values = pad_to_max(values, 0)
+            else:
+                values = torch.stack(values, dim=0)
+
+        # Stack the values
+        collated[key] = values
+
+    return collated
+
+
+class PredictionDataset(torch.utils.data.Dataset):
+    """Base iterable dataset."""
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        canonicals: dict[str, Mol],
+        moldir: str,
+        backbone_only: bool = False,
+        atom14: bool = False,
+        atom37: bool = False,
+        extra_features: Optional[List[str]] = None,
+        design: bool = True,
+        compute_affinity: bool = False,
+        disulfide_prob: float = 1.0,
+        disulfide_on: bool = False,
+        skip_offset: int = 0,
+    ) -> None:
+        """Initialize the training dataset.
+
+        Parameters
+        ----------
+        datasets : List[Dataset]
+            The datasets to sample from.
+
+        """
+        super().__init__()
+        self.dataset = dataset
+        self.moldir = moldir
+        self.canonicals = canonicals
+        self.backbone_only = backbone_only
+        self.atom14 = atom14
+        self.atom37 = atom37
+        self.skip_offset = skip_offset
+        path = dataset.yaml_path
+        self.yaml_paths = [path] if isinstance(path, str) else path
+
+        for path in self.yaml_paths:
+            filename = Path(path).name
+            if re.search(r"_\d+\.yaml$", filename):
+                raise ValueError(
+                    f"Illegal YAML filename for '{str(path)}': names must not end with the pattern _\\d+\\.yaml so, e.g., the ends '_001.yaml' or '_4.yaml' are not allowed."
+                    "This pattern is reserved for internal file indexing. Sorry :)"
+                )
+            if "_native" in filename:
+                raise ValueError(
+                    f"Illegal YAML filename for '{str(path)}': names must not contain '_native' because this substring is reserved for native structure companions. Sorry :)"
+                )
+        self.extra_features = (
+            set(extra_features) if extra_features is not None else set()
+        )
+        self.selector = (
+            ProteinSelector(  
+                design_neighborhood_sizes=[2, 4, 6, 8, 10, 12, 14, 16, 18],
+                substructure_neighborhood_sizes=[2, 4, 6, 8, 10, 12, 24],
+                structure_condition_prob=1.0,
+                distance_noise_std=1,
+                run_selection=True,
+                specify_binding_sites=True,
+                ss_condition_prob=0.1,
+                select_all=False,
+                chain_reindexing=False,
+            )
+        )
+        self.design = design
+        self.compute_affinity = compute_affinity
+        self.disulfide_prob = disulfide_prob
+        self.disulfide_on = disulfide_on
+
+        self.mols = {}
+        self.parser = YamlDesignParser(mol_dir=self.moldir)
+
+    def __getitem__(self, idx: int) -> Dict:
+        """Get an item from the dataset.
+
+        Returns
+        -------
+        Dict[str, Tensor]
+            The sampled data features.
+
+        """
+        path = Path(self.yaml_paths[idx % len(self.yaml_paths)])
+        feat = self.get_sample(path)
+        data_sample_idx = idx // len(self.yaml_paths) + self.skip_offset
+        if self.dataset.multiplicity > 1:
+            feat["data_sample_idx"] = data_sample_idx
+        return feat
+
+    def get_sample(self, path: Path, sample_id: Optional[str] = None) -> Dict:
+        # Get itemn also needs to take a smaple id as input
+        parsed = self.parser.parse_yaml(
+            path, mol_dir=self.moldir, mols=self.mols
+        )
+        structure = parsed.structure
+        design_info = parsed.design_info
+
+        # Tokenize structure
+        tokenized = self.dataset.tokenizer.tokenize(structure)
+
+        # Transfer conditioning information that is stored in tokens
+        token_to_res = tokenized.token_to_res
+        tokenized.tokens["design_mask"] = design_info.res_design_mask[token_to_res]
+        tokenized.tokens["binding_type"] = design_info.res_binding_type[token_to_res]
+        tokenized.tokens["structure_group"] = design_info.res_structure_groups[
+            token_to_res
+        ]
+        # Transfer per-residue amino acid constraints (shape: num_tokens x 20)
+        tokenized.tokens["aa_constraint_mask"] = design_info.res_aa_constraint_mask[
+            token_to_res
+        ]
+
+        # Propagate design mask to obtain chain_design_mask (True whenever something is covalently bound to any residue that is in a chain that contains a design residue).
+        chain_design_mask = tokenized.tokens["design_mask"].astype(bool)
+        asym_id = tokenized.tokens["asym_id"]
+        while True:
+            design_chains = np.unique(asym_id[chain_design_mask])
+            chain_propagated = np.isin(asym_id, design_chains)
+            for i, j, _ in tokenized.bonds:
+                if any([chain_propagated[i], chain_propagated[j]]):
+                    chain_propagated[i] = True
+                    chain_propagated[j] = True
+            if np.equal(chain_propagated, chain_design_mask).all():
+                break
+            chain_design_mask = chain_propagated.astype(bool)
+
+        # Try to find molecules in the dataset moldir if provided
+        # Find missing ones in global moldir and check if all found
+        molecules = {}
+        molecules.update(self.canonicals)
+        mol_names = set(tokenized.tokens["res_name"].tolist())
+        mol_names = mol_names - set(self.canonicals.keys())
+        mol_names = mol_names - set(parsed.extra_mols.keys())
+        if self.moldir is not None:
+            molecules.update(load_molecules(self.moldir, mol_names))
+
+        mol_names = mol_names - set(molecules.keys())
+        molecules.update(load_molecules(self.moldir, mol_names))
+        molecules.update(parsed.extra_mols)
+
+        # Finalize input data
+        input_data = Input(
+            tokens=tokenized.tokens,
+            bonds=tokenized.bonds,
+            token_to_res=token_to_res,
+            structure=structure,
+            msa={},
+            templates=None,
+        )
+
+        # Compute features
+        features = self.dataset.featurizer.process(
+            input_data,
+            molecules=molecules,
+            random=np.random.default_rng(None),
+            training=False,
+            max_seqs=1,
+            backbone_only=self.backbone_only,
+            atom14=self.atom14,
+            atom37=self.atom37,
+            design=self.design,
+            override_method="X-RAY DIFFRACTION",
+            compute_affinity=self.compute_affinity,
+            disulfide_prob=self.disulfide_prob,
+            disulfide_on=self.disulfide_on,
+        )
+
+        # transfer secondary structure conditioning
+        ss_type = design_info.res_ss_types[token_to_res]
+        features["ss_type"] = torch.from_numpy(ss_type).to(features["ss_type"])
+        features["design_ss_mask"][ss_type != const.ss_type_ids["UNSPECIFIED"]] = 1
+
+        # set chain_design_mask
+        features["chain_design_mask"] = torch.from_numpy(chain_design_mask)
+
+        # Compute template features
+        templates_features = load_dummy_templates(
+            tdim=1, num_tokens=len(features["res_type"])
+        )
+        features.update(templates_features)
+
+        # set last necessary features
+        features["idx_dataset"] = torch.tensor(1)
+
+        # If a smaple id is provided then this should be the sample id instead of path.stem
+        if sample_id is not None:
+            features["id"] = sample_id
+        else:
+            features["id"] = path.stem
+        if "structure" in self.extra_features:
+            features["structure"] = structure
+        if "tokenized" in self.extra_features:
+            features["tokenized"] = tokenized
+
+        return features
+
+    def __len__(self) -> int:
+        """Get the length of the dataset.
+
+        Returns
+        -------
+        int
+            The length of the dataset.
+
+        """
+        total = len(self.yaml_paths) * (self.dataset.multiplicity - self.skip_offset)
+        return max(total, 0)
+
+
+class FromYamlDataModule(pl.LightningDataModule):
+    """DataModule for BoltzGen."""
+
+    def __init__(
+        self, cfg: DataConfig, batch_size, num_workers, pin_memory, extra_features=None
+    ) -> None:
+        """Initialize the DataModule.
+
+        Parameters
+        ----------
+        config : DataConfig
+            The data configuration.
+
+        """
+        super().__init__()
+
+        if cfg.skip_existing and cfg.output_dir is not None:
+            design_dir = Path(cfg.output_dir)
+            max_idx: int = -1
+            if design_dir.exists():
+                pattern = re.compile(r"_(\d+)(?:\.[^.]+)$")
+                max_idx = max(
+                    (
+                        int(m.group(1))
+                        for fp in design_dir.iterdir()
+                        if fp.suffix in {".cif", ".pdb"}
+                        and not any(s in fp.name for s in ("_native.cif", "_metadata.npz"))
+                        for m in [pattern.search(fp.name)]
+                        if m
+                    ),
+                    default=-1,
+                )
+            n_samples = getattr(cfg, "diffusion_samples", 1)
+            cfg.skip_offset = (max_idx // max(n_samples, 1)) + 1 if max_idx >= 0 else 0
+            
+        self.cfg = cfg
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.pin_memory = pin_memory
+        self.collate = collate
+
+        dataset = Dataset(
+            yaml_path=cfg.yaml_path,
+            multiplicity=cfg.multiplicity,
+            tokenizer=cfg.tokenizer,
+            featurizer=cfg.featurizer,
+        )
+
+        # Load canonical molecules
+        canonicals = load_canonicals(cfg.moldir)
+
+        self.predict_set = PredictionDataset(
+            dataset=dataset,
+            canonicals=canonicals,
+            moldir=Path(cfg.moldir),
+            backbone_only=cfg.backbone_only,
+            atom14=cfg.atom14,
+            atom37=cfg.atom37,
+            extra_features=extra_features,
+            design=cfg.design,
+            compute_affinity=cfg.compute_affinity,
+            disulfide_prob=cfg.disulfide_prob,
+            disulfide_on=cfg.disulfide_on,
+            skip_offset=cfg.skip_offset,
+        )
+
+    def predict_dataloader(self) -> DataLoader:
+        """Get the training dataloader.
+
+        Returns
+        -------
+        DataLoader
+            The training dataloader.
+
+        """
+        return DataLoader(
+            self.predict_set,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            shuffle=False,
+            collate_fn=collate,
+        )
+
+    def transfer_batch_to_device(
+        self,
+        batch: Dict,
+        device: torch.device,
+        dataloader_idx: int = 0,  # noqa: ARG002
+    ) -> Dict:
+        """Transfer a batch to the given device.
+
+        Parameters
+        ----------
+        batch : Dict
+            The batch to transfer.
+        device : torch.device
+            The device to transfer to.
+
+        Returns
+        -------
+        np.Any
+            The transferred batch.
+
+        """
+        for key in batch:
+            if key not in [
+                "all_coords",
+                "all_resolved_mask",
+                "crop_to_all_atom_map",
+                "chain_symmetries",
+                "amino_acids_symmetries",
+                "ligand_symmetries",
+                "activity_name",
+                "activity_qualifier",
+                "sid",
+                "cid",
+                "normalized_protein_accession",
+                "pair_id",
+                "record",
+                "id",
+                "structure",
+                "tokenized",
+                "structure_bonds",
+                "extra_mols",
+                "data_sample_idx",
+            ]:
+                batch[key] = batch[key].to(device)
+        return batch

model/boltzgen/task/predict/data_ligands.py

@@ -0,0 +1,410 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List
+import re
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from torch import Tensor
+from torch.utils.data import DataLoader
+from rdkit.Chem import Mol
+from boltzgen.data.data import (
+    Input,
+    Structure,
+)
+from boltzgen.data.feature.featurizer import Featurizer
+from boltzgen.data.pad import pad_to_max
+from boltzgen.data.mol import (
+    load_canonicals,
+    load_molecules,
+)
+from boltzgen.data.template.features import load_dummy_templates
+from boltzgen.data.tokenize.tokenizer import Tokenizer
+from boltzgen.data.template.features import (
+    load_dummy_templates,
+)
+
+
+@dataclass
+class DataConfig:
+    """Data configuration."""
+
+    target_dir: str
+    msa_dir: str
+    moldir: str
+    min_len: int
+    max_len: int
+    target_ids: str
+    tokenizer: Tokenizer
+    featurizer: Featurizer
+    backbone_only: bool = False
+    atom14: bool = False
+    atom37: bool = False
+    design: bool = False
+    target_structure_condition: bool = False
+    multiplicity: int = 1
+    disulfide_prob: float = 1.0
+    disulfide_on: bool = False
+
+
+@dataclass
+class Dataset:
+    """Data holder."""
+
+    struct_dir: Path
+    record_dir: Path
+    target_ids: List[str]
+    min_len: int
+    max_len: int
+    tokenizer: Tokenizer
+    featurizer: Featurizer
+    multiplicity: int = 1
+
+
+def collate(data: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
+    """Collate the data.
+
+    Parameters
+    ----------
+    data : List[Dict[str, Tensor]]
+        The data to collate.
+
+    Returns
+    -------
+    Dict[str, Tensor]
+        The collated data.
+
+    """
+    # Get the keys
+    keys = data[0].keys()
+
+    # Collate the data
+    collated = {}
+    for key in keys:
+        values = [d[key] for d in data]
+
+        if key not in [
+            "all_coords",
+            "all_resolved_mask",
+            "crop_to_all_atom_map",
+            "chain_symmetries",
+            "amino_acids_symmetries",
+            "ligand_symmetries",
+            "activity_name",
+            "activity_qualifier",
+            "sid",
+            "cid",
+            "aid",
+            "normalized_protein_accession",
+            "pair_id",
+            "record",
+            "id",
+            "tokenized",
+            "structure",
+            "structure_bonds",
+            "extra_mols",
+            "data_sample_idx",
+        ]:
+            # Check if all have the same shape
+            shape = values[0].shape
+            if not all(v.shape == shape for v in values):
+                values = pad_to_max(values, 0)
+            else:
+                values = torch.stack(values, dim=0)
+
+        # Stack the values
+        collated[key] = values
+
+    return collated
+
+
+class PredictionDataset(torch.utils.data.Dataset):
+    """Base iterable dataset."""
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        canonicals: dict[str, Mol],
+        moldir: str,
+        backbone_only: bool = False,
+        atom14: bool = False,
+        atom37: bool = False,
+        design: bool = False,
+        target_structure_condition: bool = False,
+        disulfide_prob: float = 1.0,
+        disulfide_on: bool = False,
+    ) -> None:
+        """Initialize the training dataset.
+
+        Parameters
+        ----------
+        datasets : List[Dataset]
+            The datasets to sample from.
+
+        """
+        super().__init__()
+        self.dataset = dataset
+        self.moldir = moldir
+        self.canonicals = canonicals
+        self.backbone_only = backbone_only
+        self.atom14 = atom14
+        self.design = design
+        self.target_structure_condition = target_structure_condition
+        self.atom37 = atom37
+        self.disulfide_prob = disulfide_prob
+        self.disulfide_on = disulfide_on
+
+    def __getitem__(self, idx: int) -> Dict:
+        """Get an item from the dataset.
+
+        Returns
+        -------
+        Dict[str, Tensor]
+            The sampled data features.
+
+        """
+        # Get a sample from the dataset
+        data_sample_idx = idx // len(self.dataset.target_ids)
+        target_id = self.dataset.target_ids[idx % len(self.dataset.target_ids)]
+
+        # Get the structure
+        try:
+            target_id = target_id.lower()
+            ccd = target_id.split("_")[0]
+            pdb_id = target_id.split("_")[1].lower()
+            str_native = Structure.load(self.dataset.struct_dir / f"{pdb_id}.npz")
+            res_selection = str_native.residues["name"] == ccd.upper()
+            if res_selection.sum() > 1:
+                print(f"CCD {ccd} occurs multiple times in {pdb_id}. Using first.")
+                res_selection = np.arange(len(res_selection))[res_selection][[0]]
+            str_target = Structure.extract_residues(str_native, res_selection)
+
+            if self.dataset.min_len == self.dataset.max_len:
+                length = self.dataset.min_len
+            else:
+                length = np.random.random_integers(
+                    self.dataset.min_len, self.dataset.max_len
+                )
+            str_prot = Structure.empty_protein(seq_len=length)
+            structure = Structure.concatenate(str_prot, str_target)
+
+        except Exception as e:  # noqa: BLE001
+            print(f"Failed to load input for {target_id}  with error {e}. Skipping.")  # noqa: T201
+            return self.__getitem__(0)
+
+        # Check if the structure is valid
+        try:
+            tokenized = self.dataset.tokenizer.tokenize(structure)
+        except Exception as e:  # noqa: BLE001
+            print(f"Tokenizer failed on {target_id} with error {e}. Skipping.")  # noqa: T201
+            return self.__getitem__(0)
+
+        # Design mask logic
+        tokenized.tokens["design_mask"] = (
+            tokenized.tokens["asym_id"] == tokenized.tokens["asym_id"][0]
+        ).astype(tokenized.tokens["asym_id"].dtype)
+
+        # Propagate design mask to obtain chain_design_mask (True whenever something is covalently bound to any residue that is in a chain that contains a design residue).
+        chain_design_mask = tokenized.tokens["design_mask"].astype(bool)
+        asym_id = tokenized.tokens["asym_id"]
+        while True:
+            design_chains = np.unique(asym_id[chain_design_mask])
+            chain_propagated = np.isin(asym_id, design_chains)
+            for i, j, _ in tokenized.bonds:
+                if any([chain_propagated[i], chain_propagated[j]]):
+                    chain_propagated[i] = True
+                    chain_propagated[j] = True
+            if np.equal(chain_propagated, chain_design_mask).all():
+                break
+            chain_design_mask = chain_propagated.astype(bool)
+
+        try:
+            # Try to find molecules in the dataset moldir if provided
+            # Find missing ones in global moldir and check if all found
+            molecules = {}
+            molecules.update(self.canonicals)
+            mol_names = set(tokenized.tokens["res_name"].tolist())
+            mol_names = mol_names - set(self.canonicals.keys())
+            if self.moldir is not None:
+                molecules.update(load_molecules(self.moldir, mol_names))
+
+            mol_names = mol_names - set(molecules.keys())
+            molecules.update(load_molecules(self.moldir, mol_names))
+        except Exception as e:  # noqa: BLE001
+            print(f"Molecule loading failed for {target_id} with error {e}. Skipping.")
+            return self.__getitem__(0)
+
+        # Finalize input data
+        input_data = Input(
+            tokens=tokenized.tokens,
+            bonds=tokenized.bonds,
+            token_to_res=tokenized.token_to_res,
+            structure=structure,
+            msa={},
+            templates=None,
+        )
+
+        # Compute features
+        try:
+            features = self.dataset.featurizer.process(
+                input_data,
+                molecules=molecules,
+                random=np.random.default_rng(None),
+                training=False,
+                max_seqs=1,
+                backbone_only=self.backbone_only,
+                atom14=self.atom14,
+                atom37=self.atom37,
+                design=self.design,
+                override_method="X-RAY DIFFRACTION",
+                disulfide_prob=self.disulfide_prob,
+                disulfide_on=self.disulfide_on,
+            )
+        except Exception as e:  # noqa: BLE001
+            print(f"Featurizer failed on {target_id} with error {e}. Skipping.")  # noqa: T201
+            return self.__getitem__(0)
+
+        # set chain_design_mask
+        features["chain_design_mask"] = torch.from_numpy(chain_design_mask)
+
+        # Compute template features
+        templates_features = load_dummy_templates(
+            tdim=1, num_tokens=len(features["res_type"])
+        )
+        features.update(templates_features)
+
+        features["idx_dataset"] = torch.tensor(1)
+
+        def sanitize_filename(name):
+            return re.sub(r"[^A-Za-z0-9._-]", "_", name)
+
+        features["id"] = sanitize_filename(target_id)
+        if self.dataset.multiplicity > 1:
+            features["data_sample_idx"] = data_sample_idx
+        return features
+
+    def __len__(self) -> int:
+        """Get the length of the dataset.
+
+        Returns
+        -------
+        int
+            The length of the dataset.
+
+        """
+        return len(self.dataset.target_ids) * self.dataset.multiplicity
+
+
+class LigandBinderDataModule(pl.LightningDataModule):
+    """DataModule for BoltzGen."""
+
+    def __init__(self, cfg: DataConfig, batch_size, num_workers, pin_memory) -> None:
+        """Initialize the DataModule.
+
+        Parameters
+        ----------
+        config : DataConfig
+            The data configuration.
+
+        """
+        super().__init__()
+        self.cfg = cfg
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.pin_memory = pin_memory
+
+        with Path(cfg.target_ids).open("r") as f:
+            target_ids = [x for x in f.read().splitlines()]
+            print("split", target_ids)
+
+        dataset = Dataset(
+            struct_dir=Path(cfg.target_dir) / "structures",
+            record_dir=Path(cfg.target_dir) / "records",
+            target_ids=target_ids,
+            multiplicity=cfg.multiplicity,
+            min_len=cfg.min_len,
+            max_len=cfg.max_len,
+            tokenizer=cfg.tokenizer,
+            featurizer=cfg.featurizer,
+        )
+
+        # Load canonical molecules
+        canonicals = load_canonicals(cfg.moldir)
+
+        self.predict_set = PredictionDataset(
+            dataset=dataset,
+            canonicals=canonicals,
+            moldir=Path(cfg.moldir),
+            backbone_only=cfg.backbone_only,
+            atom14=cfg.atom14,
+            design=cfg.design,
+            target_structure_condition=cfg.target_structure_condition,
+            disulfide_prob=cfg.disulfide_prob,
+            disulfide_on=cfg.disulfide_on,
+        )
+
+    def predict_dataloader(self) -> DataLoader:
+        """Get the training dataloader.
+
+        Returns
+        -------
+        DataLoader
+            The training dataloader.
+
+        """
+        return DataLoader(
+            self.predict_set,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            shuffle=False,
+            collate_fn=collate,
+        )
+
+    def transfer_batch_to_device(
+        self,
+        batch: Dict,
+        device: torch.device,
+        dataloader_idx: int,  # noqa: ARG002
+    ) -> Dict:
+        """Transfer a batch to the given device.
+
+        Parameters
+        ----------
+        batch : Dict
+            The batch to transfer.
+        device : torch.device
+            The device to transfer to.
+        dataloader_idx : int
+            The dataloader index.
+
+        Returns
+        -------
+        np.Any
+            The transferred batch.
+
+        """
+        for key in batch:
+            if key not in [
+                "all_coords",
+                "all_resolved_mask",
+                "crop_to_all_atom_map",
+                "chain_symmetries",
+                "amino_acids_symmetries",
+                "ligand_symmetries",
+                "activity_name",
+                "activity_qualifier",
+                "sid",
+                "cid",
+                "normalized_protein_accession",
+                "pair_id",
+                "record",
+                "id",
+                "tokenized",
+                "structure",
+                "structure_bonds",
+                "extra_mols",
+                "data_sample_idx",
+            ]:
+                batch[key] = batch[key].to(device)
+        return batch

model/boltzgen/task/predict/data_protein_binder.py

@@ -0,0 +1,608 @@
+from dataclasses import dataclass
+from boltzgen.data.data import StructureInfo
+from pathlib import Path
+import re
+from typing import Dict, List, Optional
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from torch import Tensor
+from torch.utils.data import DataLoader
+from rdkit.Chem import Mol
+from boltzgen.data import const
+from boltzgen.data.data import Input, Structure, Record, MSA
+from boltzgen.data.feature.featurizer import Featurizer
+from boltzgen.data.pad import pad_to_max
+from boltzgen.data.mol import load_canonicals, load_molecules
+from boltzgen.data.tokenize.tokenizer import Tokenizer
+from boltzgen.data.template.features import (
+    load_dummy_templates,
+)
+from boltzgen.task.predict.loading_utils import load_record, load_structure
+
+
+@dataclass
+class DataConfig:
+    """Data configuration."""
+
+    target_dir: str
+    msa_dir: str
+    moldir: str
+    seq_len: int
+    target_ids: str
+    tokenizer: Tokenizer
+    featurizer: Featurizer
+    backbone_only: bool = False
+    atom14: bool = False
+    atom37: bool = False
+    design: bool = False
+    target_structure_condition: bool = False
+    msa_condition: bool = False
+    max_seq: int = 1024
+    ss_single_condition: bool = False
+    ss_double_condition: bool = False
+    ss_short_motif: bool = False
+    multiplicity: int = 1
+    inverse_fold: bool = False
+    num_targets: int = 100000
+    disulfide_prob: float = 1.0
+    disulfide_on: bool = False
+
+
+@dataclass
+class Dataset:
+    """Data holder."""
+
+    struct_dir: Path
+    record_dir: Path
+    target_ids: List[str]
+    seq_len: int
+    tokenizer: Tokenizer
+    featurizer: Featurizer
+    chain_ids: List[str] = None
+
+
+def ss_single(feats, random):
+    feats["design_ss_mask"].zero_()
+    feats["ss_type"].zero_()
+    design_mask = feats["design_mask"].bool()
+    seq_len = design_mask.sum().item()
+
+    feats["design_ss_mask"][design_mask] = 0
+    design_indices = torch.where(design_mask)[0]
+
+    motif_length = random.integers(5, min(30, seq_len) + 1)
+    start_pos = random.choice(design_indices[: seq_len - motif_length + 1].cpu())
+
+    motif_type = random.choice([2, 3])  # helix or sheet
+    feats["design_ss_mask"][start_pos : start_pos + motif_length] = 1
+    feats["ss_type"][start_pos : start_pos + motif_length] = motif_type
+
+
+def ss_double(feats, random):
+    feats["design_ss_mask"].zero_()
+    feats["ss_type"].zero_()
+    design_mask = feats["design_mask"].bool()
+    seq_len = design_mask.sum().item()
+
+    feats["design_ss_mask"][design_mask] = 0
+    design_indices = torch.where(design_mask)[0]
+
+    motif_len1 = random.integers(5, min(30, seq_len - 10) + 1)
+    motif_len2 = random.integers(5, min(30, seq_len - motif_len1 - 5) + 1)
+    total_required = motif_len1 + 5 + motif_len2
+    start_pos = random.choice(design_indices[: seq_len - total_required + 1].cpu())
+
+    motif_type1 = random.choice([2, 3])
+    motif_type2 = random.choice([2, 3])
+
+    # first motif
+    feats["design_ss_mask"][start_pos : start_pos + motif_len1] = 1
+    feats["ss_type"][start_pos : start_pos + motif_len1] = motif_type1
+
+    # loop
+    loop_start = start_pos + motif_len1
+    feats["design_ss_mask"][loop_start : loop_start + 5] = 1
+    feats["ss_type"][loop_start : loop_start + 5] = 1  # loop
+
+    # second motif
+    second_start = loop_start + 5
+    feats["design_ss_mask"][second_start : second_start + motif_len2] = 1
+    feats["ss_type"][second_start : second_start + motif_len2] = motif_type2
+
+
+def ss_short_motif(feats, random, helix_prob=0.5):
+    feats["design_ss_mask"].zero_()
+    feats["ss_type"].zero_()
+    design_mask = feats["design_mask"].bool()
+    seq_len = design_mask.sum().item()
+    design_indices = torch.where(design_mask)[0]
+
+    start_pos = random.choice(design_indices[: seq_len - 4 + 1].cpu())
+    motif_type = 2 if random.random() < helix_prob else 3
+
+    feats["design_ss_mask"][start_pos : start_pos + 4] = 1
+    feats["ss_type"][start_pos : start_pos + 4] = motif_type
+
+
+def load_msas(record: Record, chain_ids: set[int], msa_dir: Path) -> Dict[int, MSA]:
+    """Load the given input data.
+
+    Parameters
+    ----------
+    record : Record
+        The record to load.
+    chain_ids : set[int]
+        The chain ids to load.
+    msa_dir : Path
+        The path to the MSA directory.
+
+    Returns
+    -------
+    Input
+        The loaded input.
+
+    """
+    # Load the relevant MSAs
+    msas = {}
+    for chain in record.chains:
+        if chain.chain_id in chain_ids:
+            msa_id = chain.msa_id
+            if msa_id != -1:
+                msa = np.load(msa_dir / f"{msa_id}.npz")
+                msas[chain.chain_id] = MSA(**msa)
+    return msas
+
+
+def collate(data: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
+    """Collate the data.
+
+    Parameters
+    ----------
+    data : List[Dict[str, Tensor]]
+        The data to collate.
+
+    Returns
+    -------
+    Dict[str, Tensor]
+        The collated data.
+
+    """
+    # Get the keys
+    keys = data[0].keys()
+
+    # Collate the data
+    collated = {}
+    for key in keys:
+        values = [d[key] for d in data]
+
+        if key not in [
+            "all_coords",
+            "all_resolved_mask",
+            "crop_to_all_atom_map",
+            "chain_symmetries",
+            "amino_acids_symmetries",
+            "ligand_symmetries",
+            "activity_name",
+            "activity_qualifier",
+            "sid",
+            "cid",
+            "aid",
+            "normalized_protein_accession",
+            "pair_id",
+            "record",
+            "id",
+            "structure_bonds",
+            "extra_mols",
+            "structure",
+            "tokenized",
+            "data_sample_idx",
+        ]:
+            # Check if all have the same shape
+            shape = values[0].shape
+            if not all(v.shape == shape for v in values):
+                values = pad_to_max(values, 0)
+            else:
+                values = torch.stack(values, dim=0)
+
+        # Stack the values
+        collated[key] = values
+
+    return collated
+
+
+class PredictionDataset(torch.utils.data.Dataset):
+    """Base iterable dataset."""
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        canonicals: dict[str, Mol],
+        moldir: str,
+        backbone_only: bool = False,
+        atom14: bool = False,
+        atom37: bool = False,
+        design: bool = False,
+        target_structure_condition: bool = False,
+        msa_condition: bool = False,
+        ss_single_condition: bool = False,
+        ss_double_condition: bool = False,
+        ss_short_motif: bool = False,
+        max_seq: int = 1024,
+        msa_dir: str = "/data/rbg/shared/projects/foldeverything/rcsb/msa",
+        inverse_fold: bool = False,
+        multiplicity: int = 1,
+        disulfide_prob: float = 1.0,
+        disulfide_on: bool = False,
+        extra_features: Optional[List[str]] = None,
+    ) -> None:
+        """Initialize the training dataset.
+
+        Parameters
+        ----------
+        datasets : List[Dataset]
+            The datasets to sample from.
+
+        """
+        super().__init__()
+        self.dataset = dataset
+        self.moldir = moldir
+        self.canonicals = canonicals
+        self.backbone_only = backbone_only
+        self.atom14 = atom14
+        self.atom37 = atom37
+        self.design = design
+        self.target_structure_condition = target_structure_condition
+        self.msa_condition = msa_condition
+        self.max_seq = max_seq
+        self.msa_dir = msa_dir
+        self.inverse_fold = inverse_fold
+        self.ss_single_condition = ss_single_condition
+        self.ss_double_condition = ss_double_condition
+        self.ss_short_motif = ss_short_motif
+        self.multiplicity = multiplicity
+        self.disulfide_prob = disulfide_prob
+        self.disulfide_on = disulfide_on
+        self.extra_features = (
+            set(extra_features) if extra_features is not None else set()
+        )
+
+    def __getitem__(self, idx: int) -> Dict:
+        """Get an item from the dataset.
+
+        Returns
+        -------
+        Dict[str, Tensor]
+            The sampled data features.
+
+        """
+        # Get a sample from the dataset
+        data_sample_idx = idx // len(self.dataset.target_ids)
+        pdb_id = self.dataset.target_ids[idx % len(self.dataset.target_ids)]
+        if not self.dataset.chain_ids is None:
+            chain_id = self.dataset.chain_ids[idx % len(self.dataset.target_ids)]
+        else:
+            chain_id = None
+
+        # Load record
+        record = load_record(pdb_id, self.dataset.record_dir)
+        
+        # Get the structure
+        try:
+            if self.inverse_fold:
+                structure = load_structure(record, self.dataset.struct_dir)
+            else:
+                str_native = load_structure(record, self.dataset.struct_dir)
+
+                # If chain id is specified, extract chain
+                if not chain_id is None:
+                    chain_id = chain_id.lower()
+                    chain = None
+                    chain_names = []
+                    for _chain in str_native.chains:
+                        chain_names.append(_chain[0].lower())
+                        if _chain[0].lower() == chain_id:
+                            chain = _chain
+                    if chain is None:
+                        chain_names = "\n".join(chain_names)
+                        msg = f"Could not find chain {chain_id}. Structure contains chains:\n{chain_names}"
+                        e = ValueError(msg)
+                    res_idx, res_num = chain[7], chain[8]
+                    res_idxs = np.arange(res_idx, res_idx + res_num)
+                    str_native = str_native.extract_residues(str_native, res_idxs)
+
+                str_prot = Structure.empty_protein(seq_len=self.dataset.seq_len)
+                structure = Structure.concatenate(str_native, str_prot)
+        except Exception as e:  # noqa: BLE001
+            print(f"Failed to load input for {pdb_id} with error {e}. Skipping.")  # noqa: T201
+            return self.__getitem__(0)
+
+        # Tokenize structure
+        try:
+            tokenized = self.dataset.tokenizer.tokenize(structure)
+        except Exception as e:  # noqa: BLE001
+            print(f"Tokenizer failed on {pdb_id} with error {e}. Skipping.")  # noqa: T201
+            return self.__getitem__(0)
+
+        if self.inverse_fold:
+            tokenized.tokens["design_mask"] = (
+                tokenized.tokens["is_standard"]
+                & (tokenized.tokens["mol_type"] == const.chain_type_ids["PROTEIN"])
+                & tokenized.tokens["resolved_mask"]
+            )
+            tokenized.tokens["structure_group"] = 1
+        else:
+            # Set the last chain to be designed
+            tokenized.tokens["design_mask"] = (
+                tokenized.tokens["asym_id"] == tokenized.tokens["asym_id"][-1]
+            ).astype(tokenized.tokens["asym_id"].dtype)
+
+            if self.target_structure_condition:
+                tokenized.tokens["structure_group"][
+                    ~tokenized.tokens["design_mask"].astype(bool)
+                ] = 1
+
+        # Propagate design mask to obtain chain_design_mask (True whenever something is covalently bound to any residue that is in a chain that contains a design residue).
+        chain_design_mask = tokenized.tokens["design_mask"].astype(bool)
+        asym_id = tokenized.tokens["asym_id"]
+        while True:
+            design_chains = np.unique(asym_id[chain_design_mask])
+            chain_propagated = np.isin(asym_id, design_chains)
+            for i, j, _ in tokenized.bonds:
+                if any([chain_propagated[i], chain_propagated[j]]):
+                    chain_propagated[i] = True
+                    chain_propagated[j] = True
+            if np.equal(chain_propagated, chain_design_mask).all():
+                break
+            chain_design_mask = chain_propagated.astype(bool)
+
+        # Find the record with the matching pdb_id
+        msas = {}
+        if self.msa_condition:
+            chain_ids = set(tokenized.tokens["asym_id"])
+            msas = load_msas(record=record, chain_ids=chain_ids, msa_dir=self.msa_dir)
+
+        try:
+            # Try to find molecules in the dataset moldir if provided
+            # Find missing ones in global moldir and check if all found
+            molecules = {}
+            molecules.update(self.canonicals)
+            mol_names = set(tokenized.tokens["res_name"].tolist())
+            mol_names = mol_names - set(self.canonicals.keys())
+            if self.moldir is not None:
+                molecules.update(load_molecules(self.moldir, mol_names))
+
+            mol_names = mol_names - set(molecules.keys())
+            molecules.update(load_molecules(self.moldir, mol_names))
+        except Exception as e:  # noqa: BLE001
+            print(f"Molecule loading failed for {record.id} with error {e}. Skipping.")
+            return self.__getitem__(0)
+
+        # Finalize input data
+        input_data = Input(
+            tokens=tokenized.tokens,
+            bonds=tokenized.bonds,
+            token_to_res=tokenized.token_to_res,
+            structure=structure,
+            msa=msas,
+            templates=None,
+            record=record,
+        )
+        # Compute features
+        try:
+            features = self.dataset.featurizer.process(
+                input_data,
+                molecules=molecules,
+                random=np.random.default_rng(None),
+                training=False,
+                max_seqs=self.max_seq,  # if self.msa_condition else 1,
+                backbone_only=self.backbone_only,
+                atom14=self.atom14,
+                atom37=self.atom37,
+                design=self.design,
+                pad_to_max_seqs=True,
+                override_method="X-RAY DIFFRACTION",
+                disulfide_prob=self.disulfide_prob,
+                disulfide_on=self.disulfide_on,
+            )
+
+        except Exception as e:  # noqa: BLE001
+            print(f"Featurizer failed on {pdb_id} with error {e}. Skipping.")  # noqa: T201
+            return self.__getitem__(0)
+
+        # set chain_design_mask
+        features["chain_design_mask"] = torch.from_numpy(chain_design_mask)
+
+        # Compute template features
+        templates_features = load_dummy_templates(
+            tdim=1, num_tokens=len(features["res_type"])
+        )
+        features.update(templates_features)
+
+        rng = np.random.default_rng(None)
+        if self.ss_single_condition:
+            ss_single(features, rng)
+        if self.ss_double_condition:
+            ss_double(features, rng)
+        if self.ss_short_motif:
+            ss_short_motif(features, rng, helix_prob=0.8)
+
+        features["idx_dataset"] = torch.tensor(1)
+        features["id"] = pdb_id
+        if "structure" in self.extra_features:
+            features["structure"] = structure
+        if "tokenized" in self.extra_features:
+            features["tokenized"] = tokenized
+        if self.multiplicity > 1:
+            features["data_sample_idx"] = data_sample_idx
+        return features
+
+    def __len__(self) -> int:
+        """Get the length of the dataset.
+
+        Returns
+        -------
+        int
+            The length of the dataset.
+
+        """
+        return len(self.dataset.target_ids) * self.multiplicity
+
+
+class ProteinBinderDataModule(pl.LightningDataModule):
+    """DataModule for BoltzGen."""
+
+    def __init__(
+        self,
+        cfg: DataConfig,
+        batch_size,
+        num_workers,
+        pin_memory,
+        extra_features: Optional[List[str]] = None,
+    ) -> None:
+        """Initialize the DataModule.
+
+        Parameters
+        ----------
+        config : DataConfig
+            The data configuration.
+
+        """
+        super().__init__()
+        self.cfg = cfg
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.pin_memory = pin_memory
+
+        with Path(cfg.target_ids).open("r") as f:
+            regex = re.compile(r"([^_]+)(?:_(.*))?")
+            target_ids = [
+                regex.match(x.lower()).groups() for x in f.read().splitlines()
+            ]
+            chain_ids = [tup[1] for tup in target_ids]
+            target_ids = [tup[0] for tup in target_ids]
+            if cfg.num_targets is not None:
+                target_ids = target_ids[: cfg.num_targets]
+            print("split", target_ids)
+
+        dataset = Dataset(
+            struct_dir=Path(cfg.target_dir) / "structures",
+            record_dir=Path(cfg.target_dir) / "records",
+            target_ids=target_ids,
+            seq_len=cfg.seq_len,
+            tokenizer=cfg.tokenizer,
+            featurizer=cfg.featurizer,
+            chain_ids=chain_ids,
+        )
+
+        # Load canonical molecules
+        canonicals = load_canonicals(cfg.moldir)
+
+        self.predict_set = PredictionDataset(
+            dataset=dataset,
+            canonicals=canonicals,
+            moldir=Path(cfg.moldir),
+            backbone_only=cfg.backbone_only,
+            atom14=cfg.atom14,
+            atom37=cfg.atom37,
+            design=cfg.design,
+            target_structure_condition=cfg.target_structure_condition,
+            msa_condition=cfg.msa_condition,
+            ss_single_condition=cfg.ss_single_condition,
+            ss_double_condition=cfg.ss_double_condition,
+            ss_short_motif=cfg.ss_short_motif,
+            max_seq=cfg.seq_len,
+            msa_dir=Path(cfg.msa_dir),
+            multiplicity=cfg.multiplicity,
+            inverse_fold=cfg.inverse_fold,
+            disulfide_prob=cfg.disulfide_prob,
+            disulfide_on=cfg.disulfide_on,
+            extra_features=extra_features,
+        )
+
+    def predict_dataloader(self) -> DataLoader:
+        """Get the training dataloader.
+
+        Returns
+        -------
+        DataLoader
+            The training dataloader.
+
+        """
+        return DataLoader(
+            self.predict_set,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            shuffle=False,
+            collate_fn=collate,
+        )
+
+    def transfer_batch_to_device(
+        self,
+        batch: Dict,
+        device: torch.device,
+        dataloader_idx: int,  # noqa: ARG002
+    ) -> Dict:
+        """Transfer a batch to the given device.
+
+        Parameters
+        ----------
+        batch : Dict
+            The batch to transfer.
+        device : torch.device
+            The device to transfer to.
+        dataloader_idx : int
+            The dataloader index.
+
+        Returns
+        -------
+        np.Any
+            The transferred batch.
+
+        """
+        for key in batch:
+            if key not in [
+                "all_coords",
+                "all_resolved_mask",
+                "crop_to_all_atom_map",
+                "chain_symmetries",
+                "amino_acids_symmetries",
+                "ligand_symmetries",
+                "activity_name",
+                "activity_qualifier",
+                "sid",
+                "cid",
+                "normalized_protein_accession",
+                "pair_id",
+                "record",
+                "id",
+                "structure",
+                "tokenized",
+                "structure_bonds",
+                "extra_mols",
+                "data_sample_idx",
+            ]:
+                batch[key] = batch[key].to(device)
+        return batch
+
+
+if __name__ == "__main__":
+    # debugging code
+    from omegaconf import OmegaConf
+    from hydra.utils import instantiate
+    import tqdm
+
+    cfg = OmegaConf.load("configs/predict/prot.yaml")
+    # cfg = OmegaConf.load("configs/predict/lig.yaml")
+    # datamodule = instantiate(cfg.data)
+
+    cfg.data.cfg.target_ids = "/data/scratch/faltings/data/bgen/hard_test_set_ids.txt"
+    datamodule = instantiate(cfg.data)
+    for i in tqdm.tqdm(range(len(datamodule.predict_set))):
+        try:
+            entry = datamodule.predict_set[i]
+        except Exception as e:
+            print(e)
+            continue

model/boltzgen/task/predict/loading_utils.py

@@ -0,0 +1,45 @@
+from pathlib import Path
+from boltzgen.data.data import Record, Structure
+
+
+
+def load_record(record_id: str, record_dir: Path) -> Record:
+    """Load the given record.
+
+    Parameters
+    ----------
+    record_id : str
+        The record id to load.
+    record_dir : Path
+        The path to the record directory.
+
+    Returns
+    -------
+    Record
+        The loaded record.
+    """
+    return Record.load(record_dir / f"{record_id}.json")
+
+
+def load_structure(record: Record, struct_dir: Path) -> Structure:
+    """Load the given input data.
+
+    Parameters
+    ----------
+    record : str
+        The record to load.
+    target_dir : Path
+        The path to the data directory.
+
+    Returns
+    -------
+    Input
+        The loaded input.
+
+    """
+    if (struct_dir / f"{record.id}.npz").exists():
+        structure_path = struct_dir / f"{record.id}.npz"
+    else:
+        structure_path = struct_dir / f"{record.id}" / f"{record.id}_model_0.npz"
+    return Structure.load(structure_path)
+

model/boltzgen/task/predict/predict.py

@@ -0,0 +1,194 @@
+from boltzgen.utils.quiet import quiet_startup
+
+quiet_startup()
+
+import os
+
+# Disable Triton auto-tuning during inference
+os.environ.setdefault("CUEQ_DEFAULT_CONFIG", "1")
+os.environ.setdefault("CUEQ_DISABLE_AOT_TUNING", "1")
+
+from typing import List, Optional, Union
+
+import torch
+from omegaconf import OmegaConf, listconfig
+from pytorch_lightning import LightningModule, Trainer
+
+from pytorch_lightning.strategies import DDPStrategy
+
+from boltzgen.task.predict.data_from_generated import FromGeneratedDataModule
+from boltzgen.task.predict.writer import (
+    DesignWriter,
+    FoldingWriter,
+)
+from boltzgen.task.task import Task
+from boltzgen.utils.pipeline_progress_bar import PipelineProgressBar
+from boltzgen.model.models.boltz import Boltz
+
+
+class Predict(Task):
+    """A task to run model inference."""
+
+    def __init__(
+        self,
+        data: Union[FromGeneratedDataModule],
+        writer: Union[DesignWriter, FoldingWriter],
+        checkpoint: str,
+        output: str,
+        name: str,
+        recycling_steps: int,
+        sampling_steps: int,
+        diffusion_samples: int = 1,
+        keys_dict_out: Optional[List] = None,
+        keys_dict_batch: Optional[List] = None,
+        slurm: bool = False,
+        matmul_precision: Optional[str] = None,
+        trainer: Optional[dict] = None,
+        override: Optional[dict] = None,
+        debug: bool = False,
+        use_ema: bool = False,
+        write_manifest: bool = False,
+        compile_pairformer: bool = False,
+        compile_structure: bool = False,
+        checkpoint_diffusion_conditioning: bool = False,
+    ) -> None:
+        """Initialize the task.
+
+        Parameters
+        ----------
+        checkpoint : str
+            The path to the model checkpoint.
+        output : str
+            The path to save the inference results.
+        slurm : bool, optional
+            Whether to run on SLURM, by default False
+        matmul_precision : Optional[str], optional
+            The matmul precision, by default None
+        trainer : Optional[dict], optional
+            The configuration for the trainer, by default None
+        override : Optional[dict], optional
+            The override configuration for the model, by default None
+
+        """
+        self.data = data
+        self.checkpoint = checkpoint
+        self.output = output
+        self.slurm = slurm
+        self.matmul_precision = matmul_precision
+        self.trainer = trainer
+        self.override = override if override is not None else {}
+        self.predict_args = {
+            "recycling_steps": recycling_steps,
+            "sampling_steps": sampling_steps,
+            "diffusion_samples": diffusion_samples,
+        }
+        if keys_dict_batch is not None:
+            self.predict_args["keys_dict_batch"] = keys_dict_batch
+        if keys_dict_out is not None:
+            self.predict_args["keys_dict_out"] = keys_dict_out
+        self.debug = debug
+        self.use_ema = use_ema
+        self.write_manifest = write_manifest
+        self.writer = writer
+        self.compile_pairformer = compile_pairformer
+        self.compile_structure = compile_structure
+        self.checkpoint_diffusion_conditioning = checkpoint_diffusion_conditioning
+
+    def run(self, config: OmegaConf = None, run_prediction=True) -> None:  # noqa: ARG002
+        # Silence warnings and pytorch lightning tips
+        quiet_startup()
+
+        # Exit quickly if no predictions are needed
+        if len(self.data.predict_set) == 0:
+            print("No predictions required")
+            return
+
+        # Set no grad
+        torch.set_grad_enabled(False)
+
+        # Experiment with this during training (high or medium)
+        if self.matmul_precision is not None:
+            torch.set_float32_matmul_precision(self.matmul_precision)
+
+        # Create trainer dict
+        if self.trainer is None:
+            self.trainer = {}
+
+        # Flip some arguments in debug mode
+        devices = self.trainer.get("devices", 1)
+
+        if self.debug:
+            if isinstance(devices, int):
+                devices = 1
+            elif isinstance(devices, (list, listconfig.ListConfig)):
+                devices = [devices[0]]
+            self.trainer["devices"] = devices
+            self.data.num_workers = 0
+
+        # slurm
+        if self.slurm:
+            self.trainer["devices"] = int(
+                os.environ.get("SLURM_NTASKS_PER_NODE", "auto")
+            )
+            self.trainer["num_nodes"] = int(os.environ.get("SLURM_NNODES", 1))
+
+        # Load model
+        self.model_module: LightningModule = Boltz.load_from_checkpoint(
+            self.checkpoint,
+            strict=True,
+            use_ema=self.use_ema,
+            checkpoint_diffusion_conditioning=self.checkpoint_diffusion_conditioning,
+            map_location="cpu",
+            weights_only=False,
+            predict_args=self.predict_args,
+            **self.override,
+        )
+        self.model_module.eval()
+
+        if self.compile_pairformer:
+            self.model_module.is_pairformer_compiled = True
+            self.model_module.pairformer_module = torch.compile(
+                self.model_module.pairformer_module, dynamic=True, fullgraph=False
+            )
+        if self.compile_structure:
+            self.model_module.structure_module.score_model.is_token_transformer_compiled = True
+            self.model_module.structure_module.score_model.token_transformer = (
+                torch.compile(
+                    self.model_module.structure_module.score_model.token_transformer,
+                    dynamic=True,
+                    fullgraph=False,
+                )
+            )
+
+        # Set up trainer
+        strategy = "auto"
+        num_devices = (
+            len(devices)
+            if isinstance(devices, (list, listconfig.ListConfig))
+            else devices
+        )
+        if num_devices > 1:
+            strategy = DDPStrategy()
+            if num_devices > len(self.data.predict_set):
+                devices = max(1, len(self.data.predict_set))
+                msg = f"Fewer designs than devices. Setting devices to {devices}."
+                print(msg)
+                self.trainer["devices"] = devices
+
+        self.lightning_trainer = Trainer(
+            default_root_dir=self.output,
+            strategy=strategy,
+            callbacks=[self.writer]
+            + (
+                [PipelineProgressBar()]
+                if os.environ.get("BOLTZGEN_PIPELINE_STEP")
+                else []
+            ),
+            **self.trainer,
+        )
+        if run_prediction:
+            # Run training
+            self.lightning_trainer.predict(
+                self.model_module, datamodule=self.data, return_predictions=False
+            )
+            del self.model_module

model/boltzgen/task/predict/writer.py

@@ -0,0 +1,534 @@
+import pickle
+from pathlib import Path
+from typing import Dict, List
+
+import numpy as np
+import torch
+from pytorch_lightning import LightningModule, Trainer
+from pytorch_lightning.callbacks import BasePredictionWriter
+from torch import Tensor
+from tqdm import tqdm
+
+from boltzgen.data import const
+from boltzgen.data.data import (
+    Structure,
+    convert_ccd,
+)
+from boltzgen.data.feature.featurizer import (
+    res_from_atom14,
+    res_from_atom37,
+    res_all_gly,
+)
+from boltzgen.data.write.mmcif import to_mmcif
+from boltzgen.data.write.pdb import to_pdb
+from boltzgen.model.loss.diffusion import weighted_rigid_align
+from boltzgen.model.modules.masker import BoltzMasker
+
+
+class FoldingWriter(BasePredictionWriter):
+    """Custom writer for predictions."""
+
+    def __init__(self, design_dir: str, designfolding: bool = False) -> None:
+        super().__init__(write_interval="batch")
+        self.designfolding = designfolding
+        if design_dir is not None:
+            self.init_outdir(design_dir)
+
+    def init_outdir(self, design_dir):
+        self.outdir = Path(design_dir) / (
+            const.folding_design_dirname
+            if self.designfolding
+            else const.folding_dirname
+        )
+        self.refold_cif_dir = Path(design_dir) / (
+            const.refold_design_cif_dirname
+            if self.designfolding
+            else const.refold_cif_dirname
+        )
+        self.refold_cif_dir.mkdir(parents=True, exist_ok=True)
+        self.outdir.mkdir(exist_ok=True, parents=True)
+        self.failed = 0
+
+    def write_on_batch_end(  # noqa: PLR0915
+        self,
+        trainer: Trainer = None,  # noqa: ARG002
+        pl_module: LightningModule = None,  # noqa: ARG002
+        prediction: Dict[str, Tensor] = None,
+        batch_indices: List[int] = None,  # noqa: ARG002
+        batch: Dict[str, Tensor] = None,
+        batch_idx: int = None,  # noqa: ARG002
+        dataloader_idx: int = 0,  # noqa: ARG002
+        sample_id: str = None,
+    ) -> None:
+        """Write the predictions to disk."""
+        pred_dict = {}
+        for key, value in prediction.items():
+            # check object is tensor
+            if key in const.eval_keys:
+                pred_dict[key] = value.cpu().numpy()
+        np.savez_compressed(self.outdir / f"{batch['id'][0]}.npz", **pred_dict)
+
+        # Get best sample
+        confidence = 0.8 * pred_dict["iptm"] + 0.2 * pred_dict["ptm"]
+        best_idx = np.argmax(confidence)
+        best_sample_coords = pred_dict["coords"][best_idx]
+
+        prediction_out = {}
+        for k in prediction:
+            if k == "coords":
+                prediction_out[k] = torch.from_numpy(best_sample_coords)
+            else:
+                prediction_out[k] = prediction[k][0]
+
+        # Write structure
+        structure, _, _ = Structure.from_feat(prediction_out)
+        plddt_atom = (
+            prediction_out["atom_to_token"].float() @ prediction_out["plddt"].float()
+        )
+        structure.atoms["bfactor"] = (
+            plddt_atom[prediction_out["atom_pad_mask"].bool()].float().cpu().numpy()
+        )
+        cif_text = to_mmcif(structure)
+        open(self.refold_cif_dir / f"{batch['id'][0]}.cif", "w").write(cif_text)
+
+        # Failed prediction handling
+        if isinstance(prediction["exception"], bool):
+            if prediction["exception"]:
+                self.failed += 1
+        elif isinstance(prediction["exception"], list):
+            if prediction["exception"][0]:
+                self.failed += 1
+
+    def on_predict_epoch_end(
+        self,
+        trainer: Trainer,  # noqa: ARG002
+        pl_module: LightningModule,  # noqa: ARG002
+    ) -> None:
+        print(f"Number of failed structure predictions: {self.failed}")  # noqa: T201
+
+
+class AffinityWriter(BasePredictionWriter):
+    """Custom writer for predictions."""
+
+    def __init__(
+        self,
+        design_dir: str,
+    ) -> None:
+        super().__init__(write_interval="batch")
+        if design_dir is not None:
+            self.init_outdir(design_dir)
+
+    def init_outdir(self, design_dir):
+        self.outdir = Path(design_dir) / const.affinity_dirname
+        self.outdir.mkdir(exist_ok=True, parents=True)
+        self.failed = 0
+
+    def write_on_batch_end(  # noqa: PLR0915
+        self,
+        trainer: Trainer = None,  # noqa: ARG002
+        pl_module: LightningModule = None,  # noqa: ARG002
+        prediction: Dict[str, Tensor] = None,
+        batch_indices: List[int] = None,  # noqa: ARG002
+        batch: Dict[str, Tensor] = None,
+        batch_idx: int = None,  # noqa: ARG002
+        dataloader_idx: int = 0,  # noqa: ARG002
+        sample_id: str = None,
+    ) -> None:
+        """Write the predictions to disk."""
+        pred_dict = {}
+        for key, value in prediction.items():
+            # check object is tensor
+            if key in const.eval_keys:
+                pred_dict[key] = value.cpu().numpy()
+        np.savez_compressed(self.outdir / f"{batch['id'][0]}.npz", **pred_dict)
+
+        if isinstance(prediction["exception"], bool):
+            if prediction["exception"]:
+                self.failed += 1
+        elif isinstance(prediction["exception"], list):
+            if prediction["exception"][0]:
+                self.failed += 1
+
+    def on_predict_epoch_end(
+        self,
+        trainer: Trainer,  # noqa: ARG002
+        pl_module: LightningModule,  # noqa: ARG002
+    ) -> None:
+        print(f"Number of failed affinity predictions: {self.failed}")  # noqa: T201
+
+
+class DesignWriter(BasePredictionWriter):
+    """Custom writer for predictions."""
+
+    def __init__(
+        self,
+        output_dir: str,
+        res_atoms_only: bool,
+        save_traj: bool = False,
+        save_x0_traj: bool = False,
+        atom14: bool = True,
+        atom37: bool = False,
+        backbone_only: bool = False,
+        inverse_fold: bool = False,
+        file_suffix: str = "",
+        write_native: bool = True,
+        design: bool = True,
+    ) -> None:
+        """Initialize the writer.
+
+        Parameters
+        ----------
+        output_dir : str
+            The directory to save the predictions.
+
+        """
+        super().__init__(write_interval="batch")
+        self.mol_dir = Path(output_dir) / const.molecules_dirname
+        self.mol_dir.mkdir(parents=True, exist_ok=True)
+        self.save_traj = save_traj
+        self.save_x0_traj = save_x0_traj
+        self.res_atoms_only = res_atoms_only
+        self.file_suffix = file_suffix
+        self.failed = 0
+        self.write_native = write_native
+        self.design = design
+
+        # Create the output directories
+        self.atom14 = atom14
+        self.atom37 = atom37
+        self.inverse_fold = inverse_fold
+        self.backbone_only = backbone_only
+        self.used_stems = set()
+        self.init_outdir(output_dir)
+
+    def init_outdir(self, outdir):
+        self.outdir = Path(outdir)
+        self.outdir.mkdir(parents=True, exist_ok=True)
+
+    def write_on_batch_end(  # noqa: PLR0915
+        self,
+        trainer: Trainer = None,  # noqa: ARG002
+        pl_module: LightningModule = None,  # noqa: ARG002
+        prediction: Dict[str, Tensor] = None,
+        batch_indices: List[int] = None,  # noqa: ARG002
+        batch: Dict[str, Tensor] = None,
+        batch_idx: int = None,  # noqa: ARG002
+        dataloader_idx: int = 0,  # noqa: ARG002
+        sample_id: str = None,
+    ) -> None:
+        if prediction["exception"]:
+            self.failed += 1
+            return
+        n_samples, _, _ = prediction["coords"].shape
+
+        # TODO: remove this which is only here for temporary backward compatibility
+        masker = BoltzMasker(mask=True, mask_backbone=False)
+        feat_masked = masker(batch)
+        prediction["ref_element"] = feat_masked["ref_element"]
+        prediction["ref_atom_name_chars"] = feat_masked["ref_atom_name_chars"]
+        """Write the predictions to disk."""
+        # Check for extra molecules
+        if batch["extra_mols"] is not None:
+            extra_mols = batch["extra_mols"][0]
+            for k, v in extra_mols.items():
+                with open(self.mol_dir / f"{k}.pkl", "wb") as f:
+                    pickle.dump(v, f)
+
+        # write samples to disk
+        for n in range(n_samples):
+            # get structure for all generated coords
+            sample, native = {}, {}
+
+            for k in set(prediction.keys()) & set(batch.keys()):
+                if k == "coords":
+                    native[k] = batch[k][0][0].unsqueeze(0)
+                    sample[k] = prediction[k][n]
+                    
+                if k in const.token_features:
+                    sample[k] = prediction[k][0]
+                    native[k] = batch[k][0]
+                elif k in const.atom_features:
+                    if k == "coords":
+                        native[k] = batch[k][0][0].unsqueeze(0)
+                        sample[k] = prediction[k][n]
+                    else:
+                        native[k] = batch[k][0]
+                        sample[k] = prediction[k][0]
+                elif k == "exception":
+                    sample[k] = prediction[k]
+                    native[k] = batch[k]
+                else:
+                    native[k] = batch[k][0]
+                    sample[k] = prediction[k][0]
+                    native[k] = batch[k][0]
+
+            if self.atom14:
+                sample = res_from_atom14(sample)
+            elif self.atom37:
+                sample = res_from_atom37(sample)
+            elif self.backbone_only:
+                sample = res_all_gly(sample)
+
+
+            design_mask = batch["design_mask"][0].bool()
+            assert design_mask.sum() == sample["design_mask"].sum()
+
+            if self.inverse_fold:
+                token_ids = torch.argmax(sample["res_type"], dim=-1)
+                tokens = [const.tokens[i] for i in token_ids]
+                ccds = [convert_ccd(token) for token in tokens]
+
+                ccds = torch.tensor(ccds).to(sample["res_type"])
+                sample["ccd"][design_mask] = ccds[design_mask]
+
+            try:
+                structure, _, _ = Structure.from_feat(sample)
+                str_native, _, _ = Structure.from_feat(native)
+
+                # write structure to cif
+                if sample_id is not None:
+                    file_name = f"{sample_id}_{n}{self.file_suffix}"
+                else:
+                    stem = str(batch["id"][0])
+                    multiplicity = getattr(trainer.datamodule.cfg, "multiplicity", 1)
+                    total_files = multiplicity * n_samples
+                    sample_idx = (
+                        int(batch["data_sample_idx"][0])
+                        if "data_sample_idx" in batch
+                        else 0
+                    )
+                    global_idx = sample_idx * n_samples + n
+
+                    if total_files > 1:
+                        num_digits = len(str(total_files - 1))
+                        file_name = (
+                            f"{stem}_{global_idx:0{num_digits}d}{self.file_suffix}"
+                        )
+                    else:
+                        file_name = f"{stem}{self.file_suffix}"
+
+                native_path = f"{self.outdir}/{file_name}_native.cif"
+                gen_path = f"{self.outdir}/{file_name}.cif"
+
+                # design mask bfactor
+                design_mask = batch["design_mask"][0].float()
+                atom_design_mask = (
+                    sample["atom_to_token"].float() @ design_mask.unsqueeze(-1).float()
+                )
+                design_mask = native["design_mask"].float()
+
+                atom_design_mask = atom_design_mask.squeeze().bool()
+                bfactor = atom_design_mask * 100
+
+                # binding type bfactor
+                binding_type = batch["binding_type"][0].float()
+                atom_binding_type = (
+                    sample["atom_to_token"].float() @ binding_type.unsqueeze(-1).float()
+                )
+
+                atom_binding_type = atom_binding_type.squeeze().int()
+                bfactor[atom_binding_type == const.binding_type_ids["BINDING"]] = 60
+
+                bfactor = atom_design_mask[sample["atom_pad_mask"].bool()].float()
+                str_native.atoms["bfactor"] = bfactor.cpu().numpy()
+                structure.atoms["bfactor"] = bfactor.cpu().numpy()
+
+                # Add dummy (0-coord) design side chains if inverse fold
+                if self.inverse_fold:
+                    atom_design_mask_no_pad = atom_design_mask[
+                        native["atom_pad_mask"].bool()
+                    ]
+                    res_design_mask = np.array(
+                        [
+                            all(
+                                atom_design_mask_no_pad[
+                                    res["atom_idx"] : res["atom_idx"] + res["atom_num"]
+                                ]
+                            )
+                            for res in structure.residues
+                        ]
+                    )
+                    structure = Structure.add_side_chains(
+                        structure, residue_mask=res_design_mask
+                    )
+
+                if self.write_native:
+                    open(native_path, "w").write(to_mmcif(str_native))
+
+                pred_binding_mask = prediction["binding_type"][0].cpu().bool().numpy()
+                if self.design:
+                    chain_design_mask = (
+                        prediction["chain_design_mask"][0].cpu().bool().numpy()
+                    )
+                pred_design_mask = prediction["design_mask"][0].cpu().bool().numpy()
+                design_color_features = np.ones_like(pred_binding_mask) * 0.8
+                design_color_features[pred_binding_mask] = 1.0
+                if self.design:
+                    design_color_features[chain_design_mask] = 0.0
+                design_color_features[pred_design_mask] = 0.6
+
+                # Create a mask to identify unique token-to-res mappings.
+                # This is for small molecules where multiple tokens can be mapped to the same residue.
+                token_to_res = prediction["token_to_res"][0].cpu().numpy()
+                unique_mask = np.ones_like(token_to_res, dtype=bool)
+                unique_mask[1:] = token_to_res[1:] != token_to_res[:-1]
+                design_color_features = design_color_features[unique_mask]
+                open(gen_path, "w").write(
+                    to_mmcif(
+                        structure,
+                        design_coloring=True,
+                        color_features=design_color_features,
+                    )
+                )
+
+                # Write metadata
+                metadata_path = f"{self.outdir}/{file_name}.npz"
+                token_mask = sample["token_pad_mask"].bool()
+
+                # Build metadata dict with required fields
+                metadata_dict = {
+                    "design_mask": design_mask[token_mask].cpu().numpy(),
+                    "mol_type": sample["mol_type"][token_mask].cpu().numpy(),
+                    "ss_type": sample["ss_type"][token_mask].cpu().numpy(),
+                    "token_resolved_mask": sample["token_resolved_mask"][token_mask].cpu().numpy(),
+                    "binding_type": binding_type[token_mask].cpu().numpy(),
+                }
+
+                # Add optional fields only if they have valid values (avoid None -> object array)
+                if "inverse_fold_design_mask" in sample:
+                    metadata_dict["inverse_fold_design_mask"] = (
+                        sample["inverse_fold_design_mask"][token_mask].cpu().numpy()
+                    )
+
+                # Per-residue amino acid constraints (for inverse folding step)
+                # Only save if constraints exist AND have non-zero values
+                if "aa_constraint_mask" in batch:
+                    aa_mask = batch["aa_constraint_mask"][0]
+                    if aa_mask.any():  # Only save if there are actual constraints
+                        metadata_dict["aa_constraint_mask"] = aa_mask[token_mask].cpu().numpy()
+
+                np.savez_compressed(metadata_path, **metadata_dict)
+
+                # Write trajectories
+                if self.save_traj:
+                    trajs = torch.stack(prediction["coords_traj"], dim=1)
+                    traj = trajs[n]
+                    aligned = [traj[0]]
+                    for frame in traj[1:]:
+                        with torch.autocast("cuda", enabled=False):
+                            aligned.append(
+                                weighted_rigid_align(
+                                    frame.float().unsqueeze(0),
+                                    aligned[-1].float().unsqueeze(0),
+                                    sample["atom_pad_mask"].float().unsqueeze(0),
+                                    sample["atom_pad_mask"].float().unsqueeze(0),
+                                )
+                                .to(frame)
+                                .squeeze()
+                            )
+
+                    pdbs = []
+                    all_coords = []
+                    ensemble = []
+                    atom_idx = 0
+                    for idx, frame in tqdm(
+                        enumerate(aligned), desc="Writing traj.", total=len(aligned)
+                    ):
+                        sample["coords"] = frame
+                        if self.atom14:
+                            sample = res_from_atom14(sample)
+                        elif self.atom37:
+                            sample = res_from_atom37(sample)
+                        else:
+                            raise ValueError("Either atom14 or atom37 must be true")
+
+                        str_frame, _, _ = Structure.from_feat(sample)
+                        pdbs.append(to_pdb(str_frame))
+                        all_coords.append(str_frame.coords)
+                        ensemble.append(
+                            (
+                                atom_idx,
+                                len(str_frame.coords),
+                            )
+                        )
+                        atom_idx += len(str_frame.coords)
+
+                    open(self.outdir / f"{file_name}_traj.pdb", "w").write(
+                        self.combine_pdb_models(pdbs)
+                    )
+
+                # Write x0 trajectories
+                if self.save_x0_traj:
+                    trajs = torch.stack(prediction["x0_coords_traj"], dim=1)
+                    traj = trajs[n]
+                    aligned = [traj[0]]
+                    for frame in traj[1:]:
+                        with torch.autocast("cuda", enabled=False):
+                            aligned.append(
+                                weighted_rigid_align(
+                                    frame.float().unsqueeze(0),
+                                    aligned[-1].float().unsqueeze(0),
+                                    sample["atom_pad_mask"].float().unsqueeze(0),
+                                    sample["atom_pad_mask"].float().unsqueeze(0),
+                                )
+                                .to(frame)
+                                .squeeze()
+                            )
+
+                    pdbs = []
+                    all_coords = []
+                    ensemble = []
+                    atom_idx = 0
+                    for idx, frame in tqdm(
+                        enumerate(aligned), desc="Writing x0 traj.", total=len(aligned)
+                    ):
+                        sample["coords"] = frame
+                        if self.atom14:
+                            sample = res_from_atom14(sample)
+                        elif self.atom37:
+                            sample = res_from_atom37(sample)
+                        else:
+                            raise ValueError("Either atom14 or atom37 must be true")
+
+                        str_frame, _, _ = Structure.from_feat(sample)
+                        pdbs.append(to_pdb(str_frame))
+                        all_coords.append(str_frame.coords)
+                        ensemble.append(
+                            (
+                                atom_idx,
+                                len(str_frame.coords),
+                            )
+                        )
+                        atom_idx += len(str_frame.coords)
+
+                    open(self.outdir / f"{file_name}_x0_traj.pdb", "w").write(
+                        self.combine_pdb_models(pdbs)
+                    )
+
+            except Exception as e:  # noqa: BLE001
+                import traceback
+
+                traceback.print_exc()  # noqa: T201
+                msg = f"predict/writer.py: Validation structure writing failed on {batch['id'][0]} with error {e}. Skipping."
+                print(msg)
+
+    def combine_pdb_models(self, pdb_strings):
+        combined_pdb = ""
+        model_number = 1
+
+        for pdb in pdb_strings:
+            # Add a model number at the start of each model
+            combined_pdb += f"MODEL     {model_number}\n"
+            combined_pdb += pdb.split("\nEND")[0]
+            combined_pdb += "\nENDMDL\n"  # End of model marker
+            model_number += 1
+
+        return combined_pdb
+
+    def on_predict_epoch_end(
+        self,
+        trainer: Trainer,  # noqa: ARG002
+        pl_module: LightningModule,  # noqa: ARG002
+    ) -> None:
+        """Print the number of failed examples."""
+        print(f"Number of failed examples: {self.failed}")  # noqa: T201

model/boltzgen/task/task.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+
+from omegaconf import OmegaConf
+
+
+class Task(ABC):
+    """A task to be executed."""
+
+    @abstractmethod
+    def run(self, config: OmegaConf) -> None:
+        """Run the task.
+
+        Parameters
+        ----------
+        config : OmegaConf
+            The configuration for the task.
+
+        """
+        raise NotImplementedError

model/boltzgen/task/train/data.py

@@ -0,0 +1,1259 @@
+from collections import defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+import traceback
+from typing import Dict, List, Optional
+
+import numpy as np
+import pandas as pd
+import pytorch_lightning as pl
+import torch
+from rdkit.Chem import Mol
+from torch import Tensor
+from torch.utils.data import DataLoader
+
+from boltzgen.data import const
+from boltzgen.data.crop.cropper import Cropper
+from boltzgen.data.select.selector import Selector
+from boltzgen.data.data import (
+    MSA,
+    Input,
+    Manifest,
+    Record,
+    Structure,
+)
+from boltzgen.data.feature.featurizer import Featurizer
+from boltzgen.data.filter.dynamic.filter import DynamicFilter
+from boltzgen.data.mol import load_canonicals, load_molecules
+from boltzgen.data.pad import pad_to_max
+from boltzgen.data.sample.sampler import Sample, Sampler
+from boltzgen.data.template.features import load_dummy_templates
+from boltzgen.data.tokenize.tokenizer import Tokenizer
+from boltzgen.task.predict import data_ligands, data_protein_binder
+
+
+@dataclass
+class DatasetConfig:
+    """Dataset configuration."""
+
+    target_dir: str
+    msa_dir: str
+    prob: Optional[float]
+    sampler: Sampler
+    cropper: Cropper
+    selector: Optional[Selector] = None
+    manifest_path: Optional[str] = None
+    filters: Optional[list[DynamicFilter]] = None
+    split: Optional[str] = None
+    symmetry_correction: bool = True
+    val_group: Optional[str] = "RCSB"
+    use_train_subset: Optional[float] = None
+    moldir: Optional[str] = None
+    override_bfactor: Optional[bool] = False
+    override_method: Optional[str] = None
+
+
+@dataclass
+class DataConfig:
+    """Data configuration."""
+
+    datasets: List[DatasetConfig]
+    featurizer: Featurizer
+    tokenizer: Tokenizer
+    selector: Selector
+    max_atoms: int
+    max_tokens: int
+    max_seqs: int
+    samples_per_epoch: int
+    batch_size: int
+    num_workers: int
+    random_seed: int
+    pin_memory: bool
+    atoms_per_window_queries: int
+    min_dist: float
+    max_dist: float
+    num_bins: int
+    overfit: Optional[int] = None
+    pad_to_max_tokens: bool = False
+    pad_to_max_atoms: bool = False
+    pad_to_max_seqs: bool = False
+    return_train_symmetries: bool = False
+    return_val_symmetries: bool = True
+    val_batch_size: int = 1
+    single_sequence_prop_training: float = 0.0
+    msa_sampling_training: bool = False
+    moldir: Optional[str] = None
+    compute_frames: bool = True
+    backbone_only: bool = False
+    atom14: bool = False
+    atom37: bool = False
+    design: bool = False
+    monomer_split: str = None
+    monomer_target_dir: str = None
+    monomer_seq_len: int = 100
+    monomer_target_structure_condition: bool = True
+    inverse_fold: bool = False
+    ligand_split: str = None
+    ligand_target_dir: str = None
+    ligand_seq_len: int = 100
+    use_msa: bool = True
+    disulfide_prob: float = 1.0
+    disulfide_on: bool = False
+
+
+@dataclass
+class Dataset:
+    """Data holder."""
+
+    samples: pd.DataFrame
+    struct_dir: Path
+    msa_dir: Path
+    record_dir: Path
+    prob: float
+    cropper: Cropper
+    tokenizer: Tokenizer
+    featurizer: Featurizer
+    val_group: str
+    selector: Selector
+    symmetry_correction: bool = True
+    moldir: Optional[str] = None
+    override_bfactor: Optional[bool] = False
+    override_method: Optional[str] = None
+
+
+def load_record(record_id: str, record_dir: Path) -> Record:
+    """Load the given record.
+
+    Parameters
+    ----------
+    record_id : str
+        The record id to load.
+    record_dir : Path
+        The path to the record directory.
+
+    Returns
+    -------
+    Record
+        The loaded record.
+    """
+    return Record.load(record_dir / f"{record_id}.json")
+
+
+def load_structure(record: Record, struct_dir: Path) -> Structure:
+    """Load the given input data.
+
+    Parameters
+    ----------
+    record : str
+        The record to load.
+    target_dir : Path
+        The path to the data directory.
+
+    Returns
+    -------
+    Input
+        The loaded input.
+
+    """
+    if (struct_dir / f"{record.id}.npz").exists():
+        structure_path = struct_dir / f"{record.id}.npz"
+    else:
+        structure_path = struct_dir / f"{record.id}" / f"{record.id}_model_0.npz"
+    return Structure.load(structure_path)
+
+
+def load_msas(chain_ids: set[int], record: Record, msa_dir: Path) -> Input:
+    """Load the given input data.
+
+    Parameters
+    ----------
+    chain_ids : set[int]
+        The chain ids to load.
+    record : Record
+        The record to load.
+    msa_dir : Path
+        The path to the MSA directory.
+
+    Returns
+    -------
+    Input
+        The loaded input.
+
+    """
+    msas = {}
+    for chain in record.chains:
+        if chain.chain_id not in chain_ids:
+            continue
+
+        msa_id = chain.msa_id
+        if msa_id != -1:
+            msa_path = msa_dir / f"{msa_id}.npz"
+            msa = MSA.load(msa_path)
+            msas[chain.chain_id] = msa
+
+    return msas
+
+
+def collate(data: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
+    """Collate the data.
+
+    Parameters
+    ----------
+    data : List[Dict[str, Tensor]]
+        The data to collate.
+
+    Returns
+    -------
+    Dict[str, Tensor]
+        The collated data.
+
+    """
+    # Get the keys
+    keys = data[0].keys()
+
+    # Collate the data
+    collated = {}
+    for key in keys:
+        values = [d[key] for d in data]
+
+        if key not in [
+            "all_coords",
+            "all_resolved_mask",
+            "crop_to_all_atom_map",
+            "chain_symmetries",
+            "chain_swaps",
+            "amino_acids_symmetries",
+            "ligand_symmetries",
+            "activity_name",
+            "activity_qualifier",
+            "sid",
+            "cid",
+            "normalized_protein_accession",
+            "pair_id",
+            "ligand_edge_index",
+            "ligand_edge_lower_bounds",
+            "ligand_edge_upper_bounds",
+            "ligand_edge_bond_mask",
+            "ligand_edge_angle_mask",
+            "connections_edge_index",
+            "ligand_chiral_atom_index",
+            "ligand_chiral_check_mask",
+            "ligand_chiral_atom_orientations",
+            "ligand_stereo_bond_index",
+            "ligand_stereo_check_mask",
+            "ligand_stereo_bond_orientations",
+            "ligand_aromatic_5_ring_index",
+            "ligand_aromatic_6_ring_index",
+            "ligand_planar_double_bond_index",
+            "pdb_id",
+            "id",
+            "structure_bonds",
+            "extra_mols",
+        ]:
+            if values[0] is not None:
+                # Check if all have the same shape
+                shape = values[0].shape
+                if not all(v.shape == shape for v in values):
+                    values = pad_to_max(values, 0)
+                else:
+                    values = torch.stack(values, dim=0)
+
+        # Stack the values
+        collated[key] = values
+
+    return collated
+
+
+class TrainingDataset(torch.utils.data.Dataset):
+    """Base iterable dataset."""
+
+    def __init__(
+        self,
+        datasets: List[Dataset],
+        canonicals: dict[str, Mol],
+        moldir: str,
+        samples_per_epoch: int,
+        max_atoms: int,
+        max_tokens: int,
+        max_seqs: int,
+        pad_to_max_atoms: bool = False,
+        pad_to_max_tokens: bool = False,
+        pad_to_max_seqs: bool = False,
+        atoms_per_window_queries: int = 32,
+        min_dist: float = 2.0,
+        max_dist: float = 22.0,
+        num_bins: int = 64,
+        overfit: Optional[int] = None,
+        return_symmetries: Optional[bool] = False,
+        single_sequence_prop: Optional[float] = 0.0,
+        msa_sampling: bool = False,
+        compute_frames: bool = True,
+        backbone_only: bool = False,
+        atom14: bool = False,
+        atom37: bool = False,
+        design: bool = False,
+        disulfide_prob: float = 1.0,
+        disulfide_on: bool = False,
+        use_msa: bool = True,
+        inverse_fold: bool = False,
+    ) -> None:
+        """Initialize the training dataset.
+
+        Parameters
+        ----------
+        datasets : List[Dataset]
+            The datasets to sample from.
+        samplers : List[Sampler]
+            The samplers to sample from each dataset.
+        probs : List[float]
+            The probabilities to sample from each dataset.
+        samples_per_epoch : int
+            The number of samples per epoch.
+        max_tokens : int
+            The maximum number of tokens.
+
+        """
+        super().__init__()
+        self.datasets = datasets
+        self.canonicals = canonicals
+        self.moldir = moldir
+        self.probs = [d.prob for d in datasets]
+        self.samples_per_epoch = samples_per_epoch
+        self.max_tokens = max_tokens
+        self.max_seqs = max_seqs
+        self.max_atoms = max_atoms
+        self.pad_to_max_tokens = pad_to_max_tokens
+        self.pad_to_max_atoms = pad_to_max_atoms
+        self.pad_to_max_seqs = pad_to_max_seqs
+        self.atoms_per_window_queries = atoms_per_window_queries
+        self.min_dist = min_dist
+        self.max_dist = max_dist
+        self.num_bins = num_bins
+        self.return_symmetries = return_symmetries
+        self.backbone_only = backbone_only
+        self.atom14 = atom14
+        self.atom37 = atom37
+        self.design = design
+        self.disulfide_prob = disulfide_prob
+        self.disulfide_on = disulfide_on
+        self.single_sequence_prop = single_sequence_prop
+        self.msa_sampling = msa_sampling
+        self.use_msa = use_msa
+        self.overfit = overfit
+        self.compute_frames = compute_frames
+        self.inverse_fold = inverse_fold
+
+        self.samples: list[list[Dict]] = []
+        self.samples_weight: list[list[float]] = []
+        for d in self.datasets:
+            if self.overfit:
+                samples = d.samples[: self.overfit]
+            else:
+                samples = d.samples
+            self.samples.append(
+                [
+                    samples.iloc[sample_idx].to_dict()
+                    for sample_idx in range(len(samples))
+                ]
+            )
+            self.samples_weight.append(samples["weight"].tolist())
+
+    def __getitem__(self, idx: int) -> Dict[str, Tensor]:
+        """Get an item from the dataset.
+
+        Returns
+        -------
+        Dict[str, Tensor]
+            The sampled data features.
+
+        """
+        # Set a random state
+        random = np.random.default_rng()
+
+        # Pick a random dataset
+        dataset_idx = random.choice(len(self.datasets), p=self.probs)
+
+        dataset = self.datasets[dataset_idx]
+
+        # Get a sample from the dataset
+        samples = self.samples[dataset_idx]
+        sample_idx = random.choice(
+            len(samples),
+            p=(
+                self.samples_weight[dataset_idx]
+                / np.sum(self.samples_weight[dataset_idx])
+                if self.overfit
+                else self.samples_weight[dataset_idx]
+            ),
+        )
+
+        sample = samples[sample_idx]
+        sample: Sample = Sample(
+            record_id=str(sample["record_id"]),
+            chain_id=(
+                int(sample["chain_id"]) if sample["chain_id"] is not None else None
+            ),
+            interface_id=(
+                int(sample["interface_id"])
+                if sample["interface_id"] is not None
+                else None
+            ),
+            weight=float(sample["weight"]),
+        )
+
+        # Load record
+        record = load_record(sample.record_id, dataset.record_dir)
+
+        # Get the structure
+        try:
+            structure = load_structure(record, dataset.struct_dir)
+        except Exception as e:  # noqa: BLE001
+            print(f"Failed to load input for {record.id} with error {e}. Skipping.")
+            return self.__getitem__(random.integers(0, len(self)))
+
+        # Tokenize structure
+        try:
+            tokenized = dataset.tokenizer.tokenize(
+                structure, inverse_fold=self.inverse_fold
+            )
+        except Exception as e:  # noqa: BLE001
+            print(f"Tokenizer failed on {record.id} with error {e}. Skipping.")
+            traceback.print_exc()  # noqa: T201
+            return self.__getitem__(random.integers(0, len(self)))
+
+        # Compute crop
+        try:
+            if self.max_tokens is not None and len(tokenized.tokens) > self.max_tokens:
+                tokenized = dataset.cropper.crop(
+                    tokenized,
+                    max_atoms=self.max_atoms,
+                    max_tokens=self.max_tokens,
+                    chain_id=sample.chain_id,
+                    interface_id=sample.interface_id,
+                    random=random,
+                    prefer_protein_queries=self.inverse_fold,
+                )
+                if len(tokenized.tokens) == 0:
+                    msg = "No tokens in cropped structure."
+                    raise ValueError(msg)  # noqa: TRY301
+        except Exception as e:  # noqa: BLE001
+            print(f"Cropper failed on {record.id} with error {e}. Skipping.")
+            traceback.print_exc()  # noqa: T201
+            return self.__getitem__(random.integers(0, len(self)))
+
+        # Select which tokens to design
+        try:
+            tokenized, design_task = dataset.selector.select(
+                tokenized,
+                random=random,
+            )
+        except Exception as e:  # noqa: BLE001
+            print(f"Selector failed on {record.id} with error {e}. Skipping.")  # noqa: T201
+            traceback.print_exc()  # noqa: T201
+            return self.__getitem__(random.integers(0, len(self)))
+        structure = tokenized.structure
+
+        # Get unique chain ids
+        chain_ids = set(tokenized.tokens["asym_id"])
+
+        # Load msas and templates
+        try:
+            if self.use_msa:
+                msas = load_msas(
+                    chain_ids=chain_ids,
+                    record=record,
+                    msa_dir=dataset.msa_dir,
+                )
+            else:
+                msas = {}
+        except Exception as e:  # noqa: BLE001
+            print(f"MSA loading failed for {record.id} with error {e}. Skipping.")
+            return self.__getitem__(random.integers(0, len(self)))
+
+        # Load molecules
+        try:
+            # Try to find molecules in the dataset moldir if provided
+            # Find missing ones in global moldir and check if all found
+            molecules = {}
+            molecules.update(self.canonicals)
+            mol_names = set(tokenized.tokens["res_name"].tolist())
+            mol_names = mol_names - set(self.canonicals.keys())
+            if dataset.moldir is not None:
+                molecules.update(load_molecules(dataset.moldir, mol_names))
+
+            mol_names = mol_names - set(molecules.keys())
+            molecules.update(load_molecules(self.moldir, mol_names))
+        except Exception as e:  # noqa: BLE001
+            print(f"Molecule loading failed for {record.id} with error {e}. Skipping.")
+            return self.__getitem__(random.integers(0, len(self)))
+
+        # Finalize input data
+        input_data = Input(
+            tokens=tokenized.tokens,
+            bonds=tokenized.bonds,
+            token_to_res=tokenized.token_to_res,
+            structure=tokenized.structure,
+            msa=msas,
+            templates=None,
+            record=record,
+        )
+
+        # Compute features
+        try:
+            features: dict = dataset.featurizer.process(
+                input_data,
+                molecules=molecules,
+                random=random,
+                training=True,
+                max_atoms=self.max_atoms if self.pad_to_max_atoms else None,
+                max_tokens=self.max_tokens if self.pad_to_max_tokens else None,
+                max_seqs=self.max_seqs,
+                pad_to_max_seqs=self.pad_to_max_seqs,
+                atoms_per_window_queries=self.atoms_per_window_queries,
+                min_dist=self.min_dist,
+                max_dist=self.max_dist,
+                num_bins=self.num_bins,
+                compute_symmetries=self.return_symmetries,
+                single_sequence_prop=self.single_sequence_prop,
+                msa_sampling=self.msa_sampling,
+                override_bfactor=dataset.override_bfactor,
+                override_method=dataset.override_method,
+                compute_frames=self.compute_frames,
+                backbone_only=self.backbone_only,
+                atom14=self.atom14,
+                atom37=self.atom37,
+                design=self.design,
+                disulfide_prob=self.disulfide_prob,
+                inverse_fold=self.inverse_fold,
+            )
+        except Exception as e:  # noqa: BLE001
+            print(f"Featurizer failed on {record.id} with error {e}. Skipping.")
+            traceback.print_exc()
+            return self.__getitem__(random.integers(0, len(self)))
+
+        # Check that there is enough stuff to design in the inverse folding case so we have no nan losses
+        if self.inverse_fold and features["design_mask"].sum() < 3:
+            print(f"Skipping {record.id}. Fewer than 3 design residues.")
+            return self.__getitem__(random.integers(0, len(self)))
+
+        # Set template features
+        template_features = load_dummy_templates(
+            tdim=1, num_tokens=len(features["res_type"])
+        )
+        features.update(template_features)
+
+        features.update({"id": sample.record_id})
+        features["pdb_id"] = record.id
+
+        # Assert that all design tokens make sense
+        bad_protein_mask = (
+            (~features["is_standard"].bool())
+            & features["design_mask"].bool()
+            & (features["mol_type"] == const.chain_type_ids["PROTEIN"])
+        )
+        assert not bad_protein_mask.any()
+
+        return features
+
+    def __len__(self) -> int:
+        """Get the length of the dataset.
+
+        Returns
+        -------
+        int
+            The length of the dataset.
+
+        """
+        return self.samples_per_epoch
+
+
+class ValidationDataset(torch.utils.data.Dataset):
+    """Base iterable dataset."""
+
+    def __init__(
+        self,
+        datasets: List[Dataset],
+        canonicals: dict[str, Mol],
+        moldir: str,
+        seed: int,
+        max_atoms: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        max_seqs: Optional[int] = None,
+        pad_to_max_atoms: bool = False,
+        pad_to_max_tokens: bool = False,
+        pad_to_max_seqs: bool = False,
+        atoms_per_window_queries: int = 32,
+        min_dist: float = 2.0,
+        max_dist: float = 22.0,
+        num_bins: int = 64,
+        overfit: Optional[int] = None,
+        return_symmetries: Optional[bool] = False,
+        compute_frames: bool = True,
+        backbone_only: bool = False,
+        atom14: bool = False,
+        atom37: bool = False,
+        design: bool = False,
+        inverse_fold: bool = False,
+        disulfide_prob: float = 1.0,
+        disulfide_on: bool = False,
+    ) -> None:
+        """Initialize the training dataset.
+
+        Parameters
+        ----------
+        datasets : List[Dataset]
+            The datasets to sample from.
+        seed : int
+            The random seed.
+        max_tokens : int
+            The maximum number of tokens.
+        overfit : bool
+            Whether to overfit the dataset
+
+        """
+        super().__init__()
+        self.datasets = datasets
+        self.canonicals = canonicals
+        self.moldir = moldir
+        self.max_atoms = max_atoms
+        self.max_tokens = max_tokens
+        self.max_seqs = max_seqs
+        self.seed = seed
+        self.pad_to_max_tokens = pad_to_max_tokens
+        self.pad_to_max_atoms = pad_to_max_atoms
+        self.pad_to_max_seqs = pad_to_max_seqs
+        self.overfit = overfit
+        self.atoms_per_window_queries = atoms_per_window_queries
+        self.min_dist = min_dist
+        self.max_dist = max_dist
+        self.num_bins = num_bins
+        self.return_symmetries = return_symmetries
+        self.compute_frames = compute_frames
+        self.backbone_only = backbone_only
+        self.atom14 = atom14
+        self.atom37 = atom37
+        self.design = design
+        self.inverse_fold = inverse_fold
+        self.disulfide_prob = disulfide_prob
+        self.disulfide_on = disulfide_on
+
+    def __getitem__(self, idx: int) -> Structure:
+        """Get an item from the dataset.
+
+        Returns
+        -------
+        Dict[str, Tensor]
+            The sampled data features.
+
+        """
+        # Set random state
+        seed = self.seed if self.overfit is None else None
+        random = np.random.default_rng(seed)
+
+        # Pick dataset based on idx
+        for idx_dataset, dataset in enumerate(self.datasets):  # noqa: B007
+            size = len(dataset.samples)
+            if self.overfit is not None:
+                size = min(size, self.overfit)
+            if idx < size:
+                break
+            idx -= size
+
+        # Get a sample from the dataset
+        sample = Sample(**dataset.samples.iloc[idx].to_dict())
+        record = load_record(sample.record_id, dataset.record_dir)
+
+        # Get the structure
+        try:
+            structure = load_structure(record, dataset.struct_dir)
+        except Exception as e:  # noqa: BLE001
+            print(f"Failed to load input for {record.id} with error {e}. Skipping.")
+            return self.__getitem__(0)
+
+        # Tokenize structure
+        try:
+            tokenized = dataset.tokenizer.tokenize(structure)
+        except Exception as e:  # noqa: BLE001
+            print(f"Tokenizer failed on {record.id} with error {e}. Skipping.")  # noqa: T201
+            return self.__getitem__(0)
+
+        # Compute crop
+        try:
+            if self.max_tokens is not None:
+                tokenized = dataset.cropper.crop(
+                    tokenized,
+                    max_atoms=self.max_atoms,
+                    max_tokens=self.max_tokens,
+                    chain_id=sample.chain_id,
+                    interface_id=sample.interface_id,
+                    random=random,
+                    prefer_protein_queries=self.inverse_fold,
+                )
+                if len(tokenized.tokens) == 0:
+                    msg = "No tokens in cropped structure."
+                    raise ValueError(msg)  # noqa: TRY301
+        except Exception as e:  # noqa: BLE001
+            print(f"Cropper failed on {record.id} with error {e}. Skipping.")
+            return self.__getitem__(0)
+
+        # Get unique chains
+        chain_ids = set(np.unique(tokenized.tokens["asym_id"]).tolist())
+
+        # Load msas and templates
+        try:
+            msas = load_msas(chain_ids, record, dataset.msa_dir)
+        except Exception as e:  # noqa: BLE001
+            print(f"MSA loading failed for {record.id} with error {e}. Skipping.")
+            return self.__getitem__(0)
+
+        # Select which tokens to design
+        try:
+            tokenized, design_task = dataset.selector.select(
+                tokenized,
+                random=random,
+            )
+        except Exception as e:  # noqa: BLE001
+            print(f"Selector failed on {sample.record_id} with error {e}. Skipping.")  # noqa: T201
+            traceback.print_exc()  # noqa: T201
+            return self.__getitem__(0)
+        structure = tokenized.structure
+
+        try:
+            # Try to find molecules in the dataset moldir if provided
+            # Find missing ones in global moldir and check if all found
+            molecules = {}
+            molecules.update(self.canonicals)
+            mol_names = set(tokenized.tokens["res_name"].tolist())
+            mol_names = mol_names - set(self.canonicals.keys())
+            if dataset.moldir is not None:
+                molecules.update(load_molecules(dataset.moldir, mol_names))
+
+            mol_names = mol_names - set(molecules.keys())
+            molecules.update(load_molecules(self.moldir, mol_names))
+        except Exception as e:  # noqa: BLE001
+            print(f"Molecule loading failed for {record.id} with error {e}. Skipping.")
+            return self.__getitem__(0)
+
+        # Finalize input data
+        input_data = Input(
+            tokens=tokenized.tokens,
+            bonds=tokenized.bonds,
+            token_to_res=tokenized.token_to_res,
+            structure=tokenized.structure,
+            msa=msas,
+            templates=None,
+            record=record,
+        )
+
+        # Compute features
+        try:
+            features: dict = dataset.featurizer.process(
+                input_data,
+                molecules=molecules,
+                random=random,
+                training=False,
+                max_atoms=None,
+                max_tokens=None,
+                max_seqs=self.max_seqs,
+                pad_to_max_seqs=self.pad_to_max_seqs,
+                atoms_per_window_queries=self.atoms_per_window_queries,
+                min_dist=self.min_dist,
+                max_dist=self.max_dist,
+                num_bins=self.num_bins,
+                compute_symmetries=self.return_symmetries,
+                single_sequence_prop=0.0,
+                override_method=dataset.override_method,
+                compute_frames=self.compute_frames,
+                backbone_only=self.backbone_only,
+                atom14=self.atom14,
+                atom37=self.atom37,
+                design=self.design,
+                inverse_fold=self.inverse_fold,
+                disulfide_prob=self.disulfide_prob,
+                disulfide_on=self.disulfide_on,
+            )
+        except Exception as e:  # noqa: BLE001
+            print(f"Featurizer failed on {record.id} with error {e}. Skipping.")
+            return self.__getitem__(0)
+
+        # Check that there is enough stuff to design in the inverse folding case so we have no nan losses
+        if self.inverse_fold and features["design_mask"].sum() < 3:
+            print(f"Skipping {record.id}. Fewer than 3 design residues.")
+            return self.__getitem__(0)
+
+        # Set template features
+        template_features = load_dummy_templates(
+            tdim=1, num_tokens=len(features["res_type"])
+        )
+        features.update(template_features)
+
+        # Add dataset idx
+        idx_dataset = torch.tensor([idx_dataset], dtype=torch.long)
+        features.update({"idx_dataset": idx_dataset})
+        features.update({"id": record.id})
+        bad_protein_mask = (
+            (~features["is_standard"].bool())
+            & features["design_mask"].bool()
+            & (features["mol_type"] == const.chain_type_ids["PROTEIN"])
+        )
+        assert not bad_protein_mask.any()
+        return features
+
+    def __len__(self) -> int:
+        """Get the length of the dataset.
+
+        Returns
+        -------
+        int
+            The length of the dataaset.
+
+        """
+        if self.overfit is not None:
+            length = sum(len(d.samples[: self.overfit]) for d in self.datasets)
+        else:
+            length = sum(len(d.samples) for d in self.datasets)
+
+        return length
+
+
+class TrainingDataModule(pl.LightningDataModule):
+    """DataModule for BoltzGen training."""
+
+    def __init__(
+        self,
+        cfg: DataConfig,
+    ) -> None:
+        """Initialize the DataModule.
+
+        Parameters
+        ----------
+        config : DataConfig
+            The data configuration.
+
+        """
+        super().__init__()
+        self.cfg = cfg
+        self.inverse_fold = cfg.inverse_fold
+
+        assert self.cfg.val_batch_size == 1, "Validation only works with batch size=1."
+
+        # Load datasets
+        train: List[Dataset] = []
+        val: List[Dataset] = []
+
+        for data_config in cfg.datasets:
+            # Get relevant directories
+            if data_config.manifest_path is not None:
+                manifest_path = Path(data_config.manifest_path)
+            else:
+                manifest_path = Path(data_config.target_dir) / "manifest.json"
+            struct_dir = Path(data_config.target_dir) / "structures"
+            record_dir = Path(data_config.target_dir) / "records"
+            msa_dir = Path(data_config.msa_dir)
+
+            # Get moldir, if any
+            moldir = data_config.moldir
+            moldir = Path(moldir) if moldir is not None else None
+
+            # Load all records
+            manifest: Manifest = Manifest.load(manifest_path)
+
+            # Split records if givens
+            if data_config.split is not None:
+                with Path(data_config.split).open("r") as f:
+                    split = {x.lower() for x in f.read().splitlines()}
+
+                train_records = []
+                val_records = []
+                for record in manifest.records:
+                    if record.id.lower() in split:
+                        val_records.append(record)
+                    else:
+                        train_records.append(record)
+            else:
+                train_records = manifest.records
+                if cfg.overfit is None:
+                    val_records = []
+                else:
+                    print("Warning: modified overfit val behavior.")
+                    val_records = manifest.records[: cfg.overfit]
+
+            print("train_records before filter", len(train_records))
+
+            # Apply dataset-specific filters
+            if data_config.filters is not None:
+                train_records = [
+                    record
+                    for record in train_records
+                    if all(f.filter(record) for f in data_config.filters)
+                ]
+
+            # Train with subset of data
+            if data_config.use_train_subset is not None:
+                # Shuffle train_records list
+                assert 0 < data_config.use_train_subset < 1.0
+                rng = np.random.default_rng(cfg.random_seed)
+                rng.shuffle(train_records)
+                train_records = train_records[
+                    0 : int(len(train_records) * data_config.use_train_subset)
+                ]
+            print("train_records after filter", len(train_records))
+            print("val_records after filter", len(val_records))
+
+            # Get samples
+            train_samples: list[Sample] = data_config.sampler.sample(train_records)
+            val_samples: list[Sample] = [Sample(r.id) for r in val_records]
+
+            # Convert samples to pandas dataframe to avoid copy-on-write behavior
+            train_samples = pd.DataFrame(
+                [
+                    (
+                        r.record_id,
+                        r.chain_id,
+                        r.interface_id,
+                        r.weight,
+                    )
+                    for r in train_samples
+                ],
+                columns=["record_id", "chain_id", "interface_id", "weight"],
+            )
+            val_samples = pd.DataFrame(
+                [s.record_id for s in val_samples], columns=["record_id"]
+            )
+
+            # Use appropriate string type
+            train_samples = train_samples.replace({np.nan: None})
+            val_samples = val_samples.replace({np.nan: None})
+            train_samples["record_id"] = train_samples["record_id"].astype("string")
+            val_samples["record_id"] = val_samples["record_id"].astype("string")
+
+            del manifest, train_records, val_records
+            # Create train dataset
+            if data_config.prob > 0:
+                train.append(
+                    Dataset(
+                        samples=train_samples,
+                        record_dir=record_dir,
+                        struct_dir=struct_dir,
+                        msa_dir=msa_dir,
+                        moldir=moldir,
+                        prob=data_config.prob,
+                        cropper=data_config.cropper,
+                        tokenizer=cfg.tokenizer,
+                        featurizer=cfg.featurizer,
+                        val_group=data_config.val_group,
+                        symmetry_correction=data_config.symmetry_correction,
+                        override_bfactor=data_config.override_bfactor,
+                        override_method=data_config.override_method,
+                        selector=cfg.selector,
+                    )
+                )
+
+            # Create validation dataset
+            if len(val_samples) > 0:
+                val.append(
+                    Dataset(
+                        samples=val_samples,
+                        record_dir=record_dir,
+                        struct_dir=struct_dir,
+                        msa_dir=msa_dir,
+                        moldir=moldir,
+                        prob=data_config.prob,
+                        cropper=data_config.cropper,
+                        tokenizer=cfg.tokenizer,
+                        featurizer=cfg.featurizer,
+                        val_group=data_config.val_group,
+                        symmetry_correction=data_config.symmetry_correction,
+                        selector=cfg.selector,
+                    )
+                )
+
+        # Print dataset sizes
+        for dataset in train:
+            dataset: Dataset
+            print(f"Training dataset size: {len(dataset.samples)}")
+
+        self.val_group_mapper = defaultdict(dict)
+
+        for i, dataset in enumerate(train if cfg.overfit is not None else val):
+            dataset: Dataset
+            print(f"Validation dataset size: {len(dataset.samples)}")
+            self.val_group_mapper[i]["label"] = dataset.val_group
+            self.val_group_mapper[i]["symmetry_correction"] = (
+                # If overfit, use symmetry_correction from val dataset instead of training dataset
+                dataset.symmetry_correction
+                if cfg.overfit is None
+                else data_config.symmetry_correction
+            )
+
+        # Load canonical molecules
+        canonicals = load_canonicals(cfg.moldir)
+
+        # Create wrapper datasets
+        self._train_set = TrainingDataset(
+            datasets=train,
+            canonicals=canonicals,
+            moldir=cfg.moldir,
+            samples_per_epoch=cfg.samples_per_epoch,
+            max_atoms=cfg.max_atoms,
+            max_tokens=cfg.max_tokens,
+            max_seqs=cfg.max_seqs,
+            pad_to_max_atoms=cfg.pad_to_max_atoms,
+            pad_to_max_tokens=cfg.pad_to_max_tokens,
+            pad_to_max_seqs=cfg.pad_to_max_seqs,
+            atoms_per_window_queries=cfg.atoms_per_window_queries,
+            min_dist=cfg.min_dist,
+            max_dist=cfg.max_dist,
+            num_bins=cfg.num_bins,
+            overfit=cfg.overfit,
+            return_symmetries=cfg.return_train_symmetries,
+            single_sequence_prop=cfg.single_sequence_prop_training,
+            msa_sampling=cfg.msa_sampling_training,
+            use_msa=cfg.use_msa,
+            compute_frames=cfg.compute_frames,
+            backbone_only=cfg.backbone_only,
+            atom14=cfg.atom14,
+            atom37=cfg.atom37,
+            design=cfg.design,
+            inverse_fold=cfg.inverse_fold,
+            disulfide_prob=cfg.disulfide_prob,
+            disulfide_on=cfg.disulfide_on,
+        )
+        self._val_set = ValidationDataset(
+            datasets=train if cfg.overfit is not None else val,
+            canonicals=canonicals,
+            moldir=cfg.moldir,
+            seed=cfg.random_seed,
+            max_atoms=cfg.max_atoms,
+            max_tokens=cfg.max_tokens,
+            max_seqs=cfg.max_seqs,
+            pad_to_max_atoms=cfg.pad_to_max_atoms,
+            pad_to_max_tokens=cfg.pad_to_max_tokens,
+            pad_to_max_seqs=cfg.pad_to_max_seqs,
+            atoms_per_window_queries=cfg.atoms_per_window_queries,
+            min_dist=cfg.min_dist,
+            max_dist=cfg.max_dist,
+            num_bins=cfg.num_bins,
+            overfit=cfg.overfit,
+            return_symmetries=cfg.return_val_symmetries,
+            compute_frames=cfg.compute_frames,
+            backbone_only=cfg.backbone_only,
+            atom14=cfg.atom14,
+            atom37=cfg.atom37,
+            design=cfg.design,
+            inverse_fold=cfg.inverse_fold,
+            disulfide_prob=cfg.disulfide_prob,
+            disulfide_on=cfg.disulfide_on,
+        )
+
+        self.monomer_split = cfg.monomer_split
+        print("monomer_split", self.monomer_split)
+        if self.monomer_split is not None:
+            with Path(self.monomer_split).open("r") as f:
+                monomer_ids = [x.lower() for x in f.read().splitlines()]
+                print("monomer_split", monomer_ids)
+
+            dataset = data_protein_binder.Dataset(
+                struct_dir=Path(cfg.monomer_target_dir) / "structures",
+                record_dir=Path(cfg.monomer_target_dir) / "records",
+                target_ids=monomer_ids,
+                seq_len=cfg.monomer_seq_len,
+                tokenizer=cfg.tokenizer,
+                featurizer=cfg.featurizer,
+            )
+
+            # Load canonical molecules
+            canonicals = load_canonicals(cfg.moldir)
+
+            self.monomer_val_set = data_protein_binder.PredictionDataset(
+                dataset=dataset,
+                canonicals=canonicals,
+                moldir=Path(cfg.moldir),
+                backbone_only=cfg.backbone_only,
+                atom14=cfg.atom14,
+                atom37=cfg.atom37,
+                design=cfg.design,
+                target_structure_condition=cfg.monomer_target_structure_condition,
+                inverse_fold=cfg.inverse_fold,
+                disulfide_prob=cfg.disulfide_prob,
+                disulfide_on=cfg.disulfide_on,
+            )
+
+        self.ligand_split = cfg.ligand_split
+        print("ligand_split", self.ligand_split)
+        if self.ligand_split is not None:
+            with Path(self.ligand_split).open("r") as f:
+                ligand_ids = [x.lower() for x in f.read().splitlines()]
+                print("ligand_split", ligand_ids)
+
+            dataset = data_ligands.Dataset(
+                struct_dir=Path(cfg.ligand_target_dir) / "structures",
+                record_dir=Path(cfg.ligand_target_dir) / "records",
+                target_ids=ligand_ids,
+                min_len=cfg.ligand_seq_len,
+                max_len=cfg.ligand_seq_len,
+                tokenizer=cfg.tokenizer,
+                featurizer=cfg.featurizer,
+            )
+
+            # Load canonical molecules
+            canonicals = load_canonicals(cfg.moldir)
+
+            self.ligand_val_set = data_ligands.PredictionDataset(
+                dataset=dataset,
+                canonicals=canonicals,
+                moldir=Path(cfg.moldir),
+                backbone_only=cfg.backbone_only,
+                atom14=cfg.atom14,
+                atom37=cfg.atom37,
+                design=cfg.design,
+                disulfide_prob=cfg.disulfide_prob,
+                disulfide_on=cfg.disulfide_on,
+            )
+
+    def setup(self, stage: Optional[str] = None) -> None:  # noqa: ARG002 (unused)
+        """Run the setup for the DataModule.
+
+        Parameters
+        ----------
+        stage : str, optional
+            The stage, one of 'fit', 'validate', 'test'.
+
+        """
+        return
+
+    def train_dataloader(self) -> DataLoader:
+        """Get the training dataloader.
+
+        Returns
+        -------
+        DataLoader
+            The training dataloader.
+
+        """
+        return DataLoader(
+            self._train_set,
+            batch_size=self.cfg.batch_size,
+            num_workers=self.cfg.num_workers,
+            pin_memory=self.cfg.pin_memory,
+            shuffle=False,
+            collate_fn=collate,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        """Get the validation dataloader.
+
+        Returns
+        -------
+        DataLoader
+            The validation dataloader.s
+
+        """
+        val_loaders = []
+        val_loaders.append(
+            DataLoader(
+                self._val_set,
+                batch_size=self.cfg.val_batch_size,
+                num_workers=self.cfg.num_workers if not self.inverse_fold else 1,
+                pin_memory=self.cfg.num_workers if not self.inverse_fold else False,
+                shuffle=False,
+                collate_fn=collate,
+            )
+        )
+        if self.monomer_split is not None:
+            val_loaders.append(
+                DataLoader(
+                    self.monomer_val_set,
+                    batch_size=self.cfg.val_batch_size,
+                    num_workers=self.cfg.num_workers if not self.inverse_fold else 1,
+                    pin_memory=self.cfg.pin_memory if not self.inverse_fold else False,
+                    shuffle=False,
+                    collate_fn=data_protein_binder.collate,
+                )
+            )
+        if self.ligand_split is not None:
+            val_loaders.append(
+                DataLoader(
+                    self.ligand_val_set,
+                    batch_size=self.cfg.val_batch_size,
+                    num_workers=self.cfg.num_workers if not self.inverse_fold else 1,
+                    pin_memory=self.cfg.pin_memory if not self.inverse_fold else False,
+                    shuffle=False,
+                    collate_fn=data_ligands.collate,
+                )
+            )
+        return val_loaders
+
+    def predict_dataloader(self) -> DataLoader:
+        return DataLoader(
+            self._val_set,
+            batch_size=self.cfg.val_batch_size,
+            num_workers=self.cfg.num_workers,
+            pin_memory=self.cfg.pin_memory,
+            shuffle=False,
+            collate_fn=collate,
+        )
+
+    def transfer_batch_to_device(
+        self,
+        batch: Dict,
+        device: torch.device,
+        dataloader_idx: int,  # noqa: ARG002
+    ) -> Dict:
+        """Transfer a batch to the given device.
+
+        Parameters
+        ----------
+        batch : Dict
+            The batch to transfer.
+        device : torch.device
+            The device to transfer to.
+        dataloader_idx : int
+            The dataloader index.
+
+        Returns
+        -------
+        np.Any
+            The transferred batch.
+
+        """
+        for key in batch:
+            if key not in [
+                "all_coords",
+                "all_resolved_mask",
+                "crop_to_all_atom_map",
+                "chain_symmetries",
+                "chain_swaps",
+                "amino_acids_symmetries",
+                "ligand_symmetries",
+                "activity_name",
+                "activity_qualifier",
+                "sid",
+                "cid",
+                "normalized_protein_accession",
+                "pair_id",
+                "ligand_edge_index",
+                "ligand_edge_lower_bounds",
+                "ligand_edge_upper_bounds",
+                "ligand_edge_bond_mask",
+                "ligand_edge_angle_mask",
+                "connections_edge_index",
+                "ligand_chiral_atom_index",
+                "ligand_chiral_check_mask",
+                "ligand_chiral_atom_orientations",
+                "ligand_stereo_bond_index",
+                "ligand_stereo_check_mask",
+                "ligand_stereo_bond_orientations",
+                "ligand_aromatic_5_ring_index",
+                "ligand_aromatic_6_ring_index",
+                "ligand_planar_double_bond_index",
+                "pdb_id",
+                "id",
+                "tokenized",
+                "structure",
+                "structure_bonds",
+                "extra_mols",
+            ]:
+                if hasattr(batch[key], "to"):
+                    batch[key] = batch[key].to(device)
+        return batch

model/boltzgen/task/train/train.py

@@ -0,0 +1,250 @@
+import datetime
+import os
+import warnings
+from pathlib import Path
+import time
+from typing import Optional
+import pytorch_lightning as pl
+import torch
+from omegaconf import OmegaConf, listconfig
+from pytorch_lightning import LightningModule
+from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
+from pytorch_lightning.loggers import WandbLogger
+from pytorch_lightning.strategies import DDPStrategy
+
+from boltzgen.task.task import Task
+from boltzgen.task.train.data import DataConfig, TrainingDataModule
+
+
+class Training(Task):
+    """Training configuration."""
+
+    def __init__(
+        self,
+        data: DataConfig,
+        model: LightningModule,
+        output: str,
+        name: str,
+        trainer: Optional[dict] = None,
+        resume: Optional[str] = None,
+        pretrained: Optional[str] = None,
+        wandb: Optional[dict] = None,
+        disable_checkpoint: bool = False,
+        slurm: bool = False,
+        matmul_precision: Optional[str] = None,
+        find_unused_parameters: Optional[bool] = False,
+        ddp_timeout_seconds: Optional[int] = 6400,
+        save_every_n_train_steps: Optional[int] = None,
+        save_top_k: Optional[int] = 1,
+        validation_only: bool = False,
+        debug: bool = False,
+        strict_loading: bool = True,
+        metric_mode: str = "max",
+    ) -> None:
+        """Initialize training configuration.
+
+        Parameters
+        ----------
+        data : DataConfig
+            The data configuration.
+        model : LightningModule
+            The model.
+        output : str
+            The output directory.
+        name : str
+            The output subdirectory and wandb run name.
+        trainer : Optional[dict], optional
+            The trainer configuration, by default None.
+        resume : Optional[str], optional
+            The resume checkpoint, by default None
+        pretrained : Optional[str], optional
+            The pretrained model, by default None
+        wandb : Optional[dict], optional
+            The wandb configuration, by default None
+        disable_checkpoint : bool, optional
+            Disable checkpoint, by default False
+
+        """
+        if not isinstance(data, DataConfig):
+            data = DataConfig(**data)
+
+        self.data = data
+        self.model = model
+        self.output = output
+        self.name = name
+        self.trainer = trainer
+        self.resume = resume
+        self.pretrained = pretrained
+        self.wandb = wandb
+        self.disable_checkpoint = disable_checkpoint
+        self.slurm = slurm
+        self.matmul_precision = matmul_precision
+        self.find_unused_parameters = find_unused_parameters
+        self.ddp_timeout_seconds = ddp_timeout_seconds
+        self.save_every_n_train_steps = save_every_n_train_steps
+        self.save_top_k = save_top_k
+        self.validation_only = validation_only
+        self.debug = debug
+        self.strict_loading = strict_loading
+        self.metric_mode = metric_mode
+
+    def run(self, config: OmegaConf) -> None:
+        """Run training.
+
+        Parameters
+        ----------
+        config : OmegaConf
+            The configuration for the task, for bookkeeping.
+
+        """
+        # Disable some warnings
+        warnings.filterwarnings(
+            "ignore", ".*when logging on epoch level in distributed setting.*"
+        )
+
+        # Experiment with this during training (high or medium)
+        if self.matmul_precision is not None:
+            torch.set_float32_matmul_precision(self.matmul_precision)
+
+        # Create trainer dict
+        if self.trainer is None:
+            self.trainer = {}
+
+        # Flip some arguments in debug mode
+        devices = self.trainer.get("devices", 1)
+
+        if self.debug:
+            self.slurm = False
+            if isinstance(devices, int):
+                devices = 1
+            elif isinstance(devices, (list, listconfig.ListConfig)):
+                devices = [devices[0]]
+            self.trainer["devices"] = devices
+            self.data.num_workers = 0
+            if self.wandb:
+                self.wandb = None
+
+        # slurm
+        if self.slurm:
+            self.trainer["devices"] = int(
+                os.environ.get("SLURM_NTASKS_PER_NODE", "auto")
+            )
+            self.trainer["num_nodes"] = int(os.environ.get("SLURM_NNODES", 1))
+
+        # Create objects
+        data_module = TrainingDataModule(self.data)
+        model_module = self.model
+
+        if self.pretrained and not self.resume:
+            file_path = self.pretrained
+
+            print(f"Loading model from {file_path}")
+            model_module = type(model_module).load_from_checkpoint(
+                file_path, map_location="cpu", strict=False, weights_only=False, **(model_module.hparams)
+            )
+
+        # Create checkpoint callback
+        callbacks = self.trainer.get("callbacks", [])
+        dirpath = f"{self.output}/{self.name}"
+        os.makedirs(dirpath, exist_ok=True)
+
+        if not self.disable_checkpoint:
+            if self.slurm:
+                jobid = os.environ.get("SLURM_JOB_ID", "")
+                dirpath = Path(self.output) / jobid
+                dirpath.mkdir(parents=True, exist_ok=True)
+                print(
+                    "Configuring ModelCheckpoint for SLURM to directory: ", str(dirpath)
+                )
+            mc = ModelCheckpoint(
+                monitor="val/lddt",
+                save_top_k=self.save_top_k,
+                save_last=True,
+                mode=self.metric_mode,
+                every_n_epochs=1,
+            )
+            callbacks.append(mc)
+
+        # Create wandb logger
+        loggers = []
+        if self.wandb:
+            extra = {}
+            if self.slurm:
+                ckpt_path = (
+                    "hpc"
+                    if len([f for f in os.listdir(dirpath) if f.startswith("hpc")])
+                    else None
+                )
+                extra = {"resume": "must"} if ckpt_path else {}
+
+            wdb_logger = WandbLogger(
+                name=self.name,
+                group=self.wandb["group"],
+                save_dir=dirpath,
+                dir=f".tmp_wandb/{time.time()}",
+                project=self.wandb["project"],
+                entity=self.wandb["entity"],
+                log_model=False,
+                id=os.environ.get("SLURM_JOB_ID", None) if self.slurm else None,
+                **extra,
+            )
+            loggers.append(wdb_logger)
+            # Save the config to wandb
+            # This fails when not on rank0, so just catching
+            try:
+                config_out = Path(wdb_logger.experiment.dir) / "run.yaml"
+                with Path.open(config_out, "w") as f:
+                    OmegaConf.save(config, f)
+                wdb_logger.experiment.save(str(config_out))
+            except Exception:  # noqa: BLE001, S110
+                pass
+
+        # Set up trainer
+        strategy = "auto"
+        if (isinstance(devices, int) and devices > 1) or (
+            isinstance(devices, (list, listconfig.ListConfig)) and len(devices) > 1
+        ):
+            strategy = DDPStrategy(
+                find_unused_parameters=self.find_unused_parameters,
+                timeout=datetime.timedelta(seconds=self.ddp_timeout_seconds),
+            )
+
+        trainer = pl.Trainer(
+            default_root_dir=str(dirpath),
+            strategy=strategy,
+            callbacks=callbacks,
+            logger=loggers,
+            enable_checkpointing=not self.disable_checkpoint,
+            reload_dataloaders_every_n_epochs=1,
+            **self.trainer,
+        )
+
+        # Run training
+        if self.slurm:
+            # If we are on SLURM, we need to check if we have a checkpoint for this job
+            ckpt_path = (
+                "hpc"
+                if len([f for f in os.listdir(dirpath) if f.startswith("hpc")])
+                else None
+            )
+            # If none check with we have a resume checkpoint passed in
+            if ckpt_path is None:
+                ckpt_path = self.resume
+        else:
+            ckpt_path = self.resume
+
+        if not self.strict_loading:
+            model_module.strict_loading = False
+
+        if self.validation_only:
+            trainer.validate(
+                model_module,
+                datamodule=data_module,
+                ckpt_path=ckpt_path,
+            )
+        else:
+            trainer.fit(
+                model_module,
+                datamodule=data_module,
+                ckpt_path=ckpt_path,
+            )

model/boltzgen/utils/pipeline_progress_bar.py

@@ -0,0 +1,171 @@
+"""
+Custom PyTorch Lightning progress bar callback that displays pipeline step information.
+"""
+import os
+
+from pytorch_lightning.callbacks import TQDMProgressBar
+from tqdm import tqdm
+
+
+class PipelineProgressBar(TQDMProgressBar):
+    """Custom progress bar that shows the current pipeline step."""
+    
+    def __init__(self, refresh_rate: int = 1, process_position: int = 0):
+        """Initialize the pipeline progress bar.
+        
+        Parameters
+        ----------
+        refresh_rate : int
+            Refresh rate for the progress bar
+        process_position : int
+            Position of the process for multi-processing
+        """
+        super().__init__(refresh_rate=refresh_rate, process_position=process_position)
+        self._original_tqdm = None
+        self._patch_tqdm()
+    
+    def _get_pipeline_info(self) -> str:
+        """Get pipeline step information from environment variables."""
+        pipeline_step = os.environ.get("BOLTZGEN_PIPELINE_STEP", "")
+        pipeline_progress = os.environ.get("BOLTZGEN_PIPELINE_PROGRESS", "")
+        
+        if pipeline_step:
+            if pipeline_progress:
+                return f"[{pipeline_progress}] {pipeline_step}"
+            else:
+                return f"[Pipeline] {pipeline_step}"
+        return ""
+    
+    def _update_bar_description(self, bar):
+        """Helper method to update a progress bar description with pipeline info."""
+        if bar is not None:
+            pipeline_info = self._get_pipeline_info()
+            if pipeline_info:
+                # Update the description to include pipeline step info
+                current_desc = getattr(bar, 'desc', '') or ''
+                if current_desc:
+                    new_desc = f"{pipeline_info} - {current_desc}"
+                    bar.set_description(new_desc)
+                else:
+                    bar.set_description(pipeline_info)
+        return bar
+    
+    def init_predict_tqdm(self) -> None:
+        """Initialize the prediction progress bar."""
+        bar = super().init_predict_tqdm()
+        return self._update_bar_description(bar)
+    
+    def init_train_tqdm(self) -> None:
+        """Initialize the training progress bar."""
+        bar = super().init_train_tqdm()
+        return self._update_bar_description(bar)
+    
+    def init_validation_tqdm(self) -> None:
+        """Initialize the validation progress bar."""
+        bar = super().init_validation_tqdm()
+        return self._update_bar_description(bar)
+    
+    def init_test_tqdm(self) -> None:
+        """Initialize the test progress bar."""
+        bar = super().init_test_tqdm()
+        return self._update_bar_description(bar)
+
+    def init_sanity_tqdm(self) -> None:
+        """Initialize the sanity check progress bar."""
+        bar = super().init_sanity_tqdm()
+        return self._update_bar_description(bar)
+    
+    def on_predict_start(self, trainer, pl_module):
+        """Called when prediction starts - update all progress bars."""
+        super().on_predict_start(trainer, pl_module)
+        self._update_all_progress_bars()
+    
+    def on_predict_batch_start(self, trainer, pl_module, batch, batch_idx):
+        """Called before each prediction batch - update progress bars."""
+        super().on_predict_batch_start(trainer, pl_module, batch, batch_idx)
+        self._update_all_progress_bars()
+    
+    def _update_all_progress_bars(self):
+        """Update all existing progress bars with pipeline info."""
+        pipeline_info = self._get_pipeline_info()
+        if not pipeline_info:
+            return
+            
+        # Try to find and update all progress bars
+        for attr_name in dir(self):
+            if ('progress' in attr_name.lower() or 
+                'tqdm' in attr_name.lower() or 
+                attr_name.endswith('_bar')):
+                try:
+                    bar = getattr(self, attr_name)
+                    if (bar is not None and 
+                        hasattr(bar, 'set_description') and 
+                        hasattr(bar, 'desc')):
+                        current_desc = getattr(bar, 'desc', '') or ''
+                        if pipeline_info not in current_desc:
+                            if current_desc:
+                                new_desc = f"{pipeline_info} - {current_desc}"
+                            else:
+                                new_desc = pipeline_info
+                            bar.set_description(new_desc)
+                except (AttributeError, TypeError):
+                    pass
+    
+    def print(self, *args, **kwargs):
+        """Override print method to intercept and modify all progress updates."""
+        # Check if this is updating a progress bar description and modify it
+        if args:
+            # Try to update any active progress bars before printing
+            self._update_all_progress_bars()
+        
+        return super().print(*args, **kwargs)
+    
+    def get_metrics(self, trainer, pl_module):
+        """Override get_metrics to update progress bars on each metric update."""
+        metrics = super().get_metrics(trainer, pl_module)
+        self._update_all_progress_bars()
+        return metrics
+    
+    def _patch_tqdm(self):
+        """Patch tqdm to automatically add pipeline info to all progress bars."""
+        if self._original_tqdm is not None:
+            return  # Already patched
+            
+        # Store the original tqdm class
+        self._original_tqdm = tqdm
+        
+        # Create a wrapper that adds pipeline info
+        def patched_tqdm(*args, **kwargs):
+            # Create the original tqdm instance
+            instance = self._original_tqdm(*args, **kwargs)
+            
+            # Update its description with pipeline info
+            pipeline_info = self._get_pipeline_info()
+            if pipeline_info:
+                current_desc = kwargs.get('desc', '') or getattr(instance, 'desc', '') or ''
+                if pipeline_info not in current_desc:
+                    if current_desc and current_desc.strip():
+                        # Remove any trailing colons from current description to avoid double colons
+                        current_desc = current_desc.rstrip(':').strip()
+                        new_desc = f"{pipeline_info} - {current_desc}"
+                    else:
+                        new_desc = pipeline_info
+                    instance.set_description(new_desc)
+            
+            return instance
+        
+        # Monkey patch tqdm globally during our callback's lifetime
+        import tqdm as tqdm_module
+        tqdm_module.tqdm = patched_tqdm
+        
+    def _unpatch_tqdm(self):
+        """Restore original tqdm."""
+        if self._original_tqdm is not None:
+            import tqdm as tqdm_module
+            tqdm_module.tqdm = self._original_tqdm
+            self._original_tqdm = None
+    
+    def teardown(self, trainer, pl_module, stage):
+        """Called when the callback is being torn down."""
+        super().teardown(trainer, pl_module, stage)
+        self._unpatch_tqdm()

model/boltzgen/utils/quiet.py

@@ -0,0 +1,16 @@
+import logging, warnings
+
+
+def quiet_startup() -> None:
+    warnings.filterwarnings("ignore", message=r".*predict_dataloader.*num_workers.*")
+    warnings.filterwarnings(
+        "ignore", message=r".*tensorboardX.*removed as a dependency.*"
+    )
+    warnings.filterwarnings(
+        "ignore",
+        message=r"The pynvml package is deprecated",
+        category=FutureWarning,
+        module=r"torch\.cuda",
+    )
+
+    logging.getLogger("pytorch_lightning").setLevel(logging.ERROR)

requirements.txt

@@ -0,0 +1,53 @@
+# BoltzGen 0.3.2 on OneScience / Python 3.11 / DTK 25.04 / PyTorch 2.5.1
+# Install only after OneScience is installed:
+#   python -m pip install --no-deps -r requirements.txt
+
+# Additional dependencies absent from the OneScience base inventory.
+hydride
+pydssp
+logomaker
+frozendict
+
+# Compatibility upgrade validated on SCNet. Lightning 1.8.6 rejects the
+# precision=bf16-mixed value used by BoltzGen. Install with --no-deps so this
+# upgrade does not replace the OneScience/DTK PyTorch build.
+pytorch-lightning==2.5.6
+
+# Supplied by OneScience; retained as the complete dependency inventory.
+# Do not reinstall these packages over OneScience/DTK builds.
+# torch==2.5.1                    # DTK/DCU build supplied by OneScience
+# numpy==1.26.3                   # validated despite upstream pinning 2.0.2
+# numba
+# matplotlib
+# biotite==1.0.1
+# gemmi
+# mashumaro==3.16
+# einx==0.3.0
+# einops>=0.7.0
+# scikit-learn>=1.2.2,<=1.6.0
+# hydra-core>=1.2.0
+# omegaconf>=2.3.0
+# rdkit>=2025.3.5
+# pandas>=2.2.2
+# biopython==1.84
+# huggingface_hub
+# pyyaml>=6.0
+# scipy==1.14.1
+# tqdm>=4.60.0
+# torchmetrics
+# pytorch-lightning==2.0.6        # OneScience inventory; upgraded above for BoltzGen
+# wandb>=0.13.7                 # training logging; disabled in smoke tests
+# redis==7.0.0
+# nvidia-ml-py>=12.535.133      # NVIDIA monitoring only; not needed on DCU
+
+# Declared upstream but not imported by the validated paths.
+# edit-distance
+# pdbeccdutils
+
+# Optional: antibody/nanobody CDR logo generation during filtering only.
+# abnumber
+
+# NVIDIA CUDA 12-only extensions: intentionally disabled on DCU.
+# cuequivariance_ops_cu12>=0.5.0
+# cuequivariance_ops_torch_cu12>=0.5.0
+# cuequivariance_torch>=0.5.0

scripts/boltzgen.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python3
+"""Run the bundled BoltzGen CLI without installing the source package."""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parents[1]
+MODEL_DIR = ROOT / "model"
+sys.path.insert(0, str(MODEL_DIR))
+os.environ["PYTHONPATH"] = os.pathsep.join(
+    [str(MODEL_DIR), os.environ.get("PYTHONPATH", "")]
+).rstrip(os.pathsep)
+
+from boltzgen.cli.boltzgen import main  # noqa: E402
+
+
+if __name__ == "__main__":
+    main()

scripts/check_environment.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+"""Report core OneScience/DTK and BoltzGen dependency versions."""
+
+from __future__ import annotations
+
+import importlib
+import sys
+from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / "model"))
+
+
+def package_version(distribution: str) -> str:
+    try:
+        return version(distribution)
+    except PackageNotFoundError:
+        return "NOT INSTALLED"
+
+
+def main() -> int:
+    print("python:", sys.version.split()[0])
+    try:
+        import torch
+
+        print("torch:", torch.__version__)
+        print("HIP:", torch.version.hip)
+        print("DCU available:", torch.cuda.is_available())
+        print("device count:", torch.cuda.device_count())
+        if torch.cuda.is_available() and torch.cuda.device_count():
+            print("device 0:", torch.cuda.get_device_name(0))
+    except Exception as exc:  # noqa: BLE001
+        print("torch check failed:", repr(exc))
+        return 1
+
+    for distribution in (
+        "pytorch-lightning", "torchmetrics", "hydra-core", "omegaconf",
+        "biotite", "gemmi", "mashumaro", "einx", "hydride", "pydssp",
+        "logomaker", "frozendict",
+    ):
+        print(f"{distribution}:", package_version(distribution))
+
+    try:
+        importlib.import_module("boltzgen.cli.boltzgen")
+        print("boltzgen import: OK")
+    except Exception as exc:  # noqa: BLE001
+        print("boltzgen import failed:", repr(exc))
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/slurm-example/README.md

@@ -0,0 +1,28 @@
+slurm-example
+=============
+
+This directory contains an example for running BoltzGen on SLURM.
+
+The approach we use is to launch a job array that runs many single-GPU jobs. Each job calls "boltzgen run" to go through all pipeline steps (design generation, refolding, etc.) for a small number of designs. After all the jobs have completed, on a login node we use "boltzgen merge" to merge all individual task results together, then "boltzgen run --steps filtering" to apply
+filters to the combined set.
+
+Note that other parallelization strategies could also work here. In particular, the individual BoltzGen pipeline steps (e.g. design generation, refolding) support
+parallelization across multiple GPUs within the step.
+
+You will need to modify [run.sh](run.sh) and probably also [run_job_array.slurm](run_job_array.slurm) for your site and analysis task.
+
+After you've made your modifications, you can run:
+
+```
+$ bash run.sh submit
+```
+
+To submit the job array.
+
+After that has finished, you can run:
+
+```
+$ bash run.sh process
+```
+
+to merge and filter the results.

scripts/slurm-example/run.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+set -e
+set -x
+
+DESIGN_SPEC=/path/to/design/spec.yaml
+MERGED_OUT=/path/to/your/output/directory
+NUM_TASKS=20
+NUM_DESIGNS_PER_TASK=1000
+CONDA_ENVIRONMENT=/path/to/conda/environment/with/boltzgen
+CACHE_DIR=/path/where/models/are/saved
+ACCOUNT=FILL_THIS_IN
+TIME=05:00:00
+
+OUT="${MERGED_OUT}/task-outputs"
+LOGS="${MERGED_OUT}/task-logs"
+
+if [[ $# -ne 1 ]]; then
+    echo "Usage: $0 {submit|process}"
+    exit 1
+fi
+
+
+MODE="$1"
+
+if [[ "$MODE" == "submit" ]]; then
+
+    mkdir -p "$OUT"
+    mkdir -p "$LOGS"
+
+    sbatch -A "$ACCOUNT" -t "$TIME" --export=ALL --array=1-$NUM_TASKS -o $LOGS/stdout.%A-%a.log -e $LOGS/stderr.%A-%a.log run_job_array.slurm \
+        "$DESIGN_SPEC" "$OUT" "$NUM_DESIGNS_PER_TASK" "$CONDA_ENVIRONMENT" --protocol protein-anything --cache "$CACHE_DIR"
+    squeue --me
+
+elif [[ "$MODE" == "process" ]]; then
+
+    boltzgen merge "$OUT"/task-* --output "$MERGED_OUT"
+    boltzgen run "$DESIGN_SPEC" --steps filtering --protocol protein-anything --output "$MERGED_OUT"
+
+else
+    echo "Usage: $0 {submit|process}"
+    exit 1
+fi

scripts/slurm-example/run_job_array.slurm

@@ -0,0 +1,53 @@
+#!/bin/bash
+#SBATCH -C gpu                    # Use GPU nodes
+#SBATCH -q regular                # Queue/partition
+#SBATCH -t 03:00:00               # Walltime per array task
+#SBATCH -N 1                      # Number of nodes per task
+#SBATCH --gpus-per-node=1         # 1 GPU per task
+#SBATCH --array=1-2               # Job array indices. Override this.
+#SBATCH --mem=64G
+#SBATCH -J boltzgen
+
+set -euo pipefail
+set -x
+
+echo "Starting job on host: $(hostname)"
+echo "SLURM job ID:        ${SLURM_JOB_ID:-unknown}"
+echo "Array job ID:        ${SLURM_ARRAY_JOB_ID:-unknown}"
+echo "Array task ID:       ${SLURM_ARRAY_TASK_ID:-unknown}"
+echo "Nodes:               ${SLURM_JOB_NUM_NODES}"
+
+if [[ $# -lt 3 ]]; then
+    echo "Usage: $0 <design_spec> <outdir> <num_designs_per_job> <conda_environment>" >&2
+    exit 1
+fi
+
+design_spec="$1"
+outdir="$2"
+num_designs="${3}"
+conda_environment="${4}"
+shift 4                 # Remove parsed args
+extra_args=("$@")       # Store everything else
+
+echo "Design spec: $design_spec"
+echo "Output dir:  $outdir"
+echo "Num designs: $num_designs"
+echo "Conda environment: $conda_environment"
+
+module load conda
+conda activate "$conda_environment"
+
+which python
+which boltzgen
+nvidia-smi
+
+job_outdir="${outdir}/task-${SLURM_ARRAY_JOB_ID}-${SLURM_ARRAY_TASK_ID}"
+mkdir -p "${job_outdir}"
+
+srun --ntasks=1 --gpus-per-task=1 \
+  boltzgen run "$design_spec" \
+    --output "${job_outdir}" \
+    --num_designs "$num_designs" \
+    "${extra_args[@]}"
+
+

scripts/train.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+"""Run an official BoltzGen Hydra training configuration."""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parents[1]
+MODEL_DIR = ROOT / "model"
+sys.path.insert(0, str(MODEL_DIR))
+os.environ["PYTHONPATH"] = os.pathsep.join(
+    [str(MODEL_DIR), os.environ.get("PYTHONPATH", "")]
+).rstrip(os.pathsep)
+
+from boltzgen.resources.main import main  # noqa: E402
+
+
+def run() -> None:
+    if len(sys.argv) < 2:
+        raise SystemExit("Usage: python scripts/train.py CONFIG.yaml [key=value ...]")
+    main(sys.argv[1], sys.argv[2:])
+
+
+if __name__ == "__main__":
+    run()