spb/metrika Public

1//2//  ElasticNetTests.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation11import Testing12import ZQEngine13import ZQStats1415/// Elastic net vs glmnet fixtures. Both sides are iterative coordinate16/// descent (glmnet thresh 1e-15, ours 1e-14), so agreement is asserted at17/// 1e-6 relative — the documented tolerance for penalized solvers.18@Suite("Elastic net", .serialized)19struct ElasticNetTests {20    let fixtures: Fixtures21    let session: ZQSession2223    init() async throws {24        self.fixtures = try Fixtures()25        self.session = try ZQSession(discoverUserCommands: false)26        _ = try await session.execute("use \(fixtures.datasetURL.path)")27        _ = try await session.execute("gen log_rev = ln(revenue)")28    }2930    @Test("lasso matches glmnet and reproduces its selection")31    func lasso() async throws {32        let result = try await session.execute(33            "lasso log_rev price z1 z2 orders, lambda(0.05)"34        )35        expectClose(36            try #require(result.scalars["b_price"]),37            fixtures["lasso_b_price"], rtol: 1e-6, "b[price]"38        )39        expectClose(40            try #require(result.scalars["b_z1"]),41            fixtures["lasso_b_z1"], rtol: 1e-4, "b[z1] (tiny, near the threshold)"42        )43        expectClose(44            try #require(result.scalars["b__cons"]),45            fixtures["lasso_b_cons"], rtol: 1e-6, "intercept"46        )47        // glmnet zeroed z2 and orders — the selection must agree exactly.48        #expect(result.scalars["b_z2"] == 0)49        #expect(result.scalars["b_orders"] == 0)50        #expect(result.scalars["k_nonzero"] == 2)51    }5253    @Test("elastic net (alpha 0.4) matches glmnet")54    func elasticNet() async throws {55        let result = try await session.execute(56            "elasticnet log_rev price z1 z2 orders, alpha(0.4) lambda(0.02)"57        )58        expectClose(59            try #require(result.scalars["b_price"]),60            fixtures["enet_b_price"], rtol: 1e-6, "b[price]"61        )62        expectClose(63            try #require(result.scalars["b_z1"]),64            fixtures["enet_b_z1"], rtol: 1e-5, "b[z1]"65        )66        expectClose(67            try #require(result.scalars["b__cons"]),68            fixtures["enet_b_cons"], rtol: 1e-6, "intercept"69        )70        #expect(result.scalars["b_orders"] == 0)71    }7273    @Test("lambda 0 reproduces OLS")74    func lambdaZeroIsOLS() async throws {75        let ols = try await session.execute("reg log_rev price orders")76        let penalized = try await session.execute(77            "elasticnet log_rev price orders, alpha(1) lambda(0)"78        )79        expectClose(80            try #require(penalized.scalars["b_price"]),81            try #require(ols.scalars["b_price"]),82            rtol: 1e-8, "lambda 0 == OLS"83        )84    }8586    @Test("lambda at lambda_max zeroes everything")87    func lambdaMaxZeroes() throws {88        let n = 10089        let x = (0..<n).map { Double($0) }90        let y = x.map { 2 + 3 * $0 }91        let lambdaMax = ZQElasticNet.lambdaMax(92            y: y, predictors: [("x", x)], alpha: 193        )94        let result = try ZQElasticNet.fit(95            y: y, predictors: [("x", x)], alpha: 1, lambda: lambdaMax * 1.000000196        )97        #expect(result.nonzeroCount == 0)98        expectClose(result.intercept, y.reduce(0, +) / Double(n), "null intercept")99    }100101    @Test("predict works after a penalized fit; margins refuses")102    func postEstimation() async throws {103        _ = try await session.execute("lasso log_rev price z1 z2 orders, lambda(0.05)")104        _ = try await session.execute("predict lhat")105        let summary = try await session.execute("summarize lhat")106        #expect(summary.scalars["N"] == 57)107        await #expect(throws: ZQEngineError.self) {108            _ = try await session.execute("margins, dydx(price)")   // no VCE109        }110        _ = try await session.execute("drop lhat")111    }112113    @Test("missing lambda reports lambda_max as a hint")114    func lambdaHint() async throws {115        do {116            _ = try await session.execute("lasso log_rev price")117            Issue.record("expected an error")118        } catch let error as ZQEngineError {119            #expect(error.message.contains("lambda_max"))120        }121    }122}123