Repository navigation
refactor: simplify data module and condition layer - #837
Merged
Merged
Conversation
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
There are confirmed runtime-breaking issues in the new condition/data path (notably graph input equation evaluation and slice indexing for graph conditions) that must be fixed before merging.
Review effort: Lite
Findings: 2
Open (2)
What changed in this PR
This PR refactors PINA’s data pipeline by removing the legacy data-manager/creator/aggregator stack and shifting batching/materialization responsibilities into condition implementations, with DataModule now operating on per-condition sample-id tensors and materializing condition batches on-device.
Changes:
- Replaced
_Aggregator+_CreatorwithBatcher+MultiLoader, and rewiredDataModuleto split per-condition id tensors and callcondition.materialize(...)during device transfer. - Introduced
TensorCondition/GraphConditionbase classes to unify data storage + batching for tensor and graph conditions; updated condition variants to use these. - Updated
TrainerAPI to dropautomatic_batchingand add per-conditiondataloader_cls/collate_fnresolution; adapted callbacks and tests accordingly.
| File | Description |
|---|---|
| tests/test_trainer.py | Updates Trainer constructor tests and adds coverage for per-condition dataloader/collate option validation & warnings. |
| tests/test_data/test_tensor_data_manager.py | Removes tests for deleted tensor data manager abstraction. |
| tests/test_data/test_single_batch_data_loader.py | Removes tests for deleted single-batch loader. |
| tests/test_data/test_loader.py | Replaces _Aggregator tests with MultiLoader tests. |
| tests/test_data/test_graph_data_manager.py | Removes tests for deleted graph data manager abstraction. |
| tests/test_data/test_data_module.py | Updates DataModule tests to assert id-tensor subsets instead of _ConditionSubset. |
| tests/test_data/test_creator.py | Removes tests for deleted dataloader creator. |
| tests/test_data/test_condition_subset.py | Removes tests for deleted condition subset wrapper. |
| tests/test_condition/test_time_series_condition.py | Updates tests to use materialize and new tensor-backed condition storage. |
| tests/test_condition/test_input_target_condition.py | Updates tests for dict-based batches and materialize semantics across tensor/graph variants. |
| tests/test_condition/test_input_equation_condition.py | Updates tests to use materialize for tensor/graph inputs. |
| tests/test_condition/test_graph_time_series_condition.py | Updates tests for GraphTimeSeriesCondition and materialize behavior. |
| tests/test_condition/test_data_condition.py | Updates tests for dict-based batches and conditional variables under new materialization. |
| tests/test_callback/test_data_normalizer.py | Aligns DataNormalizer tests with condition-owned data storage. |
| pina/data/manager.py | Removes top-level exports for deleted data manager layer. |
| pina/data/__init__.py | Exposes new Batcher/MultiLoader APIs and removes old internal exports. |
| pina/condition/__init__.py | Updates public condition exports to include new graph/tensor condition base types. |
| pina/_src/data/single_batch_data_loader.py | Deletes single-batch loader implementation. |
| pina/_src/data/manager/tensor_data_manager.py | Deletes tensor data manager implementation. |
| pina/_src/data/manager/graph_data_manager.py | Deletes graph data manager implementation. |
| pina/_src/data/manager/data_manager.py | Deletes data manager factory. |
| pina/_src/data/manager/data_manager_interface.py | Deletes data manager interface. |
| pina/_src/data/manager/batch_manager.py | Deletes legacy batch container. |
| pina/_src/data/loader.py | Adds MultiLoader to aggregate per-condition dataloaders. |
| pina/_src/data/data_module.py | Refactors DataModule to id-based subsets + on-device condition.materialize batching. |
| pina/_src/data/creator.py | Deletes old multi-condition dataloader creator. |
| pina/_src/data/condition_subset.py | Deletes old subset wrapper used for cyclic indexing/auto-batching. |
| pina/_src/data/batcher.py | Adds Batcher to construct per-condition dataloaders over id tensors. |
| pina/_src/data/aggregator.py | Deletes old _Aggregator implementation. |
| pina/_src/core/trainer.py | Drops automatic_batching, adds per-condition dataloader/collate option resolution, and passes into DataModule. |
| pina/_src/condition/time_series_condition.py | Refactors time-series condition to tensor-backed storage with shared unroll helper. |
| pina/_src/condition/tensor_condition.py | Introduces tensor-backed storage/materialization implementation for conditions. |
| pina/_src/condition/input_target_condition.py | Switches to tensor/graph condition variants and dict-based batches. |
| pina/_src/condition/input_equation_condition.py | Switches to tensor/graph condition variants and shared evaluation path. |
| pina/_src/condition/graph_time_series_condition.py | Refactors graph time-series condition to graph-backed storage/materialization. |
| pina/_src/condition/graph_condition.py | Introduces graph-backed storage/materialization implementation for conditions. |
| pina/_src/condition/domain_equation_condition.py | Clarifies non-materializable semantics for domain-sampled conditions. |
| pina/_src/condition/data_condition.py | Switches to tensor/graph condition variants and updates semantics/docs. |
| pina/_src/condition/condition_interface.py | Updates condition contract to center on materialize(...). |
| pina/_src/condition/base_condition.py | Removes legacy dataloader/collate mechanisms from the base class and adds _unwrap_single. |
| pina/_src/callback/refinement/r3_refinement.py | Avoids autograd graph retention by detaching selected refinement points. |
| pina/_src/callback/refinement/base_refinement.py | Updates refinement bookkeeping for id-tensor datasets. |
| pina/_src/callback/processing/data_normalizer.py | Refactors normalizer to operate on condition-held data rather than dataset wrappers. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
GiovanniCanali
requested changes
Oct 1, 2026
GiovanniCanali
left a comment
GiovanniCanali
left a comment
Collaborator
There was a problem hiding this comment.
Great job, @ndem0! This significantly simplifies the logic connecting conditions and batching. However, there are still a few points that require attention:
- The documentation does not currently reflect the changes introduced in the source code. Please update it accordingly.
- Is the
Conditionfactory class affected in any way by the proposed changes? Please verify that its behavior remains consistent. - Tests for
Batcherare currently missing. - Each condition currently follows a common structure, with a base class that is then specialized into
TensorandGraphimplementations.TimeSeriesCondition, however, does not follow this pattern, as it develops two independent classes. Would it be possible to introduce a common base class with corresponding specializations, consistently with the other condition classes? - The
materializemethod appears to be identical in bothTensorConditionandGraphCondition. Consider moving this implementation toBaseConditionto avoid duplication.
GiovanniCanali
added
enhancement
dario-coscia
requested changes
Oct 8, 2026
dario-coscia
left a comment
dario-coscia
left a comment
Member
There was a problem hiding this comment.
Thanks @ndem0 ! I have checked the src files only and left some comments. I did not check the tests, I will do once the PR is ready for review and not draft anymore.
Anyway, I like the approach as it also more readable and reduces the number of lines of code!
ndem0
marked this pull request as ready for review
November 3, 2026 11:20
Contributor
|
Code Coverage SummaryResults for commit: d9623f5 Minimum allowed coverage is ♻️ This comment has been updated with latest results |
ndem0
added
pr-to-review
ndem0
force-pushed
the
data_module_simplification
branch
from
November 3, 2026 12:18
d1601c3 to
bed8c9f
Compare
GiovanniCanali
requested changes
Nov 4, 2026
GiovanniCanali
left a comment
GiovanniCanali
left a comment
Collaborator
There was a problem hiding this comment.
Thank you @ndem0, this PR considerably simplifies the logic governing conditions and batch management.
There are four main points I would like to address before approval:
- The previous implementation supported
InputEquationConditionwithGraphinput, but the support appears to have been lost. Can we restore it, or is there a specific reason it is no longer supported? - Several docstrings do not meet our documentation standards. I’ve left comments highlighting where changes are needed.
- I like the idea of grouping
.rstfiles by module to make the API landing page cleaner. However, the current_code.rststructure is confusing. Please see the comment below. - Please restore
conf.pyto use PINA’s documentation extension. If required dependencies are missing, add them topyproject.tomlrather than switching to local resources.
Member
Author
@GiovanniCanali Totally not! Many issues were also because of that: several classes were managing more than one implementation internally, making the structure overcomplicated and quite impossibile to debug/check. How about batching in this mixed condition tensor/graph? We have to support new method again, the user is probably unaware about the real class he's using, and I do not see any advantage. The real difference is in the problem statement. |
Collaborator
I totally agree with you, and I probably did not explain myself well: correct me if I am wrong, but the current implementation is lacking |
Member
Author
|
Already in this PR with the unittest |
Collaborator
No, this is |
Member
Author
|
Yeah right, no not implemented yet. Ill add so we can merge. |
Collaborator
|
Ready to merge @ndem0! Consider squashing some commits. Also, the first two commits are very similar. |
GiovanniCanali
approved these changes
Nov 5, 2026
Member
|
Looks great for me as well! We cut many lines (and a lot of glue code!). Thanks Nico 🚀 |
Remove data managers (Tensor/Graph/Batch), aggregator, condition subset, creator and single-batch loader. Conditions now store data directly in tensor/graph variants and materialize into id-based batches; graph tensor fields are attached per-graph so batching concatenates node-wise. Trainer drops automatic_batching in favour of per-condition dataloader_cls and collate_fn resolution. Adapt DataNormalizer and R3Refinement callbacks and rewrite tests for the new API.
Delegate subgraph creation to Graph.create_subgraph (k-hop, LabelTensor-label-preserving). Add GraphInputTargetCondition as a public supervised graph condition. Update SupervisedProblem to detect Graph/Data input and wrap it automatically.
BaseWeighting.aggregate merges per-condition weight updates to support one-condition-per-step batching. DataNormalizer and BaseRefinement/ R3Refinement updated to use train_dataset tensor and batch-dict API. Trainer warns when num_workers/pin_memory are set with batch_size=None. All tests updated to the new condition and DataModule interface.
ndem0
force-pushed
the
data_module_simplification
branch
from
November 6, 2026 08:59
36baee0 to
f0c1e6b
Compare
Add repo root to sys.path so Sphinx uses local source. Remove stale RST files for deleted data classes. Create per-module index.rst files with hidden toctrees matching the pina.* hierarchy for a collapsible sidebar. Remove automodule directives from class pages to fix right-column TOC. Clean the src directory, remove unnecessary code and file. Co-authored-by: Giovanni Canali <GiovanniCanali@users.noreply.github.com>
ndem0
force-pushed
the
data_module_simplification
branch
from
November 6, 2026 09:10
f0c1e6b to
d9623f5
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Remove data managers (Tensor/Graph/Batch), aggregator, condition subset, creator and single-batch loader. Conditions now store data directly in tensor/graph variants and materialize into id-based batches; graph tensor fields are attached per-graph so batching concatenates node-wise. Trainer drops automatic_batching in favour of per-condition dataloader_cls and collate_fn resolution. Adapt DataNormalizer and R3Refinement callbacks and rewrite tests for the new API.
This PR fixes #833
Checklist