Julia: Broadcasting for 2 Dimensional Array based on each element's indices

Viewed 85

I have a 1920 by 1080 array canvas with each element acting as a pixel, with default value RGB(0, 0, 0), for an image.

What my program (a very simple raytracer) currently does is:

  • Use the index of each "pixel" to get the corresponding position on the camera frame
  • store the ray from the center of the camera to that position as a three-dimensional vector
  • find the closest object that the ray intersects with
  • set the value of that pixel to match the color of that object (which again require the indices of that pixel)

The full code is too large to share, but most of the code relevant to the question is as follows:

world = World(RGB(0, 0, 0), 5e-6, objects, lights, 0.2, 4)
camera = Camera((0, -5000, -5000), 1000, (0, 0, 0), 1920, 1080)
canvas = CUDA.fill(world.background, camera.height, camera.width)

for i in 1:camera.width
    for j in 1:camera.height
        ray = [([i, j] - [camera.width / 2, camera.height / 2])..., camera.depth]

        (ray[2], ray[3]) = (cos(camera.rotation[1] + atan(ray[3], ray[2])), sin(camera.rotation[1] + atan(ray[3], ray[2]))) .* sqrt(ray[2]^2 + ray[3]^2)
        (ray[1], ray[3]) = (cos(camera.rotation[2] + atan(ray[3], ray[1])), sin(camera.rotation[2] + atan(ray[3], ray[1]))) .* sqrt(ray[1]^2 + ray[3]^2)
        (ray[1], ray[2]) = (cos(camera.rotation[3] + atan(ray[2], ray[1])), sin(camera.rotation[3] + atan(ray[2], ray[1]))) .* sqrt(ray[2]^2 + ray[1]^2)

        v = (Inf, nothing, nothing)

        for object in world.objects
            # traceray returns a tuple containing the distance of the point of intersection, and the normal to the surface at that point if the ray intersects the object, or nothing otherwise
            t = traceray(ray, camera.position, object, mindistance=camera.depth)
            t !== nothing && t[1] < v[1] && (v = (t[1], t[2], object))
        end

        # computecolor gets the color of the object based on lighting etc.
        v[1] != Inf && (canvas[j, i] = computecolor(v[3].material, ray, v[1], v[2], world, camera.position .+ v[1] * ray, v[3]))
    end
end

Currently, I do this by using nested for loops to loop over all the indices. However, I wanted to improve performance by using CUDA.jl. I would like to avoid writing a GPU kernel and make do with Array Programming, if possible. However, in that case, indexing over the array is not possible since scalar indexing has to be avoided, and from what it seems to me, broadcasting over an array is the preferred method.

From my understanding, a broadcasted function can return a value based on the elements' value (in this case RGB(0, 0, 0)) and not the indices (which is what I require to compute the color).

Is there any way to have a function that wraps the steps mentioned above receive the indices of the element it is acting upon via broadcasting or any alternative method that achieves what I want?

PS: I did look into Computing array values based on their index using Julia Broadcasting in julia but that approach alters only a row (which is a vector by itself), and the expression is quite simple, and I could not get something similar working for a matrix and the large number of statements my program uses.

1 Answers

Let's first create a matrix and get its CartesianIndices:

julia> M = rand(10, 10)
10×10 Matrix{Float64}:
 0.343127  0.0461174  0.000208743  0.739489  0.0881014  …  0.744987   0.150881   0.769042   0.951578
 0.415609  0.919787   0.917708     0.31243   0.659583      0.579605   0.428906   0.880912   0.0689765
 0.887541  0.308289   0.306524     0.10873   0.0731005     0.0678841  0.560323   0.684129   0.601884
 0.614483  0.817356   0.196868     0.863195  0.364332      0.0697723  0.599979   0.709638   0.598244
 0.856793  0.37702    0.357844     0.65256   0.254954      0.0564696  0.49687    0.704261   0.488966
 0.233914  0.0196481  0.201217     0.560289  0.327611   …  0.957368   0.41253    0.0796238  0.774057
 0.905422  0.227343   0.048688     0.374488  0.145598      0.0120923  0.959117   0.83804    0.413381
 0.518771  0.298739   0.0497088    0.947807  0.511978      0.0194308  0.0520221  0.40886    0.316115
 0.150333  0.3123     0.564887     0.869185  0.406472      0.822917   0.611867   0.396063   0.255475
 0.628552  0.634305   0.484303     0.452939  0.0510651     0.207345   0.143286   0.869137   0.626775

julia> ind = CartesianIndices(M)
10×10 CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}:
 CartesianIndex(1, 1)   CartesianIndex(1, 2)   …  CartesianIndex(1, 9)   CartesianIndex(1, 10)
 CartesianIndex(2, 1)   CartesianIndex(2, 2)      CartesianIndex(2, 9)   CartesianIndex(2, 10)
 CartesianIndex(3, 1)   CartesianIndex(3, 2)      CartesianIndex(3, 9)   CartesianIndex(3, 10)
 CartesianIndex(4, 1)   CartesianIndex(4, 2)      CartesianIndex(4, 9)   CartesianIndex(4, 10)
 CartesianIndex(5, 1)   CartesianIndex(5, 2)      CartesianIndex(5, 9)   CartesianIndex(5, 10)
 CartesianIndex(6, 1)   CartesianIndex(6, 2)   …  CartesianIndex(6, 9)   CartesianIndex(6, 10)
 CartesianIndex(7, 1)   CartesianIndex(7, 2)      CartesianIndex(7, 9)   CartesianIndex(7, 10)
 CartesianIndex(8, 1)   CartesianIndex(8, 2)      CartesianIndex(8, 9)   CartesianIndex(8, 10)
 CartesianIndex(9, 1)   CartesianIndex(9, 2)      CartesianIndex(9, 9)   CartesianIndex(9, 10)
 CartesianIndex(10, 1)  CartesianIndex(10, 2)     CartesianIndex(10, 9)  CartesianIndex(10, 10)

Now let's define a function that will take a CartesianIndex and an element:

f(I::CartesianIndex, m::Float64) = I[1] + I[2] + m

We can then apply our new function to ind and M:

julia> f.(ind, M)
10×10 Matrix{Float64}:
  2.34313   3.04612   4.00021   5.73949   6.0881    7.75464   8.74499   9.15088  10.769   11.9516
  3.41561   4.91979   5.91771   6.31243   7.65958   8.04954   9.57961  10.4289   11.8809  12.069
  4.88754   5.30829   6.30652   7.10873   8.0731    9.79731  10.0679   11.5603   12.6841  13.6019
  5.61448   6.81736   7.19687   8.86319   9.36433  10.7867   11.0698   12.6      13.7096  14.5982
  6.85679   7.37702   8.35784   9.65256  10.255    11.3884   12.0565   13.4969   14.7043  15.489
  7.23391   8.01965   9.20122  10.5603   11.3276   12.7932   13.9574   14.4125   15.0796  16.7741
  8.90542   9.22734  10.0487   11.3745   12.1456   13.8769   14.0121   15.9591   16.838   17.4134
  9.51877  10.2987   11.0497   12.9478   13.512    14.9054   15.0194   16.052    17.4089  18.3161
 10.1503   11.3123   12.5649   13.8692   14.4065   15.5465   16.8229   17.6119   18.3961  19.2555
 11.6286   12.6343   13.4843   14.4529   15.0511   16.0112   17.2073   18.1433   19.8691  20.6268

Importantly ind and M have exactly the same dimensions.

Related