Fix solving linear and square normalizer

[mahaitongdae]

Align normalizer implementations with batch target mean and lower bound semantics

Overview

This PR fixes mathematical inconsistencies in the linear and squared normalizer functions. Previously, the normalizers were solving for a target sum instead of a target mean, and the squared normalizer lacked support for a non-zero lower bound.

These updates align the standalone normalizer functions (solve_normalizer_linear and solve_normalizer_square) with the correct batch implementations (solve_v_batch and solve_v_squared_batch), ensuring they solve for a target mean of 1 across $N$ particles.

Changes Made

1. Linear Normalizer Fixes

• The Issue: The original implementation solved for a target sum of 1 instead of a target mean of 1 (which requires a target sum of $N$). Furthermore, it did not scale the lower bound threshold correctly relative to $N$.
• The Fix: Updated the active element threshold and the final scalar calculation.
• Math Details: The target sum was updated from 1 to N. The temperature/scale ($\tau$) is now properly scaled as $N\tau$, and the lower bound floor ($\gamma$) is scaled as $\gamma/N$ to match the standard batch formula.

2. Squared Normalizer Fixes

• The Issue: The original implementation was a hardcoded special case that assumed a clipping floor of 0.0. It also solved for a target sum of $\tau^2$ instead of the required batch target sum of $N\tau^2$.
• The Fix: Completely rewrote the energy check and quadratic root calculation to handle non-zero lower bounds.
• Math Details: * Replaced the simple excess calculation with a full quadratic expansion that accounts for the energy of clipped elements (clamped at the floor $B$).
  • Updated the target energy from temp**2 to N * (temp**2).
  • Updated the discriminant ($\Delta$) to correctly subtract the fixed energy of the inactive floor.

Impact

This is a mathematical correctness fix. Any downstream losses, projections, or layers relying on these normalizers will now correctly target a mean of 1 across the particle dimension.