This section contains the detailed API reference for the modules in the General Python Utilities library. The documentation below is generated automatically by Sphinx using autodoc, listing the classes, functions, and constants available in each module.
Most public APIs accept array-like vectors and matrices that are converted to the
active backend where possible. Shapes follow linear-algebra conventions, for
example A has shape (n,n) and b has shape (n,) or (n,k).
Dtype promotion follows backend rules; explicit float64 or complex128 is
recommended for ill-conditioned problems.
Stability depends on solver choice, conditioning, and preconditioning quality.
For reproducibility, set random seeds via algebra.ran_wrapper and keep backend
selection fixed in a run. NumPy and JAX results should agree up to floating-point
roundoff; small differences can appear due to kernel fusion and reduction order.
The module uses lazy imports to keep import time low.
This module provides utilities for importing and managing the linear algebra backend (NumPy/JAX),
random number generation, JIT compilation, and backend configuration for the general_python package.
from..algebra.utilsimportget_backend,backend_mgr,maybe_jitxp=backend_mgr.np# NumPy or JAX numpyrng=backend_mgr.default_rng@maybe_jitdefmy_func(x):returnxp.sum(x)
Return backend modules based on the provided specifier.
Delegates to the global backend_mgr.get_backend_modules.
Parameters:
backend_spec (str or module or None, optional) – Backend specifier (“numpy”, “jax”, np, jnp, “default”, None).
Defaults to the globally active backend.
random (bool, optional) – If True, include the random module/state. For JAX, also return a PRNG key.
For NumPy, returns a seeded RNG instance. Default is False.
For JAX backend, ensure you split the returned PRNG key before use.
seed (int, optional) – Seed for the random component. If None, uses the global default seed
to generate the component for this call. Providing a seed here creates
a new RNG/Key for this call based on that specific seed.
scipy (bool, optional) – If True, also return the associated SciPy module. Default is False.
Returns:
module or tuple – Requested backend components. See BackendManager.get_backend_modules docs.
**Example for using JAX backend with random number generation (****)
Prints backend configuration and library versions in a table.
Displays the active backend, available libraries, and their versions.
Also shows the active integer, float, and complex types.
backend_spec – Backend identifier. Can be:
- String: “numpy”, “np”, “jax”, “jnp”, “default”.
- Module: numpy or jax.numpy.
- None: Uses the manager’s active backend.
use_random – If True, include the random module/state. For JAX,
returns (jax.random, key). For NumPy, returns
(rng_instance, None).
seed – Seed for the random number generator. If None, uses
the manager’s default seed. If provided, creates a
new RNG/Key for this request, independent of the
manager’s default state.
use_scipy – If True, include the SciPy module.
Returns:
The requested module(s) as a single module or a tuple.
Format:
(main_module) or (main_module, random_part, scipy_module)
where random_part is (rng_instance, None) for numpy or (jrn_module, key) for jax.
Returns the globally configured default backend modules from the manager.
Uses the manager’s current active backend (self.name). If a specific seed
is provided, it generates a random component based on that seed for this
request, otherwise uses the manager’s default seed/key framework.
Parameters:
use_random – If True, include the random module/state.
seed – Optional seed for this specific request’s random component.
use_scipy – If True, include the SciPy module.
Returns:
The global default backend module(s). See get_backend_modules for format.
Abstract interfaces for backend-aware linear-system solvers.
The module defines shared types and base classes for algorithms that solve
linear systems of the form Ax=b. Implementations may operate on explicit
matrices, matrix-vector callables, or Fisher/Gram factors and can use optional
left preconditioning through callables of the form r->M^{-1}r.
The base interface is intentionally compatible with NumPy and JAX. Static
solver kernels are preferred for algorithm implementations because they are
easier to JIT-compile and reuse from configured solver instances.
Primarily defines the static interface solve that concrete algorithm
implementations (like CG, MINRES) must provide.
Also includes static helpers create_matvec_from_* to construct the
matrix-vector product function and an optional instance method solve_instance
for convenience when working with configured Solver objects.
Normally, one should focus on using the implementation provided
by the constructor-set function due to the optimized call. Nevertheless,
instance based setups are convenient for the less-time consuming tasks.
The solver can be initialized with a matrix A, a matrix-vector multiplication
function, or a Fisher matrix S. The solver can also be initialized with
a preconditioner M.
The solver can be used to solve the linear system
Ax = b or M^{-1}Ax = M^{-1}b.
Creates a matrix-vector multiplication function (MatVecFunc) based on a Fisher information inspired formulation.
This function constructs a custom matrix-vector product operator using the input arrays s and s_p.
It first computes the normalization constant n as the number of rows of s.
Depending on the flag create_full, it either:
Computes the full matrix as (s_p @ s) / n and passes it along with sigma to Solver.create_matvec_from_matrix_jax,
or
Constructs a matvec function that, for a given vector x, computes:
s_dot_x = dot(s, x)
sp_dot_s_dot_x = dot(s_p, s_dot_x)
The output as sp_dot_s_dot_x / n, with an additional term sigma * x if sigma is not None and non-zero.
The resulting matvec function is then compiled (presumably via JAX JIT) using Solver._compile_helper_jax.
:param s: A JAX array representing the first component used in constructing the operator.
:type s: Array
:param s_p: A JAX array representing the second component used alongside s.
:type s_p: Array
:param sigma: A scalar value to be added to the diagonal (identity) part of the operation.
Defaults to None.
Parameters:
create_full (Optional[bool], optional) – A flag to determine whether to create a full matrix operator using
Solver.create_matvec_from_matrix_jax. Defaults to False.
Returns:
A function that computes the matrix-vector product corresponding to the constructed operator.
This operator is compiled using Solver._compile_helper_jax.
Creates a matrix-vector multiplication function (matvec) based on the Fisher information matrix.
This function generates a matvec function that computes the product of a vector with a matrix
derived from the input arrays s and s_p. Optionally, a regularization term sigma can be
added to the diagonal of the matrix. The function can also return a full matrix-based matvec
if create_full is set to True.
:param s: A 2D array representing the first operand in the Fisher information matrix computation.
:type s: np.ndarray
:param s_p: A 2D array representing the second operand in the Fisher information matrix computation.
:type s_p: np.ndarray
:param sigma: A regularization parameter added to the diagonal of the matrix.
Defaults to None, which means no regularization is applied.
Parameters:
create_full (Optional[bool], optional) – If True, creates a matvec function based on the full matrix
s @ s_p / n. Defaults to False.
Returns:
A function that performs matrix-vector multiplication with the derived matrix.
Function implementing the matrix-vector product A @ x.
It must accept a vector of shape (N,) and return a vector of shape (N,).
Must be compatible with backend_module (NumPy or JAX).
bArray
Right-hand side vector of shape (N,). Must be a backend_module array.
x0Array
Initial guess vector of shape (N,). Must be a backend_module array.
Additional solver-specific keyword arguments (e.g., restart for GMRES).
returns:
A named tuple containing:
- x (Array): The computed solution vector of shape (N,).
- converged (bool): True if the solver reached the desired tolerance.
- iterations (int): The number of iterations performed.
- residual_norm (float): The norm of the final residual (||b - Ax||).
rtype:
SolverResult
raises NotImplementedError:
If a subclass hasn’t implemented this method.
raises SolverError:
If convergence fails catastrophically or inputs are invalid.
Sets up matvec and precond_apply based on instance configuration
(if provided during __init__) or arguments, then calls the static solve
method of this solver’s class. Stores the result in instance attributes.
Parameters:
b (Array) – Right-hand side vector.
x0 (Optional[Array]) – Initial guess. Defaults to zeros.
tol (Optional[float]) – Tolerance override. Uses instance default if None.
maxiter (Optional[int]) – Max iterations override. Uses instance default if None.
sigma (Optional[float]) – Regularization for matvec creation if matvec
is not already defined for the instance. Uses
instance default _conf_sigma if None.
**kwargs – Additional arguments passed directly to the static solve.
For real inputs, r = sqrt(a^2 + b^2) is nonnegative.
For complex inputs, r preserves the phase of a (if b==0) or b (if a==0),
and the reflectors are computed in a stable manner.
Parameters:
a (scalar (real or complex)) – The first element of the two-vector [a; b].
b (scalar (real or complex)) – The second element of the two-vector [a; b].
backend (str, optional (default "default")) – Specifies which backend to use. If set to “jax”, the function uses
jax.numpy and is jitted for speed.
Returns:
(c, s, r) (tuple of scalars) –
The computed reflection parameters satisfying:
c = a / r and s = b / r,
with r = sqrt(a^2 + b^2) for real numbers (or the appropriately phased value for complex).
Numerical stability
——————-
This function avoids overflow and underflow by scaling by the larger magnitude component
(either |a| or |b|). This ensures that the intermediate calculations of tau
and the hypotenuse do not exceed floating-point range limits unnecessarily.
Solver module for various linear algebra solvers.
Initialization file for the solvers module. Exports solver classes,
the SolverType enum, and the choose_solver factory function.
—————————————————————-
File : general_python/algebra/solvers/__init__.py
Author : Maksymilian Kliczkowski
License : MIT
Description : This module provides a factory function to choose and instantiate
different solver types based on user input. It includes various
solver implementations (direct, iterative, etc.) and allows for
customization through keyword arguments. The module also includes
a utility function to generate test matrix-vector pairs for
testing purposes. The solvers are designed to work with different
numerical backends (e.g., NumPy, SciPy, JAX) and support various
data types. The module is part of a larger algebra library and
aims to provide a flexible and extensible framework for solving
linear algebra problems.
Primarily defines the static interface solve that concrete algorithm
implementations (like CG, MINRES) must provide.
Also includes static helpers create_matvec_from_* to construct the
matrix-vector product function and an optional instance method solve_instance
for convenience when working with configured Solver objects.
Normally, one should focus on using the implementation provided
by the constructor-set function due to the optimized call. Nevertheless,
instance based setups are convenient for the less-time consuming tasks.
The solver can be initialized with a matrix A, a matrix-vector multiplication
function, or a Fisher matrix S. The solver can also be initialized with
a preconditioner M.
The solver can be used to solve the linear system
Ax = b or M^{-1}Ax = M^{-1}b.
Creates a matrix-vector multiplication function (MatVecFunc) based on a Fisher information inspired formulation.
This function constructs a custom matrix-vector product operator using the input arrays s and s_p.
It first computes the normalization constant n as the number of rows of s.
Depending on the flag create_full, it either:
Computes the full matrix as (s_p @ s) / n and passes it along with sigma to Solver.create_matvec_from_matrix_jax,
or
Constructs a matvec function that, for a given vector x, computes:
s_dot_x = dot(s, x)
sp_dot_s_dot_x = dot(s_p, s_dot_x)
The output as sp_dot_s_dot_x / n, with an additional term sigma * x if sigma is not None and non-zero.
The resulting matvec function is then compiled (presumably via JAX JIT) using Solver._compile_helper_jax.
:param s: A JAX array representing the first component used in constructing the operator.
:type s: Array
:param s_p: A JAX array representing the second component used alongside s.
:type s_p: Array
:param sigma: A scalar value to be added to the diagonal (identity) part of the operation.
Defaults to None.
Parameters:
create_full (Optional[bool], optional) – A flag to determine whether to create a full matrix operator using
Solver.create_matvec_from_matrix_jax. Defaults to False.
Returns:
A function that computes the matrix-vector product corresponding to the constructed operator.
This operator is compiled using Solver._compile_helper_jax.
Creates a matrix-vector multiplication function (matvec) based on the Fisher information matrix.
This function generates a matvec function that computes the product of a vector with a matrix
derived from the input arrays s and s_p. Optionally, a regularization term sigma can be
added to the diagonal of the matrix. The function can also return a full matrix-based matvec
if create_full is set to True.
:param s: A 2D array representing the first operand in the Fisher information matrix computation.
:type s: np.ndarray
:param s_p: A 2D array representing the second operand in the Fisher information matrix computation.
:type s_p: np.ndarray
:param sigma: A regularization parameter added to the diagonal of the matrix.
Defaults to None, which means no regularization is applied.
Parameters:
create_full (Optional[bool], optional) – If True, creates a matvec function based on the full matrix
s @ s_p / n. Defaults to False.
Returns:
A function that performs matrix-vector multiplication with the derived matrix.
Function implementing the matrix-vector product A @ x.
It must accept a vector of shape (N,) and return a vector of shape (N,).
Must be compatible with backend_module (NumPy or JAX).
bArray
Right-hand side vector of shape (N,). Must be a backend_module array.
x0Array
Initial guess vector of shape (N,). Must be a backend_module array.
Additional solver-specific keyword arguments (e.g., restart for GMRES).
returns:
A named tuple containing:
- x (Array): The computed solution vector of shape (N,).
- converged (bool): True if the solver reached the desired tolerance.
- iterations (int): The number of iterations performed.
- residual_norm (float): The norm of the final residual (||b - Ax||).
rtype:
SolverResult
raises NotImplementedError:
If a subclass hasn’t implemented this method.
raises SolverError:
If convergence fails catastrophically or inputs are invalid.
Sets up matvec and precond_apply based on instance configuration
(if provided during __init__) or arguments, then calls the static solve
method of this solver’s class. Stores the result in instance attributes.
Parameters:
b (Array) – Right-hand side vector.
x0 (Optional[Array]) – Initial guess. Defaults to zeros.
tol (Optional[float]) – Tolerance override. Uses instance default if None.
maxiter (Optional[int]) – Max iterations override. Uses instance default if None.
sigma (Optional[float]) – Regularization for matvec creation if matvec
is not already defined for the instance. Uses
instance default _conf_sigma if None.
**kwargs – Additional arguments passed directly to the static solve.
Factory function to select and instantiate a solver based on identifier.
Uses lazy loading to import specific solver classes only when requested.
Parameters:
solver_id (Union[str, int, SolverType, Type[Solver]]) – Identifier for the solver. Can be a string name, integer code,
SolverType enum, or a Solver subclass.
backend (str, optional) – Numerical backend to use (e.g., “numpy”, “scipy”, “jax”). Default is “default”.
sigma (Optional[float], optional) – Shift parameter for the solver, if applicable. Default is None.
is_gram (bool, optional) – Whether to treat the problem as a Gram matrix problem. Default is False.
default_precond (Optional[Preconditioner], optional) – Default preconditioner to use if none is specified. Default is None.
**kwargs – Additional keyword arguments to pass to the solver constructor.
Initial-value ODE integrators with NumPy and optional JAX backends.
The module provides a compact interface for explicit one-step methods used in
simulation loops where the caller owns time evolution state. Integrators return
(y_next,dt,info) and do not impose a global solver object or trajectory
storage.
Create a Runge-Kutta method instance from a specified order.
:param order: The order of the Runge-Kutta method. Supported values are 1 (Euler), 2 (RK2), and 4 (RK4).
:type order: int
:param dt: The time step size. Defaults to 1e-3.
:type dt: float, optional
:param backend: The computational backend to use (e.g., ‘numpy’). Defaults to ‘numpy’.
:type backend: str, optional
Returns:
An instance of the class initialized with the appropriate Butcher tableau for the specified order.
Return type:
cls
Raises:
ValueError – If the specified order is not supported.
The right-hand side prefactor is treated as part of the ODE,
dy/dt=af(t,y), not as a final post-processing multiplier.
This matters for complex real-time evolution in TDVP, where the stage
states must be evaluated at
y_i = y_n + h a sum_{j < i} a_{ij} k_j
in order for RK2 / RK4 to integrate the same physical time variable as
Euler and exact diagonalization.
This module contains the implementation of preconditioners for iterative solvers
of linear systems Ax = b. Preconditioners transform the system into an
equivalent one,
M^{-1}Ax = M^{-1}b (left preconditioning)
or
A M^{-1} y = b, x = M^{-1}y (right preconditioning),
where M is the preconditioner matrix.
The goal is for the transformed system to have more
favorable spectral properties (e.g., eigenvalues clustered around 1, lower
condition number), leading to faster convergence of iterative methods like CG,
MINRES, GMRES. The matrix M should approximate A in some sense, while the
operation M^{-1}r should be computationally inexpensive.
Abstract base class for preconditioners M used in iterative solvers.
Provides a framework for setting up the preconditioner based on a matrix A
(or its factors S, Sp for Gram matrices) and applying the inverse operation
M^{-1}r efficiently. Supports different computational backends (NumPy, JAX).
Returns the potentially JIT-compiled function apply(r).
Uses precomputed data stored by the last call to set().
Raises RuntimeError if called before set() to match test expectations.
Returns a potentially JIT-compiled function apply_mat(r, A, sigma, **override_kwargs)
that computes preconditioner data from A and applies it on the fly.
Returns a potentially JIT-compiled function apply_gram(r, S, Sp, sigma, **override_kwargs)
that computes preconditioner data from S, Sp and applies it on the fly.
Params:
**default_setup_kwargs: Default keyword arguments for the setup kernel.
Returns:
The compiled function.
Return type:
Callable
Note
For JAX compatibility, kwargs are frozen at function creation time.
The returned function does NOT accept runtime kwargs to avoid
dictionary operations inside JIT-traced code.
Sets up the preconditioner using the provided matrix A and optional parameters.
This method computes the preconditioner data and prepares the apply function.
Params:
a (Array):
The matrix to be used for setting up the preconditioner.
sigma (float, optional):
The regularization parameter. Defaults to 0.0.
ap (Optional[Array], optional):
An optional second matrix for Gram setup. Defaults to None.
backend (Optional[str], optional):
The backend to use for computations. Defaults to None.
Jacobi (Diagonal) Preconditioner. M = diag(A + sigma*I).
Uses the diagonal of the (potentially regularized) matrix A as the
preconditioner M. Applying the inverse M^{-1}r involves element-wise
division by the diagonal entries.
Math:
M = diag(A) + sigma*I = D + sigma*I
M^{-1}r = (D + sigma*I)^{-1} r = [1 / (A_ii + sigma)] * r_i
References
Saad, Y. (2003). Iterative Methods for Sparse Linear Systems (2nd ed.). SIAM. Chapter 10.
Cholesky Preconditioner using complete Cholesky decomposition.
Suitable for symmetric positive-definite matrices A. The preconditioner M
is defined by the Cholesky factorization of the regularized matrix:
M = L @ L.T (or L @ L.conj().T for complex), where L is the lower
Cholesky factor of (A + sigma*I).
Applying the inverse M^{-1}r involves solving two triangular systems:
1. Solve L y = r for y (forward substitution)
2. Solve L.T z = y for z (backward substitution)
(or L.conj().T z = y for complex)
Note
This performs a complete Cholesky factorization. For incomplete
Cholesky (suitable for large sparse matrices), specialized routines
(often from sparse linear algebra libraries) are required.
References
Golub, G. H., & Van Loan, C. F. (2013). Matrix Computations (4th ed.). JHU Press. Chapter 4.
Saad, Y. (2003). Iterative Methods for Sparse Linear Systems (2nd ed.). SIAM. Chapter 10.
Suitable for symmetric matrices A, particularly those that are positive definite.
It involves forward and backward Gauss-Seidel-like sweeps with a relaxation parameter omega.
Math:
Let A = D + L + U, where D is diag(A), L is strictly lower part, U is strictly upper part.
The SSOR preconditioner matrix M is defined implicitly by its application M^{-1}r:
1. Solve (D/omega + L) y = r (Forward sweep)
2. Solve (D/omega + U) z = D y / omega (Backward sweep - check formulation)
A common formulation for the inverse application M^{-1}r is:
M^{-1} = omega * (2 - omega) * (D/omega + U)^{-1} D (D/omega + L)^{-1}
Applying M^{-1}r involves solving the two triangular systems mentioned above.
The choice of omega (0 < omega < 2) is crucial for performance. omega=1 gives
Symmetric Gauss-Seidel (SGS).
References
Saad, Y. (2003). Iterative Methods for Sparse Linear Systems (2nd ed.). SIAM. Chapter 10.
Axelsson, O. (1996). Iterative Solution Methods. Cambridge University Press.
Incomplete Cholesky Preconditioner (Approximation using ILU(0)).
Suitable for large, sparse, symmetric positive-definite matrices A.
This implementation uses SciPy’s sparse ILU(0) decomposition as a proxy
for IC(0), as a direct IC(0) is not available in SciPy/NumPy/JAX.
ILU(0) computes factors L and U such that L @ U approximates A,
maintaining the sparsity pattern of A (zero fill-in).
Applying the inverse M^{-1}r involves solving L y = P r and U z = y, where P
is a permutation matrix handled internally by the ILU object.
Note:
- Requires SciPy and the NumPy backend. Not JAX compatible.
- Designed for sparse matrices. Dense input matrices will be converted.
- Uses ILU(0) factorization as a proxy for IC(0).
Args for set:
fill_factor (float): See scipy.sparse.linalg.spilu. Default 1.
drop_tol (float): See scipy.sparse.linalg.spilu. Default None.
References
Saad, Y. (2003). Iterative Methods for Sparse Linear Systems (2nd ed.). SIAM. Chapter 10 (discusses both IC and ILU).
SciPy documentation for scipy.sparse.linalg.spilu.
Suitable for large, sparse, non-symmetric matrices A.
This implementation uses SciPy’s sparse ILU(0) decomposition (scipy.sparse.linalg.spilu).
ILU(0) computes factors L and U such that L @ U approximates A,
maintaining the sparsity pattern of A (zero fill-in).
Applying the inverse M^{-1}r involves solving L y = P r and U z = y, where P
is a permutation matrix handled internally by the ILU object.
Note:
- Requires SciPy and the NumPy backend. Not JAX compatible.
- Designed for sparse matrices. Dense input matrices will be converted.
Args for set:
fill_factor (float): See scipy.sparse.linalg.spilu. Default 1.
drop_tol (float): See scipy.sparse.linalg.spilu. Default None.
References
Saad, Y. (2003). Iterative Methods for Sparse Linear Systems (2nd ed.). SIAM. Chapter 10.
SciPy documentation for scipy.sparse.linalg.spilu.
Shared utilities for IO, plotting, logging, and runtime helpers.
The common package contains cross-cutting infrastructure used by all major
subpackages, including directory management, plotting, HDF5 helpers, timers,
and low-level binary utilities.
Utilities are intentionally lightweight: they accept standard Python objects or
NumPy-compatible arrays and return plain Python or NumPy outputs where possible.
Plotting helpers expect Matplotlib-compatible inputs.
Most helpers are deterministic for fixed inputs. Logging or timing output may
vary with runtime scheduling. Binary helper functions operate on integer bit
patterns and are deterministic by construction.
Submodules are loaded lazily on first access to reduce import overhead.
List files (not directories) in this directory.
- include_empty : if False, skip files of size zero.
- filters : a list of callables Path->bool; all must pass.
- sort_key : key function for sorting.
include_empty (bool) – if False, skip empty directories. If True, include all directories.
This checks only if the directory has any entries, not if they are files or directories.
filters (list of callables Path -> bool) – A list of callables; all must return True for a directory to be included.
sort_key (callable, optional) – Key function for sorting the results.
Publication-quality plotting utilities for scientific computing.
This class provides static methods for creating, customizing, and saving
Matplotlib figures suitable for scientific journals (Nature, Science, PRL, etc.).
All methods are @staticmethod, so you can use them without instantiation:
Print help information about available plotting methods.
Parameters:
topic (str, optional) – Specific topic to get help on. Options:
- ‘plot’: Basic plotting methods
- ‘axis’: Axis configuration
- ‘color’: Colors and colorbars
- ‘layout’: Subplots and grids
- ‘save’: Saving figures
- None: Print overview of all topics
This method extends Python str.format-style placeholders with simple
math filters for scientific labels.
Parameters:
label (str) – Template string. Standard format fields are supported, e.g.
{J:.3g}, and can be combined with filters, e.g.
{point|vec:.2f}, {kpoint|sym}.
Greek-name values (for example Gamma or omega) are
automatically rendered as LaTeX variables.
auto_wrap (bool, default=True) – If True, wrap the final string with $...$ when no dollar sign is
present in the rendered output.
escape_text (bool, default=True) – If True, plain substituted strings are LaTeX-escaped by default.
Use |raw or |tex to bypass escaping for a specific field.
**values (Any) – Values used by template fields.
filters (Supported)
-----------------
tex (raw /) – Insert value as-is (no escaping).
sym – Force symbol conversion for common Greek names
(e.g. Gamma -> \Gamma).
num – Numeric formatting helper. Uses format spec if provided, otherwise
.6g.
vec – Render iterable as \left(v_1,v_2,...\right).
set – Render iterable as \left\{v_1,v_2,...\right\}.
palette (str or list of colors, default='tableau') – Named palette string (see palette()) or an explicit list of
any Matplotlib-compatible color specs.
Examples
>>> Plotter.set_color_cycle(ax,'wong')>>> ax.plot(x1,y1)# first wong color>>> ax.plot(x2,y2)# second wong color
Each parameter shifts the corresponding channel by a fraction of the
remaining headroom, so operations compose gracefully and values are
always clamped to [0, 1].
Parameters:
color (color spec) – Any Matplotlib-compatible color.
cmap (str or Colormap, default='viridis') – Source colormap.
vmin (float, default 0.0, 1.0) – Fraction range to sample from (allows using only a sub-range of the
colormap, e.g. vmin=0.1,vmax=0.9 avoids the near-white ends of
sequential maps).
vmax (float, default 0.0, 1.0) – Fraction range to sample from (allows using only a sub-range of the
colormap, e.g. vmin=0.1,vmax=0.9 avoids the near-white ends of
sequential maps).
as_hex (bool, default=False) – If True, return hex strings instead of RGBA tuples.
Get a dictionary with color properties for matplotlib patches.
:param - color [str or tuple]: Color to use, can be a named color or an RGB tuple.
:param - alpha [float]: Transparency level (0 to 1).
:param - edgecolor [tuple]: Edge color as an RGB tuple.
:param - facecolor [tuple]: Face color as an RGB tuple.
If array-like: a new ScalarMappable is built from cmap/norm (and scale, vmin, vmax).
If ScalarMappable: it is used directly. vmin/vmax update its clim; norm is taken from it
when not provided. Note: in this case discrete/boundaries resampling is not applied.
cmap (str | Colormap, default='viridis') – Colormap name or object. If mappable is a ScalarMappable, its cmap is used unless cmap
is explicitly different from the default and a new mappable is constructed (array-like path).
norm (matplotlib.colors.Normalize, optional) – Normalization to map data to 0-1. Ignored if mappable is ScalarMappable and norm is None
(then the mappable’s norm is used).
vmin (float, optional) – Data limits. When scale=’log’, non-positive vmin is clamped internally.
vmax (float, optional) – Data limits. When scale=’log’, non-positive vmin is clamped internally.
scale ({'linear', 'log', 'symlog'}, default='linear') – Creates a suitable Normalize when mappable is array-like and norm is None.
- ‘linear’ -> Normalize
- ‘log’ -> LogNorm (vmin<=0 clamped to ~1e-10)
- ‘symlog’ -> SymLogNorm with linthresh=0.1
extend ({'neither','both','min','max','neutral'}, default='neutral') – Colorbar extension behavior. Standard Matplotlib values are ‘neither’, ‘both’, ‘min’, ‘max’.
‘neutral’ is treated as a pass-through here and may behave like ‘neither’ depending on Matplotlib.
format (str | matplotlib.ticker.Formatter, optional) – Tick formatting. If str (e.g., ‘%.2e’), uses FormatStrFormatter.
discrete (bool | int, default=False) – Discretize colormap when building from array-like:
- True -> 10 bins
- int N -> N bins
Ignored when mappable is a ScalarMappable.
boundaries (list[float], optional) – Discrete bin edges. Enables BoundaryNorm and passes boundaries to fig.colorbar
(default spacing=’proportional’, overridable via kwargs[‘spacing’]).
invert (bool, default=False) – If True, invert the colorbar axis direction.
remove_pdf_lines (bool, default=True) – Set solids edgecolor to ‘face’ to avoid white hairlines in vector exports (PDF/SVG).
**kwargs – Additional arguments forwarded to fig.colorbar, e.g.:
- alpha, spacing (‘uniform’|’proportional’), fraction, pad, shrink, aspect, drawedges, etc.
Parameters:
- values (array-like): The values to map to colors.
- cmap (str, optional): The colormap to use. Defaults to ‘PuBu’.
- elsecolor (str, optional): The color to use if there is only one value. Defaults to ‘blue’.
- get_mappable (bool, optional): If True, also return a ScalarMappable as
the 4th item, ready to pass into Plotter.add_colorbar(…, mappable=…).
return_mappable (bool, optional): Alias for get_mappable.
Returns:
- getcolor (function): A function that maps a value to a color.
- colors (Colormap): The colormap object.
- norm (Normalize): The normalization object.
- mappable (ScalarMappable, optional): Returned when get_mappable=True
Apply a colormap to the given data and plot it on the provided axis.
Parameters:
- ax (object): The axis object to plot on.
- data (array-like): The data to plot.
- cmap (str, optional): The colormap to use. Defaults to ‘PuBu’.
- colorbar (bool, optional): Whether to add a colorbar. Defaults to True.
Make an annotation on the plot.
- ax : axis to annotate on
- elem : annotation string
- x : x coordinate (ignored if xycoords=’best’)
- y : y coordinate (ignored if xycoords=’best’)
- fontsize : fontsize of the annotation
- xycoords : how to interpret the coordinates (from MPL), or ‘best’ to find the best corner
- cond : condition to make the annotation
Make an annotation on the plot.
- ax : axis to annotate on
- start_T : start text
- end_T : end text
- xystart : x coordinate start
- xystart_T : x coordinate start text
- xyend : x coordinate end
- xyend_T : x coordinate end text
- arrowprops: properties of the arrow
- startcolor: color of the arrow at the start
- endcolor : color of the arrow in the end
- kwargs : additional arguments for annotation
Plots the fitting function provided by the user on a given axis using
the **kwargs provider afterwards.
- ax : axis to annotate on
- funct : function to use for the fitting
- x : arguments to the function
Plot a scalar field sampled on irregular planar points using triangulation.
This helper is meant for 2D scattered data where imshow is not
appropriate because the samples do not lie on a regular rectangular grid.
Matplotlib first builds a triangulation of the point cloud and then
interpolates values inside each triangle.
Typical use cases:
- real-space lattice-site scalar fields
- Brillouin-zone data on irregular planar k-point sets
- any scattered 2D measurement data
Parameters:
ax (matplotlib.axes.Axes) – Target axis.
points (array-like) – Sample positions shaped like (N, 2) or (N, D) with D >= 2.
Only the first two Cartesian components are used.
values (array-like) – Scalar values of length N.
triangles (array-like, optional) – Explicit connectivity passed to matplotlib.tri.Triangulation.
mask (array-like of bool, optional) – Triangle mask. True hides the corresponding triangle.
Sets tickparams to the desired ones.
- ax : axis to use
- labelsize : fontsize
- left : whether to show the left side
- right : whether to show the right side
- top : whether to show the top side
- bottom : whether to show the bottom side
- xticks : list of xticks
- yticks : list of yticks
Comprehensive axis configuration method for publication-quality plots.
This method provides centralized control over all major axis properties
with sensible defaults and advanced options for fine-tuning. It integrates
with other Plotter methods for a cohesive styling experience.
Parameters:
ax (matplotlib.axes.Axes) – The axis object to modify.
which ({'both', 'x', 'y'}, default='both') – Specifies which axes to update. Allows independent configuration
of x and y axes.
Titles** (**Labels and)
xlabel (str, optional) – Axis labels. Set to ‘’ to hide labels while maintaining formatting.
ylabel (str, optional) – Axis labels. Set to ‘’ to hide labels while maintaining formatting.
Add labels to lines with a given slope. Uses labelLines package.
:param - ax: Matplotlib axis object.
:param - align: Align the label with the slope of the line.
:param - xvals: The x values to place the labels at.
:param - yoffsets: The y offsets for the labels.
:param - zorder: The zorder of the labels.
Remove tick labels (and optionally tick marks) from the axis.
Useful for creating clean shared-axis plots where inner panels
don’t need redundant tick labels.
Parameters:
ax (matplotlib.axes.Axes) – The axes to modify.
xticks (bool, default=False) – If True, REMOVE x-tick labels. If False, KEEP them.
yticks (bool, default=False) – If True, REMOVE y-tick labels. If False, KEEP them.
xlabel (bool, default=False) – If True, also REMOVE the x-axis label.
ylabel (bool, default=False) – If True, also REMOVE the y-axis label.
remove_labels_only (bool, default=True) – If True, only remove the text labels, keeping tick marks visible.
If False, remove both the tick marks and labels.
Examples
# Remove x-tick labels for stacked plots with shared x-axis
>>> for ax in axes[:-1]: # All except bottom
… Plotter.unset_ticks(ax, xticks=True, xlabel=True)
# Remove tick marks AND labels (completely clean)
>>> Plotter.unset_ticks(ax, xticks=True, yticks=True, remove_labels_only=False)
# Remove y-ticks and y-axis label for side-by-side shared y-axis
>>> for ax in axes[1:]: # All except leftmost
… Plotter.unset_ticks(ax, yticks=True, ylabel=True)
Notes
This function is commonly used in combination with sharex/sharey in
multi-panel figures to avoid redundant labels.
Sets the formatter for the given axis on the plot.
:param ax: The axis object on which to set the formatter.
:type ax: object
:param formater: The format string for the axis labels. Defaults to “%.1e”.
:type formater: str, optional
:param axis: The axis on which to set the formatter. Defaults to ‘xy’.
:type axis: str, optional
Sets the formatter for the given axis on the plot.
:param ax: The axis object on which to set the formatter.
:type ax: object
:param axis: The axis on which to set the formatter. Defaults to ‘xy’.
:type axis: str, optional
wspace (float, optional) – Width space between columns (0.0 to 1.0, fraction of average axis width).
Recommended: 0.2-0.4 for labels, 0.05-0.1 for tight layouts.
hspace (float, optional) – Height space between rows (0.0 to 1.0, fraction of average axis height).
Recommended: 0.2-0.4 for titles, 0.05-0.1 for tight layouts.
width_ratios (list of float, optional) – Relative widths of columns. E.g., [2, 1, 1] makes first column 2x wider.
Length must equal ncols.
height_ratios (list of float, optional) – Relative heights of rows. E.g., [1, 2] makes second row 2x taller.
Length must equal nrows.
ax_sub (SubplotSpec, optional) – If provided, creates a nested GridSpec within this subplot.
Use for complex layouts with grids inside grids.
left (float, optional) – Figure margins (0.0 to 1.0). Controls space for labels.
right (float, optional) – Figure margins (0.0 to 1.0). Controls space for labels.
top (float, optional) – Figure margins (0.0 to 1.0). Controls space for labels.
bottom (float, optional) – Figure margins (0.0 to 1.0). Controls space for labels.
**kwargs – Additional arguments passed to GridSpec.
Returns:
The grid specification object.
Return type:
GridSpec or GridSpecFromSubplotSpec
Examples
Basic 2x3 grid:
>>> fig=plt.figure(figsize=(12,8))>>> gs=Plotter.get_grid(2,3,wspace=0.3,hspace=0.4)>>> ax0=fig.add_subplot(gs[0,0])# Row 0, Col 0>>> ax1=fig.add_subplot(gs[0,1:])# Row 0, Cols 1-2 (span)>>> ax2=fig.add_subplot(gs[1,:])# Row 1, all columns (span)
Unequal widths (main panel + sidebar):
>>> gs=Plotter.get_grid(1,2,width_ratios=[3,1])>>> # First column is 3x wider than second
Create a subplot from a GridSpec at the specified index.
Parameters:
gs (GridSpec) – The GridSpec object.
fig (matplotlib.figure.Figure) – The figure to add the subplot to.
index (int, tuple, or slice) – Position in the grid. Can be:
- int: Linear index (0, 1, 2, …)
- tuple: (row, col) for single cell
- slice/tuple with slices: For spanning multiple cells
sharex (Axes, optional) – Share axis with another subplot. Use for aligned multi-panel figures.
sharey (Axes, optional) – Share axis with another subplot. Use for aligned multi-panel figures.
**kwargs – Additional arguments passed to fig.add_subplot.
Returns:
The created subplot.
Return type:
matplotlib.axes.Axes
Examples
Single cell by linear index:
ax0=Plotter.get_grid_subplot(gs,fig,0)# First cellax1=Plotter.get_grid_subplot(gs,fig,1)# Second cell
Single cell by (row, col):
ax=fig.add_subplot(gs[1,2])# Row 1, Col 2
Span multiple cells:
ax_wide=fig.add_subplot(gs[0,:])# Entire first rowax_tall=fig.add_subplot(gs[:,0])# Entire first columnax_block=fig.add_subplot(gs[0:2,1:3])# 2x2 block
Shared axes (for aligned panels):
ax0=Plotter.get_grid_subplot(gs,fig,0)ax1=Plotter.get_grid_subplot(gs,fig,1,sharex=ax0)ax2=Plotter.get_grid_subplot(gs,fig,2,sharex=ax0,sharey=ax0)# ax1 and ax2 share x-axis with ax0; ax2 also shares y-axis
>>> gmap=Plotter.get_grid_map(2,3)>>> gmap['by_letter']['c']# Get index for panel 'c'2>>> gmap['by_index'][4]# Get (row, col) for index 4(1, 1)>>> gmap['grid']# 2D layout[[0, 1, 2], [3, 4, 5]]
Configure axis visibility, spines, ticks, and labels in one call.
This is a convenience function for common axis customizations.
Parameters:
ax (matplotlib.axes.Axes) – The axis to configure.
visible (bool, default=True) – If False, hide the entire axis (ax.axis(‘off’)).
spines (bool, dict, or str, default=True) – Control spine visibility:
- True: Show all spines
- False: Hide all spines
- ‘left’: Hide all except left
- ‘bottom’: Hide all except bottom
- ‘minimal’: Hide top and right (Nature-style)
- dict: {‘top’: False, ‘right’: False, …}
ticks (bool, dict, or str, default=True) – Control tick visibility:
- True/False: Show/hide all ticks
- ‘x’/’y’: Show only x/y ticks
- dict: {‘x’: True, ‘y’: False}
tick_labels (bool, dict, or str, default=True) – Control tick label visibility (same format as ticks).
Disable axis components for clean images/heatmaps.
Parameters:
ax (matplotlib.axes.Axes) – The axis to modify.
which (str, default='both') – What to disable:
- ‘both’: Disable x and y (full axis off)
- ‘x’: Disable x-axis only
- ‘y’: Disable y-axis only
- ‘labels’: Keep ticks but hide labels
- ‘ticks’: Keep labels but hide ticks
- ‘spines’: Hide all spines
Sets the legend with a preferred style for publication-quality plots.
Parameters:
- ax : Axis to which the legend will be added.
- fontsize : Font size of the legend labels.
- frameon : Whether to draw a frame around the legend.
- loc : Location of the legend (‘best’, ‘upper right’, etc.).
- alignment : Text alignment (‘left’, ‘center’, ‘right’).
- markerfirst : Whether the marker or label appears first in the legend.
- framealpha : Transparency of the legend frame (1.0 is opaque).
- reverse : Reverse the order of legend items.
- style : Predefined style for the legend (‘minimal’, ‘boxed’, etc.).
- labelspacing : Vertical space between legend entries.
- handlelength : Length of the legend markers.
- handletextpad : Space between legend markers and text.
- borderpad : Padding inside the legend box.
- columnspacing : Spacing between legend columns.
- ncol : Number of columns in the legend.
- kwargs : Additional arguments passed to ax.legend().
Set the legend with custom conditions for the labels
- ax : axis to use
- conditions: list of conditions
- fontsize : fontsize
- frameon : frame on or off
- loc : location of the legend
- alignment : alignment of the legend
- markerfirst: marker first or not
- framealpha: alpha of the frame
axes (AxesList or matplotlib.axes.Axes) – AxesList wrapper (list-compatible, grid-aware, named-panel
access). Returns single axis only when single_if_1=True.
Save figure to a specific directory.
- directory : directory to save the file
- filename : name of the file
- format : format of the file
- dpi : dpi of the file
- adjust : adjust the figure
staticsavefig(directory, filename, format, dpi, adjust, fig=None, **kwargs)[source]
Alias for save_fig() with the historical lowercase name.
Based on the specified keys, load the dictionaries from the json files
The keys are the names of the files as well!
staticsingleColumnData(directory:str, fileName:str, y, typ='.npy')[source]
Stores the values as a single vector
statictwoColumnsData(directory:str, fileName:str, x, y, typ='.npy')[source]
Stores the x, y vectors in 2D form (multiple rows and two columns)
staticmatrixData(directory:str, fileName:str, x, y, typ='.npy')[source]
Stores the x, y vectors in matrix form (appending single column
at start for x values)
staticapp_df(df, colname:str, y, fill_value=nan)[source]
Appends the data to the dataframe.
Parameters:
- df (pd.DataFrame): The dataframe to append data to.
- colname (str): The column name to append data under.
- y (array-like): The data to append.
- fill_value: The value to use for filling if resizing is needed.
DataHandler class provides static methods for handling and processing data arrays, including filtering,
interpolating, aggregating, concatenating, and cutting matrices based on specific criteria.
.. method:: _filter_typical_values(current_x, current_y, typical, threshold=1.0) -> tuple
Initializes and combines arrays from given lists. If the typical flag is set to True,
it filters the combined arrays to include only elements where the values in y_combined
are less than the threshold.
Concatenates and averages y values across multiple histograms.
:param y_list : List of y matrices (each one corresponding to a realization).
:param x_list : List of x vectors (each one corresponding to a realization).
:param typical : If True, filter y values less than 1.0.
:param use_interpolation: If True, interpolate y values for non-matching bins.
If False, aggregate only exact matches and append unique bins.
:param threshold : The threshold value for filtering y values (default: 1.0).
:returns : Combined y values and x bins after averaging.
Cut off the slices (along any specified axis) in matrix M where all elements are close to zero.
If a 1D vector is provided, it returns the vector unless all elements are close to zero,
in which case it returns an empty array.
Parameters:
- M (numpy.ndarray) : The input matrix or vector.
- axis (int) : The axis along which to check for zero elements.
For example, 0 for rows, 1 for columns, etc.
Ignored if M is a 1D vector.
tol (float) : The tolerance for considering elements as zero.
check_limit (int) : The maximum number of elements along the axis to check for zeros.
Returns:
- numpy.ndarray: The resulting matrix after removing slices (along the specified axis)
that are close to zero, or the vector after removing if all elements are close to zero.
Cut off the rows or columns in matrix M where the first check_limit elements are all below a threshold.
Parameters:
- M (numpy.ndarray): The input matrix.
- axis (int): The axis along which to check for elements below the threshold (0 for rows, 1 for columns).
- threshold (float): The threshold value.
- check_limit (int, optional): The number of elements to check from each row or column.
Returns:
- numpy.ndarray: The resulting matrix after removing rows or columns where the first check_limit elements are below the threshold.
Yields data for a specific key from a list of already loaded HDF5 data dictionaries.
Each dictionary in loaded_hdf5_data_list is expected to be the output of ‘load_file_data’.
Parameters:
loaded_hdf5_data_list – List of dictionaries, where each dict contains data from an HDF5 file.
key – The dataset key to extract from each dictionary.
Yields:
numpy.ndarray – The dataset corresponding to the key.
filename – Name of the HDF5 file (extension .h5 or .hdf5 will be ensured).
data_to_save – Data to save. Can be a single np.ndarray, a list of np.ndarrays,
or a dictionary {name: np.ndarray}.
target_shape – If specified, datasets will be reshaped to this shape before saving.
dataset_names_config: Names for datasets if ‘data_to_save’ is a list/ndarray.
If a string, used as a prefix. If a list, used as names.
overwrite – If True (default), overwrites the file if it exists.
Updates only provided datasets in an HDF5 file and keeps all other datasets unchanged.
Example
Existing file has ‘/a’ and ‘/b’.
Calling with data_to_update={‘/a’: new_data} updates only ‘/a’ and leaves ‘/b’ intact.
Parameters:
directory – Directory where the HDF5 file is located.
filename – Name of the HDF5 file.
data_to_update – Data to update. Can be dict {name: array}, list of arrays, or single ndarray.
target_shape – If provided, each updated dataset is reshaped before writing.
dataset_names_config – Names for datasets when data_to_update is list/ndarray.
create_if_missing – If True, creates a missing file and/or missing dataset paths as needed.
If False, update only existing datasets in an existing file.
filename – Name of the HDF5 file (extension .h5 or .hdf5 will be ensured).
data_to_save – Data to save. Can be a single np.ndarray, a list of np.ndarrays,
or a dictionary {name: np.ndarray}.
target_shape – If specified, datasets will be reshaped to this shape before saving.
dataset_names_config: Names for datasets if ‘data_to_save’ is a list/ndarray.
If a string, used as a prefix. If a list, used as names.
overwrite – If True (default), overwrites the file if it exists.
Updates only provided datasets in an HDF5 file and keeps all other datasets unchanged.
Example
Existing file has ‘/a’ and ‘/b’.
Calling with data_to_update={‘/a’: new_data} updates only ‘/a’ and leaves ‘/b’ intact.
Parameters:
directory – Directory where the HDF5 file is located.
filename – Name of the HDF5 file.
data_to_update – Data to update. Can be dict {name: array}, list of arrays, or single ndarray.
target_shape – If provided, each updated dataset is reshaped before writing.
dataset_names_config – Names for datasets when data_to_update is list/ndarray.
create_if_missing – If True, creates a missing file and/or missing dataset paths as needed.
If False, update only existing datasets in an existing file.
Returns a list of HDF5 files in the specified directories matching given conditions.
:param directories: A list of directory paths (str) or Directories objects, or a single one.
:param *args: Additional arguments passed to Directories constructor if directories are str.
:param conditions: A list of callables that take a filename and return True if it matches the condition.
:param check_hdf5_condition: If True (default), adds a condition to only include files ending with .h5 or .hdf5.
:param as_string: If True (default), returns file paths as strings. If False, returns as Path objects.
Returns:
A sorted list of file paths matching the conditions.
For each directory in ‘list_of_directory_paths’, loads and concatenates data for ‘key_to_extract’.
Returns a list of concatenated numpy arrays, one for each input directory.
(This was ‘read_hdf5_extract_and_concat_list’ before)
Removes slices (e.g., rows or columns) from a matrix where initial elements are all close to zero.
For 1D vector, removes elements close to zero from the beginning up to check_limit.
Parameters:
matrix – Input numpy array.
axis – Axis along which to check for zero elements and remove slices.
tolerance – Tolerance for considering an element as zero.
check_limit – Max number of elements along the slice (or vector) to check.
If None, checks all elements in the slice.
Aligns multiple y-value series (histograms) to a common x-grid by interpolation,
filling values for x-points not present in an original series.
Parameters:
y_arrays_list – A list where each element can be a list of y-arrays (if multiple realizations share an x_array)
or a single y-array. The structure should align with group_lengths.
Example: [[y1_real1, y1_real2], [y2_real1]]
x_arrays_list – List of x-value arrays (bins), one for each group of y_arrays.
Example: [x1_bins, x2_bins]
group_lengths – List indicating how many y-arrays in ‘y_arrays_list’ correspond to each x-array
in ‘x_arrays_list’. Example: [2, 1] means y_arrays_list[0] (a list of 2 y-arrays)
uses x_arrays_list[0], and y_arrays_list[1] (a list of 1 y-array) uses x_arrays_list[1].
If y_arrays_list elements are single y-arrays, then group_lengths would be [1, 1, …].
fill_value – Value used for points in the common x-grid that are outside an original series’ x-range.
Returns:
Tuple (y_aligned_all, x_common_grid).
y_aligned_all: List of 1D numpy arrays, each y-series interpolated to x_common_grid.
Return a dict of datasets for specified keys, loading them if necessary. If keys is None, return all datasets. If copy is True, return a new dict; otherwise return the internal cache dict (which may be shared).
Load lazy entries from a directory (or single file) and apply filters.
This function scans for supported files, extracts parameters from filenames and paths,
creates lazy entries, and applies filtering based on provided conditions.
It also supports optional post-processing of parameters and custom parameter extraction for filtering.
Parameters:
data_dir – Directory path (or single file) to scan for results.
filters –
None: return all
callable: filters(result) -> bool
dict: {param: condition} where condition supports:
lx – Optional shortcuts for filtering by common size parameters. If provided, they are applied as additional filters on top of filters.
ly – Optional shortcuts for filtering by common size parameters. If provided, they are applied as additional filters on top of filters.
lz – Optional shortcuts for filtering by common size parameters. If provided, they are applied as additional filters on top of filters.
Ns – Optional shortcuts for filtering by common size parameters. If provided, they are applied as additional filters on top of filters.
post_process_func – Optional function f(params: dict) -> None that can modify the extracted parameters in-place. Useful for computing derived parameters or converting units before filtering.
get_params_func – Optional function f(result) -> dict to extract parameters from a result entry for filtering. If None, the function will look for a .params attribute or use the entry itself if it’s a dict.
logger – Optional logger for progress and error messages. Should have methods like logger.info(msg, color=…) and logger.error(msg, color=…).
recursive – Whether to scan directories recursively. Default is True.
sort_files – Whether to sort the list of files before processing. Default is True.
**kwargs – Additional keyword arguments for future extensions or specific filtering needs.
Returns:
List-like container of lazy entries with convenience methods such as
filtered(...), show(...), and show_filtered(...).
The Directories wrapper keeps legacy convenience methods available
while exposing a path-like object that can be passed to standard-library APIs.
It covers path joining, directory creation, file discovery, copying, and common
serialization helpers used by analysis scripts.
List files (not directories) in this directory.
- include_empty : if False, skip files of size zero.
- filters : a list of callables Path->bool; all must pass.
- sort_key : key function for sorting.
include_empty (bool) – if False, skip empty directories. If True, include all directories.
This checks only if the directory has any entries, not if they are files or directories.
filters (list of callables Path -> bool) – A list of callables; all must return True for a directory to be included.
sort_key (callable, optional) – Key function for sorting the results.
Collects directories across multiple machines and stores them in dictionaries.
Only directories that exist are included in existing.
Example:
>>> dirs = DirectoriesData(
>>> klimak_um_only_f=(“/media/…/klimak_um_only_f_t100000/uniform”, “503”),
>>> klimak_all=(“/media/…/klimak_um_plrb_all_t100000/uniform”, “503”),
>>> locally=(“data_project/uniform”, “local”)
>>> )
Get directories for a specific machine.
:param machine: The machine name to filter directories.
:type machine: str
:param only_existing: If True, only return existing directories. Default is True.
:type only_existing: bool, optional
Returns:
A dictionary of directory names to Directories objects.
This module provides a comprehensive set of tools for creating publication-quality
plots using Matplotlib. It is designed for scientific computing workflows with
support for Nature/Science journal formatting.
Configure matplotlib rcParams for publication-quality figures.
This function sets up consistent styling across all plots. Call it once
at the start of your script or notebook.
Parameters:
style (str, default='publication') – Style preset to apply.
- ‘publication’ : compact Nature/Science-like defaults
- ‘presentation’ : larger text and strokes for talks
- ‘poster’ : very large sizes for posters
- ‘minimal’ : stripped-down axes visuals
- ‘default’ : reset to Matplotlib defaults
font_size (int, default=10) – Base typographic scale. Label/tick/title/legend sizes are derived
from this value in each preset.
use_latex (bool, default=False) – If True, tries LaTeX-backed rendering via scienceplots profile.
If unavailable, falls back gracefully to non-LaTeX settings.
dpi (int, default=150) – Screen/display DPI used in interactive rendering.
**overrides (dict) – Additional rcParams overrides. Underscores are accepted and converted
to dots, e.g. axes_linewidth=1.0 -> axes.linewidth.
Typical high-impact keys:
- ‘savefig.dpi’
- ‘axes.prop_cycle’
- ‘figure.constrained_layout.use’
- ‘font.family’
Examples
>>> # Standard publication setup>>> configure_style('publication',font_size=10)
>>> # Presentation with larger fonts>>> configure_style('presentation',font_size=14,dpi=100)
This function modifies global matplotlib rcParams. To reset to defaults,
use configure_style(‘default’) or mpl.rcParams.update(mpl.rcParamsDefault).
Recommended figure sizes for journals:
- Nature: single column = 3.5 in, double column = 7 in
- Science: single column = 3.5 in, double column = 7.25 in
- PRL: single column = 3.4 in, double column = 7 in
List-like container for subplot axes with optional grid-aware helpers.
Behaviors:
- Inherits from list (all list operations remain available).
- Supports 2D indexing when grid metadata is available: axes[row,col].
- Forwards unknown attribute access to the first axis, enabling
Remove duplicated axis labels/ticklabels for multi-panel layouts.
Parameters:
same ({'x', 'y', 'xy'}, default='xy') – Which directions should be de-duplicated.
hide ({'both', 'labels', 'ticklabels'}, default='both') – What to hide on interior panels.
keep_x ({'bottom', 'top', 'all', 'vbottom', 'vtop'}, default='bottom') – Which row keeps x-axis labels/ticklabels.
vbottom/vtop additionally force visible labels/ticks to be
drawn at the corresponding side (not only by grid-edge selection).
keep_y ({'left', 'right', 'all', 'vleft', 'vright'}, default='left') – Which column keeps y-axis labels/ticklabels.
vleft/vright additionally force visible labels/ticks to be
drawn at the corresponding side (not only by grid-edge selection).
xlabel (str, optional) – Label text applied to kept outer x/y axes.
ylabel (str, optional) – Label text applied to kept outer x/y axes.
xlabel_kwargs (dict, optional) – Forwarded to ax.set_xlabel / ax.set_ylabel on kept axes.
ylabel_kwargs (dict, optional) – Forwarded to ax.set_xlabel / ax.set_ylabel on kept axes.
x_label_position (str, optional) – Manual label side position (x: top|bottom, y: left|right).
y_label_position (str, optional) – Manual label side position (x: top|bottom, y: left|right).
x_label_coords_system ({'axes', 'data'}) – Coordinate system used for manual label coordinates.
y_label_coords_system ({'axes', 'data'}) – Coordinate system used for manual label coordinates.
x_tick_params (dict, optional) – Tick styling applied to kept x/y axes via ax.tick_params.
y_tick_params (dict, optional) – Tick styling applied to kept x/y axes via ax.tick_params.
interior_x_tick_params (dict, optional) – Tick styling for interior x/y axes after de-duplication. Useful to
keep ticks but hide labels, adjust lengths, etc.
interior_y_tick_params (dict, optional) – Tick styling for interior x/y axes after de-duplication. Useful to
keep ticks but hide labels, adjust lengths, etc.
Publication-quality plotting utilities for scientific computing.
This class provides static methods for creating, customizing, and saving
Matplotlib figures suitable for scientific journals (Nature, Science, PRL, etc.).
All methods are @staticmethod, so you can use them without instantiation:
Print help information about available plotting methods.
Parameters:
topic (str, optional) – Specific topic to get help on. Options:
- ‘plot’: Basic plotting methods
- ‘axis’: Axis configuration
- ‘color’: Colors and colorbars
- ‘layout’: Subplots and grids
- ‘save’: Saving figures
- None: Print overview of all topics
This method extends Python str.format-style placeholders with simple
math filters for scientific labels.
Parameters:
label (str) – Template string. Standard format fields are supported, e.g.
{J:.3g}, and can be combined with filters, e.g.
{point|vec:.2f}, {kpoint|sym}.
Greek-name values (for example Gamma or omega) are
automatically rendered as LaTeX variables.
auto_wrap (bool, default=True) – If True, wrap the final string with $...$ when no dollar sign is
present in the rendered output.
escape_text (bool, default=True) – If True, plain substituted strings are LaTeX-escaped by default.
Use |raw or |tex to bypass escaping for a specific field.
**values (Any) – Values used by template fields.
filters (Supported)
-----------------
tex (raw /) – Insert value as-is (no escaping).
sym – Force symbol conversion for common Greek names
(e.g. Gamma -> \Gamma).
num – Numeric formatting helper. Uses format spec if provided, otherwise
.6g.
vec – Render iterable as \left(v_1,v_2,...\right).
set – Render iterable as \left\{v_1,v_2,...\right\}.
palette (str or list of colors, default='tableau') – Named palette string (see palette()) or an explicit list of
any Matplotlib-compatible color specs.
Examples
>>> Plotter.set_color_cycle(ax,'wong')>>> ax.plot(x1,y1)# first wong color>>> ax.plot(x2,y2)# second wong color
Each parameter shifts the corresponding channel by a fraction of the
remaining headroom, so operations compose gracefully and values are
always clamped to [0, 1].
Parameters:
color (color spec) – Any Matplotlib-compatible color.
cmap (str or Colormap, default='viridis') – Source colormap.
vmin (float, default 0.0, 1.0) – Fraction range to sample from (allows using only a sub-range of the
colormap, e.g. vmin=0.1,vmax=0.9 avoids the near-white ends of
sequential maps).
vmax (float, default 0.0, 1.0) – Fraction range to sample from (allows using only a sub-range of the
colormap, e.g. vmin=0.1,vmax=0.9 avoids the near-white ends of
sequential maps).
as_hex (bool, default=False) – If True, return hex strings instead of RGBA tuples.
Get a dictionary with color properties for matplotlib patches.
:param - color [str or tuple]: Color to use, can be a named color or an RGB tuple.
:param - alpha [float]: Transparency level (0 to 1).
:param - edgecolor [tuple]: Edge color as an RGB tuple.
:param - facecolor [tuple]: Face color as an RGB tuple.
If array-like: a new ScalarMappable is built from cmap/norm (and scale, vmin, vmax).
If ScalarMappable: it is used directly. vmin/vmax update its clim; norm is taken from it
when not provided. Note: in this case discrete/boundaries resampling is not applied.
cmap (str | Colormap, default='viridis') – Colormap name or object. If mappable is a ScalarMappable, its cmap is used unless cmap
is explicitly different from the default and a new mappable is constructed (array-like path).
norm (matplotlib.colors.Normalize, optional) – Normalization to map data to 0-1. Ignored if mappable is ScalarMappable and norm is None
(then the mappable’s norm is used).
vmin (float, optional) – Data limits. When scale=’log’, non-positive vmin is clamped internally.
vmax (float, optional) – Data limits. When scale=’log’, non-positive vmin is clamped internally.
scale ({'linear', 'log', 'symlog'}, default='linear') – Creates a suitable Normalize when mappable is array-like and norm is None.
- ‘linear’ -> Normalize
- ‘log’ -> LogNorm (vmin<=0 clamped to ~1e-10)
- ‘symlog’ -> SymLogNorm with linthresh=0.1
extend ({'neither','both','min','max','neutral'}, default='neutral') – Colorbar extension behavior. Standard Matplotlib values are ‘neither’, ‘both’, ‘min’, ‘max’.
‘neutral’ is treated as a pass-through here and may behave like ‘neither’ depending on Matplotlib.
format (str | matplotlib.ticker.Formatter, optional) – Tick formatting. If str (e.g., ‘%.2e’), uses FormatStrFormatter.
discrete (bool | int, default=False) – Discretize colormap when building from array-like:
- True -> 10 bins
- int N -> N bins
Ignored when mappable is a ScalarMappable.
boundaries (list[float], optional) – Discrete bin edges. Enables BoundaryNorm and passes boundaries to fig.colorbar
(default spacing=’proportional’, overridable via kwargs[‘spacing’]).
invert (bool, default=False) – If True, invert the colorbar axis direction.
remove_pdf_lines (bool, default=True) – Set solids edgecolor to ‘face’ to avoid white hairlines in vector exports (PDF/SVG).
**kwargs – Additional arguments forwarded to fig.colorbar, e.g.:
- alpha, spacing (‘uniform’|’proportional’), fraction, pad, shrink, aspect, drawedges, etc.
Parameters:
- values (array-like): The values to map to colors.
- cmap (str, optional): The colormap to use. Defaults to ‘PuBu’.
- elsecolor (str, optional): The color to use if there is only one value. Defaults to ‘blue’.
- get_mappable (bool, optional): If True, also return a ScalarMappable as
the 4th item, ready to pass into Plotter.add_colorbar(…, mappable=…).
return_mappable (bool, optional): Alias for get_mappable.
Returns:
- getcolor (function): A function that maps a value to a color.
- colors (Colormap): The colormap object.
- norm (Normalize): The normalization object.
- mappable (ScalarMappable, optional): Returned when get_mappable=True
Apply a colormap to the given data and plot it on the provided axis.
Parameters:
- ax (object): The axis object to plot on.
- data (array-like): The data to plot.
- cmap (str, optional): The colormap to use. Defaults to ‘PuBu’.
- colorbar (bool, optional): Whether to add a colorbar. Defaults to True.
Make an annotation on the plot.
- ax : axis to annotate on
- elem : annotation string
- x : x coordinate (ignored if xycoords=’best’)
- y : y coordinate (ignored if xycoords=’best’)
- fontsize : fontsize of the annotation
- xycoords : how to interpret the coordinates (from MPL), or ‘best’ to find the best corner
- cond : condition to make the annotation
Make an annotation on the plot.
- ax : axis to annotate on
- start_T : start text
- end_T : end text
- xystart : x coordinate start
- xystart_T : x coordinate start text
- xyend : x coordinate end
- xyend_T : x coordinate end text
- arrowprops: properties of the arrow
- startcolor: color of the arrow at the start
- endcolor : color of the arrow in the end
- kwargs : additional arguments for annotation
Plots the fitting function provided by the user on a given axis using
the **kwargs provider afterwards.
- ax : axis to annotate on
- funct : function to use for the fitting
- x : arguments to the function
Plot a scalar field sampled on irregular planar points using triangulation.
This helper is meant for 2D scattered data where imshow is not
appropriate because the samples do not lie on a regular rectangular grid.
Matplotlib first builds a triangulation of the point cloud and then
interpolates values inside each triangle.
Typical use cases:
- real-space lattice-site scalar fields
- Brillouin-zone data on irregular planar k-point sets
- any scattered 2D measurement data
Parameters:
ax (matplotlib.axes.Axes) – Target axis.
points (array-like) – Sample positions shaped like (N, 2) or (N, D) with D >= 2.
Only the first two Cartesian components are used.
values (array-like) – Scalar values of length N.
triangles (array-like, optional) – Explicit connectivity passed to matplotlib.tri.Triangulation.
mask (array-like of bool, optional) – Triangle mask. True hides the corresponding triangle.
Sets tickparams to the desired ones.
- ax : axis to use
- labelsize : fontsize
- left : whether to show the left side
- right : whether to show the right side
- top : whether to show the top side
- bottom : whether to show the bottom side
- xticks : list of xticks
- yticks : list of yticks
Comprehensive axis configuration method for publication-quality plots.
This method provides centralized control over all major axis properties
with sensible defaults and advanced options for fine-tuning. It integrates
with other Plotter methods for a cohesive styling experience.
Parameters:
ax (matplotlib.axes.Axes) – The axis object to modify.
which ({'both', 'x', 'y'}, default='both') – Specifies which axes to update. Allows independent configuration
of x and y axes.
Titles** (**Labels and)
xlabel (str, optional) – Axis labels. Set to ‘’ to hide labels while maintaining formatting.
ylabel (str, optional) – Axis labels. Set to ‘’ to hide labels while maintaining formatting.
Add labels to lines with a given slope. Uses labelLines package.
:param - ax: Matplotlib axis object.
:param - align: Align the label with the slope of the line.
:param - xvals: The x values to place the labels at.
:param - yoffsets: The y offsets for the labels.
:param - zorder: The zorder of the labels.
Remove tick labels (and optionally tick marks) from the axis.
Useful for creating clean shared-axis plots where inner panels
don’t need redundant tick labels.
Parameters:
ax (matplotlib.axes.Axes) – The axes to modify.
xticks (bool, default=False) – If True, REMOVE x-tick labels. If False, KEEP them.
yticks (bool, default=False) – If True, REMOVE y-tick labels. If False, KEEP them.
xlabel (bool, default=False) – If True, also REMOVE the x-axis label.
ylabel (bool, default=False) – If True, also REMOVE the y-axis label.
remove_labels_only (bool, default=True) – If True, only remove the text labels, keeping tick marks visible.
If False, remove both the tick marks and labels.
Examples
# Remove x-tick labels for stacked plots with shared x-axis
>>> for ax in axes[:-1]: # All except bottom
… Plotter.unset_ticks(ax, xticks=True, xlabel=True)
# Remove tick marks AND labels (completely clean)
>>> Plotter.unset_ticks(ax, xticks=True, yticks=True, remove_labels_only=False)
# Remove y-ticks and y-axis label for side-by-side shared y-axis
>>> for ax in axes[1:]: # All except leftmost
… Plotter.unset_ticks(ax, yticks=True, ylabel=True)
Notes
This function is commonly used in combination with sharex/sharey in
multi-panel figures to avoid redundant labels.
Sets the formatter for the given axis on the plot.
:param ax: The axis object on which to set the formatter.
:type ax: object
:param formater: The format string for the axis labels. Defaults to “%.1e”.
:type formater: str, optional
:param axis: The axis on which to set the formatter. Defaults to ‘xy’.
:type axis: str, optional
Sets the formatter for the given axis on the plot.
:param ax: The axis object on which to set the formatter.
:type ax: object
:param axis: The axis on which to set the formatter. Defaults to ‘xy’.
:type axis: str, optional
wspace (float, optional) – Width space between columns (0.0 to 1.0, fraction of average axis width).
Recommended: 0.2-0.4 for labels, 0.05-0.1 for tight layouts.
hspace (float, optional) – Height space between rows (0.0 to 1.0, fraction of average axis height).
Recommended: 0.2-0.4 for titles, 0.05-0.1 for tight layouts.
width_ratios (list of float, optional) – Relative widths of columns. E.g., [2, 1, 1] makes first column 2x wider.
Length must equal ncols.
height_ratios (list of float, optional) – Relative heights of rows. E.g., [1, 2] makes second row 2x taller.
Length must equal nrows.
ax_sub (SubplotSpec, optional) – If provided, creates a nested GridSpec within this subplot.
Use for complex layouts with grids inside grids.
left (float, optional) – Figure margins (0.0 to 1.0). Controls space for labels.
right (float, optional) – Figure margins (0.0 to 1.0). Controls space for labels.
top (float, optional) – Figure margins (0.0 to 1.0). Controls space for labels.
bottom (float, optional) – Figure margins (0.0 to 1.0). Controls space for labels.
**kwargs – Additional arguments passed to GridSpec.
Returns:
The grid specification object.
Return type:
GridSpec or GridSpecFromSubplotSpec
Examples
Basic 2x3 grid:
>>> fig=plt.figure(figsize=(12,8))>>> gs=Plotter.get_grid(2,3,wspace=0.3,hspace=0.4)>>> ax0=fig.add_subplot(gs[0,0])# Row 0, Col 0>>> ax1=fig.add_subplot(gs[0,1:])# Row 0, Cols 1-2 (span)>>> ax2=fig.add_subplot(gs[1,:])# Row 1, all columns (span)
Unequal widths (main panel + sidebar):
>>> gs=Plotter.get_grid(1,2,width_ratios=[3,1])>>> # First column is 3x wider than second
Create a subplot from a GridSpec at the specified index.
Parameters:
gs (GridSpec) – The GridSpec object.
fig (matplotlib.figure.Figure) – The figure to add the subplot to.
index (int, tuple, or slice) – Position in the grid. Can be:
- int: Linear index (0, 1, 2, …)
- tuple: (row, col) for single cell
- slice/tuple with slices: For spanning multiple cells
sharex (Axes, optional) – Share axis with another subplot. Use for aligned multi-panel figures.
sharey (Axes, optional) – Share axis with another subplot. Use for aligned multi-panel figures.
**kwargs – Additional arguments passed to fig.add_subplot.
Returns:
The created subplot.
Return type:
matplotlib.axes.Axes
Examples
Single cell by linear index:
ax0=Plotter.get_grid_subplot(gs,fig,0)# First cellax1=Plotter.get_grid_subplot(gs,fig,1)# Second cell
Single cell by (row, col):
ax=fig.add_subplot(gs[1,2])# Row 1, Col 2
Span multiple cells:
ax_wide=fig.add_subplot(gs[0,:])# Entire first rowax_tall=fig.add_subplot(gs[:,0])# Entire first columnax_block=fig.add_subplot(gs[0:2,1:3])# 2x2 block
Shared axes (for aligned panels):
ax0=Plotter.get_grid_subplot(gs,fig,0)ax1=Plotter.get_grid_subplot(gs,fig,1,sharex=ax0)ax2=Plotter.get_grid_subplot(gs,fig,2,sharex=ax0,sharey=ax0)# ax1 and ax2 share x-axis with ax0; ax2 also shares y-axis
>>> gmap=Plotter.get_grid_map(2,3)>>> gmap['by_letter']['c']# Get index for panel 'c'2>>> gmap['by_index'][4]# Get (row, col) for index 4(1, 1)>>> gmap['grid']# 2D layout[[0, 1, 2], [3, 4, 5]]
Configure axis visibility, spines, ticks, and labels in one call.
This is a convenience function for common axis customizations.
Parameters:
ax (matplotlib.axes.Axes) – The axis to configure.
visible (bool, default=True) – If False, hide the entire axis (ax.axis(‘off’)).
spines (bool, dict, or str, default=True) – Control spine visibility:
- True: Show all spines
- False: Hide all spines
- ‘left’: Hide all except left
- ‘bottom’: Hide all except bottom
- ‘minimal’: Hide top and right (Nature-style)
- dict: {‘top’: False, ‘right’: False, …}
ticks (bool, dict, or str, default=True) – Control tick visibility:
- True/False: Show/hide all ticks
- ‘x’/’y’: Show only x/y ticks
- dict: {‘x’: True, ‘y’: False}
tick_labels (bool, dict, or str, default=True) – Control tick label visibility (same format as ticks).
Disable axis components for clean images/heatmaps.
Parameters:
ax (matplotlib.axes.Axes) – The axis to modify.
which (str, default='both') – What to disable:
- ‘both’: Disable x and y (full axis off)
- ‘x’: Disable x-axis only
- ‘y’: Disable y-axis only
- ‘labels’: Keep ticks but hide labels
- ‘ticks’: Keep labels but hide ticks
- ‘spines’: Hide all spines
Sets the legend with a preferred style for publication-quality plots.
Parameters:
- ax : Axis to which the legend will be added.
- fontsize : Font size of the legend labels.
- frameon : Whether to draw a frame around the legend.
- loc : Location of the legend (‘best’, ‘upper right’, etc.).
- alignment : Text alignment (‘left’, ‘center’, ‘right’).
- markerfirst : Whether the marker or label appears first in the legend.
- framealpha : Transparency of the legend frame (1.0 is opaque).
- reverse : Reverse the order of legend items.
- style : Predefined style for the legend (‘minimal’, ‘boxed’, etc.).
- labelspacing : Vertical space between legend entries.
- handlelength : Length of the legend markers.
- handletextpad : Space between legend markers and text.
- borderpad : Padding inside the legend box.
- columnspacing : Spacing between legend columns.
- ncol : Number of columns in the legend.
- kwargs : Additional arguments passed to ax.legend().
Set the legend with custom conditions for the labels
- ax : axis to use
- conditions: list of conditions
- fontsize : fontsize
- frameon : frame on or off
- loc : location of the legend
- alignment : alignment of the legend
- markerfirst: marker first or not
- framealpha: alpha of the frame
axes (AxesList or matplotlib.axes.Axes) – AxesList wrapper (list-compatible, grid-aware, named-panel
access). Returns single axis only when single_if_1=True.
Save figure to a specific directory.
- directory : directory to save the file
- filename : name of the file
- format : format of the file
- dpi : dpi of the file
- adjust : adjust the figure
staticsavefig(directory, filename, format, dpi, adjust, fig=None, **kwargs)[source]
Alias for save_fig() with the historical lowercase name.
Based on the specified keys, load the dictionaries from the json files
The keys are the names of the files as well!
staticsingleColumnData(directory:str, fileName:str, y, typ='.npy')[source]
Stores the values as a single vector
statictwoColumnsData(directory:str, fileName:str, x, y, typ='.npy')[source]
Stores the x, y vectors in 2D form (multiple rows and two columns)
staticmatrixData(directory:str, fileName:str, x, y, typ='.npy')[source]
Stores the x, y vectors in matrix form (appending single column
at start for x values)
staticapp_df(df, colname:str, y, fill_value=nan)[source]
Appends the data to the dataframe.
Parameters:
- df (pd.DataFrame): The dataframe to append data to.
- colname (str): The column name to append data under.
- y (array-like): The data to append.
- fill_value: The value to use for filling if resizing is needed.
This module provides a DataHandler class for handling and processing data arrays.
It includes methods for filtering, initializing, interpolating, aggregating,
concatenating, and averaging data arrays.
Classes:
DataHandler: A class containing static methods for data handling and processing.
DataHandler class provides static methods for handling and processing data arrays, including filtering,
interpolating, aggregating, concatenating, and cutting matrices based on specific criteria.
.. method:: _filter_typical_values(current_x, current_y, typical, threshold=1.0) -> tuple
Initializes and combines arrays from given lists. If the typical flag is set to True,
it filters the combined arrays to include only elements where the values in y_combined
are less than the threshold.
Concatenates and averages y values across multiple histograms.
:param y_list : List of y matrices (each one corresponding to a realization).
:param x_list : List of x vectors (each one corresponding to a realization).
:param typical : If True, filter y values less than 1.0.
:param use_interpolation: If True, interpolate y values for non-matching bins.
If False, aggregate only exact matches and append unique bins.
:param threshold : The threshold value for filtering y values (default: 1.0).
:returns : Combined y values and x bins after averaging.
Cut off the slices (along any specified axis) in matrix M where all elements are close to zero.
If a 1D vector is provided, it returns the vector unless all elements are close to zero,
in which case it returns an empty array.
Parameters:
- M (numpy.ndarray) : The input matrix or vector.
- axis (int) : The axis along which to check for zero elements.
For example, 0 for rows, 1 for columns, etc.
Ignored if M is a 1D vector.
tol (float) : The tolerance for considering elements as zero.
check_limit (int) : The maximum number of elements along the axis to check for zeros.
Returns:
- numpy.ndarray: The resulting matrix after removing slices (along the specified axis)
that are close to zero, or the vector after removing if all elements are close to zero.
Cut off the rows or columns in matrix M where the first check_limit elements are all below a threshold.
Parameters:
- M (numpy.ndarray): The input matrix.
- axis (int): The axis along which to check for elements below the threshold (0 for rows, 1 for columns).
- threshold (float): The threshold value.
- check_limit (int, optional): The number of elements to check from each row or column.
Returns:
- numpy.ndarray: The resulting matrix after removing rows or columns where the first check_limit elements are below the threshold.
Small logging helpers with verbosity, indentation, color, and file output.
The Logger wrapper keeps the public API lightweight while avoiding
duplicate notebook handlers and optionally mirroring logs to a file. Console
colors can be disabled with PYLOGCOLORS=0; file logging is enabled by
setting PYLOGFILE to a non-empty path-like value.
[summary]
Function that can print a list of elements creating indents
The separator also can be used to clean the indents.
Arguments:
- width : governing the width of each column.
- endline : if true, puts an endline after the last element of the list.
- scientific: allows for scientific printing.
Prints the arguments of a parser in a formatted table using the provided logger.
:param parser: The argument parser containing the script’s options.
:type parser: argparse.ArgumentParser
:param logger: The logger instance used to print the formatted output.
:type logger: Optional[Logger]
:param title: The title to display above the options table. Defaults to “Options for the script”.
:type title: str, optional
:param columnsize: The width of the “Option” column in the table. Defaults to 30.
:type columnsize: int, optional
>>> logger=get_global_logger()>>> logger.info("This is an informational message.")>>> logger.debug("This is a debug message.",color='blue')
binary.py
This module implements binary - manipulation routines that allow you to work with binary representations of integers
or vectors.
It includes functions to:
- check if a number is a power of 2,
- check a given bit in an integer or an indexable vector,
- convert an integer to a base representation (spin: +/- value or binary 0/1),
- convert a base representation back to an integer or a binary string or a vector of values,
- flip bits (all at once or a single bit),
- reverse the bits of a 64 - bit integer,
- extract bits (by ordinal position or via a mask),
- and prepare a bit mask from a list of positions.
Functions are written so that they work with plain Python
integers as well as with NumPy or JAX arrays.
You can choose the backend (np or jnp) by passing the corresponding module.
DescriptionThis module implements binary - manipulation routines that allow you to work with binary representations of integers
size (int) – The number of bits in the integer representation.
shift (int, optional) – The number of positions to shift. Default is 1.
backend (str, optional) – The backend to use (e.g., ‘default’, ‘numpy’, or ‘jax’).
axis (int, optional) – The axis along which to rotate for array-like states.
Returns:
Rotated state (same type as n).
Timing utilities for benchmarking and profiling.
This module provides:
- A Timer class for context-manager and decorator-based timing with support for laps and deadlines.
- Functional wrappers (timeit) for simple benchmarking.
- JAX-aware synchronization helpers to ensure accurate timing of asynchronous operations.
Features:
- High-precision monotonic clock (nanoseconds).
- Automatic synchronization for JAX arrays (blocks until ready).
- Detailed reporting with laps and mean/std stats.
Factory functions return subclasses of Lattice with explicit geometry
metadata (dimensions, boundary conditions, primitive vectors, and neighbor maps).
Typical constructor inputs are integer sizes (lx,ly,lz), a boundary mode,
and optional flux or graph descriptors.
Coordinate arrays are expected as real-valued arrays with shape (ns,dim).
Index-based neighbor structures are integer arrays or lists over site ids in
[0,ns). Plotting helpers consume NumPy-compatible arrays.
Topology construction is deterministic for fixed parameters. Floating-point
roundoff can affect reciprocal-space formatting or plotting labels but should not
change connectivity.
Collection of magnetic fluxes piercing lattice boundary loops.
The value associated with a direction is interpreted as the phase
phi
(in radians) acquired upon wrapping around the boundary once along that
direction. The corresponding hopping phase factor is exp(1j*phi).
The fluxes are stored as a mapping from LatticeDirection to corresponding
complex phase values.
Options for specifying fluxes include:
- Uniform flux in all directions (single float value).
- Direction-specific fluxes (mapping from direction to float).
- Zero flux (empty mapping).
Physically, these fluxes correspond to magnetic fluxes threading
the holes of a torus formed by periodic boundary conditions.
Return the fractional k-grid offset induced by the boundary fluxes.
With flux \(\phi_\mu\) in direction \(\mu\) of length
\(L_\mu\), the allowed Bloch momenta shift from
\(f_\mu = n_\mu / L_\mu\) to
\(f_\mu + \phi_\mu / (2\pi L_\mu)\).
Returns:
(delta_fx, delta_fy, delta_fz) – Fractional coordinate shifts (add to the standard grid).
Abstract Base Class for defining lattice structures.
This class serves as the foundation for all lattice implementations in the lattices module.
It handles geometry, connectivity, boundary conditions, and k-space properties.
Lattice sites are indexed linearly from 0 to Ns-1.
The mapping from spatial coordinates to linear index depends on the concrete implementation,
but typically follows a row-major (lexicographic) order:
General Lattice class. This class contains the general lattice model.
Parameters:
dim (int, optional) – Dimension of the lattice (1, 2, or 3). If None, inferred from lx, ly, lz.
lx (int, optional) – Length of the lattice in the x-direction.
ly (int, optional) – Length of the lattice in the y-direction.
lz (int, optional) – Length of the lattice in the z-direction.
bc (str, optional) – Boundary conditions (e.g., ‘PBC’, ‘OBC’).
adj_mat (np.ndarray, optional) – Adjacency matrix for the lattice.
flux (np.ndarray, optional) – Flux piercing the boundaries. This can be a dictionary specifying the
flux in each direction, or a single value applied to all directions. Importantly,
this automatically implies TWISTED boundary conditions, so the bc parameter can be left as None or set to ‘TWISTED’ for clarity.
Initializes the lattice object by calculating coordinates, reciprocal vectors, and neighbor lists.
This method performs the following steps:
1. Calculates the real-space coordinates, r-vectors, and k-vectors of the lattice.
2. If the number of sites (self.Ns) is less than 100, computes the discrete Fourier transform (DFT) matrix.
3. If an adjacency matrix (self._adj_mat) is provided:
Determines the number of sites (Ns) from the adjacency matrix.
For each site, identifies nearest neighbors (nn) as those connected by the highest weight in the adjacency matrix, and next-nearest neighbors (nnn) as those connected by the next highest distinct weight.
Stores forward neighbors (indices greater than the current site) for both nn and nnn.
If no adjacency matrix is provided, calculates nearest and next-nearest neighbors using default methods.
5. Calculates normalization or symmetry properties of the lattice.
This method sets up all necessary neighbor lists and lattice properties required for further computations.
Return a list of site indices defining a spatial region.
Parameters:
kind (str or RegionType) – Type of region: ‘half’, ‘disk’, ‘sublattice’, ‘graph’, ‘plaquette’, ‘custom’.
We also support specific half cuts like ‘half_x’, ‘half_y’, ‘half_z’ for convenience.
origin (int or list[float], optional) – Center of the region. Can be a site index or coordinate vector.
radius (float, optional) – Radius for ‘disk’ regions.
direction (str, optional) – Direction for ‘half’ cuts (‘x’, ‘y’, ‘z’).
sublattice (int, optional) – Sublattice index for ‘sublattice’ regions.
sites (list[int], optional) – Explicit list of sites for ‘custom’ regions.
depth (int, optional) – Depth/distance for ‘graph’ regions.
plaquettes (list[int], optional) – List of plaquette indices for ‘plaquette’ regions.
Returns:
Sorted list of site indices belonging to the region.
Generate space-group permutation table for this lattice.
Delegates to generate_space_group_perms().
When TWISTED boundary conditions are active, the point-group part is
disabled (only translations are returned) because a generic flux
breaks point-group symmetry unless the flux respects it.
Parameters:
point_group (str) – 'full' for maximal point group, 'translations' for translations only.
Return a dictionary describing the spatial symmetries of this lattice.
The information is consistent for both single-particle and many-body
representations. When TWISTED boundary conditions are present the
point-group part is absent (flux generically breaks it).
Returns:
Keys:
- 'lattice_type' : LatticeType enum
- 'sites_per_cell' : int
- 'n_cells' : number of unit cells
- 'dim' : spatial dimension
- 'bc' : boundary condition enum
- 'is_periodic' : (bool, bool, bool) per direction
- 'is_twisted' : bool
- 'translation_group' : ZL_x x ZL_y (as tuple (Lx,Ly))
- 'point_group' : str or None ('D4' for square Lx==Ly, etc.)
- 'space_group_order' : total number of space-group elements
- 'flux' : BoundaryFlux or None
Return the sublattice indices of the lattice sites.
For a Bravais lattice, this would simply be an array of zeros.
For a non-Bravais lattice, this would indicate which sublattice each site belongs to.
Return the unit cell coordinates of the lattice sites. For a Bravais lattice,
this would simply be the integer coordinates of the unit cells.
For a non-Bravais lattice, this would include the basis vectors as well.
Return the sublattice index for a given site.
By default, returns 0 for all sites (single sublattice).
Override in subclasses for multi-sublattice lattices.
Extend k-space points and optional data across translated Brillouin zones.
The helper works for arbitrary k-space dimensions and any number of
reciprocal translation vectors. Legacy b1/b2/b3 with
nx/ny/nz remain supported for existing callers.
Allows to generate extended k-point grids for plotting band structures along high-symmetry paths…
Parameters:
k_points (ndarray, shape (N, dim)) – Array of k-points in reciprocal space to be extended.
data (ndarray, shape (N, ...) or None) – Optional data associated with each k-point (e.g. eigenvalues) to be extended alongside the k-points. Must have the same leading dimension as k_points.
copies (int or iterable of ints, optional) – Number of translated copies to generate in each reciprocal direction. If an integer is provided, the same number of copies will be generated in all directions. If an iterable is
provided, it should have a length equal to the number of reciprocal lattice vectors (e.g. 3 for 3D), specifying the number of copies in each direction separately.
**kwargs – Additional keyword arguments to pass to the underlying ws_extend function. See its documentation for details.
Returns:
extended_k_points (ndarray, shape (M, dim)) – Array of extended k-points in reciprocal space, including the original points and their translated copies
extended_data (ndarray, shape (M, …) or None) – Extended data associated with each k-point, if the input data was provided. Otherwise, None
Return a boolean mask for the Wigner-Seitz cell in reciprocal space.
This can be used to identify which k-points lie within the first Brillouin zone.
Parameters:
Kx (array-like) – Arrays of k-point coordinates in reciprocal space. This is a grid
of k-points for which we want to determine if they lie within the Wigner-Seitz cell.
Ky (array-like) – Arrays of k-point coordinates in reciprocal space. This is a grid
of k-points for which we want to determine if they lie within the Wigner-Seitz cell.
Kz (array-like) – Arrays of k-point coordinates in reciprocal space. This is a grid
of k-points for which we want to determine if they lie within the Wigner-Seitz cell.
shells (int) – Number of shells of Wigner-Seitz cell to include in the mask.
tol (float) – Tolerance for determining if a point is within the Wigner-Seitz cell, accounting for numerical precision issues.
**kwargs – Additional keyword arguments to pass to the underlying ws_bz_mask function. See its documentation for details.
Resolve path input to a list of (label, fractional_coord) pairs.
Parameters:
path (list[(label, coords)], StandardBZPath, str, List[str], or HighSymmetryPoints) – Path definition (fractional coordinates), standard enum, enum name string,
list of point labels, or HighSymmetryPoints object.
lattice (Lattice, optional) – Lattice object used to resolve labels if path is a list of strings.
Returns:
resolved_path – Resolved path as a list of (label, fractional_coord) pairs.
Return True if the lattice momentum grid contains a special point.
This method helps to check whether a finite lattice
contains a particular high-symmetry point in the Brillouin zone,
which is important for band structure calculations and topological analyses.
Parameters:
point – Special point identifier. Accepted forms:
- label string (e.g. "Gamma", "K", "K'"),
- HighSymmetryPoint,
- explicit fractional coordinate tuple/array.
tol (float) – Absolute tolerance used in the coordinate match.
Notes
The check is done in fractional reciprocal coordinates and naturally
includes flux-induced shifts from twisted boundary conditions because it
uses self.kvectors_frac.
path (list of str, str, StandardBZPath, or None) – Path specification. Can be:
- List of high-symmetry point names: [‘Gamma’, ‘X’, ‘M’, ‘Gamma’]
- StandardBZPath enum or string: ‘SQUARE_2D’
- None: use default path for this lattice
points_per_seg (int) – Number of interpolated points per path segment.
Returns:
k_path (np.ndarray, shape (Npath, 3)) – Cartesian k-points along the path.
k_dist (np.ndarray, shape (Npath,)) – Cumulative distance for plotting x-axis.
labels (List[Tuple[int, str]]) – Indices and labels for high-symmetry points.
k_path_frac (np.ndarray, shape (Npath, 3)) – Fractional k-coordinates along the path.
Example
>>> lattice=SquareLattice(dim=2,lx=4,ly=4)>>> k_path,k_dist,labels,k_frac=lattice.bz_path()>>> # Or with custom path:>>> k_path,k_dist,labels,k_frac=lattice.bz_path(['Gamma','M','Gamma'])
Build an ideal Brillouin-zone path and optionally match it to an existing k-grid.
If no k-grid is provided, the returned object still contains the continuous
path geometry, which is useful for plotting or for constructing a path that
is not constrained to the sampled reciprocal mesh. When a sampled grid is
provided, reciprocal-lattice copies are generated automatically as needed
so paths in extended Brillouin-zone regions can still match the existing
data.
Parameters:
path (list of str, str, StandardBZPath, or None) – Path specification. Can be:
- List of high-symmetry point names: [‘Gamma’, ‘X’, ‘M’, ‘Gamma’]
- StandardBZPath enum or string: ‘SQUARE_2D’
- None: use default path for this lattice
points_per_seg (int) – Number of interpolated points per path segment.
k_vectors (np.ndarray, shape (Nk, 3), optional) – Cartesian k-vectors of the existing grid to match against.
k_vectors_frac (np.ndarray, shape (Nk, 3), optional) – Fractional k-vectors of the existing grid to match against. Required if k_vectors is provided.
tol (float) – Tolerance for matching path points to the existing k-grid. With
periodic=True it is interpreted in fractional reciprocal
coordinates. With periodic=False it is interpreted in plotted
Cartesian reciprocal coordinates.
periodic (bool, default=True) – If True, allow reciprocal-translation-equivalent points to match.
Set to False for visual matching in the displayed Brillouin-zone copy.
Extract k-path data from a k-grid using fractional coordinate matching.
This function finds the closest k-points on the actual grid to an ideal path
through high-symmetry points. It handles periodic boundary conditions in k-space
and automatically reuses reciprocal-lattice copies of the sampled grid when the
requested path lies in an extended Brillouin-zone region.
It also allows to return a structured KPathResult dataclass or a tuple…
Parameters:
lattice (Lattice) – Lattice object with reciprocal lattice vectors
k_vectors_frac (np.ndarray, shape (..., 3)) – Fractional coordinates of k-points (will be flattened)
values (np.ndarray) – Data values sampled on the k-grid. The k-grid axes may appear as
(Lx,Ly,Lz,...) or after leading batch axes such as time or
frequency, e.g. (Nw,Lx,Ly,Lz) or (Nw,Lx,Ly,Lz,...).
A single flattened k-grid axis of length Nk is also supported.
path (various, optional) – Path specification. Can be:
- StandardBZPath enum value (e.g., StandardBZPath.SQUARE_2D)
- String name (e.g., ‘SQUARE_2D’)
- List of (label, [f1,f2,f3]) tuples
- HighSymmetryPoints object (uses default path)
- None: uses lattice’s default path if available
points_per_seg (int) – Number of interpolated points per path segment
return_result (bool) – If True (default), return KPathResult dataclass.
If False, return tuple for backwards compatibility.
Returns:
If return_result=True: KPathResult dataclass with all path data.
The returned values preserve any leading batch axes and replace
the k-grid axes with a path axis.
If return_result=False: (k_cart, k_frac, k_dist, labels, values) tuple
>>> # Using default path from HighSymmetryPoints>>> result=bz_path_data(lattice,k_grid,k_frac,energies,HighSymmetryPoints.square_2d())>>> plt.plot(result.k_dist,result.values)
>>> # Using standard path enum>>> result=bz_path_data(lattice,k_grid,k_frac,energies,'SQUARE_2D')
Check if a set of sites spans the lattice (non-contractible on a torus).
This method uses a BFS-based winding number tracking on the induced subgraph of
the provided site indices. If any loop with a non-zero winding number
along a periodic direction is found, the set is considered spanning.
Return the complex hopping phase factor for the bond \(i \to j\).
For bonds that do not cross a periodic boundary, this is 1.
For boundary-crossing bonds under TWISTED BC, the phase is
\(\exp(i\,\phi_\mu)\) for each direction \(\mu\) in which
the bond wraps.
This is the factor that should multiply the bare hopping amplitude
in real-space Hamiltonian construction.
Build an \(N_s \times N_s\) matrix of complex hopping amplitudes
that includes the Peierls phases from boundary fluxes.
Diagonal is zero. Off-diagonal H[i,j]=t_{ij}*phase(i->j)
where t_{ij}=1 for all connected pairs and phase is the
product of boundary phases along directions that the bond wraps.
Parameters:
include_nnn (bool) – If True, include next-nearest-neighbor hoppings as well.
orientation ({'anticlockwise', 'clockwise', None}, optional) – If provided, will sort/choose based on geometric angle.
Default: None (return all middle sites).
Returns:
List of middle-site indices (can be 0, 1, or 2 elements).
Return the displacement i->j with optional PBC minimum-image wrapping.
Parameters:
i (int or tuple) – Site indices or explicit coordinates.
j (int or tuple) – Site indices or explicit coordinates.
minimum_image (bool, default=False) – If True, wrap each periodic direction to the shortest displacement.
real_space (bool, default=False) – If True and i,j are site indices, return displacement in real-space
vectors (uses displacement()). Otherwise use lattice coordinates.
i = (R, beta) real-space cell R and sublattice beta
n = (k, alpha) k-point k and sublattice alpha
Elements:
$$
F_{(k,alpha),(R,beta)} =
1/sqrt(Nc) * delta_{alpha,beta} * exp(-i k . R).
$$
This is unitary:
$$
F^dagger F = I_{Ns}, F F^dagger = I_{Ns},
$$
where Ns = Nc * Nb is the total number of sites, Nc is the number of
unit cells, and Nb is the number of sublattices.
Important
When boundary fluxes are present (TWISTED BC), the k-grid used to
build the DFT matrix is shifted by phi_mu/(2piL_mu) in each
direction, exactly as in calculate_k_vectors().
Note that this DFT matrix does not include basis-dependent phases
(i.e., exp(-i k . r_basis)).
Calculates the Discrete Fourier Transform (DFT) matrix for the lattice.
This method can be optimized using FFT (Fast Fourier Transform) in the future.
Reference: https://en.wikipedia.org/wiki/DFT_matrix
Returns the matrix with indcies corresponding to a slice from the QMC.
A usefull function for reading the position Green’s function saved from:
@url https://github.com/makskliczkowski/DQMC
The Green’s functions are saved in the following manner. If cut is True, data
has (2L_i - 1) possible position differences, otherwise we skip the negative ones and use L_i.
For 1D simulation: 1 column and (2 * Lx - 1) rows for possition differences (-Lx, -Lx + 1, …, 0, …, Lx)
For 2D simulation: (2 * Lx - 1) rows for possition differences (-Lx, -Lx + 1, …, 0, …, Lx) and (2 * Ly - 1) columns for possition differences (-Ly, -Ly + 1, …, 0, …, Ly)
For 3D simulation: Same as in 2D but after (2 * Lx - 1) x (2 * Ly - 1) matrix has finished, a new slice for Lz appears for next columns Lz * (2*Ly - 1)
- cut : if true (2L_i - 1) possible position differences, otherwise we skip the negative ones and use L_i.
Calculates the allowed reciprocal-space k-vectors (momentum grid)
consistent with the lattice size and primitive reciprocal vectors.
When boundary fluxes are present (TWISTED BC), the fractional
coordinates are shifted by \(\phi_\mu / (2\pi L_\mu)\) in
each direction, so that the Bloch condition matches the twisted
boundary.
The sampling follows the same fftfreq ordering used by the Bloch
transform (Γ at index [0,0,0], followed by positive frequencies and
finally the negative branch). This keeps the analytic grids aligned
with the numerically constructed H(k) blocks.
For TWISTED boundary conditions the neighbor connectivity is
identical to PBC — the flux phases are applied separately when
building the Hamiltonian or the DFT matrix.
Calculates the Wilson loops (non-contractible loops) for the lattice based on its boundary conditions.
Returns a list of lists, where each inner list contains the site indices of a Wilson loop.
Assumes standard lexicographic site indexing (x + y*Lx + z*Lx*Ly).
Like calculate_nn(), each calculate_nnn_in implementation
is expected to set self._nnn (and optionally self._nnn_forward)
directly. The return value—if any—is stored as a fallback.
Logs the neighbors of each site in the lattice using the provided logger.
For each site in the lattice, this method retrieves its nearest neighbors and logs their indices.
Additionally, for each neighbor, it logs detailed information using a higher verbosity level.
Parameters:
logger – An object with an info method for logging messages. The info method should accept
parameters lvl (int) for verbosity level and color (str) for message color.
Logs the forward nearest neighbors for each site in the lattice.
For each site in the lattice, this method retrieves the number of forward nearest neighbors
and logs their indices using the provided logger. The method outputs two levels of information:
- Level 1 (green): Lists the neighbors of each site.
- Level 2 (blue): Details each neighbor’s index for the site.
Parameters:
logger (A logging object with an info method that accepts a message,) – a logging level (lvl), and a color (color).
Inverse Bloch transform: reconstruct real-space matrix from k-space blocks.
This is the exact inverse of kspace_from_realspace(). It reconstructs the
real-space Hamiltonian from momentum-space blocks using the inverse Fourier transform:
Example 2: Inverse transform without explicit kgrid
>>> # If kgrid is omitted, it's reconstructed using fftfreq convention>>> H_real_recon=lattice.realspace_from_kspace(H_k)>>> np.allclose(H_real,H_real_recon)# True
Custom k-point sampling (only used when block_diag=True):
If None: Uses automatic full Brillouin zone grid based on lattice size
(recommended for most use cases)
If provided: Array of shape (Nk,3) with custom k-points in Cartesian coordinates
Returns (H_k,kpoints) with H_k shape (Nk,Nb,Nb)
unitary_norm (bool, default=True) – Use unitary normalization \(1/\sqrt{N_c}\) for the Bloch transform.
If False, uses normalization \(1/N_c\) instead. Keep True for
standard quantum mechanics convention preserving operator norms.
If True, also return the Bloch unitary matrix W used for the transformation.
This is useful for transforming additional operators or computing correlation functions.
Note: Only available when block_diag=True. The unitary is returned as a 4th
output value with shape (Lx,Ly,Lz,Ns,Nb) or (Nk,Ns,Nb) if custom
k-points are provided.
Returns:
**Case 1 (block_diag=False (default)**) –
H_k_fullnp.ndarray
Full transformed matrix with shape (Ns,Ns). This is the complete DFT
of the input matrix, preserving all information.
Example 4: Get Bloch unitary for operator transforms
>>> H_k,k_grid,k_frac,W=lattice.kspace_from_realspace(... H_real,block_diag=True,return_transform=True... )>>> # Transform another operator using the same W>>> O_k=np.einsum('kia,ij,kjb->kab',W.conj(),O_real,W)
Notes
Periodic boundary conditions (PBC) are assumed for the Bloch transform.
The method assumes translational invariance of the system, which ensures
the spectrum of H_real equals the union of spectra of H(k) blocks.
- For the full grid (kpoints=None), the k-points follow the fftfreq convention
with the Γ-point initially at index [0,0,0], then shifted to the center.
- Site ordering is arbitrary; the method uses the lattice geometry (coordinates + basis)
to correctly identify sublattices and apply phases.
- For sparse input matrices, automatic conversion to dense format is performed.
Convert a real-space correlation matrix into a momentum-resolved structure factor.
This is a convenience wrapper around the basis-aware Bloch projector in
QES.general_python.lattices.tools.lattice_kspace.kspace_from_realspace.
The real-space input mat is first transformed into the multipartite
k-space block representation evaluated on self.kvectors
where R,R' label unit cells and alpha,beta label basis sites
inside the unit cell. The reduction argument then decides how this
multipartite object is converted into a scalar structure factor at each
sampled momentum q.
Parameters:
mat (np.ndarray) – Real-space correlation or operator matrix with shape (Ns,Ns)
or batched shape (...,Ns,Ns). Any leading axes, e.g. time,
frequency, disorder sample, or state index, are preserved.
C(q) with shape (Lx,Ly,Lz,Nb,Nb) (i.e., no reduction).
"sum":
return sum_{alpha,beta}C_{alphabeta}(q) (i.e., sum over all entries of each block).
"trace":
return sum_alphaC_{alphaalpha}(q) (i.e., sum over diagonal entries of each block).
"mean":
return the arithmetic mean of all multipartite block entries at each q (i.e., sum over all entries and divide by Nb^2).
"diag":
return the eigenvalues of each block, which can be useful for identifying dominant modes or instabilities. The output shape will be (Lx,Ly,Lz,Nb) since each block’s eigenvalues are returned as a vector of length Nb.
norm ({"none", "cell", "site"}, default="none") –
Optional post-normalization of the returned k-space quantity:
"none":
keep the raw Bloch-projector normalization, i.e. the blocks
C(q) defined above with the prefactor 1/N_c.
"cell":
alias for "none" kept for readability when you want to
emphasize unit-cell normalization.
"site":
divide the returned blocks or reduced values by the number of
basis sites N_b. For scalar reductions such as "sum", this
converts the default unit-cell normalization into the more common
site normalization 1/N_s used in
\(S(q) = \langle O_{-q} O_q \rangle\).
Returns:
values (np.ndarray) – Momentum-resolved structure factor. For input shape (...,Ns,Ns)
the output shape is:
(...,Lx,Ly,Lz,Nb,Nb) for reduction="none"
(...,Lx,Ly,Lz,Nb) for reduction="diag"
(...,Lx,Ly,Lz) for "sum", "trace", or "mean"
For a single input matrix (Ns,Ns), the leading ... is absent.
k_grid (np.ndarray) – Cartesian sampled k-grid with shape (Lx,Ly,Lz,3).
k_frac (np.ndarray) – Fractional sampled k-grid with shape (Lx,Ly,Lz,3).
Notes
Use reduction="none" when sublattice-resolved information matters.
Use one of the scalar reductions when you want a single value per
momentum that can be fed directly into bz_path_data.
The default norm="none" preserves the existing unit-cell
normalization. For comparisons against structure factors built from
Fourier-transformed site operators, norm="site" is typically the
physically relevant choice.
Examples
>>> Sq,k_grid,k_frac=lattice.structure_factor(corr_zz,reduction="sum")>>> path=lattice.bz_path_data(k_grid,k_frac,Sq,path=['Gamma','K','M','Gamma'])>>>>>> # Frequency-resolved data with shape (Nw, Ns, Ns)>>> Sqw,k_grid,k_frac=lattice.structure_factor(corr_zz_w,reduction="sum")>>> # Sqw has shape (Nw, Lx, Ly, Lz)
Include:
- point_edgecolor: Color of the marker edges (default “white”).
- point_zorder: Z-order for the scatter points (default 5).
- color_extended: Color for translated copies (default “C2”).
- edgecolor_extended: Edge color for translated copies (default “gray”).
- marker_extended: Marker for translated copies (default marker).
- Any other valid arguments for ax.scatter.
Convenience wrapper returning the matplotlib figure and axes for a detailed lattice structure plot.
Parameters:
show_indices (bool) – If True, annotates nodes with their site indices.
highlight_boundary (bool) – If True, draws boundary nodes with a distinct color/edge.
show_axes (bool) – If False, hides the coordinate axes for a cleaner diagram.
partition_colors (tuple of str, optional) – Colors to use for bipartite/sublattice coloring. If provided, nodes are
colored based on sublattice parity.
show_periodic_connections (bool) – If True, indicates wrap-around connections textually or graphically.
show_primitive_cell (bool) – If True, overlays the primitive unit cell vectors/box.
(e.g. (... other kwargs passed to the underlying plotting function)
show_bz (bool, default=True) – Draw the first Brillouin zone.
show_path (bool, default=True) – Draw the ideal high-symmetry path.
show_matched_kpoints (bool, default=True) – Highlight sampled k-points whose distance to the path is within the
matching tolerance.
points_per_seg (int, default=40) – Number of interpolation points per path segment for the ideal path.
path_match_tol (float, optional) – Distance tolerance used when highlighting mesh points near the
drawn path.
extend (bool, default=False) – Draw translated copies of the sampled k-mesh.
extend_copies (int or iterable[int], optional) – Number of reciprocal-cell copies per direction. In 2D,
extend_copies=1 includes the first shell around the first Brillouin
zone and extend_copies=2 includes the second shell as well.
coordinates: Real-valued array of shape (Ns, dim).
kvectors: Real-valued array of shape (Ns, 3) (or dim).
Neighbor indices are integers in range [0, Ns).
High-symmetry points in the Brillouin zone:
- 1D: Γ (0) -> X (Pi) -> Γ (2Pi)
- 2D: Γ (0,0) -> X (Pi,0) -> M (Pi,Pi) -> Γ (0,0)
- 3D: Γ -> X -> M -> Γ -> R -> X
Return high-symmetry points for the square/cubic lattice.
Returns:
High-symmetry points with default path based on dimension:
- 1D: Γ -> X -> Γ (zone boundary at Pi)
- 2D: Γ -> X -> M -> Γ (standard square BZ path)
- 3D: Γ -> X -> M -> Γ -> R -> X (standard cubic BZ path)
Convert (x, y, z) coordinates to a site index.
:param x: x-coordinate
:type x: int
:param y: y-coordinate
:type y: int
:param z: z-coordinate
:type z: int
Returns nearest neighbors in a given direction (X, Y, Z).
:param site: Site index
:type site: int
:param direction: Direction to get the nearest neighbors
:type direction: LatticeDirection
Simple nearest-neighbour tight-binding/spin-wave-like dispersion for the
square lattice. Accepts k as (2,) or (…,2) array and returns same-shaped
scalar or array of energies.
Armchair-oriented hexagonal (honeycomb) lattice up to 3 dimensions.
The lattice is constructed so that armchair edges lie along the
horizontal (x) axis, giving a rectangular bounding box aligned
with the coordinate system. Two sites per unit cell (A / B
sublattices).
General Lattice class. This class contains the general lattice model.
Parameters:
dim (int, optional) – Dimension of the lattice (1, 2, or 3). If None, inferred from lx, ly, lz.
lx (int, optional) – Length of the lattice in the x-direction.
ly (int, optional) – Length of the lattice in the y-direction.
lz (int, optional) – Length of the lattice in the z-direction.
bc (str, optional) – Boundary conditions (e.g., ‘PBC’, ‘OBC’).
adj_mat (np.ndarray, optional) – Adjacency matrix for the lattice.
flux (np.ndarray, optional) – Flux piercing the boundaries. This can be a dictionary specifying the
flux in each direction, or a single value applied to all directions. Importantly,
this automatically implies TWISTED boundary conditions, so the bc parameter can be left as None or set to ‘TWISTED’ for clarity.
Real-space position for stored coordinate tuple (x,y,z).
The base class calculate_coordinates already stores proper
vectors via _a1,_a2,_a3,_basis. This helper is kept for
backwards compatibility and any custom coordinate look-ups.
Calculate next-nearest neighbours for the armchair hexagonal lattice.
NNN connect sites within the same sublattice. Each site has 6
NNN in the full 2D honeycomb; finite clusters may have fewer
depending on boundary conditions.
NNN displacements (same sublattice) in cell coordinates:
The honeycomb lattice is a 2D lattice with a hexagonal structure. The lattice consists of
two sublattices (A and B) arranged in a hexagonal pattern. Nearest and next-nearest neighbors
are computed based on a hexagonal unit cell.
High-symmetry points in the Brillouin zone:
- Gamma:
Zone center at (0, 0)
K:
Dirac point at (2/3, 1/3) - hosts linear band crossings in graphene
Return high-symmetry points for the honeycomb lattice.
Returns:
High-symmetry points for the hexagonal Brillouin zone:
- Γ (Gamma): Zone center (0, 0)
- K: Dirac point at (2/3, 1/3) - hosts linear band crossings
- K’: Other Dirac point at (1/3, 2/3)
- M: Edge midpoint at (1/2, 0)
Returns the nearest neighbor in the specified direction.
For the honeycomb lattice, we choose a mapping:
LatticeDirection.X -> neighbor at index 0 of _nn[site]
LatticeDirection.Y -> neighbor at index 1 of _nn[site]
LatticeDirection.Z -> neighbor at index 2 of _nn[site]
Calculates the nearest neighbors (NN) using boundary conditions.
The implementation uses a helper function to apply periodic or open boundary conditions.
For 2D, for example, we use a different treatment on even and odd indices.
Calculates the next-nearest neighbors (NNN) of the honeycomb lattice.
NNN are second-nearest neighbors, connecting sites on the same sublattice.
For sublattice A (even sites), the three NNN directions are obtained by
composing two consecutive NN hops:
NNN_1: Y-bond then X-bond^{-1} → cell shift (0, -1) [down in y]
NNN_2: Z-bond then Y-bond^{-1} → cell shift (-1, 0) [left in x]
NNN_3: X-bond then Z-bond^{-1} → cell shift (+1, +1) [diagonal]
For sublattice B (odd sites), the shifts are inverted.
Honeycomb (graphene-like) nearest-neighbour dispersion magnitude.
Computes |f(k)| where f = sum_{δ} exp(-i k·δ) for the three A->B vectors
used by this lattice implementation.
Lattice backed by an adjacency matrix instead of structural formulas.
Parameters:
adjacency – Square (NsimesNs) adjacency matrix. Non-zero entries denote
connected pairs. The absolute value of the weight is used to rank
neighbor order (highest -> nearest).
coords – Optional vertex embedding of shape (Ns,dim). Defaults to a 1D
chain ordering. Used for plotting and distance heuristics.
bc – Boundary condition descriptor. Defaults to open boundaries.
flux – Optional boundary flux phases - forwarded to Lattice.
General Lattice class. This class contains the general lattice model.
Parameters:
dim (int, optional) – Dimension of the lattice (1, 2, or 3). If None, inferred from lx, ly, lz.
lx (int, optional) – Length of the lattice in the x-direction.
ly (int, optional) – Length of the lattice in the y-direction.
lz (int, optional) – Length of the lattice in the z-direction.
bc (str, optional) – Boundary conditions (e.g., ‘PBC’, ‘OBC’).
adj_mat (np.ndarray, optional) – Adjacency matrix for the lattice.
flux (np.ndarray, optional) – Flux piercing the boundaries. This can be a dictionary specifying the
flux in each direction, or a single value applied to all directions. Importantly,
this automatically implies TWISTED boundary conditions, so the bc parameter can be left as None or set to ‘TWISTED’ for clarity.
Calculates the allowed reciprocal-space k-vectors (momentum grid)
consistent with the lattice size and primitive reciprocal vectors.
When boundary fluxes are present (TWISTED BC), the fractional
coordinates are shifted by \(\phi_\mu / (2\pi L_\mu)\) in
each direction, so that the Bloch condition matches the twisted
boundary.
The sampling follows the same fftfreq ordering used by the Bloch
transform (Γ at index [0,0,0], followed by positive frequencies and
finally the negative branch). This keeps the analytic grids aligned
with the numerically constructed H(k) blocks.
i = (R, beta) real-space cell R and sublattice beta
n = (k, alpha) k-point k and sublattice alpha
Elements:
$$
F_{(k,alpha),(R,beta)} =
1/sqrt(Nc) * delta_{alpha,beta} * exp(-i k . R).
$$
This is unitary:
$$
F^dagger F = I_{Ns}, F F^dagger = I_{Ns},
$$
where Ns = Nc * Nb is the total number of sites, Nc is the number of
unit cells, and Nb is the number of sublattices.
Important
When boundary fluxes are present (TWISTED BC), the k-grid used to
build the DFT matrix is shifted by phi_mu/(2piL_mu) in each
direction, exactly as in calculate_k_vectors().
Note that this DFT matrix does not include basis-dependent phases
(i.e., exp(-i k . r_basis)).
Calculates the Discrete Fourier Transform (DFT) matrix for the lattice.
This method can be optimized using FFT (Fast Fourier Transform) in the future.
Reference: https://en.wikipedia.org/wiki/DFT_matrix
Enumeration of standard high-symmetry paths in the Brillouin zone.
We define the k-space paths in a general representation of momentum vectors:
[
k = f1 * b1 + f2 * b2 + f3 * b3,
]
where (b1, b2, b3) are the reciprocal lattice vectors, and (f1, f2, f3) are the fractional coordinates,
f_i = n_i / N_i, with n_i = 0, 1, …, N_i - 1 for each direction i.
Each value returns a list of (label, fractional_coord) pairs.
The fractional coordinates are expressed in units of reciprocal lattice vectors.
lz (int) – Number of sites in z-direction (ignored if dim < 3)
bc – Boundary condition (e.g., LatticeBC.PBC or LatticeBC.OBC)
flux – Optional boundary flux specification forwarded to the lattice
constructor. Accepts a scalar phase, BoundaryFlux, or a
mapping from directions to phases (in radians).
Include:
- point_edgecolor: Color of the marker edges (default “white”).
- point_zorder: Z-order for the scatter points (default 5).
- color_extended: Color for translated copies (default “C2”).
- edgecolor_extended: Edge color for translated copies (default “gray”).
- marker_extended: Marker for translated copies (default marker).
- Any other valid arguments for ax.scatter.
title_kwargs (dict, optional) – Additional keyword arguments for the title.
tight_layout (bool) – If True, applies tight layout to the figure.
elev (float, optional) – Elevation and azimuth angles for 3D plots.
azim (float, optional) – Elevation and azimuth angles for 3D plots.
partition_colors (tuple of str, optional) – Colors to use for bipartite/sublattice coloring. If provided, nodes are
colored based on sublattice parity.
show_periodic_connections (bool) – If True, indicates wrap-around connections textually or graphically.
show_primitive_cell (bool) – If True, overlays the primitive unit cell vectors/box.
**scatter_kwargs – Additional arguments passed to ax.scatter.
show_bz (bool, default=True) – Draw the first Brillouin zone.
show_path (bool, default=True) – Draw the ideal high-symmetry path.
show_matched_kpoints (bool, default=True) – Highlight sampled k-points whose distance to the path is within the
matching tolerance.
points_per_seg (int, default=40) – Number of interpolation points per path segment for the ideal path.
path_match_tol (float, optional) – Distance tolerance used when highlighting mesh points near the
drawn path.
extend (bool, default=False) – Draw translated copies of the sampled k-mesh.
extend_copies (int or iterable[int], optional) – Number of reciprocal-cell copies per direction. In 2D,
extend_copies=1 includes the first shell around the first Brillouin
zone and extend_copies=2 includes the second shell as well.
Plot subsystem sweep showing cuts with different boundary sizes.
Creates a grid of subplots showing subsystems grouped by ∂A.
Parameters:
direction (str, optional) – Direction for sweep (‘x’, ‘y’, ‘z’). Creates full-width cuts.
rectangular (bool, default=False) – If True and direction is None, use rectangular subsystems (various shapes).
If False, use lexicographic sweep (sequential site addition).
max_panels (int, default=6) – Maximum number of panels to show.
figsize (tuple, optional) – Figure size. Auto-computed if None.
figsize (tuple, optional) – Figure size in inches (width, height).
fix_aspect (bool, default=True) – If True, preserve equal axis scaling in 2D plots. Set to False to
let the requested figsize control the on-screen aspect.
Include:
- point_edgecolor: Color of the marker edges (default “white”).
- point_zorder: Z-order for the scatter points (default 5).
- color_extended: Color for translated copies (default “C2”).
- edgecolor_extended: Edge color for translated copies (default “gray”).
- marker_extended: Marker for translated copies (default marker).
- Any other valid arguments for ax.scatter.
Plot the Brillouin Zone approximation based on sampled k-vectors.
Parameters:
lattice (Lattice) – The lattice object containing k-vectors.
ax (Axes, optional) – Matplotlib axes to plot on. If None, a new figure is created.
facecolor (str, default="tab:blue") – Color to fill the Brillouin Zone area.
edgecolor (str, default="black") – Color for the Brillouin Zone boundary.
alpha (float, default=0.25) – Transparency level for the Brillouin Zone fill.
figsize (tuple, optional) – Figure size in inches (width, height).
fix_aspect (bool, default=True) – If True, preserve equal axis scaling in 2D plots. Set to False to
let the requested figsize control the on-screen aspect.
Handles region definitions and extractions for a Lattice.
Usage:
handler=LatticeRegionHandler(lattice)# or equivalently via the lattice shortcut:sites=lattice.regions.get_region('disk',origin=10,radius=3.0)kp=lattice.regions.get_region('kitaev_preskill',radius=5.0)lw=lattice.regions.get_region('levin_wen',inner_radius=2,outer_radius=5)
origin (int or list[float], optional) – Center of the region (site index or coordinate).
radius (float, optional) – Radius for 'disk' and 'kitaev_preskill' regions.
direction (str, optional) – Direction for 'half' cuts ('x', 'y', 'z').
sublattice (int, optional) – Sublattice index for 'sublattice' regions.
sites (list[int], optional) – Explicit list of sites for 'custom' regions.
depth (int, optional) – Hop-distance for 'graph' balls.
plaquettes (list[int], optional) – Plaquette indices for 'plaquette' regions.
configuration (int, optional) – Legacy predefined configuration index (1-based) for the given
lattice and kind.
predefined (bool | int | str, optional) – Convenience selector for predefined regions.
- True : list available predefined entries for given lattice/kind.
- int : select predefined entry by 0-based index.
- str : select predefined entry by label.
as_region (bool, optional) – If True (default), return Region-class objects for supported kinds.
If False, keep legacy list/dict return shapes.
min_size (int, optional) – If provided, return None if the region (or any of its A, B, C parts)
is smaller than this value.
n_sectors (int) – Number of angular sectors for 'kitaev_preskill' (default 3).
rotation (float) – Rotation of sector boundaries (radians) for KP (default 0).
use_pbc (bool) – Whether to use PBC-wrapped distances for KP/LW. Default
False — regions must not wrap around the torus boundary.
region (str) – Which single region key to return for KP/LW (e.g. 'A',
'AB'). If given, a flat list is returned instead of a dict.
Returns:
Sorted site list for simple regions, a dict of labelled site
lists for topological partitions, metadata entries when predefined=True,
or Region objects when as_region=True.
The rotation parameter (radians) rotates all sector boundaries
counter-clockwise.
Warning
use_pbc defaults to False. For the TEE construction the
regions must not wrap around the boundary — otherwise a
single angular sector may pick up sites from the opposite side of
the torus and the linear combination becomes ill-defined.
Parameters:
origin (int or array-like, optional) – Center of the disk. Defaults to the centroid of all sites.
radius (float, optional) – Disk radius. Defaults to min(Lx,Ly)/2-0.5 (or 1.0).
n_sectors (int) – Number of angular sectors (default 3).
rotation (float) – Global rotation of sector boundaries in radians (default 0).
use_pbc (bool) – If True, use minimum-image (PBC-wrapped) displacements when
computing distances and angles. Default False — raw
Euclidean distances are used so the regions cannot wrap around
the boundary.
Returns:
Keys: single-sector labels 'A', 'B', … and all pairwise
/ triple unions ('AB', 'BC', 'ABC', …).
The sectors must tile a disk without overlapping interiors so that
the linear combination isolates the topological contribution -γ.
The rotation parameter lets you sweep the partition to check
independence.
B (annulus) : inner_radius ≤ distance < outer_radius
C (exterior) : distance ≥ outer_radius
Warning
use_pbc defaults to False. Wrapping would make sites
from the far side of the torus appear close to the origin,
contaminating the annular regions and invalidating the TEE
linear combination.
Parameters:
origin (int or array-like, optional) – Centre of the annuli. Defaults to the centroid of all sites.
inner_radius (float, optional) – Boundary between A and B. Default 1.0.
outer_radius (float, optional) – Boundary between B and C. Default min(Lx,Ly)/2-0.5.
use_pbc (bool) – If True, use minimum-image (PBC-wrapped) displacements.
Default False — raw Euclidean distances prevent wrap-around.
Returns:
Keys: 'A', 'B', 'C' and unions 'AB', 'BC',
'AC', 'ABC'.
The boundary dA is the number of bonds (edges) that cross between
subsystem A and its complement B. This is the quantity that appears
in the area law of entanglement entropy: S ~ dA for gapped systems.
Parameters:
subsystem (int, list, or array) – Specification of subsystem A:
- If int: First subsystem sites (contiguous from 0)
- If list/array: Explicit site indices
include_nnn (bool, default=False) – If True, include next-nearest-neighbor bonds in the boundary count.
return_bonds (bool, default=False) – If True, also return the list of boundary bonds.
Returns:
dA (int) – Number of bonds crossing the boundary (|dA|).
boundary_bonds (list of (int, int), optional) – If return_bonds=True, the list of (i, j) bonds where i in A, j in B.
Examples
>>> # Compute boundary for half-system cut>>> lattice=SquareLattice(Lx=4,Ly=4,bc='obc')>>> half_sites=lattice.regions.region_half('x').A>>> dA=lattice.regions.subsystem_boundary(half_sites)>>> print(f"Boundary size: {dA}")# Should be Ly for x-cut
>>> # Get boundary bonds explicitly>>> dA,bonds=lattice.regions.subsystem_boundary([0,1,2],return_bonds=True)>>> print(f"Boundary bonds: {bonds}")
Notes
For area-law scaling in d dimensions, S ~ L^{d-1} where L is linear size.
The boundary dA counts bonds, which for regular lattices is proportional
to the surface area of the subsystem.
For periodic boundary conditions, bonds wrapping around the system are
counted correctly using the lattice connectivity.
Creates subsystems and organizes them by their boundary size
(number of bonds cut). Useful for area-law scaling studies.
Parameters:
direction (str, optional) – Direction for sweep: ‘x’, ‘y’, ‘z’. Creates full-width cuts.
If None, uses rectangular or lexicographic mode.
rectangular (bool, default=False) – If True (and direction is None), generate rectangular subsystems
of all sizes (1x1, 1x2, 2x1, 2x2, …). Gives variety in shapes.
If False, use lexicographic sweep (sequential site addition).
by_unit_cell (bool, default=True) – If True, grow by unit cells (for multi-sublattice lattices).
Abstract Base Class for defining lattice structures.
This class serves as the foundation for all lattice implementations in the lattices module.
It handles geometry, connectivity, boundary conditions, and k-space properties.
Lattice sites are indexed linearly from 0 to Ns-1.
The mapping from spatial coordinates to linear index depends on the concrete implementation,
but typically follows a row-major (lexicographic) order:
General Lattice class. This class contains the general lattice model.
Parameters:
dim (int, optional) – Dimension of the lattice (1, 2, or 3). If None, inferred from lx, ly, lz.
lx (int, optional) – Length of the lattice in the x-direction.
ly (int, optional) – Length of the lattice in the y-direction.
lz (int, optional) – Length of the lattice in the z-direction.
bc (str, optional) – Boundary conditions (e.g., ‘PBC’, ‘OBC’).
adj_mat (np.ndarray, optional) – Adjacency matrix for the lattice.
flux (np.ndarray, optional) – Flux piercing the boundaries. This can be a dictionary specifying the
flux in each direction, or a single value applied to all directions. Importantly,
this automatically implies TWISTED boundary conditions, so the bc parameter can be left as None or set to ‘TWISTED’ for clarity.
Initializes the lattice object by calculating coordinates, reciprocal vectors, and neighbor lists.
This method performs the following steps:
1. Calculates the real-space coordinates, r-vectors, and k-vectors of the lattice.
2. If the number of sites (self.Ns) is less than 100, computes the discrete Fourier transform (DFT) matrix.
3. If an adjacency matrix (self._adj_mat) is provided:
Determines the number of sites (Ns) from the adjacency matrix.
For each site, identifies nearest neighbors (nn) as those connected by the highest weight in the adjacency matrix, and next-nearest neighbors (nnn) as those connected by the next highest distinct weight.
Stores forward neighbors (indices greater than the current site) for both nn and nnn.
If no adjacency matrix is provided, calculates nearest and next-nearest neighbors using default methods.
5. Calculates normalization or symmetry properties of the lattice.
This method sets up all necessary neighbor lists and lattice properties required for further computations.
Return a list of site indices defining a spatial region.
Parameters:
kind (str or RegionType) – Type of region: ‘half’, ‘disk’, ‘sublattice’, ‘graph’, ‘plaquette’, ‘custom’.
We also support specific half cuts like ‘half_x’, ‘half_y’, ‘half_z’ for convenience.
origin (int or list[float], optional) – Center of the region. Can be a site index or coordinate vector.
radius (float, optional) – Radius for ‘disk’ regions.
direction (str, optional) – Direction for ‘half’ cuts (‘x’, ‘y’, ‘z’).
sublattice (int, optional) – Sublattice index for ‘sublattice’ regions.
sites (list[int], optional) – Explicit list of sites for ‘custom’ regions.
depth (int, optional) – Depth/distance for ‘graph’ regions.
plaquettes (list[int], optional) – List of plaquette indices for ‘plaquette’ regions.
Returns:
Sorted list of site indices belonging to the region.
Generate space-group permutation table for this lattice.
Delegates to generate_space_group_perms().
When TWISTED boundary conditions are active, the point-group part is
disabled (only translations are returned) because a generic flux
breaks point-group symmetry unless the flux respects it.
Parameters:
point_group (str) – 'full' for maximal point group, 'translations' for translations only.
Return a dictionary describing the spatial symmetries of this lattice.
The information is consistent for both single-particle and many-body
representations. When TWISTED boundary conditions are present the
point-group part is absent (flux generically breaks it).
Returns:
Keys:
- 'lattice_type' : LatticeType enum
- 'sites_per_cell' : int
- 'n_cells' : number of unit cells
- 'dim' : spatial dimension
- 'bc' : boundary condition enum
- 'is_periodic' : (bool, bool, bool) per direction
- 'is_twisted' : bool
- 'translation_group' : ZL_x x ZL_y (as tuple (Lx,Ly))
- 'point_group' : str or None ('D4' for square Lx==Ly, etc.)
- 'space_group_order' : total number of space-group elements
- 'flux' : BoundaryFlux or None
Return the sublattice indices of the lattice sites.
For a Bravais lattice, this would simply be an array of zeros.
For a non-Bravais lattice, this would indicate which sublattice each site belongs to.
Return the unit cell coordinates of the lattice sites. For a Bravais lattice,
this would simply be the integer coordinates of the unit cells.
For a non-Bravais lattice, this would include the basis vectors as well.
Return the sublattice index for a given site.
By default, returns 0 for all sites (single sublattice).
Override in subclasses for multi-sublattice lattices.
Extend k-space points and optional data across translated Brillouin zones.
The helper works for arbitrary k-space dimensions and any number of
reciprocal translation vectors. Legacy b1/b2/b3 with
nx/ny/nz remain supported for existing callers.
Allows to generate extended k-point grids for plotting band structures along high-symmetry paths…
Parameters:
k_points (ndarray, shape (N, dim)) – Array of k-points in reciprocal space to be extended.
data (ndarray, shape (N, ...) or None) – Optional data associated with each k-point (e.g. eigenvalues) to be extended alongside the k-points. Must have the same leading dimension as k_points.
copies (int or iterable of ints, optional) – Number of translated copies to generate in each reciprocal direction. If an integer is provided, the same number of copies will be generated in all directions. If an iterable is
provided, it should have a length equal to the number of reciprocal lattice vectors (e.g. 3 for 3D), specifying the number of copies in each direction separately.
**kwargs – Additional keyword arguments to pass to the underlying ws_extend function. See its documentation for details.
Returns:
extended_k_points (ndarray, shape (M, dim)) – Array of extended k-points in reciprocal space, including the original points and their translated copies
extended_data (ndarray, shape (M, …) or None) – Extended data associated with each k-point, if the input data was provided. Otherwise, None
Return a boolean mask for the Wigner-Seitz cell in reciprocal space.
This can be used to identify which k-points lie within the first Brillouin zone.
Parameters:
Kx (array-like) – Arrays of k-point coordinates in reciprocal space. This is a grid
of k-points for which we want to determine if they lie within the Wigner-Seitz cell.
Ky (array-like) – Arrays of k-point coordinates in reciprocal space. This is a grid
of k-points for which we want to determine if they lie within the Wigner-Seitz cell.
Kz (array-like) – Arrays of k-point coordinates in reciprocal space. This is a grid
of k-points for which we want to determine if they lie within the Wigner-Seitz cell.
shells (int) – Number of shells of Wigner-Seitz cell to include in the mask.
tol (float) – Tolerance for determining if a point is within the Wigner-Seitz cell, accounting for numerical precision issues.
**kwargs – Additional keyword arguments to pass to the underlying ws_bz_mask function. See its documentation for details.
Resolve path input to a list of (label, fractional_coord) pairs.
Parameters:
path (list[(label, coords)], StandardBZPath, str, List[str], or HighSymmetryPoints) – Path definition (fractional coordinates), standard enum, enum name string,
list of point labels, or HighSymmetryPoints object.
lattice (Lattice, optional) – Lattice object used to resolve labels if path is a list of strings.
Returns:
resolved_path – Resolved path as a list of (label, fractional_coord) pairs.
Return True if the lattice momentum grid contains a special point.
This method helps to check whether a finite lattice
contains a particular high-symmetry point in the Brillouin zone,
which is important for band structure calculations and topological analyses.
Parameters:
point – Special point identifier. Accepted forms:
- label string (e.g. "Gamma", "K", "K'"),
- HighSymmetryPoint,
- explicit fractional coordinate tuple/array.
tol (float) – Absolute tolerance used in the coordinate match.
Notes
The check is done in fractional reciprocal coordinates and naturally
includes flux-induced shifts from twisted boundary conditions because it
uses self.kvectors_frac.
path (list of str, str, StandardBZPath, or None) – Path specification. Can be:
- List of high-symmetry point names: [‘Gamma’, ‘X’, ‘M’, ‘Gamma’]
- StandardBZPath enum or string: ‘SQUARE_2D’
- None: use default path for this lattice
points_per_seg (int) – Number of interpolated points per path segment.
Returns:
k_path (np.ndarray, shape (Npath, 3)) – Cartesian k-points along the path.
k_dist (np.ndarray, shape (Npath,)) – Cumulative distance for plotting x-axis.
labels (List[Tuple[int, str]]) – Indices and labels for high-symmetry points.
k_path_frac (np.ndarray, shape (Npath, 3)) – Fractional k-coordinates along the path.
Example
>>> lattice=SquareLattice(dim=2,lx=4,ly=4)>>> k_path,k_dist,labels,k_frac=lattice.bz_path()>>> # Or with custom path:>>> k_path,k_dist,labels,k_frac=lattice.bz_path(['Gamma','M','Gamma'])
Build an ideal Brillouin-zone path and optionally match it to an existing k-grid.
If no k-grid is provided, the returned object still contains the continuous
path geometry, which is useful for plotting or for constructing a path that
is not constrained to the sampled reciprocal mesh. When a sampled grid is
provided, reciprocal-lattice copies are generated automatically as needed
so paths in extended Brillouin-zone regions can still match the existing
data.
Parameters:
path (list of str, str, StandardBZPath, or None) – Path specification. Can be:
- List of high-symmetry point names: [‘Gamma’, ‘X’, ‘M’, ‘Gamma’]
- StandardBZPath enum or string: ‘SQUARE_2D’
- None: use default path for this lattice
points_per_seg (int) – Number of interpolated points per path segment.
k_vectors (np.ndarray, shape (Nk, 3), optional) – Cartesian k-vectors of the existing grid to match against.
k_vectors_frac (np.ndarray, shape (Nk, 3), optional) – Fractional k-vectors of the existing grid to match against. Required if k_vectors is provided.
tol (float) – Tolerance for matching path points to the existing k-grid. With
periodic=True it is interpreted in fractional reciprocal
coordinates. With periodic=False it is interpreted in plotted
Cartesian reciprocal coordinates.
periodic (bool, default=True) – If True, allow reciprocal-translation-equivalent points to match.
Set to False for visual matching in the displayed Brillouin-zone copy.
Extract k-path data from a k-grid using fractional coordinate matching.
This function finds the closest k-points on the actual grid to an ideal path
through high-symmetry points. It handles periodic boundary conditions in k-space
and automatically reuses reciprocal-lattice copies of the sampled grid when the
requested path lies in an extended Brillouin-zone region.
It also allows to return a structured KPathResult dataclass or a tuple…
Parameters:
lattice (Lattice) – Lattice object with reciprocal lattice vectors
k_vectors_frac (np.ndarray, shape (..., 3)) – Fractional coordinates of k-points (will be flattened)
values (np.ndarray) – Data values sampled on the k-grid. The k-grid axes may appear as
(Lx,Ly,Lz,...) or after leading batch axes such as time or
frequency, e.g. (Nw,Lx,Ly,Lz) or (Nw,Lx,Ly,Lz,...).
A single flattened k-grid axis of length Nk is also supported.
path (various, optional) – Path specification. Can be:
- StandardBZPath enum value (e.g., StandardBZPath.SQUARE_2D)
- String name (e.g., ‘SQUARE_2D’)
- List of (label, [f1,f2,f3]) tuples
- HighSymmetryPoints object (uses default path)
- None: uses lattice’s default path if available
points_per_seg (int) – Number of interpolated points per path segment
return_result (bool) – If True (default), return KPathResult dataclass.
If False, return tuple for backwards compatibility.
Returns:
If return_result=True: KPathResult dataclass with all path data.
The returned values preserve any leading batch axes and replace
the k-grid axes with a path axis.
If return_result=False: (k_cart, k_frac, k_dist, labels, values) tuple
>>> # Using default path from HighSymmetryPoints>>> result=bz_path_data(lattice,k_grid,k_frac,energies,HighSymmetryPoints.square_2d())>>> plt.plot(result.k_dist,result.values)
>>> # Using standard path enum>>> result=bz_path_data(lattice,k_grid,k_frac,energies,'SQUARE_2D')
Check if a set of sites spans the lattice (non-contractible on a torus).
This method uses a BFS-based winding number tracking on the induced subgraph of
the provided site indices. If any loop with a non-zero winding number
along a periodic direction is found, the set is considered spanning.
Return the complex hopping phase factor for the bond \(i \to j\).
For bonds that do not cross a periodic boundary, this is 1.
For boundary-crossing bonds under TWISTED BC, the phase is
\(\exp(i\,\phi_\mu)\) for each direction \(\mu\) in which
the bond wraps.
This is the factor that should multiply the bare hopping amplitude
in real-space Hamiltonian construction.
Build an \(N_s \times N_s\) matrix of complex hopping amplitudes
that includes the Peierls phases from boundary fluxes.
Diagonal is zero. Off-diagonal H[i,j]=t_{ij}*phase(i->j)
where t_{ij}=1 for all connected pairs and phase is the
product of boundary phases along directions that the bond wraps.
Parameters:
include_nnn (bool) – If True, include next-nearest-neighbor hoppings as well.
orientation ({'anticlockwise', 'clockwise', None}, optional) – If provided, will sort/choose based on geometric angle.
Default: None (return all middle sites).
Returns:
List of middle-site indices (can be 0, 1, or 2 elements).
Return the displacement i->j with optional PBC minimum-image wrapping.
Parameters:
i (int or tuple) – Site indices or explicit coordinates.
j (int or tuple) – Site indices or explicit coordinates.
minimum_image (bool, default=False) – If True, wrap each periodic direction to the shortest displacement.
real_space (bool, default=False) – If True and i,j are site indices, return displacement in real-space
vectors (uses displacement()). Otherwise use lattice coordinates.
i = (R, beta) real-space cell R and sublattice beta
n = (k, alpha) k-point k and sublattice alpha
Elements:
$$
F_{(k,alpha),(R,beta)} =
1/sqrt(Nc) * delta_{alpha,beta} * exp(-i k . R).
$$
This is unitary:
$$
F^dagger F = I_{Ns}, F F^dagger = I_{Ns},
$$
where Ns = Nc * Nb is the total number of sites, Nc is the number of
unit cells, and Nb is the number of sublattices.
Important
When boundary fluxes are present (TWISTED BC), the k-grid used to
build the DFT matrix is shifted by phi_mu/(2piL_mu) in each
direction, exactly as in calculate_k_vectors().
Note that this DFT matrix does not include basis-dependent phases
(i.e., exp(-i k . r_basis)).
Calculates the Discrete Fourier Transform (DFT) matrix for the lattice.
This method can be optimized using FFT (Fast Fourier Transform) in the future.
Reference: https://en.wikipedia.org/wiki/DFT_matrix
Returns the matrix with indcies corresponding to a slice from the QMC.
A usefull function for reading the position Green’s function saved from:
@url https://github.com/makskliczkowski/DQMC
The Green’s functions are saved in the following manner. If cut is True, data
has (2L_i - 1) possible position differences, otherwise we skip the negative ones and use L_i.
For 1D simulation: 1 column and (2 * Lx - 1) rows for possition differences (-Lx, -Lx + 1, …, 0, …, Lx)
For 2D simulation: (2 * Lx - 1) rows for possition differences (-Lx, -Lx + 1, …, 0, …, Lx) and (2 * Ly - 1) columns for possition differences (-Ly, -Ly + 1, …, 0, …, Ly)
For 3D simulation: Same as in 2D but after (2 * Lx - 1) x (2 * Ly - 1) matrix has finished, a new slice for Lz appears for next columns Lz * (2*Ly - 1)
- cut : if true (2L_i - 1) possible position differences, otherwise we skip the negative ones and use L_i.
Calculates the allowed reciprocal-space k-vectors (momentum grid)
consistent with the lattice size and primitive reciprocal vectors.
When boundary fluxes are present (TWISTED BC), the fractional
coordinates are shifted by \(\phi_\mu / (2\pi L_\mu)\) in
each direction, so that the Bloch condition matches the twisted
boundary.
The sampling follows the same fftfreq ordering used by the Bloch
transform (Γ at index [0,0,0], followed by positive frequencies and
finally the negative branch). This keeps the analytic grids aligned
with the numerically constructed H(k) blocks.
For TWISTED boundary conditions the neighbor connectivity is
identical to PBC — the flux phases are applied separately when
building the Hamiltonian or the DFT matrix.
Calculates the Wilson loops (non-contractible loops) for the lattice based on its boundary conditions.
Returns a list of lists, where each inner list contains the site indices of a Wilson loop.
Assumes standard lexicographic site indexing (x + y*Lx + z*Lx*Ly).
Like calculate_nn(), each calculate_nnn_in implementation
is expected to set self._nnn (and optionally self._nnn_forward)
directly. The return value—if any—is stored as a fallback.
Logs the neighbors of each site in the lattice using the provided logger.
For each site in the lattice, this method retrieves its nearest neighbors and logs their indices.
Additionally, for each neighbor, it logs detailed information using a higher verbosity level.
Parameters:
logger – An object with an info method for logging messages. The info method should accept
parameters lvl (int) for verbosity level and color (str) for message color.
Logs the forward nearest neighbors for each site in the lattice.
For each site in the lattice, this method retrieves the number of forward nearest neighbors
and logs their indices using the provided logger. The method outputs two levels of information:
- Level 1 (green): Lists the neighbors of each site.
- Level 2 (blue): Details each neighbor’s index for the site.
Parameters:
logger (A logging object with an info method that accepts a message,) – a logging level (lvl), and a color (color).
Inverse Bloch transform: reconstruct real-space matrix from k-space blocks.
This is the exact inverse of kspace_from_realspace(). It reconstructs the
real-space Hamiltonian from momentum-space blocks using the inverse Fourier transform:
Example 2: Inverse transform without explicit kgrid
>>> # If kgrid is omitted, it's reconstructed using fftfreq convention>>> H_real_recon=lattice.realspace_from_kspace(H_k)>>> np.allclose(H_real,H_real_recon)# True
Custom k-point sampling (only used when block_diag=True):
If None: Uses automatic full Brillouin zone grid based on lattice size
(recommended for most use cases)
If provided: Array of shape (Nk,3) with custom k-points in Cartesian coordinates
Returns (H_k,kpoints) with H_k shape (Nk,Nb,Nb)
unitary_norm (bool, default=True) – Use unitary normalization \(1/\sqrt{N_c}\) for the Bloch transform.
If False, uses normalization \(1/N_c\) instead. Keep True for
standard quantum mechanics convention preserving operator norms.
If True, also return the Bloch unitary matrix W used for the transformation.
This is useful for transforming additional operators or computing correlation functions.
Note: Only available when block_diag=True. The unitary is returned as a 4th
output value with shape (Lx,Ly,Lz,Ns,Nb) or (Nk,Ns,Nb) if custom
k-points are provided.
Returns:
**Case 1 (block_diag=False (default)**) –
H_k_fullnp.ndarray
Full transformed matrix with shape (Ns,Ns). This is the complete DFT
of the input matrix, preserving all information.
Example 4: Get Bloch unitary for operator transforms
>>> H_k,k_grid,k_frac,W=lattice.kspace_from_realspace(... H_real,block_diag=True,return_transform=True... )>>> # Transform another operator using the same W>>> O_k=np.einsum('kia,ij,kjb->kab',W.conj(),O_real,W)
Notes
Periodic boundary conditions (PBC) are assumed for the Bloch transform.
The method assumes translational invariance of the system, which ensures
the spectrum of H_real equals the union of spectra of H(k) blocks.
- For the full grid (kpoints=None), the k-points follow the fftfreq convention
with the Γ-point initially at index [0,0,0], then shifted to the center.
- Site ordering is arbitrary; the method uses the lattice geometry (coordinates + basis)
to correctly identify sublattices and apply phases.
- For sparse input matrices, automatic conversion to dense format is performed.
Convert a real-space correlation matrix into a momentum-resolved structure factor.
This is a convenience wrapper around the basis-aware Bloch projector in
QES.general_python.lattices.tools.lattice_kspace.kspace_from_realspace.
The real-space input mat is first transformed into the multipartite
k-space block representation evaluated on self.kvectors
where R,R' label unit cells and alpha,beta label basis sites
inside the unit cell. The reduction argument then decides how this
multipartite object is converted into a scalar structure factor at each
sampled momentum q.
Parameters:
mat (np.ndarray) – Real-space correlation or operator matrix with shape (Ns,Ns)
or batched shape (...,Ns,Ns). Any leading axes, e.g. time,
frequency, disorder sample, or state index, are preserved.
C(q) with shape (Lx,Ly,Lz,Nb,Nb) (i.e., no reduction).
"sum":
return sum_{alpha,beta}C_{alphabeta}(q) (i.e., sum over all entries of each block).
"trace":
return sum_alphaC_{alphaalpha}(q) (i.e., sum over diagonal entries of each block).
"mean":
return the arithmetic mean of all multipartite block entries at each q (i.e., sum over all entries and divide by Nb^2).
"diag":
return the eigenvalues of each block, which can be useful for identifying dominant modes or instabilities. The output shape will be (Lx,Ly,Lz,Nb) since each block’s eigenvalues are returned as a vector of length Nb.
norm ({"none", "cell", "site"}, default="none") –
Optional post-normalization of the returned k-space quantity:
"none":
keep the raw Bloch-projector normalization, i.e. the blocks
C(q) defined above with the prefactor 1/N_c.
"cell":
alias for "none" kept for readability when you want to
emphasize unit-cell normalization.
"site":
divide the returned blocks or reduced values by the number of
basis sites N_b. For scalar reductions such as "sum", this
converts the default unit-cell normalization into the more common
site normalization 1/N_s used in
\(S(q) = \langle O_{-q} O_q \rangle\).
Returns:
values (np.ndarray) – Momentum-resolved structure factor. For input shape (...,Ns,Ns)
the output shape is:
(...,Lx,Ly,Lz,Nb,Nb) for reduction="none"
(...,Lx,Ly,Lz,Nb) for reduction="diag"
(...,Lx,Ly,Lz) for "sum", "trace", or "mean"
For a single input matrix (Ns,Ns), the leading ... is absent.
k_grid (np.ndarray) – Cartesian sampled k-grid with shape (Lx,Ly,Lz,3).
k_frac (np.ndarray) – Fractional sampled k-grid with shape (Lx,Ly,Lz,3).
Notes
Use reduction="none" when sublattice-resolved information matters.
Use one of the scalar reductions when you want a single value per
momentum that can be fed directly into bz_path_data.
The default norm="none" preserves the existing unit-cell
normalization. For comparisons against structure factors built from
Fourier-transformed site operators, norm="site" is typically the
physically relevant choice.
Examples
>>> Sq,k_grid,k_frac=lattice.structure_factor(corr_zz,reduction="sum")>>> path=lattice.bz_path_data(k_grid,k_frac,Sq,path=['Gamma','K','M','Gamma'])>>>>>> # Frequency-resolved data with shape (Nw, Ns, Ns)>>> Sqw,k_grid,k_frac=lattice.structure_factor(corr_zz_w,reduction="sum")>>> # Sqw has shape (Nw, Lx, Ly, Lz)
Include:
- point_edgecolor: Color of the marker edges (default “white”).
- point_zorder: Z-order for the scatter points (default 5).
- color_extended: Color for translated copies (default “C2”).
- edgecolor_extended: Edge color for translated copies (default “gray”).
- marker_extended: Marker for translated copies (default marker).
- Any other valid arguments for ax.scatter.
Convenience wrapper returning the matplotlib figure and axes for a detailed lattice structure plot.
Parameters:
show_indices (bool) – If True, annotates nodes with their site indices.
highlight_boundary (bool) – If True, draws boundary nodes with a distinct color/edge.
show_axes (bool) – If False, hides the coordinate axes for a cleaner diagram.
partition_colors (tuple of str, optional) – Colors to use for bipartite/sublattice coloring. If provided, nodes are
colored based on sublattice parity.
show_periodic_connections (bool) – If True, indicates wrap-around connections textually or graphically.
show_primitive_cell (bool) – If True, overlays the primitive unit cell vectors/box.
(e.g. (... other kwargs passed to the underlying plotting function)
show_bz (bool, default=True) – Draw the first Brillouin zone.
show_path (bool, default=True) – Draw the ideal high-symmetry path.
show_matched_kpoints (bool, default=True) – Highlight sampled k-points whose distance to the path is within the
matching tolerance.
points_per_seg (int, default=40) – Number of interpolation points per path segment for the ideal path.
path_match_tol (float, optional) – Distance tolerance used when highlighting mesh points near the
drawn path.
extend (bool, default=False) – Draw translated copies of the sampled k-mesh.
extend_copies (int or iterable[int], optional) – Number of reciprocal-cell copies per direction. In 2D,
extend_copies=1 includes the first shell around the first Brillouin
zone and extend_copies=2 includes the second shell as well.
coordinates: Real-valued array of shape (Ns, dim).
kvectors: Real-valued array of shape (Ns, 3) (or dim).
Neighbor indices are integers in range [0, Ns).
High-symmetry points in the Brillouin zone:
- 1D: Γ (0) -> X (Pi) -> Γ (2Pi)
- 2D: Γ (0,0) -> X (Pi,0) -> M (Pi,Pi) -> Γ (0,0)
- 3D: Γ -> X -> M -> Γ -> R -> X
Return high-symmetry points for the square/cubic lattice.
Returns:
High-symmetry points with default path based on dimension:
- 1D: Γ -> X -> Γ (zone boundary at Pi)
- 2D: Γ -> X -> M -> Γ (standard square BZ path)
- 3D: Γ -> X -> M -> Γ -> R -> X (standard cubic BZ path)
Convert (x, y, z) coordinates to a site index.
:param x: x-coordinate
:type x: int
:param y: y-coordinate
:type y: int
:param z: z-coordinate
:type z: int
Returns nearest neighbors in a given direction (X, Y, Z).
:param site: Site index
:type site: int
:param direction: Direction to get the nearest neighbors
:type direction: LatticeDirection
Simple nearest-neighbour tight-binding/spin-wave-like dispersion for the
square lattice. Accepts k as (2,) or (…,2) array and returns same-shaped
scalar or array of energies.
This module provides the HexagonalLattice class which implements an
armchair-oriented honeycomb lattice. Unlike the zig-zag
HoneycombLattice, the primitive vectors here are chosen so that
the lattice grows vertically and horizontally (aligned with x and y
coordinate axes) with armchair edges. There are no dangling bonds at the
boundary when periodic boundary conditions are used.
Armchair-oriented hexagonal (honeycomb) lattice up to 3 dimensions.
The lattice is constructed so that armchair edges lie along the
horizontal (x) axis, giving a rectangular bounding box aligned
with the coordinate system. Two sites per unit cell (A / B
sublattices).
General Lattice class. This class contains the general lattice model.
Parameters:
dim (int, optional) – Dimension of the lattice (1, 2, or 3). If None, inferred from lx, ly, lz.
lx (int, optional) – Length of the lattice in the x-direction.
ly (int, optional) – Length of the lattice in the y-direction.
lz (int, optional) – Length of the lattice in the z-direction.
bc (str, optional) – Boundary conditions (e.g., ‘PBC’, ‘OBC’).
adj_mat (np.ndarray, optional) – Adjacency matrix for the lattice.
flux (np.ndarray, optional) – Flux piercing the boundaries. This can be a dictionary specifying the
flux in each direction, or a single value applied to all directions. Importantly,
this automatically implies TWISTED boundary conditions, so the bc parameter can be left as None or set to ‘TWISTED’ for clarity.
Real-space position for stored coordinate tuple (x,y,z).
The base class calculate_coordinates already stores proper
vectors via _a1,_a2,_a3,_basis. This helper is kept for
backwards compatibility and any custom coordinate look-ups.
Calculate next-nearest neighbours for the armchair hexagonal lattice.
NNN connect sites within the same sublattice. Each site has 6
NNN in the full 2D honeycomb; finite clusters may have fewer
depending on boundary conditions.
NNN displacements (same sublattice) in cell coordinates:
Simple triangular-lattice dispersion approximation:
ω(k) = 2J * [3 - cos(k·a1) - cos(k·a2) - cos(k·(a1 - a2))]
where a1=(a,0), a2=(a/2, √3 a/2).
Accepts k as (2,) or (…,2).
Contains the Honeycomb lattice implementation.
This module defines the HoneycombLattice class, which extends the base Lattice class
to represent a 2D honeycomb lattice structure. It includes methods for calculating
nearest and next-nearest neighbors, as well as lattice vectors and coordinates.
The honeycomb lattice is a 2D lattice with a hexagonal structure. The lattice consists of
two sublattices (A and B) arranged in a hexagonal pattern. Nearest and next-nearest neighbors
are computed based on a hexagonal unit cell.
High-symmetry points in the Brillouin zone:
- Gamma:
Zone center at (0, 0)
K:
Dirac point at (2/3, 1/3) - hosts linear band crossings in graphene
Return high-symmetry points for the honeycomb lattice.
Returns:
High-symmetry points for the hexagonal Brillouin zone:
- Γ (Gamma): Zone center (0, 0)
- K: Dirac point at (2/3, 1/3) - hosts linear band crossings
- K’: Other Dirac point at (1/3, 2/3)
- M: Edge midpoint at (1/2, 0)
Returns the nearest neighbor in the specified direction.
For the honeycomb lattice, we choose a mapping:
LatticeDirection.X -> neighbor at index 0 of _nn[site]
LatticeDirection.Y -> neighbor at index 1 of _nn[site]
LatticeDirection.Z -> neighbor at index 2 of _nn[site]
Calculates the nearest neighbors (NN) using boundary conditions.
The implementation uses a helper function to apply periodic or open boundary conditions.
For 2D, for example, we use a different treatment on even and odd indices.
Calculates the next-nearest neighbors (NNN) of the honeycomb lattice.
NNN are second-nearest neighbors, connecting sites on the same sublattice.
For sublattice A (even sites), the three NNN directions are obtained by
composing two consecutive NN hops:
NNN_1: Y-bond then X-bond^{-1} → cell shift (0, -1) [down in y]
NNN_2: Z-bond then Y-bond^{-1} → cell shift (-1, 0) [left in x]
NNN_3: X-bond then Z-bond^{-1} → cell shift (+1, +1) [diagonal]
For sublattice B (odd sites), the shifts are inverted.
Honeycomb (graphene-like) nearest-neighbour dispersion magnitude.
Computes |f(k)| where f = sum_{δ} exp(-i k·δ) for the three A->B vectors
used by this lattice implementation.
Most routines accept NumPy array-like inputs and return NumPy arrays or scalars.
APIs typically operate on 1D or 2D arrays where axis conventions are documented
per function. Use floating dtypes for interpolation and filtering paths.
Math utilities for various mathematical operations, fittings, and distributions.
It includes functions for finding maxima, nearest values, and fitting data to various models.
Find the nearest idx to the value given
- x : a DataFrame or numpy array
- val : a scalar
- col : a string on which column to find the nearest
Returns the index of the nearest value to the given value
Class that stores only the parameters of the fit function
- popt : parameters of the fit
- pcov : covariance matrix of the fit
- funct : function of the fit
Skips a certain part of the values for the fit
- x : arguments to trim
- y : values to trim
- skipF : number of first elements to skip
- skipL : number of last elements to skip
Fits [a * x + b * x ** 2]
- x : arguments to the fit
- y : values to fit
- skipF : number of elements to skip on the left
- skipR : number of elements to skip on the right
Fits function [a*x**b]
- x : arguments to the fit
- y : values to the fit
- skipF : number of elements to skip on the left
- skipR : number of elements to skip on the right
Fits function [a*x**b]
- x : arguments to the fit
- y : values to the fit
- skipF : number of elements to skip on the left
- skipR : number of elements to skip on the right
Fits function [any]
- funct : function to fit to
- skipF : number of elements to skip on the left
- skipR : number of elements to skip on the right
staticfitAny(x, y, funct, skipF=0, skipL=0, bounds=[])[source]
Fits function [any]
- x : arguments to fit
- y : values to fit
- funct : function to fit to
- skipF : number of elements to skip on the left
- skipR : number of elements to skip on the right
Generalized Cauchy distribution
- v is the normalization factor
- alpha is the stability parameter, often referred to as the shape parameter,
- beta is the scale parameter,
- gamma is a scale parameter related to the width of the distribution.
Get the next power of a number (base) that is greater than x - x : number to get the next power
- base : base of the power (default 2 for binary, can be 10 for decimal)
Get the previous power of a number (base) that is smaller than x - x : number to get the next power
- base : base of the power (default 2 for binary, can be 10 for decimal)
Compute the modified rounding division of a divided by b.
This function ensures that the result has the same sign as a.
a : integer dividend
b : integer divisor
Statistical analysis tools and data aggregation utilities.
This module contains:
- Binning and averaging routines (bin_avg, rebin).
- Histogram classes with support for merging and weighted averages.
- Statistical moments, fluctuations, and distribution fitting.
- Helpers for spectral function analysis (fractional statistics).
Numba-optimized routines (_bin_avg_numba) are provided for performance on large datasets.
Robust handling of NaN values and empty bins is included in averaging functions.
Compute the bin average of data over multiple realizations.
This method aggregates data points that fall within [center-delta,center+delta]
for each specified center. It supports both arithmetic mean and typical average (geometric mean).
Parameters:
data (np.ndarray) – Input data array.
Shape: (n_realizations,n_points).
x (np.ndarray) – X-coordinates corresponding to the data.
Shape: (n_realizations,n_points), matching data.
centers (np.ndarray) – Array of bin centers to compute averages at.
Shape: (n_centers,).
delta (float, default=0.05) – Half-width of the bin interval. Bins are [c-delta,c+delta].
typical (bool, default=False) – If True, computes the typical average (geometric mean).
Process: exp(mean(log(data))).
The input data is assumed to be positive if this is set.
cutoffNum (int, default=10) – Minimum number of data points required in a bin to consider it valid without expansion.
If the count is lower, the method attempts to expand the window to find nearest neighbors.
func (callable, optional) – Aggregation function to apply to values in a bin.
Signature: func(values:array)->scalar.
Defaults to np.mean.
Note: The Numba optimized path is only used if func is the default.
verbose (bool, default=False) – If True, prints debug information (currently unused).
Returns:
Array of averaged values corresponding to each center.
Shape: (n_valid_centers,).
Note: Invalid centers (where no data could be found even after expansion) are skipped in the output logic of the Numba path, but np.nan_to_num is applied at the end.
This method reshapes the input array and computes the mean over blocks of size av_num.
It randomly shuffles the array before binning to ensure unbiased sampling if the data
is ordered.
Parameters:
arr (np.ndarray) – Input array to rebin.
Shape depends on d.
av_num (int) – Number of elements to average into a single bin.
d (int) – Dimensionality of the array (1, 2, or 3).
If d=1: expects shape (N,).
If d=2: expects shape (N,M).
If d=3: expects shape (N,M,K).
rng (np.random.Generator, optional) – Random number generator for shuffling. If None, a default generator is used.
Returns:
Re-binned array.
The first dimension size is divided by av_num.
staticget_cdf(x, y, gammaval=0.5, BINVAL=21)[source]
Calculate the cumulative distribution function (CDF) and find the gamma value.
Parameters:
x (array-like): The independent variable values.
y (array-like): The dependent variable values, which may contain NaNs.
gammaval (float, optional): The target CDF value to find the corresponding gamma value. Default is 0.5.
Returns:
tuple: A tuple containing:
x (array-like): The input independent variable values.
y (array-like): The input dependent variable values with NaNs removed.
cdf (array-like): The cumulative distribution function values.
gammaf (float): The value of the independent variable corresponding to the target CDF value.
Find the peak value in the given data and interpolate to improve peak precision.
This function removes NaN values from the input arrays, performs spline interpolation
to improve the precision of the peak detection, and then finds the maximum value and
its corresponding alpha. A fine-grained search around the peak is performed to find
a more precise maximum.
Parameters:
- alphas (array-like): The array of alpha values.
- values (array-like): The array of corresponding values.
Returns:
- tuple: A tuple containing the refined alpha value at the peak and the refined peak value.
Initialize the histogram with either a specified number of bins or specific edges.
:param n_bins: Number of bins (if edges is None).
:param edges: Specific bin edges (if n_bins is None).
:param dtype: Data type for the bin edges.
Raises:
ValueError – If both n_bins and edges are None, or if edges is not a one-dimensional array with at least two elements.
Notes
If both n_bins and edges are None, a histogram with one bin (0 to 0) is created.
If edges are provided, the number of bins is determined from the length of edges.
For the specified values, set the histogram counts.
If set_bins is True, the bin edges will be determined from the minimum and maximum of the data.
For complex-valued inputs, only the real part is used.
Create a uniform distribution of bins between v_min and v_max.
Parameters:
- v_max: Maximum value for the histogram.
- v_min: Minimum value for the histogram.
Create a logarithmic distribution of bins between v_min and v_max.
Parameters:
- n_bins: Number of bins.
- v_max: Maximum value for the histogram.
- v_min: Minimum value for the histogram.
- base: Logarithm base (default 10).
Returns:
- bin_edges: Array of bin edges (length n_bins + 1).
Append a value to the histogram by determining its bin and incrementing the corresponding count.
Returns the bin index.
Parameters:
- value: The value to append to the histogram.
Returns:
- bin indices: The indices of the bin where the value was appended.
Additional properties for the histogram class, adding bin averages.
This class allows one to have a function f(x) averaged over the bins.
The binAverages are the sum of the function evaluated in each bin,
and they can be normalized by the bin counts.
Initialize the histogram with either a specified number of bins or specific edges.
:param n_bins: Number of bins (if edges is None).
:param edges: Specific bin edges (if n_bins is None).
:param dtype: Data type for the bin edges.
Raises:
ValueError – If both n_bins and edges are None, or if edges is not a one-dimensional array with at least two elements.
Notes
If both n_bins and edges are None, a histogram with one bin (0 to 0) is created.
If edges are provided, the number of bins is determined from the length of edges.
The bin counts and averages are initialized to zero.
Get the average of the function over the bins normalized by the counts.
If is_typical is True, exponentiate the normalized averages (useful if the averages
represent logarithms).
Add precomputed sums and counts to the histogram averages and counts.
:param - sums: Array of sums to add to the bin averages.
:param - counts: Array of counts to add to the bin counts.
Raises:
- ValueError – If the shapes of sums or counts do not match the histogram’s bin counts.
Remove precomputed sums and counts from the histogram averages and counts.
:param - sums: Array of sums to remove from the bin averages.
:param - counts: Array of counts to remove from the bin counts.
Raises:
- ValueError – If the shapes of sums or counts do not match the histogram’s bin counts.
Calculate the number of states to take based on a fraction of the total size.
Parameters:
fraction : fraction of the total size to take.
size : total size of the Hilbert space.
Returns:
The number of states to take.
Get the specific indices in a range around a given index in the Hilbert space.
Checks for boundaries.
Parameters:
l : number of elements to the left of idx.
r : number of elements to the right of idx.
idx : center index.
size : total size of the Hilbert space.
Returns:
A tuple (min_index, max_index) with the allowed index range.
Parameters:
frac (float) : The fraction of the data to take. If frac is less than 1.0, it is treated as a fraction of the total data size.
If frac is greater than 1.0, it is treated as the number of elements to take.
data (list) : The list of data from which to take the fraction.
around (float, optional) : The index around which to take the fraction. If None, it defaults to half the size of the data.
fraction_left (float, optional) : The fraction of the left side to take. Default is 0.5.
fraction_right (float, optional): The fraction of the right side to take. Default is 0.5.
around_idx (int, optional) : The index around which to take the fraction. If None, it defaults to half the size of the data.
Returns:
list: A list containing the central portion of the original data, based on the specified fraction.
If the calculated number of elements to take is less than or equal to 1, or equal to the size of the data,
the original data is returned.
Get the off-diagonal Hilbert-space fraction information.
Iterates over the energy spectrum (from index mn to max_val) and for each pair (i, j) with j > i,
if the average energy is within tol of target_en then store a tuple of
(energy difference, j, i)
in the output list. Finally, sort the list by the energy difference (first element).
Parameters:
mn : starting index (inclusive).
max_val : ending index (exclusive).
hilbert_size: size of the Hilbert space (not used in computation here, but kept for consistency).
energies : 1D NumPy array of energies.
target_en : target energy for the mean.
tol : tolerance for closeness.
sort : whether to sort the output list by the energy difference.
Returns:
A list of tuples (omega, j, i) sorted by omega if sort is True.
Public factories typically accept model identifiers, shape metadata (for
example input_shape=(n_features,)), and optional dtype and seed arguments.
Returned objects are model instances or callables compatible with the training
helpers in general_python.ml.training_phases.
Input batches are conventionally rank-2 arrays with shape (batch,features)
unless a model documents an image or sequence layout. For stable optimization,
float32 is the practical default on accelerators, while float64 may be
required for high-precision experiments.
Training trajectories depend on initialization, optimizer state, and operation
ordering. For reproducibility, fix random seeds, keep backend/device constant,
and avoid mixing precision policies within one experiment.
This module provides a centralized factory function, choose_network, for
instantiating various neural network architectures used in the general_python framework.
It uses a lazy-loading mechanism to improve startup performance.
Import and use the factory to create a network. The factory takes the
network type and common parameters like input_shape and dtype.
Network-specific parameters are passed as keyword arguments.
from general_python.ml.networks import choose_network
# Create an RBM using ‘alpha’ (hidden unit density)
rbm_net = choose_network(
‘rbm’,
input_shape=(10,),
alpha=2.0, # Creates 2*10=20 hidden units
dtype=’complex64’
)
# Create a CNN for an 8x8 lattice
cnn_net = choose_network(
This factory can create networks by name (e.g., ‘rbm’, ‘cnn’), wrap raw Flax modules,
or instantiate custom GeneralNet subclasses. It handles network-specific arguments
passed via **kwargs and provides conveniences like alpha for RBMs.
‘input_shape’: (100,), # 100 spins
‘alpha’: 2, # Density of hidden units
‘use_bias’: True,
‘dtype’: ‘complex128’ # Essential for quantum phases
}
# 2. Create the Network
# ———————
# ‘rbm’ key triggers the RBM class factory
net = choose_network(‘rbm’, **rbm_params)
# 3. Initialize & Run
# ——————-
# Initialize with a random key (handled internally or explicitly)
# params = net.init(jax.random.PRNGKey(0))
# log_psi = net(params, sample_configuration)
Convolutional Neural Network (CNN) for Lattice Systems.
A deep architecture that respects the locality of physical interactions.
Essential for 2D frustrated systems (like Kitaev or J1-J2 models) where
local correlations are complex.
Features:
- Periodic Boundary Conditions (Torus geometry).
- Sum Pooling: Ensures energy is extensive (scales with N).
- Complex Weights: Captures the sign structure of the wavefunction.
from general_python.ml.networks import choose_network
import jax.numpy as jnp
Callable]]) (- input_activation (Optional[Union[str,) – Activation function applied to the input layer.
Useful for preprocessing inputs (e.g., scaling or encoding).
An initialized or wrapped network instance compatible with the general_python framework.
The returned object is callable with signature:
log_psi=net(params,inputs) where inputs has shape (batch,features)
and output has shape (batch,).
Return type:
GeneralNet
Scheduler implementations for machine learning training.
It includes various learning rate schedulers and an early stopping mechanism.
This function can either accept:
- A string or SchedulerType enum to specify the type of scheduler to create.
- An existing Parameters instance to reconfigure.
:param scheduler_type: Type of scheduler or existing instance.
:param initial_lr: Initial learning rate.
:param max_epochs: Maximum number of epochs.
:param logger: Optional logger for the scheduler.
:param **kwargs: Additional arguments specific to the scheduler type.
lr_decay: Decay rate for exponential/step/adaptive schedulers.
step_size: Step size for step scheduler.
min_lr: Minimum learning rate for cosine/linear/adaptive schedulers.
patience: Patience for adaptive scheduler.
cooldown: Cooldown period for adaptive scheduler.
min_delta: Minimum improvement for adaptive scheduler.
early_stopping_patience: Patience for early stopping.
early_stopping_min_delta: Minimum improvement for early stopping.
Learning phase framework for Neural Quantum State training.
This module implements a multi-phase training system for NQS, allowing:
- Phase transitions with configurable parameters
- Phase-specific callbacks and hooks
- Adaptive learning rates per phase
- Regularization scheduling per phase
- Progress tracking and reporting
Learning phases represent different stages of optimization:
Pre-training: Initialize network with simple loss, high learning rate
Main Optimization: Full Hamiltonian, adaptive learning rate
Refinement: Fine-tune observables, low learning rate, high regularization
Each phase defines learning rate and regularization schedules that are
active for a specific number of epochs. Phases are processed sequentially
by the PhaseScheduler.
Scheduler type for LR. Options:
- ‘constant’: Fixed lr throughout phase
- ‘exponential’: lr * exp(-lr_decay * local_epoch)
- ‘step’: lr * lr_decay^floor(local_epoch/step_size)
- ‘cosine’: Cosine annealing from lr to min_lr
- ‘linear’: Linear decay from lr to min_lr
- ‘adaptive’: ReduceLROnPlateau (requires loss)
The PhaseScheduler orchestrates multi-phase training by:
1. Tracking the current phase based on global epoch count
2. Instantiating appropriate low-level schedulers for each phase
3. Firing callbacks on phase transitions
4. Returning scheduled values via __call__
Parameters:
phases (List[LearningPhase]) – Ordered list of training phases to execute.
param_type (str, default='lr') – Which parameter to schedule (‘lr’ or ‘reg’).
logger (Logger, optional) – Logger for phase transition messages.
>>> fromgeneral_python.ml.training_phasesimportLearningPhase,PhaseScheduler>>>>>> phases=[... LearningPhase(name='warmup',epochs=50,lr=0.1,lr_schedule='exponential',... lr_kwargs={'lr_decay':0.05}),... LearningPhase(name='main',epochs=200,lr=0.02,lr_schedule='constant'),... ]>>>>>> lr_scheduler=PhaseScheduler(phases,param_type='lr')>>> reg_scheduler=PhaseScheduler(phases,param_type='reg')>>>>>> # Use in training loop>>> forepochinrange(250):... lr=lr_scheduler(epoch,loss=current_loss)# Auto phase transition... reg=reg_scheduler(epoch,loss=current_loss)
Factory function to create LR and Reg schedulers from a preset.
Parameters:
preset (str, default='default') – Preset name. Available:
- ‘default’: 3-phase training (350 total epochs)
- ‘kitaev’: Specialized for frustrated systems (550 total epochs)
logger (Logger, optional) – Logger for scheduler messages.
>>> lr_sched,reg_sched=create_phase_schedulers('default')>>>>>> # Pass to NQSTrainer>>> trainer=NQSTrainer(nqs,phases=(lr_sched,reg_sched))>>>>>> # Or use preset string directly>>> trainer=NQSTrainer(nqs,phases='default')# Equivalent
Functions generally accept array-like states, operators, or spectra and return
NumPy or JAX-compatible arrays or scalar observables. Shape requirements follow
physics conventions, e.g. state vectors (d,), operators (d,d), and grids
for momentum or frequency response evaluations.
Entropy and spectral routines include tolerance or regularization knobs to reduce
instability near zero eigenvalues or narrow broadenings. Reproducibility depends
on deterministic eigensolver settings and fixed random seeds in upstream code.
This module contains functions for manipulating and analyzing density matrices in quantum mechanics.
It provides optimized implementations using NumPy and Numba for computing reduced density matrices,
Schmidt decompositions, and entanglement spectra.
QES Convention:
- State vector index i = s0 + d*s1 + d^2*s2 + … (Little-endian / Fortran order)
- Subsystem A site order: [a0, a1, a2, …] -> Row index I = sa0 + d*sa1 + …
- Subsystem B site order: [b0, b1, b2, …] -> Col index J = sb0 + d*sb1 + …
where d is the local_dim.
Fermionic Systems:
For fermionic systems mapped via Jordan-Wigner transformation, non-local string operators
create additional correlations between non-contiguous subsystem sites. The fermionic=True
flag applies sign corrections that account for the fermionic exchange statistics when
reordering sites. This ensures correct reduced density matrices for arbitrary subsystem
geometries.
version : 2.1
copyright : (c) 2026 by Maksymilian Kliczkowski. All rights reserved.
——————————–
Process the subsystem specification to extract site indices and the permutation order.
The order tuple specifies how to permute the state vector to bring subsystem A sites to the front.
Parameters:
va (Union[int, np.ndarray, List[int]]) – Subsystem specification. Can be:
- An integer (if contiguous=True) specifying the number of contiguous sites in A starting from site 0.
- A bitmask integer where bits set to 1 indicate sites in A (if contiguous=False).
- A list or array of site indices in A.
local_dim (int) – Local Hilbert space dimension (default is 2 for qubits).
contiguous (bool) – If True, treat va as the number of contiguous sites in A starting from site 0. If False, treat va as a bitmask or list of site indices.
Reshape and reorder a quantum state vector into a matrix Psi_{A,B} using NumPy.
This representation is used to compute the reduced density matrix rho_A = Psi @ Psi^dagger.
Parameters:
state (np.ndarray) – The input state vector of shape (local_dim**ns,).
order (Tuple[int, ...]) – The permutation order of sites to bring subsystem A sites to the front.
size_a (int) – The number of sites in subsystem A.
local_dim (int) – Local Hilbert space dimension (default is 2 for qubits).
fermionic (bool) – If True, apply fermionic sign corrections for site permutation.
This accounts for the anticommutation of fermionic operators when
reordering sites, essential for correct RDMs of non-contiguous
subsystems in fermionic systems mapped via Jordan-Wigner.
Returns:
Reshaped state matrix Psi of shape (dA, dB) where dA = local_dim^size_a
and dB = local_dim^(ns - size_a).
Return type:
np.ndarray
Notes
For fermionic systems (fermionic=True):
When we permute the site ordering, fermionic operators anticommute:
c_i c_j = -c_j c_i
For a basis state |n_0, n_1, …, n_{ns-1}> represented as:
If True, apply fermionic sign corrections for non-contiguous subsystems.
Essential for correct entanglement entropy of Jordan-Wigner mapped fermions.
For fermionic systems, permuting sites requires accounting for the
anticommutation of creation operators. Each pair of occupied sites
that gets inverted in the permutation contributes a factor of -1.
Use fermionic=True when:
- Computing RDM of non-contiguous subsystem in a fermionic system
- Working with Slater determinants or their superpositions
- Comparing with correlation matrix entropy results
For contiguous subsystems (sites 0,1,…,k-1), the fermionic flag
has no effect since no site permutation is needed.
la (Optional[int]) – Deprecated alias for specifying contiguous subsystem size.
order (Optional[Tuple[int, ...]]) – Explicit site permutation order. If provided, overrides va.
Returns:
Reduced density matrix rho_A of shape (dA, dA) where dA = local_dim^|A|.
The diagonal cutoff for the spectral statistics
- Ns: system size
- Nh: Hilbert space dimension
- nu: number of eigenvalues - if it’s less than 1, then it’s a fraction of the Hilbert space dimension
- minimal_frac: minimal fraction of the Hilbert space dimension
Returns:
- The diagonal cutoff for the spectral statistics
Take a fraction of the data for the spectral statistics.
- nu: number of eigenvalues - if it’s less than 1, then it’s a fraction of the Hilbert space dimension
- data: data to take the fraction from
- middle: middle of the data (if None, then it’s the middle of the data)
Returns:
- The fraction of the data
Take a fraction of the data for the spectral statistics.
- nu: number of eigenvalues - if it’s less than 1, then it’s a fraction of the Hilbert space dimension
- data: data to take the fraction from
- middle: middle of the data (if None, then it’s the middle of the data)
Returns:
- The fraction of the data
Take a fraction of the data for the spectral statistics. Calculates the mean of the data.
- nu: number of eigenvalues - if it’s less than 1, then it’s a fraction of the Hilbert space dimension
- data: data to take the fraction from
- middle: middle of the data (if None, then it’s the middle of the data)
- axis: axis to calculate the mean
Returns:
- The mean of the fraction of the data
Parses and resolves a site identifier string into a numeric index or value.
This method handles special symbolic representations used in operator strings,
such as ‘L’ (system size), ‘pi’, and arithmetic operations like division or subtraction.
Parameters:
site (str) – String representation of the site (e.g., “0”, “L”, “L_2”, “L-1”).
_dimension (int, default=1) – The dimension or size of the system (L), used to resolve ‘L’ or ‘l’.
Returns:
The resolved numeric value of the site or parameter.
>>> Operators.resolveSite("0",10)0>>> Operators.resolveSite("L",10)# If OPERATOR_SITE_M_1 is True (default)9>>> Operators.resolveSite("L_2",10)# L / 25
Resolves a full operator string by parsing the site component.
For an operator string like “Sz/L_2”, this splits the string by the operator separator (‘/’),
resolves the site part (“L_2” -> dimension/2), and reconstructs the string.
Convert an encoded operator name into a LaTeX-style title.
general_python/physics/statistical.py
Statistical analysis utilities for quantum systems.
This module provides tools for:
- Finite window averages and time series analysis
- Local density of states (LDOS) and strength functions
- Spectral histograms and binning
- Windowed matrix element calculations
- Generic histogram and scatter analysis
Entropy calculations handle small eigenvalues by clipping or conditional checks to avoid log(0).
Schmidt decomposition is preferred for reduced density matrices when possible for stability and efficiency.
Calculate the gaussianity <|Oab|^2>/<|Oab|>^2 -> for normal == pi/2
- arr : array to calculate the gaussianity
- axis : axis to calculate the gaussianity
For quadratic (non-interacting) Hamiltonians, the entanglement entropy can be computed
efficiently from the single-particle correlation matrix C_ij = <c_i^dag c_j>:
S = -Tr[C log C + (1-C) log(1-C)]
This scales as O(L^3) compared to O(2^L) for exact diagonalization.
For many-body states, we use Schmidt decomposition of the wavefunction:
>>> masks=MaskGenerator.contiguous(ns=12,size_a=4)# First 4 sites>>> masks=MaskGenerator.alternating(ns=12)# Even/odd sites>>> masks=MaskGenerator.random(ns=12,size_a=6)# Random 6 sites>>> masks=MaskGenerator.kitaev_preskill(ns=12)# ABC regions for TEE
Utility class for generating subsystem masks for entanglement calculations.
Provides convenient methods to create site masks for various bipartition
geometries, including contiguous, alternating, random, and topological
(Kitaev-Preskill) constructions.
>>> results=ent.entropy_multipartition(... [[0,1],[0,1,2],[0,1,2,3]],... orbitals... )# computes entropies for 3 bipartitions>>> entropies=results['entropies']# array of 3 entropies for each bipartition>>> C_matrices=results['correlation_matrices']# list of 3 matrices
>>> results=ent.entropy_scan(orbitals,sizes=[1,2,3,4,5])# entropies for subsystems of sizes 1 to 5, consqutive sites starting from 0
Many-body Hamiltonian (interacting):
Manual bipartition entropy:
>>> hamil=ManyBodyHamiltonian(ns=8,...)>>> hamil.diagonalize()>>> ent=hamil.entanglement>>> bipart=ent.bipartition([0,1,2,3])# subsystem A>>> state=hamil.eig_vec[:,0]# ground state wavefunction>>> S_manual=ent.entropy_manybody(bipart,state)# entropy from reduced density matrix, it happens internally
Density matrix access:
>>> rho_A=ent.reduced_density_matrix(bipart,state)# reduced density matrix for subsystem A>>> rho_B=ent.reduced_density_matrix(bipart,state,subsystem='B')# for subsystem B
Get single-particle correlation matrix C_ij = <c_i^\dag c_j>.
Computes the correlation matrix for a free fermion state defined by
occupied orbitals. Uses spin-unpolarized convention (factor of 2).
Parameters:
occupied_orbitals (array-like) – Indices of occupied orbitals (in eigenstate basis).
For ground state, use [0, 1, …, N-1] for N particles.
bipartition (BipartitionInfo, optional) – If provided, returns correlation matrix restricted to subsystem A.
If None, returns full correlation matrix for all sites.
subtract_identity (bool) – Whether to subtract identity from the correlation matrix.
backend (str) – ‘numpy’ or ‘jax’ for GPU acceleration.
Returns:
Correlation matrix C_ij = <c_i^\dag c_j>.
Shape: (size_a, size_a) if bipartition given, else (ns, ns).
Diagonal elements are site occupations (in [0,2] range with factor 2).
Calculate entanglement entropy from single-particle correlation matrix.
SINGLE-PARTICLE METHOD - Fast O(L_A³) method for non-interacting (quadratic)
Hamiltonians. Computes entropy from correlation matrix eigenvalues.
Works for ANY bipartition (contiguous or non-contiguous) of free fermion states.
Parameters:
bipartition (BipartitionInfo) – Bipartition of the system (use ent.bipartition() to create).
Works for both contiguous and non-contiguous subsystems.
occupied_orbitals (array-like) – Indices of occupied orbitals (in eigenstate basis).
For ground state with N particles, use [0, 1, …, N-1].
C_A (np.ndarray or jax.numpy.ndarray, optional) – Precomputed correlation matrix for subsystem A.
If provided, uses this instead of computing from occupied_orbitals.
subtract_identity (bool) – Whether to subtract identity from correlation matrix (advanced)
backend (str) – ‘numpy’ or ‘jax’ for computation backend
Returns:
Entanglement entropy (in natural log units, always positive)
Algorithm:
1. Compute full correlation matrix C_ij = <c_i^dag c_j> from occupied orbitals. For BdG, use all <c_i c_j> etc.
2. Extract subblock C_A for sites in subsystem A (handles non-contiguous)
3. Diagonalize C_A to get eigenvalues (occupations in [0,1])
4. Apply single-particle entropy formula:
S = - sum_k [ n_k log(n_k) + (1-n_k) log(1-n_k) ]
This gives the EXACT entanglement entropy for ANY bipartition of
non-interacting (quadratic) Hamiltonians and matches entropy_many_body().
Limitations:
- Requires diagonalized Hamiltonian
- Only works for quadratic (non-interacting) Hamiltonians
- For interacting systems, use entropy_many_body()
Calculate entanglement entropy from many-body state.
MANY-BODY METHOD - Exact method that works for ANY quantum state,
including interacting systems. Performs Schmidt decomposition of the
many-body wavefunction.
Parameters:
bipartition (BipartitionInfo) – Bipartition of the system (use ent.bipartition() to create)
q (float) – Renyi index (default: 1.0 for von Neumann entropy)
rho_a (np.ndarray, optional) – Precomputed reduced density matrix for subsystem A.
If provided, uses this instead of computing from state.
state (np.ndarray, optional) – Many-body state vector (length 2^ns).
If None, occupied_orbitals must be provided to construct the state (for free fermions).
method (str) – ‘auto’ : Choose best method based on bipartition geometry
‘schmidt’ : Use Schmidt decomposition with mask (for non-contiguous)
‘numpy’ : Use direct numpy Schmidt (for contiguous, faster)
use_eig (bool) – Whether to use eigenvalue decomposition (True) or SVD (False)
hilbert (HilbertSpace, optional) – Hilbert space with symmetries. If provided and has symmetries,
symmetry-based reduced density matrix computation is used.
This is only available when general_python is installed correctly, as it is not
a part of General Python.
occupied_orbitals (array-like, optional) – Indices of occupied orbitals. Required only if state is None.
Returns:
Von Neumann entanglement entropy (always positive)
Calculate entanglement entropy for multiple subsystem sizes.
Parameters:
occupied_orbitals (array-like, optional) – Occupied orbitals for the state (for free fermions).
subsystem_sizes (list of int, optional) – Sizes of subsystem A to scan. If None, scans all sizes from 1 to ns-1
q (float) – Renyi index (default: 1.0 for von Neumann entropy)
method (str) – ‘auto’ : Use correlation matrix for quadratic, many-body otherwise
‘correlation’ : Force correlation matrix method
‘many_body’ : Force many-body method
contiguous (bool) – If True, use contiguous partitions [0:size_a]
If False, use random partitions
state (np.ndarray, optional) – Many-body state vector. Required if occupied_orbitals is None and method is ‘many_body’ or ‘auto’ (for interacting systems).
Returns:
Dictionary with keys:
- ‘sizes’ : Subsystem sizes
- ‘entropies’ : Entanglement entropies
- ‘method’ : Method used
Calculate mutual information I(A:B) = S(A) + S(B) - S(AB). Importantly, regions A and B can be
smaller than total system size, for instance, a common use case is to compute mutual information
for single site subsystems, for instance I(i:j) between sites i and j.
In such case A = {i}, B = {j}, AB = {i,j}.
Parameters:
mask_a (array-like) – Indices for subsystems A and B
mask_b (array-like) – Indices for subsystems A and B
Calculate entanglement entropy for multiple bipartitions simultaneously.
Efficient batch calculation that computes correlation matrix once and
reuses it for all bipartitions, or computes many-body state once for
all Schmidt decompositions.
Parameters:
bipartitions (list) – List of BipartitionInfo objects or site masks (will create BipartitionInfo).
occupied_orbitals (array-like, optional) – Occupied orbitals for correlation method.
method (str) – ‘auto’, ‘correlation’, or ‘many_body’.
backend (str) – ‘numpy’ or ‘jax’ (for correlation method).
state (np.ndarray, optional) – Pre-computed many-body state (for many_body method).
If None and method is many_body, will be computed from occupied_orbitals.
Returns:
Results dictionary containing:
- ‘entropies’: array of entropies for each bipartition
- ‘bipartitions’: list of BipartitionInfo objects
- ‘method’: method used (‘correlation’ or ‘many_body’)
- ‘correlation_matrices’: list of C_A matrices (if method=’correlation’)
This combination cancels the area law contribution and extracts
the universal topological term. For topological phases like the
toric code, γ = log(D) where D is the total quantum dimension.
A state satisfies Wick’s theorem if and only if it is a Slater determinant
(or a mixture thereof for finite temperature). This is equivalent to
the state being a Gaussian state for fermions.