spb/metrika Public

1//2//  PredictTests.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import Foundation11import Testing12import ZQEngine1314@Suite("predict", .serialized)15struct PredictTests {16    let fixtures: Fixtures17    let session: ZQSession1819    init() async throws {20        self.fixtures = try Fixtures()21        self.session = try ZQSession(discoverUserCommands: false)22        _ = try await session.execute("use \(fixtures.datasetURL.path)")23        _ = try await session.execute("gen log_rev = ln(revenue)")24    }2526    @Test("xb + residuals reconstruct the response after regress")27    func fittedPlusResiduals() async throws {28        _ = try await session.execute("reg log_rev price")29        _ = try await session.execute("predict yhat")30        _ = try await session.execute("predict e, residuals")31        _ = try await session.execute("gen check = yhat + e - log_rev")32        let summary = try await session.execute("summarize check")33        #expect(abs(try #require(summary.scalars["min"])) < 1e-12)34        #expect(abs(try #require(summary.scalars["max"])) < 1e-12)35        // Missing prices → missing predictions.36        #expect(summary.scalars["N"] == 57)37        _ = try await session.execute("drop yhat e check")38    }3940    @Test("predict evaluates factor-variable models on all observations")41    func factorModel() async throws {42        _ = try await session.execute("reg log_rev price i.region")43        _ = try await session.execute("predict fhat")44        _ = try await session.execute("predict fe2, residuals")45        _ = try await session.execute("gen fcheck = fhat + fe2 - log_rev")46        let summary = try await session.execute("summarize fcheck")47        #expect(abs(try #require(summary.scalars["max"])) < 1e-12)48        _ = try await session.execute("drop fhat fe2 fcheck")49    }5051    @Test("logit default is predicted probability; fitted mean matches ybar")52    func logitProbability() async throws {53        _ = try await session.execute("logit purchase price")54        _ = try await session.execute("predict p")55        let pSummary = try await session.execute("summarize p")56        let ySummary = try await session.execute("summarize purchase if !missing(price)")57        // With an intercept, mean fitted probability equals the sample58        // mean of the outcome (logit score equation).59        expectClose(60            try #require(pSummary.scalars["mean"]),61            try #require(ySummary.scalars["mean"]),62            rtol: 1e-9, "mean fitted probability"63        )64        #expect(try #require(pSummary.scalars["min"]) > 0)65        #expect(try #require(pSummary.scalars["max"]) < 1)66        _ = try await session.execute("drop p")67    }6869    @Test("poisson default is the predicted mean count")70    func poissonMean() async throws {71        _ = try await session.execute("poisson orders price")72        _ = try await session.execute("predict mu")73        let muSummary = try await session.execute("summarize mu")74        let ySummary = try await session.execute("summarize orders if !missing(price)")75        // Poisson with intercept: mean fitted count equals ybar.76        expectClose(77            try #require(muSummary.scalars["mean"]),78            try #require(ySummary.scalars["mean"]),79            rtol: 1e-9, "mean fitted count"80        )81        _ = try await session.execute("drop mu")82    }8384    @Test("predict before any estimation errors clearly")85    func requiresEstimation() async throws {86        let fresh = try ZQSession(discoverUserCommands: false)87        _ = try await fresh.execute("use \(fixtures.datasetURL.path)")88        await #expect(throws: ZQEngineError.self) {89            _ = try await fresh.execute("predict yhat")90        }91    }9293    @Test("statistic options are validated against the model kind")94    func statisticValidation() async throws {95        _ = try await session.execute("reg log_rev price")96        await #expect(throws: ZQEngineError.self) {97            _ = try await session.execute("predict p, pr")   // pr needs logit/probit98        }99        await #expect(throws: ZQEngineError.self) {100            _ = try await session.execute("predict m, n")    // n needs poisson101        }102    }103}104