#A script to illustrate Laplace transforms in Sage, and using Laplace transforms
#to solve ODEs.

#To Laplace transform an expression, e.g., t^2, declare "s" and "t" as symbolic variable and use the "laplace" command:
var('s t')
laplace(t^2,t,s)

#More generally, to compute the Laplace transform of f(t) execute
#"laplace(f(t),t,s)". The second argument "t" indicates the independent
#variable in f(t) and the third argument "s" indicates that the transform
#should have independent variable "s".

#To inverse transform an expression use "inverse_laplace()", for example,
inverse_laplace(2/s^3,s,t)

#Example: Solving a first-order ODE. Consider the ODE
#u'(t) = -2*u(t)+t with initial data u(0) = 3. To solve

#Step 1: Define u(t) as a symbolic function and define the ODE
var('s t');
u = function('u')(t);
ode = diff(u,t) == -2*u + t

#Step 2: Laplace transform both sides of the ODE (which Sage does automatically):
lapode = laplace(ode,t,s);
#Sage uses the notation "laplace(u(t),t,s)" for the Laplace transform of u(t).
#Let us replace this with the notation "U" (not necessary, just more
#aesthetically pleasing) and also substitute in the initial data u(0) = 3.
var('U');
lapode2 = lapode.subs(u(0)==3,laplace(u(t),t,s)==U);

#Sage may complain that the substitutions are deprecated and will not work
#in a future release, so be aware. You can always cut-and-paste to do the
#substitutions manually.

#Step 3: We can now solve for U, which Sage returns in a list:
Usol = solve(lapode2, U);

#Step 4: The last step is to inverse Laplace transform "Usol" to find the
#solution, which we call "usol(t)".
sol = inverse_laplace(Usol[0],s,t);
usol(t) = sol.rhs();
print("The solution is ", usol(t))

#An easy alternative to the above steps is to use the desolve_laplace() command:
usol(t) = desolve_laplace(ode,u,[0,3])

#Example: A second-order ODE. Consider the second-order ODE
ode = diff(u,t,2) + 4*diff(u,t) + 3*u == sin(t)
#with u(0) = 1 and u'(0) = 2. To solve, follows these steps:

#Step 1: Laplace transform both sided of the ODE
lapode = laplace(ode,t,s)

#Step 2: Substitute in the initial data and (optionally) use "U" to replace
#"laplace(u(t),t,s)". Again, Sage may complain.
Du = diff(u,t); #Define derivative of u as a function
lapode2 = lapode.subs(u(0)==1,Du(0)==2,laplace(u(t),t,s)==U);

#Step 3: Solve for the Laplace transform U:
Usol = solve(lapode2, U);

#Step 4: Inverse Transform "Usol" to obtain the solution.
sol = inverse_laplace(Usol[0],s,t);
usol(t) = sol.rhs();
print("The solution is ", usol(t))

#Alternatively,
usol(t) = desolve_laplace(ode, u, [0,1,2]);
print("The solution really is ", usol(t));

#The Heaviside and Dirac Delta Functions:  The Heaviside function in Sage is
heaviside(t)

#For example,
plot(heaviside(t),(t,-2,2)).show()

#HOWEVER, Sage does not (currently) have the internal ability to Laplace transform Heaviside
#or Dirac delta functions. To do this we can call an external prorgram, e.g.,
#Scientific Python. Try

laplace(heaviside(t),t,s,algorithm='sympy')

#The laplace transform returns the inverse transform as a function of s as well as conditions on s under
#which the improper integral defining the transform converges. We are interested
#in the transform (the first argument).


#The Dirac delta function is
dirac_delta(t)

#Example: The mass in a damped spring-mass system with mass 1 kg, damping
#constant 2 newtons per meter per second, and spring constant 5 newtons
#per meter is subjected to a constant force of 2 newtons for 0 < t < 4
#seconds, at which point the force drops to zero. At time t = 6 seconds
#the mass is subjected to a hammer blow with total impulse 8 newton-meters.
#Find the motion of the mass if the mass starts at rest and at equilibrium,
#and plot this motion for 0 < t < 20 seconds.

#The relevant ODE is
ode = diff(u,t,2) + 2*diff(u,t) + 5*u == 2 - 2*heaviside(t-4) + 8*dirac_delta(t-6)

#We can transform the left side using Sage's usual laplace commmand. This
#command recognizes the linearity of the ODE
lapode_left = laplace(ode.lhs(),t,s);

#Then subsitute in the initial data, and use "U" for the transform
Du = diff(u,t);
lapode_left2 = lapode_left.subs(u(0)==0,Du(0)==0,laplace(u(t),t,s)==U)

#Transform the right side of the ODE using Python:
lapode_right = laplace(ode.rhs(),t,s,algorithm='sympy')

#And solve for U (using lapode_right[0] to pick off the transform)
sol = solve(lapode_left2==lapode_right[0],U);
Usol = U.subs(sol[0]);

#Finally, inverse transform using the Scientific Python option (this can be a bit slow).
usol(t) = inverse_laplace(Usol,s,t,algorithm="sympy")

#A plot
plot(usol(t), (t,0,20)).show()