Improve regularisation documentation and add ADMT demo

[jacklovell]

• Add a Regularisation section to the tomography documentation, describing the functionality inside the admt_utils module.
• Fixup the docstrings in admt_utils.
• Add calculation of "skewed" second derivative operators, which operate along the diagonals of the grid.
• Add a bolometer diagonstic system to Generomak.
• Add a new example using the Generomak bolometers and the regularisation operators to perform isotropic and ADMT inversions.
• Add a regularised NNLS function which uses sparse matrices.
• Enable the functions in admt_utils to optionally output sparse matrices depending on the input.

Fixes #382 and #380

[vsnever]

Thanks @jacklovell, very valuable demo. A similar regularisation is used in ITER divertor bolometry to reconstruct the radiation power density profile. Generomak also receives the first diagnostics, which in itself is great.

I think the changelog should reflect that this demo has been added, and also that the bolometric diagnostic model has been added to Generomak.

I also have some comments and a few suggestions.

[vsnever]

I find the interface and implementation of the generate_derivative_operators() a bit complex. In the case of a rectangular grid, it is enough to pass to the function arrays of grid nodes along the x and y axes, and a 2D voxel map. Also, this function can be written in a loop-free way. I propose the following implementation, which seems simpler to me:

def generate_derivative_operators(x, y, voxel_map):
    r"""
    Generate the first and second derivative operators for a regular rectangular grid.

    :param ndarray x: 1D array containing coordinates of grid nodes along the X-axis
    :param ndarray y: 1D array containing coordinates of grid nodes along the Y-axis.
    :param ndarray voxel_map: 2D integer array of shape (x.size - 1, y.size - 1) containing
        the indices of the voxels. This array maps spacially arranged cells to the voxels.
        Mapping multiple cells to the same voxel is not currently supported.
        Cells where `voxel_map == -1` are not mapped to any voxel.

    :return: a dictionary containing the derivative operators: Dij for
        i, j ∊ (x, y) and Di for i ∊ (x, y), Dsp and Dsm.

    Currently, all voxels are assumed to have the same width and height.
    If this is not the case, the results will be nonsense.

    The return dict contains all the first and second derivative
    operators:

    .. math::
        D_{xx} \equiv \frac{\partial^2}{\partial x^2}\\
        D_{xy} \equiv \frac{\partial^2}{\partial x \partial y}

    etc. It also produces two additional operators, Dsp and Dsm, for
    second derivatives on the dy/dx = 1 and dy/dx = -1 diagonals
    respectively.

    Note that the standard 2D laplacian (for isotropic regularisation)
    can be trivially calculated as follows:

    .. math::
        L = (1 - \alpha) (D_{xx} + D_{yy}) + (\alpha / 2) (D_{sp} + D_{sm})

    α = 2/3 produces the operator used in Carr et. al. RSI 89, 083506 (2018).
    α = 1/3 produces the operator with optimal isotropy.
    """

    x = np.asarray(x)
    y = np.asarray(y)

    if x.ndim != 1:
        raise ValueError('Attribute x must be a 1D array-like.')
    if y.ndim != 1:
        raise ValueError('Attribute y must be a 1D array-like.')

    voxel_map = np.asarray(voxel_map, dtype=int)

    dx = np.diff(x)
    dy = np.diff(y)

    if np.any(dx <= 0):
        raise ValueError('X-axis grid nodes must be given in ascending order.')
    if np.any(dy <= 0):
        raise ValueError('Y-axis grid nodes must be given in ascending order.')

    grid_shape = (dx.size, dy.size)
    if voxel_map.shape != grid_shape:
        raise ValueError('Shape mismatch: voxel_map shape {} does not match grid shape {}'.format(voxel_map.shape, grid_shape))

    num_cells = voxel_map.max() + 1

    if np.any(np.sort(voxel_map[voxel_map > -1]) != np.arange(num_cells, dtype=int)):
        raise ValueError('Voxel map contains duplicated cell numbers or some cell numbers are missing.')

    # assume regular grid
    # TODO: support non-regular rectangular grids in the future
    dx = dx.min()
    dy = dy.min()

    # add boundary rows and columns to avoid an out-of-bounds error
    voxel_map_ext = np.full((voxel_map.shape[0] + 2, voxel_map.shape[1] + 2), -1, dtype=int)
    voxel_map_ext[1:-1, 1:-1] = voxel_map

    # get cell indices for each voxel
    indx = np.where(voxel_map_ext > -1)
    voxel_indices = voxel_map_ext[indx]
    isort = np.argsort(voxel_indices)
    cell_indx = (indx[0][isort], indx[1][isort])

    # Individual derivative operators
    Dx = np.zeros((num_cells, num_cells))
    Dy = np.zeros((num_cells, num_cells))
    Dxx = np.zeros((num_cells, num_cells))
    Dxy = np.zeros((num_cells, num_cells))
    Dyy = np.zeros((num_cells, num_cells))
    Dsp = np.zeros((num_cells, num_cells))
    Dsm = np.zeros((num_cells, num_cells))

    # neighbouring voxel indices around each voxel
    n_below = voxel_map_ext[cell_indx[0], cell_indx[1] - 1]
    n_above = voxel_map_ext[cell_indx[0], cell_indx[1] + 1]
    n_left = voxel_map_ext[cell_indx[0] - 1, cell_indx[1]]
    n_right = voxel_map_ext[cell_indx[0] + 1, cell_indx[1]]
    n_below_left = voxel_map_ext[cell_indx[0] - 1, cell_indx[1] - 1]
    n_below_right = voxel_map_ext[cell_indx[0] + 1, cell_indx[1] - 1]
    n_above_left = voxel_map_ext[cell_indx[0] - 1, cell_indx[1] + 1]
    n_above_right = voxel_map_ext[cell_indx[0] + 1, cell_indx[1] + 1]

    # special voxel locations
    at_left = (n_left == -1) & (n_right > -1)
    at_right = (n_right == -1) & (n_left > -1)
    at_bottom = (n_below == -1) & (n_above > -1)
    at_top = (n_above == -1) & (n_below > -1)

    at_top_left = at_left & at_top & (n_below_right > -1)
    at_bottom_left = at_left & at_bottom & (n_above_right > -1)
    at_top_right = at_right & at_top & (n_below_left > -1)
    at_bottom_right = at_right & at_bottom & (n_above_left > -1)

    at_left_not_at_corner = at_left & ~(at_top_left | at_bottom_left)
    at_right_not_at_corner = at_right & ~(at_top_right | at_bottom_right)
    at_bottom_not_at_corner = at_bottom & ~(at_bottom_left | at_bottom_right)
    at_top_not_at_corner = at_top & ~(at_top_left | at_top_right)

    # filling derivative operator matrices
    np.fill_diagonal(Dxx, -2)
    np.fill_diagonal(Dyy, -2)

    indx, = np.where(n_below > -1)
    Dy[indx, n_below[indx]] = -1 / 2
    Dyy[indx, n_below[indx]] = 1

    indx, = np.where(n_above > -1)
    Dy[indx, n_above[indx]] = 1 / 2
    Dyy[indx, n_above[indx]] = 1

    indx, = np.where(n_left > -1)
    Dx[indx, n_left[indx]] = -1 / 2
    Dxx[indx, n_left[indx]] = 1

    indx, = np.where(n_right > -1)
    Dx[indx, n_right[indx]] = 1 / 2
    Dxx[indx, n_right[indx]] = 1

    indx, = np.where(n_below_left > -1)
    Dxy[indx, n_below_left[indx]] = 1 / 4

    indx, = np.where(n_above_left > -1)
    Dxy[indx, n_above_left[indx]] = -1 / 4

    indx, = np.where(n_below_right > -1)
    Dxy[indx, n_below_right[indx]] = -1 / 4

    indx, = np.where(n_above_right > -1)
    Dxy[indx, n_above_right[indx]] = 1 / 4

    indx, = np.where(at_left)
    Dx[indx, indx] = -1
    Dx[indx, n_right[indx]] = 1
    Dxx[indx, indx] = -1
    Dxx[indx, n_right[indx]] = 1

    indx, = np.where(at_left_not_at_corner)
    Dxy[indx, n_above_right[indx]] = 1 / 2
    Dxy[indx, n_below[indx]] = 1 / 2
    Dxy[indx, n_above[indx]] = -1 / 2
    Dxy[indx, n_below_right[indx]] = -1 / 2

    indx, = np.where(at_right)
    Dx[indx, n_left[indx]] = -1
    Dx[indx, indx] = 1
    Dxx[indx, n_left[indx]] = -1
    Dxx[indx, indx] = 1

    indx, = np.where(at_right_not_at_corner)
    Dxy[indx, n_above[indx]] = 1 / 2
    Dxy[indx, n_below_left[indx]] = 1 / 2
    Dxy[indx, n_above_left[indx]] = -1 / 2
    Dxy[indx, n_below[indx]] = -1 / 2

    indx, = np.where(at_top)
    Dy[indx, n_below[indx]] = -1
    Dy[indx, indx] = 1
    Dyy[indx, n_below[indx]] = -1
    Dyy[indx, indx] = 1

    indx, = np.where(at_top_not_at_corner)
    Dxy[indx, n_right[indx]] = 1 / 2
    Dxy[indx, n_below_left[indx]] = 1 / 2
    Dxy[indx, n_left[indx]] = -1 / 2
    Dxy[indx, n_below_right[indx]] = -1 / 2

    indx, = np.where(at_bottom)
    Dy[indx, n_above[indx]] = 1
    Dy[indx, indx] = -1
    Dyy[indx, n_above[indx]] = 1
    Dyy[indx, indx] = -1

    indx, = np.where(at_bottom_not_at_corner)
    Dxy[indx, n_above_right[indx]] = 1 / 2
    Dxy[indx, n_left[indx]] = 1 / 2
    Dxy[indx, n_above_left[indx]] = -1 / 2
    Dxy[indx, n_right[indx]] = -1 / 2

    indx, = np.where(at_top_left)
    Dxy[indx, n_below[indx]] = 1
    Dxy[indx, n_right[indx]] = 1
    Dxy[indx, indx] = -1
    Dxy[indx, n_below_right[indx]] = -1

    indx, = np.where(at_top_right)
    Dxy[indx, indx] = 1
    Dxy[indx, n_below_left[indx]] = 1
    Dxy[indx, n_left[indx]] = -1
    Dxy[indx, n_below[indx]] = -1

    indx, = np.where(at_bottom_left)
    Dxy[indx, n_above_right[indx]] = 1
    Dxy[indx, indx] = 1
    Dxy[indx, n_above[indx]] = -1
    Dxy[indx, n_right[indx]] = -1

    indx, = np.where(at_bottom_right)
    Dxy[indx, n_above[indx]] = 1
    Dxy[indx, n_left[indx]] = 1
    Dxy[indx, indx] = -1
    Dxy[indx, n_above_left[indx]] = -1

    indx, = np.where((n_above_left > -1) & (n_below_right > -1))
    Dsm[indx, indx] = -2
    Dsm[indx, n_above_left[indx]] = 1
    Dsm[indx, n_below_right[indx]] = 1
    indx, = np.where((n_above_left == -1) & (n_below_right > -1))
    Dsm[indx, indx] = -1
    Dsm[indx, n_below_right[indx]] = 1
    indx, = np.where((n_above_left > -1) & (n_below_right == -1))
    Dsm[indx, indx] = -1
    Dsm[indx, n_above_left[indx]] = 1

    indx, = np.where((n_above_right > -1) & (n_below_left > -1))
    Dsp[indx, indx] = -2
    Dsp[indx, n_above_right[indx]] = 1
    Dsp[indx, n_below_left[indx]] = 1
    indx, = np.where((n_above_right == -1) & (n_below_left > -1))
    Dsp[indx, indx] = -1
    Dsp[indx, n_below_left[indx]] = 1
    indx, = np.where((n_above_right > -1) & (n_below_left == -1))
    Dsp[indx, indx] = -1
    Dsp[indx, n_above_right[indx]] = 1

    Dx = Dx / dx
    Dy = Dy / dy
    Dxx = Dxx / dx**2
    Dyy = Dyy / dy**2
    Dxy = Dxy / (dx * dy)
    Dsp = Dsp / (dx**2 + dy**2)
    Dsm = Dsm / (dx**2 + dy**2)

    # Package all operators up into a dictionary
    operators = dict(Dx=Dx, Dy=Dy, Dxx=Dxx, Dyy=Dyy, Dxy=Dxy, Dsp=Dsp, Dsm=Dsm)
    return operators

[jacklovell]

The interface you propose is indeed cleaner for the specific case of a regular rectangular grid, but we need to preserve backwards compatibility and also support non-regular grids such as the ToroidalVoxelGrid - which doesn't have to be rectangular - albeit with rectangular cross sections for individual voxels.

I'll look at vectorising the function but in order to support a ragged grid it might not be possible: for a ragged grid it was significantly simpler to just write out the loop. And it's only O(ncells) not O(ncells^2) so the performance hit isn't that bad for an operation which only needs to be done once.

At the very least, the derivative operators can be made sparse which will help with memory usage for large ncells.

[skuba31]

Would it make sense to have a simpler easily readable function for regular grids as Vlad suggests and another one with irregular voxels in mind?

[jacklovell]

I don't know that 2 functions will be any easier to read than 1 function.

[skuba31]

It is not only about having two function, I should have explained better. In his comments, Vlad is suggesting that a function assuming regular grid would make the demo simpler and therefore the use on regular grids in general. In my experience, regular grids are by far the most common choice these days.

So the question is, whether having a function with limited scope tailored for the most common use case outweighs the longer source code with two functions, where the second would be more general.

Or a wrapper function that would handle the mapping from regular grid into general could be created to simplify the most common use case.

Alternatively I have a rather longer term idea for consideration. Given that inversions are not the primary aim of Cherab, inversion related parts (all or the more advanced ones) could be relocated to a dedicated package.

[jacklovell]

I've remembered now why I implemented it this way.

We've historically used "masked" voxel grids for doing bolometry inversions in the machine-specific subpackages: examples for JET, MAST-U and AUG. For a voxel grid there is a measurable performance advantage to omitting voxels which aren't within the first wall of the machine, which leads to "ragged" arrays of voxels which roughly follow the contours of the first wall/limiting surface.

Since this formulism is already in widespread use I think it's important to support it for both new and existing users. And since the Ray Transfer Matrix formulism specialises this more general formulism to a regular grid it's easy to support it too (as the demo shows). I don't think it would be as straightforward to do the reverse: write a function expecting a regular grid and have end users convert ragged grids into the correct form.

I could I suppose also show the usage with a ToroidalVoxelGrid, which would even more closely match what the existing packages do. But that case is sufficiently trivial (all these machine packages give exactly the inputs required for this function) that I think users would be able to work it our for themselves.

[jacklovell]

The other big advantage of the looping method over the vectorized method is memory efficiency when sparse=True is passed through, which is important for computing the operators for large grids on memory-limited machines. This alone I think makes it worthwhile keeping the existing looping function. Hopefully with the added comments it'll be a bit easier to follow now, which I think was the essence of Vlad's original comment.

[skuba31]

I see the backward compatibility and existing codes as the most important reason for keeping the approach you proposed. The memory efficiency is a nice bonus.

I think that having an explicit example for trivial code is worth it should the code be used frequently. Do you see the use with ToroidalVoxelGrid as a sort of "default" approach? If so, I would be in favor of including that example. It could be a comment in existing demo if that is possible, not necessarily entirely new one.

[vsnever]

All this code can be skipped if the proposed implementation of the generate_derivative_operators function is accepted.

derivative_operators = admt.generate_derivative_operators(
    cell_vertices_r, cell_vertices_z, ray_transfer_grid.voxel_map.squeeze()
)

[jacklovell]

TBF quite a bit of this code could have been skipped if I hadn't missed the obvious fact that the 2D-to-1D map is simply the inverse of the 1D-to-2D map:

grid_index_2d_to_1d_map = {rz: k for (k, rz) in grid_index_1d_to_2d_map.items()}

And vice versa.

So actually, only one of the 2D-to-1D or 1D-to-2D maps is required and the other can always be computed trivially, without loss of generality. I'll make generate_derivative_operators optionally accept only one of these two arguments, and automatically compute the other if the user don't already provide both. That'll simplify the demo too while maintaining backwards compatibility.

[jacklovell]

8e2a399 allows passing only 1 voxel mapping to the derivative operators generator function, which significantly simplifies this for the end user even in the Ray Transfer case.

[skuba31]

I think it is fine to leave these few extra lines here now. I suppose that this code could be simplified by adding properties or methods to grid classes in the future.

[jacklovell]

I've implemented the changes from both reviews now: thanks to @vsnever and @skuba31. I am going to hold firm on the formulism in generate_derivative_operators though, as I think the benefits of sparseness and explicitness outweigh any potential performance gains or loss of generality of the vectorized version Vlad proposed before the iterative improvements from the reviews. As a compromise the function is now better commented and the interface is a bit easier to use for callers as reflected in the demo.

Once this is approved I'll do a squash merge as the commit history is longer than necessary to keep going forwards.

[skuba31]

I see the advantages of grouping a property for all cameras, but it is a bit confusing for me to navigate in. I found it more natural to group properties by camera, so that it is simpler to see where the camera is, without need to search through multiple sections of the code.

[skuba31]

I would suggest renaming nx to nr and ny to nz. Using ny for number of voxels in vertical direction could cause some confusion as in tokamaks the y axis typically lies in midplane. It also results in somewhat awkward naming in pairs cell_r, cell_dx defined by nx and cell_z, cell_dz defined by ny.

[skuba31]

Nice work @jacklovell . I think that the documentation and code readability is significantly improved. The sparse support and ADMT demo are very useful additions. I have added a few minor comments for consideration before final approval.