from typing import TYPE_CHECKING, Optional
import numpy as np
if TYPE_CHECKING:
from ryd_numerov.units import NDArray
[docs]
class Grid:
"""A grid object storing all relevant information about the grid points.
We store the grid in the dimensionless form x = r/a_0, as well as in the scaled dimensionless form z = sqrt{x}.
The benefit of using z is that the nodes of the wavefunction are equally spaced in z-space,
allowing for a computational better choice of choosing the constant step size during the integration.
"""
def __init__(
self,
z_min: float,
z_max: float,
dz: float,
) -> None:
"""Initialize the grid object.
Args:
z_min: The minimum value of the scaled dimensionless coordinate z = sqrt{x}.
z_max: The maximum value of the scaled dimensionless coordinate z = sqrt{x}.
dz: The step size of the grid in the scaled dimensionless coordinate z = sqrt{x}
(exactly one of dz or steps must be provided).
steps: The number of steps in the grid (exactly one of dz or steps must be provided).
"""
self._dz = dz
# put all grid points on a standard grid, i.e. [dz, 2*dz, 3*dz, ...]
# this is necessary to allow integration of two different wavefunctions
# Note: using np.arange((z_min // dz) * dz, z_max + dz / 2, dz)
# would lead to 'quite big' inprecisions (1e-10) between grid points of different grids,
# because of floating point errors
self._z_list: NDArray = np.arange(0, z_max + dz / 2, dz)[round(z_min / dz) :]
def __len__(self) -> int:
return self.steps
def __repr__(self) -> str:
return f"Grid({self.z_min}, {self.z_max}, dz={self.dz}, steps={self.steps})"
@property
def steps(self) -> int:
"""The number of steps in the grid."""
return len(self.z_list)
@property
def dz(self) -> float:
"""The step size of the grid in the scaled dimensionless coordinate z = sqrt{x}."""
return self._dz
@property
def z_min(self) -> float:
"""The minimum value of the scaled dimensionless coordinate z = sqrt{x}."""
return self.z_list[0] # type: ignore [no-any-return] # FIXME: numpy indexing
@property
def z_max(self) -> float:
"""The maximum value of the scaled dimensionless coordinate z = sqrt{x}."""
return self.z_list[-1] # type: ignore [no-any-return] # FIXME: numpy indexing
@property
def z_list(self) -> "NDArray":
"""The grid in the scaled dimensionless coordinate z = sqrt{x}.
In this coordinate the grid points are chosen equidistant,
because the nodes of the wavefunction are equally spaced in this coordinate.
"""
return self._z_list
@property
def x_min(self) -> float:
"""The minimum value of the dimensionless coordinate x = r/a_0."""
return self.z_min**2
@property
def x_max(self) -> float:
"""The maximum value of the dimensionless coordinate x = r/a_0."""
return self.z_max**2
@property
def x_list(self) -> "NDArray":
"""The grid in the dimensionless coordinate x = r/a_0."""
return self.z_list**2
[docs]
def set_grid_range(self, step_start: Optional[int] = None, step_stop: Optional[int] = None) -> None:
"""Restrict the grid to the range [step_start, step_stop]."""
if step_start is None:
step_start = 0
if step_stop is None:
step_stop = self.steps
self._z_list = self._z_list[step_start:step_stop]