Это упрощенный подход. Если хотите точности, плавности на всех занчениях и прочего - нужно реализовывать модель Брауна-Конрада итеративно проходов в 4-5, высчитывая куда прилетит луч, и после - откуда. Формулы будут +/- те же за исключением того, что нужно будет считать промежуточный вектор умножая на коэффициент (или прибавляя тангенциальное смещение), а потом приращивать результат. То есть что-то вроде такого:

static inline float2 distort(float2 uv, constant DecenteringUniforms& params) {
    // Fixed-point: xu_{n+1} = xu_n + (xd - distort_forward(xu_n))
    float2 d = uv - params.center;
    float2 scale = float2(params.aspect.x, params.aspect.y);
    float2 xd = d * scale;

    float2 xu = xd;
    for (int i=0; i<5; ++i) {
        float2 xh = forward(xu, params);
        xu += (xd - xh);
    }

    return params.center + xu / scale;
}