Add depth_max keyword to solve_system! (#41)

* Add depth_max keyword to solve_system!

* Improve branching code a bit and add a test case

Author: TacHawkes

src/System.jl

@@ -433,33 +433,47 @@ function retrace_system!(system::AbstractSystem, gauss::GaussianBeamlet{T}) wher
 end
 
 """
-    solve_system!(system::System, beam::AbstractBeam; r_max=100, retrace=true)
+    solve_system!(system::System, beam::AbstractBeam; r_max=100, retrace=true, depth_max=typemax(Int))
 
 Manage the tracing of an `AbstractBeam` through an optical `system`. The function retraces the `beam` if possible and then proceeds to trace each leaf of the beam tree through the system.
 The condition to stop ray tracing is that the last `beam` intersection is `nothing` or the beam interaction is `nothing`. Then, the system is considered to be solved.
-A maximum number of rays per `beam` (`r_max`) can be specified in order to avoid infinite calculations under resonant conditions, i.e. two facing mirrors.
+A maximum number of rays per `beam` (`r_max`) can be specified in order to avoid infinite calculations under resonant conditions, i.e. two facing mirrors. Likewise, `depth_max` limits how many branching levels are explored when new sub-beams are generated (for example, by beamsplitters) so that the tree cannot grow without bound.
 
 # Arguments
 
 - `system::System`: The optical system in which the beam will be traced.
 - `beam::AbstractBeam`: The beam object to be traced through the system.
 - `r_max::Int=100` (optional): Maximum number of tracing iterations for each leaf. Default is 100.
 - `retrace::Bool=true` (optional): Flag to indicate if the system should be retraced. Default is true.
+- `depth_max::Int=typemax(Int)` (optional): Maximum number of branching levels explored from the root beam. Default allows unlimited depth.
 """
-function solve_system!(system::AbstractSystem, beam::B; r_max::Int = 100, retrace::Bool = true) where {B <: AbstractBeam}
-    # Initialize a queue for BFS with the root beam.
-    queue = [beam]
+function solve_system!(
+        system::AbstractSystem,
+        beam::B;
+        r_max::Int = 100,
+        retrace::Bool = true,
+        depth_max::Int = typemax(Int),
+    ) where {B <: AbstractBeam}
+    queue = Tuple{B, Int}[(beam, 1)]
     while !isempty(queue)
-        current = popfirst!(queue)  # Process beams in FIFO order.
-        # Optionally retrace the current beam.
+         # Process beams in FIFO order.
+        current, depth = popfirst!(queue)
+         # Optionally retrace the current beam.
         if retrace
             retrace_system!(system, current)
         end
         # Process the current leaf beam.
         solve_leaf!(system, current; r_max=r_max)
-        # Enqueue all child beams for subsequent processing.
-        for child in children(current)  # 'children' returns an iterable of sub-beams.
-            push!(queue, child)
+        # Check if the maximum branching depth has been reached.
+        if depth <= depth_max
+             # Enqueue all child beams for subsequent processing.
+            for child in children(current) # 'children' returns an iterable of sub-beams.
+                push!(queue, (child, depth + 1))
+            end
+        else
+            # Maximum braching depth is reached. Remove childrens of the current beam because they will not be solved.
+            _drop_beams!(current)
+            @debug lazy"Maximum branching depth of $depth_max levels reached."
         end
     end
     return nothing

test/Components/TestBeamsplitters.jl

@@ -122,6 +122,20 @@ const mm = 1e-3
             @test BMO.direction(last(r)) ≈ [-1, 0, 0]
         end
     end
+
+    @testset "depth_max branch limiting" begin
+        beamsplitter = CubeBeamsplitter(25e-3, n -> N0)
+        translate3d!(beamsplitter, [0, 50mm, 0])
+        system = System([beamsplitter])
+
+        beam = Beam([0, 0, 0], [0, 1, 0], 1e-6)
+        solve_system!(system, beam; depth_max=0)
+        @test isempty(beam.children)
+
+        beam = Beam([0, 0, 0], [0, 1, 0], 1e-6)
+        solve_system!(system, beam; depth_max=1)
+        @test length(beam.children) == 2
+    end
 end
 
 end