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 Parameters 

25 ---------- 

26 pos : Real ndarray, shape (n, space_dim) 

27 Positions. 

28 tags : String ndarray, shape (n,) 

29 Sublattice tags. 

30 index : Integer ndarray, shape (n,) 

31 Site indices. 

32 ''' 

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

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

35 for k, name in enumerate(names): 

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

37 rec['tag'] = tags 

38 rec['index'] = index 

39 return rec 

40 

41 

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

43 ''' 

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

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

46 ''' 

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

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

49 return dict(params) 

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

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

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

53 

54 

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

56 ''' 

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

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

59 

60 Returns 

61 ------- 

62 values 

63 Complex ndarray, shape (size,), or (size, 2, 2) if 

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

65 ''' 

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

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

68 if not spin: 

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

70 if val.ndim <= 1: 

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

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

73 

74 

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

76 r''' 

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

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

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

80 ''' 

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

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

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

84 return value