Refactor raytrace functions to streamline geometry handling

dominic-chang · Jan 26, 2025 · aff0910 · aff0910
1 parent d0fa0b7
commit aff0910
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 20 deletions.
diff --git a/src/geometries/geometry_types.jl b/src/geometries/geometry_types.jl
@@ -43,31 +43,12 @@ struct ConeGeometry{T, A} <: AbstractGeometry
     ConeGeometry(opening_angle::T, attributes::A) where {T,A} = new{T, A}(opening_angle,attributes)
 end
 
-function raytrace(pixel::AbstractPixel, mesh::Mesh{<:ConeGeometry{T,A}, <:AbstractMaterial}) where {T,A}
-    return_trait = returnTrait(mesh.material)
-    return raytrace(return_trait, pixel, mesh)
-end
-
-function raytrace(::AbstractReturnTrait, pix::AbstractPixel, mesh::Mesh{<:ConeGeometry{T,A}, <:AbstractMaterial}) where {T,A}
-    observation = zero(T)
-    return _raytrace(observation, pix, mesh)
-end
-
-function raytrace(::SimplePolarizationTrait, pix::AbstractPixel, mesh::Mesh{<:ConeGeometry{T,A}, <:AbstractMaterial}) where {T,A}
-    observation = StokesParams(zero(T), zero(T), zero(T), zero(T))
-    return _raytrace(observation, pix, mesh)
-end
-
-function _raytrace(observation, pix::AbstractPixel, mesh::Mesh{<:ConeGeometry{T,A}, <:AbstractMaterial}) where {T,A}
-    #(;magnetic_field, fluid_velocity, spectral_index, R, p1, p2, subimgs) = linpol
-
+function _raytrace(observation, pix::AbstractPixel, mesh::Mesh{<:ConeGeometry{T,A}, <:AbstractMaterial}; res) where {T,A}
     geometry = mesh.geometry
     material = mesh.material
     θs = geometry.opening_angle
     subimgs= material.subimgs
 
-    #observation = yield(material)#StokesParams(zero(T), zero(T), zero(T), zero(T))
-
     isindir = false
     for _ in 1:2 # Looping over isindir this way is needed to get Metal to work
         isindir ⊻= true

diff --git a/src/schemes/RayTrace.jl b/src/schemes/RayTrace.jl
@@ -87,3 +87,25 @@ function dot(a,b)
     return a[1]*b[1] + a[2]*b[2] + a[3]*b[3]
 end
 
+function raytrace(pixel::AbstractPixel, mesh::Mesh; res=100)
+    return_trait = returnTrait(mesh.material)
+    return raytrace(return_trait, pixel, mesh; res=res)
+end
+
+function raytrace(::AbstractReturnTrait, pix::AbstractPixel{T}, mesh::Mesh; res=100) where {T}
+    observation = zero(T)
+    return _raytrace(observation, pix, mesh; res=res)
+end
+
+function raytrace(::SimplePolarizationTrait, pix::AbstractPixel{T}, mesh::Mesh; res=100) where {T}
+    observation = StokesParams(zero(T), zero(T), zero(T), zero(T))
+    return _raytrace(observation, pix, mesh; res=res)
+end
+
+function raytrace(camera::AbstractCamera, mesh::Mesh{A}; res=100) where A
+    intersections = Array{typeof(yield(mesh.material))}(undef, size(camera.screen.pixels)...)#zeros(yield(mesh.material), size(camera.screen.pixels))
+    Threads.@threads for I in CartesianIndices(camera.screen.pixels)
+        intersections[I] = raytrace(camera.screen.pixels[I], mesh; res=res)
+    end
+    return intersections
+end