feat: add Sobol distribution (#529)

Author: HarrisonKramer

optiland/_types.py

@@ -54,6 +54,7 @@
     "hexapolar",
     "cross",
     "ring",
+    "sobol",
 ]
 ApertureType = Literal["EPD", "imageFNO", "objectNA", "float_by_stop_size"]
 Fields = Literal["all"] | Sequence[tuple[float, float]]

optiland/distribution.py

@@ -12,6 +12,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+from scipy.stats import qmc
 
 import optiland.backend as be
 
@@ -375,6 +376,42 @@ def generate_points(self, num_points: int):
         self.y = be.sin(theta)
 
 
+class SobolDistribution(BaseDistribution):
+    """A class representing a Sobol distribution.
+
+    Generates `num_points` points using a Sobol low-discrepancy sequence
+    within the unit disk.
+
+    Attributes:
+        seed (int | None): Seed for the Sobol sequence generator.
+        x: The x-coordinates of the generated points.
+        y: The y-coordinates of the generated points.
+    """
+
+    def __init__(self, seed: int | None = None):
+        super().__init__()
+        self.seed = seed
+
+    def generate_points(self, num_points: int):
+        """Generates Sobol points.
+
+        Args:
+            num_points (int): The number of points to generate.
+
+        """
+        sampler = qmc.Sobol(d=2, scramble=True, seed=self.seed)
+        sample = sampler.random(num_points)
+
+        u1 = be.array(sample[:, 0])
+        u2 = be.array(sample[:, 1])
+
+        r = be.sqrt(u1)
+        theta = 2 * be.pi * u2
+
+        self.x = r * be.cos(theta)
+        self.y = r * be.sin(theta)
+
+
 def create_distribution(distribution_type: DistributionType) -> BaseDistribution:
     """Create a distribution based on the given distribution type.
 
@@ -400,6 +437,7 @@ def create_distribution(distribution_type: DistributionType) -> BaseDistribution
         "hexapolar": HexagonalDistribution,
         "cross": CrossDistribution,
         "ring": RingDistribution,
+        "sobol": SobolDistribution,
     }
 
     if distribution_type not in distribution_classes:

tests/test_distribution.py

@@ -4,6 +4,7 @@
 import pytest
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
+from scipy.stats import qmc
 
 import optiland.backend as be
 from optiland import distribution
@@ -474,3 +475,23 @@ def test_gaussian_quad_weights(set_test_backend):
             weights / scale[k],
             be.array([0.04283, 0.09019, 0.11698, 0.11698, 0.09019, 0.04283]),
         )
+
+
+@pytest.mark.parametrize("num_points", [16, 64, 256, 1024])
+def test_sobol_distribution(set_test_backend, num_points):
+    seed = 42
+    d = distribution.SobolDistribution(seed=seed)
+    d.generate_points(num_points=num_points)
+
+    sampler = qmc.Sobol(d=2, scramble=True, seed=seed)
+    sample = sampler.random(num_points)
+    u1 = be.array(sample[:, 0])
+    u2 = be.array(sample[:, 1])
+    r = be.sqrt(u1)
+    theta = 2 * be.pi * u2
+    x = r * be.cos(theta)
+    y = r * be.sin(theta)
+
+    assert_allclose(d.x, x)
+    assert_allclose(d.y, y)
+    assert be.all(d.x**2 + d.y**2 <= 1.0 + 1e-7)

tests/test_distribution_points.py

@@ -12,6 +12,7 @@
     LineYDistribution,
     RandomDistribution,
     RingDistribution,
+    SobolDistribution,
     UniformDistribution,
 )
 
@@ -219,3 +220,14 @@ def test_ring_distribution_zero_points(self, backend):
         )
         assert len(dist.x) == 0
         assert len(dist.y) == 0
+
+    def test_sobol_distribution(self, backend):
+        be.set_backend(backend)
+        num_points = 100
+        dist = SobolDistribution(seed=42)
+        dist.generate_points(num_points=num_points)
+        # For Sobol, all points should be unique
+        unique_count = count_unique_points(dist.x, dist.y)
+        assert unique_count == num_points, (
+            f"SobolDistribution expected {num_points} unique points, got {unique_count}"
+        )