The trial solutions to additional_loss function are different depending on parameters

[affans]

Describe the bug 🐞

The PhysicsInformedNN discretizer has two arguments param_estim and additional_loss for parameter estimation. The additional_loss is a function which accepts three arguments, i.e., additional_loss(phi, θ, p) where θ is a trial solution.

When param_estim = true, the trial solution θ in the additional loss function is a ComponentArray and the dependent variables (suppose x, y z) can be accessed as θ[:x], θ[:y], θ[:z].

When param_estim = false, where we are not interested in parameter estimation but still want the model to fit to additional data, the trial solution θ is still a ComponentArray but the keys are different. That is, the dependent variables are accessed as θ.depvar[:x], θ.depvar[:y], θ.depvar[:z].

Is there a particular reason for why this is?

The docstring also seems incorrect to me. The docs say

additional_loss: a function additional_loss(phi, θ, p_) where phi are the neural network trial solutions, θ are the weights of the neural network(s), and p_ are the hyperparameters of the OptimizationProblem. If param_estim = true, then θ additionally contains the parameters of the differential equation appended to the end of the vector.

but this is not true. The parameters are not appended to the trial solution but passed in as a third argument p. Do the docs need to be updated here?

Expected behavior

I would expect the trial solution θ to still have the same structure (i.e., access to the dependent variables) regardless of the param_estim argument.

Output of versioninfo()

⌃ [b2108857] Lux v0.5.42
  [961ee093] ModelingToolkit v9.13.0
  [315f7962] NeuralPDE v5.16.0