Skip to content

Add DirectionReconstructionWithLogKappa - #929

Open
sevmag wants to merge 1 commit into
graphnet-team:mainfrom
sevmag:feature/vmf-logkappa-head
Open

sevmag wants to merge 1 commit into
graphnet-team:mainfrom
sevmag:feature/vmf-logkappa-head

Conversation

@sevmag

@sevmag sevmag commented Oct 4, 2026

Copy link
Copy Markdown
Collaborator

Closes #928.

Adds DirectionReconstructionWithLogKappa: direction from the normalised first three outputs, kappa = kappa_min + exp(x[:, 3]) (kappa_min default 1). Same prediction layout as DirectionReconstructionWithKappa, so it works with VonMisesFisher3DLoss unchanged. The docstring of DirectionReconstructionWithKappa now notes its kappa ceiling and points to the new task.

Tests added in tests/models/test_task.py.

🤖 Generated with Claude Code

Predicts kappa as kappa_min + exp(x[:, 3]) instead of the output norm,
removing the kappa ceiling of DirectionReconstructionWithKappa.

Refs graphnet-team#928

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@christianlocatelli
christianlocatelli self-requested a review October 12, 2026 08:24

@christianlocatelli christianlocatelli left a comment •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR @sevmag. Everything looks good and reasonable overall. I discussed the test functions with Aske though, and we think most of them aren't really necessary and are unlikely to catch future bugs.
test_log_kappa_direction_output, test_log_kappa_no_ceiling_and_no_overflow, and test_log_kappa_floor_must_be_positive mostly restate the _forward function and are clear from the code, would it be possible to remove them?
test_log_kappa_loss_and_gradients could arguably still be kept, because it checks the new head together with VonMisesFisher3DLoss, and a change in either of them could break the other. Let us know what you think about it.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Direction task: predict log κ as a separate output instead of κ = ‖x‖

2 participants