Source code for scqubits.core.discretization

# discretization.py
#
# This file is part of scqubits: a Python package for superconducting qubits,
# Quantum 5, 583 (2021). https://quantum-journal.org/papers/q-2021-11-17-583/
#
#    Copyright (c) 2019 and later, Jens Koch and Peter Groszkowski
#    All rights reserved.
#
#    This source code is licensed under the BSD-style license found in the
#    LICENSE file in the root directory of this source tree.
############################################################################

from typing import Any, Dict, List, Tuple, Union

import numpy as np

from numpy import ndarray
from scipy import sparse
from scipy.sparse import csc_matrix

import scqubits.core.central_dispatch as dispatch
import scqubits.core.descriptors as descriptors
import scqubits.io_utils.fileio_serializers as serializers
import scqubits.settings as settings
import scqubits.utils.misc as utils


FIRST_STENCIL_COEFFS: Dict[int, List[float]] = {
    3: [-1 / 2, 0.0, 1 / 2],
    5: [1 / 12, -2 / 3, 0.0, 2 / 3, -1 / 12],
    7: [-1 / 60, 3 / 20, -3 / 4, 0.0, 3 / 4, -3 / 20, 1 / 60],
    9: [1 / 280, -4 / 105, 1 / 5, -4 / 5, 0.0, 4 / 5, -1 / 5, 4 / 105, -1 / 280],
}

SECOND_STENCIL_COEFFS: Dict[int, List[float]] = {
    3: [1, -2, 1],
    5: [-1 / 12, 4 / 3, -5 / 2, 4 / 3, -1 / 12],
    7: [1 / 90, -3 / 20, 3 / 2, -49 / 18, 3 / 2, -3 / 20, 1 / 90],
    9: [-1 / 560, 8 / 315, -1 / 5, 8 / 5, -205 / 72, 8 / 5, -1 / 5, 8 / 315, -1 / 560],
}


[docs]def band_matrix( band_coeffs: Union[List[float], List[complex], ndarray], band_offsets: Union[List[int], ndarray], dim: int, dtype: Any = None, has_corners: bool = False, ) -> csc_matrix: """ Returns a dim x dim sparse matrix with constant diagonals of values `band_coeffs[ 0]`, `band_coeffs[1]`, ... along the (off-)diagonals specified by the offsets `band_offsets[0]`, `band_offsets[1]`, ... The `has_corners` option allows generation of band matrices with corner elements, in which lower off-diagonals wrap into the top right corner and upper off-diagonals wrap into the bottom left corner. Parameters ---------- band_coeffs: each element of band_coeffs is a number to be assigned as a constant to the (off-)diagonals band_offsets: offsets specifying the positions of the (off-)diagonals dim: dimension of the matrix dim: (linear) dimension of the matrix dtype: if not specified, dtype is inferred from the dtype of `band_vecs` has_corners: if set to True, the off diagonals are wrapped into the opposing corners of the matrix """ ones_vector = np.ones(dim) vectors = [ones_vector * number for number in band_coeffs] matrix = sparse.dia_matrix((vectors, band_offsets), shape=(dim, dim), dtype=dtype) if not has_corners: return matrix.tocsc() for index, offset in enumerate(band_offsets): if offset < 0: corner_offset = dim + offset corner_band = vectors[index] corner_band = corner_band[offset:] elif offset > 0: corner_offset = -dim + offset corner_band = vectors[index][:-offset] corner_band = corner_band[-offset:] else: # when offset == 0 continue matrix.setdiag(corner_band, k=corner_offset) return matrix.tocsc()
[docs]class Grid1d(dispatch.DispatchClient, serializers.Serializable): """Data structure and methods for setting up discretized 1d coordinate grid, generating corresponding derivative matrices. Parameters ---------- min_val: minimum value of the discretized variable max_val: maximum value of the discretized variable pt_count: number of grid points """ min_val = descriptors.WatchedProperty(float, "GRID_UPDATE") max_val = descriptors.WatchedProperty(float, "GRID_UPDATE") pt_count = descriptors.WatchedProperty(int, "GRID_UPDATE") def __init__(self, min_val: float, max_val: float, pt_count: int) -> None: self.min_val = min_val self.max_val = max_val self.pt_count = pt_count def __repr__(self) -> str: init_dict = self.get_initdata() return type(self).__name__ + f"({init_dict!r})" def __str__(self) -> str: output = "Grid1d -----[ " for param_name, param_val in sorted( utils.drop_private_keys(self.__dict__).items() ): output += f"{param_name}: {param_val}, " output = output[:-3] + " ]" return output def __eq__(self, other: Any) -> bool: if not isinstance(other, type(self)): return False return self.__dict__ == other.__dict__ def __hash__(self): return super().__hash__()
[docs] def get_initdata(self) -> Dict[str, Any]: """Returns dict appropriate for creating/initializing a new Grid1d object. Returns ------- dict """ return self.__dict__
[docs] def grid_spacing(self) -> float: """ Returns ------- spacing between neighboring grid points """ return (self.max_val - self.min_val) / (self.pt_count - 1)
[docs] def make_linspace(self) -> ndarray: """Returns a numpy array of the grid points Returns ------- ndarray """ return np.linspace(self.min_val, self.max_val, self.pt_count)
[docs] def first_derivative_matrix( self, prefactor: Union[float, complex] = 1.0, periodic: bool = False ) -> csc_matrix: """Generate sparse matrix for first derivative of the form :math:`\\partial_{x_i}`. Uses STENCIL setting to construct the matrix with a multi-point stencil. Parameters ---------- prefactor: prefactor of the derivative matrix (default value: 1.0) periodic: set to True if variable is a periodic variable Returns ------- sparse matrix in `dia` format """ if isinstance(prefactor, complex): dtp = np.complex_ else: dtp = np.float_ delta_x = self.grid_spacing() matrix_diagonals = [ coefficient * prefactor / delta_x for coefficient in FIRST_STENCIL_COEFFS[settings.STENCIL] ] offset = [i - (settings.STENCIL - 1) // 2 for i in range(settings.STENCIL)] derivative_matrix = band_matrix( matrix_diagonals, offset, self.pt_count, dtype=dtp, has_corners=periodic ) return derivative_matrix.tocsc()
[docs] def second_derivative_matrix( self, prefactor: Union[float, complex] = 1.0, periodic: bool = False ) -> csc_matrix: """Generate sparse matrix for second derivative of the form :math:`\\partial^2_{x_i}`. Uses STENCIL setting to construct the matrix with a multi-point stencil. Parameters ---------- prefactor: optional prefactor of the derivative matrix (default value = 1.0) periodic: set to True if variable is a periodic variable (default value = False) Returns ------- sparse matrix in `dia` format """ if isinstance(prefactor, complex): dtp = np.complex_ else: dtp = np.float_ delta_x = self.grid_spacing() matrix_diagonals = [ coefficient * prefactor / delta_x**2 for coefficient in SECOND_STENCIL_COEFFS[settings.STENCIL] ] offset = [i - (settings.STENCIL - 1) // 2 for i in range(settings.STENCIL)] derivative_matrix = band_matrix( matrix_diagonals, offset, self.pt_count, dtype=dtp, has_corners=periodic ) return derivative_matrix.tocsc()
[docs]class GridSpec(dispatch.DispatchClient, serializers.Serializable): """Class for specifying a general discretized coordinate grid (arbitrary dimensions). Parameters ---------- minmaxpts_array: array of with entries [minvalue, maxvalue, number of points] """ min_vals = descriptors.WatchedProperty(ndarray, "GRID_UPDATE") max_vals = descriptors.WatchedProperty(ndarray, "GRID_UPDATE") var_count = descriptors.WatchedProperty(int, "GRID_UPDATE") pt_counts = descriptors.WatchedProperty(ndarray, "GRID_UPDATE") def __init__(self, minmaxpts_array: ndarray) -> None: self.min_vals = minmaxpts_array[:, 0] self.max_vals = minmaxpts_array[:, 1] self.var_count = len(self.min_vals) self.pt_counts = minmaxpts_array[:, 2].astype(int) # used as int indices def __str__(self) -> str: output = " GridSpec ......" for param_name, param_val in sorted(self.__dict__.items()): output += f"\n{param_name}\t: {param_val}" return output
[docs] def unwrap(self) -> Tuple[ndarray, ndarray, Union[List[int], ndarray], int]: """Auxiliary routine that yields a tuple of the parameters specifying the grid.""" return self.min_vals, self.max_vals, self.pt_counts, self.var_count