Merge pull request #3595 from aml5600/andrew/change-of-vars-units

Change of Variables with units

Author: ChrisRackauckas

src/systems/diffeqs/basic_transformations.jl

@@ -126,10 +126,23 @@ function change_independent_variable(
     # Set up intermediate and final variables for the transformation
     iv1name = nameof(iv1) # e.g. :t
     iv2name = nameof(operation(iv2_of_iv1)) # e.g. :u
-    iv2, = @independent_variables $iv2name # e.g. u
-    iv1_of_iv2, = GlobalScope.(@variables $iv1name(iv2)) # inverse, e.g. t(u), global because iv1 has no namespacing in sys
     D1 = Differential(iv1) # e.g. d/d(t)
-    div2_of_iv1 = GlobalScope(default_toterm(D1(iv2_of_iv1))) # e.g. uˍt(t)
+
+    # construct new terms, e.g:
+    #   iv2 -> u
+    #   iv1_of_iv2 -> t(u), (inverse, global because iv1 has no namespacing in sys)
+    #   div2_of_iv1 -> uˍt(t)
+    iv2_unit = getmetadata(iv2_of_iv1, VariableUnit, nothing)
+    if isnothing(iv2_unit)
+        iv2, = @independent_variables $iv2name
+        iv1_of_iv2, = GlobalScope.(@variables $iv1name(iv2))
+        div2_of_iv1 = GlobalScope(default_toterm(D1(iv2_of_iv1)))
+    else
+        iv2, = @independent_variables $iv2name [unit = iv2_unit]
+        iv1_of_iv2, = GlobalScope.(@variables $iv1name(iv2) [unit = get_unit(iv1)])
+        div2_of_iv1 = GlobalScope(diff2term_with_unit(D1(iv2_of_iv1), iv1))
+    end
+
     div2_of_iv2 = substitute(div2_of_iv1, iv1 => iv2) # e.g. uˍt(u)
     div2_of_iv2_of_iv1 = substitute(div2_of_iv2, iv2 => iv2_of_iv1) # e.g. uˍt(u(t))
 

test/basic_transformations.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, OrdinaryDiffEq, DataInterpolations, Test
+using ModelingToolkit, OrdinaryDiffEq, DataInterpolations, DynamicQuantities, Test
 
 @independent_variables t
 D = Differential(t)
@@ -215,3 +215,19 @@ end
     M = ODESystem([D(x(t)) ~ x(t - 1)], t; name = :M)
     @test_throws "DDE" change_independent_variable(M, x(t))
 end
+
+@testset "Change independent variable w/ units (free fall with 2nd order horizontal equation)" begin
+    @independent_variables t_units [unit = u"s"]
+    D_units = Differential(t_units)
+    @variables x(t_units) [unit = u"m"] y(t_units) [unit = u"m"]
+    @parameters g = 9.81 [unit = u"m * s^-2"] # gravitational acceleration
+    Mt = ODESystem([D_units(D_units(y)) ~ -g, D_units(D_units(x)) ~ 0], t_units; name = :M) # gives (x, y) as function of t, ...
+    Mx = change_independent_variable(Mt, x; add_old_diff = true) # ... but we want y as a function of x
+    Mx = structural_simplify(Mx; allow_symbolic = true)
+    Dx = Differential(Mx.x)
+    u0 = [Mx.y => 0.0, Dx(Mx.y) => 1.0, Mx.t_units => 0.0, Mx.xˍt_units => 10.0]
+    prob = ODEProblem(Mx, u0, (0.0, 20.0), []) # 1 = dy/dx = (dy/dt)/(dx/dt) means equal initial horizontal and vertical velocities
+    sol = solve(prob, Tsit5(); reltol = 1e-5)
+    # compare to analytical solution (x(t) = v*t, y(t) = v*t - g*t^2/2)
+    @test all(isapprox.(sol[Mx.y], sol[Mx.x - g * (Mx.t_units)^2 / 2]; atol = 1e-10))
+end