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

41 statements  

« prev     ^ index     » next       coverage.py v7.10.7, created at 2026-05-13 02:47 +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 

33 # if unitop property 

34 if portIndex == custom_property_port_index: 

35 prop = PropertyTypeObj( 

36 displayName=propertyKey, 

37 value="", 

38 unitType="ratio", 

39 type="numeric", 

40 ) 

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

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

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

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

45 

46 bulk_create_objects.append(MLColumnMapping( 

47 order=i, 

48 model=ml_model, 

49 column=column, 

50 portIndex=portIndex, 

51 propertyKey=propertyKey, 

52 propertyType=property_type, 

53 flowsheet=flowsheet 

54 )) 

55 

56 MLColumnMapping.objects.bulk_create(bulk_create_objects) 

57 ml_model.progress = 2 

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

59 

60 

61def train(user, model_instance: MLModel): 

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

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

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

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

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

67 

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

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

70 

71 return start_ml_training_event( 

72 csv_bucket=model_instance.csv_bucket, 

73 csv_key=model_instance.csv_object_key, 

74 csv_delimiter=model_instance.csv_delimiter or None, 

75 input_labels=input_labels, 

76 output_labels=output_labels, 

77 user=user, 

78 flowsheet_id=model_instance.flowsheet.id, 

79 model_id=model_instance.id, 

80 )