oqtant .schemas .job
1# Copyright 2023 Infleqtion 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# https://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15from datetime import datetime 16from math import floor, log10, pi 17from uuid import UUID 18 19import matplotlib.pyplot as plt 20import numpy as np 21import scipy.optimize as opt 22from bert_schemas import job as job_schema 23from matplotlib.ticker import FormatStrFormatter, LinearLocator 24from pydantic import Field 25 26from oqtant.util.exceptions import ( 27 JobPlotFitError, 28 JobPlotFitMismatchError, 29 JobReadError, 30) 31 32OMEGA_X = 2 * pi * 80 33OMEGA_Y = 2 * pi * 1000 34 35 36def round_sig(x: float, sig: int = 2): 37 return round(x, sig - int(floor(log10(abs(x)))) - 1) 38 39 40def TF_dist_2D(xy_mesh, TFpOD, xc, yc, rx, ry, os): 41 """ 42 Defines 2D Thomas-Fermi distribution characteristic of zero-temperature Bose-gas 43 Requires function(s): get_image_space 44 :param xy_mesh: 45 :type xy_mesh: (2,N,M) Matrix of floats containing mesh grid of image coordinates 46 :param TFpOD: 47 :type TFpOD: float - Thomas-Fermi peak Optical Density (OD) 48 :param rx: 49 :type rx: float - Thomas-Fermi radius along the x-direction 50 :param ry: 51 :type ry: float - Thomas-Fermi radius along the y-direction (along gravity) 52 :param xc: 53 :type xc: float - Cloud center along the x-direction (along gravity) 54 :param yc: 55 :type yc: float - Cloud center along the y-direction 56 :param os: 57 :type os: float - Constant offset 58 """ 59 60 # unpack 1D list into 2D x and y coords 61 (x, y) = xy_mesh 62 63 # Simplify Thomas-Fermi expression 64 A = 1 - ((y - yc) / rx) ** 2 - ((x - xc) / ry) ** 2 65 66 # make 2D Thomas-Fermi distribution 67 OD = np.real(TFpOD * np.maximum(np.sign(A) * (np.abs(A)) ** (3 / 2), 0)) + os 68 69 # flatten the 2D Gaussian down to 1D 70 return OD.ravel() 71 72 73def Gaussian_dist_2D(xy_mesh, GpOD, xc, yc, sigx, sigy, os): 74 """ 75 Defines 2D gaussian distribution characteristic of a thermal ensemble of atoms 76 Requires function(s): get_image_space 77 :param xy_mesh: 78 :type xy_mesh: (2,N,M) Matrix of floats containing meshgrid of image coordinates 79 :param GpOD: 80 :type GpOD: float - Gaussian peak Optical Density (OD) 81 :param sigx: 82 :type sigx: float - Gaussian spread along the x-direction 83 :param sigy: 84 :type sigy: float - Gaussian spread along the y-direction (along gravity) 85 :param xc: 86 :type xc: float - Cloud center along the x-direction (along gravity) 87 :param yc: 88 :type yc: float - Cloud center along the y-direction 89 :param os: 90 :type os: float - Constant offset 91 """ 92 93 (x, y) = xy_mesh 94 95 OD = ( 96 GpOD * np.exp(-0.5 * ((y - yc) / sigy) ** 2 - 0.5 * ((x - xc) / sigx) ** 2) + os 97 ) 98 return OD.ravel() 99 100 101def bimodal_dist_2D(xy_mesh, GpOD, sigx, sigy, TFpOD, rx, ry, xc, yc, os): 102 """ 103 Defines 2D bimodal distribution characteristic of finite-temperature Bose-gas 104 Requires functions: Gaussian_dist_2D, TF_dist_2D, get_image_space 105 :param xy_mesh: 106 :type xy_mesh: (2,N,M) Matrix of floats containing meshgrid of image coordinates 107 :param GpOD: 108 :type GpOD: float - Gaussian peak Optical Density (OD) 109 :param sigx: 110 :type sigx: float - Gaussian spread along the x-direction 111 :param sigy: 112 :type sigy: float - Gaussian spread along the y-direction (along gravity) 113 :param TFpOD: 114 :type TFpOD: float - Thomas-Fermi peak Optical Density (OD) 115 :param rx: 116 :type rx: float - Thomas-Fermi radius along the x-direction 117 :param ry: 118 :type ry: float - Thomas-Fermi radius along the y-direction (along gravity) 119 :param xc: 120 :type xc: float - Cloud center along the x-direction (along gravity) 121 :param yc: 122 :type yc: float - Cloud center along the y-direction 123 :param os: 124 :type os: float - Constant offset 125 """ 126 127 return Gaussian_dist_2D(xy_mesh, GpOD, xc, yc, sigx, sigy, os) + TF_dist_2D( 128 xy_mesh, TFpOD, xc, yc, rx, ry, os 129 ) 130 131 132class OqtantJob(job_schema.JobCreate): 133 external_id: UUID | None 134 time_submit: datetime | None 135 pix_cal: float = Field(8.71, const=True) 136 sig_abs: float = Field(0.297, const=True) 137 omega_x: float = Field(OMEGA_X, const=True) 138 omega_y: float = Field(OMEGA_Y, const=True) 139 140 class Config: 141 validate_assignment = True 142 143 def TOF(self): 144 """ 145 Returns shaped TOF image if it exists 146 :returns: reshaped pixels numpy array (100,100) 147 148 """ 149 if self.status != job_schema.JobStatus.COMPLETE: 150 raise JobReadError("Job results not available.") 151 try: 152 reshaped_pixels = np.array( 153 self.inputs[0].output.values.tof_image.pixels 154 ).reshape(100, 100) 155 except Exception as exc: 156 raise JobReadError("no TOF results") from exc 157 return reshaped_pixels 158 159 def IF(self): 160 """ 161 Returns shaped IT image if it exists 162 :returns: reshaped pixels numpy array (100,100) 163 164 """ 165 if self.status != job_schema.JobStatus.COMPLETE: 166 raise JobReadError("Job results not available.") 167 168 try: 169 reshaped_pixels = np.array( 170 self.inputs[0].output.values.it_plot.pixels 171 ).reshape(148, 512) 172 except Exception as exc: 173 raise JobReadError("no IT results") from exc 174 return reshaped_pixels 175 176 @staticmethod 177 def get_image_space(datafile=np.zeros((100, 100)), centered="y"): 178 """ 179 Returns meshgrid of image coordinates 180 :param datafile: 181 :type datafile: (N,M) Matrix of Optical Density (OD) Data 182 """ 183 184 lx, ly = np.shape(datafile) 185 x, y = np.arange(lx), np.arange(ly) 186 187 if centered == "y": 188 x = x - round(lx / 2) 189 y = y - round(ly / 2) 190 191 xy_mesh = np.meshgrid(x, y) 192 193 return xy_mesh, lx, ly 194 195 def fit_bimodal_data2D(self, xi=None, lb=None, ub=None): 196 """ 197 Performs fit via trust region reflective algorithm. 198 Requires functions: bimodal_dist_2D, Gaussian_dist_2D, TF_dist_2D, get_image_space 199 For better fit performance, tune initial guess 'xi' and lower/upper bounds, 'lb' and 'ub' 200 :param xy_mesh: 201 :type xy_mesh: (2,N,M) Matrix containing meshgrid of image data coordinates 202 :param data2D: 203 :type data2D: (N,M) Matrix containing image data 204 :param xi: 205 :type xi: (1,9) List of fit parameter initial guesses 206 :param lb: 207 :type lb: (1,9) List of fit parameter lower bounds 208 :param ub: 209 :type ub: (1,9) List of fit parameter upper bounds 210 """ 211 xi = xi if xi else [0.25, 8, 8, 1, 4, 6, 0, 0, 0.02] 212 lb = lb if lb else [0, 7, 7, 0, 2, 2, -20, -20, 0] 213 ub = ub if ub else [2, 20, 20, 2, 20, 20, 20, 20, 1] 214 215 TOF_data = self.TOF() 216 xy_mesh, _, _ = self.get_image_space() # TOF_data) 217 218 (X, Y) = xy_mesh 219 x = X[0] 220 y = Y[:, 0] 221 222 fit_params, cov_mat = opt.curve_fit( 223 bimodal_dist_2D, xy_mesh, np.ravel(TOF_data), p0=xi, bounds=(lb, ub) 224 ) 225 fit_residual = TOF_data - bimodal_dist_2D(xy_mesh, *fit_params).reshape( 226 np.outer(x, y).shape 227 ) 228 fit_Rsquared = 1 - np.var(fit_residual) / np.var(TOF_data) 229 230 return fit_params, cov_mat, fit_residual, fit_Rsquared 231 232 def plot_fit_results( 233 self, fit_params, model="bimodal", file_name=None, plot_title=None 234 ): 235 """ 236 Plot the results of a fit operation 237 238 :param fit_params: 239 :type fit_params: list of parameters from a fit operation 240 :param model: 241 :type model: string "bimodal", "TF", or "gaussian". default "bimodal" 242 :param output: 243 :type output: valid filename 244 :param plot_title: 245 :type plot_title: string title for the plot. 246 default "job: "+str(self.name)+"\nTOF fit: "+str(model) 247 248 """ 249 250 xy_mesh, _, _ = self.get_image_space() # TOF_data) 251 252 (X, Y) = xy_mesh 253 254 if model == "bimodal": 255 try: 256 m = bimodal_dist_2D(xy_mesh, *fit_params) 257 except TypeError as exc: 258 raise JobPlotFitMismatchError() from exc 259 except Exception as exc: 260 raise JobPlotFitError() from exc 261 262 elif model == "gaussian": 263 try: 264 m = Gaussian_dist_2D(xy_mesh, *fit_params) 265 except TypeError as exc: 266 raise TypeError( 267 "PLOT FIT RESULTS: mismatched parameters and model type" 268 ) from exc 269 except Exception as exc: 270 raise JobPlotFitError() from exc 271 elif model == "TF": 272 try: 273 m = TF_dist_2D(xy_mesh, *fit_params) 274 except TypeError as exc: 275 raise JobPlotFitMismatchError() from exc 276 except Exception as exc: 277 raise JobPlotFitError() from exc 278 else: 279 print( 280 f"PLOT FIT RESULTS: Invalid model specified: {model}.", 281 " Select 'bimodal', 'gaussian', or 'TF'", 282 ) 283 return 284 285 m = m.reshape(100, 100) 286 plt.figure() 287 plt.imshow( 288 m, 289 origin="lower", 290 cmap="nipy_spectral", 291 extent=[ 292 np.min(X) * self.pix_cal, 293 np.max(X) * self.pix_cal, 294 np.min(Y) * self.pix_cal, 295 np.max(Y) * self.pix_cal, 296 ], 297 ) 298 299 if plot_title is None: 300 plot_title = f"job: {self.name}\nTOF fit: {model}" 301 302 plt.title(plot_title) 303 304 if file_name: 305 self._save_plot_file(plt, file_name) 306 plt.show() 307 308 @staticmethod 309 def _save_plot_file(plot, file_name): 310 file = f"{file_name}.png" 311 try: 312 plot.savefig(file) 313 print(f"plot saved to file: {file}") 314 except (FileNotFoundError, Exception): 315 print(f"failed to save plot at {file}") 316 317 def atoms_2dplot(self, file_name=None, figsize=(12, 12), gridon=False): 318 """ 319 Generate a 2D plot of atom OD (save or show) 320 321 :param output: how to output the information 322 :type output: string "show" or valid filename 323 :param figsize: 324 :type figsize: tuple. default is (12,12) 325 :param gridon: grid lines on plot on/off 326 :type gridon: Boolean. default is False 327 328 """ 329 330 TOF_data = self.TOF() 331 xy_mesh, _, _ = self.get_image_space() # TOF_data) 332 (X, Y) = xy_mesh 333 334 fig2D = plt.figure(figsize=figsize) 335 ax = fig2D.gca() 336 plt2D = plt.imshow( 337 TOF_data, 338 origin="lower", 339 cmap="nipy_spectral", 340 extent=[ 341 np.min(X) * self.pix_cal, 342 np.max(X) * self.pix_cal, 343 np.min(Y) * self.pix_cal, 344 np.max(Y) * self.pix_cal, 345 ], 346 ) 347 plt.grid(b=gridon) 348 plt.colorbar(plt2D, shrink=0.8) 349 350 ax.set_xlabel("x-position, micron", labelpad=15, fontsize=16) 351 ax.set_ylabel("y-position, micron", labelpad=15, fontsize=16) 352 plt.title("Time of Flight Optical Depth", fontsize=16) 353 354 print("Peak OD: ", np.max(TOF_data)) 355 356 if file_name: 357 self._save_plot_file(plt, file_name) 358 plt.show() 359 360 def atoms_sliceplot(self, file_name=None, sliceaxis="x", gridon=False): 361 """ 362 Generate a 1D slice plot of atom OD in x or y 363 364 :param output: how to output the information 365 :type output: string "show" or valid filename 366 :param sliceaxis: 367 :type sliceaxis: string 'x' or 'y' 368 :param figsize: 369 :type figsize: tuple. default is (12,12) 370 :param gridon: grid lines on plot on/off 371 :type gridon: Boolean. default is False 372 373 """ 374 TOF_data = self.TOF() 375 xy_mesh, lx, ly = self.get_image_space(TOF_data) 376 (X, Y) = xy_mesh 377 378 params, *_ = self.fit_bimodal_data2D() 379 fitOD = bimodal_dist_2D(xy_mesh, *params) 380 381 Gfit_params = [params[0], params[6], params[7], params[1], params[2], params[8]] 382 fitODG = Gaussian_dist_2D(xy_mesh, *Gfit_params) 383 384 # Reshape Fit Distributions to 2D form 385 fitOD2D = fitOD.reshape(lx, ly) 386 fitODG2D = fitODG.reshape(lx, ly) 387 388 # Define Central slices 389 xslice = fitOD2D[int(lx / 2), :] 390 yslice = fitOD2D[:, int(ly / 2)] 391 xsliceG = fitODG2D[int(lx / 2), :] 392 ysliceG = fitODG2D[:, int(ly / 2)] 393 394 if sliceaxis == "x": 395 xsliceD = TOF_data[int(len(X[1]) / 2), :] 396 xslice = fitOD2D[int(len(X[1]) / 2), :] 397 xsliceG = fitODG2D[int(len(X[1]) / 2), :] 398 plt.title("X-Slice", fontsize=16) 399 plt.plot(X[1] * self.pix_cal, xsliceD, "ok") 400 plt.plot(X[1] * self.pix_cal, xslice, "b") 401 plt.plot(X[1] * self.pix_cal, xsliceG, "r") 402 elif sliceaxis == "y": 403 ysliceD = TOF_data[:, int(len(Y[1]) / 2)] 404 yslice = fitOD2D[:, int(len(Y[1]) / 2)] 405 ysliceG = fitODG2D[:, int(len(Y[1]) / 2)] 406 plt.title("Y-Slice", fontsize=16) 407 plt.plot(Y[:, 1] * self.pix_cal, ysliceD, "ok") 408 plt.plot(Y[:, 1] * self.pix_cal, yslice, "b") 409 plt.plot(Y[:, 1] * self.pix_cal, ysliceG, "r") 410 else: 411 raise ValueError("Input either x or y") 412 413 plt.grid(b=gridon) 414 plt.xlabel("x-position, micron", labelpad=15, fontsize=16) 415 plt.ylabel("Optical Depth", labelpad=15, fontsize=16) 416 417 if file_name: 418 self._save_plot_file(plt, file_name) 419 plt.show() 420 421 # This function plots the optical depth as a 3D surface with projected density contours 422 def atoms_3dplot(self, file_name=None, view_angle=-45, figsize=(10, 10)): 423 """ 424 Generate a 3D slice plot of atom OD 425 426 :param output: how to output the information 427 :type output: string "show" or valid filename 428 :param view_angle: 429 :type view_angle: int (-180, 180). default -45 430 :param figsize: 431 :type figsize: tuple. default is (10,10) 432 433 """ 434 435 fig3d = plt.figure(figsize=figsize) 436 ax = fig3d.gca(projection="3d") 437 438 ax.zaxis.set_major_locator(LinearLocator(10)) 439 ax.zaxis.set_major_formatter(FormatStrFormatter("%.02f")) 440 ax.xaxis.pane.fill = False 441 ax.yaxis.pane.fill = False 442 ax.zaxis.pane.fill = False 443 444 # Set axis labels 445 ax.set_xlabel("x-position, micron", labelpad=10) 446 ax.set_ylabel("y-position, micron", labelpad=10) 447 ax.set_zlabel("Optical Depth", labelpad=10) 448 449 # rotate the axes and update 450 ax.view_init(30, view_angle) 451 452 if file_name: 453 self._save_plot_file(plt, file_name) 454 plt.show() 455 456 def atom_numbers(self, bimodalfit_params, print_results=True): 457 xy_mesh, lx, ly = self.get_image_space() 458 459 # lmfit Toolbox Results 460 # Remove Background 461 bimodalfit_params[8] = bimodalfit_params[8] * 0 462 463 # Define Fit Distributions in 1D form 464 fitOD = bimodal_dist_2D(xy_mesh, *bimodalfit_params) 465 Gfit_params = [ 466 bimodalfit_params[0], 467 bimodalfit_params[6], 468 bimodalfit_params[7], 469 bimodalfit_params[1], 470 bimodalfit_params[2], 471 bimodalfit_params[8], 472 ] 473 TFfit_params = [ 474 bimodalfit_params[3], 475 bimodalfit_params[6], 476 bimodalfit_params[7], 477 bimodalfit_params[4], 478 bimodalfit_params[5], 479 bimodalfit_params[8], 480 ] 481 fitODG = Gaussian_dist_2D(xy_mesh, *Gfit_params) 482 fitODTF = TF_dist_2D(xy_mesh, *TFfit_params) 483 lmfitbg = lx * ly * bimodalfit_params[8] 484 485 # Reshape Fit Distributions to 2D form 486 fitOD2D = fitOD.reshape(lx, ly) 487 fitODG2D = fitODG.reshape(lx, ly) 488 fitODTF2D = fitODTF.reshape(lx, ly) 489 490 ODsumtot = np.sum(fitOD2D) - lmfitbg * 2 491 ODsumcond = np.sum(fitODTF2D) - lmfitbg 492 ODsumtherm = np.sum(fitODG2D) - lmfitbg 493 Ncond = self.pix_cal**2 / self.sig_abs * ODsumcond # Condensed atom number 494 Ntherm = self.pix_cal**2 / self.sig_abs * ODsumtherm # Thermal atom number 495 Ntot = self.pix_cal**2 / self.sig_abs * ODsumtot # Thermal atom number 496 497 if print_results: 498 print("Total Atom Number: ", int(Ntot)) 499 print("Condensed Atom Number: ", int(Ncond)) 500 print("Thermal Atom Number: ", int(Ntherm)) 501 502 return [Ntot, Ncond, Ntherm] 503 504 def calculate_temperature(self, bimodalfit_params): 505 Gfit_params = [ 506 bimodalfit_params[0], 507 bimodalfit_params[6], 508 bimodalfit_params[7], 509 bimodalfit_params[1], 510 bimodalfit_params[2], 511 bimodalfit_params[8], 512 ] 513 514 TOF = 0 515 TOF = self.input.time_of_flight_ms * 1e-3 516 # Time of flight 517 tau_x = self.omega_x * TOF # Dimensionless expansion coefficient 518 tau_y = self.omega_y * TOF 519 pixcal = 8.71 * 1e-6 520 # pixel calibration - input sizes as # of pixels 521 mRb = 87 * 1.66 * 1e-27 522 # 87Rb mass 523 kb = 1.3806 * 1e-23 524 525 sigma_x = Gfit_params[3] * pixcal 526 sigma_y = Gfit_params[4] * pixcal 527 528 T_x = ( 529 (mRb / 2 / kb) 530 * (np.square(self.omega_x) * np.square(sigma_x)) 531 / (1 + np.square(self.omega_x) * np.square(TOF)) 532 ) 533 T_y = ( 534 (mRb / 2 / kb) 535 * (np.square(self.omega_y) * np.square(sigma_y)) 536 / (1 + np.square(self.omega_y) * np.square(TOF)) 537 ) 538 539 T = ( 540 2 * np.square(tau_x) / (1 + 3 * np.square(tau_x)) * T_x 541 + (1 + np.square(tau_y)) / (1 + 3 * np.square(tau_y)) * T_y 542 ) 543 T_nK = T * 1e9 544 545 return T_nK
41def TF_dist_2D(xy_mesh, TFpOD, xc, yc, rx, ry, os): 42 """ 43 Defines 2D Thomas-Fermi distribution characteristic of zero-temperature Bose-gas 44 Requires function(s): get_image_space 45 :param xy_mesh: 46 :type xy_mesh: (2,N,M) Matrix of floats containing mesh grid of image coordinates 47 :param TFpOD: 48 :type TFpOD: float - Thomas-Fermi peak Optical Density (OD) 49 :param rx: 50 :type rx: float - Thomas-Fermi radius along the x-direction 51 :param ry: 52 :type ry: float - Thomas-Fermi radius along the y-direction (along gravity) 53 :param xc: 54 :type xc: float - Cloud center along the x-direction (along gravity) 55 :param yc: 56 :type yc: float - Cloud center along the y-direction 57 :param os: 58 :type os: float - Constant offset 59 """ 60 61 # unpack 1D list into 2D x and y coords 62 (x, y) = xy_mesh 63 64 # Simplify Thomas-Fermi expression 65 A = 1 - ((y - yc) / rx) ** 2 - ((x - xc) / ry) ** 2 66 67 # make 2D Thomas-Fermi distribution 68 OD = np.real(TFpOD * np.maximum(np.sign(A) * (np.abs(A)) ** (3 / 2), 0)) + os 69 70 # flatten the 2D Gaussian down to 1D 71 return OD.ravel()
Defines 2D Thomas-Fermi distribution characteristic of zero-temperature Bose-gas Requires function(s): get_image_space
Parameters
- xy_mesh:
- TFpOD:
- rx:
- ry:
- xc:
- yc:
- os:
74def Gaussian_dist_2D(xy_mesh, GpOD, xc, yc, sigx, sigy, os): 75 """ 76 Defines 2D gaussian distribution characteristic of a thermal ensemble of atoms 77 Requires function(s): get_image_space 78 :param xy_mesh: 79 :type xy_mesh: (2,N,M) Matrix of floats containing meshgrid of image coordinates 80 :param GpOD: 81 :type GpOD: float - Gaussian peak Optical Density (OD) 82 :param sigx: 83 :type sigx: float - Gaussian spread along the x-direction 84 :param sigy: 85 :type sigy: float - Gaussian spread along the y-direction (along gravity) 86 :param xc: 87 :type xc: float - Cloud center along the x-direction (along gravity) 88 :param yc: 89 :type yc: float - Cloud center along the y-direction 90 :param os: 91 :type os: float - Constant offset 92 """ 93 94 (x, y) = xy_mesh 95 96 OD = ( 97 GpOD * np.exp(-0.5 * ((y - yc) / sigy) ** 2 - 0.5 * ((x - xc) / sigx) ** 2) + os 98 ) 99 return OD.ravel()
Defines 2D gaussian distribution characteristic of a thermal ensemble of atoms Requires function(s): get_image_space
Parameters
- xy_mesh:
- GpOD:
- sigx:
- sigy:
- xc:
- yc:
- os:
102def bimodal_dist_2D(xy_mesh, GpOD, sigx, sigy, TFpOD, rx, ry, xc, yc, os): 103 """ 104 Defines 2D bimodal distribution characteristic of finite-temperature Bose-gas 105 Requires functions: Gaussian_dist_2D, TF_dist_2D, get_image_space 106 :param xy_mesh: 107 :type xy_mesh: (2,N,M) Matrix of floats containing meshgrid of image coordinates 108 :param GpOD: 109 :type GpOD: float - Gaussian peak Optical Density (OD) 110 :param sigx: 111 :type sigx: float - Gaussian spread along the x-direction 112 :param sigy: 113 :type sigy: float - Gaussian spread along the y-direction (along gravity) 114 :param TFpOD: 115 :type TFpOD: float - Thomas-Fermi peak Optical Density (OD) 116 :param rx: 117 :type rx: float - Thomas-Fermi radius along the x-direction 118 :param ry: 119 :type ry: float - Thomas-Fermi radius along the y-direction (along gravity) 120 :param xc: 121 :type xc: float - Cloud center along the x-direction (along gravity) 122 :param yc: 123 :type yc: float - Cloud center along the y-direction 124 :param os: 125 :type os: float - Constant offset 126 """ 127 128 return Gaussian_dist_2D(xy_mesh, GpOD, xc, yc, sigx, sigy, os) + TF_dist_2D( 129 xy_mesh, TFpOD, xc, yc, rx, ry, os 130 )
Defines 2D bimodal distribution characteristic of finite-temperature Bose-gas Requires functions: Gaussian_dist_2D, TF_dist_2D, get_image_space
Parameters
- xy_mesh:
- GpOD:
- sigx:
- sigy:
- TFpOD:
- rx:
- ry:
- xc:
- yc:
- os:
133class OqtantJob(job_schema.JobCreate): 134 external_id: UUID | None 135 time_submit: datetime | None 136 pix_cal: float = Field(8.71, const=True) 137 sig_abs: float = Field(0.297, const=True) 138 omega_x: float = Field(OMEGA_X, const=True) 139 omega_y: float = Field(OMEGA_Y, const=True) 140 141 class Config: 142 validate_assignment = True 143 144 def TOF(self): 145 """ 146 Returns shaped TOF image if it exists 147 :returns: reshaped pixels numpy array (100,100) 148 149 """ 150 if self.status != job_schema.JobStatus.COMPLETE: 151 raise JobReadError("Job results not available.") 152 try: 153 reshaped_pixels = np.array( 154 self.inputs[0].output.values.tof_image.pixels 155 ).reshape(100, 100) 156 except Exception as exc: 157 raise JobReadError("no TOF results") from exc 158 return reshaped_pixels 159 160 def IF(self): 161 """ 162 Returns shaped IT image if it exists 163 :returns: reshaped pixels numpy array (100,100) 164 165 """ 166 if self.status != job_schema.JobStatus.COMPLETE: 167 raise JobReadError("Job results not available.") 168 169 try: 170 reshaped_pixels = np.array( 171 self.inputs[0].output.values.it_plot.pixels 172 ).reshape(148, 512) 173 except Exception as exc: 174 raise JobReadError("no IT results") from exc 175 return reshaped_pixels 176 177 @staticmethod 178 def get_image_space(datafile=np.zeros((100, 100)), centered="y"): 179 """ 180 Returns meshgrid of image coordinates 181 :param datafile: 182 :type datafile: (N,M) Matrix of Optical Density (OD) Data 183 """ 184 185 lx, ly = np.shape(datafile) 186 x, y = np.arange(lx), np.arange(ly) 187 188 if centered == "y": 189 x = x - round(lx / 2) 190 y = y - round(ly / 2) 191 192 xy_mesh = np.meshgrid(x, y) 193 194 return xy_mesh, lx, ly 195 196 def fit_bimodal_data2D(self, xi=None, lb=None, ub=None): 197 """ 198 Performs fit via trust region reflective algorithm. 199 Requires functions: bimodal_dist_2D, Gaussian_dist_2D, TF_dist_2D, get_image_space 200 For better fit performance, tune initial guess 'xi' and lower/upper bounds, 'lb' and 'ub' 201 :param xy_mesh: 202 :type xy_mesh: (2,N,M) Matrix containing meshgrid of image data coordinates 203 :param data2D: 204 :type data2D: (N,M) Matrix containing image data 205 :param xi: 206 :type xi: (1,9) List of fit parameter initial guesses 207 :param lb: 208 :type lb: (1,9) List of fit parameter lower bounds 209 :param ub: 210 :type ub: (1,9) List of fit parameter upper bounds 211 """ 212 xi = xi if xi else [0.25, 8, 8, 1, 4, 6, 0, 0, 0.02] 213 lb = lb if lb else [0, 7, 7, 0, 2, 2, -20, -20, 0] 214 ub = ub if ub else [2, 20, 20, 2, 20, 20, 20, 20, 1] 215 216 TOF_data = self.TOF() 217 xy_mesh, _, _ = self.get_image_space() # TOF_data) 218 219 (X, Y) = xy_mesh 220 x = X[0] 221 y = Y[:, 0] 222 223 fit_params, cov_mat = opt.curve_fit( 224 bimodal_dist_2D, xy_mesh, np.ravel(TOF_data), p0=xi, bounds=(lb, ub) 225 ) 226 fit_residual = TOF_data - bimodal_dist_2D(xy_mesh, *fit_params).reshape( 227 np.outer(x, y).shape 228 ) 229 fit_Rsquared = 1 - np.var(fit_residual) / np.var(TOF_data) 230 231 return fit_params, cov_mat, fit_residual, fit_Rsquared 232 233 def plot_fit_results( 234 self, fit_params, model="bimodal", file_name=None, plot_title=None 235 ): 236 """ 237 Plot the results of a fit operation 238 239 :param fit_params: 240 :type fit_params: list of parameters from a fit operation 241 :param model: 242 :type model: string "bimodal", "TF", or "gaussian". default "bimodal" 243 :param output: 244 :type output: valid filename 245 :param plot_title: 246 :type plot_title: string title for the plot. 247 default "job: "+str(self.name)+"\nTOF fit: "+str(model) 248 249 """ 250 251 xy_mesh, _, _ = self.get_image_space() # TOF_data) 252 253 (X, Y) = xy_mesh 254 255 if model == "bimodal": 256 try: 257 m = bimodal_dist_2D(xy_mesh, *fit_params) 258 except TypeError as exc: 259 raise JobPlotFitMismatchError() from exc 260 except Exception as exc: 261 raise JobPlotFitError() from exc 262 263 elif model == "gaussian": 264 try: 265 m = Gaussian_dist_2D(xy_mesh, *fit_params) 266 except TypeError as exc: 267 raise TypeError( 268 "PLOT FIT RESULTS: mismatched parameters and model type" 269 ) from exc 270 except Exception as exc: 271 raise JobPlotFitError() from exc 272 elif model == "TF": 273 try: 274 m = TF_dist_2D(xy_mesh, *fit_params) 275 except TypeError as exc: 276 raise JobPlotFitMismatchError() from exc 277 except Exception as exc: 278 raise JobPlotFitError() from exc 279 else: 280 print( 281 f"PLOT FIT RESULTS: Invalid model specified: {model}.", 282 " Select 'bimodal', 'gaussian', or 'TF'", 283 ) 284 return 285 286 m = m.reshape(100, 100) 287 plt.figure() 288 plt.imshow( 289 m, 290 origin="lower", 291 cmap="nipy_spectral", 292 extent=[ 293 np.min(X) * self.pix_cal, 294 np.max(X) * self.pix_cal, 295 np.min(Y) * self.pix_cal, 296 np.max(Y) * self.pix_cal, 297 ], 298 ) 299 300 if plot_title is None: 301 plot_title = f"job: {self.name}\nTOF fit: {model}" 302 303 plt.title(plot_title) 304 305 if file_name: 306 self._save_plot_file(plt, file_name) 307 plt.show() 308 309 @staticmethod 310 def _save_plot_file(plot, file_name): 311 file = f"{file_name}.png" 312 try: 313 plot.savefig(file) 314 print(f"plot saved to file: {file}") 315 except (FileNotFoundError, Exception): 316 print(f"failed to save plot at {file}") 317 318 def atoms_2dplot(self, file_name=None, figsize=(12, 12), gridon=False): 319 """ 320 Generate a 2D plot of atom OD (save or show) 321 322 :param output: how to output the information 323 :type output: string "show" or valid filename 324 :param figsize: 325 :type figsize: tuple. default is (12,12) 326 :param gridon: grid lines on plot on/off 327 :type gridon: Boolean. default is False 328 329 """ 330 331 TOF_data = self.TOF() 332 xy_mesh, _, _ = self.get_image_space() # TOF_data) 333 (X, Y) = xy_mesh 334 335 fig2D = plt.figure(figsize=figsize) 336 ax = fig2D.gca() 337 plt2D = plt.imshow( 338 TOF_data, 339 origin="lower", 340 cmap="nipy_spectral", 341 extent=[ 342 np.min(X) * self.pix_cal, 343 np.max(X) * self.pix_cal, 344 np.min(Y) * self.pix_cal, 345 np.max(Y) * self.pix_cal, 346 ], 347 ) 348 plt.grid(b=gridon) 349 plt.colorbar(plt2D, shrink=0.8) 350 351 ax.set_xlabel("x-position, micron", labelpad=15, fontsize=16) 352 ax.set_ylabel("y-position, micron", labelpad=15, fontsize=16) 353 plt.title("Time of Flight Optical Depth", fontsize=16) 354 355 print("Peak OD: ", np.max(TOF_data)) 356 357 if file_name: 358 self._save_plot_file(plt, file_name) 359 plt.show() 360 361 def atoms_sliceplot(self, file_name=None, sliceaxis="x", gridon=False): 362 """ 363 Generate a 1D slice plot of atom OD in x or y 364 365 :param output: how to output the information 366 :type output: string "show" or valid filename 367 :param sliceaxis: 368 :type sliceaxis: string 'x' or 'y' 369 :param figsize: 370 :type figsize: tuple. default is (12,12) 371 :param gridon: grid lines on plot on/off 372 :type gridon: Boolean. default is False 373 374 """ 375 TOF_data = self.TOF() 376 xy_mesh, lx, ly = self.get_image_space(TOF_data) 377 (X, Y) = xy_mesh 378 379 params, *_ = self.fit_bimodal_data2D() 380 fitOD = bimodal_dist_2D(xy_mesh, *params) 381 382 Gfit_params = [params[0], params[6], params[7], params[1], params[2], params[8]] 383 fitODG = Gaussian_dist_2D(xy_mesh, *Gfit_params) 384 385 # Reshape Fit Distributions to 2D form 386 fitOD2D = fitOD.reshape(lx, ly) 387 fitODG2D = fitODG.reshape(lx, ly) 388 389 # Define Central slices 390 xslice = fitOD2D[int(lx / 2), :] 391 yslice = fitOD2D[:, int(ly / 2)] 392 xsliceG = fitODG2D[int(lx / 2), :] 393 ysliceG = fitODG2D[:, int(ly / 2)] 394 395 if sliceaxis == "x": 396 xsliceD = TOF_data[int(len(X[1]) / 2), :] 397 xslice = fitOD2D[int(len(X[1]) / 2), :] 398 xsliceG = fitODG2D[int(len(X[1]) / 2), :] 399 plt.title("X-Slice", fontsize=16) 400 plt.plot(X[1] * self.pix_cal, xsliceD, "ok") 401 plt.plot(X[1] * self.pix_cal, xslice, "b") 402 plt.plot(X[1] * self.pix_cal, xsliceG, "r") 403 elif sliceaxis == "y": 404 ysliceD = TOF_data[:, int(len(Y[1]) / 2)] 405 yslice = fitOD2D[:, int(len(Y[1]) / 2)] 406 ysliceG = fitODG2D[:, int(len(Y[1]) / 2)] 407 plt.title("Y-Slice", fontsize=16) 408 plt.plot(Y[:, 1] * self.pix_cal, ysliceD, "ok") 409 plt.plot(Y[:, 1] * self.pix_cal, yslice, "b") 410 plt.plot(Y[:, 1] * self.pix_cal, ysliceG, "r") 411 else: 412 raise ValueError("Input either x or y") 413 414 plt.grid(b=gridon) 415 plt.xlabel("x-position, micron", labelpad=15, fontsize=16) 416 plt.ylabel("Optical Depth", labelpad=15, fontsize=16) 417 418 if file_name: 419 self._save_plot_file(plt, file_name) 420 plt.show() 421 422 # This function plots the optical depth as a 3D surface with projected density contours 423 def atoms_3dplot(self, file_name=None, view_angle=-45, figsize=(10, 10)): 424 """ 425 Generate a 3D slice plot of atom OD 426 427 :param output: how to output the information 428 :type output: string "show" or valid filename 429 :param view_angle: 430 :type view_angle: int (-180, 180). default -45 431 :param figsize: 432 :type figsize: tuple. default is (10,10) 433 434 """ 435 436 fig3d = plt.figure(figsize=figsize) 437 ax = fig3d.gca(projection="3d") 438 439 ax.zaxis.set_major_locator(LinearLocator(10)) 440 ax.zaxis.set_major_formatter(FormatStrFormatter("%.02f")) 441 ax.xaxis.pane.fill = False 442 ax.yaxis.pane.fill = False 443 ax.zaxis.pane.fill = False 444 445 # Set axis labels 446 ax.set_xlabel("x-position, micron", labelpad=10) 447 ax.set_ylabel("y-position, micron", labelpad=10) 448 ax.set_zlabel("Optical Depth", labelpad=10) 449 450 # rotate the axes and update 451 ax.view_init(30, view_angle) 452 453 if file_name: 454 self._save_plot_file(plt, file_name) 455 plt.show() 456 457 def atom_numbers(self, bimodalfit_params, print_results=True): 458 xy_mesh, lx, ly = self.get_image_space() 459 460 # lmfit Toolbox Results 461 # Remove Background 462 bimodalfit_params[8] = bimodalfit_params[8] * 0 463 464 # Define Fit Distributions in 1D form 465 fitOD = bimodal_dist_2D(xy_mesh, *bimodalfit_params) 466 Gfit_params = [ 467 bimodalfit_params[0], 468 bimodalfit_params[6], 469 bimodalfit_params[7], 470 bimodalfit_params[1], 471 bimodalfit_params[2], 472 bimodalfit_params[8], 473 ] 474 TFfit_params = [ 475 bimodalfit_params[3], 476 bimodalfit_params[6], 477 bimodalfit_params[7], 478 bimodalfit_params[4], 479 bimodalfit_params[5], 480 bimodalfit_params[8], 481 ] 482 fitODG = Gaussian_dist_2D(xy_mesh, *Gfit_params) 483 fitODTF = TF_dist_2D(xy_mesh, *TFfit_params) 484 lmfitbg = lx * ly * bimodalfit_params[8] 485 486 # Reshape Fit Distributions to 2D form 487 fitOD2D = fitOD.reshape(lx, ly) 488 fitODG2D = fitODG.reshape(lx, ly) 489 fitODTF2D = fitODTF.reshape(lx, ly) 490 491 ODsumtot = np.sum(fitOD2D) - lmfitbg * 2 492 ODsumcond = np.sum(fitODTF2D) - lmfitbg 493 ODsumtherm = np.sum(fitODG2D) - lmfitbg 494 Ncond = self.pix_cal**2 / self.sig_abs * ODsumcond # Condensed atom number 495 Ntherm = self.pix_cal**2 / self.sig_abs * ODsumtherm # Thermal atom number 496 Ntot = self.pix_cal**2 / self.sig_abs * ODsumtot # Thermal atom number 497 498 if print_results: 499 print("Total Atom Number: ", int(Ntot)) 500 print("Condensed Atom Number: ", int(Ncond)) 501 print("Thermal Atom Number: ", int(Ntherm)) 502 503 return [Ntot, Ncond, Ntherm] 504 505 def calculate_temperature(self, bimodalfit_params): 506 Gfit_params = [ 507 bimodalfit_params[0], 508 bimodalfit_params[6], 509 bimodalfit_params[7], 510 bimodalfit_params[1], 511 bimodalfit_params[2], 512 bimodalfit_params[8], 513 ] 514 515 TOF = 0 516 TOF = self.input.time_of_flight_ms * 1e-3 517 # Time of flight 518 tau_x = self.omega_x * TOF # Dimensionless expansion coefficient 519 tau_y = self.omega_y * TOF 520 pixcal = 8.71 * 1e-6 521 # pixel calibration - input sizes as # of pixels 522 mRb = 87 * 1.66 * 1e-27 523 # 87Rb mass 524 kb = 1.3806 * 1e-23 525 526 sigma_x = Gfit_params[3] * pixcal 527 sigma_y = Gfit_params[4] * pixcal 528 529 T_x = ( 530 (mRb / 2 / kb) 531 * (np.square(self.omega_x) * np.square(sigma_x)) 532 / (1 + np.square(self.omega_x) * np.square(TOF)) 533 ) 534 T_y = ( 535 (mRb / 2 / kb) 536 * (np.square(self.omega_y) * np.square(sigma_y)) 537 / (1 + np.square(self.omega_y) * np.square(TOF)) 538 ) 539 540 T = ( 541 2 * np.square(tau_x) / (1 + 3 * np.square(tau_x)) * T_x 542 + (1 + np.square(tau_y)) / (1 + 3 * np.square(tau_y)) * T_y 543 ) 544 T_nK = T * 1e9 545 546 return T_nK
144 def TOF(self): 145 """ 146 Returns shaped TOF image if it exists 147 :returns: reshaped pixels numpy array (100,100) 148 149 """ 150 if self.status != job_schema.JobStatus.COMPLETE: 151 raise JobReadError("Job results not available.") 152 try: 153 reshaped_pixels = np.array( 154 self.inputs[0].output.values.tof_image.pixels 155 ).reshape(100, 100) 156 except Exception as exc: 157 raise JobReadError("no TOF results") from exc 158 return reshaped_pixels
Returns shaped TOF image if it exists :returns: reshaped pixels numpy array (100,100)
160 def IF(self): 161 """ 162 Returns shaped IT image if it exists 163 :returns: reshaped pixels numpy array (100,100) 164 165 """ 166 if self.status != job_schema.JobStatus.COMPLETE: 167 raise JobReadError("Job results not available.") 168 169 try: 170 reshaped_pixels = np.array( 171 self.inputs[0].output.values.it_plot.pixels 172 ).reshape(148, 512) 173 except Exception as exc: 174 raise JobReadError("no IT results") from exc 175 return reshaped_pixels
Returns shaped IT image if it exists :returns: reshaped pixels numpy array (100,100)
177 @staticmethod 178 def get_image_space(datafile=np.zeros((100, 100)), centered="y"): 179 """ 180 Returns meshgrid of image coordinates 181 :param datafile: 182 :type datafile: (N,M) Matrix of Optical Density (OD) Data 183 """ 184 185 lx, ly = np.shape(datafile) 186 x, y = np.arange(lx), np.arange(ly) 187 188 if centered == "y": 189 x = x - round(lx / 2) 190 y = y - round(ly / 2) 191 192 xy_mesh = np.meshgrid(x, y) 193 194 return xy_mesh, lx, ly
Returns meshgrid of image coordinates
Parameters
- datafile:
196 def fit_bimodal_data2D(self, xi=None, lb=None, ub=None): 197 """ 198 Performs fit via trust region reflective algorithm. 199 Requires functions: bimodal_dist_2D, Gaussian_dist_2D, TF_dist_2D, get_image_space 200 For better fit performance, tune initial guess 'xi' and lower/upper bounds, 'lb' and 'ub' 201 :param xy_mesh: 202 :type xy_mesh: (2,N,M) Matrix containing meshgrid of image data coordinates 203 :param data2D: 204 :type data2D: (N,M) Matrix containing image data 205 :param xi: 206 :type xi: (1,9) List of fit parameter initial guesses 207 :param lb: 208 :type lb: (1,9) List of fit parameter lower bounds 209 :param ub: 210 :type ub: (1,9) List of fit parameter upper bounds 211 """ 212 xi = xi if xi else [0.25, 8, 8, 1, 4, 6, 0, 0, 0.02] 213 lb = lb if lb else [0, 7, 7, 0, 2, 2, -20, -20, 0] 214 ub = ub if ub else [2, 20, 20, 2, 20, 20, 20, 20, 1] 215 216 TOF_data = self.TOF() 217 xy_mesh, _, _ = self.get_image_space() # TOF_data) 218 219 (X, Y) = xy_mesh 220 x = X[0] 221 y = Y[:, 0] 222 223 fit_params, cov_mat = opt.curve_fit( 224 bimodal_dist_2D, xy_mesh, np.ravel(TOF_data), p0=xi, bounds=(lb, ub) 225 ) 226 fit_residual = TOF_data - bimodal_dist_2D(xy_mesh, *fit_params).reshape( 227 np.outer(x, y).shape 228 ) 229 fit_Rsquared = 1 - np.var(fit_residual) / np.var(TOF_data) 230 231 return fit_params, cov_mat, fit_residual, fit_Rsquared
Performs fit via trust region reflective algorithm. Requires functions: bimodal_dist_2D, Gaussian_dist_2D, TF_dist_2D, get_image_space For better fit performance, tune initial guess 'xi' and lower/upper bounds, 'lb' and 'ub'
Parameters
- xy_mesh:
- data2D:
- xi:
- lb:
- ub:
233 def plot_fit_results( 234 self, fit_params, model="bimodal", file_name=None, plot_title=None 235 ): 236 """ 237 Plot the results of a fit operation 238 239 :param fit_params: 240 :type fit_params: list of parameters from a fit operation 241 :param model: 242 :type model: string "bimodal", "TF", or "gaussian". default "bimodal" 243 :param output: 244 :type output: valid filename 245 :param plot_title: 246 :type plot_title: string title for the plot. 247 default "job: "+str(self.name)+"\nTOF fit: "+str(model) 248 249 """ 250 251 xy_mesh, _, _ = self.get_image_space() # TOF_data) 252 253 (X, Y) = xy_mesh 254 255 if model == "bimodal": 256 try: 257 m = bimodal_dist_2D(xy_mesh, *fit_params) 258 except TypeError as exc: 259 raise JobPlotFitMismatchError() from exc 260 except Exception as exc: 261 raise JobPlotFitError() from exc 262 263 elif model == "gaussian": 264 try: 265 m = Gaussian_dist_2D(xy_mesh, *fit_params) 266 except TypeError as exc: 267 raise TypeError( 268 "PLOT FIT RESULTS: mismatched parameters and model type" 269 ) from exc 270 except Exception as exc: 271 raise JobPlotFitError() from exc 272 elif model == "TF": 273 try: 274 m = TF_dist_2D(xy_mesh, *fit_params) 275 except TypeError as exc: 276 raise JobPlotFitMismatchError() from exc 277 except Exception as exc: 278 raise JobPlotFitError() from exc 279 else: 280 print( 281 f"PLOT FIT RESULTS: Invalid model specified: {model}.", 282 " Select 'bimodal', 'gaussian', or 'TF'", 283 ) 284 return 285 286 m = m.reshape(100, 100) 287 plt.figure() 288 plt.imshow( 289 m, 290 origin="lower", 291 cmap="nipy_spectral", 292 extent=[ 293 np.min(X) * self.pix_cal, 294 np.max(X) * self.pix_cal, 295 np.min(Y) * self.pix_cal, 296 np.max(Y) * self.pix_cal, 297 ], 298 ) 299 300 if plot_title is None: 301 plot_title = f"job: {self.name}\nTOF fit: {model}" 302 303 plt.title(plot_title) 304 305 if file_name: 306 self._save_plot_file(plt, file_name) 307 plt.show()
Plot the results of a fit operation
:param fit_params:
:type fit_params: list of parameters from a fit operation
:param model:
:type model: string "bimodal", "TF", or "gaussian". default "bimodal"
:param output:
:type output: valid filename
:param plot_title:
:type plot_title: string title for the plot.
default "job: "+str(self.name)+"
TOF fit: "+str(model)
318 def atoms_2dplot(self, file_name=None, figsize=(12, 12), gridon=False): 319 """ 320 Generate a 2D plot of atom OD (save or show) 321 322 :param output: how to output the information 323 :type output: string "show" or valid filename 324 :param figsize: 325 :type figsize: tuple. default is (12,12) 326 :param gridon: grid lines on plot on/off 327 :type gridon: Boolean. default is False 328 329 """ 330 331 TOF_data = self.TOF() 332 xy_mesh, _, _ = self.get_image_space() # TOF_data) 333 (X, Y) = xy_mesh 334 335 fig2D = plt.figure(figsize=figsize) 336 ax = fig2D.gca() 337 plt2D = plt.imshow( 338 TOF_data, 339 origin="lower", 340 cmap="nipy_spectral", 341 extent=[ 342 np.min(X) * self.pix_cal, 343 np.max(X) * self.pix_cal, 344 np.min(Y) * self.pix_cal, 345 np.max(Y) * self.pix_cal, 346 ], 347 ) 348 plt.grid(b=gridon) 349 plt.colorbar(plt2D, shrink=0.8) 350 351 ax.set_xlabel("x-position, micron", labelpad=15, fontsize=16) 352 ax.set_ylabel("y-position, micron", labelpad=15, fontsize=16) 353 plt.title("Time of Flight Optical Depth", fontsize=16) 354 355 print("Peak OD: ", np.max(TOF_data)) 356 357 if file_name: 358 self._save_plot_file(plt, file_name) 359 plt.show()
Generate a 2D plot of atom OD (save or show)
Parameters
- output: how to output the information
- figsize:
- gridon: grid lines on plot on/off
361 def atoms_sliceplot(self, file_name=None, sliceaxis="x", gridon=False): 362 """ 363 Generate a 1D slice plot of atom OD in x or y 364 365 :param output: how to output the information 366 :type output: string "show" or valid filename 367 :param sliceaxis: 368 :type sliceaxis: string 'x' or 'y' 369 :param figsize: 370 :type figsize: tuple. default is (12,12) 371 :param gridon: grid lines on plot on/off 372 :type gridon: Boolean. default is False 373 374 """ 375 TOF_data = self.TOF() 376 xy_mesh, lx, ly = self.get_image_space(TOF_data) 377 (X, Y) = xy_mesh 378 379 params, *_ = self.fit_bimodal_data2D() 380 fitOD = bimodal_dist_2D(xy_mesh, *params) 381 382 Gfit_params = [params[0], params[6], params[7], params[1], params[2], params[8]] 383 fitODG = Gaussian_dist_2D(xy_mesh, *Gfit_params) 384 385 # Reshape Fit Distributions to 2D form 386 fitOD2D = fitOD.reshape(lx, ly) 387 fitODG2D = fitODG.reshape(lx, ly) 388 389 # Define Central slices 390 xslice = fitOD2D[int(lx / 2), :] 391 yslice = fitOD2D[:, int(ly / 2)] 392 xsliceG = fitODG2D[int(lx / 2), :] 393 ysliceG = fitODG2D[:, int(ly / 2)] 394 395 if sliceaxis == "x": 396 xsliceD = TOF_data[int(len(X[1]) / 2), :] 397 xslice = fitOD2D[int(len(X[1]) / 2), :] 398 xsliceG = fitODG2D[int(len(X[1]) / 2), :] 399 plt.title("X-Slice", fontsize=16) 400 plt.plot(X[1] * self.pix_cal, xsliceD, "ok") 401 plt.plot(X[1] * self.pix_cal, xslice, "b") 402 plt.plot(X[1] * self.pix_cal, xsliceG, "r") 403 elif sliceaxis == "y": 404 ysliceD = TOF_data[:, int(len(Y[1]) / 2)] 405 yslice = fitOD2D[:, int(len(Y[1]) / 2)] 406 ysliceG = fitODG2D[:, int(len(Y[1]) / 2)] 407 plt.title("Y-Slice", fontsize=16) 408 plt.plot(Y[:, 1] * self.pix_cal, ysliceD, "ok") 409 plt.plot(Y[:, 1] * self.pix_cal, yslice, "b") 410 plt.plot(Y[:, 1] * self.pix_cal, ysliceG, "r") 411 else: 412 raise ValueError("Input either x or y") 413 414 plt.grid(b=gridon) 415 plt.xlabel("x-position, micron", labelpad=15, fontsize=16) 416 plt.ylabel("Optical Depth", labelpad=15, fontsize=16) 417 418 if file_name: 419 self._save_plot_file(plt, file_name) 420 plt.show()
Generate a 1D slice plot of atom OD in x or y
Parameters
- output: how to output the information
- sliceaxis:
- figsize:
- gridon: grid lines on plot on/off
423 def atoms_3dplot(self, file_name=None, view_angle=-45, figsize=(10, 10)): 424 """ 425 Generate a 3D slice plot of atom OD 426 427 :param output: how to output the information 428 :type output: string "show" or valid filename 429 :param view_angle: 430 :type view_angle: int (-180, 180). default -45 431 :param figsize: 432 :type figsize: tuple. default is (10,10) 433 434 """ 435 436 fig3d = plt.figure(figsize=figsize) 437 ax = fig3d.gca(projection="3d") 438 439 ax.zaxis.set_major_locator(LinearLocator(10)) 440 ax.zaxis.set_major_formatter(FormatStrFormatter("%.02f")) 441 ax.xaxis.pane.fill = False 442 ax.yaxis.pane.fill = False 443 ax.zaxis.pane.fill = False 444 445 # Set axis labels 446 ax.set_xlabel("x-position, micron", labelpad=10) 447 ax.set_ylabel("y-position, micron", labelpad=10) 448 ax.set_zlabel("Optical Depth", labelpad=10) 449 450 # rotate the axes and update 451 ax.view_init(30, view_angle) 452 453 if file_name: 454 self._save_plot_file(plt, file_name) 455 plt.show()
Generate a 3D slice plot of atom OD
Parameters
- output: how to output the information
- view_angle:
- figsize:
457 def atom_numbers(self, bimodalfit_params, print_results=True): 458 xy_mesh, lx, ly = self.get_image_space() 459 460 # lmfit Toolbox Results 461 # Remove Background 462 bimodalfit_params[8] = bimodalfit_params[8] * 0 463 464 # Define Fit Distributions in 1D form 465 fitOD = bimodal_dist_2D(xy_mesh, *bimodalfit_params) 466 Gfit_params = [ 467 bimodalfit_params[0], 468 bimodalfit_params[6], 469 bimodalfit_params[7], 470 bimodalfit_params[1], 471 bimodalfit_params[2], 472 bimodalfit_params[8], 473 ] 474 TFfit_params = [ 475 bimodalfit_params[3], 476 bimodalfit_params[6], 477 bimodalfit_params[7], 478 bimodalfit_params[4], 479 bimodalfit_params[5], 480 bimodalfit_params[8], 481 ] 482 fitODG = Gaussian_dist_2D(xy_mesh, *Gfit_params) 483 fitODTF = TF_dist_2D(xy_mesh, *TFfit_params) 484 lmfitbg = lx * ly * bimodalfit_params[8] 485 486 # Reshape Fit Distributions to 2D form 487 fitOD2D = fitOD.reshape(lx, ly) 488 fitODG2D = fitODG.reshape(lx, ly) 489 fitODTF2D = fitODTF.reshape(lx, ly) 490 491 ODsumtot = np.sum(fitOD2D) - lmfitbg * 2 492 ODsumcond = np.sum(fitODTF2D) - lmfitbg 493 ODsumtherm = np.sum(fitODG2D) - lmfitbg 494 Ncond = self.pix_cal**2 / self.sig_abs * ODsumcond # Condensed atom number 495 Ntherm = self.pix_cal**2 / self.sig_abs * ODsumtherm # Thermal atom number 496 Ntot = self.pix_cal**2 / self.sig_abs * ODsumtot # Thermal atom number 497 498 if print_results: 499 print("Total Atom Number: ", int(Ntot)) 500 print("Condensed Atom Number: ", int(Ncond)) 501 print("Thermal Atom Number: ", int(Ntherm)) 502 503 return [Ntot, Ncond, Ntherm]
505 def calculate_temperature(self, bimodalfit_params): 506 Gfit_params = [ 507 bimodalfit_params[0], 508 bimodalfit_params[6], 509 bimodalfit_params[7], 510 bimodalfit_params[1], 511 bimodalfit_params[2], 512 bimodalfit_params[8], 513 ] 514 515 TOF = 0 516 TOF = self.input.time_of_flight_ms * 1e-3 517 # Time of flight 518 tau_x = self.omega_x * TOF # Dimensionless expansion coefficient 519 tau_y = self.omega_y * TOF 520 pixcal = 8.71 * 1e-6 521 # pixel calibration - input sizes as # of pixels 522 mRb = 87 * 1.66 * 1e-27 523 # 87Rb mass 524 kb = 1.3806 * 1e-23 525 526 sigma_x = Gfit_params[3] * pixcal 527 sigma_y = Gfit_params[4] * pixcal 528 529 T_x = ( 530 (mRb / 2 / kb) 531 * (np.square(self.omega_x) * np.square(sigma_x)) 532 / (1 + np.square(self.omega_x) * np.square(TOF)) 533 ) 534 T_y = ( 535 (mRb / 2 / kb) 536 * (np.square(self.omega_y) * np.square(sigma_y)) 537 / (1 + np.square(self.omega_y) * np.square(TOF)) 538 ) 539 540 T = ( 541 2 * np.square(tau_x) / (1 + 3 * np.square(tau_x)) * T_x 542 + (1 + np.square(tau_y)) / (1 + 3 * np.square(tau_y)) * T_y 543 ) 544 T_nK = T * 1e9 545 546 return T_nK
Inherited Members
- pydantic.main.BaseModel
- BaseModel
- dict
- json
- parse_obj
- parse_raw
- parse_file
- from_orm
- construct
- copy
- schema
- schema_json
- validate
- update_forward_refs
- bert_schemas.job.JobCreate
- name
- job_type
- inputs
- set_input_count_and_run
- validate_inputs_match_job_type
- bert_schemas.job.JobBase
- origin
- status
- display
- qpu_name
- input_count