batch rays together, decimate allocations and speedup ~3x

.gitignore

@@ -1,4 +1,6 @@
 Manifest.toml
 **/*.vtk
 **/*.stl
-**/*.csv
+**/*.csv
+**/*.txt
+**/*.gz

Project.toml

@@ -3,6 +3,7 @@ authors = ["Freddie Barter <fjbarter@outlook.com>"]
 uuid = "e89368bc-bc31-4f60-a9df-df669e2c49e6"
 
 [deps]
+AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"

README.md

@@ -10,13 +10,14 @@ using Pkg
 Pkg.develop(url="https://github.com/fjbarter/PRICK.jl")
 ```
 
-Then you can simply load in your sphere and triangle-wall data, define a start point and direction as SVectors, as below:
+Then you can load sphere and triangle-wall data, define one start point, and trace a batch of directions:
 ```julia
 using PRICK
 
 # Load mesh and tag each surface as a mirror or a sink
 vessel_tm = TriangleMesh("RAMCylinder.stl"; units=u"m")
 surfaces = [mirror(vessel_tm)]  # or sink(vessel_tm)
+surface_bvh = build_surface_bvh(surfaces)  # build once and reuse
 
 # Load spheres (example using Packing3D VTK reader)
 data = read_vtk_file("particles_0.vtk")
@@ -27,28 +28,24 @@ radii = Float64.(r)
 # Build the spheres together and construct BVH
 spheres = build_sphere_bvh(X, radii)
 
-# Trace a ray
-p0 = SVector(0.0079, 0.0, 0.022)
-d = SVector(1.0, 0.0, 0.0)
-res = trace_ray_geometric(p0, d, surfaces, spheres)
-```
-
-## Visualisation (optional)
-
-`visualise_trace` requires GLMakie to be loaded by the user:
-
-```julia
-using GLMakie
-fig = visualise_trace(
-    res;
-    X=X, radii=radii,
-    vessel_tm=vessel_tm
+# Trace N rays from a single start point p0 with a (3, N) direction matrix
+p0 = (0.0079, 0.0, 0.022)
+N = 10_000
+D = randn(3, N)  # directions are normalized internally
+
+res = trace_rays(
+    p0,
+    D,
+    surface_bvh,
+    spheres;
+    max_bounces=20_000,
+    max_length=1e6,
 )
-display(fig)
 ```
 
 ## Notes
 
 - `mirror` surfaces reflect specularly.
 - `sink` surfaces terminate rays.
 - The union mesh bounding box is used as an unmeshed sink and is expanded by 10%.
+- Batched tracing returns per-ray `escaped`, `total_length`, `nsteps`, and `termination`.

prototyping/cylinder/Project.toml

@@ -1,3 +1,4 @@
 [deps]
-GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+PProf = "e4faabce-9ead-11e9-39d9-4379958e3056"
 PRICK = "e89368bc-bc31-4f60-a9df-df669e2c49e6"

prototyping/cylinder/allocation_flamegraph.jl

@@ -0,0 +1,71 @@
+
+using Profile
+using PRICK
+
+using PProf
+
+function random_unit_directions(n::Integer)
+    n > 0 || throw(ArgumentError("n must be > 0"))
+    D = Matrix{Float64}(undef, 3, n)
+    @inbounds for i in 1:n
+        phi = 2pi * rand()
+        u = 2rand() - 1
+        s = sqrt(1 - u * u)
+        D[1, i] = s * cos(phi)
+        D[2, i] = s * sin(phi)
+        D[3, i] = u
+    end
+    return D
+end
+
+function main(;
+    n_rays::Int=1_000,
+    max_bounces::Int=20_000,
+    max_length::Float64=1e6,
+    sample_rate::Float64=0.01,
+    flat_report_path::String="alloc_profile_flat.txt",
+    tree_report_path::String="alloc_profile_tree.txt",
+)
+    println("Building geometry...")
+    vessel_trianglemesh = TriangleMesh("RAMCylinder.stl"; units=u"m")
+    surfaces = [sink(vessel_trianglemesh)]
+    surface_bvh = build_surface_bvh(surfaces)
+
+    println("Reading particles...")
+    data = read_vtk_file("particles_0.vtk")
+    x, y, z, r = retrieve_coordinates(data)
+    X = permutedims(Float64.(hcat(x, y, z)))
+    radii = Float64.(r)
+    spheres = build_sphere_bvh(X, radii)
+
+    p0 = (PRICK.mean(x), PRICK.mean(y), PRICK.mean(z))
+    D = random_unit_directions(n_rays)
+
+    println("Warmup run...")
+    trace_rays(p0, D, surface_bvh, spheres; max_bounces=max_bounces, max_length=max_length)
+
+    println("Profiling allocations (sample_rate=$sample_rate)...")
+    Profile.Allocs.clear()
+    Profile.Allocs.@profile sample_rate=sample_rate begin
+        trace_rays(p0, D, surface_bvh, spheres; max_bounces=max_bounces, max_length=max_length)
+    end
+
+    prof = Profile.Allocs.fetch()
+    println("Collected allocation samples: ", length(prof.allocs))
+
+    println("Writing text allocation reports...")
+    open(flat_report_path, "w") do io
+        Profile.Allocs.print(io, prof; format=:flat, C=false, maxdepth=12, mincount=1)
+    end
+    open(tree_report_path, "w") do io
+        Profile.Allocs.print(io, prof; format=:tree, C=false, maxdepth=20, mincount=1)
+    end
+    println("Wrote: ", abspath(flat_report_path))
+    println("Wrote: ", abspath(tree_report_path))
+
+    println("Opening pprof web UI...")
+    PProf.Allocs.pprof(prof; from_c=false)
+    println("If browser does not open automatically, use the URL printed by PProf.")
+end
+
+main()

prototyping/cylinder/cylinder.jl

@@ -1,92 +1,63 @@
-using PRICK
-
 # ============================================================
-# example() demonstrating:
-# - provide X, radii
-# - provide vessel
-# - trace a ray and return results
+# Example batched ray tracing in a sphere bed
 # ============================================================
-function example()
-    # Vessel from stl file
-    vessel_trianglemesh = TriangleMesh("RAMCylinder.stl"; units=u"m")
-
-    # Construct surface (each is either a sink or mirror)
-    surfaces = [
-        sink(vessel_trianglemesh),
-    ]
-
-    # Spheres from VTK file
-    println("Reading file...")
-    file = "particles_0.vtk"
-    data = read_vtk_file(file)
-    x, y, z, r = retrieve_coordinates(data)
-
-    X = permutedims(Float64.(hcat(x, y, z)))   # 3×N
-    radii = Float64.(r)
-    spheres = build_sphere_bvh(X, radii)
-
-    # Build p0 from mean sphere position
-    p0x = PRICK.mean(x)
-    p0y = PRICK.mean(y)
-    p0z = PRICK.mean(z)
-    p0 = SVector{3,Float64}(p0x, p0y, p0z)
-
-    # p0 = SVector{3,Float64}(0.0079, 0.0, 0.022)
-
-    N = 1
-    path_lengths = zeros(Float64, N)
-    n_steps = zeros(Float64, N)
-    println("Tracing $N rays...")
-    for i in 1:N
-        # Fixed start point
-
-        φ = 2π * rand()          # azimuth
-        u = 2rand() - 1          # u = cos(θ) uniformly in [-1,1]
-        s = sqrt(1 - u*u)
-        d = SVector(s*cos(φ), s*sin(φ), u)
-
-        res = trace_ray_geometric(p0, d, surfaces, spheres;
-                                  record_path=false,
-                                  max_steps=20000,
-                                  max_length=1e6)
-
-        path_lengths[i] = res.escaped ? res.total_length : NaN
-        n_steps[i] = res.escaped ? res.nsteps : NaN
+using PRICK
+using BenchmarkTools
+
+function random_unit_directions(n::Integer)
+    n > 0 || throw(ArgumentError("n must be > 0"))
+    D = Matrix{Float64}(undef, 3, n)
+    @inbounds for i in 1:n
+        phi = 2pi * rand()
+        u = 2rand() - 1
+        s = sqrt(1 - u * u)
+        D[1, i] = s * cos(phi)
+        D[2, i] = s * sin(phi)
+        D[3, i] = u
     end
-
-    path_lengths_dense = filter(!isnan, path_lengths)
-
-    mean_length = PRICK.mean(path_lengths_dense)
-    println("Mean path length over $N rays (escaped: $(length(path_lengths_dense))): $(mean_length) m")
-
-    n_steps_dense = filter(!isnan, n_steps)
-    mean_number_of_steps = PRICK.mean(n_steps_dense)
-    println("Mean number of steps over escaped rays: $(mean_number_of_steps)")
-
-    # Single ray trace example
-
-    # Start point and direction
-    p0 = SVector{3,Float64}(0.0079, 0.0, 0.022)
-    φ = 2π * rand()          # azimuth
-        u = 2rand() - 1          # u = cos(θ) uniformly in [-1,1]
-        s = sqrt(1 - u*u)
-        d = SVector(s*cos(φ), s*sin(φ), u)
-
-    res = trace_ray_geometric(p0, d, surfaces, spheres;
-                              record_path=true,
-                              max_steps=20000,
-                              max_length=1e6)
-
-    println("escaped       = ", res.escaped)
-    println("total_length  = ", res.total_length)
-    println("nsteps        = ", res.nsteps)
-    println("path vertices = ", length(res.path))
-
-    # If you want to plot later:
-    P = path_to_matrix(res.path)  # 3xK
-    return (res=res, X=X, radii=r, vessel_tm=vessel_trianglemesh, path_matrix=P)
+    return D
 end
 
-out = example()
-# fig = visualise_trace(out.res; X=out.X, radii=out.radii, vessel_tm=out.vessel_tm)
-# display(fig)
+vessel_trianglemesh = TriangleMesh("RAMCylinder.stl"; units=u"m")
+surfaces = [sink(vessel_trianglemesh)]
+surface_bvh = build_surface_bvh(surfaces)
+
+println("Reading file...")
+data = read_vtk_file("particles_0.vtk")
+x, y, z, r = retrieve_coordinates(data)
+
+X = permutedims(Float64.(hcat(x, y, z)))
+radii = Float64.(r)
+spheres = build_sphere_bvh(X, radii)
+
+p0 = (PRICK.mean(x), PRICK.mean(y), PRICK.mean(z))
+n_rays = 1000
+D = random_unit_directions(n_rays)
+
+println("Tracing $n_rays rays...")
+t_trace0 = time()
+
+@benchmark trace_rays(
+    $p0,
+    $D,
+    $surface_bvh,
+    $spheres;
+    max_bounces=20_000,
+    max_length=1e6,
+)
+# t_trace = time() - t_trace0
+# println("Traced $n_rays rays in $t_trace seconds")
+
+# escaped_lengths = [res.total_length[i] for i in eachindex(res.total_length) if res.escaped[i]]
+# escaped_steps = [res.nsteps[i] for i in eachindex(res.nsteps) if res.escaped[i]]
+# escaped_term_dists = [res.p0_termination_distance[i] for i in eachindex(res.p0_termination_distance) if res.escaped[i]]
+
+# mean_length = isempty(escaped_lengths) ? NaN : PRICK.mean(escaped_lengths)
+# mean_steps = isempty(escaped_steps) ? NaN : PRICK.mean(escaped_steps)
+# mean_tortuosity = sum(escaped_lengths)/sum(escaped_term_dists)
+
+# println("Escaped rays: $(length(escaped_lengths)) / $n_rays")
+# println("Mean path length over escaped rays: $(mean_length) m")
+# println("Mean number of reflections over escaped rays: $(mean_steps)")
+# println("Mean tortuosity: $mean_tortuosity")
+