Using one model to predict a parameter and then solving the equation

[camillae00]

Suppose I have a dataset containing measurements of $\theta$. I know that the system which I have measured can be described by the following ODE:

$$ A \dot{\theta} = k(t) \Big[\theta - B\Big] + C, $$

where $A$, $B$ and $C$ are constants (I have simplified it for the sake of clarity).

In the dataset I have $n$ fixed-length time series containing measurements of $\theta$, which are all described by the aforementioned ODE, but $k(t)$ is different for each of the time series.

$k(t)$ depends on the history of the system and can be predicted from other measurement data, let us call this measurement data $\mathbf{x}_n(t)$, where $n$ indicates that $\mathbf{x}_n(t)$ is different for each of the $n$ time series.

Now, I could train a "standard" PINN-model with DeepXDE for all $n$ time series individually. I have tried this and it works well, but then I am left with $n$ different models, and in my case $n$ is on the order of thousands, so not very practical.

Is there a way to do this by only training one large model which after training ca be used as model.predict(x_n, t) and then output $\theta_n$?

I hope I managed to state my problem somewhat clearly.

[lululxvi]

You can treat k(t) as another state variable, but I am not very sure.