Merge pull request #34 from viljarjf/feat/non-array-1D-params

Non-array 1 d params

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

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
 

examples/02_quadratic_potential.py

@@ -8,16 +8,16 @@
 
 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
 
 # Set hamiltonian
 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

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)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "qm_sim"
-version = "0.1.0"
+version = "0.1.1"
 authors = [
   { name="Viljar Femoen", email="[email protected]" },
 ]

src/qm_sim/hamiltonian/spatial_hamiltonian.py

@@ -4,6 +4,7 @@
 
 """
 
+from collections.abc import Iterable
 from typing import Any, Callable
 
 import numpy as np
@@ -22,8 +23,8 @@
 class SpatialHamiltonian:
     def __init__(
         self,
-        N: tuple,
-        L: tuple,
+        N: tuple[int] | int,
+        L: tuple[float] | float,
         m: float | np.ndarray,
         spatial_scheme: str = "three-point",
         temporal_scheme: str = "leapfrog",
@@ -34,9 +35,9 @@ def __init__(
         """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
+        :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
@@ -83,14 +84,29 @@ def __init__(
             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:
@@ -100,6 +116,9 @@ def __init__(
         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]

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