259 lines
14 KiB
Plaintext
259 lines
14 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Try to use cupy and use GPU for faster compute (It does not work on my PC, since I don't have a NVIDIA GPU):"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import matplotlib.pyplot as plt\n",
|
|
"import matplotlib.animation as animation\n",
|
|
"#import numpy as np\n",
|
|
"import cupy as np\n",
|
|
"import sympy as sp\n",
|
|
"from IPython.display import Math, display\n",
|
|
"from matplotlib.axes import Axes\n",
|
|
"from scipy import constants as const\n",
|
|
"from scipy.integrate import quad\n",
|
|
"from scipy.optimize import root_scalar\n",
|
|
"from scipy.signal import argrelmax,argrelmin,find_peaks\n",
|
|
"from tqdm import tqdm\n",
|
|
"from typing import Union\n",
|
|
"\n",
|
|
"import fewfermions.analysis.units as si\n",
|
|
"from fewfermions.simulate.traps.twod.trap import DoubleTweezer"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"initial_power = 50* si.uW\n",
|
|
"initial_waist = 1.1*si.uW\n",
|
|
"initial_distance = 2*si.um\n",
|
|
"\n",
|
|
"trap: DoubleTweezer = DoubleTweezer(\n",
|
|
" power=0, # Set pancake laser power to 0, no 2D trap\n",
|
|
" grad_z= 0*si.G/si.cm,\n",
|
|
" grad_r=0,\n",
|
|
" power_tweezer1 = initial_power, #stationary\n",
|
|
" power_tweezer2 = initial_power, #transfer tweezer\n",
|
|
" waist_tweezer1 = initial_waist, #stationary\n",
|
|
" waist_tweezer2 = initial_waist, #transfer tweezer\n",
|
|
" distance_tweezers = initial_distance,\n",
|
|
"\n",
|
|
" a=180*(4 * np.pi * const.epsilon_0 * const.value(\"Bohr radius\")**3)/(2 * const.epsilon_0 * const.c),\n",
|
|
" wvl = 532 * si.nm,\n",
|
|
"\n",
|
|
" g = 0,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Try with cupy:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 16,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import cupy as np\n",
|
|
"\n",
|
|
"n_pot_steps = [40,40,40]\n",
|
|
"n_levels = 4\n",
|
|
"\n",
|
|
"left_cutoff = -0.5*initial_distance-2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
|
|
"right_cutoff = 0.5*initial_distance+2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
|
|
"back_cutoff = -2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
|
|
"front_cutoff = 2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
|
|
"bottom_cutoff = -2*max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
|
|
"top_cutoff = 2*max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
|
|
"\n",
|
|
"extend = [(left_cutoff,right_cutoff),\n",
|
|
" (back_cutoff,front_cutoff),\n",
|
|
" (bottom_cutoff,top_cutoff)]\n",
|
|
"\n",
|
|
"\n",
|
|
"# Solve the hamiltonian numerically\n",
|
|
"energies, states, potential, coords = trap.nstationary_solution(\n",
|
|
" [trap.x,trap.y,trap.z], extend, n_pot_steps, k=n_levels)\n",
|
|
"\n",
|
|
"#x3D,y3D,z3D = np.meshgrid(coords[trap.x],coords[trap.y],coords[trap.z])\n",
|
|
"np.savez(\"data/test_3D.npz\",energies=energies, states=states, pot=potential(x3D,y3D,z3D), x=coords[trap.x],y=coords[trap.y],z=coords[trap.z])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Try with regular numpy:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"ename": "",
|
|
"evalue": "",
|
|
"output_type": "error",
|
|
"traceback": [
|
|
"\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
|
|
"\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
|
|
"\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
|
|
"\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"import numpy as np\n",
|
|
"\n",
|
|
"n_pot_steps = [40,40,40]\n",
|
|
"n_levels = 4\n",
|
|
"\n",
|
|
"left_cutoff = -0.5*initial_distance-2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
|
|
"right_cutoff = 0.5*initial_distance+2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
|
|
"back_cutoff = -2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
|
|
"front_cutoff = 2*max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])\n",
|
|
"bottom_cutoff = -2*max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
|
|
"top_cutoff = 2*max([float(trap.subs(trap.get_tweezer_rayleigh1())),float(trap.subs(trap.get_tweezer_rayleigh2()))])\n",
|
|
"\n",
|
|
"extend = [(left_cutoff,right_cutoff),\n",
|
|
" (back_cutoff,front_cutoff),\n",
|
|
" (bottom_cutoff,top_cutoff)]\n",
|
|
"\n",
|
|
"\n",
|
|
"# Solve the hamiltonian numerically\n",
|
|
"energies, states, potential, coords = trap.nstationary_solution(\n",
|
|
" [trap.x,trap.y,trap.z], extend, n_pot_steps, k=n_levels)\n",
|
|
"\n",
|
|
"#x3D,y3D,z3D = np.meshgrid(coords[trap.x],coords[trap.y],coords[trap.z])\n",
|
|
"np.savez(\"data/test_3D.npz\",energies=energies, states=states, pot=potential(x3D,y3D,z3D), x=coords[trap.x],y=coords[trap.y],z=coords[trap.z])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 12,
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"ename": "CUDARuntimeError",
|
|
"evalue": "cudaErrorInsufficientDriver: CUDA driver version is insufficient for CUDA runtime version",
|
|
"output_type": "error",
|
|
"traceback": [
|
|
"\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
|
|
"\u001b[1;31mCUDARuntimeError\u001b[0m Traceback (most recent call last)",
|
|
"Cell \u001b[1;32mIn[12], line 2\u001b[0m\n\u001b[0;32m 1\u001b[0m \u001b[38;5;66;03m# Generate spatial grid\u001b[39;00m\n\u001b[1;32m----> 2\u001b[0m x \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mlinspace(\u001b[38;5;241m*\u001b[39mextend[\u001b[38;5;241m0\u001b[39m], n_pot_steps[\u001b[38;5;241m0\u001b[39m])\n\u001b[0;32m 3\u001b[0m y \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mlinspace(\u001b[38;5;241m*\u001b[39mextend[\u001b[38;5;241m1\u001b[39m], n_pot_steps[\u001b[38;5;241m1\u001b[39m])\n\u001b[0;32m 4\u001b[0m z \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mlinspace(\u001b[38;5;241m*\u001b[39mextend[\u001b[38;5;241m2\u001b[39m], n_pot_steps[\u001b[38;5;241m2\u001b[39m])\n",
|
|
"File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\cupy\\_creation\\ranges.py:161\u001b[0m, in \u001b[0;36mlinspace\u001b[1;34m(start, stop, num, endpoint, retstep, dtype, axis)\u001b[0m\n\u001b[0;32m 159\u001b[0m scalar_stop \u001b[38;5;241m=\u001b[39m cupy\u001b[38;5;241m.\u001b[39misscalar(stop)\n\u001b[0;32m 160\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m scalar_start \u001b[38;5;129;01mand\u001b[39;00m scalar_stop:\n\u001b[1;32m--> 161\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _linspace_scalar(start, stop, num, endpoint, retstep, dtype)\n\u001b[0;32m 163\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m scalar_start:\n\u001b[0;32m 164\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28misinstance\u001b[39m(start, cupy\u001b[38;5;241m.\u001b[39mndarray) \u001b[38;5;129;01mand\u001b[39;00m start\u001b[38;5;241m.\u001b[39mdtype\u001b[38;5;241m.\u001b[39mkind \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m):\n",
|
|
"File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\cupy\\_creation\\ranges.py:91\u001b[0m, in \u001b[0;36m_linspace_scalar\u001b[1;34m(start, stop, num, endpoint, retstep, dtype)\u001b[0m\n\u001b[0;32m 87\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m dtype \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m 88\u001b[0m \u001b[38;5;66;03m# In actual implementation, only float is used\u001b[39;00m\n\u001b[0;32m 89\u001b[0m dtype \u001b[38;5;241m=\u001b[39m dt\n\u001b[1;32m---> 91\u001b[0m ret \u001b[38;5;241m=\u001b[39m cupy\u001b[38;5;241m.\u001b[39mempty((num,), dtype\u001b[38;5;241m=\u001b[39mdt)\n\u001b[0;32m 92\u001b[0m div \u001b[38;5;241m=\u001b[39m (num \u001b[38;5;241m-\u001b[39m \u001b[38;5;241m1\u001b[39m) \u001b[38;5;28;01mif\u001b[39;00m endpoint \u001b[38;5;28;01melse\u001b[39;00m num\n\u001b[0;32m 93\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m div \u001b[38;5;241m<\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m:\n",
|
|
"File \u001b[1;32mc:\\Users\\naeve\\anaconda3\\Lib\\site-packages\\cupy\\_creation\\basic.py:32\u001b[0m, in \u001b[0;36mempty\u001b[1;34m(shape, dtype, order)\u001b[0m\n\u001b[0;32m 13\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mempty\u001b[39m(\n\u001b[0;32m 14\u001b[0m shape: _ShapeLike,\n\u001b[0;32m 15\u001b[0m dtype: DTypeLike \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mfloat\u001b[39m,\n\u001b[0;32m 16\u001b[0m order: _OrderCF \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mC\u001b[39m\u001b[38;5;124m'\u001b[39m,\n\u001b[0;32m 17\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m NDArray[Any]:\n\u001b[0;32m 18\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Returns an array without initializing the elements.\u001b[39;00m\n\u001b[0;32m 19\u001b[0m \n\u001b[0;32m 20\u001b[0m \u001b[38;5;124;03m Args:\u001b[39;00m\n\u001b[1;32m (...)\u001b[0m\n\u001b[0;32m 30\u001b[0m \n\u001b[0;32m 31\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[1;32m---> 32\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m cupy\u001b[38;5;241m.\u001b[39mndarray(shape, dtype, order\u001b[38;5;241m=\u001b[39morder)\n",
|
|
"File \u001b[1;32mcupy\\\\_core\\\\core.pyx:151\u001b[0m, in \u001b[0;36mcupy._core.core.ndarray.__new__\u001b[1;34m()\u001b[0m\n",
|
|
"File \u001b[1;32mcupy\\\\_core\\\\core.pyx:239\u001b[0m, in \u001b[0;36mcupy._core.core._ndarray_base._init\u001b[1;34m()\u001b[0m\n",
|
|
"File \u001b[1;32mcupy\\\\cuda\\\\memory.pyx:738\u001b[0m, in \u001b[0;36mcupy.cuda.memory.alloc\u001b[1;34m()\u001b[0m\n",
|
|
"File \u001b[1;32mcupy\\\\cuda\\\\memory.pyx:1424\u001b[0m, in \u001b[0;36mcupy.cuda.memory.MemoryPool.malloc\u001b[1;34m()\u001b[0m\n",
|
|
"File \u001b[1;32mcupy\\\\cuda\\\\memory.pyx:1444\u001b[0m, in \u001b[0;36mcupy.cuda.memory.MemoryPool.malloc\u001b[1;34m()\u001b[0m\n",
|
|
"File \u001b[1;32mcupy\\\\cuda\\\\device.pyx:40\u001b[0m, in \u001b[0;36mcupy.cuda.device.get_device_id\u001b[1;34m()\u001b[0m\n",
|
|
"File \u001b[1;32mcupy_backends\\\\cuda\\\\api\\\\runtime.pyx:202\u001b[0m, in \u001b[0;36mcupy_backends.cuda.api.runtime.getDevice\u001b[1;34m()\u001b[0m\n",
|
|
"File \u001b[1;32mcupy_backends\\\\cuda\\\\api\\\\runtime.pyx:146\u001b[0m, in \u001b[0;36mcupy_backends.cuda.api.runtime.check_status\u001b[1;34m()\u001b[0m\n",
|
|
"\u001b[1;31mCUDARuntimeError\u001b[0m: cudaErrorInsufficientDriver: CUDA driver version is insufficient for CUDA runtime version"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# Generate spatial grid\n",
|
|
"x = np.linspace(*extend[0], n_pot_steps[0])\n",
|
|
"y = np.linspace(*extend[1], n_pot_steps[1])\n",
|
|
"z = np.linspace(*extend[2], n_pot_steps[2])\n",
|
|
"\n",
|
|
"x3D, y3D, z3D = np.meshgrid(x, y, z) # Ensure correct indexing\n",
|
|
"\n",
|
|
"# Compute potential (Replace with actual function)\n",
|
|
"pot = potential(x3D, y3D, z3D)\n",
|
|
"\n",
|
|
"state_number = 0\n",
|
|
"\n",
|
|
"# Create figure and axis\n",
|
|
"fig, ax = plt.subplots()\n",
|
|
"im = ax.imshow(states[state_number, :, :, 0], extent=[*extend[0], *extend[1]], origin=\"lower\",\n",
|
|
" vmin=np.min(states[state_number]), vmax=np.max(states[state_number]))\n",
|
|
"\n",
|
|
"plt.xlabel(\"x\")\n",
|
|
"plt.ylabel(\"y\")\n",
|
|
"plt.colorbar(im)\n",
|
|
"\n",
|
|
"# Initialize contour as None before defining it globally\n",
|
|
"contour = None\n",
|
|
"\n",
|
|
"# Animation update function\n",
|
|
"def update(frame):\n",
|
|
" global contour # Ensure we're modifying the global variable\n",
|
|
"\n",
|
|
" im.set_data(states[state_number, :, :, frame]) # Update image data\n",
|
|
" ax.set_title(f\"z={z[frame]/si.um:.3f}um\") # Update title\n",
|
|
"\n",
|
|
" # Remove old contours if they exist\n",
|
|
" if contour is not None:\n",
|
|
" for c in contour.collections:\n",
|
|
" c.remove()\n",
|
|
"\n",
|
|
" # Redraw contour plot\n",
|
|
" contour = ax.contour(pot[:, :, frame], levels=10, colors='white', linewidths=0.7, extent=[*extend[0], *extend[1]])\n",
|
|
"\n",
|
|
"# Create the first contour plot after defining update()\n",
|
|
"contour = ax.contour(pot[:, :, 0], levels=10, colors='white', linewidths=0.7, extent=[*extend[0], *extend[1]])\n",
|
|
"\n",
|
|
"# Create animation\n",
|
|
"frames = n_pot_steps[2] # Number of slices\n",
|
|
"ani = animation.FuncAnimation(fig, update, frames=frames, interval=100)\n",
|
|
"\n",
|
|
"ani.save(f\"animations/test{state_number}.gif\", writer=\"pillow\", fps=frames/5) # Save as GIF\n",
|
|
"\n",
|
|
"plt.show()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "base",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.12.7"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|