Repository: HarrisonKramer/optiland

Issue:
Quick Focus
**Is your feature request related to a problem? Please describe.**
A quick focus for the image surface may be added. 

**Describe the solution you'd like**
_Option 1 for Function Placement:_
    Implement the functionality to add_surface function. In the add_surface definition, "quick_focus=True" is present and adjusts the surface length minimizing the RMS spot size for the primary wavelength. 
    Pro(s): Easier implementation for the user, and more apparent for readability.
    Con(s): Harder for GUI implementation?

_Option 2 for Function Placement:_
    Implement the functionality to optic class. After all the surfaces are defined, the user can call a lens.minimize_spot_size function, updating the last surface length.
  Pro(s): It will also be implemented easily by the user. 
  Con(s): Easier for GUI implementation. It will not be more apparent for readability (The optical system is already defined), however if GUI is used, then it is irrelevant.

_Options for RMS Spot Size Minimization:_
    1. Directly use Scipy minimize function, easier to implement, and seems to be secure. However, if Option 1 is preferred, the user can't control the iteration number. 
    2. By the angles of the rays, we can analytically calculate the arrival position, and then take the mean position. Might be tricky for some geometries, such as rotated image plane.

**Additional context**
The function can quickly get out of hand if used improperly. Strict caution should be added. For instance, if the user wants to use quick focus for a beam expander system. It simply can not. An error could be thrown, or an iteration number handles it anyway.

Hints from issue comments:
@CAClaveau: (Not sure, but it might help: [link](https://github.com/singer-yang/DeepLens/blob/24ea83c8d37ec2eab96930c25ac1a2925e47b514/deeplens/geolens.py#L1583))@HarrisonKramer: Thanks for opening this issue, @SeckinBerkay. And thanks for sharing this link, @CAClaveau!

I like the idea of a quick focus adjustment, but I think it would fit better as a `Solve` rather than being directly integrated into `Optic`. This delegates the functionality and keeps the `Optic` class cleaner IMO.

Solves are already set up and can be added to an optic, so we would just need to implement a new Solve class for this. I’m not certain of all the details yet, but using a quick focus calculation like @CAClaveau referenced, could work well.

I think we'll need to:

- Implement a new class, e.g., `QuickFocusSolve`, which subclasses `BaseSolve`.
- Update the `SolveFactory` with the new class.
- Add tests and docs.

Feel free to open a PR if you want to work on this. I can also support wherever needed.

Reference: [API section on solves](https://optiland.readthedocs.io/en/latest/api/core/solves.html#module-solves)@SeckinBerkay: All right, I will try to come up with a working code demo. Then I can create a PR, and then we can change the logic or fix issues if there are any.@HarrisonKramer: I'm thinking a bit more about this actually. I think adding a new solve for this makes sense, but we may also want something like your original proposal - if only a simple, one-time adjustment is needed to focus the system, then adding it as a solve may not always be ideal. A solve enforces some constraint on a surface and it is called in every optimization iteration. If the user really only wants to focus the system once, then adding it as a solve might be overkill, and it could considerably slow optimization. In any case, having it as an option in the solves.py module might still be useful in general.

Perhaps we can consider adding a refinements.py module (or similar), which defines functions to apply one-time adjustments to an Optic instance. That could include a "quick_focus" function, for example. We can maybe add a method to the Optic class that calls this function and passes self as an argument, or it could be completely separate from the class. I want to avoid bloating the Optic class and want to keep new functionalities logically separated from the core Optic class, if possible. Not positive on the best approach here, but open to suggestions. Thanks again!

Kramer@SeckinBerkay: Quick Focus might be useful for automated design. However, as used in LensAI example 1, image_solve is already efficient. Does this enhancement add more to image_solve, or do they serve more or less the same purpose? If new quick focus doesn't add more to image_solve, we can close this issue. @HarrisonKramer: `image_solve` moves the image plane to the paraxial focus position, i.e., where the marginal ray intersects the optical axis. This might be good enough for many cases, but you may want a functionality that uses real rays, considers all fields, all wavelengths, etc. I think it makes sense to eventually add this as a "nice to have". If we don't want to focus on this issue for now, then we can close this and come back to it later. I leave it up to you.@SeckinBerkay: Yeah, it is nice to have it as an optionality, considering it will use real rays rather than the paraxial approximation, and may have the optionality to focus over all the wavelengths and fields.@SeckinBerkay: After a long thoughtful process, I implemented it inside solves.py as you suggested earlier, a simple presentation may be seen in the figure. I will open a PR after writing the tests, and we can fix further if there are possible pitfalls. I used the default OptimizerGeneric module.

![Image](https://github.com/user-attachments/assets/5df21390-7ad3-4b80-b847-c974ccb5c961)@HarrisonKramer: Implemented in #82.

Solution patch:
diff --git a/docs/api/core/solves.rst b/docs/api/core/solves.rst
index 1ede03e5..3a12fc81 100644
--- a/docs/api/core/solves.rst
+++ b/docs/api/core/solves.rst
@@ -10,6 +10,7 @@
    
       BaseSolve
       MarginalRayHeightSolve
+      QuickFocusSolve
       SolveFactory
       SolveManager
    
\ No newline at end of file
diff --git a/optiland/solves.py b/optiland/solves.py
index 98e738d0..b87c5696 100644
--- a/optiland/solves.py
+++ b/optiland/solves.py
@@ -9,7 +9,7 @@ Kramer Harrison, 2024
 """
 
 from abc import ABC, abstractmethod
-
+from numpy import mean, errstate, where, nan, nanmean
 
 class BaseSolve(ABC):
     """
@@ -60,8 +60,7 @@ class BaseSolve(ABC):
             raise ValueError(f'Unknown solve type: {solve_type}')
         solve_class = BaseSolve._registry[data['type']]
         return solve_class.from_dict(optic, data)
-
-
+    
 class MarginalRayHeightSolve(BaseSolve):
     """
     Initializes a MarginalRayHeightSolve object.
@@ -113,6 +112,51 @@ class MarginalRayHeightSolve(BaseSolve):
         """
         return cls(optic, data['surface_idx'], data['height'])
 
+class QuickFocusSolve(BaseSolve):
+    """ Quick Focus
+    Args:
+        optic (Optic): The optic object.
+
+    Raises:
+            ValueError: If the optical system is not defined.
+    """
+    def __init__(self, optic):
+        self.optic = optic
+        self.num_surfaces = self.optic.surface_group.num_surfaces
+        if self.num_surfaces <= 2:
+            raise ValueError('Can not optimize for an empty optical system')
+        
+    def optimal_focus_distance(self, Hx=0, Hy=0, wavelength=0.55, num_rays=5, distribution='hexapolar'):
+        """
+        Compute the optimal location of the image plane where the RMS spot size is minimized.
+        This is based on solving the quadratic equation that describes the RMS spot size as a function
+        of the propagation distance.
+        
+        Args:
+            ...
+            
+        Returns:
+            t_opt : The propagation distance from the image plane that minimizes the RMS spot size.
+        """
+        # Trace rays to the image plane
+        # Trace rays to the image plane
+
+        rays = self.optic.trace(Hx=Hx, Hy=Hy, wavelength=wavelength, num_rays=num_rays, distribution=distribution)
+
+        A = rays.L**2 + rays.M**2
+        B = rays.L * rays.x + rays.M * rays.y
+        with errstate(divide='ignore', invalid='ignore'):
+            t_opt = where(A != 0, -B/A, nan)
+        z_focus = nanmean(rays.z + t_opt * rays.N)
+        
+        return z_focus
+        
+    def apply(self):
+        """Applies QuickFocusSolve to the optic"""
+        z_focus = self.optimal_focus_distance(wavelength=self.optic.wavelengths.primary_wavelength.value)
+        
+        self.optic.surface_group.surfaces[-1].geometry.cs.z = z_focus
+
 
 class SolveFactory:
     """
@@ -126,7 +170,8 @@ class SolveFactory:
             based on the given solve type.
     """
     _solve_map = {
-        'marginal_ray_height': MarginalRayHeightSolve
+        'marginal_ray_height': MarginalRayHeightSolve,
+        'quick_focus': QuickFocusSolve
     }
 
     @staticmethod


Test patch:
diff --git a/tests/test_solves.py b/tests/test_solves.py
index b0288efe..656a9b7f 100644
--- a/tests/test_solves.py
+++ b/tests/test_solves.py
@@ -3,7 +3,6 @@ import numpy as np
 from optiland import solves
 from optiland.samples.objectives import CookeTriplet
 
-
 class TestMarginalRayHeightSolve:
     def test_marginal_ray_height_solve_constructor(self):
         optic = CookeTriplet()
@@ -72,6 +71,37 @@ class TestMarginalRayHeightSolve:
         with pytest.raises(ValueError):
             solves.BaseSolve.from_dict(optic, data)
 
+class TestQuickfocusSolve:
+    def test_quick_focus_solve_constructor(self):
+        optic = CookeTriplet()
+        optic.surface_group.surfaces[-1].geometry.cs.z -= 10
+        solve = solves.QuickFocusSolve(optic)
+
+        assert solve.optic == optic
+
+    def test_quick_focus_solve_apply(self):
+        optic = CookeTriplet()
+        optic.surface_group.surfaces[-1].geometry.cs.z -= 10
+        thickness = 42.21812063592369
+        solve = solves.QuickFocusSolve(optic)
+        solve.apply()
+
+        # Check that surface has been shifted
+        pos2 = optic.surface_group.positions[-1].astype(float)[0]
+        pos1 = optic.surface_group.positions[-2].astype(float)[0]
+
+        assert pos2 - pos1 == pytest.approx(thickness, rel=1e-3)
+
+        # Implementing the extreme shift case.
+        optic.surface_group.surfaces[-1].geometry.cs.z += 1000
+        solve = solves.QuickFocusSolve(optic)
+        solve.apply()
+
+        # Check that surface has been shifted
+        pos2 = optic.surface_group.positions[-1].astype(float)[0]
+        pos1 = optic.surface_group.positions[-2].astype(float)[0]
+
+        assert pos2 - pos1 == pytest.approx(thickness, rel=1e-3)
 
 class TestSolveFactory:
     def test_create_solve(self):