diff --git a/.github/workflows/build_package.yml b/.github/workflows/build_package.yml new file mode 100644 index 0000000..5e064db --- /dev/null +++ b/.github/workflows/build_package.yml @@ -0,0 +1,26 @@ +name: Upload Python Package + +on: + release: + types: [published] + +jobs: + deploy: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.10' + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install build + - name: Build package + run: python -m build + - name: Publish package + uses: pypa/gh-action-pypi-publish@release/v1 + with: + user: __token__ + password: ${{ secrets.PYPI_API_TOKEN }} diff --git a/.github/workflows/linting.yml b/.github/workflows/linting.yml new file mode 100644 index 0000000..bf0ea65 --- /dev/null +++ b/.github/workflows/linting.yml @@ -0,0 +1,44 @@ +name: Lint the sourcecode + +on: + pull_request: + branches: + - dev + - main + workflow_dispatch: + +permissions: + contents: write + +env: + SRC_DIR: src/qm_sim + +jobs: + lint: + + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v3 + - uses: actions/setup-python@v3 + - name: setup + run: | + pip install .[linting] + - name: run black + run: | + black "$SRC_DIR" --check + - name: run isort + run: | + isort "$SRC_DIR" -m 3 --trailing-comma -c + - name: run pylint + run: | + pylint "$SRC_DIR" --fail-under=7 + # Remove -c and --check if we want to force linting in the future + # - name: Commit changes + # run: | + # git config user.name github-actions + # git config user.email github-actions@github.com + # git add . + # git commit -m "[Auto-generated] Linting" + # git push + diff --git a/README.md b/README.md index b7a864e..61495e8 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ Python library for simulation of quantum mechanical systems. To be able to use the [PyTorch](https://pytorch.org/) backend for eigenvalue calculations, run the following command: -`pip install qm-sim .[torch]` +`pip install qm-sim[torch]` This will install the cpu-version of the package. To run GPU calculations, install the version for your system at the [PyTorch website](https://pytorch.org/get-started/locally/) instead. @@ -38,18 +38,31 @@ These are enumerated with increasing level of simulation complexity. To contribute, please open a pull request to the `dev`-branch on [GitHub](https://www.github.com/viljarjf/QM_sim/pulls). +### Linting + +When opening a PR, a linting check is performed. +To ensure your contribution passes the checks, you can run the following + +~~~bash +$ pip install .[linting] +$ black src/qm_sim +$ isort src -m 3 --trailing-comma +$ pylint src --fail-under=7 +~~~ + +### Setup + The following is an example of how to set up VS Code for development, adapt to your IDE of choice. TL;DR: - `pip install -e .` to install in an editable state -- `pip install .` to (re)compile the C++ (subsequent python file edits will not be recognized before another reinstall) -### Requirements +**Requirements** - VS Code - Python extension - Python 3.10 or above -### Setup +**Steps** 1. Clone the repo recursively and open the repo in VS Code. If not cloned recursively, initialize the submodules with `git submodule update --init` 2. Press f1, and run `Python: Create Environment`. Select `.venv` 3. Open a new terminal, which should automatically use the virtual environment. If not, run `.venv\Scripts\activate` on Windows, or `source .venv/bin/activate` on Unix diff --git a/docs/source/conf.py b/docs/source/conf.py index 0717c45..19094da 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -9,7 +9,7 @@ project = 'QM sim' copyright = '2023, Viljar Femoen' author = 'Viljar Femoen' -release = '0.0.3' +release = '0.1.1' # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration diff --git a/examples/01_zero_potential.py b/examples/01_zero_potential.py index f0a31f8..2204420 100644 --- a/examples/01_zero_potential.py +++ b/examples/01_zero_potential.py @@ -6,8 +6,8 @@ from qm_sim.hamiltonian import Hamiltonian from qm_sim.nature_constants import m_e -N = (100,) # Discretisation point count -L = (1e-9,) # System size in meters +N = 100 # Discretisation point count +L = 1e-9 # System size in meters m = m_e # Mass of the particle, here chosen as electron mass n = 4 # The amount of eigenstates to find diff --git a/examples/02_quadratic_potential.py b/examples/02_quadratic_potential.py index 0c028cb..349008c 100644 --- a/examples/02_quadratic_potential.py +++ b/examples/02_quadratic_potential.py @@ -8,8 +8,8 @@ import numpy as np -N = (1000,) # Discretisation point count -L = (1e-9,) # System size in meters +N = 1000 # Discretisation point count +L = 1e-9 # System size in meters m = m_e # Mass of the particle, here chosen as electron mass n = 4 # The amount of eigenstates to find @@ -17,7 +17,7 @@ H = Hamiltonian(N, L, m) # Set potential -x = np.linspace(-L[0]/2, L[0]/2, N[0]) +x = np.linspace(-L/2, L/2, N) def V(t): k = 200 diff --git a/examples/05_temporal_evolution.py b/examples/05_temporal_evolution.py index 5bc1a06..69f9322 100644 --- a/examples/05_temporal_evolution.py +++ b/examples/05_temporal_evolution.py @@ -4,8 +4,8 @@ from qm_sim.nature_constants import e_0, m_e -N = (200,) # Discretisation point count -L = (2e-9,) # System size in meters +N = 200 # Discretisation point count +L = 2e-9 # System size in meters m = m_e # Mass of the particle, here chosen as electron mass t_end = 10e-15 # Simulation time dt = 1e-17 # Simulation data storage stepsize @@ -14,9 +14,8 @@ H = Hamiltonian(N, L, m, temporal_scheme="leapfrog") # Set the potential to a quadratic potential oscilating from side to side -z = np.linspace(-L[0]/2, L[0]/2, N[0]) -Vt = lambda t: 6*z**2 + 3*z*np.abs(z)*np.sin(4e15*t) -H.V = Vt +z = np.linspace(-L/2, L/2, N) +H.V = lambda t: 6*z**2 + 3*z*np.abs(z)*np.sin(4e15*t) # Plot H.plot_temporal(t_end, dt) diff --git a/pyproject.toml b/pyproject.toml index a302aa2..27951a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "qm_sim" -version = "0.1.0" +version = "0.1.1" authors = [ { name="Viljar Femoen", email="author@example.com" }, ] @@ -28,9 +28,10 @@ dependencies = [ "Source" = "https://github.com/viljarjf/QM_sim" [project.optional-dependencies] -test = ["pytest", "matplotlib"] -torch = ["torch"] -docs = ["sphinx", "sphinx-rtd-theme", "sphinx-mdinclude"] +test = ["pytest", "matplotlib",] +torch = ["torch",] +docs = ["sphinx", "sphinx-rtd-theme", "sphinx-mdinclude",] +linting = ["black", "isort", "pylint",] [build-system] requires = [ @@ -46,3 +47,19 @@ wheel.expand-macos-universal-tags = true [tool.pytest.ini_options] testpaths = ["tests"] addopts = "-raP" + +[tool.black] +# this is a regex, not a list +exclude = "nature_constants.py" + +[tool.pylint.'Basic'] +good-names = [ + "N", "L", "m", "x", "y", "z", "i", "j", "t", "t0", + "tn", "dt", "H", "h", "l", "n", "v", "w", "psi", "E", + "V", +] + +[tool.pylint.'Main'] +ignore = [ + "nature_constants.py", +] diff --git a/src/qm_sim/cpp/__init__.py b/src/qm_sim/cpp/__init__.py index 61bbce5..cf59d31 100644 --- a/src/qm_sim/cpp/__init__.py +++ b/src/qm_sim/cpp/__init__.py @@ -3,4 +3,4 @@ C++ code for faster eigenproblem solution. Currently unused, as the performance was just the same at best -""" \ No newline at end of file +""" diff --git a/src/qm_sim/cpp/eigen/__init__.py b/src/qm_sim/cpp/eigen/__init__.py index 0040ffe..821bdcf 100644 --- a/src/qm_sim/cpp/eigen/__init__.py +++ b/src/qm_sim/cpp/eigen/__init__.py @@ -1,4 +1,4 @@ try: from .eigen_wrapper import * except ImportError: - print("Warning: eigen_wrapper C++ module not found") \ No newline at end of file + print("Warning: eigen_wrapper C++ module not found") diff --git a/src/qm_sim/cpp/eigen/scripts/generate_init.py b/src/qm_sim/cpp/eigen/scripts/generate_init.py index b8a3918..c0ecce3 100644 --- a/src/qm_sim/cpp/eigen/scripts/generate_init.py +++ b/src/qm_sim/cpp/eigen/scripts/generate_init.py @@ -4,15 +4,21 @@ """ -import sys, pathlib +import pathlib +import sys + def main(module_name: str): path = pathlib.Path(__file__).parent.parent with open(path / "__init__.py", "w") as file: - file.write(f"""try: + file.write( + f"""try: from .{module_name} import * except ImportError: - print("Warning: {module_name} C++ module not found")""") + print("Warning: {module_name} C++ module not found") +""" + ) + if __name__ == "__main__": if len(sys.argv) < 2: diff --git a/src/qm_sim/eigensolvers/__init__.py b/src/qm_sim/eigensolvers/__init__.py index bd695e1..02a9592 100644 --- a/src/qm_sim/eigensolvers/__init__.py +++ b/src/qm_sim/eigensolvers/__init__.py @@ -9,6 +9,7 @@ """ from typing import Callable + import numpy as np from scipy import sparse as sp @@ -16,11 +17,13 @@ __SOLVERS = {} from .scipy_eigen import scipy_get_eigen + __SOLVERS["scipy"] = scipy_get_eigen # Import guard the pytorch backend since it is optional try: from .pytorch_eigen import torch_get_eigen + __SOLVERS["pytorch"] = torch_get_eigen __SOLVERS["torch"] = torch_get_eigen except (ImportError, RuntimeError): @@ -30,5 +33,6 @@ # Function signature of the solver Eigensolver = Callable[[sp.spmatrix, int, tuple[int]], tuple[np.ndarray, np.ndarray]] + def get_eigensolver(solver: str) -> Eigensolver | None: return __SOLVERS.get(solver) diff --git a/src/qm_sim/eigensolvers/pytorch_eigen.py b/src/qm_sim/eigensolvers/pytorch_eigen.py index 05e1482..80ec0e6 100644 --- a/src/qm_sim/eigensolvers/pytorch_eigen.py +++ b/src/qm_sim/eigensolvers/pytorch_eigen.py @@ -6,7 +6,8 @@ import torch from scipy import sparse as sp -PYTORCH_DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +PYTORCH_DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + def torch_get_eigen(mat: sp.dia_matrix, n: int, shape: tuple[int], **kwargs): """Calculate :code:`n` eigenvalues of :code:`mat`. Reshape output to :code:`shape` diff --git a/src/qm_sim/eigensolvers/scipy_eigen.py b/src/qm_sim/eigensolvers/scipy_eigen.py index 1778fb1..34ff488 100644 --- a/src/qm_sim/eigensolvers/scipy_eigen.py +++ b/src/qm_sim/eigensolvers/scipy_eigen.py @@ -3,13 +3,16 @@ Use the scipy backend eigen solver to skip some overhead """ -from scipy._lib._threadsafety import ReentrancyLock -from scipy.sparse.linalg._eigen.arpack.arpack import _SymmetricArpackParams +import numpy as np from scipy import sparse as sp +from scipy._lib._threadsafety import ReentrancyLock from scipy.sparse.linalg import eigsh -import numpy as np +from scipy.sparse.linalg._eigen.arpack.arpack import _SymmetricArpackParams + -def scipy_get_eigen(mat: sp.dia_matrix, n: int, shape: tuple[int], **kwargs) -> tuple[np.ndarray, np.ndarray]: +def scipy_get_eigen( + mat: sp.dia_matrix, n: int, shape: tuple[int], **kwargs +) -> tuple[np.ndarray, np.ndarray]: """Calculate :code:`n` eigenvalues of :code:`mat`. Reshape output to :code:`shape` :param mat: Matrix to calculate eigenvectors and -values for @@ -22,8 +25,7 @@ def scipy_get_eigen(mat: sp.dia_matrix, n: int, shape: tuple[int], **kwargs) -> :rtype: tuple[np.ndarray(shape = (:code:`n`)), np.ndarray(shape = (:code:`n`, :code:`shape`)] """ if kwargs.get("sigma") is None: - v, w = _eigsh(mat._mul_vector, mat.shape[0], - mat.dtype, k=n, **kwargs) + v, w = _eigsh(mat._mul_vector, mat.shape[0], mat.dtype, k=n, **kwargs) else: # Fallback to default solver v, w = eigsh(mat, k=n, which="SA", **kwargs) @@ -34,22 +36,45 @@ def scipy_get_eigen(mat: sp.dia_matrix, n: int, shape: tuple[int], **kwargs) -> return v, w -def _eigsh(A_OP, n, dtype, k=6, sigma=None, which='SA', v0=None, - ncv=None, maxiter=None, tol=0, return_eigenvectors=True, - ): +def _eigsh( + A_OP, + n, + dtype, + k=6, + sigma=None, + which="SA", + v0=None, + ncv=None, + maxiter=None, + tol=0, + return_eigenvectors=True, +): """Copied from the scipy sourcecode, removing some overhead. See https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.eigsh.html """ mode = 1 M_matvec = None Minv_matvec = None - params = _SymmetricArpackParams(n, k, dtype.char, A_OP, mode, - M_matvec, Minv_matvec, sigma, - ncv, v0, maxiter, which, tol) + params = _SymmetricArpackParams( + n, + k, + dtype.char, + A_OP, + mode, + M_matvec, + Minv_matvec, + sigma, + ncv, + v0, + maxiter, + which, + tol, + ) - with ReentrancyLock("Nested calls to eigs/eighs not allowed: " - "ARPACK is not re-entrant"): + with ReentrancyLock( + "Nested calls to eigs/eighs not allowed: ARPACK is not re-entrant" + ): while not params.converged: params.iterate() - return params.extract(return_eigenvectors) + return params.extract(return_eigenvectors) diff --git a/src/qm_sim/hamiltonian/__init__.py b/src/qm_sim/hamiltonian/__init__.py index fa06632..b328432 100644 --- a/src/qm_sim/hamiltonian/__init__.py +++ b/src/qm_sim/hamiltonian/__init__.py @@ -9,5 +9,5 @@ """ # The spatial Hamiltonian is the default -from .spatial_hamiltonian import SpatialHamiltonian as Hamiltonian from .spatial_hamiltonian import SpatialHamiltonian +from .spatial_hamiltonian import SpatialHamiltonian as Hamiltonian diff --git a/src/qm_sim/hamiltonian/spatial_hamiltonian.py b/src/qm_sim/hamiltonian/spatial_hamiltonian.py index dfa2511..455be31 100644 --- a/src/qm_sim/hamiltonian/spatial_hamiltonian.py +++ b/src/qm_sim/hamiltonian/spatial_hamiltonian.py @@ -4,39 +4,47 @@ """ -from typing import Callable, Any +from collections.abc import Iterable +from typing import Any, Callable import numpy as np from scipy.sparse import dia_matrix +from scipy.sparse.linalg import eigsh as adiabatic_eigsh from tqdm import tqdm -from ..nature_constants import h_bar from .. import plot +from ..eigensolvers import get_eigensolver +from ..nature_constants import h_bar from ..spatial_derivative import get_scheme_order -from ..spatial_derivative.cartesian import nabla, laplacian +from ..spatial_derivative.cartesian import laplacian, nabla from ..temporal_solver import TemporalSolver, get_temporal_solver -from ..eigensolvers import get_eigensolver -from scipy.sparse.linalg import eigsh as adiabatic_eigsh class SpatialHamiltonian: - - def __init__(self, N: tuple, L: tuple, m: float | np.ndarray, - spatial_scheme: str = "three-point", temporal_scheme: str = "leapfrog", - eigensolver: str = "scipy", verbose: bool = True, boundary_condition: str = "zero"): + def __init__( + self, + N: tuple[int] | int, + L: tuple[float] | float, + m: float | np.ndarray, + spatial_scheme: str = "three-point", + temporal_scheme: str = "leapfrog", + eigensolver: str = "scipy", + verbose: bool = True, + boundary_condition: str = "zero", + ): """Non-stationary Hamiltonian in real space. :param N: Discretization count along each axis - :type N: tuple + :type N: tuple[int] | int :param L: System size along each axis - :type L: tuple - :param m: Mass of the particle in the system. + :type L: tuple[float] | float + :param m: Mass of the particle in the system. Can be constant (float) or vary in the simulation area (array). If an array is used, :code:`m.shape == N` must hold :type m: float | np.ndarray - :param spatial_scheme: Finite difference scheme for spatial derivative. - Options are: - + :param spatial_scheme: Finite difference scheme for spatial derivative. + Options are: + - three-point - five-point - seven-point @@ -76,15 +84,30 @@ def __init__(self, N: tuple, L: tuple, m: float | np.ndarray, Defaults to "zero" :type boundary_condition: str, optional """ - + # 1D inputs + if isinstance(N, int): + N = (N,) + if isinstance(L, (float, int)): + L = (L,) + + # Allow any iterable that can be converted to a tuple + if isinstance(N, Iterable): + N = tuple(N) + if isinstance(L, Iterable): + L = tuple(L) + + # Check type + if not isinstance(N, tuple) or not all(isinstance(i, int) for i in N): + raise ValueError(f"Param `N` must be int or tuple of ints, got {type(N)}") + if not isinstance(L, tuple) or not all(isinstance(i, (float, int)) for i in L): + raise ValueError(f"Param `L` must be float or tuple, got {type(L)}") + if len(N) != len(L): raise ValueError("`N`and `L`must have same length") - - self.N = tuple(N) - self.L = tuple(L) - self._dim = len(N) - self.delta = [Li / Ni for Li, Ni in zip(L, N)] - + + self.N = N + self.L = L + self.eigensolver = get_eigensolver(eigensolver) if self.eigensolver is None: raise ValueError(f"Eigensolver {eigensolver} not found") @@ -93,44 +116,46 @@ def __init__(self, N: tuple, L: tuple, m: float | np.ndarray, if order is None: raise ValueError("Requested finite difference is invalid") + self._dim = len(N) + self.delta = [Li / Ni for Li, Ni in zip(L, N)] + # Handle non-isotropic effective mass if isinstance(m, np.ndarray) and np.all(m == m.flat[0]): m = m.flatten()[0] if isinstance(m, np.ndarray): print("Warning: Continuity is NOT satisfied (yet) with non-isotropic mass") if m.shape != self.N: - raise ValueError(f"Inconsistent shape of `m`: {m.shape}, should be {self.N}") + raise ValueError( + f"Inconsistent shape of `m`: {m.shape}, should be {self.N}" + ) m_inv = 1 / m.flatten() _n = nabla(N, L, order=order, boundary_condition=boundary_condition) _n2 = laplacian(N, L, order=order, boundary_condition=boundary_condition) # nabla m_inv nabla + m_inv nabla^2 - _n.data *= _n @ m_inv # First term - _n2.data *= m_inv # Second term + _n.data *= _n @ m_inv # First term + _n2.data *= m_inv # Second term self.mat = _n + _n2 else: - self.mat = laplacian(N, L, order=order, boundary_condition=boundary_condition) - self.mat *= 1/m - + self.mat = laplacian( + N, L, order=order, boundary_condition=boundary_condition + ) + self.mat *= 1 / m + self._centerline_index = list(self.mat.offsets).index(0) self._default_data = None - # Prefactor in hamiltonian. - # Is either float or array, depending on `m` - h0 = -h_bar**2 / 2 - - # Multiplying the diagonal data directly is easier - # if we have non-constant mass - self.mat.data *= h0 + # Multiply the laplacian with the remaining factors + self.mat.data *= -(h_bar**2) / 2 # static zero potential by default - self._V = lambda t: np.zeros(shape = N) + self._V = lambda t: np.zeros(shape=N) self.verbose = verbose self._temporal_solver = get_temporal_solver(temporal_scheme) - + @property def V(self) -> Callable[[float], np.ndarray]: """Potential as a function of time @@ -144,31 +169,34 @@ def V(self) -> Callable[[float], np.ndarray]: def V(self, V: np.ndarray | Callable[[float], np.ndarray]): """Set a (potentially time dependent) potential for the QM-system's Hamiltonian - :param V: The energetic value of the potential at a given point associated with given array indices, - if callable, the call variable will represent a changable parameter (usually time) with a - return type identical to the static case where V is an np.ndarray + :param V: The energetic value of the potential at a given point associated with + given array indices, + if callable, the call variable will represent a changable parameter + (usually time) with a return type identical to the static case where V is an np.ndarray :type V: np.ndarray | Callable[[float], np.ndarray] """ if not callable(V): # Avoid the lambda func referring to itself array_V = V V = lambda t: array_V - + if V(0).shape != self.shape: raise ValueError(f"Inconsistent shape. Shape must be {self.shape}") self._V = V def eigen(self, n: int, t: float = 0, **kwargs) -> tuple[np.ndarray, np.ndarray]: - """Calculate the n smallest eigenenergies and the corresponding eigenstates of the hamiltonian + """Calculate the n smallest eigenenergies and the corresponding eigenstates of + the hamiltonian :param n: Amount of eigenenergies/states to output :type n: int - :param t: If the potential is time-dependent, solves the Time-independent Schrödinger eq. as if it was frozen at time t. + :param t: If the potential is time-dependent, solves the Time-independent + Schrödinger eq. as if it was frozen at time t. Does nothing if potential is time-independent. Defaults to 0 :type t: float, optional - :return: + :return: - Eigenenergies - Normalised eigenstates @@ -184,27 +212,30 @@ def eigen(self, n: int, t: float = 0, **kwargs) -> tuple[np.ndarray, np.ndarray] E, psi = self.eigensolver(self(t), n, self.N) # calculate normalisation factor - nf = [psi[i, :]**2 for i in range(n)] + normalisation_factor = [psi[i, :] ** 2 for i in range(n)] for i, (L, N) in enumerate(zip(self.L, self.N)): dx = L / N for j in range(n): - nf[j] = np.trapz(nf[j], dx=dx) + normalisation_factor[j] = np.trapz(normalisation_factor[j], dx=dx) # normalise for i in range(n): - psi[i] /= nf[i]**0.5 + psi[i] /= normalisation_factor[i] ** 0.5 return E, psi - - def adiabatic_evolution(self, E_n: float, t0: float, dt : float, steps: int) -> tuple[np.ndarray, np.ndarray]: + def adiabatic_evolution( + self, E_n: float, t0: float, dt: float, steps: int + ) -> tuple[np.ndarray, np.ndarray]: """Adiabatically evolve an eigenstate with a slowly varying time-dependent potential with - energy (close to) E_n using the Adiabatic approximation. + energy (close to) :code:`E_n` using the Adiabatic approximation. https://en.wikipedia.org/wiki/Adiabatic_theorem - Note: This is only valid given that the adiabatic theorem holds, ie. no degeneracy and a gap between - eigenvalues. Current implementation assumes this holds and does not check if it does (yet?). - There is no mathematical guarantee (yet?) that the iterative solver will "hug" the correct eigenvector - at every step, but it should be good if V varies smoothly enough and dt is small enough. - + Note: This is only valid given that the adiabatic theorem holds, ie. no degeneracy and a + gap betweeneigenvalues. Current implementation assumes this holds and does not check if + it does (yet?). + There is no mathematical guarantee (yet?) that the iterative solver will "hug" the + correct eigenvector at every step, but it should be good if V varies smoothly enough and + :code:`dt` is small enough. + :param E_n: The Eigenvalue for you want to adiabatically evolve :type E_n: float @@ -217,34 +248,42 @@ def adiabatic_evolution(self, E_n: float, t0: float, dt : float, steps: int) -> :return: (E(t), Psi(t)), Time evolution of the eigenstate. :rtype: tuple[np.ndarray(shape = steps), np.ndarray(shape = (N, steps))] """ - Psi_t = np.empty(shape = (*self.N,steps+1)) - En_t = np.empty(steps+1) + Psi_t = np.empty(shape=(*self.N, steps + 1)) + En_t = np.empty(steps + 1) t = t0 # initialize _, Psi_t[:, :, 0] = self.eigen(1, t, sigma=E_n, is_adiabatic=True) En_t[0] = E_n - for i in tqdm(range(1, steps+1), desc="Adiabatic evolution", disable=not self.verbose): + for i in tqdm( + range(1, steps + 1), desc="Adiabatic evolution", disable=not self.verbose + ): t += dt - En_t[i], Psi_t[:,:,i] = self.eigen( 1, t, - # smartly condition eigsolver to "hug" the single eigenvalue solution; eigenvector and eigenvalue should be - # far closer to the previous one than any other if the adiabatic theorem is fulfilled - sigma=En_t[i-1], - v0=-Psi_t[:,:,i-1], - is_adiabatic=True + En_t[i], Psi_t[:, :, i] = self.eigen( + 1, + t, + # smartly condition eigsolver to "hug" the single eigenvalue solution; + # eigenvector and eigenvalue should be far closer to the previous one + # than any other if the adiabatic theorem is fulfilled + sigma=En_t[i - 1], + v0=-Psi_t[:, :, i - 1], + is_adiabatic=True, ) return En_t, Psi_t - def temporal_evolution(self, t0: float, t_final: float, dt_storage: float = None, - psi_0: np.ndarray = None) -> tuple[np.ndarray, np.ndarray]: - - + def temporal_evolution( + self, + t0: float, + t_final: float, + dt_storage: float = None, + psi_0: np.ndarray = None, + ) -> tuple[np.ndarray, np.ndarray]: # Default: superposition of 1st and 2nd eigenstate if psi_0 is None: _, psi = self.eigen(2) psi_0 = 2**-0.5 * (psi[0] + psi[1]) - + # Calculate a good dt from von Neumann analysis of the leapfrog scheme # NOTE: this assumes the temporal part is at most 4x the static part, # and that the potential takes its maximum somwhere at t=t0 @@ -254,18 +293,19 @@ def temporal_evolution(self, t0: float, t_final: float, dt_storage: float = None E_max = max( abs(V_min), abs(V_max + 4 * np.max(self.mat.data)), - ) + ) dt = 0.25 * h_bar / E_max # solve - f = lambda t: 1/(1j * h_bar) * self.__call__(t) - solver = self._temporal_solver(f, self.shape) - t, psi = solver.iterate(psi_0.astype(np.complex128), t0, - t_final, dt, dt_storage, self.verbose) + func = lambda t: 1 / (1j * h_bar) * self.__call__(t) + solver = self._temporal_solver(func, self.shape) + t, psi = solver.iterate( + psi_0.astype(np.complex128), t0, t_final, dt, dt_storage, self.verbose + ) return t, psi - temporal_evolution.__doc__ = TemporalSolver.iterate.__doc__ + temporal_evolution.__doc__ = TemporalSolver.iterate.__doc__ def plot_eigen(self, n: int, t: float = 0): """Calculate and plot :code:`n` eigenstates at time :code:`t` @@ -278,8 +318,9 @@ def plot_eigen(self, n: int, t: float = 0): E, psi = self.eigen(n, t) plot.eigen(E, psi) - - def plot_temporal(self, t_final: float, dt: float, psi_0: np.ndarray = None, t0: float = 0): + def plot_temporal( + self, t_final: float, dt: float, psi_0: np.ndarray = None, t0: float = 0 + ): """Plot the temporal evolution of the eigenstates :param t_final: Simulation end time @@ -294,7 +335,6 @@ def plot_temporal(self, t_final: float, dt: float, psi_0: np.ndarray = None, t0: t, psi = self.temporal_evolution(t0, t_final, dt, psi_0) plot.temporal(t, psi, self.V) - def plot_potential(self, t: float = 0): """Plot the potential at a given time @@ -313,11 +353,11 @@ def __add__(self, other: np.ndarray) -> dia_matrix: """ if self._default_data is None: self._default_data = self.mat.data.copy() - + self.mat.data = self._default_data.copy() self.mat.data[self._centerline_index, :] += other.flatten() return self.mat - + def __call__(self, t: float) -> dia_matrix: """Get the matrix at a given time @@ -327,7 +367,7 @@ def __call__(self, t: float) -> dia_matrix: :rtype: dia_matrix """ return self + self.V(t) - + @property def shape(self) -> tuple[int]: """System shape @@ -348,7 +388,7 @@ def N_total(self) -> int: for j in self.N: i *= j return i - + @property def ndim(self) -> int: """System dimensionallity @@ -367,9 +407,9 @@ def __matmul__(self, other: Any) -> Any: :rtype: Any """ return self.mat @ other - + def _fast_matmul_op(self, t: float = 0) -> Callable[[Any], np.ndarray]: - """Evaluate the Hamiltonian at a given time, + """Evaluate the Hamiltonian at a given time, and return a matrix-vector multiplication function :param t: Time at which to evaluate the Hamiltonian, defaults to 0 @@ -379,7 +419,7 @@ def _fast_matmul_op(self, t: float = 0) -> Callable[[Any], np.ndarray]: """ mat = self(t) return mat._mul_vector - + def asarray(self) -> np.ndarray: """Return the data matrix as a dense array @@ -387,5 +427,3 @@ def asarray(self) -> np.ndarray: :rtype: np.ndarray """ return self.mat.toarray() - - \ No newline at end of file diff --git a/src/qm_sim/plot.py b/src/qm_sim/plot.py index 333db78..fbdf82c 100644 --- a/src/qm_sim/plot.py +++ b/src/qm_sim/plot.py @@ -27,25 +27,26 @@ def _get_plot_fun(ndim: int, ax: plt.Axes = None): # It would be impressive if this is ever executed raise ValueError(f"Invalid system dimensionality: {ndim}D") + def _shape_from_int(n: int) -> tuple[int, int]: """Get plot shape from n (plot count)""" if n < 1: raise ValueError("Plot count must be positive") shapes = { - 1: (1,1), - 2: (1,2), - 3: (1,3), - 4: (2,2), - 5: (2,3), - 6: (2,3), - 7: (2,4), - 8: (2,4), - 9: (3,3), - 10: (3,4), - 11: (3,4), - 12: (3,4), + 1: (1, 1), + 2: (1, 2), + 3: (1, 3), + 4: (2, 2), + 5: (2, 3), + 6: (2, 3), + 7: (2, 4), + 8: (2, 4), + 9: (3, 3), + 10: (3, 4), + 11: (3, 4), + 12: (3, 4), } - return shapes.get(n, (1,n)) + return shapes.get(n, (1, n)) def eigen(E: np.ndarray, psi: np.ndarray): @@ -59,18 +60,19 @@ def eigen(E: np.ndarray, psi: np.ndarray): n = E.shape[0] ndim = len(psi.shape) - 1 shape = _shape_from_int(n) - + plt.figure() - plt.suptitle("$|\Psi|^2$") + plt.suptitle("$|\\Psi|^2$") plot = _get_plot_fun(ndim) for i in range(n): - plt.subplot(*shape, i+1) + plt.subplot(*shape, i + 1) plt.title(f"E{i} = {E[i] / e_0 :.3f} eV") - plot(abs(psi[i])**2) + plot(abs(psi[i]) ** 2) plt.tight_layout() plt.show() return + def temporal(t: np.ndarray, psi: np.ndarray, Vt: Callable[[float], np.ndarray]): """Create an animation of the wave function, alongside the potential @@ -87,17 +89,17 @@ def temporal(t: np.ndarray, psi: np.ndarray, Vt: Callable[[float], np.ndarray]): V = np.array([Vt(tn) for tn in t]) # We want to plot the probability distribution - psi2 = np.abs(psi)**2 + psi2 = np.abs(psi) ** 2 # Plot the results fig, (ax1, ax2) = plt.subplots(2, 1) ax1_plot = _get_plot_fun(ndim, ax1) ax2_plot = _get_plot_fun(ndim, ax2) - psi_plot, = ax1_plot(psi2[0, :]) - V_plot, = ax2_plot(V[0, :] / e_0) + (psi_plot,) = ax1_plot(psi2[0, :]) + (V_plot,) = ax2_plot(V[0, :] / e_0) - ax1.set_title(f"$|\Psi|^2$") + ax1.set_title(f"$|\\Psi|^2$") ax2.set_title("Potential [eV]") if ndim == 1: ax1.set_ylim(0, np.max(psi2) * 1.1) @@ -111,9 +113,14 @@ def temporal(t: np.ndarray, psi: np.ndarray, Vt: Callable[[float], np.ndarray]): ims = [(psi_plot, V_plot)] for n in range(psi2.shape[0]): - psi_plot, = ax1_plot(psi2[n, ...], animated=True) - V_plot, = ax2_plot(V[n, ...] / e_0, animated=True) - ims.append((psi_plot, V_plot,)) + (psi_plot,) = ax1_plot(psi2[n, ...], animated=True) + (V_plot,) = ax2_plot(V[n, ...] / e_0, animated=True) + ims.append( + ( + psi_plot, + V_plot, + ) + ) ani = ArtistAnimation(fig, ims, blit=True, interval=50) plt.show() diff --git a/src/qm_sim/spatial_derivative/__init__.py b/src/qm_sim/spatial_derivative/__init__.py index 1345b2f..60521ce 100644 --- a/src/qm_sim/spatial_derivative/__init__.py +++ b/src/qm_sim/spatial_derivative/__init__.py @@ -11,5 +11,6 @@ "nine-point": 8, } + def get_scheme_order(scheme: str) -> int | None: return _SCHEMES.get(scheme) diff --git a/src/qm_sim/spatial_derivative/cartesian.py b/src/qm_sim/spatial_derivative/cartesian.py index 1298c49..f804bbc 100644 --- a/src/qm_sim/spatial_derivative/cartesian.py +++ b/src/qm_sim/spatial_derivative/cartesian.py @@ -1,10 +1,14 @@ - import numpy as np from scipy import sparse as sp -def nabla(N: tuple[int], L: tuple[float], order: int = 2, dtype: type = np.float64, - boundary_condition: str = "zero") -> sp.dia_matrix: +def nabla( + N: tuple[int], + L: tuple[float], + order: int = 2, + dtype: type = np.float64, + boundary_condition: str = "zero", +) -> sp.dia_matrix: """Finite difference derivative in cartesian coordinates. Uses central stencil. @@ -34,7 +38,7 @@ def nabla(N: tuple[int], L: tuple[float], order: int = 2, dtype: type = np.float Defaults to 2. :type order: int, optional - :param dtype: datatype of matrix. + :param dtype: datatype of matrix. Defaults to np.float64 :type dtype: type, optional :raises NotImplementedError: if requested order is not available @@ -49,28 +53,35 @@ def nabla(N: tuple[int], L: tuple[float], order: int = 2, dtype: type = np.float Defaults to "zero" :type boundary_condition: str, optional """ - # Lookup for first half of stencil. - # A 0 is appended for the central element, + # Lookup for first half of stencil. + # A 0 is appended for the central element, # and the rest is mirrored with a sign change later match order: case 2: - stencil = [-1/2] + stencil = [-1 / 2] case 4: - stencil = [1/12, -2/3] + stencil = [1 / 12, -2 / 3] case 6: - stencil = [-1/60, 3/20, -3/4] + stencil = [-1 / 60, 3 / 20, -3 / 4] case 8: - stencil = [1/280, -4/105, 1/5, -4/5] + stencil = [1 / 280, -4 / 105, 1 / 5, -4 / 5] case _: - raise NotImplementedError(f"Finite difference scheme not found for {order = }") + raise NotImplementedError( + f"Finite difference scheme not found for {order = }" + ) stencil += [0] stencil = _mirror_sign_list(stencil) return _matrix_from_central_stencil(stencil, 1, N, L, dtype, boundary_condition) -def laplacian(N: tuple[int], L: tuple[float], order: int = 2, dtype: type = np.float64, - boundary_condition: str = "zero") -> sp.dia_matrix: +def laplacian( + N: tuple[int], + L: tuple[float], + order: int = 2, + dtype: type = np.float64, + boundary_condition: str = "zero", +) -> sp.dia_matrix: """Finite difference double derivative in cartesian coordinates. Uses central stencil @@ -100,7 +111,7 @@ def laplacian(N: tuple[int], L: tuple[float], order: int = 2, dtype: type = np.f Defaults to 2. :type order: int, optional - :param dtype: datatype of matrix. + :param dtype: datatype of matrix. Defaults to np.float64 :type dtype: type, optional :raises NotImplementedError: if requested order is not available @@ -120,29 +131,38 @@ def laplacian(N: tuple[int], L: tuple[float], order: int = 2, dtype: type = np.f case 2: stencil = [1, -2] case 4: - stencil = [-1/12, 4/3, -5/2] + stencil = [-1 / 12, 4 / 3, -5 / 2] case 6: - stencil = [1/90, -3/20, 3/2, -49/18] + stencil = [1 / 90, -3 / 20, 3 / 2, -49 / 18] case 8: - stencil = [-1/560, 8/315, -1/5, 8/5, -205/72] + stencil = [-1 / 560, 8 / 315, -1 / 5, 8 / 5, -205 / 72] case _: - raise NotImplementedError(f"Finite difference scheme not found for {order = }") - + raise NotImplementedError( + f"Finite difference scheme not found for {order = }" + ) + stencil = _mirror_list(stencil) return _matrix_from_central_stencil(stencil, 2, N, L, dtype, boundary_condition) - -def _matrix_from_central_stencil(stencil: list[float], power: int, - N: tuple[int], L: tuple[float], dtype: type, - boundary_condition: str = "zero") -> sp.dia_matrix: + +def _matrix_from_central_stencil( + stencil: list[float], + power: int, + N: tuple[int], + L: tuple[float], + dtype: type, + boundary_condition: str = "zero", +) -> sp.dia_matrix: """ Creates a full matrix from a central stencil. Determines indices from stencil """ available_boundary_conditions = ["zero", "periodic"] if boundary_condition not in available_boundary_conditions: - raise ValueError(f"Invalid boundary condition: {boundary_condition}. Options are: " - + ", ".join(available_boundary_conditions)) + raise ValueError( + f"Invalid boundary condition: {boundary_condition}. Options are: " + + ", ".join(available_boundary_conditions) + ) if boundary_condition == "zero": indices = np.arange(len(stencil)) - len(stencil) // 2 @@ -173,16 +193,17 @@ def _matrix_from_central_stencil(stencil: list[float], power: int, axis_indices = np.array(periodic_indices) stencil = np.array(periodic_stencil) - mat = np.zeros((1,1)) + mat = np.zeros((1, 1)) for l, n, indices in zip(L, N, axis_indices): h = l / n - next_mat = 1/h**power * sp.diags( + next_mat = sp.diags( stencil, indices, shape=(n, n), dtype=dtype, - format="dia" - ) + format="dia", + ) + next_mat *= 1 / h**power mat = sp.kronsum(mat, next_mat, format="dia") return mat diff --git a/src/qm_sim/temporal_solver/__init__.py b/src/qm_sim/temporal_solver/__init__.py index bf10845..2f4a022 100644 --- a/src/qm_sim/temporal_solver/__init__.py +++ b/src/qm_sim/temporal_solver/__init__.py @@ -5,10 +5,7 @@ """ -from .base import get_temporal_solver -from .base import TemporalSolver - -# Tell python where to find subclasses +from .base import TemporalSolver, get_temporal_solver from .crank_nicolson import CrankNicolson from .leapfrog import Leapfrog -from .scipy_solvers import RungeKutta45, RungeKutta23, DOP853, BDF +from .scipy_solvers import BDF, DOP853, RungeKutta23, RungeKutta45 diff --git a/src/qm_sim/temporal_solver/base.py b/src/qm_sim/temporal_solver/base.py index 18b6084..d7ed39a 100644 --- a/src/qm_sim/temporal_solver/base.py +++ b/src/qm_sim/temporal_solver/base.py @@ -22,14 +22,16 @@ class TemporalSolver(ABC): #: Is the method explicit or implicit? explicit: bool - #: Is the method stable? - #: If only conditionally stable, this will be true + #: Is the method stable? + #: If only conditionally stable, this will be true #: and :code:`dt` will be forced into its stable range stable: bool _skip_registration: bool = False - def __init__(self, H: Callable[[float], np.ndarray], output_shape: tuple[int] = None): + def __init__( + self, H: Callable[[float], np.ndarray], output_shape: tuple[int] = None + ): """Initialize a temporal solver :param H: Function of time, representing the temporal derivative at that time @@ -45,7 +47,9 @@ def __init__(self, H: Callable[[float], np.ndarray], output_shape: tuple[int] = self.output_shape = output_shape def tqdm(self, t_start: float, t_end: float, enable: bool) -> tqdm: - pbar = tqdm(desc=self.name + " solver", total=t_end - t_start, disable=not enable) + pbar = tqdm( + desc=self.name + " solver", total=t_end - t_start, disable=not enable + ) pbar.bar_format = "{l_bar}{bar}| {n:#.02g}/{total:#.02g}" # Add func to update with t, not dt @@ -53,19 +57,25 @@ def tqdm(self, t_start: float, t_end: float, enable: bool) -> tqdm: return pbar - @abstractmethod - def iterate(self, v_0: np.ndarray, t0: float, t_final: float, - dt: float, dt_storage: float = None, verbose: bool = True) -> tuple[np.ndarray, np.ndarray]: + def iterate( + self, + v_0: np.ndarray, + t0: float, + t_final: float, + dt: float, + dt_storage: float = None, + verbose: bool = True, + ) -> tuple[np.ndarray, np.ndarray]: """ Iterate the time propagation scheme. Store the current state every :code:`dt_storage` Args: - t_final (float): + t_final (float): End time for calculations - dt_storage (float, optional): - Data storage period. + dt_storage (float, optional): + Data storage period. If None, store each calculation :code:`dt` Defaults to None. @@ -78,9 +88,7 @@ def iterate(self, v_0: np.ndarray, t0: float, t_final: float, pass def __init_subclass__(cls): - """Register subclasses of :class:`TemporalSolver` - - """ + """Register subclasses of :class:`TemporalSolver`""" if cls._skip_registration: return # Register new subclasses of TemporalSolver @@ -101,5 +109,6 @@ def get_temporal_solver(scheme: str) -> type[TemporalSolver]: """ if scheme in _SCHEMES.keys(): return _SCHEMES[scheme] - raise ValueError(f"Scheme {scheme} not found. Options are:\n" - + "\n".join(_SCHEMES.keys())) + raise ValueError( + f"Scheme {scheme} not found. Options are:\n" + "\n".join(_SCHEMES.keys()) + ) diff --git a/src/qm_sim/temporal_solver/crank_nicolson.py b/src/qm_sim/temporal_solver/crank_nicolson.py index b1a9201..a5709af 100644 --- a/src/qm_sim/temporal_solver/crank_nicolson.py +++ b/src/qm_sim/temporal_solver/crank_nicolson.py @@ -15,25 +15,34 @@ def I_plus_aH(a: complex, H: dia_matrix) -> dia_matrix: out.data[zero_ind] += 1 return out + class CrankNicolson(TemporalSolver): order = 2 explicit = False stable = True name = "crank-nicolson" - def __init__(self, H: Callable[[float], np.ndarray], output_shape: tuple[int] = None): + def __init__( + self, H: Callable[[float], np.ndarray], output_shape: tuple[int] = None + ): TemporalSolver.__init__(self, H, output_shape) if len(self.output_shape) != 1: raise ValueError("Crank-Nicolson solver only supports 1D systems") - def iterate(self, v_0: np.ndarray, t0: float, t_final: float, - dt: float, dt_storage: float = None, verbose: bool = True) -> tuple[np.ndarray, np.ndarray]: - + def iterate( + self, + v_0: np.ndarray, + t0: float, + t_final: float, + dt: float, + dt_storage: float = None, + verbose: bool = True, + ) -> tuple[np.ndarray, np.ndarray]: if dt_storage is None: dt_storage = dt H = self.H - prefactor = dt/2 + prefactor = dt / 2 tn = t0 psi_n = v_0 @@ -43,7 +52,6 @@ def iterate(self, v_0: np.ndarray, t0: float, t_final: float, with self.tqdm(t0, t_final, verbose) as pbar: while tn < t_final: - # psi^n+1 = psi^n + dt/2*(F^n+1 + F^n) # F^n = 1/ihbar * H^n @ psi^n # => psi^n+1 = psi^n + dt/2ihbar*(H^n+1 @ psi^n+1 + H^n @ psi^n) @@ -51,13 +59,13 @@ def iterate(self, v_0: np.ndarray, t0: float, t_final: float, lhs = I_plus_aH(-prefactor, H(tn + dt)) rhs = I_plus_aH(prefactor, H(tn)) @ psi_n - + # Reformulate the solve_banded function in terms of the dia_matrix class psi_n = solve_banded( (-lhs.offsets[0], lhs.offsets[-1]), lhs.data[::-1, :], rhs, - ) + ) pbar.update(dt) tn += dt diff --git a/src/qm_sim/temporal_solver/leapfrog.py b/src/qm_sim/temporal_solver/leapfrog.py index 1c56119..5aa3bcc 100644 --- a/src/qm_sim/temporal_solver/leapfrog.py +++ b/src/qm_sim/temporal_solver/leapfrog.py @@ -6,11 +6,18 @@ class Leapfrog(TemporalSolver): order = 2 explicit = True - stable = True # conditionally stable, dt is chosen accordingly + stable = True # conditionally stable, dt is chosen accordingly name = "leapfrog" - def iterate(self, v_0: np.ndarray, t0: float, t_final: float, - dt: float, dt_storage: float = None, verbose: bool = True) -> tuple[np.ndarray, np.ndarray]: + def iterate( + self, + v_0: np.ndarray, + t0: float, + t_final: float, + dt: float, + dt_storage: float = None, + verbose: bool = True, + ) -> tuple[np.ndarray, np.ndarray]: # Override initial condition to preserve 2nd order accuracy # psi_0 = psi_half - c*H_half*psi_half # psi_1 = psi_half + c*H_half*psi_half @@ -22,21 +29,20 @@ def iterate(self, v_0: np.ndarray, t0: float, t_final: float, H = self.H psi_half = v_0.flatten() - psi_0 = psi_half - dt / 2 * (H(t0 + dt/2) @ psi_half) - psi_1 = psi_half + dt / 2 * (H(t0 + dt/2) @ psi_half) + psi_0 = psi_half - dt / 2 * (H(t0 + dt / 2) @ psi_half) + psi_1 = psi_half + dt / 2 * (H(t0 + dt / 2) @ psi_half) - tn = t0 + dt/2 + tn = t0 + dt / 2 psi = [psi_1.reshape(self.output_shape)] t = [tn] with self.tqdm(t0, t_final, verbose) as pbar: while tn < t_final: - # psi^n+1 = psi^n-1 + 2*dt*F^n # F^n = 1/ihbar * H^n @ psi^n # H^n = H0 + V^n - psi_2 = 2*dt * (H(tn) @ psi_1) + psi_0 + psi_2 = 2 * dt * (H(tn) @ psi_1) + psi_0 psi_0, psi_1 = psi_1, psi_2 pbar.update(dt) diff --git a/src/qm_sim/temporal_solver/scipy_solvers.py b/src/qm_sim/temporal_solver/scipy_solvers.py index 0db0c64..1aa05c6 100644 --- a/src/qm_sim/temporal_solver/scipy_solvers.py +++ b/src/qm_sim/temporal_solver/scipy_solvers.py @@ -12,9 +12,15 @@ class ScipySolver(TemporalSolver): _skip_registration = True - def iterate(self, v_0: np.ndarray, t0: float, t_final: float, - dt: float, dt_storage: float = None, verbose: bool = True) -> tuple[np.ndarray, np.ndarray]: - + def iterate( + self, + v_0: np.ndarray, + t0: float, + t_final: float, + dt: float, + dt_storage: float = None, + verbose: bool = True, + ) -> tuple[np.ndarray, np.ndarray]: pbar = self.tqdm(t0, t_final, verbose) # Precalculate the coefficient for negligible speedup @@ -24,23 +30,24 @@ def ode_fun(t, y): return self.H(t) @ y sol = solve_ivp( - ode_fun, - [t0, t_final], - v_0.flatten(), + ode_fun, + [t0, t_final], + v_0.flatten(), t_eval=np.arange(t0, t_final, dt_storage), first_step=dt, method=self.method, - ) + ) t = sol.t psi = [sol.y[:, i].reshape(*self.output_shape) for i in range(len(t))] return t, np.array(psi) - + def __init_subclass__(cls): cls.name = "scipy-" + cls.name cls._skip_registration = False return super().__init_subclass__() + class RungeKutta45(ScipySolver): order = 5 explicit = True @@ -48,6 +55,7 @@ class RungeKutta45(ScipySolver): name = "Runge-Kutta 5(4)" method = "RK45" + class RungeKutta23(ScipySolver): order = 3 explicit = True @@ -55,6 +63,7 @@ class RungeKutta23(ScipySolver): name = "Runge-Kutta 3(2)" method = "RK23" + class DOP853(ScipySolver): order = 8 explicit = True @@ -62,6 +71,7 @@ class DOP853(ScipySolver): name = "DOP853" method = "DOP853" + # Solver does not support complex numbers # class Radau(ScipySolver): # order = 5 @@ -70,6 +80,7 @@ class DOP853(ScipySolver): # name = "radau" # method = "Radau" + class BDF(ScipySolver): order = None explicit = False @@ -77,6 +88,7 @@ class BDF(ScipySolver): name = "backwards-differentiation" method = "BDF" + # Solver does not support complex numbers # class LSODA(ScipySolver): # order = None diff --git a/tests/test_hamiltonian.py b/tests/test_hamiltonian.py index 4136e73..5500f4a 100644 --- a/tests/test_hamiltonian.py +++ b/tests/test_hamiltonian.py @@ -10,7 +10,7 @@ def test_constant_mass(): - H = Hamiltonian((10,), (1,), 1, spatial_scheme="nine-point", verbose=False) + H = Hamiltonian(10, 1, 1, spatial_scheme="nine-point", verbose=False) plt.figure() plt.title("Constant mass hamiltonian") plt.imshow(H.asarray()) @@ -26,7 +26,7 @@ def test_constant_mass_2D(): plt.show() def test_array_mass(): - N = (200,) + N = 200 # Source: # https://www.researchgate.net/publication/260544509_Band-gap_shift_in_heavily_doped_n-type_Al03Ga07As_alloys @@ -39,16 +39,16 @@ def test_array_mass(): kernel = np.ones((n)) / n m = m_algaas * np.ones(N) - m[N[0] // 2 - N[0] // 6 : N[0] // 2 + N[0] // 6] = m_gaas + m[N // 2 - N // 6 : N // 2 + N // 6] = m_gaas m *= const.m_e m = convolve(m, kernel, mode="same") V = V_algaas * np.ones(N) - V[N[0] // 2 - N[0] // 6 : N[0] // 2 + N[0] // 6] = V_gaas + V[N // 2 - N // 6 : N // 2 + N // 6] = V_gaas V *= const.e_0 V = convolve(V, kernel, mode="same") - H = Hamiltonian(N, (20e-9, ), m, verbose=False) + H = Hamiltonian(N, 20e-9, m, verbose=False) H.V = V H.plot_potential() @@ -70,15 +70,15 @@ def test_array_mass(): plt.show() def test_potential(): - N = (100,) - L = (2e-9,) + N = 100 + L = 2e-9 r = np.ones(N) m = const.m_e * r n = 5 H0 = Hamiltonian(N, L, m, spatial_scheme="three-point", verbose=False) - z = np.linspace(-L[0]/2, L[0]/2, N[0]) + z = np.linspace(-L/2, L/2, N) V = 100*z**2 H0.V = V @@ -98,8 +98,8 @@ def test_potential(): plt.show() def test_eigen(): - N = (100,) - L = (2e-9,) + N = 100 + L = 2e-9 r = np.ones(N) m = const.m_e * r @@ -121,8 +121,8 @@ def test_eigen(): plt.show() def test_temporal(): - N = (200,) - L = (2e-9,) + N = 200 + L = 2e-9 m = const.m_e t_end = 10e-15 dt = 1e-17 @@ -132,7 +132,7 @@ def test_temporal(): for scheme in schemes: H.append(Hamiltonian(N, L, m, temporal_scheme=scheme, verbose=False)) - z = np.linspace(-L[0]/2, L[0]/2, N[0]) + z = np.linspace(-L/2, L/2, N) Vt = lambda t: 6*z**2 + 3*z*np.abs(z)*np.sin(4e15*t) for Hi in H: Hi.V = Vt