kindred/solver
1
0
Fork 0
Code Issues 10 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

feat(datasets): PyG dataset adapter for assembly shard files #17

New Issue
Closed
opened 2026-02-04 20:22:10 +00:00 by forbes · 0 comments
forbes commented 2026-02-04 20:22:10 +00:00
Owner
Copy Link

Summary

Implement solver/datasets/assembly.py with an AssemblyDataset(torch_geometric.data.Dataset) that loads shard files from data/synthetic/shards/ and converts each example to a torch_geometric.data.Data object using feature encoding from #16.

Depends on #16.

Data object fields

Field Shape Description
x [n_bodies, 22] Node features
edge_index [2, 2*n_joints] Bidirectional edges from joint body_a/body_b
edge_attr [2*n_joints, 22] Edge features (duplicated for both directions)
y_edge [2*n_joints] Per-edge independence labels
y_graph [1] Graph classification label
y_joint_type [2*n_joints] Per-edge joint type labels
y_dof [1] Scalar DOF target
y_body_dof [n_bodies, 2] Per-body translational + rotational DOF

Requirements

  • Load both .pt and .json shard formats
  • Support train/val/test splits via index files
  • Compatible with torch_geometric.loader.DataLoader for batching

Files

  • solver/datasets/assembly.py
  • solver/datasets/__init__.py
  • tests/datasets/__init__.py
  • tests/datasets/test_assembly.py

Acceptance criteria

  • Loads a shard file and produces valid Data objects
  • Bidirectional edges with correct feature duplication
  • Compatible with DataLoader (batching works)
  • Passes ruff, mypy, pytest
## Summary Implement `solver/datasets/assembly.py` with an `AssemblyDataset(torch_geometric.data.Dataset)` that loads shard files from `data/synthetic/shards/` and converts each example to a `torch_geometric.data.Data` object using feature encoding from #16. Depends on #16. ## Data object fields | Field | Shape | Description | |-------|-------|-------------| | `x` | `[n_bodies, 22]` | Node features | | `edge_index` | `[2, 2*n_joints]` | Bidirectional edges from joint body_a/body_b | | `edge_attr` | `[2*n_joints, 22]` | Edge features (duplicated for both directions) | | `y_edge` | `[2*n_joints]` | Per-edge independence labels | | `y_graph` | `[1]` | Graph classification label | | `y_joint_type` | `[2*n_joints]` | Per-edge joint type labels | | `y_dof` | `[1]` | Scalar DOF target | | `y_body_dof` | `[n_bodies, 2]` | Per-body translational + rotational DOF | ## Requirements - Load both `.pt` and `.json` shard formats - Support train/val/test splits via index files - Compatible with `torch_geometric.loader.DataLoader` for batching ## Files - `solver/datasets/assembly.py` - `solver/datasets/__init__.py` - `tests/datasets/__init__.py` - `tests/datasets/test_assembly.py` ## Acceptance criteria - [ ] Loads a shard file and produces valid `Data` objects - [ ] Bidirectional edges with correct feature duplication - [ ] Compatible with `DataLoader` (batching works) - [ ] Passes ruff, mypy, pytest
Sign in to join this conversation.
No Branch/Tag Specified
Branches Tags
Labels
Clear labels
Compat/Breaking
Breaking change that won't be backward compatible
 Kind/Bug
Something is not working
 Kind/Documentation
Documentation changes
 Kind/Enhancement
Improve existing functionality
 Kind/Feature
New functionality
 Kind/Security
This is security issue
 Kind/Testing
Issue or pull request related to testing
Priority
Critical
1
The priority is critical
Priority
High
2
The priority is high
Priority
Low
4
The priority is low
Priority
Medium
3
The priority is medium
Reviewed
Confirmed
1
Issue has been confirmed
Reviewed
Duplicate
2
This issue or pull request already exists
Reviewed
Invalid
3
Invalid issue
Reviewed
Won't Fix
3
This issue won't be fixed
Status
Abandoned
3
Somebody has started to work on this but abandoned work
Status
Blocked
1
Something is blocking this issue or pull request
Status
Need More Info
2
Feedback is required to reproduce issue or to continue work
feature
New feature or enhancement
 phase:1
Phase 1: Synthetic Data Generation
 port
Porting existing code into repository
 testing
Tests and test infrastructure
No Label
Milestone
No items
No Milestone
Projects
Clear projects
No project
Assignees
Clear assignees
forbes
No Assignees
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: kindred/solver#17
Delete Branch "%!s()"

Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?