Optimize raytracing core: device allocation, sqrt avoidance, polynomial eval

[ejin-photonium]

• Ray: create tensors directly on target device to avoid CPU->GPU copies
• Aspheric: refactor polynomial evaluation with helper methods and multiplication chains instead of pow()
• Boundary checks: use squared comparisons instead of sqrt() in base.py, spheric.py, plane.py, phase.py

speeds up autolens training by ~15% on GPU, mostly from better raytracing device initialization. other improvements are minor (~2% improvement or so)

[Copilot]

Pull request overview

This PR optimizes the raytracing core for GPU performance through three main strategies: direct device allocation, avoiding expensive square root operations, and refactoring polynomial evaluation with multiplication chains.

Changes:

• Ray initialization creates tensors directly on target device to avoid CPU→GPU transfers (~15% speedup)
• Boundary checks use squared distance comparisons instead of sqrt() for performance
• Aspheric surface polynomial evaluation refactored with helper methods using multiplication chains instead of pow()

Reviewed changes

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

Show a summary per file

File	Description
deeplens/optics/ray.py	Creates all ray tensors directly on target device, stores device attribute instead of calling to(device)
deeplens/optics/phase_surface/phase.py	Boundary check uses squared comparison (r²) instead of sqrt
deeplens/optics/geometric_surface/spheric.py	Boundary check uses squared comparison (r²) instead of sqrt
deeplens/optics/geometric_surface/plane.py	Boundary check uses squared comparison (r²) instead of sqrt
deeplens/optics/geometric_surface/base.py	is_within_boundary uses squared comparison (r²) instead of sqrt
deeplens/optics/geometric_surface/aspheric.py	Refactored polynomial evaluation with helper methods and multiplication chains for performance

---

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

[ejin-photonium]

The copilot code review wants to differentiate the cases where wvln is a float vs tensor. however, looking at the codebase, it seems that wvln is always passed in as a float and the docstring is wrong (hence why the assertion never fails).

[singer-yang]

Hi @ejin-photonium,

Thank you very much for your interest and contribution!

I am happy with the Ray construction and boundary checks. For the aspheric polynomial function refactor, I prefer keeping the original version, which is easier for users to read since it remains in a single function. Could you please roll this back?

I use Copilot for pre-review, and I will perform a manual review before merging your PR.

Thanks again, and have a nice day!

Best,
Xinge

[ejin-photonium]

Hi Xinge. To be honest the many if-elif statements had bothered me and seemed a bit more unreadable/overengineered. I refactored the commit inlining the functions and added a comment so that it is easier for users to read as per your suggestion. Let me know if that works; if not, I'm happy to revert given the performance gains were not substantial.

[singer-yang]

Thank you very much! I merged it into the dev branch which will later merge to main