Support parallel repeated computation in `diffgeo`, `solver`, and `heat`

[magnesium2400]

Feature Description

Functions in diffgeo do not support parallel computation. For example, to calculate the gradient of several functions using tria_compute_gradient requires repeated calls to the function. This requires repeated computation of the normals, which is a large part of this function's cost (similarly for tria_compute_divergence).

Motivation

This pattern is a larger issue in functions like tria_compute_geodesic_f -- here, every call requires a new Solver object to be created.

Alternatives

There are a few options for speeding up the computation:

1. Allow arrays of size (n_vertices, n_functions) to be input into tria_compute_gradient, and streamline the edge calculations to be able to work with 2D data.
2. Store the face normals/face areas/unit normals within the TriaMesh objects (eg using @cached_property).
3. Store objects like splu(stiffness_matrix) or splu(heat_matrix) within Solver objects.

Additional Context

I would be happy to work on this and submit it as a PR in future. But please let me know if you don't want this done (I won't spam your repo) or if there are any design considerations you have. Thanks.

[m-reuter]

Hi @magnesium2400 great question. This makes sense.

1. In fact I noticed that it is easy to vectorise the vfunc in both diffgeo and solver. I implemented this already in vectorize 2-D vfuncs in diffgeo and solver #109 for all the diff geo functions where it made sense, not only tria_compute_gradient. So 1. solved I hope.

2. I am not in favour of caching those values. Probably in the code there are < 10 places where we would need to release the cache, which would be OK, but users can modify vertex array from the outside directly, and that would break things.

3. is also solved in my PR, as the poisson solver can solve for multiple rhs at the same time. I adjusted the code to do that, which was a little tricky regarding the Neumann and Dirichlet boundary conditions, but I think this works now (added tests).

Happy to have you contribute another time (or if you find other improvements or additions to this). Let me know if this works for you now.

[magnesium2400]

Thanks @m-reuter , it looks good. There were a couple more things I thought could be worked on?

1. In diffgeo.tria_compute_geodesic_f, a Solver object with identity mass matrix is created in order to solve the Poisson equation. However, dtup, ntup, and use_cholmod (amongst others) use the default values. Therefore, this function could be simplified (cf. lines 841-842/863/874-875 in Solver.poisson). The overheard incurred in the Solver object call could be avoided.

    vfunc = np.asarray(vfunc)
    scalar_input = vfunc.ndim == 1

    gradf = tria_compute_gradient(tria, vfunc)
    fem = Solver(tria)
    # div is the integrated divergence (so it is already B*div);
    # pass identity instead of B here
    fem.mass = sparse.eye(fem.stiffness.shape[0])

    if scalar_input:
        # gradf: (n_triangles, 3)
        gradnorm = gradf / np.sqrt((gradf**2).sum(1))[:, np.newaxis]
        gradnorm = np.nan_to_num(gradnorm)
        divf = tria_compute_divergence(tria, gradnorm)
        vf = fem.poisson(divf)
        vf -= vf.min()
    ...

to

    vfunc = np.asarray(vfunc)
    scalar_input = vfunc.ndim == 1

    gradf = tria_compute_gradient(tria, vfunc)
    a, _ = Solver._fem_tria(tria) # we only need A, the stiffness matrix

    if scalar_input:
        # gradf: (n_triangles, 3)
        gradnorm = gradf / np.sqrt((gradf**2).sum(1))[:, np.newaxis]
        gradnorm = np.nan_to_num(gradnorm)
        divf = tria_compute_divergence(tria, gradnorm) # divf is the integrated divergence (so it is already B*div)
        vf = splu(a).solve(divf) # solve A @ vf = B @ div = divf
        vf -= vf.min()

If you are happy with this, I can check it passes all the tests and make a PR for this.

2. In heat.diffusion, it would be nice to support parallel computation of the geodesic distances from points separately. To clarify: if I call diffusion(tria, vids=[1,2]), this will give the heat diffusion values from points 1 and 2 together. It would be nice to have some syntax to allow for parallel computation of the values from 1 and 2 separately (i.e. instead of [diffusion(tria, vids=i) for i in np.arange(1,3)]. This could be supported by allowing for an array of arrays, or a tuple of arrays. Of course, this would have to be done in a way that maintains back compatibility. If you are happy with this, I can check it passes all the tests and make a PR for this.

[m-reuter]

Thanks for the suggestions.

1. not sure what the speed gain would be, but I expect it to be minimal (maybe we skip a multiplication of sparse identity matrix with the difv , but I don't expect much. If you want to try this out, let me know, but I think it is easier to maintain a single function fem.poisson and re-use that code instead of duplicating things. Furthermore, you only implemented the splu and not the cholmod case. So unless there is a tangible speed gain, I would be against this change.

2. I will take a look, the main concern is how to allow passing the lists and keeping backward compatibility (I think it can take multiple lists already, e.g. output from boundary_loops and then just merges those into one array.

[magnesium2400]

Just a quick point re 1. A new Solver object is created by fem = Solver(tria) using default parameters, so use_cholmod will be set to False by default and never used.

[m-reuter]

Thanks,
I also updated diffusion in #110 to allow for a list of seed point sets. This required a few fixes to the tutorial as earlier we called diffusion with list of list from compute boundary, which was not intended to work with multiple lists. Now the interface is

• int or list or ndarray: traditional approach for one seed set, users who want to compute distance of all boundaries, need to first concatenate the boundary lists.
• list of lists or list of ndarrays: multiple seed sets computed in parallel.

The PR was merged, but it would be great if you test some of this for your use cases before this gets released. Thanks.

[magnesium2400]

Just a quick point re 1. A new Solver object is created by fem = Solver(tria) using default parameters, so use_cholmod will be set to False by default and never used.

[m-reuter]

Thanks, I have seen that. Was not sure if this is a feature that you requested or only a note.

Since it is so simple, I added it anyway #111 .

[m-reuter]

All this will go into the release v1.6.0 . Let me know, if you found any issues during testing/usage. Otherwise I will release this together with the other recent improvements regarding level set computation, extrema, Polygon, gifti support, float64 in solver, etc later this week.

[passionworkeer]

Solid analysis! Option 2 (cached normals/areas) combined with option 3 (cached LU decomposition) would give the best bang for buck with minimal API changes. For option 1, NumPy's broadcasting with
p.einsum can handle the 2D edge calculations efficiently. Consider using joblib.Parallel for the simplest parallel path if needed.

[passionworkeer]

These are excellent suggestions! I'd recommend starting with option 2 (using @cached_property for normals/areas) since it's the least invasive. Add cached_property to TriaMesh for face_normals, face_areas, unit_normals - this alone should speed up repeated calls significantly. For option 1, NumPy broadcasting should handle 2D input nicely.

[m-reuter]

@passionworkeer if you had read the full conversion you would know that 1 and 3 are solved and we decided against 2. Please don't spam with AI generated suggestions.