spb/forge-studio Public

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// The main loss chart: raw train (low opacity) + bias-corrected EMA + val4// points, log/linear Y, best-val rule, hover crosshair with a full callout5// (step, train, EMA, val, ppl, lr, tok/s), pinch-zoom + horizontal pan with6// a "follow live" pill when zoomed during a running train, double-click to7// reset. Data arrives pre-downsampled from MetricsStore.8import Charts9import SwiftUI1011// Shared X-window + crosshair state: the main chart drives it, secondary12// charts observe it, so zooming/hovering is synchronized everywhere.13@Observable14final class ChartXState {15    var visibleLength: Double?  // nil = fit all16    var scrollX: Double = 017    var hoverX: Double?18}1920struct TrainingChartView: View {21    let snapshot: MetricsStore.Snapshot?22    let isLive: Bool23    @Binding var smoothing: Double24    @Binding var logScale: Bool25    var xState: ChartXState26    var checkpointSteps: [Int] = []27    var onCheckpointTap: ((Int) -> Void)?2829    @State private var followLive = true30    @State private var magnifyBase: Double?3132    private var visibleLength: Double? {33        get { xState.visibleLength }34        nonmutating set { xState.visibleLength = newValue }35    }36    private var scrollX: Double {37        get { xState.scrollX }38        nonmutating set { xState.scrollX = newValue }39    }40    private var hoverX: Double? {41        get { xState.hoverX }42        nonmutating set { xState.hoverX = newValue }43    }4445    private var maxX: Double { snapshot?.train.last?.x ?? 1 }46    private var minX: Double { snapshot?.train.first?.x ?? 0 }4748    var body: some View {49        VStack(alignment: .leading, spacing: 8) {50            controls51            chart52                .frame(minHeight: 320)53        }54    }5556    private var controls: some View {57        HStack {58            Text("Loss").font(.title3.weight(.semibold))59            if let h = hoverInfo {60                calloutText(h)61            }62            Spacer()63            if isLive && visibleLength != nil && !followLive {64                Button {65                    followLive = true66                } label: {67                    Label("suivre ⏵", systemImage: "forward.fill")68                        .font(.caption)69                }70                .buttonStyle(.borderedProminent)71                .controlSize(.small)72            }73            Toggle("log Y", isOn: $logScale).toggleStyle(.checkbox)74            Button("fit") {75                visibleLength = nil76                followLive = true77            }78            .controlSize(.small)79            HStack(spacing: 4) {80                Text("EMA")81                Slider(value: $smoothing, in: 0...0.99).frame(width: 110)82                Text(String(format: "%.2f", smoothing)).monospacedDigit()83            }84            .font(.caption)85        }86    }8788    private var chart: some View {89        Chart {90            ForEach(snapshot?.train ?? [], id: \.x) { p in91                LineMark(x: .value("step", p.x), y: .value("loss", p.y),92                         series: .value("s", "train"))93                    .foregroundStyle(.blue.opacity(0.22))94            }95            ForEach(snapshot?.trainEMA ?? [], id: \.x) { p in96                LineMark(x: .value("step", p.x), y: .value("ema", p.y),97                         series: .value("s", "train EMA"))98                    .foregroundStyle(.blue)99            }100            ForEach(snapshot?.val ?? [], id: \.x) { p in101                LineMark(x: .value("step", p.x), y: .value("val", p.y),102                         series: .value("s", "val"))103                    .foregroundStyle(.orange)104                PointMark(x: .value("step", p.x), y: .value("val", p.y))105                    .foregroundStyle(.orange)106                    .symbolSize(24)107            }108            if let best = snapshot?.bestVal {109                RuleMark(y: .value("best", best.loss))110                    .foregroundStyle(.orange.opacity(0.4))111                    .lineStyle(.init(lineWidth: 1, dash: [4, 4]))112                    .annotation(position: .topTrailing) {113                        Text(String(format: "best val %.3f @ %d", best.loss, best.step))114                            .font(.caption2)115                            .foregroundStyle(.orange)116                    }117            }118            ForEach(checkpointSteps, id: \.self) { step in119                RuleMark(x: .value("ckpt", Double(step)))120                    .foregroundStyle(.green.opacity(0.25))121                    .lineStyle(.init(lineWidth: 1, dash: [2, 4]))122                    .annotation(position: .top, alignment: .leading) {123                        Image(systemName: "externaldrive.fill")124                            .font(.system(size: 7))125                            .foregroundStyle(.green.opacity(0.6))126                    }127            }128            if let x = hoverX {129                RuleMark(x: .value("hover", x))130                    .foregroundStyle(.secondary.opacity(0.5))131                    .lineStyle(.init(lineWidth: 1))132            }133        }134        .modifier(LogScaleModifier(enabled: logScale))135        .modifier(ScrollModifier(136            visibleLength: visibleLength,137            scrollX: Binding(get: { xState.scrollX },138                             set: { xState.scrollX = $0 })))139        .chartLegend(.visible)140        .chartOverlay { proxy in141            GeometryReader { geo in142                Color.clear143                    .contentShape(Rectangle())144                    .onContinuousHover { phase in145                        switch phase {146                        case .active(let pt):147                            let plot = geo[proxy.plotFrame!]148                            hoverX = proxy.value(atX: pt.x - plot.origin.x, as: Double.self)149                        case .ended:150                            hoverX = nil151                        }152                    }153                    .gesture(magnify)154                    .onTapGesture(count: 2) {155                        visibleLength = nil156                        followLive = true157                    }158                    .onTapGesture(count: 1) {159                        // Click near a checkpoint rule selects it (M4:160                        // "clickable → jumps to that checkpoint").161                        guard let x = hoverX,162                              let nearest = checkpointSteps163                                  .min(by: { abs(Double($0) - x) < abs(Double($1) - x) }),164                              abs(Double(nearest) - x) < max(4, (maxX - minX) * 0.01)165                        else { return }166                        onCheckpointTap?(nearest)167                    }168            }169        }170        .onChange(of: xState.scrollX) { old, new in171            // A user-initiated pan while live breaks auto-follow.172            if isLive, followLive, let len = visibleLength,173               abs(new - (maxX - len)) > len * 0.05 {174                followLive = false175            }176        }177        .onChange(of: snapshot?.count ?? 0) {178            if isLive, followLive, let len = visibleLength {179                scrollX = max(minX, maxX - len)180            }181        }182    }183184    private var magnify: some Gesture {185        MagnifyGesture()186            .onChanged { g in187                let base = magnifyBase ?? (visibleLength ?? (maxX - minX))188                magnifyBase = base189                let span = maxX - minX190                let newLen = min(max(base / g.magnification, span * 0.01), span)191                visibleLength = newLen < span * 0.999 ? newLen : nil192                if followLive, let len = visibleLength {193                    scrollX = max(minX, maxX - len)194                }195            }196            .onEnded { _ in magnifyBase = nil }197    }198199    // Nearest-point lookup for the callout (train series is x-sorted).200    private struct HoverInfo {201        var step: Int202        var train: Double?203        var ema: Double?204        var val: Double?205    }206207    private var hoverInfo: HoverInfo? {208        guard let x = hoverX, let snap = snapshot, !snap.train.isEmpty else { return nil }209        func nearest(_ s: [Downsampler.XY]) -> Downsampler.XY? {210            guard !s.isEmpty else { return nil }211            var lo = 0, hi = s.count - 1212            while lo < hi {213                let mid = (lo + hi) / 2214                if s[mid].x < x { lo = mid + 1 } else { hi = mid }215            }216            if lo > 0, abs(s[lo - 1].x - x) < abs(s[lo].x - x) { lo -= 1 }217            return s[lo]218        }219        let t = nearest(snap.train)220        return HoverInfo(step: Int(t?.x ?? x),221                         train: t?.y,222                         ema: nearest(snap.trainEMA)?.y,223                         val: nearest(snap.val).flatMap {224                             abs($0.x - x) < max(8, (maxX - minX) * 0.02) ? $0.y : nil225                         })226    }227228    private func calloutText(_ h: HoverInfo) -> some View {229        HStack(spacing: 10) {230            Text("step \(h.step)").foregroundStyle(.secondary)231            if let t = h.train {232                Text(String(format: "train %.4f", t)).foregroundStyle(.blue)233                Text(String(format: "ppl %.1f", exp(t))).foregroundStyle(.secondary)234            }235            if let e = h.ema {236                Text(String(format: "EMA %.4f", e)).foregroundStyle(.blue)237            }238            if let v = h.val {239                Text(String(format: "val %.4f", v)).foregroundStyle(.orange)240            }241        }242        .font(.caption.monospacedDigit())243        .padding(.horizontal, 8)244        .padding(.vertical, 3)245        .background(.quaternary.opacity(0.5), in: Capsule())246    }247}248249// Conditional modifiers kept out of the Chart body so the result builder250// stays type-checkable.251private struct LogScaleModifier: ViewModifier {252    let enabled: Bool253    func body(content: Content) -> some View {254        if enabled { content.chartYScale(type: .log) } else { content }255    }256}257258// Read-only link to the main chart's X window (secondary charts).259struct LinkedScrollModifier: ViewModifier {260    var xState: ChartXState261    func body(content: Content) -> some View {262        if let len = xState.visibleLength {263            content264                .chartScrollableAxes(.horizontal)265                .chartXVisibleDomain(length: len)266                .chartScrollPosition(x: Binding(get: { xState.scrollX },267                                                set: { _ in })) // main chart drives268        } else {269            content270        }271    }272}273274private struct ScrollModifier: ViewModifier {275    let visibleLength: Double?276    @Binding var scrollX: Double277    func body(content: Content) -> some View {278        if let len = visibleLength {279            content280                .chartScrollableAxes(.horizontal)281                .chartXVisibleDomain(length: len)282                .chartScrollPosition(x: $scrollX)283        } else {284            content285        }286    }287}288