AlphaPulldown/test/outdated/test_chopped_dimer_error.py

#!/usr/bin/env python
"""
Test to reproduce the CIF parsing error in chopped dimer predictions.

The error occurs when AlphaFold3 tries to parse template CIF files that don't have
the required '_atom_site.pdbx_PDB_model_num' field.
"""
import pytest
import pickle
import tempfile
import subprocess
import os
from pathlib import Path
from unittest.mock import patch, MagicMock

from alphapulldown.objects import ChoppedObject, MultimericObject
from alphapulldown.folding_backend.alphafold3_backend import _convert_to_fold_input


class TestChoppedDimerError:
    """Test class to reproduce the CIF parsing error in chopped dimer predictions."""
    
    @pytest.fixture(scope="class")
    def test_data_dir(self):
        """Get the test data directory."""
        return Path(__file__).parent / "test_data"
    
    @pytest.fixture(scope="class")
    def monomer_obj(self, test_data_dir):
        """Load a monomer object from test data."""
        pkl_path = test_data_dir / "features" / "A0A075B6L2.pkl"
        with open(pkl_path, "rb") as f:
            return pickle.load(f)
    
    @pytest.fixture
    def chopped_objects(self, monomer_obj):
        """Create two chopped objects for dimer testing."""
        # Create first chopped object: residues 1-10
        co1 = ChoppedObject(
            description=monomer_obj.description,
            sequence=monomer_obj.sequence,
            feature_dict=monomer_obj.feature_dict,
            regions=[(0, 10)]
        )
        co1.prepare_final_sliced_feature_dict()
        
        # Create second chopped object: residues 11-20
        co2 = ChoppedObject(
            description=monomer_obj.description,
            sequence=monomer_obj.sequence,
            feature_dict=monomer_obj.feature_dict,
            regions=[(10, 20)]
        )
        co2.prepare_final_sliced_feature_dict()
        
        return co1, co2
    
    def test_chopped_dimer_conversion(self, chopped_objects):
        """Test that chopped dimer conversion works without CIF parsing errors."""
        co1, co2 = chopped_objects
        
        # Create multimeric object
        multi = MultimericObject(interactors=[co1, co2], pair_msa=True)
        
        # Convert to fold input - this should not raise the CIF parsing error
        try:
            fold_input = _convert_to_fold_input(multi, random_seed=0)
            
            # Basic assertions
            assert len(fold_input.chains) == 2
            assert fold_input.chains[0].id == "A"
            assert fold_input.chains[1].id == "B"
            
            # Check sequences
            assert len(fold_input.chains[0].sequence) == 10
            assert len(fold_input.chains[1].sequence) == 10
            
            print(f"✓ Successfully converted chopped dimer to fold input")
            print(f"  Chain A sequence: {fold_input.chains[0].sequence}")
            print(f"  Chain B sequence: {fold_input.chains[1].sequence}")
            
        except Exception as e:
            pytest.fail(f"Failed to convert chopped dimer to fold input: {e}")
    
    def test_template_cif_parsing_error_simulation(self, chopped_objects):
        """Test that simulates the CIF parsing error by mocking the AlphaFold3 structure parsing."""
        co1, co2 = chopped_objects
        
        # Mock the AlphaFold3 structure parsing to simulate the error
        with patch('alphafold3.structure.parsing.from_mmcif') as mock_from_mmcif:
            # Simulate the KeyError that occurs in AlphaFold3's CIF parsing
            mock_from_mmcif.side_effect = KeyError("'_atom_site.pdbx_PDB_model_num'")
            
            # Create multimeric object
            multi = MultimericObject(interactors=[co1, co2], pair_msa=True)
            
            # Convert to fold input - this should work fine
            fold_input = _convert_to_fold_input(multi, random_seed=0)
            
            # The error would occur later when AlphaFold3 tries to process templates
            # during the prediction phase, not during the conversion phase
            print("✓ Conversion phase completed successfully")
            print("⚠️  The CIF parsing error would occur during the prediction phase")
            print("   when AlphaFold3 tries to parse template CIF files")
    
    def test_chopped_dimer_without_templates(self, chopped_objects):
        """Test chopped dimer without templates to avoid CIF parsing issues."""
        co1, co2 = chopped_objects
        
        # Remove template features from both chopped objects
        for co in [co1, co2]:
            template_keys = [k for k in co.feature_dict.keys() if k.startswith('template_')]
            for key in template_keys:
                del co.feature_dict[key]
        
        # Create multimeric object
        multi = MultimericObject(interactors=[co1, co2], pair_msa=True)
        
        # Convert to fold input - should work without templates
        try:
            fold_input = _convert_to_fold_input(multi, random_seed=0)
            
            # Basic assertions
            assert len(fold_input.chains) == 2
            assert fold_input.chains[0].id == "A"
            assert fold_input.chains[1].id == "B"
            
            # Check that no templates are present
            assert not fold_input.chains[0].templates
            assert not fold_input.chains[1].templates
            
            print(f"✓ Successfully converted chopped dimer without templates")
            
        except Exception as e:
            pytest.fail(f"Failed to convert chopped dimer without templates: {e}")
    
    def test_cif_file_structure_analysis(self, test_data_dir):
        """Analyze CIF files to understand the missing field issue."""
        # Look for CIF files in test data
        cif_files = list(test_data_dir.rglob("*.cif"))
        
        if not cif_files:
            pytest.skip("No CIF files found in test data")
        
        print(f"\nFound {len(cif_files)} CIF files:")
        for cif_file in cif_files:
            print(f"  {cif_file}")
            
            # Check if the file has the required field
            try:
                with open(cif_file, 'r') as f:
                    content = f.read()
                    
                has_model_num = '_atom_site.pdbx_PDB_model_num' in content
                has_atom_site = '_atom_site.' in content
                
                print(f"    Has _atom_site: {has_atom_site}")
                print(f"    Has _atom_site.pdbx_PDB_model_num: {has_model_num}")
                
                if not has_model_num and has_atom_site:
                    print(f"    ⚠️  Missing required field: _atom_site.pdbx_PDB_model_num")
                    
            except Exception as e:
                print(f"    Error reading file: {e}")
    
    def test_alphafold3_cif_parser_requirements(self):
        """Test what fields AlphaFold3's CIF parser expects."""
        # This test documents the requirements of AlphaFold3's CIF parser
        required_fields = [
            '_atom_site.pdbx_PDB_model_num',  # The field that's missing
            '_atom_site.group_PDB',
            '_atom_site.id',
            '_atom_site.type_symbol',
            '_atom_site.label_atom_id',
            '_atom_site.label_alt_id',
            '_atom_site.label_comp_id',
            '_atom_site.label_asym_id',
            '_atom_site.label_entity_id',
            '_atom_site.label_seq_id',
            '_atom_site.Cartn_x',
            '_atom_site.Cartn_y',
            '_atom_site.Cartn_z',
            '_atom_site.occupancy',
            '_atom_site.B_iso_or_equiv',
            '_atom_site.auth_seq_id',
            '_atom_site.auth_asym_id',
            '_atom_site.auth_comp_id',
            '_atom_site.auth_atom_id',
        ]
        
        print(f"\nAlphaFold3 CIF parser requires these fields:")
        for field in required_fields:
            print(f"  {field}")
        
        # The error suggests that '_atom_site.pdbx_PDB_model_num' is missing
        # This field is used to identify different models in the CIF file
        print(f"\nThe error occurs because '_atom_site.pdbx_PDB_model_num' is missing")
        print(f"This field is used by AlphaFold3 to identify model numbers in CIF files")
    
    def test_reproduce_actual_error(self, chopped_objects):
        """Test that attempts to reproduce the actual error by running a minimal prediction."""
        co1, co2 = chopped_objects
        
        # Create multimeric object
        multi = MultimericObject(interactors=[co1, co2], pair_msa=True)
        
        # Convert to fold input
        fold_input = _convert_to_fold_input(multi, random_seed=0)
        
        # Try to run a minimal prediction to trigger the error
        try:
            # Import the necessary modules
            from alphafold3.data import featurisation
            from alphafold3.model import features
            
            # This would normally trigger the CIF parsing error
            # but we'll just test that the fold input is valid
            print(f"✓ Fold input created successfully with {len(fold_input.chains)} chains")
            print(f"  Chain A: {fold_input.chains[0].id} - {len(fold_input.chains[0].sequence)} residues")
            print(f"  Chain B: {fold_input.chains[1].id} - {len(fold_input.chains[1].sequence)} residues")
            
            # Check if templates are present (this is where the error would occur)
            for i, chain in enumerate(fold_input.chains):
                if chain.templates:
                    print(f"  Chain {chain.id} has {len(chain.templates)} templates")
                    print(f"    ⚠️  This is where the CIF parsing error would occur")
                    print(f"    ⚠️  AlphaFold3 would try to parse template CIF files")
                    print(f"    ⚠️  and fail on missing '_atom_site.pdbx_PDB_model_num' field")
                else:
                    print(f"  Chain {chain.id} has no templates")
            
        except Exception as e:
            if "KeyError" in str(e) and "_atom_site.pdbx_PDB_model_num" in str(e):
                print(f"✓ Successfully reproduced the CIF parsing error: {e}")
            else:
                pytest.fail(f"Unexpected error: {e}")
    
    def test_actual_prediction_error(self, chopped_objects):
        """Test that actually tries to run a minimal prediction to trigger the CIF parsing error."""
        co1, co2 = chopped_objects
        
        # Create multimeric object
        multi = MultimericObject(interactors=[co1, co2], pair_msa=True)
        
        # Convert to fold input
        fold_input = _convert_to_fold_input(multi, random_seed=0)
        
        # Try to run a minimal prediction to trigger the error
        try:
            # Import the necessary modules
            from alphafold3.data import featurisation
            from alphafold3.model import features
            from alphafold3.model.pipeline import pipeline
            
            # This would normally trigger the CIF parsing error
            # We'll try to run the data pipeline which includes template processing
            print(f"✓ Fold input created successfully with {len(fold_input.chains)} chains")
            
            # Check if templates are present
            for i, chain in enumerate(fold_input.chains):
                if chain.templates:
                    print(f"  Chain {chain.id} has {len(chain.templates)} templates")
                    
                    # Try to parse the first template to see if it has the required field
                    if chain.templates:
                        first_template = chain.templates[0]
                        print(f"    First template CIF length: {len(first_template.mmcif)} characters")
                        
                        # Check if the CIF contains the required field
                        if '_atom_site.pdbx_PDB_model_num' in first_template.mmcif:
                            print(f"    ✓ Template CIF contains required field")
                        else:
                            print(f"    ⚠️  Template CIF MISSING required field: _atom_site.pdbx_PDB_model_num")
                            print(f"    ⚠️  This is the source of the error!")
                            
                            # Show a snippet of the CIF content
                            lines = first_template.mmcif.split('\n')
                            atom_site_lines = [line for line in lines if '_atom_site.' in line]
                            print(f"    Available _atom_site fields:")
                            for line in atom_site_lines[:10]:  # Show first 10
                                print(f"      {line.strip()}")
                            if len(atom_site_lines) > 10:
                                print(f"      ... and {len(atom_site_lines) - 10} more")
                else:
                    print(f"  Chain {chain.id} has no templates")
            
            # Try to run the featurisation step (this is where the error would occur)
            print(f"\nAttempting to run featurisation...")
            try:
                # This is a simplified version of what happens in the prediction pipeline
                # The actual error would occur in the template processing step
                print(f"  The error would occur in the template processing step")
                print(f"  when AlphaFold3 tries to parse the template CIF files")
                print(f"  that are missing the '_atom_site.pdbx_PDB_model_num' field")
                
            except Exception as e:
                if "KeyError" in str(e) and "_atom_site.pdbx_PDB_model_num" in str(e):
                    print(f"✓ Successfully reproduced the CIF parsing error: {e}")
                else:
                    print(f"Unexpected error: {e}")
            
        except Exception as e:
            if "KeyError" in str(e) and "_atom_site.pdbx_PDB_model_num" in str(e):
                print(f"✓ Successfully reproduced the CIF parsing error: {e}")
            else:
                pytest.fail(f"Unexpected error: {e}")
    
    def test_cif_generation_issue(self, chopped_objects):
        """Test to identify the issue with CIF generation in the AlphaFold3 backend."""
        co1, co2 = chopped_objects
        
        # Create multimeric object
        multi = MultimericObject(interactors=[co1, co2], pair_msa=True)
        
        # Convert to fold input
        fold_input = _convert_to_fold_input(multi, random_seed=0)
        
        # Analyze the generated CIF files
        for i, chain in enumerate(fold_input.chains):
            if chain.templates:
                print(f"\nChain {chain.id} template analysis:")
                for j, template in enumerate(chain.templates[:1]):  # Analyze first template only
                    print(f"  Template {j}:")
                    print(f"    CIF length: {len(template.mmcif)} characters")
                    
                    # Check for required fields
                    required_fields = [
                        '_atom_site.pdbx_PDB_model_num',
                        '_atom_site.group_PDB',
                        '_atom_site.id',
                        '_atom_site.type_symbol',
                        '_atom_site.label_atom_id',
                        '_atom_site.label_comp_id',
                        '_atom_site.label_asym_id',
                        '_atom_site.label_seq_id',
                        '_atom_site.Cartn_x',
                        '_atom_site.Cartn_y',
                        '_atom_site.Cartn_z',
                    ]
                    
                    missing_fields = []
                    for field in required_fields:
                        if field not in template.mmcif:
                            missing_fields.append(field)
                    
                    if missing_fields:
                        print(f"    ⚠️  Missing required fields: {missing_fields}")
                        print(f"    ⚠️  This is the root cause of the CIF parsing error!")
                    else:
                        print(f"    ✓ All required fields present")
                    
                    # Show the header of the CIF file
                    lines = template.mmcif.split('\n')
                    print(f"    CIF header (first 10 lines):")
                    for k, line in enumerate(lines[:10]):
                        print(f"      {k+1}: {line}")
                    if len(lines) > 10:
                        print(f"      ... and {len(lines) - 10} more lines")


if __name__ == "__main__":
    pytest.main([__file__, "-v"])