"""Interface Definitions.
To help keep the code modular and predictable, we provide several
interfaces. The idea is that generalized code should only these
methods and attributes defined in the interface so that it is easy for
one to extend the code by implementing a different version of the same
interface.
"""
from zope.interface import implementer, Interface, Attribute
import mmfutils.math.bases.interfaces
from pytimeode.interfaces import IStateFlat, IStateWithNormalize, IStateWithBraket
__all__ = [
"IMinimizeState",
"IExperimentMinimal",
"IExperiment",
"IStateDFT",
"IStateMinimizeDFT",
"implementer",
]
[docs]
class IMinimizeState(Interface):
"""Interface provided by state minimizers."""
def __init__(state, real=False, fix_N=True):
"""Constructor.
Arguments
---------
state : IStateMinimizeDFT
Initial state providing the IStateMinimizeDFT interface.
real : bool
If `True`, then the solver will only consider real ansatz,
assuming that the ground state is real, or that any phases are
exactly provided by `P`. Note: if the actual state should also
be real (i.e. if `P` is also real) then this should be
indicated in `state.dtype`.
fix_N : bool
If `True`, then keep `N` fixed while minimizing.
"""
[docs]
def minimize(psi_tol=1e-8, E_tol=1e-12, callback=None, **kw):
"""Return the `state` minimizing the energy.
Arguments
---------
psi_tol : float
Desired relative tolerance for the wavefunction.
E_tol : float
Desired relative tolerance for the energy.
callback : function
Optional callback function called as `callback(state)`
during iterations for debugging.
kw : dict
Additional arguments to pass to the solver.
"""
[docs]
class IExperimentMinimal(Interface):
"""Minimal Interface for Experiments.
The full IExperiment interface can be defined from these minimal
attributes by the `ExperimentBase` class.
"""
[docs]
State = Attribute(
"""IState class used. This class will be used to create
states and copies of states."""
)
[docs]
t_unit = Attribute(
"""Conversion factor between physical times `t` and
dimensionless times `t_`. Internally, the experiment should
use `t_` which should ideally be round numbers (these are used
when saving data to disk for example) while the states use the
physical time `t = t_ * t_unit`. The name `t_` should be used
consistently for this dimensionless time, thus `t__final =
t_final/t_unit` etc."""
)
[docs]
t_name = Attribute("""Name of t_unit for plotting, messages etc.""")
[docs]
image_ts_ = Attribute("Times in units of t_unit used for imaging.")
[docs]
def init():
"""Performs any complex initialization other that assigning
parameters (which is done by the constructor). Subclasses
should not overload the constructor.
"""
[docs]
def get_Vext(state, fiducial=False, expt=False):
"""Return the external potential(s) as relevant for the
experiment.
The potential should return three different types of
potential depending on the arguments:
* If `expt==False`, then return the potential appropriate for
the simulation. This may differ from the actual physical
potential if `cells_x is not None` since then only a portion
at the center of the trap will be simulation. In this case,
the potential should be a periodic extension of the
experimental potential matching the experiment in the
middle, but becoming periodic at the boundaries so as not to
introduce cusps.
The simulation may also have different parameters such as
the trapping frequency along x since one might like to
consider a homogeneous state.
* If `expt==True`, then return the full experimental
potential with the experimental parameters. This potential
should not be a periodic modification nor should it use
modified parameters such as the trapping frequencies. The
idea here is to return the potential that should be used to
defined the initial state by determining the chemical
potential at the center of the system.
This potential may depend on one more flag:
* If `fiducial==True`, then return a potential such that
`V_ext(x_TF) = V_TF` at the edge of the cloud `x_TF` in
the Thomas Fermi approximation. This is generally how the
initial state is specified, however, this potential might
not be the one actually used to initialize the simulation
as described below.
* If `fiducial==False`, then return the actual potential
used to initialize the simulation. This may differ from
the previous case if the initial experimental state is
obtained after adiabatic evolution from
`V_ext(fiducial=True, expt=True)` to
`V_ext(fiducial=False, expt=True)`. Strictly speaking,
our simulation should also perform this adiabatic
evolution, but this can be very costly. Instead, it is
often better to minimize into the initial state directly,
but to do this, one needs to specify and fix the
appropriate particle number or chemical potential from the
fiducial state.
Arguments
---------
state : IState
Current state. Use this to get the time `state.t` and
abscissa `state.basis.xyz`.
fiducial : bool
If `True`, then return the potential that should be used to
define the initial state in terms of the Thomas Fermi
radius of the cloud `x_TF`.
expt : bool
If `True`, then return the proper experimental potential
rather than the potential used in the simulation.
"""
[docs]
def get_state(initialize=True):
"""Quickly return a valid `State` object.
Should not set `state.initializing = False` unless `initialize == True`.
Arguments
---------
initialize : bool
If `False`, then the initial data need not be provided.
(Used by simulations when data is loaded from a file.)
If returning a fully initialized state, then may set
`state.initializing=True`, but users should not generally rely on this flag,
and should still call `pre_evolve_hook()` before evolving.
"""
[docs]
def get_initial_state():
"""Return the valid `t=0` state to initialize the simulations.
The returned state should have `state.initializing == False`.
"""
[docs]
def get_initialized_state(state):
"""Return a valid state initialized from `state`.
This is used in chained simulations where a specified state of one
simulation is used to initialize a state for further use. For example,
for expansion.
"""
[docs]
class IExperiment(IExperimentMinimal):
"""Full interface for Experiment classes.
States should delegate to the experiment allowing the experiment
to control external potentials etc.
Note: All times are expressed in terms of `t_unit` such that `t_ =
t/t_unit`. We try to consistently use the name `t_` for such dimensionless
quantities except in class variables which are assumed to be in the
specified `t_unit`.
Simulations and Imaging
=======================
The idea of an "Experiment" is some sort of simulation run defined by a set
of parameters (attributes of this class) that is evolved through a set of
`image_ts_` under a set of "normal" experimental conditions. These states would be
what is observed in "in situ imaging". Typically, however, from these states, one
evolves for an additional time `t__image` without any interactions or traps to allow
the clouds to "expand", after which an "expansion image" is taken, usually resolving
better details like vortices and domain walls.
The assumption is that one experiment object is shared between all states
belonging to a simulation. For example, this means that
`experiment.t__image` will be the same for each state (along with every
other parameter).
"""
[docs]
image_ts_ = Attribute(
"image_ts_",
"""Tuple of times (in `t_unit`) with experimental images.
These are the times before expansion at which the images are taken. This list
will be used by the Simulation class to ensure that all the required data is
taken to match against the experiments.
During simulations, the state `t_final` will be set to one of these times,
optionally followed by expansion by `t__image` before imaging.
""",
)
[docs]
t__image = Attribute(
"t__image", "Imaging time for expansion imaging. Set to 0 for in situ."
)
####################
# These attributes relate to file storage.
[docs]
max_key_length = Attribute("Limit key lengths (if not directory_per_key)")
[docs]
directory_per_key = Attribute(
"If `True`, use a directory for each key (shorter filenames)"
)
[docs]
dir_name = Attribute("""Name of the directory in which to store checkpoint data""")
def __init__(_local_dict=None, **kw):
"""Constructor.
The constructor takes a dictionary of the local variables
passed to the subclass. These will be assigned to variables
of the same name in `__dict__` and stored to generate the
`dir_name` where data will be stored. Names starting with `_`
will be ignored.
The `kw` argument is provided to allow subclasses to pass in
additional parameters.
Note: We generally recommend that users DO NOT overload the
constructor for several reasons. Although one can use this approach to
define additional parameters, this use is generally discouraged for the
following reasons:
1. These will override any parameters defined at the class level -
EVEN IF OVERRIDDEN IN SUBCLASSES. If you want subclasses to
override these parameters, they MUST do so in their `__init__()`
method.
2. These parameters will ALWAYS be included in the directory name,
even if the default values are used. This behavior can be achived
anyway by simply passing the default value to the default
constructor.
If you proceed with this approach, use the following model::
def __init__(self, a=0.0, **_kw):
super().__init__(locals(), **_kw)
Make sure you pass through `_kw` with an underscore so it does not get
set as an attribute. The following approach of using `locals()` allows
you to forgo repeating the keyword arguments you specify in the signature.
Do not do any initialization here - only set parameters. Initialization
should be performed in the `init()` method which will be called by
`ExperimentBase.__init__()`.
"""
[docs]
def items():
"""Provides support for Archivable."""
# return [(_k, getattr(self, _k)) for _k in self._keys]
[docs]
def copy():
"""Return a copy of the experiment."""
[docs]
class IStateDFT(IStateFlat, 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.
Attributes
----------
experiment : IExperiment
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
:mod:`pytimeode`.
t_final : float, None
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.
"""
[docs]
xyz = Attribute(
"xyz",
"""Physical coordinates: `(x,y,z)`, `(x,y)`, `(x,r)`, etc.
This should be a tuple of abscissa shaped appropriately for
the basis to represent the physical state. This attribute
should be physical - thus, if the state is expressing a
problem in a transforming frame such as during expansion,
translation, or rotation, then these are the physics or "lab"
coordinates, not the computational coordinates.
They should be shaped appropriately that functions like
external trapping potentials can be implemented by using
these. A common example is a harmonic oscillator potential
which might be implemented as::
V = 0.5*m*w**2 * sum(_x**2 for _x in self.get_xyz())
""",
)
[docs]
t__scale = Attribute(
"t_scale",
"""Typical timescale used for the maximum timestep in evolvers.
Often set as `hbar/E_max`, evolvers should converge for `t < 0.1 t_scale`
""",
)
[docs]
t_final = Attribute("t_final", "Final time for normal evolution (not imaging).")
[docs]
experiment = Attribute("experiment", "IExperiment")
[docs]
initializing = Attribute(
"initializing",
"""True if initializing the state (before evolution). This signals that the
potentials should be set for determining the initial state, for example,
subtracting the chemical potential from Vext. This is `True` until
`pre_evolve_hook()` is called.
This is similar to the behavior for `t<0` where one might run cooling for some
time to initialize the state.
""",
)
[docs]
def set_psi(psi):
"""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.
"""
[docs]
def get_psi(psi):
"""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.
"""
[docs]
def apply_V(V, exp=False):
"""Apply V as a potential to the state and return self.
Arguments
---------
V : array-like
Diagonal potential to be applied by multiplication to the
state.
exp : bool
If `True` then `exp(V)` will be applied instead.
"""
[docs]
def apply_laplacian(factor, exp=False):
"""Apply the laplacian multiplied by `factor` to the state.
Arguments
---------
factor : array-like
The result will be multiplied by this factor.
exp : bool
If `True` then `exp(factor*laplacian)(y)` will be computed instead.
"""
[docs]
class IStateMinimizeDFT(IStateDFT, IStateWithNormalize):
"""Extension of IStateDFT that permits minimization.
These extensions are not needed for direct evolution, but helpful
for ground-state preparation.
"""
[docs]
asnumpy = Attribute(
"asnumpy", "Method to convert the output of `ravel() to a numpy array."
)
[docs]
xp = Attribute("xp", "NumPy- or CuPy-like module to work with output of `ravel()`.")
[docs]
metric = Attribute(
"metric",
"""Metric.
For a linear Hamiltonian, the energy should be::
E = sum(y.conj().T*Hy * metric)
""",
)
[docs]
def pre_minimize_hook():
"""Called like `pre_evolve_hook()` but should not set `initializing=False`."""
[docs]
def fill(value):
"""Fill state will constant value."""
[docs]
def normalize(s=None):
"""Normalize the state and return `(s, N)`self.
This is a generalization of the method required by the
IStateWithNormalize interface which returns the scale-factor
`s` effecting the normalization through multiplication::
state *= s
Arguments
---------
s : float or array-like, None
The scale factor to scale the state by such that `y*s` is
normalized. If `None`, then the scale factor should be
computed and returned along with the final normalization.
Returns
-------
s : float or array-like
The scale factor such that `y*s` is normalized.
N : float or array-like
Final normalization. This is used in two cases: 1) if some
of the normalization values are zero `np.any(N == 0)`, then
if any values are zero, then the minimizers expect the
non-zero entries of `state*N` to be zero and will not
optimize over these values. 2) When fixing the constraints,
`N` is used to remove the singular directions from the Hamiltonian.
"""
######################################################################
# Energy-related methods.
#
# These methods are used by the minimizer to find the initial
# state, and for testing convergence. They are not needed for
# direct evolution.
[docs]
def get_energy():
"""Return the energy of the state.
Used by minimizers.
"""
[docs]
def get_Hy(subtract_mu=False):
"""Return `H*psi = dE/dpsi.T.conj()`.
This determines the descent direction to minimize the energy.
Usually it is implemented simply as `1j*hbar*dy_dt` and the
default implementation in `mixins` does this.
Arguments
---------
subtract_mu : bool
If `True`, then the appropriate chemical potentials should
be subtracted from the Hamiltonian so that real-time
evolution would fix the particle number.
"""
######################################################################
# Modified Interfaces
# These come from other packages like mmfutils, but we need slightly modified interfaces
# here, so we define these. Eventually, these changes should be rolled back into the
# original packages.
class IBasis(mmfutils.math.bases.interfaces.IBasis):
pass
class IBasisLz(IBasis, mmfutils.math.bases.interfaces.IBasisLz):
pass
class IBasisWithConvolution(IBasis, mmfutils.math.bases.interfaces.IBasisWithConvolution):
pass
class IBasisCutoff(IBasis, mmfutils.math.bases.interfaces.IBasisCutoff):
pass