"""Contains all interacting functionality and solvers."""
import copy
import functools
import itertools
import string
from typing import Union
import numpy as np
import scipy.sparse as sps
import scipy.sparse.linalg as spsla
from tqdm import tqdm
import iDEA.methods.non_interacting
import iDEA.state
import iDEA.system
name = "interacting"
[docs]def kinetic_energy_operator(s: iDEA.system.System, GPU: bool = False):
r"""
Compute many-particle kinetic energy operator as a matrix.
This is built using a given number of finite differences to represent the second derivative.
The number of differences taken is defined in s.stencil.
| Args:
| s: iDEA.system.System, System object.
| GPU: bool, Compute on GPU using cupy. If false will use scipy on CPU. (default = False)
| Returns:
| K: sparse matrix, Kintetic energy operator.
"""
if GPU:
import cupyx.scipy.sparse as csps
sp = csps
fmt = "csr"
else:
sp = sps
fmt = "dia"
k = sps.csr_matrix(iDEA.methods.non_interacting.kinetic_energy_operator(s))
k = sp.csr_matrix(k) if GPU else sps.dia_matrix(k)
I = sp.identity(s.x.shape[0], format=fmt)
partial_operators = lambda A, B, k, n: (A if i + k == n - 1 else B for i in range(n))
fold_partial_operators = lambda f, po: functools.reduce(lambda acc, val: f(val, acc, format=fmt), po)
generate_terms = lambda f, A, B, n: (fold_partial_operators(f, partial_operators(A, B, k, n)) for k in range(n))
K = functools.reduce(lambda a, b: a + b, generate_terms(sp.kron, k, I, s.count))
return K
[docs]def external_potential_operator(s: iDEA.system.System, GPU: bool = False):
r"""
Compute many-particle external potential energy operator as a matrix.
| Args:
| s: iDEA.system.System, System object.
| GPU: bool, Compute on GPU using cupy. If false will use scipy on CPU. (default = False)
| Returns:
| Vext: sparse matrix, External potential operator.
"""
if GPU:
import cupyx.scipy.sparse as csps
sp = csps
fmt = "csr"
else:
sp = sps
fmt = "dia"
vext = sps.csr_matrix(iDEA.methods.non_interacting.external_potential_operator(s))
vext = sp.csr_matrix(vext) if GPU else sps.dia_matrix(vext)
I = sp.identity(s.x.shape[0], format=fmt)
partial_operators = lambda A, B, k, n: (A if i + k == n - 1 else B for i in range(n))
fold_partial_operators = lambda f, po: functools.reduce(lambda acc, val: f(val, acc, format=fmt), po)
generate_terms = lambda f, A, B, n: (fold_partial_operators(f, partial_operators(A, B, k, n)) for k in range(n))
Vext = functools.reduce(lambda a, b: a + b, generate_terms(sp.kron, vext, I, s.count))
return Vext
[docs]def hamiltonian(s: iDEA.system.System, GPU: bool = False):
r"""
Compute the many-body Hamiltonian.
| Args:
| s: iDEA.system.System, System object.
| GPU: bool, Compute on GPU using cupy. If false will use scipy on CPU. (default = False)
| Returns:
| H: sparse matrix, Hamiltonian.
"""
if GPU:
import cupy as cp
import cupyx.scipy.sparse as csps
sp = csps
xp = cp
fmt = "csr"
else:
sp = sps
xp = np
fmt = "dia"
# Construct the non-interacting part of the many-body Hamiltonian
h = sps.csr_matrix(iDEA.methods.non_interacting.hamiltonian(s)[0])
h = sp.csr_matrix(h) if GPU else sps.dia_matrix(h)
I = sp.identity(s.x.shape[0], format=fmt)
partial_operators = lambda A, B, k, n: (A if i + k == n - 1 else B for i in range(n))
fold_partial_operators = lambda f, po: functools.reduce(lambda acc, val: f(val, acc, format=fmt), po)
generate_terms = lambda f, A, B, n: (fold_partial_operators(f, partial_operators(A, B, k, n)) for k in range(n))
H0 = functools.reduce(lambda a, b: a + b, generate_terms(sp.kron, h, I, s.count))
# Add the interaction part of the many-body Hamiltonian
symbols = string.ascii_lowercase + string.ascii_uppercase
if s.count > 1:
indices = ",".join(["".join(c) for c in itertools.combinations(symbols[: s.count], 2)])
v_int = xp.asarray(s.v_int) if GPU else s.v_int
U = xp.log(xp.einsum(indices + "->" + symbols[: s.count], *(xp.exp(v_int),) * int(s.count * (s.count - 1) / 2)))
U = sp.diags(U.reshape(H0.shape[0]), format=fmt)
else:
U = 0.0
# Construct the total many-body Hamiltonian
H = H0 + U
return H
[docs]def total_energy(s: iDEA.system.System, state: iDEA.state.ManyBodyState) -> float:
r"""
Compute the total energy of an interacting state.
| Args:
| s: iDEA.system.System, System object.
| state: iDEA.state.ManyBodyState, State.
| Returns:
| E: float, Total energy.
"""
return state.energy
def _permutation_parity(p):
r"""
Compute the permulation paritiy of a given permutation.
| Args:
| p: tuple, Permutation.
| Returns:
| parity: float, Permutation parity.
"""
p = list(p)
parity = 1
for i in range(0, len(p) - 1):
if p[i] != i:
parity *= -1
mn = min(range(i, len(p)), key=p.__getitem__)
p[i], p[mn] = p[mn], p[i]
return parity
[docs]def antisymmetrize(s, spaces, spins, energies):
r"""
Antisymmetrize the solution to the Schrodinger equation.
| Args:
| s: iDEA.system.System, System object.
| spaces: np.ndarray, Spatial parts of the wavefunction.
| spins: np.ndarray, Spin parts of the wavefunction.
| energies: np.ndarray, Energies.
| Returns:
| fulls: np.ndarray, Full anantisymmetrized wavefunction.
| spaces: np.ndarray, Spatial parts of the wavefunction.
| spins: np.ndarray, Spin parts of the wavefunction.
| energies: np.ndarray, Energies.
"""
# Perform antisymmetrization.
l = string.ascii_lowercase[: s.count]
L = string.ascii_uppercase[: s.count]
st = l + "Y," + L + "Y->" + "".join([i for sub in list(zip(l, L)) for i in sub]) + "Y"
fulls = np.einsum(st, spaces, spins)
L = list(zip(list(range(0, s.count * 2, 2)), list(range(1, s.count * 2, 2))))
perms = itertools.permutations(list(range(s.count)))
fulls_copy = copy.deepcopy(fulls)
fulls = np.zeros_like(fulls)
for p in perms:
indices = list(itertools.chain(*[L[e] for e in p]))
fulls += _permutation_parity(p) * np.moveaxis(fulls_copy, list(range(s.count * 2)), indices)
# Filter out zeros.
allowed_fulls = []
allowed_energies = []
allowed_spaces = []
allowed_spins = []
for n in range(fulls.shape[-1]):
if np.allclose(fulls[..., n], np.zeros(fulls.shape[:-1])):
pass
else:
allowed_fulls.append(fulls[..., n])
allowed_energies.append(energies[n])
allowed_spaces.append(spaces[..., n])
allowed_spins.append(spins[..., n])
fulls = np.moveaxis(np.array(allowed_fulls), 0, -1)
spaces = np.moveaxis(np.array(allowed_spaces), 0, -1)
spins = np.moveaxis(np.array(allowed_spins), 0, -1)
energies = np.array(allowed_energies)
# Normalise.
for k in range(fulls.shape[-1]):
fulls[..., k] = fulls[..., k] / np.sqrt(np.sum(abs(fulls[..., k]) ** 2) * s.dx**s.count)
# Filter out duplicates.
allowed_fulls = []
allowed_energies = []
for n in range(fulls.shape[-1] - 1):
if np.allclose(abs(fulls[..., n]), abs(fulls[..., n + 1])):
pass
else:
allowed_fulls.append(fulls[..., n])
allowed_energies.append(energies[n])
allowed_fulls.append(fulls[..., -1])
allowed_energies.append(energies[-1])
fulls = np.moveaxis(np.array(allowed_fulls), 0, -1)
spaces = spaces[..., : fulls.shape[-1]]
spins = spins[..., : fulls.shape[-1]]
energies = np.array(allowed_energies)
return fulls, spaces, spins, energies
def _estimate_level(s: iDEA.system.System, k: int) -> int:
r"""
Estimate the solution to the Schrodinger equation needed to eachive given antisymetric energy state.
| Args:
| s: iDEA.system.System, System object.
| k: int, Target energy state.
| Returns:
| level: int, Extimate of level of excitement.
"""
return (abs(s.up_count - s.down_count) + 1) ** 2 * s.count * (k + 1)
[docs]def solve(
s: iDEA.system.System, H: np.ndarray = None, k: int = 0, level=None, GPU=False
) -> Union[iDEA.state.ManyBodyState, iDEA.state.ManyBodyStates]:
r"""
Solves the interacting Schrodinger equation of the given system.
| Args:
| s: iDEA.system.System, System object.
| H: np.ndarray, Hamiltonian [If None this will be computed from s]. (default = None)
| k: int, Energy state to solve for. If -1 will return all states. (default = 0, the ground-state)
| level: int. Max level of excitation to use when solving the Schrodinger equation.
| GPU: bool, Solve on GPU using cupy. If false will use scipy on CPU.
| Returns:
| state: iDEA.state.ManyBodyState or iDEA.state.ManyBodyStates, Solved state, or collection of all solved states if k=-1.
"""
# Construct the Hamiltonian.
if H is None:
H = hamiltonian(s, GPU=GPU)
# Estimate the level of excitation.
if level is None:
if k == -1:
raise ValueError("level must be provided explicitly when k=-1 (all states).")
level = _estimate_level(s, k)
# Solve the many-body Schrodinger equation.
if GPU:
import cupy as cp
import cupyx.scipy.sparse as csps
import cupyx.scipy.sparse.linalg as cspsla
name = cp.cuda.runtime.getDeviceProperties(cp.cuda.Device().id)["name"].decode()
print(f"iDEA.methods.interacting.solve: solving eigenproblem on GPU: {name}...")
energies, spaces = cspsla.eigsh(csps.csr_matrix(H), k=level, which="SA")
energies = energies.get()
spaces = spaces.get()
else:
print("iDEA.methods.interacting.solve: solving eigenproblem on CPU...")
energies, spaces = spsla.eigsh(H.tocsr(), k=level, which="SA")
# Reshape and normalise the solutions.
spaces = spaces.reshape((s.x.shape[0],) * s.count + (spaces.shape[-1],))
for j in range(spaces.shape[-1]):
spaces[..., j] = spaces[..., j] / np.sqrt(np.sum(abs(spaces[..., j]) ** 2) * s.dx**s.count)
# Construct the spin part.
symbols = string.ascii_lowercase + string.ascii_uppercase
u = np.array([1, 0])
d = np.array([0, 1])
spin_state = tuple([u if spin == "u" else d for spin in s.electrons])
spin = np.einsum(",".join(symbols[: s.count]) + "->" + "".join(symbols[: s.count]), *spin_state)
spins = np.zeros(shape=((2,) * s.count + (spaces.shape[-1],)))
for i in range(spaces.shape[-1]):
spins[..., i] = spin
# Antisymmetrize.
fulls, spaces, spins, energies = antisymmetrize(s, spaces, spins, energies)
# Populate the state.
if k != -1:
state = iDEA.state.ManyBodyState()
state.space = spaces[..., k]
state.spin = spins[..., k]
state.full = fulls[..., k]
state.energy = energies[k]
return state
else:
states = iDEA.state.ManyBodyStates()
states.spaces = spaces
states.spins = spins
states.fulls = fulls
states.energies = energies
return states
[docs]def propagate_step(
s: iDEA.system.System,
evolution: iDEA.state.ManyBodyEvolution,
H: sps.dia_matrix,
v_ptrb: np.ndarray,
j: int,
dt: float,
objs: tuple,
) -> iDEA.state.ManyBodyEvolution:
r"""
Propagate a many body state forward in time, one time-step, due to a local pertubation.
| Args:
| s: iDEA.system.System, System object.
| evolution: iDEA.state.ManyBodyEvolution, time-dependent evolution.
| H: np.ndarray, Static Hamiltonian [If None this will be computed from s]. (default = None)
| v_ptrb: np.ndarray, Local perturbing potential on the grid of t and x values, indexed as v_ptrb[time,space].
| j: int, Time index.
| dt: float, Time-step.
| objs: tuple. Tuple of objects needed to construct many-body operator (I, generate_terms).
| Returns:
| evolution: iDEA.state.ManyBodyEvolution, time-dependent evolution one time-step evolved.
"""
# Construct the pertubation potential.
vptrb = sps.dia_matrix(np.diag(v_ptrb[j, :]))
terms = objs[1](sps.kron, vptrb, objs[0], s.count)
Vptrb = sps.dia_matrix((s.x.shape[0] ** s.count,) * 2, dtype=float)
for term in terms:
Vptrb += term
# Contruct the perturbed Hamiltonian.
Hp = H + Vptrb
# Evolve.
wavefunction = evolution.td_space[j - 1, ...].reshape(s.x.shape[0] ** s.count)
wavefunction = spsla.expm_multiply(-1.0j * dt * Hp, wavefunction)
evolution.td_space[j, ...] = wavefunction.reshape((s.x.shape[0],) * s.count)
return evolution
[docs]def propagate(
s: iDEA.system.System,
state: iDEA.state.ManyBodyState,
v_ptrb: np.ndarray,
t: np.ndarray,
H: sps.dia_matrix = None,
) -> iDEA.state.ManyBodyEvolution:
r"""
Propagate a many body state forward in time due to a local pertubation.
| Args:
| s: iDEA.system.System, System object.
| state: iDEA.state.ManyBodyState, State to be propigated.
| v_ptrb: np.ndarray, Local perturbing potential on the grid of t and x values, indexed as v_ptrb[time,space].
| t: np.ndarray, Grid of time values.
| H: np.ndarray, Static Hamiltonian [If None this will be computed from s]. (default = None)
| Returns:
| evolution: iDEA.state.ManyBodyEvolution, Solved time-dependent evolution.
"""
# Construct the unperturbed Hamiltonian.
if H is None:
H = hamiltonian(s)
# Compute timestep.
dt = t[1] - t[0]
# Initilise time-dependent wavefunction.
evolution = iDEA.state.ManyBodyEvolution(initial_state=state)
evolution.td_space = np.zeros(shape=t.shape + state.space.shape, dtype=complex)
evolution.td_space[0, ...] = copy.deepcopy(evolution.space)
# Construct objects needed to update potential.
I = sps.identity(s.x.shape[0], format="dia")
partial_operators = lambda A, B, k, n: (A if i + k == n - 1 else B for i in range(n))
fold_partial_operators = lambda f, po: functools.reduce(lambda acc, val: f(val, acc, format="dia"), po)
generate_terms = lambda f, A, B, n: (fold_partial_operators(f, partial_operators(A, B, k, n)) for k in range(n))
objs = (I, generate_terms)
# Propagate.
for j, _ti in enumerate(tqdm(t, desc="iDEA.methods.interacting.propagate: propagating state")):
if j != 0:
propagate_step(s, evolution, H, v_ptrb, j, dt, objs)
# Populate the many-body time-dependent evolution.
evolution.v_ptrb = v_ptrb
evolution.t = t
return evolution