# Discovery Model Example

Here we will attempt to learn the parameters provided in [this example](../../model/compiling-example/index.md) (i.e. .0001 and 5.0) from the 
data provided. We consider that data, in this case, to be "experimental" data, however we do know it to be high-fidelity data from a  simulation of the 
AC PDE system. 

First, we present the whole script to train a discovery model, then we will break it apart and examine the major chunks. 

## Full Example
```{code} python 
# Put params into a list
params = [tf.Variable(0.0, dtype=tf.float32), tf.Variable(0.0, dtype=tf.float32)]


# Define f_model, note the `vars` argument. Inputs must follow this order!
def f_model(u_model, var, x, t):
    u = u_model(tf.concat([x, t], 1))
    u_x = tf.gradients(u, x)
    u_xx = tf.gradients(u_x, x)
    u_t = tf.gradients(u, t)
    c1 = var[0]  # tunable param 1
    c2 = var[1]  # tunable param 2
    f_u = u_t - c1 * u_xx + c2 * u * u * u - c2 * u
    return f_u


# Import data, same data as Raissi et al

data = scipy.io.loadmat('AC.mat')

t = data['tt'].flatten()[:, None]
x = data['x'].flatten()[:, None]
Exact = data['uu']
Exact_u = np.real(Exact)

# generate all combinations of x and t
X, T = np.meshgrid(x, t)

X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))
u_star = Exact_u.T.flatten()[:, None]

x = X_star[:, 0:1]
t = X_star[:, 1:2]

print(np.shape(x))
# append to a list for input to model.fit
X = [x, t]

# define MLP depth and layer width
layer_sizes = [2, 128, 128, 128, 128, 1]

# initialize, compile, train model
model = DiscoveryModel()
model.compile(layer_sizes, f_model, X, u_star, params) 

# train loop
model.fit(tf_iter=10000)

```

Let's break this apart and look at its pieces.

### Defining parameters and `f_model` for estimation
First we define the `tf.Variable` objects for the parameters and the new `f_model`. Note that the structure and syntax is largely the same as the [`CollocationSolverND` example](../../model/compiling-example/index.md), with a few notable exceptions.

```{code} python
# Put params into a list
params = [tf.Variable(0.0, dtype=tf.float32), tf.Variable(0.0, dtype=tf.float32)]


# Define f_model, note the `vars` argument. Inputs must follow this order!
def f_model(u_model, var, x, t):
    u = u_model(tf.concat([x, t], 1))
    u_x = tf.gradients(u, x)
    u_xx = tf.gradients(u_x, x)
    u_t = tf.gradients(u, t)
    c1 = var[0]  # tunable param 1
    c2 = var[1]  # tunable param 2
    f_u = u_t - c1 * u_xx + c2 * u * u * u - c2 * u
    return f_u
```

Above, we can see that the parameters must be `tf.Variables,` initialized as above, with a `tf.float32` data type. 
You must initialize as many of these as there are parameters to estimate. Those variables 
must then be added to a `list` for training at a later step. Here we initialize the parameters to `0.0` to start. As a heuristic, this 
works well since the $u$ network also needs to get somewhat close to a `0.0` residual solution before the trianing of the 
parameters can really start to take root.

Concurrently, we generate the new `f_model`. As discussed earlier, the new `f_model` contains an additional input from its 
[`CollocationSolverND` cousin](../../model/compiling-example/index.md) - the `var` input. This input is where the `list` of 
`tf.Variables` goes. Inside the `f_model` definition, that list is then partitioned out piecewise into the PDE. This allows the 
tensorflow tracing to reach into the `f_model` function and backpropagate against those values, resulting in training of the parameters 
as well as the $u$ network itself.

### Importing data and generating input
```{code} python
# Import data, same data as Raissi et al
data = scipy.io.loadmat('AC.mat')

t = data['tt'].flatten()[:, None]
x = data['x'].flatten()[:, None]
Exact = data['uu']
Exact_u = np.real(Exact)

# generate all combinations of x and t
X, T = np.meshgrid(x, t)

X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))
u_star = Exact_u.T.flatten()[:, None]

x = X_star[:, 0:1]
t = X_star[:, 1:2]

# append to a list for input to model.fit
X = [x, t]
```

In this case, the input `x` and `t` sequences were held in the file. These took a similar form to [`np.linspace` objects](https://numpy.org/doc/stable/reference/generated/numpy.linspace.html), i.e. were vectors of even spacing across the `x` and `t` 
dimensions, independently. Therefore, we needed to generate all possible combinations 
of `x` and `t` to use these. If you are in need of a multidimensional meshgrid generator (past 2D) that returns a list of all possible combinations
of `np.linepace` type arrays, check out [this github gist](https://gist.github.com/levimcclenny/e87dd0979e339ea89a9885ec05fe7c10) to get the input in the format tensordiffeq requires. This multimesh generation is 
included in tdq base and is available by combining the function `multimesh` with `flatten_and_stack`, both in `tensordiffeq.utils`. Note that you need 
an `X, u_sol` pair for each of your data points. So, if you have a 1D (with time) problem, then you need an input pair that of the form `[x,t]` and 
a target `u_sol` value. Essentially, we are performing supervised learning of the parameters, therefore we need some target value for each input coordinate in 
the domain where we have data available. 

### Defining the network and training

Next we define the `layer_size`, similar to the [`CollocationSolverND` example](../../model/compiling-example/index.md)

```{code} python 
# define MLP depth and layer width
layer_sizes = [2, 128, 128, 128, 128, 1]
```

Finally, we can compile with all the parameters we defined above and begin training!

```{code} python 
# initialize, compile, train model
model = DiscoveryModel()
model.compile(layer_sizes, f_model, X, u_star, params) 

# train loop
model.fit(tf_iter=10000)
```