Tomography with Tents

Reconstructing from projections

A tomograph sees an unknown density $\rho$ of the plane only through its projections: for each angle ${\theta }_a$, the profile ${\nu }_a$ of the mass of $\rho$ along the lines orthogonal to $n_a=(\cos {\theta }_a,\sin {\theta }_a)$, measured on a detector. To reconstruct $\rho$, we look for a cloud of $n$ points $x_1,\dots ,x_n$ of the plane, each carrying a mass $1/n$, whose projections reproduce the measured profiles. A point $x_i$ lands at $x_i\cdot n_a$ on the detector of angle $a$, and what "reproduce" means is a transport: the $W_2$ distance on the line,

$$L(x)=\sum _a{W}_2^2({\mu }_a(x),{\nu }_a),{\mu }_a(x)=\frac{1}{n}\sum _{i}f(\cdot -x_i\cdot n_a).$$

What a point projects to

If $f$ is a dirac, ${\mu }_a$ is a sum of diracs on the line, and $L$ has a kink wherever two projected points cross. We take instead a tent of half-width $w$, $f(s)=(1-|s|/w{)}_{+}/w$: two tents that cross only overlap, and $L$ is smooth in $x$. The transport between the tents and the profile of one angle is exact; the transports of all the angles are one batch. This is sdot.ProjectedSumOfTents against a 1D sdot.Image, transported by sdot.ContinuousOtPlan1d ( see Projected Sources ). The projected positions $x_i\cdot n_a$ are computed on the fly: the [ angles, n ] array is never built, which is what lets $n$ grow to millions.

The loss

The loss is plain jax: the tents are a source, the profiles of all the angles a batched 1D Image, ContinuousOtPlan1d the transport between them, and jax.grad differentiates it with respect to the points. The half-width is a parameter, since it changes during the reconstruction:

def tents_loss( angles, profiles, s_min, dw ):
    """`loss( x, half_width )`, `x` `[ n, 2 ]`, `half_width` a number or `[ n ]`: the sum over the angles of W2^2( the tents projected
    at the angle, the measured profile )"""
    normals = jnp.asarray( np.stack( [ np.cos( angles ), np.sin( angles ) ], 1 ) )

    def loss( x, half_width ):
        num_angle = Axis( ShapeVar( len( angles ) ), name = "num_angle" )           # the batch: one transport per angle
        tents = sdot.ProjectedSumOfTents( points = x, half_widths = jnp.broadcast_to( half_width, ( len( x ), ) ), normal = normals,
                                          batch_axes = [ num_angle ] )
        target = sdot.Image( values = jnp.asarray( profiles ), origin = [ s_min ], pixel_frame = [ [ dw ] ], batch_axes = [ num_angle ] )
        return sdot.ContinuousOtPlan1d( tents, target ).cost.sum()
    return loss


def stretched_tents_loss( angles, profiles, s_min, dw ):
    """`loss( x, M )`, `x` `[ n, 2 ]`, `M` `[ n, 2, 2 ]`: the same, for tents whose half-width at the angle `n` is `| M^T n |` ( the
    shadow of an ellipse: `M = w * identity` is the round tent of half-width `w` )"""
    normals = jnp.asarray( np.stack( [ np.cos( angles ), np.sin( angles ) ], 1 ) )

    def loss( x, M ):
        num_angle = Axis( ShapeVar( len( angles ) ), name = "num_angle" )
        tents = sdot.ProjectedSumOfAnisotropicTents( points = x, matrices = M, normal = normals, batch_axes = [ num_angle ] )
        target = sdot.Image( values = jnp.asarray( profiles ), origin = [ s_min ], pixel_frame = [ [ dw ] ], batch_axes = [ num_angle ] )
        return sdot.ContinuousOtPlan1d( tents, target ).cost.sum()
    return loss

Optimizing: coarse to fine

L-BFGS on a large cloud of random tents ends in a poor local minimum. The remedy is a multiscale: converge on a few tents, split each one into several children a little apart, converge again. The tents narrow along the way, with the spacing of the cloud, down to half a bin of the detector: wide tents see the large scales only, and they smooth the loss where the points are still far from their place. sdot.device_optim.lbfgs runs a whole descent on the device, as one compiled loop, and calls back every few iterations. A descent stops when the points have settled: drift_tol bounds the root mean square of their displacement over the last snapshot_every iterations. It is the points that must converge, not the loss: the loss falls by 99 % in the first 30 iterations and then creeps, while the tents keep sliding along the walls, where it is nearly flat, for hundreds of iterations more. Stopping on the loss leaves them a few bins from their place. Neither more memory for L-BFGS, nor more levels, nor a ramp-filter preconditioner ( the inverse of the smoothing that projections apply to a displacement, as in filtered back projection ) makes it faster: the slow directions are local to the walls, where a filter that is the same everywhere does not reach them.

def coarse_to_fine( loss, dw, nb_tents = ( 1024, 4096 ), max_iter = 1500, on_step = None ):
    """the centers of the tents, coarse to fine: L-BFGS on `nb_tents[ 0 ]` tents thrown at random in the cylinder, then every tent
    split in `n_next / n` children, a little apart, and L-BFGS again. The tents narrow with their spacing, `R / sqrt( n )`. A descent
    stops when the points have settled: over 20 iterations they move by less than a twentieth of a bin, in root mean square.
    `on_step( stage, k, x, half_width )` sees every 20 to 100 iterations."""
    rng = np.random.default_rng( 1 )
    r, a = R * np.sqrt( rng.random( nb_tents[ 0 ] ) ), 2 * np.pi * rng.random( nb_tents[ 0 ] )
    x = jnp.asarray( np.stack( [ r * np.cos( a ), r * np.sin( a ) ], 1 ) )
    for level, n in enumerate( nb_tents ):
        if len( x ) != n:
            x = jnp.repeat( x, n // len( x ), axis = 0 ) + ( rng.random( ( n, 2 ) ) - 0.5 ) * R / np.sqrt( n )
        width = R / np.sqrt( n )
        value_and_grad = jax.jit( jax.value_and_grad( lambda x: loss( x, width ) ) )
        x, info = lbfgs( value_and_grad, x, max_iter = max_iter, ftol = 0, min_iter = 20, drift_tol = dw / 20, snapshot_every = 20 if n <= 1024 else 40,
                         on_snapshot = None if on_step is None else lambda k, x: on_step( level, k, np.asarray( x ), width ) )
    return x, width

Stretching the tents

Once the points have settled, the shape of the tents is the next unknown: one half-width for all is a guess, and a round tent cannot follow a wall that is not round. Let each tent have a matrix: sdot.ProjectedSumOfAnisotropicTents gives it, at the direction $n$, the half-width $|M^{T}n|$, the shadow of the ellipse $M(\text{disk})$ ( see Projected Sources ). $M=L{L}^T/w+(dw/4)I$ keeps it symmetric and positive definite, with the three entries of the lower triangular $L$ free ( divided by the half-width $w$ of the last level so that $L$ is a length, like the points ). The points stay where they are: L-BFGS runs on the matrices alone, starting from the round tents of the last level.

def stretch_tents( loss, x, width, dw, max_iter = 2000, on_step = None ):
    """every tent gets a matrix of its own, the points staying where they are: `M = L L^T / width + ( dw / 4 ) identity` ( symmetric,
    positive definite, with `L` lower triangular and free ), starting from the round tents of half-width `width`"""
    wmin = dw / 4
    def matrix( l ):
        L = jnp.zeros( ( len( l ), 2, 2 ) ).at[ :, 0, 0 ].set( l[ :, 0 ] ).at[ :, 1, 0 ].set( l[ :, 1 ] ).at[ :, 1, 1 ].set( l[ :, 2 ] )
        return L @ jnp.swapaxes( L, 1, 2 ) / width + wmin * jnp.eye( 2 )
    report = on_step or ( lambda *args: None )

    r = jnp.full( len( x ), np.sqrt( width * ( width - wmin ) ) )           # L = r identity, M = width identity
    l0 = jnp.stack( [ r, jnp.zeros_like( r ), r ], 1 )
    value_and_grad = jax.jit( jax.value_and_grad( lambda l: loss( x, matrix( l ) ) ) )
    l, _ = lbfgs( value_and_grad, l0, max_iter = max_iter, ftol = 0, min_iter = 20, drift_tol = dw / 40,
                  snapshot_every = 20, on_snapshot = lambda k, l: report( 2, k, x, matrix( l ) ) )
    return matrix( l )

The problem is nearly convex in the matrices: the descent settles by itself in about 500 iterations, some twenty seconds. On this object the loss goes from $2.0\cdot {10}^{-6}$ ( round tents of half-width 6.7 bins ) to $6.0\cdot {10}^{-7}$, a factor 3.4. The median semi-axes are 5.0 and 7.4 bins: a typical tent is elongated by a factor 1.5, along the walls.

A cylinder pierced by bubbles

The object is a cylinder of material pierced by 63 bubbles of various sizes, packed as tight as can be, with $1/120$ of the diameter of the cylinder between two neighbours. It is the kind of data where the method pays: three quarters of the section are empty, and the material left is a network of thin walls that a good reconstruction must follow without filling the bubbles. The data are 600 profiles on a detector of 1024 bins.

The page shows the iterations recorded by the callback ( one in twenty ), from the 1024 tents thrown at random in the cylinder to the 4096 of the last level, then the matrices at fixed points ( play, or drag the slider ); each stage stops by itself when it has settled. Each tent is drawn as a disk of its half-width, which shrinks from level to level, then as the outline of its ellipse in the last stage.

Open in a separate page ↗

The whole takes a little over a minute on a laptop once the kernels are compiled. The card and 3D from 2D projections are on the roadmap.

The complete script is docs/examples/tomography_tents.py.