diff --git a/src/body_aerodynamics.jl b/src/body_aerodynamics.jl
index b55e2828..1596b85f 100644
--- a/src/body_aerodynamics.jl
+++ b/src/body_aerodynamics.jl
@@ -1078,14 +1078,8 @@ function set_va!(body_aero::BodyAerodynamics, va::AbstractVector, omega=zeros(MV
     
     # Update wake elements
     frozen_wake!(body_aero, va_distribution)
-    if all(iszero, omega)
-        body_aero._va .= va
-        body_aero.has_distributed_va = false
-    else
-        body_aero._va .= _compute_reference_velocity_from_distribution(
-            va_distribution, n_panels)
-        body_aero.has_distributed_va = true
-    end
+    body_aero._va .= va
+    body_aero.has_distributed_va = false
     return nothing
 end
 
diff --git a/src/solver.jl b/src/solver.jl
index e18b6ff2..10ff3e50 100644
--- a/src/solver.jl
+++ b/src/solver.jl
@@ -1044,8 +1044,9 @@ function linearize(solver::Solver, body_aero::BodyAerodynamics, y::Vector{T};
         throw(ArgumentError("Cannot use theta_idxs or delta_idxs when wing has no unrefined sections"))
     end
 
-    init_va = body_aero.cache[1][body_aero.va]
-    init_va .= body_aero.va
+    init_va = body_aero.cache[1][body_aero._va]
+    init_va .= body_aero._va
+    init_omega = copy(body_aero.omega)
     last_theta_ref = Ref{Vector{T}}(Vector{T}(undef, 0))
     if !isnothing(theta_idxs)
         @views last_theta_ref[] = body_aero.cache[2][y[theta_idxs]]
@@ -1085,15 +1086,10 @@ function linearize(solver::Solver, body_aero::BodyAerodynamics, y::Vector{T};
             end
         end
 
-        if !isnothing(va_idxs) && isnothing(omega_idxs)
-            set_va!(body_aero, y[va_idxs])
-        elseif !isnothing(va_idxs) && !isnothing(omega_idxs)
-            set_va!(body_aero, y[va_idxs], y[omega_idxs])
-        elseif isnothing(va_idxs) && !isnothing(omega_idxs)
-            set_va!(body_aero, init_va, y[omega_idxs])
-        else
-            set_va!(body_aero, init_va)
-        end
+        va = isnothing(va_idxs) ? init_va : y[va_idxs]
+        om = isnothing(omega_idxs) ? init_omega :
+            y[omega_idxs]
+        set_va!(body_aero, va, om)
 
         solve!(solver, body_aero; kwargs...)
         if !aero_coeffs
diff --git a/test/body_aerodynamics/test_body_aerodynamics.jl b/test/body_aerodynamics/test_body_aerodynamics.jl
index aa57f537..7ecfd6f8 100644
--- a/test/body_aerodynamics/test_body_aerodynamics.jl
+++ b/test/body_aerodynamics/test_body_aerodynamics.jl
@@ -494,15 +494,15 @@ end
         @test panel.va ≈ expected_va atol=1e-12
     end
     @test body_aero.omega ≈ omega
-    @test body_aero.has_distributed_va
-    @test_throws ArgumentError body_aero.va
+    @test !body_aero.has_distributed_va
+    @test body_aero.va ≈ va
 
     new_omega = [0.0, 0.0, 2.0]
-    reference_va = copy(body_aero._va)
+    @test body_aero._va ≈ va
     body_aero.omega = new_omega
 
     for panel in body_aero.panels
-        expected_va = reference_va .+ (-new_omega × panel.control_point)
+        expected_va = va .+ (-new_omega × panel.control_point)
         @test panel.va ≈ expected_va atol=1e-12
     end
     @test body_aero.omega ≈ new_omega
diff --git a/test/body_aerodynamics/test_results.jl b/test/body_aerodynamics/test_results.jl
index e67d279d..e278a10a 100644
--- a/test/body_aerodynamics/test_results.jl
+++ b/test/body_aerodynamics/test_results.jl
@@ -36,10 +36,10 @@ end
 
 @testset "Nonlinear vs Linear - Comprehensive Input Testing" begin
     # Initialize with base parameters
-    va = [15.0, 0.0, 0.0]
-    theta = zeros(4)
-    delta = zeros(4)
-    omega = zeros(3)
+    va = [15.0, 1.0, 0.5]
+    theta = deg2rad.([2.0, 1.0, -1.0, -2.0])
+    delta = deg2rad.([1.0, 0.5, -0.5, -1.0])
+    omega = [0.0, 0.1, 0.0]
     
     # Define perturbation magnitudes
     dva_magnitudes = [0.01, 0.01, 0.01]  # Velocity perturbations (m/s)
@@ -123,7 +123,7 @@ end
                     reset_va = copy(va)
                     reset_theta = copy(theta)
                     reset_delta = copy(delta)
-                    reset_omega = zeros(3)
+                    reset_omega = copy(omega)
                     
                     # Apply the perturbation to the nonlinear model
                     if input_name == "va"
@@ -133,7 +133,7 @@ end
                     elseif input_name == "delta"
                         reset_delta = delta + perturbation[11:14]
                     elseif input_name == "omega"
-                        reset_omega = perturbation[8:10]
+                        reset_omega = omega + perturbation[8:10]
                     else
                         throw(ArgumentError())
                     end
@@ -302,8 +302,8 @@ end
                 @info "$combo_name error metrics" prediction_error baseline_difference error_ratio
                 
                 # Validate the prediction
-                @test lin_prediction ≈ nonlin_res rtol=0.01 atol=1e-3
-                @test error_ratio < 0.005
+                @test lin_prediction ≈ nonlin_res rtol=0.05 atol=1e-3
+                @test error_ratio < 0.05
             end
         end
     end