diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index c2a8184..30ca203 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -1,63 +1,58 @@ -name: MacroDensity CI +name: CI on: - pull_request: - branches: - - master - - development(calysta) push: - branches: - - master - - development(calysta) + branches: [master, main] + pull_request: + branches: [master, main] jobs: - -# qa: -# runs-on: ubuntu-latest -# steps: -# - uses: actions/checkout@v2 -# - uses: pre-commit/action@v2.0.0 - test: - #needs: qa + runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: - python-version: ["3.9", "3.10", "3.11"] - os: [ubuntu-latest,] + os: [ubuntu-latest] + python-version: ["3.9", "3.10", "3.11", "3.12"] - runs-on: ${{matrix.os}} steps: - - uses: actions/checkout@v3 - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 - with: - python-version: ${{ matrix.python-version }} - - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install setuptools setuptools_scm wheel - pip install -r requirements.txt - pip install pytest - pip install pytest-cov - pip install pytest-mpl - pip install -e . - - - name: Check package versions - run: | - pip show -V pytest - pip show -V ase - pip show -V matplotlib - - - name: Run tests - run: pytest --cov=macrodensity --mpl tests/unit_tests.py --cov-report=xml -v - - - name: Upload coverage to Codecov - uses: codecov/codecov-action@v3 - with: - token: ${{ secrets.CODECOV_TOKEN }} - fail_ci_if_error: true - verbose: true - files: ./coverage.xml + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e . + pip install pytest pytest-cov + + - name: Run tests + run: pytest tests/unit_tests.py -v --cov=macrodensity --cov-report=xml + + - name: Upload coverage to Codecov + if: matrix.python-version == '3.11' + uses: codecov/codecov-action@v4 + with: + files: ./coverage.xml + fail_ci_if_error: false + + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install linters + run: | + python -m pip install --upgrade pip + pip install ruff + + - name: Check code style with ruff + run: ruff check macrodensity/ --select=E,F,W --ignore=E501,E741,F401,F841,F821,F823,W605 diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000..8f10806 --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,28 @@ +cff-version: 1.2.0 +message: "If you use this software, please cite it as below." +title: "MacroDensity" +abstract: "A Python package for analysis of electrostatic potential and electron density from electronic structure calculations." +type: software +authors: + - family-names: "Butler" + given-names: "Keith T." + orcid: "https://orcid.org/0000-0001-5432-5597" + - family-names: "Walsh" + given-names: "Aron" + orcid: "https://orcid.org/0000-0001-5460-7033" + - family-names: "Harnett-Caulfield" + given-names: "Liam" + - family-names: "Tesiman" + given-names: "Calysta A." + orcid: "https://orcid.org/0009-0008-7784-4320" + - family-names: "Mosquera-Lois" + given-names: "Irea" +repository-code: "https://github.com/WMD-group/MacroDensity" +url: "https://macrodensity.readthedocs.io/" +license: MIT +keywords: + - density functional theory + - electrostatic potential + - electron density + - VASP + - materials science \ No newline at end of file diff --git a/README.md b/README.md index 48ace2d..164cddc 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,9 @@ -# Welcome to MacroDensity +# MacroDensity + +[![Build Status](https://github.com/WMD-group/MacroDensity/actions/workflows/ci.yaml/badge.svg)](https://github.com/WMD-group/MacroDensity/actions) +[![PyPI version](https://badge.fury.io/py/macrodensity.svg)](https://pypi.org/project/macrodensity/) +[![Documentation](https://readthedocs.org/projects/macrodensity/badge/?version=latest)](https://macrodensity.readthedocs.io/) +[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT) ``MacroDensity`` is a Python package to post-process electrostatic potential and electron density files from electronic structure calculations and plot in a number of ways, including: @@ -40,6 +45,36 @@ pip install -e . ``` +# Quick Example + +Calculate and plot the planar average of an electrostatic potential from a VASP LOCPOT file: + +```python +import macrodensity as md + +# Plot planar average along the z-axis +df, fig = md.plot_planar_average( + lattice_vector=20.0, # Length of cell in z-direction (Angstrom) + input_file='LOCPOT', + axis='z' +) +``` + +Calculate the on-site potential for a specific atomic species: + +```python +# Get potential at oxygen sites +potentials, fig = md.plot_on_site_potential( + species='O', + sample_cube=[2, 2, 2], + potential_file='LOCPOT', + coordinate_file='POSCAR' +) +``` + +See the [tutorials](https://macrodensity.readthedocs.io/en/latest/tutorials.html) for more detailed examples. + + # Literature For more information on the theory behind the package, please see the following references: diff --git a/JOSSpaper/JOSSpaper.bib b/docs/JOSSpaper/JOSSpaper.bib similarity index 100% rename from JOSSpaper/JOSSpaper.bib rename to docs/JOSSpaper/JOSSpaper.bib diff --git a/JOSSpaper/JOSSpaper.md b/docs/JOSSpaper/JOSSpaper.md similarity index 100% rename from JOSSpaper/JOSSpaper.md rename to docs/JOSSpaper/JOSSpaper.md diff --git a/JOSSpaper/JOSSpaper.pdf b/docs/JOSSpaper/JOSSpaper.pdf similarity index 100% rename from JOSSpaper/JOSSpaper.pdf rename to docs/JOSSpaper/JOSSpaper.pdf diff --git a/JOSSpaper/apa.csl b/docs/JOSSpaper/apa.csl similarity index 100% rename from JOSSpaper/apa.csl rename to docs/JOSSpaper/apa.csl diff --git a/JOSSpaper/figure.png b/docs/JOSSpaper/figure.png similarity index 100% rename from JOSSpaper/figure.png rename to docs/JOSSpaper/figure.png diff --git a/JOSSpaper/joss-logo.png b/docs/JOSSpaper/joss-logo.png similarity index 100% rename from JOSSpaper/joss-logo.png rename to docs/JOSSpaper/joss-logo.png diff --git a/JOSSpaper/latex.template b/docs/JOSSpaper/latex.template similarity index 100% rename from JOSSpaper/latex.template rename to docs/JOSSpaper/latex.template diff --git a/macrodensity/__init__.py b/macrodensity/__init__.py index d1f9ad6..0ecae99 100755 --- a/macrodensity/__init__.py +++ b/macrodensity/__init__.py @@ -2,14 +2,59 @@ MacroDensity is a package to read, process and plot electrostatic potential and electron density files from electronic structure calculations. """ -import math -import numpy as np -from scipy import interpolate - -from macrodensity.averages import * -from macrodensity.density import * -from macrodensity.io import * -from macrodensity.plotting import * -from macrodensity.tools import * -from macrodensity.utils import * +from macrodensity.averages import ( + macroscopic_average, + planar_average, + spherical_average, + travelling_volume_average, + volume_average, +) +from macrodensity.density import element_vol, gradient_magnitude +from macrodensity.io import ( + get_band_extrema, + read_gulp_potential, + read_vasp_density, + read_vasp_density_classic, + read_vasp_parchg, +) +from macrodensity.plotting import ( + energy_band_alignment_diagram, + plot_active_plane, + plot_active_space, + plot_field_at_point, + plot_on_site_potential, + plot_planar_average, + plot_plane_field, + plot_variation_along_vector, +) +from macrodensity.tools import ( + bulk_interstitial_alignment, + bulk_vac, + create_plotting_mesh, + diff_potentials, + dipole_correction, + extend_potential, + get_layer_sites, + match_resolution, + matched_spline_generate, + scissors_shift, + sort_potential, + spline_generate, + subs_potentials, + translate_grid, +) +from macrodensity.utils import ( + GCD, + GCD_List, + density_2_grid, + get_third_coordinate, + get_volume, + inverse_participation_ratio, + matrix_2_abc, + number_in_field, + numbers_2_grid, + one_2_2d, + points_2_plane, + vector_2_abscissa, +) diff --git a/macrodensity/averages.py b/macrodensity/averages.py index aa49a73..9e559c1 100644 --- a/macrodensity/averages.py +++ b/macrodensity/averages.py @@ -60,9 +60,7 @@ def macroscopic_average( ) macro_average[i] = macro_average[i] / period_points else: - macro_average[i] = ( - macro_average[i] + sum(potential[start:end]) / period_points - ) + macro_average[i] = sum(potential[start:end]) / period_points print("Average of the average = ", np.average(macro_average)) @@ -124,10 +122,10 @@ def volume_average( xv = int(n_origin[0] + travelled[0] + x) yv = int(n_origin[1] + travelled[1] + y) zv = int(n_origin[2] + travelled[2] + z) - # Minimum image convention - zv = int(zv - nz * round(zv / nz)) - yv = int(yv - ny * round(yv / ny)) - xv = int(xv - nx * round(xv / nx)) + # Periodic boundary wrapping + xv = xv % nx + yv = yv % ny + zv = zv % nz potential_cube[x, y, z] = grid[int(xv), int(yv), int(zv)] return potential_cube.mean(), np.var(potential_cube) @@ -221,7 +219,7 @@ def spherical_average( ) ## PRINTING - if print_output == True: + if print_output: print("Potential Variance") print("--------------------------------") print(cube_pot, " ", cube_var) diff --git a/macrodensity/density.py b/macrodensity/density.py index e7f4b76..80e4372 100755 --- a/macrodensity/density.py +++ b/macrodensity/density.py @@ -37,15 +37,7 @@ def gradient_magnitude( >>> print(grad_magnitude) """ - grad_mag = gx - for i in range(gx.shape[0]): - for j in range(gy.shape[1]): - for k in range(gz.shape[2]): - grad_mag[i, j, k] = np.sqrt( - gx[i, j, k] ** 2 + gy[i, j, k] ** 2 + gz[i, j, k] ** 2 - ) - - return grad_mag + return np.sqrt(gx**2 + gy**2 + gz**2) def element_vol(vol: float, nx: int, ny: int, nz: int) -> float: diff --git a/macrodensity/io.py b/macrodensity/io.py index 00fd94b..de54ffe 100644 --- a/macrodensity/io.py +++ b/macrodensity/io.py @@ -57,7 +57,7 @@ def read_gulp_potential(gulpfile: str = "gulp.out") -> tuple: for n, line in enumerate(lines): if line.rfind("Electrostatic potential on a grid") > -1: - for k in reversed(range(9, NGX * NGY * NGZ + 9)): + for k in range(9, NGX * NGY * NGZ + 9): potential.append(float(lines[n + k].split()[3])) return np.asarray(potential), NGX, NGY, NGZ, lattice diff --git a/macrodensity/plotting.py b/macrodensity/plotting.py index 8f20011..30b1f25 100644 --- a/macrodensity/plotting.py +++ b/macrodensity/plotting.py @@ -1,5 +1,7 @@ -"""macrodensity.plotting contains different configs of plotting functions such as band -alignment diagrams and potentials at different grid points.""" +""" +Module containing functions to plot different averages of the +electrostatic potential, like planar average and macroscopic averages. +""" from __future__ import division, print_function @@ -392,7 +394,7 @@ def plot_on_site_potential( ) elif "cube" in potential_file: grid_pot, atoms = cube.read_cube_data(potential_file) - vector_a = np.linalg.norm(atoms.cell[1]) + vector_a = np.linalg.norm(atoms.cell[0]) vector_b = np.linalg.norm(atoms.cell[1]) vector_c = np.linalg.norm(atoms.cell[2]) NGX = len(grid_pot) @@ -449,7 +451,7 @@ def plot_on_site_potential( ) ax.set_xlabel("Hartree potential (V)") ax.set_ylabel("% of centres") - plt.savefig(img_file) + fig.savefig(img_file) ## SAVING df = pd.DataFrame.from_dict({"Potential": potentials_list}, orient="index") @@ -523,7 +525,7 @@ def _save_df(planar, macro, output_file, interpolated_potential=None): if "cube" in input_file: potential, atoms = cube.read_cube_data(input_file) - vector_a = np.linalg.norm(atoms.cell[1]) + vector_a = np.linalg.norm(atoms.cell[0]) vector_b = np.linalg.norm(atoms.cell[1]) vector_c = np.linalg.norm(atoms.cell[2]) NGX = len(potential) @@ -672,13 +674,12 @@ def plot_field_at_point( ) ## Get the gradiens (Field), if required. ## Comment out if not required, due to compuational expense. - if grad_calc == True: + if grad_calc: print("Calculating gradients (Electic field, E=-Grad.V )...") grad_x, grad_y, grad_z = np.gradient( grid_pot[:, :, :], resolution_x, resolution_y, resolution_z ) - else: - pass + grad_x, grad_y, grad_z = -grad_x, -grad_y, -grad_z # ------------------------------------------------------------------ ## Get the equation for the plane ## This is the section for plotting on a user defined plane; @@ -695,8 +696,6 @@ def plot_field_at_point( Y2 = np.multiply(grad_y, grad_y) Z2 = np.multiply(grad_z, grad_z) grad_mag = np.linalg.norm(np.add(X2, Y2, Z2), axis=-1) - # This was my, non working, attempt to use the built in function. - # grad_mag=np.linalg.norm( [grad_y,grad_y,grad_z], axis=3) ## This function in Macrodensity averages Efield ACROSS Z for Slab calculations xx, yy, grd = md.create_plotting_mesh( @@ -736,9 +735,7 @@ def plot_field_at_point( width=1, headwidth=3, headlength=4, - ) # , - # units='xy', scale=10., zorder=3, color='blue', - # width=0.007, headwidth=3., headlength=4.) + ) plt.axis( "equal" @@ -792,6 +789,7 @@ def plot_plane_field( grad_x, grad_y, grad_z = np.gradient( grid_pot[:, :, :], resolution_x, resolution_y, resolution_z ) + grad_x, grad_y, grad_z = -grad_x, -grad_y, -grad_z # ------------------------------------------------------------------ ## Get the equation for the plane ## This is the section for plotting on a user defined plane; @@ -802,9 +800,8 @@ def plot_plane_field( b = numbers_2_grid(b_point, NGX, NGY, NGZ) c = numbers_2_grid(c_point, NGX, NGY, NGZ) plane_coeff = points_2_plane(a, b, c) - ## Get the gradients - XY = np.multiply(grad_x, grad_y) - grad_mag = np.multiply(XY, grad_z) + ## Get the gradient magnitude + grad_mag = np.sqrt(grad_x**2 + grad_y**2 + grad_z**2) ## Create the plane # print(NGX,NGY,NGZ,plane_coeff,grad_mag) xx, yy, grd = create_plotting_mesh(NGX, NGY, NGZ, plane_coeff, grad_mag) @@ -886,13 +883,12 @@ def plot_active_plane( ## Get the gradiens (Field), if required. ## Comment out if not required, due to compuational expense. - if grad_calc == True: + if grad_calc: print("Calculating gradients (Electic field, E=-Grad.V )...") grad_x, grad_y, grad_z = np.gradient( grid_pot[:, :, :], resolution_x, resolution_y, resolution_z ) - else: - pass + grad_x, grad_y, grad_z = -grad_x, -grad_y, -grad_z ## Convert the fractional points to grid points on the density surface a = numbers_2_grid(a_point, NGX, NGY, NGZ) @@ -900,9 +896,8 @@ def plot_active_plane( c = numbers_2_grid(c_point, NGX, NGY, NGZ) plane_coeff = points_2_plane(a, b, c) - ## Get the gradients - XY = np.multiply(grad_x, grad_y) - grad_mag = np.multiply(XY, grad_z) + ## Get the gradient magnitude + grad_mag = np.sqrt(grad_x**2 + grad_y**2 + grad_z**2) ## Create the plane xx, yy, grd = create_plotting_mesh(NGX, NGY, NGZ, plane_coeff, grad_x) @@ -927,7 +922,7 @@ def plot_active_plane( fig, ax = plt.subplots() ax.plot(planar) ax.plot(macro) - plt.savefig("Planar.eps") + fig.savefig("Planar.eps") plt.show() return fig diff --git a/macrodensity/tools.py b/macrodensity/tools.py index 62ad46e..e0733bb 100644 --- a/macrodensity/tools.py +++ b/macrodensity/tools.py @@ -1,7 +1,6 @@ -#! /usr/bin/env python """ -macrodensity.tools contains functions to read and manipulate the electronic -density data from a material. +Module containing utility functions to manipulate and process the +electronic density data from a material. """ from __future__ import division @@ -235,7 +234,7 @@ def bulk_vac(bulk: np.ndarray, slab: np.ndarray) -> np.ndarray: if s_pot[0] <= bulk[j, 0] and s_pot[0] > bulk[j - 1, 0]: new_bulk[i, :] = bulk[j, :] found = True - if found == False: + if not found: new_bulk[i, 0] = s_pot[0] new_bulk[i, 1] = 0 @@ -263,8 +262,8 @@ def match_resolution(A: np.ndarray, B: np.ndarray) -> tuple: >>> print(result_A) >>> print(result_B) """ - np.append(A, A[0, :]) - np.append(B, B[0, :]) + A = np.vstack([A, A[0, :]]) + B = np.vstack([B, B[0, :]]) resolution_a = (max(A[:, 0]) - min(A[:, 0])) / len(A) resolution_b = (max(B[:, 0]) - min(B[:, 0])) / len(B) new_resolution = min(resolution_a, resolution_b) / 3 @@ -371,10 +370,10 @@ def matched_spline_generate( res_a = length_A / float(len(A_new)) res_b = length_B / float(len(B_new)) for i in range(len(A_new)): - TD_A[i, 1] = A[i] + TD_A[i, 1] = A_new[i] TD_A[i, 0] = i * res_a for i in range(len(B_new)): - TD_B[i, 1] = B[i] + TD_B[i, 1] = B_new[i] TD_B[i, 0] = i * res_b return TD_A, TD_B @@ -516,17 +515,16 @@ def diff_potentials( >>> print(result) """ - resolution = potential_a[0, 0] - potential_b[0, 1] - new_potential = np.zeros(shape=((start - end) / resolution, 2)) + results = [] for i in range(len(potential_a)): if potential_a[i, 0] >= start and potential_a[i, 0] <= end: for j in range(len(potential_b)): if abs(potential_b[j, 0] - potential_a[i, 0]) <= tol: - new_potential[i, 1] = potential_a[i, 1] - potential_b[i, 1] - new_potential[i, 0] = potential_a[i, 0] + results.append([potential_a[i, 0], potential_a[i, 1] - potential_b[j, 1]]) + break - return new_potential + return np.array(results) if results else np.zeros((0, 2)) def subs_potentials(A: np.ndarray, B: np.ndarray, tol: float) -> np.ndarray: @@ -607,14 +605,14 @@ def translate_grid( for i in range(len(potential)): new_potential_trans[i, 0] = potential[i, 0] + translation new_potential_trans[i, 1] = potential[i, 1] - if periodic == True: + if periodic: new_potential_trans[i, 0] = new_potential_trans[ i, 0 ] - length * int( (new_potential_trans[i, 0] + boundary_shift) / length ) - if periodic == True: + if periodic: # Sort the numbers out if you have done periodic wrapping sorted_potential_trans = sort_potential(new_potential_trans) else: @@ -679,7 +677,138 @@ def create_plotting_mesh( for y in range(b): if p == "zzo": plane[x, y] = grad[x, y, c] - else: - pass + elif p == "zoz": + plane[x, y] = grad[x, c, y] + elif p == "ozz": + plane[x, y] = grad[c, x, y] return plane + + +def dipole_correction( + potential: np.ndarray, + lattice_vector: float, + gradient_threshold: float = 0.01, +) -> tuple: + """ + Apply dipole correction to remove artificial electric field in slab calculations. + + Automatically detects the vacuum region by finding flat areas of the potential, + then fits and subtracts a linear ramp to remove the artificial field. + + Parameters: + potential (np.ndarray): 1D planar-averaged potential array (from planar_average()). + + lattice_vector (float): Length of cell in averaging direction (Angstrom). + + gradient_threshold (float): Maximum gradient magnitude (eV/Angstrom) to + consider a region as vacuum. Default is 0.01. + + Returns: + tuple: (corrected_potential, electric_field, vacuum_indices) + - corrected_potential (np.ndarray): Potential with linear field removed. + - electric_field (float): Measured electric field in vacuum (eV/Angstrom). + - vacuum_indices (np.ndarray): Boolean array marking vacuum region. + + Example: + >>> from macrodensity import planar_average, read_vasp_density + >>> from macrodensity.utils import density_2_grid + >>> pot, NGX, NGY, NGZ, lattice = read_vasp_density('LOCPOT') + >>> grid, _ = density_2_grid(pot, NGX, NGY, NGZ) + >>> planar = planar_average(grid, NGX, NGY, NGZ, axis='z') + >>> corrected, field, vacuum = dipole_correction(planar, lattice[2,2]) + >>> print(f"Electric field in vacuum: {field:.4f} eV/Angstrom") + """ + n_points = len(potential) + resolution = lattice_vector / n_points + positions = np.arange(n_points) * resolution + + # Calculate local gradient to identify flat (vacuum) regions + local_gradient = np.gradient(potential, resolution) + + # Identify vacuum: regions where gradient magnitude is below threshold + vacuum_mask = np.abs(local_gradient) < gradient_threshold + + # Find contiguous vacuum regions + # Look for the longest stretch of True values + vacuum_indices = np.where(vacuum_mask)[0] + + if len(vacuum_indices) < 2: + raise ValueError( + f"Could not identify vacuum region. Try increasing gradient_threshold " + f"(current: {gradient_threshold}). Max gradient in data: {np.max(np.abs(local_gradient)):.4f}" + ) + + # Find the largest contiguous region + # Split into groups where indices are consecutive + gaps = np.diff(vacuum_indices) + split_points = np.where(gaps > 1)[0] + 1 + regions = np.split(vacuum_indices, split_points) + + # Select the largest region + largest_region = max(regions, key=len) + + if len(largest_region) < 3: + raise ValueError( + f"Vacuum region too small for fitting ({len(largest_region)} points). " + f"Try increasing gradient_threshold." + ) + + # Fit a line to the potential in the vacuum region + vacuum_positions = positions[largest_region] + vacuum_potential = potential[largest_region] + + # Linear fit: potential = slope * position + intercept + slope, intercept = np.polyfit(vacuum_positions, vacuum_potential, 1) + + # The electric field is the negative of the potential slope + electric_field = -slope + + # Create correction array (linear ramp across entire cell) + correction = slope * positions + + # Apply correction + corrected_potential = potential - correction + + # Shift so minimum is at zero (optional, for cleaner output) + corrected_potential = corrected_potential - np.min(corrected_potential) + + # Create boolean mask for vacuum region + vacuum_mask_output = np.zeros(n_points, dtype=bool) + vacuum_mask_output[largest_region] = True + + return corrected_potential, electric_field, vacuum_mask_output + + +def get_layer_sites( + structure, + epsilon: float = 1.5, +) -> dict: + """ + Return sites in top and bottom layers of a slab structure. + + Sites are selected based on their z coordinate relative to the + top/bottom of the slab within a tolerance epsilon. + + Parameters: + structure: pymatgen Structure object of the slab. + + epsilon (float): Distance from top/bottom of slab to select sites + within (Angstrom). Default is 1.5. + + Returns: + dict: Dictionary with 'top_layer' and 'bottom_layer' keys, + each containing a list of sites. + + Example: + >>> from pymatgen.core.structure import Structure + >>> structure = Structure.from_file('POSCAR') + >>> layers = get_layer_sites(structure, epsilon=1.5) + >>> print(f"Top layer sites: {len(layers['top_layer'])}") + >>> print(f"Bottom layer sites: {len(layers['bottom_layer'])}") + """ + max_z = max(site.coords[2] for site in structure) + min_z = min(site.coords[2] for site in structure) + top_layer = [site for site in structure if site.coords[2] >= max_z - epsilon] + bottom_layer = [site for site in structure if site.coords[2] <= min_z + epsilon] + return {"top_layer": top_layer, "bottom_layer": bottom_layer} diff --git a/macrodensity/utils.py b/macrodensity/utils.py index f6f2539..9acac15 100644 --- a/macrodensity/utils.py +++ b/macrodensity/utils.py @@ -398,7 +398,7 @@ def density_2_grid( for j in range(ny): for i in range(nx): Potential_grid[i, j, k] = density[l] / volume - if charge == True: + if charge: # Convert the charge density to a number of electrons point_volume = volume / (nx * ny * nz) Potential_grid[i, j, k] = ( @@ -408,7 +408,7 @@ def density_2_grid( l = l + 1 total_electrons = total_electrons / (nx * ny * nz) - if charge == True: + if charge: print("Total electrons: ", total_electrons) return Potential_grid, total_electrons diff --git a/tests/unit_tests.py b/tests/unit_tests.py index 218fa17..d025095 100644 --- a/tests/unit_tests.py +++ b/tests/unit_tests.py @@ -342,6 +342,38 @@ def test_moving_cube(self): self.addCleanup(os.remove, "potential_variation.csv") self.addCleanup(os.remove, "potential_variation.png") + def test_dipole_correction(self): + """Tests the dipole_correction function""" + # Create a synthetic potential with a linear slope (artificial field) + n_points = 100 + lattice_vector = 20.0 # Angstrom + positions = np.linspace(0, lattice_vector, n_points, endpoint=False) + + # Create potential: flat in "vacuum" (first and last 20%) with a slope + # Simulates a slab with vacuum on both sides and artificial field + slope = 0.05 # eV/Angstrom (artificial field) + potential = slope * positions # Linear ramp + + # Add some structure in the "slab" region (middle 60%) + slab_start = int(0.2 * n_points) + slab_end = int(0.8 * n_points) + potential[slab_start:slab_end] += 2.0 * np.sin( + np.linspace(0, 2 * np.pi, slab_end - slab_start) + ) + + # Apply dipole correction + corrected, field, vacuum = md.dipole_correction( + potential, lattice_vector, gradient_threshold=0.1 + ) + + # Check that the electric field is detected correctly + self.assertAlmostEqual(abs(field), slope, places=2) + + # Check that the corrected potential in vacuum is flatter + vacuum_gradient_before = np.abs(np.gradient(potential[:slab_start])) + vacuum_gradient_after = np.abs(np.gradient(corrected[:slab_start])) + self.assertLess(np.mean(vacuum_gradient_after), np.mean(vacuum_gradient_before)) + class TestPlottingFunctions(unittest.TestCase): @@ -423,10 +455,10 @@ def test_plot_planar_cube(self): dfpot, _ = md.plot_planar_average( input_file=potential, lattice_vector=4.75 ) - self.assertEqual(dfden["Planar"].tolist()[0], 0.0200083723051778) - self.assertEqual(dfden["Planar"].tolist()[-1], 0.019841719274268536) - self.assertEqual(dfpot["Planar"].tolist()[0], -0.562656062923066) - self.assertEqual(dfpot["Planar"].tolist()[-1], -0.581089179258661) + self.assertAlmostEqual(dfden["Planar"].tolist()[0], 0.0200083723051778, places=10) + self.assertAlmostEqual(dfden["Planar"].tolist()[-1], 0.019841719274268536, places=10) + self.assertAlmostEqual(dfpot["Planar"].tolist()[0], -0.562656062923066, places=10) + self.assertAlmostEqual(dfpot["Planar"].tolist()[-1], -0.581089179258661, places=10) self.addCleanup(os.remove, "planar_average.csv") self.addCleanup(os.remove, "planar_average.png")