gpe.utils

Contents

gpe.utils#

Various utilities used throughout the project

Exceptions#

PerformanceWarning

Warning for potential performance issues.

Classes#

IStateDFT

Extension to the IState interface for orbital-free DFT equations.

Frames

Represents a series of frames (arrays) on disk for checkpointing and

GPUHelper

Class for helping with GPU and other accelerator support.

AsNumpyMixin

Mixin providing asnumpy method.

StateWithExperimentMixin

Mixing to delegate to self.experiment.

Functions#

use_wisdom(**kw)

Load the fftw_wisdom file (if it exists) and save on exit.

step(t, t1[, alpha])

Smooth step function that goes from 0 at time t=0 to 1 at time

x2_2(x[, order])

Return a periodic approximation for $x^2/2$ with period $2pi$.

good_Ns([Nmax])

Return a list of good N's for the FFT (powers of 2, 3, and 5).

get_good_N(N)

Get the lowest good size N greater than or equal to N for the FFT.

get_smooth_transition(fs, durations, transitions[, alphas])

Return a C(inf) smooth transition as a function of t.

Module Contents#

class IStateDFT(name, bases=(), attrs=None, __doc__=None, __module__=None)[source]#

Bases: pytimeode.interfaces.IStateFlat, pytimeode.interfaces.IStateWithBraket

Extension to the IState interface for orbital-free DFT equations.

The specialization here is to provide a basis capable of applying a Schrodinger-like Hamiltonian to a wavefunction, including a Laplacian. Various extensions can be supported, including using changing basis to account for expansion, rotation, and motion, as well as to multiple components.

This basic interface should be provided by all such classes, and will be relied upon by various utilities.

experiment#

All states should delegate to an Experiment object for control. Experiment objects can be shared, and help provide mechanisms for archiving to disk etc. This is an additional requirement of this project above the requirements of pytimeode.

Type:

IExperiment

t_final#

This is a special parameter used to signify how a state should be run. From t=0 to t==t_final the experiment should be run under “normal” conditions, while for t_final < t the experiment should be run under “imaging” conditions, which might include letting a trapped gas expand, or ramping the coupling constants etc.

In general, the state could have many different stages of evolution, but the need for expansion images is so common that we make a special case here – storing a single “normal” evolution up to different t_final times followed by chains of imaging evolution. In these cases, the values of these parameters (which should be considered volatile) will be manipulated by the Simulation class.

States often also use t_ < 0 to prepare the initial state, which might be prepared (minimized) under different conditions than the run.

Type:

float, None

xyz#
t__scale#
t_final#
experiment#
initializing#
set_psi()[source]#

Set the state from a physical wavefunction.

This should perform any corrections required such as un-applying Bloch twists so that the underlying computational function is periodic, or performing any adjustments to account for changing coordinates.

This is used by the Simulation class to restore data from disk, so this should accept a contiguous numpy array as input.

get_psi()[source]#

Return the physical wavefunction.

This should perform any corrections required such as applying Bloch twists so that the underlying computational function is periodic, or performing any adjustments to account for changing coordinates.

This is used by the Simulation class to archive data to disk, so this should return a contiguous numpy array as input.

apply_V(exp=False)[source]#

Apply V as a potential to the state and return self.

Parameters:
  • V (array-like) – Diagonal potential to be applied by multiplication to the state.

  • exp (bool) – If True then exp(V) will be applied instead.

apply_laplacian(exp=False)[source]#

Apply the laplacian multiplied by factor to the state.

Parameters:
  • factor (array-like) – The result will be multiplied by this factor.

  • exp (bool) – If True then exp(factor*laplacian)(y) will be computed instead.

exception PerformanceWarning[source]#

Bases: Warning

Warning for potential performance issues.

use_wisdom(**kw)[source]#

Load the fftw_wisdom file (if it exists) and save on exit.

step(t, t1, alpha=3.0)[source]#

Smooth step function that goes from 0 at time t=0 to 1 at time t=t1. This step function is \(C_\infty\):

x2_2(x, order=6)[source]#

Return a periodic approximation for \(x^2/2\) with period \(2\pi\).

order == 0: 1-cos(x) order == 2: cos(x)**2/6 - 4*cos(x)/3 + 7/6 …

good_Ns(Nmax=2**15)[source]#

Return a list of good N’s for the FFT (powers of 2, 3, and 5).

get_good_N(N)[source]#

Get the lowest good size N greater than or equal to N for the FFT.

Examples

>>> get_good_N(600)
600
>>> get_good_N(601)
625
get_smooth_transition(fs, durations, transitions, alphas=None)[source]#

Return a C(inf) smooth transition as a function of t.

Smoothly transition from fs[0] to fs[1] to fs[2] etc. and hold these for times ts[0], ts[1], respectively starting from t=0. The transitions take time dts[0] etc.

Parameters:
  • fs ([float]) – List of N function values.

  • durations ([float]) – List of N-1 durations for each function values (the last value will be held indefinitely).

  • transitions ([float]) – List of N-1 transition durations.

  • alphas ([float]) – List of N-1 alpha values for each transition.

class Frames(data_dir, mode='', prefix='frame_', sep='_image_', checkpoint_prefix='check_', checkpoints_to_retain=1, immediate=None, mem_limit_bytes=None, decimal_precision=4)[source]#

Represents a series of frames (arrays) on disk for checkpointing and making movies. Each frame is stored in a file with a key that is usually the time at which the frame is valid.

The main interface is through item access, i.e.:

state.set_data(frames[key])
with frames as frames:
    frames[key] = state.get_data()

Checkpoints should be created explicitly:

with frames as frames:
    frames.checkpoint(key, state.get_data())

These will be treated as regular frames, but will be deleted when a new checkpoint is created.

Context#

Frames instances should generally be used as a context if writing to disk. This will suspend immediate mode until the end of the context, improving performance.

Key Conversions#

To ensure safe comparisons between frame keys, we convert to and from an ikey with the methods key_to_ikey() and ikey_to_key(). These should perform appropriate manipulations like rounding so that equality comparison between keys is meaningful. (Direct comparison of floating point values is dangerous since round-off error could cause a key failure between keys that are practically the same, but obtained by slightly different orders of operations.)

keyThis is what the user supplies as a key. In the default

implementation this is either a floating point number or a tuple of floating point numbers.

ikeyThis is the internal representation of the key, use for

comparison, indexing, etc. In the default implementation, we use the Decimal() class to truncate and provide an exact representation of the floating point number.

Finally, the ikey needs to be converted to a string for use in the filenames. These conversions are done by the ikey_to_str() and str_to_ikey() methods. All four conversion methods should be redefined if a different type of key is used.

The default implementation produces filenames such as:

frame_0.1000_image_0.0500.npy

which would represent a frame evolved after 0.1 time units and then imaged after an additional 0.05 time units of expansion.

data_dir[source]#

The frames will be stored in this directory.

Type:

str

mode[source]#

Read (‘r’), write (‘w’), or in-memory (‘m’). Data not written to disk unless ‘w’ in mode. Setting a frame without ‘w’ issues a warning that the data will not be stored on disk. This can be suppressed by setting the mode to ‘m’ which indicates in-memory mode. Setting mode=’m’ will not access the disk at all (even read access will be disabled).

The default mode=’’ is ‘w’ in a context and ‘r’ outside a context.

Type:

str containing ‘r’, ‘w’, ‘m’

prefix[source]#

Frame filenames start with this.

Type:

str

sep[source]#

When the key is iterable (as defined by self.isiterable()), then keys are joined by this separator. The default is “_image_”.

Type:

str

immediate[source]#

If True, then frames will be saved to disk upon assignment, otherwise they will be saved upon flush() (also called at the end of a context.) The default (None) is False if used in a context, but True otherwise.

Type:

bool, None

checkpoints_to_retain[source]#

When saving a checkpoint, keep this many previous checkpoints and delete the rest.

Type:

int

mem_limit_bytes[source]#

If not None, limit in bytes imposed by self.limit_memory() on the total size the frames may take up in memory. If self.immediate is None (default), limit_memory is called upon exiting a context. If self.immediate, limit_memory is automatically called whenever a frame is set or loaded. Otherwise, limit_memory must be invoked manually.

Type:

None (default) or int

decimal_precision[source]#

This is a specialized argument for the default version of the class. Internally, we use the Decimal class for keys, normalized to this precision. If precision is lost (according to np.allclose, then a ValueError is raised. When loading data from a file, this may be increased to prevent loss of precision.

Type:

int

Examples

Notes

  • If del is called, the underlying file will be removed immediately, even if the immediate flag is False.

  • The list of keys() will only be updated at the start of a context, the first time it is called, or after flush().

  • One can only set an item if it is not already existing. (Call del first if needed.)

data_dir[source]#
prefix = 'frame_'[source]#
checkpoint_prefix = 'check_'[source]#
checkpoints_to_retain = 1[source]#
sep = '_image_'[source]#
property mode[source]#

Return the mode, making context corrections for defaults.

immediate = None[source]#
mem_limit_bytes = None[source]#
decimal_precision = 4[source]#
_data[source]#
_context = False[source]#
__enter__()[source]#

Enter the context.

__exit__(exc_type, exc_val, exc_tb)[source]#
isiterable(key)[source]#

Return True if the key represents a list or tuple of subkeys.

Decimal(key, increase_precision=False)[source]#

Return Decimal(key) rounded to self.decimal_precision.

Parameters:

increase_precision (bool) – If True, then increase as needed self.decimal_precision to ensure accuracy of the keys, otherwise, raise ValueError if precision is lost.

key_to_ikey(key)[source]#

Convert user-supplied key to internal format.

ikey_to_key(ikey)[source]#

Convert internal key format to user format.

ikey_to_str(ikey)[source]#

Convert internal key to string for filename.

str_to_ikey(key_string)[source]#

Convert string to internal key, increasing decimal_precision if required.

__contains__(key)[source]#
ikeys(update=False)[source]#

Return a list of available ikeys.

Parameters:

update (bool) – If True, then reset the list of previously loaded files, and reread data from disk. If False, then only previously detected frames and computed frames will be seen - any additional frames (i.e. saved by another process) will not be detected.

keys(update=False)[source]#

Return a list of available keys.

Parameters:

update (bool) – If True, then reset the list of previously loaded files, and reread data from disk. If False, then only previously detected frames and computed frames will be seen - any additional frames (i.e. saved by another process) will not be detected.

__getitem__(key)[source]#
checkpoint(key, value)[source]#

Save data to file, but as a checkpoint.

_is_checkpoint(ikey)[source]#

Return True if key is saved as a checkpoint.

_convert_checkpoint(ikey)[source]#

Convert a saved checkpoint to a real frame.

Return True if if checkpoint file was converted.

__setitem__(key, value, checkpoint=False)[source]#
__delitem__(key, ikey=None)[source]#
flush()[source]#

Write current data to disk.

get_mem_size()[source]#

Return the size of the frames stored in memory in bytes.

Note that this only calls sys.getsizeof() on each value of self._data; the actual size of a Frames instance will be slightly bigger.

limit_memory(update_files=True)[source]#

Drop the oldest frames from memory to fit under the memory limit.

property _prefixes[source]#

Return a tuple of the possible filename prefixes.

_filename_to_ikey(filename)[source]#

Extract the time key from the filename.

property _files[source]#
_get_data_dir()[source]#

Return data_dir, checking that it exists, is a directory, and making it if needed and ‘w’ in mode.

get_filename(key, checkpoint=False)[source]#

Return the filename associated with key.

_get_filename(ikey, data_dir, checkpoint=False)[source]#

Return the filename associated with key.

_save(ikey, value, data_dir=None, checkpoint=False)[source]#
_remove_checkpoints(checkpoints_to_retain=None)[source]#

Remove old checkpoints in key order.

class GPUHelper[source]#

Class for helping with GPU and other accelerator support.

Examples

>>> class B(GPUHelper):
...     def asnumpy(self, result):
...         return result + ['from B.asnumpy']  # Simulate a custom asnumpy method
>>> class A(GPUHelper):
...     def asnumpy(self, result):
...         return result + ['from A.asnumpy']  # Simulate a custom asnumpy method
...
...     @staticmethod
...     def get_list1_GPU():
...         "Return a list.  Non-GPU method will return an array."
...         return [1, 2]
...
...     def get_list2_GPU(self):
...         "Return a list.  Non-GPU will use self.asnumpy."
...         return [1, 2]
...
...     def get_list3_GPU(self, state):
...         "Return a list.  None-GPU will use state.asnumpy."
...         return [1, 2]
>>> a = A()
>>> a.get_list1_GPU()  # GPU method returns list (representing GPU array)
[1, 2]
>>> a.get_list2_GPU()
[1, 2]
>>> a.get_list3_GPU(state=B())
[1, 2]
>>> a.get_list1()  # Non GPU method returns an array.
array([1, 2])
>>> a.get_list2()  # Non GPU instance method uses A.asnumpy.
[1, 2, 'from A.asnumpy']
>>> a.get_list3(state=B())  # Non GPU instance method uses B.asnumpy.
[1, 2, 'from B.asnumpy']
>>> set(['get_list1', 'get_list2', 'get_list3']).issubset(dir(a))
True
classmethod __init_subclass__()[source]#

Prepare the class by defining all non-GPU/GPU pairs “*” and “*_GPU”.

class AsNumpyMixin[source]#

Bases: gpe.mixins.StateMixin, GPUHelper

Mixin providing asnumpy method.

asnumpy[source]#
get_data_GPU()[source]#

Partner to set_data().

This allows access to the data via self.get_data() which is guaranteed to be a number array.

class StateWithExperimentMixin(experiment, **kw)[source]#

Bases: gpe.mixins.StateMixin, GPUHelper

Mixing to delegate to self.experiment.

experiment[source]#
get_Vext_GPU()[source]#
get_Vint_GPU()[source]#
get_Eint()[source]#
get_n_TF(V_TF)[source]#
get_ns_TF(Vs_TF)[source]#