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))