feat: implement vector mask bit shifting to create Region Mask #345

jjjermiah · 2025-04-30T22:00:10Z

This is a proof of concept to implement the NNUNet Region Mask implementation

I realized that the VectorMask implementation inherently supports the RegionMask approach by being one-hot encoded vectors, where the total possible regions is 2^n where n is number of masks stored.

if we consider a vector voxel in a VectorMask:

[gtv, brain, heart, lungs]

and each individual mask has either 0s or 1s, i.e:

[0,1,1,0]

the values together, represent a unique number in binary, that can define what masks are part of that region:

0110 = 6

so all regions where only brain and heart overlap can be assigned the label 6

The code in this pull request creates some dummy shapes to show this implementation, and the result looks like:

you can see in this image, the cursor is hovering where Clover and Square overlap, and it has a value of 10

References

See this conversation

Summary by CodeRabbit

New Features
- Added support for creating and manipulating multi-component vector masks with distinct overlapping shapes.
- Introduced functionality to collapse vector masks into single-channel label images, enabling easier interpretation of overlapping regions.
- Provided a demonstration for generating, saving, and visualizing vector masks and their collapsed representations.

coderabbitai · 2025-04-30T22:00:16Z

📝 Walkthrough

Walkthrough

A new module, vector_shift.py, has been added under src/imgtools/coretypes. This module introduces functions and data structures for creating, manipulating, and collapsing vector masks represented as multi-component SimpleITK images. The main functionalities include generating a demo vector mask with distinct geometric patterns, collapsing a vector mask into a single-channel label image using bitmask encoding, and providing a lookup table for label interpretation. The module also contains a demonstration block that showcases these features and outputs results using the rich library.

Changes

File(s)	Change Summary
src/imgtools/coretypes/vector_shift.py	New module with functions for creating demo vector masks, collapsing vector masks to label images, and a demo block. Introduces `create_demo_vector_mask`, `collapse_vector_mask`, and `PartitionResult` named tuple.

✨ Finishing Touches

📝 Generate Docstrings

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

Review comments: Directly reply to a review comment made by CodeRabbit. Example:
- I pushed a fix in commit <commit_id>, please review it.
- Generate unit testing code for this file.
- Open a follow-up GitHub issue for this discussion.
Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
- @coderabbitai generate unit testing code for this file.
- @coderabbitai modularize this function.
PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
- @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
- @coderabbitai read src/utils.ts and generate unit testing code.
- @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
- @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

@coderabbitai pause to pause the reviews on a PR.
@coderabbitai resume to resume the paused reviews.
@coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
@coderabbitai full review to do a full review from scratch and review all the files again.
@coderabbitai summary to regenerate the summary of the PR.
@coderabbitai generate docstrings to generate docstrings for this PR.
@coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
@coderabbitai resolve resolve all the CodeRabbit review comments.
@coderabbitai configuration to show the current CodeRabbit configuration for the repository.
@coderabbitai help to get help.

Other keywords and placeholders

Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
Add @coderabbitai anywhere in the PR title to generate the title automatically.

Documentation and Community

Visit our Documentation for detailed information on how to use CodeRabbit.
Join our Discord Community to get help, request features, and share feedback.
Follow us on X/Twitter for updates and announcements.

codecov · 2025-04-30T22:01:42Z

Codecov Report

Attention: Patch coverage is 0% with 85 lines in your changes missing coverage. Please review.

Project coverage is 57.39%. Comparing base (4e64438) to head (fd5ecbb).

Files with missing lines	Patch %	Lines
src/imgtools/coretypes/vector_shift.py	0.00%	85 Missing ⚠️

Additional details and impacted files

@@               Coverage Diff               @@
##           development     #345      +/-   ##
===============================================
- Coverage        58.66%   57.39%   -1.28%     
===============================================
  Files               61       62       +1     
  Lines             3832     3917      +85     
===============================================
  Hits              2248     2248              
- Misses            1584     1669      +85

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

… function

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (6)

src/imgtools/coretypes/vector_shift.py (6)

11-26: Function signature and docstring need improvements.

The docstring needs type information in the Parameters section and better alignment with PEP 257.

 def create_demo_vector_mask(
     size: tuple[int, int] = (128, 128),
 ) -> tuple[sitk.Image, dict[str, int]]:
     """
     Create a 4-component VectorUInt8 image with overlapping synthetic shapes.
 
+    Parameters
+    ----------
+    size : tuple[int, int], optional
+        Size of the image in (width, height) format, by default (128, 128)
+
     Returns
     -------
     sitk.Image
         - Component 0: Circle
         - Component 1: Square
         - Component 2: stripe
         - Component 3: clover pattern
     dict[str, int]
         Mapping from component name to index.
     """

🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 11-11: src/imgtools/coretypes/vector_shift.py#L11
Added line #L11 was not covered by tests

37-43: Circle generation could be more efficient.

The nested loops approach for creating a circle is readable but inefficient. Consider using SimpleITK's built-in functions or NumPy for better performance.

-    # Draw circle: set pixels inside radius
-    center = [s // 2 for s in img_size]
-    radius = min(img_size) // 3
-    for y in range(img_size[1]):
-        for x in range(img_size[0]):
-            if (x - center[0]) ** 2 + (y - center[1]) ** 2 < radius**2:
-                circle[x, y] = 1
+    # Draw circle: set pixels inside radius
+    center = [s // 2 for s in img_size]
+    radius = min(img_size) // 3
+    
+    # Create indices for all pixels
+    import numpy as np
+    indices_x, indices_y = np.meshgrid(
+        np.arange(img_size[0]), np.arange(img_size[1]), indexing='xy'
+    )
+    
+    # Calculate distances from center
+    distances = np.sqrt(
+        (indices_x - center[0])**2 + (indices_y - center[1])**2
+    )
+    
+    # Create circle mask and convert to SimpleITK image
+    circle_array = (distances < radius).astype(np.uint8)
+    circle = sitk.GetImageFromArray(circle_array.T)

🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 37-42: src/imgtools/coretypes/vector_shift.py#L37-L42
Added lines #L37 - L42 were not covered by tests

77-84: Simplify angle calculation and improve readability for clover pattern.

The angle calculation contains a confusing expression ((dx / distance) * 0) + ((dy / distance) * 1) that simplifies to just dy / distance. Also, consider using math.pi for mathematical constants.

-            # Calculate angle in radians
-            angle = abs(((dx / distance) * 0) + ((dy / distance) * 1))
-            angle = angle % (2 * 3.14159 / num_points)
+            # Calculate angle in radians
+            import math
+            angle = abs(dy / distance)  # Simplified from previous calculation
+            angle = angle % (2 * math.pi / num_points)

             # Threshold changes based on angle to create clover points
             threshold = clover_radius_inner + (
                 clover_radius_outer - clover_radius_inner
-            ) * abs(angle - 3.14159 / num_points) / (3.14159 / num_points)
+            ) * abs(angle - math.pi / num_points) / (math.pi / num_points)

🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 77-78: src/imgtools/coretypes/vector_shift.py#L77-L78
Added lines #L77 - L78 were not covered by tests

[warning] 81-81: src/imgtools/coretypes/vector_shift.py#L81
Added line #L81 was not covered by tests

105-123: Docstring doesn't match implementation.

The docstring mentions vector_mask parameter as sitk.Image but the type annotation correctly shows VectorMask. Also, the return value description doesn't match the actual structure.

 def collapse_vector_mask(
     vector_mask: VectorMask,
 ) -> PartitionResult:
     """
     Encodes a VectorUInt8 image (with binary 0/1 components) into a single-channel
     image where each voxel value is a unique integer representing the bitmask
     of active components. Names are used in the lookup table.
 
     Parameters
     ----------
-    vector_mask : sitk.Image
+    vector_mask : VectorMask
         A VectorUInt8 image where each component is 0 or 1.
 
     Returns
     -------
     PartitionResult
         - label_image: scalar UInt16 image of unique bitmask-encoded labels
-        - lookup_table: dict[label_value] = joined component name string
+        - lookup_table: dict mapping label values to lists of ROI image IDs
     """

🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 105-105: src/imgtools/coretypes/vector_shift.py#L105
Added line #L105 was not covered by tests

153-169: Example usage should be better documented.

The main block demonstrates usage but lacks comments explaining the purpose of each step. It should also handle potential file writing errors.

 if __name__ == "__main__":
     output_dir = Path(__file__).parent
 
-    # --- Example usage
+    # --- Example: Create a demo vector mask with synthetic shapes
     img, component_names = create_demo_vector_mask()
 
+    # Create a VectorMask by wrapping the SimpleITK image with proper ROI mappings
     vm = VectorMask(
         image=img,
         roi_mapping={
             i: ROIMaskMapping(
                 roi_key=name, roi_names=name, image_id=name.upper()
             )
             for name, i in component_names.items()
         },
         metadata={},
     )
 
-    sitk.WriteImage(vm, output_dir / "vector_mask.nii.gz")
+    # Save the vector mask for visualization
+    try:
+        output_path = output_dir / "vector_mask.nii.gz"
+        sitk.WriteImage(vm, str(output_path))
+        print(f"Saved vector mask to {output_path}")
+    except Exception as e:
+        print(f"Failed to save vector mask: {e}")

170-173: Add error handling for image writing.

File writing operations should have error handling.

-    sitk.WriteImage(result.label_image, output_dir / "collapsed_mask.nii.gz")
+    # Save the collapsed mask for visualization
+    try:
+        output_path = output_dir / "collapsed_mask.nii.gz"
+        sitk.WriteImage(result.label_image, str(output_path))
+        print(f"Saved collapsed mask to {output_path}")
+    except Exception as e:
+        print(f"Failed to save collapsed mask: {e}")

📜 Review details

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 4e64438 and 6c3ecdb.

📒 Files selected for processing (1)

src/imgtools/coretypes/vector_shift.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (1)

`src/**/*.py`: Review the Python code for compliance with PEP 8 and PEP 257 (docstring conventions). Ensure the following: - Variables and functions follow meaningful naming conven...

src/**/*.py: Review the Python code for compliance with PEP 8 and PEP 257 (docstring conventions). Ensure the following: - Variables and functions follow meaningful naming conventions. - Docstrings are present, accurate, and align with the implementation. - Code is efficient and avoids redundancy while adhering to DRY principles. - Consider suggestions to enhance readability and maintainability. - Highlight any potential performance issues, edge cases, or logical errors. - Ensure all imported libraries are used and necessary.

src/imgtools/coretypes/vector_shift.py

🧬 Code Graph Analysis (1)

src/imgtools/coretypes/vector_shift.py (1)

src/imgtools/coretypes/base_masks.py (1)

VectorMask (27-382)

🪛 GitHub Check: codecov/patch

src/imgtools/coretypes/vector_shift.py

[warning] 1-2: src/imgtools/coretypes/vector_shift.py#L1-L2
Added lines #L1 - L2 were not covered by tests

[warning] 4-6: src/imgtools/coretypes/vector_shift.py#L4-L6
Added lines #L4 - L6 were not covered by tests

[warning] 8-8: src/imgtools/coretypes/vector_shift.py#L8
Added line #L8 was not covered by tests

[warning] 11-11: src/imgtools/coretypes/vector_shift.py#L11
Added line #L11 was not covered by tests

[warning] 27-28: src/imgtools/coretypes/vector_shift.py#L27-L28
Added lines #L27 - L28 were not covered by tests

[warning] 31-34: src/imgtools/coretypes/vector_shift.py#L31-L34
Added lines #L31 - L34 were not covered by tests

[warning] 37-42: src/imgtools/coretypes/vector_shift.py#L37-L42
Added lines #L37 - L42 were not covered by tests

[warning] 45-46: src/imgtools/coretypes/vector_shift.py#L45-L46
Added lines #L45 - L46 were not covered by tests

[warning] 49-52: src/imgtools/coretypes/vector_shift.py#L49-L52
Added lines #L49 - L52 were not covered by tests

[warning] 55-58: src/imgtools/coretypes/vector_shift.py#L55-L58
Added lines #L55 - L58 were not covered by tests

[warning] 61-64: src/imgtools/coretypes/vector_shift.py#L61-L64
Added lines #L61 - L64 were not covered by tests

[warning] 66-70: src/imgtools/coretypes/vector_shift.py#L66-L70
Added lines #L66 - L70 were not covered by tests

[warning] 72-74: src/imgtools/coretypes/vector_shift.py#L72-L74
Added lines #L72 - L74 were not covered by tests

[warning] 77-78: src/imgtools/coretypes/vector_shift.py#L77-L78
Added lines #L77 - L78 were not covered by tests

[warning] 81-81: src/imgtools/coretypes/vector_shift.py#L81
Added line #L81 was not covered by tests

[warning] 85-86: src/imgtools/coretypes/vector_shift.py#L85-L86
Added lines #L85 - L86 were not covered by tests

[warning] 89-89: src/imgtools/coretypes/vector_shift.py#L89
Added line #L89 was not covered by tests

[warning] 91-91: src/imgtools/coretypes/vector_shift.py#L91
Added line #L91 was not covered by tests

[warning] 97-97: src/imgtools/coretypes/vector_shift.py#L97
Added line #L97 was not covered by tests

[warning] 100-102: src/imgtools/coretypes/vector_shift.py#L100-L102
Added lines #L100 - L102 were not covered by tests

[warning] 105-105: src/imgtools/coretypes/vector_shift.py#L105
Added line #L105 was not covered by tests

[warning] 124-126: src/imgtools/coretypes/vector_shift.py#L124-L126
Added lines #L124 - L126 were not covered by tests

[warning] 129-129: src/imgtools/coretypes/vector_shift.py#L129
Added line #L129 was not covered by tests

[warning] 131-132: src/imgtools/coretypes/vector_shift.py#L131-L132
Added lines #L131 - L132 were not covered by tests

[warning] 134-135: src/imgtools/coretypes/vector_shift.py#L134-L135
Added lines #L134 - L135 were not covered by tests

[warning] 138-139: src/imgtools/coretypes/vector_shift.py#L138-L139
Added lines #L138 - L139 were not covered by tests

[warning] 141-142: src/imgtools/coretypes/vector_shift.py#L141-L142
Added lines #L141 - L142 were not covered by tests

[warning] 144-145: src/imgtools/coretypes/vector_shift.py#L144-L145
Added lines #L144 - L145 were not covered by tests

[warning] 147-148: src/imgtools/coretypes/vector_shift.py#L147-L148
Added lines #L147 - L148 were not covered by tests

[warning] 150-150: src/imgtools/coretypes/vector_shift.py#L150
Added line #L150 was not covered by tests

⏰ Context from checks skipped due to timeout of 90000ms (3)

GitHub Check: Integration-Tests (windows-latest, py313, public)
GitHub Check: Integration-Tests (windows-latest, py310, public)
GitHub Check: Linting

🔇 Additional comments (3)

src/imgtools/coretypes/vector_shift.py (3)

1-9: Imports are properly organized.

The imports follow good Python practices with standard library imports first, followed by third-party imports, and then local imports with appropriate spacing between groups. The imports are relevant to the functionality provided in this module.

🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 1-2: src/imgtools/coretypes/vector_shift.py#L1-L2
Added lines #L1 - L2 were not covered by tests

[warning] 4-6: src/imgtools/coretypes/vector_shift.py#L4-L6
Added lines #L4 - L6 were not covered by tests

[warning] 8-8: src/imgtools/coretypes/vector_shift.py#L8
Added line #L8 was not covered by tests

1-194: Overall: Missing unit tests for the implementation.

The static analysis hints indicate that none of the code is covered by tests. Given that this implements core functionality for mask operations, it should have proper unit tests.

Would you like me to help create a basic test suite for this module? I can provide a skeleton for testing the key functions and behaviors.

🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 1-2: src/imgtools/coretypes/vector_shift.py#L1-L2
Added lines #L1 - L2 were not covered by tests

[warning] 4-6: src/imgtools/coretypes/vector_shift.py#L4-L6
Added lines #L4 - L6 were not covered by tests

[warning] 8-8: src/imgtools/coretypes/vector_shift.py#L8
Added line #L8 was not covered by tests

[warning] 11-11: src/imgtools/coretypes/vector_shift.py#L11
Added line #L11 was not covered by tests

[warning] 27-28: src/imgtools/coretypes/vector_shift.py#L27-L28
Added lines #L27 - L28 were not covered by tests

[warning] 31-34: src/imgtools/coretypes/vector_shift.py#L31-L34
Added lines #L31 - L34 were not covered by tests

[warning] 37-42: src/imgtools/coretypes/vector_shift.py#L37-L42
Added lines #L37 - L42 were not covered by tests

[warning] 45-46: src/imgtools/coretypes/vector_shift.py#L45-L46
Added lines #L45 - L46 were not covered by tests

[warning] 49-52: src/imgtools/coretypes/vector_shift.py#L49-L52
Added lines #L49 - L52 were not covered by tests

[warning] 55-58: src/imgtools/coretypes/vector_shift.py#L55-L58
Added lines #L55 - L58 were not covered by tests

[warning] 61-64: src/imgtools/coretypes/vector_shift.py#L61-L64
Added lines #L61 - L64 were not covered by tests

[warning] 66-70: src/imgtools/coretypes/vector_shift.py#L66-L70
Added lines #L66 - L70 were not covered by tests

[warning] 72-74: src/imgtools/coretypes/vector_shift.py#L72-L74
Added lines #L72 - L74 were not covered by tests

[warning] 77-78: src/imgtools/coretypes/vector_shift.py#L77-L78
Added lines #L77 - L78 were not covered by tests

[warning] 81-81: src/imgtools/coretypes/vector_shift.py#L81
Added line #L81 was not covered by tests

[warning] 85-86: src/imgtools/coretypes/vector_shift.py#L85-L86
Added lines #L85 - L86 were not covered by tests

[warning] 89-89: src/imgtools/coretypes/vector_shift.py#L89
Added line #L89 was not covered by tests

[warning] 91-91: src/imgtools/coretypes/vector_shift.py#L91
Added line #L91 was not covered by tests

[warning] 97-97: src/imgtools/coretypes/vector_shift.py#L97
Added line #L97 was not covered by tests

[warning] 100-102: src/imgtools/coretypes/vector_shift.py#L100-L102
Added lines #L100 - L102 were not covered by tests

[warning] 105-105: src/imgtools/coretypes/vector_shift.py#L105
Added line #L105 was not covered by tests

[warning] 124-126: src/imgtools/coretypes/vector_shift.py#L124-L126
Added lines #L124 - L126 were not covered by tests

[warning] 129-129: src/imgtools/coretypes/vector_shift.py#L129
Added line #L129 was not covered by tests

[warning] 131-132: src/imgtools/coretypes/vector_shift.py#L131-L132
Added lines #L131 - L132 were not covered by tests

[warning] 134-135: src/imgtools/coretypes/vector_shift.py#L134-L135
Added lines #L134 - L135 were not covered by tests

[warning] 138-139: src/imgtools/coretypes/vector_shift.py#L138-L139
Added lines #L138 - L139 were not covered by tests

[warning] 141-142: src/imgtools/coretypes/vector_shift.py#L141-L142
Added lines #L141 - L142 were not covered by tests

[warning] 144-145: src/imgtools/coretypes/vector_shift.py#L144-L145
Added lines #L144 - L145 were not covered by tests

[warning] 147-148: src/imgtools/coretypes/vector_shift.py#L147-L148
Added lines #L147 - L148 were not covered by tests

[warning] 150-150: src/imgtools/coretypes/vector_shift.py#L150
Added line #L150 was not covered by tests

147-147:
✅ Verification successful

Potential error in index-to-name mapping logic.

There seems to be an assumption that indices in the loop (0 to n_components-1) should map to ROI indices starting from 1, with 0 being a background, but this may not always be true.

🏁 Script executed:
#!/bin/bash
# Verify how ROI mapping is typically structured in the codebase
# and whether index 0 is consistently used for background

# Check usage patterns of ROIMaskMapping
rg -A 3 -B 3 "roi_mapping\s*=\s*{" 

# Look at how VectorMask ROI indices are usually assigned
rg -A 5 "VectorMask\(" --type py
Length of output: 2516

I need to inspect how index_to_roi_mapping and the ROI mapping are built to confirm whether the +1 offset is correct. Let me pull in the surrounding code.

🏁 Script executed:
#!/bin/bash
# Check file size and relevant mapping definitions
wc -l src/imgtools/coretypes/vector_shift.py
rg -n "index_to_roi_mapping" src/imgtools/coretypes/vector_shift.py || true
rg -n "roi_mapping" src/imgtools/coretypes/vector_shift.py || true

# Show the block where index_to_roi_mapping is created and used
sed -n '130,200p' src/imgtools/coretypes/vector_shift.py
Length of output: 2566

Gathering how create_demo_vector_mask defines component indices to verify mapping…

🏁 Script executed:
#!/bin/bash
# Locate where create_demo_vector_mask is defined
rg -n "def create_demo_vector_mask" -C 3 src/imgtools/coretypes

# Show its implementation context
sed -n '1,200p' src/imgtools/coretypes/vector_shift.py
Length of output: 6561

Digging into how VectorMask stores its roi_mapping to confirm whether it’s 0- or 1-based:
#!/bin/bash
# Inspect VectorMask to see how roi_mapping keys are defined
rg -n "class VectorMask" -A 50 src/imgtools/coretypes/base_masks.py
Inspecting how VectorMask processes roi_mapping keys:

🏁 Script executed:
#!/bin/bash
rg -n "class VectorMask" -A200 src/imgtools/coretypes/base_masks.py
Length of output: 8752
No off-by-one error in ROI mapping – the +1 offset is intentional and correct

The collapse_vector_mask function uses i + 1 to look up each component’s ROIMaskMapping because VectorMask always reserves key 0 for “Background” and shifts user-provided indices up by 1. Your mapping of bits 0–3 to keys 1–4 in roi_mapping lines up perfectly with this design, so no change is needed.

🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 147-148: src/imgtools/coretypes/vector_shift.py#L147-L148
Added lines #L147 - L148 were not covered by tests

src/imgtools/coretypes/vector_shift.py

…nd component

coderabbitai

Actionable comments posted: 1

🔭 Outside diff range comments (1)

src/imgtools/coretypes/vector_shift.py (1)
1-196: 💡 Verification agent
❓ Verification inconclusive

Add unit tests for the module

The static analysis reports indicate that most of the code is not covered by tests. For a critical component like this, consider adding comprehensive unit tests that cover:

The create_demo_vector_mask function

The collapse_vector_mask function with various inputs

Edge cases like empty masks or single-component masks

This would improve robustness and help catch regressions in future changes.

Would you like me to help generate some unit tests for this module?

🏁 Script executed:
#!/bin/bash
# Check if there are any existing tests for this module
fd -e py 'test.*vector_shift' -t f
Length of output: 36
Add comprehensive unit tests for vector_shift module

It looks like there are currently no tests covering src/imgtools/coretypes/vector_shift.py. Please add unit tests to ensure correctness and guard against regressions:

Tests for create_demo_vector_mask
• Verify output image dimensions and pixel-component count
• Check that each synthetic shape (circle, square, stripe, clover) sets the expected pixels to 1
• Confirm the returned mapping matches component indices

Tests for collapse_vector_mask
• Input a mask with all zeros → expect label image of all zeros and an empty lookup table
• Single-component masks → ensure bitmask encoding equals the component index (e.g. only component 0 set yields label 1)
• Multi-component combinations → verify label values equal sum of 2**i for active components and lookup_table entries match image IDs
• Round-trip test: compose a known vector mask, collapse it, then reconstruct and compare

Edge cases
• Zero-sized or minimal masks
• Masks with non-binary values (should raise or sanitize)

Place your tests under tests/test_vector_shift.py (or similar) using pytest. This will greatly improve maintainability and catch future bugs early.

🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 1-2: src/imgtools/coretypes/vector_shift.py#L1-L2
Added lines #L1 - L2 were not covered by tests

[warning] 4-6: src/imgtools/coretypes/vector_shift.py#L4-L6
Added lines #L4 - L6 were not covered by tests

[warning] 8-8: src/imgtools/coretypes/vector_shift.py#L8
Added line #L8 was not covered by tests

[warning] 11-11: src/imgtools/coretypes/vector_shift.py#L11
Added line #L11 was not covered by tests

[warning] 27-28: src/imgtools/coretypes/vector_shift.py#L27-L28
Added lines #L27 - L28 were not covered by tests

[warning] 31-34: src/imgtools/coretypes/vector_shift.py#L31-L34
Added lines #L31 - L34 were not covered by tests

[warning] 37-42: src/imgtools/coretypes/vector_shift.py#L37-L42
Added lines #L37 - L42 were not covered by tests

[warning] 45-46: src/imgtools/coretypes/vector_shift.py#L45-L46
Added lines #L45 - L46 were not covered by tests

[warning] 49-52: src/imgtools/coretypes/vector_shift.py#L49-L52
Added lines #L49 - L52 were not covered by tests

[warning] 55-58: src/imgtools/coretypes/vector_shift.py#L55-L58
Added lines #L55 - L58 were not covered by tests

[warning] 61-64: src/imgtools/coretypes/vector_shift.py#L61-L64
Added lines #L61 - L64 were not covered by tests

[warning] 66-70: src/imgtools/coretypes/vector_shift.py#L66-L70
Added lines #L66 - L70 were not covered by tests

[warning] 72-74: src/imgtools/coretypes/vector_shift.py#L72-L74
Added lines #L72 - L74 were not covered by tests

[warning] 77-78: src/imgtools/coretypes/vector_shift.py#L77-L78
Added lines #L77 - L78 were not covered by tests

[warning] 81-81: src/imgtools/coretypes/vector_shift.py#L81
Added line #L81 was not covered by tests

[warning] 85-86: src/imgtools/coretypes/vector_shift.py#L85-L86
Added lines #L85 - L86 were not covered by tests

[warning] 89-89: src/imgtools/coretypes/vector_shift.py#L89
Added line #L89 was not covered by tests

[warning] 91-91: src/imgtools/coretypes/vector_shift.py#L91
Added line #L91 was not covered by tests

[warning] 97-97: src/imgtools/coretypes/vector_shift.py#L97
Added line #L97 was not covered by tests

[warning] 100-102: src/imgtools/coretypes/vector_shift.py#L100-L102
Added lines #L100 - L102 were not covered by tests

[warning] 105-105: src/imgtools/coretypes/vector_shift.py#L105
Added line #L105 was not covered by tests

[warning] 124-126: src/imgtools/coretypes/vector_shift.py#L124-L126
Added lines #L124 - L126 were not covered by tests

[warning] 131-131: src/imgtools/coretypes/vector_shift.py#L131
Added line #L131 was not covered by tests

[warning] 133-134: src/imgtools/coretypes/vector_shift.py#L133-L134
Added lines #L133 - L134 were not covered by tests

[warning] 136-137: src/imgtools/coretypes/vector_shift.py#L136-L137
Added lines #L136 - L137 were not covered by tests

[warning] 140-141: src/imgtools/coretypes/vector_shift.py#L140-L141
Added lines #L140 - L141 were not covered by tests

[warning] 143-144: src/imgtools/coretypes/vector_shift.py#L143-L144
Added lines #L143 - L144 were not covered by tests

[warning] 146-147: src/imgtools/coretypes/vector_shift.py#L146-L147
Added lines #L146 - L147 were not covered by tests

[warning] 149-150: src/imgtools/coretypes/vector_shift.py#L149-L150
Added lines #L149 - L150 were not covered by tests

[warning] 152-152: src/imgtools/coretypes/vector_shift.py#L152
Added line #L152 was not covered by tests

🧹 Nitpick comments (6)

src/imgtools/coretypes/vector_shift.py (6)
81-83: Replace magic numbers with math.pi constant

Using hardcoded π value (3.14159) isn't ideal. Python's math module provides a more precise math.pi constant.
+ import math
...
- threshold = clover_radius_inner + (
-     clover_radius_outer - clover_radius_inner
- ) * abs(angle - 3.14159 / num_points) / (3.14159 / num_points)
+ threshold = clover_radius_inner + (
+     clover_radius_outer - clover_radius_inner
+ ) * abs(angle - math.pi / num_points) / (math.pi / num_points)
🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 81-81: src/imgtools/coretypes/vector_shift.py#L81
Added line #L81 was not covered by tests

105-123: Inconsistency between docstring parameter and actual parameter type

The docstring lists vector_mask parameter as sitk.Image, but the type hint specifies VectorMask (a specific class from your codebase). Ensure the docstring matches the actual parameter type for clarity.
    Parameters
    ----------
-   vector_mask : sitk.Image
+   vector_mask : VectorMask
        A VectorUInt8 image where each component is 0 or 1.
🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 105-105: src/imgtools/coretypes/vector_shift.py#L105
Added line #L105 was not covered by tests

123-123: Docstring return description is inaccurate

The docstring says the lookup table maps label values to "joined component name string", but the implementation maps label values to a list of strings (names). Update the docstring to match the actual implementation.
    Returns
    -------
    PartitionResult
        - label_image: scalar UInt16 image of unique bitmask-encoded labels
-       - lookup_table: dict[label_value] = joined component name string
+       - lookup_table: dict[label_value] = list of component image_ids
146-150: Add comment explaining the bitmask decoding logic

The bit manipulation logic for decoding the bitmask is correct but could benefit from an explanatory comment for future maintainers.
    for value in range(1, max_val):
+       # Extract component indices from the bitmask value using bitwise operations
        indices = [i for i in range(n_components) if (value >> i) & 1]
        # skip Bckround
        names = [index_to_roi_mapping[i + 1].image_id for i in indices]
        lookup_table[value] = names
🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 146-147: src/imgtools/coretypes/vector_shift.py#L146-L147
Added lines #L146 - L147 were not covered by tests

[warning] 149-150: src/imgtools/coretypes/vector_shift.py#L149-L150
Added lines #L149 - L150 were not covered by tests

155-196: Consider moving demo code to a separate module or function

The demo code in the if __name__ == "__main__" block writes files to the module's directory, which might not be writable in production environments. Consider:

Moving the demo to a separate module intended for examples

Using a configurable output directory or a temporary directory
if __name__ == "__main__":
-   output_dir = Path(__file__).parent
+   # Use a temporary directory or make output configurable
+   import tempfile
+   output_dir = Path(tempfile.gettempdir())
91-96: Consider using an Enum for component mappings

For better type safety and self-documentation, consider using an Enum for the component mappings:
+ from enum import IntEnum
+
+ class DemoComponent(IntEnum):
+     CIRCLE = 0
+     SQUARE = 1
+     STRIPE = 2
+     CLOVER = 3
+
  mapping = {
-     "circle": 0,
-     "square": 1,
-     "stripe": 2,
-     "clover": 3,
+     "circle": DemoComponent.CIRCLE,
+     "square": DemoComponent.SQUARE,
+     "stripe": DemoComponent.STRIPE,
+     "clover": DemoComponent.CLOVER,
  }
This makes the code more robust against errors from mistyped string values.

🧰 Tools

🪛 GitHub Check: codecov/patch

[warning] 91-91: src/imgtools/coretypes/vector_shift.py#L91
Added line #L91 was not covered by tests

📜 Review details

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 6c3ecdb and fd5ecbb.

📒 Files selected for processing (1)

src/imgtools/coretypes/vector_shift.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (1)

`src/**/*.py`: Review the Python code for compliance with PEP 8 and PEP 257 (docstring conventions). Ensure the following: - Variables and functions follow meaningful naming conven...

src/**/*.py: Review the Python code for compliance with PEP 8 and PEP 257 (docstring conventions). Ensure the following: - Variables and functions follow meaningful naming conventions. - Docstrings are present, accurate, and align with the implementation. - Code is efficient and avoids redundancy while adhering to DRY principles. - Consider suggestions to enhance readability and maintainability. - Highlight any potential performance issues, edge cases, or logical errors. - Ensure all imported libraries are used and necessary.

src/imgtools/coretypes/vector_shift.py