feat: implement modified stage memory reset and scalar plotting for tensors - #9
Conversation
mattia-spider
commented
Apr 28, 2026
mattia-spider
commented
- Moved data cleaning at the beginning of stages
- Added _plot_combined_all_stages function to get parity plot with all stages combined at end on testing if requested.
- Added TensorsGraphParityPlotter to compute and plot requested properties of tensors, building up on ParityPlotter and GraphParityPlotter logic.
- Generalized _plot_parity e _plot_combined_all_stages in TensorsGraphParityPlotter to plot dictionaries with multiple keys.
…mbined plotting for tensors
|
Codecov Report❌ Patch coverage is
🚀 New features to boost your workflow:
|
muhrin
left a comment
muhrin
left a comment
There was a problem hiding this comment.
I think this needs some thinking. I would start with generalising out the plotting of individual plots onto a single matplotlib.Figure and then go from there, I think you will find that the code gets a lot shorter and a lot simpler.
| filename = f"{stage_name}_epoch_{epoch}.pdf" if epoch is not None else f"{stage_name}.pdf" | ||
| plt.savefig(str(save_dir / filename), bbox_inches="tight") | ||
| # full_path = save_dir / filename | ||
| # print(f"DEBUG: Saving plot to {full_path.absolute()}") |
|
|
||
| # 2. Clear the stage data for the next run (e.g., next 'fit' call) | ||
| self.data_store[stage_name] = ([], []) | ||
| # self.data_store[stage_name] = ([], []) |
|
|
||
| y_true, y_pred = self.get_target_predicted(batch, outputs) | ||
|
|
||
| # _LOGGER.info(f"Collected {len(y_true)} points for {stage_name}") # Log temporaneo |
| } | ||
| self._last_plotted_epoch.clear() | ||
|
|
||
| def _plot_combined_all_stages(self, trainer: "reax.Trainer", key_name: str | None = None): |
There was a problem hiding this comment.
This function violates the DRY principle (https://en.wikipedia.org/wiki/Don't_repeat_yourself). If you look above, you have code that does pretty much the same thing in _plot_parity and here you replicate it to plot multiple datasets on the same plot. Instead, you should have a single function that adds a lot to an existing figure (taking a fig object and anything else it needs), and then when you want to plot a 'combined' plot, you just loop over calling that and finally call savefig().
| """ | ||
| def __init__( | ||
| self, | ||
| targets: str = "nodes.nmr_tensors", |
There was a problem hiding this comment.
This cannot be hardcoded to NMR tensors. Remember, this is a general class that should not even know what an NMR tensor is.
| **kwargs | ||
| ) | ||
|
|
||
| self.scalar_keys = scalar_keys or [ |
There was a problem hiding this comment.
As we discussed, this is problematic because you've hard coded it to rank-2 tensors...ideally it would just plot a series of irreps.
3 participants
