pspinor package

Submodules

pspinor.pspinor module

class spinor_gpe.pspinor.pspinor.PSpinor(path, omeg=None, g_sc=None, mesh_points=(256, 256), r_sizes=(16, 16), atom_num=10000.0, pop_frac=(0.5, 0.5), **kwargs)

Bases: object

A GPU-compatible simulator of the pseudospin-1/2 GPE.

Contains the functionality to run a real- or imaginary-time propataion of the pseudospin-1/2 Gross-Pitaevskii equation. Contains methods to generate the required energy and spatial grids. Also has methods to generate the grids for the momentum-(in)dependent coupling between spin components, corresponding to an (RF) Raman coupling interaction.

The dominant length scale is in terms of the harmonic oscillator length along the x-direction a_x. The dominant energy scale is the harmonic trapping energy along the x-direction [hbar * omeg[‘x’]].

paths

Essential paths and directory names for a given trial.

  • ‘folder’ : The name of the topmost trial directory.

  • ‘data’ : Path to the simulation data & results.

  • ‘code’ : Path to the subdirectory containing the trial code.

  • ‘trial’ : Path to the subdirectory containing raw trial data.

Type

dict

atom_num

Total atom number.

Type

float

space

Spatial arrays, meshes, spacings, volume elements, and sizes:

KEYS

Arrays:

‘x’

‘y’

‘kx’

‘ky’

Meshes:

‘x_mesh’

‘y_mesh’

‘kx_mesh’

‘ky_mesh’

Spacings:

‘dr’

‘dk’

Vol. elem.:

‘dv_r’

‘dv_k’

Sizes:

‘r_sizes’

‘k_sizes’

Other:

‘mesh_points’
Type

dict

pop_frac

The initial population fraction in each spin component.

Type

iterable

omeg

Angular trapping frequencies, {‘x’, ‘y’, ‘z’}. omeg[‘x’] multiplied by \hbar is the characteristic energy scale.

Type

dict

g_sc

Relative scattering interaction strengths, {‘uu’, ‘dd’, ‘ud’}.

Type

dict

a_x

Harmonic oscillator length along the x-axis; this is the characteristic length scale.

Type

float

chem_pot

Chemical potential, [\hbar * omeg[‘x’]].

Type

float

rad_tf

Thomas-Fermi radius along the x-axis, [a_x].

Type

float

time_scale

Inverse x-axis trapping frequency - the characteristic time scale, [1 / omeg[‘x’]].

Type

float

pot_eng_spin

A list of 2D potential energy grids for each spin component, [\hbar * omeg[‘x’]].

Type

list of array

kin_eng_spin

A list of 2D kinetic energy grids for each spin component, [\hbar * omeg[‘x’]].

Type

list of array

psi

A list of 2D real-space spin component wavefunctions; generally complex.

Type

list of array

psik

A list of 2D momentum-space spin component wavefunctions; generally complex.

Type

list of array

is_coupling

Signals the presence of direct coupling between spin components.

Type

bool

kL_recoil

The value of the single-photon recoil momentum, [1 / a_x].

Type

float

EL_recoil

The energy scale corresponding to the single-photon recoil momentum, [1 / omeg[‘x’]].

Type

float

rand_seed

Value to seed the pseudorandom number generator.

Type

int

prop

Object containing the results of propagation, along with analysis methods.

Type

PropResult

rot_coupling

Option to place coupling in a rotating reference frame, i.e. no momentum shift on the coupling operation.

Type

bool, default=True

compute_energy_grids()

Compute the initial potential and kinetic energy grids.

Assumes that the BEC is in a harmonic trap. This harmonic potential determines the initial ‘Thomas-Fermi’ density profile of the BEC. pot_eng can be modified prior to progation to have any arbitrary potential energy landscape.

Also assumes that the BEC has a simple free-particle kinetic energy dispersion. If using a momentum-dependent spin coupling, this grid will be modified later.

compute_spatial_grids(mesh_points=(256, 256), r_sizes=(16, 16))

Compute the real and momentum space grids.

Stored in the dict space are the real- and momentum-space mesh grids, mesh sizes, mesh spacings, volume elements, and the corresponding linear arrays.

Parameters

  • mesh_points (iterable of int, default=(256, 256)) – The number of grid points along the x- and y-axes, respectively.

  • r_sizes (iterable of int, default=(16, 16)) – The half size of the grid along the real x- and y-axes, respectively,in units of [a_x].

compute_tf_params(species='Rb87')

Compute parameters and scales for the Thomas-Fermi solution.

Parameters

sepecies (str, default=’Rb87’) – Designates the atomic species and corresponding physical data used in the simulations.

compute_tf_psi(phase_factor=1.0)

Compute the intial pseudospinor wavefunction psi and FFT psik.

The psuedospinor wavefunction psi that is generated is a list of 2D NumPy arrays. The Thomas-Fermi solution is real and has the form of an inverted parabaloid. psik is a list of the 2D FFT of psi’s components.

Parameters

phase_factor (complex, default=1.0) – Unit complex number; initial relative phase factor between the two spin components.

property coupling

Get the coupling attribute.

2D coupling array [\hbar * omeg[‘x’]].

coupling_grad(slope, offset, axis=1)

Generate a linear gradient of the interspin coupling strength.

Convenience function for generating linear gradients of the coupling. coupling can also be set to any arbitrary NumPy array directly:

>>> ps = PSpinor()
>>> ps.coupling_setup()
>>> ps.coupling = np.exp(-ps.x_mesh**2 / 2)  # Gaussian function

Note

When working with Raman recoil units [E_L], they will first need to be converted to [hbar*omeg_x] units.

Parameters

  • slope (float) – The slope of the coupling gradient, in [hbar*omeg_x/a_x].

  • offset (float) – The origin offset of the coupling gradient, in [hbar*omeg_x].

  • axis (int, optional) – The axis along which the coupling gradient runs.

coupling_setup(wavel=7.901e-07, scale=1.0, kin_shift=False)

Calculate parameters for the momentum-(in)dependent coupling.

If kin_shift=True, the kinetic energy grid receives a spin-dependent shift in momentum.

Parameters

  • wavel (float, default=790.1e-9) – Wavelength of Raman coupling. Note that you must specify the wavelength in meters.

  • scale (float, default=1.0) – The relative scale of recoil momentum. Interesting physics may be simulated by considering a recoil momentum that is hypothetically much larger or much smaller than the native wavelength recoil momentum.

  • kin_shift (bool, default=False) – Option for a momentum-(in)dependent coupling.

coupling_uniform(value)

Generate a uniform interspin coupling strength.

Convenience function for generating unirom gradients of the coupling. coupling can also be set to any arbitrary NumPy array directly.

Parameters

value (float) – The value of the coupling, in [hbar*omega_x].

See also

coupling_grad

Coupling gradient

property detuning

Get the detuning attribute.

2D detuning array [\hbar * omeg[‘x’]].

detuning_grad(slope, offset=0.0, axis=1)

Generate a linear gradient of the interspin coupling strength.

Convenience function for generating linear gradients of the coupling. detuning can also be set to any arbitrary NumPy array directly:

>>> ps = PSpinor()
>>> ps.coupling_setup()
>>> ps.detuning = np.sin(2 * np.pi * ps.x_mesh)  # Sin function

Note

when working with Raman recoil units [E_L], they will first need to be converted to [hbar*omeg_x] units.

Parameters

  • slope (float) – The slope of the detuning gradient, in [hbar*omeg_x/a_x].

  • offset (float) – The origin offset of the detuning gradient, in [hbar*omeg_x].

  • axis (int, optional) – The axis along which the detuning gradient runs.

See also

coupling_grad

Coupling gradient

detuning_uniform(value)

Generate a uniform coupling detuning.

Convenience function for generating unirom gradients of the coupling. detuning can also be set to any arbitrary NumPy array directly.

Parameters

value (float) – The value of the coupling, in [hbar*omega_x].

See also

coupling_grad

Coupling gradient

imaginary(t_step, n_steps=1000, device='cpu', is_sampling=False, n_samples=1)

Perform imaginary-time propagation.

Propagation is carried out in a TensorPropagator object. The results are stored and returned in the PropResult object for further analysis.

Parameters

  • t_step (float) – The propagation time step.

  • n_steps (int, optional) – The total number of propagation steps.

  • device (str, optional) – {‘cpu’, ‘cuda’}

  • is_sampling (bool, optional) – Option to sample wavefunctions throughout the propagation.

  • n_samples (int, optional) – The number of samples to collect.

property kin_eng

Get the kin_eng attribute.

2D kinetic energy grid, [\hbar * omeg[‘x’]]0

no_coupling_setup()

Provide the default parameters for no coupling.

plot_kdens(psik=None, spin=None, cmap='viridis', scale=1.0)

Plot the momentum-space density of the wavefunction.

Shows the momentum-space density of either the up (spin=0), down (spin=1), or both (spin=None) spin components.

Parameters

  • psik (list of Numpy array, optional.) – The wavefunction to plot. If no psik is supplied, then it uses the object attribute self.psik.

  • spin (int or None, optional) – Which spin to plot. None plots both spins. 0 or 1 plots only the up or down spin, respectively.

  • cmap (str, default=’viridis’) – The matplotlib colormap to use for the plots.

  • scale (float, default=1.0) – A factor by which to scale the spatial dimensions, e.g. Thomas- Fermi radius.

See also

plotting_tools.plot_dens

Density plots.

plot_rdens(psi=None, spin=None, cmap='viridis', scale=1.0)

Plot the real-space density of the wavefunction.

Shows the real-space density of either the up (spin=0), down (spin=1), or both (spin=None) spin components.

Parameters

  • psi (list of Numpy array, optional.) – The wavefunction to plot. If no psi is supplied, then it uses the object attribute self.psi.

  • spin (int or None, optional) – Which spin to plot. None plots both spins. 0 or 1 plots only the up or down spin, respectively.

  • cmap (str, default=’viridis’) – The matplotlib colormap to use for the plots.

  • scale (float, default=1.0) – A factor by which to scale the spatial dimensions, e.g. Thomas- Fermi radius.

See also

plotting_tools.plot_dens

Density plots.

plot_rphase(psi=None, spin=None, cmap='twilight_shifted', scale=1.0)

Plot the real-space phase of the wavefunction.

Shows the real-space phase of either the up (spin=0), down (spin=1), or both (spin=None) spin components.

Parameters

  • psi (list of Numpy array, optional.) – The wavefunction to plot. If no psi is supplied, then it uses the object attribute self.psi.

  • spin (int or None, optional) – Which spin to plot. None plots both spins. 0 or 1 plots only the up or down spin, respectively.

  • cmap (str, default=’twilight_shifted’) – The matplotlib colormap to use for the plots.

  • scale (float, default=1.0) – A factor by which to scale the spatial dimensions, e.g. Thomas- Fermi radius.

See also

plotting_tools.plot_phase

Phase plots.

plot_spins(rscale=1.0, kscale=1.0, cmap='viridis', save=True, ext='.pdf', zoom=1.0)

Plot the densities (both real & k) and phases of spin components.

Parameters

  • rscale (float, optional) – Real-space length scale. The default of 1.0 corresponds to the naturatl harmonic length scale along the x-axis.

  • kscale (float, optional) – Momentum-space length scale. The default of 1.0 corresponds to the inverse harmonic length scale along the x-axis.

  • cmap (str, optional) – Color map name for the real- and momentum-space density plots.

  • save (bool, optional) – Saves the figure as a .pdf file (default). The filename has the format “/data_path/pop_evolution%s-trial_name.pdf”.

  • ext (str, optional) – Saved plot image file extension.

  • zoom (float, optional) – A zoom factor for the k-space density plot.

property pot_eng

Get the pot_eng attribute.

2D potential energy grid, [\hbar * omeg[‘x’]].

real(t_step, n_steps=1000, device='cpu', is_sampling=False, n_samples=1)

Perform real-time propagation.

Propagation is carried out in a TensorPropagator object. The results are stored and returned in the PropResult object for further analysis.

Parameters

  • t_step (float) – The propagation time step.

  • n_steps (int, optional) – The total number of propagation steps.

  • device (str, optional) – {‘cpu’, ‘cuda’}

  • is_sampling (bool, optional) – Option to sample wavefunctions throughout the propagation.

  • n_samples (int, optional) – The number of samples to collect.

seed_random_vortices(N)

Seed randomly-arranged vortices into the wavefunction.

seed_regular_vortices()

Seed regularly-arranged vortices into the wavefunction.

These seed-vortex functions might be moved to the ttools module.

seed_vortices(positions, windings)

Seeds vortices at the positions specified.

Pass a list of tuple coordinates

Parameters

  • positions (list of tuple) – The positions at which to seed the vortices.

  • windings (int or list of int) – Must have values of +1 or -1. If a single int is supplied, all vortices will have the same winding. If a list is supplied, it must have the same length as positions.

  • TODO (#) –

  • component. (# in each spinor) –

  • ??? (#) –

setup_data_path(path, overwrite)

Create new data directory to store simulation data & results.

Parameters

  • path (str) – The name of the directory to save the simulation. If path does not represent an absolute path, then the data is stored in spinor-gpe/data/path.

  • overwrite (bool) – Gives the option to overwrite existing data sub-directories

shift_momentum(psik=None, scale=1.0, frac=(0.5, 0.5))

Shifts momentum components of psi by a fracion of +/- kL_recoil.

The ground-state solutions of Raman-coupled spinor systems in general have spinor components with both left- and right-moving momentum peaks. Providing a manual shift on the momentum-space wavefunction components better approximates these solutions, i.e. faster convergence in imaginary time propagation.

Parameters

  • psik (list of NumPy array, optional.) – The momentum-space pseudospinor wavefunction. If psik is not provided, then this function uses the current class attribute self.psik.

  • scale (float, default=1.0) – By default, the function shifts the momentum peaks by a single unit of recoil momenta kL_recoil. scale gives the option of scaling the shift larger or smaller for faster convergence.

  • frac (iterable, default=(0.5, 0.5)) – The fraction of each spinor component’s momentum peak to shift in either direction. frac=(0.5, 0.5) splits into two equal peaks, while (0.0, 1.0) and (1.0, 0.0) move the entire peak one direction or the other.

pspinor.tensor_propagator module

class spinor_gpe.pspinor.tensor_propagator.TensorPropagator(spin, t_step, n_steps, device='cpu', time='imag', is_sampling=False, n_samples=1)

Bases: object

CPU- or GPU-compatible propagator of the GPE, with tensors.

n_steps

The total number of full time steps in propagation.

Type

int

device

The computing device on which propagation is performed

Type

str

paths

See pspinor.Pspinor.

Type

dict

t_step

Duration of the full time step.

Type

float or complex

dt_out

Duration of the outer time sub-step.

Type

float or complex

dt_in

Duration of the inner time sub-step.

Type

float or complex

rand_seed

See pspinor.Pspinor.

Type

int

is_sampling

Option to sample the wavefunction periodically throughout propagation.

Type

bool

atom_num

See pspinor.Pspinor.

Type

float

is_coupling

See pspinor.Pspinor.

Type

bool

g_sc

See pspinor.Pspinor.

Type

dict of Tensor

kin_eng_spin

See pspinor.Pspinor.

Type

list of Tensor

pot_eng_spin

See pspinor.Pspinor.

Type

list of Tensor

psik

See pspinor.Pspinor.

Type

list of Tensor

space
See pspinor.Pspinor. Contains only keys:

{‘dr’, ‘dk’, ‘x_mesh’, ‘y_mesh’, ‘dv_r’, ‘dv_k’}

Type

dict of Tensor

coupling

See pspinor.Pspinor.

Type

Tensor

kL_recoil

See pspinor.Pspinor.

Type

float

expon

The exponential argument on the coupling operator off-diagonals. If the coupling is in a rotated reference frame, then `expon`=0.0.

Type

Tensor

sample_rate

How often wavefunctions are sampled.

Type

int

eng_out

Pre-computed energy evolution operators for the outer time sub-step.

Type

dict of Tensor

eng_in

Pre-computed energy evolution operators for the inner time sub-step.

Type

dict of Tensor

eng_expect(psik)

Compute the energy expectation value of the wavefunction.

To calculate the kinetic energy potion of the energy expectation value, spatial gradients of the phase and square root density must be obtained.

Parameters

psik (list of NumPy array) – The k-space representation of wavefunction to evaluate.

Notes

While spatial gradients of the wavefunction’s phase can be computed with PyTorch tensors, there is currently not an implementation of the 2D phase-unwrapping algorithm. For this this reason, the energy expectation value needs to be computed with NumPy arrays.

full_step()

Full step forward in real or imaginary time.

For accuracy, divide the full propagation step into three single steps using the magic gamma time steps.

prop_loop(n_steps)

Evaluate the propagation steps in a for-loop.

Saves the spin populations at every time step. If wavefunctions are sampled throughout the propagation, they are saved with the associated sampled times in trial_data/psik_sampled%s_`folder_name.npz.

Parameters

n_steps (int) – The number of propagation steps.

Returns

result – Contains the propagation results and analysis methods.

Return type

PropResult

See also

spinor_gpe.prop_results

Propagation results

single_step(t_step, eng)

Single step forward in real or imaginary time with spectral method.

The kinetic, interaction, and coupling time-evolution operators are symmetrically split into two half-single steps around the full-single step potential energy operator.

Parameters

  • t_step (float) – The sub-time step.

  • eng (dict) – The kinetic and potential energy evolution operators corresponding to the given sub-time step.

pspinor.prop_result module

class spinor_gpe.pspinor.prop_result.PropResult(psi_final, psik_final, eng_final, pops, sampled_path=None)

Bases: object

The result of propagation, with plotting and analysis tools.

psi

The final real-space wavefunctions.

Type

list of array

psik

The final momentum-space wavefunctions.

Type

list of array

eng_final

The energy expectation values: [<total>, <kin.>, <pot.>, <int.>].

Type

list

pops

Times and populations at every time step, {‘times’, ‘vals’}.

Type

dict of array

sampled_path

Path to the .npz file where the sampled wavefunctions and times are stored for this result.

Type

str

dens

The final real-space densities.

Type

list of array

densk

The final momentum-space densities.

Type

list of array

phase

The final real-space phases.

Type

list of array

paths

See pspinor.PSpinor.

Type

dict

time_scale

See pspinor.PSpinor.

Type

float

space

See tensor_propagator.TensorPropagator.

Type

dict of array

analyze_vortex()

Compute the total vorticity in each spin component.

calc_separation()

Calculate the phase separation of the two spin components.

make_movie(rscale=1.0, kscale=1.0, cmap='viridis', play=False, zoom=1.0, norm_type='all')

Generate a movie of the wavefunctions’ densities and phases.

Parameters

  • rscale (float, optional) – Real-space length scale. The default of 1.0 corresponds to the naturatl harmonic length scale along the x-axis.

  • kscale (float, optional) – Momentum-space length scale. The default of 1.0 corresponds to the inverse harmonic length scale along the x-axis.

  • cmap (str, optional) – Color map name for the real- and momentum-space density plots.

  • play (bool, default=False) – If True, the movie is opened in the computer’s default media player after it is saved.

  • kzoom (float, optional) – A zoom factor for the k-space density plot.

  • norm_type (str, optional) – {‘all’, ‘half’} Normalizes the colormaps to the full or half sum of the max densites. ‘half’ is useful for visualizing situations where the population is equally divided between the two spins.

plot_eng()

Plot the sampled energy expectation values.

plot_pops(scaled=True, save=True, ext='.pdf')

Plot the spin populations as a function of propagation time.

Parameters

  • scaled (bool, optional) – If scaled is True then the time-axis will be rescaled into proper time units. Otherwise, it’s left in dimensionless time units.

  • save (bool, optional) – Saves the figure as a .pdf file (default). The filename has the format “/data_path/pop_evolution%s-trial_name.pdf”.

  • ext (str, optional) – File extension for the saved plot image.

plot_spins(rscale=1.0, kscale=1.0, cmap='viridis', save=True, ext='.pdf', show=True, zoom=1.0)

Plot the densities (real & k) and phases of spin components.

Parameters

  • rscale (float, default=1.0) – Real-space length scale. The default of 1.0 corresponds to the naturatl harmonic length scale along the x-axis.

  • kscale (float, default=1.0) – Momentum-space length scale. The default of 1.0 corresponds to the inverse harmonic length scale along the x-axis.

  • cmap (str, default=’viridis’) – Matplotlib color map name for the real- and momentum-space density plots.

  • save (bool, default=True) – Saves the figure as a .pdf file (default). The filename has the format “/data_path/spin_dens_phase%s-trial_name.pdf”.

  • ext (str, default=’.pdf’) – File extension for the saved plot image.

  • zoom (float, default=1.0) – A zoom factor for the k-space density plot.

plot_total(rscale=1.0, kscale=1.0, cmap='viridis', save=True, ext='.pdf', show=True, zoom=1.0)

Plot the total real-space density and phase of the wavefunction.

Parameters

  • rscale (float, default=1.0) – Real-space length scale. The default of 1.0 corresponds to the naturatl harmonic length scale along the x-axis.

  • kscale (float, default=1.0) – Momentum-space length scale. The default of 1.0 corresponds to the inverse harmonic length scale along the x-axis.

  • cmap (str, default=’viridis’) – Color map name for the real- and momentum-space density plots.

  • save (bool, default=True) – Saves the figure as a .pdf file (default). The filename has the format “/data_path/pop_evolution%s-trial_name.pdf”.

  • ext (str, default=’.pdf’) – File extension for the saved plot image.

  • show (bool, default=True) – Option to display the generated image.

  • zoom (float, default = 1.0) – A zoom factor for the k-space density plot.

rebin(arr, new_shape=(256, 256))

Rebin a 2D arr to shape new_shape by averaging.

This may be used when generating movies of sampled wavefunctions. By down-sampling the density grids, the movie is generated much faster.

Parameters

  • arr (2D list or NumPy array) – The input 2D array to rebin.

  • new_shape (iterable, default=(256, 256)) – The target rebinned shape.

pspinor.tensor_tools module

tensor_tools.py module.

spinor_gpe.pspinor.tensor_tools.calc_atoms(psi, vol_elem=1.0)

Calculate the total number of atoms.

Parameters

  • psi (list of 2D NumPy array or PyTorch Tensor) – The input spinor wavefunction.

  • vol_elem (float) – 2D volume element of the space.

Returns

atom_num – The total atom number in both spin components.

Return type

float

spinor_gpe.pspinor.tensor_tools.calc_pops(psi, vol_elem=1.0)

Calculate the populations in each spin component.

Parameters

  • psi (list of 2D NumPy array or PyTorch Tensor) – The input spinor wavefunction.

  • vol_elem (float) – 2D volume element of the space.

Returns

pops – The atom number in each spin component.

Return type

list of float

spinor_gpe.pspinor.tensor_tools.conj(psi)

Complex conjugate of a complex tensor.

spinor_gpe.pspinor.tensor_tools.conj_comp(psi_comp)

Complex conjugate of a single wavefunction component.

spinor_gpe.pspinor.tensor_tools.coupling_op(t_step, coupling=None, expon=tensor(0))

Compute the time-evolution operator for the coupling term.

Parameters

  • t_step (float) – Sub-time step.

  • coupling (2D PyTorch complex Tensor, optional.) – The coupling mesh. Default is none if self.is_coupling = False. In this case, the evolution operator returned will the identity matrix.

  • expon (2D PyTorch real Tensor, optional.) – The exponential argument in the bare coupling term. If there is no coupling, then this is 0 by default.

Returns

coupl_ev_op

Return type

2D PyTorch Tensor, optional.

spinor_gpe.pspinor.tensor_tools.density(psi)

Compute the density of a spinor wavefunction.

Parameters

psi (list of 2D NumPy array or PyTorch Tensor) – The input spinor wavefunction.

Returns

dens – The density of each component’s wavefunction.

Return type

NumPy array, PyTorch Tensor, or list thereof

spinor_gpe.pspinor.tensor_tools.evolution_op(t_step, energy)

Compute the unitary time-evolution operator for the given energy.

Parameters

  • energy –

  • time_step –

spinor_gpe.pspinor.tensor_tools.expect_val(psi)

Compute the expectation value of the supplied spatial operator.

spinor_gpe.pspinor.tensor_tools.fft_1d(psi, delta_r=(1, 1), axis=0) → list

Compute the forward 1D FFT of psi along a single axis.

Parameters

  • psi (list of NumPy array or PyTorch Tensor) – The input wavefunction.

  • delta_r (NumPy array, default=(1,1)) – A two-element list of the real-sapce x- and y-mesh spacings, respectively.

  • axis (int, default=0) – The axis along which to transform; note that 0 -> y-axis, and 1 -> x-axis.

Returns

psik_axis – The FFT of psi along axis.

Return type

list of NumPy array or PyTorch Tensor

spinor_gpe.pspinor.tensor_tools.fft_2d(psi, delta_r=(1, 1)) → list

Compute the forward 2D FFT of psi.

Parameters

  • psi (list of NumPy array or PyTorch Tensor) – The input wavefunction.

  • delta_r (NumPy array, default=(1,1)) – A two-element list of the real-space x- and y-mesh spacings, respectively. Typically, use ps.space[‘dr’].

Returns

psik – The k-space FFT of the input wavefunction.

Return type

list of NumPy array or PyTorch Tensor

spinor_gpe.pspinor.tensor_tools.grad(psi, delta_r)

Compute the spatial gradient of a wavefunction list.

Parameters

  • psi –

  • delta_r –

spinor_gpe.pspinor.tensor_tools.grad_comp(psi_comp, delta_r)

Spatial gradient of a single wavefunction component.

Raises

TypeError – If psi_comp is neither an array or a Tensor of the correct shape.

spinor_gpe.pspinor.tensor_tools.grad_sq(psi, delta_r)

Compute the gradient squared of a wavefunction.

spinor_gpe.pspinor.tensor_tools.grad_sq_comp(psi_comp, delta_r)

Take a list of tensors or np arrays; checks type.

spinor_gpe.pspinor.tensor_tools.ifft_1d(psik, delta_r=(1, 1), axis=0) → list

Compute the inverse 1D FFT of psi along a single axis.

Parameters

  • psik (list of NumPy array or PyTorch Tensor) – The input wavefunction.

  • delta_r (NumPy array, default=(1,1)) – A two-element list of the real-sapce x- and y-mesh spacings, respectively.

  • axis (int, default=0) – The axis along which to transform; note that 0 -> x-axis, and 1 -> y-axis.

Returns

psi_axis – The FFT of psi along axis.

Return type

list of NumPy array or PyTorch Tensor

spinor_gpe.pspinor.tensor_tools.ifft_2d(psik, delta_r=(1, 1)) → list

Compute the inverse 2D FFT of psik.

Parameters

  • psik (list of NumPy array or PyTorch Tensor) – The input wavefunction.

  • delta_r (NumPy array, default=(1,1)) – A two-element list of the real-sapce x- and y-mesh spacings, respectively. Typically, use ps.space[‘dr’].

Returns

psi – The real-space FFT of the input wavefunction.

Return type

list of NumPy array or PyTorch Tensor

spinor_gpe.pspinor.tensor_tools.inner_prod()

Calculate the inner product of two wavefunctions.

spinor_gpe.pspinor.tensor_tools.norm(psi, vol_elem, atom_num, pop_frac=None)

Normalize spinor wavefunction to the expected atom numbers and populations.

This function normalizes to the total expected atom number atom_num, and to the expected population fractions pop_frac. Normalization is essential in processes where the total atom number is not conserved, (e.g. imaginary time propagation).

Parameters

  • psi (list of NumPy arrays or PyTorch Tensors.) – The wavefunction to normalize.

  • vol_elem (float) – Volume element for either real- or k-space.

  • atom_num (int) – The total expected atom number.

  • pop_frac (array-like, optional) – The expected population fractions in each spin component.

Returns

  • psi_norm (list of NumPy arrays or PyTorch Tensors.) – The normalized wavefunction.

  • dens_norm (list of NumPy arrays or PyTorch Tensors.) – The densities of the normalized wavefunction’s components

spinor_gpe.pspinor.tensor_tools.norm_sq(psi_comp)

Compute the density (norm-squared) of a single wavefunction component.

Parameters

psi_comp (NumPy array or PyTorch Tensor) – A single wavefunction component.

Returns

psi_sq – The norm-square of the wavefunction.

Return type

NumPy array or PyTorch Tensor

Raises

TypeError – If psi_comp is neither an array or a Tensor of the correct shape.

spinor_gpe.pspinor.tensor_tools.phase(psi, uwrap=False, dens=None)

Compute the phase of a real-space spinor wavefunction.

Parameters

psi (list of 2D NumPy array or PyTorch Tensor) – The input spinor wavefunction.

Returns

phase – The phase of each component’s wavefunction.

Return type

NumPy array, PyTorch Tensor, or list thereof

spinor_gpe.pspinor.tensor_tools.phase_comp(psi_comp, uwrap=False, dens=None)

Compute the phase (angle) of a single complex wavefunction component.

Parameters

psi_comp (NumPy array or PyTorch Tensor) – A single wavefunction component.

Returns

angle – The phase (angle) of the component’s wavefunction.

Return type

NumPy array or PyTorch Tensor

spinor_gpe.pspinor.tensor_tools.prod(factors)

General function for multiplying the elements of a 1D data structure.

Operates similar to the sum function from the standard library.

spinor_gpe.pspinor.tensor_tools.to_cpu(input_tens)

Transfers input_tens from GPU to CPU memory.

Parameters

input_tens (PyTorch Tensor or list of PyTorch Tensor) – Input tensor stored on GPU memory.

Returns

output_tens – Output tensor stored on CPU memory.

Return type

PyTorch Tensor or list of PyTorch Tensor

spinor_gpe.pspinor.tensor_tools.to_gpu(input_tens, dev='cuda')

Transfers input_tens from CPU to GPU memory.

Parameters

  • input_tens (PyTorch Tensor or list of PyTorch Tensor) – Input tensor stored on GPU memory.

  • dev (str, default=’cuda’) – The name of the device on which to store the tensor, e.g. {‘cuda’, ‘cuda:0’}

Returns

output_tens

Return type

PyTorch Tensor or list of PyTorch Tensor

spinor_gpe.pspinor.tensor_tools.to_numpy(input_tens)

Convert from PyTorch Tensor to NumPy arrays.

Accepts a single PyTorch Tensor, or a list of PyTorch Tensor, as in the wavefunction objects.

Parameters

input_tens (PyTorch Tensor, or list of PyTorch Tensor) – Input tensor, or list of tensor, to be converted to array, on CPU memory.

Returns

output_arr – Output array stored on CPU memory.

Return type

NumPy array or list of NumPy array

spinor_gpe.pspinor.tensor_tools.to_tensor(input_arr, dev='cpu', dtype=64)

Convert from NumPy arrays to Tensors.

Accepts a single NumPy array, or a list of NumPy arrays, as in the wavefunction objects.

Parameters

  • input_arr (NumPy array, or list of NumPy array) – Input array, or list of arrays, to be converted to a Tensor, on either CPU or GPU memory.

  • dev (str, default=’cpu’) – The name of the device on which to store the tensor, e.g. {‘cpu’, ‘cuda’, ‘cuda:0’}

  • dtype (int, default=64) –

    Designator for the torch dtype -

    • 32 : torch.float32;

    • 64 : torch.float64;

    • 128 : torch.complex128

Returns

output_tens – Output tensor of dtype stored on dev memory.

Return type

PyTorch Tensor or list of PyTorch Tensor

pspinor.plotting_tools module

plotting_tools.py module.

spinor_gpe.pspinor.plotting_tools.next_available_path(file_name, trial_name, ext='')

Test for the next available path for a given file.

Parameters

  • file_name (str) – The base file path to test.

  • trial_name (str) – The name of the trial to append to the end of the file name.

  • ext (str, default=’’) – File extention.

Returns

test_path – The file path with the next available index.

Return type

str

spinor_gpe.pspinor.plotting_tools.plot_dens(psi, spin=None, cmap='viridis', scale=1.0, extent=None)

Plot the real or k-space density of the wavefunction.

Based on the value of the spin parameter, this function will plot either the up (0), down (1), or both (None) spin components of the spinor wavefunction.

Parameters

  • psi (list of Numpy array, optional.) – The wavefunction to plot. If no psi is supplied, then it uses the object attribute self.psi.

  • spin (int or None, optional) – Which spin to plot. None plots both spins. 0 or 1 plots only the up or down spin, respectively.

  • cmap (str, default=’viridis’) – The matplotlib colormap to use for the plots.

  • scale (float, optional) – A factor to scale the spatial dimensions by, e.g. Thomas-Fermi radius.

  • extent (iterable) – The spatial extent of the wavefunction components, in the format np.array([x_min, x_max, y_min, y_max]). Determines the natural spatial scale of the plot.

spinor_gpe.pspinor.plotting_tools.plot_phase(psi, spin=None, cmap='twilight_shifted', scale=1, extent=None)

Plot the phase of the real wavefunction.

Based on the value of the spin parameter, this function will plot either the up (0), down (1), or both (None) spin components of the spinor wavefunction.

Parameters

  • psi (list of Numpy array, optional.) – The wavefunction to plot.

  • spin (int or None, optional) – Which spin to plot. None plots both spins. 0 or 1 plots only the up or down spin, respectively.

  • cmap (str, optional) – The colormap to use for the plots.

  • scale (float, optional) – A factor to scale the spatial dimensions by, e.g. Thomas-Fermi radius.

  • extent (iterable) – The spatial extent of the wavefunction components, in the format np.array([x_min, x_max, y_min, y_max]). Determines the natural spatial scale of the plot.

spinor_gpe.pspinor.plotting_tools.plot_spins(psi, psik, extents, paths, cmap='viridis', save=True, ext='.pdf', show=True, zoom=1.0)

Plot the densities (real & k) and phases of spin components.

In total, six subplots are generated. Each pair of axes are stored together in a list, which is returned in all_plots.

Parameters

  • psi (list of Numpy array, optional.) – The real-space wavefunction to plot.

  • psik (list of Numpy array, optional.) – The momentum-space wavefunction to plot.

  • extents (dict of iterable) – The dictionary keys are {‘r’, ‘k’}, and each value is a 4-element iterables giving the x- (kx-) and y- (ky-) spatial extents of the plot area, e.g. [x_min, x_max, y_min, y_max]

  • paths (dict of str) – The dictionary keys contain {‘data’, ‘folder’}, and the values are absolute paths to the saved data path and its containing folder.

  • cmap (str, default=’viridis’) – Matplotlib color map name for the real- and momentum-space density plots.

  • save (bool, default=True) – Saves the figure as a .pdf file (default). The filename has the format “/data_path/spin_dens_phase%s-trial_name.pdf”.

  • ext (str, default=’.pdf’) – File extension for the saved density plots.

  • zoom (float, default=1.0) – A zoom factor for the k-space density plot.

Returns

  • fig (plt.Figure) – The matplotlib figure for the plot.

  • all_plots (dict of list) – The keys are {‘r’, ‘ph’, ‘k’}. Each value is a pair of matplotlib.image.AxesImage for both spins.

spinor_gpe.pspinor.plotting_tools.plot_total(psi, psik, extents, paths, cmap='viridis', save=True, ext='.pdf', show=True, zoom=1.0)

Plot the total densities and phase of the wavefunction.

Parameters

  • psi (list of Numpy array, optional.) – The real-space wavefunction to plot.

  • psik (list of Numpy array, optional.) – The momentum-space wavefunction to plot.

  • extents (dict of iterable) – The dictionary keys are {‘r’, ‘k’}, and each value is a 4-element iterables giving the x- (kx-) and y- (ky-) spatial extents of the plot area, e.g. [x_min, x_max, y_min, y_max]

  • paths (dict of str) – The dictionary keys contain {‘data’, ‘folder’}, and the values are absolute paths to the saved data path and its containing folder.

  • cmap (str, default=’viridis’) – Matplotlib color map name for the real- and momentum-space density plots.

  • save (bool, default=True) – Saves the figure as a .pdf file (default). The filename has the format “/data_path/spin_dens_phase%s-trial_name.pdf”.

  • ext (str, default=’.pdf’) – File extension for the saved density plots.

  • zoom (float, default=1.0) – A zoom factor for the k-space density plot.

Returns

  • fig (plt.Figure) – The matplotlib figure for the plot.

  • all_plots (dict of matplotlib.image.AxesImage) – The keys are {‘r’, ‘ph’, ‘k’}. Each value is a separate matplotlib.image.AxesImage.

spinor_gpe.pspinor.plotting_tools.progress_message(frame, n_total)

Display an updating progress message while the animation is saving.

This function produces an output similar to what the tqdm package gives. This one works in situations where tqdm cannot be applied.

Parameters

  • frame (int) – The current frame/index number in the loop.

  • n_total (int) – The total number of frames in the loop.

spinor_gpe.pspinor.plotting_tools.time_remaining(frame, n_total, its)

Calculate completion time in the progress_message function.

Parameters

  • frame (int) – The current frame/index number in the loop.

  • n_total (int) – The total number of frames in the loop.

  • its (float) – The time in seconds between successive iterations.

Module contents

Created on Fri Apr 9 13:35:19 2021

@author: benjamin