Coverage for python / aubellhop / models.py: 94%

1"""The 'Model Registry' for aubellhop to allow multiple BellhopSimulators to be run. 

2 

3This model defines the class and then initialises it. 

4As this is a utility class, the initialised class operates 

5as a global container of methods and the `_models` register 

6of `BellhopSimulator` models. 

7 

8See file `bellhop.py` for the class definition of `BellhopSimulator`. 

9""" 

10 

11 

12from __future__ import annotations 

13 

14from typing import Any 

15 

16from .constants import ModelDefaults 

17from .environment import Environment 

18from .bellhop import BellhopSimulator 

19 

20class Models: 

21 """Registry for BellhopSimulator models. 

22 

23 This is a *Utility Class* which consists of only class methods and a global registry 

24 of defined models. 

25 """ 

26 

27 _models: list[BellhopSimulator] = [] # class-level storage for all models 

28 

29 @classmethod 

30 def init(cls) -> None: 

31 """Create default 2D and 3D models.""" 

32 cls.new(name=ModelDefaults.name_2d, exe=ModelDefaults.exe_2d, dim=ModelDefaults.dim_2d) 

33 cls.new(name=ModelDefaults.name_3d, exe=ModelDefaults.exe_3d, dim=ModelDefaults.dim_3d) 

34 

35 @classmethod 

36 def reset(cls) -> None: 

37 """Clear all models from the registry.""" 

38 cls._models.clear() 

39 

40 @classmethod 

41 def new(cls, name: str, **kwargs: Any) -> BellhopSimulator: 

42 """Instantiate a new Bellhop model and add it to the registry. 

43 

44 Parameters 

45 ---------- 

46 name : str 

47 The (unique) name of the BellhopSimulator model 

48 kwargs : Any 

49 Arguments to pass onto the BellhopSimulator constructor 

50 

51 Returns 

52 ------- 

53 BellhopSimulator 

54 The defined model which was just added to the registry. 

55 """ 

56 for m in cls._models: 

57 if name == m.name: 

58 raise ValueError(f"Bellhop model with this name ('{name}') already exists.") 

59 model = BellhopSimulator(name=name, **kwargs) 

60 cls._models.append(model) 

61 return model 

62 

63 @classmethod 

64 def supported(cls, env: Environment | None = None, task: str | None = None, dim: int | None = None) -> list[str]: 

65 """List available models by name, maybe narrowed by env, task, and/or dimension.""" 

66 if env is not None: 

67 env.check() 

68 rv: list[str] = [] 

69 for m in cls._models: 

70 if m.supports(env=env, task=task, dim=dim): 

71 rv.append(m.name) 

72 return rv 

73 

74 @classmethod 

75 def get(cls, name: str) -> BellhopSimulator: 

76 """Get a model by name. 

77 

78 Parameters 

79 ---------- 

80 name : str 

81 The name of the BellhopSimulator model 

82 

83 Returns 

84 ------- 

85 BellhopSimulator 

86 The first model in the registry which matches the specified name. 

87 """ 

88 for m in cls._models: 

89 if m.name == name: 

90 return m 

91 raise KeyError(f"Unknown model: '{name}'") 

92 

93 @classmethod 

94 def select( cls, 

95 env: Environment, 

96 task: str, 

97 model: str | None = None, 

98 debug: bool = False, 

99 ) -> BellhopSimulator: 

100 """Finds a model to use, or if a model is requested validate it. 

101 

102 Parameters 

103 ---------- 

104 env : dict 

105 The environment dictionary 

106 task : str 

107 The task to be computed 

108 model : str, optional 

109 Specified model to use 

110 debug : bool, default=False 

111 Whether to print diagnostics 

112 

113 Returns 

114 ------- 

115 BellhopSimulator 

116 The first model in the list which satisfies the input parameters. 

117 

118 """ 

119 if model is not None: 

120 return cls.get(model) 

121 models = cls.supported(env=env, task=task, dim=env._dimension) 

122 if debug: 

123 print(f'Models found: {models}') 

124 if len(models) > 0: 

125 return cls.get(models[0]) 

126 raise ValueError('No suitable propagation model available') 

127 

128 def __new__(cls, *args: Any, **kwargs: Any) -> "Models": 

129 raise TypeError("This utility class cannot be instantiated") 

130 

131Models.init()