linreg: Include training loop in module

Author: jonnor

src/emlearn_linreg/Makefile

@@ -19,7 +19,7 @@ DIST_DIR := ../../dist/$(ARCH)_$(MPY_ABI_VERSION)
 MOD = emlearn_linreg
 
 # Source files (.c or .py)
-SRC = linreg.c
+SRC = linreg.c linreg.py
 
 # Include to get the rules for compiling and linking the module
 include $(MPY_DIR)/py/dynruntime.mk

src/emlearn_linreg/linreg.py

@@ -0,0 +1,47 @@
+
+
+def train(model, X_train, y_train,
+        max_iterations=100, 
+        tolerance=1e-6,
+        check_interval=10,
+        divergence_factor=10.0,
+        verbose=0,
+        ):
+    """
+    Simple training loop
+    """
+    prev_mse = float('inf')
+    
+    for iteration in range(max_iterations):
+        # Perform one gradient descent step
+        model.step(X_train, y_train)
+
+        # Only check progress at intervals        
+        if iteration % check_interval != 0:
+            continue
+
+        # Calculate current MSE
+        current_mse = model.score_mse(X_train, y_train)
+        change = abs(prev_mse - current_mse)
+
+        if verbose >= 2:
+            print(f'Iteration {iteration} mse={current_mse}')        
+
+        # Check convergence
+        converged = change < tolerance and iteration > check_interval * 2
+        
+        # Check divergence
+        diverged = current_mse > prev_mse * divergence_factor or not (current_mse == current_mse)  # NaN check
+        
+        if converged:
+            if verbose:
+                print(f"Converged at iteration {iteration}")
+            break
+            
+        if diverged:
+            if verbose:
+                print(f"Diverged at iteration {iteration}")
+            break
+            
+        prev_mse = current_mse
+

tests/test_linreg.py

@@ -2,15 +2,14 @@
 import emlearn_linreg
 import array
 
-# Create model: 4 features, alpha=0.1, l1_ratio=0.5, lr=0.01
 model = emlearn_linreg.new(4, 0.1, 0.5, 0.01)
 
 # Training data (float32 arrays)
 X = array.array('f', [1,2,3,4, 2,3,4,5])  # flattened
 y = array.array('f', [10, 15])
 
 # Train
-model.train(X, y, 1000, 1e-6)
+emlearn_linreg.train(model, X, y, max_iterations=100, tolerance=1e-6, verbose=0)
 
 # Predict
 prediction = model.predict(array.array('f', [1,2,3,4]))

tests/test_linreg_california.py

@@ -13,81 +13,6 @@ def load_npy_as_array(filename, dtype='f'):
     return shape, data
 
 
-def train_with_monitoring(model, X_train, y_train, max_iterations=1000, 
-                         tolerance=1e-6, check_interval=10, verbose=True):
-    """
-    Train model with progress monitoring using generator pattern.
-    
-    Yields: (iteration, mse, change, converged, diverged)
-    """
-    prev_mse = float('inf')
-    n_features = model.get_n_features()
-    
-    for iteration in range(max_iterations):
-        # Perform one gradient descent step
-        model.step(X_train, y_train)
-        
-        # Check progress at intervals
-        if iteration % check_interval == 0:
-            # Calculate current MSE
-            current_mse = model.score_mse(X_train, y_train)
-            change = abs(prev_mse - current_mse)
-            
-            # Check convergence
-            converged = change < tolerance and iteration > check_interval * 2
-            
-            # Check divergence
-            diverged = current_mse > prev_mse * 10.0 or not (current_mse == current_mse)  # NaN check
-            
-            # Yield progress
-            yield iteration, current_mse, change, converged, diverged
-            
-            if converged:
-                if verbose:
-                    print(f"Converged at iteration {iteration}")
-                break
-                
-            if diverged:
-                if verbose:
-                    print(f"Diverged at iteration {iteration}")
-                break
-                
-            prev_mse = current_mse
-        else:
-            # Yield minimal info for non-check iterations
-            yield iteration, None, None, False, False
-
-def train_model(model, X_train, y_train, max_iterations=1000, tolerance=1e-6, 
-                check_interval=10, verbose=True):
-    """
-    Simple training function with progress monitoring.
-    """
-    if verbose:
-        print(f"Training model for up to {max_iterations} iterations...")
-        print(f"Checking convergence every {check_interval} iterations")
-        print("Iter    MSE       Change    Status")
-        print("-" * 40)
-    
-    for iter_num, mse, change, converged, diverged in train_with_monitoring(
-        model, X_train, y_train, max_iterations, tolerance, check_interval, verbose=False
-    ):
-        
-        if verbose and mse is not None:  # Only print on check intervals
-            status = ""
-            if converged:
-                status = "CONVERGED"
-            elif diverged:
-                status = "DIVERGED"
-            
-            print(f"{iter_num:4d}    {mse:8.6f}  {change:8.6f}  {status}")
-        
-        if converged or diverged:
-            break
-    
-    if verbose:
-        print("Training completed.\n")
-
-
 def compare_regularization():
     """Compare different regularization settings."""
     print("\n=== Regularization Comparison ===")
@@ -98,7 +23,7 @@ def compare_regularization():
 
     n_features = X_shape[1]
     n_samples = min(500, y_shape[0])
-    
+
     X_subset = array.array('f', X_data[:n_samples * n_features])
     y_subset = array.array('f', y_data[:n_samples])
 
@@ -108,13 +33,14 @@ def compare_regularization():
         ("LASSO (L1)", 0.01, 1.0),
         ("ElasticNet", 0.01, 0.5),
     ]
-    
+
     for name, alpha, l1_ratio in configs:
         print(f"\nTesting {name} (alpha={alpha}, l1_ratio={l1_ratio}):")
 
         model = emlearn_linreg.new(n_features, alpha, l1_ratio, 0.01)
-        train_model(model, X_subset, y_subset, max_iterations=300, check_interval=50, verbose=False)
-        
+        emlearn_linreg.train(model, X_subset, y_subset,
+            max_iterations=2000, check_interval=50, verbose=True)
+
         # Calculate predictions and MSE
         mse = model.score_mse(X_subset, y_subset)
 
@@ -123,7 +49,6 @@ def compare_regularization():
 
         # Count non-zero weights (sparsity)
         non_zero = sum(1 for w in weights if abs(w) > 1e-6)
-        
         print(f"  MSE: {mse:.6f}")
         print(f"  Non-zero weights: {non_zero}/{n_features}")
         print(f"  Max weight magnitude: {max(abs(w) for w in weights):.6f}")