# 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