remove DiameterBandSurface

Author: victorpoughon

src/torchlensmaker/core/collision_detection.py

@@ -208,7 +208,8 @@ def __call__(
         history: bool = False,
     ) -> CollisionDetectionResult:
         """
-        tmin and tmax must be such that all points within P+tminV and P+tmaxV are within the valid domain of the implicit surface F function
+        tmin and tmax must be such that all points within P+tminV and P+tmaxV
+        are within the valid domain of the implicit surface F function
 
         TODO check we are safe for any ordering of tmin, tmax, i.e. if tmax < tmin
         """
@@ -223,9 +224,7 @@ def __call__(
         assert tmin.dim() == tmax.dim() == 1
 
         # Tensor dimensions
-        N = P.shape[0]  # Number of rays
-        D = P.shape[1]  # Rays dimension (2 or 3)
-        B = self.B
+        (N, _), B = P.shape, self.B
 
         # Initialize solutions t
         t_sample = torch.linspace(0., 1., B, dtype=P.dtype)
@@ -240,15 +239,19 @@ def __call__(
             history_coarse = torch.zeros((B, N, self.num_iterA), dtype=surface.dtype)
             history_fine = torch.zeros((N, self.num_iterB), dtype=surface.dtype)
 
-        br = surface.bounding_radius()
+        max_delta = torch.abs(tmax - tmin) / B
 
         with torch.no_grad():
             # Iteration tensor t
             t = init_t
 
             # Coarse phase (multiple beams)
             for ia in range(self.num_iterA):
-                t = t - self.algoA.delta(surface, P, V, t, max_delta=br / B)
+                t = torch.clamp(
+                    t - self.algoA.delta(surface, P, V, t, max_delta),
+                    min=tmin,
+                    max=tmax,
+                )
                 if history:
                     history_coarse[:, :, ia] = t.clone()
 
@@ -281,14 +284,20 @@ def __call__(
 
             # Fine phase (single beam)
             for ib in range(self.num_iterB):
-                t = t - self.algoB.delta(
-                    surface, P, V, t, max_delta=br / (B * self.num_iterA)
+                t = torch.clamp(
+                    t - self.algoB.delta(surface, P, V, t, max_delta),
+                    min=tmin,
+                    max=tmax,
                 )
                 if history:
                     history_fine[:, ib] = t
 
         # Differentiable phase: one iteration for backwards pass
-        t = t - self.algoC.delta(surface, P, V, t, max_delta=br)
+        t = torch.clamp(
+            t - self.algoC.delta(surface, P, V, t, max_delta=max_delta),
+            min=tmin,
+            max=tmax,
+        )
 
         return CollisionDetectionResult(
             t[0, :],

src/torchlensmaker/core/surfaces.py

@@ -228,19 +228,6 @@ def __init__(
     def bcyl(self) -> Tensor:
         raise NotImplementedError
 
-    def contains(self, points: Tensor, tol: Optional[float] = None) -> Tensor:
-        if tol is None:
-            tol = {torch.float32: 1e-4, torch.float64: 1e-7}[self.dtype]
-
-        # TODO if removing diameterbandsurface, check bounds/bcyl before calling F
-        # because points might be out of g() domain
-
-        dim = points.shape[1]
-
-        F = self.F if dim == 3 else self.f
-
-        return torch.abs(F(points)) < tol
-
     def rmse(self, points: Tensor) -> float:
         N = sum(points.shape[:-1])
         return torch.sqrt(torch.sum(self.Fd(points) ** 2) / N).item()
@@ -337,54 +324,6 @@ def F_grad(self, points: Tensor) -> Tensor:
         raise NotImplementedError
 
 
-class DiameterBandSurfaceSq(ImplicitSurface):
-    "Distance to edge points"
-
-    def __init__(self, Ax: Tensor, Ar: Tensor, dtype: torch.dtype = torch.float64):
-        super().__init__(dtype=dtype)
-        self.Ax = Ax
-        self.Ar = Ar
-
-    def f(self, points: Tensor) -> Tensor:
-        assert points.shape[-1] == 2
-        X, R = points.unbind(-1)
-        Ax, Ar = self.Ax, self.Ar
-        return torch.sqrt((X - Ax) ** 2 + (torch.abs(R) - Ar) ** 2)
-
-    def f_grad(self, points: Tensor) -> Tensor:
-        assert points.shape[-1] == 2
-        X, R = points.unbind(-1)
-        Ax, Ar = self.Ax, self.Ar
-        sq = self.f(points)
-        return torch.stack(
-            ((X - Ax) / sq, torch.sign(R) * (torch.abs(R) - Ar) / sq), dim=-1
-        )
-
-    def F(self, points: Tensor) -> Tensor:
-        assert points.shape[-1] == 3
-        X, Y, Z = points.unbind(-1)
-        R2 = Y**2 + Z**2
-        Ax, Ar = self.Ax, self.Ar
-        return torch.sqrt((X - Ax) ** 2 + (torch.sqrt(R2) - Ar) ** 2)
-
-    def F_grad(self, points: Tensor) -> Tensor:
-        assert points.shape[-1] == 3
-        X, Y, Z = points.unbind(-1)
-        R2 = Y**2 + Z**2
-        Ax, Ar = self.Ax, self.Ar
-        sq = self.F(points)
-        sqr2 = torch.sqrt(R2)
-        quot = (sqr2 - Ar) / (sqr2 * sq)
-        return torch.stack(
-            (
-                (X - Ax) / sq,
-                Y * quot,
-                Z * quot,
-            ),
-            dim=-1,
-        )
-
-
 class SagSurface(ImplicitSurface):
     """
     Axially symmetric implicit surface defined by a sag function.
@@ -416,16 +355,10 @@ def __init__(
     def mask_function(self, points: Tensor) -> Tensor:
         return within_radius(self.diameter / 2, points)
 
-    def fallback_surface(self) -> DiameterBandSurfaceSq:
-        return DiameterBandSurfaceSq(
-            Ax=self.extent_x(),
-            Ar=torch.as_tensor(self.diameter / 2, dtype=self.dtype),
-            dtype=self.dtype,
-        )
-
     def parameters(self) -> dict[str, nn.Parameter]:
         return self.sag_function.parameters()
 
+    # TODO remove?
     def bounding_radius(self) -> float:
         """
         Any point on the surface has a distance to the center that is less
@@ -438,47 +371,28 @@ def tau(self) -> Tensor:
         return torch.as_tensor(self.diameter / 2, dtype=self.dtype)
 
     def f(self, points: Tensor) -> Tensor:
+        "points are assumed to be within the bcyl domain"
         assert points.shape[-1] == 2
         x, r = points.unbind(-1)
-        sag_f = self.sag_function.g(r, self.tau()) - x
-        mask = self.mask_function(points)
-        fallback = self.fallback_surface()
-        return torch.where(mask, sag_f, fallback.f(points))
+        return self.sag_function.g(r, self.tau()) - x
 
     def f_grad(self, points: Tensor) -> Tensor:
         assert points.shape[-1] == 2
         x, r = points.unbind(-1)
-        sag_f_grad = torch.stack(
+        return torch.stack(
             (-torch.ones_like(x), self.sag_function.g_grad(r, self.tau())), dim=-1
         )
-        mask = self.mask_function(points)
-        fallback = self.fallback_surface()
-        return torch.where(
-            mask.unsqueeze(-1).expand(*mask.size(), 2),
-            sag_f_grad,
-            fallback.f_grad(points),
-        )
 
     def F(self, points: Tensor) -> Tensor:
         assert points.shape[-1] == 3
         x, y, z = points.unbind(-1)
-        sag_F = self.sag_function.G(y, z, self.tau()) - x
-        mask = self.mask_function(points)
-        fallback = self.fallback_surface()
-        return torch.where(mask, sag_F, fallback.F(points))
+        return self.sag_function.G(y, z, self.tau()) - x
 
     def F_grad(self, points: Tensor) -> Tensor:
         assert points.shape[-1] == 3
         x, y, z = points.unbind(-1)
         grad_y, grad_z = self.sag_function.G_grad(y, z, self.tau())
-        sag_F_grad = torch.stack((-torch.ones_like(x), grad_y, grad_z), dim=-1)
-        mask = self.mask_function(points)
-        fallback = self.fallback_surface()
-        return torch.where(
-            mask.unsqueeze(-1).expand(*mask.size(), 3),
-            sag_F_grad,
-            fallback.F_grad(points),
-        )
+        return torch.stack((-torch.ones_like(x), grad_y, grad_z), dim=-1)
 
     def extent_x(self) -> Tensor:
         return torch.max(torch.abs(self.sag_function.bounds(self.tau())))
@@ -497,6 +411,34 @@ def bcyl(self) -> Tensor:
             dim=0,
         )
 
+    def contains(self, points: Tensor, tol: Optional[float] = None) -> Tensor:
+        if tol is None:
+            tol = {torch.float32: 1e-4, torch.float64: 1e-7}[self.dtype]
+
+        N, dim = points.shape
+
+        # Check points are within the diameter
+        r2 = points[:, 1] if dim == 2 else points[:, 1] ** 2 + points[:, 2] ** 2
+        within_diameter = r2 <= self.diameter**2
+
+        tau = self.tau()
+        zeros1d = torch.zeros_like(points[:, 1])
+        zeros2d = torch.zeros_like(r2)
+
+        # If within diameter, check the sag equation x = g(r)
+        if dim == 2:
+            safe_input = torch.where(within_diameter, torch.sqrt(r2), zeros2d)
+            sagG = self.sag_function.g(safe_input, tau)
+            G = torch.where(within_diameter, sagG, zeros2d)
+        else:
+            safe_input_y = torch.where(within_diameter, points[:, 1], zeros1d)
+            safe_input_z = torch.where(within_diameter, points[:, 2], zeros1d)
+            sagG = self.sag_function.G(safe_input_y, safe_input_z, tau)
+            G = torch.where(within_diameter, sagG, zeros2d)
+
+        within_tol = torch.abs(G - points[:, 0]) < tol
+        return torch.logical_and(within_diameter, within_tol)
+
     def samples2D_full(self, N, epsilon):
         start = -(1 - epsilon) * self.diameter / 2
         end = (1 - epsilon) * self.diameter / 2

src/torchlensmaker/optimize.py

@@ -112,7 +112,7 @@ def optimize(
 
         if i % show_every == 0 or i == num_iter - 1:
             iter_str = f"[{i + 1:>3}/{num_iter}]"
-            L_str = f"L= {loss.item():>6.3f} | grad norm= {torch.linalg.norm(grad)}"
+            L_str = f"L= {loss.item():>6.5f} | grad norm= {torch.linalg.norm(grad)}"
             print(f"{iter_str} {L_str}")
 
     return OptimizationRecord(num_iter, parameters_record, loss_record, optics)

test_notebooks/XYPolynomial.ipynb

@@ -85,7 +85,7 @@
     "\n",
     "optics = tlm.Sequential(\n",
     "    tlm.PointSourceAtInfinity(6.0),\n",
-    "    tlm.Gap(0.5),\n",
+    "    tlm.Gap(5),\n",
     "    *elements\n",
     ")\n",
     "\n",
@@ -109,6 +109,7 @@
     "import matplotlib.pyplot as plt\n",
     "import json\n",
     "\n",
+    "\n",
     "C = tlm.parameter(torch.tensor([[0.0, 0.0], [0.0, 0.0]], dtype=torch.float64))\n",
     "#surface = tlm.SagSurface(10, tlm.XYPolynomial(C, normalize=True))\n",
     "\n",
@@ -117,9 +118,69 @@
     "# nan in grad with XYPolynomial\n",
     "# nan in grad with Parabola\n",
     "\n",
-    "# surface = tlm.Sphere(10, R=tlm.parameter(15))\n",
+    "#surface = tlm.SphereR(10, R=tlm.parameter(15))\n",
     "surface = tlm.Parabola(10, tlm.parameter(0.02))\n",
     "\n",
+    "x = tlm.parameter(50.)\n",
+    "\n",
+    "optics = tlm.Sequential(\n",
+    "    tlm.PointSourceAtInfinity(8.0),\n",
+    "    tlm.Gap(10),\n",
+    "    tlm.RefractiveSurface(surface, material=\"water-nd\"),\n",
+    "    tlm.Gap(x),\n",
+    "    tlm.FocalPoint()\n",
+    ")\n",
+    "\n",
+    "tlm.show3d(optics, sampling={\"base\": 100})\n",
+    "\n",
+    "### get gradient\n",
+    "sampling = {\"base\": 100}\n",
+    "default_input = tlm.default_input(sampling, dim=3, dtype=torch.float64)\n",
+    "outputs = optics(default_input)\n",
+    "loss = outputs.loss\n",
+    "loss.backward()\n",
+    "\n",
+    "params = list(surface.parameters().values())\n",
+    "\n",
+    "print(surface.to_dict(3))\n",
+    "for param in params:\n",
+    "    print(param, param.grad)\n",
+    "    print()\n",
+    "\n",
+    "print(x, x.grad)\n",
+    "\n",
+    "\n",
+    "tlm.show3d(optics, sampling={\"base\": 100})\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b49cfbb7-81f8-42c9-a53a-446d8725ae0b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## XYPolynomial optimization\n",
+    "\n",
+    "import torchlensmaker as tlm\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import json\n",
+    "\n",
+    "\n",
+    "C = tlm.parameter(torch.zeros((13,13), dtype=torch.float64))\n",
+    "surface = tlm.SagSurface(10, tlm.XYPolynomial(C, normalize=True))\n",
+    "\n",
+    "fixed_mask = torch.zeros_like(C, dtype=torch.bool)\n",
+    "fixed_mask[0, 0] = True  # Freeze position (0,0)\n",
+    "\n",
+    "C.register_hook(lambda grad: grad.masked_fill(fixed_mask, 0.))\n",
+    "\n",
+    "surface = tlm.Sphere(10, R=tlm.parameter(15))\n",
+    "\n",
     "optics = tlm.Sequential(\n",
     "    tlm.PointSourceAtInfinity(8.0),\n",
     "    tlm.Gap(10),\n",
@@ -132,19 +193,30 @@
     "\n",
     "tlm.optimize(\n",
     "    optics,\n",
-    "    optimizer = optim.Adam(optics.parameters(), lr=.1),\n",
-    "    sampling = {\"base\": 100},\n",
+    "    optimizer = tlm.optim.Adam(optics.parameters(), lr=2e-3),\n",
+    "    sampling = {\"base\": 10},\n",
     "    dim = 3,\n",
-    "    num_iter = 100\n",
+    "    num_iter = 50\n",
     ").plot()\n",
     "\n",
-    "tlm.show3d(optics, sampling={\"base\": 100})\n"
+    "tlm.optimize(\n",
+    "    optics,\n",
+    "    optimizer = tlm.optim.Adam(optics.parameters(), lr=1e-4),\n",
+    "    sampling = {\"base\": 10},\n",
+    "    dim = 3,\n",
+    "    num_iter = 500\n",
+    ").plot()\n",
+    "\n",
+    "print(surface.parameters())\n",
+    "\n",
+    "tlm.show3d(optics, sampling={\"base\": 100})\n",
+    "\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b49cfbb7-81f8-42c9-a53a-446d8725ae0b",
+   "id": "cd2555b0-878c-445d-b58e-9bee05bb112c",
    "metadata": {},
    "outputs": [],
    "source": []