Coverage for backend/ahuora-builder/src/ahuora_builder/unit_model_manager.py: 92%

1from .methods.load_unit_model import add_unit_model_to_flowsheet 

2from .methods.adapter_library import ModelConstructor, AdapterLibrary 

3from .flowsheet_manager_type import FlowsheetManager 

4from ahuora_builder_types import UnitModelId, UnitModelSchema 

5from idaes.core import UnitModelBlock 

6 

7 

8 

9class UnitModelManager: 

10 """ 

11 Manages the unit models in the flowsheet 

12 """ 

13 

14 def __init__(self, flowsheet_manager: FlowsheetManager) -> None: 

15 """ 

16 Create a new unit model manager 

17 """ 

18 self.flowsheet_manager = flowsheet_manager 

19 self._unit_models: dict[UnitModelId, UnitModelBlock] = {} 

20 self._has_initial_values_by_object_id: dict[int, bool] = {} 

21 

22 def load(self) -> None: 

23 """ 

24 load all unit models from the schema into the flowsheet 

25 """ 

26 schema = self.flowsheet_manager.schema 

27 

28 for unit_model_def in schema.unit_models: 

29 unit_model_type = unit_model_def.type 

30 if unit_model_type not in AdapterLibrary: 30 ↛ 31line 30 didn't jump to line 31 because the condition on line 30 was never true

31 raise Exception( 

32 f"Unit model type '{unit_model_type}' is not in the adapter library." 

33 ) 

34 model_constructor: ModelConstructor = AdapterLibrary[unit_model_type] 

35 self.load_from_def(unit_model_def, model_constructor) 

36 

37 # run post_load methods (e.g to build surrogate models) after all models have been loaded 

38 for unit_model_def in schema.unit_models: 

39 unit_model = self._unit_models[unit_model_def.id] 

40 if hasattr(unit_model, "post_load"): 

41 unit_model.post_load() 

42 

43 def load_from_def( 

44 self, unit_model_def: UnitModelSchema, model_constructor: ModelConstructor 

45 ) -> None: 

46 """ 

47 Deserialise a unit model from a JSON object into a unit model adapter 

48 """ 

49 # add the unit model to the flowsheet 

50 unit_model = add_unit_model_to_flowsheet( 

51 unit_model_def, model_constructor, self.flowsheet_manager 

52 ) 

53 

54 # store a reference to the adapter 

55 self._unit_models[unit_model_def.id] = unit_model 

56 has_initial_values = bool(unit_model_def.initial_values) 

57 self._has_initial_values_by_object_id[id(unit_model)] = has_initial_values 

58 

59 def has_initial_values(self, unit_model: UnitModelBlock) -> bool: 

60 return self._has_initial_values_by_object_id.get(id(unit_model), False) 

61 

62 def count_units_with_initial_values(self) -> int: 

63 return sum(1 for has_values in self._has_initial_values_by_object_id.values() if has_values)