DiffPointRasterisation

Differentiable rasterisation of point clouds in julia

About

This package provides a rasterisation routine for arbitrary-dimensional point cloud data that is fully (auto-)differentiable. The implementation uses multiple threads on CPU or GPU hardware if available.

The roots of this package are in single-particle 3d reconstruction from tomographic data, and as such it comes with the following ins and out:

Currently only a single "channel" (e.g gray scale) is supported
If data is projected to a lower dimension (meaning the dimensionality of the output grid is lower than the dimensionality of the points)
- Projections are always orthographic (currently no support for perspective projection)
- no "z-blending": The contribution of each point to its output voxel is independent of the point's coordinate along the projection axis.

Rasterisation interface

The interface consists of a single function raster that accepts a point cloud (as a vector of m-dimensional vectors) and pose/projection parameters, (as well as optional weight and background parameters) and returns a n-dimensional (n <= m) array into which the points are rasterized, each point by default with a weight of 1 that is mulit-linearly interpolated into the neighboring grid cells.

Simple Example

julia> using DiffPointRasterisation, LinearAlgebra

julia> grid_size = (5, 5)  # 2d grid with 5 x 5 pixels
(5, 5)

julia> rotation, translation = I(2), zeros(2)  # pose parameters
(Bool[1 0; 0 1], [0.0, 0.0])

julia> raster(grid_size, [zeros(2)], rotation, translation)  # single point at center
5×5 Matrix{Float64}:
 0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0
 0.0  0.0  1.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0

julia> raster(grid_size, [[0.2, 0.0]], rotation, translation)  # single point half a pixel below center
5×5 Matrix{Float64}:
 0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.5  0.0  0.0
 0.0  0.0  0.5  0.0  0.0
 0.0  0.0  0.0  0.0  0.0

julia> raster(grid_size, [[0.2, -0.2]], I(2), zeros(2))  # single point half a pixel below and left of center
5×5 Matrix{Float64}:
 0.0  0.0   0.0   0.0  0.0
 0.0  0.0   0.0   0.0  0.0
 0.0  0.25  0.25  0.0  0.0
 0.0  0.25  0.25  0.0  0.0
 0.0  0.0   0.0   0.0  0.0

Differentiability

Both, an explicit function that calculates derivatives of raster, as well as an integration to common automatic differentiation libraries in julia are provided.

Automatic differentiation libraries

This package provides rules for reverse-mode automatic differentiation libraries that are based on the ChainRules.jl ecosystem. So using raster(args...) in a program that uses any of the ChainRules-based reverse-mode autodiff libraries should just work™. Gradients with respect to all parameters (except grid_size) are supported:

Example using Zygote.jl

we define a simple scalar "loss" that we want to differentiate:

julia> using Zygote

julia> target_image = rand(grid_size...)
5×5 Matrix{Float64}:
 0.345889  0.032283  0.589178   0.0625972  0.310929
 0.404836  0.573265  0.350633   0.0417926  0.895955
 0.174528  0.127555  0.0906833  0.639844   0.832502
 0.189836  0.360597  0.243664   0.825484   0.667319
 0.672631  0.520593  0.341701   0.101026   0.182172

julia> loss(params...) = sum((target_image .- raster(grid_size, params...)).^2)
loss (generic function with 1 method)

some input parameters to raster:

julia> points = [2 * rand(2) .- 1 for _ in 1:5]  # 5 random points
5-element Vector{Vector{Float64}}:
 [0.8457397177007744, 0.3482756109584688]
 [-0.6028188536164718, -0.612801322279686]
 [-0.47141692007256464, 0.6098964840013308]
 [-0.74526926786903, 0.6480225109030409]
 [-0.4044384373422192, -0.13171854413805173]

julia> rotation = [  # explicit matrix for rotation (to satisfy Zygote)
           1.0 0.0
           0.0 1.0
       ]
2×2 Matrix{Float64}:
 1.0  0.0
 0.0  1.0

and let Zygote calculate the gradient of loss with respect to those parameters

julia> d_points, d_rotation, d_translation = Zygote.gradient(loss, points, rotation, translation);

ulia> d_points
5-element Vector{StaticArraysCore.SVector{2, Float64}}:
 [-2.7703628931165025, 3.973371400200988]
 [-0.70462225282373, 1.0317734946448016]
 [-1.7117138793471494, -3.235178706903591]
 [-2.0683933141077886, -0.6732149105779637]
 [2.6278388385655904, 1.585621066861592]

julia> d_rotation
2×2 reshape(::StaticArraysCore.SMatrix{2, 2, Float64, 4}, 2, 2) with eltype Float64:
 -0.632605  -3.26353
  4.12402   -1.86668

julia> d_translation
2-element StaticArraysCore.SVector{2, Float64} with indices SOneTo(2):
 -4.62725350082958
  2.6823723442258274

Explicit interface

The explicit interface for calculating derivatives of raster with respect to its arguments again consists of a single function called raster_pullback!:

The function raster_pullback!(ds_dout, raster_args...) takes as input the sensitivity of some scalar quantity to the output of raster(grid_size, raster_args...), ds_dout, and returns the sensitivity of said quantity to the input arguments raster_args of raster (hence the name pullback).

Example

We can repeat the above calculation using the explicit interface. To do that, we first have to calculate the sensitivity of the scalar output of loss to the output of raster. This could be done using an autodiff library, but here we do it manually:

julia> ds_dout = 2 .* (target_image .- raster(grid_size, points, rotation, translation))  # gradient of loss \w respect to output of raster
5×5 Matrix{Float64}:
 0.152276  -0.417335   1.16347   -0.700428  -0.63595
 0.285167   0.033845  -0.625258  -0.801198   0.760124
 0.349055   0.25511    0.181367   1.27969    1.665
 0.379672   0.721194   0.487329   1.65097    1.33464
 1.34526    1.04119    0.454354  -1.3402     0.364343

Then we can feed that into raster_pullback! together with the same arguments that were fed to raster:

julia> ds_din = raster_pullback!(ds_dout, points, rotation, translation);

We see that the result is the same as obtained via Zygote:

julia> ds_din.points
2×5 Matrix{Float64}:
  2.77036   0.704622  1.71171  2.06839   -2.62784
 -3.97337  -1.03177   3.23518  0.673215  -1.58562

julia> ds_din.rotation
2×2 StaticArraysCore.SMatrix{2, 2, Float64, 4} with indices SOneTo(2)×SOneTo(2):
  0.632605  3.26353
 -4.12402   1.86668

julia> ds_din.translation
2-element StaticArraysCore.SVector{2, Float64} with indices SOneTo(2):
  4.62725350082958
 -2.6823723442258274

Timings

Points	Images	Pixel	Mode	Fwd time CPU	Fwd time CPUx8*	Fwd time GPU	Bwd time CPU	Bwd time CPUx8*	Bwd time GPU**
10⁴	64	128²	3D → 2D	341 ms	73 ms	15 ms	37 ms	10 ms	1 ms
10⁴	64	1024²	3D → 2D	387 ms	101 ms	16 ms	78 ms	24 ms	2 ms
10⁵	64	128²	3D → 2D	3313 ms	741 ms	153 ms	374 ms	117 ms	9 ms
10⁵	64	1024²	3D → 2D	3499 ms	821 ms	154 ms	469 ms	173 ms	10 ms
10⁵	1	1024³	3D → 3D	493 ms	420 ms	24 ms	265 ms	269 ms	17 ms

* 8 julia threads on 4 hardware threads with hyperthreading

** 1 Nvidia HGX A100 GPU