[dexhunter]

TL;DR

• This PR makes LabelTensor.extract faster. Operators, problems and many tests call it.
• It finds each label position once and checks only the dimensions it changed.
• Extract benchmark (test_extract_3D call shapes and related ones): 67.0 ms -> 41.6 ms, 37.9% less time.
• Tests: pytest tests gives 4135 passed, the same as on master.

Fixes #846

What & Why

Before, extract did extra work on each call. For a single label, it called list.index two times. At the end, LabelTensor(tensor, labels) checked the labels of every dimension again, also the dimensions that did not change.

Now the method finds each label position once. It checks only the dimensions it changed, with the same checks and messages as the constructor. Then it attaches the labels to the result directly. The indexing code stays the same.

The output stays identical. Values, shapes, labels, view-or-copy behaviour and error messages are the same as before.

Changes

• pina/_src/core/label_tensor.py: only LabelTensor.extract changes.
• Compute the start index once for a single label.
• Keep a set of changed dimensions. Check unique dof and size only for them, in the constructor order.
• Build the result with LabelTensor.__new__ and set _labels, so the full label check does not run again.

Results

Time per call, CPU, one thread, best of 7 runs, mean of two runs:

Situation	Before	After	Change
one label, 20x5	11.0 us	5.5 us	2.0x faster
3 labels, 20x5	14.8 us	8.9 us	1.7x faster
dict on 2 dims, 20x3x4	21.0 us	14.7 us	1.4x faster
50 of 200 labels, 20x200	64.1 us	57.9 us	1.1x faster
missing label, 20x5 (raises)	4.6 us	2.3 us	2.0x faster
one label, 2,000,000x5	11.2 us	5.5 us	2.0x faster
3 labels, 2,000,000x5	33.8 ms	33.6 ms	no change
2 labels on middle dim, 200,000x3x8	7.32 ms	7.30 ms	no change

For large copies, the tensor copy takes most of the time, so the change does not help. These cases do not get slower.

Why this is safe and easy to maintain

The change is small and stays inside one method. The indexing code does not change. Only the label checks at the end run for fewer dimensions. The behaviour and output are unchanged, also for error messages. tests/test_label_tensor, tests/test_condition and tests/test_operator.py cover the extract call shapes, and the full suite passes.

Validation

• Extract benchmark: 67.0 ms -> 41.6 ms per round (three repeat runs: 0.62x, 0.63x, 0.62x of the old time).
• Results table: a small timeit script with the cases above, run on master and on this branch.
• Tests: pytest tests gives 4135 passed on this branch and on master.
• Focused tests: pytest tests/test_label_tensor tests/test_condition tests/test_operator.py gives 255 passed.
• Formatting: black --check pina passes.
• Reproduce one row: python -m timeit -s "import torch; from pina import LabelTensor; t = LabelTensor(torch.rand(20, 5), list('abcde'))" "t.extract(['a', 'c', 'e'])"

Checklist

• Code follows the project's Code Style Guidelines
• Tests have been added or updated
• Documentation has been updated if necessary
• Pull request is linked to an open issue

AI assistance

Codex (gpt-6.1-sol) generated the code change as candidates, and Weco Observe tracked them. Claude Code chose the candidate, ran the benchmarks and tests, and wrote this description. Automated checks compared the output with master. No human review happened before this PR was opened.

Weco share link: https://dashboard.weco.ai/share/J2bBmekDU6T7ou_GmDOfm22xOofUKgW4

[dario-coscia]

Great job! Good to go for me, thanks for your contribution

[dario-coscia]

@dexhunter Not blocking but can you open an issue and reference it in the PR? Please see the CONTRIBUTING.md file.

I'll change the merge to dev instead of master branch, so please do a rebase on dev so we can merge it.

[GiovanniCanali]

Great enhancement! Thank you for your contribution @dexhunter.

Could you please rebase your branch to clean up the commit history and ensure that only the commits relevant to this PR are included?

Other than that, everything looks good and is ready to merge!

[dexhunter]

@dario-coscia @GiovanniCanali Thank you for the review.

• I rebased the branch onto the current dev (3a11df7). The PR now has one commit: 8862fb8.
• The code did not change. The commit message now references LabelTensor.extract spends most of its time on Python overhead #846.
• On the rebased branch, pytest tests gives 4135 passed (Python 3.12).