### Example Output: Solver and Initialization Progress in Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/GLV_JointEstimation_Alternate.ipynb

This snippet shows example output from the Python script that tests solvers and initializations. It indicates the solver and initialization types being processed and whether convergence issues were encountered. This output is useful for debugging and understanding the execution flow.

```text
solver_type='Dogleg' + init_type='zeros'...
with log_obs[:10] -> has convergence issues.

solver_type='Dogleg' + init_type='true'...
has convergence issues.



```

--------------------------------

### Train HyperPINN using jinns.solve

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

This snippet demonstrates how to initiate the training of a HyperPINN model using the `jinns.solve` method. It takes initial parameters, training data, and an optimizer as input, returning trained parameters and loss information.

```python
start = time.time()
params, total_loss_list, loss_by_term_dict, train_data, loss, _, _, _, _, _ = (
    jinns.solve(
        init_params=params,
        data=train_data,
        param_data=param_train_data,
        optimizer=tx,
        loss=loss,
        n_iter=n_iter,
    )
)
end = time.time()
```

--------------------------------

### Perform Sanity Check on Losses in Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Calculates and prints the total loss and individual component losses for the initial parameters. This is a sanity check to verify the magnitudes of the losses before training. It uses JAX for automatic differentiation.

```python
losses_and_grad = jax.value_and_grad(loss, 0, has_aux=True)
losses, grads = losses_and_grad(
    init_params,
    append_param_batch(train_data.get_batch()[1], param_train_data.get_batch()[1]),
)

# Sanity check
l_tot, d = losses
print(f"total loss: {l_tot:.2f}")
print(f"Individual losses: { {key: f'{val:.2f}' for key, val in d.items()} }")
```

--------------------------------

### Set Up Optimizer with Multiple Transformations (Python)

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/Tutorials/1D_non_stationary_Burgers_JointEstimation_Vanilla.ipynb

Configures the optimizer for training the PINN. This example demonstrates using `optax.multi_transform` to apply different optimizers (AdamW and Rprop) to different sets of parameters (neural network parameters and theta).

```python
params = init_params

param_labels = jinns.parameters.Params(
    nn_params="adamw",  # in this simple example we will use adamw for the nu
    eq_params={"theta": "rprop"},  # in this simple example we will use rprop for theta
)

tx_nu_and_theta = optax.multi_transform(
    {
        "adamw": optax.adamw(learning_rate=1e-5),
        "rprop": optax.rprop(1e-4),
    },  # all the keys in this dict must be found as leaves in the PyTree param_labels
    param_labels,
)

key, subkey = jax.random.split(key, 2)

```

--------------------------------

### Set Up Training Data for Reaction-Diffusion Model

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_heterogenous_model.ipynb

Initializes the data generator for a non-stationary partial differential equation (PDENonStatio) using Jinns library. This setup defines the domain, mesh size, batch sizes for different data types (domain, border, initial), and the time interval for training.

```python
n = 2048
nb = 500
ni = 500
domain_batch_size = 32
border_batch_size = 32
initial_batch_size = 32
dim = 2
xmin = 0
xmax = 1
ymin = 0
ymax = 1
tmin = 0
tmax = 1
method = "uniform"

Tmax = 2
key, subkey = random.split(key)
train_data = jinns.data.CubicMeshPDENonStatio(
    key=subkey,
    n=n,
    nb=nb,
    ni=ni,
    domain_batch_size=domain_batch_size,
    border_batch_size=border_batch_size,
    initial_batch_size=initial_batch_size,
    dim=dim,
    min_pts=(xmin, ymin),
    max_pts=(xmax, ymax),
    tmin=tmin,
    tmax=tmax,
    method=method,
)
```

--------------------------------

### Generating Parameter Grid for r and K

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Creates a grid of r and K values using jax.numpy.linspace and jax.numpy.meshgrid. The resulting R and K arrays are then flattened and stacked to form a list of (r, K) pairs, suitable for vectorized computations.

```python
r = jnp.linspace(0.05, 0.5, 50)
k = jnp.linspace(8, 15, 50)
R, K = jnp.meshgrid(r, k)
rk = jnp.stack([R.ravel(), K.ravel()], axis=1)
```

--------------------------------

### Initialize Neural Network Parameter Learning in Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Sets up the initial parameters, optimizer, and number of iterations for learning neural network parameters. It uses the AdamW optimizer with a specified learning rate and a fixed number of training iterations. The process focuses on updating only the neural network parameters.

```python
params = init_params
tx = optax.adamw(learning_rate=1e-3)
n_iter = 100000
```

--------------------------------

### Load Trained HyperPINN Model

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Loads a pre-trained HyperPINN model from a specified file path. This function is used to retrieve the trained neural network parameters and associated model configuration, enabling subsequent use of the trained model for predictions or analysis.

```python
u_trained, params = load_pinn(
    filename="./metamodel-logistic-100000iter", type_="hyperpinn"
)
```

--------------------------------

### Create HyperPINN Model Instance with Jinns

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Instantiates the HyperPINN model using the `jinns.nn.HyperPINN.create` function. This function takes the network architectures, ODE type, hyperparameter names, and hypernetwork input size to build the physics-informed hypernetwork model.

```python
u, init_nn_params = jinns.nn.HyperPINN.create(
    key=subkey,
    eqx_list=eqx_list,
    eq_type="ODE",
    hyperparams=hyperparams,
    hypernet_input_size=hypernet_input_size,
    eqx_list_hyper=eqx_list_hyper,
)
```

--------------------------------

### Setup Optimizer with Custom Updates for Joint Training

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/Tutorials/GLV_JointEstimation_Alternate.ipynb

Prepares the optimizer for joint training by defining parameter labels and configuring `optax.multi_transform`. This includes setting up AdamW for neural network parameters and custom chained optimizers with update and projection functions for 'g' and 'i' parameters.

```python
from functools import partial
from utils_GLV_JointEstimation_Alternate import update_and_project

update_and_project_g = partial(update_and_project, param="g")
update_and_project_i = partial(update_and_project, param="i")

param_labels = jinns.parameters.Params(
    nn_params="nn_adam",  # in this simple example we will use adamw for the nu
    eq_params={
        "g": "g_adam",
        "i": "i_adam",
    },  # in this simple example we will use rprop for theta
)

tx_nu_and_theta = optax.multi_transform(
    {
        "nn_adam": optax.adamw(learning_rate=1e-3),
        "g_adam": optax.chain(
            optax.adam(learning_rate=1e-3),
            update_and_project_g(),
        ),
        "i_adam": optax.chain(
            optax.adam(learning_rate=1e-6),
            update_and_project_i(),
        ),
    },  # all the keys in this dict must be found as leaves in the PyTree param_labels
    param_labels,
)

key, subkey = jax.random.split(key, 2)
n_iter_nu_and_theta = (n_iter * sum(jax.tree.leaves(n_iter_by_solver))) // (
    len(jax.tree.leaves(n_iter_by_solver))
)
```

--------------------------------

### Check JAX Device Configuration

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/Tutorials/tuto_GLV_inverse_problem.ipynb

This code snippet checks which device (CPU or GPU) JAX is configured to use. It includes a note that GPU support requires installing a CUDA-enabled jaxlib. The output indicates whether JAX is using the CPU or GPU.

```python
# Find out which device (GPU vs CPU) JAX is currently using for computations.
# For GPU, you need to install JAX+GPU prior to jinns
# see https://docs.jax.dev/en/latest/installation.html
jax.devices()
```

--------------------------------

### Log Posterior PDF Calculation and Timing (Python)

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Calculates and times the log posterior probability density function using different solver options. This snippet demonstrates performance comparison between solvers 1, 2, and 3, highlighting execution time differences.

```python
print(log_posterior_pdf(theta_init, Y, times, solver=1))
%timeit log_posterior_pdf(theta_init, Y ,times, solver = 1).block_until_ready()

print(log_posterior_pdf(theta_init, Y, times, solver=2))
%timeit log_posterior_pdf(theta_init, Y ,times, solver = 2).block_until_ready()

print(log_posterior_pdf(theta_init, Y, times, solver=3))
%timeit log_posterior_pdf(theta_init, Y ,times, solver = 3).block_until_ready()
```

--------------------------------

### Setup Burgers Equation Loss for SPINN - Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/1D_non_stationary_Burgers.ipynb

Configures the loss function for the Burgers equation using the SPINN. Similar to the PINN setup, but with potentially different loss weights to tune the training process for the SPINN architecture.

```python
import jinns.loss


loss_weights_spinn = jinns.loss.LossWeightsPDENonStatio(
    dyn_loss=1, initial_condition=10, boundary_loss=1
)

loss_spinn = jinns.loss.LossPDENonStatio(
    u=u_spinn,
    loss_weights=loss_weights_spinn,
    dynamic_loss=be_loss,
    boundary_condition=jinns.loss.Dirichlet(),
    initial_condition_fun=u0,
    params=init_params_spinn,
)
```

--------------------------------

### Initialize JAX and PRNG Key

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/2D_Heat_inverse_problem.ipynb

Initializes the JAX library for numerical computation and sets up a pseudo-random number generator (PRNG) key for reproducible randomness.

```python
import jax
from jax import random, vmap
import jax.numpy as jnp
import equinox as eqx
import matplotlib.pyplot as plt

key = random.PRNGKey(2)
key, subkey = random.split(key)
```

--------------------------------

### Set Training Parameters and Initialize Network

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/2D_non_stationary_OU.ipynb

Sets the number of training iterations and initializes the network parameters. These are essential configurations before starting the optimization process.

```python
n_iter = 25000
```

```python
params = init_params
```

--------------------------------

### Set Up Simulation Domain and Batch Sizes

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/2D_Navier_Stokes_PipeFlow_Metamodel_hyperpinn.ipynb

Configures parameters for the simulation domain and data batching. This includes the number of data points (n), batch sizes for training (omega_batch_size, omega_border_batch_size), and the spatial dimensions (dim) and boundaries (xmin, xmax, ymin, ymax) of the simulation area.

```python
n = 10000
b = None
omega_batch_size = 128
omega_border_batch_size = None
dim = 2
xmin = 0
xmax = xmin + L
ymin = -R
ymax = ymin + 2 * R
```

--------------------------------

### Initialize Model Parameters in Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_homogeneous_metamodel_hyperpinn_diffrax.ipynb

Sets up initial values for model parameters D, g, and r. These parameters are fundamental for defining the behavior of the physical system being modeled.

```python
D = 0.5
g = 1.0
r = 0.1
```

--------------------------------

### Measure and Print Training Time

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_homogeneous_metamodel_hyperpinn_diffrax.ipynb

Calculates the duration of the training process by subtracting the start time from the end time and prints the result. This helps in evaluating the computational cost of training.

```python
time_training = end - start
print("Time training the metamodel", time_training)
```

--------------------------------

### Calculate and Print Training Time

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/2D_Poisson_inverse_problem.ipynb

Calculates the total training time by subtracting the start time from the end time and prints the result. This helps in assessing the performance of the JINNS solver.

```python
time_training = end - start
print("Time training", time_training)
```

--------------------------------

### Initialize and Run Optimizer with Adam - Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_heterogenous_model.ipynb

This snippet demonstrates how to initialize the Adam optimizer from the optax library and use it to solve a problem with the jinns.solve function. It takes initial parameters, training data, the optimizer, and the loss function as input, returning updated parameters and training metrics.

```python
import optax

params = init_params

tx = optax.adam(learning_rate=1e-3)
n_iter = 50000

params, total_loss_list, loss_by_term_dict, train_data, loss, _, _, _, _, _, _, _ = (
    jinns.solve(
        init_params=params,
        data=train_data,
        optimizer=tx,
        loss=loss,
        n_iter=n_iter,
    )
)
```

--------------------------------

### Import necessary libraries for jinns

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/Tutorials/load_save_model.ipynb

Imports the jinns library and other essential tools like JAX for numerical computations and Equinox for neural network definitions. This setup is crucial for defining and training jinns models.

```python
%load_ext autoreload
%autoreload 2
%matplotlib inline

import jinns

import jax
from jax import random
import jax.numpy as jnp
import equinox as eqx
```

--------------------------------

### Initialize JAX and Equinox Environment

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/1D_Generalized_Lotka_Volterra.ipynb

Initializes JAX for numerical computation and Equinox for neural network definitions. It sets up the PRNG key for random number generation and splits it for subsequent uses.

```python
import jax
from jax import random, vmap
import jax.numpy as jnp
import equinox as eqx

import matplotlib.pyplot as plt

key = random.PRNGKey(2)
key, subkey = random.split(key)
```

--------------------------------

### Initialize DataGenerator for ODE Parameters

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Initializes a `DataGeneratorParameter` object to sample ODE parameters (r and K) within specified ranges. This is used for the inverse problem and for updating the initial parameters of the model.

```python
method = "uniform"
key, subkey = random.split(key)
np = 100000
param_batch_size = 1000
param_train_data = jinns.data.DataGeneratorParameter(
    key=subkey,
    n=np,
    param_batch_size=param_batch_size,
    param_ranges={"r": (0.05, 0.5), "K": (8, 15)},
    method=method,
)
```

--------------------------------

### Import JAX and Related Modules

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/1D_non_stationary_Fisher_KPP_Bounded_Domain.ipynb

Imports JAX, a high-performance numerical computation library, along with its random number generation, automatic vectorization (vmap), Equinox for neural network components, Optax for optimizers, and Matplotlib for plotting. These are fundamental for building and training the neural networks.

```python
import jax
from jax import random, vmap
import jax.numpy as jnp
import equinox as eqx
import optax
import matplotlib.pyplot as plt

key = random.PRNGKey(2)
```

--------------------------------

### Configure Adam Optimizer with Learning Rate Decay

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/1D_GLV_adaptative_loss_weights.ipynb

Sets up an Adam optimizer with a linearly decaying learning rate using `optax`. The scheduler defines the learning rate's behavior over time, starting at 1e-3 and decaying to 1e-4.

```python
import optax

start_learning_rate = 1e-3
# Exponential decay of the learning rate.
scheduler = optax.linear_schedule(
    init_value=start_learning_rate,
    transition_begin=20000,
    transition_steps=30000,
    end_value=1e-4,
)

# tx = optax.chain(
#     optax.scale_by_adam(),  # Use the updates from adam.
#     optax.scale_by_schedule(scheduler),  # Use the learning rate from the scheduler.
#     # Scale updates by -1 since optax.apply_updates is additive and we want to descend on the loss.
#     optax.scale(-1.0),
```

--------------------------------

### Initialize Adam Optimizer with Optax

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_heterogenous_model.ipynb

Initializes the Adam optimizer from the optax library with a specified learning rate. This is a common first step in setting up the optimization process for a model.

```python
import optax

tx = optax.adam(learning_rate=1e-3)
```

--------------------------------

### Initialize JAX and Utilities

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_homogeneous_metamodel_hyperpinn_diffrax.ipynb

Initializes JAX, a high-performance numerical computation library, along with its random number generation and NumPy interface. Equinox, a library for neural networks in JAX, and matplotlib for plotting are also imported. A JAX random key is generated for reproducibility.

```python
import jax
from jax import random
import jax.numpy as jnp
import equinox as eqx
import time
import matplotlib.pyplot as plt

key = random.PRNGKey(2)
```

--------------------------------

### Initialize DataGenerator for ODE Collocation Points

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Sets up a `DataGeneratorODE` object to generate collocation points for solving the ODE. It specifies the number of points, sampling method (uniform), time interval, and batch size for training.

```python
nt = 100000
batch_size = 1000
method = "uniform"
tmin = 0
tmax = 1

Tmax = 50
train_data = jinns.data.DataGeneratorODE(
    key=subkey,
    nt=nt,
    tmin=tmin,
    tmax=tmax,
    temporal_batch_size=batch_size,
    method=method,
)
```

--------------------------------

### Initialize AdamW Optimizer with Optax

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_homogeneous_metamodel_hyperpinn_diffrax.ipynb

Sets up the AdamW optimizer from the optax library with a specified learning rate. This is a common first step in configuring the optimization process for training neural networks.

```python
import optax

tx = optax.adamw(learning_rate=1e-3)
```

--------------------------------

### Check JAX Device Configuration

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/linear_fo_equation.ipynb

Determines the device (CPU or GPU) JAX is using for computations. It includes a note about installing JAX+GPU for CUDA support if a GPU is present. The output shows the currently active device.

```python
# Find out which device (GPU vs CPU) JAX is currently using for computations.
# For GPU, you need to install JAX+GPU prior to jinns
# see https://docs.jax.org/en/latest/installation.html
jax.devices()
```

--------------------------------

### Setup Burgers Equation Loss for PINN - Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/1D_non_stationary_Burgers.ipynb

Configures the loss function for the Burgers equation using the PINN. It defines the dynamic loss, loss weights (emphasizing initial conditions), and boundary/initial condition functions.

```python
be_loss = jinns.loss.BurgersEquation(Tmax=Tmax)
loss_weights = jinns.loss.LossWeightsPDENonStatio(
    dyn_loss=1, initial_condition=100, boundary_loss=1
)

loss_pinn = jinns.loss.LossPDENonStatio(
    u=u_pinn,
    loss_weights=loss_weights,
    dynamic_loss=be_loss,
    boundary_condition=jinns.loss.Dirichlet(),
    initial_condition_fun=u0,
    params=init_params_pinn,
)
```

--------------------------------

### Initialize and Update Equation Parameters in Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/2D_Navier_Stokes_PipeFlow_Metamodel_hyperpinn.ipynb

This snippet demonstrates initializing equation parameters for a JINNS model and updating them using batch data. It imports necessary functions and sets up initial parameters including neural network parameters and equation-specific values like density and viscosity.

```python
from jinns.parameters import update_eq_params

param_train_data, param_batch = param_train_data.get_batch()
init_params = jinns.parameters.Params(
    nn_params=u_p_init_nn_params,
    eq_params={"rho": rho, "nu": None},
)
init_params = update_eq_params(init_params, param_batch)
```

--------------------------------

### Configure ODE Solvers and Initial Parameters for Optimization

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/GLV_JointEstimation_Alternate.ipynb

Sets up various optimization solvers from Optimistix, including Levenberg-Marquardt and Dogleg, along with different initializations for the parameters to be optimized. It specifies solver tolerances and verbosity levels.

```python
import optimistix as optx
import diffrax as dfx  # https://github.com/patrick-kidger/diffrax
from jaxtyping import Array, Float  # https://github.com/google/jaxtyping
import time  # for timing

key, subkey = jax.random.split(key)

solvers = {
    # "Levenberg": optx.LevenbergMarquardt(
    #     rtol=1e-6,
    #     atol=1e-6,
    #     # verbose=frozenset({"step", "accepted", "loss", "step_size"}),
    # ),
    # Errors with small sample sizes
    "Dogleg": optx.Dogleg(
        rtol=1e-4,
        atol=1e-4,
        # verbose=frozenset({"step", "accepted", "loss", "step_size"}),
    ),
    # "GD": optx.OptaxMinimiser(optax.rprop(learning_rate=0.0001), rtol=1e-6, atol=1e-6),
    #       optx.GradientDescent(learning_rate=1e-5, rtol=1e-4, atol=1e-4),
}

init_thetas = {
    "zeros": {"g": jnp.zeros((Ns,)), "i": jnp.zeros((Ns, Ns))},
    "true": eq_params,
    # "random": 1e-1 * jax.random.normal(subkey, (Ns, Ns + 1)),
    # "close": theta + 1e-1 * jax.random.normal(subkey, (Ns, Ns + 1)),
}
```

--------------------------------

### Solve PDE with u0_alt3 Initial Condition using Jinns

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_heterogenous_model.ipynb

This code snippet demonstrates solving a PDE with the `u0_alt3` initial condition using the Jinns library. It sets up the loss function, optimizer, and training parameters, then proceeds to solve the system and visualize the resulting estimated solution over a range of times.

```python
loss = jinns.loss.LossPDENonStatio(
    u=u_spinn,
    loss_weights=loss_weights,
    dynamic_loss=fisher_dynamic_loss,
    boundary_condition=jinns.loss.Neumann(),
    initial_condition_fun=u0_alt3,
    params=init_params,
)
params = init_params
params, total_loss_list, loss_by_term_dict, train_data, loss, _, _, _, _, _, _, _ = (
    jinns.solve(
        init_params=params,
        data=train_data,
        optimizer=tx,
        loss=loss,
        n_iter=n_iter,
    )
)
u_est = lambda t_x: u_spinn(t_x, params)
jinns.plot.plot2d(
    fun=u_est,
    xy_data=val_xy_data,
    times=times,
    Tmax=Tmax,
    figsize=(10, 10 * len(times)),
    spinn=True,
)
```

--------------------------------

### Evaluate Initialized Neural Network

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/Tutorials/implementing_your_own_PDE_problem.ipynb

Demonstrates how to evaluate the initialized neural network. A lambda function `u_init` is created to wrap the network evaluation, taking a combined time-space input `t_x`. The example shows evaluating the network at the origin {0, 0}.

```python
u_init = lambda t_x: u(t_x, init_params)
u_init(jnp.zeros((2,)))  # u_init({0, 0})
```

--------------------------------

### Initialize JAX and Random Keys

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/2D_Poisson_inverse_problem.ipynb

Imports necessary JAX modules for automatic differentiation, numerical operations, and random number generation. It also initializes a pseudo-random number generator key and splits it for subsequent use.

```python
import jax
from jax import random, vmap
import jax.numpy as jnp
import equinox as eqx
import matplotlib.pyplot as plt
import time

key = random.PRNGKey(2)
key, subkey = random.split(key)
```

--------------------------------

### Define and Plot Initial Condition - Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/1D_non_stationary_Burgers.ipynb

Defines the initial condition function u0(x) for the Burgers equation and plots it using the training data's initial points. This visualizes the starting state of the system.

```python
def u0(x):
    return -jnp.sin(jnp.pi * x)


plt.plot(
    train_data.initial.sort(axis=0), vmap(u0, (0), 0)(train_data.initial.sort(axis=0))
)
plt.title(r"Init condition $u_0$")
```

--------------------------------

### Initialize and Evaluate PDE Loss Function

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/2D_Navier_Stokes_PipeFlow_SoftConstraints.ipynb

Sets up the main PDE loss function using initialized parameters, dynamic loss, boundary conditions, and loss weights. It then evaluates this loss function along with its gradients.

```python
loss = jinns.loss.LossPDEStatio(
    u=u_p,
    loss_weights=loss_weights,
    dynamic_loss=dyn_loss,
    boundary_condition=boundary_condition,
    params=init_params,
)

losses_and_grad = jax.value_and_grad(loss.evaluate, 0, has_aux=True)
key, subkey = random.split(key)
train_data, batch = train_data.get_batch()
losses, grads = losses_and_grad(init_params, batch=batch)

l_tot, dict_losses = losses
print(f"total loss: {l_tot:.2f}")
print(f"Individual losses: { {key: f'{val}' for key, val in dict_losses.items()} }")
```

--------------------------------

### Log-Normal PDF and Vectorized Version in Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Defines the logarithm of the Log-normal probability density function (PDF) and its vectorized version using JAX's `vmap`. This function is crucial for calculating the likelihood of observed data under a Log-normal distribution.

```python
@jax.jit
def log_lognorm_pdf(x, log_mu, sigma):
    """Logarithm of the Log-normal PDF"""
    return jnp.where(
        x > 0,
        -jnp.log(x * sigma * jnp.sqrt(2 * jnp.pi))
        - (jnp.log(x) - log_mu) ** 2 / (2 * sigma**2),
        -jnp.inf,
    )


## Vectorized version using vmap
Vectorized_log_lognorm_pdf = lambda X, log_Mu, fixed_sigma: jnp.sum(
    jax.vmap(log_lognorm_pdf, in_axes=(0, 0, None))(X, log_Mu, fixed_sigma)
)
```

--------------------------------

### Generate and Plot Noisy Observations in Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Creates synthetic noisy data by adding Gaussian noise to the exact solution of the ODE. It then visualizes these noisy observations alongside the true population sizes, illustrating the effect of noise on the data.

```python
## Synthetic data with N_0 = 1, r = 0.1, K = 10, sigma = 0.5
eps = jax.random.normal(key, shape=sizes.shape)
log_Y = jnp.log(sizes) + sigma * eps
Y = jnp.exp(log_Y)

plt.scatter(times, Y, s=3, label="Noisy observations")
plt.plot(times, sizes, "red", label="Real observations")
plt.legend(loc="upper left")
plt.title("Noisy observations of the exact solution")
plt.xlabel("time")
plt.ylabel("population size")
plt.show()
```

--------------------------------

### Update Initial Model Parameters with Sampled ODE Parameters

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Combines the initial neural network parameters with sampled ODE parameters (r and K) to create the initial parameter set for the model. The `update_eq_params` function is used for this purpose.

```python
init_params = jinns.parameters.Params(
    nn_params=init_nn_params, eq_params={"r": None, "K": None}
)
init_params = update_eq_params(init_params, param_train_data.get_batch()[1])
```

--------------------------------

### Solve PDE with u0_alt1 Initial Condition using Jinns

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_heterogenous_model.ipynb

Sets up and solves a partial differential equation using the Jinns library with the `u0_alt1` initial condition. It configures the loss function, optimizer, and training iterations, then solves the system and visualizes the estimated solution over time.

```python
loss = jinns.loss.LossPDENonStatio(
    u=u_spinn,
    loss_weights=loss_weights,
    dynamic_loss=fisher_dynamic_loss,
    boundary_condition=jinns.loss.Neumann(),
    initial_condition_fun=u0_alt1,
    params=init_params,
)
params = init_params
params, total_loss_list, loss_by_term_dict, train_data, loss, _, _, _, _, _, _, _ = (
    jinns.solve(
        init_params=params,
        data=train_data,
        optimizer=tx,
        loss=loss,
        n_iter=n_iter,
    )
)
u_est = lambda t_x: u_spinn(t_x, params)
jinns.plot.plot2d(
    fun=u_est,
    xy_data=val_xy_data,
    times=times,
    Tmax=Tmax,
    figsize=(10, 10 * len(times)),
    spinn=True,
)
```

--------------------------------

### Import JAX and Related Libraries - Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/Tutorials/GLV_JointEstimation_Alternate.ipynb

This code block imports essential libraries for numerical computation and machine learning using JAX. It includes `jax` for automatic differentiation and JIT compilation, `jax.numpy` for numerical operations similar to NumPy, `optax` for optimization algorithms, `equinox` for building neural network models, and `matplotlib.pyplot` for plotting. The commented-out lines show how to enable 64-bit precision in JAX if needed.

```python
import jax

# from jax import config
# config.update("jax_enable_x64", True)
import jax.numpy as jnp
import optax
import equinox as eqx
import matplotlib.pyplot as plt

figsize = (4, 4)
plt.rcParams["figure.figsize"] = figsize
```

--------------------------------

### Import JAX and Equinox for Scientific Computing

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/2D_Navier_Stokes_PipeFlow_Metamodel_hyperpinn.ipynb

Imports essential libraries for numerical computation and machine learning: JAX for automatic differentiation and GPU acceleration, and Equinox for building neural networks. Matplotlib is imported for plotting.

```python
import jax
from jax import random, vmap
import jax.numpy as jnp
import equinox as eqx
import matplotlib.pyplot as plt
```

--------------------------------

### Define and Plot Initial Condition u0

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_heterogenous_model.ipynb

Defines the initial condition function `u0(x)` for the PINN. It also sets up parameters for plotting this initial condition over a 2D spatial domain using `jinns.plot.plot2d`. This helps visualize the starting state of the system.

```python
sigma_init = 1 * jnp.ones((2))
mu_init = jnp.array([0.7, 0.15])

def u0(x):
    print((x - mu_init).shape)
    return jnp.exp(-jnp.linalg.norm(x - mu_init, axis=-1))

nx, ny = 100, 100
val_xy_data = [jnp.linspace(xmin, xmax, nx), jnp.linspace(ymin, ymax, ny)]

jinns.plot.plot2d(fun=u0, xy_data=val_xy_data, title="u0(x,y)", figsize=(5, 5))
```

--------------------------------

### Generate Training Data for SPINN

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/1D_non_stationary_Fisher_KPP_Bounded_Domain.ipynb

Generates training data specifically for the SPINN using `jinns.data.CubicMeshPDENonStatio`. This configuration allows for finer control over batch sizes for different types of data points (domain, border, initial) which can be beneficial for SPINN training.

```python
key, subkey = random.split(key)
train_data_spinn = jinns.data.CubicMeshPDENonStatio(
    key=subkey,
    n=n,
    nb=nb,
    ni=ni,
    domain_batch_size=32,
    border_batch_size=32,
    initial_batch_size=32,
    dim=dim,
    min_pts=(xmin,),
    max_pts=(xmax,),
    tmin=tmin,
    tmax=tmax,
    method=method,
)
```

--------------------------------

### Integrate Observation Data into Loss for Parameter Inference (Python)

Source: https://context7.com/mia-jinns/jinns/llms.txt

This snippet demonstrates how to use `jinns.data.DataGeneratorObservations` to create observation data and integrate it into the loss function for parameter inference. It shows the setup for simulated observations and the configuration of observation loss weights.

```python
import jax.numpy as jnp
import jax.random as random
import jinns

key = random.PRNGKey(42)

# Simulated observations: (time_points, observed_values)
obs_times = jnp.linspace(0.1, 1.0, 20)[:, None]  # Shape (20, 1)
obs_values = jnp.sin(obs_times)  # Shape (20, 1)

# Create observation data generator
obs_data = jinns.data.DataGeneratorObservations(
    key=key,
    obs_batch_size=10,
    observed_pinn_in=obs_times,      # PINN input locations
    observed_values=obs_values,       # Target values
    observed_eq_params=None,          # Optional: varying eq params per observation
)

# Loss with observations
loss_weights = jinns.loss.LossWeightsODE(
    dyn_loss=jnp.array(1.0),
    initial_condition=jnp.array(1.0),
    observations=jnp.array(10.0),  # Weight for observation loss
)

loss = jinns.loss.LossODE(
    u=u,
    dynamic_loss=dynamic_loss,
    initial_condition=(t0, u0),
    loss_weights=loss_weights,
    params=params,
)

# Train with observations
params, *_ = jinns.solve(
    n_iter=5000,
    init_params=params,
    data=train_data,
    loss=loss,
    optimizer=optimizer,
    obs_data=obs_data,  # Include observations
    key=key,
)
```

--------------------------------

### Generate and Plot Exact ODE Solution in Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Generates a time series and calculates the corresponding population sizes using the exact solution. It then plots the results, showing the population size over time. This helps visualize the theoretical growth curve.

```python
nt = 100  # 1000
Tmax = 50  # 100
times = jnp.linspace(0, Tmax, nt)  # t1, ...,tI
sizes = N_exact(times, beta=(r, K))  # N(t1), ...,N(tI)

assert len(times) == len(sizes)
plt.plot(times, sizes, "red")
plt.xlabel("time")
plt.ylabel("Population size")
plt.title("Exact solution of the ODE")
plt.show()
```

--------------------------------

### Define Loss Weights and Total Loss (Python)

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_heterogenous_model.ipynb

This snippet shows how to define the weights for different components of the loss function (dynamic, initial condition, boundary loss) using `LossWeightsPDENonStatio`. It then constructs the total loss using `LossPDENonStatio`, incorporating the dynamic loss, boundary conditions, initial condition function, and parameters.

```python
loss_weights = jinns.loss.LossWeightsPDENonStatio(
    dyn_loss=1,
    initial_condition=1 * Tmax,
    boundary_loss=1 * Tmax
)

loss = jinns.loss.LossPDENonStatio(
    u=u,
    loss_weights=loss_weights,
    dynamic_loss=fisher_dynamic_loss,
    boundary_condition=jinns.loss.Neumann(),
    initial_condition_fun=u0,
    params=init_params,
)
```

--------------------------------

### Define Logistic ODE Loss Function using Jinns

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Defines a custom ODE loss class, `Logistic_reparametrized`, inheriting from `jinns.loss.ODE`. This class implements the equation for the Verhulst model, calculating the time derivative and the ODE residual for loss computation.

```python
class Logistic_reparametrized(ODE):
    """
    Return the Logistic loss term.
    """

    def equation(self, t, u, params):
        u_ = lambda t: u(t, params)[0]
        du_dt = jax.grad(u_)(t)[0]
        return du_dt + self.Tmax * (
            -params.eq_params.r * u(t, params) * (1 - u(t, params) / params.eq_params.K)
        )


Logistic_dynamic_loss = Logistic_reparametrized(Tmax=Tmax)
```

--------------------------------

### Initialize PINN and SPINN Parameters

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/1D_non_stationary_Fisher_KPP_Bounded_Domain.ipynb

Initializes parameter objects for both Physics-Informed Neural Networks (PINNs) and Spectral Physics-Informed Neural Networks (SPINNs). These objects hold neural network parameters and equation-specific parameters.

```python
init_params_pinn = jinns.parameters.Params(
    nn_params=init_nn_params_pinn,
    eq_params={"D": jnp.array([D]), "r": jnp.array([r]), "g": jnp.array([g])},
)
init_params_spinn = jinns.parameters.Params(
    nn_params=init_nn_params_spinn,
    eq_params={"D": jnp.array([D]), "r": jnp.array([r]), "g": jnp.array([g])},
)
```

--------------------------------

### Import Jax and Jinns Libraries for Mechanistic-Statistical Modeling

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Imports necessary libraries from JAX, equinox, matplotlib, seaborn, pandas, time, and jinns. These libraries are essential for numerical computations, neural network definitions, plotting, and the physics-informed hypernetwork framework.

```python
# Jax utils
import jax
from jax import random, lax
import jax.numpy as jnp
from functools import partial
from typing import Tuple
from jaxtyping import Array, Float
import optax

# For the neuronal architectures
import equinox as eqx

import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import time
from tqdm import trange

# The jinns library
import jinns
from jinns.loss._DynamicLossAbstract import ODE
from jinns.data import append_param_batch
from jinns.parameters import update_eq_params
from jinns.nn._save_load import load_pinn, save_pinn

# Finite Differences Jax based solver
from diffrax import diffeqsolve, Dopri5, ODETerm, SaveAt, PIDController

# Initialize random key
key = jax.random.PRNGKey(48)
key, subkey = random.split(key)
```

--------------------------------

### Example Usage: Parameter Error Calculation (Python)

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/Tutorials/GLV_JointEstimation_Alternate.ipynb

Demonstrates the usage of the `errors_on_parameters` function by calculating RMSE for different sets of estimated parameters against true parameters. It initializes `EqParams` and then calls the error function for various configurations.

```python
from jinns.parameters import EqParams

true_params = EqParams({"g": growth_rates, "i": interactions})

rmse_parameters_alternate = errors_on_parameters(true_params, params_alternate)
rmse_parameters_alternate_regularized = errors_on_parameters(
    true_params, params_alternate_regularized
)
rmse_parameters_vanilla = errors_on_parameters(true_params, params_vanilla)
rmse_parameters_vanilla_regularized = errors_on_parameters(
    true_params, params_vanilla_regularized
)

print(f"{rmse_parameters_alternate=}")
print(f"{rmse_parameters_alternate_regularized=}")
print(f"{rmse_parameters_vanilla=}")
print(f"{rmse_parameters_vanilla_regularized=}")
```

--------------------------------

### Define and Plot Initial Condition

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/1D_non_stationary_Fisher_KPP_Bounded_Domain.ipynb

Sets up the initial condition function `u0` for the Fisher-KPP equation using a Gaussian distribution from `jax.scipy.stats.norm`. It then plots this initial condition over the spatial domain defined by `train_data` to visualize the starting state of the population.

```python
from jax.scipy.stats import norm

# true solution N(0,1)
sigma_init = 0.2 * jnp.ones((1))
mu_init = 0 * jnp.ones((1))


def u0(x):
    return jnp.squeeze(norm.pdf(x, loc=mu_init, scale=sigma_init))


plt.plot(
    train_data.domain[:, 1:2].sort(axis=0),
    vmap(u0, (0), 0)(train_data.domain[:, 1:2].sort(axis=0)),
)
plt.title(r"Init condition $u_0$")
```

--------------------------------

### Visualize Posterior Samples Trace Plots

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/MS_model_Verhulst.ipynb

Generates trace plots to visualize the convergence of Markov chain samples for population dynamics parameters. It compares samples obtained from a diffrax solver against those from a HyperPINN model. Dependencies include matplotlib and the posterior samples.

```python
# samples.shape == (num_samples, 4)
param_names = ["$r$", "$K$", "$\sigma$"]

fig, axs = plt.subplots(3, 1, figsize=(10, 8), sharex=True)
for i in range(3):
    axs[i].plot(samples_Vectorised[:, i], label="diffrax solver")
    axs[i].plot(samples_Metamodel[:, i], label="HyperPINN", linestyle="dashed")
    axs[i].set_ylabel(param_names[i])
    axs[i].legend()
axs[-1].set_xlabel("Iteration")
fig.suptitle("Trace Plots of Posterior Samples")
plt.tight_layout()
plt.show()
```

--------------------------------

### Import JAX and Jinns Libraries

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/1D_GLV_adaptative_loss_weights.ipynb

Imports necessary libraries for JAX, Equinox, and jinns, including random number generation and neural network components. This is a foundational step for setting up the PINN.

```python
import jinns
import jax
from jax import random, vmap
import jax.numpy as jnp
import equinox as eqx

import matplotlib.pyplot as plt

key = random.PRNGKey(2)
key, subkey = random.split(key)
```

--------------------------------

### Iterate Solvers and Initializations for Optimization in Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/GLV_JointEstimation_Alternate.ipynb

This Python code iterates through different solver types and initialization types to perform optimization. It uses `optx.least_squares` for the optimization process and prints progress and results. Dependencies include `time`, `optx`, and `jax`. Inputs are dictionaries of solvers and initializations, and outputs are stored solutions and performance metrics.

```python
sols = {}

for solver_type in solvers.keys():
    sols[solver_type] = {}
    for init_type in init_thetas.keys():
        print(f"{solver_type=} + {init_type=}...")
        solver = solvers[solver_type]
        init_parameters = init_thetas[init_type]

        try:
            start = time.time()

            init_parameters_ = init_parameters.copy()
            for idx_end in idx_ends:
                print(f"with log_obs[:{idx_end}] -> ", end="")
                mb_lobs = log_obs_subsample_noisy[:idx_end]
                keep_ts = t_subsample[:idx_end]

                residuals_ = partial(residuals_mb, save_t=keep_ts)
                sol_ = optx.least_squares(
                    residuals_,
                    solver,
                    init_parameters_,
                    args=(y0, mb_lobs),
                    max_steps=MAX_STEPS,
                )

                print(
                    f"{solver_type} converged in {sol_.stats['num_steps']} iterations."
                )

            end = time.time()
            sols[solver_type][init_type] = sol_  # store
            print(f"Total time {end - start:.2f} s.")

            max_norm = optx.max_norm(
                residuals_(sol_.value, (y0, log_obs_subsample_noisy))
            )
            mse = jax.tree.map(
                lambda a, b: jnp.sqrt(jnp.mean((a - b) ** 2)), eq_params, sol_.value
            )
            print(f"max_residual={max_norm}, MSE on theta={mse}\n")
        except:
            print("has convergence issues.\n")
```

--------------------------------

### Example Usage: Curve Error Calculation (Python)

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/Tutorials/GLV_JointEstimation_Alternate.ipynb

Demonstrates calculating curve errors using the `error_on_the_curves` function. It first prepares estimated functions for different parameter sets and then calls the error calculation function, storing the results.

```python
trained_u_est_alternate = jax.vmap(lambda t_x: jnp.exp(u(t_x, params_alternate)))
trained_u_est_alternate_regularized = jax.vmap(
    lambda t_x: jnp.exp(u(t_x, params_alternate_regularized))
)
trained_u_est_vanilla = jax.vmap(lambda t_x: jnp.exp(u(t_x, params_vanilla)))
trained_u_est_vanilla_regularized = jax.vmap(
    lambda t_x: jnp.exp(u(t_x, params_vanilla_regularized))
)


rmse_curves_alternate, rec_sol_log_alternate = error_on_the_curves(
    y, trained_u_est_alternate, params_alternate
)
rmse_curves_alternate_regularized, rec_sol_log_alternate_regularized = (
    error_on_the_curves(
        y, trained_u_est_alternate_regularized, params_alternate_regularized
    )
)
rmse_curves_vanilla, rec_sol_log_vanilla = error_on_the_curves(
    y, trained_u_est_vanilla, params_vanilla
)
rmse_curves_vanilla_regularized, rec_sol_log_vanilla_regularized = error_on_the_curves(
    y, trained_u_est_vanilla_regularized, params_vanilla_regularized
)

print(f"{rmse_curves_alternate=}")
print(f"{rmse_curves_alternate_regularized=}")
print(f"{rmse_curves_vanilla=}")
print(f"{rmse_curves_vanilla_regularized=}")
```

--------------------------------

### Initialize FisherKPP Loss Function in Python

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/Reaction_Diffusion_2D_heterogenous_model.ipynb

Sets up the `FisherKPP` loss function for a non-stationary problem. It configures parameters like maximum time (Tmax), heterogeneous equation parameters, dimensionality, and initial parameters.

```python
fisher_dynamic_loss = jinns.loss.FisherKPP(
    Tmax=Tmax,
    eq_params_heterogeneity={"D": None, "r": r_fun, "g": None, "beta": None},
    dim_x=2,
    params=init_params,
)
```

--------------------------------

### Configure Loss Function Weights and Initialization

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/ODE/linear_fo_equation.ipynb

Sets up the weights for different components of the loss function (dynamic loss and initial condition). It then initializes the complete loss function, combining the neural network, loss weights, the defined ODE loss, initial conditions, and parameters.

```python
loss_weights = jinns.loss.LossWeightsODE(dyn_loss=1.0, initial_condition=1.0)

loss = jinns.loss.LossODE(
    u=u,
    loss_weights=loss_weights,
    dynamic_loss=fo_loss,
    initial_condition=(float(tmin), jnp.log(x0)),
    params=init_params,
)
```

--------------------------------

### Assign Initial Parameters for HyperPINN Validation (Python)

Source: https://github.com/mia-jinns/jinns/blob/main/Notebooks/PDE/2D_Navier_Stokes_PipeFlow_Metamodel_hyperpinn.ipynb

This is a simple assignment statement that sets the `params_hyper` variable to the `init_params_hyper`. This is typically done before starting a validation or training loop to ensure the validation process uses the correct initial model parameters.

```python
params_hyper = init_params_hyper
```