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
OMEGA_X = 502.6548245743669
OMEGA_Y = 6283.185307179586
def round_sig(x: float, sig: int = 2):
37def round_sig(x: float, sig: int = 2):
38    return round(x, sig - int(floor(log10(abs(x)))) - 1)
def TF_dist_2D(xy_mesh, TFpOD, xc, yc, rx, ry, os):
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:
def Gaussian_dist_2D(xy_mesh, GpOD, xc, yc, sigx, sigy, 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:
def bimodal_dist_2D(xy_mesh, GpOD, sigx, sigy, TFpOD, rx, ry, 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:
class OqtantJob(bert_schemas.job.JobCreate):
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
external_id: uuid.UUID | None
time_submit: datetime.datetime | None
pix_cal: float
sig_abs: float
omega_x: float
omega_y: float
def TOF(self):
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)

def IF(self):
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)

@staticmethod
def get_image_space( datafile=array([[0., 0., 0., ..., 0., 0., 0.], [0., 0., 0., ..., 0., 0., 0.], [0., 0., 0., ..., 0., 0., 0.], ..., [0., 0., 0., ..., 0., 0., 0.], [0., 0., 0., ..., 0., 0., 0.], [0., 0., 0., ..., 0., 0., 0.]]), centered='y'):
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:
def fit_bimodal_data2D(self, xi=None, lb=None, ub=None):
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:
def plot_fit_results(self, fit_params, model='bimodal', file_name=None, plot_title=None):
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)

def atoms_2dplot(self, file_name=None, figsize=(12, 12), gridon=False):
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
def atoms_sliceplot(self, file_name=None, sliceaxis='x', gridon=False):
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
def atoms_3dplot(self, file_name=None, view_angle=-45, figsize=(10, 10)):
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:
def atom_numbers(self, bimodalfit_params, print_results=True):
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]
def calculate_temperature(self, bimodalfit_params):
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
class OqtantJob.Config:
141    class Config:
142        validate_assignment = True
validate_assignment = True