nn/test_polynomial_system.py at main · makingmusic/nn · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
#!/usr/bin/env python3
"""
Simple test script to verify the polynomial system works correctly.
This script tests the core functionality without requiring the full training environment.
"""

def test_polynomial_configuration():
    """Test the polynomial coefficient configuration system."""
    print("Testing polynomial configuration system...")

    # Test basic functionality
    from dataset import (
        set_polynomial_coefficients,
        get_polynomial_coefficients,
        get_polynomial_degree,
        print_polynomial_function
    )

    # Test 1: Linear polynomial
    print("\nTest 1: Linear polynomial")
    set_polynomial_coefficients([10, 5])  # y = 10 + 5x
    coeffs = get_polynomial_coefficients()
    degree = get_polynomial_degree()
    print(f"Expected: [10, 5], degree 1")
    print(f"Got: {coeffs}, degree {degree}")
    print_polynomial_function()

    # Test 2: Quadratic polynomial
    print("\nTest 2: Quadratic polynomial")
    set_polynomial_coefficients([2, 3, 4])  # y = 2 + 3x + 4x²
    coeffs = get_polynomial_coefficients()
    degree = get_polynomial_degree()
    print(f"Expected: [2, 3, 4], degree 2")
    print(f"Got: {coeffs}, degree {degree}")
    print_polynomial_function()

    # Test 3: Cubic polynomial
    print("\nTest 3: Cubic polynomial")
    set_polynomial_coefficients([1, 2, 3, 4])  # y = 1 + 2x + 3x² + 4x³
    coeffs = get_polynomial_coefficients()
    degree = get_polynomial_degree()
    print(f"Expected: [1, 2, 3, 4], degree 3")
    print(f"Got: {coeffs}, degree {degree}")
    print_polynomial_function()

    # Test 4: Polynomial with negative coefficients
    print("\nTest 4: Polynomial with negative coefficients")
    set_polynomial_coefficients([100, -10, 2])  # y = 100 - 10x + 2x²
    coeffs = get_polynomial_coefficients()
    degree = get_polynomial_degree()
    print(f"Expected: [100, -10, 2], degree 2")
    print(f"Got: {coeffs}, degree {degree}")
    print_polynomial_function()

    print("\n✅ All polynomial configuration tests passed!")

def test_dataset_creation():
    """Test polynomial dataset creation."""
    print("\nTesting polynomial dataset creation...")

    from dataset import create_polynomial_dataset, set_polynomial_coefficients

    # Test with different polynomial configurations
    test_cases = [
        ([10, 5], "Linear: y = 10 + 5x"),
        ([2, 3, 4], "Quadratic: y = 2 + 3x + 4x²"),
        ([1, 2, 3, 4], "Cubic: y = 1 + 2x + 3x² + 4x³"),
    ]

    for coeffs, description in test_cases:
        print(f"\n{description}")
        set_polynomial_coefficients(coeffs)
        x_train, y_train = create_polynomial_dataset()

        print(f"  Dataset shape: x={x_train.shape}, y={y_train.shape}")
        print(f"  X range: {x_train.min().item():.0f} to {x_train.max().item():.0f}")
        print(f"  Y range: {y_train.min().item():.0f} to {y_train.max().item():.0f}")

        # Test a few sample points
        for i in range(0, min(3, len(x_train))):
            x_val = x_train[i].item()
            y_val = y_train[i].item()

            # Calculate expected value
            expected = 0
            for j, coeff in enumerate(coeffs):
                expected += coeff * (x_val ** j)

            print(f"  x={x_val}, y={y_val:.2f}, expected={expected:.2f}")

    print("\n✅ All dataset creation tests passed!")

def test_model_creation():
    """Test polynomial model creation."""
    print("\nTesting polynomial model creation...")

    from dataset import create_model, set_polynomial_coefficients, get_polynomial_degree

    # Test with different polynomial degrees
    test_cases = [
        ([10, 5], "Linear model"),
        ([2, 3, 4], "Quadratic model"),
        ([1, 2, 3, 4], "Cubic model"),
    ]

    for coeffs, description in test_cases:
        print(f"\n{description}")
        set_polynomial_coefficients(coeffs)
        model = create_model('polynomial')

        degree = get_polynomial_degree()
        expected_params = degree + 1

        print(f"  Expected degree: {degree}")
        print(f"  Expected parameters: {expected_params}")
        print(f"  Model coefficients shape: {model.coefficients.shape}")
        print(f"  Actual parameters: {model.coefficients.numel()}")

        # Test forward pass
        import torch
        test_x = torch.tensor([[1.0], [2.0], [3.0]])
        with torch.no_grad():
            output = model(test_x)
        print(f"  Forward pass output shape: {output.shape}")

    print("\n✅ All model creation tests passed!")

def main():
    """Run all tests."""
    print("🧪 Testing Flexible Polynomial System")
    print("=" * 50)

    try:
        test_polynomial_configuration()
        test_dataset_creation()
        test_model_creation()

        print("\n🎉 All tests passed! The polynomial system is working correctly.")

    except Exception as e:
        print(f"\n❌ Test failed with error: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()