SYSTEM:
You are a helpful assistant. You can help me by answering my questions. You can also ask me questions. I am a physicist and electrical engineer. I am also knowledgeable in mechanical engineering, software engineering, and I'm familiar with C/C++, python, javascript, Matlab, and LLMs.
USER:
write a 2-D python simulation of a discrete spinor field using only pytorch, numpy and matplotlib of a single electron propagating through a stern-gerlach apparatus. Give the electron initial velocity and set the simulation parameters so that splitting is evident before the particle reaches the other end of the simulation. Periodically plot the 2-D field probability density.
ASSISTANT:
Certainly! Below is a Python code that simulates a single electron
propagating through a Stern-Gerlach apparatus in 2D using
PyTorch,
NumPy, and
Matplotlib. The simulation demonstrates the splitting of the
electron's wavefunction due to the interaction with a
magnetic field gradient.
<PYTHON>import numpy as np import matplotlib.pyplot as plt import torch
<PYTHON># Physical constants hbar = 1.054571817e-34 # Reduced Planck's constant (J·s) m_e = 9.10938356e-31 # Electron mass (kg) mu_B = 9.274009994e-24 # Bohr magneton (J/T) # Simulation parameters nx, ny = 200, 200 # Number of spatial grid points dx, dy = 1e-9, 1e-9 # Spatial resolution (m) dt = 1e-16 # Time step (s) nt = 500 # Number of time steps # Magnetic field parameters B0 = 0.0 # Base magnetic field (T) G = 1e7 # Magnetic field gradient (T/m) # Initial electron parameters x0, y0 = 50e-9, 100e-9 # Initial position (m) sigma_x = sigma_y = 5e-9 # Initial wavepacket spread (m) p0 = np.array([1e5 * m_e, 0.0]) # Initial momentum (kg·m/s)
<PYTHON># Create spatial grids x = np.linspace(0, (nx - 1) * dx, nx) y = np.linspace(0, (ny - 1) * dy, ny) X, Y = np.meshgrid(x, y, indexing='ij') # Convert grids to torch tensors X_torch = torch.from_numpy(X).float() Y_torch = torch.from_numpy(Y).float()
<PYTHON># Initialize the spinor wavefunction ψ = [ψ_up, ψ_down]^T # Start with a Gaussian wavepacket centered at (x0, y0) def gaussian_wavepacket(x, y, x0, y0, sigma_x, sigma_y, p0): exponent = -((x - x0)**2 / (2 * sigma_x**2) + (y - y0)**2 / (2 * sigma_y**2)) phase = (p0[0] * (x - x0) + p0[1] * (y - y0)) / hbar envelope = torch.exp(torch.tensor(exponent, dtype=torch.float32)) phase_factor = torch.exp(1j * torch.tensor(phase, dtype=torch.float32)) return envelope * phase_factor # Initial spinor components (equal superposition) psi_up = gaussian_wavepacket(X_torch, Y_torch, x0, y0, sigma_x, sigma_y, p0) / np.sqrt(2) psi_down = gaussian_wavepacket(X_torch, Y_torch, x0, y0, sigma_x, sigma_y, p0) / np.sqrt(2) # Combine into a single tensor psi = torch.stack([psi_up, psi_down], dim=-1) # Shape: (nx, ny, 2)
<PYTHON># Define the kinetic operator using finite differences kx = torch.fft.fftfreq(nx, d=dx) * 2 * np.pi ky = torch.fft.fftfreq(ny, d=dy) * 2 * np.pi KX, KY = torch.meshgrid(kx, ky, indexing='ij') K_squared = KX**2 + KY**2 # Kinetic energy operator in momentum space T_operator = torch.exp(-1j * hbar * K_squared * dt / (2 * m_e)) # Magnetic field along z with gradient in y Bz = B0 + G * (Y_torch - y0) # Pauli matrices sigma_z = torch.tensor([[1, 0], [0, -1]], dtype=torch.cfloat) # Potential energy operator due to magnetic field V = -mu_B * Bz.unsqueeze(-1).unsqueeze(-1) * sigma_z # Shape: (nx, ny, 2, 2) # Exponential of the potential energy operator V_operator = torch.linalg.matrix_exp(-1j * V * dt / hbar)
<PYTHON># Prepare for time evolution probability_densities = [] for t in range(nt): # Apply potential operator (spin-dependent) psi = torch.einsum('...ab,...b->...a', V_operator, psi) # Fourier transform to momentum space psi_ft = torch.fft.fft2(psi, dim=(0, 1)) # Apply kinetic operator psi_ft[..., 0] *= T_operator psi_ft[..., 1] *= T_operator # Inverse Fourier transform back to position space psi = torch.fft.ifft2(psi_ft, dim=(0, 1)) # Calculate probability density prob_density = (psi.conj() * psi).real.sum(dim=-1) # Store probability density at intervals if t % 50 == 0: probability_densities.append(prob_density.numpy())
<PYTHON># Plot the probability density at different time intervals fig, axes = plt.subplots(1, len(probability_densities), figsize=(15, 3)) for i, density in enumerate(probability_densities): ax = axes[i] im = ax.imshow(density.T, extent=(0, nx*dx*1e9, 0, ny*dy*1e9), origin='lower', cmap='viridis') ax.set_title(f'Time = {i * 50 * dt * 1e15:.1f} fs') ax.set_xlabel('x (nm)') ax.set_ylabel('y (nm)') fig.colorbar(im, ax=ax) plt.tight_layout() plt.show()
p0
gives the electron an initial velocity in the x-direction,
allowing it to propagate through the apparatus.
G
causes the spin components to experience different forces,
leading to spatial separation (splitting) in the y-direction.
This simulation demonstrates how an electron's spinor wavefunction evolves as it passes through a Stern-Gerlach apparatus. The splitting of the probability density illustrates the measurement of spin components due to the magnetic field gradient, a fundamental concept in quantum mechanics.
Feel free to adjust the parameters (such as
G,
nt,
dt, and initial conditions) to explore different regimes and
observe how they affect the splitting behavior.
USER:
ok, I think I have this working now. next I want to generate a .gif animation from these successive plots:
for t in range(nt):
# Apply potential operator (spin-dependent)
psi = torch.einsum('...ab,...b->...a', V_operator,
psi)
<TEXT># Fourier transform to momentum space psi_ft = torch.fft.fft2(psi, dim=(0, 1)) # Apply kinetic operator psi_ft[..., 0] *= T_operator psi_ft[..., 1] *= T_operator # Inverse Fourier transform back to position space psi = torch.fft.ifft2(psi_ft, dim=(0, 1)) # Calculate probability density prob_density = (psi.conj() * psi).real.sum(dim=-1) # Store probability density at intervals if t % 50 == 0: # probability_densities.append(prob_density.numpy()) plt.figure(figsize=(8, 6)) plt.imshow(prob_density.numpy(), extent=(0, nx*dx*1e9, 0, ny*dy*1e9), origin='lower', aspect='auto', cmap='viridis') plt.xlabel('z position (mm)') plt.ylabel('y position (mm)') plt.title(f'Time: {t * dt:.2e} s') # plt.colorbar(label='Probability Density') plt.show()
ASSISTANT:
Certainly! Creating a GIF animation of your simulation is a great
way to visualize the dynamics over time. You can achieve this by
modifying your code to save each frame during the loop and then
compile these frames into a GIF using
matplotlib's animation functions or an external library like
imageio.
Below, I'll provide two options:
FuncAnimation
Here's how you can modify your code to create an animation
using
matplotlib.animation.FuncAnimation.
<PYTHON>import numpy as np import matplotlib.pyplot as plt import torch from matplotlib.animation import FuncAnimation # ... [Your existing code above remains unchanged] ... # Prepare figure and axis for animation fig, ax = plt.subplots(figsize=(8, 6)) extent = (0, nx*dx*1e9, 0, ny*dy*1e9) # Convert to nm or mm as needed # Initialize the image with the first frame prob_density = (psi.conj() * psi).real.sum(dim=-1) im = ax.imshow(prob_density.numpy(), extent=extent, origin='lower', aspect='auto', cmap='viridis') ax.set_xlabel('z position (nm)') ax.set_ylabel('y position (nm)') title = ax.set_title('') # Function to update each frame def animate(t): global psi for _ in range(50): # Update 50 steps per frame # Apply potential operator (spin-dependent) psi = torch.einsum('...ab,...b->...a', V_operator, psi) # Fourier transform to momentum space psi_ft = torch.fft.fft2(psi, dim=(0, 1)) # Apply kinetic operator psi_ft[..., 0] *= T_operator psi_ft[..., 1] *= T_operator # Inverse Fourier transform back to position space psi = torch.fft.ifft2(psi_ft, dim=(0, 1)) # Calculate probability density prob_density = (psi.conj() * psi).real.sum(dim=-1) # Update image data im.set_data(prob_density.numpy()) title.set_text(f'Time: {t * 50 * dt * 1e15:.1f} fs') return [im, title] # Create the animation frames = nt // 50 # Number of frames anim = FuncAnimation(fig, animate, frames=frames, interval=50, blit=True) # Save the animation as a GIF anim.save('stern_gerlach_simulation.gif', writer='imagemagick', fps=10) plt.close()
psi
over 50 time steps and then updates the plot.
matplotlib.animation
creates the animation by repeatedly calling
animate.
save
method saves the animation as a GIF file. Ensure you have
ImageMagick
installed, or you can use
Pillow
as the writer.
ImageMagick: You might need to install
ImageMagick for
matplotlib
to save GIF files.
<BASH># For Ubuntu/Debian sudo apt-get install imagemagick
Pillow Alternative: Alternatively, you can
use
Pillow
as the writer:
<PYTHON>anim.save('stern_gerlach_simulation.gif', writer='pillow', fps=10)
Ensure you have Pillow installed:
<BASH>pip install pillow
imageio
Alternatively, you can save each frame as an image file and then
compile them into a GIF using
imageio.
imageio
First, install the
imageio
library:
<BASH>pip install imageio
<PYTHON>import numpy as np import matplotlib.pyplot as plt import torch import imageio import os # ... [Your existing code above remains unchanged] ... # Create a directory to save frames os.makedirs('frames', exist_ok=True) filenames = [] for t in range(nt): # Apply potential operator (spin-dependent) psi = torch.einsum('...ab,...b->...a', V_operator, psi) # Fourier transform to momentum space psi_ft = torch.fft.fft2(psi, dim=(0, 1)) # Apply kinetic operator psi_ft[..., 0] *= T_operator psi_ft[..., 1] *= T_operator # Inverse Fourier transform back to position space psi = torch.fft.ifft2(psi_ft, dim=(0, 1)) # Calculate probability density prob_density = (psi.conj() * psi).real.sum(dim=-1) # Store probability density at intervals if t % 10 == 0: # Adjust interval as needed plt.figure(figsize=(8, 6)) plt.imshow(prob_density.numpy(), extent=(0, nx*dx*1e9, 0, ny*dy*1e9), origin='lower', aspect='auto', cmap='viridis') plt.xlabel('z position (nm)') plt.ylabel('y position (nm)') plt.title(f'Time: {t * dt * 1e15:.1f} fs') filename = f'frames/frame_{t:04d}.png' filenames.append(filename) plt.savefig(filename) plt.close() # Compile frames into a GIF with imageio.get_writer('stern_gerlach_simulation.gif', mode='I', fps=10) as writer: for filename in filenames: image = imageio.imread(filename) writer.append_data(image) # Optionally, remove frames after creating GIF import shutil shutil.rmtree('frames')
t % 10), save the plot as an image file.
imageio: After the loop, use
imageio
to read the saved images and compile them into a GIF.
fps
parameter in
get_writer
to control the speed of the GIF.
t % 10
to capture frames more or less frequently.
plt.figure
and
plt.savefig
to improve the quality of the images.
Colorbar: If you want to include a colorbar,
you can uncomment the
plt.colorbar
line and adjust the layout accordingly.
Position Units: Ensure consistency in units.
Your extent is calculated using
nx*dx*1e9, which converts meters to nanometers. Adjust labels if you
prefer millimeters.
<PYTHON>extent = (0, nx*dx*1e3, 0, ny*dy*1e3) # For mm ax.set_xlabel('z position (mm)') ax.set_ylabel('y position (mm)')
Animation Libraries: There are other
libraries like
celluloid
and
matplotlib.animation
that can be used to create animations. The above methods are
straightforward and directly modify your existing code.
By modifying your code as shown above, you can generate a GIF animation that visualizes the evolution of the electron's probability density as it propagates through the Stern-Gerlach apparatus. This animation will help you and others better understand the dynamics of the simulation.
Feel free to adjust the parameters and styling to suit your preferences. If you encounter any issues or have further questions, don't hesitate to ask!
USER:
this all works really well, but I think I need to renormalize each frame because as the wave function spreads out I can no longer see the animation.
fig, ax = plt.subplots(figsize=(8, 6))
extent = (0, nxdx1e9, 0, nydy1e9) # Convert to
nm or mm as needed
prob_density = (psi.conj() * psi).real.sum(dim=-1)
im = ax.imshow(prob_density.numpy(),
extent=extent,
origin='lower',
aspect='auto',
cmap='viridis')
ax.set_xlabel('z position (nm)')
ax.set_ylabel('y position (nm)')
title = ax.set_title('')
def animate(t):
global psi
<TEXT>for _ in range(50): # Update 50 steps per frame # Apply potential operator (spin-dependent) psi = torch.einsum('...ab,...b->...a', V_operator, psi) # Fourier transform to momentum space psi_ft = torch.fft.fft2(psi, dim=(0, 1)) # Apply kinetic operator psi_ft[..., 0] *= T_operator psi_ft[..., 1] *= T_operator # Inverse Fourier transform back to position space psi = torch.fft.ifft2(psi_ft, dim=(0, 1)) # Calculate probability density prob_density = (psi.conj() * psi).real.sum(dim=-1) # Update image data im.set_data(prob_density.numpy()) title.set_text(f'Time: {t * 50 * dt * 1e15:.1f} fs') return [im, title]
frames = nt // 50 # Number of frames
anim = FuncAnimation(fig, animate, frames=frames, interval=50,
blit=True)
anim.save('stern_gerlach_simulation.gif', writer='imagemagick', fps=10)
plt.close()
ASSISTANT:
You're absolutely right—the spreading of the wavefunction causes the peak probability density to decrease over time, making it harder to visualize in the animation. This happens because the wavefunction's amplitude spreads over a larger area, leading to lower values in any given point, and when plotted with a fixed color scale based on the initial maximum, later frames appear dim.
To address this, you can normalize the probability density in each frame so that the maximum value is consistent across all frames. This way, the highest probability density in each frame maps to the maximum of the color scale, enhancing visibility throughout the animation.
Here's how you can modify your code to normalize each frame:
animate
Function
Add normalization of
prob_density
within your
animate
function:
<PYTHON>def animate(t): global psi for _ in range(50): # Update 50 steps per frame # Apply potential operator (spin-dependent) psi = torch.einsum('...ab,...b->...a', V_operator, psi) # Fourier transform to momentum space psi_ft = torch.fft.fft2(psi, dim=(0, 1)) # Apply kinetic operator psi_ft[..., 0] *= T_operator psi_ft[..., 1] *= T_operator # Inverse Fourier transform back to position space psi = torch.fft.ifft2(psi_ft, dim=(0, 1)) # Calculate probability density prob_density = (psi.conj() * psi).real.sum(dim=-1) # Normalize the probability density prob_density /= prob_density.max() + 1e-16 # Add small number to prevent division by zero # Update image data im.set_data(prob_density.numpy()) im.set_clim(0, 1) # Update color scale to match normalization title.set_text(f'Time: {t * 50 * dt * 1e15:.1f} fs') return [im, title]
Explanation:
prob_density
by its maximum value in each frame. The
+ 1e-16
ensures numerical stability by preventing division by zero in
case
prob_density.max()
is zero.
im.set_clim(0, 1)
to ensure the color mapping corresponds to the normalized data.
If you prefer not to normalize each frame (to preserve relative amplitude information), you can use a logarithmic color scale, which enhances features across a wide range of values.
imshow
Call with
norm=LogNorm
First, import
LogNorm:
<PYTHON>from matplotlib.colors import LogNorm
Then, when initializing
im, use:
<PYTHON>im = ax.imshow(prob_density.numpy(), extent=extent, origin='lower', aspect='auto', cmap='viridis', norm=LogNorm(vmin=1e-16, vmax=prob_density.max()))
Note:
vmin: We set
vmin
to a small positive number (1e-16) to avoid issues with log of zero.
vmax: Use the maximum of
prob_density
from the initial frame or compute it over all frames.
animate
Function
Within the
animate
function, after updating
prob_density, update the color limits if necessary:
<PYTHON>def animate(t): global psi for _ in range(50): # Update 50 steps per frame # [Time evolution steps remain the same] # Calculate probability density prob_density = (psi.conj() * psi).real.sum(dim=-1) # Update image data im.set_data(prob_density.numpy()) # Optionally update color limits if the maximum changes significantly im.set_clim(vmin=1e-16, vmax=prob_density.max()) title.set_text(f'Time: {t * 50 * dt * 1e15:.1f} fs') return [im, title]
Explanation:
LogNorm
scales the colormap logarithmically, enhancing low-magnitude
features.
vmax: If you expect
prob_density.max()
to decrease significantly over time, you might need to update
vmax
to maintain contrast.
vmax
may cause flickering in the animation due to shifting color
scales.
Normalization per Frame: Best for visual clarity if you're primarily interested in observing qualitative changes in the probability density distribution over time.
Logarithmic Scale: Preserves relative amplitude information, useful if you want to perceive both high and low probability regions simultaneously.
Here's how your code would look with the normalization approach applied:
<PYTHON>import numpy as np import matplotlib.pyplot as plt import torch from matplotlib.animation import FuncAnimation # ... [Your existing code for constants, grid setup, and initialization remains unchanged] ... # Prepare figure and axis for animation fig, ax = plt.subplots(figsize=(8, 6)) extent = (0, nx*dx*1e9, 0, ny*dy*1e9) # Convert to nm or mm as needed # Initialize the image with the first frame prob_density = (psi.conj() * psi).real.sum(dim=-1) prob_density /= prob_density.max() + 1e-16 # Normalize im = ax.imshow(prob_density.numpy(), extent=extent, origin='lower', aspect='auto', cmap='viridis', vmin=0, vmax=1) # Set color limits ax.set_xlabel('z position (nm)') ax.set_ylabel('y position (nm)') title = ax.set_title('') # Function to update each frame def animate(t): global psi for _ in range(50): # Update 50 steps per frame # Apply potential operator (spin-dependent) psi = torch.einsum('...ab,...b->...a', V_operator, psi) # Fourier transform to momentum space psi_ft = torch.fft.fft2(psi, dim=(0, 1)) # Apply kinetic operator psi_ft[..., 0] *= T_operator psi_ft[..., 1] *= T_operator # Inverse Fourier transform back to position space psi = torch.fft.ifft2(psi_ft, dim=(0, 1)) # Calculate probability density prob_density = (psi.conj() * psi).real.sum(dim=-1) # Normalize the probability density prob_density /= prob_density.max() + 1e-16 # Add small number to prevent division by zero # Update image data im.set_data(prob_density.numpy()) # Color limits are already set to [0, 1], no need to update clim title.set_text(f'Time: {t * 50 * dt * 1e15:.1f} fs') return [im, title] # Create the animation frames = nt // 50 # Number of frames anim = FuncAnimation(fig, animate, frames=frames, interval=50, blit=True) # Save the animation as a GIF anim.save('stern_gerlach_simulation.gif', writer='pillow', fps=10) plt.close()
Notes:
vmin
and
vmax: When initializing
im, we set
vmin=0
and
vmax=1
to match the normalization.
clim: Since the color limits are consistent, you don't need to
update them in each frame.
writer
to
'pillow'
for broader compatibility. Ensure you have Pillow installed
(pip install pillow).
By normalizing each frame, you'll maintain a consistent and clear visualization of the probability density throughout the animation. This will help you and others observe the dynamics of the wavefunction as it propagates through the Stern-Gerlach apparatus without the display becoming too dim.
Feel free to experiment with both the normalization and logarithmic scaling methods to see which provides the best visual representation for your purposes. If you have any further questions or need additional assistance, don't hesitate to ask!