"""Defines the structures to describe the system states"""
import copy
import pickle
from abc import ABC as Interface
import numpy as np
import iDEA.utilities
__all__ = [
"State",
"States",
"ManyBodyState",
"ManyBodyStates",
"SingleBodyState",
"Evolution",
"ManyBodyEvolution",
"SingleBodyEvolution",
]
[docs]class State(Interface):
"""Interface class representing a static state."""
[docs]class States(Interface):
"""Interface class representing a group of static states."""
[docs]class Evolution(Interface):
"""Interface class representing a time-dependent evolution of a state."""
[docs]class ManyBodyState(State):
"""State of interacting particles."""
def __init__(
self,
space: np.ndarray = None,
spin: np.ndarray = None,
full=None,
energy=None,
):
r"""
State of particles in a many-body state.
This is described by a spatial part
.. math:: \psi(x_1,x_2,\dots,x_N)
on the spatial grid, and a spin
part on the spin grid
.. math:: \chi(\sigma_1,\sigma_2,\dots,\sigma_N).
These are NOT necessarily antisymmetric states,
they can be combined using the antisymmetrisation operaration to produce the full
wavefunction
.. math:: \Psi(x_1,\sigma_1,x_2,\sigma_2,\dots,x_N,\sigma_N).
| Args:
| space: np.ndarray, Spatial part of the wavefunction on the spatial grid \psi(x_1,x_2,\dots,x_N). (default = None)
| spin: np.ndarray, Spin part of the wavefunction on the spin grid \chi(\sigma_1,\sigma_2,\dots,\sigma_N). (default = None)
| full: np.ndarray, Total antisymmetrised wavefunction \Psi(x_1,\sigma_1,x_2,\sigma_2,\dots,x_N,\sigma_N). (default = None)
| energy: float, Total energy of the state.
"""
if space is None:
self.space = iDEA.utilities.ArrayPlaceholder()
else:
self.space = space
if spin is None:
self.spin = iDEA.utilities.ArrayPlaceholder()
else:
self.spin = spin
if full is None:
self.full = iDEA.utilities.ArrayPlaceholder()
else:
self.full = full
if energy is None:
self.energy = 0.0
else:
self.energy = energy
[docs]class ManyBodyStates(States):
"""A collection of many-body states over a range of energy levels from a single system of interacting particles."""
def __init__(self, spaces: np.ndarray = None, spins: np.ndarray = None, fulls=None, energies=None):
r"""
Collection of many-body states indexed by energy level n.
Each state is described by a spatial part
.. math:: \psi(x_1,x_2,\dots,x_N,n)
on the spatial grid, and a spin part on the spin grid
.. math:: \chi(\sigma_1,\sigma_2,\dots,\sigma_N,n).
These are NOT necessarily antisymmetric states;
they can be combined using the antisymmetrisation operation to produce the full
wavefunction
.. math:: \Psi(x_1,\sigma_1,x_2,\sigma_2,\dots,x_N,\sigma_N,n).
| Args:
| spaces: np.ndarray, Spatial part of each wavefunction on the spatial grid \psi(x_1,x_2,\dots,x_N,n). (default = None)
| spins: np.ndarray, Spin part of each wavefunction on the spin grid \chi(\sigma_1,\sigma_2,\dots,\sigma_N,n). (default = None)
| fulls: np.ndarray, Total antisymmetrised wavefunction of each state \Psi(x_1,\sigma_1,x_2,\sigma_2,\dots,x_N,\sigma_N,n). (default = None)
| energies: np.ndarray, Total energy of each state, indexed as energies[n]. (default = None)
"""
if spaces is None:
self.spaces = iDEA.utilities.ArrayPlaceholder()
else:
self.spaces = spaces
if spins is None:
self.spins = iDEA.utilities.ArrayPlaceholder()
else:
self.spins = spins
if fulls is None:
self.fulls = iDEA.utilities.ArrayPlaceholder()
else:
self.fulls = fulls
if energies is None:
self.energies = iDEA.utilities.ArrayPlaceholder()
else:
self.energies = energies
[docs]class SingleBodyState(State):
r"""
State of particles in a single-body state.
This is described by three arrays for each spin channel:
| up.energies: np.ndarray, Array of single-body energies, indexed as energies[orbital_number].
| up.orbitals: np.ndarray, Array of single-body orbitals, indexed as orbitals[space,orbital_number].
| up.occupations: np.ndarray, Array of single-body occupations, indexed as occupations[orbital_number].
| up.occupied: np.ndarray, Indices of up.occupations that are non-zero, to indicate occupied orbitals.
| down.energies: np.ndarray, Array of single-body energies, indexed as energies[orbital_number].
| down.orbitals: np.ndarray, Array of single-body orbitals, indexed as orbitals[space,orbital_number].
| down.occupations: np.ndarray, Array of single-body occupations, indexed as occupations[orbital_number].
| down.occupied: np.ndarray, Indices of down.occupations that are non-zero, to indicate occupied orbitals.
"""
def __init__(self):
self.up = iDEA.utilities.Container()
self.down = iDEA.utilities.Container()
self.up.energies = iDEA.utilities.ArrayPlaceholder()
self.up.orbitals = iDEA.utilities.ArrayPlaceholder()
self.up.occupations = iDEA.utilities.ArrayPlaceholder()
self.up.occupied = iDEA.utilities.ArrayPlaceholder()
self.down.energies = iDEA.utilities.ArrayPlaceholder()
self.down.orbitals = iDEA.utilities.ArrayPlaceholder()
self.down.occupations = iDEA.utilities.ArrayPlaceholder()
self.down.occupied = iDEA.utilities.ArrayPlaceholder()
[docs]class ManyBodyEvolution(Evolution):
r"""
Time-dependent evolution of particles in a many-body state.
In addition to the arrays defined within the initial ManyBodyState, this state is described by three additional arrays:
| td_space: np.ndarray, Spatial part of the wavefunction on the spatial grid \psi(t,x_1,x_2,\dots,x_N).
| v_ptrb: np.ndarray, Perturbation potential that this time-dependence was driven by. indexed as v_ptrb[space] if static, and v_ptrb[time,space] if dynamic.
| t: np.ndarray, Time grid used during evolution.
"""
def __init__(self, initial_state: ManyBodyState):
self.space = copy.deepcopy(initial_state.space)
self.spin = copy.deepcopy(initial_state.spin)
self.full = copy.deepcopy(initial_state.full)
self.td_space = iDEA.utilities.ArrayPlaceholder()
self.v_ptrb = iDEA.utilities.ArrayPlaceholder()
self.t = iDEA.utilities.ArrayPlaceholder()
[docs]class SingleBodyEvolution(Evolution):
r"""
Time-dependent evolution of particles in a single-body state.
In addition to the arrays defined within the initial SingleBodyState, this state is described by four additional arrays:
| up.td_orbitals: np.ndarray, Array of single-body time-dependend orbitals, indexed as orbitals[time,space,orbital_number].
| down.td_orbital: np.ndarray, Array of single-body time-dependend orbitals, indexed as orbitals[time,space,orbital_number].
| v_ptrb: np.ndarray, Perturbation potential that this time-dependence was driven by. indexed as v_ptrb[space] if static, and v_ptrb[time,space] if dynamic.
| t: np.ndarray, Time grid used during evolution.
In this case, only the occupied time-dependent orbitals are stored, as only these are propigated.
"""
def __init__(self, initial_state: SingleBodyState):
self.up = copy.deepcopy(initial_state.up)
self.down = copy.deepcopy(initial_state.down)
self.up.td_orbitals = iDEA.utilities.ArrayPlaceholder()
self.down.td_orbitals = iDEA.utilities.ArrayPlaceholder()
self.v_ptrb = iDEA.utilities.ArrayPlaceholder()
self.t = iDEA.utilities.ArrayPlaceholder()
def save_many_body_state(state: ManyBodyState, file_name: str) -> None:
r"""
Save a many body state to a system file.
| Args:
| state: iDEA.state.ManyBodyState, State object to save.
| file_name: str, file name.
"""
with open(file_name, "wb") as f:
pickle.dump(state, f)
def save_single_body_state(state: SingleBodyState, file_name: str) -> None:
r"""
Save a single body state to a system file.
| Args:
| state: iDEA.state.SingleBodyState, State object to save.
| file_name: str, file name.
"""
with open(file_name, "wb") as f:
pickle.dump(state, f)
def load_many_body_state(file_name: str) -> ManyBodyState:
r"""
Load a many body state from an system file.
| Args:
| file_name: str, file name.
| Returns
| system: iDEA.state.ManyBodyState, Loaded State object.
"""
return pickle.load(open(file_name, "rb"))
def load_single_body_state(file_name: str) -> SingleBodyState:
r"""
Load a single body state from an system file.
| Args:
| file_name: str, file name.
| Returns
| system: iDEA.state.SingleBodyState, Loaded State object.
"""
return pickle.load(open(file_name, "rb"))