test_cgo_arrows.py

#!/usr/bin/env python3
"""
Test script for CGO arrow functionality.

This script tests the new CGO arrow creation to replace pseudoatom arrows.
"""

import sys
import os
import numpy as np
import pandas as pd

# Add the current directory to the Python path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

def test_cgo_arrow_creation():
    """Test CGO arrow creation functionality"""
    print("🏹 TESTING CGO ARROW CREATION")
    print("=" * 50)
    
    # Mock CGO constants for testing
    class MockCGO:
        BEGIN = 0
        END = 1
        VERTEX = 2
        NORMAL = 3
        COLOR = 4
        LINEWIDTH = 5
        LINES = 6
        TRIANGLES = 7
        CYLINDER = 22
        CONE = 23
    
    # Create a simple test class with the CGO arrow method
    class CGOArrowTester:
        def create_cgo_arrow(self, start_pos, end_pos, magnitude):
            """Create a CGO arrow from start to end position"""
            try:
                cgo = MockCGO()
                arrow_objects = []
                
                # Calculate arrow properties
                arrow_vector = np.array(end_pos) - np.array(start_pos)
                arrow_length = np.linalg.norm(arrow_vector)
                
                if arrow_length < 1e-6:
                    return []
                
                # Normalize vector for direction
                arrow_direction = arrow_vector / arrow_length
                
                # Calculate shaft end (90% of total length)
                shaft_length = arrow_length * 0.9
                shaft_end = np.array(start_pos) + arrow_direction * shaft_length
                
                # Color based on magnitude and direction
                if np.sum(arrow_vector) > 0:
                    color = [1.0, 0.65, 0.0]  # Orange
                else:
                    color = [0.0, 1.0, 1.0]  # Cyan
                
                # Cylinder shaft
                shaft_radius = min(0.05, arrow_length * 0.02)
                arrow_objects.extend([
                    cgo.CYLINDER,
                    start_pos[0], start_pos[1], start_pos[2],
                    shaft_end[0], shaft_end[1], shaft_end[2],
                    shaft_radius,
                    color[0], color[1], color[2],
                    color[0], color[1], color[2]
                ])
                
                # Cone arrowhead
                head_radius = shaft_radius * 2.5
                arrow_objects.extend([
                    cgo.CONE,
                    shaft_end[0], shaft_end[1], shaft_end[2],
                    end_pos[0], end_pos[1], end_pos[2],
                    head_radius, 0.0,
                    color[0], color[1], color[2],
                    color[0], color[1], color[2],
                    1.0, 1.0
                ])
                
                return arrow_objects
                
            except Exception as e:
                print(f"Error creating CGO arrow: {e}")
                return []
    
    # Test arrow creation
    tester = CGOArrowTester()
    
    # Test case 1: Standard arrow
    start_pos = [0.0, 0.0, 0.0]
    end_pos = [1.0, 1.0, 1.0]
    magnitude = np.linalg.norm(np.array(end_pos) - np.array(start_pos))
    
    arrow_cgo = tester.create_cgo_arrow(start_pos, end_pos, magnitude)
    
    print(f"✅ Test 1: Standard arrow")
    print(f"   Start: {start_pos}")
    print(f"   End: {end_pos}")
    print(f"   Magnitude: {magnitude:.3f}")
    print(f"   CGO elements: {len(arrow_cgo)}")
    
    # Test case 2: Negative direction arrow
    start_pos2 = [2.0, 2.0, 2.0]
    end_pos2 = [1.0, 1.0, 1.0]
    magnitude2 = np.linalg.norm(np.array(end_pos2) - np.array(start_pos2))
    
    arrow_cgo2 = tester.create_cgo_arrow(start_pos2, end_pos2, magnitude2)
    
    print(f"\n✅ Test 2: Negative direction arrow")
    print(f"   Start: {start_pos2}")
    print(f"   End: {end_pos2}")
    print(f"   Magnitude: {magnitude2:.3f}")
    print(f"   CGO elements: {len(arrow_cgo2)}")
    
    # Test case 3: Zero-length arrow (should return empty)
    start_pos3 = [0.0, 0.0, 0.0]
    end_pos3 = [0.0, 0.0, 0.0]
    magnitude3 = 0.0
    
    arrow_cgo3 = tester.create_cgo_arrow(start_pos3, end_pos3, magnitude3)
    
    print(f"\n✅ Test 3: Zero-length arrow")
    print(f"   Start: {start_pos3}")
    print(f"   End: {end_pos3}")
    print(f"   Magnitude: {magnitude3:.3f}")
    print(f"   CGO elements: {len(arrow_cgo3)} (should be 0)")
    
    # Validate structure (CYLINDER has 13 elements, CONE has 18 elements = 31 total)
    expected_elements = 31  # CYLINDER (13 elements) + CONE (18 elements)
    if len(arrow_cgo) == expected_elements and len(arrow_cgo2) == expected_elements and len(arrow_cgo3) == 0:
        print(f"\n✅ CGO ARROW TESTS PASSED!")
        print(f"   - Correct number of CGO elements generated")
        print(f"   - Zero-length arrows properly handled")
        print(f"   - Color assignment working")
        return True
    else:
        print(f"\n❌ CGO ARROW TESTS FAILED!")
        print(f"   Expected {expected_elements} elements, got {len(arrow_cgo)}")
        return False

def test_displacement_vector_workflow():
    """Test the complete displacement vector workflow with CGO"""
    print(f"\n🎯 TESTING COMPLETE DISPLACEMENT VECTOR WORKFLOW")
    print("=" * 60)
    
    # Create test data for displacement vectors
    test_data = [
        {
            'chain': 'A', 'seqid': 100, 'residue': 'ALA', 'atom': 'CA', 'altloc': None,
            'coordx': 10.0, 'coordy': 15.0, 'coordz': 20.0,
            'peakx': 11.0, 'peaky': 16.0, 'peakz': 21.0,
            'markx': 10.5, 'marky': 15.5, 'markz': 20.5,
            'peak': 5.2, 'deltax': 0.5, 'deltay': 0.5, 'deltaz': 0.5
        },
        {
            'chain': 'A', 'seqid': 100, 'residue': 'ALA', 'atom': 'CB', 'altloc': None,
            'coordx': 10.0, 'coordy': 15.0, 'coordz': 20.0,
            'peakx': 9.5, 'peaky': 14.5, 'peakz': 19.5,
            'markx': 10.0, 'marky': 15.0, 'markz': 20.0,
            'peak': -3.8, 'deltax': -0.5, 'deltay': -0.5, 'deltaz': -0.5
        }
    ]
    
    peaks_df = pd.DataFrame(test_data)
    
    # Mock the vector calculation (simplified version)
    def calculate_per_atom_vectors(peaks_df):
        """Calculate weighted average displacement vectors per atom"""
        atom_vectors = {}
        
        # Group by atom coordinates - fix the issue with string None
        grouped = peaks_df.groupby(['coordx', 'coordy', 'coordz', 'chain', 'seqid', 'residue', 'atom'])
        
        for atom_coords, group in grouped:
            atom_key = f"{atom_coords[3]}:{atom_coords[4]}:{atom_coords[5]}:{atom_coords[6]}"
            
            # Use mark coordinates as the starting point
            atom_pos = [group.iloc[0]['markx'], group.iloc[0]['marky'], group.iloc[0]['markz']]
            
            # Calculate weighted average displacement
            weights = np.abs(group['peak'].values)
            displacements = group[['deltax', 'deltay', 'deltaz']].values
            
            if np.sum(weights) > 0:
                avg_vector = np.average(displacements, weights=weights, axis=0)
            else:
                avg_vector = np.mean(displacements, axis=0)
            
            atom_vectors[atom_key] = {
                'atom_pos': atom_pos,
                'avg_vector': avg_vector
            }
        
        return atom_vectors
    
    # Test vector calculation
    atom_vectors = calculate_per_atom_vectors(peaks_df)
    
    print(f"✅ Calculated displacement vectors for {len(atom_vectors)} atoms:")
    for atom_key, vector_data in atom_vectors.items():
        pos = vector_data['atom_pos']
        vec = vector_data['avg_vector']
        magnitude = np.linalg.norm(vec)
        print(f"   {atom_key}: pos=({pos[0]:.1f}, {pos[1]:.1f}, {pos[2]:.1f}), "
              f"vec=({vec[0]:.2f}, {vec[1]:.2f}, {vec[2]:.2f}), mag={magnitude:.3f}")
    
    print(f"\n✅ DISPLACEMENT VECTOR WORKFLOW TEST COMPLETED!")
    print(f"   - Successfully grouped peaks by atom")
    print(f"   - Calculated weighted displacement vectors")
    print(f"   - Ready for CGO arrow creation")
    
    return True

if __name__ == "__main__":
    success1 = test_cgo_arrow_creation()
    success2 = test_displacement_vector_workflow()
    
    if success1 and success2:
        print(f"\n🎉 ALL CGO ARROW TESTS PASSED!")
    else:
        print(f"\n❌ SOME TESTS FAILED!")
    
    sys.exit(0 if (success1 and success2) else 1)