Hello!
I am working on a problem where I want to propagate a trajectory using diffrax, use a discrete terminating event to end that trajectory, and then get the derivative of that final time with respect to a variable. Regardless of the problems I have tried to do this on, the gradient of the final time always ends up as zero. Below is an example where I model the simplified dynamics of a cannonball and I try to evaluate the gradient of the final time with respect to the coefficient of drag.
import jax
import jax.numpy as jnp
import diffrax as dfx
jax.config.update('jax_enable_x64', True)
def dynamics(t, y, CD):
g = 9.80665 # m/s^2
rho = 1.225 # kg/m^3
S = 0.005 # m^2
m = 1. # kg
v = y[2]
gam = y[3]
D = (0.5*rho*v**2)*S*CD
sin_gam = jnp.sin(gam)
cos_gam = jnp.cos(gam)
ydot = jnp.array([v*cos_gam, # rdot
v*sin_gam, # hdot
-D/m - g*sin_gam, # vdot
-(g*cos_gam)/v]) # gamdot
return ydot
def event(state, **kwargs):
h = state.y[1]
return h < 0.
def obj(CD):
t0 = 0.
t1 = 20.
dt0 = None
y0 = jnp.array([0., # r0
100., # h0
5., # v0
jnp.deg2rad(45.)]) # gam0
solver = dfx.Tsit5()
stepsize_controller = dfx.ConstantStepSize()
saveat = dfx.SaveAt(ts=jnp.linspace(t0, t1, 100))
sol = dfx.diffeqsolve(dfx.ODETerm(dynamics),
solver,
t0,
t1,
0.1,
y0,
args=CD,
stepsize_controller=stepsize_controller,
max_steps=None,
saveat=saveat,
discrete_terminating_event=event,
adjoint=dfx.RecursiveCheckpointAdjoint(checkpoints=100)
)
ts = sol.ts[jnp.where(jnp.isfinite(sol.ts))]
return ts[-1]
if __name__ == '__main__':
CD = 0.5
print(jax.grad(obj)(CD))Is there a correct/better way to do this? Thanks in advance!
Hello!
I am working on a problem where I want to propagate a trajectory using
diffrax, use a discrete terminating event to end that trajectory, and then get the derivative of that final time with respect to a variable. Regardless of the problems I have tried to do this on, the gradient of the final time always ends up as zero. Below is an example where I model the simplified dynamics of a cannonball and I try to evaluate the gradient of the final time with respect to the coefficient of drag.Is there a correct/better way to do this? Thanks in advance!