diff --git a/src/physics/transport.jl b/src/physics/transport.jl
index 124cbfab..3b42c020 100644
--- a/src/physics/transport.jl
+++ b/src/physics/transport.jl
@@ -61,11 +61,11 @@ If `rho_ped > transport_grid[end]` then rotation shear is linearly interpolated
 
 If `rho_ped < transport_grid[end]` then boundary condition is at `transport_grid[end]`
 """
-function profile_from_rotation_shear_transport(
+function profile_from_grad(
     profile_old::AbstractVector{<:Real},
     rho::AbstractVector{<:Real},
     transport_grid::AbstractVector{<:Real},
-    rotation_shear_transport_grid::AbstractVector{<:Real},
+    grad_transport_grid::AbstractVector{<:Real},
     rho_ped::Real=0.0)
 
     transport_indices = [IMAS.argmin_abs(rho, rho_x) for rho_x in transport_grid]
@@ -74,17 +74,17 @@ function profile_from_rotation_shear_transport(
 
     if index_ped > index_last
         # If rho_ped is beyond transport_grid[end], extend interpolation
-        dw_dr_old = gradient(rho, profile_old; method=:backward)
+        grad_old = grad(rho, profile_old; method=:backward)
         transport_indices = vcat(1, transport_indices, index_ped)
-        rotation_shear_transport_grid = vcat(0.0, rotation_shear_transport_grid, dw_dr_old[index_ped])
+        grad_transport_grid = vcat(0.0, grad_transport_grid, grad_old[index_ped])
     else
         # If rho_ped is within transport_grid, use transport_grid[end] as boundary
         transport_indices = vcat(1, transport_indices)
-        rotation_shear_transport_grid = vcat(0.0, rotation_shear_transport_grid)
+        grad_transport_grid = vcat(0.0, grad_transport_grid)
     end
 
     # Interpolate rotation shear to the integration range
-    dw_dr = IMAS.interp1d(transport_indices, rotation_shear_transport_grid).(1:index_last)
+    grad = IMAS.interp1d(transport_indices, grad_transport_grid).(1:index_last)
 
     # Set up the profile
     profile_new = similar(profile_old)
@@ -95,9 +95,9 @@ function profile_from_rotation_shear_transport(
     profile_new[index_last] = profile_old[index_last]
     for i in (index_last-1):-1:1
         # Trapezoidal integration: negative because we integrate inward
-        dw_avg = (dw_dr[i] + dw_dr[i+1]) / 2.0
+        grad_avg = (grad[i] + grad[i+1]) / 2.0
         dr = rho[i+1] - rho[i]
-        profile_new[i] = profile_new[i+1] - dw_avg * dr
+        profile_new[i] = profile_new[i+1] - grad_avg * dr
     end
 
     return profile_new
@@ -281,4 +281,4 @@ end
 
 
 @compat public total_fluxes
-push!(document[Symbol("Physics transport")], :total_fluxes)
\ No newline at end of file
+push!(document[Symbol("Physics transport")], :total_fluxes)