refactor(loss): extract shift_mu function and simplify onsite shift calculation

The shift_mu function was extracted to avoid code duplication across multiple loss classes. The onsite shift calculation was simplified by using a more accurate formula that accounts for both node and edge features.

Author: QG-phy