fix external variable initialization

[bonneted]

I've faced two bugs when trying to implement :
https://github.com/lu-group/sbinn/blob/b2c1c94d6564732189722f6e6772af0f63cb0d8c/sbinn/sbinn_tf.py#L8

• in model.py it's because the external variables were not initialized on the second compile as the parameters of the net were already

• in pde.py if I don't compile with external variables I still want the code to work with the default values of unknowns
  I think this code can be safely modified only for jax because the line after was already only for jax

[lululxvi]

Could you point out an example for using this code?

[bonneted]

Here : https://github.com/bonneted/sbinn/blob/main/sbinn/sbinn_jax.py

The implementation of sbinn using JAX.
We first train the model without the external variables :

    def ODE(t, y, unknowns=[var.value for var in var_list_]):
    ...

    model.compile("adam", lr=1e-3, loss_weights=[0, 0, 0, 0, 0, 0, 1e-2])
    model.train(epochs=firsttrain, display_every=1000)
    model.compile(
        "adam",
        lr=1e-3,
        loss_weights=[1, 1, 1e-2, 1, 1, 1, 1e-2],
        external_trainable_variables=var_list_,
    )
    variablefilename = "variables.csv"
    variable = dde.callbacks.VariableValue(
        var_list_, period=callbackperiod, filename=variablefilename
    )
    losshistory, train_state = model.train(
        epochs=maxepochs, display_every=1000, callbacks=[variable]
    )

For this first train, we want to use the default unknowns argument for the ODE

[lululxvi]

The code modification seems necessary. But there is another example https://github.com/lululxvi/deepxde/blob/master/examples/pinn_inverse/Lorenz_inverse.py , which works well (at least worked earlier).

[bonneted]

This one was already working well because there is no pertaining without the external variables.
The model is only compiled with the external trainable variables :

model.compile(
    "adam", lr=0.001, external_trainable_variables=external_trainable_variables
)
losshistory, train_state = model.train(iterations=20000, callbacks=[variable])

The problem occurs when we compile without the external trainable variables, which is when we want the PDE to use the default unknowns argument.

[lululxvi]

The code seems OK. But the underlying logic becomes extremely complicated now.

In fact, you can simply add external_trainable_variables in the first compile. As the PDE loss weight is 0, those variables won't get updated any way.

[bonneted]

That's true in that case, but it can be interesting to start training the model with frozen parameters (for example https://doi.org/10.1126/sciadv.abk0644)
Moreover, it would mean that putting default unknowns values for the PDE is useless and misleading as they can never be used.

[lululxvi]

Please resolve the conflicts.

[bonneted]

I've resolved the conflict based on your improved logic.
In the JAX backend conditional I added the possibility that there are no external trainable variables but a default value available.

[lululxvi]

What if self.external_trainable_variables is []?

[bonneted]

it's working fine, var.value is never evaluated and the output is []

[lululxvi]

In this case,

self.params = [self.net.params, []]

which is strange.

[bonneted]

It was already like that before the introduction external_trainable_variables_val in a previous commit.

It used to be: (see bonneted@3810a98)
self.params = [self.net.params, self.external_trainable_variables]
with self.external_trainable_variables = [] by default

We can see how it works in the original commit by @ZongrenZou : bonneted@09c74ab
self.params is split in nn_params and ext_params in outputs_losses function:

  def outputs_losses(params, training, inputs, targets, losses_fn):
      nn_params, ext_params = params
      # TODO: Add auxiliary vars
      def outputs_fn(inputs):
          return self.net.apply(nn_params, inputs, training=training)
      outputs_ = self.net.apply(nn_params, inputs, training=training)
      # Data losses
      # We use aux so that self.data.losses is a pure function.
      aux = [outputs_fn, ext_params] if ext_params else [outputs_fn]
      losses = losses_fn(targets, outputs_, loss_fn, inputs, self, aux=aux)