%%{init: {"theme": "default", "themeVariables": {"fontSize": "13px", "primaryColor": "#2d6a4f", "lineColor": "#52b788"}, "flowchart": {"nodeSpacing": 15, "rankSpacing": 40, "curve": "basis"}}}%%
flowchart TD
subgraph Input["Input Layer (1)"]
x1(["x"])
end
subgraph H1["Hidden Layer (2)"]
h1(["h1"])
h2(["h2"])
end
subgraph A1["Tanh"]
t1(["tanh"])
end
subgraph Out["Output Layer (1)"]
o1(["o1"])
end
yhat(["y-hat"])
x1 -->|"W11=0.5"| h1
x1 -->|"W12=-1.2"| h2
h1 & h2 --> t1
t1 --> o1
o1 --> yhat
style x1 fill:#52b788,stroke:#2d6a4f,color:#0d0900
style h1 fill:#e6a817,stroke:#c98f0a,color:#0d0900
style h2 fill:#e6a817,stroke:#c98f0a,color:#0d0900
style t1 fill:#a2d2ff,stroke:#6da9e4,color:#0d0900
style o1 fill:#e6a817,stroke:#c98f0a,color:#0d0900
style yhat fill:#52b788,stroke:#2d6a4f,color:#0d0900
style Input fill:none,stroke:#2d6a4f,color:#2d6a4f
style H1 fill:none,stroke:#c98f0a,color:#c98f0a
style A1 fill:none,stroke:#6da9e4,color:#6da9e4
style Out fill:none,stroke:#c98f0a,color:#c98f0a
Introduction
A neural network takes an input, passes it through layers of simple math operations, and produces an output. Each layer multiplies the input by a set of weights, adds a bias, and applies an activation function. The weights and biases are the learnable parameters, where they start random and get tuned during training.
In this notebook we walk through every step by hand, starting with the smallest possible network: one input, two hidden neurons, and one output. Once you see how the numbers flow through a single forward pass, the remaning steps of training, loss functions, and physics-informed constraints will feel much more concrete.
How to Compute the Forward Pass of a Neural Network by Hand?
Forward-pass steps for a single input x = 3.9000
Step 1: Pre-activation (linear transform in hidden layer)
\begin{aligned} z_1 &= x \, W_1 + b_1 \\ &= \begin{bmatrix} 3.9000 \end{bmatrix} \begin{bmatrix} 0.5000 & -1.2000 \end{bmatrix} + \begin{bmatrix} 0.1000 & 0.3000 \end{bmatrix} \\ &= \begin{bmatrix} 1.9500 & -4.6800 \end{bmatrix} + \begin{bmatrix} 0.1000 & 0.3000 \end{bmatrix} \\ &= \begin{bmatrix} 2.0500 & -4.3800 \end{bmatrix} \end{aligned}
Step 2: Activation (non-linearity)
\begin{aligned} a_1 &= \tanh(z_1) \\ &= \tanh\! \begin{bmatrix} 2.0500 & -4.3800 \end{bmatrix} \\ &= \begin{bmatrix} \tanh(2.0500) & \tanh(-4.3800) \end{bmatrix} \\ &= \begin{bmatrix} 0.9674 & -0.9997 \end{bmatrix} \end{aligned}
The element-wise \tanh squashes each pre-activation value into (-1, 1), introducing the non-linearity that lets the network approximate complex functions.
Step 3: Output (linear transform, no activation)
\begin{aligned} \hat{y} &= a_1 \, W_2 + b_2 \\ &= \begin{bmatrix} 0.9674 & -0.9997 \end{bmatrix} \begin{bmatrix} 0.8000 \\ -0.6000 \end{bmatrix} + \begin{bmatrix} 0.0000 \end{bmatrix} \\ &= \begin{bmatrix} (0.9674)(0.8000) + (-0.9997)(-0.6000) \end{bmatrix} + \begin{bmatrix} 0.0000 \end{bmatrix} \\ &= \begin{bmatrix} 0.7739 + 0.5998 \end{bmatrix} \\ &= \begin{bmatrix} 1.3737 \end{bmatrix} \end{aligned}
The hidden activations are linearly combined to produce the final scalar prediction. No activation is applied at the output layer, so \hat{y} \in \mathbb{R}.
NoteWhat is a hyperbolic tangent?
The hyperbolic tangent (\tanh) is a mathematical function defined as the ratio of hyperbolic sine (\sinh) to hyperbolic cosine (\cosh), calculated as
\tanh = \frac{e^x - e^{-x}}{e^x + e^{-x}}
It maps any real number input to an S-shaped curve with outputs strictly between -1 and 1, which is commonly used in neural networks to normalize data.
Implementation in Numpy
# A simple feedforward neural network example to illustrate the forward pass
# Architecture: 1 input → 2 hidden neurons (Tanh) → 1 output
# Weights and biases
# normally these are learned parameters but we fix them manually here
# shape (1, 2): one input, two hidden neurons
W1 = np.array([[0.5, -1.2]])
print(f"W1:\n{W1}")
print(f"W1 shape: {W1.shape}")
print() # blank line for readability
# bias for each hidden neuron
b1 = np.array([0.1, 0.3])
print(f"b1:\n{b1}")
print(f"b1 shape: {b1.shape}")
print()
# shape (2, 1): two hidden neurons, one output
W2 = np.array([[0.8], [-0.6]])
print(f"W2:\n{W2}")
print(f"W2 shape: {W2.shape}")
print()
# bias for output neuron
b2 = np.array([0.0])
print(f"b2:\n{b2}")
print(f"b2 shape: {b2.shape}")
print()
# Forward pass for a single input x = 3.9
x = np.array([[3.9]]) # shape (1, 1)
print(f"x:\n{x}")
print(f"x shape: {x.shape}")
print()
# Step 1: linear transformation in hidden layer
# (1,1)@(1,2) + (2,) = (1,2) pre-activation
z1 = x @ W1 + b1
print(f"Pre-activation (z1): {z1}")
# Step 2: apply Tanh activation
# element-wise Tanh
a1 = np.tanh(z1)
print(f"Post-activation (a1): {a1}")
# Step 3: linear transformation to output (no activation on output layer)
y_hat = a1 @ W2 + b2
print(f"Network output: {y_hat[0,0]:.4f}")W1:
[[ 0.5 -1.2]]
W1 shape: (1, 2)
b1:
[0.1 0.3]
b1 shape: (2,)
W2:
[[ 0.8]
[-0.6]]
W2 shape: (2, 1)
b2:
[0.]
b2 shape: (1,)
x:
[[3.9]]
x shape: (1, 1)
Pre-activation (z1): [[ 2.05 -4.38]]
Post-activation (a1): [[ 0.967395 -0.99968628]]
Network output: 1.3737Implementation in PyTorch
# Same architecture: 1 input → 2 hidden neurons (Tanh) → 1 output
# Weights and biases (same values as the NumPy version)
W1 = torch.tensor([[0.5, -1.2]]) # shape (1, 2)
b1 = torch.tensor([0.1, 0.3]) # shape (2,)
W2 = torch.tensor([[0.8], [-0.6]]) # shape (2, 1)
b2 = torch.tensor([0.0]) # shape (1,)
# Forward pass for a single input x = 3.9
x = torch.tensor([[3.9]]) # shape (1, 1)
# Step 1: linear transformation in hidden layer
z1 = x @ W1 + b1
print(f"Pre-activation (z1): {z1}")
# Step 2: apply Tanh activation
a1 = torch.tanh(z1)
print(f"Post-activation (a1): {a1}")
# Step 3: linear transformation to output (no activation on output layer)
y_hat = a1 @ W2 + b2
print(f"Network output: {y_hat.item():.4f}")Pre-activation (z1): tensor([[ 2.0500, -4.3800]])
Post-activation (a1): tensor([[ 0.9674, -0.9997]])
Network output: 1.3737
NoteWhat does
.item() do?
.item() is a PyTorch tensor method that extracts the single element from a tensor as a plain Python scalar (float or int). It only works on tensors containing exactly one element; calling it on a multi-element tensor raises a ValueError.
y = torch.tensor([0.8415])
y.item() # → 0.8414709848078965 (a Python float)
z = torch.tensor([1, 2, 3])
z.item() # → ValueError: only one element tensors can be convertedIn the code above, y_hat has shape (1, 1), which is a single scalar wrapped in a 2D tensor, so y_hat.item() unwraps it to a Python float for clean printing with :.4f formatting.
What is Physics-Informed Neural Network?
A Physics-Informed Neural Network (PINN) is a neural network that learns to solve differential equations by embedding the governing physics directly into its loss function. In simple terms:
- It is still a neural network. The architecture (layers, weights, activations) is exactly the same as any standard feedforward network. It takes inputs (spatial coordinates x_1, \dots, x_n and/or time t) and produces an output u_{\text{NN}}.
- The key difference is in the loss function. Instead of minimising the error against labelled data alone, a PINN minimises how badly its output violates a known physical law (an ODE or PDE).
- Little or no data is needed. Because the physics equation itself provides the supervision signal, a PINN can find the solution without ever seeing measured data.
- Derivatives come from automatic differentiation. PyTorch’s
autogradcomputes exact partial derivatives of the network’s output with respect to its inputs (\partial u / \partial x_1, \dots, \partial u / \partial t), which are then substituted into the governing equation.
The total loss is a weighted sum of up to four terms:
\mathcal{L}_{\text{Total}} = \lambda_1 \mathcal{L}_{\text{Data}} + \lambda_2 \mathcal{L}_{\text{DE}} + \lambda_3 \mathcal{L}_{\text{IC}} + \lambda_4 \mathcal{L}_{\text{BC}}
| Loss term | What it measures |
|---|---|
| \mathcal{L}_{\text{Data}} | Data loss: mismatch between the network’s prediction and any available measured/labelled data. This term is optional; a PINN can work with no data at all. |
| \mathcal{L}_{\text{DE}} | Differential equation loss: the ODE/PDE residual evaluated at collocation points. If the network perfectly satisfies the governing equation, this term is zero. |
| \mathcal{L}_{\text{IC}} | Initial condition loss: error at t = 0 (e.g. x(0) = 1, \dot{x}(0) = 0). Anchors the solution to its known starting state. |
| \mathcal{L}_{\text{BC}} | Boundary condition loss: error at the spatial boundaries of the domain (e.g. u = 0 at a wall). For time-only ODEs this may be absent. |
The weights \lambda_1, \lambda_2, \lambda_3, \lambda_4 control the relative importance of each term. The optimiser drives all active terms toward zero simultaneously, forcing the network to output a function that satisfies the equation everywhere, matches the boundary/initial values, and (if available) agrees with measured data.
Standard ML vs PINNs
The following table summarizes the similarities and differences of standard ML versus PINNs.
| Aspect | Standard Supervised ML | Physics-Informed Neural Network (PINN) |
|---|---|---|
| Training data | Requires labelled input-output pairs (x_i, y_i) | Can work with little or no labelled data |
| Loss function | Measures prediction error against labels: \mathcal{L} = \sum (f_\theta(x_i) - y_i)^2 | Measures how well the output satisfies a differential equation (PDE/ODE residual) |
| What is learned | A mapping from inputs to outputs that fits the data | A mapping whose output obeys the governing physics |
| Training loop | Forward pass → compute loss → backprop → update weights | Identical loop, only the loss term changes |
| Generalisation | Interpolates well within the training data distribution | Extrapolates better because physics constrains the solution space |
Standard ML Training Example: Fitting \sin(2\pi t)
The cell below demonstrates a complete standard ML training loop. Here, the network learns to approximate f(t) = \sin(2\pi t) purely from labelled data without involving any physics.
Network architecture:
1 \;\text{input} \;\xrightarrow{} \;16 \;\text{neurons (Tanh)} \;\xrightarrow{} \;16 \;\text{neurons (Tanh)} \;\xrightarrow{} \;1 \;\text{output}
| Layer | Code | Role |
|---|---|---|
| Hidden layer 1 | nn.Linear(1, 16) + nn.Tanh() |
16 neurons with Tanh activation |
| Hidden layer 2 | nn.Linear(16, 16) + nn.Tanh() |
16 neurons with Tanh activation |
| Output layer | nn.Linear(16, 1) |
1 output, no activation |
This is a deeper network than our hand-calculated example: two hidden layers of 16 neurons each, giving it enough capacity to capture the sine wave’s shape.
Training setup:
- Data: 100 evenly spaced points t \in [0, 1] with labels y = \sin(2\pi t)
- Loss function: Mean Squared Error (MSE) \mathcal{L} = \frac{1}{N}\sum_{i=1}^{N}(\hat{y}_i - y_i)^2
- Optimiser: Adam with learning rate \alpha = 10^{-3}
- Epochs: 3000 passes through the full dataset
What happens at each epoch:
optimiser.zero_grad(): reset all gradients to zero (PyTorch accumulates gradients by default)y_pred = simple_nn(t_train): forward pass that pushes all 100 inputs through the networkloss = loss_fn(y_pred, y_train): compute MSE between predictions and true labelsloss.backward(): backpropagation that computes \nabla_\theta \mathcal{L} for every weightoptimiser.step(): update weights such that \theta \leftarrow \theta - \alpha \cdot \nabla_\theta \mathcal{L}
Outputs:
- Left-plot: The network’s prediction overlaid on the true sine curve
- Right-plot: Training loss over epochs on a log scale, showing convergence.
torch.manual_seed(42)
# A tiny network:
# 1 input → 16 neurons (Tanh) → 16 neurons (Tanh) → 1 output
# 2 hidden layers and 1 output layer
simple_nn = nn.Sequential(
nn.Linear(1, 16), nn.Tanh(),
nn.Linear(16, 16), nn.Tanh(),
nn.Linear(16, 1)
)
# Training data: 100 points uniformly in [0, 1]
t_train = torch.linspace(0, 1, 100).unsqueeze(1) # shape (100, 1)
y_train = torch.sin(2 * torch.pi * t_train) # true values
# We use the Adam optimiser and mean squared error (MSE) loss
# lr is the learning rate, which controls how big the weight updates are
# The optimiser will adjust the weights of simple_nn to minimize the loss
optimiser = torch.optim.Adam(
simple_nn.parameters(),
lr=1e-3)
loss_fn = nn.MSELoss()
# Training loop
losses = []
for epoch in range(3000):
optimiser.zero_grad()
y_pred = simple_nn(t_train)
loss = loss_fn(y_pred, y_train)
loss.backward() # compute gradients
optimiser.step() # update weights
losses.append(loss.item())
# Plot
fig, (ax1, ax2) = plt.subplots(
1, 2,
figsize=(9, 4))
ax1.grid(alpha=0.3)
ax1.plot(
t_train.numpy(),
y_train.numpy(),
c="black",
label=r"True $\sin(2\pi t)$",
lw=2)
ax1.plot(
t_train.detach().numpy(),
simple_nn(t_train).detach().numpy(),
"--",
c="red",
label=r"Network prediction",
lw=2)
ax1.set_xlabel(
r"Input $t$",
fontsize=12,
)
ax1.set_ylabel(
r"Output $\sin(2\pi t)$",
fontsize=12,
)
ax1.legend()
ax1.set_title(
r"Neural network fit of $\sin(2\pi t)$",
fontsize=14)
ax2.grid(alpha=0.3)
ax2.semilogy(
losses,
"--",
c="red",
label=r"Final training loss: {:.2e}".format(losses[-1]))
ax2.set_xlabel(
r"Epoch",
fontsize=12)
ax2.set_ylabel(
r"MSE Loss (log scale)",
fontsize=12)
ax2.set_title(
r"Training loss",
fontsize=14)
ax2.legend()
plt.tight_layout()
plt.show()
ImportantWhy do we use
t_train.detach().numpy() for plotting?
.detach()removes the tensor from PyTorch’s computation graph (so no gradients will be tracked), and.numpy()converts the resulting tensor to a NumPy array. Matplotlib requires NumPy arrays, not PyTorch tensors, so this chain is needed for plotting.In this specific case,
t_traindoesn’t actually haverequires_grad=True, so.detach()is technically unnecessary.numpy()alone would work. However, calling.detach().numpy()is a safe habit because PyTorch will raise an error if we call.numpy()on a tensor that requires gradients without detaching first.
Automatic Differentiation for PINNs
The fundamental question: if the network predicts \hat{x}(t) (a displacement), how do we get \dot{\hat{x}}(t) and \ddot{\hat{x}}(t) (velocity and acceleration)?
We need these because the equation of motion involves them. For example, for a spring-mass system: m\ddot{x} + kx = 0 We cannot check whether the network satisfies this equation without computing \ddot{x}.
Automatic differentiation (AD) solves this exactly. It computes the exact derivative of any computation with respect to any input (not numerically and not approximately, but exactly). This is because the network is just a chain of arithmetic operations, and the derivative of a chain is computable by the chain rule.
Example (1): differentiating y = \sin(t^2)
The cell below computes \frac{dy}{dt} for y = \sin(t^2) at t = 1.0 using PyTorch’s autograd.
By the chain rule, the analytical derivative is:
\frac{dy}{dt} = 2t \cdot \cos(t^2)
What Happens Here:
requires_grad=True: tells PyTorch to record every operation ontso it can later differentiate through them.y = torch.sin(t**2): the forward computation. PyTorch silently builds a computation graph: t \to t^2 \to \sin(t^2).torch.autograd.grad(y, t, create_graph=True): walks the graph backwards and applies the chain rule to return \frac{dy}{dt} exactly. Thecreate_graph=Trueflag keeps the derivative itself differentiable, which is needed when computing higher-order derivatives (e.g. \frac{d^2y}{dt^2} for acceleration).
At t = 1.0 the result is 2 \times 1 \times \cos(1) = 1.0806, matching the analytical value to machine precision. This is not a finite-difference approximation, but an exact differentiation based on computational graph.
# Let's compute d/dt [sin(t^2)] at t = 1.0
t = torch.tensor(
[1.0],
requires_grad=True) # requires_grad=True tells PyTorch to track this
print("Input t:")
print(t.item())
print() # blank line for readability
# forward pass
y = torch.sin(t**2)
print(f"Output y: {y.item():.6f}")
print(y.shape)
print()
# Compute dy/dt using autograd
dy_dt = torch.autograd.grad(
outputs=y,
inputs=t,
create_graph=True)[0]
# Compute the analytical solution for comparison
y_analytical = 2 * t * torch.cos(t**2)
print("Output y and its derivative dy/dt:")
print(f"y: {y.item()}")
print(f"dy/dt: {dy_dt.item():.6f}")
print()
print(f"Analytical solution: {y_analytical.item():.6f}")Input t:
1.0
Output y: 0.841471
torch.Size([1])
Output y and its derivative dy/dt:
y: 0.8414709568023682
dy/dt: 1.080605
Analytical solution: 1.080605Example (2): Differentiating a Neural Network’s Output
The previous example differentiated a known expression. The cell below does the same thing but through an actual neural network, showing that autograd works on arbitrarily complex computation graphs, not just textbook formulas.
Setup: A small network (1 \to 16 \to 1 with Tanh) is evaluated at 200 time points t \in [0, 1], producing predictions \hat{x}(t).
Computing derivatives:
First derivative dy/dt: obtained by calling
torch.autograd.gradon the network output with respect to the inputt_test. Thegrad_outputs=torch.ones_like(y_pred)argument is needed becausey_predis a vector (200 values), not a scalar. This argument tells PyTorch to sum the gradients across all outputs (equivalent to differentiating each output independently).Second derivative d²y/dt²: obtained by calling
autograd.gradagain, this time ondydt. This is only possible becausecreate_graph=Truewas set in the first call, keeping the derivative in the computation graph.
Key point: The network is untrained (random weights), so the curves have no physical meaning. But the derivatives are still exact. AD differentiates the actual arithmetic the network performs, regardless of whether the weights have been optimised yet. This is the mechanism a PINN uses to compute \dot{x} and \ddot{x} inside its loss function.
# Setting a random seed for reproducibility
torch.manual_seed(0)
# A simple feedforward network to illustrate AD for derivatives w.r.t. input
diff_nn = nn.Sequential(
nn.Linear(1,16),
nn.Tanh(),
nn.Linear(16,1))
# Evaluate at multiple time points
t_test = torch.linspace(0, 1, 200).unsqueeze(1)
t_test.requires_grad_(True)
print(f"t_test shape: {t_test.shape}")
y_pred = diff_nn(t_test)
print(f"y_pred shape: {y_pred.shape}")
print()
# First derivative:
dydt = torch.autograd.grad(
outputs=y_pred,
inputs=t_test,
# needed because y_pred is a vector
# set grad_outputs to a vector of ones with the same shape as y_pred
grad_outputs=torch.ones_like(y_pred),
# needed so we can differentiate again
create_graph=True
)[0]
# Second derivative:
d2x_dt2 = torch.autograd.grad(
outputs=dydt,
inputs=t_test,
# needed because dydt is a vector
# set grad_outputs to a vector of ones with the same shape as dydt
grad_outputs=torch.ones_like(dydt),
# needed so we can differentiate again
create_graph=True
)[0]
fig, axes = plt.subplots(
3, 1,
figsize=(8, 10))
ts = t_test.detach().numpy()
axes[0].plot(
ts,
y_pred.detach().numpy(),
marker="o",
markersize=3,
markerfacecolor="red",
linestyle="--",
c="black")
axes[0].grid(alpha=0.5)
axes[0].set_ylabel(
r"$y(t)$",
fontsize=12)
axes[0].set_title(r"$y(t)$")
axes[1].plot(
ts,
dydt.detach().numpy(),
marker="o",
markersize=3,
markerfacecolor="red",
linestyle="--",
c="black")
axes[1].grid(alpha=0.5)
axes[1].set_ylabel(
r"$dy/dt$",
fontsize=12)
axes[1].set_title(
"First derivative using Automatic Differentiation:"
+ r"$dy/dt$")
axes[2].plot(
ts,
d2x_dt2.detach().numpy(),
marker="o",
markersize=3,
markerfacecolor="red",
linestyle="--",
c="black")
axes[2].grid(alpha=0.5)
axes[2].set_ylabel(
r"$d^2y/dt^2$",
fontsize=12)
axes[2].set_title(
"Second derivative using Automatic Differentiation:"
+ r"$d^2y/dt^2$")
for ax in axes:
ax.set_xlabel(
r"$t$",
fontsize=12,
)
plt.suptitle(
r"AD gives exact derivatives of the network w.r.t. its input",
fontsize=12)
plt.tight_layout()
plt.show()t_test shape: torch.Size([200, 1])
y_pred shape: torch.Size([200, 1])
NoteWhy do we need
grad_outputs and create_graph?
grad_outputs=torch.ones_like(y_pred): In Example (1), the output y was a scalar (a single number), so torch.autograd.grad(y, t) directly returns dy/dt. Here, y_pred is a vector of 200 values, which is one prediction per time point. PyTorch’s autograd.grad computes a vector-Jacobian product (VJP), not element-wise derivatives, so it needs a “vector” to multiply against the Jacobian. Setting grad_outputs to a tensor of ones is equivalent to asking: “give me the sum \sum_i \partial y_i / \partial t_i”. Because each output y_i depends only on its own input t_i (there is no cross-dependence between different time points), this sum decomposes into independent element-wise derivatives, and the result is a vector where each entry is dy_i/dt_i. Without this argument, PyTorch would raise an error for non-scalar outputs.
create_graph=True: By default, autograd.grad computes the derivative and then discards the computation graph used to produce it. With create_graph=True, the derivative operation itself is recorded in a new graph. This is essential because we need to differentiate twice: once to get dy/dt, and again to get d^2y/dt^2. If create_graph were False on the first call, the second autograd.grad call would fail because there would be no graph left to differentiate through.
Summary of torch.autograd.grad Parameters
torch.autograd.grad(outputs, inputs, grad_outputs=None, retain_graph=None,
create_graph=False, only_inputs=True, allow_unused=None,
is_grads_batched=False, materialize_grads=False)| Parameter | Type | Default | Description |
|---|---|---|---|
outputs |
sequence of Tensor | (required) | The output tensors to differentiate (e.g. y_pred or dydt). |
inputs |
sequence of Tensor | (required) | The input tensors to differentiate with respect to (e.g. t_test). Gradients are returned, not accumulated into .grad. |
grad_outputs |
sequence of Tensor or None | None |
The “vector” in the vector-Jacobian product. Required when outputs is non-scalar. Set to torch.ones_like(outputs) to get the element-wise derivative. |
retain_graph |
bool | value of create_graph |
If False, the computation graph is freed after use. Rarely needs to be set manually. |
create_graph |
bool | False |
If True, the derivative itself becomes part of the computation graph, allowing higher-order derivatives (e.g. computing \ddot{x} from \dot{x}). |
allow_unused |
bool or None | value of materialize_grads |
If False, raises an error when an input was not used in computing the output. |
is_grads_batched |
bool | False |
If True, treats the first dimension of each grad_outputs tensor as a batch dimension for batched vector-Jacobian products. |
materialize_grads |
bool | False |
If True, returns zero tensors instead of None for unused inputs. Useful for higher-order derivatives. |
Return value
torch.autograd.grad returns a tuple of Tensors, one element per input in inputs. Each element is the gradient of outputs with respect to the corresponding input.
| Index | Value | Description |
|---|---|---|
[0] |
\frac{\partial \text{outputs}}{\partial \text{inputs[0]}} | Gradient with respect to the first input tensor. Same shape as inputs[0]. |
[1] |
\frac{\partial \text{outputs}}{\partial \text{inputs[1]}} | Gradient with respect to the second input tensor (only present if inputs has two or more tensors). |
[i] |
\frac{\partial \text{outputs}}{\partial \text{inputs[i]}} | Gradient with respect to the i-th input tensor. |
Each element can be:
- A Tensor with the same shape as the corresponding input, containing the computed gradient.
- None if the input was not used in computing
outputsandallow_unused=True(or zero ifmaterialize_grads=True).
In this notebook we always pass a single tensor as inputs (e.g. t_test), so the returned tuple has length 1 and we index with [0] to get the gradient tensor directly.
Example from this notebook:
# First derivative: dy/dt
dydt = torch.autograd.grad(
outputs=y_pred,
inputs=t_test,
# needed because y_pred is a vector
grad_outputs=torch.ones_like(y_pred),
# needed so we can differentiate again
create_graph=True
)[0]
# Second derivative: d²y/dt²
d2x_dt2 = torch.autograd.grad(
# now differentiating the derivative
outputs=dydt,
inputs=t_test,
grad_outputs=torch.ones_like(dydt),
create_graph=True
)[0]The [0] at the end extracts the first (and only) gradient tensor from the returned tuple, since inputs contains a single tensor.
Why Does ReLU Give Poor Derivatives?
The cell below repeats the same AD demonstration but replaces Tanh with ReLU. By comparing the derivative plots, we can observe that the results are very different.
ReLU is piecewise linear:
\text{ReLU}(z) = \max(0, z) = \begin{cases} 0 & z < 0 \\ z & z \geq 0 \end{cases}
This means:
- The first derivative \frac{d}{dz}\text{ReLU}(z) is either 0 or 1, which is a step function, not a smooth curve.
- The second derivative \frac{d^2}{dz^2}\text{ReLU}(z) is zero everywhere (except at z = 0 where it is undefined).
Because a network with ReLU is just a chain of piecewise linear functions, the overall output \hat{x}(t) is also piecewise linear. Its first derivative is piecewise constant (flat steps), and its second derivative is zero almost everywhere.
Why this matters for PINNs: A PINN loss function involves a second derivative \ddot{x} (or \nabla^2 u for PDEs). If the second derivative is always zero, the network cannot learn to satisfy any differential equation that requires non-zero curvature. This is why PINNs use Tanh (or other smooth activations like \sin). They are infinitely differentiable, producing meaningful gradients at every order. It is a hard requirement for PINNs involving second-order ODEs.
# Setting a random seed for reproducibility
torch.manual_seed(0)
# A simple feedforward network to illustrate AD for derivatives w.r.t. input
diff_nn = nn.Sequential(
nn.Linear(1,16),
nn.ReLU(),
nn.Linear(16,1))
# Evaluate at multiple time points
t_test = torch.linspace(0, 1, 200).unsqueeze(1)
t_test.requires_grad_(True)
y_pred = diff_nn(t_test)
# First derivative:
dydt = torch.autograd.grad(
outputs=y_pred,
inputs=t_test,
# needed because y_pred is a vector
grad_outputs=torch.ones_like(y_pred),
# needed so we can differentiate again
create_graph=True
)[0]
# Second derivative:
d2x_dt2 = torch.autograd.grad(
outputs=dydt,
inputs=t_test,
grad_outputs=torch.ones_like(dydt),
create_graph=True
)[0]
fig, axes = plt.subplots(
3, 1,
figsize=(8, 10))
ts = t_test.detach().numpy()
axes[0].plot(
ts,
y_pred.detach().numpy(),
marker="o",
markersize=3,
markerfacecolor="red",
linestyle="--",
c="black")
axes[0].grid(alpha=0.5)
axes[0].set_ylabel(
r"$y(t)$",
fontsize=12)
axes[0].set_title(r"$y(t)$")
axes[1].plot(
ts,
dydt.detach().numpy(),
marker="o",
markersize=3,
markerfacecolor="red",
linestyle="--",
c="black")
axes[1].grid(alpha=0.5)
axes[1].set_ylabel(
r"$dy/dt$",
fontsize=12)
axes[1].set_title(
"First derivative using Automatic Differentiation:"
+ r"$dy/dt$")
axes[2].plot(
ts,
d2x_dt2.detach().numpy(),
marker="o",
markersize=3,
markerfacecolor="red",
linestyle="--",
c="black")
axes[2].grid(alpha=0.5)
axes[2].set_ylabel(
r"$d^2y/dt^2$",
fontsize=12)
axes[2].set_title(
"Second derivative using Automatic Differentiation:"
+ r"$d^2y/dt^2$")
for ax in axes:
ax.set_xlabel(
r"$t$",
fontsize=12,
)
plt.suptitle(
r"Derivatives of a ReLU network using AD (note the non-smoothness)",
fontsize=12)
plt.tight_layout()
plt.show()
Putting It All Together: A PINN for the Undamped Spring-Mass System
The cell below trains a PINN to solve the ODE for an undamped spring-mass system without any labelled data with only the governing equation and initial conditions.
The Problem
\ddot{x} + x = 0, \qquad x(0) = 1, \quad \dot{x}(0) = 0
This is the equation of motion for an undamped spring-mass system with unit mass (m = 1) and unit stiffness (k = 1). The general form is m\ddot{x} + kx = 0, which simplifies to \ddot{x} + x = 0 when m = k = 1.
- x(t) is the displacement of the mass from its equilibrium position at time t.
- \ddot{x} = d^2x/dt^2 is the acceleration: the second derivative of displacement with respect to time.
- The term +x is the restoring force from the spring (Hooke’s law), which always pulls the mass back toward equilibrium.
The initial conditions specify the state of the system at t = 0:
- x(0) = 1: the mass starts displaced 1 unit from equilibrium (pulled and held).
- \dot{x}(0) = 0: the mass is released from rest (zero initial velocity).
Together, the ODE and these two initial conditions uniquely determine the solution. The exact analytical answer is x(t) = \cos(t) as the mass oscillates back and forth with period 2\pi and amplitude 1. The PINN must discover this from physics alone.
NoteAngular frequency vs the oscillation period T
The angular frequency \omega (in rad/s) and the period T (in seconds) are related by:
T = \frac{2\pi}{\omega}
For this ODE, \omega = \sqrt{k/m} = \sqrt{1/1} = 1 rad/s, so:
T = \frac{2\pi}{1 \;\text{rad/s}} = 2\pi \;\text{s} \approx 6.28 \;\text{s}
The angular frequency \omega tells us how many radians the oscillation sweeps per second. Since one full cycle is 2\pi radians, dividing 2\pi by \omega gives the time for one complete oscillation. That is why t_end = 2 * torch.pi in the code as the PINN is trained over exactly one full cycle of the cosine solution.
Step-by-Step Walkthrough
1. Network Architecture
1 \;\text{input } (t) \;\xrightarrow{}\; 32 \;\text{(Tanh)} \;\xrightarrow{}\; 32 \;\text{(Tanh)} \;\xrightarrow{}\; 32 \;\text{(Tanh)} \;\xrightarrow{}\; 1 \;\text{output } (\hat{x})
Three hidden layers of 32 neurons each with Tanh activation (i.e., smooth and infinitely differentiable, as required for computing \ddot{x}).
2. Training Setup: Key Terminology
Collocation points (t_phys): 500 time values evenly spaced in [0, 2\pi] where the physics residual is evaluated. These are not data points with known answers, but the locations where we ask: “does the network’s output satisfy the ODE here?”. The word “collocation” comes from numerical methods, meaning “points where we enforce the equation”.
Boundary point (t_ic): The single point t = 0 where the initial conditions are enforced.
NoteWhat are collocation points?
The image below (panel a) shows how the computational domain is divided into three types of collocation points, each associated with a different loss term:
| Point type | Colour | Location | Loss term | Purpose |
|---|---|---|---|---|
| IC (Initial Condition) | Blue | Bottom edge (t = t_0) | \mathcal{L}_{\text{IC}} | Enforce the known initial state, e.g. x(0) = 1, \dot{x}(0) = 0. |
| BC (Boundary Condition) | Orange | Left and right edges (x = x_0, x = x_f) | \mathcal{L}_{\text{BC}} | Enforce boundary values at the spatial edges of the domain. |
| PDE (Collocation) | Green | Interior of the domain | \mathcal{L}_{\text{PDE}} | Enforce the governing equation. The PDE residual is evaluated at these points. |
Panel b shows the full PINN workflow: the inputs (x, t) are fed into a deep neural network (DNN), which outputs the solution u and auxiliary variables. Automatic differentiation (AD) computes the required partial derivatives (\partial_x, \partial_t). These are then used to evaluate the three loss terms (\mathcal{L}_{\text{IC}}, \mathcal{L}_{\text{BC}}, \mathcal{L}_{\text{PDE}}), which are summed into a total loss. The network weights \boldsymbol{\theta} are updated until the loss converges.
In our spring-mass ODE example, the domain is 1D (time only), so there are no spatial boundaries. The IC points reduce to t = 0 and the collocation points are the 500 evenly spaced values in [0, 2\pi] where we evaluate the ODE residual \ddot{x} + x = 0.
3. The Loss Function
The total loss has two parts, which is what makes it a PINN instead of standard ML:
IC loss (Initial Condition loss):
\mathcal{L}_{\text{IC}} = \bigl(\hat{x}(0) - x_0\bigr)^2 + \bigl(\dot{\hat{x}}(0) - v_0\bigr)^2
This penalises the network if its output at t = 0 does not match x(0) = 1 and \dot{x}(0) = 0. The derivative \dot{\hat{x}}(0) is computed via autograd, exactly the AD technique demonstrated earlier.
This is a special case of the Mean Squared Error (MSE) loss. The general form of MSE is:
\text{MSE} = \frac{1}{N}\sum_{i=1}^{N}\bigl(\hat{y}_i - y_i\bigr)^2
Here we have just N = 1 boundary point and two conditions to enforce (position and velocity), so the IC loss is simply the sum of two squared errors, no averaging is needed. If we had multiple boundary points (e.g. for a boundary value problem with conditions at both ends), the IC loss would become a proper mean over all of them:
\mathcal{L}_{\text{BC}} = \frac{1}{N_{\text{bc}}}\sum_{i=1}^{N_{\text{bc}}}\bigl(\hat{x}(t_i^{\text{bc}}) - x_i^{\text{bc}}\bigr)^2
Physics residual loss:
\mathcal{L}_{\text{phys}} = \frac{1}{N_\textrm{phys}} \sum_{i=1}^{N_\text{phys}} r(t_i)^2, \qquad r(t) = \ddot{\hat{x}}(t) + \hat{x}(t)
The residual r(t) is what remains when you substitute the network’s prediction into the ODE. If the network perfectly satisfies \ddot{x} + x = 0, the residual is zero everywhere. The loss is the mean squared residual across all collocation points. Both \dot{\hat{x}} and \ddot{\hat{x}} are computed via two successive calls to autograd.grad with create_graph=True.
Total loss:
\mathcal{L} = \mathcal{L}_{\text{IC}} + \mathcal{L}_{\text{phys}}
4. The Training Loop
Each of the 5000 epochs performs:
optimiser.zero_grad(): reset gradients- Compute \hat{x}(0) and \dot{\hat{x}}(0) via the network + AD → IC loss
- Compute \hat{x}(t_\text{phys}), \dot{\hat{x}}(t_\text{phys}), \ddot{\hat{x}}(t_\text{phys}) via two AD passes → physics residual loss
- Sum the losses →
loss.backward()→optimiser.step()
Notice: there are no target labels anywhere. The only supervision comes from the differential equation and the initial conditions.
5. Outputs
- Top plot: The PINN’s prediction vs the exact solution \cos(t). The two curves should overlap closely, confirming that the network has learned the correct solution from physics alone.
- Bottom plot: Training loss over epochs on a log scale, showing three curves:
- Total loss (black): the sum \mathcal{L}_{\text{IC}} + \mathcal{L}_{\text{phys}} that the optimiser minimises. Its convergence indicates overall training progress.
- IC loss (red): measures how well the network satisfies the initial conditions x(0) = 1 and \dot{x}(0) = 0. This typically drops fastest because the IC constraint involves only a single point and two simple equality conditions, making it easy for the optimiser to satisfy early in training.
- Physics loss (blue): measures how well the network satisfies the ODE \ddot{x} + x = 0 across all 500 collocation points. This is the harder constraint because it must hold everywhere simultaneously, so it tends to dominate the total loss and converge more slowly than the IC loss.
# Setting a random seed for reproducibility
torch.manual_seed(42)
# Check if GPU is available and set device accordingly
device = torch.device("cpu")
# ----------------------------------------------------------------------------
# 1. Define the Network
class PINN(nn.Module):
def __init__(self):
super().__init__()
# A simple feedforward network with 3 hidden layers of 32 neurons each
self.net = nn.Sequential(
nn.Linear(1, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 1)
)
# Forward pass: input is time t, output is predicted displacement x(t)
def forward(self, t):
return self.net(t)
# Instantiate the PINN and move to device
pinn = PINN().to(device)
# ----------------------------------------------------------------------------
# 2. Training setup
# Physical parameters
# Initial conditions for the undamped spring-mass system
# ODE: d²x/dt² + x = 0
# x0 = 1.0 means the mass starts at x=1 at t=0
# x(0) = 1 (Initial displacement)
x0 = 1.0
# v0 = 0.0 means the mass starts at rest (dx/dt=0 at t=0)
# dx/dt(0) = 0 (Initial velocity)
v0 = 0.0
# We will train the PINN to learn the solution over one
# full period of oscillation
t_end = 2 * torch.pi
# We will enforce the initial conditions at t=0, so we create
# a tensor for that
t_ic = torch.tensor([[0.0]], device=device, requires_grad=True)
# Collocation points (physics residual evaluated here)
# We use 500 points uniformly spaced in [0, t_end]
N_phys = 500
t_phys = torch.linspace(
0,
t_end,
N_phys, device=device).unsqueeze(1)
t_phys.requires_grad_(True)
# We use the Adam optimiser to train the PINN
optimiser = torch.optim.Adam(
pinn.parameters(),
lr=1e-3)
# ----------------------------------------------------------------------------
# 3. Training loop
ic_loss = []
phys_loss = []
total_loss = []
for epoch in range(5000):
# Reset gradients from the previous step
optimiser.zero_grad()
# IC loss ----------------------------------------------------------------
# We need to compute x(0) and dx/dt(0) from the PINN
# to enforce the initial conditions
# x_ic means x when t=0
x_ic = pinn(t_ic)
# dxdt_ic means dx/dt when t=0, computed using autograd
dxdt_ic = torch.autograd.grad(
outputs=x_ic,
inputs=t_ic,
grad_outputs=torch.ones_like(x_ic),
create_graph=True)[0]
# IC loss: (x(0) - x0)^2 + (dx/dt(0) - v0)^2
# This loss term ensures the PINN's prediction at t=0 matches the
# known initial conditions
loss_ic = (x_ic - x0)**2 + (dxdt_ic - v0)**2
# Physics residual loss --------------------------------------------------
# ODE: d²x/dt² + x = 0 → residual r(t) = d²x/dt² + x
# To compute the residual, we need
# x(t), dx/dt, and d²x/dt² at the collocation points t_phys
# x_phys means x at the collocation points
x_phys = pinn(t_phys)
# dxdt_phys means dx/dt at the collocation points
dxdt_phys = torch.autograd.grad(
outputs=x_phys,
inputs=t_phys,
grad_outputs=torch.ones_like(x_phys),
create_graph=True)[0]
# d2xdt2_phys means d²x/dt² at the collocation points
d2xdt2_phys = torch.autograd.grad(
outputs=dxdt_phys,
inputs=t_phys,
grad_outputs=torch.ones_like(dxdt_phys),
create_graph=True)[0]
# The physics residual is the left-hand side of the ODE:
# r(t) = d²x/dt² + x
residual = d2xdt2_phys + x_phys
# Physics loss: mean squared residual at the collocation points
loss_phys = torch.mean(residual**2)
# Total loss ------------------------------------------------------------
# The total loss is a combination of the initial condition loss and the
# physics residual loss. The optimiser will try to minimize this
# total loss.
loss = loss_ic + loss_phys
loss.backward()
optimiser.step()
ic_loss.append(loss_ic.item())
phys_loss.append(loss_phys.item())
total_loss.append(loss.item())
# ----------------------------------------------------------------------------
# 4. Evaluate and plot
# We evaluate the trained PINN at 300 points in [0, t_end] to
# compare with the exact solution x(t) = cos(t)
t_eval = torch.linspace(
0, t_end, 300, device=device).unsqueeze(1)
# We use torch.no_grad() to tell PyTorch we don't need gradients
# for this part, which saves memory and computation
# since we're just evaluating the PINN.
with torch.no_grad():
x_eval_np = pinn(t_eval).cpu().numpy()
# The exact solution for the undamped spring-mass system with
# the given initial conditions is: x(t) = cos(t)
t_eval_np = t_eval.cpu().numpy()
x_exact = np.cos(t_eval_np)
# We compute the maximum absolute error between the PINN's prediction
# and the exact solution
print("Max error: "
+ f"{np.max(np.abs(x_eval_np - x_exact)):.4f} (should be small)")
fig, (ax1, ax2) = plt.subplots(
2, 1, figsize=(6, 8))
# Plot the exact solution and the PINN's prediction on the same graph
ax1.plot(
t_eval_np,
x_exact,
"k-",
lw=2.5,
label="Exact: cos(t)")
ax1.plot(
t_eval_np,
x_eval_np,
"r--",
lw=2,
label="PINN prediction")
ax1.set_xlabel(
r"Timestep $t(s)$",
fontsize=12)
ax1.set_ylabel(
r"Displacement $x(t)$",
fontsize=12)
ax1.set_title(
"Undamped spring-mass: PINN vs exact analytical solution",
fontsize=14)
ax1.legend()
ax1.grid(True, alpha=0.3)
# Plot the training loss curve on a logarithmic scale to show convergence
ax2.semilogy(
total_loss,
"k--",
lw=2,
label=f"Final loss: {total_loss[-1]:.2e}")
ax2.semilogy(
ic_loss,
"r--",
lw=2,
label=f"Final IC loss: {ic_loss[-1]:.2e}")
ax2.semilogy(
phys_loss,
"b--",
lw=2,
label=f"Final physics loss: {phys_loss[-1]:.2e}")
ax2.set_xlabel(
"Epoch",
fontsize=12)
ax2.set_ylabel(
"Total loss (log scale)",
fontsize=12)
ax2.set_title(
"Training loss curve",
fontsize=14)
ax2.legend()
ax2.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()Max error: 0.0059 (should be small)
Important
t_icandt_physboth haverequires_grad=True, which allowstorch.autograd.gradto differentiate w.r.t. tcreate_graph=Trueon the first derivative is essential as it allows us to differentiate again to get the second derivative- The residual is
d2xdt2_phys + x_phys, which is exactly the left-hand side of \ddot{x} + x = 0 with the right-hand side (zero) moved over - We never tell the network what the true solution is, PINN figures it out by being penalised for violating the ODE and the ICs
References
Francis Fernandes (2026). Mastering Dynamic PINNs.
Wu, Yuandi, Brett Sicard, and Stephen Andrew Gadsden. “Physics-informed machine learning: A comprehensive review on applications in anomaly detection and condition monitoring.” Expert Systems with Applications 255 (2024): 124678.
Dazzi, Susanna. “Physics‐informed neural networks for the augmented system of shallow water equations with topography.” Water Resources Research 60, no. 10 (2024): e2023WR036589.
NoteAcknowledgement
Recently, as I started to spend more time on neural-networks and specifically, physics-informed neural networks (PINNs), I came across a very good and informative resource published by Francis Fernandes on Mastering Dynamic PINNs. I took some time to work through his material. This post and the entire series on Dynamic PINNs represent my understanding with additional notes after studying PINNs based on his tutorials. If you are interested to study his original content, you can purchase his packages from here:
https://topmate.io/pinnsformechanicalengineers