class ExtendedHuckelHamiltonian(Operator):
S: np.array = attrib()
molecule: Molecule = attrib()
hermitian = True
def __attrs_post_init__(self):
invalid = [a.element for a in self.molecule if a.element.voie is None]
if invalid:
raise ValueError(
f"Could not make EHT Hamiltonian from elements {invalid}"
f". They are not configured with VOIE."
)
def matrix(self, basis=None) -> np.array:
"""Create the Hamiltonian under the Extended Hueckel Approximation."""
h = np.zeros(self.S.shape)
n = self.S.shape[0]
for i in range(n):
h[i, i] = -self.molecule.atoms[i].element.voie
for i in range(n):
for j in range(n):
if i != j:
h[i, j] = 1.75 * (h[i, i] + h[j, j]) * self.S[i, j] / 2.0
return h
class Parabola:
a: float = attrib(validator=[
not_nan, not_inf,
Range(-conf.large_number, conf.large_number)
])
b: float = attrib(validator=[
not_nan, not_inf,
Range(-conf.large_number, conf.large_number)
])
c: float = attrib(validator=[
not_nan, not_inf,
Range(-conf.large_number, conf.large_number)
])
@document_me
def __call__(self, x):
return self.a * x**2 + self.b * x + self.c
@property
def has_vertex(self) -> bool:
return not isclose(self.a, 0.0)
@property
def vertex(self) -> float:
if not self.has_vertex:
raise ValueError(f"{self} is a line and has no vertex.")
return -self.b / (2.0 * self.a)
@staticmethod
def from_points(point1: Tuple[float, float], point2: Tuple[float, float],
point3: Tuple[float, float]) -> Parabola:
"""Create a Parabola passing through 3 points.
Parameters
-----------
point1
x, y point
point2
x, y point
point3
x, y point
Returns
--------
Parabola
Parabola with coefficients fit to the points.
"""
x1, y1 = point1
x2, y2 = point2
x3, y3 = point3
# The polynomial coefficients are the solution to the
# linear equation Ax = b where
# x = [a, b, c]
A = np.array([[x1**2, x1, 1], [x2**2, x2, 1], [x3**2, x3, 1]])
b = np.array([y1, y2, y3])
coeffs = np.linalg.solve(A, b)
return Parabola(a=coeffs[0], b=coeffs[1], c=coeffs[2])
class Potential(Operator):
potential: Callable[[float], float] = attrib()
overlap: Overlap = attrib(default=Overlap())
hermitian = True
@document_me
def __call__(self, first, second) -> float:
def v(x):
return self.potential(x) * second(x)
return self.overlap(first, v)
class Overlap(Operator):
lower_limit: float = attrib(default=-np.inf)
upper_limit: float = attrib(default=np.inf)
hermitian: bool = True
"""Compute the overlap integral in 1 dimension over the specified range
Parameters
-----------
first
The first function
second
The second function
lower_limit
The lower limit of integration
upper_limit
The upper limit of integration
Returns
--------
overlap
The value of the overlap integral.
Examples
---------
>>> from quantized import overlap
>>> from numpy import sin, pi
>>>
>>> def f1(x):
... return sin(x)
...
>>> def f2(x):
... return sin(2*x)
...
>>> overlap(f1, f2, lower_limit=0, upper_limit=2*pi)
6.869054119163646e-17
>>> overlap(f1, f1, lower_limit=0, upper_limit=2*pi)
3.141592653589793
"""
@document_me
def __call__(self, first: Callable, second: Callable) -> float:
try:
val = first.__overlap__(second, self.lower_limit, self.upper_limit)
return val
except (AttributeError, NotImplementedError):
pass
def integrand(x):
return first(x) * second(x)
return integrate.quad(integrand, a=self.lower_limit, b=self.upper_limit)[0]
class Hamiltonian(Operator):
potential: Callable[[float], float] = attrib()
kinetic: Kinetic = attrib(default=Kinetic())
hermitian = True
def __attrs_post_init__(self):
self.Potential = Potential(self.potential)
@document_me
def __call__(self, first, second) -> float:
return self.Potential(first, second) + self.kinetic(first, second)
class Harmonic:
center: float = attrib(
validator=[Range(-conf.large_number, conf.large_number)])
mass: float = attrib(default=1.0,
validator=Range(conf.small_number, conf.large_number))
omega: float = attrib(default=1.0,
validator=Range(conf.small_number,
conf.large_number))
@document_me
def __call__(self, x: float) -> float:
"""Return the value of the harmonic potential at coordinate x"""
return 0.5 * self.mass * (self.omega**2) * ((x - self.center)**2)
class Element:
"""Class to hold element information"""
z: int = attrib(repr=False, validator=[Range(1, 118)], desc="The atomic number of the element")
name: str = attrib(desc="The symbol of the element, e.g. H, He, O")
voie: Optional[float] = attrib(
default=None,
repr=False,
desc="The valence orbital ionization energy in eV",
validator=optional(positive),
)
alpha: Optional[float] = attrib(
default=None, repr=False, desc="The orbital decay coefficient", validator=optional(positive)
)
class Kinetic(Operator):
overlap: Overlap = attrib(default=Overlap())
hermitian = True
@document_me
def __call__(self, first, second) -> float:
return Overlap()(first, second.__kinetic__())
class OccupancyProbabilites:
s: Tuple[TimeEvolvingObservable] = attrib(converter=tuple)
@property
def initial(self) -> np.array:
return self.__call__(0.0)
@document_me
def __call__(self, t: float) -> np.array:
return np.array([s(t) for s in self.s])
@document_me
def __getitem__(self, item) -> Callable:
return self.s[item]
@document_me
def __iter__(self):
return iter(self.s)
@staticmethod
def from_1d_state(state: TimeEvolvingState,
borders: Tuple[float, ...]) -> "OccupancyProbabilites":
bounds = pairwise(borders)
overlaps = [Overlap(a, b) for a, b in bounds]
return OccupancyProbabilites(
tuple(state.observable(ovlp, hermitian=True) for ovlp in overlaps))
class EigenBasis:
"""A class for representing an eigenbasis for a hamiltonian"""
# TODO: validate shapes, properties
states: Tuple[Callable, ...] = attrib(converter=tuple,
desc="A set of eigen states")
energies: Tuple[float,
...] = attrib(converter=tuple,
desc="The energies of the eigen states")
ao_S: np.ndarray = attrib(desc="The overlap matrix in the original basis")
eigvecs: np.ndarray = attrib(
desc="The eigenvectors of the hamiltonian. Each column is a vector.")
ao_basis: List[Callable] = attrib(desc="The original basis")
@document_me
def __len__(self):
"""The size of the eigenbasis"""
return len(self.states)
@staticmethod
def from_basis(basis: List[Callable], H: np.ndarray,
S: np.ndarray) -> EigenBasis:
"""Create an eigenbasis from another basis, given a hamiltonian and overlap matrix
**H (np.ndarray)**
The hamiltonian matrix in the basis
**S (np.ndarray)**
The overlap matrix in the basis
"""
# Sort vals/vecs
eigvals, eigvecs = eigh(H, S)
idx = np.argsort(eigvals)
eigvals = eigvals[idx]
eigvecs = eigvecs[:, idx]
states = [LinearComb(c, basis) for c in eigvecs.T]
return EigenBasis(states=tuple(states),
energies=tuple(eigvals),
eigvecs=eigvecs,
ao_S=S,
ao_basis=basis)
def transformed(self, matrix: np.ndarray) -> np.ndarray:
"""Given a matrix in the original basis, return the matrix in the Eigen basis."""
return inv(self.eigvecs) @ inv(self.ao_S) @ matrix @ self.eigvecs
class Config:
"""Class to hold configuration for quantized"""
harmonic_oscillator_max_n: int = attrib(
default=50,
converter=int,
desc="Limit of n for the harmonic oscillator")
small_number: float = attrib(
default=1e-8,
converter=float,
desc="A value to be used to ensure quantities are nonzero")
large_number: float = attrib(
default=1000.0,
converter=float,
desc=
"A cutoff value that's used to guard against errors, where a number is not normally expected to be large.",
)
float_tol: float = attrib(
default=1e-6,
converter=float,
desc="Two floats are considered equal if they are within this value.",
)
enable_progressbar: bool = attrib(
default=False,
converter=to_bool,
desc="If True, certain methods will print a progress bar to the screen",
)
cache_dir: Path = attrib(
default=".quantized/cache",
converter=Path,
desc="The directory where cached objects are stored",
)
joblib_verbosity: int = attrib(default=0,
converter=int,
desc="Verbosity level for joblib's cache")
class TimeEvolvingObservable:
time_evolving_state: TimeEvolvingState = attrib()
operator: Operator = attrib()
hermitian: bool = attrib()
def __attrs_post_init__(self):
e = self.time_evolving_state.eigen_basis.energies
N = len(self.time_evolving_state.eigen_basis)
c = self.time_evolving_state.expansion_coeffs
# TODO: this is awkward
ao_matrix = self.operator.matrix(
self.time_evolving_state.eigen_basis.ao_basis)
eigen_basis_matrix = self.time_evolving_state.eigen_basis.transformed(
ao_matrix)
# The following code is equivalent to the definition of f below.
# f is written that way to optimize for speed.
# On my machine, it was roughly 7x faster
#
# def f(t: float):
# return sum(
# c[i] * c[j] * np.exp(1j * (e[j] - e[i]) * t) * eigen_basis_matrix[i, j]
# for i in range(N)
# for j in range(N)
# )
# Precalculate reused terms
# P is a real matrix, but we cast it to complex here so that it doesn't
# get casted inside the function during the multiplication.
self.P = np.zeros_like(eigen_basis_matrix, dtype=np.complex128)
self.W = np.zeros_like(eigen_basis_matrix, dtype=np.complex128)
for i in range(N):
for j in range(N):
self.P[i, j] = c[i] * c[j] * eigen_basis_matrix[i, j]
self.W[i, j] = np.exp(1j * (e[j] - e[i]))
@document_me
def __call__(self, t: Union[float, np.array]) -> Union[float, np.array]:
if isinstance(t, Real):
return np.abs(np.sum(self.P * (self.W**t)))
elif isinstance(t, np.ndarray):
return np.array(
[np.abs(np.sum(self.P * (self.W**t_))) for t_ in t])
class TimeEvolvingState:
initial_state: Callable = attrib()
eigen_basis: EigenBasis = attrib(repr=False)
def __attrs_post_init__(self):
self.expansion_coeffs = get_expansion_coeffs(self.initial_state,
self.eigen_basis.states)
@document_me
def __call__(self, x: Union[float, np.ndarray], t: float) -> float:
return sum([
c * np.exp(-1j * e * t) * state(x)
for state, e, c in zip(self.eigen_basis.states, self.eigen_basis.
energies, self.expansion_coeffs)
])
def observable(self, operator: Callable,
hermitian: bool) -> "TimeEvolvingObservable":
return TimeEvolvingObservable(self, operator, hermitian=hermitian)
class Range:
"""Attrs Validator to ensure values between a min and a max value.
"""
min = attrib(desc="The minimum allowed value")
max = attrib(desc="The maximum allowed value")
@min.validator
def check(self, attribute, value):
if not self.min <= self.max:
raise ValidationError(
f"Min must be below max got min={self.min}, max={self.max}")
def __call__(self, instance: Any, attribute: Attribute, value: Any) -> Any:
if not self.min <= value <= self.max:
raise ValidationError(
f"Validation error for {instance}. "
f"{attribute.name} out of range, must be "
f"between {self.min} and {self.max}, got {value}")
class HarmonicOscillator:
"""A 1D quantum harmonic oscillator wave function.
```
>>> ho = HarmonicOscillator(n=2, center=0.5)
>>> ho
HarmonicOscillator(n=2, center=0.5, omega=1, mass=1.0)
>>> ho(0.5)
-0.5311259660135984
>>> ho(1000)
0.0
>>> ho(-1000)
0.0
```
"""
n: int = attrib(validator=[Range(0, conf.harmonic_oscillator_max_n)],
desc="The quantum number")
center: float = attrib(
validator=[Range(-conf.large_number, conf.large_number)],
desc="The center of the function")
mass: float = attrib(
default=1.0,
validator=Range(conf.small_number, conf.large_number),
desc="Mass of the particle",
)
omega: float = attrib(
default=1.0,
validator=Range(conf.small_number, conf.large_number),
desc="Angular frequency of the oscillator",
)
@staticmethod
def from_parabola(p: Parabola,
n: int,
mass: float = 1.0) -> HarmonicOscillator:
"""Create a harmonic oscillator, who's potential is defined by the given parabola"""
# a = m/2 * w**2
# 2a / m = w**2
# sqrt(2a/m) = w
w = np.sqrt(2 * p.a / mass)
return HarmonicOscillator(n=n, center=p.vertex, omega=w, mass=mass)
@staticmethod
def from_potential_points(
point1: Tuple[float, float],
point2: Tuple[float, float],
point3: Tuple[float, float],
n: int,
mass: float = 1.0,
) -> HarmonicOscillator:
"""Create a harmonic oscillator wave function from 3 samples of the potential.
The three points are fit to a parabola (harmonic potential), then the parameters
for the harmonic oscillator are determined, and the corresponding wave function
generated and returned.
**point1**
A sample point of the potential
**point2**
A sample point of the potential
**point3**
A sample point of the potential
**n**
The quantum number of the resulting wave function
**mass**
The mass of the particle
**Examples**
```python
ho = HarmonicOscillator.from_potential_points(
... point1=(0.5, 1),
... point2=(2.0, 0.5),
... point3=(3.0, 1.5),
... n=0
... )
ho
HarmonicOscillator(n=0, center=1.5624999999999998, omega=1.0327955589886444, mass=1.0)
```
"""
p = Parabola.from_points(point1, point2, point3)
return HarmonicOscillator.from_parabola(p, n=n, mass=mass)
@document_me
def __kinetic__(self) -> Callable:
"""Return kinetic energy operator applied on this."""
# K = \frac{p^2}{2m}
# p = i * \sqrt{m w hbar/2}(a+ - a)
#
# k = -1/2 * (m w hbar / 2)[(a+ - a)^2]
# [] = (a+^2 + a+a + aa+ - a^2)
#
# [] = sqrt(n+1)sqrt(n+2)| n + 2 > always
# sqrt(n+1)sqrt(n+1)| n > always
# sqrt(n) sqrt(n) | n > if n == 0 0
# sqrt(n) sqrt(n-1)| n - 2 > if n <= 1 0
#
# k = - (m w hbar) / 4 * []
def k(x):
first = np.sqrt(self.n + 1) * np.sqrt(self.n + 2) * attr.evolve(
self, n=self.n + 2)(x)
if self.n == 0:
# Lowering operator on n=0
# means the second term is zero
inner = (self.n + 1) * self(x)
else:
inner = (2 * self.n + 1) * self(x)
if self.n <= 1:
# Lowering operator twice on n=0 or n=1 will be 0
# if x is numpy array, we want to return a numpy
# array of zeros, so multiply x by zeros will work on
# numpy array or float.
last = 0.0 * x
else:
last = np.sqrt(self.n) * np.sqrt(self.n - 1) * attr.evolve(
self, n=self.n - 2)(x)
return -self.mass * self.omega * (first - inner + last) / 4.0
return k
@document_me
def __overlap__(self, other, lower_limit: float,
upper_limit: float) -> float:
"""Determine the overlap with some other function
This specializes a generic overlap integral, and short circuits integral
calculations if the integral is analytically known.
"""
if lower_limit != -np.inf or upper_limit != np.inf:
raise NotImplementedError
if not isinstance(other, HarmonicOscillator):
raise NotImplementedError
if (isclose(self.center, other.center)
and isclose(self.mass, other.mass)
and isclose(self.omega, other.omega)):
return 1.0 if self.n == other.n else 0.0
raise NotImplementedError
@property
def energy(self):
"""The energy of harmonic oscillator"""
return (self.n + 0.5) * self.omega
@property
def potential(self) -> Harmonic:
"""The potential for this oscillator"""
return Harmonic(center=self.center, mass=self.mass, omega=self.omega)
@property
def N(self):
"""The normalization constant"""
return (1.0 / math.sqrt(2**self.n * math.factorial(self.n)) *
((self.mass * self.omega) / math.pi)**0.25)
@property
def _hermite(self):
return special.hermite(self.n)
@document_me
def __call__(self, x: Union[float,
np.ndarray]) -> Union[float, np.ndarray]:
"""Return"""
y = (np.sqrt(self.mass * self.omega)) * (x - self.center)
return self.N * np.exp(-(y**2) / 2.0) * self._hermite(y)
class Molecule:
atoms: List[Atom] = attrib()
def __len__(self):
return len(self.atoms)
@property
def coords(self):
return np.asarray([a.coords for a in self.atoms])
@property
def R(self):
"""Calculate the distances for each atom-atom pair."""
return pairwise_array_from_func(self.atoms, Atom.distance)
@property
def mass(self):
return sum(a.mass for a in self.atoms)
@property
def center_of_mass(self):
"""Determine the center of mass of the molecule."""
mx = sum(a.x * a.mass for a in self.atoms)
my = sum(a.y * a.mass for a in self.atoms)
mz = sum(a.z * a.mass for a in self.atoms)
return mx / self.mass, my / self.mass, mz / self.mass
def translated(self,
x: float = 0.0,
y: float = 0.0,
z: float = 0.0) -> "Molecule":
f = partial(Atom.translated, x=x, y=y, z=z)
return self.map(f)
@staticmethod
def from_xyz(xyz: str) -> "Molecule":
"""Create a molecule from an xyz-file formatted string"""
text = (line.split() for line in xyz.splitlines()[2:] if line.strip())
atoms = [Atom(element=e, x=x, y=y, z=z) for e, x, y, z in text]
return Molecule(atoms)
def com_as_origin(self) -> "Molecule":
x, y, z = self.center_of_mass
return self.translated(x=-x, y=-y, z=-z)
def with_atom_aligned_to(self, atom: Atom, x: float, y: float,
z: float) -> "Molecule":
r = atom.rotation_matrix_to(x=x, y=y, z=z)
return self.rotated(r)
def rotated(self, r: np.array) -> "Molecule":
if not r.shape == (3, 3):
raise ValueError(f"Rotation matrix has invalid shape: {r.shape}")
rotate = partial(Atom.rotated, r=r)
return self.map(rotate)
def map(self, f: Callable) -> "Molecule":
return attr.evolve(self, atoms=[f(a) for a in self.atoms])
def scaled(self, factor: float) -> "Molecule":
return self.map(partial(Atom.scaled, factor=factor))
def flipped_x(self) -> "Molecule":
return self.map(Atom.flipped_x)
def sorted(self, atomic_key: Callable) -> "Molecule":
return attr.evolve(self, atoms=sorted(self.atoms, key=atomic_key))
def rotated_about_x(self, angle: float) -> "Molecule":
f = partial(Atom.rotated_about_x, angle=angle)
return self.map(f)
def rotated_about_y(self, angle: float) -> "Molecule":
f = partial(Atom.rotated_about_y, angle=angle)
return self.map(f)
def rotated_about_z(self, angle: float) -> "Molecule":
f = partial(Atom.rotated_about_z, angle=angle)
return self.map(f)
def __iter__(self):
return iter(self.atoms)
def __eq__(self, other):
return all(a1 == a2 for a1, a2 in zip(self, other))
class HoppingMatrix:
occ_probs: OccupancyProbabilites = attrib()
N: int = 3
def at_time(self, t: float, delta_t: float) -> np.array:
d_occ_probs = self.occ_probs(t) - self.occ_probs(t - delta_t)
increasing = d_occ_probs > 0
two_states_increasing = np.count_nonzero(increasing) == 2
A = np.zeros((self.N, self.N))
if two_states_increasing:
# If the state is decreasing in probability, the diagonal is
# determined
for j in range(3):
if not increasing[j]:
A[j, j] = self.occ_probs[j](t) / self.occ_probs[j](t -
delta_t)
for k in range(3):
if k != j:
A[k, j] = d_occ_probs[k] / self.occ_probs[j](t)
# If the state is increasing, diagonal is 1. Off diagonal
# columns = 0
elif increasing[j] > 0:
A[j, j] = 1
for k in range(3):
if k != j:
A[k, j] = 0
else:
for j in range(3):
# If the state is decreasing in probability, the diagonal
# is determined
if not increasing[j]:
A[j, j] = self.occ_probs[j](t) / self.occ_probs[j](t -
delta_t)
for k in range(3):
if k != j:
A[j, k] = 0
# If the state is increasing, diagonal is 1. Off diagonal
# column element is 1 - diagonal
for j in range(3):
if increasing[j]:
A[j, j] = 1
for k in range(3):
if k != j:
A[j, k] = 1 - A[k, k]
return A
@document_me
def __call__(self, t: Union[float, np.ndarray], delta_t: float):
if isinstance(t, Iterable):
result = [self.at_time(t_, delta_t) for t_ in t]
if len(result) == 0:
raise ValueError(
"Call to hopping matrix had no elements in the iterable.")
return np.stack(result)
else:
return self.at_time(t, delta_t)
class TripleWell:
"""Construct a well potential from the points.
```
x x
x x
x x
x x
x x
x barrier12 x
x x x barrier23 x
x x x x x x
center1 x x x x
x x center3
center2
```
"""
center1: Tuple[float, float] = attrib()
barrier12: Tuple[float, float] = attrib()
center2: Tuple[float, float] = attrib()
barrier23: Tuple[float, float] = attrib()
center3: Tuple[float, float] = attrib()
well1: Parabola = attrib()
well2: Parabola = attrib()
well3: Parabola = attrib()
def __attrs_post_init__(self):
# Wells are in expected left to right ordering.
if (not self.center1[0] < self.barrier12[0] < self.center2[0] <
self.barrier23[0] < self.center3[0]):
raise ValueError("Points are not in ascending x-value order.")
# Wells are below the barriers
assert self.center1[1] < self.barrier12[1]
assert self.center2[1] < self.barrier12[1]
assert self.center2[1] < self.barrier23[1]
assert self.center3[1] < self.barrier23[1]
def _call_numpy(self, x: np.ndarray) -> np.ndarray:
y = np.zeros_like(x)
mask1 = np.where(x < self.barrier12[0])
y[mask1] = self.well1(x[mask1])
mask2 = np.where((self.barrier12[0] <= x) & (x <= self.barrier23[0]))
y[mask2] = self.well2(x[mask2])
mask3 = np.where(x > self.barrier23[0])
y[mask3] = self.well3(x[mask3])
return y
def _call_scalar(self, x: float) -> float:
if x < self.barrier12[0]:
return self.well1(x)
elif self.barrier12[0] <= x <= self.barrier23[0]:
return self.well2(x)
elif x > self.barrier23[0]:
return self.well3(x)
else:
raise ValueError(
f"Value {x} is not valid for {self}. "
"Probably the well/barrier parameters are invalid")
@document_me
def __call__(self, x: Union[np.ndarray,
float]) -> Union[np.ndarray, float]:
if isinstance(x, np.ndarray):
return self._call_numpy(x)
else:
return self._call_scalar(x)
@staticmethod
def from_params(
well1_depth: float,
well1_halfwidth: float,
bridge_length: float,
bridge_depth: float,
well3_halfwidth: float,
well3_depth: float,
):
center1 = (0, 0)
barrier1 = (center1[0] + well1_halfwidth, center1[1] + well1_depth)
center2 = (barrier1[0] + 0.5 * bridge_length,
barrier1[1] - bridge_depth)
barrier2 = (barrier1[0] + bridge_length, barrier1[1])
center3 = (barrier2[0] + well3_halfwidth, barrier2[1] - well3_depth)
def fit_well(center, barrier):
center_x, center_y = center
barrier_x, barrier_y = barrier
# Third point is reflecting barrier about center
x = -barrier_x + 2 * center_x
y = barrier_y
return Parabola.from_points(center, barrier, (x, y))
well1 = fit_well(center1, barrier1)
well2 = Parabola.from_points(barrier1, center2, barrier2)
well3 = fit_well(center3, barrier2)
return TripleWell(center1, barrier1, center2, barrier2, center3, well1,
well2, well3)
class Pnot:
acceptor: int = attrib()
hopping_matrix: HoppingMatrix = attrib(hash=False)
times: List[float] = attrib()
values: List[float] = attrib()
delta_t: float = attrib()
"""Determine P_not, the probability than acceptor state has never been occupied.
Parameters
-----------
acceptor
The index of the acceptor well
hopping_matrix
A (N, 3, 3) array containing the hopping matrix over time
p0
An (3,) array containing probabilities of occupying each region initially
"""
@property
def min_time(self) -> float:
return self.times[0]
@property
def max_time(self) -> float:
return self.times[-1]
@property
def tau90(self) -> float:
return self.time_when_equal_to(0.1)
@property
def initial_value(self) -> float:
return self(self.min_time)
@property
def final_value(self) -> float:
return self(self.max_time)
@lru_cache()
def time_when_equal_to(self, x: float) -> float:
def f(t):
return self(t) - x
if not self.final_value < x < self.initial_value:
raise ValueError(
f"Can't solve for Pnot = {x}. It is out the calculated range: {self.initial_value}, {self.final_value}"
)
result = root_scalar(f, bracket=[self.min_time, self.max_time])
return result.root
@lru_cache()
def _interpolate(self):
return interp1d(self.times, self.values, bounds_error=False)
@document_me
def __call__(self, t: float) -> float:
return self._interpolate()(t)
@property
def acceptor_occ_prob(self) -> Callable:
return self.hopping_matrix.occ_probs[self.acceptor]
@staticmethod
def gen(delta_t: float, hopping_matrix: HoppingMatrix,
acceptor: int) -> Generator[Tuple, None, None]:
times = (delta_t * i for i in count())
matrices = (hopping_matrix(t, delta_t) for t in times)
A_tilde = (np.delete(np.delete(a, acceptor, axis=0), acceptor, axis=1)
for a in matrices)
p0_tilde = np.delete(hopping_matrix.occ_probs.initial, acceptor)
# Taking [a, nb, c, d, ...] -> [(b @ a), (c @ b @ a), (d @ c @ b @ a), ...]
a_prod = accumulate(A_tilde, lambda a, b: b @ a, initial=p0_tilde)
pnot = map(np.sum, a_prod)
return ((delta_t * i, p) for i, p in enumerate(pnot))
@staticmethod
def gen_until_time(t: float, delta_t: float, hopping_matrix: HoppingMatrix,
acceptor: int) -> Iterator[Tuple[float, float]]:
gen = Pnot.gen(delta_t=delta_t,
hopping_matrix=hopping_matrix,
acceptor=acceptor)
while_lt = takewhile(lambda x: x[0] < t, gen)
extra_terms = islice(gen, 100)
return chain(while_lt, extra_terms)
@staticmethod
def gen_until_prob(p: float, delta_t: float, hopping_matrix: HoppingMatrix,
acceptor: int) -> Iterator[Tuple[float, float]]:
gen = Pnot.gen(delta_t=delta_t,
hopping_matrix=hopping_matrix,
acceptor=acceptor)
while_gt = takewhile(lambda x: x[1] > p, gen)
extra_terms = islice(gen, 100)
return chain(while_gt, extra_terms)
@staticmethod
def converged_with_timestep(
a: HoppingMatrix,
acceptor: int,
until_equals: float,
tolerance: float,
max_dt: float = 1.0,
min_dt: float = 1e-4,
) -> "Pnot":
if not max_dt > min_dt:
raise ValueError(
f"Min dt must be below max dt min_dt={min_dt}, max_dt={max_dt}"
)
dts = [max_dt]
last_dt = dts[-1]
while last_dt > min_dt:
dts.append(last_dt / 1.61)
last_dt = dts[-1]
last_t = -(10.0**6)
for dt in dts:
times, p_not_list = zip(*Pnot.gen_until_prob(
until_equals, acceptor=acceptor, delta_t=dt, hopping_matrix=a))
p_not = Pnot(acceptor=acceptor,
times=times,
values=p_not_list,
delta_t=dt,
hopping_matrix=a)
t = p_not.time_when_equal_to(until_equals)
logger.info(f"For delta_t = {dt:.5f}, Tau = {t:.5f}")
if isclose(t, last_t, abs_tol=tolerance):
break
last_t = t
else:
raise ValueError(
f"Time for Pnot to reach {until_equals} did not converge with timestep to tolerance={tolerance}"
)
logger.info(
f"Yes!! Time for pnot to reach ({until_equals:6f}) successfully converged to {t:6f}"
f" with difference of {t - last_t:3f}"
f", which is within the given tolerance, {tolerance}")
return p_not
class Atom:
"""Atom class containing coordinates, basis and mass.
The atom will generally not be used in isolation, but will
likely be part of a molecule. This is expected to be used
as a structured container for atomic information, but
it does contain logic to transform atoms in space.
"""
element: Element = attrib(converter=element_from_string,
desc="The element")
x: float = attrib(converter=float, desc="The x coordinate of the atom")
y: float = attrib(converter=float, desc="The y coordinate of the atom")
z: float = attrib(converter=float, desc="The z coordinate of the atom")
@property
def mass(self):
"""The mass of the atom"""
return self.element.z
@property
def coords(self) -> np.array:
"""Three dimensional array of coordinates, [x, y, z]"""
return np.array([self.x, self.y, self.z])
def with_coords(self, x: float, y: float, z: float) -> "Atom":
"""Return an equivalent atom at these coordinates"""
return attr.evolve(self, x=x, y=y, z=z)
def translated(self, x=0.0, y=0.0, z=0.0) -> Atom:
"""Return an equivalent atom translated in the direction given"""
return attr.evolve(self, x=self.x + x, y=self.y + y, z=self.z + z)
def scaled(self, factor: float) -> Atom:
"""Return an equivalent atom with all coordinates scaled by some factor"""
return attr.evolve(self,
x=factor * self.x,
y=factor * self.y,
z=factor * self.z)
def rotated(self, r: np.ndarray) -> Atom:
"""Return an equivalent atom rotated by the given rotation matrix
The matrix must have shape (3, 3)
"""
if not r.shape == (3, 3):
raise ValueError(f"Rotation matrix R must be 3x3, got {r.shape}")
x, y, z = r @ np.array([self.x, self.y, self.z])
return attr.evolve(self, x=x, y=y, z=z)
def flipped_x(self) -> Atom:
"""Return an equivalent atom, but the x coordinate is the opposite"""
return attr.evolve(self, x=-self.x)
def distance(self, other: Atom) -> float:
"""Determine the distance between this atom and another atom."""
return sqrt((self.x - other.x)**2 + (self.y - other.y)**2 +
(self.z - other.z)**2)
@property
def normalized_coords(self) -> np.array:
"""Return a unit vector pointing towards the atom"""
return self.coords / np.linalg.norm(self.coords)
def rotation_matrix_to(self, x: float, y: float, z: float) -> np.ndarray:
"""Get a matrix that would rotate this atom to align with the given coordinates
https://math.stackexchange.com/questions/180418/calculate-rotation-matrix-to-align-vector-a-to-vector-b-in-3d
"""
a = self.normalized_coords
b = np.array([x, y, z])
b = b / np.linalg.norm(b)
v = np.cross(a, b)
s = np.linalg.norm(v)
c = np.dot(a, b)
I = np.eye(*a.shape) # noqa: E741
v_x = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
return I + v_x + v_x @ v_x * (1 - c) / s**2
def angle_to_xy_plane(self) -> float:
"""Angle, in radians, between this atom's coordinate vector and the xy plane"""
return angle(self.coords, [self.x, self.y, 0])
def angle_to_xz_plane(self) -> float:
"""Angle, in radians, between this atom's coordinate vector and the xz plane"""
return angle(self.coords, [self.x, 0, self.z])
def angle_to_yz_plane(self) -> float:
"""Angle, in radians, between this atom's coordinate vector and the yz plane"""
return angle(self.coords, [0, self.y, self.z])
def rotated_about_x(self, angle: float) -> "Atom":
"""Return an equivalent atom, rotated by `angle` radians about the x axis"""
r = np.array([[1, 0, 0], [0, cos(angle), -sin(angle)],
[0, sin(angle), cos(angle)]])
return self.with_coords(*r @ self.coords)
def rotated_about_z(self, angle: float) -> "Atom":
"""Return an equivalent atom, rotated by `angle` radians about the z axis"""
r = np.array([[cos(angle), -sin(angle), 0],
[sin(angle), cos(angle), 0], [0, 0, 1]])
return self.with_coords(*r @ self.coords)
def rotated_about_y(self, angle: float) -> "Atom":
"""Return an equivalent atom, rotated by `angle` radians about the y axis"""
r = np.array([[cos(angle), 0, -sin(angle)], [0, 1, 0],
[sin(angle), 0, cos(angle)]])
return self.with_coords(*r @ self.coords)
@document_me
def __eq__(self, other: "Atom") -> bool:
"""Return True if the other atom is the same element, and very close"""
return (self.element is other.element
and isclose(self.x, other.x, abs_tol=conf.small_number)
and isclose(self.y, other.y, abs_tol=conf.small_number)
and isclose(self.z, other.z, abs_tol=conf.small_number))
