"""Contains all functionality to define and manage definitions of model systems."""
import warnings
import numpy as np
import iDEA.utilities
import iDEA.interactions
__all__ = ["System", "save_system", "load_system", "systems"]
[docs]class System:
r"""Model system, containing all defining properties."""
def __init__(
self,
x: np.ndarray,
v_ext: np.ndarray,
v_int: np.ndarray,
electrons: str,
stencil: int = 13,
):
r"""
Model system, containing all defining properties.
| Args:
| x: np.ndarray, Grid of x values in 1D space.
| v_ext: np.ndarray, External potential on the grid of x values.
| v_int: np.ndarray, Interaction potential on the grid of x values.
| electrons: string, Electrons contained in the system.
| stencil: int, Stencil to use for derivatives on the grid of x values. (default = 13)
| Raises:
| AssertionError.
"""
self.__x = x
self.__dx = self.x[1] - self.x[0]
self.v_ext = v_ext
self.v_int = v_int
self.__electrons = electrons
self.count = len(electrons)
self.up_count = electrons.count("u")
self.down_count = electrons.count("d")
self.stencil = stencil
self.check()
[docs] def check(self):
r"""Performs checks on system properties. Raises AssertionError if any check fails."""
assert (
type(self.x) == np.ndarray
), f"x grid is not of type np.ndarray, got {type(self.x)} instead."
assert (
type(self.v_ext) == np.ndarray
), f"v_ext is not of type np.ndarray, got {type(self.v_ext)} instead."
assert (
type(self.v_int) == np.ndarray
), f"v_int is not of type np.ndarray, got {type(self.v_int)} instead."
assert (
type(self.count) == int
), f"count is not of type int, got {type(self.NE)} instead."
assert (
len(self.x.shape) == 1
), f"x grid is not a 1D array, got {len(self.x.shape)}D array instead."
assert (
len(self.v_ext.shape) == 1
), f"v_ext is not a 1D array, got {len(self.v_ext.shape)}D array instead."
assert (
len(self.v_int.shape) == 2
), f"v_int is not a 2D array, got {len(self.v_int.shape)}D array instead."
assert (
self.x.shape == self.v_ext.shape
), f"x grid and v_ext arrays are not the same shape, got x.shape = {self.x.shape} and v_ext.shape = {self.v_ext.shape} instead."
assert (
self.x.shape[0] == self.v_int.shape[0]
and self.x.shape[0] == self.v_int.shape[1]
), "v_int is not of the correct shape, got shape {self.v_int.shape} instead."
assert self.count >= 0, f"count is not positive."
assert set(self.electrons).issubset(
set(["u", "d"])
), f"Electrons must have only up or down spin, e.g 'uudd'. Got {self.electrons} instead"
assert (
self.count == self.up_count + self.down_count
), f"Electrons must obay up_count + down_count = count."
assert self.stencil in [
3,
5,
7,
9,
11,
13,
], f"stencil must be one of [3,5,7,9,11,13], got {self.stencil} instead."
@property
def x(self):
return self.__x
@x.setter
def x(self, value):
self.__x = value
self.__dx = self.__x[1] - self.__x[0]
warnings.warn(
"x grid has been changed: dx has been recomputed, please update v_ext and v_int on this grid."
)
@x.deleter
def x(self):
del self.__x
@property
def dx(self):
return self.__dx
@dx.setter
def dx(self, value):
raise AttributeError(
"cannot set dx directly: set the x grid and dx will be updated automatically."
)
@dx.deleter
def dx(self):
del self.__dx
@property
def electrons(self):
return self.__electrons
@electrons.setter
def electrons(self, value):
self.__electrons = value
self.count = len(value)
self.up_count = value.count("u")
self.down_count = value.count("d")
@electrons.deleter
def electrons(self):
del self.__electrons
def __str__(self):
return f"iDEA.system.System: x = np.array([{self.x[0]:.3f},...,{self.x[-1]:.3f}]), dx = {self.dx:.4f}..., v_ext = np.array([{self.v_ext[0]:.3f},...,{self.v_ext[-1]:.3f}]), electrons = {self.electrons}"
[docs]def save_system(s: System, file_name: str) -> None:
r"""
Save a system to an system file.
| Args:
| system: iDEA.system.System, System object to save.
| file_name: str, file name.
"""
pickle.dump(s, open(file_name, "wb"))
[docs]def load_system(file_name: str) -> System:
r"""
Load a system from an system file.
| Args:
| file_name: str, file name.
| Returns
| system: iDEA.system.System, Loaded System object.
"""
return pickle.load(open(file_name, "rb"))
# Define some default built in systems.
__x1 = np.linspace(-10, 10, 300)
systems = iDEA.utilities.Container()
systems.qho = System(
__x1,
0.5 * (0.25**2) * (__x1**2),
iDEA.interactions.softened_interaction(__x1),
"uu",
)
__x2 = np.linspace(-20, 20, 300)
systems.atom = System(
__x2, -2.0 / (abs(__x2) + 1.0), iDEA.interactions.softened_interaction(__x2), "ud"
)