spb/metrika Public

1//2//  Planner.swift3//  Metrika4//5//  Author:  Simon-Pierre Boucher6//  Contact: contact@spboucher.ai7//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.8//910import ZQData11import ZQParser1213/// Execution backend. The planner chooses automatically; the user never14/// does (CLAUDE.md §5).15public enum ZQBackend: String, Equatable, Sendable {16    case cpu, gpu, hybrid17}1819/// A planned command: the parsed AST plus dispatch decision. Factor20/// variables expand here (plan time), not at parse time.21public struct ZQExecutionPlan: Equatable, Sendable {22    public var command: ZQCommand23    public var backend: ZQBackend2425    public init(command: ZQCommand, backend: ZQBackend) {26        self.command = command27        self.backend = backend28    }29}3031public struct ZQPlanner: Sendable {32    /// Single-estimation row-count threshold below which the CPU LAPACK33    /// path wins (heuristic 1, tuned by Tests/Bench).34    public var cpuRowThreshold: Int35    /// Independent-replicate threshold at which batched GPU solves win36    /// (heuristic 2).37    public var gpuReplicateThreshold: Int38    /// Whether a GPU backend is available at all. The pure-Swift Philox39    /// reference keeps plans reproducible even when this is false.40    public var gpuAvailable: Bool4142    public init(43        cpuRowThreshold: Int = 5_000_000,44        gpuReplicateThreshold: Int = 500,45        gpuAvailable: Bool = false46    ) {47        self.cpuRowThreshold = cpuRowThreshold48        self.gpuReplicateThreshold = gpuReplicateThreshold49        self.gpuAvailable = gpuAvailable50    }5152    public func plan(_ command: ZQCommand, rowCount: Int) -> ZQExecutionPlan {53        ZQExecutionPlan(command: command, backend: backend(for: command, rowCount: rowCount))54    }5556    private func backend(for command: ZQCommand, rowCount: Int) -> ZQBackend {57        guard gpuAvailable else { return .cpu }5859        // Heuristic 2: ≥ threshold independent replicates → batched GPU.60        // (permute joins once its GPU argsort path lands; the permutation61        // definition is already argsort-of-Philox-keys for that reason.)62        if command.verb == "bootstrap" {63            if let reps = command.option("reps")?.firstArgument,64               let count = Int(reps), count >= gpuReplicateThreshold {65                return .gpu66            }67        }6869        // Heuristic 1: single estimation under the row threshold → CPU.70        if rowCount < cpuRowThreshold { return .cpu }7172        // Large single estimations stay on CPU until the batched GPU73        // estimators land (v0.2); MCMC-style workloads will plan .hybrid.74        return .cpu75    }76}7778/// Plan-time factor-variable expansion (`i.region` → one indicator column79/// per level, base level omitted).80public enum ZQFactorExpansion {8182    /// Distinct sorted levels of a numeric variable, ignoring missing.83    public static func levels(values: [Double], missing: [Bool]) -> [Double] {84        var seen = Set<Double>()85        for i in 0..<values.count where !missing[i] {86            seen.insert(values[i])87        }88        return seen.sorted()89    }9091    /// Expands one factor spec into named indicator columns against the92    /// given data. The lowest level is the omitted base (Stata default).93    /// The numeric level rides along so post-estimation commands can94    /// recompute the indicator on other observations.95    public static func expand(96        name: String, values: [Double], missing: [Bool]97    ) -> [(name: String, level: Double, values: [Double])] {98        let allLevels = levels(values: values, missing: missing)99        guard allLevels.count > 1 else { return [] }100        return allLevels.dropFirst().map { level in101            let rendered = level == level.rounded()102                ? String(Int(level))103                : String(level)104            let indicator = values.enumerated().map { index, value in105                (!missing[index] && value == level) ? 1.0 : 0.0106            }107            return (name: "\(rendered).\(name)", level: level, values: indicator)108        }109    }110}111