Source code for iDEA.observables

import copy
import string
import itertools
from typing import Union
import numpy as np
import iDEA.system
import iDEA.state
import iDEA.methods.non_interacting
import iDEA.methods.interacting


[docs]def observable( s: iDEA.system.System, observable_operator: np.ndarray, state: Union[iDEA.state.SingleBodyState, iDEA.state.ManyBodyState] = None, evolution: Union[ iDEA.state.SingleBodyEvolution, iDEA.state.ManyBodyEvolution ] = None, return_spins: bool = False, ) -> Union[float, np.ndarray]: r""" Compute an observable based on a given operator and state or evolution. | Args: | s: iDEA.system.System, System object. | observable_operator: np.ndarray, Obserbable operator. | state: iDEA.state.SingleBodyState or iDEA.state.ManyBodyState, State. (default = None) | evolution: iDEA.state.SingleBodyEvolution or iDEA.state.ManyBodyEvolution, Evolution. (default = None) | return_spins: bool, True to also return the spin observables: total, up, down. (default = False) | Returns: | observable: float or np.ndarray, Observable. """ if state is not None and type(state) == iDEA.state.ManyBodyState: raise NotImplementedError() if state is not None and type(state) == iDEA.state.SingleBodyState: up_O = 0.0 for i in range(state.up.orbitals.shape[1]): up_O += ( np.vdot( state.up.orbitals[:, i], np.dot(observable_operator, state.up.orbitals[:, i]), ) * state.up.occupations[i] * s.dx ) down_O = 0.0 for i in range(state.down.orbitals.shape[1]): down_O += ( np.vdot( state.down.orbitals[:, i], np.dot(observable_operator, state.down.orbitals[:, i]), ) * state.down.occupations[i] * s.dx ) O = up_O + down_O if return_spins: return O, up_O, down_O else: return O if evolution is not None and type(evolution) == iDEA.state.ManyBodyEvolution: raise NotImplementedError() if evolution is not None and type(evolution) == iDEA.state.SingleBodyEvolution: up_O = np.zeros(shape=evolution.t.shape, dtype=complex) for i, I in enumerate(evolution.up.occupied): for j, ti in enumerate(evolution.t): up_O[j] += ( np.vdot( evolution.up.td_orbitals[j, :, i], np.dot(observable_operator, evolution.up.td_orbitals[j, :, i]), ) * evolution.up.occupations[I] * s.dx ) down_O = np.zeros(shape=evolution.t.shape, dtype=complex) for i, I in enumerate(evolution.down.occupied): for j, ti in enumerate(evolution.t): down_O[j] += ( np.vdot( evolution.down.td_orbitals[j, :, i], np.dot( observable_operator, evolution.down.td_orbitals[j, :, i] ), ) * evolution.down.occupations[I] * s.dx ) O = up_O + down_O if return_spins: return O.real, up_O.real, down_O.real else: return O.real else: raise AttributeError(f"State or Evolution must be provided.")
[docs]def density( s: iDEA.system.System, state: Union[iDEA.state.SingleBodyState, iDEA.state.ManyBodyState] = None, evolution: Union[ iDEA.state.SingleBodyEvolution, iDEA.state.ManyBodyEvolution ] = None, time_indices: np.ndarray = None, return_spins: bool = False, ) -> np.ndarray: r""" Compute the charge density of a non_interacting state. | Args: | s: iDEA.system.System, System object. | state: iDEA.state.SingleBodyState or iDEA.state.ManyBodyState, State. (default = None) | evolution: iDEA.state.SingleBodyEvolution or iDEA.state.ManyBodyEvolution, Evolution. (default = None) | time_indices: np.ndarray, Time indices to compute observable if given evolution. If None will perform for all time indices. (default = None) | return_spins: bool, True to also return the spin densities: total, up, down. (default = False) | Returns: | n: np.ndarray, Charge density, or evolution of charge density. """ if state is not None and type(state) == iDEA.state.ManyBodyState: spin_densities = np.zeros(shape=(s.x.shape[0], 2)) for i in range(s.x.shape[0]): for k in range(2): spin_densities[i, k] = ( np.sum(abs(state.full[i, k, ...]) ** 2) * s.dx ** (s.count - 1) * s.count ) up_n = spin_densities[:, 0] down_n = spin_densities[:, 1] n = up_n + down_n if return_spins: return n, up_n, down_n else: return n if state is not None and type(state) == iDEA.state.SingleBodyState: up_n = np.zeros(shape=s.x.shape[0]) down_n = np.zeros(shape=s.x.shape[0]) for i in range(state.up.orbitals.shape[1]): up_n += abs(state.up.orbitals[:, i]) ** 2 * state.up.occupations[i] for i in range(state.down.orbitals.shape[1]): down_n += abs(state.down.orbitals[:, i]) ** 2 * state.down.occupations[i] n = up_n + down_n if return_spins: return n, up_n, down_n else: return n if evolution is not None and type(evolution) == iDEA.state.ManyBodyEvolution: if time_indices is None: time_indices = np.array(range(evolution.t.shape[0])) spin_densities = np.zeros(shape=(time_indices.shape[0], s.x.shape[0], 2)) for j, ti in enumerate(time_indices): l = string.ascii_lowercase[: s.count] L = string.ascii_uppercase[: s.count] st = ( l + "," + L + "->" + "".join([i for sub in list(zip(l, L)) for i in sub]) ) full = np.einsum(st, evolution.td_space[ti, ...], evolution.spin) L = list( zip(list(range(0, s.count * 2, 2)), list(range(1, s.count * 2, 2))) ) perms = itertools.permutations(list(range(s.count))) full_copy = copy.deepcopy(full) full = np.zeros_like(full_copy) for p in perms: indices = list(itertools.chain(*[L[e] for e in p])) full += iDEA.methods.interacting._permutation_parity(p) * np.moveaxis( full_copy, list(range(s.count * 2)), indices ) full = full / np.sqrt(np.sum(abs(full) ** 2) * s.dx**s.count) for i in range(s.x.shape[0]): for k in range(2): spin_densities[j, i, k] = ( np.sum(abs(full[i, k, ...]) ** 2) * s.dx ** (s.count - 1) * s.count ) up_n = spin_densities[:, :, 0] down_n = spin_densities[:, :, 1] n = up_n + down_n if return_spins: return n, up_n, down_n else: return n if evolution is not None and type(evolution) == iDEA.state.SingleBodyEvolution: if time_indices is None: time_indices = np.array(range(evolution.t.shape[0])) up_n = np.zeros(shape=(time_indices.shape[0], s.x.shape[0])) for i, I in enumerate(evolution.up.occupied): for j, ti in enumerate(time_indices): up_n[j, :] += ( abs(evolution.up.td_orbitals[ti, :, i]) ** 2 * evolution.up.occupations[I] ) down_n = np.zeros(shape=(time_indices.shape[0], s.x.shape[0])) for i, I in enumerate(evolution.down.occupied): for j, ti in enumerate(time_indices): down_n[j, :] += ( abs(evolution.down.td_orbitals[ti, :, i]) ** 2 * evolution.down.occupations[I] ) n = up_n + down_n if return_spins: return n, up_n, down_n else: return n else: raise AttributeError(f"State or Evolution must be provided.")
[docs]def density_matrix( s: iDEA.system.System, state: Union[iDEA.state.SingleBodyState, iDEA.state.ManyBodyState] = None, evolution: Union[ iDEA.state.SingleBodyEvolution, iDEA.state.ManyBodyEvolution ] = None, time_indices: np.ndarray = None, return_spins: bool = False, ) -> np.ndarray: r""" Compute the charge density matrix of a non_interacting state. | Args: | s: iDEA.system.System, System object. | state: iDEA.state.SingleBodyState or iDEA.state.ManyBodyState, State. (default = None) | evolution: iDEA.state.SingleBodyEvolution or iDEA.state.ManyBodyEvolution, Evolution. (default = None) | time_indices: np.ndarray, Time indices to compute observable if given evolution. If None will perform for all time indices. (default = None) | return_spins: bool, True to also return the spin density matrices: total, up, down. (default = False) | Returns: | p: np.ndarray, Charge density matrix, or evolution of charge density matrix. """ if state is not None and type(state) == iDEA.state.ManyBodyState: tosum = list(range(2, s.count * 2)) spin_p = ( np.tensordot(state.full, state.full.conj(), axes=(tosum, tosum)).diagonal( axis1=1, axis2=3 ) * s.dx ** (s.count - 1) * s.count ) up_p = spin_p[:, :, 0] down_p = spin_p[:, :, 1] p = up_p + down_p if return_spins: return p, up_p, down_p else: return p if state is not None and type(state) == iDEA.state.SingleBodyState: up_p = np.zeros(shape=s.x.shape * 2) down_p = np.zeros(shape=s.x.shape * 2) for i in range(state.up.orbitals.shape[1]): up_p += ( np.tensordot( state.up.orbitals[:, i], state.up.orbitals[:, i].conj(), axes=0 ) * state.up.occupations[i] ) for i in range(state.down.orbitals.shape[1]): down_p += ( np.tensordot( state.down.orbitals[:, i], state.down.orbitals[:, i].conj(), axes=0 ) * state.down.occupations[i] ) p = up_p + down_p if return_spins: return p, up_p, down_p else: return p if evolution is not None and type(evolution) == iDEA.state.ManyBodyEvolution: if time_indices is None: time_indices = np.array(range(evolution.t.shape[0])) tosum = list(range(2, s.count * 2)) spin_density_matrices = np.zeros( shape=(time_indices.shape[0], s.x.shape[0], s.x.shape[0], 2), dtype=complex ) for j, ti in enumerate(time_indices): l = string.ascii_lowercase[: s.count] L = string.ascii_uppercase[: s.count] st = ( l + "," + L + "->" + "".join([i for sub in list(zip(l, L)) for i in sub]) ) full = np.einsum(st, evolution.td_space[ti, ...], evolution.spin) L = list( zip(list(range(0, s.count * 2, 2)), list(range(1, s.count * 2, 2))) ) perms = itertools.permutations(list(range(s.count))) full_copy = copy.deepcopy(full) full = np.zeros_like(full_copy) for p in perms: indices = list(itertools.chain(*[L[e] for e in p])) full += iDEA.methods.interacting._permutation_parity(p) * np.moveaxis( full_copy, list(range(s.count * 2)), indices ) full = full / np.sqrt(np.sum(abs(full) ** 2) * s.dx**s.count) spin_density_matrices[j, :, :, :] = ( np.tensordot(full, full.conj(), axes=(tosum, tosum)).diagonal( axis1=1, axis2=3 ) * s.dx ** (s.count - 1) * s.count ) up_p = spin_density_matrices[:, :, :, 0] down_p = spin_density_matrices[:, :, :, 1] p = up_p + down_p if return_spins: return p, up_p, down_p else: return p if evolution is not None and type(evolution) == iDEA.state.SingleBodyEvolution: if time_indices is None: time_indices = np.array(range(evolution.t.shape[0])) up_p = np.zeros( shape=(time_indices.shape[0], s.x.shape[0], s.x.shape[0]), dtype=complex ) for i, I in enumerate(evolution.up.occupied): for j, ti in enumerate(time_indices): up_p[j, :] += ( np.tensordot( evolution.up.td_orbitals[ti, :, i], evolution.up.td_orbitals[ti, :, i].conj(), axes=0, ) * evolution.up.occupations[I] ) down_p = np.zeros( shape=(time_indices.shape[0], s.x.shape[0], s.x.shape[0]), dtype=complex ) for i, I in enumerate(evolution.down.occupied): for j, ti in enumerate(time_indices): down_p[j, :] += ( np.tensordot( evolution.down.td_orbitals[ti, :, i], evolution.down.td_orbitals[ti, :, i].conj(), axes=0, ) * evolution.down.occupations[I] ) p = up_p + down_p if return_spins: return p, up_p, down_p else: return p else: raise AttributeError(f"State or Evolution must be provided.")
[docs]def kinetic_energy( s: iDEA.system.System, state: iDEA.state.SingleBodyState = None, evolution: iDEA.state.SingleBodyEvolution = None, ) -> Union[float, np.ndarray]: r""" Compute the kinetic energy of a non_interacting state. | Args: | s: iDEA.system.System, System object. | state: iDEA.state.SingleBodyState, State. (default = None) | evolution: iDEA.state.SingleBodyEvolution, Evolution. (default = None) | Returns: | E_k: float or np.ndarray, Kinetic energy, or evolution of kinetic energy. """ if state is not None and type(state) == iDEA.state.ManyBodyState: K = iDEA.methods.interacting.kinetic_energy_operator(s) return observable(s, K, state=state) if state is not None and type(state) == iDEA.state.SingleBodyState: K = iDEA.methods.non_interacting.kinetic_energy_operator(s) return observable(s, K, state=state) if evolution is not None and type(evolution) == iDEA.state.ManyBodyEvolution: K = iDEA.methods.interacting.kinetic_energy_operator(s) return observable(s, K, evolution=evolution) if evolution is not None and type(evolution) == iDEA.state.SingleBodyEvolution: K = iDEA.methods.non_interacting.kinetic_energy_operator(s) return observable(s, K, evolution=evolution) else: raise AttributeError(f"State or Evolution must be provided.")
[docs]def external_potential(s: iDEA.system.System) -> np.ndarray: r""" Compute the external potential. | Args: | s: iDEA.system.System, System object. | Returns: | v_ext: np.ndarray, External potential of the system. """ return s.v_ext
[docs]def external_energy( s: iDEA.system.System, n: np.ndarray, v_ext: np.ndarray ) -> Union[float, np.ndarray]: r""" Compute the external energy from a density and external potential. Args: | s: iDEA.system.System, System object. | n: np.ndarray, Charge density of the system. | v_ext: np.ndarray, External potential of the system. Returns: | E_ext: float or np.ndarray, External energy, or evolution of external energy. """ if len(n.shape) == 1: E_ext = np.dot(n, v_ext) * s.dx return E_ext elif len(n.shape) == 2: E_ext = np.zeros(shape=n.shape[0]) for j in range(E_ext.shape[0]): E_ext[j] = np.dot(n[j, :], v_ext[:]) * s.dx return E_ext else: raise AttributeError(f"Expected array of shape 1 or 2, got {n.shape} instead.")
[docs]def hartree_potential(s: iDEA.system.System, n: np.ndarray) -> np.ndarray: r""" Compute the Hartree potential from a density. | Args: | s: iDEA.system.System, System object. | n: np.ndarray, Charge density of the system. | Returns: | v_h: np.ndarray, Hartree potential, or evolution of Hartree potential. """ if len(n.shape) == 1: v_h = np.dot(n, s.v_int) * s.dx return v_h elif len(n.shape) == 2: v_h = np.zeros_like(n) for j in range(v_h.shape[0]): v_h[j, :] = np.dot(n[j, :], s.v_int[:, :]) * s.dx return v_h else: raise AttributeError(f"Expected array of shape 1 or 2, got {n.shape} instead.")
[docs]def hartree_energy( s: iDEA.system.System, n: np.ndarray, v_h: np.ndarray ) -> Union[float, np.ndarray]: r""" Compute the Hartree energy from a density and Hartree potential. | Args: | s: iDEA.system.System, System object. | n: np.ndarray, Charge density of the system. | v_h: np.ndarray, Hartree potential of the system. | Returns: | E_h: float or np.ndarray, Hartree energy, or evolution of Hartree energy. """ if len(n.shape) == 1: E_h = 0.5 * np.dot(n, v_h) * s.dx return E_h elif len(n.shape) == 2: E_h = np.zeros(shape=n.shape[0]) for j in range(E_h.shape[0]): E_h[j] = 0.5 * np.dot(n[j, :], v_h[j, :]) * s.dx return E_h else: raise AttributeError(f"Expected array of shape 1 or 2, got {n.shape} instead.")
[docs]def exchange_potential(s: iDEA.system.System, p: np.ndarray) -> np.ndarray: r""" Compute the exchange potential from a density matrix. | Args: | s: iDEA.system.System, System object. | p: np.ndarray, Density matrix of the system. | Returns: | v_x: np.ndarray, Exchange potential, or evolution of exchange potential. """ if len(p.shape) == 2: v_x = -p * s.v_int return v_x elif len(p.shape) == 3: v_x = np.zeros_like(p) for j in range(v_x.shape[0]): v_x[j, :, :] = -p[j, :, :] * s.v_int[:, :] return v_x else: raise AttributeError(f"Expected array of shape 1 or 2, got {n.shape} instead.")
[docs]def exchange_energy( s: iDEA.system.System, p: np.ndarray, v_x: np.ndarray ) -> Union[float, np.ndarray]: r""" Compute the exchange energy from a density matrix and exchange potential. | Args: | s: iDEA.system.System, System object. | p: np.ndarray, Density matrix of the system. | v_x: np.ndarray, Exchange potential of the system. | Returns: | E_x: float or np.ndarray, Exchange energy, or evolution of exchange energy. """ if len(p.shape) == 2: E_x = 0.5 * np.tensordot(p, v_x, axes=2) * s.dx * s.dx return E_x elif len(p.shape) == 3: E_x = np.zeros(shape=p.shape[0], dtype=complex) for j in range(E_x.shape[0]): E_x[j] = 0.5 * np.tensordot(p[j, :, :].T, v_x[j, :, :], axes=2) * s.dx * s.dx return E_x.real else: raise AttributeError(f"Expected array of shape 1 or 2, got {n.shape} instead.")
[docs]def single_particle_energy( s: iDEA.system.System, state: iDEA.state.SingleBodyState ) -> float: r""" Compute the single particle energy of a single particle state. | Args: | s: iDEA.system.System, System object. | state: iDEA.state.SingleBodyState, State. | Returns: | E: float, Single particle energy. """ return np.sum(state.up.energies[:] * state.up.occupations[:]) + np.sum( state.down.energies[:] * state.down.occupations[:] )
def _placeholder( s: iDEA.system.System, state: Union[iDEA.state.SingleBodyState, iDEA.state.ManyBodyState] = None, evolution: Union[ iDEA.state.SingleBodyEvolution, iDEA.state.ManyBodyEvolution ] = None, return_spins: bool = False, ) -> Union[float, np.ndarray]: r""" Placeholer function. Use this as a template when constructing observable methods. | Args: | s: iDEA.system.System, System object. | state: iDEA.state.SingleBodyState or iDEA.state.ManyBodyState, State. (default = None) | evolution: iDEA.state.SingleBodyEvolution or iDEA.state.ManyBodyEvolution, Evolution. (default = None) | return_spins: bool, True to also return the spin placeholer: total, up, down. (default = False) | Returns: | observable: float or np.ndarray, Placeholer. """ if state is not None and type(state) == iDEA.state.SingleBodyState: raise NotImplementedError() if state is not None and type(state) == iDEA.state.ManyBodyState: raise NotImplementedError() if evolution is not None and type(evolution) == iDEA.state.SingleBodyEvolution: raise NotImplementedError() if evolution is not None and type(evolution) == iDEA.state.ManyBodyEvolution: raise NotImplementedError() else: raise AttributeError(f"State or Evolution must be provided.")