SIENTIAPDE-1321: Added equation related methods

model_manager/utils/models/train_model_result.py

@@ -28,6 +28,8 @@ class TrainModelResult:
         mse_val (float | None): The Mean Squared Error (MSE) of the predictions. Default is None.
         mae_val (float | None): The Mean Absolute Error (MAE) of the predictions. Default is None.
         r2_val (float | None): The R-squared (R²) value of the predictions. Default is None.
+        equation (dict | None): The equation of the model. Default is None.
+        equation_path (str | None): The path to the equation file. Default is None.
         run_name (str | None): The name of the MLFlow run. Default is None.
         report_path (str | None): The path to the generated HTML report file. Default is None.
         train_data_path (str | None): The path to the training dataset CSV file. Default is None.
@@ -47,6 +49,8 @@ class TrainModelResult:
     mse_val: float | None = None
     mae_val: float | None = None
     r2_val: float | None = None
+    equation: dict | None = None
+    equation_path: str | None = None
     run_name: str | None = None
     report_path: str | None = None
     train_data_path: str | None = None

model_manager/utils/repository/model_repository.py

@@ -9,6 +9,7 @@ and logging model runs to MLFlow.
 
 """
 
+import json
 import os
 import shutil
 import warnings
@@ -223,6 +224,10 @@ class ModelRepository:
             self.model_serving.log_artifact(data.report_path)
             self.model_serving.log_artifact(data.train_data_path)
             self.model_serving.log_artifact(data.test_data_path)
+            
+            # Log equation artifact if available
+            if data.equation_path and path.exists(data.equation_path):
+                self.model_serving.log_artifact(data.equation_path)
 
     def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]:
         """
@@ -412,6 +417,12 @@ class ModelRepository:
             data.test_data_path = path.join(data.run_dir, 'test_data.csv')
             current_data.to_csv(data.test_data_path, index=False)
 
+            # Save equation as JSON
+            if data.equation is not None:
+                data.equation_path = path.join(data.run_dir, 'model_equation.json')
+                with open(data.equation_path, 'w', encoding='utf-8') as f:
+                    json.dump(data.equation, f, indent=2, ensure_ascii=False)
+
             return data
         except ValueError as e:
             error_msg = f'Failed to convert data to float64 for report generation: {str(e)}'

model_manager/utils/repository/training_repository.py

@@ -175,6 +175,10 @@ class TrainingRepository:
         )
 
         tmr.r2_val = round(r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)), 2)
+        
+        # Extract model equation
+        tmr.equation = self._extract_model_equation(tmr.regr, params)
+        
         self.logger.info(
             f'Model metrics calculated successfully - experiment run id: {params.experiment_run_id}'
         )
@@ -249,3 +253,48 @@ class TrainingRepository:
             scaler_params={} if params.use_scaler else None,
             ar_var=params.target_variable if params.include_ar else None,
         )
+
+    def _extract_model_equation(self, regr: LinearRegressionModel, params: TrainModelParams) -> dict:
+        """
+        Extract the linear regression equation coefficients and create equation metadata.
+        
+        This method extracts the coefficients and intercept from the trained model
+        and creates a structured dictionary containing the equation information
+        for serialization as JSON artifact.
+        
+        Args:
+            regr: Trained LinearRegressionModel object
+            params: Training parameters containing variable information
+            
+        Returns:
+            dict: Equation metadata containing:
+                - target_variable: Name of the target variable
+                - coefficients: Dictionary mapping variable names to coefficients
+                - intercept: Model intercept value
+                - equation_string: Human-readable equation string
+                - latex_equation: LaTeX formatted equation
+        """
+        coefficients = regr.regr.coef_
+        intercept = regr.regr.intercept_
+        
+        # Create coefficients dictionary
+        coefficients_dict = {}
+        for i, var in enumerate(params.variable_columns):
+            coefficients_dict[var] = float(coefficients[i])
+        
+        # Create equation string
+        equation_parts = [f"{coef:.6f} * {var}" for var, coef in coefficients_dict.items()]
+        equation_string = f"{params.target_variable} = {intercept:.6f} + " + " + ".join(equation_parts)
+        
+        # Create LaTeX equation
+        latex_parts = [f"{coef:.6f} \\cdot {var}" for var, coef in coefficients_dict.items()]
+        latex_equation = f"{params.target_variable} = {intercept:.6f} + " + " + ".join(latex_parts)
+        
+        return {
+            'target_variable': params.target_variable,
+            'coefficients': coefficients_dict,
+            'intercept': float(intercept),
+            'equation_string': equation_string,
+            'latex_equation': latex_equation,
+            'model_type': 'Linear Regression'
+        }