add edge sorting/flipping to reproduce acorn workflow - #103
Conversation
| for (std::size_t n = 0; n < numNodes; ++n) { | ||
| const float r = inputValues[n * fullNumFeatures + static_cast<std::size_t>(indices[eR])]; | ||
| const float z = inputValues[n * fullNumFeatures + static_cast<std::size_t>(indices[eZ])]; | ||
| distances[n] = r * r + z * z; |
There was a problem hiding this comment.
This can to go out of sync with the edge walking / track building above quite easily. Is it possible to lift at least parts of this to some common functionality?
There was a problem hiding this comment.
Good point !
Pulled it in a header-only GNNTrackFinding/include/EdgeDirection.h, which now owns the whole thing.
CCAndWalkTrackBuilding uses it directly: its three lambdas collapse to one nodeDistancesSq() call plus a thin pointsOutward() wrapper (which also stops the distance being recomputed on every comparison inside the sort). Its Config now takes radiusFeatureIndices in the same shape as OnnxMetricLearning::Config, so
GNNTrackFinder hands the identical vector to both stages instead of unpacking it.
The flip in orderEdgesByRadius() stays in torch ops, since it has to run on whatever device the edge building used. However, it no longer computes a metric of its own.
Added test/unittests/EdgeDirectionTests.cpp to pin the shared behaviour (column/stride handling, padding rows ignored, bad indices, the tie-break, and that it is a strict total order). It needs no torch or ONNX, so it hangs off the existing unittest_gnntrackfinding target.
Today we have ~2 pipelines and they both come from acorn, so they both expect the edges to be oriented like that. However in principle acorn is not the only way to set up a GNN to do tracking, so I'd have a weak preference to keep this configurable. Yes #94 is about shipping tensors back and forth from CPU to GPU, this is a separate feature. |
| /// tensor. The edges are directed by the distance from the interaction | ||
| /// point these two form, see the class documentation. Same convention (and | ||
| /// same value) as OnnxMetricLearning::Config::radiusFeatureIndices. | ||
| std::vector<int> radiusFeatureIndices{}; |
There was a problem hiding this comment.
This could already be a std::vector<size_t>? That could remove some of the static_casts that are currently necessary, I think.
There was a problem hiding this comment.
Done.
std::vector<std::size_t> throughout converted once where the names are resolved. The casts and the negative-index check are also gone.
| .radiusFeatureIndices = | ||
| m_sortEdges.value() ? m_radiusFeatureIndices : std::vector<int>{}, |
There was a problem hiding this comment.
Just to make sure: We cannot simply use m_radiusFeatureIndices here, because they are only initialized to the empty vector in case ccAndWalk is also true above?
There was a problem hiding this comment.
Almost the other way round: the indices are filled when SortEdges or cc-and-walk is on. With SortEdges=False + cc-and-walk they're non-empty, and since an empty vector is what switches the ordering off in OnnxMetricLearning, passing them unconditionally would turn it on.
I reworded the comment to say so.
| // Point every edge away from the interaction point. This runs before the | ||
| // guard below because it can only ever remove edges, and before the edge | ||
| // features because those are signed differences along the edge. | ||
| edgeList = orderEdgesByRadius(inputValues, numNodes, fullNumFeatures, std::move(edgeList)); |
There was a problem hiding this comment.
| edgeList = orderEdgesByRadius(inputValues, numNodes, fullNumFeatures, std::move(edgeList)); | |
| edgeList = orderEdgesByDistance(inputValues, numNodes, fullNumFeatures, std::move(edgeList)); |
? Depending a bit on what people usually denote as radius. The 2D sqrt(x**2 + y**2), or the 3D version which is used here.
There was a problem hiding this comment.
Agreed, it's the 3D distance. Renamed it and the rest of the family (distanceFeatureIndices, kNumDistanceFeatures, …) for consistency.
| /// | ||
| /// @throws std::invalid_argument if @p radiusFeatureIndices does not name two | ||
| /// columns of the buffer | ||
| [[nodiscard]] inline std::vector<float> nodeDistancesSq(const float* nodeFeatures, std::size_t numNodes, |
There was a problem hiding this comment.
I guess this is a more an interface to torch issue, but the const float*, size_t pair is in principle exactly what std::span is for. Not sure if it's easily fixed here though.
There was a problem hiding this comment.
Turned out to be easy: the torch side only concerns the output buffer. It takes a std::span now, which also lets it check the buffer actually holds numNodes rows.
2 participants
BEGINRELEASENOTES
ENDRELEASENOTES