```python
#import autograd.numpy as np
#from autograd.scipy.integrate import odeint
#from autograd import jacobian
#from autograd.builtins import tuple
from torchdiffeq import odeint, odeint_adjoint
from torch import nn
import numpy as np
import torch as tch

#odeint(func, y0, t)
import matplotlib.pyplot as plt

#Ca0 = tch.tensor(1.0)
#k1 = k_1 = tch.tensor(3.0, dtype=t.float64)
#odeint(func, y0, t)
#tspan = tch.from_numpy(np.linspace(0, 0.5, 50))
tspan = tch.linspace(0, 0.5, 50, requires_grad=True)

def f(t,z):
    z0,z1,z2 = tch.unbind(z)
#    return tch.stack((z[0]**2 - z[1], z[1]**1, z[2]**1-z[1]) )
    return tch.stack((z0**2 - z1, z1**1, z2**1-z1) )

class Lambda(nn.Module):
    #def forward(self, t, y):
        #return torch.mm(y**3, true_A)
    def forward(self, t,z):
        return tch.sin(z)
        #z0,z1,z2 = tch.unbind(z)
        #return tch.stack((z[0]**2 - z[1], z[1]**1, z[2]**1-z[1]) )

def C(z0):
    #z0 = t.tensor(z0, dtype=t.float64, requires_grad=True )

    #sol = odeint(f, z0, tspan)

    return sol
z0 =[1.,2.,3.0]
z0 = tch.tensor(z0, dtype=t.float64 , requires_grad=True)
#ol = C( z0)
sol = odeint_adjoint(Lambda(), z0, tspan)
#sol = odeint(f, z0, tspan)

```

```python
loss = tch.sum(sol[-1])
loss.backward(tch.ones(loss.shape) )
loss.grad
```

```python
sol.backward(tch.ones(sol.shape))
sol.grad
```

```python
sol.grad.
```

```python
plt.plot(tspan.detach().numpy(), sol.detach().numpy() )
#plt.xlim([tspan.min(), tspan.max()])
plt.xlabel('t')
plt.ylabel('C');
```

```python
#sol[-1].backward(z0)

```

```python
L.backward(torch.ones(L.shape))
```

```python
odeint?
```

    Signature: odeint(func, y0, t, rtol=1e-07, atol=1e-09, method=None, options=None)
    Docstring:
    Integrate a system of ordinary differential equations.

    Solves the initial value problem for a non-stiff system of first order ODEs:
        ```
        dy/dt = func(t, y), y(t[0]) = y0
        ```
    where y is a Tensor of any shape.

    Output dtypes and numerical precision are based on the dtypes of the inputs `y0`.

    Args:
        func: Function that maps a Tensor holding the state `y` and a scalar Tensor
            `t` into a Tensor of state derivatives with respect to time.
        y0: N-D Tensor giving starting value of `y` at time point `t[0]`. May
            have any floating point or complex dtype.
        t: 1-D Tensor holding a sequence of time points for which to solve for
            `y`. The initial time point should be the first element of this sequence,
            and each time must be larger than the previous time. May have any floating
            point dtype. Converted to a Tensor with float64 dtype.
        rtol: optional float64 Tensor specifying an upper bound on relative error,
            per element of `y`.
        atol: optional float64 Tensor specifying an upper bound on absolute error,
            per element of `y`.
        method: optional string indicating the integration method to use.
        options: optional dict of configuring options for the indicated integration
            method. Can only be provided if a `method` is explicitly set.
        name: Optional name for this operation.

    Returns:
        y: Tensor, where the first dimension corresponds to different
            time points. Contains the solved value of y for each desired time point in
            `t`, with the initial value `y0` being the first element along the first
            dimension.

    Raises:
        ValueError: if an invalid `method` is provided.
        TypeError: if `options` is supplied without `method`, or if `t` or `y0` has
            an invalid dtype.
    File:      c:\users\morgan\anaconda3\envs\pymbe_py37\lib\site-packages\torchdiffeq\_impl\odeint.py
    Type:      function

```python
sol.grad_fn
```

    <torch.autograd.function.OdeintAdjointMethodBackward at 0x16553256c88>

```python
t.tensor(5)
```

    tensor(5)

```python
import time
t0 = time.time()
dCdk = jacobian(C, 0)

k_sensitivity = dCdk(np.array([k1, k_1]))

k1_sensitivity = k_sensitivity[:, 0, 0]
k_1_sensitivity = k_sensitivity[:, 0, 1]

plt.plot(tspan, np.abs(k1_sensitivity), label='dC/dk1')
plt.plot(tspan, np.abs(k_1_sensitivity), label='dC/dk_1')
plt.legend(loc='best')
plt.xlabel('t')
plt.ylabel('sensitivity')
print(f'Elapsed time = {time.time() - t0:1.1f} seconds')
```

```python
import jax
```

```python
import torch
# Creating the graph
x = torch.tensor(1.0, requires_grad = True)
z = x ** 3
z.backward() #Computes the gradient
print(x.grad.data) #Prints '3' which is dz/dx
```

    tensor(3.)

# TF
test differentiating an ODE solves using "tfdiffeq".

```python
%load_ext autoreload
%autoreload 2
import tfdiffeq as td

from tfdiffeq import odeint
#from tfdiffeq import cast_double
import tensorflow as tf
tf.keras.backend.set_floatx('float64') # Presumably can get more speed, especially on a GPU,
#if we used float32
import numpy as np
import pylab as pl
#import torch as tch

#odeint(func, y0, t)
import matplotlib.pyplot as plt

@tf.function
def f(t, z, pars): # The actual function, dz_dt
    # y now reprents the a vector of [u, v]
    #u, v = y[0], y[1]

    #self.val.assign_add( 1)
    return z**2*pars* tf.exp(-z/10)
    #z0,z1,z2 = z
    #z0_dt = z[0]**2 - z[1]
    #z1_dt = pars[0]*tf.abs(t-pars[1])**2*z[1]**1
    #z2_dt = z[2]**1-z[1]
    #du_dt = v
    #dv_dt = 5 * v - 6 * u
    #return tf.stack([z0_dt, z1_dt, z2_dt])

class Lambda(tf.keras.Model):

    def __init__(self, pars):#,f): #pars are parameters to the ode
        self.pars=pars
        #self.Nevals = tf.Variable(0)#tf.convert_to_tensor(0, dtype=tf.int64)
        #self.f = tf.function(f)
        super().__init__()

    @tf.function
    def call(self, t, z):
        z0,z1 = z[0], z[1]
        #self.Nevals.assign_add(1)
        #return self.f(t,z,self.pars)
        #self.val.assign_add( 1)
        pow_step  = tf.pow(z0,tf.constant(2, dtype='float64') )
        tf.print(pow_step)
        tf.print(t)
        exp_step = pow_step*tf.exp(-z0)*t
        z0_dt = self.pars* exp_step
        z1_dt = z1+z0
        return tf.stack([z0_dt, z1_dt])

#     def call(self, t, z): # The actual function, dz_dt
#         #return f(t,z, self.pars)
#         # y now reprents the a vector of [u, v]
#         #u, v = y[0], y[1]
#         #z0,z1,z2 = z
#         z0_dt = z[0]**2 - z[1]
#         z1_dt = self.pars[0]*tf.abs(t-self.pars[1])**2*z[1]**1
#         z2_dt = z[2]**1-z[1]
#         #du_dt = v
#         #dv_dt = 5 * v - 6 * u
#         self.val.assign_add( 1)
#         return tf.stack([z0_dt, z1_dt, z2_dt])

    #return tf.stack([du_dt, dv_dt])  # vector of shape [2]

```

    The autoreload extension is already loaded. To reload it, use:
      %reload_ext autoreload

```python
with tf.device("/cpu:0"):
    tspan = tf.linspace(0.0, 5.0, 400)
    z0 = np.linspace(0,2,2000).reshape(2,1000) # initial values
    z0 = tf.convert_to_tensor(z0, dtype=tf.float64)
    p0 = tf.linspace(-4, 4,1000) #parameters for the ODE
    p0 = tf.convert_to_tensor(p0, dtype=tf.float64)
    ode = Lambda(p0)
    #tspan =td.cast_double(tspan)
    #ol = C( z0)
    %time sol = odeint(ode, z0, tspan, method='rk4')
    #sol = odeint(f, z0, tspan)
    #L = tf.reduce_sum(sol[-1])
```

```python
tspan = tf.linspace(0.0, 5.0, 400)
z0 = np.linspace(0,2,2000).reshape(2,1000) # initial values
z0 = tf.convert_to_tensor(z0, dtype=tf.float64)
p0 = tf.linspace(-4, 4,1000) #parameters for the ODE
p0 = tf.convert_to_tensor(p0, dtype=tf.float64)
ode = Lambda(p0)
#tspan =td.cast_double(tspan)
#ol = C( z0)
def run(z0, tspan):
    sol = odeint(ode, z0, tspan, method='rk4')
    return sol
run2=tf.function(run)
```

```python
tf.function(ode)(0.0, z0)
```

    WARNING:tensorflow:Layer lambda_15 is casting an input tensor from dtype float32 to the layer's dtype of float64, which is new behavior in TensorFlow 2.  The layer has dtype float64 because its dtype defaults to floatx.

    If you intended to run this layer in float64, you can safely ignore this warning. If in doubt, this warning is likely only an issue if you are porting a TensorFlow 1.X model to TensorFlow 2.

    To change all layers to have dtype float32 by default, call `tf.keras.backend.set_floatx('float32')`. To change just this layer, pass dtype='float32' to the layer constructor. If you are the author of this layer, you can disable autocasting by passing autocast=False to the base Layer constructor.

    [0 1.0010007505003125e-06 4.00400300200125e-06 ... 0.99500375500406524 0.99700075150131318 0.99899975000006236]
    0

    <tf.Tensor: shape=(2, 1000), dtype=float64, numpy=
    array([[0.        , 0.        , 0.        , ..., 0.        , 0.        ,
            0.        ],
           [1.00050025, 1.00250125, 1.00450225, ..., 2.99549775, 2.99749875,
            2.99949975]])>

```python
run2(z0, tspan)
```

```python

```

```python

with tf.device("/cpu:0"):
    with tf.GradientTape() as g:
        g.watch(p0)
        ode = Lambda(p0)

        #g.watch(z0)

        #tspan = tf.linspace(0.0, 2, 1000)
        #tspan =td.cast_double(tspan)
        #ol = C( z0)
        %time sol = odeint(ode, z0, tspan, method='adams')
        #sol = odeint(f, z0, tspan)
        L = tf.reduce_sum(sol[-1])

```

```python
ode.val
```

```python
L
```

    <tf.Tensor: shape=(), dtype=float64, numpy=37932.76585269482>

```python
%time dsol_dz = g.gradient(L, p0)
```

    CPU times: user 2.67 s, sys: 64.8 ms, total: 2.73 s
    Wall time: 2.51 s

```python
pl.plot(sol[:,::30])
```

```python
pl.plot(sol[:,::30])
```

```python
#pl.plot(sol)
```

```python
dsol_dz
```

    <tf.Tensor: shape=(1000,), dtype=float64, numpy=
    array([ 0.00000000e+00,  7.34075084e+03,  8.58407524e+04,  3.30432847e+05,
            8.18119238e+05,  1.60137469e+06,  2.71170689e+06,  4.16399596e+06,
            5.95996869e+06,  8.09075296e+06,  1.05387291e+07,  1.32788861e+07,
            1.62798363e+07,  1.95045950e+07,  2.29111975e+07,  2.64532001e+07,
            3.00801001e+07,  3.37376967e+07,  3.73684074e+07,  4.09115525e+07,
            4.43036147e+07,  4.74784791e+07,  5.03676588e+07,  5.29005076e+07,
            5.50044235e+07,  5.66050436e+07,  5.76264322e+07,  5.79912625e+07,
            5.76209939e+07,  5.64360431e+07,  5.43559521e+07,  5.12995517e+07,
            4.71851212e+07,  4.19305447e+07,  3.54534637e+07,  2.76714270e+07,
            1.85020363e+07,  7.86309018e+06, -4.32727653e+06, -1.81504553e+07,
           -3.36872300e+07, -5.10176461e+07, -7.02208823e+07, -9.13751251e+07,
           -1.14557448e+08, -1.39843691e+08, -1.67308346e+08, -1.97024440e+08,
           -2.29063429e+08, -2.63495089e+08, -3.00387412e+08, -3.39806502e+08,
           -3.81816483e+08, -4.26479393e+08, -4.73855104e+08, -5.24001223e+08,
           -5.76973011e+08, -6.32823293e+08, -6.91602388e+08, -7.53358021e+08,
           -8.18135253e+08, -8.85976413e+08, -9.56921019e+08, -1.03100572e+09,
           -1.10826425e+09, -1.18872731e+09, -1.27242261e+09, -1.35937469e+09,
           -1.44960500e+09, -1.54313174e+09, -1.63996989e+09, -1.74013114e+09,
           -1.84362384e+09, -1.95045298e+09, -2.06062018e+09, -2.17412359e+09,
           -2.29095793e+09, -2.41111445e+09, -2.53458086e+09, -2.66134139e+09,
           -2.79137673e+09, -2.92466399e+09, -3.06117676e+09, -3.20088502e+09,
           -3.34375523e+09, -3.48975020e+09, -3.63882925e+09, -3.79094806e+09,
           -3.94605878e+09, -4.10410999e+09, -4.26504672e+09, -4.42881046e+09,
           -4.59533920e+09, -4.76456740e+09, -4.93642608e+09, -5.11084278e+09,
           -5.28774158e+09, -5.46704326e+09, -5.64866509e+09, -5.83252110e+09,
           -6.01852196e+09, -6.20657519e+09, -6.39658486e+09, -6.58845207e+09,
           -6.78207466e+09, -6.97734746e+09, -7.17416216e+09, -7.37240755e+09,
           -7.57196933e+09, -7.77273055e+09, -7.97457111e+09, -8.17736851e+09,
           -8.38099712e+09, -8.58532908e+09, -8.79023359e+09, -8.99557754e+09,
           -9.20122521e+09, -9.40703869e+09, -9.61287755e+09, -9.81859925e+09,
           -1.00240589e+10, -1.02291099e+10, -1.04336032e+10, -1.06373881e+10,
           -1.08403118e+10, -1.10422196e+10, -1.12429559e+10, -1.14423627e+10,
           -1.16402804e+10, -1.18365488e+10, -1.20310053e+10, -1.22234864e+10,
           -1.24138283e+10, -1.26018647e+10, -1.27874290e+10, -1.29703535e+10,
           -1.31504703e+10, -1.33276098e+10, -1.35016024e+10, -1.36722779e+10,
           -1.38394659e+10, -1.40029944e+10, -1.41626926e+10, -1.43183892e+10,
           -1.44699121e+10, -1.46170905e+10, -1.47597526e+10, -1.48977269e+10,
           -1.50308433e+10, -1.51589313e+10, -1.52818208e+10, -1.53993424e+10,
           -1.55113283e+10, -1.56176104e+10, -1.57180220e+10, -1.58123976e+10,
           -1.59005725e+10, -1.59823829e+10, -1.60576672e+10, -1.61262655e+10,
           -1.61880169e+10, -1.62427659e+10, -1.62903554e+10, -1.63306318e+10,
           -1.63634431e+10, -1.63886395e+10, -1.64060719e+10, -1.64155962e+10,
           -1.64170674e+10, -1.64103453e+10, -1.63952900e+10, -1.63717668e+10,
           -1.63396409e+10, -1.62987815e+10, -1.62490627e+10, -1.61903575e+10,
           -1.61225435e+10, -1.60455027e+10, -1.59591194e+10, -1.58632802e+10,
           -1.57578756e+10, -1.56428008e+10, -1.55179534e+10, -1.53832334e+10,
           -1.52385450e+10, -1.50837990e+10, -1.49189059e+10, -1.47437822e+10,
           -1.45583471e+10, -1.43625274e+10, -1.41562476e+10, -1.39394415e+10,
           -1.37120451e+10, -1.34739996e+10, -1.32252490e+10, -1.29657425e+10,
           -1.26954335e+10, -1.24142808e+10, -1.21222477e+10, -1.18193002e+10,
           -1.15054104e+10, -1.11805539e+10, -1.08447131e+10, -1.04978721e+10,
           -1.01400239e+10, -9.77116237e+09, -9.39128978e+09, -9.00040890e+09,
           -8.59853162e+09, -8.18567206e+09, -7.76184998e+09, -7.32709317e+09,
           -6.88142806e+09, -6.42489438e+09, -5.95752800e+09, -5.47937563e+09,
           -4.99048906e+09, -4.49092502e+09, -3.98073973e+09, -3.46000171e+09,
           -2.92878328e+09, -2.38715986e+09, -1.83520992e+09, -1.27302590e+09,
           -7.00691744e+08, -1.18308544e+08,  4.74020640e+08,  1.07618958e+09,
            1.68809238e+09,  2.30960197e+09,  2.94060370e+09,  3.58096690e+09,
            4.23056235e+09,  4.88924880e+09,  5.55688854e+09,  6.23333027e+09,
            6.91842619e+09,  7.61201642e+09,  8.31394035e+09,  9.02403270e+09,
            9.74211805e+09,  1.04680279e+10,  1.12015764e+10,  1.19425796e+10,
            1.26908499e+10,  1.34461896e+10,  1.42084031e+10,  1.49772856e+10,
            1.57526335e+10,  1.65342302e+10,  1.73218677e+10,  1.81153143e+10,
            1.89143557e+10,  1.97187674e+10,  2.05283100e+10,  2.13427449e+10,
            2.21618408e+10,  2.29853584e+10,  2.38130405e+10,  2.46446441e+10,
            2.54799137e+10,  2.63185917e+10,  2.71604197e+10,  2.80051337e+10,
            2.88524768e+10,  2.97021665e+10,  3.05539375e+10,  3.14075117e+10,
            3.22626187e+10,  3.31189732e+10,  3.39762950e+10,  3.48343036e+10,
            3.56927035e+10,  3.65512131e+10,  3.74095387e+10,  3.82673917e+10,
            3.91244802e+10,  3.99804910e+10,  4.08351456e+10,  4.16881398e+10,
            4.25391698e+10,  4.33879378e+10,  4.42341420e+10,  4.50774774e+10,
            4.59176366e+10,  4.67543193e+10,  4.75872285e+10,  4.84160388e+10,
            4.92404634e+10,  5.00601886e+10,  5.08749053e+10,  5.16843120e+10,
            5.24881038e+10,  5.32859752e+10,  5.40776218e+10,  5.48627344e+10,
            5.56410168e+10,  5.64121672e+10,  5.71758780e+10,  5.79318557e+10,
            5.86797969e+10,  5.94194113e+10,  6.01503964e+10,  6.08724650e+10,
            6.15853236e+10,  6.22886821e+10,  6.29822479e+10,  6.36657514e+10,
            6.43388927e+10,  6.50014067e+10,  6.56530085e+10,  6.62934206e+10,
            6.69223839e+10,  6.75396191e+10,  6.81448688e+10,  6.87378711e+10,
            6.93183661e+10,  6.98861068e+10,  7.04408336e+10,  7.09823091e+10,
            7.15102853e+10,  7.20245341e+10,  7.25248144e+10,  7.30109032e+10,
            7.34825751e+10,  7.39396127e+10,  7.43818037e+10,  7.48089364e+10,
            7.52208064e+10,  7.56172224e+10,  7.59979884e+10,  7.63629150e+10,
            7.67118293e+10,  7.70445464e+10,  7.73609035e+10,  7.76607351e+10,
            7.79438881e+10,  7.82102088e+10,  7.84595526e+10,  7.86917841e+10,
            7.89067739e+10,  7.91043942e+10,  7.92845331e+10,  7.94470763e+10,
            7.95919198e+10,  7.97189764e+10,  7.98281460e+10,  7.99193580e+10,
            7.99925335e+10,  8.00476055e+10,  8.00845181e+10,  8.01032195e+10,
            8.01036683e+10,  8.00858323e+10,  8.00496825e+10,  7.99951934e+10,
            7.99223631e+10,  7.98311913e+10,  7.97216742e+10,  7.95938366e+10,
            7.94476904e+10,  7.92832755e+10,  7.91006263e+10,  7.88997937e+10,
            7.86808314e+10,  7.84438120e+10,  7.81888012e+10,  7.79158841e+10,
            7.76251529e+10,  7.73167103e+10,  7.69906611e+10,  7.66471282e+10,
            7.62862317e+10,  7.59081114e+10,  7.55129088e+10,  7.51007815e+10,
            7.46718854e+10,  7.42263982e+10,  7.37644922e+10,  7.32863623e+10,
            7.27922017e+10,  7.22822154e+10,  7.17566199e+10,  7.12156389e+10,
            7.06595044e+10,  7.00884539e+10,  6.95027385e+10,  6.89026142e+10,
            6.82883497e+10,  6.76602171e+10,  6.70184974e+10,  6.63634835e+10,
            6.56954717e+10,  6.50147719e+10,  6.43216967e+10,  6.36165699e+10,
            6.28997191e+10,  6.21714835e+10,  6.14322089e+10,  6.06822460e+10,
            5.99219552e+10,  5.91517051e+10,  5.83718657e+10,  5.75828204e+10,
            5.67849548e+10,  5.59786628e+10,  5.51643431e+10,  5.43424030e+10,
            5.35132528e+10,  5.26773090e+10,  5.18349970e+10,  5.09867434e+10,
            5.01329812e+10,  4.92741471e+10,  4.84106856e+10,  4.75430417e+10,
            4.66716678e+10,  4.57970179e+10,  4.49195511e+10,  4.40397276e+10,
            4.31580114e+10,  4.22748723e+10,  4.13907789e+10,  4.05062015e+10,
            3.96216159e+10,  3.87374956e+10,  3.78543172e+10,  3.69725579e+10,
            3.60926947e+10,  3.52152052e+10,  3.43405658e+10,  3.34692534e+10,
            3.26017431e+10,  3.17385087e+10,  3.08800214e+10,  3.00267506e+10,
            2.91791638e+10,  2.83377234e+10,  2.75028891e+10,  2.66751169e+10,
            2.58548572e+10,  2.50425555e+10,  2.42386521e+10,  2.34435809e+10,
            2.26577686e+10,  2.18816357e+10,  2.11155940e+10,  2.03600472e+10,
            1.96153904e+10,  1.88820090e+10,  1.81602780e+10,  1.74505622e+10,
            1.67532146e+10,  1.60685766e+10,  1.53969767e+10,  1.47387301e+10,
            1.40941383e+10,  1.34634875e+10,  1.28470493e+10,  1.22450785e+10,
            1.16578136e+10,  1.10854749e+10,  1.05282649e+10,  9.98636666e+09,
            9.45994328e+09,  8.94913705e+09,  8.45406877e+09,  7.97483671e+09,
            7.51151584e+09,  7.06415701e+09,  6.63278578e+09,  6.21740196e+09,
            5.81797816e+09,  5.43445926e+09,  5.06676125e+09,  4.71477024e+09,
            4.37834152e+09,  4.05729844e+09,  3.75143156e+09,  3.46049731e+09,
            3.18421718e+09,  2.92227642e+09,  2.67432300e+09,  2.43996648e+09,
            2.21877685e+09,  2.01028327e+09,  1.81397294e+09,  1.62928984e+09,
            1.45563349e+09,  1.29235767e+09,  1.13876911e+09,  9.94126190e+08,
            8.57637606e+08,  7.28460965e+08,  6.05701421e+08,  4.88410247e+08,
            3.75583394e+08,  2.66160023e+08,  1.59021008e+08,  5.29874172e+07,
           -5.31810401e+07, -1.60787585e+08, -2.71199977e+08, -3.85852144e+08,
           -5.06245845e+08, -6.33952362e+08, -7.70614221e+08, -9.17946940e+08,
           -1.07774082e+09, -1.25186274e+09, -1.44225804e+09, -1.65095236e+09,
           -1.88005356e+09, -2.13175370e+09, -2.40833095e+09, -2.71215172e+09,
           -3.04567258e+09, -3.41144241e+09, -3.81210454e+09, -4.25039888e+09,
           -4.72916416e+09, -5.25134011e+09, -5.81996978e+09, -6.43820178e+09,
           -7.10929287e+09, -7.83660996e+09, -8.62363308e+09, -9.47395736e+09,
           -1.03912959e+10, -1.13794822e+10, -1.24424725e+10, -1.35843495e+10,
           -1.48093235e+10, -1.61217366e+10, -1.75260651e+10, -1.90269211e+10,
           -2.06290577e+10, -2.23373697e+10, -2.41568997e+10, -2.60928356e+10,
           -2.81505215e+10, -3.03354536e+10, -3.26532879e+10, -3.51098422e+10,
           -3.77110989e+10, -4.04632092e+10, -4.33724960e+10, -4.64454598e+10,
           -4.96887767e+10, -5.31093090e+10, -5.67141038e+10, -6.05103988e+10,
           -6.45056279e+10, -6.87074225e+10, -7.31236138e+10, -7.77622448e+10,
           -8.26315627e+10, -8.77400371e+10, -9.30963473e+10, -9.87094085e+10,
           -1.04588355e+11, -1.10742562e+11, -1.17181628e+11, -1.23915408e+11,
           -1.30953996e+11, -1.38307754e+11, -1.45987268e+11, -1.54003428e+11,
           -1.62367360e+11, -1.71090487e+11, -1.80184476e+11, -1.89661323e+11,
           -1.99533277e+11, -2.09812897e+11, -2.20513052e+11, -2.31646907e+11,
           -2.43227944e+11, -2.55269957e+11, -2.67787076e+11, -2.80793764e+11,
           -2.94304836e+11, -3.08335420e+11, -3.22901016e+11, -3.38017512e+11,
           -3.53701096e+11, -3.69968403e+11, -3.86836409e+11, -4.04322481e+11,
           -4.22444408e+11, -4.41220337e+11, -4.60668908e+11, -4.80809121e+11,
           -5.01660407e+11, -5.23242684e+11, -5.45576246e+11, -5.68681901e+11,
           -5.92580935e+11, -6.17295082e+11, -6.42846535e+11, -6.69257999e+11,
           -6.96552747e+11, -7.24754489e+11, -7.53887508e+11, -7.83976568e+11,
           -8.15047054e+11, -8.47124884e+11, -8.80236515e+11, -9.14409010e+11,
           -9.49670055e+11, -9.86047918e+11, -1.02357145e+12, -1.06227019e+12,
           -1.10217428e+12, -1.14331459e+12, -1.18572259e+12, -1.22943038e+12,
           -1.27447096e+12, -1.32087787e+12, -1.36868548e+12, -1.41792877e+12,
           -1.46864373e+12, -1.52086674e+12, -1.57463546e+12, -1.62998795e+12,
           -1.68696315e+12, -1.74560117e+12, -1.80594272e+12, -1.86802932e+12,
           -1.93190359e+12, -1.99760894e+12, -2.06518979e+12, -2.13469153e+12,
           -2.20616042e+12, -2.27964383e+12, -2.35519017e+12, -2.43284878e+12,
           -2.51267016e+12, -2.59470577e+12, -2.67900867e+12, -2.76563220e+12,
           -2.85463169e+12, -2.94606322e+12, -3.03998412e+12, -3.13645324e+12,
           -3.23553044e+12, -3.33727674e+12, -3.44175493e+12, -3.54902873e+12,
           -3.65916382e+12, -3.77222650e+12, -3.88828559e+12, -4.00740999e+12,
           -4.12967146e+12, -4.25514274e+12, -4.38389782e+12, -4.51601342e+12,
           -4.65156624e+12, -4.79063644e+12, -4.93330484e+12, -5.07965427e+12,
           -5.22976942e+12, -5.38373759e+12, -5.54164623e+12, -5.70358660e+12,
           -5.86965084e+12, -6.03993392e+12, -6.21453214e+12, -6.39354439e+12,
           -6.57707157e+12, -6.76521722e+12, -6.95808711e+12, -7.15578891e+12,
           -7.35843271e+12, -7.56613208e+12, -7.77900235e+12, -7.99716121e+12,
           -8.22072963e+12, -8.44983095e+12, -8.68459141e+12, -8.92514074e+12,
           -9.17161013e+12, -9.42413522e+12, -9.68285371e+12, -9.94790692e+12,
           -1.02194394e+13, -1.04975988e+13, -1.07825366e+13, -1.10744062e+13,
           -1.13733672e+13, -1.16795824e+13, -1.19932144e+13, -1.23144370e+13,
           -1.26434200e+13, -1.29803440e+13, -1.33253899e+13, -1.36787442e+13,
           -1.40405985e+13, -1.44111477e+13, -1.47905920e+13, -1.51791373e+13,
           -1.55769911e+13, -1.59843700e+13, -1.64014913e+13, -1.68285812e+13,
           -1.72658662e+13, -1.77135841e+13, -1.81719719e+13, -1.86412772e+13,
           -1.91217488e+13, -1.96136464e+13, -2.01172338e+13, -2.06327803e+13,
           -2.11605655e+13, -2.17008722e+13, -2.22539937e+13, -2.28202315e+13,
           -2.33998927e+13, -2.39932932e+13, -2.46007530e+13, -2.52226061e+13,
           -2.58591895e+13, -2.65108509e+13, -2.71779398e+13, -2.78608223e+13,
           -2.85598666e+13, -2.92754518e+13, -3.00079655e+13, -3.07578075e+13,
           -3.15253833e+13, -3.23111154e+13, -3.31154311e+13, -3.39387767e+13,
           -3.47816032e+13, -3.56443800e+13, -3.65275838e+13, -3.74317096e+13,
           -3.83572589e+13, -3.93047514e+13, -4.02747102e+13, -4.12676855e+13,
           -4.22842249e+13, -4.33248972e+13, -4.43902909e+13, -4.54809980e+13,
           -4.65976277e+13, -4.77408029e+13, -4.89111713e+13, -5.01093846e+13,
           -5.13361165e+13, -5.25920615e+13, -5.38779261e+13, -5.51944353e+13,
           -5.65423358e+13, -5.79224009e+13, -5.93354061e+13, -6.07821629e+13,
           -6.22634994e+13, -6.37802714e+13, -6.53333422e+13, -6.69236255e+13,
           -6.85520296e+13, -7.02195140e+13, -7.19270379e+13, -7.36756108e+13,
           -7.54662583e+13, -7.73000297e+13, -7.91780027e+13, -8.11012895e+13,
           -8.30710217e+13, -8.50883506e+13, -8.71544885e+13, -8.92706361e+13,
           -9.14380432e+13, -9.36580081e+13, -9.59318422e+13, -9.82608768e+13,
           -1.00646496e+14, -1.03090135e+14, -1.05593238e+14, -1.08157306e+14,
           -1.10783868e+14, -1.13474503e+14, -1.16230839e+14, -1.19054548e+14,
           -1.21947324e+14, -1.24910948e+14, -1.27947209e+14, -1.31057979e+14,
           -1.34245151e+14, -1.37510686e+14, -1.40856594e+14, -1.44284917e+14,
           -1.47797788e+14, -1.51397349e+14, -1.55085840e+14, -1.58865510e+14,
           -1.62738715e+14, -1.66707829e+14, -1.70775322e+14, -1.74943694e+14,
           -1.79215549e+14, -1.83593504e+14, -1.88080312e+14, -1.92678740e+14,
           -1.97391694e+14, -2.02222096e+14, -2.07172974e+14, -2.12247448e+14,
           -2.17448714e+14, -2.22780062e+14, -2.28244876e+14, -2.33846609e+14,
           -2.39588822e+14, -2.45475206e+14, -2.51509495e+14, -2.57695578e+14,
           -2.64037426e+14, -2.70539112e+14, -2.77204833e+14, -2.84038888e+14,
           -2.91045696e+14, -2.98229767e+14, -3.05595724e+14, -3.13148335e+14,
           -3.20892414e+14, -3.28832909e+14, -3.36974848e+14, -3.45323367e+14,
           -3.53883708e+14, -3.62661182e+14, -3.71661170e+14, -3.80889160e+14,
           -3.90350745e+14, -4.00051600e+14, -4.09997537e+14, -4.20194439e+14,
           -4.30648372e+14, -4.41365511e+14, -4.52352313e+14, -4.63615329e+14,
           -4.75161385e+14, -4.86997579e+14, -4.99131305e+14, -5.11570220e+14,
           -5.24322326e+14, -5.37395989e+14, -5.50799948e+14, -5.64543237e+14,
           -5.78635381e+14, -5.93086117e+14, -6.07905572e+14, -6.23104284e+14,
           -6.38693015e+14, -6.54682761e+14, -6.71084817e+14, -6.87910607e+14,
           -7.05171852e+14, -7.22880258e+14, -7.41047761e+14, -7.59686315e+14,
           -7.78807974e+14, -7.98424833e+14, -8.18549183e+14, -8.39193182e+14,
           -8.60369261e+14, -8.82089959e+14, -9.04367868e+14, -9.27215885e+14,
           -9.50647257e+14, -9.74675209e+14, -9.99313633e+14, -1.02457658e+15,
           -1.05047871e+15, -1.07703502e+15, -1.10426108e+15, -1.13217295e+15,
           -1.16078741e+15, -1.19012155e+15, -1.22019336e+15, -1.25102112e+15,
           -1.28262406e+15, -1.31502180e+15, -1.34823464e+15, -1.38228348e+15,
           -1.41718987e+15, -1.45297593e+15, -1.48966448e+15, -1.52727873e+15,
           -1.56584269e+15, -1.60538086e+15, -1.64591841e+15, -1.68748116e+15,
           -1.73009540e+15, -1.77378808e+15, -1.81858692e+15, -1.86452000e+15,
           -1.91161624e+15, -1.95990488e+15, -2.00941613e+15, -2.06018045e+15,
           -2.11222918e+15, -2.16559393e+15, -2.22030714e+15, -2.27640167e+15,
           -2.33391105e+15, -2.39286917e+15, -2.45331063e+15, -2.51527051e+15,
           -2.57878443e+15, -2.64388868e+15, -2.71062014e+15, -2.77901622e+15,
           -2.84911519e+15, -2.92095583e+15, -2.99457799e+15, -3.07002211e+15,
           -3.14732971e+15, -3.22654330e+15, -3.30770632e+15, -3.39086348e+15,
           -3.47606066e+15, -3.56334502e+15, -3.65276504e+15, -3.74437084e+15,
           -3.83821352e+15, -3.93434618e+15, -4.03282334e+15, -4.13370101e+15,
           -4.23703675e+15, -4.34288995e+15, -4.45132162e+15, -4.56239412e+15,
           -4.67617178e+15, -4.79272035e+15, -4.91210701e+15, -5.03440105e+15,
           -5.15967247e+15, -5.28799358e+15, -5.41943773e+15, -5.55408006e+15,
           -5.69199689e+15, -5.83326616e+15, -5.97796735e+15, -6.12618132e+15,
           -6.27799025e+15, -6.43347816e+15, -6.59273015e+15, -6.75583291e+15,
           -6.92287508e+15, -7.09394607e+15, -7.26913780e+15, -7.44854298e+15,
           -7.63225638e+15, -7.82037441e+15, -8.01299527e+15, -8.21021861e+15,
           -8.41214644e+15, -8.61888193e+15, -8.83053079e+15, -9.04720022e+15,
           -9.26899960e+15, -9.49604039e+15, -9.72843624e+15, -9.96630293e+15,
           -1.02097581e+16, -1.04589223e+16, -1.07139181e+16, -1.09748705e+16,
           -1.12419069e+16, -1.15151575e+16, -1.17947549e+16, -1.20808344e+16,
           -1.23735341e+16, -1.26729944e+16, -1.29793590e+16, -1.32927734e+16,
           -1.36133873e+16, -1.39413519e+16, -1.42768214e+16, -1.46199527e+16,
           -1.49709057e+16, -1.53298424e+16, -1.56969275e+16, -1.60723279e+16,
           -1.64562137e+16, -1.68487566e+16, -1.72501310e+16, -1.76605137e+16,
           -1.80800838e+16, -1.85090231e+16, -1.89475158e+16, -1.93957492e+16,
           -1.98539125e+16, -2.03221996e+16, -2.08008074e+16, -2.12899362e+16,
           -2.17897919e+16, -2.23005837e+16, -2.28225282e+16, -4.40866780e+15])>

```python
sol[-1]
```

    <tf.Tensor: shape=(3,), dtype=float64, numpy=array([0.21025278, 2.09066148, 3.77962996])>

```python
x = tf.constant(3.0)
with tf.GradientTape() as g:
  g.watch(x)
  y = x * x
dy_dx = g.gradient(y, x) # Will compute to 6.0
```

```python
sol
```

    <tf.Tensor: shape=(10, 3), dtype=float64, numpy=
    array([[1.        , 2.        , 3.1       ],
           [0.94080566, 2.01698345, 3.1623503 ],
           [0.87416031, 2.03199163, 3.22734861],
           [0.79990744, 2.04511868, 3.29525646],
           [0.71805308, 2.05646788, 3.36634463],
           [0.62880523, 2.06615045, 3.44089327],
           [0.53259905, 2.07428458, 3.51919197],
           [0.43011572, 2.0809944 , 3.60153989],
           [0.32228107, 2.08640896, 3.68824615],
           [0.21025278, 2.09066148, 3.77962996]])>

```python
@tf.function
def f(a):
    return tf.constant(1j)*a

def tf_realify(tf_func):
    tf.print(tf_func.function_spec.input_signature)
    @tf.function
    def f_wrapped(a,b):
        return tf_func(tf.cast(a, tf.complex128)+1j*tf.cast(b, tf.complex128) )
    return f_wrapped

f2=tf_realify(f)
f2(3,1)
```

    None

    <tf.Tensor: shape=(), dtype=complex128, numpy=(-1+3j)>

```python
a, b = tf.constant(3, dtype=tf.float64), tf.constant(1, dtype=tf.float64)
%timeit f2(a,b)
```

    The slowest run took 4.76 times longer than the fastest. This could mean that an intermediate result is being cached.
    10000 loops, best of 3: 146 µs per loop

```python
c = tf.constant(1+1j, dtype=tf.complex128)
%timeit f(c)
```

    The slowest run took 4.39 times longer than the fastest. This could mean that an intermediate result is being cached.
    10000 loops, best of 3: 134 µs per loop

```python
tf.constant(1j)
```

    <tf.Tensor: shape=(), dtype=complex128, numpy=1j>

```python
cf=f.get_concrete_function(0.5,z0, p0)
```

```python
cf.
```

    [<tf.Tensor 'z:0' shape=(1000,) dtype=float64>,
     <tf.Tensor 'pars:0' shape=(1000,) dtype=float64>]

```python
f.function_spec.input_signature
```

```python
cf.structured_input_signature
```

    ((0.5,
      TensorSpec(shape=(1000,), dtype=tf.float64, name='z'),
      TensorSpec(shape=(1000,), dtype=tf.float64, name='pars')),
     {})

```python
@tf.function
def f1(a,b,c):
    return a+b+c
@tf.function
def f2(a,b,c):
    return a+b
@tf.function
def f3(a,b,c):
    return a

fL = [f1,f2,f3]
```

```python
def get_wrapped(fL):
    @tf.function(experimental_compile=True)
    def f_c(p):
        return tf.stack([f(a,b,c) for f in fL])
    return f_c

f_c = get_wrapped(fL)
f_c(1.5,2.5,3.5)
```

    <tf.Tensor: shape=(3,), dtype=float32, numpy=array([7.5, 4. , 1.5], dtype=float32)>

```python
%timeit f_c(1.5,2.5,3.5)
```

    The slowest run took 6.29 times longer than the fastest. This could mean that an intermediate result is being cached.
    10000 loops, best of 3: 159 µs per loop

```python
@tf.function
def f(a,b,c):
    return (a+b+c)

@tf.function
def tf_unpack(vals):
    #for k in range(100):
    #    f(*tf.unstack(vals))
    return f(*tf.unstack(vals))

```

```python
vals = tf.convert_to_tensor(np.arange(3)*0.5)

```

    The slowest run took 477.05 times longer than the fastest. This could mean that an intermediate result is being cached.
    10000 loops, best of 3: 139 µs per loop

```python
%timeit tf_unpack(vals)
```

    The slowest run took 6.06 times longer than the fastest. This could mean that an intermediate result is being cached.
    10000 loops, best of 3: 138 µs per loop

```python
fp = lambda vals:tf.math.real(tf_unpack(vals))
fp = tf.function(fp)
```

```python
%timeit fp(vals)
```

    The slowest run took 8.05 times longer than the fastest. This could mean that an intermediate result is being cached.
    10000 loops, best of 3: 138 µs per loop

```python
T=tf.convert_to_tensor(1 + 0.5j)
```

```python
tf.math.real(T)
```

    <tf.Tensor: shape=(), dtype=float64, numpy=1.0>

# Testing pyIVP and tensorflow

```python
%load_ext autoreload
%autoreload 2
import tensorflow as tf
import sympy as sm
from spylind import spyIVP as so
from pylab import *
xS, yS, vxS, vyS = sm.symbols("x, y, v_x, v_y", real=True)
g= -10.0
dy_dtD = {xS:vxS,
       yS:vyS,
       vxS:0,
       vyS:g }
ode_s = so.ODESolver(dy_dtD,  dims={}, backend='tensorflow', default_dtype=np.float64)
ode_s.set_initial_conditions({xS:0,yS:0,vxS:10,vyS:10})
ode_s.setup()
tSteps = np.linspace(0,1,10000)
%time arr =ode_s.integrate(tSteps, 1.0)

sm.ode
plot(tSteps, arr.real)
xlabel('t')
legend(dy_dtD.keys())
```

```python
import tensorflow as tf
xS, yS, vxS, vyS = sm.symbols("x, y, v_x, v_y", real=True)
zS = sm.symbols("z", real=True)
g= tf.constant(-10.0, dtype='float64')
dy_dtD = {xS:vxS,
       yS:vyS ,
       vxS:0 +zS,
       vyS: g  }

ode_s = so.ODESolver(dy_dtD,  dims={zS:linspace(-2,2,10000)}, backend='numpy', default_dtype='float64')
ode_s.set_initial_conditions({xS:0,yS:0,vxS:10,vyS:10})
ode_s.setup(bForceStateDimensions=True)
tSteps = linspace(0,200,4)
%time arr =ode_s.integrate(tSteps, 1.0)
if 0:
    plot(tSteps, arr[:,:,0].real)
    xlabel('t')
    legend(dy_dtD.keys())
    figure()
    plot(tSteps, arr[:,:,-1].real)
    xlabel('t')
    legend(dy_dtD.keys())
    figure()
    imshow(arr[:,1,:].real, aspect='auto')
```

    Munch({'tSym': t, 'dimSyms': [z], 'prop_state_syms': [x, y, v_x, v_y], 'stationary_state_syms': [], 'driving_syms': [], 'state_dep_syms': [z]})
    state dependent functions should have signature [t, z, x, y, v_x, v_y]
    Not integrating first step (it's just the initial state)
    CPU times: user 242 ms, sys: 149 ms, total: 391 ms
    Wall time: 259 ms

```python
plot(arr[:,:,5])
plot(arr[:,:,9])
```

```python
xRef =tf.linspace(0,1,1000)
yRef = tf.sin(5*pi*xRef)
```

```python
def tf_interpolator(xOrig, yOrig):
    xOrig = tf.convert_to_tensor(xOrig, dtype=tf.float64)
    yOrig = tf.convert_to_tensor(yOrig, dtype=tf.float64)
    N = xOrig.shape[0]
    dx = xOrig[1]-xOrig[0]
    x0 = xOrig[0]
    xMax = xOrig[-1]
    zero = tf.constant(0, dtype=tf.float64)
    #@tf.function(experimental_compile=True)
    def f(x):
        #pdb.set_trace()
        #ind = tf.math.floormod((x-x0), dx)

        if x>=xMax or x<x0:
            return zero#tf.constant(0., dtype=tf.float64);
        else:
            ind = tf.math.floordiv(x-x0, dx)
            remainder = x- ind*dx
            ind=tf.cast(ind, tf.int64)

            return (1.-remainder)*yOrig[ind] + remainder*yOrig[ind+1]
    return f

intpF = tf.function(tf_interpolator(xRef, yRef))
f = tf.function(intpF, experimental_compile=True)
```

```python
xOut =linspace(-0.01, 1.01, 1000)
yOut = np.array([intpF(val) for val in xOut])

```

```python
#randL =np.random.normal(size=10000)
f = tf.function(lambda x: x**2)
f = intpF
randL = tf.random.uniform(minval=0.01, maxval=0.9, shape=(10000,), dtype=tf.float64 )
randL2 = tf.random.uniform(minval=0.01, maxval=0.9, shape=(10000,), dtype=tf.float64 )
randL3 = tf.random.uniform(minval=0.01, maxval=0.9, shape=(10000,), dtype=tf.float64 )

def test(randL):
    for val in randL:
        a = f(val)
t= tf.function(test, experimental_compile=True)

```

```python
%timeit intpF(tf.random.uniform(minval=-0.01, maxval=1.01, dtype=tf.float64, shape=[]))
```

    426 µs ± 49.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

```python
tf.timeit tf.random.unifi
```

```python
#%time t(randL3)
%time t(randL)
```

```python
%time t(randL)
%time t(randL2)
%time t(randL3)
```

    CPU times: user 0 ns, sys: 1e+03 µs, total: 1e+03 µs
    Wall time: 671 µs
    CPU times: user 1e+03 µs, sys: 0 ns, total: 1e+03 µs
    Wall time: 1.49 ms
    CPU times: user 1 ms, sys: 0 ns, total: 1 ms
    Wall time: 609 µs

```python
%timeit intpF(tf.constant(0.5, dtype='float64'))
```

    230 µs ± 50.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

```python
%timeit f(c)
```

```python

#f = tf.function(intpF, experimental_compile=True)
f = intpF
def test(randL):
    return tf.map_fn(f, randL)
t= tf.function(test, experimental_compile=True)
out = t(randL)
%time out=t(randL2)
%time out=t(randL3)
```

    CPU times: user 40 ms, sys: 1 ms, total: 41 ms
    Wall time: 41.2 ms
    CPU times: user 42 ms, sys: 0 ns, total: 42 ms
    Wall time: 41.3 ms

```python
plot(xRef, yRef)
```

```python

```

# Autograd and interpolating

```python
%load_ext autoreload
%autoreload 2
import tensorflow as tf
import sympy as sm
from spylind import spyIVP as so
from pylab import *
import time
from spylind.tf_misc import tf_interpolator2
if 0:
    def tf_interpolator(xOrig, yOrig):
        xOrig = tf.constant(tf.convert_to_tensor(xOrig, dtype=tf.float64))
        yOrig = tf.constant(tf.convert_to_tensor(yOrig, dtype=tf.float64))

        N = xOrig.shape[0]
        dx = xOrig[1]-xOrig[0]
        x0 = xOrig[0]
        xMax = xOrig[-1]
        zero = tf.constant(0, dtype=tf.float64)
        #@tf.function(experimental_compile=True)
        def f(x):
            #pdb.set_trace()
            #ind = tf.math.floormod((x-x0), dx)

            if x>=xMax or x<x0:
                return zero#tf.constant(0., dtype=tf.float64);
            else:
                ind_f = tf.math.floordiv(x-x0, dx)
                remainder = (x/dx- ind_f)
                ind=tf.cast(ind_f, tf.int64)
                #tf.print(remainder)
                return (1.-remainder)*yOrig[ind] + remainder*yOrig[ind+1]
        return f
    def tf_interpolator2(xOrig, yOrig):
        xOrig = tf.constant(tf.convert_to_tensor(xOrig, dtype=tf.float64))
        #yOrig = tf.constant(tf.convert_to_tensor(yOrig, dtype=tf.float64))

        N = xOrig.shape[0]
        dx = xOrig[1]-xOrig[0]
        x0 = xOrig[0]
        xMax = xOrig[-1]
        zero = tf.constant(0, dtype=tf.float64)
        mask = tf.concat([tf.constant([0.5,0.5], dtype='float64'), tf.zeros(xOrig.shape[0]-2, dtype='float64') ], axis=0)
        def f(x):
            #pdb.set_trace()
            #ind = tf.math.floormod((x-x0), dx)

            if x>=xMax or x<x0:
                return zero#tf.constant(0., dtype=tf.float64);
            else:
                ind = tf.math.floordiv(x-x0, dx)
                remainder = x- ind*dx
                ind=tf.cast(ind, tf.int64)
                res = tf.roll(mask, ind, axis=0)*yOrig
                return tf.reduce_sum(res)
        return f

```

```python

xS, yS, vxS, vyS = sm.symbols("x, y, v_x, v_y", real=True)
drvS = sm.symbols("D", real=True)

zS = sm.symbols("z", real=True)
#g= tf.constant(-10.0, dtype='float64')
g = -10.
dy_dtD = {xS:vxS,
       yS:vyS ,
       vxS:0 +0*zS,
       vyS: g + drvS  }

@tf.function(experimental_compile=True)
def f_driving(t):
    return 30*tf.sin(2*pi*t)

tDrv = linspace(0,2,100)
yDrv0 = 0*tf.sin(2*pi*tDrv) + 10
#yDrv0 = tf.convert_to_tensor(yN.repeat(4), dtype='float64')
ode_s = so.ODESolver(dy_dtD,  dims={zS:linspace(-2,2,10)}, backend='tensorflow', default_dtype='float64')
ode_s.set_initial_conditions({xS:0,yS:0,vxS:10,vyS:10})

def calc_g(yDrv):
    yDrv =tf.convert_to_tensor(yDrv, dtype='float64')
    f_driving = tf.function(tf_interpolator2(tDrv, yDrv), experimental_compile=False)

    ode_s.set_driving({drvS:f_driving})
    ode_s.setup(bForceStateDimensions=True)
    tSteps = linspace(0,1.8,20)

    with tf.device("/cpu:0"):
        with tf.GradientTape(watch_accessed_variables=False) as g:
            #g.watch(p0)
            g.watch(yDrv)
            #arr =ode_s.integrate(tSteps, 1.0)
            tStart = time.time()
            arr = ode_s.integrate(tDrv,1.0, method='rk4', solver_options=dict(step_size = (tDrv[1]-tDrv[0])/2 ) )
            print("time taken: ", tStart-time.time())
            #print(arr.shape)
            #L=tf.pow(tf.reduce_mean(arr[-1,1]), 2)
            L=tf.reduce_mean(tf.pow(arr[-30:,1],2))
            #L = arr[-8:,1, 1]**2
            #L = arr[-5:,1]

    print(L.numpy())
    grad = g.gradient(L, yDrv)
    return arr[:,1],grad, L
```

    Munch({'tSym': t, 'dimSyms': [z], 'prop_state_syms': [x, y, v_x, v_y], 'stationary_state_syms': [], 'driving_syms': [], 'state_dep_syms': [D]})
    state dependent functions should have signature [t, z, x, y, v_x, v_y]

```python
from scipy import optimize
def f(yDrv):
    y, grad, L =calc_g(yDrv)
    return L.numpy().squeeze(), grad.numpy()
res=optimize.minimize(f, yDrv0, jac=True)

```

```python
y,grad,L= calc_g(res.x)
```

    Nevals: <tf.Variable 'Variable:0' shape=() dtype=int32, numpy=792>
    time taken:  -1.4418601989746094
    3.0830106257856106e-06

```python
plot(y)
```

```python
grad, yVals =calc_g(1*yDrv0)
plot(yVals)
figure()
plot(grad)
```

```python
#yDrv = tf.identity(yDrv0)
yDrvL = []
gradL = []
yValsL = []
for k in range(30):
    grad, yVals =calc_g(yDrv)
    yDrv = yDrv - 4*grad
    yDrvL.append(yDrv)
    gradL.append(grad)
    yValsL.append(yVals)
```

```python
#plot(yDrv0)
slc = slice(0,30)
for y in yDrvL[slc]:
    plot(y)
legend(arange(10))
title('drv')
figure()
for y in gradL[slc]:
    plot(y)
legend(arange(10))
title('grad')
if 1:
    figure()
    for y in yValsL[slc]:
        plot(tf.reduce_mean(y, axis=-1))
    legend(arange(10))
    title('val')
    k =-1
    figure()
    plot(tf.reduce_mean(yValsL[-2], axis=-1)-tf.reduce_mean(yValsL[-1], axis=-1))
    title('last')
```

```python
plot(grad)
```

```python
plot(tSteps, arr[:,:,50])
```

```python
tf.constant(tf.linspace(0,1,1000))
```

    <tf.Tensor: shape=(1000,), dtype=float64, numpy=
    array([0.        , 0.001001  , 0.002002  , 0.003003  , 0.004004  ,
           0.00500501, 0.00600601, 0.00700701, 0.00800801, 0.00900901,
           0.01001001, 0.01101101, 0.01201201, 0.01301301, 0.01401401,
           0.01501502, 0.01601602, 0.01701702, 0.01801802, 0.01901902,
           0.02002002, 0.02102102, 0.02202202, 0.02302302, 0.02402402,
           0.02502503, 0.02602603, 0.02702703, 0.02802803, 0.02902903,
           0.03003003, 0.03103103, 0.03203203, 0.03303303, 0.03403403,
           0.03503504, 0.03603604, 0.03703704, 0.03803804, 0.03903904,
           0.04004004, 0.04104104, 0.04204204, 0.04304304, 0.04404404,
           0.04504505, 0.04604605, 0.04704705, 0.04804805, 0.04904905,
           0.05005005, 0.05105105, 0.05205205, 0.05305305, 0.05405405,
           0.05505506, 0.05605606, 0.05705706, 0.05805806, 0.05905906,
           0.06006006, 0.06106106, 0.06206206, 0.06306306, 0.06406406,
           0.06506507, 0.06606607, 0.06706707, 0.06806807, 0.06906907,
           0.07007007, 0.07107107, 0.07207207, 0.07307307, 0.07407407,
           0.07507508, 0.07607608, 0.07707708, 0.07807808, 0.07907908,
           0.08008008, 0.08108108, 0.08208208, 0.08308308, 0.08408408,
           0.08508509, 0.08608609, 0.08708709, 0.08808809, 0.08908909,
           0.09009009, 0.09109109, 0.09209209, 0.09309309, 0.09409409,
           0.0950951 , 0.0960961 , 0.0970971 , 0.0980981 , 0.0990991 ,
           0.1001001 , 0.1011011 , 0.1021021 , 0.1031031 , 0.1041041 ,
           0.10510511, 0.10610611, 0.10710711, 0.10810811, 0.10910911,
           0.11011011, 0.11111111, 0.11211211, 0.11311311, 0.11411411,
           0.11511512, 0.11611612, 0.11711712, 0.11811812, 0.11911912,
           0.12012012, 0.12112112, 0.12212212, 0.12312312, 0.12412412,
           0.12512513, 0.12612613, 0.12712713, 0.12812813, 0.12912913,
           0.13013013, 0.13113113, 0.13213213, 0.13313313, 0.13413413,
           0.13513514, 0.13613614, 0.13713714, 0.13813814, 0.13913914,
           0.14014014, 0.14114114, 0.14214214, 0.14314314, 0.14414414,
           0.14514515, 0.14614615, 0.14714715, 0.14814815, 0.14914915,
           0.15015015, 0.15115115, 0.15215215, 0.15315315, 0.15415415,
           0.15515516, 0.15615616, 0.15715716, 0.15815816, 0.15915916,
           0.16016016, 0.16116116, 0.16216216, 0.16316316, 0.16416416,
           0.16516517, 0.16616617, 0.16716717, 0.16816817, 0.16916917,
           0.17017017, 0.17117117, 0.17217217, 0.17317317, 0.17417417,
           0.17517518, 0.17617618, 0.17717718, 0.17817818, 0.17917918,
           0.18018018, 0.18118118, 0.18218218, 0.18318318, 0.18418418,
           0.18518519, 0.18618619, 0.18718719, 0.18818819, 0.18918919,
           0.19019019, 0.19119119, 0.19219219, 0.19319319, 0.19419419,
           0.1951952 , 0.1961962 , 0.1971972 , 0.1981982 , 0.1991992 ,
           0.2002002 , 0.2012012 , 0.2022022 , 0.2032032 , 0.2042042 ,
           0.20520521, 0.20620621, 0.20720721, 0.20820821, 0.20920921,
           0.21021021, 0.21121121, 0.21221221, 0.21321321, 0.21421421,
           0.21521522, 0.21621622, 0.21721722, 0.21821822, 0.21921922,
           0.22022022, 0.22122122, 0.22222222, 0.22322322, 0.22422422,
           0.22522523, 0.22622623, 0.22722723, 0.22822823, 0.22922923,
           0.23023023, 0.23123123, 0.23223223, 0.23323323, 0.23423423,
           0.23523524, 0.23623624, 0.23723724, 0.23823824, 0.23923924,
           0.24024024, 0.24124124, 0.24224224, 0.24324324, 0.24424424,
           0.24524525, 0.24624625, 0.24724725, 0.24824825, 0.24924925,
           0.25025025, 0.25125125, 0.25225225, 0.25325325, 0.25425425,
           0.25525526, 0.25625626, 0.25725726, 0.25825826, 0.25925926,
           0.26026026, 0.26126126, 0.26226226, 0.26326326, 0.26426426,
           0.26526527, 0.26626627, 0.26726727, 0.26826827, 0.26926927,
           0.27027027, 0.27127127, 0.27227227, 0.27327327, 0.27427427,
           0.27527528, 0.27627628, 0.27727728, 0.27827828, 0.27927928,
           0.28028028, 0.28128128, 0.28228228, 0.28328328, 0.28428428,
           0.28528529, 0.28628629, 0.28728729, 0.28828829, 0.28928929,
           0.29029029, 0.29129129, 0.29229229, 0.29329329, 0.29429429,
           0.2952953 , 0.2962963 , 0.2972973 , 0.2982983 , 0.2992993 ,
           0.3003003 , 0.3013013 , 0.3023023 , 0.3033033 , 0.3043043 ,
           0.30530531, 0.30630631, 0.30730731, 0.30830831, 0.30930931,
           0.31031031, 0.31131131, 0.31231231, 0.31331331, 0.31431431,
           0.31531532, 0.31631632, 0.31731732, 0.31831832, 0.31931932,
           0.32032032, 0.32132132, 0.32232232, 0.32332332, 0.32432432,
           0.32532533, 0.32632633, 0.32732733, 0.32832833, 0.32932933,
           0.33033033, 0.33133133, 0.33233233, 0.33333333, 0.33433433,
           0.33533534, 0.33633634, 0.33733734, 0.33833834, 0.33933934,
           0.34034034, 0.34134134, 0.34234234, 0.34334334, 0.34434434,
           0.34534535, 0.34634635, 0.34734735, 0.34834835, 0.34934935,
           0.35035035, 0.35135135, 0.35235235, 0.35335335, 0.35435435,
           0.35535536, 0.35635636, 0.35735736, 0.35835836, 0.35935936,
           0.36036036, 0.36136136, 0.36236236, 0.36336336, 0.36436436,
           0.36536537, 0.36636637, 0.36736737, 0.36836837, 0.36936937,
           0.37037037, 0.37137137, 0.37237237, 0.37337337, 0.37437437,
           0.37537538, 0.37637638, 0.37737738, 0.37837838, 0.37937938,
           0.38038038, 0.38138138, 0.38238238, 0.38338338, 0.38438438,
           0.38538539, 0.38638639, 0.38738739, 0.38838839, 0.38938939,
           0.39039039, 0.39139139, 0.39239239, 0.39339339, 0.39439439,
           0.3953954 , 0.3963964 , 0.3973974 , 0.3983984 , 0.3993994 ,
           0.4004004 , 0.4014014 , 0.4024024 , 0.4034034 , 0.4044044 ,
           0.40540541, 0.40640641, 0.40740741, 0.40840841, 0.40940941,
           0.41041041, 0.41141141, 0.41241241, 0.41341341, 0.41441441,
           0.41541542, 0.41641642, 0.41741742, 0.41841842, 0.41941942,
           0.42042042, 0.42142142, 0.42242242, 0.42342342, 0.42442442,
           0.42542543, 0.42642643, 0.42742743, 0.42842843, 0.42942943,
           0.43043043, 0.43143143, 0.43243243, 0.43343343, 0.43443443,
           0.43543544, 0.43643644, 0.43743744, 0.43843844, 0.43943944,
           0.44044044, 0.44144144, 0.44244244, 0.44344344, 0.44444444,
           0.44544545, 0.44644645, 0.44744745, 0.44844845, 0.44944945,
           0.45045045, 0.45145145, 0.45245245, 0.45345345, 0.45445445,
           0.45545546, 0.45645646, 0.45745746, 0.45845846, 0.45945946,
           0.46046046, 0.46146146, 0.46246246, 0.46346346, 0.46446446,
           0.46546547, 0.46646647, 0.46746747, 0.46846847, 0.46946947,
           0.47047047, 0.47147147, 0.47247247, 0.47347347, 0.47447447,
           0.47547548, 0.47647648, 0.47747748, 0.47847848, 0.47947948,
           0.48048048, 0.48148148, 0.48248248, 0.48348348, 0.48448448,
           0.48548549, 0.48648649, 0.48748749, 0.48848849, 0.48948949,
           0.49049049, 0.49149149, 0.49249249, 0.49349349, 0.49449449,
           0.4954955 , 0.4964965 , 0.4974975 , 0.4984985 , 0.4994995 ,
           0.5005005 , 0.5015015 , 0.5025025 , 0.5035035 , 0.5045045 ,
           0.50550551, 0.50650651, 0.50750751, 0.50850851, 0.50950951,
           0.51051051, 0.51151151, 0.51251251, 0.51351351, 0.51451451,
           0.51551552, 0.51651652, 0.51751752, 0.51851852, 0.51951952,
           0.52052052, 0.52152152, 0.52252252, 0.52352352, 0.52452452,
           0.52552553, 0.52652653, 0.52752753, 0.52852853, 0.52952953,
           0.53053053, 0.53153153, 0.53253253, 0.53353353, 0.53453453,
           0.53553554, 0.53653654, 0.53753754, 0.53853854, 0.53953954,
           0.54054054, 0.54154154, 0.54254254, 0.54354354, 0.54454454,
           0.54554555, 0.54654655, 0.54754755, 0.54854855, 0.54954955,
           0.55055055, 0.55155155, 0.55255255, 0.55355355, 0.55455455,
           0.55555556, 0.55655656, 0.55755756, 0.55855856, 0.55955956,
           0.56056056, 0.56156156, 0.56256256, 0.56356356, 0.56456456,
           0.56556557, 0.56656657, 0.56756757, 0.56856857, 0.56956957,
           0.57057057, 0.57157157, 0.57257257, 0.57357357, 0.57457457,
           0.57557558, 0.57657658, 0.57757758, 0.57857858, 0.57957958,
           0.58058058, 0.58158158, 0.58258258, 0.58358358, 0.58458458,
           0.58558559, 0.58658659, 0.58758759, 0.58858859, 0.58958959,
           0.59059059, 0.59159159, 0.59259259, 0.59359359, 0.59459459,
           0.5955956 , 0.5965966 , 0.5975976 , 0.5985986 , 0.5995996 ,
           0.6006006 , 0.6016016 , 0.6026026 , 0.6036036 , 0.6046046 ,
           0.60560561, 0.60660661, 0.60760761, 0.60860861, 0.60960961,
           0.61061061, 0.61161161, 0.61261261, 0.61361361, 0.61461461,
           0.61561562, 0.61661662, 0.61761762, 0.61861862, 0.61961962,
           0.62062062, 0.62162162, 0.62262262, 0.62362362, 0.62462462,
           0.62562563, 0.62662663, 0.62762763, 0.62862863, 0.62962963,
           0.63063063, 0.63163163, 0.63263263, 0.63363363, 0.63463463,
           0.63563564, 0.63663664, 0.63763764, 0.63863864, 0.63963964,
           0.64064064, 0.64164164, 0.64264264, 0.64364364, 0.64464464,
           0.64564565, 0.64664665, 0.64764765, 0.64864865, 0.64964965,
           0.65065065, 0.65165165, 0.65265265, 0.65365365, 0.65465465,
           0.65565566, 0.65665666, 0.65765766, 0.65865866, 0.65965966,
           0.66066066, 0.66166166, 0.66266266, 0.66366366, 0.66466466,
           0.66566567, 0.66666667, 0.66766767, 0.66866867, 0.66966967,
           0.67067067, 0.67167167, 0.67267267, 0.67367367, 0.67467467,
           0.67567568, 0.67667668, 0.67767768, 0.67867868, 0.67967968,
           0.68068068, 0.68168168, 0.68268268, 0.68368368, 0.68468468,
           0.68568569, 0.68668669, 0.68768769, 0.68868869, 0.68968969,
           0.69069069, 0.69169169, 0.69269269, 0.69369369, 0.69469469,
           0.6956957 , 0.6966967 , 0.6976977 , 0.6986987 , 0.6996997 ,
           0.7007007 , 0.7017017 , 0.7027027 , 0.7037037 , 0.7047047 ,
           0.70570571, 0.70670671, 0.70770771, 0.70870871, 0.70970971,
           0.71071071, 0.71171171, 0.71271271, 0.71371371, 0.71471471,
           0.71571572, 0.71671672, 0.71771772, 0.71871872, 0.71971972,
           0.72072072, 0.72172172, 0.72272272, 0.72372372, 0.72472472,
           0.72572573, 0.72672673, 0.72772773, 0.72872873, 0.72972973,
           0.73073073, 0.73173173, 0.73273273, 0.73373373, 0.73473473,
           0.73573574, 0.73673674, 0.73773774, 0.73873874, 0.73973974,
           0.74074074, 0.74174174, 0.74274274, 0.74374374, 0.74474474,
           0.74574575, 0.74674675, 0.74774775, 0.74874875, 0.74974975,
           0.75075075, 0.75175175, 0.75275275, 0.75375375, 0.75475475,
           0.75575576, 0.75675676, 0.75775776, 0.75875876, 0.75975976,
           0.76076076, 0.76176176, 0.76276276, 0.76376376, 0.76476476,
           0.76576577, 0.76676677, 0.76776777, 0.76876877, 0.76976977,
           0.77077077, 0.77177177, 0.77277277, 0.77377377, 0.77477477,
           0.77577578, 0.77677678, 0.77777778, 0.77877878, 0.77977978,
           0.78078078, 0.78178178, 0.78278278, 0.78378378, 0.78478478,
           0.78578579, 0.78678679, 0.78778779, 0.78878879, 0.78978979,
           0.79079079, 0.79179179, 0.79279279, 0.79379379, 0.79479479,
           0.7957958 , 0.7967968 , 0.7977978 , 0.7987988 , 0.7997998 ,
           0.8008008 , 0.8018018 , 0.8028028 , 0.8038038 , 0.8048048 ,
           0.80580581, 0.80680681, 0.80780781, 0.80880881, 0.80980981,
           0.81081081, 0.81181181, 0.81281281, 0.81381381, 0.81481481,
           0.81581582, 0.81681682, 0.81781782, 0.81881882, 0.81981982,
           0.82082082, 0.82182182, 0.82282282, 0.82382382, 0.82482482,
           0.82582583, 0.82682683, 0.82782783, 0.82882883, 0.82982983,
           0.83083083, 0.83183183, 0.83283283, 0.83383383, 0.83483483,
           0.83583584, 0.83683684, 0.83783784, 0.83883884, 0.83983984,
           0.84084084, 0.84184184, 0.84284284, 0.84384384, 0.84484484,
           0.84584585, 0.84684685, 0.84784785, 0.84884885, 0.84984985,
           0.85085085, 0.85185185, 0.85285285, 0.85385385, 0.85485485,
           0.85585586, 0.85685686, 0.85785786, 0.85885886, 0.85985986,
           0.86086086, 0.86186186, 0.86286286, 0.86386386, 0.86486486,
           0.86586587, 0.86686687, 0.86786787, 0.86886887, 0.86986987,
           0.87087087, 0.87187187, 0.87287287, 0.87387387, 0.87487487,
           0.87587588, 0.87687688, 0.87787788, 0.87887888, 0.87987988,
           0.88088088, 0.88188188, 0.88288288, 0.88388388, 0.88488488,
           0.88588589, 0.88688689, 0.88788789, 0.88888889, 0.88988989,
           0.89089089, 0.89189189, 0.89289289, 0.89389389, 0.89489489,
           0.8958959 , 0.8968969 , 0.8978979 , 0.8988989 , 0.8998999 ,
           0.9009009 , 0.9019019 , 0.9029029 , 0.9039039 , 0.9049049 ,
           0.90590591, 0.90690691, 0.90790791, 0.90890891, 0.90990991,
           0.91091091, 0.91191191, 0.91291291, 0.91391391, 0.91491491,
           0.91591592, 0.91691692, 0.91791792, 0.91891892, 0.91991992,
           0.92092092, 0.92192192, 0.92292292, 0.92392392, 0.92492492,
           0.92592593, 0.92692693, 0.92792793, 0.92892893, 0.92992993,
           0.93093093, 0.93193193, 0.93293293, 0.93393393, 0.93493493,
           0.93593594, 0.93693694, 0.93793794, 0.93893894, 0.93993994,
           0.94094094, 0.94194194, 0.94294294, 0.94394394, 0.94494494,
           0.94594595, 0.94694695, 0.94794795, 0.94894895, 0.94994995,
           0.95095095, 0.95195195, 0.95295295, 0.95395395, 0.95495495,
           0.95595596, 0.95695696, 0.95795796, 0.95895896, 0.95995996,
           0.96096096, 0.96196196, 0.96296296, 0.96396396, 0.96496496,
           0.96596597, 0.96696697, 0.96796797, 0.96896897, 0.96996997,
           0.97097097, 0.97197197, 0.97297297, 0.97397397, 0.97497497,
           0.97597598, 0.97697698, 0.97797798, 0.97897898, 0.97997998,
           0.98098098, 0.98198198, 0.98298298, 0.98398398, 0.98498498,
           0.98598599, 0.98698699, 0.98798799, 0.98898899, 0.98998999,
           0.99099099, 0.99199199, 0.99299299, 0.99399399, 0.99499499,
           0.995996  , 0.996997  , 0.997998  , 0.998999  , 1.        ])>

```python
tf.math.interp_regular_1d_grid(
    x, x_ref_min, x_ref_max, y_ref, axis=-1, fill_value='constant_extension',
    fill_value_below=None, fill_value_above=None, grid_regularizing_transform=None,
    name=None
)
```

```python
mask=tf.concat([tf.constant([1,1], dtype='float64'), tf.zeros(98, dtype='float64') ], axis=0)
```

```python
tf.reduce_sum(mask)
```

    <tf.Tensor: shape=(), dtype=float64, numpy=2.0>

# Scratch test

```python

from tfdiffeq import odeint
#from tfdiffeq import cast_double
import tensorflow as tf
tf.keras.backend.set_floatx('float64') # Presumably can get more speed, especially on a GPU,
#if we used float32
import numpy as np
from matplotlib.pyplot import *
#odeint(func, y0, t)
import matplotlib.pyplot as plt
tf.ceil=tf.math.ceil

#def interpolator(xRef, yRef):

def tf_interpolator3(xOrig, yOrig):
    #xOrig = tf.constant(tf.convert_to_tensor(xOrig, dtype=tf.float64))
    #yOrig = tf.constant(tf.convert_to_tensor(yOrig, dtype=tf.float64))

    N = xOrig.shape[0]
    dx = xOrig[1]-xOrig[0]
    x0 = xOrig[0]
    xMax = xOrig[-1]
    zero = tf.constant(0, dtype=tf.float64)
    #@tf.function(experimental_compile=True)
    def f(x):
        #pdb.set_trace()
        #ind = tf.math.floormod((x-x0), dx)

        if x>=xMax or x<x0:
            return zero#tf.constant(0., dtype=tf.float64);
        else:
            ind_f = tf.math.floordiv(x-x0, dx)
            remainder = (x/dx- ind_f)
            ind=tf.cast(ind_f, tf.int64)
            w0 = (1.-remainder)
            w1 = remainder
            res=tf.reduce_sum(tf.stack([w0,w1])*tf.gather(yOrig, [ind,ind+1]))
            #tf.print(res)
            return res
            #tf.print(remainder)
            #return *yOrig[ind] + remainder*yOrig[ind+1]
    return f

class Lambda(tf.keras.Model):

    def __init__(self, tDrv, yDrv): #pars are parameters to the ode
        #self.pars=pars
        #self.Nevals = tf.Variable(0)#tf.convert_to_tensor(0, dtype=tf.int64)
        #self.f = tf.function(f)
        super().__init__()

        self.f = tf.function(tf_interpolator2(tDrv, yDrv), experimental_compile=False)

    @tf.function
    def call(self, t, z):

        y, vy= z[0], z[1]
        d_y = vy
        d_vy = self.f(t)
        return tf.stack([d_y, d_vy])

with tf.device("/cpu:0"):
    with tf.GradientTape(watch_accessed_variables=False) as g:
        tDrv = tf.cast(tf.linspace(0.,20.1, 100), dtype=tf.float64)
        yDrv = tf.Variable(-3*tf.ones(tDrv.shape[0], dtype=tf.float64))
        g.watch(yDrv)

        tspan = tf.linspace(0.0, 20., 10)
        z0 = tf.convert_to_tensor([0,10.], dtype=tf.float64)
        #f = tf.function(lambda t: tf.constant(-1.0, dtype='float64') )
        ode = Lambda(tDrv, yDrv)
        #tspan =td.cast_double(tspan)
        #ol = C( z0)
        #sol = odeint(ode, z0, tspan, method='dopri5', options=dict(max_step_size=0.1 ) )
        sol = odeint(ode, z0, tspan, method='rk4', options=dict(step_size=tDrv[1]-tDrv[0]))#, options=dict(step_size=0.08080808073282242 ) )

        #%time sol = odeint(ode, z0, tspan, method='rk4')
        #sol = odeint(f, z0, tspan)
        L = sol[-1,0]
    %time grad=g.gradient(L, yDrv)
```

    CPU times: user 292 ms, sys: 13 ms, total: 305 ms
    Wall time: 276 ms

```python

plot(sol)
figure()
plot(grad)
```

```python

plot(sol)
figure()
plot(grad)
```

```python
plot(grad.values)
```

```python
tf.gather(yDrv, 8)
```

    <tf.Tensor: shape=(), dtype=float64, numpy=-3.0>

```python
import tfdiffeq as tfd
tfd.solvers.FixedGridODESolver?
```

```python
def lor(center):
    return 1/(0.4+ (x-center)**2 )

y = linspace(0,10,20)
x = arange(20)
msk = lor(5)
#M = np.diag(  )
plot([y.dot(lor(val)) for val in linspace(0,20,100) ])
```

```python
odeint
```

```python
tf.concat([tf.range(0,5,1.5), [0]], axis=0)
```

    <tf.Tensor: shape=(5,), dtype=float32, numpy=array([0. , 1.5, 3. , 4.5, 0. ], dtype=float32)>

```python
import qutip as q
```

```python

q.sigmax()
K1, K2 = q.basis(2,0), q.basis(2,1)
```

```python
q.sigmay()
```

Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = True\begin{equation*}\left(\begin{array}{*{11}c}0.0 & -1.0j\\1.0j & 0.0\\\end{array}\right)\end{equation*}

```python
a=q.destroy(2)
1j*(a.dag()-a)
```

Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = True\begin{equation*}\left(\begin{array}{*{11}c}0.0 & -1.0j\\1.0j & 0.0\\\end{array}\right)\end{equation*}

```python
a = K1*K2.dag()# + K2*K1.dag()
sx = a + a.dag()
sy = 1j*(a.dag()-a)
sx,sy
```

    (Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = True
     Qobj data =
     [[0. 1.]
      [1. 0.]],
     Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = True
     Qobj data =
     [[0.+0.j 0.-1.j]
      [0.+1.j 0.+0.j]])

```python
a.dag()
```

Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = False\begin{equation*}\left(\begin{array}{*{11}c}0.0 & 0.0\\1.0 & 0.0\\\end{array}\right)\end{equation*}

```python
a
```

Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = False\begin{equation*}\left(\begin{array}{*{11}c}0.0 & 1.0\\0.0 & 0.0\\\end{array}\right)\end{equation*}

# TFP Odeint

```python
%load_ext autoreload
%autoreload 2
import tensorflow as tf
from spylind.tf_misc import tf_interpolator2
tf.keras.backend.set_floatx('float64') # Presumably can get more speed, especially on a GPU,
import tensorflow_probability as tfp

import numpy as np
import pylab as pl
import matplotlib.pyplot as plt
```

```python

class Interpolator(tf.keras.Model):

    def __init__(self, xOrig, yOrig): #pars are parameters to the ode
        super().__init__()
        self.xOrig = tf.constant(tf.convert_to_tensor(xOrig, dtype=tf.float64))
        self.yOrig = tf.Variable(tf.convert_to_tensor(yOrig, dtype=tf.float64))

        self.N = xOrig.shape[0]
        self.dx = xOrig[1]-xOrig[0]
        self.x0 = xOrig[0]
        self.xMax = xOrig[-1]
        self.zero = tf.constant(0, dtype=tf.float64)
        self.mask = tf.concat([tf.constant([0.5,0.5], dtype='float64'), tf.zeros(xOrig.shape[0]-2, dtype='float64') ], axis=0)

    def call(self, x):
        #pdb.set_trace()
        #ind = tf.math.floormod((x-x0), dx)
        if x>=self.xMax or x<self.x0:
            return self.zero#tf.constant(0., dtype=tf.float64);
        else:
            ind = tf.math.floordiv(x-self.x0, self.dx)
            remainder = x- ind*self.dx
            ind = tf.cast(ind, tf.int64)
            res = tf.roll(self.mask, ind, axis=0)*self.yOrig
            return tf.reduce_sum(res)
        return f
class Interpolator2(tf.keras.Model):

    def __init__(self, xOrig): #pars are parameters to the ode
        super().__init__()
        self.xOrig = tf.constant(tf.convert_to_tensor(xOrig, dtype=tf.float64))
        self.N = xOrig.shape[0]
        self.dx = xOrig[1]-xOrig[0]
        self.x0 = xOrig[0]
        self.xMax = xOrig[-1]
        self.zero = tf.constant(0, dtype=tf.float64)
        self.mask = tf.concat([tf.constant([0.5,0.5], dtype='float64'), tf.zeros(xOrig.shape[0]-2, dtype='float64') ], axis=0)

    def call(self, x, yOrig):
        #pdb.set_trace()
        #ind = tf.math.floormod((x-x0), dx)
        if x>=self.xMax or x<self.x0:
            return self.zero#tf.constant(0., dtype=tf.float64);
        else:
            ind = tf.math.floordiv(x-self.x0, self.dx)
            remainder = x- ind*self.dx
            ind = tf.cast(ind, tf.int64)
            res = tf.roll(self.mask, ind, axis=0)*yOrig
            return tf.reduce_sum(res)
        return f

def interpolate(x, xOrig, yOrig):
    dx = xOrig[1]-xOrig[0]
    x0 = xOrig[0]
    xMax = xOrig[-1]
    zero = tf.constant(0, dtype=tf.float64)
    mask = tf.concat([tf.constant([0.5,0.5], dtype='float64'), tf.zeros(xOrig.shape[0]-2, dtype='float64') ], axis=0)
    if x>=xMax or x<x0:
        return zero#tf.constant(0., dtype=tf.float64);
    else:
        ind = tf.math.floordiv(x- x0, dx)
        remainder = x- ind*dx
        ind = tf.cast(ind, tf.int64)
        res = tf.roll(mask, ind, axis=0)*yOrig
        return tf.reduce_sum(res)
    return f

t_init =0;#
#, t0, t1 = 0., 0.5, 1.
tSteps = tf.linspace(0,5,200) #= [t0, t1]
y_init = tf.constant([1., 3.], dtype=tf.float64)
#A = tf.constant([[-1., -2.], [-3., -4.]], dtype=tf.float64)

#@tf.function
#def ode_fn(t, y):
#  return tf.linalg.matvec(A, y)
tDrv = tf.convert_to_tensor(np.linspace(0,5.1, 100), dtype='float64')
yDrv = -tf.ones(tDrv.shape[0], dtype='float64')
#f = tf_interpolator2(tDrv, yDrv)
f = Interpolator(tDrv, yDrv)
@tf.function
def ode_fn(t, z):
    y, vy= z[0], z[1]
    d_y = vy
    d_vy =  f(t)
    return tf.stack([d_y, d_vy])

with tf.GradientTape(watch_accessed_variables=False) as tape:
    tape.watch(yDrv)

    results = tfp.math.ode.DormandPrince().solve(ode_fn, t_init, y_init, solution_times=tSteps)
    y0 = results.states[0]  # == dot(matrix_exp(A * t0), y_init)
    y1 = results.states[1]  # == dot(matrix_exp(A * t1), y_init)
#tape.gradient(results.states[-1,0], yDrv)
```

```python
import sympy as sm
from spylind.tf_misc import InterpolatorMask
ES = sm.symbols("E")

yDrv = tf.Variable(yDrv, name='var', dtype='float64')
class Lambda(tf.keras.Model):

    def __init__(self, tDrv,yDrv): #pars are parameters to the ode
        super().__init__()

        #self.f = tf.function(tf_interpolator2(tDrv, yDrv), experimental_compile=False)
        self.f = InterpolatorMask(tDrv, yDrv)
        self.tDrv = tDrv;
        #self.yDrv = yDrv
        #self.yDrv = tf.Variable(yDrv)
    @tf.function
    def call(self, t, z, **kwargs):
        #yDrv = kwargs[str(ES)][0]
        #tf.print(t, z)
        y, vy= z[0], z[1]
        d_y = vy
        #d_vy = interpolate(t, self.tDrv, yDrv)
        d_vy = self.f(t)

        return tf.stack([d_y, d_vy])

ode_fn = Lambda(tDrv, yDrv)

@tf.function(experimental_compile=True)
def run(tSteps, yDrv):
    with tf.GradientTape(watch_accessed_variables=False) as tape:
        #ode_fn.f.yOrig.assign(yDrv)
        ode_fn.f.set_y(yDrv)

        #tape.watch(yDrv)
        tape.watch(y_init)
        #tape.watch(var)
        #ode_fn.yDrv.assign(yDrv)
        tape.watch(ode_fn.f.yOrig)
        #ode_fn.f.yOrig.assign(yDrv)
        #results = tfp.math.ode.DormandPrince().solve(ode_fn, t_init, y_init, solution_times=tSteps, constants={'yDrv': yDrv})
        #results = tfp.math.ode.DormandPrince().solve(ode_fn, t_init, y_init, solution_times=tSteps)#, constants={str(ES):[yDrv]} )
        results = tfp.math.ode.FixedStep(step_size=1e-4).solve(ode_fn, t_init, y_init, solution_times=tSteps)
        y0 = results.states[:,0]  # == dot(matrix_exp(A * t0), y_init)
        y1 = results.states[:,1]  # == dot(matrix_exp(A * t1), y_init)
        L = y0[-2]
        #L=
    return results, tape.gradient(L, ode_fn.f.yOrig)
```

```python
tf.range(0,1.05,0.1)
```

    <tf.Tensor: shape=(11,), dtype=float32, numpy=
    array([0.        , 0.1       , 0.2       , 0.3       , 0.4       ,
           0.5       , 0.6       , 0.70000005, 0.8000001 , 0.9000001 ,
           1.0000001 ], dtype=float32)>

```python
ode_fn.f.set_y(yDrv*0.1)
#tSteps = np.linspace(0,1, 20)
tSteps = tfp.math.ode.ChosenBySolver(1.0)
res=tfp.math.ode.FixedStep(step_size=1e-2).solve(ode_fn, t_init, y_init, solution_times= tSteps)
```

```python
res=tfp.math.ode.FixedStep(step_size=1e-2).solve(ode_fn, t_init, y_init, solution_times= tSteps)
```

```python
plt.plot(res.states[:-1])
```

```python
#yDrvV = tf.Variable(yDrv)
#tDrvV = tf.Variable(tDrv)
tSteps = np.linspace(0,5,200) #= [t0, t1]
tSteps2 = np.linspace(0,5,20) #= [t0, t1]

drv= np.linspace(0,1, yDrv.shape[0])
%time res1a= run(tSteps, drv[::-1])
#var.assign(5)
#%time res2a= run(tSteps2, 10*drv)

#var.assign(3)
%time res3a= run(tSteps, 100*-drv)
#%time res3= run(1*yDrv)
```

    CPU times: user 128 ms, sys: 4 ms, total: 132 ms
    Wall time: 139 ms
    CPU times: user 126 ms, sys: 0 ns, total: 126 ms
    Wall time: 126 ms

```python
plt.plot(res1a[1])
```

```python
#yDrvV = tf.Variable(yDrv)
#tDrvV = tf.Variable(tDrv)
tSteps = np.linspace(0,5,200) #= [t0, t1]
tSteps2 = np.linspace(0,5,20) #= [t0, t1]

drv= np.linspace(0,1, yDrv.shape[0])
%time res1= run(tSteps, drv[::-1])
#var.assign(5)
%time res2= run(tSteps2, 10*drv)

#var.assign(3)
%time res3= run(tSteps, 100*-drv)
#%time res3= run(1*yDrv)
```

    CPU times: user 5.27 s, sys: 31 ms, total: 5.3 s
    Wall time: 5.43 s
    CPU times: user 5.24 s, sys: 15 ms, total: 5.25 s
    Wall time: 5.36 s
    CPU times: user 7 ms, sys: 0 ns, total: 7 ms
    Wall time: 6.58 ms

```python
plt.plot(res1a[0].states)
```

```python
plt.plot(res1[0].states)
plt.figure()
plt.plot(res1a[0].states)
```

```python
plt.plot(res1[1])
```

```python
dp.ODE_FN_EVALS_PER_STEP
```

    6

```python
%time res1= run(tDrvV, yDrvV)
%time res1= run(tDrvV, yDrvV)
```

    CPU times: user 1.42 s, sys: 79 ms, total: 1.49 s
    Wall time: 1.2 s
    CPU times: user 1.45 s, sys: 75 ms, total: 1.53 s
    Wall time: 1.28 s

```python
plt.plot(res1[0].states)#
```

```python
plt.plot(res1[0].states)#
```

```python
yDrvV = tf.Variable(yDrv)
ode_fn = Lambda(tDrv, yDrvV)
#yDrvV.assigna
ode_fn2 = Lambda(tDrv, 2*yDrv)
ode_fn3 = Lambda(tDrv, 3*yDrv)

%time res1= run(yDrv, ode_fn)
ode_fn.f.yOrig.assign_add(tf.ones(100, dtype='float64'))
%time res2= run(yDrv, ode_fn)
ode_fn.f.yOrig.assign_add(tf.ones(100, dtype='float64'))
%time res3= run(yDrv, ode_fn)
```

    CPU times: user 1.04 s, sys: 3 ms, total: 1.04 s
    Wall time: 1.04 s
    CPU times: user 32 ms, sys: 1 ms, total: 33 ms
    Wall time: 30.1 ms
    CPU times: user 39 ms, sys: 1 ms, total: 40 ms
    Wall time: 57.6 ms

```python
res1[1]
```

    <tf.Tensor: shape=(2,), dtype=float64, numpy=array([1.        , 5.00000006])>

```python
plt.plot(res1.states)
plt.plot(res2.states)
plt.plot(res3.states)
```

```python
tf.GradientTape(watch_accessed_variables=False)
```

```python
plt.plot(results.states.numpy())
```

```python

```


# Using NumPy and SciPy modules

In addition to using Cantera and Pint to help solve thermodynamics problems, we will need to use some additional packages in the scientific Python ecosystem to make plots, solve systems of equations, integrate ordinary differential equations, and more.

```{margin}
The [*SciPy Lecture Notes*](https://scipy-lectures.org) are excellent, detailed resources on all these topics, and Python programming in general {cite}`scipylecture`.
```

The examples contained in this electronic book will integrate these techniques as needed, but this notebook contains some specific examples.

## Index

1. [Plotting](#plotting)
2. [Solving systems of equations](#solving-systems-of-equations)
3. [Integrating ODE systems](#integrating-ode-systems)
4. [Optimization](#optimization)
5. [Differentiation](#differentiation)

## Plotting

We can use [Matplotlib](https://matplotlib.org) to produce nice plots of our results. If you used Anaconda to set up your computing environment, you likely already have Matplotlib installed; if not, see their [installation instructions](https://matplotlib.org/users/installing.html).

Matplotlib provides an interface that is very similar to what you might already know from Matlab: `pyplot`. You can import this in Python files or a Jupyter notebook with the standard abbreviation `plt`:

```python
# this line makes figures interactive in Jupyter notebooks
%matplotlib inline
from matplotlib import pyplot as plt
```

```python
# these lines are only for helping improve the display
import matplotlib_inline.backend_inline
matplotlib_inline.backend_inline.set_matplotlib_formats('pdf', 'png')
plt.rcParams['figure.dpi']= 150
plt.rcParams['savefig.dpi'] = 150
```

For example, let's generate some values of an independent variable $x$ linearly spaced between 0 and 10 (using the NumPy function [`linspace()`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.linspace.html)), and then plot the function $y(x) = \sin(x)$. We can also add labels to the axes, a legend, and a helpful grid.

```python
import numpy as np

x = np.linspace(0, 10, num=50, endpoint=True)
y = np.sin(x)

plt.plot(x, y, label='y(x)')
plt.xlabel('x axis')
plt.ylabel('y axis')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
```

We can also plot multiple data series in a single figure:

```python
plt.plot(x, y, label='sin(x)')
plt.plot(x, np.cos(x), label='cos(x)')
plt.xlabel('x axis')
plt.ylabel('y axis')
plt.grid(True)
plt.legend()
plt.show()
```

Or, we can use subplots to plot multiple axes in the same overall figure:

```python
# 2 rows, 1 column
fig, axes = plt.subplots(2, 1)

axes[0].plot(x, y, label='sin(x)')
axes[0].set_ylabel('sin(x)')
axes[0].grid(True)

axes[1].plot(x, np.cos(x), label='cos(x)')
axes[1].set_xlabel('x axis')
axes[1].set_ylabel('cos(x)')
axes[1].grid(True)

plt.show()
```

## Solving systems of equations

Frequently we will encounter a system of one or more equations that involves an equal number of unknowns. If this is a linear system of equations, we can use linear algebra and the NumPy [`linalg.solve()`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.solve.html) function, but more often in thermodynamics we encounter complex and/or nonlinear systems.

In cases like this, we will need to set up our problems to find the roots, or zeroes, of the function(s); in other words, given a function $ f(x) $, finding the root means to find the value of $x$ such that $ f(x) = 0 $. If we are dealing with a system of equations and the same number of unknowns, then these would be vectors: $\mathbf{f}(\mathbf{x}) = 0$.

(You might be wondering what to do about equations that don't equal zero... for example, if you have something like $f(x) = g(x)$. In this case, you just need to manipulate the equation to be in the form $f(x) - g(x) = 0$.)

The [SciPy optimization module](https://docs.scipy.org/doc/scipy/reference/optimize.html) provides functions to find roots of equations; for scalar equations, we can use [`root_scalar()`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.root_scalar.html#scipy.optimize.root_scalar), and for vector equations, we can use [`root()`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.root.html).

### Scalar equations

Let's first look at an example of a scalar function: one equation, one unknown.

Find the root of this equation:
\begin{equation}
\cos(x) = x^3
\end{equation}

We need to create a Python function that returns $f(x) = 0$, so that the function returns zero when the input value of $x$ is the (correct) root. Then, we can use the `root_scalar` function with some initial guesses.

```python
import numpy as np
from scipy import optimize

def func(x):
    return np.cos(x) - x**3

sol = optimize.root_scalar(func, x0=1.0, x1=2.0)

print(f'Root: x ={sol.root: .3f}')

print(f'Function evaluated at root: {func(sol.root)}')
```

    Root: x = 0.865
    Function evaluated at root: -2.220446049250313e-16

### Systems of equations / vector functions

We can also solve systems of equations in a similar fashion using the SciPy [`root()`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.root.html) functino, where we find the roots $\mathbf{x}$ that satisfy $\mathbf{f} (\mathbf{x}) = 0$.

For example, let's try to find the values of $x$ and $y$ that satisfy these equations:
$$
x \ln (x) = y^3 \\
\sqrt{x} = \frac{1}{y}
$$

We have two equations and two unknowns, so we should be able to find the roots.

To solve, we create a function that evaluates these equations when made equal to zero, or

$$
x \ln (x) - y^3 = 0 \\
\sqrt{x} - \frac{1}{y} = 0
$$

then we call `root` specifying this function and two initial guesses for $x$ and $y$:

```python
import numpy as np
from scipy import optimize

def system(vars):
    x = vars[0]
    y = vars[1]

    return [
        x*np.log(x) - y**3,
        np.sqrt(x) - (1/y)
        ]

sol = optimize.root(system, [1.0, 1.0])

x = sol.x[0]
y = sol.x[1]

print(f'Roots: x = {x: .3f}, y = {y: .3f}')
```

    Roots: x =  1.467, y =  0.826

## Integrating ODE systems

In some cases, we encounter problems that require integrating one or more ordinary different equations in time. Depending on the form of the problem, we may need to integrate a function between two points (definite integral), or we may have a system of ordinary differential equations.

### Numerical integral of samples

In some cases we have a set of $(x,y)$ data that we want to integrate numerically.
We can do this using the NumPy [`trapz()` function](https://docs.scipy.org/doc/numpy/reference/generated/numpy.trapz.html), which implements the composite trapezoidal rule.

For example, let's consider a situation where a reciprocating compressor is being used to compress ammonia vapor during a refrigeration cycle. We have some experimental measurements of the pressure-volume data during the compression stroke (see table), and we want to determine the work done on the ammonia by the piston.

| Pressure (psi) | Volume (in^3) |
|----------------|---------------|
| 65.1           | 80.0          |
| 80.5           | 67.2          |
| 93.2           | 60.1          |
| 110            | 52.5          |
| 134            | 44.8          |
| 161            | 37.6          |
| 190            | 32.5          |

To find the work done by the piston to the ammonia, we can integrate pressure with respect to volume:
\begin{equation}
W_{\text{in}} = -\int_{V_1}^{V_2} P \, dV
\end{equation}

```python
import numpy as np
from pint import UnitRegistry
ureg = UnitRegistry()
Q_ = ureg.Quantity

pressure = Q_([65.1, 80.5, 93.2, 110, 134, 161, 190], 'psi')
volume = Q_([80.0, 67.2, 60.1, 52.5, 44.8, 37.6, 32.5], 'in^3')

# convert to SI units
pressure.ito('Pa')
volume.ito('m^3')

work = -Q_(
    np.trapz(pressure.magnitude, volume.magnitude),
    pressure.units * volume.units
    )
print(f'Work done on fluid: {work.to("J"): .2f}')
```

    Work done on fluid: 589.45 joule

### Numerical integral of expression

We can also numerically integrate expressions/functions using the SciPy [`quad()` function](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.quad.html), also in the `integrate` module.

Let's build on the previous example: first, let's fit the compressor data to the polytropic form $P V^n = c$, where $c$ and $n$ are constants, then we can integrate the resulting function to calculate work.

To fit the data, we need to rearrange the equation to the form $y = f(x)$:
\begin{equation}
P = c V^{-n}
\end{equation}
We'll use the [`scipy.optimize.curve_fit()`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.curve_fit.html) function for that.

```python
import numpy as np
from scipy.optimize import curve_fit
from scipy.integrate import quad

from pint import UnitRegistry
ureg = UnitRegistry()
Q_ = ureg.Quantity

pressure = Q_([65.1, 80.5, 93.2, 110, 134, 161, 190], 'psi')
volume = Q_([80.0, 67.2, 60.1, 52.5, 44.8, 37.6, 32.5], 'in^3')

# convert to SI units
pressure.ito('Pa')
volume.ito('m^3')

def fit(V, n, c):
    '''Evaluate P = c * V**(-n).

    n and c will be found by the curve_fit functino.
    '''
    return c * np.power(V, -n)

# this function will automatically fit the unknown constants
params, cov = curve_fit(fit, volume.magnitude, pressure.magnitude)

print(f'Parameters: n={params[0]: 5.3f}, c={params[1]: 5.3f}')
```

    Parameters: n= 1.170, c= 194.200

    /var/folders/ng/d9rd9fb92c7bxz9fy7vwpj7m0000gp/T/ipykernel_14239/2169371920.py:21: RuntimeWarning: overflow encountered in power
      return c * np.power(V, -n)

```python
plt.plot(volume.magnitude, pressure.magnitude, 'o', label='Data')

x = np.linspace(volume[0].magnitude, volume[-1].magnitude, 100)

# plot data and fit
plt.plot(x, fit(x, *params), 'r-',
         label=f'fit: n={params[0]: 5.3f}, c={params[1]: 5.3f}'
         )
plt.xlabel('Volume (m^3)')
plt.ylabel('Pressure (Pa)')
plt.legend()
plt.tight_layout()
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
plt.show()
```

That looks like a great fit to the data, so now let's evaluate the work by integrating:

```python
y, err = quad(
    fit, volume[0].magnitude, volume[-1].magnitude,
    args=(params[0], params[1])
    )

work = -Q_(y, pressure.units * volume.units)

print(f'Work done on fluid: {work.to("J"): .2f}')
```

    Work done on fluid: 586.26 joule

### Initial value problems

In other words, we may have a system like
\begin{equation}
\frac{d \mathbf{y}}{dt} = \mathbf{f} (t, \mathbf{y})
\end{equation}
where $\mathbf{y}$ is the vector of state variables, for which we should have the initial values ($\mathbf{y}(t=0)$). In this case, we can use the SciPy function [`solve_ivp()` function](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html), part of the `integrate` module.

For example, let's consider a problem where we want to find the volume of water, volume of air, and pressure in a tank as a function of time, as the tank releases water to an environment. The volumetric flow rate through the tank's valve is based on its instantaneous pressure:

$$
\dot{V}_f = \frac{P - P_{\text{atm}}}{R_v} \;,
$$

where $P_{\text{atm}}$ is atmospheric pressure and $R_v = 10$ psi/gpm is the valve resistance parameter. The tank has volume 50 gal, and is initially filled with water (volume fraction $f$ = 0.8); the air in the tank is initially at 100 psi. We will treat the air as an ideal gas, and the water as an incompressible substance.

We have the initial air pressure, and we can calculate the initial volumes of water and air:

$$
V_{f,0} = f \, V \\
V_{g,0} = (1 - f) V \\
$$

We can perform a mass balance on the water in the tank to obtain a rate equation for the volume of water:

$$
0 = \dot{m}_f + \frac{d m_f}{dt} = \dot{V}_f \rho_f + \rho_f \frac{d V_f}{dt} \\
0 = \dot{V}_f + \frac{dV_f}{dt} \\
\therefore \frac{dV_f}{dt} = -\frac{(P - P_{\text{atm}})}{R_v}
$$

Since the overall volume of the tank is constant, we can obtain the rate of change of the volume of air:

$$
V = V_g + V_f \\
\frac{dV_g}{dt} + \frac{dV_f}{dt} = 0 \\
\therefore \frac{dV_g}{dt} = -\frac{dV_f}{dt}
$$

Finally, a mass balance on the air in the tank allows us to find the rate of change of pressure:

$$
0 = \frac{dm_a}{dt} \\
0 = \frac{d}{dt} \left( \frac{P V_g}{R T} \right) \\
0 = V_g \frac{dP}{dt} + P \frac{dV_g}{dt} \\
\therefore \frac{dP}{dt} = -\frac{P}{V_g} \frac{dV_g}{dt}
$$

Now, we can integrate this system of ODEs:

```python
from pint import UnitRegistry
ureg = UnitRegistry()
Q_ = ureg.Quantity

# initial conditions and constants
valve_resistance = Q_(10, 'psi/(gal/min)')
pressure_atmosphere = Q_(1, 'atm')

volume_tank = Q_(50, 'gal')
water_fraction = 0.8

pressure_initial = Q_(100, 'psi')
volume_water_initial = water_fraction * volume_tank
volume_air_initial = (1.0 - water_fraction) * volume_tank
```

```python
def tank_equations(t, y, valve_resistance, pressure_atmosphere):
    '''Rates of change for water volume, air volume, and air pressure in tank.

    Input values in SI units.
    '''
    volume_water = Q_(y[0], 'm^3')
    volume_air = Q_(y[1], 'm^3')
    pressure_air = Q_(y[2], 'Pa')

    dVf_dt = -(pressure_air - pressure_atmosphere) / valve_resistance
    dVg_dt = -dVf_dt
    dP_dt = -(pressure_air / volume_air) * dVg_dt

    return [
        dVf_dt.to('m^3/s').magnitude,
        dVg_dt.to('m^3/s').magnitude,
        dP_dt.to('Pa/s').magnitude
        ]
```

Now that we have the initial conditions and rate function set up, we can integrate in time. Let's do this for 500 seconds, and then plot the water volume and tank pressure as functions of time:

```python
from scipy.integrate import solve_ivp

# now integrate for 500 seconds, specifying the function, time interval,
#initial conditions, and additional arguments to the function

sol = solve_ivp(
    tank_equations, [0, 500.0],
    [volume_water_initial.to('m^3').magnitude,
     volume_air_initial.to('m^3').magnitude,
     pressure_initial.to('Pa').magnitude
     ],
    args=(valve_resistance, pressure_atmosphere,),
    method='BDF'
    )
```

```python
time = sol.t
volume_water = Q_(sol.y[0], 'm^3')
volume_air = Q_(sol.y[1], 'm^3')
pressure = Q_(sol.y[2], 'Pa')

plt.plot(time, pressure.to('psi').magnitude, label='Tank pressure (psi)')
plt.plot(time, volume_water.to('gal').magnitude, label='Water volume (gal)')
plt.grid(True)
plt.xlabel('Time (s)')
plt.ylabel('Water volume (gal) and tank pressure (psi)')
plt.ylim(ymin=0.0) # ensure lower y-axis bound is zero
plt.legend()
plt.tight_layout()
plt.show()
```

## Optimization

For some problems, we may want to find the input parameter(s) that minimize or maximize some function. In these cases, we can use the SciPy [`minimize_scalar()`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize_scalar.html) function for a scalar function with one input variable, or the [`minimize()`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html) function for a scalar function of one or more input variables.

(Note: this is a fairly complicated topic, and we will only consider relatively simple optimization problems.)

For example, consider a spring with a horizontal force applied; we can calculate the potential energy of the spring:
\begin{equation}
PE(x) = 0.5 k x^2 - F x \;.
\end{equation}
At equilibrium, the potential energy will be at a minimum, so given a particular force we can find the displacement by finding the value that minimizes the potential energy.

```python
from scipy.optimize import minimize_scalar

spring_constant = 2.0 # N/cm
force = 5.0 # N

def spring(x, k, F):
    '''Calculates potential energy of spring pulled by force.
    '''
    return 0.5 * k * x**2 - F * x

sol = minimize_scalar(spring, args=(spring_constant, force))

print(f'Equilibrium displacement: {sol.x: .2f} cm')
```

    Equilibrium displacement:  2.50 cm

A more complicated problem is a system with two springs, connected at one end with a force in some general direction (with $x$ and $y$ components); the system has both horizontal and vertical components. The springs have spring constants $k_a$ and $k_b$, and unloaded lengths $L_a$ and $L_b$.

The potential energy for this system is
\begin{equation}
PE(x,y) = 0.5 k_a \left( \sqrt{x^2 + (L_a - y)^2} - L_a \right)^2 + 0.5 k_b \left( \sqrt{x^2 + (L_b + y)^2} - L_b \right)^2 - F_x x - F_y y
\end{equation}
where $x$ and $y$ are the horizontal and vertical deformations from the unloaded state.
We can minimize this function of two variables to find the displacement based on the force:

```python
from scipy.optimize import minimize

spring_constant_a = 9.0 # N/cm
spring_constant_b = 2.0 # N/cm
length_a = 10.0 # cm
length_b = 10.0 # cm

force_x = 2.0 # N
force_y = 4.0 # N

def spring_system(xvec, ka, kb, La, Lb, Fx, Fy):
    '''Calculates potential energy of springs pulled by force.
    '''
    x = xvec[0]
    y = xvec[1]
    return (
        0.5*ka * (np.sqrt(x**2 + (La - y)**2) - La)**2 +
        0.5*kb * (np.sqrt(x**2 + (Lb + y)**2) - Lb)**2
        - (Fx * x) - (Fy * y)
        )

guesses = [1.0, 1.0]
sol = minimize(
    spring_system, guesses,
    args=(spring_constant_a, spring_constant_b,
          length_a, length_b, force_x, force_y
          )
    )

x = sol.x[0]
y = sol.x[1]
print(f'Equilibrium displacement: x={x: .2f} cm, y={y: .2f} cm')
```

    Equilibrium displacement: x= 4.95 cm, y= 1.28 cm

## Differentiation

Many thermodynamic properties are derivatives of other properties, so you may find the need to take derivatives of either data or of an expression. We can use the NumPy function [`gradient()`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.gradient.html) to take numerical derivatives of data using finite differences, and we can use [SymPy](https://www.sympy.org/en/index.html) to find analytical derivatives of expressions.

### Numerical derivatives

Let's say we have data for some function $f(x) = x^2$ at some locations $x$:

```python
x = np.linspace(0, 1, 10)
f = x**2

plt.plot(x, f)
plt.xlabel('x')
plt.ylabel('f(x)')
plt.grid(True)
plt.tight_layout()
plt.show()
```

We do not need to know the functional form of $f(x)$ to take the numerical derivatives—this example just uses a simple function for ease of checking the results.

The `gradient()` function uses second-order central differences to evaluate the derivative of input data, using forward and backward differences at the boundaries. Let's use that to get the derivative, and compare against the exact derivative ($\frac{df}{dx} = 2x$):

```python
dfdx = np.gradient(f, x)
dfdx_exact = 2*x

plt.plot(x, dfdx_exact, label='Exact')
plt.plot(x, dfdx, 'o', label='Numerical')
plt.xlabel('x')
plt.ylabel('df/dx')

plt.grid(True)
plt.legend()
plt.show()
```

The numerical derivative we obtain is very accurate, thanks to the linear nature of the derivative. However, at the boundaries the approximate derivative is a bit off, due to the first-order finite differences used there.

### Analytical derivative

In some cases we may be given (or obtain) analytical expressions that we need to differentiate. Given a function, we have two options for calculating the derivative:

1. Use the function to calculate values over the desired range, then numerically differentiate.
2. Obtain the exact derivative by differentiating analytically.

We'll now focus on the latter case; we can use SymPy to construct a function symbolically, and then find the exact derivative.

```python
import sympy
sympy.init_printing(use_latex='mathjax')

x = sympy.symbols('x', real=True)

f = x**2

# take derivative of f with respect to x
dfdx = sympy.diff(f, x)

# this complicated expression is only necessary for printing.
# It creates an equation object, that sets the unevaluated
# derivative of f equal to the calculated derivative.
display(sympy.Eq(sympy.Derivative(f), dfdx))
```

$\displaystyle \frac{d}{d x} x^{2} = 2 x$

Once we have evaluated the analytical derivative, we can even turn it into a Python function for evaluation! Passing the `'numpy'` argument creates a NumPy array-compatible function.

```python
calc_derivative = sympy.lambdify(x, dfdx, 'numpy')

x_vals = np.linspace(0, 1, 10)

plt.plot(x_vals, 2*x_vals, label='Exact derivative')
plt.plot(x_vals, calc_derivative(x_vals), 'o', label='SymPy derivative')

plt.xlabel('x')
plt.ylabel('df/dx')
plt.grid(True)
plt.legend()
plt.show()
```

As expected, the exact derivative we calculated using SymPy matches the known derivative perfectly—this was obtained symbolically, so there are no approximations/errors involved.

When taking the derivative of more complicated expressions that involve other functions (e.g., $\sin$, $\log$), you will need to use the SymPy-provided versions: `sympy.sin`, `sympy.log`, etc. We also need to explicitly define symbolic variables, either using `x = sympy.Symbol('x')` to create one variable at a time or `x, y = sympy.symbols('x y', real=True)` to create multiple variables at once. (The string you specify is the displayed representation of the variable, and can include complex formatting such as subscripts/superscripts and Greek letters.)

For example, let's take the derivative of
\begin{equation}
f(x) = \log(x) + x^3
\end{equation}

```python
x = sympy.Symbol('x')
f = sympy.log(x) + x**3

dfdx = sympy.diff(f, x)

display(dfdx)
```

$\displaystyle 3 x^{2} + \frac{1}{x}$

This matches what we expect for this derivative.


[Sascha Spors](https://orcid.org/0000-0001-7225-9992),
Professorship Signal Theory and Digital Signal Processing,
[Institute of Communications Engineering (INT)](https://www.int.uni-rostock.de/),
Faculty of Computer Science and Electrical Engineering (IEF),
[University of Rostock, Germany](https://www.uni-rostock.de/en/)

# Tutorial Signals and Systems (Signal- und Systemtheorie)

Summer Semester 2021 (Bachelor Course #24015)

- lecture: https://github.com/spatialaudio/signals-and-systems-lecture
- tutorial: https://github.com/spatialaudio/signals-and-systems-exercises

WIP...
The project is currently under heavy development while adding new material for the summer semester 2021

Feel free to contact lecturer [frank.schultz@uni-rostock.de](https://orcid.org/0000-0002-3010-0294)

## Numerisches Rechnen mit numpy Paket

```python
import numpy as np
import matplotlib.pyplot as plt

sigma, omega = -0.2, 5
s = sigma + 1j*omega
t = np.linspace(0,10,1000)
x = np.exp(s*t)

plt.plot(t,np.real(x), label=r'$\Re(x)$', color='C0')
plt.plot(t,np.imag(x), label=r'$\Im(x)$', color='C3')
plt.xlabel('t /s')
plt.ylabel('x(t)')
plt.title(r'$x(t) = \mathrm{e}^{\mathrm{j} \omega t}$')
plt.legend()
plt.grid(True)

```

## Symbolisches Rechnen mit sympy Paket

```python
import sympy as sp

t, sigma, omega = sp.symbols('t sigma omega', real=True)
s = sigma + 1j*omega
x = sp.exp(s*t)
x_num = x.subs({sigma: -0.2, omega: 5})

print(x)
print(x_num)

sp.plot(sp.re(x_num), (t, 0, 10), xlabel='t/s', ylabel=r'$\Re(x)$', title=r'$x(t) = \mathrm{e}^{\mathrm{j} \omega t}$')
sp.plot(sp.im(x_num), (t, 0, 10), xlabel='t/s', ylabel=r'$\Im(x)$');

```

## Copyright

This tutorial is provided as Open Educational Resource (OER), to be found at
https://github.com/spatialaudio/signals-and-systems-exercises
accompanying the OER lecture
https://github.com/spatialaudio/signals-and-systems-lecture.
Both are licensed under a) the Creative Commons Attribution 4.0 International
License for text and graphics and b) the MIT License for source code.
Please attribute material from the tutorial as *Frank Schultz,
Continuous- and Discrete-Time Signals and Systems - A Tutorial Featuring
Computational Examples, University of Rostock* with
``main file, github URL, commit number and/or version tag, year``.