Schrödinger Equation Demystified using Tensorflow. | by Devmallya Karar | Sep, 2024

The Schrödinger equation is a fundamental equation in quantum mechanics that describes how the quantum state of a physical system changes over time. In the context of quantum computing, it is crucial because it governs the behaviour of quantum bits or qubits, which are the fundamental units of information in a quantum computer.

To understand its relevance, let’s start by briefly explaining the Schrödinger equation itself. In its simplest form, the Schrödinger equation is a partial differential equation that provides a way to calculate the evolution of a quantum system’s wave function over time. The wave function is a mathematical object that encapsulates all the information about a quantum system, such as the probability distribution of particles.

In quantum mechanics, the Schrödinger equation is often written as:

Here, Ψ(x,t) represents the wave function of the system, ℏ is the reduced Planck constant, and H^ is the Hamiltonian operator, which represents the total energy of the system. The equation essentially states that the change in the wave function over time (on the left side) is equal to the action of the Hamiltonian operator on the wave function (on the right side).

In quantum computing, the Schrödinger equation plays a crucial role in describing the behaviour of qubits. A qubit is a quantum analoge of a classical bit but can exist in a superposition of states. This means a qubit can represent both 0 and 1 simultaneously, unlike a classical bit that can be only 0 or 1 at any given time. The evolution of these quantum states is governed by the Schrödinger equation, which dictates how the probability amplitudes of qubit states change.

When a quantum computer runs an algorithm, it manipulates qubits through various quantum gates. These gates are unitary operations that transform qubit states in ways that are determined by the Schrödinger equation. Essentially, the Schrödinger equation ensures that the quantum state evolves in a way that preserves the total probability (or norm) and reflects the underlying physics.

Now, let’s explore how the Schrödinger equation intersects with machine learning, particularly in the context of quantum machine learning. Quantum machine learning seeks to leverage the principles of quantum computing to enhance or transform machine learning algorithms. Here are a few ways the Schrödinger equation is relevant to this field:

1. Quantum State Representation: In quantum machine learning, quantum states can represent complex data structures more efficiently than classical representations. For example, quantum algorithms can encode high-dimensional data into a quantum state. The Schrödinger equation governs how these quantum states evolve during computations, which can potentially lead to new methods for data representation and manipulation.

2. Quantum Algorithms: Many quantum algorithms rely on the Schrödinger equation to execute operations. For instance, algorithms like the Quantum Fourier Transform and Grover’s Search utilize quantum superposition and entanglement, which are directly tied to the principles described by the Schrödinger equation. These algorithms can solve certain problems more efficiently than their classical counterparts, offering potential speedups for machine learning tasks.

3. Variational Quantum Algorithms: Variational quantum algorithms, such as the Variational Quantum Eigensolver (VQE) and Quantum Approximate Optimization Algorithm (QAOA), are designed to find approximate solutions to optimization problems. These algorithms involve preparing and evolving quantum states (or wave functions) according to the Schrödinger equation and then measuring their properties. In the context of machine learning, these algorithms can be used to optimize models or solve complex optimization problems more efficiently.

4. Quantum Neural Networks: Quantum neural networks are a type of quantum machine learning model that applies quantum operations to neural network architectures. The Schrödinger equation helps describe the evolution of quantum states within these networks. Quantum neural networks can exploit quantum phenomena like entanglement and superposition to potentially enhance learning capabilities and solve problems that are challenging for classical neural networks.

5. Quantum Data Analysis: Quantum algorithms can process and analyze large datasets in ways that classical algorithms cannot easily achieve. The Schrödinger equation governs the evolution of quantum states during these processes, ensuring that the quantum computations adhere to the laws of quantum mechanics. This can lead to new approaches for data analysis and pattern recognition.

Before moving further let me explain you about the famous Schrödinger cat thought experiment !!

The Schrödinger’s cat experiment is a thought experiment proposed by physicist Erwin Schrödinger in 1935 to illustrate the concept of quantum superposition and the paradoxes of quantum mechanics. It is a central example in discussions about the interpretation of quantum mechanics.

The Thought Experiment:

Imagine a sealed box containing a cat, a radioactive atom, a Geiger counter, a vial of poison, and a hammer. The setup is as follows:

1. Radioactive Atom: The atom has a 50% chance of decaying within an hour and a 50% chance of not decaying.
   2. Geiger Counter: This device detects radiation. If it detects radiation (indicating that the atom has decayed), it triggers a mechanism.
   3. Trigger Mechanism: The mechanism releases the hammer, which then breaks the vial of poison.
   4. Poison: The poison, if released, kills the cat.

Quantum Mechanics and Superposition:

According to quantum mechanics, particles like the radioactive atom exist in a state of superposition, where they can be in multiple states simultaneously until measured. This means that, until someone opens the box to observe the system, the atom is considered to be in a superposition of both decayed and not decayed states.

Given this, the Geiger counter is also in a superposition of detecting radiation and not detecting radiation. Consequently, the trigger mechanism is in a superposition of having both released and not released the hammer. The vial of poison is similarly in a superposition of both broken and unbroken states.

Because the entire system — including the cat — is entangled with the state of the radioactive atom, the cat is also in a superposition of being both dead and alive until someone opens the box and observes the situation.

Paradox and Implications:

Schrödinger proposed this scenario to highlight the strange and seemingly paradoxical nature of quantum mechanics. The thought experiment questions the validity of applying quantum mechanical principles to everyday objects. In quantum mechanics, superposition applies to microscopic particles, but applying it to macroscopic objects like a cat leads to seemingly absurd conclusions.

The paradox raises questions about the role of observation in quantum mechanics. According to the Copenhagen interpretation, the act of measurement collapses the superposition into one of the possible definite states. Until observed, the cat is in a mixture of both states, which challenges our intuitive understanding of reality.

Interpretations and Modern Views:

Several interpretations of quantum mechanics address Schrödinger’s cat paradox:

1. Copenhagen Interpretation: Asserts that the cat is in a superposition of states until observed.
2. Many-Worlds Interpretation: Suggests that all possible outcomes occur in separate, branching universes. Thus, there are worlds where the cat is alive and others where it is dead.
3. Objective Collapse Theories: Propose that wave function collapse occurs independently of observation, with certain mechanisms causing this collapse.

The Schrödinger’s cat thought experiment remains a cornerstone for discussions about the nature of reality and measurement in quantum mechanics, illustrating the complexity and counterintuitive aspects of quantum theory.

To solve the Schrödinger equation using TensorFlow, you can use neural networks to approximate the wave function or to learn the dynamics of quantum systems. Here’s a basic outline of how you might approach this problem:

1. Define the Schrödinger Equation: For a one-dimensional system, the time-dependent Schrödinger equation is:

where ψ(x,t) is the wave function, V(x) is the potential, ℏ is the reduced Planck constant, and m is the mass of the particle.

2. Prepare the Dataset:

• Generate or use an existing dataset of initial conditions and potential functions.
• You might simulate the wave function evolution using numerical methods to generate training data.

3. Define the Neural Network Model:

• Create a neural network model in TensorFlow to approximate the wave function ψ(x,t).
• Use the model to predict ψ(x,t) given x and t.

4. Define Loss Function:

• The loss function can be based on the residual of the Schrödinger equation when evaluated using the neural network predictions. This involves calculating the difference between the left and right sides of the Schrödinger equation using automatic differentiation.

5. Train the Model:

• Use TensorFlow’s training routines to optimize the neural network parameters by minimizing the loss function.

Here’s a basic TensorFlow implementation for solving a simplified version of the Schrödinger equation:

1. Importing Libraries

import tensorflow as tf
import numpy as np

• TensorFlow: A popular machine learning library used here for creating and training the neural network.
• NumPy: A library for numerical operations, used here for generating synthetic data.

2. Defining Parameters

hbar = 1.0  # Reduced Planck constant
m = 1.0     # Mass of the particle

• hbar: The reduced Planck constant, set to 1.0 for simplicity. In a more realistic scenario, this would be set to its actual value.
• m: The mass of the particle, also set to 1.0 for simplicity.

3. Defining the Neural Network Model

def create_model():
model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(2,)),  # Input is (x, t)
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(1)  # Output is the wave function ψ(x, t)
])
return model

Model Definition: This function creates a simple feedforward neural network with:

• Input Layer: Accepts a 2-dimensional input (x, t), where x is position and t is time.
• Hidden Layers: Two dense (fully connected) layers with 64 neurons each and ReLU activation functions.
• Output Layer: Produces a single value representing the wave function ψ(x,t).

4. Defining the Loss Function

def loss_function(model, x, t, V):
with tf.GradientTape() as tape:
tape.watch(x)
with tf.GradientTape() as tape2:
tape2.watch(t)
psi = model(tf.stack([x, t], axis=1))
dpsi_dt = tape2.gradient(psi, t)
dpsi_dx = tape.gradient(psi, x)
d2psi_dx2 = tf.gradients(dpsi_dx, x)[0]
H_psi = (-hbar**2 / (2 * m)) * d2psi_dx2 + V * psi
schrodinger_eq = -1j * hbar * dpsi_dt - H_psi
loss = tf.reduce_mean(tf.square(tf.abs(schrodinger_eq)))
return loss

• GradientTape: TensorFlow’s way of recording operations for automatic differentiation.
• Computing Gradients:
• dpsi_dt: The time derivative of the wave function.
• dpsi_dx: The spatial derivative of the wave function.
• d2psi_dx2: The second spatial derivative, approximated using TensorFlow’s gradients function.
• Hamiltonian: Computed as Hψ, and V as the potential energy.
• Schrödinger Equation: The goal is to minimize the difference between the left and right sides of the Schrödinger equation.
• Loss Calculation: Measures how well the neural network approximates the Schrödinger equation by squaring the absolute value of the equation’s residual and taking the mean.

5. Generating Synthetic Data

def generate_data(num_samples):
x = np.linspace(-1, 1, num_samples).astype(np.float32)
t = np.linspace(0, 1, num_samples).astype(np.float32)
V = np.zeros(num_samples, dtype=np.float32)  # Potential function
return x, t, V

• x: Array of spatial coordinates.
• t: Array of time coordinates.
• V: Array of potential values (set to zero here).

6. Training the Model

def train_model(model, x, t, V, epochs=1000):
optimizer = tf.keras.optimizers.Adam()
for epoch in range(epochs):
with tf.GradientTape() as tape:
loss = loss_function(model, x, t, V)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
if epoch % 100 == 0:
print(f'Epoch {epoch}, Loss: {loss.numpy()}')

• Optimizer: Adam optimizer is used to adjust the model weights.
• Training Loop: Runs for a specified number of epochs, computing the loss and updating the model weights to minimize the loss.

7. Main Code Execution

x, t, V = generate_data(1000)
x_tf = tf.convert_to_tensor(x)
t_tf = tf.convert_to_tensor(t)
V_tf = tf.convert_to_tensor(V)
model = create_model()
train_model(model, x_tf, t_tf, V_tf)

• Data Preparation: Generates synthetic data and converts it into TensorFlow tensors.
• Model Creation: Instantiates the neural network.
• Model Training: Trains the model using the synthetic data.

And here is the Joke:

Why did Schrödinger’s cat stay out of the casino?

Because it couldn’t decide if it was winning or losing !!