add, test code

Author: singer-yang

test/README.md

@@ -0,0 +1,9 @@
+# Run all tests
+cd /home/yangx0i/deeplens_proj/debug/DeepLens
+pytest test/ -v
+
+# Run specific test file
+pytest test/test_ray.py -v
+
+# Run with coverage
+pytest test/ --cov=deeplens --cov-report=term-missing

test/conftest.py

@@ -0,0 +1,164 @@
+"""
+Shared pytest fixtures for DeepLens test suite.
+"""
+
+import os
+import sys
+
+import pytest
+import torch
+
+# Add deeplens to path
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+
+# =============================================================================
+# Device fixtures
+# =============================================================================
+@pytest.fixture(scope="session")
+def device():
+    """Return CUDA device if available, otherwise CPU."""
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    else:
+        pytest.skip("CUDA not available, skipping GPU tests")
+
+
+@pytest.fixture(scope="session")
+def device_cpu():
+    """Return CPU device."""
+    return torch.device("cpu")
+
+
+@pytest.fixture(scope="session")
+def device_auto():
+    """Return CUDA if available, otherwise CPU (for tests that should run on either)."""
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+# =============================================================================
+# Lens fixtures
+# =============================================================================
+@pytest.fixture(scope="function")
+def sample_singlet_lens(device_auto):
+    """Load a simple singlet lens for testing."""
+    from deeplens import GeoLens
+
+    lens_path = os.path.join(
+        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        "datasets/lenses/singlet/example1.json",
+    )
+    lens = GeoLens(filename=lens_path)
+    lens.to(device_auto)
+    return lens
+
+
+@pytest.fixture(scope="function")
+def sample_cellphone_lens(device_auto):
+    """Load a cellphone lens with aspheric surfaces for testing."""
+    from deeplens import GeoLens
+
+    lens_path = os.path.join(
+        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        "datasets/lenses/cellphone/cellphone68deg.json",
+    )
+    lens = GeoLens(filename=lens_path)
+    lens.to(device_auto)
+    return lens
+
+
+@pytest.fixture(scope="function")
+def sample_camera_lens(device_auto):
+    """Load a camera lens for testing."""
+    from deeplens import GeoLens
+
+    lens_path = os.path.join(
+        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        "datasets/lenses/camera/ef50mm_f1.8.json",
+    )
+    lens = GeoLens(filename=lens_path)
+    lens.to(device_auto)
+    return lens
+
+
+# =============================================================================
+# Image fixtures
+# =============================================================================
+@pytest.fixture(scope="function")
+def sample_image(device_auto):
+    """Create a simple test image tensor [B, C, H, W]."""
+    # Create a gradient image for testing
+    H, W = 256, 256
+    x = torch.linspace(0, 1, W, device=device_auto)
+    y = torch.linspace(0, 1, H, device=device_auto)
+    yy, xx = torch.meshgrid(y, x, indexing="ij")
+    
+    img = torch.stack([xx, yy, (xx + yy) / 2], dim=0)  # [3, H, W]
+    img = img.unsqueeze(0)  # [1, 3, H, W]
+    return img
+
+
+@pytest.fixture(scope="function")
+def sample_image_small(device_auto):
+    """Create a small test image tensor for fast tests."""
+    H, W = 64, 64
+    img = torch.rand(1, 3, H, W, device=device_auto)
+    return img
+
+
+# =============================================================================
+# Ray fixtures
+# =============================================================================
+@pytest.fixture(scope="function")
+def sample_ray(device_auto):
+    """Create a sample ray for testing."""
+    from deeplens.optics.ray import Ray
+
+    o = torch.tensor([[0.0, 0.0, -100.0]], device=device_auto)
+    d = torch.tensor([[0.0, 0.0, 1.0]], device=device_auto)
+    ray = Ray(o, d, wvln=0.55, device=device_auto)
+    return ray
+
+
+@pytest.fixture(scope="function")
+def sample_rays_batch(device_auto):
+    """Create a batch of rays for testing."""
+    from deeplens.optics.ray import Ray
+
+    # Create 100 rays in a grid pattern
+    n = 10
+    x = torch.linspace(-1, 1, n, device=device_auto)
+    y = torch.linspace(-1, 1, n, device=device_auto)
+    yy, xx = torch.meshgrid(y, x, indexing="ij")
+    
+    o = torch.stack([xx.flatten(), yy.flatten(), torch.full((n*n,), -100.0, device=device_auto)], dim=-1)
+    d = torch.zeros_like(o)
+    d[..., 2] = 1.0
+    
+    ray = Ray(o, d, wvln=0.55, device=device_auto)
+    return ray
+
+
+# =============================================================================
+# Path helpers
+# =============================================================================
+@pytest.fixture(scope="session")
+def project_root():
+    """Return the project root directory."""
+    return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+@pytest.fixture(scope="session")
+def lenses_dir(project_root):
+    """Return the lenses dataset directory."""
+    return os.path.join(project_root, "datasets/lenses")
+
+
+@pytest.fixture(scope="session")
+def test_output_dir(project_root):
+    """Return a directory for test outputs."""
+    output_dir = os.path.join(project_root, "test/test_outputs")
+    os.makedirs(output_dir, exist_ok=True)
+    return output_dir

test/test_basics.py

@@ -0,0 +1,210 @@
+"""
+Tests for deeplens/basics.py - Basic utilities and constants.
+"""
+
+import pytest
+import torch
+
+
+from deeplens.basics import (
+    DEPTH,
+    DEFAULT_WAVE,
+    EPSILON,
+    PSF_KS,
+    SPP_PSF,
+    WAVE_RGB,
+    DeepObj,
+    init_device,
+    wave_rgb,
+)
+
+
+class TestConstants:
+    """Test default constants are properly defined."""
+
+    def test_depth_constant(self):
+        """DEPTH should be a large negative value representing infinity."""
+        assert DEPTH == -20000.0
+        assert DEPTH < 0
+
+    def test_wave_rgb(self):
+        """WAVE_RGB should contain R, G, B wavelengths in micrometers."""
+        assert len(WAVE_RGB) == 3
+        assert WAVE_RGB[0] > WAVE_RGB[1] > WAVE_RGB[2]  # R > G > B
+        # All wavelengths should be in visible range (0.38 - 0.78 um)
+        for wvln in WAVE_RGB:
+            assert 0.38 < wvln < 0.78
+
+    def test_default_wave(self):
+        """DEFAULT_WAVE should be green wavelength."""
+        assert 0.5 < DEFAULT_WAVE < 0.6  # Green light
+
+    def test_spp_psf(self):
+        """SPP_PSF should be a power of 2."""
+        assert SPP_PSF > 0
+        assert (SPP_PSF & (SPP_PSF - 1)) == 0  # Check power of 2
+
+    def test_psf_ks(self):
+        """PSF_KS should be a reasonable kernel size."""
+        assert PSF_KS > 0
+        assert PSF_KS < 256
+
+    def test_epsilon(self):
+        """EPSILON should be a small positive value."""
+        assert EPSILON > 0
+        assert EPSILON < 1e-6
+
+
+class TestInitDevice:
+    """Test device initialization."""
+
+    def test_init_device_returns_device(self):
+        """init_device should return a torch device."""
+        device = init_device()
+        assert isinstance(device, torch.device)
+
+    def test_init_device_cuda_or_cpu(self):
+        """init_device should return cuda or cpu."""
+        device = init_device()
+        assert device.type in ["cuda", "cpu"]
+
+    def test_init_device_matches_availability(self):
+        """init_device result should match CUDA availability."""
+        device = init_device()
+        if torch.cuda.is_available():
+            assert device.type == "cuda"
+        else:
+            assert device.type == "cpu"
+
+
+class TestWaveRgb:
+    """Test random wavelength sampling."""
+
+    def test_wave_rgb_returns_three(self):
+        """wave_rgb should return 3 wavelengths."""
+        waves = wave_rgb()
+        assert len(waves) == 3
+
+    def test_wave_rgb_order(self):
+        """wave_rgb should return [R, G, B] in decreasing wavelength order."""
+        waves = wave_rgb()
+        assert waves[0] > waves[1] > waves[2]
+
+    def test_wave_rgb_range(self):
+        """All wavelengths should be in visible range."""
+        waves = wave_rgb()
+        for w in waves:
+            assert 0.4 < w < 0.75
+
+    def test_wave_rgb_randomness(self):
+        """wave_rgb should produce different results (probabilistic)."""
+        results = [tuple(wave_rgb()) for _ in range(10)]
+        # At least some should be different
+        assert len(set(results)) > 1
+
+
+class TestDeepObj:
+    """Test DeepObj base class functionality."""
+
+    def test_deep_obj_init(self):
+        """DeepObj should initialize with default dtype."""
+        obj = DeepObj()
+        assert obj.dtype == torch.get_default_dtype()
+
+    def test_deep_obj_init_custom_dtype(self):
+        """DeepObj should accept custom dtype."""
+        obj = DeepObj(dtype=torch.float64)
+        assert obj.dtype == torch.float64
+
+    def test_deep_obj_str(self):
+        """DeepObj should have string representation."""
+        obj = DeepObj()
+        s = str(obj)
+        assert "DeepObj" in s
+
+    def test_deep_obj_clone(self):
+        """DeepObj clone should create independent copy."""
+        obj = DeepObj()
+        obj.test_attr = torch.tensor([1.0, 2.0, 3.0])
+        cloned = obj.clone()
+        
+        # Modify original
+        obj.test_attr[0] = 999.0
+        
+        # Clone should be unchanged
+        assert cloned.test_attr[0] != 999.0
+
+    def test_deep_obj_to_device(self, device_auto):
+        """DeepObj.to() should move tensors to device."""
+        obj = DeepObj()
+        obj.tensor_attr = torch.tensor([1.0, 2.0, 3.0])
+        
+        obj.to(device_auto)
+        
+        assert obj.device.type == device_auto.type
+        assert obj.tensor_attr.device.type == device_auto.type
+
+    def test_deep_obj_to_device_nested(self, device_auto):
+        """DeepObj.to() should handle nested DeepObj."""
+        outer = DeepObj()
+        inner = DeepObj()
+        inner.data = torch.tensor([1.0, 2.0])
+        outer.child = inner
+        
+        outer.to(device_auto)
+        
+        assert inner.data.device.type == device_auto.type
+
+    def test_deep_obj_to_device_list(self, device_auto):
+        """DeepObj.to() should handle tensor lists."""
+        obj = DeepObj()
+        obj.tensor_list = [torch.tensor([1.0]), torch.tensor([2.0])]
+        
+        obj.to(device_auto)
+        
+        for t in obj.tensor_list:
+            assert t.device.type == device_auto.type
+
+    def test_deep_obj_astype_float32(self):
+        """DeepObj.astype() should convert to float32."""
+        obj = DeepObj(dtype=torch.float64)
+        obj.data = torch.tensor([1.0, 2.0], dtype=torch.float64)
+        
+        obj.astype(torch.float32)
+        
+        assert obj.dtype == torch.float32
+        assert obj.data.dtype == torch.float32
+
+    def test_deep_obj_astype_float64(self):
+        """DeepObj.astype() should convert to float64."""
+        obj = DeepObj(dtype=torch.float32)
+        obj.data = torch.tensor([1.0, 2.0], dtype=torch.float32)
+        
+        obj.astype(torch.float64)
+        
+        assert obj.dtype == torch.float64
+        assert obj.data.dtype == torch.float64
+
+    def test_deep_obj_astype_none(self):
+        """DeepObj.astype(None) should be no-op."""
+        obj = DeepObj(dtype=torch.float32)
+        original_dtype = obj.dtype
+        
+        result = obj.astype(None)
+        
+        assert obj.dtype == original_dtype
+        assert result is obj
+
+    def test_deep_obj_astype_invalid(self):
+        """DeepObj.astype() should reject invalid dtypes."""
+        obj = DeepObj()
+        
+        with pytest.raises(AssertionError):
+            obj.astype(torch.int32)
+
+    def test_deep_obj_call_raises(self):
+        """DeepObj.__call__() should raise if forward not implemented."""
+        obj = DeepObj()
+        
+        with pytest.raises(AttributeError):
+            obj(torch.tensor([1.0]))