requirements

remove files

pre-commit cleanup

ruff updates

plane waves

planewaves

palane waves examples

scipy fucntions

pyright off

plane wave operators

wavefunctions init

Plabe wave operators example

fourier methods

init file for fourier methods

.pre-commit-config.yaml

@@ -1,27 +1,4 @@
 repos:
-  - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
-    hooks:
-      - id: trailing-whitespace
-      - id: end-of-file-fixer
-      - id: check-toml
-      - id: check-yaml
-      - id: check-added-large-files
-      # Python-specific
-      - id: check-ast
-      - id: check-docstring-first
-      - id: debug-statements
-
-  - repo: https://github.com/crate-ci/typos
-    rev: v1.16.23
-    hooks:
-      - id: typos
-        args: []
-
-  - repo: https://github.com/psf/black
-    rev: 23.11.0
-    hooks:
-      - id: black
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
@@ -33,7 +10,12 @@ repos:
       # Run the formatter.
       - id: ruff-format
 
-  - repo: https://github.com/RobertCraigie/pyright-python
-    rev: v1.1.357
+  - repo: https://github.com/psf/black
+    rev: 23.11.0
     hooks:
-      - id: pyright
+      - id: black
+
+  # - repo: https://github.com/RobertCraigie/pyright-python
+  #   rev: v1.1.357
+  #   hooks:
+  #     - id: pyright

grid1q/operators/__init__.py

@@ -0,0 +1 @@
+"""Init file for the operators module."""

grid1q/operators/plane_wave.py

@@ -0,0 +1,130 @@
+"""Plane wave Hamiltonian for a ND system."""
+
+import multiprocessing
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def kenetic(
+    k_points: NDArray[np.int32],
+    hbar: float = 1.0,
+    m: float = 1.0,
+) -> NDArray[np.complex128]:
+    """Return the diagonal kinetic energy matrix for a ND plane wave system.
+
+    Args:
+        k_points: The k points of the system.
+        hbar: The reduced Planck constant.
+        m: The mass of the particle.
+
+    Returns:
+        The digonal kinetic energy matrix.
+    """
+    return np.diag([(hbar**2 * np.dot(k, k) ** 2) / (2 * m) for k in k_points])
+
+
+def elec_nuc_potential_element(
+    args: tuple[
+        int,
+        int,
+        NDArray[np.int32],
+        NDArray[np.int32],
+        NDArray[np.float64],
+        float,
+    ],
+) -> tuple[int, int, complex]:
+    """Return the matrix element of the electron-nucleus potential.
+
+    This functions is parallelized and should be used with the multiprocessing
+    module. Where the each matrix element will be calculated in parallel.
+
+    Args:
+        args: A tuple containing the following elements:
+            i: The row index of the matrix element.
+            j: The column index of the matrix element.
+            k_bra: The bra k point.
+            k_ket: The ket k point.
+            r_pos: The position of the nucleus.
+            cell_area: The area of the cell.
+
+    Returns:
+        The matrix element of the electron-nucleus potential.
+    """
+    i, j, k_bra, k_ket, r_pos, cell_area = args
+    mat_element = 0
+    for r in r_pos:
+        if np.array_equal(k_bra, k_ket):
+            mat_element = 0
+        else:
+            mat_element += (
+                (4 * np.pi) / (cell_area * np.linalg.norm(k_bra - k_ket) ** 2)
+            ) * np.exp(-1j * np.dot(k_bra - k_ket, r))
+    return i, j, mat_element
+
+
+def elec_nuc_potential(
+    k_points: NDArray[np.int32],
+    cell_area: float,
+    r_pos: NDArray[np.float64],
+) -> NDArray[np.complex128]:
+    """Return the electron-nucleus potential matrix for a ND plane wave system.
+
+    This function is parallelized and should be used with the multiprocessing.
+
+    Args:
+        k_points: The k points of the system.
+        cell_area: The area of the cell.
+        r_pos: The position of the nucleus.
+
+    Returns:
+        The electron-nucleus potential matrix.
+    """
+    mat = np.zeros((len(k_points), len(k_points)), dtype=complex)
+    args_list: list[
+        tuple[
+            tuple[
+                int,
+                int,
+                NDArray[np.int32],
+                NDArray[np.int32],
+                NDArray[np.float64],
+                float,
+            ]
+        ]
+    ] = []
+    for i, k_bra in enumerate(k_points):
+        for j, k_ket in enumerate(k_points):
+            args_list.append((i, j, k_bra, k_ket, r_pos, cell_area))
+
+    with multiprocessing.Pool() as pool:
+        results = pool.map(elec_nuc_potential_element, args_list)
+
+    for i, j, mat_element in results:
+        mat[i, j] = mat_element
+
+    return mat
+
+
+def plane_wave_hamiltonian(
+    k_points: NDArray[np.int32],
+    cell_area: float,
+    r_pos: NDArray[np.float64],
+    hbar: float = 1.0,
+    m: float = 1.0,
+) -> NDArray[np.complex128]:
+    """Return the Hamiltonian matrix for a ND plane wave system.
+
+    This function is parallelized and should be used with the multiprocessing.
+
+    Args:
+        k_points: The k points of the system.
+        cell_area: The area of the cell.
+        r_pos: The position of the nucleus.
+        hbar: The reduced Planck constant.
+        m: The mass of the particle.
+
+    Returns:
+    The Hamiltonian matrix.
+    """
+    return kenetic(k_points, hbar, m) + elec_nuc_potential(k_points, cell_area, r_pos)

grid1q/utils/fourier_methods/__init__.py

@@ -0,0 +1,3 @@
+"""init file for fourier_methods module."""
+
+from .fourier import dft, dft2, idft, idft2

grid1q/utils/fourier_methods/fourier.py

@@ -0,0 +1,121 @@
+"""Fourier methods for 1D and 2D signals."""
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def dft(
+    signal: NDArray[np.complex128],
+    k_min: int,
+    k_max: int,
+) -> NDArray[np.complex128]:
+    """Compute the Discrete Fourier Transform of a signal.
+
+    Args:
+        signal: The input signal.
+        k_min: The minimum index of the DFT.
+        k_max: The maximum index of the DFT.
+
+    Returns:
+        The DFT of the signal.
+    """
+    truncated_dft_result = np.zeros(k_max - k_min, dtype=complex)
+
+    for k in range(k_min, k_max):
+        sum_result = 0
+        for n in range(len(signal)):
+            sum_result += signal[n] * np.exp(-1j * 2 * np.pi * k * n / len(signal))
+        truncated_dft_result[k - k_min] = sum_result / len(signal)
+
+    return truncated_dft_result
+
+
+def idft(
+    truncated_dft_result: NDArray[np.complex128],
+    sig_len: int,
+    k_min: int,
+    k_max: int,
+) -> NDArray[np.complex128]:
+    """Compute the Inverse Discrete Fourier Transform of a signal.
+
+    Args:
+        truncated_dft_result: The truncated DFT of the signal.
+        sig_len: The length of the original signal.
+        k_min: The minimum index of the DFT.
+        k_max: The maximum index of the DFT.
+
+    Returns:
+        The IDFT of the signal.
+    """
+    reconstructed_signal = np.zeros(sig_len, dtype=complex)
+
+    for n in range(sig_len):
+        sum_result = 0
+        for k in range(k_min, k_max):
+            sum_result += truncated_dft_result[k - k_min] * np.exp(
+                1j * 2 * np.pi * k * n / sig_len,
+            )
+        reconstructed_signal[n] = sum_result
+
+    return reconstructed_signal
+
+
+def dft2(
+    signal_2d: NDArray[np.complex128],
+    k_min: int,
+    k_max: int,
+) -> NDArray[np.complex128]:
+    """Compute the 2D Discrete Fourier Transform of a signal.
+
+    Args:
+        signal_2d: The input 2D signal.
+        k_min: The minimum index of the DFT.
+        k_max: The maximum index of the DFT.
+
+    Returns:
+        The 2D DFT of the signal.
+    """
+    # Apply DFT to each row
+    dft_rows = np.array([dft(row, k_min, k_max) for row in signal_2d])
+
+    # Apply DFT to each column of the result
+    dft_columns = np.array(
+        [dft(dft_rows[:, col], k_min, k_max) for col in range(dft_rows.shape[1])],
+    )
+
+    # Transpose the result back to the original orientation
+    return dft_columns.T
+
+
+def idft2(
+    truncated_dft_2d: NDArray[np.complex128],
+    original_shape: tuple[int, int],
+    k_min: int,
+    k_max: int,
+) -> NDArray[np.complex128]:
+    """Compute the 2D Inverse Discrete Fourier Transform of a signal.
+
+    Args:
+        truncated_dft_2d: The truncated 2D DFT of the signal.
+        original_shape: The shape of the original signal.
+        k_min: The minimum index of the DFT.
+        k_max: The maximum index of the DFT.
+
+    Returns:
+        The 2D IDFT of the signal.
+    """
+    # Apply IDFT to each column
+    idft_cols = np.array(
+        [
+            idft(truncated_dft_2d[:, col], original_shape[0], k_min, k_max)
+            for col in range(truncated_dft_2d.shape[1])
+        ],
+    ).T
+
+    # Apply IDFT to each row of the result
+    return np.array(
+        [
+            idft(idft_cols[row], original_shape[1], k_min, k_max)
+            for row in range(idft_cols.shape[0])
+        ],
+    )

grid1q/wavefunctions/__init__.py

@@ -0,0 +1 @@
+"""Wavefunctions package."""

grid1q/wavefunctions/plane_wave.py

@@ -0,0 +1,70 @@
+"""This module provides functions to represent plane waves in ND space."""
+
+from collections.abc import Callable
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def plane_wave(
+    coeffs_dict: dict[tuple[int, ...], np.complex128],
+) -> Callable[[NDArray[np.float64]], NDArray[np.complex128]]:
+    """Return a function that represents a plane wave in ND space.
+
+    The plane wave is represented as a linear combination of complex exponentials.
+    The tuple keys are used to determine the dimension of the k space.
+
+    Args:
+        coeffs_dict: A dictionary where the keys are tuples of integers representing the
+        wave vector and the values are the coefficients of the complex exponentials.
+
+    Returns:
+        A function that takes a vector r in ND space and returns the value of the plane
+        wave at r. Where the x,y ... are stacked in a vector r = [x,y,...]
+
+
+    Example:
+        plane_wave({(1,0):
+        0.4343, (0,1): 0.343434}) returns a function that represents a plane wave in 2D
+        space with wave vector (1,0) and (0,1), with expansion coefficients 0.4343
+        and 0.343434 respectively.
+    """
+    dim = len(next(iter(coeffs_dict.keys())))
+
+    def space_function(
+        r: NDArray[np.float64],
+    ) -> NDArray[np.complex128]:  # Vectorised function
+        """Return the value of the plane wave at r.
+
+        It is not normalised and you must use plane_wave_renorm to normalise it.
+        Where the x,y ... are stacked in a vector r = [x,y,...].
+
+        Args:
+            r: A vector in ND space.
+
+        Returns:
+            The value of the plane wave at r.
+        """
+        return sum(
+            [
+                coeffs_dict[k]
+                * np.exp(np.pi * 1j * np.tensordot(r, k, axes=([dim], [0])))
+                for k in coeffs_dict
+            ],
+        )  # type: ignore  # noqa: PGH003
+
+    return space_function
+
+
+def plane_wave_renorm(plane_wave_r: NDArray[np.complex128]) -> NDArray[np.complex128]:
+    """Return the normalised plane wave.
+
+    Args:
+        plane_wave_r: The plane wave to normalise.
+
+    Returns:
+        The normalised plane wave.
+    """
+    pw_flat = plane_wave_r.flatten()
+    pw_flat_norm = pw_flat / np.linalg.norm(pw_flat)
+    return pw_flat_norm.reshape(plane_wave_r.shape)