perf: GPU memory optimization, torch_scatter compat, NaN handling, and robustness fixes

[liu-687]

Summary

This PR includes a series of optimizations and robustness improvements to the DeepH-pack codebase, mainly focused on memory efficiency, portability, and inference reliability.

Changes

🚀 GPU Memory Optimization (kernel.py)

• make_mask now moves only necessary tensors (index_to_Z, x, edge_index) to GPU individually instead of the entire data object
• mask and label tensors are kept on CPU after computation, dramatically reducing GPU memory usage
• Explicit del + torch.cuda.empty_cache() after each sample to prevent memory fragmentation
• Impact: Prevents OOM errors when training on large systems (e.g., twisted vdW materials with hundreds of atoms)

🔧 Portability — Native PyTorch Scatter Fallback (kernel.py, model.py, graph_norm.py)

• torch_scatter imports now use try/except ImportError with fallback to native PyTorch via new compat_scatter.py
• compat_scatter.py provides drop-in replacements for scatter_add, scatter, scatter_mean using built-in scatter_add_
• Impact: Package works without the torch_scatter C++ extension on platforms where compilation is problematic

🛡️ Robustness Improvements

• pred_ham.py: Replaces NaN entries with zeros for orbital pairs not covered by the trained model, instead of crashing with assertion error. Logs count of replaced entries for visibility
• graph.py: Gracefully handles single-sample inference without DFT subgraph data (returns None for subgraph instead of asserting)

🐛 Bug Fixes

• evaluate.py: Fixed device mismatch — reshapes output before moving to label's device instead of hardcoding .cpu()
• data.py: Cleaned up redundant imports, added explicit GC calls to recover pool worker memory after multiprocessing

Files Changed

File	Change
deeph/kernel.py	GPU memory management + torch_scatter fallback
deeph/data.py	Delayed collation, GC, redundant import cleanup
deeph/model.py	torch_scatter fallback
deeph/graph.py	Single-sample inference LCMP handling
deeph/inference/pred_ham.py	NaN → zero for uncovered orbital pairs
deeph/scripts/evaluate.py	Device handling fix
deeph/from_PyG_future/graph_norm.py	torch_scatter fallback
deeph/compat_scatter.py	[new] Native PyTorch scatter implementations
.gitignore	[new] Exclude __pycache__/, *.egg-info/

Backward Compatibility

• All changes are backward compatible
• Existing graph files load correctly (new collation-at-load-time format is auto-detected)
• Old graph files without subgraph_dict are handled gracefully

[Copilot]

Pull request overview

This PR improves DeepH-pack’s robustness and portability by reducing GPU memory pressure during mask/label preparation, adding a native PyTorch fallback for torch_scatter, and hardening inference utilities against missing/invalid values.

Changes:

• Added torch_scatter import fallbacks to a new compat_scatter.py implementation for broader platform support.
• Reduced GPU memory usage in DeepHKernel.make_mask via targeted tensor moves and per-sample cleanup.
• Improved inference robustness (NaN filling in pred_ham.py, relaxed LCMP subgraph assumptions in graph.py) and fixed a device handling issue in evaluate.py.

Reviewed changes

Copilot reviewed 8 out of 9 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
deeph/scripts/evaluate.py	Fixes output reshaping/device placement to avoid device mismatch.
deeph/model.py	Adds torch_scatter fallback imports to compat_scatter.
deeph/kernel.py	Adjusts mask/label creation workflow for lower GPU memory usage and adds scatter fallback import.
deeph/inference/pred_ham.py	Replaces NaN assertion with NaN→0 filling + logging during block restoration.
deeph/graph.py	Changes LCMP collation behavior when subgraph data is missing.
deeph/from_PyG_future/graph_norm.py	Adds scatter_mean fallback import.
deeph/data.py	Saves un-collated data_list and collates on load; adds GC and dataset slicing logic.
deeph/compat_scatter.py	New: native PyTorch scatter/scatter_add/scatter_mean compatibility layer.
.gitignore	New: ignores common Python build/cache artifacts.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.