Coverage for physped/core/lattice.py: 100%

1"""Module for the Latice class.""" 

2 

3import logging 

4from pprint import pformat 

5from typing import Dict, Tuple 

6 

7import numpy as np 

8from omegaconf import OmegaConf 

9 

10from physped.utils.functions import get_bin_middle 

11 

12log = logging.getLogger(__name__) 

13 

14 

15class Lattice: 

16 def __init__(self, bins: Dict[str, np.ndarray]): 

17 """A class for the lattice. 

18 

19 Args: 

20 bins: A dictionary containing the bin edges for each dimension. 

21 """ 

22 self.bins = bins 

23 self.dimensions = tuple(bins.keys()) 

24 self.bin_centers = self.get_bin_centers() 

25 self.shape = self.get_lattice_shape() 

26 # self.cell_volume = self.compute_cell_volume() 

27 

28 def __repr__(self): 

29 formatted_output = pformat( 

30 OmegaConf.to_container(self.bins, resolve=True), depth=1 

31 ) 

32 return f"Lattice(bins={formatted_output})" 

33 

34 def get_bin_centers(self) -> Dict[str, np.ndarray]: 

35 """Return the middle of the input bins. 

36 

37 Returns: 

38 The middle of the input bins. 

39 """ 

40 return {key: get_bin_middle(self.bins[key]) for key in self.bins} 

41 

42 def get_lattice_shape(self) -> Tuple[int]: 

43 """Return the shape of the lattice. 

44 

45 Returns: 

46 The shape of the lattice. 

47 """ 

48 return tuple(len(self.bin_centers[key]) for key in self.bin_centers) 

49 

50 def compute_cell_volume(self) -> np.ndarray: 

51 """Compute the volume of each cell in the lattice. 

52 

53 Returns: 

54 The volume of each cell in the lattice. 

55 """ 

56 dx = np.diff(self.bins["x"]) 

57 dy = np.diff(self.bins["y"]) 

58 dr = np.diff(self.bins["r"]) 

59 r = self.bin_centers["r"] 

60 dtheta = np.diff(self.bins["theta"]) 

61 dk = np.diff(self.bins["k"]) 

62 

63 i, j, k, l, m = np.meshgrid( 

64 np.arange(len(self.bins["x"]) - 1), 

65 np.arange(len(self.bins["y"]) - 1), 

66 np.arange(len(self.bins["r"]) - 1), 

67 np.arange(len(self.bins["theta"]) - 1), 

68 np.arange(len(self.bins["k"]) - 1), 

69 indexing="ij", 

70 ) 

71 

72 # return the volume for each cell using broadcasting 

73 return dx[i] * dy[j] * dr[k] * r[k] * dtheta[l] * dk[m]