Coverage for backend/django/core/auxiliary/models/MLWizard.py: 96%

42 statements  

« prev     ^ index     » next       coverage.py v7.10.7, created at 2026-06-23 21:51 +0000

1from core.auxiliary.models.MLModel import MLModel 

2from core.auxiliary.models.MLColumnMapping import MLColumnMapping, custom_property_port_index 

3from core.auxiliary.enums.generalEnums import PropertyType 

4from core.auxiliary.models.PropertyInfo import PropertyInfo 

5from core.auxiliary.models.PropertyValue import PropertyValue 

6from flowsheetInternals.unitops.config.config_methods import get_property_fields 

7from common.config_types import PropertyType as PropertyTypeObj 

8from typing import TypedDict, List 

9from idaes_factory.endpoints import start_ml_training_event 

10 

11class ColumnMapping(TypedDict): 

12 portIndex: int 

13 propertyKey: str 

14 column: str 

15 

16 

17def create_column_mapping(flowsheet, model: int, inlet_mappings: List[ColumnMapping], outlet_mappings: List[ColumnMapping]): 

18 bulk_create_objects = [] 

19 ml_model = MLModel.objects.select_related("simulationObject", "flowsheet").get(id=model) 

20 property_set = ml_model.simulationObject.properties 

21 

22 mappings = [ 

23 (inlet_mappings, PropertyType.InletProperty), 

24 (outlet_mappings, PropertyType.OutletProperty), 

25 ] 

26 

27 for mappings_list, property_type in mappings: 

28 for i, mapping in enumerate(mappings_list): 

29 propertyKey = mapping.get("propertyKey") 

30 portIndex = mapping.get("portIndex") 

31 column = mapping.get("column") 

32 propertyInfo = None 

33 

34 # if unitop property 

35 if portIndex == custom_property_port_index: 

36 prop = PropertyTypeObj( 

37 displayName=propertyKey, 

38 value="", 

39 unitType="ratio", 

40 type="numeric", 

41 ) 

42 fields = get_property_fields(propertyKey, prop, property_set, 0) 

43 value = fields.pop("value") # since property value is separated from property info 

44 propertyInfo = PropertyInfo.objects.create(**fields, flowsheet=flowsheet) 

45 PropertyValue.objects.create(value=value, property=propertyInfo, enabled=property_type==PropertyType.InletProperty, flowsheet=flowsheet) 

46 

47 bulk_create_objects.append(MLColumnMapping( 

48 order=i, 

49 model=ml_model, 

50 column=column, 

51 portIndex=portIndex, 

52 propertyKey=propertyKey, 

53 propertyType=property_type, 

54 propertyInfo=propertyInfo, 

55 flowsheet=flowsheet 

56 )) 

57 

58 MLColumnMapping.objects.bulk_create(bulk_create_objects) 

59 ml_model.progress = 2 

60 ml_model.save(update_fields=["progress"]) 

61 

62 

63def train(user, model_instance: MLModel): 

64 """Queue ML training for the mapped model using its uploaded object-storage CSV.""" 

65 input_columns = MLColumnMapping.objects.filter(model=model_instance, propertyType=PropertyType.InletProperty) 

66 output_columns = MLColumnMapping.objects.filter(model=model_instance, propertyType=PropertyType.OutletProperty) 

67 input_labels = [item.column for item in input_columns] 

68 output_labels = [item.column for item in output_columns] 

69 

70 if not model_instance.csv_bucket or not model_instance.csv_object_key: 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true

71 raise ValueError("ML training requires a completed object-storage CSV upload.") 

72 

73 return start_ml_training_event( 

74 csv_bucket=model_instance.csv_bucket, 

75 csv_key=model_instance.csv_object_key, 

76 csv_delimiter=model_instance.csv_delimiter or None, 

77 input_labels=input_labels, 

78 output_labels=output_labels, 

79 user=user, 

80 flowsheet_id=model_instance.flowsheet.id, 

81 model_id=model_instance.id, 

82 )