From 9a1cad6e021073350b3d1c8365c769dc363653e6 Mon Sep 17 00:00:00 2001 From: Matheus Cordeiro Date: Wed, 25 Sep 2024 21:41:34 -0300 Subject: [PATCH] v1.4.0 - cython module --- requirements.txt | 3 +- setup.cfg | 3 +- setup.py | 3 +- src/fast_wave/setup_cython.py | 8 + src/fast_wave/wavefunction_cython.pyx | 458 ++++++++++++++++++ ...ion_arb_prec.py => wavefunction_mpmath.py} | 133 ++--- ...{wavefunction.py => wavefunction_numba.py} | 216 ++------- tests/test_wavefunction.py | 117 ----- tests/test_wavefunction_arb_prec.py | 44 -- tests/test_wavefunction_cython.py | 52 ++ tests/test_wavefunction_mpmath.py | 36 ++ tests/test_wavefunction_numba.py | 109 +++++ 12 files changed, 745 insertions(+), 437 deletions(-) create mode 100644 src/fast_wave/setup_cython.py create mode 100644 src/fast_wave/wavefunction_cython.pyx rename src/fast_wave/{wavefunction_arb_prec.py => wavefunction_mpmath.py} (63%) rename src/fast_wave/{wavefunction.py => wavefunction_numba.py} (64%) delete mode 100644 tests/test_wavefunction.py delete mode 100644 tests/test_wavefunction_arb_prec.py create mode 100644 tests/test_wavefunction_cython.py create mode 100644 tests/test_wavefunction_mpmath.py create mode 100644 tests/test_wavefunction_numba.py diff --git a/requirements.txt b/requirements.txt index dcc2645..9dfef1e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,5 @@ mpmath==1.3.0 numba==0.59.1 numpy==1.26.4 scipy==1.13.0 -sympy==1.12 \ No newline at end of file +sympy==1.12 +cython==3.0.10 \ No newline at end of file diff --git a/setup.cfg b/setup.cfg index 2ebe1b0..98ccf04 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = fast_wave -version = 1.3.11 +version = 1.4.0 description = Package for the calculation of the time-independent wavefunction. author = Matheus Gomes Cordeiro author_email = matheusgomescord@gmail.com @@ -17,6 +17,7 @@ numba==0.59.1 numpy==1.26.4 scipy==1.13.0 sympy==1.12 +cython==3.0.10 [test_requires] pytest = ^7.1.2 diff --git a/setup.py b/setup.py index 91ceb5c..37bea6b 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ long_description = fh.read() name = "fast_wave" -version = "1.3.11" +version = "1.4.0" description = "Package for the calculation of the time-independent wavefunction." author_email = "matheusgomescord@gmail.com" url = "https://github.com/pikachu123deimos/fast-wave" @@ -17,6 +17,7 @@ "numpy==1.26.4", "scipy==1.13.0", "sympy==1.12", + "cython==3.0.10", ] test_requires = [ diff --git a/src/fast_wave/setup_cython.py b/src/fast_wave/setup_cython.py new file mode 100644 index 0000000..204b5d5 --- /dev/null +++ b/src/fast_wave/setup_cython.py @@ -0,0 +1,8 @@ +from setuptools import setup +from Cython.Build import cythonize +import numpy as np + +setup( + ext_modules=cythonize("wavefunction_cython.pyx"), + include_dirs=[np.get_include()], +) diff --git a/src/fast_wave/wavefunction_cython.pyx b/src/fast_wave/wavefunction_cython.pyx new file mode 100644 index 0000000..68b0cec --- /dev/null +++ b/src/fast_wave/wavefunction_cython.pyx @@ -0,0 +1,458 @@ +# BSD 3-Clause License +# +# Copyright (c) 2024, Matheus Gomes Cordeiro +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +# cython: language_level=3 +#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION + +cimport cython +from cython cimport double, complex +cimport numpy as np +import numpy as np +from libc.math cimport exp, sqrt, pi +from cmath import exp as cexp, sqrt as csqrt, pi as cpi + +@cython.nogil +@cython.cfunc +@cython.locals(index=int, r0=double, r1=double, r2=double) +@cython.boundscheck(False) +cpdef double psi_n_single_fock_single_position(int n, double x): + + """ + Compute the wavefunction to a real scalar x using adapted recurrence relation. + + Parameters + ---------- + n : int + Quantum state number. + x : double + Position(s) at which to evaluate the wavefunction. + + + Returns + ------- + double + The evaluated wavefunction. + + Examples + -------- + ```python + >>> psi_n_single_fock_single_position(0, 1.0) + 0.45558067201133257 + >>> psi_n_single_fock_single_position(61, 1.0) + -0.2393049199171131 + ``` + + References + ---------- + - Pérez-Jordá, J. M. (2017). On the recursive solution of the quantum harmonic oscillator. *European Journal of Physics*, 39(1), + 015402. doi:10.1088/1361-6404/aa9584 + """ + + cdef np.ndarray[np.float64_t, ndim=1] n_coeffs + + r0 = 0.0 + r1 = (pi ** (-0.25)) * exp(-(x ** 2) / 2) + + for index in range(n): + if index == 0: + r2 = 2 * x * (r1 / sqrt(2 * (index + 1))) + r0 = r1 + r1 = r2 + else: + r2 = 2 * x * (r1 / sqrt(2 * (index + 1))) - sqrt(index / (index + 1)) * r0 + r0 = r1 + r1 = r2 + + return r1 + +@cython.nogil +@cython.cfunc +@cython.locals(index=int, r0=complex, r1=complex, r2=complex) +@cython.boundscheck(False) +cpdef double complex psi_n_single_fock_single_position_complex(int n, double complex x): + + """ + Compute the wavefunction to a complex scalar x using adapted recurrence relation. + + Parameters + ---------- + n : int + Quantum state number. + x : double + Position(s) at which to evaluate the wavefunction. + + + Returns + ------- + double + The evaluated wavefunction. + + Examples + -------- + ```python + >>> psi_n_single_fock_single_position_complex(0,1.0+2.0j) + (-1.4008797330262455-3.0609780602975003j) + >>> psi_n_single_fock_single_position_complex(61,1.0+2.0j) + (-511062135.47555304+131445997.75753704j) + ``` + + References + ---------- + - Pérez-Jordá, J. M. (2017). On the recursive solution of the quantum harmonic oscillator. *European Journal of Physics*, 39(1), + 015402. doi:10.1088/1361-6404/aa9584 + """ + + cdef np.ndarray[np.float64_t, ndim=1] n_coeffs + + r0 = 0.0 + 0.0j + r1 = (cpi ** (-0.25)) * cexp(-(x ** 2) / 2) + + for index in range(n): + if index == 0: + r2 = 2 * x * (r1 / csqrt(2 * (index + 1))) + r0 = r1 + r1 = r2 + else: + r2 = 2 * x * (r1 / csqrt(2 * (index + 1))) - csqrt(index / (index + 1)) * r0 + r0 = r1 + r1 = r2 + + return r1 + +@cython.nogil +@cython.cfunc +@cython.locals(x_size=np.npy_intp, j=int, i=int, k=int, temp1=double, temp2=double) +@cython.boundscheck(False) +cpdef np.ndarray[np.float64_t, ndim=1] psi_n_single_fock_multiple_position(int n, np.ndarray[np.float64_t, ndim=1] x): + + """ + Compute the wavefunction to a real vector x using adapted recurrence relation. + + Parameters + ---------- + n : int + Quantum state number. + x : np.ndarray[np.float64_t] + Position(s) at which to evaluate the wavefunction. + + + Returns + ------- + np.ndarray[np.float64_t] + The evaluated wavefunction. + + Examples + -------- + ```python + >>> psi_n_single_fock_multiple_position(0, np.array([1.0, 2.0])) + array([0.45558067, 0.10165379]) + >>> psi_n_single_fock_multiple_position(61, np.array([1.0, 2.0])) + array([-0.23930492, -0.01677378]) + ``` + + References + ---------- + - Pérez-Jordá, J. M. (2017). On the recursive solution of the quantum harmonic oscillator. *European Journal of Physics*, 39(1), + 015402. doi:10.1088/1361-6404/aa9584 + """ + + x_size = x.shape[0] + cdef np.ndarray[np.float64_t, ndim=2] result = np.zeros((n + 1, x_size), dtype=np.float64) + + for j in range(x_size): + result[0, j] = (pi ** (-0.25)) * exp(-(x[j] ** 2) / 2) + + for i in range(n): + temp1 = sqrt(2 * (i + 1)) + temp2 = sqrt(i / (i + 1)) + if(i == 0): + for k in range(x_size): + result[i + 1, k] = 2 * x[k] * (result[i, k] / temp1) + else: + for k in range(x_size): + result[i + 1, k] = 2 * x[k] * (result[i, k] / temp1) - temp2 * result[i - 1, k] + + return result[-1, :] + +@cython.nogil +@cython.cfunc +@cython.locals(x_size=np.npy_intp, j=int, i=int, k=int, temp1=complex, temp2=complex) +@cython.boundscheck(False) +cpdef np.ndarray[np.complex128_t, ndim=1] psi_n_single_fock_multiple_position_complex(int n, np.ndarray[np.complex128_t, ndim=1] x): + + """ + Compute the wavefunction to a complex vector x using adapted recurrence relation. + + Parameters + ---------- + n : int + Quantum state number. + x : np.ndarray[np.complex128_t] + Position(s) at which to evaluate the wavefunction. + + + Returns + ------- + np.ndarray[np.complex128_t] + The evaluated wavefunction. + + Examples + -------- + ```python + >>> psi_n_single_fock_multiple_position_complex(0, np.array([1.0 + 1.0j, 2.0 + 2.0j])) + array([ 0.40583486-0.63205035j, -0.49096842+0.56845369j]) + >>> psi_n_single_fock_multiple_position_complex(61, np.array([1.0 + 1.0j, 2.0 + 2.0j])) + array([-7.56548941e+03+9.21498621e+02j, -1.64189542e+08-3.70892077e+08j]) + ``` + + References + ---------- + - Pérez-Jordá, J. M. (2017). On the recursive solution of the quantum harmonic oscillator. *European Journal of Physics*, 39(1), + 015402. doi:10.1088/1361-6404/aa9584 + """ + + x_size = x.shape[0] + cdef np.ndarray[np.complex128_t, ndim=2] result = np.zeros((n + 1, x_size), dtype=np.complex128) + + for j in range(x_size): + result[0, j] = (cpi ** (-0.25)) * cexp(-(x[j] ** 2) / 2) + + for i in range(n): + temp1 = csqrt(2 * (i + 1)) + temp2 = csqrt(i / (i + 1)) + if(i == 0): + for k in range(x_size): + result[i + 1, k] = 2 * x[k] * (result[i, k] / temp1) + else: + for k in range(x_size): + result[i + 1, k] = 2 * x[k] * (result[i, k] / temp1) - temp2 * result[i - 1, k] + + return result[-1, :] + + +@cython.nogil +@cython.cfunc +@cython.locals(index=int) +@cython.boundscheck(False) +cpdef np.ndarray[np.float64_t, ndim=1] psi_n_multiple_fock_single_position(int n, double x): + + """ + Compute the wavefunction to a real scalar x to all fock states until n using adapted recurrence relation. + + Parameters + ---------- + n : int + Quantum state number. + x : double + Position(s) at which to evaluate the wavefunction. + + + Returns + ------- + np.ndarray[np.float64_t] + The evaluated wavefunction. + + Examples + -------- + ```python + >>> psi_n_multiple_fock_single_position(0, 1.0) + array([0.45558067, 0.64428837]) + ``` + + References + ---------- + - Pérez-Jordá, J. M. (2017). On the recursive solution of the quantum harmonic oscillator. *European Journal of Physics*, 39(1), + 015402. doi:10.1088/1361-6404/aa9584 + """ + + cdef np.ndarray[np.float64_t, ndim=1] result = np.zeros((n + 1), dtype=np.float64) + result[0] = (pi ** (-0.25)) * exp(-(x ** 2) / 2) + + for index in range(n): + if(index == 0): + result[index + 1] = 2 * x * (result[index] / sqrt(2 * (index + 1))) + else: + result[index + 1] = 2 * x * (result[index] / sqrt(2 * (index + 1))) - sqrt(index / (index + 1)) * result[index - 1] + + return result + +@cython.nogil +@cython.cfunc +@cython.locals(index=int) +@cython.boundscheck(False) +cpdef np.ndarray[np.complex128_t, ndim=1] psi_n_multiple_fock_single_position_complex(int n, double complex x): + + """ + Compute the wavefunction to a complex scalar x to all fock states until n using adapted recurrence relation. + + Parameters + ---------- + n : int + Quantum state number. + x : double complex + Position(s) at which to evaluate the wavefunction. + + + Returns + ------- + np.ndarray[np.complex128_t] + The evaluated wavefunction. + + Examples + -------- + ```python + >>> psi_n_multiple_fock_single_position_complex(0, 1.0 +2.0j) + array([-1.40087973-3.06097806j, 6.67661026-8.29116292j]) + ``` + + References + ---------- + - Pérez-Jordá, J. M. (2017). On the recursive solution of the quantum harmonic oscillator. *European Journal of Physics*, 39(1), + 015402. doi:10.1088/1361-6404/aa9584 + """ + + cdef np.ndarray[np.complex128_t, ndim=1] result = np.zeros((n + 1), dtype=np.complex128) + result[0] = (cpi ** (-0.25)) * cexp(-(x ** 2) / 2) + + for index in range(n): + if(index == 0): + result[index + 1] = 2 * x * (result[index] / csqrt(2 * (index + 1))) + else: + result[index + 1] = 2 * x * (result[index] / csqrt(2 * (index + 1))) - csqrt(index / (index + 1)) * result[index - 1] + + return result + +@cython.nogil +@cython.cfunc +@cython.locals(x_size=np.npy_intp, j=int, i=int, k=int, temp1=double, temp2=double) +@cython.boundscheck(False) +cpdef np.ndarray[np.float64_t, ndim=2] psi_n_multiple_fock_multiple_position(int n, np.ndarray[np.float64_t, ndim=1] x): + + """ + Compute the wavefunction to a real vector x to all fock states until n using adapted recurrence relation. + + Parameters + ---------- + n : int + Quantum state number. + x : np.ndarray[np.float64_t] + Position(s) at which to evaluate the wavefunction. + + + Returns + ------- + np.ndarray[np.ndarray[np.float64_t]] + The evaluated wavefunction. + + Examples + -------- + ```python + >>> psi_n_multiple_fock_multiple_position(1, np.array([1.0, 2.0])) + array([[0.45558067, 0.10165379], + [0.64428837, 0.28752033]]) + ``` + + References + ---------- + - Pérez-Jordá, J. M. (2017). On the recursive solution of the quantum harmonic oscillator. *European Journal of Physics*, 39(1), + 015402. doi:10.1088/1361-6404/aa9584 + """ + + x_size = x.shape[0] + cdef np.ndarray[np.float64_t, ndim=2] result = np.zeros((n + 1, x_size), dtype=np.float64) + + for j in range(x_size): + result[0, j] = (pi ** (-0.25)) * exp(-(x[j] ** 2) / 2) + + for i in range(n): + temp1 = sqrt(2 * (i + 1)) + temp2 = sqrt(i / (i + 1)) + if(i == 0): + for k in range(x_size): + result[i + 1, k] = 2 * x[k] * (result[i, k] / temp1) + else: + for k in range(x_size): + result[i + 1, k] = 2 * x[k] * (result[i, k] / temp1) - temp2 * result[i - 1, k] + + return result + +@cython.nogil +@cython.cfunc +@cython.locals(x_size=np.npy_intp, j=int, i=int, k=int, temp1=complex, temp2=complex) +@cython.boundscheck(False) +cpdef np.ndarray[np.complex128_t, ndim=2] psi_n_multiple_fock_multiple_position_complex(int n, np.ndarray[np.complex128_t, ndim=1] x): + + """ + Compute the wavefunction to a complex vector x to all fock states until n using adapted recurrence relation. + + Parameters + ---------- + n : int + Quantum state number. + x : np.ndarray[np.float64_t] + Position(s) at which to evaluate the wavefunction. + + + Returns + ------- + np.ndarray[np.ndarray[np.float64_t]] + The evaluated wavefunction. + + Examples + -------- + ```python + >>> psi_n_multiple_fock_multiple_position_complex(1,np.array([1.0 + 1.0j, 2.0 + 2.0j])) + array([[ 0.40583486-0.63205035j, -0.49096842+0.56845369j], + [ 1.46779135-0.31991701j, -2.99649822+0.21916143j]]) + ``` + + References + ---------- + - Pérez-Jordá, J. M. (2017). On the recursive solution of the quantum harmonic oscillator. *European Journal of Physics*, 39(1), + 015402. doi:10.1088/1361-6404/aa9584 + """ + + x_size = x.shape[0] + cdef np.ndarray[np.complex128_t, ndim=2] result = np.zeros((n + 1, x_size), dtype=np.complex128) + + for j in range(x_size): + result[0, j] = (cpi ** (-0.25)) * cexp(-(x[j] ** 2) / 2) + + for i in range(n): + temp1 = csqrt(2 * (i + 1)) + temp2 = csqrt(i / (i + 1)) + if(i == 0): + for k in range(x_size): + result[i + 1, k] = 2 * x[k] * (result[i, k] / temp1) + else: + for k in range(x_size): + result[i + 1, k] = 2 * x[k] * (result[i, k] / temp1) - temp2 * result[i - 1, k] + + return result \ No newline at end of file diff --git a/src/fast_wave/wavefunction_arb_prec.py b/src/fast_wave/wavefunction_mpmath.py similarity index 63% rename from src/fast_wave/wavefunction_arb_prec.py rename to src/fast_wave/wavefunction_mpmath.py index 36e757f..504661d 100644 --- a/src/fast_wave/wavefunction_arb_prec.py +++ b/src/fast_wave/wavefunction_mpmath.py @@ -28,16 +28,12 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -from functools import lru_cache - from mpmath import mp, matrix import mpmath import numpy as np - - -def wavefunction_smod_arb_prec(n: np.uint64, x: np.float64, prec: np.uint64) -> mpmath.ctx_mp_python.mpf: +def psi_n_single_fock_single_position(n: np.uint64, x: np.float64, prec: np.uint64) -> mpmath.ctx_mp_python.mpf: """ Calculates the nth wavefunction to an real scalar x with arbitrary precision using mpmath. @@ -60,9 +56,9 @@ def wavefunction_smod_arb_prec(n: np.uint64, x: np.float64, prec: np.uint64) -> Examples -------- ```python - >>> wavefunction_smod_arb_prec(0,1.0,60) + >>> psi_n_single_fock_single_position(0,1.0,60) mpf('0.45558067201133253483370525689785138607662639040929439687915331') - >>> wavefunction_smod_arb_prec(61,1.0,60) + >>> psi_n_single_fock_single_position(61,1.0,60) mpf('-0.239304919917113097789996116536717211865611421191819349290628243') ``` @@ -80,7 +76,7 @@ def wavefunction_smod_arb_prec(n: np.uint64, x: np.float64, prec: np.uint64) -> -def c_wavefunction_smod_arb_prec(n: np.uint64, x: np.complex128, prec: np.uint64) -> mpmath.ctx_mp_python.mpc: +def psi_n_single_fock_single_position_complex(n: np.uint64, x: np.complex128, prec: np.uint64) -> mpmath.ctx_mp_python.mpc: """ Calculates the nth wavefunction to a complex scalar x with arbitrary precision using mpmath. @@ -103,9 +99,9 @@ def c_wavefunction_smod_arb_prec(n: np.uint64, x: np.complex128, prec: np.uint64 Examples -------- ```python - >>> c_wavefunction_smod_arb_prec(0,1.0+2.0j,60) + >>> psi_n_single_fock_single_position_complex(0,1.0+2.0j,60) mpc(real='-1.40087973302624535996319358379185603705205815719366827159881527', imag='-3.06097806029750039193292973729038840279841978760336147713769087') - >>> c_wavefunction_smod_arb_prec(61,1.0+2.0j,60) + >>> psi_n_single_fock_single_position_complex(61,1.0+2.0j,60) mpc(real='-511062135.475553070892329856229109412939170026007243421420322129', imag='131445997.757536932748911867174534983962121585813389430606204944') ``` @@ -124,7 +120,7 @@ def c_wavefunction_smod_arb_prec(n: np.uint64, x: np.complex128, prec: np.uint64 -def wavefunction_smmd_arb_prec(n: np.uint64, X: np.ndarray[np.float64], prec: np.uint64) -> mpmath.matrices.matrices._matrix: +def psi_n_single_fock_multiple_position(n: np.uint64, X: np.ndarray[np.float64], prec: np.uint64) -> mpmath.matrices.matrices._matrix: """ @@ -148,10 +144,10 @@ def wavefunction_smmd_arb_prec(n: np.uint64, X: np.ndarray[np.float64], prec: np Examples -------- ```python - >>> wavefunction_smmd_arb_prec(0,np.array([1.0,2.0]),20) + >>> psi_n_single_fock_multiple_position(0,np.array([1.0,2.0]),20) matrix( [['0.45558067201133253483', '0.10165378830641791152']]) - >>> wavefunction_smmd_arb_prec(61,np.array([1.0,2.0]),20) + >>> psi_n_single_fock_multiple_position(61,np.array([1.0,2.0]),20) matrix( [['-0.23930491991711309779', '-0.016773782204892582343']]) ``` @@ -162,14 +158,14 @@ def wavefunction_smmd_arb_prec(n: np.uint64, X: np.ndarray[np.float64], prec: np """ - return matrix([wavefunction_smod_arb_prec(n, x, prec) for x in X]).T + return matrix([psi_n_single_fock_single_position(n, x, prec) for x in X]).T -def c_wavefunction_smmd_arb_prec(n: np.uint64, X: np.ndarray[np.complex128], prec: np.uint64) -> mpmath.matrices.matrices._matrix: +def psi_n_single_fock_multiple_position_complex(n: np.uint64, X: np.ndarray[np.complex128], prec: np.uint64) -> mpmath.matrices.matrices._matrix: """ Calculates the nth wavefunction to a complex vector x with arbitrary precision using mpmath. @@ -192,10 +188,10 @@ def c_wavefunction_smmd_arb_prec(n: np.uint64, X: np.ndarray[np.complex128], pre Examples -------- ```python - >>> c_wavefunction_smmd_arb_prec(0,np.array([1.0 + 1.0j, 2.0 + 2.0j]),20) + >>> psi_n_single_fock_multiple_position_complex(0,np.array([1.0 + 1.0j, 2.0 + 2.0j]),20) matrix( [[mpc(real='0.40583486367087033308603', imag='-0.63205035161528260798606'), mpc(real='-0.49096842060721693717778', imag='0.56845368634059468652777')]]) - >>> c_wavefunction_smmd_arb_prec(61,np.array([1.0 + 1.0j, 2.0 + 2.0j]),20) + >>> psi_n_single_fock_multiple_position_complex(61,np.array([1.0 + 1.0j, 2.0 + 2.0j]),20) matrix( [[mpc(real='-7565.4894098859360141926', imag='921.4986211518276840917'), mpc(real='-164189541.53192908120809', imag='-370892077.23796911662203')]]) ``` @@ -205,16 +201,16 @@ def c_wavefunction_smmd_arb_prec(n: np.uint64, X: np.ndarray[np.complex128], pre - The mpmath development team. (2023). mpmath: a Python library for arbitrary-precision floating-point arithmetic (version 1.3.0). Retrieved from http://mpmath.org/ """ - return matrix([c_wavefunction_smod_arb_prec(n, x, prec) for x in X]).T + return matrix([psi_n_single_fock_single_position_complex(n, x, prec) for x in X]).T -def wavefunction_mmod_arb_prec(n: np.uint64, x: np.float64, prec: np.uint64) -> mpmath.matrices.matrices._matrix: +def psi_n_multiple_fock_single_position(n: np.uint64, x: np.float64, prec: np.uint64) -> mpmath.matrices.matrices._matrix: """ - Determines the wavefunction for a real scalar x up to mode n, employing mpmath for arbitrary-precision calculations. + Determines the wavefunction for a real scalar x to all fock states until n, employing mpmath for arbitrary-precision calculations. Parameters: ---------- @@ -234,7 +230,7 @@ def wavefunction_mmod_arb_prec(n: np.uint64, x: np.float64, prec: np.uint64) -> Examples -------- ```python - >>> wavefunction_mmod_arb_prec(1,1.0,60) + >>> psi_n_multiple_fock_single_position(1,1.0,60) matrix( [['0.455580672011332534833705256897851386076626390409294396879153', '0.644288365113475181510837645362740498634994248687269122618738']]) ``` @@ -245,16 +241,16 @@ def wavefunction_mmod_arb_prec(n: np.uint64, x: np.float64, prec: np.uint64) -> """ - return matrix([wavefunction_smod_arb_prec(i, x, prec) for i in range(n+1)]).T + return matrix([psi_n_single_fock_single_position(i, x, prec) for i in range(n+1)]).T -def c_wavefunction_mmod_arb_prec(n: np.uint64, x: np.complex128, prec: np.uint64) -> mpmath.matrices.matrices._matrix: +def psi_n_multiple_fock_single_position_complex(n: np.uint64, x: np.complex128, prec: np.uint64) -> mpmath.matrices.matrices._matrix: """ - Determines the wavefunction for a complex scalar x up to mode n, employing mpmath for arbitrary-precision calculations. + Determines the wavefunction for a complex scalar x to all fock states until n, employing mpmath for arbitrary-precision calculations. Parameters: ---------- @@ -284,16 +280,16 @@ def c_wavefunction_mmod_arb_prec(n: np.uint64, x: np.complex128, prec: np.uint64 - The mpmath development team. (2023). mpmath: a Python library for arbitrary-precision floating-point arithmetic (version 1.3.0). Retrieved from http://mpmath.org/ """ - return matrix([c_wavefunction_smod_arb_prec(i, x, prec) for i in range(n+1)]).T + return matrix([psi_n_single_fock_single_position_complex(i, x, prec) for i in range(n+1)]).T -def wavefunction_mmmd_arb_prec(n: np.uint64, X: np.ndarray[np.float64], prec: np.uint64) -> mpmath.matrices.matrices._matrix: +def psi_n_multiple_fock_multiple_position(n: np.uint64, X: np.ndarray[np.float64], prec: np.uint64) -> mpmath.matrices.matrices._matrix: """ - Determines the wavefunction for a real vector x up to mode n, employing mpmath for arbitrary-precision calculations. + Determines the wavefunction for a real vector x to all fock states until n, employing mpmath for arbitrary-precision calculations. Parameters: ---------- @@ -324,16 +320,16 @@ def wavefunction_mmmd_arb_prec(n: np.uint64, X: np.ndarray[np.float64], prec: np - The mpmath development team. (2023). mpmath: a Python library for arbitrary-precision floating-point arithmetic (version 1.3.0). Retrieved from http://mpmath.org/ """ - return matrix([[wavefunction_smod_arb_prec(i, x, prec) for x in X] for i in range(n+1)]) + return matrix([[psi_n_single_fock_single_position(i, x, prec) for x in X] for i in range(n+1)]) -def c_wavefunction_mmmd_arb_prec(n: np.uint64, X: np.ndarray[np.complex128], prec: np.uint64) -> mpmath.matrices.matrices._matrix: +def psi_n_multiple_fock_multiple_position_complex(n: np.uint64, X: np.ndarray[np.complex128], prec: np.uint64) -> mpmath.matrices.matrices._matrix: """ - Determines the wavefunction for a complex vector x up to mode n, employing mpmath for arbitrary-precision calculations. + Determines the wavefunction for a complex vector x to all fock states until n, employing mpmath for arbitrary-precision calculations. Parameters: ---------- @@ -363,79 +359,4 @@ def c_wavefunction_mmmd_arb_prec(n: np.uint64, X: np.ndarray[np.complex128], pre - The mpmath development team. (2023). mpmath: a Python library for arbitrary-precision floating-point arithmetic (version 1.3.0). Retrieved from http://mpmath.org/ """ - return matrix([[c_wavefunction_smod_arb_prec(i, x, prec) for x in X] for i in range(n+1)]) - - -def wavefunction_arb_prec(s_mode: bool = True, o_dimensional: bool = True, complex_bool: bool = False, cache: bool = False, cache_size: np.uint64 = 128): - - """ - Computes the wavefunction of a quantum harmonic oscillator with arbitrary precision. This function dispatches to different implementations of the wavefunction - depending on the specified parameters. - - Parameters: - ---------- - s_mode : bool - Whether to use the single-mode (True) or the multi-mode (False). Defaults to True. - o_dimensional : bool - Whether to evaluate the wavefunction for a single coordinate (True) or multiple coordinates (False). Defaults to True. - complex_boll : bool - Whether to return the complex wavefunction (True) or the real wavefunction (False). Defaults to False. - cache : bool - Whether to cache the results of the wavefunction computation using least-recently-used caching. Defaults to False. - cache_size : np.uint64 - The maximum size of the cache to use. Defaults to 128. - - Returns: - ------- - function: - A dispatcher object that resolves to the appropriate wavefunction implementation based on the input parameters. - """ - - - - if(s_mode): - if(o_dimensional): - if(not(complex_bool)): - if(not(cache)): - return wavefunction_smod_arb_prec - else: - return lru_cache(maxsize=cache_size)(wavefunction_smod_arb_prec) - else: - if(not(cache)): - return c_wavefunction_smod_arb_prec - else: - return lru_cache(maxsize=cache_size)(c_wavefunction_smod_arb_prec) - else: - if(not(complex_bool)): - if(not(cache)): - return wavefunction_smmd_arb_prec - else: - return lru_cache(maxsize=cache_size)(wavefunction_smmd_arb_prec) - else: - if(not(cache)): - return c_wavefunction_smmd_arb_prec - else: - return lru_cache(maxsize=cache_size)(c_wavefunction_smmd_arb_prec) - else: - if(o_dimensional): - if(not(complex_bool)): - if(not(cache)): - return wavefunction_mmod_arb_prec - else: - return lru_cache(maxsize=cache_size)(wavefunction_mmod_arb_prec) - else: - if(not(cache)): - return c_wavefunction_mmod_arb_prec - else: - return lru_cache(maxsize=cache_size)(c_wavefunction_mmod_arb_prec) - else: - if(not(complex_bool)): - if(not(cache)): - return wavefunction_mmmd_arb_prec - else: - return lru_cache(maxsize=cache_size)(wavefunction_mmmd_arb_prec) - else: - if(not(cache)): - return c_wavefunction_mmmd_arb_prec - else: - return lru_cache(maxsize=cache_size)(c_wavefunction_mmmd_arb_prec) \ No newline at end of file + return matrix([[psi_n_single_fock_single_position_complex(i, x, prec) for x in X] for i in range(n+1)]) diff --git a/src/fast_wave/wavefunction.py b/src/fast_wave/wavefunction_numba.py similarity index 64% rename from src/fast_wave/wavefunction.py rename to src/fast_wave/wavefunction_numba.py index 5567dd1..ab58039 100644 --- a/src/fast_wave/wavefunction.py +++ b/src/fast_wave/wavefunction_numba.py @@ -28,10 +28,7 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -import os -import pickle import math -from functools import lru_cache from sympy import symbols, diff, exp, Poly import numpy as np @@ -117,44 +114,13 @@ def create_normalized_hermite_coefficients_matrix(n_max: np.uint64) -> np.ndarra return C_s -def delete_normalized_hermite_coefficients_matrix(): - - """ - Delete the 'C_s_matrix.pickle' file from the local directory. - - Parameters - ---------- - None - This function doesn't receive anything. - - Returns - ------- - None - This function doesn't return anything. It prints a success message if the file is deleted, - or an error message if the file was not found or there was an issue with deletion. - - """ - - package_dir = os.path.dirname(__file__) - matrix_filename = 'C_s_matrix.pickle' - matrix_path = os.path.join(package_dir, matrix_filename) - - if os.path.exists(matrix_path): - try: - os.remove(matrix_path) - print(f"File {matrix_path} successfully removed.") - except Exception as e: - print(f"Error while trying to remove the file: {e}") - else: - print(f"File {matrix_path} not found.") - @nb.jit(nopython=True, looplift=True, nogil=True, boundscheck=False, cache=True) -def wavefunction_smod(n: np.uint64, x:np.float64, more_fast:bool = True) -> np.float64: +def psi_n_single_fock_single_position(n: np.uint64, x:np.float64, more_fast:bool = True) -> np.float64: """ Compute the wavefunction to a real scalar x using a pre-computed matrix of normalized Hermite polynomial coefficients until n=60 and - then use the adapted recursion relation for multidimensional M-mode wavefunction for higher orders. + then use the adapted recurrence relation for higher orders. Parameters ---------- @@ -175,9 +141,9 @@ def wavefunction_smod(n: np.uint64, x:np.float64, more_fast:bool = True) -> np.f Examples -------- ```python - >>> wavefunction_smod(0, 1.0) + >>> psi_n_single_fock_single_position(0, 1.0) 0.45558067201133257 - >>> wavefunction_smod(61, 1.0) + >>> psi_n_single_fock_single_position(61, 1.0) -0.2393049199171131 ``` @@ -209,11 +175,11 @@ def wavefunction_smod(n: np.uint64, x:np.float64, more_fast:bool = True) -> np.f @nb.jit(nopython=True, looplift=True, nogil=True, boundscheck=False, cache=True) -def c_wavefunction_smod(n: np.uint64, x: np.complex128, more_fast:bool = True) -> np.complex128: +def psi_n_single_fock_single_position_complex(n: np.uint64, x: np.complex128, more_fast:bool = True) -> np.complex128: """ Compute the wavefunction to a complex scalar x using a pre-computed matrix of normalized Hermite polynomial coefficients until n=60 and - then use the adapted recursion relation for multidimensional M-mode wavefunction for higher orders. + then use the adapted recurrence relation for higher orders. Parameters ---------- @@ -234,9 +200,9 @@ def c_wavefunction_smod(n: np.uint64, x: np.complex128, more_fast:bool = True) - Examples -------- ```python - >>> c_wavefunction_smod(0,1.0+2.0j) + >>> psi_n_single_fock_single_position_complex(0,1.0+2.0j) (-1.4008797330262455-3.0609780602975003j) - >>> c_wavefunction_smod(61,1.0+2.0j) + >>> psi_n_single_fock_single_position_complex(61,1.0+2.0j) (-511062135.47555304+131445997.75753704j) ``` @@ -268,11 +234,11 @@ def c_wavefunction_smod(n: np.uint64, x: np.complex128, more_fast:bool = True) - @nb.jit(nopython=True, looplift=True,nogil=True, boundscheck=False, cache=True) -def wavefunction_smmd(n: np.uint64, x: np.ndarray[np.float64], more_fast: bool = True) -> np.ndarray[np.float64]: +def psi_n_single_fock_multiple_position(n: np.uint64, x: np.ndarray[np.float64], more_fast: bool = True) -> np.ndarray[np.float64]: """ Compute the wavefunction to a real vector x using a pre-computed matrix of normalized Hermite polynomial coefficients until n=60 and - then use the adapted recursion relation for multidimensional M-mode wavefunction for higher orders. + then use the adapted recurrence relation for higher orders. Parameters ---------- @@ -294,9 +260,9 @@ def wavefunction_smmd(n: np.uint64, x: np.ndarray[np.float64], more_fast: bool = Examples -------- ```python - >>> wavefunction_smmd(0,np.array([1.0, 2.0])) + >>> psi_n_single_fock_multiple_position(0,np.array([1.0, 2.0])) array([0.45558067, 0.10165379]) - >>> wavefunction_smmd(61,np.array([1.0, 2.0])) + >>> psi_n_single_fock_multiple_position(61,np.array([1.0, 2.0])) array([-0.23930492, -0.01677378]) ``` @@ -329,11 +295,11 @@ def wavefunction_smmd(n: np.uint64, x: np.ndarray[np.float64], more_fast: bool = @nb.jit(nopython=True, looplift=True,nogil=True, boundscheck=False, cache=True) -def c_wavefunction_smmd(n: np.uint64, x: np.ndarray[np.complex128], more_fast: bool = True) -> np.ndarray[np.complex128]: +def psi_n_single_fock_multiple_position_complex(n: np.uint64, x: np.ndarray[np.complex128], more_fast: bool = True) -> np.ndarray[np.complex128]: """ Compute the wavefunction to a complex vector x using a pre-computed matrix of normalized Hermite polynomial coefficients until n=60 and - then use the adapted recursion relation for multidimensional M-mode wavefunction for higher orders. + then use the adapted recurrence relation for higher orders. Parameters ---------- @@ -355,9 +321,9 @@ def c_wavefunction_smmd(n: np.uint64, x: np.ndarray[np.complex128], more_fast: b Examples -------- ```python - >>> c_wavefunction_smmd(0,np.array([1.0 + 1.0j, 2.0 + 2.0j])) + >>> psi_n_single_fock_multiple_position_complex(0,np.array([1.0 + 1.0j, 2.0 + 2.0j])) array([ 0.40583486-0.63205035j, -0.49096842+0.56845369j]) - >>> c_wavefunction_smmd(61,np.array([1.0 + 1.0j, 2.0 + 2.0j])) + >>> psi_n_single_fock_multiple_position_complex(61,np.array([1.0 + 1.0j, 2.0 + 2.0j])) array([-7.56548941e+03+9.21498621e+02j, -1.64189542e+08-3.70892077e+08j]) ``` @@ -389,10 +355,10 @@ def c_wavefunction_smmd(n: np.uint64, x: np.ndarray[np.complex128], more_fast: b @nb.jit(nopython=True, looplift=True,nogil=True, boundscheck=False, cache=True) -def wavefunction_mmod(n: np.uint64, x:np.float64) -> np.ndarray[np.float64]: +def psi_n_multiple_fock_single_position(n: np.uint64, x:np.float64) -> np.ndarray[np.float64]: """ - Compute the wavefunction to a real scalar x to all modes until the mode n using the recursion relation for multidimensional M-mode wavefunction. + Compute the wavefunction to a real scalar x to all fock states until n using the recurrence relation. Parameters ---------- @@ -410,7 +376,7 @@ def wavefunction_mmod(n: np.uint64, x:np.float64) -> np.ndarray[np.float64]: Examples -------- ```python - >>> wavefunction_mmod(1,1.0) + >>> psi_n_multiple_fock_single_position(1,1.0) array([0.45558067, 0.64428837]) ``` @@ -430,10 +396,10 @@ def wavefunction_mmod(n: np.uint64, x:np.float64) -> np.ndarray[np.float64]: @nb.jit(nopython=True, looplift=True,nogil=True, boundscheck=False, cache=True) -def c_wavefunction_mmod(n: np.uint64, x: np.complex128) -> np.ndarray[np.complex128]: +def psi_n_multiple_fock_single_position_complex(n: np.uint64, x: np.complex128) -> np.ndarray[np.complex128]: """ - Compute the wavefunction to a complex scalar x to all modes until the mode n using the recursion relation for multidimensional M-mode wavefunction. + Compute the wavefunction to a complex scalar x to all fock states until n using the recurrence relation. Parameters ---------- @@ -451,7 +417,7 @@ def c_wavefunction_mmod(n: np.uint64, x: np.complex128) -> np.ndarray[np.complex Examples -------- ```python - >>> c_wavefunction_mmod(1,1.0 +2.0j) + >>> psi_n_multiple_fock_single_position_complex(1,1.0 +2.0j) array([-1.40087973-3.06097806j, 6.67661026-8.29116292j]) ``` @@ -472,10 +438,10 @@ def c_wavefunction_mmod(n: np.uint64, x: np.complex128) -> np.ndarray[np.complex @nb.jit(nopython=True, looplift=True,nogil=True, boundscheck=False, cache=True) -def wavefunction_mmmd(n: np.uint64, x: np.ndarray[np.float64]) -> np.ndarray[np.ndarray[np.float64]]: +def psi_n_multiple_fock_multiple_position(n: np.uint64, x: np.ndarray[np.float64]) -> np.ndarray[np.ndarray[np.float64]]: """ - Compute the wavefunction to a real vector x to all modes until the mode n using the recursion relation for multidimensional M-mode wavefunction. + Compute the wavefunction to a real vector x to all fock states until n using the recurrence relation. Parameters ---------- @@ -493,7 +459,7 @@ def wavefunction_mmmd(n: np.uint64, x: np.ndarray[np.float64]) -> np.ndarray[np. Examples -------- ```python - >>> wavefunction_mmmd(1,np.array([1.0, 2.0])) + >>> psi_n_multiple_fock_multiple_position(1,np.array([1.0, 2.0])) array([[0.45558067, 0.10165379], [0.64428837, 0.28752033]]) ``` @@ -515,10 +481,10 @@ def wavefunction_mmmd(n: np.uint64, x: np.ndarray[np.float64]) -> np.ndarray[np. @nb.jit(nopython=True, looplift=True,nogil=True, boundscheck=False, cache=True) -def c_wavefunction_mmmd(n: np.uint64, x: np.ndarray[np.complex128]) -> np.ndarray[np.ndarray[np.float64]]: +def psi_n_multiple_fock_multiple_position_complex(n: np.uint64, x: np.ndarray[np.complex128]) -> np.ndarray[np.ndarray[np.float64]]: """ - Compute the wavefunction to a complex vector x to all modes until the mode n using the recursion relation for multidimensional M-mode wavefunction. + Compute the wavefunction to a complex vector x to all fock states until n using the recurrence relation. Parameters ---------- @@ -536,7 +502,7 @@ def c_wavefunction_mmmd(n: np.uint64, x: np.ndarray[np.complex128]) -> np.ndarra Examples -------- ```python - >>> c_wavefunction_mmmd(1,np.array([1.0 + 1.0j, 2.0 + 2.0j])) + >>> psi_n_multiple_fock_multiple_position_complex(1,np.array([1.0 + 1.0j, 2.0 + 2.0j])) array([[ 0.40583486-0.63205035j, -0.49096842+0.56845369j], [ 1.46779135-0.31991701j, -2.99649822+0.21916143j]]) ``` @@ -555,83 +521,9 @@ def c_wavefunction_mmmd(n: np.uint64, x: np.ndarray[np.complex128]) -> np.ndarra result[index+1] = 2*x*(result[index]/np.sqrt(2*(index+1))) - np.sqrt(index/(index+1)) * result[index-1] return result - -def wavefunction(s_mode: bool = True, o_dimensional: bool = True, complex_bool: bool = False, cache: bool = False, cache_size: np.uint64 = 128) -> nb.core.registry.CPUDispatcher: - - """ - Computes the wavefunction of a quantum harmonic oscillator .This function dispatches to different implementations of the wavefunction depending on the specified parameters. - - Parameters: - ---------- - s_mode : bool - Whether to use the single-mode recursion (True) or the multi-mode recursion (False). Defaults to True. - o_dimensional : bool - Whether to evaluate the wavefunction for a single coordinate (True) or multiple coordinates (False). Defaults to True. - complex_boll : bool - Whether to return the complex wavefunction (True) or the real wavefunction (False). Defaults to False. - cache : bool - Whether to cache the results of the wavefunction computation using least-recently-used caching. Defaults to False. - cache_size : np.uint64 - The maximum size of the cache to use. Defaults to 128. - - Returns: - ------- - nb.core.registry.CPUDispatcher: - A dispatcher object that resolves to the appropriate wavefunction implementation based on the input parameters. - """ - - if(s_mode): - if(o_dimensional): - if(not(complex_bool)): - if(not(cache)): - return wavefunction_smod - else: - return lru_cache(maxsize=cache_size)(wavefunction_smod) - else: - if(not(cache)): - return c_wavefunction_smod - else: - return lru_cache(maxsize=cache_size)(c_wavefunction_smod) - else: - if(not(complex_bool)): - if(not(cache)): - return wavefunction_smmd - else: - return lru_cache(maxsize=cache_size)(wavefunction_smmd) - else: - if(not(cache)): - return c_wavefunction_smmd - else: - return lru_cache(maxsize=cache_size)(c_wavefunction_smmd) - else: - if(o_dimensional): - if(not(complex_bool)): - if(not(cache)): - return wavefunction_mmod - else: - return lru_cache(maxsize=cache_size)(wavefunction_mmod) - else: - if(not(cache)): - return c_wavefunction_mmod - else: - return lru_cache(maxsize=cache_size)(c_wavefunction_mmod) - else: - if(not(complex_bool)): - if(not(cache)): - return wavefunction_mmmd - else: - return lru_cache(maxsize=cache_size)(wavefunction_mmmd) - else: - if(not(cache)): - return c_wavefunction_mmmd - else: - return lru_cache(maxsize=cache_size)(c_wavefunction_mmmd) """ Main execution block to initialize the coefficient matrix and test the wavefunction computation. -This block checks for the existence of the precomputed normalized Hermite polynomial coefficients matrix. If it doesn't exist, -it computes the matrix and saves it for future use. Then, it performs a basic test to verify that the wavefunction -computation works as expected. References ---------- @@ -639,42 +531,32 @@ def wavefunction(s_mode: bool = True, o_dimensional: bool = True, complex_bool: https://docs.python.org/3/tutorial/modules.html . """ -package_dir = os.path.dirname(__file__) -matrix_filename = 'C_s_matrix.pickle' -matrix_path = os.path.join(package_dir, matrix_filename) - -if os.path.isfile(matrix_path): - with open(matrix_path, 'rb') as file: - c_s_matrix = pickle.load(file) -else: - c_s_matrix = create_normalized_hermite_coefficients_matrix(60) - with open(matrix_path, 'wb') as file: - pickle.dump(c_s_matrix, file) +c_s_matrix = create_normalized_hermite_coefficients_matrix(60) try: # Basic functionality test - test_output_smod_2 = wavefunction_smod(2, 10.0) - test_output_smod_less_fast_2 = wavefunction_smod(2, 10.0, more_fast=False) - test_output_smod_61 = wavefunction_smod(61, 10.0) - test_output_smod_less_fast_61 = wavefunction_smod(61, 10.0, more_fast=False) - test_output_mmod = wavefunction_mmod(2, 10.0) - test_output_smmd_2 = wavefunction_smmd(2, np.array([10.0,4.5])) - test_output_smmd_less_fast_2 = wavefunction_smmd(2, np.array([10.0,4.5]), more_fast=False) - test_output_smmd_61 = wavefunction_smmd(61, np.array([10.0,4.5])) - test_output_smmd_less_fast_61 = wavefunction_smmd(61, np.array([10.0,4.5]), more_fast=False) - test_output_mmmd = wavefunction_mmmd(2, np.array([10.0,4.5])) - test_output_c_smod = c_wavefunction_smod(2, 10.0 + 0.0j) - test_output_c_smod_less_fast = c_wavefunction_smod(2, 10.0 + 0.0j, more_fast=False) - test_output_c_smod_61 = c_wavefunction_smod(61, 10.0 + 0.0j) - test_output_c_smod_less_fast_61 = c_wavefunction_smod(61, 10.0 + 0.0j, more_fast=False) - test_output_c_mmod = c_wavefunction_mmod(2, 10.0 + 0.0j) - test_output_c_smmd = c_wavefunction_smmd(2, np.array([10.0 + 0.0j,4.5 + 0.0j])) - test_output_c_smmd_less_fast = c_wavefunction_smmd(2, np.array([10.0 + 0.0j,4.5 + 0.0j]), more_fast=False) - test_output_c_smmd_61 = c_wavefunction_smmd(61, np.array([10.0 + 0.0j,4.5 + 0.0j])) - test_output_c_smmd_less_fast_61 = c_wavefunction_smmd(61, np.array([10.0 + 0.0j,4.5 + 0.0j]), more_fast=False) - test_output_c_mmmd = c_wavefunction_mmmd(2, np.array([10.0 + 0.0j,4.5 + 0.0j])) + test_output_sfsp_2 = psi_n_single_fock_single_position(2, 10.0) + test_output_sfsp_less_fast_2 = psi_n_single_fock_single_position(2, 10.0, more_fast=False) + test_output_sfsp_61 = psi_n_single_fock_single_position(61, 10.0) + test_output_sfsp_less_fast_61 = psi_n_single_fock_single_position(61, 10.0, more_fast=False) + test_output_mfsp = psi_n_multiple_fock_single_position(2, 10.0) + test_output_sfmp_2 = psi_n_single_fock_multiple_position(2, np.array([10.0,4.5])) + test_output_sfmp_less_fast_2 = psi_n_single_fock_multiple_position(2, np.array([10.0,4.5]), more_fast=False) + test_output_sfmp_61 = psi_n_single_fock_multiple_position(61, np.array([10.0,4.5])) + test_output_sfmp_less_fast_61 = psi_n_single_fock_multiple_position(61, np.array([10.0,4.5]), more_fast=False) + test_output_mfmp = psi_n_multiple_fock_multiple_position(2, np.array([10.0,4.5])) + test_output_sfsp_c = psi_n_single_fock_single_position_complex(2, 10.0 + 0.0j) + test_output_sfsp_c_less_fast = psi_n_single_fock_single_position_complex(2, 10.0 + 0.0j, more_fast=False) + test_output_sfsp_c_61 = psi_n_single_fock_single_position_complex(61, 10.0 + 0.0j) + test_output_sfsp_c_less_fast_61 = psi_n_single_fock_single_position_complex(61, 10.0 + 0.0j, more_fast=False) + test_output_mfsp_c = psi_n_multiple_fock_single_position_complex(2, 10.0 + 0.0j) + test_output_sfmp_c = psi_n_single_fock_multiple_position(2, np.array([10.0 + 0.0j,4.5 + 0.0j])) + test_output_sfmp_c_less_fast = psi_n_single_fock_multiple_position(2, np.array([10.0 + 0.0j,4.5 + 0.0j]), more_fast=False) + test_output_sfmp_c_61 = psi_n_single_fock_multiple_position(61, np.array([10.0 + 0.0j,4.5 + 0.0j])) + test_output_sfmp_c_less_fast_61 = psi_n_single_fock_multiple_position(61, np.array([10.0 + 0.0j,4.5 + 0.0j]), more_fast=False) + test_output_mfmp_c = psi_n_multiple_fock_multiple_position_complex(2, np.array([10.0 + 0.0j,4.5 + 0.0j])) compilation_test = True print(f"Functionality Test Passed: {compilation_test}") except Exception as e: diff --git a/tests/test_wavefunction.py b/tests/test_wavefunction.py deleted file mode 100644 index 1b5f011..0000000 --- a/tests/test_wavefunction.py +++ /dev/null @@ -1,117 +0,0 @@ -import pytest -import numpy as np -from src.fast_wave.wavefunction import * - -@pytest.fixture(scope="module", autouse=True) -def initialize_c_s_matrix(): - """ - Fixture to initialize the global variable c_s_matrix before running tests. - """ - global c_s_matrix - c_s_matrix = create_normalized_hermite_coefficients_matrix(60) - -def test_hermite_sympy(): - """ - Tests the hermite_sympy function to verify the accuracy of Hermite polynomial computation. - """ - x = symbols("x") - h0 = hermite_sympy(0) - h1 = hermite_sympy(1) - h2 = hermite_sympy(2) - - assert h0 == 1 - assert h1 == 2 * x - assert h2 == 4 * x**2 - 2 - -def test_create_hermite_coefficients_table(): - """ - Tests the create_normalized_hermite_coefficients_table function to verify if the normalized coefficient matrix is correct. - """ - n_max = 2 - coeffs_table = create_normalized_hermite_coefficients_matrix(n_max) - - expected_table = np.zeros((3, 3)) - expected_table[0, 2] = 0.75112554 - expected_table[1, 1] = 1.06225193 - expected_table[2, 0] = 1.06225193 - expected_table[2, 2] = -0.53112597 - - assert np.allclose(coeffs_table, expected_table) - -def test_wavefunction_computation(): - """ - Tests the basic functionality of all wavefunction functions. - """ - - wave_smod = wavefunction(s_mode = True, o_dimensional = True, complex_bool = False, cache = False, cache_size = 128) - wave_smmd = wavefunction(s_mode = True, o_dimensional = False, complex_bool = False, cache = False, cache_size = 128) - wave_mmod = wavefunction(s_mode = False, o_dimensional = True, complex_bool = False, cache = False, cache_size = 128) - wave_mmmd = wavefunction(s_mode = False, o_dimensional = False, complex_bool = False, cache = False, cache_size = 128) - c_wave_smod = wavefunction(s_mode = True, o_dimensional = True, complex_bool = True, cache = False, cache_size = 128) - c_wave_smmd = wavefunction(s_mode = True, o_dimensional = False, complex_bool = True, cache = False, cache_size = 128) - c_wave_mmod = wavefunction(s_mode = False, o_dimensional = True, complex_bool = True, cache = False, cache_size = 128) - c_wave_mmmd = wavefunction(s_mode = False, o_dimensional = False, complex_bool = True, cache = False, cache_size = 128) - - # Testing basic functionality - test_output_odsm_2 = wave_smod(2, 10.0) - assert isinstance(test_output_odsm_2, float) - - test_output_odsm_61 = wave_smod(61, 10.0) - assert isinstance(test_output_odsm_61, float) - - test_output_odsm_less_fast_2 = wave_smod(2, 10.0, more_fast = False) - assert isinstance(test_output_odsm_less_fast_2, float) - - test_output_odsm_less_fast_61 = wave_smod(61, 10.0, more_fast = False) - assert isinstance(test_output_odsm_less_fast_61, float) - - test_output_odmm = wave_mmod(2, 10.0) - assert isinstance(test_output_odmm, np.ndarray) - - test_output_mdsm_2 = wave_smmd(2, np.array([10.0, 4.5])) - assert isinstance(test_output_mdsm_2, np.ndarray) - - test_output_mdsm_61 = wave_smmd(61, np.array([10.0, 4.5])) - assert isinstance(test_output_mdsm_61, np.ndarray) - - test_output_mdsm_less_fast_2 = wave_smmd(2, np.array([10.0, 4.5]), more_fast = False) - assert isinstance(test_output_mdsm_less_fast_2, np.ndarray) - - test_output_mdsm_less_fast_61 = wave_smmd(61, np.array([10.0, 4.5]), more_fast = False) - assert isinstance(test_output_mdsm_less_fast_61, np.ndarray) - - test_output_mdmm = wave_mmmd(2, np.array([10.0, 4.5])) - assert isinstance(test_output_mdmm, np.ndarray) - - test_output_c_odsm_2 = c_wave_smod(2, 10.0 + 0.0j) - assert isinstance(test_output_c_odsm_2, complex) - - test_output_c_odsm_61 = c_wave_smod(61, 10.0 + 0.0j) - assert isinstance(test_output_c_odsm_61, complex) - - test_output_c_odsm_less_fast_2 = c_wave_smod(2, 10.0 + 0.0j, more_fast = False) - assert isinstance(test_output_c_odsm_less_fast_2, complex) - - test_output_c_odsm_less_fast_61 = c_wave_smod(61, 10.0 + 0.0j, more_fast = False) - assert isinstance(test_output_c_odsm_less_fast_61, complex) - - test_output_c_odmm = c_wave_mmod(2, 10.0 + 0.0j) - assert isinstance(test_output_c_odmm, np.ndarray) - - test_output_c_mdsm_2 = c_wave_smmd(2, np.array([10.0 + 0.0j, 4.5 + 0.0j])) - assert isinstance(test_output_c_mdsm_2, np.ndarray) - - test_output_c_mdsm_61 = c_wave_smmd(61, np.array([10.0 + 0.0j, 4.5 + 0.0j])) - assert isinstance(test_output_c_mdsm_61, np.ndarray) - - test_output_c_mdsm_less_fast_2 = c_wave_smmd(2, np.array([10.0 + 0.0j, 4.5 + 0.0j]), more_fast = False) - assert isinstance(test_output_c_mdsm_less_fast_2, np.ndarray) - - test_output_c_mdsm_less_fast_61 = c_wave_smmd(61, np.array([10.0 + 0.0j, 4.5 + 0.0j]), more_fast = False) - assert isinstance(test_output_c_mdsm_less_fast_61, np.ndarray) - - test_output_c_mdmm = c_wave_mmmd(2, np.array([10.0 + 0.0j, 4.5 + 0.0j])) - assert isinstance(test_output_c_mdmm, np.ndarray) - - print("All functionality tests passed.") - diff --git a/tests/test_wavefunction_arb_prec.py b/tests/test_wavefunction_arb_prec.py deleted file mode 100644 index dbc5df0..0000000 --- a/tests/test_wavefunction_arb_prec.py +++ /dev/null @@ -1,44 +0,0 @@ -import numpy as np -from src.fast_wave.wavefunction_arb_prec import * - -def test_wavefunction_computation(): - """ - Tests the basic functionality of all wavefunction_arb_prec functions. - """ - - wave_smod_ap = wavefunction_arb_prec(s_mode = True, o_dimensional = True, complex_bool = False, cache = False, cache_size = 128) - wave_smmd_ap = wavefunction_arb_prec(s_mode = True, o_dimensional = False, complex_bool = False, cache = False, cache_size = 128) - wave_mmod_ap = wavefunction_arb_prec(s_mode = False, o_dimensional = True, complex_bool = False, cache = False, cache_size = 128) - wave_mmmd_ap = wavefunction_arb_prec(s_mode = False, o_dimensional = False, complex_bool = False, cache = False, cache_size = 128) - c_wave_smod_ap = wavefunction_arb_prec(s_mode = True, o_dimensional = True, complex_bool = True, cache = False, cache_size = 128) - c_wave_smmd_ap = wavefunction_arb_prec(s_mode = True, o_dimensional = False, complex_bool = True, cache = False, cache_size = 128) - c_wave_mmod_ap = wavefunction_arb_prec(s_mode = False, o_dimensional = True, complex_bool = True, cache = False, cache_size = 128) - c_wave_mmmd_ap = wavefunction_arb_prec(s_mode = False, o_dimensional = False, complex_bool = True, cache = False, cache_size = 128) - - # Testing basic functionality - test_output_odsm = wave_smod_ap(2, 10.0, 20) - assert isinstance(test_output_odsm, mpmath.ctx_mp_python.mpf) - - test_output_odmm = wave_mmod_ap(2, 10.0, 20) - assert isinstance(test_output_odmm, mpmath.matrices.matrices._matrix) - - test_output_mdsm = wave_smmd_ap(2, np.array([10.0, 4.5]), 20) - assert isinstance(test_output_mdsm, mpmath.matrices.matrices._matrix) - - test_output_mdmm = wave_mmmd_ap(2, np.array([10.0, 4.5]), 20) - assert isinstance(test_output_mdmm, mpmath.matrices.matrices._matrix) - - test_output_c_odsm = c_wave_smod_ap(2, 10.0 + 0.0j, 20) - assert isinstance(test_output_c_odsm, mpmath.ctx_mp_python.mpc) - - test_output_c_odmm = c_wave_mmod_ap(2, 10.0 + 0.0j, 20) - assert isinstance(test_output_c_odmm, mpmath.matrices.matrices._matrix) - - test_output_c_mdsm = c_wave_smmd_ap(2, np.array([10.0 + 0.0j, 4.5 + 0.0j]), 20) - assert isinstance(test_output_c_mdsm, mpmath.matrices.matrices._matrix) - - test_output_c_mdmm = c_wave_mmmd_ap(2, np.array([10.0 + 0.0j, 4.5 + 0.0j]), 20) - assert isinstance(test_output_c_mdmm, mpmath.matrices.matrices._matrix) - - print("All functionality tests passed.") - diff --git a/tests/test_wavefunction_cython.py b/tests/test_wavefunction_cython.py new file mode 100644 index 0000000..fd73df2 --- /dev/null +++ b/tests/test_wavefunction_cython.py @@ -0,0 +1,52 @@ +import numpy as np +import src.fast_wave.wavefunction_cython as wc + +def test_wavefunction_computation(): + """ + Tests the basic functionality of all wavefunction functions. + """ + + # Testing basic functionality + test_output_sfsp_2 = wc.psi_n_single_fock_single_position(2, 10.0) + assert isinstance(test_output_sfsp_2, float) + + test_output_sfsp_61 = wc.psi_n_single_fock_single_position(61, 10.0) + assert isinstance(test_output_sfsp_61, float) + + test_output_mfsp = wc.psi_n_multiple_fock_single_position(2, 10.0) + assert isinstance(test_output_mfsp , np.ndarray) + + test_output_sfmp_2 = wc.psi_n_single_fock_multiple_position(2, np.array([10.0, 4.5])) + assert isinstance(test_output_sfmp_2, np.ndarray) + + test_output_sfmp_61 = wc.psi_n_single_fock_multiple_position(61, np.array([10.0, 4.5])) + assert isinstance(test_output_sfmp_61, np.ndarray) + + test_output_mfmp = wc.psi_n_multiple_fock_multiple_position(2, np.array([10.0, 4.5])) + assert isinstance(test_output_mfmp, np.ndarray) + + test_output_sfsp_c_2 = wc.psi_n_single_fock_single_position_complex(2, 10.0 + 0.0j) + assert isinstance(test_output_sfsp_c_2, complex) + + test_output_sfsp_c_61 = wc.psi_n_single_fock_single_position_complex(61, 10.0 + 0.0j) + assert isinstance(test_output_sfsp_c_61, complex) + + test_output_mfsp_c = wc.psi_n_multiple_fock_single_position_complex(2, 10.0 + 0.0j) + assert isinstance(test_output_mfsp_c, np.ndarray) + + test_output_sfmp_c_2 = wc.psi_n_single_fock_multiple_position_complex(2, np.array([10.0 + 0.0j, 4.5 + 0.0j])) + assert isinstance(test_output_sfmp_c_2, np.ndarray) + + test_output_sfmp_c_61 = wc.psi_n_single_fock_multiple_position_complex(61, np.array([10.0 + 0.0j, 4.5 + 0.0j])) + assert isinstance(test_output_sfmp_c_61, np.ndarray) + + test_output_sfmp_c_less_fast_2 = wc.psi_n_single_fock_multiple_position_complex(2, np.array([10.0 + 0.0j, 4.5 + 0.0j])) + assert isinstance(test_output_sfmp_c_less_fast_2, np.ndarray) + + test_output_sfmp_c_less_fast_61 = wc.psi_n_single_fock_multiple_position_complex(61, np.array([10.0 + 0.0j, 4.5 + 0.0j])) + assert isinstance(test_output_sfmp_c_less_fast_61, np.ndarray) + + test_output_mfmp_c = wc.psi_n_multiple_fock_multiple_position_complex(2, np.array([10.0 + 0.0j, 4.5 + 0.0j])) + assert isinstance(test_output_mfmp_c, np.ndarray) + + print("All functionality tests passed.") \ No newline at end of file diff --git a/tests/test_wavefunction_mpmath.py b/tests/test_wavefunction_mpmath.py new file mode 100644 index 0000000..d78e992 --- /dev/null +++ b/tests/test_wavefunction_mpmath.py @@ -0,0 +1,36 @@ +import numpy as np +import mpmath +import src.fast_wave.wavefunction_mpmath as wm + +def test_wavefunction_computation(): + """ + Tests the basic functionality of all wavefunction_arb_prec functions. + """ + + # Testing basic functionality + test_output_sfsp = wm.psi_n_single_fock_single_position(2, 10.0, 20) + assert isinstance(test_output_sfsp, mpmath.ctx_mp_python.mpf) + + test_output_mfsp = wm.psi_n_multiple_fock_single_position(2, 10.0, 20) + assert isinstance(test_output_mfsp , mpmath.matrices.matrices._matrix) + + test_output_sfmp = wm.psi_n_single_fock_multiple_position(2, np.array([10.0, 4.5]), 20) + assert isinstance(test_output_sfmp, mpmath.matrices.matrices._matrix) + + test_output_mfmp = wm.psi_n_multiple_fock_multiple_position(2, np.array([10.0, 4.5]), 20) + assert isinstance(test_output_mfmp, mpmath.matrices.matrices._matrix) + + test_output_sfsp_c = wm.psi_n_single_fock_single_position_complex(2, 10.0 + 0.0j, 20) + assert isinstance(test_output_sfsp_c, mpmath.ctx_mp_python.mpc) + + test_output_mfsp_c = wm.psi_n_multiple_fock_single_position_complex(2, 10.0 + 0.0j, 20) + assert isinstance(test_output_mfsp_c, mpmath.matrices.matrices._matrix) + + test_output_sfmp_c = wm.psi_n_single_fock_multiple_position_complex(2, np.array([10.0 + 0.0j, 4.5 + 0.0j]), 20) + assert isinstance(test_output_sfmp_c, mpmath.matrices.matrices._matrix) + + test_output_mfmp_c = wm.psi_n_multiple_fock_multiple_position_complex(2, np.array([10.0 + 0.0j, 4.5 + 0.0j]), 20) + assert isinstance(test_output_mfmp_c, mpmath.matrices.matrices._matrix) + + print("All functionality tests passed.") + diff --git a/tests/test_wavefunction_numba.py b/tests/test_wavefunction_numba.py new file mode 100644 index 0000000..aa4ac32 --- /dev/null +++ b/tests/test_wavefunction_numba.py @@ -0,0 +1,109 @@ +import pytest +import numpy as np +from sympy import symbols +import src.fast_wave.wavefunction_numba as wn + +@pytest.fixture(scope="module", autouse=True) +def initialize_c_s_matrix(): + """ + Fixture to initialize the global variable c_s_matrix before running tests. + """ + global c_s_matrix + c_s_matrix = wn.create_normalized_hermite_coefficients_matrix(60) + +def test_hermite_sympy(): + """ + Tests the hermite_sympy function to verify the accuracy of Hermite polynomial computation. + """ + x = symbols("x") + h0 = wn.hermite_sympy(0) + h1 = wn.hermite_sympy(1) + h2 = wn.hermite_sympy(2) + + assert h0 == 1 + assert h1 == 2 * x + assert h2 == 4 * x**2 - 2 + +def test_create_hermite_coefficients_table(): + """ + Tests the create_normalized_hermite_coefficients_table function to verify if the normalized coefficient matrix is correct. + """ + n_max = 2 + coeffs_table = wn.create_normalized_hermite_coefficients_matrix(n_max) + + expected_table = np.zeros((3, 3)) + expected_table[0, 2] = 0.75112554 + expected_table[1, 1] = 1.06225193 + expected_table[2, 0] = 1.06225193 + expected_table[2, 2] = -0.53112597 + + assert np.allclose(coeffs_table, expected_table) + +def test_wavefunction_computation(): + """ + Tests the basic functionality of all wavefunction functions. + """ + + # Testing basic functionality + test_output_sfsp_2 = wn.psi_n_single_fock_single_position(2, 10.0) + assert isinstance(test_output_sfsp_2, float) + + test_output_sfsp_61 = wn.psi_n_single_fock_single_position(61, 10.0) + assert isinstance(test_output_sfsp_61, float) + + test_output_sfsp_less_fast_2 = wn.psi_n_single_fock_single_position(2, 10.0, more_fast = False) + assert isinstance(test_output_sfsp_less_fast_2, float) + + test_output_sfsp_less_fast_61 = wn.psi_n_single_fock_single_position(61, 10.0, more_fast = False) + assert isinstance(test_output_sfsp_less_fast_61, float) + + test_output_mfsp = wn.psi_n_multiple_fock_single_position(2, 10.0) + assert isinstance(test_output_mfsp , np.ndarray) + + test_output_sfmp_2 = wn.psi_n_single_fock_multiple_position(2, np.array([10.0, 4.5])) + assert isinstance(test_output_sfmp_2, np.ndarray) + + test_output_sfmp_61 = wn.psi_n_single_fock_multiple_position(61, np.array([10.0, 4.5])) + assert isinstance(test_output_sfmp_61, np.ndarray) + + test_output_sfmp_less_fast_2 = wn.psi_n_single_fock_multiple_position(2, np.array([10.0, 4.5]), more_fast = False) + assert isinstance( test_output_sfmp_less_fast_2, np.ndarray) + + test_output_sfmp_less_fast_61 = wn.psi_n_single_fock_multiple_position(61, np.array([10.0, 4.5]), more_fast = False) + assert isinstance(test_output_sfmp_less_fast_61, np.ndarray) + + test_output_mfmp = wn.psi_n_multiple_fock_multiple_position(2, np.array([10.0, 4.5])) + assert isinstance(test_output_mfmp, np.ndarray) + + test_output_sfsp_c_2 = wn.psi_n_single_fock_single_position_complex(2, 10.0 + 0.0j) + assert isinstance(test_output_sfsp_c_2, complex) + + test_output_sfsp_c_61 = wn.psi_n_single_fock_single_position_complex(61, 10.0 + 0.0j) + assert isinstance(test_output_sfsp_c_61, complex) + + test_output_sfsp_c_less_fast_2 = wn.psi_n_single_fock_single_position_complex(2, 10.0 + 0.0j, more_fast = False) + assert isinstance(test_output_sfsp_c_less_fast_2, complex) + + test_output_sfsp_c_less_fast_61 = wn.psi_n_single_fock_single_position_complex(61, 10.0 + 0.0j, more_fast = False) + assert isinstance(test_output_sfsp_c_less_fast_61, complex) + + test_output_mfsp_c = wn.psi_n_multiple_fock_single_position_complex(2, 10.0 + 0.0j) + assert isinstance(test_output_mfsp_c, np.ndarray) + + test_output_sfmp_c_2 = wn.psi_n_single_fock_multiple_position_complex(2, np.array([10.0 + 0.0j, 4.5 + 0.0j])) + assert isinstance(test_output_sfmp_c_2, np.ndarray) + + test_output_sfmp_c_61 = wn.psi_n_single_fock_multiple_position_complex(61, np.array([10.0 + 0.0j, 4.5 + 0.0j])) + assert isinstance(test_output_sfmp_c_61, np.ndarray) + + test_output_sfmp_c_less_fast_2 = wn.psi_n_single_fock_multiple_position_complex(2, np.array([10.0 + 0.0j, 4.5 + 0.0j]), more_fast = False) + assert isinstance(test_output_sfmp_c_less_fast_2, np.ndarray) + + test_output_sfmp_c_less_fast_61 = wn.psi_n_single_fock_multiple_position_complex(61, np.array([10.0 + 0.0j, 4.5 + 0.0j]), more_fast = False) + assert isinstance(test_output_sfmp_c_less_fast_61, np.ndarray) + + test_output_mfmp_c = wn.psi_n_multiple_fock_multiple_position_complex(2, np.array([10.0 + 0.0j, 4.5 + 0.0j])) + assert isinstance(test_output_mfmp_c, np.ndarray) + + print("All functionality tests passed.") +