Coverage for tbkit/values.py: 100%

1''' 

2Value functions: hoppings and onsite energies given as callables of the 

3sites and of named parameters, ``t(site_i, site_j, **params)`` and 

4``onsite(site, **params)`` (see *System.set_hopping*, *System.set_onsite* 

5and *KSpace.set_hopping*). The model stores which bonds or sites each one 

6applies to; ``get_ham(**params)`` calls it once per Hamiltonian, with 

7arrays of sites (one entry per bond or site). 

8''' 

9from __future__ import annotations 

10 

11import inspect 

12from typing import Callable 

13 

14import numpy as np 

15from numpy.typing import NDArray 

16import tbkit.error_handling as error_handling 

17 

18 

19def sites(pos: NDArray[np.float64], tags: NDArray, index: NDArray) -> NDArray: 

20 ''' 

21 Private function. The sites passed to a value function: a structured 

22 array with fields 'x', 'y' (and 'z' in 3D space), 'tag' and 'index'. 

23 

24 :param pos: Real ndarray, shape (n, space_dim). Positions. 

25 :param tags: String ndarray, shape (n,). Sublattice tags. 

26 :param index: Integer ndarray, shape (n,). Site indices. 

27 ''' 

28 names = ('x', 'y', 'z')[:pos.shape[1]] 

29 rec = np.zeros(len(pos), dtype=[(name, 'f8') for name in names] + [('tag', 'U1'), ('index', 'i8')]) 

30 for k, name in enumerate(names): 

31 rec[name] = pos[:, k] 

32 rec['tag'] = tags 

33 rec['index'] = index 

34 return rec 

35 

36 

37def _kwargs(func: Callable, n_args: int, params: dict) -> dict: 

38 ''' 

39 Private function. The entries of *params* that *func* takes (all of them 

40 if it has ``**kwargs``), past its first *n_args* positional arguments. 

41 ''' 

42 args = list(inspect.signature(func).parameters.values()) 

43 if any(p.kind is p.VAR_KEYWORD for p in args): 

44 return dict(params) 

45 named = [p for p in args[n_args:] if p.kind in (p.POSITIONAL_OR_KEYWORD, p.KEYWORD_ONLY)] 

46 error_handling.value_params([p.name for p in named if p.default is p.empty], params, func) 

47 return {p.name: params[p.name] for p in named if p.name in params} 

48 

49 

50def evaluate(func: Callable, args: tuple, params: dict, size: int, spin: bool = False) -> NDArray[np.complex128]: 

51 ''' 

52 Private function. Call the value function *func* on the site arrays 

53 *args*, with the entries of *params* it takes. 

54 

55 :returns: 

56 * **values** -- Complex ndarray, shape (size,), or (size, 2, 2) if 

57 *spin* (a number per bond then multiplies the identity). 

58 ''' 

59 val = np.asarray(func(*args, **_kwargs(func, len(args), params)), dtype='c16') 

60 error_handling.value_shape(val, size, spin, func) 

61 if not spin: 

62 return np.broadcast_to(val, (size,)) 

63 if val.ndim <= 1: 

64 return np.broadcast_to(val, (size,))[:, None, None] * np.eye(2, dtype='c16') 

65 return np.broadcast_to(val, (size, 2, 2)) 

66 

67 

68def reversed_conj(func: Callable) -> Callable: 

69 r''' 

70 Private function. The value function of the reversed bonds: given 

71 :math:`t(i, j) = H_{ij}`, the function :math:`(i, j) \mapsto H_{ji}^\dagger 

72 = t(j, i)^\dagger` (conjugate transpose for 2x2 spin blocks). 

73 ''' 

74 def value(site_i, site_j, **params): 

75 val = np.asarray(func(site_j, site_i, **_kwargs(func, 2, params)), dtype='c16') 

76 return val.conj() if val.ndim < 2 else np.swapaxes(val.conj(), -1, -2) 

77 return value